# Iterative Solvers 4 - Preconditioning

## The basic idea

For both the GMRES method and CG we have seen that the eigenvalue distribution is crucial for fast convergence. In both cases we would like the eigenvalues of the matrix be clustered close together and be well separated from zero. Unfortunately, in many applications the matrices that arise naturally are badly behaved in this respect. Preconditioning is a strategy to try to modify a linear system so that it becomes more ameniable to iterative solvers. Preconditioning is one of the most active research areas and crucial for the solution of linear systems of equations with millions or even billions of unknowns.

Consider the linear system of equations

$$
Ax = b
$$

In its basic form the idea of preconditioning is to multiply the system with a matrix $P^{-1}$ that is some kind of approximation to $A^{-1}$, that is $P^{-1}\approx A^{-1}$ in some sense (making this precise is the difficult bit).

We obtain either the left-preconditioned system

$$
P^{-1}Ax = P^{-1}b
$$

or the right-preconditioned system

$$
AP^{-1}y = b,
$$

where in the latter case we then additionally need to solve $Px=y$.

Classes of preconditioners include

* SPAI (Sparse Approximate Inverses)
* Incomplete LU Decomposition
* Incomplete Cholesky Decomposition
* Splitting Preconditioners
* Algebraic Multigrid Methods

These are also known as **algebraic preconditioners**. They consider the matrix $A$ and try to find an approximation to $A$ that is easily invertible.

A different class of preconditioners are **analytic preconditioners**. These are preconditioners that are often constructed as problems to PDEs that are easier to solve than the original PDE but still approximate to some sense the underlying physics of the problem.

## Sparse Approximate Inverse

As example of the class of algebraic preconditioners we consider here the Sparse Approximate Inverse (SPAI). Incomplete LU decompositions will be discussed later on. We note that SPAI is a technique that works well in certain cases, but is not suitable in others. **There is no general preconditioning technique that always works well.**

We denote by $\|A\|_F$ the Frobenious norm of a matrix $A$ defined by

$$
\|A\|_F^2 := \sum_{i, j}|a_{ij}|^2.
$$

The idea of SPAI is now to try to find a matrix $M := P^{-1}$ such that

$$
F(M) := \|I - AM\|_F
$$

is small. The minimum of this function is obviously reached for $M = A^{-1}$. But this usually not practical. Instead, we try to find a successive sequence of matrices $M_k$ that are approaching the minimum of the function $F$.

There are many ways to define an approximate minimization procedure to minimize $F$. The following is a global minimum residual algorithm, described by Saad in [Iterative Methods for Sparse Linear Systems](https://www-users.cs.umn.edu/~saad/IterMethBook_2ndEd.pdf).

$$
\begin{align}
C_k &= A M_k\nonumber\\
G_k &= I - C_k\nonumber\\
\alpha_k &=\text{tr}(G_k^TAG_k) / \|AG_k\|_F^2\nonumber\\
M_{k+1} &= M_k + \alpha_k G_k\nonumber\\
\end{align}
$$

In each step of the algorithm the matrix $M_k$ becomes slightly denser. Hence, in practice this is often combined with a numerical drop strategy for entries of $M$.

The following code implements this method. As starting matrix $M_0$ we choose $M_0 = \frac{2}{\|AA^T\|_1}A$, which was recommended by Chow and Saad in [Approximate Inverse Preconditioners via Sparse-Sparse Iterations](https://dl.acm.org/doi/10.1137/S1064827594270415). Note that in the following code we did not implement dropping of values. For practical purposes this is essential and strategies are discussed in the paper by Chow and Saad.

```python
def spai(A, m):
    """Perform m step of the SPAI iteration."""
    from scipy.sparse import identity
    from scipy.sparse import diags
    from scipy.sparse.linalg import onenormest

    n = A.shape[0]

    ident = identity(n, format='csr')
    alpha = 2 / onenormest(A @ A.T)
    M = alpha * A

    for index in range(m):
        C = A @ M
        G = ident - C
        AG = A @ G
        trace = (G.T @ AG).diagonal().sum()
        alpha = trace / np.linalg.norm(AG.data)**2
        M = M + alpha * G

    return M
```

We run the code with the following matrix, which is a slightly shifted version of the discrete 3-point second order differential operator.

```python
import numpy as np
from scipy.sparse import diags

n = 1000

data = [2.001 * np.ones(n),
        -1. * np.ones(n - 1),
        -1. * np.ones(n - 1)]

offsets = [0, 1, -1]

A = diags(data, offsets=offsets, shape=(n, n), format='csr')
```

The condition number without the preconditioner is

```python
np.linalg.cond(A.todense())
```

    3961.9652414689454

Let us now generate the preconditioner.

```python
M = spai(A, 50)
```

```python
%matplotlib inline
from matplotlib import pyplot as plt

fig = plt.figure(figsize=(8 ,8))
ax = fig.add_subplot(111)
ax.spy(M, markersize=1)
```

Let us check the condition number of the right-preconditioned system.

```python
np.linalg.cond(A.todense() @ M.todense())
```

    40.18659718436073

It has been reduced by a factor 100. This is a good sign. We are now running CG for the preconditioned and non-preconditioned system.

```python
%matplotlib inline
from matplotlib import pyplot as plt
from scipy.sparse.linalg import cg

n = A.shape[0]
b = np.ones(n)

residuals = []
callback = lambda x: residuals.append(np.linalg.norm(A @ x - b))
x, _ = cg(A, b, callback=callback)

residuals_preconditioned = []
callback = lambda x: residuals_preconditioned.append(np.linalg.norm(A @ x - b) / np.linalg.norm(b))
x, _ = cg(A, b, M=M, callback=callback)

fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111)
ax.semilogy(residuals, 'k--')
ax.semilogy(residuals_preconditioned, 'b--')
ax.set_ylabel('relative residual')
fig.legend(['no preconditioning', 'SPAI'], loc='lower center', fancybox=True, shadow=True)
```

We can see a significant improvement in the number of iterations. However, the comparison is not quite fair. Setting up the preconditioner took time, which should be taken into account. Just considering number of iterations is not sufficient. In practice, the overall computation time is the much more important measure. We have implemented here a very primitive variant of SPAI that should not be used in practice. For practical alternatives see the cited literature. For the type of matrix that $A$ is there are also much better precondtioners available, some of which we will encounter later.

## A note on preconditioned Conjugate Gradient

We have passed here the preconditioner as a matrix into the CG algorithm. This is possible if the preconditioner is also symmetric, positive definite. This is the case in our example. We are not going to discuss the details of preconditioned conjugate gradients in more details here and refer to the book by Saad.


#  Step-by-step NMO correction

Devito is equally useful as a framework for other stencil computations in general; for example, computations where all array indices are affine functions of loop variables. The Devito compiler is also capable of generating
arbitrarily nested, possibly irregular, loops. This key feature is needed to support many complex algorithms that are used in engineering and scientific practice, including applications from image processing, cellular automata, and machine-learning. This tutorial, a step-by-step NMO correction, is an example of it.

In reflection seismology, normal moveout (NMO) describes the effect that the distance between a seismic source and a receiver (the offset) has on the arrival time of a reflection in the form of an increase of time with offset. The relationship between arrival time and offset is hyperbolic.

Based on the field geometry information, each individual trace is assigned to the midpoint between the shot and receiver locations associated with that trace. Those traces with the same midpoint location are grouped together, making up a common midpoint gather (CMP).

Consider a reflection event on a CMP gather. The difference between the two-way time at a given offset and the two-way zero-offset time is called normal moveout (NMO). Reflection traveltimes must be corrected for NMO prior to summing the traces in the CMP gather along the offset axis. The normal moveout depends on velocity above the reflector, offset, two-way zero-offset time associated with the reflection event, dip of the reflector, the source-receiver azimuth with respect to the true-dip direction, and the degree of complexity of the near-surface and the medium above the reflector.

# Seismic modelling with devito

Before the NMO corretion 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 multiple 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

from devito import *
```

We will create a simple velocity model here by hand for demonstration purposes. This model essentially consists of three layers, each with a different velocity: 1.5km/s in the top layer, 2.5km/s in the middle layer and 4.5 km/s in the bottom layer.

```python
#NBVAL_IGNORE_OUTPUT
from examples.seismic import Model, plot_velocity

shape = (301, 501)  # Number of grid point (nx, ny, nz)
spacing = (10., 10)  # Grid spacing in m. The domain size is now 3km by 5km
origin = (0., 0)  # What is the location of the top left corner.

# Define a velocity profile. The velocity is in km/s
v = np.empty(shape, dtype=np.float32)
v[:,:100] = 1.5
v[:,100:350] = 2.5
v[:,350:] = 4.5

# With the velocity and model size defined, we can create the seismic model that
# encapsulates these properties. We also define the size of the absorbing layer as 10 grid points
model = Model(vp=v, origin=origin, shape=shape, spacing=spacing, space_order=4, nbl=40, bcs="damp")

plot_velocity(model)
```

Next we define the positioning and the wave signal of our source, as well as the location of our receivers. To generate the wavelet for our sources we require the discretized values of time that we are going to use to model a multiple "shot", which depends on the grid spacing used in our model. We will use one source and eleven receivers. The source is located in the position (550, 20). The receivers start at (550, 20) with an even horizontal spacing of 100m at consistent depth.

```python
from examples.seismic import TimeAxis

t0 = 0.  # Simulation starts a t=0
tn = 2400.  # Simulation last 2.4 second (2400 ms)
dt = model.critical_dt  # Time step from model grid spacing

time_range = TimeAxis(start=t0, stop=tn, step=dt)

nrcv = 250 # Number of Receivers
```

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

f0 = 0.010  # Source peak frequency is 10Hz (0.010 kHz)
src = RickerSource(name='src', grid=model.grid, f0=f0,
                   npoint=1, time_range=time_range)

# Define the wavefield with the size of the model and the time dimension
u = TimeFunction(name="u", grid=model.grid, time_order=2, space_order=4)

# We can now write the PDE
pde = model.m * u.dt2 - u.laplace + model.damp * u.dt
stencil = Eq(u.forward, solve(pde, u.forward))

src.coordinates.data[:, 0] = 400 # Source coordinates
src.coordinates.data[:, -1] = 20.  # Depth is 20m
```

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

rec = Receiver(name='rec', grid=model.grid, npoint=nrcv, time_range=time_range)
rec.coordinates.data[:,0] = np.linspace(src.coordinates.data[0, 0], model.domain_size[0], num=nrcv)
rec.coordinates.data[:,-1] = 20.  # Depth is 20m

# Finally we define the source injection and receiver read function to generate the corresponding code
src_term = src.inject(field=u.forward, expr=src * dt**2 / model.m)

# Create interpolation expression for receivers
rec_term = rec.interpolate(expr=u.forward)

op = Operator([stencil] + src_term + rec_term, subs=model.spacing_map)
op(time=time_range.num-1, dt=model.critical_dt)
```

    Operator `Kernel` run in 0.85 s

    PerformanceSummary([(PerfKey(name='section0', rank=None),
                         PerfEntry(time=0.7779790000000012, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),
                        (PerfKey(name='section1', rank=None),
                         PerfEntry(time=0.028586000000000156, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),
                        (PerfKey(name='section2', rank=None),
                         PerfEntry(time=0.03306400000000026, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])

How we are modelling a horizontal layers, we will group this traces and made a NMO correction using this set traces.

```python
offset = []
data = []
for i, coord in enumerate(rec.coordinates.data):
    off = (src.coordinates.data[0, 0] - coord[0])
    offset.append(off)
    data.append(rec.data[:,i])
```

Auxiliary function for plotting traces:

```python
#NBVAL_IGNORE_OUTPUT
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.axes_grid1 import make_axes_locatable

mpl.rc('font', size=16)
mpl.rc('figure', figsize=(8, 6))

def plot_traces(rec, xb, xe, t0, tn, colorbar=True):
    scale = np.max(rec)/100
    extent = [xb, xe, 1e-3*tn, t0]
    plot = plt.imshow(rec, cmap=cm.gray, vmin=-scale, vmax=scale, extent=extent)
    plt.xlabel('X position (km)')
    plt.ylabel('Time (s)')

    # Create aligned colorbar on the right
    if colorbar:
        ax = plt.gca()
        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.05)
        plt.colorbar(plot, cax=cax)
    plt.show()
```

# Common Midpoint Gather

At this point, we have a dataset composed of the receivers. "If our model wasn't purely horizontal, we would have to sort these traces by common midpoints prior to NMO correction."

```python
plot_traces(np.transpose(data), rec.coordinates.data[0][0]/1000, rec.coordinates.data[nrcv-1][0]/1000, t0, tn)
```

# NMO Correction

We can correct the measured traveltime of a reflected wave $t$ at a given offset $x$ to obtain the traveltime at normal incidence $t_0$ by applying the following equation:

\begin{equation*}
t = \sqrt{t_0^2 + \frac{x^2}{V_{nmo}^2}}
\end{equation*}

in which $V_{nmo}$ is the NMO velocity. This equation results from the Pythagorean theorem, and is only valid for horizontal reflectors. There are variants of this equation with different degrees of accuracy, but we'll use this one for simplicity.

For the NMO Correction we use a grid of size samples x traces.

```python
ns = time_range.num # Number of samples in each trace
grid = Grid(shape=(ns, nrcv)) # Construction of grid with samples X traces dimension
```

In this example we will use a constant velocity guide. The guide will be arranged in a SparseFunction with the number of points equal to number of samples in the traces.

```python
vnmo = 1500
vguide = SparseFunction(name='v', grid=grid, npoint=ns)
vguide.data[:] = vnmo
```

The computed offset for each trace will be arraged in another SparseFunction with number of points equal to number of traces.

```python
off = SparseFunction(name='off', grid=grid, npoint=nrcv)
off.data[:] = offset
```

The previous modelled traces will be arranged in a SparseFunction with the same dimensions as the grid.

```python
amps = SparseFunction(name='amps', grid=grid, npoint=ns*nrcv, dimensions=grid.dimensions, shape=grid.shape)
amps.data[:] = np.transpose(data)
```

Now, we define SparseFunctions with the same dimensions as the grid, describing the NMO traveltime equation. The $t_0$ SparseFunction isn't offset dependent, so the number of points is equal to the number of samples.

```python
sample, trace = grid.dimensions

t_0 = SparseFunction(name='t0', grid=grid, npoint=ns, dimensions=[sample], shape=[grid.shape[0]])
tt = SparseFunction(name='tt', grid=grid, npoint=ns*nrcv, dimensions=grid.dimensions, shape=grid.shape)
snmo = SparseFunction(name='snmo', grid=grid, npoint=ns*nrcv, dimensions=grid.dimensions, shape=grid.shape)
s = SparseFunction(name='s', grid=grid, dtype=np.intc, npoint=ns*nrcv, dimensions=grid.dimensions,
                   shape=grid.shape)
```

The Equation relates traveltimes: the one we can measure ($t_0$) and the one we want to know (t). But the data in our CMP gather are actually a matrix of amplitudes measured as a function of time ($t_0$) and offset. Our NMO-corrected gather will also be a matrix of amplitudes as a function of time (t) and offset. So what we really have to do is transform one matrix of amplitudes into the other.

With Equations we describe the NMO traveltime equation, and use the Operator to compute the traveltime and the samples for each trace.

```python
#NBVAL_IGNORE_OUTPUT

dtms = model.critical_dt/1000 # Time discretization in ms
E1 = Eq(t_0, sample*dtms)
E2 = Eq(tt, sp.sqrt(t_0**2 + (off[trace]**2)/(vguide[sample]**2) ))
E3 = Eq(s, sp.floor(tt/dtms))
op1 = Operator([E1, E2, E3])
op1()
```

    Operator `Kernel` run in 0.01 s

    PerformanceSummary([(PerfKey(name='section0', rank=None),
                         PerfEntry(time=0.001971, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])

With the computed samples, we remove all that are out of the samples range, and shift the amplitude for the correct sample.

```python
#NBVAL_IGNORE_OUTPUT

s.data[s.data >= time_range.num] = 0
E4 = Eq(snmo, amps[s[sample, trace], trace])

op2 = Operator([E4])
op2()

stack = snmo.data.sum(axis=1) # We can stack traces and create a ZO section!!!

plot_traces(snmo.data, rec.coordinates.data[0][0]/1000, rec.coordinates.data[nrcv-1][0]/1000, t0, tn)
```

# References:

    https://library.seg.org/doi/full/10.1190/tle36020179.1
    https://wiki.seg.org/wiki/Normal_moveout
    https://en.wikipedia.org/wiki/Normal_moveout