"""Extended Thomas Fermi model of the Unitary Fermi Gas."""
import time
import numpy as np
# import cupy as cp
import matplotlib.pyplot as plt
# from mmfutils.math.bases import PeriodicBasis
# from mmfutils.math.special import mstep, step
# from mmfutils.performance.fft import fft, ifft, fftn, ifftn, resample
# from mpl_toolkits.mplot3d import Axes3D
# from mmfutils.interface import Interface, Attribute, implementer
from mmfutils.contexts import NoInterrupt
from mmfutils.plot import imcontourf
from pytimeode.evolvers import EvolverABM
import gpe.bec
from . import utils
from .utils import GPUHelper
from .minimize import MinimizeState
import gpe.gpu.bec
# Should give the user the choice - not just force it if available.
gpu = gpe.gpu if gpe.gpu.cupy else None
[docs]
class DimerStateMixinBase(GPUHelper):
"""Dimer state.
This class represents a BEC of dimers.
Attributes
----------
m : float
Dimer mass
Note: Everything in the code, except for the functional, refers
to the dimer properties. Thus m is the dimer mass, get_density()
returns the dimer density, etc.
"""
@property
[docs]
def t_unit(self):
return self.experiment.t_unit
@property
[docs]
def pxyz(self):
"""Return pxyz as list of numpy arrays."""
return list(map(self.asnumpy, self.basis._pxyz))
[docs]
def init(self):
super().init()
self.K_factor = -((self.hbar) ** 2) / 2.0 / self.m
# Rotating frame is now implemented in the basis through this
# factor. See gpe.bec.StateBase.apply_laplacian
if self.experiment.omega == 0:
self.kwz2 = None
else:
self.kwz2 = self.m * self.experiment.omega / self.hbar
# cooling_phase = self.cooling_phase/abs(self.cooling_phase)
# self._phase = 1. / 1j / self.hbar / self.cooling_phase
self._N = self.get_N()
@property
[docs]
def cooling_phase(self):
t_ = self.t / self.t_unit
cooling_phase = self.experiment.get("cooling_phase", t_=t_)
return cooling_phase
@cooling_phase.setter
def cooling_phase(self, cooling_phase):
assert np.allclose(self.cooling_phase, cooling_phase)
@property
[docs]
def _phase(self):
cooling_phase = self.cooling_phase
if self.experiment.normalized_cooling:
_phase = 1.0 / 1j / self.hbar / (cooling_phase / abs(cooling_phase))
else:
_phase = 1.0 / 1j / self.hbar / cooling_phase
return _phase
@_phase.setter
def _phase(self, _phase):
if self.experiment.normalized_cooling:
assert np.allclose(self._phase, _phase)
else:
assert np.allclose(self._phase * abs(self.cooling_phase), _phase)
"""
https://www.python-course.eu/python3_properties.php
@property
def _phase(self):
return self.__phase
@_phase.setter
def _phase(self, _phase):
if self.experiment.normalized_cooling:
self.__phase = _phase / abs(_phase)
else:
self.__phase = _phase
"""
[docs]
def get_V_GPU(self):
"""Return the complete potential `V` - internal and external.
This function defines the non-linear interaction of the equations.
"""
V_int = self.get_Vint_GPU()
V_ext = self.get_Vext_mu_GPU()
return V_int + V_ext
[docs]
def get_energy_density(self, energy=False):
"""Return the energy density.
Arguments
---------
energy : bool
If `True`, then set `self.mu = None` during the computation
so that the actual energy is computed. Otherwise, if
`self.mu` is provided, then this is actually the grand
potential including the `mu*N` subtraction. (See
`get_Vext_GPU` which subtracts `self.mu` if it is not `None`.)
In this energy density function, we calculate the energy density
for the dimers in two ways. First, we do it from equation
of state (please see the method edf below) and then do it
from explicit analytical expression with _factors. Both
gives the same result and it has been tested carefully. Hence,
we have two Vint expressions. We use the one from edf as it is
easier to keep track. Both expressions can be used interchangably.
"""
# Warning: this is not correct. It may not be real until summed. The
# correct energy density requires abs(grad psi)^2
y = self
psi = self.get_psi()
nF = self.get_nF()
nD = nF / 2
Ky = y.copy()
Ky.apply_laplacian(factor=self.K_factor)
K = psi.conj() * Ky.get_psi()
Eint = self.get_Eint(nF=nF)
# _factor = (
# self.experiment.hbar**2 * (3.0) ** (5.0 / 3.0) * (np.pi) ** (4.0 / 3.0)
# ) / (10.0 * (self.m / 2.0))
if energy:
Eext = self.get_Vext() * nD
else:
Eext = self.get_Vext_mu() * nD
return K + Eint + Eext
# TODO: Why is this default energy=False? This is inconsistent with the rest of the
# code!
[docs]
def get_energy(self, energy=False):
"""Return the energy of the state. Useful for minimization.
Arguments
---------
energy : bool
If `True`, then set `self.mu = None` during the computation
so that the actual energy is computed. Otherwise, if
`self.mu` is provided, then this is actually the grand
potential including the `mu*N` subtraction. (See
`get_Vext_GPU` which subtracts `self.mu` if it is not `None`.)
"""
mu_ = self.mu
try:
if energy:
self.mu = None
E = self.integrate(self.get_energy_density(energy=energy))
assert np.allclose(0, E.imag)
E = E.real
finally:
if energy:
self.mu = mu_
return E
[docs]
def get_Eint(self, nF=None):
if nF is None:
nF = self.get_nF()
Eint = self.edf(nF, d=0)
return Eint
[docs]
def get_velocity(self):
"""Returns superfluid velocity field in position space
as an array."""
phase = np.angle(self.get_psi())
vel = (self.hbar / self.m) * np.ascontiguousarray(np.gradient(phase))
return vel
[docs]
def get_current(self):
r"""Calculating the probability current.
Attributes:
-----------
_j : Probability current (shape = 3*N)
$\frac{\hbar}{2*m*i} {\psi^* \nabla \psi - \psi \nabla \psi^*}$.
"""
psi = self.get_psi()
_j = (self.hbar / 2 / self.m / 1.0j) * (
np.conj(psi) * np.asarray(np.gradient(psi))
- psi * np.asarray(np.gradient(np.conj(psi)))
)
return _j
[docs]
def get_omegas(self):
"""Returns irrotational and compressible effective velocity field.
Tsubota 2017, https://arxiv.org/abs/1704.02566, Eq. 79, 80.
"""
j_ = self.get_current()
n_ = self.get_density()
A_ = j_ / np.sqrt(n_)
A_k_ = np.fft.fftn(A_, axes=(1, 2, 3))
kx, ky, kz = map(self.asnumpy, self.basis._pxyz)
zero = 0 * kx + 0 * ky + 0 * kz
k_ = np.asarray([kx + zero, ky + zero, kz + zero])
# norm_k2 = np.linalg.norm(k_, axis=0)**2
k2 = kx**2 + ky**2 + kz**2
# assert np.allclose(k2, norm_k2)
# print(k_.shape, norm_k2)
k_dot_A = (kx * A_k_[0] + ky * A_k_[1] + kz * A_k_[2]) / k2
A_c_k_ = k_ * k_dot_A[np.newaxis, ...]
# A_c_k_ = np.asarray([
# (kx.ravel() * A_k_[0]),
# (ky.ravel() * A_k_[1]),
# (kz.ravel() * A_k_[2])]) / k2 * k_
# A_c_x_ = np.fft.ifftn(A_c_k_)
A_i_k_ = A_k_ - A_c_k_
# A_i_x_ = np.fft.ifftn(A_i_k_)
return (k2, A_i_k_, A_c_k_)
[docs]
def get_Vint_GPU(self, nF=None):
if nF is None:
nF = self.get_nF_GPU()
Vint = 2.0 * self.edf(nF, d=1)
return Vint
[docs]
def get_Vext_GPU(self):
"""Return the external potential.
This delegates to the experiment get_VFext_GPU() and then
includes an appropriate factor of 2 for the dimers.
"""
if not self._Vext or self.t != self._Vext[0] or not self._memoize:
Vext = 2.0 * self.experiment.get_VFext_GPU(state=self)
self._Vext = (self.t, Vext)
return self._Vext[1]
[docs]
def get_Vext_mu_GPU(self):
"""Return Vext with the chemical potential subtracted if initializing.
The chemical potential should be subtracted if initializing or minimizing to get
the initial state. Minimization might be done with imaginary time evolution,
which should be done with negative times.
Thus, we check for `self.initializing`, `self.t < 0`, and make sure that
`self.mu` is valid.
"""
Vext = self.get_Vext_GPU()
# Use getattr rather than hasattr here to allow mu = None.
mu = getattr(self, "mu")
if (self.initializing or self.t < 0) and bool(mu):
# One might consider inplace operations for performance if needed - but be
# careful that the array is not shared. Remember, this only gets called if
# we are initializing, so the minor performance hit is probably justified.
Vext = Vext - mu
return Vext
######################################################################
# Fermionic properties.
#
# These methods refer to fermionic properties for use in the
# equation of state.
@property
[docs]
def mF(self):
"""Fermion mass."""
return self.m / 2
[docs]
def get_nF_GPU(self):
ns = self.get_density_GPU()
nF = 2.0 * ns
return nF
[docs]
def edf(self, nF, d=0):
"""Return the energy density of homogeneous matter.
Arguments
---------
nF : float
Total fermion density = 2*dimer density
d : int
Derivative.
"""
xi = self.experiment.xi
kF = (3 * np.pi**2 * nF) ** (1.0 / 3.0)
pF = self.hbar * kF
eF = pF**2 / 2 / self.mF
e_FG = (3.0 / 5.0) * nF * eF
if d == 0:
return xi * e_FG
elif d == 1:
return xi * eF
else:
raise NotImplementedError("Only d=0 or 1 supported")
#### I don't understand why there are two functions here, why they use
#### experiment.get_VFext rather than state.get_Vext_GPU() and thus how they should
#### use the new initialization scheme with mu. I have tried to fix the second one
#### so that tests pass, but this needs a review.
[docs]
def get_nF_TF(self, muF=None):
"""Invert the edf to get nF(muF).
Returns Thomas-Fermi density of the dimers.
Arguments
---------
muF : float
Effective fermionic chemical potential. Usually muF_0 - V_F(x)
"""
if muF is None:
muF = self.mu / 2
V_F = self.experiment.get_VFext(self)
muF_eff = np.maximum(0, muF - V_F)
nF = (2 * self.mF * muF_eff / self.experiment.xi / self.hbar**2) ** (
3.0 / 2
) / (3 * np.pi**2)
return nF
[docs]
def get_n_TF(self, V_TF=None, V_ext=None):
"""Return the Thomas Fermi density profile n (for dimers) from V_TF.
Arguments
---------
V_TF : float
Value of V(x_TF) where the density should vanish in the TF limit.
"""
if V_TF is None:
V_TF = self.get_V_TF_from_mu(mu=self.mu)
if V_ext is None:
# V_ext = self.experiment.get_VFext(self)
V_ext = self.get_Vext_mu()
muF_eff = np.maximum(0, V_TF - V_ext) / 2
n = (
(2 * self.mF * muF_eff / self.experiment.xi / self.hbar**2) ** (3.0 / 2)
/ (3 * np.pi**2)
) / 2
return n
######################################################################
# Utilities
[docs]
def get_healing_length(self):
"""Coherence length in dimensionless units."""
return self.hbar / np.sqrt(2 * self.m * self.mu)
[docs]
def get_c(self):
"""Speed of sound."""
return np.sqrt(2 * self.mu / 3 / self.m)
@property
[docs]
def txt(self):
"""Text widget."""
if not hasattr(self, "_txt"):
from ipywidgets import Text
self._txt = Text()
return self._txt
[docs]
class DimerStateMixin(DimerStateMixinBase):
"""Class with method to evolve and plot."""
[docs]
def plot(self, log=False, slice=True):
if self.dim < 3:
gpe.bec.State.plot(self)
else:
fig = plt.figure(figsize=(10, 10))
ns = self.get_density()
if log:
ns = np.log10(ns + 1e-12)
x, y, z = self.xyz
args = dict(aspect=1, vmin=0)
if slice:
Nx, Ny, Nz = ns.shape
ax = plt.subplot(221)
imcontourf(x, y, ns[..., Nz // 2], **args)
ax.set(xlabel="x", ylabel="y")
ax = plt.subplot(222)
imcontourf(z, y, ns[Nx // 2, ...].T, **args)
ax.set(xlabel="z", ylabel="y")
ax = plt.subplot(223)
imcontourf(x, z, ns[:, Ny // 2, :], **args)
ax.set(xlabel="x", ylabel="z")
else:
ax = plt.subplot(221)
imcontourf(x, y, ns.sum(axis=2), **args)
ax.set(xlabel="x", ylabel="y")
ax = plt.subplot(222)
imcontourf(z, y, ns.sum(axis=0).T, **args)
ax.set(xlabel="z", ylabel="y")
ax = plt.subplot(223)
imcontourf(x, z, ns.sum(axis=1), **args)
ax.set(xlabel="x", ylabel="z")
N_f = 2 * self.get_N()
PV = -self.get_energy(energy=False)
E = self.get_energy(energy=True)
np.allclose(self.mu, 2 * self.experiment.xi * self.experiment.eF)
##### This is wrong! Must recompute eF given density...
# alpha = self.hbar**2 * (3*np.pi**2)**(2./3.) / self.m
# eF_2 = self.integrate(self.edf(2*self.get_density(), d=1))/self.experiment.xi
# eF_ = alpha * (self.get_density().max())**(2./3.)
# print(alpha, eF_, eF_2, self.experiment.eF)
E_FG = 3 / 5 * self.experiment.eF * N_f
PV_FG = 2 / 3 * E_FG
t_ = self.t / self.t_unit
plt.suptitle(
f"t_={t_:.2f}, N_F={N_f:.2f}, "
+ f"|E/E_FG|={E/E_FG:.4f}, "
+ f"|P/P_FG|={PV/PV_FG:.4f}"
)
[docs]
def plot_pv(self, n_nmax=0.5):
# Defer importing pyvista so that we can use the remaining code
# without it since it is a large dependency.
import pyvista as pv
x, y, z = np.meshgrid(*self.xyz)
grid = pv.StructuredGrid(x, y, z)
n = self.get_density()
grid["vol"] = n.flatten()
n_max = n.max()
print(n.min() / n.max())
contours = grid.contour([n_nmax * n_max])
pv.set_plot_theme("document")
pl = pv.PlotterITK()
pl.add_mesh(contours, scalars=contours.points[:, 2]) # , show_scalar_bar=False)
return pl.show()
[docs]
def evolve_to(self, t, dt_t_scale=0.1):
state = self
dt = dt_t_scale * state.t_scale
steps = max(2, int(np.ceil(t / dt)))
print("Evolving for {steps}steps")
dt = t / steps
ev = EvolverABM(state, dt=dt)
ev.evolve(steps)
return ev.get_y()
[docs]
def evolve(
self,
steps=500,
t__max=10000,
dt_t_scale=0.1,
display=True,
pv=False,
hist=False,
):
"""Time evolves a state.
Arguments:
----------
pv: bool
If True, does PyVista 3d plotting.
hist: bool
If True, returns a list of evolved states at different times.
t__max : float
Evolve up to this time (units of t_unit).
"""
if display:
from IPython.display import display, clear_output
state = self.copy()
# Set dt and make sure we end at t_max
t_max = t__max * self.t_unit
dt = dt_t_scale * state.t_scale
Nsteps = int(np.ceil(t_max / dt))
dt = t_max / Nsteps
history = [state]
ev = EvolverABM(state, dt=dt)
NoInterrupt.unregister()
with NoInterrupt() as interrupted:
while Nsteps > 0 and not interrupted:
tic = time.time()
if Nsteps <= steps:
steps_ = Nsteps
else:
# Ensure that we always leave at least 2 steps
steps_ = min(steps, Nsteps - 2)
ev.evolve(steps_)
Nsteps -= steps_
history.append(ev.get_y())
self.txt.value = ", ".join(
[
f"{(time.time() - tic):.4f}s/{steps}step",
f"E={ev.y.get_energy():.4f}",
f"N={ev.y.get_N():.4f}",
f"c={ev.y.get_c():.4f}",
]
)
if display:
plt.clf()
if pv:
ev.y.plot_pv()
else:
ev.y.plot()
display(plt.gcf())
clear_output(wait=True)
plt.close("all")
if hist:
return history
return ev.get_y()
[docs]
class DimerStateCPU(DimerStateMixin, gpe.bec.StateBase):
def __init__(self, experiment, **kw):
[docs]
self.experiment = experiment
kw.setdefault("cooling_phase", self.cooling_phase)
super().__init__(**kw)
[docs]
class DimerState(DimerStateMixin, gpu.bec.StateBase if gpu else gpe.bec.StateBase):
def __init__(self, experiment, **kw):
[docs]
self.experiment = experiment
kw.setdefault("cooling_phase", self.cooling_phase)
super().__init__(**kw)
[docs]
class Experiment(utils.ExperimentBase, GPUHelper):
"""Experimental setup.
Attributes
----------
xi : float
Bertsch parameter.
Everything here refers to fermions. External potentials
are defined for fermions, etc.
"""
[docs]
mF = 1.0 # Fermion mass
# xi = 0.3705
[docs]
normalized_cooling = True # Flag to normalize the cooling phase
# eta = 0.651 # equation (6) of Phys. Rev. A 86, 053603 (2012),
# from SLDA fitting
[docs]
eta = 0.5 # Gabriel uses this value (Aug, 4, 2020),
# Book chapter equation (93)
[docs]
xi = 0.40 # Book chapter equation (93)
[docs]
def init(self):
self.eF = (self.hbar * self.kF) ** 2 / 2 / self.mF
self.muF = self.xi * self.eF
self.dxyz = np.array((self.dx_kF / self.kF,) * self.dim)
self.Lxyz = self.dxyz * self.Nxyz
self.omega = self.omega_eF * self.eF
super().init()
@property
[docs]
def dim(self):
return len(self.Nxyz)
@property
[docs]
def t_unit(self):
return self.hbar / self.eF
[docs]
def get_state(
self,
t_=0,
minimize=False,
initialize=False, # utils.py line 1521 initilize (dummy)
_E_tol=1e-12,
_psi_tol=1e-12,
):
"""Quickly return a valid `State` object."""
"""This method prepares a 3D state, to use for minimization."""
t = t_ * self.t_unit
args = dict(Nxyz=self.Nxyz, Lxyz=self.Lxyz, mu=2 * self.muF, m=2 * self.mF)
state = self.State(experiment=self, t=t, **args)
if minimize:
minimizer = MinimizeState(state, fix_N=self.fix_N)
state = minimizer.minimize(E_tol=_E_tol, psi_tol=_psi_tol)
return state
[docs]
def get_initial_state(
self, psi0=None, _E_tol=1e-12, _psi_tol=1e-12, callback=None, **kw
):
"""Return the valid `t=0` state to initialize the simulations."""
# t = t_*self.t_unit
state = self.get_state(**kw)
if psi0 is not None:
state.set_psi(psi0)
minimizer = MinimizeState(state, fix_N=self.fix_N)
psi0 = minimizer.minimize(E_tol=_E_tol, psi_tol=_psi_tol, callback=callback)
state.set_psi(psi0)
return state
[docs]
def get_initialized_state(self, state):
"""Return a valid state initialized from `state`.
This is used in chained simulations where a specified state of one
simulation is used to initialize a state for further use. For example,
for expansion."""
state_ = self.get_state()
state_[...] = state
return state_
# End of methods required by IExperiment
######################################################################
[docs]
def get_VFext_GPU(self, state):
"""Return the external potential for fermions."""
return 0.0