r"""Experimental setup for Peter Engels' 2-component Rb 87 counterflow
experiments with spin-orbit coupling. The results are summarized in the
SOC Soliton.ipynb notebook.
The design of this module is as follows:
State1, State2: These are two state objects for single component and two
component SOC states respectively. The single component state uses the
Dispersion class to provide a modified dispersion. These maintain a
reference to the relevant Experiment object (see below) and overload the
get_Vext() function of the base class to use the appropriate experimental
protocol.
Experiment*: Each experimental setup is characterized by a subclass of
Experiment which specifies the experimental parameters and protocol in terms
of physical dimensions. These experimental objects are responsible for
returning an appropriate State* object to implement the dynamics.
Experiment objects can define a set of time-dependent parameters by defining
methods with the name `*_t_()` that take an argument `t_` and returns the
value at a specified time. To facilitate plotting, another method `*_info()`
can be defined which returns `(unit, label)` with the associated numerical
unit, and a text label. These are used for plotting. For example, for a
varying Omega, one might define::
def Omega_t_(self, t_):
return np.sin(t_)
@property
def Omega_info(self):
return (self.E_R, r'$\Omega/E_R$')
The underlying Experiment class designed here allows for the following
manipulations.
* Time dependent parameters:
* Omega: SOC coupling strength
* detuning: Detuning
* B_gradient: counterflow
* xyz0: Location of trap center
* ws: Trap frequencies
* Functions (which may contain time dependence):
* get_Vext(state, fiducial, expt, single_band):
External trapping potential. The default version of this combines
the following pieces:
* Vtrap: (Always)
* Vt: (If not fiducial)
* Magnetic field gradient: (If not single band)
* SOC: (If not single band)
* get_Vtrap_expt(state, xyz):
Trapping potential used in the experiment. This is used to
compute the fiducial_V_TF.
Unless you need to modify this structure you should customize the
following pieces instead.
* get_Vtrap(state, xyz):
Trapping potential for your simulation.
* get_Vt(state):
Time dependent trap. This is only used for initial state
preparation but not computing the fiducial_V_TF. The need for
this is the case where the initial state is specified with an x_TF
without this potential, but then the system is prepared by cooling
with this potential. Unless you are matching an experiment where
they report x_TF this way, you probably don't need to modify
this.
* barrier_xyz0: Position of a barrier potential
* barrier_depth: Depth of a barrier potential
Initial States
--------------
The initial state is defined in terms of the chemical potentials which define
the initial wavefunction through the Thomas Fermi approximation. From the
experimental standpoint, the initial state is typically expressed in terms of
the Thomas Fermi "radius" `x_TF` at which point in space the "chemical
potential" should match the external potential `V_TF = V_ext(x_TF)`. In the
presence of axial degrees of freedom and the various tube approximations, the
chemical potential required to reproduce this extent of the cloud can differ
from `V_TF` (for example, due to a shift $\hbar\omega_\perp$) so we provide the
following functions to convert between the two::
mu = State.get_mu_from_V_TF(V_TF)
V_TF = State.get_V_TF_from_mu(mu)
To specify the initial state then one should define `V_TF`. This is
the fundamental quantity determining the initial state.
The alternative - used by default for the experiments - is to specify `x_TF`
and then to compute `V_TF` from the potential. This is typically determined as
the extent of a cloud without any spin-orbit coupling, barrier potential, etc.
The actual initial state is usually obtained by then evolving this state
adiabatically into the ground state with SOC, possibly including the barrier
potential. This adiabatic preparation should be thought of as roughly
performed at constant total particle number (though there may be some particle
losses).
Thus, the external chemical potential that should be used to compute `V_TF`
from `x_TF` may differ from the potentials defined for evolution. Finally, the
state may actually be too small to represent the entire and may not even
include `x_TF` in the range of abscissa. This leads us to define a
`fiducial_V_TF` which is the value of `V_TF` that should be used to initialize
the state. This may differ from the `V_TF = V_ext(x_TF)` because of these
factors. The function `get_fiducial_V_TF()` is responsible for computing the
appropriate `V_TF` to match the experiment and may need to construct another
larger state in order to solve for the value of `V_TF` that will appropriately
match the experiment.
For these reasons, `Experiment.get_Vext(state, fiducial, expt)` and
`State.set_initial_state(V_TF, fiducial, expt)` take two flags `fiducial` and
`expt` in order to accommodate this functionality (see their docstrings for
details).
"""
from argparse import Namespace
from collections import namedtuple
import logging
import time
import warnings
import attr
import numpy as np
from scipy.stats.mstats import gmean
from scipy.integrate import odeint
from scipy.optimize import brentq
from mmfutils.contexts import coroutine, NoInterrupt
from mmfutils.math.bases import CylindricalBasis
import mmfutils.optimize
from pytimeode.evolvers import EvolverABM
from gpe import bec, bec2, axial
from gpe import tube, tube2
from gpe import utils
from gpe.bec2 import u
from gpe.minimize import MinimizeState
from gpe.utils import GPUHelper
from gpe.interfaces import implementer, IExperiment
[docs]
_EPS = np.finfo(float).eps
@attr.s
[docs]
class Dispersion:
r"""Tools for computing properties of the lower band dispersion.
Everything is expressed in dimensionless units with $k/k_r$,
$d=\delta/4E_R$, $w=\Omega/4E_R$, and $E_\pm/2E_R$.
Examples
--------
>>> E = Dispersion(w=1.5/4, d=0.543/4)
>>> ks = np.linspace(-3, 3, 500)
>>> ks_ = (ks[1:] + ks[:-1])/2.0
>>> ks__ = ks[1:-1]
>>> Es = E(ks).T
>>> dEs = E(ks, d=1).T
>>> ddEs = E(ks, d=2).T
>>> d3Es = E(ks, d=3).T
>>> d4Es = E(ks, d=4).T
>>> dEs_ = np.diff(Es, axis=0)/np.diff(ks)[:, None]
>>> ddEs_ = np.diff(dEs, axis=0)/np.diff(ks)[:, None]
>>> d3Es_ = np.diff(ddEs, axis=0)/np.diff(ks)[:, None]
>>> d4Es_ = np.diff(d3Es, axis=0)/np.diff(ks)[:, None]
>>> np.allclose((dEs[1:] + dEs[:-1])/2, dEs_, rtol=0.01)
True
>>> np.allclose((ddEs[1:] + ddEs[:-1])/2, ddEs_, rtol=0.02)
True
>>> np.allclose((d3Es[1:] + d3Es[:-1])/2, d3Es_, rtol=0.01)
True
>>> np.allclose((d4Es[1:] + d4Es[:-1])/2, d4Es_, rtol=0.1)
True
Here are some regression tests
>>> E = Dispersion(w=1.0, d=0)
>>> float(E.get_k0())
0.0
"""
[docs]
E0 = attr.ib(default=0)
[docs]
k0 = attr.ib(default=0)
[docs]
def Es(self, k, d=0):
k = k + self.k0
if self.w == 0:
# Short circuit for no SOC. Just return a single band
if d == 0:
res = k**2 / 2.0
elif d == 1:
res = k
elif d == 2:
res = 1.0
else:
res = 0.0
return np.asarray((res, res))
D = np.sqrt((k - self.d) ** 2 + self.w**2)
if d == 0:
res = (k**2 + 1) / 2.0 - D - self.E0, (k**2 + 1) / 2.0 + D - self.E0
elif d == 1:
res = k - (k - self.d) / D, k + (k - self.d) / D
elif d == 2:
res = 1.0 - self.w**2 / D**3, 1.0 + self.w**2 / D**3
elif d == 3:
res = 3 * (k - self.d) * self.w**2 / D**5
res = (res, -res)
elif d == 4:
dres = (-4 * (k - self.d) ** 2 + self.w**2) * 3 * self.w**2 / D**7
res = (dres, -dres)
else:
raise NotImplementedError("Only d<5 supported. (got d={})".format(d))
return np.asarray(res)
[docs]
def newton(self, k):
return k - np.ma.divide(self.Es(k, d=1), self.Es(k, d=2)).filled(0)[0]
[docs]
def get_k0(self, N=5):
"""Return the minimum of the lower dispersion branch."""
if self.w == 0:
# Short circuit
return 0.0
# Note that sign(k_0) = -sign(d) for the lower minimum
k0 = -np.sign(self.d) / 2.0
for n in range(N):
k0 = self.newton(k0)
return k0
@attr.s
[docs]
class Dispersion3:
r"""Tools for computing properties of the lower band dispersion for
a three-component system.
Everything is expressed in dimensionless units with $k/k_r$,
$d=\delta/4E_R$, $w=\Omega/4E_R$, $e=\epsilon/2E_R$ and $E_\pm/2E_R$.
"""
[docs]
E0 = attr.ib(default=0)
[docs]
k0 = attr.ib(default=0)
[docs]
def Es(self, k, d=0):
k = k + self.k0
w, d_, e = self.w, self.d, self.e
z = np.zeros_like(k)
H = np.array(
[
[(k + 1) ** 2 / 2 - d_, w + z, z],
[w + z, (k - 1) ** 2 / 2 + d_, w + z],
[z, w + z, (k - 3) ** 2 / 2 + 3 * d_ + e],
]
).T
if d == 0:
Es = np.linalg.eigvalsh(H)
res = (Es - self.E0).T
else:
dH = np.array([[(k + 1), z, z], [z, (k - 1), z], [z, z, (k - 3)]]).T
Es, Psis = np.linalg.eigh(H)
if d == 1:
res = np.array(
[
np.einsum(
"...b, ...bc, ...c->...",
Psis[..., _n].conj(),
dH,
Psis[..., _n],
)
for _n in range(3)
]
)
elif d == 2:
# Hadamard second variation formula
vAv2 = (
abs(np.einsum("...bj, ...bc, ...ck->...jk", Psis.conj(), dH, Psis))
** 2
)
factor = np.ma.divide(2, Es[..., :, None] - Es[..., None, :]).filled(0)
vAv2 *= factor
res = 1 + vAv2.sum(axis=-1).T
else:
raise NotImplementedError("Only d=0, 1 supported. (got d={})".format(d))
return res
[docs]
def newton(self, k):
return k - np.ma.divide(self.Es(k, d=1), self.Es(k, d=2)).filled(0)[0]
[docs]
def get_k0(self, N=5):
"""Return the minimum of the lower dispersion branch."""
# Note that sign(k_0) = -sign(d) for the lower minimum
k0 = -np.sign(self.d) / 2.0
for n in range(N):
k0 = self.newton(k0)
return k0
[docs]
class SOCMixin(GPUHelper):
"""State for spin-orbit coupled experiments with a counterflow induced by a
magnetic field gradient.
This class depends on an Experiment() instance to define the time-dependent
properties such as ramping on of the magnetic field gradient, Raman
lasers, and ramping off everything for expansion imaging.
To do this, the experiment can define the following methods, which will be
used if they exist. Note: for performance reasons, these should return
`NotImplemented` if they are not used.
`experiment.delta_t_(t_)`: `experiment.delta`
`experiment.Omega_t_(t_)`: `experiment.Omega`
`experiment.B_gradient_t_(t_)`: `experiment.B_gradient`
In addition, the following attributes are expected in experiments:
`experiment.t_unit`:
All experiment times are specified in this unit `t_ = t/t_unit`.
`experiment.t__final`:
For `t_ > self.experiment.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`.)
"""
E_R = self.get("E_R")
E = self.get_dispersion()
E0 = E(E.get_k0())[0] * 2 * E_R
mu = V_TF + E0
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).
"""
E_R = self.get("E_R")
E = self.get_dispersion()
E0 = E(E.get_k0())[0] * 2 * E_R
V_TF = mu - E0
return V_TF
######################################################################
# Required by base State classes.
[docs]
def get_Vext_GPU(self, mu=None, fiducial=False, expt=False, single_band=None):
"""Return the external potentials `(V_a, V_b, V_ab)` (or `V_ext` in the
single-band case).
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 single_band is None:
single_band = self.single_band
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)
):
Va, Vb, Vab = self._Vext[1]
else:
_Vext = self.experiment.get_Vext(state=self, fiducial=False, expt=False)
self._Vext = [self.t, _Vext]
Va, Vb, Vab = _Vext
if mu is not None:
if np.isscalar(mu):
mu_a = mu_b = mu
else:
mu_a, mu_b = mu
Va, Vb = Va - mu_a, Vb - mu_b
if single_band:
return (Va + Vb) / 2.0
else:
return (Va, Vb, Vab)
######################################################################
# Initial State
#
# These functions differ from the underlying bec2 versions implementing the
# protocol we discuss in the introduction.
[docs]
def get_ns_TF(self, Vs_TF, fiducial=False, expt=False):
"""Return the TF densities. Assumes gaa = gbb = gab and populates only
the lower band.
"""
g = np.mean(self.gs)
Va, Vb, Vab = self.get_Vext(fiducial=fiducial, expt=expt, single_band=False)
if not np.isscalar(Vs_TF):
ga, gb, gab = self.gs
na = self._get_n_TF(V_TF=Vs_TF[0], V_ext=Va, g=ga)
nb = self._get_n_TF(V_TF=Vs_TF[1], V_ext=Vb, g=gb)
ns = na, nb
else:
V = (Va + Vb) / 2.0
# E = self.experiment.get_dispersion()
psi_a, psi_b = self.get_branch_psi_ab(
n0=self._get_n_TF(V_TF=Vs_TF, V_ext=V, g=g)
)
ns = na, nb = abs(psi_a) ** 2, abs(psi_b) ** 2
# 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 + ns
[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
or SOC and 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._Ns = (0, 0)
return
V_TF = self.experiment.fiducial_V_TF
ns = self.get_ns_TF(Vs_TF=V_TF, fiducial=fiducial, expt=expt)
k_r = self.get("k_r")
E = self.get_dispersion()
if self.single_band:
ns = ns.sum(axis=0)
ks = E.get_k0() * k_r
else:
k0 = E.get_k0() * k_r
ka = k0 + k_r
kb = k0 - k_r
ks = np.array([ka, kb])[self.bcast]
x = self.get_xyz_GPU()[0]
self.data[...] = np.exp(1j * ks * x) * np.sqrt(ns)
self._Ns = self.get_Ns_GPU()
[docs]
def get_k0(self, t_=0, branch=-1):
"""Return the momentum on the lattice that has minimum dispersion.
Parameters
----------
branch : -1 or 1
1 for upper branch, -1 for lower branch (default).
"""
E = self.experiment.get_dispersion(t_=t_)
k_r = self.get("k_r")
i_branch = 0 if branch == -1 else 1
return self.basis.kx[np.argmin(E(self.basis.kx / k_r)[i_branch])] / k_r
######################################################################
# 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 conversions help support single-band states by providing
# some attributes and methods of bec2.State for bec.State.
@property
[docs]
def ms(self):
if self.single_band:
return self.m * np.ones(2)
else:
return self._ms
@ms.setter
def ms(self, ms):
if self.single_band:
self.m = np.mean(ms)
else:
self._ms = ms
@property
[docs]
def bcast(self):
"""Return a set of indices suitable for broadcasting masses etc."""
return (slice(None),) + (None,) * self.dim
[docs]
rotating_phases = False
######################################################################
# 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 t_final(self):
# Convenience access to expansion time.
return self.experiment.t__final * self.experiment.t_unit
@property
[docs]
def xyz_trap(self):
"""Return the abscissa for use in the trapping potentials."""
return self.experiment._xyz_trap(state=self)
@property
[docs]
def k_r(self):
return self.get("k_r")
@k_r.setter
def k_r(self, k_r):
assert np.allclose(self.k_r, self.experiment.k_r)
@property
[docs]
def gs(self):
"""Get gs from experiment so it can manipulate them."""
return self.get("gs")
@gs.setter
def gs(self, gs):
assert np.allclose(self.gs, self.experiment.gs)
@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 not self.single_band:
if self.experiment.basis_type == "axial":
na, nb = axial.State2Axial.get_density_x(self)
elif self.experiment.basis_type in set(["2D", "3D"]):
na, nb = self.get_density()
while len(na.shape) > 1:
na = na.mean(axis=-1)
nb = nb.mean(axis=-1)
else:
na, nb = self.get_density()
else:
# Rely on the spin-quasimomentum map
if self.experiment.basis_type == "axial":
n_x = axial.StateAxial.get_density_x(self)
n = self.get_density()
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()
y = self[...]
k_r = self.get("k_r")
disp = self.experiment.get_dispersion()
kxs = self.basis.kx
ks = np.ma.divide(
(y.conj() * np.fft.ifft(kxs * np.fft.fft(y, axis=0), axis=0)).real, n
).filled(0)
k, d, w = ks / k_r, disp.d, disp.w
K = (k - d) / np.sqrt((k - d) ** 2 + w**2)
if self.experiment.basis_type == "axial":
na = axial.StateAxial.get_density_x(self, n * (1 - K) / 2)
nb = axial.StateAxial.get_density_x(self, n * (1 + K) / 2)
else:
na, nb = n * (1 - K) / 2, n * (1 + K) / 2
return np.asarray([na, nb])
[docs]
def get_dispersion(self, t_=None):
if t_ is None:
t_ = self.t / self.t_unit
return self.experiment.get_dispersion(t_=t_)
[docs]
def get_homogeneous_ground_state(self):
"""Return the properties of the homogeneous ground state."""
B_gradient = self.get("B_gradient")
if not np.allclose(B_gradient, 0):
mu_Bs = np.array(self.get("mu_Bs"))
mu_Bs * B_gradient
# Compute shifts to mu and detuning
raise NotImplementedError("B_gradient not yet supported for ground state")
k_r = self.get("k_r")
E_R = self.get("E_R")
E = self.get_dispersion()
k0_ = E.get_k0()
k0 = k0_ * k_r
E0 = E(k0)[0] * 2 * E_R
K = (k0_ - E.d) / np.sqrt((k0_ - E.d) ** 2 + E.w**2)
# na = cos(theta)*n, nb = sin(theta)*n
theta = np.arctan2(1 + K, 1 - K)
HomogeneousState = namedtuple("HomogeneousState", ["k0", "E0", "theta"])
return HomogeneousState(k0=k0, E0=E0, theta=theta)
[docs]
def get_branch_mixture(self):
"""Return the population [n_+, n_-] of the two energy branches assuming
the TF approximation is valid."""
if self.single_band:
zero = np.zeros(self.shape)
return zero, zero + 1.0
k_r = self.get("k_r")
# Omega = self.get('Omega')
# detuning = self.get('detuning')
B_gradient = self.get("B_gradient")
ga, gb, gab = self.get("gs")
ns = na, nb = self.get_density_x()
y = self[...]
# Will fail if there is a twist
assert np.allclose(self.twist, 0)
xs = self.get_xyz()[0].ravel()
kxs = self.basis.kx
ks = ka, kb = (
y.conj() * np.fft.ifft(kxs * np.fft.fft(y, axis=1), axis=1)
).real / ns
ma, mb = self.ms
Va, Vb, Vab = self.get_Vext(single_band=False)
a_ = self.hbar**2 * ka**2 / 2 / ma + Va + ga * na + gab * nb
b_ = self.hbar**2 * kb**2 / 2 / mb + Vb + gb * nb + gab * na
# A = (a_ + b_)/2.0
B = (a_ - b_) / 2.0
C = abs(Vab)
D = np.sqrt(B**2 + C**2)
up = (B + D) / np.sqrt(2 * B * (B + D) + 2 * C**2)
um = (B - D) / np.sqrt(2 * B * (B - D) + 2 * C**2)
vp = C / np.sqrt(2 * B * (B + D) + 2 * C**2)
vm = C / np.sqrt(2 * B * (B - D) + 2 * C**2)
psi_a, psi_b = np.exp(np.array([-1j, 1j])[self.bcast] * k_r * xs) * self[...]
psi_p = up.conj() * psi_a + vp.conj() * psi_b
psi_m = um.conj() * psi_a + vm.conj() * psi_b
self._l = locals()
n_p, n_m = abs(psi_p) ** 2, abs(psi_m) ** 2
return n_p, n_m
[docs]
def get_branch_psi_ab(self, n0, k=None, branch=-1, tol=10 * _EPS):
"""Return the wavefunction factors `(psi_a, psi_b)` for the
specified band. In some sense, this is the inverse of
`get_branch_mixture()` but for purely homogeneous states.
These include the appropriate phase factors of
`exp(1j*(k+k_r)*x)` and `exp(1j*(k-k_r)*x)`.
Parameters
----------
k : float
Momentum of homogeneous state (in the single band model)
If None, then get the minimium of the dispersion.
n0 : float
Background density (3D)
experiment : Experiment
Defines `Omega`, `detuning`, `gs`, `E_R` and `k_r`.
branch : -1 or 1
1 for upper branch, -1 for lower branch (default).
"""
Omega = self.get("Omega")
detuning = self.get("detuning")
gaa, gbb, gab = self.get("gs")
k_r = self.get("k_r")
E_R = self.get("E_R")
x = self.get_xyz()[0]
# First assume that gs are equal
d0_ = detuning / 4 / E_R
C_ = w_ = Omega / 4 / E_R
if k is None:
k = k_r * self.get_k0()
k_ = k / k_r
gn1_ = (gaa - gab) * n0 / 4 / E_R
gn2_ = (gbb - gab) * n0 / 4 / E_R
# Iterate until we converge to the solution to account for
# the non-linear interactions
theta = 0
theta0 = 1
while np.max(abs(theta0 - theta)) > tol:
theta0 = theta
d_ = d0_ - gn1_ * np.cos(theta) ** 2 + gn2_ * np.sin(theta) ** 2
D_ = -branch * np.sqrt((k_ - d_) ** 2 + w_**2)
B_ = k_ - d_
theta = np.arctan2(-B_ - D_, C_)
return np.array(
[
np.cos(theta) * np.sqrt(n0) * np.exp(1j * (k + k_r) * x),
np.sin(theta) * np.sqrt(n0) * np.exp(1j * (k - k_r) * x),
]
)
######################################################################
# 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.
@attr.s
[docs]
class PlotElements:
"""Collection of elements for redrawing a plot."""
[docs]
densities = attr.ib(attr.Factory(Namespace))
[docs]
history = attr.ib(attr.Factory(Namespace))
[docs]
momenta = attr.ib(attr.Factory(Namespace))
[docs]
master_grid = attr.ib(None)
[docs]
def _repr_html_(self):
"""Provide display capability in IPython."""
from IPython.display import display
return display(self.fig)
[docs]
def plot(
self,
plot_elements=None,
fig=None,
data=None,
grid=None,
show_n=True,
show_na=True,
show_nb=True,
show_log_n=True,
show_momenta=False,
show_mixtures=False,
show_V=False,
show_history_ab=True,
show_history_a=False,
show_history_b=False,
show_log_history=False,
combine_momenta_and_history=False,
dynamic_range=100,
log_dynamic_range=5,
parity=False,
# kwargs for get_momenta_data
clip_momenta_modes=1,
k_range_k_r=(-2.5, 2.5),
nk_amplitude_factor=0.5,
mu_background=None,
symmetric=True,
fig_width=15,
pane_height=2,
space=0.1,
history=None,
split=False,
rescale=False,
plot_dim=None,
subplot_spec=None,
# These are full-panel components.
): # pragma: nocover
"""Plot the state.
Parameters
----------
data : PlotData
Instance returned by self.get_plot_data() of the data to be plotted.
plot_elements : PlotElements
If defined, then update the plot here (faster, but will not rescale).
history : [State]
List of states defining the frames.
parity : bool
If `True`, then plot abs(x) on the abscissa of the density plots to
show parity violations.
clip_momenta_modes : int
Clip this many momenta modes (peaks). If most of the cloud occupies
a single background momentum state, then there will be a large peak
that obscures features. This many peaks will be clipped.
plot_dim : None, int
If provided, then try to reduce plots to this dimension. If None,
then use self.dim. This is particularly useful for plot_dim = 1
which will plot the line density along the x axis, averaging or
integrating over the other dimensions.
combine_momenta_and_history : bool
If True, then combine the momenta and history plots in a single pane.
The following "boolean" flags can actually be numbers which will
specify the relative height of the corresponding pane in the plot:
show_n, show_mixtures, show_momenta, show_history_ab, show_history_a,
show_history_b
"""
from matplotlib import pyplot as plt
from gpe.plot_utils import MPLGrid
if grid is not None:
assert fig is None
fig = grid.fig
if isinstance(fig, self.PlotElements):
# Backwards compatibility
fig, plot_elements = None, plot_elements
if plot_elements is None:
pe_ = self.PlotElements(fig=fig)
elif isinstance(plot_elements, self.PlotElements):
pe_ = plot_elements
elif fig is None:
# Assume a fig was passed instead (old behaviour)
pe_ = self.PlotElements(fig=plot_elements)
plot_elements = None
fig = pe_.fig
if data is None:
get_momenta_kw = dict(
k_range_k_r=k_range_k_r,
clip_momenta_modes=clip_momenta_modes,
log_dynamic_range=log_dynamic_range,
mu_background=mu_background,
)
data = self.get_plot_data(
states=history,
dynamic_range=dynamic_range,
get_momenta_data=show_momenta,
**get_momenta_kw,
)
# from mmfutils.plot import imcontourf
show_histories = (
show_history_ab + show_history_a + show_history_b + show_log_history
)
if plot_dim is None:
if show_histories:
dim = 1
else:
dim = self.dim
else:
dim = plot_dim
if dim == 1:
if not history:
show_histories = show_log_history = False
show_history_ab = show_history_a = show_history_b = False
if show_histories and show_momenta and combine_momenta_and_history:
show_momenta_history = max(show_momenta, show_histories)
else:
show_momenta_history = show_momenta + show_histories
if fig is None:
panes = sum(
[
not split and (show_n or show_na or show_nb),
split and show_n,
split and show_na,
split and show_nb,
show_mixtures,
show_momenta_history,
]
)
fig = pe_.fig = plt.figure(figsize=(fig_width, pane_height * panes))
if split:
show_densities = show_n + show_na + show_nb + show_mixtures
else:
show_densities = max(show_n, show_na, show_nb) + show_mixtures
density_grid = None
momenta_grid = None
history_grid = None
if not plot_elements:
# Allocate all grids:
if grid is None:
master_grid = MPLGrid(
direction="down", share=False, fig=fig, space=space
)
else:
master_grid = grid
density_grid = master_grid.grid(show_densities)
if combine_momenta_and_history:
momenta_history_grid = master_grid.grid(
show_momenta_history, direction="right"
)
momenta_grid = momenta_history_grid.grid()
history_grid = momenta_history_grid.grid()
else:
momenta_grid = master_grid.grid(show_momenta)
history_grid = master_grid.grid(show_histories)
pe_.master_grid = master_grid
with NoInterrupt(ignore=False):
if show_densities:
self.plot_densities(
data=data,
plot_elements=pe_.densities,
grid=density_grid,
show_n=show_n,
show_na=show_na,
show_nb=show_nb,
show_mixtures=show_mixtures,
show_V=show_V,
show_log_n=show_log_n,
split=split,
parity=parity,
symmetric=symmetric,
dynamic_range=dynamic_range,
log_dynamic_range=log_dynamic_range,
)
if show_momenta:
self.plot_momenta(
data=data,
plot_elements=pe_.momenta,
grid=momenta_grid,
amplitude_factor=nk_amplitude_factor,
)
if show_histories:
self.plot_history(
data=data,
plot_elements=pe_.history,
grid=history_grid,
show_history_ab=show_history_ab,
show_history_a=show_history_a,
show_history_b=show_history_b,
show_log_history=show_log_history,
)
E = self.get_energy()
N = self.get_N()
if self.single_band:
msg = "t={:.4g}t0, N={:.4g}, E={:.4g}".format(self.t / self.t0, N, E)
else:
Na, Nb = self.get_Ns()
msg = "t={:.4g}t0, Ns={:.4g}+{:.4g}={:.4g}, E={:.4g}".format(
self.t / self.t0, Na, Nb, N, E
)
if not plot_elements:
pe_.title = plt.suptitle(msg)
# Add space for title:
# https://stackoverflow.com/a/45161551
plt.tight_layout(rect=[0, 0.03, 1, 0.95])
else:
pe_.title.set_text(msg)
if rescale:
pe_.master_grid.rescale(tight=None, scalex=True, scaley=True)
return pe_
elif dim == 2:
from mmfutils import plot as mmfplt
if fig is None:
fig = plt.figure(figsize=(10, 5))
x, y = self.xyz[:2]
plt.subplot(121)
z = self[0]
vmax = 1.1 * abs(self[...]).max()
mmfplt.imcontourf(x, y, mmfplt.colors.color_complex(z, vmax=vmax), aspect=1)
mmfplt.phase_contour(x, y, z, aspect=1, linewidths=0.5)
plt.subplot(122)
z = self[1]
mmfplt.imcontourf(x, y, mmfplt.colors.color_complex(z, vmax=vmax), aspect=1)
mmfplt.phase_contour(x, y, z, aspect=1, linewidths=0.5)
else:
raise NotImplementedError("Only 1D and 2D plotting supported")
E = self.get_energy()
N = self.get_N()
if self.single_band:
plt.suptitle("t={}t0, N={:.4f}, E={:.4f}".format(self.t / self.t0, N, E))
else:
Na, Nb = self.get_Ns()
plt.suptitle(
"t={}t0, Ns={:.4f}+{:.4f}={:.4f}, E={:.4f}".format(
self.t / self.t0, Na, Nb, N, E
)
)
pe_.fig = fig
return pe_
[docs]
def plot_densities(
self,
data,
plot_elements,
grid=None,
split=False,
show_n=True,
show_na=True,
show_nb=True,
show_mixtures=False,
show_log_n=False,
show_V=False,
parity=False,
symmetric=False,
dynamic_range=100,
log_alpha=0.3,
log_dynamic_range=5,
):
# Determine plot limits.
from matplotlib import pyplot as plt
xs = np.ma.MaskedArray(data.xs, mask=data.mask)
log_xs = np.ma.MaskedArray(data.xs, mask=data.log_mask)
if grid:
# Drawing for the first time.
# Record these to ensure that updates make sense
plot_elements.show_n = show_n
plot_elements.show_na = show_na
plot_elements.show_nb = show_nb
plot_elements.show_log_n = show_log_n
plot_elements.show_V = show_V
plot_elements.show_mixtures = show_mixtures
x_min = xs.min()
x_max = xs.max()
x_log_min = log_xs.min()
x_log_max = log_xs.max()
if symmetric:
x_min = min(x_min, -x_max)
x_max = max(x_max, -x_min)
x_log_min = min(x_log_min, -x_log_max)
x_log_max = max(x_log_max, -x_log_min)
if split:
size = show_n + show_na + show_nb + show_mixtures
grid = grid.grid(size, direction="down", share=True)
ax_n = grid.next(show_n)
ax_na = grid.next(show_na)
ax_nb = grid.next(show_nb)
else:
n_size = max(show_n, show_na, show_nb)
size = n_size + show_mixtures
grid = grid.grid(size, direction="down", share=True)
ax = grid.next(n_size)
ax_n = show_n and ax
ax_na = show_na and ax
ax_nb = show_nb and ax
lines = plot_elements.lines = []
cs = plot_elements.colors = []
for _ax, _n, _label, _ls in [
(ax_n, data.n, "a+b", ("-", "--")),
(ax_na, data.na, "a", ("--", "-.")),
(ax_nb, data.nb, "b", ("--", "-.")),
]:
if _ax:
ax = _ax
if parity:
left, right = np.where(xs <= 0), np.where(xs >= 0)
lines.extend(
_ax.plot(
-xs[left], _n[left], _ls[0], label=_label + " (left)"
)
)
cs.append(lines[-1].get_c())
lines.extend(
_ax.plot(
xs[right],
_n[right],
_ls[1],
label=_label + " (right)",
c=cs[-1],
)
)
_ax.set_xlim(0, x_max, auto=None)
else:
lines.extend(_ax.plot(xs, _n, _ls[0], label=_label))
cs.append(lines[-1].get_c())
_ax.set_xlim(x_min, x_max, auto=None)
if _label == "a+b":
_ax.set_ylim(*data.n_lim, auto=None)
_ax.set_ylabel("n")
if not split:
ax.legend()
if show_log_n:
lines = plot_elements.lines_log = []
ax_log = None
cs = list(plot_elements.colors)
for _ax, ys in [
(ax_n, data.log_n),
(ax_na, data.log_na),
(ax_nb, data.log_nb),
]:
if _ax:
c = cs.pop(0)
if not ax_log or split:
ax_log = _ax.twinx()
ax_log.set_ylim([-log_dynamic_range, 1], auto=None)
plt.ylabel(r"$\log_{10}(n/n_\max)$")
args = dict(c=c, alpha=log_alpha)
if parity:
lines.extend(ax_log.plot(-xs[left], ys[left], "-", **args))
lines.extend(ax_log.plot(xs[right], ys[right], "-.", **args))
else:
lines.extend(ax_log.plot(xs, ys, "-", **args))
if show_mixtures:
grid.size = show_mixtures
plot_elements.mixtures = Namespace()
self.plot_mixtures(
data=data,
plot_elements=plot_elements.mixtures,
grid=grid,
show_V=show_V,
)
ax = plot_elements.mixtures.ax
ax.set_xlabel("x [micron]")
else:
# Just update lines
xs = data.xs
ls = list(plot_elements.lines)
if parity:
left, right = np.where(xs <= 0), np.where(xs >= 0)
if plot_elements.show_n:
ls.pop(0).set_data(-xs[left], data.n[left])
ls.pop(0).set_data(xs[right], data.n[right])
if plot_elements.show_na:
ls.pop(0).set_data(-xs[left], data.na[left])
ls.pop(0).set_data(xs[right], data.na[right])
if plot_elements.show_nb:
ls.pop(0).set_data(-xs[left], data.nb[left])
ls.pop(0).set_data(xs[right], data.nb[right])
if plot_elements.show_log_n:
ls = list(plot_elements.lines_log)
if plot_elements.show_n:
ys = data.log_n
ls.pop(0).set_data(-xs[left], ys[left])
ls.pop(0).set_data(xs[right], ys[right])
if plot_elements.show_na:
ys = data.log_na
ls.pop(0).set_data(-xs[left], ys[left])
ls.pop(0).set_data(xs[right], ys[right])
if plot_elements.show_nb:
ys = data.log_nb
ls.pop(0).set_data(-xs[left], ys[left])
ls.pop(0).set_data(xs[right], ys[right])
else:
left, right = np.where(xs <= 0), np.where(xs >= 0)
if plot_elements.show_n:
ls.pop(0).set_data(xs, data.n)
if plot_elements.show_na:
ls.pop(0).set_data(xs, data.na)
if plot_elements.show_nb:
ls.pop(0).set_data(xs, data.nb)
if plot_elements.show_log_n:
ls = list(plot_elements.lines_log)
if plot_elements.show_n:
ls.pop(0).set_data(xs, data.log_n)
if plot_elements.show_na:
ls.pop(0).set_data(xs, data.log_na)
if plot_elements.show_nb:
ls.pop(0).set_data(xs, data.log_nb)
if plot_elements.show_mixtures:
self.plot_mixtures(data=data, plot_elements=plot_elements.mixtures)
return plot_elements
@property
[docs]
def time_dependent_quantities(self):
time_dependent_quantities = []
for name in self.experiment.time_dependent_parameters:
method = self.experiment.time_dependent_parameters[name]
info_name = name + "_info"
unit, label = getattr(self.experiment, info_name)
time_dependent_quantities.append((method, unit, label))
return time_dependent_quantities
[docs]
def plot_time_dependent_parameters(
self, fig=None, subplot_spec=None, share=False
): # pragma: nocover
"""Plot all the time-dependent parameters stacked on top of each
other. Used as part of other plotting routines.
"""
from gpe.plot_utils import MPLGrid
ts_ = np.linspace(0, self.experiment.t__final, 200)
t_ = self.t / self.t_unit
grid = MPLGrid(fig=fig, subplot_spec=subplot_spec, right=True, hspace=0.0)
for _n, (_f_t_, _unit, _label) in enumerate(self.time_dependent_quantities):
ax = grid.next(share=share)
share = True
ys = np.array(map(_f_t_, ts_)) / _unit
ls = ax.plot(ts_, ys)
ys = np.ravel(_f_t_(t_)) / _unit # Ravel to make scalars a 1D array
for _y, _l in zip(ys, ls):
ax.plot(t_, _y, "o", c=_l.get_c())
ax.set_ylabel(_label)
ax.set_xlabel("t [ms]")
return ax
[docs]
def plot_mixtures(
self, data, plot_elements, grid=None, dynamic_range=100, show_V=False
):
"""Plot the occupation of the two bands."""
x = data.xs
n_p, n_m = self.get_branch_mixture()
y_m = np.ma.masked_array(n_m / (n_p + n_m), mask=data.mask)
y_p = np.ma.masked_array(n_p / (n_p + n_m), mask=data.mask)
if getattr(plot_elements, "show_V", show_V):
Va, Vb, Vab = [
np.ma.masked_array(_V, data.mask)
for _V in self.get_Vext(single_band=False)
]
if grid:
# Drawing for the first time.
ax = plot_elements.ax = grid.next()
plot_elements.lines = ax.plot(x, y_m, "b-", label="Lower branch") + ax.plot(
x, y_p, "r:", label="Upper branch"
)
ax.legend()
plot_elements.show_V = show_V
if show_V:
ax_V = ax.twinx()
zero = np.zeros_like(x)
plot_elements.lines_V = ax_V.plot(
x, Va + zero, "b--", label="Va"
) + ax_V.plot(x, Vb + zero, "g--", label="Vb")
ax_V.legend()
else:
plot_elements.lines[0].set_data(x, y_m)
plot_elements.lines[1].set_data(x, y_p)
if plot_elements.show_V:
plot_elements.lines_V[0].set_data(x, Va)
plot_elements.lines_V[1].set_data(x, Vb)
return plot_elements
[docs]
def plot_history(
self,
data,
plot_elements,
grid=None,
rescale=False,
show_history_ab=True,
show_history_a=False,
show_history_b=False,
show_log_history=False,
dynamic_range=100,
log_dynamic_range=5,
**kw,
):
"""Show an imcountorf plot of of the states.
Parameters
----------
states : list
List of states to show.
rescale : bool
If True, then scale each time-slice so the maximum is 1.
"""
from matplotlib import pyplot as plt
from mmfutils.plot import imcontourf
ts = data.ts
xs = data.xs
t = data.t
if not grid:
# Just update the line positions
for line in plot_elements.lines:
line.set_data([[t] * 2, [xs.min(), xs.max()]])
return plot_elements
# Drawing for the first time.
grid.share = True
if rescale:
ns = data.ns / data.ns.max(axis=1)[:, None]
nas = data.nas / data.nas.max(axis=1)[:, None]
nbs = data.nbs / data.nbs.max(axis=1)[:, None]
else:
ns, nas, nbs = data.ns, data.nas, data.nbs
n_max = ns.max()
x_log_min = np.ma.masked_array(xs, data.log_mask).min()
x_log_max = np.ma.masked_array(xs, data.log_mask).max()
x_min = np.ma.masked_array(xs, data.mask).min()
x_max = np.ma.masked_array(xs, data.mask).max()
histories = []
if show_history_ab:
histories.append((show_history_ab, ns, "$n$"))
if show_history_a:
histories.append((show_history_a, nas, "$n_a$"))
if show_history_b:
histories.append((show_history_b, nbs, "$n_b$"))
plot_elements.lines = []
plot_elements.imgs = []
for _size, _ns, _label in histories:
_grid = grid.grid(size=_size, direction="right", space=0)
ax = _grid.next()
plot_elements.imgs.append(imcontourf(ts, xs, _ns))
ax.set_xlim([ts.min(), ts.max()], auto=None)
ax.set_ylim([x_min, x_max], auto=None)
plot_elements.lines.append(plt.axvline([data.t], c="r", ls=":"))
plt.ylabel("x [micron]")
ax.xaxis.set_visible(False)
cax_hist = _grid.next(0.01)
plt.colorbar(cax=cax_hist, label=_label)
if show_log_history:
_grid = grid.grid(size=show_log_history, direction="right", space=0)
ax = _grid.next()
plot_elements.imgs.append(
imcontourf(ts, xs, np.log10(ns / n_max), vmin=-log_dynamic_range)
)
ax.set_xlim([ts.min(), ts.max()], auto=None)
ax.set_ylim([x_log_min, x_log_max], auto=None)
plot_elements.lines.append(plt.axvline([data.t], c="r", ls=":"))
plt.ylabel("x [micron]")
cax_log_hist = _grid.next(0.01)
plt.colorbar(cax=cax_log_hist, label="log10(n/n_max)")
ax.xaxis.set_visible(False)
ax.xaxis.set_visible(True)
ax.set_xlabel("t [ms]")
plot_elements.ax = ax
return plot_elements
[docs]
def get_momenta_data(
self,
clip_momenta_modes=1,
log_dynamic_range=5,
mu_background=None,
k_range_k_r=(-2.5, 2.5),
):
"""Return (ks, Em, Ep, E_ph, nk, log_nk).
Arguments
---------
clip_momenta_modes : int
Clip this many momenta modes (peaks). If most of the cloud occupies
a single background momentum state, then there will be a large peak
that obscures features. This many peaks will be clipped.
log_dynamic_range : int
Return the log of the density scaled so that 0 corresponds to a
10**log_dynamic_range suppression.
k_range_k_r : (float, float)
Range of momenta (in units of k_r).
mu_background: float, None
Background chemical potential used for computing the phonon
dispersion. If this is None, the the maximum density will be used.
Returns
-------
ks : array
Momenta abscissa.
Em, Ep : array
Dispersion bands.
Ph_m, Ph_p : array
Phonon dispersions
nk : array
Momentum distribution ranging from 0 to the `Em.max()-Em.min()`.
This is intended to be added to Em for plotting.
log_nk : array
Log of the momentum distribution ranging from 0 to the
`Em.max()-Em.min()`. This is intended to be added to Em for
plotting.
"""
Omega = self.get("Omega")
detuning = self.get("detuning")
k_r = self.get("k_r")
E_R = self.get("E_R")
d = detuning / 4 / E_R
w = Omega / 4 / E_R
if self.single_band:
# Single-band models are already in the rotating phase basis.
nk = np.fft.fftshift(abs(np.fft.fft(self[...], axis=0) ** 2))
else:
# Go to rotating phase basis
ytilde = self[...] * np.exp(
-1j * k_r * self.get_xyz()[0] * np.array([1, -1])[self.bcast]
)
na_k, nb_k = np.fft.fftshift(abs(np.fft.fft(ytilde, axis=1) ** 2))
nk = na_k + nb_k
ka = np.fft.fftshift(self.basis.kx.ravel())
def get_E(q):
D = np.sqrt((q - d) ** 2 + w**2)
Ep = (q**2 + 1) / 2.0 + D
Em = (q**2 + 1) / 2.0 - D
return Em, Ep
km, kp = np.asarray(k_range_k_r) * k_r
_ks = np.linspace(km, kp, 100)
Em, Ep = get_E(_ks / k_r)
Em_range = Em.max() - Em.min()
inds = np.where(np.logical_and(km <= ka, ka <= kp))[0]
ks = ka[inds]
Em, Ep = get_E(ks / k_r)
_inds = np.argsort(nk)
k0 = ka[_inds[-1]]
_inds = _inds[-clip_momenta_modes - 1 :]
nk[_inds] = nk[_inds].min()
log_nk = (
np.maximum(np.log10(nk / nk.max()), -log_dynamic_range) + log_dynamic_range
)
nk *= Em_range / nk.max()
log_nk *= Em_range / log_nk.max()
nk = nk[inds]
log_nk = log_nk[inds]
if mu_background is None:
# Guess background chemical potential
mu_background = np.mean(self.gs) * self.get_density().max()
# Compute phonon dispersion relationships
def get_E_ph(k, k0, mu_=mu_background / 2 / E_R):
q = k - k0
Em_0, Ep_0 = get_E(k0)
Em_p, Ep_p = get_E(k0 + q)
Em_m, Ep_m = get_E(k0 - q)
Em = (Em_p - Em_m) / 2.0
Ep = (Em_p + Em_m - 2 * Em_0) / 2.0
E_ph = Em + np.sqrt(Ep * (Ep + 2 * mu_))
# E_ph_m = Em - np.sqrt(Ep*(Ep + 2*mu_))
return E_ph + Em_0
E_ph = get_E_ph(ks / k_r, k0=k0 / k_r)
return self.MomentaData(ks, Em, Ep, E_ph, nk, log_nk, d, w)
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
[docs]
PlotData = namedtuple(
"PlotData",
(
"frame",
"frames",
"xs",
"ts",
"t",
"na",
"nb",
"n",
"nas",
"nbs",
"ns",
"log_na",
"log_nb",
"log_n",
"mask",
"log_mask",
"n_lim",
"state",
)
+ MomentaData._fields,
)
[docs]
def get_plot_data(
self,
states=None,
frame=None,
t_unit=u.ms,
x_unit=u.micron,
n_unit=1.0 / u.micron**3,
dynamic_range=100,
log_dynamic_range=5,
get_momenta_data=False,
**get_momenta_kw,
):
"""Return a PlotData namedtuple with the data needed for plotting.
Arguments
---------
states : None, [State]
If a list of states is provided, then this is used to generate data
for a history vs time contour plot.
frame : int
Frame when making movies.
**get_momenta_kw :
Additional keywords are passed to get_momenta_data.
"""
nas = []
nbs = []
ns = []
ts = []
frames = None
if states:
frames = len(states)
for state in states:
na, nb = state.get_density_x()
nas.append(na)
nbs.append(nb)
ts.append(state.t)
nas, nbs, ts = map(np.array, [nas, nbs, ts])
ts /= t_unit
nas /= n_unit
nbs /= n_unit
ns = nas + nbs
_n = ns.max(axis=0)
n_min = ns.min()
if frame is None:
frame = states.index(self) if self in states else 0
if states[frame] is self:
na, nb = nas[frame], nbs[frame]
else:
na, nb = self.get_density_x()
warnings.warn(
"Inconsistent frame={} found: self != states[frame]".format(frame)
)
else:
na, nb = self.get_density_x() / n_unit
_n = na + nb
n_min = _n.min()
n_max = _n.max()
mask = _n < n_max / dynamic_range
log_mask = _n < n_max / 10**log_dynamic_range
na = np.ma.MaskedArray(na, mask=mask)
nb = np.ma.MaskedArray(nb, mask=mask)
n = np.ma.MaskedArray(na + nb, mask=mask)
log_na = np.ma.MaskedArray(np.log10(na / n_max), mask=log_mask)
log_nb = np.ma.MaskedArray(np.log10(nb / n_max), mask=log_mask)
log_n = np.ma.MaskedArray(np.log10(n / n_max), mask=log_mask)
xs = self.xyz[0].ravel() / x_unit
if get_momenta_data:
momenta_data = self.get_momenta_data(**get_momenta_kw)
else:
momenta_data = self.MomentaData(*(None,) * len(self.MomentaData._fields))
return self.PlotData(
frame=frame,
frames=frames,
xs=xs,
ts=ts,
t=self.t / t_unit,
na=na,
nb=nb,
n=n,
nas=nas,
nbs=nbs,
ns=ns,
log_na=log_na,
log_nb=log_nb,
log_n=log_n,
mask=mask,
log_mask=log_mask,
n_lim=(n_min, n_max),
state=self,
**momenta_data._asdict(),
)
[docs]
def plot_momenta(self, data, plot_elements, grid=None, amplitude_factor=0.5, ax=None):
"""Plot the dispersion and momentum distribution. Return plot_elements.
Arguments
---------
amplitude_factor : float
The occupation is shown added on top of the dispersion such that the
maximum occupation is this fraction of the dispersion range.
I.e. If amplitude_factor=0.5, then the maximum momentum occupation
will be half of the dispersion range.
"""
Em, Ep, E_ph = data.Em, data.Ep, data.E_ph
ks_kr = data.ks / self.k_r
_i = np.where(np.logical_and(ks_kr < 1.5, -2.5 < ks_kr))
ylim = (min(Em.min(), E_ph[_i].min()), max(Em.max(), E_ph[_i].max()))
if grid:
# Drawing for the first time.
ax = plot_elements.ax = grid.next()
plot_elements.lines = (
ax.plot(ks_kr, Em, "b-")
+ ax.plot(ks_kr, E_ph, "b:")
+ ax.plot(ks_kr, Ep, "r-")
+ ax.plot(ks_kr, Em + amplitude_factor * data.nk, "--")
+ ax.plot(ks_kr, Em + amplitude_factor * data.log_nk, ":y")
)
ax.set_ylim(ylim, auto=None)
ax.set_xlabel("$k$ [$k_R$]")
else:
# Update existing plot
lines = plot_elements.lines
lines[0].set_data(ks_kr, Em)
lines[1].set_data(ks_kr, E_ph)
lines[2].set_data(ks_kr, Ep)
lines[3].set_data(ks_kr, Em + amplitude_factor * data.nk)
lines[4].set_data(ks_kr, Em + amplitude_factor * data.log_nk)
return plot_elements
[docs]
def make_movie(
self,
filename,
states,
skip=1,
fps=20,
skip_frames=10,
show_frames=2,
dpi=None,
display=True,
extra_args=("-vcodec", "libx264", "-pix_fmt", "yuv420p"),
**kw,
):
"""Make a movie of the states.
Arguments
---------
filename : str
Name of movie file such as "movie.mp4".
states : [State]
List of states defining the frames.
skip : int
If this is >1, then skip this many states when making the movie.
This allows one to have more states for smooth history plots, but
speeds making the movie.
fps : int
Frames-per-second for output movie
show_frames : int
Display this many initial frames to aid debugging.
Does not affect the generated movie.
skip_frames : int
After the initial frames, draw only these frames. (Drawing frames
greatly slows down the making of the movie so we limit the display).
Does not affect the generated movie.
display : bool
If `False`, then do not display anything. (Use for offline generation
of movies for example.)
"""
if display:
from IPython.display import display, clear_output
from matplotlib import pyplot as plt
from mmfutils.plot import animation
movie_states = states[::skip]
def get_data(states):
if movie_states[-1] is not states[-1]:
movie_states.append(states[-1])
for frame, state in enumerate(movie_states):
yield frame, state
fig_ = [None]
@coroutine
def get_plot_data(fig_=fig_):
tic = time.time()
tocs = []
plot_elements = states[0].plot(history=movie_states, **kw)
fig_[0] = plot_elements.fig
with NoInterrupt(ignore=True) as interrupted:
while not interrupted:
frame, state = yield plot_elements.fig
frames = len(movie_states)
plot_elements = state.plot(
plot_elements=plot_elements, history=movie_states, **kw
)
tocs.append(time.time() - tic)
avg_time_for_frame = np.mean(np.diff(tocs))
time_to_completion = (frames - frame) * avg_time_for_frame
msg = "{} of {}: {:1g}s remaining ({}ms per frame)".format(
frame, frames, time_to_completion, avg_time_for_frame
)
if display:
if frame < show_frames or frame % skip_frames == 0:
clear_output(wait=True)
display(plot_elements.fig)
display(msg)
else:
if frame < show_frames or frame % skip_frames == 0:
logging.info(msg)
self._l = locals()
with get_plot_data(fig_) as plot_data:
anim = animation.MyFuncAnimation(
fig_[0], plot_data, get_data(states), save_count=len(states)
)
anim.save(filename, fps=fps, dpi=dpi, extra_args=extra_args)
plt.close("all")
[docs]
class State2(SOCMixin, bec2.State):
def __init__(self, experiment, **kw):
[docs]
self.experiment = experiment
super().__init__(**kw)
[docs]
self._time_independent_Vext = False
[docs]
self._Vext = [None, None] # Cache for performance
[docs]
class State2Tube(SOCMixin, tube2.StateGPEdrZ):
def __init__(self, experiment, **kw):
[docs]
self.experiment = experiment
super().__init__(**kw)
[docs]
self._time_independent_Vext = False
[docs]
self._Vext = [None, None] # Cache for performance
[docs]
class State2Axial(SOCMixin, axial.State2Axial):
def __init__(self, experiment, **kw):
[docs]
self.experiment = experiment
super().__init__(**kw)
[docs]
self._time_independent_Vext = False
[docs]
self._Vext = [None, None] # Cache for performance
[docs]
class State1Mixin:
"""Mixin for single-band states."""
# The main task here is to replace the dispersion relationship K with an
# appropriately modified form. This is done by modifying the laplacian
# function.
[docs]
def laplacian(self, y, factor, exp=False):
"""Return the laplacian applied to the state `y` multiplied by
`factor`.
Arguments
---------
y : State
The laplacian will be applied to this state (not `self`).
factor : array-like
The result will be multiplied by this factor.
exp : bool
If `True` then `exp(factor*laplacian)(y)` will be computed instead.
"""
return self.basis.laplacian(y=y, factor=factor, kx2=self.kx2)
@property
[docs]
def kx2(self):
"""Return the effective `kx**2` for use in the laplacian."""
# dispersion()*2*E_R = dispersion()*2 * (hbar**2 * k_r**2 / 2/ m)
# = hbar**2 * kx**2 / 2 / m
t_ = self.t / self.t_unit
dispersion = self.experiment.get_dispersion(t_=t_)
k_r = self.get("k_r")
return 2 * k_r**2 * dispersion(self.kx / k_r)[0]
[docs]
class State1(State1Mixin, SOCMixin, bec.StateTwist_x):
def __init__(self, experiment, **kw):
[docs]
self.experiment = experiment
super().__init__(**kw)
[docs]
self._time_independent_Vext = False
[docs]
class State1Tube(State1Mixin, SOCMixin, tube.StateGPEdrZ):
"""Tube state for single band model."""
def __init__(self, experiment, **kw):
[docs]
self.experiment = experiment
super().__init__(**kw)
[docs]
self._time_independent_Vext = False
[docs]
self._Vext = [None, None] # Cache for performance
[docs]
class State1Axial(State1Mixin, SOCMixin, axial.StateAxial):
"""Axial state for single_band models."""
def __init__(self, experiment, **kw):
[docs]
self.experiment = experiment
super().__init__(**kw)
[docs]
self._time_independent_Vext = False
@implementer(IExperiment)
[docs]
class Experiment(utils.ExperimentBase):
"""Represents a set of experimental parameters.
1. Responsible for translating these into a proper State object. This
generally amounts to converting experimental parameters into values
useful to the state class. These are stored in self.state_args for use
in `self.get_state` and `self.get_initial_state`.
2. Provides get_Vext(state) with a (possibly time dependent) external
potential.
3. Implements an experimental protocol. To do this, subclass and define
the run_experiment method. This can be used to get the initial state,
then to manipulate it. (Not implemented yet.)
Note: these experiments are designed for SOC along the x-axis only. The
following basis options are supported:
'1D' : Periodic basis along the x-axis.
`tube==False` : Assume a constant density along y and z. This
should be thought of as modeling a 3D gas with translation symmetry
along the transverse directions. (The coupling constants are not
modified for modeling 1D systems.) The density here will be the
central density.
`tube==True` : Use the the tube code to model the transverse trap. The
density here will be the line density obtained by integrating over
the transverse directions. This will effectively modify the coupling
constants.
'axial' : Periodic basis along the x-axis with radial excitations assuming
axial symmetry. The transverse trapping frequency will be the average
of the y and z frequencies. The radial extent and lattice spacing will
be inferred from `Nx`, `dx` and `Lx` using the ratio of the trapping
frequencies at time `t_=0`.
'2D', '3D' : Full 2D or 3D simulation assuming no symmetry. The transverse
extents and lattice spacing should be provided in `Nxyz` and `Lxyz`,
otherwise it will be inferred from `Nx`, `dx` and `Lx` using the ratio
of the trapping frequencies at time `t_=0`.
Note: All times in the Experiment classes are specified in units of
`t_unit`. In arguments this is `t_ = t/t_unit` and we use `t_` as a kwarg
to ensure that this is properly used.
Parameters
----------
cells_x : None, int
If specified (not `None`), then the length of the box in the x direction
`Lxyz[0]` will be computed to be an integral number of cells to hold
`cells_x` periods of a wave with frequency `k_r`. In addition, the
trap will be reset to `'cos'` which will replace the HO potential with a
compatible, but periodic cosine. This allows one to simulate the center
of the trapped system with matching gradients and densities. If this is
defined, then the original `Lx` will be stored as `Lx_expt` which is used
to obtain initial states.
dx : None, (float,)
If specified (not `None`), then `Nx` will be computed such that his
minimum lattice spacing is realized, guaranteeing a certain UV cutoff.
initial_imbalance : float
Initial state is prepared in (1,-1) then this fraction is transfered to
the (1, 0) state.
Omega : float
This is the strength of the spin-orbit coupling.
delta : float
This is the value of detuning in energy units.
B_gradient : float
This is the strength of the magnetic field gradient which induces a
counterflow.
x0 : float
Location of the center of the trapping potential.
The following attributes are guaranteed to exist for use by the State
class.
Attributes
----------
k_r : float
Recoil wave-vector.
E_R : float
Recoil energy `E_R = (hbar*k_r)**2/2/m`.
"""
######################################################################
# Attributes required by IExperiment
[docs]
t_unit = u.ms # All times specified in this unit
[docs]
t__final = np.inf # Time when potentials are turned off for expansion.
[docs]
t__image = 7 # Expansion time for imaging
# End of attributes required by IExperiment
######################################################################
# Default values. These are not expected to change, but could be changed
# by passing kwargs to the base class __init__ method.
[docs]
species = ((1, -1), (1, 0)) # Spin states of the species
[docs]
gs = None # Coupling constants: determined from states
[docs]
trapping_frequencies_Hz = (3, 273, 277)
trapping_frequencies_Hz = (3.07, 278.08, 278.04)
[docs]
trapping_wavelength = 1064 * u.nm
[docs]
x_TF = 234.6 * u.micron # Thomas Fermi "radius" (where V(x_TF) = mu)
[docs]
mu = None # Alternative: specify mu
[docs]
detuning_kHz = 0.1 # 100.0*u.Hz up to kHz
[docs]
detuning_E_R = None # Optionally specify in E_R
[docs]
recoil_frequency_Hz = 3683.8 # Sets E_R (in Hz)
[docs]
rabi_frequency_E_R = 2.9 # Rabi frequency 1.0 - 2.9
[docs]
rabi_frequency_kHz = None # Optionally specify in kHz
[docs]
B_gradient_mG_cm = 10 # Gradient in -x direction.
[docs]
Nx = None # 2**14 # Default lattice and box size
[docs]
dx = 0.061 * 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)
[docs]
cells_x = None # Can be used instead to fix Lx
[docs]
old_cells_x = True # Use old definition of cells_x (2x larger)
[docs]
initial_imbalance = 0.5 # Fractional populations in initial state.
[docs]
twist = 0 # Twist phase over x
[docs]
x0 = 0.0 # Location of trap center
[docs]
basis_type = "1D" # '1D', '2D', '3D', `axial`
[docs]
gab_attenuation_t_ = None # Timescale for attenuating gab during
# imaging (see gs_t_()).
# The following are for single band models using either Disperion
# (if active_components == 2) or Dispersion3 (if active_components == 3).
[docs]
single_band = False # Use single-band (one-component) model.
[docs]
epsilon_E_R = 3.8 # epsilon parameter in Dispersion3
@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 = []
self.E_R = u.hbar * 2 * np.pi * self.recoil_frequency_Hz * u.Hz
if not hasattr(self, "detuning"):
if self.detuning_kHz is not None and self.detuning_E_R is None:
# Get detuning from detuning_E_kHz
self.detuning = self.detuning_kHz * u.kHz * (2 * np.pi * u.hbar)
self.detuning_E_R = self.detuning / self.E_R
elif self.detuning_E_R is not None and self.detuning_kHz is None:
# Get detuning from detuning_E_R
self.detuning = self.detuning_E_R * self.E_R
self.detuning_kHz = self.detuning / (2 * np.pi * u.hbar) / u.kHz
else:
raise ValueError(
"Both detuning_E_R={} and detuning_kHz={} set. Set one to None".format(
self.detuning_E_R, self.detuning_kHz
)
)
else:
self.detuning_E_R = self.detuning / self.E_R
self.detuning_kHz = self.detuning / (2 * np.pi * u.hbar) / u.kHz
if not hasattr(self, "rabi_frequency"):
if self.rabi_frequency_kHz is not None and self.rabi_frequency_E_R is None:
self.rabi_frequency = self.rabi_frequency_kHz * u.kHz
self.rabi_frequency_E_R = (
self.rabi_frequency * 2 * np.pi * u.hbar / self.E_R
)
elif self.rabi_frequency_E_R is not None and self.rabi_frequency_kHz is None:
self.rabi_frequency = (
self.rabi_frequency_E_R * self.E_R / (2 * np.pi * u.hbar)
)
self.rabi_frequency_kHz = self.rabi_frequency / u.kHz
else:
raise ValueError(
"Both rabi_frequency_E_R={} and rabi_frequency_kHz={} set.".format(
self.rabi_frequency_E_R, self.rabi_frequency_kHz
)
+ " Set one to None"
)
else:
self.rabi_frequency_E_R = self.rabi_frequency * 2 * np.pi * u.hbar / self.E_R
self.rabi_frequency_kHz = self.rabi_frequency / u.kHz
if not hasattr(self, "epsilon"):
self.epsilon = self.epsilon_E_R * self.E_R
else:
self.epsilon_E_R = self.epsilon / self.E_R
self.Omega = u.hbar * (2 * np.pi * self.rabi_frequency)
self.delta = self.detuning
if not hasattr(self, "B_gradient"):
self.B_gradient = self.B_gradient_mG_cm * u.mG / u.cm
_a, _b = self.species
scattering_lengths = np.array(
[u.scattering_lengths[_k] for _k in [(_a, _a), (_b, _b), (_a, _b)]]
)
# scattering_lengths = np.array([scattering_lengths[0]]*3)
self.mu_Bs = (u.magnetic_moment[_a], u.magnetic_moment[_b])
if not hasattr(self, "ms"):
self.ms = (u.masses[_a], u.masses[_b])
m = np.mean(self.ms)
self.k_r = np.sqrt(2 * m * self.E_R / u.hbar**2)
recoil_speed = u.hbar * self.k_r / m
msgs.append("recoil speed = {}mm/s".format(recoil_speed / (u.mm / u.s)))
msgs.append("recoil wavelength = {}nm".format(2 * np.pi / self.k_r / u.nm))
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.cells_x is not None:
# Special case where the length of the box along x will be an
# integral number of k_r. To ensure that homogeneous states remain
# eigenstates, we need to make sure that 2*pi*n/Lx = k_r for some
# n = cells_x. ****THINK***
if self.old_cells_x:
# This version is backward compatible...
self.Lx = 2 * np.pi * self.cells_x / (2 * self.k_r)
else:
self.Lx = 2 * np.pi * self.cells_x / self.k_r
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
if False and self.dim == 1:
msgs.append("Using 1D coupling constants")
# Convert to 1D values from Olshani
w_perp = np.prod(ws[1:]) ** (1.0 / len(ws[1:]))
a_perp = np.sqrt(u.hbar / m / w_perp)
zeta_1_2 = -1.46035450880958681288949915251529801246722933101258
gs = (
a_perp
* u.hbar
* w_perp
* 2
* scattering_lengths
/ a_perp
/ (1 - scattering_lengths / a_perp * abs(zeta_1_2))
)
else:
msgs.append("Using 3D coupling constants")
gs = 4 * np.pi * u.hbar**2 / m * scattering_lengths
if self.gs is None:
self.gs = gs
# Collect all time-dependent methods.
# Note: getmembers() accesses fiducial_V_TF which used to trigger an
# unwanted call to the potentially expensive get_fiducial_V_TF(). We
# preempt this here by setting it to None, then deleting it later.
# This also resets it which is consistent with init()'s semantics of
# resetting the experiment (parameters may have changed).
self._fiducial_V_TF = None
super().init()
del self._fiducial_V_TF
self.msgs = msgs
[docs]
def get_Vext(self, state, fiducial=False, expt=False):
"""Return (V_a, V_b, V_ab), the external potentials.
For `t_ > self.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
"""
# Convert times into experiment units
t_ = state.t / self.t_unit
Omega = self.get("Omega", t_=t_)
delta = self.get("delta", t_=t_)
zero = np.zeros(state.shape)
if t_ > self.t__final:
_Vext = (zero,) * 3
return _Vext
x_periodic = self.cells_x is not None and not fiducial and not expt
xyz = self._xyz_trap(state, x_periodic=x_periodic)
if expt:
Va, Vb = self.get_Vtrap_expt(state=state, xyz=xyz)
else:
Va, Vb = self.get_Vtrap(state=state, xyz=xyz)
if not fiducial:
Va_t, Vb_t = self.get_Vt(state=state)
Va += Va_t
Vb += Vb_t
if state.single_band:
V_ab = zero
V_Bs = (zero, zero)
else:
x = self._xyz_trap(state)[0]
B_gradient = self.get("B_gradient", t_=t_)
V_Bs = (B_gradient * np.asarray(state.mu_Bs)[state.bcast]) * (x[None, ...])
if state.rotating_phases:
V_ab = Omega / 2.0
else:
k_r = self.get("k_r", t_=t_)
V_ab = Omega / 2.0 * np.exp(2j * x * k_r)
_Vext = (Va - delta / 2.0 + V_Bs[0], Vb + delta / 2.0 + V_Bs[1], V_ab)
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,
k_r=self.k_r,
mu_Bs=self.mu_Bs,
gs=self.gs,
ms=self.ms,
twist=(self.twist, self.twist),
constraint="N",
)
if hasattr(self, "v_x"):
state_args["v_x"] = self.v_x
if self.single_band and not expt:
# Single band model is more limited
assert self.B_gradient == 0 # No way to simulate this with 1 component
del state_args["k_r"]
del state_args["mu_Bs"]
del state_args["ms"]
del state_args["gs"]
state_args.update(dict(twist=self.twist, m=self.ms[0], g=self.gs[0]))
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:
if self.single_band:
state = State1Tube(Nxyz=[Nx], Lxyz=[Lx], **state_args)
else:
state = State2Tube(Nxyz=[Nx], Lxyz=[Lx], **state_args)
elif self.basis_type == "axial":
from mmfutils.math.bases import CylindricalBasis
if self.R is None:
ws = self.ws_expt
wx, wr = ws[0], gmean(ws[1:])
a_perp = u.hbar / wr / min(self.ms)
# 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)
else:
R = self.R
if self.Nr is None:
Nr = int(np.ceil(R / self.dx))
else:
Nr = self.Nr
Nxr = [Nx, Nr]
Lxr = [Lx, R]
basis = CylindricalBasis(Nxr=Nxr, Lxr=Lxr, symmetric_x=False)
if self.single_band:
state = State1Axial(basis=basis, **state_args)
else:
state = State2Axial(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:]]
if self.single_band:
state = State1(Nxyz=Nxyz, Lxyz=Lxyz, **state_args)
else:
state = State2(Nxyz=Nxyz, Lxyz=Lxyz, **state_args)
# These should be provided by the Experiment so we remove them from the
# state to avoid any confusion.
if not state.single_band:
if hasattr(state, "Omega"):
del state.Omega
if hasattr(state, "delta"):
del state.delta
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
if state.rotating_phases:
k_r = self.get("k_r", t_=state.t)
state[0, ...] *= np.exp(1j * state.get_xyz()[0] * k_r)
if self.initial_imbalance is None:
# Cool to minimum with fixed chemical potential
state.mu = state.get_mu_from_V_TF(V_TF=self.get_fiducial_V_TF())
state.constraint = None
fix_N = False
else:
# Cool to minimum with only component 0 then RF transfered
state.constraint = "Nab"
state[0, :] = np.sqrt((abs(state[...]) ** 2).sum(axis=0))
state[1, :] = 0.0
fix_N = True
state.init()
if perturb > 0:
# Add a phase perturbation so that the state is not an eigenstate.
# This is useful in cases where the TF initial state is so good
# that the minimizer cannot initialize. (This is a bug with the
# minimizer).
state[:] *= np.exp(
1j
* perturb
* np.sin(2 * np.pi * state.get_xyz()[0] / self._get_Lx(state))
)
if minimize:
m = MinimizeState(state, fix_N=fix_N)
print("Minimize check:")
m.check()
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, use_scipy=True, **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
# Transfer particles to the second state
if self.initial_imbalance is not None:
state[1, ...] = np.sqrt(self.initial_imbalance) * state[0]
state[0, ...] *= np.sqrt(1 - self.initial_imbalance)
if state.rotating_phases:
state[1, ...] *= np.exp(-2j * state.get_xyz()[0] * k_r)
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.ms[state.bcast] * V_m
[docs]
def get_Vtrap(self, state, xyz):
"""Return the external trapping potential used in the simulation.
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
"""
return 0.0, 0.0
[docs]
def _xyz_trap(self, state, x_periodic=False):
"""Return the abscissa for use in the trapping potentials.
This has the abscissa shifted by x0 so that this value is zero, with
periodic wrapping. I.e. it assumes that the trapping potentials are
periodic with center at zero.
Parameters
----------
x_periodic : bool
If `True`, then transform the abscissa so that any potential
applied is smooth and periodic as discussed in the
Utils.ipynb notebook. This assumes additionally that the
potentials are even functions about zero.
"""
t_ = state.t / self.t_unit
xyz = list(state.get_xyz())
x = xyz[0]
Lx = self._get_Lx(state)
x_ = (x - self.get("x0", t_=t_) - x.min()) % Lx + x.min()
if x_periodic:
x_ = Lx / 2 * utils.x_periodic(2 * x_ / Lx)
xyz[0] = x_
return xyz
[docs]
def _get_Lx(self, state):
"""Return the actual length of the box represented by the state.
This is needed because self.Lx might not be the same as self.Lx_expt and
the latter is sometimes needed for fiducial states.
"""
x = state.get_xyz()[0]
dx = np.diff(x.ravel())
assert np.allclose(0, np.diff(dx))
dx = dx[0]
Lx = x.max() - x.min() + dx
return Lx
@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]
# E = self.get_dispersion()
# E0 = E(E.get_k0())[0]*4*E_R
# mu0 = mu0 - E0
# E_R = self.get('E_R', t_=t_)
expt_basis = CylindricalBasis(Nxr=Nxr, Lxr=Lxr, symmetric_x=False)
expt_state = State2(
basis=expt_basis,
experiment=self,
gs=self.gs,
ms=self.ms,
k_r=self.k_r,
x_TF=None,
mu_Bs=self.mu_Bs,
cooling_phase=self.cooling_phase,
constraint="N",
twist=(0, 0),
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 = np.mean(
self.get_Vtrap_expt(
state=expt_state, xyz=(x_TF,) + (0,) * (expt_state.dim - 1)
)
)
# Sanity check that V_TF is correctly defined
Va, Vb, Vab = self.get_Vext(state=expt_state, fiducial=True, expt=True)
# Use average potential for initial state.
V = (Va + Vb) / 2
assert np.all(np.where(expt_state.get_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_TF0 = V_TF
V_TF = brentq(f, V0, V1)
self._expt_state = expt_state
del self._computing_fiducial_V_TF
return V_TF
[docs]
def get_dispersion(self, t_=0.0, active_components=None):
"""Return Dispersion instance.
Arguments
---------
t_ : float
Time at which to get parameters.
active_components : 2, 3
Specify which dispersion to use - Dispersion or Dispersion3
"""
Omega = self.get("Omega", t_=t_)
delta = self.get("delta", t_=t_)
E_R = self.get("E_R", t_=t_)
if active_components is None:
active_components = self.active_components
if active_components == 2:
disp = Dispersion(w=Omega / 4 / E_R, d=delta / 4 / E_R)
elif active_components == 3:
epsilon = self.get("epsilon", t_=t_)
disp = Dispersion3(w=Omega / 4 / E_R, d=delta / 4 / E_R, e=epsilon / 2 / E_R)
# If we want, we can center this on zero as follows, but this makes it
# harder to compare with the two-band model.
# k0 = disp.get_k0()
# E0 = disp(k0)[0]
# disp = Dispersion(k0=k0, E0=E0, w=Omega/4/E_R, d=delta/4/E_R)
return disp
[docs]
def gs_t_(self, t_):
gs = np.array(self.gs)
if self.gab_attenuation_t_ is not None and t_ > self.t__final:
# Turn off inter-species coupling for expansion since the clouds
# rapidly drop away from each other due to Sturn-Gurlach imaging
gs[-1] *= np.exp(-t_ / self.gab_attenuation_t_)
return gs
@coroutine
[docs]
def plotter(self, fig=None, plot=True):
from matplotlib import pyplot as plt
if fig is None:
fig = plt.figure(figsize=(20, 5))
plt.clf()
lines = [
plt.plot([], [], "k-")[0],
plt.plot([], [], "b:")[0],
plt.plot([], [], "g:")[0],
]
title = plt.title("")
x = None
plt.xlabel("x [micron]")
plt.ylabel("n")
plt.xlim(-300, 300, auto=None)
plt.ylim(0, 300, auto=None)
while True:
state = yield lines
na, nb = state.get_density_x()
E = state.get_energy()
N = state.get_N()
t = state.t
title_text = "t={:.4f}ms, N={:.4f}, E={:.4f}".format(t / u.ms, N, E)
if x is None:
x = state.get_xyz()[0] / u.micron
# lines.extend(plt.plot(x, na, 'b'))
# lines.extend(plt.plot(x, nb, 'g'))
lines[0].set_data(x, na + nb)
lines[1].set_data(x, na)
lines[2].set_data(x, nb)
title.set_text(title_text)
if plot:
import IPython.display
IPython.display.display(fig)
IPython.display.clear_output(wait=True)
else:
print("{} < 1000ms".format(t / u.ms))
[docs]
def run(
self,
get_data,
fig=None,
figsize=(20, 5),
plot=True,
factor=8.0,
filename="generated_plots.mp4",
):
if plot:
import IPython.display
from matplotlib import pyplot as plt
from mmfutils.plot.animation import MyFuncAnimation
if fig is None:
fig = plt.figure(figsize=figsize)
with self.plotter(fig=fig, plot=plot) as plot_data:
anim = MyFuncAnimation(
fig,
plot_data,
frames=get_data(factor=factor, steps=1000),
save_count=None,
)
anim.save(
filename=filename,
fps=20,
extra_args=["-vcodec", "libx264", "-pix_fmt", "yuv420p"],
)
plt.close("all")
IPython.display.clear_output()
self.video = IPython.display.HTML(
'<video src="{0}" width="100%", type="video/mp4" controls/>'.format(filename)
)
[docs]
def plot_dispersion(self, t_=0.0, klims=(-2, 2), hv=None):
k = np.linspace(-2.5, 2.5, 100)
d = self.get("delta", t_=t_) / self.E_R / 4
w = self.get("Omega", t_=t_) / self.E_R / 4
D = np.sqrt((k - d) ** 2 + w**2)
Em = 0.5 * (k**2 + 1) - D
Ep = 0.5 * (k**2 + 1) + D
if hv:
return hv.Curve((k, Em)) + hv.Curve((k, Ep))
else:
from matplotlib import pyplot as plt
plt.plot(k, Em, k, Ep)
[docs]
class ExperimentStub:
[docs]
def get(self, key, t_):
return getattr(self, key)
[docs]
class ExperimentBarrier(Experiment):
"""Adds to the basic Experiment class a moving barrier potential (along the
`x` axis only) with the following parameters:
Attributes
----------
barrier_x : float
Position of barrier potential.
barrier_depth : float
Strength of the barrier potential (energy).
barrier_width : float
Width of the barrier potential (length)
barrier_type : 'gaussian', 'harmonic', 'constant'
Type of barrier potential. The 'constant' barrier will just provide the
cosine modulation.
barrier_k : float
If this is non-zero, then the barrier has an underlying cosine
modulation with this wavelength.
These use the function get_barrier_potential(state).
"""
[docs]
barrier_type = "gaussian"
[docs]
def _barrier_potential(self, x_barrier_width):
"""Return the barrier potential without the power factor.
x_barrier_width : x/barrier_width
"""
if self.barrier_type == "gaussian":
# Gaussian barrier
V_barrier = np.exp(-2.0 * (x_barrier_width) ** 2)
elif self.barrier_type == "harmonic":
# Harmonic barrier
V_barrier = np.maximum(0, 1 - 2.0 * (x_barrier_width) ** 2)
elif self.barrier_type == "step":
# Harmonic barrier
V_barrier = np.exp(-2.0 * (x_barrier_width) ** 4)
else:
# Constant barrier
V_barrier = 1 + 0 * x_barrier_width
return V_barrier, V_barrier
[docs]
def get_Vt(self, state):
t_ = state.t / self.t_unit
barrier_depth = self.get("barrier_depth", t_=t_)
barrier_width = self.get("barrier_width", t_=t_)
x = state.xyz_trap[0]
x0 = self.get("barrier_x", t_=t_)
# Make barrier potential periodic in box
Lx = self._get_Lx(state)
x = (x - x0 - x.min()) % Lx + x.min()
k = self.barrier_k
return (
barrier_depth
* np.cos(k * x)
* self._barrier_potential(x_barrier_width=x / barrier_width)
)
[docs]
class Bloch:
"""Classes and Glasses"""
K0 = Delta = 0.0
def __init__(self, experiment=ExperimentStub()):
[docs]
self.experiment = experiment
@property
[docs]
def k_r(self):
return self.experiment.k_r
@property
[docs]
def t_unit(self):
return self.experiment.t_unit
@property
[docs]
def Omega(self):
return self.experiment.get("Omega", t_=self._t / self.t_unit)
@property
[docs]
def delta(self):
return self.experiment.get("delta", t_=self._t / self.t_unit)
[docs]
def pack(self, psi):
return np.asarray(psi).view(dtype=float)
[docs]
def unpack(self, q):
return np.asarray(q).view(dtype=complex)
[docs]
def get_su2_generators(self):
return utils.pauli_matrices / 2.0
@property
[docs]
def beta(self):
return np.asarray(
[
self.Omega,
0,
self.hbar**2 / 2.0 / self.m * 4.0 * self.K0 * (self.k_r + self.Delta)
- self.delta,
]
)
[docs]
def get_beta_norm(self, beta=None):
if beta is None:
return np.linalg.norm(self.beta)
else:
return np.linalg.norm(beta)
[docs]
def get_Bloch_period(self, t_ms=None):
if t_ms is not None:
self._t = t_ms * self.t_unit
return 4.0 * np.pi * self.hbar / np.linalg.norm(self.beta) / self.t_unit
[docs]
def rhs(self, q, t):
self._t = t
_L2 = self.get_su2_generators()
psi = self.unpack(q=q)
_beta = self.beta
Hpsi = np.einsum("abc,a...,c...->b...", _L2, _beta, psi)
return self.pack(psi=Hpsi / 1j / self.hbar)
[docs]
def get_initial_state(self, theta=np.pi, phi=np.pi / 2.0):
"""returns an initial state packed for odeint.
theta is the azimuthal angle, phi is in the xy-plane
in the bloch sphere"""
q0 = self.pack(psi=[np.cos(theta / 2.0), np.sin(theta / 2.0) * np.exp(1j * phi)])
return q0
[docs]
def get_psi_t(self, t_=20.0, Nt_=5000):
ts = np.linspace(0, t_ * self.t_unit, Nt_)
q0 = self.get_initial_state(theta=np.pi, phi=np.pi / 2.0)
psis = self.unpack(q=odeint(self.rhs, q0, ts)).T
return ts / self.t_unit, psis
[docs]
def plot(self):
from matplotlib import pyplot as plt
ts_, psis = self.get_psi_t()
# plt.subplot(211)
plt.plot(ts_, (abs(psis) ** 2).T)
plt.plot(ts_, psis.real.T)
# plt.subplot(212)
# plt.plot(ts_, psis.imag.T)
[docs]
def get_Bloch_vector(self, psi):
"""returns the theta and phi of the Bloch vector
theta is the azimuthal angle, phi is in the xy-plane"""
theta1 = 2.0 * np.arccos(abs(psi[0]))
theta2 = 2.0 * np.arcsin(abs(psi[1]))
assert np.allclose(theta1, theta2)
phi = np.angle(psi[1]) - np.angle(psi[0]) - np.pi
return theta1, phi
@staticmethod
[docs]
def get_psi(a):
"""Return the two-component state from the vector `a` where the total
density is the length of `a` and the state points in the direction
specified with the top component real and positive."""
n = np.linalg.norm(a)
a = a / n
theta = np.arccos(a[2])
phi = np.angle(a[0] + 1j * a[1])
return np.sqrt(n) * np.array(
[np.cos(theta / 2.0), np.exp(1j * phi) * np.sin(theta / 2.0)]
)