---
jupytext:
  cell_metadata_json: true
  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:
  argv: [/usr/bin/python3, -m, ipykernel, --HistoryManager.enabled=False, --matplotlib=inline,
    -c, '%config InlineBackend.figure_formats = set([''retina''])

      import matplotlib; matplotlib.rcParams[''figure.figsize''] = (12, 7)', -f, '{connection_file}']
  display_name: Python 3 (system-wide)
  env: {}
  language: python
  metadata:
    cocalc:
      description: Python 3 programming language
      priority: 100
      url: https://www.python.org/
  name: python3
  resource_dir: /ext/jupyter/kernels/python3
---

```{code-cell}
:init_cell: true
:tags: [hide-cell]

import numpy as np, matplotlib.pyplot as plt
import mmf_setup;mmf_setup.nbinit()
import logging;logging.getLogger('matplotlib').setLevel(logging.CRITICAL)
```

(sec:drGPE)=
# Dynamically Rescaled GPE (dr-GPE)

```{toctree}
---
maxdepth: 2
titlesonly:
caption: "Prerequisites"
---
CoordinateTransformations
NPSEQ
```


## Introduction

The dynamically-rescaled GPE (dr-GPE) is a combination of {ref}`sec:NPSEQ` and
{ref}`sec:CoordinateTransformations`.  As with the NPSEQ, the primary
use-case is a quasi-1D system with tight harmonic radial confinement along the $y$ and
$z$ directions, but with non-trivial dynamics along $x$.  The dr-GPE enhances this by
allowing for simple scaling dynamics in the radial direction when the radial trapping
frequency is adjusted.

This formalism is particularly useful for simulating expansion dynamics in cold atom
experiments where the radial confinement is removed to allow for expansion before
imaging.  The NPSEQ cannot account for this due to the assumption that radial dynamics
are rapid.  In the bare NPSEQ formalism, turning off the radial confinement for
expansion would cause the radial wavefunction to suddenly have infinite extent with zero
local density.

The dr-GPE instead uses scaled coordinates that naturally take into account such
dynamics, allowing for a realistic simulation with time-dependence in the radial
direction.

The reader should first make sure they understand the content of {ref}`sec:NPSEQ`, upon
which we will expand.

## Formalism

The first formalism of the NPSEQ posits a wavefunction factor as
\begin{gather*}
  \Phi(\vect{x}, t) = \phi(\vect{x}, t) \psi(z, t), \qquad
  \phi(\vect{x}, t) = \frac{1}{\sqrt{\pi}\sigma(z, t)}e^{-(x^2+y^2)/2\sigma^2(z, t)},
\end{gather*}
and makes assumptions about the transverse wavefunction $\phi$ that yield approximate 1D
effective theories for $\psi(z, t)$.

The second formalism is to work with dynamically scaled coordinates:
\begin{gather*}
  \Psi(\vect{X}, t) = e^{\I\theta(\vect{X}, t)}
                      \frac{\Phi(\vect{x}, t)}
                           {\sqrt{\lambda_x(t)\lambda_y(t)\lambda_z(t)}}, \qquad
  X_i = \lambda_i(t)x_i.
\end{gather*}
Here the physical coordinates $\vect{X}$ are rescaled by time-dependent factors
$\vect{\lambda}(t)$.  No approximations are made -- this simply amounts to working in a
non-inertial frame -- but by carefully choosing the scaling factors, we can absorb the
bulk effects of expansion in a time-dependent harmonic trapping potential so that
$\Phi(\vect{x}, t)$ changes minimally and remains in a similar co-moving volume as the
original cloud before expansion.

The dr-GPE combines these, applying the carefully chosen scale-factors so that the
approximations made in the NPSEQ are accurate when applied to $\Phi(\vect{x}, t)$ in the
**co-moving** coordinates $\vect{x}$ rather than the lab frame $\vect{X}$.

# Summary

## Time-Dependent Harmonic Trap

The effects of a time-dependent harmonic trap can be reduced by the following modified
wavefunction: 
\begin{gather*}
  \Psi(\vect{X}, t)
    = e^{\I\theta(\vect{X}, t)}
    \frac{\Phi(\vect{x}, t)}{\sqrt{\lambda_x(t)\lambda_y(t)\lambda_z(t)}}, \qquad 
  X_i = \lambda_i(t)x_i, \qquad
    \theta = \frac{m}{2\hbar} \sum_i \frac{X_i^2\dot{\lambda}_i(t)}{\lambda_i(t)}.
\end{gather*}
With this transformation, the following theories are equivalent:
\begin{gather*}
  \I\hbar\dot{\Psi}(\vect{X}, t) = 
  \left(
    \frac{-\hbar^2\nabla_{\vect{X}}^2}{2m} 
    + \frac{m \sum_{i}\omega_i^2(t)X_i^2}{2}
    + \mathcal{E}'(n) + V(\vect{X})
    \right)\Psi(\vect{X}, t), \\
  \I\hbar\dot{\Phi}(\vect{x}, t) = 
  \left(
    \frac{-\hbar^2\nabla_{\vect{X}}^2}{2m} 
        + \frac{m \sum_{i}\bigl(
      \overbrace{\omega_{i}^{2}(t)\lambda_{i}^{2} + \lambda_{i}\ddot{\lambda}_{i}}^{\tilde{\omega}_i^2(t)}\bigr)x_i^2}{2}
    + \mathcal{E}'(n) + V(\vect{X})
    \right)\Phi(\vect{x}, t).
\end{gather*}
Note that four parts must be added to the code:
1. The derivatives in the kinetic energy must be rescaled.  Note that they are still
   expressed in the original variables $\vect{\nabla}_{X_i} =
   \vect{\nabla}_{x_i}/\lambda_i(t)$.
2. The trapping frequencies must be modified: $\tilde{\omega}_i(t)^2 =
   \omega_i^2\lambda_i^2 + \lambda_i\ddot{\lambda}_i$.
3. The density must be scaled to include the factor $\Lambda = \lambda_x\lambda_y\lambda_z$:
   $n = \abs{\Psi}^2/\Lambda$.
4. The corrections to the non-polynomial interaction induced by the NPSEQ formalism must
   be included.  Here is a schematic of the








We keep this general form in the code, but note that the main utility of this
transformation is that the classical behavior of an expanding gas can be modeled by
setting:
\begin{gather*}
  \ddot{\lambda}_i(t) =
  \frac{\omega_i^2(0)}{\lambda_i(t)\lambda_x(t)\lambda_y(t)\lambda_z(t)} 
  - \omega_i^2(t)\lambda_i(t).
\end{gather*}
With this transformation, $\Phi(\vect{x}, t)$ evolves minimally (see
{cite}`Castin:1996`) and the effective trapping frequencies become
\begin{gather*}
  \tilde{\omega}_i^2(t) = \frac{\omega_i^2(0)}{\Lambda(t)}, 
  \qquad \Lambda(t) = \lambda_x(t)\lambda_y(t)\lambda_z(t).
\end{gather*}
Crucial for the dr-GPE discussed below is the fact that, even if the traps are
completely switched off $\omega_i(t>0) = 0$, there remains a trapping potential in the
expanding frame.  This is crucial because an approximation made in the NPSEQ is that the
radial extent $\sigma$ responds instantaneously to changes in the potential.  Without
this correction, an expanding cloud in the NPSEQ would have $\sigma \rightarrow \infty$
immediately, destroying the approximation.

:::{admonition} Full Transform

If we apply the full transform to all coordinates, we have
\begin{gather*}
  \I\hbar\dot{\Phi}(\vect{x}, t) = 
  \left(
    \frac{-\hbar^2\nabla_{\vect{X}}^2}{2m} 
    + \frac{m \sum_{i}\omega_i^2(0)x_i^2}{2\lambda_x\lambda_y\lambda_z}
    + \mathcal{E}'(n) + V(\vect{X})
    \right)\Phi(\vect{x}, t).\\
\end{gather*}
:::

:::{admonition} Partial Transforms

What happens if we want to scale only two of the coordinates?  This might be necessary
if the dispersion relation along $x$ is not quadratic, for example -- see
{ref}`sec:SOC`.  *(Why? The formulae derived here are require quadratic dispersion.  The
formulae would be much more complicated with higher-order dispersion.)*

The differential equation here describes the expansion of a classical gas, and the
coupling of the directions accounts for the mean-field pressure.  It is thus important
to consider the full solution, even if we will ultimately only apply it to the $y$ and
$z$ coordinates.  Thus we solve all three coupled equations, even if only partially
applying them.

If we then apply only the scaling to $x$ and $y$, leaving $z$ un-scaled, we have
\begin{gather*}
  \Psi(\vect{X}, t)
    = e^{\I\theta(\vect{X}, t)}
    \frac{\Phi(\vect{x}, t)}{\sqrt{\lambda_x(t)\lambda_y(t)}},\\
  X = \lambda_x(t)x, \qquad
  Y = \lambda_y(t)y, \qquad
  Z = z,\\
  \theta = \frac{m}{2\hbar} \left(
    \frac{X^2\dot{\lambda}_x(t)}{\lambda_x(t)}
    +
    \frac{Y^2\dot{\lambda}_y(t)}{\lambda_y(t)}
  \right),\\
  \I\hbar\dot{\Phi}(\vect{x}, t) = 
  \left(
    \frac{-\hbar^2\nabla_{\vect{X}}^2}{2m} 
    + \frac{m \sum_{i\neq z}\omega_i^2(0)x_i^2}{2\lambda_x\lambda_y\lambda_z}
    + \mathcal{E}'(n) 
    + V(\vect{X}) + \frac{m\omega_z^2(t)Z^2}{2}
    \right)\Phi(\vect{x}, t).
\end{gather*}
I.e. we can just keep the harmonic $Z$ potential in $V(\vect{X})$, but we should still
include the scaling for the potential.
:::

:::{admonition} Implementation

Our general implementation consists of the following components:
* {py:class}`gpe.bec.StateScaleBase`: This implements the dynamical rescaling of the
  coordinates through a user-provided arbitrary scaling $\vect{\lambda}(t)$ returned by
  `get_lambdas()`.  This provides a properly rescaled set of "lab" coordinates
  $\vect{X}$ through the method `get_xyz()` *(or properties `xyz` and `_xyz_` in the old code)*.
  
* {py:class}`gpe.bec.StateScaleHO`: Provides an version of `get_lambdas()` that solves
  the classical equations of motion for harmonically trapped clouds.
* {py:class}`gpe.tube`: Implements the NPSEQ and uses these scaling classes to implement
  the dr-GPE.
:::

Note that we can rewrite these in terms of $\omega_\perp = \sqrt{\omega_x\omega_y}$:

\begin{gather*}
  \ddot{\lambda}_z(t) =
  \frac{\omega_z^2(0)}{\lambda_\perp^2(t)\lambda_z(t)} 
  - \omega_z^2(t)\lambda_z(t),\\
  \lambda_\perp^2 = \lambda_x\lambda_y\\
  2\lambda_{\perp}\dot{\lambda}_{\perp} = \dot{\lambda}_x\lambda_y + \lambda_x\dot{\lambda}_y\\
  2\dot{\lambda}_{\perp}^2 + 2\lambda_{\perp}\ddot{\lambda}_{\perp} 
  = \ddot{\lambda}_x\lambda_y + \lambda_x\ddot{\lambda}_y + 2\dot{\lambda}_x\dot{\lambda}_y\\
  \ddot{\lambda}_x(t) =
  \frac{\omega_i^2(0)}{\lambda_i(t)\lambda_x(t)\lambda_y(t)\lambda_z(t)} 
  - \omega_i^2(t)\lambda_i.
\end{gather*}

```{code-cell}
import sympy;sympy.init_printing()
t = sympy.var('t', real=True)
Xs = sympy.var('X, Y, Z', real=True)
xs_ = sympy.var('x, y, z', real=True)
ws = sympy.var('omega_x, omega_y, omega_z', real=True)
Ve = sympy.Function('V_{ext}')
ls = [sympy.Function('_l')(t) for _l in sympy.var('lambda_x, lambda_y, lambda_z')]
xs = [_X/_l for _X, _l in zip(Xs, ls)]
m, h = sympy.var('m, hbar', positive=True)
theta = m/2/h * sum(_X**2*_l.diff(t)/_l for _X, _l in zip(Xs, ls))
Phi = sympy.Function('Phi')(*(xs + [t]))
s = sympy.sqrt(sympy.prod(ls))
Psi = sympy.exp(sympy.I*theta)*Phi/s
V = sum(m*_w**2*_X**2/2 for _w, _X in zip(ws, Xs)) + Ve(*Xs)
L = (sympy.I*h*Psi.diff(t) 
       - (-h**2*(sum(Psi.diff(_X, _X) for _X in Xs))/2/m 
          + V*Psi)
      )/Psi*Phi
#display(sympy.simplify(sympy.I*h*Psi.diff(t).doit()))
subs = {_X:_x*_l for _X, _x, _l in zip(Xs, xs_, ls)}
L = L.subs(subs)
xs = xs_

# Here is the Lagrangian for Phi
Phi = Phi.subs(subs)
L0 = (sympy.I*h*Phi.diff(t)
      - (-h**2*(sum(Phi.diff(_x, _x)/_l**2 for _x, _l in zip(xs, ls)))/2/m
         + Ve(*Xs)*Phi
        )
     ).subs(subs)

sympy.simplify((L-L0).doit())
#sympy.simplify((L).doit())
```

## + NPSEQ = dr-GPE

Including the {ref}`sec:NPSEQ`, we obtain the dr-GPE -- an effective 1D theory for
expanding clouds.  We express this in two forms.  The first excludes scaling along the
$z$ axis because the scaling transformation assumes a quadratic dispersion and we would
like a form that we can use with theories with modified dispersions along
the $z$ direction.  The second assumes quadratic dispersion in all directions, and
allows for dynamic rescaling along $z$ too.  Note: As per the discussion above, we solve
for full expansion in all coordinates, but simply include or exclude the scaling in $z$.

\begin{gather*}
  \Psi(\vect{X}, Z, t) = e^{\I\theta(\vect{X}, t)}
    \frac{\Phi(\vect{x}, Z, t)}
         {\sqrt{\lambda_x(t)\lambda_y(t)}},\\
  X = \lambda_x(t)x, \qquad
  Y = \lambda_y(t)y, \qquad
  Z = z,\\
  \theta = \frac{m}{2\hbar} \sum_{i=x,y} \frac{X_i^2\dot{\lambda}_i(t)}{\lambda_i(t)}, 
  \qquad
  \ddot{\lambda}_i(t) = \frac{\omega_i^2(0)}{\lambda_i(t)\lambda_x(t)\lambda_y(t)\lambda_z(t)} 
  - \omega_i^2(t)\lambda_i,\\
  \I\hbar\dot{\Phi}(\vect{x}, Z, t) = 
  \left(
    - \frac{\hbar^2}{2m}\left[
      \sum_{i=x,y}\frac{\hbar^2\nabla_{i}^2}{2m\lambda_i^2} + \nabla_Z^2
    \right]
    + \frac{m}{2}\frac{\omega_x^2(0)x^2 + \omega_y^2(0)y^2}
                      {\lambda_x\lambda_y\lambda_z}
    + \frac{g}{\lambda_x\lambda_y}\abs{\Phi}^2 
    + V(Z, t)
    \right)\Phi(\vect{x}, t)\\
    =
  \left(
    - \frac{\hbar^2}{2}\left[
      \sum_{i=x,y}\frac{\hbar^2\nabla_{i}^2}{2\tilde{m}_i} + \nabla_Z^2
    \right]
    + \frac{\tilde{m}}{2}\sum_{i=x,y}\tilde{\omega}_i^2(0)x_i^2
    + \tilde{g}\abs{\Phi}^2 + V(Z, t)
    \right)\Phi(\vect{x}, t).
\end{gather*}




Using only time-dependent scaling only for $x$ and $z$ to obtain the following, where we
assume $\omega_x = \omega_y = \omega_\perp$:
%
\begin{gather*}
  \Psi(\vect{X}, Z, t) = e^{\I\theta(\vect{X}, t)}\frac{\Phi(\vect{x}, Z, t)}{\sqrt{\lambda_x(t)\lambda_y(t)}},
  \qquad
  X_i = \lambda_i(t)x_i, \\
    \theta = \frac{m}{2\hbar} \sum_{i=x,y} \frac{X_i^2\dot{\lambda}_i(t)}{\lambda_i(t)}, 
  \qquad
  \ddot{\lambda}_i(t) = \frac{\omega_i^2(0)}{\lambda_i(t)\lambda_x(t)\lambda_y(t)} 
  - \omega_i^2(t)\lambda_i,\\
  \I\hbar\dot{\Phi}(\vect{x}, Z, t) = 
  \left(
    - \frac{\hbar^2\nabla_{\vect{x}}^2}{2m\lambda_\perp^2} 
    - \frac{-\hbar^2\nabla_Z^2}{2m}     
    + \frac{m}{2}\frac{\omega_\perp^2(0)(x^2 + y^2)}{\lambda_\perp^2}
    + \frac{g}{\lambda_\perp^2}\abs{\Phi}^2 + V(Z, t)
    \right)\Phi(\vect{x}, t)\\
    =
  \left(
    - \frac{\hbar^2\nabla_{\vect{x}}^2}{2\tilde{m}}
    - \frac{-\hbar^2\nabla_Z^2}{2m}
    + \frac{\tilde{m}}{2}\tilde{\omega}_\perp^2(0)(x^2 + y^2)
    + \tilde{g}\abs{\Phi}^2 + V(Z, t
    \right)\Phi(\vect{x}, t).
\end{gather*}

We match this to the base form of the effective 1D model by rescaling $\tilde{m} =
m\lambda_\perp^2$, $\tilde{g} = g/\lambda_\perp^2$ and $\tilde{\omega}_\perp =
\omega_\perp/\lambda_\perp^2$ to obtain the following effecting 1D model for $\Phi(Z,
t)$:
\begin{gather*}
  \I\hbar \dot{\Phi}(Z, t) = 
  \left(
    \frac{-\hbar^2\nabla_Z^2}{2m}+ V(Z, t)
    + \frac{\hbar^2}{2\tilde{m}\sigma^2}
    + \frac{\tilde{m}\tilde{\omega}_\perp^2\sigma^2}{2}
    + \frac{\tilde{g}\abs{\Phi}^2}{2\pi\sigma^2}
    \right)\Phi,\\
  \frac{\tilde{m}\tilde{\omega}_\perp^2}{2}\sigma^4 = \frac{\tilde{g}\abs{\Phi}^2}{4\pi}+\frac{\hbar^2}{2\tilde{m}},\\
  \I\hbar \dot{\Phi}(Z, t) = 
  \left(
    \frac{-\hbar^2\nabla_Z^2}{2m} + V(Z, t)
    + \frac{\hbar^2}{2m\lambda_\perp^2\sigma^2}
    + \frac{m\omega_\perp^2\sigma^2}{2\lambda_\perp^2}
    + \frac{g\abs{\Phi}^2}{2\pi\lambda_\perp^2\sigma^2}
    \right)\Phi,\\
  \sigma^2 = \frac{\hbar}{m\omega_\perp}\sqrt{1 + \frac{mg\abs{\Phi}^2}{2\pi\hbar^2}},\\
\end{gather*}

## 3D Dynamics

Here we explore a test case where all three oscillator potentials change.  This was used
to develop an extension that does the dr-GPE with dynamical scaling along $x$ too.

```{code-cell}
%matplotlib inline
import numpy as np, matplotlib.pyplot as plt
from importlib import reload
from gpe import utils, bec, tube, minimize
from pytimeode.evolvers import EvolverABM
reload(bec)
reload(tube)

class Mixin:
    ws = np.array([1.0, 1.0, 1.0])
    ws_factors = np.array([1.2, 1.5, 0.7])

    def get_ws(self, t=None):
        if t is None:
            t = self.t
            
        if t <= 0:
            return self.ws
        else:
            return self.ws * self.ws_factors    
    
class State(Mixin, bec.StateBase):
    """Full 3D"""
    hbar = m = 1.0
    n0 = 1.0
    healing_length = 0.5
    Rxyz = np.array([3.0, 2.0, 2.0])
    dx = 0.25
    
    def __init__(self, dim=3, trap_T=1.0, **kw):
        for key in kw:
            if not hasattr(self, key):
                raise ValueError(f"Unknown {key=}")
            setattr(self, key, kw[key])
        
        mu = self.hbar**2/2/self.m/self.healing_length**2
        g = mu / self.n0
        Lxyz = 3.0*self.Rxyz
        Nxyz = list(map(utils.get_good_N, Lxyz / self.dx))
        args = dict(g=g, m=self.m, Lxyz=Lxyz[:dim], Nxyz=Nxyz[:dim],
                    mu=mu,
                   )
        super().__init__(**args)
        self.init()

    def get_density_1d(self, axis=0):
        """Return the integrated density the specified axis."""
        dxyz = np.divide(self.basis.Lxyz, self.basis.Nxyz)
        n = self.get_density()
        for ax in range(self.dim):
            if ax != axis:
                n = dxyz[ax] * n.sum(axis=ax, keepdims=True)
        return np.squeeze(n)
        
    def init(self):
        self.ws = np.sqrt(2*self.mu/self.m)/self.Rxyz
        super().init()

    def _get_Vext_(self):
        V_ext = 0.5*self.m * sum((_w*_x)**2
                                 for _w, _x in zip(self.get_ws(), self.xyz))
        if self.t <= 0:
            # For initial state
            V_ext -= self.mu
        return V_ext

    
class State1(Mixin, tube.StateGPEdrZ):
    """dr-GPE in y and z, but full x dynamics."""
    def __init__(self, state):
        args = dict(Nxyz=state.basis.Nxyz[:1],
                    Lxyz=state.basis.Lxyz[:1],
                    ws=state.ws,
                    g=state.g,
                    m=state.m,
                    mu=state.mu
                   )
        super().__init__(**args)

     
        
def get_states(use_scipy=True):
    """Returns two states `(s, s_tube)` for comparison:
    
    Returns
    -------
    s : Full 3D state.
    s_tube : 1D state but in the dr-GPE-Z approximation where the cross-section
       of the cloud is modelled as a Gaussian.
    """
    s = minimize.MinimizeState(State(), fix_N=False).minimize(use_scipy=use_scipy)
    Nx, Ny, Nz = s.basis.Nxyz
    n_central = s.get_density()[:, Ny//2, Nz//2]
    
    s_tube = State1(state=s)
    s_tube[...] *= np.sqrt(s.get_N()/s_tube.get_N())
    s_tube = minimize.MinimizeState(s_tube, fix_N=True).minimize(use_scipy=use_scipy)
    return s, s_tube
```

```{code-cell}
s0 = State(healing_length=0.5)
print(s0.get_density().max())
s = minimize.MinimizeState(s0, fix_N=False).minimize(use_scipy=True)
print(s.get_density().max())
```

```{code-cell}
s3, s1 = get_states()
plt.plot(s3.xyz[0].ravel(), s3.get_density_1d())
plt.plot(s1.xyz[0].ravel(), s1.get_density())
```

```{code-cell}
# from tqdm import tqdm
from mmfutils.contexts import FPS
states = get_states()
#s = minimize.MinimizeState(State(), fix_N=False).minimize(use_scipy=True)
#print(s.shape)
evs = [EvolverABM(s, dt=0.1*s.t_scale) for s in states]
ys = [states]
for frame in FPS(tqdm(range(10)), timeout=100):
    for ev in evs:
        ev.evolve(100)
    ys.append([ev.get_y() for ev in evs])
```

```{code-cell}
s = ys[0][1]
plt.plot(s.get_density())
```

```{code-cell}
n = s.get_density()
n.sum(axis=1, keepdims=True).shape
```

```{code-cell}
s.x_lab
```

```{code-cell}
s.basis.Lxyz
def get_n1d(s, axis):
    dxyz = np.divide(s.basis.Lxyz, s.basis.Nxyz)
    n = s.get_density()
    for ax in range(s.dim):
        if ax == axis:
            continue
        n = dxyz[ax] * n.sum(axis=ax, keepdims=True)
    return np.squeeze(n)
```

```{code-cell}
fig, ax = plt.subplots()
for n, label in enumerate(['x', 'y', 'z']):
    ax.plot(*np.transpose([(s.t, (s.xyz[n]*get_n1d(s, axis=n)).std()) for (s, s2) in ys]),
            label=label)

for n, label in enumerate(['x']):
    ax.plot(*np.transpose([(s.t, (s.x_lab*s.get_density()).std()) for (s1, s) in ys]),
            label="tube")
ax.legend()
```

```{code-cell}
s1, s2 = ys[0]
print(s1.get_N(), s2.get_N())
axis = 0
plt.plot(s1.xyz[axis].ravel(), get_n1d(s1, axis))

plt.plot(s2.x, s2.get_density())
```

## Incomplete/Broken
Here we add some dynamics in the form of a sinusoidally varying $\omega(t)$.

```{code-cell}
s1 = State1(state=s)
s1[...] *= np.sqrt(s.get_N()/s1.get_N())
m = minimize.MinimizeState(s1, fix_N=True)
s1 = m.minimize()
```

```{code-cell}
assert np.allclose(s.get_N(), s1.get_N())
```

```{code-cell}
# Here is the 1D density across the cloud.
dxyz = s.basis.Lxyz/s.basis.Nxyz
plt.plot(s0.basis.xyz[0], s0.get_density(), label='1D')
plt.plot(s.basis.xyz[0].ravel(), 
         s.get_density().sum(axis=-1).sum(axis=-1)*np.prod(dxyz[1:]), 
         label='3D')
plt.plot(s1.basis.xyz[0], s1.get_density(), label='tube')
plt.ylabel('1D integrated density')
plt.legend(loc='best')
```

```{code-cell}
# Now compare the central densities.
# This works well for the tube if the cross-section is small 
# compared with the healing length.
plt.plot(s0.basis.xyz[0], s0.get_density(), label='1D')
plt.plot(s.basis.xyz[0].ravel(), n_central, label='3D')
plt.plot(s1.basis.xyz[0].ravel(), s1.get_central_density(), label='tube')
plt.legend(loc='best')
plt.ylabel('central density')
```

```{code-cell}
Nx, Ny, Nz = s.Nxyz
x, y, z = s.xyz
n = s.get_density()
sigma2 = s1.get_sigma2()
plt.plot(x.ravel(), 
         ((n*(y**2 + z**2)).reshape((Nx, Ny*Nz)).mean(axis=1) /
          (n).reshape((Nx, Ny*Nz)).mean(axis=1)))
plt.plot(s1.xyz[0].ravel(), sigma2)
```

# GPE

+++

In the following, we consider various effective GPE-like equations of the form:

\begin{gather*}
  \I\hbar\dot{\Phi} = \op{H}[\Phi]\cdot\Phi = \frac{\delta E[\Phi]}{\Phi^\dagger}.
\end{gather*}

It is important in the code that this relationship between a "energy" functional $E[\Phi]$ and the equation of motion be maintained for minimization utilities to work properly.  In the case of the standard GPE we have:

\begin{gather*}
  E[\Psi] = \int\left(
    \frac{\I\hbar \Psi^\dagger\dot{\Psi} + \text{h.c.}}{2}
    + \frac{\hbar^2\abs{\nabla\Psi}^2}{2m}
    + V(\vect{x}, t)n
    + \frac{g}{2}n^2
  \right)\d^3{x}, \qquad
  n = \abs{\Psi}^2.
\end{gather*}

We can also consider a more general non-linear interaction $\mathcal{E}(n)$ where for the GPE this has the form $\mathcal{E}(n) = gn^2/2$.  The equations of motion will then have $\epsilon(n) = \mathcal{E}'(n)$ as the non-linear interaction.

+++

# Harmonic Traps

+++

Let $\vect{X} = (X, Y, Z)$ be the physical coordinates.  In a time-varying harmonic
potential, {cite}`Castin:1996` provide the following alternative description in terms of
scaled coordinates: 

\begin{gather*}
  x_i = \frac{X_i}{\lambda_i(t)}
\end{gather*}

using the rescaled function $\Phi(\vect{x}, t)$:

\begin{gather*}
  \Psi(\vect{X}, t) = e^{\I\theta}
  \frac{\Phi(\vect{x}, t)}{\sqrt{\lambda_x(t)\lambda_y(t)\lambda_z(t)}}, \qquad
  \theta = m\sum_{i}\frac{X_i^2}{2\hbar}\diff{}{t}\ln\lambda_i(t).
\end{gather*}

Let $s = \sqrt{\lambda_x(t)\lambda_y(t)\lambda_z(t)}$, then

\begin{gather*}
  \dot{s} = \frac{s}{2}\sum_i\diff{}{t}\ln\lambda_i(t), \qquad
  \partial_t \frac{1}{s} = -\frac{1}{2s}\sum_i\diff{}{t}\ln\lambda_i(t), \qquad
\end{gather*}

\begin{gather*}
  \dot{\Psi}(\vect{X}, t) = 
  e^{\I\theta}\frac{
    \dot{\Phi}(\vect{x}, t) 
    + \sum_i (x_i\partial_t\ln \lambda_i) \nabla_{i}\Phi
  }{\sqrt{\lambda_x(t)\lambda_y(t)\lambda_z(t)}}
  +
  \sum_{i}\left(
    \I m\frac{X_i^2}{2\hbar}\diff[2]{}{t}\ln\lambda_i(t)
    - \frac{1}{2}\diff{}{t}\ln\lambda_i(t)
  \right)\Psi\\
  \nabla_{\vect{X}}\Psi(\vect{X}, t) = 
  e^{\I\theta}\frac{\nabla_{\vect{X}}\Phi(\vect{x}, t)}{\sqrt{\lambda_x(t)\lambda_y(t)\lambda_z(t)}}
  + \I m\frac{X_i}{\hbar}\diff{}{t}\ln\lambda_i(t)\Psi.
\end{gather*}

The energy is 

\begin{gather*}
  E[\Phi] = \int\left(
    \frac{\I\hbar \Psi^\dagger\dot{\Psi} + \text{h.c.}}{2}
    + \frac{\hbar^2\abs{\nabla\Psi}^2}{2m}
    + V(\vect{X}, t)n
    + \mathcal{E}(n)
  \right)\d^3{X}, \qquad
  n = \frac{\abs{\Phi}^2}{\lambda_x(t)\lambda_y(t)\lambda_z(t)}
\end{gather*}

+++

Consider the following GPE in an axially confining harmonic trap:

\begin{gather*}
  \I\hbar\dot{\Psi} = \left\{\frac{-\hbar^2(\nabla_X^2+\nabla_Y^2)}{2m} + E(p_Z) + V(Z, t) + \frac{m}{2}\left(\omega_x^2(t)X^2 + \omega_y^2(t)Y^2\right)
  +\epsilon(\abs\Psi^2)
  \right\}\Psi.
\end{gather*}

Now consider the following transformation:

\begin{gather*}
  \Psi(X, Y, Z; t) = e^{\I\theta}\frac{\Phi(x, y, Z;t)}{\sqrt{\lambda_x(t)\lambda_y(t)}}, \qquad
  x = \frac{X}{\lambda_x(t)}, \qquad
  y = \frac{X}{\lambda_x(t)}.
\end{gather*}

The resulting equations of motion for $\Phi(x, y, Z;t)$ can be expressed in one of the following two forms:

\begin{gather*}
  \I\hbar\dot{\Phi} = \left\{
    \frac{-\hbar^2(\nabla_X^2+\nabla_Y^2)}{2m} + E(\op{p}_Z) + V(Z, t)
    +\epsilon(\abs\Psi^2)
  \right\}\Phi, \qquad
  \ddot{\lambda}_j = - \lambda_j \omega_i^2(t), \qquad
  \lambda_i(0) = 1, \qquad
  \dot{\lambda}_j(0) = 0.
\end{gather*}

This form, however, is not of particular use.  Consider for example an expanding cloud with $\omega_j(t) = 0$ for $t>0$.  The resulting theory simply has no scaling, $\lambda_j(t) = 1$ and degenerates to the raw GPE.

Instead, we keep the $t=0$ potentials, applying the scaling as follows:

\begin{gather*}
  \I\hbar\dot{\Phi} = \left\{
    \frac{-\hbar^2(\nabla_X^2+\nabla_Y^2)}{2m} + E(\op{p}_Z) + V(Z, t)
    +\epsilon(\abs\Psi^2)
    + \frac{1}{\lambda_x(t)\lambda_y(t)}
      \frac{m}{2}\left(\omega_x^2(0)x^2 + m\omega_y^2(0)y^2\right)
  \right\}\Phi, \qquad
  \ddot{\lambda}_j = - \lambda_j \omega_i^2(t) + \frac{\omega_j^2(0)}{\lambda_j(t)\lambda_1(t)\lambda_2(t)}, \qquad
  \lambda_i(0) = 1, \qquad
  \dot{\lambda}_j(0) = 0.
\end{gather*}

This is similar to the form in {cite}`Castin:1996`, but does not include the dynamics along the $z$ axis in the equation.  The reason is that the scaling transform depends on the dispersion being quadratic and we would like to explore modified dispersions.  Thus, the dynamics along $Z$ remain unscaled and explicitly managed.

+++

## Dynamics Traps

See the discussion in {ref}`sec:CoordinateTransformations`.


## Separable Ansatz

+++

The other portion of the analysis is to factorize the wavefunction as

\begin{gather*}
  \Psi(x, y, z; t) = \phi(x, y, z, t)\psi(z, t), \qquad
  \int \abs{\phi}^2(x, y, z, t)\;\d{x}\d{y} = 1.
\end{gather*}

The normalization condition gives the interpretation of $\abs{\psi}^2(z,t)$ as the density of particles per unit length along $z$.

Noting that the GPE derives from a principle of minimal action:

\begin{gather*}
  S = \int \Psi^\dagger \left(
  \I\hbar\partial_t 
  - \frac{\hbar^2\overleftarrow\nabla\overrightarrow\nabla}{2m}
  - V_\perp(x, y) - V_z(z) - \frac{g}{2}\abs{\Psi}^2\right)\d{x}\d{y}\d{z},
\end{gather*}

we obtain the following equation for $\psi(z, t)$:

\begin{gather*}
  \newcommand{\Braket}[1]{\left\langle{#1}\right\rangle}
  \I\hbar\dot{\psi} = \left(
  - \frac{\hbar^2\nabla_z^2}{2m}
  + V_{1D}(z) + g_{1D} \abs{\psi}^2\right)\psi\\
  \braket{A} = \frac{\int \phi^\dagger A\phi\; \d{x}\d{y}}
                    {\int \abs{\phi}^2\;\d{x}\d{y}}, \qquad
  V_{1D}(z) = V_z(z) + \Braket{
    -\I\partial_t 
    - \frac{\hbar^2\nabla^2}{2m}
    + V_\perp(x, y)}, \qquad
  g_{1D} = \braket{\abs{\phi}^2}.
\end{gather*}

The normalization condition for $\phi$ ensures that the linear term $C\nabla_z$ with $C\propto\braket{\nabla_Z} = 0$ vanishes.

+++

Likewise, minimizing wrt $\phi$ gives:

\begin{gather*}
  \I\hbar \dot{\phi} = \left(
    -\frac{\hbar^2\nabla^2}{2m} + V_\perp(x, y) + V_z(z) - \mu + g\abs{\Psi}^2
    +\frac{\hbar^2\abs{\nabla_z\psi}^2}{2m\abs{\psi}^2}
    \right)\phi,\\
  V_{2D}(x, y, z) = V_\perp(x, y) + V_z(z).
\end{gather*}

+++

So far this is exact.  To proceed, one assumes that the last term above is small, and that $\phi$ instantly adjusts so that the time-dependence can be dropped.

+++

Proceeding further, we factor the wavefunction as:

\begin{gather*}
  \Phi(x, y, Z;t) = \phi(x, y, t; Z)\psi(Z, t)
\end{gather*}

and make the "slowly varying approximation" that the dependence of $\phi$ on $Z$ is slowly varying so that we may ignore derivatives.  This allows $\phi$ to commute with the term $E(\op{p}_Z)$ so that the equations separate:

\begin{gather*}
  \I\hbar\frac{\dot{\phi}}{\phi} =
    \frac{\frac{-\hbar^2(\nabla_X^2+\nabla_Y^2)}{2m}\phi}{\phi} + \epsilon(\abs\Psi^2) 
    + \frac{E(p_Z)\psi}{\psi} + V(Z, t) - \I\hbar \frac{\dot{\psi}}{\psi}.
\end{gather*}

If we make the ansatz 

\begin{gather*}
  \phi\bigl(x, y; \sigma(Z, t)\bigr) = \frac{1}{\sqrt{\pi}\sigma(Z, t)}e^{-(x^2+y^2)/2\sigma^2}, \\
  \frac{\frac{-\hbar^2(\nabla_X^2+\nabla_Y^2)}{2m}\phi}{\phi} = \frac{\hbar^2}{2m}\frac{2\sigma^2-x^2-y^2}{\sigma^4}
\end{gather*}

+++

## dr-GPE

+++

Here we consider the problem of simulating an expanded cloud with a 1D simulation.  The goal is to treat the transverse directions in some variational manner in order to derive an effective 1D model that can capture the evolution of a quasi-1D cloud as it undergoes transverse expansion.

We expect the result to be an effective model where the density scales by some time-dependent factor accounting for the transverse expansion.  Here is the model from {cite}`Massignan:2003` that they call the dynamically-rescalled GPE (dr-GPE):

In their notations:

\begin{gather*}
  \I\hbar \dot{\tilde{\psi}} = \left\{
    -\frac{\hbar^2\nabla^2_z}{2m}\frac{1}{\lambda_z^2}
    + V(\lambda_z z, t)
    - \tilde{\mu} + \tilde{g}\abs{\tilde{\psi}}^2
    + \frac{m\tilde{\omega}_z^2 z^2}{2}
  \right\}\tilde{\psi}, \qquad
  \int \abs{\tilde{\psi}}^2\d{z} = 1,\\
  \tilde{\mu} = -\frac{\hbar^2}{2m\sigma^2\lambda_\perp^2}
  - \frac{m\omega_\perp^2(0)\sigma^2}{2\lambda_z\lambda_\perp^2}, \qquad
  \tilde{g} = \frac{gN}{2\pi\sigma^2\lambda_z\lambda_\perp^2}, \qquad
  g = \frac{4\pi\hbar^2 a}{m}, \qquad
  \tilde{\omega}_z = \frac{\omega_z(0)}{\sqrt{\lambda_z\lambda_\perp^2}},\qquad
  \sigma = a_{\perp}\sqrt[4]{\lambda_z + 2aN\abs{\tilde{\psi}}^2},\\
  \ddot{\lambda}_j(t) = \frac{\omega_j^2(0)}{\lambda_j(t)\lambda_1(t)\lambda_2(t)\lambda_3(t)}
  - \omega_j^2(t)\lambda_j(t), \qquad
  \vect{\lambda}(0) = \vect{1}, \qquad
  \dot{\vect{\lambda}}(0) = \vect{0}.
\end{gather*}

Note that their wavefunction is related to the GPE wavefunction $\Psi$ with coordinates $\vect{r} = (X, Y, Z)$ as follows:

\begin{gather*}
  \Psi(\vect{r}, t) = e^{\I\theta} \frac{\tilde{\phi}\bigl(X/\lambda_1, Y/\lambda_2;\sigma(Z/\lambda_3, t)\bigr)\tilde{\psi}(Z/\lambda_3, t)}
  {\sqrt{\lambda_1\lambda_2\lambda_3}},
  \qquad
  \theta = \frac{m}{2\hbar} \sum_i \frac{r_i^2\dot{\lambda}_i(t)}{\lambda_i(t)}
  \qquad
  \tilde{\phi}\bigl(x, y;\sigma\bigr) = \frac{e^{-(x^2+y^2)/2\sigma^2}}{\sigma \sqrt{\pi}},\\
  \int\abs{\tilde{\phi}}^2\d{x}\d{y} = 1.
\end{gather*}

+++

The central density is thus:

\begin{gather*}
  n(x=y=0) = \frac{N\abs{\tilde{\psi}}^2}{\pi\sigma^2\lambda_1\lambda_2\lambda_3}.
\end{gather*}

+++

### dr-GPE-Z

To enable us to work with spin-orbit coupled BECs, we reformulate the problem without
scaling in the $z=Z$ direction.  We shall simulate the expansion along axis directly,
obtaining the following equations: 

\begin{gather*}
  \I\hbar \dot{\psi} = \left\{
    -\frac{\hbar^2\nabla^2_z}{2m}
    + V(z, t) + \frac{m\omega_z^2 z^2}{2}
    + \frac{1}{\lambda_x\lambda_y}\left(
    \frac{\hbar^2}{2m\sigma^2}
    +\frac{m\omega_\perp^2(0)\sigma^2}{2}
    + \frac{g}{2\pi\sigma^2}\abs{\psi}^2
    \right)
  \right\}\psi, \qquad
  \int \abs{\psi}^2\d{z} = N,\\
  g = \frac{4\pi\hbar^2 a}{m}, \qquad
  \sigma^2 = \frac{\hbar}{m\omega_\perp}\sqrt{1 + 2a\abs{\psi}^2},\\
  \ddot{\lambda}_j(t) = \frac{\omega_j^2(0)}{\lambda_j(t)\lambda_1(t)\lambda_2(t)}
  - \omega_j^2(t)\lambda_j(t), \\
  \vect{\lambda}(0) = \vect{1}, \qquad
  \dot{\vect{\lambda}}(0) = \vect{0}.
\end{gather*}

To minimize these equations, we must derive an energy density that gives us this form.

\begin{gather*}
  \pdiff{\sigma^2}{\psi^\dagger} = \frac{aa_\perp^4}{\sigma^2}\psi, \qquad
  \pdiff{\sigma^4}{\psi^\dagger} = 2aa_\perp^4 \psi, \qquad
  \pdiff{\sigma^6}{\psi^\dagger} = 3aa_\perp^4\sigma^2\psi.  
\end{gather*}

To express this as the derivative of an energy, we eliminate the non-linear piece: 

\begin{gather*}
  2a\abs{\psi}^2 = \frac{m^2\omega_\perp^2 \sigma^4}{\hbar^2} - 1, \qquad
  \frac{g}{2\pi \sigma^2}\abs{\psi}^2 = m\omega_\perp^2 \sigma^2 - \frac{\hbar^2}{m\sigma^2}\\
\end{gather*}

\begin{gather*}
  \I\hbar \dot{\psi} = \left\{
    -\frac{\hbar^2\nabla^2_Z}{2m}
    + V(Z, t) + \frac{m\omega_z^2 z^2}{2}
    + \frac{1}{\lambda_x\lambda_y}\left(
    \frac{3m\omega_\perp^2(0)}{2}\sigma^2 
    - \frac{\hbar^2}{2m \sigma^2}
    \right)
  \right\}\psi
\end{gather*}

This follows variationally from the following energy density:

\begin{gather*}
   \mathcal{E} =
    \frac{\hbar^2\abs{\nabla_Z\psi}^2}{2m}
    + \left(V(Z, t) + \frac{m\omega_z^2 z^2}{2}\right)\abs{\psi}^2
    + \frac{1}{\lambda_x\lambda_y}\left(
      \frac{m\omega_\perp^2}{2aa_\perp^4}\sigma^6
      -
      \frac{\hbar^2}{2 m aa_\perp^4}\sigma^2
    \right)\\
    =
    \frac{\hbar^2\abs{\nabla_Z\psi}^2}{2m}
    + \left(V(Z, t) + \frac{m\omega_z^2 z^2}{2}\right)\abs{\psi}^2
    + \frac{1}{\lambda_x\lambda_y}\frac{m\omega_\perp^2}{2a}\left(
      \frac{\sigma^6}{a_\perp^4}
      - \sigma^2
    \right)
\end{gather*}

+++

Here we reformulate this to now include the factors of $\lambda_\perp$ in $\sigma$:

\begin{gather*}
  \I\hbar \dot{\psi} = \left\{
    -\frac{\hbar^2\nabla^2_z}{2m}
    + V(z, t) + \frac{m\omega_z^2 z^2}{2}
    + \frac{1}{\lambda_x\lambda_y}\left(
    \frac{\hbar^2\lambda_\perp}{2m\sigma^2}
    +\frac{m\omega_\perp^2(0)\sigma^2}{2\lambda_\perp}
    + \frac{g\lambda_\perp}{2\pi\sigma^2}\abs{\psi}^2
    \right)
  \right\}\psi, \qquad
  \int \abs{\psi}^2\d{z} = N,\\
  g = \frac{4\pi\hbar^2 a}{m}, \qquad
  \sigma^2 = \frac{\hbar \lambda_\perp }{m\omega_\perp}\sqrt{1 + 2a\abs{\psi}^2},\\
  \ddot{\lambda}_j(t) = \frac{\omega_j^2(0)}{\lambda_j(t)\lambda_1(t)\lambda_2(t)}
  - \omega_j^2(t)\lambda_j(t), \qquad
  \vect{\lambda}(0) = \vect{1}, \qquad
  \dot{\vect{\lambda}}(0) = \vect{0}.
\end{gather*}

To minimize these equations, we must derive an energy density that gives us this form.

\begin{gather*}
  \pdiff{\sigma^2}{\psi^\dagger} = \frac{aa_\perp^4}{\sigma^2}\psi, \qquad
  \pdiff{\sigma^4}{\psi^\dagger} = 2aa_\perp^4 \psi, \qquad
  \pdiff{\sigma^6}{\psi^\dagger} = 3aa_\perp^4\sigma^2\psi.  
\end{gather*}

To express this as the derivative of an energy, we eliminate the non-linear piece: 

\begin{gather*}
  2a\abs{\psi}^2 = \frac{m^2\omega_\perp^2 \sigma^4}{\hbar^2} - 1, \qquad
  \frac{g}{2\pi \sigma^2}\abs{\psi}^2 = m\omega_\perp^2 \sigma^2 - \frac{\hbar^2}{m\sigma^2}\\
\end{gather*}

\begin{gather*}
  \I\hbar \dot{\psi} = \left\{
    -\frac{\hbar^2\nabla^2_Z}{2m}
    + V(Z, t) + \frac{m\omega_z^2 z^2}{2}
    + \frac{1}{\lambda_x\lambda_y}\left(
    \frac{3m\omega_\perp^2(0)}{2}\sigma^2 
    - \frac{\hbar^2}{2m \sigma^2}
    \right)
  \right\}\psi
\end{gather*}

This follows variationally from the following energy density:

\begin{gather*}
   \mathcal{E} =
    \frac{\hbar^2\abs{\nabla_Z\psi}^2}{2m}
    + \left(V(Z, t) + \frac{m\omega_z^2 z^2}{2}\right)\abs{\psi}^2
    + \frac{1}{\lambda_x\lambda_y}\left(
      \frac{m\omega_\perp^2}{2aa_\perp^4}\sigma^6
      -
      \frac{\hbar^2}{2 m aa_\perp^4}\sigma^2
    \right)\\
    =
    \frac{\hbar^2\abs{\nabla_Z\psi}^2}{2m}
    + \left(V(Z, t) + \frac{m\omega_z^2 z^2}{2}\right)\abs{\psi}^2
    + \frac{1}{\lambda_x\lambda_y}\frac{m\omega_\perp^2}{2a}\left(
      \frac{\sigma^6}{a_\perp^4}
      - \sigma^2
    \right)
\end{gather*}

+++

In this formulation, the interpretation of $\abs{\psi}^2$ is the number density per unit length along the $z$ axis.  As above, this is related to the central density by:

\begin{gather*}
  n(x=y=0, z) = \frac{\abs{\psi}^2}{\pi \sigma^2 \lambda_1\lambda_2}.
\end{gather*}

```{code-cell}
%pylab inline --no-import-all
from gpe import tube;reload(tube)

s = tube.State(Nxyz=(2**9,), Lxyz=(10.0,))
m = minimize.MinimizeState(s)
print m.check()
s0 = m.minimize()
s0.plot()
```

## Static Trap

+++

If the radial trapping potential is constant, $\omega_j(t) = \omega_j(0)$, then $\lambda_j(t)=1$ and these equations reduce to the NPSE:

\begin{gather*}
  \I\hbar \dot{\psi} = \left\{
    -\frac{\hbar^2\nabla^2_z}{2m}
    + V(z, t)
    - \tilde{\mu} + \tilde{g}n
    + \frac{m\omega_z^2 z^2}{2}
  \right\}\psi, \\
  a_\perp = \sqrt{\frac{\hbar}{m\omega_\perp(0)}},\qquad
  \tilde{\mu} = -\frac{\hbar^2}{2m\sigma^2}
  - \frac{m\omega_\perp^2(0)\sigma^2}{2}, \qquad
  \tilde{g} = \frac{gN}{2\pi a_{\perp}^2\sqrt{1 + 2aN\abs{\psi}^2}}, \qquad
  \tilde{\omega}_z = \omega_z(0),\qquad
  \sigma = a_{\perp}\sqrt[4]{1 + 2aN\abs{\psi}^2},\\
\end{gather*}

+++

To compare with Salasnich, we match their units with $\hbar = m = a_\perp = 1$ and note that $\tilde{\psi} = \sqrt{N/2\pi}\psi$:

\begin{gather*}
  \I\dot{\psi} = \left\{
    -\frac{\nabla^2_z}{2}
    + V(z, t)
    + \frac{1}{2}\left(
      \frac{1}{\sqrt{1 + g\abs{\psi}^2}} + \sqrt{1 + g\abs{\psi}^2}
    \right)
    + \frac{g\abs{\psi}^2}{\sqrt{1 + g\abs{\psi}^2}}
  \right\}\psi.
  \tag*{(35) of arXiv/1504.07517v4}
\end{gather*}

Here $V(z, t)$ includes the harmonic trap.

```{code-cell}
%pylab inline
import numpy as np
from scipy.integrate import odeint

wx, wy, wz = 1.0, 2.0, 3.0

def ws(t):
    if t <= 0:
        res = [wx, wy, wz]
    else:
        res = [0,0,wz]
    return np.array(res)
    

def rhs(q, t):
    ls, dls = q.reshape(2, 3)
    ddls = ws(0)**2 / np.prod(ls) / ls - ws(t)**2*ls
    return np.ravel([dls, ddls])

q0 = np.ravel([[1,1,1], [0,0,0]])

ts = np.linspace(0, 10, 100)
res = odeint(rhs, q0, ts).T
#plt.plot(ts, res.T)
plt.plot(ts, res[4], ':')
```

In the case where all the traps are suddenly turned off, we have $\omega_j(t) = \omega_j\Theta(-t)$:

\begin{gather*}
  \frac{\ddot{\lambda}_j(t)\lambda_j(t)}{\omega_j} = \frac{1}{\lambda_1(t)\lambda_2(t)\lambda_3(t)}
\qquad
  \vect{\lambda}(0) = \vect{1}, \qquad
  \dot{\vect{\lambda}}(0) = \vect{0}.
\end{gather*}

+++

# Effective $g$ in Axial Confinement

+++

[Olshanii:1998] gives the relationship:

\begin{gather*}
  a_{1D} = \frac{-a_\perp^2}{2a}\left(1 - C\frac{a}{a_{\perp}}\right), \qquad
  a_\perp = \sqrt{\frac{\hbar}{m\omega_\perp}}.\\
  g_{1D} = \frac{\hbar \omega_\perp 2 a}{1-C\frac{a}{a_\perp}}
         = \frac{2\hbar^2}{m}\frac{a}{a_\perp^2}\frac{1}{1-C\frac{a}{a_\perp}}
         = \frac{-a_\perp^2\hbar\omega_\perp}{a_{1D}}
         = \frac{-\hbar^2}{ma_{1D}},\\
         C = -\zeta(1/2) = 1.46035450880958681288949915251529801246722933101258.
\end{gather*}


but see [Zhang:2014] for a correction with SOC.

[Olshanii:1998]: http://dx.doi.org/10.1103/PhysRevLett.81.938 (M. Olshanii, "Atomic Scattering in the Presence of an External Confinement and a Gas of Impenetrable Bosons", Phys. Rev. Lett. 81, 938--941 (1998) )

[Zhang:2014]: http://dx.doi.org/10.1038/srep04992 (Yi-Cai Zhang, Shu-Wei Song, and Wu-Ming Liu, "The confinement induced resonance in spin-orbit coupled cold atoms with Raman coupling", Sci. Rep. 4, 4992 (2014) )

+++

In the weak-confinement limit $a \ll a_\perp$ this becomes:

\begin{gather*}
  g_{1D} = 2 a \hbar \omega_\perp = \frac{2 \hbar^2}{m}\frac{a}{a_\perp^2}.
\end{gather*}

```{code-cell}
        zeta_1_2 = -1.46035450880958681288949915251529801246722933101258
        gs_1d = (a_perp * u.hbar * w_perp * 2 * a_s/a_perp
                 / (1 + a_s/a_perp * abs(zeta_1_2)))
```