S = 6
import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import expm
%config InlineBackend.figure_format = "svg"
adj = lambda u : np.conj(u.T)
import tenpy

Set Up Spin Operators

s0 = np.array([[1,0],[0,1]],dtype=complex)
sx = np.array([[0,1],[1,0]],dtype=complex)
sy = np.array([[0,-1j],[1j,0]],dtype=complex)
sz = np.array([[1,0],[0,-1]],dtype=complex)

pauli = {"0" : s0,
         "x" : sx,
         "y" : sy,
         "z" : sz,
         "+" : (sx+1j*sy)/2,
         "-" : (sx-1j*sy)/2}

def spin(polarization,s,S):
    if(s==0):
        temp = pauli[polarization]
    else:
        temp = s0
    for i in range(1,S,1):
        if(i==s):
            temp = np.kron(temp,pauli[polarization])
        else:
            temp = np.kron(temp,s0)
    return temp

Functions to Create States

def vacuum(S):
    vac = np.zeros(2**S,dtype=complex)
    vac[-1] = 1
    return vac

Set Up Hamiltonian

def H_s(S,dir): #open boundaries
    H = np.zeros((2**S,2**S),dtype=complex)
    for i in range(0,S-1,1):
        H += spin(dir,i,S)@spin(dir,i+1,S)
    return H

HX = H_s(S,"x")
HY = H_s(S,"y")
HZ = H_s(S,"z")

def h_s(S,dir):
    #create a dictionary for the Sz operators
    numbers = {}
    for i in range(0,S,1):
        numbers[str(i)] = spin(dir,i,S)
    return numbers

hX = h_s(S,"x")
hY = h_s(S,"y")
hZ = h_s(S,"z")

def Ham(S,Jx=1,Jy=0,Jz=0,hx=0,hy=0,hz=0):
    #build the Hamiltonian
    H = Jx*HX + Jy*HY + Jz*HZ
    for i in range(0,S,1): #we could do a random transverse field here
        H += hx*hX[str(i)] + hy*hY[str(i)] + hz*hZ[str(i)]
    return H

Model Parameters

L = S
Jx,Jy,Jz=1,0,0#0.5,1,0.3     #random spin couplings
hx,hy,hz=0,0,0.5#-2,0.2,0.5    #random B fields
G = 1

nT = 1001
Tmax = 20
Ts = np.linspace(0,Tmax,nT)
dt = Tmax/(nT-1)
dirs = ["x","y","z"]

psi0 = vacuum(S) #or make another state

H = Ham(S,Jx=Jx,Jy=Jy,Jz=Jz,hx=hx,hy=hy,hz=hz)
J1 = spin("+",0,S) #raising operator at left boundary
J2 = spin("-",S-1,S) #lowering operator at right boundary
G1,G2 = G,G

Exponential

Vectorize the density matrix by columns $$ \rho = \begin{pmatrix} \rho_{11}&\rho_{21}\\ \rho_{12}&\rho_{22} \end{pmatrix} \quad \mapsto \quad \vec{\rho} = \begin{pmatrix} \rho_{11}\\ \rho_{12}\\ \rho_{21}\\ \rho_{22} \end{pmatrix} $$

Time evolution is now an eigenvalue operator (not superoperator) equation Lvecρ⃗ = i∂tρ⃗ with

$$ L_\text{vec} = 1\otimes H - H^\top \otimes 1 -i\sum_m \frac{\gamma_m}{2} ((J_m^\dagger J_m)^\top\otimes 1 + 1 \otimes J_m^\dagger J_m - 2(J_m^\dagger)^\top \otimes J_m) $$ where we note the switched order of the dagger term versus the master equation. With U(δt) = e−iδtLvec we have ρ⃗(t+δt) = U(δt)ρ⃗(t) where here δt need not be small.

%%time
one = np.identity(2**S)

L_vec = np.kron(one,H) - np.kron(H.T,one)\
        - 1j*G1/2 * (np.kron((adj(J1)@J1).T,one) + np.kron(one,adj(J1)@J1) - 2*np.kron(adj(J1).T,J1))\
        - 1j*G2/2 * (np.kron((adj(J2)@J2).T,one) + np.kron(one,adj(J2)@J2) - 2*np.kron(adj(J2).T,J2))

rho = np.einsum("a,b->ab",adj(psi0),psi0)
rho_vec = np.zeros(2**S * 2**S,dtype=complex)
for c in range(2**S):
    rho_vec[c*2**S:(c+1)*2**S] = rho[:,c]

mags = np.zeros((nT,3))
norm_vec = np.zeros(nT,dtype=complex)
U = expm(-1j*dt*L_vec)

for t in range(nT):
    #expectation values
    norm_vec[t] = np.trace(rho)
    for s in range(S):
        for d in range(3):
            mags[t,d] += np.real(np.trace(rho@spin(dirs[d],s,S))) / S
    rho_vec = U@rho_vec
    for c in range(2**S):
        rho[:,c] = rho_vec[c*2**S:(c+1)*2**S]
CPU times: user 5min 55s, sys: 49.7 s, total: 6min 45s
Wall time: 36.8 s

Naive Integrate

Numerically integrate the master equation in time using the Euler method $$ \rho(t+\delta t) = \rho(t) - i\, \delta t \bigg([H,\rho(t)]-i\sum_m \frac{\gamma_m}{2} (J_m^\dagger J_m\rho + \rho J_m^\dagger J_m - 2J_m\rho J_m^\dagger)\bigg) $$

Making the timestep smaller seems to help.

%%time
rho = np.einsum("a,b->ab",adj(psi0),psi0)
magsNaive = np.zeros((nT,3))
trace = np.zeros(nT)
for t in range(nT):
    for s in range(S):
        for d in range(3):
            magsNaive[t,d] += np.real(np.trace(rho@spin(dirs[d],s,S))) / S
    rho = rho + -1j*dt * ((H@rho-rho@H)\
                       -1j*G1/2*(adj(J1)@J1@rho + rho@adj(J1)@J1 - 2*J1@rho@adj(J1))\
                       -1j*G2/2*(adj(J2)@J2@rho + rho@adj(J2)@J2 - 2*J2@rho@adj(J2)))
CPU times: user 40 s, sys: 360 ms, total: 40.3 s
Wall time: 3.4 s

TEBD

class VectorizedSpinHalfSite(tenpy.networks.site.Site):
    def __init__(self,conserve="None"):
        dim = 2 #SU(2), local hilbert space dimension before vectorization
        #conserved quantities
        self.conserve = conserve
        
        #basic operators
        id = np.array([[1,0],[0,1]],dtype=complex) #no factor of 1/2 here
        sx = np.array([[0,1],[1,0]],dtype=complex)
        sy = np.array([[0,-1j],[1j,0]],dtype=complex)
        sz = np.array([[1,0],[0,-1]],dtype=complex)
        sp = (sx + 1j*sy)/2 #[[0,1],[0,0]]
        sm = (sx - 1j*sy)/2 #[[0,0],[1,0]]

        #composite operators (with transposes!)
        names = ["0","x","y","z","p","m"]
        operators = [id,sx,sy,sz,sp,sm]
        my_operators = {}
        for i in range(0,6,1):
            for j in range(0,6,1):
                my_operators["s"+str(names[i])+str(names[j])] = np.kron(operators[i].T,operators[j])

        #dissipative operators
        my_operators["diss_m"] = np.kron((adj(sm)@sm).T,id) + np.kron(id,adj(sm)@sm) - 2*np.kron(adj(sm).T,sm)
        my_operators["diss_p"] = np.kron((adj(sp)@sp).T,id) + np.kron(id,adj(sp)@sp) - 2*np.kron(adj(sp).T,sp)
        
        #make it!
        tenpy.networks.site.Site.__init__(self, tenpy.linalg.charges.LegCharge.from_trivial(dim**2), **my_operators)

        #label the states - this needs to be done AFTER Site.__init__() is run
        S = (dim-1)/2
        my_spin = {}
        for i in range(0,dim,1):
            my_spin[str(-S+i)] = i*(dim+1)
        my_spin["up"] = 0
        my_spin["down"] = dim**2-1
        self.state_labels = my_spin
class VectorizedXYZChain(tenpy.models.model.CouplingMPOModel,tenpy.models.model.NearestNeighborModel):
    def init_sites(self, model_params):
        return VectorizedSpinHalfSite(conserve="None")

    def init_terms(self, model_params):
        # read out parameters
        Jx = model_params.get("Jx", 0.)
        Jy = model_params.get("Jy", 0.)
        Jz = model_params.get("Jz", 0.)
        hx = model_params.get("hx", 0.)
        hy = model_params.get("hy", 0.)
        hz = model_params.get("hz", 0.)
        Gp = model_params.get("Gp", 0.)
        Gm = model_params.get("Gm", 0.)
        # add terms
        for i in range(len(self.lat.unit_cell)):
            self.add_onsite(-hx, i, "sx0")
            self.add_onsite(-hy, i, "sy0")
            self.add_onsite(-hz, i, "sz0")
            #
            self.add_onsite(hx, i, "s0x")
            self.add_onsite(hy, i, "s0y")
            self.add_onsite(hz, i, "s0z")
        for i1, i2, dx in self.lat.pairs["nearest_neighbors"]:
            self.add_coupling(-Jx, i1, "sx0", i2, "sx0", dx, plus_hc=False)
            self.add_coupling(-Jy, i1, "sy0", i2, "sy0", dx, plus_hc=False)
            self.add_coupling(-Jz, i1, "sz0", i2, "sz0", dx, plus_hc=False)
            #
            self.add_coupling(Jx, i1, "s0x", i2, "s0x", dx, plus_hc=False)
            self.add_coupling(Jy, i1, "s0y", i2, "s0y", dx, plus_hc=False)
            self.add_coupling(Jz, i1, "s0z", i2, "s0z", dx, plus_hc=False)
        #add boundary dissipation
        self.add_local_term(-1j*Gp/2,[("diss_p",(0,0))])
        self.add_local_term(-1j*Gm/2,[("diss_m",(L-1,0))])
#simulation parameters
tebd_params = {
    "trunc_params": {
        "chi_max": 20, #bond dimension
        "svd_min": 1*10**-10
    },
    "N_steps": 1,
    "dt": Tmax/(nT-1),
    "order": 2,
    "preserve_norm": False
}
%%time

mzs = np.zeros(nT)
model_params = dict(L=L, Jx=Jx,Jy=Jy,Jz=Jz, hx=hx,hy=hy,hz=hz,Gp=G1,Gm=G2, bc_MPS="finite")

M = VectorizedXYZChain(model_params)
psi = tenpy.networks.mps.MPS.from_lat_product_state(M.lat,[["down"]]) #initial ferromagnetic state
engine = tenpy.algorithms.tebd.TEBDEngine(psi, M, tebd_params)

id = np.array([1,0,0,1]).reshape((4,1,1)) #vectorized
my_tuple = (id,id) * int(L//2) # only works for even length
psi_id = tenpy.networks.mps.MPS.from_Bflat(M.lat.mps_sites(), my_tuple)

norms = np.zeros(nT)
for step in range(nT):
    #norms[step] = np.abs(psi.overlap(psi))
    braket = tenpy.networks.mps.MPSEnvironment(psi_id,psi)
    mzs[step] = np.real(np.mean(braket.expectation_value("sz0")))
    engine.run()
CPU times: user 22min 38s, sys: 37min 28s, total: 1h 7s
Wall time: 5min 5s

plt.plot(Ts,mags[:,2],lw=4,label="Reference: vectorized exp")
#plt.plot(Ts,magsNaive[:,2],label="Reference: euler integrate")
plt.plot(Ts,mzs,ls="dashed",label=r"TEBD $\langle 1111|S^{z0}|\rho\rangle$")
plt.legend()
plt.xlabel(r"$t$")
plt.ylabel(r"$\langle S^z\rangle$")
Text(0, 0.5, '$\\langle S^z\\rangle$')