"""Dynamics in flattened harmonic traps."""
##############################################################################
# IMPORTS
from __future__ import division
from operator import truediv
import numpy as np
from scipy.integrate import odeint
# from scipy.optimize import brentq
from scipy.stats.mstats import gmean
import gpe.bec
from gpe.utils import get_good_N, ExperimentBase, GPUHelper
from gpe.bec import HOMixin, StateTwist_x
from gpe.minimize import MinimizeState
# from mmfutils.math.bases import PeriodicBasis
# from mmfutils.optimize import bracket_monotonic
from pytimeode.evolvers import EvolverABM
import pytest
piston = pytest.importorskip("piston")
from piston.mixins import EvolveMixin
u = gpe.bec.u
__all__ = ["StatePan", "Experiment"]
##############################################################################
# States
class StateGPEdxy(HOMixin, StateTwist_x):
r"""Effective 2D model for a pancake cloud implementing a modified form
of the dr-GPE with dynamic rescaling in the x direction but not in the y
and z directions. The state here is $\psi(x, t)$ but `get_density()` has
been modified to include the correct scaling.
"""
experiment = None
def init(self):
lams = np.ones(len(self.ws) - 2) # -2 since we only expand along x.
dlams = np.zeros(len(self.ws) - 2) # -2 since we only expand along x.
self._lambda_cache = (0, np.ravel([lams, dlams]))
super().init()
########################################################################
# Methods defining NPSEQ (pancake) formalism
def get_ws_perp(self, t=None):
"""Return the perpendicular frequencies at time t.
Required method. User must define.
"""
ws_perp = None
if t is None:
t = self.t
try:
ws_perp = self.experiment.get("ws", t_=t / self.experiment.t_unit)[-1:]
except:
raise NotImplementedError
return ws_perp
@property
def w0_perp(self):
"""Average perpendicular frequency at time t=0."""
return gmean(self.get_ws_perp(t=0)) # FIXME: Maybe redundant
def get_sigma2(self, abs_Phi2=None):
if abs_Phi2 is None:
abs_Phi2 = abs(self[...]) ** 2
a = self.m * self.g / 4 / np.pi / self.hbar**2
a_perp2 = self.hbar / self.m / self.w0_perp
return a_perp2 * np.sqrt(1 + 4 * np.sqrt(2) * a * abs_Phi2)
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_2D = self.get_density()
Lz = self.basis.Lxyz[-1]
lam_perp = self.get_lambdas()
if TF:
n = np.sqrt(n_2D * self.m * self.w0_perp / 2 / Lz / self.g)
return n / lam_perp
else:
sigma2 = self.get_sigma2()
return n_2D / lam_perp / np.sqrt(sigma2 * np.pi)
def get_V_GPU(self):
"""Return the complete potential `V` - internal and external."""
m = self.m
# n_2D = self.get_density()
# hbar = self.hbar
sigma2 = self.get_sigma2()
lam_perp2 = self.xp.prod(self.get_lambdas())
# lam_perp2 = self.get_lambdas()
V_ext = self.get_Vext_GPU()
# V_int_etc_ = (
# hbar**2 / 4 / m / sigma2
# + m * self.w0_perp**2 * sigma2 / 4
# + self.g * n_2D / 2 / np.sqrt(2) / np.pi / sigma2
# ) / lam_perp2
V_int_etc = (m * self.w0_perp**2 / 2 * sigma2) / lam_perp2
# assert np.allclose(V_int_etc, V_int_etc_)
return V_ext + V_int_etc
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
a = m * g / 4 / np.pi / hbar**2
y = self
psi = self.get_psi()
w0_perp2 = self.w0_perp**2
abs_Phi2 = abs(psi) ** 2
a_perp4 = hbar**2 / m**2 / w0_perp2
sigma2 = self.get_sigma2(abs_Phi2=abs_Phi2)
lam_perp2 = self.xp.prod(self.get_lambdas())
Ky = y.copy()
Ky.apply_laplacian(factor=self.K_factor)
K = psi.conj() * Ky.get_psi()
Vext = self.get_Vext_mu() * abs_Phi2
# Should we call self.get_Eint()?
Vint = m * w0_perp2 / 2.0 / a / lam_perp2 * (sigma2**3 / a_perp4 / 6 / np.sqrt(2))
# Vint_alt = (
# (
# hbar ** 2 / 4.0 / m / sigma2
# + m * w0_perp2 * sigma2 / 4.0
# + g * abs_Phi2 / 2.0 / np.sqrt(2) / np.pi / sigma2
# )
# * abs_Phi2
# / lam_perp2
# )
# assert np.allclose(Vint, Vint_alt)
self.Vint = Vint
# self.Vint_alt = Vint_alt
self._l = locals()
return K + Vint + Vext
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_perp(t=0)
ws = self.get_ws_perp(t=t)
ddlams = -lams * ws**2 + w0s**2 / np.prod(lams) / lams
return np.ravel([dlams, ddlams])
def get_lambdas(self):
if self.t != self._lambda_cache[0]:
t0, q0 = self._lambda_cache
q1 = odeint(self._rhs, q0, [t0, self.t])[-1]
self._lambda_cache = (self.t, q1)
return self._lambda_cache[1][: len(self._lambda_cache[1]) // 2]
def get_n_TF(self, V_TF, V_ext=None, g=None, **kw):
"""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.
"""
zero = np.zeros(self.shape)
if g is None:
g = self.g
if V_ext is None:
V_ext = self.get_Vext_mu(**kw)
V = V_ext + zero
h = self.hbar
m = self.m
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 * m)
n_1D = 2 * np.pi * m * np.maximum(zero, sigma2w**2 - (h / m) ** 2) / g
return n_1D
# End of methods required by the model
##########################################################################
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
return mu
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).
"""
V_TF = mu - self.hbar * self.w0_perp
return V_TF
def get_Vext_GPU(self, mu=None, fiducial=False, expt=False):
"""Return the external potential).
This method just delegates to the experiment, and provides some simple
memoization for performance.
Arguments
---------
mu : float, None
If None, then subtract `self.mu`. This will minimize the phase
rotations during time evolution, but is undesirable when computing
initial states. In the latter case, set `mu=0`.
"""
if fiducial or expt:
return self.experiment.get_Vext(state=self, fiducial=fiducial, expt=expt)
if mu is None and hasattr(self, "mu"):
mu = self.mu
if (
self._Vext
and self._Vext[1] is not None
and (self.t == self._Vext[0] or self._time_independent_Vext)
):
Vext = self._Vext[1]
else:
Vext = self.experiment.get_Vext(state=self, fiducial=False, expt=False)
self._Vext = [self.t, Vext]
if mu is not None:
Vext = Vext - mu
return Vext
[docs]
class StatePan(EvolveMixin, StateGPEdxy):
def __init__(self, experiment, **kw):
[docs]
self.experiment = experiment
StateGPEdxy.__init__(self, **kw)
[docs]
self._time_independent_Vext = False
[docs]
self._Vext = [None, None] # Cache for performance
[docs]
class Experiment(ExperimentBase):
"""Base experiment class for pancake formulation."""
######################################################################
# Attributes required by IExperiment
[docs]
t_unit = u.ms # All times specified in this unit
[docs]
t__image = 1 # Expansion time for imaging
#
######################################################################
# Attributes characterizing particles/media
[docs]
species = (1, -1) # Spin states of the species
[docs]
gs = None # Coupling constants: determined from states
[docs]
x_TF = 15 * u.micron # Alternatively x_TF can be set
[docs]
mu = None # Chemcal potential
#
######################################################################
# Attributes defining symmetry of the system
[docs]
trapping_frequencies_Hz = (2.4, 2.4, 222.0)
[docs]
basis_type = "pan" # "2D", "pan", "3D"
[docs]
Nxy = [None, None] # Default lattice and box size
[docs]
Lxy = [40 * u.micron, 40 * u.micron]
[docs]
dxy = [0.1 * u.micron, 0.1 * u.micron] # Can be used instead to fix Nx
#
######################################################################
# Attributes defining external potentials
[docs]
x_p0 = 0 # Initial x position
[docs]
y_p0 = 0 # Initial y position
[docs]
vx_p = 0 # x velocity of potential
[docs]
vy_p = 0 # y velocity of potential
[docs]
sigmax_p = 1 * u.micron # x-width of potential
[docs]
sigmay_p = 1 * u.micron # y-width of potential
#
######################################################################
# Methods required by IExperiment
[docs]
def init(self):
msgs = []
key = (self.species, self.species)
self.scattering_length = u.scattering_lengths[key]
self.m = u.masses[self.species]
if self.basis_type == "pan":
self.dim = 2
self.State = StatePan
elif self.basis_type == "2D":
self.dim = 2
self.State = State2D
elif self.basis_type == "3D":
self.dim = 3
self.State = State3D
else:
raise ValueError(
"Unknown basis_type={} (use one of '2D', 'pan', or '3D')".format(
self.basis_type
)
)
if getattr(self, "trapping_frequencies_Hz", None) is not None:
if not hasattr(self, "ws"):
# Allow subclasses to change this
self.ws = 2 * np.pi * np.array(self.trapping_frequencies_Hz) * u.Hz
self.ws = np.asarray(self.ws)
if not hasattr(self, "ws_expt"):
# These may be set externally, but if not, then define them
# here. They are used in get_fiducial_V_TF().
self.ws_expt = self.ws
if not hasattr(self, "Lxy_expt"):
self.Lxy_expt = self.Lxy
if self.dxy is not None:
# Special case to calculate Nxyz in terms of a lattice spacing
_N = np.fromiter(map(truediv, self.Lxy, self.dxy), dtype=np.float)
self.Nxy = [get_good_N(x) for x in _N]
else:
self.dxy = list(map(truediv, self.Lxy, self.Nxy))
msgs.append("Using 3D coupling constants")
if self.gs is None:
self.g = 4 * np.pi * u.hbar**2 / self.m * self.scattering_length
else:
self.g = self.gs
# Collect all time-dependent methods.
# Note: getmembers() accesses fiducial_V_TF which used to trigger an
# unwanted call to the potentially expensive get_fiducial_V_TF(). We
# preempt this here by setting it to None, then deleting it later.
# This also resets it which is consistent with init()'s semantics of
# resetting the experiment (parameters may have changed).
self._fiducial_V_TF = None
ExperimentBase.init(self)
del self._fiducial_V_TF
self.msgs = msgs
# Defining the scaling parameter for length
self.x_s = np.sqrt(1.0 / self.m / self.ws[1]) * u.micron # FIXME
self.x_s = u.micron
self.t_s = 1 / self.ws[1] # FIXME
self.A = np.sqrt(1.0 * self.ws[1] / self.g) # FIXME
[docs]
def get_Vtrap(self, state, xyz, expt=False):
"""Return the experimental trapping potential and various
approximations (background potential, not including the
time-dependent potential like the bucket).
"""
if expt:
ws = self.ws_expt
else:
t_ = state.t / self.t_unit
ws = self.get("ws", t_=t_)
V_m = 0.5 * sum((_w * _x) ** 2 for _w, _x in zip(ws, xyz))
return state.m * V_m
[docs]
def get_Vext(self, state, fiducial=False, expt=False):
"""Return the external potential.
For `t > state.t_final`, all the potentials are set to zero.
Arguments
---------
state : IState
Current state. Use this to get the time `state.t` and
abscissa `state.basis.xyz`.
fiducial : bool
If `True`, then return the potential that should be used to
define the initial state in terms of the Thomas Fermi
radius of the cloud `x_TF`.
expt : bool
If `True`, then return the proper experimental potential
rather than the potential used in the simulation.
See Also
--------
* interface.IExperiment.get_Vext
"""
zero = np.zeros(state.shape)
if state.t > state.t_final:
_Vext = zero
return _Vext
xyz = state.basis.xyz
_Vext = self.get_Vtrap(state=state, xyz=xyz, expt=expt)
if not fiducial:
_Vext += self.get_Vt(state=state)
return _Vext
[docs]
def get_Vt(self, state):
"""Optional time-dependent trapping potentials.
These potentials are not included in the fiducial state used
to determine the initial conditions, however, if `Vt` is
non-zero at time `t=0`, then this *will* be included in the
initial state preparation.
See Also
--------
* interface.IExperiment.get_Vext
"""
x_p = self.x_p0 - self.vx_p * state.t
y_p = self.y_p0 - self.vy_p * state.t
x, y = state.basis.xyz
return self.V_p * np.exp(
-((x - x_p) ** 2) / self.sigmax_p**2 / 2
- ((y - y_p) ** 2) / self.sigmay_p**2 / 2
)
# End of methods required by IExperiment
######################################################################
@property
[docs]
def fiducial_V_TF(self):
"""This may be slow to calculate, so we defer calculation until we
really need it."""
if not hasattr(self, "_fiducial_V_TF"):
self._fiducial_V_TF = self.get_fiducial_V_TF()
return self._fiducial_V_TF
[docs]
def get_fiducial_V_TF(self, t_=0.0, Nx=2**12, Lx_factor=1.1):
"""Return the V_TF required to initialize the state.
If V_TF is None or not defined, then compute the V_TF that defines the
state in terms of the Thomas-Fermi radius x_TF along the x axis using a
Harmonic trapping potential with frequencies `ws_expt` as follows:
"""
V_TF = getattr(self, "V_TF", None)
x_TF = getattr(self, "x_TF", None)
mu = getattr(self, "mu", None)
if mu is not None:
state = self.get_state(initialize=False)
return state.get_V_TF_from_mu(mu)
if V_TF is not None:
if x_TF is not None:
raise ValueError(
"Both V_TF={} and x_TF={} set. Set one to None".format(V_TF, x_TF)
)
return V_TF
return V_TF
[docs]
def get_state(self, expt=False, initialize=True):
"""Quickly return an appropriate initial state."""
state_args = dict(
experiment=self,
x_TF=self.x_TF,
mu=self.mu,
cooling_phase=self.cooling_phase,
t=0.0,
g=self.g,
m=self.m,
constraint="N",
)
if expt:
Lxy = self.Lxy_expt
Nxy = [get_good_N(a / b) for (a, b) in zip(Lxy, self.dxy)]
else:
Lxy = self.Lxy
Nxy = self.Nxy
state = self.State(Nxyz=Nxy, Lxyz=Lxy, **state_args)
if initialize:
# V_TF = self.fiducial_V_TF
if np.allclose(state[...], 0):
state[...] = 1.0
return state
[docs]
def get_initial_state(
self,
perturb=0.0,
E_tol=1e-12,
psi_tol=1e-12,
disp=1,
tries=20,
cool_steps=100,
cool_dt_t_scale=0.1,
minimize=True,
**kw,
):
"""Return an initial state with the specified population fractions.
This initial state is prepared in state[0] with the potentials as
they are at time `t=0`, then the `initial_imbalance` is transferred as
specified simulating an RF pulse by simply the appropriate
fraction in each state. Phases are kept the same as in the
state[0].
"""
state = self.get_state()
state.cooling_phase = 1.0
state.init()
if minimize:
m = MinimizeState(state, fix_N=True)
self._debug_state = m # Store in case minimize fails
if "use_scipy" not in kw:
kw["tries"] = tries
state = m.minimize(E_tol=E_tol, psi_tol=psi_tol, disp=disp, **kw)
self._debug_state = state # Store in case evolve fails
if cool_steps > 1:
# Cool a bit to remove any fluctuations.
state.cooling_phase = 1j
dt = cool_dt_t_scale * state.t_scale
state.t = -dt * cool_steps
evolver = EvolverABM(state, dt=dt)
evolver.evolve(cool_steps)
state = evolver.get_y()
del self._debug_state
psi0 = state[...]
# Rely on get_state for all other parameters like t, cooling_phase etc.
self._state = state = self.get_state()
state[...] = psi0
return state