1##################################################################
  2#                            example 1                           #
  3#    In this script we demonstrate how to use QuSpin's           #
  4#    functionality to solve the time-dependent Schroedinger      #
  5#    equation with a time-dependent operator. We also showcase   #
  6#    a function which is used to calculate the entanglement      #
  7#    entropy of a pure state.                                    #
  8##################################################################
  9from quspin.operators import quantum_operator  # Hamiltonians and operators
 10from quspin.basis import spin_basis_1d  # Hilbert space spin basis
 11from quspin.tools.measurements import ent_entropy, diag_ensemble  # entropies
 12from numpy.random import uniform, seed  # pseudo random numbers
 13from joblib import delayed, Parallel  # parallelisation
 14import numpy as np  # generic math functions
 15from time import time  # timing package
 16
 17#
 18##### define simulation parameters #####
 19n_real = 100  # number of disorder realisations
 20n_jobs = 2  # number of spawned processes used for parallelisation
 21#
 22##### define model parameters #####
 23L = 10  # system size
 24Jxy = 1.0  # xy interaction
 25Jzz_0 = 1.0  # zz interaction at time t=0
 26h_MBL = 3.9  # MBL disorder strength
 27h_ETH = 0.1  # delocalised disorder strength
 28vs = np.logspace(-3.0, 0.0, num=20, base=10)  # log_2-spaced vector of ramp speeds
 29
 30
 31#
 32##### set up Heisenberg Hamiltonian with linearly varying zz-interaction #####
 33# define linear ramp function
 34def ramp(t, v):
 35    return 0.5 + v * t
 36
 37
 38# compute basis in the 0-total magnetisation sector (requires L even)
 39basis = spin_basis_1d(L, m=0, pauli=False)
 40# define operators with OBC using site-coupling lists
 41J_zz = [[Jzz_0, i, i + 1] for i in range(L - 1)]  # OBC
 42J_xy = [[Jxy / 2.0, i, i + 1] for i in range(L - 1)]  # OBC
 43# dictionary of operators for quantum_operator.
 44# define XXZ chain parameters
 45op_dict = dict(J_xy=[["+-", J_xy], ["-+", J_xy]], J_zz=[["zz", J_zz]])
 46# define operators for local disordered field.
 47for i in range(L):
 48    op = [[1.0, i]]
 49    op_dict["hz" + str(i)] = [["z", op]]
 50# costruct the quantum_operator
 51H_XXZ = quantum_operator(op_dict, basis=basis, dtype=np.float64)
 52
 53
 54#
 55##### calculate diagonal and entanglement entropies #####
 56def realization(vs, H_XXZ, real):
 57    """
 58    This function computes the entropies for a single disorder realisation.
 59    --- arguments ---
 60    vs: vector of ramp speeds
 61    H_XXZ: time-dep. Heisenberg Hamiltonian with driven zz-interactions
 62    basis: spin_basis_1d object containing the spin basis
 63    n_real: number of disorder realisations; used only for timing
 64    """
 65    ti = time()  # get start time
 66    #
 67    seed()  # the random number needs to be seeded for each parallel process
 68    N = H_XXZ.basis.N
 69    basis = H_XXZ.basis
 70    hz = uniform(-1, 1, size=N)
 71    # define parameters to pass into quantum_operator for
 72    # hamiltonian at end of ramp
 73    pars_MBL = {"hz" + str(i): h_MBL * hz[i] for i in range(N)}
 74    pars_ETH = {"hz" + str(i): h_ETH * hz[i] for i in range(N)}
 75    #
 76    pars_MBL["J_xy"] = 1.0
 77    pars_ETH["J_xy"] = 1.0
 78    # J_zz = 1 at end of the ramp for all velocities
 79    pars_MBL["J_zz"] = 1.0
 80    pars_ETH["J_zz"] = 1.0
 81    # diagonalize
 82    E_MBL, V_MBL = H_XXZ.eigh(pars=pars_MBL)
 83    E_ETH, V_ETH = H_XXZ.eigh(pars=pars_ETH)
 84    # reset J_zz to be initial value:
 85    pars_MBL["J_zz"] = 0.5
 86    # get many-body bandwidth at t=0
 87    eigsh_args = dict(
 88        k=2, which="BE", maxiter=1e4, return_eigenvectors=False, pars=pars_MBL
 89    )
 90    Emin, Emax = H_XXZ.eigsh(**eigsh_args)
 91    # calculating middle of spectrum
 92    E_inf_temp = (Emax + Emin) / 2.0
 93    # calculate nearest eigenstate to energy at infinite temperature
 94    E, psi_0 = H_XXZ.eigsh(pars=pars_MBL, k=1, sigma=E_inf_temp, maxiter=1e4)
 95    psi_0 = psi_0.reshape((-1,))
 96
 97    run_MBL = []
 98
 99    for v in vs:
100        # update J_zz to be time-dependent operator
101        pars_MBL["J_zz"] = (ramp, (v,))
102        # get hamiltonian
103        H = H_XXZ.tohamiltonian(pars=pars_MBL)
104        # evolve state and calculate oberservables.
105        run_MBL.append(_do_ramp(psi_0, H, basis, v, E_MBL, V_MBL))
106
107    run_MBL = np.vstack(run_MBL).T
108    # reset J_zz to be initial value:
109    pars_ETH["J_zz"] = 0.5
110    # get many-body bandwidth at t=0
111    eigsh_args = dict(
112        k=2, which="BE", maxiter=1e4, return_eigenvectors=False, pars=pars_ETH
113    )
114    Emin, Emax = H_XXZ.eigsh(**eigsh_args)
115    # calculating middle of spectrum
116    E_inf_temp = (Emax + Emin) / 2.0
117    # calculate nearest eigenstate to energy at infinite temperature
118    E, psi_0 = H_XXZ.eigsh(pars=pars_ETH, k=1, sigma=E_inf_temp, maxiter=1e4)
119    psi_0 = psi_0.reshape((-1,))
120
121    run_ETH = []
122
123    for v in vs:
124        # update J_zz to be time-dependent operator
125        pars_ETH["J_zz"] = (ramp, (v,))
126        # get hamiltonian
127        H = H_XXZ.tohamiltonian(pars=pars_ETH)
128        # evolve state and calculate oberservables.
129        run_ETH.append(_do_ramp(psi_0, H, basis, v, E_ETH, V_ETH))
130    run_ETH = np.vstack(run_ETH).T
131    # show time taken
132    print("realization {0}/{1} took {2:.2f} sec".format(real + 1, n_real, time() - ti))
133    #
134    return run_MBL, run_ETH
135
136
137#
138##### evolve state and evaluate entropies #####
139def _do_ramp(psi_0, H, basis, v, E_final, V_final):
140    """
141    Auxiliary function to evolve the state and calculate the entropies after the
142    ramp.
143    --- arguments ---
144    psi_0: initial state
145    H: time-dependent Hamiltonian
146    basis: spin_basis_1d object containing the spin basis (required for Sent)
147    E_final, V_final: eigensystem of H(t_f) at the end of the ramp t_f=1/(2v)
148    """
149    # determine total ramp time
150    t_f = 0.5 / v
151    # time-evolve state from time 0.0 to time t_f
152    psi = H.evolve(psi_0, 0.0, t_f)
153    # calculate entanglement entropy
154    subsys = range(basis.L // 2)  # define subsystem
155    Sent = basis.ent_entropy(psi, sub_sys_A=subsys)["Sent_A"]
156    # calculate diagonal entropy in the basis of H(t_f)
157    S_d = diag_ensemble(basis.L, psi, E_final, V_final, Sd_Renyi=True)["Sd_pure"]
158    #
159    return np.asarray([S_d, Sent])
160
161
162#
163##### produce data for n_real disorder realisations #####
164# __name__ == '__main__' required to use joblib in Windows.
165if __name__ == "__main__":
166
167    # alternative way without parallelisation
168    data = np.asarray([realization(vs, H_XXZ, i) for i in range(n_real)])
169    """
170	data = np.asarray(Parallel(n_jobs=n_jobs)(delayed(realization)(vs,H_XXZ,basis,i) for i in range(n_real)))
171	"""
172    #
173    run_MBL, run_ETH = zip(*data)  # extract MBL and data
174    # average over disorder
175    mean_MBL = np.mean(run_MBL, axis=0)
176    mean_ETH = np.mean(run_ETH, axis=0)
177    #
178    ##### plot results #####
179    import matplotlib.pyplot as plt
180
181    ### MBL plot ###
182    fig, pltarr1 = plt.subplots(2, sharex=True)  # define subplot panel
183    # subplot 1: diag enetropy vs ramp speed
184    pltarr1[0].plot(vs, mean_MBL[0], label="MBL", marker=".", color="blue")  # plot data
185    pltarr1[0].set_ylabel("$s_d(t_f)$", fontsize=22)  # label y-axis
186    pltarr1[0].set_xlabel("$v/J_{zz}(0)$", fontsize=22)  # label x-axis
187    pltarr1[0].set_xscale("log")  # set log scale on x-axis
188    pltarr1[0].grid(True, which="both")  # plot grid
189    pltarr1[0].tick_params(labelsize=16)
190    # subplot 2: entanglement entropy vs ramp speed
191    pltarr1[1].plot(vs, mean_MBL[1], marker=".", color="blue")  # plot data
192    pltarr1[1].set_ylabel("$s_\\mathrm{ent}(t_f)$", fontsize=22)  # label y-axis
193    pltarr1[1].set_xlabel("$v/J_{zz}(0)$", fontsize=22)  # label x-axis
194    pltarr1[1].set_xscale("log")  # set log scale on x-axis
195    pltarr1[1].grid(True, which="both")  # plot grid
196    pltarr1[1].tick_params(labelsize=16)
197    # save figure
198    fig.savefig("example1_MBL.pdf", bbox_inches="tight")
199    #
200    ### ETH plot ###
201    fig, pltarr2 = plt.subplots(2, sharex=True)  # define subplot panel
202    # subplot 1: diag enetropy vs ramp speed
203    pltarr2[0].plot(vs, mean_ETH[0], marker=".", color="green")  # plot data
204    pltarr2[0].set_ylabel("$s_d(t_f)$", fontsize=22)  # label y-axis
205    pltarr2[0].set_xlabel("$v/J_{zz}(0)$", fontsize=22)  # label x-axis
206    pltarr2[0].set_xscale("log")  # set log scale on x-axis
207    pltarr2[0].grid(True, which="both")  # plot grid
208    pltarr2[0].tick_params(labelsize=16)
209    # subplot 2: entanglement entropy vs ramp speed
210    pltarr2[1].plot(vs, mean_ETH[1], marker=".", color="green")  # plot data
211    pltarr2[1].set_ylabel("$s_\\mathrm{ent}(t_f)$", fontsize=22)  # label y-axis
212    pltarr2[1].set_xlabel("$v/J_{zz}(0)$", fontsize=22)  # label x-axis
213    pltarr2[1].set_xscale("log")  # set log scale on x-axis
214    pltarr2[1].grid(True, which="both")  # plot grid
215    pltarr2[1].tick_params(labelsize=16)
216    # save figure
217    fig.savefig("example1_ETH.pdf", bbox_inches="tight")
218    #
219    # plt.show() # show plots