tV model ground state entanglent as a function of bipartition length

S = 14 #sites

import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import expm as matrixexp
from scipy.linalg import logm as matrixlog
from scipy.sparse.linalg import eigsh as sparse_eigs
import scipy.sparse as sparse
from functools import lru_cache
%config InlineBackend.figure_format = "svg"

# New version with sparse matrices. These are defined via Jordan-Wigner
@lru_cache(maxsize=None)
def cD_(n, L=10):
    S = sparse.csc_matrix([[0,1],[0,0]])
    pauli_z = sparse.csc_matrix([[1,0],[0,-1]])
    id2 = sparse.identity(2,format="csc")
    
    # This construction is written explicitly in the slides. I am just copying them
    # matrices to the left of the index n are pauli z matrices, and to the right are identity.
    # at index n, it is the S^+ operator.
    temp = S if n==0 else pauli_z
    
    for idx in range(1, L):
        if idx < n:
            temp = sparse.kron(temp, pauli_z)
        elif idx == n:
            temp = sparse.kron(temp, S)
        else:
            temp = sparse.kron(temp, id2)
    
    # re-type casting to a sparse matrix removes extraneous zeroes and greatly speeds up the code
    return sparse.csc_matrix(temp.toarray())


@lru_cache(maxsize=None)
# Identical to cD_ except each matrix is replaced with its transpose
def c_(n, L=10):
    S = sparse.csc_matrix([[0,0],[1,0]])
    pauli_z = sparse.csc_matrix([[1,0],[0,-1]])
    id2 = sparse.identity(2,format="csc")
    
    temp = S if n==0 else pauli_z
    
    for idx in range(1, L):
        if idx < n:
            temp = sparse.kron(temp, pauli_z)
        elif idx == n:
            temp = sparse.kron(temp, S)
        else:
            temp = sparse.kron(temp, id2)
        
    return sparse.csc_matrix(temp.toarray())

#Helper function for later that defines a number operator to save space in the code & help readability
@lru_cache(maxsize=None)
def n_(n,L=10):
    return cD_(n,L) @ c_(n,L)

def indicesFixedN(N, L):
    # My testing suggests that the eigenvectors of the number operators are the pure cartesian bases vectors and the
    # number operator eigenvalue of eigenvector n is the bit sum of (2**L-1-n) in binary. 
    temp = []
    for idx in range(2**L):
        if bit_sum(2**L-1-idx) == N:
            temp.append(idx)
    return np.array(temp)

def bit_sum(n):
    return bin(n).count("1")

def vacuum(L):
    temp = np.zeros(2**L)
    temp[-1]=1
    return temp

%%time

def rot(sites,fermions):
    #rotate into the fixed particle number sector
    indices = indicesFixedN(fermions,sites)
    num = np.shape(indices)[0]
    rot = sparse.csc_matrix(np.zeros((2**sites,2**sites)))
    for i in range(0,num,1):
        rot[i,indices[i]] = 1
    return num,rot

dim,rotate = rot(S,S//2)
identity = sparse.identity(2**S, format = "csc")#*****dim, format = "csc")

def number(sites):
    #create a dictionary for the number operators rotated into the fixed particle number sector
    numbers = {}
    for i in range(0,sites,1):
        numbers[str(i)] = n_(i,sites) #*****(rotate @ n_(i,sites) @ rotate.T)[0:dim,0:dim]
    return numbers
        
numbers = number(S)

def t_Hamiltonian(sites,t=1):
    #build the Hamiltonian without disorder
    t_total = sparse.csc_matrix(np.zeros((2**sites,2**sites)))
    for i in range(0,sites-1,1):
        t_total += - t * (cD_(i,sites)@c_(i+1,sites) + cD_(i+1,sites)@c_(i,sites))
    t_total += (-1)**((sites/2)%2) * t * (cD_(sites-1,sites)@c_(0,sites) + cD_(0,sites)@c_(sites-1,sites)) #(anti)periodic BCs (for (even) odd number of fermions)
    t_reduced = t_total #*****(rotate @ t_total @ rotate.T)[0:dim,0:dim]
    return t_reduced

def V_Hamiltonian(sites,V=1):
    #build the Hamiltonian without disorder
    V_total = sparse.csc_matrix(np.zeros((2**S,2**S))) #******np.zeros((dim,dim))
    for i in range(0,sites-1,1):
        V_total += V*(numbers[str(i)])@(numbers[str(i+1)])
    V_total += V * (numbers[str(sites-1)])@(numbers[str(0)]) #(anti)periodic BCs (for (even) odd number of fermions)
    return V_total

H_t = t_Hamiltonian(S)
H_V = V_Hamiltonian(S)

def partial_basis(s):
    basis_states = np.zeros((2**s,2**s),dtype=complex)
    for i in range(0,2**s,1):
        state = vacuum(s)
        b = bin(2**s+i)
        for j in range(1,s+1,1):
            if(int(b[-j])==1):
                state = cD_(j-1,s)@state
        basis_states[i] = state
    return basis_states

def partial_trace(psi,s,S):
    Psi = np.zeros((2**s,2**(S-s)),dtype=complex)
    for i in range(0,2**s,1):
        for j in range(0,2**(S-s),1):
            Psi[i,j] = np.kron(pbA[i],pbB[j])@psi
    rho_A = Psi @ np.conjugate(Psi.T)
    return rho_A

%%time

t = 1
V = 2

S_EE = np.zeros(S+1)
idxs = indicesFixedN(S//2,S)

H = t*H_t + V*H_V #create a Hamiltonian
H = (rotate @ H @ rotate.T)[0:dim,0:dim] #*********rotate into fixed particle basis
#find ground state, express in big basis
eigs, vecs = sparse_eigs(H,k=1,sigma=-0.5*S) #make sure sigma is less than the lowest eigenvalue for convergence! -- an exact value is available for the Heisenberg model
psi_gs = vecs[:,0]
psi_gs_big = np.zeros(2**S,dtype=complex)
psi_gs_big[idxs] = psi_gs
    
for s in range(1,S//2+1,1):
    pbA = partial_basis(s)
    pbB = partial_basis(S-s)
    rho_A = partial_trace(psi_gs_big,s,S)
    S_EE[s] = -np.real(np.trace(rho_A@matrixlog(rho_A)))
    S_EE[-s-1] = S_EE[s]

plt.scatter(np.linspace(0,S,S+1),np.log2(np.exp(1))*S_EE,marker="o",label=r"$S_{GS}$ for $L=$"+str(S))
a = 1.0
c = np.pi
x = np.linspace(0,S,100)
plt.plot(x,a + c/6 * np.log(S/np.pi * np.sin(np.pi*x/S)),linestyle="dashed",label=r"$A+\frac{c}{6}\ln(\frac{L}{\pi}\sin(\pi\frac{L_A}{L}))$")
plt.ylim(0,2.3)
plt.title(r"tV Model Ground State Entanglement ($V={}$, $c={}$)".format(V,np.round(c,2)))
plt.ylabel(r"$S=-\mathrm{Tr}[\rho^A_{GS}\log_{2}(\rho^A_{GS})]$")
plt.xlabel(r"Bipartition Size ($L_A$)")
plt.legend()