import mmf_setup;mmf_setup.nbinit()
from gpe.imports import *
This cell adds /home/docs/checkouts/readthedocs.org/user_builds/gpe/checkouts/latest/src to your path, and contains some definitions for equations and some CSS for styling the notebook. If things look a bit strange, please try the following:
- Choose "Trust Notebook" from the "File" menu.
- Re-execute this cell.
- Reload the notebook.
[I 03:35:35 numexpr.utils] NumExpr defaulting to 2 threads.
[I 03:35:35 root] Patching zope.interface.document.asReStructuredText to format code
2D Turbulence#
Here we reproduce the results of [Zhao et al., 2025], where they stir a 2D BEC and claim to measure a Kolmogorov decay spectrum.
\(N = 2\times 10^{5}\) atoms.
\(R = r_{\mathrm{tr}} = \qty{22}{micron}\)
from gpe.imports import *
from gpe.bec import StateBase, u
from gpe.utils import ExperimentBase, get_good_N, _GPU
from gpe.minimize import MinimizeState
from mmfutils.math.special import mstep
@_GPU.add_non_GPU_methods
class State(StateBase):
def __init__(self, experiment, kcut, kappa0, **kw):
self.experiment = experiment
self.kcut = kcut
self.kappa0 = kappa0
super().__init__(**kw)
def init(self):
super().init()
ks = s.basis.kx
mask = np.logical_and(ks <= self.kcut, ks >= -self.kcut)
self._phase_k = np.exp(-1j * self.kappa0 * np.where(mask, 0, 1))
def get_Vext_GPU(self):
V_ext = self.experiment.get_Vext(state=self)
if (self.initializing or self.t < 0) and getattr(self, "mu", None):
V_ext -= self.mu
return V_ext
def _apply_dissipation(self, dy):
"""Return applying dissipation."""
fftn, ifftn = self.basis.fftn, self.basis.ifftn
dy.set_psi(ifftn(self._phase_k * fftn(dy.get_psi())))
def compute_dy_dt(self, dy, subtract_mu=True):
super().compute_dy_dt(dy, subtract_mu=subtract_mu)
if not (self.initializing or self.t < 0):
self._apply_dissipation(dy)
return dy
class Experiment(ExperimentBase):
hbar = 1
m = u.m_Rb87
w_z = 2*np.pi * (220*u.Hz)
R_micron = 22.0
R_Rx = 0.9
g = 1
V0_mu = 100.0
V_dr_healing_length = 1.0
mu = None
healing_length_micron = 2.0
dx_healing_length = 0.1
n0_micron2 = 200.0 # Background density
# Stirring Potential
Vstir_Hz = 25
Vstir_R_micron = 3.5
Vstir_dr_healing_length = 1.0
Vstir_mu = 10.0
t_stir_ms = 16
rng_seed = 0
# Dissipation
kcut_invmicron = 5
kappa0 = 0.02
t_final_ms = 56
State = State
def init(self):
if self.mu is not None and self.healing_length_micron is None:
self.healing_length = np.sqrt(self.hbar**2 / 2 / self.m / self.mu)
mu = self.mu
elif self.mu is None and self.healing_length_micron is not None:
self.healing_length = self.healing_length_micron * u.micron
mu = self.hbar**2 / 2 / self.m / self.healing_length**2
else:
raise ValueError("Must specify (only) one of `mu` or `healing_length_micron`")
dx = self.dx_healing_length * self.healing_length
self.V_R = self.R_micron * u.micron
Rx = self.V_R / self.R_Rx
Nx = get_good_N(2*Rx/dx)
n0 = self.n0_micron2 / u.micron**2
kcut = self.kcut_invmicron / u.micron
g = mu / n0
self.state_args = dict(kcut=kcut, kappa0=self.kappa0, Nxyz=(Nx, Nx), Lxyz=(2*Rx, 2*Rx),
g=g, mu=mu, m=self.m, hbar=self.hbar)
self.V0 = self.V0_mu * mu
self.V_dr = self.V_dr_healing_length * self.healing_length
self._Vtrap = None
self.Vstir_w = 2*np.pi * self.Vstir_Hz * u.Hz
self.Vstir_R = self.Vstir_R_micron * u.micron
self.Vstir_dr = self.Vstir_dr_healing_length * self.healing_length
self.Vstir_V0 = self.Vstir_mu * mu
self._n_R_changes = 0
self.rng = np.random.default_rng(seed=self.rng_seed)
self._R_stir = 13.5 * u.micron # initial radius of stirring
self.t_stir = self.t_stir_ms * u.ms
self.t_final = self.t_final_ms * u.ms
def get_Vext(self, state):
"""Return the external potential."""
if self._Vtrap is None:
self._Vtrap = self.get_Vtrap(state=state)
V_ext = self._Vtrap + self.get_Vstir(state=state)
return V_ext
def get_Vtrap(self, state):
"""Return the static trapping potential."""
x, y = state.get_xyz()
r = np.sqrt(x**2 + y**2)
return self._Vstep(r, R=self.V_R, dr=self.V_dr)
def _Vstep(self, r, R, dr):
"""Smooth step at R of width dr."""
return mstep(r - (R - dr/2), dr)
def _Vrod(self, r):
"""Stirring rod potential (central)."""
return 1-self._Vstep(r, R=self.Vstir_R, dr=self.Vstir_dr)
def get_Vstir(self, state):
"""Return the dynamic stirring potential."""
t = state.t
if t > self.t_stir:
return 0
th0 = -np.pi/4
dth = self.Vstir_w * t
# TODO: Need to make these random transitions smooth?
if t >= 400 * u.s * 1e-6 * self._n_R_changes:
self._n_R_changes += 1
self._R_stir = self.rng.uniform(12, 15) * u.micron
z0, z1 = self._R_stir * np.exp([1j*(th0 + dth), 1j*(th0 - dth)])
x, y = state.get_xyz()
z = x + 1j*y
return self.Vstir_V0 * (self._Vrod(abs(z-z0)) + self._Vrod(abs(z-z1)))
def get_state(self):
state = State(experiment=self, **self.state_args)
return state
def get_initialized_state(self):
state0 = self.get_state()
m = MinimizeState(state0)
print(m.check())
state1 = m.minimize(use_scipy=True)
state = self.get_state()
state.set_psi(state1.get_psi())
return state
e = Experiment()
s = e.get_initialized_state()
s.plot()
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[2], line 165
161 state.set_psi(state1.get_psi())
162 return state
163
164 e = Experiment()
--> 165 s = e.get_initialized_state()
166 s.plot()
Cell In[2], line 156, in Experiment.get_initialized_state(self)
155 def get_initialized_state(self):
--> 156 state0 = self.get_state()
157 m = MinimizeState(state0)
158 print(m.check())
159 state1 = m.minimize(use_scipy=True)
Cell In[2], line 152, in Experiment.get_state(self)
151 def get_state(self):
--> 152 state = State(experiment=self, **self.state_args)
153 return state
Cell In[2], line 14, in State.__init__(self, experiment, kcut, kappa0, **kw)
10 def __init__(self, experiment, kcut, kappa0, **kw):
11 self.experiment = experiment
12 self.kcut = kcut
13 self.kappa0 = kappa0
---> 14 super().__init__(**kw)
File ~/checkouts/readthedocs.org/user_builds/gpe/checkouts/latest/src/gpe/bec.py:1134, in GPEMixin.__init__(self, g, **kw)
1132 if g is not None:
1133 self.g = g
-> 1134 super().__init__(**kw)
File ~/checkouts/readthedocs.org/user_builds/gpe/checkouts/latest/src/gpe/bec.py:464, in _StateBase.__init__(self, basis, Nxyz, Lxyz, symmetric_grid, t, hbar, m, mu, x_TF, cooling_phase, constraint, PGPE, kc_kmax, basis_args, t_final, Omega, **kw)
461 if t_final is not None:
462 self.t_final = t_final
--> 464 super().__init__(**kw)
466 # Once the state is initialized, we can set the initial state.
467 self.set_initial_state()
File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/mmfutils/containers.py:114, in ObjectBase.__init__(self, **kw)
112 if "picklable_attributes" not in self.__dict__:
113 self.picklable_attributes = sorted(_k for _k in self.__dict__)
--> 114 self.init()
Cell In[2], line 19, in State.init(self)
16 def init(self):
17 super().init()
18
---> 19 ks = s.basis.kx
20 mask = np.logical_and(ks <= self.kcut, ks >= -self.kcut)
21 self._phase_k = np.exp(-1j * self.kappa0 * np.where(mask, 0, 1))
NameError: name 's' is not defined
from pytimeode.evolvers import EvolverABM
from mmfutils.contexts import FPS
ev = EvolverABM(s, dt=0.1*s.t_scale)
steps = 200
loops = int(e.t_final / ev.dt / steps)
for frame in FPS(loops, timeout=10):
ev.evolve(steps)
plt.clf()
ev.y.plot()
clear_output(wait=True)
display(plt.gcf())
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[3], line 4
1 from pytimeode.evolvers import EvolverABM
2 from mmfutils.contexts import FPS
3
----> 4 ev = EvolverABM(s, dt=0.1*s.t_scale)
5
6 steps = 200
7 loops = int(e.t_final / ev.dt / steps)
NameError: name 's' is not defined
s0 = e.get_initialized_state()
rng = np.random.default_rng(seed=2)
def get_psi(s0=s0, p=0.1, steps=200, rng=rng):
s = e.get_state()
s.set_psi(s0.get_psi() * np.exp(1j*p*rng.normal(size=s0.shape)))
ev = EvolverABM(s, dt=0.2*s.t_scale)
ev.evolve(steps)
return ev.y.get_psi()
psis = [get_psi() for n in range(10)]
n_ = np.abs(np.mean(psis, axis=0))**2
print(s0.get_N(), s0.integrate(n_))
#n = get_n()
#plt.imshow(n)
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[4], line 1
----> 1 s0 = e.get_initialized_state()
2 rng = np.random.default_rng(seed=2)
3 def get_psi(s0=s0, p=0.1, steps=200, rng=rng):
4 s = e.get_state()
Cell In[2], line 156, in Experiment.get_initialized_state(self)
155 def get_initialized_state(self):
--> 156 state0 = self.get_state()
157 m = MinimizeState(state0)
158 print(m.check())
159 state1 = m.minimize(use_scipy=True)
Cell In[2], line 152, in Experiment.get_state(self)
151 def get_state(self):
--> 152 state = State(experiment=self, **self.state_args)
153 return state
Cell In[2], line 14, in State.__init__(self, experiment, kcut, kappa0, **kw)
10 def __init__(self, experiment, kcut, kappa0, **kw):
11 self.experiment = experiment
12 self.kcut = kcut
13 self.kappa0 = kappa0
---> 14 super().__init__(**kw)
File ~/checkouts/readthedocs.org/user_builds/gpe/checkouts/latest/src/gpe/bec.py:1134, in GPEMixin.__init__(self, g, **kw)
1132 if g is not None:
1133 self.g = g
-> 1134 super().__init__(**kw)
File ~/checkouts/readthedocs.org/user_builds/gpe/checkouts/latest/src/gpe/bec.py:464, in _StateBase.__init__(self, basis, Nxyz, Lxyz, symmetric_grid, t, hbar, m, mu, x_TF, cooling_phase, constraint, PGPE, kc_kmax, basis_args, t_final, Omega, **kw)
461 if t_final is not None:
462 self.t_final = t_final
--> 464 super().__init__(**kw)
466 # Once the state is initialized, we can set the initial state.
467 self.set_initial_state()
File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/mmfutils/containers.py:114, in ObjectBase.__init__(self, **kw)
112 if "picklable_attributes" not in self.__dict__:
113 self.picklable_attributes = sorted(_k for _k in self.__dict__)
--> 114 self.init()
Cell In[2], line 19, in State.init(self)
16 def init(self):
17 super().init()
18
---> 19 ks = s.basis.kx
20 mask = np.logical_and(ks <= self.kcut, ks >= -self.kcut)
21 self._phase_k = np.exp(-1j * self.kappa0 * np.where(mask, 0, 1))
NameError: name 's' is not defined