Calculate Numerical Diffusivity
============================

This ipython notebook calculates the numerical diffusivity experienced by a tracer, associated to a specific configuration of the MITgcm. In particular, it was developed to calculate the equivalent diffusivity $\kappa$, defined (here) as $\kappa = \kappa_{pres}+\kappa_{num}$, where $\kappa_{pres}$
is the prescibed or explicit tracer diffusivity one imposes on the model and $k_{num}$ is the additional diffusivity due to numerical truncation errors. Note that there are two $\kappa_{pres}$ and therefore two $\kappa$, one for the horizontal dimensions and one for the vertical one.

These calculations try to reproduce the method used by [1] Abernathy et al. 2010, [2] Hill et al. 2011, and [3] Leibensperger and Plumb, 2013 to determine the numerical diffusivity MITgcm Southern Ocean configurations [1,2] and a baroclinic flow simulation simulation [3].

The method goes as follows:

From the evolution equation for the variance of the tracer concentration in the model output

\begin{equation}
\frac{1}{2}\frac{\partial{\overline{q^{2}}}}{\partial{t}}=-\kappa_{h} \overline{|\nabla_h q|^2}-\kappa_{v} \overline{(\frac{\partial{q}}{\partial {z}})^{2}}
\end{equation}

one can fit by a least squares regression, suitable values of $\kappa_h$ and $\kappa_v$ that satisfy the equation.

```python
#KRM
import os

import numpy as np

import matplotlib.pyplot as plt

from scipy.stats import nanmean

from math import *

%matplotlib inline

import numpy as np

import matplotlib.colors as mcolors

from netCDF4 import Dataset

import pylab as pl

from decimal import getcontext, Decimal, Inexact
```

```python
#'''
#NAME
#    Custom Colormaps for Matplotlib
#PURPOSE
#    This program shows how to implement make_cmap which is a function that
#    generates a colorbar.  If you want to look at different color schemes,
#    check out https://kuler.adobe.com/create.
#PROGRAMMER(S)
#    Chris Slocum
#REVISION HISTORY
#    20130411 -- Initial version created
#    20140313 -- Small changes made and code posted online
#    20140320 -- Added the ability to set the position of each color
#'''

def make_cmap(colors, position=None, bit=False):
    #'''
    #make_cmap takes a list of tuples which contain RGB values. The RGB
    #values may either be in 8-bit [0 to 255] (in which bit must be set to
    #rue when called) or arithmetic [0 to 1] (default). make_cmap returns
    #a cmap with equally spaced colors.
    #Arrange your tuples so that the first color is the lowest value for the
    #colorbar and the last is the highest.
    #position contains values from 0 to 1 to dictate the location of each color.
    #'''
    import matplotlib as mpl
    import numpy as np
    bit_rgb = np.linspace(0,1,256)
    if position == None:
        position = np.linspace(0,1,len(colors))
    else:
        if len(position) != len(colors):
            sys.exit("position length must be the same as colors")
        elif position[0] != 0 or position[-1] != 1:
            sys.exit("position must start with 0 and end with 1")
    if bit:
        for i in range(len(colors)):
            colors[i] = (bit_rgb[colors[i][0]],
                         bit_rgb[colors[i][1]],
                         bit_rgb[colors[i][2]])
    cdict = {'red':[], 'green':[], 'blue':[]}
    for pos, color in zip(position, colors):
        cdict['red'].append((pos, color[0], color[0]))
        cdict['green'].append((pos, color[1], color[1]))
        cdict['blue'].append((pos, color[2], color[2]))

    cmap = mpl.colors.LinearSegmentedColormap('my_colormap',cdict,256)
    return cmap
```

```python
# Get field from MITgcm netCDF output
#
# statefile : string with /path/to/state.0000000000.t001.nc
# fieldname : string with the variable name as written on the netCDF file ('Temp', 'S','Eta', etc.)
#
def getField(statefile, fieldname):

    StateOut = Dataset(statefile)

    Fld = StateOut.variables[fieldname][:]

    shFld = np.shape(Fld)

    if len(shFld) == 1:

        return Fld

    if len(shFld) == 2:

        Fld2 = np.reshape(Fld,(shFld[0],shFld[1])) # reshape to pcolor order
        return Fld2

    elif len(shFld) == 3:

        Fld2 = np.zeros((shFld[0],shFld[1],shFld[2]))
        Fld2 = np.reshape(Fld,(shFld[0],shFld[1],shFld[2])) # reshape to pcolor order
        return Fld2

    elif len(shFld) == 4:

        Fld2 = np.zeros((shFld[0],shFld[1],shFld[2],shFld[3]))
        Fld2 = np.reshape(Fld,(shFld[0],shFld[1],shFld[2],shFld[3])) # reshape to pcolor order
        return Fld2

    else:

        print (' Check size of field ')

```

```python

def dsum(iterable):

    "Full precision summation using Decimal objects for intermediate values"
    # Transform (exactly) a float to m * 2 ** e where m and e are integers.
    # Convert (mant, exp) to a Decimal and add to the cumulative sum.
    # If the precision is too small for exact conversion and addition,
    # then retry with a larger precision.
    from decimal import getcontext, Decimal, Inexact
    getcontext().traps[Inexact] = True

    total = Decimal(0)
    for x in iterable:
        mant, exp = frexp(x)
        mant, exp = int(mant * 2.0 ** 53), exp-53
        while True:
            try:
                total += mant * Decimal(2) ** exp
                break
            except Inexact:
                getcontext().prec += 1
    return float(total)

```

Important input parameters:
------------------------

```python
#nx = 120  # number of x points in domain
#ny = 90  # number of y points in domain
#nz = 35   # number of z levels in domain
#nt = 10   # number of output files (n time)

nx = 180  # number of x points in domain
ny = 180  # number of y points in domain
nz = 35   # number of z levels in domain
nt = 55   # number of output files (n time)
```

Load output filenames:

```python
### Load output filenames ###

filename = '/ocean/kramosmu/MITgcm/CanyonUpwelling/180x180x35_BodyForcing_1Tr_LinProfiles/run05/output_0001/state.0000000000.glob.nc'#'/ocean/kramosmu/MITgcm/CanyonUpwelling/120x90x35_UpwLinRampWind_1Tr_LinProfiles/run30/output_0001/state.0000000000.t001.nc'

StateOut = Dataset(filename)

filename2 = '/ocean/kramosmu/MITgcm/CanyonUpwelling/180x180x35_BodyForcing_1Tr_LinProfiles/run05/output_0001/grid.glob.nc'#'/ocean/kramosmu/MITgcm/CanyonUpwelling/120x90x35_UpwLinRampWind_1Tr_LinProfiles/run30/output_0001/grid.t001.nc'
GridOut = Dataset(filename2)
#for dimobj in StateOut.variables.values():
#    print dimobj
filename3 = '/ocean/kramosmu/MITgcm/CanyonUpwelling/180x180x35_BodyForcing_1Tr_LinProfiles/run05/output_0001/ptracers.0000000000.glob.nc'#'/ocean/kramosmu/MITgcm/CanyonUpwelling/120x90x35_UpwLinRampWind_1Tr_LinProfiles/run30/output_0001/ptracers.0000000000.t001.nc'
PtracersOut = Dataset(filename3)

```

Make some plots

```python
plt.rcParams.update({'font.size':13})

tracer1=getField(filename3,'NO3')

tracer1msk = np.ma.masked_values(tracer1, 0)

MASK = np.ma.getmask(tracer1msk)

### PLOT TRACER 1 at zlev and time###

zlev =20
time =9

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

ax=plt.gca()
ax.set_axis_bgcolor((205/255.0, 201/255.0, 201/255.0))

plt.contourf(tracer1msk[time,zlev,:,:],30, cmap='rainbow')
cb = plt.colorbar()
cb.set_label('C', position=(0, 1),rotation=0)

z=StateOut.variables['Z']

plt.title("Tracer at z=%d m" % -z[zlev])

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

ax=plt.gca()
ax.set_axis_bgcolor((205/255.0, 201/255.0, 201/255.0))

plt.contourf(tracer1msk[time,:,:,60],30, cmap='rainbow')
plt.gca().invert_yaxis()
cb = plt.colorbar()
cb.set_label('C', position=(0, 1),rotation=0)

z=StateOut.variables['Z']

plt.title("Tracer at x=60 nx")

```

Calculate the volume of the domain
-------------------------------
The volume of a tracer cell (remember we have an Arakawa C grid, so this changes depending on which kind of cell we are thinking about) is given by

$V(i,j,k)=depth \times area = (hfacC(i,j,k)\times dRf(k)) \times rA(i,j) = (hfacC(i,j,k)\times dRf(k)) \times dXg(i,j) \times dYg(i,j)$,

where hfacC is the fraction of the cell that is open (not occupied with land). So, the total volume of the domain is

$\sum\limits_{i=1}^{nx}{\sum\limits_{j=1}^{ny}{\sum\limits_{k=1}^{nz}{(hfacC(i,j,k)\times dRf(k)) \times rA(i,j)}}}$

```python
DomVolume = 0.0

hFacC = getField(filename2,'HFacC')
drF   = getField(filename2,'drF')
rAC   = getField(filename2,'rA')

for ii in np.arange(nx):

    for jj in np.arange(ny):

        for kk in np.arange(nz):

            DomVolume = DomVolume + hFacC[kk,jj,ii]*drF[kk]*rAC[jj,ii]

print('\n The total volume is %e m^3 \n' %DomVolume)

```

     The total volume is 5.938188e+12 m^3

1st Term: The volume-weighted average of the squared concentration
-------------------------------------------------------------
The first term in the variance evolution equation is $\frac{1}{2}\frac{\partial{\overline{q^{2}}}}{\partial{t}}$. Note that we care about the time derivative of the variance, so that the mean concentration that usually appears in the definition of variance will not play a role here, since it is constant in time (we are not putting in or letting out any tracer).

We are going to calculate $\overline{q^2}$, the volume-weighted average of the squared concentration, and then the time derivative of that using a centered difference scheme.

```python
### Calculate the volume-weighted mean of q^2

q2mean = np.zeros((nt,1))

for tt in np.arange(nt):
    q2sum = 0.0
    for kk in np.arange(nz):

        for jj in np.arange(ny):

            q2sum = q2sum + np.sum(((hFacC[kk,jj,:]*drF[kk]*rAC[jj,:])/DomVolume)*(tracer1[tt,kk,jj,:])**2)

    q2mean[tt] = q2sum

    print('volume-weighted mean of q^2 nt=%d is %e  \n' %(tt, q2mean[tt]))

```

    volume-weighted mean of q^2 nt=0 is 6.364882e+02

    volume-weighted mean of q^2 nt=1 is 6.364880e+02

    volume-weighted mean of q^2 nt=2 is 6.364876e+02

    volume-weighted mean of q^2 nt=3 is 6.364869e+02

    volume-weighted mean of q^2 nt=4 is 6.364852e+02

    volume-weighted mean of q^2 nt=5 is 6.364800e+02

    volume-weighted mean of q^2 nt=6 is 6.364705e+02

    volume-weighted mean of q^2 nt=7 is 6.364582e+02

    volume-weighted mean of q^2 nt=8 is 6.364449e+02

    volume-weighted mean of q^2 nt=9 is 6.364312e+02

    volume-weighted mean of q^2 nt=10 is 6.364161e+02

    volume-weighted mean of q^2 nt=11 is 6.363992e+02

    volume-weighted mean of q^2 nt=12 is 6.363808e+02

    volume-weighted mean of q^2 nt=13 is 6.363611e+02

    volume-weighted mean of q^2 nt=14 is 6.363401e+02

    volume-weighted mean of q^2 nt=15 is 6.363172e+02

    volume-weighted mean of q^2 nt=16 is 6.362921e+02

    volume-weighted mean of q^2 nt=17 is 6.362660e+02

    volume-weighted mean of q^2 nt=18 is 6.362397e+02

    volume-weighted mean of q^2 nt=19 is 6.362125e+02

    volume-weighted mean of q^2 nt=20 is 6.361835e+02

    volume-weighted mean of q^2 nt=21 is 6.361527e+02

    volume-weighted mean of q^2 nt=22 is 6.361203e+02

    volume-weighted mean of q^2 nt=23 is 6.360869e+02

    volume-weighted mean of q^2 nt=24 is 6.360526e+02

    volume-weighted mean of q^2 nt=25 is 6.360174e+02

    volume-weighted mean of q^2 nt=26 is 6.359816e+02

    volume-weighted mean of q^2 nt=27 is 6.359454e+02

    volume-weighted mean of q^2 nt=28 is 6.359095e+02

    volume-weighted mean of q^2 nt=29 is 6.358747e+02

    volume-weighted mean of q^2 nt=30 is 6.358408e+02

    volume-weighted mean of q^2 nt=31 is 6.358069e+02

    volume-weighted mean of q^2 nt=32 is 6.357730e+02

    volume-weighted mean of q^2 nt=33 is 6.357393e+02

    volume-weighted mean of q^2 nt=34 is 6.357063e+02

    volume-weighted mean of q^2 nt=35 is 6.356739e+02

    volume-weighted mean of q^2 nt=36 is 6.356415e+02

    volume-weighted mean of q^2 nt=37 is 6.356090e+02

    volume-weighted mean of q^2 nt=38 is 6.355767e+02

    volume-weighted mean of q^2 nt=39 is 6.355447e+02

    volume-weighted mean of q^2 nt=40 is 6.355127e+02

    volume-weighted mean of q^2 nt=41 is 6.354809e+02

    volume-weighted mean of q^2 nt=42 is 6.354489e+02

    volume-weighted mean of q^2 nt=43 is 6.354167e+02

    volume-weighted mean of q^2 nt=44 is 6.353843e+02

    volume-weighted mean of q^2 nt=45 is 6.353518e+02

    volume-weighted mean of q^2 nt=46 is 6.353192e+02

    volume-weighted mean of q^2 nt=47 is 6.352865e+02

    volume-weighted mean of q^2 nt=48 is 6.352532e+02

    volume-weighted mean of q^2 nt=49 is 6.352193e+02

    volume-weighted mean of q^2 nt=50 is 6.351849e+02

    volume-weighted mean of q^2 nt=51 is 6.351501e+02

    volume-weighted mean of q^2 nt=52 is 6.351153e+02

    volume-weighted mean of q^2 nt=53 is 6.350805e+02

    volume-weighted mean of q^2 nt=54 is 6.350460e+02

```python
### Calculate time derivative
dt = (3600.*12.) # half day

dqdt = np.zeros(nt-2)

for tt in np.arange(nt-2):

    dqdt[tt] = (q2mean[tt+2]-q2mean[tt])/(2.*dt) # approximation of time derivative of v-w mean q^2

```

2nd Term: The volume-weighted average of the squared horizontal gradient
------------------------------------------------------------------
The second term in the variance evolution equation is $-\kappa_{h} \overline{|\nabla_h q|^2}$. Next, we calculate the square of the horizontal gradient $|\nabla_h q|^2=(\frac{\partial{q}}{\partial{x}})^2+(\frac{\partial{q}}{\partial{y}})^2$.

Spatial derivatives are approximated using a centered-difference scheme.

```python

### Calculate mean of (dq/dx)^2 + (dq/dy)^2)

dxG = getField(filename2,'dxG')
dyG = getField(filename2,'dyG')

qmeanDh = np.zeros((nt,1))
qmeanDx = np.zeros((nt,1))
qmeanDy = np.zeros((nt,1))

for tt in np.arange(nt):

    sumDx = 0.0

    tracer2dx = tracer1[tt,:,:,2:nx]
    tracer0dx = tracer1[tt,:,:,0:nx-2]
    tracerDerX = ((tracer2dx-tracer0dx)/(2.0*dxG[:,1:nx-1]))**2
    #print tracerDerX

    for kk in np.arange(nz):
        for jj in np.arange(ny):
            sumDx = sumDx + np.sum((hFacC[kk,jj,1:nx-1]*drF[kk]*rAC[jj,1:nx-1])*(tracerDerX[kk,jj,:]))

        #print(sumDx)

    qmeanDx[tt] = sumDx/DomVolume

    print('The v-w mean of (dq/dx)^2 at nt=%d is %e \n' % (tt,qmeanDx[tt]))

for tt in np.arange(nt):

    sumDy = 0.0

    tracer2dy = tracer1[tt,:,2:ny,:]
    tracer0dy = tracer1[tt,:,0:ny-2,:]
    tracerDerY = ((tracer2dy-tracer0dy)/(2.0*dyG[1:ny-1,:]))**2

    for kk in np.arange(nz):
        for ii in np.arange(nx):
            sumDy = sumDy + np.sum((hFacC[kk,1:ny-1,ii]*drF[kk]*rAC[1:ny-1,ii])*(tracerDerY[kk,:,ii]))

    qmeanDy[tt] = sumDy/DomVolume

    print('The v-w mean of (dq/dy)^2 at nt=%d is %e \n' % (tt,qmeanDy[tt]))

qmeanDh = qmeanDx + qmeanDy

```

    The v-w mean of (dq/dx)^2 at nt=0 is 9.899452e-07

    The v-w mean of (dq/dx)^2 at nt=1 is 9.899126e-07

    The v-w mean of (dq/dx)^2 at nt=2 is 9.897773e-07

    The v-w mean of (dq/dx)^2 at nt=3 is 9.893966e-07

    The v-w mean of (dq/dx)^2 at nt=4 is 9.887390e-07

    The v-w mean of (dq/dx)^2 at nt=5 is 9.878681e-07

    The v-w mean of (dq/dx)^2 at nt=6 is 9.868427e-07

    The v-w mean of (dq/dx)^2 at nt=7 is 9.856734e-07

    The v-w mean of (dq/dx)^2 at nt=8 is 9.843457e-07

    The v-w mean of (dq/dx)^2 at nt=9 is 9.828474e-07

    The v-w mean of (dq/dx)^2 at nt=10 is 9.811924e-07

    The v-w mean of (dq/dx)^2 at nt=11 is 9.794027e-07

    The v-w mean of (dq/dx)^2 at nt=12 is 9.774727e-07

    The v-w mean of (dq/dx)^2 at nt=13 is 9.753741e-07

    The v-w mean of (dq/dx)^2 at nt=14 is 9.731656e-07

    The v-w mean of (dq/dx)^2 at nt=15 is 9.709826e-07

    The v-w mean of (dq/dx)^2 at nt=16 is 9.688804e-07

    The v-w mean of (dq/dx)^2 at nt=17 is 9.668396e-07

    The v-w mean of (dq/dx)^2 at nt=18 is 9.648065e-07

    The v-w mean of (dq/dx)^2 at nt=19 is 9.627306e-07

    The v-w mean of (dq/dx)^2 at nt=20 is 9.605777e-07

    The v-w mean of (dq/dx)^2 at nt=21 is 9.583650e-07

    The v-w mean of (dq/dx)^2 at nt=22 is 9.561357e-07

    The v-w mean of (dq/dx)^2 at nt=23 is 9.539121e-07

    The v-w mean of (dq/dx)^2 at nt=24 is 9.517169e-07

    The v-w mean of (dq/dx)^2 at nt=25 is 9.495911e-07

    The v-w mean of (dq/dx)^2 at nt=26 is 9.476585e-07

    The v-w mean of (dq/dx)^2 at nt=27 is 9.459764e-07

    The v-w mean of (dq/dx)^2 at nt=28 is 9.444784e-07

    The v-w mean of (dq/dx)^2 at nt=29 is 9.430512e-07

    The v-w mean of (dq/dx)^2 at nt=30 is 9.416445e-07

    The v-w mean of (dq/dx)^2 at nt=31 is 9.403471e-07

    The v-w mean of (dq/dx)^2 at nt=32 is 9.391999e-07

    The v-w mean of (dq/dx)^2 at nt=33 is 9.382285e-07

    The v-w mean of (dq/dx)^2 at nt=34 is 9.373768e-07

    The v-w mean of (dq/dx)^2 at nt=35 is 9.366105e-07

    The v-w mean of (dq/dx)^2 at nt=36 is 9.359248e-07

    The v-w mean of (dq/dx)^2 at nt=37 is 9.351967e-07

    The v-w mean of (dq/dx)^2 at nt=38 is 9.345276e-07

    The v-w mean of (dq/dx)^2 at nt=39 is 9.338675e-07

    The v-w mean of (dq/dx)^2 at nt=40 is 9.332508e-07

    The v-w mean of (dq/dx)^2 at nt=41 is 9.326224e-07

    The v-w mean of (dq/dx)^2 at nt=42 is 9.320873e-07

    The v-w mean of (dq/dx)^2 at nt=43 is 9.315647e-07

    The v-w mean of (dq/dx)^2 at nt=44 is 9.310202e-07

    The v-w mean of (dq/dx)^2 at nt=45 is 9.304620e-07

    The v-w mean of (dq/dx)^2 at nt=46 is 9.301469e-07

    The v-w mean of (dq/dx)^2 at nt=47 is 9.294878e-07

    The v-w mean of (dq/dx)^2 at nt=48 is 9.288994e-07

    The v-w mean of (dq/dx)^2 at nt=49 is 9.279982e-07

    The v-w mean of (dq/dx)^2 at nt=50 is 9.272485e-07

    The v-w mean of (dq/dx)^2 at nt=51 is 9.265554e-07

    The v-w mean of (dq/dx)^2 at nt=52 is 9.260354e-07

    The v-w mean of (dq/dx)^2 at nt=53 is 9.253516e-07

    The v-w mean of (dq/dx)^2 at nt=54 is 9.246576e-07

    The v-w mean of (dq/dy)^2 at nt=0 is 1.408839e-06

    The v-w mean of (dq/dy)^2 at nt=1 is 1.408818e-06

    The v-w mean of (dq/dy)^2 at nt=2 is 1.408760e-06

    The v-w mean of (dq/dy)^2 at nt=3 is 1.408635e-06

    The v-w mean of (dq/dy)^2 at nt=4 is 1.408405e-06

    The v-w mean of (dq/dy)^2 at nt=5 is 1.408077e-06

    The v-w mean of (dq/dy)^2 at nt=6 is 1.407666e-06

    The v-w mean of (dq/dy)^2 at nt=7 is 1.407177e-06

    The v-w mean of (dq/dy)^2 at nt=8 is 1.406595e-06

    The v-w mean of (dq/dy)^2 at nt=9 is 1.405886e-06

    The v-w mean of (dq/dy)^2 at nt=10 is 1.405034e-06

    The v-w mean of (dq/dy)^2 at nt=11 is 1.404046e-06

    The v-w mean of (dq/dy)^2 at nt=12 is 1.402917e-06

    The v-w mean of (dq/dy)^2 at nt=13 is 1.401638e-06

    The v-w mean of (dq/dy)^2 at nt=14 is 1.400219e-06

    The v-w mean of (dq/dy)^2 at nt=15 is 1.398674e-06

    The v-w mean of (dq/dy)^2 at nt=16 is 1.397022e-06

    The v-w mean of (dq/dy)^2 at nt=17 is 1.395282e-06

    The v-w mean of (dq/dy)^2 at nt=18 is 1.393440e-06

    The v-w mean of (dq/dy)^2 at nt=19 is 1.391490e-06

    The v-w mean of (dq/dy)^2 at nt=20 is 1.389442e-06

    The v-w mean of (dq/dy)^2 at nt=21 is 1.387312e-06

    The v-w mean of (dq/dy)^2 at nt=22 is 1.385114e-06

    The v-w mean of (dq/dy)^2 at nt=23 is 1.382836e-06

    The v-w mean of (dq/dy)^2 at nt=24 is 1.380478e-06

    The v-w mean of (dq/dy)^2 at nt=25 is 1.378047e-06

    The v-w mean of (dq/dy)^2 at nt=26 is 1.375600e-06

    The v-w mean of (dq/dy)^2 at nt=27 is 1.373185e-06

    The v-w mean of (dq/dy)^2 at nt=28 is 1.370834e-06

    The v-w mean of (dq/dy)^2 at nt=29 is 1.368549e-06

    The v-w mean of (dq/dy)^2 at nt=30 is 1.366331e-06

    The v-w mean of (dq/dy)^2 at nt=31 is 1.364200e-06

    The v-w mean of (dq/dy)^2 at nt=32 is 1.362200e-06

    The v-w mean of (dq/dy)^2 at nt=33 is 1.360361e-06

    The v-w mean of (dq/dy)^2 at nt=34 is 1.358652e-06

    The v-w mean of (dq/dy)^2 at nt=35 is 1.357071e-06

    The v-w mean of (dq/dy)^2 at nt=36 is 1.355592e-06

    The v-w mean of (dq/dy)^2 at nt=37 is 1.354215e-06

    The v-w mean of (dq/dy)^2 at nt=38 is 1.352917e-06

    The v-w mean of (dq/dy)^2 at nt=39 is 1.351674e-06

    The v-w mean of (dq/dy)^2 at nt=40 is 1.350494e-06

    The v-w mean of (dq/dy)^2 at nt=41 is 1.349328e-06

    The v-w mean of (dq/dy)^2 at nt=42 is 1.348239e-06

    The v-w mean of (dq/dy)^2 at nt=43 is 1.347207e-06

    The v-w mean of (dq/dy)^2 at nt=44 is 1.346185e-06

    The v-w mean of (dq/dy)^2 at nt=45 is 1.345213e-06

    The v-w mean of (dq/dy)^2 at nt=46 is 1.344305e-06

    The v-w mean of (dq/dy)^2 at nt=47 is 1.343318e-06

    The v-w mean of (dq/dy)^2 at nt=48 is 1.342362e-06

    The v-w mean of (dq/dy)^2 at nt=49 is 1.341408e-06

    The v-w mean of (dq/dy)^2 at nt=50 is 1.340468e-06

    The v-w mean of (dq/dy)^2 at nt=51 is 1.339589e-06

    The v-w mean of (dq/dy)^2 at nt=52 is 1.338763e-06

    The v-w mean of (dq/dy)^2 at nt=53 is 1.337871e-06

    The v-w mean of (dq/dy)^2 at nt=54 is 1.336983e-06

```python

```

3rd Term: The volume-weighted average of the squared vertical derivative
------------------------------------------------------------------
The third term in the variance evolution equation is $-\kappa_{v} \overline{(\frac{\partial{q}}{\partial{z}})^2}$. Next, we calculate the square of the vertical gradient $(\frac{\partial{q}}{\partial{z}})^2$.

The vertical derivative is approximated using a centered-difference scheme.

```python
### Calculate mean of (dq/dz)^2

qmeanDz = np.zeros((nt,1))

for tt in np.arange(nt):
    qsumDz = 0.0
    for ii in np.arange(nx):

        for jj in np.arange(ny):

            for kk in np.arange(nz-2):

                if (hFacC[kk,jj,ii]*drF[kk])+(hFacC[kk+1,jj,ii]*drF[kk+1]) == 0:

                    continue

                else :

                   qsumDz =qsumDz +(hFacC[kk,jj,ii]*drF[kk]*rAC[jj,ii])*(((tracer1[tt,kk+2,jj,ii]-tracer1[tt,kk,jj,ii])/((hFacC[kk+1,jj,ii]*drF[kk+1])+(hFacC[kk,jj,ii]*drF[kk])))**2)

    qmeanDz[tt] = qsumDz/DomVolume

    print('Mean (dq/dz)^2 at nt=%d is %e \n' %(tt, qmeanDz[tt]))

```

    Mean (dq/dz)^2 at nt=0 is 1.576420e-02

    Mean (dq/dz)^2 at nt=1 is 1.576244e-02

    Mean (dq/dz)^2 at nt=2 is 1.576011e-02

    Mean (dq/dz)^2 at nt=3 is 1.575648e-02

    Mean (dq/dz)^2 at nt=4 is 1.575064e-02

    Mean (dq/dz)^2 at nt=5 is 1.574196e-02

    Mean (dq/dz)^2 at nt=6 is 1.573010e-02

    Mean (dq/dz)^2 at nt=7 is 1.571468e-02

    Mean (dq/dz)^2 at nt=8 is 1.569504e-02

    Mean (dq/dz)^2 at nt=9 is 1.567057e-02

    Mean (dq/dz)^2 at nt=10 is 1.564097e-02

    Mean (dq/dz)^2 at nt=11 is 1.560613e-02

    Mean (dq/dz)^2 at nt=12 is 1.556597e-02

    Mean (dq/dz)^2 at nt=13 is 1.552039e-02

    Mean (dq/dz)^2 at nt=14 is 1.546964e-02

    Mean (dq/dz)^2 at nt=15 is 1.541422e-02

    Mean (dq/dz)^2 at nt=16 is 1.535479e-02

    Mean (dq/dz)^2 at nt=17 is 1.529188e-02

    Mean (dq/dz)^2 at nt=18 is 1.522583e-02

    Mean (dq/dz)^2 at nt=19 is 1.515700e-02

    Mean (dq/dz)^2 at nt=20 is 1.508559e-02

    Mean (dq/dz)^2 at nt=21 is 1.501196e-02

    Mean (dq/dz)^2 at nt=22 is 1.493639e-02

    Mean (dq/dz)^2 at nt=23 is 1.485893e-02

    Mean (dq/dz)^2 at nt=24 is 1.477951e-02

    Mean (dq/dz)^2 at nt=25 is 1.469811e-02

    Mean (dq/dz)^2 at nt=26 is 1.461502e-02

    Mean (dq/dz)^2 at nt=27 is 1.453113e-02

    Mean (dq/dz)^2 at nt=28 is 1.444774e-02

    Mean (dq/dz)^2 at nt=29 is 1.436620e-02

    Mean (dq/dz)^2 at nt=30 is 1.428714e-02

    Mean (dq/dz)^2 at nt=31 is 1.421094e-02

    Mean (dq/dz)^2 at nt=32 is 1.414015e-02

    Mean (dq/dz)^2 at nt=33 is 1.407518e-02

    Mean (dq/dz)^2 at nt=34 is 1.401607e-02

    Mean (dq/dz)^2 at nt=35 is 1.396266e-02

    Mean (dq/dz)^2 at nt=36 is 1.391518e-02

    Mean (dq/dz)^2 at nt=37 is 1.387294e-02

    Mean (dq/dz)^2 at nt=38 is 1.383519e-02

    Mean (dq/dz)^2 at nt=39 is 1.380138e-02

    Mean (dq/dz)^2 at nt=40 is 1.377061e-02

    Mean (dq/dz)^2 at nt=41 is 1.374278e-02

    Mean (dq/dz)^2 at nt=42 is 1.371727e-02

    Mean (dq/dz)^2 at nt=43 is 1.369384e-02

    Mean (dq/dz)^2 at nt=44 is 1.367188e-02

    Mean (dq/dz)^2 at nt=45 is 1.365105e-02

    Mean (dq/dz)^2 at nt=46 is 1.363219e-02

    Mean (dq/dz)^2 at nt=47 is 1.361403e-02

    Mean (dq/dz)^2 at nt=48 is 1.359667e-02

    Mean (dq/dz)^2 at nt=49 is 1.357998e-02

    Mean (dq/dz)^2 at nt=50 is 1.356431e-02

    Mean (dq/dz)^2 at nt=51 is 1.354949e-02

    Mean (dq/dz)^2 at nt=52 is 1.353553e-02

    Mean (dq/dz)^2 at nt=53 is 1.352195e-02

    Mean (dq/dz)^2 at nt=54 is 1.350939e-02

Tracer variance plots:
--------------------

```python

fig1=plt.figure(figsize=(21,3))

plt.subplot(1,3,1)
ax=plt.gca()
ax.set_axis_bgcolor((205/255.0, 201/255.0, 201/255.0))

plt.plot(qmeanDh[1:len(qmeanDh)-1],dqdt)

plt.xlabel('$<|\partial_{h}q|^2>$')

plt.ylabel('$\partial{<q^2>}/\partial{t}$')

plt.subplot(1,3,2)
ax=plt.gca()
ax.set_axis_bgcolor((205/255.0, 201/255.0, 201/255.0))

plt.plot(qmeanDz[1:len(qmeanDz)-1],dqdt)

plt.xlabel('$<\partial{q^2}/\partial{z}>$')

plt.ylabel('$\partial{<q^2>}/\partial{t}$')

plt.subplot(1,3,3)
ax=plt.gca()
ax.set_axis_bgcolor((205/255.0, 201/255.0, 201/255.0))

plt.plot(np.arange(0,nt*0.5,0.5),q2mean)

plt.xlabel('days')

plt.ylabel('$<q^2>$')

plt.tight_layout()
```

Linear regression for $\kappa_h$ and $\kappa_v$
------------------------------------------

In the regression results box, x1 corresponds to $\kappa_h$ and x2 to $\kappa_v$.

```python
import pandas as pd
import statsmodels.api as sm

dt = 1*3600.0*2.0 # 2hrs

x = np.column_stack((qmeanDh[1:len(qmeanDh)-1], qmeanDz[1:len(qmeanDz)-1])) #variance of space derivatives

y = 0.5*dqdt[0:len(dqdt)]

## fit a OLS model with intercept
#x = sm.add_constant(x)
#est = sm.OLS(y, x).fit()
est = sm.GLS(y,x).fit()
est.summary()

```

<table class="simpletable">
<caption>GLS Regression Results</caption>
<tr>
  <th>Dep. Variable:</th>            <td>y</td>        <th>  R-squared:         </th> <td>   0.948</td>
</tr>
<tr>
  <th>Model:</th>                   <td>GLS</td>       <th>  Adj. R-squared:    </th> <td>   0.946</td>
</tr>
<tr>
  <th>Method:</th>             <td>Least Squares</td>  <th>  F-statistic:       </th> <td>   465.9</td>
</tr>
<tr>
  <th>Date:</th>             <td>Tue, 04 Nov 2014</td> <th>  Prob (F-statistic):</th> <td>1.72e-33</td>
</tr>
<tr>
  <th>Time:</th>                 <td>15:30:35</td>     <th>  Log-Likelihood:    </th> <td>  793.66</td>
</tr>
<tr>
  <th>No. Observations:</th>      <td>    53</td>      <th>  AIC:               </th> <td>  -1583.</td>
</tr>
<tr>
  <th>Df Residuals:</th>          <td>    51</td>      <th>  BIC:               </th> <td>  -1579.</td>
</tr>
<tr>
  <th>Df Model:</th>              <td>     2</td>      <th>                     </th>     <td> </td>
</tr>
</table>
<table class="simpletable">
<tr>
   <td></td>     <th>coef</th>     <th>std err</th>      <th>t</th>      <th>P>|t|</th> <th>[95.0% Conf. Int.]</th>
</tr>
<tr>
  <th>x1</th> <td>   -1.2900</td> <td>    0.125</td> <td>  -10.334</td> <td> 0.000</td> <td>   -1.541    -1.039</td>
</tr>
<tr>
  <th>x2</th> <td>    0.0002</td> <td> 1.99e-05</td> <td>    9.276</td> <td> 0.000</td> <td>    0.000     0.000</td>
</tr>
</table>
<table class="simpletable">
<tr>
  <th>Omnibus:</th>       <td> 2.295</td> <th>  Durbin-Watson:     </th> <td>   0.030</td>
</tr>
<tr>
  <th>Prob(Omnibus):</th> <td> 0.317</td> <th>  Jarque-Bera (JB):  </th> <td>   1.745</td>
</tr>
<tr>
  <th>Skew:</th>          <td> 0.443</td> <th>  Prob(JB):          </th> <td>   0.418</td>
</tr>
<tr>
  <th>Kurtosis:</th>      <td> 3.079</td> <th>  Cond. No.          </th> <td>1.72e+05</td>
</tr>
</table>

```python
import pandas as pd
import statsmodels.api as sm

dt = 1*3600.0*2.0 # 2hrs

x = np.column_stack((qmeanDh[1:len(qmeanDh)-1], qmeanDz[1:len(qmeanDz)-1])) #variance of space derivatives

y = 0.5*dqdt[0:len(dqdt)]

x1=qmeanDh[1:len(qmeanDh)-1]
x2=qmeanDz[1:len(qmeanDz)-1]

## fit a OLS model with intercept
#x = sm.add_constant(x)
#est = sm.OLS(y, x).fit()
est = sm.GLS(y,x1).fit()
est.summary()

```

<table class="simpletable">
<caption>GLS Regression Results</caption>
<tr>
  <th>Dep. Variable:</th>            <td>y</td>        <th>  R-squared:         </th> <td>   0.861</td>
</tr>
<tr>
  <th>Model:</th>                   <td>GLS</td>       <th>  Adj. R-squared:    </th> <td>   0.858</td>
</tr>
<tr>
  <th>Method:</th>             <td>Least Squares</td>  <th>  F-statistic:       </th> <td>   320.9</td>
</tr>
<tr>
  <th>Date:</th>             <td>Tue, 04 Nov 2014</td> <th>  Prob (F-statistic):</th> <td>6.74e-24</td>
</tr>
<tr>
  <th>Time:</th>                 <td>15:30:37</td>     <th>  Log-Likelihood:    </th> <td>  767.47</td>
</tr>
<tr>
  <th>No. Observations:</th>      <td>    53</td>      <th>  AIC:               </th> <td>  -1533.</td>
</tr>
<tr>
  <th>Df Residuals:</th>          <td>    52</td>      <th>  BIC:               </th> <td>  -1531.</td>
</tr>
<tr>
  <th>Df Model:</th>              <td>     1</td>      <th>                     </th>     <td> </td>
</tr>
</table>
<table class="simpletable">
<tr>
   <td></td>     <th>coef</th>     <th>std err</th>      <th>t</th>      <th>P>|t|</th> <th>[95.0% Conf. Int.]</th>
</tr>
<tr>
  <th>x1</th> <td>   -0.1328</td> <td>    0.007</td> <td>  -17.913</td> <td> 0.000</td> <td>   -0.148    -0.118</td>
</tr>
</table>
<table class="simpletable">
<tr>
  <th>Omnibus:</th>       <td>13.849</td> <th>  Durbin-Watson:     </th> <td>   0.012</td>
</tr>
<tr>
  <th>Prob(Omnibus):</th> <td> 0.001</td> <th>  Jarque-Bera (JB):  </th> <td>  15.626</td>
</tr>
<tr>
  <th>Skew:</th>          <td> 1.311</td> <th>  Prob(JB):          </th> <td>0.000404</td>
</tr>
<tr>
  <th>Kurtosis:</th>      <td> 3.450</td> <th>  Cond. No.          </th> <td>    1.00</td>
</tr>
</table>

```python
dt = 1*3600.0*2.0 # 2hrs

x = np.column_stack((qmeanDh[1:len(qmeanDh)-1], qmeanDz[1:len(qmeanDz)-1])) #variance of space derivatives

y = 0.5*dqdt[0:len(dqdt)]

x1=qmeanDh[1:len(qmeanDh)-1]
x2=qmeanDz[1:len(qmeanDz)-1]

## fit a OLS model with intercept
#x = sm.add_constant(x)
#est = sm.OLS(y, x).fit()
est = sm.GLS(y,x2).fit()
est.summary()

```

<table class="simpletable">
<caption>GLS Regression Results</caption>
<tr>
  <th>Dep. Variable:</th>            <td>y</td>        <th>  R-squared:         </th> <td>   0.839</td>
</tr>
<tr>
  <th>Model:</th>                   <td>GLS</td>       <th>  Adj. R-squared:    </th> <td>   0.836</td>
</tr>
<tr>
  <th>Method:</th>             <td>Least Squares</td>  <th>  F-statistic:       </th> <td>   271.9</td>
</tr>
<tr>
  <th>Date:</th>             <td>Tue, 04 Nov 2014</td> <th>  Prob (F-statistic):</th> <td>2.66e-22</td>
</tr>
<tr>
  <th>Time:</th>                 <td>15:30:38</td>     <th>  Log-Likelihood:    </th> <td>  763.74</td>
</tr>
<tr>
  <th>No. Observations:</th>      <td>    53</td>      <th>  AIC:               </th> <td>  -1525.</td>
</tr>
<tr>
  <th>Df Residuals:</th>          <td>    52</td>      <th>  BIC:               </th> <td>  -1524.</td>
</tr>
<tr>
  <th>Df Model:</th>              <td>     1</td>      <th>                     </th>     <td> </td>
</tr>
</table>
<table class="simpletable">
<tr>
   <td></td>     <th>coef</th>     <th>std err</th>      <th>t</th>      <th>P>|t|</th> <th>[95.0% Conf. Int.]</th>
</tr>
<tr>
  <th>x1</th> <td>-2.087e-05</td> <td> 1.27e-06</td> <td>  -16.489</td> <td> 0.000</td> <td>-2.34e-05 -1.83e-05</td>
</tr>
</table>
<table class="simpletable">
<tr>
  <th>Omnibus:</th>       <td>12.293</td> <th>  Durbin-Watson:     </th> <td>   0.011</td>
</tr>
<tr>
  <th>Prob(Omnibus):</th> <td> 0.002</td> <th>  Jarque-Bera (JB):  </th> <td>  13.676</td>
</tr>
<tr>
  <th>Skew:</th>          <td> 1.239</td> <th>  Prob(JB):          </th> <td> 0.00107</td>
</tr>
<tr>
  <th>Kurtosis:</th>      <td> 3.231</td> <th>  Cond. No.          </th> <td>    1.00</td>
</tr>
</table>

```python

```

```python

```


## 4.3 Decentralized Exact Diffusion

In Section 3 we discussed ATC/AWC-DGD that can solve the distributed optimization problem in a decentralized manner. However, it was shown that ATC/AWC-DGD can only converge to a neighborhood of the solution, but not to the solution itself. This section will discuss another decentralized algorithm that can converge to the **exact** solution.

### 4.3.1 Problem and decentralized gradient descent

Consider $n$ computing nodes collaborate to solve the problem:

$$\min_{x\in \mathbb{R}^d} \quad \frac{1}{n}\sum_{i=1}^n f_i(x)$$

where $f_i(x)$ is a local and private function held by node $i$. Each node $i$ can access its own variable $x$ or gradient $\nabla f_i(x)$, but it has to communicate to access information from other nodes.

With combination matrix $W$, the decentralized gradient descent has recursions as follows:

$$\begin{align}
x_i^{(k+1)} = \sum_{j=1}^n w_{ij} \Big(x_j^{(k)} - \alpha \nabla f_j(x_j^{(k)}) \Big) \quad \mbox{(ATC-DGD)}
\end{align}$$

Now we briefly discuss how ATC-DGD cannot converge to the exact solution. Suppose $x^\star$ is the global solution, and each node has converged to $x^\star$, i.e., $x_i^{(k)} = x^\star$. Now we examine whether $x_i^{(k+1)}$ can stay at $x^\star$ or not. Note that while $\sum_{i=1}^n \nabla f_i(x^\star) =0$, each single $\nabla f_i(x^\star)$ does not equal to $0$ in general. With the fact that $x_i^{(k)} = x^\star$ and $W$ is doubly stochastic, we have

$$\begin{align}
x_i^{(k+1)} = \sum_{j=1}^n w_{ij} \Big(x^\star - \alpha \nabla f_j(x^\star) \Big) = x^\star - \alpha \sum_{j=1}^n w_{ij} \nabla f_j(x^\star)  \neq x^\star.
\end{align}$$

The above relation implies that while ATC-DGD has converged to the global solution $x^\star$ at iteration $k$, ATC-DGD will jump away from $x^\star$ at iteration $k+1$. In other words, ATC-DGD cannot converge to the exact solution $x^\star$, but to its neighborhood.

### 4.3.2 Decentralized exact diffusion

We now discuss the exact diffusion algorithm that can converge to the global solution exactly. It was inspired by a novel algorithm EXTRA \[SLWY-2015\]. We let $\bar{W} = (I+W)/2$. Apparently, $\bar{W}$ is also doubly-stochastic. The exact diffusion has the following recursion:

$$\begin{align}
x_i^{(k+1)} = \sum_{j=1}^n \bar{w}_{ij} \Big(2 x_j^{(k)} - x_j^{(k-1)} - \alpha\big(\nabla f_j(x_j^{(k)}) - \nabla f_j(x_j^{(k-1)})\big) \Big) \quad \mbox{(exact-diffusion)}
\end{align}$$
where the initialization is $x_i^{(1)} = \sum_{j=1}^n \bar{w}_{ij}(x_j^{(0)} - \alpha \nabla f_j(x_j^{(0)})).$

Suppose $x_i^{(k)}$ stays at $x^\star$. According to the above exact diffusion recursion, we have

$$x_i^{(k+1)} = \sum_{j=1}^n \bar{w}_{ij} \Big(2 x^\star - x^\star - \alpha\big(\nabla f_j(x^\star) - \nabla f_j(x^\star)\big) \Big) = \sum_{j=1}^n \bar{w}_{ij} x^\star = x^\star,$$

which implies that $x_i^{(k+1)}$ will stay at $x^\star$ if $x_i^k = x^\star$. This derivation indicates that exact diffusion can converge to the exact solution. Another widely-used formulation of exact diffusion is (for $k=0,1,\cdots$):

\begin{align}
\psi_i^{(k)} &= x_i^{(k)} - \alpha \nabla f_i(x_i^{(k)})  \hspace{1cm} \mbox{(adaptation)}\\
\phi_i^{(k)} &= \psi_i^{(k)} + x_i^{(k)} - \psi_i^{(k-1)} \hspace{0.45cm} \mbox{(correction)}\\
x_i^{(k+1)} &= \sum_{j=1}^n \bar{w}_{ij} \phi_j^{(k)} \hspace{2.1cm} \mbox{(combination)}
\end{align}

As the initialization, we let $x_i^{(0)} = \psi_i^{(-1)}$ for each node $i$.

### 4.3.3 Convergence properties

When each $f_i(x)$ is smooth and strongly-convex, it is established in \[YYZS\] that exact diffusion can converge to the global solution $x^\star$ linearly.

### 4.3.4 An example: least-square problem

In this section, we will show a demo on how to solve a decentralized least-square problem with Exact Diffusion. Suppose $n$ computing nodes collaborate to solve the following problem:

$$\min_x \quad \frac{1}{n}\sum_{i=1}^n \|A_i x - b_i\|^2$$

where $\{A_i, b_i\}$ are local data held in node $i$.

#### 4.3.4.1 Set up Environment

```
import ipyparallel as ipp

rc = ipp.Client(profile="bluefog")
dview = rc[:]  # A DirectView of all engines
dview.block = True
```

```
%%px
import numpy as np
import bluefog.torch as bf
import torch
from bluefog.common import topology_util
import networkx as nx
import matplotlib.pyplot as plt

%matplotlib inline

bf.init()
```

#### 4.3.4.2 Generate local data $A_i$ and $b_i$

```
%%px

def generate_data(m, n):

    A = torch.randn(m, n).to(torch.double)
    x_o = torch.randn(n, 1).to(torch.double)
    ns = 0.1 * torch.randn(m, 1).to(torch.double)
    b = A.mm(x_o) + ns

    return A, b
```

```
%%px

def distributed_grad_descent(A, b, maxite=5000, alpha=1e-1):

    x_opt = torch.zeros(n, 1, dtype=torch.double)

    for _ in range(maxite):
        # calculate local gradient
        grad_local = A.t().mm(A.mm(x_opt) - b)

        # global gradient
        grad = bf.allreduce(grad_local, name="gradient")

        # distributed gradient descent
        x_opt = x_opt - alpha * grad

    grad_local = A.t().mm(A.mm(x_opt) - b)
    grad = bf.allreduce(grad_local, name="gradient")  # global gradient

    # evaluate the convergence of distributed gradient descent
    # the norm of global gradient is expected to 0 (optimality condition)
    global_grad_norm = torch.norm(grad, p=2)
    if bf.rank() == 0:
        print(
            "[Distributed Grad Descent] Rank {}: global gradient norm: {}".format(
                bf.rank(), global_grad_norm
            )
        )

    return x_opt
```

In the following code we run distributed gradient descent to achieve the global solution $x^\star$ to the optimization problem. To validate whether $x^\star$ is optimal, it is enough to examine $\frac{1}{n}\sum_{i=1}^n \nabla f_i(x^\star) = 0$.

```
%%px

m, n = 2000, 25
A, b = generate_data(m, n)
x_opt = distributed_grad_descent(A, b, maxite=1000, alpha=1e-4)
```

    [stdout:3] [Distributed Grad Descent] Rank 0: global gradient norm: 1.0325386419283952e-12

#### 4.3.4.3 Exact Diffusion
In this section, we depict the convergence curve of the  Exact Diffusion algorithm in a decentralized environment. We will utilize the $x^\star$ achieved by distributed gradient descent as the optimal solution. First, we define one step of the Exact Diffusion method.

```
%%px

def exact_diffusion_one_step(x, A, b, pre_psi, alpha=1e-2):
    # Compute local gradient
    grad = A.t().mm(A.mm(x)-b)
    # Correct the bias
    psi = x-alpha*grad
    phi = psi + x - pre_psi
    # Communication with neighbors
    x_opt = bf.neighbor_allreduce(phi)
    return x_opt, psi
```

Next we run the Exact Diffusion algorithm.

```
%%px

# Set topology as ring topology.
G = topology_util.RingGraph(bf.size())
bf.set_topology(G)

max_iter = 200
alpha = 2e-4

x = torch.zeros(n, 1, dtype=torch.double)  # Initialize x for exact diffusion
psi = x
rel_error = torch.zeros((max_iter, 1))

for ite in range(max_iter):
    if ite % 10 == 0:
        if bf.rank() == 0:
            print("Progress {}/{}".format(ite, max_iter))
    x, psi = exact_diffusion_one_step(x, A, b, psi, alpha=alpha)
    rel_error[ite] = torch.norm(x - x_opt, p=2) / torch.norm(x_opt, p=2)
```

    [stdout:3]
    Progress 0/200
    Progress 10/200
    Progress 20/200
    Progress 30/200
    Progress 40/200
    Progress 50/200
    Progress 60/200
    Progress 70/200
    Progress 80/200
    Progress 90/200
    Progress 100/200
    Progress 110/200
    Progress 120/200
    Progress 130/200
    Progress 140/200
    Progress 150/200
    Progress 160/200
    Progress 170/200
    Progress 180/200
    Progress 190/200

```
# collect relative error
rel_error = dview.pull("rel_error", block=True)
rel_error_avg = sum(rel_error) / len(rc.ids)
```

```
import matplotlib.pyplot as plt

%matplotlib inline

plt.semilogy(rel_error_avg)

plt.xlabel("Iteration", fontsize=16)
plt.ylabel("Relative error", fontsize=16)
```

It can be observed from the above figure that the Exact Diffusion can precisely converge to the optimal solution, demonstrating that the Exact Diffusion resolves the bias introduced in a decentralized environment.

```

```


# Intro to Seismology: Programming for Homework 3
## Name:

## Introduction

The goal of this assignment is to locate an earthquake based on travel times.  To do this will require three ingredients

  1. A function that generates travel times from any point in the media to all receivers.
     - We will use a closed-form solution to the eikonal equation in a homogeneous medium.
  2. A way to compare the estimated and observed travel times.
     - We will use least-squares.
  3. A strategy for finding a model that better predicts the correct observed travel times.
     - We will use gradient-based nonlinear optimization (Geiger's method) and grid-search.

```python
# includes and convenience functions
import matplotlib.pyplot as plt
from numpy import amin, amax
from numpy import argmin, argmax, unravel_index
from numpy import zeros
from numpy import median
from numpy import linspace
from numpy import exp
from numpy import meshgrid
from numpy import average
from numpy import copy
from numpy import sqrt
from numpy.linalg import lstsq

def unpackModel(model):
    """!
    Convenience function to unpack model stored as a vector.

    Parameters
    ----------
    model : array_like
        Model to unpack.

    Returns
    -------
    xs : double
        x source position in meters.
    ys : double
        y source position in meters.
    t0 : double
        Origin time in seconds.
    """
    xs = model[0]
    ys = model[1]
    t0 = model[2]
    return xs, ys, t0

def packModel(xs, ys, t0):
    """!
    Convenience function to pack the model as a vector.

    Parameters
    ----------
    xs : double
        x source position in meters.
    ys : double
        y source position in meters.
    t0 : double
        Origin time in seconds.

    Returns
    -------
    model : array_like
        Model parameters packed as a vector.
    """
    model = zeros(3)
    model[0] = xs
    model[1] = ys
    model[2] = t0
    return model
```

## Modeling Travel Times

In a homogeneous medium the estimated travel time, $t_{est}$, is simply the distance, $d$, divided by the medium velocity, $v$, i.e., $t_{est} = \frac{d}{v}$.  Moreover, observe that earthquakes have unknown origin times, $t_0$.  Thus, the estimated travel time function must add some bias corresponding to the origin time, i.e., $t_{est} = t_0 = \frac{d}{v}$.  For ease of implementation, we will consider only a $2D$ media, so that distance between the source, $\textbf{x}_s$, and receiver, $\textbf{x}_j$, is $d = \sqrt{ (x_j - x_s)^2 + (y_j - y_s)^2}$.  Therefore, our travel time function modeled at the $j$'th receiver is

\begin{equation}
  T_{est}^{(j)}(t_0, \textbf{x}_s)
 =t_0 + \frac{\sqrt{ (x_j - x_s)^2 + (y_j - y_s)^2}}{v}
\end{equation}

In general it will be more convenient to define the model as $\textbf{m} = \{x_s, y_s, t_0 \}$ so that

\begin{equation}
  T_{est}^{(j)}(\textbf{m})
 =t_0 + \frac{\sqrt{ (x_j - x_s)^2 + (y_j - y_s)^2}}{v}
\end{equation}

```python
def computeTravelTimes(xr, yr, vel, model):
    """
    Computes travel times in a 2D homogeneous medium with contant velocity.

    Parameters
    ----------
    xr : array_like
        The x positions of the receivers (meters).  This has dimension [nrec]
        where nrec is the number of receivers.
    yr : array_like
        The y positions of the receivers (meters).  This has dimension [nrec]
        where nrec is the number of receivers.
    vel : double
        The constant velocity in the medium (meters/seconds).
    model : array_like
        The model that contains the source location in (x,y) as well as the origin time.

    Returns
    -------
    test : array_like
        The estimate travel times from the source to each receiver (seconds).
    """
    xs, ys, t0 = unpackModel(model)
    test = # Put your code here
    return test
```

## Measuring Misfit

In all inverse problems a measure of how well the model explains the data is required.  Such measures are typically called objective or cost funtcions.  Many objective functions exist and each has its own merits and drawbacks.  Arguably, the most common objective function is the sum of the residuals squared.  Least-squares typically yields quick convergence in non-linear inversions and results in simple very algebra.  Unfortunately, outliers can have an undo influence on the solution and therefore, least-squares is not robust.  To mitigate the last point, it is generally intelligent to screen data and manually remove outlier data prior to inversion.  That said, the misfit (cost) function to minimize looks like

\begin{equation}
  \mathcal{C}(\textbf{m}) = \sum_{j=1}^{n_{obs}} (T_{obs}^{(j)} - T_{est}^{(j)}(\textbf{m}) )^2
\end{equation}

where we are simply squaring and summing the residual travel times for all $n_{obs}$ observations.  If only locating with $P$ waves then the number of observations would be the number of receivers in the network.  Notice, the residual is the difference between the observed travel time, $T_{obs}^{(j)})$, and estimate travel time for the $j$'th observation.

```python
def computeResiduals(tobs, test):
    """
    Computes the residuals.

    Parameters
    ----------
    tobs : array_like
         The observations in seconds.  This has dimension [nrec] which is
         the number of receivers.
    test : array_like
         The estimates in seconds.  This has dimension [nrec] which is the
         number of recievers.

    Returns
    -------
    res : array_like
         The travel time residuals in seconds.
    """
    res = # Put your code here
    return res

def computeCostFunction(model, tobs, xr, yr, vel):
    """
    Computes the cost (objective) function corresponding to the current model.

    Parameters
    ----------
    model : array_like
         Contains the current source position and origin time.
    tobs : array_like
         The observed travel times (seconds).  This has dimension [nrec] which is
         the number of receivers (observations).
    xr : array_like
         The receiver positions in x in meters.  This has dimension [nrec].
    yr : array_like
         The receiver positions in y in meters.  This has dimension [nrec].
    vel : double
         The constant velocity in the medium in meters/second.

    Returns
    -------
    cost : double
         The sum of the squared residuals.
    """
    # Compute the estimate travel times with computeTravelTimes
    test = # Put your code here
    # Compute the residuals with computeResiduals
    res = # Put your code here
    # Compute the sum of the squared residuals (e.g., sum(res*res))
    cost = # Put your code here
    return cost
```

## Geiger's Method

The iterative solution to minimizing the previous cost function was developed well before the advent of computers in 1910 by Geiger.  In principle the proposed solution is no different than iteratively minimizing a function in Calculus I.  The difference between Geiger's method and a Calculus I optimization problem is that there are four parameters (a hypocenter and origin) comprising an optimum instead of one.

Ideally, we desire an optimal model, $\textbf{m}^*$ that yields a produces a cost function whose every residual is identically $0$
\begin{equation}
  \mathbf{r}(\textbf{m}^*) = \textbf{0}
\end{equation}

Here each row of the residual vector $\mathbf{r}(\textbf{m}^*) = \textbf{0}$ is an equation of the form
$T_{obs}^{(j)} - T_{est}^{(j)} = 0$.  By inspection, we note that $\mathcal{C} = \mathbf{r}^T \cdot \mathbf{r}$ is indeed the sum of the residuals squared.

In general, we don't know $\textbf{m}^*$.  Instead, we guess it to be some initial model, $\textbf{m}_0$.  Then, we look for a perturbation, $\textbf{m}_1 \leftarrow \textbf{m}_0 + \alpha \textbf{p}$, so that $\textbf{m}_1$ is a better approximation to $\textbf{m}^*$.  Here $\alpha$ is a unitless 'step-length' that scales the search-direction, $\textbf{p}$.  The search direction necessarily has units equivalent to the model, i.e., $\textbf{p} = \{\Delta x, \Delta y, \Delta t\}$.  For the time being, let's ignore $\alpha$ though in practice it is a very important parameter whose estimation requires a difficult class of algorithms called line-search methods.  Thus, we seek a solution to the over-determined system

\begin{equation}
  \mathbf{r}(\textbf{m}_0 + \textbf{p})
 =\textbf{0}
\end{equation}

Such a formulation is of little practical value so we linearize the above with a Taylor expansion about $ \textbf{p}$

\begin{align}
   \mathbf{r}(\textbf{m}_0 + \textbf{p})
&=\mathbf{r}(\textbf{m}_0)
 + \nabla_{\textbf{m}}  \mathbf{r}(\textbf{m}_0) \textbf{p}
 +\cdots \\
&\approx
  \mathbf{r}(\textbf{m}_0)
 +\nabla_{\textbf{m}}  \mathbf{r}(\textbf{m}_0) \textbf{p}
%+\frac{1}{2} \alpha^2 \textbf{p}^T \nabla  \mathbf{r}(\textbf{m}_0)^T \cdot \nabla  \mathbf{r}(\textbf{m}_0) \textbf{p}
\end{align}

Recalling that the above should equal $\textbf{0}$ we can solve for the search direction

\begin{equation}
  \mathbf{r}(\textbf{m}_0) + \nabla_{\textbf{m}} \mathbf{r}(\textbf{m}_0) \cdot \textbf{p} = 0
\end{equation}

Note, that the gradient of a vector involves computing an outer-product.  This yields a matrix.  The matrix is known as the Jacobian matrix, $\nabla_{\textbf{m}} r(\textbf{m}) = J(\textbf{m})$ and is of great consequence to inverse problems.  To obtain a search direction we solve a system of over-determined equations

\begin{equation}
 J(\textbf{m}_0) \cdot \textbf{p} =-\mathbf{r}(\textbf{m}_0)
\end{equation}

whose solution is conventionally written as the normal equations

\begin{equation}
  \textbf{p} =-(J^T J)^{-1} J^T \textbf{r}(\textbf{m}_0)
\end{equation}

The term $J^T \textbf{r}$ is the gradient of the objective function, $\nabla_{\textbf{m}} \mathcal{C}$.  Indeed, if we approximate $(J^T J)^{-1}$ as the identity matrix, then we are left with something of the form $\textbf{p} =-J^T \textbf{r} =-\textbf{g}$.  This says that the update direction is the negative of the gradient.  This result is very intuitive and indicates that to optimize this nonlinear function we should simply march down the gradient of the objective function.  Recall, that the gradient points in the direction of maximum increase so a negative gradient points in the direction in the direction of maximum decrease.  This particularly optimization scheme is steepest descent.

Finally, the location algorithm proceeds as follows:

   1.  Specify an initial model $\textbf{m}_0 = (x_s, y_s, t_0)$.
   2.  Compute the objective function, $\mathcal{C}(\textbf{m}_0)$.  If the objective function is small enough then stop.
   3.  Compute the search direction either by solving $J \textbf{p} =-\textbf{f}$ (Gauss-Newton) or $\textbf{p} =-J^T \textbf{r}$ (steepest-descent).
   4.  Update the model $\textbf{m}_1 = \textbf{m}_0 - \alpha \textbf{p}$ where we will take $\alpha = 1$ for Gauss-Newton or some small number $\ll 1$ for steepest-descent.
   5.  Iterate

```python
def computeGaussNewtonSearchDirection(model, tobs, xr, yr, vel):
    """
    Computes the Gauss-Newton search direction from the Jacobian and residual
    in a numerically stable way by solving the normal equations, Jp=r, with the
    singular value decomposition instead of the normal equations.

    Parameters
    ----------
    model : Array of length 3
        This is an array containing the model.
    tobs : array_like
        These are the observed travel times in second.  This has dimension [nrec]
        which corresponds to the number observations.
    xr : array_like
        These are the x positions (meters) of each reciever.  This has dimension [nrec]
        which corresponds to the number of receivers and observations.
    yr : array_like
        These are the y positions (meters) of each receiver.  This has dimension [nrec]
        which corresponds to the number of receivers and observations.
    vel : double
        This is the constant velocity in the medium.

    Returns
    -------
    p : array_like
        The Gauss-Newton search direction.  This has dimension [3] and is paked
        in the same order as the model so that m1 = m0 - alpha*p.
    """
    # Compute the Jacobian with computeJacobian
    J = # Put your code here
    # Compute the estimate travel times with computeTravelTimes
    test = # Put your code here
    # Compute the residuals with computeResidiuals
    res = # Put your code here
    # The normal equations are equivalent to solving the overdetermined
    # system J p =-r
    p = lstsq(J, -res, rcond=-1)[0]
    return p

def computeSteepestDescentSearchDirection(model, tobs, xr, yr, vel):
    """
    Computes the steepest descent search direction from the Jacobian and residual.

    Parameters
    ----------
    model : Array of length 3
        This is an array containing the model.
    tobs : array_like
        These are the observed travel times in second.  This has dimension [nrec]
        which corresponds to the number observations.
    xr : array_like
        These are the x positions (meters) of each reciever.  This has dimension [nrec]
        which corresponds to the number of receivers and observations.
    yr : array_like
        These are the y positions (meters) of each receiver.  This has dimension [nrec]
        which corresponds to the number of receivers and observations.
    vel : double
        This is the constant velocity in the medium.

    Returns
    -------
    p : array_like
        The Gauss-Newton search direction.  This has dimension [3] and is paked
        in the same order as the model so that m1 = m0 - alpha*p.
    """
    # Compute the Jacobian with computeJacobian
    J = # Put your code here
    # Compute the estimate travel times with computeTravelTimes
    test = # Put your code here
    # Compute the residuals with computeResidiuals
    p =-J.T@res # negative of gradient
    return p
```

### Computing Derivatives

Geiger's method is a gradient-descent technique that naturally requires derivatives of the cost function.  Recall, our objective function looks like

\begin{equation}
  \mathcal{C}(t_0, \textbf{x}_s) = \sum_{j=1}^{n_{obs}} (T_{obs}^{(j)} - T_{est}^{(j)}(\textbf{m}) )^2 = \sum_{r=1}^{n_{rec}} \left (T_{obs}^{(r)} - \left ( t_0 + \frac{\sqrt{ (x_j - x_s)^2 + (y_j - y_s)^2}}{v} \right ) \right )^2
\end{equation}

Here what is required are the derivatives of the cost-function that comprise the gradient:

\begin{equation}
   \frac{\partial \mathcal{C}}{\partial x_s}
 =\sum_{j=1}^{n_{obs}} -2(T_{obs}^{(j)} - T_{est}^{(j)}) \frac{\partial T_{est}^{(j)}}{\partial x_s}
 = 2 \sum_{j=1}^{n_{obs}} (T_{obs}^{(j)} - T_{est}^{(j)}) \frac{(x_j - x_s)}{v \sqrt{(x_j-x_s)^2 + (y_j-y_s)^2}}
\end{equation}
\begin{equation}
   \frac{\partial \mathcal{C}}{\partial y_s}
 =\sum_{j=1}^{n_{obs}} -2(T_{obs}^{(j)} - T_{est}^{(j)}) \frac{\partial T_{est}^{(j)}}{\partial y_s}
 = 2 \sum_{j=1}^{n_{obs}} (T_{obs}^{(j)} - T_{est}^{(j)}) \frac{(y_j - y_s)}{v \sqrt{(x_j-x_s)^2 + (y_j-y_s)^2}}
\end{equation}
\begin{equation}
   \frac{\partial \mathcal{C}}{\partial t_0}
 =\sum_{j=1}^{n_{obs}} -2(T_{obs}^{(j)} - T_{est}^{(j)}) \frac{\partial T_{est}^{(j)}}{\partial t_0}
 =-2 \sum_{j=1}^{n_{obs}} (T_{obs}^{(j)} - T_{est}^{(j)})
\end{equation}

Defining the residual, $r_j$, for the $j$'th observation as $r_j \equiv T_{obs}^{(j)} - T_{est}^{(j)}$, and the distance between the $j$'th source and receiver, $d_j$, as $d_j \equiv \sqrt{ (x_j - x_s)^2 + (y_j - y_s)^2}$ allows us to write the compactly as
\begin{equation}
   \frac{\partial \mathcal{C}}{\partial x_s}
 = 2 \sum_{j=1}^{n_{obs}} r_j \frac{(x_j - x_s)}{v d_j}
\end{equation}
\begin{equation}
   \frac{\partial \mathcal{C}}{\partial y_s}
 = 2 \sum_{j=1}^{n_{obs}} r_j \frac{(y_j - y_s)}{v d_j}
\end{equation}
\begin{equation}
   \frac{\partial \mathcal{C}}{\partial t_0}
 =-2 \sum_{j=1}^{n_{obs}} r_j
\end{equation}

### The Jacobian Matrix

Notice that the gradient has a common $\sum_{j=1}^{n_{obs}} r_j$ term scaled by some coefficient.  We can rewrite the above has a matrix-vector multiplication

\begin{equation}
  \nabla \mathcal{C}
 =
  \left [
   \begin{array}{cccc}
     \frac{2(x_1 - x_s)}{v d_1} & \frac{2(x_2 - x_s)}{v d_2} & \cdots & \frac{2(x_{n_{obs}} - x_s)}{v d_{n_{obs}}} \\
     \frac{2(y_1 - y_s)}{v d_1} & \frac{2(y_2 - y_s)}{v d_2} & \cdots & \frac{2(y_{n_{obs}} - y_s)}{v d_{n_{obs}}} \\
     -2 & -2 & \cdots & -2
   \end{array}
  \right ]
  \left \{
    \begin{array}{c}
      T_{obs}^{(1)} - T_{est}^{(1)} \\
      T_{obs}^{(2)} - T_{est}^{(2)} \\
      \vdots \\
      T_{obs}^{(n_{rec})} - T_{est}^{(n_{obs})} \\
    \end{array}
  \right \}
 =\left [
   \begin{array}{cccc}
     \frac{2(x_1 - x_s)}{v d_1} & \frac{2(x_2 - x_s)}{v d_2} & \cdots & \frac{2(x_{n_{obs}} - x_s)}{v d_{n_{obs}}} \\
     \frac{2(y_1 - y_s)}{v d_1} & \frac{2(y_2 - y_s)}{v d_2} & \cdots & \frac{2(y_{n_{obs}} - y_s)}{v d_{n_{obs}}} \\
     -2 & -2 & \cdots & -2
   \end{array}
  \right ]
  \left \{
    \begin{array}{c}
      r_1 \\
      r_2 \\
      \vdots \\
      r_{n_{obs}} \\
    \end{array}
  \right \}\end{equation}

This form has a well-known form and is compactly written as something we've already seen

\begin{equation}
  \nabla \mathcal{C} = J^T \textbf{r}
\end{equation}

Again, the matrix, $J$, is the Jacobian matrix and gives us insights into the sensitivity of the model parameters to different observations.  The Jacobian matrix is essential in higher-order methods with faster convegence like Gauss-Newton optimization are second order.  Recall, the Gauss-Newton search direction, $\textbf{p}$, is found by computing

\begin{equation}
  J^T J \textbf{p} =-\nabla J^T \textbf{r}
\end{equation}

or

\begin{equation}
  \textbf{p} =-(J^T J)^{-1} J^T \textbf{r}
\end{equation}

with model update given by $\textbf{m}_{k+1} = \textbf{m}_k + \alpha \textbf{p}$.  For (Gauss)-Newton steps $\alpha$ is typically $1$.  The algorithm should try small values of $\alpha$ should $\alpha=1$ prove to not reduce the objective function - i.e., we've stepped too far.  Something to keep in mind, $J$ is usually ill-conditioned. Computing $J^T J$ then squares the condition number of $J$.  Hence, inverting $J^T J$ is very difficult.  It is advised to use the QR decomposition or SVD instead of the Cholesky decomposition of the normal equations.  Shifting the small singular values (eigenvalues) of $J^T J$ away from zero by adding a diagonal regularization matrix like $J^T J + \gamma I$ is a common strategy to mitigate the dramatic effect that small singular values can have on the solution.

As shown previously, if we set $J^T J$ to the identity matrix, $I$, so that $(J^T J)^{-1} = I^{-1} = I$, then we obtain the steepest-descent method from the Gauss-Newton method as the search direction is simply

\begin{equation}
  \textbf{p} =-J^T \textbf{r}
\end{equation}

Note, for problems of appreciable size the Jacobian matrix is expensive to store and explicity compute.  For this reason, like in full-waveform inversion, we typically compute the action of $J$ on a vector of residuals and avoid its explicit construction.  More sophisticated algorithms like conjugate-gradient or the the limited memory BFGS algorithm can approximate the inverse of $J^T J$ and considerably increase convergence rates for a very small cost in system random access memory (RAM).  For this reasion l-BFGS and conjugate gradient are typically used for large-scale optimization problems in seismology as they balance low-memory footprint of steepest-descent with the convergence rate of Gauss-Newton.

```python
def computeJacobian(model, xr, yr, vel):
    """
    Computes the Jacobian for the 2D travel time location in a homogeneous medium.
    The derivatives in the Jacobian are analytically computed.

    Parameters
    ----------
    model : Array of length 3
        This is an array containing the model.
    xr : array_like
        These are the x positions (meters) of each reciever.  This has dimension [nrec]
        which corresponds to the number of receivers and observations.
    yr : array_like
        These are the y positions (meters) of each receiver.  This has dimension [nrec]
        which corresponds to the number of receivers and observations.
    vel : double
        This is the constant velocity in the medium.

    Returns
    -------
    J : Matrix
        The [nobs x 3] Jacobian matrix.  Here, nobs is the number of observations which
        equals the number of receivers.  The gradient is readily obtained by computing
    """
    # Unpack the model
    xs, ys, t0 = unpackModel(model)
    # Allocate space for the Jacobian matrix
    nrows = len(xr)    # The number of rows in the Jacobian is the number of observations
    ncols = len(model) # The number of columns in the Jacobian is the dimension of the model spcae
    J = zeros([nrows, ncols])
    # Loop on observations
    for irow in range(nrows):
        # It's possible to get a division by zero.  Technically the singularity goes to infinity.
        # But, for numerical reasons I'll force it to 0.  The trickery here is that the denominator
        # will be shifted away from zero and the numerator will evaluate to 0
        d = max(1.e-14, sqrt( (xr[irow] - xs)**2 + (yr[irow] - ys)**2 )) # Can get division by 0
        J[irow,0] = # Put your code here for dC_{irow}/dx
        J[irow,1] = # Put your code here for dC_{irow}/dy
        J[irow,2] = # Put your code here for dC_{irow}/dt
    return J

def computeGradient(model, tobs, xr, yr, vel):
    """
    Computes the gradient of the objective function for the 2D
    homogeneous velocity model.

    Parameters
    ----------
    model : Array of length 3
        This is an array containing the model.
    tobs : array_like
        These are the observed travel times in second.  This has dimension [nrec]
        which corresponds to the number observations.
    xr : array_like
        These are the x positions (meters) of each reciever.  This has dimension [nrec]
        which corresponds to the number of receivers and observations.
    yr : array_like
        These are the y positions (meters) of each receiver.  This has dimension [nrec]
        which corresponds to the number of receivers and observations.
    vel : double
        This is the constant velocity in the medium.

    Returns
    -------
    grad : array_like
        The gradient of the objective function.  This has dimension [3] and is packed
        in the same order as the model so that m1 = m0 - alpha*grad.
    """
    # Compute the estimate travel times
    test = computeTravelTimes(xr, yr, vel, model)
    # Compute the residuals
    res = computeResiduals(tobs, test)
    # Compute the Jacobian with computeJacobian
    J = # Put your code here
    # Compute the gradient
    grad = J.T@res # The `at' symbol in numpy is shorthand for matrix-vector and matrix-matrix multiplication
    return grad

```

Below are a series of unit tests that you should pass before proceeding to the next section.  Realize that there is a much more general way to compute gradient entries without explicitly calculating derivatives.  To do this we appeal to finite differencing the cost function so that the gradient can be approximated by

\begin{equation}
  \nabla \mathcal{C}
  \approx
  \left \{
    \begin{array}{c}
      \frac{\mathcal{C}(\textbf{m} + \delta \textbf{e}_1) - \mathcal{C}(\textbf{m})}{\delta} \\
      \frac{\mathcal{C}(\textbf{m} + \delta \textbf{e}_2) - \mathcal{C}(\textbf{m})}{\delta} \\
      \frac{\mathcal{C}(\textbf{m} + \delta \textbf{e}_3) - \mathcal{C}(\textbf{m})}{\delta}
    \end{array}
  \right \}
\end{equation}

where $\textbf{e}_i$ is a unit vector in the $i$'th component of the model and $\delta$ some perturbation.  Realize, that choosing $\delta$ requires some trial and error.  There is a strange tension between a $\delta$ that is too large and does not well represent the concept of a derivative and a $\delta$ too small that succumbs to cancellation error (a form of numerical error).

```python
def unit_test():
    # Test the travel time function calculator
    xs = 2
    ys = 1
    t0 =-1
    # Distance of a (3,4,5) triangle is 5
    xr = [5, 5, 5]
    yr = [5, 5, 5]
    # All times should be -1 + 5/5 = 0
    vel = 5
    model = packModel(xs, ys, t0)
    test = computeTravelTimes(xr, yr, vel, model)
    assert( max(abs(test)) < 1.e-14 ), "Failed travel time calculation"
    print("Passed travel time calculation test")

    # Test the residual calculation: tobs - test = 1 - 0 = 1
    tobs = zeros(4) + 1
    test = zeros(4)
    res = computeResiduals(tobs, test)
    for i in range(len(res)):
        assert( abs(res[i] - 1.0) < 1.e-14), 'Failed residual test'
    print("Passed residual test")

    # Begin making a little more substantial test - emplace a real source
    xs_true = 1200
    ys_true = 1300
    t0_true = 5.0
    model_true = packModel(xs_true, ys_true, t0_true) # pack the model
    # Put the estimate source in the ballpark of the true source
    xs_est  = 1100
    ys_est  = 900
    ts_est  = 0.0
    model_est = packModel(xs_est, ys_est, ts_est) # pack the model
    # Surround the source with receivers
    xr = [0,    1000, 2000, 0,    1000, 2000, 0,    1000, 2000]
    yr = [0,    0,    0,    1000, 1000, 1000, 2000, 2000, 2000]
    vel = 5000.0
    tobs = computeTravelTimes(xr, yr, vel, model_true)
    test = computeTravelTimes(xr, yr, vel, model_est)
    cost = computeCostFunction(model_est, tobs, xr, yr, vel)
    r2 = 0
    for i in range(len(tobs)):
        r2 = r2 + (tobs[i] - test[i])**2
    assert(abs(cost - r2) < 1.e-10), "Failed to compute cost function"

    # check scipy's algebra
    J = computeJacobian(model_est, xr, yr, vel)
    res = computeResiduals(tobs, test)
    grad = computeGradient(model_est, tobs, xr, yr, vel)
    assert(max(abs(J.T@res - grad)) < 1.e-14), "Strange algebraic mistake"
    print("Passed grad = J^T*res test")
    # For more general cases where the travel times have no analytic oslution we would
    # use a finite difference.  Here, compute the trusty finite difference to verify
    # our algebra.
    cost = computeCostFunction(model_est, tobs, xr, yr, vel)
    pert = 1.e-7
    model_pert_x = model_est + [pert, 0, 0]
    model_pert_y = model_est + [0, pert, 0]
    model_pert_t = model_est + [0, 0, pert]
    # compute perturbed cost functions
    cost_pert_x = computeCostFunction(model_pert_x, tobs, xr, yr, vel)
    cost_pert_y = computeCostFunction(model_pert_y, tobs, xr, yr, vel)
    cost_pert_t = computeCostFunction(model_pert_t, tobs, xr, yr, vel)
    # compute the forward finite differnces
    dcdx = (cost_pert_x - cost)/pert
    dcdy = (cost_pert_y - cost)/pert
    dcdt = (cost_pert_t - cost)/pert
    grad_fd = [dcdx, dcdy, dcdt]
    # these are never that accurate
    assert(max(abs(grad_fd - grad)) < 1.e-4), "Gradient is wrong; Jacobian is incorrect"
    print("Passed residual test")

    return True

if (not unit_test()):
    print("Failed unit tests")
else:
    print("Passed unit tests")
```

    Passed travel time calculation test
    Passed residual test
    Passed grad = J^T*res test
    Passed residual test
    Passed unit tests

## Implement Geiger's Method

Here, the nonlinear inversion strategy is implemented with the gradient-descent technique and a more sophisticad Gauss-Newton step.

```python
# Now implement the optimization
def geiger(tobs, xr, yr, vel, maxit=100, alpha=1, tol=5.e-5, lgn = False, model_specified = None):
    # Initialize a solution
    if (model_specified is None):
        xs0 = average(xr) + 20
        ys0 = average(yr) + 20
        ts0 = average(tobs)
        model0 = packModel(xs0, ys0, ts0)
    else:
        model0 = copy(model_specified)
    models = []
    lconv = False
    # Begin iterative method
    for k in range(maxit):
        # Save the current model
        models.append(model0)
        # Compute the cost function with computeCostFunction and check if the method converged
        cost = # Put your code here
        print("Iteration %d, Objective function: %f (s^2), (xs,ys,ts)=(%f,%f,%f)"%(k, cost, model0[0], model0[1], model0[2]))
        if (cost < tol):
            print("Convergence achieved; terminating", cost)
            lconv = True
            break
        # Compute the new search direction either with a Gauss-Newton or gradient step
        if (lgn):
            # Gauss-Newton search direction from computeGaussNewtonSearchDirection
            p = # Put your code here
        else:
            p = computeSteepestDescentSearchDirection(model0, tobs, xr, yr, vel)
        # This is where the line search would go.  For now, recommend alpha=1 for Gauss
        # Newton and alpha << 1 for gradient-descent.
        model1 = model0 + alpha*p
        # Save model for final iteration
        model0 = copy(model1)
    # Check final convergence
    if (not lconv):
        cost = computeCostFunction(model0, tobs, xr, yr, vel)
        if (cost < tol):
            print("Convergence achieved", cost)
        else:
            print("Method did not converge to given tolerance")
    return models

# Begin making a little more substantial test - emplace a real source
xs_true = 1200
ys_true = 1300
t0_true = 5.0
model_true = packModel(xs_true, ys_true, t0_true) # pack the model
# Make up some initial source estimate
xs_est  = 400 #1100
ys_est  = 400# 900
ts_est  = 0.0
model_est = packModel(xs_est, ys_est, ts_est) # pack the model
# Surround the source with receivers
xr = [0,    1000, 2000, 0,    1000, 2000, 0,    1000, 2000]
yr = [0,    0,    0,    1000, 1000, 1000, 2000, 2000, 2000]
vel = 5000.0
# Compute the observed travel times
tobs = computeTravelTimes(xr, yr, vel, model_true)
# Apply Geiger's method with a gradient descent step
print("Geiger Gradient Descent")
models_grad = geiger(tobs, xr, yr, vel, maxit=20, alpha=1.e-1, lgn=False, model_specified=model_est)
print("Geiger Gauss Newton")
models_gn = geiger(tobs, xr, yr, vel, maxit=20, alpha=1, lgn=True, model_specified=model_est)
```

    Geiger Gradient Descent
    Iteration 0, Objective function: 222.036394 (s^2), (xs,ys,ts)=(400.000000,400.000000,0.000000)
    Iteration 1, Objective function: 142.178724 (s^2), (xs,ys,ts)=(399.999424,399.999427,8.936312)
    Iteration 2, Objective function: 91.069813 (s^2), (xs,ys,ts)=(399.999921,399.999927,1.787262)
    Iteration 3, Objective function: 58.360109 (s^2), (xs,ys,ts)=(399.999560,399.999568,7.506502)
    Iteration 4, Objective function: 37.425898 (s^2), (xs,ys,ts)=(399.999885,399.999896,2.931110)
    Iteration 5, Objective function: 24.028002 (s^2), (xs,ys,ts)=(399.999661,399.999675,6.591424)
    Iteration 6, Objective function: 15.453348 (s^2), (xs,ys,ts)=(399.999876,399.999893,3.663173)
    Iteration 7, Objective function: 9.965570 (s^2), (xs,ys,ts)=(399.999740,399.999760,6.005774)
    Iteration 8, Objective function: 6.453391 (s^2), (xs,ys,ts)=(399.999885,399.999907,4.131693)
    Iteration 9, Objective function: 4.205597 (s^2), (xs,ys,ts)=(399.999805,399.999830,5.630958)
    Iteration 10, Objective function: 2.767009 (s^2), (xs,ys,ts)=(399.999905,399.999933,4.431546)
    Iteration 11, Objective function: 1.846312 (s^2), (xs,ys,ts)=(399.999861,399.999892,5.391075)
    Iteration 12, Objective function: 1.257066 (s^2), (xs,ys,ts)=(399.999932,399.999966,4.623452)
    Iteration 13, Objective function: 0.879949 (s^2), (xs,ys,ts)=(399.999912,399.999948,5.237551)
    Iteration 14, Objective function: 0.638594 (s^2), (xs,ys,ts)=(399.999964,400.000003,4.746272)
    Iteration 15, Objective function: 0.484126 (s^2), (xs,ys,ts)=(399.999958,400.000000,5.139295)
    Iteration 16, Objective function: 0.385267 (s^2), (xs,ys,ts)=(399.999999,400.000044,4.824876)
    Iteration 17, Objective function: 0.321997 (s^2), (xs,ys,ts)=(400.000002,400.000050,5.076411)
    Iteration 18, Objective function: 0.281505 (s^2), (xs,ys,ts)=(400.000036,400.000086,4.875183)
    Iteration 19, Objective function: 0.255589 (s^2), (xs,ys,ts)=(400.000045,400.000098,5.036166)
    Method did not converge to given tolerance
    Geiger Gauss Newton
    Iteration 0, Objective function: 222.036394 (s^2), (xs,ys,ts)=(400.000000,400.000000,0.000000)
    Iteration 1, Objective function: 54.844127 (s^2), (xs,ys,ts)=(787.187988,833.894086,2.537039)
    Iteration 2, Objective function: 13.640387 (s^2), (xs,ys,ts)=(952.971538,1033.257485,3.781577)
    Iteration 3, Objective function: 3.366715 (s^2), (xs,ys,ts)=(1051.203950,1181.090681,4.396462)
    Iteration 4, Objective function: 0.839053 (s^2), (xs,ys,ts)=(1125.080016,1240.356071,4.699161)
    Iteration 5, Objective function: 0.209477 (s^2), (xs,ys,ts)=(1162.328049,1269.871869,4.849808)
    Iteration 6, Objective function: 0.052334 (s^2), (xs,ys,ts)=(1181.100231,1284.839511,4.924962)
    Iteration 7, Objective function: 0.013079 (s^2), (xs,ys,ts)=(1190.532891,1292.393519,4.962496)
    Iteration 8, Objective function: 0.003269 (s^2), (xs,ys,ts)=(1195.261985,1296.189950,4.981252)
    Iteration 9, Objective function: 0.000817 (s^2), (xs,ys,ts)=(1197.629858,1298.093242,4.990627)
    Iteration 10, Objective function: 0.000204 (s^2), (xs,ys,ts)=(1198.814643,1299.046184,4.995314)
    Iteration 11, Objective function: 0.000051 (s^2), (xs,ys,ts)=(1199.407250,1299.522982,4.997657)
    Iteration 12, Objective function: 0.000013 (s^2), (xs,ys,ts)=(1199.703607,1299.761464,4.998828)
    Convergence achieved; terminating 1.2768187063160215e-05

# Question

The gradient descent method does not converge to a satisfactory solution while Gauss-Newton converges very quickly.  Both methods use the same Jacobian.  Why is gradient descent is having difficulty?

__TYPE YOUR ANSWER HERE__

## Grid Search

An alternative strategy for locating an earthquake is to exhaustively search every point in a grid and tabulate the cost function.  Then, the earthquake location is assigned to the grid whose point yields the smallest objective function.  With modern computers exhaustive search strategies have replaced gradient-based non-linear optimization strategies because exhaustive searches do not fall into local minima and exhaustive searches naturally lead to uncertainty quantification.

### Optimizing For An Origin Time In The Grid Search
In a conventional grid-search earthquake location we look to minimize a likelihood function of the form

\begin{equation}
  \mathcal{L}(\textbf{m})
 \propto e^{-\sum_{j=1}^{n_{obs}} (T_{obs}^{(j)} - T_{est}^{(j)})^2 }
\end{equation}

Note that the argument of the exponential is the sum-squared residuals.  Additionally, $e^{-r^2}$ is simply a Gaussian.  This indicates that maximizing the likelihood of an earthquake location corresponds to minimizing the argument of the exponential - i.e., minimizing the residual.

In minimizing the likelihood function recall that the estimate travel times are a function of the origin time and the candidate source position.  In this instance we will simply specify source positions (i.e., search through a grid of source positions) and at each source position optimize for an origin time.  Thus, we optimize the likelihood function by computing the origin time, $t_0$, that corresonds to

\begin{equation}
  \frac{\partial \mathcal{L}(t_0, \textbf{x}_s)}{\partial t_0} = 0
\end{equation}

To simplify the algebra, let's call the estimate travel time $T_{est}^{(j)} = t_0 + \mathcal{T}^{(j)}$ where $\mathcal{T}^{(j)}$, is a solution to the eikonal equation.  In this instance,
$\mathcal{T}^{(j)} = \frac{\sqrt{ (x_j - x_s)^2 + (y_j - y_s)^2 }}{v}$,

\begin{align*}
   \frac{\partial \mathcal{L}(t_0, \textbf{x}_s)}{\partial t_0}
 &=\frac{\partial }{\partial t_0}
    \left [
       e^{-\sum_{j=1}^{n_{obs}} \left (T_{obs}^{(j)} - (t_0 + \mathcal{T}^{(j)}) \right )^2 }
    \right ] \\
 &=-e^{-\sum_{j=1}^{n_{obs}} \left (T_{obs}^{(j)} - (t_0 + \mathcal{T}^{(j)}) \right )^2 }
   \frac{\partial }{\partial t_)}
   \left [
     \left (T_{obs}^{(j)} - (t_0 + \mathcal{T}^{(j)}) \right )^2
   \right ] \\
 &=2e^{-\sum_{j=1}^{n_{obs}} \left (T_{obs}^{(j)} - (t_0 + \mathcal{T}^{(j)}) \right )^2 } \cdot
   \left (
     \sum_{j=1}^{n_{obs}} \left (T_{obs}^{(j)} - (t_0 + \mathcal{T}^{(j)}) \right )
   \right )
\end{align*}

Next, the above is set to $0$,
\begin{equation}
  2e^{-\sum_{j=1}^{n_{obs}} \left (T_{obs}^{(j)} - (t_0 + \mathcal{T}^{(j)}) \right )^2 } \cdot
   \left (
     \sum_{j=1}^{n_{obs}} \left (T_{obs}^{(j)} - (t_0 + \mathcal{T}^{(j)}) \right )
   \right ) = 0
\end{equation}

The exponential is always greater than $0$, so we instead solve for $t_0$ in the second term
\begin{equation}
 \sum_{j=1}^{n_{obs}} \left (T_{obs}^{(j)} - (t_0 + \mathcal{T}^{(j)}) \right ) = 0
\end{equation}

Using linearity of the summation operator,
\begin{align*}
0&=\sum_{j=1}^{n_{obs}} \left (T_{obs}^{(j)} - (t_0 + \mathcal{T}^{(j)}) \right ) \\
 &=\sum_{j=1}^{n_{obs}} \left (T_{obs}^{(j)} -  \mathcal{T}^{(j)}) \right ) - \sum_{j=1}^{n_{obs}} t_0 \\
 &=\sum_{j=1}^{n_{obs}} \left (T_{obs}^{(j)} -  \mathcal{T}^{(j)}) \right ) - n_{obs} t_0 \\
\end{align*}

Finally, the optimal origin time is found by solving for $t_0$ so that
\begin{equation}
  t_0 = \frac{\sum_{j=1}^{n_{obs}} \left (T_{obs}^{(j)} -  \mathcal{T}^{(j)}) \right )}{n_{obs}}
\end{equation}

This says that the origin time is simply the weighted sum of residual travel times.  The algorithm is then as follows,

   1. Define a grid of locations (candidate sources).
   2. At each grid point compute the estimate travel times with a $0$ origin time.
   2. Optimize for the origin time by computing the average residual.
   3. Add the origin time in the travel time estimates and tabulate the sum of the squared residuals.

```python
def gridSearch(tobs, xr, yr, x, y, vel, objfn = 'l2'):
    # Tabulate the x's and y's in the grid
    nx = len(x)
    ny = len(y)
    nxy = nx*ny
    xv, yv = meshgrid(x, y, sparse=False, indexing='ij')
    # Set space for ther esult
    t0s = zeros(xv.shape)
    costfn = zeros(xv.shape)
    # Compute the travel times
    t0s = zeros(xv.shape)
    for ix in range(nx):
        for iy in range(ny):
            # Pack the model
            model = packModel(xv[ix,iy], yv[ix,iy], t0s[ix,iy])
            # Compute the travel times
            test = computeTravelTimes(xr, yr, vel, model)
            # Optimize for the origin time, i.e., the average residual
            res = computeResiduals(tobs, test)
            if (objfn == 'l2'):
                t0s[ix,iy] = average(res)
            else:
                # graduate students optimize l1 objective function
                t0s[ix,iy] = median(res)
            # Update the origin times with the origin time
            testWithOT = test + t0s[ix,iy]
            # Tabulate the cost function
            res = computeResiduals(tobs, testWithOT)
            if (objfn == 'l2'):
                costfn[ix,iy] = sum(res**2)
            else:
                # graduate students write l1 objective function
                costfn[ix,iy] = sum(abs(res))
        # Loop on y grid points
    # Loop on x grid points
    # This is a likelihood function so the area under the curve should sum to unity
    norm = sum(exp(-costfn.flatten()))
    likelihood = exp(-costfn)/norm # For numerical reasons this is not a good idea; typically use log likelihood
    return t0s, costfn, likelihood

x = linspace(-500,2500,301)
y = linspace(-500,2500,301)
t0s, costfn, likelihood = gridSearch(tobs, xr, yr, x, y, vel, objfn='l2')
# Get the optimum location
[ixopt, iyopt] = unravel_index(argmin(costfn), dims=[len(x), len(y)])
print("Estimate optimum (x,y,t0)=(%f,%f,%f); true optimum (x,y)=(%f,%f,%f)"%
      (x[ixopt], y[iyopt], t0s[ixopt,iyopt], xs_true, ys_true, t0_true))
```

    Estimate optimum (x,y,t0)=(1200.000000,1300.000000,5.000000); true optimum (x,y)=(1200.000000,1300.000000,5.000000)

```python

```

```python
# Plot the results from the Gauss-Newton and grid-search optimization
plt.figure(figsize=(10,10))
plt.imshow(likelihood.T, extent=[min(x), max(x), min(y), max(y)], origin='lower')
plt.xlabel('x offset (m)')
plt.ylabel('y offset (m)')
plt.title("Grid Search Location and Gauss-Newton Convergence History")
plt.grid(True)
plt.colorbar()
# We typically plot stations as (inverted) triangles
plt.scatter(xr, yr, c='red', marker='v', s=60)
# Extract the convergence history
niter = len(models_gn)
xconv = zeros(niter)
yconv = zeros(niter)
for i in range(niter):
    xconv[i] = models_gn[i][0]
    yconv[i] = models_gn[i][1]
# Plot the convergence history
plt.plot(xconv, yconv, c='black', linewidth=0.5)
plt.scatter(xconv, yconv, c='blue', marker='o', s=10)
# Plot the true source location as a black star
plt.scatter(xs_true, ys_true, c='black', marker='*', s=120)
plt.show()
```

# Questions

Does the Gauss-Newton method or grid-search take less computational time?

In this instance, both methods identified the true source location.  Do you find the Gauss-Newton method more informative than the result of the grid-search?

Notice that the Gauss-Newton method quickly converges on a minimum.  For more interesting acquisition geometries and velocity models the objective function can have multiple minima.  In general, would you prefer Gauss-Newton or grid-search?

## Graduate Students

Rewrite the grid-search function to optimize an $L_1$ misfit function.  In this case the origin time optimization amounts to computing the median residual travel time.  Plot the result.  Comment on the robustness of $L_1$ vs $L_2$ norms.

```python

```


## Markov switching autoregression models

This notebook provides an example of the use of Markov switching models in statsmodels to replicate a number of results presented in Kim and Nelson (1999). It applies the Hamilton (1989) filter the Kim (1994) smoother.

This is tested against the Markov-switching models from E-views 8, which can be found at http://www.eviews.com/EViews8/ev8ecswitch_n.html#MarkovAR or the Markov-switching models of Stata 14 which can be found at http://www.stata.com/manuals14/tsmswitch.pdf.

```python
%matplotlib inline

import numpy as np
import pandas as pd
import statsmodels.api as sm
import matplotlib.pyplot as plt
import requests
from io import BytesIO

# NBER recessions
from pandas_datareader.data import DataReader
from datetime import datetime
usrec = DataReader('USREC', 'fred', start=datetime(1947, 1, 1), end=datetime(2013, 4, 1))
```

### Hamilton (1989) switching model of GNP

This replicates Hamilton's (1989) seminal paper introducing Markov-switching models. The model is an autoregressive model of order 4 in which the mean of the process switches between two regimes. It can be written:

$$
y_t = \mu_{S_t} + \phi_1 (y_{t-1} - \mu_{S_{t-1}}) + \phi_2 (y_{t-2} - \mu_{S_{t-2}}) + \phi_3 (y_{t-3} - \mu_{S_{t-3}}) + \phi_4 (y_{t-4} - \mu_{S_{t-4}}) + \varepsilon_t
$$

Each period, the regime transitions according to the following matrix of transition probabilities:

$$ P(S_t = s_t | S_{t-1} = s_{t-1}) =
\begin{bmatrix}
p_{00} & p_{10} \\
p_{01} & p_{11}
\end{bmatrix}
$$

where $p_{ij}$ is the probability of transitioning *from* regime $i$, *to* regime $j$.

The model class is `MarkovAutoregression` in the time-series part of `statsmodels`. In order to create the model, we must specify the number of regimes with `k_regimes=2`, and the order of the autoregression with `order=4`. The default model also includes switching autoregressive coefficients, so here we also need to specify `switching_ar=False` to avoid that.

After creation, the model is `fit` via maximum likelihood estimation. Under the hood, good starting parameters are found using a number of steps of the expectation maximization (EM) algorithm, and a quasi-Newton (BFGS) algorithm is applied to quickly find the maximum.

```python
# Get the RGNP data to replicate Hamilton
dta = pd.read_stata('https://www.stata-press.com/data/r14/rgnp.dta').iloc[1:]
dta.index = pd.DatetimeIndex(dta.date, freq='QS')
dta_hamilton = dta.rgnp

# Plot the data
dta_hamilton.plot(title='Growth rate of Real GNP', figsize=(12,3))

# Fit the model
mod_hamilton = sm.tsa.MarkovAutoregression(dta_hamilton, k_regimes=2, order=4, switching_ar=False)
res_hamilton = mod_hamilton.fit()
```

```python
res_hamilton.summary()
```

Markov Switching Model Results

  Dep. Variable:	           rgnp	           No. Observations:  	    131

  Model:	           MarkovAutoregression	   Log Likelihood     	 -181.263

  Date:	              Tue, 02 Feb 2021	     AIC                	  380.527

  Time:	                  06:52:20	         BIC                	  406.404

  Sample:	               04-01-1951	        HQIC               	  391.042

  	                     - 10-01-1984

  Covariance Type:	        approx

Regime 0 parameters

    	       coef	     std err	      z	      P>|z|	  [0.025	    0.975]

  const	    -0.3588	     0.265	    -1.356	  0.175	    -0.877	     0.160

Regime 1 parameters

    	       coef	     std err	      z	      P>|z|	  [0.025	    0.975]

  const	     1.1635	     0.075	    15.614	  0.000	     1.017	     1.310

Non-switching parameters

     	       coef	     std err	      z	      P>|z|	  [0.025	    0.975]

  sigma2	     0.5914	     0.103	     5.761	  0.000	     0.390	     0.793

  ar.L1	      0.0135	     0.120	     0.112	  0.911	    -0.222	     0.249

  ar.L2	     -0.0575	     0.138	    -0.418	  0.676	    -0.327	     0.212

  ar.L3	     -0.2470	     0.107	    -2.310	  0.021	    -0.457	    -0.037

  ar.L4	     -0.2129	     0.111	    -1.926	  0.054	    -0.430	     0.004

Regime transition parameters

     	        coef	     std err	      z	      P>|z|	  [0.025	    0.975]

  p[0->0]	     0.7547	     0.097	     7.819	  0.000	     0.565	     0.944

  p[1->0]	     0.0959	     0.038	     2.542	  0.011	     0.022	     0.170

Warnings:
[1] Covariance matrix calculated using numerical (complex-step) differentiation.

We plot the filtered and smoothed probabilities of a recession. Filtered refers to an estimate of the probability at time $t$ based on data up to and including time $t$ (but excluding time $t+1, ..., T$). Smoothed refers to an estimate of the probability at time $t$ using all the data in the sample.

For reference, the shaded periods represent the NBER recessions.

```python
fig, axes = plt.subplots(2, figsize=(7,7))
ax = axes[0]
ax.plot(res_hamilton.filtered_marginal_probabilities[0])
ax.fill_between(usrec.index, 0, 1, where=usrec['USREC'].values, color='k', alpha=0.1)
ax.set_xlim(dta_hamilton.index[4], dta_hamilton.index[-1])
ax.set(title='Filtered probability of recession')

ax = axes[1]
ax.plot(res_hamilton.smoothed_marginal_probabilities[0])
ax.fill_between(usrec.index, 0, 1, where=usrec['USREC'].values, color='k', alpha=0.1)
ax.set_xlim(dta_hamilton.index[4], dta_hamilton.index[-1])
ax.set(title='Smoothed probability of recession')

fig.tight_layout()
```

From the estimated transition matrix we can calculate the expected duration of a recession versus an expansion.

```python
print(res_hamilton.expected_durations)
```

    [ 4.07604747 10.42589381]

In this case, it is expected that a recession will last about one year (4 quarters) and an expansion about two and a half years.

### Kim, Nelson, and Startz (1998) Three-state Variance Switching

This model demonstrates estimation with regime heteroskedasticity (switching of variances) and no mean effect. The dataset can be reached at http://econ.korea.ac.kr/~cjkim/MARKOV/data/ew_excs.prn.

The model in question is:

$$
\begin{align}
y_t & = \varepsilon_t \\
\varepsilon_t & \sim N(0, \sigma_{S_t}^2)
\end{align}
$$

Since there is no autoregressive component, this model can be fit using the `MarkovRegression` class. Since there is no mean effect, we specify `trend='nc'`. There are hypothesized to be three regimes for the switching variances, so we specify `k_regimes=3` and `switching_variance=True` (by default, the variance is assumed to be the same across regimes).

```python
# Get the dataset
ew_excs = requests.get('http://econ.korea.ac.kr/~cjkim/MARKOV/data/ew_excs.prn').content
raw = pd.read_table(BytesIO(ew_excs), header=None, skipfooter=1, engine='python')
raw.index = pd.date_range('1926-01-01', '1995-12-01', freq='MS')

dta_kns = raw.loc[:'1986'] - raw.loc[:'1986'].mean()

# Plot the dataset
dta_kns[0].plot(title='Excess returns', figsize=(12, 3))

# Fit the model
mod_kns = sm.tsa.MarkovRegression(dta_kns, k_regimes=3, trend='nc', switching_variance=True)
res_kns = mod_kns.fit()
```

```python
res_kns.summary()
```

Markov Switching Model Results

  Dep. Variable:	           0	          No. Observations:  	    732

  Model:	           MarkovRegression	   Log Likelihood     	 1001.895

  Date:	            Tue, 02 Feb 2021	   AIC                	 -1985.790

  Time:	                06:52:26	       BIC                	 -1944.428

  Sample:	             01-01-1926	      HQIC               	 -1969.834

  	                   - 12-01-1986

  Covariance Type:	      approx

Regime 0 parameters

     	       coef	     std err	      z	      P>|z|	  [0.025	    0.975]

  sigma2	     0.0012	     0.000	     7.136	  0.000	     0.001	     0.002

Regime 1 parameters

     	       coef	     std err	      z	      P>|z|	  [0.025	    0.975]

  sigma2	     0.0040	     0.000	     8.489	  0.000	     0.003	     0.005

Regime 2 parameters

     	       coef	     std err	      z	      P>|z|	  [0.025	    0.975]

  sigma2	     0.0311	     0.006	     5.461	  0.000	     0.020	     0.042

Regime transition parameters

     	        coef	     std err	      z	      P>|z|	  [0.025	    0.975]

  p[0->0]	     0.9747	     0.000	  7857.416	  0.000	     0.974	     0.975

  p[1->0]	     0.0195	     0.010	     1.949	  0.051	    -0.000	     0.039

  p[2->0]	  2.354e-08	       nan	       nan	    nan	       nan	       nan

  p[0->1]	     0.0253	  3.97e-05	   637.835	  0.000	     0.025	     0.025

  p[1->1]	     0.9688	     0.013	    75.528	  0.000	     0.944	     0.994

  p[2->1]	     0.0493	     0.032	     1.551	  0.121	    -0.013	     0.112

Warnings:
[1] Covariance matrix calculated using numerical (complex-step) differentiation.

Below we plot the probabilities of being in each of the regimes; only in a few periods is a high-variance regime probable.

```python
fig, axes = plt.subplots(3, figsize=(10,7))

ax = axes[0]
ax.plot(res_kns.smoothed_marginal_probabilities[0])
ax.set(title='Smoothed probability of a low-variance regime for stock returns')

ax = axes[1]
ax.plot(res_kns.smoothed_marginal_probabilities[1])
ax.set(title='Smoothed probability of a medium-variance regime for stock returns')

ax = axes[2]
ax.plot(res_kns.smoothed_marginal_probabilities[2])
ax.set(title='Smoothed probability of a high-variance regime for stock returns')

fig.tight_layout()
```

### Filardo (1994) Time-Varying Transition Probabilities

This model demonstrates estimation with time-varying transition probabilities. The dataset can be reached at http://econ.korea.ac.kr/~cjkim/MARKOV/data/filardo.prn.

In the above models we have assumed that the transition probabilities are constant across time. Here we allow the probabilities to change with the state of the economy. Otherwise, the model is the same Markov autoregression of Hamilton (1989).

Each period, the regime now transitions according to the following matrix of time-varying transition probabilities:

$$ P(S_t = s_t | S_{t-1} = s_{t-1}) =
\begin{bmatrix}
p_{00,t} & p_{10,t} \\
p_{01,t} & p_{11,t}
\end{bmatrix}
$$

where $p_{ij,t}$ is the probability of transitioning *from* regime $i$, *to* regime $j$ in period $t$, and is defined to be:

$$
p_{ij,t} = \frac{\exp\{ x_{t-1}' \beta_{ij} \}}{1 + \exp\{ x_{t-1}' \beta_{ij} \}}
$$

Instead of estimating the transition probabilities as part of maximum likelihood, the regression coefficients $\beta_{ij}$ are estimated. These coefficients relate the transition probabilities to a vector of pre-determined or exogenous regressors $x_{t-1}$.

```python
# Get the dataset
filardo = requests.get('http://econ.korea.ac.kr/~cjkim/MARKOV/data/filardo.prn').content
dta_filardo = pd.read_table(BytesIO(filardo), sep=' +', header=None, skipfooter=1, engine='python')
dta_filardo.columns = ['month', 'ip', 'leading']
dta_filardo.index = pd.date_range('1948-01-01', '1991-04-01', freq='MS')

dta_filardo['dlip'] = np.log(dta_filardo['ip']).diff()*100
# Deflated pre-1960 observations by ratio of std. devs.
# See hmt_tvp.opt or Filardo (1994) p. 302
std_ratio = dta_filardo['dlip']['1960-01-01':].std() / dta_filardo['dlip'][:'1959-12-01'].std()
dta_filardo['dlip'][:'1959-12-01'] = dta_filardo['dlip'][:'1959-12-01'] * std_ratio

dta_filardo['dlleading'] = np.log(dta_filardo['leading']).diff()*100
dta_filardo['dmdlleading'] = dta_filardo['dlleading'] - dta_filardo['dlleading'].mean()

# Plot the data
dta_filardo['dlip'].plot(title='Standardized growth rate of industrial production', figsize=(13,3))
plt.figure()
dta_filardo['dmdlleading'].plot(title='Leading indicator', figsize=(13,3));
```

The time-varying transition probabilities are specified by the `exog_tvtp` parameter.

Here we demonstrate another feature of model fitting - the use of a random search for MLE starting parameters. Because Markov switching models are often characterized by many local maxima of the likelihood function, performing an initial optimization step can be helpful to find the best parameters.

Below, we specify that 20 random perturbations from the starting parameter vector are examined and the best one used as the actual starting parameters. Because of the random nature of the search, we seed the random number generator beforehand to allow replication of the result.

```python
mod_filardo = sm.tsa.MarkovAutoregression(
    dta_filardo.iloc[2:]['dlip'], k_regimes=2, order=4, switching_ar=False,
    exog_tvtp=sm.add_constant(dta_filardo.iloc[1:-1]['dmdlleading']))

np.random.seed(12345)
res_filardo = mod_filardo.fit(search_reps=20)
```

```python
res_filardo.summary()
```

Markov Switching Model Results

  Dep. Variable:	           dlip	           No. Observations:  	    514

  Model:	           MarkovAutoregression	   Log Likelihood     	 -586.572

  Date:	              Tue, 02 Feb 2021	     AIC                	 1195.144

  Time:	                  06:52:39	         BIC                	 1241.808

  Sample:	               03-01-1948	        HQIC               	 1213.433

  	                     - 04-01-1991

  Covariance Type:	        approx

Regime 0 parameters

    	       coef	     std err	      z	      P>|z|	  [0.025	    0.975]

  const	    -0.8659	     0.153	    -5.658	  0.000	    -1.166	    -0.566

Regime 1 parameters

    	       coef	     std err	      z	      P>|z|	  [0.025	    0.975]

  const	     0.5173	     0.077	     6.706	  0.000	     0.366	     0.668

Non-switching parameters

     	       coef	     std err	      z	      P>|z|	  [0.025	    0.975]

  sigma2	     0.4844	     0.037	    13.172	  0.000	     0.412	     0.556

  ar.L1	      0.1895	     0.050	     3.761	  0.000	     0.091	     0.288

  ar.L2	      0.0793	     0.051	     1.552	  0.121	    -0.021	     0.180

  ar.L3	      0.1109	     0.052	     2.136	  0.033	     0.009	     0.213

  ar.L4	      0.1223	     0.051	     2.418	  0.016	     0.023	     0.221

Regime transition parameters

        	           coef	     std err	      z	      P>|z|	  [0.025	    0.975]

  p[0->0].tvtp0	     1.6494	     0.446	     3.702	  0.000	     0.776	     2.523

  p[1->0].tvtp0	    -4.3595	     0.747	    -5.833	  0.000	    -5.824	    -2.895

  p[0->0].tvtp1	    -0.9945	     0.566	    -1.758	  0.079	    -2.103	     0.114

  p[1->0].tvtp1	    -1.7702	     0.508	    -3.484	  0.000	    -2.766	    -0.775

Warnings:
[1] Covariance matrix calculated using numerical (complex-step) differentiation.

Below we plot the smoothed probability of the economy operating in a low-production state, and again include the NBER recessions for comparison.

```python
fig, ax = plt.subplots(figsize=(12,3))

ax.plot(res_filardo.smoothed_marginal_probabilities[0])
ax.fill_between(usrec.index, 0, 1, where=usrec['USREC'].values, color='gray', alpha=0.2)
ax.set_xlim(dta_filardo.index[6], dta_filardo.index[-1])
ax.set(title='Smoothed probability of a low-production state');
```

Using the time-varying transition probabilities, we can see how the expected duration of a low-production state changes over time:

```python
res_filardo.expected_durations[0].plot(
    title='Expected duration of a low-production state', figsize=(12,3));
```

During recessions, the expected duration of a low-production state is much higher than in an expansion.