# Customizing Potentials in JAX MD

```
#@title Imports & Utils
!pip install --upgrade -q https://storage.googleapis.com/jax-releases/cuda$(echo $CUDA_VERSION | sed -e 's/\.//' -e 's/\..*//')/jaxlib-$(pip search jaxlib | grep -oP '[0-9\.]+' | head -n 1)-cp36-none-linux_x86_64.whl
!pip install --upgrade -q jax
!pip install -q git+https://www.github.com/conference-submitter/jax-md

import numpy as onp

import jax.numpy as np
from jax import random
from jax import jit, grad, vmap, value_and_grad
from jax import lax
from jax import ops

from jax.config import config
config.update("jax_enable_x64", True)

from jax_md import space, smap, energy, minimize, quantity, simulate, partition

from functools import partial
import time

f32 = np.float32
f64 = np.float64

import matplotlib
import matplotlib.pyplot as plt
plt.rcParams.update({'font.size': 16})
#import seaborn as sns
#sns.set_style(style='white')

def format_plot(x, y):
  plt.grid(True)
  plt.xlabel(x, fontsize=20)
  plt.ylabel(y, fontsize=20)

def finalize_plot(shape=(1, 0.7)):
  plt.gcf().set_size_inches(
    shape[0] * 1.5 * plt.gcf().get_size_inches()[1],
    shape[1] * 1.5 * plt.gcf().get_size_inches()[1])
  plt.tight_layout()

def calculate_bond_data(displacement_or_metric, R, dr_cutoff, species=None):
  if( not(species is None)):
    assert(False)

  metric = space.map_product(space.canonicalize_displacement_or_metric(displacement))
  dr = metric(R,R)

  dr_include = np.triu(np.where(dr<1, 1, 0)) - np.eye(R.shape[0],dtype=np.int32)
  index_list=np.dstack(np.meshgrid(np.arange(N), np.arange(N), indexing='ij'))

  i_s = np.where(dr_include==1, index_list[:,:,0], -1).flatten()
  j_s = np.where(dr_include==1, index_list[:,:,1], -1).flatten()
  ij_s = np.transpose(np.array([i_s,j_s]))

  bonds = ij_s[(ij_s!=np.array([-1,-1]))[:,1]]
  lengths = dr.flatten()[(ij_s!=np.array([-1,-1]))[:,1]]

  return bonds, lengths

def plot_system(R,box_size,species=None,ms=20):
  R_plt = onp.array(R)

  if(species is None):
    plt.plot(R_plt[:, 0], R_plt[:, 1], 'o', markersize=ms)
  else:
    for ii in range(np.amax(species)+1):
      Rtemp = R_plt[species==ii]
      plt.plot(Rtemp[:, 0], Rtemp[:, 1], 'o', markersize=ms)

  plt.xlim([0, box_size])
  plt.ylim([0, box_size])
  plt.xticks([], [])
  plt.yticks([], [])

  finalize_plot((1,1))

key = random.PRNGKey(0)
```

##Prerequisites

This cookbook assumes a working knowledge of Python and Numpy. The concept of broadcasting is particularly important both in this cookbook and in JAX MD.

We also assume a basic knowlege of [JAX](https://github.com/google/jax/), which JAX MD is built on top of. Here we briefly review a few JAX basics that are important for us:

*   ```jax.vmap``` allows for automatic vectorization of a function. What this means is that if you have a function that takes an input ```x``` and returns an output ```y```, i.e. ```y = f(x)```, then ```vmap``` will transform this function to act on an array of ```x```'s and return an array of ```y```'s, i.e. ```Y = vmap(f)(X)```, where ```X=np.array([x1,x2,...,xn])``` and ```Y=np.array([y1,y2,...,yn])```.

*   ```jax.grad``` employs automatic differentiation to transform a function into a new function that calculates its gradient, for example: ```dydx = grad(f)(x)```.

*   ```jax.lax.scan``` allows for efficient for-loops that can be compiled and differentiated over. See [here](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html#jax.lax.scan) for more details.

*   [Random numbers are different in JAX.](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#%F0%9F%94%AA-Random-Numbers) The details aren't necessary for this cookbook, but if things look a bit different, this is why.

##The basics of user-defined potentials

###Create a user defined potential function to use throughout this cookbook

Here we create a custom potential that has a short-ranged, non-diverging repulsive interaction and a medium-ranged Morse-like attractive interaction. It takes the following form:
\begin{equation}
V(r) =
\begin{cases}
    \frac{1}{2} k (r-r_0)^2 - D_0,&  r < r_0\\
    D_0\left( e^{-2\alpha (r-r_0)} -2 e^{-\alpha(r-r_0)}\right),              & r \geq r_0
\end{cases}
\end{equation}
and has 4 parameters: $D_0$, $\alpha$, $r_0$, and $k$.

```
def harmonic_morse(dr, D0=5.0, alpha=5.0, r0=1.0, k=50.0, **kwargs):
  U = np.where(dr < r0,
               0.5 * k * (dr - r0)**2 - D0,
               D0 * (np.exp(-2. * alpha * (dr - r0)) - 2. * np.exp(-alpha * (dr - r0)))
               )
  return np.array(U, dtype=dr.dtype)
```

plot $V(r)$.

```
drs = np.arange(0,3,0.01)
U = harmonic_morse(drs)
plt.plot(drs,U)
format_plot(r'$r$', r'$V(r)$')
finalize_plot()
```

###Calculate the energy of a system of interacting particles

We now want to calculate the energy of a system of $N$ spheres in $d$ dimensions, where each particle interacts with every other particle via our user-defined function $V(r)$. The total energy is
\begin{equation}
E_\text{total} = \sum_{i<j}V(r_{ij}),
\end{equation}
where $r_{ij}$ is the distance between particles $i$ and $j$.

Our first task is to set up the system by specifying the $N$, $d$, and the size of the simulation box. We then use JAX's internal random number generator to pick positions for each particle.

```
N = 50
dimension = 2
box_size = 6.8

key, split = random.split(key)
R = random.uniform(split, (N,dimension), minval=0.0, maxval=box_size, dtype=f64)

plot_system(R,box_size)
```

At this point, we could manually loop over all particle pairs and calculate the energy, keeping track of boundary conditions, etc. Fortunately, JAX MD has machinery to automate this.

First, we must define two functions, ```displacement``` and ```shift```, which contain all the information of the simulation box, boundary conditions, and underlying metric. ```displacement``` is used to calculate the vector displacement between particles, and ```shift``` is used to move particles. For most cases, it is recommended to use JAX MD's built in functions, which can be called using:
*   ``` displacement, shift = space.free()```
*   ``` displacement, shift = space.periodic(box_size)```
*   ``` displacement, shift = space.periodic_general(T)```

For demonstration purposes, we will define these manually for a square periodic box, though without proper error handling, etc. The following should have the same functionality as ```displacement, shift = space.periodic(box_size)```.

```
def setup_periodic_box(box_size):
  def displacement_fn(Ra, Rb, **unused_kwargs):
    dR = Ra - Rb
    return np.mod(dR + box_size * f32(0.5), box_size) - f32(0.5) * box_size

  def shift_fn(R, dR, **unused_kwargs):
    return np.mod(R + dR, box_size)

  return displacement_fn, shift_fn

displacement, shift = setup_periodic_box(box_size)
```

We now set up a function to calculate the total energy of the system. The JAX MD function ```smap.pair``` takes a given potential and promotes it to act on all particle pairs in a system. ```smap.pair``` does not actually return an energy, rather it returns a function that can be used to calculate the energy.

For convenience and readability, we wrap ```smap.pair``` in a new function called ```harmonic_morse_pair```. For now, ignore the species keyword, we will return to this later.

```
def harmonic_morse_pair(
    displacement_or_metric, species=None, D0=5.0, alpha=10.0, r0=1.0, k=50.0):
  D0 = np.array(D0, dtype=f32)
  alpha = np.array(alpha, dtype=f32)
  r0 = np.array(r0, dtype=f32)
  k = np.array(k, dtype=f32)
  return smap.pair(
      harmonic_morse,
      space.canonicalize_displacement_or_metric(displacement_or_metric),
      species=species,
      D0=D0,
      alpha=alpha,
      r0=r0,
      k=k)
```

Our helper function can be used to construct a function to compute the energy of the entire system as follows.

```
# Create a function to calculate the total energy with specified parameters
energy_fn = harmonic_morse_pair(displacement,D0=5.0,alpha=10.0,r0=1.0,k=500.0)

# Use this to calculate the total energy
print(energy_fn(R))

# Use grad to calculate the net force
force = -grad(energy_fn)(R)
print(force[:5])
```

We are now in a position to use our energy function to manipulate the system. As an example, we perform energy minimization using JAX MD's implementation of the FIRE algorithm.

We start by defining a function that takes an energy function, a set of initial positions, and a shift function and runs a specified number of steps of the minimization algorithm. The function returns the final set of positions and the maximum absolute value component of the force. We will use this function throughout this cookbook.

```
def run_minimization(energy_fn, R_init, shift, num_steps=5000):
  dt_start = 0.001
  dt_max   = 0.004
  init,apply=minimize.fire_descent(jit(energy_fn),shift,dt_start=dt_start,dt_max=dt_max)
  apply = jit(apply)

  @jit
  def scan_fn(state, i):
    return apply(state), 0.

  state = init(R_init)
  state, _ = lax.scan(scan_fn,state,np.arange(num_steps))

  return state.position, np.amax(np.abs(-grad(energy_fn)(state.position)))
```

Now run the minimization with our custom energy function.

```
Rfinal, max_force_component = run_minimization(energy_fn, R, shift)
print('largest component of force after minimization = {}'.format(max_force_component))
plot_system( Rfinal, box_size )
```

###Create a truncated potential

It is often desirable to have a potential that is strictly zero beyond a well-defined cutoff distance. In addition, MD simulations require the energy and force (i.e. first derivative) to be continuous. To easily modify an existing potential $V(r)$ to have this property, JAX MD follows the approach [taken by HOOMD Blue](https://hoomd-blue.readthedocs.io/en/stable/module-md-pair.html#hoomd.md.pair.pair).

Consider the function
\begin{equation}
S(r) =
\begin{cases}
    1,& r<r_\mathrm{on} \\
    \frac{(r_\mathrm{cut}^2-r^2)^2 (r_\mathrm{cut}^2 + 2r^2 - 3 r_\mathrm{on}^2)}{(r_\mathrm{cut}^2-r_\mathrm{on}^2)^3},&  r_\mathrm{on} \leq r < r_\mathrm{cut}\\
    0,& r \geq r_\mathrm{cut}
\end{cases}
\end{equation}

Here we plot both $S(r)$ and $\frac{dS(r)}{dr}$, both of which are smooth and strictly zero above $r_\mathrm{cut}$.

```
dr = np.arange(0,3,0.01)
S = energy.multiplicative_isotropic_cutoff(lambda dr: 1, r_onset=1.5, r_cutoff=2.0)(dr)
ngradS = vmap(grad(energy.multiplicative_isotropic_cutoff(lambda dr: 1, r_onset=1.5, r_cutoff=2.0)))(dr)
plt.plot(dr,S,label=r'$S(r)$')
plt.plot(dr,ngradS,label=r'$\frac{dS(r)}{dr}$')
plt.legend()
format_plot(r'$r$','')
finalize_plot()
```

We then use $S(r)$ to create a new function
\begin{equation}\tilde V(r) = V(r) S(r),
\end{equation}
which is exactly $V(r)$ below $r_\mathrm{on}$, strictly zero above $r_\mathrm{cut}$ and is continuous in its first derivative.

This is implemented in JAX MD through ```energy.multiplicative_isotropic_cutoff```, which takes in a potential function $V(r)$ (e.g. our ```harmonic_morse``` function) and returns a new function $\tilde V(r)$.

```
harmonic_morse_cutoff = energy.multiplicative_isotropic_cutoff(
    harmonic_morse, r_onset=1.5, r_cutoff=2.0)

dr = np.arange(0,3,0.01)
V = harmonic_morse(dr)
V_cutoff = harmonic_morse_cutoff(dr)
F = -vmap(grad(harmonic_morse))(dr)
F_cutoff = -vmap(grad(harmonic_morse_cutoff))(dr)
plt.plot(dr,V, label=r'$V(r)$')
plt.plot(dr,V_cutoff, label=r'$\tilde V(r)$')
plt.plot(dr,F, label=r'$-\frac{d}{dr} V(r)$')
plt.plot(dr,F_cutoff, label=r'$-\frac{d}{dr} \tilde V(r)$')
plt.legend()
format_plot('$r$', '')
plt.ylim(-13,5)
finalize_plot()
```

As before, we can use ```smap.pair``` to promote this to act on an entire system.

```
def harmonic_morse_cutoff_pair(
    displacement_or_metric, D0=5.0, alpha=5.0, r0=1.0, k=50.0,
    r_onset=1.5, r_cutoff=2.0):
  D0 = np.array(D0, dtype=f32)
  alpha = np.array(alpha, dtype=f32)
  r0 = np.array(r0, dtype=f32)
  k = np.array(k, dtype=f32)
  return smap.pair(
      energy.multiplicative_isotropic_cutoff(
          harmonic_morse, r_onset=r_onset, r_cutoff=r_cutoff),
      space.canonicalize_displacement_or_metric(displacement_or_metric),
      D0=D0,
      alpha=alpha,
      r0=r0,
      k=k)
```

This is implemented as before

```
# Create a function to calculate the total energy
energy_fn = harmonic_morse_cutoff_pair(displacement, D0=5.0, alpha=10.0, r0=1.0,
                                       k=500.0, r_onset=1.5, r_cutoff=2.0)

# Use this to calculate the total energy
print(energy_fn(R))

# Use grad to calculate the net force
force = -grad(energy_fn)(R)
print(force[:5])

# Minimize the energy using the FIRE algorithm
Rfinal, max_force_component = run_minimization(energy_fn, R, shift)
print('largest component of force after minimization = {}'.format(max_force_component))
plot_system( Rfinal, box_size )
```

##Specifying parameters

###Dynamic parameters

In the above examples, the strategy is to create a function ```energy_fn``` that takes a set of positions and calculates the energy of the system with all the parameters (e.g. ```D0```, ```alpha```, etc.) baked in. However, JAX MD allows you to override these baked-in values dynamically, i.e. when ```energy_fn``` is called.

For example, we can print out the minimized energy and force of the above system with the truncated potential:

```
print(energy_fn(Rfinal))
print(-grad(energy_fn)(Rfinal)[:5])
```

This uses the baked-in values of the 4 parameters: ```D0=5.0,alpha=10.0,r0=1.0,k=500.0```. If, for example, we want to dynamically turn off the attractive part of the potential, we simply pass ```D0=0``` to ```energy_fn```:

```
print(energy_fn(Rfinal, D0=0))
```

Since changing the potential moves the minimum, the force will not be zero:

```
print(-grad(energy_fn)(Rfinal, D0=0)[:5])
```

This ability to dynamically pass parameters is very powerful. For example, if you want to shrink particles each step during a simulation, you can simply specify a different ```r0``` each step.

This is demonstrated below, where we run a Brownian dynamics simulation at zero temperature with continuously decreasing ```r0```. The details of ```simulate.brownian``` are beyond the scope of this cookbook, but the idea is that we pass a new value of ```r0``` to the function ```apply``` each time it is called. The function ```apply``` takes a step of the simulation, and internally it passes any extra parameters like ```r0``` to ```energy_fn```.

```
def run_brownian(energy_fn, R_init, shift, key, num_steps):
  init, apply = simulate.brownian(energy_fn, shift, dt=0.00001, T_schedule=0.0, gamma=0.1)
  apply = jit(apply)

  # Define how r0 changes for each step
  r0_initial = 1.0
  r0_final = .5
  def Get_r0(t):
    return r0_final + (r0_initial-r0_final)*(num_steps-t)/num_steps

  @jit
  def scan_fn(state, t):
    # Dynamically pass r0 to apply, which passes it on to energy_fn
    return apply(state, t, r0=Get_r0(t)), 0

  key, split = random.split(key)
  state = init(split, R_init)

  state, _ = lax.scan(scan_fn,state,np.arange(num_steps))
  return state.position, np.amax(np.abs(-grad(energy_fn)(state.position)))
```

If we use the previous result as the starting point for the Brownian Dynamics simulation, we find exactly what we would expect, the system contracts into a finite cluster, held together by the attractive part of the potential.

```
key, split = random.split(key)
Rfinal2, max_force_component = run_brownian(energy_fn, Rfinal, shift, split, num_steps=6000)
plot_system( Rfinal2, box_size )
```

###Particle-specific parameters

Our example potential has 4 parameters: ```D0```, ```alpha```, ```r0```, and ```k```. The usual way to pass these parameters is as a scalar (e.g. ```D0=5.0```), in which case that parameter is fixed for every particle pair. However, Python broadcasting allows for these parameters to be specified separately for every different particle pair by passing an $(N,N)$ array rather than a scalar.

As an example, let's do this for the parameter ```r0```, which is an effective way of generating a system with continuous polydispersity in particle size. Note that the polydispersity disrupts the crystalline order after minimization.

```
# Draw the radii from a uniform distribution
key, split = random.split(key)
radii = random.uniform(split, (N,), minval=1.0, maxval=2.0, dtype=f64)

# Rescale to match the initial volume fraction
radii = np.array([radii * np.sqrt(N/(4.*np.dot(radii,radii)))])

# Turn this into a matrix of sums
r0_matrix = radii+radii.transpose()

# Create the energy function using r0_matrix
energy_fn = harmonic_morse_pair(displacement, D0=5.0, alpha=10.0, r0=r0_matrix,
                                       k=500.0)

# Minimize the energy using the FIRE algorithm
Rfinal, max_force_component = run_minimization(energy_fn, R, shift)
print('largest component of force after minimization = {}'.format(max_force_component))
plot_system( Rfinal, box_size )
```

In addition to standard Python broadcasting, JAX MD allows for the special case of additive parameters. If a parameter is passed as a (N,) array ```p_vector```, JAX MD will convert this into a (N,N) array ```p_matrix``` where ```p_matrix[i,j] = 0.5 (p_vector[i] + p_vector[j])```. This is a JAX MD specific ability and not a feature of Python broadcasting.

As it turns out, our above polydisperse example falls into this category. Therefore, we could achieve the same result by passing ```r0=2.0*radii```.

```
# Create the energy function the radii array
energy_fn = harmonic_morse_pair(displacement, D0=5.0, alpha=10.0, r0=2.*radii,
                                       k=500.0)

# Minimize the energy using the FIRE algorithm
Rfinal, max_force_component = run_minimization(energy_fn, R, shift)
print('largest component of force after minimization = {}'.format(max_force_component))
plot_system( Rfinal, box_size )
```

### Species

It is often important to specify parameters differently for different particle pairs, but doing so with full ($N$,$N$) matrices is both inefficient and obnoxious. JAX MD allows users to create species, i.e. $N_s$ groups of particles that are identical to each other, so that parameters can be passed as much smaller ($N_s$,$N_s$) matrices.

First, create an array that specifies which particles belong in which species. We will divide our system into two species.

```
N_0 = N // 2  # Half the particles in species 0
N_1 = N - N_0 # The rest in species 1
species = np.array([0] * N_0 + [1] * N_1, dtype=np.int32)
print(species)
```

Next, create the $(2,2)$ matrix of ```r0```'s, which are set so that the overall volume fraction matches our monodisperse case.

```
rsmall=0.41099747 # Match the total volume fraction
rlarge=1.4*rsmall
r0_species_matrix = np.array([[2*rsmall, rsmall+rlarge],
                              [rsmall+rlarge, 2*rlarge]])
print(r0_species_matrix)
```

```
energy_fn = harmonic_morse_pair(displacement, species=species, D0=5.0,
                                alpha=10.0, r0=r0_species_matrix, k=500.0)

Rfinal, max_force_component = run_minimization(energy_fn, R, shift)
print('largest component of force after minimization = {}'.format(max_force_component))

plot_system( Rfinal, box_size, species=species )
```

###Dynamic Species

Just like standard parameters, the species list can be passed dynamically as well. However, unlike standard parameters, you have to tell ```smap.pair``` that the species will be specified dynamically. To do this, set  ```species=quantity.Dynamic``` when creating your energy function.

The following sets up an energy function where the attractive part of the interaction only exists between members of the first species, but where the species will be defined dynamically.

```
D0_species_matrix = np.array([[ 5.0, 0.0],
                              [0.0,  0.0]])

energy_fn = harmonic_morse_pair(displacement, species=quantity.Dynamic,
                                D0=D0_species_matrix, alpha=10.0, r0=0.5, k=500.0)
```

Now we set up a finite temperature Brownian Dynamics simulation where, at every step, particles on the left half of the simulation box are assigned to species 0, while particles on the right half are assigned to species 1.

```
def run_brownian(energy_fn, R_init, shift, key, num_steps):
  init, apply = simulate.brownian(energy_fn, shift, dt=0.00001, T_schedule=1.0, gamma=0.1)
  # apply = jit(apply)

  # Define a function to recalculate the species each step
  def get_species(R):
    return np.where(R[:,0] < box_size / 2, 0, 1)

  #@jit
  def scan_fn(state, t):
    # Recalculate the species list
    species = get_species(state.position)
    # Dynamically pass species to apply, which passes it on to energy_fn
    return apply(state, species=species, species_count=2), 0

  key, split = random.split(key)
  state = init(split, R_init)

  state, _ = lax.scan(scan_fn,state,np.arange(num_steps))
  return state.position,np.amax(np.abs(-grad(energy_fn)(state.position,
                                                        species=get_species(state.position),
                                                        species_count=2)))
```

When we run this, we see that particles on the left side form clusters while particles on the right side do not.

```
key, split = random.split(key)
Rfinal, max_force_component = run_brownian(energy_fn, R, shift, split, num_steps=10000)
plot_system( Rfinal, box_size )
```

##Efficeiently calculating neighbors

The most computationally expensive part of most MD programs is calculating the force between all pairs of particles. Generically, this scales with $N^2$. However, for systems with isotropic pairwise interactions that are strictly zero beyond a cutoff, there are techniques to dramatically improve the efficiency. The two most common methods are cell list and neighbor lists.

**Cell lists**

The technique here is to divide space into small cells that are just larger than the largest interaction range in the system. Thus, if particle $i$ is in cell $c_i$ and particle $j$ is in cell $c_j$, $i$ and $j$ can only interact if $c_i$ and $c_j$ are neighboring cells. Rather than searching all $N^2$ combinations of particle pairs for non-zero interactions, you only have to search the particles in the neighboring cells.

**Neighbor lists**

Here, for each particle $i$, we make a list of *potential* neighbors: particles $j$ that are within some threshold distance $r_\mathrm{threshold}$. If $r_\mathrm{threshold} = r_\mathrm{cutoff} + \Delta r_\mathrm{threshold}$ (where $r_\mathrm{cutoff}$ is the largest interaction range in the system and $\Delta r_\mathrm{threshold}$ is an appropriately chosen buffer size), then all interacting particles will appear in this list as long as no particles moves by more than $\Delta r_\mathrm{threhsold}/2$. There is a tradeoff here: smaller $\Delta r_\mathrm{threhsold}$ means fewer particles to search over each MD step but the list must be recalculated more often, while larger $\Delta r_\mathrm{threhsold}$ means slower force calculates but less frequent neighbor list calculations.

In practice, the most efficient technique is often to use cell lists to calculate neighbor lists. In JAX MD, this occurs under the hood, and so only calls to neighbor-list functionality are necessary.

To implement neighbor lists, we need two functions: 1) a function to create and update the neighbor list, and 2) an energy function that uses a neighbor list rather than operating on all particle pairs. We create these functions with ```partition.neighbor_list``` and ```smap.pair_neighbor_list```, respectively.

```partition.neighbor_list``` takes basic box information as well as the maximum interaction range ```r_cutoff``` and the buffer size ```dr_threshold```.

```
 def harmonic_morse_cutoff_neighbor_list(
    displacement_or_metric,
    box_size,
    species=None,
    D0=5.0,
    alpha=5.0,
    r0=1.0,
    k=50.0,
    r_onset=1.0,
    r_cutoff=1.5,
    dr_threshold=2.0,
    **kwargs):

  D0 = np.array(D0, dtype=np.float32)
  alpha = np.array(alpha, dtype=np.float32)
  r0 = np.array(r0, dtype=np.float32)
  k = np.array(k, dtype=np.float32)
  r_onset = np.array(r_onset, dtype=np.float32)
  r_cutoff = np.array(r_cutoff, np.float32)
  dr_threshold = np.float32(dr_threshold)

  neighbor_fn = partition.neighbor_list(
        displacement_or_metric,
        box_size,
        r_cutoff,
        dr_threshold)

  energy_fn = smap.pair_neighbor_list(
    energy.multiplicative_isotropic_cutoff(harmonic_morse, r_onset, r_cutoff),
    space.canonicalize_displacement_or_metric(displacement_or_metric),
    species=species,
    D0=D0,
    alpha=alpha,
    r0=r0,
    k=k)

  return neighbor_fn, energy_fn
```

To test this, we generate our new ```neighbor_fn``` and ```energy_fn```, as well as a comparison energy function using the default approach.

```
r_onset  = 1.5
r_cutoff = 2.0
dr_threshold = 1.0

neighbor_fn, energy_fn = harmonic_morse_cutoff_neighbor_list(
    displacement, box_size, D0=5.0, alpha=10.0, r0=1.0, k=500.0,
    r_onset=r_onset, r_cutoff=r_cutoff, dr_threshold=dr_threshold)

energy_fn_comparison = harmonic_morse_cutoff_pair(
    displacement, D0=5.0, alpha=10.0, r0=1.0, k=500.0,
    r_onset=r_onset, r_cutoff=r_cutoff)
```

Next, we use ```neighbor_fn``` and the current set of positions to populate the neighbor list.

```
nbrs = neighbor_fn(R)
```

To calculate the energy, we pass ```nbrs.idx``` to ```energy_fn```. The energy matches the comparison.

```
print(energy_fn(R,neighbor_idx=nbrs.idx))
print(energy_fn_comparison(R))
```

Note that by default ```neighbor_fn``` uses a cell list internally to populate the neighbor list. This approach fails when the box size in any dimension is less than 3 times $r_\mathrm{threhsold} = r_\mathrm{cutoff} + \Delta r_\mathrm{threshold}$. In this case, ```neighbor_fn``` automatically turns off the use of cell lists, and instead searches over all particle pairs. This can also be done manually by passing ```disable_cell_list=True``` to ```partition.neighbor_list```. This can be useful for debugging or for small systems where the overhead of cell lists outweighs the benefit.

###Updating neighbor lists

The function ```neighbor_fn``` has two different usages, depending on how it is called. When used as above, i.e. ```nbrs = neighbor_fn(R)```, a new neighbor list is generated from scratch. Internally, JAX MD uses the given positions ```R``` to estimate a maximum capacity, i.e. the maximum number of neighbors any particle will have at any point during the use of the neighbor list. This estimate can be adjusted by passing a value of  ```capacity_multiplier``` to ```partition.neighbor_list```, which defaults to ```capacity_multiplier=1.25```.

Since the maximum capacity is not known ahead of time, this construction of the neighbor list cannot be compiled. However, once a neighbor list is created in this way, repopulating the list with the same maximum capacity is a simpler operation that *can* be compiled. This is done by calling ```nbrs = neighbor_fn(R, nbrs)```. Internally, this checks if any particle has moved more than $\Delta r_\mathrm{threshold}/2$ and, if so, recomputes the neighbor list. If the new neighbor list exceeds the maximum capacity for any particle, the boolean variable ```nbrs.did_buffer_overflow``` is set to ```True```.

These two uses together allow for safe and efficient neighbor list calculations. The example below demonstrates a typical simulation loop that uses neighbor lists.

```
def run_brownian_neighbor_list(energy_fn, neighbor_fn, R_init, shift, key, num_steps):
  nbrs = neighbor_fn(R_init)

  init, apply = simulate.brownian(energy_fn, shift, dt=0.00001, T_schedule=1.0, gamma=0.1)
  #apply = jit(apply)

  #@jit
  def body_fn(state, t):
    state, nbrs = state
    nbrs = neighbor_fn(state.position, nbrs)
    state = apply(state, neighbor_idx=nbrs.idx)
    return (state, nbrs), 0

  key, split = random.split(key)
  state = init(split, R_init)

  step = 0
  step_inc=100
  while step < num_steps/step_inc:
    rtn_state, _ = lax.scan(body_fn,(state,nbrs), np.arange(step_inc))
    new_state, nbrs = rtn_state
    # If the neighbor list overflowed, rebuild it and repeat part of
    # the simulation.
    if nbrs.did_buffer_overflow:
      nbrs = neighbor_fn(state.position)
    else:
      state = new_state
      step += 1

  return state.position
```

To run this, we consider a much larger system than we have to this point. Warning: running this may take a few minutes.

```
Nlarge = 100*N
box_size_large = 10*box_size
displacement_large, shift_large = setup_periodic_box(box_size_large)

key, split1, split2 = random.split(key,3)
Rlarge = random.uniform(split1, (Nlarge,dimension), minval=0.0, maxval=box_size_large, dtype=f64)

dr_threshold = 1.5
neighbor_fn, energy_fn = harmonic_morse_cutoff_neighbor_list(
    displacement_large, box_size_large, D0=5.0, alpha=10.0, r0=1.0, k=500.0,
    r_onset=r_onset, r_cutoff=r_cutoff, dr_threshold=dr_threshold)
energy_fn = jit(energy_fn)

start_time = time.clock()
Rfinal = run_brownian_neighbor_list(energy_fn, neighbor_fn, Rlarge, shift_large, split2, num_steps=4000)
end_time = time.clock()
print('run time = {}'.format(end_time-start_time))

plot_system( Rfinal, box_size_large, ms=2 )
```

##Bonds

Bonds are a way of specifying potentials between specific pairs of particles that are "on" regardless of separation. For example, it is common to employ a two-sided spring potential between specific particle pairs, but JAX MD allows the user to specify arbitrary potentials with static or dynamic parameters.

### Create and implement a bond potential

We start by creating a custom potential that corresponds to a bistable spring, taking the form
\begin{equation}
V(r) = a_4(r-r_0)^4 - a_2(r-r_0)^2.
\end{equation}
$V(r)$ has two minima, at $r = r_0 \pm \sqrt{\frac{a_2}{2a_4}}$.

```
def bistable_spring(dr, r0=1.0, a2=2, a4=5, **kwargs):
  return a4*(dr-r0)**4 - a2*(dr-r0)**2
```

Plot $V(r)$

```
drs = np.arange(0,2,0.01)
U = bistable_spring(drs)
plt.plot(drs,U)
format_plot(r'$r$', r'$V(r)$')
finalize_plot()
```

The next step is to promote this function to act on a set of bonds. This is done via ```smap.bond```, which takes our ```bistable_spring``` function, our displacement function, and a list of the bonds. It returns a function that calculates the energy for a given set of positions.

```
def bistable_spring_bond(
    displacement_or_metric, bond, bond_type=None, r0=1, a2=2, a4=5):
  """Convenience wrapper to compute energy of particles bonded by springs."""
  r0 = np.array(r0, f32)
  a2 = np.array(a2, f32)
  a4 = np.array(a4, f32)
  return smap.bond(
    bistable_spring,
    space.canonicalize_displacement_or_metric(displacement_or_metric),
    bond,
    bond_type,
    r0=r0,
    a2=a2,
    a4=a4)
```

However, in order to implement this, we need a list of bonds. We will do this by taking a system minimized under our original ```harmonic_morse``` potential:

```
R_temp, max_force_component = run_minimization(harmonic_morse_pair(displacement,D0=5.0,alpha=10.0,r0=1.0,k=500.0), R, shift)
print('largest component of force after minimization = {}'.format(max_force_component))
plot_system( R_temp, box_size )
```

We now place a bond between all particle pairs that are separated by less than 1.3. ```calculate_bond_data``` returns a list of such bonds, as well as a list of the corresponding current length of each bond.

```
bonds, lengths = calculate_bond_data(displacement, R_temp, 1.3)

print(bonds[:5])   # list of particle index pairs that form bonds
print(lengths[:5]) # list of the current length of each bond
```

We use this length as the ```r0``` parameter, meaning that initially each bond is at the unstable local maximum $r=r_0$.

```
bond_energy_fn = bistable_spring_bond(displacement, bonds, r0=lengths)
```

We now use our new ```bond_energy_fn``` to minimize the energy of the system. The expectation is that nearby particles should either move closer together or further apart, and the choice of which to do should be made collectively due to the constraint of constant volume. This is exactly what we see.

```
Rfinal, max_force_component = run_minimization(bond_energy_fn, R_temp, shift)
print('largest component of force after minimization = {}'.format(max_force_component))
plot_system( Rfinal, box_size )
```

###Specifying bonds dynamically

As with species or parameters, bonds can be specified dynamically, i.e. when the energy function is called. Importantly, note that this does NOT override bonds that were specified statically in ```smap.bond```.

```
# Specifying the bonds dynamically ADDS additional bonds.
#  Here, we dynamically pass the same bonds that were passed statically, which
#  has the effect of doubling the energy
print(bond_energy_fn(R))
print(bond_energy_fn(R,bonds=bonds, r0=lengths))
```

We won't go thorugh a further example as the implementation is exactly the same as specifying species or parameters dynamically, but the ability to employ bonds both statically and dynamically is a very powerful and general framework.

## Combining potentials

Most JAX MD functionality (e.g. simulations, energy minimizations) relies on a function that calculates energy for a set of positions. Importantly, while this cookbook focus on simple and robust ways of defining such functions, JAX MD is not limited to these methods; users can implement energy functions however they like.

As an important example, here we consider the case where the energy includes both a pair potential and a bond potential. Specifically, we combine ```harmonic_morse_pair``` with ```bistable_spring_bond```.

```
# Note, the code in the "Bonds" section must be run prior to this.
energy_fn = harmonic_morse_pair(displacement,D0=0.,alpha=10.0,r0=1.0,k=1.0)
bond_energy_fn = bistable_spring_bond(displacement, bonds, r0=lengths)
def combined_energy_fn(R):
  return energy_fn(R) + bond_energy_fn(R)
```

Here, we have set $D_0=0$, so the pair potential is just a one-sided repulsive harmonic potential. For particles connected with a bond, this raises the energy of the "contracted" minimum relative to the "extended" minimum.

```
drs = np.arange(0,2,0.01)
U = harmonic_morse(drs,D0=0.,alpha=10.0,r0=1.0,k=1.0)+bistable_spring(drs)
plt.plot(drs,U)
format_plot(r'$r$', r'$V(r)$')
finalize_plot()
```

This new energy function can be passed to the minimization routine (or any other JAX MD simulation routine) in the usual way.

```
Rfinal, max_force_component = run_minimization(combined_energy_fn, R_temp, shift)
print('largest component of force after minimization = {}'.format(max_force_component))
plot_system( Rfinal, box_size )
```

##Specifying forces instead of energies

So far, we have defined functions that calculate the energy of the system, which we then pass to JAX MD. Internally, JAX MD uses automatic differentiation to convert these into functions that calculate forces, which are necessary to evolve a system under a given dynamics. However, JAX MD has the option to pass force functions directly, rather than energy functions. This creates additional flexibility because some forces cannot be represented as the gradient of a potential.

As a simple example, we create a custom force function that zeros out the force of some particles. During energy minimization, where there is no stochastic noise, this has the effect of fixing the position of these particles.

First, we break the system up into two species, as before.

```
N_0 = N // 2  # Half the particles in species 0
N_1 = N - N_0 # The rest in species 1
species = np.array([0]*N_0 + [1]*N_1, dtype=np.int32)
print(species)
```

Next, we we creat our custom force function. Starting with our ```harmonic_morse``` pair potential, we calculate the force manually (i.e. using built-in automatic differentiation), and then multiply the force by the species id, which has the desired effect.

```
energy_fn = harmonic_morse_pair(displacement,D0=5.0,alpha=10.0,r0=1.0,k=500.0)
force_fn = quantity.force(energy_fn)

def custom_force_fn(R, **kwargs):
  return vmap(lambda a,b: a*b)(force_fn(R),species)
```

To run with our custom force function, we have to tell JAX MD that we are giving it a force rather than an energy. To do this, we define a new minimization function with a parameter ```quant``` that is passed to the JAX MD minimizer. If we set ```quant=quantity.Force```, this tells JAX MD that we are specifying a force function rather than an energy function.

```
def run_minimization_general(energy_or_force, R_init, shift, num_steps=5000, quant=quantity.Energy):
  dt_start = 0.001
  dt_max   = 0.004
  init,apply=minimize.fire_descent(jit(energy_or_force),shift,quant=quant,dt_start=dt_start,dt_max=dt_max)
  apply = jit(apply)

  @jit
  def scan_fn(state, i):
    return apply(state), 0.

  state = init(R_init)
  state, _ = lax.scan(scan_fn,state,np.arange(num_steps))

  return state.position, np.amax(np.abs(quantity.canonicalize_force(energy_or_force,quant)(state.position)))
```

We run this as usual, passing ```quant=quantity.Force```.

```
key, split = random.split(key)
Rfinal, _ = run_minimization_general(custom_force_fn, R, shift, quant=quantity.Force)
plot_system( Rfinal, box_size, species )
```

After the above minimization, the blue particles have the same positions as they did initially:

```
plot_system( R, box_size, species )
```

Note, this method for fixing particles only works when there is no stochastic noise (e.g. in Langevin or Brownian dynamics) because such noise affects partices whether or not they have a net force. A safer way to fix particles is to create a custom ```shift``` function.

##Coupled ensembles

For a final example that demonstrates the flexibility within JAX MD, lets do something that is particularly difficult in most standard MD packages. We will create a "coupled ensemble" -- i.e. a set of two identical systems that are connected via a $Nd$ dimensional spring. An extension of this idea is used, for example, in the Doubly Nudged Elastic Band method for finding transition states.

If the "normal" energy of each system is
\begin{equation}
U(R) = \sum_{i,j} V( r_{ij} ),
\end{equation}
where $r_{ij}$ is the distance between the $i$th and $j$th particles in $R$ and the $V(r)$ is a standard pair potential, and if the two sets of positions, $R_0$ and $R_1$, are coupled via the potential
\begin{equation}
U_\mathrm{spr}(R_0,R_1) = \frac 12 k_\mathrm{spr} \left| R_1 - R_0 \right|^2,
\end{equation}
so that the total energy of the system is
\begin{equation}
U_\mathrm{total} = U(R_0) + U(R_1) + U_\mathrm{spr}(R_0,R_1).
\end{equation}

```
energy_fn = harmonic_morse_pair(displacement,D0=5.0,alpha=10.0,r0=0.5,k=500.0)
def spring_energy_fn(Rall, k_spr=50.0, **kwargs):
  metric = vmap(space.canonicalize_displacement_or_metric(displacement), (0, 0), 0)
  dr = metric(Rall[0],Rall[1])
  return 0.5*k_spr*np.sum((dr)**2)
def total_energy_fn(Rall, **kwargs):
  return np.sum(vmap(energy_fn)(Rall)) + spring_energy_fn(Rall)
```

We now have to define a new shift function that can handle arrays of shape $(2,N,d)$. In addition, we make two copies of our initial positions ```R```, one for each system.

```
def shift_all(Rall, dRall, **kwargs):
  return vmap(shift)(Rall, dRall)
Rall = np.array([R,R])
```

Now, all we have to do is pass our custom energy and shift functions, as well as the $(2,N,d)$ dimensional initial position, to JAX MD, and proceed as normal.

As a demonstration, we define a simple and general Brownian Dynamics simulation function, similar to the simulation routines above except without the special cases (e.g. chaning ```r0``` or species).

```
def run_brownian_simple(energy_or_force, R_init, shift, key, num_steps, quant=quantity.Energy):
  init, apply = simulate.brownian(energy_or_force, shift, quant=quant, dt=0.00001, T_schedule=1.0, gamma=0.1)
  apply = jit(apply)

  @jit
  def scan_fn(state, t):
    return apply(state), 0

  key, split = random.split(key)
  state = init(split, R_init)

  state, _ = lax.scan(scan_fn,state,np.arange(num_steps))
  return state.position
```

Note that nowhere in this function is there any indication that we are simulating an ensemble of systems. This comes entirely form the inputs: i.e. the energy function, the shift function, and the set of initial positions.

```
key, split = random.split(key)
Rall_final = run_brownian_simple(total_energy_fn, Rall, shift_all, split, num_steps=10000)
```

The output also has shape $(2,N,d)$. If we display the results, we see that the two systems are in similar, but not identical, positions, showing that we have succeeded in simulating a coupled ensemble.

```
for Ri in Rall_final:
  plot_system( Ri, box_size )
finalize_plot((0.5,0.5))
```


# Efficient Estimation using 1 simulation

1. This notebook shows how to **estimate** a simple model using Simulated Minimum Distance (SMD)
2. It illustrates how an **efficient** estimator can be constructed using only 1 simulatoin, following the idea proposed by [Kirill Evdokimov](https://www.mit.edu/~kevdokim/ESMSM_sep16.pdf "Efficient Estimation with a Finite Number of Simulation Draws per Observation")

## Recap: Simulated Minimum Distance

**Data:** We assume that we have data available for $N$ households over $T$ periods, collected in $\{w_i\}_i^N$.

**Goal:** We wish to estimate the true, unknown, parameter vector $\theta_0$. We assume our model is correctly specified in the sense that the observed data stems from the model.

The **Simulated Minimum Distance (SMD)** estimator is

$$
\hat{\theta} = \arg\min_{\theta} g(\theta)'Wg(\theta)
$$

where $W$ is a $J\times J$ positive semidefinite **weighting matrix** and

$$
g(\theta)=\Lambda_{data}-\Lambda_{sim}(\theta)
$$

is the distance between $J\times1$ vectors of moments calculated in the data and the simulated data, respectively. Concretely,

$$
\Lambda_{data} = \frac{1}{N}\sum_{i=1}^N m(\theta_0|w_i) \\
\Lambda_{sim}(\theta) = \frac{1}{N_{sim}}\sum_{s=1}^{N_{sim}} m(\theta|w_s)
$$

are $J\times1$ vectors of moments calculated in the data and the simulated data, respectively.

**Variance of the estimator:** Recall that  the variance of the estimator was
$$
\begin{align}
\text{Var}(\hat{\theta})&=(1+S^{-1})\Gamma\Omega\Gamma'/N \\
\Gamma &= -(G'WG)^{-1}G'W \\
\Omega & = \text{Var}(m(\theta_0|w_i))
\end{align}
$$
where we implicitly used that $Var(m(\theta_0|w_i))=Var(m(\theta|w_s))$ and $Cov(m(\theta_0|w_i),m(\theta|w_s))=0$

**Efficient Estimator:** Using the "optimal" weighting matrix, $W=\Omega^{-1}$, gives the *lowest variance* for a given number of simulations, $S$, as
$$
\begin{align}
\text{Var}(\hat{\theta})&=(1+S^{-1})(G'\Omega^{-1}G)^{-1}/N
\end{align}
$$

> **Observation:** Only as $S\rightarrow\infty$ does the minimum variance of the SMD estimator approach the minimum variance of the GMM estimator.

> **Solution:** [Kirill Evdokimov](https://www.mit.edu/~kevdokim/ESMSM_sep16.pdf "Efficient Estimation with a Finite Number of Simulation Draws per Observation") shows how we can use an augmented set of moments related to the assumptions related to simulation to basically remove the factor $(1+S^{-1})$ on the asymptotic variance of the SMD estimator using only one(!) simulation, $S=1$!

# Model and Estimators
We use the same example as Kirill Evdokimov. Imagine the simple setup where we have the data-generating process (DGP):
$$
\begin{align}
Y_i &= \theta_0 + \varepsilon_i \\
\varepsilon_i &\sim N(0,1)
\end{align}
$$
**SMD:** We can use the moment function with only $S=1$ simulatin of $\varepsilon$ per individual
$$
g_i(\theta|w_i) = Y_i - \theta -\varepsilon_i
$$
to estimate $\theta$. We will call that $\hat{\theta}_{SMD}$. The moment vector would be
$$
g(\theta) =
\bigg( \begin{array}{c}
\overline{Y} - \theta -\overline{\varepsilon} \\
\end{array} \bigg)
$$
where $\overline{Y} = \frac{1}{N}\sum_{i=1}^{N} Y_i$ and $\overline{\varepsilon} = \frac{1}{N}\sum_{i=1}^{N} \varepsilon_i$.

**ES-SMD:** We can use the efficient SMD to augment the moment conditions with the fact that the simulated $\varepsilon$'s should have mean-zero and get the vector of moments in this augmented situation as
$$
g_{aug}(\theta) =
\bigg( \begin{array}{c}
\overline{Y} - \theta -\overline{\varepsilon} \\
0-\overline{\varepsilon} \\
\end{array} \bigg)
$$
where we use the optimal weighting matrix $W=\Omega^{-1}$ where
$$
\Omega = Var(g_{i,aug}(\theta|w_i)) =
\bigg( \begin{array}{cc}
2 & 1\\
1 & 1 \\
\end{array} \bigg)
$$
and
$$
\Omega^{-1} = \bigg( \begin{array}{cc}
1 & -1\\
-1 & 2 \\
\end{array} \bigg)
$$

We will call this estimator $\hat{\theta}_{ES-SMD}$.

**Asymptotic Variances:**
1. In the standard SMD estimator, the weighting matrix does not matter and we have
$$
\begin{align}
AVar(\hat{\theta}_{SMD}) &= Var(g_i(\theta|w_i)) \\
 &= Var(Y_i - \theta -\varepsilon_i)\\
 &= Var(Y_i) +Var(\varepsilon_i) \\
 &= 2
\end{align}
$$
2. In the augmented ES-SMD estmator, we have
$$
\begin{align}
AVar(\hat{\theta}_{ES-SMD}) &= Var((G'WG)^{-1}G'Wg_{i,aug}(\theta|w_i)) \\
 &= Var(-Y_i + \theta)\\
 &= 1
\end{align}
$$
bacause
$$
(G'WG)^{-1}G'Wg_{i,aug}(\theta|w_i) =  - (Y_i - \theta -\varepsilon) - \varepsilon.
$$
3. We thus have that the asymptotic variance of the ES-SMD estimator is lower that the SMD estimator!

We will now illustrate this result through a **Monte Carlo experiment** too!

# Setup

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

import numpy as np
from types import SimpleNamespace

import sys
sys.path.append('../')

from SimulatedMinimumDistance import SimulatedMinimumDistanceClass
```

# Model construction

```python
class ModelClass():
    def __init__(self,**kwargs):
        self.par = SimpleNamespace()
        self.sim = SimpleNamespace()

        self.par.theta = 0.5

        self.par.simN = 5000

        for key,val in kwargs.items():
            setattr(self.par,key,val)

    def solve(self,do_print=False): pass

    def simulate(self,seed=None,do_print=False):
        if seed is not None:
            np.random.seed(seed)
        self.sim.e = np.random.normal(size=self.par.simN)

        self.sim.Y = self.par.theta + self.sim.e
```

# Estimation choices

```python
# a. model settings
N = 100_000
N_sim = N

par = {'theta':0.2,'simN':N_sim}

par_true = par.copy()
par_true['simN'] = N

# b. parameters to estimate
est_par = {
    'theta': {'guess':0.5,'lower':0.0,'upper':1.0,},
}

# c. moment function used in estimation.
def mom_func(data,ids=None):
    """ returns the average Y """

    if ids is None:
        mean_Y = np.mean(data.Y)
    else:
        mean_Y = np.mean(data.Y[ids])

    return np.array([mean_Y]) # alwaus give a zero

# c. augmented moment function used in efficient estimation.
def mom_func_aug(data,ids=None):
    """ returns the average Y and the average of the simulations"""

    if ids is None:
        mean_Y_e = np.mean([data.Y,data.e],axis=1)
    else:
        mean_Y_e = np.mean([data.Y[ids],data[ids].e],axis=1)

    return mean_Y_e
```

# Monte Carlo Estimation results

```python
num_boot = 1_000

theta_est = np.empty(num_boot)
theta_est_aug = theta_est.copy()

model = ModelClass(**par)

for b in range(num_boot):
    # a. setup model to simulate data
    true = ModelClass(**par_true)
    true.simulate(seed=2050+b) # this seed is different from the default

    # b. data moments
    datamoms = mom_func(true.sim)
    datamoms_aug = np.array([datamoms[0],0.0])

    # c. setup estimators
    smd = SimulatedMinimumDistanceClass(est_par,mom_func,datamoms=datamoms)
    smd_aug = SimulatedMinimumDistanceClass(est_par,mom_func_aug,datamoms=datamoms_aug)

    # d. weighting matrix
    W = np.ones((datamoms.size,datamoms.size)) # does not matter here
    Omega = np.array([[2.0,1.0],[1.0,1.0]]) # covariance matrix of augmentet moments.
    W_aug = np.linalg.inv(Omega)

    # e. estimate the model (can take several minutes)
    est = smd.estimate(model,W,do_print_initial=False)
    est_aug = smd_aug.estimate(model,W_aug,do_print_initial=False)

    # f. store the estimates
    theta_est[b] = est['theta']
    theta_est_aug[b] = est_aug['theta']

```

```python
print(f'Variance, SMD:    {np.var(theta_est-par_true["theta"])*N:2.6f}')
print(f'Variance, ES-SMD: {np.var(theta_est_aug-par_true["theta"])*N:2.6f}')
```

    Variance, SMD:    2.036992
    Variance, ES-SMD: 1.096781

DynamicProgramming/extra/Efficient estimation with 1 simulation.ipynb


.. meta::
   :description: An implementation of the famous Particle Swarm Optimization (PSO) algorithm which is inspired by the behavior of the movement of particles represented by their position and velocity. Each particle is updated considering the cognitive and social behavior in a swarm.

.. meta::
   :keywords: Particle Swarm Optimization, Nature-inspired Algorithm, Single-objective Optimization, Python

# Particle Swarm Optimization (PSO)

Particle Swarm Optimization was proposed 1995 by Kennedy and Eberhart  based on the simulating of social behavior. The algorithm uses a *swarm* of particles to guide its search. Each particle has a velocity and is influenced by locally and globally best found solutions. Many different implementations have been proposed in the past and, therefore, it is quite difficult to refer to THE correct implementation of PSO. However, the general concepts shall be explained in the following.

Given the following variables:

- $X_{d}^{(i)}$ d-th coordinate of i-th particle's position
- $V_{d}^{(i)}$ d-th coordinate of i-th particle's velocity
- $\omega$ Inertia weight
- $P_{d}^{(i)}$ d-th coordinate of i-th particle's *personal* best
- $G_{d}^{(i)}$ d-th coordinate of the globally (sometimes also only locally) best solution found
- $c_1$ and $c_2$ Two weight values to balance exploiting the particle's best $P_{d}^{(i)}$ and swarm's best $G_{d}^{(i)}$
- $r_1$ and $r_2$ Two random values being create for the velocity update

The velocity update is given by:

\begin{equation}
V_{d}^{(i)} = \omega \, V_{d}^{(i)} \;+\; c_1 \, r_1 \, \left(P_{d}^{(i)} - X_{d}^{(i)}\right) \;+\; c_2 \, r_2 \, \left(G_{d}^{(i)} - X_{d}^{(i)}\right)
\end{equation}

The corresponding position value is then updated by:

\begin{equation}
X_{d}^{(i)} = X_{d}^{(i)} + V_{d}^{(i)}
\end{equation}

The social behavior is incorporated by using the *globally* (or locally) best found solution in the swarm for the velocity update. Besides the social behavior, the cognitive behavior of the swarm is determined by the particle's *personal* best solution found.
The cognitive and social components need to be well balanced to ensure the algorithm is performing well on a variety of optimization problems.
Thus, some effort has been made to determine suitable values for $c_1$ and $c_2$. In **pymoo** both values are updated as proposed in . Our implementation deviates in some implementation details (e.g. fuzzy state change) but follows the general principles proposed in the paper.

Assuming you have our third party library `pyrecorder` installed you can create an animation for a two-dimensional problem easily. Below we provide the code to observe the behavior of the swarm. We have reduced to set the maximum velocity to `max_velocity_rate=0.025` for illustration purposes. Otherwise, the algorithm converges even quicker to the global optimum of the `Ackley` function with two variables.

```python
from pymoo.algorithms.so_pso import PSO, PSOAnimation
from pymoo.factory import Ackley
from pymoo.optimize import minimize

problem = Ackley()

algorithm = PSO(max_velocity_rate=0.025)

res = minimize(problem,
               algorithm,
               callback=PSOAnimation(fname="pso.mp4"),
               seed=1,
               save_history=True,
               verbose=False)

print("Best solution found: \nX = %s\nF = %s" % (res.X, res.F))
```

In general, the PSO algorithm can be used by execute the following code. For the available parameters please see the API description below.

```python
from pymoo.algorithms.so_pso import PSO
from pymoo.factory import Rastrigin
from pymoo.optimize import minimize

problem = Rastrigin()

algorithm = PSO()

res = minimize(problem,
               algorithm,
               seed=1,
               verbose=False)

print("Best solution found: \nX = %s\nF = %s" % (res.X, res.F))
```

    Best solution found:
    X = [-3.07972689e-07  3.05136161e-07]
    F = [3.72857301e-11]

### API


<link rel="stylesheet" href="../../styles/theme_style.css">

<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/4.7.0/css/font-awesome.min.css">

         Parameter Extraction - Temporal/Statistical Parameters

        Difficulty Level:

                    Tags
                    extract|statistics|temporal signals

When we are working with a signal recording, it is important to perform an preliminar analysis to understand how it behaves and obtain its main characteristics. One of the most prevalent ways to do this is by extracting descriptive parameters. There are different ways to approach this, and this notebook will cover some of the most common statistical and temporal parameters that can be used to analyze time series.

The use of these descriptive charateristics will help us decide the most effective strategy to follow during data preparation, model selection and model tuning depending on the use we will give to the data.

**List of statistical parameters:**
+ Maximum, minimum and range
+ Mean
+ Mode
+ Median and quantiles
+ Variance and standard deviation
+ Skewness and kurtosis

**List of temporal parameters:**
+ Autocorrelation
+ Stationarity
+ Seasonal decomposition

1 - Importation of the needed packages

```python
# biosignalsnotebooks own package for loading and plotting the acquired data
import biosignalsnotebooks as bsnb

# Scientific packages
import numpy
import math
import scipy.integrate as integrate
import scipy.stats as stats
import statsmodels.tsa.stattools as stattools
import statsmodels.tsa.seasonal as seasonal
import pandas as pd
```

```python
# Base packages used in OpenSignals Tools Notebooks for plotting data
from bokeh.plotting import figure, output_file, show, curdoc
from bokeh.io import output_notebook
from bokeh.layouts import column
from bokeh.models import ColumnDataSource, Plot, LinearAxis, BoxAnnotation, Arrow, VeeHead, LinearAxis, Range1d
output_notebook(hide_banner=True)
```

2 - Load of sample signal data

```python
# Load of data
data, header = bsnb.load("../../signal_samples/ecg_20_sec_100_Hz.h5", get_header=True)
channel = list(data.keys())[0]
print(header)
```

    {'channels': array([1]), 'comments': '', 'date': b'2018-9-28', 'device': 'biosignalsplux', 'device connection': b'BTH00:07:80:D8:A7:F9', 'device name': b'00:07:80:D8:A7:F9', 'digital IO': array([0, 1]), 'firmware version': 773, 'resolution': array([16]), 'sampling rate': 100, 'sync interval': 2, 'time': b'14:36:39.886', 'sensor': [b'ECG'], 'column labels': {1: 'channel_1'}}

3 - Storage of sampling frequency and acquired data inside variables

Since in this case we know that the original signal is a ECG recording, we can also convert its units to mV following the method explained in the ECG Sensor - Unit Conversion  notebook.

```python
# Sampling frequency and acquired data
fs = header["sampling rate"]

# Signal Samples
signal_raw = data[channel]
time = numpy.linspace(0, len(signal_raw) / fs, len(signal_raw))

# Let's convert the signal's units, since we know it is a ECG signal
vcc = 3000 # mV
gain = 1000
ch = "CH1" # Channel
sr = header['sampling rate'] # Sampling rate
resolution = header['resolution'] # Resolution (number of available bits)
signal = (((numpy.array(signal_raw) / 2**resolution) - 0.5) * vcc) / gain

# Let's plot our raw signal
p = figure(plot_width=1000, plot_height=200)
p.ygrid.grid_line_alpha=0.5

# add a circle renderer with x and y coordinates, size, color, and alpha
p.line(time, signal)

show(p) # show the results
```

4 - Extraction of statistical parameters in a time signal

A fundamental task in many statistical analysis is to characterize the location (a central value that best describes the data) and variability (or spread) of a dataset, in order to describe and understand it better. These descriptors are independent of the domain of the signal, and can be used to describe any type of information or dataset. Due to their simplicity and the fact that they are easy to calculate, they are widely used during statistical analyses.

4.1 - Minimum, maximum and range values

The minimum and maximum values are, respectively, the lowest and highest power values across the signal duration. There may be more than one occurrence of the minimum and maximum value across the time series.

The range is the difference between the maximum and the minimum value.

```python
# Let's calculate the minimum, maximum and range of our signal
max_value = max(signal)
min_value = min(signal)
signal_range = max_value - min_value

print("Maximum: ", max_value, "\nMinimum: ", min_value, "\nRange (maximum - minimum): ", signal_range)
```

    Maximum:  1.10101318359375
    Minimum:  -0.170013427734375
    Range (maximum - minimum):  1.271026611328125

```python
# Let's draw this values in our signal
index_max = numpy.where(signal == max_value)[0]
index_min = numpy.where(signal == min_value)[0]

# Let's plot our raw signal
p_maxmin = figure(plot_width=1000, plot_height=200)
p_maxmin.ygrid.grid_line_alpha=0.5

p_maxmin.line(time, signal)
p_maxmin.circle([time[i] for i in index_max],[max_value for i in range(len(index_max))], size=15, line_color="navy", fill_color="orange", fill_alpha=0.5, legend_label='Maximum')
p_maxmin.circle([time[i] for i in index_min],[min_value for i in range(len(index_min))], size=15, line_color="orange", fill_color="navy", fill_alpha=0.5, legend_label='Minimum')

show(p_maxmin)
```

4.2 - Mean

The arithmetic mean ($\bar{x}$) of a temporal data set represents the average power value across the signal duration.

\begin{align}
\bar{x}=\frac{1}{n}\sum_{i=1}^n x_i=\frac{x_1+x_2+\cdots+x_n}{n} \\
\end{align}
With $x_1, x_2, \ldots, x_n$ being the values in the dataset and $n$ the number of elements in the dataset.

```python
# Let's calculate the mean value of our signal
mean = numpy.mean(signal)

print('Mean:', mean)
```

    Mean: 0.00021564978679627862

```python
# Let's draw the mean line in our plot
p.line(time, mean, line_alpha=0.5, color="firebrick", line_width=3, legend_label='Mean')

show(p)
```

4.3 - Mode

The mode represents the power value with the highest number of occurrences in the dataset across the signal duration.

There may be more than one mode, if two or more different values have the same number of occurrences.

```python
# Let's find the mode (or modes) in our signal
mode, occurrences = stats.mode(signal)

print('Mode: ', mode, '. Occurrences: ', occurrences)
```

    Mode:  [-0.04655457] . Occurrences:  [9]

```python
# Let's print the mode values in a separate new graph
index_mode = numpy.where(signal == mode)

p_modes = figure(plot_width=1000, plot_height=200)
p_modes.line(time, signal)
p_modes.ygrid.grid_line_alpha=0.5

index_aux = 0
for i in index_mode:
    p_modes.line(time, mode[index_aux], color="green", line_width=3, line_alpha=0.5) # we can draw the line marking the mode value in the y-axis
    p_modes.circle([time[idx] for idx in i], mode[index_aux],
             size=10, line_color="green", fill_color="green", fill_alpha=0.5) # we can also mark the points where the function has the mode value
    index_aux += 1

show(p_modes)
```

4.4 - Median value and quantiles

The median is the value that separates the top half from the lower half in a dataset. It is a popular summary statistic because it is simple to calculate and gives a measure that is more robust than the mean in the presence of outlier values.

[//]: <> (\begin{align})
[//]: <> (\mathrm{median}(a) = \frac{a_{\lfloor(\#x+1) \div 2\rfloor} + a_{\lceil(\#x+1) \div 2\rceil}}{2} \\)
[//]: <> (\end{align})

The median is also known as the 2-quantile because it marks the middle point of the dataset. However, other quantiles  can also be explored.

For example, 4-quantiles (or quartiles) divide the dataset in four sections:
- the first quartile (Q1), also known as the 25th percentile, splits off the lowest 25% of data from the highest 75%.
- the second quartile (Q2), also known as the 50th percentile, is the same as the median.
- the third quartile (Q3), also known as the 75th percentile splits off the lowest 75% of data from the highest 25%.

```python
# Let's calculate the median, 25 and 75 quartiles
median = numpy.median(signal) # 50% of the values of the signal will be lower than the median or 50th quantile.
Q1 = numpy.quantile(signal, 0.25) # 25% of the values of the signal will be lower than the 25th quantile
Q3 = numpy.quantile(signal, 0.75) # 75% of the values of the signal will be lower than the 75th quantile

print('Mean: ', mean)
print('| Q1 |\t\t\t| Q2 (Median) |\t\t\t| Q3 |\n', Q1, '\t', median, '\t\t', Q3)
```

    Mean:  0.00021564978679627862
    | Q1 |			| Q2 (Median) |			| Q3 |
     -0.0703582763671875 	 -0.0411529541015625 		 0.026275634765625

```python
# Let's see them in a separated new graph to avoid cluttering the main one
p_quartiles = figure(plot_width=1000, plot_height=200)
p_quartiles.line(time, signal)
p_quartiles.ygrid.grid_line_alpha=0.5

# Redrawing the mean line...
p_quartiles.line(time, mean, color="red", line_width=2, legend_label='Mean')

p_quartiles.line(time, Q3, color="yellow", line_width=2, line_dash='dashed', legend_label='Q3')
p_quartiles.line(time, median, color="deeppink", line_width=2, line_dash='dashed', legend_label='Q2 (Median)')
p_quartiles.line(time, Q1, color="cyan", line_width=2, line_dash='dashed', legend_label='Q1')

p_quartiles.legend.location = "top_left"
p_quartiles.y_range = Range1d(-0.15, 0.15)
p_quartiles.x_range = Range1d(-1, 6)

show(p_quartiles)
```

In this example we can observe that the distance between the Q3 and Q2 is bigger than the distance between Q2 and Q1.
Quartiles give us and idea about how the signal's values are distributed.

Even though the signal has extreme values (peaks), their frequency of occurrence is low, and that is why the distance between the Q3 line and the top of the graph is much bigger than the distance between any of the other quartiles.

4.5 - Standard deviation and variance

Both the standard deviation and the variance are statistics determined by using the mean, and used to measure the degree of dispersion in a dataset. Informally, they measure how close a group of numbers tend to be to the mean.

The variance ($\sigma^2$) is defined as the average of the squared differences from the mean of the dataset. It is obtained by calculating the difference between each point of the dataset and the mean, squaring and averaging the results. Squaring the distance from the mean gives greater weight to values that are further from the mean.
However, because of the squaring, the variance is no longer in the same measurement units as the dataset. This makes measuring variance a bit difficult, so many times the standard deviation is used instead.

The standard deviation ($\sigma$) is the square root of the variation. A low standard deviation value indicates that the values tend to be more concentrated and closer to the mean, whilst a high value indicates that the values tend to spread out over a wider range.

\begin{align}
\sigma = \sqrt{\frac{1}{N-1}\sum_{i=1}^N (x_i - \bar{x})^2 }, \\
\end{align}

```python
# Let's calculate the standard deviation and the variance
std = numpy.std(signal)
variance = numpy.var(signal)

print('Mean: ', mean)
print('Standard deviation: ', std, '\nVariance: ', variance)
```

    Mean:  0.00021564978679627862
    Standard deviation:  0.16587464942332372
    Variance:  0.02751439932131055

```python
# Let's draw the std line in our graph
p.line(x=time, y=mean+std, line_color='orange', line_dash='dashed', line_width=2, legend_label='Standard deviation')
p.line(x=time, y=mean-std, line_color='orange', line_dash='dashed', line_width=2)

show(p)
```

4.6 - Skewness and kurtosis

The skewness is a metric of symmetry (or the lack of) in a dataset. A dataset is symmetric if it looks the same to the left and right of the center point. A normal distribution has a skewness of 0.

- Negative (<0) skew value: longer left tail; distribution concentrates on the right.
- Positive (>=0) skew value: longer right tail; distribution concentrates on the left.

Kurtosis indicates if the dataset is heavy-tailed or light-tailed in comparison to a normal distribution. Datasets with high kurtosis ('heavy-tailed') tend to have outliers. The lower the kurtosis value, the less common it is to find outliers in a dataset.

```python
# Let's calculate the skewness and kurtosis
skewness = stats.skew(signal)
kurtosis = stats.kurtosis(signal)

print('Skewness: ', skewness, '\nKurtosis: ', kurtosis)
```

    Skewness:  4.271322272366648
    Kurtosis:  21.250186573425474

```python
arr_hist, edges = numpy.histogram(signal, bins = int(180/5))

# Put the information in a dataframe
delays = pd.DataFrame({'arr_delay': arr_hist, 'left': edges[:-1], 'right': edges[1:]})

p_histogram = figure(plot_height = 250, plot_width = 1000, x_axis_label="mV", y_axis_label="bins")

# Add a quad glyph
p_histogram.quad(bottom=0, top=delays['arr_delay'],
       left=delays['left'], right=delays['right'],
       fill_color='red', line_color='black')

# Show the plot
show(p_histogram)
```

The histogram shows clearly that the time series is right tailed, which is consistent with the positive skew value obtained.
Some outliers are also present on the right side of the histogram, around the 1mv value. This is consistent with the kurtosis result, because as we previously mentioned in the analysis of the quartiles, the highest peaks in the heartbeat cycle are, althought with a low frequency, common across the time signal.

5 - Temporal parameters

Parameters for analyzing time series and extract meaningful statistics or other characteristics of the signal.
These are parameters specific for data series located in the time domain. They extract relevant and descriptive information about a dataset in virtue of its temporal nature, which should be taken into account when analyzing temporal series.

5.1 - Autocorrelation

Autocorrelation, or serial correlation, is often used in time domain signals. It is the relationship of a signal with a delayed copy of itself as a function of delay. The analysis of autocorrelation finds repeating patterns, such as the presence of a periodic signal obscured by noise.

The ACF (autocorrelation function) shows the correlation between points separated by various time lags. So, in other words, it is the degree of association between points based on how many time steps apart they are. Normally, the autocorrelation function falls towards 0 as points become more separated, because the bigger the separation, the less correlation between the points. This is not a rule, but it is the most typical scenario.

```python
# Let's obtain the autocorrelation function for our signal
autocorrelation_function = stattools.acf(signal, fft=True, nlags=round(len(signal)*0.75))
```

```python
def get_autocorrelation_plot_params(series):
    n = len(series)
    data = numpy.asarray(series)
    mean = numpy.mean(data)
    c0 = numpy.sum((data - mean) ** 2) / float(n)

    def r(h):
        return ((data[:n - h] - mean) *
                (data[h:] - mean)).sum() / float(n) / c0
    x = numpy.arange(n) + 1
    y = list(map(r, x))
    z95 = 1.959963984540054 # confidence interval 95%
    z99 = 2.5758293035489004 # confidence interval 99%
    return n, x, y, z95, z99

n, x, y, z95, z99 = get_autocorrelation_plot_params(autocorrelation_function)
x = x/100
auto_correlation_plot2 = figure(title='Time Series Auto-Correlation', plot_width=1000,
                                plot_height=400, x_axis_label="Lag (s)", y_axis_label="Autocorrelation")

auto_correlation_plot2.line(x, y=z99 / numpy.sqrt(n), line_dash='dashed', line_color='red', legend_label='99% confidence band')
auto_correlation_plot2.line(x, y=z95 / numpy.sqrt(n), line_color='red', legend_label='95% confidence band')
auto_correlation_plot2.line(x, y=0.0, line_color='black')
auto_correlation_plot2.line(x, y=-z95 / numpy.sqrt(n), line_color='red')
auto_correlation_plot2.line(x, y=-z99 / numpy.sqrt(n), line_dash='dashed', line_color='red')

auto_correlation_plot2.line(x, y, line_width=1)
#auto_correlation_plot2.circle(x, y, fill_color="white", size=8)  # optional

curdoc().add_root(column(auto_correlation_plot2))

show(auto_correlation_plot2)
```

As expected, the ACF tends to 0 the longer the distance between the points. However the analysis highlights the presence of a periodic pattern, with high correlation, characterized by the sucesion of a negative, positive and negative peak with a period of less than 1 second.

In this case, since this signal is a ECG recording, it makes sense that there is a periodicy associated with the heartbeat.

5.2 - Stationarity

A time series is stationary when its statistic properties, such as the mean and variance, do not change over time.

The stationarity of a dataset can usually be predicted by looking at the signal's plot, histogram or ACF. It can also be checked by calculating the mean and variance of the signal in different time intervals; if the results are similar then the signal is probably stationary. There are also statistical tests to determine if a time series is stationary, for instance the Dickey-Fuller test or the KPSS (Kwiatkowski-Phillips-Schmidt-Shin) test.

```python
# Let's try to check the stationarity of our signal
print('Length of the signal: ', len(signal))
# We can measure the mean and variance of different intervals of the signal and see if they are different
mean_1 = numpy.mean(signal[1:300])
mean_2 = numpy.mean(signal[800:1100])
mean_3 = numpy.mean(signal[1600:1900])

var_1 = numpy.var(signal[1:300])
var_2 = numpy.var(signal[800:1100])
var_3 = numpy.var(signal[1600:1900])

print('Means: ', mean_1, mean_2, mean_3)
print('Variances: ', var_1, var_2, var_3, '\n')

# Let's try the Dickey-Fuller test
from statsmodels.tsa.stattools import adfuller

adf_test = adfuller(signal)

# Let's try the KPSS test
from statsmodels.tsa.stattools import kpss

kpsstest = kpss(signal, regression='c')
```

    Length of the signal:  1965
    Means:  -0.0026015980187865805 0.006810760498046875 -0.00485137939453125
    Variances:  0.0260274864826267 0.03005196664656978 0.025713346455339344

    c:\users\gui_s\appdata\local\programs\python\python37-32\lib\site-packages\statsmodels\tsa\stattools.py:1685: FutureWarning: The behavior of using lags=None will change in the next release. Currently lags=None is the same as lags='legacy', and so a sample-size lag length is used. After the next release, the default will change to be the same as lags='auto' which uses an automatic lag length selection method. To silence this warning, either use 'auto' or 'legacy'
      warn(msg, FutureWarning)
    c:\users\gui_s\appdata\local\programs\python\python37-32\lib\site-packages\statsmodels\tsa\stattools.py:1710: InterpolationWarning: p-value is greater than the indicated p-value
      warn("p-value is greater than the indicated p-value", InterpolationWarning)

```python
# Using pandas Series, we can print the ADF results in a readable way.
print ('Results of ADF Test:')
dfoutput = pd.Series(adf_test[0:4], index=['Test Statistic','p-value','#Lags Used','Number of Observations Used'])
for key,value in adf_test[4].items():
   dfoutput['Critical Value (%s)'%key] = value
print (dfoutput)

# Using pandas Series, we can print the KPSS results in a readable way.
print ('\nResults of KPSS Test:')
kpss_output = pd.Series(kpsstest[0:3], index=['Test Statistic','p-value','Lags Used'])
for key,value in kpsstest[3].items():
    kpss_output['Critical Value (%s)'%key] = value
print (kpss_output)
```

    Results of ADF Test:
    Test Statistic                -9.868093e+00
    p-value                        4.063043e-17
    #Lags Used                     2.600000e+01
    Number of Observations Used    1.938000e+03
    Critical Value (1%)           -3.433729e+00
    Critical Value (5%)           -2.863033e+00
    Critical Value (10%)          -2.567565e+00
    dtype: float64

    Results of KPSS Test:
    Test Statistic            0.014125
    p-value                   0.100000
    Lags Used                26.000000
    Critical Value (10%)      0.347000
    Critical Value (5%)       0.463000
    Critical Value (2.5%)     0.574000
    Critical Value (1%)       0.739000
    dtype: float64

Although the means remain pretty similar in the three chosen intervals, the variances do show some significant changes.

Let's analyze the results of the statistical tests:

- The Dickey-Fuller test returned a very small p-value. Since this value is much smaller than 0.05 (95% confidence interval) and 0.01 (99% confidence interval), then we can reject the null hypothesis of non-stationarity. This test deems the signal as stationary.
- In the case of the KPSS test, the results have to be interpreted in the opposite way: a high p-value would point towards a stationary signal. In our case, however, the p-value is 0.1, which means that the signal may be non-stationary.

So we got opposite results in each test. What does this mean?
In reality, there is more than one type of non-stationarity. As a short introduction, there are three possible types of signals based in stationarity:

- Strict stationary: signals where the mean, variance and covariance do not vary with time. The aim is to convert non-stationary series into strict stationary series for making predictions.
- Trend stationary: a time series that can be made strict stationary by removing its trend. KPSS is a trend stationarity test.
- Difference stationary: a time series that can be made strict stationary by differencing. ADF is a difference stationarity test.

Therefore, the results of the KPSS and ADF tests are useful to test wether a signal es stationary, trend stationary or difference stationary:

- If both tests conclude that the series is not stationary: the series is not stationary.
- If both tests conclude that the series is stationary: the series is stationary.
- If KPSS = stationary and ADF = not stationary: the signal is trend stationary.
- If KPSS = not stationary and ADF = stationary: the signal difference stationary.

So in this case our signal would fall under the difference stationarity definition. To understand more about signal stationarity, there are some interesting online resources  that go deeper on how to transform non-stationary signals and different ways to handle them.

5.3 - Seasonality decomposition

The seasonality decomposition of a signal allows dividing the main time series in different components, each of which has a specific characteristic. This decomposition can be done following two different models:

- additive, for linear distributions, where the time signal is decomposed following the formula: $Signal=trend+seasonality+residue$
- multiplicative, for exponential series, where the time signal is decomposed following the formula: $Signal=trend*seasonality*residue$

The resulting components are the following:
- Trend: increasing or decreasing tendency in the time series.
- Seasonality: the repeating short-term cycles in the series.
- Residue: random noise present in the signal that cannot be attributed to trend or seasonality

```python
# Let's decompose our signal into trend, seasonal and residual components using seasonal decomposition
signal = numpy.array(signal) - numpy.mean(signal)
filter_signal_1 = bsnb.lowpass(signal, f=40, order=1, fs=sr)
filter_signal_2 = bsnb.lowpass(filter_signal_1, f=40, order=3, fs=sr)
results = seasonal.seasonal_decompose(filter_signal_2, model='additive', freq=sr, extrapolate_trend='freq')
```

```python
decomposition_plot2 = figure(title='Seasonal decomposition - Observed and trend', plot_width=1200, plot_height=200)
decomposition_plot2.line(numpy.arange(0,len(results.observed)), results.observed, legend_label='Observed')
decomposition_plot2.line(numpy.arange(0,len(results.trend)), results.trend, line_color='red', legend_label='Trend')
show(decomposition_plot2)

decomposition_plot3 = figure(title='Seasonal decomposition - Seasonality and residuals', plot_width=1200, plot_height=200)
decomposition_plot3.line(numpy.arange(0,len(results.seasonal)), results.seasonal, line_color='orange', legend_label='Seasonality')
decomposition_plot3.line(numpy.arange(0,len(results.resid)), results.resid, line_color='green', legend_label='Residuals')
show(decomposition_plot3)
```

The decomposition in this case is unable to separate the cyclic element (the heartbeat) from the signal. That is why the most significant part of the signal ends up as residues. This is not a good result, and it is probably caused by the naive approach of the seasonal_decompose function in the statsmodels package.

There are more complex and robust approaches to this problem, like Loess or STL decomposition, which would however require a more in depth explanation.

**Auxiliary Code Segment (should not be replicated by the user)**

```python
bsnb.css_style_apply()
```

    .................... CSS Style Applied to Jupyter Notebook .........................


# Non-linear Gaussian filtering and smoothing

Provided are two examples of nonlinear state-space models on which one can perform Bayesian filtering and smoothing in order to obtain
a posterior distribution over a latent state trajectory based on noisy observations.
In order to understand the theory behind these methods in detail we refer to [1] and [2].

We provide examples for two different types of state-space model:
1. [Nonlinear, Discrete State-Space Model](#1.-Non-Linear-Discrete-State-Space-Model:-Pendulum): Pendulum
2. [Nonlinear, Continuous-Discrete State-Space Model](#2.-Non-Linear-Continuous-Discrete-State-Space-Model:-Benes-Daum-Filter): Benes-Daum Filter

In order to perform Bayesian Filtering and Smoothing on non-linear models, we show how to use functionalities provided by ProbNum to
either linearise the respective model (Extended Kalman Filter, EKF) or to use an Unscented Transform (Unscented Kalman Filter, UKF) which does not require a Jacobian of the nonlinear model.

**References**:
> [1] Särkkä, Simo, and Solin, Arno. Applied Stochastic Differential Equations. Cambridge University Press, 2019.
>
> [2] Särkkä, Simo. Bayesian Filtering and Smoothing. Cambridge University Press, 2013.

```python
import numpy as np

import probnum as pn
from probnum import filtsmooth, randvars, statespace, randprocs
from probnum.problems import TimeSeriesRegressionProblem
```

```python
rng = np.random.default_rng(seed=123)
```

```python
# Make inline plots vector graphics instead of raster graphics
%matplotlib inline
from IPython.display import set_matplotlib_formats

set_matplotlib_formats("pdf", "svg")

# Plotting
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

plt.style.use("../../probnum.mplstyle")
```

## 1. **Non-Linear Discrete** State-Space Model: Pendulum

---

For nonlinear models, we assume that at least one of the components of our state-space model (i.e. either the dynamics, the measurement model, or both) is a nonlinear function of the latent state.

Consider nonlinear functions $f: \mathbb{R}^d \rightarrow \mathbb{R}^d$ and $h: \mathbb{R}^d \rightarrow \mathbb{R}^m$ where $d$ is the dimension of the latent state and
$m$ is the dimension of the measurements.

For $k = 1, \dots, K$ and $x_0 \sim \mathcal{N}(\mu_0, \Sigma_0)$:

$$
\begin{align}
    \boldsymbol{x}_k &\sim \mathcal{N}(f(\boldsymbol{x}_{k-1}), \boldsymbol{Q}) \\
    \boldsymbol{y}_k &\sim \mathcal{N}(h(\boldsymbol{x}_k), \boldsymbol{R})
\end{align}
$$

This defines a dynamics model that assumes a state $x_k$ in a **discrete** sequence of states arising from a nonlinear map of the previous state $x_{k-1}$ corrupted with additive Gaussian noise under a **process noise** covariance matrix $Q$.
Similarly, the measurements $y_k$ are assumed to be nonlinear transformations of the latent state under additive Gaussian noise according to a **measurement noise** covariance $R$.
In the described case, we cannot use the Kalman Filtering equations since these assume linear relationships. One method is to linearise the nonlinear parts of the model by a first-order Taylor Expansion (Extended Kalman Filter, EKF) or employing the Unscented Kalman Filter (UKF) which is based on the unscented transform and does not require the computation of a Jacobian.
The `probnum` packages provides both methods by wrapping linearizing "Components" around nonlinear dynamcis and/or measurement models, as we will see below.

Note that this can be generalized to a nonlinear time-varying state-space model, as well. Then we would have $f: \mathbb{T} \times \mathbb{R}^d \rightarrow \mathbb{R}^d$ and $h: \mathbb{T} \times \mathbb{R}^d \rightarrow \mathbb{R}^m$ where $\mathbb{T}$ is the "time dimension".

### Define State-Space Model

#### I. Discrete Dynamics Model: Nonlinear Gaussian Transitions

```python
state_dim = 2
observation_dim = 1
```

```python
# approx. gravitational constant
g = 9.81
# dt
delta_t = 0.0075

def pendulum_rhs(state):
    """Right-hand side of an ODE that defines pendulum dynamics"""
    x1, x2 = state
    y1 = x1 + x2 * delta_t
    y2 = x2 - g * np.sin(x1) * delta_t
    return np.array([y1, y2])

def pendulum_jacobian(state):
    """Jacobian of the pendulum ODE"""
    x1, x2 = state
    dy1_dx = [1.0, delta_t]
    dy2_dx = [-g * np.cos(x1) * delta_t, 1.0]
    return np.array([dy1_dx, dy2_dx])

dynamics_transition_function = lambda t, state: pendulum_rhs(state)
dynamics_transition_jacobian_function = lambda t, state: pendulum_jacobian(state)

dynamics_diffusion_matrix = 1.0 * (
    np.diag(np.array([delta_t ** 3 / 3, delta_t]))
    + np.diag(np.array([delta_t ** 2 / 2]), 1)
    + np.diag(np.array([delta_t ** 2 / 2]), -1)
)
```

We here define the non-linear dynamics of a pendulum. Note (in the cell directly above) that `probnum` expects `dynamicsfun` to be a function of time and state. Even though we here consider a nonlinear but time-**in**variant model, we still have to fulfill the contract that the interface requires.
```python
dynamics_transition_function = lambda t, state: pendulum_rhs(state)
```
In other words, we simply ignore the time variable `t` in case of time-invariant models.

To create a discrete, (non-linear) Gaussian dynamics model, `probnum` provides the `DiscreteGaussian` class that takes
- `input_dim` and `output_dim`
- `state_trans_fun` : the function that returns the next state given the current one
- `proc_noise_cov_mat_fun` : the function that returns the covariance matrix of the Gaussian process noise
- `jacob_state_trans_fun` : the function that returns the Jacobian matrix of the state transition function

```python
# Create discrete, non-linear Gaussian dynamics model
dynamics_model = statespace.DiscreteGaussian(
    input_dim=state_dim,
    output_dim=state_dim,
    state_trans_fun=dynamics_transition_function,
    proc_noise_cov_mat_fun=lambda t: dynamics_diffusion_matrix,
    jacob_state_trans_fun=dynamics_transition_jacobian_function,
)
```

#### II. Discrete Measurement Model: Nonlinear Gaussian Measurements

```python
def pendulum_measurement(state):
    x1, x2 = state
    return np.array([np.sin(x1)])

def pendulum_measurement_jacobian(state):
    x1, x2 = state
    return np.array([[np.cos(x1), 0.0]])

measurement_function = lambda t, state: pendulum_measurement(state)
measurement_jacobian_function = lambda t, state: pendulum_measurement_jacobian(state)

measurement_variance = 0.32 ** 2
measurement_covariance = measurement_variance * np.eye(observation_dim)
```

```python
# Create discrete, non-linear Gaussian measurement model
measurement_model = statespace.DiscreteGaussian(
    input_dim=state_dim,
    output_dim=observation_dim,
    state_trans_fun=measurement_function,
    proc_noise_cov_mat_fun=lambda t: measurement_covariance,
    jacob_state_trans_fun=measurement_jacobian_function,
)
```

#### III. Initial State Random Variable

```python
mu_0 = np.ones(state_dim)
sigma_0 = measurement_variance * np.eye(state_dim)
initial_state_rv = randvars.Normal(mean=mu_0, cov=sigma_0)
```

### Generate Data for the State-Space Model

`statespace.generate_samples()` is used to sample both latent states and noisy observations from the specified state space model.

```python
time_grid = np.arange(0.0, 5.0, step=delta_t)
```

```python
latent_states, observations = statespace.generate_samples(
    rng=rng,
    dynmod=dynamics_model,
    measmod=measurement_model,
    initrv=initial_state_rv,
    times=time_grid,
)
```

### Kalman Filtering

#### I. Linearize model (Extended Kalman Filter)

Since we could easily derive the Jacobian of our nonlinear model components above, we can use the Extended Kalman Filter (EKF) to obtain a filtering posterior over the states. Therefore, in `probnum` we just wrap the dynamics model and the measurement model in a `DiscreteEKFComponent` object.
In cases in which the Jacobian is not as easy to write down, `probnum` also provides the `DiscreteUKFComponent` class that implements the linearization by means of an
_unscented Kalman filter_ that does not require a Jacobian.
Alternatively, the Jacobian can be calculated via automatic differentiation.

```python
linearised_dynamics_model = filtsmooth.DiscreteEKFComponent(dynamics_model)
regression_problem = TimeSeriesRegressionProblem(
    observations=observations,
    locations=time_grid,
    measurement_models=[filtsmooth.DiscreteEKFComponent(measurement_model)]
    * len(time_grid),
)
```

```python
prior_process = randprocs.MarkovProcess(
    transition=linearised_dynamics_model, initrv=initial_state_rv, initarg=0.0
)
```

#### II. Kalman Filter

Then, we simply pass these linearized models to the Kalman Filter (exactly the same interface as in the linear case) and proceed.

```python
kalman_filter = filtsmooth.Kalman(prior_process)
```

#### III. Perform Kalman Filtering + Rauch-Tung-Striebel Smoothing

```python
state_posterior, _ = kalman_filter.filtsmooth(regression_problem)
```

The method `filtsmooth` returns a `KalmanPosterior` object which provides convenience functions for e.g. sampling and interpolation.
We can also extract the just computed posterior smoothing state variables by querying the `.state_rvs` property.
This yields a list of Gaussian Random Variables from which we can extract the statistics in order to visualize them.

```python
posterior_state_rvs = (
    state_posterior.states
)  # List of <num_time_points> Normal Random Variables
posterior_state_means = posterior_state_rvs.mean  # Shape: (num_time_points, state_dim)
posterior_state_covs = (
    posterior_state_rvs.cov
)  # Shape: (num_time_points, state_dim, state_dim)
```

### Visualize Results

```python
state_fig = plt.figure()
state_fig_gs = gridspec.GridSpec(ncols=2, nrows=1, figure=state_fig)

ax_00 = state_fig.add_subplot(state_fig_gs[0, 0])
ax_01 = state_fig.add_subplot(state_fig_gs[0, 1])

# Plot means
mu_x_1, mu_x_2 = [posterior_state_means[:, i] for i in range(state_dim)]

ax_00.plot(time_grid, mu_x_1, label="posterior mean")
ax_01.plot(time_grid, mu_x_2)

# Plot marginal standard deviations
std_x_1, std_x_2 = [np.sqrt(posterior_state_covs[:, i, i]) for i in range(state_dim)]

ax_00.fill_between(
    time_grid,
    mu_x_1 - 1.96 * std_x_1,
    mu_x_1 + 1.96 * std_x_1,
    alpha=0.2,
    label="1.96 marginal stddev",
)
ax_01.fill_between(
    time_grid, mu_x_2 - 1.96 * std_x_2, mu_x_2 + 1.96 * std_x_2, alpha=0.2
)

# Plot 5 samples on a subsampled grid
samples = state_posterior.sample(rng=rng, t=state_posterior.locations[::5], size=5)
for smp in samples:
    ax_00.plot(
        state_posterior.locations[::5],
        smp[:, 0],
        color="gray",
        alpha=0.75,
        linewidth=1,
        linestyle="dashed",
        label="sample",
    )
    ax_01.plot(
        state_posterior.locations[::5],
        smp[:, 1],
        color="gray",
        alpha=0.75,
        linewidth=1,
        linestyle="dashed",
        label="sample",
    )

# Plot groundtruth
ax_00.scatter(time_grid, observations, marker=".", label="measurements")

# Add labels etc.
ax_00.set_xlabel("t")
ax_01.set_xlabel("t")

ax_00.set_title(r"$x_1$")
ax_01.set_title(r"$x_2$")

# The following two lines are just to remove duplicate labels (caused by samples) from the legend
handles, labels = ax_00.get_legend_handles_labels()
by_label = dict(zip(labels, handles))

state_fig.legend(
    by_label.values(), by_label.keys(), loc="center left", bbox_to_anchor=(1, 0.5)
)

state_fig.tight_layout()
```

## 2. **Non-Linear Continuous-Discrete** State-Space Model: Benes-Daum Filter

---

Now, consider we assume **continuous** dynamics. We assume that there is a continuous process that defines the nonlinear dynamics of our latent space from which we collect discrete nonlinear-Gaussian measurements (as above). Only the dynamics model changes. In particular, we formulate the dynamics as a stochastic process in terms of a Nonlinear Time-Invariant Stochastic Differential Equation. We refer to [1] for more details.
Again, consider non-linear functions $f: \mathbb{R}^d \rightarrow \mathbb{R}^d$ and $h: \mathbb{R}^d \rightarrow \mathbb{R}^m$ where $d$ is the dimension of the latent state and
$m$ is the dimension of the measurements.
We define the following nonlinear **continuous-discrete** state-space model:

Let $x(t_0) \sim \mathcal{N}(\mu_0, \Sigma_0)$.

$$
\begin{align}
    d\boldsymbol{x} &= f(\boldsymbol{x}) \, dt + \boldsymbol{L} \, d \boldsymbol{\omega} \\
    \boldsymbol{y}_k &\sim \mathcal{N}(h(\boldsymbol{x}(t_k)), \boldsymbol{R}), \qquad k = 1, \dots, K
\end{align}
$$

where $\boldsymbol{\omega}$ denotes a vector of driving forces (often Brownian Motion).

Note that this can be generalized to a nonlinear time-varying state-space model, as well. Then we would have $f: \mathbb{T} \times \mathbb{R}^d \rightarrow \mathbb{R}^d$ and $h: \mathbb{T} \times \mathbb{R}^d \rightarrow \mathbb{R}^m$ where $\mathbb{T}$ is the "time dimension".

### Define State-Space Model

#### I. Continuous Dynamics Model: Nonlinear Stochastic Differential Equation

```python
state_dim = 1
observation_dim = 1
```

```python
drift_function = lambda t, state: np.tanh(state)
drift_function_jacobian = lambda t, state: 1.0 - np.tanh(state) ** 2

dispersion_matrix_function = lambda t: np.ones(state_dim).reshape(1, 1)
```

For continous, non-linear dynamics, `probnum` provides a general stochastic differential equation (SDE) interface that takes
- `driftfun`: a (possibly non-linear) drift function,
- `dispmatfun`: a (possibly non-linear) dispersion function,
- `jacobfun`: the Jacobian of the drift function.

```python
# Create continuous, non-linear SDE dynamics model
dynamics_model = statespace.SDE(
    dimension=state_dim,
    driftfun=drift_function,
    dispmatfun=dispersion_matrix_function,
    jacobfun=drift_function_jacobian,
)
```

#### II. Discrete Measurement Model: Linear, Time-Invariant, Gaussian Measurements

```python
measurement_marginal_variance = 1.0
measurement_matrix = np.eye(observation_dim, state_dim)
measurement_noise_matrix = measurement_marginal_variance * np.eye(observation_dim)
```

```python
measurement_model = statespace.DiscreteLTIGaussian(
    state_trans_mat=measurement_matrix,
    shift_vec=np.zeros(observation_dim),
    proc_noise_cov_mat=measurement_noise_matrix,
)
```

#### III. Initial State Random Variable

```python
mu_0 = np.zeros(state_dim)
sigma_0 = 3.0 * np.eye(state_dim)
initial_state_rv = randvars.Normal(mean=mu_0, cov=sigma_0)
```

### Generate Data for the State-Space Model

`statespace.generate_samples()` is used to sample both latent states and noisy observations from the specified state space model.

```python
time_grid = np.arange(0.0, 10.0, step=10 * delta_t)
```

```python
linearised_dynamics_model = filtsmooth.ContinuousEKFComponent(dynamics_model)
```

```python
latent_states, observations = statespace.generate_samples(
    rng=rng,
    dynmod=linearised_dynamics_model,
    measmod=measurement_model,
    initrv=initial_state_rv,
    times=time_grid,
)
```

```python
regression_problem = TimeSeriesRegressionProblem(
    observations=observations,
    locations=time_grid,
    measurement_models=[measurement_model] * len(time_grid),
)
```

### Kalman Filtering

#### I. Linearize model (Extended Kalman Filter)

#### II. Kalman Filter

```python
prior_process = randprocs.MarkovProcess(
    transition=linearised_dynamics_model, initrv=initial_state_rv, initarg=0.0
)
kalman_filter = filtsmooth.Kalman(prior_process)
```

#### III. Perform Kalman Filtering + Rauch-Tung-Striebel Smoothing

```python
state_posterior, _ = kalman_filter.filter(regression_problem)
```

```python
smoothed_state_posterior = kalman_filter.smooth(filter_posterior=state_posterior)
```

The method `filtsmooth` returns a `KalmanPosterior` object which provides convenience functions for e.g. sampling and prediction.
We can also extract the just computed posterior smoothing state variables by querying the `.state_rvs` property.
This yields a list of Gaussian Random Variables from which we can extract the statistics in order to visualize them.

```python
grid = np.linspace(0, 11, 500)

posterior_state_rvs = state_posterior(
    grid
)  # List of <num_time_points> Normal Random Variables
posterior_state_means = (
    posterior_state_rvs.mean.squeeze()
)  # Shape: (num_time_points, state_dim)
posterior_state_covs = (
    posterior_state_rvs.cov
)  # Shape: (num_time_points, state_dim, state_dim)
```

```python
smoothed_posterior_state_rvs = (
    smoothed_state_posterior.states
)  # List of <num_time_points> Normal Random Variables
smoothed_posterior_state_means = (
    smoothed_posterior_state_rvs.mean.squeeze()
)  # Shape: (num_time_points, state_dim)
smoothed_posterior_state_covs = (
    smoothed_posterior_state_rvs.cov
)  # Shape: (num_time_points, state_dim, state_dim)
```

### Visualize Results

```python
state_fig = plt.figure()

ax = state_fig.add_subplot()

# Plot means
ax.plot(grid, posterior_state_means, label="posterior mean")

# Plot marginal standard deviations
std_x = np.sqrt(np.abs(posterior_state_covs)).squeeze()
ax.fill_between(
    grid,
    posterior_state_means - 1.96 * std_x,
    posterior_state_means + 1.96 * std_x,
    alpha=0.2,
    label="1.96 marginal stddev for filter",
)

# Plot smoothed means
ax.plot(time_grid, smoothed_posterior_state_means, label="smoothed posterior mean")

# Plot marginal standard deviations
std_x = np.sqrt(smoothed_posterior_state_covs).squeeze()
ax.fill_between(
    time_grid,
    smoothed_posterior_state_means - 1.96 * std_x,
    smoothed_posterior_state_means + 1.96 * std_x,
    alpha=0.2,
    label="1.96 marginal stddev for smoother",
)

ax.scatter(time_grid, observations, marker=".", label="measurements")

# Add labels etc.
ax.set_xlabel("t")
ax.set_title(r"$x$")
ax.legend()

state_fig.tight_layout()
```

To get a closer look at filter and smoother, we inspect the interval $[0, 1]$ in greater detail.

```python
idx = time_grid < 1.0

ax.set_xlim([0, 1])
ax.set_ylim(np.min(observations[idx]), np.max(observations[idx]))

state_fig.tight_layout()
state_fig
```