# Single Particle Model (SPM)

## Model Equations

The SPM consists of two spherically symmetric diffusion equations: one within a representative negative particle ($\text{k}=\text{n}$) and one within a representative positive particle ($\text{k}=\text{p}$). In the centre of the particle the standard no-flux condition is imposed. Since the SPM assumes that all particles in an electrode behave in exactly the same way, the flux on the surface of a particle is simply the current $I$ divided by the thickness of the electrode $L_{\text{k}}$. The concentration of lithium in electrode $\text{k}$ is denoted $c_{\text{k}}$ and the current is denoted by $I$. All parameters in the model stated here are dimensionless and are given in terms of dimensional parameters at the end of this notebook. The model equations for the SPM are then:
\begin{align}
\mathcal{C}_{\text{k}} \frac{\partial c_{\text{s,k}}}{\partial t} &= -\frac{1}{r_{\text{k}}^2} \frac{\partial}{\partial r_{\text{k}}} \left(r_{\text{k}}^2 N_{\text{s,k}}\right), \\
N_{\text{s,k}} &= -D_{\text{s,k}}(c_{\text{s,k}}) \frac{\partial c_{\text{s,k}}}{\partial r_{\text{k}}}, \quad \text{k} \in \text{n, p}, \end{align}
$$
N_{\text{s,k}}\big|_{r_{\text{k}}=0} = 0, \quad \text{k} \in \text{n, p}, \quad \ \ - \frac{a_{R, \text{k}}\gamma_{\text{k}}}{\mathcal{C}_{\text{k}}} N_{\text{s,k}}\big|_{r_{\text{k}}=1} =
\begin{cases}
		  \frac{I}{L_{\text{n}}}, \quad &\text{k}=\text{n}, \\
		  -\frac{I}{L_{\text{p}}}, \quad &\text{k}=\text{p},
\end{cases} \\
c_{\text{s,k}}(r_{\text{k}},0) = c_{\text{s,k,0}}, \quad \text{k} \in \text{n, p},$$
where $D_{\text{s,k}}$ is the diffusion coefficient in the solid, $N_{\text{s,k}}$ denotes the flux of lithium ions in the solid particle within the region $\text{k}$, and $r_{\text{k}} \in[0,1]$ is the radial coordinate of the particle in electrode $\text{k}$.

### Voltage Expression
The terminal voltage is obtained from the expression:
$$
V = U_{\text{p}}(c_{\text{p}})\big|_{r_{\text{p}}=1} - U_{\text{n}}(c_{\text{n}})\big|_{r_{\text{n}}=1} -2\sinh^{-1}\left(\frac{I}{j_{\text{0,p}} L_{\text{p}}}\right) - 2\sinh^{-1}\left(\frac{I}{j_{\text{0,n}} L_{\text{n}}}\right)
$$
with the exchange current densities given by
$$j_{\text{0,k}} =  \frac{\gamma_{\text{k}}}{\mathcal{C}_{\text{r,k}}}(c_{\text{k}})^{1/2}(1-c_{\text{k}})^{1/2}  $$

More details can be found in [[3]](#References).

## Example solving SPM using PyBaMM

Below we show how to solve the Single Particle Model, using the default geometry, mesh, parameters, discretisation and solver provided with PyBaMM. In this notebook we explicitly handle all the stages of setting up, processing and solving the model in order to explain them in detail. However, it is often simpler in practice to use the `Simulation` class, which handles many of the stages automatically, as shown [here](../simulation-class.ipynb).

First we need to import `pybamm`, and then change our working directory to the root of the pybamm folder.

```python
%pip install pybamm -q    # install PyBaMM if it is not installed
import pybamm
import numpy as np
import os
import matplotlib.pyplot as plt
os.chdir(pybamm.__path__[0]+'/..')
```

    Note: you may need to restart the kernel to use updated packages.

We then create an instance of the SPM:

```python
model = pybamm.lithium_ion.SPM()
```

The model object is a subtype of [`pybamm.BaseModel`](https://pybamm.readthedocs.io/en/latest/source/models/base_models/base_model.html), and contains all the equations that define this particular model. For example, the `rhs` dict contained in `model` has a dictionary mapping variables such as $c_n$ to the equation representing its rate of change with time (i.e. $\partial{c_n}/\partial{t}$). We can see this explicitly by visualising this entry in the `rhs` dict:

```python
variable = next(iter(model.rhs.keys()))
equation = next(iter(model.rhs.values()))
print('rhs equation for variable \'',variable,'\' is:')
path = 'examples/notebooks/models/'
equation.visualise(path+'spm1.png')
```

    rhs equation for variable ' Discharge capacity [A.h] ' is:

We need a geometry in which to define our model equations. In pybamm this is represented by the [`pybamm.Geometry`](https://pybamm.readthedocs.io/en/latest/source/geometry/geometry.html) class. In this case we use the default geometry object defined by the model

```python
geometry = model.default_geometry
```

This geometry object defines a number of domains, each with its own name, spatial variables and min/max limits (the latter are represented as equations similar to the rhs equation shown above). For instance, the SPM has the following domains:

```python
print('SPM domains:')
for i, (k, v) in enumerate(geometry.items()):
    print(str(i+1)+'.',k,'with variables:')
    for var, rng in v.items():
        if 'min' in rng:
            print('  -(',rng['min'],') <=',var,'<= (',rng['max'],')')
        else:
            print(var, '=', rng['position'])
```

    SPM domains:
    1. negative electrode with variables:
      -( 0 ) <= x_n <= ( Negative electrode thickness [m] / (Negative electrode thickness [m] + Separator thickness [m] + Positive electrode thickness [m]) )
    2. separator with variables:
      -( Negative electrode thickness [m] / (Negative electrode thickness [m] + Separator thickness [m] + Positive electrode thickness [m]) ) <= x_s <= ( Negative electrode thickness [m] / (Negative electrode thickness [m] + Separator thickness [m] + Positive electrode thickness [m]) + Separator thickness [m] / (Negative electrode thickness [m] + Separator thickness [m] + Positive electrode thickness [m]) )
    3. positive electrode with variables:
      -( Negative electrode thickness [m] / (Negative electrode thickness [m] + Separator thickness [m] + Positive electrode thickness [m]) + Separator thickness [m] / (Negative electrode thickness [m] + Separator thickness [m] + Positive electrode thickness [m]) ) <= x_p <= ( 1 )
    4. negative particle with variables:
      -( 0 ) <= r_n <= ( 1 )
    5. positive particle with variables:
      -( 0 ) <= r_p <= ( 1 )
    6. current collector with variables:
    z = 1

Both the model equations and the geometry include parameters, such as $\gamma_p$ or $L_p$. We can substitute these symbolic parameters in the model with values by using the [`pybamm.ParameterValues`](https://pybamm.readthedocs.io/en/latest/source/parameters/parameter_values.html) class, which takes either a python dictionary or CSV file with the mapping between parameter names and values. Rather than create our own instance of `pybamm.ParameterValues`, we will use the default parameter set included in the model

```python
param = model.default_parameter_values
```

We can then apply this parameter set to the model and geometry

```python
param.process_model(model)
param.process_geometry(geometry)
```

The next step is to mesh the input geometry. We can do this using the [`pybamm.Mesh`](https://pybamm.readthedocs.io/en/latest/source/meshes/meshes.html) class. This class takes in the geometry of the problem, and also two dictionaries containing the type of mesh to use within each domain of the geometry (i.e. within the positive or negative electrode domains), and the number of mesh points.

The default mesh types and the default number of points to use in each variable for the SPM are:

```python
for k, t in model.default_submesh_types.items():
    print(k,'is of type',t.__repr__())
for var, npts in model.default_var_pts.items():
    print(var,'has',npts,'mesh points')
```

    negative electrode is of type Generator for Uniform1DSubMesh
    separator is of type Generator for Uniform1DSubMesh
    positive electrode is of type Generator for Uniform1DSubMesh
    negative particle is of type Generator for Uniform1DSubMesh
    positive particle is of type Generator for Uniform1DSubMesh
    current collector is of type Generator for SubMesh0D
    x_n has 20 mesh points
    x_s has 20 mesh points
    x_p has 20 mesh points
    r_n has 30 mesh points
    r_p has 30 mesh points
    y has 10 mesh points
    z has 10 mesh points

With these defaults, we can then create our mesh of the given geometry:

```python
mesh = pybamm.Mesh(geometry, model.default_submesh_types, model.default_var_pts)
```

The next step is to discretise the model equations using this mesh. We do this using the [`pybamm.Discretisation`](https://pybamm.readthedocs.io/en/latest/source/spatial_methods/discretisation.html) class, which takes both the mesh we have already created, and a dictionary of spatial methods to use for each geometry domain. For the case of the SPM, we use the following defaults for the spatial discretisation methods:

```python
for k, method in model.default_spatial_methods.items():
    print(k,'is discretised using',method.__class__.__name__,'method')
```

    macroscale is discretised using FiniteVolume method
    negative particle is discretised using FiniteVolume method
    positive particle is discretised using FiniteVolume method
    current collector is discretised using ZeroDimensionalSpatialMethod method

We then create the `pybamm.Discretisation` object, and use this to discretise the model equations

```python
disc = pybamm.Discretisation(mesh, model.default_spatial_methods)
disc.process_model(model);
```

After this stage, all of the variables in `model` have been discretised into `pybamm.StateVector` objects, and spatial operators have been replaced by matrix-vector multiplications, ready to be evaluated within a time-stepping algorithm of a given solver. For example, the rhs expression for $\partial{c_n}/\partial{t}$ that we visualised above is now represented by:

```python
model.concatenated_rhs.children[0].visualise(path+'spm2.png')
```

Now we are ready to run the time-stepping routine to solve the model. Once again we use the default ODE solver.

```python
# Solve the model at the given time points (in seconds)
solver = model.default_solver
n = 250
t_eval = np.linspace(0, 3600, n)
print('Solving using',type(solver).__name__,'solver...')
solution = solver.solve(model, t_eval)
print('Finished.')
```

    Solving using CasadiSolver solver...
    Finished.

Each model in pybamm has a list of relevant variables defined in the model, for use in visualising the model solution or for comparison with other models. The SPM defines the following variables:

```python
print('SPM model variables:')
for v in model.variables.keys():
    print('\t-',v)
```

    SPM model variables:
    	- Time
    	- Time [s]
    	- Time [min]
    	- Time [h]
    	- x
    	- x [m]
    	- x_n
    	- x_n [m]
    	- x_s
    	- x_s [m]
    	- x_p
    	- x_p [m]
    	- Sum of electrolyte reaction source terms
    	- Sum of negative electrode electrolyte reaction source terms
    	- Sum of positive electrode electrolyte reaction source terms
    	- Sum of x-averaged negative electrode electrolyte reaction source terms
    	- Sum of x-averaged positive electrode electrolyte reaction source terms
    	- Sum of interfacial current densities
    	- Sum of negative electrode interfacial current densities
    	- Sum of positive electrode interfacial current densities
    	- Sum of x-averaged negative electrode interfacial current densities
    	- Sum of x-averaged positive electrode interfacial current densities
    	- r_n
    	- r_n [m]
    	- r_p
    	- r_p [m]
    	- Total current density
    	- Total current density [A.m-2]
    	- Current [A]
    	- C-rate
    	- Discharge capacity [A.h]
    	- Porosity
    	- Negative electrode porosity
    	- Separator porosity
    	- Positive electrode porosity
    	- X-averaged negative electrode porosity
    	- X-averaged separator porosity
    	- X-averaged positive electrode porosity
    	- Leading-order porosity
    	- Leading-order negative electrode porosity
    	- Leading-order separator porosity
    	- Leading-order positive electrode porosity
    	- Leading-order x-averaged negative electrode porosity
    	- Leading-order x-averaged separator porosity
    	- Leading-order x-averaged positive electrode porosity
    	- Porosity change
    	- Negative electrode porosity change
    	- Separator porosity change
    	- Positive electrode porosity change
    	- X-averaged negative electrode porosity change
    	- X-averaged separator porosity change
    	- X-averaged positive electrode porosity change
    	- Leading-order x-averaged negative electrode porosity change
    	- Leading-order x-averaged separator porosity change
    	- Leading-order x-averaged positive electrode porosity change
    	- Negative electrode active material volume fraction
    	- X-averaged negative electrode active material volume fraction
    	- Negative particle radius
    	- Negative particle radius [m]
    	- Negative electrode surface area to volume ratio
    	- Negative electrode surface area to volume ratio [m-1]
    	- X-averaged negative electrode surface area to volume ratio
    	- X-averaged negative electrode surface area to volume ratio [m-1]
    	- Negative electrode active material volume fraction change
    	- X-averaged negative electrode active material volume fraction change
    	- Positive electrode active material volume fraction
    	- X-averaged positive electrode active material volume fraction
    	- Positive particle radius
    	- Positive particle radius [m]
    	- Positive electrode surface area to volume ratio
    	- Positive electrode surface area to volume ratio [m-1]
    	- X-averaged positive electrode surface area to volume ratio
    	- X-averaged positive electrode surface area to volume ratio [m-1]
    	- Positive electrode active material volume fraction change
    	- X-averaged positive electrode active material volume fraction change
    	- Negative electrode volume-averaged velocity
    	- Positive electrode volume-averaged velocity
    	- Negative electrode volume-averaged velocity [m.s-1]
    	- Positive electrode volume-averaged velocity [m.s-1]
    	- Negative electrode volume-averaged acceleration
    	- Positive electrode volume-averaged acceleration
    	- Negative electrode volume-averaged acceleration [m.s-1]
    	- Positive electrode volume-averaged acceleration [m.s-1]
    	- X-averaged negative electrode volume-averaged acceleration
    	- X-averaged positive electrode volume-averaged acceleration
    	- X-averaged negative electrode volume-averaged acceleration [m.s-1]
    	- X-averaged positive electrode volume-averaged acceleration [m.s-1]
    	- Negative electrode pressure
    	- Positive electrode pressure
    	- X-averaged negative electrode pressure
    	- X-averaged positive electrode pressure
    	- Separator pressure
    	- X-averaged separator pressure
    	- Negative electrode transverse volume-averaged velocity
    	- Separator transverse volume-averaged velocity
    	- Positive electrode transverse volume-averaged velocity
    	- Negative electrode transverse volume-averaged velocity [m.s-2]
    	- Separator transverse volume-averaged velocity [m.s-2]
    	- Positive electrode transverse volume-averaged velocity [m.s-2]
    	- X-averaged negative electrode transverse volume-averaged velocity
    	- X-averaged separator transverse volume-averaged velocity
    	- X-averaged positive electrode transverse volume-averaged velocity
    	- X-averaged negative electrode transverse volume-averaged velocity [m.s-2]
    	- X-averaged separator transverse volume-averaged velocity [m.s-2]
    	- X-averaged positive electrode transverse volume-averaged velocity [m.s-2]
    	- Transverse volume-averaged velocity
    	- Transverse volume-averaged velocity [m.s-2]
    	- Negative electrode transverse volume-averaged acceleration
    	- Separator transverse volume-averaged acceleration
    	- Positive electrode transverse volume-averaged acceleration
    	- Negative electrode transverse volume-averaged acceleration [m.s-2]
    	- Separator transverse volume-averaged acceleration [m.s-2]
    	- Positive electrode transverse volume-averaged acceleration [m.s-2]
    	- X-averaged negative electrode transverse volume-averaged acceleration
    	- X-averaged separator transverse volume-averaged acceleration
    	- X-averaged positive electrode transverse volume-averaged acceleration
    	- X-averaged negative electrode transverse volume-averaged acceleration [m.s-2]
    	- X-averaged separator transverse volume-averaged acceleration [m.s-2]
    	- X-averaged positive electrode transverse volume-averaged acceleration [m.s-2]
    	- Transverse volume-averaged acceleration
    	- Transverse volume-averaged acceleration [m.s-2]
    	- Negative particle concentration
    	- Negative particle concentration [mol.m-3]
    	- X-averaged negative particle concentration
    	- X-averaged negative particle concentration [mol.m-3]
    	- R-averaged negative particle concentration
    	- R-averaged negative particle concentration [mol.m-3]
    	- Average negative particle concentration
    	- Average negative particle concentration [mol.m-3]
    	- Negative particle surface concentration
    	- Negative particle surface concentration [mol.m-3]
    	- X-averaged negative particle surface concentration
    	- X-averaged negative particle surface concentration [mol.m-3]
    	- Negative electrode extent of lithiation
    	- X-averaged negative electrode extent of lithiation
    	- Minimum negative particle concentration
    	- Maximum negative particle concentration
    	- Minimum negative particle concentration [mol.m-3]
    	- Maximum negative particle concentration [mol.m-3]
    	- Minimum negative particle surface concentration
    	- Maximum negative particle surface concentration
    	- Minimum negative particle surface concentration [mol.m-3]
    	- Maximum negative particle surface concentration [mol.m-3]
    	- Positive particle concentration
    	- Positive particle concentration [mol.m-3]
    	- X-averaged positive particle concentration
    	- X-averaged positive particle concentration [mol.m-3]
    	- R-averaged positive particle concentration
    	- R-averaged positive particle concentration [mol.m-3]
    	- Average positive particle concentration
    	- Average positive particle concentration [mol.m-3]
    	- Positive particle surface concentration
    	- Positive particle surface concentration [mol.m-3]
    	- X-averaged positive particle surface concentration
    	- X-averaged positive particle surface concentration [mol.m-3]
    	- Positive electrode extent of lithiation
    	- X-averaged positive electrode extent of lithiation
    	- Minimum positive particle concentration
    	- Maximum positive particle concentration
    	- Minimum positive particle concentration [mol.m-3]
    	- Maximum positive particle concentration [mol.m-3]
    	- Minimum positive particle surface concentration
    	- Maximum positive particle surface concentration
    	- Minimum positive particle surface concentration [mol.m-3]
    	- Maximum positive particle surface concentration [mol.m-3]
    	- Electrolyte concentration
    	- Electrolyte concentration [mol.m-3]
    	- Electrolyte concentration [Molar]
    	- X-averaged electrolyte concentration
    	- X-averaged electrolyte concentration [mol.m-3]
    	- X-averaged electrolyte concentration [Molar]
    	- Negative electrolyte concentration
    	- Negative electrolyte concentration [mol.m-3]
    	- Negative electrolyte concentration [Molar]
    	- Separator electrolyte concentration
    	- Separator electrolyte concentration [mol.m-3]
    	- Separator electrolyte concentration [Molar]
    	- Positive electrolyte concentration
    	- Positive electrolyte concentration [mol.m-3]
    	- Positive electrolyte concentration [Molar]
    	- X-averaged negative electrolyte concentration
    	- X-averaged negative electrolyte concentration [mol.m-3]
    	- X-averaged separator electrolyte concentration
    	- X-averaged separator electrolyte concentration [mol.m-3]
    	- X-averaged positive electrolyte concentration
    	- X-averaged positive electrolyte concentration [mol.m-3]
    	- Electrolyte flux
    	- Electrolyte flux [mol.m-2.s-1]
    	- Negative current collector temperature
    	- Negative current collector temperature [K]
    	- X-averaged negative electrode temperature
    	- X-averaged negative electrode temperature [K]
    	- Negative electrode temperature
    	- Negative electrode temperature [K]
    	- X-averaged separator temperature
    	- X-averaged separator temperature [K]
    	- Separator temperature
    	- Separator temperature [K]
    	- X-averaged positive electrode temperature
    	- X-averaged positive electrode temperature [K]
    	- Positive electrode temperature
    	- Positive electrode temperature [K]
    	- Positive current collector temperature
    	- Positive current collector temperature [K]
    	- Cell temperature
    	- Cell temperature [K]
    	- X-averaged cell temperature
    	- X-averaged cell temperature [K]
    	- Volume-averaged cell temperature
    	- Volume-averaged cell temperature [K]
    	- Ambient temperature [K]
    	- Ambient temperature
    	- Inner SEI thickness
    	- Inner SEI thickness [m]
    	- X-averaged inner SEI thickness
    	- X-averaged inner SEI thickness [m]
    	- Outer SEI thickness
    	- Outer SEI thickness [m]
    	- X-averaged outer SEI thickness
    	- X-averaged outer SEI thickness [m]
    	- Total SEI thickness
    	- Total SEI thickness [m]
    	- X-averaged total SEI thickness
    	- X-averaged total SEI thickness [m]
    	- X-averaged negative electrode resistance [Ohm.m2]
    	- Inner SEI concentration [mol.m-3]
    	- X-averaged inner SEI concentration [mol.m-3]
    	- Outer SEI concentration [mol.m-3]
    	- X-averaged outer SEI concentration [mol.m-3]
    	- Negative SEI concentration [mol.m-3]
    	- X-averaged SEI concentration [mol.m-3]
    	- Loss of lithium to SEI [mol]
    	- Inner SEI interfacial current density
    	- Inner SEI interfacial current density [A.m-2]
    	- X-averaged inner SEI interfacial current density
    	- X-averaged inner SEI interfacial current density [A.m-2]
    	- Outer SEI interfacial current density
    	- Outer SEI interfacial current density [A.m-2]
    	- X-averaged outer SEI interfacial current density
    	- X-averaged outer SEI interfacial current density [A.m-2]
    	- SEI interfacial current density
    	- SEI interfacial current density [A.m-2]
    	- X-averaged SEI interfacial current density
    	- X-averaged SEI interfacial current density [A.m-2]
    	- Inner positive electrode SEI thickness
    	- Inner positive electrode SEI thickness [m]
    	- X-averaged inner positive electrode SEI thickness
    	- X-averaged inner positive electrode SEI thickness [m]
    	- Outer positive electrode SEI thickness
    	- Outer positive electrode SEI thickness [m]
    	- X-averaged outer positive electrode SEI thickness
    	- X-averaged outer positive electrode SEI thickness [m]
    	- Total positive electrode SEI thickness
    	- Total positive electrode SEI thickness [m]
    	- X-averaged total positive electrode SEI thickness
    	- X-averaged total positive electrode SEI thickness [m]
    	- X-averaged positive electrode resistance [Ohm.m2]
    	- Inner positive electrode SEI concentration [mol.m-3]
    	- X-averaged inner positive electrode SEI concentration [mol.m-3]
    	- Outer positive electrode SEI concentration [mol.m-3]
    	- X-averaged outer positive electrode SEI concentration [mol.m-3]
    	- Positive SEI concentration [mol.m-3]
    	- X-averaged positive electrode SEI concentration [mol.m-3]
    	- Loss of lithium to positive electrode SEI [mol]
    	- Inner positive electrode SEI interfacial current density
    	- Inner positive electrode SEI interfacial current density [A.m-2]
    	- X-averaged inner positive electrode SEI interfacial current density
    	- X-averaged inner positive electrode SEI interfacial current density [A.m-2]
    	- Outer positive electrode SEI interfacial current density
    	- Outer positive electrode SEI interfacial current density [A.m-2]
    	- X-averaged outer positive electrode SEI interfacial current density
    	- X-averaged outer positive electrode SEI interfacial current density [A.m-2]
    	- Positive electrode SEI interfacial current density
    	- Positive electrode SEI interfacial current density [A.m-2]
    	- X-averaged positive electrode SEI interfacial current density
    	- X-averaged positive electrode SEI interfacial current density [A.m-2]
    	- Electrolyte tortuosity
    	- Negative electrolyte tortuosity
    	- Positive electrolyte tortuosity
    	- X-averaged negative electrolyte tortuosity
    	- X-averaged positive electrolyte tortuosity
    	- Separator tortuosity
    	- X-averaged separator tortuosity
    	- Electrode tortuosity
    	- Negative electrode tortuosity
    	- Positive electrode tortuosity
    	- X-averaged negative electrode tortuosity
    	- X-averaged positive electrode tortuosity
    	- Separator volume-averaged velocity
    	- Separator volume-averaged velocity [m.s-1]
    	- Separator volume-averaged acceleration
    	- Separator volume-averaged acceleration [m.s-1]
    	- X-averaged separator volume-averaged acceleration
    	- X-averaged separator volume-averaged acceleration [m.s-1]
    	- Volume-averaged velocity
    	- Volume-averaged velocity [m.s-1]
    	- Volume-averaged acceleration
    	- X-averaged volume-averaged acceleration
    	- Volume-averaged acceleration [m.s-1]
    	- X-averaged volume-averaged acceleration [m.s-1]
    	- Pressure
    	- Negative particle flux
    	- X-averaged negative particle flux
    	- Negative electrode volume-averaged concentration
    	- Negative electrode volume-averaged concentration [mol.m-3]
    	- Total lithium in negative electrode [mol]
    	- Positive particle flux
    	- X-averaged positive particle flux
    	- Total lithium in electrolyte [mol]
    	- Positive electrode volume-averaged concentration
    	- Positive electrode volume-averaged concentration [mol.m-3]
    	- Total lithium in positive electrode [mol]
    	- Total concentration in electrolyte [mol]
    	- Ohmic heating
    	- Ohmic heating [W.m-3]
    	- X-averaged Ohmic heating
    	- X-averaged Ohmic heating [W.m-3]
    	- Volume-averaged Ohmic heating
    	- Volume-averaged Ohmic heating [W.m-3]
    	- Irreversible electrochemical heating
    	- Irreversible electrochemical heating [W.m-3]
    	- X-averaged irreversible electrochemical heating
    	- X-averaged irreversible electrochemical heating [W.m-3]
    	- Volume-averaged irreversible electrochemical heating
    	- Volume-averaged irreversible electrochemical heating[W.m-3]
    	- Reversible heating
    	- Reversible heating [W.m-3]
    	- X-averaged reversible heating
    	- X-averaged reversible heating [W.m-3]
    	- Volume-averaged reversible heating
    	- Volume-averaged reversible heating [W.m-3]
    	- Total heating
    	- Total heating [W.m-3]
    	- X-averaged total heating
    	- X-averaged total heating [W.m-3]
    	- Volume-averaged total heating
    	- Volume-averaged total heating [W.m-3]
    	- Negative current collector potential
    	- Negative current collector potential [V]
    	- Current collector current density
    	- Current collector current density [A.m-2]
    	- Leading-order current collector current density
    	- SEI interfacial current density
    	- SEI interfacial current density [A.m-2]
    	- SEI interfacial current density per volume [A.m-3]
    	- X-averaged negative electrode total interfacial current density
    	- X-averaged negative electrode total interfacial current density [A.m-2]
    	- X-averaged negative electrode total interfacial current density per volume [A.m-3]
    	- Negative electrode exchange current density
    	- X-averaged negative electrode exchange current density
    	- Negative electrode exchange current density [A.m-2]
    	- X-averaged negative electrode exchange current density [A.m-2]
    	- Negative electrode exchange current density per volume [A.m-3]
    	- X-averaged negative electrode exchange current density per volume [A.m-3]
    	- Negative electrode reaction overpotential
    	- X-averaged negative electrode reaction overpotential
    	- Negative electrode reaction overpotential [V]
    	- X-averaged negative electrode reaction overpotential [V]
    	- Negative electrode surface potential difference
    	- X-averaged negative electrode surface potential difference
    	- Negative electrode surface potential difference [V]
    	- X-averaged negative electrode surface potential difference [V]
    	- SEI film overpotential
    	- X-averaged SEI film overpotential
    	- SEI film overpotential [V]
    	- X-averaged SEI film overpotential [V]
    	- Negative electrode open circuit potential
    	- Negative electrode open circuit potential [V]
    	- X-averaged negative electrode open circuit potential
    	- X-averaged negative electrode open circuit potential [V]
    	- Negative electrode entropic change
    	- X-averaged negative electrode entropic change
    	- X-averaged positive electrode total interfacial current density
    	- X-averaged positive electrode total interfacial current density [A.m-2]
    	- X-averaged positive electrode total interfacial current density per volume [A.m-3]
    	- Positive electrode exchange current density
    	- X-averaged positive electrode exchange current density
    	- Positive electrode exchange current density [A.m-2]
    	- X-averaged positive electrode exchange current density [A.m-2]
    	- Positive electrode exchange current density per volume [A.m-3]
    	- X-averaged positive electrode exchange current density per volume [A.m-3]
    	- Positive electrode reaction overpotential
    	- X-averaged positive electrode reaction overpotential
    	- Positive electrode reaction overpotential [V]
    	- X-averaged positive electrode reaction overpotential [V]
    	- Positive electrode surface potential difference
    	- X-averaged positive electrode surface potential difference
    	- Positive electrode surface potential difference [V]
    	- X-averaged positive electrode surface potential difference [V]
    	- Positive electrode SEI film overpotential
    	- X-averaged positive electrode SEI film overpotential
    	- Positive electrode SEI film overpotential [V]
    	- X-averaged positive electrode SEI film overpotential [V]
    	- Positive electrode open circuit potential
    	- Positive electrode open circuit potential [V]
    	- X-averaged positive electrode open circuit potential
    	- X-averaged positive electrode open circuit potential [V]
    	- Positive electrode entropic change
    	- X-averaged positive electrode entropic change
    	- Negative electrode interfacial current density
    	- X-averaged negative electrode interfacial current density
    	- Negative electrode interfacial current density [A.m-2]
    	- X-averaged negative electrode interfacial current density [A.m-2]
    	- Negative electrode interfacial current density per volume [A.m-3]
    	- X-averaged negative electrode interfacial current density per volume [A.m-3]
    	- Positive electrode interfacial current density
    	- X-averaged positive electrode interfacial current density
    	- Positive electrode interfacial current density [A.m-2]
    	- X-averaged positive electrode interfacial current density [A.m-2]
    	- Positive electrode interfacial current density per volume [A.m-3]
    	- X-averaged positive electrode interfacial current density per volume [A.m-3]
    	- Interfacial current density
    	- Interfacial current density [A.m-2]
    	- Interfacial current density per volume [A.m-3]
    	- Exchange current density
    	- Exchange current density [A.m-2]
    	- Exchange current density per volume [A.m-3]
    	- Negative electrode oxygen interfacial current density
    	- X-averaged negative electrode oxygen interfacial current density
    	- Negative electrode oxygen interfacial current density [A.m-2]
    	- X-averaged negative electrode oxygen interfacial current density [A.m-2]
    	- Negative electrode oxygen interfacial current density per volume [A.m-3]
    	- X-averaged negative electrode oxygen interfacial current density per volume [A.m-3]
    	- Negative electrode oxygen exchange current density
    	- X-averaged negative electrode oxygen exchange current density
    	- Negative electrode oxygen exchange current density [A.m-2]
    	- X-averaged negative electrode oxygen exchange current density [A.m-2]
    	- Negative electrode oxygen exchange current density per volume [A.m-3]
    	- X-averaged negative electrode oxygen exchange current density per volume [A.m-3]
    	- Negative electrode oxygen reaction overpotential
    	- X-averaged negative electrode oxygen reaction overpotential
    	- Negative electrode oxygen reaction overpotential [V]
    	- X-averaged negative electrode oxygen reaction overpotential [V]
    	- Negative electrode oxygen open circuit potential
    	- Negative electrode oxygen open circuit potential [V]
    	- X-averaged negative electrode oxygen open circuit potential
    	- X-averaged negative electrode oxygen open circuit potential [V]
    	- Positive electrode oxygen interfacial current density
    	- X-averaged positive electrode oxygen interfacial current density
    	- Positive electrode oxygen interfacial current density [A.m-2]
    	- X-averaged positive electrode oxygen interfacial current density [A.m-2]
    	- Positive electrode oxygen interfacial current density per volume [A.m-3]
    	- X-averaged positive electrode oxygen interfacial current density per volume [A.m-3]
    	- Positive electrode oxygen exchange current density
    	- X-averaged positive electrode oxygen exchange current density
    	- Positive electrode oxygen exchange current density [A.m-2]
    	- X-averaged positive electrode oxygen exchange current density [A.m-2]
    	- Positive electrode oxygen exchange current density per volume [A.m-3]
    	- X-averaged positive electrode oxygen exchange current density per volume [A.m-3]
    	- Positive electrode oxygen reaction overpotential
    	- X-averaged positive electrode oxygen reaction overpotential
    	- Positive electrode oxygen reaction overpotential [V]
    	- X-averaged positive electrode oxygen reaction overpotential [V]
    	- Positive electrode oxygen open circuit potential
    	- Positive electrode oxygen open circuit potential [V]
    	- X-averaged positive electrode oxygen open circuit potential
    	- X-averaged positive electrode oxygen open circuit potential [V]
    	- Oxygen interfacial current density
    	- Oxygen interfacial current density [A.m-2]
    	- Oxygen interfacial current density per volume [A.m-3]
    	- Oxygen exchange current density
    	- Oxygen exchange current density [A.m-2]
    	- Oxygen exchange current density per volume [A.m-3]
    	- Negative electrode potential
    	- Negative electrode potential [V]
    	- X-averaged negative electrode potential
    	- X-averaged negative electrode potential [V]
    	- Negative electrode ohmic losses
    	- Negative electrode ohmic losses [V]
    	- X-averaged negative electrode ohmic losses
    	- X-averaged negative electrode ohmic losses [V]
    	- Gradient of negative electrode potential
    	- Negative electrode current density
    	- Negative electrode current density [A.m-2]
    	- Negative electrolyte potential
    	- Negative electrolyte potential [V]
    	- Separator electrolyte potential
    	- Separator electrolyte potential [V]
    	- Positive electrolyte potential
    	- Positive electrolyte potential [V]
    	- Electrolyte potential
    	- Electrolyte potential [V]
    	- X-averaged electrolyte potential
    	- X-averaged electrolyte potential [V]
    	- X-averaged negative electrolyte potential
    	- X-averaged negative electrolyte potential [V]
    	- X-averaged separator electrolyte potential
    	- X-averaged separator electrolyte potential [V]
    	- X-averaged positive electrolyte potential
    	- X-averaged positive electrolyte potential [V]
    	- X-averaged electrolyte overpotential
    	- X-averaged electrolyte overpotential [V]
    	- Gradient of negative electrolyte potential
    	- Gradient of separator electrolyte potential
    	- Gradient of positive electrolyte potential
    	- Gradient of electrolyte potential
    	- Electrolyte current density
    	- Electrolyte current density [A.m-2]
    	- Negative electrolyte current density
    	- Negative electrolyte current density [A.m-2]
    	- Positive electrolyte current density
    	- Positive electrolyte current density [A.m-2]
    	- X-averaged concentration overpotential
    	- X-averaged electrolyte ohmic losses
    	- X-averaged concentration overpotential [V]
    	- X-averaged electrolyte ohmic losses [V]
    	- Positive electrode potential
    	- Positive electrode potential [V]
    	- X-averaged positive electrode potential
    	- X-averaged positive electrode potential [V]
    	- Positive electrode ohmic losses
    	- Positive electrode ohmic losses [V]
    	- X-averaged positive electrode ohmic losses
    	- X-averaged positive electrode ohmic losses [V]
    	- Gradient of positive electrode potential
    	- Positive electrode current density
    	- Positive electrode current density [A.m-2]
    	- Electrode current density
    	- Positive current collector potential
    	- Positive current collector potential [V]
    	- Local voltage
    	- Local voltage [V]
    	- Terminal voltage
    	- Terminal voltage [V]
    	- X-averaged open circuit voltage
    	- Measured open circuit voltage
    	- X-averaged open circuit voltage [V]
    	- Measured open circuit voltage [V]
    	- X-averaged reaction overpotential
    	- X-averaged reaction overpotential [V]
    	- X-averaged SEI film overpotential
    	- X-averaged SEI film overpotential [V]
    	- X-averaged solid phase ohmic losses
    	- X-averaged solid phase ohmic losses [V]
    	- X-averaged battery open circuit voltage [V]
    	- Measured battery open circuit voltage [V]
    	- X-averaged battery reaction overpotential [V]
    	- X-averaged battery solid phase ohmic losses [V]
    	- X-averaged battery electrolyte ohmic losses [V]
    	- X-averaged battery concentration overpotential [V]
    	- Battery voltage [V]
    	- Change in measured open circuit voltage
    	- Change in measured open circuit voltage [V]
    	- Local ECM resistance
    	- Local ECM resistance [Ohm]
    	- Terminal power [W]

To help visualise the results, pybamm provides the `pybamm.ProcessedVariable` class, which takes the output of a solver and a variable, and allows the user to evaluate the value of that variable at any given time or $x$ value. These processed variables are automatically created by the solution dictionary.

```python
voltage = solution['Terminal voltage [V]']
c_s_n_surf = solution['Negative particle surface concentration']
c_s_p_surf = solution['Positive particle surface concentration']
```

One we have these variables in hand, we can begin generating plots using a library such as Matplotlib. Below we plot the terminal voltage and surface particle concentrations versus time

```python
t = solution["Time [s]"].entries
x = solution["x [m]"].entries[:, 0]
f, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(13,4))

ax1.plot(t, voltage(t))
ax1.set_xlabel(r'$Time [s]$')
ax1.set_ylabel('Terminal voltage [V]')

ax2.plot(t, c_s_n_surf(t=t, x=x[0]))  # can evaluate at arbitrary x (single representative particle)
ax2.set_xlabel(r'$Time [s]$')
ax2.set_ylabel('Negative particle surface concentration')

ax3.plot(t, c_s_p_surf(t=t, x=x[-1]))  # can evaluate at arbitrary x (single representative particle)
ax3.set_xlabel(r'$Time [s]$')
ax3.set_ylabel('Positive particle surface concentration')

plt.tight_layout()
plt.show()
```

Some of the output variables are defined over space as well as time. Once option to visualise these variables is to use the `interact` slider widget. Below we plot the negative/positive particle concentration over $r$, using a slider to change the current time point

```python
c_s_n = solution['Negative particle concentration']
c_s_p = solution['Positive particle concentration']
r_n = solution["r_n [m]"].entries[:, 0]
r_p = solution["r_p [m]"].entries[:, 0]
```

```python
c_s_n = solution['Negative particle concentration']
c_s_p = solution['Positive particle concentration']
r_n = solution["r_n [m]"].entries[:, 0, 0]
r_p = solution["r_p [m]"].entries[:, 0, 0]

def plot_concentrations(t):
    f, (ax1, ax2) = plt.subplots(1, 2 ,figsize=(10,5))
    plot_c_n, = ax1.plot(r_n, c_s_n(r=r_n,t=t,x=x[0]))  # can evaluate at arbitrary x (single representative particle)
    plot_c_p, = ax2.plot(r_p, c_s_p(r=r_p,t=t,x=x[-1]))  # can evaluate at arbitrary x (single representative particle)
    ax1.set_ylabel('Negative particle concentration')
    ax2.set_ylabel('Positive particle concentration')
    ax1.set_xlabel(r'$r_n$ [m]')
    ax2.set_xlabel(r'$r_p$ [m]')
    ax1.set_ylim(0, 1)
    ax2.set_ylim(0, 1)
    plt.show()

import ipywidgets as widgets
widgets.interact(plot_concentrations, t=widgets.FloatSlider(min=0,max=3600,step=10,value=0));

```

    interactive(children=(FloatSlider(value=0.0, description='t', max=3600.0, step=10.0), Output()), _dom_classes=…

The QuickPlot class can be used to plot the common set of useful outputs which should give you a good initial overview of the model. The method `Quickplot.dynamic_plot` employs the slider widget.

```python
quick_plot = pybamm.QuickPlot(solution)
quick_plot.dynamic_plot();
```

    interactive(children=(FloatSlider(value=0.0, description='t', max=1.0, step=0.01), Output()), _dom_classes=('w…

## Dimensionless Parameters

In the table below, we provide the dimensionless parameters in the SPM in terms of the dimensional parameters in LCO.csv. We use a superscript * to indicate dimensional quantities.

| Parameter                 | Expression                              |Interpretation                             |
|:--------------------------|:----------------------------------------|:------------------------------------------|
| $L_{\text{k}}$            | $L_{\text{k}}^*/L^*$                    | Ratio of region thickness to cell thickness|
|$\mathcal{C}_{\text{k}}$   | $\tau_{\text{k}}^*/\tau_{\text{d}}^*$   | Ratio of solid diffusion and discharge timescales |
|$\mathcal{C}_{\text{r,k}}$ |$\tau_{\text{r,k}}^*/\tau_{\text{d}}^*$  |Ratio of reaction and discharge timescales|
|$a_{R, \text{k}}$             |$a_{\text{k}}^* R_{\text{k}}^*$          | Product of particle radius and surface area to volume ratio|
|$\gamma_{\text{k}}$        |$c_{\text{k,max}}^*/c_{\text{n,max}}^*$  |Ratio of maximum lithium concentrations in solid|

## References

The relevant papers for this notebook are:

```python
pybamm.print_citations()
```

    [1] Joel A. E. Andersson, Joris Gillis, Greg Horn, James B. Rawlings, and Moritz Diehl. CasADi – A software framework for nonlinear optimization and optimal control. Mathematical Programming Computation, 11(1):1–36, 2019. doi:10.1007/s12532-018-0139-4.
    [2] Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, and others. Array programming with NumPy. Nature, 585(7825):357–362, 2020. doi:10.1038/s41586-020-2649-2.
    [3] Scott G. Marquis, Valentin Sulzer, Robert Timms, Colin P. Please, and S. Jon Chapman. An asymptotic derivation of a single particle model with electrolyte. Journal of The Electrochemical Society, 166(15):A3693–A3706, 2019. doi:10.1149/2.0341915jes.
    [4] Valentin Sulzer, Scott G. Marquis, Robert Timms, Martin Robinson, and S. Jon Chapman. Python Battery Mathematical Modelling (PyBaMM). ECSarXiv. February, 2020. doi:10.1149/osf.io/67ckj.


```python
%pylab inline
```

    Populating the interactive namespace from numpy and matplotlib

```python
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
from numpy import *
from matplotlib.patches import FancyArrowPatch
from mpl_toolkits.mplot3d import proj3d
font = {'size': 14}
matplotlib.rc('font', **font)
```

```python
class Arrow3D(FancyArrowPatch):
    def __init__(self, xs, ys, zs, *args, **kwargs):
        FancyArrowPatch.__init__(self, (0,0), (0,0), *args, **kwargs)
        self._verts3d = xs, ys, zs
    def draw(self, renderer):
        xs3d, ys3d, zs3d = self._verts3d
        xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
        self.set_positions((xs[0],ys[0]),(xs[1],ys[1]))
        FancyArrowPatch.draw(self, renderer)
```

```python
from sympy import symbols, Matrix
from sympy.parsing.sympy_parser import parse_expr
```

```python
vec2 = parse_expr('Matrix([1, 1])')
vec3 = parse_expr('Matrix([1, 1, 1])')
vec2, vec3
```

    (Matrix([
     [1],
     [1]]), Matrix([
     [1],
     [1],
     [1]]))

```python
V = np.array([vec2[0].evalf(), vec2[1].evalf()]).astype('float')
V
```

    array([1., 1.])

```python
U = np.array([vec3[0].evalf(), vec3[1].evalf(), vec3[2].evalf()]).astype('float')
U
```

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

```python
V[0]
```

    array([3.])

```python
vec2 = parse_expr('Matrix([3,2])')
V = np.array(vec2).astype('float')
if V.shape == (2, 1): V = np.array([V])

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

plt.xlim([-2,2])
plt.ylim([-2,2])

plt.axis('equal')
X = [np.cos(x) for x in np.linspace(0, 2*np.pi, 64)]
Y = [np.sin(y) for y in np.linspace(0, 2*np.pi, 64)]

plot(X,Y, c='darkgray', lw=3)

for i in range(V.shape[0]):
    x = V[i][0]; y = V[i][1]
    magn = np.sqrt(x**2 + y**2)
    annotate("", xy=(x/magn, y/magn), xytext=(0, 0),arrowprops=dict(arrowstyle="->", color="purple", lw=4))

axhline(0, ls='dashed', alpha=0.33, c='gray', lw=3)
axvline(0, ls='dashed', alpha=0.33, c='gray', lw=3)
plt.axis('off')
fig.tight_layout()
```

```python
vec3 = parse_expr('Matrix([[1, 0, 0], [0, 1, 0], [0, 0, 1]])')
U = np.array(vec3).astype('float')
U.shape
```

    (3, 3)

```python
U
```

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

```python
vec3 = parse_expr('Matrix([0, 0, 1])')
U = np.array(vec3).astype('float')
if U.shape == (3, 1): U = np.array([U])

fig = plt.figure(figsize=(10,10))

ax = fig.add_subplot(111, projection='3d')
plt.xlim([-2,2])
plt.ylim([-2,2])

plt.axis('equal')
X = np.array([np.cos(x) for x in np.linspace(0, 2*np.pi, 64)])
Y = np.array([np.sin(y) for y in np.linspace(0, 2*np.pi, 64)])
zeros = np.zeros(X.shape)

plot(X,zeros,zeros,c='darkgray', lw=2, alpha=0.33, ls='dashed')
plot(zeros,X,zeros,c='darkgray', lw=2, alpha=0.33, ls='dashed')

#annotate("", xy=(V[0]/magn, V[1]/magn), xytext=(0, 0),arrowprops=dict(arrowstyle="->", color="purple", lw=3))
plot(zeros,zeros,X,c='darkgray', lw=2, alpha=0.33, ls='dashed')

for i in range(U.shape[0]):
    x = U[i][0]; y = U[i][1]; z = U[i][2]
    magn = np.sqrt(x**2+y**2+z**2)
    ax.add_artist(Arrow3D([0, x], [0, y], [0,z], mutation_scale=15, lw=3, arrowstyle="-|>", color="purple"))
    #ax.add_artist(Arrow3D([0, 0], [0, -1], [0,1], mutation_scale=15, lw=3, arrowstyle="-|>", color="purple"))

plot(X,Y,zeros,c='darkgray', lw=3, alpha=0.33)
plot(zeros,X,Y,c='darkgray', lw=3, alpha=0.33)
plt.axis('off')
fig.tight_layout()
```

```python

```


In this notebook we will walk through some of the main tools for optimizing your code in order to efficiently solve DifferentialEquations.jl. User-side optimizations are important because, for sufficiently difficult problems, most of the time will be spent inside of your `f` function, the function you are trying to solve. "Efficient" integrators are those that reduce the required number of `f` calls to hit the error tolerance. The main ideas for optimizing your DiffEq code, or any Julia function, are the following:

- Make it non-allocating
- Use StaticArrays for small arrays
- Use broadcast fusion
- Make it type-stable
- Reduce redundant calculations
- Make use of BLAS calls
- Optimize algorithm choice

We'll discuss these strategies in the context of small and large systems. Let's start with small systems.

## Optimizing Small Systems (<100 DEs)

Let's take the classic Lorenz system from before. Let's start by naively writing the system in its out-of-place form:
# Optimizing DiffEq Code
### Chris Rackauckas

```julia
function lorenz(u,p,t)
 dx = 10.0*(u[2]-u[1])
 dy = u[1]*(28.0-u[3]) - u[2]
 dz = u[1]*u[2] - (8/3)*u[3]
 [dx,dy,dz]
end
```

    lorenz (generic function with 1 method)

Here, `lorenz` returns an object, `[dx,dy,dz]`, which is created within the body of `lorenz`.

This is a common code pattern from high-level languages like MATLAB, SciPy, or R's deSolve. However, the issue with this form is that it allocates a vector, `[dx,dy,dz]`, at each step. Let's benchmark the solution process with this choice of function:

```julia
using DifferentialEquations, BenchmarkTools
u0 = [1.0;0.0;0.0]
tspan = (0.0,100.0)
prob = ODEProblem(lorenz,u0,tspan)
@benchmark solve(prob,Tsit5())
```

    ┌ Info: Precompiling BenchmarkTools [6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf]
    └ @ Base loading.jl:1186

    BenchmarkTools.Trial:
      memory estimate:  10.94 MiB
      allocs estimate:  102469
      --------------
      minimum time:     3.455 ms (0.00% GC)
      median time:      4.810 ms (0.00% GC)
      mean time:        4.788 ms (12.08% GC)
      maximum time:     56.229 ms (89.52% GC)
      --------------
      samples:          1043
      evals/sample:     1

The BenchmarkTools package's `@benchmark` runs the code multiple times to get an accurate measurement. The minimum time is the time it takes when your OS and other background processes aren't getting in the way. Notice that in this case it takes about 5ms to solve and allocates around 11.11 MiB. However, if we were to use this inside of a real user code we'd see a lot of time spent doing garbage collection (GC) to clean up all of the arrays we made. Even if we turn off saving we have these allocations.

```julia
@benchmark solve(prob,Tsit5(),save_everystep=false)
```

    BenchmarkTools.Trial:
      memory estimate:  9.57 MiB
      allocs estimate:  89577
      --------------
      minimum time:     2.993 ms (0.00% GC)
      median time:      4.178 ms (0.00% GC)
      mean time:        4.231 ms (10.06% GC)
      maximum time:     55.648 ms (89.30% GC)
      --------------
      samples:          1180
      evals/sample:     1

The problem of course is that arrays are created every time our derivative function is called. This function is called multiple times per step and is thus the main source of memory usage. To fix this, we can use the in-place form to ***make our code non-allocating***:

```julia
function lorenz!(du,u,p,t)
 du[1] = 10.0*(u[2]-u[1])
 du[2] = u[1]*(28.0-u[3]) - u[2]
 du[3] = u[1]*u[2] - (8/3)*u[3]
end
```

    lorenz! (generic function with 1 method)

Here, instead of creating an array each time, we utilized the cache array `du`. When the inplace form is used, DifferentialEquations.jl takes a different internal route that minimizes the internal allocations as well. When we benchmark this function, we will see quite a difference.

```julia
u0 = [1.0;0.0;0.0]
tspan = (0.0,100.0)
prob = ODEProblem(lorenz!,u0,tspan)
@benchmark solve(prob,Tsit5())
```

    BenchmarkTools.Trial:
      memory estimate:  1.37 MiB
      allocs estimate:  12964
      --------------
      minimum time:     719.300 μs (0.00% GC)
      median time:      863.950 μs (0.00% GC)
      mean time:        996.328 μs (7.60% GC)
      maximum time:     52.938 ms (97.40% GC)
      --------------
      samples:          5004
      evals/sample:     1

```julia
@benchmark solve(prob,Tsit5(),save_everystep=false)
```

    BenchmarkTools.Trial:
      memory estimate:  6.86 KiB
      allocs estimate:  92
      --------------
      minimum time:     388.800 μs (0.00% GC)
      median time:      424.500 μs (0.00% GC)
      mean time:        456.487 μs (0.16% GC)
      maximum time:     3.320 ms (86.45% GC)
      --------------
      samples:          10000
      evals/sample:     1

There is a 4x time difference just from that change! Notice there are still some allocations and this is due to the construction of the integration cache. But this doesn't scale with the problem size:

```julia
tspan = (0.0,500.0) # 5x longer than before
prob = ODEProblem(lorenz!,u0,tspan)
@benchmark solve(prob,Tsit5(),save_everystep=false)
```

    BenchmarkTools.Trial:
      memory estimate:  6.86 KiB
      allocs estimate:  92
      --------------
      minimum time:     1.997 ms (0.00% GC)
      median time:      2.255 ms (0.00% GC)
      mean time:        2.366 ms (0.00% GC)
      maximum time:     4.923 ms (0.00% GC)
      --------------
      samples:          2108
      evals/sample:     1

since that's all just setup allocations.

#### But if the system is small we can optimize even more.

Allocations are only expensive if they are "heap allocations". For a more in-depth definition of heap allocations, [there are a lot of sources online](http://net-informations.com/faq/net/stack-heap.htm). But a good working definition is that heap allocations are variable-sized slabs of memory which have to be pointed to, and this pointer indirection costs time. Additionally, the heap has to be managed and the garbage controllers has to actively keep track of what's on the heap.

However, there's an alternative to heap allocations, known as stack allocations. The stack is statically-sized (known at compile time) and thus its accesses are quick. Additionally, the exact block of memory is known in advance by the compiler, and thus re-using the memory is cheap. This means that allocating on the stack has essentially no cost!

Arrays have to be heap allocated because their size (and thus the amount of memory they take up) is determined at runtime. But there are structures in Julia which are stack-allocated. `struct`s for example are stack-allocated "value-type"s. `Tuple`s are a stack-allocated collection. The most useful data structure for DiffEq though is the `StaticArray` from the package [StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl). These arrays have their length determined at compile-time. They are created using macros attached to normal array expressions, for example:

```julia
using StaticArrays
A = @SVector [2.0,3.0,5.0]
```

    3-element SArray{Tuple{3},Float64,1,3}:
     2.0
     3.0
     5.0

Notice that the `3` after `SVector` gives the size of the `SVector`. It cannot be changed. Additionally, `SVector`s are immutable, so we have to create a new `SVector` to change values. But remember, we don't have to worry about allocations because this data structure is stack-allocated. `SArray`s have a lot of extra optimizations as well: they have fast matrix multiplication, fast QR factorizations, etc. which directly make use of the information about the size of the array. Thus, when possible they should be used.

Unfortunately static arrays can only be used for sufficiently small arrays. After a certain size, they are forced to heap allocate after some instructions and their compile time balloons. Thus static arrays shouldn't be used if your system has more than 100 variables. Additionally, only the native Julia algorithms can fully utilize static arrays.

Let's ***optimize `lorenz` using static arrays***. Note that in this case, we want to use the out-of-place allocating form, but this time we want to output a static array:

```julia
function lorenz_static(u,p,t)
 dx = 10.0*(u[2]-u[1])
 dy = u[1]*(28.0-u[3]) - u[2]
 dz = u[1]*u[2] - (8/3)*u[3]
 @SVector [dx,dy,dz]
end
```

    lorenz_static (generic function with 1 method)

To make the solver internally use static arrays, we simply give it a static array as the initial condition:

```julia
u0 = @SVector [1.0,0.0,0.0]
tspan = (0.0,100.0)
prob = ODEProblem(lorenz_static,u0,tspan)
@benchmark solve(prob,Tsit5())
```

    BenchmarkTools.Trial:
      memory estimate:  472.02 KiB
      allocs estimate:  2662
      --------------
      minimum time:     423.700 μs (0.00% GC)
      median time:      475.600 μs (0.00% GC)
      mean time:        519.208 μs (2.59% GC)
      maximum time:     1.817 ms (48.73% GC)
      --------------
      samples:          9591
      evals/sample:     1

```julia
@benchmark solve(prob,Tsit5(),save_everystep=false)
```

    BenchmarkTools.Trial:
      memory estimate:  6.16 KiB
      allocs estimate:  73
      --------------
      minimum time:     340.901 μs (0.00% GC)
      median time:      370.400 μs (0.00% GC)
      mean time:        395.819 μs (0.13% GC)
      maximum time:     2.881 ms (86.59% GC)
      --------------
      samples:          10000
      evals/sample:     1

And that's pretty much all there is to it. With static arrays you don't have to worry about allocating, so use operations like `*` and don't worry about fusing operations (discussed in the next section). Do "the vectorized code" of R/MATLAB/Python and your code in this case will be fast, or directly use the numbers/values.

#### Exercise 1

Implement the out-of-place array, in-place array, and out-of-place static array forms for the [Henon-Heiles System](https://en.wikipedia.org/wiki/H%C3%A9non%E2%80%93Heiles_system) and time the results.

## Optimizing Large Systems

### Interlude: Managing Allocations with Broadcast Fusion

When your system is sufficiently large, or you have to make use of a non-native Julia algorithm, you have to make use of `Array`s. In order to use arrays in the most efficient manner, you need to be careful about temporary allocations. Vectorized calculations naturally have plenty of temporary array allocations. This is because a vectorized calculation outputs a vector. Thus:

```julia
A = rand(1000,1000); B = rand(1000,1000); C = rand(1000,1000)
test(A,B,C) = A + B + C
@benchmark test(A,B,C)
```

    BenchmarkTools.Trial:
      memory estimate:  7.63 MiB
      allocs estimate:  3
      --------------
      minimum time:     3.257 ms (0.00% GC)
      median time:      4.102 ms (0.00% GC)
      mean time:        4.647 ms (17.26% GC)
      maximum time:     60.949 ms (92.17% GC)
      --------------
      samples:          1074
      evals/sample:     1

That expression `A + B + C` creates 2 arrays. It first creates one for the output of `A + B`, then uses that result array to `+ C` to get the final result. 2 arrays! We don't want that! The first thing to do to fix this is to use broadcast fusion. [Broadcast fusion](https://julialang.org/blog/2017/01/moredots) puts expressions together. For example, instead of doing the `+` operations separately, if we were to add them all at the same time, then we would only have a single array that's created. For example:

```julia
test2(A,B,C) = map((a,b,c)->a+b+c,A,B,C)
@benchmark test2(A,B,C)
```

    BenchmarkTools.Trial:
      memory estimate:  7.63 MiB
      allocs estimate:  5
      --------------
      minimum time:     4.263 ms (0.00% GC)
      median time:      5.274 ms (0.00% GC)
      mean time:        5.760 ms (14.21% GC)
      maximum time:     60.352 ms (91.60% GC)
      --------------
      samples:          867
      evals/sample:     1

Puts the whole expression into a single function call, and thus only one array is required to store output. This is the same as writing the loop:

```julia
function test3(A,B,C)
    D = similar(A)
    @inbounds for i in eachindex(A)
        D[i] = A[i] + B[i] + C[i]
    end
    D
end
@benchmark test3(A,B,C)
```

    BenchmarkTools.Trial:
      memory estimate:  7.63 MiB
      allocs estimate:  2
      --------------
      minimum time:     3.068 ms (0.00% GC)
      median time:      4.069 ms (0.00% GC)
      mean time:        4.573 ms (17.14% GC)
      maximum time:     60.806 ms (93.12% GC)
      --------------
      samples:          1091
      evals/sample:     1

However, Julia's broadcast is syntactic sugar for this. If multiple expressions have a `.`, then it will put those vectorized operations together. Thus:

```julia
test4(A,B,C) = A .+ B .+ C
@benchmark test4(A,B,C)
```

    BenchmarkTools.Trial:
      memory estimate:  7.63 MiB
      allocs estimate:  2
      --------------
      minimum time:     3.100 ms (0.00% GC)
      median time:      4.090 ms (0.00% GC)
      mean time:        4.611 ms (17.00% GC)
      maximum time:     60.075 ms (91.42% GC)
      --------------
      samples:          1083
      evals/sample:     1

is a version with only 1 array created (the output). Note that `.`s can be used with function calls as well:

```julia
sin.(A) .+ sin.(B)
```

    1000×1000 Array{Float64,2}:
     1.03364    0.98507   1.27358    0.877     …  0.573536  1.23032   0.911319
     1.63118    1.53896   0.71938    0.982407     1.18889   0.666447  0.456239
     1.326      0.635384  1.00153    0.872837     0.57143   0.957094  0.644016
     1.20394    0.252596  0.586836   0.706981     0.484221  1.16159   0.320898
     0.602921   0.942712  0.245305   1.04536      0.811623  1.01748   1.49487
     0.967846   0.933214  0.425036   0.763434  …  0.339187  0.863011  0.697906
     0.672935   0.702386  0.815461   1.4289       1.43948   1.23604   0.991722
     0.57422    1.16769   0.553201   0.327915     0.303723  0.17508   1.62917
     0.16162    1.48327   0.394958   1.32529      0.73105   0.559688  0.996632
     0.664391   0.808939  1.36914    0.981928     0.702237  0.539255  1.432
     0.811794   1.08456   0.632371   1.07589   …  1.45691   1.66896   0.456393
     0.951618   0.533518  0.78531    1.27613      0.887392  0.883319  1.43496
     1.14518    1.27886   0.719942   0.939989     0.399901  0.768837  0.314349
     ⋮                                         ⋱
     1.35936    1.06333   0.885049   1.31827      1.1435    0.501145  0.358912
     1.30287    1.31587   1.1749     0.975842     1.61852   1.10175   0.68206
     0.749153   1.35504   1.4116     1.15027   …  1.20905   0.864016  1.1477
     1.27482    0.576958  1.17769    0.77052      0.532696  0.833744  1.28555
     1.04075    0.844666  0.360267   1.27131      0.679665  0.533756  1.4078
     0.674019   0.831317  1.5796     0.122812     0.585702  1.10292   1.46805
     0.921581   1.47563   1.16279    1.26375      0.535886  1.04286   0.698426
     0.508372   1.27464   0.0699981  1.29666   …  1.59411   1.02171   1.14568
     0.0579072  0.884255  0.277504   0.957518     0.521798  1.38383   1.38164
     0.958096   1.00757   1.52481    1.49831      1.37391   0.829876  1.52642
     0.627583   1.14951   0.796797   1.43778      1.17557   0.490823  0.480095
     0.234199   0.78984   1.28475    1.36424      1.11589   0.335778  0.113732

Also, the `@.` macro applys a dot to every operator:

```julia
test5(A,B,C) = @. A + B + C #only one array allocated
@benchmark test5(A,B,C)
```

    BenchmarkTools.Trial:
      memory estimate:  7.63 MiB
      allocs estimate:  3
      --------------
      minimum time:     3.213 ms (0.00% GC)
      median time:      4.115 ms (0.00% GC)
      mean time:        4.570 ms (15.95% GC)
      maximum time:     8.086 ms (38.82% GC)
      --------------
      samples:          1092
      evals/sample:     1

Using these tools we can get rid of our intermediate array allocations for many vectorized function calls. But we are still allocating the output array. To get rid of that allocation, we can instead use mutation. Mutating broadcast is done via `.=`. For example, if we pre-allocate the output:

```julia
D = zeros(1000,1000);
```

Then we can keep re-using this cache for subsequent calculations. The mutating broadcasting form is:

```julia
test6!(D,A,B,C) = D .= A .+ B .+ C #only one array allocated
@benchmark test6!(D,A,B,C)
```

    BenchmarkTools.Trial:
      memory estimate:  0 bytes
      allocs estimate:  0
      --------------
      minimum time:     799.699 μs (0.00% GC)
      median time:      902.050 μs (0.00% GC)
      mean time:        962.816 μs (0.00% GC)
      maximum time:     3.425 ms (0.00% GC)
      --------------
      samples:          5152
      evals/sample:     1

If we use `@.` before the `=`, then it will turn it into `.=`:

```julia
test7!(D,A,B,C) = @. D = A + B + C #only one array allocated
@benchmark test7!(D,A,B,C)
```

    BenchmarkTools.Trial:
      memory estimate:  0 bytes
      allocs estimate:  0
      --------------
      minimum time:     798.400 μs (0.00% GC)
      median time:      892.550 μs (0.00% GC)
      mean time:        950.436 μs (0.00% GC)
      maximum time:     2.693 ms (0.00% GC)
      --------------
      samples:          5218
      evals/sample:     1

Notice that in this case, there is no "output", and instead the values inside of `D` are what are changed (like with the DiffEq inplace function). Many Julia functions have a mutating form which is denoted with a `!`. For example, the mutating form of the `map` is `map!`:

```julia
test8!(D,A,B,C) = map!((a,b,c)->a+b+c,D,A,B,C)
@benchmark test8!(D,A,B,C)
```

    BenchmarkTools.Trial:
      memory estimate:  32 bytes
      allocs estimate:  1
      --------------
      minimum time:     1.049 ms (0.00% GC)
      median time:      1.167 ms (0.00% GC)
      mean time:        1.238 ms (0.00% GC)
      maximum time:     3.015 ms (0.00% GC)
      --------------
      samples:          4011
      evals/sample:     1

Some operations require using an alternate mutating form in order to be fast. For example, matrix multiplication via `*` allocates a temporary:

```julia
@benchmark A*B
```

    BenchmarkTools.Trial:
      memory estimate:  7.63 MiB
      allocs estimate:  2
      --------------
      minimum time:     12.283 ms (0.00% GC)
      median time:      13.676 ms (0.00% GC)
      mean time:        14.268 ms (5.67% GC)
      maximum time:     20.365 ms (12.72% GC)
      --------------
      samples:          350
      evals/sample:     1

Instead, we can use the mutating form `mul!` into a cache array to avoid allocating the output:

```julia
using LinearAlgebra
@benchmark mul!(D,A,B) # same as D = A * B
```

    BenchmarkTools.Trial:
      memory estimate:  0 bytes
      allocs estimate:  0
      --------------
      minimum time:     11.668 ms (0.00% GC)
      median time:      12.447 ms (0.00% GC)
      mean time:        12.623 ms (0.00% GC)
      maximum time:     16.679 ms (0.00% GC)
      --------------
      samples:          396
      evals/sample:     1

For repeated calculations this reduced allocation can stop GC cycles and thus lead to more efficient code. Additionally, ***we can fuse together higher level linear algebra operations using BLAS***. The package [SugarBLAS.jl](https://github.com/lopezm94/SugarBLAS.jl) makes it easy to write higher level operations like `alpha*B*A + beta*C` as mutating BLAS calls.

### Example Optimization: Gierer-Meinhardt Reaction-Diffusion PDE Discretization

Let's optimize the solution of a Reaction-Diffusion PDE's discretization. In its discretized form, this is the ODE:

$$
\begin{align}
du &= D_1 (A_y u + u A_x) + \frac{au^2}{v} + \bar{u} - \alpha u\\
dv &= D_2 (A_y v + v A_x) + a u^2 + \beta v
\end{align}
$$

where $u$, $v$, and $A$ are matrices. Here, we will use the simplified version where $A$ is the tridiagonal stencil $[1,-2,1]$, i.e. it's the 2D discretization of the LaPlacian. The native code would be something along the lines of:

```julia
# Generate the constants
p = (1.0,1.0,1.0,10.0,0.001,100.0) # a,α,ubar,β,D1,D2
N = 100
Ax = Array(Tridiagonal([1.0 for i in 1:N-1],[-2.0 for i in 1:N],[1.0 for i in 1:N-1]))
Ay = copy(Ax)
Ax[2,1] = 2.0
Ax[end-1,end] = 2.0
Ay[1,2] = 2.0
Ay[end,end-1] = 2.0

function basic_version!(dr,r,p,t)
  a,α,ubar,β,D1,D2 = p
  u = r[:,:,1]
  v = r[:,:,2]
  Du = D1*(Ay*u + u*Ax)
  Dv = D2*(Ay*v + v*Ax)
  dr[:,:,1] = Du .+ a.*u.*u./v .+ ubar .- α*u
  dr[:,:,2] = Dv .+ a.*u.*u .- β*v
end

a,α,ubar,β,D1,D2 = p
uss = (ubar+β)/α
vss = (a/β)*uss^2
r0 = zeros(100,100,2)
r0[:,:,1] .= uss.+0.1.*rand.()
r0[:,:,2] .= vss

prob = ODEProblem(basic_version!,r0,(0.0,0.1),p)
```

    ODEProblem with uType Array{Float64,3} and tType Float64. In-place: true
    timespan: (0.0, 0.1)
    u0: [11.0394 11.002 … 11.013 11.0926; 11.0837 11.0584 … 11.0908 11.0092; … ; 11.0918 11.0966 … 11.0582 11.0298; 11.0307 11.0959 … 11.0378 11.0558]

    [12.1 12.1 … 12.1 12.1; 12.1 12.1 … 12.1 12.1; … ; 12.1 12.1 … 12.1 12.1; 12.1 12.1 … 12.1 12.1]

In this version we have encoded our initial condition to be a 3-dimensional array, with `u[:,:,1]` being the `A` part and `u[:,:,2]` being the `B` part.

```julia
@benchmark solve(prob,Tsit5())
```

    BenchmarkTools.Trial:
      memory estimate:  186.88 MiB
      allocs estimate:  8589
      --------------
      minimum time:     80.334 ms (8.83% GC)
      median time:      167.756 ms (48.87% GC)
      mean time:        137.745 ms (37.82% GC)
      maximum time:     192.902 ms (47.32% GC)
      --------------
      samples:          37
      evals/sample:     1

While this version isn't very efficient,

#### We recommend writing the "high-level" code first, and iteratively optimizing it!

The first thing that we can do is get rid of the slicing allocations. The operation `r[:,:,1]` creates a temporary array instead of a "view", i.e. a pointer to the already existing memory. To make it a view, add `@view`. Note that we have to be careful with views because they point to the same memory, and thus changing a view changes the original values:

```julia
A = rand(4)
@show A
B = @view A[1:3]
B[2] = 2
@show A
```

    A = [0.22581, 0.809546, 0.344768, 0.640214]
    A = [0.22581, 2.0, 0.344768, 0.640214]

    4-element Array{Float64,1}:
     0.2258101941951307
     2.0
     0.3447681198957038
     0.6402135416553185

Notice that changing `B` changed `A`. This is something to be careful of, but at the same time we want to use this since we want to modify the output `dr`. Additionally, the last statement is a purely element-wise operation, and thus we can make use of broadcast fusion there. Let's rewrite `basic_version!` to ***avoid slicing allocations*** and to ***use broadcast fusion***:

```julia
function gm2!(dr,r,p,t)
  a,α,ubar,β,D1,D2 = p
  u = @view r[:,:,1]
  v = @view r[:,:,2]
  du = @view dr[:,:,1]
  dv = @view dr[:,:,2]
  Du = D1*(Ay*u + u*Ax)
  Dv = D2*(Ay*v + v*Ax)
  @. du = Du + a.*u.*u./v + ubar - α*u
  @. dv = Dv + a.*u.*u - β*v
end
prob = ODEProblem(gm2!,r0,(0.0,0.1),p)
@benchmark solve(prob,Tsit5())
```

    BenchmarkTools.Trial:
      memory estimate:  119.55 MiB
      allocs estimate:  7119
      --------------
      minimum time:     69.464 ms (6.19% GC)
      median time:      119.849 ms (35.84% GC)
      mean time:        119.770 ms (37.59% GC)
      maximum time:     175.505 ms (55.22% GC)
      --------------
      samples:          42
      evals/sample:     1

Now, most of the allocations are taking place in `Du = D1*(Ay*u + u*Ax)` since those operations are vectorized and not mutating. We should instead replace the matrix multiplications with `mul!`. When doing so, we will need to have cache variables to write into. This looks like:

```julia
Ayu = zeros(N,N)
uAx = zeros(N,N)
Du = zeros(N,N)
Ayv = zeros(N,N)
vAx = zeros(N,N)
Dv = zeros(N,N)
function gm3!(dr,r,p,t)
  a,α,ubar,β,D1,D2 = p
  u = @view r[:,:,1]
  v = @view r[:,:,2]
  du = @view dr[:,:,1]
  dv = @view dr[:,:,2]
  mul!(Ayu,Ay,u)
  mul!(uAx,u,Ax)
  mul!(Ayv,Ay,v)
  mul!(vAx,v,Ax)
  @. Du = D1*(Ayu + uAx)
  @. Dv = D2*(Ayv + vAx)
  @. du = Du + a*u*u./v + ubar - α*u
  @. dv = Dv + a*u*u - β*v
end
prob = ODEProblem(gm3!,r0,(0.0,0.1),p)
@benchmark solve(prob,Tsit5())
```

    BenchmarkTools.Trial:
      memory estimate:  29.76 MiB
      allocs estimate:  5355
      --------------
      minimum time:     54.087 ms (2.33% GC)
      median time:      57.128 ms (2.28% GC)
      mean time:        59.910 ms (5.30% GC)
      maximum time:     142.458 ms (51.42% GC)
      --------------
      samples:          84
      evals/sample:     1

But our temporary variables are global variables. We need to either declare the caches as `const` or localize them. We can localize them by adding them to the parameters, `p`. It's easier for the compiler to reason about local variables than global variables. ***Localizing variables helps to ensure type stability***.

```julia
p = (1.0,1.0,1.0,10.0,0.001,100.0,Ayu,uAx,Du,Ayv,vAx,Dv) # a,α,ubar,β,D1,D2
function gm4!(dr,r,p,t)
  a,α,ubar,β,D1,D2,Ayu,uAx,Du,Ayv,vAx,Dv = p
  u = @view r[:,:,1]
  v = @view r[:,:,2]
  du = @view dr[:,:,1]
  dv = @view dr[:,:,2]
  mul!(Ayu,Ay,u)
  mul!(uAx,u,Ax)
  mul!(Ayv,Ay,v)
  mul!(vAx,v,Ax)
  @. Du = D1*(Ayu + uAx)
  @. Dv = D2*(Ayv + vAx)
  @. du = Du + a*u*u./v + ubar - α*u
  @. dv = Dv + a*u*u - β*v
end
prob = ODEProblem(gm4!,r0,(0.0,0.1),p)
@benchmark solve(prob,Tsit5())
```

    BenchmarkTools.Trial:
      memory estimate:  29.66 MiB
      allocs estimate:  1090
      --------------
      minimum time:     46.933 ms (2.46% GC)
      median time:      49.269 ms (2.41% GC)
      mean time:        51.746 ms (5.35% GC)
      maximum time:     137.551 ms (57.03% GC)
      --------------
      samples:          97
      evals/sample:     1

We could then use the BLAS `gemmv` to optimize the matrix multiplications some more, but instead let's devectorize the stencil.

```julia
p = (1.0,1.0,1.0,10.0,0.001,100.0,N)
function fast_gm!(du,u,p,t)
  a,α,ubar,β,D1,D2,N = p

  @inbounds for j in 2:N-1, i in 2:N-1
    du[i,j,1] = D1*(u[i-1,j,1] + u[i+1,j,1] + u[i,j+1,1] + u[i,j-1,1] - 4u[i,j,1]) +
              a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end

  @inbounds for j in 2:N-1, i in 2:N-1
    du[i,j,2] = D2*(u[i-1,j,2] + u[i+1,j,2] + u[i,j+1,2] + u[i,j-1,2] - 4u[i,j,2]) +
            a*u[i,j,1]^2 - β*u[i,j,2]
  end

  @inbounds for j in 2:N-1
    i = 1
    du[1,j,1] = D1*(2u[i+1,j,1] + u[i,j+1,1] + u[i,j-1,1] - 4u[i,j,1]) +
            a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end
  @inbounds for j in 2:N-1
    i = 1
    du[1,j,2] = D2*(2u[i+1,j,2] + u[i,j+1,2] + u[i,j-1,2] - 4u[i,j,2]) +
            a*u[i,j,1]^2 - β*u[i,j,2]
  end
  @inbounds for j in 2:N-1
    i = N
    du[end,j,1] = D1*(2u[i-1,j,1] + u[i,j+1,1] + u[i,j-1,1] - 4u[i,j,1]) +
           a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end
  @inbounds for j in 2:N-1
    i = N
    du[end,j,2] = D2*(2u[i-1,j,2] + u[i,j+1,2] + u[i,j-1,2] - 4u[i,j,2]) +
           a*u[i,j,1]^2 - β*u[i,j,2]
  end

  @inbounds for i in 2:N-1
    j = 1
    du[i,1,1] = D1*(u[i-1,j,1] + u[i+1,j,1] + 2u[i,j+1,1] - 4u[i,j,1]) +
              a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end
  @inbounds for i in 2:N-1
    j = 1
    du[i,1,2] = D2*(u[i-1,j,2] + u[i+1,j,2] + 2u[i,j+1,2] - 4u[i,j,2]) +
              a*u[i,j,1]^2 - β*u[i,j,2]
  end
  @inbounds for i in 2:N-1
    j = N
    du[i,end,1] = D1*(u[i-1,j,1] + u[i+1,j,1] + 2u[i,j-1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end
  @inbounds for i in 2:N-1
    j = N
    du[i,end,2] = D2*(u[i-1,j,2] + u[i+1,j,2] + 2u[i,j-1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]
  end

  @inbounds begin
    i = 1; j = 1
    du[1,1,1] = D1*(2u[i+1,j,1] + 2u[i,j+1,1] - 4u[i,j,1]) +
              a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[1,1,2] = D2*(2u[i+1,j,2] + 2u[i,j+1,2] - 4u[i,j,2]) +
              a*u[i,j,1]^2 - β*u[i,j,2]

    i = 1; j = N
    du[1,N,1] = D1*(2u[i+1,j,1] + 2u[i,j-1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[1,N,2] = D2*(2u[i+1,j,2] + 2u[i,j-1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]

    i = N; j = 1
    du[N,1,1] = D1*(2u[i-1,j,1] + 2u[i,j+1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[N,1,2] = D2*(2u[i-1,j,2] + 2u[i,j+1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]

    i = N; j = N
    du[end,end,1] = D1*(2u[i-1,j,1] + 2u[i,j-1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[end,end,2] = D2*(2u[i-1,j,2] + 2u[i,j-1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]
   end
end
prob = ODEProblem(fast_gm!,r0,(0.0,0.1),p)
@benchmark solve(prob,Tsit5())
```

    BenchmarkTools.Trial:
      memory estimate:  29.63 MiB
      allocs estimate:  505
      --------------
      minimum time:     7.409 ms (8.40% GC)
      median time:      8.683 ms (8.87% GC)
      mean time:        9.342 ms (12.98% GC)
      maximum time:     91.148 ms (78.92% GC)
      --------------
      samples:          535
      evals/sample:     1

Lastly, we can do other things like multithread the main loops, but these optimizations get the last 2x-3x out. The main optimizations which apply everywhere are the ones we just performed (though the last one only works if your matrix is a stencil. This is known as a matrix-free implementation of the PDE discretization).

This gets us to about 8x faster than our original MATLAB/SciPy/R vectorized style code!

The last thing to do is then ***optimize our algorithm choice***. We have been using `Tsit5()` as our test algorithm, but in reality this problem is a stiff PDE discretization and thus one recommendation is to use `CVODE_BDF()`. However, instead of using the default dense Jacobian, we should make use of the sparse Jacobian afforded by the problem. The Jacobian is the matrix $\frac{df_i}{dr_j}$, where $r$ is read by the linear index (i.e. down columns). But since the $u$ variables depend on the $v$, the band size here is large, and thus this will not do well with a Banded Jacobian solver. Instead, we utilize sparse Jacobian algorithms. `CVODE_BDF` allows us to use a sparse Newton-Krylov solver by setting `linear_solver = :GMRES` (see [the solver documentation](http://docs.juliadiffeq.org/latest/solvers/ode_solve.html#Sundials.jl-1), and thus we can solve this problem efficiently. Let's see how this scales as we increase the integration time.

```julia
prob = ODEProblem(fast_gm!,r0,(0.0,10.0),p)
@benchmark solve(prob,Tsit5())
```

    BenchmarkTools.Trial:
      memory estimate:  2.76 GiB
      allocs estimate:  41689
      --------------
      minimum time:     2.255 s (13.38% GC)
      median time:      2.399 s (19.28% GC)
      mean time:        2.826 s (30.95% GC)
      maximum time:     3.823 s (48.64% GC)
      --------------
      samples:          3
      evals/sample:     1

```julia
using Sundials
@benchmark solve(prob,CVODE_BDF(linear_solver=:GMRES))
```

    BenchmarkTools.Trial:
      memory estimate:  306.28 MiB
      allocs estimate:  87141
      --------------
      minimum time:     1.819 s (0.00% GC)
      median time:      1.831 s (4.41% GC)
      mean time:        1.889 s (4.83% GC)
      maximum time:     2.018 s (9.57% GC)
      --------------
      samples:          3
      evals/sample:     1

```julia
prob = ODEProblem(fast_gm!,r0,(0.0,100.0),p)
# Will go out of memory if we don't turn off `save_everystep`!
@benchmark solve(prob,Tsit5(),save_everystep=false)
```

    BenchmarkTools.Trial:
      memory estimate:  2.91 MiB
      allocs estimate:  113
      --------------
      minimum time:     5.422 s (0.00% GC)
      median time:      5.422 s (0.00% GC)
      mean time:        5.422 s (0.00% GC)
      maximum time:     5.422 s (0.00% GC)
      --------------
      samples:          1
      evals/sample:     1

```julia
@benchmark solve(prob,CVODE_BDF(linear_solver=:GMRES))
```

    BenchmarkTools.Trial:
      memory estimate:  306.28 MiB
      allocs estimate:  87141
      --------------
      minimum time:     1.849 s (4.01% GC)
      median time:      1.871 s (3.96% GC)
      mean time:        1.893 s (4.25% GC)
      maximum time:     1.959 s (8.54% GC)
      --------------
      samples:          3
      evals/sample:     1

Now let's check the allocation growth.

```julia
@benchmark solve(prob,CVODE_BDF(linear_solver=:GMRES),save_everystep=false)
```

    BenchmarkTools.Trial:
      memory estimate:  4.07 MiB
      allocs estimate:  77881
      --------------
      minimum time:     1.646 s (0.00% GC)
      median time:      1.694 s (0.00% GC)
      mean time:        1.705 s (0.00% GC)
      maximum time:     1.774 s (0.00% GC)
      --------------
      samples:          3
      evals/sample:     1

```julia
prob = ODEProblem(fast_gm!,r0,(0.0,500.0),p)
@benchmark solve(prob,CVODE_BDF(linear_solver=:GMRES),save_everystep=false)
```

    BenchmarkTools.Trial:
      memory estimate:  5.31 MiB
      allocs estimate:  108359
      --------------
      minimum time:     2.306 s (0.00% GC)
      median time:      2.337 s (0.00% GC)
      mean time:        2.340 s (0.00% GC)
      maximum time:     2.375 s (0.00% GC)
      --------------
      samples:          3
      evals/sample:     1

Notice that we've elimated almost all allocations, allowing the code to grow without hitting garbage collection and slowing down.

Why is `CVODE_BDF` doing well? What's happening is that, because the problem is stiff, the number of steps required by the explicit Runge-Kutta method grows rapidly, whereas `CVODE_BDF` is taking large steps. Additionally, the `GMRES` linear solver form is quite an efficient way to solve the implicit system in this case. This is problem-dependent, and in many cases using a Krylov method effectively requires a preconditioner, so you need to play around with testing other algorithms and linear solvers to find out what works best with your problem.

## Conclusion

Julia gives you the tools to optimize the solver "all the way", but you need to make use of it. The main thing to avoid is temporary allocations. For small systems, this is effectively done via static arrays. For large systems, this is done via in-place operations and cache arrays. Either way, the resulting solution can be immensely sped up over vectorized formulations by using these principles.


# BatchBALD

# Introduction

BALD is useful in the setting where only one sample is acquired at a time; however, when multiple samples are acquired at once, BALD would select many similar points as it has no "memory" that it had acquired previous points (assuming the usual method where the top-b scoring BALD points are taken). BatchBALD considers the _joint_ mutual information of all the points in the batch. Further, because it is _submodular_, it is at worst a $1 - \frac{1}{\epsilon}$ approximation of the optimal choice when a greedy algorithm is used to choose the batch (to avoid exponential search). Hence, this paper's contributions is threefold:

1. derivation of BatchBALD.
2. a greedy algorithm to select the batch of $b$ points effeciently.
3. estimator for BatchBALD that scales to larger acquisition sizes and datasets with many classes.

# Background

## BALD
BALD estimates the mutual information (MI) between the model's predictions on the unlabelled pool and the model's parameters. Finding out the true label of the model's predictions with high MI should inform us of the true model parameters:

$$
\mathbb{I}[y; \mathbf{\omega}\ |\ \mathbf x, \mathcal{D}_{train}] = \mathbb{H}(y\ |\ \mathbf x, \mathcal{D}_{train}) - \mathbb{E}_{p(\mathbf \omega|\mathcal{D}_{train})}[\mathbb{H}(y\ |\ \mathbf x, \mathbf \omega, \mathcal{D}_{train})]
$$

To maximise mutual information, the left term must be high (implies that the model is uncertain in its predictions over all possible parameters $\omega$) and the right term must be low (implies that the model is very certain in its predictions for each _individual_ realisation of parameters $\omega$). Consider the extreme case where given $k$ realisations of $\omega$ and $k$ classes, each of the $k$ predictions have softmax probability of 1 at different classes (i.e. a diagonal matrix of shape $k$x$k$).

BALD is not suitable for deep learning as acquiring one point at a time is terribly inefficient and costly. A simple heuristic is to choose b points at once; this comes at the cost of batch diversity as the top $b$ points are likely to be similar.

Furthermore, increasing the number of inference samples (MC dropout samples) may actually deteriorate the performance (think: the more accurate the BALD scores are, the more likely it is going to pick points that are similar in a single acquisition round for $b > 1$). As noted in section 5:
> ... the lack of diversity in batches of data points acquired using the BALD objective has been noted by Janz et al., who propose to resolve it by _limiting the number of MC dropout samples and relying on noisy estimates_.

## Bayesian neural networks

Performing exact inference in BNNs is intractable for any non-trivial model; MC dropout is used as variational approximation.

# Methods

## BatchBALD

In essence, to avoid double counting, BatchBALD estimates the mutual information of _joint data points_ and the model parameters (omitting conditioning on $x_1, \dots, x_n$ and $\mathcal{D}_{train}$ in both terms on RHS and $p(\omega)$ denotes $p(\omega|\mathcal{D}_{train}))$:

$$
\mathbb{I}[y_1, y_2, \dots, y_n; \mathbf{\omega}\ |\ x_1, x_2, \dots, x_n, \mathcal{D}_{train}] = \mathbb{H}(y_1, \dots, y_n) - \mathbb{E}_{p(\omega)}[\mathbb{H}(y_1, \dots, y_n | \omega)]
$$

where:

$$
\begin{align}
\mathbb{E}_{p(\omega)}[\mathbb{H}(y_1, \dots, y_n | \omega)] &= \displaystyle\sum_{i=1}^{n}\mathbb{E}_{p(\omega)}[\mathbb{H}(y_i|\omega)] \ \ \ \text{($y_i$s are independent conditioned on $\omega$)} \\
&\approx \frac{1}{k}\displaystyle\sum_{i=1}^{n}\sum_{j=1}^{k}\mathbb{H}(y_i|\hat \omega_j)\ \ \text{($k$ monte-carlo estimates $\hat \omega_j \sim p(\omega)$)}
\end{align}
$$

and

$$
\begin{align}
\mathbb{H}(y_1, \dots, y_n) &= -\displaystyle\sum_{y_1}\dots\sum_{y_n}p(y_1, \dots, y_n)log\ p(y_1, \dots, y_n) \\
&\approx  -\displaystyle\sum_{y_1}\dots\sum_{y_n}\left(\frac{1}{k}\sum_{j=1}^{k}p(y_1, \dots, y_n | \hat \omega_j)\right)log\ \left(\frac{1}{k}\sum_{j=1}^{k}p(y_1, \dots, y_n | \hat \omega_j)\right) \ \ \text{(1)}
\end{align}
$$

where in (1), $p(y_1, \dots, y_n) = \int p(y_1, \dots, y_n | \omega)\ p(\omega)\ d\omega$ is approximated using monte-carlo estimate.

## Greedy algorithm

Selecting the best batch $\{x_1, \dots, x_B\}^*$ from $\mathcal{D}_{pool}$ according to $\alpha_{BatchBALD}$ requires an exponential running time algorithm. However, we could settle for a greedy $\left(1 - \frac{1}{e}\right)$-approximate algorithm with a guarantee that:
$$\alpha_{BatchBALD}(\{x_1, \dots, x_B\}^{greedy}, p(\omega | \mathcal{D}_{train})) \geq 0.63 \cdot \alpha_{BatchBALD}(\{x_1, \dots, x_B\}^*, p(\omega|\mathcal{D}_{train}))$$

where $0.63 \approx 1 - \frac{1}{e}$, hence, the name. The greedy algorithm below only has this guarantee if $\alpha_{BatchBALD}$ is submodular, which is proven in appendix C. Therefore,

1. $A_0 \leftarrow \varnothing$
2. for $i = 1 \dots b$:
3.        $A_i \leftarrow A_{i-1} \cup \underset{ x \in \mathcal{D}_{pool} \backslash A_{i-1}}{\text{argmax}} \alpha_{BatchBALD}(A_{i-1} \cup \{x\})$
4. **Return** $A_b$

is $\left(1 - \frac{1}{e}\right)$-approximate. In practice these greedy solutions are often within 90% of optimal (Krause, 2008).

## Efficient implementation

In the above greedy algorithm, notice that $x_1, \dots, x_{n-1}$ ($A_{n-1}$) are always fixed when greedily searching for $x_n$. Define:

$$
\begin{align}
\left(\hat P_{1:n-1}\right)_{i, j} &= p(y_1^{(i)}, \dots, y_{n-1}^{(i)} | \hat \omega_j) \ \ &\text{(matrix of shape $\mathbb{R}^{c^{n-1}\times k}$)} \\
\left(\hat P_{n}\right)_{i, j} &= p(y_n^{(i)} | \hat \omega_j) \ \ &\text{(matrix of shape $\mathbb{R}^{c\times k}$)} \\
\frac{1}{k}\displaystyle\sum_{j=1}^{k}p(y_1, \dots, y_n | \hat \omega_j) &= \frac{1}{k} \displaystyle \sum_{j=1}^{k} p(y_1, \dots, y_{n-1} | \hat \omega_j)p(y_n | \hat \omega_j) = \left(\frac{1}{k}\hat P_{1:n-1}\hat P^{T}_{n}\right)_{y_{1:n-1}, y_n}\ \ &\text{(matrix of shape $\mathbb{R}^{c^{n-1}\times c}$)}
\end{align}
$$

where $p(y_1^{(i)}, \dots, y_{n-1}^{(i)} | \hat \omega_j)$ denotes the $i^{th}$ permutation of the random vector $\mathbf y_{1:n-1}^{(i)}$ (total of $c^{n-1}$ possible permutations) and the $j^{th}$ MC dropout sample.

So:

$$
\mathbb{H}(y_1, \dots, y_n) = \sum_{i}^{c^{n-1}}\sum_{j}^{c}\left(\left[\frac{1}{k}\hat P_{1:n-1}\hat P^{T}_{n}\right] \odot log\left[\frac{1}{k}\hat P_{1:n-1}\hat P^{T}_{n}\right]\right)_{i,j} (2)
$$

Therefore, in each iteration of the greedy algorithm, $\hat P_{1:n-1}$ only needs to be computed once. When finding the $n^{th}$ addition ($x_n \in \mathcal{D}_{pool}$ \ $A_{n-1}$), $l$ candidates can be considered at once by using batch matrix multiplication (i.e. concatenating mutiple candidate $x_n$s to the rows of $\hat P_n \in \mathbb{R}^{lc \times k}$).

### The problem, and the solution

#### Problem
However, $c^{n-1}$ is quite memory-demanding; a reasonable dataset with 10 classes and batch size 10 has 10 giga-rows!

#### Solution
To overcome intractable $c^{n-1}$ configurations, the authors proposed MC-sampling $m$ configurations as an _approximation_ instead. Equation 2 now becomes:

$$
\mathbb{H}(y_1, \dots, y_n) \approx -\frac{1}{m} \sum_{i}^{m}\sum_{j}^{c}\left(\left[\frac{\hat P_{1:n-1}\hat P^{T}_{n}}{\hat P_{1:n-1} \mathbb{1}_{k,1} }\right] \odot log\left[\frac{1}{k}\hat P_{1:n-1}\hat P^{T}_{n}\right]\right)_{i,j}
$$

> See appendix C for working.

TODO: assuming $p(y_1, \dots, y_n | \hat \omega_j)$ factorises to $\prod_i^n p(y_i | \hat \omega_j)$, wouldn't the product of probabilities underflow for large $n$? (Logspace?)

# Experiments

## Practical considerations
* Labelled dataset is usually small in the beginning (e.g. 20 points); it is important to avoid overfitting. The authors used _early stopping after 3 epochs of declining accuracy on the validation set_.
* The models are reinitialised after each acquisition rounds. Empirically, it helped the model improve even when very small batches are acquired & decorrelates subsequent acquisitions.
* _Dropout masks_ are kept consistent. This reduces variance and is necessary to capture dependencies between the inputs of BatchBALD ($\mathbb{I}$[$\mathbf \omega$ ; $y_1, \dots, y_n$]) since the dropout masks directly affect $\mathbf \omega$. In other words, we ensure that we're taking the MI of the same set of model parameters $\mathbf \omega$ each time.
* In the following experiments, the authors used either 10,000 MC samples ($m$ above) or directly computed the joint entropy by enumerating all $c^{n-1}$ configurations (they used MC samples when $b > 4$).

## Repeated MNIST
The MNIST dataset is replicated twice (yielding 3x the size) where the duplicates are normalised and isotropic gaussian noise ($\sigma = 0.1$) added to it.

BatchBALD had no problems overcoming the replicated dataset; BALD performs worse than random acquisition.

(acquisition batch size = 10; 10 MC dropout samples)

## MNIST

* Increasing batch acquisision size for BALD quickly deteriorates its performance whilst BatchBALD maintains its performance (only slightly worse when batch size is 40, the authors postulates that it's due to estimator noise)
* BALD with acquisition size 1 and BatchBALD with acquisition size 10 performs very similarly, unsurprisingly outperforming BALD with size 10. Further, the time needed to reach 95% accuracy using BALD with size 1 is significantly longer than size 10, of which BatchBALD is only slightly slower than BALD.

(varying acquisition batch size; 100 MC dropout samples)

## EMNIST

* BatchBALD > random acquisition > BALD
* Entropy (of categorical distribution of classes) of BatchBALD is much higher than BALD, implying a more diverse set of acquired points. (cf. ICAL's postulation about lower entropy when number of classes increases)

(__no initial dataset__, model randomly initialised when performing the initial acquisition step; acquisition batch size = 5; 10 MC dropout samples)

## CINIC-10

* BatchBALD > BALD

(pretrained ImageNet model; acquisition batch size = 10; 50 MC dropout samples)

## Comments

The authors mostly used acquisition batch size of less than 40 in all experiments while ICAL's authors used generous acquisition sizes of up to 3000 (CIFAR-10/100 experiments). They were able to train the models (for CIFAR-10) up to 80% accuracy that way while BatchBALD's authors trained theirs (CINIC-10) up to 60% accuracy because the final dataset size is significantly smaller (1,400 in CINIC-10 vs 35,000 in CIFAR-10). ICAL's authors mentioned that BatchBALD ran out of memory (not sure if they tried sampling the joint entropies instead of enumerating all configurations).

# Scopes & Limitations

1. __Noisy estimator__: As the experiments showed, larger acquisition sizes deteriorate the performance of BatchBALD. The authors postulate that it might be caused by the sampling of the joint entropies ($m$ samples of $c^{n-1}$ configurations) and a significant amount of noise introduced by MC-dropout's VI approximation. _Does antithetic sampling reduce the noise of the latter?_

2. __Unlabelled data__: Semi-supervised extension to fully leverage the unlabelled dataset

3. __Unbalanced dataset__: Both BALD and BatchBALD is expected to not perform well when the test set is unbalanced as they aim to learn well about all classes and do not follow the density of the dataset.

# Appendix

A. Proved the submodularity of BatchBALD used one of the three definitions of submodular functions and the general subadditivity and additivity of entropies.

B.1. Proved that BatchBALD is upper-bounded by BALD using subadditivity of entropies.

B.2. By manipulating identities in the greedy algorithm above, one can show that it is maximising the mutual information of the **individual** points from $\mathcal{D}_U$ with the model parameters **conditioned** on the already-picked batch $A_{n-1}$ and the training dataset $\mathcal{D}_U$.

Furthermore, this MI is greater than or equal to the MI when conditioned on $A_{n-1} \cup \mathcal{D}_U$. In other words, each step of this greedy algorithm is at most as good as retraining the model on the new training set $A_{n-1} \cup \mathcal{D}_U$ and picking $x_n$ using $\alpha_{\text{BALD}}$ with acquisition size 1! (in the idealised setting, of course.)

C. Showed that it's possible to take $m$ MC samples as an approximate to the joint entropy $\mathbb{H}(y_1, \dots, y_n)$.

D. Behavour of BatchBALD is consistent regardless of the number of times repeated MNIST is repeated. Meanwhile, BALD is affected by the degree of repretitions, since it does not account for the joint MI, as expected.

E. BatchBALD outperforms Var Ratios and Mean STD (acquisition size 10, 10 mc-dropout samples)

F. -

G. Measuring diversity of points acquired: entropy of target class (acquired by BatchBALD) is higher than other methods.

# Closing remarks (relatability to AL project)

Possible extensions include:
1. Antithetic sampling of MC-dropout masks to reduce variance in inference samples.
2. Explore Semi-supervised learning + active learning combination.
3. Alternatives to MC-dropout? Any other methods that have lesser variance?
4. While measuring diversity of acquired points using entropy seems to be a natural thing to do (and should generally be expected?), occasionally, the best points could be specific classes (esp. if the current training pool has a deficit) that reduces the uncertainty in model parameters $\mathbf \omega$. In that case, class diversity is moot -- but perhaps this is a contrived example.

## TODO
1. find out how $\prod_{i=1}^np(y_i|\omega)$ was calculated without underflowing.

## Identities/properties used in proofs

1. **Subadditivity of Entropies**
    - $\mathbb{H}[x, y | \omega] \sim \mathbb{H}[x | \omega] + \mathbb{H}[y | \omega]$ where ~ is = (additive) iff $x \perp y\ |\ \omega$ else ~ is $\leq$ (subadditive).
2. **Conditional MI**
    - $\mathbb{I}[A \cup B; C] - \mathbb{I}[B; C] = \mathbb{I}[A; C|B]$ (more intuitive if you think about the information diagram)


# Chapter 10: Biological Membranes

(c) 2021 Manuel Razo. This work is licensed under a [Creative Commons Attribution License CC-BY 4.0](https://creativecommons.org/licenses/by/4.0/). All code contained herein is licensed under an [MIT license](https://opensource.org/licenses/MIT).

---

```python
# Our numerical workhorses
import numpy as np
import scipy as sp
from scipy.special import kv as kv # Bessel function
import pandas as pd
import sympy

# Import matplotlib stuff for plotting
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib as mpl

# Seaborn, useful for graphics
import seaborn as sns

# Import the default plotting style
import plotting_style as pstyle

# Function to import and display images from the Garland website
from IPython.display import Image

# This enables SVG graphics inline.
%config InlineBackend.figure_formats = {'png', 'retina'}

# Define directories
figdir = '../fig/'
```

```python
# Set plotting style
pstyle.set_plotting_style()
# Set dpi for inline figures
mpl.rcParams['figure.dpi'] = 110
```

## 11.10 Two-dimensional treatment of MscL

For this problem we have solved a PDE that defines the height profile of a membrane surrounding a mechanosensitive channel. The solution is written in terms of Hyperbolic Bessel functions and takes the form
$$
u(r) = C_1 K_0(k_1 r) + C_2 K_0(k_1^* r),
$$
where $r$ is the distance from the center of the channel, $K_o(\cdot)$ is the exponentially decaying hyperbolic Bessel function of order zero. The constants $k_1$ and $k_1^*$ are defined as
$$
k_{1}=\sqrt[4]{\frac{K_{t}}{K_{b} \omega_{0}^{2}}}\left(\frac{\sqrt{2}(1+i)}{2}\right),
$$
and
$$
k_{1}^=\sqrt[4]{\frac{K_{t}}{K_{b} \omega_{0}^{2}}}\left(\frac{\sqrt{2}(1-i)}{2}\right),
$$
where $K_b$ is the is the bending rigidity--a parameter characterizing the energy needed to bend the membrane. $K_t$ characterizes the energy cost of having a hydrophobic mismatch between the channel residues and the membrane, and $\omega_0$ is the membrane mid-plane height.

After using the boundary conditions $u(R) = U_o$, meaning that the height of the membrane next to the channel perimeter is given by $U_o$, and $u'(R) = 0$, meaning that the membrane is perfectly aligned with the channel (no weird kinks or anything next to the channel) we find that the two constants $C_1$ and $C_2$ are given by
$$
C_{1}=U_{0} \frac{K_{1}\left(k_{1}^{*} R\right) k_{1}^{*}}{K_{0}\left(k_{1} R\right) K_{1}\left(k_{1}^{*} R\right) k_{1}^{*}-K_{0}\left(k_{1}^{*} R\right) K_{1}\left(k_{1} R\right) k_{1}},
$$
and
$$
C_{2}=-U_{0} \frac{K_{1}\left(k_{1} R\right) k_{1}}{K_{0}\left(k_{1} R\right) K_{1}\left(k_{1}^{*} R\right) k_{1}^{*}-K_{0}\left(k_{1}^{*} R\right) K_{1}\left(k_{1} R\right) k_{1}},
$$
where $K_1(\cdot)$ is the hyperbolic Bessel function of order zero.

Knowing all of this, our task is now to plot the solution $u(r)$ as a function of the distance from the channel. For numerical values we will say that $K_b \approx 20\;k_BT$, $K_t \approx 60\; k_BT/nm^2$, and $\omega_0 \approx 2\; nm$, which gives
$$
\sqrt[4]{\frac{K_{t}}{K_{b} \omega_0^2}} \approx 1 \; nm^{-1}.
$$
Let us first compute the two constants $C_1$ and $C_2$ up to the $U_0$ factor.

```python
# Define channel radius
R = 2  # nm

# Define constants k1 and k1_star
k1 = np.sqrt(2) * (1 - (1j)) / 2
k1_star = np.conjugate(k1)

# Compute C1
C1 = (
    kv(1, k1_star * R)
    * k1_star
    / (
        kv(0, k1 * R) * kv(1, k1_star * R) * k1_star
        - kv(0, k1_star * R) * kv(1, k1 * R) * k1
    )
)

# Compute C2
C2 = -(
    kv(1, k1 * R)
    * k1
    / (
        kv(0, k1 * R) * kv(1, k1_star * R) * k1_star
        - kv(0, k1_star * R) * kv(1, k1 * R) * k1
    )
)
```

Now we can compute the height profile $u(r)$

```python
# Define range of values for r
r_range = np.linspace(R, 10, 100)  # nm

# Compute height profile
u_r = C1 * kv(0, k1 * r_range) + C2 * kv(0, k1_star * r_range)
```

And now let's plot the solution!

```python
# Plot solution
plt.plot(r_range, u_r)

# Label axis
plt.xlabel(r"$r$ (nm)")
plt.ylabel(r"$u(r) / U_0$ (nm)")

plt.savefig(figdir + "problem_11_10_03.png", bbox_inches="tight", dpi=300)
```