"""This file contains a set of exact solutions useful for testing."""
from __future__ import absolute_import, division, print_function
import numpy as np
import scipy.integrate
import scipy.optimize
from scipy.special import ellipj, ellipk
import scipy as sp
from . import bec
from . import bec2
[docs]
class HarmonicOscillator(bec.StateBase):
"""Modified 1D harmonic oscillator potential with analytic solution.
Here we add a piece to the harmonic oscillator potential to ensure that a
Gaussian is an eigenstate, even in the presence of interactions.
"""
def __init__(self, Nx=46, Lx=17.0, sigma=1.0, g=1.0, n0=1.0, **kw):
m = 1.0
[docs]
self.w = bec.u.hbar / m / sigma**2 # Trap frequency to give specified sigma
super().__init__(Nxyz=(Nx,), Lxyz=(Lx,), m=m, g=g, **kw)
self.set_psi(self.psi_exact)
self.pre_evolve_hook()
[docs]
def get_psi_exact(self, r):
return np.sqrt(self.n0) * np.exp(-((r / self.sigma) ** 2) / 2)
[docs]
def get_Vext_r(self, r):
k = 1.0 / 2.0 / self.sigma**2
psi_0 = self.get_psi_exact(r)
n_0 = abs(psi_0) ** 2
return self.hbar**2 / 2.0 / self.m * (4 * (k * r) ** 2 - 2 * k) - self.g * n_0
@property
[docs]
def psi_exact(self):
x = self.get_xyz()[0]
return self.get_psi_exact(r=abs(x))
[docs]
def get_Vext_GPU(self):
x = self.get_xyz_GPU()[0]
return self.get_Vext_r(r=abs(x))
[docs]
class HarmonicOscillators(bec.StateBase):
"""Modified 2D harmonic oscillator potential with analytic solution.
Here we add a piece to the harmonic oscillator potential to ensure that a
Gaussian is an eigenstate, even in the presence of interactions.
"""
def __init__(
self, Nxyz=(46, 46), Lxyz=(17.0, 17.0), sigmas=(1.0, 1.0), g=1.0, n0=1.0, **kw
):
m = 1.0
[docs]
self.ws = bec.u.hbar / m / np.asarray(sigmas) ** 2
super().__init__(Nxyz=Nxyz, Lxyz=Lxyz, m=m, g=g, **kw)
self.set_psi(self.psi_exact)
self.pre_evolve_hook()
[docs]
def get_psi_exact(self, xyz=None):
if xyz is None:
xyz = self.get_xyz()
r2_sigma = sum((x / sigma) ** 2 for (x, sigma) in zip(xyz, self.sigmas))
return np.sqrt(self.n0) * np.exp(-r2_sigma / 2)
[docs]
def get_Vext_GPU(self):
xyz = self.get_xyz_GPU()
psi0 = self.get_psi_exact(xyz=xyz)
return (
sum(
self.m * (w * x) ** 2 - self.hbar**2 / self.m / s**2
for w, s, x in zip(self.ws, self.sigmas, xyz)
)
- self.g * abs(psi0) ** 2
)
@property
[docs]
def psi_exact(self):
return self.get_psi_exact()
[docs]
class BrightSoliton(bec.StateTwist_x):
"""Analytic moving bright soliton (for attractive interactions)."""
def __init__(self, Nx=2**4 * 3**3, Lx=58.0, sigma=1.0, g=None, n0=1.0, v=0.0, **kw):
m = hbar = 1.0
if g is None:
g = -(hbar**2) / m / sigma**2
super().__init__(Nxyz=(Nx,), Lxyz=(Lx,), m=m, g=g, hbar=hbar, **kw)
self.set_psi(self.psi0)
self.pre_evolve_hook()
@property
[docs]
def x(self):
return self.get_xyz()[0]
@property
[docs]
def psi0(self):
x, L = self.get_xyz()[0], self.basis.Lxyz[0]
eta = 1.0 / self.sigma
v, t = self.v, self.t
x = (x - v * t + L / 2.0) % L - L / 2.0 + v * t
mu = -((self.hbar * eta) ** 2) / (2 * self.m)
arg1 = eta * (x - v * t)
arg2 = -v * x + (mu + self.m * v**2 / 2.0) * t
return np.sqrt(self.n0) / np.cosh(arg1) * np.exp(arg2 / 1j)
[docs]
def sn(u, m):
"""Jacobi elliptic function sn."""
return ellipj(u, m)[0]
[docs]
def K(m):
"""Jacobi elliptic integral"""
return ellipk(m)
[docs]
class TravellingWaves(bec.StateTwist_x):
"""Analytic traveling wave solutions for the GPE without a
potential.
The solutions are characterized by the following parameters:
n0, n1 : float
Minimum and maximum density of the soliton. The amplitude is
a=n1-n0.
Lx : float
The period of the wave, and the length of our box. This is
specified in terms of Lx = 2lK(m_), where m_=al^2.
v_p : float
Phase velocity.
v_x : float
Velocity of the frame. If None, this is set to be the phase
velocity so that the solution should be stationary.
twist : float
Twisted boundary condition. If None, this is computed by
integrating the continuity equation.
"""
def __init__(
self,
Nx=64,
Lx=10.0,
n0=0.1,
n1=1.0,
m=1.0,
g=1.0,
hbar=1.0,
v_p=0.0,
v_x=None,
twist=None,
**kw,
):
self.l_unit = l_unit = m * g / hbar**2
self.t_unit = t_unit = m**3 * g**2 / hbar**5
self.n_unit = n_unit = 1.0 / l_unit**3
[docs]
self.v_unit = l_unit / t_unit
self.n0, self.n1 = n0, n1
# Internally we store dimensionless parameters starting with
# an underscore.
self._a = _a = (n1 - n0) / n_unit
[docs]
self._m = self.get_m(_a=_a, _L=Lx / l_unit)
self._l = _l = np.sqrt(self._m / self._a)
assert np.allclose(_l, Lx / l_unit / (2 * K(self._m)))
[docs]
self._C = -np.sqrt(n0 * n1 * (_l ** (-2) + n0 / n_unit)) / n_unit
[docs]
self.mu = g * (n0 + (n1 + (n1 - n0) / m) / 2.0) - m * v_p**2 / 2.0
if twist is None:
x = np.linspace(-Lx / 2.0, Lx / 2.0, 10000)
_psi, twist = self.psi_exact(x=x, Lx=Lx)
if v_x is None:
v_x = v_p
super().__init__(
Nxyz=(Nx,), Lxyz=(Lx,), ws=(0,), m=m, g=g, twist=twist, v_x=v_x, **kw
)
self.set_psi(self.psi_exact()[0])
[docs]
def n_exact(self, x=None):
if x is None:
x = self.get_xyz()[0]
l = self._l * self.l_unit # noqa: E741
return (self._n0 + self._a * sn(x / l, self._m) ** 2) * self.n_unit
[docs]
def psi_exact(self, x=None, Lx=None):
"""Return the exact solution.
Arguments
---------
x, Lx : array, float
Abscissa and box size - used by the constructor for getting
the exact solution before the state has been properly initialized.
"""
if x is None:
x = self.get_xyz()[0]
if Lx is None:
Lx = self.Lx
_a = self.v_p / self.v_unit
C = self._C * self.n_unit
def dtheta(x):
return _a + C / self.n_exact(x)
theta = [0]
for _n in range(1, len(x)):
theta.append(theta[-1] + scipy.integrate.quad(dtheta, x[_n - 1], x[_n])[0])
theta = np.array(theta)
theta -= theta[len(x) // 2]
twist = scipy.integrate.quad(dtheta, 0, Lx)[0]
psi = np.sqrt(self.n_exact(x)) * np.exp(1j * theta)
return psi, twist
@staticmethod
[docs]
def get_m(_a, _L):
"""Return the solution to _L = 2*sqrt(m/_a)*K(m)."""
def f(m):
return 2 * np.sqrt(m) * K(m) - _L * np.sqrt(_a)
m1 = 0.0
f1 = f(0.0)
while f1 < 0:
m1 = (m1 + 1) / 2.0
f1 = f(m1)
return sp.optimize.brentq(f, 0, m1)
[docs]
class DarkBrightSolitons(bec2.State):
def __init__(self, kappa=1.0, bs=[-1.0, -2.0, -3.0], lambdas=(0.0, 0.0), **kw):
b11, b22, b12 = bs
[docs]
self.Cs = np.divide((b22 - kappa * b12, b12 - kappa * b11), b12**2 - b12 * b22)
super().__init__(ms=(1.0, 1.0 / self.kappa), ws=self.lambdas, gs=bs, **kw)
[docs]
def get_initial_state(self, x0=0.0, eta=1.0, v=1.0, t=0.0, mu=1.0):
x = self.get_xyz()[0]
L = self.Lxyz[0]
kappa = self.kappa
arg = eta * (x - v * t)
b11, b22, b12 = self.bs
C1, C2 = self.Cs
C = np.sqrt(abs(self.Cs))[self.bcast]
psi_B = eta / np.cosh(arg)
f1 = -(eta**2) - v**2 * (-b11 * C1 + b12 * C2 / mu**2) + 0 * x
f2 = -mu * eta**2 - v**2 * (b12 * C1 + b22 * C2 / mu**2) + 0 * x
twist = [0, 0]
if C1 < 0:
if C2 < 0:
print("DB")
phi = np.array([f2 * t + 0 * x, v * x + f1 * t])
psi = C * [1j * v + eta * np.tanh(arg), psi_B]
twist = [np.pi, v * L - np.pi]
else:
print("DD")
phi = np.array([f1 * t, f2 * t])
psi = C * [
1j * v + eta * np.tanh(arg),
1j * v / kappa + eta * np.tanh(arg),
]
else:
if C2 < 0:
print("BB")
phi = np.array(
[
v * x + (eta**2 - v**2) * t / 2.0,
v * x / kappa + (kappa * eta**2 - v**2 / kappa) * t / 2.0,
]
)
psi = C * [psi_B, psi_B]
else:
print("BD")
phi = np.array([v * x + f1 * t, f2 * t])
psi = C * [psi_B, 1j * v / kappa + eta * np.tanh(arg)]
theta = np.pi
twist = [0, theta]
self.twist = np.asarray(twist)
return psi * np.exp(1j * phi)