# KinMS galaxy fitting tutorial

This tutorial aims at getting you up and running with galaxy kinematic modelling using _KinMS_! To start you will need to download the KinMSpy code and have it in your python path.

To do this you can simply call `pip install kinms`

To get started with kinematic modelling we will complete the following steps:
1. Generate a model to fit (can be skipped if you have your own observed data cube)
2. Read in that cube, and extract the important information from the header
3. Fit the data using an MCMC code

We will start by importing a variety of modules we will need to work with KinMS, and plot its output.

```python
from kinms import KinMS
import numpy as np
from astropy.io import fits
from kinms.utils.KinMS_figures import KinMS_plotter
```

## Generate a model

First we will generate a simple galaxy model using KinMS itself, that we can attempt to determine the parameters of later. If you have your own observed galaxy to fit then of course this step can be skipped!

The `make_model` function below creates a simple exponential disc:

$
\begin{align}
\large \Sigma_{H2}(r) \propto e^{\frac{-r}{d_{scale}}}
\end{align}
$

with a circular velocity profile which is parameterized using an arctan function:

$
\begin{align}
\large V(r) = \frac{2V_{flat}}{\pi} \arctan\left(\frac{r}{r_{turn}}\right)
\end{align}
$

```python
def make_model(param,obspars,rad,filename=None,plot=False):
    '''
    This function takes in the `param` array (along with obspars; the observational setup,
    and a radius vector `rad`) and uses it to create a KinMS model.
    '''

    total_flux=param[0]
    posAng=param[1]
    inc=param[2]
    v_flat=param[3]
    r_turn=param[4]
    scalerad=param[5]

    ### Here we use an exponential disk model for the surface brightness of the gas ###
    sbprof = np.exp((-1)*rad/scalerad)

    ### We use a very simple arctan rotation curve model with two free parameters. ###
    vel=(v_flat*2/np.pi)*np.arctan(rad/r_turn)

    ### This returns the model
    return KinMS(obspars['xsize'],obspars['ysize'],obspars['vsize'],obspars['cellsize'],obspars['dv'],\
                 obspars['beamsize'],inc,sbProf=sbprof,sbRad=rad,velRad=rad,velProf=vel,\
                 intFlux=total_flux,posAng=posAng,fixSeed=True,fileName=filename).model_cube(toplot=plot)

```

Note that we have set `fixSeed=True` in the KinMS call - this is crucial if you are fitting with KinMS. It ensures if you generate two models with the same input parameters you will get an identical output model!

Now we have our model function, lets use it to generate a model which we will later fit. The first thing we need is to define the setup of our desired datacube (typically if you are fitting real data this will all be determined from the header keywords- see below).

```python
### Setup cube parameters ###
obspars={}
obspars['xsize']=64.0 # arcseconds
obspars['ysize']=64.0 # arcseconds
obspars['vsize']=500.0 # km/s
obspars['cellsize']=1.0 # arcseconds/pixel
obspars['dv']=20.0 # km/s/channel
obspars['beamsize']=np.array([4.0,4.0,0]) # [bmaj,bmin,bpa] in (arcsec, arcsec, degrees)
```

We also need to create a radius vector- you ideally want this to oversample your pixel grid somewhat to avoid interpolation errors!

```python
rad=np.arange(0,100,0.3)
```

Now we have all the ingredients we can create our data to fit. Here we will also output the model to disc, so we can demonstrate how to read in the header keywords from real ALMA/VLA etc data.

```python
'''
True values for the flux, posang, inc etc, as defined in the model function
'''

guesses=np.array([30.,270.,45.,200.,2.,5.])

'''
RMS of data. Here we are making our own model so this is arbitary.
When fitting real data this should be the observational RMS
'''
error=np.array(1e-3)

fdata=make_model(guesses,obspars,rad, filename="Test",plot=True)
```

## Read in the data

In this example we already have our data in memory. But if you are fitting a real datacube this wont be the case!  Here we read in the model we just created from a FITS file to make it clear how to do this.

```python
### Load in your observational data ###
hdulist = fits.open('Test_simcube.fits',ignore_blank=True)
fdata = hdulist[0].data.T

### Setup cube parameters ###
obspars={}
obspars['cellsize']=np.abs(hdulist[0].header['cdelt1']*3600.) # arcseconds/pixel
obspars['dv']=np.abs(hdulist[0].header['cdelt3']/1e3) # km/s/channel
obspars['xsize']=hdulist[0].header['naxis1']*obspars['cellsize'] # arcseconds
obspars['ysize']=hdulist[0].header['naxis2']*obspars['cellsize'] # arcseconds
obspars['vsize']=hdulist[0].header['naxis3']*obspars['dv'] # km/s
obspars['beamsize']=np.array([hdulist[0].header['bmaj']*3600.,hdulist[0].header['bmin']*3600.,hdulist[0].header['bpa']])# [bmaj,bmin,bpa] in (arcsec, arcsec, degrees)

```

## Fit the model

Now we have our 'observational' data read into memory, and a model function defined, we can fit one to the other! As our fake model is currently noiseless, lets add some gaussian noise (obviously dont do this if your data is from a real telecope!):

```python
fdata+=(np.random.normal(size=fdata.shape)*error)
```

Below we will proceed using the MCMC code GAStimator which was specifically designed to work with KinMS, however any minimiser should work in principle. For full details of how this code works, and a tutorial, see https://github.com/TimothyADavis/GAStimator .

```python
from gastimator import gastimator,corner_plot

mcmc = gastimator(make_model,obspars,rad)

mcmc.labels=np.array(['Flux','posAng',"Inc","VFlat","R_turn","scalerad"])
mcmc.min=np.array([30.,1.,10,50,0.1,0.1])
mcmc.max=np.array([30.,360.,80,400,20,10])
mcmc.fixed=np.array([True,False,False,False,False,False])
mcmc.precision=np.array([1.,1.,1.,10,0.1,0.1])
mcmc.guesses=np.array([30.,275.,55.,210.,2.5,4.5]) #starting guesses, purposefully off!
```

Setting good priors on the flux of your source is crucial to ensure the model outputs are physical. Luckily the integrated flux of your source should be easy to measure from your datacube! If you have a good measurement of this, then I would recommend forcing the total flux to that value by fixing it in the model (set `mcmc.fixed=True` for that parameter). If you can only get a guess then set as tight a prior as you can. This stops the model hiding bad fitting components below the noise level.

Its always a good idea to plot your model over your data before you start a fitting processes. That allows you to check that the model is reasonable, and tweak the parameters by hand to get good starting guesses. Firs you should generate a cube from your model function, then you can overplot it on your data using the simple plotting tool included with KinMS:

```python
model=make_model(mcmc.guesses,obspars,rad) # make a model from your guesses

KinMS_plotter(fdata, obspars['xsize'], obspars['ysize'], obspars['vsize'], obspars['cellsize'],\
              obspars['dv'], obspars['beamsize'], posang=guesses[1],overcube=model,rms=error).makeplots()
```

As you can see, the black contours of the model arent a perfect match to the moment zero, spectrum and position-velocity diagram extracted from our "observed" datacube. One could tweak by hand, but as these are already close we can go on to do a fit!

If you are experimenting then running until convergence should be good enough to get an idea if the model is physical (setting a low number of iterations, ~3000 works for me).

```python
outputvalue, outputll= mcmc.run(fdata,error,3000,plot=False)
```

    Doing chain 1
         Chain has not converged - Accept rate: 0.08666666666666667
         Chain has not converged - Accept rate: 0.1
         Still varying: ['posAng' 'Inc' 'R_turn' 'scalerad']
         Chain has not converged - Accept rate: 0.2866666666666667
         Still varying: ['Inc']
         Chain has not converged - Accept rate: 0.45666666666666667
         Still varying: ['Inc']
    Chain converged: LL: -5417.6915586532605 - Accept rate:0.43
    Best chain so far!
    Best fit:
      Flux: 30.0
      posAng: 269.90007276088943
      Inc: 47.854655719347996
      VFlat: 190.0066572057255
      R_turn: 1.9306907362857078
      scalerad: 4.931147911070421
    Starting final chain

    [Parallel(n_jobs=3)]: Using backend LokyBackend with 3 concurrent workers.

    Final best fit values and 1sigma errors:
      Flux: 30.0 (Fixed)
      posAng: 269.8973861297394 + 0.046268424097263505 - 0.017317013075341947
      Inc: 46.69701598839461 + 0.10982841217780503 - 0.20773477781962413
      VFlat: 193.9123324657629 ± 0.34437214188913856
      R_turn: 1.9498716985274913 + 0.013209686719802827 - 0.010186707854953791
      scalerad: 4.8752275278090975 ± 0.013733314328538526

    [Parallel(n_jobs=3)]: Done   3 out of   3 | elapsed:  6.6min remaining:    0.0s
    [Parallel(n_jobs=3)]: Done   3 out of   3 | elapsed:  6.6min finished

As you can see, the final parameters (listed in the output with their 1sigma errors) are pretty close to those we input! One could use the cornor_plot routine shipped with GAStimator to visualize our results, but with only 3000 steps (and a $\approx$30% acceptance rate) these wont be very pretty. If you need good error estimates/nice looking cornor plots for publication then I recommend at least 30,000 iterations, which may take several hours/days depending on your system, and the size of your datacube.

One can visualize the best-fit model again to check how we did - turns out pretty well! (Note the flux in the integrated spectrum isnt perfect, this is because of the masking of the noisy data).

```python
bestmodel=make_model(np.median(outputvalue,1),obspars,rad) # make a model from your guesses

KinMS_plotter(fdata, obspars['xsize'], obspars['ysize'], obspars['vsize'], obspars['cellsize'],\
              obspars['dv'], obspars['beamsize'], posang=guesses[1],overcube=bestmodel,rms=error).makeplots()
```

## Tiny error problem

I have found that fitting whole datacubes with kinematic modelling tools such as KinMS can yield unphysically small uncertanties, for instance constraining inclination to $\pm\approx0.1^{\circ}$ in the fit example performed above. This is essentially a form of model mismatch - you are finding the very best model _of a given type_ that fits the data - and as you have a large number of free-parameters in a data cube you can find the best model (no matter how bad it is at actually fitting the data!) really well.

In works such as [Smith et al. (2019)](https://ui.adsabs.harvard.edu/abs/2019MNRAS.485.4359S/abstract) we have attempted to get around by taking into account the variance of the $\chi^2$ statistic.

As observed data are noisy, the $\chi^2$ statistic has an additional uncertainty associated with it, following the chi-squared distribution ([Andrae 2010](https://ui.adsabs.harvard.edu/abs/2010arXiv1009.2755A/abstract)). This distribution has a variance of $2(N - P)$, where $N$ is the number of constraints and $P$ the number of inferred parameters. For fitting datacubes $N$ is very large, so the variance becomes $\approx2N$.
Systematic effects can produce variations of $\chi^2$ of the order of this variance, and ignoring this effect yields unrealistically small uncertainty estimates. In order to mitigate this effect [van
den Bosch & van de Ven (2009)](https://ui.adsabs.harvard.edu/abs/2009MNRAS.398.1117V/abstract) proposed to increase the $1\sigma$ confidence interval to $\Delta\chi^2=\sqrt{2N}$. To achieve the same effect within the Bayesian MCMC approach discussed above we need to scale the log-likelihood, by increasing the RMS estimate provided to GAStimator  by $(2N)^{1/4}$. This approach appears to yield physically credible formal uncertainties in the inferred parameters, whereas otherwise these uncertainties are unphysically small.

Lets try that with the example above:

```python
error*=((2.0*fdata.size)**(0.25))
outputvalue, outputll= mcmc.run(fdata,error,3000,plot=False)
```

    Doing chain 1
         Chain has not converged - Accept rate: 0.11
    Chain converged: LL: -112.97430961201346 - Accept rate:0.2966666666666667
    Best chain so far!
    Best fit:
      Flux: 30.0
      posAng: 270.0
      Inc: 44.957062230395515
      VFlat: 201.01285763281135
      R_turn: 2.03929515740805
      scalerad: 5.013047475148751
    Starting final chain
    Final best fit values and 1sigma errors:
      Flux: 30.0 (Fixed)
      posAng: 269.8723305066128 + 0.946321323228176 - 1.1107270271547236
      Inc: 45.781525045651954 + 2.1113884290902405 - 0.9361234945474166
      VFlat: 199.09454514602965 + 3.294369458631081 - 5.327156135062893
      R_turn: 2.110158720751331 + 0.15726926408759345 - 0.18741126338718428
      scalerad: 5.029909360644904 + 0.20293482085984138 - 0.2480078138126771

As you can see we now get a much more reasonable error estimates, for instance a 1$\sigma$ inclination error of $\pm2^{\circ}$.

If you want to implement this fix yourself there is a wrinkle to consider. You need to be careful choosing $N$. Formally this should be the number of constraints- i.e. the total number of pixels in your cube. But consider a large datacube with signal only in a small section (although fitting such a datacube would be inefficient anyway; see speed tips below), all of the actual constraints are coming from a small number of pixels.  If you find yourself in this situation I would recommend setting $N$ to the number of pixels with actual detected flux in your datacube.

## Speed tips

There are some common ways to make sure you don't spend a lot of time MCMCing rather than doing science.

- Cut down your observed data to only include spaxels and frequency channels near signal. Ideally you want some padding around your observed signal so the model knows it must not include flux in those positions, but not so much as to drastically increase runtime! On a similar note...

- Make sure your observed data (and thus models) have spatial dimensions that are $2^n$. If this is impossible then $3^n$ and $6^n$ are pretty good too. This is because convolving the KinMS model with the beam takes the majority of the computation time, and FFTs are faster when working with such dimensions.

- Don't provide a radius vector that is very oversampled/overlong well beyond the projected dimensions of the cube. This can slow down internal interpolation routines.


```python
from sympy import init_session
init_session()
```

    IPython console for SymPy 1.5.1 (Python 3.6.9-64-bit) (ground types: gmpy)

    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 https://docs.sympy.org/1.5.1/

```python
init_printing?
```

```python
# this is for rac[3,1]
#
#  A=a**2, B=b**2, and D=d**2
#
A, B, D, k, l0 = symbols('A, B, D, k, l0')
```

```python
f = k**2 + 2*A*k + A**2 + B
g = 1 + D*k
h = A**2 + B + 2*A*k + D*A**2*k + D*B*k
y = l0*f*g/h
y
```

```python
diff(y,l0)
```

```python
dA=diff(y,A)
dA.factor()
```

```python
dB=diff(y,B)
dB.factor()
```

```python
dD=diff(y,D)
dD.factor()
```

```python

```

```python
E, G, A, beta = symbols('E, G, A, beta', real=True)
```

```python
nonlinsolve([E - beta**2 + A**2, G - 4*A*beta], [E, G])
```

```python
solve([E - beta**2 + A**2, G - 4*A*beta], [A,beta])
```

```python
symbols?
```

```python

```


# Hypersolvers for Optimal Control

### Basic usage of the `Hypersolver` API

We showcase the `torchdyn` API for [hypersolvers](https://arxiv.org/pdf/2007.09601.pdf). Hypersolvers, or hybrid ODE solvers equipped with a neural network to approximate residuals, are seamlessly integrated into the `odeint` API. A major design of the API is exactly preserving a persistent state for the solver, which in the case of standard ODE solvers contains the Tableau. For hypersolvers, the state also contains their hypernetwork parameters.

```python
import time
import torch
import torch.nn as nn
import matplotlib.pyplot as plt

from torchdyn.core import NeuralODE
from torchdyn.datasets import *
from torchdyn.numerics import odeint, Euler, HyperEuler
```

    /home/botu/.local/lib/python3.8/site-packages/pytorch_lightning/metrics/__init__.py:43: LightningDeprecationWarning: `pytorch_lightning.metrics.*` module has been renamed to `torchmetrics.*` and split off to its own package (https://github.com/PyTorchLightning/metrics) since v1.3 and will be removed in v1.5
      rank_zero_deprecation(

## Optimal Control

We want to control an inverted pendulum and stabilize it in the upright position. The equations in Hamiltonian form describing an inverted pendulum with a torsional spring are as following:

\begin{equation}
    \begin{bmatrix} \dot{q}\\ \dot{p}\\ \end{bmatrix} =
    \begin{bmatrix}
    0& 1/m \\
    -k& -\beta/m\\
    \end{bmatrix}
    \begin{bmatrix} q\\ p\\ \end{bmatrix} -
    \begin{bmatrix}
    0\\
    mgl \sin{q}\\
    \end{bmatrix}+
    \begin{bmatrix}
    0\\
    1\\
    \end{bmatrix} u
\end{equation}

```python
class ControlledPendulum(nn.Module):
    """
    Inverted pendulum with torsional spring
    """
    def __init__(self, u, m=1., k=.5, l=1., qr=0., β=.01, g=9.81):
        super().__init__()
        self.u = u # controller (nn.Module)
        self.nfe = 0 # number of function evaluations
        self.cur_f = None # current function evaluation
        self.cur_u = None # current controller evaluation
        self.m, self.k, self.l, self.qr, self.β, self.g = m, k, l, qr, β, g # physics

    def forward(self, t, x):
        self.nfe += 1
        q, p = x[..., :1], x[..., 1:]
        self.cur_u = self.u(t, x)
        dq = p/self.m
        dp = -self.k*(q - self.qr) - self.m*self.g*self.l*torch.sin(q) \
            -self.β*p/self.m + self.cur_u
        self.cur_f = torch.cat([dq, dp], -1)
        return self.cur_f
```

In order to control the pendulum, we have to define a proper _integral cost function_ which will be our loss to be minimized during training. In a general form, it can be defined as:

\begin{equation}
        J = x^\top(t_f)\mathbf{P} x(t_f) + \int_{t_0}^{t_f} \left[ x^\top(t) \mathbf{Q} x(t) + u^\top(t) \mathbf{R} u(t) \right] dt
\end{equation}

where $ x = \begin{bmatrix} q\\ p\\ \end{bmatrix}$ is the state and $\mathbf{u}$ is the controller and matrices $\mathbf{P},~\mathbf{Q}, ~ \mathbf{R}$ are weights for controlling the performance.

```python
class IntegralCost(nn.Module):
    '''Integral cost function
    Args:
        x_star: torch.tensor, target position
        P: float, terminal cost weights
        Q: float, state weights
        R: float, controller regulator weights
    '''
    def __init__(self, x_star, P=0, Q=1, R=0):
        super().__init__()
        self.x_star = x_star
        self.P, self.Q, self.R, = P, Q, R

    def forward(self, x, u=torch.Tensor([0.])):
        """
        x: trajectory
        u: control input
        """
        cost = self.P*torch.norm(x[-1] - self.x_star, p=2, dim=-1).mean()
        cost += self.Q*torch.norm(x - self.x_star, p=2, dim=-1).mean()
        cost += self.R*torch.norm(u - 0, p=2).mean()
        return cost
```

```python
# Change device according to your configuration
device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
```

```python
# The controller is a simple MLP with one hidden layer with bounded output
class NeuralController(nn.Module):
    def __init__(self, model, u_min=-20, u_max=20):
        super().__init__()
        self.model = model
        self.u_min, self.u_max = u_min, u_max

    def forward(self, t, x):
        x = self.model(x)
        return torch.clamp(x, self.u_min, self.u_max)

model = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 1)).to(device)
u = NeuralController(model)
for p in u.model[-1].parameters(): torch.nn.init.zeros_(p)

# Controlled system
sys = ControlledPendulum(u).to(device)
```

```python
from math import pi as π

# Loss function declaration
x_star = torch.Tensor([0., 0.]).to(device)
cost_func = IntegralCost(x_star)

# Time span
t0, tf = 0, 2 # initial and final time for controlling the system
steps = 20 + 1 # so we have a time step of 0.1s
t_span = torch.linspace(t0, tf, steps).to(device)

# Initial distribution
x0 = π # limit of the state distribution (in rads and rads/second)
init_dist = torch.distributions.Uniform(torch.Tensor([-x0, -x0]), torch.Tensor([x0, x0]))
```

```python
# Hyperparameters
lr = 3e-3
epochs = 500
bs = 1024
opt = torch.optim.Adam(u.parameters(), lr=lr)

# Training loop
t0 = time.time(); losses=[]
for e in range(epochs):
    x0 = init_dist.sample((bs,)).to(device)
    _, trajectory = odeint(sys, x0, t_span, solver='tsit5', atol=1e-7, rtol=1e-7)
    loss = cost_func(trajectory); losses.append(loss.detach().cpu().item())
    loss.backward(); opt.step(); opt.zero_grad()
    print('Loss {:.4f} , epoch {}'.format(loss.item(), e), end='\r')
timing = time.time() - t0; print('\nTraining time: {:.4f} s'.format(timing))
```

    Loss 0.9749 , epoch 499
    Training time: 228.4546 s

```python
# Testing the controller on the real system
x0 = init_dist.sample((100,)).to(device)
t0, tf, steps = 0, 2, 20*10 + 1
t_span_fine = torch.linspace(t0, tf, steps).to(device)
_, traj = odeint(sys, x0, t_span_fine, solver='tsit5', atol=1e-7, rtol=1e-7)
t_span_fine = t_span_fine.detach().cpu(); traj = traj.detach().cpu()

# Plotting
fig, ax = plt.subplots(1, 1, figsize=(12,6))
for i in range(len(x0)):
    ax.plot(t_span_fine, traj[:,i,0], 'k-.', alpha=.3, label='p' if i==0 else None)
    ax.plot(t_span_fine, traj[:,i,1], 'b', alpha=.3, label='q' if i==0 else None)
ax.legend()
ax.set_title('Controlled trajectories')
ax.set_xlabel(r"$t~[s]$")
```

Now we have to address one question: can we speed up the simulation?

## Hypersolvers for Controlled Systems

By training the hypersolver on residuals between the ground-truth solution and the base solver (we will use Euler in this case) we can render the training more efficient while retaining solution accuracy.

The residuals can be defined by:
\begin{equation}
\begin{aligned}
    R \left(t_k,  x(t_k), x(t_{k+1}) \right) = \Phi(x(t_k), t_k, t_{k+1}) - x(t_k) - \epsilon f(x_k, u_k)
    \label{eq:residual_hypereuler}
\end{aligned}
\end{equation}

where $\Phi$ is an accurate solver and $\epsilon$ is the time stepping.

The hypernet $g_w$ should have as input data important features for predicting future trajectories: we can use the current state $x$, its derivative $f$ and the control input $u$. The training will be carried out by minimizing a loss function like:

\begin{equation}
    \ell = \frac{1}{K} \sum_{k=0}^{K-1} \left|\left|{ R \left( t_k,  x(t_k), x(t_{k+1}) \right) - g_\omega\left(x(t_k), f(t_k), u(t_{k}) \right) }\right|\right|_2
    \label{eq:loss_hypersolver}
\end{equation}

```python
# We consider the controller fixed during each solver step
class RandConstController(nn.Module):
    def __init__(self):
        super().__init__()
        self.u0 = torch.Tensor(1024, 1).uniform_(-10,10).to(device)

    def forward(self, t, x):
        return self.u0

# Save previously learned controller
u_no_hypersolver = sys.u
sys.u = RandConstController() # modify controller for training
```

```python
class VanillaHyperNet(nn.Module):
    """Simple hypernetwork for controlled systems
    Input: current x, f and u from the controlled system
    Output: p-th order residuals"""
    def __init__(self, net):
        super().__init__()
        self.net = net

    def forward(self, t, x):
        xfu = torch.cat([x, sys.cur_f, sys.cur_u], -1)
        return self.net(xfu)

net = nn.Sequential(nn.Linear(5, 64), nn.Softplus(), nn.Linear(64, 64), nn.Softplus(), nn.Linear(64, 2))
hypersolver = HyperEuler(VanillaHyperNet(net))
model = nn.DataParallel(hypersolver, device_ids=[1]) # feel free to change here according to your setup and GPU available.
model = model.to(device)
```

## Training the Hypernetwork

Training the hypernetwork $g_w$ requires us to know which set its input values are contained in.
If we want the pendulum with initial states in $\pm \pi$, we can sample the input states for the hypernetwork from a distribution which includes possible states which the system can visit (i.e., the hypernetwork should be able to predict outside of the initial boundaries in a certain range). For example, we can draw the states $x$ from a Uniform distribution:
$$x = \begin{bmatrix} q\\ p\\ \end{bmatrix} \sim Uniform(-2\pi, 2\pi) \times Uniform(-2\pi, 2\pi)$$
Similarly, the controller values can be sampled from a distribution including all the possible controller values, such as
$$u \sim Uniform(-20, 20)$$

```python
# Initial distribution
x0 = 2*π # limit of the state distribution (in rads and rads/second)
init_dist = torch.distributions.Uniform(torch.Tensor([-x0, -x0]), torch.Tensor([x0, x0]))

base_solver = Euler()
# Time span
t0, tf = 0, 2 # initial and final time for controlling the system
steps = 20 + 1 # so we have a time step of 0.1s
t_span = torch.linspace(t0, tf, steps).to(device)
dt = (t_span[1] - t_span[0]).detach().cpu().item()
```

```python
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_func = nn.MSELoss()
epochs = 10000
bs = 2048
hypernet = model.module.hypernet
span = torch.linspace(0, dt, 2)
losses = []

for i in range(epochs):
    # Sample random intial states and controllers
    x0 = init_dist.sample((bs,)).to(device)
    sys.u.u0 = torch.Tensor(bs, 1).uniform_(-20, 20).to(device)

     # Compute residuals
    _, sol_gt = odeint(sys, x0, span, solver='tsit5', atol=1e-8, rtol=1e-8)[-1]
    sol = x0 + sys(0., x0)*dt
    residuals = (sol_gt - sol) / dt**2
    residuals_hypersolver = hypernet(0., x0)
    loss = loss_func(residuals, residuals_hypersolver)

    # Optimization step
    loss.backward(); opt.step(); opt.zero_grad()
    print(f'Step: {i}, Residual loss: {loss:.3f}', end='\r')
    losses.append(loss.detach().cpu().item())
```

    Step: 9999, Residual loss: 0.0254

```python
# Sample and plot some trajectories
x0 = π # limit of the state distribution (in rads and rads/second)
init_dist = torch.distributions.Uniform(torch.Tensor([-x0, -x0]), torch.Tensor([x0, x0]))
x0 = init_dist.sample((10000,)).to(device)
sys = ControlledPendulum(RandConstController()).to(device)
sys.u.u0 = torch.Tensor(10000, 1).uniform_(0,0).to(device)
t_span = torch.linspace(0, 5, 101).to(device)

_, traj_euler = odeint(sys, x0, t_span, solver='euler')
_, traj_hyper = odeint(sys, x0, t_span, solver=hypersolver)
_, traj_gt = odeint(sys, x0, t_span, solver='tsit5', atol=1e-7, rtol=1e-7)
traj_euler = traj_euler.detach().cpu(); traj_hyper = traj_hyper.detach().cpu(); traj_gt = traj_gt.detach().cpu()
t_span = t_span.cpu()

fig, axs = plt.subplots(2, 5, figsize=(16,6))

for i in range(5):
    for j in range(2):
        axs[j, i].plot(t_span, traj_euler[:,i,j], 'r')
        axs[j, i].plot(t_span, traj_hyper[:,i,j], 'k')
        axs[j, i].plot(t_span, traj_gt[:,i,j], 'b-.')
        axs[j, i].legend(['Euler', 'HyperEuler', 'Ground Truth'])
        axs[j, i].label_outer()
    axs[0, i].set_ylabel(r'Positions $q$'); axs[1, i].set_ylabel(r'Momenta $p$')

fig.suptitle('Uncontrolled Trajectories')
```

```python
# Error analysis
def smape(yhat, y):
    return torch.abs(yhat - y) / (torch.abs(yhat) + torch.abs(y)) / 2

err_euler, err_hyper = smape(traj_gt, traj_euler).detach().cpu(), smape(traj_gt, traj_hyper).detach().cpu()
mean_euler, std_euler = err_euler.mean(1), err_euler.std(1)
mean_hyper, std_hyper = err_hyper.mean(1), err_hyper.std(1)

fig = plt.figure(figsize=(16,5))
ax = fig.add_subplot(1,2,1)
ax.plot(t_span, mean_euler[:,0], 'r', label='Euler')
ax.plot(t_span, mean_hyper[:,0], 'b', label='HyperEuler')
ax.fill_between(t_span, mean_euler[:,0]-std_euler[:,0], mean_euler[:,0]+std_euler[:,0], alpha=.2, color='r')
ax.fill_between(t_span, mean_hyper[:,0]-std_hyper[:,0], mean_hyper[:,0]+std_hyper[:,0], alpha=.2, color='b')
ax.legend()
ax.set_title('Positions')

ax = fig.add_subplot(1,2,2)
ax.plot(t_span, mean_euler[:,1], 'r', label='Euler')
ax.plot(t_span, mean_hyper[:,1], 'b', label='HyperEuler')
ax.fill_between(t_span, mean_euler[:,1]-std_euler[:,1], mean_euler[:,1]+std_euler[:,1], alpha=.2, color='r')
ax.fill_between(t_span, mean_hyper[:,1]-std_hyper[:,1], mean_hyper[:,1]+std_hyper[:,1], alpha=.2, color='b')
ax.legend()
ax.set_title('Momenta')
fig.suptitle('SMAPE Propagation')
```

Nice! We managed obtain a more accurate solver with respect to the Euler baseline. Now, let's see how this performs with optimal control.

## Optimal Control with Hypersolvers
We can use the trained hypernetwork to generate trajectories of the pendulum: then, we cast the same optimal control problem defined at the beginning of the notebook and train the neural controller.

```python
# Reinstantiate controller
model = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 1)).to(device)
u = NeuralController(model)
for p in u.model[-1].parameters(): torch.nn.init.zeros_(p)
sys.u = u

# Time span
t0, tf = 0, 2 # initial and final time for controlling the system
steps = 20 + 1 # so we have a time step of 0.1s
t_span = torch.linspace(t0, tf, steps).to(device)

# Initial distribution
x0 = π # limit of the state distribution (in rads and rads/second)
init_dist = torch.distributions.Uniform(torch.Tensor([-x0, -x0]), torch.Tensor([x0, x0]))
```

```python
# Hyperparameters
lr = 3e-3
epochs = 500
bs = 1024
opt = torch.optim.Adam(u.parameters(), lr=lr)

# Training loop
t0 = time.time(); losses=[]
for e in range(epochs):
    x0 = init_dist.sample((bs,)).to(device)
    _, trajectory = odeint(sys, x0, t_span, solver=hypersolver, atol=1e-7, rtol=1e-7)
    loss = cost_func(trajectory); losses.append(loss.detach().cpu().item())
    loss.backward(); opt.step(); opt.zero_grad()
    print('Loss {:.4f} , epoch {}'.format(loss.item(), e), end='\r')
timing = time.time() - t0; print('\nTraining time: {:.4f} s'.format(timing))
```

    Loss 1.1300 , epoch 499
    Training time: 21.2143 s

```python
# Comparing inference times between Euler, HyperEuler, RK4 (fixed-step, 4th order solver) and the base solver
t_span_test = torch.linspace(0, 100, 1001)
x0 = init_dist.sample((bs,)).to(device)
t0 = time.time(); odeint(sys, x0, t_span_test, solver='euler'); timing_euler = time.time() - t0
t0 = time.time(); odeint(sys, x0, t_span_test, solver=hypersolver); timing_hyper = time.time() - t0
t0 = time.time(); odeint(sys, x0, t_span_test, solver='rk4'); timing_rk4 = time.time() - t0
t0 = time.time(); odeint(sys, x0, t_span_test, solver='tsit5', atol=1e-7, rtol=1e-7); timing_tsit5 = time.time() - t0
print('Inference times\nEuler: {:.4f} s\nHyperEuler: {:.4f} s\nRunge-Kutta-4: {:.4f} s\nTsitouras: {:.4f} s'.format(timing_euler, timing_hyper, timing_rk4, timing_tsit5))
```

    Inference times
    Euler: 0.2584 s
    HyperEuler: 0.4956 s
    Runge-Kutta-4: 1.0962 s
    Tsitouras: 2.6975 s

```python
# Testing the controller on the real system
x0 = init_dist.sample((100,)).to(device)
t0, tf, steps = 0, 2, 20*10 + 1 # nominal trajectory
t_span_fine = torch.linspace(t0, tf, steps).to(device)
_, traj = odeint(sys, x0, t_span_fine, solver='tsit5', atol=1e-7, rtol=1e-7)
t_span_fine = t_span_fine.detach().cpu(); traj = traj.detach().cpu()

# Plotting
fig, ax = plt.subplots(1, 1, figsize=(12,6))
for i in range(len(x0)):
    ax.plot(t_span_fine, traj[:,i,0], 'k-.', alpha=.3, label='p' if i==0 else None)
    ax.plot(t_span_fine, traj[:,i,1], 'b', alpha=.3, label='q' if i==0 else None)
ax.legend()
ax.set_title('Controlled trajectories')
ax.set_xlabel(r"$t~[s]$")
```

Great 🎉

Training the controller with `HyperEuler` resulted in a working controller stabilizing the system in the given time; more importantly, while in the first part we used a high accuracy solver to obtain similar results, which took us 220 seconds, training with the same number of epochs and all the same hyperparameters but using the hypersolver took less than 20 seconds : more than $10\times$ faster!


# Implementation of linear discriminant contrast (LDC)
After the RSA workshop this year (2018), I wanted to figure out how the linear discriminant contrast (LDC) works and how it can be implemented in Python. The notebook is the result of that; it's a bit messy, but you'll get the idea hopefully.

I've based the content of this notebook on [this blog](http://www.johancarlin.com/pilab-tutorial-2-linear-discriminant-contrast.html) and [this presentation](http://www.mrc-cbu.cam.ac.uk/wp-content/uploads/2014/10/rsaworkshop_alex.pdf). I also included the calculation of the LD-t, but I'm not sure that's correct.

```python
import numpy as np
from sklearn.discriminant_analysis import _cov
from nistats.hemodynamic_models import spm_hrf
import matplotlib.pyplot as plt
import seaborn as sns
import scipy

%matplotlib inline
```

## Generating (relatively) realistic fMRI data
First, let's generate some data. We'll sample onsets of stimuli from 4 conditions ($P$), each with 5 trials/instances ($I$) across a run lasting 300 seconds (5 minutes; $dur$).

```python
P = 4  # conditions
I = 5   # instances
dur = 400  # duration (in seconds)

X = np.zeros((dur, P))
onsets = np.random.choice(np.arange(dur - 20), size=I*P, replace=False)

i = 0
for p in range(P):

    for inst in range(I):
        X[onsets[i], p] = 1
        i += 1

plt.figure(figsize=(15, 3))
plt.plot(X)
plt.title('Onsets conditions', fontsize=20)
plt.xlabel('Time (seconds)')
plt.ylabel('Amplitude (A.U.)')
plt.xlim(0, dur)
plt.ylim(0, 1.2)
plt.legend([f'Condition {p+1}' for p in range(P)], fontsize=10)
sns.despine()
plt.show()
```

Now, we'll convolve it with an HRF. We'll oversample the HRF by a factor 20. Given that our TR is 2, that means we'll define our HRF at the resolution of 10 Hz.

```python
TR = 2

oversampling = 20
hrf = spm_hrf(TR, oversampling=oversampling, time_length=32)

# Scale HRF such that range is [min(hrf), 1]
hrf = hrf / np.max(hrf)

plt.figure(figsize=(10, 3))
plt.plot(np.arange(0, 32, (1/(oversampling/TR))), hrf)
plt.axhline(0, ls='--', c='k', lw=1)
plt.ylim(-0.2, 1.2)
plt.xlabel('Time (seconds)')
plt.ylabel('Amplitude (A.U.)')
plt.title("HRF", fontsize=20)
sns.despine()
plt.show()
```

To actually convolve the design, we need to upsample our onsets as well (from 1 Hz to 10 Hz):

```python
plt.figure(figsize=(15, 3))

us_factor = int(oversampling / TR)
Xus = np.repeat(X, axis=0, repeats=us_factor)
Xconv = np.hstack([np.convolve(Xus[:, i], hrf, 'full')[:dur*us_factor, np.newaxis] for i in range(P)])

# Downsample again to seconds (1 Hz)
Xconv = Xconv[::us_factor, :]
plt.plot(Xconv)
plt.title('Convolved regressors', fontsize=20)
plt.xlabel('Time (seconds)')
plt.ylabel('Amplitude (A.U.)')
plt.xlim(0, dur)

plt.tight_layout()
sns.despine()
plt.show()
```

We're almost there: we still need to downsample our design-matrix to the resolution of our scans, i.e., the resolution of the TR (0.5 Hz) and add an intercept:

```python
X_final = Xconv[::2, :]
X_final = np.c_[np.ones(X_final.shape[0]), X_final]
print(X_final.shape)
```

    (200, 5)

Alright, now, let's generate the signal ($y$) as follows:

\begin{align}
y = X\beta + \epsilon \sim N(0, V)
\end{align}

First, let's generate the noise (for $K = 4$ voxels):

```python
K = 4

# Somewhat realistic covariance matrix of signal
V = 1 / (1 + scipy.linalg.toeplitz(np.arange(dur/TR), np.arange(dur/TR)))
noise = np.random.multivariate_normal(np.zeros(int(dur/TR)), V, size=K).T
print(noise.shape)

plt.imshow(V, vmax=0.5)
plt.title('Autocorrelation matrix')
plt.axis('off')
plt.show()
```

Now, let's generate the true parameters ($\beta$):

\begin{align}
\beta \sim N(0, 0.15)
\end{align}

```python
# Including intercept
true_params = np.random.normal(0, 0.15, (P+1, K))
```

Now we can create the signals!

```python
y = X_final.dot(true_params) + noise

plt.figure(figsize=(15, 15))

for k in range(K):
    plt.subplot(K, 1, (k+1))
    plt.plot(y[:, k], lw=2)
    plt.plot(X_final.dot(true_params[:, k]), ls='--')
    plt.xlim(0, 200)
    plt.ylim(-5, 5)
    plt.axhline(0, ls='-.', lw=1)
    plt.title(f"Voxel {k+1}", fontsize=15)

sns.despine()
plt.tight_layout()
plt.show()
```

## Calculating the linear discriminant contrast
The LDC a.k.a. Mahalanobis distance. Note, this is (not yet) the cross-validated version ("crossnobis"); this is explained later.

First, we have to compute the parameters from the linear model:

\begin{align}
\beta = (X^{T}X)^{-1}X^{T}y
\end{align}

```python
est_params = np.linalg.lstsq(X_final, y, rcond=None)[0]
```

Let's compare it to the true parameters:

```python
plt.figure(figsize=(10, 3))
plt.scatter(est_params, true_params)
plt.plot([-0.5, 0.5], [-0.5, 0.5], ls="--", c=".3")
plt.xlim(-0.5, 0.5)
plt.ylim(-0.5, 0.5)
plt.xlabel('Estimated parameters', fontsize=15)
plt.ylabel('True parameters', fontsize=15)
sns.despine()
plt.show()
```

Not bad. ANyway, now we need to calculate the covariance matrix of the features (voxels), $\Omega$. (Here, we use $\Omega$ instead of the more common notation of $V$, because we've used that to denote the covariance matrix across time!) Usually, this covariance matrix is calculated on the residuals of the model ($y - \hat{y}$):

```python
residuals = y - X_final.dot(est_params)
omega = np.cov(residuals.T)

plt.title("Covariance matrix $\Omega$")
plt.imshow(omega)
plt.axis('off')
plt.show()
```

The next step in LDA is calculate the difference in parameters between two conditions, let' say condition 1 and 2:

```python
diff_c1c2 = est_params[1, :] - est_params[2, :]
print(diff_c1c2)
```

    [-0.13557863 -0.16938876  0.17502224  0.15755565]

A different way to do this is using contrast-vectors. For example, with contrast vector $c = [0, 1, -1, 0, 0]$, we can calculate the same difference of parameters as above (the first 0 in $c$ refers to the intercept):

```python
c = np.array([0, 1, -1, 0, 0])
diff_c1c2 = c.dot(est_params)
diff_c1c2
```

    array([-0.13557863, -0.16938876,  0.17502224,  0.15755565])

Now, the last thing we need to do is calculate the actual LDC/Mahalanobis distance:

\begin{align}
ldc = (\beta_{1} - \beta_{2})\Omega^{-1}(\beta_{1} - \beta_{2})^{T}
\end{align}

```python
diff_c1c2.dot(np.linalg.inv(omega)).dot(diff_c1c2).T
```

    0.10494452961363823

## Cross-validation of ldc
Also known as cross-validated Mahalanobis distance or crossnobis.

```python
def generate_data(P, I, dur, TR, K, true_params):
    """ Generates data.

    P : int
        Number of conditions
    I : int
        Number of instances per condition
    dur : int
        Duration of run in seconds
    TR : int
        Time to repetition (for simplicity, can only be an integer)
    K : int
        Number of voxels
    true : list
        True parameters (must be of length P + 1)

    Returns
    -------
    y : numpy array
        Array reflecting signal of K voxels of length dur/TR
    X_final : numpy array
        Array reflecting final (convolved) design of length dur/TR by K
    """
    X = np.zeros((dur, P))
    onsets = np.random.choice(np.arange(dur - 20), size=I*P, replace=False)

    i = 0
    for p in range(P):

        for inst in range(I):
            X[onsets[i], p] = 1
            i += 1

    oversampling = 20
    hrf = spm_hrf(TR, oversampling=oversampling, time_length=32)

    # Scale HRF such that range is [min(hrf), 1]
    hrf = hrf / np.max(hrf)

    us_factor = int(oversampling / TR)
    Xus = np.repeat(X, axis=0, repeats=us_factor)
    Xconv = np.hstack([np.convolve(Xus[:, i], hrf, 'full')[:dur*us_factor, np.newaxis] for i in range(P)])

    # Downsample again to seconds (1 Hz)
    Xconv = Xconv[::us_factor, :]

    X_final = Xconv[::TR, :]
    X_final = np.c_[np.ones(X_final.shape[0]), X_final]

    # Somewhat realistic covariance matrix of signal
    V = 1 / (1 + scipy.linalg.toeplitz(np.arange(dur/TR), np.arange(dur/TR)))
    noise = np.random.multivariate_normal(np.zeros(int(dur/TR)), V, size=K).T

    true_params = np.array(true_params)[:, np.newaxis]
    y = X_final.dot(true_params) + noise
    return y, X_final
```

We'll simulate two runs:

```python
#true_params = np.random.normal(0, 0.15, 3)
true_params = np.array([0, 0, 0])
y1, X1 = generate_data(P=2, I=6, dur=300, TR=2, K=10, true_params=true_params)
y2, X2 = generate_data(P=2, I=6, dur=300, TR=2, K=10, true_params=true_params)
```

So, the idea is that for two runs, let's say R1 and R2, the cross-validated Mahalanobis distance between condition A and B is:

\begin{align}
\delta_{Mahalanobis}(A, B) = (\beta_{A} - \beta_{B})_{R1}\Omega^{-1}_{R1}(\beta_{A} - \beta_{B})_{R2}^{T}
\end{align}

While we're at it, let's also apply shrinkage to our covariance matrix of the residuals ($\Omega$), which is often recommended. We'll do that using the `cov` function from the scikit-learn package:

```python
from sklearn.discriminant_analysis import _cov

params_R1 = np.linalg.lstsq(X1, y1, rcond=None)[0]
residuals_R1 = y1 - X1.dot(params_R1)
omega_R1 = _cov(residuals_R1, shrinkage='auto')

params_R2 = np.linalg.lstsq(X2, y2, rcond=None)[0]

cv_ldc = (params_R1[1, :] - params_R2[2, :]).dot(omega_R1).dot((params_R1[1, :] - params_R2[2, :]).T)
cv_ldc
```

    0.026289009884526548

So, supposedly you can also calculate a t-value of your LDC statistic as follows:

\begin{align}
LD_{t} = \frac{ldc}{\sqrt(\mathbf{w}\Omega_{R2}\mathbf{w}^{T}) * c(\mathbf{X}^{T}\mathbf{X})^{-1}c^{T}})
\end{align}

Here, $\mathbf{w}$ is the weight vector for optimal separation of the two classes:

\begin{align}
\mathbf{w} = (\beta_{A} - \beta_{B})\Omega^{-1}
\end{align}

and $c$ is the contrast-vector defining which conditions are compared. So, we can calculate this LD-t as follows:

```python
residuals_R2 = y2 - X2.dot(params_R2)
omega_R2 = _cov(residuals_R2, shrinkage='auto')
w = (params_R1[1, :] - params_R2[2, :]).dot(omega_R1)
c = np.array([0, 1, -1])
ldT = cv_ldc / np.sqrt(w.dot(omega_R2).dot(w.T) * c.dot(np.linalg.inv(X1.T.dot(X1))).dot(c.T))
ldT # this is probably not correct
```

    4.203573427129181


# ***Introduction to Radar Using Python and MATLAB***
## Andy Harrison - Copyright (C) 2019 Artech House
<br/>

# Shnidman Approximation
***

Referring to Section 6.3.6, Shnidman  developed a set of equations, based on empirical data, for calculating the required signal-to-noise ratio for a specified probability of detection and probability of false alarm.  These equations are applicable to all of the Swerling target models for either a single pulse or the noncoherent integration of $N$ pulses. Shnidman's equation proceeds as follows.  First, the parameter $K$ is chosen based on the Swerling target type, then the parameter $\alpha$ is chosen based on the number of pulses as (Equation 6.45)

$$
    K = \begin{cases}
    \infty &\text{Swerling 0}\\
    1 &\text{Swerling I}\\
    N &\text{Swerling II}\\
    2 &\text{Swerling III}\\
    2N &\text{Swerling IV}
    \end{cases}, \hspace{0.4in}
    \alpha = \begin{cases}
    0 &N<40\\
    \frac{1}{4} &N\ge40
    \end{cases}.
$$

Then the two parameters $\eta$ and $X$ are computed from (Equation 6.46)

$$
    \eta = \sqrt{-0.8\ln(4 P_{fa}(1-P_{fa}))} + \text{sign}(P_d - 0.5)\sqrt{-0.8\ln(4 P_d(1-P_d))}
$$

and

$$
    X = \eta\Bigg[ \eta + 2\sqrt{\frac{N}{2} + \Big( \alpha - \frac{1}{4}\Big)} \,\Bigg].
$$

The following constants are now computed (Equation 6.48 - 6.51)

\begin{align}
    C_1 &= \left( \big[  (17.7006 P_d - 18.4496) P_d  + 14.5339\big]P_d - 3.525 \right) / K, \\ \nonumber \\
    C_2 &= \frac{1}{K}\bigg( \exp\big( 27.31 P_D - 25.14\big) + \big( P_d-0.8\big) \nonumber \\ &\hspace{1.5in}\times \left[ 0.7\ln\left( \frac{10^{-5}}{P_{fa}}\right) + \frac{(2N-20)}{80}\right]\bigg), \\ \nonumber \\
    C_{dB} &= \begin{cases} C_1 &\text{for}\hspace{5pt}0.1\le P_d \le 0.872,\\[5pt] C_1 + C_2 &\text{for}\hspace{5pt}0.872\le P_d \le 0.99,\end{cases} \\ \nonumber \\
    C &= 10^{(C_{dB}/10)}.\nonumber
\end{align}

Finally, the signal-to-noise ratio is (Equation 6.52)

$$
    {SNR} = 10\log_{10}\left(\frac{C\cdot X}{N}\right) \hspace{0.5in} \text{(dB)}.
$$

Shnidman's equation has been shown to be accurate to within $0.5$ dB within the following bounds

$$
    0.1 \le \,P_d \le 0.99, \\ \nonumber \\
    10^{-9} \le \,P_{fa} \le 10^{-3}, \\ \nonumber \\
    1 \le \,N \le 10.
$$

Shnidman's equation is valid for all five Swerling target types.  Shnidman's equation provides accurate results for first cut type of radar system calculations.
***

Begin by getting the library path

```python
import lib_path
```

Set the start and end probability of detection

```python
pd_start = 0.8

pd_end = 0.99
```

Create the prbability of detection array using the `linspace` routine from `scipy`

```python
from numpy import linspace

pd_all = linspace(pd_start, pd_end, 200)
```

Set the probability of false alarm and the number_of_pulses

```python
pfa = 1e-6

number_of_pulses = 10
```

Set the target type (Swerling 0 - Swerling 4)

```python
target_type = 'Swerling 2'
```

Calculate the error in the Shnidman approximation of signal to noise using the `single_pulse_snr` and `signal_to_noise` routines

```python
from Libs.detection.shnidman import signal_to_noise
from Libs.detection.non_coherent_integration import single_pulse_snr

from numpy import log10

error = [10.0 * log10(single_pulse_snr(p, pfa, number_of_pulses, target_type)) - signal_to_noise(p, pfa, number_of_pulses, target_type) for p in pd_all]
```

Display the error in the Shnidman approximation

```python
from matplotlib import pyplot as plt

# Set the figure size

plt.rcParams["figure.figsize"] = (15, 10)

# Display the results

plt.plot(pd_all, error, '')

# Set the plot title and labels

plt.title('Shnidman\'s Approximation', size=14)

plt.xlabel('Probability of Detection', size=12)

plt.ylabel('Signal to Noise Error (dB)', size=12)

plt.ylim(-1, 1)

# Set the tick label size

plt.tick_params(labelsize=12)

# Turn on the grid

plt.grid(linestyle=':', linewidth=0.5)
```

```python

```


```python
from sympy import *
init_printing()
```

```python
def skew(l):
    l1, l2, l3 = l
    return Matrix([
        [0, -l3, l2],
        [l3, 0, -l1],
        [-l2, l1, 0]
    ])
```

```python
# define state variables
x, y, z, eta0, eps1, eps2, eps3, u, v, w, p, q, r = symbols('x y z et0 eps1 eps2 eps3 u v w p q r', real=True)
s = Matrix([x, y, z, eta0, eps1, eps2, eps3, u, v, w, p, q, r])

# position and orientation
eta = Matrix([x, y, z, eta0, eps1, eps2, eps3])
nu = Matrix([u, v, w, p, q, r])

# centre of gravity
xg, yg, zg = symbols('xg yg zg', real=True)
rg = Matrix([xg, yg, zg])

# centre of bouyancy
xb, yb, zb = symbols('xb yb zb', real=True)
rb = Matrix([xb, yb, zb])

# center of pressure
xcp, ycp, zcp = symbols('xcp ycp zcp', real=True)
rcp = Matrix([xcp, ycp, zcp])

# mass matrix
m = symbols('m', real=True, positive=True)
Ixx, Iyy, Izz = symbols('Ixx Iyy Izz')
I0 = diag(Ixx, Iyy, Izz)
M = BlockMatrix([
    [m*eye(3), -m*skew(rg)],
    [m*skew(rg), I0]
])
M = Matrix(M)
# M = simplify(M)

# Coriolis and centripetal matrix
nu1 = Matrix([u, v, w])
nu2 = Matrix([p, q, r])
crb = BlockMatrix([
    [zeros(3), -m*skew(nu1)-m*skew(nu2)*skew(rg)],
    [-m*skew(nu1)+m*skew(rg)*skew(nu2), -skew(I0*nu2)]
])
crb = Matrix(crb)
# crb = simplify(crb)

# damping matrix
Xuu, Yvv, Zww, Kpp, Mqq, Nrr = symbols(
    'Xuu Yvv Zww Kpp Mqq Nrr', real=True
)
D = Matrix([
    [Xuu*abs(u), 0, 0, 0, 0, 0],
    [0, Yvv*abs(v), 0, 0, 0, 0],
    [0, 0, Zww*abs(w), 0, 0, 0],
    [0, -zcp*Yvv*abs(v), ycp*Zww*abs(w), Kpp*abs(p), 0, 0],
    [zcp*Xuu*abs(u), 0, -xcp*Zww*abs(w), 0, Mqq*abs(q), 0],
    [-ycp*Xuu*abs(u), xcp*Yvv*abs(v), 0, 0, 0, Nrr*abs(r)]
])
# D = simplify(D)

# rotational transform between body and NED quaternions
Tq = Rational(1,2)*Matrix([
    [-eps1, -eps2, -eps3],
    [eta0, -eps3, eps2],
    [eps3, eta0, -eps1],
    [-eps2, eps1, eta0]
])
# Tq = simplify(Tq)

Rq = Matrix([
    [1-2*(eps2**2+eps3**2), 2*(eps1*eps2-eps3*eta0), 2*(eps1*eps3+eps2*eta0)],
    [2*(eps1*eps2+eps3*eta0), 1-2*(eps1**2+eps3**2), 2*(eps2*eps3-eps1*eta0)],
    [2*(eps1*eps3-eps2*eta0), 2*(eps2*eps3+eps1*eta0), 1-2*(eps1**2+eps2**2)]
])
Jeta = BlockMatrix([
    [Rq, zeros(3)],
    [zeros(4,3), Tq]
])
Jeta = Matrix(Jeta)
# Jeta = simplify(Jeta)

# bouyancy in quaternions
W, B = symbols('W B', real=True)
fg = Matrix([0, 0, W])
fb = Matrix([0, 0, -B])
Rqinv = Rq.inv()
geta = Matrix([
    Rqinv*(fg+fb),
    skew(rg)*Rqinv*fg + skew(rb)*Rqinv*fb
])
# geta = simplify(geta)
```

```python
print(cse(Jeta))
```

    ([(x0, -2*eps2**2), (x1, 1 - 2*eps3**2), (x2, 2*eps2), (x3, eps1*x2), (x4, 2*eps3), (x5, et0*x4), (x6, eps1*x4), (x7, et0*x2), (x8, -2*eps1**2), (x9, 2*eps1*et0), (x10, eps2*x4), (x11, eps1/2), (x12, -x11), (x13, eps2/2), (x14, -x13), (x15, eps3/2), (x16, -x15), (x17, et0/2)], [Matrix([
    [x0 + x1,  x3 - x5,     x6 + x7,   0,   0,   0],
    [x3 + x5,  x1 + x8,    x10 - x9,   0,   0,   0],
    [x6 - x7, x10 + x9, x0 + x8 + 1,   0,   0,   0],
    [      0,        0,           0, x12, x14, x16],
    [      0,        0,           0, x17, x16, x13],
    [      0,        0,           0, x15, x17, x12],
    [      0,        0,           0, x14, x11, x17]])])

```python
# thrust model
Kt0, Kt1 = symbols('Kt0 Kt1', real=True)
Kt = Matrix([Kt0, Kt1])
Qt0, Qt1 = symbols('Qt0 Qt1', real=True)
Qt = Matrix([Qt0, Qt1])

# control inputs
rpm0, rpm1 = symbols('rpm0 rpm1', real=True)
rpm = Matrix([rpm0, rpm1])
de, dr = symbols('de dr', real=True)
control_vector = Matrix([rpm0, rpm1, de, dr])

# control force vector
Ft = Kt.dot(rpm)
Mt = Qt.dot(rpm)

# coefficient for each element in cost function
tauc = Matrix([
    Ft*cos(de)*cos(dr),
    -Ft*sin(dr),
    Ft*sin(de)*cos(dr),
    Mt*cos(de)*cos(dr),
    -Mt*sin(dr),
    Mt*sin(de)*cos(dr)
])
```

```python
etadot = Jeta*nu
nudot = M.inv()*(tauc - (crb + D)*nu - geta)
```

```python
sdot = Matrix([
   etadot,
   nudot
])
```

```python
print(list(set(sdot.free_symbols) - set(s.free_symbols) - set(control_vector.free_symbols)))
```

```python
# Lagrangian
alpha = symbols('\\alpha', real=True, positive=True)
L = alpha + (1-alpha)*tauc.norm()
```

```python
l = Matrix([symbols('lambda_{}'.format(var)) for var in s])
```

```python
H = l.dot(sdot) + L
```

```python
eq = H.diff(control_vector)
```

```python
sol = solve(eq, control_vector)
```

```python
L.diff(control_vector)
```

```python
tauc.transpose()*tauc
```

```python

```


# Numerical integration: introduction

This notebook is based on Chapters 7 and 10 of

> Süli, Endre and Mayers, David F. _An introduction to numerical analysis_. Cambridge University Press, Cambridge, 2003.
<https://doi.org/10.1017/CBO9780511801181> (ebook in [Helka](https://helka.helsinki.fi/permalink/358UOH_INST/1h3k2rg/alma9926836783506253))

We know how to evaluate some simple integrals of elementary functions by pen and paper but, in general, symbolic evaluation of integrals is a complicated problem, see for example the Wikipedia article of on the [Risch algorithm](https://en.wikipedia.org/wiki/Risch_algorithm). Here we consider numerical evaluation of integrals, a much simpler problem.

```python
import sympy as sp
from sympy.abc import x
from IPython.display import Math

def display_risch(f):
    int_tex = r'\int ' + sp.latex(f)  + ' dx'
    F = sp.integrate(f, x, risch=True)
    if not isinstance(F, sp.integrals.Integral):
        # The integral is elementary and we display the antiderivative
        display(Math(int_tex + '=' + sp.latex(F)))
    elif isinstance(F, sp.integrals.risch.NonElementaryIntegral):
        # The integral is proven to be nonelementary by the algorithm,
        # meaning the integral cannot be represented using a combination of
        # exponentials, logarithms, trig functions, powers, rational functions,
        # algebraic functions, and function composition
        display(Math(int_tex + r'\,\text{ is not an elementary function}'))
    else:
        # The version implemented in SymPy only supports a small subset of
        # the full Risch algorithm, no conclusion can be drawn from this case
        display(Math(int_tex + r'\,\text{ can not be handled by the algorithm}'))
```

```python
# Example from pp. 82-84 (Examples 2 and 3) of
#   Grozin, A. Introduction to Mathematica for Physicists. Springer, 2014.
# shows that a small change in the integrand makes it impossible to express
# an integral by using elementary functions
display_risch(sp.log(x)/x)
display_risch(sp.log(x)/(x + 1))
# Example in the Wikipedia article on the Risch algorithm
display_risch(x / sp.sqrt(x**4 + 10*x**2 - 96*x - 72))
```

# Newton-Cotes formulae

The _Newton-Cotes formula_ of order $n = 1, 2, \dots$ is obtained by letting $p \in \mathbb P_n$ be the Lagrange interpolation polynomial of a function $f$ at $a = x_0 < x_1 < \dots < x_n = b$ on an interval $[a,b]$, and approximating

$$
\int_a^b f(x) dx \approx \int_a^b p(x) dx.
$$

Recall that

$$
p(x) = \sum_{k=0}^n f(x_k) L_k(x),
$$

where $L_k \in \mathbb P_n$ are the polynomials satisfying $L_k(x_i) = \delta_{ki}$.
Hence

$$
\int_a^b p(x) dx = \sum_{k=0}^n w_k f(x_k), \qquad w_k = \int_a^b L_k(x) dx.
$$

```python
import numpy as np
import scipy.interpolate as interp

def nc_weights_demo(a, b, n):
    '''Compute the quadrature weights for Newton-Cotes formula of order n on [a, b]'''
    N = n+1
    xs = np.linspace(a, b, N) # uniform grid
    ws = np.zeros(N)
    for k in range(N):
        ys = np.zeros(N)
        ys[k] = 1
        l = interp.lagrange(xs, ys) # L_k
        L = np.polyint(l)
        ws[k] = L(b) - L(a)
    return ws, xs

def integrate(ws, xs, f):
    '''Integrate f using the quadrature weights ws and points xs'''
    return np.sum(ws * f(xs))
```

```python
a, b = 0, 1
def test_integrate(ws, xs, expr):
    # Convert a symbolic expression with x as a free variable to a function
    f = sp.lambdify(x, expr)
    n = len(ws)-1
    display(Math(r'\text{' +
        f'Integrate gives {integrate(ws, xs, f):.2f} '
        r' for }\int' + f'_{a}^{b}' + sp.latex(expr)  +
        f' dx = {sp.integrate(expr, (x, a, b))}'))

n = 1
ws_demo, xs = nc_weights_demo(a, b, n)
test_integrate(ws_demo, xs, 1)
test_integrate(ws_demo, xs, x)
test_integrate(ws_demo, xs, x**2)
```

```python
n = 2
ws_demo, xs = nc_weights_demo(a, b, n)
fs = [x**n for n in range(5)]
for f in fs:
    test_integrate(ws_demo, xs, f)

from scipy import integrate as integ
ws_scipy, _ = integ.newton_cotes(n, equal=1)
ws_scipy *= (b-a)/n
print('Using SciPy')
test_integrate(ws_scipy, xs, fs[-1])
```

## Theorem: error in Newton-Cotes
>
> Let $n=1,2,\dots$ and $f \in C^{n+1}(a,b)$. Let $w_i, x_i$, $i=0,\dots,n$, be the quadrature weights and points in the Newton-Cotes formula of order $n$ on $[a,b]$. Then
>
>\begin{align*}
&\left|\int_a^b f(x) dx - \sum_{k=0}^n w_k f(x_k) \right|
\\&\quad\le
\frac{\|f^{(n+1)}\|_\infty}{(n+1)!} \int_a^b |x-x_0|\dots|x-x_n| dx.
\end{align*}

For a proof, see Theorem 7.1 in [the book](#thebook). The Newton-Cotes formula of order $n$ is exact (that is, it gives the exactly correct integral) for polynomials of degree $n$. When $n$ is even and the grid is uniform, it is exact also for polynomials of degree $n+1$, leading to a better estimate than that in the theorem. For the sharp estimate in the case $n=2$, see Theorem 7.2 in [the book](#thebook).

```python
def nc_weights_sym(n):
    '''Compute symbolically the weights for Newton-Cotes formula of order n'''
    from sympy.abc import a, b, x
    ws = []
    h = (b - a)/n
    xs = [sp.simplify(a + k*h) for k in range(n+1)]
    for k in range(len(xs)):
        L = 1
        for i in range(len(xs)):
            if i != k:
                L *= (x - xs[i]) / (xs[k] -xs[i])
        ws.append(sp.simplify(sp.integrate(L, (x, a, b))))
    return ws, xs

def nc_formula(n):
    '''Compute symbolic expression for Newton-Cotes formula of order n'''
    w, x = nc_weights_sym(n)
    f = sp.Function('f')
    nc_form = 0
    for k in range(len(w)):
        nc_form += w[k] * f(x[k])
    return sp.simplify(nc_form)
```

```python
print('Trapezium rule')
display(Math(sp.latex(nc_formula(1))))
print("Simpson's rule")
display(Math(sp.latex(nc_formula(2))))
```

Recall that the Lagrange interpolation polynomial of

$$
f(x) = \frac{1}{1+x^2}
$$

gives a poor approximation of $f$ on $[-5, 5]$ due to the Runge phenomenon.

Thus the Newton-Cotes formulae are expected to work poorly for

$$
\int_{-5}^5 f(x) dx.
$$

```python
def f(x):
    return 1 / (1 + x**2)
ns = range(2, 13, 2)
N = np.size(ns)
data = np.zeros(N)
for k in range(N):
    ws, xs = nc_weights_demo(-5, 5, ns[k])
    data[k] = integrate(ws, xs, f)

import pandas as pd
df = pd.DataFrame(data)
df.columns = ['$I_n$']
df.index = ns
df.index.name = 'n'
df.style.format('{:.5f}')
```

# Composite formulae

Analogously to splines, we will use a division to small subintervals.

Let $a = x_0 < x_1 < \dots < x_n = b$ be equally spaced
and write $h = (b-a)/n$ for the spacing. For $f \in C(a,b)$ there holds

$$
\int_a^b f(x) dx = \sum_{i=1}^n \int_{x_{i-1}}^{x_i} f(x) dx.
$$

We approximate using the trapezium rule

$$
\int_{x_{i-1}}^{x_i} f(x) dx \approx \frac{x_i - x_{i-1}}{2}(f(x_{i-1}) + f(x_i))
=
\frac h 2 (f(x_{i-1}) + f(x_i)),
$$

and obtain the _composite trapezium rule_

$$
\int_a^b f(x) dx
\approx
h \left( \frac12 f(x_0) + f(x_1) + \dots + f(x_{n-1}) + \frac12 f(x_n) \right).
$$

## Theorem: error in composite trapezium rule
> Let $f \in C^2(a,b)$ and write $I_h f$ for the composite trapezium rule applied to $f$ with the spacing $h>0$. Then there is $C>0$, independent of $f$ and $h$, such that
>
>$$
\left|\int_a^b f dx - I_h f \right| \le C \|(h\partial)^2 f\|_\infty.
$$

_Proof_. By the error in Newton-Cotes theorem,

\begin{align*}
\left| \int_a^b f(x) dx - I_h f \right|
\le
\sum_{i=1}^n \left|\int_{x_{i-1}}^{x_i} f(x) dx - \frac h 2 (f(x_{i-1}) + f(x_i)) \right|
\le n \|f''\|_\infty h^3.
\end{align*}

The claim follows from $h = (a-b)/n$.
$\blacksquare$

# Gaussian quadrature

The Gaussian quadrature is an approximation of the form

\begin{align*}
\int_a^b f(x) dx = \sum_{k=0}^n w_k f(x_k)
\end{align*}

where the quarature points $x_k \in [a,b]$ and weights $w_k$ are chosen so that the formula is exact for polynomials up to degree $2n + 1$.

In order to construct the points and weights we need to study Hermite interpolation that reads

> Let $x_i, y_i, z_i \in \mathbb R$, $i=0,\dots,n$, and suppose that $x_i \ne x_j$ for $i \ne j$.
>
> Find $p \in \mathbb P_{2n +1}$ such that $p(x_i) = y_i$ and $p'(x_i) = z_i$ for all $i=0,\dots,n$.

## Theorem: Hermite interpolation
> Let $n \ge 0$ be an integer let $x_i \in \mathbb R$, $i=0,\dots,n$, be distinct.
> Then for any $y_i, z_i \in \mathbb R$, $i=0,\dots,n$, there is a unique $p \in \mathbb P_{2n+1}$
> such that $p(x_i) = y_i$ and $p'(x_i) = z_i$ for all $i=0,\dots,n$.

For a proof, see Theorem 6.3 in [the book](#thebook).

## Theorem: error in Hermite interpolation
> Let $x_0, \dots x_n \in [a,b]$ be distinct, let $f \in C^{2n+2}(a,b)$, and
> let $p \in \mathbb P_{2n+1}$ be the Hermite interpolation polynomial of $f$, that is,
>
>\begin{align*}
p(x_i) = f(x_i), \quad p'(x_i) = f'(x_i), \qquad i=0,\dots,n.
\end{align*}
>
> Then for all $x \in [a,b]$ there is $\xi \in (a,b)$ such that
>
>\begin{align*}
f(x) - p(x) = \frac{f^{(2n+2)}(\xi)}{(2n+2)!} \prod_{i=0}^n (x-x_i)^2.
\end{align*}

For a proof, see Theorem 6.4 in [the book](#thebook).

Let $f \in C^{2n+2}(a,b)$, let

\begin{align*}
p(x) = \sum_{k=0}^n f(x_k) H_k(x) + \sum_{k=0}^n f'(x_k) K_k(x)
\end{align*}

be its Hermite interpolation polynomial, and consider the approximation

\begin{align*}
\int_a^b f(x) dx \approx \int_a^b p(x) dx = \sum_{k=0}^n w_k f(x_k) + \sum_{k=0}^n \tilde w_k f'(x_k),
\end{align*}

where

\begin{align*}
w_k = \int_a^b H_k(x) dx, \quad \tilde w_k = \int_a^b K_k(x) dx.
\end{align*}

The idea of the Gaussian quadrature is to choose the points $x_k$ so that $\tilde w_k = 0$ for all $k=0,\dots,n.$

Using

\begin{align*}
K_k(x) = L_k^2(x) (x - x_k), \quad L_k(x) = \prod_{i=0, i \ne k}^n \frac{x-x_i}{x_k-x_i}, \qquad k=0,\dots,n,
\end{align*}

we have

\begin{align*}
\tilde w_k = \int_a^b K_k(x) dx = c_n \int_a^b L_k(x) \pi(x) dx,
\end{align*}

where

\begin{align*}
\pi(x) = \prod_{i=0}^n (x-x_i),\quad c_n = \prod_{i=0, i \ne k}^n \frac{1}{x_k-x_i}.
\end{align*}

(In the case $n=0$, the formula holds with $c_n = 1$.) We see that $\tilde w_k = 0$ if $\pi \in \mathbb P_{n+1}$ is orthogonal to all polynomials in $\mathbb P_{n}$ in the sense of $L^2(a,b)$.

In the case that $a=-1$ and $b=1$, the [Legendre polynomials](https://en.wikipedia.org/wiki/Legendre_polynomials) $P_n \in \mathbb P_n$, $n=0,1,\dots$, are orthogonal with respect to the inner product $L^2(a,b)$. The points $x_i$, $i=0,\dots,n$ are chosen as the roots of $P_{n+1}$.

```python
from numpy.polynomial import legendre as leg
from matplotlib import pyplot as plt

xs_plot = np.linspace(-1,1)
for n in range(1, 4):
    cs = np.zeros(n + 2)
    cs[-1] = 1
    plt.plot(xs_plot, leg.legval(xs_plot, cs), label=f'$P_{n+1}$')
    roots = leg.legroots(cs)
    plt.scatter(roots, np.zeros(len(roots)))
plt.legend();
```

The case of arbitrary $a < b$ can be reduced to the case $a = -1$ and $b = 1$ by the change of variable

\begin{align*}
y = \frac{b - a} 2 x + \frac{a + b}2.
\end{align*}

Note that $x=-1$ corresponds to $y = a$ and $x = 1$ to $y = b$. Hence

\begin{align*}
\int_a^b f(y) dy
&= \frac{b - a} 2 \int_{-1}^1 f(\frac{b - a} 2 x + \frac{a + b}2) dx
\\&\approx \sum_{k=0}^n \frac{b - a} 2 w_k f(\frac{b - a} 2 x_k + \frac{a + b}2),
\end{align*}

where $w_k$ and $x_k$ are the quadrature weights and points on $[-1,1]$.

```python
def change_of_interval(x, a, b):
    return (b - a)/2 * x + (a + b)/2

a, b = 0, 1
n = 1

xs, ws = leg.leggauss(n + 1)
xs = change_of_interval(xs, a, b)
ws *= (b - a) / 2

fs = [x**n for n in range(5)]
for f in fs:
    test_integrate(ws, xs, f)
```

## Theorem: error in Gaussian quadrature

> Let $n=1,2,\dots$ and $f \in C^{2n+2}(a,b)$.
> Let $w_i, x_i$, $i=0,\dots,n$, be the quadrature weights and points in the Gaussian quadrature of order $n$ on $[a,b]$. Then
>
>\begin{align*}
&\left|\int_a^b f(x) dx - \sum_{k=0}^n w_k f(x_k) \right|
\\&\quad\le
\frac{\|f^{(2n+2)}\|_\infty}{(2n+2)!} \int_a^b (x-x_0)^2\dots(x-x_n)^2 dx.
\end{align*}

# On the integration sub-package of SciPy

The composite formulae are well-suited for approximation of integrals when the integrand is known only at some points $x_i$, $i=0,\dots,n$. The composite trapezium rule is implemented by [trapezoid](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.trapezoid.html) and its second order analogue, the composite Simpson's rule, by [simpson](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.simpson.html#scipy.integrate.simpson).

For highly regular grids $x_i$, $i=0,\dots,n$, there are methods that outperform the composite rules. If the samples are equally-spaced and the number of samples is $2^k + 1$ for some integer $k=1,2,\dots$, then Romberg integration [romb](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.romb.html#scipy.integrate.romb) can be used to obtain high-precision estimates of the integral using the available samples. This method is described in Section 7.7 of [the book](#thebook).

If the integrand can be evaluated at any point on the region of integration, then [quad](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.quad.html) can be used. This calls the Fortran library [QUADPACK](https://en.wikipedia.org/wiki/QUADPACK). The main focus of this library is on automatic integration in which the algorithm attempts to perform an integration with an error no larger than that requested. The library uses an [adaptive quadrature](https://en.wikipedia.org/wiki/Adaptive_quadrature) that divides the region of integration into adaptively refined subintervals. Generally, adaptive algorithms are just as efficient as composite rules for smooth integrands, but they may work for non-smooth integrands for which traditional algorithms may fail.

For more details on integration with SciPy see the [tutorial](https://docs.scipy.org/doc/scipy/tutorial/integrate.html).

```python
import scipy.integrate as integ
def f(x):
    return np.log(x)/(x+1)
a, b = 1, 2

xs = np.linspace(a, b)
ys = f(xs)
print(f'trapezoid gives {integ.trapezoid(ys, xs)}')
print(f'simpson gives   {integ.simpson(ys, xs)}')

xs, dx = np.linspace(a, b, 33, retstep=True)
ys = f(xs)
print(f'romb gives      {integ.romb(ys, dx)}')

# quad returns also an estimate of the absolute error in the result
print(f'quad gives     {integ.quad(f, a, b)}')
```

```python
def f(x):
    return 1/np.sqrt(x)
a, b = np.finfo(float).eps, 1

xs = np.linspace(a, b)
ys = f(xs)
print(f'trapezoid gives {integ.trapezoid(ys, xs)}')
print(f'simpson gives   {integ.simpson(ys, xs)}')

xs, dx = np.linspace(a, b, 33, retstep=True)
ys = f(xs)
print(f'romb gives      {integ.romb(ys, dx)}')

print(f'quad gives     {integ.quad(f, a, b)}')

print(f'correct value   2 (up to machine precision)')
```

```python

```


# Tutorial on Matrix Product States for Generative Learning
[Pan Zhang](http://lib.itp.ac.cn/html/panzhang/)
Institute of Theoretical Physics, Chinese Academy of Sciences
Reference: Z. Han, J. Wang, H. Fan, L. Wang and P. Zhang, [arXiv:1709.01662](https://arxiv.org/abs/1709.01662)   (to appear in *Phys. Rev. X* )
Date: 06.11.2018

_Abstract_:
Generative modeling is an important task in machine learning and artificial intelligence. It asks to learn joint probability distribution from training data and generates samples according to it. Many models have been proposed for generative modeling. In this tutorial, we introuduce a new model that is inspired by probabilistic interpretation of quantum physics, which we call _Born Machine_. In particular we introduce a generative model using matrix product states, which is a tensor network originally proposed for describing (particularly one-dimensional) entangled quantum states. The model enjoys efficient learning by utilizing the density matrix renormalization group method, and offers an efficient direct sampling for generative tasks.

The tutorial contains three parts:
1. An introducion to tensor networks and matrix product states (tensor train format).
2. An Introducion to generative modeling, and how to do this using matrix product states, i.e. MPS Born Machine.
3. In the last part I will give simple python code for training a MPS.

## 1.   An mini introduction to tensor networks and matrix product states
### 1.1 Data and tensors

* Data can be expressed (or mapped to) as a vector in a Hilbert space $\mathcal {H}$ of dimension $n$. Different mappings result to different spaces.
* Linear opearations in the space $\mathcal{H}$ can be expressed by an operator. Linear operators in ab high-dimensional space correpond to non-linear operators in a low-dimensional space: Kernel methods, representer theorem.
* In practice, we use **tensors** (multi-way arrays) to represent both operators and vectors.
* **Curse of dimensionality** $\Longleftrightarrow$ A large effective dimension $\Longleftrightarrow$ Tensors contain too many elements
* Low-dimensional approximation to a tensor $\Longleftrightarrow$ Compression of the Hilbert space $\Longleftrightarrow$ **Tensor networks**

### 1.2  Applications of tensor networks
* Tensor networks and graphical models: tensor contractions correpond to summation (or integration) of hidden (i.e. latent) variables
* In Physics, tensor networks are used for
  * expressing a wave function, by contracting virtual indices.
  * expressing the partition function of a classic system, by contracting all indices.
* In Applied Mathematics, tensor networks are used for
  * fitting a function or another tensor
  * Data compression
* In Machine Learning (new), tensor networks are used for
  * predicting unknown elements of a tensor using known its known elements.
  * As a classifier
  * Generative model via Born's rule (the topic of this tutorial)

  A biased conclusion: **Optimization** in physics, **Fitting** in applied mathematics, and **Generalization** in machine learning.

### 1.3  Tensor Networks
#### 1.3.1 Tensor diagrams

#### 1.3.2 Rank-1 tensors: product states
* The simplest tensor network: rank-one tensor, or factorized state, or factorized pure state. The number of elemens is proportional to the number of mode. The different modes are _independent_.
  * A qubit has two eigen states:
    * The first eigen state (analogous to pixel _black_) is expressed as    $| {\uparrow   } \rangle =|0\rangle=\left(\begin{matrix}1\\0\end{matrix}\right)$,
    * The second eigen state (analogous to pixel _white_) is expressed as $|{\downarrow }\rangle=|1\rangle=\left(\begin{matrix}0\\1\end{matrix}\right)$，
    * A pixel with gray level $0.8$ can be expressed as $|\Psi\rangle=\sqrt{0.8}|{\uparrow }\rangle+\sqrt{0.2}|{\downarrow }\rangle=\left(\begin{matrix}\sqrt{0.8}\\\sqrt{0.2}\end{matrix}\right)$
* Product of two qubits: $|\Psi\rangle =|{\uparrow }\rangle|{\downarrow }\rangle=\left(\begin{matrix}1\\0\end{matrix}\right)\otimes \left(\begin{matrix}0\\1\end{matrix}\right)=\left(\begin{matrix}0\\1\\0\\0\end{matrix}\right)=\left(\begin{matrix}0&0\\1&0\end{matrix}\right)$
* Product of two qubits: $|\Psi\rangle =|{\uparrow }\rangle\left(\frac{1}{\sqrt{2}}|{\uparrow }\rangle+\frac{1}{\sqrt{2}}|{\downarrow }\rangle\right)$
* Rank-one tensor can be used to fit any tensor i.e. extracting _key_ information of the tensor.  However the error could be very large, as the number of parameters in the rank-one tensor is very limited, hence rank-one tensor has very _limited representation power_. Consider a general system with $n$ qubits, it has $2^{n}-1$ free parameters, its rank-one fitting contains only $2n$ free parameters.

* Generalization of rank-one tensor: $n$ denotes number of modes, $d$ denotes local dimension, and $r$ denotes rank
  * CP decomposition，Canonical rank, number of parameters $ndr$
  * Tucker decomposition, number of parameters $ndr+r^n$
  * Tensor product states, also knon as _Tensor Train_, number of parameters $ndr^2$

#### 1.3.3 Tensors with rank $> 1$: entangled states
* Non-factorized pure states:
  * entanglement state of two qubits:   $\frac{1}{\sqrt{2}}|{\uparrow }\rangle|{\downarrow }\rangle-\frac{1}{\sqrt{2}}|{\downarrow }\rangle|{\uparrow }\rangle $, it can not be written as product state in any basis. The entanglements come from (anti) correlations.
  * EPR pair:    $\frac{1}{\sqrt{2}}|{\uparrow }\rangle|{\uparrow }\rangle+\frac{1}{\sqrt{2}}|{\downarrow }\rangle|{\downarrow }\rangle $, what is its rank?
  * $|\textrm{GHZ}\rangle=\frac{1}{\sqrt{2}}\left( |{\uparrow }\rangle^{\otimes n}+ | {\downarrow }\rangle^{\otimes n}\right)$

* Rank and engtanglement:
 * Rank-one $\Longleftrightarrow$ no entanglements
 * Larger rank, more enganglements. Quantitatively characterized by _entanglement entropy_.
 * Entanglments mean that some information of the system is shared by two qubits. Looking at one isolated qubit looses information!
 * Unfold a tensor into a <math>2\times 2</math>matrix, rank is the number of non-zero singular values of the matrix. It is also equivalent to the number of non-zero eigenvalues of the _reduced density matrix_.
 * Why SVD? Because SVD = Schmidt decomposition, a natural orthogonal bisection method!
 * How to determine whether there are engantlements between two sub-systems of a system with multiple (say $7$) qubits?
   * The same principle: unfoled the size-$2^7$ tensor to a matrix of size $2^3\times 2^4$, then look at its rank.
   * Squared singular values of this matrix ＝eigenvalues of the reduced density matrix.
   * Recall the definition of singular values: it is actually equivalent to find eigenvalues of the related covariance matrix.
   * There are many methods for unfolding this size-$2^7$ tesor, each of which corresponds to different _permute_$\to$_reshape_ operations. If we can do this for any tensor (with unlimited memory and computational power), it gives exact estimate of entaglement entropy, as well as the bond dimension of DMRG.
   * A quick test: is $|\Psi\rangle=\frac{1}{ {2}}|{\uparrow }\rangle|{\uparrow }\rangle+\frac{1}{ {2}}|{\downarrow }\rangle|{\downarrow }\rangle+\frac{1}{ {2}}|{\uparrow }\rangle|{\downarrow }\rangle+\frac{1}{ {2}}|{\downarrow }\rangle|{\uparrow }\rangle$ entangled?

* Entanglment entropy
  * The entanglement entropy defined on eigenvalues of the reduced density matrix.
  * The density matrix $\rho=|\Psi\rangle\langle\Psi|$ of pure state $\Psi$ is rank-1.
  * Whether the reduced density matrix $\rho_A=\textrm{Tr}_B|\Psi\rangle\langle\Psi|$ has rank larger than $1$ depends on whether sub-system A is entangled with sub system B.
  * $\psi_{AB}$ is a matrix obtained by unfolding tensor $\Psi$, and $\rho_A=\psi_{AB}\psi_{AB}^\dagger$. Thus rank of $\rho_A$ equals $\rho_B$, which is also identical to rank of $\psi_{AB}$.
  * Eigenvalues of $\rho_A$ are coinside with eigenvalues of $\rho_B$, which are identical to squared singular values of $\psi_{AB}$.
  * Entanglement entropy is defined as $S(\rho_A)=S(U\rho_A U^\dagger)=S(P(\{\lambda_i\}))=\sum_i\lambda_i\log(\lambda_i)$, where U is a (column-orthogonal) unitary matrix, $P(\{\lambda_i)\}$ denotes distribution of eigenvalues of the reduced density matrix.

#### 1.3.4 Tensor networks used in physics

### 1.4  Tensor decompositions
I want to emphasize that tensor networks are indeed tensor decompositions. This section focus on the mathematical aspects of them.
* Canonical Polyadic (CP) decomposition
to be update
* Tucker decomposition, also known as _higher order SVD_, _n-mode PCA_.
to be update
* Tensor train decomposition, which is named _matrix product states_.
to be update
* Hierarchical Tucker decomposition, which is known as _tree tensor netorks_ in physics.
to be update

## 2.   An introduction to generative modeling using matrix product states
### 2.1 Generative modeling using Born Machine
There are mainly two kinds of tasks in marchine learning
* Supervised learning (discriminative learning), where data $x$ come with labels $t$. It asks to model conditional distributions $P(t|x)$.
* Unsupervised learning (generative learning), where data $x$ does not have label. It asks to model joint distribution of data $P(x)$.

Many models have been proposed for generative learning, for example, the Hidden Markov model, Boltzmann machines, Gaussian mixture models, Variational autoencoders, generative adversary networks... Among them, the Boltzmann machine (BM) is closely related to the Ising model in statistical physics, and ultilizes Boltzmann distribution for modeling data joint distribution.
$$P(x)=\frac{1}{Z}e^{-E(x)},$$
where $E(x)$ denotes energy of configuration $x$ and $Z$ is the partition function.

In this tutorial, we model joint distribution using _Born's_ rule of quantum mechanics
$$P(x)=\frac{1}{\mathcal{N}}{|\Psi(x)|^2},$$
where $\Psi$ denotes wave function, $\mathcal{N}$ is the normalization factor. We call the model _Born Machine_.

### 2.2 Data mapping to Hilbert space
* Each image (among totoally $m$ images) is mapped to a product state $\mathbf{x}=v_1\otimes v_2\otimes v_3\otimes\cdots\otimes v_n$
  * A white pixel $i$ is represented as $v_i=| {\uparrow   } \rangle =|0\rangle=\left(\begin{matrix}1\\0\end{matrix}\right)$,
  * A black pixel $j$ is represented as $v_j=|{\downarrow }\rangle=|1\rangle=\left(\begin{matrix}0\\1\end{matrix}\right)$.
* The data distribution becomes $$ Q_{\mathrm{data}}(\mathbf{x}) = \frac{1}{m}\sum_{t=1}^m\delta(x - v_1^t\otimes v_2^t\otimes\cdots\otimes v_n^t)$$
* This can be seen as rank-$m$ CP decomposition of the original (unknown) tensor.

### 2.3 Matrix product state as wavefunction in Born machine
\begin{align}
P(\mathbf{x})&=\frac{|\Psi(\mathbf{x})|^2}{Z},
\end{align}
and the wave function is parametrized by MPS

and is formulated as
\begin{align}
\Psi(\mathbf{x})&=\Psi_{\mathrm{mps}} \cdot \left( v_1^t\otimes v_2^t\otimes \cdots\otimes v_n^t  \right)\\
&=\sum_{w_1,w_2,...w_{n-1}} A^{(1)}_{v_1,w_1}A^{(2)}_{w_1,v_2,w_2}\cdots A^{(n-1)}_{w_{n-2},v_{n-1},w_{n-1}}A^{(n)}_{w_{n-1},v_{n}}\\
&= A^{(1)}_{v_1}A^{(2)}_{v_2}\cdots A^{(n)}_{v_n}
\end{align}
and $Z$ represent the partition function
\begin{align}
Z&=\sum_{\mathbf{x}}\Psi(\mathbf{x})
\end{align}
* The number of parameters is roughly $2nD^2$, rather than $2^n$.
* Training via an algorithm analogous to the Density Matrix Renormalization Group.
* Easy computation of the partition function (thanks to the data mapping).
* Direct sampling.
* Exact representation of whole data requires $D=m$ (recall the MPS and CP section), but this is not what we want:
  * $m$ is too large as a bound dimension.
  * It is _overfitting_!

### 2.4 Learning of MPS
#### 2.4.1 Maximum-lieklihood learning
The goal of learning is to make the joint distribution given by the Born's rule as close as the data distribution. First one needs to quantify _distance_ between two probability distributions. Here we use the Kullback-Leibler divergence between data distribution $Q_{\mathrm{data}}(\mathbf{x})$ and the MPS distribution $P_\theta(\mathbf{x})$ which is parametrized by $\theta$ (tensor elements).
\begin{align}
D_{\mathrm{KL}}\left(Q_{\mathrm{data}}(\mathbf{x})|P_\theta(\mathbf{x})\right)&=\sum_{\mathbf{x}}Q_{\mathrm{data}}(\mathbf{x})\log\frac{Q_{\mathrm{data}}(\mathbf{x})}{P_\theta(\mathbf{x})}\\
&=\sum_{\mathbf{x}}Q_{\mathrm{data}}(\mathbf{x})\log Q_{\mathrm{data}}(\mathbf{x})-\sum_{\mathbf{x}}Q_{\mathrm{data}}(\mathbf{x}) \log P_\theta(\mathbf{x})\\
&=-S_Q+\mathcal{L_\theta},
\end{align}
where $S_Q$ is the entropy of distribution $Q_{\mathrm{data}}$, and $\mathcal{L}$ is the so-called _Negative Log-Likelihood_ (NLL).

It is easy to show (e.g. using Jenson's inequality) that KL divergence is non-negative. This indicates a lower-bound for the NLL $\mathcal{L_\theta}\geq S_Q$.
Notice that this lower bound is achieved only when $P_\theta(\mathbf{x}) = Q_{\mathrm{data}}(\mathbf{x})$.

The learning amounts to finding parameters $\theta$, i.e. tensor elements, such that the KL divergence is minimized.
$$\hat \theta=\arg\min_\theta D_{\mathrm{KL}}\left(Q_{\mathrm{data}}(\mathbf{x})|P_\theta(\mathbf{x})\right) =\arg\max_\theta \mathcal{L_\theta}$$
#### 2.4.2 Learning of MPS
Let us rewrite the NLL as
\begin{align}
\mathcal{L_\theta}&=-\sum_{t=1}^mQ_{\mathrm{data}}(\mathbf{x}_t) \log P_\theta(\mathbf{x})\\
&=-\frac{1}{m}\sum_{t=1}^m\log |\Psi(\mathbf{x}_t)|^2 +\log Z\\
\end{align}
then update tensor elements using gradient descent.

The gradients are computed as
\begin{align}
\frac{\partial \mathcal{L}}{A^{(i)}_{w_{i-1},v_i,w_{i}}}&=-\frac{2}{m}\sum_{t=1}^m\frac{\Psi'_{w_{i-1},v_i,w_i}(\mathbf{x_t})}{\Psi(\mathbf{x_t})}+\frac{2A^{(i)}_{w_{i-1},v_i,w_{i}}}{Z},
\end{align}
where $\Psi'_{w_{i-1},v_i,w_{i}}(\mathbf{x_t})$ is the derivative of $\Psi$ with respect to tensor element  $A^{(i)}_{w_{i-1},v_i,w_{i}}$ evaluated at $\mathbf{x_t}$. It is evaluated as $$\Psi'_{w_{i-1},v_i,w_{i}}(\mathbf{x_t}) =u^l_{w_{i-1},v_i}(\mathbf{x_t})\cdot u^r_{w_i,v_i}(\mathbf{x_t})=\left(\sum_{w_1,w_2,...w_{i-2}}A^{(1)}_{v_1,w_1}A^{(2)}_{w_1,v_2,w_2}\cdots A^{(i-1)}_{w_{i-2},v_{i-1},w_{i-1}}\right )\cdot \left ( \sum_{w_{i+1},w_{i+2},...w_{n}}A^{(i+1)}_{v_i,w_i,v_{i+1}}\cdots A^{(n)}_{w_{n-1},v_{n}}\right).$$
The gradients can actaully be written in a more compact form:
$$\bigtriangledown_{A^{(i)}}\mathcal{L}=-\frac{2}{m}\sum_{t=1}^m\frac{ u^l(\mathbf{x_t})\circ u^r(\mathbf{x_t}) }{\Psi(\mathbf{x_t})}+\frac{2A^{(i)}}{Z},$$
where $u^l$ and $u^r$ are vectors contracted from left and from right respectively, as shown in the function above. $u^l(\mathbf{x_t})\circ u^r(\mathbf{x_t})$ denotes a rank-one tensor obtained by tensor product of vectors $u^l(\mathbf{x_t})$ and $u^r(\mathbf{x_t})$, it is also known as _environment tensor_.

Another important part in computing gradients is that the partition function $Z$ is easy to compute, thanks to the orthogonalization.

### 2.5 Generating new images using _direct sampling_
The most important feature of MPS is that the joint probability distribution can be computed exactly,

as well as conditional probabilities

This property makes pixel-by-pixel sampling easy and exact.

### 2.6 Numerical results
A typical behavior of the log-likelihood look like

where the NLL decreases monotonically, untill it reaches the lower-bound $S(Q)=\log(m)$ for non-overlapping training data.
The lower-bound is always achiveable by increasing the bond-dimension

However usually one does not want to reach the lower-bound, because it is overfitting: training NLL does not make too much sense, test NLL does!

After the MPS is well trained, it could find many unsupervised applications, e.g. reconstructing partial images

to

## 3. A simple implementation of MPS Born Machine

```python
import sys
import numpy as np
import torch
import math
%matplotlib inline
import matplotlib.pyplot as plt

torch.manual_seed(1) # Fix seed of the random number generators
np.random.seed(1)
```

The function below is used to plot mnist images

```python
def show_imgs(imgs,l1=4,l2=5,s1=6,s2=6):
    """    Plot images    """
    plt.rcParams['figure.figsize']=(s1,s2)
    imgs=imgs.reshape([-1,28,28])
    g, ax = plt.subplots(l1,l2)
    for i in range(l1):
        for j in range(l2):
            a=i*l1+j
            if(a>=imgs.shape[0]):
                break
            ax[i][j].imshow(imgs[a,:,:],cmap='summer')
            ax[i][j].set_xticks([])
            ax[i][j].set_yticks([])
    plt.show()
```

### Data loading
$100$ MNIST images have been stored as "mnist_100_28x28_p0.5.npy".
Each image contains $n=28\times 28=784$ pixels, each of which takes value $0$ or $1$.
In our settings, each image is viewed as a product state in the Hilbert space of  dimension $2^n$.

```python
n=784 # number of qubits
m=20 # m images
data=np.load("imgs/mnist_100_28x28_p0.5.npy")
data=data[:m,:,:]
data=torch.LongTensor(data)
data=data.view(-1,784) # m images, each of which is reshaped into a vector of length 784
show_imgs(data,2,10,10,2)
```

### MPS initialization
Define the mps, which is a list of 3-way tensors containing random values

```python
Dmax=30 # maximum bond dimension
bond_dims=[Dmax for i in range(n-1)]+[1]
tensors= [ torch.randn(bond_dims[i-1],2,bond_dims[i]) for i in range(n)]
```

Now check the bond dimensions and tensors

```python
print("There are len(tensors) tensors")
print(bond_dims)
print(tensors[5].shape)
```

    There are len(tensors) tensors
    [30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 1]
    torch.Size([30, 2, 30])

Question: does the contration with one image give a probability of the image? Why?

Canonicalization using QR decompositions

```python
def orthogonalize(site,going_right):
    dl=bond_dims[site-1] # left bond dimension
    d=bond_dims[site]   # current bond dimension
    if(going_right):
        A=tensors[site].view(dl*2,d) # A is a matrix unfolded from the current tensor
        Q,R=torch.qr(A)
        R/=R.norm() # devided by norm
        tensors[site] = Q.contiguous().view(dl,2,-1)
        tensors[site+1] = (R@tensors[site+1].view(d,-1)).view(-1,2,bond_dims[site+1])
        bond_dims[site] = Q.shape[1] # economy QR, so the right dimension could be either dl or d
    else: # going left
        A=tensors[site].view(dl,d*2).t()
        Q,R=torch.qr(A)
        R/=R.norm()
        tensors[site]=Q.t().contiguous().view(-1,2,d)
        tensors[site-1] = (tensors[site-1].view(-1,dl)@R.t()).view(bond_dims[site-2],2,-1)
        bond_dims[site-1] = Q.shape[1]

```

### MPS left canonicalization
Canonicalization gives several advantages:
* It make the partition function of the model equals to $1$.
* The isometries have condition number $1$, preserving very well the computation precisions.

```python
for site in range(n-1):
    orthogonalize(site,True)
```

```python
tensors[783].shape
```

    torch.Size([30, 2, 1])

Canonicalization changes the bond dimensions, to see it

```python
print(bond_dims)
```

    [2, 4, 8, 16, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 1]

Now contracting mps with one image gives the probability amplitude of the image.

```python
def get_psi():
    psi=torch.ones([m,1,1])
    for site in range(n):
        psi = psi @ tensors[site][:,data[:,site],:].permute(1,0,2)
    return(psi)
```

```python
def gen_samples(ns):
    samples=np.zeros([ns,n])
    for site in range(n-1):# left canonicalize
        orthogonalize(site,True)
    for s in range(ns):
        vec=torch.ones(1,1)
        for site in range(n-1,-1,-1):
            vec = (tensors[site].view(-1,bond_dims[site])@vec).view(-1,2)
            p0 = vec[:,0].norm()**2/ (vec.norm()**2)
            x = (0 if np.random.rand() < p0 else 1)
            vec = vec[:,x]
            samples[s][site]=x
    return samples
```

### Initialize cache for MPS
Computing probability of a image consists of contracting bonds from the first tensor to the last tensor. Lots of computation results can be re-used in the future computations, so we would like to store the contraction results in the cache.
Notice that the cache is for all images.

```python
cache=[]
cache.append( torch.ones([m,1,1])) # The initial elements, all images have cache 1
for site in range(n-1):
    B=cache[site] @ tensors[site][:,data[:,site],:].permute(1,0,2)
    B /= B.abs().max()
    cache.append(  B  ) # batched matrix multiplications
cache.append( torch.ones(m,1,1)) # the last element, matrix [1,1] for all images
```

Notice that code ```B /= B.abs().max()``` is used to preserve computational precision, avoiding maximum value of the cache becomes too small. We can do this because the gradient is invariant by multiplying the cache by a constant, for each image.
Length of the cache is $n+1$. In the caceh, for an image, each pixel has a correponding vector, denoting the temporary results in tensor contractions (from left to right or from right to left).
Let us look at the content of cache for image alpha=1:

```python
alpha=1 # the image 1
print("cache site 1 ",cache[0][alpha])
print("cache site 2 ",cache[1][alpha])
print("cache site 3 ",cache[2][alpha])
```

    cache site 1  tensor([[ 1.]])
    cache site 2  tensor([[-1.0000, -0.5231]])
    cache site 3  tensor([[ 0.6791,  0.3302, -0.6933,  1.0000]])

Then the probability amplitude $\psi$ for images can be obtained by

```python
psi=get_psi()
```

Let us output the probablity for an image, which equals $|\psi|^2$

```python
print("Probability of generating image 3 = %.5f"%(psi*psi)[3])
```

    Probability of generating image 3 = 0.00000

```python
ax=plt.bar(range(1,psi.shape[0]+1),psi**2)
```

**Hey, the probability is $0$, what is wrong?**

Because the space is too large and the model is randomly initialized!

Actually the purpose of training is exactly to increase the probability (or equivalently, log probability) of given images. This is the so-called *maximum likelihood learning*

The basic procedure is sweeping back and force, from right to left, then from left to right. During each sweep, the visited tensor is updated according to the *gradients* of the log-probability with respect to tensor elements.

## Training MPS

Some tips on the code:
* Only images with $i^{\mathrm{th}}$ element $v_i$ contribute to tensor element $A_{w_{i-1},v_i,w_i}$.
* ```torch.sum(left_vec.permute(0,2,1) @ right_vec.permute(0,2,1)``` computes $\Psi'(\mathbf{x})$.
* ```psi = left_vec @ A.permute(1,0,2) @right_vec ``` does not compute _real_ $\Psi$, as cache is rescaled to preserve precision. Thus one needs to call ```psi=get_psi()``` in order to compute correct probability amplitude $\Psi$.
* ```tensors_bak=tensors.copy()``` is to backup the training environment for restoring after generating images.
* ```@``` operator in ```left_vec @ A.permute(1,0,2) @right_vec ``` and in ```cache[site] @ tensors[site][:,data[:,site],:].permute(1,0,2)``` works as batched matrix-matrix multiplications.
* In ```orthogonalize()```, imposing matrix $R$ to be norm-one using ```R/=R.norm() ``` actually let the partition function $Z=1$.

```python
learning_rate=0.08
for epoch in range(9):    # one sweep, from right to left, then from left to right
    going_right=False
    for site in [i for i in range(n-1,0,-1)]+[i for i in range(n-1)]:
        # to update tensors[site] which is a 3-way tensor of size [dl,2,dr]
        sys.stdout.write("\r Epoch #%d, site #%d / %d           "%(epoch, site+1,n)); sys.stdout.flush()
        if(site==0): going_right=True
        gradients = torch.zeros_like(tensors[site])
        for i in [0,1]: # the pixel could be either 0 or 1
            idx=(data[:,site]==i).nonzero().type(torch.LongTensor).squeeze() # this returns indices of non-zero elements
            if(idx.numel()==0): continue
            left_vec = cache[site][idx,:,:] # a vector on the left of the site
            right_vec = cache[site+1][idx,:,:] # a vector on the right of the site
            A=tensors[site][:,data[idx,site],:]
            if(idx.numel()==1):
                A=A.view(A.shape[0],1,A.shape[1])
                left_vec=left_vec.view(1,left_vec.shape[0],left_vec.shape[1])
                right_vec=right_vec.view(1,right_vec.shape[0],right_vec.shape[1])
            psi = left_vec @ A.permute(1,0,2) @right_vec
            gradients[:,i,:] = torch.sum(left_vec.permute(0,2,1) @ right_vec.permute(0,2,1) / psi,0)
        gradients = 2.0*(gradients/m-tensors[site])
        tensors[site] += learning_rate * gradients/gradients.norm()
        orthogonalize(site,going_right)
        if(going_right):
            cache[site+1] = cache[site] @ tensors[site][:,data[:,site],:].permute(1,0,2)
            cache[site+1] /= cache[site+1].abs().max()
        else:
            cache[site] = tensors[site][:,data[:,site],:].permute(1,0,2) @ cache[site+1]
            cache[site] /= cache[site].abs().max()
    psi=get_psi()
    tensors_bak=tensors.copy()
    sys.stdout.write("NLL=%.3f, LowerBound=%.3f "%(-torch.mean(torch.log(psi*psi)),math.log(m)))
    sys.stdout.write("  generating samples...")
    ax=plt.bar(range(1,psi.shape[0]+1),psi**2)
    imgs=gen_samples(30)
    show_imgs(imgs,2,15,15,2)
    tensors=tensors_bak.copy()
    if(epoch < 8):
        input("Press Enter to continue  ")

```

### References for further reading:
* Ulrich Schollwock, “The density-matrix renormalization group in the age of matrix product states,” Annals of Physics 326, 96– 192 (2011).
* Ivan V Oseledets, “Tensor-train decomposition,” SIAM Journal on Scientific Computing 33, 2295–2317 (2011).
* Edwin Miles Miles Stoudenmire and David J. Schwab, “Supervised Learning with Quantum-Inspired Tensor Networks,” Advances in Neural Information Processing Systems 29, 4799 (2016), arXiv:1605.05775.
* Alexander Novikov, Mikhail Trofimov, and Ivan Oseledets, “Exponential machines,” arXiv:1605.05775 (2016).
* Jinguo Liu, and Lei Wang,Differentiable Learning of Quantum Circuit Born Machine, arXiv:1804.04168 (2018).
* Edwin Miles Stoudenmire, “Learning relevant features of data with multi-scale tensor networks,” Quantum Science and Technology (2018).

```python

```