# Neuromatch Academy: Week 1, Day 5, Tutorial 2
# Dimensionality Reduction: Principal component analysis

__Content creators:__ Alex Cayco Gajic, John Murray

__Content reviewers:__ Roozbeh Farhoudi, Matt Krause, Spiros Chavlis, Richard Gao, Michael Waskom

---
# Tutorial Objectives

In this notebook we'll learn how to perform PCA by projecting the data onto the eigenvectors of its covariance matrix.

Overview:
- Calculate the eigenvectors of the sample covariance matrix.
- Perform PCA by projecting data onto the eigenvectors of the covariance matrix.
- Plot and explore the eigenvalues.

To quickly refresh your knowledge of eigenvalues and eigenvectors, you can watch this [short video](https://www.youtube.com/watch?v=kwA3qM0rm7c) (4 minutes) for a geometrical explanation. For a deeper understanding, this [in-depth video](https://www.youtube.com/watch?v=PFDu9oVAE-g&list=PLZHQObOWTQDPD3MizzM2xVFitgF8hE_ab&index=14) (17 minutes) provides an excellent basis and is beautifully illustrated.

```python
# @title Video 1: PCA
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="-f6T9--oM0E", 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=-f6T9--oM0E

---
# Setup
Run these cells to get the tutorial started.

```python
# Imports
import numpy as np
import matplotlib.pyplot as plt
```

```python
# @title Figure Settings
import ipywidgets as widgets  # interactive display
%config InlineBackend.figure_format = 'retina'
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle")
```

```python
# @title Helper functions

def plot_eigenvalues(evals):
  """
  Plots eigenvalues.

  Args:
      (numpy array of floats) : Vector of eigenvalues

  Returns:
    Nothing.

  """
  plt.figure(figsize=(4, 4))
  plt.plot(np.arange(1, len(evals) + 1), evals, 'o-k')
  plt.xlabel('Component')
  plt.ylabel('Eigenvalue')
  plt.title('Scree plot')
  plt.xticks(np.arange(1, len(evals) + 1))
  plt.ylim([0, 2.5])

def sort_evals_descending(evals, evectors):
  """
  Sorts eigenvalues and eigenvectors in decreasing order. Also aligns first two
  eigenvectors to be in first two quadrants (if 2D).

  Args:
    evals (numpy array of floats)    : Vector of eigenvalues
    evectors (numpy array of floats) : Corresponding matrix of eigenvectors
                                        each column corresponds to a different
                                        eigenvalue

  Returns:
    (numpy array of floats)          : Vector of eigenvalues after sorting
    (numpy array of floats)          : Matrix of eigenvectors after sorting
  """

  index = np.flip(np.argsort(evals))
  evals = evals[index]
  evectors = evectors[:, index]
  if evals.shape[0] == 2:
    if np.arccos(np.matmul(evectors[:, 0],
                           1 / np.sqrt(2) * np.array([1, 1]))) > np.pi / 2:
      evectors[:, 0] = -evectors[:, 0]
    if np.arccos(np.matmul(evectors[:, 1],
                           1 / np.sqrt(2) * np.array([-1, 1]))) > np.pi / 2:
      evectors[:, 1] = -evectors[:, 1]
  return evals, evectors

def plot_data(X):
  """
  Plots bivariate data. Includes a plot of each random variable, and a scatter
  scatter plot of their joint activity. The title indicates the sample
  correlation calculated from the data.

  Args:
    X (numpy array of floats) : Data matrix each column corresponds to a
                                different random variable

  Returns:
    Nothing.
  """

  fig = plt.figure(figsize=[8, 4])
  gs = fig.add_gridspec(2, 2)
  ax1 = fig.add_subplot(gs[0, 0])
  ax1.plot(X[:, 0], color='k')
  plt.ylabel('Neuron 1')
  ax2 = fig.add_subplot(gs[1, 0])
  ax2.plot(X[:, 1], color='k')
  plt.xlabel('Sample Number (sorted)')
  plt.ylabel('Neuron 2')
  ax3 = fig.add_subplot(gs[:, 1])
  ax3.plot(X[:, 0], X[:, 1], '.', markerfacecolor=[.5, .5, .5],
           markeredgewidth=0)
  ax3.axis('equal')
  plt.xlabel('Neuron 1 activity')
  plt.ylabel('Neuron 2 activity')
  plt.title('Sample corr: {:.1f}'.format(np.corrcoef(X[:, 0], X[:, 1])[0, 1]))
  plt.show()

def get_data(cov_matrix):
  """
  Returns a matrix of 1000 samples from a bivariate, zero-mean Gaussian

  Note that samples are sorted in ascending order for the first random
  variable.

  Args:
    var_1 (scalar)                     : variance of the first random variable
    var_2 (scalar)                     : variance of the second random variable
    cov_matrix (numpy array of floats) : desired covariance matrix

  Returns:
    (numpy array of floats)            : samples from the bivariate Gaussian,
                                          with each column corresponding to a
                                          different random variable
  """

  mean = np.array([0, 0])
  X = np.random.multivariate_normal(mean, cov_matrix, size=1000)
  indices_for_sorting = np.argsort(X[:, 0])
  X = X[indices_for_sorting, :]
  return X

def calculate_cov_matrix(var_1, var_2, corr_coef):
  """
  Calculates the covariance matrix based on the variances and
  correlation coefficient.

  Args:
    var_1 (scalar)         :  variance of the first random variable
    var_2 (scalar)         :  variance of the second random variable
    corr_coef (scalar)     :  correlation coefficient

  Returns:
    (numpy array of floats) : covariance matrix
  """
  cov = corr_coef * np.sqrt(var_1 * var_2)
  cov_matrix = np.array([[var_1, cov], [cov, var_2]])
  return cov_matrix

def define_orthonormal_basis(u):
  """
  Calculates an orthonormal basis given an arbitrary vector u.

  Args:
    u (numpy array of floats) : arbitrary 2D vector used for new basis

  Returns:
    (numpy array of floats)   : new orthonormal basis columns correspond to
                                basis vectors
  """

  u = u / np.sqrt(u[0] ** 2 + u[1] ** 2)
  w = np.array([-u[1], u[0]])
  W = np.column_stack((u, w))
  return W

def plot_data_new_basis(Y):
  """
  Plots bivariate data after transformation to new bases. Similar to plot_data
  but with colors corresponding to projections onto basis 1 (red) and
  basis 2 (blue).
  The title indicates the sample correlation calculated from the data.

  Note that samples are re-sorted in ascending order for the first random
  variable.

  Args:
    Y (numpy array of floats) : Data matrix in new basis each column
                                corresponds to a different random variable

  Returns:
    Nothing.
  """

  fig = plt.figure(figsize=[8, 4])
  gs = fig.add_gridspec(2, 2)
  ax1 = fig.add_subplot(gs[0, 0])
  ax1.plot(Y[:, 0], 'r')
  plt.ylabel('Projection \n basis vector 1')
  ax2 = fig.add_subplot(gs[1, 0])
  ax2.plot(Y[:, 1], 'b')
  plt.xlabel('Projection \n basis vector 1')
  plt.ylabel('Projection \n basis vector 2')
  ax3 = fig.add_subplot(gs[:, 1])
  ax3.plot(Y[:, 0], Y[:, 1], '.', color=[.5, .5, .5])
  ax3.axis('equal')
  plt.xlabel('Projection basis vector 1')
  plt.ylabel('Projection basis vector 2')
  plt.title('Sample corr: {:.1f}'.format(np.corrcoef(Y[:, 0], Y[:, 1])[0, 1]))
  plt.show()

def change_of_basis(X, W):
  """
  Projects data onto a new basis.

  Args:
    X (numpy array of floats) : Data matrix each column corresponding to a
                                different random variable
    W (numpy array of floats) : new orthonormal basis columns correspond to
                                basis vectors

  Returns:
    (numpy array of floats)   : Data matrix expressed in new basis
  """

  Y = np.matmul(X, W)
  return Y

def plot_basis_vectors(X, W):
  """
  Plots bivariate data as well as new basis vectors.

  Args:
    X (numpy array of floats) : Data matrix each column corresponds to a
                                different random variable
    W (numpy array of floats) : Square matrix representing new orthonormal
                                basis each column represents a basis vector

  Returns:
    Nothing.
  """

  plt.figure(figsize=[4, 4])
  plt.plot(X[:, 0], X[:, 1], '.', color=[.5, .5, .5], label='Data')
  plt.axis('equal')
  plt.xlabel('Neuron 1 activity')
  plt.ylabel('Neuron 2 activity')
  plt.plot([0, W[0, 0]], [0, W[1, 0]], color='r', linewidth=3,
           label='Basis vector 1')
  plt.plot([0, W[0, 1]], [0, W[1, 1]], color='b', linewidth=3,
           label='Basis vector 2')
  plt.legend()
  plt.show()

```

---
# Section 1: Calculate the eigenvectors of the the sample covariance matrix

As we saw in the lecture, PCA represents data in a new orthonormal basis defined by the eigenvectors of the covariance matrix. Remember that in the previous tutorial, we generated bivariate normal data with a specified covariance matrix $\bf \Sigma$, whose $(i,j)$th element is:
\begin{equation}
\Sigma_{ij} = E[ x_i x_j ] - E[ x_i] E[ x_j ] .
\end{equation}
However, in real life we don't have access to this ground-truth covariance matrix. To get around this, we can use the sample covariance matrix, $\bf\hat\Sigma$, which is calculated directly from the data. The $(i,j)$th element of the sample covariance matrix is:
\begin{equation}
 \hat \Sigma_{ij} =  \frac{1}{N_\text{samples}}{\bf x}_i^T {\bf x}_j - \bar {\bf x}_i \bar{\bf x}_j ,
\end{equation}
where ${\bf x}_i = [ x_i(1), x_i(2), \dots,x_i(N_\text{samples})]^T$ is a column vector representing all measurements of neuron $i$, and  $\bar {\bf x}_i$ is the mean of neuron $i$ across samples:
\begin{equation}
\bar {\bf x}_i = \frac{1}{N_\text{samples}} \sum_{k=1}^{N_\text{samples}} x_i(k).
\end{equation}
If we assume that the data has already been mean-subtracted, then we can write the sample covariance matrix in a much simpler matrix form:
\begin{align}
{\bf \hat \Sigma}
&= \frac{1}{N_\text{samples}} {\bf X}^T {\bf X}.
\end{align}

## Exercise 1: Calculation of the covariance matrix

Before calculating the eigenvectors, you must first calculate the sample covariance matrix.

**Steps**
* Complete the function `get_sample_cov_matrix` by first subtracting the sample mean of the data, then calculate $\bf \hat \Sigma$ using the equation above.
* Use `get_data` to generate bivariate normal data, and calculate the sample covariance matrix with your finished `get_sample_cov_matrix`. Compare this estimate to the true covariate matrix using `calculate_cov_matrix`.

```python
help(get_data)
help(calculate_cov_matrix)
```

    Help on function get_data in module __main__:

    get_data(cov_matrix)
        Returns a matrix of 1000 samples from a bivariate, zero-mean Gaussian

        Note that samples are sorted in ascending order for the first random
        variable.

        Args:
          var_1 (scalar)                     : variance of the first random variable
          var_2 (scalar)                     : variance of the second random variable
          cov_matrix (numpy array of floats) : desired covariance matrix

        Returns:
          (numpy array of floats)            : samples from the bivariate Gaussian,
                                                with each column corresponding to a
                                                different random variable

    Help on function calculate_cov_matrix in module __main__:

    calculate_cov_matrix(var_1, var_2, corr_coef)
        Calculates the covariance matrix based on the variances and
        correlation coefficient.

        Args:
          var_1 (scalar)         :  variance of the first random variable
          var_2 (scalar)         :  variance of the second random variable
          corr_coef (scalar)     :  correlation coefficient

        Returns:
          (numpy array of floats) : covariance matrix

```python
def get_sample_cov_matrix(X):
  """
  Returns the sample covariance matrix of data X

  Args:
    X (numpy array of floats) : Data matrix each column corresponds to a
                                different random variable

  Returns:
    (numpy array of floats)   : Covariance matrix
  """

  #################################################
  ## TODO for students: calculate the covariance matrix
  # Fill out function and remove
#   raise NotImplementedError("Student excercise: calculate the covariance matrix!")
  #################################################

  # Subtract the mean of X
  X = X-X.mean(axis=0)
  # Calculate the covariance matrix (hint: use np.matmul)
  cov_matrix = 1/X.shape[0]*(X.T@X)

  return cov_matrix

##########################################################
## TODO for students: generate bivariate Gaussian data
# with variances of 1 and a correlation coefficient of 0.8
# compare the true and sample covariance matrices
##########################################################
variance_1 = 1
variance_2 = 1
corr_coef = 0.8

np.random.seed(2020)  # set random seed
# Uncomment below code to test your function
cov_matrix = calculate_cov_matrix(variance_1, variance_2, corr_coef)
print(cov_matrix)

X = get_data(cov_matrix)
sample_cov_matrix = get_sample_cov_matrix(X)
print(sample_cov_matrix)
```

    [[1.  0.8]
     [0.8 1. ]]
    [[0.99315313 0.82347589]
     [0.82347589 1.01281397]]

SAMPLE OUTPUT
```
[[1.  0.8]
 [0.8 1. ]]
[[0.99315313 0.82347589]
 [0.82347589 1.01281397]]
```

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

## Exercise 2: Eigenvectors of the Covariance matrix

Next you will calculate the eigenvectors of the covariance matrix. Plot them along with the data to check that they align with the geometry of the data.

**Steps:**
* Calculate the eigenvalues and eigenvectors of the sample covariance matrix. (**Hint:** use `np.linalg.eigh`, which finds the eigenvalues of a symmetric matrix).
* Use the provided code to sort the eigenvalues in descending order.
* Plot the eigenvectors on a scatter plot of the data, using the function `plot_basis_vectors`.

```python
help(sort_evals_descending)
help(plot_basis_vectors)
```

    Help on function sort_evals_descending in module __main__:

    sort_evals_descending(evals, evectors)
        Sorts eigenvalues and eigenvectors in decreasing order. Also aligns first two
        eigenvectors to be in first two quadrants (if 2D).

        Args:
          evals (numpy array of floats)    : Vector of eigenvalues
          evectors (numpy array of floats) : Corresponding matrix of eigenvectors
                                              each column corresponds to a different
                                              eigenvalue

        Returns:
          (numpy array of floats)          : Vector of eigenvalues after sorting
          (numpy array of floats)          : Matrix of eigenvectors after sorting

    Help on function plot_basis_vectors in module __main__:

    plot_basis_vectors(X, W)
        Plots bivariate data as well as new basis vectors.

        Args:
          X (numpy array of floats) : Data matrix each column corresponds to a
                                      different random variable
          W (numpy array of floats) : Square matrix representing new orthonormal
                                      basis each column represents a basis vector

        Returns:
          Nothing.

```python
#################################################
## TO DO for students: Calculate and sort the eigenvalues in descending order
#################################################
# Calculate the eigenvalues and eigenvectors
evals, evectors = np.linalg.eigh(sample_cov_matrix)
# Sort the eigenvalues in descending order
evals, evectors = sort_evals_descending(evals,evectors)

plot_basis_vectors(X, evectors)
```

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

*Example output:*

---
# Section 2: Perform PCA by projecting data onto the eigenvectors

To perform PCA, we will project the data onto the eigenvectors of the covariance matrix, i.e.:
\begin{equation}
\bf S = X W
\end{equation}
where $\bf S$ is an $N_\text{samples} \times N$ matrix representing the projected data (also called *scores*), and $W$ is an $N\times N$ orthogonal matrix, each of whose columns represents the eigenvectors of the covariance matrix (also called *weights* or *loadings*).

## Exercise 3: PCA implementation

You will now perform PCA on the data using the intuition you have developed so far. Fill in the function below to carry out the steps to perform PCA by projecting the data onto the eigenvectors of its covariance matrix.

**Steps:**
* First subtract the mean.
* Then calculate the sample covariance matrix.
* Then find the eigenvalues and eigenvectors and sort them in descending order.
* Finally project the mean-centered data onto the eigenvectors.

```python
help(change_of_basis)
help(plot_data_new_basis)
```

    Help on function change_of_basis in module __main__:

    change_of_basis(X, W)
        Projects data onto a new basis.

        Args:
          X (numpy array of floats) : Data matrix each column corresponding to a
                                      different random variable
          W (numpy array of floats) : new orthonormal basis columns correspond to
                                      basis vectors

        Returns:
          (numpy array of floats)   : Data matrix expressed in new basis

    Help on function plot_data_new_basis in module __main__:

    plot_data_new_basis(Y)
        Plots bivariate data after transformation to new bases. Similar to plot_data
        but with colors corresponding to projections onto basis 1 (red) and
        basis 2 (blue).
        The title indicates the sample correlation calculated from the data.

        Note that samples are re-sorted in ascending order for the first random
        variable.

        Args:
          Y (numpy array of floats) : Data matrix in new basis each column
                                      corresponds to a different random variable

        Returns:
          Nothing.

```python
def pca(X):
  """
  Sorts eigenvalues and eigenvectors in decreasing order.

  Args:
    X (numpy array of floats): Data matrix each column corresponds to a
                               different random variable

  Returns:
    (numpy array of floats)  : Data projected onto the new basis
    (numpy array of floats)  : Vector of eigenvalues
    (numpy array of floats)  : Corresponding matrix of eigenvectors

  """

  #################################################
  ## TODO for students: calculate the covariance matrix
  # Fill out function and remove
#   raise NotImplementedError("Student excercise: sort eigenvalues/eigenvectors!")
  #################################################

  # Subtract the mean of X
  X = X - X.mean(axis=0)
  # Calculate the sample covariance matrix
  cov_matrix = 1/X.shape[0]*(X.T@X)
  # Calculate the eigenvalues and eigenvectors
  evals, evectors = np.linalg.eigh(cov_matrix)
  # Sort the eigenvalues in descending order
  evals, evectors = sort_evals_descending(evals,evectors)
  # Project the data onto the new eigenvector basis
  score = X@evectors

  return score, evectors, evals

#################################################
## TODO for students: Call the function to calculate the eigenvectors/eigenvalues
#################################################

# Perform PCA on the data matrix X
score, evectors, evals = pca(X)
# Plot the data projected into the new basis
plot_data_new_basis(score)
```

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

*Example output:*

## Plot and explore the eigenvalues

Finally, we will examine the eigenvalues of the covariance matrix. Remember that each eigenvalue describes the variance of the data projected onto its corresponding eigenvector. This is an important concept because it allows us to rank the PCA basis vectors based on how much variance each one can capture. First run the code below to plot the eigenvalues (sometimes called the "scree plot"). Which eigenvalue is larger?

```python
plot_eigenvalues(evals)
evals
```

## Interactive Demo: Exploration of the correlation coefficient

Run the following cell and use the slider to change the correlation coefficient in the data. You should see the scree plot and the plot of basis vectors update.

**Questions:**
* What happens to the eigenvalues as you change the correlation coefficient?
* Can you find a value for which both eigenvalues are equal?
* Can you find a value for which only one eigenvalue is nonzero?

```python
# @title

# @markdown Make sure you execute this cell to enable the widget!

def refresh(corr_coef=.8):
  cov_matrix = calculate_cov_matrix(variance_1, variance_2, corr_coef)
  X = get_data(cov_matrix)
  score, evectors, evals = pca(X)
  plot_eigenvalues(evals)
  plot_basis_vectors(X, evectors)

_ = widgets.interact(refresh, corr_coef=(-1, 1, .1))
```

    interactive(children=(FloatSlider(value=0.8, description='corr_coef', max=1.0, min=-1.0), Output()), _dom_clas…

---
# Summary
- In this tutorial, we learned that goal of PCA is to find an orthonormal basis capturing the directions of maximum variance of the data. More precisely, the $i$th basis vector is the direction that maximizes the projected variance, while being orthogonal to all previous basis vectors. Mathematically, these basis vectors are the eigenvectors of the covariance matrix (also called *loadings*).
- PCA also has the useful property that the projected data (*scores*) are uncorrelated.
- The projected variance along each basis vector is given by its corresponding eigenvalue. This is important because it allows us rank the "importance" of each basis vector in terms of how much of the data variability it explains. An eigenvalue of zero means there is no variation along that direction so it can be dropped without losing any information about the original data.
- In the next tutorial, we will use this property to reduce the dimensionality of high-dimensional data.

---
# Bonus: Mathematical basis of PCA properties

```python
# @title Video 2: Properties of PCA
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="p56UrMRt6-U", 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=p56UrMRt6-U

```python

```


# Regression simulation

```python
import warnings
warnings.filterwarnings("ignore")
```

```python
%matplotlib inline
%load_ext autoreload
%autoreload 2
```

```python
from jupyterthemes import jtplot
jtplot.style(theme='onedork', context='notebook', ticks=True, grid=False)
```

```python
import pandas as pd
pd.options.display.max_rows = 999
pd.options.display.max_columns = 999
pd.set_option("display.max_columns", None)
import numpy as np
import os
import matplotlib.pyplot as plt
from pylab import rcParams
rcParams['figure.figsize'] = 15, 5

import data
import copy
from rolldecay.bis_system import BisSystem

from rolldecay import database
from mdldb.tables import Run, RolldecayQuadraticBRegression
from rolldecayestimators.direct_estimator_cubic import EstimatorQuadraticB
from rolldecayestimators.ikeda_estimator import IkedaQuadraticEstimator
from rolldecayestimators import symbols, equations
from rolldecayestimators.substitute_dynamic_symbols import lambdify
from rolldecayestimators.transformers import CutTransformer, LowpassFilterDerivatorTransformer, ScaleFactorTransformer, OffsetTransformer

import sympy as sp

from sklearn.pipeline import Pipeline
from sklearn.model_selection import cross_val_score
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.feature_selection import VarianceThreshold
from sklearn.ensemble import ExtraTreesClassifier
from sklearn.feature_selection import SelectKBest
from sklearn.feature_selection import chi2
from sklearn.feature_selection import f_regression

from sklearn.linear_model import LinearRegression
from sklearn.model_selection import cross_validate
from sklearn.preprocessing import PolynomialFeatures
from sklearn.pipeline import Pipeline
from sklearn.pipeline import make_pipeline

import signal_lab
from sqlalchemy.inspection import inspect
import seaborn as sns
import docs
import pickle
import gc
import dill
from rolldecayestimators.substitute_dynamic_symbols import lambdify,run
from rolldecayestimators.polynom_estimator import Polynom
```

## omega0

```python
omega_hat_polynom = Polynom.load('omega0_hat_polynom.sym')

data = pd.read_pickle('data.sav')
data['omega_hat'] = omega_hat_polynom.predict(data)
y_s = pd.read_pickle('y.sav')
```

## $\hat{B_1}$

```python
filename = 'B_1_hat_model.sav'
B_1_hat_model = pickle.load(open(filename, 'rb'))
```

```python
y = y_s['B_1_hat']
#X=data[zeta_model.keys].copy()  # Handle categorical data later
#zeta_model.score(X=X,y=y)
```

```python
B_1_hat_polynom = Polynom.load('B_1_hat_polynom.sym')
B_1_hat_polynom.score(X=data,y=y)
```

## $\hat{B_2}$

```python
filename = 'B_2_hat_model.sav'
B_2_hat_model = pickle.load(open(filename, 'rb'))
```

```python
B_2_hat_polynom = Polynom.load('B_2_hat_polynom.sym')
B_2_hat_polynom.score(X=data,y=y)
```

## Predict parameters

```python
predict = data.copy()
#predict['omega_hat'] = omega0_hat_model.predict(data[omega0_hat_model.keys])
predict['omega_hat'] = omega_hat_polynom.predict(data)

#predict['zeta'] = zeta_model.predict(data[zeta_model.keys])
predict['B_1_hat'] = B_1_hat_polynom.predict(data)

#predict['d'] = d_model.predict(data[d_model.keys])
predict['B_2_hat'] = B_2_hat_polynom.predict(data)

```

## Simulate

```python
df_rolldecay = database.load(rolldecay_table_name='rolldecay_quadratic_b', limit_score=0.90)
db = database.get_db()
```

```python
df_rolldecay_ikeda = database.load(rolldecay_table_name='rolldecay_simplified_ikeda', limit_score=0.0)
```

```python
predict.head()
```

```python
meta_data = predict.iloc[0]
```

```python
meta_data_full_scale = df_rolldecay.loc[meta_data.name]
```

```python
db_run = db.session.query(Run).get(int(meta_data.name))
df = database.load_run(db_run)
```

```python
df.plot(y='phi')
```

```python
lowpass_filter = LowpassFilterDerivatorTransformer(cutoff=2, minimum_score=0.99)
#scaler = ScaleFactorTransformer(scale_factor=db_run.model.scale_factor)  # dummy value None for now
cutter = CutTransformer(phi_max=np.deg2rad(9), phi_min=np.deg2rad(0.25), phi1d_start_tolerance=0.015)
offset_transformer = OffsetTransformer()

steps = [
    ('filter',lowpass_filter),
#    ('scaler',scaler),  # Is froude scaling a good idea??
    ('cutter', cutter),
    ('offset_transformer',offset_transformer)
]

preprocessor = Pipeline(steps) # define the pipeline object.
preprocessor.fit(df)
X = preprocessor.transform(df)
```

```python
X.plot(y='phi')
```

```python
omega_hat_equation = equations.omega_hat_equation.subs(symbols.omega,symbols.omega0)
omega_hat_equation
```

```python
omega0_equation = sp.Eq(symbols.omega0,sp.solve(omega_hat_equation,symbols.omega0)[0])
omega0_equation
```

```python
meta_data['omega_hat']
```

```python
omega0_lambda = lambdify(sp.solve(omega0_equation,symbols.omega0)[0])
```

```python
C_1A = equations.new_symbols[symbols.C_1]
C_1A_equation = sp.Eq(C_1A,sp.solve(EstimatorQuadraticB.omgea0_equation,C_1A)[0])
C_1A_equation
```

```python
C_1A_lambda = lambdify(sp.solve(C_1A_equation,C_1A)[0])
```

```python
B_1_equation= sp.Eq(symbols.B_1,sp.solve(equations.B_1_hat_equation,symbols.B_1)[0])
B_1_equation
```

```python
B_2_equation= sp.Eq(symbols.B_2,sp.solve(equations.B_2_hat_equation,symbols.B_2)[0])
B_2_equation
```

```python
B_1_lambda=lambdify(sp.solve(B_1_equation,symbols.B_1)[0])
B_2_lambda=lambdify(sp.solve(B_2_equation,symbols.B_2)[0])
```

```python
equations.normalize_equations[symbols.B_1]
```

```python
B_1A = equations.new_symbols[symbols.B_1]
B_1A_lambda = lambdify(sp.solve(equations.normalize_equations[symbols.B_1],B_1A)[0])

B_2A = equations.new_symbols[symbols.B_2]
B_2A_lambda = lambdify(sp.solve(equations.normalize_equations[symbols.B_2],B_2A)[0])
```

```python
run_data = meta_data.copy()
run_data['g']=9.81
run_data['rho']=1000
run_data['m']=run_data['Disp']*run_data['rho']
run_data['omega0']=run(function=omega0_lambda, inputs=run_data)
run_data['C_1A']=run(function=C_1A_lambda, inputs=run_data)
run_data['A_44']=run(function=EstimatorQuadraticB.functions['A44'], inputs=run_data)
run_data['B_1']=run(function=B_1_lambda, inputs=run_data)
run_data['B_2']=run(function=B_2_lambda, inputs=run_data)
run_data['B_1A']=run(function=B_1A_lambda, inputs=run_data)
run_data['B_2A']=run(function=B_2A_lambda, inputs=run_data)

```

```python
run_data['Disp']
```

```python
direct_estimator = EstimatorQuadraticB.load(data=run_data, X=X)
```

```python
direct_estimator.plot_fit()
```

```python
run_data['omega0']
```

```python
direct_estimator.omega0
```

```python
direct_estimator.plot_error()
```

```python
direct_estimator.score()
```

### Cutting

```python
direct_estimator = EstimatorQuadraticB.load(data=run_data, X=X.loc[0:400])
```

```python
direct_estimator.plot_fit()
```

```python
direct_estimator.plot_error()
```

```python
direct_estimator.score()
```

### Using "real" omega0

```python
run_data['omega0']=meta_data_full_scale.omega0
run_data['C_1A']=run(function=C_1A_lambda, inputs=run_data)
run_data['A_44']=run(function=EstimatorQuadraticB.functions['A44'], inputs=run_data)
run_data['B_1']=run(function=B_1_lambda, inputs=run_data)
run_data['B_2']=run(function=B_2_lambda, inputs=run_data)
run_data['B_1A']=run(function=B_1A_lambda, inputs=run_data)
run_data['B_2A']=run(function=B_2A_lambda, inputs=run_data)
```

```python
direct_estimator = EstimatorQuadraticB.load(data=run_data, X=X)
```

```python
fig,ax=plt.subplots()
direct_estimator.plot_fit(ax=ax)

meta_data_ikeda = df_rolldecay_ikeda.loc[meta_data.name]
meta_data_ikeda['omega0']=meta_data_full_scale.omega0
ikeda_estimator = IkedaQuadraticEstimator.load(data=meta_data_ikeda, X=X)
ikeda_estimator.plot_fit(ax=ax, model_test=False)
```

```python
direct_estimator.score()
```

```python
ikeda_estimator.score()
```

## Do it for many...

```python
def evaluate(meta_data):

    if not meta_data.name in df_rolldecay.index:
        return

    meta_data_full_scale = df_rolldecay.loc[meta_data.name]
    db_run = db.session.query(Run).get(int(meta_data.name))
    df = database.load_run(db_run)
    preprocessor.fit(df)
    X = preprocessor.transform(df)

    direct_estimator = EstimatorQuadraticB.load(data=meta_data, X=X)
    meta_data['Volume']=meta_data['Disp']
    s=direct_estimator.result_for_database(meta_data=meta_data)
    rolldecay_db=RolldecayQuadraticBRegression(run_id=db_run.id,**s)
    db.session.merge(rolldecay_db)
    db.session.commit()

    return direct_estimator.score(X=X)

```

```python
predict['omega0']=df_rolldecay.omega0
predict['g']=9.81
predict['rho']=1000
predict['m']=predict['Disp']*predict['rho']
predict['C_1A']=run(function=C_1A_lambda, inputs=predict)
predict['A_44']=run(function=EstimatorQuadraticB.functions['A44'], inputs=predict)
predict['B_1']=run(function=B_1_lambda, inputs=predict)
predict['B_2']=run(function=B_2_lambda, inputs=predict)
predict['B_1A']=run(function=B_1A_lambda, inputs=predict)
predict['B_2A']=run(function=B_2A_lambda, inputs=predict)
```

```python
predict['score'] = predict.apply(func=evaluate, axis=1)
```

```python
predict['score'].hist(bins=40)
```

```python
df_rolldecay['score'].hist(bins=40)
```

## Faster way to calculate score

```python
predict.head()
```

```python
B_e_equation =sp.Eq(symbols.B_e,sp.solve(equations.B_e_equation, symbols.B_e)[0])
B_e_equation
```

```python
B_ep, B_1p, B_2p = sp.symbols('B_ep B_1p B_2p')

B_ep_equation = B_e_equation.subs([
    (symbols.B_e,B_ep),
    (symbols.B_1, B_1p),
    (symbols.B_2, B_2p),
])
B_ep_equation
```

```python
Be_lambda = lambdify(sp.solve(B_e_equation,symbols.B_e)[0])
```

```python
Be_lambda
```

```python
predict['phi_a'] = np.rad2deg(df_rolldecay['phi_start'].abs())
df_rolldecay['phi_a'] = np.rad2deg(df_rolldecay['phi_start'].abs())
```

```python
predict['B_e']=run(function=Be_lambda, inputs=predict)
```

```python
df_rolldecay['B_e']=run(function=Be_lambda, inputs=df_rolldecay)
```

```python
epsilon,SSR,SST,R = sp.symbols('epsilon SSR SST R')
phi_min, phi_max, = sp.symbols('phi_min phi_max')
residual_equation = sp.Eq(epsilon, symbols.B_e-B_ep)
residual_equation
```

```python
SSR_equation = sp.Eq(SSR,sp.Integral(epsilon**2,(symbols.phi_a,phi_min,phi_max)))
SSR_equation
```

```python
SSR_equation = SSR_equation.subs(epsilon, sp.solve(residual_equation, epsilon)[0]).subs(
    symbols.B_e,sp.solve(B_e_equation, symbols.B_e)[0]).subs(
    B_ep,sp.solve(B_ep_equation, B_ep)[0])
SSR_equation
```

```python
SSR_equation2 = SSR_equation.doit()
SSR_equation2
```

```python
B_mean = sp.symbols('B_mean')

B_mean_equation=sp.Eq(B_mean,
      sp.Integral(sp.solve(B_e_equation, symbols.B_e)[0]**2,
                  (symbols.phi_a,phi_min,phi_max)
                 )/(phi_max-phi_min)
     )
B_mean_equation
```

```python
B_mean_equation2 = B_mean_equation.doit()
B_mean_equation2
```

```python
SST_equation = sp.Eq(SST,sp.Integral((symbols.B_e-B_mean)**2,(symbols.phi_a,phi_min,phi_max)))
SST_equation
```

```python
SST_equation2 = SST_equation.subs([
                                (B_mean, sp.solve(B_mean_equation2,B_mean)[0]),
                                (symbols.B_e, sp.solve(B_e_equation,symbols.B_e)[0])
                                  ]
                                  ).doit()
SST_equation2
```

```python
r2_equation = sp.Eq(R**2, 1-SSR/SST)
r2_equation
```

```python
r2_equation2 = r2_equation.subs(
    SSR,sp.solve(SSR_equation2,SSR)[0]).subs(
    SST, sp.solve(SST_equation2,SST)[0])

r2_equation2
```

```python
r2_lambda = lambdify(r2_equation2.rhs)
```

```python
r2_lambda
```

```python
phi_min, phi_max, Be_tot = sp.symbols('phi_min phi_max B_e_tot')
Be_integral = sp.Eq(Be_tot,sp.Integral(B_e_equation.rhs,(symbols.phi_a,phi_min,phi_max)))
Be_integral
```

```python
Be_tot_equation = sp.Eq(Be_tot,Be_integral.rhs.simplify())
Be_tot_equation
```

```python
Be_tot_lambda = lambdify(sp.solve(Be_tot_equation, Be_tot)[0])
```

```python
predict['phi_min'] = 0
predict['phi_max'] = predict['phi_a']
predict['Be_tot'] = run(function=Be_tot_lambda, inputs=predict)

df_rolldecay['phi_min'] = 0
df_rolldecay['phi_max'] = df_rolldecay['phi_a']
df_rolldecay['Be_tot'] = run(function=Be_tot_lambda, inputs=df_rolldecay)
```

```python
mask = predict['score'] > 0.5
predict_good = predict.loc[mask].copy()
df_compare = pd.merge(left=predict_good, right=df_rolldecay, how='left', left_index=True, right_index=True,
                      suffixes=('_regression','_model_test'))

Be_tot_true = df_compare['Be_tot_model_test']
Be_tot_predicted = df_compare['Be_tot_regression']

df_compare['B_1'] = df_compare['B_1_model_test']
df_compare['B_2'] = df_compare['B_2_model_test']
df_compare['B_1p'] = df_compare['B_1_regression']
df_compare['B_2p'] = df_compare['B_2_regression']
df_compare['omega0'] = df_compare['omega0_model_test']
df_compare['phi_max'] = df_compare['phi_max_model_test']
df_compare['phi_min'] = df_compare['phi_min_model_test']

df_compare['score_simple'] = run(function=r2_lambda, inputs=df_compare)
mask=df_compare['score_simple']>0.7
df_compare=df_compare.loc[mask].copy()

fig,ax=plt.subplots()
df_compare.plot(x='B_e_model_test', y='B_e_regression', ax=ax, style='.')

fig,ax=plt.subplots()
df_compare.plot(x='score_regression', y='score_simple', ax=ax, style='.')
```

```python
fig,ax=plt.subplots()
df_compare.plot(x='Be_tot_model_test', y='Be_tot_regression', ax=ax, style='.')

fig,ax=plt.subplots()
df_compare.plot(x='Be_tot_regression', y='score_regression', ax=ax, style='.')
```

```python
def evaluate_plot(meta_data):

    if not meta_data.name in df_rolldecay.index:
        return

    meta_data_full_scale = df_rolldecay.loc[meta_data.name]
    db_run = db.session.query(Run).get(int(meta_data.name))
    df = database.load_run(db_run)

    preprocessor.fit(df)
    X = preprocessor.transform(df)

    direct_estimator = EstimatorQuadraticB.load(data=meta_data, X=X)
    fig,ax=plt.subplots()
    direct_estimator.plot_fit(ax=ax)
    ax.set_title('Score:%0.2f' % direct_estimator.score(X=X))

    if meta_data.name in df_rolldecay_ikeda.index:
        meta_data_ikeda = df_rolldecay_ikeda.loc[meta_data.name]
        meta_data_ikeda['omega0']=meta_data_full_scale.omega0
        ikeda_estimator = IkedaQuadraticEstimator.load(data=meta_data_ikeda, X=X)
        ikeda_estimator.plot_fit(ax=ax, model_test=False)

```

```python
for run_id,meta_data in predict.iloc[0:50].iterrows():
    evaluate_plot(meta_data)
    gc.collect()

```

```python

```

```python

```

```python

```


# Autograd (4)：PyTorch 高阶自动导数

> 创建时间：2019-12-22

PyTorch 的 Autograd 不仅可以作为一阶导数的计算工具，还可以对任意阶的高阶导数作计算。这篇文档我们简单地讨论这个问题。

这篇文档若没有 Stackoverflow 回答 [^stackoverflow-qa]，恐怕是不会写出来了。

```python
%matplotlib notebook

import torch
from torch.autograd import grad
import numpy as np
import scipy
import scipy.linalg
from matplotlib import pyplot as plt

torch.set_printoptions(precision=5, sci_mode=False, linewidth=120)
```

## 一元函数的高阶导数

我们仍然拿以下的函数

$$
y (b) = x^3 + 10 \exp \left( - \frac{x^2}{10} \right)
$$

其中，当 $x = 3$ 时，$y \simeq 31.07$。

```python
x = torch.tensor(3., requires_grad=True)
y = x**3 + 10 * torch.exp(-x**2 / 10)
float(y)
```

    31.065696716308594

根据 Stackoverflow 回答 [^stackoverflow-qa]，定义如下高阶求导函数 `nth_derivative`：

```python
def nth_derivative(f, wrt, n):
    for i in range(n):
        if not f.requires_grad:
            return torch.zeros_like(wrt)
        grads = grad(f, wrt, create_graph=True)[0]
        f = grads.sum()
    return grads
```

其中输入的三个参数分别是 `y` 因变量，`wrt` 自变量，`n` 不小于零的导数阶数。

我们可以很容易地验证其一阶导数结果确实是如文档 [Autograd (1)](./Autograd_TensorContract.ipynb) 所述的那样。

$$
\frac{\partial y}{\partial x} = 3 x^2 - 2 x \exp \left( - \frac{x^2}{10} \right)
$$

```python
float(nth_derivative(y, x, 1))
```

    24.56058120727539

```python
float(3 * x**2 - 2 * x * torch.exp(-x**2 / 10))
```

    24.56058120727539

但不仅如此，我们还可以求出其二阶导数：

$$
\begin{align}
u &= \exp \left( - \frac{x^2}{10} \right) \\
\frac{\partial^2 y}{\partial x^2} &= \frac{2}{5} u x^2 + 6 x - 2 u
\end{align}
$$

```python
float(nth_derivative(y, x, 2))
```

    18.6505126953125

```python
u = torch.exp(-x**2 / 10)
float(2 / 5 * u * x**2 + 6 * x - 2 * u)
```

    18.650510787963867

甚至是三阶导数：

$$
\frac{\partial^3 y}{\partial x^3} = - \frac{2}{25} u x^3 + \frac{6}{5} u x + 6
$$

```python
float(nth_derivative(y, x, 3))
```

    6.585460186004639

```python
float(-2 / 25 * u * x**3 + 6 / 5 * u * x + 6)
```

    6.585460186004639

## 向量二阶自动求导：Newton 法解极值点

### 问题的定义

下面我们以 Newton 法解极小值点为例，来讨论与向量的二阶自动求导有关的问题。

首先，我们的目标是求取下述标量 $y$ 关于向量 $x_i$ 函数的极小值：

$$
y = \sum_i A_i x_i + \sum_{ij} B_{ij} x_i x_j + \sum_{ijk} C_{ijk} x_i x_j x_k
$$

其中，`d` $\dim(i) = \dim(j) = \dim(k) = 2$。

```python
d = 2
```

为了让极小值基本上确实存在，我们不能随意构造矩阵 $B_{ij}$，并且让关于 $i, j$ 角标的矩阵 $B_{ij} + B_{ji}$ 为正定矩阵。一种构建非对称的 `B` $B_{ij}$ 或写作 $\mathbf{B}$ 的方式是

$$
\mathbf{B} = \exp(-\mathbf{K}) \mathbf{\Lambda} \exp(\mathbf{K}) + \mathbf{K}
$$

其中，

- `K` $\mathbf{K}$ 为反对称矩阵；

- `eK` $\exp(-\mathbf{K})$ 为正交矩阵，具有性质 $\exp(-\mathbf{K})^{-1} = \exp(-\mathbf{K})^\mathrm{T}$；

- `eig` $\mathbf{\Lambda}$ 为对角矩阵，其每个对角元的值是正值。

上述的向量 `x` $x_i$，向量 `A` $A_i$，反对称矩阵 `K` $K_{ij}$，对角矩阵 `eig` $\Lambda_{ij} \delta_{ij}$，张量 `C` $C_{ijk}$ 的取法都是有条件地任意的。

矩阵 `B` $B_{ij}$ 的构造过程如下：

```python
np.random.seed(0)
K = np.random.randn(d, d)
K -= K.T
eK = scipy.linalg.expm(K)
eig = np.abs(np.random.randn(d))
B = eK.T @ np.diag(eig) @ eK + K
```

其余量的构造过程如下 (其中 `C` 取了较小的值)：

```python
torch.random.manual_seed(0)
A = torch.randn(d)
B = torch.tensor(B, dtype=torch.float)
C = torch.randn(d, d, d) / 25
```

```python
A
```

    tensor([ 1.54100, -0.29343])

```python
B
```

    tensor([[ 1.60134, -0.98618],
            [ 0.17098,  1.24350]])

```python
C
```

    tensor([[[-0.08715,  0.02274],
             [-0.04338, -0.05594]],

            [[ 0.01613,  0.03352],
             [-0.02877, -0.01613]]])

这里我们指出，极值点问题的解分别为 `xmin_0`；同时存在鞍点 `xmin_1`。该问题可以表示如下：

```python
xmin_0 = np.array([-0.4719677, -0.0355982])
xmin_1 = np.array([ 9.3234644,  6.9193769])
```

```python
y_func = lambda x: A @ x + B @ x @ x + C @ x @ x @ x
x_list = x_list = torch.arange(-5, 15, 0.05)
y_list = [[y_func(torch.tensor([x0, x1]))
           for x0 in x_list]
           for x1 in x_list]
x_list = x_list.numpy()
y_list = torch.tensor(y_list).detach().numpy()
```

```python
fig, ax = plt.subplots()
cont = ax.contourf(x_list, x_list, y_list, 25)
plt.scatter(*xmin_0, s=100, marker="+")
plt.scatter(*xmin_1, s=100, marker="+")
ax.set_xlabel("dimension 0")
ax.set_ylabel("dimension 1")
ax = fig.colorbar(cont)
```

    <IPython.core.display.Javascript object>

### 二阶自动求导：程序

事实上，上述函数 $y$ 并没有下界，这也能从上图中右上角的趋势中看出；因此简单的梯度下降法在初猜偏离极值点附近时，很可能会在学习率 `lr` 设置得较大时出现这种不收敛的情况：

```python
x = torch.tensor([-10., 10.], requires_grad=True)
print("Initial x:", x)

lr = 0.5
for epoch in range(7):
    y = A @ x + B @ x @ x + C @ x @ x @ x
    print(float(y))
    g = grad(y, x)[0]
    x = x - g * lr
    x = x.detach().clone().requires_grad_()
```

    Initial x: tensor([-10.,  10.], requires_grad=True)
    465.36041259765625
    -69.391845703125
    -13590.681640625
    -7374935.0
    -1465419038720.0
    -5.513727300032158e+22
    nan

事实上以上述初猜，即使调变学习率也很难出现任何收敛的趋势。但对于下述使用了 Newton 二阶收敛的做法，一般来说只要学习率 `lr` 不超过 1，可以迅速收到极值点 (但也可能收到鞍点)：

```python
x = torch.tensor([-1., 1.], requires_grad=True) * 10
print("Initial x:", x)

lr = 1
for epoch in range(7):
    y = A @ x + B @ x @ x + C @ x @ x @ x
    print(float(y))
    g = grad(y, x, create_graph=True)[0]
    H = torch.stack([grad(g, x, grad_outputs=arr, create_graph=True)[0] for arr in torch.eye(d)]).detach()
    g = g.detach()
    x = x - torch.inverse(H) @ g * lr
    x = x.detach().clone().requires_grad_()
print("Optimized x:", x)
```

    Initial x: tensor([-10.,  10.], grad_fn=<MulBackward0>)
    465.36041259765625
    23.22213363647461
    0.332691490650177
    -0.36113929748535156
    -0.3630421757698059
    -0.3630422055721283
    -0.3630422055721283
    Optimized x: tensor([-0.47197, -0.03560], requires_grad=True)

### $y$ 的一阶导数

我们现在对一个任意确定的 `x` $x_i$ 来分析函数 `y` $y$ 求导的问题。

$$
y = \sum_i A_i x_i + \sum_{ij} B_{ij} x_i x_j + \sum_{ijk} C_{ijk} x_i x_j x_k
$$

```python
torch.random.manual_seed(0)
x = torch.randn(d, requires_grad=True)
y = A @ x + B @ x @ x + C @ x @ x @ x
```

任何优化问题需要求取一阶导数。一阶导数的求取我们应该已经比较熟悉了：

$$
\frac{\partial y}{\partial x_i} = A_i + \sum_j (B_{ij} + B_{ji}) x_j + \sum_{jk} (C_{ijk} + C_{jik} + C_{kji}) x_j x_k
$$

上述推导过程中需要利用到 $i, j, k$ 角标等价的技巧。

```python
A + (B + B.T) @ x + (C + C.transpose(0, 1) + C.transpose(0, 2)) @ x @ x
```

    tensor([ 6.09431, -2.24798], grad_fn=<AddBackward0>)

PyTorch 的自动求导程序可以通过如下方式求取一阶导数 `g`：

```python
g = grad(y, x, create_graph=True)[0]
g
```

    tensor([ 6.09431, -2.24798], grad_fn=<AddBackward0>)

### $y$ 的二阶导数

二阶导数通常也称为 Hessian 矩阵。

$$
\frac{\partial^2 y}{\partial x_i \partial x_j} = (B_{ij} + B_{ji}) + \sum_{k} (C_{ijk} + C_{ikj} + C_{jik} + C_{jki} + C_{kji} + C_{kij}) x_k
$$

上式仍然利用了 $j, k$ 角标等价的技巧。我们用 `Cp2` 表示六个 $\mathbf{C}$ 的转置张量之和：

```python
Cp2 = (
    C.permute(0, 1, 2) + C.permute(0, 2, 1)
  + C.permute(1, 0, 2) + C.permute(1, 2, 0)
  + C.permute(2, 0, 1) + C.permute(2, 1, 0)
)
```

那么 Hessian 矩阵可以写为

```python
B + B.T + Cp2 @ x
```

    tensor([[ 2.39952, -0.79906],
            [-0.79906,  2.35762]], grad_fn=<AddBackward0>)

我们 [曾经](./Autograd_TensorContract.ipynb#矩阵-向量乘积的自动求导) 指出过，在 PyTorch 中，真正的二阶梯度 `H` 是可以获得的，但需要依赖矩阵 $\delta_{ii_0}$ `torch.eye(d)`。获得二阶梯度的前提是一阶梯度 `g` 已经被求出，并且一阶梯度的导数图 (`create_graph=True` 选项) 已经被绘制。

```python
H = torch.stack([grad(g, x, grad_outputs=arr, retain_graph=True)[0] for arr in torch.eye(d)])
H
```

    tensor([[ 2.39952, -0.79906],
            [-0.79906,  2.35762]])

[^stackoverflow-qa]: <https://stackoverflow.com/questions/50322833/higher-order-gradients-in-pytorch>


```
%load_ext autoreload
%autoreload 2

from __future__ import division

import numpy as np
import matplotlib.pyplot as plt

import pydy
import sympy as sym
import sympy.physics.mechanics as me

from scipy.linalg import solve, lstsq
from scipy.integrate import odeint

import time
import functools

me.Vector.simp = False  # to increase computation speed

from sympy.physics.vector import init_vprinting
init_vprinting()

from matplotlib import animation
from matplotlib.patches import Ellipse, Rectangle

# from IPython.html import widgets
# from IPython.html.widgets import interact, interactive

import nchain

# %matplotlib inline
```

```
## 2-snake ##
n = 2

# group variables
x, y, z, phi, theta, psi = me.dynamicsymbols('x y z phi theta psi')  # generalized coordinates
xd, yd, zd, phid, thetad, psid = me.dynamicsymbols('x y z  phi theta psi', 1)
u, v, w, p, q, r = me.dynamicsymbols('u v w p q r')  # generalized speeds

# shape or configuration variables
qt = me.dynamicsymbols('q_t:' + str(n - 1))  # transverse
# qs = me.dynamicsymbols('q_s:' + str(n - 1))  # sagittal
# qc = me.dynamicsymbols('q_c:' + str(n - 1))  # coronal

ut = me.dynamicsymbols('u_t:' + str(n - 1))
# us = me.dynamicsymbols('u_s:' + str(n - 1))
# uc = me.dynamicsymbols('u_c:' + str(n - 1))

# constraints
cqt = me.dynamicsymbols('c_qt:' + str(n - 1))  # holonomic constraint
cut = me.dynamicsymbols('c_ut:' + str(n - 1), 1)  # non-holonomic constraint
cst = me.dynamicsymbols('c_st:' + str(n - 1), 2)  # used in matrix substitution later

# cqs = me.dynamicsymbols('c_qs:' + str(n - 1))  # holonomic constraint
# cus = me.dynamicsymbols('c_us:' + str(n - 1), 1)  # non-holonomic constraint
# css = me.dynamicsymbols('c_ss:' + str(n - 1), 2)  # used in matrix substitution later

# cqc = me.dynamicsymbols('c_qc:' + str(n - 1))  # holonomic constraint
# cuc = me.dynamicsymbols('c_uc:' + str(n - 1), 1)  # non-holonomic constraint
# csc = me.dynamicsymbols('c_sc:' + str(n - 1), 2)  # used in matrix substitution later

# forces and torques
F = me.dynamicsymbols('F:' + str(n + 1))  # n + 1 joints to apply forces/torques
T = me.dynamicsymbols('T:' + str(n + 1))

# variables for bodies
g, t = sym.symbols('g t')
m0, m1, m2 = sym.symbols('m_0 m_1 m_2')
l0, l1, l2 = sym.symbols('l_0 l_1 m_2')
if n > 3:
    Ixx0, Ixx1, Ixx2 = sym.symbols('I_xx_:' + str(n))
    Iyy0, Iyy1, Iyy2 = sym.symbols('I_yy_:' + str(n))
    Izz0, Izz1, Izz2 = sym.symbols('I_zz_:' + str(n))
Ixx0, Ixx1 = sym.symbols('I_xx_:' + str(n))
Iyy0, Iyy1 = sym.symbols('I_yy_:' + str(n))
Izz0, Izz1 = sym.symbols('I_zz_:' + str(n))

# inertial reference frame
frN = me.ReferenceFrame('fr_N')

# body-fixed frames
# fr0 = frN.orientnew('fr_0', 'Axis', [q0, frN.z])
# fr1 = fr0.orientnew('fr_1', 'Axis', [q1, fr0.z])
fr0 = frN.orientnew('fr_0', 'Body', [phi, theta, psi], '123')
#fr1 = fr0.orientnew('fr_1', 'Body', [qt[0], 0, 0], '123')  # not sure if this is the correct axis to rotate around
fr1 = fr0.orientnew('fr_1', 'Axis', [qt[0], fr0.z])  # this is actually the coronal axis from my notes on 2014-07-23
# fr2 = fr0.orientnew('fr_2', 'Body', [qt[1], qs[1], qc[1]], '123')

# angular velocities of frames
fr0.set_ang_vel(frN, p * frN.x + q * frN.y + r * frN.z)
# fr1.set_ang_vel(fr0, ut[0] * fr0.x + us[0] * fr0.y + uc[0] * fr0.z)
fr1.set_ang_vel(fr0, ut[0] * fr0.x)
#f2s.set_ang_vel(fr1, ...)

# origin of inertial frame
Nor = me.Point('N_or')

# joints
jt0 = Nor.locatenew('jt_0', x * frN.x + y * frN.y + z * frN.z)
jt1 = jt0.locatenew('jt_1', l0 * fr0.x)  # along transverse direction
jt2 = jt1.locatenew('jt_2', l1 * fr1.x)
# jt3 = jt2.locatenew('jt_3', l2 * fr2.x)

# mass centers
mc0 = jt0.locatenew('mc_0', l0 / 2 * fr0.x)
mc1 = jt1.locatenew('mc_1', l1 / 2 * fr1.x)
# mc2 = jt2.locatenew('mc_2', l2 / 2 * fr1.x)

# constrained point velocities
Nor.set_vel(frN, 0)
jt0.set_vel(frN, u * frN.x + v * frN.y + w * frN.z)

# unconstrained point velocities
jt1.v2pt_theory(jt0, frN, fr0)
jt2.v2pt_theory(jt1, frN, fr1)
# jt3.v2pt_theory(jt2, frN, ft2)

mc0.v2pt_theory(jt0, frN, fr0)
mc1.v2pt_theory(jt1, frN, fr1)
# mc2.v2pt_theory(jt2, frN, fr2)

# inertia dyads
# in0 = (me.inertia(fr0, Ixx0, Iyy0, Izz0), mc0)
# in1 = (me.inertia(fr1, Ixx1, Iyy1, Izz1), mc1)
in0 = (me.inertia(fr0, 0, 0, Izz0), mc0)
in1 = (me.inertia(fr1, 0, 0, Izz1), mc1)
# in2 = (me.inertia(fr2, Ixx2, Iyy2, Izz2), mc2)

# bodies
bd0 = me.RigidBody('bd_0', mc0, fr0, m0, in0)
bd1 = me.RigidBody('bd_1', mc1, fr1, m1, in1)
# bd2 = me.RigidBody('bd_2', mc2, fr2, m2, in2)

# kinematic differential equations
kd = [xd - u,
      yd - v,
      zd - w,
      phid - p,
      thetad - q,
      psid - r]

#TODO not sure if we are geting the order wrong here...
for i in range(n - 1):
    kd.append(qt[i].diff(t) - ut[i])
#     kd.append(qs[i].diff(t) - us[i])
#     kd.append(qc[i].diff(t) - uc[i])

# forces
fl = [(mc0, -m0 * g * frN.z),
      (mc1, -m1 * g * frN.z)]
#     (mc2, -m2 * g * frN.z)]

# fl = [(mc0, -m0 * g * frN.y),
#       (mc1, -m1 * g * frN.y)]

# bodies
bl = [bd0, bd1]  # bd2]

# q_ind and u_ind
qind = [x, y, z, phi, theta, psi]
uind = [u, v, w, p, q, r]

qdep = qt #+ qs + qc
udep = ut #+ us + uc

# constraints (they must equal zero)
# Note: qcons (holonomic constraints) only matter for linearization
cons_qt = [qq - cc for (qq, cc) in zip(qt, cqt)]
# cons_qs = [qq - cc for (qq, cc) in zip(qs, cqs)]
# cons_qc = [qq - cc for (qq, cc) in zip(qc, cqc)]

cons_ut = [uu - cc for (uu, cc) in zip(ut, cut)]
# cons_us = [uu - cc for (uu, cc) in zip(us, cus)]
# cons_uc = [uu - cc for (uu, cc) in zip(uc, cuc)]

qcons = cons_qt #+ cons_qs + cons_qc
ucons = cons_ut #+ cons_us + cons_uc

# equations of motion using Kane's method
kane = me.KanesMethod(frame=frN,
                      q_ind=qind,
                      u_ind=uind,
                      kd_eqs=kd,
                      q_dependent=qdep,
                      configuration_constraints=qcons,
                      u_dependent=udep,
                      velocity_constraints=ucons)

(fr, frstar) = kane.kanes_equations(fl, bl)
kanezero = fr + frstar

mass = kane.mass_matrix_full
forcing = kane.forcing_full
```

```
vel0 = mc0.vel(frN).express(fr0).simplify()
```

```
vel0
```

```
vel0y = vel0.dot(fr0.y)
vel0y
```

```
vel0z = vel0.dot(fr0.z)
vel0z
```

```
sym.atan2(vel0z, vel0y)
```

```
kane._k_ku
```

```
mm = kane.mass_matrix
```

```
np.shape(mm)
```

```
mm[-1, :]
```

```
kane._k_kqdot
```

```
kane._q
```

```
k_d = kane._k_d
```

```
np.shape(k_d)
```

```
kane._k_dnh
```

```
k_qdot = kane._k_kqdot
k_qdot
```

```
o = len(kane._u)  # 7
n = len(kane._q)  # 7
```

```
fst = ((kane._k_kqdot).row_join(sym.zeros(n, o)))
fst
```

```
snd = (sym.zeros(o, n)).row_join(kane.mass_matrix)
```

```
mm = fst.col_join(snd)
```

```
mm.shape
```

```
._partial_velocity([v.masscenter.vel(N),
                                                     v.frame.ang_vel_in(N)],
                                                    u, N)]
```

```
v = bl[0]

N = kane._inertial
vlist = [v.masscenter.vel(N), v.frame.ang_vel_in(N)]
ulist = kane._u
ulist
pv = kane._partial_velocity(vlist, ulist, N)
```

```
len(pv)
```

```
pv[0][3]
```

```
pv[0][4]
```

```

```

```
vlist[0].simplify().express(N).simplify()
```

```
vlist[0]
```

```
vlist[1]
```

```

```

```
((self._k_kqdot).row_join(zeros(n, o))).col_join()
```

```

```

```
aux = kane.auxiliary_eqs
```

```
kane.kanes_equations?
```

```
len(kanezero)
```

```
exp = kanezero[0]
```

```
sym.trigsimp(kanezero[0])
```

```
N = kane._inertial
vel_list = []
f_list = []
for i in fl:
    if isinstance(i[0], me.ReferenceFrame):
        vel_list += [i[0].ang_vel_in(N)]
    elif isinstance(i[0], me.Point):
        vel_list += [i[0].vel(N)]
    else:
        raise TypeError('First entry in pair must be point or frame.')
    f_list += [i[1]]
```

```
len(vel_list)
```

```

```

```
N
```

    fr_N

```
part_vel = kane._partial_velocity(vel_list, kane._u, kane._inertial)
```

```
part_vel[0][3]
```

```
map(len, part_vel)
```

```

```


###### Content under Creative Commons Attribution license CC-BY 4.0, code under MIT license (c)2014 L.A. Barba, G.F. Forsyth, I. Hawke. Partly based on [HyperPython](http://nbviewer.ipython.org/github/ketch/HyperPython/tree/master/) by D.I. Ketcheson, also under CC-BY.

# Riding the wave

This is the fourth and final lesson of Module 3, _Riding the wave: convection problems_, of the course **"Practical Numerical Methods with Python"** (a.k.a., [#numericalmooc](https://twitter.com/hashtag/numericalmooc)). We learned about conservation laws and the traffic-flow model in the [first lesson](http://nbviewer.ipython.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/03_wave/03_01_conservationLaw.ipynb), and then about better numerical schemes for convection in [lesson 2](http://nbviewer.ipython.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/03_wave/03_02_convectionSchemes.ipynb).

By then, you should have started to recognize that both mathematical models and numerical schemes work together to give us a good solution to a problem. To drive the point home, [lesson 3](http://nbviewer.ipython.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/03_wave/03_03_aBetterModel.ipynb) deals only with an improved model—and showed you some impressive SymPy tricks!

In this lesson, we'll learn about a new class of discretization schemes, known as finite-volume methods. They are the _most widely used_ methods in computational fluid dynamics, and for good reasons! Let's get started ...

## Finite-volume method

Are you curious to find out why the finite-volume method (FVM) is the _most popular method_ in computational fluid dynamics? In fact, almost all of the commercial CFD software packages are based on the finite-volume discretization. Here are some reasons:

* FVM discretizations are very general and have no requirement that the grid be structured, like in the finite-difference method. This makes FVM very _flexible_.

* FVM gives a _conservative discretization_ automatically by using directly the conservation laws in integral form.

### Conservative discretization

Let's go right back to the start of this module, where we explained conservation laws looking at a tiny control volume. To simplify the discussion, we just looked at flow in one dimension, with velocity $u$. Imagining a tiny cylindrical volume, like the one shown in Figure 1, there is flux on the left face and right face and we easily explained conservation of mass in that case.

#### Figure 1. Tiny control volume in the shape of a cylinder.

The law of conservation of mass says that the rate of change of mass in the control volume, plus the net rate of flow of mass across the control surfaces must be zero. The same idea works for other conserved quantities.

Conservation means that any change in the quantity within a volume is due to the amount of that quantity that crosses the boundary. Sounds simple enough. (Remember that we are ignoring possible internal sources of the quantity.) The amount crossing the boundary is the flux. A general conservation law for a quantity $e$ is thus:

\begin{equation}
\frac{\partial}{\partial t}\int_{\text{cv}}e \, dV + \oint_{\text{cs}}\vec{F}\cdot d\vec{A} =0
\end{equation}

where $\vec{F}$ is the flux, and $\text{cv}$ denotes the control volume with control surface $\text{cs}$.

**Why not make the control volume itself our computational cell?**

Imagine that the one-dimensional domain of interest is divided using grid points $x_i$. But instead of using local values at the grid points, like we did before, we now want to use _average_ values within each one-dimensional cell of width $\Delta x$ with center at $x_i$.

Define $e_i$ as the integral average across the little control volume on the cell with center at $x_i$ (see Figure 2).

$$
\begin{equation}
  e_i = \frac{1}{\Delta x} \int_{x_i - \Delta x / 2}^{x_i + \Delta x / 2} e(x, t) \, dx.
\end{equation}
$$

If we know the flux terms at the boundaries of the control volume, which are at $x_{i-1/2}$ and $x_{i+1/2}$, the general conservation law for this small control volume gives:

$$
\begin{equation}
  \frac{\partial}{\partial t} e_i + \frac{1}{\Delta x} \left[ F \left( x_{i+1/2}, t \right) - F \left( x_{i - 1 / 2}, t \right) \right] = 0.
\end{equation}
$$

This now just requires a time-stepping scheme, and is easy to solve *if* we can find $F$ on the control surfaces.

####Figure 2. Discretizing a 1D domain into finite volumes.

We've seen with the traffic model that the flux can depend on the conserved quantity (in that case, the traffic density). That is generally the case, so we write that $F = F(e)$. We will need to compute, or approximate, the flux terms at the cell edges (control surfaces) from the integral averages, $e_i$.

If we had a simple convection equation with $c>0$, then the flux going into the cell centered at $x_i$ from the left would be $F(e_{i-1})$ and the flux going out the cell on the right side would be $F(e_{i})$ (see Figure 2). Applying these fluxes in Equation (3) results in a scheme that is equivalent to our tried-and-tested backward-space (upwind) scheme!

We know from previous lessons that the backward-space scheme is first order and the error introduces numerical diffusion. Also, remember what happened when we tried to use it with the non-linear [traffic model](http://nbviewer.ipython.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/03_wave/03_01_conservationLaw.ipynb) in the green-light problem? *It blew up!* That was because the problem contains both right-moving and left-moving waves (if you don't remember that discussion, go back and review it; it's important!).

To skirt this difficulty in the green-light problem, we chose initial conditions that don't produce negative wave speeds. But that's cheating! A genuine solution would be to have a scheme that can deal with both positive and negative wave speeds. Here is where Godunov comes in.

### Godunov's method

Godunov proposed a first-order method in 1959 that uses the integral form of the conservation laws, Equation (1), and a piecewise constant representation of the solution, as shown in Figure (2). Notice that representing the solution in this way is like having a myriad of little shocks at the cell boundaries (control surfaces).

For each control surface, we have *two values* for the solution $e$ at a given time: the constant value to the left, $e_L$, and the constant value to the right, $e_R$. A situation where you have a conservation law with a constant initial condition, except for a single jump discontinuity is called a **Riemann problem**.

The Riemann problem has an exact solution for the Euler equations (as well as for any scalar conservation law). The [shock-tube problem](http://nbviewer.ipython.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/03_wave/03_05_Sods_Shock_Tube.ipynb), subject of your assignment for this course module, is a Riemann problem! And because it has an analytical solution, we can use it for testing the accuracy of numerical methods.

But Godunov had a better idea. With the solution represented as piecewise constant (Figure 2), why not use the analytical solution of the Riemann problem at each cell boundary? Solving a Riemann problem gives all the information about the characteristic structure of the solution, including the sign of the wave speed. The full solution can then be reconstructed from the union of all the Riemann solutions at cell boundaries. *Neat idea, Godunov!*

Figure 3 illustrates a Riemann problem for the Euler equations, associated to the shock tube. The space-time plot shows the characteristic lines for the left-traveling expansion wave, and the right-traveling contact discontinuity and shock.

####Figure 3. The shock tube: a Riemann problem for Euler's equations. Physical space (top) and $x, t$ space (bottom).

We need to solve many Riemann problems from $t$ to $\Delta t$, one on each cell boundary (illustrated in Figure 4). The numerical flux on $x_{i+1/2}$ is

\begin{equation}
  F_{i+1/2}= \frac{1}{\Delta t} \int_{t^n}^{t^{n+1}} F\left(e(x_{i+1/2},t) \right)\,dt
\end{equation}

To be able to solve each Riemann problem independently, they should not interact, which imposes a limit on $\Delta t$. Looking at Figure 4, you might conclude that we must require a CFL number of 1/2 to avoid interactions between the Riemann solutions, but the numerical flux above only depends on the state at $x_{i+1/2}$, so we're fine as long as the solution there is not affected by that at $x_{i-1/2}$—i.e., the CFL limit is really 1.

####Figure 4. Riemann problems on each cell boundary.

The Riemann solution, even though known analytically, can get quite hairy for non-linear systems of conservation laws (like the Euler equations). And we need as many Riemann solutions as there are finite-volume cell boundaries, and again at each time step! This gets really cumbersome.

Godunov solved the Riemann problems exactly, but many after him proposed *approximate* Riemann solutions instead. We'll be calculating the full solution numerically, after all, so some controlled approximations can be made. You might imagine a simple approximation for the flux at a cell boundary that is just the average between the left and right values, for example: $\frac{1}{2}\left[F(e_L)+F(e_R)\right]$. But that leads to a central scheme and, on its own, is unstable. Adding a term proportional to the difference between left and right states, $e_R-e_L$, supplies artificial dissipation and gives stability (see van Leer et al., 1987).

One formula for the numerical flux at $x_{i+1/2}$ called the Rusanov flux, a.k.a. Lax-Friedrichs flux, is given by

\begin{equation}
  F_{i+1/2}= \frac{1}{2} \left[ F \left( e_L \right) + F \left( e_R \right)  \right] -  \frac{1}{2}  \max \left|F'(e)\right| \left( e_R - e_L \right)
\end{equation}

where $F'(e)$ is the Jacobian of the flux function and $\max\left|F'(e)\right|$ is the local propagation speed of the fastest traveling wave. The Riemann solutions at each cell boundary do not interact if $\max|F'(e)|\leq\frac{\Delta x}{\Delta t}$, which leads to a flux formula we can now use:

\begin{equation}
  F_{i+1/2}= \frac{1}{2} \left( F \left( e_{i} \right) + F \left( e_{i+1} \right) -  \frac{\Delta x}{\Delta t} \left( e_{i+1} - e_{i} \right) \right)
\end{equation}

### Let's try it!

Let's apply Godunov's method to the [LWR traffic model](http://nbviewer.ipython.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/03_wave/03_01_conservationLaw.ipynb). In [lesson 2](http://nbviewer.ipython.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/03_wave/03_02_convectionSchemes.ipynb) we already wrote functions to set the initial conditions for a red-light problem and to compute the fluxes.  To save us from  writing this out again, we saved those functions into a Python file named `traffic.py` (found in the same directory of the course repository).  Now, we can use those functions by importing them in the same way that we import NumPy or any other library. Like this:

```python
from traffic import rho_red_light, computeF
```

You've probably noticed that we have the habit of writing detailed explanations of what a function does after defining it, in comments.  These comments are called *docstrings* and it is good practice to include them in all your functions.  It can be very useful when loading a function that you aren't familiar with (or don't remember!), because the `help` command will print them out for you.  Check it out:

```python
help(rho_red_light)
```

    Help on function rho_red_light in module traffic:

    rho_red_light(nx, rho_max, rho_in)
        Computes "red light" initial condition with shock

        Parameters
        ----------
        nx        : int
            Number of grid points in x
        rho_max   : float
            Maximum allowed car density
        rho_in    : float
            Density of incoming cars

        Returns
        -------
        rho: array of floats
            Array with initial values of density

Now, we can write some code to set up our notebook environment, and set the calculation parameters, with the functions imported above readily available.

```python
%matplotlib inline
import numpy
from matplotlib import pyplot
from matplotlib import rcParams
rcParams['font.family'] = 'serif'
rcParams['font.size'] = 16
from matplotlib import animation
from JSAnimation.IPython_display import display_animation
```

```python
#Basic initial condition parameters
#defining grid size, time steps, CFL condition, etc...
nx = 101
nt = 30
dx = 4.0/(nx-2)

rho_in = 5.
rho_max = 10.

V_max = 1.

x = numpy.linspace(0,4,nx-1)

rho = rho_red_light(nx-1, rho_max, rho_in)
```

```python
def animate(data):
    x = numpy.linspace(0,4,nx-1)
    y = data
    line.set_data(x,y)
    return line,
```

The cells above are code that you are already familiar with from lesson 2.  Below is a new function for applying Godunov's method with Lax-Friedrichs fluxes. Study it carefully.

```python
def godunov(rho, nt, dt, dx, rho_max, V_max):
    """ Computes the solution with the Godunov scheme using the Lax

    Parameters
    ----------
    rho    : array of floats
            Density at current time-step
    nt     : int
            Number of time steps
    dt     : float
            Time-step size
    dx     : float
            Mesh spacing
    rho_max: float
            Maximum allowed car density
    V_max  : float
            Speed limit

    Returns
    -------
    rho_n : array of floats
            Density after nt time steps at every point x
    """

    #initialize our results array with dimensions nt by nx
    rho_n = numpy.zeros((nt,len(rho)))
    #copy the initial u array into each row of our new array
    rho_n[:,:] = rho.copy()

    #setup some temporary arrays
    rho_plus = numpy.zeros_like(rho)
    rho_minus = numpy.zeros_like(rho)
    flux = numpy.zeros_like(rho)

    for t in range(1,nt):

        rho_plus[:-1] = rho[1:] # Can't do i+1/2 indices, so cell boundary
        rho_minus = rho.copy() # arrays at index i are at location i+1/2
        flux = 0.5 * (computeF(V_max, rho_max, rho_minus) +
                      computeF(V_max, rho_max, rho_plus) +
                      dx / dt * (rho_minus - rho_plus))
        rho_n[t,1:-1] = rho[1:-1] + dt/dx*(flux[:-2]-flux[1:-1])
        rho_n[t,0] = rho[0]
        rho_n[t,-1] = rho[-1]
        rho = rho_n[t].copy()

    return rho_n
```

We can run using a CFL of $1$, to start with, but you should experiment with different values.

```python
sigma = 1.0
dt = sigma*dx/V_max

rho = rho_red_light(nx-1, rho_max, rho_in)      #make sure that u is set to our expected initial conditions
rho_n = godunov(rho, nt, dt, dx, rho_max, V_max)
```

```python
fig = pyplot.figure();
ax = pyplot.axes(xlim=(0,4),ylim=(4.5,11),xlabel=('Distance'),ylabel=('Traffic density'));
line, = ax.plot([],[],color='#003366', lw=2);

anim = animation.FuncAnimation(fig, animate, frames=rho_n, interval=50)
display_animation(anim, default_mode='once')
```

    <input id="_anim_sliderPVUIZJWGJPMJQGNL" type="range" style="width:350px" name="points" min="0" max="1" step="1" value="0" onchange="animPVUIZJWGJPMJQGNL.set_frame(parseInt(this.value));"></input>

    <button onclick="animPVUIZJWGJPMJQGNL.slower()">–</button>
    <button onclick="animPVUIZJWGJPMJQGNL.first_frame()"></button>
    <button onclick="animPVUIZJWGJPMJQGNL.previous_frame()"></button>
    <button onclick="animPVUIZJWGJPMJQGNL.reverse_animation()"></button>
    <button onclick="animPVUIZJWGJPMJQGNL.pause_animation()"></button>
    <button onclick="animPVUIZJWGJPMJQGNL.play_animation()"></button>
    <button onclick="animPVUIZJWGJPMJQGNL.next_frame()"></button>
    <button onclick="animPVUIZJWGJPMJQGNL.last_frame()"></button>
    <button onclick="animPVUIZJWGJPMJQGNL.faster()">+</button>
  <form action="#n" name="_anim_loop_selectPVUIZJWGJPMJQGNL" class="anim_control">
    <input type="radio" name="state" value="once" checked> Once </input>
    <input type="radio" name="state" value="loop" > Loop </input>
    <input type="radio" name="state" value="reflect" > Reflect </input>
  </form>

You'll see that the result is very similar to the original Lax-Friedrichs method, and with good reason: they're essentially the same! But this is only because we are using a uniform grid. In the finite-volume approach, using the integral form of the equations, we were free to use a spatially varying grid spacing, if we wanted to.

The original Godunov method is first-order accurate, due to representing the conserved quantity by a piecewise-constant approximation. That is why you see considerable numerical diffusion in the solution. But Godunov's method laid the foundation for all finite-volume methods to follow and it was a milestone in numerical solutions of hyperbolic conservation laws. A whole industry developed inventing "high-resolution" methods that offer second-order accuracy and higher.

##### Dig deeper

* Godunov's method works in problems having waves moving with positive or negative wave speeds. Try it on the green-light problem introduced in [lesson 1](http://nbviewer.ipython.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/03_wave/03_01_conservationLaw.ipynb) using the initial condition containing waves traveling in both directions.

* Investigate two or three different numerical flux schemes (you can start with van Leer et al., 1987, or Google for other references. Implement the different flux schemes and compare!

## MUSCL schemes

Godunov's method is first-order accurate, which we already know is not appropriate for hyperbolic conservation laws, due to the high numerical diffusion. This poses particular difficulty near sharp gradients in the solution.

To do better, we can replace the piecewise constant representation of the solution with a piecewise linear version (still discontinuous at the edges). This leads to the MUSCL scheme (for Monotonic Upstream-Centered Scheme for Conservation Laws), invented by van Leer (1979).

### Reconstruction in space

The piecewise linear reconstruction consists of representing the solution inside each cell with a *straight line* (see Figure 5). Define the cell representation as follows:

\begin{equation}
  e(x) = e_i + \sigma_i (x - x_i).
\end{equation}

where $\sigma_i$ is the *slope* of the approximation within the cell (to be defined), and $e_i$ is the Godunov cell average. The choice $\sigma_i=0$ gives Godunov's method.

Standard central differencing would give

\begin{equation}
  \sigma_i = \frac{e_{i+1} - e_{i-1}}{2 \Delta x}.
\end{equation}

####Figure 5. Piecewise linear approximation of the solution.

But we saw with the results [in the second lesson](http://nbviewer.ipython.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/03_wave/03_02_convectionSchemes.ipynb) that this can lead to oscillations near shocks. These [Gibbs oscillations](http://en.wikipedia.org/wiki/Gibbs_phenomenon) will always appear (according to [Godunov's theorem](http://en.wikipedia.org/wiki/Godunov's_theorem)) unless we use constant reconstruction. So we have to modify, or *limit* the slope, near shocks.

The easiest way to limit is to compute one-sided slopes

\begin{equation}
  \Delta e^- = \frac{e_i - e_{i-1}}{\Delta x}, \quad \Delta e^+ = \frac{e_{i+1} - e_{i}}{\Delta x},
\end{equation}

#### Figure 6. One-sided slopes

Now build the *minmod* slope

\begin{align}
  \sigma_i & = \text{minmod}(\Delta e^-, \Delta e^+) \\
  & = \begin{cases} \min(\Delta e^-, \Delta e^+) & \text{ if } \Delta e^-, \Delta e^+ > 0 \\
  \max(\Delta e^-, \Delta e^+) & \text{ if } \Delta e^-, \Delta e^+ < 0 \\
0 & \text{ if } \Delta e^- \cdot \Delta e^+ \leq 0
  \end{cases}
\end{align}

That is, use the *smallest* one-sided slope in magnitude, unless the slopes have different sign, in which cases it uses the constant reconstruction (i.e., Godunov's method).

Once the *minmod* slope is calculated, we can use it to obtain the values at the interfaces between cells.

\begin{align}
e^{R}_{i-1/2} &= e_i - \sigma_i \frac{\Delta x}{2}\\
e^{L}_{i+1/2} &= e_i + \sigma_i \frac{\Delta x}{2}
\end{align}

where $e^R$ and $e^L$ are the local interpolated values of the conserved quantity immediately to the right and left of the cell boundary, respectively.

##### Index headache

Notice that for the cell with index $i$, we calculate $e^R_{i-1/2}$ and $e^L_{i+1/2}$.  Look at Figure 5: those are the two local values of the solution are at opposite cell boundaries.

However, when we calculate the local flux at the cell boundaries, we use the local solution values on either side of that cell boundary. That is:

\begin{equation}
F_{i+1/2} = f(e^L_{i+1/2}, e^R_{i+1/2})
\end{equation}

You can calculate two flux vectors; one for the right-boundary values and one for the left-boundary values. Be careful that you know which boundary value a given index in these two vectors might refer to!

_____

Here is a Python function implementing minmod.

```python
def minmod(e, dx):
    """
    Compute the minmod approximation to the slope

    Parameters
    ----------
    e : array of float
        input data
    dx : float
        spacestep

    Returns
    -------
    sigma : array of float
            minmod slope
    """

    sigma = numpy.zeros_like(e)
    de_minus = numpy.ones_like(e)
    de_plus = numpy.ones_like(e)

    de_minus[1:] = (e[1:] - e[:-1])/dx
    de_plus[:-1] = (e[1:] - e[:-1])/dx

    # The following is inefficient but easy to read
    for i in range(1, len(e)-1):
        if (de_minus[i] * de_plus[i] < 0.0):
            sigma[i] = 0.0
        elif (numpy.abs(de_minus[i]) < numpy.abs(de_plus[i])):
            sigma[i] = de_minus[i]
        else:
            sigma[i] = de_plus[i]

    return sigma
```

### Evolution in time

Since we are aiming for second-order accuracy in space, we might as well try for second-order in time, as well. We need a method to evolve the *ordinary* differential equation forwards in time:

\begin{equation}
  \frac{\partial}{\partial t} e_i + \frac{1}{\Delta x} \left[ F \left( x_{i+1/2}, t \right) - F \left( x_{i - 1 / 2}, t \right) \right] = 0
\end{equation}

A second-order Runge-Kutta method with special characteristics (due to Shu & Osher, 1988) gives the following scheme:

\begin{align}
e^*_i & = e^n_i + \frac{\Delta t}{\Delta x}\left( F^n_{i-1/2} - F^n_{i+1/2} \right) \\
e^{n+1}_i & = \frac{1}{2} e^n_i + \frac{1}{2}\left( e^*_i + \frac{\Delta t}{\Delta x}\left( F^*_{i-1/2} - F^*_{i+1/2} \right) \right)
\end{align}

Recall that the Rusanov flux is defined as

$$
F_{i+1/2}= \frac{1}{2} \left[ F \left( e_L \right) + F \left( e_R \right)  \right] -  \frac{1}{2}  \max \left|F'(e)\right| \left( e_R - e_L \right)
$$

Armed with the interpolated values of $e$ at the cell boundaries we can generate a more accurate Rusanov flux. At cell boundary $i+1/2$, for example, this is:

\begin{equation}
F_{i+1/2} = \frac{1}{2} \left( F \left( e^L_{i+1/2} \right) + F \left( e^R_{i+1/2} \right)   -  \frac{\Delta x}{\Delta t} \left( e^R_{i+1/2} - e^L_{i+1/2} \right) \right)
\end{equation}

Now we are ready to try some MUSCL!

```python
def muscl(rho, nt, dt, dx, rho_max, V_max):
    """ Computes the solution with the MUSCL scheme using the Lax-Friedrichs flux,
    RK2 in time and minmod slope limiting.

    Parameters
    ----------
    rho    : array of floats
            Density at current time-step
    nt     : int
            Number of time steps
    dt     : float
            Time-step size
    dx     : float
            Mesh spacing
    rho_max: float
            Maximum allowed car density
    V_max  : float
            Speed limit

    Returns
    -------
    rho_n : array of floats
            Density after nt time steps at every point x
    """

    #initialize our results array with dimensions nt by nx
    rho_n = numpy.zeros((nt,len(rho)))
    #copy the initial u array into each row of our new array
    rho_n[:,:] = rho.copy()

    #setup some temporary arrays
    rho_plus = numpy.zeros_like(rho)
    rho_minus = numpy.zeros_like(rho)
    flux = numpy.zeros_like(rho)
    rho_star = numpy.zeros_like(rho)

    for t in range(1,nt):

        sigma = minmod(rho,dx) #calculate minmod slope

        #reconstruct values at cell boundaries
        rho_left = rho + sigma*dx/2.
        rho_right = rho - sigma*dx/2.

        flux_left = computeF(V_max, rho_max, rho_left)
        flux_right = computeF(V_max, rho_max, rho_right)

        #flux i = i + 1/2
        flux[:-1] = 0.5 * (flux_right[1:] + flux_left[:-1] - dx/dt *\
                          (rho_right[1:] - rho_left[:-1] ))

        #rk2 step 1
        rho_star[1:-1] = rho[1:-1] + dt/dx * (flux[:-2] - flux[1:-1])

        rho_star[0] = rho[0]
        rho_star[-1] = rho[-1]

        sigma = minmod(rho_star,dx) #calculate minmod slope

        #reconstruct values at cell boundaries
        rho_left = rho_star + sigma*dx/2.
        rho_right = rho_star - sigma*dx/2.

        flux_left = computeF(V_max, rho_max, rho_left)
        flux_right = computeF(V_max, rho_max, rho_right)

        flux[:-1] = 0.5 * (flux_right[1:] + flux_left[:-1] - dx/dt *\
                          (rho_right[1:] - rho_left[:-1] ))

        rho_n[t,1:-1] = .5 * (rho[1:-1] + rho_star[1:-1] + dt/dx * (flux[:-2] - flux[1:-1]))

        rho_n[t,0] = rho[0]
        rho_n[t,-1] = rho[-1]
        rho = rho_n[t].copy()

    return rho_n
```

```python
sigma = 1.
dt = sigma*dx/V_max
rho = rho_red_light(nx-1, rho_max, rho_in)      #make sure that u is set to our expected initial conditions
rho_n = muscl(rho, nt, dt, dx, rho_max, V_max)
```

```python
fig = pyplot.figure();
ax = pyplot.axes(xlim=(0,4),ylim=(4.5,11),xlabel=('Distance'),ylabel=('Traffic density'));
line, = ax.plot([],[],color='#003366', lw=2);

anim = animation.FuncAnimation(fig, animate, frames=rho_n, interval=50)
display_animation(anim, default_mode='once')
```

    <input id="_anim_sliderMRLEOWBXXJCQHKGO" type="range" style="width:350px" name="points" min="0" max="1" step="1" value="0" onchange="animMRLEOWBXXJCQHKGO.set_frame(parseInt(this.value));"></input>

    <button onclick="animMRLEOWBXXJCQHKGO.slower()">–</button>
    <button onclick="animMRLEOWBXXJCQHKGO.first_frame()"></button>
    <button onclick="animMRLEOWBXXJCQHKGO.previous_frame()"></button>
    <button onclick="animMRLEOWBXXJCQHKGO.reverse_animation()"></button>
    <button onclick="animMRLEOWBXXJCQHKGO.pause_animation()"></button>
    <button onclick="animMRLEOWBXXJCQHKGO.play_animation()"></button>
    <button onclick="animMRLEOWBXXJCQHKGO.next_frame()"></button>
    <button onclick="animMRLEOWBXXJCQHKGO.last_frame()"></button>
    <button onclick="animMRLEOWBXXJCQHKGO.faster()">+</button>
  <form action="#n" name="_anim_loop_selectMRLEOWBXXJCQHKGO" class="anim_control">
    <input type="radio" name="state" value="once" checked> Once </input>
    <input type="radio" name="state" value="loop" > Loop </input>
    <input type="radio" name="state" value="reflect" > Reflect </input>
  </form>

This MUSCL scheme does not show any of the oscillations you might see with MacCormack or Lax-Wendroff, but the features are not as sharp. Using the _minmod_ slopes led to some smearing of the shock, which motivated many researchers to investigate other options. Bucketloads of so-called _shock-capturing_ schemes exist and whole books are written on this topic. Some people dedicate their lives to developing numerical methods for hyperbolic equations!

##### Challenge task

* Go back to Sod! Calculate the shock-tube problem using the MUSCL scheme and compare with your previous results. What do think?

## References

* Godunov, S.K. (1959), "A difference scheme for numerical computation of discontinuous solutions of equations of fluid dynamics," _Math. Sbornik_, Vol. 47, pp. 271–306.

* van Leer, Bram (1979), "Towards the ultimate conservative difference scheme, V. A second-order sequel to Godunov's method," _J. Comput. Phys._, Vol. 32, pp. 101–136

* van Leer, B., J.L. Thomas, P.L. Roe, R.W. Newsome (1987). "A comparison of numerical flux formulas for the Euler and Navier-Stokes equations," AIAA paper 87-1104 // [PDF from umich.edu](http://deepblue.lib.umich.edu/bitstream/handle/2027.42/76365/AIAA-1987-1104-891.pdf), checked 11/01/14.

* Shu, Chi-Wang and Osher, Stanley (1988). "Efficient implementation of essentially non-oscillatory shock-capturing schemes," _J. Comput. Phys._, Vol. 77, pp. 439–471 // [PDF from NASA Tech. Report server](http://ntrs.nasa.gov/archive/nasa/casi.ntrs.nasa.gov/19870013797.pdf)

---

######The cell below loads the style of the notebook.

```python
from IPython.core.display import HTML
css_file = '../../styles/numericalmoocstyle.css'
HTML(open(css_file, "r").read())
```

<link href='http://fonts.googleapis.com/css?family=Alegreya+Sans:100,300,400,500,700,800,900,100italic,300italic,400italic,500italic,700italic,800italic,900italic' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Arvo:400,700,400italic' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=PT+Mono' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Shadows+Into+Light' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Nixie+One' rel='stylesheet' type='text/css'>


```python
# %load science.py
%matplotlib notebook
from sympy import init_session, symbols
import numpy as np
import matplotlib.pyplot as plt
import scipy as sp
import pandas as pd
import theano as T
#from sympy.geometry import *
init_session()
a,b,c = symbols('a b c')
r, theta, phi = symbols('r theta phi', positive=True)
print('Loaded a b c theta and phi')
print('Load Theano as T')
```

    IPython console for SymPy 1.0 (Python 3.5.2-64-bit) (ground types: python)

    These commands were executed:
    >>> from __future__ import division
    >>> from sympy import *
    >>> x, y, z, t = symbols('x y z t')
    >>> k, m, n = symbols('k m n', integer=True)
    >>> f, g, h = symbols('f g h', cls=Function)
    >>> init_printing()

    Documentation can be found at http://docs.sympy.org/1.0/
    Loaded a b c theta and phi
    Load Theano as T

## 11.2: Cantor Set

Highly recommended is Daniel Shiffman's [The Nature of Code](http://natureofcode.com/book/chapter-8-fractals/), which uses Processing to generate some awesome reusable code-based visualisations to explore.

```python
#What is the limit of  a set?
limit(Rational(2,3)**n, z=n, z0=oo, dir='+')
```

```python
#Dimension of Self-similar Fractals.
n_copies = float(input("How many copies does one rescaling create? "))
r_scale_factor = float(input("What is the rescaling factor? "))
dimension = np.log(n_copies)/np.log(r_scale_factor)
print("The dimension, d, is {0:.2f}".format(dimension))
```

    How many copies does one rescaling create? 9
    What is the rescaling factor? 3
    The dimension, d is 2.00

## 11.4: Box and Hausdorff Dimension

Noting that Box dimension is not used in practice, Hausdorff dimension can be done using a prewritten module [py-hausdorff](https://github.com/mavillan/py-hausdorff), whihc uses a validated method. The method leverages of Cython and there is a C-code for use, should one be bothered (which for large sets one should indeed be bothered!).

## 11.5: Pointwise and Correlation Dimensions (p.418)

Again, the problem is not at all [trivial](http://www.scholarpedia.org/article/Grassberger-Procaccia_algorithm) in terms of computational requirements. One approach is to use an already robust package, such as [PyDSTool](http://www.ni.gsu.edu/~rclewley/PyDSTool/FrontPage.html).

[PyDSTool code](http://pydoc.net/PyDSTool/0.90.2/PyDSTool.Toolbox.fracdim/)

The basic process is:
A python approach to correlation of mutli-arrays...[Correlation on array](Chttps://stackoverflow.com/questions/30143417/computing-the-correlation-coefficient-between-two-multi-dimensional-arrays)
A readable article on background ...[Algorithm](http://media.wiley.com/product_data/excerpt/25/07803601/0780360125.pdf)
And this by [Sprott](http://thor.physics.wisc.edu/pubs/paper255.pdf)

```python
#out = np.dot(arr_one,arr_two.T) #MxT; NxT arrays
def corr2_coeff(A,B):
    # Rowwise mean of input arrays & subtract from input arrays themeselves
    A_mA = A - A.mean(1)[:,None]
    B_mB = B - B.mean(1)[:,None]

    # Sum of squares across rows
    ssA = (A_mA**2).sum(1);
    ssB = (B_mB**2).sum(1);

    # Finally get corr coeff
    return np.dot(A_mA,B_mB.T)/np.sqrt(np.dot(ssA[:,None],ssB[None]))
```

```python
a1 = np.random.rand(10,3)
a2 = np.random.rand(10,3)
ans = corr2_coeff(a1,a2)
ans.shape
```

```python

```