---
jupytext:
  encoding: '# -*- coding: utf-8 -*-'
  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}
:tags: [hide-input]
import mmf_setup; mmf_setup.nbinit()
```

# The Gross-Pitaevskii Equation (GPE)

:::{margin}
This document describes the code implemented in {mod}`gpe.bec`, which forms the basis
for the remaining code-base.
:::
Here we demonstrate the solution of the Gross-Pitaevskii equation (GPE) (or non-linear
Schrödinger equation (NLSEQ)) describing a weakly interacting Bose-Einsten condensate
(BEC):
\begin{gather*}
  \I\hbar\pdiff{}{t}\psi(\vect{r}, t) =
    \left(
      \frac{-\hbar^2\vect{\nabla}^2}{2m} + V(\vect{r}, t)
      + gn(\vect{r}, t)
   \right)\psi(\vect{r}, t), \qquad
   n(\vect{r}, t) = \abs{\psi(\vect{r}, t)}^2.
\end{gather*}

The idea here is that at very cold temperatures, a collection of bosons can all occupy
the same quantum state.  Without any interactions between the particles ($g=0$), the
system of particles can be described by a single quantum wavefunction $\psi(\vect{x},
t)$ which will evolve through the Schrödinger equation.

Once particle-particle interactions are included ($g \neq 0$) the system must formally
be described by a manybody quantum wavefunction $\Psi(\vect{x}_0, \vect{x}_1, \cdots,
\vect{x}_{N-1}, t)$, but if the interactions are weak and temperature low, this can be
well approximated again by a single wavefunction, suitably normalized so as to represent
the local density $n(\vect{x}, t)$.  The interactions are codified in this theory as a
non-linear term $gn(\vect{x}, t)$ that acts as a density-dependent effective potential
for the wavefunction.

:::{margin}
To simplify the following presentation, we shall now suppress the arguments
$n(\vect{x}, t)$ and $\psi(\vect{x}, t)$.  Unless explicitly stated, densities $n$ and
wavefunctions $\psi$ will depend on both space and time.
:::

This is the essense of the "mean field" approximation - the effects of the interactions
are averaged over all of the particles an replaced by this non-linear term.  For
example, if the particles are repulsive $g>0$, then this term $gn$ acts as a repulsive
potential, pushing other particles away from regions of high density.  Likewise, if the
interaction is attractive, then $g<0$ and this term will act as a potential well,
drawing more particles into regions of high density.

This mean-field description works well at $T\approx 0$ when the interactions are weak in
the sense that $gn^3 \ll 1$.  In this case, the coupling constant $g$ is related to the
two-body s-wave scattering length $a_s$ as:
\begin{gather*}
  g = \frac{4\pi \hbar^2 a_s}{m}.
\end{gather*}

:::{margin}
Note that to obtain this, you must do some integrations by parts to isolated terms like
$\dot{\psi}^\dagger$ in the action.  Check that doing so you obtain the correct signs.
Also note that the non-linear term in the energy-density $\mathcal{E}[\psi]$ is $gn^2/2$
whereas the effective potential is $gn$ which is obtained by differentiating the energy
appropriately.
:::
The GPE follows from the following principle of minimal action:
\begin{gather*}
  S[\psi] = \int\d^{3}{\vect{x}}\d{t} L[\psi], \qquad
  \mathcal{L}[\psi] = \psi^\dagger \I\hbar \dot{\psi} - E[\psi],\\
  \mathcal{E}[\psi] = \frac{\hbar^2\abs{\vect{\nabla}\psi}^2}{2m} + Vn + \frac{gn^2}{2},
\end{gather*}
with the GPE resulting from the Euler-Lagrange equations:
\begin{gather*}
  \pdiff{}{t}\frac{\delta \mathcal{L}[\psi]}{\delta\dot{\psi}^\dagger}
  = \frac{\delta L[\psi]}{\delta\psi^\dagger} 
  = \frac{\delta \mathcal{E}[\psi]}{\delta\psi^\dagger}.
\end{gather*}

## Projections and Cooling

The standard GPE can be expressed as
\begin{gather*}
  \I\hbar \dot{\ket{\psi}} = \op{H}\ket{\psi}
\end{gather*}
where the Hamiltonian operator depends on the state $\ket{\psi}$ through the density
(hence, the non-linear Schr\"odinger equation), and effects the operation 
\begin{gather*}
  \op{H}\ket{\psi} = \pdiff{E}{\bra{\psi}}.
\end{gather*}
Several modifications to this standard evolution are used in the code:

### Imaginary-Time Cooling and Quantum Friction

The first is complex-time evolution, or cooling with **imaginary-time** evolution.  This
is typically denoted through a phase
\begin{gather*}
  \I e^{\I\eta} \hbar \dot{\ket{\psi}} = \op{H}\ket{\psi}
\end{gather*}
where $e^{\I \eta}$ is the {attr}`.StateBase.cooling_phase` parameter.  If
`cooling_phase = 1j`, then we have what is referred to as imaginary-time cooling and the
evolution implements a gradient descent:
\begin{gather*}
  \dot{\ket{\psi}} = -\hbar^{-1}\pdiff{E}{\bra{\psi}}.
\end{gather*}
This is guaranteed to reduce the energy.
:::{admonition} Details including Quantum Friction
:class: dropdown

Consider evolving the state with an operator $\op{A}$:
\begin{gather*}
  \dot{\ket{\psi}} = \op{A}\ket{\psi}.
\end{gather*}
Assuming that there is no explicit time dependence in $E[\ket{\psi}]$, the energy will
evolve as 
\begin{gather*}
  \dot{E} = \pdiff{E}{\ket{\psi}}\ket{\dot{\psi}} + \text{h.c.}
          = \braket{\psi|\op{H}\op{A}|\psi} + \text{h.c.}.
\end{gather*}
More generally -- with explicit time-dependence -- this can be expressed as
\begin{gather*}
  \dot{E} = \braket{\psi|\dot{\op{H}}|\psi} 
          + \braket{\psi|(\op{H}\op{A} + \op{A}^\dagger\op{H})|\psi}.
\end{gather*}
Complex time evolution evolves with $\op{A} = \op{H}/(\I \hbar e^{\I\eta})$.  Thus, if
$\eta = 0$, then $\op{A} = \op{H}/\I\hbar$ is anti-hermitian, and we obtain Heisenberg's
equations of motion
\begin{gather*}
  \dot{E} = \braket{\psi|\dot{\op{H}}|\psi} 
          + \underbrace{\braket{\psi|[\op{H},\op{A}]|\psi}}_{0}
\end{gather*}
where the energy is conserved, unless there is an explicit time-dependence (e.g., a
stirring potential) doing work on the system.  For imaginary-time evolution, $\op{A} =
-\op{H}/\hbar$ is hermitian:
\begin{gather*}
  \dot{E} = \braket{\psi|\dot{\op{H}}|\psi} 
          - \frac{1}{2\hbar}\braket{\psi|\op{H}^\dagger\op{H}|\psi}.
\end{gather*}
Since the second term is the expectation of a positive-semi-definite operator
$\op{H}^\dagger\op{H}$, the energy is guaranteed to never increase unless external work
is being done on the system through the first term.

This formalism allows for some generalizations as discussed in {cite}`Bulgac:2013c` that
we call **quantum friction**.  The idea here is to express
\begin{gather*}
  \op{A} = \op{H} + \op{U}
\end{gather*}
where $\op{U}$ is a local operator (i.e. a time-dependent potential).  In
{cite}`Bulgac:2013c` we show that if we take $\op{U}\propto \dot{n}(\op{x})$, then we
can lower the energy with a local potential.  This is not so useful for the GPE, but
becomes very important for distributed multi-orbital functionals (like the SLDA used for
fermions) where communication prohibits the action of non-local operators.
:::

### The dissipative GPE (dGPE)

### Constraints

The previously mentioned cooling techniques will not only remove energy from the system,
but may change the particle number.  To fix this, one simply subtracts the chemical
potential from the Hamiltonian
\begin{gather*}
  \op{H} - \mu \op{1}, \qquad
  \mu = \braket{\psi|\op{H}|\psi}.
\end{gather*}
More generically, if we implement an evolution
\begin{gather*}
  \dot{\ket{\psi}} = \op{A}\ket{\psi}
\end{gather*}
and want to constrain a quantity $\op{C}= \op{C}^\dagger$:
\begin{gather*}
  c = \braket{\psi|\op{C}|\psi}, \qquad
  \op{C}\ket{\psi} = \pdiff{c}{\bra{\psi}},
\end{gather*}
then we want to make sure that the evolution is orthogonal to the gradient of the
constraint.  This can be done by projecting out the change along this direction:
\begin{gather*}
  \op{\Pi}_{\perp} = \frac{\op{C}\ket{\psi}\bra{\psi}\op{C}^\dagger}
                          {\sqrt{\braket{\psi|\op{C}^\dagger\op{C}\psi}}},\qquad
  \dot{\ket{\psi}} = \Bigl(\op{1} - \op{\Pi}_{\perp}\Bigr)\op{A}\ket{\psi}.
\end{gather*}

## Ground State

To obtain the ground state, we find the wavefunction that minimizes the energy $E[\psi]
= \int\d^3{\vect{x}} \mathcal{E}[\psi]$.  However, this minimization must be
constrained: $\psi = 0$ in general minimizes the energy but with no particles.  Thus, we
obtain the constrained minimization problem:

\begin{gather*}
  \min_{\psi} E[\psi] \text{ where } N[\psi] = \int \d^3{\vect{x}} \abs{\psi}^2 = N_0.
\end{gather*}
This can be implemented using Lagrange multipliers:
\begin{gather*}
  \frac{\delta}{\delta\psi^\dagger} \left(E[\psi] - \mu N[\psi]\right) = 0,
\end{gather*}
which gives the time-independent non-linear Schrödinger equation:
\begin{gather*}
  \left(
      \frac{-\hbar^2\vect{\nabla}^2}{2m} + V
      + gn
   \right)\psi = \mu \psi.
\end{gather*}
Solutions to this equation are "stationary" in the sense that their only time-evolution is a continual phase change (unobservable):
\begin{gather*}
  \psi(\vect{x}, t) = e^{-\I\mu t/\hbar}\psi(\vect{x}) = e^{-\I\omega t}\psi(\vect{x}),
\end{gather*}
where the phase oscillates with frequency:
\begin{gather*}
  \omega = \frac{\mu}{\hbar}.
\end{gather*}

## Some Solutions

### Homogeneous States

In the absence of an external potential, the system has translational invariance, hence
momentum is a conserved and provides good quantum number $\vect{k}$.  These homogeneous
states are solutions of the GPE with constant density $n = n_0$:
\begin{gather*}
  \psi_{\vect{k}}(\vect{x}, t) = \sqrt{n_0}e^{\I(\vect{k}\cdot\vect{x} - \omega t)},
  \qquad
  \mu = \hbar \omega = \frac{\hbar^2 k^2}{2m} + gn_0.
\end{gather*}

If the potential varies slowly, then it can be a good approximation to treat each region
of the gas locally as being homogeneous.  This approximation is called the Thomas-Fermi
approximation and amounts to solving the GPE in the limit where the kinetic energy can
be neglected.  These Thomas-Fermi states thus satisfy:
\begin{gather*}
  [gn(\vect{x}) + V(\vect{x}) - \mu]\psi_{TF}(\vect{x}) = 0, \qquad
  \psi_{TF}(\vect{x}) = \begin{cases}
    \sqrt{\frac{\mu - V(\vect{x})}{g}} & V(\vect{x}) < \mu\\
    0 & V(\vect{x}) > \mu
  \end{cases}.
\end{gather*}

When there are lots of particles in a smoothly varying system, this approximation is good starting point, and often describes cold-atom experiments well.

:::{todo} **Khalid: Discuss Pressure etc., speed of sound, etc.**
:::

Let's start by considering a homogeneous gas with constant $\psi(\vect{x}, t) =
\sqrt{n}$.  The energy density is thus:

\begin{gather*}
  \mathcal{E}(n) = \frac{g}{2} n^2.
\end{gather*}

Using thermodynamics, we can derive the chemical potential and pressure:

\begin{gather*}
  \mu = \pdiff{\mathcal{E}(n)}{n} = gn, \qquad
  P(\mu) = \mu n - \mathcal{E} = \frac{gn^2}{2} = \frac{\mu^2}{2g}.
\end{gather*}

### Bright Solitons

In the case of an attractive interaction $g<0$ there is an analytic self-bound solution
with variations along one dimension (i.e. this forms a flat wall): 
\begin{gather*}
  \psi(x, t) = \sqrt{n_0}\sech[\eta(x-vt)]e^{\I(kx - \omega t)}\\
  \eta = \sqrt{\frac{-gmn_0}{\hbar^2}},\qquad
  k = \frac{mv}{\hbar},\qquad
  \mu = \hbar \omega = \hbar^2\frac{k^2-\eta^2}{2m}
\end{gather*}

### Dark Solitons

In a background of constant density, there is an excited state that is an analytical
solution to the GPE called a dark (or grey) soliton, which is an extended object with
structure in one dimension, forming an "anti-wall" with a density depletion in a plane:
\begin{gather*}
  \psi(x, t) = \sqrt{n_0}\left(
    \I\frac{v}{c} + \frac{u}{c}\tanh[\eta(x-vt)]
  \right)e^{-\I\omega t},\\
  \eta = \frac{m u}{\hbar}, \qquad
  \mu = \hbar\omega = gn_0 = mc^2, \qquad
  u = \sqrt{c^2-v^2},\qquad
\end{gather*}

This soliton moves with speed $v$, and $c = \sqrt{gn_0/m}$ is the speed of sound.

**Khalid: Discuss healing length etc.**

### Harmonic Oscillator

A common example is a BEC trapped in a 1D harmonic oscillator potential:
\begin{gather*}
  V(x) = \frac{m\omega^2 x^2}{2}.
\end{gather*}
The introduction of this potential provides another length scale in the problem:
\begin{gather*}
  a_{HO} = \sqrt{\frac{\hbar}{m\omega}}.
\end{gather*}

This potential is ubiquitous because almost all functions can be expanded about their
minimum as a Taylor series and the quadratic term being dominant for small deviations.
For example, most optical trapping potentials provide a Gaussian potential which effects
a harmonic trap if the power is high enough:
\begin{gather*}
  V(x) = -Ae^{-x^2/2\sigma^2} \approx -A + \frac{A x^2}{2\sigma^2} + \order(x^3).
\end{gather*}

Unfortunately, this potential does not admit an analytic solution in general, however,
if we set $g=0$, then the GPE reduces to the Schrödinger equation for the [quantum
harmonic oscillator](https://en.wikipedia.org/wiki/Quantum_harmonic_oscillator) whose
solutions are well known and can be used to test the code:
\begin{gather*}
  \psi(x, t) = \frac{1}{\sqrt{2^n n!}\sqrt[4]{\pi}}\frac{1}{\sqrt{a_{n}}}
  \exp\left\{-\frac{x^2}{2a_{HO}^2}\right\}
  H_n\left(\frac{x}{a_{HO}}\right)
  e^{-\I\omega_n t}, \qquad
  \mu_n = \hbar\omega_n = \hbar\omega\left(n + \frac{1}{2}\right).
\end{gather*}

## Length Scales

We start by considering the dimensions of various quantities:

\begin{gather*}
  [m] = M, \qquad
  [\hbar] = [Et] = [mv^2t] = \frac{M D^2}{T}, \qquad
  [a_s] = D, \qquad
  [gn] = [E] = \frac{M D^2}{T^2}, \qquad
  [g] = [EV] = \frac{M D^5}{T^2}, \qquad
  [n] = [1/V] = \frac{1}{D^3}, \qquad
  [\psi] = [\sqrt{n}] = \frac{1}{D^{3/2}}.
\end{gather*}

To understand the relevant scales in this problem we need to identify the unique dimensionless parameters.  To do this, we can do the following: let $m$ defined the mass scale by setting $m=1$ and let $\hbar = 1$ define the time scale.  This allows us to express everything in units of length:

\begin{gather*}
  [m] = M = 1, \qquad
  [\hbar] = 1 = \frac{D^2}{T}, \qquad
  [t] = T = D^2, \qquad
  [g] = [EV] = D, \qquad
  [n] = [1/V] = \frac{1}{D^3}.
\end{gather*}

+++

Thus, there are only two fundamental dimensionless quantities - the interaction $g$ and the density $n$.  The corresponding length scales are:

\begin{gather}
  a_s = \frac{gm}{4\pi\hbar^2}, \tag{scattering length}\\
  r_0 = \left(\frac{3}{4\pi n}\right)^{1/3}. \tag{interparticle separation}
\end{gather}

To use this, one can form dimensionless quantities by first cancelling the length scales, then using factors of $m$ and $\hbar$ to get rid of factors of $M$ and $T$ respectively.  The only dimensionful parameter in the free theory is thus:

\begin{gather}
  \gamma = n a_s^{3} = n\left(\frac{mg}{4\pi \hbar^2}\right)^3. \tag{gas parameter}
\end{gather}

The GPE is generally a valid approximation if $\gamma \ll 0.01$.

*Occasionally this is expressed as $1/k_Fa_s = 1/(3\pi^2 n)^{1/3}a_s$ where $n = k_F^3/3\pi^2$ which comes from an analogy with a Fermi gas.*

+++

If we were to consider the BEC at finite temperature we would have another scale, temperature, which can be equated to an energy through the Boltzmann constant $k_B=1$ so $[T] = [E] = 1/D^2$.  This introduces another length scale – the [thermal de Broglie wavelength](https://en.wikipedia.org/wiki/Thermal_de_Broglie_wavelength):

\begin{gather*}
  \lambda_{\text{th}} = \sqrt{\frac{2\pi \hbar^2}{m k_B T}}.
\end{gather*}

Thus we can introduce another dimensionless constant:

\begin{gather*}
  \ln z = \frac{\lambda_{\text{th}}^2}{a_s^2} = \frac{gn}{k_B T}
        = \frac{4(2\pi)^3 \hbar^6}{m^2g^2 k_B T}
\end{gather*}

The exponential of this $z$ is called the "fugacity".  Again, the GPE is valid only for $z \gg 1$

## Moving Frames and Twisted Boundaries

Two related extensions that can be considered as part of our basis allow us to work in a
moving frame and to apply twisted boundary conditions.  These features are implemented
in the {class}`.StateTwist_x` class.  We start with moving frames.

### Boosts

Consider a solution to the GPE in 1D:
\begin{gather*}
  \I\hbar \dot{\psi}(x, t) =
  -\frac{\hbar^2\psi''(x, t)}{2m}
  + g\abs{\psi(x, t)}^2\psi(x, t).
\end{gather*}
Suppose we wish to study a solution that travels with velocity $v$.  I.e. something
which should be static when considered as a function of $x_v = x - vt$.  As discussed in
[Galilean Covariance][], we may consider several transformations of the form:
\begin{gather*}
  \psi(x, t) = e^{\I\phi(x_v, t)}\psi_v(x_v, t)
\end{gather*}

* Simple boost:

  \begin{gather*}
    \hbar\phi(x_v, t) = 0, \\
    \I\hbar\dot{\psi}_v(x_v, t) = -\frac{\hbar^2\psi_v''(x_v, t)}{2m}
    + \overbrace{\I\hbar v \psi_v'(x_v)}^{-v\op{p}\psi_v}
    + g\abs{\psi_v(x_v, t)}^2\psi_v(x_v, t).
  \end{gather*}

* Full Galilean Transformation

  \begin{gather*}
    \hbar\phi(x_v, t) = mvx_v + \tfrac{m}{2}v^2t, \\
    \I\hbar\dot{\psi}_v(x_v, t) = -\frac{\hbar^2\psi_v''(x_v, t)}{2m} 
    + g\abs{\psi_v(x_v, t)}^2\psi_v(x_v, t).
  \end{gather*}

The full transformation makes the covariance of the original equation manifest - the
form of the equation stays the same.  However, it does not simplify if the original
equations are not Galilean covariant (e.g. if there is a modified dispersion due to a
spin-orbit coupling.)

[Galilean Covariance]: <http://swan.physics.wsu.edu/forbes/public/student_resources/galilean-covariance/>

### Twists (Bloch waves)

[Bloch's theorem](https://en.wikipedia.org/wiki/Bloch_wave#Proof_of_Bloch's_theorem)
tells us that if we have a periodic potential with period $L$, then we can classify the
eigenstates in terms of their Bloch quasi-momentum $\hbar k_B \in [-\hbar\pi/L,
\hbar\pi/L)$ where the wavefunction is:
\begin{gather*}
  \psi(x) = e^{\I k_B x}u_{k_B}(x), \qquad \frac{-\pi}{L} \leq k_B < \frac{\pi}{L},
\end{gather*}
and $u_{k_B}(x)$ is periodic.  In our code, we store the periodic function $u_{k_B}(x)$.
The implementation must keep track of this twist phase in several places:

* When setting or getting the physical wavefunction $\psi(x)$, the twist must be applied or removed.  This is done with `get_psi()` and `set_psi()` functions.
* When applying the momentum operator, we must shift the momenta:

  \begin{gather*}
    \op{p}\psi(x) = -\I\hbar\pdiff{}{x}e^{\I k_B x}u(x)
    = e^{\I k_B x}\left(-\I\hbar\pdiff{}{x} + \hbar k_B\right)u(x)
    = e^{\I k_B x}(\op{p} + \hbar k_B)u(x).
  \end{gather*}

  When applying the momentum operator directly to $u(x)$, we must therefore shift the
  discrete momenta $k_n \rightarrow k_n + k_B$.
  
:::{note}
If one works with a box of twice the size, then one has only half of the possible range
of twists.  This represents a general feature of Bloch waves: the twist allows you to
consider solutions that will not fit in your periodic box, but that would fit in a
larger box containing multiple unit cells.  Thus, allowing for twists allows you to
consider solutions to a periodic potential in an infinite universe.  If you want to
actually consider physics in a periodic universe, you should not allow twists.
:::

### Twists and Boosts

Consider the GPE with general dispersion:
\begin{gather*}
  \I\hbar \dot{\psi}(x, t) = \bigl(E(\op{p}) + g\abs{\psi(x, t)}^2 + V(x, t)\bigr)\psi(x, t).
\end{gather*}
When combining both twists $k_B$ and boosts $x_v = x-vt$, we boost the twisted
wavefunction.  Thus we have:
\begin{gather*}
  \psi(x, t) = \psi_v(x_v, t) = e^{\I k_B x_v}u_v(x_v, t) = e^{\I k_B x}u_v(x_v, t)\\
  \I \hbar \dot{u}_v(x_v, t)
  = \bigl(E(\op{p} + \hbar k_B) - (\op{p}+\hbar k_B) v + g\abs{u_v(x_v)}^2 
          + V(x_v + vt, t)\bigr)u_v(x_v, t).
\end{gather*}

The {class}`.StateTwist_x` thus works internally with $u_v(x_v, t)$ on the
abscissa $x_v = x - vt$ where $x$ are the abscissa in the "lab frame" (accessible from
the property {attr}`.StateTwist_x.x_lab`).

For comparison, here is what happens if we twist the boosted wavefunction:
\begin{gather*}
  \psi(x, t) = e^{\I k_B x}\tilde{u}(x, t) = e^{\I k_B x}\tilde{u}_v(x_v, t)\\
  \I \hbar \dot{\tilde{u}}_v(x_v, t)
  = \bigl(E(\op{p} + \hbar k_B) - \op{p} v + g\abs{\tilde{u}_v(x_v)}^2 
          + V(x_v + vt, t)\bigr)\tilde{u}_v(x_v, t),\\
  \tilde{u}_v(x_v, t) = e^{-\I k_B vt}u_v(x_v, t).
\end{gather*}
The resulting equation does not have $\op{p}$ shifted in the boost term.  These two
solutions are related by a phase.

## Numerical Examples

The base code for solving the GPE in 1D with a single component is provided in the file
[bec.py](bec.py).  It uses the [pytimeode](https://bitbucket.org/mforbes/pytimeode)
project to implement the time dependence.  This requires providing a `State` class with
come methods that tell the evolvers how to compute the time derivative.  The
`compute_dy_dt()` method allows one to use generic evolvers like the
Adams-Bashforth-Milne (ABM) predictor corrector method, and the
`apply_exp_K()`/`apply_exp_V()` pair allow one to use the faster (but less accurate)
Split-Operator evolvers.  See the code and
[pytimeode](https://bitbucket.org/mforbes/pytimeode) documentation for details.

### Default Example

The default example in the code is a BEC in a Harmonic Oscillator potential.  We we
obtain the initial state using the Thomas-Fermi approximation: 

**Note:** *Here we use the code in {mod}`gpe.bec_basic` which is a simplified version of
the code for pedagogical purposes.  Production code should use {mod}`gpe.bec`.
One major difference is the introduction of {attr}`.StateBase.basis` which handles the
discretization (allowing, for example, axial symmetric abscissa).*

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

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

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

```{code-cell}
from gpe import bec, minimize
from pytimeode.evolvers import EvolverABM
from mmfutils.contexts import FPS
s = bec.State(Nxyz=(256,), Lxyz=(10,), x_TF=5/2)
s = minimize.MinimizeState(s, fix_N=True).minimize()
h = s.hbar / np.sqrt(2*s.m*s.g*s.get_density().max())
s.cooling_phase = 1+0.1j
psi = s.get_psi()
x = s.get_xyz()[0]
x0 = 2.5/2*0
theta = np.where(x < x0, 0, np.pi)
s.set_psi(np.exp(1j*theta)*psi*np.tanh((x-x0)/np.sqrt(2)/h))
ev = EvolverABM(s, dt=0.1*s.t_scale)
steps = 1000
for frame in FPS(1000):
    ev.evolve(steps)
    plt.clf()
    ev.y.plot()
    ax = plt.twinx()
    ax.plot(x, np.angle(ev.y/ev.y[s.basis.Nxyz[0]//2]), 'C1')
    clear_output(wait=True)
    display(plt.gcf())
    
```

### 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,))
    s.g = 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
    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 = list(map(get_err, Ns))
plt.semilogy(Ns, errs, "-+")
```

#### Question

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

+++

#### Partial Solution

```{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, s.kxyz[0].max())
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]), np.fft.fftshift(abs(np.fft.fft(psi_0))), "-+")
```

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

+++

### Exact Solution with Interactions

+++

Although there is no analytic solution for the harmonic oscillator with an interaction, we can construct a potential so as to have an exact solution for an interacting system.  In particular, we solve the time-independent GPE for the potential to obtain:

\begin{gather*}
  V_\psi(x) - \mu
  = \frac{\frac{\hbar^2\nabla^2\psi(x)}{2m} - gn(x)\psi(x)}{\psi(x)}
  = \frac{\hbar^2\nabla^2\psi(x)}{2m \psi(x)} - gn(x).
\end{gather*}

Given any real wavefunction $\psi(x) > 0$ without nodes, we can simply apply this formula to determine an external potential in which this state will be a stationary state with chemical potential $\mu$.  This will also work for functions with nodes as long as $\nabla^2\psi(x)$ has nodes in the same location.

+++

For example, we can construct a potential with a Gaussian ground state and zero chemical potential:

\begin{gather*}
  \psi(x) = \sqrt{n_0}\exp\left\{\frac{-x^2}{2a^2}\right\}, \qquad
  \psi''(x) = \frac{x^2-a^2}{a^4}\psi(x),\qquad
  V_\psi(x) = \frac{\hbar^2}{2m}\frac{x^2-a^2}{a^4} - gn_0e^{-x^2/a^2}.
\end{gather*}

+++

Notice that if $g=0$ we end up with shifted harmonic oscillator potential (check that the coefficients agree with the solution given above).

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

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

s = State(Nxyz=(64,), Lxyz=(23 * u.micron,))
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):
        self._V_psi = 0  # Set this because State.__init__() needs a potential
        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
        self._V_psi = u.hbar**2 / 2.0 / u.m * (4 * (k * x) ** 2 - 2 * k) - self.g * n_0
        self.data[...] = psi_0
        self.pre_evolve_hook()

    def get_Vext(self):
        return self._V_psi


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

#### Question

We explicitly constructed the state to have $\mu = \hbar\omega = 0$.  Why then is the energy $E\approx -3.3348$ negative?

+++

Here we start from another state ($\psi(x) = 1$) and evolve with imaginary time to see if we recover the state we initially constructed the potential from.

```{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)
```

#### Question

Although the state $\psi(x)$ is a node-less stationary solution of the GPE with external potential $V_\psi(x)$ constructed above, it is not necessarily the ground state in this potential.  Non-linear problems do not necessarily follow the same rules as the linear Schrödinger equation where one can prove that the ground state (and only the ground state) must have no nodes.  Here multiple states can be nodeless.

* Is this the ground state?
* Can you construct a different problem where the solution $\psi(x)$ in the potential $V_\psi(x)$ is different (i.e. if this is the ground state, find a different potential where the constructed solution is not a ground state, or vice versa.)
* What are the necessary and sufficient conditions for $\psi(x)$ to be the ground state of the GPE with external potential $V_\psi(x)$? *(Hard)*

+++

Here is an example of real-time evolution of a state with domain wall phase imprinted.  Locally this domain wall has the structure of a dark soliton, but now since we are in a trap, the solution is only approximate.  We evolve this state with some imaginary time cooling to allow energy to be removed from the system.  As a result, we see the soliton oscillate back and forth with larger and larger amplitude but shallower and shallower depth until it eventually is lost.

```{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)
```

## 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*}
  \DeclareMathOperator{\sech}{sech}
  \psi(x) = 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
\end{gather*}

```{code-cell}
import numpy as np
import sympy
```

```{code-cell}
sympy.init_printing()
m, hbar, b, g = sympy.var("m, hbar, b, g", positive=True)
x = sympy.var("x", real=True)
psi = b / sympy.cosh(b * sympy.sqrt(g * m) * x / hbar)
n = psi**2
E = sympy.trigsimp((-(hbar**2) * psi.diff(x, x) / 2 / m - g * n * psi) / psi).together()
N_1D = (psi**2).integrate((x, -sympy.oo, sympy.oo))
sympy.Eq(sympy.S("E_"), E), sympy.Eq(sympy.S("N_1D"), N_1D)
```

One might ask if similar objects are possible in 3D.  Consider a self-bound object with wavefunction $\psi_{D}(\vect{x})$ and total particle number $N = \int \d^{d}\vect{x}\abs{\psi_D(x)}^2$.  Now scale it as follows so that the particle number is unchanged:

\begin{gather*}
  \psi_{s} = s^{D/2} \psi_{D}(s\vect{x}).
\end{gather*}

The kinetic energy and potential energy thus scale as:

\begin{gather*}
  K_s = s^{-2}K_D, \qquad
  V_s = s^{D}V_D.
\end{gather*}

Thus, such an object will have a minimum energy when $2s^{-3}K_D = Ds^{D-1}V_D$, or

\begin{gather*}
  s^{D+2} = \frac{2K_D}{DV_D}.
\end{gather*}

+++

## 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_basic

reload(bec_basic)
from pytimeode.evolvers import EvolverABM
from mmfutils.contexts import NoInterrupt
from gpe.bec_basic 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
```

## Interacting with a BEC

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

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

s = State(Nxyz=(64,), Lxyz=(23 * u.micron,))
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 get_Vext(self):
        V_ext = State.get_Vext(self)
        if self.t > 0:
            V_ext += 0.01 * np.cos(s.ws[0] * self.t) * self.xyz[0]
        return V_ext


s = State1(Nxyz=(64,), Lxyz=(23 * u.micron,))
s.plot()
plt.plot(x, s.get_Vext())
```

```{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.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()
        plt.plot(s.xyz[0], e.y.get_Vext())
        display(plt.gcf())
        clear_output(wait=True)
```

## Errors in Derivatives

```{code-cell}
def f(x=0.5):
    return np.cos(x)


def df(x=0.5):
    return -np.sin(x)


def df1(h, x=0.5):
    return (f(x + h) - f(x)) / h


def err(h):
    return abs(df1(h) - df())


hs = 10 ** (np.linspace(-16, -1, 100))
errs = list(map(err, hs))
```

```{code-cell}
%pylab inline --no-import-all
plt.loglog(hs, errs, "-+")
plt.plot(hs, abs(hs * f() / 2.0), "--")
plt.plot(hs, abs(np.finfo(np.double).eps / hs), ":")
```