# 実験B

```python
import numpy as np
import scipy as sp
import scipy.constants as const
import matplotlib.pyplot as plt
import japanize_matplotlib
import pandas as pd
import sympy as sym
x=[1,2,3,4,5,6,7]
y=[0.085, 0.3, 0.77, 0.41, 0.33, 0.15, 0.11]
label = ['10:0', '9:1', '8:2', '7:3', '6:4', '5:5', '4:6']
plt.bar(x, y)
plt.ylabel('吸光度')
plt.xlabel('A液、B液の体積比')
plt.xticks(x, label)
plt.savefig('1.pdf')
plt.show()
```

```python
from scipy import optimize
from sklearn.linear_model import LinearRegression as LR
```

```python
bunsi = [10,15,20, 25, 37, 50, 75, 100, 150, 250][::-1]
kyori = np.array([22, 41, 67, 95, 155, 215, 318, 372, 457, 526])
x = np.array([85, 289, 437]) / kyori[0]
kyori = kyori / kyori[0]
reg = LR().fit(kyori[3:10].reshape((-1, 1)), np.log(bunsi)[3:10])
plt.yscale('log')
plt.xlim(0, 25)
plt.ylim(9, 300)
plt.scatter(kyori, bunsi)
X = np.linspace(3.5, 30)
plt.plot(X, np.exp(reg.predict(X.reshape((-1, 1)))))
plt.xlabel('相対移動度')
plt.ylabel('分子量 [kDa]')
plt.savefig('kenryo.pdf')
plt.show()
```

```python
np.exp(reg.predict(x.reshape((-1, 1))))
```

    array([71.11821952, 29.00111299, 15.12823281])

```python
a = {}
for i in range(len(bunsi) +  1):
    for j in range(i + 2, len(bunsi) + 1):
        a[(i, j)] = LR().fit(kyori[2:].reshape((-1, 1)), np.log(bunsi)[2:]).score(kyori[i:j].reshape((-1, 1)), np.log(bunsi)[i:j])
```

```python
a[max(a, key=a.get)]
```

    0.9861703447030934

```python
max(a, key=a.get)
```

    (3, 10)

```python
a
```

    {(0, 2): -6.511948431152308,
     (0, 3): -1.3948796944155566,
     (0, 4): -0.23260129325809187,
     (0, 5): 0.33798663982872185,
     (0, 6): 0.5817137178939801,
     (0, 7): 0.731913066101846,
     (0, 8): 0.8073015109453174,
     (0, 9): 0.8550947200493135,
     (0, 10): 0.8928892388618701,
     (1, 3): -1.8952481384425446,
     (1, 4): 0.0177726851430261,
     (1, 5): 0.6118111307585004,
     (1, 6): 0.7813661404896894,
     (1, 7): 0.8765860380926409,
     (1, 8): 0.9171181331468986,
     (1, 9): 0.9395165121680846,
     (1, 10): 0.9577125545000336,
     (2, 4): 0.27673166274474637,
     (2, 5): 0.8267219142055817,
     (2, 6): 0.8973470943326054,
     (2, 7): 0.949779654569567,
     (2, 8): 0.9685813532119547,
     (2, 9): 0.9759480200080157,
     (2, 10): 0.9841659604460699,
     (3, 5): 0.8499616944367776,
     (3, 6): 0.8827232538542915,
     (3, 7): 0.9518166920593597,
     (3, 8): 0.9719821665034415,
     (3, 9): 0.9774365157154317,
     (3, 10): 0.9861703447030934,
     (4, 6): 0.35472436602189855,
     (4, 7): 0.8715755432372102,
     (4, 8): 0.9373663615423117,
     (4, 9): 0.956083904576191,
     (4, 10): 0.9765649532162581,
     (5, 7): 0.753664360733005,
     (5, 8): 0.901514731645618,
     (5, 9): 0.934939058894631,
     (5, 10): 0.9711619611133956,
     (6, 8): 0.9216759562391929,
     (6, 9): 0.9149127975771352,
     (6, 10): 0.9758876984519066,
     (7, 9): 0.773219709995751,
     (7, 10): 0.9606490306330246,
     (8, 10): 0.8860065777957079}

# 実験A

```python
# 試験管番号、吸光度洗浄あり、洗浄なし
idx, a1, a2 = list(zip(*[
[16, 0.044, 0.045],
[15, 0.067, 0.061],
[14, 0.069, 0.072],
[13, 0.098, 0.099],
[12, 0.153, 0.172],
[11, 0.233, 0.288],
[10, 0.372, 0.445],
[9, 0.59, 0.716],
[8, 0.849, 0.918],
[7, 1.014, 0.877],
[6, 1.176, 0.595],
[5, 1.218, 0.411],
[4, 1.209, 0.299],
[3, 1.226, 0.22],
[2, 1.219, 0.157],
[1, 1.241, 0.118],
]))
# 濃度 ug / mL
c = 32 / (2 ** np.array(idx))
c[0] = 1e-10
```

```python
c
```

    array([1.000000e-10, 9.765625e-04, 1.953125e-03, 3.906250e-03,
           7.812500e-03, 1.562500e-02, 3.125000e-02, 6.250000e-02,
           1.250000e-01, 2.500000e-01, 5.000000e-01, 1.000000e+00,
           2.000000e+00, 4.000000e+00, 8.000000e+00, 1.600000e+01])

```python
plt.scatter(c, a1, label='洗浄あり')
# plt.plot(X, np.poly1d(np.polyfit(np.log(c[1:-7]), a1[1:-7], 2))(np.log(X)))
X = np.logspace(-4, 2)
sig = lambda x, a, b, c, d: d / (1 + np.exp(-a * np.log(x) - b)) + c
param = sp.optimize.curve_fit(sig, c, a1)[0]
plt.plot(X, sig(X, *param))

plt.scatter(c, a2, label='洗浄なし')

plt.xlabel('ヒトIgG濃度 [ug / mL]')
plt.ylabel('吸光度')
plt.xscale('log')
plt.xlim((1e-4, 1e2))
plt.ylim((0, 1.3))
plt.legend()
plt.savefig('A-kenryo.pdf')
plt.show()
```

```python
sig = lambda x, a, b, c, d: d / (1 + np.exp(-a * np.log(x) - b)) + c
param = sp.optimize.curve_fit(sig, c, a1)
param
```

    (array([1.22130073, 3.2227649 , 0.06091429, 1.18055615]),
     array([[ 2.32512190e-03,  6.08497694e-03,  2.39194056e-04,
             -4.57408298e-04],
            [ 6.08497694e-03,  1.78633283e-02,  4.48922993e-04,
             -1.18010142e-03],
            [ 2.39194056e-04,  4.48922993e-04,  8.78645868e-05,
             -1.03461178e-04],
            [-4.57408298e-04, -1.18010142e-03, -1.03461178e-04,
              1.95828179e-04]]))

```python

```


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

(For global files, output glued from miltiprocessor runs)

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 = 360  # number of x points in domain
ny = 360  # number of y points in domain
nz = 35   # number of z levels in domain
nt = 10  # number of output files (n time)
delt= 12*3600.0 # time interval between model output
```

Load output filenames:

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

filename = '/ocean/kramosmu/MITgcm/CanyonUpwelling/360x360x90_BodyForcing_6Tr_LinProfiles/run02/tave.0000000000.glob.nc'#'/ocean/kramosmu/MITgcm/CanyonUpwelling/120x90x35_UpwLinRampWind_1Tr_LinProfiles/run30/output_0001/state.0000000000.t001.nc'
StateOut = Dataset(filename)
#for dimobj in StateOut.variables.values():
#    print dimobj

filename2 = '/ocean/kramosmu/MITgcm/CanyonUpwelling/360x360x90_BodyForcing_1Tr/run01/mnc_0001/grid.t001.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/360x360x90_BodyForcing_6Tr_LinProfiles/run02/ptr_tave.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
tracer1 = getField(filename3,'Tr3')

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

MASK = np.ma.getmask(tracer1msk)

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

zlev = 20
time = 0

plt.rcParams.update({'font.size':13})
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
Dom=0.0
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):

        Dom = Dom + hFacC[:,jj,ii]*drF[:]*rAC[jj,ii]

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

```

     The total volume is 6.102853e+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 5.674921e+02

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

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

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

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

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

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

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

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

    volume-weighted mean of q^2 nt=9 is 5.544492e+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,1:nx-1]))**2

    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,:-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 3.168370e-06

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    The v-w mean of (dq/dy)^2 at nt=9 is 8.438798e-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 jj in np.arange(ny):

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

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

    qmeanDz[tt] = (sum(qsumDz))/DomVolume

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

```

    Mean (dq/dz)^2 at nt=0 is nan

    Mean (dq/dz)^2 at nt=1 is nan

    Mean (dq/dz)^2 at nt=2 is nan

    Mean (dq/dz)^2 at nt=3 is nan

    Mean (dq/dz)^2 at nt=4 is nan

    Mean (dq/dz)^2 at nt=5 is nan

    Mean (dq/dz)^2 at nt=6 is nan

    Mean (dq/dz)^2 at nt=7 is nan

    Mean (dq/dz)^2 at nt=8 is nan

    Mean (dq/dz)^2 at nt=9 is nan

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*2,2.0),q2mean)

plt.xlabel('hours')

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*24.0 # 2hrs

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

y = 0.5*dqdt[3: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()

```

GLS Regression Results

  Dep. Variable:	            y	          R-squared:         	    1.000

  Model:	                   GLS	         Adj. R-squared:    	    0.999

  Method:	             Least Squares	    F-statistic:       	    4271.

  Date:	             Fri, 05 Dec 2014	   Prob (F-statistic):	 6.58e-06

  Time:	                 13:00:50	       Log-Likelihood:    	   67.959

  No. Observations:	           5	        AIC:               	   -131.9

  Df Residuals:	               3	        BIC:               	   -132.7

  Df Model:	                   2

   	     coef	     std err	      t	      P>|t|	 [95.0% Conf. Int.]

  x1	   -90.9932	    26.382	    -3.449	  0.041	  -174.954    -7.033

  x2	     0.0179	     0.005	     3.397	  0.043	     0.001     0.035

  Omnibus:	          nan	   Durbin-Watson:     	    2.713

  Prob(Omnibus):	    nan	   Jarque-Bera (JB):  	    0.703

  Skew:	          -0.279	   Prob(JB):          	    0.704

  Kurtosis:	       1.250	   Cond. No.          	 8.82e+06

```python

dt = 1*3600.0*24.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]
print(np.shape(dqdt))
print(np.shape(qmeanDh))

## 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()

```

    (8,)
    (10, 1)

GLS Regression Results

  Dep. Variable:	            y	          R-squared:         	    0.998

  Model:	                   GLS	         Adj. R-squared:    	    0.998

  Method:	             Least Squares	    F-statistic:       	    3227.

  Date:	             Fri, 05 Dec 2014	   Prob (F-statistic):	 1.37e-10

  Time:	                 13:00:51	       Log-Likelihood:    	   101.15

  No. Observations:	           8	        AIC:               	   -200.3

  Df Residuals:	               7	        BIC:               	   -200.2

  Df Model:	                   1

   	     coef	     std err	      t	      P>|t|	 [95.0% Conf. Int.]

  x1	    -1.4173	     0.025	   -56.807	  0.000	    -1.476    -1.358

  Omnibus:	        1.648	   Durbin-Watson:     	    0.492

  Prob(Omnibus):	  0.439	   Jarque-Bera (JB):  	    0.258

  Skew:	           0.440	   Prob(JB):          	    0.879

  Kurtosis:	       3.025	   Cond. No.          	     1.00

```python
dt = 1*3600.0*24.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()

```

GLS Regression Results

  Dep. Variable:	            y	          R-squared:         	    0.998

  Model:	                   GLS	         Adj. R-squared:    	    0.998

  Method:	             Least Squares	    F-statistic:       	    3788.

  Date:	             Fri, 05 Dec 2014	   Prob (F-statistic):	 7.85e-11

  Time:	                 13:00:53	       Log-Likelihood:    	   101.79

  No. Observations:	           8	        AIC:               	   -201.6

  Df Residuals:	               7	        BIC:               	   -201.5

  Df Model:	                   1

   	     coef	     std err	      t	      P>|t|	 [95.0% Conf. Int.]

  x1	    -0.0003	  4.58e-06	   -61.545	  0.000	    -0.000    -0.000

  Omnibus:	        2.684	   Durbin-Watson:     	    0.559

  Prob(Omnibus):	  0.261	   Jarque-Bera (JB):  	    0.651

  Skew:	           0.696	   Prob(JB):          	    0.722

  Kurtosis:	       3.121	   Cond. No.          	     1.00

```python

```

```python

```


```python
%matplotlib inline
```

# Multiway Cluster Robust DML

This example shows how the multiway cluster roboust DML (Chiang et al. 2020) can be implemented with the DoubleML
package.
Chiang et al. (2020) consider double-indexed data

\begin{equation}
\lbrace W_{ij}: i \in \lbrace 1, \ldots, N \rbrace, j \in \lbrace 1, \ldots, M \rbrace \rbrace
\end{equation}

and the partially linear IV regression model (PLIV)

$$\begin{aligned}
Y_{ij} = D_{ij} \theta_0 +  g_0(X_{ij}) + \epsilon_{ij}, & &\mathbb{E}(\epsilon_{ij} | X_{ij}, Z_{ij}) = 0, \\
Z_{ij} = m_0(X_{ij}) + v_{ij}, & &\mathbb{E}(v_{ij} | X_{ij}) = 0.
\end{aligned}$$

TODO: Add a few more details and the reference!
https://arxiv.org/pdf/1909.03489.pdf

```python
import numpy as np
import pandas as pd

import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
import seaborn as sns

from sklearn.model_selection import KFold, RepeatedKFold
from sklearn.base import clone

from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import LinearRegression

from doubleml import DoubleMLData, DoubleMLPLIV
from doubleml.double_ml_resampling import DoubleMLMultiwayResampling

from doubleml.datasets import make_pliv_multiway_cluster_CKMS2019
```

## Simulate multiway cluster data

We use the PLIV data generating process described in Section 4.1 of Chiang et al. (2020).

```python
# Set the simulation parameters
N = 25  # number of observations (first dimension)
M = 25  # number of observations (second dimension)
dim_X = 100  # dimension of X

obj_dml_data = make_pliv_multiway_cluster_CKMS2019(N, M, dim_X)
```

```python
# The data comes with multi index for rows (tuples with two entries)
obj_dml_data.data.head(30)
```

## Initialize the objects of class DoubleMLData and DoubleMLPLIV

```python
# Set machine learning methods for m & g
learner = RandomForestRegressor(max_depth=2, n_estimators=10)
ml_g = clone(learner)
ml_m = clone(learner)
ml_r = clone(learner)

# initialize the DoubleMLPLIV object
dml_pliv_obj = DoubleMLPLIV(obj_dml_data,
                            ml_g,
                            ml_m,
                            ml_r,
                            score='partialling out',
                            dml_procedure='dml1',
                            draw_sample_splitting=False)
```

## Split samples and transfer the sample splitting to the object

```python
K = 3  # number of folds
smpl_sizes = [N, M]
obj_dml_multiway_resampling = DoubleMLMultiwayResampling(K, smpl_sizes)
smpls_multi_ind, smpls_lin_ind = obj_dml_multiway_resampling.split_samples()

dml_pliv_obj.set_sample_splitting([smpls_lin_ind])
```

## Fit the model and show a summary

```python
dml_pliv_obj.fit()
print(dml_pliv_obj.summary)
```

## Visualization of sample splitting with tuple and linear indexing

```python
#discrete color scheme
x = sns.color_palette("RdBu_r", 7)
cMap = ListedColormap([x[0], x[3], x[6]])
plt.rcParams['figure.figsize'] = 15, 12
sns.set(font_scale=1.3)
```

### Visualize sample splitting with tuples (one plot per fold)

```python
for i_split, this_split_ind in enumerate(smpls_multi_ind):
    plt.subplot(K, K, i_split + 1)
    df = pd.DataFrame(np.zeros([N, M]))
    ind_array_train = np.array([*this_split_ind[0]])
    ind_array_test = np.array([*this_split_ind[1]])
    df.loc[ind_array_train[:, 0], ind_array_train[:, 1]] = -1.
    df.loc[ind_array_test[:, 0], ind_array_test[:, 1]] = 1.

    ax = sns.heatmap(df, cmap=cMap);
    ax.invert_yaxis();
    ax.set_ylim([0, M]);
    colorbar = ax.collections[0].colorbar
    colorbar.set_ticks([-0.667, 0, 0.667])
    if i_split % K == (K - 1):
        colorbar.set_ticklabels(['Nuisance', '', 'Score'])
    else:
        colorbar.set_ticklabels(['', '', ''])
```

### Visualize sample splitting with linear indexing (one column per fold)

```python
df = pd.DataFrame(np.zeros([N*M, K*K]))
for i_split, this_split_ind in enumerate(smpls_lin_ind):
    df.loc[this_split_ind[0], i_split] = -1.
    df.loc[this_split_ind[1], i_split] = 1.

ax = sns.heatmap(df, cmap=cMap);
ax.invert_yaxis();
ax.set_ylim([0, N*M]);
colorbar = ax.collections[0].colorbar
colorbar.set_ticks([-0.667, 0, 0.667])
colorbar.set_ticklabels(['Nuisance', '', 'Score'])
```


# Neuromatch Academy: Week 3, Day 2, Tutorial 1

# Neuronal Network Dynamics: Neural Rate Models

## Background

The brain is a complex system, not because it is composed of a large number of diverse types of neurons, but mainly because of how neurons are connected to each other. The brain is a very large network of densely interconnected neurons.

The activity of neurons is constantly evolving in time. For this reason, neurons can be modeled as dynamical systems. The dynamical system approach is only one of the many modeling approaches that computational neuroscientists have developed (other points of views include information processing, network science, and statistical models). How the dynamics of neuronal networks affect the representation and processing of information in the brain is an open question. However, signatures of altered brain dynamics present in many brain diseases (e.g., in epilepsy or Parkinson's disease) tell us that it is crucial to study neuronal dynamics if we want to understand the brain.

In this tutorial, we will simulate and study one of the simplest models of biological neuronal networks. Instead of modeling and simulating individual excitatory neurons (e.g., LIF models that you implemented yesterday), we will treat them as a single homogeneous population and approximate their dynamics using a single one-dimensional equation describing the evolution of their average spiking rate in time.

## Objectives
In this tutorial we will learn how to build a firing rate model of a single population of excitatory neurons.

Steps:
- Write the equation for the firing rate dynamics of a 1D excitatory population.
- Visualize the response of the population as a function of parameters such as threshold level and gain, using the frequency-current (F-I) curve.
- Numerically simulate the dynamics of the excitatory population and find the fixed points of the system.
- Investigate the stability of the fixed points by linearizing the dynamics around them.

# Setup

```python
# Imports
import matplotlib.pyplot as plt    # import matplotlib
import numpy as np                 # import numpy
import scipy.optimize as opt       # import root-finding algorithm
import ipywidgets as widgets       # interactive display
```

```python
#@title Figure Settings
%matplotlib inline

fig_w, fig_h = 6, 4
my_fontsize = 16
my_params = {'axes.labelsize': my_fontsize,
          'axes.titlesize': my_fontsize,
          'figure.figsize': [fig_w, fig_h],
          'font.size': my_fontsize,
          'legend.fontsize': my_fontsize-4,
          'lines.markersize': 8.,
          'lines.linewidth': 2.,
          'xtick.labelsize': my_fontsize-2,
          'ytick.labelsize': my_fontsize-2}

plt.rcParams.update(my_params)
```

```python
# @title Helper functions

def plot_fI(x, f):
  plt.figure(figsize=(6,4)) # plot the figure
  plt.plot(x, f, 'k')
  plt.xlabel('x (a.u.)', fontsize=14.)
  plt.ylabel('F(x)', fontsize=14.)
  plt.show()
```

```python
#@title Helper functions
def plot_dE_E(E, dEdt):
  plt.figure()
  plt.plot(E_grid, dEdt, 'k')
  plt.plot(E_grid, 0.*E_grid, 'k--')
  plt.xlabel('E activity')
  plt.ylabel(r'$\frac{dE}{dt}$', fontsize=20)
  plt.ylim(-0.1, 0.1)

def plot_dFdt(x,dFdt):
  plt.figure()
  plt.plot(x, dFdt, 'r')
  plt.xlabel('x (a.u.)', fontsize=14.)
  plt.ylabel('dF(x)', fontsize=14.)
  plt.show()

```

# Neuronal network dynamics

```python
#@title Video: Dynamic networks
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="ZSsAaeaG9ZM", width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

    Video available at https://youtube.com/watch?v=ZSsAaeaG9ZM

## Dynamics of a single excitatory population

Individual neurons respond by spiking. When we average the spikes of neurons in a population, we can define the average firing activity of the population. In this model, we are interested in how the population-averaged firing varies as a function of different network parameters.

\begin{align}
\tau_E \frac{dE}{dt} &= -E + F(w_{EE}E + I^{\text{ext}}_E)  \quad\qquad (1)
\end{align}

$E(t)$ represents the average firing rate of the excitatory population at time $t$, $\tau_E$ controls the timescale of the evolution of the average firing rate, $w_{EE}$ denotes the strength (synaptic weight) of the recurrent excitatory input to the population, $I^{\text{ext}}_E$ represents the external input, and the transfer function $F(\cdot)$ (which can be related to f-I curve of individual neurons described in the next sections) represents the population activation function in response to all received inputs.

To start building the model, please execute the cell below to initialize the simulation parameters.

```python
#@title Default parameters for a single excitatory population model
def default_parsE( **kwargs):
  pars = {}

  ### Excitatory parameters ###
  pars['tau_E']   = 1.    # Timescale of the E population [ms]
  pars['a_E']     = 1.2   # Gain of the E population
  pars['theta_E'] = 2.8   # Threshold of the E population

  ### Connection strength ###
  pars['wEE'] = 0. # E to E, we first set it to 0

  ### External input ###
  pars['I_ext_E']  = 0.

  ### simulation parameters ###
  pars['T'] = 20. # Total duration of simulation [ms]
  pars['dt'] = .1  # Simulation time step [ms]
  pars['E_init'] = 0.2 # Initial value of E

  ### External parameters if any ###
  for k in kwargs:
      pars[k] = kwargs[k]

  pars['range_t'] = np.arange(0, pars['T'], pars['dt']) # Vector of discretized time points [ms]

  return pars
```

You can use:
- `pars = default_parsE()` to get all the parameters, and then you can execute `print(pars)` to check these parameters.
- `pars = default_parsE(T=T_sim, dt=time_step)` to set new simulation time and time step
- After `pars = default_parsE()`, use `pars['New_para'] = value` to add an new parameter with its value

## F-I curves
In electrophysiology, a neuron is often characterized by its spike rate output in response to input currents. This is often called the **F-I** curve, denoting the spike frequency (**F**) in response to different injected currents (**I**). We estimated this for an LIF neuron in yesterday's tutorial.

The transfer function $F(\cdot)$ in Equation (1) represents the gain of the population as a function of the total input. The gain is often modeled as a sigmoidal function, i.e., more input drive leads to a nonlinear increase in the population firing rate. The output firing rate will eventually saturate for high input values.

A sigmoidal $F(\cdot)$ is parameterized by its gain $a$ and threshold $\theta$.

$$ F(x;a,\theta) = \frac{1}{1+\text{e}^{-a(x-\theta)}} - \frac{1}{1+\text{e}^{a\theta}}  \quad(2)$$

The argument $x$ represents the input to the population. Note that the second term is chosen so that $F(0;a,\theta)=0$.

Many other transfer functions (generally monotonic) can be also used. Examples are the rectified linear function $ReLU(x)$ or the hyperbolic tangent $tanh(x)$.

### Exercise 1: Implement F-I curve

Let's first investigate the activation functions before simulating the dynamics of the entire population.

In this exercise, you will implement a sigmoidal **F-I** curve or transfer function $F(x)$, with gain $a$ and threshold level $\theta$ as parameters.

```python
# Excercise 1
def F(x,a,theta):
  """
  Population activation function.

  Args:
    x (float): the population input
    a (float): the gain of the function
    theta (float): the threshold of the function

  Returns:
    float: the population activation response F(x) for input x
  """

  #################################################################################
  ## TODO for students: compute f = F(x), remove the NotImplementedError once done#
  #################################################################################
  # the exponential function: np.exp(.)
  # f = ...
  raise NotImplementedError("Student excercise: implement the f-I function")

  return f

# Uncomment these lines when you've filled the function, then run the cell again
# to plot the f-I curve.
pars = default_parsE() # get default parameters
# print(pars)          # print out pars to get familiar with parameters
x = np.arange(0,10,.1) # set the range of input

# Uncomment this when you fill the exercise, and call the function
# plot_fI(x,  F(x,pars['a_E'],pars['theta_E']))
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D2_DynamicNetworks/solutions/W3D2_Tutorial1_Solution_3d4f01e4.py)

*Example output:*

### Interactive Demo: Parameter exploration of F-I curve
Here's an interactive demo that shows how the F-I curve is changing for different values of the gain and threshold parameters.

**Remember to enable the demo by running the cell.**

```python
#@title F-I curve Explorer
def interactive_plot_FI(a, theta):
  '''
  Population activation function.

  Expecxts:
  a     : the gain of the function
  theta : the threshold of the function

  Returns:
  plot the F-I curve with give parameters
  '''
  # set the range of input
  x = np.arange(0,10,.1)
  plt.figure()
  plt.plot(x, F(x, a, theta), 'k')
  plt.xlabel('x (a.u.)', fontsize=14.)
  plt.ylabel('F(x)', fontsize=14.)
  plt.show()

_ = widgets.interact(interactive_plot_FI, a = (0.3, 3., 0.3), \
                                              theta = (2., 4., 0.2))
```

## Simulation scheme of E dynamics

Because $F(\cdot)$ is a nonlinear function, the exact solution of Equation $1$ can not be determined via analytical methods. Therefore, numerical methods must be used to find the solution. In practice, the derivative on the left-hand side of Equation (1) can be approximated using the Euler method on a time-grid of stepsize $\Delta t$:

\begin{align}
&\frac{dE}{dt} \approx \frac{E[k+1]-E[k]}{\Delta t}
\end{align}
where $E[k] = E(k\Delta t)$.

Thus,

$$\Delta E[k] = \frac{\Delta t}{\tau_E}[-E[k] + F(w_{EE}E[k] + I^{\text{ext}}_E(k;a_E,\theta_E)]$$

Hence, Equation (1) is updated at each time step by:

$$E[k+1] = E[k] + \Delta E[k]$$

**_Please execute the following cell to enable the WC simulator_**

```python
#@title E population simulator: `simulate_E`
def simulate_E(pars):

  """
  Simulate an excitatory population of neurons

  Args:
    pars : Parameter dictionary

  Returns:
    E    : Activity of excitatory population (array)
  """

  # Set parameters
  tau_E, a_E, theta_E = pars['tau_E'], pars['a_E'], pars['theta_E']
  wEE = pars['wEE']
  I_ext_E = pars['I_ext_E']
  E_init = pars['E_init']
  dt, range_t = pars['dt'], pars['range_t']
  Lt = range_t.size

  # Initialize activity
  E = np.zeros(Lt)
  E[0] = E_init
  I_ext_E = I_ext_E*np.ones(Lt)

  # Update the E activity
  for k in range(Lt-1):
      dE = dt/tau_E * (-E[k] + F(wEE*E[k]+I_ext_E[k], a_E, theta_E))
      E[k+1] = E[k] + dE

  return E

print(help(simulate_E))

```

    Help on function simulate_E in module __main__:

    simulate_E(pars)
        Simulate an excitatory population of neurons

        Args:
          pars : Parameter dictionary

        Returns:
          E    : Activity of excitatory population (array)

    None

#### Interactive Demo: Parameter Exploration of single population dynamics

Note that $w_{EE}=0$, as in the default setting, means no recurrent input to the excitatory population in Equation (1). Hence, the dynamics is entirely determined by the external input $I_{E}^{\text{ext}}$. Try to explore how $E_{sim}(t)$ changes with different $I_{E}^{\text{ext}}$ and $\tau_E$ parameter values, and investigate the relationship between $F(I_{E}^{\text{ext}}; a_E, \theta_E)$ and the steady value of E. Note that, $E_{ana}(t)$ denotes the analytical solution.

```python
#@title Mean-field model Explorer

# get default parameters
pars = default_parsE(T=20.)

def Myplot_E_diffI_difftau(I_ext, tau_E):
  # set external input and time constant
  pars['I_ext_E'] = I_ext
  pars['tau_E'] = tau_E

  # simulation
  E = simulate_E(pars)

  # Analytical Solution
  E_ana = pars['E_init'] + (F(I_ext,pars['a_E'],pars['theta_E'])-pars['E_init'])*\
  (1.-np.exp(-pars['range_t']/pars['tau_E']))

  # plot
  plt.figure()
  plt.plot(pars['range_t'], E, 'b', label=r'$E_{\mathrm{sim}}$(t)', alpha=0.5, zorder=1)
  plt.plot(pars['range_t'], E_ana, 'b--', lw=5, dashes=(2,2),\
            label=r'$E_{\mathrm{ana}}$(t)', zorder=2)
  plt.plot(pars['range_t'], F(I_ext,pars['a_E'],pars['theta_E'])\
          *np.ones(pars['range_t'].size), 'k--', label=r'$F(I_E^{\mathrm{ext}})$')
  plt.xlabel('t (ms)', fontsize=16.)
  plt.ylabel('E activity', fontsize=16.)
  plt.legend(loc='best', fontsize=14.)
  plt.show()

_ = widgets.interact(Myplot_E_diffI_difftau, I_ext = (0.0, 10., 1.),\
                                             tau_E = (1., 5., 0.2))

```

### Think!
Above, we have numerically solved a system driven by a positive input and that, if $w_{EE} \neq 0$, receives an excitatory recurrent input (**try changing the value of $w_{EE}$ to a positive number**). Yet, $E(t)$ either decays to zero or reaches a fixed non-zero value.
- Why doesn't the solution of the system "explode" in a finite time? In other words, what guarantees that E(t) stays finite?
- Which parameter would you change in order to increase the maximum value of the response?

## Fixed points of the E system

```python
#@title Video: Fixed point
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="B31fX6V0PZ4", width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

    Video available at https://youtube.com/watch?v=B31fX6V0PZ4

As you varied the two parameters in the last Interactive Demo, you noticed that, while at first the system output quickly changes, with time, it reaches its maximum/minimum value and does not change anymore. The value eventually reached by the system is called the **steady state** of the system, or the **fixed point**. Essentially, in the steady states the derivative with respect to time of the activity ($E$) is zero, i.e. $\frac{dE}{dt}=0$.

We can find that the steady state of the Equation $1$ by setting $\displaystyle{\frac{dE}{dt}=0}$ and solve for $E$:

$$E_{\text{steady}} = F(w_{EE}E_{\text{steady}} + I^{\text{ext}}_E;a_E,\theta_E) = 0, \qquad (3)$$

When it exists, the solution of Equation $3$ defines a **fixed point** of the dynamics which satisfies $\displaystyle{\frac{dE}{dt}=0}$ (and determines steady state of the system). Notice that the right-hand side of the last equation depends itself on $E_{steady}$. If $F(x)$ is nonlinear it is not always possible to find an analytical solution that can instead be found via numerical simulations, as we will do later.

From the Interactive Demo one could also notice that the value of $\tau_E$ influences how quickly the activity will converge to the steady state from its initial value.

In the specific case of $w_{EE}=0$, we can also analytically compute  the analytical solution of Equation $1$ (i.e., the thick blue dashed line) and deduce the role of $\tau_E$ in determining the convergence to the fixed point:

$$\displaystyle{E(t) = \big{[}F(I^{\text{ext}}_E;a_E,\theta_E) -E(t=0)\big{]} (1-\text{e}^{-\frac{t}{\tau_E}})} + E(t=0)$$ \\

We can now numerically calculate the fixed point with the `scipy.optimize.root` function.

_(note that at the very beginning, we `import scipy.optimize as opt` )_.

\\

Please execute the cell below to define the functions `my_fpE`, `check_fpE`, and `plot_fpE`

```python
#@title Function of calculating the fixed point
def my_fpE(pars, E_init):

  # get the parameters
  a_E, theta_E = pars['a_E'], pars['theta_E']
  wEE = pars['wEE']
  I_ext_E = pars['I_ext_E']

  # define the right hand of E dynamics
  def my_WCr(x):
    E = x[0]
    dEdt=(-E + F(wEE*E+I_ext_E,a_E,theta_E))
    y = np.array(dEdt)

    return y

  x0 = np.array(E_init)
  x_fp = opt.root(my_WCr, x0).x

  return x_fp

def check_fpE(pars, x_fp):

  a_E, theta_E = pars['a_E'], pars['theta_E']
  wEE = pars['wEE']
  I_ext_E = pars['I_ext_E']

  # calculate Equation(3)
  y = x_fp- F(wEE*x_fp+I_ext_E, a_E, theta_E)

  return np.abs(y)<1e-4

def plot_fpE(pars, x_fp, mycolor):

  wEE = pars['wEE']
  I_ext_E = pars['I_ext_E']

  plt.plot(wEE*x_fp+I_ext_E, x_fp, 'o', color=mycolor)

```

#### Exercise 2: Visualization of the fixed point

When no analytical solution of Equation $3$ can be found, it is often useful to plot $\displaystyle{\frac{dE}{dt}=0}$ as a function of $E$. The values of E for which the plotted function crosses zero on the y axis correspond to fixed points.

Here, let us, for example, set $w_{EE}=5.0$ and $I^{\text{ext}}_E=0.5$. Define $\displaystyle{\frac{dE}{dt}}$ using Equation $1$, plot the result, and check for the presence of fixed points.

We will now try to find the fixed points using the previously defined function `my_fpE(pars, E_init)` with different initial values ($E_{\text{init}}$). Use the previously defined function `check_fpE(pars, x_fp)` to verify that the values of $E$ for which $\displaystyle{\frac{dE}{dt}} = 0$ are the true fixed points.

```python
# Exercise 2
pars = default_parsE() # get default parameters

# set your external input and wEE
pars['I_ext_E'] = 0.5
pars['wEE'] = 5.0

E_grid = np.linspace(0, 1., 1000)# give E_grid

#figure, line (E, dEdt)
###############################
## TODO for students:         #
## Calculate dEdt = -E + F(.) #
## Then plot the lines        #
###############################
# Calculate dEdt
# dEdt = ...

# Uncomment this to plot the dEdt across E
# plot_dE_E(E_grid, dEdt)

# Add fixed point
#####################################################
## TODO for students:                               #
# Calculate the fixed point with your initial value #
# verify your fixed point and plot the corret ones  #
#####################################################

# Calculate the fixed point with your initial value
x_fp_1 = my_fpE(pars, 1)

#check if x_fp is the intersection of the lines with the given function check_fpE(pars, x_fp)
#vary different initial values to find the correct fixed point (Should be 3)
# Use blue, red and yellow colors, respectively ('b', 'r', 'y' codenames)

# if check_fpE(pars, x_fp_1):
#   plt.plot(x_fp_1, 0, 'bo', ms=8)

# Replicate the code above (lines 35-36) for all fixed points.

```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D2_DynamicNetworks/solutions/W3D2_Tutorial1_Solution_e231986d.py)

*Example output:*

#### Interactive Demo: fixed points as a function of recurrent and external inputs.

You can now explore how the previous plot changes when the recurrent coupling $w_{\text{EE}}$ and the external input $I_E^{\text{ext}}$ take different values.

```python
#@title Fixed point Explorer

def plot_intersection_E(wEE, I_ext_E):
  #set your parameters
  pars['wEE'] = wEE
  pars['I_ext_E'] = I_ext_E

  #note that wEE !=0
  if wEE>0:
    # find fixed point
    x_fp_1 = my_fpE(pars, 0.)
    x_fp_2 = my_fpE(pars, 0.4)
    x_fp_3 = my_fpE(pars, 0.9)

    plt.figure()
    E_grid = np.linspace(0, 1., 1000)
    dEdt = -E_grid + F(wEE*E_grid+I_ext_E, pars['a_E'],  pars['theta_E'])

    plt.plot(E_grid, dEdt, 'k')
    plt.plot(E_grid, 0.*E_grid, 'k--')

    if check_fpE(pars, x_fp_1):
        plt.plot(x_fp_1, 0, 'bo', ms=8)
    if check_fpE(pars, x_fp_2):
        plt.plot(x_fp_2, 0, 'bo', ms=8)
    if check_fpE(pars, x_fp_3):
        plt.plot(x_fp_3, 0, 'bo', ms=8)

    plt.xlabel('E activity', fontsize=14.)
    plt.ylabel(r'$\frac{dE}{dt}$', fontsize=18.)

    plt.show()

_ = widgets.interact(plot_intersection_E, wEE = (1., 7., 0.2), \
                                          I_ext_E = (0., 3., 0.1))
```

## Summary

In this tutorial, we have investigated the dynamics of a rate-based single excitatory population of neurons.

We learned about:
- The effect of the input parameters and the time constant of the network on the dynamics of the population.
- How to find the fixed point(s) of the system.

Next, we have two Bonus, but important concepts in dynamical system analysis and simulation. If you have time left, watch the next video and proceed to solve the exercises. You will learn:

- How to determine the stability of a fixed point by linearizing the system.
- How to add realistic inputs to our model.

## Bonus 1: Stability of a fixed point

```python
#@title Video: Stability of fixed points
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="nvxxf59w2EA", width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

    Video available at https://youtube.com/watch?v=nvxxf59w2EA

#### Initial values and trajectories

Here, let us first set $w_{EE}=5.0$ and $I^{\text{ext}}_E=0.5$, and investigate the dynamics of $E(t)$ starting with different initial values $E(0) \equiv E_{\text{init}}$. We will plot the trajectories of $E(t)$ with $E_{\text{init}} = 0.0, 0.1, 0.2,..., 0.9$.

```python
#@title Initial values
pars = default_parsE()
pars['wEE'] = 5.0
pars['I_ext_E'] = 0.5

plt.figure(figsize=(10,6))
for ie in range(10):
  pars['E_init'] =  0.1*ie # set the initial value
  E = simulate_E(pars)   # run the simulation
  # plot the activity with given initial
  plt.plot(pars['range_t'], E, 'b', alpha=0.1 + 0.1*ie, label= r'E$_{\mathrm{init}}$=%.1f' % (0.1*ie))

plt.xlabel('t (ms)')
plt.title('Two steady states?')
plt.ylabel('E(t)')
plt.legend(loc=[0.72, 0.13], fontsize=14)
plt.show()

```

#### Interactive Demo: dynamics as a function of the initial value.

Let's now set $E_{init}$ to a value of your choice in this demo. How does the solution change? What do you observe?

```python
#@title Initial value Explorer
pars = default_parsE()
pars['wEE'] = 5.0
pars['I_ext_E'] = 0.5

def plot_E_diffEinit(E_init):
  pars['E_init'] =  E_init
  E = simulate_E(pars)

  plt.figure()
  plt.plot(pars['range_t'], E, 'b', label='E(t)')
  plt.xlabel('t (ms)', fontsize=16.)
  plt.ylabel('E activity', fontsize=16.)
  plt.show()

_ = widgets.interact(plot_E_diffEinit, E_init = (0., 1., 0.02))
```

### Stability analysis via linearization of the dynamics

Just like Equation $1$ in the case ($w_{EE}=0$) discussed above, a generic linear system
$$\frac{dx}{dt} = \lambda (x - b),$$
has a fixed point for $x=b$. The analytical solution of such a system can be found to be:
$$x(t) = b + \big{(} x(0) - b \big{)} \text{e}^{\lambda t}.$$
Now consider a small perturbation of the activity around the fixed point: $x(0) = b+ \epsilon$, where $|\epsilon| \ll 1$. Will the perturbation $\epsilon(t)$ grow with time or will it decay to the fixed point? The evolution of the perturbation with time can be written, using the analytical solution for $x(t)$, as:
 $$\epsilon (t) = x(t) - b = \epsilon \text{e}^{\lambda t}$$

- if $\lambda < 0$, $\epsilon(t)$ decays to zero, $x(t)$ will still converge to $b$ and the fixed point is "**stable**".

- if $\lambda > 0$, $\epsilon(t)$ grows with time, $x(t)$ will leave the fixed point $b$ exponentially and the fixed point is, therefore, "**unstable**" .

### Compute the stability of Equation (1)

Similar to what we did in the linear system above, in order to determine the stability of a fixed point $E_{\rm fp}$ of the excitatory population dynamics, we perturb Equation $1$ around $E_{\rm fp}$ by $\epsilon$, i.e. $E = E_{\rm fp} + \epsilon$. We can plug in Equation $1$ and obtain the equation determining the time evolution of the perturbation $\epsilon(t)$:

\begin{align}
\tau_E \frac{d\epsilon}{dt} \approx -\epsilon + w_{EE} F'(w_{EE}E_{\text{fp}} + I^{\text{ext}}_E;a_E,\theta_E) \epsilon
\end{align}
where $F'(\cdot)$ is the derivative of the transfer function $F(\cdot)$. We can rewrite the above equation as:
\begin{align}
\frac{d\epsilon}{dt} \approx \frac{\epsilon}{\tau_E }[-1 + w_{EE} F'(w_{EE}E_{\text{fp}} + I^{\text{ext}}_E;a_E,\theta_E)]
\end{align}

That is, as in the linear system above, the value of $\lambda = [-1+ w_{EE}F'(w_{EE}E_{\text{fp}} + I^{\text{ext}}_E;a_E,\theta_E)]/\tau_E$ determines whether the perturbation will grow or decay to zero, i.e., $\lambda$ defines the stability of the fixed point. This value is called the **eigenvalue** of the dynamical system.

### Exercise 4: Compute $dF$ and Eigenvalue

The derivative of the sigmoid transfer function is:
\begin{align}
\frac{dF}{dx} & = \frac{d}{dx} (1+\exp\{-a(x-\theta)\})^{-1}  \\
& = a\exp\{-a(x-\theta)\} (1+\exp\{-a(x-\theta)\})^{-2}.
\end{align}

Let's now find the expression for the derivative $\displaystyle{\frac{dF}{dx}}$ in the following cell and plot it.

```python
# Exercise 4

def dF(x,a,theta):
    """
    Population activation function.

    Args:
      x     : the population input
      a     : the gain of the function
      theta : the threshold of the function

    Returns:
      dFdx  : the population activation response F(x) for input x
    """

    #####################################################################
    ## TODO for students: compute dFdx, then remove NotImplementedError #
    #####################################################################
    # dFdx = ...
    raise NotImplementedError("Student excercise: compute the deravitive of F(x)")

    return dFdx

pars = default_parsE() # get default parameters
x = np.arange(0,10,.1) # set the range of input

# Uncomment below lines after completing the dF function
# plot_dFdt(x,dF(x,pars['a_E'],pars['theta_E']))

```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D2_DynamicNetworks/solutions/W3D2_Tutorial1_Solution_1b43f8f1.py)

*Example output:*

### Exercise 5: Compute eigenvalues

As discussed above, for the case with $w_{EE}=5.0$ and $I^{\text{ext}}_E=0.5$, the system displays **3** fixed points. However, when we simulated the dynamics and varied the initial conditions $E_{\rm init}$, we could only obtain **two** steady states. In this exercise, we will now check the stability of each of the $3$ fixed points by calculating the corresponding eigenvalues with the function `eig_E` defined above. Check the sign of each eigenvalue (i.e., stability of each fixed point). How many of the fixed points are stable?

```python
# Exercise 5
pars = default_parsE()
pars['wEE'] = 5.0
pars['I_ext_E'] = 0.5

def eig_E(pars, fp):
  """
  Args:
    pars : Parameter dictionary
    fp   : fixed point E

  Returns:
    eig : eigevalue of the linearized system
  """

  #get the parameters
  tau_E, a_E, theta_E = pars['tau_E'], pars['a_E'], pars['theta_E']
  wEE, I_ext_E = pars['wEE'],  pars['I_ext_E']
  # fixed point
  E = fp

  #######################################################################
  ## TODO for students: compute eigenvalue, remove NotImplementedError  #
  #######################################################################
  # eig = ...
  raise NotImplementedError("Student excercise: compute the eigenvalue")

  return eig

# Uncomment below lines after completing the eigE function.
# x_fp_1 = fpE(pars, 0.)
# eig_fp_1 = eig_E(pars, x_fp_1)
# print('Fixed point1=%.3f, Eigenvalue=%.3f' % (x_fp_1, eig_E1))

# Continue by finding the eigenvalues for all fixed points of Exercise 2
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D2_DynamicNetworks/solutions/W3D2_Tutorial1_Solution_3979c031.py)

### Think!
Throughout the tutorial, we have assumed $w_{\rm EE}> 0 $, i.e., we considered a single population of **excitatory** neurons. What do you think will be the behavior of a population of inhibitory neurons, i.e., where $w_{\rm EE}> 0$ is replaced by $w_{\rm II}< 0$?

## Bonus 2: Noisy input drives transition between two stable states

### Ornstein-Uhlenbeck (OU) process

As discussed in several previous tutorials, the OU process is usually used to generate a noisy input into the neuron. The OU input $\eta(t)$ follows:

$$\tau_\eta \frac{d}{dt}\eta(t) = -\eta (t) + \sigma_\eta\sqrt{2\tau_\eta}\xi(t)$$

Execute the following function `my_OU(pars, sig, myseed=False)` to generate an OU process.

```python
#@title OU process `my_OU(pars, sig, myseed=False)`
def my_OU(pars, sig, myseed=False):
  """
  A functions that generates Ornstein-Uhlenback process

  Args:
    pars       : parameter dictionary
    sig        : noise amplitute
    myseed     : random seed. int or boolean

  Returns:
    I          : Ornstein-Uhlenbeck input current
  """

  # Retrieve simulation parameters
  dt, range_t = pars['dt'], pars['range_t']
  Lt = range_t.size
  tau_ou = pars['tau_ou'] # [ms]

  # set random seed
  if myseed:
      np.random.seed(seed=myseed)
  else:
      np.random.seed()

  # Initialize
  noise =  np.random.randn(Lt)
  I = np.zeros(Lt)
  I[0] = noise[0] * sig

  #generate OU
  for it in range(Lt-1):
      I[it+1] = I[it] + dt/tau_ou*(0.-I[it]) + np.sqrt(2.*dt/tau_ou) * sig * noise[it+1]

  return I

pars = default_parsE(T=100)
pars['tau_ou'] = 1. #[ms]
sig_ou = 0.1
I_ou = my_OU(pars, sig=sig_ou, myseed=1998)
plt.figure(figsize=(10, 4))
plt.plot(pars['range_t'], I_ou, 'b')
plt.xlabel('Time (ms)')
plt.ylabel(r'$I_{\mathrm{OU}}$');
```

### Bonus Example: Up-Down transition

In the presence of two or more fixed points, noisy inputs can drive a transition between the fixed points! Here, we stimulate an E population for 1,000 ms applying OU inputs.

```python
#@title Simulation of an E population with OU inputs
pars = default_parsE(T = 1000)
pars['wEE'] = 5.0
sig_ou = 0.7
pars['tau_ou'] = 1. #[ms]
pars['I_ext_E'] = 0.56 + my_OU(pars, sig=sig_ou, myseed=2020)

E = simulate_E(pars)

plt.figure(figsize=(10, 4))
plt.plot(pars['range_t'], E, 'r', alpha=0.8)
plt.xlabel('t (ms)')
plt.ylabel('E activity')
plt.show()
```


# A complete use case

In this section we present a complete use case, based on the meaning classification dataset introduced in [Lorenz et al. (2021)](https://arxiv.org/abs/2102.12846) QNLP paper. The goal is to classify simple sentences (such as "skillful programmer creates software" and "chef prepares delicious meal") into two categories, food or IT. The dataset consists of 130 sentences created using a simple context-free grammar.

We will use a [SpiderAnsatz](../lambeq.rst#lambeq.tensor.SpiderAnsatz) to split large tensors into chains of smaller ones. For differentiation we will use JAX, and we will apply simple gradient-descent optimisation to train the tensors.

## Preparation

We start with a few essential imports.

```python
import warnings
warnings.filterwarnings('ignore')  # Ignore warnings

from discopy.tensor import Tensor
from jax import numpy as np
import numpy

np.random = numpy.random
Tensor.np = np

np.random.seed(123458)  # Fix the seed
```

<div class="alert alert-info">

**Note**

Note the `Tensor.np = np` assignment in the above code. This is required to let `discopy` know that from now on we use JAX's version of `numpy`.

</div>

Let's read the datasets:

```python
# Read data
def read_data(fname):
    with open(fname, 'r') as f:
        lines = f.readlines()
    data, targets = [], []
    for ln in lines:
        t = int(ln[0])
        data.append(ln[1:].strip())
        targets.append(np.array([t, not(t)], dtype=np.float32))
    return data, np.array(targets)

train_data, train_targets = read_data('datasets/mc_train_data.txt')
test_data, test_targets = read_data('datasets/mc_test_data.txt')
```

The first few lines of the train dataset:

```python
train_data[:10]
```

    ['skillful man prepares sauce',
     'skillful man bakes dinner',
     'woman cooks tasty meal',
     'man prepares meal',
     'skillful woman debugs program',
     'woman prepares tasty meal',
     'person runs program',
     'person runs useful application',
     'woman prepares sauce',
     'woman prepares dinner']

Targets are represented as 2-dimensional arrays:

```python
train_targets
```

    DeviceArray([[1., 0.],
                 [1., 0.],
                 [1., 0.],
                 ...,
                 [0., 1.],
                 [1., 0.],
                 [0., 1.]], dtype=float32)

## Creating and parameterising diagrams

First step is to convert sentences into string diagrams:

```python
# Parse sentences to diagrams

from lambeq.ccg2discocat import DepCCGParser

parser = DepCCGParser()
train_diagrams = parser.sentences2diagrams(train_data)
test_diagrams = parser.sentences2diagrams(test_data)

train_diagrams[0].draw(figsize=(8,4), fontsize=13)
```

The produced diagrams need to be parameterised by a specific ansatz. For this experiment we will use a [SpiderAnsatz](../lambeq.rst#lambeq.tensor.SpiderAnsatz).

```python
# Create ansatz and convert to tensor diagrams

from lambeq.tensor import SpiderAnsatz
from lambeq.core.types import AtomicType
from discopy import Dim

N = AtomicType.NOUN
S = AtomicType.SENTENCE

# Create an ansatz by assigning 2 dimensions to both
# noun and sentence spaces
ansatz = SpiderAnsatz({N: Dim(2), S: Dim(2)})

train_circuits = [ansatz(d) for d in train_diagrams]
test_circuits = [ansatz(d) for d in test_diagrams]

all_circuits = train_circuits + test_circuits

all_circuits[0].draw(figsize=(8,4), fontsize=13)
```

## Creating a vocabulary

We are now ready to create a vocabulary.

```python
# Create vocabulary

from sympy import default_sort_key

vocab = sorted(
   {sym for circ in all_circuits for sym in circ.free_symbols},
    key=default_sort_key
)
tensors = [np.random.rand(w.size) for w in vocab]

tensors[0]
```

    array([0.35743395, 0.45764418])

## Defining a loss function

This is a binary classification task, so we will use binary cross entropy as the loss.

```python
def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def loss(tensors):
    # Lambdify
    np_circuits = [c.lambdify(*vocab)(*tensors) for c in train_circuits]
    # Compute predictions
    predictions =  sigmoid(np.array([c.eval().array for c in np_circuits]))

    # binary cross-entropy loss
    cost = -np.sum(train_targets * np.log2(predictions)) / len(train_targets)
    return cost
```

The loss function follows the steps below:

1. The symbols in the train diagrams are replaced with concrete ``numpy`` arrays.
2. The resulting tensor networks are evaluated and produce results.
3. Based on the predictions, an average loss is computed for the specific iteration.

We use JAX in order to get a gradient function on the loss, and "just-in-time" compile it to improve speed:

```python
from jax import jit, grad

training_loss = jit(loss)
gradient = jit(grad(loss))
```

## Training loop

We are now ready to start training. The following loop computes gradients and  uses them to update the tensors associated with the symbols.

```python
training_losses = []

epochs = 90

for i in range(epochs):

    gr = gradient(tensors)
    for k in range(len(tensors)):
        tensors[k] = tensors[k] - gr[k] * 1.0

    training_losses.append(float(training_loss(tensors)))

    if (i + 1) % 10 == 0:
        print(f"Epoch {i + 1} - loss {training_losses[-1]}")
```

    Epoch 10 - loss 0.07233709841966629
    Epoch 20 - loss 0.015333528630435467
    Epoch 30 - loss 0.00786149874329567
    Epoch 40 - loss 0.00515687046572566
    Epoch 50 - loss 0.0037753921933472157
    Epoch 60 - loss 0.0029438300989568233
    Epoch 70 - loss 0.002392344642430544
    Epoch 80 - loss 0.0020021884702146053
    Epoch 90 - loss 0.001713048666715622

## Testing

Finally, we use the trained model on the test dataset:

```python
# Testing

np_test_circuits = [c.lambdify(*vocab)(*tensors) for c in test_circuits]
test_predictions =  sigmoid(np.array([c.eval().array for c in np_test_circuits]))

hits = 0
for i in range(len(np_test_circuits)):
    target = test_targets[i]
    pred = test_predictions[i]
    if np.argmax(target) == np.argmax(pred):
        hits += 1

print("Accuracy on test set:", hits / len(np_test_circuits))
```

    Accuracy on test set: 0.9

## Working with quantum circuits

The process when working with quantum circuits is very similar, with two important differences:

1. The parameterisable part of the circuit is an array of parameters, as described in Section [Circuit Symbols](training-symbols.ipynb#Circuit-symbols), instead of tensors associated to words.
2. If optimisation takes place on quantum hardware, standard automatic differentiation cannot be used. An alternative is to use a gradient-approximation technique, such as [Simultaneous Perturbation Stochastic Approximation](https://en.wikipedia.org/wiki/Simultaneous_perturbation_stochastic_approximation) (SPSA).

Complete examples in training quantum circuits can be found in the following notebooks:

- [Quantum pipeline with JAX](../examples/quantum_pipeline_jax.ipynb)
- [Quantum pipeline with tket](../examples/quantum_pipeline_tket.ipynb)

**See also:**

- [Classical pipeline with PyTorch](../examples/classical_pipeline.ipynb)


```python
#@title Imports and utility code
!pip install jax-md

import numpy as onp

import jax.numpy as jnp
from jax.config import config
config.update('jax_enable_x64', True)

from jax import random
from jax import jit, lax, grad, vmap
import jax.scipy as jsp

from jax_md import space, energy, smap, minimize, util, elasticity, quantity
from jax_md.colab_tools import renderer

f32 = jnp.float32
f64 = jnp.float64

from functools import partial

import matplotlib
import matplotlib.pyplot as plt
plt.rcParams.update({'font.size': 16})

def format_plot(x, y):
  plt.grid(True)
  plt.xlabel(x, fontsize=20)
  plt.ylabel(y, fontsize=20)

def finalize_plot(shape=(1, 0.7)):
  plt.gcf().set_size_inches(
    shape[0] * 1.5 * plt.gcf().get_size_inches()[1],
    shape[1] * 1.5 * plt.gcf().get_size_inches()[1])

def run_minimization_while(energy_fn, R_init, shift, max_grad_thresh = 1e-12, max_num_steps=1000000, **kwargs):
  init,apply=minimize.fire_descent(jit(energy_fn), shift, **kwargs)
  apply = jit(apply)

  @jit
  def get_maxgrad(state):
    return jnp.amax(jnp.abs(state.force))

  @jit
  def cond_fn(val):
    state, i = val
    return jnp.logical_and(get_maxgrad(state) > max_grad_thresh, i<max_num_steps)

  @jit
  def body_fn(val):
    state, i = val
    return apply(state), i+1

  state = init(R_init)
  state, num_iterations = lax.while_loop(cond_fn, body_fn, (state, 0))

  return state.position, get_maxgrad(state), num_iterations

def run_minimization_while_neighbor_list(energy_fn, neighbor_fn, R_init, shift,
                                         max_grad_thresh = 1e-12, max_num_steps = 1000000,
                                         step_inc = 1000, verbose = False, **kwargs):
  nbrs = neighbor_fn.allocate(R_init)

  init,apply=minimize.fire_descent(jit(energy_fn), shift, **kwargs)
  apply = jit(apply)

  @jit
  def get_maxgrad(state):
    return jnp.amax(jnp.abs(state.force))

  @jit
  def body_fn(state_nbrs, t):
    state, nbrs = state_nbrs
    nbrs = neighbor_fn.update(state.position, nbrs)
    state = apply(state, neighbor=nbrs)
    return (state, nbrs), 0

  state = init(R_init, neighbor=nbrs)

  step = 0
  while step < max_num_steps:
    if verbose:
      print('minimization step {}'.format(step))
    rtn_state, _ = lax.scan(body_fn, (state, nbrs), step + jnp.arange(step_inc))
    new_state, nbrs = rtn_state
    # If the neighbor list overflowed, rebuild it and repeat part of
    # the simulation.
    if nbrs.did_buffer_overflow:
      print('Buffer overflow.')
      nbrs = neighbor_fn.allocate(state.position)
    else:
      state = new_state
      step += step_inc
      if get_maxgrad(state) <= max_grad_thresh:
        break

  if verbose:
    print('successfully finished {} steps.'.format(step*step_inc))

  return state.position, get_maxgrad(state), nbrs, step

def run_minimization_scan(energy_fn, R_init, shift, num_steps=5000, **kwargs):
  init,apply=minimize.fire_descent(jit(energy_fn), shift, **kwargs)
  apply = jit(apply)

  @jit
  def scan_fn(state, i):
    return apply(state), 0.

  state = init(R_init)
  state, _ = lax.scan(scan_fn,state,jnp.arange(num_steps))

  return state.position, jnp.amax(jnp.abs(state.force))

key = random.PRNGKey(0)

```

#Linear elasticity in athermal systems

## The elastic modulus tensor

An global affine deformation is given to lowest order by a symmetric strain tensor $\epsilon$, which transforms any vector $r$ according to
\begin{equation}
r \rightarrow (1 + \epsilon) \cdot r.
\end{equation}
Note that in $d$ dimensions, the strain tensor has $d(d + 1)/2$ independent elements. Now, when a mechanically stable system (i.e. a system at a local energy minimum where there is zero net force on every particle) is subject to an affine deformation, it usually does not remain in mechanical equilibrium. Therefore, there is a secondary, nonaffine response that returns the system to mechanical equilibrium, though usually at a different energy than the undeformed state.

The change of energy can be written to quadratic order as
\begin{equation}
\frac{ \Delta U}{V^0} = \sigma^0_{ij}\epsilon_{ji} + \frac 12 C_{ijkl} \epsilon_{ij} \epsilon_{kl} + O\left( \epsilon^3 \right)
\end{equation}
where $C_{ijkl}$ is the $d × d × d × d$ elastic modulus tensor, $\sigma^0$ is the $d × d$ symmetric stress tensor describing residual stresses in the initial state, and $V^0$ is the volume of the initial state. The symmetries of $\epsilon_{ij}$ imply the following:
\begin{equation}
C_{ijkl} = C_{jikl} = C_{ijlk} = C_{klij}
\end{equation}
When no further symmetries are assumed, the number of independent elastic constants becomes $\frac 18 d(d + 1)(d^2 + d + 2)$, which is 6 in two dimensions and 21 in three dimensions.

##Linear response to an external force

Consider a set of $N$ particles in $d$ dimensions with positions $R_0$. Using $u \equiv R - R_0$ and assuming fixed boundary conditions, we can expand the energy about $R_0$:
\begin{equation}
U = U^0 - F^0 u + \frac 12 u H^0 u + O(u^3),
\end{equation}
where $U^0$ is the energy at $R_0$, $F^0$ is the force, $F^0_\mu \equiv \left. \frac {\partial U}{\partial u_\mu} \right |_{u=0}$, and $H^0$ is the Hessian, $H^0 \equiv \left. \frac{ \partial^2 U}{\partial u_\mu \partial u_\nu}\right|_{u=0}$.
Note that here we are expanding in terms of the particle positions, where as above we were expanding in the global strain degrees of freedom.
If we assume that $R_0$ corresponds to a local energy minimum, then $F^0=0$. Dropping higher order terms, we have a system of coupled harmonic oscillators given by
\begin{equation}
\Delta U \equiv U - U^0 = \frac 12 u H^0 u.
\end{equation}
This is independent of the form or details of $U$.

Hooke's law for this system gives the net force $f$ as a result of displacing the particles by $u$:
\begin{equation}
f = -H^0 u.
\end{equation}
Thus, if an *external* force $f_\mathrm{ext}$ is applied, the particles will respond so that the total force is zero, i.e. $f = -f_\mathrm{ext}$. This response is obtained by solving for $u$:
\begin{equation}
u = (H^0)^{-1} f_\mathrm{ext}.
\end{equation}

## Response to an affine strain

Now consider a strain tensor $\epsilon = \tilde \epsilon \gamma$, where $\gamma$ is a scalar and will be used to explicitly take the limit of small strain for fixed $\tilde \epsilon$. Importantly, the strain tensor represents a deformation of the underlying space that the particles live in and thus is a degree of freedom that is independent of the $Nd$ particle degrees of freedom. Therefore, knowing the particle positions $R$ is not sufficient to describe the energy, we also need to know $\gamma$ to specify the correct boundary conditions:
\begin{equation}
U = U(R, \gamma).
\end{equation}

We now have a system with $Nd+1$ variables $\{R, \gamma\}$ that, like before, form a set of coupled harmonic oscillators. We can describe this using the so-called "generalized Hessian" matrix of second derivatives of the energy with respect to both $R$ and $\gamma$. Specifically, Hooke's law reads
\begin{equation}
  \left( \begin{array}{ ccccc|c}
	&&&&&\\
	&&H^0 &&& -\Xi \\
	&&&&& \\ \hline
	&&-\Xi^T &&&\frac{\partial ^2U}{\partial \gamma^2}
	\end{array}\right)
  \left( \begin{array}{ c}
	\\
	u \\
	\\ \hline
	\gamma
	\end{array}\right)
  =
  \left( \begin{array}{ c}
	\\
	0 \\
	\\ \hline
	\tilde \sigma
	\end{array}\right),
\end{equation}

where $u = R - R_0$ is the displacement of every particle, $\Xi = -\frac{ \partial^2 U}{\partial R \partial \gamma}$, and $\tilde \sigma$ is the induced stress caused by the deformation. (If there is prestress in the system, i.e. $\sigma^0 = \frac{\partial U}{\partial \gamma} \neq 0$, the total stress is $\sigma = \sigma^0 + \tilde \sigma$.) In this equation, $\gamma$ is held fixed and the zero in the top of the right-hand-side imposes force balance after the deformation and resulting non-affine displacement of every particle. The non-affine displacement itself, $u$, and the induced stress $\sigma$, are both unknown but can be solved for. First, the non-affine response is
\begin{equation}
u = (H^0)^{-1} \Xi \; \gamma,
\end{equation}
where we note that in the limit of small $\gamma$, the force induced on every particle due to the affine deformation is $\Xi \; \gamma$. Second, the induced stress is
\begin{equation}
\tilde \sigma = \frac{\partial ^2U}{\partial \gamma^2} \gamma - \Xi^T u = \left(\frac{\partial ^2U}{\partial \gamma^2} - \Xi^T (H^0)^{-1} \Xi  \right) \gamma.
\end{equation}

Similarly, the change in energy is
\begin{equation}
\frac{\Delta U}{V^0} = \sigma^0 \gamma + \frac 1{2V^0} \left(\frac{\partial ^2U}{\partial \gamma^2} - \Xi^T (H^0)^{-1} \Xi  \right) \gamma^2,
\end{equation}
where $\sigma^0$ is the prestress in the system per unit volume. Comparing this to the above definition of the the elastic modulus tensor, we see that the elastic constant associated with the deformation $\tilde \epsilon$ is
\begin{equation}
C(\tilde \epsilon) = \frac 1{V^0} \left( \frac{\partial^2 U}{\partial \gamma^2} - \Xi^T (H^0)^{-1} \Xi \right).
\end{equation}

$C(\tilde \epsilon)$ is related to $C_{ijkl}$ by summing $C(\tilde \epsilon) = C_{ijkl}\tilde \epsilon_{ij} \tilde \epsilon_{kl}$. So, if $\tilde \epsilon_{ij} = \delta_{0i}\delta_{0j}$, then $C_{0000} = C(\tilde \epsilon)$.

The internal code in `jax_md.elasticity` repeats this calculation for different $\tilde \epsilon$ to back out the different independent elastic constants.

#First example

As a first example, let's consider a 3d system of 128 soft spheres. The elastic modulus tensor is only defined for systems that are at a local energy minimum, so we start by minimizing the energy.

```python
N = 128
dimension = 3

box_size = quantity.box_size_at_number_density(N, 1.4, dimension)
displacement, shift = space.periodic(box_size)

energy_fn = energy.soft_sphere_pair(displacement)

key, split = random.split(key)
R_init = random.uniform(split, (N,dimension), minval=0.0, maxval=box_size, dtype=f64)

R, max_grad, niters = run_minimization_while(energy_fn, R_init, shift)
print('Minimized the energy in {} minimization steps and reached a final \
maximum gradient of {}'.format(niters, max_grad))
```

We can now calculate the elastic modulus tensor

```python
emt_fn = jit(elasticity.athermal_moduli(energy_fn, check_convergence=True))
C, converged = emt_fn(R,box_size)
print(converged)
```

The elastic modulus tensor gives a quantitative prediction for how the energy should change if we deform the system according to a strain tensor
\begin{equation}
\frac{ \Delta U}{V^0} = \sigma^0\epsilon + \frac 12 \epsilon C  \epsilon + O\left(\epsilon^3\right)
\end{equation}
To test this, we define $\epsilon = \tilde \epsilon \gamma$ for a randomly chosen strain tensor $\tilde \epsilon$ and for $\gamma << 1$. Ignoring terms of order $\gamma^3$ and higher, we have
\begin{equation}
\frac{ \Delta U}{V^0} - \sigma^0\epsilon = \left[\frac 12 \tilde \epsilon C  \tilde \epsilon \right] \gamma^2
\end{equation}
Thus, we can test our calculation of $C$ by plotting $\frac{ \Delta U}{V^0} - \sigma^0\epsilon$ as a function of $\gamma$ for our randomly chosen $\tilde \epsilon$ and comparing it to the line $\left[\frac 12 \tilde \epsilon C  \tilde \epsilon \right] \gamma^2$.

First, generate a random $\tilde \epsilon$ and calculate $U$ for different $\gamma$.

```python
key, split = random.split(key)
#Pick a random (symmetric) strain tensor
strain_tensor = random.uniform(split, (dimension,dimension), minval=-1, maxval=1, dtype=f64)
strain_tensor = (strain_tensor + strain_tensor.T) / 2.0

#Define a function to calculate the energy at a given strain
def get_energy_at_strain(gamma, strain_tensor, R_init, box):
  R_init = space.transform(space.inverse(box),R_init)
  new_box = jnp.matmul(jnp.eye(strain_tensor.shape[0]) + gamma * strain_tensor, box)
  displacement, shift = space.periodic_general(new_box, fractional_coordinates=True)
  energy_fn = energy.soft_sphere_pair(displacement, sigma=1.0)
  R_final, _, _ = run_minimization_while(energy_fn, R_init, shift)
  return energy_fn(R_final)

gammas = jnp.logspace(-7,-4,50)
Us = vmap(get_energy_at_strain, in_axes=(0,None,None,None))(gammas, strain_tensor, R, box_size * jnp.eye(dimension))
```

Plot $\frac{ \Delta U}{V^0} - \sigma^0\epsilon$ and $\left[\frac 12 \tilde \epsilon C  \tilde \epsilon \right] \gamma^2$ as functinos of $\gamma$. While there may be disagreements for very small $\gamma$ due to numerical precision or at large $\gamma$ due to higher-order terms becoming relevant, there should be a region of quantitative agreement.

```python
U_0 = energy_fn(R)
stress_0 = -quantity.stress(energy_fn, R, box_size)
V_0 = quantity.volume(dimension, box_size)

#Plot \Delta E/V - sigma*epsilon
y1 = (Us - U_0)/V_0 - gammas * jnp.einsum('ij,ji->',stress_0,strain_tensor)
plt.plot(jnp.abs(gammas), y1, lw=3, label=r'$\Delta U/V^0 - \sigma^0 \epsilon$')

#Plot 0.5 * epsilon*C*epsilon
y2 = 0.5 * jnp.einsum('ij,ijkl,kl->',strain_tensor, C, strain_tensor) * gammas**2
plt.plot(jnp.abs(gammas), y2, ls='--', lw=3, label=r'$(1/2) \epsilon C \epsilon$')

plt.xscale('log')
plt.yscale('log')
plt.legend()
format_plot('$\gamma$','')
finalize_plot()
```

To test the accuracy of this agreement, we first define:
\begin{equation}
T(\gamma) = \frac{ \Delta U}{V^0} - \sigma^0\epsilon - \frac 12 \epsilon C  \epsilon \sim O\left(\gamma^3\right)
\end{equation}
which should be proportional to $\gamma^3$ for small $\gamma$ (note that this expected scaling should break down when the y-axis approaches machine precision). This is a prediction of scaling only, so we plot a line proportional to $\gamma^3$ to compare the slopes.

```python
#Plot the difference, which should scales as gamma**3
plt.plot(jnp.abs(gammas), jnp.abs(y1-y2), label=r'$T(\gamma)$')
#Plot gamma**3 for reference
plt.plot(jnp.abs(gammas), jnp.abs(gammas**3), 'black', label=r'slope = $\gamma^3$ (for reference)')

plt.xscale('log')
plt.yscale('log')
plt.legend()
format_plot('$\gamma$','')
finalize_plot()
```

Save `C` for later testing.

```python
C_3d = C
```

#Example with neighbor lists

As a second example, consider a much larger systems that is implemented using neighbor lists.

```python
N = 5000
dimension = 2

box_size = quantity.box_size_at_number_density(N, 1.3, dimension)
box = box_size * jnp.eye(dimension)
displacement, shift = space.periodic_general(box, fractional_coordinates=True)

sigma = jnp.array([[1.0, 1.2], [1.2, 1.4]])
N_2 = int(N / 2)
species = jnp.where(jnp.arange(N) < N_2, 0, 1)

neighbor_fn, energy_fn = energy.soft_sphere_neighbor_list(
    displacement, box_size, species=species, sigma=sigma, dr_threshold = 0.1,
    fractional_coordinates = True)

key, split = random.split(key)
R_init = random.uniform(split, (N,dimension), minval=0.0, maxval=1.0, dtype=f64)

R, max_grad, nbrs, niters = run_minimization_while_neighbor_list(energy_fn, neighbor_fn, R_init, shift)
print('Minimized the energy in {} minimization steps and reached a final \
maximum gradient of {}'.format(niters, max_grad))
```

We have to pass the neighbor list to `emt_fn`.

```python
emt_fn = jit(elasticity.athermal_moduli(energy_fn, check_convergence=True))
C, converged = emt_fn(R,box,neighbor=nbrs)
print(converged)
```

We can time the calculation of the compiled function.

```python
%timeit emt_fn(R,box,neighbor=nbrs)
```

Repeat the same tests as above. NOTE: this may take a few minutes.

```python
key, split = random.split(key)
#Pick a random (symmetric) strain tensor
strain_tensor = random.uniform(split, (dimension,dimension), minval=-1, maxval=1, dtype=f64)
strain_tensor = (strain_tensor + strain_tensor.T) / 2.0

def get_energy_at_strain(gamma, strain_tensor, R_init, box):
  new_box = jnp.matmul(jnp.eye(strain_tensor.shape[0]) + gamma * strain_tensor, box)
  displacement, shift = space.periodic_general(new_box, fractional_coordinates=True)

  neighbor_fn, energy_fn = energy.soft_sphere_neighbor_list(
    displacement, box_size, species=species, sigma=sigma, dr_threshold = 0.1,
    fractional_coordinates = True, capacity_multiplier = 1.5)

  R_final, _, nbrs, _ = run_minimization_while_neighbor_list(energy_fn, neighbor_fn, R_init, shift)
  return energy_fn(R_final, neighbor=nbrs)

gammas = jnp.logspace(-7,-3,20)
Us = jnp.array([ get_energy_at_strain(gamma, strain_tensor, R, box) for gamma in gammas])
```

```python
U_0 = energy_fn(R, neighbor=nbrs)
stress_0 = -quantity.stress(energy_fn, R, box, neighbor=nbrs)
V_0 = quantity.volume(dimension, box)

#Plot \Delta E/V - sigma*epsilon
y1 = (Us - U_0)/V_0 - gammas * jnp.einsum('ij,ji->',stress_0,strain_tensor)
plt.plot(jnp.abs(gammas), y1, lw=3, label=r'$\Delta U/V^0 - \sigma^0 \epsilon$')

#Plot 0.5 * epsilon*C*epsilon
y2 = 0.5 * jnp.einsum('ij,ijkl,kl->',strain_tensor, C, strain_tensor) * gammas**2
plt.plot(jnp.abs(gammas), y2, ls='--', lw=3, label=r'$(1/2) \epsilon C \epsilon$')

plt.xscale('log')
plt.yscale('log')
plt.legend()
format_plot('$\gamma$','')
finalize_plot()
```

```python
#Plot the difference, which should scales as gamma**3
plt.plot(jnp.abs(gammas), jnp.abs(y1-y2), label=r'$T(\gamma)$')
#Plot gamma**3 for reference
plt.plot(jnp.abs(gammas), jnp.abs(gammas**3), 'black', label=r'slope = $\gamma^3$ (for reference)')

plt.xscale('log')
plt.yscale('log')
plt.legend()
format_plot('$\gamma$','')
finalize_plot()
```

Save `C` for later testing.

```python
C_2d = C
```

#Mandel notation

Mandel notation is a way to represent symmetric second-rank tensors and fourth-rank tensors with so-called "minor symmetries", i.e. $T_{ijkl} = T_{ijlk} = T_{jilk}$. The idea is to map pairs of indices so that $(i,i) \rightarrow i$ and $(i,j) \rightarrow K - i - j$ for $i\neq j$, where $K = d(d+1)/2$ is the number of independent pairs $(i,j)$ for tensors with $d$ elements along each axis. Thus, second-rank tensors become first-rank tensors, and fourth-rank tensors become second-rank tensors, according to:

\begin{align}
M_{m(i,j)} &= T_{ij} w(i,j) \\
M_{m(i,j),m(k,l)} &= T_{ijkl} w(i,j) w(k,l).
\end{align}

Here, $m(i,j)$ is the mapping function described above, and w(i,j) is a weight that preserves summation rules and is given by
\begin{align}
w(i,j) = \delta_{ij} + \sqrt{2} (\delta_{ij}-1).
\end{align}

We can convert strain tensors, stress tensors, and elastic modulus tensors to and from Mandel notation using the functions `elasticity.tensor_to_mandel` and `elasticity.mandel_to_tensor`.

First, lets copy one of the previously calculated elastic modulus tensors and define a random strain tensor.

```python
#This can be 2 or 3 depending on which of the above solutions has been calculated
dimension = 3
if dimension == 2:
  C = C_2d
else:
  C = C_3d
key, split = random.split(key)
e = random.uniform(key, (dimension,dimension), minval=-1, maxval=1, dtype=f64)
e = (e + e.T)/2.
```

Convert `e` and `C` to Mental notation

```python
e_m = jit(elasticity.tensor_to_mandel)(e)
C_m = jit(elasticity.tensor_to_mandel)(C)
print(e_m)
print(C_m)
```

Using "bar" notation to represent Mandel vectors and matrices, we have
\begin{equation}
\frac{ \Delta U}{V^0} = \bar \sigma_i^0 \bar\epsilon_i + \frac 12 \bar \epsilon_i \bar C_{ij} \bar\epsilon_j + O\left(\bar \epsilon^3\right)
\end{equation}
We can explicity test that the sums are equivalent to the sums involving the original tensors

```python
sum_m = jnp.einsum('i,ij,j->',e_m, C_m, e_m)
sum_t = jnp.einsum('ij,ijkl,kl->',e, C, e)
print('Relative error is {}, which should be very close to 0'.format((sum_t-sum_m)/sum_t))
```

Finally, we can convert back to the full tensors and check that they are unchanged.

```python
C_new = jit(elasticity.mandel_to_tensor)(C_m)
print('Max error in C is {}, which should be very close to 0.'.format(jnp.max(jnp.abs(C-C_new))))

e_new = jit(elasticity.mandel_to_tensor)(e_m)
print('Max error in e is {}, which should be very close to 0.'.format(jnp.max(jnp.abs(e-e_new))))
```

# Isotropic elastic constants

The calculation of the elastic modulus tensor does not make any assumptions about the underlying symmetries in the material. However, for isotropic systems, only two constants are needed to completely describe the elastic behavior. These are often taken to be the bulk modulus, $B$, and the shear modulus, $G$, or the Young's modulus, $E$, and the Poisson's ratio, $\nu$. The function `elasticity.extract_isotropic_moduli` extracts these values, as well as the longitudinal modulus, $M$, from an elastic modulus tensor.

Importantly, since there is not guarantee that `C` is calculated from a truely isotropic systems, these are "orientation-averaged" values. For example, there are many directions in which you can shear a system, and the shear modulus that is returned represents and average over all these orientations. This can be an effective way to average over small fluctuations in an "almost isotropic" system, but the values lose their typical meaning when the systems is highly anisotropic.

```python
elasticity.extract_isotropic_moduli(C)
```

# Gradients

The calculation of the elastic modulus tensor is fully differentiable:

```python
def setup(N,dimension,key):
  box_size = quantity.box_size_at_number_density(N, 1.4, dimension)
  box = box_size * jnp.eye(dimension)
  displacement, shift = space.periodic_general(box, fractional_coordinates=True)
  R_init = random.uniform(key, (N,dimension), minval=0.0, maxval=1.0, dtype=f64)

  def run(sigma):
    energy_fn = energy.soft_sphere_pair(displacement, sigma=sigma)
    R, max_grad = run_minimization_scan(energy_fn, R_init, shift, num_steps=1000)
    emt_fn = jit(elasticity.athermal_moduli(energy_fn))
    C = emt_fn(R,box)
    return elasticity.extract_isotropic_moduli(C)['G']
  return run

key, split = random.split(key)
N = 50
dimension = 2
run = setup(N, dimension, split)
```

```python
sigma = jnp.linspace(1.0,1.4,N)
print(run(sigma))
print(grad(run)(sigma))
```