Writing Tests#
Here we will walk through the process of writing some tests for code. Specifically, we will test the code in that implements the Dynamically Rescaled GPE (dr-GPE) and Non-Polynomial Schrödinger Equation (NPSEQ).
In this case, we are testing a rather complete system of code that generates a 1D approximation of a 3D simulation. A reasonably strategy to start might be to simulate both the 3D system and then the 1D approximation, making sure that they agree.
Of course, for tests, we want this to run quickly, so we will push the limits of the calculation, but first we just get something working. The most general code allows for dynamic rescaling of the trapping frequencies as well as dynamics in the \(x\) direction, so we start with a simple 3D simulation that models this.
import numpy as np, matplotlib.pyplot as plt
import gpe.utils, gpe.bec, gpe.minimize
class StateMixin(gpe.bec.StateHOMixin):
def __init__(self, experiment, **kw):
self.experiment = experiment
super().__init__(**kw)
def get_ws(self, t=None):
"""Return the trapping frequencies."""
if t is None:
t = self.t
e = self.experiment
ws = [
w * (1 + e.dw * np.sin(dw * t))
for w, dw in zip(e.ws, e.dws)
]
return ws
def _get_Vext_(self, gpu=True):
Vext = super()._get_Vext_()
xyz = self._xyz_ if gpu else self.xyz
# Could add pokey here.
return Vext
class State(StateMixin, gpe.bec.StateBase):
pass
class Experiment(gpe.utils.ExperimentBase):
hbar = m = 1
ws = (2.0, 10.0, 15.0) # Trapping frequencies
dws = (10.0, -12.0, 14.0) # Frequency to change trapping frequencies
dw = 0.1 # Fraction of amplitude of trapping frequency change.
# These choices are not obvious: see below for how we selected them.
Nxyz = (80, 45, 45)
Lxyz = (7.0, 4.0, 4.0)
mu = 15
g = 1.0
State = State
def init(self):
self.state_args = dict(
experiment=self,
hbar=self.hbar,
m=self.m,
Lxyz=self.Lxyz,
Nxyz=self.Nxyz,
mu=self.mu,
g=self.g)
def get_state(self):
return self.State(**self.state_args)
def get_initialized_state(self):
s0 = self.get_state()
s = gpe.minimize.MinimizeState(s0, fix_N=False).minimize(use_scipy=True)
state = self.get_state()
state.set_psi(s.get_psi())
return state
KeyboardInterrupt
Setting Parameters
We need to choose reasonable parameters to make sure the tests work and make sense. I want the tube to be long (quasi-1D), so choose \(\omega_x \ll \omega_{y,z}\). I also want to make sure that at least one radial mode is occupied in each direction:
For UV convergence, the lattice spacing should be at most half the healing length:
Finally, the box must be large enough so that it can hold the cloud: we can estimate this with the TF radius in each direction, but need to include a margin for the decay. If we want the trap ground state to decay by a factor of \(\epsilon \approx 10^{-12}\), then we need to make sure that \(e^{-x^2/a^2} < \epsilon\) or \(x < \sqrt{-\ln\epsilon}a \approx 5a\) where \(a = \sqrt{\hbar/m\omega}\) is the corresponding trap length. We should also include several healing lengths.
Starting with a numerical guess for the frequencies, we arrive at the following estimates:
m, hbar = Experiment.m, Experiment.hbar
ws = Experiment.ws
E0 = hbar * sum(ws) / 2
V_TF = 0.1*E0
mu = V_TF + E0
h = hbar / np.sqrt(2 * m * mu)
dx = h / 2.0
Rxyz = np.divide(np.sqrt(2 * V_TF / m), ws)
axyz = np.sqrt(np.divide(hbar/m, ws))
Lxyz = 2*np.maximum(Rxyz + 10*h, 5*axyz)
Nxyz = list(map(gpe.utils.get_good_N, Lxyz/dx))
print(f"{Lxyz=}, {Nxyz=}, {mu=}, {V_TF=}")
e = Experiment()
s0 = e.get_state()
s0.plot()
%time s = e.get_initialized_state()
s.plot(log=True)
This works, and we see we have a good estimate of the box size, but the lattice is pretty big, and this might take too long for an actual test. One strategy is to run a long simulation, then save the results and use those to test against.
from tqdm import tqdm
from pytimeode.evolvers import EvolverABM
T = 2*np.pi / min(np.abs(e.dws)) / 4
Nt = 20
dT = T/Nt
dt = 0.1*s.t_scale
steps = int(max(np.ceil(dT/dt), 2))
dt = dT / steps
states = []
ev = EvolverABM(s, dt=dt)
for frame in tqdm(range(Nt)):
ev.evolve(steps)
states.append(ev.get_y())
ev.y.plot()
np.save('test_4.npy', ev.y.get_psi())
fig, ax = plt.subplots()
x = s.xyz[0].ravel()
t = [s.t for s in states]
n = np.array([s.get_density_x() for s in states])
mesh = ax.pcolormesh(t, x, (n - n[0]).T)
fig.colorbar(mesh, ax=ax)
Tube Code#
%load_ext autoreload
%autoreload 2
import gpe.tube
class StateTube(StateMixin, gpe.tube.StateGPEdrZ):
pass
et = Experiment(State=StateTube)
st0 = et.get_state()
st0.plot()
%time st = et.get_initialized_state()
st.plot()
The first thing we should test is that the particle numbers are similar:
print(s.get_N()/st.get_N() - 1)
assert np.allclose(s.get_N(), st.get_N(), rtol=0.008)
The second test is that the integrated densities match:
fig, ax = plt.subplots()
x = st.xyz[0]
ax.plot(x, s.get_density_x(), "-", label='3D')
ax.plot(x, st.get_density_x(), ":", label='Tube')
ax.set(xlabel="x", ylabel="$n_{1D}$")
ax.legend()
assert np.allclose(s.get_density_x(), st.get_density_x(), atol=2e-4, rtol=0.008)
The third test is the central densities. Note: this is an approximation unless we solve exactly for the radial wavefunction. One must choose whether the result is more close to the TF approximation (large number of radial modes) or the HO approximation. In this case, we have only a single mode occupied, so the HO approximation is much better:
Nx, Ny, Nz = s.basis.Nxyz
x = s.xyz[0].ravel()
fig, ax = plt.subplots()
ax.plot(x, s.get_density()[:, Ny//2, Nz//2], '-', label="3D")
ax.plot(x, st.get_central_density(TF=True), '--', label="TF")
ax.plot(x, st.get_central_density(TF=False), ':', label="HO")
ax.set(xlabel="x", ylabel="central density")
ax.legend()
statest = []
ev = EvolverABM(st, dt=dt)
for frame in tqdm(range(Nt)):
ev.evolve(steps)
statest.append(ev.get_y())
fig, axs = plt.subplots(2, 1)
x = s.xyz[0].ravel()
t = [s.t for s in states]
n = np.array([s.get_density_x() for s in states])
mesh = axs[0].pcolormesh(t, x, (n - n[0]).T)
fig.colorbar(mesh, ax=axs[0])
x = st.xyz[0].ravel()
t = [st.t for st in statest]
n = np.array([st.get_density_x() for st in statest])
mesh = axs[1].pcolormesh(t, x, (n - n[0]).T)
fig.colorbar(mesh, ax=axs[1])
from scipy.interpolate import InterpolatedUnivariateSpline
fig, axs = plt.subplots(2, 1)
x = s.xyz[0].ravel()
t = [s.t for s in states]
n = np.array([s.get_density_x() for s in states])
mesh = axs[0].pcolormesh(t, x, (n - n[0]).T)
fig.colorbar(mesh, ax=axs[0])
X = []
t = [st.t for st in statest]
st = statest[0]
x0 = st.xyz[0]
n0 = st.get_density_x()
get_n0 = InterpolatedUnivariateSpline(x0, n0)
n_n0 = []
for st in statest:
x = st.xyz[0]
n_n0.append(st.get_density_x() - get_n0(x))
X.append(x)
n_n0 = np.array(n_n0)
T = np.array([t for x in st.xyz[0]])
X = np.transpose(X)
mesh = axs[1].pcolormesh(T, X, n_n0.T, shading='auto')
fig.colorbar(mesh, ax=axs[1])
st = statest[-1]
st.get_density() - st.get_density_x()
N0 = [np.trapz(st.get_density_x(), st.xyz[0]) for st in statest]
N1 = [np.trapz(st.get_density_x(), x0) for st in statest]
N0, N1
Strange Issue#
%load_ext autoreload
%autoreload 2
import gpe.tube
class StateTube(StateMixin, gpe.tube.StateGPEdrZ):
pass
et = Experiment(State=StateTube)
st = et.get_state()
ts = np.linspace(0, 3, 200)
st._tsqs = []
lams = np.array([st.get_lambdas(t)[0] for t in ts])
ws = np.array([st.get_ws(t) for t in ts])
ts_, lams_ = ts, lams
plt.plot(ts_, lams_, ':')
Previous Issues#
While developing the tests in this notebook, we discovered several issues in the code (see Notes.md). Here are some snipits of code that demonstrated these issues - they will be included as unit tests.
import numpy as np, matplotlib.pyplot as plt
import gpe.utils, gpe.bec, gpe.minimize
class State(gpe.bec.StateHOMixin, gpe.bec.StateBase):
def get_ws(self, t=None):
return (2.0, 10.0, 15.0)
def _get_Vext_(self, gpu=True):
Vext = super()._get_Vext_()
xyz = self._xyz_ if gpu else self.xyz
x = xyz[0]
# Previously this was true
Vext.flags['WRITEABLE'] = True
Vext += x**2 # This is dangerous because of mutation... Now should raise exception.
return Vext
s = State()
assert np.allclose(s.get_energy(), s.get_energy())