"""Dynamics in elongated harmonic traps.
See Expansion.md for details.
"""
import warnings
import numpy as np
from scipy.integrate import odeint
from scipy.stats.mstats import gmean
from . import bec2
from .bec2 import u
__all__ = ["StateGPEdrZ", "u"]
class Sigma2s:
"""Helper class to solve the equations for sigma using Newton's method."""
def __init__(self, Cs, Bs):
self.Cs = np.asarray(Cs)
self.Bs = np.asarray(Bs)
def f(self, xs):
xa, xb, xc = xs
return [
(1.0 - self.Cs[0] * xc) * xa**2 - self.Bs[0],
(1.0 - self.Cs[1] * xc) * xb**2 - self.Bs[1],
xc * (xa + xb) ** 2 - 1.0,
]
def J(self, xs):
xa, xb, xc = xs
zero = np.zeros_like(xa)
return [
[2 * (1.0 - self.Cs[0] * xc) * xa, zero, -self.Cs[0] * xa**2],
[zero, 2 * (1.0 - self.Cs[1] * xc) * xb, -self.Cs[1] * xb**2],
[2 * xc * (xa + xb), 2 * xc * (xa + xb), (xa + xb) ** 2],
]
def step(self, x0, f0=None):
if f0 is None:
f0 = self.f(x0)
x1 = x0 - np.einsum(
"...a->a...",
np.linalg.solve(
np.einsum("ab...->...ab", self.J(x0)), np.einsum("a...->...a", f0)
),
)
return x1
def solve3(self, x0=None):
if x0 is None:
x0 = np.ones((3,) + self.Cs.shape[1:], dtype=float)
else:
x0 = np.asarray(x0)
self.x0 = x0.copy()
for n in range(20):
f0 = self.f(x0)
if np.allclose(f0, 0, atol=1e-12):
x1 = x0
break
x1 = self.step(x0=x0, f0=f0)
x0 = x1
try:
assert np.allclose(self.f(x1), 0)
except Exception:
self.x1 = x1
import mmfutils
mmfutils._s = self
raise
return x1
def solve(self):
Bs = self.Bs
Cs = self.Cs
bad_order = Cs[0] < Cs[1]
inds = None
if np.any(bad_order):
inds = np.where(bad_order)[0]
Bs[0, inds], Bs[1, inds] = Bs[1, inds], Bs[0, inds]
Cs[0, inds], Cs[1, inds] = Cs[1, inds], Cs[0, inds]
b = (Cs[0] - Cs[1]) / Bs[1]
a = Bs[0] / Bs[1]
x = np.sqrt(a + b)
for n in range(10):
f0 = -a + x * (-2 * a + x * (1 - a - b + x * (2 + x)))
if np.allclose(f0, 0):
break
df0 = -2 * a + x * (2 * (1 - a - b) + x * (6 + x * 4))
x = x - f0 / df0
assert n < 9
if inds is not None:
x[inds] = 1.0 / x[inds]
Bs[0, inds], Bs[1, inds] = Bs[1, inds], Bs[0, inds]
Cs[0, inds], Cs[1, inds] = Cs[1, inds], Cs[0, inds]
sigma2_a = np.sqrt(Bs[0] + Cs[0] / (1 + 1.0 / x) ** 2)
sigma2_b = sigma2_a / x
return np.array([sigma2_a, sigma2_b])
[docs]
class StateGPEdrZ(bec2.State):
r"""Effective 1D model for an elongated cloud implementing a modified form
of the dr-GPE with dynamic rescaling in the y and z directions but not in
the x direction. The state here is $\Phi(Z, t)$ but `get_density()` has
been modified to include the correct scaling.
"""
[docs]
def get_ws(self, t=None):
"""Return the trapping frequencies at time t."""
if t is None:
t = self.t
# For backward compatibility - constant wx
if hasattr(self, "get_ws_perp"):
ws = (1.0,) + tuple(self.get_ws_perp(t=t))
return ws
raise NotImplementedError
[docs]
def init(self):
lams = np.ones(len(self.ws))
dlams = np.zeros(len(self.ws))
self._lambda_cache = (0, np.ravel([lams, dlams]))
super().init()
@property
[docs]
def w0_perp(self):
"""Average perpendicular frequency at time t=0."""
return gmean(self.get_ws(t=0)[1:])
[docs]
def get_sigma2s(self, abs_Phi2s=None):
if abs_Phi2s is None:
abs_Phi2s = self.get_density()
ms = self.ms[self.bcast]
gs = np.asarray(self.gs[:2])[self.bcast]
A = ms * self.w0_perp**2 / 2.0
ga, gb, gab = self.gs
Bs = gs * abs_Phi2s / 4 / np.pi + self.hbar**2 / 2.0 / ms
Cs = gab * abs_Phi2s[::-1, ...] / np.pi
return Sigma2s(Bs=Bs / A, Cs=Cs / A).solve()
xa, xb = self._get_sigma2s(abs_Phi2s=abs_Phi2s)
xc = 1.0 / (xa + xb) ** 2
xs = [xa, xb, xc]
self._sigma_helper = Sigma2s(Bs=Bs / A, Cs=Cs / A)
try:
sigma2s = self._sigma_helper.solve(xs)[:2]
except Exception:
sigma2s = self._sigma_helper.solve()[:2]
return sigma2s
[docs]
def get_central_density(self, TF=False):
"""Return the physical density (3d) along the central axis of the
trap.
Arguments
---------
TF : bool
If True, then assume the transverse cloud is a TF profile (otherwise
use the internal Gaussian anzatz.)
"""
n_1D = na_1D, nb_1D = self.get_density()
lam_perp2 = np.prod(self.get_lambdas()[1:])
if TF:
g = self.gs.mean()
n = np.sqrt(
n_1D * self.ms[self.bcast] * np.prod(self.get_ws(t=0)[1:]) / np.pi / g
)
return n / lam_perp2
else:
sigma2s = self.get_sigma2s()
return n_1D / sigma2s / lam_perp2 / np.pi
[docs]
def get_V_GPU(self):
"""Return the complete potential `V` - internal and external."""
ms = self.ms[self.bcast]
hbar = self.hbar
w0_perp2 = self.w0_perp**2
sigma2s = self.get_sigma2s()
lam_perp2 = self.xp.prod(self.get_lambdas()[1:])
ga, gb, gab = self.gs
gs = self.xp.asarray(self.gs[:2])[self.bcast]
psi = self.get_psi_GPU()
abs_Phi2s = abs(psi) ** 2
Va, Vb, Vab = self.get_Vext_mu_GPU()
V_int_etc = (
hbar**2 / 2 / ms / sigma2s
+ ms * w0_perp2 * sigma2s / 2.0
+ gs * abs_Phi2s / 2 / np.pi / sigma2s
+ gab * abs_Phi2s[::-1, ...] / np.pi / sigma2s.sum(axis=0)
) / lam_perp2
return (Va + V_int_etc[0], Vb + V_int_etc[1], Vab)
[docs]
def get_energy_density(self, a=True, b=True, ab=True):
"""Return the energy density.
Arguments
---------
a, b : bool
Include the energies associated with this species. This includes the kinetic
energy and the self-interaction, but no inter-species interactions.
ab : bool
Include only the interaction energy.
"""
# Warning: this is not correct. It may not be real until summed. The
# correct energy density requires abs(grad psi)^2
hbar = self.hbar
ms = self.ms[self.bcast]
y = self
psi_a, psi_b = psi = self.get_psi()
w0_perp2 = self.w0_perp**2
lam_perp2 = np.prod(self.get_lambdas()[1:])
gaa, gbb, gab = self.gs
gs = np.asarray(self.gs[:2])[self.bcast]
na, nb = abs_Phi2s = abs(psi) ** 2
sigma_a, sigma_b = sigma2s = self.get_sigma2s(abs_Phi2s=abs_Phi2s)
Ky = y.copy()
Ky.apply_laplacian(factor=self.K_factor)
Ka, Kb = psi.conj() * Ky.get_psi()
psi2_a, psi2_b = abs_Phi2s
Va_int, Vb_int = (
(
hbar**2 / 2.0 / ms / sigma2s
+ ms * w0_perp2 * sigma2s / 2.0
+ gs * abs_Phi2s / 4.0 / np.pi / sigma2s
)
* abs_Phi2s
) / lam_perp2
Vab_int = gab * abs_Phi2s.prod(axis=0) / np.pi / sigma2s.sum(axis=0) / lam_perp2
Va, Vb, Vab = self.get_Vext_mu()
energy_density = 0
Ea = Ka + Va_int + Va * na
Eb = Kb + Vb_int + Vb * nb
Eab = Vab_int + 2 * (psi_a.conj() * Vab * psi_b).real
energy_density = 0.0
if a:
energy_density += Ea
if b:
energy_density += Eb
if ab:
energy_density += Eab
return energy_density
[docs]
def _rhs(self, q, t):
"""RHS for lambda(t) ODE."""
lams, dlams = np.reshape(q, (2, len(self._lambda_cache[1]) // 2))
w0s = self.get_ws(t=0)
ws = self.get_ws(t=t)
ddlams = -lams * ws**2 + w0s**2 / np.prod(lams) / lams
return np.ravel([dlams, ddlams])
[docs]
def get_lambdas(self):
t = self.t
if t != self._lambda_cache[0]:
t0, q0 = self._lambda_cache
q1 = odeint(self._rhs, q0, [t0, t])[-1]
self._lambda_cache = (t, q1)
return self._lambda_cache[1][: len(self._lambda_cache[1]) // 2]
[docs]
def _get_n_TF(self, V_TF, V_ext=None, m=None, g=None, state=None):
"""Return the total TF density. Assumes gaa = gbb = gab and
populates only the lower band.
Arguments
---------
V_g : float
(V_TF - V)/g. The density will vanish where V=V_TF.
"""
if state is None:
state = self
else:
assert state is self
if not self.initializing:
warnings.warn(
f"In _get_n_TF() while not initializing: `{V_TF=}` might be incorrect."
)
if g is None:
g = np.mean(self.gs)
if m is None:
# Different masses not supported here.
assert np.allclose(self.ms[0], self.ms[1])
m = np.mean(self.ms)
if V_ext is None:
Va, Vb, Vab = self.get_Vext_mu()
V_ext = (Va + Vb) / 2.0
V = V_ext
h = self.hbar
w = self.w0_perp
hw = h * w
mu_eff_hw = (V_TF - V) / hw
mu_eff_hw += 1.0 # This is the extra hbar*w0_perp piece
sigma2w = h * (mu_eff_hw + np.sqrt(mu_eff_hw**2 + 3.0)) / (3.0 * m)
n_1D = 2 * np.pi * m * np.maximum(0, sigma2w**2 - (h / m) ** 2) / g
return n_1D
[docs]
def get_mus_from_Vs_TF(self, Vs_TF):
"""Return the corrected chemical potential from Vs_TF.
Arguments
---------
Vs_TF : (float, float)
External potentials at the Thomas Fermi "radius". (The external
potential is evaluated at this position and this is used to get
`mu`.)
"""
mus = np.asarray(Vs_TF) + self.hbar * self.w0_perp
# Here we assume that self.mus is subtracted in get_Vext_mu().
if self.mus is not None:
mus = mus + self.mus
if not self.initializing:
warnings.warn(
f"In get_mus_from_Vs_TF() while not initializing: `{mus=}` might be incorrect."
)
return mus
[docs]
def get_Vs_TF_from_mus(self, mus):
"""Return Vs_TF from the chemical potentials mus.
Arguments
---------
mus : (float, float)
Physical chemical potentials (i.e. what you would pass to the minimizer).
"""
if not self.initializing:
warnings.warn(
f"In get_Vs_TF_from_mus() while not initializing: `{mus=}` might be incorrect."
)
Vs_TF = np.asarray(mus) - self.hbar * self.w0_perp
# Here we assume that self.mu is subtracted in get_V().
if self.mus is not None:
Vs_TF = Vs_TF - self.mus
return Vs_TF