"""Code for simulating a single component BEC.
This is the reference implementation of the GPE for a single component
BEC. The current parameters are designed for easily working with the
Rb87 experiments in Peter Engels lab at WSU, but the framework is
quite general and should be easy to adapt to other experiments.
Note: for performance, the classes here use the AsNumpyMixin class
which turns "sunder" methods - those starting and ending with an
underscore - into proper public methods using either numpy or cupy arrays.
"""
# The start of every python files should contain a multi-line string like that
# above which describes the purpose of the file. This is called a docstring.
# The first line is the summary, followed by more details. You can see this
# documentation from the IPython prompt by adding a question mark after the
# module: `bec?` for example
# Now should follow all the imports. Use the following order:
# 1. Standard libraries (sorted)
# 2. Other libraries like numpy, scipy, etc.
# 3. Our libraries
# 4. Local packages for this project.
import warnings
from functools import partial
import numpy as np
# We use the scipy odeint package to numerically integrate the scaling
# equations from Castin and Dum to implement expanding traps.
from scipy.integrate import odeint
# Here I use by bases libraries for defining the abscissa. This allows us to
# generalize to more complicated bases later if we use only the methods
# defined in mmfutils.math.base.interface.IBasis
from mmfutils.math import bases
# I also use my ObjectBase as a base which provides support for the
# split between __init__() and init().
from mmfutils.containers import ObjectBase
# The basis for time evolution in our systems is my pytimeode package. This
# provides a core interface which we must code to in order to use the
# functionality of evolution code.
from pytimeode import interfaces, mixins
from zope.interface import implementer
# Mixin to help generate public methods from high-performance
# GPU-based methods with `_GPU()` extension.
from .utils import GPUHelper, _GPU, AsNumpyMixin, PerformanceWarning
from .interfaces import IStateMinimizeDFT
# The special variable __all__ is used in python to list which functions,
# classes, etc. should be used by users. It is not strictly needed, but tells
# people what portions of the module they should look at
__all__ = [
"Units",
"u",
"StateBase",
"StateScaleBase",
"GPEMixin",
"HOMixin",
"StateTwist_x",
]
[docs]
class Units:
"""Units and physical constants.
This class is simply a container for physical constants and units. We have
chosen values here that are relevant for 87Rb which is studied at WSU in
Peter Engels' lab.
"""
[docs]
kg = u / 1.660539040e-27
m = m_Rb87 = 86.909187 * u # 87Rb mass
[docs]
m_K39 = 38.96370648 * u # 39K mass
[docs]
m_Li6 = 6.0151228874 * u # 6Li mass
[docs]
a_B = 5.2917721054980885238e-5 * micron # Bohr radius
# hbar/micron^2/u = 63507.799258914903398 Hz
[docs]
Hz = hbar / micron**2 / u / 63507.799258914903398
[docs]
c = 299792458 * meter / s # Speed of light
[docs]
mW = 1e-3 * meter**2 * kg / s**3 # milliwatt
[docs]
microK = 2.06148374391870375 * hbar**2 / micron**2 / u # microK*k_b
# Scattering lengths for 87Rb hyperfine states:
# |1,-1> = |F, m_F>
# |1, 0>
# |2, 0>
# |2, -2>
# [1] https://iopscience.iop.org/article/10.1088/1367-2630/18/7/073029/pdf
# [2] Private Communication (Table from S. Kokkelmanns)
[docs]
scattering_lengths = {
((1, -1), (1, -1)): 100.40 * a_B, # [1]
((1, +1), (1, +1)): 100.40 * a_B,
((2, 0), (2, 0)): 94.57 * a_B,
((2, -2), (2, -2)): 98.98 * a_B, # [1]
((1, 0), (1, 0)): 100.86 * a_B, # [1]
((1, -1), (1, +1)): 101.32 * a_B,
((1, -1), (2, 0)): 98.13 * a_B, # immiscible
((1, -1), (2, -2)): 98.98 * a_B, # miscible
((2, -2), (1, 0)): 97.4 * a_B, # [1]
# ((1, -1), (1, 0)): 101.09 * a_B, # immiscible [1]: Peter says this is wrong.
((1, -1), (1, 0)): 100.41 * a_B, # miscible [2]
((2, 2), (2, 2)): 109 * a_B, # From Cass Sackett's email 6 Sept 2019
}
for _k in list(scattering_lengths):
# Include all symmetric entries
scattering_lengths[_k[::-1]] = scattering_lengths[_k]
masses[_k[0]] = masses[_k[1]] = m
[docs]
a = scattering_lengths[(1, -1), (1, -1)]
# Bohr magneton
[docs]
mu_B = 1.399624504 * 2 * np.pi * hbar * MHz / G
@staticmethod
[docs]
def get_magnetic_moment_mu_B(species=(1, -1)):
"""Return the magnetic moment for the specified species in
units of mu_B.
See Also
--------
https://steck.us/alkalidata/
https://link.aps.org/doi/10.1103/RevModPhys.49.31
Arguments
---------
species : (int, int)
`(F, mF)` specifying the hypdrfine states of the 87Rb
ground state (5^2 S_{1/2}).
Examples
--------
>>> u.get_magnetic_moment_mu_B((1, -1)), u.magnetic_moment[(1, -1)]/u.mu_B
(0.5018..., 0.5018...)
>>> u.get_magnetic_moment_mu_B((1, 0))
-0.0
>>> u.get_magnetic_moment_mu_B((2, -2))
-0.99967...
>>> u.get_magnetic_moment_mu_B((2, 0))
0.0
>>> u.get_magnetic_moment_mu_B((2, 2))
0.99967...
"""
# L = 0 # S...
# S = 1 / 2 # ..._{1/2}
J = 1 / 2 # L+S
I = 3 / 2 # Nuclear spin of 87Rb # noqa: E741
F, mF = species
# https://link.aps.org/doi/10.1103/RevModPhys.49.31
g_J = 2.00233113 # 2.00233113(20)
g_I = -0.0009951414 # −0.0009951414(10)
g_F = ( # (23) of https://steck.us/alkalidata/rubidium87numbers.pdf
g_J * (F * (F + 1) - I * (I + 1) + J * (J + 1))
+ g_I * (F * (F + 1) + I * (I + 1) - J * (J + 1))
) / (2 * F * (F + 1))
return g_F * mF # (24) of https://steck.us/alkalidata/rubidium87numbers.pdf
Units.magnetic_moment = {
(F, mF): Units.get_magnetic_moment_mu_B((F, mF)) * Units.mu_B
for F in [1, 2]
for mF in range(-F, F + 1)
}
# Here is a convenience instance for use in the code.
u = units = Units()
# In the future, we might be able to just use mixins.ArrayStateWithBraketMixin, but
# currently this fails with _GPU methods since it defines a default braket
# method. Once we switch to the gpu argument technique, it will probably work.
# NEED TO CHECK IF THIS IS STILL VALID with new _GPU method.
class StateDFTBase(AsNumpyMixin, mixins.ArrayStateMixin, GPUHelper, ObjectBase):
"""Underlying implementation of features needed for IStateDFT."""
asnumpy = staticmethod(np.asarray)
initializing = True
def pre_minimize_hook(self):
"""This method is called by the minimizers at the start of minimization to
ensure that the state is properly initialized.
"""
self.init() # We just defer everything to init()
def pre_evolve_hook(self):
"""This method is called by the evolvers at the start of evolution to
ensure that the state is properly initialized.
"""
self.pre_minimize_hook()
self.initializing = False
def apply_V(self, V, exp=False):
if exp:
V = self.xp.exp(V)
self.data *= V
def fill(self, value):
self.data.fill(value)
def braket_GPU(self, y):
"""Return the dot product of `self.conj()` with `y`.
Arguments
---------
y : State
The states for which the braket will be computed.
"""
return (self * self.get_metric_GPU()).ravel().conj().dot(y.ravel())
def get_Vext_mu_GPU(self):
"""Return Vext with the chemical potential subtracted if initializing.
The chemical potential should be subtracted if initializing or minimizing to get
the initial state. Minimization might be done with imaginary time evolution,
which should be done with negative times.
Thus, we check for `self.initializing`, `self.t < 0`, and make sure that
`self.mu` is valid.
"""
Vext = self.get_Vext_GPU()
# Use getattr rather than hasattr here to allow mu = None.
mu = getattr(self, "mu")
if (self.initializing or self.t < 0) and bool(mu):
# One might consider inplace operations for performance if needed - but be
# careful that the array is not shared. Remember, this only gets called if
# we are initializing, so the minor performance hit is probably justified.
Vext = Vext - mu
return Vext
######################################################################
# User functions.
# These can be overloaded by the user to implement different
# potentials, functionals, etc.
def get_Vext_GPU(self):
"""Return the external potential. This default version is zero."""
# TODO: Consider caching self.xp.zeros(self.shape) if performance is an issue.
# One will have to make sure not to mutate this in e.g. `get_Vext_mu_GPU()` if
# we do this.
Vext = self.xp.zeros(self.shape)
return Vext
def get_Vint_GPU(self):
"""Return the "internal" mean-field potential.
Consider using the `GPEMixin`.
"""
raise NotImplementedError
def get_Eint(self):
"""Return the "internal" mean-field energy-density.
Consider using the `GPEMixin`.
"""
raise NotImplementedError
# End user functions.
######################################################################
@implementer(
interfaces.IStateForABMEvolvers,
interfaces.IStateForSplitEvolvers,
interfaces.IStateWithNormalize,
IStateMinimizeDFT,
)
class _StateBase(StateDFTBase):
"""State class implementing the DFT for a single-component BEC.
To provide time-evolution functionality, the `pytimeode` package requires
that you provide a `State` class that implements the required behaviour of
at least one of the `IStateFor*Evolvers` interfaces. Most of the required
behaviour can be provided by the `mixins.ArrayStateMixin` class which we
inherit from as long as we provide an array called `data` which represents
the wavefunction.
On top of this, we need to provide the following three methods used by the
`pytimeode.evolvers`:
* `compute_dy_dt(dy)`: Required by interface IStateForABMEvolvers
This computes the rhs of the time-dependent differential equation and is
all that is needed by the `pytimeode.evolvers.EvolverABM` evolver. This
is a high-order generic evolver. It requires a rather small step size to
ensure that it behaves well, but once the step size is small enough, it
is converges to high accuracy. For example: there is often a critical
step size above which the evolution will blow up. If you reduce this by
a factor of 2, then you will typically obtain accuracy to 5 digits or so,
at least for modest evolution lengths.
The ABM evolver requires about 10 copies of the state to start up, and 8
copies of the state to be maintained for evolution. This can be a memory
issue if the states are very large.
* `apply_exp_K(dt)`, `apply_exp_V(dt, state)`: Required by interface
`IStateForSplitEvolvers`
These two methods allow you to use `pytimeode.evolvers.EvolverSplit`
based on the Trotter decomposition. This is rather specialized to systems
like the GPE where the exponential of the kinetic and potential terms can
be computed exactly. The ODE solver is not as high order as the ABM
method, but is manifestly unitary. Thus, you can often get away with
very large step sizes and the system will still behave "reasonably"
meaning that the evolution will not be accurate quantitatively, but will
be qualitatively reasonable so you can get an idea of what is happening.
Unfortunately, to achieve quantitative accuracy, you must generally go to
very small step sizes.
Another advantage of the split operator methods is that they require use
only 2 or 3 copies of the state, and so are better for large states that
might be a memory issue
* `normalize()`: Required by interface `IStateWithNormalize`
This method allows the evolvers to continually normalize the state, which
can improve numerical performance, or be of use when finding initial
states.
This base class provides an implementation for these in terms of the following
functions, which need to be provided by the subclass. Default functionality can be
provided by various *Mixin classes.
* `get_Vext_GPU()`: Returns the external potential.
* `get_Vint_GPU()`: Returns the self-consistent (mean-field) potential.
* `get_Eint()`: Returns the internal energy. `Vint` should be the derivative of
this with respect to the density. Not used by the evolvers, but used by the
minimizers.
The `*_GPU()` methods should use `self.xp` instead of `np` and `get_xyz_GPU()` to
allow the state to function well on a GPU where `self.data` may be a GPU-array.
These will be wrapped by `GPUHelper.__init_subclass__` to provide user-facing
methods without the `_GPU` extensions which fetch the data from the GPU for plotting
etc. The default implementations provided here delegate to the non-GPU versions,
raising a performance warning.
The reset of the methods in the code play a supporting role to these
methods required by the interfaces and for our convenience (for example,
plotting the state).
Arguments
---------
basis : IBasis
If provided, the basis should implement the
mmfutils.math.bases.interface.IBasis interface. If not provided, a
PeriodicBasis will be used instead.
mu : float, None
Chemical potential. Used for finding the initial state. If this
is `None`, then x_TF will be used instead. (If both `mu` and `x_TF`
are not `None`, a warning will be raise and `x_TF` will be used.
x_TF : float, None
Thomas Fermi "radius" for setting the initial state. The initial state
will be set so that the density will fall to zero at this point
`x=x_TF` (with y and z in the middle of the trap). If `None`, then we
default to 80% of the length in the x direction.
PGPE : bool
If `True`, then the code implements the PGPE, imposing a cutoff at `kc`
during the evolution. See `kc_kmax`.
kc_kmax : float, 0.5
Value of the momentum cutoff to use when implementing the PGPE. Note: the
value of 0.5 is treated differently than other values in that the internal
projection of the density in the non-linear interaction is not necessary.
The internal projection will be used for other values.
Omega : float
Angular velocity. If provided, we work in a rotating frame.
"""
t0 = u.ms # Base unit for times.
def __init__(
self,
basis=None,
# Specify either basis or the following
Nxyz=(2**5, 2**5, 2**5),
Lxyz=(30 * u.micron, 50 * u.micron, 50 * u.micron),
symmetric_grid=False,
# Basis independent parameters
t=0,
hbar=u.hbar,
m=1.0,
mu=None,
x_TF=None,
cooling_phase=1.0,
# In principle, different types of constraints might
# be applied - such as fixed momentum, fixed angular
# momentum etc. Here we currently only support
# constraint='N'
constraint="N",
PGPE=False,
kc_kmax=2 / 3,
basis_args={},
t_final=None,
Omega=None,
**kw,
):
if basis is None:
args = dict(
Nxyz=Nxyz,
Lxyz=Lxyz,
symmetric_lattice=symmetric_grid,
smoothing_cutoff=kc_kmax,
)
args.update(basis_args)
basis = bases.PeriodicBasis(**args)
self.basis = basis
# Here is the data for the state. The ArrayStateMixin class uses this
# to provide most of the required functionality required by the IState
# interface. Use np.zeros here so that we don't get overflow errors.
# (Using np.empty is slightly faster, but the call to get_N() may then
# raise an overflow error.)
self.data = self.xp.zeros(self.shape, dtype=complex)
# We defer all other calculations to the init() method so that they
# will use the current values of the various parameters. This allows
# the user to change the values of the parameters and they will take
# effect the next time init() is called.
self.t = t
self.hbar, self.m = hbar, m
self.mu, self.x_TF = mu, x_TF
self.cooling_phase = cooling_phase
self.constraint = constraint
self.Omega = Omega
self.PGPE = PGPE
# t_final is set by StateMixin (parent of AsNumpyMixin); override it here
if t_final is not None:
self.t_final = t_final
super().__init__(**kw)
# Once the state is initialized, we can set the initial state.
self.set_initial_state()
def init(self):
"""Initialize the state.
This method defines the basis positions, momenta, etc. for use later
on. We define these here rather than in the constructor `__init__()`
so that the user can change them later and the reinitialize the state.
We also call this function from the `pre_evolve_hook()` so that it is
called before any evolution takes place. For this reason, we should
not modify the state here.
"""
# Here we compute the kinetic energy factor which multiplies the
# laplacian: note the minus sign which converts the laplacian
# into k^2 = kx2
self.K_factor = -((self.hbar) ** 2) / 2.0 / self.m
if self.PGPE:
assert interfaces.verifyObject(bases.interfaces.IBasisCutoff, self.basis)
# Here we precompute the "phase" factor appearing in the GPE relating
# H(psi) with dpsi/dt. We include the value of hbar here and the
# cooling_phase. A potential optimization here is to allow the state
# to be real if the cooling phase is purely imaginary
cooling_phase = self.cooling_phase / abs(self.cooling_phase)
self._phase = 1.0 / 1j / self.hbar / cooling_phase
# We also record the current particle number so that normalize() can
# restore it if requested during evolution.
self._N = self.get_N_GPU()
# Always call inherited init methods.
super().init()
@property
def x(self):
"""Flat x abscissa as a numpy array."""
return self.get_xyz()[0].ravel()
@property
def kc_kmax(self):
return self.basis.smoothing_cutoff
def set_initial_state(self, mu=None, x_TF=None):
"""Set the state using the Thomas Fermi (TF) approximation from either
`mu` or `x_TF` (pick only one or the other).
Arguments
---------
mu : float
Fixed chemical potential.
x_TF : float
Position defining the Thomas Fermi "radius". (The external potential is
evaluated at this position and this is used to set `mu`.)
"""
if mu is not None and x_TF is not None:
raise ValueError(f"Got both {mu=} and {x_TF=} (specify only one).")
if mu is None and x_TF is None:
mu = self.mu
x_TF = self.x_TF
if self.mu is not None and self.x_TF is not None:
warnings.warn(f"Both {self.mu=} and {self.x_TF=} set. Using x_TF.")
if x_TF is not None:
mu = None # x_TF overrides mu
else:
x = self.x
x_TF = 0.2 * x[0] + 0.8 * x[-1] # Choose point 80% along
if mu is None:
mu = self.get_mu_from_V_TF(V_TF=self.get_V_TF(x_TF=x_TF))
V_TF = self.get_V_TF_from_mu(mu=mu)
n = self.get_n_TF(V_TF=V_TF)
# 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.
n = n + np.zeros(self.shape)
self.set_psi(np.sqrt(n))
self._N = self.get_N_GPU()
def get_mu_from_V_TF(self, V_TF):
"""Return the corrected chemical potential from V_TF.
In some cases, the chemical potential may differ from the value of the
external potentials at V(x_TF) due to kinetic energy shifts (in the SOC
case for example) or due to the energy of radial excitations (see the
tube codes). This function adds the appropriate correction.
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`.)
"""
# Here we assume that self.mu is subtracted in get_V().
mu = V_TF
self_mu = getattr(self, "mu", None)
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
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
# Here we assume that self.mu is subtracted in get_V()
self_mu = getattr(self, "mu", None)
if self_mu:
V_TF -= self.mu
return V_TF
def get_V_TF(self, x_TF, V_ext=None):
"""Return the Thomas Fermi chemical potential at x_TF.
Arguments
---------
x_TF : float
Position defining the Thomas Fermi "radius". (The external potential is
evaluated at this position and this is used to get `mu`.)
"""
zero = np.zeros(self.shape)
if V_ext is None:
# Don't include mu...
initializing, self.initializing = self.initializing, False
V_ext = self.get_Vext_mu()
self.initializing = initializing
V = V_ext + zero
# Minimize along all axes except x which is the 0th axis.
while len(V.shape) > 1:
V = np.min(V, axis=-1)
x = self.x.ravel()
# Find the closest lattice points an perform a polynomial fit
# so we can interpolate to the closest V even if x_TF does not
# lie on a lattice point.
i = np.argmin(abs(x - x_TF))
# Make sure slice has at least 3 points.
i = min(max(i, 1), len(x) - 2)
inds = slice(i - 1, i + 2)
order = 2
V_TF = np.polyval(np.polyfit(x[inds], V[inds], order), x_TF)
return V_TF
######################################################################
# Required menthods
#
# Thise should be defined by subclasses or mixins.
def get_n_TF(self, V_TF, V_ext=None, g=None):
"""Return the Thomas Fermi density profile n from V_TF.
Arguments
---------
V_TF : float
Value of V(x_TF) where the density should vanish in the TF limit.
"""
raise NotImplementedError
######################################################################
# State functions and attributes
#
# These functions likely need to be modified by more complicated
# state implementations. See StateTwist_x below for an example.
# In most cases, if these are properly defined, then the other
# functions below will work properly.
def get_xyz_GPU(self):
return tuple(self.basis.xyz)
@property
def xyz(self):
warnings.warn("Property xyz is Deprecated: use get_xyz()", DeprecationWarning)
return self.get_xyz()
def get_metric_GPU(self):
return self.basis.metric
@property
def metric(self):
return self.get_metric()
def get_psi_GPU(self):
"""Return the physical wavefunction. Use this so classes can overload."""
return self.data
def set_psi(self, psi):
"""Set the state from a physical wavefunction."""
self.set_data(psi)
def get_density_GPU(self):
"""Return the density of the state."""
psi = self.get_psi_GPU()
return abs(psi) ** 2
def get_V_GPU(self):
"""Return the complete potential `V` - internal and external."""
# This is the function that is actually used in the code. If
# you need to modify the potentials, then overload this
# function (see StateScaleBase below for an example).
V_ext = self.get_Vext_mu_GPU()
V_int = self.get_Vint_GPU()
return V_ext + V_int
# End state functions
######################################################################
@property
def dim(self):
"""Dimension of the state."""
return len(self.shape)
@property
def shape(self):
"""Shape of the state Nxyz."""
return tuple(self.basis.shape)
@property
def E_max(self):
"""Return the maximum kinetic energy in the basis.
This is useful when using evolvers as convergence should be obtained
when the time-step is roughly::
dt = 0.1 * state.hbar / state.E_max
See `t_scale`.
"""
# The use of np.asarray().max() allows for the possibility that
# self.basis.k_max is a float or an array (in 2D or 3D). This
# should be np.asarray, not self.xp.asarray since it needs to
# support objects.
return self.asnumpy(np.asarray(abs(self.K_factor) * self.basis.k_max**2).max())
@property
def t_scale(self):
"""Return the smallest time-scale for the problem.
Evolvers - especially the ABM evolvers - should use a `dt=0.1*t_scale` or
so. If much smaller `dt` values are required for convergence, then it
usually indicates that your lattice spacing is too larger. Likewise,
if you can get away with much larger `dt` values, then your lattice
spacing might be unnecessarily small.
Required by `Simulation`.
"""
return self.hbar / self.E_max
def get_density_x(self):
"""Return integrated density along x axis."""
n = self.get_density()
for _x in reversed(self.get_xyz()[1:]):
n = np.trapezoid(n, _x.ravel(), axis=-1)
return n
def get_N_GPU(self):
"""Return the total D-dimensional particle number in the state."""
n = self.get_density_GPU()
return self.integrate_GPU(n)
def integrate_GPU(self, a, **kw):
"""Return the integral of `a` over the box."""
return self.xp.sum(self.get_metric_GPU() * self.xp.asarray(a), **kw)
######################################################################
# Required by interface IStateForABMEvolvers
#
def compute_dy_dt(self, dy, subtract_mu=True):
"""Return `dy_dt` storing the results in `dy`.
Arguments
---------
subtract_mu : bool
If `True`, then subtract the chemical potential such that `dy_dt` is
orthogonal to the original state `y`. This will minimize the
evolution of the overall phase during real-time evolution (which can
reduce numerical errors) and will ensure that evolution under
imaginary or complex time will preserve particle number.
This should not be set if computing physical energy of the state,
however, which is why it is a parameter.
"""
y = self
Ky = y.copy()
Ky.apply_laplacian(factor=self.K_factor)
Vy = y.copy()
Vy.apply_V(V=self.get_V_GPU())
Hy = Ky + Vy
if subtract_mu:
if self.constraint == "N":
mu = y.braket_GPU(Hy) / y.braket_GPU(y)
if self.Omega:
mu = mu.real
else:
assert np.allclose(0, mu.imag)
elif not self.constraint:
mu = 0
else:
raise ValueError("constraint={} not recognized".format(self.constraint))
mu = mu.real
Hy.axpy(y, -mu)
self._mu = mu
if self.PGPE:
Hy.set_psi(self.basis.smooth(Hy.get_psi()))
dy.copy_from(Hy)
dy.scale(self._phase)
return dy
######################################################################
# Required by interface IStateForSplitEvolvers
#
linear = False
def apply_exp_K(self, dt):
r"""Apply $e^{K dt/i\hbar}$ in place."""
# TODO: Does this work with PGPE?
dt_ihbar = dt * self._phase
self.apply_laplacian(factor=self.K_factor * dt_ihbar, exp=True)
def apply_exp_V(self, dt, state):
r"""Apply $e^{-i V dt}$ in place using `state` for any
nonlinear dependence in V. (Linear problems should ignore
`state`.)"""
# TODO: Does this work with PGPE?
dt_ihbar = dt * self._phase
V = self.get_V_GPU()
_tmp = V * dt_ihbar
self.apply_V(_tmp, exp=True)
######################################################################
# Required by interface IStateWithNormalize
#
def normalize(self, s=None):
"""Normalize the state, return the scale factors and number `(s, N)`."""
if self.constraint == "N":
N = self._N
if s is None:
s = self.xp.sqrt(N / self.get_N_GPU())
else:
raise ValueError("constraint={} not recognized".format(self.constraint))
self *= s
assert self.xp.allclose(N, self.get_N_GPU())
return s, N
######################################################################
# Required by interface IStateDFT
#
def apply_laplacian(self, factor, exp=False, **_kw):
"""Apply the laplacian multiplied by `factor` to the state.
Arguments
---------
factor : array-like
The result will be multiplied by this factor.
exp : bool
If `True` then `exp(factor*laplacian)(y)` will be computed instead.
"""
# Note: twist_phase_x need not be applied since we store the
# periodic function.
extras = ["kx2", "kwz2"]
for _k in extras:
_v = getattr(self, _k, None)
if _v is not None:
_kw[_k] = _v
self.data[...] = self.basis.laplacian(self.data, factor=factor, exp=exp, **_kw)
if self.Omega:
assert not exp
self.data += self.Omega * self.hbar * self.basis.apply_Lz_hbar(self.data)
######################################################################
# Required by IStateMinimizeDFT
#
def get_energy(self):
"""Return the energy of the state. Useful for minimization."""
E = self.integrate(self.get_energy_density())
atol = abs(E) * 1e-4 if self.Omega else 1e-7
assert np.allclose(0, E.imag, atol=atol)
return E.real
def get_Hy(self, subtract_mu=False):
"""Return `H(y)` for convenience only."""
dy = self.empty()
self.compute_dy_dt(dy=dy, subtract_mu=subtract_mu)
Hy = dy / self._phase
return Hy
def get_H(self, subtract_mu=False):
"""Return the Hamiltonian. Warning, slow and big."""
shape = self.shape
state = self.copy()
H = []
for n in range(np.prod(shape)):
y = np.zeros(shape).ravel()
y[n] = 1
state.set_psi(y.reshape(shape))
H.append(state.get_Hy().ravel())
return np.array(H)
# End of interface definitions
######################################################################
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
y = self
psi = self.get_psi()
n = self.get_density()
Ky = y.copy()
Ky.apply_laplacian(factor=self.K_factor)
K = psi.conj() * Ky.get_psi()
Eint = self.get_Eint()
Eext = self.get_Vext_mu() * n
return K + Eint + Eext
######################################################################
# The remaining methods are not needed for evolution or ground state
# preparation, but may be helpful for analysis.
def get_mu(self):
"""Compute the chemical potential for convenience only."""
y = self
dy = self.empty()
self.compute_dy_dt(dy, subtract_mu=False)
Hy = dy / self._phase
mu = y.braket(Hy) / y.braket(y)
assert np.allclose(0, mu.imag)
return mu.real
def get_convergence(self, full_output=False):
"""Return `(ir, uv)` convergence factors.
These are the ratios of the density at the edge vs maximum density in position
(IR) and momentum (UV) space.
Arguments
---------
full_output : bool
If `True`, return `(ir, uv, x, nx, k, nk)`.
"""
nx = abs(self.get_psi())**2
nk = np.fft.fftshift(abs(self.basis.fftn(self.get_psi()))**2)
ir = uv = 0
nx_, nk_ = nx, nk
for _N in nx.shape:
ir = max(ir, nx_[[0, -1], ...].max())
uv = max(ir, nk_[[0, -1], ...].max())
nx_, nk_ = (np.rollaxis(n, -1) for n in [nx_, nk_])
ir /= nx.max()
uv /= nk.max()
if full_output:
r = np.sqrt(sum(x**2 for x in self.get_xyz())).ravel()
r_inds = np.argsort(r)
k = np.fft.fftshift(np.sqrt(sum(x**2 for x in self.basis._pxyz))).ravel()
k_inds = np.argsort(k)
return (ir, uv, r[r_inds], nx.ravel()[r_inds], k[k_inds], nk.ravel()[k_inds])
else:
return (ir, uv)
def plot_convergence(self, ax=None, legend=True, **kw):
if ax is None:
from matplotlib import pyplot as plt
fig, ax = plt.subplots()
ir, uv, x, nx, k, nk = self.get_convergence(full_output=True)
ax.semilogy(x/abs(x.max()), nx/nx.max(), label='IR', **kw)
ax.semilogy(k/abs(k.max()), nk/nk.max(), label='UV', **kw)
if legend:
ax.legend()
return ax
def plot(
self,
log=False,
label=None,
phase=True,
fig=None,
axs=None,
convergence=True,
plot_convergence=False,
): # pragma: nocover
"""Plot 1D, 2D, or 3D data and return axs.
Arguments
---------
log : bool
If `True`, plot the log of the densities.
label : str
For 1D data, use this label.
phase : bool
If `True`, draw the phase plot in 2D.
fig : Figure
Figure in which to draw plots. Only used if `ax` or `axs` are `None`.
axs : [Axes]
List of axes used in plot. This is what is returned and should be passed in
during an animation loop
convergence : bool
If `True`, display UV and IR convergence factors.
plot_convergence : bool
If `True`, plot converge in inset.
"""
from matplotlib import pyplot as plt
from matplotlib.gridspec import GridSpec
from mmfutils.plot import imcontourf
n = self.get_density()
def scale(n):
if log:
return np.log10(n)
return n
axc = None
if self.dim == 1:
if axs is not None:
ax = axs[0]
if len(axs) > 1:
axc = axs[1]
elif fig is not None:
if len(fig.axes) > 0:
ax = fig.axes[0]
else:
ax = fig.add_subplot()
else:
fig = plt.gcf()
ax = fig.add_subplot()
axs = [ax]
n = scale(n)
x = self.x / u.micron
ax.plot(x, n, label=label)
elif self.dim == 2:
n = scale(n)
from mmfutils import plot as mmfplt
x, y = [self.asnumpy(_x).ravel() / u.micron for _x in self.get_xyz()[:2]]
psi = self.get_psi()
n = scale(self.get_density())
if axs is not None:
ax0, ax1 = axs[:2]
cax0 = cax1 = None
if len(axs) >= 4:
cax0, cax1 = axs[2:5]
if len(axs) in {3, 5}:
axc = axs[-1]
else:
if fig is None:
fig = plt.gcf()
ax0, ax1 = fig.subplots(1, 2)
cax0 = cax1 = None
plt.sca(ax0)
mesh = ax0.pcolormesh(x, y, n.T)
ax0.set_aspect(1)
cb0 = plt.colorbar(mesh, cax=cax0, ax=ax0)
cax0 = cb0.ax
if phase:
plt.sca(ax1)
mesh = ax1.pcolormesh(x, y, np.angle(psi).T, cmap="huslp")
cb1 = plt.colorbar(mesh, cax=cax1, ax=ax1)
cax1 = cb1.ax
mmfplt.phase_contour(x, y, psi, linewidths=0.5)
ax1.set_aspect(1)
axs = [ax0, ax1, cax1, cax1]
elif self.dim == 3:
x, y, z = [self.asnumpy(_x).ravel() / u.micron for _x in self.get_xyz()]
Lx, Ly, Lz = x[-1] - x[0], y[-1] - y[0], z[-1] - z[0]
nxy = scale(n.sum(axis=2))
nxz = scale(n.sum(axis=1))
nyz = scale(n.sum(axis=0))
# nmax = max(nxy.max(), nxz.max(), nyz.max())
# nmin = max(nxy.min(), nxz.min(), nyz.min())
xs = dict(x=x, y=y, z=z)
ns = dict(xy=nxy, xz=nxz, yz=nyz)
if axs is None:
if fig is None:
fig = plt.figure(figsize=(3.1 * (2 * Lx + Ly) / Lz, 3))
gs = GridSpec(1, 3, width_ratios=[Lx, Lx, Ly])
axs = [fig.add_subplot(gs[0, i]) for i in range(3)]
elif len(axs) == 4 or len(axs) == 7:
axc = axs[-1]
fig = axs[0].figure
caxs = axs[3:]
for i, (x_, y_) in enumerate([("x", "y"), ("x", "z"), ("y", "z")]):
ax = axs[i]
cax = caxs[i] if len(caxs) > i else None
x, y = xs[x_], xs[y_]
n = ns[x_ + y_]
mesh = ax.pcolormesh(x.ravel(), y.ravel(), n.T)
ax.set(xlabel=f"${x_}$", ylabel=f"${y_}$")
cb = fig.colorbar(
mesh,
ax=ax,
cax=cax,
location="top",
label=r"$\log n$" if log else "$n$",
)
if cax is None:
cax = cb.ax
axs.append(cax)
caxs.append(cax)
ax.set_aspect(1)
E = self.get_energy()
N = self.get_N()
t = np.ravel(self.t)[0]
t_unit = getattr(self, "t_unit", u.ms)
t_unit_name = getattr(self, "t_unit_name", "ms")
title = f"t={t/t_unit:.4f}{t_unit_name}, {N=:.4f}, {E=:.4f}"
if convergence:
ir, uv = map(np.log10, self.get_convergence())
title = title + f" {ir=:.2g}, {uv=:.2g}"
plt.suptitle(title)
if plot_convergence:
if axc is None:
axc = ax.inset_axes((0,0.75, 0.25, 0.25))
axc.yaxis.tick_right()
self.plot_convergence(ax=axc, legend=False)
axs.append(axc)
return axs
[docs]
class GPEMixin(GPUHelper):
"""Mixing providing the GPE equation of state."""
# The following attributes are used by this mixin: These are convenient defaults.
[docs]
g = 4 * np.pi * u.hbar**2 * u.a / u.m
def __init__(self, g=None, **kw):
if g is not None:
self.g = g
super().__init__(**kw)
[docs]
def init(self):
assert getattr(self, "g", None) is not None
super().init()
[docs]
def get_n_TF(self, V_TF, V_ext=None, g=None, state=None):
"""Return the Thomas Fermi density profile n from V_TF.
Arguments
---------
V_TF : float
Value of V(x_TF) where the density should vanish in the TF limit.
"""
if state is None:
state = self
if not state.initializing:
warnings.warn(
f"In get_n_TF() while not initializing: `{V_TF=}` might be incorrect."
)
if g is None:
if hasattr(state, "g"):
assert np.allclose(state.g, self.g)
g = self.g
if V_ext is None:
V_ext = state.get_Vext_mu()
# In some applications, the external potential may be complex, so we
# consider only the real part here
V_ext = V_ext.real + np.zeros(state.shape)
return np.maximum(0, V_TF - V_ext) / g
[docs]
def get_Vint_GPU(self, state=None):
"""Return the "internal" mean-field potential.
This version implements the standard GPE where the
energy-density has $gn^2/2$, so we have the derivative `gn` here.
"""
if state is None:
state = self
n = state.get_density_GPU()
if state.PGPE and state.basis.smoothing_cutoff != 0.5:
n = state.basis.smooth(n)
V_int = self.g * n
return V_int
[docs]
def get_Eint(self, state=None):
"""Return the "internal" mean-field energy-density.
This version implements the standard GPE where the
energy-density has $gn^2/2$. The method get_Vint() should
return the appropriate derivative of this.
"""
# TODO Make this work with PGPE.
if state is None:
state = self
n = state.get_density()
if state.PGPE and state.basis.smoothing_cutoff != 0.5:
n = state.basis.smooth(n)
return self.g * n**2 / 2.0
[docs]
class StateBase(GPEMixin, _StateBase):
"""State for the GPE in a homogeneous box."""
class StateGPEBase(GPEMixin, _StateBase):
"""State for the GPE in a homogeneous box."""
## To Do: This class should use self.experiment.ws or self.experiment.ws_t_ to deal with
## memoization etc. We will upgrade things when we switch to this interface.
[docs]
class HOMixin(GPUHelper):
"""Helper mixin class for harmonically trapped systems.
It is common in experiments for the clouds to be trapped in an external harmonic
oscillator potential. Often, components of this trapping potential are turned off
or relaxed to allow for expansion before imaging. This class provides an
implementation of `get_Vext_GPU()` that implements the trapping potential. The user
should call this (with `super().get_Vext_GPU()`) and then add any additional
potentials.
This provides a default `get_Vext_GPU()`.
"""
# The following attributes are used by this mixin: These are convenient defaults,
# however, we should eschew defining parameters in a mixin because it creates
# confusing side-cases... not sure how to fix. Maybe properties?
[docs]
g = 4 * np.pi * u.hbar**2 * u.a / u.m
[docs]
ws = 2 * np.pi * np.array([np.sqrt(8) * 126.0, 126.0, 126.0]) * u.Hz
ws.flags["WRITEABLE"] = False
[docs]
memoize = True # If True, memoize V_ext for performance.
def __init__(self, ws=None, **kw):
if ws is not None:
self.ws = ws
super().__init__(**kw)
[docs]
def get_ws(self, t=None):
"""Return the trapping frequencies at time t."""
if t is None:
t = self.t
return self.ws
[docs]
def init(self):
super().init()
# This is used to cache the external potential for performance reasons. It
# should be reset here so that when pre_evolve_hook() is called (which calls
# init()), the cache is invalidated as the chemical potential will be dropped.
self._Vext_cache = None
[docs]
def get_Vext_GPU(self, state=None):
"""Return the external potential.
Overload this method if you want to change the external potential. If
the potential should be time dependent, use `self.t` which will be
updated by the evolvers.
If a chemical potential `self.mu` is defined, then this is subtracted
from `Vext`. This allows `gpe.minimize` to find states at a constant
chemical potential. Note: for general evolution, it is better not to
set the chemical potential as this is automatically set by
`compute_dy_dt` and will then cause `get_energy` to return the
thermodynamic potential instead.
"""
if state is None:
state = self
ws = self.get_ws(t=state.t)
if (
not self.memoize
or self._Vext_cache is None
or not state.initializing is self._Vext_cache[0]
or not np.allclose(self._Vext_cache[1], ws)
):
xyz = state.get_xyz_GPU()
V_trap = 0.5 * self.m * sum((_w * _x) ** 2 for _w, _x in zip(ws, xyz))
if self.memoize:
# Try not to let users change this array if memoizing!.
try:
V_trap.flags["WRITEABLE"] = False
except:
pass
self._Vext_cache = (state.initializing, ws, V_trap)
return self._Vext_cache[2]
[docs]
def get_mu_HO(self, N):
"""Return the chemical potential required to set the particle number in
the Thomas Fermi (TF) approximation for a 3D harmonic oscillator.
"""
assert self.g > 0
ws = self.get_ws(t=self.t)
w3 = (np.prod(ws) ** (1.0 / len(ws))) ** 3
mu = ((15.0 * self.g * N * w3 / (16 * np.pi)) ** 2 * self.m**3 / 2) ** (1.0 / 5.0)
return mu
class State(HOMixin, StateGPEBase):
"""State class for backwards compatibility."""
def __init__(self, *v, **kw):
warnings.warn(
"State class is Deprecated: use HOMixin and StateBase",
DeprecationWarning,
)
super().__init__(*v, **kw)
[docs]
class StateTwist_x(StateBase):
"""Minimal modification of the State class that implements twisted boundary
conditions along the x direction and a boost with velocity `v_x`.
Arguments
---------
twist : float
Twisted boundary conditions. This is the twist angle from left
to right over the length of the box. If there is a twist, then
the wavefunction stored in the state is the periodic version with
the twist removed. For this reason, we include the functions
`get_psi()` and `set_psi()` which allow the user to access the
physical non-periodic wavefunction.
v_x : float
Boost velocity along the x axis.
kwz2 : float
Angular velocity of about z-axis in expressed as
`kwz2 = m*omega_z/hbar`.
"""
def __init__(self, twist=None, v_x=0, kwz2=None, **kw):
if twist is None:
twist = 0
[docs]
self.twist = (twist + np.pi) % (2 * np.pi) - np.pi
super().__init__(**kw)
@property
[docs]
def Lx(self):
return self.basis.Lx
@property
[docs]
def k_B(self):
"""Bloch wave-vector. This can be redefined by subclasses if needed."""
return self.twist / self.Lx
@property
[docs]
def kx(self):
"""Momentum along x direction.
Twisted boundary conditions are implemented as a shift in
the momenta by the "Bloch" momentum
"""
return self.basis.kx + self.k_B
@property
[docs]
def twist_phase_x(self):
"""Return the current twist phase."""
return np.exp(1j * self.k_B * self.x_v)
[docs]
def get_psi_GPU(self):
"""Return the physical wavefunction (applying any twist)."""
return self.data * self.twist_phase_x
[docs]
def set_psi(self, psi):
"""Set the state from a physical wavefunction (removing any twist)."""
self.set_data(psi / self.twist_phase_x)
@property
[docs]
def x_lab(self):
"""Return the abscissa in the non-moving (lab) frame for
comparison."""
x = self.get_xyz()[0]
x_lab = x + self.v_x * self.t
Lx = self.Lx
return (x_lab - x.min()) % Lx + x.min()
@property
[docs]
def x_v(self):
"""Return the abscissa in the moving frame. This is just a
shortcut for the abscissa since the problem is formulated in
this frame."""
return self.get_xyz()[0]
@property
[docs]
def kx2(self):
"""Return the equivalent of $k_x^2$ which enters the
dispersion.
For normal dispersion, we have:
K = (hbar*k)**2/2/m = - K_factor * kx2
Hence, for a modified dispersion E(kx2) we should have:
K = E(kx2) = - K_factor * kx2
kx2 = -E(kx2) / K_factor
"""
# For a boost we have (note that kx already has the twist included):
#
# K = E(kx2) - hbar*kx*v
# kx2 = -E(kx2) / K_factor + hbar * kx * v / K_factor
#
# thus, for the usual dispersion, we thus have:
#
# kx2 = kx2 + hbar * kx * v / K_factor
return self.kx**2 + self.hbar * self.kx * self.v_x / self.K_factor
[docs]
class StateScaleBase(StateBase):
"""This state implements the scaling from Castin and Dum.
To use this class, you must provide `get_lambdas()` which defines
the scaling factors at the current time.
"""
[docs]
def get_lambdas(self, t=None):
"""Return `(lams, dlams, ddlams)`: the scale factors and derivatives.
These should be computed at time `t` which should be `self.t`
if `t` is None. There should be exactly `self.dim` scale
factors - one for each dimension. If a dimension is not to be
scaled, then the scale factor should be kept constant at 1.
"""
if t is None:
t = self.t
raise NotImplementedError
[docs]
def _get_scale_factor_and_phase_GPU(self):
lams, dlams, ddlams = self.get_lambdas()
theta = (
self.m
/ 2
/ self.hbar
* sum(
_X**2 * _dlam / _lam
for (_X, _lam, _dlam) in zip(self.get_xyz_GPU(), lams, dlams)
)
)
return np.exp(1j * theta) / np.sqrt(np.prod(lams))
[docs]
def get_psi_GPU(self):
# Includes scaling factor and phase
return self.data * self._get_scale_factor_and_phase_GPU()
[docs]
def set_psi(self, psi):
# Includes scaling factor and phase
self.set_data(psi / self._get_scale_factor_and_phase_GPU())
[docs]
def get_xyz_GPU(self):
"""Return the scaled physical coordinates (XYZ in notes)."""
xyz = super().get_xyz_GPU()
lams, dlams, ddlams = self.get_lambdas()
return tuple(_x * _l for (_l, _x) in zip(lams, xyz))
@property
[docs]
def metric(self):
lams, dlams, ddlams = self.get_lambdas()
return self.basis.metric * np.prod(lams)
[docs]
def apply_laplacian(self, factor, exp=False, **_kw):
"""Apply the laplacian multiplied by `factor` to the state."""
# Here we need to add the scale factors. We do this with the
# k2 argument.
lams, dlams, ddlams = self.get_lambdas()
k2 = sum((_k / _l) ** 2 for _l, _k in zip(lams, self.basis._pxyz))
super().apply_laplacian(factor=factor, exp=exp, k2=k2, **_kw)
[docs]
def _get_Vcorr_GPU(self):
"""Return the correction to the potential induced by
coordinate transform.
"""
# Correction to external potentials:
lams, dlams, ddlams = self.get_lambdas()
xyz = self.get_xyz_GPU()
V_corr = (
self.m
/ 2.0
* sum(_ddlam / _lam * _X**2 for _ddlam, _lam, _X in zip(ddlams, lams, xyz))
)
return V_corr
[docs]
def get_V_GPU(self):
"""Return the complete potential `V` - internal and external.
This version includes the correction from the scaled
coordinates.
"""
V_int = self.get_Vint_GPU()
V_ext = self.get_Vext_mu_GPU()
V_corr = self._get_Vcorr_GPU()
return V_int + V_ext + V_corr
class StateScaleHO(StateScaleBase):
"""This state implements the scaling from Castin and Dum for
harmonic oscillators. This provides the required get_lambdas()
function but now requires the user specify the time-dependence of
the trapping frequencies in get_ws().
"""
# Integration tolerances etc. for odeint
_odeint_params = dict(rtol=1e-12, atol=1e-12, mxstep=100000)
def get_ws(self, t=None):
"""Return the trapping frequencies at time t."""
if t is None:
t = self.t
raise NotImplementedError
######################################################################
# Lambda computation required by StateScaleBase
#
# We use odeint here to solve the ODE setup for harmonic
# potentials by Castin and Dum
def init(self):
dim = len(self.get_ws(t=0))
lams = (1.0,) * dim
dlams = (0.0,) * dim
ddlams = (0.0,) * dim
self._lambda_cache = (0, lams, dlams, ddlams)
super().init()
def get_lambdas(self, t=None):
store_in_cache = False
if t is None:
# Only store in cache if not provided. This prevents user
# calls with specific times from munging the cache.
store_in_cache = True
t = self.t
if t == self._lambda_cache[0]:
t, lams, dlams, ddlams = self._lambda_cache
else:
# Compute new values.
dim = len(self._lambda_cache[1])
t0, lams, dlams, ddlams = self._lambda_cache
q0 = np.ravel([lams, dlams])
qs, infodict = odeint(
self._rhs, q0, [t0, t], full_output=True, **self._odeint_params
)
q = qs[-1]
if not infodict["message"] == "Integration successful.":
raise ValueError(infodict["message"])
dq = self._rhs(q, t=t)
lams, dlams = np.reshape(q, (2, dim))
dlams, ddlams = np.reshape(dq, (2, dim))
if store_in_cache:
self._lambda_cache = (t, lams, dlams, ddlams)
return (lams, dlams, ddlams)
def _rhs(self, q, t):
"""RHS for lambda(t) ODE."""
# Don't use self.dim since we might evolve more scaling factors for
# use in tube.py
dim = len(self._lambda_cache[1])
lams, dlams = np.reshape(q, (2, dim))
w0s = np.asarray(self.get_ws(t=0)[:dim])
ws = np.asarray(self.get_ws(t=t)[:dim])
ddlams = -lams * ws**2 + w0s**2 / lams / np.prod(lams)
return np.ravel([dlams, ddlams])
# End of required functions.
######################################################################
class StateScaleCass(StateBase):
"""This state implements the scaling for Cass' experiment.
- explain lambds and rhos
To use this class, you must provide `get_lambdas()` which defines
the scaling factors at the current time.
"""
def get_lambdas(self, t=None):
"""Return `(lams, dlams, ddlams)`: the scale factors and derivatives.
These should be computed at time `t` which should be `self.t`
if `t` is None. There should be exactly `self.dim` scale
factors - one for each dimension. If a dimension is not to be
scaled, then the scale factor should be kept constant at 1.
"""
if t is None:
t = self.t
raise NotImplementedError
def get_rhos(self, t=None):
"""Return `(rho, drho, ddrho)`: the origin of the computational basis.
These should be computed at time `t` which should be `self.t`
if `t` is None. There should be exactly `self.dim` scale
factors - one for each dimension. If a dimension is not to be
scaled, then the scale factor should be kept constant at 1.
"""
if t is None:
t = self.t
raise NotImplementedError
def _get_scale_factor_and_phase_GPU(self):
lams, dlams, ddlams = self.get_lambdas()
rhos, drhos, ddrhos = self.get_rhos()
xyz = self.get_xyz_GPU()
theta = (
self.m
/ 2
/ self.hbar
* sum(
_X**2 * _dlam / _lam - 2 * _lam * _drho * _X
for (_X, _lam, _dlam, _drho) in zip(xyz, lams, dlams, drhos)
)
)
return np.exp(1j * theta) / np.sqrt(np.prod(lams))
def get_psi_GPU(self):
# Includes scaling factor and phase
return self.data * self._get_scale_factor_and_phase_GPU()
def set_psi(self, psi):
# Includes scaling factor and phase
self.set_data(psi / self._get_scale_factor_and_phase_GPU())
def get_xyz_GPU(self):
# Rescale to return physical coordinates.
lams, dlams, ddlams = self.get_lambdas()
rhos, drhos, ddrhos = self.get_rhos()
xyz = super().get_xyz_GPU()
return tuple(_x * _lam + _rho for (_lam, _rho, _x) in zip(lams, rhos, xyz))
def get_metric_GPU(self):
lams, dlams, ddlams = self.get_lambdas()
return self.basis.metric * np.prod(lams)
def apply_laplacian(self, factor, exp=False, **_kw):
"""Apply the laplacian multiplied by `factor` to the state."""
# Here we need to add the scale factors. We do this with the
# k2 argument.
lams, dlams, ddlams = self.get_lambdas()
if "k2" not in _kw:
_kw["k2"] = sum((_k / _l) ** 2 for _l, _k in zip(lams, self.basis._pxyz))
super().apply_laplacian(factor=factor, exp=exp, **_kw)
def _get_Vcorr_GPU(self):
"""Return the correction to the potential induced by
coordinate transform.
"""
lams, dlams, ddlams = self.get_lambdas()
rhos, drhos, ddrhos = self.get_rhos()
xyz = self.get_xyz_GPU()
V_corr = (
self.m
/ 2.0
* sum(
(_ddlam / _lam) * _X**2
+ 2 * (_ddrho - _rho * _ddlam / _lam) * _X
+ (_drho - _rho * _dlam / _lam) ** 2
for _lam, _dlam, _ddlam, _rho, _drho, _ddrho, _X in zip(
lams, dlams, ddlams, rhos, drhos, ddrhos, xyz
)
)
)
return V_corr
def get_V_GPU(self):
"""Return the complete potential `V` - internal and external.
This version includes the correction from the scaled
coordinates.
"""
V_int = self.get_Vint_GPU()
V_ext = self.get_Vext_mu_GPU()
V_corr = self._get_Vcorr_GPU()
return V_int + V_ext + V_corr
class StateScaleCassHO(StateScaleCass):
"""This state implements the scaling from Castin and Dum for
harmonic oscillators. This provides the required get_lambdas()
function but now requires the user specify the time-dependence of
the trapping frequencies in get_ws().
"""
# Integration tolerances etc. for odeint
_odeint_params = dict(rtol=1e-12, atol=1e-12, mxstep=100000)
def get_ws(self, t=None):
"""Return the trapping frequencies at time t."""
if t is None:
t = self.t
raise NotImplementedError
def get_x0s(self, t=None):
"""Return the center of the trap at time t.
Note: the code will initialize the coordinates so that the
origin of the computational basis sits at x0s(t=0).
"""
# Default version - trap stays at the origin.
if t is None:
t = self.t
return (0.0,) * self.dim
def get_Vext_GPU(self, t=None):
"""Return the external potential. This is the moving HO."""
xyz = self.get_xyz_GPU()
ws = self.get_ws(t=t)
x0s = self.get_x0s(t=t)
V_trap = (
0.5
* self.m
* sum((_w * (_X - _x0)) ** 2 for _w, _x0, _X in zip(ws, x0s, xyz))
)
# Subtract chemical potential if defined
if self.initializing and getattr(self, "mu", None):
V_trap -= self.mu
self._Vext = V_trap
return self._Vext
######################################################################
# Lambda computation required by StateScaleBase
#
# We use odeint here to solve the ODE setup for harmonic
# potentials by Castin and Dum
def init(self):
t0 = 0.0
lams = np.array((1.0,) * self.dim)
dlams = np.array((0.0,) * self.dim)
ddlams = np.array((0.0,) * self.dim)
rhos = np.asarray(self.get_x0s(t=t0))
drhos = np.array((0.0,) * self.dim)
ddrhos = np.array((0.0,) * self.dim)
self._lam_cache = (t0, lams, dlams, ddlams, rhos, drhos, ddrhos)
super().init()
def _get_lams_rhos(self, t=None):
store_in_cache = False
if t is None:
# Only store in cache if not provided. This prevents user
# calls with specific times from munging the cache.
store_in_cache = True
t = self.t
if t == self._lam_cache[0]:
t, lams, dlams, ddlams, rhos, drhos, ddrhos = self._lam_cache
else:
# Compute new values.
t0, lams, dlams, ddlams, rhos, drhos, ddrhos = self._lam_cache
q0 = np.ravel([lams, dlams, rhos, drhos])
qs, infodict = odeint(
self._rhs, q0, [t0, t], full_output=True, **self._odeint_params
)
q = qs[-1]
if not infodict["message"] == "Integration successful.":
raise ValueError(infodict["message"])
self._l = locals()
dq = self._rhs(q, t=t)
lams, dlams, rhos, drhos = np.reshape(q, (4, self.dim))
dlams, ddlams, drhos, ddrhos = np.reshape(dq, (4, self.dim))
if store_in_cache:
self._lam_cache = (t, lams, dlams, ddlams, rhos, drhos, ddrhos)
return ((lams, dlams, ddlams), (rhos, drhos, ddrhos))
def get_lambdas(self, t=None):
return self._get_lams_rhos(t=t)[0]
def get_rhos(self, t=None):
return self._get_lams_rhos(t=t)[1]
def _rhs(self, q, t):
"""RHS for lambda(t) ODE."""
lams, dlams, rhos, drhos = np.reshape(q, (4, self.dim))
w0s = np.asarray(self.get_ws(t=0))[: self.dim]
ws = np.asarray(self.get_ws(t=t))[: self.dim]
x0s = np.asarray(self.get_x0s(t=t))[: self.dim]
ddlams = -lams * ws**2 + w0s**2 / lams / np.prod(lams)
ddrhos = ws**2 * (x0s - rhos)
return np.ravel([dlams, ddlams, drhos, ddrhos])
# End of required functions.
######################################################################