---
jupytext:
  formats: md:myst,ipynb
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.13
    jupytext_version: 1.16.4
kernelspec:
  display_name: Python 3 (ipykernel)
  language: python
  name: python3
---

```{code-cell}
import mmf_setup
```

```{code-cell}
mmf_setup.nbinit()
```

# Two-Component BECs (BEC2)

Here and in the file [`bec2.py`](../src/gpe/bec2.py) we present code that solves the GPE
for a two-component system.  The energy density and equations of motion are:
\begin{gather*}
  \mathcal{E}[\Psi] = \Psi^\dagger\begin{pmatrix}
    \frac{-\hbar^2\nabla^2}{2m} + V_{a} & V_{ab}\\
    V_{ab}^\dagger & \frac{-\hbar^2\nabla^2}{2m} + V_{b}\\
  \end{pmatrix}\Psi
  +
  \mathcal{E}(n_a, n_b),
  \qquad
  \Psi = \begin{pmatrix}
    \psi_a\\
    \psi_b
  \end{pmatrix}\\
  -\I\hbar \pdiff{}{t}\Psi =
  \begin{pmatrix}
    \frac{-\hbar^2\nabla^2}{2m} + V_{a}
    + \mathcal{E}_{,n_a}(n_a, n_b)  & V_{ab}\\
    V_{ab}^\dagger & \frac{-\hbar^2\nabla^2}{2m} + V_{b}
    + \mathcal{E}_{,n_b}(n_a, n_b)\\
  \end{pmatrix}\Psi.
\end{gather*}
For the standard GPE, the equation of state is:
\begin{gather*}
  \mathcal{E}(n_a, n_b) = \frac{g_{aa}}{2} n_a^2 + \frac{g_{aa}}{2} n_a^2 + g_{ab} n_an_b.
\end{gather*}

This structure is supported in an augmented `state.data` array of shape `(2, N_x, N_y, N_z)` where the first index refers to indices `a` or `b`.  Note now that functions like `integrate()`, `get_N()`, and `braket()` now return arrays of shape `(2,)` instead of scalars.


## Twisted Boundary Conditions

Instead of a periodic box, one can work with twisted boundary conditions whereby

\begin{gather*}
  \psi(x+L) = e^{\I\phi}\psi(x).
\end{gather*}
To implement such boundary conditions, we use the periodic function
\begin{gather*}
  \tilde{\psi}(x) = e^{-\I\phi x/L}\psi(x)
\end{gather*}
when computing the FFT.  To compute derivatives, we must thus first apply this shift to
the wavefunction to get a periodic function, then take the FFT, then use momenta that
have been shifted by $-\phi/L$ etc.

```{code-cell}
%pylab inline --no-import-all
from IPython.display import display, clear_output
from gpe import bec2
```

```{code-cell}
reload(bec2)
from gpe.bec2 import State, u
```

```{code-cell}
s = State(Nxyz=(128,), Lxyz=(40 * u.micron,), N=1e5)
s.plot()
```

### Testing

+++

To test the code, we will set $g=0$ and use the exact solution for the Harmonic Oscillator:

\begin{gather*}
  \psi(x) \propto e^{-(x/a)^2/2}, \qquad
  a^2 = \frac{\hbar}{m\omega}
\end{gather*}

```{code-cell}
def get_err(N=128, L=24 * u.micron):
    s = State(Nxyz=(N,), Lxyz=(L,), N=1e5)
    s.gs = 0, 0, 0
    a = np.sqrt(u.hbar / u.m / s.ws[0])
    x = s.xyz[0]
    psi_0 = np.exp(-((x / a) ** 2) / 2.0)
    s[...] = psi_0[None, :]
    s.normalize()
    dy = s.empty()
    s.compute_dy_dt(dy=dy, subtract_mu=True)
    return abs(dy[...]).max()


Ns = 2 ** np.arange(2, 8)
errs = map(get_err, Ns)
plt.semilogy(Ns, errs, "-+")
```

Why are $L=23$microns and $N=2^6$ optimal?

```{code-cell}
s = State(Nxyz=(46,), Lxyz=(23 * u.micron,))
a = np.sqrt(u.hbar / u.m / s.ws[0])
L, N = s.Lxyz[0], s.Nxyz[0]
k_max = np.pi * (N - 2) / L  # For Khalid...
print k_max, np.max(s.kxyz[0])
print np.exp(-((L / 2 / a) ** 2) / 2)  # Wavefunction drops by factor of macheps
```

```{code-cell}
psi_0 = s.xyz[0] * np.exp(-((s.xyz[0] / a) ** 2) / 2)
plt.semilogy(
    np.fft.fftshift(s.kxyz[0]).ravel(),
    np.fft.fftshift(abs(np.fft.fft(psi_0))).ravel(),
    "-+",
)
```

So we see that for the ground state $k$ needs to go up to $6$.

### Exact Solution with Interactions

Consider now a state with $\psi_a = e^{-(x/a)^2/2}$ and $\psi_b = e^{-(x/b)^2/2}$.

\begin{gather*}
  \psi_a''(x)
\end{gather*}

```{code-cell}
%pylab inline --no-import-all
from IPython.display import display, clear_output
from gpe import bec2

reload(bec2)
from gpe.bec2 import State, u

s = State(Nxyz=(64,), Lxyz=(23 * u.micron,), N=1e5)
a = np.sqrt(u.hbar / u.m / s.ws[0])
x = s.xyz[0]
psi_0 = np.exp(-((x / a) ** 2) / 2)


class State1(State):
    def __init__(self, *v, **kw):
        State.__init__(self, *v, **kw)
        a = np.sqrt(u.hbar / u.m / self.ws[0])
        x = self.xyz[0]
        k = 1.0 / 2.0 / a ** 2
        psi_0 = 4.0 * np.exp(-((x / a) ** 2) / 2)
        n_0 = abs(psi_0) ** 2
        g_aa, g_bb, g_ab = self.gs
        self._V_ext = (
            u.hbar ** 2 / 2.0 / u.m * (4 * (k * x) ** 2 - 2 * k) - g_aa * n_0,
            u.hbar ** 2 / 2.0 / u.m * (4 * (k * x) ** 2 - 2 * k) - g_bb * n_0,
            0 * n_0,
        )
        self.data[...] = psi_0
        self.get_Vext = lambda: self._V_ext
        self.pre_evolve_hook()


s = State1(Nxyz=(64,), Lxyz=(23 * u.micron,))
s.plot()
plt.plot(x, s.get_Vext()[0])
dy = s.empty()
s.compute_dy_dt(dy=dy, subtract_mu=False)
abs(dy[...]).max()
```

```{code-cell}
from mmfutils.contexts import NoInterrupt
from pytimeode.evolvers import EvolverSplit, EvolverABM
from IPython.display import display, clear_output

s = State1(Nxyz=(64 * 4,), Lxyz=(23 * u.micron,))
assert np.allclose(s._N, s.get_N())

s[...] = 1.0
s.normalize()
s.cooling_phase = 1j

E_max = u.hbar ** 2 * np.abs(s.kxyz).max() ** 2 / 2.0 / u.m

# e = EvolverSplit(s, dt=0.01*u.hbar/E_max, normalize=True)
e = EvolverABM(s, dt=0.1 * u.hbar / E_max, normalize=True)

with NoInterrupt(ignore=True) as interrupted:
    while e.y.t < 4 * u.ms and not interrupted:
        e.evolve(100)
        plt.clf()
        e.y.plot()
        display(plt.gcf())
        clear_output(wait=True)
```

```{code-cell}
from mmfutils.contexts import NoInterrupt
from pytimeode.evolvers import EvolverSplit, EvolverABM
from IPython.display import display, clear_output

s = State1(Nxyz=(64 * 4,), Lxyz=(23 * u.micron,))
s *= np.sign(s.xyz[0] - 0.5)
s.cooling_phase = 1 + 0.01j

E_max = u.hbar ** 2 * np.abs(s.kxyz).max() ** 2 / 2.0 / u.m
# e = EvolverSplit(s, dt=0.01*u.hbar/E_max, normalize=True)
e = EvolverABM(s, dt=0.5 * u.hbar / E_max, normalize=True)

with NoInterrupt(ignore=True) as interrupted:
    while e.y.t < 40 * u.ms and not interrupted:
        e.evolve(100)
        plt.clf()
        e.y.plot()
        display(plt.gcf())
        clear_output(wait=True)
```

## 2D Scissor Modes (incomplete)

+++

We base the physical parameters on a system like [Marago:2001] so we can explore damping of the normal modes of the system.  They have trapping frequencies of $\omega_y=\omega_z = 128$Hz and $\omega_x = \sqrt{8}\omega_y$ and $N = 2\times 10^4$ particles.  In the TF approximation at $T=0$ this corresponds to $\mu = $

[Marago:2001]: http://dx.doi.org/10.1103/PhysRevLett.86.3938 (Onofrio Marag\`o, Gerald Hechenblaikner, Eleanor Hodby, and Christopher Foot, "Temperature Dependence of Damping and Frequency Shifts of the Scissors Mode of a Trapped Bose-Einstein Condensate", Phys. Rev. Lett. 86, 3938--3941 (2001) )

```{code-cell}
%pylab inline --no-import-all
from IPython.display import display, clear_output
```

```{code-cell}
from gpe import bec

reload(bec)
from pytimeode.evolvers import EvolverABM
from mmfutils.contexts import NoInterrupt
from gpe.bec import State, u

s = State()
s.cooling_phase = 1j
s.t = -100 * u.ms
e = EvolverABM(s, dt=0.001)
with NoInterrupt(ignore=True) as interrupted:
    while not interrupted:
        e.evolve(500)
        plt.clf()
        e.y.plot()
        display(plt.gcf())
        clear_output(wait=True)
```

```{code-cell}
s = e.get_y()
n = s.get_density()
x, y, z = s.xyz
plt.plot(x.ravel(), n.sum(axis=-1).sum(axis=-1))
s.mu, s._mu
```

## Bright Solitons

+++

There is nice integrable 1D model that is useful for testing code in the case of attractive interactions (the so-called focusing NLSEQ):

\begin{gather*}
  \I\hbar \partial_t \Psi = - \frac{\hbar^2\nabla^2}{2m}\Psi
  + \begin{pmatrix}
    -gn_{+} & \Omega(t)\\
    \Omega(t) & -gn_{+}
  \end{pmatrix}\cdot\Psi
\end{gather*}

where $n_{+} = n_a + n_b$.  This corresponds to a coupled set of GPEs with equal attractive interactions $g_{aa} = g_{bb} = g_{ab} = -g$.


\begin{gather*}
  \DeclareMathOperator{\sech}{sech}
  \Psi(x, t) = \frac{1}{2}\begin{pmatrix}
      \cos\theta e^{-\I\Gamma(t)} + \sin\theta e^{\I\Gamma(t)}\\
      \cos\theta e^{-\I\Gamma(t)} - \sin\theta e^{\I\Gamma(t)}
  \end{pmatrix}e^{\I g b^2 t/4} b\sech\frac{b \sqrt{gm} x}{\hbar}, \qquad
  -\frac{\hbar^2 \psi''(x)}{2m} - g\abs{\psi}^2\psi = -\frac{b^2 g}{2}\psi, \\
  \Gamma(t) = \int_{0}^{t}\Omega(t')\d{t'}
\end{gather*}

\begin{gather*}
  \Psi(x, t) = \left(\frac{\cos\theta e^{-\I\Gamma(t)}}{2}
  \begin{pmatrix}1\\1\end{pmatrix}
  +
  \frac{\sin\theta e^{\I\Gamma(t)}}{2}
  \begin{pmatrix}1\\ -1\end{pmatrix} \right)
  e^{\I g b^2 t/4} b\sech(b x\sqrt{gm/\hbar^2}), \qquad
  -\frac{\hbar^2 \psi''(x)}{2m} - g\abs{\psi}^2\psi = -\frac{b^2 g}{2}\psi, \\
  \Gamma(t) = \int_{0}^{t}\Omega(t')\d{t'}
\end{gather*}

+++

We have the following dimensions:

\begin{gather*}
  [\hbar] = \frac{M D^2}{T}, \qquad
  [2m] = M, \qquad
  [gn] = [E] = \frac{M D^2}{T^2}, \qquad
  [g] = [VE] = \frac{M D^5}{T^2}, \qquad
  [\Omega] = [E] = \frac{M D^2}{T^2}\\
  \left[\frac{2mg}{\hbar^2}\right] = D, \qquad
  \left[\frac{2m\Omega}{\hbar^2}\right] = \frac{1}{D^2}, \qquad
 \left[\frac{t\hbar}{2m}\right] = D^2.
\end{gather*}

In the paper, units are scaled so that $2m = \hbar = 1$.  This means that $g$ provides the length scale for the problem.

```{code-cell}
from numpy import cos, sin, exp
```

```{code-cell}
theta, Gamma = np.random.random(2)
(cos(theta) * exp(-1j * Gamma) + sin(theta) * np.exp(1j * Gamma)) / 2, (
    cos(theta + Gamma) - 1j * sin(theta - Gamma)
) / 4
```

```{code-cell}
%pylab inline --no-import-all
from IPython.display import display, clear_output
from pytimeode.evolvers import EvolverABM
from mmfutils.contexts import NoInterrupt
from gpe import bec2

reload(bec2)
from gpe.bec2 import State2, u


class StateBS(State2):
    t0 = u.micron ** 2 * 2 * u.m / u.hbar

    def get_Vext(self):
        Va = Vb = 0 * self.xyz[0]
        return (Va, Vb, self.Omega)


# Paper dimensions
g_ = 2.0
Omega_ = 2.0
theta_ = np.pi / 4.0
b_ = 2.0
s = StateBS(Nxyz=(128,), Lxyz=(4 * u.micron,), N=1e5)
s = StateBS(
    Nxyz=(64 * 4,),
    Lxyz=(16 * u.micron,),
    N=1e5,
    gs=(-g_ * u.hbar ** 2 / 2.0 / u.m,) * 3,
    Omega=Omega_ * u.hbar ** 2 / 2.0 / u.m,
    cooling_phase=1.0j,
)

# Fig 9
g = -g_ * u.hbar ** 2 / 2.0 / u.m
s = StateBS(
    Nxyz=(64,),
    Lxyz=(16 * u.micron,),
    N=1e5,
    gs=(g * 1.01, g * 1.01, g * 0.99),
    Omega=Omega_ * u.hbar ** 2 / 2.0 / u.m,
    cooling_phase=1.0,
)
x = s.xyz[0]
s[0, ...] = (
    (np.cos(theta_) + np.sin(theta_)) / 2.0 * b_ / np.cosh(b_ * np.sqrt(g_) * x / 2.0)
)
s[1, ...] = (
    (np.cos(theta_) - np.sin(theta_)) / 2.0 * b_ / np.cosh(b_ * np.sqrt(g_) * x / 2.0)
)
# s[...] += np.random.random(s.shape)*0.1
s.plot()
```

```{code-cell}
e = EvolverABM(s, dt=0.5 / s.E_max / u.hbar)
with NoInterrupt(ignore=True) as interrupted:
    while not interrupted:
        e.evolve(10)
        plt.clf()
        e.y.plot()
        display(plt.gcf())
        clear_output(wait=True)
```

```{code-cell}
import sympy
```

```{code-cell}
sympy.init_session()
var("a, theta, b, g, x")
psi = a / cosh(x * b * sqrt(g) / 2)
n = psi ** 2
r1, r2 = (-psi.diff(x, x) / 2).trigsimp().simplify(), (g * n * psi).simplify()
r1, r2
```

```{code-cell}
sympy.expand_trig
```