"""Two component BEC in 1D.
Here we demonstrate the dynamics a two component BEC such as 87Rb.
Simplifications:
* g_aa = g_bb = g_ab = g
* m_a = m_b = m
* Assume that only total particle number is conserved.
"""
import warnings
import numpy as np
from scipy.integrate import odeint
import scipy as sp
from mmfutils.math import bases
from mmfutils.contexts import NoInterrupt
from pytimeode import interfaces, mixins, evolvers
from zope.interface import implementer
from . import bec
from .utils import GPUHelper, AsNumpyMixin, PerformanceWarning
from .interfaces import IStateMinimizeDFT
__all__ = [
"GPEMixin",
"V0Mixin",
"HOMixin",
"StateDFTBase",
"_StateBase",
"StateScaleBase",
"StateScaleHO",
"StateGPEBase",
"u",
]
_TINY = np.finfo(float).tiny
[docs]
class StateDFTBase(bec.StateDFTBase):
"""Underlying implementation of features needed for IStateDFT."""
[docs]
def axpy1(self, x, a):
"""axpy where a is a scalar."""
assert self.xp.isscalar(a)
return self._axpy(self.data, x, a)
[docs]
def axpy2(self, x, a):
"""Version of axpy where a has 2 components.
Uses `self._axpy` for each component.
"""
a = self.xp.ravel(a)
assert a.shape == (2,)
self._axpy(self.data[0], x[0], a[0])
self._axpy(self.data[1], x[1], a[1])
[docs]
def axpy(self, x, a=1):
"""Version of axpy where a has 2 components.
Uses `self._axpy` for each component.
"""
if self.xp.isscalar(a):
return self.axpy1(x, a)
a = self.xp.asarray(a)
if self.xp.prod(a.shape) == 2:
return self.axpy2(x, a=a)
self.data += a * x
[docs]
def _axpy_numpy(self, y, x, a=1):
"""Default numpy version of axpy."""
y += a * x
[docs]
def apply_V(self, Va, Vb, Vab, exp=False):
if exp:
# Matrix exponential over first two indices
# See: http://en.wikipedia.org/wiki/Matrix_exponential
a, b, c, d = Va, Vab, Vab.conjugate(), Vb
s = (a + d) / 2.0
q = self.xp.sqrt(0j + b * c - (a - s) * (d - s))
exp_s = self.xp.exp(s)
exp_s_cosh_q = exp_s * self.xp.cosh(q)
exp_s_sinch_q = exp_s * self.xp.sinc(q * 1j / self.xp.pi)
_tmp = exp_s_cosh_q - s * exp_s_sinch_q
A = _tmp + exp_s_sinch_q * a
B = exp_s_sinch_q * b
C = exp_s_sinch_q * c
D = _tmp + exp_s_sinch_q * d
else:
A, B, C, D = Va, Vab, Vab.conjugate(), Vb
# CuPy does not support arbitrary array production:
# https://docs-cupy.chainer.org/en/stable/reference/difference.html\
# #array-creation-from-python-objects
# V = self.xp.asarray([[A,B],[C,D]])
V = self.xp.stack([self.xp.stack([A, B]), self.xp.stack([C, D])])
# Einsum does not yet support out=self.data.
self.data[...] = self.xp.einsum(
"ab...,b...->a...", V, self.data, dtype=self.dtype
)
[docs]
def fill(self, value):
self.data.fill(value)
[docs]
def braket_GPU(self, y):
"""Return the dot product of `self.conjugate()` with `y`.
Arguments
---------
y : State
The states for which the braket will be computed.
"""
return self.xp.asarray(
[
(y1_ * self.get_metric_GPU()).ravel().conjugate().dot(y2_.ravel())
for y1_, y2_ in zip(self, y)
]
)
[docs]
def get_Vext_mu_GPU(self):
"""Return `(Va, Vb, Vab)` 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.
"""
Va, Vb, Vab = self.get_Vext_GPU()
# Use getattr rather than hasattr here to allow mu = None.
mus = getattr(self, "mus", None)
if (self.initializing or self.t < 0) and bool(mus):
mu_a, mu_b = mus
# One might consider inplace operations for performance if needed - but be
# careful! The default has Va is Vb, which will break if we mutate.
# Remember, this only gets called if we are initializing, so the minor
# performance hit is probably justified.
Va = Va - mu_a
Vb = Vb - mu_b
return Va, Vb, Vab
######################################################################
# User functions.
# These can be overloaded by the user to implement different
# potentials, functionals, etc.
[docs]
def get_Vext_GPU(self):
"""Return (Va, Vb, Vab), the external potentials."""
Va = Vb = Vab = self.xp.zeros(self.shape)
return Va, Vb, Vab
@implementer(
interfaces.IStateForABMEvolvers,
interfaces.IStateForSplitEvolvers,
interfaces.IStateWithNormalize,
IStateMinimizeDFT,
)
[docs]
class _StateBase(StateDFTBase):
"""Two component state.
This state is designed for a spin-orbit (SO) coupled BEC in a harmonic
confining potential with the SO field along the x-axis, and a possible
magnetic field gradient also along the x-axis.
Arguments
---------
basis : IBasis, None
If provided, then this is used to specify the basis. If not, then a
periodic basis is constructed from `Nxyz` and `Lxyz`.
constraint : 'N' | 'Nab'
If 'N', then constrain the total density allowing particles to convert
from one species to the other. If 'Nab' Then independently constrain
`Na` and `Nb`.
twist : (float, float)
Twisted boundary conditions for species a and b. This is the twist
angle.
x_TF : float
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.
PGPE : bool
If `True`, then the code implements the PGPE, imposing a cutoff at `kc`
"""
[docs]
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,
twist=None,
v_x=0,
kwz2=None,
# Basis independent parameters
t=0,
hbar=u.hbar,
ms=(u.m, u.m),
mus=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' or 'Nab'
constraint="N",
PGPE=False,
mu_Bs=(0 * u.mu_B, 0 * u.mu_B), # Magnetic moments
k_r=0.0, # Recoil wave-vector
rotating_phases=False, # If True, work in rotating phase basis
basis_args={},
t_final=None,
):
if basis is None:
# In principle, since we have two components, we should specify
# that the basis should only compute the laplacian over the last
# set of dimension and not over the first dimension which
# corresponds to species. This can be done with the axes=
# argument, but the default behaviour is to use the last dimensions
# so this will work.
args = dict(Nxyz=Nxyz, Lxyz=Lxyz, symmetric_lattice=symmetric_grid)
args.update(basis_args)
basis = bases.PeriodicBasis(**args)
# 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.)
[docs]
self.data = self.xp.zeros((2,) + 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.
[docs]
self.ms = np.asarray(ms)
[docs]
self.cooling_phase = cooling_phase
# In principle, different types of constraints might be applied - such
# as fixed momentum, fixed angular momentum etc. Here we currently
# only support constraint in ('N', 'Nab').
[docs]
self.constraint = constraint
if twist is None:
twist = (0, 0)
[docs]
self.twist = (np.asarray(twist) + np.pi) % (2 * np.pi) - np.pi
# Setting the following will trigger the update of self.data's phases.
self.rotating_phases = rotating_phases
# t_final is set by StateMixin (parent of AsNumpyMixin); override it here
if t_final is not None:
self.t_final = t_final
self.init()
# Once the state is initialized, we can set the initial state.
self.set_initial_state()
[docs]
_rotating_phases = False
[docs]
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.
"""
if (
self.ms[0] == self.ms[1]
and self.twist[0] == self.twist[1]
and (self.k_r == 0 or not self.rotating_phases)
):
# In this case, it is safe to treat m as a scalar. This saves some
# computations. Sometimes, this means we must have code like
# if self._use_m: V = (V, V)
m = self._m = self.ms[0]
self._use_m = True
else:
# In this case, we must use two components, so we broadcast.
m = self._m = np.asarray(self.ms)[self.bcast]
self._use_m = False
self.K_factor = -((self.hbar) ** 2) / 2.0 / m
if self.PGPE:
assert interfaces.verifyObject(bases.interfaces.IBasisCutoff, self.basis)
self.kx2, self._twist_phase_x = self._get_kx2()
# Here we precompute the "phase" factor appearing in the DFT 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._Ns = self.get_Ns_GPU()
# Always call inherited init methods.
super().init()
@property
[docs]
def x(self):
"""Flat x abscissa as a numpy array."""
return self.asnumpy(self.get_xyz()[0]).ravel()
[docs]
def set_initial_state(self, mus=None, x_TF=None):
"""Set the state using the Thomas Fermi (TF) approximation from either
`mus` or `x_TF` (pick only one or the other).
Arguments
---------
mus : (float, float)
Fixed chemical potentials.
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 mus is not None and x_TF is not None:
raise ValueError(f"Got both {mus=} and {x_TF=} (specify only one)")
if mus is None and x_TF is None:
mus = self.mus
x_TF = self.x_TF
if mus is None:
if x_TF is None:
x = self.get_xyz()[0].ravel()
x_TF = 0.2 * x[0] + 0.8 * x[-1] # Choose point 80% along
mus = self.get_mus_from_Vs_TF(Vs_TF=self.get_Vs_TF(x_TF=x_TF))
Vs_TF = self.get_Vs_TF_from_mus(mus)
ns = self.get_ns_TF(Vs_TF=Vs_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.
ns = ns + np.zeros(self.shape)
self.set_psi(np.sqrt(ns))
self._Ns = self.get_Ns_GPU()
[docs]
def get_mus_from_Vs_TF(self, Vs_TF):
"""Return the corrected chemical potential from Vs_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
---------
Vs_TF : (float, float)
External potentials at the Thomas Fermi "radius". (The external
potential is evaluated at this position and this is used to get
`mus`.)
"""
# Here we assume that self.mus is subtracted in get_Vext().
mus = np.asarray(Vs_TF)
if self.mus is not None:
mus = mus + self.mus
if not self.initializing:
warnings.warn(
f"In get_mus_from_Vs_TF() while not initializing: `{mus=}` might be incorrect."
)
return mus
[docs]
def get_Vs_TF_from_mus(self, mus):
"""Return Vs_TF from the chemical potentials mus.
Arguments
---------
mus : (float, float)
Physical chemical potentials (i.e. what you would pass to the minimizer).
"""
if not self.initializing:
warnings.warn(
f"In get_Vs_TF_from_mus() while not initializing: `{mus=}` might be incorrect."
)
Vs_TF = np.asarray(mus)
# Here we assume that self.mu is subtracted in get_V().
if self.mus is not None:
Vs_TF = Vs_TF - self.mus
return Vs_TF
[docs]
def get_Vs_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:
V_ext = self.get_Vext_mu()
Va, Vb, Vab = V_ext
Va = Va + zero
Vb = Vb + zero
while 1 < len(Va.shape):
Va = Va.min(axis=-1)
Vb = Vb.min(axis=-1)
x = self.get_xyz()[0].ravel()
i = np.argmin(abs(x - x_TF))
inds = slice(i - 1, i + 2)
order = min(len(x[inds]) - 1, 2)
Vs_TF = [
np.polyval(np.polyfit(x[inds], _V[inds], order), x_TF) for _V in [Va, Vb]
]
return Vs_TF
@property
[docs]
def rotating_phases(self):
"""Get rotating_phase."""
return self._rotating_phases
@rotating_phases.setter
def rotating_phases(self, rotating_phases):
"""Set rotating_phase and correct state phases if needed."""
if self._rotating_phases != rotating_phases:
x = self.get_xyz_GPU()[0]
phases = np.array([np.exp(-1j * self.k_r * x), np.exp(1j * self.k_r * x)])
if self._rotating_phases and not rotating_phases:
self.data *= phases.conjugate()
else:
self.data *= phases
self._rotating_phases = rotating_phases
self.kx2, self._twist_phase_x = self._get_kx2()
[docs]
def _get_kx2(self, kx=None, Lx=None):
"""Return the effective `kx**2` and `twist_phase_x` for the kinetic
energy matrix. This will be passed as the `kx2` and `twist_phase_x`
arguments to the basis.lagrangian function.
"""
x = self.basis.xyz[0]
if kx is None:
kx = self.basis.kx
if Lx is None:
Lx = self.basis.Lx
if self._use_m:
k_B = self.twist[0] / Lx
kx2 = (kx + k_B) ** 2
twist_phase_x = np.exp(1j * k_B * x)
else:
k_B = np.divide(self.twist, Lx)
k_r = self.k_r if self.rotating_phases else 0.0
kx2 = [(kx + k_r + k_B[0]) ** 2, (kx - k_r + k_B[1]) ** 2]
twist_phase_x = [self.xp.exp(1j * _kb * x) for _kb in k_B]
if self.v_x != 0:
kx2 += self.hbar * kx * self.v_x / self.K_factor
return kx2, twist_phase_x
######################################################################
# 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.
[docs]
def get_xyz_GPU(self):
return tuple(self.basis.xyz)
@property
[docs]
def xyz(self):
warnings.warn("Property xyz is Deprecated: use get_xyz()", DeprecationWarning)
return self.get_xyz()
[docs]
def get_metric_GPU(self):
return self.basis.metric
@property
[docs]
def metric(self):
return self.get_metric()
[docs]
def get_psi_GPU(self):
"""Return the physical wavefunction (applying any twist)."""
return self._twist_phase_x * self.data
[docs]
def set_psi(self, psi):
"""Set the state from a physical wavefunction (removing any twist)."""
self.set_data(psi / self.asnumpy(self._twist_phase_x))
[docs]
def get_density_GPU(self):
psi = self.get_psi_GPU()
return (psi.conjugate() * psi).real
[docs]
def get_V_GPU(self):
"""Return the complete potentials `(Va, Vb, Vab)` - 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).
Va, Vb, Vab = self.get_Vext_mu_GPU()
Va_int, Vb_int = self.get_Vint_GPU()
return (Va_int + Va, Vb_int + Vb, Vab)
# End state functions
######################################################################
@property
[docs]
def dim(self):
"""Dimension of the state."""
return len(self.shape)
@property
[docs]
def shape(self):
"""Shape of the state Nxyz (not including components)."""
return tuple(self.basis.shape)
@property
[docs]
def bcast(self):
"""Return a set of indices suitable for broadcasting masses etc."""
return (slice(None),) + (None,) * self.dim
@property
[docs]
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
[docs]
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
[docs]
def get_Ns_GPU(self):
return self.integrate_GPU(self.get_density_GPU())
[docs]
def get_N_GPU(self):
return sum(self.get_Ns_GPU())
[docs]
def integrate_GPU(self, a):
"""Integrate over both components, but do not sum."""
res = [(self.get_metric_GPU() * self.xp.asarray(_a)).sum() for _a in a]
return self.xp.asarray(res, dtype=res[0].dtype)
######################################################################
# Required by interface IStateForABMEvolvers
[docs]
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)
Va, Vb, Vab = self.get_V_GPU()
Vy = y.copy()
Vy.apply_V(Va=Va, Vb=Vb, Vab=Vab)
Hy = Ky + Vy
if subtract_mu:
if self.constraint == "N":
mu = y.braket_GPU(Hy).sum() / y.braket_GPU(y).sum()
# Check this - does not work if Vab != 0
mu = mu.real
Hy.axpy(y, -mu)
elif self.constraint == "Nab":
# We use _TINY to deal with special cases where
# one component may be zero (which would lead to mu =
# nan).
mu = mus = y.braket_GPU(Hy) / (y.braket_GPU(y) + _TINY)
# Check this - does not work if Vab != 0
mus = mus.real
Hy.axpy2(y, -mus)
elif not self.constraint:
mu = 0
else:
raise ValueError("constraint={} not recognized".format(self.constraint))
# Check this - does not work if Vab != 0
self._mu = mu
if self.PGPE:
Hy.set_psi(map(self.basis.smooth, Hy.get_psi()))
dy.copy_from(Hy)
dy.scale(self._phase)
return dy
######################################################################
# Required by interface IStateForSplitEvolvers
#
[docs]
def apply_exp_K(self, dt):
r"""Apply $e^{-i K dt}$ in place."""
# TODO: Does this work with PGPE?
dt_ihbar = dt * self._phase
self.apply_laplacian(factor=self.K_factor * dt_ihbar, exp=True)
[docs]
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
Va, Vb, Vab = [_V * dt_ihbar for _V in self.get_V_GPU()]
self.apply_V(Va=Va, Vb=Vb, Vab=Vab, exp=True)
######################################################################
# Required by interface IStateWithNormalize
#
[docs]
def normalize(self, s=None):
"""Normalize the state, return the scale factors and number `(s, N)`."""
if self.constraint == "N":
N = sum(self._Ns)
if s is None:
s = self.xp.sqrt(N / sum(self.get_Ns_GPU()))
elif self.constraint == "Nab":
N = self._Ns
if s is None:
s = self.xp.sqrt((N + _TINY) / (self.get_Ns_GPU() + _TINY))[self.bcast]
N = N[self.bcast]
else:
raise ValueError("constraint={} not recognized".format(self.constraint))
self *= s
if self.constraint == "N":
assert self.xp.allclose(N, self.get_N_GPU())
elif self.constraint == "Nab":
assert self.xp.allclose(N.ravel(), self.get_Ns_GPU().ravel())
return s, N
######################################################################
# Required by interface IStateDFT
#
[docs]
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)
######################################################################
# Required by IStateMinimizeDFT
#
[docs]
def get_energy(self, **kw):
E = sum(self.integrate([self.get_energy_density(**kw)]))
assert np.allclose(0, E.imag / (np.abs(E) + _TINY))
return E.real
[docs]
def get_Hy(self, subtract_mu=False):
dy = self.empty()
self.compute_dy_dt(dy=dy, subtract_mu=subtract_mu)
Hy = dy / self._phase
return Hy
# End of interface definitions
######################################################################
[docs]
def get_energy_density(self, a=True, b=True, ab=True):
"""Return the energy density.
Arguments
---------
a, b : bool
Include the energies associated with this species. This includes the kinetic
energy and the self-interaction, but no inter-species interactions.
ab : bool
Include only the interaction energy.
"""
# Warning: this is not correct. It may not be real until summed. The
# correct energy density requires abs(grad psi)^2
y = self
psi_a, psi_b = psi = self.get_psi()
na, nb = self.get_density()
Ky = y.copy()
Ky.apply_laplacian(factor=self.K_factor)
Ka, Kb = psi.conjugate() * Ky.get_psi()
Va, Vb, Vab = self.get_Vext_mu()
########## Correct for expansion etc.
Ea_int, Eb_int, Eab_int = self.get_Eint()
Ea = Ka + Ea_int + Va * na
Eb = Kb + Eb_int + Vb * nb
Eab = Eab_int + 2 * (psi_a.conjugate() * Vab * psi_b).real
energy_density = 0.0
if a:
energy_density += Ea
if b:
energy_density += Eb
if ab:
energy_density += Eab
return energy_density
######################################################################
# The remaining methods are not needed for evolution or ground state
# preparation, but may be helpful for analysis.
[docs]
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 = sum(y.braket(Hy)) / sum(y.braket(y))
assert np.allclose(0, mu.imag)
return mu
[docs]
def plot(self, log=False): # pragma: nocover
from matplotlib import pyplot as plt
from matplotlib.gridspec import GridSpec
from mmfutils.plot import imcontourf
na, nb = self.get_density()
if log:
na, nb = np.log10(na), np.log10(nb)
if self.dim == 1:
x = self.x / u.micron
plt.plot(x, na, "b:")
plt.plot(x, nb, "g:")
plt.plot(x, na + nb, "k-")
elif self.dim == 2:
from mmfutils import plot as mmfplt
x, y = [self.asnumpy(_x) / u.micron for _x in self.get_xyz()[:2]]
psi = self.get_psi()
n = self.get_density()
ax = plt.subplot(121)
imcontourf(x, y, n)
ax.set_aspect(1)
plt.colorbar()
ax = plt.subplot(122)
mmfplt.imcontourf(x, y, np.angle(psi), cmap="huslp", aspect=1)
plt.colorbar()
mmfplt.phase_contour(x, y, psi, aspect=1, linewidths=0.5)
ax.set_aspect(1)
elif self.dim == 3:
x, y, z = [self.asnumpy(_x) / u.micron for _x in self.get_xyz()]
nxy = n.sum(axis=2)
nxz = n.sum(axis=1)
nyz = n.sum(axis=0)
gs = GridSpec(1, 3)
ax = plt.subplot(gs[0])
imcontourf(x, y, nxy)
ax.set_aspect(1)
ax = plt.subplot(gs[1])
imcontourf(x, z, nxz)
ax.set_aspect(1)
ax = plt.subplot(gs[2])
imcontourf(y, z, nyz)
ax.set_aspect(1)
E = self.get_energy()
Na, Nb = self.get_Ns()
N = Na + Nb
plt.title(
"t={:.4f}t0, Ns={:.4f}+{:.4f}={:.4f}, E={:.4f}".format(
self.t / self.t0, Na, Nb, N, E
)
)
[docs]
def plot_k(self): # pragma: nocover
"""Plot the momentum distribution."""
from matplotlib import pyplot as plt
# from matplotlib.gridspec import GridSpec
# from mmfutils.plot import imcontourf
na_k, nb_k = np.fft.fftshift(abs(np.fft.fft(self[...]) ** 2))
if self.dim == 1:
ka = np.fft.fftshift(self.asnumpy(self.kxyz[0][0]).ravel())
kb = np.fft.fftshift(self.asnumpy(self.kxyz[1][0]).ravel())
k_r = self.k_r
plt.semilogy(ka / k_r, na_k, "b:")
plt.semilogy(kb / k_r, nb_k, "g:")
else:
raise NotImplementedError
E = self.get_energy()
Na, Nb = self.get_Ns()
N = Na + Nb
plt.title(
"t={}t0, Ns={:.4f}+{:.4f}={:.4f}, E={:.4f}".format(
self.t / self.t0, Na, Nb, N, E
)
)
[docs]
def evolve_to(
self,
t_end,
dt_t_scale=0.2,
Evolver=evolvers.EvolverABM,
evolve_steps=200,
callback=None,
):
"""Evolve state to `self.t = t_end`.
Arguments
---------
t_end : float
Evolves the state for `t_end` time units. Remember to
convert times by `self.t_unit`.
dt_t_scale : float
Scales dt.
Evolver : IEvolver
Pick the relevent evolver, either `evolvers.EvolverABM` or
`evolvers.EvolverSplit`.
evolve_steps : float
Number of evolution steps between callbacks if defined.
callback : function
Any function that takes the state as an argument, like
plotting or visualization.
"""
dt = dt_t_scale * self.t_scale
T = t_end - self.t
Nt = int(np.ceil(T / dt))
dt = T / Nt
evolver = Evolver(self, dt=dt)
with NoInterrupt() as interrupted:
while not interrupted and Nt > 1:
_steps = min(evolve_steps, Nt)
evolver.evolve(_steps)
Nt -= _steps
if callback:
callback(evolver.y)
self[...] = evolver.y[...]
self.t = evolver.y.t
[docs]
class GPEMixin:
"""Mixing providing the 2-component GPE equation of state."""
[docs]
def init(self):
assert getattr(self, "gs", None) is not None
super().init()
[docs]
def get_Vint_GPU(self, state=None):
"""Return (Va, Vb), the "internal" mean-field potential."""
if state is None:
state = self
na, nb = state.get_density_GPU()
if state.PGPE and state.basis.smoothing_cutoff != 0.5:
na, nb = map(state.basis.smooth, (na, nb))
gaa, gbb, gab = self.gs
Va = gaa * na + gab * nb
Vb = gbb * nb + gab * na
return Va, Vb
[docs]
def get_Eint(self, state=None):
"""Return the "internal" mean-field energy-densities (Ea, Eb, Eab).
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.
"""
if state is None:
state = self
na, nb = state.get_density()
# TODO Make this work with PGPE.
if state.PGPE and state.basis.smoothing_cutoff != 0.5:
na, nb = map(state.basis.smooth, (na, nb))
gaa, gbb, gab = self.gs
Ea, Eb, Eab = gaa * na**2 / 2.0, gbb * nb**2 / 2.0, gab * na * nb
return Ea, Eb, Eab
[docs]
def get_ns_TF(self, Vs_TF, V_ext=None, gs=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.
Note: Ignores `gab`.
"""
if state is None:
state = self
if not state.initializing:
warnings.warn(
f"In get_ns_TF() while not initializing: `{Vs_TF=}` might be incorrect."
)
if gs is None:
gs = self.gs
if V_ext is None:
V_ext = state.get_Vext_mu()
# 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(state.shape)
Va, Vb, Vab = V_ext
gaa, gbb, gab = gs
# Allow for a single V_TF to be passed in.
try:
V_TF_a, V_TF_b = Vs_TF
except (TypeError, ValueError):
V_TF_a = V_TF_b = Vs_TF
ns = zero + [
self._get_n_TF(V_TF=V_TF_a, V_ext=Va, g=gaa, state=state),
self._get_n_TF(V_TF=V_TF_b, V_ext=Vb, g=gbb, state=state),
]
return ns
[docs]
def _get_n_TF(self, V_TF, V_ext, g, state=None):
"""Helper the total TF density for a single component.
Arguments
---------
V_TF : float
Value of V(x_TF) where the density should vanish in the TF limit.
"""
if state is None:
state = self
# 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[1:])
return np.maximum(0, V_TF - V_ext) / g
[docs]
class StateGPEBase(GPEMixin, _StateBase):
"""Two component state with GPE equation of state."""
# The following attributes are used by this mixin: These are convenient defaults.
def __init__(self, gs=None, **kw):
if gs is not None:
self.gs = gs
super().__init__(**kw)
class StateBase(StateGPEBase):
"""State for the 2-component GPE in a homogeneous box."""
[docs]
class HOMixin(GPUHelper):
"""Helper mixin class for backwards compatibility.
This provides a default get_Vext_GPU().
"""
def __init__(self, ws=None, Omega=None, delta=None, **kw):
if ws is not None:
self.ws = ws
if Omega is not None:
self.Omega = Omega
if delta is not None:
self.delta = delta
super().__init__(**kw)
[docs]
def init(self):
super().init()
self._Vext = None
[docs]
def get_Vext_GPU(self, state=None):
"""Return (V_a, V_b, V_ab), the external potentials."""
if state is None:
state = self
if self._Vext is None:
zero = self.xp.zeros(self.shape)
m = self._m
xyz = self.get_xyz_GPU()
V_HO = 0.5 * m * sum((w * x) ** 2 for w, x in zip(self.ws, xyz))
if self.rotating_phases:
V_ab = self.Omega
else:
V_ab = self.Omega * np.exp(2j * self.k_r * xyz[0])
mu_a = mu_b = 0.0
if self.initializing and getattr(self, "mus", None):
mu_a, mu_b = self.mus
if self._use_m:
V_HO = (V_HO, V_HO)
self._Vext = (
V_HO[0] - mu_a - self.delta / 2.0,
V_HO[1] - mu_b + self.delta / 2.0,
V_ab + zero,
)
return self._Vext
class State(HOMixin, StateBase):
"""State with HO potential for backward compatibility."""
[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 init(self):
super().init()
for key in ["kx2", "k2"]:
# These need to be recomputed as we scale, so we must delete any cached values.
if hasattr(self, key):
delattr(self, key)
[docs]
def _get_scale_factor_and_phase_GPU(self):
lams, dlams, ddlams = self.get_lambdas()
m = self._m
xyz = self.get_xyz_GPU()
theta = (
m
/ 2
/ self.hbar
* sum(_X**2 * _dlam / _lam for (_X, _lam, _dlam) in zip(xyz, lams, dlams))
)
return self._twist_phase_x * np.exp(1j * theta) / np.sqrt(np.prod(lams))
[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_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.asnumpy(self._get_scale_factor_and_phase_GPU()))
[docs]
def get_xyz_GPU(self):
# Rescale to return physical coordinates.
lams, dlams, ddlams = self.get_lambdas()
return tuple(_x * _l for (_l, _x) in zip(lams, self.basis.xyz))
[docs]
def get_metric_GPU(self):
lams, dlams, ddlams = self.get_lambdas()
return self.basis.metric * np.prod(lams)
######################################################################
# Required by interface IStateDFT
#
[docs]
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.
"""
# Here we need to add the scale factors. We do this with the
# k2 argument.
if "k2" not in _kw:
lams, dlams, ddlams = self.get_lambdas()
kx = self.basis._pxyz[0] / lams[0]
Lx = self.basis.Lxyz[0] * lams[0]
kx2, twist_phase_x = self._get_kx2(kx=kx, Lx=Lx)
k2 = sum(
((_k / _l) ** 2 for _l, _k in zip(lams[1:], self.basis._pxyz[1:])), kx2
)
_kw["k2"] = k2
super().apply_laplacian(factor=factor, exp=exp, **_kw)
[docs]
def get_V_GPU(self):
"""Return the complete potential `V` - internal and external.
This version includes the correction from the scaled
coordinates.
"""
Va, Vb, Vab = self.get_Vext_mu_GPU()
Va_int, Vb_int = self.get_Vint_GPU()
# Correction to external potentials:
lams, dlams, ddlams = self.get_lambdas()
m = self._m
xyz = self.get_xyz_GPU()
V_corr = (
m
/ 2.0
* sum(_ddlam / _lam * _X**2 for _ddlam, _lam, _X in zip(ddlams, lams, xyz))
)
if self._use_m:
V_corr = (V_corr, V_corr)
return (Va_int + Va + V_corr[0], Vb_int + Vb + V_corr[0], Vab)
[docs]
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().
"""
[docs]
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
[docs]
def init(self):
lams = (1.0,) * self.dim
dlams = (0.0,) * self.dim
ddlams = (0.0,) * self.dim
self._lambda_cache = (0, lams, dlams, ddlams)
super().init()
[docs]
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.
t0, lams, dlams, ddlams = self._lambda_cache
q0 = np.ravel([lams, dlams])
q = odeint(self._rhs, q0, [t0, t], rtol=1e-8, atol=1e-8)[-1]
dq = self._rhs(q, t=t)
lams, dlams = np.reshape(q, (2, self.dim))
dlams, ddlams = np.reshape(dq, (2, self.dim))
if store_in_cache:
self._lambda_cache = (t, lams, dlams, ddlams)
return (lams, dlams, ddlams)
[docs]
def _rhs(self, q, t):
"""RHS for lambda(t) ODE."""
lams, dlams = np.reshape(q, (2, self.dim))
w0s = self.get_ws(t=0)[: self.dim]
ws = self.get_ws(t=t)[: self.dim]
ddlams = -lams * ws**2 + w0s**2 / lams / np.prod(lams)
return np.ravel([dlams, ddlams])
# End of required functions.
######################################################################