"""Demonstration of Travelling Wave solutions of the GPE.
For the plain GPE these have analytic solutions as discussed in the
`Docs/Travelling Wave.ipynb` notebook.
"""
import gpe.bec
import gpe.minimize
from gpe.bec import StateTwist_x
from gpe.exact_solutions import TravellingWaves
import numpy as np
[docs]
class StateTravellingWave(StateTwist_x):
"""Class for solving for traveling waves.
Here we hold one point of the wavefunction fixed. To do this, we
need to hijack some of the functions from the original State
class. For example, the total particle number and braket
functions exclude this point so that particle number projection
can be properly implemented.
To do this, we set the metric to exclude this point, and overload
some functions to ensure that a constant value is maintained.
"""
def __init__(
self,
Nx=128,
Lx=10.0,
psi_0=1.0,
ind=None,
n_avg=1.0,
v_p=0.0,
twist=0.0,
m=1.0,
hbar=1.0,
g=1.0,
):
"""
Arguments
---------
v_p : float
Phase velocity. This is implemented as a boost.
twist : float
Additional twist (radians). For example, to form a dark
soliton, one would set `twist=np.pi` with `v_p=0`.
psi_0 : float
The wavefunction value at this point is held fixed during
minimization.
ind : int, None
Index where `psi[ind] = psi_0` is held fixed. (Default is the
midpoint of the grid.)
n_avg : float
Average density in the box. Used to fix the total particle number.
"""
if ind is None:
ind = Nx // 2
StateTwist_x.__init__(
self, Nxyz=(Nx,), Lxyz=(Lx,), m=m, hbar=hbar, g=g, twist=twist, v_x=v_p
)
self.set_psi(np.sqrt((Nx * self.n_avg - abs(self.psi_0) ** 2) / (Nx - 1)))
self[self.ind] = psi_0
assert np.allclose(self.get_density().mean(), self.n_avg)
[docs]
def init(self):
# Exclude fixed point for integration and brakets.
self._metric = np.ones(self.shape) * self.basis.metric
self._metric[self.ind] = 0
StateTwist_x.init(self)
@property
[docs]
def metric(self):
return self._metric
[docs]
def exact_psi(self):
"""Return the analytic soliton"""
n1 = abs(self.psi_0) ** 2
tw = TravellingWaves(Nx=self.Nxyz[0], Lx=self.Lxyz[0])
v = abs(self.psi_0)
u = np.sqrt(self.mu - v**2)
return 1j * v + u * np.tanh(u * self.xyz[0])
######################################################################
# Overload these functions to ensure that psi[ind] = psi_0
[docs]
def compute_dy_dt(self, dy, subtract_mu=False):
dy = StateTwist_x.compute_dy_dt(self, dy=dy, subtract_mu=subtract_mu)
dy[self.ind] = 0
return dy
[docs]
def scale(self, f):
"""Perform `self *= f` as efficiently as possible."""
StateTwist_x.scale(self, f)
self.data[self.ind] = self.psi_0
[docs]
def axpy(self, x, a=1):
"""Perform `self += a*x` as efficiently as possible."""
StateTwist_x.axpy(self, x=x, a=a)
self.data[self.ind] = self.psi_0
[docs]
def get_energy(self):
"""Return the energy of the state. Useful for minimization."""
# This needs to be real...
E = self.basis.metric * sum(self.get_energy_density())
assert np.allclose(0, E.imag)
return E.real
[docs]
def plot(self, log=False, label=None): # pragma: nocover
from matplotlib import pyplot as plt
n = self.get_density()
psi = self.get_psi()
theta = np.angle(psi)
if log:
n = np.log10(n)
x = self.xyz[0] / u.micron
ax = plt.gca()
plt.plot(x, n, label=label)
plt.xlabel("x [micron]")
plt.ylabel("n")
plt.twinx()
plt.plot(x, theta, "y:")
plt.ylim(-np.pi, np.pi)
plt.ylabel("theta")
plt.sca(ax)
E = self.get_energy()
N = self.get_N()
t = np.ravel(self.t)[0]
plt.suptitle("t={:.4f}ms, N={:.4f}, E={:.4f}".format(t / u.ms, N, E))
[docs]
class MinimizeStateTravellingWave(gpe.minimize.MinimizeState):
def __init__(self, state, fix_N=True, **kw):
[docs]
self.psi_0 = state.psi_0
gpe.minimize.MinimizeState.__init__(self, state, fix_N=fix_N, **kw)
[docs]
def pack(self, psi):
psi[self.ind] = self.psi_0
fact = 1 if self.real else 2
x = gpe.minimize.MinimizeState.pack(self, psi)
return np.concatenate([x[: self.ind * fact], x[(self.ind + 1) * fact :]])
[docs]
def unpack(self, x, state=None):
fact = 1 if self.real else 2
x = np.concatenate([x[: self.ind * fact], [0] * fact, x[self.ind * fact :]])
state = gpe.minimize.MinimizeState.unpack(self, x, state=state)
assert np.allclose(state[self.ind], 0)
state[self.ind] = self.psi_0
return state