1#
  2import sys, os
  3
  4os.environ["KMP_DUPLICATE_LIB_OK"] = (
  5    "True"  # uncomment this line if omp error occurs on OSX for python 3
  6)
  7os.environ["OMP_NUM_THREADS"] = "1"  # set number of OpenMP threads to run in parallel
  8os.environ["MKL_NUM_THREADS"] = "1"  # set number of MKL threads to run in parallel
  9#
 10
 11###########################################################################
 12#                            example 16                                   #
 13#  In this script we demonstrate how to apply the user_basis to reduce    #
 14#  user-imported arrays of bases states (in integer representation)       #
 15#  to user-defined symmetry-reduced subspaces.                            #
 16###########################################################################
 17from quspin.basis import (
 18    spin_basis_1d,
 19    spin_basis_general,
 20)  # Hilbert space spin basis_1d
 21from quspin.basis.user import user_basis  # Hilbert space user basis
 22from quspin.basis.user import (
 23    next_state_sig_32,
 24    op_sig_32,
 25    map_sig_32,
 26)  # user basis data types
 27from numba import carray, cfunc  # numba helper functions
 28from numba import uint32, int32  # numba data types
 29import numpy as np
 30from scipy.special import comb
 31
 32#
 33#####
 34N_half = 10  # number of sites for each leg of the ladder
 35N = 2 * N_half  # total number of lattice sites
 36
 37
 38#
 39def make_basis(N_half):
 40    """Generates a list of integers to represent external, user-imported basis"""
 41    old_basis = spin_basis_general(N_half, m=0)
 42    #
 43    states = old_basis.states
 44    shift_states = np.left_shift(states, N_half)
 45    #
 46    shape = states.shape + states.shape
 47    #
 48    states_b = np.broadcast_to(states, shape)
 49    shift_states_b = np.broadcast_to(shift_states, shape)
 50    # this does the kronecker sum in a more memory efficient way.
 51    return (states_b + shift_states_b.T).ravel()
 52
 53
 54#
 55external_basis = make_basis(N_half)
 56#
 57Np = ()  # dummy argument, could be any value (particle conservation should've been
 58
 59
 60# taken into account when constructing the basis object)
 61#
 62######  function to call when applying operators
 63@cfunc(op_sig_32, locals=dict(s=int32, b=uint32))
 64def op(op_struct_ptr, op_str, ind, N, args):
 65    # using struct pointer to pass op_struct_ptr back to C++ see numba Records
 66    op_struct = carray(op_struct_ptr, 1)[0]
 67    err = 0
 68    ind = N - ind - 1  # convention for QuSpin for mapping from bits to sites.
 69    s = (((op_struct.state >> ind) & 1) << 1) - 1
 70    b = 1 << ind
 71    #
 72    if op_str == 120:  # "x" is integer value 120 (check with ord("x"))
 73        op_struct.state ^= b
 74    elif op_str == 121:  # "y" is integer value 120 (check with ord("y"))
 75        op_struct.state ^= b
 76        op_struct.matrix_ele *= 1.0j * s
 77    elif op_str == 122:  # "z" is integer value 120 (check with ord("z"))
 78        op_struct.matrix_ele *= s
 79    else:
 80        op_struct.matrix_ele = 0
 81        err = -1
 82    #
 83    return err
 84
 85
 86#
 87op_args = np.array([], dtype=np.uint32)
 88
 89
 90#
 91######  function to read user-imported basis into QuSpin
 92#
 93# function to call when generating next_state
 94@cfunc(next_state_sig_32)
 95def next_state(s, counter, N, args):
 96    # return pre-calculated basis state.
 97    # add one to counter because the first state is already checked.
 98    return args[counter + 1]  # = basis
 99
100
101#
102next_state_args = (
103    external_basis  # this has to be an array of same dtype as the user_basis
104)
105
106
107#
108class function_wrapper(object):
109    """
110    This class provides a wrapper for the user-imported basis,
111    as well as the functions required for the `user_basis` functionality.
112    #
113    This is needed to easily pass parameters (defined as class attributes) to the
114    functions `get_so_pcon()` and `get_Ns_pcon`.
115    """
116
117    def __init__(self, basis):
118        self.basis = basis
119
120    #
121    # python function to calculate the starting state to generate the particle conserving basis
122    def get_s0_pcon(self, N, Np):
123        """calculates the starting state to generate the particle conserving basis."""
124        # ignore input arguments as basis is already calculated.
125        return self.basis[0]
126
127    #
128    # python function to calculate the size of the particle-conserved basis,
129    # i.e. BEFORE applying pre_check_state and symmetry maps
130    def get_Ns_pcon(self, N, Np):
131        """calculates the size of the particle conservation basis (ignoring symmetries at this stage)."""
132        # ignore input arguments as basis is already calculated.
133        return self.basis.size
134
135
136#
137######  define symmetry maps
138#
139if N_half != 10:
140    print(
141        "symmetry masks are hard-coded and work only for N=10; \n\
142To do a different system size, it is required to update the masks accordingly.\n\
143exiting..."
144    )
145    exit()
146
147
148#
149@cfunc(map_sig_32)
150def translation(x, N, sign_ptr, args):
151    # bit permutation target bits: 1 2 3 4 5 6 7 8 9 0 11 12 13 14 15 16 17 18 19 10
152    # code generated here: http://programming.sirrida.de/calcperm.php
153    # only works for N_half=10
154    return ((x & 0x0007FDFF) << 1) | ((x & 0x00080200) >> 9)
155
156
157T_args = np.array([], dtype=np.uint32)
158
159
160#
161@cfunc(map_sig_32)
162def parity(x, N, sign_ptr, args):
163    # bit permutation target bits: 9 8 7 6 5 4 3 2 1 0 19 18 17 16 15 14 13 12 11 10
164    # code generated here: http://programming.sirrida.de/calcperm.php
165    # only works for N_half=10
166    return (
167        ((x & 0x00004010) << 1)
168        | ((x & 0x00002008) << 3)
169        | ((x & 0x00001004) << 5)
170        | ((x & 0x00000802) << 7)
171        | ((x & 0x00000401) << 9)
172        | ((x & 0x00080200) >> 9)
173        | ((x & 0x00040100) >> 7)
174        | ((x & 0x00020080) >> 5)
175        | ((x & 0x00010040) >> 3)
176        | ((x & 0x00008020) >> 1)
177    )
178
179
180P_args = np.array([], dtype=np.uint32)
181#
182######  construct user_basis
183# define maps dict
184maps = dict(
185    T_block=(translation, N_half, 0, T_args),
186    P_block=(parity, 2, 0, P_args),
187)
188# define particle conservation and op dicts
189FW = function_wrapper(external_basis)
190pcon_dict = dict(
191    Np=Np,
192    next_state=next_state,
193    next_state_args=next_state_args,
194    get_Ns_pcon=FW.get_Ns_pcon,
195    get_s0_pcon=FW.get_s0_pcon,
196)
197op_dict = dict(op=op, op_args=op_args)
198# create user basis
199basis = user_basis(
200    np.uint32, N, op_dict, allowed_ops=set("xyz"), sps=2, pcon_dict=pcon_dict, **maps
201)
202# print basis
203print(basis)