```python
%matplotlib inline
```

Linear Regression of multivariate data
======================================

In this example, we demonstrate how to use sklearn_xarray classes to solve a
simple linear regression problem on synthetic dataset.

This class demonstrates the use of :py:class:`~sklearn_xarray.Stacker` and
:py:class:`~sklearn_xarray.Select`.

```python
import numpy as np
import xarray as xr
from sklearn.linear_model import LinearRegression
from sklearn.pipeline import make_pipeline, make_union

from sklearn_xarray import Stacker, Select

# Make synthetic data
lat, lon = np.ogrid[-45:45:50j, 0:360:100j]
noise = np.random.randn(lat.shape[0], lon.shape[1])

data_vars = {
    'a': (['lat', 'lon'], np.sin(lat/90 + lon/100)),
    'b': (['lat', 'lon'], np.cos(lat/90 + lon/100)),
    'noise': (['lat', 'lon'], noise)
}

coords = {'lat': lat.ravel(), 'lon': lon.ravel()}
dataset = xr.Dataset(data_vars, coords)
```

make a simple linear model for the output

\begin{align}y = a + .5 b + 1\end{align}

```python
x = dataset[['a', 'b']]
y = dataset.a + dataset.b * .5 + .3 * dataset.noise  + 1

y.plot()
```

now we want to fit a linear regression model using these data

```python
mod = make_pipeline(
    make_union(
        make_pipeline(Select('a'), Stacker()),
        make_pipeline(Select('b'), Stacker())),
    LinearRegression())
```

for now we have to use Stacker manually to transform the output data
into a 2d array

```python
y_np = Stacker().fit_transform(y)
print(y_np)
```

fit the model

```python
mod.fit(x, y_np)

# print the coefficients
lm = mod.named_steps['linearregression']
coefs = tuple(lm.coef_.flat)
print("The exact regression model is y = 1 + a + .5 b + noise")
print("The estimated coefficients are a: {}, b: {}".format(*coefs))
print("The estimated intercept is {}".format(lm.intercept_[0]))
```


# Computing the implicit gradient w.r.t hyperparameter

## Elementary theory

Consider $\mathbb{R}$-normed spaces $(E, \|\cdot\|)$ and $(F, \|\cdot\|)$,
and a map $\phi\colon U \to F$, where $U$ is open set in $E$.

The map $\phi$ is Frechet differentiadle at $a\in U$ if there exists a
bounded linear map $\ell\colon E \to F$ such that for any $\varepsilon > 0$
there is $\delta > 0$ such that for any $\|h\| \leq \delta$
with $a+h \in U$ we have
$$
    \bigl\|
        \phi(a+h) - \phi(a)
        - \ell(h)
    \bigr\| \leq \varepsilon \|h\|
    \,. $$

This map is unique and called the differential of $\phi$ at $a$, i.e. $d\phi(a) = \ell$.

Note that, since $U$ is open (in the norm topology) and $a\in U$ we can
always choose a smaller $\delta$ so that for all $\|h\|\leq \delta$ it
holds that $a + h \in U$.

As an operator, $d\phi\colon U \to \mathcal{B}(E, F)$, is a mapping
from the open domain of $\phi$ into the space of linear operators.
Note that the space of linear operators $\mathcal{B}(E, F)$ is also
normed: for $\ell \in \mathcal{B}(E, F)$ we have
$$
    \|\ell\|
        = \sup\bigl\{
            \tfrac{\|\ell x\|}{\|x\|}
            \colon x \in E\,,\,x\neq 0
        \bigr\}
        = \sup\bigl\{
            \|\ell x\|
            \colon x \in E\,,\,\|x\|\leq 1
        \bigr\}
        = \inf\bigl\{
            K \geq 0
            \colon \|\ell x \| \leq K \|x\|\,,\, \forall x\in E
        \bigr\}
    \,. $$

Therefore applying the definition to $a \mapsto d\phi(a)$, we get
that the second order differential of $\phi$, aka the differential
of $a\mapsto d\phi(a)$, denoted by $d^2\phi$ at $a$ is an element
of $\mathcal{B}(E, \mathcal{B}(E, F))$, i.e. a bilinear map. The
definition goes: $\phi$ is twice differentiable at $a$ if $\phi$
is differentiable on $U$ and the map $d\phi$ is differentiable as
$a$, i.e. for any $\varepsilon > 0$ there is $\delta > 0$ such that
for all $\|h\| \leq \delta$ such that $a + h \in U \cap E$ we have

$$
\bigl\|
    d\phi(a+h) - d\phi(a)
    - d(d\phi(a))(h)
\bigr\| \leq \varepsilon \|h\|
    \,. $$

The linear map $d(d\phi(a))$ takes linear maps as values, and so
$h\mapsto d(d\phi(a))(h)$ is shortened to $h\mapsto d^2\phi(a)(h, \cdot)$.

The trick is to notice that for any fixed $h\in E$ we can
put $\psi_h \colon U \to F \colon a\mapsto d\phi(a)(h)$.
Then $d\psi_h(a)$ is $d^2\phi(a)(\cdot, h)$.

Indeed fix some $h \in E$. Then, formally, for any $\varepsilon > 0$
there exists $\delta > 0$ such that for any $\|v\| \leq \delta$ with
$a+v \in U\cap E$ we have the linearization for $\psi_h(a)$ and
$d\phi(a)$. Then

$$
\begin{align}
    \bigl\|
        d\psi_h(a)(v) - d^2\phi(a)(v, h)
    \bigr\|
        &\leq \bigl\|
            d\phi(a + v)(h) - d\phi(a)(h) - d^2\phi(a)(v, h)
        \bigr\| + \bigl\|
            d\phi(a + v)(h) - d\phi(a)(h) - d\psi_h(a)(v)
        \bigr\|
        \\
        &\leq \bigl\|
            d\phi(a + v) - d\phi(a) - d^2\phi(a)(v, \cdot)
        \bigr\| \| h \| + \bigl\|
            \psi_h(a + v) - \psi_h(a) - d\psi_h(a)(v)
        \bigr\|
        \\
        &\leq \varepsilon \|v\| (1 + \| h \|)
\end{align}
    \,. $$

Then for any $\|x\| \leq 1$ we have $\|\delta x\|\leq \delta$ and

$$
\bigl\|
    d\psi_h(a)(x) - d^2\phi(a)(x, h)
\bigr\|
    = \tfrac1\delta \bigl\|
        d\psi_h(a)(\delta x) - d^2\phi(a)(\delta x, h)
    \bigr\|
    \leq \tfrac1\delta \varepsilon (1 + \| h \|) \|\delta x \|
    \,. $$

Since $\delta > 0$ the term $\varepsilon (1 + \| h \|)$ upper bounds
the norm for all $\|x\| \leq 1$, and thus the operator norm is also
bounded. Therefore for any $\varepsilon > 0$ we have

$$
\bigl\|
    d\psi_h(a) - d^2\phi(a)(\cdot, h)
\bigr\|
    \leq \varepsilon (1 + \| h \|)
    \,, $$
which implies that $d\psi_h(a) = d^2\phi(a)(\cdot, h)$

**(duh)**
This implies that we cat take $a\mapsto d\phi(a)$, evaluate it at
some constant vector $h$ to get $a\mapsto d\phi(a)(h)$, and
differentiale. In $\mathbb{R}^n$ this means that differentiating
$a \mapsto \nabla \phi(a)^\top h$ yields $\nabla^2 \phi(a) h$.

To find $\nabla^2_{\lambda\omega} f(\omega, \lambda) v$ do:
1. fix $v$ and compute $
\phi_v
\colon (\omega, \lambda) \mapsto
    \nabla_\omega f(\omega, \lambda)^\top v
$
2. compute $
\nabla_\lambda \phi_v(\omega, \lambda)
    = \nabla_\lambda \bigl( \nabla_\omega f(\omega, \lambda)^\top v \bigr)
    = \nabla^2_{\lambda\omega} f(\omega, \lambda) v
$

```python
import numpy as np

import torch

import torch.nn.functional as F
```

## Primitives

linear operatons in the direct-product space of inner product spaces:
$( E, [\cdot, \cdot])$ with $E = \prod_i E_i$ and $
[\cdot, \cdot]
    = \oplus_i \langle \cdot, \cdot \rangle_i
$
with each $(E_i, \langle \cdot, \cdot \rangle_i)$ being an inner product
space over $\mathbb{K}$:

* (`daxpy`, `dscal`) $a, b \in E$, and $\lambda \in \mathbb{K}$ we have $\lambda a + b \in E$
given by $(\lambda a + b)_i = \lambda a_i + b_i$ in each $E_i$.

* (`ddot`, `dnorm`) for $a, b \in E$ we have $[a, b] = \sum_i \langle a_i, b_i \rangle_i$
and $\|a\| = \sqrt{[a, a]}$

```python
from tools.trcg import clone, dzero, daxpy, ddot

from tools.trcg import trcg
```

Flat/unflat for `scipy.sparse.linalg` api

```python
from tools.trcg import dflat as pack

def numel(a):
    return sum(map(torch.numel, a))

def shapes(a):
    return [u.shape for u in a]

def unpack(array, shapes):
    slices = torch.split(array, [*map(np.prod, shapes)])
    return [*map(torch.reshape, slices, shapes)]

omega = [torch.randn(10, 10), torch.randn(100), torch.randn(10, 10, 12)]
assert all(map(torch.allclose, unpack(pack(omega), shapes(omega)), omega))
```

Get the data:

```python
import torch.utils.data

X = torch.randn(1024, 10).double()
y = X @ torch.randn(10, 1).double()

train = torch.utils.data.TensorDataset(X[:25], y[:25])
test = torch.utils.data.TensorDataset(X[250:], y[250:])
```

setup losses.

```python
alpha = torch.tensor(1., requires_grad=True).double()
```

#### What about within-model hyperparameter?

```python
def mse_loss(model, X, y):
    l2_reg = sum(p.norm()**2 for n, p in model.named_parameters()
                 if "bias" not in n)
    return 0.5 * F.mse_loss(model(X), y) + 0.5 * alpha * l2_reg
```

```python
loss = mse_loss
```

```python
from torch.nn import Sequential, Linear, LeakyReLU, Tanh

model = Sequential(
    Linear(10, 64),
    LeakyReLU(),
    Linear(64, 64),
    LeakyReLU(),
    Linear(64, 1),
#     Linear(10, 1, bias=False)
).double()
```

## Hyperparameter optimization with approximate gradient

Observations
* at $\omega^*(\lambda) = \arg\min_\omega \mathcal{L}_\mathrm{train}(\omega, \lambda)$ we have $\nabla_\omega \mathcal{L}_\mathrm{train}(\omega^*(\lambda), \lambda) \equiv 0$, i.e.

$$
    \nabla_\lambda \omega^*(\lambda)
        \nabla^2_{\omega\omega} \mathcal{L}_\mathrm{train}(\omega^*(\lambda), \lambda)
    + \nabla^2_{\lambda\omega} \mathcal{L}_\mathrm{train}(\omega^*(\lambda), \lambda)
        = 0
    \,, $$

* using the chain rule we get for $
F\colon \lambda \mapsto \mathcal{L}_\mathrm{test}(\omega^*(\lambda), \lambda)
$

$$
    \nabla_\lambda F(\lambda)
        = \nabla_\lambda \mathcal{L}_\mathrm{test}(\omega^*(\lambda), \lambda)
        + \nabla_\lambda \omega^*(\lambda)
            \, \nabla_\omega \mathcal{L}_\mathrm{test}(\omega^*(\lambda), \lambda)
    \,. $$

[HOAG](https://arxiv.org/pdf/1602.02355.pdf) Algorithm:
1. (approx) solve $\omega^*(\lambda) = \arg\min \mathcal{L}_\mathrm{train}(\omega, \lambda)$
2. (approx) find such $q_\lambda$ that
$
    \nabla^2_{\omega\omega} \mathcal{L}_\mathrm{train}(\omega^*(\lambda), \lambda)
        \, q_\lambda
    = \nabla_\omega \mathcal{L}_\mathrm{test}(\omega^*(\lambda), \lambda)
$ (can add $c I_\omega$ for $0 < c \ll 1$ to the hessian so stabilize CG)
3. compute
$$
    \nabla_\lambda F(\lambda)
        = \nabla_\lambda \mathcal{L}_\mathrm{test}(\omega^*(\lambda), \lambda)
        - \nabla^2_{\lambda\omega} \mathcal{L}_\mathrm{train}(\omega^*(\lambda), \lambda) \, q_\lambda
    \,. $$

Therefore

$$
    \nabla_\lambda F(\lambda)
        = \nabla_\lambda \mathcal{L}_\mathrm{test}(\omega^*(\lambda), \lambda)
        - \nabla^2_{\lambda\omega} \mathcal{L}_\mathrm{train}(\omega^*(\lambda), \lambda)
        \underbrace{
            \bigl(
                \nabla^2_{\omega\omega} \mathcal{L}_\mathrm{train}(\omega^*(\lambda), \lambda)
            \bigr)^{-1}
            \nabla_\omega \mathcal{L}_\mathrm{test}(\omega^*(\lambda), \lambda)
        }_{q_\lambda}
    \,.
    \tag{$\lambda$-deriv}
$$

### Steps

#### 1. (approx) find stationary point

$\omega^*(\lambda)$ sits in `model.parameters()`

**(note)**
we have to train the model to $
\nabla_\omega \mathcal{L}_\mathrm{train}
    (\omega^*(\lambda), \lambda)
    \approx 0
$
(almost zero gradient), and preferably to a local minimum (so that CG
below uses a positive definite hessian).

```python
from mlss2019bdl import fit

model.zero_grad()

# 1. find (approx) \arg \min_\omega L_train(\omega, \lambda)
fit(model, train, criterion=loss, batch_size=32, n_epochs=1000, verbose=True)
```

```python
loss_h = loss(model, *train.tensors)  # L_train(\omega^*(\lambda), \lambda)
grad_h_omega = torch.autograd.grad(loss_h, model.parameters(), create_graph=False)

sum(map(torch.norm, grad_h_omega))
```

#### 2. (approx) find $q_\lambda$

Observe
$$
    \nabla^2_{\omega\omega}
        h(\omega, \lambda) \delta
        = \tfrac{\partial}{\partial \omega}
            \bigl\{ \nabla_\omega^\top h(\omega, \lambda) \delta \bigr \}
        = \tfrac{\partial}{\partial \omega} \ell_\delta(\omega)
        = d\bigl\{ dh(\omega, \lambda)(\delta) \bigr \}
    \,. $$

```python
# 2.2 create a hess-vect closure
def get_hesv_op(model, dataset):
    # get \nabla_\omega L_train(\omega^*(\lambda), \lambda)
    loss_h = loss(model, *dataset.tensors)
    grad_h_omega = torch.autograd.grad(loss_h, model.parameters(), create_graph=True)

    def _closure(delta, nugget=1e-6):  # adding a tiny diagonal helps a lot
        with torch.enable_grad():
            # \omega \mapsto \nabla_\omega^\top h(\omega, \lambda) \delta
            grad_h_vect = ddot(grad_h_omega, delta)

            # get \nabla^2_{\omega\omega} L_train(\omega^*(\lambda), \lambda) \delta
            hesv_h_omega = torch.autograd.grad(grad_h_vect, model.parameters(), retain_graph=True)

            return daxpy(nugget, delta, hesv_h_omega)

    return _closure
```

Evaluate test loss and find the omega-gradient at the star-point $\omega^*(\lambda)$

```python
loss_g = loss(model, *test.tensors)  # L_test(\omega^*(\lambda), \lambda)

grad_g_omega = torch.autograd.grad(loss_g, model.parameters(), create_graph=False)
```

Solve $
\nabla^2_{\omega\omega}
    \mathcal{L}_\mathrm{train}(\omega^*(\lambda), \lambda)
    \,\, \mathbf{x}
    = \nabla_\omega \mathcal{L}_\mathrm{test}(\omega^*(\lambda), \lambda)
$

```python
Ax = get_hesv_op(model, train)  # lambda v: [-p for p in v]
```

```python
x, r = dzero(clone(grad_g_omega)), clone(grad_g_omega)
print(trcg(Ax, r, x, rtol=1e-8, atol=1e-9, verbose=True))

assert all(map(torch.allclose, Ax(x), grad_g_omega))
```

#### check the gradinet norms

```python
loss_h = loss(model, *train.tensors)  # L_train(\omega^*(\lambda), \lambda)
grad_h_omega = torch.autograd.grad(loss_h, model.parameters(), create_graph=False)
```

```python
sum(map(torch.norm, grad_h_omega)), sum(map(torch.norm, grad_g_omega))
```

Let's use the interface from `scipy.sparse`

```python
from scipy.sparse.linalg import LinearOperator
from torch.autograd import grad

def get_linear_operator(model, dataset):
    # get \nabla_\omega L_train(\omega^*(\lambda), \lambda)
    loss_h = loss(model, *dataset.tensors)
    grad_h = pack(grad(loss_h, model.parameters(), create_graph=True))

    def matvec(delta, nugget=1e-6):
        with torch.enable_grad():
            # \omega \mapsto \nabla_\omega^\top h(\omega, \lambda) \delta
            # grad_h_delta = torch.dot(grad_h, torch.from_numpy(delta))

            # get \nabla^2_{\omega\omega} L_train(\omega^*(\lambda), \lambda) \delta
            # hesv_h = pack(grad(grad_h_delta, model.parameters(), retain_graph=True))

            # Supply torch.autograd.grad with the delta-vector through
            #  `output_grads` for the Jacobian-vector product
            hesv_h = pack(grad(grad_h, model.parameters(), torch.from_numpy(delta),
                               retain_graph=True))

        # pack and back to numpy
        hesv_h = hesv_h.numpy()

        # adding a tiny diagonal helps a lot
        hesv_h += nugget * delta
        return hesv_h

    return LinearOperator(shape=(len(grad_h), len(grad_h)), matvec=matvec)
```

The following needlessly unpacks and packs stuff, while
the one above pushes these out of the matvec `kernel`
```python
def get_linear_operator(model, dataset):
    loss_h = loss(model, *dataset.tensors)

    grad_h = grad(loss_h, model.parameters(), create_graph=True)

    def matvec(delta, nugget=1e-6):
        output_grads = unpack(torch.from_numpy(delta), shapes(grad_h))
        with torch.enable_grad():
            hesv_h = grad(grad_h, model.parameters(), output_grads, retain_graph=True)

        hesv_h = pack(hesv_h).numpy()

        hesv_h += nugget * delta
        return hesv_h

    return LinearOperator(shape=2*[numel(grad_h)], matvec=matvec)
```

**(loose notation)**
If we define a function $
\Phi
    \colon \Omega \times C_2(\Omega \times A)
    \to A
$ (it is possible that $\Lambda = \Omega$)

$$
\Phi_\alpha(\delta \, \vert f)
    = \nabla_\alpha \bigl(
        \nabla_\omega f^\top \delta
    \bigr)
    = \nabla^2_{\alpha\omega} f \, \delta
    \,. $$

Then the last term in eq. ($\lambda$-deriv) for the HOAG can be computed by using

* use cg to solve $
\Phi_\omega(q_\lambda \, \vert \mathcal{L}_\mathrm{train})
    = \nabla_\omega \mathcal{L}_\mathrm{test}(\omega^*(\lambda), \lambda)
$ at $\omega = \omega^*(\lambda)$ for $q_\lambda$

* compute $\Phi_\lambda(q_\lambda \, \vert \mathcal{L}_\mathrm{train})$ at $\omega = \omega^*(\lambda)$

$$
\nabla^2_{\lambda\omega} \mathcal{L}_\mathrm{train}(\omega^*(\lambda), \lambda)
    \underbrace{
        \bigl(
            \nabla^2_{\omega\omega} \mathcal{L}_\mathrm{train}(\omega^*(\lambda), \lambda)
        \bigr)^{-1}
        \nabla_\omega \mathcal{L}_\mathrm{test}(\omega^*(\lambda), \lambda)
    }_{q_\lambda}
    \,. $$

Observe that for a twice-differentiable $f$ we can approximate
$$
\partial^2_{\alpha\omega} f
    = \lim_{\delta \to 0}
        \frac{\partial_\alpha f(\omega+\delta, \alpha) - \partial_\alpha f(\omega-\delta, \alpha)}{2\delta}
\,, $$
where $\lim$ is understood as convergence in $\|\cdot\|$ op-norm of
linear functionals $\mathcal{B}(A, \mathbb{R})$. Here $
    \partial^2_{\omega\alpha} f \equiv \partial^2_{\alpha\omega} f
$
is a bilinear form $\mathcal{B}(\Omega, \mathcal{B}(A, \mathbb{R}))$

A smart-function approach

```python
class LossHessVect():
    """Computes a hessian-vector product of a computed loss."""
    def __init__(self, output, params):
        self.params = list(params)
        self.grad_h = pack(grad(output, self.params,
                                create_graph=True,
                                allow_unused=False))

    def linop(self):
        shape = len(self.grad_h), len(self.grad_h)
        return LinearOperator(shape=shape, matvec=self)

    def hessvec(self, delta, params=None):
        r"""get \nabla^2_{\omega\omega} L_train(\omega^*(\lambda), \lambda) \delta"""
        params = self.params if params is None else params
        if not isinstance(delta, torch.Tensor):
            delta = torch.from_numpy(delta)

        with torch.enable_grad():
            # Supply torch.autograd.grad with the delta-vector through
            #  `grad_outputs` for the Jacobian-vector product
            return grad(self.grad_h, params, delta.to(self.grad_h),
                        retain_graph=True, allow_unused=True)

    def __call__(self, delta, nugget=1e-6):
        hesv_h = pack(self.hessvec(delta, self.params))

        # pack and back to numpy
        hesv_h = hesv_h.cpu().numpy()

        # adding a tiny diagonal helps a lot
        hesv_h += nugget * delta
        return hesv_h
```

Invoke scipy's `cg`

```python
from scipy.sparse.linalg import cg

lhv = LossHessVect(loss(model, *train.tensors), model.parameters())

# A = get_linear_operator(model, train)
A = lhv.linop()
b = pack(grad_g_omega).numpy()

res, info = cg(A, b, M=None, tol=1e-7)

print(info)
```

```python
unp = unpack(torch.from_numpy(res), shapes(grad_g_omega))

print(all(map(torch.allclose, Ax(unp), grad_g_omega)))

print(all(map(torch.allclose, unp, x)))
```

```python
[*map(lambda d: float(abs(d).max()), (abs(a-b) for a, b in zip(unp, x)))]
```

#### 3. (approx) find the gradient

$$
    \nabla_\lambda F(\lambda)
        = \nabla_\lambda \mathcal{L}_\mathrm{test}(\omega^*(\lambda), \lambda)
        - \nabla^2_{\lambda\omega} \mathcal{L}_\mathrm{train}(\omega^*(\lambda), \lambda)
        \underbrace{
            \bigl(
                \nabla^2_{\omega\omega} \mathcal{L}_\mathrm{train}(\omega^*(\lambda), \lambda)
            \bigr)^{-1}
            \nabla_\omega \mathcal{L}_\mathrm{test}(\omega^*(\lambda), \lambda)
        }_{q_\lambda}
    \,. $$

Observing
$$
    \nabla^2_{\lambda\omega}
        h(\omega, \lambda) \delta
        = \tfrac{\partial}{\partial \lambda}
            \bigl\{ \nabla_\omega^\top h(\omega, \lambda) \delta \bigr \}
        = \tfrac{\partial}{\partial \lambda} \ell_\delta(\omega)
        = d\bigl\{ dh(\omega, \lambda)(\delta) \bigr \}
    \,. $$

```python
from torch.autograd import grad

def grad_lambda(q):
    # get \nabla_\omega L_train(\omega^*(\lambda), \lambda)
    loss_h = loss(model, *train.tensors)
    grad_h_omega = grad(loss_h, model.parameters(), create_graph=True)

    # \lambda \mapsto \nabla_\omega^\top h(\omega, \lambda) q
    # grad_h_vect = ddot(grad_h_omega, q)

    # get \nabla^2_{\lambda\omega} L_train(\omega^*(\lambda), \lambda) q
    # hesv_h_lambda = grad(grad_h_vect, alpha, retain_graph=False)

    # Supply `grad()` with `q` as the `output_grads` for
    #  the Jacobian-vector product (see grad's docs).
    hesv_h_lambda = grad(grad_h_omega, alpha, q, retain_graph=False)

    # get \nabla_\lambda L_test(\omega^*(\lambda), \lambda)
    loss_g = loss(model, *test.tensors)
    grad_g_lambda = grad(loss_g, alpha, create_graph=False)

    # return daxpy(-1, hesv_h_lambda, grad_g_lambda)
    return grad_g_lambda, hesv_h_lambda
```

```python
t_grad, t_hesv = grad_lambda(x)
t_grad, t_hesv
```

```python
lhv.hessvec(torch.from_numpy(res), alpha)
```

```python
eta = 1e-3
# alpha.data -= eta * t_hesv[0]
```

```python
alpha
```

```python
assert False
```

# Trunk

$$
\mathcal{L}(\omega, \lambda)
    = \tfrac12 \| y - X \omega \|^2_2
    + \tfrac\lambda2 \|\omega\|^2_2
    \,. $$

$$
\nabla_\omega \mathcal{L}
    = - X^\top(y - X\omega) + \omega \lambda
    = (X^\top X + \lambda I) \, \omega - X^\top y
    \,. $$

$$
\nabla_\lambda \mathcal{L}
    = \tfrac12 \|\omega \|^2_2
    \,. $$

$$
\nabla_{\omega \lambda} \mathcal{L}
    = \nabla_{\lambda \omega} \mathcal{L}
    = \omega
    \,. $$

Now $\omega^*(\lambda) = (X^\top X + \lambda I)^{-1} X^\top y$, whence

$$
F(\lambda)
    = \mathcal{L}(\omega^*(\lambda), \lambda)
    = \tfrac12 \| y - X \omega^*(\lambda) \|^2_2
    + \tfrac\lambda2 \| \omega^*(\lambda) \|^2_2
    \,, $$

and

$$
\nabla_\lambda F(\lambda)
    = \tfrac12 \|\omega^*(\lambda)\|^2_2
    \,. $$

```python
loss_h = loss(model, *train.tensors)

grad_h_omega = torch.autograd.grad(loss_h, model.parameters(), create_graph=True)

fn = ddot(grad_h_omega, [p.data for p in model.parameters()])

torch.autograd.grad(fn, alpha, retain_graph=True)
```

```python
sum(map(torch.norm, grad_h_omega)), sum(map(torch.norm, model.parameters()))
```

Consider the following setup:
$$
    l^*(\lambda)
        = \min_{\omega} l(\omega, \lambda)
        = l(\omega^*(\lambda), \lambda)
    \,, $$
where
$$
    \omega^*(\lambda) = \arg \min_{\omega} l(\omega, \lambda)
    \,. $$

Suppose everything is sufficiently smooth, and $l \neq \ell$. Then differentiating
$f(\lambda) = \ell(\omega^*(\lambda), \lambda)$:
$$
    \partial_\lambda f(\lambda)
        = \partial_1 \ell(\omega^*, \lambda) \circ \partial_\lambda \omega^*
        + \partial_2 \ell(\omega^*, \lambda)
    \,. $$

In finite dimensions we have
$$
    \nabla_\lambda f
        = \nabla_\lambda \omega^*(\lambda) \nabla_\omega \ell(\omega, \lambda)
            \big\vert_{\omega^*(\lambda), \lambda}
        + \nabla_\lambda \ell(\omega, \lambda)
            \big\vert_{\omega^*(\lambda), \lambda}
    \,. $$

Otoh

$$
    \nabla_\lambda l^*(\lambda)
        = \nabla_\lambda \omega^*(\lambda) \nabla_\omega l(\omega, \lambda)
            \big\vert_{\omega^*(\lambda), \lambda}
        + \nabla_\lambda l(\omega, \lambda)
            \big\vert_{\omega^*(\lambda), \lambda}
        = 0
    \,, $$
whence
$$
    \nabla_\lambda \omega^*(\lambda) \nabla_\omega l(\omega, \lambda)
            \big\vert_{\omega^*(\lambda), \lambda}
        = - \nabla_\lambda l(\omega, \lambda)
            \big\vert_{\omega^*(\lambda), \lambda}
    \,. $$

Gateaux derivative
$$
    \tfrac{d}{d\eta } \omega^*(\lambda + \eta v) \big \vert_{\eta = 0}
        = \lim_{\eta \to 0} \frac{
            \arg \min_{\omega} l(\omega, \lambda + \eta v)
            - \arg \min_{\omega} l(\omega, \lambda)
        }{\eta}
\,. $$

# Using a class

```python
from numpy import prod

class vector(tuple):
    def __new__(cls, *tensors):
        assert tensors
        if isinstance(tensors[0], cls):
            return tensors[0]

        if isinstance(tensors[0], (tuple, list)):
            tensors = tensors[0]

        self = super().__new__(cls, tensors)
        self.shape = tuple(u.shape for u in self)
        return self

    def __pos__(self):
        return self

    def __neg__(self):
        return type(self)(*map(torch.neg, self))

    def __add__(self, other):
        return type(self)(*map(torch.add, self, other))
    __radd__ = __add__

    def __iadd__(self, other):
        return type(self)(*map(torch.Tensor.add_, self, other))

    def __sub__(self, other):
        return type(self)(*map(torch.sub, self, other))

    def __isub__(self, other):
        return type(self)(*map(torch.Tensor.sub_, self, other))

    def __mul__(self, alpha):
        return type(self)([torch.mul(u, alpha) for u in self])
    __rmul__ = __mul__

    def __imul__(self, alpha):
        return type(self)([u.mul_(alpha) for u in self])

    def __div__(self, other):
        return type(self)([torch.div(u, alpha) for u in self])

    def __idiv__(self, other):
        return type(self)([u.div_(alpha) for u in self])

    def clone(self):
        return type(self)(*map(torch.clone, self))

    def zero_(self):
        return type(self)(*map(torch.zero_, self))

    def detach(self):
        return type(self)(*map(torch.detach, self))

    def requires_grad_(self, requires_grad=True):
        return type(self)([u.requires_grad_(requires_grad) for u in self])

    def to(self, *args, **kwargs):
        return type(self)([u.to(*args, **kwargs) for u in self])

    def tensor(self):
        return torch.cat([u.flatten() for u in self])

    @classmethod
    def from_tensor(cls, shapes):
        flat = torch.split(a, [*map(prod, shape)])
        return cls([u.reshape(s) for u, s in zip(flat, shape)])

    def __matmul__(self, other):
        return torch.dot(self.tensor(), other.tensor())

    def __abs__(self):
        return torch.norm(self.tensor())

    def __len__(self):
        return sum(u.numel() for u in x)
```

```python
from math import sqrt

def trcg_vec(Ax, r, x, n_iterations=1000, tr_delta=0, rtol=1e-5, atol=1e-8,
         args=(), verbose=False):
    if n_iterations > 0:
        n_iterations = min(n_iterations, len(x))

    p, iteration = r.clone(), 0
    tr_delta_sq = tr_delta ** 2

    rtr, rtr_old = float(r @ r), 1.0
    cg_tol = sqrt(rtr) * rtol + atol
    region_breached = False
    while (iteration < n_iterations) and (sqrt(rtr) > cg_tol):
        Ap = vector(Ax(p, *args))
        iteration += 1
        if verbose:
            print("""iter %2d |Ap| %5.3e |p| %5.3e """
                  """|r| %5.3e |x| %5.3e beta %5.3e""" %
                  (iteration, abs(Ap), abs(p), abs(r), abs(x), rtr / rtr_old))
        # end if

        alpha = rtr / float(p @ Ap)
        x += alpha * p
        r -= alpha * Ap

        # check trust region (diverges from tron.cpp in liblinear and leml-imf)
        if tr_delta_sq > 0:
            xTx = float(x @ x)
            if xTx > tr_delta_sq:
                xTp = float(x @ p)
                if xTp > 0:
                    # backtrack into the trust region
                    p_nrm = abs(p)

                    q = xTp / p_nrm
                    eta = (q - sqrt(max(q * q + tr_delta_sq - xTx, 0))) / p_nrm

                    # reproject onto the boundary of the region
                    r += eta * Ap
                    x -= eta * p
                else:
                    # this never happens maybe due to CG iteration properties
                    pass
                # end if

                region_breached = True
                break
            # end if
        # end if

        rtr, rtr_old = float(r @ r), rtr
        p *= rtr / rtr_old
        p += r
    # end while

    return iteration, region_breached
```

```python
x, r = vector(grad_g_omega).clone().zero_(), vector(grad_g_omega).clone()
print(trcg_vec(Ax, r, x, rtol=1e-9, atol=1e-9, verbose=True))

assert all(map(torch.allclose, Ax(x), grad_g_omega))
```

```python

```

```python

```

```python
class Term(object):
    def __init__(self):
        self.losses = set()

    def __add__(self, other):
        self.losses.add(other)
        return self

    def __call__(self, model, X, y):
        if self.losses:
            return sum(loss.forward(model, X, y) for loss in self.losses)
        return self.forward(model, X, y)
```

```python
class MSELoss(Term):
    def forward(self, model, X, y):
        return 0.5 * F.mse_loss(model(X), y)
```

# Splitting alphas and weights

```python
def get_parameters(module, prefix=""):
    for name, par in module.named_parameters(prefix):
        pass
    pass
```

# Gumbel-softmax trick

```python
    prob_t = tt.nnet.softmax(logit_t)

    # Gumbel-softmax sampling: Gumbel (e^{-e^{-x}}) distributed random noise
    gumbel = -tt.log(-tt.log(theano_random_state.uniform(size=logit_t.shape) + eps) + eps)
#     logit_t = theano.ifelse.ifelse(tt.gt(tau, 0), gumbel + logit_t, logit_t)
#     inv_temp = theano.ifelse.ifelse(tt.gt(tau, 0), 1.0 / tau, tt.constant(1.0))
    logit_t = tt.switch(tt.gt(tau, 0), gumbel + logit_t, logit_t)
    inv_temp = tt.switch(tt.gt(tau, 0), 1.0 / tau, tt.constant(1.0))

    # Get the softmax: x_t is `BxV`
    x_t = tt.nnet.softmax(logit_t * inv_temp)
```

```python
temp = 1e-1
logit = torch.randn(10, 10)

gumbel = -torch.log(-torch.log(torch.rand_like(logit)))

proba = torch.softmax((logit + gumbel) / temp, dim=-1)
```


# Section 4.3.2

This is an implementation of the second order problem in Sec. 4.3.2 of the paper

    M Mortensen, "A generic and strictly banded spectral Petrov-Galerkin method for differential equations with polynomial coefficients"

# Helmholtz equation

The Helmholtz Dirichlet problem is described as

$$
u''(x) - \mu u(x) = f(x), \quad u(\pm 1) = 0,  x \in [-1, 1]
$$

where $\mu \ge 0$ is a constant, $f(x)$ is some right hand side function and $u(x)$ is the solution. For the solution we will use either one of the orthogonal polynomials $Q_n(x)$ defined in [Section-4.3.1](Section-4.3.1.ipynb). That is, we use $T_n, L_n, U_n$ or $Q^{(\alpha)}_n$.

We solve the Helmholtz equation with a Petrov-Galerkin method and use the trial functions

$$
\psi_n = Q_n-Q_{n+2}, \quad n=0,1, \ldots, N-2,
$$

for the three families $Q_n=T_n, L_n,$ or $Q^{(\alpha)}_n$. For Chebyshev polynomials of the second kind the trial functions are

$$
\psi_n = {U_n}-\frac{n+1}{n+3}U_{n+2}, \quad n=0,1, \ldots, N-2
$$

which is actually equal to $\psi_n=(n+1)(Q^{(1/2)}_n-Q^{(1/2)}_{n+2})$.

The trial space is $S=\{\psi_n\}_{n=0}^{N-2}$. For the test space we use $V=\text{span}\left\{\phi^{(2)}_n\right\}_{n=0}^{N-2}$, where $\phi^{(2)}_n$ also was defined in [Section-4.3.1](Section-4.3.1.ipynb). Note that we use

$$
\phi^{(k)}_n = \frac{(1-x^2)^k}{h^{(k)}_{n+k}} \frac{d^k Q_{n+k}}{dx^k},
$$

for $Q_n=T_n, L_n, Q^{(\alpha)}_n$, or $U_n$, and that $h^{(k)}_{n+k}$ includes the scaling function $g_n$.

The method is to find $u \in S$ such that

$$
(u'', v)_{\omega^{\alpha}} - \mu(u, v)_{\omega^{\alpha}} = (f, v)_{\omega^{\alpha}}, \quad \text{for all } v \in V
$$

We start by importing shenfun's functionality, and specify some parameters with obvious meaning.

```python
from shenfun import *
import matplotlib.pyplot as plt

N = 60
mu = 1
```

Choose family of polynomials by uncommenting only one of the following four lines.

```python
#family = 'Ultraspherical'; kw = {'alpha': 1}
#family = 'Legendre'; kw = {}
family = 'Chebyshev'; kw = {}
#family = 'ChebyshevU'; kw = {}
```

Create function spaces and test and trial functions. Assemble the coefficient matrix A, that will consist of both mass and stiffness matrices.

```python
S = FunctionSpace(N, family, bc=(0, 0), **kw)
V = FunctionSpace(N+2, family, basis='Phi2', **kw)
u = TrialFunction(S)
v = TestFunction(V)
A = inner(v, div(grad(u)) - mu*u)
```

```python
fig, (ax1, ax2) = plt.subplots(1, 2)
ax1.spy(A[0].diags(), ms=2)
ax2.spy(A[1].diags(), ms=2)
fig.suptitle('Sparsity pattern stiffness (left) and mass (right) matrix')
```

## Alternative formulation

An alternative formulation (see Sec. 4.1.4) for this problem  is to find $u\in \text{V}^{(1)}_N$ such that

\begin{equation}
    (u'', v)_{\omega^{(\alpha+1)}} - \mu (u, v)_{\omega^{(\alpha+1)}} = (f, v)_{\omega^{(\alpha+1)}}, \quad \forall \, v \in \text{V}^{(1)}_{N} = \text{span}\{\overline{\phi}^{(2,\alpha)}_m\}_{m=0}^{N-2},
\end{equation}

where

\begin{equation}
    \overline{\phi}^{(2,\alpha,\beta)}_m = \gamma^{(2,\alpha,\beta)}_{m} \phi^{(1, \alpha+1, \beta+1)}_{m}, \label{eq:phiover}
\end{equation}

and $\gamma^{(2,\alpha,\beta)}_m$ is given by Eq. (4.18) in the paper.

We use Chebyshev polynomials of the first kind for trial function ($\psi_n=T_n-T_{n+2}$) and the second kind for the test functions that become

$$
\overline{\phi}^{(2,-1/2, -1/2)}_m = \frac{1}{m+2} \phi^{(1,1/2,1/2)}_m =  \frac{1}{\pi (m+2)}\left(\frac{U_m}{m+1}-\frac{U_{m+2}}{m+3}\right)
$$

It is implemented in shenfun as

```python
S = FunctionSpace(N, 'C', bc=(0, 0))
V = FunctionSpace(N, 'U', basis='Phi1', scaled=True)
u = TrialFunction(S)
v = TestFunction(V)
A2 = inner(v, div(grad(u)) - mu*u)
```

Verify that the matrix is the same as `A` assembled above. Need to make sure Chebyshev is used, so recompute `A` here

```python
S = FunctionSpace(N, 'C', bc=(0, 0), **kw)
V = FunctionSpace(N+2, 'C', basis='Phi2', **kw)
u = TrialFunction(S)
v = TestFunction(V)
A = inner(v, div(grad(u)) - mu*u)
```

```python
M = A[0].diags('csr') + A[1].diags('csr')
M2 = A2[0].diags('csr') + A2[1].diags('csr')
assert np.linalg.norm((M-M2).toarray()) < 1e-14
```

If the assert passes, then the matrices are the same.

# Condition numbers

Check the condition number of the coefficient matrix. Use a relatively low number since Jacobi is not implemented very robustly, with gamma functions that overflow. Also, Numpy's routine for computing the condition number makes use of dense matrices, which limits the size.

```python
def cond(N, family='Chebyshev', alpha=0):
    """Return condition number of Helmholtz coefficient matrix

    Parameters
    ----------
    N : int
        The number of quadrature points. For a Dirichlet space there is N-2 dofs.
    family : str, optional
        Either one of
        - 'Chebyshev'      (or 'C')
        - 'ChebyshevU'     (or 'U')
        - 'Legendre'       (or 'L')
        - 'Ultraspherical' (or 'Q')
    alpha : number, optional
        Parameters used only by the Ultraspherical family
    """
    assert N < 1000
    S = FunctionSpace(N, family, bc=(0, 0), alpha=alpha)
    V = FunctionSpace(N+2, family, basis='Phi2', alpha=alpha)
    u = TrialFunction(S)
    v = TestFunction(V)
    A = inner(v, div(grad(u)) - mu*u)
    M = np.sum(np.array(A, dtype=object)).diags()
    return np.linalg.cond(M.toarray())
```

Run over all families and compute the condition numbers for a range of matrix sizes.

```python
alpha = 0.5
con = {'C': [], 'L': [], 'U': [], 'Q': []}
for family in 'CULQ':
    for n in range(3, 8):
        con[family].append(cond(2**n, family, alpha)) # alpha neglected by C, U, L
```

Plot the condition numbers using loglog axes to show that Chebyshev, Legendre and any ultraspherical basis using scaling $g_n^{\alpha}=(P^{(\alpha,\alpha)}_n(1))^{-1}$ obtain condition numbers that are $\mathcal{O}(N)$, whereas Chebyshev of the second kind (which uses a slightly different scaling) obtain condition numbers scaling as $\mathcal{O}(N^{3/2})$. Note that an ultraspherical basis with $\alpha=1/2$ results in $\mathcal{O}(N)$.

```python
for family in 'CULQ':
    plt.loglog(2**np.array(range(3, 8)), con[family])

plt.loglog(2**np.array(range(3, 8)), 2**np.array(range(3, 8)), '--')
plt.loglog(2**np.array(range(3, 8)), (2**np.array(range(3, 8)))**1.5, '--')
plt.legend(['Chebyshev 1', 'Chebyshev 2', 'Legendre', rf'Ultrasph. $(\alpha={alpha})$', r'$\mathcal{O}(N)$', r'$\mathcal{O}(N^{3/2})$'])
plt.title('Condition numbers Helmholtz coefficient matrix')
plt.show()
```

# The Airy equation

Another second order problem is the Airy differential equation

\begin{equation}
    \epsilon u'' - x u = 0, \quad u(-1) = \text{Ai}\left(-\sqrt[3]{\tfrac{1}{\epsilon}}\right), u(1) = \text{Ai}\left(\sqrt[3]{\tfrac{1}{\epsilon}}\right), \label{eq:airy}
\end{equation}

which has the Airy function $u(x) = \text{Ai}\left(\sqrt[3]{\tfrac{1}{\epsilon}} x\right)$ as solution.

```python
import sympy as sp
x = sp.Symbol('x', real=True)

def main(N, family, alpha=0, returnmat=False, e=sp.Rational(1, 1e9)):
    ue = sp.airyai((1/e)**(sp.Rational(1, 3))*x)
    SN = FunctionSpace(N, family, bc=(ue.subs(x, -1), ue.subs(x, 1)))
    VN = FunctionSpace(N+2, family, basis='Phi2')
    u = TrialFunction(SN)
    v = TestFunction(VN)
    A = [inner(div(grad(u)), e*v)] + inner(u, -x*v)
    sol = la.Solver(A)
    if returnmat:
        return sol.mat
    u_hat = Function(SN)
    f_hat = Function(SN)
    u_hat = sol(f_hat, u_hat)
    ua = Array(SN.get_orthogonal(), buffer=ue)
    return np.sqrt(inner(1, (u_hat.backward()-ua)**2))

```

```python
error = []
#e, M = sp.Rational(1, 1e6), np.arange(500, 900, 10)
e, M = sp.Rational(1, 1e9), np.hstack([np.arange(200, 19201, 1000), np.arange(19220, 20781, 20), np.arange(20800, 40000, 1000)])
for N in M:
    error.append(main(N, 'C', e=e))
```

```python
ue = sp.airyai((1/e)**(sp.Rational(1, 3))*x)
fig, ax = plt.subplots(1, 2, figsize=(10, 4))
ax[1].semilogy(M, error, 'k')
ax[1].set(xlabel='N', ylabel='$||u_N-u||$')
M = np.linspace(-1, 1, 40000)
ax[0].plot(M, sp.lambdify(x, ue)(M), 'k')
ax[0].set(xlabel='x', ylabel='u(x)')
left, bottom, width, height = [0.35, 0.6, 0.1, 0.2]
ax2 = fig.add_axes([left, bottom, width, height])
M = np.linspace(-0.05, 0.05, 1000)
ax2.plot(M, sp.lambdify(x, ue)(M), 'k')
plt.show()
```

```python
alpha = 0.5
con = {'C': [], 'L': [], 'U': [], 'Q': []}
M = range(3, 12)
e = sp.Rational(1, 1e3)
for family in 'CULQ':
    for n in M:
        con[family].append(np.linalg.cond(main(2**n, family, alpha, True, e=e).toarray())) # alpha neglected by C, U, L
```

```python
for family in 'CULQ':
    plt.loglog(2**np.array(M), con[family])

plt.loglog(2**np.array(M), 2**np.array(M), '--')
plt.loglog(2**np.array(M), (2**np.array(M))**1.5, '--')
plt.legend(['Chebyshev 1', 'Chebyshev 2', 'Legendre', rf'Ultrasph. $(\alpha={alpha})$', r'$\mathcal{O}(N)$', r'$\mathcal{O}(N^{3/2})$'])
plt.title('Condition numbers Airy coefficient matrix')
plt.show()
```


# Change points

*Joseph Hall (October 2019)*

This notebook demonstrates the use of the `ChangePoints` kernel, which can be used to describe one-dimensional functions that contain a number of change-points, or regime changes. The kernel makes use of sigmoids ($\sigma$) to blend smoothly between different kernels. For example, a single change-point kernel is defined by:

\begin{equation}
\textrm{cov}(f(x), f(y)) = k_1(x, y)\cdot\bar{\sigma}(x, y) + k_2(x, y)\cdot\sigma(x, y)
\end{equation}

where $\sigma(x, y) = \sigma(x)\cdot\sigma(y)$ and $\bar{\sigma}(x, y) = (1 - \sigma(x))\cdot(1 - \sigma(y))$. The sigmoid ($\sigma$) is parameterized by a location ($l$) and a width ($w$).

```python
import gpflow
import numpy as np
import matplotlib.pyplot as plt

np.random.seed(123)  # for reproducibility of this notebook

plt.style.use("ggplot")
%matplotlib inline

def plotkernelsample(k, ax, xmin=-3, xmax=3, title=None):
    xx = np.linspace(xmin, xmax, 100)[:, None]
    ax.plot(xx, np.random.multivariate_normal(np.zeros(100), k(xx), 3).T)
    ax.set_title(title)
```

We demonstrate the use of the kernel by drawing a number of samples from different parameterizations. Firstly, a simple single change-point between two kernels of differing lengthscales.

```python
np.random.seed(1)

base_k1 = gpflow.kernels.Matern32(lengthscales=0.2)
base_k2 = gpflow.kernels.Matern32(lengthscales=2.0)
k = gpflow.kernels.ChangePoints([base_k1, base_k2], [0.0], steepness=5.0)

f, ax = plt.subplots(1, 1, figsize=(10, 3))
plotkernelsample(k, ax)
```

Secondly, an implementation of a "change window" in which we change from one kernel to another, then back to the original.

```python
np.random.seed(3)

base_k1 = gpflow.kernels.Matern32(lengthscales=0.3)
base_k2 = gpflow.kernels.Constant()
k = gpflow.kernels.ChangePoints([base_k1, base_k2, base_k1], locations=[-1, 1], steepness=10.0)

f, ax = plt.subplots(1, 1, figsize=(10, 3))
plotkernelsample(k, ax)
```

And finally, allowing different change-points to occur more or less abruptly by defining different steepness parameters.

```python
np.random.seed(2)

base_k1 = gpflow.kernels.Matern32(lengthscales=0.3)
base_k2 = gpflow.kernels.Constant()
k = gpflow.kernels.ChangePoints(
    [base_k1, base_k2, base_k1], locations=[-1, 1], steepness=[5.0, 50.0]
)

f, ax = plt.subplots(1, 1, figsize=(10, 3))
plotkernelsample(k, ax)
```


```python
%config InlineBackend.figure_format = 'retina'
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
```

# [International Mathmatics Olympiad 2021](https://www.imo-official.org/problems.aspx) - Problem 2

Show that the inequality

$
\begin{equation}
\sum_{i=1}^{n} \sum_{j=1}^{n} \sqrt{|x_i - x_j|} \le \sum_{i=1}^{n} \sum_{j=1}^{n} \sqrt{|x_i + x_j|}
\end{equation}
$

holds for all real numbers $x_1, ...,  x_n$.

Shubhanshu Mishra (shubhanshu.com)

$
\begin{equation}
\sum_{i=1}^{n} \sum_{j=1}^{n} \color{red}{\sqrt{|x_i - x_j|}} \le \sum_{i=1}^{n} \sum_{j=1}^{n} \color{blue}{\sqrt{|x_i + x_j|}} \\
\Rightarrow \sum_{i=1}^{n} \sum_{j=1}^{n} \color{red}{\sqrt{|x_i - x_j|}} - \sum_{i=1}^{n} \sum_{j=1}^{n} \color{blue}{\sqrt{|x_i + x_j|}} \le 0 \\
\Rightarrow \sum_{i=1}^{n} \sum_{j=1}^{n} \color{blue}{\sqrt{|x_i + x_j|}} - \sum_{i=1}^{n} \sum_{j=1}^{n} \color{red}{\sqrt{|x_i - x_j|}} \ge 0 \\
\end{equation}
$

$
\begin{equation}
f(X) = \sum_{i=1}^{n} \sum_{j=1}^{n} \color{blue}{\sqrt{|x_i + x_j|}} - \sum_{i=1}^{n} \sum_{j=1}^{n} \color{red}{\sqrt{|x_i - x_j|}}, \\
X = \{x_1, x_2, ...., x_n\}
\end{equation}
$

$
\begin{equation}
F(X) = \sum_{i=1}^{n} \sum_{j=1}^{n} \color{blue}{\sqrt{|x_i + x_j|}} - \color{red}{\sqrt{|x_i - x_j|}}
\end{equation}
$

Our goal is to show:

$
\begin{equation}
F(X) \ge 0; \forall X \in \mathbb{R}^n
\end{equation}
$

Now let us consider the term $\color{blue}{\sqrt{|x_i + x_j|}} - \color{red}{\sqrt{|x_i - x_j|}}$, which we are going to call $f(i,j)$, i.e.

$
\begin{equation}
f(i,j) = \color{blue}{\sqrt{|x_i + x_j|}} - \color{red}{\sqrt{|x_i - x_j|}}
\end{equation}
$

This now converts the original $F(X)$ as follows:

$
\begin{equation}
F(X) = \sum_{i=1}^{n} \sum_{j=1}^{n} f(i,j)
\end{equation}
$

We can note that $f(i,j) = f(j,i)$

$
\begin{equation}
f(j,i) = \color{blue}{\sqrt{|x_i + x_j|}} - \color{red}{\sqrt{|x_j - x_i|}} \\
= \color{blue}{\sqrt{|x_i + x_j|}} - \color{red}{\sqrt{|-(x_i - x_j)|}} \\
\because |-x| = |x| \\
\Rightarrow f(j,i) = \color{blue}{\sqrt{|x_i + x_j|}} - \color{red}{\sqrt{|x_i - x_j|}} = f(i,j)
\end{equation}
$

We can further utilize this decomposition of $F(X)$ as follows:

$
\begin{equation}
F(X) = \sum_{i=1}^{n} f(i,i) + \sum_{i=1}^{n} \sum_{j=i+1}^{n} f(i,j) + \sum_{j=1}^{n} \sum_{i=j+1}^{n} f(i, j)
\end{equation}
$

Using $f(i,j) = f(j,i)$, we get this simplified form:

$
\begin{equation}
F(X) = \sum_{i=1}^{n} f(i,i) + 2*\sum_{i=1}^{n} \sum_{j=i+1}^{n} f(i,j) \\
\because f(i,i) = \color{blue}{\sqrt{|x_i + x_i|}} - \color{red}{\sqrt{|x_i - x_i|}} = \color{blue}{\sqrt{2|x_i|}} - \color{red}{0} = \color{blue}{\sqrt{2}\sqrt{|x_i|}} \ge 0 \\
\Rightarrow F(X) = \sum_{i=1}^{n} \color{blue}{\sqrt{2}\sqrt{|x_i|}} + 2*\sum_{i=1}^{n} \sum_{j=i+1}^{n} \color{blue}{\sqrt{|x_i + x_i|}} - \color{red}{\sqrt{|x_i - x_i|}}
\end{equation}
$

```python
np.random.seed(1337)

cmap = "coolwarm"
N = 10

def get_F(X):
    f1 = np.sqrt(np.abs(X + X.T))
    f2 = np.sqrt(np.abs(X - X.T))
    F = f1 - f2
    return F, f1, f2

def plot_heatmaps(F, f1, f2):
    fig, ax = plt.subplots(1,3, sharex=True, sharey=True, figsize=(18,6))
    sns.heatmap(f1, ax=ax[0], square=True, cbar=False, cmap=cmap, vmin=0)
    sns.heatmap(f2, ax=ax[1], square=True, cbar=False, cmap=cmap, vmin=0)
    sns.heatmap(F, ax=ax[2], square=True, cbar=False, cmap=cmap, center=0, annot=True, fmt=".2f")

    for i in range(3):
        eqn = "$\sqrt{|x_i + x_j|}$"
        if i == 1:
            eqn = "$\sqrt{|x_i - x_j|}$"
        if i == 2:
            eqn = "$\sqrt{|x_i + x_j|} - \sqrt{|x_i - x_j|}$"
            eqn = f"{eqn}; $F(X)={F.sum():.2f}$"
        ax[i].set_title(eqn, fontsize=20)
        ax[i].set_xlabel("i", fontsize=20)
        ax[i].set_ylabel("j", fontsize=20)

    fig.tight_layout()
```

```python
X = (np.random.rand(N) - 0.5)*100
X = X[:, np.newaxis]
X
```

    array([[ 33.596788  ],
           [ 11.02373322],
           [  6.45189795],
           [ 49.73891432],
           [-24.52758636],
           [-48.5622143 ],
           [-41.09549328],
           [ 43.89826005],
           [ 47.34593994],
           [ -0.85184702]])

```python
F, f1, f2 = get_F(X)
plot_heatmaps(F, f1, f2)
```

Now without loss of generality we can assume that $x_i$ are sorted in ascending order, i.e.

$x_1 \le x_2 \le ... \le x_n$

```python
X = np.sort(X, axis=0)
F, f1, f2 = get_F(X)
plot_heatmaps(F, f1, f2)
```

```python
a = np.ones_like(X)
a.T@F@a
```

    array([[40.7975076]])

```python
## Simulation to check psd
evals = []
for N in range(2,100):
    for i in range(1000):
        X = (np.random.rand(N) - 0.5)*100
        X = X[:, np.newaxis]
        X = np.sort(X, axis=0)
        F, f1, f2 = get_F(X)
        e = np.linalg.eigvals(F).min()
        evals.append(e)
min(evals)
```

    0.00025519921280370284

```python
np.linalg.eigvals(F)
```

    array([47.08346898,  7.78699484,  4.74105443,  3.41421932,  2.54024474,
            0.93296417,  2.00557604,  1.63489447,  1.39446185,  1.25605382])

```python
X1 = X - X.min()
F, f1, f2 = get_F(X1)
plot_heatmaps(F, f1, f2)
```

```python
print(X.max()*2*N)
X1 = X - X.max()
F, f1, f2 = get_F(X1)
plot_heatmaps(F, f1, f2)
```

Lets now come back to

$
\begin{equation}
F(X) = \sum_{i=1}^{n} \color{blue}{\sqrt{2}\sqrt{|x_i|}} + 2*\sum_{i=1}^{n} \sum_{j=i+1}^{n} \color{blue}{\sqrt{|x_i + x_i|}} - \color{red}{\sqrt{|x_i - x_i|}} \\
\end{equation}
$

And focus on the term
$
\begin{equation}
2*\sum_{i=1}^{n} \sum_{j=i+1}^{n} \color{blue}{\sqrt{|x_i + x_i|}} - \color{red}{\sqrt{|x_i - x_i|}} \\
= 2*\sum_{i=1}^{n} \sum_{j=i+1}^{n} f(i,j); i < j, x_i \le x_j
\end{equation}
$

$
\begin{equation}
f(i,j) = \color{blue}{\sqrt{|x_i + x_j|}} - \color{red}{\sqrt{|x_i - x_j|}} \\
= \sqrt{|x_i|}(\color{blue}{\sqrt{|1 + \frac{x_j}{x_i}|}} - \color{red}{\sqrt{|1 - \frac{x_j}{x_i}|}}) \\
\end{equation}
$

Let $
\begin{equation}
r_{ij} = \Large{\frac{x_j}{x_i}} \\
\because x_j \ge x_i \\
\Rightarrow \Large{\frac{x_j}{x_i}} \ge 1; x_i > 0 (this~is~wrong~when~x_i \le 0)\\
\Rightarrow r_{ij} \ge 1 \\
f(i,j) = \sqrt{|x_i|}\left[\color{blue}{\sqrt{|1 + r_{ij}|}} - \color{red}{\sqrt{|1 - r_{ij}|}}\right]
\end{equation}
$

$
\begin{equation}
f(i,j) = \sqrt{|x_i|}\left[\color{blue}{\sqrt{|1 + r_{ij}|}} - \color{red}{\sqrt{|1 - r_{ij}|}}\right] \\
\because \sqrt{a} - \sqrt{b} = \Large{\frac{(\sqrt{a} - \sqrt{b})*(\sqrt{a} + \sqrt{b})}{\sqrt{a} + \sqrt{b}}} = \Large{\frac{a - b}{\sqrt{a} + \sqrt{b}}}\\
f(i,j) = \sqrt{|x_i|}\Large{\frac{\left[\color{blue}{{|1 + r_{ij}|}} - \color{red}{{|1 - r_{ij}|}}\right]}{\left[\color{blue}{\sqrt{|1 + r_{ij}|}} + \color{red}{\sqrt{|1 - r_{ij}|}}\right]}} \\
\end{equation}
$

$
\begin{equation}
f(i,j) = \sqrt{|x_i|}\Large{\frac{\left[\color{blue}{{|1 + r_{ij}|}} - \color{red}{{|1 - r_{ij}|}}\right]}{\left[\color{blue}{\sqrt{|1 + r_{ij}|}} + \color{red}{\sqrt{|1 - r_{ij}|}}\right]}} \\
\because r_{ij} \ge 1 \\
\Rightarrow 1 + r_{ij} \ge 2 \\
\Rightarrow |1 + r_{ij}| = 1 + r_{ij} \\
\end{equation}
$
Similarly,
$
\begin{equation}
-r_{ij} \le -1 \\
\Rightarrow 1 - r_{ij} \le 0 \\
\Rightarrow |1 - r_{ij}| = r_{ij} - 1 \\
\end{equation}
$

Using $|1 + r_{ij}| = 1 + r_{ij}$ and $|1 - r_{ij}| = r_{ij} - 1$

$
\begin{equation}
f(i,j) = \sqrt{|x_i|}\Large{\frac{\left[\color{blue}{{1 + r_{ij}}} - \color{red}{{r_{ij} + 1}}\right]}{\left[\color{blue}{\sqrt{|1 + r_{ij}|}} + \color{red}{\sqrt{|1 - r_{ij}|}}\right]}} \\
= \sqrt{|x_i|}\Large{\frac{2}{\left[\color{blue}{\sqrt{|1 + r_{ij}|}} + \color{red}{\sqrt{|1 - r_{ij}|}}\right]}} \\
\end{equation}
\Rightarrow f(i,j) \ge 0
$

Substituting $f(i,j) \ge 0$ in $F(X)$, we get the following:

$
\begin{equation}
F(X) = \sum_{i=1}^{n} \color{blue}{\sqrt{2}\sqrt{|x_i|}} + 2*\sum_{i=1}^{n} \sum_{j=i+1}^{n} f(i,j) \\
\ge \sum_{i=1}^{n} \color{blue}{\sqrt{2}\sqrt{|x_i|}} \ge 0 \\
\Rightarrow \Large{F(X) \ge 0}
\end{equation}
$

Hence, proved.

## Recap of core concepts

* Symmetrix Matrix view of the problem and decomposition of $F(X)$ using $f(i,j) = f(j,i)$
* Ordering of the $x_i$ to make the problem simpler
* Simplification of $\sqrt{a} - \sqrt{b} = \Large{\frac{(\sqrt{a} - \sqrt{b})*(\sqrt{a} + \sqrt{b})}{\sqrt{a} + \sqrt{b}}} = \Large{\frac{a - b}{\sqrt{a} + \sqrt{b}}}$
* Simplication of ${|1+r_{ij}|} = 1 + r_{ij}$ and ${|1 - r_{ij}|} = r_{ij} - 1$, if $r_{ij} \ge 1$.

# More videos to come

Shubhanshu Mishra (shubhanshu.com)

## References

* International Mathematics Olympiad 2021 Problems - https://www.imo-official.org/problems.aspx

Slide Notebooks/IMO 2021 - Problem 2.ipynb


# SciPy: Python tools for science
*Monica Rizzo, 2019*

Among the many useful python packages for data analysis, SciPy is one of the most commonly used and widely applicable. In this notebook we're going to be going over some of the functions in Scipy you can use to fit and model data.

First things first, lets import all of the important stuff.

```python
import numpy as np
from math import factorial
from scipy.optimize import curve_fit
from scipy.interpolate import interp1d
from scipy.optimize import brentq

from scipy.integrate import quad
from scipy.misc import derivative

import matplotlib.pyplot as plt

#this makes the plot labels bigger (looks nicer)
plt.rcParams.update({'font.size': 18})
```

## 1. Regression

There are a lot of things you can do with data. Often times, we will want to either 1) see how well data agrees with a given model or 2) without assuming anything, see what sort of model agrees the best with some data.
One way to test this is by using regression techniques. "Regression" is just the math word for fitting data to a certain function. You may have done something similar to this before with Excel, or by hand using graph paper.

Lucky for us, SciPy can fit things for you! First, lets generate some data that follows a known function.
1. **Using numpy, create an array of n equally spaced x values within some range - doesn't really matter what range**
2. **Then choose two values for the slope m and y intercept b**
3. **Finally, create a numpy array of y values which follows the equation for a line (*hint: take advantage of numpy array addition/multiplication*)**

```python
#number of data points to generate
n = #complete

#create a numpy array of x values
x_values = #complete

#define slope and y intercept
m = #complete
b = #complete

#make an array of y values that follow the equation of a line
y_values = #complete
```

Now let's plot the data to see what it looks like (always a good first step).
1. **Make a scatter plot of the x and y values you created above**

```python
plt.figure(figsize=(10, 8))
#complete
plt.grid(True)
#ALWAYS LABEL YOUR AXES
plt.xlabel(#complete
plt.ylabel(#complete
plt.show()
```

Obviously, we already know what this function looks like, so we already know what Scipy should return as the answer.

The Scipy function we are going to use to do the fitting is called *curve_fit*. Using a user-defined python function that describes the fit (in this case a line), *curve_fit* will attempt to find the best fitting parameters for some input data. So, in this case, we want to create a fit function that returns a line and takes 2 free parameters: the slope and the y intercept. Then, we want to hand *curve_fit* that fit function, the x data, and the y data, and it's going to tell us what it thinks the best fitting parameters are, along with the uncertainties on those parameters (which we're not going to worry about right now).

1. **Define a function which takes as inputs x, m, and b, and returns the y values corresponding to a line**
2. **Pass *curve_fit* the name of your fit function, your x data, and your y data (see documentation here: https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.curve_fit.html)**
3. **Print the result of *curve_fit* (note: result should be a 2 element array. The values are in the same order as the arguments of your fit function)**

```python
def fit_function():
    #complete

```

```python
result, cov = curve_fit(#complete

#print the result

```

Nice.

Now let's look at a more realistic example. In real life, data is noisy and usually has some associated uncertainty. Lets add a little bit of noise to our existing data, and assign some uncertainty.

1. **Add the noise values to your array of y values**
2. **Create an array - which will be your y error - of length n where all of the values are 3**

```python
#numpy has built-in capabilities to generate random numbers :)
#this will generate random numbers drawn from a bell curve
y_noise = np.random.normal(0, 3, n)

#array of y values plus noise
y_noisy_values = #complete

#array uncertainty values
y_error = #complete
```

Now, let's plot the data again.

***Hint: in order to plot the errors, instead of using the 'plot' function, use 'errorbar': https://matplotlib.org/3.1.0/api/_as_gen/matplotlib.pyplot.errorbar.html***

```python
plt.figure(figsize=(10, 8))
#try using fmt='o' and markersize=6.0
plt.errorbar(#complete
#complete
```

Not so pretty anymore, but maybe still looks like a line? Let's see what SciPy thinks.
1. **Using the same syntax as before, pass *curve_fit* your fit function, your x data, your noisy y data, and also the uncertainty on your y values - this can be done using the 'sigma' option.**
2. **As before, print the result**

```python
#use curve fit to find the fit parameters for your noisy data
result, cov = curve_fit(#complete

#print the results
```

As you can see, this answer is pretty different from the actual values of the parameters. We can plot what the fit looks like over our data to visually assess how bad/good it is.
1. **Create a numpy array of *n_test* x values to plot, that has a range that is slightly bigger that the range of your x values**
2. **Pass this array, along with the best fit m and b values, to your fit function, and store the result as a separate array**
3. **Copy the cell you used to plot the data with errorbars, and add a line to plot the test x and y values**

```python
#number of test values
n_test = 500

#array of x test values
x_test_values = #complete

#array of y test values
y_test_values = fit_function(#complete
```

```python
plt.figure(figsize=(10, 8))
#complete
```

Keep in mind that these answers are just guesses - they also have some uncertainty to them. The other thing that *curve_fit* returns - which we're calling *cov* - describes the uncertainites on the fit parameters. Let's see what those look like.

```python
#cov is a matrix, and we only want the diagonal values
unc = np.sqrt(np.diag(cov))

print("param1_err = %0.2f and param2_err = %0.2f" % (unc[0], unc[1]))
```

**Given what your data looks like, do you think these numbers are reasonable? Do your fit parameters to the noisy data agree with the real parameters within this uncertainty?**

## 2. Interpolation

Now let's switch gears a little bit. Suppose we have some data where instead of the function parameters (slopes, y-intercepts, etc.), we want to approximate what the values between our data points are. So in this case, we don't care about the specifics of the fit, we just want to connect the dots. This can be achieved using *interpolation* techniques.

First, lets again create some data that follows a known function:
1. **Using numpy arrays, create an array of 10 x values that ranges from 0 to 2$\pi$**
2. **Create an array of y values which is the sine of the x values**

```python
#array of x values
x_values = #complete

#array of y values
y_values = #complete
```

As per *always*, plot your data before doing anything else.

```python
plt.figure(figsize=(10, 8))
#complete
```

This could be sine curve - and I'm sure if you connected the dots by hand you would draw something that resembles a sine function. But let's see how well we can approximate this by numerically interpolating. To start, let's use the default settings of *interp_1d*. When passed x and y values, *interp_1d* will return a function which interpolates the data at every point between the maximum and minimum x value. We can then plot the values of this interpolation function over a range of test x values.

1. **Pass *interp1d* your x and y values (see documentation here: https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.interp1d.html#scipy.interpolate.interp1d)**
2. **Create an array of *n_test* test x values between 0 and 2$\pi$**
3. **Pass your interpolation function your array of test x values and store the result as test y values**
4. **Plot your test values as a line over a scatter plot of your data**

```python
#create an interpolation function
interp_func = interp1d(#complete

#array of test x values
x_test_values = #complete

#recall that interp_func is a function, so it can be used like any other function
y_test_values = #complete
```

```python
plt.figure(figsize=(10, 8))
#complete
```

As you can see, *interp_1d* connects the data points with straight lines by default. This isn't always the best way to interpolate data; we can tell just by looking at what we've plotted that our data most likely has some curvature to it. It's also generally better to interpolate data smoothly, because this makes performing operations like differentiation easier/more accurate. Lets try out some different types of interpolation.

To change the way *interp_1d* connects our data points, we can use the 'kind' option.
1. **Try setting kind to each of the following: 'zero', 'quadratic', 'cubic'. For each different kind, create a new variable name to store the interpolating function in.**
2. **Using the same method as above, plot the data and the interpolated approximation on top of each other. You can use the same test x values, but will need to create new test y values for each plot**

```python
interp_func_zero = interp1d(x_values, y_values, kind='zero')
#repeat for all interpolation kinds

```

```python
#test y values

#zero plot
```

```python
#test y values

#quadratic plot
```

```python
#test y values

#cubic plot
```

How can we judge which of these kinds of interpolation is the best? We know what the values of the function *should* be, so we can just pick a point and see how different the approximate answer is from the real answer. In the cell below, check to see how different each of the approximate y values is from the real value where x = 4.

1. **Compute the real value - sin(4) - and store in a variable**
2. **For each of the interpolation functions, compute and store in different variables the approximate y value at an x value of 4**
3. **Calculate and print the percent error of each of the approximate y values**

*Hint: calculate percent error using: \begin{equation}\frac{|approximation - real|}{|real|} \times 100\end{equation}*

*Hint: python has a built in absolute value function called 'abs'*

```python
real_value = np.sin(4)

approx_value_linear = interp_func(4)
#repeat for all interpolation functions

#calculate percent error for each interpolation function
```

Based on this, which appears to be the best kind of interpolation? This happens to be the best kind in general.

*Optional*: We can also check to make sure the area under the curve is what we expect it to be. For a sine curve in the range zero to 2$\pi$, the total area should be zero - this is because half the the curve is above zero and half is below. In order to find the area under the curve, we can *integrate* it - SciPy also has functions to do this. The particular algorithm we're going to use is called *quad* (which is all around the most accurate in SciPy). *quad* takes as arguments the name of the function you want to integrate, the upper bound of the integration, and the lower bound of the integration.

1. **For each of the interpolation functions, use quad to integrate from 0 to 2$\pi$ (see documentation here: https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.quad.html#scipy.integrate.quad)**
2. **Print out the result of each. Why do you think the results are what they are?**

*Note: when dealing with computers, numbers smaller than or close to $10^{-16}$ are considered zero*

```python
res_linear, err = quad(interp_func, 0, 2*np.pi)
#repeat for all interpolation functions

```

## 3. Root Finding

Now let's talk - briefly - about something completely different. In math classes, you may have needed to find the zeros of functions before (set y equal to zero and solve for x). Finding zeros is very important in the area of numerical optimization - which we won't talk about explicity - but can often be more complicated than just setting y equal to zero and solving for x. In fact, when you have just data points, you often can't do this at all.

And that's where SciPy comes in. Suppose we have some interpolating function which passes through zero, and we want to know approximately where. One of the best algorithms to do this is something called *Brent's method* (which you may or may not learn about in college). *Brent's method* will attempt to find where a function crosses zero in a given interval; it needs to know the endpoints to look between in order to work.

First, lets create some data to try this out on:
1. **Using numpy, create an array of 10 equally spaced x values between -2 and 2**
2. **Create an array of y values that is the x values cubed plus 3**

```python
#array of x values
x_values = #complete

#array of y values
y_values = #complete
```

Plot your data.

```python
plt.figure(figsize=(10, 8))
#complete
```

Based on the plot, at what x value do you expect the zero point to be?

Let's see what SciPy thinks. We're going to use a function called *brentq*. This function takes the name of the function you want the zero of, the upper bound of that function, and the lower bound of that function. So first, we need to interpolate the data, and then pass that interpolant to *brentq*.

1. **Using the best interpolation kind from the previous section, create and store an interpolating function for your x and y data**
2. **Pass *brentq* the name of your interpolating function, the lower bound of the function, and the upper bound of the function (see documentation here: https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.brentq.html#scipy.optimize.brentq)**
3. **Print the result. Does this answer seem correct?**

```python
#interpolating function
interp_func = interp1d(#complete

#find zero using brentq
x0 = brentq(#complete

#print result

```

## 4. Fun With Real Data

Now let's tackle a (sort of) real-life example. Using the provided data file, we will attempt to fit the data it contains with a more complicated fit function. First, we will load and plot the data.

1. **Load the data and plot it (with error bars)**

```python
#load data and store it to an array
#columns: x values, y values, y error
data = #complete

#split the data into separate x, y, and y_err arrays using numpy array slicing
x_values = #complete
y_values = #complete
y_error = #complete
```

```python
#plot. that. data
plt.figure(figsize=(10, 8))
#complete
```

What do you think this function looks like? Using your knowledge of regression, try to think of a fitting function which might model this data, and then fit that model to the data.

1. **Define a fitting function, and fit it to the data using *curve_fit***
2. **Print the best fit parameters and their uncertainties**
3. **Plot the data and the fit on top of each other**

*Hint: try to fit a higher order polynomial*

```python
#create a fitting function
def :

#use curve fit to fit the function to the data
result, cov = #complete

#print the result and the uncertainties
#you can probably just copy and paste code from the first section, with minor adjustments

```

```python
#create test x values, and test y values from the fit

#plot the data and the fit
plt.figure(figsize=(10, 8))

#complete

plt.show()
```

This data actually has a function which fits it best, for science reasons. In order to fit that function we're going to cheat a little and find the location of the peak, i.e. where the slope of the function equals zero. To do this, we need to first interpolate our data, and then find where the slope (or *derivative*) equals zero.

1. **Interpolate the data using the best kind of interpolation**
2. **Define a function that takes x values and returns the derivative of your interpolation function at those x values (already completed for you)**
3. **Pass the derivative function to *brentq* (remember you need to pass it a lower and upper bound), and store the result. Print this result. Does it make sense?**

```python
#interpolation function
interp_func = interp1d(#complete

#function which takes x values and returns derivative
def deriv_func(x):

    deriv = derivative(interp_func, x, dx = 1e-6)

    return deriv

#find root of derivative function
x0 = brentq(#complete

#print result
```

The function we want to fit to is:

\begin{equation}f(x) = a ((x/x0)^{b1/-0.8} + (x/x0)^{b2/-0.8})^{-0.8},\end{equation}

where a, b1, and b2 are the three fit parameters. Now that we know x0, we can use *curve_fit* to find the remaining 3 parameters.

1. **Define a function called 'broken_power_law' which takes x values, a, b1, and b2 as input and returns the values corresponding to the above expression**
2. **Use curve fit to fit this function to the data, *including the error***
3. **Plot the data and the fit on top of each other. Would you say this fit is good or bad? Why or why not?**

```python
#create a function to return the expression above
def broken_power_law:
    #complete
```

```python
#fit the above function using curve fit
result, cov = curve_fit(broken_power_law, x_values, y_values, sigma=y_error)
```

```python
#create test x and y values

#plot the data with error and the fit
plt.figure(figsize=(10, 8))
#complete
```

## More Fun Stuff: Bonus

As far as cool and useful tools go, SciPy has a *ton* of other nifty functions (with decent documentation), for almost every application. The data above is from a course I took on radiative processes in astrophysics (which is a complicated way of saying "light curves: the class"). You can learn various things about the object which produced this light curve by finding the slope of each side of the logarithm of the data. This is just one idea of something you can do with this data besides just fitting it.

See what else you can do with it!

Some ideas:

* try a different fitting function
* interpolate the data
* find the point where the slope of the data is zero
* integrate/differentiate the data

```python
##play around with the data and see what you can do with it
##if you need help or suggestions, don't hesitate to ask!

```

## Resources

A set of SciPy tutorials and documentation can be found here: https://docs.scipy.org/doc/scipy/reference/index.html

The Scipy Lectures site also has a pretty comprehensive guide to all of the modules: https://scipy-lectures.org/intro/scipy.html

Here is a nifty explanation of the difference between regression and interpolation (which I found useful): https://stats.stackexchange.com/questions/33659/how-is-interpolation-related-to-the-concept-of-regression

Here is a lecture (which is honestly pretty heavy, but if you want to look through it you can) on numerical root finding methods: http://butler.cc.tut.fi/~piche/numa/lecture0506.pdf

```python

```


# Data Fitting Exercises

Assessment

This exercise is part of your assessment. When you have finished, please save and upload your .ipynb file to Moodle.

To save, use **File > Save and Checkpoint** in the Jupyter menu.

Before saving and uploading, be sure to check that your notebook runs as expected from a fresh start (using **Kernel > Restart and Run All** from the Jupyter menu).

# 3. Non-Linear Least Squares

In Exercise 1, you used `scipy.stats.linregress` to perform a linear regression to some experimental kinetics data, modelled using the linear form of the Van't Hoff equation.

In Exercise 2, you learned how to solve the same problem by writing two functions: a model function, and an error function; then finding an optimised set of model parameters that minimise your errors, using `scipy.optimize.minimize`.

The second approach is not restricted to fitting linear equations, and can generally be used to fit any model function to some reference dataset. In this, third, exercise, you will use the same technique to fit a **non-linear** model describing some flash photolysis data.

## Flash Photolysis

Flash photolysis is a technique for investigating fast photochemical reactions, by exposing the reactant to very brief and intense flashes of light, and then spectroscopically analysing the resulting products. The data collected give a measure of fluorescence intensity with respect to time. The time scale is usually of the order of picoseconds, while the measure of intensity is usually the photon count.

The equation that governs a fluorescence lifetime is a biexponential — it depends on two exponential terms, and for this reason cannot be converted to a linear form by taking logs.

\begin{equation}
I \propto A_1\mathrm{e}^{-k_{1}t} + A_2\mathrm{e}^{-k_{2}t}\tag{1}.
\end{equation}

The pre-exponential factors $A_1$ and $A_2$ govern the contribution of each exponential decay (green and red line) to the overall signal. $t$ is the time, and $k_{1}$ and $k_{2}$ are the respective radiative rate constants for each lifetime.

In a bi-exponential plot such as the one displayed in the Figure below, there are three key features. The blue curve is the flash or pump signal. It is a very brief flash from the light source and will always show up in the signal. The green curve is the exponential decay signal from a short fluorescence lifetime. The red curve is the exponential decay signal from a longer fluorescence lifetime.

The radiative rate constant for each component is related to the natural lifetime, $\tau_r$, by

\begin{equation}\tau_r = \frac{1}{k_r}\end{equation}

where the natural lifetime may be thought of as the average amount of time a molecule will remain in its excited state before it decays. By calculating $k_1$ and $k_2$, one will be able to determine the natural lifetime as

\begin{equation}k_r = k_1 + k_2\end{equation}

For an given experiment, the parameters $A_1$, $k_1$, $A_2$, and $k_2$ can be obtained by fitting Equation 1 to the experimental data. The fitted values of $k_1$ and $k_2$ can then be used to calculate the natural lifetime.

## Performing a least-squares fit to flash photolysis data

The file `data/flash_photolysis.dat` contains a set of experimental data from a flash photolysis experiment. This file looks like

```
Time (ps)  Signal (counts)
0.0  -0.002740138
0.2  0.001982640
0.4  0.014379181
0.6  0.010439134
0.8  -0.014530330
1.0  0.009097397
1.2  -0.000735803
1.4  0.011072823
1.6  -0.004858392
…
```

By copying and editing the appropriate code from Exercise 2, or writing it from scratch, read in the experimental data from 'data/flash_photolysis.dat' and plot it.

To skip the first line in the data file, you can use

np.loadtxt('data/flash_photolysis.dat', skiprows=1)

```python

```

Your plot should show an initial fast rise in intensity, from zero, followed by a slower decay curve, that is approximately exponential.

By copying and editing the appropriate code from Exercise 2, or writing it from scratch, define a function model_function() that takes a set of parameters $P$ and the time $t$ as inputs, and returns the intensity as a function of time, according to equation 1.

Test that your model function can approximate the experimental data by plotting the original data, and your model intensity using

$A_1=0.5$

$A_2=0.5$

$k_1=5\times10^{-4}$

$k_2=5\times10^{-4}$

```python

```

By copying and editing the appropriate code from Exercise 2, or writing it from scratch, write an error function called error_function() that returns the sum of squared errors between your predicted intensity and the experimental intensity.

Use the `minimize()` function from `scipy.optimize` to find the optimal parameter set to describe the experimental data.

### Hints

* Equation 1 only describes the fluorescence decay. It does not include a description of the initial increase in intensity. When you perform your fit, you can limit your fitting data to the decay portion by slicing from the 200th element onwards; e.g. `data[200:]`.

* If you do not specify any additional options for `minimize`, the optimisation algorithm can try any values for the model parameters. Because the intensity model contains terms that look like $\mathrm{e}^{-x}$, if the algorithm tries negative numbers for $k_1$ or $k_2$ this can produce a result too large for the computer to describe, and give an error. In reality, we know that $k_1$ and $k_2$ must be positive. We can include this information in the function minimisation by including an additional setting that places **bounds** on the possible values of $P$:
```python
minimize( … , bounds=((0,1), (0,1), (0,np.inf), (0,np.inf)) )
```
The `bounds` argument takes a list of pairs of numbers, that specify the minimum and maximum allowed values to check. If $P$ is ordered as $(A_1, A_2, k_1, k_2)$, this example specifies that $A_1$ and $A_2$ must be between 0 and 1 (we know that the maximum intensity at $t=0$ is equal to 1), and that $k_1$ and $k_2$ must be between 0 and (positive) infinity.

* The `minimize()` function works by trying different parameter values until the result of the test function falls below a certain convergence tolerance. The default convergence settings for `minimize()` are not small enough to guarantee that you find the optimal set of parameter values. You can specify a lower tolerance setting by using the `tol` keyword:
```python
minimize( …, tol=1e-20)
```

```python

```

Plot the experimental data and your fitted intensity function.

Using your fitted parameters, calculate the natural lifetime.

```python

```


Computers, Waves, Simulations

Finite-Difference Method - Second Derivative

#### This exercise covers the following aspects:
* Initializing a Gaussian test function
* Calculation of numerical second derivative with 3-point operator
* Accuracy improvement of numerical derivative with 5-point operator

**Note:**
* Loop boundaries changed for 5-point operator, May 2020

```python
# Import Libraries
import numpy as np
from math import *
import matplotlib.pyplot as plt
```

We initialize a Gaussian function

\begin{equation}
f(x)=\dfrac{1}{\sqrt{2 \pi a}}e^{-\dfrac{(x-x_0)^2}{2a}}
\end{equation}

Note that this specific definition is a $\delta-$generating function. This means that $\int{f(x) dx}=1$ and in the limit $a\rightarrow0$ the function f(x) converges to a $\delta-$function.

```python
# Initialization
xmax=10.0       # physical domain (m)
nx=100          # number of space samples
a=.25           # exponent of Gaussian function
dx=xmax/(nx-1)  # Grid spacing dx (m)
x0 = xmax/2     # Center of Gaussian function x0 (m)

x=np.linspace(0,xmax,nx)   # defining space variable

# Initialization of Gaussian function
f=(1./sqrt(2*pi*a))*np.exp(-(((x-x0)**2)/(2*a)))

```

```python
# Plotting of gaussian
plt.figure(figsize=(10,6))
plt.plot(x, f)
plt.title('Gaussian function')
plt.xlabel('x, m')
plt.ylabel('Amplitude')
plt.xlim((0, xmax))
plt.grid()
plt.show()
```

Now let us calculate the second derivative using the finite-difference operator with three points

\begin{equation}
f^{\prime\prime}_{num}(x)=\dfrac{f(x+dx)-2 f(x)+f(x-dx)}{dx^2}
\end{equation}

and compare it with the analytical solution
\begin{equation}
f^{\prime\prime}(x)= \dfrac{1}{\sqrt{2\pi a}} ( \dfrac{(x-x_0)^2}{a^2}- \dfrac{1}{a} ) \ e^{-\dfrac{(x-x_0)^2}{2a}}
\end{equation}

```python
# Second derivative with three-point operator

# Initiation of numerical and analytical derivatives
nder3=np.zeros(nx)          # numerical derivative
ader=np.zeros(nx)          # analytical derivative

# Numerical second derivative of the given function
for i in range (1, nx-1):
    nder3[i]=(f[i+1] - 2*f[i] + f[i-1])/(dx**2)

# Analytical second derivative of the Gaissian function
ader=1./sqrt(2*pi*a)*((x-x0)**2/a**2 -1/a)*np.exp(-1/(2*a)*(x-x0)**2)
# Exclude boundaries
ader[0]=0.
ader[nx-1]=0.

# Calculate rms error of numerical derivative
rms = np.sqrt(np.mean((nder3-ader)**2))

```

```python
# Plotting
plt.figure(figsize=(10,6))
plt.plot (x, nder3,label="Numerical Derivative, 3 points", lw=2, color="violet")
plt.plot (x, ader, label="Analytical Derivative", lw=2, ls="--")
plt.plot (x, nder3-ader, label="Difference", lw=2, ls=":")
plt.title("Second derivative, Err (rms) = %.6f " % (rms) )
plt.xlabel('x, m')
plt.ylabel('Amplitude')
plt.legend(loc='lower left')
plt.grid()
plt.show()
```

In the cell below calculation of the first derivative with four points is provided with the following weights:

\begin{equation}
f^{\prime\prime}(x)=\dfrac{-\dfrac{1}{12}f(x-2dx)+\dfrac{4}{3}f(x-dx)-\dfrac{5}{2}f(x) +\dfrac{4}{3}f(x+dx)-\dfrac{1}{12}f(x+2dx)}{dx^2}
\end{equation}

```python
# First derivative with four points

# Initialisation of derivative
nder5=np.zeros(nx)

# Calculation of 2nd derivative
for i in range (2, nx-2):
    nder5[i] = (-1./12 * f[i - 2] + 4./3  * f[i - 1] - 5./2 * f[i] \
                       +4./3  * f[i + 1] - 1./12  * f[i + 2]) / dx ** 2

# Exclude boundaries
ader[1]=0.
ader[nx-2]=0.

# Calculate rms error of numerical derivative
rms=rms*0
rms = np.sqrt(np.mean((nder5-ader)**2))
```

```python
# Plotting
plt.figure(figsize=(10,6))
plt.plot (x, nder5,label="Numerical Derivative, 5 points", lw=2, color="violet")
plt.plot (x, ader, label="Analytical Derivative", lw=2, ls="--")
plt.plot (x, nder5-ader, label="Difference", lw=2, ls=":")
plt.title("Second derivative, Err (rms) = %.6f " % (rms) )
plt.xlabel('x, m')
plt.ylabel('Amplitude')
plt.legend(loc='lower left')
plt.grid()
plt.show()
```

### Conclusions

* 3-point finite-difference approximations can provide  estimates of the 2nd derivative of a function
* We can increase the accuracy of the approximation by using further functional values further
* A 5-point operator leads to substantially more accurate results


# Music machine learning - Advanced neural networks

### Author: Philippe Esling (esling@ircam.fr)

In this course we will cover
1. An introduction to [convolutions](#convolution) and how they can be used
2. Defining a [Convolutional Neural Network](#cnn) in Pytorch for image classification
3. Coding our own [convolutional layer](#layer)
4. An explanation on [recurrent networks](#rnn) in practice

## Convolutions

In purely mathematical terms, convolution is a function derived from two given functions by integration which expresses how the shape of one is modified by the other. In simpler (discrete) terms, the convolution product of a matrix by a smaller one can be seen as _filtering_ the large matrix. Hence we slide the small matrix over the large one and compute local products at each position. Therefore the convolution operator $\star$ computes at each position $n$

$$
(f \star g)[n]=\sum _{m=-M}^{M}f[n-m]g[m].
$$

An example of this operation is shown here

This operation can be used to _filter_ the image (as in the _gaussian blur_ operator), or _detect_ features (such as edges).

Given an 32x32 image with RGB channels, we can represent it as a tensor of shape `(32, 32, 3)` which is (height, width, channels). When we perform convolution, we need a filter that has the same channel depth as the image. For example, we can use a 5x5 filter which is of shape `(5, 5, 3)` and slide it across the image left to right, top to bottom with a stride of 1 to perform convolution. We are going to perform this in numpy, depending on a certain amount of parameters, which define the behavior of our convolution

* `height` and `width`: spatial extend of the filters
* `S`: stride size (number of steps to jump to the next position)
* `P`: amount of padding (adding zeros in the original matrix)

```python
import numpy as np
# Padding
pad = 2
stride = 1
height, width = 5, 5
# A random fake image
x = np.random.randn(3, 32, 32)
x_h, x_w = x.shape[1:]
# Our convolution kernel
weight = np.random.randn(3, 5, 5)
# Padding the original image
x_pad = np.pad(x, pad_width=((0, 0,), (pad, pad), (pad, pad)), mode='constant', constant_values=0)
# We can expect the output size to be
h_out = int(1 + (x_h + 2 * pad - height) / stride)
w_out = int(1 + (x_w + 2 * pad - width) / stride)
# So we will store our result in
y = np.zeros((1, h_out, w_out))
```

Computing the convolution itself can be performed by using the following loop (which amounts to _slide_ our kernel across the large matrix)

```python
# Sliding the kernel across the matrix
for h in range(h_out):
    for w in range(w_out):
        i, j = h * stride, w * stride
        conv_sum = np.sum(x_pad[:, i:i+height, j:j+width] * weight)
        y[0, h, w] = conv_sum
```

We can witness the effect of this operation with the following plot

```python
######################
# YOUR CODE GOES HERE
######################
```

## Convolutional Neural Networks (CNN)

Convolutional NNs (CNNs) rely on convolution in place of general matrix multiplication. They are specialized for processing data with a known grid-like topology and are among the best performing systems in classification/recognition tasks. Each layer in a CNN consists in a set of $N$ _filters_ called _kernels_, that are convolved across the input. If we denote as $\{k^l_n\}_{n\in[1;N]}$ the set of kernels for layer $l$, these all share a unique _kernel size_. By convolving each one of its $N$ kernels across a d-dimensional input $x$, a convolutional layer produces $N$ d-dimensional outputs called _feature maps_, denoted as $\{a^l_n\}_{n\in[1;N]}$. Hence, the computation of the $n$-th activation map in layer $l$ for input $x$ is defined as:

$$
\begin{equation}
a^l_n = \sum_{m=1}^{M} k^l_n \star x_m + b^l_n
\end{equation}
$$

Thus, as depicted in the following Figure, the feature map corresponding to kernel $n$ consists in the sum of the d-dimensional discrete convolutions (denoted by the $\star$ operator) between the kernel $n$ and each one of the d-dimensional data $\{x_m\}_{m\in[1;M]}$, plus a bias $b$. A convolutional layer is thus a 3-dimensional tensor $h \in \mathcal{T}_{N,I,J}(\mathbb{R})$ where $N$ is the number of features maps while $I$ and $J$ are respectively the _width_ and _height_ of the maps.

This can be seen as replacing our _neurons_ by _feature detectors_ (the convolutional kernels), which will increasingly process the image. In the following, we will first use the high-level interface of `Pytorch` to define a CNN

## Defining a CNN in Pytorch

Defining a convolutional network in Pytorch is quite easy, as we can rely on the `nn` module, which contains all the required layers

```python
import torch
import torch.nn as nn
```

In the previous course, we have seen that we could define our network in a very simple way, by using the `Sequential` model definition. Here we define a CNN followed by a MLP, as seen in the previous course

```python
# Define the input dimensions
in_size = 1000
# Use the nn package to define our model and loss function.
model = nn.Sequential(
    nn.Conv2d(3, 6, 5),
    nn.ReLU(),
    nn.MaxPool2d(2, 2),
    nn.Conv2d(6, 16, 5),
    nn.ReLU(),
    nn.MaxPool2d(2, 2),
    nn.Flatten(),
    nn.Linear(16 * 5 * 5, 120),
    nn.ReLU(),
    nn.Linear(120, 84),
    nn.ReLU(),
    nn.Linear(84, 10),
    nn.Softmax(dim=1)
)
```

Another way to implement the networks in Pytorch is to use the `functional` approach. In this version, each layer is seen as a function, that we apply on sucessive inputs. For instance, we can define one layer of fully-connected units and apply it to some inputs as follows.

```python
import torch.nn.functional as F
# Define one layer
layer = nn.Linear(100, 10)
# Define the non-linearity
activation = nn.ReLU()
# Create some random input
inputs = torch.rand(32, 100)
# Apply our layers
outputs = activation(layer(inputs))
# Equivalently, as ReLU is parameter-free
outputs = F.relu(layer(inputs))
```

Now to make this even cleaner, we can define our own `nn.Module`, which is a `Pytorch` class representing models. To do so, we can define a sub-class, and implement the functions `__init__` (defining our layers) and `forward` (explaining how our forward pass will behave)

***

**Exercise**
1. Implement a CNN in Pytorch using the `functional` library

***

```python
import torch.nn as nn
import torch.nn.functional as F

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool1 = nn.MaxPool2d(2, 2)

        ######################
        # YOUR CODE GOES HERE
        ######################

    def forward(self, x):
        l1 = F.relu(self.conv1(x))
        l2 = F.pool1(l1) ** 2
        ...
        o = ???(lN)
        ######################
        # YOUR CODE GOES HERE
        ######################

        return o

model = CNN()
```

### Training the network

In order to test our CNN, we are going to try to perform image classification. To do so, we can use the simplifications for data loading contained in `torchvision`

```python
import torchvision
import torchvision.transforms as transforms
```

The `torchvision` package provides pre-coded simplification to download and use the major image datasets, notably `MNIST` and `CIFAR`, which are the baseline datasets for testing image ML models. The output of torchvision datasets are PILImage images of range [0, 1]. We transform them to Tensors of normalized range [-1, 1].

In the following code, we are going to load the `CIFAR10` _train_ and _test_ sets. **Note that this code will automatically download the dataset if you did not have it before, and place it in the `data` folder, so this might take a bit of time**

```python
# Transforms to apply to the images
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# Import the train dataset
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)
# Import the test dataset
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2)
# Classes in the CIFAR dataset
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
```

Using your knowledge from the previous course, you can now define an optimization problem, and implement the training loop for your model

***

**Exercise**
1. Define a `criterion` and `optimizer`
2. Fill in the training loop to train your model

***

```python
import torch.optim as optim

######################
# YOUR CODE GOES HERE
######################

criterion = ...
optimizer = ...

```

```python
n_epochs = 100
# Training loop
for epoch in range(n_epochs):
    running_loss = 0.0
    # Go through all batches
    for i, data in enumerate(trainloader, 0):
        # Get the inputs; data is a list of [inputs, labels]
        inputs, labels = data

        ######################
        # YOUR CODE GOES HERE
        ######################

        # Keep track of the loss
        running_loss += loss.item()
    print(f'loss at epoch {epoch}: {running_loss}')
```

Now that your model is trained, you can test it by feeding some new (unseen) images and see if it is able to classify them correctly

## Coding our own convolutional layer

Although `Pytorch` comes packed with pre-implemented layers, we can also very easily define our own layers. This will be useful when you start doing research and propose your own way of processing the information. A large advantage of `Pytorch` is that it performs _automatic gradient differentiation_, this means that we simply have to define how the `forward` pass will work, and `Pytorch` will automatically infer the backpropagation equations, without us having to go through any complicated differentiation

In the following, we are going to redefine the `Conv2d` layer, by computing the operation ourselves.

***

**Exercise**
1. Complete the `forward` function to compute a convolution

***

```python
class MyConv2d(nn.Module):
    def __init__(self, n_channels, out_channels, kernel_size, dilation=1, padding=0, stride=1):
        super(MyConv2d, self).__init__()
        self.kernel_size = (kernel_size, kernel_size)
        self.kernal_size_number = kernel_size * kernel_size
        self.out_channels = out_channels
        self.dilation = (dilation, dilation)
        self.padding = (padding, padding)
        self.stride = (stride, stride)
        self.n_channels = n_channels
        self.weights = nn.Parameter(torch.Tensor(self.out_channels, self.n_channels, self.kernal_size_number))

    def forward(self, x):
        width = self.calculateNewWidth(x)
        height = self.calculateNewHeight(x)
        windows = self.calculateWindows(x)

        ######################
        # YOUR CODE GOES HERE
        ######################

        return result

    def calculateWindows(self, x):
        windows = F.unfold(x, kernel_size=self.kernel_size, padding=self.padding, dilation=self.dilation, stride=self.stride)
        windows = windows.transpose(1, 2).contiguous().view(-1, x.shape[1], self.kernal_size_number)
        windows = windows.transpose(0, 1)
        return windows

    def calculateNewWidth(self, x):
        return ((x.shape[2] + 2 * self.padding[0] - self.dilation[0] * (self.kernel_size[0] - 1) - 1) // self.stride[0]) + 1

    def calculateNewHeight(self, x):
        return ((x.shape[3] + 2 * self.padding[1] - self.dilation[1] * (self.kernel_size[1] - 1) - 1)// self.stride[1]) + 1

# Testing the code directly
conv = MyConv2d(3, 1, 3)
x = torch.randn(1, 3, 24, 24)
out = conv(x)
out.mean().backward()
# Check that we do have gradients
print(conv.weights.grad)
```

Now you can use your own `MyConv2d` layer, and use it in real-life scenarios, by trying to change your previous model to use your own layer instead.

## Recurrent Neural Networks (RNN)

Recurrent neural networks (RNNs) are a family of models designed to process time series and sequential data, which perform remarkably in applications such as speech recognition or machine translation. The ability of RNNs to model correlations between successive computations through recurrent connection make them efficient for temporal problem as they provide a form of _memory_.

To model structured sequential data, NNs can be augmented with recurrent loops, which allow to retain information across time steps. Considering a sequence $\mathbf{X}=\{\mathbf{x}_t\}$, dependencies between elements are managed by having a recurrent hidden state $\mathbf{h}_t$ at time $t$ in the network. The value of $\mathbf{h}_t$ at each time depends of the previous time and the input, as depicted in the following figure.

Formally, each hidden state is updated as

$$
\begin{equation}
\mathbf{h}_t =
\begin{cases}
\phi_{\mathbf{\mathbf{\theta}}}(\mathbf{x}_0) & \mbox{if } t=0 \\
\phi_{\mathbf{\mathbf{\theta}}}(\mathbf{h}_{t-1},\mathbf{x}_t), & \mbox{otherwise}
\end{cases}
\label{eq:RNNhiddenupdate}
\end{equation}
$$

### Implementing a simple rnn

Here's a simple one-to-many vanilla recurrent neural network example in functional form. If we were to produce `h[t]`, we need some weight matrices, `h[t-1]`, `x[t]` and a non-linearity `tanh`.

$$
h_{t} = tanh(W_{hh}h_{t-1} + W_{xh}x_{t} + B_{h})
$$

Since this is a **one-to-many** network, we'd want to produce an output `y[t]` at every timestep, thus, we need another weight matrix that accepts a hidden state and project it to an output.

$$
y_{t} = W_{hy}h_{t} + B_{y}
$$

Now that we know how to use the `Functional` library of `Pytorch`, we are going to implement our own simple RNN layer as previously. This time, we do not provide the content of the `__init__` function, so think carefully of what parameters you will need and how you need to define them.

***

**Exercise**
1. Complete the `__init__` function to instantiate all required parameters
1. Complete the `forward` function to compute the forward pass

***

```python
import torch.nn as nn
import torch.nn.functional as F

class RecurrentNetwork(nn.Module):

    def __init__(self, input_size, hidden_size, output_size):
        """ Defines the required variables """
        super(RecurrentNetwork, self).__init__()

        ######################
        # YOUR CODE GOES HERE
        ######################

    def forward(self, x):
        """ Computes the forward pass """

        ######################
        # YOUR CODE GOES HERE
        ######################

        return output
```

We can then forward propagate our information inside our network

```python
in_size = 4
input_vector = torch.ones(16, in_size)
simple_network = RecurrentNetwork(in_size, 10, 5)
# Notice that same input, but leads to different ouptut at every single time step.
print(simple_network(input_vector))
print(simple_network(input_vector))
print(simple_network(input_vector))
```

As usual, `Pytorch` also comes packed with some pre-coded recurrent layers. You can go check the documentation to find how to use these.