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

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

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

(sec:Viscosity2)=
# Viscosity Notes 2

These are some supporting notes for the document {ref}`sec:EffectiveViscosity`.

## Time Evolution

Here we consider the time-evolution of the system in terms of the following PDEs:
\begin{gather*}
  \dot{n} + (nu)_{,x} = 0\\
  D_tu = \dot{u} + uu_{,x} = -\frac{1}{m}\Bigl(
    V(x) + \overbrace{\mathcal{E}'(n)}^{\mu}
  \Bigr)_{,x} + \frac{\bigl(ν(n) u_{,x}\bigr)_{,x}}{n}.
\end{gather*}
We start by implementing them as follows, in terms of the density and velocity $y = (n,
u)$:
\begin{gather*}
  \diff{}{t}
  \begin{pmatrix}
    n\\
    u
  \end{pmatrix}
  =
  \begin{pmatrix}
    -(nu)_{,x},\\
    - uu_{,x} - \frac{V'(x) + \mathcal{E}''(n)n_{,x}}{m} 
    + \frac{\nu'(n)n_{,x} u_{,x} + \nu(n)u_{,xx}}{n}
  \end{pmatrix}
\end{gather*}
We start with this simple formulation, and use a finite difference scheme for computing
the spatial derivatives with boundary conditions imposed by adding ghost points,
extending the arrays so we can use simple finite-difference operators.

:::{admonition} How to compute derivatives.
There are ambiguities about how to compute the derivatives when implementing these
equations owing to the fact that the Leibniz identity (product rule) is not satisfied by
numerical derivative techniques:
\begin{gather*}
  (nu)_{,x} = n_{,x}u + nu_{,x}, \qquad
  \op{D}(nu) \neq \op{D}(n)u + n\op{D}(u).
\end{gather*}
Errors/differences introduced this way should vanish in the limit of UV convergence as
we take the lattice spacing to zero, but is there a preferred approach?  There may be
other factors to guide the choice.

1. **Stability:** As we shall see, probably the most important criteria is that the
   evolution be stable in the sense that we do not want any growth of unphysical modes
   (typically high-frequency modes).  This will be a dominant criterion.

2. **Consistency:** If the equations follow from a minimization principle or stationary
   action, then it can be worthwhile ensuring that the equations of motion are a
   numerically exact implementation of the variation.  For example, if one has an action
   with a spatially dependent effective mass in the Schrödinger equation:
   \begin{gather*}
     \mathcal{L} = \frac{\hbar^2}{2m(x)} \abs{\op{D}\psi(x)}^2,
   \end{gather*}
   then one should implement the Hamiltonian as
   \begin{gather*}
     \op{H}\psi = +\op{D}^T\left(\frac{\hbar^2 \op{D}\psi}{2m(x)}\right).
   \end{gather*}
   Depending on how the numerical derivative $\mat{D}$ is implemented, one might have
   $\op{D}^T = - \op{D}$, yielding the more familiar form from integrating by parts, but
   this should be explicitly verified.

3. **Conservation:** Adherence to the form of conservation laws might suggest computing
   $j_{,x}$ above rather than $n_{,x}u + nu_{,x}$.  Good numerical approaches will often
   ensure that conserved quantities are conserved to numerical accuracy.  Note that this
   can be misleading: solutions may now look nice (as opposed to blowing up) but may be
   quantitatively inaccurate. It can be valuable to have a technique that **does not**
   manifestly conserve particle number, energy, etc. and then to **use these** to test
   check the accuracy of the evolution.
   
4. Alternatively, if computing the derivative is costly, performance considerations
   might favour the alternative approach if $n_{,x}$ and $u_{,x}$ are needed elsewhere.
   *Remember: don't make such a decision without actually profiling the performance.*

5. Boundary conditions might also play a role.  In our case, we want to fix the current
   $j = nu$ flowing into the system, allowing the state to relax (due to the viscosity)
   to a state with fixed $j$.  This would be quite difficult to implement using ghost
   points if we were strict about using only $n$ and $u$.

We will explore some of these effects below.
:::

## Burgers' Equation

We start with a simplified version of these equations that admits an analytic solution
for testing: [Burgers' equation][]:
\begin{gather*}
  D_t u = \dot{u} + uu_{,x} = \nu u_{,xx}.
\end{gather*}
:::{margin}
The approximation we make is to drop the $\nu n_{,x}u_{,x}/n$ piece, which can be
ensured in our code by setting $n(x)=0$.
:::
Our system reduces to this in the limit $\mathcal{P}(x) = 0$ (i.e. $V(x) =
\mathcal{E}(n) = 0$) with Lamé viscosity $\nu(n) = \nu n$ with $[\nu] = D^2/T$ in the
low density limit $n(x) \rightarrow 0$.  In this case, we have a single equation for the
velocity $u(x)$.

This equation is solvable (see below).  Some nice solutions for testing are:
\begin{gather*}
  u(x, t) = c - \delta \tanh\frac{\delta(x-ct)}{2\nu}.
\end{gather*}
This is a traveling wave with constant shape, but requires Dirichlet boundary conditions
far from the transition, or open boundary conditions.

In ["Step 4: Burgers' Equation" of Lorena Barba's "12 steps to Navier-Stokes"][step 4],
they present the following approximate, but highly accurate (for small $\nu$) periodic solution:
:::{margin}
Here we choose units so that distances are measured in terms of the wavelength (box length) $L =
2\pi$ and period $T = 2\pi$, defining a wavevector $k=2\pi/L$ and frequency $\omega =
2\pi/T$.  The dimensionful solution would be
\begin{gather*}
  \phi(x, t) = \exp\left(
    \frac{-(kx - \omega t)^2}{\nu (4 + \omega t)} 
  \right) \\
  + \exp\left(
    \frac{-(kx - \omega t - 2\pi)^2}{\nu (4 + \omega t)} 
  \right).
\end{gather*}
:::
\begin{gather*}
  u(x, t) = 4 - 2\nu (\ln \varphi)_{,x}, \\
  \varphi(x, t) 
    = \exp\left(\frac{-(x-4t)^2}{4\nu(1+t)}\right) 
    + \exp\left(\frac{-(x-4t - 2\pi)^2}{4\nu(1+t)}\right)
\end{gather*}
which gives an (almost) periodic velocity profile.

```{code-cell}
:tags: [margin, hide-input]
import sympy
sympy.var('x, t, nu, c, sigma')
def phi(x):
    return sympy.exp(-x**2/(4*nu*t+sigma**2))
phi_c = phi(x-c*t) + phi(x-2*sympy.pi-c*t)
u = (c - 2*nu*sympy.log(phi_c).diff(x))
#assert (u.diff(t) + u*u.diff(x) - nu*u.diff(x,x)).simplify() == 0
u = u.subs([(c, 4), (sigma, sympy.sqrt(4*nu))])
y = u.subs([(t, 0), (nu, 1)])
sympy.plot(y);
```
:::{admonition} On the Periodicity.
:class: dropdown

This solution is not periodic (see the margin), but, for small $\nu$, it is very close.
The discontinuity at the boundary,
\begin{gather*}
  u(2\pi, 0) - u(0, 0) = \frac{4\pi}{e^{\frac{\pi^{2}}{\nu}} + 1},
\end{gather*}
becomes exponentially small as $\nu$ gets smaller.  I.e., with $\nu = 0.07$ as used in
[step 4][], the discontinuity is $\delta u \sim 10^{-61}$.
:::

:::::{admonition} Solve with the [Cole-Hopf Transformation][]: $u=-2\nu (\ln \phi)_{,x}$.
:class: dropdown

[Burgers' equation][] can be linearized by the [Cole-Hopf Transformation][]:
\begin{gather*}
  2\nu \left(\frac{\dot{\phi} - \nu \phi_{,xx}}{\phi}\right)_{,x} = 0,\\
  \dot{\phi} - \nu \phi_{,xx} = \dot{f}(t)\phi.
\end{gather*}
This is solved by transforming $\phi \rightarrow \phi e^{f(t)}$, which does not affect
$u(x, t)$ *(Why?)*, to find the heat equation
\begin{gather*}
  \dot{\phi} = \nu \phi_{,xx}.
\end{gather*}
We now have a linear problem that can be completely solved (do it!):
\begin{gather*}
  \phi(x, t) = \frac{1}{\sqrt{4\pi \nu t}}\int_{-\infty}^{\infty}
  \phi_0(\tilde{x})\exp\left[-\frac{(x-\tilde{x})^2}{4\nu t}\right]\d{\tilde{x}},\\
  \ln \phi_0(x) = -\frac{1}{2\nu}\int_0^{x} u_0(\tilde{x})\d{\tilde{x}}.
\end{gather*}

To get the stated solution, note that, if we start with a Gaussian, it remains a Gaussian:
\begin{gather*}
  \phi_0(x) = \frac{e^{- x^2 / \sigma^2}}{\sigma}, \\
  \phi_{\sigma}(x, t) = \frac{1}{\sigma\sqrt{4\pi \nu t}}\int_{-\infty}^{\infty}
  \exp\left[
    -\frac{\tilde{x}^2}{\sigma^2} - \frac{(x-\tilde{x})^2}{4\nu t}
  \right]\d{\tilde{x}}
  =
  \frac{e^{- x^{2}/ (4 \nu t + \sigma^2)}}
       {\sqrt{4 \nu t + \sigma^2}}.
\end{gather*}
Recalling that multiplying $\phi(x, t)$ by a constant function of time does not affect
the solution, we can simplify this by dropping the overall factor:
\begin{gather*}
  \phi_{\sigma}(x, t) \equiv \exp\left(\frac{-x^2}{4\nu t + \sigma^2}\right).
\end{gather*}
This is a stationary gaussian that expands with time (as per the diffusion equation).
To obtain moving solutions, we can simply boost to a moving frame $u(x, t) \rightarrow
u(x-ct, t) + c$.  *(Remember, $u$ is a velocity, so we need to add $c$.)*  Formally, 
if $u(x,t)$ solves Burgers' equation, then so does
\begin{gather*}
  \tilde{u}(x, t) = u(x - x_0 - ct, t) + c\\
  \dot{\tilde{u}} +\tilde{u}\tilde{u}_{,x} = 
  \dot{u} - cu_{,x} + (u+c)u_{,x} =
  \dot{u} + uu_{,x} = \nu u_{,xx} = \nu \tilde{u}_{,xx}.
\end{gather*}
This allows us to express moving and translated solutions in terms of stationary
solutions.
\begin{gather*}
  \phi(x, t) = \phi_{\sigma}(x - x_0 - ct, t), \qquad
  u(x, t) = c - 2\nu (\ln\phi_{\sigma})_{,x}.
\end{gather*}
Finally, since the equation for $\phi$ is linear, we can simply add solutions.  The
stated solution is such a sum with $c=4$, $\sigma^2 = 4\nu$, and $x_0 \in \{0, 2\pi\}$
(in appropriate units):
\begin{gather*}
  u(x, t) = c - 2 \nu (\ln \phi_c)_{,x}, \qquad
  \phi_{c}(x, t) = \phi_{\sigma}(x - ct, t) + \phi_{\sigma}(x - x_0 - ct, t).
\end{gather*}

:::{admonition} Solution
:class: dropdown

The solution is simple in Fourier space.  Note that this has exactly the same structure
as the Schödinger equation, up to factors of $\I$ and $\hbar$:
\begin{gather*}
  \ket{\dot{\phi}} = -\nu \op{k}^2\ket{\phi}, \qquad
  \braket{x|\op{k}|\phi} = -\I \partial_{x}\braket{x|\phi} 
  = -\I \phi'(x),
\end{gather*}
so we can express our answer in terms of plane-waves (momentum eigenstates)
\begin{gather*}
  \partial_{x}\ket{k} = - \I k \ket{k}, \qquad
  \braket{x|k} = e^{\I k x}, \qquad
  \int_{-\infty}^{\infty}\frac{\d{k}}{2\pi} \ket{k}\!\bra{k} 
  =
  \int_{-\infty}^{\infty} \d{x}\ket{x}\!\bra{x} = \op{1}.
\end{gather*}
The state at time $t$ can simply be expressed in terms of an initial state, and then
formulae obtained by inserting the identity:
\begin{gather*}
    \ket{\phi(t)} = e^{\nu \op{k}^2 t}\ket{\phi(0)}
                  = e^{\nu \op{k}^2 t}\op{1}\ket{\phi(0)}
                  = \int_{-\infty}^{\infty}\frac{\d{k}}{2\pi}
                  e^{\nu \op{k}^2 t} \ket{k} \braket{k|\phi(0)}\\
                  = \int_{-\infty}^{\infty}\frac{\d{k}}{2\pi}
                  \ket{k} e^{\nu k^2 t} \braket{k|\phi(0)}.
\end{gather*}
Converting to position space, we have
\begin{gather*}
  \phi(x, t) = \braket{x|\phi} 
             = \int_{-\infty}^{\infty}\d{\tilde{x}}
               \int_{-\infty}^{\infty}\frac{\d{k}}{2\pi}
               \overbrace{\braket{x|k}}^{e^{\I k x}}
               e^{-\nu k^2 t} 
               \overbrace{\braket{k|\tilde{x}}}^{e^{-\I k \tilde{x}}}
               \overbrace{\braket{\tilde{x}|\phi(0)}}^{\phi_0(x)},\\
  = \int_{-\infty}^{\infty}\d{\tilde{x}}
    \underbrace{\int_{-\infty}^{\infty}\frac{\d{k}}{2\pi}
                 e^{-\nu k^2 t + \I k (x - \tilde{x})}}_{G_t(x-\tilde{x})}
                 \phi_0(\tilde{x}).
\end{gather*}
The propagator (Green's function) can thus be computed by completing the square of the
Gaussian integral with $2\nu k_0 t = \I x$, or $k_0^2 = -x^2/4\nu^2$ with the
substitution $\tilde{k}^2 =\nu(k-k_0)^2 t$ so that $\d{\tilde{k}} =  \sqrt{\nu t} \d{k}$:
\begin{gather*}
  G_t(x) = \int_{-\infty}^{\infty}\frac{\d{k}}{2\pi} e^{-\nu k^2 t + \I k x}
         = \int_{-\infty}^{\infty}\frac{\d{k}}{2\pi} e^{-\nu (k-k_0)^2 t + \nu k_0^2 t}\\
  = \frac{e^{-x^2 t/4\nu}}{\sqrt{4\pi^2 \nu t}}
    \underbrace{\int_{-\infty}^{\infty}\d{\tilde{k}}\; e^{-\tilde{k}^2}}_{\sqrt{\pi}}
  = \frac{e^{-x^2 t/4\nu}}{\sqrt{4\pi \nu t}}.
\end{gather*}
:::
:::::

```{code-cell}
:tags: [margin, hide-input]
import sympy
sympy.init_printing(use_latex=True)
x, t, nu, c = sympy.var('x, t, nu, c')
phi = sympy.Function('phi')(x, t)
u = c - 2*nu*(sympy.log(phi)).diff(x)
eq = -(u.diff(t) + u*u.diff(x) - nu*u.diff(x, x)).simplify()
eq2 = (2*nu*((phi.diff(t) + c*phi.diff(x) - nu*phi.diff(x,x))/phi).diff(x)).simplify()
assert (eq - eq2).simplify() == 0
```

[Burgers' equation]: <https://en.wikipedia.org/wiki/Burgers'_equation>
[Cole-Hopf Transformation]: <https://en.wikipedia.org/wiki/Cole%E2%80%93Hopf_transformation>

### Exercise #1: Numerically Solve Burger's Equation

```{code-cell}
Nx = 256
Lx = 2*np.pi
dx = Lx / Nx
x = np.arange(Nx) * dx

# Forward differencing 
Df = (np.diag(np.ones(Nx-1), 1) - np.diag(np.ones(Nx), 0))/dx
Df[-1, 0] = 1/dx

# Backwards differencing
Db = (np.diag(np.ones(Nx), 0) - np.diag(np.ones(Nx-1), -1))/dx
Db[0, -1] = -1/dx

assert np.allclose(Df.T, -Db)

# Center difference
Dc = (Df + Db)/2

# Second derivative
D2 = -Df.T @ Df

# Check:
f = np.sin(x)
df = np.cos(x)
ddf = -np.sin(x)

# What error do we expect?  Df and Db should 

# Tolerance factor for errors.  Since f, df, etc. all have magnitude
# <= 1, we should be fine with a value of 1 here, but in other cases,
# you may need something up to the magnitude of the appropriate derivatives.
tol_factor = 1.1
assert np.allclose(df, Df @ f, atol=tol_factor * (dx/2), rtol=0)
assert np.allclose(df, Db @ f, atol=tol_factor * (dx/2), rtol=0)
assert np.allclose(df, Dc @ f, atol=tol_factor * (dx**2/6), rtol=0)
assert np.allclose(ddf, D2 @ f, atol=tol_factor * (dx**2/12), rtol=0)

# To test the error formula, these should be tight!
tol_factor = 0.9
assert not np.allclose(df, Df @ f, atol=tol_factor * (dx/2), rtol=0)
assert not np.allclose(df, Db @ f, atol=tol_factor * (dx/2), rtol=0)
assert not np.allclose(df, Dc @ f, atol=tol_factor * (dx**2/6), rtol=0)
assert not np.allclose(ddf, D2 @ f, atol=tol_factor * (dx**2/12), rtol=0)
```

:::::{admonition} Notes about the Solution.
First we consider the finite difference errors.  For large step sizes, these derive from
the Taylor expansion (so-called truncation errors, as opposed to round-off errors that
occur at small $h \lesssim \sqrt{\varepsilon}$):
\begin{align*}
  f(x\pm h) = f(x) \pm h f'(x) + &\frac{h^2}{2}f''(x) \pm \frac{h^3}{6} f'''(x) +
  \frac{h^4}{24} f^{(4)}(x) + \cdots\\
  \frac{f(x+h) - f(x)}{h} &= f'(x) + \frac{h}{2}f''(\xi),\\
  \frac{f(x) - f(x-h)}{h} &= f'(x) - \frac{h}{2}f''(\xi),\\
  \frac{f(x+h) - f(x-h)}{2h} &= f'(x) + \frac{h^2}{6}f^{(3)}(\xi),\\
  \frac{f(x+h) + f(x-h) - 2f(x)}{h^2} &= f''(x) + \frac{h^2}{12}f^{(4)}(\xi).
\end{align*}
*(It is not obvious, but one can use the mean-value theorem to show that these
error bounds are exact for some point $\xi \in [x-h, x+h]$. **Need Checking.**)*
:::::


```{code-cell}
nu = 0.1
c = 1.0

def compute_du_dt(u):
    return nu * D2 @ u - u*Df@u

x0 = Lx/4
x1 = 3*Lx/4
c0 = 1
d0 = -1
c1 = 0
d1 = 1
u0 = c0 - d0 * np.tanh(d0*(x-x0)/2/nu)
u1 = c1 + d1 * np.tanh(d1*(x-x1)/2/nu)
u = u0 + u1

from gpe.contexts import FPS

Nt = 100
dt = dx * nu * 0.1
ts = np.arange(Nt) * dt
fig, ax = plt.subplots()

u = u0 + u1
for t in FPS(ts, fig=fig, display=True):
    u += compute_du_dt(u) * dt
    ax.cla()
    ax.plot(x, u)
```

:::{admonition} About the code.

We include here as simple an implementation as possible without any bells and whistles,
but have an accompanying code that will be feature-rich. I do this for several reasons:
1. The documentation stands complete as a reference -- no need to refer to another file.
2. Simple code is easier to understand, and the features can distract from the main
   point.
3. Simple code is easy to modify to explore effects like the choice of discretization
   discussed above: adding flags for all of these features would greatly bloat the
   production code.  Once we learn what is important, we can add only the important
   features.
4. Writing simple code like this helps keep me comfortable with rapid prototyping. I
   find it empowering and reassuring to quickly test things out and often get distracted
   myself when writing feature-rich code.
5. Independently implementing code in different ways provides independent checks,
   helping me find mistakes.
6. Trying write simple, short, code like this often helps me improve my design.

Note: I actually work as follows:

1. Start with a simple version that gets me something, often develop in the markdown
   file, or sometimes switching to a notebook.
2. Once I have something working, I code this in a module that I can import to play
   with.  I generally prefer to code in a proper Python file with an editor that
   supports completion, linting, etc. so I don't make mistakes.  Also, once I have
   importable code, I can easily explore changes without having to copy or re-execute
   cells in a notebook.  Notebooks are good for demonstrating things, but not for
   extended code.
3. I then re-write the simple versions for demonstration purposes here in the markdown
   file, co-developing and improving the pure Python versions as I have ideas.

I highly recommend that you do this too!
:::

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


class Flow:
    """Simple class for solving viscous flow problems without a potential."""

    m = 1.2
    g = 1.3
    nu = 0.1  # Viscosity (Lame with power alpha = 1)

    # Lattice
    x0 = -1.0
    x1 = 1.0
    Nx = 200
    
    # Boundary conditions on the left and right.  None means Neumann.
    n0 = v0 = None  # Left
    n1 = v1 = None  # Right

    def __init__(self, **kw):
        for key in kw:
            if not hasattr(self, key):
                raise ValueError(f"Unknown {key=}")
            setattr(self, key, kw[key])
        self.init()

    def init(self):
        x = np.linspace(self.x0, self.x1, self.Nx + 2)
        self.x = x[1:-1]
        dx = np.diff(x).mean()

        def one(k):
            """Return a diagonal matrix with 1s on the kth diagonal."""
            return np.diag(np.ones(self.Nx + 2 - abs(k)), k=k)

        self._Dx = (one(1) - one(-1)) / 2 / dx
        self._D2x = (one(1) - 2*one(0) + one(-1)) / dx**2
        
    def get_Eint(self, n, d=0):
        """Equation of state (GPE here)."""
        if d == 0:
            return self.g * n**2 / 2
        elif d == 1:
            return self.g * n
        elif d == 2:
            return self.g
        else:
            return 0

    def get_Vext(self, x, t, d=0):
        """Return the external potential."""
        return 0.0

    def get_nu(self, n, d=0):
        if d == 0:
            return self.nu * n
        else:
            return self.nu

    ######################################################################
    # Dynamics solver
    def apply_D(self, f, f0=None, f1=None, d=1):
        """Return the dth derivative of f."""
        # Pad with external points (Neumann boundary condistions)
        f_ = np.concatenate([f[:1], f, f[-1:]])
        
        # If values are provided, use these instead
        if f0 is not None:
            f_[0] = f0
        if f1 is not None:
            f_[-1] = f1
        
        if d == 1:
            D = self._Dx
        elif d == 2:
            D = self._D2x
        return (D @ f_)[1:-1]

    def pack(self, n, u):
        return np.ravel([n, u])

    def unpack(self, y):
        return np.reshape(y, (2, len(y) // 2))

    def compute_dy_dt(self, t, y):
        x = self.x
        n, u = self.unpack(y)
        j = n * u
        
        n0, v0, n1, v1 = self.n0, self.v0, self.n1, self.v1
        j0 = j1 = None
        if n0 is not None and v0 is not None:
            j0 = n0 * v0
        if n1 is not None and v0 is not None:
            j1 = n1 * v1
            
        # Compute derivatives
        j_x = self.apply_D(j, j0, j1)
        u_x = self.apply_D(u, v0, v1)
        u_xx = self.apply_D(u, v0, v1, d=2)
        n_x = self.apply_D(n, n0, n1)
        n_xx = self.apply_D(n, n0, n1, d=2)
        
        dn_dt = -j_x
        du_dt = (
            -u * u_x
            - (self.get_Vext(x, t, d=1) + self.get_Eint(n, d=2) * n_x) / self.m
            # + self.apply_D((self.get_nu(n) * u_x)) / n
            + (self.get_nu(n, d=1) * n_x * u_x + self.get_nu(n) * u_xx) / n
        )
        return self.pack(dn_dt, du_dt)


class FlowStep(Flow):
    """Add a step potential."""
    V_dV = 0.1  # Potential change from right to left.

    def get_Vext(self, x, t, d=0):
        """Return the external potential."""
        if d == 0:
            return np.where(x<0, 0, self.V_dV)
        elif d == 1:
            ind = np.argmin(abs(x-0))
            dx = np.diff(x).mean()
            dV = np.zeros_like(x)
            dV[ind] = self.V_dV / (x[ind+1] - x[ind])
            return dV
```

### Dam Breaking
Let's see how well this works.  We start with a dam-breaking problem.

```{code-cell}
from scipy.integrate import solve_ivp

f = Flow(x0=-1.0, x1=1.0, Nx=200)
x = f.x

n0 = np.where(x < 0, 1.0, 0.5)
v0 = 0*n0
y0 = f.pack(n0, v0)

t_span = (0, 100.0)

res = solve_ivp(f.compute_dy_dt, y0=y0, t_span=t_span, method="BDF")
if not res.success:
    print(res.message)
    
# Unpack solution
x, ts = f.x, res.t
V = f.get_Vext(x, t=0)
ns, us = res.y.reshape((2, len(x), len(ts)))
fig, axs = plt.subplots(2, 1, height_ratios=(3, 1),
                        gridspec_kw=dict(hspace=0))
ax0, axV = axs
ax1 = ax0.twinx()
ax0.plot(x, ns[:, -1])
ax1.plot(x, us[:, -1], 'C1:')
axV.plot(x, V)
```


```{code-cell}
from scipy.integrate import solve_ivp

j0 = 1.0
n0 = 1.0
v0 = j0/n0

f = FlowStep(x0=-1.0, x1=1.0, Nx=200, n0=n0, v0=v0)

m_scale = f.m
n_scale = n0
x_scale = 1/n0
v_scale = abs(v0)
E_scale = f.m * v_scale**2

one = np.ones_like(f.x)
y0 = f.pack(n=n0*one, u=v0*one)

t_span = (0, 100.0)

# The default https://github.com/scipy/scipy/issues/18039
res = solve_ivp(f.compute_dy_dt, y0=y0, t_span=t_span, method="BDF")
if not res.success:
    print(res.message)
    
# Unpack solution
x, ts = f.x, res.t
V = f.get_Vext(x, t=0)
ns, us = res.y.reshape((2, len(x), len(ts)))
fig, axs = plt.subplots(2, 1, height_ratios=(3, 1),
                        gridspec_kw=dict(hspace=0))
ax0, axV = axs
ax1 = ax0.twinx()
ax0.plot(x/x_scale, ns[:, -1]/n_scale)
ax1.plot(x/x_scale, us[:, -1]/v_scale, 'C1:')
axV.plot(x/x_scale, V/E_scale)
```

:::{margin}
Here is some code to make a movie.  To develop this we do the following:

1. Initialize the plot (axes, subplots, etc.).
2. Make a function `draw_frame` to draw the plot elements.
3. Test this and adjust until it looks like we want.

Unlike regular plotting, we must now make arrangements to either clear the axes before
drawing each frame, or to update the lines.  We start with the former.

4. Clear the axes in the function.
5. Create the animation object.
6. Generate and display the output - in this case we choose HTML5.  Note: we modify the
   width a little here.
   
As can be seen, this is quite slow.  To speed things up, 
:::

```{code-cell}
# Here we use a custom `animate` function defined at the start of this file.
# This requires the following coroutine
def draw_frame(f=f, res=res):
    # First initialize the figure and any other computations
    fig, axs = plt.subplots(2, 1, figsize=(4, 3), 
                            height_ratios=(3, 1),
                            gridspec_kw=dict(hspace=0))
    ax0, axV = axs
    ax1 = ax0.twinx()
    
    m_scale = f.m
    n_scale = f.n0
    x_scale = 1/f.n0
    v_scale = abs(f.v0)
    E_scale = f.m * v_scale**2
    t_scale = 1 / (n_scale * v_scale)

    x, ts = f.x, res.t
    ns, us = res.y.reshape((2, len(x), len(ts)))

    x_ = x / x_scale
    ns_ = ns.T / n_scale
    us_ = us.T / v_scale
    ts_ = ts / t_scale

    # Since we will be updating the data, we must scale ylims at the start.
    ax0.set(xlabel="$x$ [$1/n_0$]",
            ylabel="$n$ [$n_0$]",
            ylim=(0, ns_.max()))
    ax1.set(ylabel="$u$ [$v_0$]", ylim=(us_.min(), us_.max()))
    axV.set(ylabel="$V$ [$mv_0^2$]")
    
    # We first return the figure as the artists first.
    artists = fig
    
    while True:
        # REQUIRED: The first time through, you must yield the figure `fig`.  After 
        # this, if `blit==True`, then you must yield a list of the artists.
        # The results yielded by `get_data` will be passed in here, so `data` must
        # recieve these properly. 
        frame = yield artists

        # Now compute the things to plot.
        n_ = ns_[frame]
        u_ = us_[frame]
        V_ = f.get_Vext(x, t=ts[frame])/E_scale
        t_ = ts_[frame]
        title = f"$t={t_:.2f}$ [$1/j_0$]"
        
        if artists is fig:
            # First time through, plot the data and save the artists
            line_n, = ax0.plot(x_, n_)
            line_u, = ax1.plot(x_, u_, 'C1:')
            line_V, = axV.plot(x_, V_)
            text_title = ax0.set_title(title)
            artists = [line_n, line_u, line_V, text_title]
            continue
    
        # Every other time, update these
        line_n.set_data(x_, n_)
        line_u.set_data(x_, u_)
        line_V.set_data(x_, V_)
        text_title.set_text(title)


%time anim = animate(draw_frame, len(res.t), interval=10, repeat=False, close=False)
```

:::{note}
There are two options for viewing these in an HTML document: as an html5 video, or as a
JS video.  The latter is larger and slowe, so we use the former with some modifications
for the width:

```{code-cell}
%time anim_js = anim.to_jshtml()
%time anim_html = anim.to_html5_video()
print(f"JS:   {sys.getsizeof(anim_js) / 1024**2:4.1f} MiB")
print(f"HTML: {sys.getsizeof(anim_html) / 1024**2:4.1f} MiB")
```

```{code-cell}
from IPython.display import HTML
display(HTML(anim_html))
display(HTML(anim_js))
```

:::{margin}
This is poorly documented: see, e.g., [this
code in
`animation.py`](https://github.com/matplotlib/matplotlib/blob/d1fb1bebc1d0a5bf13cfba3c6b775789b85d3f4e/lib/matplotlib/animation.py#L1378)
where the parameter is dispatched.
:::
If you want to do this for all animations, you can set the following parameter.

```{code-cell}
#plt.rcParams["animation.html"] = "jshtml"
plt.rcParams["animation.html"] = "html5"
display(anim)
```


## Another Approach
An alternative implementation replaces the velocity with the current $j=nu$:
\begin{gather*}
  \dot{n} = -j_{,x},\\
  \dot{j} = \dot{n}u + n\dot{u}
          = -\left(\frac{j^2}{n}\right)_{,x}
          - n\frac{V'(x) + \mathcal{E}''(n)n_{,x}}{m} 
    + \left(\nu(n) \left(\frac{j}{n}\right)_{,x}\right)_{,x}.
\end{gather*}
This is implemented as
\begin{gather*}
  \diff{}{t}
  \begin{pmatrix}
    n\\
    j
  \end{pmatrix}
  =
  \begin{pmatrix}
    -j_{,x}\\
    - \frac{2jj_{,x}}{n} + \frac{j^2}{n^2}n_{,x}
    - n\frac{V'(x) + \mathcal{E}''(n)n_{,x}}{m} 
    + \nu'(n)n_{,x}u_{,x}
    + \nu(n)u_{,xx}
  \end{pmatrix},\\
  u_{,x} = \frac{j_{,x}}{n} - \frac{j}{n^2}n_{,x}, \qquad
  u_{,xx} = \frac{j_{,xx}}{n} - \frac{2}{n^2}j_{,x}n_{,x}
          + 2\frac{j}{n^3}n_{,x}^2 - \frac{j}{n^2}n_{,xx}.
\end{gather*}


\begin{gather*}
  j = nu\\
  j_{,x} = n_{,x}u + nu_{,x}, \qquad
  u_{,x} = \frac{nj_{,x} - jn_{,x}}{n^2},\\
  \dot{j} = \dot{n}u + n\dot{u}
          = -\left(\frac{j^2}{n}\right)_{,x}
          - n\frac{V'(x) + \mathcal{E}''(n)n_{,x}}{m} 
    + \left(\nu(n) \left(\frac{j}{n}\right)_{,x}\right)_{,x}
\end{gather*}

[step 4]: <https://nbviewer.org/github/barbagroup/CFDPython/blob/master/lessons/05_Step_4.ipynb>