---
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.17.3
kernelspec:
  name: python3
  display_name: Python 3 (ipykernel)
  language: python
---

```{code-cell} ipython3
: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:CoordinateTransformations)=
# Coordinate Transformations

:::{margin}
This includes the typical GPE as a special case when the homogeneous equation-of-state
(EoS) is $\mathcal{E}(n) = gn^2/2$ and the mean-field potential $\mathcal{E}'(n) =
gn$. 
:::
Here we consider solutions to time-dependent GPE-like theories of the form:
\begin{gather*}
  \I \hbar \dot{Ψ}(\vect{X}, t) = \left(
    \frac{-\hbar^2\nabla^2}{2m}
    + \mathcal{E}'\Bigl(\abs{Ψ(\vect{X}, t)}^2\Bigr)
    + \tilde{V}(\vect{X}, t)
  \right)Ψ(\vect{X}, t).
\end{gather*}

If the cloud under consideration is moving, rotating, or expanding, then one can often
gain significant computational advantage by transforming the coordinates from the **lab
frame** $\vect{X}$ to a **co-moving frame** $\vect{x}$.  In some cases, it also helps to
scale the wavefunction by a potentially time- and space-dependent factor $f$:
\begin{gather*}
  Ψ(\vect{X}, t) = f(\vect{X}, t)Φ(\vect{x}, t), \qquad
  \vect{x} = \vect{x}(\vect{X}, t).
\end{gather*}
Our goal here will be to find transformations that facilitate the numerical problem of finding
$Φ(\vect{x}, t)$.

## General Transformations

:::{margin}
Be careful here that $Φ_{,\mu}$ refers to partials with respect to the argument
$x^{\mu}$, while $Ψ_{,\mu}$ refers to partials with respect to the argument
$X^{\mu}$.
:::
Although we will not consider changing time, we can simplify our calculations of the
general transformations by adopting Einstein's notation, letting $X^{\mu} = (t,
\vect{X})$ etc. and denoting partials $\partial Φ/\partial t \equiv Φ_{,t}$.  We
then have
\begin{gather*}
  Ψ = fΦ, \qquad
  Ψ_{,α} = f_{,α}Φ + fΦ_{,μ}x^{μ}_{,α},\\
  Ψ_{,αβ} = f_{,αβ}Φ + f_{,α}Φ_{,μ}x^{μ}_{,β} 
             + f_{,β}Φ_{,μ}x^{μ}_{,α} 
             + fΦ_{,μν}x^{μ}_{,α}x^{ν}_{,β}
             + fΦ_{,μ}x^{μ}_{,αβ}.
\end{gather*}
:::{margin}
This form is specific for quadratic dispersion $K(\vect{\op{p}}) = \op{p}^2/2m$.  We
will try to point out when our results generalize.
:::
The transformed GPE is
\begin{gather*}
  \I\hbar (f_{,t}Φ + fΦ_{,μ}x^{μ}_{,t}) =
    \frac{-\hbar^2}{2m}
    \Biggl(
      f_{,ii}Φ + 2f_{,i}Φ_{,μ}x^{μ}_{,i}
               + f(Φ_{,μν}x^{μ}_{,i}x^{ν}_{,i} + Φ_{,μ}x^{μ}_{,ii})
    \Biggr)
    +  \left(\mathcal{E}'(n) + \tilde{V}\right)fΦ.
\end{gather*}

If we simplify by considering only transformation of space $x^{t}_{,t} = t_{,t} = 1$,
$x^{t}_{,i} = t_{,i} = 0$, we can re-express this as a modified GPE for $Φ$:
\begin{multline*}
  \I\hbar Φ_{,t}
  =
  \frac{-\hbar^2}{2m}\overbrace{x^{j}_{,i}x^{k}_{,i}Φ_{,jk}}^{g^{jk}\nabla_{j}\nabla_{k}\;Φ}
  %+ \\
  + \Biggl(\overbrace{
     x^{j}_{,t}
     +
     \frac{\hbar}{2m\I}
     \Bigl(
       2\frac{f_{,i}}{f}x^{j}_{,i}
       + 
       x^{j}_{,ii}
     \Bigr)}^{\vect{v}}
  \Biggr)
  \overbrace{(-\I\hbar Φ_{,j})}^{\op{\vect{p}}\;Φ} + \\
  + \Bigl(
    \underbrace{
      \mathcal{E}'(n) + \tilde{V}
      - \frac{\hbar^2}{2m}\frac{f_{,ii}}{f}
      - \I\hbar \frac{f_{,t}}{f}
    }_{\mathcal{E}'(n) + V}
  \Bigr)Φ,\qquad
  n = \abs{fΦ}^2.
\end{multline*}
The first term is a modified kinetic energy with metric $g^{jk} = x^{j}_{,i}x^{k}_{,i}$,
while the second term contains the momentum operator $\op{p} \equiv -\I\hbar
\vect{\nabla}_{x}$ multiplied by a generalized velocity $\vect{v}$, and the last term is
a corrected potential.

Alternatively, we can write $f(X^μ) = e^{\I Θ(X^{μ})}$ (but note that $Θ$ need not be
real):
\begin{gather*}
  Ψ = e^{\I Θ}Φ, \qquad
  e^{-\I Θ}Ψ_{,α} = \I Θ_{,α}Φ + Φ_{,μ}x^{μ}_{,α},\\
  e^{-\I Θ}Ψ_{,αβ} = (\I Θ_{,αβ} - Θ_{,α}Θ_{,β})Φ 
             + \I Θ_{,α}Φ_{,μ}x^{μ}_{,β} 
             + \I Θ_{,β}Φ_{,μ}x^{μ}_{,α} 
             + Φ_{,μν}x^{μ}_{,α}x^{ν}_{,β}
             + Φ_{,μ}x^{μ}_{,αβ}.
\end{gather*}
The transformed GPE is
\begin{gather*}
  \I\hbar (\I Θ_{,t}Φ + Φ_{,μ}x^{μ}_{,t}) =
    \frac{-\hbar^2}{2m}
    \Biggl(
      (\I Θ_{,ii} - Θ_{,i}Θ_{,i})Φ + 2\I Θ_{,i}Φ_{,μ}x^{μ}_{,i}
               + Φ_{,μν}x^{μ}_{,i}x^{ν}_{,i} + Φ_{,μ}x^{μ}_{,ii}
    \Biggr) + \\
    +  \left(\mathcal{E}'(n) + \tilde{V}\right)Φ.
\end{gather*}
Considering only transformation of space $x^{t}_{,t} = t_{,t} = 1$, $x^{t}_{,i} = t_{,i}
= 0$, we have the modified GPE for $Φ$:
\begin{multline*}
  \I\hbar Φ_{,t}
  =
  \frac{-\hbar^2}{2m}\overbrace{x^{j}_{,i}x^{k}_{,i}Φ_{,jk}}^{g^{jk}\nabla_{j}\nabla_{k}\;Φ}
  %+ \\
  + \Biggl(\overbrace{
      x^{j}_{,t}
      +
      \frac{\hbar}{m}
      Θ_{,i}x^{j}_{,i}
      +
      \frac{\hbar x^{j}_{,ii}}{2m\I}
     }^{\vect{v}}
  \Biggr)
  \overbrace{(-\I\hbar Φ_{,j})}^{\op{\vect{p}}\;Φ} + \\
  + \Bigl(
    \underbrace{
      \mathcal{E}'(n) + \tilde{V}
      - \frac{\hbar^2}{2m}(\I Θ_{,ii} - Θ_{,i}Θ_{,i})
      + \hbar Θ_{,t}
    }_{\mathcal{E}'(n) + V}
  \Bigr)Φ,\qquad
  n = \abs{e^{\I Θ}Φ}^2.
\end{multline*}

:::{margin}
Recall that
\begin{gather*}
  [\vect{a}\times]_{kj} = \epsilon_{ijk}a_i,\\
  [\vect{a}\times\vect{b}]_{k} = \epsilon_{ijk}a_ib_j,\\
  (\vect{a}\times\vect{b})\cdot\vect{c} = \epsilon_{ijk}a_ib_jc_k,\\
  \vect{\op{L}} = \vect{\op{x}}\times\vect{\op{p}},\\
\end{gather*}
:::
::::{admonition} Examples

**Translations**
\begin{gather*}
  \vect{x} = \vect{X} - \vect{x}_0(t), \qquad f = 1, \\
  x^{j}_{,i} = \delta^{j}_{i}, \qquad
  x^{j}_{,ii} = 0, \qquad
  x^{i}_{,t} = -v_0^{i} = -\dot{r}_0^{i},\\
  g^{ij} = \delta^{ij}, \qquad
  \vect{v} = -\vect{v}_0,
  \\
  \I\hbar \dot{Φ}
  =
  \Bigl(
    K(\op{p})
    -\vect{v}_0\cdot\op{\vect{p}}
    + \mathcal{E}'(n) + \tilde{V}
  \Bigr)Φ.
\end{gather*}
Note: Although we have only derived our transformations for quadratic dispersion
$K(\op{p}) = \op{p}^2/2m$, in this case, the kinetic energy does not transform, and
this remains valid for arbitrary dispersion.

**Rotations**
Consider a rotation about a fixed axis $\vect{\theta}$:
\begin{gather*}
  \vect{x} = \mat{R}\vect{X}, \qquad f = 1, \qquad
  [\ln\mat{R}]_{ij} = [\vect{\theta}\times]_{ij}
                    = \epsilon_{kji}[\vect{\theta}]_k,\\
  x^{j}_{,i} = [\mat{R}]_{ji}, \qquad
  x^{j}_{,ii} = 0, \qquad
  x^{i}_{,t} = \epsilon_{kji}[\vect{\omega}]_{k}[\vect{x}]_{j}
             = [\vect{\omega}\times \vect{x}]_{i},\qquad
  g^{ij} = [\mat{R}^T\mat{R}]^{ij} = \delta^{ij},\\
  \I\hbar Φ_{,t}
  =
  \Biggl(
    K(\op{p})
    +
    \underbrace{(\vect{\omega}\times\vect{x})\cdot\vect{\op{p}}}
              _{\vect{\omega}\cdot(\vect{\op{x}}\times\vect{\op{p}}) 
                = \vect{\omega}\cdot \vect{\op{L}}}
    +
    \mathcal{E}'(n) + \tilde{V}
  \Biggr)Φ.
\end{gather*}
I.e., to work in a rotating frame, we include an angular momentum term.

**Scaling**

We start with a simple scaling of the coordinates -- we will extend this discussion
below.
\begin{gather*}
  x^{j} = \frac{X^{j}}{λ_{j}(t)}, \qquad
  x^{j}_{,t} = -x^{j}\frac{\dot{λ}_{j}}{λ_{j}}, \qquad
  x^{j}_{,i} = \frac{\delta^{j}_{i}}{λ_{i}}, \qquad
  x^{i}_{,ii} = 0, \qquad f = 1,\\
  \I\hbar Φ_{,t}
  =
  \Biggl(
    K\Bigl(\frac{\op{\vect{p}}}{\vect{λ}}\Bigr)
    -
    \sum_{j}
    \frac{\dot{λ}_{j}}{λ_{j}}x^{j}\op{p}_j
    + 
   \mathcal{E}'(n) + \tilde{V}
  \Biggr)Φ,\qquad
  n = \abs{Φ}^2.
\end{gather*}
More typically, as we shall see below, we include a factor $f$ so that the density
scales properly, and to remove the term linear in $\op{p}$, but the specialization there
requires assuming that the dispersion is quadratic.
::::

## Translation

The simplest transformation corresponds to an accelerating frame:
\begin{gather*}
    \vect{x} = \vect{X} - \vect{x}_0(t).
\end{gather*}
The minimal transformation simply tracks this through the equations of motion:
\begin{gather*}
  Ψ(\vect{X}, t) = Φ(\vect{x}, t) = Φ\bigl(\vect{X} - \vect{x}_0(t), t\bigr).
\end{gather*}
The only modifications to the equation of motion are:
1. An addition contribution to the time dependence:
   \begin{gather*}
     \I\hbar \dot{Ψ} = \I\hbar \dot{Φ} 
       - \I\hbar \underbrace{\dot{\vect{x}}_0(t)}_{\vect{v}_0(t)} \cdot \vect{\nabla}Φ
     = \I\hbar \dot{Φ} + \vect{v}_0(t) \cdot \op{\vect{p}}Φ.
   \end{gather*}
2. The potential must be evaluated in the lab-frame: $V(\vect{X}) = V\bigl(\vect{x} -
   \vect{x}_0(t)\bigr)$.
   
Thus, in the scaled coordinates, we have
\begin{gather*}
  \I\hbar \dot{Φ}(\vect{x}, t) = \Bigl(
    K(\op{p}) - \vect{v}\cdot\op{\vect{p}} + V(\vect{X}) + \mathcal{E}'(n)
  \Bigr)Φ(\vect{x}, t).
\end{gather*}
This transformation can be accounted for by simply modifying the dispersion relationship.

:::{admonition} Where is the Acceleration?
:class: dropdown

If you are familiar with classical mechanics, you might expect a force proportional to
the acceleration $\vect{a}_0(t) = \ddot{\vect{x}}_0(t)$ contributing a linear piece to the
potential $m\vect{a}\cdot\vect{x}$.  This can be recovered by including an additional
phase
\begin{gather*}
  Ψ(\vect{X}, t) = e^{\I\theta(\vect{X}, t)}Φ(\vect{x}, t)
  = e^{\I\theta(\vect{X}, t)}Φ\bigl(\vect{X} - \vect{x}_0(t), t\bigr).
\end{gather*}
Including this phase gives additional pieces to the equations of motion:
\begin{gather*}
  e^{-\I\theta}\I\hbar \left.\pdiff{Ψ}{t}\right|_{X} = 
  \I\hbar \left.\pdiff{Φ}{t}\right|_{x}
  + \Biggl(
    \vect{v}_0\cdot\op{\vect{p}} -
    \hbar\overbrace{\left.\pdiff{\theta}{t}\right|_{X}}^{\dot{\theta}}
   \Biggr)Φ,\\
  e^{-\I\theta}K(\op{P})Ψ = K\Biggl(
    \op{p}+\hbar\underbrace{\left.\pdiff{\theta}{\vect{X}}\right|_{t}}_{\vect{\nabla}\theta}
  \Biggr)Φ.
\end{gather*}
Combining these we have
\begin{gather*}
  \I\hbar \dot{Φ}(\vect{x}, t) = \Bigl(
    K\bigl(\op{p}+\hbar\vect{\nabla}\theta\bigr) 
    - \vect{v}_0(t)\cdot\op{\vect{p}} 
    + \hbar \dot{\theta}
    + V(\vect{X}) + \mathcal{E}'(n)
  \Bigr)Φ(\vect{x}, t).
\end{gather*}
To recover the linear potential, we use the phase
\begin{gather*}
  \hbar\theta(\vect{X}, t) = m\vect{v}_0(t)\cdot\vect{X} 
  %=  m\vect{v}_0(t)\cdot\bigl(\vect{X} - \vect{x}_0(t)\bigr)
  ,\\
  \hbar\vect{\nabla}\theta = m\vect{v}_0(t), \qquad
  \hbar\dot{\theta} = m\dot{\vect{v}}_0(t)\cdot\vect{X} = m\vect{a}_0(t)\cdot\vect{X},\\
  \I\hbar \dot{Φ}(\vect{x}, t) = \Bigl(
    K\bigl(\op{p}+m\vect{v}_0(t)\bigr) 
    - \vect{v}_0(t)\cdot\op{\vect{p}} 
    + m\vect{a}_0(t)\cdot\vect{X}
    + V(\vect{X}) + \mathcal{E}'(n)
  \Bigr)Φ(\vect{x}, t).
\end{gather*}
This is more difficult to implement, but is necessary to display Galilean covariance.
To see this, consider a Galilean boost $\vect{x}_0(t) = \vect{v} t$ in a system with quadratic
dispersion $K(\op{P}) = P^2/2m$.  Galilean covariance becomes manifest if we use
the full transformation that includes the kinetic energy shift in the phase
\begin{gather*}
  \hbar\theta(\vect{X}, t) = m\vect{v}\cdot\vect{X} - \tfrac{m}{2}v^2 t 
  = m\vect{v}\cdot\vect{x} + \tfrac{m}{2}v^2 t,\\
  \hbar\vect{\nabla}\theta = m\vect{v}, \qquad
  \hbar\dot{\theta} = -\tfrac{m}{2}v^2,\\
  \begin{aligned}
  \I\hbar \dot{Φ}(\vect{x}, t) &= \Bigl(
    \frac{(\op{p}+m\vect{v})^2}{2m} 
    - \vect{v}\cdot\op{\vect{p}}
    - \tfrac{m}{2}v^2
    + V(\vect{X}) + \mathcal{E}'(n)
  \Bigr)Φ(\vect{x}, t),\\
  &= \Bigl(
    \frac{\op{p}^2}{2m} + V(\vect{X}) + \mathcal{E}'(n)
  \Bigr)Φ(\vect{x}, t).
  \end{aligned}
\end{gather*}
:::

## Scaling (Expansion)

If we need to scale the coordinates $X_i = λ_i(t)x_i$ where $\vect{X}$ is the
**lab frame** and $\vect{x}$ is the **co-expanding frame**, then we can introduce the
following wavefunction:
\begin{gather*}
  Ψ(\vect{X}, t) = \frac{e^{\I\theta(\vect{X}, t)}Φ(\vect{x}, t)}
                           {\sqrt{λ_x(t)λ_y(t)λ_z(t)}}, \qquad
  X_i = λ_i(t)x_i, \qquad
  \hbar\theta(\vect{X}, t) = \frac{m}{2} \sum_i \frac{X_i^2\dot{λ}_i(t)}{λ_i(t)}.
\end{gather*}
With this transformation, the following theories are equivalent:
\begin{gather*}
  \I\hbar\dot{Ψ}(\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)Ψ(\vect{X}, t), \\
  \I\hbar\dot{Φ}(\vect{x}, t) = 
  \left(
    \frac{-\hbar^2\nabla_{\vect{X}}^2}{2m} 
    + \frac{m \sum_{i}\bigl(
      \omega_{i}^{2}(t)λ_{i}^{2} + λ_{i}\ddot{λ}_{i}\bigr)x_i^2}{2}
    + \mathcal{E}'(n) + V(\vect{X})
    \right)Φ(\vect{x}, t).
\end{gather*}
The utility of this transformation is that the classical behavior of an expanding gas
can be modeled by setting:
\begin{gather*}
  \ddot{λ}_i(t) =
  \frac{\omega_i^2(0)}{λ_i(t)λ_x(t)λ_y(t)λ_z(t)} 
  - \omega_i^2(t)λ_i(t),
\end{gather*}
which gives:
\begin{gather*}
  \I\hbar\dot{Φ}(\vect{x}, t) = 
  \left(
    \frac{-\hbar^2\nabla_{\vect{X}}^2}{2m} 
    + \frac{1}{λ_x(t)λ_y(t)λ_z(t)}\frac{m}{2} \sum_{i}
      \omega_i^2(0) x_i^2
      + \mathcal{E}'(n) + V(\vect{X})
    \right)Φ(\vect{x}, t).
\end{gather*}
With this transformation, $Φ(\vect{x}, t)$ evolves minimally (see
{cite}`Castin:1996`).  Note that in the new coordinates, the effective trapping
potential is constant, up to the overall density scaling.
I.e., the time-dependence of the harmonic potential is completely absorbed into the
scaling factor.  This is crucial for
the {ref}`dr-GPE` since in this co-expanding frame, the effective trap remains, even if the
trap is completely turned off: $\omega_i(t>0) = 0$.

A similar expansion can be performed scaling only $X$ and $Y$:
\begin{gather*}
  \I\hbar\dot{Ψ} = \left\{\frac{-\hbar^2(\nabla_X^2+\nabla_Y^2)}{2m} + E(\op{p}_Z) + V(Z, t) + \frac{m}{2}\left(\omega_x^2(t)X^2 + \omega_y^2(t)Y^2\right)
  +\epsilon(\absΨ^2)
  \right\}Ψ, \\
  Ψ(X, Y, Z; t) = e^{\I\theta}\frac{Φ(x, y, Z;t)}{\sqrt{λ_x(t)λ_y(t)}}, \qquad
  x = \frac{X}{λ_x(t)}, \qquad
  y = \frac{Y}{λ_y(t)}.
\end{gather*}

The resulting equations of motion for $Φ(x, y, Z;t)$ can be expressed as:
\begin{gather*}
  \I\hbar\dot{Φ} = \left\{
    \frac{-\hbar^2(\nabla_X^2+\nabla_Y^2)}{2m} + E(\op{p}_Z) + V(Z, t)
    +\mathcal{E}'(n)
    +\tfrac{m}{2}\frac{\omega_x^2(0)x^2 + m\omega_y^2(0)y^2}{λ_x(t)λ_y(t)}
  \right\}Φ,\\
  \ddot{λ}_j = - λ_j \omega_j^2(t) + \frac{\omega_j^2(0)}{λ_j(t)λ_x(t)λ_y(t)}, \qquad
  λ_i(0) = 1, \qquad
  \dot{λ}_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.  A reason for using this is that the scaling
transform depends on the dispersion being quadratic and we would like to explore
modified dispersions (i.e. for spin-orbit coupled (SOC) BECs).  In this expression, the dynamics along
$Z$ remain unscaled and explicitly managed. 

:::{admonition} A failed attempt.

One might try to express the equations simply by completely removing the potentials:
\begin{gather*}
  \I\hbar\dot{Φ} = \left\{
    \frac{-\hbar^2(\nabla_X^2+\nabla_Y^2)}{2m} + E(\op{p}_Z) + V(Z, t)
    +\epsilon(\absΨ^2)
  \right\}Φ, \\
  \ddot{λ}_j = - λ_j \omega_i^2(t), \qquad
  λ_i(0) = 1, \qquad
  \dot{λ}_j(0) = 0.
\end{gather*}
This form, however, is of little use.  Consider for example an expanding cloud with
$\omega_j(t) = 0$ for $t>0$.  The resulting theory simply has no scaling, $λ_j(t)
= 1$ and degenerates to the raw GPE.  We have thrown out the baby with the bath water.
:::

### Derivation

The general derivation is a bit messy, so we start with 1D.  This will also allow us to
consider simultaneous scaling and translation.  Let $Ψ(X, t)$ be the physical
wavefunction satisfying the non-linear Schrödinger equation, and relate this to a
computational wavefunction $Φ(x, t)$ in the computational coordinates $x(X, t)$ with:
\begin{gather*}
  Ψ(X, t) = e^{\IΘ(X,t)}Φ(x(X, t), t).
\end{gather*}
Expanding the Schrödinger equation
\begin{gather*}
  \I\hbar \dot{Ψ} = \frac{-\hbar^2}{2m}\nabla_X^2Ψ,
\end{gather*}
and demanding that first-order spatial derivatives $\nabla_x Φ$ cancel *(for reasons of
computational efficiency)*, we obtain:
\begin{gather*}
  Θ'(X, t) = -\frac{m}{\hbar}\frac{\dot{x}}{x'} + \frac{\I}{2}\ln(x')',\\
  Θ(X, t) = - \frac{m}{\hbar}\int\frac{\dot{x}}{x'}\d{X}
  + \frac{\I}{2}\ln(x') + \I C(t),\\
  \dot{C}(t) = \frac{\dot{x}x''-x'\dot{x}'}{x'^2} = -\left(\frac{\dot{x}}{x'}\right)',
\end{gather*}
where primes denote derivatives with respect to $X$.
:::{note}
$Θ$ is not real here -- the last two terms of $Θ$ represents an overall scaling of the
wavefunction, and the integration constant $C(t)$ is chosen here so that the evolution
remains unitary.  *(An additional real integration constant can be introduced as a
chemical potential, but it will only affect the phase of the wavefunction so we drop
this.)*  As a consistency condition, $\dot{C}(t)$ must not depend on $X$, so we are
restricted to coordinate transformations that have the form $\dot{x} = [a(t)+b(t)X]x'$.
This will suffice for our purposes.  More general transformation will require that we
allow for linear derivatives.
:::

+++

We now specialize to a coordinate scaling and translation:

\begin{gather*}
  x(X, t) = \frac{X - ρ(t)}{λ(t)}, \qquad
  X = λ(t)x + ρ(t), \qquad
  Ψ(X, t) = \frac{e^{\I\theta(X, t)}}{\sqrt{λ(t)}}Φ(x(X, t), t), \\
  \theta(X, t)
  = \frac{m}{2\hbar}\left(
    \frac{\dot{λ}}{λ}X^2
    - 2λ\dot{ρ}X\right)
  = \frac{m}{2\hbar}\left(
    \frac{\dot{λ}}{λ}\left(
    X-\frac{λ^2\dot{ρ}}{\dot{λ}}\right)^2
    - \frac{\dot{ρ}^2λ^3}{\dot{λ}}
  \right)
\end{gather*}
*The scaling factor $1/\sqrt{λ}$ comes from the imaginary terms $C(t)$ above in $Θ$.*

This leads to a effective potential of the following form:

\begin{gather*}
  V_{λ, ρ}(X, t) = \frac{m}{2}\left(
    \frac{\ddot{λ}}{λ}X^2
    + 2\left(\ddot{ρ}-\frac{\ddot{λ}}{λ}ρ\right)X
    + \left(
    \dot{ρ} - ρ \frac{\dot{λ}}{λ}
  \right)^2
  \right).
\end{gather*}

+++

\begin{gather*}
  V_{λ, ρ}(X, t) = \frac{m}{2}\left(
    λ \ddot{λ} x^2
    + 2λ\ddot{ρ}x
    + c(t)
  \right),\\
  c(t) = \dot{ρ}^2
       + 2\ddot{ρ}ρ +
       + ρ^2\left(\frac{\dot{λ}}{λ}\right)^2
       - 2\dot{ρ}ρ \frac{\dot{λ}}{λ}
       - ρ^2 \frac{\ddot{λ}}{λ}
       = \dot{ρ}^2
       + 2\ddot{ρ}ρ
       - 2\dot{ρ}ρ \frac{\dot{λ}}{λ}
       - ρ^2(t)\diff[2]{\lnλ(t)}{t}
\end{gather*}

+++

One nice use of this is to take a problem that has a time-dependent external potential and replace it with a stationary potential.  Thus, for example, suppose we have:

\begin{gather*}
  V_{\text{ext}}(X, t) = \frac{m\omega^2(t)}{2}[X-R(t)]^2
\end{gather*}

in the lab frame.  The essence of Castin and Dumm's argument is that a good strategy is to choose $λ(t)$ and $R(t)$ such that the effective potential for $Φ$ is constant after scaling.  I.e., for the GPE $\mathcal{E}'(n) = gn$ and the equations of motion for $Φ$ have the form:

\begin{align}
  \I\hbar \dot{Φ} &= \left(
    \frac{-\hbar^2\nabla^2_{\vect{X}}}{2m}
    + V_{λ,ρ}+ V_{\text{ext}}
    + g \abs{Ψ}^2
  \right)Φ, \\
 \I\hbar \dot{Φ} &=
  \left(
    \frac{-\hbar^2\nabla^2_{\vect{X}}}{2m}
    + \frac{
      (λ_xλ_yλ_z)
      (V_{λ,ρ}+ V_{\text{ext}})
    + g\abs{Φ}^2
    }{λ_xλ_yλ_z}
  \right)Φ.
\end{align}
The change in $Φ$ is minimized if we choose
\begin{gather*}
  V_{λ,ρ}(t) = \frac{V_{\text{ext}}(0)}{λ_xλ_yλ_z} - V_{\text{ext}}(t).
\end{gather*}
I.e., if we started in a stationary state at time $t=0$ with external potential
$V_{\text{ext}}(0)$, then choosing such a transformation minimizes the effects of
dynamics due to the trap expansion. 

This gives the following differential equations for $λ_i(t)$ and $ρ_i(t)$:

\begin{gather*}
  \ddot{λ}_i(t) = \frac{\omega_i^2(t)}{λ_i λ_xλ_yλ_z} - \omega^2(t)λ_i, \qquad
  \ddot{ρ}_i(t) = \omega_i^2(t)[R_1(t) - ρ_i(t)].
\end{gather*}

As a simple check, consider a stationary solution in the frame where the coordinates
scale to new values so that at some late time $t$, $λ \neq 1$ but $\dot{λ} =
\ddot{λ} = 0$.  In the computational basis we should have:
\begin{gather*}
  Φ(x) = \sqrt{λ}Ψ(λ x)\\
  \frac{-\hbar^2Ψ''(X)}{2mΨ(X)} + g \abs{Ψ}^2 + V(X) = E,\\
  \frac{-\hbar^2Φ''(x)}{2mλ^2Φ(x)} + \frac{g}{λ} \abs{Φ}^2 + V(X) = E,\\
\end{gather*}

### Radial Expansion

A specific use-case is that of an axially symmetric trap in which we describe the
expansion in the radial direction.  This is implemented in
[gpe/axial.py](../gpe/axial.py).  As before, the physical coordinate $R$ is related to
the scaled coordinate $r$ as 
\begin{gather*}
  x = X, \qquad
  r = \frac{R}{λ_\perp(t)} = \sqrt{x^2+y^2}, \\
  Ψ(X, R, t) = e^{\I\theta}
  \frac{Φ(x, r, t)}{λ_\perp}, \qquad
  \theta = \frac{mR^2}{2\hbar}\diff{}{t}\lnλ_\perp(t).
\end{gather*}

We discretize the equations for the scaled function $Φ(x, r, t)$ on the scaled
coordinates $x$ and $r$:
\begin{gather*}
  \I\hbar\dot{Φ}(x, r ,t) = \left\{
    \frac{-\hbar^2}{2mλ_\perp^2}\left(
    \frac{1}{r}\pdiff{}{r}r\pdiff{}{r} + \frac{1}{r^2}\pdiff[2]{}{\phi}
  \right) + E(\op{p}_x) 
    + V(X, R, t)
    -\frac{m\omega_\perp^2(t)}{2} R^2
    +\mathcal{E}(n)
    +
    \tfrac{m}{2}\frac{\omega_\perp^2(0)r^2}{λ_\perp^2(t)}
  \right\}Φ(x, r ,t),\\
  \ddot{λ}_\perp(t) = - λ_\perp(t) \omega_\perp^2(t) 
  + \frac{\omega_\perp^2(0)}{λ_\perp^3(t)}, \qquad
  λ_\perp(0) = 1, \qquad
  \dot{λ}_\perp(0) = 0.
\end{gather*}
Note that the scaling $λ_\perp(t)$ depends only on the trapping frequency, so can
be applied to multiple components as long as the external trapping potential behaves
the same way for both.

## Test Examples

In the case of expansion $\omega_\perp(t>0) = 0$ we have:
\begin{gather*}
  \ddot{λ}_\perp(t) = \frac{\omega_\perp^2(0)}{λ_\perp^3(t)}, \qquad
  λ_\perp(0) = 1, \qquad
  \dot{λ}_\perp(0) = 0,\\
  λ_\perp = \sqrt{1+\tau^2}, \qquad \omega_\perp(0)t.
\end{gather*}

Here we explicitly code this example so we have something to work with later.
<!-- \begin{gather*}
  \ddot{λ}_\perp(t) = \frac{\omega_\perp^2(0)}{λ_\perp^3(t)}, \qquad
  λ_\perp(0) = 1, \qquad
  \dot{λ}_\perp(0) = 0.\\
  \int_0^{\dot{λ_\perp}}\dot{λ}_\perp\d{\dot{λ}_\perp} = \int_1^{λ_\perp}\frac{\omega_\perp^2(0)}{λ_\perp^3}\d{λ_\perp},\\
  \dot{λ}_\perp = \omega_\perp(0)\sqrt{1-λ_\perp^{-2}},\\
  \int_1^{λ_\perp}\frac{1}{\sqrt{1-λ_\perp^{-2}}}\d{λ_\perp} 
  = \sqrt{λ_\perp^2-1} = \tau = \omega_\perp(0)t, \\
  λ_\perp = \sqrt{1+\tau^2}, \qquad \omega_\perp(0)t.
\end{gather*} -->

```{code-cell} ipython3
from gpe.imports import *
from pytimeode import mixins, interfaces
from pytimeode.evolvers import EvolverABM
from mmfutils.math.bases import CylindricalBasis

@interfaces.implementer(interfaces.IStateForABMEvolvers)
class State0(mixins.ArrayStateMixin):
    def __init__(self, N=2**5, R=6):
        self.basis = CylindricalBasis(Nxr=(1, N), 
                                      Lxr=(1, R))
        R = self.R
        self.data = 5*np.exp(-R[None,:]**2/2) + 0j
        
    @property
    def R(self):
        return self.basis.xyz[1].ravel()
    
    @property
    def n(self):
        return abs(self[...].ravel())**2
    
    def compute_dy_dt(self, dy):
        y = self[...]
        dy[...] = self.basis.laplacian(y, factor=-1./2)
        dy[...] += abs(y)**2*y
        dy /= 1j
        return dy
    
    def plot(self):
        plt.plot(self.R, self.n)
        plt.title("t={:.4g}".format(self.t))


class State1(State0):
    @property
    def lam(self):
        return np.sqrt(1+self.t**2)
    
    @property
    def R(self):
        return self.basis.xyz[1].ravel() * self.lam
    
    @property
    def n(self):
        return abs(self[...].ravel())**2 / self.lam**2
    
    def compute_dy_dt(self, dy):
        x, r = self.basis.xyz
        y = self[...]
        dy[...] = self.basis.laplacian(y, factor=-1./2)
        dy[...] += (r**2/2.0 + abs(y)**2)*y
        dy /= self.lam**2
        dy /= 1j
        return dy
            
s0 = State0()
s1 = State1()

dt = 0.0001
steps = 200
e0 = EvolverABM(s0, dt=dt)
e1 = EvolverABM(s1, dt=dt)
#fig = plt.figure(figsize=(10,5))
with NoInterrupt() as interrupted:
    while not interrupted and e0.t < 1:
        e0.evolve(steps)
        e1.evolve(steps)

        plt.clf()
        e0.y.plot()
        e1.y.plot()
        display(plt.gcf())
        clear_output(wait=True)
```

Here we do the same thing with the `axial.State2Axial` class:

```{code-cell} ipython3
from gpe.imports import *
from gpe import bec2, axial;reload(bec2); reload(axial)
from gpe.axial import u
from gpe.minimize import MinimizeState

l_unit = 1.0
t_unit = 1.0
basis = axial.CylindricalBasis(Nxr=(64, 32), Lxr=(2*6*l_unit, 6*l_unit))


class State(axial.State2Axial):
    t_expand = .0*t_unit
    _ws = np.array([48.78*0, 48.78])*u.Hz
    _ws = np.array([0.0/u.Hz, 1.0/u.Hz])*u.Hz

    def get_ws(self, t=None):
        if t is None:
            t = self.t
        if t < self.t_expand:
            return self._ws
        else:
            #return 0*self._ws
            return np.concatenate([self._ws[:1], 0*self._ws[1:]])
        
    def get_ws_perp(self, t=None):
        return self.get_ws(t=t)[1:]        

    @property
    def ws(self):
        return self.get_ws()
        
    @ws.setter
    def ws(self, ws):
        pass

    def _get_Vext_(self):
        """Return (V_a, V_b, V_ab), the external potentials."""
        V_HO = (0.5*self.ms[self.bcast]
                * sum((_w*_x)**2 for _w, _x in zip(self.ws, self.xyz)))
        zero = np.zeros_like(V_HO[0])
        V_ab = self.Omega*np.exp(2j*self.k_r*self.xyz[0])
        mu_a = mu_b = 0.0
        if hasattr(self, 'mus') and self.mus:
            mu_a, mu_b = self.mus

        return (V_HO[0] - mu_a - self.delta/2.0,
                V_HO[1] - mu_b + self.delta/2.0,
                V_ab + zero)

class State_(State):
    def get_ws_perp(self, t=None):
        """Do not change these - no expansion in basis"""
        return self.get_ws(t=0)[1:]

args = dict(basis=basis, x_TF=4*u.micron, gs=[0.0007682083478971494, 0.0007682083478971494, 0])
args = dict(basis=basis, x_TF=4*u.micron, gs=[1.0, 1.0, 0.0], ms=[1.0, 1.0], Omega=0, delta=0)
s_ = State_(**args)
s = State(**args)
x, r = s.xyz
s[...] = 5/(l_unit)**1.5*np.exp(-(0*x**2+r**2)/2/l_unit**2)
s_[...] = s[...]
s.cooling_phase = s_.cooling_phase = 1.0
s.plot()
```

```{code-cell} ipython3
from gpe.plot_utils import MPLGrid

s0 = State0()

dt = 0.02*s.t_scale
steps = 100
e0 = EvolverABM(s0, dt=dt/t_unit)
e = EvolverABM(s, dt=dt)
e_ = EvolverABM(s_, dt=dt)

plt.figure(figsize=(15,5))
grid = MPLGrid(direction='right', space=0.2)
e.y.plot(grid=grid)
plt.plot(e0.y.R, 2*e0.y.n)
e_.y.plot(grid=grid)        
plt.plot(e0.y.R, 2*e0.y.n)

with NoInterrupt() as interrupted:
    while e.y.t < 0.3*t_unit and not interrupted:
        e0.evolve(steps)
        e.evolve(steps)
        e_.evolve(steps)
        plt.figure(figsize=(15,5))
        grid = MPLGrid(direction='right', space=0.2)
        e.y.plot(grid=grid)
        plt.plot(e0.y.R, 2*e0.y.n)
        e_.y.plot(grid=grid)        
        plt.plot(e0.y.R, 2*e0.y.n)
        display(plt.gcf())
        plt.close('all')
        clear_output(wait=True)
```

Here is a more realistic example of an expanding spherical cloud.

```{code-cell} ipython3
from gpe.imports import *
from gpe import bec2, axial;reload(bec2); reload(axial)
from gpe.axial import u
from gpe.minimize import MinimizeState

basis = axial.CylindricalBasis(Nxr=(2*64, 2*32), Lxr=(20.0*u.micron, 10*u.micron))

class State(axial.State2Axial):
    t_expand = .02*u.ms
    _ws = np.array([20*126.0, 20*126.0])*u.Hz

    def get_ws(self, t=None):
        if t is None:
            t = self.t
        if t <= self.t_expand:
            return self._ws
        else:
            return 0*self._ws
            return np.concatenate([self._ws[:1], 0*self._ws[1:]])
        
    def get_ws_perp(self, t=None):
        return self.get_ws(t=t)[1:]

    @property
    def ws(self):
        return self.get_ws()
        
    @ws.setter
    def ws(self, ws):
        pass

    def _get_Vext_(self):
        """Return (V_a, V_b, V_ab), the external potentials."""
        V_HO = (0.5*self.ms[self.bcast]
                * sum((_w*_x)**2 for _w, _x in zip(self.ws, self.xyz)))
        zero = np.zeros_like(V_HO[0])
        V_ab = self.Omega*np.exp(2j*self.k_r*self.xyz[0])
        mu_a = mu_b = 0.0
        if hasattr(self, 'mus') and self.mus:
            mu_a, mu_b = self.mus

        return (V_HO[0] - mu_a - self.delta/2.0,
                V_HO[1] - mu_b + self.delta/2.0,
                V_ab + zero)

class State_(State):
    def get_ws_perp(self, t=None):
        """Do not change these - no expansion in basis"""
        return self.get_ws(t=0)[1:]

args = dict(basis=basis, x_TF=4*u.micron)
s_ = State_(**args)
s = State(**args)
s.plot()
m = MinimizeState(s, fix_N=True)
assert m.check()
def callback(s, n=[0]):
    if n[0] % 200 == 0:
        plt.clf()
        s.plot()
        display(plt.gcf())
        plt.close('all')
        clear_output(wait=True)
    n[0] += 1
s = m.minimize(callback=callback, use_scipy=True)
s_[...] = s[...]
s.cooling_phase = s_.cooling_phase = 1.0
xi = np.sqrt(s.hbar**2/2/s.ms[0]/(s.gs[0]*s.get_density().max()))
```

```{code-cell} ipython3
from gpe.plot_utils import MPLGrid
dt = 0.1*s.t_scale
steps = 10
e = EvolverABM(s, dt=dt)
e_ = EvolverABM(s_, dt=dt)
with NoInterrupt() as interrupted:
    while e.y.t < 0.5*u.ms and not interrupted:
        e.evolve(steps)
        e_.evolve(steps)
        plt.figure(figsize=(15,5))
        grid = MPLGrid(direction='right', space=0.2)
        e.y.plot(grid=grid)
        e_.y.plot(grid=grid)        
        display(plt.gcf())
        plt.close('all')
        clear_output(wait=True)
```

## Expansion

+++

If we consider pure expansion from a static trap with $\omega_j(t>0) = 0$ and equal radial frequencies $\omega_x = \omega_y = \omega_\perp$, then we have the following scaling parameters in the two cases:

\begin{gather*}
  λ_\perp(t) = \sqrt{1+\tau^2}, \qquad
  \tau = \omega_\perp(0)t.
\end{gather*}

```{code-cell} ipython3
from scipy.integrate import odeint

from gpe import bec, minimize


class State(bec.State):
    """Here we simulate the state Phi(X, y, z) discussed above with
    explicit scaling dependence added (in y and z).
    
    Note: self.xyz == (X, y, z)
    """
    def init(self):
        lams = np.ones(self.dim)
        dlams = np.zeros(self.dim)
        self._lambda_cache = (0, np.ravel([lams, dlams]))
        super().init()        
        
    def _get_Vext_(self):
        return sum(self.m*(_w*_x)**2/2.0 
                   for _w, _x in zip(self.ws, self.xyz))
    
    def get_ws(self, t=None):
        if t is None:
            t = self.t
        if t <= 0:
            return self.ws
        else:
            return np.zeros(self.dim)
            
    def _rhs(self, q, t):
        """RHS for lambda(t) ODE."""
        lams, dlams =np.reshape(q, (2, self.dim))
        w0s = self.get_ws(t=0)
        ws = self.get_ws(t=t)
        ddlams = -lams * ws**2 + w0s**2/np.prod(lams)/lams
        return np.ravel([dlams, ddlams])
    
    def get_lambdas(self):
        if _t != self._lambda_cache[0]:
            t0, q0 = self._lambda_cache
            q1 = odeint(self._rhs, q0, [t0, self.t])[-1]
            self._lambda_cache = (t1, q1)
        return self._lambda_cache[1][:self.dim]

s = State()
s0 = minimize.MinimizeState(s).minimize()
s0.plot()
```

```{code-cell} ipython3

```

## Summary

+++

### Time-Dependent Harmonic Trap

```{code-cell} ipython3
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.var("V_{ext}", cls=sympy.Function)
ls = [_l(t) for _l in sympy.var('lambda_x, lambda_y, lambda_z', cls=sympy.Function)]
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.var("Phi", cls=sympy.Function)(*(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())
```

## Demonstration

```{code-cell} ipython3
from gpe import bec
from gpe.utils import GPUHelper
from gpe.contexts import FPS

reload(bec)
u = bec.u

class StateMixin(GPUHelper):
    ws = 2 * np.pi * np.array([np.sqrt(8) * 126.0, 126.0, 126.0]) * u.Hz
    def get_ws(self, t=None):
        if t is None:
            t = self.t
        if t <= 0:
            return self.ws
        else:
            return 1.5 * self.ws

    def get_Vext_GPU(self):
        ws = self.get_ws()
        xyz = self.get_xyz_GPU()
        V_trap = 0.5 * self.m * sum((_w * _x) ** 2 for _w, _x in zip(ws, xyz))

        # Arbitrary time-dependence to show that this still works.
        V_trap += 0.5 * np.sin(self.t) * np.exp(-xyz[0] ** 2 / 2)
        # V_trap += 0.5*np.exp(-self.xyz[0]**2/2)
        return V_trap


class State(StateMixin, bec.State):
    pass


class StateScale(StateMixin, bec.StateScaleHO):
    pass


args = dict(Nxyz=(32, 32, 32))
args = dict(Nxyz=(128,), x_TF=10.0)
s0 = State(**args)
s1 = StateScale(**args)


def plot(s):
    n = s.get_density()
    plt.subplot(221)
    x, y, z = list(map(np.ravel, s.xyz))
    imcontourf(y, z, n.max(axis=0), aspect=1)
    plt.subplot(222)
    imcontourf(x, z, n.max(axis=1), aspect=1)
    plt.subplot(223)
    imcontourf(x, y, n.max(axis=2), aspect=1)


dT = 0.1
dt = 0.1 * s0.t_scale
steps = int(np.ceil(dT/dt))
dt = dT/steps
e0 = EvolverABM(MinimizeState(s0).minimize(), dt=dt)
e1 = EvolverABM(MinimizeState(s1).minimize(), dt=dt)
# e0 = EvolverABM(s0, dt=0.4*s0.t_scale)
# e1 = EvolverABM(s1, dt=0.4*s1.t_scale)
clear_output()
e0.evolve(900)
e1.evolve(900)
e0.y.get_energy(), e1.y.get_energy()
```

```{code-cell} ipython3
T = 800
Nt = int(T/dT)
fig, ax = plt.subplots()
for frame in FPS(range(Nt), fig=fig, display=True):
    e0.evolve(steps)
    e1.evolve(steps)
    x0 = e0.y.xyz[0]
    x1 = e1.y.xyz[0]
    ax.cla()
    axs = e0.y.plot(axs=[ax])
    e1.y.plot(axs=axs)
    plt.suptitle(f"{e0.y.get_energy(), e1.y.get_energy()}")
```

### Harmonic Traps

+++

\begin{gather*}
  -\I\hbar \nabla^X_i\Psi = (\hbar \nabla^X_i\theta)\Psi
  + \frac{e^{\I\theta}}{s}\lambda_i(-\I\hbar \nabla^x_i)\Phi(\vect{x}, t)
  = \frac{e^{\I\theta}}{s}\left(
  mX_i \frac{\dot{\lambda}_i}{\lambda_i}
  + \lambda_i(-\I\hbar \nabla^x_i)\right)\Phi(\vect{x}, t),
\end{gather*}

+++

\begin{gather*}
  \int \abs{\op{P}\Psi}^2 \d{X} = \frac{1}{s^2}
  \int \d{X}
  \left(
    \left(mX\frac{\dot{\lambda}}{\lambda}\Phi\right)^2
    + mX\frac{\dot{\lambda}}{\lambda}(2\Phi^\dagger
    \op{p}\Phi)
    + \abs{\op{p}\Phi}^2
  \right)
\end{gather*}

Integrating the middle term by parts we get:

\begin{gather*}
 -\I\hbar\frac{\dot{\lambda}}{\lambda}m  X\Phi^\dagger \Phi_X
 -\I\hbar\frac{\dot{\lambda}}{\lambda}m  \Phi^\dagger (\Phi +X\Phi_X)
\end{gather*}

+++

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*}


## Radial Expansion

+++

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, allowing for modified dispersion (SOC) for example.

## Dynamics Traps

+++

As similar expansion can be performed scaling only $X$ and $Y$:

\begin{gather*}
  \I\hbar\dot{\Psi} = \left\{\frac{-\hbar^2(\nabla_X^2+\nabla_Y^2)}{2m} + E(\op{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, \\
  \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{Y}{\lambda_y(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_j^2(t) + \frac{\omega_j^2(0)}{\lambda_j(t)\lambda_x(t)\lambda_y(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.

+++

## Radial Expansion

+++

A specific use-case is that of an axially symmetric trap in which we describe the expansion in the radial direction.  This is implemented in [gpe/axial.py](../gpe/axial.py).  As before, the physical coordinate $R$ is related to the scaled coordinate $r$ as

\begin{gather*}
  x = X, \qquad
  r = \frac{R}{\lambda_\perp(t)} = \sqrt{x^2+y^2}, \\
  \Psi(X, R, t) = e^{\I\theta}
  \frac{\Phi(x, r, t)}{\lambda_\perp}, \qquad
  \theta = \frac{mR^2}{2\hbar}\diff{}{t}\ln\lambda_\perp(t).
\end{gather*}

We discretize the equations for the scaled function $\Phi(x, r, t)$ on the scaled coordinates $x$ and $r$:

\begin{gather*}
  \I\hbar\dot{\Phi}(x, r ,t) = \left\{
    \frac{-\hbar^2}{2m\lambda_\perp^2}\left(
    \frac{1}{r}\pdiff{}{r}r\pdiff{}{r} + \frac{1}{r^2}\pdiff[2]{}{\phi}
  \right) + E(\op{p}_x)
    + V(X, R, t)
    -\frac{m\omega_\perp^2(t)}{2} R^2
    +\epsilon(\abs\Psi^2)
    +\frac{1}{\lambda_\perp^2(t)}\frac{m\omega_\perp^2(0)}{2}r^2
  \right\}\Phi(x, r ,t),\\
  \ddot{\lambda}_\perp(t) = - \lambda_\perp(t) \omega_\perp^2(t) + \frac{\omega_\perp^2(0)}{\lambda_\perp^3(t)}, \qquad
  \lambda_\perp(0) = 1, \qquad
  \dot{\lambda}_\perp(0) = 0.
\end{gather*}

Note that the scaling $\lambda_\perp(t)$ depends only on the trapping frequency, so can be applied to multiple components as long as the external trapping potential behaves the same way for both.

+++

### Test Example

+++

In the case of expansion $\omega_\perp(t>0) = 0$ we have:

\begin{gather*}
  \ddot{\lambda}_\perp(t) = \frac{\omega_\perp^2(0)}{\lambda_\perp^3(t)}, \qquad
  \lambda_\perp(0) = 1, \qquad
  \dot{\lambda}_\perp(0) = 0,\\
  \lambda_\perp = \sqrt{1+\tau^2}, \qquad \omega_\perp(0)t.
\end{gather*}

Here we explicitly code this example so we have something to work with later.

<!-- \begin{gather*}
  \ddot{\lambda}_\perp(t) = \frac{\omega_\perp^2(0)}{\lambda_\perp^3(t)}, \qquad
  \lambda_\perp(0) = 1, \qquad
  \dot{\lambda}_\perp(0) = 0.\\
  \int_0^{\dot{\lambda_\perp}}\dot{\lambda}_\perp\d{\dot{\lambda}_\perp} = \int_1^{\lambda_\perp}\frac{\omega_\perp^2(0)}{\lambda_\perp^3}\d{\lambda_\perp},\\
  \dot{\lambda}_\perp = \omega_\perp(0)\sqrt{1-\lambda_\perp^{-2}},\\
  \int_1^{\lambda_\perp}\frac{1}{\sqrt{1-\lambda_\perp^{-2}}}\d{\lambda_\perp}
  = \sqrt{\lambda_\perp^2-1} = \tau = \omega_\perp(0)t, \\
  \lambda_\perp = \sqrt{1+\tau^2}, \qquad \omega_\perp(0)t.
\end{gather*} -->

```{code-cell} ipython3
%pylab inline --no-import-all

from gpe.imports import *
from pytimeode import mixins, interfaces
from pytimeode.evolvers import EvolverABM
from mmfutils.math.bases import CylindricalBasis
from zope.interface import implementer


@implementer(interfaces.IStateForABMEvolvers)
class State0(mixins.ArrayStateMixin):
    def __init__(self, N=2**5, R=6):
        self.basis = CylindricalBasis(Nxr=(1, N), Lxr=(1, R))
        R = self.R
        self.data = 5 * np.exp(-R[None, :] ** 2 / 2) + 0j

    @property
    def R(self):
        return self.basis.xyz[1].ravel()

    @property
    def n(self):
        return abs(self[...].ravel()) ** 2

    def compute_dy_dt(self, dy):
        y = self[...]
        dy[...] = self.basis.laplacian(y, factor=-1.0 / 2)
        dy[...] += abs(y) ** 2 * y
        dy /= 1j
        return dy

    def plot(self):
        plt.plot(self.R, self.n)
        plt.title("t={:.4g}".format(self.t))


@implementer(interfaces.IStateForABMEvolvers)
class State1(State0):
    @property
    def lam(self):
        return np.sqrt(1 + self.t**2)

    @property
    def R(self):
        return self.basis.xyz[1].ravel() * self.lam

    @property
    def n(self):
        return abs(self[...].ravel()) ** 2 / self.lam**2

    def compute_dy_dt(self, dy):
        x, r = self.basis.xyz
        y = self[...]
        dy[...] = self.basis.laplacian(y, factor=-1.0 / 2)
        dy[...] += (r**2 / 2.0 + abs(y) ** 2) * y
        dy /= self.lam**2
        dy /= 1j
        return dy


s0 = State0()
s1 = State1()

dt = 0.0001
steps = 200
e0 = EvolverABM(s0, dt=dt)
e1 = EvolverABM(s1, dt=dt)
# fig = plt.figure(figsize=(10,5))
with NoInterrupt() as interrupted:
    while not interrupted and e0.t < 1:
        e0.evolve(steps)
        e1.evolve(steps)

        plt.clf()
        e0.y.plot()
        e1.y.plot()
        display(plt.gcf())
        clear_output(wait=True)
```

Here we do the same thing with the `axial.State2Axial` class:

```{code-cell} ipython3
%pylab inline --no-import-all

from gpe.imports import *
from gpe import bec2, axial

reload(bec2)
reload(axial)
from gpe.axial import u
from gpe.minimize import MinimizeState

l_unit = 1.0
t_unit = 1.0
basis = axial.CylindricalBasis(Nxr=(64, 32), Lxr=(2 * 6 * l_unit, 6 * l_unit))


class State(axial.State2Axial):
    t_expand = 0.0 * t_unit
    _ws = np.array([48.78 * 0, 48.78]) * u.Hz
    _ws = np.array([0.0 / u.Hz, 1.0 / u.Hz]) * u.Hz

    def get_ws(self, t=None):
        if t is None:
            t = self.t
        if t < self.t_expand:
            return self._ws
        else:
            # return 0*self._ws
            return np.concatenate([self._ws[:1], 0 * self._ws[1:]])

    def get_ws_perp(self, t=None):
        return self.get_ws(t=t)[1:]

    @property
    def ws(self):
        return self.get_ws()

    @ws.setter
    def ws(self, ws):
        pass

    def get_Vext(self):
        """Return (V_a, V_b, V_ab), the external potentials."""
        V_HO = (
            0.5
            * self.ms[self.bcast]
            * sum((_w * _x) ** 2 for _w, _x in zip(self.ws, self.xyz))
        )
        V_ab = self.Omega * np.exp(2j * self.k_r * self.xyz[0])
        mu_a = mu_b = 0.0
        if hasattr(self, "mus") and self.mus:
            mu_a, mu_b = self.mus

        return (
            V_HO[0] - mu_a - self.delta / 2.0,
            V_HO[1] - mu_b + self.delta / 2.0,
            V_ab,
        )


class State_(State):
    def get_ws_perp(self, t=None):
        """Do not change these - no expansion in basis"""
        return self.get_ws(t=0)[1:]


args = dict(
    basis=basis, x_TF=4 * u.micron, gs=[0.0007682083478971494, 0.0007682083478971494, 0]
)
args = dict(
    basis=basis, x_TF=4 * u.micron, gs=[1.0, 1.0, 0.0], ms=[1.0, 1.0], Omega=0, delta=0
)
s_ = State_(**args)
s = State(**args)
x, r = s.xyz
s[...] = 5 / (l_unit) ** 1.5 * np.exp(-(0 * x**2 + r**2) / 2 / l_unit**2)
s_[...] = s[...]
s.cooling_phase = s_.cooling_phase = 1.0
s.plot()
```

```{code-cell} ipython3
from gpe.plot_utils import MPLGrid

s0 = State0()

dt = 0.02 * s.t_scale
steps = 100
e0 = EvolverABM(s0, dt=dt / t_unit)
e = EvolverABM(s, dt=dt)
e_ = EvolverABM(s_, dt=dt)

plt.figure(figsize=(15, 5))
grid = MPLGrid(direction="right", space=0.2)
e.y.plot(grid=grid)
plt.plot(e0.y.R, 2 * e0.y.n)
e_.y.plot(grid=grid)
plt.plot(e0.y.R, 2 * e0.y.n)

with NoInterrupt() as interrupted:
    while e.y.t < 0.3 * t_unit and not interrupted:
        e0.evolve(steps)
        e.evolve(steps)
        e_.evolve(steps)
        plt.figure(figsize=(15, 5))
        grid = MPLGrid(direction="right", space=0.2)
        e.y.plot(grid=grid)
        plt.plot(e0.y.R, 2 * e0.y.n)
        e_.y.plot(grid=grid)
        plt.plot(e0.y.R, 2 * e0.y.n)
        display(plt.gcf())
        plt.close("all")
        clear_output(wait=True)
```

Here is a more realistic example of an expanding spherical cloud.

```{code-cell} ipython3
%pylab inline --no-import-all

from gpe.imports import *
from gpe import bec2, axial

reload(bec2)
reload(axial)
from gpe.axial import u
from gpe.minimize import MinimizeState

basis = axial.CylindricalBasis(Nxr=(2 * 64, 2 * 32), Lxr=(20.0 * u.micron, 10 * u.micron))


class State(axial.State2Axial):
    t_expand = 0.02 * u.ms
    _ws = np.array([20 * 126.0, 20 * 126.0]) * u.Hz

    def get_ws(self, t=None):
        if t is None:
            t = self.t
        if t <= self.t_expand:
            return self._ws
        else:
            return 0 * self._ws
            return np.concatenate([self._ws[:1], 0 * self._ws[1:]])

    def get_ws_perp(self, t=None):
        return self.get_ws(t=t)[1:]

    @property
    def ws(self):
        return self.get_ws()

    @ws.setter
    def ws(self, ws):
        pass

    def get_Vext(self):
        """Return (V_a, V_b, V_ab), the external potentials."""
        V_HO = (
            0.5
            * self.ms[self.bcast]
            * sum((_w * _x) ** 2 for _w, _x in zip(self.ws, self.xyz))
        )
        V_ab = self.Omega * np.exp(2j * self.k_r * self.xyz[0])
        mu_a = mu_b = 0.0
        if hasattr(self, "mus") and self.mus:
            mu_a, mu_b = self.mus

        return (
            V_HO[0] - mu_a - self.delta / 2.0,
            V_HO[1] - mu_b + self.delta / 2.0,
            V_ab,
        )


class State_(State):
    def get_ws_perp(self, t=None):
        """Do not change these - no expansion in basis"""
        return self.get_ws(t=0)[1:]


args = dict(basis=basis, x_TF=4 * u.micron)
s_ = State_(**args)
s = State(**args)
s.plot()
m = MinimizeState(s, fix_N=True)
assert m.check()


def callback(s, n=[0]):
    if n[0] % 200 == 0:
        plt.clf()
        s.plot()
        display(plt.gcf())
        plt.close("all")
        clear_output(wait=True)
    n[0] += 1


s = m.minimize(callback=callback, use_scipy=True)
s_[...] = s[...]
s.cooling_phase = s_.cooling_phase = 1.0
xi = np.sqrt(s.hbar**2 / 2 / s.ms[0] / (s.gs[0] * s.get_density().max()))
```

```{code-cell} ipython3
from gpe.plot_utils import MPLGrid
```

```{code-cell} ipython3
dt = 0.1 * s.t_scale
steps = 10
e = EvolverABM(s, dt=dt)
e_ = EvolverABM(s_, dt=dt)
with NoInterrupt() as interrupted:
    while e.y.t < 0.5 * u.ms and not interrupted:
        e.evolve(steps)
        e_.evolve(steps)
        plt.figure(figsize=(15, 5))
        grid = MPLGrid(direction="right", space=0.2)
        e.y.plot(grid=grid)
        e_.y.plot(grid=grid)
        display(plt.gcf())
        plt.close("all")
        clear_output(wait=True)
```

## Expansion

+++

If we consider pure expansion from a static trap with $\omega_j(t>0) = 0$ and equal radial frequencies $\omega_x = \omega_y = \omega_\perp$, then we have the following scaling parameters in the two cases:

\begin{gather*}
  \lambda_\perp(t) = \sqrt{1+\tau^2}, \qquad
  \tau = \omega_\perp(0)t.
\end{gather*}

```{code-cell} ipython3
%pylab inline --no-import-all
from scipy.integrate import odeint

from gpe import bec, minimize


class State(bec.State):
    """Here we simulate the state Phi(X, y, z) discussed above with
    explicit scaling dependence added (in y and z).

    Note: self.xyz == (X, y, z)
    """

    def init(self):
        lams = np.ones(self.dim)
        dlams = np.zeros(self.dim)
        self._lambda_cache = (0, np.ravel([lams, dlams]))
        bec.State.init()

    def get_Vext(self):
        return sum(self.m * (_w * _x) ** 2 / 2.0 for _w, _x in zip(self.ws, self.xyz))

    def get_ws(self, t=None):
        if t is None:
            t = self.t
        if t <= 0:
            return self.ws
        else:
            return np.zeros(self.dim)

    def _rhs(self, q, t):
        """RHS for lambda(t) ODE."""
        lams, dlams = np.reshape(q, (2, self.dim))
        w0s = self.get_ws(t=0)
        ws = self.get_ws(t=t)
        ddlams = -lams * ws**2 + w0s**2 / np.prod(lams) / lams
        return np.ravel([dlams, ddlams])

    def get_lambdas(self):
        if _t != self._lambda_cache[0]:
            t0, q0 = self._lambda_cache
            q1 = odeint(self._rhs, q0, [t0, self.t])[-1]
            self._lambda_cache = (t1, q1)
        return self._lambda_cache[1][: self.dim]


s = State()
s0 = minimize.MinimizeState(s).minimize()
s0.plot()
```