"""Dynamics in elongated harmonic traps.
See Expansion.ipynb for details.
"""
from __future__ import division
import warnings
import numpy as np
from scipy.integrate import odeint
from mmfutils.math.bases import CylindricalBasis
from . import bec, bec2
__all__ = ["StateAxialBase", "State2AxialBase", "u", "CylindricalBasis"]
[docs]
class StateAxialBase(bec.StateTwist_x):
"""3D model for an elongated cloud with axial symmetry.
For compatibility with the tube code, the user should still provide the
function `get_ws()`, but the assumption here is that the frequencies
are the same.
Expansion is implemented in the axial direction through a rescaling of the
coordinates according to the procedure described in Docs/Expansion.ipynb
developed by Y. Castin and R. Dum.
https://doi.org/10.1103/PhysRevLett.77.5315
"""
# Implementation Details
# ----------------------
# Here we assume that:
# * The basis is a CylindricalBasis, so that xyz = xr.
# * Only r is scaled
#
# Trick: The radial kinetic energy density needs to be scaled by a factor
# of 1/lam_perp**2 but the x kinetic energy should not have this. In
# order to do this we use a trick - the basis.laplacian function takes
# two arguments, an overall factor `K_factor` and `kx2` which is the
# square of the momentum along x. Here we include the factor of
# 1/lam_perp**2 in K_factor, but remove it from kx2.
#
# Note: E_max depends on K_factor, so this needs to be corrected for the
# k_x too.
def __init__(self, basis, **kw):
assert isinstance(basis, CylindricalBasis)
super().__init__(basis=basis, **kw)
[docs]
def init(self):
lam_perp = 1
dlam_perp = 0
self._lambda_cache = (0, np.ravel([lam_perp, dlam_perp]))
super().init()
@property
[docs]
def K_factor(self):
lam_perp2 = self.get_lambda() ** 2
return self._K_factor / lam_perp2
@K_factor.setter
def K_factor(self, K_factor):
self._K_factor = K_factor
@property
[docs]
def E_max(self):
# Needs to have the correction removed since k_x does not change.
lam_perp2 = self.get_lambda() ** 2
return max(abs(self.K_factor) * self.basis.k_max**2) * lam_perp2
@property
[docs]
def kx2(self):
"""Overload the bec.StateTwist_x.kx2 variable to include the factor of
lam_perp**2 to compensate for K_factor."""
lam_perp2 = self.get_lambda() ** 2
kx2 = getattr(self, "_kx2", self.basis.kx**2)
return kx2 * lam_perp2
@kx2.setter
def kx2(self, kx2):
self._kx2 = kx2
[docs]
def get_xyz_GPU(self):
# Rescale to return physical coordinates.
lam_perp = self.get_lambda()
x, r = self.basis.xyz
return (x, r * lam_perp)
@property
[docs]
def metric(self):
lam_perp2 = self.get_lambda() ** 2
return self.basis.metric * lam_perp2
[docs]
def get_ws(self, t=None):
"""Return the perpendicular frequency at time t."""
if t is None:
t = self.t
raise NotImplementedError
[docs]
def get_psi_GPU(self):
# Includes scaling factor
lam_perp = self.get_lambda()
return super().get_psi_GPU() / lam_perp
[docs]
def set_psi(self, psi):
# Includes scaling factor
lam_perp = self.get_lambda()
return super().set_psi(psi * lam_perp)
[docs]
def get_density_x(self, n=None):
# We need an additional factor of lam_perp2 to account for the changed
# metric since we use the basis.integrate1 which does not know about
# the lambda factor.
lam_perp2 = self.get_lambda() ** 2
if n is None:
n = self.get_density()
n = n * lam_perp2
n = self.basis.integrate1(n)
return n
@property
[docs]
def w0_perp(self):
"""Average perpendicular frequency at time t=0."""
# Maybe this should be the geometric mean?
ws_perp = self.get_ws(t=0)[1:]
assert np.allclose(ws_perp[0], ws_perp)
return ws_perp[0]
[docs]
def _get_Vext(self):
"""Correct the external potentials."""
w0_perp2 = self.w0_perp**2
ws_perp = self.get_ws()[1:]
assert np.allclose(ws_perp[0], ws_perp)
w_perp2 = ws_perp[0] ** 2
x, R = self.get_xyz()
lam_perp = self.get_lambda()
r = R / lam_perp
V_corr = self.m / 2.0 * (w0_perp2 * r**2 / lam_perp**2 - w_perp2 * R**2)
return self.get_Vext() + V_corr
[docs]
def _rhs(self, q, t):
"""RHS for lambda(t) ODE."""
lam, dlam = q
w0s = self.get_ws(t=0)[1:]
ws = self.get_ws(t=t)[1:]
assert np.allclose(ws[0], ws)
assert np.allclose(w0s[0], w0s)
w0 = w0s[0]
w = ws[0]
ddlam = -lam * w**2 + w0**2 / lam**3
return (dlam, ddlam)
[docs]
def get_lambda(self):
t = self.t
if t != self._lambda_cache[0]:
t0, q0 = self._lambda_cache
q1 = odeint(self._rhs, q0, [t0, t], rtol=1e-8, atol=1e-8)[-1]
self._lambda_cache = (t, q1)
return self._lambda_cache[1][0]
[docs]
def _plot(self, ratios=(1, 2), colorbar_fraction=0.05, grid=None):
from matplotlib import pyplot as plt
from mmfutils.plot import imcontourf
from .plot_utils import MPLGrid
if grid is None:
grid = MPLGrid(space=0)
else:
grid = grid.grid(space=0)
X, R = self.get_xyz()
lam_perp = self.get_lambda()
# r = R/lam_perp
n2 = self.get_density()
n1 = self.get_density_x()
V = self._get_Vext()
g1 = grid.grid(ratios[0], direction="right", space=0)
ax1 = g1.next()
ax1.plot(X / u.micron, n1)
plt.twinx()
plt.plot(R.ravel() / u.micron, V[V.shape[0] // 2, :], scalex=False)
plt.ylabel("n_1D")
E = self.get_energy()
N = self.get_N()
lam_perp = self.get_lambda()
plt.title(
"t={:2g}ms, N={:.5g}, E={:.5g}, lam={:.2g}".format(
self.t / u.ms, N, E, lam_perp
)
)
cax1 = g1.next(colorbar_fraction)
cax1.set_visible(False)
g2 = grid.grid(ratios[1], direction="right", space=0)
g2.next()
imcontourf(X / u.micron, R / u.micron, n2, aspect=1)
plt.ylim(0, 10)
plt.xlabel("x [micron]")
plt.ylabel("R [micron]")
cax2 = g2.next(colorbar_fraction)
plt.colorbar(cax=cax2, label="n_3D")
[docs]
def plot(self, grid=None, axs=None):
from matplotlib import pyplot as plt
if axs is None:
from .plot_utils import MPLGrid
if grid is None:
grid = MPLGrid(space=0)
else:
grid = grid.grid(space=0)
ax1 = grid.next()
axs = [ax1]
ax1 = axs[0]
X, R = self.get_xyz()
lam_perp = self.get_lambda()
# r = R/lam_perp
n2 = self.get_density()
ax1.plot(X / u.micron, n2[:, 0])
ax1.plot(R.ravel() / u.micron, n2[n2.shape[0] // 2, :])
ax1.set(ylabel="n")
E = self.get_energy()
N = self.get_N()
lam_perp = self.get_lambda()
plt.title(
"t={:2g}ms, N={:.5g}, E={:.5g}, lam={:.4g}".format(
self.t / u.ms, N, E, lam_perp
)
)
return axs
[docs]
class State2AxialBase(bec2.StateBase):
"""3D model for an elongated cloud with axial symmetry.
For compatibility with the tube code, the user should still provide the
function `get_ws()`, but the assumption here is that the frequencies
are the same.
Expansion is implemented in the axial direction through a rescaling of the
coordinates according to the procedure described in Docs/Expansion.ipynb
developed by Y. Castin and R. Dum.
https://doi.org/10.1103/PhysRevLett.77.5315
"""
# Implementation Details
# ----------------------
# Here we assume that:
# * The basis is a CylindricalBasis, so that xyz = xr.
# * Only r is scaled
#
# Trick: The radial kinetic energy density needs to be scaled by a factor
# of 1/lam_perp**2 but the x kinetic energy should not have this. In
# order to do this we use a trick - the basis.laplacian function takes
# two arguments, an overall factor `K_factor` and `kx2` which is the
# square of the momentum along x. Here we include the factor of
# 1/lam_perp**2 in K_factor, but remove it from kx2.
#
# Note: E_max depends on K_factor, so this needs to be corrected for the
# k_x too.
def __init__(self, basis, **kw):
assert isinstance(basis, CylindricalBasis)
super().__init__(basis=basis, **kw)
[docs]
def init(self):
lam_perp = 1
dlam_perp = 0
self._lambda_cache = (0, np.ravel([lam_perp, dlam_perp]))
super().init()
@property
[docs]
def K_factor(self):
lam_perp2 = self.get_lambda() ** 2
return self._K_factor / lam_perp2
@K_factor.setter
def K_factor(self, K_factor):
self._K_factor = K_factor
@property
[docs]
def E_max(self):
# Needs to have the correction removed since k_x does not change.
lam_perp2 = self.get_lambda() ** 2
return max(abs(self.K_factor) * self.basis.k_max**2) * lam_perp2
@property
[docs]
def kx2(self):
"""Overload the bec2.State.kx2 variable to include the factor of
lam_perp**2 to compensate for K_factor."""
lam_perp2 = self.get_lambda() ** 2
return self._kx2 * lam_perp2
@kx2.setter
def kx2(self, kx2):
self._kx2 = kx2
[docs]
def get_xyz_GPU(self):
# Rescale to return physical coordinates.
lam_perp = self.get_lambda()
x, r = self.basis.xyz
return (x, r * lam_perp)
@property
[docs]
def metric(self):
lam_perp2 = self.get_lambda() ** 2
return self.basis.metric * lam_perp2
[docs]
def get_ws(self, t=None):
"""Return the perpendicular frequency at time t."""
if t is None:
t = self.t
raise NotImplementedError
[docs]
def get_psi_GPU(self):
# Includes scaling factor
lam_perp = self.get_lambda()
return super().get_psi_GPU() / lam_perp
[docs]
def set_psi(self, psi):
# Includes scaling factor
lam_perp = self.get_lambda()
return super().set_psi(psi * lam_perp)
[docs]
def get_density_x(self, ns=None):
# We need an additional factor of lam_perp2 to account for the changed
# metric since we use the basis.integrate1 which does not know about
# the lambda factor.
lam_perp2 = self.get_lambda() ** 2
if ns is None:
ns = self.get_density()
na, nb = ns * lam_perp2
na, nb = list(map(self.basis.integrate1, (na, nb)))
return na, nb
[docs]
def _get_energy_density(self):
# Warning: this is not correct. It may not be real until summed. The
# correct energy density requires abs(grad psi)^2
# This version corrects for the lambda scaling...
y = self
psi_a, psi_b = psi = self.get_psi()
na, nb = self.get_density()
Ky = y.copy()
Ky.apply_laplacian(factor=self.K_factor)
K = (psi.conj() * Ky.get_psi()).sum(axis=0)
gaa, gbb, gab = self.gs
Va, Vb, Vab = self.get_Vext_mu()
Vint = gaa * na**2 / 2.0 + gbb * nb**2 / 2.0 + gab * na * nb
Vext = Va * na + Vb * nb + 2 * (psi_a.conj() * Vab * psi_b).real
return K + Vint + Vext
@property
[docs]
def w0_perp(self):
"""Average perpendicular frequency at time t=0."""
# Maybe this should be the geometric mean?
ws_perp = self.get_ws(t=0)[1:]
assert np.allclose(ws_perp[0], ws_perp)
return ws_perp[0]
[docs]
def _get_Vext(self):
"""Correct the external potentials."""
Va, Vb, Vab = self.get_Vext()
ms = self.ms[self.bcast]
w0_perp2 = self.w0_perp**2
ws_perp = self.get_ws()[1:]
assert np.allclose(ws_perp[0], ws_perp)
w_perp2 = ws_perp[0] ** 2
x, R = self.get_xyz()
lam_perp = self.get_lambda()
r = R / lam_perp
# V_corr = ms*(w0_perp2 - w_perp2)/2.0*R**2/lam_perp2
V_corr = ms / 2.0 * (w0_perp2 * r**2 / lam_perp**2 - w_perp2 * R**2)
return (Va + V_corr[0], Vb + V_corr[1], Vab)
[docs]
def _rhs(self, q, t):
"""RHS for lambda(t) ODE."""
lam, dlam = q
w0s = self.get_ws(t=0)[1:]
ws = self.get_ws(t=t)[1:]
assert np.allclose(ws[0], ws)
assert np.allclose(w0s[0], w0s)
w0 = w0s[0]
w = ws[0]
ddlam = -lam * w**2 + w0**2 / lam**3
return (dlam, ddlam)
[docs]
def get_lambda(self):
t = self.t
if t != self._lambda_cache[0]:
t0, q0 = self._lambda_cache
q1 = odeint(self._rhs, q0, [t0, t], rtol=1e-8, atol=1e-8)[-1]
self._lambda_cache = (t, q1)
return self._lambda_cache[1][0]
[docs]
def _plot(self, ratios=(1, 2), colorbar_fraction=0.05, grid=None):
from matplotlib import pyplot as plt
from mmfutils.plot import imcontourf
from .plot_utils import MPLGrid
if grid is None:
grid = MPLGrid(space=0)
else:
grid = grid.grid(space=0)
X, R = self.get_xyz()
lam_perp = self.get_lambda()
# r = R / lam_perp
na2, nb2 = self.get_density()
na1, nb1 = self.get_density_x()
Va, Vb, Vab = self._get_Vext()
V = Va + Vb
g1 = grid.grid(ratios[0], direction="right", space=0)
ax1 = g1.next()
ax1.plot(X / u.micron, na1 + nb1)
plt.twinx()
plt.plot(R.ravel() / u.micron, V[V.shape[0] // 2, :], scalex=False)
plt.ylabel("n_1D")
E = self.get_energy()
N = self.get_N()
lam_perp = self.get_lambda()
plt.title(
"t={:2g}ms, N={:.5g}, E={:.5g}, lam={:.2g}".format(
self.t / u.ms, N, E, lam_perp
)
)
cax1 = g1.next(colorbar_fraction)
cax1.set_visible(False)
g2 = grid.grid(ratios[1], direction="right", space=0)
g2.next()
imcontourf(X / u.micron, R / u.micron, na2 + nb2, aspect=1)
plt.ylim(0, 10)
plt.xlabel("x [micron]")
plt.ylabel("R [micron]")
cax2 = g2.next(colorbar_fraction)
plt.colorbar(cax=cax2, label="n_3D")
[docs]
def plot(self, grid=None):
from matplotlib import pyplot as plt
# from mmfutils.plot import imcontourf
from .plot_utils import MPLGrid
if grid is None:
grid = MPLGrid(space=0)
else:
grid = grid.grid(space=0)
X, R = self.get_xyz()
lam_perp = self.get_lambda()
# r = R / lam_perp
na2, nb2 = self.get_density()
na1, nb1 = self.get_density_x()
Va, Vb, Vab = self._get_Vext()
# V = Va + Vb
ax1 = grid.next()
ax1.plot(X / u.micron, (na2 + nb2)[:, 0])
ax1.plot(R.ravel() / u.micron, (na2 + nb2)[na2.shape[0] // 2, :])
plt.ylabel("n")
E = self.get_energy()
N = self.get_N()
lam_perp = self.get_lambda()
plt.title(
"t={:2g}ms, N={:.5g}, E={:.5g}, lam={:.4g}".format(
self.t / u.ms, N, E, lam_perp
)
)
class StateAxial(bec.HOMixin, StateAxialBase):
"""State with HO potential for backward compatibility."""
def __init__(self, *v, **kw):
warnings.warn(
"StateAxial is Deprecated: use HOMixin and StateAxialBase",
DeprecationWarning,
)
super().__init__(*v, **kw)
class State2Axial(bec2.HOMixin, State2AxialBase):
"""State with HO potential for backward compatibility."""
def __init__(self, *v, **kw):
warnings.warn(
"State2Axial is Deprecated: use HOMixin and State2AxialBase",
DeprecationWarning,
)
super().__init__(*v, **kw)