---
## 02. Linear Fit Example 01.

Eduard Larrañaga (ealarranaga@unal.edu.co)

---

### About this notebook

In this worksheet, we consdier the discovery of the Expansion of the Universe by Edwin Hubble as example of a linear fit of experimental data.

---

```python
import numpy as np
from matplotlib import pyplot as plt
#import seaborn as sns
#sns.set()
import pandas as pd
%matplotlib inline
```

### Edwin Hubble's data

Around the 1920's, Edwin Hubble show that the "nebulae" were external galaxies and not part of our own Galaxy, the Milky Way. In a seminal paper,

https://ui.adsabs.harvard.edu/abs/1931ApJ....74...43H

E. Hubble and M. Humason determined that some of these galaxies moved away from Earth with a velocity $v$ that is proportional to their distance $d$, i.e.

$v=H_0 d.$

This relation is now known as *Hubble's law* and quantity $H_0$ is called the *Hubble constant*. It is usual to give the value of $H_0$ in units of $\textrm{km}\, \textrm{s}^{-1} \, \textrm{Mpc}^{-1}$.
The original data of Hubble and Humason is summarized in the data file `hubble.csv`.

```python
df = pd.read_csv("data/hubble.csv")

df
```

      Name
      N_measurements
      velocity
      mean_m

      0
      Virgo
      7
      890
      12.5

      1
      Pegasus
      5
      3810
      15.5

      2
      Pisces
      4
      4630
      15.4

      3
      Cancer
      2
      4820
      16.0

      4
      Perseus
      4
      5230
      16.4

      5
      Coma
      3
      7500
      17.0

      6
      Ursa_Major
      1
      11800
      18.0

      7
      Leo
      1
      19600
      19.0

      8
      [No_Name]
      16
      2350
      13.8

      9
      [No_Name]
      16
      630
      11.6

```python
df.describe()
```

      N_measurements
      velocity
      mean_m

      count
      10.000000
      10.000000
      10.000000

      mean
      5.900000
      6126.000000
      15.520000

      std
      5.626327
      5750.565576
      2.331332

      min
      1.000000
      630.000000
      11.600000

      25%
      2.250000
      2715.000000
      14.200000

      50%
      4.000000
      4725.000000
      15.750000

      75%
      6.500000
      6932.500000
      16.850000

      max
      16.000000
      19600.000000
      19.000000

The data in the data frame includes 10 samples (Nebulae) with 4 features:

**Name** : Name of the nebula \
**N_measurement** : Number of velocities measured by Hubble and Humason \
**velocity** : Mean velocity of the nebula measured in km/s \
**mean_m** : Apparent magnitude of the nebula

___

The apaprent magnitude is related to the distance in parsecs through the relation

$\log_{10} d = \frac{m-M+5}{5}$

where $M=-13.8$ is the absolute magnitude reported by Hubble and considered as a constant in the paper.

The relation of velocity vs. apparent magnitude gives the plot

```python
plt.scatter(df['mean_m'], df['velocity'])
plt.xlabel(r'$m$')
plt.ylabel(r'velocity [km/s]')
plt.show()
```

### Covariance

We will calculate the covariance between these features.

```python
def cov(x,y):
    N = len(x)
    mu_x = sum(x)/N
    mu_y = sum(y)/N
    return sum((x - mu_x)*(y - mu_y))/N

cov(df['mean_m'],df['velocity'])
```

    10565.68

This result shows that there is some kidn of (positive) covariance. however, it is not clear if the features have a linear relation between them.

### Correlation Coefficient
Now, we calculate the correlation coefficient  of these features.

```python
def corr(x,y):
    N = len(x)
    num = N*sum(x*y)- sum(x)*sum(y)
    den = np.sqrt(N*sum(x*x) - sum(x)*sum(x))*np.sqrt(N*sum(y*y) - sum(y)*sum(y))
    return num/den

corr(df['mean_m'],df['velocity'])
```

    0.8756693926772317

This result, together with the plot above, shows that the relation is not completly linear. However, we will introduce a column in the dataframe including the logarithm of the velocity to show that it is possible to obtain a good linear behavior.

```python
df['log10_velocity'] = np.log10(df['velocity'])
df
```

      Name
      N_measurements
      velocity
      mean_m
      log10_velocity

      0
      Virgo
      7
      890
      12.5
      2.949390

      1
      Pegasus
      5
      3810
      15.5
      3.580925

      2
      Pisces
      4
      4630
      15.4
      3.665581

      3
      Cancer
      2
      4820
      16.0
      3.683047

      4
      Perseus
      4
      5230
      16.4
      3.718502

      5
      Coma
      3
      7500
      17.0
      3.875061

      6
      Ursa_Major
      1
      11800
      18.0
      4.071882

      7
      Leo
      1
      19600
      19.0
      4.292256

      8
      [No_Name]
      16
      2350
      13.8
      3.371068

      9
      [No_Name]
      16
      630
      11.6
      2.799341

```python
df.describe()
```

      N_measurements
      velocity
      mean_m
      log10_velocity

      count
      10.000000
      10.000000
      10.000000
      10.000000

      mean
      5.900000
      6126.000000
      15.520000
      3.600705

      std
      5.626327
      5750.565576
      2.331332
      0.461869

      min
      1.000000
      630.000000
      11.600000
      2.799341

      25%
      2.250000
      2715.000000
      14.200000
      3.423532

      50%
      4.000000
      4725.000000
      15.750000
      3.674314

      75%
      6.500000
      6932.500000
      16.850000
      3.835921

      max
      16.000000
      19600.000000
      19.000000
      4.292256

```python
plt.scatter(df['mean_m'], df['log10_velocity'])
plt.xlabel(r'apparent magnitude')
plt.ylabel(r'logarithm of velocity')
plt.show()
```

```python
cov(df['mean_m'],df['log10_velocity'])
```

    0.9620255525348563

```python
corr(df['mean_m'],df['log10_velocity'])
```

    0.9927066870715695

This result and the plot above, show that there is a linear relation between these features. Hence, we will obtain a linear fit.

### Linear Fit

Since the plot of the logarithm of the velocity vs. apparent magnitude has a linear tendence, we will create a linear fit for this data (including the possibility of having error in the y measurements).

```python
import LinearFit as lf
```

Since we have no information about the observational error, we use a constant value of $\sigma_i = 1$

```python
number_rows = df['log10_velocity'].count()
error_y = np.ones(number_rows)
```

Now we make the linear fit

```python
a_1, a_2 , sigma_a1, sigma_a2, chi2, R2 = lf.linear_fit(df['mean_m'],df['log10_velocity'],error_y)
a_1, a_2, chi2, R2
```

    (0.5484040402255217,
     0.19666889208743046,
     0.027902852936639286,
     0.9922507256121249)

The obtained relation fits very well the observational data, with a value of $\chi ^2 = 0.028 $ and a score of $R^2 = 0.99$. The plot shows the good fit

```python
m_range= np.linspace(10,20,40)
logV = a_1 + a_2*m_range

plt.figure(figsize=(5,5))
plt.scatter(df['mean_m'], df['log10_velocity'],label='observational data')
plt.plot(m_range,logV,'--k', label='linear fit')

plt.xlabel(r'apparent magnitude')
plt.ylabel(r'logarithm of velocity')
plt.legend()
plt.show()
```

The obtained linear model gives the relation between velocity and apparent magnitude as

$\log_{10} v = a_1 + a_2 m = 0.548 + 0.197 m$.

At this point we will assume, as Hubble and Humason did, that the coefficient of $m$ in this equation will be $0.2$, and therefore the equation becomes

$\log_{10} v = a_1 + a_2 m = 0.548 + 0.2 m$.

From the expression for the distance, we have

$\log_{10} d = \frac{m-M+5}{5} \,  \rightarrow \, \log_{10} d = 0.2m - 0.2M + 1$

and therefore we can write

$\log_{10} \left( \frac{v}{d} \right) = \log_{10}v - \log_{10} d = a_1  - 1+ 0.2M $.

Using  the obtained value for $a_1$ and $M=-13.8$ we get

$\log_{10} \left( \frac{v}{d} \right) = -3.212$,

which gives the Hubble constant

$H_0 = \frac{v}{d} = 10^{-3.212} = 614 \times 10^{-6} \textrm{ km } \textrm{s}^{-1} \textrm{ pc}^{-1} = 614 \textrm{ km } \textrm{s}^{-1} \textrm{ Mpc}^{-1}$

---

### Alternative Linear Fit

We can perform another linear fit, introducing directly the distance as a feature in the dataframe,

```python
M = -13.8
df['log10_distance'] = (df['mean_m'] - M + 5.)/5.
df
```

      Name
      N_measurements
      velocity
      mean_m
      log10_velocity
      log10_distance

      0
      Virgo
      7
      890
      12.5
      2.949390
      6.26

      1
      Pegasus
      5
      3810
      15.5
      3.580925
      6.86

      2
      Pisces
      4
      4630
      15.4
      3.665581
      6.84

      3
      Cancer
      2
      4820
      16.0
      3.683047
      6.96

      4
      Perseus
      4
      5230
      16.4
      3.718502
      7.04

      5
      Coma
      3
      7500
      17.0
      3.875061
      7.16

      6
      Ursa_Major
      1
      11800
      18.0
      4.071882
      7.36

      7
      Leo
      1
      19600
      19.0
      4.292256
      7.56

      8
      [No_Name]
      16
      2350
      13.8
      3.371068
      6.52

      9
      [No_Name]
      16
      630
      11.6
      2.799341
      6.08

The linear behavior can be seen in a plot of $\log_{10} v$ vs. $\log_{10} d$,

```python
plt.scatter(df['log10_distance'], df['log10_velocity'])
plt.xlabel(r'logarithm of distance')
plt.ylabel(r'logarithm of velocity')
plt.show()
```

Hence, we create the linear fit between these variables.

```python
number_rows = df['log10_velocity'].count()
error_y = np.ones(number_rows)

a_1, a_2 , sigma_a1, sigma_a2, chi2, R2 = lf.linear_fit(df['log10_distance'],df['log10_velocity'],error_y)
a_1, a_2, chi2, R2
```

    (-3.1489711310180266,
     0.9833444604370973,
     0.027902852936639126,
     0.992250725612125)

One again, the fit has a value of $\chi ^2 = 0.028 $ and a score of $R^2 = 0.99$. The plot shows the good fit:

```python
distance_range= np.linspace(6,8,20)
logV = a_1 + a_2*distance_range

plt.figure(figsize=(5,5))
plt.scatter(df['log10_distance'], df['log10_velocity'],label='observational data')
plt.plot(distance_range,logV,'--k', label='linear fit')

plt.xlabel(r'logarithm of distance')
plt.ylabel(r'logarithm of velocity')
plt.legend()
plt.show()
```

The obtained linear model gives the relation between velocity and distance as

$\log_{10} v = a_1 + a_2 \log_{10} d = -3.15 + 0.98 \log_{10} d$.

At this point we will assume that the coefficient of $\log_{10} d$ in this equation will be approximately $1$, and therefore the equation becomes

$\log_{10} v = a_1 + a_2 m = -3.15 +\log_{10}d$.

From this expression we have

$\log_{10} \left( \frac{v}{d} \right) = \log_{10}v - \log_{10} d = -3.15 $,

and therefore the Hubble constant is

$H_0 = \frac{v}{d} = 10^{-3.15} = 708 \times 10^{-6} \textrm{ km } \textrm{s}^{-1} \textrm{ pc}^{-1} = 708 \textrm{ km } \textrm{s}^{-1} \textrm{ Mpc}^{-1}$

---

## Linear Regression with `Numpy`

The function [numpy.linalg.lstsg](https://numpy.org/doc/stable/reference/generated/numpy.linalg.lstsq.html) provides an efficient algorithm to obtain a linear regression.

In order to use it, we need to transform our dataframe's columns into numpy arrays,

```python
x = df['mean_m'].to_numpy()
y = df['log10_velocity'].to_numpy()
x
```

    array([12.5, 15.5, 15.4, 16. , 16.4, 17. , 18. , 19. , 13.8, 11.6])

Now, we will put the linear system

\begin{equation}
y = a_1 + a_2 x
\end{equation}

in the form

\begin{equation}
y = Ap
\end{equation}

where $A$ is the matrix

\begin{equation}
A = \pmatrix{x_1 & 1\\
             x_2 & 1\\
             x_3 & 1\\
             ... & ...}
\end{equation}

and $p$ is the 2-dimensional vector

\begin{equation}
p = \pmatrix{a_2 \\ a_1}.
\end{equation}

Then, we define $A$ as

```python
A = np.vstack([x, np.ones(len(x))]).T
A
```

    array([[12.5,  1. ],
           [15.5,  1. ],
           [15.4,  1. ],
           [16. ,  1. ],
           [16.4,  1. ],
           [17. ,  1. ],
           [18. ,  1. ],
           [19. ,  1. ],
           [13.8,  1. ],
           [11.6,  1. ]])

and apply the function `np.linalg.lstsq` to $A$ and $y$ to obtain the slope and the intercept of the regression line,

```python
a_2, a_1 = np.linalg.lstsq(A, y, rcond=None)[0] # slope, intercept
a_2, a_1
```

    (0.1966688920874267, 0.548404040225562)

The complete set of returns of this function include
- Least-squares solution: If b is two-dimensional, the solutions are in the K columns of x.
- Residuals: Sums of squared residuals ($\chi^2$)
- Rank: Rank of matrix A.
- s(min(M, N),): Singular values of A.

```python
np.linalg.lstsq(A, y, rcond=None)
```

    (array([0.19666889, 0.54840404]),
     array([0.02790285]),
     2,
     array([49.67314921,  0.44524984]))

Here you can see that this fit gives a value of $\chi^2 = 0.02790285 $

```python
distance_range= np.linspace(min(x)-1, max(x)+1, 20)
logV = a_1 + a_2*distance_range

plt.figure(figsize=(5,5))
plt.scatter(x, y, label='observational data')
plt.plot(distance_range, logV,'--k', label='linear fit')

plt.xlabel(r'logarithm of distance')
plt.ylabel(r'logarithm of velocity')
plt.legend()
plt.show()
```

```python

```


# InfoNCE

## Setup

The following cell imports the necessary packages:

```python
import torch
import torch.nn.functional as F
import os
import numpy as np
import matplotlib.pyplot as plt
import argparse
from src.utils import *
from src.mi_estimators import *
from torch.utils.tensorboard import SummaryWriter
```

Set random seed for reproducibility.

```python
SEED = 0
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
np.random.seed(SEED)

cuda = True if torch.cuda.is_available() else False
FloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
LongTensor = torch.cuda.LongTensor if cuda else torch.LongTensor
torch.set_default_tensor_type(FloatTensor)
```

Specify the hyper-parameters:

```python
parser = argparse.ArgumentParser()
parser.add_argument("--rho", type=float, default=0.9, help="coefficient of Gaussian")
parser.add_argument("--d", type=int, default=20, help="dimension of X & Y")
parser.add_argument("--sample_size", type=int, default=400, help="sample size")
parser.add_argument("--gamma", type=float, default=1e-10, help="clipping parameter")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--batch_size", type=int, default=40, help="size of the batches")
parser.add_argument("--lr", type=float, default=1e-3, help="adam: learning rate")
parser.add_argument("--hidden_dim", type=int, default=100, help="Hidden dimension")
parser.add_argument("--hidden_layer", type=int, default=2, help="the number of hidden layers")
parser.add_argument("--ma_rate", type=float, default=0.1, help="move average rate")
parser.add_argument("--ma_ef", type=float, default=1, help="move average ef")
parser.add_argument("--exc", type=int, default=1, help="whether exclude the data sample from the reference samples")
parser.add_argument("--reg", type=int, default=1, help="if apply regularization")
parser.add_argument("--n_epoch", type=int, default=2000, help="number of epochs of training")
parser.add_argument("--n_iters_1epoch", type=int, default=10, help="number of epochs of training")

opt, unknown = parser.parse_known_args()
opt.n_iters = opt.n_epoch * opt.n_iters_1epoch
ma_rate = 0.01  # moving average rate
```

```python
continue_train = False  # set to True to continue to train
load_available = True # set to False to prevent loading previous results
overwrite = False  # set to True to overwrite previously stored results
```

Generate the data.

```python
data = GaussianData(opt.sample_size, d=opt.d, rho=opt.rho)
X, Y, XY, Ground_truth = data.X, data.Y, torch.cat((data.X, data.Y), dim=1), data.mutual_information()
```

## Initialize the Model

```python
name = f'./results/InfoNCE_dim{opt.d}_reg{opt.reg}'   # filename
chkpt_name = name+'.pt'      # checkpoint

from datetime import datetime
TIMESTAMP = "{0:%Y-%m-%dT%H-%M-%S/}".format(datetime.now())
writer = SummaryWriter(f'./results/log/InfoNCE_dim{opt.d}_reg{opt.reg}/{TIMESTAMP}')
discriminator = Net(input_size=opt.d*2, hidden_layers=opt.hidden_layer, hidden_size=opt.hidden_dim)

# move NN model to GPU if GPU is available
if cuda:
    discriminator.cuda()

# Adam optimizer
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))
```

Define the original InfoNCE loss function and its exclusion version

```python
def infonce_loss(net, x_samples, y_samples, writer=None, epoch=None, reg=False):
    sample_size = x_samples.shape[0]
    x_tile = x_samples.unsqueeze(0).repeat((sample_size, 1, 1))
    y_tile = y_samples.unsqueeze(1).repeat((1, sample_size, 1))

    T0 = net(torch.cat([x_samples,y_samples], dim = -1))
    T1 = net(torch.cat([x_tile, y_tile], dim = -1))
    a = T0.mean()
    b = T1.logsumexp(dim = 1).mean() - np.log(sample_size)
    if reg:
        lower_bound = a - b - b**2
    else:
        lower_bound = a - b

    if writer is not None:
        writer.add_scalar('a', a, epoch)
        writer.add_scalar('b', b, epoch)
        writer.add_scalar('log ratio', a-b-np.log(sample_size), epoch)
    return lower_bound
```

```python
def infonce_exc_loss(net, x_samples, y_samples, writer=None, epoch=None, reg=False):
    sample_size = x_samples.shape[0]
    x_tile = x_samples.unsqueeze(0).repeat((sample_size, 1, 1))
    y_tile = y_samples.unsqueeze(1).repeat((1, sample_size, 1))
    mask =  torch.ones([sample_size]).diag().unsqueeze(-1).bool().cuda()
    T0 = net(torch.cat([x_samples,y_samples], dim = -1))
    T1 = net(torch.cat([x_tile, y_tile], dim = -1))
    a = T0.mean()
    b = T1[~mask].reshape(-1, sample_size-1, 1).logsumexp(dim = 0).mean() - np.log(sample_size-1)

    if reg:
        lower_bound = a - b - b**2
    else:
        lower_bound = a - b

    if writer is not None:
        writer.add_scalar('a', a, epoch)
        writer.add_scalar('b', b, epoch)
    return lower_bound
```

Load previous results.

```python
if load_available and os.path.exists(chkpt_name):
    checkpoint = torch.load(
        chkpt_name, map_location='cuda' if torch.cuda.is_available() else 'cpu')
    mi_list = checkpoint['mi_list']
    model_state = checkpoint['model_state']
    discriminator.load_state_dict(model_state)
    print('Previous results loaded.')
else:
    mi_list = [] # storing the mi estimation of each iteration
```

    Previous results loaded.

## Training

Since maximizing the infoNCE bound over $f$ attains the mutual information

\begin{align}
I(X\wedge Y) = \sup_{f} I_{\text{infoNCE}}(f),
\end{align}

we train the neural network to minimize the loss $-I_{\text{infoNCE}}(f)$. The optimal solution $f$ satisfies

\begin{align}
f(x,y)=\log \frac{p_{Y|X}(y|x)}{p_{Y}(y)} + c \quad \forall x\in \mathcal{X}, y \in \mathcal{Y}
\end{align}

for any constant $c$.

The following cell train the neural network using the data samples.

```python
# continue_train = True  # set to True to continue to train
if continue_train:
    iter_ = 0
    for i in range(opt.n_epoch*opt.n_iters_1epoch):
        data_joint = resample(XY, batch_size=opt.batch_size)
        train_x = data_joint[:,0:opt.d]
        train_y = data_joint[:,opt.d:]
        optimizer_D.zero_grad()
        loss = -infonce_loss(discriminator, train_x, train_y, reg=opt.reg) # negative infonce_bound as the loss
        loss.backward()

        optimizer_D.step()
        with torch.no_grad():
            # mi_est = infonce_loss(discriminator, X, Y, writer, i)
            mi_est = infonce_loss(discriminator, X, Y)
        mi_list.append(mi_est.item())
        writer.add_scalar('mi_list', mi_est.item(), i)
        writer.add_scalar('loss', loss, i)
        iter_ += 1

        if i%1000==0:
            print("Iternation: %d, loss: %f, mi_est: %f"%(iter_, loss.item(), mi_est))
            fig = plot_fig(discriminator, X, Y, d= opt.d if opt.d<6 else 6 )
            writer.add_figure('heatmap', fig, i)
            writer.add_histogram('first layer', discriminator.fc[0].weight.data, i)
            writer.add_histogram('second layer', discriminator.fc[1].weight.data, i)
            writer.add_histogram('third layer', discriminator.fc[2].weight.data, i)
            writer.add_histogram('first layer (grad)', discriminator.fc[0].weight.grad.data, i)
            writer.add_histogram('second layer (grad)', discriminator.fc[1].weight.grad.data, i)
            writer.add_histogram('third layer (grad)', discriminator.fc[2].weight.grad.data, i)

# writer.add_graph(discriminator, (XY,))
writer.close()
```

```python
if overwrite or not os.path.exists(chkpt_name):
    model_state = discriminator.state_dict()
    torch.save({
        'mi_list': mi_list,
        'model_state': model_state
    }, chkpt_name)
    writer.close()
    print('Current results saved.')
```

## MI estimates

Use moving average to smooth the estimated MI in each training iteration.

```python
mi_copy = mi_list.copy()
for k in range(1,len(mi_list)):
    mi_copy[k] = (1-ma_rate) * mi_copy[k-1] + ma_rate * mi_copy[k]
```

Plot the MI estimation curve against the training iteration, together with the ground truth.

```python
plt.plot(mi_copy, label='MI estimate')
plt.axhline(Ground_truth,label='Ground Truth',linestyle='--',color='red')
plt.axhline(np.log(opt.sample_size),label='log $n$',linestyle='--',color='grey')
# for t in range(len(mi_copy)):
#     if (mi_copy[t]>.8*Ground_truth):
#         plt.axvline(t,label='80% reached',linestyle=':',color='green')
#         break
plt.xlabel('number of iterations')
plt.ylabel('MI estimation')
plt.legend()
# plt.savefig(f'./results/InfoNCE_dim{opt.d}_reg{opt.reg}.pdf')
```

```python
heatmap = plot_fig(discriminator, X, Y, vmin=-40, vmax=5, d=6)
# plt.savefig(f'./results/heatmap_InfoNCE_dim{opt.d}_reg{opt.reg}.pdf')
```

## InfoNCE w\o data samples (exclusion)

```python
name = f'./results/InfoNCE_exc_dim{opt.d}_reg{opt.reg}'   # filename
chkpt_name = name+'.pt'      # checkpoint
# opt.reg = 0
from datetime import datetime
TIMESTAMP = "{0:%Y-%m-%dT%H-%M-%S/}".format(datetime.now())
writer = SummaryWriter(f'./results/log/InfoNCE_exc_dim{opt.d}_reg{opt.reg}/{TIMESTAMP}')
discriminator2 = Net(input_size=opt.d*2, hidden_layers=opt.hidden_layer, hidden_size=opt.hidden_dim)

# move NN model to GPU if GPU is available
if cuda:
    discriminator2.cuda()

# Adam optimizer
optimizer_D = torch.optim.Adam(discriminator2.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))
```

```python
# load_available = True # set to False to prevent loading previous results
if load_available and os.path.exists(chkpt_name):
    checkpoint = torch.load(
        chkpt_name, map_location='cuda' if torch.cuda.is_available() else 'cpu')
    mi_list = checkpoint['mi_list']
    model_state = checkpoint['model_state']
    discriminator2.load_state_dict(model_state)
    print('Previous results loaded.')
else:
    mi_list = [] # storing the mi estimation of each iteration
```

    Previous results loaded.

```python
if continue_train:
    for i in range(opt.n_epoch*opt.n_iters_1epoch):
        data_joint = resample(XY, batch_size=opt.batch_size)
        train_x = data_joint[:,0:opt.d]
        train_y = data_joint[:,opt.d:]

        optimizer_D.zero_grad()
        loss = -infonce_exc_loss(discriminator2, train_x, train_y, reg=opt.reg) # negative infonce_bound as the loss
        loss.backward()

        optimizer_D.step()
        with torch.no_grad():
            mi_est = infonce_exc_loss(discriminator2, X, Y, writer, i)
        mi_list.append(mi_est.item())
        writer.add_scalar('mi_list', mi_est.item(), i)
        writer.add_scalar('loss', loss, i)

        if i%1000==0:
            print("Iternation: %d, loss: %f, mi_est: %f"%(i, loss.item(), mi_est))
            fig = plot_fig(discriminator2, X, Y, d=opt.d if opt.d<6 else 6)
            writer.add_figure('heatmap', fig, i)
            writer.add_histogram('first layer', discriminator2.fc[0].weight.data, i)
            writer.add_histogram('second layer', discriminator2.fc[1].weight.data, i)
            writer.add_histogram('third layer', discriminator2.fc[2].weight.data, i)
            writer.add_histogram('first layer (grad)', discriminator2.fc[0].weight.grad.data, i)
            writer.add_histogram('second layer (grad)', discriminator2.fc[1].weight.grad.data, i)
            writer.add_histogram('third layer (grad)', discriminator2.fc[2].weight.grad.data, i)

# writer.add_graph(discriminator2, (XY,))
writer.close()
```

```python
mi_copy2 = mi_list.copy()
for k in range(1,len(mi_list)):
    mi_copy2[k] = (1-ma_rate) * mi_copy2[k-1] + ma_rate * mi_copy2[k]
```

```python
# overwrite = True  # set to True to overwrite previously stored results
if overwrite or not os.path.exists(chkpt_name):
    model_state = discriminator2.state_dict()
    torch.save({
        'mi_list': mi_list,
        'model_state': model_state
    }, chkpt_name)
    writer.close()
    print('Current results saved.')
```

```python
plt.plot(mi_copy2, label='MI estimate')
plt.axhline(Ground_truth,label='Ground Truth',linestyle='--',color='red')
plt.axhline(np.log(opt.sample_size),label='log $n$',linestyle='--',color='grey')
plt.xlabel('number of iterations')
plt.ylabel('MI estimation')
plt.legend()
# plt.savefig(f'./results/InfoNCE_exc_dim{opt.d}_reg{opt.reg}.pdf')
```

```python
heatmap = plot_fig(discriminator2, X, Y, vmin=-8000, vmax=10000, d=6)
# plt.savefig(f'./results/heatmap_InfoNCE_exc_dim{opt.d}_reg{opt.reg}.pdf')
```

## Ground Truth Heapmap

```python
def GT(x, y):
    mi = -0.5 * np.log(1 - opt.rho**2) * opt.d
    c = opt.rho / (2 * (1 - opt.rho**2))
    return mi + c * (2 * x * y - opt.rho *(x**2 + y**2)).sum(dim=1)
```

```python
d = 6
DEVICE = "cuda"
f, axs = plt.subplots(nrows=d,ncols=d,sharex=True, sharey=True)
for i in range(d):
    for j in range(d):
        im = plot_subfigure_GT(GT, X, Y, dimX=i, dimY=j, x0=None, y0=None, xmin=-5, xmax=5, ymin=-5, ymax=5, xgrids=50, ygrids=50, ax=axs[i, j], show_details=False)

f.colorbar(im, ax=axs.ravel().tolist())
# plt.savefig(f'./results/GT_heatmap_dim{opt.d}.pdf')
plt.show()
```


```python
import numpy
from matplotlib import pyplot

# Imports from the pyfoomb package
import pyfoomb
print(f'Current package version of {pyfoomb.__name__}: {pyfoomb.__version__}')

from pyfoomb import BioprocessModel
from pyfoomb import Caretaker
from pyfoomb import Visualization
from pyfoomb import TimeSeries
from pyfoomb import Helpers
```

    C:\Users\hemmerich\AppData\Local\Continuum\miniconda3\envs\pyfoomb-dev\lib\site-packages\assimulo\solvers\odepack.py:701: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray
      N.array([[0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.],  # 5th order
    Could not find GLIMDA.

    Current package version of pyfoomb: 2.17.7

# 1. Example: Simple growth model

Implements a simple exponential growth model, with the specific growth rate described by the well-known Monod equation.

### Model equations
$$
\begin{align}
    \mu &= \mu_{max} \frac{c_S}{c_S + k_S}\\
    \frac{dc_X}{dt} &= \mu \cdot c_X & c_X(t_0) = c_{X,0}\\
    \frac{dc_S}{dt} &= -\frac{1}{Y_{X/S}} \cdot \frac{dc_X}{dt}  & c_S(t_0) = c_{S,0}\\
\end{align}
$$

### Parametrization
$$
\mathbf{\theta} = \left[ k_S \; \mu_{max} \; Y_{X/S} \; c_{S,0} \; c_{X,0}\right]^T
$$

```python
# Defines the model class
class BatchGrowth(BioprocessModel):

    # The Monod equation is defined as classmethod
    @classmethod
    def Monod(cls, S, mu_max, kS):
        return mu_max * S / (kS + S)

    def rhs(self, t, y):

        # Unpacks the state vector. The states are alphabetically case-insensitive ordered.
        S, X = y

        # Unpacks the model parameters.
        kS = self.model_parameters['kS']
        mu_max = self.model_parameters['mu_max']
        YXS = self.model_parameters['YXS']

        # Calculate the current growth rate
        mu = self.Monod(S, mu_max, kS)

        # Calculate state derivatives
        dXdt = mu * X
        dSdt = -1/YXS * dXdt

        return [dSdt, dXdt]

model_parameters = {
    'kS' : 0.05,
    'mu_max' : 0.4,
    'YXS' : 0.5,
}

initial_values = {
    'S0' : 20.0,
    'X0' : 0.01,
}

caretaker = Caretaker(
    bioprocess_model_class=BatchGrowth,
    model_parameters=model_parameters,
    initial_values=initial_values,
)
```

```python
simulation = caretaker.simulate(t=24)
_ = Visualization.show_kinetic_data(simulation, ncols=2)
```

# 2. Example: Simple growth with product formation

Extends the first example by product formation, which is modelled as purely growth coupled.

### Model equations
$$
\begin{align}
    \mu &= \mu_{max} \frac{c_S}{c_S + k_S}\\
    \frac{dc_X}{dt} &= \mu \cdot c_X & c_X(t_0) = c_{X,0}\\
    \frac{dc_S}{dt} &= -\frac{1}{Y_{X/S}} \cdot \frac{dc_X}{dt}  & c_S(t_0) = c_{S,0}\\
    \frac{dc_P}{dt} &= Y_{P/X} \cdot \frac{dc_X}{dt}  & c_P(t_0) = c_{P,0}\\
\end{align}
$$

### Parametrization
$$
\mathbf{\theta} = \left[ k_S \; \mu_{max} \; Y_{X/S} \; Y_{P/X} \; c_{S,0} \; c_{X,0} \; c_{P,0}\right]^T
$$

```python
# Defines the model class
class GrowthCoupledProduction(BioprocessModel):

    # The Monod equation is defined as classmethod
    @classmethod
    def Monod(cls, S, mu_max, kS):
        return mu_max * S / (kS + S)

    def rhs(self, t, y):

        # Unpacks the state vector. The states are alphabetically ordered.
        P, S, X = y

        # Unpacks the model parameters.
        kS = self.model_parameters['kS']
        mu_max = self.model_parameters['mu_max']
        YPX = self.model_parameters['YPX']
        YXS = self.model_parameters['YXS']

        # Calculate the current growth rate
        mu = self.Monod(S, mu_max, kS)

        # Calculate state derivatives
        dXdt = mu * X
        dSdt = -1/YXS * dXdt
        dPdt = YPX * dXdt

        return [dPdt, dSdt, dXdt]

model_parameters = {
    'kS' : 0.05,
    'mu_max' : 0.4,
    'YXS' : 0.5,
    'YPX' : 0.2,
}

initial_values = {
    'S0' : 20.0,
    'X0' : 0.01,
    'P0' : 0.0,
}

caretaker = Caretaker(
    bioprocess_model_class=GrowthCoupledProduction,
    model_parameters=model_parameters,
    initial_values=initial_values,
)
```

```python
simulation = caretaker.simulate(t=24)
_ = Visualization.show_kinetic_data(simulation)
```

# 3. Example: Modelling induction with events

Implements a variant of model 2, which activates product formation (i.e., the corresponding term) at a explicitly defined timepoint (parametrized by `t_ind`).

### Model equations
$$
\begin{align}
    \mu &= \mu_{max} \frac{c_S}{c_S + k_S}\\
    \frac{dc_X}{dt} &= \mu \cdot c_X\\
    \frac{dc_S}{dt} &= -\frac{1}{Y_{X/S}} \cdot \frac{dc_X}{dt} \\
    \frac{dc_P}{dt} &=
    \begin{cases}
         0 & t < t_{ind}\\
         Y_{P/X} \cdot \frac{dc_X}{dt} & t \ge t_{ind}\\
    \end{cases}\\
    c_X(t_0) &= c_{X,0}\\
    c_S(t_0) &= c_{S,0}\\
    c_P(t_0) &= c_{P,0}\\
\end{align}
$$

### Parametrization
$$
\mathbf{\theta} = \left[ k_S \; \mu_{max} \; Y_{X/S} \; Y_{P/X} \; t_{ind} \; c_{S,0} \; c_{X,0} \; c_{P,0}\right]^T
$$

```python
# Defines the model class
class InductionEvent(BioprocessModel):

    # The Monod equation is defined as classmethod
    @classmethod
    def Monod(cls, S, mu_max, kS):
        return mu_max * S / (kS + S)

    def rhs(self, t, y, sw):

        # Unpacks the state vector. The states are alphabetically ordered.
        P, S, X = y

        # Unpacks the model parameters.
        kS = self.model_parameters['kS']
        mu_max = self.model_parameters['mu_max']
        YPX = self.model_parameters['YPX']
        YXS = self.model_parameters['YXS']

        # Calculate the current growth rate
        mu = self.Monod(S, mu_max, kS)

        # Calculate state derivatives
        dXdt = mu * X
        dSdt = -1/YXS * dXdt

        if sw[0]:
            dPdt = YPX * dXdt
        else:
            dPdt = 0

        return [dPdt, dSdt, dXdt]

    def state_events(self, t, y, sw):
        t_ind = self.model_parameters['t_ind']
        event_t = t - t_ind
        return [event_t]

model_parameters = {
    'kS' : 0.05,
    'mu_max' : 0.4,
    'YXS' : 0.5,
    'YPX' : 0.2,
    't_ind' : 12.0,
}

initial_values = {
    'S0' : 20.0,
    'X0' : 0.01,
    'P0' : 0.0,
}

caretaker = Caretaker(
    bioprocess_model_class=InductionEvent,
    model_parameters=model_parameters,
    initial_values=initial_values,
)
```

```python
simulation = caretaker.simulate(t=24)
_ = Visualization.show_kinetic_data(simulation)
```

# 4. Example: A more complex bioprocess model

This model introduces the use of more complex description of the biological process.

Product formation is induced by the substrate, modelled as a Monod relation.

The growth rate is calculated according to a substrate inhibition, with a rather high inhibition constant meaning that takes severe effect for high substrate concentration. In addtion, the product is modelled as growth inhibitory.

Substrate consumption is composed from substrate needed for biomass formation, as well as needed for product formation. See also Takors R: *Kommentierte Formelsammlung Bioverfahrenstechnik* (2013) p.109

The use of classmethods is also demonstrated: First, they are used to "outsource" calculations during the integration. Second, they are used afterwards to calculate the trajectories of the specific rates from the integrated time series of the model states.

### Model equations

#### Specific rates
$$
\begin{align}
    \mu &= \mu_{max} \cdot \frac{c_S}{c_S + k_S + \frac{c_S^2}{k_{SI}}} \cdot \frac{c_P}{c_P + k_{PI}}\\
    q_P &= q_{P,max} \cdot \frac{c_S}{c_S + k_{SP}}\\
    q_S &= \frac{\mu}{Y_{X/S}} + \frac{q_P}{Y_{P/S}}\\
\end{align}
$$

#### State derivatives
$$
\begin{align}
    \frac{dc_X}{dt} &= \mu \cdot c_X\\
    \frac{dc_S}{dt} &= -q_S \cdot c_X\\
    \frac{dc_P}{dt} &= q_P \cdot c_X\\
\end{align}
$$

#### Initial values
$$
\begin{align}
     c_X(t_0) &= c_{X,0}\\
     c_S(t_0) &= c_{S,0}\\
     c_P(t_0) &= c_{P,0}\\
\end{align}
$$

### Parametrization
$$
\mathbf{\theta} = \left[ k_S \; k_{SI} \; k_{PI} \; \mu_{max} \; q_{P,max} \; k_{SP} \; Y_{X/S} \; Y_{P/S} \; c_{S,0} \; c_{X,0} \; c_{P,0}\right]^T
$$

```python
# Defines the model class
class InducingSubstrate(BioprocessModel):

    # The Monod equation with the growth-inhibiting product concentration
    @classmethod
    def growth_rate(cls, S, mu_max, kS, kSI, P, kPI):
        return mu_max * S / (kS + S + S**2/kSI) * kPI / (P + kPI)

    # The cell specific substrate consumption
    @classmethod
    def substrate_consumption(cls, mu, YXS, qP, YPS):
        return mu/YXS + qP/YPS

    # The cell specific catabolite repressed productivity, modelled as substrate inhibition
    @classmethod
    def product_formation(cls, S, qP_max, kSP):
        return qP_max * S / (kSP + S)

    def rhs(self, t, y):

        # Unpacks the state vector. The states are alphabetically ordered.
        P, S, X = y

        # Unpacks the model parameters.
        mu_max = self.model_parameters['mu_max']
        kS = self.model_parameters['kS']
        kPI = self.model_parameters['kPI']
        kSI = self.model_parameters['kSI']
        qP_max = self.model_parameters['qP_max']
        kSP = self.model_parameters['kSP']
        YXS = self.model_parameters['YXS']
        YPS = self.model_parameters['YPS']

        # Calculate the current growth rate
        mu = self.growth_rate(S, mu_max, kS, kSI, P, kPI)

        # Calculate the current cell specific productivity
        qP = self.product_formation(S, qP_max, kSP)

        # Calculate the current cell specific substrate consumption rate
        qS = self.substrate_consumption(mu, YXS, qP, YPS)

        # Calculate state derivatives
        dXdt = mu * X
        dSdt = -qS * X
        dPdt = qP * X

        return [dPdt, dSdt, dXdt]

# Define model parameters and initial values for Caretaker instantiation
model_parameters = {
    'kS' : 0.05,
    'kSP' : 0.1,
    'kSI' : 10.0,
    'kPI' : 0.1,
    'mu_max' : 0.4,
    'YXS' : 0.55,
    'YPS' : 0.2,
    'qP_max' : 0.05,
}

initial_values = {
    'S0' : 20.0,
    'X0' : 1.0,
    'P0' : 0.0,
}

# Get Caretaker instance
caretaker = Caretaker(
    bioprocess_model_class=InducingSubstrate,
    model_parameters=model_parameters,
    initial_values=initial_values,
)

# Run a simulation and visualize
simulation = caretaker.simulate(t=48)
_ = Visualization.show_kinetic_data(simulation)
```

## 4.1. Visualize specific rates (which are no model states) over time

This is where the classmethods come in handy, as mentioned before. The classmethod can be called without the need for an instance of the bioprocess model class, since such instance is handled by the caretaker object.

```python
# Extract necessary trajectory values from the simulation
t = Helpers.extract_time_series(simulation, name='S', replicate_id=None).timepoints
S = Helpers.extract_time_series(simulation, name='S', replicate_id=None).values
P = Helpers.extract_time_series(simulation, name='P', replicate_id=None).values

# Calculate the trajectory of specific growth rate and create a corresponding TimeSeries object
_mu = InducingSubstrate.growth_rate(
    S,
    model_parameters['mu_max'],
    model_parameters['kS'],
    model_parameters['kSI'],
    P,
    model_parameters['kPI'],
)
mu = TimeSeries(name='mu', timepoints=t, values=_mu)

# Calculate the trajectory of biomass specific product formation rate and create a corresponding TimeSeries object
_qP = InducingSubstrate.product_formation(
    S,
    model_parameters['qP_max'],
    model_parameters['kSP'],
)
qP = TimeSeries(name='qP', timepoints=t, values=_qP)

# Calculate the trajectory of biomass specific substrate formation rate and create a corresponding TimeSeries object
_qS = InducingSubstrate.substrate_consumption(
    _mu,
    model_parameters['YXS'],
    _qP,
    model_parameters['YPS'],
)
qS = TimeSeries(name='qS', timepoints=t, values=_qS)

# Put the additional TimeSeries objects into a dictonary
specific_rates = [mu, qS, qP]
```

```python
# Have a look at the additional TimeSeries objects
_ = Visualization.show_kinetic_data(specific_rates)
```

## 4.2. What can be done further, what purpose may serve the model?
__Get an idea on the functional relationship between initial substrate concentration $S_0$ and maximum Space-Time-Yield $STY_{max}$__

For each bioprocess that is run, the Space-Time-Yield over time has an optimium (maximum). This optimium varies, among others, with the initial substrate concentration $S_0$. The model can now be run for a variation of $S_0$ values, and for each variation the value of $STY_{max}$ can be calculated.

```python
variation_S0 = numpy.linspace(0, 100, 1000)
# Choose a final simulation time that allows for complete substrate consumption for all variants
t_final = 480
STYmax = []
for S0 in variation_S0:
    # Runs a simulation for each S0
    _simulation = caretaker.simulate(t=t_final, parameters={'S0' : S0})
    # Calculate the corresponding trajectory of STY
    _t = Helpers.extract_time_series(_simulation, name='P', replicate_id=None).timepoints[1:] # omit first timepoint 0 to avoid zero division error
    _P = Helpers.extract_time_series(_simulation, name='P', replicate_id=None).values[1:]
    _STY = _P / _t
    # Get the maximum STY for this simulation
    STYmax.append(numpy.nanmax(_STY))
STYmax = numpy.array(STYmax)
```

```python
# Visualize the results
fig, ax = pyplot.subplots(dpi=100)
ax.plot(variation_S0, STYmax, linestyle='-', marker='', color='black', zorder=1, label='$STY_{max} = f\\left(S_0\\right)$')
_max_idx = numpy.argmax(STYmax)
ax.scatter(variation_S0[_max_idx], STYmax[_max_idx], zorder=2, color='red', label='Optimum')
ax.set_xlabel('Initial substrate concentration $S_0$')
ax.set_ylabel('Maximum Space-Time-Yield $STY_{max}$')
ax.legend(frameon=False)
fig.tight_layout()
```

# 5. Example: Two-substrates growth,  one acting as inducer and inhibitor

```python
# Defines the model class
class TwoSubstrateGrowth(BioprocessModel):

    # The Monod equation is defined as classmethod
    @classmethod
    def Monod(cls, S, mu_max, kS):
        return mu_max * S / (kS + S)

    @classmethod
    def substrate_inhibition(cls, S, mu_max, kS, kI):
        return mu_max * S / (kS + S + S**2/kI)

    @classmethod
    def product_formation(cls, S, qP_max, kSP, P, kPI):
        return qP_max * S / (kSP + S) * kPI/(kPI + P)

    def rhs(self, t, y, sw):

        # Unpacks the state vector. The states are alphabetically ordered.
        P, S1, S2, X = y

        # Unpacks the model parameters.
        kS1 = self.model_parameters['kS1']
        mu_max1 = self.model_parameters['mu_max1']
        YXS1 = self.model_parameters['YXS1']

        kS2 = self.model_parameters['kS2']
        kS2I = self.model_parameters['kS2I']
        mu_max2 = self.model_parameters['mu_max2']
        YXS2 = self.model_parameters['YXS2']

        qP_max = self.model_parameters['qP_max']
        kSP = self.model_parameters['kSP']
        kPI = self.model_parameters['kPI']

        # Calculate specific growth rate
        mu1 = self.Monod(S1, mu_max1, kS1)
        mu2 = self.substrate_inhibition(S2, mu_max2, kS2, kS2I)
        mu = mu1 + mu2

        # Calculate substrate consumption rates.
        qS1 = 1/YXS1*mu1
        qS2 = 1/YXS2*mu2

        # Calculate product formation rate.
        qP = self.product_formation(S2, qP_max, kSP, P, kPI)

        # Calculate state derivatives.
        dXdt = mu*X
        dS1dt = -qS1*X
        dS2dt = -qS2*X
        dPdt = qP*X

        # Return list of state derivatives in the same order as the state vector was unpacked.
        return [dPdt, dS1dt, dS2dt, dXdt]

    def state_events(self, t, y, sw):
        t_add = self.model_parameters['t_add']
        event_t = t - t_add
        return [event_t]

    def change_states(self, t, y, sw):
        P, S1, S2, X = y
        S2_add = self.model_parameters['S2_add']
        if sw[0]:
            S2 = S2 + S2_add
        return [P, S1, S2, X]

model_parameters = {
    'kS1' : 0.02,
    'mu_max1' : 0.4,
    'YXS1' : 0.55,
    'kS2' : 0.2,
    'kS2I' : 0.05,
    'mu_max2' : 0.15,
    'YXS2' : 0.35,
    'qP_max' : 0.02,
    'kSP' : 0.2,
    'kPI' : 0.1,
    't_add' : 10.0,
    'S2_add' : 5.0,
}

initial_values = {
    'S10' : 50.0,
    'S20' : 0.0,
    'P0' : 0.0,
    'X0' : 0.1,
}

caretaker = Caretaker(
    bioprocess_model_class=TwoSubstrateGrowth,
    model_parameters=model_parameters,
    initial_values=initial_values,
)

simulation = caretaker.simulate(t=48)
_ = Visualization.show_kinetic_data(simulation, ncols=4)
```

```python

```


#  CS207 Systems Development for Computational Science
## Final Project
## Milestone 1

**Harvard University**

**Fall 2019**

**Instructor**: David Sondak

**Team #14**: Fantastic Four

**Students**: Daniel Cox, Anna Davydova, Stephen Moon, Valentina Toll Villagra

**Git Repository**:https://github.com/IACS-CS-207-FantasticFour/cs207-FinalProject

*“Nothing takes place in the world whose meaning is not that of some maximum or minimum.”*
― Leonhard Euler

**INTRODUCTION:**

This project aims to deliver an Automatic Differentiation Software DeltaPi that will calculate derivatives efficiently and accurately, making it useful across a wide spectrum of applications. In this package, we will execute the forward mode of Automatic Differentiation and provide an extension for pricing of financial options. We are especially excited about the applications of this software on Wall Street.

Derivatives are powerful and ubiquitous. Their use spans from gradients and Hessians in machine learning applications to hedge sensitivities in financial markets. However, of the four existing methods for derivative calculation only Automatic  Differentiation (AD) combines interpretability, efficiency and accuracy. Manual differentiation is inefficient and susceptible to error.(Baydin et. al., 2018). Numerical differentiation, while easy and fairly quick in its implementation, is also inaccurate, prone to rounding and truncation errors (Jerrell, 1997).  It does not scale well, which makes it a poor choice for machine learning models. Symbolic differentiation, while accurate, can become incredibly complex facing "expression swell" issues (Corliss, 1988). AD overcomes these issues, via its application of the chain rule and a step by step approach to differentiation, and accurately computes derivatives with asymptotic efficiency. (Baydin et.al, 2018).

Our software DeltaPi will implement AD methods, allowing the end user to benefit from its accuracy and efficiency.  Our goal is to produce a package that can handle a wide variety of uses beyond simple scalar functions with an additional extension for option trading applications.

**BACKGROUND:**

Generally speaking, in any computer program a function can be broken down into its elementary function components (unary or binary) such as addition, subtraction  log, sin, sqrt etc (Heath, 2018).  Since the value of partial derivatives of these elementary functions can be easily calculated, then the value of the entire function can be calculated via the application of the chain rule.

Let us walk through this process in more detail. Fundamentally, Chain Rule is the key pillar to the AD process. Recall that:

$$\dfrac{\partial h}{\partial t} = \dfrac{\partial h}{\partial u}\dfrac{\partial u}{\partial t}$$

We can extend the definition of the Chain Rule to a function that contains to functions as follows: h(u(t),v(t)) (Sondak, 2019).:

\begin{align}
  \displaystyle
  \frac{\partial h}{\partial t} = \frac{\partial h}{\partial u}\frac{\partial u}{\partial t} + \frac{\partial h}{\partial v}\frac{\partial v}{\partial t}.
\end{align}

This result, from a perspetive of a gradient, leads us to this general rule:

\begin{align}
  \nabla_{x}h = \sum_{i=1}^{n}{\frac{\partial h}{\partial y_{i}}\nabla y_{i}\left(x\right)}.
\end{align}

Thus, in a nutshell, the AD process consist of breaking down the function into its elementary components and carrying out the differentiation process in sequential order of these operations while multiplying through with the chain rule. The output of this process is a dual number containing the value of the function along with its derivative.  We can visualize this process with a computational graph and a table that contains the trace of the calculations.

For example, for a simple function (*adapted from class exercises*): $$f\left(x,y\right) = (sin(x)-cos(y))^2.$$

The forward mode AD graph looks as follows:

The corresponding computational table for values $x=\frac{\pi}{2}$ and $y=\frac{\pi}{3}$looks as follows:

| Trace    | Elementary Operation    | Current Function Value           | Derivative         | Partial Derivative w.r.t. x                            | Partial Derivative w.r.t. y                       |evaluate f(x) and its derivative: (f(x), f'(x))                   |
| :------: | :-------------------------------------:| :--------------------------------------------------------:|:-------------------------------------: | :-------------------------------------------------:|:----------------------------------------------:|:---------------------------------------:
| $x_{1}$  | $x_{1}$                                |          $\frac{\pi}{2}$                                              |$\dot {x_1}$                                     |      $1$                                           |      $0$                                       |$\left(\frac{\pi}{2}, 1\right)$                                                                  |
| $x_{2}$  | $x_{2}$                                |          $\frac{\pi}{3}$                                              |$\dot {x_2}$                                     |      $0$                                           |      $1$                                       |$\left(\frac{\pi}{3},1 \right)$                                                                  |
| $x_{3}$  | $sin(x_1)$                             |          $1$                                          |$cos(x_1)\dot{x_1}$                          |      $0$                                          |      $0$                                       |$\left(1,0)\right)$                                                        |
| $x_{4}$  | $cos(x_2)$                         |               $\frac{1}{2}$                                            |$-sin(x_2)\dot{x_2}$                     |      $0$                                          |    $-\frac{\sqrt{3}}{2}$                             |$\left(\frac{1}{2},-\frac{\sqrt{3}}{2}\right )$                                                         |
| $x_{5}$  | $x_3-x_4$                          |          $\frac{1}{2}$                                        |$\dot{x_3}-\dot{x_5}$           |     $0$                                        |   $\frac{\sqrt{3}}{2}$                            |$\left(\frac{1}{2},\frac{\sqrt{3}}{2}\right)$                                                         |
| $x_{6}$  | $x_5^2$                             |    $\frac{1}{4}$                                     |  $2x_5\dot{x_5}$                         |                        $0$                            |       $\frac{\sqrt{3}}{2}$ |           $\left(\frac{1}{4},\frac{\sqrt{3}}{2}\right)$      |

We note not only the simplicity of differentiation process on elementary functions carried along by the chain rule but relatively interpretable process that could be useful to our end-users.

**HOW TO USE OUR PACKAGE**

####  Pre-requisites:

import numpy as np

--------------------------------------------------------------------------------------------------------------

 Be x,y,z,... input variables
 Be f(x,y,z,...) the output of function f at (x,y,z,...)
 Be AutoDiff the class

1) Create an instance of the AutoDiff class
x = AutoDiff(val, derv)

where:
 val = the value of x at the desired point
 derv = the initial value of the derivative of x
     1 for calculating the partial derivative df/dx
     0 for calculating the partial derivative df/dy,df/dz,...

2) Repeat the process of creating an instance of the AutoDiff class for each input variable
y = AutoDiff(val, derv)
z = AutoDiff(val, derv)
...

where:
 val = the value of the instance variable (y,z,...) at the desired point
 derv = the initial value of the derivative of the instance variable (y,z,...)
     1 for calculating the partial derivative of the instance variable (df/dy for y, df/dz for z,...)
     0 for calculating the partial derivative of another variable that is not the instance one (df/dx,df/dz for y, df/dx,df/dy for z,...)

3) Enter the function of interest
f = function of x,y,z,...

where function can contain the following operations:
 +
 -
 *
 /
 **
 exp()
 sqrt()
 sin()
 cos()
 tan()

4) Print f.val to get the calculated value of f at the specified point (x,y,z,...)
print(f.val)

5) Print f.derv to get the calculated value of df/dx or df/dy or df/dz or ... at the specified point (x,y,z,...)
print(f.derv)

#### Example

 For values:
x = 3
y = 4
f = x + 2 * y

 Be the goal to calculate the partial derivatives df/dx at (x,y)=(3,4)

1) Create an instance of the AutoDiff class
x = AutoDiff(3, 1)

2) Repeat the process of creating an instance of the AutoDiff class for y
y = AutoDiff(4, 0)

3) Enter the function of interest
f = x + 2 * y

4) Print f.val to get the calculated value of f at (3,4)
print(f.val)

5) Print f.derv to get the calculated value of df/dx at (3,4)
print(f.derv)

 Be the goal to calculate the partial derivatives df/dy at the same (x,y)=(3,4)

1) Create an instance of the AutoDiff class
x = AutoDiff(3, 0)

2) Repeat the process of creating an instance of the AutoDiff class for y
y = AutoDiff(4, 1)

3) Enter the function of interest
f = x + 2 * y

4) Print f.val to get the calculated value of f at (3,4) - will be the same from df/dx
print(f.val)

5) Print f.derv to get the calculated value of df/dy at (3,4) print(f.derv)

--------------------------------------------------------------------------------------------------------------
#### To use all_derivatives() function

 Be x,y,z,... input variables
 Be func(x,y,z,...) the output of function f at (x,y,z,...)
 Be all_derivatives the function that uses AutoDiff class

1) Create a vector containing the input variables
in_vars = [x,y,z,...]

where:
 x,y,z,... are variable names

2) Create a vector containing the variable values
in_vals = [val_x,val_y,val_z,...]

where:
 val_x = the value of x at which to evaluate the function
 val_y = the value of y at which to evaluate the function
 val_z = the value of z at which to evaluate the function
...

3) Enter the function to be evaluated
func = function of x,y,z,...

where function can contain the following operations:
 +
 -
 *
 /
 **
 exp()
 sqrt()
 sin()
 cos()
 tan()

4) Run all_derivatives()
result = allderivatives(func, in_vars, in_vals)

5) Print result.val to get the calculated value of func at (in_vars)
print(result.val)

6) Print result.derv_vals to get vector of calculated partial derivative values, one for each variable in in_vars print(result.derv_vals)

#### Example

 For values:
x = 3
y = 4
z = 5
func = x + 2 * y - z

1) Create a vector containing the input variables
in_vars = [x,y,z]

2) Create a vector containing the variable values
in_vals = [3,4,5]

3) Enter the function to be evaluated
func =  x + 2 * y - z

4) Run all_derivatives()
result = all_derivatives(func, in_vars, in_vals)

5) Print result.val to get the calculated value of func at (in_vars)
print(result.val)

6) Print result.derv_vals to get vector of calculated partial derivative values, one for each variable in in_vars print(result.derv_vals)

--------------------------------------------------------------------------------------------------------------
#### To use multi_func_all_derivatives() function

 Be x,y,z,... input variables
 Be func(x,y,z,...) the output of function f at (x,y,z,...)
 Be multi_func_all_derivatives the function that uses AutoDiff class

1) Create a vector containing the input variables
in_vars = [x,y,z,...]

where:
 x,y,z,... are variable names

2) Create a vector containing the variable values
in_vals = [val_x,val_y,val_z,...]

where:
 val_x = the value of x at which to evaluate the function
 val_y = the value of y at which to evaluate the function
 val_z = the value of z at which to evaluate the function
...

3) Create a vector containing the functions to be evaluated
funcs = [func_1,func_2,func_3,...]

where:
func_n = function of x,y,z,...

where func_n can contain the following operations:
 +
 -
 *
 /
 **
 exp()
 sqrt()
 sin()
 cos()
 tan()

4) Run multi_func_all_derivatives()
result = multi_func_all_derivatives(funcs, in_vars, in_vals)

5) Print result.out_vals to get vector of calculated values of all funcs each at (in_vars)
print(result.out_vals)

6) Print result.out_derv_matrix to get the Hessian matrix of partial derivatives evaluated at (in_vars)
print(result.out_derv_matrix)

#### Example
 For values:
x = 3
y = 4
z = 5
func_1 = x + 2 * y - z
func_2 = x - exp(y) + z

1) Create a vector containing the input variables
in_vars = [x,y,z]

2) Create a vector containing the variable values
in_vals = [3,4,5]

3) Create a vector containing the functions to be evaluated
func_1 = x + 2 * y - z
func_2 = x - exp(y) + z
funcs = [func_1,func_2]

4) Run multi_func_all_derivatives()
result = multi_func_all_derivatives(funcs, in_vars, in_vals)

5) Print result.out_vals to get vector of calculated values of all funcs each at (in_vars)
print(result.out_vals)

6) Print result.out_derv_matrix to get the Hessian matrix of partial derivatives evaluated at (in_vars) print(result.out_derv_matrix)

**SOFTWARE ORGANIZATION**

While some elements of our package’s directory structure could be subject to change moving forward, as a general outline, our package will be comprised of two sub-packages - one for our implementation of the forward mode of auto-differentiation and one for our application of auto-differentiation to the pricing of stock options.

Within our forward mode sub-package, in addition to an __init__.py file, we plan to have one module within which we will include the entirety of our forward mode implementation. After importing this module, users will be able to evaluate both the value and derivatives of any multivariate equation. Specifics on usage are included in the following section. Since we have yet to work out the details of our project’s extension, it is difficult to provide a comprehensive map of our second sub-package’s structure. In an ideal world, we would be able to implement a complete application of auto-differentiation to the pricing of stock options, but this may prove to be too large in scope.

If we were able to accomplish this task, we would need at least two modules within our option-pricing sub-package: one for the reverse mode of auto-differentiation (since it is required for our application) and one that actually extends auto differentiation to option-pricing. The latter module - for option-pricing - would need to import from the modules for both the forward and reverse mode. Additionally, depending on some of the more specific design choices for our extension, a module containing helper functions could also be included in the sub-package implementing our extension.

Due to the group nature of our project, we think it makes the most sense to test our code using a combination of Travis CI and Codecov as our code will automatically be tested on every commit that we push to our group’s repo. We plan to write tests for our code as we develop, so Codecov will be another useful resource as it will provide us with immediate feedback on how much of our code is being tested.

We plan to distribute our package on PyPI as it is easy to use with pip install. We will package the software using the standard package files (Setup.py, README.md, and LICENSE).

Preliminary Directory Structure:

Final_Project\
    __init__.py
    Auto_Diff\
        __init__.py
        ForwardMode.py
    Option_Pricing\
        __init__.py
        ReverseMode.py
        OptionPricing.py

**IMPEMENTATION**

Our software will provide the user with a package he or she can use to evaluate functions of multiple variables at a specified point while at the same time evaluating the partial derivatives of the function at that same point. A class will be provided to perform simple calculations, and two functions will be provided to handle vector inputs and outputs.

Our strategy will be to create input-variable objects in python that once defined may be used in python calculations as if they were standard python variables of type float. However, when calculations are done with these Autdiff objects, a partial derivative with respect to one of the input variables will be calculated and carried along. The object representing the result, then, will contain an attribute that contains the value of the function at the specified point (val) and an attribute that contains a partial derivative with respect to one of the input variables (derv). Which partial derivative to calculate will be specified by the user by intiating the derv attribute of one input variable to 1 and the derv attributes of the other input variables to 0.  The calculation can then be repeated to get partial derivatives with respect to other input variables by changing which input variable’s derv attribute is intiated to 1.

*The AutDiff() class*

Most of the functionality of our package will be implemented in a single class called AutoDiff( ). Below is a skeleton of that class. As mentioned above, it will have two attributes self.val and self.derv that respectively will hold the value of the function being evaluated and the partial derivative of the function being evaluated up to some point in the calculation.

Autodiff() will have many methods. Each will be short and designed to replace a standard mathematical operation. Thus, for example there will be a method \_\_add\_\_ that will overload the standard addition operator (+) and thus specify how objects of type AutoDiff will be added to scalar values or to each other. And similar methods will be included to overload the other standard mathematical operators (- * / ^). Also, methods will be included to replace the
numpy functions (exp, ln, log, sin, cos, tan, sqrt) so that these functions can be used as well in any calculations involving AutoDiff objects.

In the skeleton below, two of AutoDiff’s methods, multiplication and exponentiation, have been fully coded to illustrate how the desired functionality will be implemented generally. Central to all methods is the principle that each method will return a new AutDiff object that will represent the next stage in the function’s evaluation. For example, if x is an AutDiff object, and it is multiplied by the scalar 2 (x\*2), then x’s \_\_mul\_\_ method will be called and it will return a new object that will represent x\*2. The AutoDiff.\_\_mul\_\_ method will perform this calculation as follows. It will assume that 2 is another object of type AutoDiff, and it will try to multiply the two AutoDiff objects together. It will find that 2 is not an AutoDiff object, fall to its except block, and then make two calculations: 2 * x.val and 2 * x.derv and store these values in the returned object’s val and derv attributes. If the 2 had been to the right of x (2*x), then the \_\_rmul\_\_ method would have been called, and it would have executed the same operation.

More subtle is the case where two AutoDiff objects, say x and y, are multiplied together x\*y. Here x’s \_\_mul\_\_ method will be called. It will find that the two objects are instances of AutoDiff, and it will 1) multiply x.val and y.val together and store the result as the returned object’s val attribute, and 2) calculate (x.derv\*y.val + y.derv\*x.val) and store this value as the returned object’s derv attribute. Notice in the second calculation the derivatives from both objects x and y are carried along as a sum. Recall, however, that the derv value of either x or y will have been initiated to 0, so one term in the sum (x.derv\*y.val + y.derv\*x.val) will always evaluate to 0. In this way, the algorithm can automatically choose the proper derivative to carry along for subsequent calculations (here either that of x or that of y) when two Autodiff objects collide. This same approach can be used more generally and will be used when overloading the division and power operators.

Also fully coded in the AutoDiff skeleton below is the exponential function. This works as follows. For a given AutoDiff object, say x, when exp(x) is executed, x’s exp() function will be called on itself. It will use numpy’s exp() function to evaluate np.exp(x.val) and store this value in the return object’s val attribute. It will evaluate the derivative of exp(x.value) —which in this case is just exp(x.value)— and store this value in the return object’s derv attribute. Our preliminary tests indicate that this will work for all the numpy fuctions: np.exp(), np.ln(), np.log(), np.sqrt(), np.sin(), np.cos(), np.tan(), so this is how they will be implemented.

Autdiff objects will be fully functional replacements for python float variables such that once the AutoDiff class is defined, these objects can be used to perform forward autodifferentiation on any multi-input, single output function at any specified point.

*Skeleton of AutoDiff()*
class AutoDiff():
    '''Class to do automatic differentiation'''

    def __init__(self, val, derv=1):
        self.val = float(val)
        self.derv = float(derv)

    def __mul__(self, other):
        try:
            return AutoDiff(self.val * other.val, self.derv * other.val +
                                                          other.derv* self.val)
        except:
            return AutoDiff(self.val * other, self.derv * other)

    def __rmul__(self, other):
        return AutoDiff(self.val * other, self.derv * other)

    def __div__(self, other):

    def __rdiv__(self, other):

    def __add__(self, other):

    def __radd__(self, other):

    def __sub__(self, other):

    def __rsub__(self, other):

    def __pow__(self, other):

    def __rpow__(self, other):

    def exp(self):
		return AutoDiff(np.exp(self.val), np.exp(self.val)*self.derv)

	def sin(self):

    def cos(self):

    def tan(self):

    def sqrt(self):

    def ln(self):

    def log(self):
*The all_derivatives() function*

This function will allow the user to calculate all partial derivatives at once without having to manually initiate the derv attributes of the input variables several times and each time repeat the calculation.

Psuedo code for this function is listed below. It takes advantage of the AutoDiff() class.

    def all_derivatives(func, in_vars, in_vals):
        ***
        Inputs:

            func():   —the function to be evaluated
            in_vars:  —a vector containing the input variables.
            in_vals:  —a vector of values, one for each variable in in_vars that specify
                         a point where the function and its partial derivatives are to be evaluated.
        Returns:

            func():    —the function evaluated
            val:       —the value of func( ) at the input point given by in_vals
            derv_vals: —a vector of partial derivative values, one for each input variable in in_vars.
        ***

        Steps:
        * make a vector diff_vars which  coorepsonds to in_vars but all
          values are AutoDiff objects.
        * Set the derv value of the first element of diff_vars to 1 and the derv
          values of the rest of the diff_vars elements to 0.

        * Create the derv_vals vector that will hold the calculated partial derivatives

        * For i in range( the number of elements in in_vars):
            * Set the derv value of each element of in_vars to 0
            * Set  the derv value of the ith element to 1.
            * Evaluate func()

            * If it is the first time thought the loop:
                  set val to the val of the AutoDiff output object
            * Else:
                * Store derivative of the output AutoDiff object as the ith element of derv_vals

        * Return  func, val, derv_vals

*The multi_func_all_derivatives() function*

This function will handle vector inputs and outputs. It will allow the user to specify: a vector of functions, a vector of input variables, and a vectors of values defining a specific point. It will return: the input vector containing the functions, a vector of output values (one for each function evaluated at the specified point), and a Hessian matrix of partial derivatives containing partial derivatives of all functions evaluated at the specified point. Consider the equation below as an example:

\begin{equation}
\begin{pmatrix} f(x,y) \\ g(x,y)\end{pmatrix} =
\begin{pmatrix} 2xy \\ x+y^2 \end{pmatrix} at \begin{pmatrix} x=a \\ y=b \end{pmatrix}
\end{equation}

Here the multi_func_all_derivatives() function would return:
1) funcs, a vector of the functions that were evaluated ($2xy$,  $x+y^2$)
2) out_vals, a vector of the values of each function after evaluation $(f(a,b), g(a,b))$
3) out_derv_matrix, a matrix containing the partial derivatives evaluated at the point (a, b)

Pseudo code for this function is below. It takes advantge of the AutoDiff() class and the all_derivatives() function:

    def multi_func_all_derivatives(funcs, in_vars, in_vals):
        ***
        Inputs:
            funcs:    —a vector of input functions
            in_vars: —a vector of input variables
            in_vals:  —a vector of input values specifying a point.

        Returns:
            funcs: 	   —the vector of input functions
            out_vals —a vector of output values, one for each function
            out_derv_matrix —the Hessian matrix of partial derivatives evaluated at the input point
        ***
        Steps:
        •	Create an empty vector out_vals
        •	Create an empty matrix out_derv_matrix
        •	For each function in funcs:
                - Run all_derivatives(func, in_vars, in_vals)
                - Store the returned val in the vector out_vals
                - Store the returned vector derv_vals in the matrix out_derv_matrix

        •	Return funcs, out_vals, out_derv_matrix

**EXTENSION**$\text{**}$

We aim to extend our DeltaPi package into a second package that would be useful to options traders as they estimate first and second order derivatives of the Black Scholes options pricing equation. Our goal is to provide a more efficient and accurate process for these calculations. As an example of this application let us consider an American option that gives its holder the right but not the obligation to buy (call) or sell (put) the underlying financial asset S at a predetermined price (strike) K during a predetermined time period  until date T.

The payoff function for this option with stock price $S_t$ at a time t and strike K is as follows:

\begin{equation}
  V(S_t,K) =
    \begin{cases}
      max(K-S_t,0) & \text{(put)}\\
      max(S_t-K,0) & \text{(call)}\\
    \end{cases}
\end{equation}

Black-Scholes Equation describes the option price overtime as follows:

$$\frac{\partial V}{\partial t}+\frac{1}{2}\sigma^2 S^2\frac{\partial^2 V}{\partial S^2}+r S\frac{\partial V}{\partial S} -r V = 0$$

The sensitivities (known as Greeks among the options traders) of an option with price V are calculated are first and second order derivatives as follows:

$$\varDelta=\frac{\partial V}{\partial S_o}$$

$$\varGamma = \frac{\partial^2 V}{\partial S^2_o}$$

$$\tau=\frac {\partial V}{\partial \sigma}$$

$$\varTheta=\frac{\partial C}{\partial t}$$

$$\text{Rho}=\frac{\partial C}{\partial r}$$

$$\text{Vanna}=\frac{\partial^2 V}{\partial S_o\sigma}$$

The value of a call option can be expressed as follows:

$$C(S_{t},t)=N(d_{1})S_{t}-N(d_{2})PV(K)$$

The value of a put option can be expressed as follows:

$$\begin{aligned}P(S_{t},t)=Ke^{-r(T-t)}-S_{t}+C(S_{t},t)\\=N(-d_{2})Ke^{-r(T-t)}-N(-d_{1})S_{t}\end{aligned}$$

where:

   $d_{1}={\frac {1}{\sigma {\sqrt {T-t}}}}\left[\ln \left({\frac {S_{t}}{K}}\right)+\left(r+{\frac {\sigma ^{2}}{2}}\right)(T-t)\right]\\d_{2}=d_{1}-\sigma {\sqrt {T-t}}\\PV(K)=Ke^{-r(T-t)}$

$N(x)=\frac{1}{\sqrt {2\pi}}\int^x_{-\inf}e^{\frac{-z^2}{2}}dz$

We aim to apply our automatic differentiation to compute the first and second order derivatives and will compare computational time (efficiency) and accuracy of our AD process vs. symbolic and numerical differentiation solutions for Black Scholes Greeks.  This will require forward and reverse mode AD (Homescu, 2011).

$\text{**}$*We understand that extension was not part of the formal assignment for this Milestone but we were hoping to get some early feedback if possible as we are very excited to start working on this idea. Thank you!*

**REFERENCES:**

Atilim Baydin. Automatic Differentiation in Machine Learning: a Survey. Journal of Machine Learning Research, 2018.

George F. Corliss. Application of differentiation arithmetic, volume 19 of Perspectives in Computing, pages 127–48. Academic Press, Boston, 1988.

Michael Heath. Scientific Computing: An Introductory Survey. Society for Industrial and Applied Mathematics. 8(6):367, 2018.

Christina Homescu. Adjoints and automatic (algorithmic) differentiation in computational finance. arXiv:1107.1831v1 10 Jul 2011.

Max E. Jerrell. Automatic differentiation and interval arithmetic for estimation of disequilibrium models. Computational Economics, 10(3):295–316, 1997.

David Sondak. Lectures 10 - 11. Harvard University. CS207 Fall 2019.


# [Shifted Kerr-Schild Solution](https://arxiv.org/pdf/1704.00599.pdf) Initial Data

## Authors: George Vopal & Zach Etienne
### Formatting improvements courtesy Brandon Clark

## This module sets up Shifted Kerr-Schild initial data [Etienne et al., 2017 GiRaFFE](https://arxiv.org/pdf/1704.00599.pdf).

**Notebook Status:**  Validated

**Validation Notes:** This module has been validated to exhibit convergence to zero of the Hamiltonian and momentum constraint violations at the expected order to the exact solution (see plots at bottom of [the exact initial data validation start-to-finish tutorial notebook](Tutorial-Start_to_Finish-BSSNCurvilinear-Setting_up_Exact_Initial_Data.ipynb)).

### NRPy+ Source Code for this module: [BSSN/ShiftedKerrSchild.py](../edit/BSSN/ShiftedKerrSchild.py)

## Introduction:
Shifted Kerr-Schild coordinates are similar to the trumpet spacetime, in that $r=0$ maps to some finite radius surface in Kerr-Schild coordinates. The radial shift $r_0$ both reduces the black hole's coordinate size and causes the very strongly-curved spacetime fields at $r<r_{0}$ to vanish deep inside the horizon, which aids in numerical stability, e.g., when evolving hydrodynamic, MHD, and FFE fields inside the horizon.

# Table of Contents:
$$\label{toc}$$

1. [Step 1](#initialize_nrpy): Set up the needed NRPy+ infrastructure and declare core gridfunctions
1. [Step 2](#kerr_schild_lapse): The Kerr-Schild Lapse, Shift, and 3-Metric
    1. [Step 2.a](#define_rho): Define $\rho^{2}$, $\alpha$, $\beta^{r}$, $\beta^{\theta}$, $\beta^{\phi}$, $\gamma_{r\theta}$, $\gamma_{\theta\phi}$
    1. [Step 2.b](#nonzero_gamma): Define and construct nonzero components of $\gamma_{ij}$
1. [Step 3](#extrinsic_curvature): The extrinsic curvature $K_{ij}$
    1. [Step 3.a](#abc): Define useful quantities $A$, $B$, $C$
    1. [Step 3.b](#nonzero_k): Define and construct nonzero components of $K_{ij}$
1. [Step 4](#code_validation): Code Validation against `BSSN.ShiftedKerrSchild` NRPy+ module
1. [Step 5](#latex_pdf_output): Output this notebook to $\LaTeX$-formatted PDF file

# Step 1: Set up the needed NRPy+ infrastructure and declare core gridfunctions \[Back to [top](#toc)\]
$$\label{initialize_nrpy}$$

First, we will import the core modules from Python/NRPy+ and specify the main gridfunctions we will need.

**Input for initial data**:

* The black hole mass $M$.
* The black hole spin parameter $a$
* The radial offset $r_0$

```python
# Step P0: Load needed modules
import sympy as sp
import NRPy_param_funcs as par
import indexedexp as ixp
import grid as gri
from outputC import *
import reference_metric as rfm

# All gridfunctions will be written in terms of spherical coordinates (r, th, ph):
r,th,ph = sp.symbols('r th ph', real=True)

thismodule = "ShiftedKerrSchild"

DIM = 3
par.set_parval_from_str("grid::DIM",DIM)

# Input parameters:
M, a, r0 = par.Cparameters("REAL", thismodule,
                           ["M", "a", "r0"],
                           [1.0, 0.9,  1.0])

# Auxiliary variables:
rho2 = sp.symbols('rho2', real=True)
```

# Step 2: The Kerr-Schild Lapse, Shift, and 3-Metric \[Back to [top](#toc)\]
$$\label{kerr_schild_lapse}$$

## Step 2.a: Define $\rho^{2}$, $\alpha$, $\beta^{r_{\rm KS}}$, $\beta^{\theta}$, $\beta^{\phi}$, $\gamma_{r_{\rm KS}\theta}$, $\gamma_{\theta\phi}$ \[Back to [top](#toc)\]
$$\label{define_rho}$$

The relationship between the Kerr-Schild radius $r_{\rm KS}$ and the radial coordinate used on our numerical grid $r$, is given by

$$
r_{\rm KS} = r + r_0,
$$
where $r_0\ge 0$ is the radial shift.

Notice that the radial shift has no impact on Jacobians since $\frac{\partial{r_{\rm KS}}}{\partial{r}}=1$. $r_0$ must be set to a value less than the horizon radius $R$, but not so close to $R$ that finite-difference stencils from outside the horizon cross $r=0$. Thus $r_0$ must be set with consideration of the numerical grid structure in mind, as nonzero values of $r_0$ will shrink the coordinate size of the black hole by exactly $r_0$.

All of these equations are as defined in the appendix of the original GiRaFFE paper ([Etienne et al., 2017 GiRaFFE](https://arxiv.org/pdf/1704.00599.pdf)).

First, we define $\rho^{2}$ as

$$ \rho^2 = r_{\rm KS} + a^{2}\cos^{2}(\theta) $$

And we then define the Kerr-Schild lapse $\alpha$ from equation (A.1)

$$ \alpha = \frac{1}{\sqrt{1 + \frac{2Mr_{\rm KS}}{\rho^2}}} $$

And the shift $\beta$ from equations (A.2) & (A.3)

$$ \beta^{r_{\rm KS}} = \alpha^2\frac{2Mr_{\rm KS}}{\rho^2} $$

$$ \beta^{\theta} = \beta^{\phi} = \gamma_{r_{\rm KS}\theta} = \gamma_{\theta\phi} = 0 $$

```python
# Step 1: Define rho^2, alpha, beta^(r_{KS}), beta^(theta), beta^(phi), gamma_{r_{KS}theta}, gamma_{theta\phi}

# r_{KS} = r + r0
rKS = r+r0

# rho^2 = rKS^2 + a^2*cos^2(theta)
rho2 = rKS*rKS + a*a*sp.cos(th)**2

# alpha = 1/sqrt{1 + M*rKS/rho^2}
alphaSph = 1/(sp.sqrt(1 + 2*M*rKS/rho2))

# Initialize the shift vector, \beta^i, to zero.
betaSphU = ixp.zerorank1()
# beta^r = alpha^2*2Mr/rho^2
betaSphU[0] = alphaSph*alphaSph*2*M*rKS/rho2

# Time derivative of shift vector beta^i, B^i, is zero.
BSphU = ixp.zerorank1()
```

## Step 2.b: Define and construct nonzero components $\gamma_{r_{\rm KS}r_{\rm KS}}$, $\gamma_{r_{\rm KS}\phi}$, $\gamma_{\theta\theta}$, $\gamma_{\phi\phi}$ \[Back to [top](#toc)\]
$$\label{nonzero_gamma}$$

From equations (A.4)-(A.7) of [Etienne et al., 2017](https://arxiv.org/pdf/1704.00599.pdf) we define the nonzero components of the 3-metric:

$$ \gamma_{r_{\rm KS}r_{\rm KS}} = 1 + \frac{2Mr_{\rm KS}}{\rho^2} $$

$$ \gamma_{r_{\rm KS}\phi} = -a\gamma_{r_{\rm KS}r_{\rm KS}}\sin^2(\theta) $$

$$ \gamma_{\theta\theta} = \rho^2 $$

$$ \gamma_{\phi\phi} = \left(r_{\rm KS}^2 + a^2 + \frac{2Mr_{\rm KS}}{\rho^2}a^{2}\sin^{2}(\theta)\right)\sin^{2}(\theta) $$

```python
# Step 2: Define and construct nonzero components gamma_{r_{KS}r_{KS}}$, gamma_{r_{KS}phi},
#         gamma_{thetatheta}, gamma_{phiphi}

# Initialize \gamma_{ij} to zero.
gammaSphDD = ixp.zerorank2()

# gammaDD{rKS rKS} = 1 +2M*rKS/rho^2
gammaSphDD[0][0] = 1 + 2*M*rKS/rho2

# gammaDD{rKS phi} = -a*gammaDD{r r}*sin^2(theta)
gammaSphDD[0][2] = gammaSphDD[2][0] = -a*gammaSphDD[0][0]*sp.sin(th)**2

# gammaDD{theta theta} = rho^2
gammaSphDD[1][1] = rho2

# gammaDD{phi phi} = (rKS^2 + a^2 + 2Mr/rho^2*a^2*sin^2(theta))*sin^2(theta)
gammaSphDD[2][2] = (rKS*rKS + a*a + 2*M*rKS*a*a*sp.sin(th)**2/rho2)*sp.sin(th)**2
```

# Step 3: The extrinsic curvature $K_{ij}$ \[Back to [top](#toc)\]
$$\label{extrinsic_curvature}$$

## Step 3.a: Define useful quantities $A$, $B$, $C$ \[Back to [top](#toc)\]
$$\label{abc}$$

From equations (A.8)-(A.10) of [Etienne et al., 2017](https://arxiv.org/pdf/1704.00599.pdf) we define the following expressions which will help simplify the nonzero extrinsic curvature components:

$$ A = \left(a^{2}\cos(2\theta) + a^{2} + 2r_{\rm KS}^{2}\right) $$

$$ B = A + 4Mr_{\rm KS} $$

$$ D = \sqrt{\frac{2Mr_{\rm KS}}{a^{2}\cos^{2}(\theta) + r_{\rm KS}^2} + 1} $$

```python
# Step 3: Define useful quantities A, B, C

# A = (a^2*cos^2(2theta) + a^2 + 2r^2)
A = (a*a*sp.cos(2*th) + a*a + 2*rKS*rKS)

# B = A + 4M*rKS
B = A + 4*M*rKS

# D = \sqrt(2M*rKS/(a^2cos^2(theta) + rKS^2) + 1)
D = sp.sqrt(2*M*rKS/(a*a*sp.cos(th)**2 + rKS*rKS) + 1)
```

## Step 3.b: Define and construct nonzero components of $K_{ij}$ \[Back to [top](#toc)\]
$$\label{nonzero_k}$$

We will now express the extrinsic curvature $K_{ij}$ in spherical polar coordinates.

From equations (A.11) - (A.13) of [Etienne et al., 2017](https://arxiv.org/pdf/1704.00599.pdf) we define the following:

$$ K_{r_{\rm KS}r_{\rm KS}} = \frac{D(A + 2Mr_{\rm KS})}{A^{2}B}\left[4M\left(a^{2}\cos(2\theta) + a^{2} - 2r_{\rm KS}^{2}\right)\right] $$

$$ K_{r_{\rm KS}\theta} = \frac{D}{AB}\left[8a^{2}Mr_{\rm KS}\sin(\theta)\cos(\theta)\right] $$

$$ K_{r_{\rm KS}\phi} = \frac{D}{A^2}\left[-2aM\sin^{2}(\theta)\left(a^{2}\cos(2\theta) + a^{2} - 2r_{\rm KS}^{2}\right)\right] $$

```python
# Step 4: Define the extrinsic curvature in spherical polar coordinates
# Establish the 3x3 zero-matrix
KSphDD = ixp.zerorank2()

# *** Fill in the nonzero components ***
# *** This will create an upper-triangular matrix ***
# K_{r r} = D(A+2Mr)/(A^2*B)[4M(a^2*cos(2theta) + a^2 - 2r^2)]
KSphDD[0][0] = D*(A+2*M*rKS)/(A*A*B)*(4*M*(a*a*sp.cos(2*th)+a*a-2*rKS*rKS))

# K_{r theta} = D/(AB)[8a^2*Mr*sin(theta)cos(theta)]
KSphDD[0][1] = KSphDD[1][0] = D/(A*B)*(8*a*a*M*rKS*sp.sin(th)*sp.cos(th))

# K_{r phi} = D/A^2[-2aMsin^2(theta)(a^2cos(2theta)+a^2-2r^2)]
KSphDD[0][2] = KSphDD[2][0] =  D/(A*A)*(-2*a*M*sp.sin(th)**2*(a*a*sp.cos(2*th)+a*a-2*rKS*rKS))
```

And from equations (A.14) - (A.17) of [Etienne et al., 2017](https://arxiv.org/pdf/1704.00599.pdf) we define the following expressions to complete the upper-triangular matrix $K_{ij}$:

$$ K_{\theta\theta} = \frac{D}{B}\left[4Mr_{\rm KS}^{2}\right] $$

$$ K_{\theta\phi} = \frac{D}{AB}\left[-8a^{3}Mr_{\rm KS}\sin^{3}(\theta)\cos(\theta)\right] $$

$$ K_{\phi\phi} = \frac{D}{A^{2}B}\left[2Mr_{\rm KS}\sin^{2}(\theta)\left(a^{4}(r_{\rm KS}-M)\cos(4\theta) + a^{4}(M + 3r_{\rm KS}) + 4a^{2}r_{\rm KS}^{2}(2r_{\rm KS} - M) + 4a^{2}r_{\rm KS}\cos(2\theta)\left(a^{2} + r_{\rm KS}(M + 2r_{\rm KS})\right) + 8r_{\rm KS}^{5}\right)\right] $$

```python
# K_{theta theta} = D/B[4Mr^2]
KSphDD[1][1] = D/B*(4*M*rKS*rKS)

# K_{theta phi} = D/(AB)*(-8*a^3*Mr*sin^3(theta)cos(theta))
KSphDD[1][2] = KSphDD[2][1] = D/(A*B)*(-8*a**3*M*rKS*sp.sin(th)**3*sp.cos(th))

# K_{phi phi} = D/(A^2*B)[2Mr*sin^2(theta)(a^4(M+3r)
#   +4a^2r^2(2r-M)+4a^2r*cos(2theta)(a^2+r(M+2r))+8r^5)]
KSphDD[2][2] = D/(A*A*B)*(2*M*rKS*sp.sin(th)**2*(a**4*(rKS-M)*sp.cos(4*th)\
                        + a**4*(M+3*rKS)+4*a*a*rKS*rKS*(2*rKS-M)\
                        + 4*a*a*rKS*sp.cos(2*th)*(a*a + rKS*(M + 2*rKS)) + 8*rKS**5))
```

# Step 4: Code Validation against `BSSN.ShiftedKerrSchild` NRPy+ module \[Back to [top](#toc)\]
$$\label{code_validation}$$

Here, as a code validation check, we verify agreement in the SymPy expressions for Shifted Kerr-Schild initial data between

1. this tutorial and
2. the NRPy+ [BSSN.ShiftedKerrSchild](../edit/BSSN/ShiftedKerrSchild.py) module.

```python
# Reset the list of gridfunctions, as registering a gridfunction
#   twice will spawn an error.
gri.glb_gridfcs_list = []

# First we import needed core NRPy+ modules
import reference_metric as rfm
import grid as gri
# Unused, but needed because BrillLindquist() also does tensor rescalings.
rfm.reference_metric()

import BSSN.ShiftedKerrSchild as sks
sks.ShiftedKerrSchild()

# It is SAFE to ignore the warning(s) from re-initializing parameters.
print("Consistency check between Brill-Lindquist tutorial and NRPy+ BSSN.BrillLindquist module. ALL SHOULD BE ZERO.")
print("alphaSph - sks.alphaSph = "+str(sp.simplify(alphaSph - sks.alphaSph)))
for i in range(DIM):
    print("betaSphU["+str(i)+"] - sks.betaSphU["+str(i)+"] = "+\
          str(sp.simplify(betaSphU[i] - sks.betaSphU[i])))
    print("BSphU["+str(i)+"] - sks.BaSphU["+str(i)+"] = "+str(sp.simplify(BSphU[i] - sks.BSphU[i])))
    for j in range(DIM):
        print("gammaSphDD["+str(i)+"]["+str(j)+"] - sks.gammaSphDD["+str(i)+"]["+str(j)+"] = "+\
             str(sp.simplify(gammaSphDD[i][j] - sks.gammaSphDD[i][j])))
        print("KSphDD["+str(i)+"]["+str(j)+"] - sks.KSphDD["+str(i)+"]["+str(j)+"] = "+\
              str(sp.simplify(KSphDD[i][j] - sks.KSphDD[i][j])))
```

    Consistency check between Brill-Lindquist tutorial and NRPy+ BSSN.BrillLindquist module. ALL SHOULD BE ZERO.
    alphaSph - sks.alphaSph = 0
    betaSphU[0] - sks.betaSphU[0] = 0
    BSphU[0] - sks.BaSphU[0] = 0
    gammaSphDD[0][0] - sks.gammaSphDD[0][0] = 0
    KSphDD[0][0] - sks.KSphDD[0][0] = 0
    gammaSphDD[0][1] - sks.gammaSphDD[0][1] = 0
    KSphDD[0][1] - sks.KSphDD[0][1] = 0
    gammaSphDD[0][2] - sks.gammaSphDD[0][2] = 0
    KSphDD[0][2] - sks.KSphDD[0][2] = 0
    betaSphU[1] - sks.betaSphU[1] = 0
    BSphU[1] - sks.BaSphU[1] = 0
    gammaSphDD[1][0] - sks.gammaSphDD[1][0] = 0
    KSphDD[1][0] - sks.KSphDD[1][0] = 0
    gammaSphDD[1][1] - sks.gammaSphDD[1][1] = 0
    KSphDD[1][1] - sks.KSphDD[1][1] = 0
    gammaSphDD[1][2] - sks.gammaSphDD[1][2] = 0
    KSphDD[1][2] - sks.KSphDD[1][2] = 0
    betaSphU[2] - sks.betaSphU[2] = 0
    BSphU[2] - sks.BaSphU[2] = 0
    gammaSphDD[2][0] - sks.gammaSphDD[2][0] = 0
    KSphDD[2][0] - sks.KSphDD[2][0] = 0
    gammaSphDD[2][1] - sks.gammaSphDD[2][1] = 0
    KSphDD[2][1] - sks.KSphDD[2][1] = 0
    gammaSphDD[2][2] - sks.gammaSphDD[2][2] = 0
    KSphDD[2][2] - sks.KSphDD[2][2] = 0

# Step 5: Output this notebook to $\LaTeX$-formatted PDF \[Back to [top](#toc)\]
$$\label{latex_pdf_output}$$

The following code cell converts this Jupyter notebook into a proper, clickable $\LaTeX$-formatted PDF file. After the cell is successfully run, the generated PDF may be found in the root NRPy+ tutorial directory, with filename
[Tutorial-ADM_Initial_Data-ShiftedKerrSchild.pdf](Tutorial-ADM_Initial_Data-ShiftedKerrSchild.pdf) (Note that clicking on this link may not work; you may need to open the PDF file through another means.)

```python
!jupyter nbconvert --to latex --template latex_nrpy_style.tplx --log-level='WARN' Tutorial-ADM_Initial_Data-ShiftedKerrSchild.ipynb
!pdflatex -interaction=batchmode Tutorial-ADM_Initial_Data-ShiftedKerrSchild.tex
!pdflatex -interaction=batchmode Tutorial-ADM_Initial_Data-ShiftedKerrSchild.tex
!pdflatex -interaction=batchmode Tutorial-ADM_Initial_Data-ShiftedKerrSchild.tex
!rm -f Tut*.out Tut*.aux Tut*.log
```

    This is pdfTeX, Version 3.14159265-2.6-1.40.18 (TeX Live 2017/Debian) (preloaded format=pdflatex)
     restricted \write18 enabled.
    entering extended mode
    This is pdfTeX, Version 3.14159265-2.6-1.40.18 (TeX Live 2017/Debian) (preloaded format=pdflatex)
     restricted \write18 enabled.
    entering extended mode
    This is pdfTeX, Version 3.14159265-2.6-1.40.18 (TeX Live 2017/Debian) (preloaded format=pdflatex)
     restricted \write18 enabled.
    entering extended mode