# Further Reading: Full-Waveform Inversion (FWI)

This notebook is an adaption of `examples/seismic/tutorials/03_fwi.ipynb`. In this example we aim to highlight how GemPy-generated models can be used in more involved applications, and the core ideas behind seismic inversion, where we create an image of the subsurface from field recorded data. This tutorial will use the model created in the main tutorial, and will reuse many of the concepts.

## Inversion requirement

Seismic inversion relies on two known parameters:

- **Field data** - or also called **recorded data**. This is a shot record corresponding to the true velocity model. In practice this data is acquired as described in the first tutorial. In order to simplify this tutorial we will fake field data by modelling it with the true velocity model.

- **Initial velocity model**. This is a velocity model that has been obtained by processing the field data. This model is a rough and very smooth estimate of the velocity as an initial estimate for the inversion. This is a necessary requirement for any optimization (method).

## Inversion computational setup

In this tutorial, we will introduce the gradient operator. This operator corresponds to the imaging condition introduced in the previous tutorial with some minor modifications that are defined by the objective function (also referred to in the tutorial series as the *functional*, *f*) and its gradient, *g*. We will define these two terms in the tutorial too.

## Notes on the operators

As we have already described the creation of a forward modelling operator, we will only call a wrapper function here. This wrapper already contains all the necessary operators for seismic modeling, imaging and inversion. Operators introduced for the first time in this tutorial will be properly described.

```python
import numpy as np
%matplotlib inline

from devito import configuration
configuration['log-level'] = 'WARNING'
```

## Computational considerations

As we will see, FWI is computationally extremely demanding, even more than RTM. To keep this tutorial as lightwight as possible we therefore again use a very small demonstration model. We also define here a few parameters for the final example runs that can be changed to modify the overall runtime of the tutorial.

```python
nshots = 9  # Number of shots to create gradient from
nreceivers = 101  # Number of receiver locations per shot
fwi_iterations = 25  # Number of outer FWI iterations
```

# True and smooth velocity models

We will use a very simple model domain, consisting of a circle within a 2D domain. We will again use the "true" model to generate our synthetic shot data and use a "smooth" model as our initial guess. In this case the smooth model is very smooth indeed - it is simply a constant background velocity without any features.

```python
from examples.seismic import Model, demo_model, plot_velocity, plot_perturbation
from scipy.ndimage import gaussian_filter

# Define true and initial model
shape = (101, 101)  # Number of grid point (nx, nz)
spacing = (10., 10.)  # Grid spacing in m. The domain size is now 1km by 1km
origin = (0., 0.)  # Need origin to define relative source and receiver locations

gempy_vp = np.load("model_data.npy")[:, ::-1]
gempy_vp = np.pad(gempy_vp[:, :-10], ((0, 0), (10, 0)), 'edge')  # Increase the water depth a little
# Load in the model we created earlier.
# Note we need to flip it upside down, as the origin is now at the top
model = Model(vp=gempy_vp, origin=origin, shape=shape, spacing=spacing,
              nbl=40, space_order=2, bcs="damp")

smooth_vp = gempy_vp
smooth_vp[:, 8:] = gaussian_filter(smooth_vp[:, 8:], sigma=(15, 8))

model0 = Model(vp=smooth_vp, dt=model.critical_dt,
               origin=origin, shape=shape, spacing=spacing, nbl=40, space_order=2,
               bcs="damp", grid=model.grid)

# Make a separate inital for final plotting
modeli = Model(vp=smooth_vp, dt=model.critical_dt,
               origin=origin, shape=shape, spacing=spacing, nbl=40, space_order=2,
               bcs="damp", grid=model.grid)

plot_velocity(model)
plot_velocity(model0)

# Note, currently has a bug meaning that padding is not trimmed off, hence the weird edges
plot_perturbation(model0, model)
```

```python
assert model.grid == model0.grid
assert model.vp.grid == model0.vp.grid
```

## Acquisition geometry

In this tutorial, we will use the easiest case for inversion, namely a transmission experiment. The sources are located on one side of the model and the receivers on the other side. This allows to record most of the information necessary for inversion, as reflections usually lead to poor inversion results.

```python
#NBVAL_IGNORE_OUTPUT
# Define acquisition geometry: source
from examples.seismic import AcquisitionGeometry

t0 = 0.
tn = 1000.
f0 = 0.008
# First, position source centrally in all dimensions, then set depth
src_coordinates = np.empty((1, 2))
src_coordinates[0, :] = np.array(model.domain_size) * .5
src_coordinates[0, 1] = 20.  # Depth is 20m

# Define acquisition geometry: receivers

# Initialize receivers for synthetic and imaging data
rec_coordinates = np.empty((nreceivers, 2))
rec_coordinates[:, 0] = np.linspace(0, model.domain_size[0], num=nreceivers)
rec_coordinates[:, 1] = 20.

# Geometry

geometry = AcquisitionGeometry(model, rec_coordinates, src_coordinates, t0, tn, f0=f0, src_type='Ricker')
# We can plot the time signature to see the wavelet
geometry.src.show()
```

```python
# Plot acquisition geometry
plot_velocity(model, source=geometry.src_positions,
              receiver=geometry.rec_positions[::4, :])
```

## True and smooth data

We can generate shot records for the true and smoothed initial velocity models, since the difference between them will again form the basis of our imaging procedure.

```python
# Compute synthetic data with forward operator
from examples.seismic.acoustic import AcousticWaveSolver

solver = AcousticWaveSolver(model, geometry, space_order=4)
true_d, _, _ = solver.forward(vp=model.vp)
```

```python
# Compute initial data with forward operator
smooth_d, _, _ = solver.forward(vp=model0.vp)
```

```python
#NBVAL_IGNORE_OUTPUT
from examples.seismic import plot_shotrecord

# Plot shot record for true and smooth velocity model and the difference
plot_shotrecord(true_d.data, model, t0, tn)
plot_shotrecord(smooth_d.data, model, t0, tn)
plot_shotrecord(smooth_d.data - true_d.data, model, t0, tn)
```

# Full-Waveform Inversion

## Formulation

Full-waveform inversion (FWI) aims to invert an accurate model of the discrete wave velocity, $\mathbf{c}$, or equivalently the square slowness of the wave, $\mathbf{m} = \frac{1}{\mathbf{c}^2}$, from a given set of measurements of the pressure wavefield $\mathbf{u}$. This can be expressed as the following optimization problem [1, 2]:

\begin{aligned}
	\mathop{\hbox{minimize}}_{\mathbf{m}} \Phi_s(\mathbf{m})&=\frac{1}{2}\left\lVert\mathbf{P}_r
	\mathbf{u} - \mathbf{d}\right\rVert_2^2 \\
	\mathbf{u} &= \mathbf{A}(\mathbf{m})^{-1} \mathbf{P}_s^T \mathbf{q}_s,
\end{aligned}

where $\mathbf{P}_r$ is the sampling operator at the receiver locations, $\mathbf{P}_s^T$ is the injection operator at the source locations, $\mathbf{A}(\mathbf{m})$ is the operator representing the discretized wave equation matrix, $\mathbf{u}$ is the discrete synthetic pressure wavefield, $\mathbf{q}_s$ is the corresponding pressure source and $\mathbf{d}$ is the measured data. It is worth noting that $\mathbf{m}$ is the unknown in this formulation and that multiple implementations of the wave equation operator $\mathbf{A}(\mathbf{m})$ are possible.

We have already defined a concrete solver scheme for $\mathbf{A}(\mathbf{m})$ in the first tutorial, including appropriate implementations of the sampling operator $\mathbf{P}_r$ and source term $\mathbf{q}_s$.

To solve this optimization problem using a gradient-based method, we use the
adjoint-state method to evaluate the gradient $\nabla\Phi_s(\mathbf{m})$:

\begin{align}
 \nabla\Phi_s(\mathbf{m})=\sum_{\mathbf{t} =1}^{n_t}\mathbf{u}[\mathbf{t}] \mathbf{v}_{tt}[\mathbf{t}] =\mathbf{J}^T\delta\mathbf{d}_s,
\end{align}

where $n_t$ is the number of computational time steps, $\delta\mathbf{d}_s = \left(\mathbf{P}_r \mathbf{u} - \mathbf{d} \right)$ is the data residual (difference between the measured data and the modelled data), $\mathbf{J}$ is the Jacobian operator and $\mathbf{v}_{tt}$ is the second-order time derivative of the adjoint wavefield solving:

\begin{align}
 \mathbf{A}^T(\mathbf{m}) \mathbf{v} = \mathbf{P}_r^T \delta\mathbf{d}.
\end{align}

We see that the gradient of the FWI function is the previously defined imaging condition with an extra second-order time derivative. We will therefore reuse the operators defined previously inside a Devito wrapper.

## FWI gradient operator

To compute a single gradient $\nabla\Phi_s(\mathbf{m})$ in our optimization workflow we again use `solver.forward` to compute the entire forward wavefield $\mathbf{u}$ and a similar pre-defined gradient operator to compute the adjoint wavefield `v`. The gradient operator provided by our `solver` utility also computes the correlation between the wavefields, allowing us to encode a similar procedure to the previous imaging tutorial as our gradient calculation:

- Simulate the forward wavefield with the background velocity model to get the synthetic data and save the full wavefield $\mathbf{u}$
- Compute the data residual
- Back-propagate the data residual and compute on the fly the gradient contribution at each time step.

This procedure is applied to multiple source positions and summed to obtain a gradient image of the subsurface. We again prepare the source locations for each shot and visualize them, before defining a single gradient computation over a number of shots as a single function.

```python
#NBVAL_IGNORE_OUTPUT

# Prepare the varying source locations sources
source_locations = np.empty((nshots, 2), dtype=np.float32)
source_locations[:, 1] = 30.
source_locations[:, 0] = np.linspace(0., 1000, num=nshots)

plot_velocity(model, source=source_locations)
```

```python
from devito import Eq, Operator

# Computes the residual between observed and synthetic data into the residual
def compute_residual(residual, dobs, dsyn):
    if residual.grid.distributor.is_parallel:
        # If we run with MPI, we have to compute the residual via an operator
        # First make sure we can take the difference and that receivers are at the
        # same position
        assert np.allclose(dobs.coordinates.data[:], dsyn.coordinates.data)
        assert np.allclose(residual.coordinates.data[:], dsyn.coordinates.data)
        # Create a difference operator
        diff_eq = Eq(residual, dsyn.subs({dsyn.dimensions[-1]: residual.dimensions[-1]}) -
                               dobs.subs({dobs.dimensions[-1]: residual.dimensions[-1]}))
        Operator(diff_eq)()
    else:
        # A simple data difference is enough in serial
        residual.data[:] = dsyn.data[:] - dobs.data[:]

    return residual
```

```python
# Create FWI gradient kernel
from devito import Function, TimeFunction, norm
from examples.seismic import Receiver

import scipy
def fwi_gradient(vp_in):
    # Create symbols to hold the gradient
    grad = Function(name="grad", grid=model.grid)
    # Create placeholders for the data residual and data
    residual = Receiver(name='residual', grid=model.grid,
                        time_range=geometry.time_axis,
                        coordinates=geometry.rec_positions)
    d_obs = Receiver(name='d_obs', grid=model.grid,
                     time_range=geometry.time_axis,
                     coordinates=geometry.rec_positions)
    d_syn = Receiver(name='d_syn', grid=model.grid,
                     time_range=geometry.time_axis,
                     coordinates=geometry.rec_positions)
    objective = 0.
    for i in range(nshots):
        # Update source location
        geometry.src_positions[0, :] = source_locations[i, :]

        # Generate synthetic data from true model
        _, _, _ = solver.forward(vp=model.vp, rec=d_obs)

        # Compute smooth data and full forward wavefield u0
        _, u0, _ = solver.forward(vp=vp_in, save=True, rec=d_syn)

        # Compute gradient from data residual and update objective function
        compute_residual(residual, d_obs, d_syn)

        objective += .5*norm(residual)**2
        solver.gradient(rec=residual, u=u0, vp=vp_in, grad=grad)
    grad.data[:, :model.nbl+3] = 0
    grad.data[:, -model.nbl-3:] = 0
    return objective, grad
```

Having defined our FWI gradient procedure we can compute the initial iteration from our starting model. This allows us to visualize the gradient alongside the model perturbation and the effect of the gradient update on the model.

```python
# Compute gradient of initial model
ff, update = fwi_gradient(model0.vp)
assert np.isclose(ff, 57283, rtol=1e0)
```

```python
#NBVAL_IGNORE_OUTPUT
from devito import mmax
from examples.seismic import plot_image

# Plot the FWI gradient
plot_image(-update.data, vmin=-1e2, vmax=1e2, cmap="jet")

# Plot the difference between the true and initial model.
# This is not known in practice as only the initial model is provided.
plot_image(model0.vp.data - model.vp.data, vmin=-1e-1, vmax=1e-1, cmap="jet")

# Show what the update does to the model
alpha = 2.5 / mmax(update)
plot_image(model0.vp.data + alpha*update.data, vmin=2.5, vmax=3.0, cmap="jet")
```

We see that the gradient and the true perturbation have the same sign, therefore, with an appropriate scaling factor, we will update the model in the correct direction.

```python
from sympy import Min, Max
# Define bounding box constraints on the solution.
def update_with_box(vp, alpha, dm, vmin=1.5, vmax=3.5):
    """
    Apply gradient update in-place to vp with box constraint

    Notes:
    ------
    For more advanced algorithm, one will need to gather the non-distributed
    velocity array to apply constrains and such.
    """
    update = vp + alpha * dm
    update_eq = Eq(vp, Max(Min(update, vmax), vmin))
    Operator(update_eq)()
```

```python
#NBVAL_SKIP

from devito import mmax, mmin

# Run FWI with gradient descent
history = np.zeros((fwi_iterations, 1))
for i in range(0, fwi_iterations):
    # Compute the functional value and gradient for the current
    # model estimate
    phi, direction = fwi_gradient(model0.vp)

    # Store the history of the functional values
    history[i] = phi

    # Artificial Step length for gradient descent
    # In practice this would be replaced by a Linesearch (Wolfe, ...)
    # that would guarantee functional decrease Phi(m-alpha g) <= epsilon Phi(m)
    # where epsilon is a minimum decrease constant
    alpha = .01 / mmax(direction)

    # Update the model estimate and enforce minimum/maximum values
    update_with_box(model0.vp , alpha , direction, vmin=mmin(model.vp), vmax=mmax(model.vp))

    # Log the progress made
    print('Objective value is %f at iteration %d' % (phi, i+1))

```

    Objective value is 9251.535154 at iteration 1
    Objective value is 8832.330789 at iteration 2
    Objective value is 8471.526218 at iteration 3
    Objective value is 8148.875188 at iteration 4
    Objective value is 7854.766629 at iteration 5
    Objective value is 7583.467015 at iteration 6
    Objective value is 7331.128043 at iteration 7
    Objective value is 7095.208614 at iteration 8
    Objective value is 6874.240803 at iteration 9
    Objective value is 6666.312804 at iteration 10
    Objective value is 6470.001998 at iteration 11
    Objective value is 6283.991342 at iteration 12
    Objective value is 6107.163216 at iteration 13
    Objective value is 5938.679140 at iteration 14
    Objective value is 5777.663463 at iteration 15
    Objective value is 5623.487279 at iteration 16
    Objective value is 5475.471031 at iteration 17
    Objective value is 5333.121586 at iteration 18
    Objective value is 5195.936796 at iteration 19
    Objective value is 5063.528480 at iteration 20
    Objective value is 4935.483044 at iteration 21
    Objective value is 4811.506605 at iteration 22
    Objective value is 4691.246826 at iteration 23
    Objective value is 4574.480522 at iteration 24
    Objective value is 4460.937443 at iteration 25

```python
#NBVAL_IGNORE_OUTPUT

# Plot inverted velocity model
plot_velocity(model0)
plot_velocity(modeli)
plot_velocity(model)
```

```python
#NBVAL_SKIP
import matplotlib.pyplot as plt

# Plot objective function decrease
plt.figure()
plt.loglog(history)
plt.xlabel('Iteration number')
plt.ylabel('Misift value Phi')
plt.title('Convergence')
plt.show()
```

## References

[1] _Virieux, J. and Operto, S.: An overview of full-waveform inversion in exploration geophysics, GEOPHYSICS, 74, WCC1–WCC26, doi:10.1190/1.3238367, http://library.seg.org/doi/abs/10.1190/1.3238367, 2009._

[2] _Haber, E., Chung, M., and Herrmann, F. J.: An effective method for parameter estimation with PDE constraints with multiple right hand sides, SIAM Journal on Optimization, 22, http://dx.doi.org/10.1137/11081126X, 2012._

This notebook is part of the tutorial "Optimised Symbolic Finite Difference Computation with Devito" presented at the Intel® HPC Developer Conference 2017.


###### Text provided under a Creative Commons Attribution license, CC-BY. All code is made available under the FSF-approved MIT license. (c) Daniel Koehn based on Jupyter notebooks by Marc Spiegelman [Dynamical Systems APMA 4101](https://github.com/mspieg/dynamical-systems) and  Kyle Mandli from his course [Introduction to numerical methods](https://github.com/mandli/intro-numerical-methods), notebook style sheet by L.A. Barba, N.C. Clementi [Engineering Computations](https://github.com/engineersCode)

```python
# Execute this cell to load the notebook's style sheet, then ignore it
from IPython.core.display import HTML
css_file = '../style/custom.css'
HTML(open(css_file, "r").read())
```

<link href="https://fonts.googleapis.com/css?family=Merriweather:300,300i,400,400i,700,700i,900,900i" rel='stylesheet' >
<link href="https://fonts.googleapis.com/css?family=Source+Sans+Pro:300,300i,400,400i,700,700i" rel='stylesheet' >
<link href='http://fonts.googleapis.com/css?family=Source+Code+Pro:300,400' rel='stylesheet' >
<style>

@font-face {
    font-family: "Computer Modern";
    src: url('http://mirrors.ctan.org/fonts/cm-unicode/fonts/otf/cmunss.otf');
}

#notebook_panel { /* main background */
    background: rgb(245,245,245);
}

div.cell { /* set cell width */
    width: 800px;
}

div #notebook { /* centre the content */
    background: #fff; /* white background for content */
    width: 1000px;
    margin: auto;
    padding-left: 0em;
}

#notebook li { /* More space between bullet points */
margin-top:0.5em;
}

/* draw border around running cells */
div.cell.border-box-sizing.code_cell.running {
    border: 1px solid #111;
}

/* Put a solid color box around each cell and its output, visually linking them*/
div.cell.code_cell {
    background-color: rgb(256,256,256);
    border-radius: 0px;
    padding: 0.5em;
    margin-left:1em;
    margin-top: 1em;
}

div.text_cell_render{
    font-family: 'Source Sans Pro', sans-serif;
    line-height: 140%;
    font-size: 110%;
    width:680px;
    margin-left:auto;
    margin-right:auto;
}

/* Formatting for header cells */
.text_cell_render h1 {
    font-family: 'Merriweather', serif;
    font-style:regular;
    font-weight: bold;
    font-size: 250%;
    line-height: 100%;
    color: #004065;
    margin-bottom: 1em;
    margin-top: 0.5em;
    display: block;
}
.text_cell_render h2 {
    font-family: 'Merriweather', serif;
    font-weight: bold;
    font-size: 180%;
    line-height: 100%;
    color: #0096d6;
    margin-bottom: 0.5em;
    margin-top: 0.5em;
    display: block;
}

.text_cell_render h3 {
    font-family: 'Merriweather', serif;
	font-size: 150%;
    margin-top:12px;
    margin-bottom: 3px;
    font-style: regular;
    color: #008367;
}

.text_cell_render h4 {    /*Use this for captions*/
    font-family: 'Merriweather', serif;
    font-weight: 300;
    font-size: 100%;
    line-height: 120%;
    text-align: left;
    width:500px;
    margin-top: 1em;
    margin-bottom: 2em;
    margin-left: 80pt;
    font-style: regular;
}

.text_cell_render h5 {  /*Use this for small titles*/
    font-family: 'Source Sans Pro', sans-serif;
    font-weight: regular;
    font-size: 130%;
    color: #e31937;
    font-style: italic;
    margin-bottom: .5em;
    margin-top: 1em;
    display: block;
}

.text_cell_render h6 { /*use this for copyright note*/
    font-family: 'Source Code Pro', sans-serif;
    font-weight: 300;
    font-size: 9pt;
    line-height: 100%;
    color: grey;
    margin-bottom: 1px;
    margin-top: 1px;
}

    .CodeMirror{
            font-family: "Source Code Pro";
			font-size: 90%;
    }
/*    .prompt{
        display: None;
    }*/

    .warning{
        color: rgb( 240, 20, 20 )
        }
</style>

# Simplified Convection Problem: The Lorenz Equations

The Lorenz Equations are a 3-D dynamical system that is a simplified model of Rayleigh-Benard thermal convection.  They are derived and described in detail in Edward Lorenz' 1963 paper [Deterministic Nonperiodic Flow](http://journals.ametsoc.org/doi/pdf/10.1175/1520-0469%281963%29020%3C0130%3ADNF%3E2.0.CO%3B2)  in the Journal of Atmospheric Science.

Here we will just sketch out the key points of the derivation.  A more complete derivation can be found [here](https://www.math.uni-hamburg.de/home/lauterbach/scripts/seminar03/prill.pdf)

The key idea is that the Lorenz Equations result from a severely truncated spectral approximation to the 2-D equations for incompressible thermal convection in stream-function/vorticity form.  These equations govern the flow of a buouyant incompressible fluid with a temperature dependent density in a layer of depth $h$, that is heated from below and cooled from the top.

## Governing Equations

The full coupled set of scaled PDE's describe the coupling of incompressible 2D Navier Stokes flow with an advection-diffusion equation for temperature, and can be written in dimensionless form as,

$$
   \frac{1}{\mathrm{Pr}}\left[ \frac{\partial \omega}{\partial t} + \vec{v}\cdot\nabla\omega\right] = \nabla^2\omega + \mathrm{Ra}\frac{\partial T}{\partial x}
$$
$$
    \nabla^2 \psi = -\omega
$$
$$
    \frac{\partial T}{\partial t} + \vec{v}\cdot\nabla T = \nabla^2 T
$$

where

$$
    \vec{v}=(u,0,w) = \nabla\times\psi\hat{j}=(-\frac{\partial\psi}{\partial z}, 0,  \frac{\partial\psi}{\partial x})
$$

is the fluid velocity field (which in this form is exactly incompressible with $\nabla\cdot\vec{v}=0$).  $\psi$ is the "Streamfunction"  whose contours are tangent to the fluid trajectories at all times.  The vorticity,

$$
    \omega = \bf{\vec{\omega} \cdot \hat{j}} = (\nabla\times\vec{v}) \cdot \hat{j}
$$

measures the local rate of rotation of the fluid, and is driven by horizontal variations in temperature (actually density). $\hat{j}=(0,1,0)^T$ denotes the unit vector in y-direction.

Boundary conditions for temperature are $T=1$ on the bottom of the layer and $T=0$ on the top.  In the absence of any fluid motion ($\omega=\vec{v}=0$),  the temperature field is just a steady conductive ramp with

$$
    T = 1 - z
$$

Thus we can also solve for the perturbation $\theta(x,z,t)$ away from this steady state by substituting

$$
    T = 1 - z + \theta(x,z,t)
$$

into the energy equation to solve for the perturbed temperature using

$$
    \frac{\partial \theta}{\partial t} + \vec{v}\cdot\nabla \theta = \nabla^2\theta + w
$$

### Parameters
In dimensionless form,  these equations have two important dimensionless numbers that control the structure and behavior of the convection.

**The Prandtl Number**

The first is the "Prandtl Number", $\mathrm{Pr} = \frac{\nu}{\kappa}$  which is the ratio of the fluid viscosity $\nu$ to the thermal diffusivitiy $\kappa$.  Since both vorticity and temperature both obey advection diffusion equations (and viscosity acts to diffuse momentum/vorticity),  the Prandtl number is a measure of whether momemtum or energy is more dissipative.

**The Rayleigh Number**

The second key parameter is the Rayleigh number

$$
    \mathrm{Ra} = \frac{g\alpha(T_1 - T_0)h^3}{\nu\kappa}
$$

which measures the balance of forces that drive convection (i.e. gravity, or temperature differences), to those that damp convection such as viscosity and thermal diffusivity.  Systems with large Rayleigh numbers are prone to vigorous convection.  However, it was shown by Rayleigh, that there is a critical value of the Rayleigh Number $\mathrm{Ra}_c$ below which there is no convection.  This value depends on the size of the convection cell and boundary conditions for stress on the fluid,  however, for the simplest case of a layer with no-slip top and bottom boundary conditions and cell with aspect ratio  $a=h/L$ (with $h$ the layer depth and $L$ the width of the convection cell), then the critical Ra number is

$$
\mathrm{Ra}_c = \pi^4 (1 + a^2)^3/a^2
$$

which has a minimum value for $a_{min}^2=1/2$.

```python
# load libraries
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from pylab import rcParams
from pylab import rcParams
from matplotlib import rc

# define font size
FSize = 18
font = {'color':  'black',
        'weight': 'normal',
        'size': FSize}
mpl.rc('xtick', labelsize=FSize)
mpl.rc('ytick', labelsize=FSize)
```

```python
# Show Ra vs a
a = np.linspace(0.01,2.)
Rc = np.pi**4*(1. + a**2)**3/a**2
plt.figure()
plt.semilogy(a,Rc)
amin = np.sqrt(1./2.)
Rcmin = np.pi**4*(1. + amin**2)**3/amin**2
plt.semilogy(amin,Rcmin,'ro')
plt.xlabel('a', fontdict=font)
plt.ylabel('Ra$_c$', fontdict=font)
plt.title('Critical Rayleigh Number', fontdict=font)
plt.show()
```

## Spectral decomposition

Next, we expand the streamfunction and temperature fields in terms of a highly truncated Fourier Series where the streamfunction contains one cellular mode

$$
    \psi(x,z,t) = X(t)\sin(a\pi x)\sin(\pi z)
$$

and temperature has two modes

$$
    \theta(x,z,t) = Y(t)\cos(a\pi x)\sin(\pi z) - Z(t)\sin(2\pi z)
$$

Here, $X(t)$, $Y(t)$ and $Z(t)$ are the time dependent amplitudes of each mode. Defining

\begin{equation}
\begin{split}
\phi(x,z) &= \sin(a\pi x)\sin(\pi z)\\
\theta_0(x,z) &= \cos(a\pi x)\sin(\pi z)\\
\theta_1(z) &= -\sin(2\pi z)
\end{split}
\notag
\end{equation}

streamfunction and temperature fields simplify to:

\begin{equation}
\begin{split}
\psi(x,z,t) &= X(t) \phi(x,z)\\
\theta(x,z,t) &= Y(t) \theta_0(x,z) + Z(t) \theta_1(z)\\
\end{split}
\notag
\end{equation}

Using $a = a_{min} = \sqrt{0.5}$ for the minimum critical Rayleigh number $Ra_c^{min}$, the spatial components of each mode looks like

```python
a = np.sqrt(0.5)
x0 = np.linspace(0,1./a)
z0 = np.linspace(0.,1.)
x,z = np.meshgrid(x0,z0)

psi = np.sin(a*np.pi*x)*np.sin(np.pi*z)
theta0 = np.cos(a*np.pi*x)*np.sin(np.pi*z)
theta1 = -np.sin(2.*np.pi*z)

plt.figure()
plt.figure(figsize=(10,8))

plt.subplot(2,2,1)
plt.contourf(x,z,psi)
plt.title(r'$\Phi = sin(a \pi x)sin(\pi z)$', fontdict=font)
plt.gca().set_aspect('equal')

plt.subplot(2,2,3)
plt.contourf(x,z,theta0)
plt.title(r'$\theta_0 = cos(a \pi x)sin(\pi z)$', fontdict=font)
plt.gca().set_aspect('equal')

plt.subplot(2,2,4)
plt.contourf(x,z,theta1)
plt.title(r'$\theta_1 = -sin(2 \pi z)$', fontdict=font)
plt.gca().set_aspect('equal')

plt.tight_layout()
plt.show()
```

For our initial conditions $X(0) = 2$, $Y(0) = 3$, $Z(0) = 4$, the streamfunction $\psi(x,z,t)$ and temperature field $\theta(x,z,t)$

\begin{equation}
\begin{split}
\psi(x,z,t) &= X(t)\sin(a\pi x)\sin(\pi z)\\
\theta(x,z,t) & = Y(t)\cos(a\pi x)\sin(\pi z) - Z(t)\sin(2\pi z)\\
\end{split}
\notag
\end{equation}

become

\begin{equation}
\begin{split}
\psi(x,z,0) &= 2\sin(a\pi x)\sin(\pi z)\\
\theta(x,z,0) &= 3\cos(a\pi x)\sin(\pi z) - 4\sin(2\pi z)\\
\end{split}
\notag
\end{equation}

```python
# Define figure size
plt.figure()
plt.figure(figsize=(20,20))

# Initial Streamfunction psi and velocity field
plt.subplot(1,2,1)
plt.contourf(x,z,2.*psi,cmap='viridis_r')

# Velocity field
U = - 2. * np.pi * np.sin(a*np.pi*x) * np.cos(np.pi*z)
V = 2. * a * np.pi * np.cos(a*np.pi*x) * np.sin(np.pi*z)
plt.quiver(x,z,U,V)
plt.gca().set_aspect('equal')
plt.title('Streamfunction $\psi$', fontdict=font)

# Initial temperature field
plt.subplot(1,2,2)
plt.contourf(x,z,3*theta0 + 4*theta1,cmap='magma')
plt.gca().set_aspect('equal')
plt.title(r'Temperature $\theta$', fontdict=font)

plt.tight_layout()
plt.show()
```

## Lorenz equations

Inserting the spectral decomposition solutions into the momentum and energy conservation equations lead to the **Lorenz equations**

\begin{equation}
\begin{split}
\frac{\partial X}{\partial t} &= \sigma( Y - X)\\
\frac{\partial Y}{\partial t} &= rX - Y - XZ \\
\frac{\partial Z}{\partial t} &= XY -b Z
\end{split}
\notag
\end{equation}

where $\sigma$ denotes the "Prandtl number",  $r = \mathrm{Ra}/\mathrm{Ra}_c$ is a scaled "Rayleigh number" and $b$ is a parameter that is related to the the aspect ratio of a convecting cell in the original derivation.

## What we learned:

- The basic concept of the Lorenz equations to approximate the governing equations for thermal convection, described by the incompressible 2D Navier Stokes flow with an advection-diffusion equation for temperature.