<form action="javascript:code_toggle()"><input type="submit" id="toggleButton" value="Show Code"></form>
# An attempt to work out different rooting depths and distributions for savanna sites in Northern Australia
----------

The below technique follows from the SPA version 1 calibration spreadsheet for calculating the root profile of a canopy. My madness on screen as displayed here is to determine whether I can build this in to a parameter estimation workflow in some emulator work saved elsewhere.

```python
# Libraries to import

from scipy.optimize import minimize
from scipy.optimize import leastsq

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

%matplotlib inline
```

```python
# Notebook plot settings

plt.style.use(['fivethirtyeight'])
plt.rcParams['figure.figsize'] = (14, 6)
```

### Describing the root profile

The distribution of root biomass follows the following conditional equation:

\begin{equation}
r_{mass, z}=
    \begin{cases}
        r_{mass, surf}\left(\frac{-\exp({k\space d_{z+1}})\space+\space\exp({k\space d_{z}})}{k}\right), & \text{if}\ z < z_{root, max} \\\\
        0, & \text{else}
    \end{cases}
\end{equation}

where $z$ is the soil layer, $r_{mass, z}$ is the root biomass (in grams) across some arbitrary depth of soil $d_{z} \rightarrow d_{z+1}$ that describes the volume of the current soil layer, $r_{mass, surf}$ is the density of root biomass at the soil surface, $k$ is the decline coefficient, and $z_{root, max}$ is the maximum rooting depth in terms of soil layer.

The density of roots for each layer is then determined as:

\begin{equation}
    r_{dens, z} = \frac{r_{mass, z}}{z_{thick}}
\end{equation}

And finally the fraction of roots at some aribtrary depth $z$ is determed as:

\begin{equation}
    r_{frac, z} = \frac{r_{dens, z}}{r_{mass, total}}
\end{equation}

The decline coefficient seems to be site-specific and this is pre-determined via optimising on it ala:

\begin{equation}
    \chi_{\min}^{2} = \min\left( r_{mass, total} - \sum_{z=0}^{N} r_{mass, z}\right)
\end{equation}

We can now use these equations to come up with a set of shallow-to-deep root profiles based on collected site information.

### Example case: Howard Springs

Very little information has been collected on root distributions along the NATT, with the only noteable cases being Howard Springs (see Chen et al. 2002; and Kelley et al. 2007) - could be wrong though. Prior information to work with is therefore limited but can be used to get a ballpark idea of the root profile:

$r_{mass, total} = 1930 \space g\space m^{-2}$

Pretty much everything else is determined by tweaking the model. From previous work the following root profile information has been used.

$r_{mass, surf} = 14000 \space g\space m^{-3}$

$z_{root, max} = 6.0 \space m$

Soil texture information is extracted from the Australian Digital Soil Atlas, and for this site A/B horizon for the soil profile is:

**Horizon A:** 55 % sand | 5 % clay

**Horizon B:** 45 % sand | 10 % clay

```python
# default number of soil layers
nlayer = 20

# create 20 soil layers
soil_layer = ["layer_{0}".format(i+1) for i in (range(nlayer))] + ['core']

# set the thicknesses of the 20 soil layers, divide accordingly to match with common field measurements
soil_thick = [0.1 if z < 4 else 0.2 if (z >= 4) & (z < 12) \
                else 0.1 if (z == len(soil_layer) - 1) else 1.0 \
                  for (z, _) in enumerate(soil_layer)]

soil_depth = np.cumsum([0] + soil_thick)

# create a dataframe to work with
soil_df = pd.DataFrame({'thick': soil_thick, 'depth': soil_depth[:-1], 'layer': soil_layer})
soil_df.set_index(['layer'], inplace=True)

print(soil_df)
```

              depth  thick
    layer
    layer_1     0.0    0.1
    layer_2     0.1    0.1
    layer_3     0.2    0.1
    layer_4     0.3    0.1
    layer_5     0.4    0.2
    layer_6     0.6    0.2
    layer_7     0.8    0.2
    layer_8     1.0    0.2
    layer_9     1.2    0.2
    layer_10    1.4    0.2
    layer_11    1.6    0.2
    layer_12    1.8    0.2
    layer_13    2.0    1.0
    layer_14    3.0    1.0
    layer_15    4.0    1.0
    layer_16    5.0    1.0
    layer_17    6.0    1.0
    layer_18    7.0    1.0
    layer_19    8.0    1.0
    layer_20    9.0    1.0
    core       10.0    0.1

Now create a function to determine the root biomass per soil layer

```python
def root_biom(rz_a, rz_b, kd, rdsurf):
    """
    rz_a = root layer at depth z + 1
    rz_b = root layer at depth z
    kd = decay coefficient
    rdsurf = root density at the surface
    """
    return rdsurf*(-np.exp(-kd*rz_a)/kd + np.exp(-kd*rz_b)/kd)
```

We now want to create the *cost-function* to determine our extinction coefficient and therefore the distribution of roots through the profile. We're using the scipy's optimise package as it is really fast.

```python
def root_biom_vx(kd, rd_surf, x_depth, zmax):
    rbm = [root_biom(x_depth[i + 1], x_depth[i], kd, rd_surf) \
            if i < zmax else 0.0 \
            for i in range(len(x_depth))]
    return rbm

def get_root_dist(rb_total, rd_surf, x_depth, zmax):

    cost_fun2 = lambda kd: abs(rb_total - sum(root_biom_vx(kd, rd_surf, x_depth, zmax)))

    # now find the coefficient
    res = leastsq(cost_fun2, x0=3)
    # optimal parameter
    kd_0 = res[0][0]

    # now we can determine the root density per soil layer
    root_biom = root_biom_vx(kd_0, rd_surf, x_depth, zmax)
    # return to user
    return (root_biom, kd_0)
```

Test that the above minimiser functions works accordingly... and it (collectively) seems to do just fine. So lets test a bunch of profile and visualise what they might look like.

```python
rd_surf = 1.4e4
rb_total = 1980
root_depth = 17

# test
rdens_list = np.linspace(1e4, 15e3, 6)
color_list = plt.get_cmap('Set1', len(rdens_list))

ax1 = plt.subplot(131)
ax2 = plt.subplot(132, sharey=ax1)
ax3 = plt.subplot(133, sharey=ax1)

for (i, rd) in enumerate(rdens_list):

    rbm, _ = get_root_dist(rb_total, rd, soil_depth[:-1], root_depth)
    rdn = np.array(rbm)/np.array(soil_thick)
    rfr = np.array(rbm)/rb_total*100

    ax1.plot(rbm, soil_depth[:-1], label=rd, color=color_list(i))

    ax2.plot(rdn, soil_depth[:-1], label=rd, color=color_list(i))

    ax3.plot(rfr, soil_depth[:-1], label=rd, color=color_list(i))

ax1.set_ylim([0, 2])
ax1.set_ylabel(r'Soil depth (m)')

ax1.set_xlabel(r'Root biomass (g m$^{-2}$)')
ax2.set_xlabel(r'Root density (g m$^{-3}$)')
ax3.set_xlabel(r'Root fraction (%)')

ax2.set_title('Total root biomass = {0} g/m-2'.format(rb_total), y=1.05)

ax1.invert_yaxis()

ax3.legend(loc='lower right', title="Surface root density (g m$^{-3}$)", fontsize=10)
```

```python
rd_surf = 1.4e4
rb_total = 1930
root_depth = 17

# test
rbtot_list = np.linspace(1e3, 3e3, 6)
color_list = plt.get_cmap('Set1', len(rdens_list))

ax1 = plt.subplot(131)
ax2 = plt.subplot(132, sharey=ax1)
ax3 = plt.subplot(133, sharey=ax1)

for (i, rb) in enumerate(rbtot_list):

    rbm, _ = get_root_dist(rb, rd_surf, soil_depth[:-1], root_depth)
    rdn = np.array(rbm)/np.array(soil_thick)
    rfr = np.array(rbm)/rb_total*100

    ax1.plot(rbm, soil_depth[:-1], label=rb, color=color_list(i))
    ax2.plot(rdn, soil_depth[:-1], label=rb, color=color_list(i))
    ax3.plot(rfr, soil_depth[:-1], label=rb, color=color_list(i))

ax1.set_ylim([0, 2])
ax1.set_ylabel(r'Soil depth (m)')

ax1.set_xlabel(r'Root biomass (g m$^{-2}$)')
ax2.set_xlabel(r'Root density (g m$^{-3}$)')
ax3.set_xlabel(r'Root fraction (%)')

ax2.set_title('Surface root density = {0} g/m-3'.format(rd_surf), y=1.05)

ax1.invert_yaxis()

ax3.legend(loc='lower right', title="Total root biomass (g m$^{-2}$)", fontsize=10)
```

### Alternative using Jackson et al. (1996), Oecologia 108: pp389-411

This is another highly cited way of determining the fraction of root biomass per soil depth; based on a global analysis of root distributions for different biomes. Data for various ecosystems or plant functional types (PFTs) were used to fit a simple decay function, with the decay coefficient determined for each PFT, as follows:

\begin{equation}
    Y = 1 - \beta^{d}
\end{equation}

where $Y$ is the root fraction at depth, $d$, and $\beta$ is the decay coefficient that has a range of approximately 0.90 to 0.98.

```python
jackson = lambda B, d: 1 - B**(d*10)
cum_rfrac = [jackson(0.9, d) for d in soil_depth]

plt.figure(figsize=(12, 6))
ax1 = plt.subplot(121)
ax1.plot(cum_rfrac, soil_depth, '-', label='cumulative root frac')
ax1.invert_yaxis()

ax2 = plt.subplot(122, sharey=ax1)
rfrac = [cum_rfrac[i+1] - cum_rfrac[i] for i in range(len(cum_rfrac)-1)]
ax2.plot(test, soil_depth[:-1], '-', label='root frac')

ax1.set_ylabel('Soil depth (m)')
ax1.set_xlabel('Cumulative root fraction')
ax2.set_xlabel('Root fraction')

```

<footer id="attribution" style="float:right; color:#999; background:#fff;">
Created with Jupyter, delivered by Fastly, rendered by Rackspace.
</footer>


# Phase Diagram
|Author  | Stanley A. Baronett|
|--------|--------------------|
|Created | 3/23/2022          |
|Updated | 3/30/2022          |

## [Beattie (1993)](https://doi.org/10.1007/BF00712982)

> **Fig. 4.** Phase diagram for the system $\mathrm{Mg}_2\mathrm{Si}\mathrm{O}_4-\mathrm{Fe}_2\mathrm{Si}\mathrm{O}_4$.
The *solid lines* represent the phase diagram calculated using Eqs. 10 and 12 for the forsterite component, *dotted lines* are the phase diagram as calculated by Ghiorso et al.(1983), assuming that the melt phase is a regular solution.
The circles are experimentally determined values (Bowen and Schairer 1935);
the solidus and liquidus compositions are represented by *solid* and *hollow* circles respectively.
Bowen and Schairer estimated the errors to be $\pm4^\circ\,\mathrm{C}$ for all but the $1890^\circ\,\mathrm{C}$ experiment.
It is clear that neither model reproduces the binary exactly, but the equations developed here do reproduce the shape of the phase diagram.
Temperatures in $^\circ\,\mathrm{C}$

### Equations

D_MgO^ol/melt = (2x/3)/(2y/3) = x/y
D_FeO^ol/melt = (1-x)/(1-y)
D_FeO = A*D_MgO+B
(1-x)/(1-y)=A*(x/y)+B
e.g., for x=0.1, solve for y

Try for x=0.2 [Mg0.2Fe0.8]SiO4
D_FeO = A_Fe*D_MgO + B_Fe

Combining activity models, at equilibrium
\begin{equation}
  T^\mathrm{Ol}=\frac{\Delta H_M^\mathrm{Ol}/R+(P-10^5)\Delta V_M^\mathrm{Ol}/R}
  {\Delta S_M^\mathrm{Ol}/R+2\ln D_{M*}^{\mathrm{Ol}/L}+2\ln1.5C_{NM*}^L+2\ln3C_{\mathrm{SiO}_2*}^L-(NF)},
\tag{10}
\end{equation}
where $\Delta H,\,\Delta S$ and $\Delta V$ are the enthalpy, entropy, and volumes of fusion, respectively; $T$ and $P$ are the temperature in Kelvin and pressure in $\mathrm{Pa}$, respectively; $R$ is the gas constant; and $(NF)=7/2\ln(1-C_{\mathrm{AlO}_{1.5*}}^L)+7\ln(1-C_{\mathrm{TiO}_{2*}}^L)$.
For pure olivine compositions, i.e., $[\mathrm{Fe,Mg}]_2\mathrm{SiO}_4$, network-modifying molar concentration $C_{NM*}^L=2/3$, silicate concentration $C_{\mathrm{SiO}_2*}^L=1/3$, and network-forming trace elements $(NF)=0$, so the last three terms in the denominator of Eq. 10 vanish.
At $1\,\mathrm{atm}\approx10^5\,\mathrm{Pa}$, Eq. 10 reduces to
\begin{equation}
  T^\mathrm{Ol}=\frac{\Delta H_M^\mathrm{Ol}/R}
  {\Delta S_M^\mathrm{Ol}/R+2\ln D_{M*}^{\mathrm{Ol}/L}},
\tag{10*}
\end{equation}
where, from Eq. 13,
\begin{equation}
  D_\mathrm{M*}^{\mathrm{Ol}/L}= A_M^\mathrm{Ol} D_\mathrm{Mg*}^{\mathrm{Ol}/L} + B_M^\mathrm{Ol}.
\tag{13*}
\end{equation}

We can parameterize the molar fraction of $\mathrm{MgO}=x$ in a solid olivine crystal, so $\mathrm{FeO}=1-x$.
Similarly, for a liquidus olivine melt, $\mathrm{MgO}=y$, and $\mathrm{FeO}=1-y$.
This gives the following molar partition coefficients:
\begin{align}
  D_\mathrm{MgO}^{\mathrm{Ol}/L}&=\frac{x}{y}\\
  D_\mathrm{FeO}^{\mathrm{Ol}/L}&=\frac{1-x}{1-y}.
\tag{A}
\end{align}
Plugging Eqs. A into Eq. 13*, and solving for $x$ gives
\begin{equation}
  x=\frac{B_\mathrm{FeO}^\mathrm{Ol}y^2-B_\mathrm{FeO}^\mathrm{Ol}y+y}
  {y-A_\mathrm{FeO}^\mathrm{Ol}y+A_\mathrm{FeO}^\mathrm{Ol}}.
\tag{B}
\end{equation}
Thus, for a given molar fraction of forsterite, $\mathrm{MgO}=y$, in an olivine melt, we can simultaneously calculate the equilibrium temperatures for liquidus and solidus phase boundary curves.

```python
import matplotlib.pyplot as plt
import numpy as np

# Physical constants, conversions
R = 8.31446261815324            # gas constant [J/K/mol] or [m³Pa/K/mol]
KtoC = -273.15                  # absolute zero in (deg) C

# Table 1. Empirical constants for Eq. 13
A_Mg, B_Mg = 1.000, 0.000
A_Fe, B_Fe = 0.299, 0.027

# Table 2. Thermodynamic parameters for olivine components in silicate melts
DH_Mg, DH_Fe = 113.1e3, 109.9e3 # [J/mol]
DS_Mg, DS_Fe = 52.05,   69.60   # [J/K/mol]

# Endmember cases
Fo0, T0 = 0, DH_Fe/DS_Fe        #   0% forsterite (100% fayalite, Fe)
Fo1, T1 = 1, DH_Mg/DS_Mg        # 100% forsterite (Mg)

# Eq. 13
y = C_Mg = np.linspace(0.001, 1, 1000, endpoint=False) # mol fraction of forsterite
x = (B_Fe*y**2-B_Fe*y+y)/(y-A_Fe*y+A_Fe)

# Eq. 12
D_Mg = x/y
D_Fe = (1-x)/(1-y)

# Eq. 10
Ts_Mg = (DH_Mg/R)/(DS_Mg/R+2*np.log(D_Mg))
Ts_Fe = (DH_Fe/R)/(DS_Fe/R+2*np.log(D_Fe))

# Plot
fig, ax = plt.subplots(dpi=150)
ax.plot(Fo0, T0+KtoC, 'ko')
ax.plot(Fo1, T1+KtoC, 'ko')
ax.plot(y, Ts_Fe+KtoC, label='Liquidus')
ax.plot(x, Ts_Mg+KtoC, label='Solidus')
ax.grid()
ax.legend()
ax.minorticks_on()
ax.set(xlabel=r'Mol fraction forsterite (Mg$_2$SiO$_4$)',
       ylabel=r'$T$ / $^\circ\mathrm{C}$',
       xlim=(0, 1), ylim=(1000, 2000))
ax.tick_params(axis='both', which='both', top=True, right=True)
```


###### Content under Creative Commons Attribution license CC-BY 4.0, code under MIT license (c)2014 L.A. Barba, G.F. Forsyth, C.D. Cooper.

# Spreading out

Welcome back! This is the third lesson of the course [Module 4](https://github.com/numerical-mooc/numerical-mooc/tree/master/lessons/04_spreadout), _Spreading out: parabolic PDEs,_ where we study the numerical solution of diffusion problems.

In the first two notebooks, we looked at the 1D heat equation, and solved it numerically using [*explicit*](https://github.com/numerical-mooc/numerical-mooc/blob/master/lessons/04_spreadout/01_Heat_Equation_1D_Explicit.ipynb) and [*implicit*](https://github.com/numerical-mooc/numerical-mooc/blob/master/lessons/04_spreadout/02_Heat_Equation_1D_Implicit.ipynb) schemes. We learned that implicit schemes are unconditionally stable, and we are free to choose any time step. —Wait: _any time step?_ Remember, we still want to capture the physics of the problem accurately. So although stability concerns do not limit the time step, it still hast to be small enough to satisfy any accuracy concerns.

We are now ready to graduate to two dimensions! In the remaining lessons of this course module, we will study the 2D heat equation and reaction-diffusion equation. Like before, we start with explicit methods (this lesson) and then move to implicit methods (next lesson). Let's get started.

## 2D Heat conduction

The equation of heat conduction in 2D is:

\begin{equation}
\rho c_p \frac{\partial T}{\partial t} = \frac{\partial}{\partial x} \left( \kappa_x \frac{\partial T}{\partial x} \right) + \frac{\partial}{\partial y} \left(\kappa_y \frac{\partial T}{\partial y} \right)
\end{equation}

where $\rho$ is the density, $c_p$ is the heat capacity and $\kappa$ is the thermal conductivity.

If the thermal conductivity $\kappa$ is constant, then we can take it outside of the spatial derivative and the equation simplifies to:

\begin{equation}
\frac{\partial T}{\partial t} = \alpha \left(\frac{\partial^2 T}{\partial x^2} + \frac{\partial^2 T}{\partial y^2} \right)
\end{equation}

where $\alpha = \frac{\kappa}{\rho c_p}$ is the thermal diffusivity.  The thermal diffusivity describes the ability of a material to conduct heat vs. storing it.

Does that equation have a familiar look to it?  That's because it's the same as the diffusion equation.  There's a reason that $\alpha$ is called the thermal *diffusivity*! We're going to set up a interesting problem where 2D heat conduction is important, and set about to solve it with explicit finite-difference methods.

### Problem statement

Removing heat out of micro-chips is a big problem in the computer industry. We are at a point in technology where computers can't run much faster because the chips might start failing due to the high temperature. This is a big deal! Let's study the problem more closely.

We want to understand how heat is dissipated from the chip with a very simplified model. Say we consider the chip as a 2D plate of size $1{\rm cm}\times 1{\rm cm}$, made of Silicon: $\kappa = 159{\rm W/m C}$, $c_p = 0.712\cdot 10^3 {\rm J/kg C}$, $\rho = 2329{\rm kg/m}^3$, and diffusivity $\alpha \approx 10^{-4}{\rm m}^2{/rm s}$. Silicon melts at $1414{\rm C}$, but chips should of course operate at much smaller temperatures. The maximum temperature allowed depends on the processor make and model; in many cases, the maximum temperature is somewhere between $60{\rm C}$ and $\sim70{\rm C}$, but better CPUs are recommended to operate at a [maximum of $80{\rm C}$](http://www.pugetsystems.com/blog/2009/02/26/intel-core-i7-temperatures/) (like the Intel Core i7, for example).

We're going to set up a somewhat artificial problem, just to demonstrate an interesting numerical solution. Say the chip is in a position where on two edges (top and right) it is in contact with insulating material.  On the other two edges the chip is touching other components that have a constant temperature of $T=100{\rm C}$ when the machine is operating. Initially, the chip is at room temperature $(20{\rm C})$. *How long does it take for the center of the chip to reach $70{\rm C}$?*

####Figure 1: Simplified microchip problem setup.

Let's use what we have learned to tackle this problem!

## 2D Finite differences

Everything you learned about finite-difference schemes in [Notebook 1 of Module 2](http://nbviewer.ipython.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/02_spacetime/02_01_1DConvection.ipynb) still applies, but now there are two spatial dimensions.  We will need to build a 2D grid of discrete points to compute the solution on.

We will use a 2D Cartesian grid: one that consists of two families of (grid) lines parallel to the two spatial directions. Two lines (of different families) intersect on one and only one grid node (this is called a _structured_ grid). In the $x$ direction, the discretization uses $i=0, \cdots N_x$ lines, and in the $y$ direction we have $j=0, \cdots N_y$ lines. A given node on the grid will now have two spatial coordinates, and we need two indices: for the two lines that intersect at that node.  For example, the middle point in the figure below would be $T_{i,j}$.

#### Figure 2. Nodal coordinates in 2 dimensions

### Explicit scheme in 2D

Recall from above that the 2D heat equation is

$$\frac{\partial T}{\partial t} = \alpha \left(\frac{\partial^2 T}{\partial x^2} + \frac{\partial^2 T}{\partial y^2} \right)$$

Let's write this out discretized using forward difference in time, and central difference in space,  using an explicit scheme. You should be able write this out yourself, without looking—if you need to look, it means you still need to write more difference equations by your own hand!

\begin{equation}
\frac{T^{n+1}_{i,j} - T^n_{i,j}}{\Delta t} = \alpha \left( \frac{T^n_{i+1, j} - 2T^n_{i,j} + T^n_{i-1,j}}{\Delta x^2} + \frac{T^n_{i, j+1} - 2T^n_{i,j} + T^n_{i,j-1}}{\Delta y^2}\right)
\end{equation}

Rearranging the equation to solve for the value at the next time step, $T^{n+1}_{i,j}$, yields

\begin{equation}
T^{n+1}_{i,j}= T^n_{i,j} + \alpha \left( \frac{\Delta t}{\Delta x^2} (T^n_{i+1, j} - 2T^n_{i,j} + T^n_{i-1,j}) + \frac{\Delta t}{\Delta y^2} (T^n_{i, j+1} - 2T^n_{i,j} + T^n_{i,j-1})\right)
\end{equation}

That's a little messier than 1D, but still recognizable.

Up until now, we've used stencils to help visualize how a scheme will advance the solution for one timestep.  Stencils in 2D are a little harder to draw, but hopefully the figure below will guide your understanding of this method: we are using five grid points at time step $n$ to obtain the solution on one point at time step $n+1$.

####Figure 3: 2D Explicit Stencil

Similar to all of the 1D explicit methods we've used, the solution at $T^{n+1}_{i,j}$ is updated using only known values from the current solution at time $n$. This is straightforward to implement in code, but will be subject to stability limitations on the time step that you can choose. We'll study an implicit method in the next lesson.

### Boundary Conditions

Whenever we reach a point that interacts with the boundary, we apply the boundary condition. As in the previous notebook, if the boundary has Dirichlet conditions, we simply impose the prescribed temperature at that point. If the boundary has Neumann conditions, we approximate them with a finite-difference scheme.

Remember, Neumann boundary conditions prescribe the derivative in the normal direction. For example, in the problem described above, we have $\frac{\partial T}{\partial y} = q_y$ in the top boundary and $\frac{\partial T}{\partial x} = q_x$ in the right boundary, with $q_y = q_x = 0$ (insulation).

Thus, at every time step, we need to enforce

\begin{equation}
T_{i,end} = q_y\cdot\Delta y + T_{i,end-1}
\end{equation}

and

\begin{equation}
T_{end,j} = q_x\cdot\Delta x + T_{end-1,j}
\end{equation}

Write the finite-difference discretization of the boundary conditions yourself, and confirm that you can get the expressions above.

### Stability

Before doing any coding, let's revisit stability constraints. We saw in the first notebook of this series that the 1D explicit discretization of the diffusion equation was stable as long as $\alpha \frac{\Delta t}{(\Delta x)^2} \leq \frac{1}{2}$. In 2D, this constraint is even tighter, as we need to add them in both directions:

\begin{equation}
\alpha \frac{\Delta t}{(\Delta x)^2} + \alpha \frac{\Delta t}{(\Delta y)^2} < \frac{1}{2}.
\end{equation}

Say that the mesh has the same spacing in $x$ and $y$, $\Delta x = \Delta y = \delta$. In that case, the stability condition is:

\begin{equation}
\alpha \frac{\Delta t}{\delta^2} < \frac{1}{4}
\end{equation}

## Code implementation

### Array storage

The physical problem has two dimensions, so we also store the temperatures in two dimensions: in a 2D array.

We chose to store it with the $y$ coordinates corresponding to the rows of the array and $x$ coordinates varying with the columns (this is just a code design decision!). If we are consistent with the stencil formula (with $x$ corresponding to index $i$ and $y$ to index $j$), then $T_{i,j}$ will be stored in array format as `T[j,i]`.

This might be a little confusing as most of us are used to writing coordinates in the format $(x,y)$, but our preference is to have the data stored so that it matches the physical orientation of the problem. Then, when we make a plot of the solution, the visualization will make sense to us, with respect to the geometry of our set-up. That's just nicer than to have the plot rotated!

####Figure 4: Row-column data storage

As you can see on Figure 4 above, if we want to access the value $18$ we would write those coordinates as $(x_2, y_3)$.  You can also see that its location is the 3rd row, 2nd column, so its array address would be `T[3,2]`.

Again, this is a design decision.  However you can choose to manipulate and store your data however you like; just remember to be consistent!

### Code time!

Now, to some coding! First, we have a little function that will advance the solution in time with a forward-time, centered-space scheme, and will monitor the center of the plate to tell us when it reaches $70{\rm C}$. Let's start by setting up our Python compute environment.

```
import numpy
from matplotlib import pyplot
%matplotlib inline
from matplotlib import rcParams
rcParams['font.family'] = 'serif'
rcParams['font.size'] = 16
```

```
def ftcs(T, nt, alpha, dt, dx, dy):

    j_mid = (numpy.shape(T)[0])/2
    i_mid = (numpy.shape(T)[1])/2

    for n in range(nt):
        Tn = T.copy()
        T[1:-1,1:-1] = Tn[1:-1,1:-1] + alpha *\
            (dt/dy**2 * (Tn[2:,1:-1] - 2*Tn[1:-1,1:-1] + Tn[:-2,1:-1]) +\
             dt/dx**2 * (Tn[1:-1,2:] - 2*Tn[1:-1,1:-1] + Tn[1:-1,:-2]))

        # Enforce Neumann BCs
        T[-1,:] = T[-2,:]
        T[:,-1] = T[:,-2]

        # Check if we reached T=70C
        if T[j_mid, i_mid] >= 70:
            print ("Center of plate reached 70C at time {0:.2f}s.".format(dt*n))
            break

    if T[j_mid, i_mid]<70:
        print ("Center has not reached 70C yet, it is only {0:.2f}C.".format(T[j_mid, i_mid]))

    return T
```

See the [`break`](https://docs.python.org/2/tutorial/controlflow.html) statement? It exits the `for` loop at the closest time iteration when the plate reaches $70{\rm C}$.

In the code cell below, we define our initial conditions according to the problem set up, and choose the discretization parameters. We start with only 20 spatial steps in each coordinate direction and advance for 500 time steps. You should later experiments with these parameters at your leisure!

```
L = 1.0e-2
H = 1.0e-2

nx = 21
ny = 21
nt = 500

dx = L/(nx-1)
dy = H/(ny-1)

x = numpy.linspace(0,L,nx)
y = numpy.linspace(0,H,ny)

alpha = 1e-4

Ti = numpy.ones((ny, nx))*20
Ti[0,:]= 100
Ti[:,0] = 100
```

We don't want our solution blowing up, so let's find a time step with $\frac{\alpha \Delta t}{\Delta x^2} = \frac{\alpha \Delta t}{\Delta y^2} = \frac{1}{4}$.

```
sigma = 0.25
dt = sigma * min(dx, dy)**2 / alpha
T=Ti.copy()
T = ftcs(T, nt, alpha, dt, dx, dy)
```

    Center of plate reached 70C at time 0.16s.

### Visualize the results

By now, you're no doubt *very* familiar with the `pyplot.plot` command.  It's great for line plots, scatter plots, etc., but what about when we have two spatial dimensions and another value (temperature) to display?

Are you thinking contour plot?  We're thinking contour plot.  Check out the documentation on [`pyplot.contourf`](http://matplotlib.org/api/pyplot_api.html#matplotlib.pyplot.contour) (the 'f' denotes "filled" contours).

First, though, we need to create a `meshgrid` purely for plotting purposes:   [`numpy.meshgrid`](http://docs.scipy.org/doc/numpy/reference/generated/numpy.meshgrid.html) is a function that takes two vectors (`x` and `y`, say) and returns two 2D arrays of $x$ and $y$ coordinates that we then use to create the contour plot.

```
#my, mx = numpy.meshgrid(y,x)
mx, my = numpy.meshgrid(x,y)
```

```
pyplot.figure(figsize=(8,5))
pyplot.contourf(my,mx,T,20)
pyplot.xlabel('$x$')
pyplot.ylabel('$y$')
pyplot.colorbar();
```

That looks pretty cool!  Note that in the call to `pyplot.contourf` you can specify the number of contour levels to display (we chose 20).  Look at that visualization: does it make physical sense to you, considering that the upper and right sides of the chip are insulated, in our problem?

##### Dig deeper

In the problem we just demonstrated, the chip reaches a temperature of $70{\rm C}$ at a given time, but will it keep increasing? That spells trouble.

Imagine that you have a heat sink instead of an insulator acting on the upper and right sides. What should be the heat flux that the heat sink achieves there, so that the temperature does not exceed $70{\rm C}$ at the center of the chip?

---
###### The cell below loads the style of the notebook

```
from IPython.core.display import HTML
css_file = '../../styles/numericalmoocstyle.css'
HTML(open(css_file, "r").read())
```

<link href='http://fonts.googleapis.com/css?family=Alegreya+Sans:100,300,400,500,700,800,900,100italic,300italic,400italic,500italic,700italic,800italic,900italic' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Arvo:400,700,400italic' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=PT+Mono' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Shadows+Into+Light' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Nixie+One' rel='stylesheet' type='text/css'>
<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: 750px;
}

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.8em;
}

/* 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: 'Alegreya Sans' sans-serif;
    line-height: 140%;
    font-size: 125%;
    font-weight: 400;
    width:600px;
    margin-left:auto;
    margin-right:auto;
}

/* Formatting for header cells */
.text_cell_render h1 {
    font-family: 'Nixie One', serif;
    font-style:regular;
    font-weight: 400;
    font-size: 45pt;
    line-height: 100%;
    color: rgb(0,51,102);
    margin-bottom: 0.5em;
    margin-top: 0.5em;
    display: block;
}
.text_cell_render h2 {
    font-family: 'Nixie One', serif;
    font-weight: 400;
    font-size: 30pt;
    line-height: 100%;
    color: rgb(0,51,102);
    margin-bottom: 0.1em;
    margin-top: 0.3em;
    display: block;
}

.text_cell_render h3 {
    font-family: 'Nixie One', serif;
    margin-top:16px;
	font-size: 22pt;
    font-weight: 600;
    margin-bottom: 3px;
    font-style: regular;
    color: rgb(102,102,0);
}

.text_cell_render h4 {    /*Use this for captions*/
    font-family: 'Nixie One', serif;
    font-size: 14pt;
    text-align: center;
    margin-top: 0em;
    margin-bottom: 2em;
    font-style: regular;
}

.text_cell_render h5 {  /*Use this for small titles*/
    font-family: 'Nixie One', sans-serif;
    font-weight: 400;
    font-size: 16pt;
    color: rgb(163,0,0);
    font-style: italic;
    margin-bottom: .1em;
    margin-top: 0.8em;
    display: block;
}

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

.CodeMirror{
        font-family: "PT Mono";
        font-size: 90%;
}

</style>


```python
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit

data = np.fromfile('fit_result_8p9uM_15022021_n_1p5883r_1p504fix_0p052_wav_532.dat', dtype=np.float64)
data = data.reshape(len(data)//10,10)

x = data[:,0]*0.0535
y = data[:,1]*0.0535
z = data[:,2]*0.0535
```

```python

```

```python
def traj_sedim(x, a, b):
    return a * x + b

```

```python

z = z[0:3000]
time = np.arange(0,len(z)) / 100

```

```python
plt.plot(time, z)
plt.xlabel("$t$ (s)")
plt.ylabel("$z$ ($\\mu$m)")
```

```python
popt, pcov = curve_fit(traj_sedim, time, z)
```

```python
plt.plot(time, z)
plt.plot(time, traj_sedim(time, *popt))
plt.xlabel("$t$ (s)")
plt.ylabel("$z$ ($\\mu$m)")
```

```python
popt
```

    array([-0.21171449, 29.22368497])

The sedimentation drift velocity for a spherical particle in bulk writes:

\begin{equation}
v_\mathrm{sed} = \frac{2}{9} \frac{\Delta \rho a^2 g}{\eta}
\end{equation}

```python
a = 1.45e-6
v_th = 2/9 * 50 * a ** 2 * 9.81 / 0.001 * 1e6
```

```python
v_th
```

    0.22917250000000006

```python
def msd(x, Dt):
	"""Function that returns the MSD for a list of time indices Dt for a trajectory x"""
	_msd = lambda x, Dt: np.mean((x[:-Dt] - x[Dt:]) ** 2)
	return [_msd(x, i) for i in Dt]
```

```python
Dt = np.arange(1,450,5)
```

```python
MSD = msd(z, Dt)

```

```python
def msd_th(x, d, v):
    return 2*d*x + v**2 * x**2

def msd_log(x, d, v):
    return np.log( 2*d*x + v**2 * x**2)
```

```python
plt.plot(np.log(Dt/100),np.log(MSD), "o")
```

```python
popt, pcov = curve_fit(msd_th, Dt/100, MSD, p0 = [0.13, 0.22], maxfev=20000, ftol=1e-2)
```

```python
popt
```

    array([0.13839701, 0.24298569])

```python
perr = np.sqrt(np.diag(pcov))
```

```python
perr
```

    array([0.00079552, 0.00094242])

```python
Do
```

    0.13690747792851213

```python
plt.plot(Dt/100,MSD, "o")
plt.plot(Dt/100, msd_th(Dt/100, *popt), color = "k")
plt.xlabel("$\Delta t$ (s)")
plt.ylabel("MDS ($\mu$m$^2$)")
```

```python
relat_err_D = (popt[0] - 0.137)/0.137
relat_err_D*100
```

    1.019712992433843

```python
relat_err_v = (popt[1] - 0.229)/0.229
relat_err_v*100
```

    6.107287317556856

```python

```


```python
import holoviews as hv
hv.extension('bokeh')
```

```python
import numpy as np
import scipy.signal
```

# Bayesian modeling

## Introduction

In this lecture we return to parametric modeling but using the bayesian approach.

A summary of the bayesian premise

- Inference is made by producing probability density functions (pdf): **posterior**
- We model the uncertainty of the data, experiment, parameters, etc. as a **joint pdf**
- The parameter vector $\theta$ is a R.V., *i.e.* it follows a distribution: **prior**

The Bayes theorem and the law of total probability tell us

$$
p(\theta| \{x\}) = \frac{p(\{x\}, \theta)}{p(\{x\})}= \frac{p(\{x\}|\theta) p(\theta)}{\int p(\{x\}|\theta) p(\theta) d\theta} \propto p(\{x\}|\theta) p(\theta),
$$

:::{note}

The posterior is build from the **likelihood**, **prior** and **evidence** (marginal data likelihood), *i.e.* the posterior can be small if either the likelihood or the prior are small

:::

**Why/When should I use the Bayesian formalism?**

In many cases bayesian inference will not differ much from frequentist techniques. Also, in general, bayesian inference is harder to compute and requires more sophisticated methods

But bayesian modeling gives us some key advantages:

- We know the uncertainty of our parameters/predictions, i.e. and we can take more informed decisions
- It gives a principled way of injecting prior knowledge (regularization)
- We can integrate unknown or missing (nuisance) parameters

The following is a summary of the Bayesian inference procedure

1. Formulate your problem: likelihood and prior
1. Build a joint distribution (relation of all parameters)
1. Determine the posterior using Bayes Theorem. Find MAP and credible regions
1. Test your hypothesis
1. **Criticize:** Evaluate how appropriate the model is and suggest improvements

We will review these steps in this lesson

## Maximum *a posteriori* (MAP) estimation

In the Bayesian setting the best "point estimate" of the parameters of the model is given by the MAP

$$
\hat \theta = \text{arg} \max_\theta p(\theta|\{x\}) =  \text{arg} \max_\theta p(\{x\}| \theta) p(\theta),
$$

where we "omit" the evidence (denominator in Bayes rule) because it does not depend on $\theta$

Applying the logarithm (monotonic) we can decouple the likelihood from the prior

$$
\hat \theta = \text{arg} \max_\theta \log p(\{x\}| \theta) + \log p(\theta),
$$

:::{note}

MAP is still a point estimate: poor's man Bayes

:::

The main difference to what we saw in previous lessons is **the prior**

### What can I do with priors?

Priors are distributions that summarize what we know about the parameters before-hand, for example

- a parameter is continuous and has no bounds: Normal
- a parameter is continuous and positive: Lognormal, Inverse gamma, Half-normal, etc
- a parameter is positive-semidefinite: Inverse Wishart, LKJ, etc
- a parameter is in the simplex: Dirichlet

Priors can be described as

- Informative: $\mathcal{N}(\theta|\mu=5.4, \sigma^2=0.1)$
- Weakly informative: $\mathcal{N}(\theta|\mu=0, \sigma^2=100.)$
- Uninformative (or objective): My parameter is positive

Of course these notions depend on the problem at hand.

We should select priors that

- add a positive weight on values that may occur
- put zero weight to impossible values
- help regularize the solution

Later we will see the case of **conjugate prior**, which are very convenient from a computational point of view

I suggest reading the [practical principles for choosing priors](https://github.com/stan-dev/stan/wiki/Prior-Choice-Recommendations) in the Stan repository

### Example: MAP estimate of the mean of a Gaussian distribution

Assuming $N$ i.i.d samples and a Gaussian likelihood with known variance we can write

$$
\log p(\{x\}|\theta) = \log L (\mu)  = - \frac{N}{2} \log 2\pi\sigma^2 - \frac{1}{2\sigma^{2}}   \sum_{i=1}^N (x_i-\mu)^2,
$$

In this particular example we will select a Gaussian prior with parameters $\mu_0$ and $\sigma_0$ for $\mu$

$$
\log p(\theta) = -\frac{1}{2} \log 2 \pi \sigma^2_0 - \frac{1}{2 \sigma^2_0}  (\mu - \mu_0)^2,
$$

Adding the log likelihood and log prior and taking the derivative

$$
\frac{d}{d\mu} \log p(\{x\}|\theta) + \log p(\theta) =   \frac{1}{\sigma^{2}}   \sum_{i=1}^N (x_i-\mu)  - \frac{1}{ \sigma^2_0}  (\mu - \mu_0),
$$

then setting the derivative equal to zero gives us the MAP estimate

$$
\hat \mu_{\text{map}} =  \left(\frac{N}{\sigma^2} + \frac{1}{\sigma^2_0} \right)^{-1} \left(\frac{N}{\sigma^2} \bar x + \frac{1}{\sigma^2_0} \mu_0 \right),
$$

where $\bar x = \frac{1}{N} \sum_{i=1}^N x_i$.

:::{important}

Do not confuse $\sigma^2$ (the likelihood/noise variance) and $\sigma^2_0$ (prior variance)

:::

(Using a bit of algebra) we can write the MAP expression as

$$
\begin{align}
\hat \mu_{\text{map}}  &=  \left(\frac{N}{\sigma^2} + \frac{1}{\sigma^2_0} \right)^{-1} \left(\frac{N\bar x}{\sigma^2}  + \frac{\mu_0}{\sigma^2_0} \right)  \nonumber \\
&=  \frac{N \bar x \sigma^2_0 + \mu_0 \sigma^2}{N\sigma^2_0+ \sigma^2} \nonumber \\
&= \frac{\bar x + \mu_0 \frac{\sigma^2}{\sigma^2_0 N}}{1 + \frac{\sigma^2}{\sigma^2_0 N}}  \nonumber \\
&= w \bar x + (1-w) \mu_0, \qquad \text{where} \quad w = \frac{1}{1 + \frac{\sigma^2}{\sigma^2_0 N}}  \nonumber
\end{align}
$$

> The MAP estimate of $\mu$ is a weighted average between $\mu_0$ (prior) and $\bar x$ (the MLE solution)

:::{note}

In the last expression:

- if either $\sigma^2_0 \to \infty$ or $N \to \infty$ then $w\to1$, i.e. the MAP converges to the MLE solution
- the prior is more relevant if have a few sample (small $N$) or a noisy samples (large $\sigma^2$)

:::

### Extra: MAP intepretation as a penalized MLE/regularized LS

We can rewrite the MAP optimization problem for a Gaussian likelihood with known variance and a zero-mean Gaussian prior as

$$
\begin{align}
\hat \mu_{\text{map}}  &= \text{arg} \max_\mu  \log p(\{x\}| \mu, \sigma^2) + \log p(\mu) \nonumber \\
&= \text{arg} \max_\mu   - \frac{N}{2} \log 2\pi\sigma^2 - \frac{1}{2\sigma^{2}}   \sum_{i=1}^N (x_i-\mu)^2 -  \frac{1}{2\sigma_0^2} \mu^2 \nonumber \\
&= \text{arg} \min_\mu \frac{1}{2\sigma^{2}}   \sum_{i=1}^N (x_i-\mu)^2 +  \frac{1}{2\sigma_0^2} \mu^2 \nonumber \\
&= \text{arg} \min_\mu \|x-\mu\|^2  + \lambda \|\mu \|^2, \nonumber
\end{align}
$$

where $\lambda = \frac{\sigma^2}{\sigma_0^2}$.

We can recognize the last equation as a regularized least squares problem. In this case a using a Gaussian priors is equivalent to using a L2 norm regularizar on the parameters (this is known as ridge regression). A Laplacian prior yields a L1 regularizer (LASSO) [^lasso]

[^lasso]: Hastie, Tibshirani, Friedman, chapter 3.4 (Shrinkage methods), page 61.

We will review ridge regression in a future lecture

## Analytical posterior with conjugate priors

Remember that the MAP is only a point estimate. In a fully-bayesian setting want we are interested in is the posterior of the parameter

In the particular case of a Gaussian likelihod and a Gaussian prior we can rearrange the terms to show that

$$
\begin{align}
p(\theta |\{x\}) &\propto p(\{x\} |\theta ) p(\theta ) \nonumber \\
&\propto \exp \left ( \frac{1}{2\sigma^2} \sum_{i=1}^N (x_i - \mu)^2 \right) \exp \left ( \frac{1}{2\sigma_0^2} (\mu - \mu_0)^2 \right) \nonumber \\
&\propto \exp \left ( -\frac{1}{2 \hat \sigma^2} (\mu - \hat \mu_{\text{map}} )^2 \right),  \nonumber
\end{align}
$$

where

$$
\hat \sigma^2 = \left(\frac{N}{\sigma^2} + \frac{1}{\sigma^2_0} \right)^{-1},
$$

i.e. the posterior has a closed analytical form and is also Gaussian [^gaussmult]

[^gaussmult]: Another way to show that the posterior is Gaussian is to use the [property of Gaussian pdf multiplication](http://www.tina-vision.net/docs/memos/2003-003.pdf)

> When the resulting posterior has the same distribution as the specified prior we say that the prior is a **conjugate prior for the specified likelihood**

In this particular case the Gaussian distribution is conjugate with itself

Other examples are:

```{list-table} Conjugacy table
:header-rows: 1

* - Likelihood
  - Conjugate prior
* - Bernoulli
  - Beta
* - Poisson
  - Gamma
* - Multinomial or categorial
  - Dirichlet
* - Exponential
  - Gamma
* - Normal with unknown variance
  - Normal-inverse gamma (NIG)
* - Multivariate normal with unknown covariance
  - Normal-inverse Wishart
```

### Interactive example

We generate Gaussian distributed data with $\mu=2$ and $\sigma=1$ and plot the asymptotic distribution of the MLE (yellow) and the analytical posterior (red) and the prior (blue)

```python
from scipy.stats import norm

def mle_mu(xi: np.array) -> float:
    return np.mean(xi)

def asymptotic_mle(x: np.array, xi: np.array, s2: float) -> np.array:
    N = len(xi)
    return norm(loc=mle_mu(xi), scale=np.sqrt(s2/N)).pdf(x)

def map_mu(xi: np.array, mu0: float, s20: float, s2: float):
    N = len(xi)
    w = (N*s20)/(N*s20 + s2)
    return mle_mu(xi)*w + mu0*(1. - w)

def prior_mu(x: np.array, mu0: float, s20: float) -> np.array:
    return norm(loc=mu0, scale=np.sqrt(s20)).pdf(x)

def posterior_mu(x: np.array, xi: np.array, mu0: float, s20: float, s2: float) -> np.array:
    N = len(xi)
    s2_pos = s2*s20/(N*s20 + s2)
    mu_pos = map_mu(xi, mu0, s20, s2)
    return norm(loc=mu_pos, scale=np.sqrt(s2_pos)).pdf(x)
```

Explore

- What happens with $N$ grows?
- What happens when $\sigma_0$ grows?

```python
mu_real, s2_real = 2., 1.
x_plot = np.linspace(-5, 5, num=1000)
true_value = hv.VLine(mu_real).opts(color='k', line_width=2, alpha=0.5)
hmap = hv.HoloMap(kdims=['N', 'mu0', 's20'])
for N in [1, 5, 10, 50, 100, 500]:
    for mu0 in np.linspace(-3, 3, num=5):
        for s20 in np.logspace(-1, 1, num=3):
            data = norm(loc=mu_real, scale=np.sqrt(s2_real)).rvs(N, random_state=1234)
            plot_prior = hv.Curve((x_plot, prior_mu(x_plot, mu0, s20)), 'x', 'density', label='prior')
            plot_mle = hv.Curve((x_plot, asymptotic_mle(x_plot, data, s2_real)), label='MLE')
            plot_post = hv.Curve((x_plot, posterior_mu(x_plot, data, mu0, s20, s2_real)), label='posterior')
            hmap[(N, mu0, s20)] = (plot_prior * plot_post * plot_mle * true_value).opts(hv.opts.Curve(width=500))

hmap
```

### Conjugate prior for Gaussian likelihood when $\sigma^2$ is unknown

Before we assumed that $\sigma^2$ was a known quantity and we focused on estimating $\mu$

If we now assume that the mean $\mu$ is known and the variance is unknown then the conjugate prior for the variance is an inverse-Gamma distribution

$$
p(\sigma^2) = \text{IG}(\sigma^2| \alpha_0, \beta_0) = \frac{\beta_0^{\alpha_0}}{\Gamma(\alpha_0)} x^{-\alpha_0-1} e^{-\frac{\beta_0}{x}}
$$

With which the resulting posterior is also

$$
\text{IG}\left(\sigma^2| \alpha_N , \beta_N  \right),
$$

where

- $ \alpha_N = \alpha_0 + N/2$
- $\beta_N = \beta_0 + \frac{1}{2} \sum_{i=1}^N (x_i - \mu)^2$

As both $\alpha$ and $\beta$ encode the strength of the prior the following parameterization is broadly used

$$
p(\sigma^2) = \text{IG}(\sigma^2| \alpha, \beta) = \text{IG}\left(\sigma^2 \bigg| \frac{\nu}{2}, \frac{\nu \sigma_0^2}{2}\right)
$$

where $\sigma_0^2$ controls the value of the prior and $\nu$ the strength. Note that this is also closely related to the [inverse chi-square distribution](https://en.wikipedia.org/wiki/Inverse-chi-squared_distribution)

### Conjugate prior for Gaussian likelihood when both $\mu$ and $\sigma^2$ are unknown

Multiplying the normal prior and the IG prior does not yield a conjugate prior (assumes independence of $\mu$ and $\sigma$). In this case the conjugate prior is hierarchical

$$
\begin{align}
p(x_i|\mu, \sigma^2) &= \mathcal{N}(\mu, \sigma^2)  \nonumber \\
p(\mu|\sigma^2) &= \mathcal{N}(\mu_0, \sigma^2/\lambda_0)  \nonumber \\
p(\sigma^2) &= \text{IG}(\alpha, \beta)  \nonumber
\end{align}
$$

which is called **normal-inverse-gamma (NIG)**, a four parameter distribution

The NIG prior is

$$
p(\mu, \sigma^2) = \text{NIG}(\mu_0, \lambda_0, \alpha_0, \beta_0) = \mathcal{N}(\mu|\mu_0 , \sigma^2/\lambda_0) \text{IG}(\sigma^2|\alpha_0, \beta_0)
$$

An the posterior is also NIG

$$
p(\mu, \sigma^2|\{x\}) =  \text{NIG}(\mu_n, \lambda_n, \alpha_n, \beta_n)
$$

where

- $\lambda_n = \lambda_0 + N$
- $\mu_n = \lambda_n^{-1} \left ( \lambda_0 \mu_0  + N \bar x \right)$
- $\alpha_n = \alpha_0 + N/2$
- $\beta_n = \beta_0 + 0.5\mu_0^2\lambda_0 + 0.5\sum_i x_i^2 - 0.5\lambda_n \mu_n^2$

## Describing the posterior using Credible Interval (CI) and the High Posterior Density (HPD) regions

One way to summarize the posterior is to measure its **width**

The $100(1-\alpha)$ % CI of $\theta$ is a contiguous region $[\theta_{l}, \theta_{u}]$ such that

$$
P(\theta_{l}< \theta < \theta_{u}) = 1 - \alpha
$$

We have to either know the functional form of the posterior (analytical) or have a posterior from which we can sample from (this is the case if we are using MCMC)

The HPD is an alternative to CI that is better when we have multiple modes. The HPD depends not only on the width but also on the height of the posterior. The following figure shows the difference between them

### Example

The 95% CI for the previous example for a given combination of $\mu_0$, $\sigma_0^2$ and $N$ is

```python
mu0, s20, N = 0., 10., 100
data = norm(loc=mu_real, scale=np.sqrt(s2_real)).rvs(N, random_state=12345)

N = len(data)
s2_pos = s2_real*s20/(N*s20 + s2_real)
mu_pos = map_mu(data, mu0, s20, s2_real)
dist = norm(loc=mu_pos, scale=np.sqrt(s2_pos))

display(f'95 % CI for mu: [{dist.ppf(0.025):0.4f}, {dist.ppf(0.975):0.4f}]')
```

### Extra: Mean of the posterior

Other point estimate that can be used to characterize the posterior is

$$
\hat \theta = \mathbb{E}[\theta|\{x\}] = \int \theta p(\theta| \{x\}) d\theta,
$$

*i.e.* the mean or expected value of the posterior

## Help: My posterior does not have an analytical form

In this case we resort to either variational inference (VI) or Markov Chain Monte Carlo (MCMC) methods

We will learn how to use MCMC to sample from intractable posterior distributions in a future lesson

```python

```


# 03 - Full-Waveform Inversion (FWI)

This notebook is the third in a series of tutorial highlighting various aspects of seismic inversion based on Devito operators. In this second example we aim to highlight the core ideas behind seismic inversion, where we create an image of the subsurface from field recorded data. This tutorial follows on the modelling tutorial and will reuse the modelling and velocity model.

## 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 = 5  # 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
#NBVAL_IGNORE_OUTPUT
from examples.seismic import demo_model, plot_velocity, plot_perturbation

# 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

model = demo_model('circle-isotropic', vp=3.0, vp_background=2.5,
                    origin=origin, shape=shape, spacing=spacing, nbl=40)

model0 = demo_model('circle-isotropic', vp=2.5, vp_background=2.5,
                     origin=origin, shape=shape, spacing=spacing, nbl=40,
                    grid = model.grid)

plot_velocity(model)
plot_velocity(model0)
plot_perturbation(model0, model)
```

## 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.010
# 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, 0] = 20.  # Depth is 20m

# Define acquisition geometry: receivers

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

# 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[:, 0] = 30.
source_locations[:, 1] = np.linspace(0., 1000, num=nshots)

plot_velocity(model, source=source_locations)
```

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

import scipy
def fwi_gradient(vp_in):
    # Create symbols to hold the gradient and residual
    grad = Function(name="grad", grid=model.grid)
    residual = Receiver(name='rec', 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
        true_d, _, _ = solver.forward(vp=model.vp)

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

        # Compute gradient from data residual and update objective function
        residual.data[:] = smooth_d.data[:] - true_d.data[:]

        objective += .5*np.linalg.norm(residual.data.flatten())**2
        solver.gradient(rec=residual, u=u0, vp=vp_in, grad=grad)

    return objective, -grad.data
```

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
#NBVAL_IGNORE_OUTPUT

# Compute gradient of initial model
ff, update = fwi_gradient(model0.vp)
print('Objective value is %f ' % ff)
```

    Objective value is 57283.085366

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

# Plot the FWI gradient
plot_image(update, vmin=-1e4, vmax=1e4, 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 = .5 / np.abs(update).max()
plot_image(model0.vp.data - alpha*update, 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
# Define bounding box constraints on the solution.
def apply_box_constraint(vp):
    # Maximum possible 'realistic' velocity is 3.5 km/sec
    # Minimum possible 'realistic' velocity is 2 km/sec
    return np.clip(vp, 2.0, 3.5)
```

```python
#NBVAL_SKIP

# 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 = .05 / np.abs(direction).max()

    # Update the model estimate and enforce minimum/maximum values
    model0.vp = apply_box_constraint(model0.vp.data - alpha * direction)

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

    Objective value is 57283.085366 at iteration 1
    Objective value is 35764.719809 at iteration 2
    Objective value is 21001.471868 at iteration 3
    Objective value is 11547.503946 at iteration 4
    Objective value is 5785.739636 at iteration 5

```python
#NBVAL_IGNORE_OUTPUT

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

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


# "Next Generation Reservoir Computing"

Daniel J. Gauthier, Erik Bollt, Aaron Grifﬁth & Wendson A. S. Barbosa

*Nature Communications*, vol. 12, no. 1, p. 5564, Sep. 2021, doi: 10.1038/s41467-021-25801-2.

```python
import matplotlib.pyplot as plt
import numpy as np

from reservoirpy.datasets import lorenz, doublescroll
from reservoirpy.observables import nrmse
from reservoirpy.nodes import Ridge, NVAR

%matplotlib inline

from IPython.core.display import HTML
HTML("""
<style>
.output_png {
    display: table-cell;
    text-align: center;
    vertical-align: middle;
}
</style>
""")
```

<style>
.output_png {
    display: table-cell;
    text-align: center;
    vertical-align: middle;
}
</style>

## Abstract

Reservoir computing is a best-in-class machine learning algorithm for processing information
generated by dynamical systems using observed time-series data. Importantly, it requires
very small training data sets, uses linear optimization, and thus requires minimal computing
resources. However, the algorithm uses randomly sampled matrices to deﬁne the underlying
recurrent neural network and has a multitude of metaparameters that must be optimized.
Recent results demonstrate the equivalence of reservoir computing to nonlinear vector
autoregression, which requires no random matrices, fewer metaparameters, and provides
interpretable results. Here, we demonstrate that nonlinear vector autoregression excels at
reservoir computing benchmark tasks and requires even shorter training data sets and
training time, heralding the next generation of reservoir computing.

## Implementation using *ReservoirPy*

This notebook is provided as a demo of the `NVAR` node in ReservoirPy, implementing the method described in the paper *Next Generation Reservoir Computing* by Gauthier *et al.*

The nonlinear vector autoregressive (NVAR) machine implements the following equations. The state $\mathbb{O}_{total}$ of the NVAR first contains a serie of linear features $\mathbb{O}_{lin}$ made of input data concatenated with delayed inputs:

$$
\mathbb{O}_{lin}[t] = \mathbf{X}[t] \oplus \mathbf{X}[t - s] \oplus \mathbf{X}[t - 2s] \oplus \dots \oplus \mathbf{X}[t - (k-1)s]
$$

where $\mathbf{X}[t]$ are the inputs at time $t$, $k$ is the delay and $s$ is the strides (only one input every $s$ inputs within the delayed inputs is used). The operator $\oplus$ denotes the concatenation.

In addition to these linear features, nonlinear representations $\mathbb{O}_{nonlin}^n$ of the inputs are contructed using all unique monomials of order $n$ of these inputs:

$$
\mathbb{O}_{nonlin}^n[t] = \mathbb{O}_{lin}[t] \otimes \mathbb{O}_{lin}[t] \overbrace{\otimes \dots \otimes}^{n-1~\mathrm{times}} \mathbb{O}_{lin}[t]
$$

where $\otimes$ is the operator denoting an outer product followed by the selection of all unique monomials generated by this outer product. Under the hood, this product is computed by ReservoirPy by finding all unique combinations of input features and multiplying each combination of terms.

Finally, all representations are gathered to form the final feature vector $\mathbb{O}_{total}$:

$$
\mathbb{O}_{total} = \mathbb{O}_{lin}[t] \oplus \mathbb{O}_{nonlin}^n[t]
$$

Tikhonov regression is used to compute the readout weights using this feature vector and some target values, in an offline way (we will simply use ReservoirPy's `Ridge` node of for this).

**Fig. 1** A traditional (Reservoir Computing machine) is implicit in an NG-RC (Next Generation Reservoir Computing machine). *(top)* A traditional RC processes time-series data associated with a strange attractor (blue, middle left)
using an artiﬁcial recurrent neural network. The forecasted strange attractor (red, middle right) is a linear weight of the reservoir states. *(bottom)* The NG-RC performs a forecast using a linear weight of time-delay states (two times shown here) of the time series data and nonlinear functionals of this data
(quadratic functional shown here). Figure and legend from Gauthier *et al.* (2021).

## 1. NVAR for Lorenz strange attractor forecating

Lorenz attractor is defined by three coupled differential equations:

$$
\begin{split}
\dot{x} &= 10(y-x)\\
\dot{y} &= x(28-z)\\
\dot{z} &= xy - \frac{8z}{3}
\end{split}
$$

```python
# time step duration (in time unit)
dt = 0.025
# training time (in time unit)
train_time  = 10.
# testing time (idem)
test_time   = 120.
# warmup time (idem): should always be > k * s
warm_time   = 5.

# discretization
train_steps = round(train_time / dt)
test_steps  = round(test_time  / dt)
warm_steps  = round(warm_time  / dt)
```

```python
x0 = [17.67715816276679, 12.931379185960404, 43.91404334248268]
n_timesteps = train_steps + test_steps + warm_steps

X = lorenz(n_timesteps, x0=x0, h=dt, method="RK23")
```

```python
N = train_steps + warm_steps + test_steps

fig = plt.figure(figsize=(10, 10))
ax  = fig.add_subplot(111, projection='3d')
ax.set_title("Lorenz attractor (1963)")
ax.set_xlabel("$x$")
ax.set_ylabel("$y$")
ax.set_zlabel("$z$")
ax.grid(False)

for i in range(N-1):
    ax.plot(X[i:i+2, 0], X[i:i+2, 1], X[i:i+2, 2], color=plt.cm.magma(255*i//N), lw=1.0)

plt.show()
```

```python

```

We define a NVAR with delay $k=2$, strides $s=1$ and order $n=2$. The final feature vector $\mathbb{O}_{total}$ will hence be:

$$
\begin{align}
\mathbb{O}_{lin}[t] &=
    \begin{bmatrix}
         x_t\\y_t\\z_t\\x_{t-1}\\y_{t-1}\\z_{t-1}
    \end{bmatrix}
&
\mathbb{O}_{nonlin}[t] &=
    \begin{bmatrix}
       x_t^2\\x_t y_t\\x_t z_t\\x_t x_{t-1}\\x_t y_{t-1}\\x_t
       z_{t-1}\\y_t^2\\y_t z_t\\y_t x_{t-1}\\y_t y_{t-1}\\y_t z_{t-1}\\z_t^2\\z_t x_{t-1}\\z_t y_{t-1}\\z_t
       z_{t-1}\\x_{t-1}^2\\x_{t-1} y_{t-1}\\x_{t-1} z_{t-1}\\y_{t-1}^2\\y_{t-1} z_{t-1}\\z_{t-1}^2
    \end{bmatrix}
\end{align}
$$

$$
\mathbb{O}_{total}[t] = \mathbb{O}_{lin}[t] \oplus \mathbb{O}_{nonlin}[t]
$$

The NVAR is connected to a readout layer with offline learning using regularized linear regression. The regularization parameter is set to $2.5\times10^{-6}$.

The model is first trained to infer the value of $\mathbf{X}[t+1]$ knowing the value of $\mathbf{X}[t]$. This training step enable the model to learn an internal representation of the local dynamics of the attractor.

```python
nvar = NVAR(delay=2, order=2, strides=1)
readout = Ridge(3, ridge=2.5e-6)

model = nvar >> readout
```

We first warmup the NVAR. The warmup time can be as short as only $k \times s$ steps. The NVAR has relevant features as soon as all the delayed signals are non zeros, i.e. as soon as at leas $k \times s$ steps has been stores in the linear vector.

```python
_ = nvar.run(X[:warm_steps-1])
```

    Running NVAR-0: 100%|██████████████████████| 199/199 [00:00<00:00, 16676.99it/s]

Then, train the model to perform one-step-ahead prediction.

```python
Xi  = X[warm_steps-1:train_steps+warm_steps-1]
dXi = X[warm_steps:train_steps+warm_steps] - X[warm_steps-1:train_steps+warm_steps-1]

model = model.fit(Xi, dXi)
```

    Running Model-0:   0%|                                    | 0/1 [00:00<?, ?it/s]
    Running SubModel-acbe96c7-f28c-476f-8e23-d40f78914be8: 400it [00:00, 9133.79it/s][A
    Running Model-0: 100%|████████████████████████████| 1/1 [00:00<00:00, 21.79it/s]

    Fitting node Ridge-0...

```python
lin = ["$x_t$", "$y_t$", "$z_t$", "$x_{t-1}$", "$y_{t-1}$", "$z_{t-1}$"]

nonlin = []
for idx in nvar._monomial_idx:
    idx = idx.astype(int)
    if idx[0] == idx[1]:
        c = lin[idx[0]][:-1] + "^2$"
    else:
        c = " ".join((lin[idx[0]][:-1], lin[idx[1]][1:]))

    nonlin.append(c)

coefs = ["$c$"] + lin + nonlin
```

In the plot below are displayed the linear coefficients learned by the model.

```python
fig = plt.figure(figsize=(10, 10))

Wout = np.r_[readout.bias, readout.Wout]
x_Wout, y_Wout, z_Wout = Wout[:, 0], Wout[:, 1], Wout[:, 2]

ax = fig.add_subplot(131)
ax.set_xlim(-0.2, 0.2)
ax.grid(axis="y")
ax.set_xlabel("$[W_{out}]_x$")
ax.set_ylabel("Features")
ax.set_yticks(np.arange(len(coefs)))
ax.set_yticklabels(coefs[::-1])
ax.barh(np.arange(x_Wout.size), x_Wout.ravel()[::-1])

ax1 = fig.add_subplot(132)
ax1.set_xlim(-0.2, 0.2)
ax1.grid(axis="y")
ax1.set_yticks(np.arange(len(coefs)))
ax1.set_xlabel("$[W_{out}]_y$")
ax1.barh(np.arange(y_Wout.size), y_Wout.ravel()[::-1])

ax2 = fig.add_subplot(133)
ax2.set_xlim(-0.2, 0.2)
ax2.grid(axis="y")
ax2.set_yticks(np.arange(len(coefs)))
ax2.set_xlabel("$[W_{out}]_z$")
ax2.barh(np.arange(z_Wout.size), z_Wout.ravel()[::-1])

plt.show()
```

```python
nvar.run(X[warm_steps+train_steps-2:warm_steps+train_steps])
u = X[warm_steps+train_steps]
res = np.zeros((test_steps, readout.output_dim))
for i in range(test_steps):
    u = u + model(u)
    res[i, :] = u
```

    Running NVAR-0: 100%|███████████████████████████| 2/2 [00:00<00:00, 2889.63it/s]

```python
N = test_steps
Y = X[warm_steps+train_steps:]

fig = plt.figure(figsize=(15, 10))
ax  = fig.add_subplot(121, projection='3d')
ax.set_title("Generated attractor")
ax.set_xlabel("$x$")
ax.set_ylabel("$y$")
ax.set_zlabel("$z$")
ax.grid(False)

for i in range(N-1):
    ax.plot(res[i:i+2, 0], res[i:i+2, 1], res[i:i+2, 2], color=plt.cm.magma(255*i//N), lw=1.0)

ax2 = fig.add_subplot(122, projection='3d')
ax2.set_title("Real attractor")
ax2.grid(False)

for i in range(N-1):
    ax2.plot(Y[i:i+2, 0], Y[i:i+2, 1], Y[i:i+2, 2], color=plt.cm.magma(255*i//N), lw=1.0)
```

## 2. NVAR for double scroll strange attractor forecasting

```python
dt = 0.25
train_time  = 100.
test_time   = 800.
warm_time   = 1.
train_steps = round(train_time / dt)
test_steps  = round(test_time  / dt)
warm_steps  = round(warm_time  / dt)
```

```python
x0 = [0.37926545, 0.058339, -0.08167691]
n_timesteps = train_steps + test_steps + warm_steps

X = doublescroll(n_timesteps, x0=x0, h=dt, method="RK23")
```

```python
N = train_steps + warm_steps + test_steps

fig = plt.figure(figsize=(10, 10))
ax  = fig.add_subplot(111, projection='3d')
ax.set_title("Double scroll attractor (1998)")
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("z")
ax.grid(False)

for i in range(N-1):
    ax.plot(X[i:i+2, 0], X[i:i+2, 1], X[i:i+2, 2], color=plt.cm.cividis(255*i//N), lw=1.0)

plt.show()
```

```python
nvar = NVAR(delay=2, order=3, strides=1)
readout = Ridge(3, ridge=2.5e-6, input_bias=False)

model = nvar >> readout
```

```python
_ = nvar.run(X[:warm_steps-1])
```

```python
Xi  = X[warm_steps-1:train_steps+warm_steps-1]
dXi = X[warm_steps:train_steps+warm_steps] - X[warm_steps-1:train_steps+warm_steps-1]

model = model.fit(Xi, dXi)
```

```python
nvar.run(X[warm_steps+train_steps-2:warm_steps+train_steps])
u = X[warm_steps+train_steps]
res = np.zeros((test_steps, readout.output_dim))
for i in range(test_steps):
    u = u + model(u)
    res[i, :] = u
```

```python
N = test_steps
Y = X[warm_steps+train_steps:]

fig = plt.figure(figsize=(15, 10))
ax  = fig.add_subplot(121, projection='3d')
ax.set_title("Generated attractor")
ax.set_xlabel("$x$")
ax.set_ylabel("$y$")
ax.set_zlabel("$z$")
ax.grid(False)

for i in range(N-1):
    ax.plot(res[i:i+2, 0], res[i:i+2, 1], res[i:i+2, 2], color=plt.cm.cividis(255*i//N), lw=1.0)

ax2 = fig.add_subplot(122, projection='3d')
ax2.set_title("Real attractor")
ax2.grid(False)

for i in range(N-1):
    ax2.plot(Y[i:i+2, 0], Y[i:i+2, 1], Y[i:i+2, 2], color=plt.cm.cividis(255*i//N), lw=1.0)
```

```python
# time step duration (in time unit)
dt = 0.05
# training time (in time unit)
train_time  = 20.
# testing time (idem)
test_time   = 45.
# warmup time (idem): should always be > k * s
warm_time   = 5.

# discretization
train_steps = round(train_time / dt)
test_steps  = round(test_time  / dt)
warm_steps  = round(warm_time  / dt)
```

```python
x0 = [17.67715816276679, 12.931379185960404, 43.91404334248268]
n_timesteps = train_steps + test_steps + warm_steps

X = lorenz(n_timesteps, x0=x0, h=dt, method="RK23")
```

```python
nvar = NVAR(delay=4, order=2, strides=5)
readout = Ridge(1, ridge=0.05)

model = nvar >> readout
```

We first warmup the NVAR. The warmup time can be as short as only $k \times s$ steps. The NVAR has relevant features as soon as all the delayed signals are non zeros, i.e. as soon as at leas $k \times s$ steps has been stores in the linear vector.

```python
_ = nvar.run(X[:warm_steps-1, :2])
```

Then, train the model to perform one-step-ahead prediction.

```python
xy  = X[warm_steps-1:train_steps+warm_steps-1, :2]
z   = X[warm_steps:train_steps+warm_steps, 2][:, np.newaxis]

model = model.fit(xy, z)
```

```python
_ = nvar.run(X[train_steps:warm_steps+train_steps, :2])

xy_test = X[warm_steps+train_steps:-1, :2]
res = model.run(xy_test)
```

```python
fig = plt.figure(figsize=(10, 5))
ax = fig.add_subplot(111)
ax.plot(res, label="Inferred")
ax.plot(X[warm_steps+train_steps+1:, 2], label="Truth", linestyle="--")
ax.plot(abs(res[:, 0] - X[warm_steps+train_steps+1:, 2]), label="Absolute deviation")
ax.set_ylabel("$z$")
ax.set_xlabel("time")
ax.set_title("Lorenz attractor $z$ component inferred value")
ax.set_xticks(np.linspace(0, 900, 5))
ax.set_xticklabels(np.linspace(0, 900, 5) * dt + train_time + warm_time)
plt.legend()
plt.show()
```

```python

```

examples/Next Generation Reservoir Computing/NG-RC.ipynb