"""Dynamics in elongated harmonic traps.
See Expansion.md for details.
"""
import warnings
import numpy as np
from scipy.stats.mstats import gmean
from . import bec
__all__ = ["StateGPEdrZ", "u"]
_EPS = np.finfo(float).eps
[docs]
class StateGPEdrZ(bec.HOMixin, bec.StateScaleHO):
r"""Effective 1D model for the GPE in harmonic cigars.
This class combines the NPSEQ to derive an effective 1D equation for the dynamics of
a harmonically trapped cloud with the scaling solution of Castin and Dum to allow
for expansion dynamics.
"""
[docs]
beta = 2.0 # EoS = g * n**beta / beta
# Allows one to use the Mateo:2009's suggestion of minimizing the chemical
# potential. Note that setting this to True might break minimization.
[docs]
minimize_chemical_potential = False
[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
def __init__(self, Nxyz, Lxyz, **kw):
# Store specified grid, but only pass 1D grid to base
kw.update(Lxyz=Lxyz[:1], Nxyz=Nxyz[:1])
super().__init__(**kw)
[docs]
def init(self):
super().init()
# Need to update these since the mixing will use self.dim == 1.
w0s = np.asarray(self.get_ws(t=0))
lams = np.ones(len(w0s))
dlams = np.zeros(len(w0s))
ddlams = np.zeros(len(w0s))
self._lambda_cache = (0, lams, dlams, ddlams)
self.w0s_perp2 = w0s[1:] ** 2
self.w0_perp = gmean(w0s[1:])
self._symmetric_w0 = len(w0s) == 2 or np.allclose(np.diff(w0s[1:]), 0)
if not self._symmetric_w0:
beta = self.beta
A = self.hbar**2 / 2 / self.m
Bx, By = Bxy = self.m * w0s[1:] ** 2 / 2
if self.minimize_chemical_potential:
a = 1 / beta / (beta - 1)
else:
a = 1 / (beta - 1)
C2_n2 = (self.g / a / beta**2 / np.pi ** (beta - 1)) ** 2
axy = A * Bxy / Bx / By * (4 * Bx * By / C2_n2) ** (1 / beta)
self._sigma2_cache = dict(A=A, Bxy=Bxy, C2_n2=C2_n2, axy=axy, a=a)
[docs]
def get_sigma2s(self, abs_Phi2=None):
"""Return (sigma2x, sigma2y)."""
lams, dlams, ddlams = self.get_lambdas()
lam_xyz = np.prod(lams)
w0s = np.asarray(self.get_ws(t=0))
if abs_Phi2 is None:
abs_Phi2 = abs(self.get_psi()) ** 2
n1 = abs_Phi2 / lam_xyz
if self._symmetric_w0:
# CHECK
if self.minimize_chemical_potential:
a = self.m * self.g / 2 / np.pi / self.hbar**2
else:
a = self.m * self.g / 4 / np.pi / self.hbar**2
a_perp2 = self.hbar / self.m / self.w0_perp
sigma2x = sigma2y = a_perp2 * np.sqrt(1 + 2 * a * abs_Phi2)
else:
# Solve for sigma2 with Newton's method. See Sigmas.md
# The formula here refer to x and y to match those notes, but these actually
# refer to y and z axes here.
A = self.hbar**2 / 2 / self.m
Ax, Ay = Axy = np.divide(A, lams[1:])
Bx, By = Bxy = self.m * w0s[1:] ** 2 / 2
beta = self.beta
if self.minimize_chemical_potential:
a = 1 / beta / (beta - 1)
else:
a = 1 / (beta - 1)
C2_n2 = (self.g / a / beta**2 / np.pi ** (beta - 1)) ** 2
C2 = C2_n2 * n1 ** (2 * (beta - 1))
axy = Axy * Bxy / Bx / By * (4 * Bx * By / C2_n2) ** (1 / beta)
tmp = beta / (beta - 1)
axy_inv = n1[None, :] ** (2 / tmp) / axy[:, None]
# Globally Convergent Newton's method to solve for
# q**tmp = (1+sqrt(1+ax*q))*(1+sqrt(1+ay*q))
# g = 0 = tmp*log(q) - log(1+sqrt(1+ax*q)*(1+sqrt(1+ay*q))
# This is very bad on a GPU!
cond = self.xp.min(axy_inv, axis=0) > _EPS
inds = self.xp.where(cond)[0]
lamxy = 1 / (1 + axy_inv[:, inds]) ** 2
power = tmp - 0.5 * np.sum(lamxy, axis=0)
ax, ay = axy = 1 / axy_inv[:, inds]
q = self.xp.prod(1 + self.xp.sqrt(1 + axy), axis=0) ** (1 / power)
for n in range(5): # 5 iterations of Newton should be enough
_x, _y = 1 + self.xp.sqrt(1 + ax * q), 1 + np.sqrt(1 + ay * q)
g = tmp * self.xp.log(q) - self.xp.log(_x * _y)
dg = tmp / q - ax / 2 / (ax * q + _x) - ay / 2 / (ay * q + _y)
q -= g / dg
assert np.allclose(g, 0)
Bx, By = Bxy
sigma2 = ((Ax * Ay) / (Bx * By)) ** (1 / 4) + self.xp.zeros_like(n1)
sigma2[inds] = np.sqrt(
q ** (1 / (beta - 1)) * (C2[inds] / 4 / Bx / By) ** (1 / beta)
)
f2 = C2 / sigma2 ** (2 * (beta - 1))
f = self.xp.sqrt(f2)
sigma2x, sigma2y = [
(f + self.xp.sqrt(f2 + 4 * A * B)) / 2 / B for (A, B) in zip(Axy, Bxy)
]
return sigma2x, sigma2y
[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 = self.get_density()
lams, dlams, ddlams = self.get_lambdas()
lam_xyz = np.prod(lams)
if TF:
beta = self.beta
coeff = self.m * beta * self.w0_perp**2 / 2 / np.pi / self.g / (beta - 1)
n = (n_1D * coeff) ** (1 / self.beta)
# n = np.sqrt(n_1D * self.m * self.w0_perp**2 / np.pi / self.g)
else:
sigma2x, sigma2y = self.get_sigma2s()
sigma2 = np.sqrt(sigma2x * sigma2y)
n = n_1D / sigma2 / np.pi
return n / lam_xyz
[docs]
def get_V_GPU(self):
"""Return the complete potential `V` - internal and external."""
# We assume here that this is the full physical external potential, including
# any external harmonic trapping potentials evaluated with the physical
# coordinates X returned by self.get_xyz(). This does not contain any
# corrections from either the expanding frame or from the
V_ext = self.get_Vext_mu_GPU()
m = self.m
hbar = self.hbar
sigma2x, sigma2y = self.get_sigma2s()
lams, dlams, ddlams = self.get_lambdas()
lam_xyz = self.xp.prod(lams)
if self._symmetric_w0:
sigma2 = self.xp.sqrt(self.xp.prod([sigma2x, sigma2y], axis=0))
# beta = 2
if self.minimize_chemical_potential:
V_int_etc = m * self.w0_perp**2 * sigma2
else:
# V_int_etc_ = (
# hbar**2 / 2 / m / sigma2
# + m*w_perp2*sigma2/2.0
# + self.g * n/2/np.pi/sigma2
# )
V_int_etc = (
3 * m * self.w0_perp**2 / 2 * sigma2 - hbar**2 / 2 / m / sigma2
)
# assert np.allclose(V_int_etc, V_int_etc_)
else:
w2x, w2y = self.w0s_perp2
beta = self.beta
a = self._sigma2_cache["a"]
_A = (1 - beta * a) * hbar**2 / 4 / m
_B = (1 + beta * a) * m / 4
V_int_etc = _B * (w2x * sigma2x + w2y * sigma2y)
if _A != 0:
# Optimization for GPE with self.minimize_chemical potential
V_int_etc += _A * (1 / sigma2x + 1 / sigma2y)
return V_ext + V_int_etc / lam_xyz
[docs]
def get_energy_density(self):
"""Return the energy density."""
# Warning: this is not correct. It may not be real until summed. The
# correct energy density requires abs(grad psi)^2
hbar = self.hbar
m = self.m
g = self.g
y = self
psi = self.get_psi()
a = self.m * self.g / 4 / np.pi / self.hbar**2
w0_perp2 = self.w0_perp**2
a_perp4 = self.hbar**2 / self.m**2 / w0_perp2
abs_Phi2 = abs(psi) ** 2
sigma2x, sigma2y = self.get_sigma2s(abs_Phi2=abs_Phi2)
sigma2 = np.sqrt(sigma2x * sigma2y)
lams, dlams, ddlams = self.get_lambdas()
lam_xyz = np.prod(lams)
Ky = y.copy()
Ky.apply_laplacian(factor=self.K_factor)
K = psi.conj() * Ky.get_psi()
Vext = self.get_Vext_mu() * abs_Phi2
if self._symmetric_w0:
if self.minimize_chemical_potential:
raise NotImplementedError()
else:
# Should we call self.get_Eint()?
Vint = m * w0_perp2 / 2.0 / a * (sigma2**3 / a_perp4 - sigma2)
Vint_alt = (
hbar**2 / 2.0 / m / sigma2
+ m * w0_perp2 * sigma2 / 2.0
+ g * abs_Phi2 / 4.0 / np.pi / sigma2
) * abs_Phi2
else:
# To Do: Implement direct version and make sure it works with chemical
# potential minimization.
w2x, w2y = self.w0s_perp2
# a, beta = self._sigma2_cache["a"], self.beta
# _A = (1 - beta * a) * hbar**2 / 4 / m
# _B = (1 + beta * a) * m / 4
# V_int = _A * (1 / sigma2x + 1 / sigma2y)
# V_int += _B * (w2x * sigma2x + w2y * sigma2y)
Vint = Vint_alt = (
hbar**2 / 4.0 / m * (1 / sigma2x + 1 / sigma2y)
+ m / 4.0 * (w2x * sigma2x + w2y * sigma2y)
+ g * abs_Phi2 / 4.0 / np.pi / sigma2
) * abs_Phi2
assert np.allclose(Vint, Vint_alt)
self.Vint = Vint
self.Vint_alt = Vint_alt
self._l = locals()
return K + Vint / lam_xyz + Vext
[docs]
def get_n_TF(self, V_TF, V_ext=None, g=None):
"""Return the Thomas Fermi density profile n_1D from mu.
Arguments
---------
V_TF : float
Value of V(x_TF) where the density should vanish in the TF limit.
"""
if not self.initializing:
warnings.warn(
f"In get_n_TF() while not initializing: `{V_TF=}` might be incorrect."
)
zero = np.zeros(self.shape)
if g is None:
g = self.g
if V_ext is None:
V_ext = self.get_Vext_mu()
V = V_ext + zero
h = self.hbar
m = self.m
# See NPSEQ.md: hw should be the sum of the zero-point energies, hence we should
# use the arithmetic mean for hw.
w = np.mean(self.get_ws()[1:])
hw = h * w
mu_eff = V_TF + hw - V
mu_eff_hw = mu_eff / hw
if False:
sigma2w = h * (mu_eff_hw + np.sqrt(mu_eff_hw**2 + 3.0)) / (3 * m)
n_1D_ = 2 * np.pi * m * np.maximum(zero, sigma2w**2 - (h / m) ** 2) / g
if not self._symmetric_w0:
warnings.warn(
"get_n_TF() called with asymmetric wy/wz. Answer will be approximate."
)
beta = self.beta
if self.minimize_chemical_potential:
a = 1 / beta / (beta - 1)
else:
a = 1 / (beta - 1)
alpha = self.g
ab = a * beta
if ab == 1:
# Optimization for GPE with self.minimize_chemical potential
mwsigma2_h = mu_eff_hw
else:
mwsigma2_h = (mu_eff_hw + np.sqrt(mu_eff_hw**2 - (1 - ab**2))) / (1 + ab)
mwsigma2_h = np.maximum(1, mwsigma2_h)
sigma2 = mwsigma2_h / m / hw
n_1D = (
np.pi
* sigma2
* (a * beta**2 * hw / 2 / alpha * (mwsigma2_h - 1 / mwsigma2_h))
** (1 / (beta - 1))
)
if False:
assert np.allclose(n_1D, n_1D_)
return n_1D
[docs]
def get_mu_from_V_TF(self, V_TF):
"""Return the Thomas Fermi chemical potential from V_TF.
Arguments
---------
V_TF : float
External potential at the Thomas Fermi "radius". (The external
potential is evaluated at this position and this is used to get
`mu`.)
"""
mu = V_TF + self.hbar * self.w0_perp
# Here we assume that self.mu is subtracted in get_V().
if self.mu:
mu = mu + self.mu
if not self.initializing:
warnings.warn(
f"In get_mu_from_V_TF() while not initializing: `{mu=}` might be incorrect."
)
return mu
[docs]
def get_V_TF_from_mu(self, mu):
"""Return V_TF from the chemical potential mu.
Arguments
---------
mu : float
Physical chemical potential (i.e. what you would pass to the
minimizer).
"""
if not self.initializing:
warnings.warn(
f"In get_V_TF_from_mu() while not initializing: `{mu=}` might be incorrect."
)
V_TF = mu - self.hbar * self.w0_perp
# Here we assume that self.mu is subtracted in get_V().
if self.mu:
V_TF -= self.mu
return V_TF
class StateGPEdrZTwist_x(StateGPEdrZ, bec.StateTwist_x):
"""This version does not perform scaling along the x direction.
This makes it suitable for modified dispersions, twists, etc.
"""
########### Needs testing.