<commit_before>import fractions
from itertools import chain
import numpy as np
from spinsys.exceptions import NotFoundError
import functools
from hamiltonians.triangular_lattice_model import SiteVector


def bond_list(Nx, Ny):
    N = Nx * Ny
    bonds = []
    vec = SiteVector((0, 0), Nx, Ny)
    for i in range(N):
        bonds.append((vec, vec.xhop(1)))
        bonds.append((vec, vec.yhop(1)))
        bonds.append((vec, vec.xhop(-1).yhop(1)))
        vec = vec.next_site()
    return bonds


def zero_momentum_states(N):
    def genfunc(i):
        bloch_func = [i]
        j = i
        while True:
            j = (j * 2) % maxdec
            if not j == i:
                bloch_func.append(j)
                sieve[j] = False
            else:
                break
        bloch_func.sort()
        try:
            l = len(bloch_func)
            states[l].append(bloch_func)
        except KeyError:
            states[l] = [bloch_func]

    sieve = np.ones(2 ** N)
    sieve[0] = sieve[-1] = 0
    maxdec = 2 ** N - 1
    states = {1: [[0], [maxdec]]}
    for i in range(maxdec + 1):
        if sieve[i]:
            genfunc(i)
    return states


def bloch_states(N, k):
    zero_k_states = zero_momentum_states(N)
    if k > N // 2:
        raise NotFoundError

    if k == 0:
        states = list(chain(*zero_k_states.values()))
    elif k > 0 or (k < 0 and ((not -k == N // 2) or (not N % 2 == 0))):
        states = []
        for l in zero_k_states.keys():
            period = N // fractions.gcd(N, k)
            if l % period == 0:
                states.extend(zero_k_states[l])
    else:
        raise NotFoundError
    return states


@functools.lru_cache(maxsize=None)
def generate_dec_to_ind_dictionary(N, k):
    states = sorted(bloch_states(N, k))
    dec_to_ind = {}
    for i, state in enumerate(states):
        state_len = len(state)
        for num in state:
            dec_to_ind[num] = (i, state_len)
    return states, dec_to_ind


def all_bloch_states(N):
    states = {}
    max_k = N // 2
    for k in range(-max_k, max_k + 1):
        try:
            states[k] = bloch_states(N, k)
        except NotFoundError:
            continue
    return states


@functools.lru_cache(maxsize=None)
def count_same_spins(N, state, l, return_loc=False):
    # edge cases that break down under the kind of binary operation we
    #  use later
    if not (state == 2 ** N - 1 or state == 0):
        # subtraction does a bitwise flip of 1's and 0's. We need this
        #  because subsequent operations are insensitive to patterns of 0's
        bit_inverted_state = 2 ** N - 1 - state
        # the mod operator accounts for periodic boundary conditions
        couplings = 2 ** np.arange(N) + 2 ** (np.arange(l, N + l) % N)
        nup = ndown = 0
        ups = None
        downs = None
        for i in couplings:
            # if a state is unchanged under bitwise OR with the given number,
            #  it has two 1's at those two sites
            if state | i == state:
                nup += 1
                ups = i
            # cannot be merged with the previous statement because bit-flipped
            #  numbers using subtraction are possibly also bit-shifted
            if bit_inverted_state | i == bit_inverted_state:
                ndown += 1
                downs = i
    elif state == 0:
        nup, ndown = 0, N
    else:
        nup, ndown = N, 0  # when state is all 0's

    if return_loc:
        return nup, ndown, ups, downs
    else:
        return nup, ndown


def H_z_elements(N, k, i, l):
    state = generate_dec_to_ind_dictionary(N, k)[0][i][0]
    nup, ndown = count_same_spins(N, state, l)
    same_dir = nup + ndown
    diff_dir = N - same_dir
    return 0.25 * (same_dir - diff_dir)


def H_pm_elements(N, k, i, l):
    states, dec_to_ind = generate_dec_to_ind_dictionary(N, k)
    state = states[i][0]
    i_len = len(states[i])
    bit1 = 2 ** np.arange(N)
    bit2 = 2 ** (np.arange(l, N + 1) % N)
    count = 0  # count of up-downs and down-ups separated by l
    for b1, b2 in zip(bit1, bit2):
        if (state | b1 == state) and (not state | b2 == state):
            count += 1
        if (not state | b1 == state) and (state | b2 == state):
            count += 1
            pm = (b1, b2)
    # bit-flip the state once to see what it is coupled to. We don't need
    #  to do it more than once since such flips will always get to
    #  constituent states of the same bloch state
    connected_state = state + pm[0] - pm[1]
    j, j_len = dec_to_ind[connected_state]
    val = i_len * count / j_len * np.sqrt(j_len / i_len)
    return val, j


def H_ppmm_elements(N, k, i, l):
    state = generate_dec_to_ind_dictionary(N, k)[0][i][0]
    nup, ndown, ups, downs = count_same_spins(N, state, l, return_loc=True)
    if ups is not None:
        connected_state = state - ups
    elif downs is not None:
        connected_state = state + downs
    else:
        raise NotFoundError
    j, j_len = dec_to_ind[connected_state]


# N = 6
# states = zero_momentum_states(N)
# totlen = 0
# for key, val in states.items():
#     print('\nlength = {}\ncount = {}\n'.format(key, len(val)))
#     for s in val:
#         print(s)
#         totlen += len(s)
# print('\nSanity check. Total number of product states:', totlen)
# print('Total number of zero momentum states:{}\n'
#       .format(len(list(chain(*states.values())))))
# print('--------------------------------------------------')
# states = all_bloch_states(N)
# for key, val in states.items():
#     print('\nk = {}\ncount = {}\n'.format(key, len(val)))
#     for s in val:
#         print(s)
# print('\nTotal number of states:', len(list(chain(*states.values()))))

# N = 5
# k = 0
# i = 0
# state = generate_dec_to_ind_dictionary(N, k)[0][i]
# print(list(map(lambda x: format(x, '0{}b'.format(N)), state)))
# print(H_z_elements(N, k, i, 1))

# i = 2
# val, j = H_pm_elements(N, k, i, 1)
# print(val, j)
<commit_msg>Made some changes<commit_after>import fractions
from itertools import chain
import numpy as np
from spinsys.exceptions import NotFoundError
import functools
from hamiltonians.triangular_lattice_model import SiteVector


def bond_list(Nx, Ny):
    N = Nx * Ny
    bonds = []
    vec = SiteVector((0, 0), Nx, Ny)
    for i in range(N):
        bonds.append((vec, vec.xhop(1)))
        bonds.append((vec, vec.yhop(1)))
        bonds.append((vec, vec.xhop(-1).yhop(1)))
        vec = vec.next_site()
    return bonds


# def roll_state_x(state, Nx, Ny):
#     maxdec = 2 ** Nx - 1
#     # slice state into Ny equal slices
#     n = np.arange(0, Nx * Ny, Nx)
#     s = state % (2 ** (n + Nx)) // (2 ** n)
#     s = (s * 2) % maxdec
#     return (2 ** np.arange(0, Nx * Ny, Nx)).dot(s)


def roll_state_x(state, Nx, Ny):
    # slice state into Ny equal slices
    s = np.array([(state % (2 ** (n + Nx))) // (2 ** n) for n in range(0, Nx * Ny, Nx)])
    print(s)
    maxdec = 2 ** Nx - 1
    s = (s * 2) % maxdec
    return (2 ** np.arange(0, Nx * Ny, Nx)).dot(s)


print(format(7, '06b'))
x = roll_state_x(7, 3, 4)
print(format(x, '06b'))


def roll_state_y(state, Nx, Ny):
    n = 2 ** (Nx * (Ny - 1))
    head = state // n
    return state % n * (2 ** Nx) + head


def zero_momentum_states(Nx, Ny):

    def find_T_invariant_set(i):
        bloch_func = [i]
        j = i
        while True:
            print(i, j)
            j = roll_state_x(j, Nx, Ny)
            if not j == i:
                bloch_func.append(j)
                sieve[j] = False
            else:
                break
        bloch_func.sort()
        try:
            l = len(bloch_func)
            states[l].append(bloch_func)
        except KeyError:
            states[l] = [bloch_func]

    N = Nx * Ny
    sieve = np.ones(2 ** N)
    # sieve[0] = sieve[-1] = 0
    invariant_states = np.kron(np.arange(Nx), 3)
    sieve[invariant_states] = 0
    maxdec = 2 ** N - 1
    states = {1: [[0], [maxdec]]}
    for i in range(maxdec + 1):
        if sieve[i]:
            find_T_invariant_set(i)
    return states


def bloch_states(N, k):
    zero_k_states = zero_momentum_states(N)
    if k > N // 2:
        raise NotFoundError

    if k == 0:
        states = list(chain(*zero_k_states.values()))
    elif k > 0 or (k < 0 and ((not -k == N // 2) or (not N % 2 == 0))):
        states = []
        for l in zero_k_states.keys():
            period = N // fractions.gcd(N, k)
            if l % period == 0:
                states.extend(zero_k_states[l])
    else:
        raise NotFoundError
    return states


@functools.lru_cache(maxsize=None)
def generate_dec_to_ind_dictionary(N, k):
    states = sorted(bloch_states(N, k))
    dec_to_ind = {}
    for i, state in enumerate(states):
        state_len = len(state)
        for num in state:
            dec_to_ind[num] = (i, state_len)
    return states, dec_to_ind


def all_bloch_states(N):
    states = {}
    max_k = N // 2
    for k in range(-max_k, max_k + 1):
        try:
            states[k] = bloch_states(N, k)
        except NotFoundError:
            continue
    return states


@functools.lru_cache(maxsize=None)
def count_same_spins(N, state, l, return_loc=False):
    # edge cases that break down under the kind of binary operation we
    #  use later
    if not (state == 2 ** N - 1 or state == 0):
        # subtraction does a bitwise flip of 1's and 0's. We need this
        #  because subsequent operations are insensitive to patterns of 0's
        bit_inverted_state = 2 ** N - 1 - state
        # the mod operator accounts for periodic boundary conditions
        couplings = 2 ** np.arange(N) + 2 ** (np.arange(l, N + l) % N)
        nup = ndown = 0
        ups = None
        downs = None
        for i in couplings:
            # if a state is unchanged under bitwise OR with the given number,
            #  it has two 1's at those two sites
            if state | i == state:
                nup += 1
                ups = i
            # cannot be merged with the previous statement because bit-flipped
            #  numbers using subtraction are possibly also bit-shifted
            if bit_inverted_state | i == bit_inverted_state:
                ndown += 1
                downs = i
    elif state == 0:
        nup, ndown = 0, N
    else:
        nup, ndown = N, 0  # when state is all 0's

    if return_loc:
        return nup, ndown, ups, downs
    else:
        return nup, ndown


def H_z_elements(N, k, i, l):
    state = generate_dec_to_ind_dictionary(N, k)[0][i][0]
    nup, ndown = count_same_spins(N, state, l)
    same_dir = nup + ndown
    diff_dir = N - same_dir
    return 0.25 * (same_dir - diff_dir)


def H_pm_elements(N, k, i, l):
    states, dec_to_ind = generate_dec_to_ind_dictionary(N, k)
    state = states[i][0]
    i_len = len(states[i])
    bit1 = 2 ** np.arange(N)
    bit2 = 2 ** (np.arange(l, N + 1) % N)
    count = 0  # count of up-downs and down-ups separated by l
    for b1, b2 in zip(bit1, bit2):
        if (state | b1 == state) and (not state | b2 == state):
            count += 1
        if (not state | b1 == state) and (state | b2 == state):
            count += 1
            pm = (b1, b2)
    # bit-flip the state once to see what it is coupled to. We don't need
    #  to do it more than once since such flips will always get to
    #  constituent states of the same bloch state
    connected_state = state + pm[0] - pm[1]
    j, j_len = dec_to_ind[connected_state]
    val = i_len * count / j_len * np.sqrt(j_len / i_len)
    return val, j


def H_ppmm_elements(N, k, i, l):
    state = generate_dec_to_ind_dictionary(N, k)[0][i][0]
    nup, ndown, ups, downs = count_same_spins(N, state, l, return_loc=True)
    if ups is not None:
        connected_state = state - ups
    elif downs is not None:
        connected_state = state + downs
    else:
        raise NotFoundError
    j, j_len = dec_to_ind[connected_state]


Nx = 2
Ny = 3

print(roll_state_x(8, Nx, Ny))

# N = Nx * Ny
# states = zero_momentum_states(Nx, Ny)
# totlen = 0
# for key, val in states.items():
#     print('\nlength = {}\ncount = {}\n'.format(key, len(val)))
#     for s in val:
#         print(s)
#         totlen += len(s)
# print('\nSanity check. Total number of product states:', totlen)
# print('Total number of zero momentum states:{}\n'
#       .format(len(list(chain(*states.values())))))

# print('--------------------------------------------------')
# states = all_bloch_states(N)
# for key, val in states.items():
#     print('\nk = {}\ncount = {}\n'.format(key, len(val)))
#     for s in val:
#         print(s)
# print('\nTotal number of states:', len(list(chain(*states.values()))))

# N = 5
# k = 0
# i = 0
# state = generate_dec_to_ind_dictionary(N, k)[0][i]
# print(list(map(lambda x: format(x, '0{}b'.format(N)), state)))
# print(H_z_elements(N, k, i, 1))

# i = 2
# val, j = H_pm_elements(N, k, i, 1)
# print(val, j)

# s = int('0110100011100110', 2)
# print(format(s, '016b'))
# print(format(roll_state_to_rightx(s, 8, 2), '016b'))

{'fraction_non_alphanumeric': 0.13514969599014853, 'fraction_numerical': 0.016855229739090278, 'mean_word_length': 2.1295761078998074, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4713250', 'n_tokens_mistral': 4801, 'n_tokens_neox': 4558, 'n_words': 1378}