---
execution:
  timeout: 60
jupytext:
  cell_metadata_filter: -all
  formats: md:myst,ipynb
  main_language: python
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.13
    jupytext_version: 1.17.2
kernelspec:
  name: python3
  display_name: Python 3 (ipykernel)
  language: python
---

```{code-cell} ipython3
import mmf_setup; mmf_setup.nbinit()
import os, sys
from pathlib import Path
FIG_DIR = Path(mmf_setup.ROOT) / '../Docs/_build/figures/'
os.makedirs(FIG_DIR, exist_ok=True)
import logging; logging.getLogger("matplotlib").setLevel(logging.CRITICAL)
%matplotlib
import numpy as np, matplotlib.pyplot as plt
try: from myst_nb import glue
except: glue = None

from matplotlib.animation import FuncAnimation
from gpe.contexts import FPS
```

# Harmonic Traps

Harmonically trapped gases play a special role in cold atom physics for several reasons.

1. All potentials look harmonic close to their minimum.  Thus, a laser with a gaussian
   intensity profile can be used to make a harmonic close to the core:
   \begin{gather*}
     V(x) = V_0 e^{-x^2/2\sigma^2} = V_0\left(
       1 - \frac{x^2}{2\sigma^2} + \frac{x^4}{8\sigma^4} 
       + \cdots\right).
   \end{gather*}
   Identifying the second term with a harmonic trapping potential gives
   \begin{gather*}
     V_0 = -m\omega^2\sigma^2, \qquad
     V(x) = \text{const} + \frac{m\omega^2}{2}x^2 
          - \underbrace{\frac{m\omega^2 x^4}{8\sigma^2}}_{\frac{m^2\omega^4 x^4}{8V_0}}
          + \cdots,
   \end{gather*}
   where the last term characterizes the anharmonicities.
2. Aspects of dynamics are after simple in harmonic traps due to the fact that all
   classical trajectories have the same period.  This leads to many cases where the
   shape of a trapped gas remains invariant, simply oscillating and vibrating with the
   trapping frequencies.  See for example {cite}`Castin:1996` where the effects of
   harmonic traps can be incorporated into scaled coordinates and phases -- see
   {ref}`sec:drGPE` which make use of to model cloud expansion for imaging,
3. Related is Kohn's Harmonic theorem (see {cite}`Wu:2014` for details) which basically
   says that, no matter what the internal interactions are, the center of mass of a
   harmonically trapped cloud will follow the classical trajectory of a particle with
   the total mass of the system.

Here we will demonstrate some of these effects as realized in the GPE.

```{code-cell}
:tags: [hide-input]
from gpe.bec import _StateBase, GPEMixin, HOMixin
from gpe import utils
from gpe.minimize import MinimizeState


# Note: StateBase is earlier in the mro, spoiling the ability to set g in StateHOMixin
@utils._GPU.add_non_GPU_methods
class State(GPEMixin, HOMixin, utils.StateWithExperimentMixin, _StateBase):
    def __init__(self, *v, g, ws, **kw):
        self.g = g
        self.ws = ws
        return super().__init__(*v, **kw)
    
    def shift(self, x0):
        """Shift the wavefunction by x0.  Relies on fourier basis."""
        kx = self.basis.kx
        self.set_psi(
            self.basis.ifft(
                np.exp(1j*kx*x0) * self.basis.fft(self.get_psi(), axis=0), axis=0))

    def get_Vext_GPU(self, state=None):
        if state is None:
            state = self
        x = state.xyz[0] 
        Vext = (sum(state.m*(w*x)**2/2 for w, x in zip(state.ws, state.xyz))
                - state.m*state.ws[0]**2*x**4/8/self.experiment.sigma)
        if state.initializing:
            Vext -= self.mu
            
        return Vext

        
class Experiment(utils.ExperimentBase):
    m = hbar = 1
    ws = 2*np.pi * np.array([1 , 3, 5])
    healing_length = 0.1
    Rx_x_TF = 2.0
    Rperp_perp_TF = 1.3
    dx_healing_length = 0.5
    n0 = 1.0
    State = State
    t_unit = 1
    dim = 1
    sigma = np.inf  # Anisotropy parameter - width of gaussian.
    
    def init(self):
        dx = self.dx_healing_length * self.healing_length
        mu = self.hbar**2 / 2 / self.m / self.healing_length**2
        wx = self.ws[0]
        xyz_TF = np.sqrt(2*mu / self.m) / self.ws
        Lx = 2 * self.Rx_x_TF * xyz_TF[0]
        Lyz = 2 * self.Rperp_perp_TF * xyz_TF[1:]
        Lxyz = [Lx] + list(Lyz)
        Lxyz = Lxyz[:self.dim]
        Nxyz = [utils.get_good_N(L/dx) for L in Lxyz]
        
        g = mu / self.n0
        self.state_args = dict(
            experiment=self,
            hbar=self.hbar,
            m=self.m,
            g=g,
            mu=mu,
            ws=self.ws,  # Maybe should delegate?
            Nxyz=Nxyz,
            Lxyz=Lxyz)
        super().init()
        
    # We should provide a default in ExperimentBase
    def get_state(self):
        return self.State(**self.state_args)
        
    def get_initial_state(self):
        s = self.get_state()
        m = MinimizeState(s)
        psi0 = m.minimize(use_scipy=True).get_psi()
        s = self.get_state()
        s.set_psi(psi0)
        s.pre_evolve_hook()
        return s
    

e = Experiment()
s = e.get_state()
s0 = e.get_initial_state()
s.plot(label="TF");
s0.plot(label="Ground State");
```
We first plot the initial states to make sure we have a large enough box.  Now we shift
the state to the right and watch it oscillate.

```{code-cell}
:tags: [hide-input]
from gpe.contexts import FPS
from pytimeode.evolvers import EvolverABM

e = Experiment()
s = e.get_initial_state()
s.shift(1)
T = 4*e.t_unit
Nt = 100
dT = T/Nt
dt = 0.4*s.t_scale
steps = int(np.ceil(dT/dt))
dt = dT/steps
ev = EvolverABM(s, dt=dt)
fig, ax = plt.subplots()
Es = [s.get_energy()]
for f in FPS(Nt, fig=fig, embed=True):
    ev.evolve(steps)
    s = ev.y
    Es.append(s.get_energy())
    V = s.get_Vext()
    ax.cla()
    ev.y.plot(fig=fig)
    ax.set(ylim=(0, 1.2))
    ax.plot(s.x, V/V.max())
assert np.allclose(Es, Es[0])
```

Notice that the oscillations preserve the shape.  The is a feature of harmonic traps
that follows from Kohn's Harmonic Trapping theorem (see {cite}`Wu:2014` and references
therein for details).  We now show that this also applies in higher dimension by solving
a 2D problem:

```{code-cell}
:tags: [hide-input]
from gpe.contexts import FPS
from pytimeode.evolvers import EvolverABM

e = Experiment(dim=2)
s = e.get_initial_state()
s.shift(1)
T = 1*e.t_unit
Nt = 100
dT = T/Nt
dt = 0.2*s.t_scale
steps = int(np.ceil(dT/dt))
dt = dT/steps
ev = EvolverABM(s, dt=dt)
fig, axs = plt.subplots(2, 1, figsize=(10,5))
axs = s.plot(axs=axs)
Es = [s.get_energy()]
for f in FPS(Nt, fig=fig, embed=True):
    ev.evolve(steps)
    [ax.cla() for ax in axs]
    Es.append(ev.y.get_energy())
    ev.y.plot(axs=axs)
assert np.allclose(Es, Es[0])
```
The bottom frame shows the evolution of the wavefunction phase.  Notice that the
contours remain linear and equally spaced -- this indicates that the velocity of the
cloud is constant, hence the motion is shape preserving.

## Anharmonicities

Here is what happens if we add a slight anharmonicity:

```{code-cell}
:tags: [hide-input]
e = Experiment(sigma=10)
s = e.get_initial_state()
s.shift(1)
T = 4*e.t_unit
Nt = 100
dT = T/Nt
dt = 0.4*s.t_scale
steps = int(np.ceil(dT/dt))
dt = dT/steps
ev = EvolverABM(s, dt=dt)
fig, ax = plt.subplots()
Es = [s.get_energy()]
for f in FPS(Nt, fig=fig, embed=True):
    ev.evolve(steps)
    s = ev.y
    Es.append(s.get_energy())
    V = s.get_Vext()
    ax.cla()
    ev.y.plot(fig=fig)
    ax.set(ylim=(0, 1.2))
    ax.plot(s.x, V/V.max())
assert np.allclose(Es, Es[0])
```

Now the cloud has some internal oscillations.

## Thomas Fermi Approximation

```{code-cell}
e = Experiment(dim=2, healing_length=0.1/4, ws=2*np.pi * np.array([1, 200, 5]),
              Rx_x_TF=1.3, Rperp_perp_TF=2.5)
s = e.get_initial_state()
mu = (s.braket(s.get_Hy(subtract_mu=False)) / s.braket(s)).real
plt.plot(s.x, s.get_density_x())
wx, wy = s.ws[:2]
x_TF0 = np.sqrt(2*mu/s.m)/wx
x_TF1 = np.sqrt(2*(mu-s.hbar*(wx+wy)/2)/s.m)/wx
plt.axvline([x_TF0], c='y', ls=':')
plt.axvline([x_TF1], c='g', ls='-')
```

```{code-cell}
plt.plot(s.get_density()[s.basis.Nxyz[0]//2, :])
```