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 14                                   #
 13#  In this script we demonstrate how to use the user_basis to construct   #
 14#  a spin-1/2 Hamiltonian on a restricted Hilbert space where a spin-up   #
 15#  on a given lattice site must be preceded and succeeded by a spin-down. #
 16###########################################################################
 17from quspin.operators import hamiltonian
 18from quspin.basis import spin_basis_1d  # Hilbert space spin basis_1d
 19from quspin.basis.user import user_basis  # Hilbert space user basis
 20from quspin.basis.user import (
 21    pre_check_state_sig_32,
 22    op_sig_32,
 23    map_sig_32,
 24)  # user basis data types
 25from numba import carray, cfunc  # numba helper functions
 26from numba import uint32, int32  # numba data types
 27import numpy as np
 28
 29#
 30N = 10  # total number of lattice sites
 31
 32
 33#
 34######  function to call when applying operators
 35@cfunc(op_sig_32, locals=dict(s=int32, b=uint32))
 36def op(op_struct_ptr, op_str, ind, N, args):
 37    # using struct pointer to pass op_struct_ptr back to C++ see numba Records
 38    op_struct = carray(op_struct_ptr, 1)[0]
 39    err = 0
 40    ind = N - ind - 1  # convention for QuSpin for mapping from bits to sites.
 41    s = (((op_struct.state >> ind) & 1) << 1) - 1
 42    b = 1 << ind
 43    #
 44    if op_str == 120:  # "x" is integer value 120 (check with ord("x"))
 45        op_struct.state ^= b
 46    elif op_str == 121:  # "y" is integer value 120 (check with ord("y"))
 47        op_struct.state ^= b
 48        op_struct.matrix_ele *= 1.0j * s
 49    elif op_str == 122:  # "z" is integer value 120 (check with ord("z"))
 50        op_struct.matrix_ele *= s
 51    else:
 52        op_struct.matrix_ele = 0
 53        err = -1
 54    #
 55    return err
 56
 57
 58#
 59op_args = np.array([], dtype=np.uint32)
 60
 61
 62#
 63######  function to filter states/project states out of the basis
 64#
 65@cfunc(
 66    pre_check_state_sig_32,
 67    locals=dict(s_shift_left=uint32, s_shift_right=uint32),
 68)
 69def pre_check_state(s, N, args):
 70    """imposes that that a bit with 1 must be preceded and followed by 0,
 71    i.e. a particle on a given site must have empty neighboring sites.
 72    #
 73    Works only for lattices of up to N=32 sites (otherwise, change mask)
 74    #
 75    """
 76    mask = 0xFFFFFFFF >> (32 - N)  # works for lattices of up to 32 sites
 77    # cycle bits left by 1 periodically
 78    s_shift_left = ((s << 1) & mask) | ((s >> (N - 1)) & mask)
 79    #
 80    # cycle bits right by 1 periodically
 81    s_shift_right = ((s >> 1) & mask) | ((s << (N - 1)) & mask)
 82    #
 83    return (((s_shift_right | s_shift_left) & s)) == 0
 84
 85
 86#
 87pre_check_state_args = None
 88#
 89######  construct user_basis
 90# define maps dict
 91maps = dict()  # no symmetries to apply.
 92# define op_dict
 93op_dict = dict(op=op, op_args=op_args)
 94# define pre_check_state
 95pre_check_state = (
 96    pre_check_state,
 97    pre_check_state_args,
 98)  # None gives a null pointer to args
 99# create user basis
100basis = user_basis(
101    np.uint32,
102    N,
103    op_dict,
104    allowed_ops=set("xyz"),
105    sps=2,
106    pre_check_state=pre_check_state,
107    Ns_block_est=300000,
108    **maps,
109)
110# print basis
111print(basis)
112#
113###### construct Hamiltonian
114# site-coupling lists
115h_list = [[1.0, i] for i in range(N)]
116# operator string lists
117static = [
118    ["x", h_list],
119]
120# compute Hamiltonian, no checks have been implemented
121no_checks = dict(check_symm=False, check_pcon=False, check_herm=False)
122H = hamiltonian(static, [], basis=basis, dtype=np.float64, **no_checks)