"""Code for simulating a single component BEC.
This is the simplest code in the project, implementing all of the basic
features without excess complications. It is somewhat excessively commented to
provide a pedagogical introduction for users unfamiliar with the techniques and
using python for scientific computing.
A more complete implementation can be found in bec.py, but this relies on
external packages to implement, for example, different bases, making it less
suitable for teaching.
"""
# The start of every python files should contain a multi-line string like that
# above which describes the purpose of the file. This is called a docstring.
# The first line is the summary, followed by more details. You can see this
# documentation from the IPython prompt by adding a question mark after the
# module: `bec?` for example
from __future__ import absolute_import, division, print_function, unicode_literals
# This should be the first code line in the module. The __future__ module
# provides some features that will be a standard part of python 3, to ease
# adapting our code later on. Importing division is important as it prevents
# one from mistakes where `1/2 = 0` because the old default was integer
# division.
# Now should follow all the imports. Use the following order:
# 1. Standard libraries (sorted)
# 2. Other libraries like numpy, scipy, etc.
# 3. Our libraries
# 4. Local packages for this project.
import numpy as np
# This is my custom utility library. Here we use wrappers to the
# high-performance FFTW libraries.
from mmfutils.performance.fft import fftn, ifftn
# The basis for time evolution in our systems is my pytimeode package. This
# provides a core interface which we must code to in order to use the
# functionality of evolution code.
from pytimeode import interfaces, mixins
# Our helpful tools
from .mixins import StateMixin
# The special variable __all__ is used in python to list which functions,
# classes, etc. should be used by users. It is not strictly needed, but tells
# people what portions of the module they should look at
__all__ = ["Units", "u", "State"]
[docs]
class Units:
"""Units and physical constants.
This class is simply a container for physical constants and units. We have
chosen values here that are relevant for 87Rb which is studied at WSU in
Peter Engels' lab.
"""
[docs]
kg = u / 1.660539040e-27
[docs]
m = 86.909187 * u # 87Rb mass
[docs]
a_B = 5.2917721054980885238e-5 * micron # Bohr radius
# hbar/micron^2/u = 63507.799258914903398 Hz
[docs]
Hz = hbar / micron**2 / u / 63507.799258914903398
[docs]
nK = 0.00206148374391870375 * hbar**2 / micron**2 / u # nK*k_b
# Scattering lengths for 87Rb hyperfine states:
# |1,-1> = |F, m_F>
# |1, 0>
# |2, 0>
# |2, -2>
[docs]
scattering_lengths = {
((1, -1), (1, -1)): 100.40 * a_B,
((2, 0), (2, 0)): 94.57 * a_B,
((2, -2), (2, -2)): 98.98 * a_B,
((1, 0), (1, 0)): 100.86 * a_B,
((1, -1), (2, 0)): 98.13 * a_B, # immiscible
((1, -1), (2, -2)): 98.98 * a_B, # miscible
((1, -1), (1, 0)): 100.41 * a_B, # miscible
}
for _k in list(scattering_lengths):
# Include all symmetric entries
scattering_lengths[_k[::-1]] = scattering_lengths[_k]
[docs]
a = scattering_lengths[(1, -1), (1, -1)]
# Bohr magneton
[docs]
mu_B = 1.399624504 * 2 * np.pi * hbar * MHz / G
# Here is a convenience instance for use in the code.
u = units = Units()
# Decorate the class telling everyone that we provide these interfaces.
@interfaces.implementer(
[
interfaces.IStateForABMEvolvers,
interfaces.IStateForSplitEvolvers,
interfaces.IStateWithNormalize,
]
)
[docs]
class State(StateMixin, mixins.ArrayStateMixin):
"""State class implementing the GPE for a single-component BEC.
To provide time-evolution functionality, the `pytimeode` package requires
that you provide a `State` class that implements the required behaviour of
at least one of the `IStateFor*Evolvers` interfaces. Most of the required
behaviour can be provided by the `mixins.ArrayStateMixin` class which we
inherit from as long as we provide an array called `data` which represents
the wavefunction.
On top of this, we need to provide the following three methods used by the
`pytimeode.evolvers`:
* `compute_dy_dt(dy)`: Required by interface IStateForABMEvolvers
This computes the rhs of the time-dependent differential equation and is
all that is needed by the `pytimeode.evolvers.EvolverABM` evolver. This
is a high-order generic evolver. It requires a rather small step size to
ensure that it behaves well, but once the step size is small enough, it
is converges to high accuracy. For example: there is often a critical
step size above which the evolution will blow up. If you reduce this by
a factor of 2, then you will typically obtain accuracy to 5 digits or so,
at least for modest evolution lengths.
The ABM evolver requires about 10 copies of the state to start up, and 8
copies of the state to be maintained for evolution. This can be a memory
issue if the states are very large.
* `apply_exp_K(dt)`, `apply_exp_V(dt, state)`: Required by interface
`IStateForSplitEvolvers`
These two methods allow you to use `pytimeode.evolvers.EvolverSplit`
based on the Trotter decomposition. This is rather specialized to systems
like the GPE where the exponential of the kinetic and potential terms can
be computed exactly. The ODE solver is not as high order as the ABM
method, but is manifestly unitary. Thus, you can often get away with
very large step sizes and the system will still behave "reasonably"
meaning that the evolution will not be accurate quantitatively, but will
be qualitatively reasonable so you can get an idea of what is happening.
Unfortunately, to achieve quantitative accuracy, you must generally go to
very small step sizes.
Another advantage of the split operator methods is that they require use
only 2 or 3 copies of the state, and so are better for large states that
might be a memory issue
* `normalize()`: Required by interface `IStateWithNormalize`
This method allows the evolvers to continually normalize the state, which
can improve numerical performance, or be of use when finding initial
states.
The reset of the methods in the code play a supporting role to these
methods required by the interfaces and for our convenience (for example,
plotting the state).
Arguments
---------
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.
"""
def __init__(
self,
Nxyz=(2**5, 2**5, 2**5),
Lxyz=(30 * u.micron, 50 * u.micron, 50 * u.micron),
ws=np.array([np.sqrt(8) * 126.0, 126.0, 126.0]) * u.Hz,
g=4 * np.pi * u.hbar**2 * u.a / u.m,
m=u.m,
mu=None,
hbar=u.hbar,
x_TF=None,
cooling_phase=1.0,
symmetric_grid=False,
t=0,
):
[docs]
self.Nxyz = np.asarray(Nxyz)
[docs]
self.Lxyz = np.asarray(Lxyz)
[docs]
self.cooling_phase = cooling_phase
[docs]
self.symmetric_grid = symmetric_grid
# 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 = np.zeros(self.Nxyz, dtype=complex)
# We defer all other calculations to the init() method so that they
# will use the current values of the various parameters. This allows
# the user to change the values of the parameters and they will take
# effect the next time init() is called.
self.init()
# Once the state is initialized, we can set the initial state.
self.set_initial_state()
[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.
"""
# Check that the user did not change the shape of the data, otherwise
# the state will be invalid. As this is a potential user error, we
# raise an exception here with a helpful message
if not np.allclose(self.Nxyz, self.data.shape):
raise ValueError(
"The dimension of the state changed from Nxyz={} to {}.".format(
self.data.shape, self.Nxyz
)
)
# The metric is used wen computing integrals etc.
self.metric = np.prod(np.divide(self.Lxyz, self.Nxyz))
# Define the center of grid to respect the symmetric_grid flag.
xyz0 = self.Lxyz / 2.0
dxyz = self.Lxyz / self.Nxyz
for _n, _d in enumerate(dxyz):
if self.symmetric_grid or self.Nxyz[_n] == 1:
# Special case for N = 1 should also always be centered
xyz0[_n] -= _d / 2.0
# Here we use the meshgrid function to construct the lattice points.
# This is a little complicated, but allows the code to work in D
# dimensions. The indexing and sparse options differ from the MATLAB
# inspired defaults which focus on plotting. Here is the equivalent
# explicit 2D code for comparison:
#
# Nx, Ny = self.Nxyz
# Lx, Ly = self.Lxyz
# x0, y0 = xyz0
# x = np.arange(Nx)*Lx/Nx - x0
# y = np.arange(Ny)*Ly/Ny - y0
# self.xyz = [x[:, None], y[None, :]]
self.xyz = np.meshgrid(
*[
(np.arange(_N)) * _L / _N - _x0
for _N, _L, _x0 in zip(self.Nxyz, self.Lxyz, xyz0)
],
sparse=True,
indexing="ij",
)
if self.symmetric_grid:
for _x in self.xyz:
# It is good practice to put assert statements in your code to
# make sure that everything is working as expected. If they
# become a performance issue, you can disable them after you
# test your code by running python in -O mode.
assert np.allclose(_x.ravel(), -_x.ravel()[::-1])
# Here we generate the momenta. The equivalent explicit 2D code is:
#
# kx = 2*np.pi * np.fft.fftfreq(Nx, Lx/Nx) + k_B[0]
# ky = 2*np.pi * np.fft.fftfreq(Ny, Ly/Ny) + k_B[1]
# self.kxyz = [kx[:, None], ky[None, :]]
self.kxyz = np.meshgrid(
*[
2 * np.pi * np.fft.fftfreq(_N, _L / _N)
for _N, _L in zip(self.Nxyz, self.Lxyz)
],
sparse=True,
indexing="ij",
)
# Here we compute the kinetic energy matrix for use with the FFT.
self.K = sum((self.hbar * _k) ** 2 / 2.0 / self.m for _k in self.kxyz)
# Here we precompute the "phase" factor appearing in the GPE relating
# H(psi) with dpsi/dt. We include the value of hbar here and the
# cooling_phase. A potential optimization here is to allow the state
# to be real if the cooling phase is purely imaginary
self._phase = 1.0 / 1j / self.hbar / self.cooling_phase
# We also record the current particle number so that normalize() can
# restore it if requested during evolution.
self._N = self.get_N()
# This is used to cache the external potential for performance reasons.
self._Vext = None
@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`.
"""
return self.K.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 pre_evolve_hook(self):
"""This method is called by the evolvers at the start of evolution to
ensure that the state is properly initialized.
"""
self.init() # We just defer everything to init()
[docs]
def set_initial_state(self, mu=None, x_TF=None):
"""Set the state using the Thomas Fermi (TF) approximation from either
`mu` or `x_TF` (pick only one or the other).
Arguments
---------
mu : float
Fixed chemical potential.
x_TF : float
Position defining the Thomas Fermi "radius". (The external potential is
evaluated at this position and this is used to set `mu`.)
"""
if mu is not None and x_TF is not None:
raise ValueError(
"Got both mu={} and x_TF={} (specify only one).".format(mu, x_TF)
)
if mu is None and x_TF is None:
mu = self.mu
x_TF = self.x_TF
if mu is None:
if x_TF is None:
x = self.xyz[0].ravel()
x_TF = 0.2 * x[0] + 0.8 * x[-1] # Choose point 80% along
mu = self.get_mu_TF(x_TF=x_TF)
n = self.get_n_TF(mu=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.
n = n + np.zeros(self.shape)
self.data[...] = np.sqrt(n)
self._N = self.get_N()
[docs]
def get_n_TF(self, mu, V_ext=None):
"""Return the Thomas Fermi density profile n from mu."""
zero = np.zeros(self.shape)
if V_ext is None:
V_ext = self.get_Vext()
V = V_ext + zero
# In some applications, the external potential may be complex, so we
# consider only the real part here
return np.maximum(mu - V.real, 0) / self.g
[docs]
def get_mu_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()
V = V_ext + zero
x = self.xyz[0].ravel()
i = np.argmin(abs(x - x_TF))
inds = list(np.asarray(self.Nxyz) // 2)
inds[0] = slice(i - 1, i + 2)
inds = tuple(inds) # Needed for future numpy indexing behaviour
order = min(len(inds) - 1, 2)
V_TF = np.polyval(np.polyfit(x[inds[0]], V[inds], order), x_TF)
mu = V_TF
return mu
[docs]
def get_mu_HO(self, N):
"""Return the chemical potential required to set the particle number in
the Thomas Fermi (TF) approximation for a 3D harmonic oscillator.
"""
assert self.g > 0
w3 = (np.prod(self.ws) ** (1.0 / len(self.ws))) ** 3
mu = ((15.0 * self.g * N * w3 / (16 * np.pi)) ** 2 * self.m**3 / 2) ** (1.0 / 5.0)
return mu
# The use of the property decorator here makes `dim` an attribute rather
# than a method. This means you treat it as if it were a variable
# `self.dim` instead of as a function `self.dim()` but it computes the
# result so it is always up to date.
@property
[docs]
def dim(self):
"""Dimension of the state."""
return len(self.Nxyz)
@property
[docs]
def shape(self):
"""Shape of the state Nxyz."""
return tuple(self.Nxyz)
[docs]
def get_Vext(self):
"""Return the external potential.
Overload this method if you want to change the external potential. If
the potential should be time dependent, use `self.t` which will be
updated by the evolvers.
If a chemical potential `self.mu` is defined, then this is subtracted
from `Vext`. This allows `gpe.minimize` to find states at a constant
chemical potential. Note: for general evolution, it is better not to
set the chemical potential as this is automatically set by
`compute_dy_dt` and will then cause `get_energy` to return the
thermodynamic potential instead.
"""
V_trap = 0.5 * self.m * sum((_w * _x) ** 2 for _w, _x in zip(self.ws, self.xyz))
# Subtract chemical potential if defined
if hasattr(self, "mu") and self.mu:
V_trap -= self.mu
return V_trap
[docs]
def get_density(self):
"""Return the density of the state."""
return abs(self[...]) ** 2
[docs]
def get_N(self):
"""Return the total D-dimensional particle number in the state."""
n = self.get_density()
return self.integrate(n)
[docs]
def integrate(self, a):
"""Return the integral of `a` over the box."""
return self.metric * np.sum(a)
[docs]
def braket(self, a, b):
"""Return the inner product of the functions `a` and `b`.
Note: `a` is conjugated and the metric is applied.
"""
return self.metric * a[...].ravel().conj().dot(b[...].ravel())
[docs]
def get_V(self):
"""Return the complete potential `V` - internal and external.
This function defines the non-linear interaction of the equations. For
the GPE, the energy-density has $gn^2/2$, so we have the derivative
`gn` here.
"""
n = self.get_density()
V_ext = self.get_Vext()
V_int = self.g * n
return V_ext + V_int
[docs]
def fft(self, y):
"""Return the FFT of y over the spatial indices."""
return self._fftn(y)
[docs]
def ifft(self, y):
"""Return the IFFT of y over the spatial indices."""
yt = self._ifftn(y)
return yt
[docs]
def _fftn(self, y):
"""Return the FFT of y over the spatial indices. These versions with
an underscore may be overloaded for performance."""
return fftn(y)
[docs]
def _ifftn(self, y):
"""Return the IFFT of y over the spatial indices. These versions with
an underscore may be overloaded for performance."""
return ifftn(y)
######################################################################
# 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 = self.ifft(self.K * self.fft(y))
Vy = self.get_V() * y
Hy = Ky + Vy
if subtract_mu:
mu = self.braket(y, Hy) / self.braket(y, y)
assert np.allclose(0, mu.imag)
Hy[...] -= mu * y
self._mu = mu
dy[...] = Hy * self._phase
return dy
######################################################################
# Required by interface IStateForSplitEvolvers
[docs]
def apply_exp_K(self, dt):
r"""Apply $e^{-i K dt}$ in place."""
y = self[...]
self[...] = self.ifft(np.exp(self.K * self._phase * dt) * self.fft(y))
[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`.)"""
self *= np.exp(self.get_V() * self._phase * dt)
######################################################################
# Required by interface IStateWithNormalize
[docs]
def normalize(self, s=None):
"""Normalize the state, return the scale factors and number `(s, N)`."""
if s is None:
s = np.sqrt(self._N / self.get_N())
self *= s
assert np.allclose(self._N, self.get_N())
return s, self._N
# End of interface definitions
######################################################################
######################################################################
# The equations of motion follow from minimizing the energy of the state.
# Here we compute the energy so that we can test the equations of motion
# and use a minimizer to find the ground state. These methods are not
# needed for evolution.
[docs]
def get_energy(self):
"""Return the energy of the state. Useful for minimization."""
E = self.integrate(self.get_energy_density())
assert np.allclose(0, E.imag)
return E.real
[docs]
def get_energy_density(self):
"""Return the energy density."""
# Warning: this is not correct. It may not be real until summed. The
# correct energy density requires abs(grad psi)^2
y = self[...]
n = self.get_density()
K = y.conj() * self.ifft(self.K * self.fft(y))
Vint = self.g * n**2 / 2.0
Vext = self.get_Vext() * n
return K + Vint + Vext
######################################################################
# The remaining methods are not needed for evolution or ground state
# preparation, but may be helpful for analysis.
[docs]
def get_H(self, subtract_mu=False):
"""Return the single-particle Hamiltonian for convenience only.
.. note:: May be huge!
This is a square matrix whose sides are of dimension `prod(Nxyz)`
which can be huge. It is not needed for general computations, but
might be useful for analysis such as exploring the BdG in 1D
systems.
"""
N = np.prod(self.Nxyz)
Q = fftn(np.eye(N).reshape(self.shape * 2), axes=range(self.dim)).reshape((N, N))
Qinv = Q.T.conj() / N
mu = 0 if not subtract_mu else self.get_mu()
K = Qinv.dot(self.K.ravel()[:, None] * Q)
V = np.diag(self.get_V().ravel() - mu)
H = K + V
return H
[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 = self.braket(y, Hy) / self.braket(y, y)
assert np.allclose(0, mu.imag)
return mu
[docs]
def get_Hy(self, subtract_mu=False):
"""Return `H(y)` for convenience only."""
dy = self.empty()
self.compute_dy_dt(dy=dy, subtract_mu=subtract_mu)
Hy = dy / self._phase
return Hy
[docs]
def plot(self, log=False, label=None): # pragma: nocover
from matplotlib import pyplot as plt
from matplotlib.gridspec import GridSpec
from mmfutils.plot import imcontourf
n = self.get_density()
if log:
n = np.log10(n)
if self.dim == 1:
x = self.xyz[0] / u.micron
plt.plot(x, n, label=label)
elif self.dim == 2:
from mmfutils import plot as mmfplt
x, y = [_x / u.micron for _x in self.xyz[:2]]
psi = self[...]
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 = [_x / u.micron for _x in self.xyz]
Lx, Ly, Lz = self.Lxyz
nxy = n.sum(axis=2)
nxz = n.sum(axis=1)
nyz = n.sum(axis=0)
gs = GridSpec(1, 3, width_ratios=[Lx, Lx, Ly])
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()
N = self.get_N()
plt.suptitle("t={:.4f}ms, N={:.4f}, E={:.4f}".format(self.t / u.ms, N, E))
class StateTwist(State):
"""Minimal modification of the State class that implements twisted boundary
conditions.
Arguments
---------
twist : (float, float, ...)
Twisted boundary conditions. This is the twist angle from left to right
over the length of the box.
"""
def __init__(self, twist=(0, 0, 0), **kw):
# Since specifying a twist is not a common use case, we are flexible
# here about the size of the twist argument. If it is longer than
# Nxyz, we just take the relevant components.
twist = np.asarray(twist[: len(kw["Nxyz"])])
# Reduce the twists to lie between -pi and pi.
self.twist = (twist + np.pi) % (2 * np.pi) - np.pi
super().__init__(**kw)
def init(self):
State.init(self)
# Twisted boundary conditions are implemented as a shift in the momenta
# by the following "Bloch" momentum:
self.k_B = k_B = self.twist / self.Lxyz
for _n, _kb in enumerate(k_B):
self.kxyz[_n] += _kb
self.twist_phase = np.exp(1j * sum(_kb * _x for _x, _kb in zip(self.xyz, k_B)))
# Here we update the kinetic energy matrix for use with the FFT.
self.K = sum((self.hbar * _k) ** 2 / 2.0 / self.m for _k in self.kxyz)
def fft(self, y):
"""Return the FFT of y over the spatial indices, unapplying the twist."""
return self._fftn(y / self.twist_phase)
def ifft(self, y):
"""Return the IFFT of y over the spatial indices, applying the twist."""
return self.twist_phase * self._ifftn(y)