# An introduction to Gaussian Processes

```python
%matplotlib inline
```

```python
%config InlineBackend.figure_format = 'retina'
```

```python
from matplotlib import rcParams
rcParams["figure.dpi"] = 100
rcParams["figure.figsize"] = 12, 4
```

## The Gaussian Distribution

In this notebook, we'll go over the very basics of Gaussian Processes (GPs) and how to construct and draw samples from them. But first, let's review some stuff about the Gausssian distribution $-$ the familiar bell curve $-$ itself.

We're probably all familiar with ``numpy``'s built-in ``np.random.randn()`` function, which draws a sample from the standard normal $\mathcal{N}(0, 1)$, i.e., a Gaussian with zero mean and unit variance:

```python
import numpy as np
import matplotlib.pyplot as plt
np.random.seed(1)
u = np.random.randn(100000)
plt.hist(u, bins=50, density=1);
plt.axvline(0, color="C1", ls="--")
y = np.exp(-0.5) / (np.sqrt(2 * np.pi))
plt.plot([-1, 1], [y, y], "C1-")
plt.text(0.1, 0.05, r"$\mu$", color="C1", fontsize=20);
plt.text(0.3, 0.26, r"$\sigma$", color="C1", fontsize=20);
plt.text(-0.5, 0.26, r"$\sigma$", color="C1", fontsize=20);
```

To draw from a Gaussian distribution with a different mean $\mu$, we simply *add* $\mu$ to $u$, and to draw from a distribution with different variance $\sigma^2$, we simply *multiply* $u$ by $\sigma$. We can therefore draw from a Gaussian with mean (say) $10.0$ and variance $5.0^2$ by running

```python
np.random.seed(1)
u = np.random.randn(100000)
plt.hist(10.0 + 5.0 * u, bins=50, density=1);
```

As I mentioned in my lecture, in a single dimension, the probability density function for a Gaussian with mean $\mu$ and variance $\sigma^2$ is

$$
p(y | \mu, \sigma^2) =
\frac{1}{\sqrt{2\pi\sigma^2}}
\exp\left[
-\frac{1}{2}\left(\frac{y - \mu}{\sigma}\right)^2
\right]
$$

In Exercise 1 below, we'll derive the expression for a Gaussian in $N$ dimensions. We'll do it for the case where there isn't any covariance across the $N$ dimensions, in which case the probability density is obtained by multiplying the expression above $N$ times.

Exercise 1

Show (numerically) that the probability function for an uncorrelated Gaussian in $N$ dimensions

$$
p(\mathbf{y} | \boldsymbol{\mu}, \mathbf{\Sigma}) =
\prod_{n=0}^{N-1}
\frac{1}{\sqrt{2\pi\sigma_n^2}}
\exp\left[
-\frac{1}{2}\left(\frac{y_n - \mu_n}{\sigma_n}\right)^2
\right]
$$

can be written in vector form as

$$
p(\mathbf{y} | \boldsymbol{\mu}, \mathbf{\Sigma}) =
\frac{1}{\sqrt{(2\pi)^N \mathrm{det}(\mathbf{\Sigma})}}
\exp\left[
-\frac{1}{2}
(\mathbf{y} - \boldsymbol{\mu})^\top
\mathbf{\Sigma}^{-1}
(\mathbf{y} - \boldsymbol{\mu})
\right]
$$

where

$$
\mathbf{\Sigma} =
\begin{pmatrix}
 \sigma_0^2 & 0 & 0 \\
 0 & \sigma_1^2 & 0 \\
 0 & 0 & \ddots
\end{pmatrix}
$$

is a diagonal covariance matrix. You can use the following random inputs to test the equivalence of the two expressions:

```python
N = 5
y = np.array([1.0, 2.0, 0.5, 0.75, 1.1])
mu = np.array([0.7, 1.5, 0.8, 0.7, 1.3])
sigma = np.array([0.3, 0.5, 0.5, 0.4, 0.5])
```

```python

```

Even though we derived the expression for the probability of a multidimensional Gaussian assuming no covariance between the different dimensions, the expression in Exercise 1 is general. In the most general case, the covariance matrix is allowed to have off-diagonal elements:

$$
\mathbf{\Sigma} =
\begin{pmatrix}
 \mathrm{Var(x_0)} & \mathrm{Cov(x_0, x_1)} & \cdots & \mathrm{Cov(x_0, x_N)} \\ \\
 \mathrm{Cov(x_1, x_0)} & \mathrm{Var(x_1)} & \cdots & \mathrm{Cov(x_1, x_N)} \\ \\
 \vdots & \vdots & \ddots & \vdots \\ \\
 \mathrm{Cov(x_N, x_0)} & \mathrm{Cov(x_N, x_1)} & \cdots & \mathrm{Var(x_N)}
\end{pmatrix}
$$

Here, $\mathrm{Var(x_i)}$ is the variance, or square of the standard deviation, along the $i^\mathrm{th}$ dimension, and $\mathrm{Cov(x_i, x_j)}$ is the covariance $-$ a measure of how two variables vary jointly $-$ between the $i^\mathrm{th}$ and $j^\mathrm{th}$ dimensions.

It's hard to picture a multivariate Gaussian, but it helps to imagine what the contours of equal probability (or equipotential surfaces) look like. For a two dimensional standard normal with no covariance ($\mathbf{\Sigma} = \mathbf{I}$), these are just circles. As I add to the off-diagonal terms of the matrix, I introduce a preferred direction in the space, since now the dimensions are correlated. The circle becomes an ellipse whose elongation increases with the covariance. In the general case of a multidimensional Gaussian, the contours form a multidimensional ellipsoid, elongated by different amounts along different dimensions.

**So.** How do we draw samples from a multivariate Gaussian given $\mathbf{\Sigma}$? We use ``np.random.multivariate_normal``.

Here's a function ``draw_from_gaussian(mu, S, ndraws=1)`` that returns ``ndraws`` samples from a multivariate Gaussian with mean ``mu`` and covariance matrix ``S``. The shape of the output is ``(ndraws, ndim)`` where ``ndim`` is the dimension of the Gaussian:

```python
def draw_from_gaussian(mu, S, ndraws=1):
    """
    Generate samples from a multivariate gaussian
    specified by covariance ``S`` and mean ``mu``.

    """
    return np.random.multivariate_normal(mu, S, (ndraws,))
```

Exercise 2

Use the function above to draw ``10,000`` samples from a zero-mean Gaussian with covariance

$$
\mathbf{\Sigma} =
\begin{pmatrix}
 1 & 0.5 \\
 0.5 & 1
\end{pmatrix}
$$

Plot the "corner" plot for your ``samples`` using the ``corner`` package (``!pip install corner``):

```python
from corner import corner
fig = corner(samples);
```

Vary the terms in the covariance matrix (recalling that it must be symmetric!) to get a sense of how they affect the joint distribution of samples.

```python

```

## Parametrizing the covariance

As the number of dimensions increases, and in particular as it becomes *infinite* (which it will, when we get to actual Gaussian Processes), it's no longer convenient to describe the covariance in terms of every single entry in the covariance matrix. Instead, it's useful to introduce the notion of a *kernel*, a function of just a few parameters that describes the overall structure of the covariance matrix.

In general, the covariance matrix can have *any* structure, but quite often in timeseries analysis *data points close to each other in time are strongly correlated*. This is true whether your timeseries contains photometric variability from a rotating star, the rise and fall of a supernova light curve, a gravitational wave signal, or PSF shape changes due to temperature fluctuations on the detector. This is because each of these processes have characteristic timescales over which they operate, and on timescales shorter than that, all measurements you make are likely to be close to relatively each other. But if you wait long enough, what your system will do down the line will be pretty decoupled from what it is doing right now.

So we can imagine defining a kernel that looks something like this:

where the covariance peaks at zero time lag and drops smoothly to zero as the time lag increases. This particular kernel is extremely useful and has a name: the *Squared Exponential Kernel*, defined as

$$
\begin{align}
    k(t_i, t_j) = A^2 \exp\left(-\frac{\left(t_i - t_j\right)^2}{2l^2}\right)
\end{align}
$$

Here is a function ``ExpSquaredCovariance`` that returns the covariance matrix described by the squared exponential kernel for a timeseries ``t`` given an amplitude ``A`` and a timescale ``l``:

```python
def ExpSquaredCovariance(t, A=1.0, l=1.0, tprime=None):
    """
    Return the ``N x M`` exponential squared
    covariance matrix.

    """
    if tprime is None:
        tprime = t
    TPrime, T = np.meshgrid(tprime, t)
    return A ** 2 * np.exp(-0.5 * (T - TPrime) ** 2 / l ** 2)
```

Don't worry about the `tprime` keyword for now -- we'll use it later in Exercise 7.

Exercise 3

For ``t = np.linspace(0, 2, 11)``, draw ``10,000`` samples from the Gaussian described by the covariance matrix above and plot the corresponding corner plot, labeling each subplot with the time point it represents. Comment on the structure of the plot and the correlations among the different dimensions.

```python

```

## From Gaussians to Gaussian Processes

So far we've been taking about Gaussians *distributions* and how they can jointly model a few random variables at a time. But what happens as the dimensionality of the Gaussian increases and eventually becomes infinite? The Gaussian no longer represents a collection of random variables, but instead the behaviour of a continuous function. And that's exactly what a Gaussian Process is: a distribution over functions with infinitely many points. We can't technically model a function with infinitely many points on a computer, but we can investigate what happens as the number of points becomes very large.

Exercise 4

Construct a ``1000 x 1000`` covariance matrix using the squared exponential kernel for a timeseries spanning 10 time units, and plot an image of it using ``plt.imshow``. Comment on the structure of the matrix.

```python

```

Below, we draw ``30`` samples from the zero-mean Gaussian described by the covariance matrix you plotted in the previous exercise. These are shown as the transluscent black curves. The mean of the process is the solid blue line, and the 1 and 2 sigma levels are indicated with the blue shading.

```python
t = np.linspace(0, 10, 1000)
mu = np.zeros_like(t)
S = ExpSquaredCovariance(t, A=1.0, l=1.0)
np.random.seed(1)
for i in range(30):
    samples = draw_from_gaussian(mu, S)
    plt.plot(t, samples.T, color="k", alpha=0.1)

mu_ = np.zeros_like(t)
plt.plot(t, mu_, color="C0")
plt.xlabel("time")
plt.ylabel("GP")
std_ = np.sqrt(S.diagonal())
plt.fill_between(t, mu_ - std_, mu_ + std_, color="C0", alpha=0.3);
plt.fill_between(t, mu_ - 2 * std_, mu_ + 2 * std_, color="C0", alpha=0.15);
```

These curves are samples from the (infinitely large) family of functions described by your GP!

Exercise 5

In the code above, vary the hyperparameters of the GP, $A$ and $l$, and comment on how they change the behavior of the GP.

## Conditioning on data

If you made it this far, you have actually implemented your own Gaussian Process! But it's still pretty simple, because so far we are drawing only from our *prior*: we know the Gaussian has a certain covariance and a certain mean (zero, in the case above), but nothing else about how it's supposed to behave. To apply this to a real dataset, we need to *condition* the GP on observations.

Say I make two perfectly noise-free observations:

$$
\begin{align}
    y(t = 2.5) &= 1.0 \\
    y(t = 7.5) &= -1.0
\end{align}
$$

Again, $y$ could be anything: the flux received from a star, the strain due to a gravitational wave, etc. I want to draw new samples from my GP, but not just *any* samples: I explicitly want to draw functions that go through those two data points. One (very inefficient) way to do this is to draw a *ton* of samples and discard any samples that don't agree with the data. (In this particular case, that would be literally impossible, since I know $y(t)$ at those two points *exactly*: I will never draw at random a function whose value is exactly what I want it to be. But note that *rejection sampling* can actually be useful in other cases. Anyways.)

A *better* way to do this is to consider the *joint distribution* of the observed data $y(t)$ (usually called the "training set") and the points where I'm trying to predict the value of the function $y_\star(t_\star)$ (the "test set"):

$$
\begin{align}
\begin{pmatrix}
 y \\
 y_\star
\end{pmatrix}
&=
\mathcal{N}
\left[
\mathbf{0},
\begin{pmatrix}
 \mathbf{\Sigma}(t, t) & \mathbf{\Sigma}(t, t_\star)\\
 \mathbf{\Sigma}(t_\star, t) & \mathbf{\Sigma}(t_\star, t_\star)
\end{pmatrix}
\right]
\end{align}
$$

where $\mathbf{\Sigma}(t, t)$ is the $N_\mathrm{train} \times N_\mathrm{train}$ covariance matrix evaluated at all pairs of training points, $\mathbf{\Sigma}(t_\star, t_\star)$ is the $N_\mathrm{test} \times N_\mathrm{test}$ covariance matrix evaluated at all pairs of test points, and the remaining two entries are the (rectangular) covariance matrices evaluated at all pairs of (test, training) points.

Given this joint distribution, we can compute the *conditional* distribution of $y_\star$ given $y$ with some linear algebra (for a derivation, see Appendix A.2 of [Rasmussen & Williams (2006)](http://www.gaussianprocess.org/gpml/chapters/RW.pdf) and references therein):

$$
\begin{align}
y_\star | y \sim \mathcal{N}\left(
    \mathbf{\Sigma}(t_\star, t) \mathbf{\Sigma}(t, t)^{-1}y,
    \mathbf{\Sigma}(t_\star, t_\star) -
        \mathbf{\Sigma}(t_\star, t) \mathbf{\Sigma}(t, t)^{-1} \mathbf{\Sigma}(t, t_\star)
\right)
\end{align}
$$

In other words, given $y$, the distribution for $y_\star$ is *still Gaussian*, but this time with mean equal to
$$
\begin{align}
    \mathbf{\Sigma}(t_\star, t) \mathbf{\Sigma}(t, t)^{-1}y
\end{align}
$$
and covariance
$$
\begin{align}
    \mathbf{\Sigma}(t_\star, t_\star) - \mathbf{\Sigma}(t_\star, t) \mathbf{\Sigma}(t, t)^{-1} \mathbf{\Sigma}(t, t_\star)
\end{align}
$$

Exercise 6

Convince yourself that when $t_\star = t$, the mean and covariance of the GP for $y_\star$ give you what you'd expect. When $t_\star \neq t$, what can you say about the variance of the GP after versus before data was collected?

```python

```

Below, we define a function ``compute_gp`` that returns the mean and covariance of the Gaussian process at the test points (``y_test(t_test)``) conditioned on the values at the training points (``y_train(t_train)``). The ``**kwargs`` are for additional optional keyword arguments passed directly to the kernel (in this simple case, ``A`` and ``l``).

```python
def compute_gp(t_train, y_train, t_test, sigma=0, **kwargs):
    """

    """
    # Compute the required matrices
    Stt = ExpSquaredCovariance(t_train, **kwargs)
    Stt += sigma ** 2 * np.eye(Stt.shape[0])
    Spp = ExpSquaredCovariance(t_test, **kwargs)
    Spt = ExpSquaredCovariance(t_test, tprime=t_train, **kwargs)

    # Compute the mean and covariance of the GP
    mu = np.dot(Spt, np.linalg.solve(Stt, y_train))
    S = Spp - np.dot(Spt, np.linalg.solve(Stt, Spt.T))

    return mu, S
```

Don't worry about ``sigma`` yet -- we'll discuss this below.

Exercise 7

Given

```python
t_train = np.array([2.5, 7.5])
y_train = np.array([1.0, -1.0])
t_test = np.linspace(0, 10, 1000)
```

compute the mean and covariance of the GP. Draw ``30`` samples from it as we did above. Plot the mean function and shade the 1- and 2-$\sigma$ levels. Comment on the behavior of the GP. Does the GP correctly predict the training points? How does the variance behave close to and far from the training points? How does the GP behave very far from the training points?

```python

```

The last thing we'll do regarding drawing from a GP is to account for noise in the observations. If the measurement error on the data points is Gaussian and uncorrelated (white), as we usually assume it to be, we simply add a term to the joint distribution of the training and test data:

$$
\begin{align}
\begin{pmatrix}
 y \\
 y_\star
\end{pmatrix}
&=
\mathcal{N}
\left[
\mathbf{0},
\begin{pmatrix}
 \mathbf{\Sigma}(t, t) + \sigma_n^2 \mathbf{I} & \mathbf{\Sigma}(t, t_\star)\\
 \mathbf{\Sigma}(t_\star, t) & \mathbf{\Sigma}(t_\star, t_\star)
\end{pmatrix}
\right]
\end{align}
$$

where $\sigma_n$ is the standard deviation for the $n^\mathrm{th}$ data point in the training set and $\mathbf{I}$ is the identiy matrix. If you look at our function ``compute_gp`` above, you'll see that it accepts a ``sigma`` keyword, and that it adds its square to the $\mathbf{\Sigma}(t, t)$ term in the covariance, as shown above.

Exercise 8

Re-plot the figure from **Exercise 7**, this time given an observational uncertainty $\sigma = 0.25$ at both training points. Comment on the new behavior of the GP.

```python

```

## GP optimization

In the previous section, we learned how to construct and sample from a simple GP. This is useful for making predictions, i.e., interpolating or extrapolating based on the data you measured. But the true power of GPs comes from their application to *regression* and *inference*: given a dataset $D$ and a model $M(\theta)$, what are the values of the model parameters $\theta$ that are consistent with $D$?

A very common use of GPs is to model things you don't have an explicit physical model for, so quite often they are used to model "nuisances" in the dataset. But just because you don't care about these nuisances doesn't mean they don't affect your inference: in fact, unmodelled correlated noise can often lead to strong biases in the parameter values you infer. In this notebook, we'll learn how to compute likelihoods of Gaussian Processes so that we can *marginalize* over the nuisance parameters (given suitable priors) and obtain unbiased estimates for the physical parameters we care about.

Let's return to the definition of the probability density function for a multivariate Gaussian:

$$
p(\mathbf{y} | \boldsymbol{\mu}, \mathbf{\Sigma}) =
\frac{1}{(2\pi)^N \mathrm{det}(\mathbf{\Sigma})}
\exp\left[
-\frac{1}{2}
(\mathbf{y} - \boldsymbol{\mu})^\top
\mathbf{\Sigma}^{-1}
(\mathbf{y} - \boldsymbol{\mu})
\right]
$$

It's usually easier to work in log-probability space since probabilities can typically be *very* small. Let's therefore take the log of this:

$$
\begin{align}
    \ln p(\mathbf{y} | \boldsymbol{\mu}, \mathbf{\Sigma}) = -\frac{1}{2}(\mathbf{y} - \boldsymbol{\mu})^\top \mathbf{\Sigma}^{-1} (\mathbf{y} - \boldsymbol{\mu}) - \frac{1}{2}\ln \mathrm{det}\,\mathbf{\Sigma} - \frac{N}{2} \ln 2\pi
\end{align}
$$

We can now define a function to compute this log probability given our Squared Exponential covariance:

```python
def ln_gp_likelihood(t, y, sigma=0, A=1.0, l=1.0):
    """
    Return the log of the GP likelihood for a datatset y(t)
    with uncertainties sigma, modeled with a Squared Exponential
    Kernel with amplitude A and lengthscale l.

    """
    # The covariance and its determinant
    npts = len(t)
    K = ExpSquaredCovariance(t, A=A, l=l) + sigma ** 2 * np.eye(npts)

    # The log marginal likelihood
    log_like = -0.5 * np.dot(y.T, np.linalg.solve(K, y))
    log_like -= 0.5 * np.linalg.slogdet(K)[1]
    log_like -= 0.5 * npts * np.log(2 * np.pi)

    return log_like
```

Since this is the probability of our data conditioned on our model, it is often referred to as a *likelihood*. Specifically, it is a *marginal likelihood*, since it's actually the likelihood of the data *marginalized* (integrated) over all the infinitely many curves in the family of functions defined by our GP.

Exercise 9

Consider the following dataset:

```python
import matplotlib.pyplot as plt
t, y, sigma = np.loadtxt("data/sample_data.txt", unpack=True)
plt.plot(t, y, "k.", alpha=0.5, ms=3)
plt.xlabel("time")
plt.ylabel("data");
```

Compute the log GP likelihood for this data, conditioned on different values of the lengthscale of the GP. Specifically, compute it for each of the following values:

```python
l = np.linspace(0.2, 0.4, 300)
```

Assume you know the amplitude `A` to be unity (the default), and use an observational uncertainty `sigma` of `0.05`. Then plot the likelihood as a function of $l$. To compute the likelihood from the log likelihood, do

```python
like = np.exp(lnlike - np.max(lnlike))
```

This will normalize things so that the maximum likelihood is unity. (If you don't do this, you might run into floating point underflow errors, since the numbers we're dealing with are extremely small!)

Comment on your results. The true timescale of the dataset is $l = 0.3$. Were you able to correctly infer that?

```python

```

You just performed your first GP inference problem! It was quite simple, since we assumed the data could be solely modeled with a GP. In practice, we usually have a base model (usually something that depends on physics, like a transit model) whose parameters we want to learn about; the GP is included as an additional model to capture the nuisance, correlated noise that's standing in the way of our measurement. We'll consider a problem like that in the next section.

## Inference with a GP

The timeseries below was generated by a linear function of time, $y(t)= mt + b$. In addition to observational uncertainty $\sigma$ (white noise), there is a fair bit of correlated (red) noise, which we will assume is well described
by the squared exponential covariance with a certain (unknown) amplitude $A$ and timescale $l$.

```python
t, y, sigma = np.loadtxt("data/sample_data_line.txt", unpack=True)
m_true, b_true, A_true, l_true = np.loadtxt("data/sample_data_line_truths.txt", unpack=True)
plt.errorbar(t, y, yerr=sigma, fmt="k.", label="observed")
plt.plot(t, m_true * t + b_true, color="C0", label="truth")
plt.legend(fontsize=12)
plt.xlabel("time")
plt.ylabel("data");
```

Our task is to estimate the values of $m$ and $b$, the slope and intercept of the line, respectively. Initially, we are going to **assume there is no correlated noise.** Our model for the $n^\mathrm{th}$ datapoint is thus

$$
\begin{align}
    y_n \sim \mathcal{N}(m t_n + b, \sigma_n\mathbf{I})
\end{align}
$$

and the probability of the data given the model can be computed by calling our GP log-likelihood function:

```python
def lnprob(params):
    m, b = params
    model = m * t + b
    return ln_gp_likelihood(t, y - model, sigma, A=0, l=1)
```

Note, importantly, that we are passing the **residual vector**, $y - (mt + b)$, to the GP, since above we coded up a zero-mean Gaussian process. We are therefore using the GP to model the **residuals** of the data after applying our physical model (the equation of the line).

To estimate the values of $m$ and $b$ we could generate a fine grid in those two parameters and compute the likelihood at every point, as we did above. But since we'll soon be fitting for four parameters (in the next part), we might as well upgrade our inference scheme and use the ``emcee`` package to do Markov Chain Monte Carlo (MCMC). If you haven't used ``emcee`` before, check out the first few tutorials on the [documentation page](https://emcee.readthedocs.io/en/latest/). The basic setup for the problem is this:

```python
import emcee
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob)

print("Running burn-in...")
p0, _, _ = sampler.run_mcmc(p0, nburn)   # nburn = 500 should do
sampler.reset()

print("Running production...")
sampler.run_mcmc(p0, nsteps);            # nsteps = 1000 should do
```

where ``nwalkers`` is the number of walkers (something like 20 or 30 is fine), ``ndim`` is the number of dimensions (2 in this case), and ``lnprob`` is the log-probability function for the data given the model. Finally, ``p0`` is a list of starting positions for each of the walkers. We're going to pick eyeballed values for $m$ and $b$, then add a small random number to each to generate different initial positions for each walker. This will initialize all walkers in a ball centered on some point, and as the chain progresses they'll diffuse out and begin to explore the posterior.

```python
def lnprob(p):
    """

    """
    m, b = p
    if (m < 0) or (m > 10):
        return -np.inf
    elif (b < 0) or (b > 30):
        return -np.inf
    model = m * t + b
    lnlike = ln_gp_likelihood(t, y - model, sigma, A=0, l=1)
    return lnlike
```

```python
import emcee
print("Using emcee version {0}".format(emcee.__version__))

initial = [4.0, 15.0]
ndim = len(initial)
nwalkers = 32
p0 = initial + 1e-3 * np.random.randn(nwalkers, ndim)
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob)

print("Running burn-in...")
p0, _, _ = sampler.run_mcmc(p0, 500, progress=True)
sampler.reset()

print("Running production...")
sampler.run_mcmc(p0, 1000, progress=True);
```

Once the chain finishes running, we're going to plot several draws from it on top of the data. We'll also plot the **true** line that generated the dataset (given by the variables ``m_true`` and ``b_true``).

```python
# Plot the data
plt.errorbar(t, y, yerr=sigma, fmt=".k", capsize=0)

# The positions where the prediction should be computed
x = np.linspace(0, 10, 500)

# Plot 24 posterior samples
samples = sampler.flatchain
for s in samples[np.random.randint(len(samples), size=24)]:
    m, b = s
    model = m * x + b
    plt.plot(x, model, color="#4682b4", alpha=0.3)

# Plot the truth
plt.plot(x, m_true * x + b_true, "C1", label="truth")

plt.ylabel("data")
plt.xlabel("time")
plt.title("fit assuming uncorrelated noise");
```

Let's also plot the corner plot to see how well we inferred the slope and the intercept:

```python
import corner
labels = ["slope", "intercept"]
truths = [m_true, b_true]
corner.corner(sampler.flatchain, truths=truths, labels=labels);
```

Not great! Our slope is actually OK, but our intercept is quite wrong -- we underestimated it by more than 5 sigma!

Exercise 10

This time, let's actually model the correlated noise. Re-define the ``lnprob`` function above to accept four parameters (slope, intercept, amplitude, and timescale). Restrict the fit to reasonable values of the amplitude (say $0 < A < 10$ and the lengthscale ($0 < l < 10$) by returning `-np.inf` from your log prior function for samples outside that range.

You'll probably want to run your chains for a bit longer this time, too. As before, plot some posterior samples for the line, as well as the corner plot. How did you do this time? Is there any bias in your inferred values? How does the variance compare to the previous estimate?

```python

```


http://nbviewer.jupyter.org/github/pydy/pydy/blob/master/examples/npendulum/n-pendulum-control.ipynb

http://www.moorepants.info/blog/npendulum.html

https://www.youtube.com/watch?v=cyN-CRNrb3E

# Introduction

Several pieces of the puzzle have come together lately to really demonstrate the power of the scientific python software packages to handle complex dynamic and controls problems (i.e. IPython notebooks, matplotlib animations, python-control, and our software packages: sympy.physics.mechanics and PyDy).

This [blog post by Wolfram](http://blog.wolfram.com/2011/03/01/stabilized-n-link-pendulum/) demonstrates Mathematica's ability to symbolically derive the equations of motion for the n-link pendulum and stabilize it with an LQR controller. This blog post inspired us to replicate the example with all free and open source software.

In this example problem, we derive the equations of motion of an n-link pendulum on a laterally sliding cart and then develop a controller to stabilize it. Balancing a single inverted pendulum is a classic problem that is often a student's first experience with non-linear dynamics and control. The problem here is extended to a general n-link pendulum in which the equations of motion quickly get messy with greater than 2 links.

The diagram below shows the general description of the problem.

```python
from IPython.display import SVG
SVG(filename='n-pendulum-with-cart.svg')
```

Setup
===

This example depends on the following software:

- IPython
- NumPy
- SciPy
- SymPy >= 0.7.6
- matplotlib
- python-control
- slycot (used by python-control)

The easiest way to install the Python packages it is to use conda:

```
$ conda install ipython-notebook numpy scipy sympy matplotlib
$ conda install -c https://conda.binstar.org/cwrowley control slycot
```

You may need to install `slycot` from source, refer to https://github.com/avventi/Slycot

To create animations you need a video encoder like ffmpeg or avconv installed.

Equations of Motion
===================

We'll start by generating the equations of motion for the system with SymPy **[mechanics](http://docs.sympy.org/dev/modules/physics/mechanics/index.html)**. The functionality that mechanics provides is much more in depth than Mathematica's functionality. In the Mathematica example, Lagrangian mechanics were implemented manually with Mathematica's symbolic functionality. **mechanics** provides an assortment of functions and classes to derive the equations of motion for arbitrarily complex (i.e. configuration constraints, nonholonomic motion constraints, etc) multibody systems in a very natural way. First we import the necessary functionality from SymPy.

```python
from __future__ import division, print_function
```

```python
import sympy as sm
import sympy.physics.mechanics as me
```

We can enable mathematical rendering of the resulting equations in the notebook with the following command.

```python
me.init_vprinting()
```

Now specify the number of links, $n$. I'll start with 5 since the Wolfram folks only showed four.

```python
n = 5
```

**mechanics** will need the generalized coordinates, generalized speeds, and the input force which are all time dependent variables and the bob masses, link lengths, and acceleration due to gravity which are all constants. Time, $t$, is also made available because we will need to differentiate with respect to time.

```python
q = me.dynamicsymbols('q:{}'.format(n + 1))  # Generalized coordinates
u = me.dynamicsymbols('u:{}'.format(n + 1))  # Generalized speeds
f = me.dynamicsymbols('f')                   # Force applied to the cart

m = sm.symbols('m:{}'.format(n + 1))         # Mass of each bob
l = sm.symbols('l:{}'.format(n))             # Length of each link
g, t = sm.symbols('g t')                     # Gravity and time
```

Now we can create and inertial reference frame $I$ and define the point, $O$, as the origin.

```python
I = me.ReferenceFrame('I')  # Inertial reference frame
O = me.Point('O')           # Origin point
O.set_vel(I, 0)             # Origin's velocity is zero
```

Secondly, we define the define the first point of the pendulum as a particle which has mass. This point can only move laterally and represents the motion of the "cart".

```python
P0 = me.Point('P0')                 # Hinge point of top link
P0.set_pos(O, q[0] * I.x)           # Set the position of P0
P0.set_vel(I, u[0] * I.x)           # Set the velocity of P0
Pa0 = me.Particle('Pa0', P0, m[0])  # Define a particle at P0
```

Now we can define the $n$ reference frames, particles, gravitational forces, and kinematical differential equations for each of the pendulum links. This is easily done with a loop.

```python
frames = [I]                              # List to hold the n + 1 frames
points = [P0]                             # List to hold the n + 1 points
particles = [Pa0]                         # List to hold the n + 1 particles
forces = [(P0, f * I.x - m[0] * g * I.y)] # List to hold the n + 1 applied forces, including the input force, f
kindiffs = [q[0].diff(t) - u[0]]          # List to hold kinematic ODE's

for i in range(n):
    Bi = I.orientnew('B' + str(i), 'Axis', [q[i + 1], I.z])   # Create a new frame
    Bi.set_ang_vel(I, u[i + 1] * I.z)                         # Set angular velocity
    frames.append(Bi)                                         # Add it to the frames list

    Pi = points[-1].locatenew('P' + str(i + 1), l[i] * Bi.x)  # Create a new point
    Pi.v2pt_theory(points[-1], I, Bi)                         # Set the velocity
    points.append(Pi)                                         # Add it to the points list

    Pai = me.Particle('Pa' + str(i + 1), Pi, m[i + 1])        # Create a new particle
    particles.append(Pai)                                     # Add it to the particles list

    forces.append((Pi, -m[i + 1] * g * I.y))                  # Set the force applied at the point

    kindiffs.append(q[i + 1].diff(t) - u[i + 1])              # Define the kinematic ODE:  dq_i / dt - u_i = 0
```

With all of the necessary point velocities and particle masses defined, the `KanesMethod` class can be used to derive the equations of motion of the system automatically.

```python
kane = me.KanesMethod(I, q_ind=q, u_ind=u, kd_eqs=kindiffs)  # Initialize the object
fr, frstar = kane.kanes_equations(particles, forces)         # Generate EoM's fr + frstar = 0
```

The equations of motion are quite long as can been seen below. This is the general nature of most non-simple mutlibody problems. That is why a SymPy is so useful; no more mistakes in algebra, differentiation, or copying hand written equations. Note that `trigsimp` can take quite a while to complete for extremely large expressions. Below we print $\tilde{M}$ and $\tilde{f}$ from $\tilde{M}\dot{u}=\tilde{f}$ to show the size of the expressions.

```python
sm.trigsimp(kane.mass_matrix)
```

$\tilde{M}$ is a function of the constant parameters and the configuration.

```python
me.find_dynamicsymbols(kane.mass_matrix)
```

```python
sm.trigsimp(kane.forcing)
```

$\tilde{f}$ is a function of the constant parameters, configuration, speeds, and the applied force.

```python
me.find_dynamicsymbols(kane.forcing)
```

Simulation
==========

Now that the symbolic equations of motion are available we can simulate the pendulum's motion. We will need some more SymPy functionality and several NumPy functions, and most importantly the integration function from SciPy, `odeint`.

```python
import numpy as np
from numpy.linalg import solve
from scipy.integrate import odeint
```

First, define some numeric values for all of the constant parameters in the problem.

```python
arm_length = 1. / n                          # The maximum length of the pendulum is 1 meter
bob_mass = 0.01 / n                          # The maximum mass of the bobs is 10 grams
parameters = [g, m[0]]                       # Parameter definitions starting with gravity and the first bob
parameter_vals = [9.81, 0.01 / n]            # Numerical values for the first two
for i in range(n):                           # Then each mass and length
    parameters += [l[i], m[i + 1]]
    parameter_vals += [arm_length, bob_mass]
```

Mathematica has a really nice `NDSolve` function for quickly integrating their symbolic differential equations. We make use of SymPy's lambdify function to do something similar, i.e. to create functions that will evaluate the "full" mass matrix, $M$, and "full" forcing vector, $f$ from $M\dot{x} = f(x, r, t)$ as a NumPy function.

```python
dynamic = q + u    # Make a list of the states
dynamic.append(f)  # Add the input force

M_func = sm.lambdify(dynamic + parameters, kane.mass_matrix_full) # Create a callable function to evaluate the mass matrix
f_func = sm.lambdify(dynamic + parameters, kane.forcing_full)     # Create a callable function to evaluate the forcing vector
```

To integrate the ODE's we need to define a function that returns the derivatives of the states given the current state and time.

```python
def right_hand_side(x, t, args):
    """Returns the derivatives of the states.

    Parameters
    ----------
    x : ndarray, shape(2 * (n + 1))
        The current state vector.
    t : float
        The current time.
    args : ndarray
        The constants.

    Returns
    -------
    dx : ndarray, shape(2 * (n + 1))
        The derivative of the state.

    """
    r = 0.0                                  # The input force is always zero
    arguments = np.hstack((x, r, args))      # States, input, and parameters
    dx = np.array(solve(M_func(*arguments),  # Solving for the derivatives
                  f_func(*arguments))).T[0]

    return dx
```

Now that we have the right hand side function, the initial conditions are set such that the pendulum is in the vertical equilibrium and a slight initial rate is set for each speed to ensure the pendulum falls. The equations can then be integrated with SciPy's `odeint` function given a time series.

```python
x0 = np.hstack((0.0,                              # q0
                np.pi / 2 * np.ones(len(q) - 1),  # q1...qn+1
                1e-3 * np.ones(len(u))))          # u0...un+1
```

```python
t = np.linspace(0.0, 10.0, num=500)  # Time vector
```

```python
x = odeint(right_hand_side, x0, t, args=(parameter_vals,))  # Numerical integration
```

Plotting
========

The results of the simulation can be plotted with matplotlib. First, load the plotting functionality.

```python
import matplotlib.pyplot as plt
```

```python
%matplotlib inline
```

```python
from IPython.core.pylabtools import figsize
figsize(8.0, 6.0)
```

The coordinate trajectories are plotted below.

```python
lines = plt.plot(t, x[:, :x.shape[1] // 2])
lab = plt.xlabel('Time [sec]')
leg = plt.legend(dynamic[:x.shape[1] // 2])
```

And the generalized speed trajectories.

```python
lines = plt.plot(t, x[:, x.shape[1] // 2:])
lab = plt.xlabel('Time [sec]')
leg = plt.legend(dynamic[x.shape[1] // 2:])
```

Animation
=========

matplotlib now includes very nice animation functions for animating matplotlib plots. First we import the necessary functions for creating the animation.

```python
from matplotlib import animation
from matplotlib.patches import Rectangle
```

The following function was modeled from Jake Vanderplas's [post on matplotlib animations](http://jakevdp.github.com/blog/2012/08/18/matplotlib-animation-tutorial/). The default animation writer is used (typically ffmpeg), you can change it by adding `writer` argument to `anim.save` call.

```python
def animate_pendulum(t, states, length, filename=None):
    """Animates the n-pendulum and optionally saves it to file.

    Parameters
    ----------
    t : ndarray, shape(m)
        Time array.
    states: ndarray, shape(m,p)
        State time history.
    length: float
        The length of the pendulum links.
    filename: string or None, optional
        If true a movie file will be saved of the animation. This may take some time.

    Returns
    -------
    fig : matplotlib.Figure
        The figure.
    anim : matplotlib.FuncAnimation
        The animation.

    """
    # the number of pendulum bobs
    numpoints = states.shape[1] // 2

    # first set up the figure, the axis, and the plot elements we want to animate
    fig = plt.figure()

    # some dimesions
    cart_width = 0.4
    cart_height = 0.2

    # set the limits based on the motion
    xmin = np.around(states[:, 0].min() - cart_width / 2.0, 1)
    xmax = np.around(states[:, 0].max() + cart_width / 2.0, 1)

    # create the axes
    ax = plt.axes(xlim=(xmin, xmax), ylim=(-1.1, 1.1), aspect='equal')

    # display the current time
    time_text = ax.text(0.04, 0.9, '', transform=ax.transAxes)

    # create a rectangular cart
    rect = Rectangle([states[0, 0] - cart_width / 2.0, -cart_height / 2],
                     cart_width, cart_height, fill=True, color='red',
                     ec='black')
    ax.add_patch(rect)

    # blank line for the pendulum
    line, = ax.plot([], [], lw=2, marker='o', markersize=6)

    # initialization function: plot the background of each frame
    def init():
        time_text.set_text('')
        rect.set_xy((0.0, 0.0))
        line.set_data([], [])
        return time_text, rect, line,

    # animation function: update the objects
    def animate(i):
        time_text.set_text('time = {:2.2f}'.format(t[i]))
        rect.set_xy((states[i, 0] - cart_width / 2.0, -cart_height / 2))
        x = np.hstack((states[i, 0], np.zeros((numpoints - 1))))
        y = np.zeros((numpoints))
        for j in np.arange(1, numpoints):
            x[j] = x[j - 1] + length * np.cos(states[i, j])
            y[j] = y[j - 1] + length * np.sin(states[i, j])
        line.set_data(x, y)
        return time_text, rect, line,

    # call the animator function
    anim = animation.FuncAnimation(fig, animate, frames=len(t), init_func=init,
            interval=t[-1] / len(t) * 1000, blit=True, repeat=False)

    # save the animation if a filename is given
    if filename is not None:
        anim.save(filename, fps=30, codec='libx264')
```

Now we can create the animation of the pendulum. This animation will show the open loop dynamics.

```python
animate_pendulum(t, x, arm_length, filename="open-loop.mp4")
```

```python
from IPython.display import HTML
html = \
"""

"""
HTML(html)
```

Controller Design
=================

The n-link pendulum can be balanced such that all of the links are inverted above the cart by applying the correct lateral force to the cart. We can design a full state feedback controller based from a linear model of the pendulum about its upright equilibrium point. We'll start by specifying the equilibrium point and parameters in dictionaries. We make sure to use SymPy types in the equilibrium point to ensure proper cancelations in the linearization.

```python
equilibrium_point = [sm.S(0)] + [sm.pi / 2] * (len(q) - 1) + [sm.S(0)] * len(u)
equilibrium_dict = dict(zip(q + u, equilibrium_point))
equilibrium_dict
```

The `KanesMethod` class has method that linearizes the forcing vector about generic state and input perturbation vectors. The equilibrium point and numerical constants can then be substituted in to give the linear system in this form: $M\dot{x}=F_Ax+F_Br$. The state and input matrices, $A$ and $B$, can then be computed by left side multiplication by the inverse of the mass matrix: $A=M^{-1}F_A$ and $B=M^{-1}F_B$.

```python
M, F_A, F_B, r = kane.linearize(new_method=True, op_point=equilibrium_dict)
```

```python
sm.simplify(M)
```

```python
sm.simplify(F_A)
```

```python
sm.simplify(F_B)
```

Now the numerical $A$ and $B$ matrices can be formed. First substitute numerical parameter values into $M$, $F_A$, and $F_B$.

```python
parameter_dict = dict(zip(parameters, parameter_vals))
parameter_dict
```

```python
M_num = sm.matrix2numpy(M.subs(parameter_dict), dtype=float)
F_A_num = sm.matrix2numpy(F_A.subs(parameter_dict), dtype=float)
F_B_num = sm.matrix2numpy(F_B.subs(parameter_dict), dtype=float)
```

```python
A = np.linalg.solve(M_num, F_A_num)
B = np.linalg.solve(M_num ,F_B_num)
```

```python
print(A)
```

    [[  0.00000000e+00   0.00000000e+00   0.00000000e+00   0.00000000e+00
        0.00000000e+00   0.00000000e+00   1.00000000e+00   0.00000000e+00
        0.00000000e+00   0.00000000e+00   0.00000000e+00   0.00000000e+00]
     [  0.00000000e+00   0.00000000e+00   0.00000000e+00   0.00000000e+00
        0.00000000e+00   0.00000000e+00   0.00000000e+00   1.00000000e+00
        0.00000000e+00   0.00000000e+00   0.00000000e+00   0.00000000e+00]
     [  0.00000000e+00   0.00000000e+00   0.00000000e+00   0.00000000e+00
        0.00000000e+00   0.00000000e+00   0.00000000e+00   0.00000000e+00
        1.00000000e+00   0.00000000e+00   0.00000000e+00   0.00000000e+00]
     [  0.00000000e+00   0.00000000e+00   0.00000000e+00   0.00000000e+00
        0.00000000e+00   0.00000000e+00   0.00000000e+00   0.00000000e+00
        0.00000000e+00   1.00000000e+00   0.00000000e+00   0.00000000e+00]
     [  0.00000000e+00   0.00000000e+00   0.00000000e+00   0.00000000e+00
        0.00000000e+00   0.00000000e+00   0.00000000e+00   0.00000000e+00
        0.00000000e+00   0.00000000e+00   1.00000000e+00   0.00000000e+00]
     [  0.00000000e+00   0.00000000e+00   0.00000000e+00   0.00000000e+00
        0.00000000e+00   0.00000000e+00   0.00000000e+00   0.00000000e+00
        0.00000000e+00   0.00000000e+00   0.00000000e+00   1.00000000e+00]
     [ -0.00000000e+00   4.90500000e+01  -1.38777878e-14   1.06974614e-14
       -3.32488666e-15  -3.61400724e-17  -0.00000000e+00  -0.00000000e+00
       -0.00000000e+00  -0.00000000e+00  -0.00000000e+00  -0.00000000e+00]
     [ -0.00000000e+00   4.90500000e+02  -1.96200000e+02   7.80625564e-14
       -1.30104261e-14   5.42101086e-16  -0.00000000e+00  -0.00000000e+00
       -0.00000000e+00  -0.00000000e+00  -0.00000000e+00  -0.00000000e+00]
     [ -0.00000000e+00  -2.45250000e+02   3.92400000e+02  -1.47150000e+02
       -0.00000000e+00  -2.03287907e-15  -0.00000000e+00  -0.00000000e+00
       -0.00000000e+00  -0.00000000e+00  -0.00000000e+00  -0.00000000e+00]
     [ -0.00000000e+00   2.16840434e-13  -1.96200000e+02   2.94300000e+02
       -9.81000000e+01   7.22801448e-15  -0.00000000e+00  -0.00000000e+00
       -0.00000000e+00  -0.00000000e+00  -0.00000000e+00  -0.00000000e+00]
     [ -0.00000000e+00  -1.08420217e-13   6.50521303e-14  -1.47150000e+02
        1.96200000e+02  -4.90500000e+01  -0.00000000e+00  -0.00000000e+00
       -0.00000000e+00  -0.00000000e+00  -0.00000000e+00  -0.00000000e+00]
     [ -0.00000000e+00  -0.00000000e+00  -0.00000000e+00  -0.00000000e+00
       -9.81000000e+01   9.81000000e+01  -0.00000000e+00  -0.00000000e+00
       -0.00000000e+00  -0.00000000e+00  -0.00000000e+00  -0.00000000e+00]]

```python
print(B)
```

    [[  0.00000000e+00]
     [  0.00000000e+00]
     [  0.00000000e+00]
     [  0.00000000e+00]
     [  0.00000000e+00]
     [  0.00000000e+00]
     [  5.00000000e+02]
     [  2.50000000e+03]
     [ -8.67361738e-13]
     [  8.67361738e-13]
     [ -4.33680869e-13]
     [ -0.00000000e+00]]

Also convert `equilibrium_point` to a numeric array:

```python
equilibrium_point = np.asarray([x.evalf() for x in equilibrium_point], dtype=float)
```

Now that we have a linear system, the python-control package can be used to design an optimal controller for the system.

```python
import control
from numpy.linalg import matrix_rank
```

First we can check to see if the system is, in fact, controllable. The rank of the controllability matrix must be equal to the number of rows in $A$, but the `matrix_rank` algorithm is numerically ill conditioned and for certain values of $n$ this will fail, as seen below for $n=5$. Nevertheless, the system is controllable, no matter the number of links.

```python
matrix_rank(control.ctrb(A, B)) == A.shape[0]
```

    False

So now we can compute the optimal gains with a linear quadratic regulator. I chose identity matrices for the weightings for simplicity.

```python
K, X, E = control.lqr(A, B, np.ones(A.shape), 1);
```

The gains can now be used to define the required input during simulation to stabilize the system. The input $r$ is simply the gain vector multiplied by the error in the state vector from the equilibrium point, $r(t)=K(x_{eq} - x(t))$.

```python
def right_hand_side(x, t, args):
    """Returns the derivatives of the states.

    Parameters
    ----------
    x : ndarray, shape(2 * (n + 1))
        The current state vector.
    t : float
        The current time.
    args : ndarray
        The constants.

    Returns
    -------
    dx : ndarray, shape(2 * (n + 1))
        The derivative of the state.

    """
    r = np.dot(K, equilibrium_point - x)    # The controller
    arguments = np.hstack((x, r, args))     # States, input, and parameters
    dx = np.array(solve(M_func(*arguments), # Solving for the derivatives
                        f_func(*arguments))).T[0]

    return dx
```

Now we can simulate and animate the system to see if the controller works.

```python
x0 = np.hstack((0,
                np.pi / 2 * np.ones(len(q) - 1),
                1 * np.ones(len(u))))
t = np.linspace(0.0, 10.0, num=500)
x = odeint(right_hand_side, x0, t, args=(parameter_vals,))
```

The plots show that we seem to have a stable system.

```python
lines = plt.plot(t, x[:, :x.shape[1] // 2])
lab = plt.xlabel('Time [sec]')
leg = plt.legend(dynamic[:x.shape[1] // 2])
```

```python
lines = plt.plot(t, x[:, x.shape[1] // 2:])
lab = plt.xlabel('Time [sec]')
leg = plt.legend(dynamic[x.shape[1] // 2:])
```

```python
animate_pendulum(t, x, arm_length, filename="closed-loop.mp4")
```

```python
from IPython.display import HTML
html = \
"""

"""
HTML(html)
```

The video clearly shows that the controller can balance all $n$ of the pendulum links. The weightings in the lqr design can be tweaked to give different performance if needed.

This example shows that the free and open source scientific Python tools for dynamics are easily comparable in ability and quality to a commercial package such as Mathematica.

The IPython notebook for this example can be downloaded from https://github.com/pydy/pydy/tree/master/examples/npendulum. You can try out different $n$ values. I've gotten the equations of motion to compute for an open loop simulation of 10 links. My computer ran out of memory when I tried to compute for $n=50$. The controller weightings and initial conditions will probably have to be adjusted for better performance for $n>5$, but it should work.

```python
# Install with pip install version_information
%load_ext version_information
%version_information numpy, sympy, scipy, matplotlib, control
```

<table><tr><th>Software</th><th>Version</th></tr><tr><td>Python</td><td>3.4.4 64bit [GCC 4.2.1 (Apple Inc. build 5577)]</td></tr><tr><td>IPython</td><td>4.1.2</td></tr><tr><td>OS</td><td>Darwin 15.4.0 x86_64 i386 64bit</td></tr><tr><td>numpy</td><td>1.11.0</td></tr><tr><td>sympy</td><td>1.0.1.dev</td></tr><tr><td>scipy</td><td>0.18.0.dev0+72bae57</td></tr><tr><td>matplotlib</td><td>1.5.1</td></tr><tr><td>control</td><td>0.7.0</td></tr><tr><td colspan='2'>Tue May 17 21:04:08 2016 YEKT</td></tr></table>

```python

```


# COOT on MNIST/USPS example

This notebook gives a simple example of the COOT problem between two well-known datasets: MNIST and USPS.

We recall that for two datasets $\mathbf{X} \in \mathbb{R}^{n \times d}, \mathbf{X'} \in \mathbb{R}^{n' \times d'}$ the COOT problem aims at finding two optimal transport maps $\boldsymbol{\pi^{s}}$ and $\boldsymbol{\pi^{v}}$ between the samples and the features that solve:

\begin{equation}
\underset{\begin{smallmatrix}\boldsymbol{\pi^{s}} \in \Pi(\mathbf{w},\mathbf{w'}) \\ \boldsymbol{\pi^{v}} \in \Pi(\mathbf{v},\mathbf{v'}) \end{smallmatrix}} {\min} \sum_{ijkl} \left|X_{ik}-X_{jl}\right|^{p} \pi^{s}_{ij} \pi^{v}_{kl}
\end{equation}

where $\mathbf{w},\mathbf{w'}$ are histograms on the samples and $\mathbf{v},\mathbf{v'}$ are histogram on the features.

In our example the samples are images and the features are the pixels of the images. We will have $n,n'=300$ and $d=784, d'=256$

```python
import numpy as np
from scipy import ndimage
import scipy as sp
import matplotlib.pylab as pl
import ot
import scipy.io
import sys
sys.path.append('../code')
import cot
```

### Load the MNIST/USPS dataset

```python
data=sp.io.loadmat('../data/mnist.mat')

Xtot1=data['xapp'].astype(np.float32)

Ytot1=data['yapp']
d1=Xtot1.shape[1]
Ytot1[Ytot1==10]=0

data=sp.io.loadmat('../data/usps.mat')

Xtot2=(data['xapp'].astype(np.float32)+1)/2

Ytot2=data['yapp']
Ytot2-=1
d2=Xtot2.shape[1]

np.random.seed(1976)
```

```python
def get_data(x,y,nbperclass):

    xr=np.zeros((0,x.shape[1]))
    yr=np.zeros((0))

    for i in range(np.max(y).astype(int)+1):

        xi=x[y.ravel()==i,:]
        idx=np.random.permutation(xi.shape[0])

        xr=np.concatenate((xr,xi[idx[:nbperclass],:]),0)
        yr=np.concatenate((yr,i*np.ones(nbperclass)))

    return xr,yr

#%% We take 300 samples per class
nbperclass=300
xs,ys=get_data(Xtot1,Ytot1,nbperclass)
xs=xs/255

selmnist=xs.sum(0)>0

ntot=nbperclass*10

xs2=np.zeros((xs.shape[0],d1))
#xs2[:,sel1]=xs

xt,ys=get_data(Xtot2,Ytot2,nbperclass)

vs=xs.sum(axis=0) # set the weights on the features
vs/=vs.sum()
vt=xt.sum(axis=0)
vt/=vt.sum()

```

```python
ot.tic()
Ts,Tv,_,log=cot.cot_numpy(xs,xt,v1=vs,v2=vt,niter=100,log=True) # solve COOT
ot.toc()

pl.figure(1,figsize=(6,4))
pl.plot(log['cost'])
pl.title('evolution of cost (no Mass correction)')
pl.show()

ot.tic()
Tsr,Tvr,_,logr=cot.cot_numpy(xs,xt,v1=vs,v2=vt,niter=100,log=True,algo2='sinkhorn',reg2=.5e-2) # solve COOT with sinkhorn
ot.toc()

```

```python
pl.figure(2,figsize=(6,6))
pl.imshow(Ts)
pl.colorbar()
pl.show()
```

```python
Tv.shape,Ts.shape
```

    ((784, 256), (3000, 3000))

### Confusion matrix on the samples

We evaluate the COOT ability to find the good assignments of the images (samples), i.e. if it aligns well the samples classes between the two datasets based on the knowledge of $\boldsymbol{\pi^{s}}$

```python
#%% confusion matrix

nbc=10

Cmat=np.zeros((nbc,nbc))
for i in range(ntot):
    #print(i)
    for j in range(ntot):
        if Ts[i,j]:
            Cmat[int(ys[i]),int(ys[j])]+=Ts[i,j]

print('Find the good class in {:.2f}% '.format(100*np.sum(np.diag(Cmat))))

#%%
pl.imshow(Cmat*10), pl.colorbar()
pl.title('Confusion matrix for COOT between samples')
pl.ylabel('Labels MNIST')
pl.xlabel('Labels USPS')

```

### Vizualize the transport on the features

We propose to vizualize the optimal oupling on the features $\boldsymbol{\pi^{v}}$. In order to do that we color code each pixel of an image of size USPS and we transfer to an image of size MNIST.

```python
#%%pix
dim_source=16
dim_target=28

image = np.zeros((dim_source,dim_source,3))

for i in range(dim_source):
    for j in range(dim_source):
        image[i,j,0]=i
        image[i,j,1]=j
        image[i,j,2]=dim_source/2
image=image.astype(np.float32)/dim_source

diag=1./Tv.sum(axis=1)
diag[diag==np.inf]=0
image_target = np.dot(np.diag(diag),np.dot(image.reshape((dim_source*dim_source,3)).T,Tv.T).T)

image_target[~selmnist,:]=np.nan #we remove non informative features

image_target=image_target.reshape((dim_target,dim_target,3))

diagr=1./Tvr.sum(axis=1)
diagr[diagr==np.inf]=0
image_targetr = np.dot(np.diag(diagr),np.dot(image.reshape((dim_source*dim_source,3)).T,Tvr.T).T)

image_targetr[~selmnist,:]=np.nan

image_targetr=image_targetr.reshape((dim_target,dim_target,3))

pl.figure(3,figsize=(16,32))
pl.subplot(1,2,1)
pl.imshow(image)
pl.title('source image')
pl.axis('off')
pl.subplot(1,2,2)
pl.imshow(image_target)
pl.title('Transfered image')
pl.axis('off')
pl.show()

#%%
import scipy.sparse
sTs= scipy.sparse.coo_matrix(Ts)

row=sTs.row
col=sTs.col

pl.figure(10,figsize=(14,3.5))
pl.clf()
pl.subplot(1,4,1)
pl.plot(col,row,'.',markersize=3,alpha=0.5)
#pl.spy(Tv,markersize=3,marker='.',alpha=0.5)

pl.title('$\pi^s$ matrix between samples')
pl.xlabel('USPS samples')
pl.ylabel('MNIST samples')
pl.xticks([300*i for i in range(11)],[' ']*11)
pl.yticks([300*i for i in range(11)],[]*11)
pl.axis('scaled')
pl.xlim((0,ntot))
pl.ylim((ntot,0))
pl.grid()

pl.subplot(1,4,2)
pl.imshow(Cmat*10,cmap='Blues'),
#pl.colorbar()
pl.title('Confusion matrix')
pl.ylabel('Labels MNIST')
pl.xlabel('Labels USPS')
#pl.xticks(*pl.yticks())
pl.yticks([i for i in range(10)],[i for i in range(10)])
pl.ylim((nbc-.5,-.5))
pl.xticks([i for i in range(10)],[i for i in range(10)])
pl.xlim((-.5,nbc-.5,))
pl.subplot(1,4,3)

pl.imshow(image)
pl.title('USPS colored pixels')
pl.axis('off')
pl.xlim([-6,22])
pl.ylim([-6,22])

pl.subplot(1,4,4)
pl.imshow(image_target)
pl.title("MNIST pixels through $\pi^v$")
pl.axis('off')
pl.show()

pl.savefig('./mnist_usps.png')
pl.savefig('./mnist_usps.pdf',bbox_inches='tight')

```

We observe that the spatial structured is preserved (without supervision): the pixel are transported coherently on the center of the image

### Vizualize the images after transformation through the optimal couplings

We can also vizualize the images after transformation via the optimal couplings

```python
#%%

nbl,nbc=5,2

#idx_sel=np.random.randint(0,ntot,n_fig)
idx_sel=np.arange(0,ntot,nbperclass)+3

xts=xt[idx_sel,:]
xss=xs[idx_sel,:]

I=np.zeros((28*nbl,28*(nbc*2+1)))+1
for i in range(nbl):
    for j in range(nbc):
        I[i*28:(i+1)*28,j*28:(j+1)*28]=xss[i+j*nbl].reshape((28,28))

        I[i*28+6:(i)*28+22,j*28+28*(nbc+1)+6:j*28+28*(nbc+1)+22]=xts[i+j*nbl].reshape((16,16))

pl.figure(15)
pl.clf()
pl.imshow(I,cmap='Blues')
pl.axis('off')
pl.ylim([ I.shape[0],-10])
pl.text(20,-7,'MNIST',fontsize=15)
pl.text(20+28*(nbc+1),-7,'USPS',fontsize=15)

```

```python
#%%
import scipy.sparse
sTv= scipy.sparse.coo_matrix(Ts)

row=sTs.row
col=sTs.col

pl.figure(11,figsize=(16,3.5))

pl.clf()

pl.subplot(1,5,1)
pl.imshow(I,cmap='gray')
pl.axis('off')
pl.ylim([ I.shape[0],-10])
pl.text(15,-9,'MNIST',fontsize=12)
pl.text(15+28*(nbc+1),-9,'USPS',fontsize=12)

pl.subplot(1,5,2)
pl.plot(col,row,'.',markersize=3,alpha=0.5)
#pl.spy(Tv,markersize=3,marker='.',alpha=0.5)

pl.title('$\pi^s$ matrix between samples')
pl.xlabel('USPS samples')
pl.ylabel('MNIST samples')
pl.xticks([300*i for i in range(11)],[' ']*11)
pl.yticks([300*i for i in range(11)],[]*11)
pl.axis('scaled')
pl.xlim((0,ntot))
pl.ylim((ntot,0))
pl.grid()

pl.subplot(1,5,3)

pl.imshow(image)
pl.title('USPS colored coded pixels')
pl.axis('off')
pl.xlim([-6,22])
pl.ylim([22,-6])

pl.subplot(1,5,4)
pl.imshow(image_target)
pl.title("MNIST pixels through $\pi^v$")
pl.axis('off')
#pl.show()

pl.subplot(1,5,5)
pl.imshow(image_targetr)
pl.title("MNIST pixels through entropic $\pi^v$")
pl.axis('off')
#pl.show()

pl.savefig('./mnist_usps.png')
pl.savefig('./mnist_usps.pdf',bbox_inches='tight')

#%%
import random
import PIL as pil

# build a rectangle in axes coords
left, width = .25, .5
bottom, height = .25, .5
right = left + width
top = bottom + height

def predict_barycenter(data,T):
    diag=1./T.sum(axis=1)
    diag[diag==np.inf]=0
    return np.dot(T,data.T).T.dot(np.diag(diag))

def predict_barycenter_reverse(data,T):
    diag=1./T.sum(axis=1)
    diag[diag==np.inf]=0
    return np.dot(T,data).T.dot(np.diag(diag)).T

random.seed(1985)
np.random.seed(1976)

n_fig=16

idx_sel=np.random.randint(0,ntot,n_fig)

xsel=xs[idx_sel,:]
xpred=np.zeros((n_fig,d2))
xpredr=np.zeros((n_fig,d2))
for i in range(n_fig):
    xpred[i,:]=predict_barycenter(xsel[i,:],Tv.T)
    xpredr[i,:]=predict_barycenter(xsel[i,:],Tvr.T)

cmap_g='gray'
pl.figure(figsize=(n_fig,4))

for i in range(n_fig):

    ax= pl.subplot(4,n_fig,i+1)

    pl.imshow(xsel[i,:].reshape((28,28)),cmap=cmap_g)
    pl.axis('off')
    pl.xticks(())
    pl.yticks(())
    #pl.xlim([-6,22])
    #pl.ylim([22,-6])
    if i==0:
        ax.text(left-.3, 0.5*(bottom+top), 'MNIST',
        horizontalalignment='right',
        verticalalignment='center',
        rotation='vertical',
        transform=ax.transAxes)
        #pl.ylabel('MNIST')

    pl.subplot(4,n_fig,i+1+n_fig)
    img = pil.Image.fromarray((xsel[i,:]* 255 / np.max(xsel[i,:])).reshape((28,28)) .astype('float32'))
    img = img.resize((16,16))
    pl.imshow(img,cmap=cmap_g)
    pl.axis('off')
    pl.xticks(())
    pl.yticks(())
    pl.xlim([-6,22])
    pl.ylim([22,-6])
    if i==0:
        ax.text(left-.3, 0.5*(bottom+top) - 1.1, 'Resize',
        horizontalalignment='right',
        verticalalignment='center',
        rotation='vertical',
        transform=ax.transAxes)

    ax = pl.subplot(4,n_fig,i+1+2*n_fig)
    ax.imshow(xpred[i,:].reshape((16,16)),cmap=cmap_g)
    pl.axis('off')
    pl.xticks(())
    pl.yticks(())
    pl.xlim([-6,22])
    pl.ylim([22,-6])
    if i==0:
        ax.text(left-.3, 0.5*(bottom+top), 'Map $\pi^v$',
        horizontalalignment='right',
        verticalalignment='center',
        rotation='vertical',
        transform=ax.transAxes)

    ax= pl.subplot(4,n_fig,i+1+3*n_fig)
    pl.imshow(xpredr[i,:].reshape((16,16)),cmap=cmap_g)
    pl.axis('off')
    pl.xticks(())
    pl.yticks(())
    pl.xlim([-6,22])
    pl.ylim([22,-6])
    if i==0:
        ax.text(left-.3, 0.5*(bottom+top), 'Map reg $\pi^v$',
        horizontalalignment='right',
        verticalalignment='center',
        rotation='vertical',
        transform=ax.transAxes)
        #pl.ylabel('Map reg $\pi^v$')

pl.savefig('./transp_piv_mnist.png')
pl.savefig('./transp_piv_mnist.pdf',bbox_inches='tight')

#%%
import random
import PIL as pil
def predict_barycenter(data,T):
    diag=1./T.sum(axis=1)
    diag[diag==np.inf]=0
    return np.dot(T,data.T).T.dot(np.diag(diag))
def predict_barycenter_reverse(data,T):
    diag=1./T.sum(axis=1)
    diag[diag==np.inf]=0
    return np.dot(T,data).T.dot(np.diag(diag)).T
random.seed(1985)
np.random.seed(1986)
n_fig=15
idx_sel=np.random.randint(0,ntot,n_fig)
xsel=xt[idx_sel,:]
xpred=np.zeros((n_fig,d1))
xpredr=np.zeros((n_fig,d1))
for i in range(n_fig):
    xpred[i,:]=predict_barycenter(xsel[i,:],Tv)
    xpredr[i,:]=predict_barycenter(xsel[i,:],Tvr)
pl.figure(figsize=(n_fig,4))
cmap_g='gray'
for i in range(n_fig):

    ax=pl.subplot(4,n_fig,i+1)

    pl.imshow(xsel[i,:].reshape((16,16)),cmap=cmap_g)
    pl.axis('off')
    pl.xticks(())
    pl.yticks(())
    pl.xlim([-6,22])
    pl.ylim([22,-6])
    if i==0:
        ax.text(left-.3, 0.5*(bottom+top), 'USPS',
        horizontalalignment='right',
        verticalalignment='center',
        rotation='vertical',
        transform=ax.transAxes)

    ax=pl.subplot(4,n_fig,i+1+n_fig)
    img = pil.Image.fromarray((xsel[i,:]* 255 / np.max(xsel[i,:])).reshape((16,16)).astype('float32'))
    img = img.resize((28,28))
    pl.imshow(np.array(img),cmap=cmap_g)
    pl.axis('off')
    pl.xticks(())
    pl.yticks(())
    if i==0:
        ax.text(left-.3, 0.5*(bottom+top) , 'Resize',
        horizontalalignment='right', verticalalignment='center',rotation='vertical',
        transform=ax.transAxes)
    ax=pl.subplot(4,n_fig,i+1+2*n_fig)
    pl.imshow(xpred[i,:].reshape((28,28)),cmap=cmap_g)
    pl.axis('off')
    pl.xticks(())
    pl.yticks(())
    if i==0:
        ax.text(left-.2, 0.5*(bottom+top) , 'Map $\pi^v$',
        horizontalalignment='right',
        verticalalignment='center',
        rotation='vertical',
        transform=ax.transAxes)

    ax=pl.subplot(4,n_fig,i+1+3*n_fig)
    pl.imshow(xpredr[i,:].reshape((28,28)),cmap=cmap_g)
    pl.axis('off')
    pl.xticks(())
    pl.yticks(())
    if i==0:
        ax.text(left-.2, 0.5*(bottom+top) , 'Map reg $\pi^v$',
        horizontalalignment='right',
        verticalalignment='center',
        rotation='vertical',
        transform=ax.transAxes)
        pl.ylabel('Map reg $\pi^v$')

pl.savefig('./transp_piv_usps.png')
pl.savefig('./transp_piv_usps.pdf',bbox_inches='tight')
```


We are working on the Burgers' Equation:

$u_t + uu_x = \nu u_{xx}$ with $x \in [-1,1]$.

Provided are the Dirichlet boundary conditions $u(t,-1)=u(t,1)=0$ and noisy measurements $\{x^0, u^0\}$ of the black-box initial function $u(0,x)=-\sin(\pi x)$.

The noise is chosen to be uniformly distributed in $[-0.25, 0.25].$

We apply the backward Euler scheme to discretize over time:

$u^{n}(x) = u^{n-1}(x) - \Delta t u^n \frac{d}{dx}u^n(x) + \Delta t \nu \frac{d^2}{dx^2}u^n(x)$

In order to be able to put Gaussian Processes as priors on $u^{n-1}$ and on $u^n$, we have to relinquish the non-linear term and approximate it with $\mu^{n-1}$, which is the posterior mean of the $n-1$-th time step. We get:

$u^n(x) = u^{n-1}(x) - \Delta t \mu^{n-1} \frac{d}{dx}u^n(x) + \Delta t \nu \frac{d^2}{dx^2}u^n(x)$.

Or:
$
u^{n-1}(x) = u^n(x) + \Delta t \mu^{n-1} \frac{d}{dx}u^n(x) - \Delta t \nu \frac{d^2}{dx^2}u^n(x) =: \mathcal{L}_x u^n(x)
$

We can then assume:

Prior: $u^n \sim \mathcal{GP}(0, k(x,x';\theta_n))$

We obtain the Numerical Gaussian Process:

$\begin{pmatrix}
u^n \\
u^{n-1} \end{pmatrix}
\sim
\mathcal{GP}\left(0,
\begin{pmatrix}
k_{u,u}^{n,n} & k_{u,u}^{n,n-1} \\
k_{u,u}^{n-1,n} & k_{u,u}^{n-1,n-1}
\end{pmatrix} \right)$, where $k_{u,u}^{n,n} = k$.

#### Step 1: Simulate data

```python
import time
import numpy as np
import sympy as sp
import matplotlib.pyplot as plt
from scipy.optimize import minimize
```

```python
n = 24
np.random.seed(int(time.time()))
noise = 0.5*np.random.rand(24) - 0.25
x0 = 2*np.random.rand(n) - 1
u0 = -np.sin(2*np.pi*x0) + noise

xb = np.array([-1,1])
ub = np.array([0,0])
```

Plotting the function

```python
plt.plot(x0,u0,'bo')
plt.plot(xb,ub,'bo', color = 'red')
xplot = np.arange(-1,1,0.002)
uplot = -np.sin(2*np.pi*xplot)
plt.plot(xplot,uplot)
```

#### Step 2: evaluate kernels and covariance matrix

Declare symbols

```python
x, xp, sig0, sig = sp.symbols('x, xp, sig0, sig')
```

We choose a non-stationary prior, since Burger's equation can develop discontinuities:

$ k_{u,u}^{n,n} = k(x, x'; \theta) = \frac{2}{\pi}\sin^{-1}\left(
\frac{2(\sigma_0^2 + \sigma^2xx')}{\sqrt{(1+2(\sigma_0^2 + \sigma^2x^2))(1+2(\sigma_0^2 + \sigma^2x'^2))}} \right)$

```python
k_sym = 2/sp.pi*sp.asin(2*(sig0**2 + sig**2*x*xp)/sp.sqrt((1+2*(sig0**2+sig**2*x**2))*(1+2*(sig0**2+sig**2*xp**2))))
k_fn = sp.lambdify((x, xp, sig0, sig), k_sym, "numpy")
def k(x, sig0, sig):
    m = np.zeros((x.size, x.size))
    for i in range(x.size):
        for j in range(x.size):
            m[i,j] = k_fn(x[i], x[j], sig0, sig)
    return m
```

Choosing parameters

```python
mu = np.mean(u0)
dt = 0.01
nu = dt/np.pi
```

$k_{u,u}^{n,n-1} = \mathcal{L}_{x'}k = k + \Delta t \mu^{n-1}(x')\frac{d}{dx'}k - \nu \Delta t \frac{d^2}{dx'^2}k$

```python
knn0_sym = k_sym + dt*mu*sp.diff(k_sym, xp) - nu*dt*sp.diff(k_sym, xp, xp)
knn0_fn = sp.lambdify((x, xp, sig0, sig), knn0_sym, "numpy")
def knn0(xb, x, sig0, sig):
    m = np.zeros((xb.size, x.size))
    for i in range(xb.size):
        for j in range(x.size):
            m[i,j] = knn0_fn(xb[i], x[j], sig0, sig)
    return m
```

$k_{u,u}^{n-1,n} = \mathcal{L}_{x}k = k + \Delta t \mu^{n-1}(x)\frac{d}{dx}k - \nu \Delta t \frac{d^2}{dx^2}k$

```python
def kn0n(xb, x, sig0, sig):
    return knn0(xb, x, sig0, sig).T
```

Calculating the kernel of $u^{n-1}$:

$k_{u,u}^{n-1,n-1} = \mathcal{L}_x \mathcal{L}_{x'} k = \mathcal{L}_x k_{u,u}^{n,n-1} \\ = k_{u,u}^{n,n-1} + \Delta t \mu^{n-1}(x)\frac{d}{dx}k_{u,u}^{n,n-1} - \nu \Delta t \frac{d^2}{dx^2}k_{u,u}^{n,n-1} \\ = k + \Delta t \mu^{n-1}(x')\frac{d}{dx'}k - \nu \Delta t \frac{d^2}{dx'^2}k + \Delta t \mu^{n-1}(x)\frac{d}{dx}k + \Delta t^2 \mu^{n-1}(x)\mu^{n-1}(x')\frac{d}{dx}\frac{d}{dx'}k - \nu \Delta t^2 \mu^{n-1}(x)\frac{d}{dx}\frac{d^2}{dx'^2}k - \nu \Delta t \frac{d^2}{dx^2}k - \nu \Delta t^2 \mu^{n-1}(x')\frac{d^2}{dx^2}\frac{d}{dx'}k + \nu^2\Delta t^2 \frac{d^2}{dx^2} \frac{d^2}{dx'^2}k$

```python
kn0n0_sym = knn0_sym + dt*mu*sp.diff(knn0_sym, x) - mu*dt*sp.diff(knn0_sym, x, x)
kn0n0_fn = sp.lambdify((x, xp, sig0, sig), kn0n0_sym, "numpy")
def kn0n0(x, sig0, sig):
    m = np.zeros((x.size, x.size))
    for i in range(x.size):
        for j in range(x.size):
            m[i,j] = kn0n0_fn(x[i], x[j], sig0, sig)
    return m
```

#### Step 3: create covariance matrix and NLML

```
params = [sig0, sig]
```

```python
def nlml(params, xb, x0, ub, u0, s):
    params = np.exp(params)
    K = np.block([
        [
            k(xb, params[0], params[1]) + s*np.identity(xb.size),
            knn0(xb, x0, params[0], params[1])
        ],
        [
            kn0n(xb, x0, params[0], params[1]),
            kn0n0(x0, params[0], params[1]) + s*np.identity(x0.size)
        ]
    ])
    y = np.concatenate((ub,u0))
    val = 0.5*(np.log(abs(np.linalg.det(K))) + np.mat(y) * np.linalg.inv(K) * np.mat(y).T)
    return val.item(0)
```

```python
nlml((1,1), xb, x0, ub, u0, 1e-8)
```

    5949290.832508889

```python
nlml_wp = lambda params: nlml(params, xb, x0, ub, u0, 1e-8)
res = minimize(nlml_wp, np.random.rand(2), method="Nelder-Mead", options={'maxiter' : 5000, 'fatol' : 0.001})
res
```

     final_simplex: (array([[1.39185121, 1.39541777],
           [1.39185121, 1.39541777],
           [1.39185121, 1.39541777]]), array([-3.23761254e+15, -3.23761254e+15, -3.23761254e+15]))
               fun: -3237612538384236.5
           message: 'Optimization terminated successfully.'
              nfev: 338
               nit: 150
            status: 0
           success: True
                 x: array([1.39185121, 1.39541777])

Predicting $u_n(x_{\ast}^n)$ at a test point $x_{\ast}^n$:

$u^n(x_{\ast}^n) \vert \begin{pmatrix} u_b^n \\ u^{n-1} \end{pmatrix} \sim \mathcal{N} \left( q^T K^{-1} \begin{pmatrix} u_b^n \\ u^{n-1} \end{pmatrix} , k_{u,u}^{n,n}(x_{\ast}^n, x_{\ast}^n) - q^T K^{-1}q \right)$, where

$ q^T := \begin{pmatrix} k_{u,b}^{n,n}(x_{\ast}^n, x_b^n) & k_{u,u}^{n, n-1}(x_{\ast}^n, x^{n-1}) \end{pmatrix} $

So the expected value is:

$u^n(x_{\ast}^n) = q^T K^{-1} \begin{pmatrix} u_b^n \\ u^{n-1} \end{pmatrix}$.

```python
sig0 = np.exp(res.x[0])
sig = np.exp(res.x[1])
s = 1e-8
xstar = 0.8                # As an example

K = np.block([
    [
        k(xb, sig0, sig) + s*np.identity(xb.size),
        knn0(xb, x0, sig0, sig)
    ],
    [
        kn0n(xb, x0, sig0, sig),
        kn0n0(x0, sig0, sig) + s*np.identity(x0.size)
    ]
])

q = np.zeros(xb.size + x0.size)

for i in range(xb.size):
    q[i] = k_fn(xb[i], xstar, sig0, sig)

for j in range(x0.size):
    q[xb.size + j] = knn0_fn(x0[j], xstar, sig0, sig)

y = np.concatenate([ub, u0])

print((q.T).dot(np.linalg.inv(K)).dot(y))

```

    4.022289251796972
    192914482.55315757

```python
#def minimize_restarts(xb, x0, ub, u0, n = 5):
#    nlml_wp = lambda params: nlml(params, xb, x0, ub, u0, 1e-8)
#    all_results = []
#    for it in range(0,n):
#        all_results.append(minimize(nlml_wp, np.random.rand(2), method="Nelder-Mead", options={'maxiter' : 5000, 'fatol' : 0.001}))
#    filtered_results = [m for m in all_results if 0 == m.status]
#    return min(filtered_results, key = lambda x: x.fun)
```

```python
#m = minimize_restarts(xb, x0, ub, u0)
#m
```


```python
import numpy as np
import math

from numpy import pi
from sympy import *
```

---

```python
class IntegralForDirichlet:

    def __init__(self, func, curve, domain, iterations, precision = 0.01):
        self.func, self.funcSymbols = func
        self.curve = np.array(curve)
        self.domain = domain
        self.N, self.n = iterations

        self.__initGradient()
        self.__initCurveDerivatives()

# Initializaton functions
#----------------------------------------------------------------------

    def __initCurveDerivatives(self):
        self.curve1D = self.curveDiff(self.curve)
        self.curve2D = self.curveDiff(self.curve1D)

    def __initGradient(self):
        self.funcGradient = self.gradientDiff(self.func)

# Main functions
#----------------------------------------------------------------------

    def solve(self):
        core = 0

        aK = 1 / self.N

        for k in range(1, self.N):
            etaK = 1 - (2*k-1)/(2*self.N)

            for j in range(2 * self.n - 1):
                tK = j*pi / self.n

                core += aK * self.getCore(self.parametrize(0.3, pi/3), self.parametrize(etaK, tK))

        return 1/(2*self.n) * abs(core)

    def getCore(self, x,y):
        eta, _ = x
        ksi, _ = y

        return self.kernel(x,y) if eta != ksi else self.kernelSingular(x,y)

    def kernel(self, x,y): return np.dot( (x-y), self.gradient(x) ) * self.jacobian(y) \
                                    / ( self.functionSubs(self.func, y) * (np.linalg.norm(x-y)**2) )

    def kernelSingular(self, x,y):
        eta, t = x
        ksi, tau = y

        self.calculateSingularFunctions(t, tau)

        firstPart = np.dot( self.gradient(eta*self.curveT), self.normalVector ) \
                        * self.innerKernelOne(t, tau)

        secondPart = (1 / self.euclideanDistance) \
                        * np.dot( self.gradient(eta*self.curveT), self.tangentVector ) \
                        * self.innerKernelTwo(t, tau)

        thirdPart = (1 / (2*self.euclideanDistance)) \
                        * np.dot( self.gradient(eta*self.curveT), self.tangentVector )

        denominator = eta*self.functionSubs(self.func, eta*self.curveTau)

        jacobian = self.jacobian(y)
        cot = self.singularCot(t, tau)

        return (firstPart - secondPart - thirdPart) * cot * jacobian  / denominator

#----------------------------------------------------------------------

    def innerKernelOne(self, t, tau):

        if t == tau:
            vector = np.flip(self.curve1DT)

            return np.multiply( self.negateAt(vector, 0), self.curve2DT ) / ( 2 * self.euclideanDistance ) ** 3

        return np.dot( self.curveT - self.curveTau, self.normalVector ) / np.linalg.norm( self.curveT - self.curveTau ) ** 2

    def innerKernelTwo(self, t, tau):

        if t == tau:
            return np.dot(self.curve1DT, self.curve2DT) / ( 2 * self.euclideanDistance ) ** 2

        return np.dot( self.curveT - self.curveTau, self.curve1DT ) / np.linalg.norm( self.curveT - self.curveTau ) ** 2  \
            - sin(tau -t) / ( 2 * (1 - cos(tau-t) ))

    def calculateSingularFunctions(self, t, tau):
        self.normalVector = self.normal(t)
        self.tangentVector = self.tangent(t)

        self.euclideanDistance = np.linalg.norm(self.curveSubs(self.curve1D, t))

        self.curveT = self.curveSubs(self.curve, t)
        self.curveTau = self.curveSubs(self.curve, tau)

        self.curve1DT = self.curveSubs(self.curve1D, t)
        self.curve2DT = self.curveSubs(self.curve2D, t)

#----------------------------------------------------------------------

    def parametrize(self, eta, t): return eta * self.curveSubs(self.curve, t)

    def normal(self, value):

        vector = np.flip(self.curveSubs(self.curve1D, value))

        vector = self.negateAt(vector, -1)

        return vector / np.linalg.norm(vector)

    def tangent(self, value):

        vector = self.curveSubs(self.curve1D, value)

        return vector / np.linalg.norm(vector)

    def gradient(self, vector):
        return np.array([ self.functionSubs(function, vector) for function in self.funcGradient ])

    def jacobian(self, vector):
        ksi, tau = vector

        innerVector = np.multiply( self.curveSubs(self.curve, tau), np.flip( self.curveSubs(self.curve1D, tau)) )

        return ksi * (innerVector[0] - innerVector[-1])

    def singularCot(self, t, tau):
        return - 1/self.n * sum([ sin(m*(t-tau)) for m in range(self.n) ]) - 1/(2*self.n) * sin(n*(t-tau))

# Helper functions
#----------------------------------------------------------------------

    def functionSubs(self, function, vector): return float(function.subs(zip(self.funcSymbols, vector)))

    def curveSubs(self, curve, value): return np.array([ float(function.subs({s:value})) for function in curve ])

    def curveDiff(self, curve): return np.array([ function.diff() for function in curve ])

    def gradientDiff(self, function): return np.array([ function.diff(symbol) for symbol in self.funcSymbols ])

    def negateAt(self, vector, index):

        vector[index] = vector[index] * -1

        return vector

```

---

## Initialization

```python
s = Symbol("s")                           # substitution value
funcSymbols = [a, b] = symbols("a, b")    # func substitution values
```

---

```python
func = (a+b+2)**2, funcSymbols
curve = x1, x2 = cos(s), sin(s)
domain = (0, 1), (0, 2*pi)
iterations = N, n = 10, 256
```

## Calculation

```python
integralSolver = IntegralForDirichlet(func, curve, domain, iterations)

print('\n==============================\n')

result = integralSolver.solve()
print(f'Result = {float(result)}')

print('\n==============================\n')
```

    ==============================

    Result = 0.6946521147089756

    ==============================

```python

```


# Solução dos problemas propostos na aula 02 - Interpolação trigonométrica

```julia
using PyPlot
using FFTW
```

```julia
#nfftutil(N) = (N % 2 == 0) ? (div(N,2) + 1) : div(N+1,2)
nfftutil(N) = div(N,2) + 1

```

<!-- TEASER_END -->

# Problemas

## Problema 1

Seja a seguinte função definida no domínio $-1 < x \le 1$:

$$
y(x) = \left\{\begin{matrix} -1 & x < 0 \\ 1 & x > 0\\ 0 & x=0\end{matrix}\right.
$$

 1. Ache os coeficientes $a_n$ e $b_n$ da série de Fourier
 2. Plote a série de Fourier para diferentes números de termos da série de Fourier. Percebe algo estranho?
 3. Qual a relação desta função com a função $y = |x|$ que estudamos anteriormente?
 4. Escolha diferentes números de pontos e interpole esta função nestes pontos

### Solução

#### Ache os coeficientes $𝑎_n$ e $𝑏_𝑛$ da série de Fourier

Esta é uma função ímpar, portanto os coeficientes $a_n$ são nulos!

$$
b_n = \frac{2}{L_0}\int_{-L_0/2}^{L_0/2} y(x) \cdot \sin\left(\frac{2\pi n x}{L_0}\right)\:dx = 2\int_0^1 \sin\left(\pi n x\right)\:dx = \frac{2}{n\pi}\left(1 - \cos n\pi\right) = \frac{2}{n\pi}\left[1 - (-1)^n\right]
$$

```julia
x = range(-1.0, 1.0, length=201);
y0 = 1.0 .* sign.(x)
bn = [2/(π*n) * (1.0 - cos(π*n)) for n in 1:500];
```

```julia
sinseries(x, b, N, L0) = sum(b[n]*sin(2π*n*x/L0) for n in 1:N);
```

#### Plote a série de Fourier para diferentes números de termos da série de Fourier

```julia
plot(x, y0, "k:")
for n in 1:2:15
    plot(x, sinseries.(x, Ref(bn), n, 2.0), label="n = $n")
end
legend()
```

```julia
x1 = range(-1, 1, length=1001)
plot(x1, sinseries.(x1, Ref(bn), 100, 2.0))
```

#### Qual a relação desta função com a função $y = |x|$ que estudamos anteriormente?

Esta função é a derivada da função $y = |x|$!

E a série de $dy/dx$ é formada pelas derivadas da série original!

Existem condições específicas onde se pode derivar a série termo a termo. Mas se as funções forem suaves, derive e integre a vontade! Funciona até para algumas descontinuidades.

#### Escolha diferentes números de pontos e interpole esta função nestes pontos

Lembram da função `myresample`? Do jeito que foi implementada originalmente, só funcionava para N ímpar! Vamos primeiro entender isso.

```julia
qpoints(Q, a=-1.0, b=1.0) = range(a, b, length=Q+1)[1:Q]

```

```julia
Q = 9
x = qpoints(Q, 0.0, 1.0)
dx  = x[2]-x[1]
f =
y = cos.(2π*f.*x) .+ 2*sin.(2π*f.*x);
Y = fft(y)./Q
freq = fftfreq(Q, 1/dx)

hcat(freq, round.(real.(Y), digits=2), round.(imag.(Y), digits=2))

```

```julia
Q = 8
x = qpoints(Q, 0.0, 1.0)
dx  = x[2]-x[1]
f = 4
y = cos.(2π*f.*x) .+ 2*sin.(2π*f.*x);
Y = rfft(y)./Q
freq = rfftfreq(Q, 1/dx)

hcat(freq, round.(real.(Y), digits=2), round.(imag.(Y), digits=2))

```

```julia
## Resample

function myresample(y, Qout)
    Qin = length(y)
    Nin = Qin ÷ 2 + 1
    Nout = Qout ÷ 2 + 1

    Yin = rfft(y)
    Yout = zeros(ComplexF64, Nout)
    nmin = min(Nout, Nin)
    for i in 1:nmin
        Yout[i] = Qout/Qin * Yin[i]
    end

    if Qout > Qin  # interpolação
        if iseven(Qin)
            Yout[Nin] /= 2
        end
    elseif Qout < Qin
        if iseven(Qin)
            Yout[Nout] = 2*real(Yout[Nout])
        end
    end

    return irfft(Yout, Qout)
end
```

```julia
Q0 = 32
x0 = range(-1, 1, length=Q0+1)[1:Q0]
y0 = 1.0 * sign.(x0);
```

```julia
Q1 =  256
x1 = qpoints(Q1, -1, 1)
y1 = myresample(y0, Q1)
plot(x1, y1)
plot(x0, y0, "r.")

```

## Problema 2

Verifique a ortogonalidade discreta da base trigonométrica, ou seja,

$$
\sum_{i=1}^Q \phi_n(x_i)\cdot \phi_k(x_i)=0 \qquad k\ne n, k+n \text{ e } k-n \text{ não são múltiplos de}\: Q/2
$$

para

\begin{align}
\phi_0(x) &= \frac{1}{2}\\
\phi_n(x) &= \cos\left(\frac{2\pi n x}{L_0}\right) \qquad 0\le n \le N \\
\phi_{n+N}(x) &= \sin\left(\frac{2\pi n x}{L_0}\right) \qquad 1\le n \le N
\end{align}

```julia
Q = 10
x = qpoints(Q)
```

```julia
L0 = 2.0
n = 3; k = 13
sum(sin.((2π*n/L0)*x) .* sin.((2π*k/L0)*x))
```

## Problema 3

Nos exemplos iniciais usando a DFT, ao se escolher os pontos, foi utilizado o seguinte código:
```julia
L₀ = 1.0; Q = 16; freq=1
x = range(0.0, L₀, length=Q+1)[1:Q]
```
O que acontece se utilzarmos a expressão mais simples a seguir?
```julia
x = range(0.0, L₀, length=Q)
```
O que está acontecendo? Você consegue explicar?

```julia
L₀ = 1.0; Q = 16; freq=1;
x = range(0.0, L₀, length=Q+1)[1:Q]
y = sin.(freq.*2π.*x); f = (0:Q-1) / L₀
Y = 1/Q * fft(y);
plot(f, imag.(Y), "ro")
plot(f, real.(Y), "xb")
```

```julia
L₀ = 1.0; Q = 16; freq=1;
x = range(0.0, L₀, length=Q)
y = sin.(freq.*2π.*x); f = (0:Q-1) / L₀
Y = 1/Q * fft(y);
plot(f, imag.(Y), "ro")
plot(f, real.(Y), "xb")

```

**O que está acontecendo?**

O que a gente quer?  Interpolar a função $\sin2\pi x$ para $0\le x\le 1$

Visualmente o que ocorre pode ser visto no gráfico a seguir:

```julia
Q = 4; x1 = range(0, 1, length=Q+1)[1:Q]; dx1 = x1[2]-x1[1];
x2 = range(0, 1, length=Q); dx2 = x2[2]-x2[1]
subplot(121); plot(x1, zeros(Q), "r.")
for i in 1:Q
    plot([x1[i], x1[i]+dx1], [0, 0], "k-")
end
subplot(122); plot(x2, zeros(Q), "r.")
for i in 1:Q
    plot([x2[i], x2[i]+dx2], [0, 0], "k-")
end
```

Temos uma inconsistência!

Como só passamos o número de pontos, na verdade estamos aproximando a função mais um trecho reto:

```julia
Q = 4; x1 = range(0, 1, length=Q+1)[1:Q]; dx1 = x1[2]-x1[1]
dx1 = 1.0 / Q; xx1 = 0:0.01:1; yy1 = sin.(2π.*xx1)
xx = [xx1; 1.0 + dx1]; yy = [yy1; 0.0]
plot(xx, yy)
```

## Problema 4

Implemente a DFT usando a definição. Verifique como aumenta o custo computacional quando o número de pontos aumenta.

Como foi mostrado, esta operação é uma multiplicação de matriz por vetor. Verifique se invertendo a ordem dos loops afeta o desempenho.

### DFT - Discrete Fourier Transform (transformada de Fourier discreta)

$$
c_n = X_n = \frac{1}{N}\sum_{k=0}^{N-1} x_k \exp\left(-\frac{2\pi i n k}{N} \right)
$$

Chamando
$$
w_k = w^k = \exp\left(-\frac{2\pi i k}{N}\right)
$$
Repare que isto pode ser escrito como uma multiplicação de matrizes:

$$
\left(\begin{matrix} X_0 \\ X_1 \\ \vdots \\ X_{N-1}\end{matrix} \right)
=
\left(\begin{matrix} 1 & 1 & \cdots & 1\\
1 & w & \cdots & w^{N-1} \\
\vdots & \vdots & \ddots &\vdots \\
1 & w^N & \cdots & w^{(N-1)^2}\\
\end{matrix}\right)
\cdot
\left(\begin{matrix} x_0 \\ x_1 \\ \vdots \\ x_{N-1}\end{matrix} \right)
$$

```julia
function mydft(x)
    N = length(x)
    X = zeros(ComplexF64,N)
    w = exp(-2π*im/N)

    for i in 1:N
        for k in 1:N
            X[i] += x[k] * w^((i-1)*(k-1))
        end
    end

    return X
end
```

```julia
x = rand(10)
X1 = mydft(x)
X2 = fft(x)
maximum(abs, X1 - X2)
```

```julia
function mydft!(x, X)
    N = length(x)
    w = exp(-2π*im/N)
    for i in 1:N
        Xi = 0.0 + 0.0im
        wi = w^(i-1)
        wk = 1.0 + 0.0im
        for k in 1:N
            Xi += x[k] * wk
            wk *= wi
        end
        X[i] = Xi
    end
    return X
end
mydft(x) = mydft!(x, zeros(ComplexF64, length(x)))
```

```julia
x = randn(20)
X1 = mydft(x)
X2 = fft(x)
maximum(abs, X1-X2)
```

## Problema 5

(Desafio) procure na net ou em algum livro como implementar a FFT. Tente!

Com $N$ divísível por 2:
\begin{align}
X[k] &= \sum_{r=0}^N x[r] W_N^{rk} \qquad W_N = \exp\left(-\frac{2\pi i}{N}\right)\\
X[k] &= \sum_{r=0}^{N/2-1} x[2r] W_N^{2rk} + \sum_{r=0}^{N/2-1} x[2r+1] W_N^{(2r+1)k} \\
X[k] &= \sum_{r=0}^{N/2-1} x[2r] W_{N/2}^{rk} + W_N^k\sum_{r=0}^{N/2-1} x[2r+1] W_{N/2}^{rk} \\
X[k] &= X_{par}[k] + W_N^k\cdot X_{ímpar}[k]
\end{align}

```julia
function myfft_rec(x)
    N = length(x)
    n1 = div(N,2)
    if N==1
        return [x[1]]
    elseif N % 2 != 0
        error("O tamanho deve ter potência de 2")
    else
        Xe = myfft_rec(x[1:2:end])
        Xo = myfft_rec(x[2:2:end])
        factor = exp.(-2π*im .* (0:N-1) ./ N)
        return[Xe .+ factor[1:n1] .* Xo;
            Xe .+ factor[(n1+1):end] .* Xo]

    end
end
```

```julia
x = randn(2^10)
X1 = myfft_rec(x)
X2 = fft(x)
maximum(abs, X1 - X2)
```

```julia
@time myfft(x);
@time fft(x);
```

### Jeito mais inteligente
(retirado do livro Random Vibratrions, Spectral & Wavelet Analysis de D. E. Newland)

```julia
function bitrev(n, bits)
    nrev = zero(n)
    N = 1 << bits
    nrev = n

    for i in 1:bits-1
        n >>= 1
        nrev <<= 1
        nrev |= n & 1
    end
    nrev &= N-1
    return nrev
end
function swaparr!(x)
    N = length(x)
    n = Int(log2(N))

    for i in 1:N
        ir = bitrev(i-1, n)+1
        if ir > i
            tmp = x[i]
            x[i] = x[ir]
            x[ir] = tmp
        end
    end
end

```

```julia
function myfft!(x)
    N = length(x)
    n = Int(log2(N))
    swaparr!(x)
    m = 1
    while true
        U = 1.0+0.0im
        j = 1
        k = 2^(m-1)
        W = exp(-π*im/k)
        while true
            l = j
            while true
                t = x[l+k] * U
                x[l+k] = x[l] - t
                x[l] = x[l] + t

                l = l + 2^m
                if l > 2^n
                    break
                end
            end
            U = U * W
            j = j + 1
            if j > k
                break
            end
        end
        m = m + 1
        if m > n
            break
        end
    end
    return x
end

```

```julia
n = 20
N = 2^n

x = rand(N) + rand(N) .* im
@time X1 = fft(x)
@time X2 = myfft!(copy(x))

maximum(abs,X1-X2)
```

## Problema 6

A transformada de Fourier está diretamente relacionada com a série de Fourier e é dada por:
$$
X(\omega) = \frac{1}{2\pi} \int_{-\infty}^\infty x(t)e^{-i\omega t}\: dt
$$

 ### 6.1 Calcula a transformada de Fourier para a função
  $$
  x(t) = \left\{\begin{matrix} 1 & -1 < x < 1 \\ 0 & x \ge 1\end{matrix}\right.
  $$
$$
X(\omega) = \frac{1}{2\pi}\int_{-\infty}^\infty x(t)e^{-i\omega t}\: dt =
\frac{1}{2\pi}\int_{-1}^1 e^{-i\omega t}\: dt = \frac{1}{2\pi}\left.\frac{i e^{-i\omega}}{\omega}\right|_{-1}^1 =
\frac{\sin\omega}{\pi\omega}
$$
$$
\text{sinc} = \frac{\sin\pi x}{\pi x}
$$

```julia
x = -15:0.01:15
y = 1/π .* sinc.(x./π)
plot(x,y)
```

### 6.2 Discretize a função e use a FFT para calcular a transformada.

$$
X(\omega) = \frac{1}{2\pi}\int_{-\infty}^\infty x(t)e^{-i\omega t}\: dt =
\frac{1}{2\pi}\int_{-T_0/2}^{T_0/2} e^{-i\omega t}\: dt \approx \frac{1}{2\pi} \sum_{k=0}^{N-1} x(t_k) e^{-i\omega t_k}\Delta t \qquad \Delta t = \frac{T_0}{N} \qquad t_k = -\frac{T_0}{2} + k\cdot \Delta t
$$

Ou seja,

$$
X(\omega) \approx \frac{1}{2\pi} \sum_{k=0}^{N-1} x(t_k) e^{-i\omega t_k}\Delta t =
\frac{\Delta t}{2\pi} \sum_{k=0}^{N-1} x_k e^{-i\omega(-T_0/2 + k\Delta t)} =
\frac{\Delta t}{2\pi}\cdot e^{\frac{i\omega T_0}{2}}\cdot \sum_{k=0}^{N-1} x_k e^{-i\omega k\Delta t} =
$$

Discretizando $\omega = \omega_j = j\cdot \omega_0 = \frac{2\pi j}{T_0}$, $0 \le j < N$

$$
X_j = X(\omega_j) = \frac{\Delta t}{2\pi}\cdot e^{\frac{2\pi j T_0}{2{T_0}}}\cdot \sum_{k=0}^{N-1} x_k e^{-\frac{2\pi i j\omega k\Delta t}{N\cdot \Delta t}} =
 \frac{\Delta t}{2\pi}\cdot e^{\pi i j}\cdot \sum_{k=0}^{N-1} x_k \exp\left(-\frac{2\pi i j k}{N}\right)
$$
ou seja
$$
X_j = X(\omega_j) =
\frac{\Delta t}{2\pi}\cdot e^{\pi i j}\cdot \sum_{i=0}^{N-1} x_k \exp\left(-\frac{2\pi i j k}{N}\right)
$$
A somatória no ultimo termo é a DFT. O termo $e^{\pi i j}$ é corresponde à translação.

```julia
Q = 32
Δt = 2.0 / Q
x = ones(Q)
X0 = rfft(x)
N = length(X)
for j in 0:N-1
    X0[j+1] = X0[j+1] * Δt/(2π) * exp(π*j*im)
end
freq0 = rfftfreq(Q, 1/Δt)
w = 2π*freq0;
```

```julia
ww = range(0, maximum(w), length=1001)
yy = 1/π .* sinc.(ww./π)
plot(ww./(2π), yy)
plot(freq0, real.(X0), "ro")
```

### 6.3 Tente verificar o que acontece se você adiciona zeros à esquerda e à direita da função discretizada.

```julia
Q1 = 32; nn = 1; x = [zeros(nn*Q1); ones(Q1); zeros(nn*Q1-1)]
Q = length(x); Δt = 2.0 / Q1
X = rfft(x); N = length(X)
for j in 0:N-1
    X[j+1] = X[j+1] * Δt/(2π) * exp(π*j*im)
end
freq = rfftfreq(Q, 1/Δt); w = 2π*freq;
```

```julia
plot(ww./(2π), yy); plot(freq, real.(X), "r.")
plot(freq0, real.(X0), "b+")
```

## Problema 7

Faça alguns exemplos do pacote `ApproxFun` (<https://github.com/JuliaApproximation/ApproxFun.jl>)


# Introduction

This tutorial reproduces part of [Thomas Fösel et al. (2020)](https://arxiv.org/abs/2004.14256),
titled **Efficient cavity control with SNAP gates**. The general
idea is to start in an initial state and apply a sequence of
operators to reach the desired target state.

In the paper, the authors use a vacuum state, $|0 \rangle$,
with a Hilbert space cutoff of 10.

Authors apply sequence of operators as blocks, $\hat{B}$,
where

\begin{equation}
\hat B = D(\alpha) \hat S(\vec \theta) D(\alpha)^{\dagger}
\end{equation}

In the equation above, $D(\alpha)$ is a Displace operation
given by

\begin{equation}
 D(\alpha) = e^{\alpha a^{\dagger} - \alpha^* \hat a}
\end{equation}

where $\alpha$ is the complex displacement parameter,
and $a^{\dagger}$ and $\hat a$ are the bosonic creation
and annihilation operators respectively.

$S(\vec \theta)$ in the block definition of $\hat{B}$
above is the selective number-dependent-snap arbitrary
phase (SNAP) gate defined as

\begin{equation}\label{snap-gate}
        \hat S(\vec{\theta}) = \sum_{0}^{n} e^{i \theta^{(n)}} |n\rangle \langle n|
\end{equation}

where $\vec{\theta}$ is the real vector containing SNAP
parameters, and $|n\rangle$ is the number state.

With the operator definitions laid out, the idea now is to
apply this sequence of $D(\alpha)$, $S(\vec{\theta})$, and
$D(\alpha)^{\dagger}$ operations to $|0 \rangle$, in order
to reach the target binomial state, which in this case is

\begin{equation}
b_{1} = \frac{\sqrt 3 |3\rangle + |9\rangle}{2}
\end{equation}

**Note:** We do not follow the _exact_ same scheme as used in
[Thomas Fösel et al. (2020)](https://arxiv.org/abs/2004.14256),
in that we use a simpler cost function (without regularization)
and do not add blocks in a Breadth-First manner.

```python
import numpy as onp # this is ordinary numpy that we will
                    # use with QuTiP for visualization only.
                    # onp should not used for gradient
                    # calculations. Instead jax.numpy
                    # should be used.

import jax.numpy as jnp
from jax import grad, jit
from jax.experimental import optimizers

# Visualization
import matplotlib.pyplot as plt
from qutip.visualization import plot_wigner, hinton
from qutip import Qobj # imported purely for visualization purposes due to QuTiP-JAX non-compatibility

from qgrad.qgrad_qutip import basis, to_dm, dag, Displace, fidelity

from tqdm.auto import tqdm
```

```python
def pad_thetas(hilbert_size, thetas):
    """
    Pads zeros to the end of a theta vector to
    fill it upto the Hilbert space cutoff.

    Args:
        hilbert_size (int): Size of the hilbert space
        thetas (:obj:`jnp.ndarray`): List of angles thetas

    Returns:
        :obj:`jnp.ndarray`: List of angles padded with zeros
        in place of Hilbert space cutoff

    """
    if len(thetas) != hilbert_size:
        thetas = jnp.pad(thetas, (0, hilbert_size - len(thetas)), mode="constant")
    return thetas

def snap(hilbert_size, thetas):
    """
    Constructs the matrix for a SNAP gate operation
    that can be applied to a state.

    Args:
        hilbert_size (int): Hilbert space cuttoff
        thetas (:obj:`jnp.ndarray`): A vector of theta values to
                apply SNAP operation

    Returns:
        :obj:`jnp.ndarray`: matrix representing the SNAP gate
    """
    op = 0 * jnp.eye(hilbert_size)
    for i, theta in enumerate(thetas):
        op += jnp.exp(1j * theta) * to_dm(basis(hilbert_size, i))
    return op
```

# Insertion of Blocks

$T$ number of displace-snap-displace blocks,
$\hat B = D(\alpha) \hat S(\vec \theta) D(\alpha)^{\dagger}$
are applied on the initial vacuum state, $|0 \rangle$

In this reconstruction we choose $T = 3$, so $3$
$\hat B$'s are applied to the initial state. The
aim is to tune $\alpha$'s and $\vec{\theta}$'s
such that three repeated applications of
$\hat B$ on $|0 \rangle$

\begin{equation}
 D(\alpha) \hat S(\vec \theta) D(\alpha)^{\dagger}
D(\alpha) \hat S(\vec \theta) D(\alpha)^{\dagger}
D(\alpha) \hat S(\vec \theta) D(\alpha)^{\dagger} |0 \rangle
\end{equation}

lead to the desired
binomial state $b_{1}$

**Note on Displace operator**: We provide a fast implementation
of the displacement operation using a diagonal decomposition
such that a `Displace` class initialises the operator which
can then be applied for different $\alpha$ repeatedly. This
construct also supports autodiff with JAX.

```python
def apply_blocks(alphas, thetas, initial_state):
    """Applies blocks of displace-snap-displace
       operators to the initial state.

    Args:
        alphas (list): list of alpha paramters
                for the Displace operation
        thetas (list): vector of thetas for the
                SNAP operation
        initial_state(:obj:`jnp.array`): initial
                state to apply the blocks on

    Returns:
        :obj:`jnp.array`: evolved state after
                applying T blocks to the initial
                state

    """

    if len(alphas) != len(thetas):
        raise ValueError("The number of alphas and theta vectors should be same")

    N = initial_state.shape[0]
    displace = Displace(N)
    x = initial_state

    for t in range(len(alphas)):
        x = jnp.dot(displace(alphas[t]), x)
        x = jnp.dot(snap(N, thetas[t]), x)
        x = jnp.dot(displace(-alphas[t]), x) # displace(alpha)^{\dagger} = displace(-alpha)
    return x
```

# Visualizing the state evolution under $\hat B$

Before we move on to make learning routines, it might be a
good idea to see what `apply_blocks` does to the initial
state, $|0 \rangle$ with randomly initialized parameters
$\alpha$ and $\vec{\theta}$

```python
def show_state(state):
    """Shows the Hinton plot and Wigner function for the state"""
    fig, ax = plt.subplots(1, 2, figsize=(11, 5))
    if state.shape[1] == 1: # State is a ket
        dm = Qobj(onp.array(jnp.dot(state, dag(state))))
        hinton(dm, ax=ax[0])
        plot_wigner(dm, ax=ax[1])
    plt.show()

N = 10 # Hilbert space cutoff

initial_state = basis(N, 0) # initial vacuum state
show_state(initial_state)

alphas = jnp.array([1., 0.5, 1.]) # Displace parameters

theta1, theta2, theta3 = [0.5], [0.5, 1.5, 0.5], [0.5, 1.5, 0.5, 1.3] # SNAP parameters
# NOTE: No input values to JAX differentiable functions should be int
thetas = jnp.array([pad_thetas(N, p) for p in [theta1, theta2, theta3]])

evolved_state = apply_blocks(alphas, thetas, initial_state)
show_state(evolved_state)
```

The target state that we aim to reach is visualized below.
The aim is to act `apply_blocks` function to the initial
state with the optimized parameters, say $\alpha_{opt}$
and $\theta_{opt}$ such that we land extremely close to the
desired binomial state $b_{1}$ as defined above.

```python
target_state = (jnp.sqrt(3) * basis(N, 3) +  basis(N, 9)) / 2.0 # target state b1 as shown above|
show_state(target_state)
```

```python
def cost(params, initial, target):
    """
    Calculates the cost between the target state and
    the one evolved by the action of three blocks.

    Args:
    -----
        params (jnp.array): alpha and theta params of Displace and SNAP respectively
        initial (jnp.array): initial state to apply the blocks on
        target (jnp.array): desired state

    Returns:
    --------
        cost (float): cost at a particular parameter vector
    """
    alphas, thetas = params[0], params[1]
    evo = apply_blocks(alphas, thetas, initial)
    return 1 - fidelity(target, evo)[0][0]
```

# Optimization using Adam -- case in point for `qgrad`

This is where the power of `qgrad` comes in.
Since qgrad's functions used in this notebook -- basis,
to_dm, dag, Displace, fidelity -- support JAX, we can
evaluate the gradient of the cost function in one line
using JAX's `grad`. This saves us painstakingly evaluating
the derivative of the cost function analytically.

```python
alphas = jnp.array([1, 0.2, 0.1 - 1j]) # Displace parameters
theta1, theta2, theta3 = ([0.5],
                          [0., 0.5],
                          [0.1, 0., 0., 0.1])
# NOTE: No input values to JAX differentiable functions should be int
thetas = jnp.array([pad_thetas(N, p) for p in [theta1, theta2, theta3]])

init_params = [alphas, thetas]
opt_init, opt_update, get_params = optimizers.adam(step_size=1e-2)
opt_state = opt_init([alphas, thetas])

def step(i, opt_state, opt_update):
    params = get_params(opt_state)
    g = grad(cost)(params, initial_state, target_state)
    return opt_update(i, g, opt_state)

epochs = 150
pbar = tqdm(range(epochs))

fidel_hist = []
params_hist = []

for i in pbar:
    opt_state = step(i, opt_state, opt_update)
    params = get_params(opt_state)
    params_hist.append(params)
    f = 1 - cost(params, initial_state, target_state)
    fidel_hist.append(f)
    pbar.set_description("Fidelity {}".format(f))
```

    HBox(children=(IntProgress(value=0, max=150), HTML(value='')))

```python
def display_evolution(parameters, num_plots = 4):
    """
    Displays the intermediate states during the learning schedule.

    Args:
        parameters (list): List of device arrays of parameters
            `alpha` and `theta`
        num_plots (int): Number of plots to make generate (except
            the initial and final state plots)

    Returns:
        Evolution of the state at `num_plots` equidistant
        times during the training
    """
    show_state(initial_state)

    diff = epochs // int(num_plots)
    for i in range(0, epochs, diff):
        alphas, thetas = parameters[i]
        x = apply_blocks(alphas, thetas, initial_state)
        show_state(x)

    # final parameters after last epoch
    alphas, thetas = parameters[-1]
    x = apply_blocks(alphas, thetas, initial_state)
    show_state(x)

display_evolution(params_hist)
```

# Conclusion

We see that starting from a vacuum state $|0 \rangle$,
we efficiently learn the target state
$b_{1} = \frac{\sqrt 3 |3> + |9>}{2}$, as is corroborated
by the fidelity plot below.

The desired target state's Hinton plot and Wigner function
are shown before the learning scedule starts. It can be
seen that the last row above is _almost_ the same as the
target state, implying we learn parameters $\alpha$ and
$\vec \theta$ efficiently so as to go from the vacuum
state to a desired target state in just **three** applications
of blocks, $\hat B$. This tutorial is just meant to show the
usage of `qgrad`. Like the authors in the original paper, one
may use more sophistacated techniques like gradient clipping,
regularized cost function, and greater number of epochs to
achieve even better learning.

```python
plt.plot(fidel_hist)
plt.ylabel("Fidelity")
plt.xlabel("epochs")
```

# References
[1] Fösel, Thomas, et al. "Efficient cavity control with SNAP gates." arXiv preprint arXiv:2004.14256 (2020).

examples/notebooks/Efficient Cavity Control with SNAP Gates.ipynb