import numpy as np
from scipy.stats.mstats import gmean
from scipy.optimize import brentq
from pytimeode.evolvers import EvolverABM
import mmfutils.optimize
from mmfutils.math.bases import CylindricalBasis
from gpe import bec, tube, axial, utils
from gpe.minimize import MinimizeState
[docs]
class Experiment(utils.ExperimentBase):
######################################################################
# Attributes required by IExperiment
[docs]
t_unit = u.ms # All times specified in this unit
[docs]
t__image = 7 # Expansion time for imaging
# End of attributes required by IExperiment
######################################################################
[docs]
species = (1, -1) # Spin states of the species
[docs]
gs = None # Coupling constants: determined from states
[docs]
trapping_frequencies_Hz = (2.4, 229.0, 222.0)
trapping_frequencies_Hz = (2.4, 225.0, 225.0)
[docs]
Nx = None # 2**14 # Default lattice and box size
[docs]
dx = 0.071 * u.micron # Can be used instead to fix Nx
[docs]
Nr = None # Size of axial basis (computed if not provided)
[docs]
R = None # Radius of axial basis (computed if not provided)
Nr = 62
R = 4.4 * u.micron
[docs]
basis_type = "1D" # '1D', '2D', '3D', `axial`
@property
[docs]
def dim(self):
if self.basis_type == "1D":
return 1
elif self.basis_type in set(["2D", "axial"]):
return 2
elif self.basis_type == "3D":
return 3
else:
raise ValueError(
"Unknown basis_type={} (use one of '1D', '2D', '3D', or 'axial')".format(
self.basis_type
)
)
######################################################################
# Methods required by IExperiment
[docs]
def init(self):
msgs = []
key = (self.species, self.species)
scattering_length = u.scattering_lengths[key]
self.m = u.masses[self.species]
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, "Lx_expt"):
self.Lx_expt = self.Lx
if self.dx is not None:
# Special case to calculate Nxyz in terms of a lattice spacing
self.Nx = utils.get_good_N(self.Lx / self.dx)
else:
self.dx = self.Lx / self.Nx
msgs.append("Using 3D coupling constants")
self.g = 4 * np.pi * u.hbar**2 / self.m * scattering_length
# 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
utils.ExperimentBase.init(self)
del self._fiducial_V_TF
self.msgs = msgs
[docs]
def get_Vext(self, state, fiducial=False, expt=False):
"""Return the external potential.
For `state.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
if expt:
_Vext = self.get_Vtrap_expt(state=state, xyz=xyz)
else:
_Vext = self.get_Vtrap(state=state, xyz=xyz)
if not fiducial:
_Vext += self.get_Vt(state=state)
return _Vext
[docs]
def get_state(self, expt=False, initialize=True):
"""Quickly return an appropriate initial state."""
state_args = dict(
experiment=self,
x_TF=None,
cooling_phase=self.cooling_phase,
t=0.0,
g=self.g,
m=self.m,
constraint="N",
)
if expt:
Lx = self.Lx_expt
Nx = utils.get_good_N(Lx / self.dx)
else:
Lx = self.Lx
Nx = self.Nx
if self.basis_type == "1D" and self.tube:
state = StateTube(Nxyz=[Nx], Lxyz=[Lx], **state_args)
elif self.basis_type == "axial":
from mmfutils.math.bases import CylindricalBasis
if self.Nr is None or self.R is None:
ws = self.ws_expt
wx, wr = ws[0], gmean(ws[1:])
a_perp = u.hbar / wr / self.m
# Maximum of radial width or 4*a_perp where a Gaussian would
# fall off by a factor of 1e-7.
R = max(4 * a_perp, self.Lx_expt / 2.0 * wx / wr)
Nr = int(np.ceil(R / self.dx))
else:
Nr, R = self.Nr, self.R
Nxr = [Nx, Nr]
Lxr = [Lx, R]
basis = CylindricalBasis(Nxr=Nxr, Lxr=Lxr, symmetric_x=False)
state = StateAxial(basis=basis, **state_args)
else:
Nxyz = getattr(self, "Nxyz", None)
Lxyz = getattr(self, "Lxyz", None)
if Nxyz is None or Lxyz is None:
ws = self.ws_expt[: self.dim]
wx, ws = ws[0], ws[1:]
if self.basis_type == "2D":
ws = [gmean(ws)]
aspect_ratios = np.divide(wx, ws)
if Lxyz is None:
Lxyz = [Lx] + [self.Lx_expt * _aspect for _aspect in aspect_ratios]
else:
Lxyz = [Lx] + list(Lxyz[1:])
if Nxyz is None:
Nxyz = [Nx] + [utils.get_good_N(_L / self.dx) for _L in Lxyz[1:]]
state = State(Nxyz=Nxyz, Lxyz=Lxyz, **state_args)
if initialize:
V_TF = self.fiducial_V_TF
state.set_initial_state(V_TF=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
# End of methods required by IExperiment
######################################################################
######################################################################
# External potentials
#
# If possible, these methods should be customized in subclasses rather than
# overloading get_Vext()
[docs]
def get_Vtrap_expt(self, state, xyz, ws=None):
"""Return the external trapping potential used in the experiment.
This is used to determine the chemical potential and
`fidcucial_V_TF` from a value of `x_TF`. It uses the
experimental trapping frequencies stored in `ws_expt`.
Arguments
---------
state : State
Current state.
xyz : [array]
Use these abscissa (not state.xyz) to compute the
potential. This is likely being applied to a much larger
state in order to compute the `fiducial_V_TF`.
ws : [float]
These are the trapping frequencies. This is just a
convenience if your get_Vtrap needs to compute a harmonic
potential.
"""
if ws is None:
ws = self.ws_expt
V_m = 0.5 * sum((_w * _x) ** 2 for _w, _x in zip(ws, xyz))
return state.m * V_m
[docs]
def get_Vtrap(self, state, xyz):
"""Return the experimental trapping potential `Vtrap(xyz)`.
This should return the trapping potential that is used in
cases, both when determining x_TF, and when cooling into the
initial state. It should not include time-dependent potential
like the bucket that are used for preparing the initial state
but NOT used when determining x_TF.
Arguments
---------
state : State
Current state.
xyz : [array]
Use these abscissa (not state.xyz) to compute the
potential. This is likely being applied to a much larger
state in order to compute the `fiducial_V_TF`.
"""
t_ = state.t / self.t_unit
ws = self.get("ws", t_=t_)
# This default version uses the same harmonic trap as the experiment.
return self.get_Vtrap_expt(state=state, xyz=xyz, ws=ws)
[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_p0 = self.x_TF + 45 * u.micron
x_p0 = self.x_TF + 10 * u.micron
x_p = x_p0 - self.v_p * state.t
x = state.basis.xyz[0]
return self.V_p * np.exp(-((x - x_p) ** 2) / self.r_p**2 / 2)
@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:
1. Construct an axially symmetric state using a CylindricalBasis with
enough points to model the harmonically trapped cloud in the TF
approximation. The assumption is made here that the extremities of
the cloud are harmonic. Inside, the potential need not be
harmonic. This construction uses `self.ws_expt`.
2. Compute `V_TF` from `self.x_TF` and `self.ws_expt` assuming harmonic
extremities.
3. Get the external potential using this state by calling
`get_Vext(expt_state, fiducial=True, expt=True)`. This allows the
function `get_Vext()` to customize what is meant by the fiducial
state (i.e. including or excluding parts of the time-varying
potential.)
4. Set the initial state and compute the total particle number in the
TF approximation using:
`expt_state.set_initial_state(V_TF=V_TF, fiducial=True, expt=True)`
5. Adjust V_TF to reproduce the total particle number in the TF
approximation using:
`expt_state.set_initial_state(V_TF=V_TF, fiducial=False, expt=True)`
This allows one to simulate a long time cooling into the ground
with a different potential, while fixing `x_TF` without that
potential.
"""
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
self._computing_fiducial_V_TF = True
# Construct fiducial axial state to get proper chemical potential.
ws = self.ws_expt
wx, wr = ws[0], gmean(ws[1:])
Lx = 2 * Lx_factor * x_TF
Nr = int(np.ceil(Nx * wx / wr))
R = Lx / 2.0 * wx / wr
Nxr = [Nx, Nr]
Lxr = [Lx, R]
expt_basis = CylindricalBasis(Nxr=Nxr, Lxr=Lxr, symmetric_x=False)
expt_state = State(
basis=expt_basis,
experiment=self,
g=self.g,
m=self.m,
cooling_phase=self.cooling_phase,
constraint="N",
t=0.0,
)
# Compute V_TF from the experimental potential at x=x_TF, y=z=0.
# Take the mean so we get (Va+Vb)/2
V_TF = self.get_Vtrap_expt(
state=expt_state, xyz=(x_TF,) + (0,) * (expt_state.dim - 1)
)
# Sanity check that V_TF is correctly defined
V = self.get_Vext(state=expt_state, fiducial=True, expt=True)
assert np.all(np.where(expt_state.xyz[0] >= x_TF, V >= V_TF, True))
expt_state.set_initial_state(V_TF=V_TF, fiducial=True, expt=True)
self._N_expt = N = expt_state.get_N()
def f(V_TF):
expt_state.set_initial_state(V_TF=V_TF, fiducial=False, expt=True)
res = expt_state.get_N() - N
return res
V0, V1 = mmfutils.optimize.bracket_monotonic(f, x0=V_TF)
V_TF = brentq(f, V0, V1)
self._expt_state = expt_state
del self._computing_fiducial_V_TF
return V_TF
[docs]
class PistonMixin:
"""State for piston experiment.
This class depends on an Experiment() instance to define the time-dependent
properties.
In addition, the following attributes are expected in experiments:
`experiment.t_unit`:
All experiment times are specified in this unit `t_ = t/t_unit`.
`state.t_final`:
For `t > state.t_final`, all the potentials are set to
zero.
`experiment.t__image`:
If imaging is performed (see `image_ts_` in the Simulation class), then
expansion occurs for this length of time (previously this was
`t__expand`).
"""
[docs]
def get_mu_from_V_TF(self, V_TF):
"""Return the corrected 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
if self.experiment.basis_type == "1D" and not self.experiment.tube:
# 1D bases represent homogeneous transverse dimensions so have no
# hbar*omega corrections.
pass
else:
w_perp = gmean(self.get_ws_perp(t=0))
mu += self.hbar * w_perp
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).
"""
return mu
######################################################################
# Required by base State classes.
[docs]
def get_Vext(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
######################################################################
# Initial State
#
# These functions differ from the underlying bec2 versions implementing the
# protocol we discuss in the introduction.
[docs]
def get_n_TF(self, V_TF, fiducial=False, expt=False):
"""Return the TF densities. Assumes gaa = gbb = gab and populates only
the lower band.
"""
V = self.get_Vext(fiducial=fiducial, expt=expt)
n = super().get_n_TF(V_TF=V_TF, V_ext=V, g=self.g)
# Here we add a full array of zeros to make sure that V_ext is full
# size. (Sometimes this might try to save memory by returning an
# object like 0.0 that does not have full shape, but here we are
# initializing the state, so we should make sure it gets expended.
zero = np.zeros(self.shape)
return zero + n
[docs]
def set_initial_state(self, V_TF=None, fiducial=False, expt=False):
"""Set the state using the Thomas Fermi (TF) approximation.
This will overload the base class versions to pass the `fiducial` and
`expt` arguments to the experiments. See `get_fiducial_V_TF()` for
details about the meaning of these parameters. This also sets more
appropriate phases for the state.
Arguments
---------
V_TF : float
Value of the external potential at the Thomas-Fermi "radius". Note:
this is the chemical potential without any corrections from the tube
is applied directly to the external potential to get the
densities. These are then transformed into components using the
spin-quasimomentum map.
fiducial : bool
If True, then use the fiducial initial state - i.e. do not
include the barrier potential.
expt : bool
If True, then initialize the state using experimental
parameters. This is needed, for example, if the desired
state is homogeneous (`wx=0`) but we want to initialize the
state with the experimental `x_TF` parameter. the
`expt=True` flag should be used when initializing a
(typically) larger state with `Lx=Lx_expt` and frequencies
`wx=wx_expt` to compute the `fiducial_V_TF` in the presence
of a barrier potential.
"""
if V_TF is None:
if hasattr(self.experiment, "_computing_fiducial_V_TF"):
# This happens from within get_fiducial_V_TF() so we cannot
# call that here. Hence, we just return zero.
self.data[...] = 0
self._N = 0
return
V_TF = self.experiment.fiducial_V_TF
n = self.get_n_TF(V_TF=V_TF, fiducial=fiducial, expt=expt)
self.data[...] = np.sqrt(n)
self._N = self.get_N()
######################################################################
# Required for using tube codes.
[docs]
def get_ws_perp(self, t=None):
"""Return the frequencies at time `t`. Required by tube2 classes."""
if t is None:
t = self.t
ws_perp = self.experiment.get("ws", t_=t / self.experiment.t_unit)[1:]
if t <= self.t_final:
return ws_perp
else:
return np.zeros_like(ws_perp)
######################################################################
# The following are for user convenience only. They should not be used
# internally.
@property
[docs]
def t_unit(self):
"""Get the time unit from the experiment."""
# Convenience method
return self.experiment.t_unit
@property
[docs]
def g(self):
"""Get g from experiment so it can manipulate them."""
return self.get("g")
@g.setter
def g(self, g):
assert np.allclose(self.g, self.experiment.g)
@property
[docs]
def ws(self):
"""Get ws from experiment so it can manipulate them."""
return self.get("ws")
@ws.setter
def ws(self, ws):
assert np.allclose(self.ws, self.experiment.ws)
[docs]
def get(self, name, t_=None):
"""Return the value of the parameter `name` at time `self.t`.
This just delegates to `self.experiment.get()`.
"""
if t_ is None:
t_ = self.t / self.t_unit
return self.experiment.get(name, t_=t_)
[docs]
def get_density_x(self):
"""Return the density along x.
For periodic boxes, this is the average density in the transverse plane
(which is the central density if the transverse plane is
translationally invariant) and the integrated 1D line density for tube
and axial codes.
"""
if self.experiment.basis_type == "axial":
n = axial.StateAxial.get_density_x(self)
elif self.experiment.basis_type in set(["2D", "3D"]):
n = self.get_density()
while len(n.shape) > 1:
n = n.mean(axis=-1)
else:
n = self.get_density()
return np.asarray(n)
######################################################################
# Plotting
#
# Here we provide various stack-able plotting components. Each plot
# function takes three arguments:
#
# data : PlotData
# Result of get_plot_data() with the data to be plotted.
# grid : MPLGrid
# If the plot is being generated from scratch, then this will be
# provided, and should be used to get the axis for plotting. New axes
# can be generated with `grid.next(size)` where size represents the
# relative size of the axis (default is 1). If more control is needed,
# one can generate a subgrid with `grid.grid(size)`.
#
# If `grid` is not provided, then it is expected that the plot function
# saved the axes in `plot_elements` during a previous call and those
# elements should be used. If `plot_elements.ax` is set, then the
# decorators will automatically make this active with
# `plt.sca(plot_elements.ax)`. If possible, previous elements should be
# updated rather than redrawn for performance, but if needed `ax.cla()`
# can be used. (Do not use `plt.clf()` as there may be other plot
# elements.)
# plot_elements : Namespace
# Each plot function should store their plot elements in this
# namespace. If the namespace passed in contains previously stored
# elements, then these should be updated instead if possible to speed
# the drawing of the plot. The latter feature is used when making
# movies and animations. In general, when updating, the ranges should
# *not* be adjusted so that movies and animations are smooth.
#
# If the plot is being generated as a subplot, then the plot_elements
# object will contain the following which should be used:
#
# fig : Figure
# Matplotlib figure instance.
[docs]
class State(PistonMixin, bec.State):
def __init__(self, experiment, **kw):
[docs]
self.experiment = experiment
bec.State.__init__(self, **kw)
[docs]
self._time_independent_Vext = False
[docs]
self._Vext = [None, None] # Cache for performance
[docs]
class StateTube(PistonMixin, tube.StateGPEdrZ):
def __init__(self, experiment, **kw):
[docs]
self.experiment = experiment
tube.StateGPEdrZ.__init__(self, **kw)
[docs]
self._time_independent_Vext = False
[docs]
self._Vext = [None, None] # Cache for performance
[docs]
class StateAxial(PistonMixin, axial.StateAxial):
def __init__(self, experiment, **kw):
[docs]
self.experiment = experiment
axial.StateAxial.__init__(self, **kw)
[docs]
self._time_independent_Vext = False
[docs]
self._Vext = [None, None] # Cache for performance
[docs]
class _State(tube.StateGPEdrZ):
def __init__(
self,
experiment,
Nxyz=(256 * 8,),
Lxyz=(250.0 * u.micron,),
x_TF=80 * u.micron,
v_p=2 * u.mm / u.s,
r_p=4 * u.micron,
V_p=10.0,
**kw,
):
[docs]
self.experiment = experiment
ws = np.array([2.4, 229.0, 222.0]) * u.Hz
[docs]
self.x_p0 = x_TF + 45 * u.micron
tube.StateGPEdrZ.__init__(self, Nxyz=Nxyz, Lxyz=Lxyz, x_TF=x_TF, ws=ws, **kw)
[docs]
def get_ws_perp(self, t):
return self.ws[1:]
[docs]
def get_Vext(self):
V_trap = tube.StateGPEdrZ.get_Vext(self)
x_p = self.x_p0 - self.v_p * self.t
x = self.basis.xyz[0]
return V_trap + self.V_p * np.exp(-((x - x_p) ** 2) / self.r_p**2 / 2)