# Parameter Identification Technique (PIT) on static VCT with a nonlinear VMM model

# Purpose
Is it possible to use a similar regression as used in [04.02_PIT_linear_VMM.ipynb](04.02_PIT_linear_VMM.ipynb) but with data from static Virtual Captive Tests (VCT)?

# Methodology
* Load static data from VCT
* Find the best parameter values in the linear model using OLS linear regression.

# Setup

```python
# %load imports.py
%matplotlib inline
%load_ext autoreload
%autoreload 2
%config Completer.use_jedi = False  ## (To fix autocomplete)

## External packages:
import pandas as pd
pd.options.display.max_rows = 999
pd.options.display.max_columns = 999
pd.set_option("display.max_columns", None)

import numpy as np
import os
import matplotlib.pyplot as plt
#if os.name == 'nt':
#    plt.style.use('presentation.mplstyle')  # Windows

import plotly.express as px
import plotly.graph_objects as go

import seaborn as sns
import sympy as sp
from sympy.physics.mechanics import (dynamicsymbols, ReferenceFrame,
                                      Particle, Point)
from sympy.physics.vector.printing import vpprint, vlatex
from IPython.display import display, Math, Latex
from src.substitute_dynamic_symbols import run, lambdify

import pyro

import sklearn
import pykalman
from statsmodels.sandbox.regression.predstd import wls_prediction_std
import statsmodels.api as sm

from scipy.integrate import solve_ivp

## Local packages:
from src.data import mdl
import src.models.nonlinear_martin_vmm as vmm
import src.nonlinear_martin_vmm_equations as eq
from src.symbols import *
import src.symbols as symbols
from src import prime_system
from src.models import regression
from src.visualization.plot import track_plot

```

## Load VCT data

```python
df_VCT_all = pd.read_csv('../data/external/vct.csv', index_col=0)
df_VCT_all.head()
```

```python
df_VCT = df_VCT_all.groupby(by=['model_name']).get_group('V2_5_MDL_modelScale')
```

```python
df_VCT.describe()
```

```python
df_VCT['test type'].unique()
```

# Subtract the resistance

```python
df_resistance = df_VCT.groupby(by='test type').get_group('resistance')

X = df_resistance[['u','fx']].copy()
X['u**2'] = X['u']**2
y = X.pop('fx')

model_resistance = sm.OLS(y,X)
results_resistance = model_resistance.fit()

X_pred = pd.DataFrame()
X_pred['u'] = np.linspace(X['u'].min(), X['u'].max(), 20)
X_pred['u**2'] = X_pred['u']**2
X_pred['fx'] = results_resistance.predict(X_pred)

fig,ax=plt.subplots()
df_resistance.plot(x='u', y='fx', style='.', ax=ax)
X_pred.plot(x='u', y='fx', style='--', ax=ax);

```

```python
df_VCT_0_resistance = df_VCT.copy()
df_VCT_0_resistance['u**2'] = df_VCT_0_resistance['u']**2
#df_VCT_0_resistance['fx']-= results_resistance.predict(df_VCT_0_resistance[['u','u**2']])
df_VCT_0_resistance['thrust']= results_resistance.predict(df_VCT_0_resistance[['u','u**2']])  # Balancing the resistance
```

## Load test

```python
df_runs = mdl.runs()
```

```python
#id=22773
#id=22616
#id=22774
id=22770

df, units, meta_data = mdl.load(id=id, dir_path='../data/processed/kalman')
df.index = df.index.total_seconds()
df = df.iloc[0:-100].copy()
df.index-=df.index[0]
df.sort_index(inplace=True)
```

```python
meta_data['rho']=1000
meta_data['mass'] = meta_data['Volume']*meta_data['rho']
meta_data.dropna()
```

```python
df.head()
```

```python
df.plot(y=['Prop/PS/Thrust','Prop/SB/Thrust','Fan/Fore/Fx','Fan/Aft/Fx'])
```

```python
df['thrust'] = df['Prop/PS/Thrust'] + df['Prop/SB/Thrust']
```

```python
from src.visualization.plot import track_plot
fig,ax=plt.subplots()
#fig.set_size_inches(10,10)
track_plot(df=df, lpp=meta_data.lpp, x_dataset='x0', y_dataset='y0',  psi_dataset='psi', beam=meta_data.beam, ax=ax);
```

# Ship parameters

```python
T_ = (meta_data.TA + meta_data.TF)/2
L_ = meta_data.lpp
m_ = meta_data.mass
rho_ = meta_data.rho
B_ = meta_data.beam
CB_ = m_/(T_*B_*L_*rho_)
volume_ = m_/rho_
I_z_ = m_*meta_data.KZZ**2

ship_parameters = {
        'T' : T_,
        'L' : L_,
        'CB' :CB_,
        'B' : B_,
        'rho' : rho_,
        'x_G' : meta_data.lcg,  # motions are expressed at CG
        'm' : m_,
        'I_z': I_z_,
        'volume':volume_,
    }

ps = prime_system.PrimeSystem(**ship_parameters)  # model

scale_factor = meta_data.scale_factor
ps_ship = prime_system.PrimeSystem(L=ship_parameters['L']*scale_factor, rho=meta_data['rho'])  # ship
```

# VCT to prime system

```python
interesting = [
    'u',
    'v',
    'r',
    'delta',
    'fx',
    'fy',
    'mz',
    'thrust',
]
df_VCT_prime = ps_ship.prime(df_VCT_0_resistance[interesting], U=df_VCT_0_resistance['V'])
```

```python
df_VCT_prime.describe()
```

# Nonlinear Vessel Manoeuvring Model (VMM)
Implementing according to:
Matusiak, Jerzy. Dynamics of a Rigid Ship. Aalto University, 2017. https://aaltodoc.aalto.fi:443/handle/123456789/24408.

## N
The yaw is described by the following Ordinary Differential Equation (ODE):

```python
Math(vlatex(eq.N_eq))
```

The hydrodynamic moment part of this model is:

```python
Math(vlatex(eq.mz_eq))
```

```python
N_eq = eq.mz_eq.subs(eq.N_qs,sp.solve(eq.N_qs_eq,N_qs)[0])
Math(vlatex(N_eq))
```

```python
N_eq = N_eq.subs([
    #(x_G,0),  # Assuming or moving to CG=0
    #(I_z,1),  # Removing inertia
    #(eq.p.Nrdot,0),  # Removing added mass

    (u1d,0),
    (v1d,0),
    (r1d,0),

])

#solution = sp.solve(N_eq,r1d)[0]
#N_eq = sp.Eq(r1d, solution*(I_z-eq.p.Nrdot))  # Putting r1d on the LHS
```

```python
diff_eq_N = regression.DiffEqToMatrix(ode=N_eq, label=N_nonlin, base_features=[delta,u,v,r])
```

If accelerations (dotted states) are dropped from this equation, the quasi-static force model is obtained:

```python
Math(vlatex(diff_eq_N.acceleration_equation))
```

This is a model that now only depend on steady velocities and rudder angle and will be the same thing as the forces/moment coming from the VCT calculations (CFD):

$ m_{z}^{VCT} = N_{lin} $

The hydrodynamic derivatives (the parameters in the matematical model) can be regressed by solving the following linear regression problem:

$ y = X \cdot \beta + \epsilon $

Where $y$ and $X$ and $\beta$ are:

```python
Math(vlatex(diff_eq_N.acceleration_equation_x))
```

```python
Math(vlatex(diff_eq_N.eq_y))
```

```python
diff_eq_N.eq_beta
```

```python
Math(vlatex(diff_eq_N.eq_X))
```

```python
diff_eq_N.X_lambda
```

```python
from statsmodels.sandbox.regression.predstd import wls_prediction_std
def show_pred(X,y,results, label):

    display(results.summary())

    X_ = X.copy()
    X_['y'] = y
    X_.sort_values(by='y', inplace=True)

    y_ = X_.pop('y')

    y_pred = results.predict(X_)

    prstd, iv_l, iv_u = wls_prediction_std(results, exog=X_, alpha=0.05)
    #iv_l*=-1
    #iv_u*=-1

    fig,ax=plt.subplots()
    #ax.plot(X_.index,y_, label='Numerical gradient from model test')
    #ax.plot(X_.index,y_pred, '--', label='OLS')

    ax.plot(y_,y_pred, '.')
    ax.plot([y_.min(),y_.max()], [y_.min(),y_.max()], 'r-')

    ax.set_ylabel(f'{label} (prediction)')
    ax.set_xlabel(label)

    ax.fill_between(y_, y1=iv_l, y2=iv_u, zorder=-10, color='grey', alpha=0.5, label=r'5% confidence')
    ax.legend();
```

```python
X = diff_eq_N.calculate_features(data=df_VCT_prime)
y = diff_eq_N.calculate_label(y=df_VCT_prime['mz'])

model_N = sm.OLS(y,X)
results_N = model_N.fit()

show_pred(X=X,y=y,results=results_N, label=r'$N$')
```

## Y

```python
Math(vlatex(eq.Y_eom))
```

```python
Y_eq = eq.fy_eq.subs(eq.Y_qs,sp.solve(eq.Y_qs_eq,Y_qs)[0])
Math(vlatex(Y_eq))
```

```python
Y_eq = Y_eq.subs([
    #(eq.p.Yvdot,1),  # Removing added mass
    #(x_G,0),
    (u1d,0),
    (v1d,0),
    (r1d,0),

])
#solution = sp.solve(Y_eq,v1d)[0]
#solution = solution*(-eq.p.Yvdot+m)
#solution = solution + U*m*r# (adding u*m*r to the measurement fy instead)
#Y_eq = sp.simplify(sp.Eq(v1d, solution))  # Putting r1d on the LHS
```

```python
Math(vlatex(Y_eq))
```

```python
diff_eq_Y = regression.DiffEqToMatrix(ode=Y_eq, label=Y_nonlin, base_features=[delta,u,v,r])
```

```python
diff_eq_Y.eq_beta
```

```python
diff_eq_Y.acceleration_equation_x
```

```python
X = diff_eq_Y.calculate_features(data=df_VCT_prime, simplify_names=True)
y = diff_eq_Y.calculate_label(y=df_VCT_prime['fy'])
#y+=df['u']*df['r']  # adding centrifugal force

model_Y = sm.OLS(y,X)
results_Y = model_Y.fit()

show_pred(X=X,y=y,results=results_Y, label=r'$Y$')
```

```python
X.head()
```

## X

```python
X_eq = eq.fx_eq.subs(eq.X_qs,sp.solve(eq.X_qs_eq,X_qs)[0])
Math(vlatex(X_eq))
```

```python

X_eq = X_eq.subs([
    #(eq.p.Xudot,1),  # Removing added mass
    (u1d,0),
    (v1d,0),
    (r1d,0),
    #(m,0),  # mass

])

#solution = sp.solve(X_eq,u1d)[0]
#X_eq = sp.Eq(u1d, solution*(-eq.p.Xudot+m))  # Putting r1d on the LHS
```

```python
Math(vlatex(X_eq))
```

```python
diff_eq_X = regression.DiffEqToMatrix(ode=X_eq, label=X_nonlin, base_features=[delta,u,v,r,thrust])
```

```python
X = diff_eq_X.calculate_features(data=df_VCT_prime)
y = diff_eq_X.calculate_label(y=df_VCT_prime['fx'])

model_X = sm.OLS(y,X)
results_X = model_X.fit()

show_pred(X=X,y=y,results=results_X, label=r'$X$')
```

```python
results_summary_X = regression.results_summary_to_dataframe(results_X)
results_summary_Y = regression.results_summary_to_dataframe(results_Y)
results_summary_N = regression.results_summary_to_dataframe(results_N)
```

# Brix parameters

```python
def calculate_prime(row, ship_parameters):
    return run(function=row['brix_lambda'], inputs=ship_parameters)

mask = df_parameters['brix_lambda'].notnull()
df_parameters.loc[mask,'brix_prime'] = df_parameters.loc[mask].apply(calculate_prime, ship_parameters=ship_parameters, axis=1)
df_parameters['brix_prime'].fillna(0, inplace=True)
```

## Add the regressed parameters
Hydrodynamic derivatives that depend on acceleration cannot be obtained from the VCT regression. They are however essential if a time simulation should be conducted. These values have then been taken from Brix semi empirical formulas for the simulations below.

```python
df_parameters_all = df_parameters.copy()
for other in [results_summary_X, results_summary_Y, results_summary_N]:
    df_parameters_all = df_parameters_all.combine_first(other)

df_parameters_all.rename(columns={'coeff':'regressed'}, inplace=True)
df_parameters_all.drop(columns=['brix_lambda'], inplace=True)

df_parameters_all['prime'] = df_parameters_all['regressed'].combine_first(df_parameters_all['brix_prime'])  # prefer regressed
```

```python
df_parameters_all
```

```python
fig,ax=plt.subplots()
fig.set_size_inches(15,5)
mask = ((df_parameters_all['brix_prime'] != 0) |
        (df_parameters_all['regressed']).notnull())

df_parameters_plot = df_parameters_all.loc[mask]
df_parameters_plot.plot.bar(y=['brix_prime','regressed'], ax=ax);
```

It seems that the regressed hydrodynamic derivatives are reasonable similar to the corresponding predicted values with semi-empirical formulas from Brix.

# Simulation

```python
df_ = df.copy()
t = df_.index
control = df_[['delta','thrust']]
#control = {'delta':0,}

df_0 = df_.iloc[0:100].median(axis=0)
y0 = {
    'u' : df_0['u'],
    'v' : df_0['v'],
    'r' : df_0['r'],
    'x0' : df_0['x0'],
    'y0' : df_0['y0'],
    'psi' : df_0['psi'],
    }

#solution = model.simulate(y0=y0, t=t, df_parameters=df_parameters, df_ship_parameters=df_ship_parameters, control=control, rtol=1e-3, atol=1e-3)
simulation = vmm.NonLinearSimulation()
solution = simulation.simulate(y0=y0, t=t, df_parameters=df_parameters_all, ship_parameters=ship_parameters, control=control)

columns = list(y0.keys())
df_result_prime = pd.DataFrame(data=solution.y.T, columns=columns)
df_result_prime.index=t[0:len(df_result_prime)]

#df_result_prime['U'] = np.sqrt(df_result_prime['u']**2 + df_result_prime['v']**2)

df_result_prime.plot(y='u')
df_result_prime.plot(y='v')
df_result_prime.plot(y='r')

```

```python
U_ = np.sqrt(df_0['u']**2 + df_0['v']**2)
df_result = ps.unprime(values=df_result_prime, U=U_)
df_result['beta'] = -np.arctan2(df_result['v'],df_result['u'])
```

```python
fig,ax=plt.subplots()
track_plot(df=df, lpp=ship_parameters['L'], beam=ship_parameters['B'],ax=ax, label='model test')
track_plot(df=df_result, lpp=ship_parameters['L'], beam=ship_parameters['B'],ax=ax, label='simulation', color='green')
ax.legend()

for key in df_result:
    fig,ax = plt.subplots()
    df.plot(y=key, label='model test', ax=ax)
    df_result.plot(y=key, label='simulation', ax=ax)
    ax.set_ylabel(key)

```

```python

```


# Polar and Cylindrical Frame of Reference

> Renato Naville Watanabe, Marcos Duarte
> [Laboratory of Biomechanics and Motor Control](http://pesquisa.ufabc.edu.br/bmclab)
> Federal University of ABC, Brazil

Table of Contents

1  Python setup

2  Versors ${\bf\hat{e_R}}$ and ${\bf\hat{e_\theta}}$

2.1  Tridimensional case

3  Time-derivative of the versors ${\bf\hat{e_R}}$ and ${\bf\hat{e_\theta}}$

4  Position, velocity and acceleration

4.1  Position

4.2  Velocity

4.3  Acceleration

5  Important to note

6  Example

6.1  Solving numerically

6.2  Solved symbolically

7  Further reading

8  Video lectures on the Internet

9  Problems

10  Reference

## Python setup

```python
import numpy as np
import sympy as sym
from sympy.plotting import plot_parametric,plot3d_parametric_line
from sympy.vector import CoordSys3D
import matplotlib.pyplot as plt
# from matplotlib import rc
# rc('text', usetex=True)
sym.init_printing()
```

## Versors ${\bf\hat{e_R}}$ and ${\bf\hat{e_\theta}}$

Consider that we have the position vector $\bf\vec{r}$ of a particle, moving in a circular path indicated in the figure below by a dashed line. This vector $\bf\vec{r(t)}$ is described in a fixed reference frame as:

\begin{equation}
{\bf\vec{r}}(t) = x(t){\bf\hat{i}} + y(t){\bf\hat{j}} + z(t){\bf\hat{k}}
\label{eq_1}
\end{equation}

Figure. Position vector of a moving particle moving in a circular path.

Naturally, we could describe all the kinematic variables in the fixed reference frame. But in circular motions, is convenient to define a basis with a vector in the direction of the position vector  $\bf\vec{r}$. So, the vector  $\bf\hat{e_R}$ is defined as:

\begin{equation}
{\bf\hat{e_R}} =  \frac{\bf\vec{r}}{\Vert{\bf\vec{r} }\Vert}
\label{eq_2}
\end{equation}

The second vector of the basis can be obtained by the cross multiplication between $\bf\hat{k}$ and $\bf\hat{e_R}$:

\begin{equation}
{\bf\hat{e_\theta}} = {\bf\hat{k}} \times {\bf\hat{e_R}}
\label{eq_3}
\end{equation}

The third vector of the basis is the conventional ${\bf\hat{k}}$ vector.

Figure. Position vector of a particle moving in a circular path and a time-varying basis>

### Tridimensional case

This basis can be used also for non-circular movements. For a 3D movement, the versor ${\bf\hat{e_R}}$ is obtained by removing the projection of the vector ${\bf\vec{r}}$ onto the versor ${\bf\hat{k}}$:

\begin{equation}
{\bf\hat{e_R}} =  \frac{\bf\vec{r} - ({\bf\vec{r}.{\bf\hat{k}}){\bf\hat{k}}}}{\Vert\bf\vec{r} - ({\bf\vec{r}.{\bf\hat{k}}){\bf\hat{k}}\Vert}}
\label{eq_4}
\end{equation}

Figure. Position vector of a moving particle and a time-varying basis>

Note that if the movement is on a plane, the expression above is equal to ${\bf\hat{e_R}} =  \frac{\bf\vec{r}}{\Vert{\bf\vec{r} }\Vert}$ since the projection of $\bf\vec{r}$ on the $\bf\hat{k}$ versor is zero.

## Time-derivative of the versors ${\bf\hat{e_R}}$ and ${\bf\hat{e_\theta}}$

To obtain the expressions of the velocity and acceleration vectors, it is necessary to obtain the expressions of the time-derivative of the vectors ${\bf\hat{e_R}}$ and ${\bf\hat{e_\theta}}$.

This can be done by noting that:

\begin{align}
{\bf\hat{e_R}} &= \cos(\theta){\bf\hat{i}} + \sin(\theta){\bf\hat{j}}\\
{\bf\hat{e_\theta}} &= -\sin(\theta){\bf\hat{i}} + \cos(\theta){\bf\hat{j}}
\label{eq_5}
\end{align}

Deriving ${\bf\hat{e_R}}$ we obtain:

\begin{equation}
\frac{d{\bf\hat{e_R}}}{dt} = -\sin(\theta)\dot\theta{\bf\hat{i}} + \cos(\theta)\dot\theta{\bf\hat{j}} = \dot{\theta}{\bf\hat{e_\theta}}
\label{eq_6}
\end{equation}

Similarly, we obtain the time-derivative of ${\bf\hat{e_\theta}}$:

\begin{equation}
    \frac{d{\bf\hat{e_\theta}}}{dt} = -\cos(\theta)\dot\theta{\bf\hat{i}} - \sin(\theta)\dot\theta{\bf\hat{j}} = -\dot{\theta}{\bf\hat{e_R}}
\label{eq_7}
\end{equation}

## Position, velocity and acceleration

### Position

The position vector $\bf\vec{r}$, from the definition of $\bf\hat{e_R}$, is:

\begin{equation}
{\bf\vec{r}} = R{\bf\hat{e_R}} + z{\bf\hat{k}}
\label{eq_8}
\end{equation}

where $R = \Vert\bf\vec{r} - ({\bf\vec{r}.{\bf\hat{k}}){\bf\hat{k}}\Vert}$.

### Velocity

The velocity vector $\bf\vec{v}$ is obtained by deriving the vector $\bf\vec{r}$:

\begin{equation}
{\bf\vec{v}} = \frac{d(R{\bf\hat{e_R}})}{dt} + \dot{z}{\bf\hat{k}} = \dot{R}{\bf\hat{e_R}}+R\frac{d\bf\hat{e_R}}{dt}=\dot{R}{\bf\hat{e_R}}+R\dot{\theta}{\bf\hat{e_\theta}}+ \dot{z}{\bf\hat{k}}
\label{eq_9}
\end{equation}

### Acceleration

The acceleration vector $\bf\vec{a}$ is obtained by deriving the velocity vector:

\begin{align}
{\bf\vec{a}} =& \frac{d(\dot{R}{\bf\hat{e_R}}+R\dot{\theta}{\bf\hat{e_\theta}}+\dot{z}{\bf\hat{k}})}{dt}=\\\nonumber
    =&\ddot{R}{\bf\hat{e_R}}+\dot{R}\frac{d\bf\hat{e_R}}{dt} + \dot{R}\dot{\theta}{\bf\hat{e_\theta}} + R\ddot{\theta}{\bf\hat{e_\theta}} + R\dot{\theta}\frac{d{\bf\hat{e_\theta}}}{dt} + \ddot{z}{\bf\hat{k}}=\\\nonumber
    =&\ddot{R}{\bf\hat{e_R}}+\dot{R}\dot{\theta}{\bf\hat{e_\theta}} + \dot{R}\dot{\theta}{\bf\hat{e_\theta}} + R\ddot{\theta}{\bf\hat{e_\theta}} - R\dot{\theta}^2{\bf\hat{e_R}}+ \ddot{z}{\bf\hat{k}} =\\
    =&\ddot{R}{\bf\hat{e_R}}+2\dot{R}\dot{\theta}{\bf\hat{e_\theta}}+ R\ddot{\theta}{\bf\hat{e_\theta}} - {R}\dot{\theta}^2{\bf\hat{e_R}}+ \ddot{z}{\bf\hat{k}} =\\\nonumber
    =&(\ddot{R}-R\dot{\theta}^2){\bf\hat{e_R}}+(2\dot{R}\dot{\theta} + R\ddot{\theta}){\bf\hat{e_\theta}}+ \ddot{z}{\bf\hat{k}}\nonumber
\label{eq_10}
\end{align}

+ The term $\ddot{R}$ is  an acceleration in the radial direction.

+ The term $R\ddot{\theta}$ is an angular acceleration.

+ The term $\ddot{z}$ is an acceleration in the $\bf\hat{k}$ direction.

+ The term $-R\dot{\theta}^2$ is the well known centripetal acceleration.

+ The term $2\dot{R}\dot{\theta}$ is known as Coriolis acceleration. This term may be difficult to understand intuitively. It appears when there is displacement in the radial and angular directions at the same time.

## Important to note

The reader must bear in mind that the use of a different basis to represent the position, velocity or acceleration vectors is only a different representation of the same vector. For example, for the acceleration vector:

\begin{equation}
{\bf\vec{a}} = \ddot{x}{\bf\hat{i}}+ \ddot{y}{\bf\hat{j}} + \ddot{z}{\bf\hat{k}}=(\ddot{R}-R\dot{\theta}^2){\bf\hat{e_R}}+(2\dot{R}\dot{\theta} + R\ddot{\theta}){\bf\hat{e_\theta}} + \ddot{z}{\bf\hat{k}}=\dot{\Vert\bf\vec{v}\Vert}{\bf\hat{e}_t}+{\Vert\bf\vec{v}\Vert}^2\Vert{\bf\vec{C}} \Vert{\bf\hat{e}_n}
\label{eq_11}
\end{equation}

In which the last equality is the acceleration vector represented in the path-coordinate of the particle (see http://nbviewer.jupyter.org/github/BMClab/bmc/blob/master/notebooks/Time-varying%20frames.ipynb).

## Example

Consider a particle following the spiral path described below:

\begin{equation}
{\bf\vec{r}}(t) = (2\sqrt{t}\cos(t)){\bf\hat{i}}+ (2\sqrt{t}\sin(t)){\bf\hat{j}}
\label{eq_12}
\end{equation}

### Solving numerically

```python
t = np.linspace(0.01,10,30).reshape(-1,1) #create a time vector and reshapes it to a column vector
R = 2*np.sqrt(t)
theta = t
rx = R*np.cos(t)
ry = R*np.sin(t)
r = np.hstack((rx, ry)) # creates the position vector by stacking rx and ry horizontally
```

```python
e_r = r/np.linalg.norm(r, axis=1, keepdims=True) # defines e_r vector
e_theta = np.cross([0,0,1],e_r)[:,0:-1] # defines e_theta vector

```

```python
dt = t[1] #defines delta_t
Rdot = np.diff(R, axis=0)/dt #find the R derivative
thetaDot = np.diff(theta, axis=0)/dt #find the angle derivative
v = Rdot*e_r[0:-1,:] +R[0:-1]*thetaDot*e_theta[0:-1,:] # find the linear velocity.
```

```python
Rddot = np.diff(Rdot, axis=0)/dt
thetaddot = np.diff(thetaDot, axis=0)/dt
```

```python
a = ((Rddot - R[1:-1]*thetaDot[0:-1]**2)*e_r[1:-1,:]
     + (2*Rdot[0:-1]*thetaDot[0:-1] + Rdot[0:-1]*thetaddot)*e_theta[1:-1,:])
```

The versors $\bf\hat{e_R}$, $\bf\hat{e_\theta}$ are plotted below at some points of the path.

```python
from matplotlib.patches import FancyArrowPatch
%matplotlib inline
plt.rcParams['figure.figsize'] = (8,8)
fig = plt.figure()

ax = fig.add_axes([0,0,1,1])
plt.plot(r[:,0],r[:,1],'.')
for i in np.arange(len(t)-2):
    vec1 = FancyArrowPatch(r[i,:],r[i,:]+e_r[i,:],mutation_scale=30,color='r', label='e_r')
    vec2 = FancyArrowPatch(r[i,:],r[i,:]+e_theta[i,:],mutation_scale=30,color='g', label='e_theta')
    ax.add_artist(vec1)
    ax.add_artist(vec2)
plt.xlim((-10,10))
plt.ylim((-10,10))
plt.grid()
plt.legend([vec1, vec2],[r'$\vec{e_r}$', r'$\vec{e_{\theta}}$'])
plt.show()
```

The velocity and acceleleration of the particle are plotted below at some points of the path.

```python
from matplotlib.patches import FancyArrowPatch
%matplotlib inline
plt.rcParams['figure.figsize'] = (8,8)
fig = plt.figure()
plt.plot(r[:,0],r[:,1],'.')
ax = fig.add_axes([0,0,1,1])
for i in np.arange(len(t)-2):
    vec1 = FancyArrowPatch(r[i,:],r[i,:]+v[i,:],mutation_scale=10,color='r')
    vec2 = FancyArrowPatch(r[i,:],r[i,:]+a[i,:],mutation_scale=10,color='g')
    ax.add_artist(vec1)
    ax.add_artist(vec2)
plt.xlim((-10,10))
plt.ylim((-10,10))
plt.grid()
plt.legend([vec1, vec2],[r'$\vec{v}$', r'$\vec{a}$'])
plt.show()
```

### Solved symbolically

```python
O = sym.vector.CoordSys3D(' ')
t = sym.symbols('t')
```

```python
r = 2*sym.sqrt(t)*sym.cos(t)*O.i+2*sym.sqrt(t)*sym.sin(t)*O.j
r
```

```python
plot_parametric(r.dot(O.i),r.dot(O.j),(t,0,10))
```

```python
e_r = r - r.dot(O.k)*O.k
e_r = e_r/sym.sqrt(e_r.dot(O.i)**2+e_r.dot(O.j)**2+e_r.dot(O.k)**2)
```

```python
e_r
```

```python
e_theta = O.k.cross(e_r)
e_theta
```

```python
from matplotlib.patches import FancyArrowPatch
plt.rcParams['figure.figsize'] = (8,8)
fig = plt.figure()
ax = fig.add_axes([0, 0, 1, 1])
ax.axis("on")
time = np.linspace(0,10,30)
for instant in time:
    vt = FancyArrowPatch([float(r.dot(O.i).subs(t,instant)),float(r.dot(O.j).subs(t,instant))],
                         [float(r.dot(O.i).subs(t,instant))+float(e_r.dot(O.i).subs(t,instant)),
                          float(r.dot(O.j).subs(t, instant))+float(e_r.dot(O.j).subs(t,instant))],
                         mutation_scale=20,
                         arrowstyle="->",color="r",label='${{e_r}}$')
    vn = FancyArrowPatch([float(r.dot(O.i).subs(t, instant)),float(r.dot(O.j).subs(t,instant))],
                         [float(r.dot(O.i).subs(t, instant))+float(e_theta.dot(O.i).subs(t, instant)),
                          float(r.dot(O.j).subs(t, instant))+float(e_theta.dot(O.j).subs(t, instant))],
                         mutation_scale=20,
                         arrowstyle="->",color="g",label='${{e_{theta}}}$')
    ax.add_artist(vn)
    ax.add_artist(vt)
plt.xlim((-10,10))
plt.ylim((-10,10))
plt.legend(handles=[vt,vn],fontsize=20)
plt.grid()
plt.show()
```

```python
R = 2*sym.sqrt(t)
```

```python
Rdot = sym.diff(R,t)
Rddot = sym.diff(Rdot,t)
Rddot
```

```python
v = Rdot*e_r + R*e_theta
```

```python
v
```

```python
a = (Rddot - R)*e_r + (2*Rdot*1+0)*e_theta
aCor = 2*Rdot*1*e_theta
aCor
```

```python
a
```

```python
from matplotlib.patches import FancyArrowPatch
plt.rcParams['figure.figsize'] = (8,8)
fig = plt.figure()
ax = fig.add_axes([0, 0, 1, 1])
ax.axis("on")
time = np.linspace(0.1,10,30)
for instant in time:
    vt = FancyArrowPatch([float(r.dot(O.i).subs(t,instant)),float(r.dot(O.j).subs(t,instant))],
                         [float(r.dot(O.i).subs(t,instant))+float(v.dot(O.i).subs(t,instant)),
                          float(r.dot(O.j).subs(t, instant))+float(v.dot(O.j).subs(t,instant))],
                         mutation_scale=20,
                         arrowstyle="->",color="r",label='${{v}}$')
    vn = FancyArrowPatch([float(r.dot(O.i).subs(t, instant)),float(r.dot(O.j).subs(t,instant))],
                         [float(r.dot(O.i).subs(t, instant))+float(a.dot(O.i).subs(t, instant)),
                          float(r.dot(O.j).subs(t, instant))+float(a.dot(O.j).subs(t, instant))],
                         mutation_scale=20,
                         arrowstyle="->",color="g",label='${{a}}$')
    vc = FancyArrowPatch([float(r.dot(O.i).subs(t, instant)),float(r.dot(O.j).subs(t,instant))],
                         [float(r.dot(O.i).subs(t, instant))+float(aCor.dot(O.i).subs(t, instant)),
                          float(r.dot(O.j).subs(t, instant))+float(aCor.dot(O.j).subs(t, instant))],
                         mutation_scale=20,
                         arrowstyle="->",color="b",label='${{a_{Cor}}}$')
    ax.add_artist(vn)
    ax.add_artist(vt)
    ax.add_artist(vc)
plt.xlim((-10,10))
plt.ylim((-10,10))
plt.legend(handles=[vt,vn,vc],fontsize=20)
plt.grid()
plt.show()
```

## Further reading

- Read pages 916-931 of the 18th chapter of the [Ruina and Rudra's book] (http://ruina.tam.cornell.edu/Book/index.html) about polar coordinates.

## Video lectures on the Internet

- Khan Academy:
    - [Polar coordinates 1](https://www.youtube.com/watch?v=B5dOy4m6I1E)
    - [Polar coordinates 2](https://www.youtube.com/watch?v=1pQXJuWbz9w)
    - [Polar coordinates 3](https://www.youtube.com/watch?v=roSG9V3zApM)

## Problems

1. Find the polar basis (using the computer) for a projectile motion of a particle following the parametric equation below:

\begin{equation}
\vec{r}(t) = (10t-51){\bf{\hat{i}}} + \left(-\frac{9,81}{2}t^2+50t+100\right){\bf{\hat{j}}}
\label{eq_13}
\end{equation}

2. Problems from 15.1.1 to 15.1.14 from Ruina and Rudra's book,
3. Problems from 18.1.1 to 18.1.8 and 18.1.10 from Ruina and Rudra's book.

## Reference

- Ruina A, Rudra P (2019) [Introduction to Mechanics for Engineers](http://ruina.tam.cornell.edu/Book/index.html). Oxford University Press.


# Implementation of a Devito viscoacoustic equations

## This tutorial is contributed by SENAI CIMATEC (2020)

This tutorial is based on:

**Linear inversion in layered viscoacoustic media using a time‐domain method** (1994)

Joakim O. Blanch and William W. Symes

SEG Technical Program Expanded Abstracts

https://doi.org/10.1190/1.1822695

**True-amplitude prestack depth migration** (2007)

Feng Deng and George A. McMechan

GEOPHYSICS Technical Papers

https://doi.org/10.1190/1.2714334

**Attenuation compensation for least-squares reverse time migration using the viscoacoustic-wave equation** (2014)

Gaurav Dutta and Gerard T. Schuster

GEOPHYSICS Technical Papers

https://doi.org/10.1190/geo2013-0414.1

**Multiscale viscoacoustic waveform inversion with the second generation wavelet transform and adaptive time–space domain finite-difference method** (2014)

Zhiming Ren, Yang Liu,and Qunshan Zhang

Geophysical Journal International, Volume 197, Issue 2, 1 May 2014, Pages 948–974

https://doi.org/10.1093/gji/ggu024

**Viscoacoustic prestack reverse time migration based on the optimal time-space domain high-order finite-difference method** (2014)

Yan Zhao, Yang Liu, and Zhi-Ming Ren

Appl. Geophys. 11, 50–62.

https://doi.org/10.1007/s11770-014-0414-8

**A stable and efficient approach of Q reverse time migration** (2018)

Yan Zhao, Ningbo Mao, and Zhiming Ren

GEOPHYSICS Technical Papers

https://doi.org/10.1190/geo2018-0022.1

## Introduction

The conversion of mechanical energy to heat, occurs during the propagation of seismic waves on the subsurface, due to the viscosity of the rocks. The presence of oil and gas in these rocks causes seismic attenuations. Thus, associated effects, such as dispersion and dissipation, can significantly affect the amplitudes, as well as the phase of the seismic pulse.  However, in the seismic exploration, the subsurface has still been considered as an ideal elastic/acoustic medium, that is, disregarding its mitigating effect. In practice, the propagation of seismic waves on the subsurface is in many ways different from propagation in an ideal solid.

For example, some subsurface rocks have anisotropic properties, are heterogeneous, porous and so on. The acoustic/elastic wave equation is not sensitive enough to describe propagation in these more complicated mediums. Generally, the viscosity of materials in the subsurface causes energy dissipation and consequently a decrease in amplitude, in addition to modifying the frequency content of the waves. This phenomenon of energy dissipation of the wave is called seismic absorption or attenuation.

The goal of this tutorial is to perform a seismic modeling taking into account the viscosity of the medium, so that it is possible to more accurately simulate the seismic data and consequently build images with better resolution in the processing of this data, in addition to extracting more detailed information on rocky materials through seismic inversion.

This tutorial follow three main viscoacoustic approaches in time-space domain:

- Blanch and Symes (1995) / Dutta and Schuster (2014)

- Ren et al. (2014)

- Deng and McMechan (2007)

Table of symbols

| Symbol         | Description
| :---         | :---
|$f$                                  |Frequency                                            |
|$f_o$                                |Reference frequency                                  |
|$\omega$                             |Angular frenquency                                   |
|$\omega_0$                           |Angular Reference Frequency                          |
|$v$                                  |Velocity model                                       |
|$v_0$                                |Reference velocity at $\omega_0$                     |
|$\kappa$                             |Bulk modulus                                         |
|$g$                                  |Absorption coefficient                               |
|$\tau$                               |Relaxation time                                      |
|$\tau_\sigma$                        |Stress relaxation parameter                          |
|$\tau_\epsilon$                      |Strain relaxation parameter                          |
|$Q$                                  |Quality factor                                       |
|$\eta$                               |Viscosity                                            |
|$\rho$                               |Density                                              |
|$\nabla$                             |Nabla operator                                       |
|$P({\bf x},t)$                       |Pressure field                                       |
|$r({\bf x},t)$                       |Memory variable                                      |
|${\bf v}({\bf x},t)$                 |Particle velocity                                    |
|$S({\bf x}_s,t)$                     |Source                                               |

# Seismic modelling with Devito

Before start with the viscoacoustic approaches we will describe a setup of seismic modelling with Devito in a simple 2D case. We will create a physical model of our domain and define a single source and an according set of receivers to model for the forward model. But first, we initialize some basic utilities.

```python
import numpy as np
import sympy as sp

import matplotlib.pyplot as plt

from devito import *
from examples.seismic.source import RickerSource, WaveletSource, TimeAxis
from examples.seismic import ModelViscoacoustic, plot_image, setup_geometry, plot_velocity
```

```python
nx = 300
nz = 300

# Define a physical size
shape = (nx, nz)
spacing = (20., 20.)
origin = (0., 0.)
nlayers = 3
nbl = 50
space_order = 8
dtype = np.float32

# Model physical parameters:
vp = np.zeros(shape)
qp = np.zeros(shape)
rho = np.zeros(shape)

# Define a velocity profile. The velocity is in km/s
vp_top = 1.5
vp_bottom = 3.5

# Define a velocity profile in km/s
v = np.empty(shape, dtype=dtype)
v[:] = vp_top  # Top velocity (background)
vp_i = np.linspace(vp_top, vp_bottom, nlayers)
for i in range(1, nlayers):
    v[..., i*int(shape[-1] / nlayers):] = vp_i[i]  # Bottom velocity

qp[:] = 3.516*((v[:]*1000.)**2.2)*10**(-6) # Li's empirical formula

rho[:] = 0.31*(v[:]*1000.)**0.25 # Gardner's relation
```

```python
#NBVAL_IGNORE_OUTPUT
model = ModelViscoacoustic(space_order=space_order, vp=v, qp=qp, rho=rho,
                           origin=origin, shape=shape, spacing=spacing,
                           nbl=nbl)
```

    Operator `initdamp` run in 0.01 s
    Operator `padfunc` run in 0.01 s
    Operator `padfunc` run in 0.01 s
    Operator `padfunc` run in 0.01 s

```python
#NBVAL_IGNORE_OUTPUT
aspect_ratio = model.shape[0]/model.shape[1]

plt_options_model = {'cmap': 'jet', 'extent': [model.origin[0], model.origin[0] + model.domain_size[0],
                                               model.origin[1] + model.domain_size[1], model.origin[1]]}
fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(15, 5))

slices = [slice(model.nbl, -model.nbl), slice(model.nbl, -model.nbl)]

img1 = ax[0].imshow(np.transpose(model.vp.data[slices]), vmin=1.5, vmax=3.5, **plt_options_model)
fig.colorbar(img1, ax=ax[0])
ax[0].set_title(r"V (km/s)", fontsize=20)
ax[0].set_xlabel('X (m)', fontsize=20)
ax[0].set_ylabel('Depth (m)', fontsize=20)
ax[0].set_aspect('auto')

img2 = ax[1].imshow(np.transpose(qp), vmin=15, vmax=220, **plt_options_model)
fig.colorbar(img2, ax=ax[1])
ax[1].set_title("Q", fontsize=20)
ax[1].set_xlabel('X (m)', fontsize=20)
ax[1].set_ylabel('Depth (m)', fontsize=20)
ax[1].set_aspect('auto')

img3 = ax[2].imshow(np.transpose(rho), vmin=1.9, vmax=2.4, **plt_options_model)
fig.colorbar(img3, ax=ax[2])
ax[2].set_title(r"Density $\rho$ (g/cm^3)", fontsize=20)
ax[2].set_xlabel('X (m)', fontsize=20)
ax[2].set_ylabel('Depth (m)', fontsize=20)
ax[2].set_aspect('auto')

plt.tight_layout()
```

```python
f0 = 0.005 # peak/dominant frequency
irho = model.irho
rho = 1./irho

# velocity model
vp = model.vp

t_s = (sp.sqrt(1.+1./model.qp**2)-1./model.qp)/f0
t_ep = 1./(f0**2*t_s)
tt = (t_ep/t_s) - 1.

s = model.grid.stepping_dim.spacing
damp = model.damp
```

```python
# Time step in ms and time range:
t0, tn = 0., 2000.
dt = model.critical_dt
time_range = TimeAxis(start=t0, stop=tn, step=dt)
```

```python
from examples.seismic import Receiver

def src_rec(p, model):
    src = RickerSource(name='src', grid=model.grid, f0=f0, time_range=time_range)
    src.coordinates.data[0, :] = np.array(model.domain_size) * .5
    src.coordinates.data[0, -1] = 8.

    # Create symbol for receivers
    rec = Receiver(name='rec', grid=model.grid, npoint=shape[0], time_range=time_range)

    # Prescribe even spacing for receivers along the x-axis
    rec.coordinates.data[:, 0] = np.linspace(0, model.domain_size[0], num=shape[0])
    rec.coordinates.data[:, 1] = 8.

    src_term = src.inject(field=p.forward, expr=(s*src))
    rec_term = rec.interpolate(expr=p)

    return src_term + rec_term, src, rec
```

Auxiliary functions for plotting data:

```python
def plot_receiver(rec):
    rec_plot = rec.resample(num=1001)
    scale_for_plot = np.diag(np.linspace(1.0, 2.5, 1001)**2.0)
    # Pressure (txx + tzz) data at sea surface
    extent = [rec_plot.coordinates.data[0, 0], rec_plot.coordinates.data[-1, 0], 1e-3*tn, t0]
    aspect = rec_plot.coordinates.data[-1, 0]/(1e-3*tn)/.5

    plt.figure(figsize=(10, 10))
    plt.imshow(np.dot(scale_for_plot, rec_plot.data), vmin=-.01, vmax=.01, cmap="seismic",
           interpolation='lanczos', extent=extent, aspect=aspect)
    plt.ylabel("Time (s)", fontsize=20)
    plt.xlabel("Receiver position (m)", fontsize=20)
```

```python
def plot_v_and_p(model, v, p):

    slices = [slice(model.nbl, -model.nbl), slice(model.nbl, -model.nbl)]
    scale = .5*1e-3

    plt_options_model = {'extent': [model.origin[0] , model.origin[0] + model.domain_size[0],
                                    model.origin[1] + model.domain_size[1], model.origin[1]]}

    fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(15, 7))

    ax[0].imshow(np.transpose(v[0].data[0][slices]), vmin=-scale, vmax=scale, cmap="RdGy", **plt_options_model)
    ax[0].imshow(np.transpose(model.vp.data[slices]), vmin=1.5, vmax=3.5, cmap="jet", alpha=.5, **plt_options_model)
    ax[0].set_aspect('auto')
    ax[0].set_xlabel('X (m)', fontsize=20)
    ax[0].set_ylabel('Depth (m)', fontsize=20)
    ax[0].set_title(r"$v_{x}$", fontsize=20)

    ax[1].imshow(np.transpose(v[1].data[0][slices]), vmin=-scale, vmax=scale, cmap="RdGy", **plt_options_model)
    ax[1].imshow(np.transpose(model.vp.data[slices]), vmin=1.5, vmax=3.5, cmap="jet", alpha=.5, **plt_options_model)
    ax[1].set_aspect('auto')
    ax[1].set_xlabel('X (m)', fontsize=20)
    ax[1].set_title(r"$v_{z}$", fontsize=20)

    ax[2].imshow(np.transpose(p.data[0][slices]), vmin=-scale, vmax=scale, cmap="RdGy", **plt_options_model)
    ax[2].imshow(np.transpose(model.vp.data[slices]), vmin=1.5, vmax=3.5, cmap="jet", alpha=.5, **plt_options_model)
    ax[2].set_aspect('auto')
    ax[2].set_xlabel('X (m)', fontsize=20)
    ax[2].set_title(r"$P$", fontsize=20)
```

## Blanch and Symes (1995) / Dutta and Schuster (2014)

The equations of motion for a viscoacoustic medium can be written as:

\begin{equation}
\left\{
\begin{array}{lcl}
\frac{\partial P}{\partial t} + \kappa (\tau + 1)(\nabla \cdot {\bf v}) + r = S({\bf x}_{s}, t) \\
\frac{\partial {\bf v}}{\partial t} + \frac{1}{\rho}\nabla{P} = 0 \\
\frac{\partial r}{\partial t} + \frac{1}{\tau_{\sigma}} [r + \tau \kappa (\nabla \cdot {\bf v})] = 0.
\end{array}
\right.
\label{first-order-dutta}
\end{equation}

Where $\tau = \tau_{\epsilon}/\tau_{\sigma} -1$ represents the magnitude of  $Q$. $\tau_{\epsilon}$ and $\tau_{\sigma}$ are, respectively, the stress and strain relaxation parameters, given by:

\begin{equation}
\tau_\sigma = \frac{\sqrt{Q^2+1}-1}{2 \pi f_0 Q}
\end{equation}
and
\begin{equation}
\tau_\epsilon= \frac{\sqrt{Q^2+1}+1}{2\pi f_0 Q}
\end{equation}

```python
# Stencil created from Blanch and Symes (1995) / Dutta and Schuster (2014)
def blanch_symes(model, p, r, v):

    # Bulk modulus
    bm = rho * (vp * vp)

    # Define PDE to v
    pde_v = v - s * irho * grad(p)
    u_v = Eq(v.forward, damp * pde_v)

    # Define PDE to r
    pde_r = r - s * (1. / t_s) * r - s * (1. / t_s) * tt * bm * div(v.forward)
    u_r = Eq(r.forward, damp * pde_r)

    # Define PDE to p
    pde_p = p - s * bm * (tt + 1.) * div(v.forward) - s * r.forward
    u_p = Eq(p.forward, damp * pde_p)

    return [u_v, u_r, u_p]
```

```python
# Seismic Modelling from Blanch and Symes (1995) / Dutta and Schuster (2014) viscoacoustic wave equation.
def modelling_blanch_symes(model):

    # Create symbols for particle velocity, pressure field, memory variable, source and receivers

    v = VectorTimeFunction(name="v", grid=model.grid, time_order=1, space_order=space_order)

    p = TimeFunction(name="p", grid=model.grid, time_order=1, space_order=space_order,
                     staggered=NODE)

    r = TimeFunction(name="r", grid=model.grid, time_order=1, space_order=space_order,
                     staggered=NODE)

    # define the source injection and create interpolation expression for receivers

    src_rec_expr, src, rec = src_rec(p, model)

    eqn = blanch_symes(model, p, r, v)

    op = Operator(eqn + src_rec_expr, subs=model.spacing_map)

    op(time=time_range.num-1, dt=dt, src=src, rec=rec)

    return rec, v, p
```

```python
#NBVAL_IGNORE_OUTPUT
rec, v, p = modelling_blanch_symes(model)
```

    Operator `Kernel` run in 2.45 s

```python
#NBVAL_IGNORE_OUTPUT
plot_receiver(rec)
```

```python
assert np.isclose(np.linalg.norm(rec.data), 16, rtol=10)
```

```python
#NBVAL_IGNORE_OUTPUT
plot_v_and_p(model, v, p)
```

```python
assert np.isclose(norm(v[0]), 1.8687493, atol=1e-3, rtol=0)
```

## Ren et al. (2014)

The viscoacoustic wave equation in time domain is written as:

\begin{equation}
\frac{\partial^{2}P}{\partial{t^2}} - v^{2}\nabla^{2}{P} - \eta\nabla^{2}\left(\frac{\partial P}{\partial t}\right) = S({\bf x}_{s}, t),
\end{equation}

where $\eta = \frac{v^2}{\omega_{0}Q}$ represents the viscosity of medium.

Considering the variable density $\rho$, the equation can be rewritten as:

\begin{equation}
\frac{\partial^{2}P}{\partial{t^2}} - \kappa \nabla \cdot \frac{1}{\rho} \nabla{P} - \eta \rho \nabla \cdot \frac{1}{\rho} \nabla \left(\frac{\partial{P}}{\partial{t}}\right) = S({\bf x}_{s}, t).
\end{equation}

The equation can be written using a first order formulation, given by:

\begin{equation}
\left\{
\begin{array}{ll}
\frac{\partial P}{\partial t} + \kappa \nabla \cdot {\bf v} - \eta \rho \nabla \cdot \frac{1}{\rho} \nabla{P} = S({\bf x}_{s}, t) \\
\frac{\partial {\bf v}}{\partial t} + \frac{1}{\rho} \nabla{P} = 0
\end{array}
\right.
\end{equation}

```python
# Stencil created from Ren et al. (2014) viscoacoustic wave equation.
def ren(model, p, v):

    # Angular frequency
    w = 2. * np.pi * f0

    # Define PDE to v
    pde_v = v - s * irho * grad(p)
    u_v = Eq(v.forward, damp * pde_v)

    # Define PDE to p
    pde_p = p - s * vp * vp * rho * div(v.forward) + s * \
            ((vp * vp * rho) / (w * model.qp)) * div(irho * grad(p))

    u_p = Eq(p.forward, damp * pde_p)

    return [u_v, u_p]
```

```python
# Seismic Modelling from Ren et al. (2014) viscoacoustic wave equation.
def modelling_ren(model):

    # Create symbols for particle velocity, pressure field, source and receivers

    v = VectorTimeFunction(name="v", grid=model.grid, time_order=1, space_order=space_order)

    p = TimeFunction(name="p", grid=model.grid, time_order=1, space_order=space_order,
                     staggered=NODE)

    # define the source injection and create interpolation expression for receivers

    src_rec_expr, src, rec = src_rec(p, model)

    eqn = ren(model, p, v)

    op = Operator(eqn + src_rec_expr, subs=model.spacing_map)

    op(time=time_range.num-1, dt=dt, src=src, rec=rec)

    return rec, v, p
```

```python
#NBVAL_IGNORE_OUTPUT
rec, v, p = modelling_ren(model)
```

    Operator `Kernel` run in 3.49 s

```python
#NBVAL_IGNORE_OUTPUT
plot_receiver(rec)
```

```python
assert np.isclose(np.linalg.norm(rec.data), 15, rtol=10)
```

```python
#NBVAL_IGNORE_OUTPUT
plot_v_and_p(model, v, p)
```

```python
assert np.isclose(norm(v[0]), 1.0587339, atol=1e-3, rtol=0)
```

## Deng and McMechan (2007)

The viscoacoustic wave equation for the propagating pressure $P$ in the time-space domain:

\begin{equation}
\frac{1}{v^2}\frac{\partial^{2}P}{\partial{t^2}} - \nabla^{2}P + \frac{g}{v}\frac{\partial P}{\partial{t}}  = S({\bf x}_{s}, t),
\label{eq-deng}
\end{equation}

where $g$ is the absorption coefficient, given by:

\begin{equation}
g = \frac{2\pi f_{0}}{vQ},
\end{equation}

The equation can be written using a first order formulation, given by:

\begin{equation}
\left\{
\begin{array}{lcl}
\frac{\partial P}{\partial t} + \kappa (\nabla \cdot {\bf v}) + \frac{2\pi f_{0}}{Q}P= S({\bf x}_{s}, t) \\
\frac{\partial {\bf v}}{\partial t} + \frac{1}{\rho}\nabla{P} = 0 \\
\end{array}
\right.
\end{equation}

```python
# Stencil created from Deng and McMechan (2007) viscoacoustic wave equation.
def deng_mcmechan(model, p, v):

    # Angular frequency
    w = 2. * np.pi * f0

    # Define PDE to v
    pde_v = v - s * irho * grad(p)
    u_v = Eq(v.forward, damp * pde_v)

    # Define PDE to p
    pde_p = p - s * vp * vp * rho * div(v.forward) - s * (w / model.qp) * p
    u_p = Eq(p.forward, damp * pde_p)

    return [u_v, u_p]
```

```python
# Seismic Modelling from Deng and McMechan (2007) viscoacoustic wave equation.
def modelling_deng_mcmechan(model):

    # Create symbols for particle velocity, pressure field, source and receivers

    v = VectorTimeFunction(name="v", grid=model.grid, time_order=1, space_order=space_order)

    p = TimeFunction(name="p", grid=model.grid, time_order=1, space_order=space_order,
                     staggered=NODE)

    # define the source injection and create interpolation expression for receivers

    src_rec_expr, src, rec = src_rec(p, model)

    eqn = deng_mcmechan(model, p, v)

    op = Operator(eqn + src_rec_expr, subs=model.spacing_map)

    op(time=time_range.num-1, dt=dt, src=src, rec=rec)

    return rec, v, p
```

```python
#NBVAL_IGNORE_OUTPUT
rec, v, p = modelling_deng_mcmechan(model)
```

    Operator `Kernel` run in 1.78 s

```python
#NBVAL_IGNORE_OUTPUT
plot_receiver(rec)
```

```python
assert np.isclose(np.linalg.norm(rec.data), 16, rtol=10)
```

```python
#NBVAL_IGNORE_OUTPUT
plot_v_and_p(model, v, p)
```

```python
assert np.isclose(norm(v[0]), 1.126901, atol=1e-3, rtol=0)
```

# More references

[1] https://academic.oup.com/gji/article/197/2/948/616510

[2] https://link.springer.com/article/10.1007/s11770-014-0414-8

[3] https://janth.home.xs4all.nl/Software/fdelmodcManual.pdf


###### Content under Creative Commons Attribution license CC-BY 4.0, code under MIT license (c) 2019 Daniel Koehn, based on (c)2018 L.A. Barba, G.F. Forsyth [CFD Python](https://github.com/barbagroup/CFDPython#cfd-python), (c)2014 L.A. Barba, I. Hawke, B. Knaepen [Practical Numerical Methods with Python](https://github.com/numerical-mooc/numerical-mooc#practical-numerical-methods-with-python), also under CC-BY.

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

# 2D Laplace Equation

So far, you've learned to solve problems dominated by advection—where solutions have a directional bias and can form shocks—in [Lecture 4](https://github.com/daniel-koehn/Differential-equations-earth-system/tree/master/04_Advection_1D\#differential-equations-in-earth-sciences-1d-nonlinear-advection). In [Lecture 5](https://github.com/daniel-koehn/Differential-equations-earth-system/tree/master/05_Diffusion_1D\#differential-equations-in-earth-sciences-1d-diffusion) we explored diffusion-dominated problems—where solutions spread in all directions. But what about situations where solutions are steady?

Many problems in physics have no time dependence, yet are rich with physical meaning: the gravitational field produced by a massive object, the electrostatic potential of a charge distribution, the displacement of a stretched membrane and the steady flow of fluid through a porous medium ... all these can be modeled by **Poisson's equation**:

$$
\begin{equation}
\nabla^2 u = f \notag
\end{equation}
$$

where the unknown $u$ and the known $f$ are functions of space, in a domain $\Omega$. To find the solution, we require boundary conditions. These could be  Dirichlet boundary conditions, specifying the value of the solution on the boundary,

$$
\begin{equation}
u = b_1 \text{ on } \partial\Omega, \notag
\end{equation}
$$

or Neumann boundary conditions, specifying the normal derivative of the solution on the boundary,

$$
\begin{equation}
\frac{\partial u}{\partial n} = b_2 \text{ on } \partial\Omega. \notag
\end{equation}
$$

A boundary-value problem consists of finding $u$, given the above information. Numerically, we can do this using *relaxation methods*, which start with an initial guess for $u$ and then iterate towards the solution. Let's find out how!

## Laplace's equation

The particular case of $f=0$ (homogeneous case) results in Laplace's equation:

$$
\begin{equation}
\nabla^2 u = 0 \notag
\end{equation}
$$

For example, the equation for steady, two-dimensional heat conduction is:

$$
\begin{equation}
\frac{\partial ^2 T}{\partial x^2} + \frac{\partial ^2 T}{\partial y^2} = 0 \notag
\end{equation}
$$

This is similar to the model we studied in [lesson 1](https://github.com/daniel-koehn/Differential-equations-earth-system/tree/master/06_Heat_2D\#differential-equations-in-earth-sciences-2d-heat-equation) of **Lecture 6**, but without the time derivative: i.e., for a temperature $T$ that has reached steady state. The Laplace equation models the equilibrium state of a system under the supplied boundary conditions.

The study of solutions to Laplace's equation is called *potential theory*, and the solutions themselves are often potential fields. Let's use $p$ from now on to represent our generic dependent variable, and write Laplace's equation again (in two dimensions):

$$
\begin{equation}
\frac{\partial ^2 p}{\partial x^2} + \frac{\partial ^2 p}{\partial y^2} = 0 \notag
\end{equation}
$$

Like in the diffusion equation of the previous course module, we discretize the second-order derivatives with *central differences*. You should be able to write down a second-order central-difference formula by heart now! On a two-dimensional Cartesian grid, it gives:

$$
\begin{equation}
\frac{p_{i+1, j}  - 2p_{i,j}  + p_{i-1,j} }{\Delta x^2} + \frac{p_{i,j+1}  - 2p_{i,j}  + p_{i, j-1} }{\Delta y^2} = 0 \notag
\end{equation}
$$

When $\Delta x = \Delta y$, we end up with the following equation:

$$
\begin{equation}
p_{i+1, j}   + p_{i-1,j} + p_{i,j+1}  + p_{i, j-1}- 4 p_{i,j} = 0 \notag
\end{equation}
$$

This tells us that the Laplacian differential operator at grid point $(i,j)$ can be evaluated discretely using the value of $p$ at that point (with a factor $-4$) and the four neighboring points to the left and right, above and below grid point $(i,j)$.

The stencil of the discrete Laplacian operator is shown in Figure 1. It is typically called the *five-point stencil*, for obvious reasons.

#### Figure 1: Laplace five-point stencil.

The discrete equation above is valid for every interior point in the domain. If we write the equations for *all* interior points, we have a linear system of algebraic equations. We *could* solve the linear system directly (e.g., with Gaussian elimination), but we can be more clever than that!

Notice that the coefficient matrix of such a linear system has mostly zeroes. For a uniform spatial grid, the matrix is *block diagonal*: it has diagonal blocks that are tridiagonal with $-4$ on the main diagonal and $1$ on two off-center diagonals, and two more diagonals with $1$. All of the other elements are zero. Iterative methods are particularly suited for a system with this structure, and save us from storing all those zeroes.

We will start with an initial guess for the solution, $p_{i,j}^{0}$, and use the discrete Laplacian to get an update, $p_{i,j}^{1}$, then continue on computing $p_{i,j}^{k}$ until we're happy.  Note that $k$ is _not_ a time index here, but an index corresponding to the number of iterations we perform in the *relaxation scheme*.

At each iteration, we compute updated values $p_{i,j}^{k+1}$ in a (hopefully) clever way so that they converge to a set of values satisfying Laplace's equation. The system will reach equilibrium only as the number of iterations tends to $\infty$, but we can approximate the equilibrium state by iterating until the change between one iteration and the next is *very* small.

The most intuitive method of iterative solution is known as the [**Jacobi method**](https://en.wikipedia.org/wiki/Jacobi_method), in which the values at the grid points are replaced by the corresponding weighted averages:

$$
\begin{equation}
p^{k+1}_{i,j} = \frac{1}{4} \left(p^{k}_{i,j-1} + p^k_{i,j+1} + p^{k}_{i-1,j} + p^k_{i+1,j} \right) \tag{1}
\end{equation}
$$

This method does indeed converge to the solution of Laplace's equation. Thank you Professor Jacobi!

##### Challenge task

Grab a piece of paper and write out the coefficient matrix for a discretization with 7 grid points in the $x$ direction (5 interior points) and 5 points in the $y$ direction (3 interior). The system should have 15 unknowns, and the coefficient matrix three diagonal blocks. Assume prescribed Dirichlet boundary conditions on all sides (not necessarily zero).

### Boundary conditions and relaxation

Suppose we want to model steady-state heat transfer on (say) a computer chip with one side insulated (zero Neumann BC), two sides held at a fixed temperature (Dirichlet condition) and one side touching a component that has a sinusoidal distribution of temperature.
We would need to solve Laplace's equation with boundary conditions like

$$
\begin{equation}
  \begin{gathered}
p=0 \text{ at } x=0\\
\frac{\partial p}{\partial x} = 0 \text{ at } x = L_x\\
p = 0 \text{ at }y = 0 \\
p = \sin \left(  \frac{\frac{3}{2}\pi x}{L_x} \right) \text{ at } y = L_y
  \end{gathered} \notag
\end{equation}
$$

We'll take $L_x=1$ and $L_y=1$ for the sizes of the domain in the $x$ and $y$ directions.

One of the defining features of elliptic PDEs is that they are "driven" by the boundary conditions.  In the iterative solution of Laplace's equation, boundary conditions are set and **the solution relaxes** from an initial guess to join the boundaries together smoothly, given those conditions.  Our initial guess will be $p=0$ everywhere. Now, let's relax!

First, we import our usual smattering of libraries (plus a few new ones!)

```python
import numpy
from matplotlib import pyplot
%matplotlib inline

# import JIT from Numba
from numba import jit
```

```python
# Set the font family and size to use for Matplotlib figures.
pyplot.rcParams['font.family'] = 'serif'
pyplot.rcParams['font.size'] = 16
```

To visualize 2D data, we can use [`pyplot.imshow()`](http://matplotlib.org/api/pyplot_api.html#matplotlib.pyplot.imshow), like we've done before, but a 3D plot can sometimes show a more intuitive view the solution. Or it's just prettier!

Be sure to enjoy the many examples of 3D plots in the `mplot3d` section of the [Matplotlib Gallery](http://matplotlib.org/gallery.html#mplot3d).

We'll import the `mplot3d` module to create 3D plots and also grab the `cm` package, which provides different colormaps for visualizing plots.

```python
from mpl_toolkits import mplot3d
from matplotlib import cm
```

Let's define a function for setting up our plotting environment, to avoid repeating this set-up over and over again. It will save us some typing.

```python
def plot_3d(x, y, p, label='$z$', elev=30.0, azim=45.0):
    """
    Creates a Matplotlib figure with a 3D surface plot
    of the scalar field p.

    Parameters
    ----------
    x : numpy.ndarray
        Gridline locations in the x direction as a 1D array of floats.
    y : numpy.ndarray
        Gridline locations in the y direction as a 1D array of floats.
    p : numpy.ndarray
        Scalar field to plot as a 2D array of floats.
    label : string, optional
        Axis label to use in the third direction;
        default: 'z'.
    elev : float, optional
        Elevation angle in the z plane;
        default: 30.0.
    azim : float, optional
        Azimuth angle in the x,y plane;
        default: 45.0.
    """
    fig = pyplot.figure(figsize=(8.0, 6.0))
    ax = mplot3d.Axes3D(fig)
    ax.set_xlabel('$x$')
    ax.set_ylabel('$y$')
    ax.set_zlabel(label)
    X, Y = numpy.meshgrid(x, y)
    ax.plot_surface(X, Y, p, cmap=cm.viridis)
    ax.set_xlim(x[0], x[-1])
    ax.set_ylim(y[0], y[-1])
    ax.view_init(elev=elev, azim=azim)
```

##### Note
This plotting function uses *Viridis*, a new (and _awesome_) colormap available in Matplotlib versions 1.5 and greater.  If you see an error when you try to plot using `cm.viridis`, just update Matplotlib using `conda` or `pip`.

### Analytical solution

The Laplace equation with the boundary conditions listed above has an analytical solution, given by

$$
\begin{equation}
p(x,y) = \frac{\sinh \left( \frac{\frac{3}{2} \pi y}{L_y}\right)}{\sinh \left(  \frac{\frac{3}{2} \pi L_y}{L_x}\right)} \sin \left( \frac{\frac{3}{2} \pi x}{L_x} \right) \notag
\end{equation}
$$

where $L_x$ and $L_y$ are the length of the domain in the $x$ and $y$ directions, respectively.

We previously used `numpy.meshgrid` to plot our 2D solutions to the heat equation in Module 4.  Here, we'll use it again as a plotting aid. Always useful, `linspace` creates 1-row arrays of equally spaced numbers: it helps for defining $x$ and $y$ axes in line plots, but now we want the analytical solution plotted for every point in our domain.  To do this, we'll use in the analytical solution the 2D arrays generated by `numpy.meshgrid`.

```python
def laplace_solution(x, y, Lx, Ly):
    """
    Computes and returns the analytical solution of the Laplace equation
    on a given two-dimensional Cartesian grid.

    Parameters
    ----------
    x : numpy.ndarray
        The gridline locations in the x direction
        as a 1D array of floats.
    y : numpy.ndarray
        The gridline locations in the y direction
        as a 1D array of floats.
    Lx : float
        Length of the domain in the x direction.
    Ly : float
        Length of the domain in the y direction.

    Returns
    -------
    p : numpy.ndarray
        The analytical solution as a 2D array of floats.
    """
    X, Y = numpy.meshgrid(x, y)
    p = (numpy.sinh(1.5 * numpy.pi * Y / Ly) /
         numpy.sinh(1.5 * numpy.pi * Ly / Lx) *
         numpy.sin(1.5 * numpy.pi * X / Lx))
    return p
```

Ok, let's try out the analytical solution and use it to test the `plot_3D` function we wrote above.

```python
# Set parameters.
Lx = 1.0  # domain length in the x direction
Ly = 1.0  # domain length in the y direction
nx = 41  # number of points in the x direction
ny = 41  # number of points in the y direction

# Create the gridline locations.
x = numpy.linspace(0.0, Lx, num=nx)
y = numpy.linspace(0.0, Ly, num=ny)

# Compute the analytical solution.
p_exact = laplace_solution(x, y, Lx, Ly)

# Plot the analytical solution.
plot_3d(x, y, p_exact)
```

It worked!  This is what the solution *should* look like when we're 'done' relaxing. (And isn't viridis a cool colormap?)

### How long do we iterate?

We noted above that there is no time dependence in the Laplace equation.  So it doesn't make a lot of sense to use a `for` loop with `nt` iterations, like we've done before.

Instead, we can use a `while` loop that continues to iteratively apply the relaxation scheme until the difference between two successive iterations is small enough.

But how small is small enough?  That's a good question.  We'll try to work that out as we go along.

To compare two successive potential fields, a good option is to use the [L2 norm](http://en.wikipedia.org/wiki/Norm_%28mathematics%29#Euclidean_norm) of the difference.  It's defined as

$$
\begin{equation}
|\textbf{x}| = \sqrt{\sum_{i=0, j=0}^k \left|p^{k+1}_{i,j} - p^k_{i,j}\right|^2} \notag
\end{equation}
$$

But there's one flaw with this formula.  We are summing the difference between successive iterations at each point on the grid. So what happens when the grid grows? (For example, if we're refining the grid, for whatever reason.) There will be more grid points to compare and so more contributions to the sum. The norm will be a larger number just because of the grid size!

That doesn't seem right.  We'll fix it by normalizing the norm, dividing the above formula by the norm of the potential field at iteration $k$.

For two successive iterations, the relative L2 norm is then calculated as

$$
\begin{equation}
|\textbf{x}| = \frac{\sqrt{\sum_{i=0, j=0}^k \left|p^{k+1}_{i,j} - p^k_{i,j}\right|^2}}{\sqrt{\sum_{i=0, j=0}^k \left|p^k_{i,j}\right|^2}}\notag
\end{equation}
$$

For this purpose, we define the `l2_norm` function:

```python
def l2_norm(p, p_ref):
    """
    Computes and returns the relative L2-norm of the difference
    between a solution p and a reference solution p_ref.

    Parameters
    ----------
    p : numpy.ndarray
        The solution as an array of floats.
    p_ref : numpy.ndarray
        The reference solution as an array of floats.

    Returns
    -------
    diff : float
        The relative L2-norm of the difference.
    """
    l2_diff = (numpy.sqrt(numpy.sum((p - p_ref)**2)) /
               numpy.sqrt(numpy.sum(p_ref**2)))
    return l2_diff
```

Now, let's define two functions that will apply Jacobi's method for Laplace's equation.
First, we put the spatial finite difference approximation of the Laplace equation (1) into a function `laplace` and apply Just-In-Time (JIT) compilation from the Numba package.

```python
@jit(nopython=True) # use Just-In-Time (JIT) Compilation for C-performance
def laplace(p,pn,nx,ny):

    for i in range(1,nx-1):
        for j in range(1,ny-1):

            # Solve Laplace eq. (1) using Jacobi method
            p[j, i] = 0.25 * (pn[j+1, i] + pn[j, i+1] +
                              pn[j-1, i] + pn[j, i-1])

    return p
```

Finally, we code the solution of the Laplace equation using the Jacobi method in a function `laplace_2d_jacobi`. Three of the boundaries are Dirichlet boundaries and so we can simply leave them alone. Only the Neumann boundary needs to be explicitly calculated at each iteration.

```python
def laplace_2d_jacobi(p0, maxiter=20000, rtol=1e-6):
    """
    Solves the 2D Laplace equation using Jacobi relaxation method.

    The function assumes Dirichlet condition with value zero
    at all boundaries except at the right boundary where it uses
    a zero-gradient Neumann condition.
    The exit criterion of the solver is based on the relative L2-norm
    of the solution difference between two consecutive iterations.

    Parameters
    ----------
    p0 : numpy.ndarray
        The initial solution as a 2D array of floats.
    maxiter : integer, optional
        Maximum number of iterations to perform;
        default: 20000.
    rtol : float, optional
        Relative tolerance for convergence;
        default: 1e-6.

    Returns
    -------
    p : numpy.ndarray
        The solution after relaxation as a 2D array of floats.
    ite : integer
        The number of iterations performed.
    diff : float
        The final relative L2-norm of the difference.
    """

    # estimate nx, ny
    ny, nx = numpy.shape(p0)

    # copy initial guess p0
    p = p0.copy()

    diff = rtol + 1.0  # initial difference
    ite = 0  # iteration index

    # Jacobi iteration
    while diff > rtol and ite < maxiter:

        # store old p-values
        pn = p.copy()

        # Update the solution at interior points of the Laplace equation
        p = laplace(p,pn,nx,ny)

        # Apply Neumann condition (zero-gradient)
        # at the right boundary.
        p[1:-1, -1] = p[1:-1, -2]

        # Compute the residual as the L2-norm of the difference.
        diff = l2_norm(p, pn)
        ite += 1

    return p, ite, diff
```

##### Rows and columns, and index order

Recall that in the [2D explicit heat equation](https://nbviewer.jupyter.org/github/daniel-koehn/Differential-equations-earth-system/blob/master/06_Heat_2D/01_Heat_Equation_2D.ipynb) we stored data with the $y$ coordinates corresponding to the rows of the array and $x$ coordinates on the columns (this is just a code design decision!). We did that so that a plot of the 2D-array values would have the natural ordering, corresponding to the physical domain ($y$ coordinate in the vertical).

We'll follow the same convention here (even though we'll be plotting in 3D, so there's no real reason), just to be consistent. Thus, $p_{i,j}$ will be stored in array format as `p[j,i]`. Don't be confused by this.

### Let's relax!

The initial values of the potential field are zero everywhere (initial guess), except at the top boundary:

$$
p = \sin \left(  \frac{\frac{3}{2}\pi x}{L_x} \right) \text{ at } y=L_y
$$

To initialize the domain, `numpy.zeros` will handle everything except that one Dirichlet condition. Let's do it!

```python
# Set the initial conditions.
p0 = numpy.zeros((ny, nx))
p0[-1, :] = numpy.sin(1.5 * numpy.pi * x / Lx)
```

Now let's visualize the initial conditions using the `plot_3D` function, just to check we've got it right.

```python
# Plot the initial conditions.
plot_3d(x, y, p0)
```

The `p` array is equal to zero everywhere, except along the boundary $y = 1$.  Hopefully you can see how the relaxed solution and this initial condition are related.

Now, run the iterative solver with a target L2-norm difference between successive iterations of $10^{-8}$.

```python
# Compute the solution using Jacobi relaxation method.
p, ites, diff = laplace_2d_jacobi(p0, rtol=1e-8)
print('Jacobi relaxation: {} iterations '.format(ites) +
      'to reach a relative difference of {}'.format(diff))
```

    Jacobi relaxation: 4473 iterations to reach a relative difference of 9.989253686324196e-09

Let's make a gorgeous plot of the final field using the newly minted `plot_3d` function.

```python
# Plot the numerical solution.
plot_3d(x, y, p)
```

Awesome!  That looks pretty good.  But we'll need more than a simple visual check, though. The "eyeball metric" is very forgiving!

## Convergence analysis

### Convergence, Take 1

We want to make sure that our Jacobi function is working properly.  Since we have an analytical solution, what better way than to do a grid-convergence analysis?  We will run our solver for several grid sizes and look at how fast the L2 norm of the difference between consecutive iterations decreases.

Now run Jacobi's method on the Laplace equation using four different grids, with the same exit criterion of $10^{-8}$ each time. Then, we look at the error versus the grid size in a log-log plot. What do we get?

```python
# List of the grid sizes to investigate.
nx_values = [11, 21, 41, 81]

# Create an empty list to record the error on each grid.
errors = []

# Compute the solution and error for each grid size.
for nx in nx_values:
    ny = nx  # same number of points in all directions.
    # Create the gridline locations.
    x = numpy.linspace(0.0, Lx, num=nx)
    y = numpy.linspace(0.0, Ly, num=ny)
    # Set the initial conditions.
    p0 = numpy.zeros((ny, nx))
    p0[-1, :] = numpy.sin(1.5 * numpy.pi * x / Lx)
    # Relax the solution.
    # We do not return the number of iterations or
    # the final relative L2-norm of the difference.
    p, _, _ = laplace_2d_jacobi(p0, rtol=1e-8)
    # Compute the analytical solution.
    p_exact = laplace_solution(x, y, Lx, Ly)
    # Compute and record the relative L2-norm of the error.
    errors.append(l2_norm(p, p_exact))
```

```python
# Plot the error versus the grid-spacing size.
pyplot.figure(figsize=(6.0, 6.0))
pyplot.xlabel(r'$\Delta x$')
pyplot.ylabel('Relative $L_2$-norm\nof the error')
pyplot.grid()
dx_values = Lx / (numpy.array(nx_values) - 1)
pyplot.loglog(dx_values, errors,
              color='black', linestyle='--', linewidth=2, marker='o')
pyplot.axis('equal');
```

Hmm. That doesn't look like 2nd-order convergence, but we're using second-order finite differences.  *What's going on?*  The culprit is the boundary conditions. Dirichlet conditions are order-agnostic (a  set value is a set value), but the scheme we used for the Neumann boundary condition is 1st-order.

Remember when we said that the boundaries drive the problem?  One boundary that's 1st-order completely tanked our spatial convergence.  Let's fix it!

### 2nd-order Neumann BCs

Up to this point, we have used the first-order approximation of a derivative to satisfy Neumann B.C.'s. For a boundary located at $x=0$ this reads,

$$
\begin{equation}
\frac{p^{k+1}_{1,j} - p^{k+1}_{0,j}}{\Delta x} = 0
\end{equation}
$$

which, solving for $p^{k+1}_{0,j}$ gives us

$$
\begin{equation}
p^{k+1}_{0,j} = p^{k+1}_{1,j}
\end{equation}
$$

Using that Neumann condition will limit us to 1st-order convergence.  Instead, we can start with a 2nd-order approximation (the central-difference approximation):

$$
\begin{equation}
\frac{p^{k+1}_{1,j} - p^{k+1}_{-1,j}}{2 \Delta x} = 0
\end{equation}
$$

That seems problematic, since there is no grid point $p^{k}_{-1,j}$.  But no matter … let's carry on. According to the 2nd-order approximation,

$$
\begin{equation}
p^{k+1}_{-1,j} = p^{k+1}_{1,j}
\end{equation}
$$

Recall the finite-difference Jacobi equation with $i=0$:

$$
\begin{equation}
p^{k+1}_{0,j} = \frac{1}{4} \left(p^{k}_{0,j-1} + p^k_{0,j+1} + p^{k}_{-1,j} + p^k_{1,j} \right)
\end{equation}
$$

Notice that the equation relies on the troublesome (nonexistent) point $p^k_{-1,j}$, but according to the equality just above, we have a value we can substitute, namely $p^k_{1,j}$. Ah! We've completed the 2nd-order Neumann condition:

$$
\begin{equation}
p^{k+1}_{0,j} = \frac{1}{4} \left(p^{k}_{0,j-1} + p^k_{0,j+1} + 2p^{k}_{1,j} \right)
\end{equation}
$$

That's a bit more complicated than the first-order version, but it's relatively straightforward to code.

##### Note

Do not confuse $p^{k+1}_{-1,j}$ with `p[-1]`:
`p[-1]` is a piece of Python code used to refer to the last element of a list or array named `p`.  $p^{k+1}_{-1,j}$ is a 'ghost' point that describes a position that lies outside the actual domain.

### Convergence, Take 2

We can copy the previous Jacobi function and replace only the line implementing the Neumann boundary condition.

##### Careful!
Remember that our problem has the Neumann boundary located at $x = L$ and not $x = 0$ as we assumed in the derivation above.

```python
def laplace_2d_jacobi_neumann(p0, maxiter=20000, rtol=1e-6):
    """
    Solves the 2D Laplace equation using Jacobi relaxation method.

    The function assumes Dirichlet condition with value zero
    at all boundaries except at the right boundary where it uses
    a zero-gradient second-order Neumann condition.
    The exit criterion of the solver is based on the relative L2-norm
    of the solution difference between two consecutive iterations.

    Parameters
    ----------
    p0 : numpy.ndarray
        The initial solution as a 2D array of floats.
    maxiter : integer, optional
        Maximum number of iterations to perform;
        default: 20000.
    rtol : float, optional
        Relative tolerance for convergence;
        default: 1e-6.

    Returns
    -------
    p : numpy.ndarray
        The solution after relaxation as a 2D array of floats.
    ite : integer
        The number of iterations performed.
    diff : float
        The final relative L2-norm of the difference.
    """

    # estimate nx, ny
    ny, nx = numpy.shape(p0)

    # copy initial guess p0
    p = p0.copy()

    diff = rtol + 1.0  # intial difference
    ite = 0  # iteration index

    while diff > rtol and ite < maxiter:

        # store old value p
        pn = p.copy()

        # Update the solution at interior points of the Laplace equation
        p = laplace(p,pn,nx,ny)

        # Apply 2nd-order Neumann condition (zero-gradient)
        # at the right boundary.
        p[1:-1, -1] = 0.25 * (2.0 * pn[1:-1, -2] +
                              pn[2:, -1] + pn[:-2, -1])
        # Compute the residual as the L2-norm of the difference.
        diff = l2_norm(p, pn)
        ite += 1
    return p, ite, diff
```

Again, this is the exact same code as before, but now we're running the Jacobi solver with a 2nd-order Neumann boundary condition. Let's do a grid-refinement analysis, and plot the error versus the grid spacing.

```python
# List of the grid sizes to investigate.
nx_values = [11, 21, 41, 81]

# Create an empty list to record the error on each grid.
errors = []

# Compute the solution and error for each grid size.
for nx in nx_values:
    ny = nx  # same number of points in all directions.
    # Create the gridline locations.
    x = numpy.linspace(0.0, Lx, num=nx)
    y = numpy.linspace(0.0, Ly, num=ny)
    # Set the initial conditions.
    p0 = numpy.zeros((ny, nx))
    p0[-1, :] = numpy.sin(1.5 * numpy.pi * x / Lx)
    # Relax the solution.
    # We do not return the number of iterations or
    # the final relative L2-norm of the difference.
    p, _, _ = laplace_2d_jacobi_neumann(p0, rtol=1e-8)
    # Compute the analytical solution.
    p_exact = laplace_solution(x, y, Lx, Ly)
    # Compute and record the relative L2-norm of the error.
    errors.append(l2_norm(p, p_exact))
```

```python
# Plot the error versus the grid-spacing size.
pyplot.figure(figsize=(6.0, 6.0))
pyplot.xlabel(r'$\Delta x$')
pyplot.ylabel('Relative $L_2$-norm\nof the error')
pyplot.grid()
dx_values = Lx / (numpy.array(nx_values) - 1)
pyplot.loglog(dx_values, errors,
              color='black', linestyle='--', linewidth=2, marker='o')
pyplot.axis('equal');
```

Nice!  That's much better.  It might not be *exactly* 2nd-order, but it's awfully close. (What is ["close enough"](http://ianhawke.github.io/blog/close-enough.html) in regards to observed convergence rates is a thorny question.)

Now, notice from this plot that the error on the finest grid is around $0.0002$. Given this, perhaps we didn't need to continue iterating until a target difference between two solutions of $10^{-8}$. The spatial accuracy of the finite difference approximation is much worse than that! But we didn't know it ahead of time, did we? That's the "catch 22" of iterative solution of systems arising from discretization of PDEs.

## Final word

The Jacobi method is the simplest relaxation scheme to explain and to apply. It is also the *worst* iterative solver! In practice, it is seldom used on its own as a solver, although it is useful as a smoother with multi-grid methods. As we will see in the next [lesson](https://nbviewer.jupyter.org/github/daniel-koehn/Differential-equations-earth-system/blob/master/07_BVP_2D/03_Iterative_Solvers.ipynb) of this module, there are much better iterative methods! But first, let's play with [Poisson's equation](http://nbviewer.ipython.org/urls/github.com/daniel-koehn/Differential-equations-earth-system/tree/master/07_BVP_2D/02_2D_Poisson_Equation.ipynb).


## Equivariance

input f(x): (F = R^n => S) (can be a graph, use R^n for generalisability)
fully connected layer w(x, x+): R^n x R^n => S
morphing operator: A {f(x)} = AL(f(AR(x))): F => F

output g(x+): R^n => S = <f(x), w(x, x+)> dx

equivariant condition: <A {f(x)}, w(x, x+)> dx = Q {<f(x), w(x, x+)> dx}    \forall f(x)

assuming only AR is in consideration (e.g. a translation):

<f(AR x), w(x, x+)> dx = <f(x), w(x, QR x+))> dx

<f(x), w(AR^-1 x, x+)> dx = <f(x), w(x, QR x+))> dx     \forall f(x)

<f(x), w(AR^-1 x, QR^-1 y+)> dx = <f(x), w(x, y+))> dx     \forall f(x)

w(AR^-1 x, QR^-1 y+) = w(x, y+)

w(AR x, QR y+) = w(x, y+)

as a result:

assuming that x+ = H x, v(x) = w(H x, QH H x) = w(x, H x), then:

g(x+) = g(H x) = <f(x), w(x, H x)> dH =

assuming that x+ = Q x0, then

\begin{align}
g(x+) &= g(Q x0) = <f(y), w(y, Q x0)> dy \\
    &= <f(y), w(P^{-1} y, x0)> dy \\
    &= conv(f(y), w(y, x0))(P^{-1})
\end{align}

QED

## Data Augmentation - Similarity

Hypothesis 2:

$$
<f_1(x), f_2(x)> dx = k <A_{ug} \circ f_1(x), A_{ug} \circ f_2(x)> dx
$$

OR: inner-product preserving (isometry) + scaling
    (I know it sounds even worse, please don't judge me)

## Data Augmentation - Put It All Together

- linear fully-connected dense

\begin{align}
& & U_{ga} \circ f_+(y) = U_{ga}|_y \circ <f(x), w(x, y)> dx &= <A_{ug} \circ f(x), w(x, y)> dx \\
\text{(similarity)} \\
& \Longrightarrow & <f(x), U_{ga}|_y \circ w(x, y)> dx &= \frac{1}{k} <A_{ug}^{-1} \circ A_{ug} \circ f(x), A_{ug}^{-1}|_x \circ w(x, y)> dx \\
& & &= \frac{1}{k} <f(x), A_{ug}^{-1}|_x \circ w(x, y)> dx \\
\forall f(x) \\
& \Longrightarrow & U_{ga}|_y \circ w(x, y) &= \frac{1}{k} A_{ug}^{-1}|_x \circ w(x, y) \\
& & & \text{(lemma of weight sharing!)}\\
\end{align}

- now assuming that there is a reference point $y_0$
    - ... such that for any $y$ there is a data augmentation $\bar{U}_{ga}$ that satisfy **transitivity**:

    $f_+(y) = \bar{U}_{ga} \circ f_+(y_0)$

    $w(x, y) = \bar{U}_{ga}|_y \circ w(x, y_0)$

\begin{align}
& & f_+(y) &= <f(x), w(x, y)> dx = <f(x), \bar{U}_{ga}|_y \circ w(x, y_0)> dx \\
& \text{(lemma of weight sharing)} & &= <f(x), \bar{A}_{ug}^{-1}|_x \circ w(x, y_0)> dx \\
& \text{(similarity)} & &= <\bar{A}_{ug} \circ f(x), w(x, z_0)> dx = <\bar{A}_{ug} \circ f(x), w_0(x)> dx
\end{align}

- looks familiar yet?

$$
conv(f(- \Delta), w_0(\Delta)) = corr(f(\Delta), w_0(\Delta)) = \int_{\text{Manifold}} f(\Delta + x) w_0(x) d x = <f(\Delta + x), w_0(x)> d x
$$

```python
import numpy as np

left = np.zeros((3, 2))
right = np.zeros((2,4))

np.matmul(left, right)
```

    array([[0., 0., 0., 0.],
           [0., 0., 0., 0.],
           [0., 0., 0., 0.]])

```python
## Glossary

- A monoid is anything compatible with MapReduce

```python
apple.merge(pen).merge(pineapple).merge(pen).merge(face).merge(palm) =
apple.merge(pen)                    # core 1
  .merge(pineapple.merge(pen))      # core 2
  .merge(face.merge(palm))          # core 3
```
```

```python
## Glossary

- A group member is an invertible monoid

```python
pineapplepen.merge(-pen) = pineapple
```

A commutative (Abelian) monoid/group member is an monoid i which order of reduce doesn't matter

```python
pineapple.merge(pen) = pen.merge(pineapple)
```
```