# Tensorflow Experiments on Neural Ordinary Differential Equations

> You can contact me on twitter as [@mandubian](http://twitter.com/mandubian)

This notebook is a sandbox to test concepts exposed in this amazing paper:

> Neural Ordinary Differential Equations
> http://arxiv.org/abs/1806.07366
>
> Authors:
>    Chen, Ricky T. Q.
>    Rubanova, Yulia
>    Bettencourt, Jesse
>    Duvenaud, David

My idea is to reproduce the concepts exposed in the paper fully in Tensorflow using Eager-Mode and GradientTape.

I didn't want to depend on Autograd and to be able to use classic Keras models.

For ODE Solver, I wanted to compare implementations in TF with Scipy very robust ones (the only one I found is Runge-Kutta Dopri5).

## A few preliminary ideas (to let you scratch your head with me)

> Skip it if you just want to see it running

### Modeling the derivative of hidden state

Deep neural networks as we know in general are aimed at simulating a function $y = f(x)$ from a set of known inputs/outputs (a supervised learning based on a groundtruth). The classic approach is to represent the function with inputs, outputs and a _hidden state_ we don't know but want to learn. This hidden state is represented by a sequence of discrete deep neural layers whose internal parameters are learnt by step by step back-propagation of gradients on a chosen derivable loss function and an optimization mechanism to reach the best minimum (ideally global). So the hidden state to be learnt is parameterized by the trainable parameters of those discrete layers.

>In _Neural ODE_ paper (http://arxiv.org/abs/1806.07366), the authors propose a new approach by modeling not the hidden state of the function but the derivative of the hidden state.

### Residual layer as Ordinary Differential Equation

The authors make this assumption starting from the case of well-known residual neural networks _ResNet_ constituted of multiple residual layers. A residual layer is meant to _keep in memory_ what you knew in previous layer and to slighly update it without forgetting it. A residual layer can be represented by the following equation:

$$
z_{t+1} = z_t + f(z_t, {\theta}_t) \\
$$

where:
- $z_t$ is the input of the layer
- $z_{t+1}$ is the output of the layer
- $f$ is the function parameterized by $θ_t$ that we want to learn
- $t$ can be seen as the depth of this layer in a full ResNet with $0, 1, ..., t, t+1, ..., t+n$ successive residual layers.

Previous equation can be rewritten as:

$$
z_{t+1} - z_t = f(z_t, {\theta}_t) * ({\delta}_t = 1) \\
or \\
\frac {z_{t+1} - z_t} {({\delta}_t = 1)} = f(z_t, {\theta}_t)
$$

${\delta}_t = 1$ can be seen as discrete depth increment going from one residual layer to the next in the ResNet.

Let's now consider this depth $t$ is no more discrete but continuous. We can easily infer the following:

$$
\begin{aligned}
{\frac {dz(t)}{dt}} = f(z(t), t, {\theta})
\end{aligned}
$$

The previous equation is called an _ordinary differential equation (ODE)_. There is a whole field of study in mathematics about solving ordinary differential equations, with so-called _ODE Solvers_, i.e solving equations of kind:

$$
\text{Find } z(t) \text{ knowing }
\begin{aligned}
{\frac {dz(t)}{dt}} & = f(z(t), t) \\
\end{aligned}
\text{ given an initial value }
\begin{aligned}
z(t_{0}) & = z_{0} \\
\end{aligned} \text{ and a discretized interval }
\begin{aligned}
{t_0, t_1, ... t_n}
\end{aligned}
$$

ODE Solvers are numerical methods to compute fast but robust approximated integrals (analytical methods exist but can't solve many cases of ODE like $\dot{y} = sin(x^2 + y^2)$). During your math study, you might have seen this Euler method https://en.wikipedia.org/wiki/Euler_method of approximating a function by using a naive construction with the tangent, an incremental step and an initial point on the curve. I'm not a mathematician so I won't teach you anything about ODE Solvers, you just need to know that this field has been studied and used intensively in multiple real-life apps (physics, geology, meteo, finance, chemistry etc...) and provide very efficient, optimized algorithms implemented in multiple libs and for example Scipy provides nice implementations of different ODE solvers.

### Computing hidden state using ODE Solver

Let's imagine we have a ResNet with multiple residual layers of depths $0, 1, ..., t, t+1, ..., t+n$. In a classic ResNet, all layers have their own function $f_t$ but if we mutualize this function for all layers, we can generalize to the whole network:

$$
\begin{aligned}
{\frac {dz(t)}{dt}} & = f(z(t), t)\hspace{0.5cm}\forall{t} \in [t_0, t_1, ...t_n]
\end{aligned}
$$

So the $f$ function of this ResNet can be seen as the derivative of the hidden state ${\frac {dz(t)}{dt}}$ we are trying to learn.

Using ODE Solver, if we know $f$ and an initial value $z(t_{0})$ and this interval discretization ${t_0, t_1, ..., t_n}$, we can compute a robust approximation $z(t)$ in the following way:

$$
\begin{aligned}
z(t_n) & = ODESolve(z(t_0), f, [t_0, t_1, ...t_n], θ)
\end{aligned}
$$

where:
- $t_0$ is the initial layer
- $[t_0, t_1, ...t_n]$ is the discretized interval
- $z(t_0)$ is the initial value of the function
- $z(t_n)$ is the output value of our network
- $t_n$ is the output layer
- ${\theta}$ are the parameters of the function $f$ we want to learn

### Depth as an input parameter of neural model

$[t_0, t_1, ...t_n]$ represents the depth distribution of the layers of our _mono-function_ ResNet. In a classic network, the depth in terms of layers is a static part of the architecture of the network which can't be changed without retraining completely the network.

In our case, depth distribution becomes a continuous space that we discretized and provide as a dynamic parameter to the ODE Solver at training or inference. We can imagine having different depth distribution at training or inference to speed computations for example. We could even imagine using more complicated depth distribution like non linear distribution or even learning an optimized shape for a given problem. Next studies...

There is also an immediate analogy between this depth distribution and time distribution. It also quickly reminds of time-series issues and recurrent neural networks... They speak about it in the paper but I still haven't implemented that part... Next steps...

### Back-propagation by backward ODE Solving

So that's cool but how do you represent and learn those ${\theta}$???

I won't go in deep demonstration as it would be too long. Don't worry if you don't catch this part :D

It starts with a classic neural network representing our derivative $f$ with internal parameters $\theta$.

Let simplify to a 1-step time interval $[t_0, t_1]$, an initial input value $z(t_0)$ and our known output $groundtruth(t_1)$.

First, it computes current estimated value $z({t_1})$ by a classic forward computation with a call to our ODE Solver using current $\theta$ and input value $z(t_0)$.

$$
\begin{aligned}
z(t_1) = ODESolve(z(t_0), f, [t_0, t_1], {\theta})
\end{aligned}
$$

Then, it computes the classic Loss function:

$$
\begin{aligned}
L(z(t_1), groundtruth(t_1)) & = L(ODESolve(z(t_0), f, [t_0, t_1], {\theta}), groundtruth(t_1))
\end{aligned}
$$

The authors then consider the so-called _adjoint state_ (adjoint state method is the math _trick_ at the origin of backward auto-differentiation _aka_ gradient back-propagation that often simplifies gradient computation as neural networks are often many-to-scalars functions with lower gradient complexity in backward mode... This is another topic)

$$
\begin{aligned}
a(t) = − {\frac {\partial L(z(t))}{\partial z(t)}}
\end{aligned}
$$

The paper appendix proves that:

$$
\begin{aligned}
{\frac {da(t)}{dt}} = - a(t){\frac {\partial{f(z(t), t, \theta)}}{\partial z}}
\end{aligned}
$$

and the same respectively for ${\frac {da(t)}{d\theta}}$ and ${\frac {da(t)}{dt}}$

${\frac {\partial{f(z(t), t, \theta)}}{\partial z}}$ is the classic gradient we already compute in normal residual layer using back-propagation.

Knowing all the derivatives of $a(t)$ with respect to $z(t), t, \theta$, we can then compute $a(t)$ with a call to ODE Solver going backward in time from $t_1$ to $t_0$. The initial value is the gradient of loss ${\frac {\partial L(z(t_1))}{\partial z}}$ at $t_1$ between our current output $z(t_1)$ and our groundtruth...

$$
\begin{aligned}
s0 & = [z(t1), {\frac {\partial L}{\partial z(t_1)}}, 0, −{\frac {\partial L}{t_1}}] \\
[z(t_0), {\frac {\partial L}{\partial z(t_0)}}, {\frac {\partial L}{\partial θ}}, {\frac {\partial L}{\partial t}}] & = ODESolve(z(t_0), augmented_f, [t_1, t_0], {\theta})
\end{aligned}
$$

By using a bit more machinery (an _augmented state_ grouping $z(t)$, ${\frac {da(t)}{dz}}$, ${\frac {da(t)}{d\theta}}$, ${\frac {da(t)}{dt}}$ together), the paper demonstrates how all required backward gradients can be computed in one single call to ODESolver and then you can train your parameters using a classic optimizer like SGD or Adam...

Not so clear? It's normal :D

### Conclusion

By representing the residual-like hidden-state through its derivative and a _continuous_ depth distribution, it's possible to train model parameters with ODE Solvers. The hidden-state of the derivative is still represented by a static human-engineered discrete sequence of neural layers but the depth distribution becomes an input parameter of the network that can be changed dynamically at training or inference.

Is it going to change your life as a datascientist drastically? Not yet... But it's bringing new cool ideas on how we can represent neural networks in another way, how it learns, how residual layers work. It opens new fields of study, drawing bridges between Deep Learning and existing older branches of mathematics. It also bring diversity in DL domain that IMHO tends to re-use the same tricks without always catching how they work.

#### Other sources of inspiration

There are more precise details in this nice article https://rkevingibson.github.io/blog/neural-networks-as-ordinary-differential-equations/ and then read the paper first pages, it's quite simple after all. If I can understand it, you can too...

This notebook is also very inspired by this cool article on Neural ODE & appealing Julia:

https://julialang.org/blog/2019/01/fluxdiffeq

If you want to know more about ODE Solvers, there is a whole litterature about it with lots of different methods like `Runge-Kutta` or `LSODA` used in this notebook and you can find more details in scipy
https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.odeint.html#scipy.integrate.odeint

#### Code under MIT License

```
Copyright 2019 Pascal Voitot

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
```

## Initialization

```python
from __future__ import absolute_import
from __future__ import print_function

import time
import os
import numpy as np
import scipy.integrate
from plotly import __version__
from plotly.offline import download_plotlyjs, init_notebook_mode, plot, iplot
#import plotly.plotly as py
import plotly.graph_objs as go
import plotly.io as pio

import tensorflow as tf
print("tensorflow version: ", tf.__version__)

```

    tensorflow version:  1.12.0

```python
tf.enable_eager_execution()
init_notebook_mode(connected=True)
```

```python
config = tf.ConfigProto()
config.gpu_options.allow_growth=True

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
```

## Test ODE Solvers (TF and Scipy)

Let a derivative function:

$
\hspace{1cm}f'(t) = 2t
$

Let a set of values $t \in [0, 4.0]$

Let the initial value be $f(t_0) = 0.0$

A good ODE Solver should integrate to a parabolic:

$
\hspace{1cm}f(t) = t^2
$

## Test Full-TF ODE Solver (Runge-Kutta - Adaptative - Explicit - Dopri5)

> `Dopri5` is an explicit runge-kutta method of order (4)5 due to Dormand & Prince (with stepsize control and dense output).
>
> Authors:
> E. Hairer and G. Wanner Universite de Geneve, Dept. de Mathematiques CH-1211 Geneve 24
>
> **Please note that this implementation works well but the method tends to be less stable computationally speaking than the scipy LSODA implicit method used below**

```python
ts = np.linspace(0., 4.0, 10)

ys, info = tf.contrib.integrate.odeint(lambda y, t: tf.to_float(2*t), 0.0, ts, rtol = 0.01, full_output=True)

iplot([go.Scatter(x=ts, y=ys.numpy())])
```

## Test Scipy ODE Solver (LSODE)

> By default Scipy is using on LSODA Solver method (there are other methods, even DOPRI5)
>
> Real-valued Variable-coefficient Ordinary Differential Equation solver, with fixed-leading-coefficient implementation. It provides automatic method switching between implicit Adams method (for non-stiff problems) and a method based on backward differentiation formulas (BDF) (for stiff problems).
>
> Source: http://www.netlib.org/odepack
>
> **Please note that this implementation is more stable computionally speaking than full-TF implementation above**

```python
ts = np.linspace(0., 4.0, 10)

ys = scipy.integrate.odeint(lambda y, t: 2*t, [0.0], ts)

iplot([go.Scatter(x=ts, y=np.squeeze(ys))])
```

## Illustrating ODE Solvers with oscillators

### Two-dimensional damped oscillator

http://www.dam.brown.edu/people/mraissi/research/7_multistep_neural_networks/

Let us consider the two-dimensional damped harmonic oscillator with cubic dynamics:

$
\begin{align}
{\frac {dx} {dt}} & = −0.1x^3 + 2.0y^3 \\
{\frac {dy} {dt}} & = −2.0x^3 − 0.1y^3 \\
z & = (x, y)
\end{align}
$

```python
# z = [x, y]
def damped_oscillator_2d(z, t, A):
    return np.dot(z**3, A)

A = np.array([
    # −0.1x^3 + 2.0y^3
    [-0.1, 2.0],
    # −2.0x^3 − 0.1y^3
    [-2.0, -0.1]])

# Initial u0 = [x0, y0]
z0 = np.array([2., 0.]).T

# Time sampling
N = 100  # Dataset size
max_T = 6.0 # Time sampling
ts = np.linspace(0., max_T, N)

```

#### TF-RungeKutta ODE Solver

```python
zs, info = tf.contrib.integrate.odeint(lambda z0, ts: tf.tensordot(z0**3, A, axes=1), z0, ts,
                                method="dopri5", options={"max_num_steps":10, "first_step": 0.1}, rtol = 0.001,
                                full_output=True)
zs = zs.numpy()
xs = zs[:, 0]
ys = zs[:, 1]
iplot([go.Scatter(x=ts, y=xs, name="x"), go.Scatter(x=ts, y=ys, name="y")])
iplot([go.Scatter(x=xs, y=ys)])
```

#### Scipy ODE Solver

```python
zs = scipy.integrate.odeint(damped_oscillator_2d, z0, ts, args=(A,))

xs = zs[:, 0]
ys = zs[:, 1]
iplot([go.Scatter(x=ts, y=xs, name="x"), go.Scatter(x=ts, y=ys, name="y")])
iplot([go.Scatter(x=xs, y=ys)])
```

## Helpers to call ODE Solvers from TF

```python
def ode_tf(f, z0, ts, args):
    """calls Full TF ODE Runge-Kutta Dopri5 Solver

    Args:
        f: a function of type (z0, ts, args) -> ys
            where z0, the initial value
                  ts, the time values
                  args, the supplementary args

        z0: initial value
        ts: time values on which to solve equation
        args: supplementary arsg

    Returns:
        array of solutions computed by ODE Solver
    Please note, TF solver params are hard-coded here but they could be managed in a cleaner way
    """
    # current TF odeint implementation doesn't manage args so we need this lambda
    zs, info = tf.contrib.integrate.odeint(
        lambda z0, ts: f(z0, ts, args), z0, ts,
        method="dopri5",
        options={"max_num_steps":1000, "first_step": 0.001},
        rtol = 0.001,
        full_output=True
    )
    return zs

def ode_scipy(f, z0, ts, args):
    """calls Python Scipy ODE Solver from TF

    Args:
        f: a function of type (z0, ts, args) -> ys
            where z0, the initial value
                  ts, the time values
                  args, the supplementary args

        y-z0: initial value
        ts: time values on which to solve equation
        args: supplementary args

    Returns:
        array of solutions computed by ODE Solver
    """
    # managing args as None... this is ugly but simplifies a bit later code
    if args is not None:
        r = tf.py_func(
            lambda z0, ts, args: scipy.integrate.odeint(f, z0, ts, args=(args,)),
            [z0, ts, args], z0.dtype)
    else:
        r = tf.py_func(
            lambda z0, ts: scipy.integrate.odeint(f, z0, ts, args=(None,)),
            [z0, ts], z0.dtype)
    return r

```

## Lotka–Volterra equations (the funny predators/preys or rabbits/foxes equations)

> This sample was taken from article https://julialang.org/blog/2019/01/fluxdiffeq

https://en.wikipedia.org/wiki/Lotka%E2%80%93Volterra_equations

$
\begin{aligned}
{\frac {dx}{dt}} & = \alpha x - \beta xy, \\
{\frac {dy}{dt}} & = \delta xy - \gamma y
\end{aligned}
$

where

$x$ is the number of preys (for example, rabbits);
$y$ is the number of predators (for example, foxes);

${\frac {dy}{dt}}$ and ${\frac {dx}{dt}}$ represent the instantaneous growth rates of the two populations;

$t$ represents time

$α, β, γ, δ$ are positive real parameters describing the interaction of the two species.

```python
# Numpy or eager-mode TF
def lotka_volterra(z, ts, args):
    x, y = z[0], z[1]
    α, β, δ, γ = args[0], args[1], args[2], args[3]
    dx = α*x - β*x*y
    dy = -δ*y + γ*x*y
    return tf.stack([dx, dy], axis=0)

# Tensorflow Variables
D = 2
# time interval from 0 to 10 divided in 100 steps
N = 100
ts = tf.constant(np.linspace(0., 10.0, N))
# initial arguments
args = tf.Variable(tf.constant([1.5, 1.0, 3.0, 1.0], shape=(4,), dtype=ts.dtype))
z0 = tf.Variable(tf.constant([1.0,1.0], shape=(2,), dtype=ts.dtype))

```

#### TF-RungeKutta ODE Solver

```python
sol = ode_tf(lotka_volterra, z0, ts, args)
xs = sol[:, 0].numpy()
ys = sol[:, 1].numpy()

iplot([go.Scatter(x=ts.numpy(), y=xs, name="x"), go.Scatter(x=ts.numpy(), y=ys, name="y")])

iplot([go.Scatter(x=xs, y=ys)])
```

#### Scipy ODE Solver

```python
sol = ode_scipy(lotka_volterra, z0, ts, args)
xs = sol[:, 0].numpy()
ys = sol[:, 1].numpy()

iplot([go.Scatter(x=ts.numpy(), y=xs, name="x"), go.Scatter(x=ts.numpy(), y=ys, name="y")])

iplot([go.Scatter(x=xs, y=ys)])
```

## The ugly Augmented State in Tensorfow Eager-mode + GradientTape

This is an implementation of the augmented-state gradient computed by ODE-Solver code sample found in Neural ODE paper made in pure TF 1.x in eager-mode and using Gradient-Tape.

To be clear, that was a true pain to implement it and it's still not stable IMHO...

I don't pretend this is nice or idiomatic code but it works now and enables experimentations...

```python

def flatten(params):
    """flatten a list of tensors/arrays into a flat tensor and returns also the unflatten function
    able to rebuild original list from flattenned tensor
    """
    if len(params) == 0:
        def unflatten(flat_params):
            return flat_params
        return np.array(params), unflatten

    original_shapes = []
    flat_params = []
    default_dtype = params[0].dtype
    for p in params:
        sh = p.shape
        original_shapes.append(sh)
        flat_sh = 1
        for s in sh:
            flat_sh *= s

        res = tf.reshape(p, (flat_sh,))
        res = tf.cast(res, default_dtype)
        flat_params.append(res)

    def unflatten(flat_params):
        cur_idx = 0
        params = []
        for sh in original_shapes:
            flat_sh = 1
            for s in sh:
                flat_sh *= s
            param = flat_params[cur_idx:cur_idx + flat_sh]
            param = tf.reshape(param, sh)
            cur_idx += int(flat_sh)
            params.append(param)

        return params

    flat_params = tf.concat(flat_params, axis=0)
    return flat_params, unflatten

# Directly inspired from function implemented with autograd at the end of the paper
# Extended from "Scalable Inference of Ordinary Differential
# Equation Models of Biochemical Processes", Sec. 2.4.2
# Fabian Froehlich, Carolin Loos, Jan Hasenauer, 2017
# https://arxiv.org/abs/1711.08079
def grad_odeint(z_ts, func, ts, args, T, D, ode_solver, variables = None):
    """
    This code is still ugly and not very robust IMHO...
    It's full of reshape that might be unneeded but at least I know what I manipulate without compilted types

    Args:
        z_ts: the values with ODE Solver in forward path on ts time values
        func: the adjoint gradient function to integrate with ODE Solver
        ts: the discretized time values
        args: the derivable arguments of func
        ode_solver: the ode_solver function (z0, ts, args) -> zs
        variables: the hidden TF/Keras tensor variables (separated from args to be manageable as they have their own life cycle I'm not sure to understand yet)

    Returns:
        grad: the custom augmentate-state TF gradient function
    """
    # let's store
    if variables is not None:
        flat_variables, unflatten_variables = flatten(variables)

    def unpack_D_2D(x, D):
        """simple extract """
        x = tf.squeeze(x)
        return (
            x[0:D],
            x[D:2 * D]
        )

    # augmented_state = [z_t, a_z_t, a_t, a_args, a_vars]
    def a_gradient_aug(augmented_state, t, args):
        augmented_state = tf.convert_to_tensor(augmented_state, dtype=z_ts.dtype)
        t = tf.convert_to_tensor(t, dtype=ts.dtype)
        t = tf.reshape(t, (1,))
        args = tf.convert_to_tensor(args, dtype=args.dtype)
        z_t, a_z_t = unpack_D_2D(augmented_state, z_ts.shape[1])

        # the internal gradient-tape computing second derivative
        with tf.GradientTape() as gg:
            gg.watch(args)
            gg.watch(a_z_t)
            gg.watch(t)
            gg.watch(z_t)
            dz_dt = func(z_t, t, args)

        if variables is not None:
            da_z_dt, da_t_dt, da_args_dargs, da_dvars = gg.gradient(
                dz_dt, [z_t, t, args, variables], output_gradients = -a_z_t)
        else:
            da_z_dt, da_t_dt, da_args_dargs = gg.gradient(
                dz_dt, [z_t, t, args], output_gradients = -a_z_t)
            da_dvars = None

        if da_z_dt is None:
            da_z_dt = tf.zeros_like(z_t, dtype=z_t.dtype)
        if da_t_dt is None:
            da_t_dt = tf.zeros_like(t, dtype=t.dtype)
        if da_args_dargs is None:
            da_args_dargs = tf.zeros_like(args, dtype=args.dtype)
        if da_dvars is None:
            da_dvars = tf.zeros_like((1,))
        else:
            da_dvars, _ = flatten(da_dvars)
        r, _ = flatten([dz_dt, da_z_dt, da_t_dt, da_args_dargs, da_dvars])
        return r

    def grad(gt):
        """the TF custom gradient function"""

        # initialization of backward gradient computation starting at the end of time interval
        dl_dz_t = gt[-1, :]
        dl_dt0 = tf.constant([0.], shape=(1,), dtype=ts.dtype)
        dl_dargs = tf.zeros_like(args, dtype=ts.dtype)
        if variables is not None:
            dl_dvars = tf.zeros_like(flat_variables, dtype=flat_variables.dtype)
        else:
            dl_dvars = tf.zeros((1,))
        dl_dt_list = []

        # compute gradient back-propagation by ODE solver calls going backward in time
        for i in range(T - 1, 0, -1):
            # Compute effect of moving back in time (just porting formulas found in paper)
            ztsi = z_ts[i, :]
            gti = gt[i, :]
            # Compute the adjoint gradient of loss at current time
            fi = func(ztsi, ts[i], args)
            dl_dt = tf.tensordot(fi, tf.transpose(gti), axes=1) #, axes=[[1], [0]])
            dl_dt = tf.reshape(dl_dt, (1,))
            dl_dt_list.append(dl_dt)
            dl_dt0 = dl_dt0 - dl_dt
            dl_dt0 = tf.reshape(dl_dt0, shape=(1,))

            # build augmented state
            aug_z0, unflatten_aug = flatten([ztsi, dl_dz_t, dl_dt0, dl_dargs, dl_dvars])
            # Run augmented system backwards to the previous observation using ODE_Solver
            cur_ts = tf.reverse(tf.slice(ts, [i-1], [2]), axis=[0])
            a_aug = ode_solver(lambda aug, t, args: a_gradient_aug(aug, t, args), aug_z0, cur_ts, args)

            # get at ts[i - 1]
            a_aug = a_aug[1]
            _, dl_dz_t, dl_dt0, dl_dargs, dl_dvars = unflatten_aug(a_aug)

            # Add gradient from current output.
            dl_dz_t = dl_dz_t + gt[i - 1, :]

        dl_dt0 = tf.reshape(dl_dt0, (1,))
        dl_dt_list.append(dl_dt0)
        dl_dts = tf.stack(dl_dt_list, axis=0)
        dl_dts = tf.reverse(dl_dts, axis=[0])

        # output z_t is not needed but can't put None in TF tensors
        n_o_n_e = tf.zeros_like(z_ts) #, dtype=z_ts.dtype)

        if variables is not None:
            return [n_o_n_e, dl_dz_t, dl_dts, dl_dargs], unflatten_variables(dl_dvars)
        else:
            return [n_o_n_e, dl_dz_t, dl_dts, dl_dargs], []
    return grad

def tf_ode_grad(f, N, D, ode_solver):
    """The plug in custom_gradient TF mechanism"""
    @tf.custom_gradient
    def custom_grad(z0, ts, args):
        z0 = tf.squeeze(z0)
        # forward values computation with ODE Solver
        ys = ode_solver(f, z0, ts, args)
        def grad_fn(dr, variables=None):
            # backward gradient computation using backward ODE Solver
            (_, dl_dz_t, dl_dts, dl_dargs), dl_dvars = grad_odeint(ys, f, ts, args, N, D, ode_solver, variables)(dr)
            return [dl_dz_t, dl_dts, dl_dargs], dl_dvars
        return ys, grad_fn
    return custom_grad

```

## A gradient-descent Neural-ODE optimizer based on loss function

The following function optimizes iteratively a function defined by its derivative f using an ODE Solver, a ground_truth and a time interval sampling.

The derivative function is parameterized by a set of custom arguments and/or a set of hidden Tensorflow variables (the weights & bias in Keras mode for example). Those parameters are those we want to learn during training by gradient descent controlled by an Adam optimizer.

I had to separate Tensorflow variables from other arguments because they are managed differently in TF `@custom_gradient`. That was the deepest pain during this study because there is no doc or samples about those features except very basic cases which were useless to me. Anyway, by forcing this separation, I could reach a runnable state that seems to optimize my parameters as you'll see below.

```python
from IPython import display

def loss_L2(y_true, y_pred):
    """The classic derivable least-square loss function"""
    return tf.reduce_sum(tf.square(y_true - y_pred))

def optimize(f, ode_solver, loss_fn, labels, z0, ts,
             args=None, variables=None,
             epochs=25,
             x_axis_min=-0.25, x_axis_max=10.25,
             y_axis_min=-0.25, y_axis_max=7.0,
             x2_axis_min=-0.25, x2_axis_max=10.25,
             y2_axis_min=-0.25, y2_axis_max=7.0,
             learning_rate=0.1, epsilon=1e-7, output_dir="images", start = 0):
    """A sample optimization process that trie to optimize args/variables parameters of a function
    defined by its derivative f integrated with ODE Solver according to a loss function loss_fn, groundtruth labels,
    initial input value and time interval sampling using Adam optimizer (with plots & image saving)"""
    def custom_grad(z0, ts, the_args):
        with tf.GradientTape() as tape:
            tape.watch(z0)
            tape.watch(the_args)
            ode = tf_ode_grad(f, N, D, ode_solver)(z0, ts, the_args)
            loss = loss_fn(labels, ode)

        # managing stupidly the abscence of args or variables...
        if variables is not None:
            grads, grad_vars = tape.gradient(loss, [[the_args], variables])
        else:
            grads = tape.gradient(loss, [the_args])
            grad_vars = []

        the_grads = []
        for grad in grads:
            grad = tf.reshape(grad, the_args.shape)
            grad = tf.cast(grad, the_args.dtype)
            the_grads.append(grad)

        the_grad_vars = []
        for i, grad in enumerate(grad_vars):
            the_grad_vars.append(tf.cast(grad, variables[i].dtype))

        return loss, ode, the_grads, the_grad_vars

    # Adam Optimizer
    optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate,
                                       epsilon=epsilon)

    # plotting layout
    layout = go.Layout(
        xaxis=dict(
            range=[x_axis_min, x_axis_max],
        ),
        yaxis=dict(
            range=[y_axis_min, y_axis_max],
            domain=[0.55, 1.0],
        ),
        xaxis2=dict(
            range=[x2_axis_min, x2_axis_max],
        ),
        yaxis2=dict(
            range=[y2_axis_min, y2_axis_max],
            domain=[0, 0.45],
        ),
        title="",
        height=680, width=1024
    )

    # no args, replacing it by a fake tf.Variable for now... to be removed later
    if args is None:
        var_args = tf.Variable(0, dtype=tf.float64)
    else:
        var_args = args

    # the training loop
    global_step = tf.Variable(start)
    best_loss = 1e400
    best_zss = []

    if not os.path.exists(output_dir):
        os.mkdir(output_dir)

    for epoch in range(epochs):
        loss, ode, grads, grad_vars = custom_grad(z0, ts, var_args)

        print("loss:", loss)
        #print("ode:", ode)
        #print("grads:", grads)

        grad_params = list(zip(grads, [var_args]))
        if variables is not None:
            grad_params.extend(list(zip(grad_vars, variables)))

        optimizer.apply_gradients(grad_params, global_step)

        #loss, ode, grads, grad_vars = custom_grad(z0, ts, var_args)
        ys = ode.numpy()
        last_ts = ts.numpy()
        last_ys = ys

        if loss < best_loss:
            best_loss = loss
            best_var_args = var_args
            #best_ts = ts.numpy()
            best_zs = ys
            best_zss.append(best_zs)

        # More Plotting
        display.clear_output(wait=True)
        data = []

        for zs in best_zss[-1:]:
            data.extend([
                go.Scatter(x=ts.numpy(), y=zs[:, 0], name="best_x",
                           mode = 'lines+markers',
                           line = dict(
                                color = ('rgb(34, 119, 158)'),
                                width = 1)),
                go.Scatter(x=ts.numpy(), y=zs[:, 1], name="best_y",
                           mode = 'lines+markers',
                           line = dict(
                                color = ('rgb(201, 79, 14)'),
                                width = 1)),
                go.Scatter(x=zs[:, 0], y=zs[:, 1], name="best_xy",
                           xaxis="x2", yaxis="y2",
                           mode = 'lines+markers',
                           line = dict(
                                color = ('rgb(8, 173, 99)'),
                                width = 1)),
            ])

        data.extend([
            go.Scatter(x=ts.numpy(), y=labels.numpy()[:, 0], name="groundtruth_x",
                       line = dict(
                            color = ('rgb(34, 119, 158)'),
                            width = 2,
                            dash = 'dot')),
            go.Scatter(x=ts.numpy(), y=labels.numpy()[:, 1], name="groundtruth_y",
                       line = dict(
                            color = ('rgb(201, 79, 14)'),
                            width = 2,
                            dash = 'dot')),
            go.Scatter(x=labels.numpy()[:, 0], y=labels.numpy()[:, 1], name="groundtruth_xy",
                       xaxis="x2", yaxis="y2",
                       line = dict(
                            color = ('rgb(173, 7, 7)'),
                            width = 2,
                            dash = 'dot')),
        ])
        if loss >= best_loss:
            data.extend([
                go.Scatter(x=ts.numpy(), y=last_ys[:, 0], name="current_x",
                           line = dict(
                                    color = ('rgb(24, 24, 24)'),
                                    width = 1,
                                    dash = 'dash')),
                go.Scatter(x=ts.numpy(), y=last_ys[:, 1], name="current_y",
                           line = dict(
                                    color = ('rgb(24, 24, 24)'),
                                    width = 1,
                                    dash = 'dash')),
                go.Scatter(x=last_ys[:, 0], y=last_ys[:, 1], name="current_xy",
                           xaxis="x2", yaxis="y2",
                           line = dict(
                                color = ('rgb(24, 24, 24)'),
                                width = 2,
                                dash = 'dash')),
            ])

        layout.title = f"Step {global_step.numpy()}<br>Best Loss {best_loss.numpy()}<br>Current Loss {loss.numpy()}"
        fig = go.Figure(data=data, layout=layout)
        iplot(fig)

        pio.write_image(fig, f"{output_dir}/step_{global_step.numpy()}.png")

        #print(f"Step: {global_step.numpy()}, Loss: {loss.numpy()}")

```

## Optimizing lotka_volterra

Let's optimize lotka_volterra parameters using our optimization function with a dumb loss function that is the squared distance of the output values from 1 to check it works.

### Scipy ODE Solver

#### Model

```python
# Tensorflow Variables
# the dimension of our input vector [x, y] => 2
D = 2
# time interval sampling
N = 25
ts = tf.constant(np.linspace(0., 10.0, N))
# the arguments we want to optimize
args = tf.Variable(tf.constant([1.5, 1.0, 3.0, 1.0], shape=(4,), dtype=ts.dtype))
z0 = tf.Variable(tf.constant([1.0,1.0], shape=(2,), dtype=ts.dtype))

def ode_scipy_solver(f, z0, ts, args):
    r = ode_scipy(f, z0, ts, args)
    r = tf.cast(r, z0.dtype)
    return r

ground_truth_ones = tf.ones(shape=(N, D), dtype=z0.dtype)

```

#### Training Video

```python
from IPython.display import HTML

HTML("""

""")
```

#### Running Optimization by yourself

```python
optimize(lotka_volterra, ode_scipy_solver, loss_L2,
         ground_truth_ones, z0, ts, args=args, variables=None, epochs=80,
         y_axis_min=-0.25, y_axis_max=7.0,
         x_axis_min=-0.25, x_axis_max=10.25,
         output_dir="lotka_volterra_optim_scipy")
```

### TF Runge-Kutta ODE Solver (Diverging)

#### Model

```python
# Tensorflow Variables
# the dimension of our input vector [x, y] => 2
D = 2
# time interval sampling
N = 25
ts = tf.constant(np.linspace(0., 10.0, N))
# the arguments we want to optimize
args = tf.Variable(tf.constant([1.5, 1.0, 3.0, 1.0], shape=(4,), dtype=ts.dtype))

# the initial input value
z0 = tf.Variable(tf.constant([1.0,1.0], shape=(2,), dtype=ts.dtype))

def ode_tf_solver(f, z0, ts, args):
    ts = tf.contrib.framework.sort(ts)
    r = ode_tf(f, z0, ts, args)
    r = tf.cast(r, z0.dtype)
    return r

# our ground_truth (all 1.0)
ground_truth_ones = tf.ones(shape=(N, D), dtype=z0.dtype)

```

#### Training Video

>Look how it diverges at a first period of optimization.
>This is due to computational instability of ODE Solver Runge-Kutta Dopri5 method.
>That's why in next parts of this notebook, I'll use on Scipy LSODA implementation that is more much stable.

```python
from IPython.display import HTML

HTML("""

""")
```

#### Running Training by yourself

```python
optimize(lotka_volterra, ode_tf_solver, loss_L2,
         ground_truth_ones, z0, ts, args=args, variables=None, epochs=150,
         y_axis_min=-0.25, y_axis_max=7.0,
         x_axis_min=-0.25, x_axis_max=10.25,
         output_dir="lotka_volterra_optim_tf")
```

## Final Step: learning time-series with Neural ODE

Finally, let's try to learn some time-series function with our neural ODE to check paper assertions are not just imaginary.

We are going to model derivative of our function with a classic one-layer dense neural layer as shown in https://julialang.org/blog/2019/01/fluxdiffeq.

Then we compute our ground-truth using ODE-Solver from true derivative, decide to use L2 loss.

Finally, we launch iterative optimization process using ODE augmented state gradient descent and observe what happens...

I won't launch it with Runge-Kutta as it is much less stable than scipy LSODA... tensors are flowing from GPU to CPU but it's OK for current experimentations...

### Learning a piece of lotka_volterra curve

#### Model

```python
D = 2
N = 20
ts = tf.constant(np.linspace(0., 3.5, N))
args = tf.Variable(tf.constant([1.5, 1.0, 3.0, 1.0], shape=(4,), dtype=ts.dtype))
z0 = tf.Variable(tf.constant([1.0, 1.0], shape=(2,), dtype=ts.dtype))

model = tf.keras.Sequential([
  tf.keras.layers.Dense(25, activation=tf.nn.relu, input_shape=(1, D)),  # input shape required
  tf.keras.layers.Dense(D)
])

def ode_scipy_solver(f, z0, ts, args):
    r = ode_scipy(f, z0, ts, args)
    r = tf.cast(r, z0.dtype)
    return r

def predict_fn(y, ts, args):
    y2 = tf.reshape(y, [1, D])
    y2 = tf.cast(y2, tf.float32)
    r = model(y2)
    r = tf.squeeze(r)
    r = tf.cast(r, y.dtype)
    return r

ground_truth = ode_scipy(lotka_volterra, z0, ts, args)

```

#### For the lazy, video of training

Run next cell to see video of a 500 epochs training I did on CPU/GPU (starting at epoch 80 because the rest was about the same).
It's quite interesting to see how it learns the curve starting from a kind of average curve and then trying to fit curves progressively...

Conclusion: Neural ODE is able to learn a tiny time-series (and certainly overfit a lot on it ;))

```python
from IPython.display import HTML

HTML("""

""")
```

#### Running the experiment by yourself

> On a 1080TI, it takes around 3h for 1000 epochs

```python

optimize(predict_fn, ode_scipy_solver, loss_L2, ground_truth, z0, ts,
         args=None, variables=model.variables, epochs=500,
         x_axis_min=-0.25, x_axis_max=3.75,
         y_axis_min=-0.25, y_axis_max=8.0,
         x2_axis_min=-0.25, x2_axis_max=8.0,
         y2_axis_min=-0.25, y2_axis_max=8.0,
         output_dir="ode_dense_lotka_volterra")

```

#### Result of training

```python
zs = ode_scipy(predict_fn, z0, ts, args)

xs = zs[:, 0].numpy()
ys = zs[:, 1].numpy()

iplot([go.Scatter(x=ts.numpy(), y=xs, name="x"), go.Scatter(x=ts.numpy(), y=ys, name="y")])

iplot([go.Scatter(x=xs, y=ys)])
```

### Learning a piece of Damped 2D oscillator curve but not converging

#### Model

```python
D = 2
N = 50
ts = tf.constant(np.linspace(0., 6.0, N))
z0 = tf.Variable(tf.constant([2.0, 0.0], shape=(2,), dtype=ts.dtype))

args = tf.Variable(tf.constant([
    # −0.1x^3 + 2.0y^3
    [-0.1, 2.0],
    # −2.0x^3 − 0.1y^3
    [-2.0, -0.1]], shape=(2, 2), dtype=ts.dtype))

model2 = tf.keras.Sequential([
  tf.keras.layers.Dense(25, activation=tf.nn.tanh, input_shape=(1, D)),  # input shape required
  tf.keras.layers.Dense(D)
])

def ode_scipy_solver2(f, z0, ts, args):
    r = ode_scipy(f, z0, ts, args)
    r = tf.cast(r, z0.dtype)
    return r

def predict_fn2(y, ts, args):
    y2 = tf.reshape(y, [1, D])
    y2 = tf.cast(y2, tf.float32)
    r = model2(y2)
    r = tf.squeeze(r)
    r = tf.cast(r, y.dtype)
    return r

ground_truth2 = ode_scipy(damped_oscillator_2d, z0, ts, args)

```

#### For the lazy, video of training

Run next cell to see video of a 1000 epochs training I did on CPU/GPU.
It's quite interesting to see how it learns first curves but fails to converge on the whole oscillator.
I need to test with a bigger Dense network and maybe another time distribution...
Anyway, it's fun to see how it learns, isnt' it? ;)

```python
from IPython.display import HTML

HTML("""

""")

```

#### Running the experiment by yourself

> On a 1080TI, it takes around 3h for 1000 epochs

```python

optimize(predict_fn2, ode_scipy_solver2, loss_L2, ground_truth2, z0, ts,
         args=None, variables=model2.variables, epochs=5,
         x_axis_min=-0.25, x_axis_max=6.25,
         y_axis_min=-2.5, y_axis_max=2.5,
         x2_axis_min=-2.5, x2_axis_max=2.5,
         y2_axis_min=-2.5, y2_axis_max=2.5,
         output_dir="ode_dense_damped_oscillator_2d", start=0)

```

## Post Work'em

This is not a post mortem as the study is not finished yet and still on-going.

After reaching this step, I'm truly happy as it works and sounds promising about the paper concepts which are reproducible... I don't pretend it's perfect yet and I can say that trainings are not very fast and tend to differ between runs. Sometimes, it seems to converge at first and then loop without converging further. Yet, it's already quite cool to me :)

But it has been painful... Ok, I'm not yet an expert in TF or Python and I'm still learning both deeper concepts and libraries. But I can say I've been developing for tens of years in complex backend production environment and distributed systems with Native, Java and Scala/FP. From this experience, current state of maturity of TF and AI Python ecosystem, in general, is far from those in terms of robustness, consistency and determinism. Ok, you can have fun with a few pre-built samples but as soon as you go out of the paved road, it becomes hairy and not very deterministic or consistent... Everything takes a lot of time and effort to reach acceptable state... Specially when dealing with custom gradients which tend to crash with no error at all so it's hard to know which part of your graph fails..

I'm (in) eager(-mode) to discover incoming TF 2.x. which should make TF a bit more practical (more like Pytorch) and consistent in terms API... I'm also looking at TF-Swift, not that I like Swift but because a multi-stage compiled auto-diff remains quite appealing to me compared to Pythonic weakness...

## Next steps...

_This notebook is the first one of a few incoming studies I'll publish when I find a few minutes as I'm also implementing Time-series models and Continuous Normalizing Flows as explained in the paper..._

I want to test more networks for derivative or time distribution and see how it can speed trainings...

I want to make a quick Pytorch implementation and compare to TF 1.x.

I want to see how they check full-GPU ODE Solvers in Julia or Pytorch...

I want to play with depth distribution as a learnable parameter...


```python
import os
import numpy as np
import numba
import numba.cffi_support
import importlib
import sunode
from sunode import lambdify
import sympy as sy
import matplotlib.pyplot as plt
```

```python
ndim = 2
nparam = 3

#p = sy.MatrixSymbol('p', nparam, 1)
#y = sy.MatrixSymbol('y', ndim, 1)

a, b, c = sy.symarray('p', (3,))
V, R = sy.symarray('y', (2,))

dV_dt = (V - V**3 / 3 + R) * c
dR_dt = (V - a + b * R) / -c

dy = sy.Matrix([dV_dt, dR_dt])
```

```python
import sympy.codegen.rewriting
```

```python
def optimize(expr):
    return sympy.codegen.rewriting.optimize(
        expr, sympy.codegen.rewriting.optims_c99)
```

```python
dy = dy.applyfunc(optimize)
dy
```

$\displaystyle \left[\begin{matrix}p_{2} \left(- \frac{y_{0}^{3}}{3} + y_{0} + y_{1}\right)\\- \frac{- p_{0} + p_{1} y_{1} + y_{0}}{p_{2}}\end{matrix}\right]$

```python
# Create native function using numba
dy_pre, dy_comp = lambdify.lambdify_consts('_rhs', [], [p, y], dy)
```

```python
v = sy.MatrixSymbol('v', ndim, 1)
Jv = dy.jacobian(y) * sy.Matrix(v)
Jv = Jv.applyfunc(optimize)
Jv
```

$\displaystyle \left[\begin{matrix}\left(1 - y_{0, 0}^{2}\right) p_{2, 0} v_{0, 0} + p_{2, 0} v_{1, 0}\\- \frac{p_{1, 0} v_{1, 0}}{p_{2, 0}} - \frac{v_{0, 0}}{p_{2, 0}}\end{matrix}\right]$

```python
# We also have dy in the matrix-vector product. Can we use it somehow?
Jv_pre, Jv_comp = lambdify.lambdify_consts('_jacobian_prod', [], [p, y, v], Jv)
```

```python
J = dy.jacobian(y)
J = J.applyfunc(optimize)
J
```

$\displaystyle \left[\begin{matrix}\left(1 - y_{0, 0}^{2}\right) p_{2, 0} & p_{2, 0}\\- \frac{1}{p_{2, 0}} & - \frac{p_{1, 0}}{p_{2, 0}}\end{matrix}\right]$

```python
J_pre, J_comp = lambdify.lambdify_consts('_jacobian', [], [p, y], J)
```

```python
df_dp = dy.jacobian(p)
df_dp = df_dp.applyfunc(optimize)
df_dp
```

$\displaystyle \left[\begin{matrix}0 & 0 & - \frac{y_{0, 0}^{3}}{3} + y_{0, 0} + y_{1, 0}\\\frac{1}{p_{2, 0}} & - \frac{y_{1, 0}}{p_{2, 0}} & \frac{- p_{0, 0} + p_{1, 0} y_{1, 0} + y_{0, 0}}{p_{2, 0}^{2}}\end{matrix}\right]$

```python
df_dp_pre, df_dp_comp = lambdify.lambdify_consts('_df_dp', [], [p, y], df_dp)
```

```python
user_dtype = np.dtype([('p', (np.float64, (nparam,)))], align=True)
user_ndtype = numba.from_dtype(user_dtype)
user_ndtype

user_ndtype_p = numba.types.CPointer(user_ndtype)
user_data = np.zeros((1,), dtype=user_dtype)[0]
user_data['p'] = 1
user_data
```

    ([1., 1., 1.],)

```python
@numba.njit()
def rhs(t, y, out, user_data):
    p = user_data.p.reshape(-1, 1)
    # TODO The optimizer doesn't seem to be able
    # to optimize the malloc away. We could
    # extend lambdify to take an out argument.
    out[:] = dy_comp((), p, y.reshape(-1, 1))[:, 0]
    return 0

@numba.njit
def jacobian_prod(t, v, y, fy, out, user_data):
    p = user_data.p
    out[:] = Jv_comp(
        (),
        p.reshape((-1, 1)),
        y.reshape((-1, 1)),
        v.reshape((-1, 1))
    )[:, 0]
    return 0

@numba.njit
def jacobian(t, y, fy, out, user_data):
    p = user_data.p
    out[:, :] = J_comp((), p.reshape((-1, 1)), y.reshape((-1, 1)))
    return 0
```

```python
ffi = sunode._cvodes.ffi
lib = sunode._cvodes.lib
```

```python
numba.cffi_support.register_module(sunode._cvodes)
numba.cffi_support.register_type(
    ffi.typeof('N_Vector').item,
    numba.types.Opaque('N_Vector'))
numba.cffi_support.register_type(
    ffi.typeof('SUNMatrix').item,
    numba.types.Opaque('SUNMatrix'))
```

```python
ndim = 2
y0 = sunode.empty_vector(ndim)
y0.data[:] = 0
```

```python
abstol = sunode.empty_vector(ndim)
abstol.data[:] = [1e-8, 1e-8]
reltol = 1e-8
```

```python
#ode = lib.CVodeCreate(lib.CV_BDF)
ode = lib.CVodeCreate(lib.CV_ADAMS)
```

```python
user_data_p = ffi.cast('void *', ffi.addressof(ffi.from_buffer(user_data.data)))
```

```python
N_VGetArrayPointer_Serial = lib.N_VGetArrayPointer_Serial

func_type = numba.cffi_support.map_type(ffi.typeof('CVRhsFn'))
func_type = func_type.return_type(*(func_type.args[:-1] + (user_ndtype_p,)))
func_type

@numba.cfunc(func_type)
def rhs_wrapper(t, y_, out_, user_data_):
    y_ptr = N_VGetArrayPointer_Serial(y_)
    out_ptr = N_VGetArrayPointer_Serial(out_)
    y = numba.carray(y_ptr, (ndim,))
    out = numba.carray(out_ptr, (ndim,))
    user_data = numba.carray(user_data_, (1,))[0]
    p = user_data.p

    return rhs(t, y, out, user_data)
```

```python
SUNDenseMatrix_Data = lib.SUNDenseMatrix_Data

func_type = numba.cffi_support.map_type(ffi.typeof('CVLsJacFn'))
args = list(func_type.args)
args[4] = user_ndtype_p
func_type = func_type.return_type(*args)
func_type

@numba.cfunc(func_type)
def Jac(t, y_, fy_, out_, user_data_, tmp1_, tmp2_, tmp3_):
    y_ptr = N_VGetArrayPointer_Serial(y_)
    out_ptr = SUNDenseMatrix_Data(out_)
    fy_ptr = N_VGetArrayPointer_Serial(fy_)
    y = numba.carray(y_ptr, (ndim,))
    out = numba.farray(out_ptr, (ndim, ndim))
    fy = numba.carray(fy_ptr, (ndim,))
    user_data = numba.carray(user_data_, (1,))[0]
    p = user_data.p

    return jacobian(t, y, fy, out, user_data)
```

```python
# Not used
SUNDenseMatrix_Data = lib.SUNDenseMatrix_Data

func_type = numba.cffi_support.map_type(ffi.typeof('CVLsJacTimesVecFn'))
args = list(func_type.args)
args[-2] = user_ndtype_p
func_type = func_type.return_type(*args)
func_type

@numba.cfunc(func_type)
def jac_prod_wrapper(v_, out_, t, y_, fy_, user_data_, tmp_):
    v_ptr = N_VGetArrayPointer_Serial(v_)
    y_ptr = N_VGetArrayPointer_Serial(y_)
    #out_ptr = SUNDenseMatrix_Data(out_)
    out_ptr = N_VGetArrayPointer_Serial(out_)
    fy_ptr = N_VGetArrayPointer_Serial(fy_)
    v = numba.carray(v_ptr, (ndim,))
    y = numba.carray(y_ptr, (ndim,))
    out = numba.carray(out_ptr, (ndim,))
    fy = numba.carray(fy_ptr, (ndim,))
    user_data = numba.carray(user_data_, (1,))[0]
    p = user_data.p

    return jacobian_prod(t, v, y, fy, out, user_data)
```

```python
print(lib.CVodeInit(ode, rhs_wrapper.cffi, 0, y0.c_ptr))
lib.CVodeSVtolerances(ode, reltol, abstol.c_ptr)
```

    0

    0

```python
A = lib.SUNDenseMatrix(ndim, ndim)
```

```python
linsolver = lib.SUNLinSol_Dense(y0.c_ptr, A)
```

```python
lib.CVodeSetLinearSolver(ode, linsolver, A)
```

    0

```python
lib.CVodeSetJacFn(ode, Jac.cffi)
```

    0

```python
lib.CVodeSetUserData(ode, user_data_p)
```

    0

```python
yS = lib.N_VCloneVectorArray(nparam, y0.c_ptr)
```

```python
vecs = [sunode.basic.Vector(yS[i]) for i in range(nparam)]
for vec in vecs:
    vec.data[:] = 0
np.array([vec.data for vec in vecs])
```

    []
    []

    array([[0., 0.],
           [0., 0.],
           [0., 0.]])

```python
func_type = numba.cffi_support.map_type(ffi.typeof('CVSensRhsFn'))
args = list(func_type.args)
args[-3] = user_ndtype_p
func_type = func_type.return_type(*args)
func_type

@numba.cfunc(func_type)
def sens_rhs(n_params, t, y_, ydot_, yS_, out_, user_data_, tmp1_, tmp2_):
    y_ptr = N_VGetArrayPointer_Serial(y_)
    y = numba.carray(y_ptr, (ndim,))
    y_dot_ptr = N_VGetArrayPointer_Serial(ydot_)
    ydot = numba.carray(y_dot_ptr, (ndim,))

    out = []
    for i in range(n_params):
        out_i_ptr = N_VGetArrayPointer_Serial(out_[i])
        out_i = numba.carray(out_i_ptr, (ndim,))
        out.append(out_i)

    yS = []
    for i in range(n_params):
        yS_i_ptr = N_VGetArrayPointer_Serial(yS_[i])
        yS_i = numba.carray(yS_i_ptr, (ndim,))
        yS.append(yS_i)

    user_data = numba.carray(user_data_, (1,))[0]
    p = user_data.p

    df_dp = df_dp_comp((), p.reshape((-1, 1)), y.reshape((-1, 1)))

    for i in range(n_params):
        s = yS[i]
        ySdot = out[i]

        jacobian_prod(t, s, y, None, ySdot, user_data)
        ySdot[:] += df_dp[:, i]

    return 0
```

```python
assert 0 == lib.CVodeSensInit(ode, nparam, lib.CV_STAGGERED, sens_rhs.cffi, yS)
assert 0 == lib.CVodeSensEEtolerances(ode)
assert 0 == lib.CVodeSetSensErrCon(ode, 1)
```

```python
pbar = np.ones(nparam)
```

```python
NULL_D = ffi.cast('double *', 0)
NULL_I = ffi.cast('int *', 0)
pbar_p = ffi.cast('double *', ffi.addressof(ffi.from_buffer(pbar.data)))
```

```python
lib.CVodeSetSensParams(ode, NULL_D, pbar_p, NULL_I)
```

    0

```python
time_p = ffi.new('double*')
time_p[0] = 0.
```

```python
tvals = np.linspace(0, 10)
y_vals = []
sens_vals = []

user_data['p'][...] = [-0.2, 0.3, 0.5]
```

```python
user_data
```

    ([-0.2,  0.3,  0.5],)

```python
%%timeit
y_vals.clear()
sens_vals.clear()

y0.data[:] = [0.5, 0.6]
for vec in vecs:
    vec.data[:] = 0

lib.CVodeReInit(ode, 0, y0.c_ptr)
for t in tvals:
    retval = -1
    while retval == -1:
        retval = lib.CVode(ode, t, y0.c_ptr, time_p, lib.CV_NORMAL)
    y_vals.append(y0.data.copy())
    lib.CVodeGetSens(ode, time_p, yS)
    sens_vals.append(np.array([vec.data.copy() for vec in vecs]).T)
```

    1.91 ms ± 6.06 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

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

```python
plt.plot(tvals, np.array(y_vals)[:, 0])
plt.plot(tvals, np.array(y_vals)[:, 1])
```

```python
plt.plot(tvals, np.array(sens_vals)[:, 0, 0])
plt.plot(tvals, np.array(sens_vals)[:, 0, 1])
plt.plot(tvals, np.array(sens_vals)[:, 0, 2])
```

```python
plt.plot(tvals, np.array(sens_vals)[:, 1, 0])
plt.plot(tvals, np.array(sens_vals)[:, 1, 1])
plt.plot(tvals, np.array(sens_vals)[:, 1, 2])
```

```python
num_evals_p = ffi.new('long*')
num_evals_p[0] = 0
```

```python
lib.CVodeGetNumRhsEvals(ode, num_evals_p)
num_evals_p[0]
```

    331

```python
lib.CVodeGetNumSteps(ode, num_evals_p)
num_evals_p[0]
```

    158

```python
time_p[0]
```

    10.0

```python
from scipy import integrate
```

```python
def rhs_(t, y):
    out = np.empty_like(y)
    assert rhs(t, y, out, user_data) == 0
    return out
```

```python
def jac_(t, y):
    out = np.empty((ndim, ndim))
    assert jacobian(t, y, None, out, user_data) == 0
    return out
```

```python
from scipy import integrate
```

```python
%%timeit
out = integrate.solve_ivp(
    rhs_,
    jac=jac_,
    t_span=(0, tvals.max()),
    t_eval=tvals, y0=np.array([0.5, 0.6]), method='BDF',
    rtol=reltol, atol=abstol.data.copy())
```

    46.6 ms ± 185 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

```python
out = integrate.solve_ivp(
    rhs_,
    jac=jac_,
    t_span=(0, tvals.max()),
    t_eval=tvals, y0=np.array([0.5, 0.6]), method='BDF',
    rtol=reltol, atol=abstol.data.copy())
```

```python
plt.plot(tvals, out.y[0, :])
plt.plot(tvals, out.y[1, :])

plt.plot(tvals, np.array(y_vals)[:, 0], '.-.')
plt.plot(tvals, np.array(y_vals)[:, 1], '.-.')
```

```python
out
```

      message: 'The solver successfully reached the end of the integration interval.'
         nfev: 414
         njev: 1
          nlu: 36
          sol: None
       status: 0
      success: True
            t: array([ 0.        ,  0.20408163,  0.40816327,  0.6122449 ,  0.81632653,
            1.02040816,  1.2244898 ,  1.42857143,  1.63265306,  1.83673469,
            2.04081633,  2.24489796,  2.44897959,  2.65306122,  2.85714286,
            3.06122449,  3.26530612,  3.46938776,  3.67346939,  3.87755102,
            4.08163265,  4.28571429,  4.48979592,  4.69387755,  4.89795918,
            5.10204082,  5.30612245,  5.51020408,  5.71428571,  5.91836735,
            6.12244898,  6.32653061,  6.53061224,  6.73469388,  6.93877551,
            7.14285714,  7.34693878,  7.55102041,  7.75510204,  7.95918367,
            8.16326531,  8.36734694,  8.57142857,  8.7755102 ,  8.97959184,
            9.18367347,  9.3877551 ,  9.59183673,  9.79591837, 10.        ])
     t_events: None
            y: array([[ 0.5       ,  0.59330775,  0.6558635 ,  0.68691887,  0.68730437,
             0.65901042,  0.60470768,  0.52735582,  0.42997359,  0.31557878,
             0.18726967,  0.04840153, -0.09720291, -0.24508888, -0.39014637,
            -0.52681342, -0.64953166, -0.75337261, -0.83465121, -0.89132052,
            -0.92303038, -0.93087671, -0.91697523, -0.88401334, -0.83488646,
            -0.77245968, -0.69944732, -0.61838069, -0.53163038, -0.44145466,
            -0.35005157, -0.25959857, -0.172268  , -0.09021141, -0.01551158,
             0.0498925 ,  0.10429556,  0.14631476,  0.17496078,  0.18967952,
             0.19036044,  0.17731603,  0.15124244,  0.11317357,  0.06443882,
             0.0066299 , -0.05842347, -0.12867506, -0.20189037, -0.27569338],
           [ 0.6       ,  0.24269343, -0.10330376, -0.42727928, -0.71978789,
            -0.97303477, -1.18103599, -1.33959005, -1.44613273, -1.49955428,
            -1.5000435 , -1.44900144, -1.34904174, -1.20406567, -1.0193672 ,
            -0.80169251, -0.55916445, -0.30100268, -0.03703116,  0.22295627,
             0.4698444 ,  0.69572922,  0.89425338,  1.06073041,  1.19209685,
             1.28675424,  1.3443617 ,  1.36562446,  1.35210615,  1.30607714,
             1.23040031,  1.12844806,  1.0040396 ,  0.86138545,  0.70502601,
             0.53975302,  0.37050806,  0.20225772,  0.03985218, -0.11212185,
            -0.24948055, -0.36856378, -0.4663299 , -0.54042293, -0.5892147 ,
            -0.61182775, -0.60814451, -0.57880707, -0.52520765, -0.44946602]])

```python

```

```python

```


# Random Signals and LTI-Systems

*This jupyter notebook is part of a [collection of notebooks](../index.ipynb) on various topics of Digital Signal Processing. Please direct questions and suggestions to [Sascha.Spors@uni-rostock.de](mailto:Sascha.Spors@uni-rostock.de).*

## Power Spectral Densitity

For a wide-sense stationary (WSS) real-valued random process $x[k]$, the [power spectral density](../random_signals/power_spectral_densities.ipynb#Power-Spectral-Density) (PSD) $\Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})$ is given as the discrete-time Fourier transformation (DTFT) of the auto-correlation function (ACF) $\varphi_{xx}[\kappa]$

\begin{equation}
\Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) = \sum_{\kappa = -\infty}^{\infty} \varphi_{xx}[\kappa] \; \mathrm{e}^{\,-\mathrm{j}\,\Omega\,\kappa}
\end{equation}

Under the assumption of a real-valued LTI system with impulse response $h[k] \in \mathbb{R}$, the PSD $\Phi_{yy}(\mathrm{e}^{\,\mathrm{j}\,\Omega})$ of the output signal of an LTI system $y[k] = \mathcal{H} \{ x[k] \}$ is derived by taking the DTFT of the [ACF of the output signal](../random_signals_LTI_systems/correlation_functions.ipynb#Auto-Correlation-Function) $\varphi_{yy}[\kappa]$

\begin{align}
\Phi_{yy}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) &= \sum_{\kappa = -\infty}^{\infty} \underbrace{h[\kappa] * h[-\kappa]}_{\varphi_{hh}[\kappa]} * \varphi_{xx}[\kappa] \; \mathrm{e}^{\,-\mathrm{j}\,\Omega\,\kappa} \\
&= H(\mathrm{e}^{\,\mathrm{j}\,\Omega}) \cdot H(\mathrm{e}^{\,-\mathrm{j}\,\Omega}) \cdot \Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) = | H(\mathrm{e}^{\,\mathrm{j}\,\Omega}) |^2 \cdot \Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})
\end{align}

The PSD of the output signal $\Phi_{yy}(\mathrm{e}^{\,\mathrm{j}\,\Omega})$ of an LTI system is given by the  PSD of the input signal $\Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})$ multiplied with the squared magnitude $| H(\mathrm{e}^{\,\mathrm{j}\,\Omega}) |^2$ of the transfer function of the system.

### Example - Pink Noise

It can be concluded from above findings, that filtering can be applied to a white noise random signal $x[k]$ with $\Phi_{yy}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) = N_0$ in order to create a random signal $y[k] = x[k] * h[k]$ with a desired PSD

\begin{equation}
\Phi_{yy}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) = N_0 \cdot | H(\mathrm{e}^{\,\mathrm{j}\,\Omega}) |^2
\end{equation}

where $N_0$ denotes the power per frequency of the white noise. Such a random signal is commonly termed as [*colored noise*](https://en.wikipedia.org/wiki/Colors_of_noise). Different application specific types of colored noise exist. One of these is [*pink noise*](https://en.wikipedia.org/wiki/Pink_noise) whose PSD is inversely proportional to the frequency. The approximation of a pink noise signal by filtering is illustrated by the following example. The PSDs $\Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})$ and $\Phi_{yy}(\mathrm{e}^{\,\mathrm{j}\,\Omega})$ are estimated from $x[k]$ and $y[k]$ using the [Welch technique](../spectral_estimation_random_signals/welch_method.ipynb).

```python
import numpy as np
import matplotlib.pyplot as plt
import scipy.signal as sig

fs = 44100
N = 5*fs

# generate uniformly distributed white noise
np.random.seed(1)
x = np.random.uniform(size=N) - .5
# filter white noise to yield pink noise
# see http://www.firstpr.com.au/dsp/pink-noise/#Filtering
a = np.poly([0.99572754, 0.94790649, 0.53567505])  # denominator coefficients
b = np.poly([0.98443604, 0.83392334, 0.07568359])  # numerator coefficients
y = 1/3 * sig.lfilter(b, a, x)
# estimate PSDs using Welch's technique
f, Pxx = sig.csd(x, x, nperseg=256)
f, Pyy = sig.csd(y, y, nperseg=256)

# PSDs
Om = f * 2 * np.pi
plt.plot(Om, 20*np.log10(np.abs(.5*Pxx)),
         label=r'$| \Phi_{xx}(e^{j \Omega}) |$ in dB')
plt.plot(Om, 20*np.log10(np.abs(.5*Pyy)),
         label=r'$| \Phi_{yy}(e^{j \Omega}) |$ in dB')
plt.title('Power Spectral Density')
plt.xlabel(r'$\Omega$')
plt.legend()
plt.axis([0, np.pi, -60, -10])
plt.grid()
```

Let's listen to white and pink noise

```python
from scipy.io import wavfile

wavfile.write('uniform_white_noise.wav', fs, np.int16(x*32768))
wavfile.write('uniform_pink_noise.wav', fs, np.int16(y*32768))
```

**White noise**

<audio src="./uniform_white_noise.wav" controls>Your browser does not support the audio element.</audio>[./uniform_white_noise.wav](./uniform_white_noise.wav)

**Pink noise**

<audio src="./uniform_pink_noise.wav" controls>Your browser does not support the audio element.</audio>[./uniform_pink_noise.wav](./uniform_white_noise.wav)

## Cross-Power Spectral Densities

The cross-power spectral densities $\Phi_{yx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})$ and $\Phi_{xy}(\mathrm{e}^{\,\mathrm{j}\,\Omega})$ between the in- and output of an LTI system are given by taking the DTFT of the [cross-correlation functions](../random_signals_LTI_systems/correlation_functions.ipynb#Cross-Correlation-Function) (CCF) $\varphi_{yx}[\kappa]$ and $\varphi_{xy}[\kappa]$. Hence,

\begin{equation}
\Phi_{yx}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) = \sum_{\kappa = -\infty}^{\infty} h[\kappa] * \varphi_{xx}[\kappa] \; \mathrm{e}^{\,-\mathrm{j}\,\Omega\,\kappa} = \Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) \cdot H(\mathrm{e}^{\,\mathrm{j}\,\Omega})
\end{equation}

and

\begin{equation}
\Phi_{xy}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) = \sum_{\kappa = -\infty}^{\infty} h[-\kappa] * \varphi_{xx}[\kappa] \; \mathrm{e}^{\,-\mathrm{j}\,\Omega\,\kappa} = \Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) \cdot H(\mathrm{e}^{\,-\mathrm{j}\,\Omega})
\end{equation}

## System Identification by Spectral Division

Using the result above for the cross-power spectral density $\Phi_{yx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})$ between out- and input, and the relation of the [CCF of finite-length signals to the convolution](../random_signals/correlation_functions.ipynb#Definition) yields

\begin{equation}
H(\mathrm{e}^{\,\mathrm{j}\,\Omega}) = \frac{\Phi_{yx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})}{\Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})} = \frac{\frac{1}{K} Y(\mathrm{e}^{\,\mathrm{j}\,\Omega}) \cdot X(\mathrm{e}^{\,-\mathrm{j}\,\Omega})}{\frac{1}{K} X(\mathrm{e}^{\,\mathrm{j}\,\Omega}) \cdot X(\mathrm{e}^{\,-\mathrm{j}\,\Omega})}
= \frac{Y(\mathrm{e}^{\,\mathrm{j}\,\Omega})}{X(\mathrm{e}^{\,\mathrm{j}\,\Omega})}
\end{equation}

holding for $\Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) \neq 0$ and $X(\mathrm{e}^{\,\mathrm{j}\,\Omega}) \neq 0$. Hence, the transfer function $H(\mathrm{e}^{\,\mathrm{j}\,\Omega})$ of an unknown system can be derived by dividing the spectrum of the output signal $Y(\mathrm{e}^{\,\mathrm{j}\,\Omega})$ through the spectrum of the input signal $X(\mathrm{e}^{\,\mathrm{j}\,\Omega})$. This is equal to the [definition of the transfer function](https://en.wikipedia.org/wiki/Transfer_function). However, care has to be taken that the spectrum of the input signal does not contain zeros.

Above relation can be realized by the discrete Fourier transformation (DFT) by taking into account that a multiplication of two spectra $X[\mu] \cdot Y[\mu]$ results in the cyclic/periodic convolution $x[k] \circledast y[k]$. Since we aim at a linear convolution, zero-padding of the in- and output signal has to be applied.

### Example

We consider the estimation of the impulse response $h[k] = \mathcal{F}_*^{-1} \{ H(\mathrm{e}^{\,\mathrm{j}\,\Omega}) \}$ of an unknown system using the spectral division method. Normal distributed white noise with variance $\sigma_n^2 = 1$ is used as wide-sense ergodic input signal $x[k]$. In order to show the effect of sensor noise, normally distributed white noise $n[k]$ with the variance $\sigma_n^2 = 0.01$ is added to the output signal $y[k] = x[k] * h[k] + n[k]$.

```python
N = 1000  # number of samples for input signal

# generate input signal
# normally distributed (zero-mean, unit-variance) white noise
np.random.seed(1)
x = np.random.normal(size=N, scale=1)
# impulse response of the system
h = np.concatenate((np.zeros(20), np.ones(10), np.zeros(20)))
# output signal by convolution
y = np.convolve(h, x, mode='full')
# add noise to the output signal
y = y + np.random.normal(size=y.shape, scale=.1)

# zero-padding of input signal
x = np.concatenate((x, np.zeros(len(h)-1)))
# estimate transfer function
H = np.fft.rfft(y)/np.fft.rfft(x)
# compute inpulse response
he = np.fft.irfft(H)
he = he[0:len(h)]

# plot impulse response
plt.figure()
plt.stem(he, label='estimated')
plt.plot(h, 'g-', label='true')
plt.title('Estimated impulse response')
plt.xlabel(r'$k$')
plt.ylabel(r'$\hat{h}[k]$')
plt.legend();
```

**Exercise**

* Change the length `N` of the input signal. What happens?
* Change the variance $\sigma_n^2$ of the additive noise. What happens?

Solution: Increasing the length `N` of the input signal lowers the uncertainty in estimating the impulse response. The higher the variance of the additive white noise, the higher the uncertainties in the estimated impulse response.

**Copyright**

This notebook is provided as [Open Educational Resource](https://en.wikipedia.org/wiki/Open_educational_resources). Feel free to use the notebook for your own purposes. The text is licensed under [Creative Commons Attribution 4.0](https://creativecommons.org/licenses/by/4.0/), the code of the IPython examples under the [MIT license](https://opensource.org/licenses/MIT). Please attribute the work as follows: *Sascha Spors, Digital Signal Processing - Lecture notes featuring computational examples*.


## Introduction
-----

In this assignment you will recursively estimate the position of a vehicle along a trajectory using available measurements and a motion model.

The vehicle is equipped with a very simple type of LIDAR sensor, which returns range and bearing measurements corresponding to individual landmarks in the environment. The global positions of the landmarks are assumed to be known beforehand. We will also assume known data association, that is, which measurment belong to which landmark.

## Motion and Measurement Models
-----

### Motion Model

The vehicle motion model recieves linear and angular velocity odometry readings as inputs, and outputs the state (i.e., the 2D pose) of the vehicle:

\begin{align}
\mathbf{x}_{k} &= \mathbf{x}_{k-1} + T
\begin{bmatrix}
\cos\theta_{k-1} &0 \\
\sin\theta_{k-1} &0 \\
0 &1
\end{bmatrix}
\left(
\begin{bmatrix}
v_k \\
\omega_k
\end{bmatrix}
+ \mathbf{w}_k
\right)
\, , \, \, \, \, \, \mathbf{w}_k = \mathcal{N}\left(\mathbf{0}, \mathbf{Q}\right)
\end{align}

- $\mathbf{x}_k = \left[ x \, y \, \theta \right]^T$ is the current 2D pose of the vehicle
- $v_k$ and $\omega_k$ are the linear and angular velocity odometry readings, which we use as inputs to the model

The process noise $\mathbf{w}_k$ has a (zero mean) normal distribution with a constant covariance $\mathbf{Q}$.

### Measurement Model

The measurement model relates the current pose of the vehicle to the LIDAR range and bearing measurements $\mathbf{y}^l_k = \left[r \, \phi \right]^T$.

\begin{align}
\mathbf{y}^l_k =
\begin{bmatrix}
\sqrt{(x_l - x_k - d\cos\theta_{k})^2 + (y_l - y_k - d\sin\theta_{k})^2} \\
atan2\left(y_l - y_k - d\sin\theta_{k},x_l - x_k - d\cos\theta_{k}\right) - \theta_k
\end{bmatrix}
+
\mathbf{n}^l_k
\, , \, \, \, \, \, \mathbf{n}^l_k = \mathcal{N}\left(\mathbf{0}, \mathbf{R}\right)
\end{align}

- $x_l$ and $y_l$ are the ground truth coordinates of the landmark $l$
- $x_k$ and $y_k$ and $\theta_{k}$ represent the current pose of the vehicle
- $d$ is the known distance between robot center and laser rangefinder (LIDAR)

The landmark measurement noise $\mathbf{n}^l_k$ has a (zero mean) normal distribution with a constant covariance $\mathbf{R}$.

## Getting Started
-----

Since the models above are nonlinear, we recommend using the extended Kalman filter (EKF) as the state estimator.
Specifically, you will need to provide code implementing the following steps:
- the prediction step, which uses odometry measurements and the motion model to produce a state and covariance estimate at a given timestep, and
- the correction step, which uses the range and bearing measurements provided by the LIDAR to correct the pose and pose covariance estimates

### Unpack the Data
First, let's unpack the available data:

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

with open('data/data.pickle', 'rb') as f:
    data = pickle.load(f)

t = data['t']  # timestamps [s]

x_init  = data['x_init'] # initial x position [m]
y_init  = data['y_init'] # initial y position [m]
th_init = data['th_init'] # initial theta position [rad]

# input signal
v  = data['v']  # translational velocity input [m/s]
om = data['om']  # rotational velocity input [rad/s]

# bearing and range measurements, LIDAR constants
b = data['b']  # bearing to each landmarks center in the frame attached to the laser [rad]
r = data['r']  # range measurements [m]
l = data['l']  # x,y positions of landmarks [m]
d = data['d']  # distance between robot center and laser rangefinder [m]
d = np.squeeze(d)
```

Note that distance from the LIDAR frame to the robot center is provided and loaded as an array into the `d` variable.

### Ground Truth
If available, it is useful to plot the ground truth position and orientation before starting the assignment.

Notice that the orientation values are wrapped to the $\left[-\pi,\pi\right]$ range in radians.

### Initializing Parameters

Now that our data is loaded, we can begin getting things set up for our solver. One of the
most important aspects of designing a filter is determining the input and measurement noise covariance matrices, as well as the initial state and covariance values. We set the values here:

```python
v_var = 0.01  # translation velocity variance
om_var = 0.01  # rotational velocity variance
r_var = 0.001  # range measurements variance
b_var = 0.001  # bearing measurement variance

Q_km = np.diag([v_var, om_var]) # input noise covariance
cov_y = np.diag([r_var, b_var])  # measurement noise covariance

x_est = np.zeros([len(v), 3])  # estimated states, x, y, and theta
P_est = np.zeros([len(v), 3, 3])  # state covariance matrices

x_est[0] = np.array([x_init, y_init, th_init]) # initial state
P_est[0] = np.diag([1, 1, 0.1]) # initial state covariance
P_check = P_est[0]
```

**Remember:** that it is neccessary to tune the measurement noise variances `r_var`, `b_var` in order for the filter to perform well!

In order for the orientation estimates to coincide with the bearing measurements, it is also neccessary to wrap all estimated $\theta$ values to the $(-\pi , \pi]$ range.

```python
# Wraps angle to (-pi,pi] range
def wraptopi(x):
    if x > np.pi:
        x = x - (np.floor(x / (2 * np.pi)) + 1) * 2 * np.pi
    elif x < -np.pi:
        x = x + (np.floor(x / (-2 * np.pi)) + 1) * 2 * np.pi
    return x
```

## Correction Step
-----
First, let's implement the measurement update function, which takes an available landmark measurement $l$ and updates the current state estimate $\mathbf{\check{x}}_k$.
For each landmark measurement received at a given timestep $k$, you should implement the following steps:

- Compute the measurement model Jacobians at $\mathbf{\check{x}}_{k}$
\begin{align}
\mathbf{y}^l_k = &\mathbf{h}(\mathbf{x}_{k}, \mathbf{n}^l_k) \\\\
\mathbf{H}_{k} = \frac{\partial \mathbf{h}}{\partial \mathbf{x}_{k}}\bigg|_{\mathbf{\check{x}}_{k},0}& \, , \, \, \, \,
\mathbf{M}_{k} = \frac{\partial \mathbf{h}}{\partial \mathbf{n}_{k}}\bigg|_{\mathbf{\check{x}}_{k},0} \, .
\end{align}
- Compute the Kalman Gain
\begin{align}
\mathbf{K}_k &= \mathbf{\check{P}}_k \mathbf{H}_k^T \left(\mathbf{H}_k \mathbf{\check{P}}_k \mathbf{H}_k^T + \mathbf{M}_k \mathbf{R}_k \mathbf{M}_k^T \right)^{-1}
\end{align}
- Correct the predicted state
\begin{align}
\mathbf{\check{y}}^l_k &= \mathbf{h}\left(\mathbf{\check{x}}_k, \mathbf{0}\right) \\
\mathbf{\hat{x}}_k &= \mathbf{\check{x}}_k + \mathbf{K}_k \left(\mathbf{y}^l_k - \mathbf{\check{y}}^l_k\right)
\end{align}
- Correct the covariance
\begin{align}
\mathbf{\hat{P}}_k &= \left(\mathbf{I} - \mathbf{K}_k \mathbf{H}_k \right)\mathbf{\check{P}}_k
\end{align}

### Measurement Model

The measurement model relates the current pose of the vehicle to the LIDAR range and bearing measurements $\mathbf{y}^l_k = \left[r \, \phi \right]^T$.

\begin{align}
\mathbf{y}^l_k =
\begin{bmatrix}
\sqrt{(x_l - x_k - d\cos\theta_{k})^2 + (y_l - y_k - d\sin\theta_{k})^2} \\
atan2\left(y_l - y_k - d\sin\theta_{k},x_l - x_k - d\cos\theta_{k}\right) - \theta_k
\end{bmatrix}
+
\mathbf{n}^l_k
\, , \, \, \, \, \, \mathbf{n}^l_k = \mathcal{N}\left(\mathbf{0}, \mathbf{R}\right)
\end{align}

- $x_l$ and $y_l$ are the ground truth coordinates of the landmark $l$
- $x_k$ and $y_k$ and $\theta_{k}$ represent the current pose of the vehicle
- $d$ is the known distance between robot center and laser rangefinder (LIDAR)

The landmark measurement noise $\mathbf{n}^l_k$ has a (zero mean) normal distribution with a constant covariance $\mathbf{R}$.

```python
def measurement_update(lk, rk, bk, P_check, x_check):

    x_k = x_check[0, 0]
    y_k = x_check[1, 0]
    theta_k = x_check[2, 0]
    x_l = lk[0]
    y_l = lk[1]

    x_dist = x_l - x_k - d * np.cos(theta_k)
    y_dist = y_l - y_k - d * np.sin(theta_k)
    q = np.linalg.norm(np.array([x_dist, y_dist]))
    dist = x_dist ** 2 + y_dist ** 2

    H = np.mat([[-x_dist/q, -y_dist/q, (x_dist*d*np.sin(theta_k) - y_dist*d*np.cos(theta_k))/q],
               [y_dist/dist, -x_dist/dist, (-x_dist*d*np.cos(theta_k) - y_dist*d*np.sin(theta_k))/dist - 1]])

    M = np.identity(2)

    K = P_check.dot(H.T).dot(np.linalg.inv(H.dot(P_check.dot(H.T)) + M.dot(cov_y.dot(M.T))))

    y_k = np.mat([[q], [wraptopi(np.arctan2(y_dist, x_dist) - theta_k)]])
    y_meas = np.mat([[rk], [wraptopi(bk)]])

    # 3. Correct predicted state (remember to wrap the angles to [-pi,pi])

    x_check = x_check.reshape(3, 1)
    x_check = x_check + K.dot(y_meas - y_k)
    x_check[2, 0] = wraptopi(x_check[2, 0])

    # 4. Correct covariance
    P_check = (np.identity(3) - K.dot(H)).dot(P_check)

    return x_check, P_check

```

## Prediction Step
-----
Now, implement the main filter loop, defining the prediction step of the EKF using the motion model provided:

\begin{align}
\mathbf{\check{x}}_k &= \mathbf{f}\left(\mathbf{\hat{x}}_{k-1}, \mathbf{u}_{k-1}, \mathbf{0} \right) \\
\mathbf{\check{P}}_k &= \mathbf{F}_{k-1}\mathbf{\hat{P}}_{k-1}\mathbf{F}_{k-1}^T + \mathbf{L}_{k-1}\mathbf{Q}_{k-1}\mathbf{L}_{k-1}^T \, .
\end{align}

Where

\begin{align}
\mathbf{F}_{k-1} = \frac{\partial \mathbf{f}}{\partial \mathbf{x}_{k-1}}\bigg|_{\mathbf{\hat{x}}_{k-1},\mathbf{u}_{k},0}  \, , \, \, \, \,
\mathbf{L}_{k-1} = \frac{\partial \mathbf{f}}{\partial \mathbf{w}_{k}}\bigg|_{\mathbf{\hat{x}}_{k-1},\mathbf{u}_{k},0} \, .
\end{align}

### Motion Model

The vehicle motion model recieves linear and angular velocity odometry readings as inputs, and outputs the state (i.e., the 2D pose) of the vehicle:

\begin{align}
\mathbf{x}_{k} &= \mathbf{x}_{k-1} + T
\begin{bmatrix}
\cos\theta_{k-1} &0 \\
\sin\theta_{k-1} &0 \\
0 &1
\end{bmatrix}
\left(
\begin{bmatrix}
v_k \\
\omega_k
\end{bmatrix}
+ \mathbf{w}_k
\right)
\, , \, \, \, \, \, \mathbf{w}_k = \mathcal{N}\left(\mathbf{0}, \mathbf{Q}\right)
\end{align}

- $\mathbf{x}_k = \left[ x \, y \, \theta \right]^T$ is the current 2D pose of the vehicle
- $v_k$ and $\omega_k$ are the linear and angular velocity odometry readings, which we use as inputs to the model

The process noise $\mathbf{w}_k$ has a (zero mean) normal distribution with a constant covariance $\mathbf{Q}$.

```python
#### 5. Main Filter Loop #######################################################################
x_check = x_est[0]
x_check = x_check.reshape((3, 1))

for k in range(1, len(t)):  # start at 1 because we've set the initial prediciton

    delta_t = t[k] - t[k - 1]  # time step (difference between timestamps)

    theta = x_check[2, 0]
    # 1. Update state with odometry readings (remember to wrap the angles to [-pi,pi])
    u = np.mat([[v[k - 1]], [om[k - 1]]])
    A = np.mat([[np.cos(theta), 0], [np.sin(theta), 0], [0, 1]])
    A = A * delta_t
    x_check = x_check + A.dot(u)
    x_check[2, 0] = wraptopi(x_check[2, 0])

    # 2. Motion model jacobian with respect to last state

    F_km = np.mat([[1, 0, -delta_t*np.sin(theta)*v[k - 1]], [0, 1, delta_t*np.cos(theta)*v[k - 1]], [0, 0, 1]])

    # 3. Motion model jacobian with respect to noise
    L_km = A

    # 4. Propagate uncertainty
    P_check = F_km.dot(P_check.dot(F_km.T)) + L_km.dot(Q_km.dot(L_km.T))

    # 5. Update state estimate using available landmark measurements
    for i in range(len(r[k])):
        x_check, P_check = measurement_update(l[i], r[k, i], b[k, i], P_check, x_check)

    # Set final state predictions for timestep
    x_est[k, 0] = x_check[0]
    x_est[k, 1] = x_check[1]
    x_est[k, 2] = x_check[2]
    P_est[k, :, :] = P_check
```

Let's plot the resulting state estimates:

```python
e_fig = plt.figure()
ax = e_fig.add_subplot(111)
ax.plot(x_est[:, 0], x_est[:, 1])
ax.set_xlabel('x [m]')
ax.set_ylabel('y [m]')
ax.set_title('Estimated trajectory')
plt.show()

e_fig = plt.figure()
ax = e_fig.add_subplot(111)
ax.plot(t[:], x_est[:, 2])
ax.set_xlabel('Time [s]')
ax.set_ylabel('theta [rad]')
ax.set_title('Estimated trajectory')
plt.show()
```

Are you satisfied wth your results? The resulting trajectory should closely resemble the ground truth, with minor "jumps" in the orientation estimate due to angle wrapping. If this is the case, run the code below to produce your solution file.

```python
with open('submission.pkl', 'wb') as f:
    pickle.dump(x_est, f, pickle.HIGHEST_PROTOCOL)
```


# Laplace approximation ( Quadratic approximation)

In this notebook we will approximate posterior of beta-bernouli model for coin toss problem using laplace approximation method

```python
try:
    from probml_utils import latexify, savefig
except:
    %pip install git+https://github.com/probml/probml-utils.git
    from probml_utils import latexify, savefig

import jax
import jax.numpy as jnp
from jax import lax

try:
    from tensorflow_probability.substrates import jax as tfp
except ModuleNotFoundError:
    %pip install -qqq tensorflow_probability
    from tensorflow_probability.substrates import jax as tfp

try:
    import optax
except ModuleNotFoundError:
    %pip install -qqq optax
    import optax

try:
    from rich import print
except ModuleNotFoundError:
    %pip install -qqq rich
    from rich import print

try:
    from tqdm import trange
except:
    %pip install -qqq tqdm
    from tqdm import trange

import seaborn as sns
import numpy as np
import matplotlib.pyplot as plt
import warnings

warnings.filterwarnings("ignore")
dist = tfp.distributions
```

```python
latexify(width_scale_factor=2, fig_height=2)  # to apply latexify, set LATEXIFY=1 in environment variable
```

```python
# Use same data as https://github.com/probml/probml-notebooks/blob/main/notebooks/beta_binom_approx_post_pymc.ipynb
key = jax.random.PRNGKey(128)
dataset = np.repeat([0, 1], (10, 1))
n_samples = len(dataset)
print(f"Dataset: {dataset}")
n_heads = dataset.sum()
n_tails = n_samples - n_heads
```

    WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

Dataset: [0 0 0 0 0 0 0 0 0 0 1]

```python
# prior distribution ~ Beta
def prior_dist():
    return dist.Beta(concentration1=1.0, concentration0=1.0)

# likelihood distribution ~ Bernoulli
def likelihood_dist(theta):
    return dist.Bernoulli(probs=theta)
```

```python
# closed form of beta posterior
a = prior_dist().concentration1
b = prior_dist().concentration0

exact_posterior = dist.Beta(concentration1=a + n_heads, concentration0=b + n_tails)

theta_range = jnp.linspace(0.01, 0.99, 100)

ax = plt.gca()
ax2 = ax.twinx()
(plt2,) = ax2.plot(theta_range, exact_posterior.prob(theta_range), "g--", label="True Posterior")
(plt3,) = ax2.plot(theta_range, prior_dist().prob(theta_range), label="Prior")

likelihood = jax.vmap(lambda x: jnp.prod(likelihood_dist(x).prob(dataset)))(theta_range)
(plt1,) = ax.plot(theta_range, likelihood, "r-.", label="Likelihood")

ax.set_xlabel("theta")
ax.set_ylabel("Likelihood")
ax2.set_ylabel("Prior & Posterior")
ax2.legend(handles=[plt1, plt2, plt3], bbox_to_anchor=(1.6, 1));
```

## Laplace approximation from scratch in JAX

As mentioned in book2 section 7.4.3, Using laplace approximation, any distribution can be approximated as normal distribution having mean $\hat{\theta}$ and standard deviation as $H^{-1}$
\begin{align}
  H = \triangledown ^2_{\theta = \hat{\theta}} \log p(\theta|\mathcal{D}) \\
  p(\theta|\mathcal{D}) = \frac{1}{Z}p(\theta|\mathcal{D}) = \mathcal{N}(\theta |\hat{\theta}, H^{-1})
\end{align}
Where H is Hessian and $\hat{\theta}$ is the mode

### Find $\hat{\theta}$

No we find $\hat{\theta}$ ($\theta$_map) by minimizing negative log prior-likelhihood.

```python
def neg_log_prior_likelihood_fn(params, dataset):
    theta = params["theta"]
    likelihood_log_prob = likelihood_dist(theta).log_prob(dataset).sum()  # log probability of likelihood
    prior_log_prob = prior_dist().log_prob(theta)  # log probability of prior
    return -(likelihood_log_prob + prior_log_prob)  # negative log_prior_liklihood
```

```python
loss_and_grad_fn = jax.value_and_grad(neg_log_prior_likelihood_fn)
params = {"theta": 0.5}
neg_joint_log_prob, grads = loss_and_grad_fn(params, dataset)

optimizer = optax.adam(0.01)
opt_state = optimizer.init(params)
```

```python
@jax.jit
def train_step(carry, data_output):

    params = carry["params"]
    neg_joint_log_prob, grads = loss_and_grad_fn(params, dataset)

    opt_state = carry["opt_state"]
    updates, opt_state = optimizer.update(grads, opt_state)
    params = optax.apply_updates(params, updates)

    carry = {"params": params, "opt_state": opt_state}
    data_output = {"params": params, "loss": neg_joint_log_prob}

    return carry, data_output
```

```python
carry = {"params": params, "opt_state": opt_state}
data_output = {"params": params, "loss": neg_joint_log_prob}

n = 100
iterator = jnp.ones(n)
last_carry, output = jax.lax.scan(train_step, carry, iterator)
```

```python
loss = output["loss"]
plt.plot(loss, label="loss")
plt.legend();
```

```python
optimized_params = last_carry["params"]
theta_map = optimized_params["theta"]
print(f"theta_map = {theta_map}")
```

theta_map = 0.09168018400669098

### loc and scale of approximated normal posterior

```python
loc = theta_map  # loc of approximate posterior
print(f"loc = {loc}")
```

loc = 0.09168018400669098

```python
# scale of approximate posterior
scale = 1 / jnp.sqrt(jax.hessian(neg_log_prior_likelihood_fn)(optimized_params, dataset)["theta"]["theta"])
print(f"scale = {scale}")
```

scale = 0.08733917027711868

### True posterior and laplace approximated posterior

```python
plt.figure()
y = jnp.exp(dist.Normal(loc, scale).log_prob(theta_range))
plt.title("Quadratic approximation")
plt.plot(theta_range, y, label="laplace approximation", color="tab:red")
plt.plot(theta_range, exact_posterior.prob(theta_range), label="true posterior", color="tab:green", linestyle="--")
plt.xlabel("$\\theta$")
plt.ylabel("$p(\\theta)$")
sns.despine()
plt.legend()
savefig("bb_laplace")  # set FIG_DIR =  "path/to/figure" enviornment variable to save figure
```

## Pymc

```python
try:
    import pymc3 as pm
except ModuleNotFoundError:
    %pip install pymc3
    import pymc3 as pm
try:
    import scipy.stats as stats
except ModuleNotFoundError:
    %pip install scipy
    import scipy.stats as stats

import scipy.special as sp

try:
    import arviz as az
except ModuleNotFoundError:
    %pip install arviz
    import arviz as az

import math
```

    WARNING:theano.link.c.cmodule:install mkl with `conda install mkl-service`: No module named 'mkl'
    WARNING:theano.tensor.blas:Using NumPy C-API based implementation for BLAS functions.

```python
# Laplace
with pm.Model() as normal_aproximation:
    theta = pm.Beta("theta", 1.0, 1.0)
    y = pm.Binomial("y", n=1, p=theta, observed=dataset)  # Bernoulli
    mean_q = pm.find_MAP()
    std_q = ((1 / pm.find_hessian(mean_q, vars=[theta])) ** 0.5)[0]
    loc = mean_q["theta"]

# plt.savefig('bb_laplace.pdf');
```

  <progress value='8' class='' max='8' style='width:300px; height:20px; vertical-align: middle;'></progress>
  100.00% [8/8 00:00<00:00 logp = -7.6246, ||grad|| = 4.5]

```python
x = theta_range

plt.figure()
plt.plot(x, stats.norm.pdf(x, loc, std_q), "--", label="Laplace")
post_exact = stats.beta.pdf(x, n_heads + 1, n_tails + 1)
plt.plot(x, post_exact, label="exact")
plt.title("Quadratic approximation")
plt.xlabel("θ", fontsize=14)
plt.yticks([])
plt.legend()
```


# NBA Team Ratings, Pt. 3

In NBA Team Ratings, Pt. 1, I built a simple model that applyied shrinkage while taking matchups and home-court into account. In part 2, I extended that model to provide more explanation for how each team achieves its results on each side of the floor, looking at the results of a possession on a more granular level by modeling individual outcomes instead of the average one.
Unfortunately, that model doesn't really speak the "language of basketball"; no one refers to a team's probability of scoring 0 points on a possession for example.
Instead, basketball analysts focus on the process of scoring with statistics like shooting percentage, rebounds, and turnovers.
What we'd really like to model is the relationship between these process metrics, and scoring outcomes.

To do that, I've modified Dean Oliver's "Four Factors of Basketball Success," replacing them with seven:
1. $ \frac{\text{3pt att}}{\text{shot}} $: percentage of shots that were attempted from three point range
2. $ \text{2pt}\% $: percentage of two point shot attempts made
3. $ \text{3pt}\% $: percentage of three point shot attempts made
4. $ \text{reb}\% $: percentage of rebounds grabbed
5. $ \text{to}\% $: percentage of possessions that end in a turnover
6. $ \frac{\text{ft att}}{\text{poss}} $: free throws attempted per possession
7. $ \text{ft}\% $: percentage of free throw attempts made

The math is a bit complicated, but with these seven factors, we can accurately estimate the points scored per possession.

\begin{align}
\frac{\text{points}}{\text{poss}} &= \frac{\text{shots}}{\text{poss}} \times \frac{\text{points}}{\text{shot}} + \frac{\text{ft att}}{\text{poss}} \times \text{ft}\% \notag \\
\text{where } \frac{\text{points}}{\text{shot}} &= 3 \times \frac{\text{3pt att}}{\text{shot}} \times \text{3pt}\% + 2 \times \left(1 - \frac{\text{3pt att}}{\text{shot}} \right) \times \text{2pt}\% \notag
\end{align}

These equations have one undefined term which must be estimated: $ \frac{\text{shots}}{\text{poss}} $.
A possession can have any number of shots thanks to the resetting nature of a missed shot attempt combined with an offensive rebound.
To clarify this, we'll need a new term: a shot opportunity. Every possession starts with a shot opportunity, and has additional shot opportunities for each offensive rebound, which acts to reset the possession.

\begin{align}
\frac{\text{shots}}{\text{poss}} &= 0 \times \mathbb{P}(0) + 1 \times \mathbb{P}(1) + 2 \times \mathbb{P}(2) + \cdots \notag \\
\text{where } \mathbb{P}(k) &\cong p_\text{last shot} \times p_\text{reset}^{k - 1} \notag \\
p_\text{last shot} &\equiv p_\text{shot} \times \left( p_\text{make} + p_\text{miss} \times \left(p_\text{def reb} + p_\text{off reb} \times \left(1 - p_\text{shot} \right) \right) \right) \notag \\
p_\text{reset} &\equiv p_\text{shot} \times p_\text{miss} \times p_\text{off reb} \notag \\
p_\text{make} &\equiv \frac{\text{3pt att}}{\text{shot}} \times \text{3pt}\% + \left(1 - \frac{\text{3pt att}}{\text{shot}} \right) \times \text{2pt}\% \notag \\
p_\text{miss} &\equiv 1 - p_\text{make} \notag
\end{align}

Note that I've labeled that $ \mathbb{P}(k) $ here is approximate. We're assuming that the events in a possession are independent, i.e. when a possession resets with a missed shot \& offensive rebound, this is the same state as the beginning of a possession.
Of course, this isn't true, since an offensive rebound can often lead to an immediate "putback" attempt, the shot clock is usually set to 14 instead of the usual 24, there aren't fast breaks, and so on. Still, empirically, this assumption seems to work.

That set aside, we're again left with one undefined: $ p_\text{shot} $, the probability of attempting a shot given a shot opportunity.
Turnovers prevent shots, but less obviously, so do shooting fouls, since they don't count as a shot attempt if the shot is missed.
Thus, to estimate this factor, I ran a simple linear regression to estimate it:

\begin{align}
p_\text{shot} &\cong k_\text{1} \times 1 + k_\text{to} \times \text{to}\% + k_\text{fta} \times \frac{\text{ft att}}{\text{poss}} \notag \\
&= 1 - 0.85 \times \text{to}\% - 0.4 \times \frac{\text{ft att}}{\text{poss}} \notag
\end{align}

With all these terms defined, we can reformulate $ \frac{\text{shots}}{\text{poss}} $ as the sum of an infinite series:

\begin{align}
\frac{\text{shots}}{\text{poss}} &= \sum_{k=0}^{\infty}{ k \times p_\text{last shot} \times p_\text{reset}^{k - 1}} \notag \\
&= p_\text{last shot} \sum_{k=1}^{\infty}{ k \times p_\text{reset}^{k - 1}} \notag \\
&= \frac{p_\text{last shot}}{p_\text{reset}} \sum_{k=1}^{\infty}{ k \times p_\text{reset}^k} \notag \\
&= \frac{p_\text{last shot}}{p_\text{reset}} \times \frac{p_\text{reset}}{(1 - p_\text{reset})^2} \notag \\
&= \frac{p_\text{last shot}}{(1 - p_\text{reset})^2} \notag
\end{align}

All that work done, we're left with an almost exact method of calculating a team's points per possession given the seven factors.
So how did I model those seven factors?
Well conveniently, they can be each modeled independently of one another.
As well, six of the seven are simple binary probabilities, allowing me to fit the data using a simple Binomial distribution and Beta priors, much easier than the Multinomial distribution and Dirichlet priors from before.
For the one continuous variable ($ \frac{\text{ft att}}{\text{poss}} $), I used a simple Normal distribution with Gamma priors.
As well, I modeled the pace of each game (seconds per possession) using a similar approach, a Gamma distribution for both the observations and priors.
Most of these parameters depend on the offense, defense, and home court, with the following exceptions:
1. $ \frac{\text{3pt att}}{\text{shot}} $ didn't depend on home court
2. pace didn't depend on home court
3. $ \text{ft}\% $ didn't depend on the defense

Running this final model on the 2019 NBA Regular Season, we see similar patterns of shrinkage of extreme values, but with more variation, due to the model factoring in more inputs:

```python
from matplotlib import pyplot as plt

from pynba import halfgames_from_file, teams_from_file, team_id_to_abb, plot_ratings, plot_paces, use_blackontrans_style

league = "nba"
year = 2019
season_type = "Regular Season"
```

```python
halfgames = halfgames_from_file(league, year, season_type)
team_stats = teams_from_file(league, year, season_type)
team_stats
```

```python
team_id_to_team_abb = team_id_to_abb(league, year)
team_abb_to_team_id = {
    team_abb: team_id
    for team_id, team_abb in team_id_to_team_abb.items()
}
points_mu = halfgames["points_scored"].sum() / halfgames["possession_num"].sum() * 100
pace_mu = halfgames["possession_num"].sum() / halfgames["duration"].sum() * 12 * 4 * 60
```

```python
use_blackontrans_style()

fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(1, 1, 1)
for team, off_aa, def_aa in zip(team_stats["team"], team_stats["off_scoring_above_average"], team_stats["def_scoring_above_average"]):
    team_id = team_abb_to_team_id[team]
    filt = halfgames["off_team_id"] == team_id
    raw_off = halfgames.loc[filt, "points_scored"].sum() / halfgames.loc[filt, "possession_num"].sum() * 100 - points_mu
    filt = halfgames["def_team_id"] == team_id
    raw_def = points_mu - halfgames.loc[filt, "points_scored"].sum() / halfgames.loc[filt, "possession_num"].sum() * 100
    ax.annotate("", xy=(off_aa, def_aa), xytext=(raw_off, raw_def), arrowprops={"arrowstyle": "->", "color": "white"})
plot_ratings(team_stats, ax)

```

Finally, we also get a nice new plot, showing the pace for each team. Here we see that the model is doing very little in comparison to the raw values:

```python
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(1, 1, 1)
for team, off_aa, def_aa in zip(team_stats["team"], team_stats["off_pace_above_average"], team_stats["def_pace_above_average"]):
    team_id = team_abb_to_team_id[team]
    filt = halfgames["off_team_id"] == team_id
    raw_off = (halfgames.loc[filt, "possession_num"].sum() / halfgames.loc[filt, "duration"].sum() * 12 * 4 * 60 - pace_mu) / 2
    filt = halfgames["def_team_id"] == team_id
    raw_def = (halfgames.loc[filt, "possession_num"].sum() / halfgames.loc[filt, "duration"].sum() * 12 * 4 * 60 - pace_mu) / 2
    ax.annotate("", xy=(off_aa, def_aa), xytext=(raw_off, raw_def), arrowprops={"arrowstyle": "->", "color": "white"})
plot_paces(team_stats, ax)

```