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
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
Cell In[1], line 5
1 import numpy as np, matplotlib.pyplot as plt
2 import gpe.utils, gpe.bec, gpe.minimize
3
4
----> 5 class StateMixin(gpe.bec.StateHOMixin):
6 def __init__(self, experiment, **kw):
7 self.experiment = experiment
8 super().__init__(**kw)
AttributeError: module 'gpe.bec' has no attribute 'StateHOMixin'
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:
\[\begin{gather*}
\mu > \hbar \frac{\omega_x + \omega_y + \omega_z}{2}, \qquad
V_{TF} = \mu - \hbar \frac{\omega_x + \omega_y + \omega_z}{2} > 0.
\end{gather*}\]
For UV convergence, the lattice spacing should be at most half the healing length:
\[\begin{gather*}
2\d{x} < h = \frac{\hbar}{\sqrt{2m\mu}}.
\end{gather*}\]
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.
\[\begin{gather*}
R_i = \omega_i^{-1}\sqrt{\frac{2V_{TF}}{m}}, \qquad
L_{i} > 2R_i + 5h + 5a_i
\end{gather*}\]
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=}")
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[2], line 1
----> 1 m, hbar = Experiment.m, Experiment.hbar
2 ws = Experiment.ws
3 E0 = hbar * sum(ws) / 2
4 V_TF = 0.1*E0
NameError: name 'Experiment' is not defined
e = Experiment()
s0 = e.get_state()
s0.plot()
%time s = e.get_initialized_state()
s.plot(log=True)
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[3], line 1
----> 1 e = Experiment()
2 s0 = e.get_state()
3 s0.plot()
4 get_ipython().run_line_magic('time', 's = e.get_initialized_state()')
NameError: name 'Experiment' is not defined
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()
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[4], line 4
1 from tqdm import tqdm
2 from pytimeode.evolvers import EvolverABM
3
----> 4 T = 2*np.pi / min(np.abs(e.dws)) / 4
5 Nt = 20
6 dT = T/Nt
7 dt = 0.1*s.t_scale
NameError: name 'e' is not defined
np.save('test_4.npy', ev.y.get_psi())
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[5], line 1
----> 1 np.save('test_4.npy', ev.y.get_psi())
NameError: name 'ev' is not defined
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)
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[6], line 2
1 fig, ax = plt.subplots()
----> 2 x = s.xyz[0].ravel()
3 t = [s.t for s in states]
4 n = np.array([s.get_density_x() for s in states])
5 mesh = ax.pcolormesh(t, x, (n - n[0]).T)
NameError: name 's' is not defined
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()
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[7], line 4
1 get_ipython().run_line_magic('load_ext', 'autoreload')
2 get_ipython().run_line_magic('autoreload', '2')
3 import gpe.tube
----> 4 class StateTube(StateMixin, gpe.tube.StateGPEdrZ):
5 pass
6
7 et = Experiment(State=StateTube)
NameError: name 'StateMixin' is not defined
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)
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[8], line 1
----> 1 print(s.get_N()/st.get_N() - 1)
2 assert np.allclose(s.get_N(), st.get_N(), rtol=0.008)
NameError: name 's' is not defined
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)
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[9], line 2
1 fig, ax = plt.subplots()
----> 2 x = st.xyz[0]
3 ax.plot(x, s.get_density_x(), "-", label='3D')
4 ax.plot(x, st.get_density_x(), ":", label='Tube')
5 ax.set(xlabel="x", ylabel="$n_{1D}$")
NameError: name 'st' is not defined
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()
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[10], line 1
----> 1 Nx, Ny, Nz = s.basis.Nxyz
2 x = s.xyz[0].ravel()
3 fig, ax = plt.subplots()
4 ax.plot(x, s.get_density()[:, Ny//2, Nz//2], '-', label="3D")
NameError: name 's' is not defined
statest = []
ev = EvolverABM(st, dt=dt)
for frame in tqdm(range(Nt)):
ev.evolve(steps)
statest.append(ev.get_y())
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[11], line 2
1 statest = []
----> 2 ev = EvolverABM(st, dt=dt)
3 for frame in tqdm(range(Nt)):
4 ev.evolve(steps)
5 statest.append(ev.get_y())
NameError: name 'st' is not defined
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])
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[12], line 2
1 fig, axs = plt.subplots(2, 1)
----> 2 x = s.xyz[0].ravel()
3 t = [s.t for s in states]
4 n = np.array([s.get_density_x() for s in states])
5 mesh = axs[0].pcolormesh(t, x, (n - n[0]).T)
NameError: name 's' is not defined
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])
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[13], line 5
1 from scipy.interpolate import InterpolatedUnivariateSpline
2
3 fig, axs = plt.subplots(2, 1)
4
----> 5 x = s.xyz[0].ravel()
6 t = [s.t for s in states]
7 n = np.array([s.get_density_x() for s in states])
8 mesh = axs[0].pcolormesh(t, x, (n - n[0]).T)
NameError: name 's' is not defined
st = statest[-1]
st.get_density() - st.get_density_x()
---------------------------------------------------------------------------
IndexError Traceback (most recent call last)
Cell In[14], line 1
----> 1 st = statest[-1]
2 st.get_density() - st.get_density_x()
IndexError: list index out of range
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()
The autoreload extension is already loaded. To reload it, use:
%reload_ext autoreload
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[16], line 4
1 get_ipython().run_line_magic('load_ext', 'autoreload')
2 get_ipython().run_line_magic('autoreload', '2')
3 import gpe.tube
----> 4 class StateTube(StateMixin, gpe.tube.StateGPEdrZ):
5 pass
6
7 et = Experiment(State=StateTube)
NameError: name 'StateMixin' is not defined
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_, ':')
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[17], line 2
1 ts = np.linspace(0, 3, 200)
----> 2 st._tsqs = []
3 lams = np.array([st.get_lambdas(t)[0] for t in ts])
4 ws = np.array([st.get_ws(t) for t in ts])
5 ts_, lams_ = ts, lams
NameError: name 'st' is not defined
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())
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
Cell In[18], line 5
1 import numpy as np, matplotlib.pyplot as plt
2 import gpe.utils, gpe.bec, gpe.minimize
3
4
----> 5 class State(gpe.bec.StateHOMixin, gpe.bec.StateBase):
6 def get_ws(self, t=None):
7 return (2.0, 10.0, 15.0)
8
AttributeError: module 'gpe.bec' has no attribute 'StateHOMixin'