1#####################################################################
  2#                            example 8                              #
  3#   In this script we demonstrate how to use QuSpin to define       #
  4#   and solve nonlinear ordinary differential equations.            #
  5#   In particular, we show real and imaginary time evolution        #
  6#   in the Gross-Pitaevskii equation of a Bose condensate in a      #
  7#   harmonic trap.                                                  #
  8#####################################################################
  9from quspin.operators import hamiltonian  # Hamiltonians and operators
 10from quspin.basis import boson_basis_1d  # Hilbert space boson basis
 11from quspin.tools.evolution import evolve  # nonlinear evolution
 12import numpy as np  # generic math functions
 13import matplotlib.pyplot as plt  # plotting library
 14from six import iteritems  # loop over elements of dictionary
 15
 16#
 17##### define model parameters #####
 18L = 300  # system size
 19# calculate centre of chain
 20if L % 2 == 0:
 21    j0 = L // 2 - 0.5  # centre of chain
 22else:
 23    j0 = L // 2  # centre of chain
 24sites = np.arange(L) - j0
 25# static parameters
 26J = 1.0  # hopping
 27U = 1.0  # Bose-Hubbard interaction strength
 28# dynamic parameters
 29kappa_trap_i = 0.001  # initial chemical potential
 30kappa_trap_f = 0.0001  # final chemical potential
 31t_ramp = 40.0 / J  # set total ramp time
 32
 33
 34# ramp protocol
 35def ramp(t, kappa_trap_i, kappa_trap_f, t_ramp):
 36    return (kappa_trap_f - kappa_trap_i) * t / t_ramp + kappa_trap_i
 37
 38
 39# ramp protocol parameters
 40ramp_args = [kappa_trap_i, kappa_trap_f, t_ramp]
 41#
 42##### construct single-particle Hamiltonian #####
 43# define site-coupling lists
 44hopping = [[-J, i, (i + 1) % L] for i in range(L - 1)]
 45trap = [[0.5 * (i - j0) ** 2, i] for i in range(L)]
 46# define static and dynamic lists
 47static = [["+-", hopping], ["-+", hopping]]
 48dynamic = [["n", trap, ramp, ramp_args]]
 49# define basis
 50basis = boson_basis_1d(L, Nb=1, sps=2)
 51# build Hamiltonian
 52Hsp = hamiltonian(static, dynamic, basis=basis, dtype=np.float64)
 53E, V = Hsp.eigsh(time=0.0, k=1, which="SA")
 54
 55
 56#
 57##### imaginary-time evolution to compute GS of GPE #####
 58def GPE_imag_time(tau, phi, Hsp, U):
 59    """
 60    This function solves the real-valued GPE in imaginary time:
 61    $$ -\dot\phi(\tau) = Hsp(t=0)\phi(\tau) + U |\phi(\tau)|^2 \phi(\tau) $$
 62    """
 63    return -(Hsp.dot(phi, time=0) + U * np.abs(phi) ** 2 * phi)
 64
 65
 66# define ODE parameters
 67GPE_params = (Hsp, U)
 68# define initial state to flow to GS from
 69phi0 = V[:, 0] * np.sqrt(L)  # initial state normalised to 1 particle per site
 70# define imaginary time vector
 71tau = np.linspace(0.0, 35.0, 71)
 72# evolve state in imaginary time
 73psi_tau = evolve(
 74    phi0,
 75    tau[0],
 76    tau,
 77    GPE_imag_time,
 78    f_params=GPE_params,
 79    imag_time=True,
 80    real=True,
 81    iterate=True,
 82)
 83#
 84# display state evolution
 85for i, psi0 in enumerate(psi_tau):
 86    # compute energy
 87    E_GS = (
 88        Hsp.matrix_ele(psi0, psi0, time=0) + 0.5 * U * np.sum(np.abs(psi0) ** 4)
 89    ).real
 90    # plot wave function
 91    plt.plot(
 92        sites,
 93        abs(phi0) ** 2,
 94        color="r",
 95        marker="s",
 96        alpha=0.2,
 97        label="$|\\phi_j(0)|^2$",
 98    )
 99    plt.plot(sites, abs(psi0) ** 2, color="b", marker="o", label="$|\\phi_j(\\tau)|^2$")
100    plt.xlabel("$\\mathrm{lattice\\ sites}$", fontsize=14)
101    plt.title(
102        "$J\\tau=%0.2f,\\ E_\\mathrm{GS}(\\tau)=%0.4fJ$" % (tau[i], E_GS), fontsize=14
103    )
104    plt.ylim([-0.01, max(abs(phi0) ** 2) + 0.01])
105    plt.legend(fontsize=14)
106    plt.draw()  # draw frame
107    plt.pause(0.005)  # pause frame
108    plt.clf()  # clear figure
109plt.close()
110
111
112#
113##### real-time evolution of GPE #####
114def GPE(time, psi):
115    """
116    This function solves the complex-valued time-dependent GPE:
117    $$ i\dot\psi(t) = Hsp(t)\psi(t) + U |\psi(t)|^2 \psi(t) $$
118    """
119    # solve static part of GPE
120    psi_dot = Hsp.static.dot(psi) + U * np.abs(psi) ** 2 * psi
121    # solve dynamic part of GPE
122    for f, Hd in iteritems(Hsp.dynamic):
123        psi_dot += f(time) * Hd.dot(psi)
124    return -1j * psi_dot
125
126
127# define real time vector
128t = np.linspace(0.0, t_ramp, 101)
129# time-evolve state according to GPE
130psi_t = evolve(psi0, t[0], t, GPE, iterate=True, atol=1e-12, rtol=1e-12)
131#
132# display state evolution
133for i, psi in enumerate(psi_t):
134    # compute energy
135    E = (Hsp.matrix_ele(psi, psi, time=t[i]) + 0.5 * U * np.sum(np.abs(psi) ** 4)).real
136    # compute trap
137    kappa_trap = ramp(t[i], kappa_trap_i, kappa_trap_f, t_ramp) * (sites) ** 2
138    # plot wave function
139    plt.plot(
140        sites,
141        abs(psi0) ** 2,
142        color="r",
143        marker="s",
144        alpha=0.2,
145        label="$|\\psi_{\\mathrm{GS},j}|^2$",
146    )
147    plt.plot(sites, abs(psi) ** 2, color="b", marker="o", label="$|\\psi_j(t)|^2$")
148    plt.plot(sites, kappa_trap, "--", color="g", label="$\\mathrm{trap}$")
149    plt.ylim([-0.01, max(abs(psi0) ** 2) + 0.01])
150    plt.xlabel("$\\mathrm{lattice\\ sites}$", fontsize=14)
151    plt.title(
152        "$Jt=%0.2f,\\ E(t)-E_\\mathrm{GS}=%0.4fJ$" % (t[i], E - E_GS), fontsize=14
153    )
154    plt.legend(loc="upper right", fontsize=14)
155    plt.tight_layout()
156    plt.draw()  # draw frame
157    plt.pause(0.00005)  # pause frame
158    plt.clf()  # clear figure
159plt.close()