# continuous2discrete

*Johannes Hörmann, Lukas Elflein, 2019*

from continuous electrochemical double layer theory to discrete coordinate sets

```python
# for dynamic module reload during testing, code modifications take immediate effect
%load_ext autoreload
%autoreload 2
```

```python
# stretching notebook width across whole window
from IPython.core.display import display, HTML
display(HTML("<style>.container { width:100% !important; }</style>"))
```

<style>.container { width:100% !important; }</style>

```python
# basics
import logging
import numpy as np
import scipy.constants as sc
import matplotlib.pyplot as plt
```

```python
# sampling
from scipy import interpolate
from continuous2discrete import generate_structure
from continuous2discrete import plot_dist
from continuous2discrete import get_histogram
```

```python
# Physical distribution generation
from continuous2discrete.poisson_boltzmann_distribution import debye
from continuous2discrete.poisson_boltzmann_distribution import gamma
from continuous2discrete.poisson_boltzmann_distribution import charge_density
from continuous2discrete.poisson_boltzmann_distribution import potential

from continuous2discrete.continuous2discrete import concat_names_structs
from continuous2discrete.continuous2discrete import export_named_struc

np.random.seed(74)
```

```python
# Poisson-Nernst-Planck solver
from continuous2discrete.poisson_nernst_planck_distribution import PoissonNernstPlanckSystem
```

```python
# 3rd party file output
import ase
import ase.io
```

# The Poisson-Boltzman Distribution
*Lukas Elflein, 2019*

In order to understand lubrication better, we simulate thin layers of lubricant on a metallic surface, solvated in water.
Different structures of lubricant films are created by varying parameters like their concentration and the charge of the surface.
The lubricant is somewhat solvable in water, thus parts of the film will diffuse into the bulk water.
Lubricant molecules are charged, and their distribution is roughly exponential.

As simplification, we  first create a solution of ions (Na+, purple; Cl-, green) in water (not shown).

Close to the positively charged metallic surface, the electric potential (red) will be highest, falling off exponentially when further away.
This potential attracts negatively charged Chlorine ions, and pushes positively charged Natrium ions away, resulting in a higher (lower) concentration of Clorine (Natrium) near the surface.

To calculate this, we first need to find out how ions are distributed in solution.
A good description of the concentration of our ion species, $c_{Na^+}$, is given by the solution to the Poisson-Boltzmann equation:

$
\begin{align}
\rho_{Na^+}(z) &= \rho_{Na^+}(\infty) e^{-e \Psi(z)/k_B T}\\
\Psi(z) &= \frac{2k_B T}{e} \log\Big(\frac{1 + \gamma e^{-\kappa z}}{1- \gamma e^{-\kappa z}}\Big)
        \approx \frac{4k_B T}{e} \gamma e^{-\kappa z} \\
\gamma &= \tanh(\frac{e\Psi(0)}{4k_B T})\\
\kappa &= 1/\lambda_D\\
\lambda_D &= \Big(\frac{\epsilon \epsilon_0 k_B T}{\sum_{i} \rho_i(\infty) e^2 z_i^2} \Big)^\frac{1}{2} [m^{-1}]
\end{align}
$

With:
* $z$: Distance from the double layer
* $\Psi(0)$: Potential at the surface
* $\Psi(z)$: Potential in the solution
* $k_B$: Boltzmann Constant
* $T$: Temperature [Kelvin]
* $e$: Elemental Charge (or Euler's constant when exponentiated)
* $\gamma$: Term from Gouy-Chapmann theory
    * $\gamma \rightarrow 1$ for high potentials
    * $\Psi(z) \approx \Psi_0 e^{-\kappa z}$ for low potentials $\Psi(0) \approx 0$
* $\lambda_D$: Debye Length (34.0 nm for NaCl, 10^-4 M, 25°C)
* $\rho_{Na^+}$: Concentration of Natrium ions
* $\rho_{Na^+}(\infty)$: Bulk Natrium concentration (at infinity, where the solution is homogeneous)
* $\epsilon$: Permittivity of the solution
* $\epsilon_0$: Electric constant aka Vacuum permittivity
* $z_i$: Charge of species i

These equations are implemented in `poisson_boltzmann_distribution.py`

```python
default_density = sc.Avogadro * 1000 * 10**-4
rho = {'Na': default_density, 'Cl': default_density}
charge = {'Na': 1, 'Cl': -1}
deb = debye(rho_bulk=rho, charge=charge) * 10**9
print('Debye Length of 10^-4 M saltwater: {} nm (Target: 30.4 nm)'.format(round(deb, 2)))

density = np.logspace(-6, 0, 50) * sc.Avogadro * 1000
debyes = [debye(rho_bulk={'Na': d, 'Cl': d}, charge=charge) * 10**9 for d in density]

plt.xlabel('Density [1/m^3]')
plt.ylabel('Debye length at 25° [nm]')
plt.semilogx(density, debyes, marker='.')
plt.show()
```

The debye length depends on the concentration of ions in solution, at low concentrations it becomes large. We can reproduce literature debye lengths with our function, so everything looks good.

## Gamma Function

Next we calculate the gamma function $\gamma = \tanh(\frac{e\Psi(0)}{4k_B T})$

```python
x = np.linspace(12, 16, 40)
gammas = [gamma(10 ** i, 300) for i in x]
plt.xlabel('Potential')
plt.ylabel('Gamma at 300K')
plt.plot(x, gammas, marker='o')
plt.show()
```

Which looks as expected, but we have no values to compare it against.

## Potential

We plug these two functions into the expression for the potential

$\Psi(z) = \frac{2k_B T}{e} \log\Big(\frac{1 + \gamma e^{-\kappa z}}{1- \gamma e^{-\kappa z}}\Big)
        \approx \frac{4k_B T}{e} \gamma e^{-\kappa z}$

```python
z = np.linspace(0, 2*10**-7, 10000)
density = sc.Avogadro * 1000 * 10**-4
rho = {'Na': density, 'Cl':density}
charge = {'Na': 1, 'Cl': -1}
pot_0 = 100
psi = [potential(location=loc, rho_bulk=rho, charge=charge, surface_potential=pot_0) for loc in z]
plt.xlabel('z [nm]')
plt.ylabel('Potential [V]')
plt.plot(z*10**9, psi, marker='')
plt.show()
```

The potential is smooth and looks roughly exponential. Everything good so far.

## Charge density

Now we obtain the charge density $\rho$ from the potential $\Psi$ via

$\rho_{Na^+}(z) = \rho_{Na^+}(\infty) e^{-e \Psi(z)/k_B T}$

```python
z = np.linspace(0, 100*10**-9, 2000)
density = sc.Avogadro * 1000 * 10**-4
rho = {'Na': density, 'Cl':density}
charge = {'Na': 1, 'Cl': -1}
pot_0 = 0.05  # Breaks if > 1

psi = [potential(location=loc, rho_bulk=rho, charge=charge, surface_potential=pot_0) for loc in z]

rho_na = np.array([charge_density(location=loc, rho_bulk=rho, charge=charge, surface_potential=pot_0, species='Na') for loc in z])
rho_cl = np.array([charge_density(location=loc, rho_bulk=rho, charge=charge, surface_potential=pot_0, species='Cl') for loc in z])

deb = debye(rho_bulk=rho, charge=charge) * 10**9

fig, ax1 = plt.subplots(figsize=[16,5])
ax1.set_xlabel('z [nm]')
ax1.plot(z*10**9, psi, marker='', color='red', label='Potential', linewidth=1, linestyle='--')
ax1.set_ylabel('Potential')
ax1.axvline(x=deb, label='Debye Length', color='orange')

ax2 = ax1.twinx()
ax2.plot(z*10**9, [density]*len(z), label='Bulk concentration', color='grey')
ax2.plot(z*10**9, rho_na, marker='', color='green', label='Na+ ions')
ax2.plot(z*10**9, rho_cl, marker='', color='blue', label='Cl- ions')
ax2.set_ylabel('Density')

#fig.legend(loc='center')
ax2.legend(loc='best', fontsize=15)
ax1.legend(loc='upper center', fontsize=15)
fig.tight_layout()
plt.show()
```

The charge density behaves as expected, it interpolates between low (high) concentration and the bulk concentration within the first few debye lengths.

## Sampling
Now let's see if we can just plug our new distribution in our existing framework.

First, we need to convert the physical distribution to the format we were using so far:

```python
def wrap_distribution(x, species):
    """Wrapper for na+ ions."""
    density = sc.Avogadro * 1000 * 10**-4
    rho = {'Na': density, 'Cl':density}
    charge = {'Na': 1, 'Cl': -1}
    pot_0 = 0.05  # Breaks if > 1

    if np.isscalar(x):
        x = [x]

    y = [ charge_density(
            location=pos,
            rho_bulk=rho,
            charge=charge,
            surface_potential=pot_0,
            species=species) for pos in x ]

    return np.array(y)

def cl_distribution(x):
    return wrap_distribution(x, species='Cl')

def na_distribution(x):
    return wrap_distribution(x, species='Na')
```

Now we can sample the distribution of our $Na^+$ ions in z-direction.

```python
x = 50 * 10**-9
z = 100 * 10**-9
box = np.array([x, x, z])
sample_size = 10000
```

```python
struc = generate_structure(distribution=na_distribution, box=box, atom_count=sample_size)
histx, histy, histz = get_histogram(struc, box=box, n_bins=51)
plot_dist(histz, 'Distribution of Na+ ions in z-direction', reference_distribution=na_distribution)
```

```python
struc = generate_structure(distribution=cl_distribution, box=box, atom_count=sample_size)
histx, histy, histz = get_histogram(struc, box=box, n_bins=51)
plot_dist(histz, 'Distribution of Cl- ions in z-direction', reference_distribution=cl_distribution)
```

## Write to file
To visualize our structure, we export it to the .xyz file format, which is basically

```
ATOM_NUMBER
Optional comment
atom_type x y z
atom_type x y z
```

Avogadro expects x, y, z to be in units of $10^{-9}~m$, so we convert our salt "solution" to this unit.

Also, we use fewer atoms, just 100 for the 100x50x50 box - this will increase fluctuations with respect to the correct distribution though.

```python
cl_struc = generate_structure(distribution=cl_distribution, box=box, atom_count=100)
na_struc = generate_structure(distribution=na_distribution, box=box, atom_count=100)

concat_list = concat_names_structs(struc_list=[cl_struc, na_struc], name_list=['Cl', 'Na'])
rescaled_list = []
for line in concat_list:
    name, x, y, z = line
    x = float(x) * 10**9
    y = float(y) * 10**9
    z = float(z) * 10**9
    rescaled_list += [[name, x, y, z]]

rescaled_list = np.array(rescaled_list)
export_named_struc(rescaled_list)

histx, histy, histz = get_histogram(cl_struc, box=box, n_bins=21)
plot_dist(histz, 'Distribution of Cl- ions in z-direction', reference_distribution=cl_distribution)

histx, histy, histz = get_histogram(na_struc, box=box, n_bins=21)
plot_dist(histz, 'Distribution of Na+ ions in z-direction', reference_distribution=na_distribution)
```

# General Poisson-Nernst-Planck System

A binary Poisson-Nernst-Planck system corresponds to the transport problem in semiconductor physics.
In this context, Debye length, charge carrier densities and potential are related as follows.

## Excursus: Transport problem in PNP junction (German)

### Debye length

Woher kommt die Debye-Länge

$$ \lambda = \sqrt{ \frac{\varepsilon \varepsilon_0 k_B T}{q^2 n_i} }$$

als natürliche Längeneinheit des Transportptoblems?

Hier ist $n_i$ eine Referenzladungsträgerdichte, in der Regel die intrinsische Ladungsträgerdichte.
In dem Beispiel mit $N^+NN^+$-dotiertem Halbleiter erzeugen wir durch unterschiedliches Doping an den Rändern die erhöhte Donatorendichte $N_D^+ = 10^{20} \mathrm{cm}^{-3}$ und im mitteleren Bereich "Standarddonatorendichte" $N_D = 10^{18} \mathrm{cm}^{-3}$. Nun können wir als Referenz $n_i = N_D$ wählen und die Donatorendichten als $N_D = 1 \cdot n_i$ und $N_D^+ = 100 \cdot n_i$ ausdrücken. Diese normierte Konzentration nennen wir einfach $\tilde{N}_D$: $N_D = \tilde{N}_D \cdot n_i$.

Ein ionisierter Donator trägt die Ladung $q$, ein Ladungsträger (in unserem Fall ein Elektron) trägt die Elementarladung $-q$. Die Raumladungsdichte $\rho$ in der Poissongleichung

$$ \nabla^2 \varphi = - \frac{\rho}{\varepsilon \varepsilon_0}$$

lässt sich also ganz einfach als $\rho = - (n - N_D) \cdot q = - (\tilde{n} - \tilde{N}_D) ~ n_i ~ q$ ausdrücken.

Konventionell wird das Potential auf $u = \frac{\phi ~ q}{k_B ~ T}$ normiert. Die Poissongleichung nimmt damit die Form

$$\frac{k_B ~ T}{q} \cdot \nabla^2 u = \frac{(\tilde{n} - \tilde{N}_D) ~ n_i ~ q }{\varepsilon \varepsilon_0}$$

oder auch

$$ \frac{\varepsilon ~ \varepsilon_0 ~ k_B ~ T}{q^2 n_i} \cdot \nabla^2 u = \lambda^2 \cdot \nabla^2 u = \tilde{n} - \tilde{N}_D$$

### Dimensionless formulation

Poisson- und Drift-Diffusionsgleichung

$$
\lambda^2 \frac{\partial^2 u}{\partial x^2} = n - N_D
$$

$$
\frac{\partial n}{\partial t} = - D_n \ \frac{\partial}{\partial x} \left( n \ \frac{\partial u}{\partial x} - \frac{\partial n}{\partial x} \right) + R
$$

Skaliert mit [l], [t]:

$$
\frac{\lambda^2}{[l]^2} \frac{\partial^2 u}{\partial \tilde{x}^2} = n - N
$$

und

$$
\frac{1}{[t]} \frac{\partial n}{\partial \tilde{t}} = - \frac{D_n}{[l]^2} \ \frac{\partial}{\partial x} \left( n \ \frac{\partial u}{\partial x} - \frac{\partial n}{\partial x} \right) + R
$$

oder

$$
\frac{\partial n}{\partial \tilde{t}} = - \tilde{D}_n \ \frac{\partial}{\partial x} \left( n \ \frac{\partial u}{\partial x} - \frac{\partial n}{\partial x} \right) +  \tilde{R}
$$

mit

$$
\tilde{D}_n = D_n \frac{[t]}{[l]^2} \Leftrightarrow [t] = [l]^2 \ \frac{ \tilde{D}_n } { D_n }
$$

und

$$ \tilde{R} = \frac{n - N_D}{\tilde{\tau}}$$

mit $\tilde{\tau} = \tau / [t]$.

$\tilde{\lambda} = 1$ und $\tilde{D_n} = 1$ werden mit
$[l] = \lambda$ und $[t] = \frac{\lambda^2}{D_n}$ erreicht:

### Discretization

Naive Diskretisierung (skaliert):

$$ \frac{1}{\Delta x^2} ( u_{i+1}-2u_i+u_{i-1} ) = n_i - N_i $$

$$ \frac{1}{\Delta t} ( n_{i,j+1} - n_{i,j} ) = - \frac{1}{\Delta x^2} \cdot \left[ \frac{1}{4} (n_{i+1} - n_{i-1}) (u_{i+1} - u_{i-1}) + n_i ( u_{i+1} - 2 u_i + u_{i-1} ) - ( n_{i+1} - 2 n_i + n_{i-1} ) \right] + \frac{ n_i - N_i}{ \tilde{\tau} } $$

Stationär:

$$
 u_{i+1}-2u_i+u_{i-1} - \Delta x^2 \cdot n_i + \Delta x^2 \cdot N_i = 0
$$

und

$$
  \frac{1}{4} (n_{i+1} - n_{i-1}) (u_{i+1} - u_{i-1}) + n_i ( u_{i+1} - 2 u_i + u_{i-1} ) - ( n_{i+1} - 2 n_i + n_{i-1} ) - \Delta x^2 \cdot \frac{ n_i - N_i}{ \tilde{\tau} } = 0
$$

### Newton-Iteration für gekoppeltes nicht-lineares Gleichungssystem

Idee: Löse nicht-lineares Finite-Differenzen-Gleichungssystem über Newton-Verfahren

$$ \vec{F}(\vec{x}_{k+1}) = F(\vec{x}_k + \Delta \vec{x}_k) \approx F(\vec{x}_k) + \mathbf{J_F}(\vec{x}_k) \cdot \Delta \vec{x}_k + \mathcal{O}(\Delta x^2)$$

mit Unbekannter $\vec{x_k} = \{u_1^k, \dots, u_N^k, n_1^k, \dots, n_N^k\}$  und damit

$$ \Rightarrow \Delta \vec{x}_k = - \mathbf{J}_F^{-1} ~ F(\vec{x}_k)$$

wobei die Jacobi-Matrix $2N \times 2N$ Einträge

$$ \mathbf{J}_{ij}(\vec{x}_k) = \frac{\partial F_i}{\partial x_j} (\vec{x}_k) $$

besitzt, die bei jedem Iterationsschritt für $\vec{x}_k$ ausgewertet werden.
Der tatsächliche Aufwand liegt in der Invertierung der Jacobi-Matrix, um in jeder Iteration $k$ den Korrekturschritt $\Delta \vec{x}_k$ zu finden.m

$F(x)$ wird wie unten definiert als:

$$
 u_{i+1}-2u_i+u_{i-1} - \Delta x^2 \cdot n_i + \Delta x^2 \cdot N_i = 0
$$

und

$$
  \frac{1}{4} (n_{i+1} - n_{i-1}) (u_{i+1} - u_{i-1}) + n_i ( u_{i+1} - 2 u_i + u_{i-1} ) - ( n_{i+1} - 2 n_i + n_{i-1} ) - \Delta x^2 \cdot \frac{ n_i - N_i}{ \tilde{\tau} } = 0
$$

### Controlled-Volume

Drücke nicht-linearen Teil der Transportgleichung (genauer, des Flusses) über Bernoulli-Funktionen

$$ B(x) = \frac{x}{\exp(x)-1} $$

aus (siehe Vorlesungsskript). Damit wir in der Nähe von 0 nicht "in die Bredouille geraten", verwenden wir hier lieber die Taylorentwicklung. In der Literatur (Selbherr, S. Analysis and Simulation of Semiconductor Devices, Spriger 1984) wird eine noch aufwendigere stückweise Definition empfohlen, allerdings werden wir im Folgenden sehen, dass unser Ansatz für dieses stationäre Problem genügt.

## Implementation for Poisson-Nernst-Planck system

Poisson-Nernst-Planck system for $k = {1 \dots M}$ ion species in dimensionless formulation

$$ \nabla^2 u + \rho(n_{1},\dots,n_{M}) = 0 $$

$$ \nabla^2 n_k + \nabla ( z_k n_k \nabla u ) = 0 \quad \text{for} \quad k = 1 \dots M $$

yields a naive finite difference discretization on $i = {1 \dots N}$ grid points for $k = {1 \dots M}$ ion species

$$ \frac{1}{\Delta x^2} ( u_{i+1}-2u_i+u_{i-1} )  + \frac{1}{2} \sum_{k=1}^M z_k n_{i,k} = 0 $$

$$ - \frac{1}{\Delta x^2} \cdot \left[ \frac{1}{4} z_k (n_{i+1,k} - n_{i-1,k}) (u_{i+1} - u_{i-1}) + z_k n_{i,k} ( u_{i+1} - 2 u_i + u_{i-1} ) + ( n_{i+1,k} - 2 n_{i,k} + n_{i-1,k} ) \right] $$

or rearranged

$$ u_{i+1}-2 u_i+u_{i-1} + \Delta x^2 \frac{1}{2} \sum_{k=1}^M z_k n_{i,k}  = 0 $$

and

$$
  \frac{1}{4} z_k (n_{i+1,k} - n_{i-1,k}) (u_{i+1,k} - u_{i-1,k}) + z_k n_{i,k} ( u_{i+1} - 2 u_i + u_{i-1} ) - ( n_{i+1,k} - 2 n_{i,k} + n_{i-1,k} ) = 0
$$

### Controlled Volumes, 1D

Finite differences do not converge in our non-linear systems. Instead, we express non-linear part of the Nernts-Planck equations with Bernoulli function (Selberherr, S. Analysis and Simulation of Semiconductor Devices, Spriger 1984)

$$ B(x) = \frac{x}{\exp(x)-1} $$

```python
def B(x):
    return np.where( np.abs(x) < 1e-9,
        1 - x/2 + x**2/12 - x**4/720, # Taylor
        x / ( np.exp(x) - 1 ) )
```

```python
xB = np.arange(-10,10,0.1)
```

```python
plt.plot( xB ,B( xB ), label="$B(x)$")
plt.plot( xB, - B(-xB), label="$-B(-x)$")
plt.plot( xB, B(xB)-B(-xB), label="$B(x)-B(-x)$")
plt.legend()
```

Looking at (dimensionless) flux $j_k$ throgh segment $k$ in between grid points $i$ and $j$,

$$ j_k = - \frac{dn}{dx} - z n \frac{du}{dx} $$

for an ion species with number charge $z$ and (dimensionless) concentration $n$,
we assume (dimensionless) potential $u$ to behave linearly within this segment. The linear expression

$$ u = \frac{u_j - u_i}{L_k} \cdot \xi_k + u_i = a_k \xi_k + u_i $$

with the segment's length $L_k = \Delta x$ for uniform discretization, $\xi_k = x - x_i$ and proportionality factor $a_k = \frac{u_j - u_i}{L_k}$ leadsd to a flux

$$ j_k = - \frac{dn}{d\xi} - z a_k n  $$

solvable for $v$ via

$$ \frac{dn}{d\xi} = - z a_k n - j_k $$

or

$$ \frac{dn}{z a_k n + j_k} = - d\xi \text{.} $$

We intergate from grid point $i$ to $j$

$$ \int_{n_i}^{n_j} \frac{1}{z a_k n + j_k} dn = - L_k $$

and find

$$ \frac{1}{(z a_k)} \left[ \ln(j_k + z a_k n) \right]_{n_i}^{n^j} = - L_k $$

or

$$ \ln(j_k + z a_k n_j) - \ln(j_k + z a_k n_i) = - z a_k L_k $$

which we solve for $j_k$ by rearranging

$$ \frac{j_k + z a_k n_j}{j_k + z a_k n_i} = e^{- z a_k L_k} $$

$$ j_k + z a_k n_j = (j_k + z a_k n_i) e^{- z a_k L_k} $$

$$ j_k ( 1 - e^{- z a_k L_k} ) = - z a_k n_j  +  z a_k n_i e^{- z a_k L_k} $$

$$j_k = \frac{z a_k n_j}{e^{- z a_k L_k} - 1}  +  \frac{ z a_k n_i e^{- z a_k L_k}}{ 1 - e^{- z a_k L_k}}$$

$$j_k = \frac{1}{L_k} \cdot \left[ \frac{z a_k L_k n_j}{e^{- z a_k L_k} - 1}  +  \frac{ z a_k L_k n_i }{ e^{z a_k L_k} - 1} \right] $$

or with $B(x) = \frac{x}{e^x-1}$ expressed as

$$j_k = \frac{1}{L_k} \cdot \left[ - n_j B( - z a_k L_k ) +  n_i B( z a_k L_k) \right] $$

and resubstituting $a_k = \frac{u_j - u_i}{L_k}$ as

$$j_k = - \frac{1}{L_k} \cdot \left[ n_j B( z [u_i - u_j] ) - n_i B( z [u_j - u_i] ) \right] \ \text{.}$$

When employing our 1D uniform grid with $j_k = j_{k-1}$ for all $k = 1 \dots N$,

$$ j_k \Delta x  = n_{i+1} B( z [u_i - u_{i+1}] ) - n_i B( z [u_{i+1} - u_i] ) $$

and

$$ j_{k-1} \Delta x  = n_i B( z [u_{i-1} - u_i] ) - n_{i-1} B( z [u_i - u_{i-1}] ) $$

require

$$ n_{i+1} B( z [u_i - u_{i+1}] ) - n_i \left( B( z [u_{i+1} - u_i] ) + B( z [u_{i-1} - u_i] ) \right) + n_{i-1} B( z [u_i - u_{i-1}] ) = 0 $$

```python
# PoissonNernstPlanckSystem makes extensive use of Python's logging module

# configure logging: verbosity level and format as desired
standard_loglevel   = logging.INFO
# standard_logformat  = ''.join(("%(asctime)s",
#  "[ %(filename)s:%(lineno)s - %(funcName)s() ]: %(message)s"))
standard_logformat  = "[ %(filename)s:%(lineno)s - %(funcName)s() ]: %(message)s"

# reset logger if previously loaded
logging.shutdown()

logging.basicConfig(level=standard_loglevel,
                    format=standard_logformat,
                    datefmt='%m-%d %H:%M')

# in Jupyter notebooks, explicitly modifying the root logger necessary
logger = logging.getLogger()
logger.setLevel(standard_loglevel)

# remove all handlers
for h in logger.handlers: logger.removeHandler(h)

# create and append custom handles
ch = logging.StreamHandler()
formatter = logging.Formatter(standard_logformat)
ch.setFormatter(formatter)
ch.setLevel(standard_loglevel)
logger.addHandler(ch)
```

```python
# Test 1
logging.info("Root logger")
```

    [ <ipython-input-53-adec08aea02d>:2 - <module>() ]: Root logger

```python
# Test 2
logger.info("Root Logger")
```

    [ <ipython-input-54-4335e316391c>:2 - <module>() ]: Root Logger

## Test case 1: 0.1 mM NaCl, positive potential u = 0.05 V

```python
L = 1e-7 # 1 mu_m
```

```python
# define desired system
pnp = PoissonNernstPlanckSystem(L=L)
# constructor takes keyword arguments
#   c=array([0.1, 0.1]), z=array([ 1, -1]), L=1e-07, T=298.15, delta_u=0.05, relative_permittivity=79, vacuum_permittivity=8.854187817620389e-12, R=8.3144598, F=96485.33289
# with default values set for 0.1 mM NaCl aqueous solution across 100 nm  and 0.05 V potential drop
```

```python
# technical settings
pnp.output = True # makes Newton solver display convergence plots
pnp.N = 200       # uniformly distanced grid points
pnp.maxit = 20    # maximum number of Newton iterations
pnp.outfreq = 2   # plot every 2nd step
pnp.e = 1e-12     # tolerance
```

```python
pnp.useStandardInterfaceBC()
```

```python
pnp.init()
```

```python
uij, nij, lamj = pnp.solve()
```

### Validation: Analytical half-space solution & Numerical finite-size PNP system

```python
# analytic Poisson-Boltzmann distribution and numerical solution to full Poisson-Nernst-Planck system

z = np.linspace(0, L, 2000)
density = sc.Avogadro * 1000 * 1e-4
rho = {'Na': density, 'Cl':density}
charge = {'Na': 1, 'Cl': -1}
pot_0 = pnp.delta_u  # Breaks if > 1

psi = [potential(location=loc, rho_bulk=rho, charge=charge, surface_potential=pot_0) for loc in z]

rho_na = np.array([charge_density(location=loc, rho_bulk=rho, charge=charge, surface_potential=pot_0, species='Na') for loc in z])
rho_cl = np.array([charge_density(location=loc, rho_bulk=rho, charge=charge, surface_potential=pot_0, species='Cl') for loc in z])

deb = debye(rho_bulk=rho, charge=charge) * 1e9

fig, (ax1,ax3) = plt.subplots(nrows=2,ncols=1,figsize=[16,10])
ax1.set_xlabel('z [nm]')
ax1.plot(z*1e9, psi, marker='', color='tomato', label='potential, PB', linewidth=1, linestyle='--')
ax1.plot(pnp.grid*1e9, pnp.potential, marker='', color='tab:red', label='potential, PNP', linewidth=1, linestyle='-')

ax1.set_ylabel('potential (V)')

ax2 = ax1.twinx()
ax2.plot(z*1e9, [density/sc.Avogadro]*len(z), label='bulk concentration', color='grey', linestyle=':')
ax2.plot(z*1e9, rho_na/sc.Avogadro, marker='', color='bisque', label='Na+, PB',linestyle='--')
ax2.plot(pnp.grid*1e9, pnp.concentration[0], marker='', color='tab:orange', label='Na+, PNP', linewidth=2, linestyle='-')

ax2.plot(z*1e9, rho_cl/sc.Avogadro, marker='', color='lightskyblue', label='Cl-, PB',linestyle='--')
ax2.plot(pnp.grid*1e9, pnp.concentration[1], marker='', color='tab:blue', label='Cl-, PNP', linewidth=2, linestyle='-')

ax2.set_ylabel('concentration (mM)')

ax1.axvline(x=deb, label='Debye Length', color='grey', linestyle=':')

ax3.semilogy(z*1e9, [density/sc.Avogadro]*len(z), label='bulk concentration', color='grey', linestyle=':')
ax3.semilogy(z*1e9, rho_na/sc.Avogadro, marker='', color='bisque', label='Na+, PB',linestyle='--')
ax3.semilogy(pnp.grid*1e9, pnp.concentration[0], marker='', color='tab:orange', label='Na+, PNP', linewidth=2, linestyle='-')

ax3.semilogy(z*1e9, rho_cl/sc.Avogadro, marker='', color='lightskyblue', label='Cl-, PB',linestyle='--')
ax3.semilogy(pnp.grid*1e9, pnp.concentration[1], marker='', color='tab:blue', label='Cl-, PNP', linewidth=2, linestyle='-')

ax3.set_ylabel('concentration (mM)')

#fig.legend(loc='center')
ax1.legend(loc='upper center', fontsize=15)
ax2.legend(loc='best', fontsize=15)
ax3.legend(loc='upper right', fontsize=15)

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

#### Potential at left and right hand side of domain

```python
(pnp.potential[0],pnp.potential[-1])
```

#### Residual cation flux at interface and at open right hand side

```python
( pnp.leftFluxBC(pnp.xij1,0), pnp.rightFluxBC(pnp.xij1,0) )
```

#### Residual anion flux at interface and at open right hand side

```python
(pnp.leftFluxBC(pnp.xij1,1),  pnp.rightFluxBC(pnp.xij1,0) )
```

#### Cation concentration at interface and at open right hand side

```python
(pnp.concentration[0,0],pnp.concentration[0,-1])
```

#### Anion concentration at interface and at open right hand side

```python
(pnp.concentration[1,0],pnp.concentration[1,-1])
```

## Test case 2: 0.1 mM NaCl, negative potential u = -0.05 V, analytical solution as initial values

```python
L = 1e-7 # 1 mu_m
delta_u=-0.05
```

```python
# define desired system
del pnp
pnp = PoissonNernstPlanckSystem(L=L,delta_u=delta_u)
# constructor takes keyword arguments
#   c=array([0.1, 0.1]), z=array([ 1, -1]), L=1e-07, T=298.15, delta_u=0.05, relative_permittivity=79, vacuum_permittivity=8.854187817620389e-12, R=8.3144598, F=96485.33289
# with default values set for 0.1 mM NaCl aqueous solution across 100 nm  and 0.05 V potential drop
```

```python
# technical settings
pnp.output = True # makes Newton solver display convergence plots
pnp.N = 200       # uniformly distanced grid points
pnp.maxit = 20    # maximum number of Newton iterations
pnp.outfreq = 2   # plot every 2nd step
pnp.e = 1e-12     # tolerance
```

```python
pnp.useStandardInterfaceBC()
```

```python
pnp.init()
```

```python

# initial config
z = np.linspace(0, pnp.L, pnp.Ni)
density = sc.Avogadro * 1000 * 1e-4
rho = {'Na': density, 'Cl':density}
charge = {'Na': 1, 'Cl': -1}
pot_0 = pnp.delta_u

rho_na = np.array([charge_density(location=loc, rho_bulk=rho, charge=charge, surface_potential=pnp.delta_u, species='Na') for loc in z])
rho_cl = np.array([charge_density(location=loc, rho_bulk=rho, charge=charge, surface_potential=pnp.delta_u, species='Cl') for loc in z])

c_na = rho_na / sc.Avogadro
c_cl = rho_cl / sc.Avogadro
```

```python
plt.semilogy(c_na)
plt.semilogy(c_cl)
```

```python
ni0 = np.vstack([ c_na, c_cl ]) / pnp.c_unit
```

```python
pnp.ni0 = ni0
```

```python
ui0 = pnp.initial_values()
```

```python
plt.plot(ui0)
```

```python
uij, nij, lamj = pnp.solve() # faster conversion, compare to 7 steps for test case 1
```

### Validation: Analytical half-space solution & Numerical finite-size PNP system

```python
# analytic Poisson-Boltzmann distribution and numerical solution to full Poisson-Nernst-Planck system

z = np.linspace(0, L, 2000)
density = sc.Avogadro * 1000 * 1e-4
rho = {'Na': density, 'Cl':density}
charge = {'Na': 1, 'Cl': -1}
pot_0 = pnp.delta_u  # Breaks if > 1

psi = [potential(location=loc, rho_bulk=rho, charge=charge, surface_potential=pot_0) for loc in z]

rho_na = np.array([charge_density(location=loc, rho_bulk=rho, charge=charge, surface_potential=pot_0, species='Na') for loc in z])
rho_cl = np.array([charge_density(location=loc, rho_bulk=rho, charge=charge, surface_potential=pot_0, species='Cl') for loc in z])

deb = debye(rho_bulk=rho, charge=charge) * 1e9

fig, (ax1,ax3) = plt.subplots(nrows=2,ncols=1,figsize=[16,10])
ax1.set_xlabel('z [nm]')
ax1.plot(z*1e9, psi, marker='', color='tomato', label='potential, PB', linewidth=1, linestyle='--')
ax1.plot(pnp.grid*1e9, pnp.potential, marker='', color='tab:red', label='potential, PNP', linewidth=1, linestyle='-')

ax1.set_ylabel('potential (V)')

ax2 = ax1.twinx()
ax2.plot(z*1e9, [density/sc.Avogadro]*len(z), label='bulk concentration', color='grey', linestyle=':')
ax2.plot(z*1e9, rho_na/sc.Avogadro, marker='', color='bisque', label='Na+, PB',linestyle='--')
ax2.plot(pnp.grid*1e9, pnp.concentration[0], marker='', color='tab:orange', label='Na+, PNP', linewidth=2, linestyle='-')

ax2.plot(z*1e9, rho_cl/sc.Avogadro, marker='', color='lightskyblue', label='Cl-, PB',linestyle='--')
ax2.plot(pnp.grid*1e9, pnp.concentration[1], marker='', color='tab:blue', label='Cl-, PNP', linewidth=2, linestyle='-')

ax2.set_ylabel('concentration (mM)')

ax1.axvline(x=deb, label='Debye Length', color='grey', linestyle=':')

ax3.semilogy(z*1e9, [density/sc.Avogadro]*len(z), label='bulk concentration', color='grey', linestyle=':')
ax3.semilogy(z*1e9, rho_na/sc.Avogadro, marker='', color='bisque', label='Na+, PB',linestyle='--')
ax3.semilogy(pnp.grid*1e9, pnp.concentration[0], marker='', color='tab:orange', label='Na+, PNP', linewidth=2, linestyle='-')

ax3.semilogy(z*1e9, rho_cl/sc.Avogadro, marker='', color='lightskyblue', label='Cl-, PB',linestyle='--')
ax3.semilogy(pnp.grid*1e9, pnp.concentration[1], marker='', color='tab:blue', label='Cl-, PNP', linewidth=2, linestyle='-')

ax3.set_ylabel('concentration (mM)')

#fig.legend(loc='center')
ax1.legend(loc='upper center', fontsize=15)
ax2.legend(loc='best', fontsize=15)
ax3.legend(loc='upper right', fontsize=15)

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

#### Potential at left and right hand side of domain

```python
(pnp.potential[0],pnp.potential[-1])
```

#### Residual cation flux at interface and at open right hand side

```python
( pnp.leftFluxBC(pnp.xij1,0), pnp.rightFluxBC(pnp.xij1,0) )
```

#### Residual anion flux at interface and at open right hand side

```python
(pnp.leftFluxBC(pnp.xij1,1),  pnp.rightFluxBC(pnp.xij1,0) )
```

#### Cation concentration at interface and at open right hand side

```python
(pnp.concentration[0,0],pnp.concentration[0,-1])
```

#### Anion concentration at interface and at open right hand side

```python
(pnp.concentration[1,0],pnp.concentration[1,-1])
```

## Test case 3: 0.1 mM NaCl, positive potential u = 0.05 V, 200 nm domain

```python
L = 2e-7 # 200 nm
```

```python
del pnp
```

```python
# define desired system
pnp = PoissonNernstPlanckSystem(L=L)
# constructor takes keyword arguments
#   c=array([0.1, 0.1]), z=array([ 1, -1]), L=1e-07, T=298.15, delta_u=0.05, relative_permittivity=79, vacuum_permittivity=8.854187817620389e-12, R=8.3144598, F=96485.33289
# with default values set for 0.1 mM NaCl aqueous solution across 100 nm  and 0.05 V potential drop
```

```python
# technical settings
pnp.output = True # makes Newton solver display convergence plots
pnp.N = 800       # uniformly distanced grid points
pnp.maxit = 20    # maximum number of Newton iterations
pnp.outfreq = 2   # plot every 2nd step
pnp.e = 1e-10     # tolerance
```

```python
pnp.useStandardInterfaceBC()
```

```python
pnp.init()
```

```python
uij, nij, lamj = pnp.solve()
```

### Validation: Analytical half-space solution & Numerical finite-size PNP system

```python
# analytic Poisson-Boltzmann distribution and numerical solution to full Poisson-Nernst-Planck system

z = np.linspace(0, pnp.L, 2000)
density = sc.Avogadro * 1000 * 1e-4
rho = {'Na': density, 'Cl':density}
charge = {'Na': 1, 'Cl': -1}
pot_0 = pnp.delta_u  # Breaks if > 1

psi = [potential(location=loc, rho_bulk=rho, charge=charge, surface_potential=pot_0) for loc in z]

rho_na = np.array([charge_density(location=loc, rho_bulk=rho, charge=charge, surface_potential=pot_0, species='Na') for loc in z])
rho_cl = np.array([charge_density(location=loc, rho_bulk=rho, charge=charge, surface_potential=pot_0, species='Cl') for loc in z])

deb = debye(rho_bulk=rho, charge=charge) * 1e9

fig, (ax1,ax3) = plt.subplots(nrows=2,ncols=1,figsize=[16,10])
ax1.set_xlabel('z [nm]')
ax1.plot(z*1e9, psi, marker='', color='tomato', label='potential, PB', linewidth=1, linestyle='--')
ax1.plot(pnp.grid*1e9, pnp.potential, marker='', color='tab:red', label='potential, PNP', linewidth=1, linestyle='-')

ax1.set_ylabel('potential (V)')

ax2 = ax1.twinx()
ax2.plot(z*1e9, [density/sc.Avogadro]*len(z), label='bulk concentration', color='grey', linestyle=':')
ax2.plot(z*1e9, rho_na/sc.Avogadro, marker='', color='bisque', label='Na+, PB',linestyle='--')
ax2.plot(pnp.grid*1e9, pnp.concentration[0], marker='', color='tab:orange', label='Na+, PNP', linewidth=2, linestyle='-')

ax2.plot(z*1e9, rho_cl/sc.Avogadro, marker='', color='lightskyblue', label='Cl-, PB',linestyle='--')
ax2.plot(pnp.grid*1e9, pnp.concentration[1], marker='', color='tab:blue', label='Cl-, PNP', linewidth=2, linestyle='-')

ax2.set_ylabel('concentration (mM)')

ax1.axvline(x=deb, label='Debye Length', color='grey', linestyle=':')

ax3.semilogy(z*1e9, [density/sc.Avogadro]*len(z), label='bulk concentration', color='grey', linestyle=':')
ax3.semilogy(z*1e9, rho_na/sc.Avogadro, marker='', color='bisque', label='Na+, PB',linestyle='--')
ax3.semilogy(pnp.grid*1e9, pnp.concentration[0], marker='', color='tab:orange', label='Na+, PNP', linewidth=2, linestyle='-')

ax3.semilogy(z*1e9, rho_cl/sc.Avogadro, marker='', color='lightskyblue', label='Cl-, PB',linestyle='--')
ax3.semilogy(pnp.grid*1e9, pnp.concentration[1], marker='', color='tab:blue', label='Cl-, PNP', linewidth=2, linestyle='-')

ax3.set_ylabel('concentration (mM)')

#fig.legend(loc='center')
ax1.legend(loc='upper center', fontsize=15)
ax2.legend(loc='best', fontsize=15)
ax3.legend(loc='upper right', fontsize=15)

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

#### Potential at left and right hand side of domain

```python
(pnp.potential[0],pnp.potential[-1])
```

#### Residual cation flux at interface and at open right hand side

```python
( pnp.leftFluxBC(pnp.xij1,0), pnp.rightFluxBC(pnp.xij1,0) )
```

#### Residual anion flux at interface and at open right hand side

```python
(pnp.leftFluxBC(pnp.xij1,1),  pnp.rightFluxBC(pnp.xij1,0) )
```

#### Cation concentration at interface and at open right hand side

```python
(pnp.concentration[0,0],pnp.concentration[0,-1])
```

#### Anion concentration at interface and at open right hand side

```python
(pnp.concentration[1,0],pnp.concentration[1,-1])
```

### Construct arbitrary initial conditions for subsequent system

```python
csup = pnp.concentration[1]-pnp.concentration[0]
```

```python
scale_factor = np.sum(pnp.ni0[0]*pnp.c_unit) / np.sum(csup)
```

```python
scale_factor
```

```python
ci0 = scale_factor*np.vstack((np.flip(csup),csup))
```

```python
ci0.shape
```

```python
plt.plot(ci0[0], label='Na+')
plt.plot(ci0[1], label='Cl-')
plt.legend()
```

## Test case 4: 1D electrochemical cell

```python
del pnp
```

```python
pnp = PoissonNernstPlanckSystem(delta_u=0.05)
```

```python
# technical settings
pnp.output = True # makes Newton solver display convergence plots
pnp.N = 800 # uniformlya distanced grid points
pnp.maxit = 20 # maximum number of Newton iterations
pnp.outfreq=2
pnp.e = 1e-10 # absolute tolerance
```

```python
pnp.useStandardCellBC()
```

```python
pnp.init()
```

```python
xij = pnp.solve()
```

```python
fig, (ax1,ax3) = plt.subplots(nrows=2,ncols=1,figsize=[16,10])
ax1.set_xlabel('z [nm]')

ax1.plot(pnp.grid*1e9, pnp.potential, marker='', color='tab:red', label='potential, PNP', linewidth=1, linestyle='-')

ax1.set_ylabel('potential (V)')

ax2 = ax1.twinx()
ax2.plot(pnp.grid*1e9, pnp.c[0]*np.ones(pnp.grid.shape), label='initial concentration', color='grey', linestyle=':')
ax2.plot(pnp.grid*1e9, pnp.concentration[0], marker='', color='tab:orange', label='Na+, PNP', linewidth=2, linestyle='-')

ax2.plot(pnp.grid*1e9, pnp.concentration[1], marker='', color='tab:blue', label='Cl-, PNP', linewidth=2, linestyle='-')

ax2.set_ylabel('concentration (mM)')

ax1.axvline(x=pnp.lambda_D()*1e9, label='Debye Length', color='grey', linestyle=':')

ax3.semilogy(pnp.grid*1e9, pnp.concentration[0], marker='', color='tab:orange', label='Na+, PNP', linewidth=2, linestyle='-')
ax3.semilogy(pnp.grid*1e9, pnp.concentration[1], marker='', color='tab:blue', label='Cl-, PNP', linewidth=2, linestyle='-')

ax3.set_ylabel('concentration (mM)')

ax2.legend(loc='upper center', fontsize=15)
ax1.legend(loc='lower center', fontsize=15)

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

#### Potential at left and right hand side of domain

```python
(pnp.potential[0],pnp.potential[-1])
```

#### Residual cation flux at interfaces

```python
( pnp.leftFluxBC(pnp.xij1,0), pnp.rightFluxBC(pnp.xij1,0) )
```

#### Residual anion flux at interfaces

```python
(pnp.leftFluxBC(pnp.xij1,1),  pnp.rightFluxBC(pnp.xij1,0) )
```

#### Cation concentration at interfaces

```python
(pnp.concentration[0,0],pnp.concentration[0,-1])
```

#### Anion concentration at interfaces

```python
(pnp.concentration[1,0],pnp.concentration[1,-1])
```

#### Equilibrium cation and anion amount

```python
( pnp.numberConservationConstraint(pnp.xij1,0,0), pnp.numberConservationConstraint(pnp.xij1,1,0) )
```

#### Initial cation and anion amount

```python
( pnp.numberConservationConstraint(pnp.xi0,0,0), pnp.numberConservationConstraint(pnp.xi0,1,0) )
```

#### Species conservation

```python
(pnp.numberConservationConstraint(pnp.xij1,0,
                                 pnp.numberConservationConstraint(pnp.xi0,0,0)),
 pnp.numberConservationConstraint(pnp.xij1,1,
                                 pnp.numberConservationConstraint(pnp.xi0,1,0)) )
```

## Test case 5: 1D electrochemical cell, non-trivial initial values

```python
del pnp
```

```python
pnp = PoissonNernstPlanckSystem(delta_u=0.05)
```

```python
# technical settings
pnp.output = True # makes Newton solver display convergence plots
pnp.N = 800 # uniformlya distanced grid points
pnp.maxit = 20 # maximum number of Newton iterations
pnp.outfreq=2
pnp.e = 1e-10 # absolute tolerance
```

```python
pnp.useStandardCellBC()
```

```python
pnp.init()
```

```python
plt.plot(ci0.T) # use somewhat arbitrary (normalized) superposition as initial values for concentrations
```

```python
pnp.ni0 = ci0 / pnp.c_unit
```

```python
ui0 = pnp.initial_values()
```

```python
plt.plot(ui0)
```

```python
xij = pnp.solve()
```

```python
fig, (ax1,ax3) = plt.subplots(nrows=2,ncols=1,figsize=[16,10])
ax1.set_xlabel('z [nm]')

ax1.plot(pnp.grid*1e9, pnp.potential, marker='', color='tab:red', label='potential, PNP', linewidth=1, linestyle='-')

ax1.set_ylabel('potential (V)')

ax2 = ax1.twinx()
ax2.plot(pnp.grid*1e9, pnp.c[0]*np.ones(pnp.grid.shape), label='initial concentration', color='grey', linestyle=':')
ax2.plot(pnp.grid*1e9, pnp.concentration[0], marker='', color='tab:orange', label='Na+, PNP', linewidth=2, linestyle='-')

ax2.plot(pnp.grid*1e9, pnp.concentration[1], marker='', color='tab:blue', label='Cl-, PNP', linewidth=2, linestyle='-')

ax2.set_ylabel('concentration (mM)')

ax1.axvline(x=pnp.lambda_D()*1e9, label='Debye Length', color='grey', linestyle=':')

ax3.semilogy(pnp.grid*1e9, pnp.concentration[0], marker='', color='tab:orange', label='Na+, PNP', linewidth=2, linestyle='-')
ax3.semilogy(pnp.grid*1e9, pnp.concentration[1], marker='', color='tab:blue', label='Cl-, PNP', linewidth=2, linestyle='-')

ax3.set_ylabel('concentration (mM)')

ax2.legend(loc='upper center', fontsize=15)
ax1.legend(loc='lower center', fontsize=15)

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

#### Potential at left and right hand side of domain

```python
(pnp.potential[0],pnp.potential[-1])
```

#### Residual cation flux at interfaces

```python
( pnp.leftFluxBC(pnp.xij1,0), pnp.rightFluxBC(pnp.xij1,0) )
```

#### Residual anion flux at interfaces

```python
(pnp.leftFluxBC(pnp.xij1,1),  pnp.rightFluxBC(pnp.xij1,0) )
```

#### Cation concentration at interfaces

```python
(pnp.concentration[0,0],pnp.concentration[0,-1])
```

#### Anion concentration at interfaces

```python
(pnp.concentration[1,0],pnp.concentration[1,-1])
```

#### Equilibrium cation and anion amount

```python
( pnp.numberConservationConstraint(pnp.xij1,0,0), pnp.numberConservationConstraint(pnp.xij1,1,0) )
```

#### Initial cation and anion amount

```python
( pnp.numberConservationConstraint(pnp.xi0,0,0), pnp.numberConservationConstraint(pnp.xi0,1,0) )
```

#### Species conservation

```python
(pnp.numberConservationConstraint(pnp.xij1,0,
                                 pnp.numberConservationConstraint(pnp.xi0,0,0)),
 pnp.numberConservationConstraint(pnp.xij1,1,
                                 pnp.numberConservationConstraint(pnp.xi0,1,0)) )
```

## Application of 1D electrochemical cell model:

We want to fill a gap of 3 nm between gold electrodes with 0.2 wt % NaCl aqueous solution, apply a small potential difference and generate an initial configuration for LAMMPS within a cubic box:

```python
box=np.array([30.,30.,30.]) # Angstrom
```

```python
vol = box.prod() # Angstrom^3
```

With a concentration of 0.2 wt %, we are close to NaCl's solubility limit in water.
We estimate molar concentrations and atom numbers in our box:

```python
# enter number between 0 ... 0.2
weight_concentration_NaCl = 0.2 # wt %
# calculate saline mass density g/cm³
saline_mass_density       = 1 + weight_concentration_NaCl * 0.15 / 0.20 # g / cm^3
# see https://www.engineeringtoolbox.com/density-aqueous-solution-inorganic-sodium-salt-concentration-d_1957.html
```

```python
# calculate total saline mass, volume in angstrom!
saline_mass = (vol * 10**(-24) * saline_mass_density)
# rough numbers: 1 gram/cm^3 / (18g/mol) = 1/18 mol/cm³ = 6/18*10^(23)/cm³ = 1/3*0.1/angstrom³=1/30 atom per cubic angstrom
water_molar_mass = 18.015 # g / mol
# isotonic solution: 9 g NaCl / liter_H2O -> 1kg H2O == 1kg/0.018kg/mole=55.5555 mole; 9g NaCl / (22.9+35.5)g/mole = 0.15411 mole -->>> 1 NaCl for every 360 waters
salt_molar_mass  = 58.44 # g / mol

#calculate mixed molar mass
mixed_molar_mass = ( (1-weight_concentration_NaCl) * water_molar_mass + weight_concentration_NaCl * salt_molar_mass )

#calculate total number of molecules
number_NaCl      = int(np.round(saline_mass * weight_concentration_NaCl / salt_molar_mass / 2 * sc.Avogadro))
number_water     = int(np.round(saline_mass * (1 - weight_concentration_NaCl) / water_molar_mass * sc.Avogadro))
```

```python
cNaCl = number_NaCl / (vol*1e-30) / sc.Avogadro # mol / m^3, mM
```

```python
cNaCl / 1000 # mol / L, M
```

    1.968046270135898

```python
try:
    del pnp
except:
    pass
```

```python
pnp = PoissonNernstPlanckSystem(delta_u=0.05,c=[20,20],L=box[2]*1e-9)
```

    [ poisson_nernst_planck_distribution.py:667 - __init__() ]: ion species 00 concentration c                 20
    [ poisson_nernst_planck_distribution.py:669 - __init__() ]: ion species 00 number charge z                  1
    [ poisson_nernst_planck_distribution.py:667 - __init__() ]: ion species 01 concentration c                 20
    [ poisson_nernst_planck_distribution.py:669 - __init__() ]: ion species 01 number charge z                 -1
    [ poisson_nernst_planck_distribution.py:672 - __init__() ]: temperature T                               298.1
    [ poisson_nernst_planck_distribution.py:674 - __init__() ]: domain size L                               3e-08
    [ poisson_nernst_planck_distribution.py:676 - __init__() ]: potential difference delta_u                 0.05
    [ poisson_nernst_planck_distribution.py:678 - __init__() ]: relative permittivity eps_R                    79
    [ poisson_nernst_planck_distribution.py:680 - __init__() ]: vacuum permittivity eps_0                 8.854e-12
    [ poisson_nernst_planck_distribution.py:682 - __init__() ]: universal gas constant R                    8.314
    [ poisson_nernst_planck_distribution.py:684 - __init__() ]: Faraday constant F                        9.649e+04
    [ poisson_nernst_planck_distribution.py:686 - __init__() ]: f = F / (RT)                                38.92
    [ poisson_nernst_planck_distribution.py:702 - __init__() ]: spatial unit [l]                          2.158e-09
    [ poisson_nernst_planck_distribution.py:704 - __init__() ]: concentration unit [c]                         20
    [ poisson_nernst_planck_distribution.py:706 - __init__() ]: potential unit [u]                        0.02569
    [ poisson_nernst_planck_distribution.py:724 - __init__() ]: reduced domain size L*                       13.9
    [ poisson_nernst_planck_distribution.py:728 - __init__() ]: ion species 00 reduced concentration c*         1
    [ poisson_nernst_planck_distribution.py:728 - __init__() ]: ion species 01 reduced concentration c*         1
    [ poisson_nernst_planck_distribution.py:730 - __init__() ]: reduced potential delta_u*                  1.946

```python
# technical settings
pnp.output = True # makes Newton solver display convergence plots
pnp.N = 1600 # uniformlya distanced grid points
pnp.maxit = 20 # maximum number of Newton iterations
pnp.outfreq=2
pnp.e = 1e-10 # absolute tolerance
```

```python
pnp.useStandardCellBC()
```

    [ poisson_nernst_planck_distribution.py:381 - useStandardCellBC() ]: Left hand side Dirichlet boundary condition:                               u0 =    0.973
    [ poisson_nernst_planck_distribution.py:382 - useStandardCellBC() ]: Right hand side Dirichlet boundary condition:                              u1 =   -0.973
    [ poisson_nernst_planck_distribution.py:393 - useStandardCellBC() ]: Ion species 00 left hand side concentration Flux boundary condition:       j0 =        0
    [ poisson_nernst_planck_distribution.py:395 - useStandardCellBC() ]: Ion species 00 number conservation constraint:                             N0 =     13.9
    [ poisson_nernst_planck_distribution.py:393 - useStandardCellBC() ]: Ion species 01 left hand side concentration Flux boundary condition:       j0 =        0
    [ poisson_nernst_planck_distribution.py:395 - useStandardCellBC() ]: Ion species 01 number conservation constraint:                             N0 =     13.9

```python
pnp.init()
```

    [ poisson_nernst_planck_distribution.py:246 - init() ]: discretization segments N                    1600
    [ poisson_nernst_planck_distribution.py:248 - init() ]: grid points N                                1601
    [ poisson_nernst_planck_distribution.py:263 - init() ]: dx                                        0.008689
    [ poisson_nernst_planck_distribution.py:265 - init() ]: dx2overtau                                     10

```python
xij = pnp.solve()
```

```python
fig, (ax1,ax3) = plt.subplots(nrows=2,ncols=1,figsize=[16,10])
ax1.set_xlabel('z [nm]')

ax1.plot(pnp.grid*1e9, pnp.potential, marker='', color='tab:red', label='potential, PNP', linewidth=1, linestyle='-')
ax1.set_ylabel('potential (V)')

ax2 = ax1.twinx()
ax2.plot(pnp.grid*1e9, pnp.c[0]*np.ones(pnp.grid.shape), label='initial concentration', color='grey', linestyle=':')
ax2.plot(pnp.grid*1e9, pnp.concentration[0], marker='', color='tab:orange', label='Na+, PNP', linewidth=2, linestyle='-')
ax2.plot(pnp.grid*1e9, pnp.concentration[1], marker='', color='tab:blue', label='Cl-, PNP', linewidth=2, linestyle='-')
ax2.set_ylabel('concentration (mM)')

ax1.axvline(x=pnp.lambda_D()*1e9, label='Debye Length', color='grey', linestyle=':')

ax3.semilogy(pnp.grid*1e9, pnp.concentration[0], marker='', color='tab:orange', label='Na+, PNP', linewidth=2, linestyle='-')
ax3.semilogy(pnp.grid*1e9, pnp.concentration[1], marker='', color='tab:blue', label='Cl-, PNP', linewidth=2, linestyle='-')

ax3.set_ylabel('concentration (mM)')

ax2.legend(loc='upper center', fontsize=15)
ax1.legend(loc='lower center', fontsize=15)

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

#### Potential at left and right hand side of domain

```python
(pnp.potential[0],pnp.potential[-1])
```

#### Residual cation flux at interfaces

```python
( pnp.leftFluxBC(pnp.xij1,0), pnp.rightFluxBC(pnp.xij1,0) )
```

#### Residual anion flux at interfaces

```python
(pnp.leftFluxBC(pnp.xij1,1),  pnp.rightFluxBC(pnp.xij1,0) )
```

#### Cation concentration at interfaces

```python
(pnp.concentration[0,0],pnp.concentration[0,-1])
```

#### Anion concentration at interfaces

```python
(pnp.concentration[1,0],pnp.concentration[1,-1])
```

#### Equilibrium cation and anion amount

```python
( pnp.numberConservationConstraint(pnp.xij1,0,0), pnp.numberConservationConstraint(pnp.xij1,1,0) )
```

#### Initial cation and anion amount

```python
( pnp.numberConservationConstraint(pnp.xi0,0,0), pnp.numberConservationConstraint(pnp.xi0,1,0) )
```

#### Species conservation

```python
(pnp.numberConservationConstraint(pnp.xij1,0,
                                 pnp.numberConservationConstraint(pnp.xi0,0,0)),
 pnp.numberConservationConstraint(pnp.xij1,1,
                                 pnp.numberConservationConstraint(pnp.xi0,1,0)) )
```

## Sampling
Create interpolation functions from concentrations on grid points

```python
na_distribution = interpolate.interp1d(pnp.X*pnp.l_unit, pnp.concentration[0])
cl_distribution = interpolate.interp1d(pnp.X*pnp.l_unit, pnp.concentration[1])
```

Now we can sample the distribution of our $Na^+$ ions in z-direction.

```python
scaled_na_distribution = lambda x: na_distribution(x) / na_distribution(0) * 0.1
```

```python
scaled_cl_distribution = lambda x: cl_distribution(x) / cl_distribution(0) * 0.1
```

```python
x = 50e-9
z = pnp.L
box = np.array([x, x, z])

n_bins = 50 # for histogram plots
```

```python
na_sample_size = int(np.round(box.prod()*pnp.c[0] * sc.Avogadro)) # [V] = m^3, [c] = mol / m^3, [N_A] = 1 / mol
```

```python
cl_sample_size = int(np.round(box.prod()*pnp.c[1] * sc.Avogadro)) # [V] = m^3, [c] = mol / m^3, [N_A] = 1 / mol
```

```python
na_sample_size
```

```python
na_struc = generate_structure(
    distribution=scaled_na_distribution, box=box, atom_count=na_sample_size )

histx, histy, histz = get_histogram(na_struc, box=box, n_bins=n_bins)
```

```python
plot_dist(histz, 'Distribution of Na+ ions in z-direction', reference_distribution=scaled_na_distribution)
```

```python
plot_dist(histx, 'Distribution of Na+ ions in x-direction', reference_distribution=(lambda x: np.array([1.0/n_bins]*len(x))))
```

```python
cl_struc = generate_structure(distribution=scaled_cl_distribution, box=box, atom_count=cl_sample_size)
histx, histy, histz = get_histogram(cl_struc, box=box, n_bins=n_bins)
```

```python
plot_dist(histz, 'Distribution of Cl- ions in z-direction', reference_distribution=scaled_cl_distribution)
```

```python
plot_dist(histy, 'Distribution of Cl- ions in y-direction',reference_distribution=(lambda x: np.array([1.0/n_bins]*len(x))))
```

## Write to file
Utilize ASE to write to several standard formats:

```python
na_atoms = ase.Atoms(
    symbols='Na'*na_sample_size,
    charges=[1]*na_sample_size,
    positions=na_struc/sc.angstrom,
    cell=box/sc.angstrom,
    pbc=[1,1,0])
```

```python
cl_atoms = ase.Atoms(
    symbols='Cl'*cl_sample_size,
    charges=[-1]*cl_sample_size,
    positions=cl_struc/sc.angstrom,
    cell=box/sc.angstrom,
    pbc=[1,1,0])
```

```python
system = na_atoms + cl_atoms
```

```python
system
```

```python
ase.io.write('NaCl_0.1mM_0.05V_50x50x100nm_nanogap.xyz',system,format='xyz')
```

```python
# LAMMPS data format, units 'real', atom style 'full'
# weird behavior of ase output
with open('NaCl_0.1mM_0.05V_50x50x100nm_nanogap.lammps','wb') as f:
    ase.io.write(f,system,format='lammps-data',units="real",atom_style='full')
```

## Excursus: Constraints (not implemented)

The following snippet on Lagrange multipliers originates from a point charge optimization context and has not been adapted yet

According to https://theochem.github.io/horton/2.1.0/user_postproc_espfit.html, cost function is constructed as

$$ c(\mathbf{q}, \Delta V_\text{ref}) = \int_V d\mathbf{r} \omega(\mathbf{r}) \cdot \left( V_\text{ab initio}(\mathbf{r}) - \sum_{i=1}^N \frac{q_i}{\mathbf{r} - \mathbf{R}_i} - \Delta V_\text{ref} \right)^2  $$

> $\Delta V_\text{ref}$: constant that may account for differences in reference between the ab initio ESP and the point charge ESP. The need for such a term was established in the REPEAT paper for ESP fitting in 3D periodic systems.

We look at an aperiodic system, neglect this term for now. The (unconstrained) cost function takes the general quadratic form

$$ c_u(\mathbf{x}) = \mathbf{x}^T A\ \mathbf{x} - 2 \mathbf{b}^T \mathbf{x} - C $$

with

$$ C = - \int_V d\mathbf{r} \omega(\mathbf{r}) \cdot \left[ V_\text{ab initio}(\mathbf{r}) \right]^2$$

entry i,j of matrix $A^{N\times N}$

$$  A_{ij} = \int_V d\mathbf{r}\ \omega(\mathbf{r}) \cdot \left( \frac{1}{|\mathbf{r} - \mathbf{R}_i|} \cdot \frac{1}{|\mathbf{r} - \mathbf{R}_j|} \right)  $$

and entry i of vector $\mathbf{b}$

$$  b_i = \int_V d\mathbf{r}\ \omega(\mathbf{r}) \cdot \frac{V_\text{ab initio}(\mathbf{r})}{|\mathbf{r} - \mathbf{R}_i|} $$

In the code below, first the miniumum of an unconstrained system is found by solving

$$ \frac{1}{2} \cdot \frac{\mathrm{d} c(\mathbf{x})}{\mathrm d  \mathbf{x}}  = \frac{1}{2} \cdot \nabla_\mathbf{x} c = A \mathbf{x} - \mathbf{b} = 0$$

with $\nabla (\mathbf{b}^T \mathbf{x}) = \mathbf{b}$ and
$\nabla \cdot (\mathbf{x}^T A \mathbf{x}) = (A + A^T) \mathbf{x}
= 2 A \mathbf{x} $ for symmetric A, as in our case.
The (unconstrained) solution

$$ \mathbf{x}_u = A^{-1} \mathbf{b} $$

is corrected for *one* total charge constraint of the form

$$ g(\mathbf{x}) = \mathbf{d} \cdot \mathbf{x} - q_\text{tot} = 0 $$

with all entries of $\mathbf{d}$ unity. Notice that in the code below, the whole system is normalized in order to have unity diagonal entries $A_{jj}$. We neglect this normalization here.

A Lagrange multiplier $\lambda$ is introduced into the *constrained* cost function

$$ c_c(\mathbf{x},\lambda) = \mathbf{x}^T A\ \mathbf{x} - 2 \mathbf{b}^T \mathbf{x} - C + \lambda \cdot  g(\mathbf{x}) $$

and the system

$$ \frac{1}{2} \cdot \nabla_\mathbf{x} c = A \mathbf{x} - \mathbf{b} + \lambda \cdot \nabla_\mathbf{x} g(\mathbf{x}) = A \mathbf{x} - \mathbf{b} +\lambda \cdot \mathbf{d} = 0 $$

$$  \nabla_\mathbf{\lambda} c = g(\mathbf{x}) = \mathbf{d} \cdot \mathbf{x} - q_\text{tot} = 0 $$

is solved by finding a correction for the unconstrained solution

$$ \mathbf{x} = \mathbf{x}_u - \lambda \cdot\delta \mathbf{x} = A^{-1} \mathbf{b} - \lambda \cdot \delta \mathbf{x} $$

as

$$ - \lambda A \delta \mathbf{x} + \lambda \cdot \mathbf{d} = 0
\Leftrightarrow \delta \mathbf{x} = A^{-1} \mathbf{d}$$

and the Lagrange multiplier by

$$ g(\mathbf{x})
    = \mathbf{d} \cdot \left( \mathbf{x}_u
        - \lambda \ \delta \mathbf{x} \right)- q_\text{tot}
    = \mathbf{d} \cdot A^{-1} \mathbf{b}
        - \lambda \ \mathbf{d} \cdot \delta \mathbf{x} - q_\text{tot}
    = 0 $$

$$ \lambda = \frac{\mathbf{d} \cdot A^{-1} \mathbf{b} - q_\text{tot}}
    { \mathbf{d} \cdot \delta \mathbf{x} }
    = \frac{\mathbf{b} \cdot \delta \mathbf{x} - q_\text{tot}}
    { \mathbf{d} \cdot \delta \mathbf{x} }$$

and thereby the constrained minimum at
$$ \mathbf{x}
    = \mathbf{x}_u
        - \frac{\mathbf{b} \cdot \delta \mathbf{x} - q_\text{tot}}
        { \mathbf{d} \cdot \delta \mathbf{x} }
        \cdot \delta \mathbf{x}$$

as implemented in HORTON

We modifiy the optimization in order to allow for an arbitrary amount of constraints.

M lagrange multipliers $\lambda_k$ are introduced into the *constrained* cost function

$$ c_c(\mathbf{x},\mathbf{\lambda}) = \mathbf{x}^T A\ \mathbf{x} - 2 \mathbf{b}^T \mathbf{x} - C + \sum_{k=1}^M \lambda_k \cdot  g_k(\mathbf{x}) $$

All our constraints (charge groups and symmetries) will be of linear form

$$ g_k(\mathbf{x}) = \mathbf{d}_k \cdot \mathbf{x} - q_k = 0 $$

and thus can be compacted into matrix form

$$ D^{(M \times N)} \mathbf{x} - \mathbf{q}^{(M\times 1)} = 0 $$

with

$$ D^T = [\mathbf{d}_1, \mathbf{d}_2, \dots , \mathbf{d}_M] $$

and hence
$$ c_c(\mathbf{x}^{(N\times1)},\mathbf{\lambda}^{(M\times1)}) = \mathbf{x}^T A\ \mathbf{x}
    - 2 \mathbf{b}^T \mathbf{x}
    - C + \mathbf{\lambda}^T \cdot \left( D \mathbf{x} - \mathbf{q} \right) $$

Derivative

$$
\begin{align}
    \nabla_\mathbf{x} \cdot c_c & = 2\ A\ \mathbf{x}
    + \sum_{k=1}^M \lambda_k \mathbf{d}_k - 2 \mathbf{b} & = 0\\
    \nabla_\mathbf{\lambda} \cdot c_c & = D\ \mathbf{x} - \mathbf{q} & = 0
\end{align}
$$

Identify

$$ D^T \mathbf{\lambda}= \sum_{k=1}^M \lambda_k \mathbf{d}_k  $$

and solve

$$ \tilde{A} \mathbf{y} - \tilde{\mathbf{b}} = 0$$

with generalized $\mathbf{y}^T = [\mathbf{x}^T, \mathbf{\lambda}^T ]$
as well as $(N+M)\times(N+M)$ matrix

$$ \tilde{A} =
    \begin{bmatrix}
        2 A & D^T \\
        D & 0
    \end{bmatrix} $$

and $(N+M)$ vector

$$ \tilde{\mathbf{b}} =
   \begin{bmatrix}
        2 \mathbf{b} \\
        \mathbf{q}
   \end{bmatrix} $$

```python

```


```python
import numpy as np
import pandas as pd
import linearsolve as ls
import matplotlib.pyplot as plt
plt.style.use('classic')
%matplotlib inline
```

# Homework 8

**Instructions:** Complete the notebook below. Download the completed notebook in HTML format. Upload assignment using Canvas.

**Due:** Mar. 12 at **12:30pm.**

## Exercise: New-Keynesian Model Stochastic Simulation

### Equilibrium Conditions and Paramter Values

The most basic version of the New-Keynesian Model can be expressed as:

\begin{align}
y_t & = E_t y_{t+1}  - \left( r_{t} - \bar{r}\right) + g_t\\
i_{t} & = r_{t} + E_t \pi_{t+1}\\
i_{t} & = \bar{r} + \pi^T + \phi_{\pi}\big(\pi_t - \pi^T\big) + \phi_{y}\big(y_t - \bar{y}\big)\\
\pi_t -\pi^T & = \beta \left( E_t\pi_{t+1} - \pi^T\right)  + \kappa (y_t -\bar{y})+ u_t,
\end{align}

where: $y_t$ is (log) output, $r_t$ is the real interest rate, $i_t$ is the nominal interest rate, $\pi_t$ is the rate of inflation between periods $t-1$ and $t$, $\bar{r}$ is the long-run average real interest rate or the *natural rate of interest*, $\beta$ is the household's subjective discount factor, and $\pi^T$ is the central bank's inflation target. The coeffieints $\phi_{\pi}$ and $\phi_{y}$ reflect the degree of intensity to which the central bank *endogenously* adjusts the nominal interest rate in response to movements in inflation and output.

The variables $g_t$ and $u_t$ represent exogenous shocks to aggregate demand and inflation. They follow AR(1) processes:

\begin{align}
g_{t+1} & = \rho_g g_{t} + \epsilon^g_{t+1}\\
u_{t+1} & = \rho_u u_{t} + \epsilon^u_{t+1}
\end{align}

### Parameter Values:

You will use the following parameter values:

| $\bar{y}$ | $\beta$ | $\bar{r}$    | $\kappa$ | $\pi^T$ | $\phi_{\pi}$ | $\phi_y$ | $\rho_g$ | $\sigma_g^2$ | $\rho_u$ | $\sigma_u^2$ |
|-----------|---------|--------------|----------|---------|--------------|----------|----------|--------------|----------|--------------|
| 0         | 0.995   | $-\log\beta$ | 0.1      | 0.02/4  | 1.5          | 0.5/4    | 0.5      | 0.0022         | 0.5      | 0.0012        |

### Input Model and Solve

Refer to the Notebook for the week 9 discussion section for a complete example of how to input the model and solve.

```python
# Create a variable called 'parameters' that stores the model parameter values in a Pandas Series.

# Create variable called 'varNames' that stores the variable names in a list with state variables ordered first

# Create variable called 'shockNames' that stores an exogenous shock name for each state variable.

# Define a function that evaluates the equilibrium conditions of the model solved for zero.

    # Parameters

    # Current variables

    # Forward variables

    # IS equation

    # Fisher_equation

    # Monetary policy

    # Phillips curve

    # Demand process

    # Inflation process

    # Stack equilibrium conditions into a numpy array

# Initialize the model into a variable named 'nk_model'

# Compute the steady state numerically using .compute_ss() method of nk_model

# Approximate and solve the model. Set loglinear argument to False since the model is already linear

```

### Stochastic Simulation

Construct a stochastic simulation of the New-Keynesian model with the follwoing properties:

1. 201 periods.
2. Seed for random number generator: 126.

Note that the shock covariance matrix is:

\begin{align}
\text{Covariance matrix} & = \left[\begin{array}{cc}\sigma_g^2 & 0\\ 0 & \sigma_u^2\end{array}\right]
\end{align}

Use the following values for $\sigma_g^2$ and $\sigma_u^2$ in the simulation:

| $\sigma_g^2$      | $\sigma_u^2$      |
|-------------------|-------------------|
| 0.0022 | 0.0012 |

Refer to the Notebook for the week 7 discussion section for an example of how to use `linearsolve` to construct a stochastic simulation.

```python
# Compute the simulation

```

### Analyze Simulation Results

Construct a plot with simulated output and inflation plotted together. Multiply simulated output by 100 and simulated inflation by 400 since, by convention, we always annualize inflation and interest rates.

```python
# Create a figure with dimensions 12x4. PROVIDED
fig = plt.figure(figsize=(12,4))

# Create the left axis. PROVIDED
ax1 = fig.add_subplot(1,1,1)

# Plot the simulated series for output (times 100) and inflation (times 400)

# Construct legend. PROVIDED
ax1.legend(loc='center left', bbox_to_anchor=(1, 0.5))
```

Compute the standard deviations of inflation $\pi_t$ (times 400), output $y_t$ (times 100), and the nominal interest rate $i_t$ (times 400) for the simulated series.

```python

```

Compute the coefficients of correlation of inflation $\pi_t$, output $y_t$, and the nominal interest rate $i_t$ for the simulated series.

```python

```

**Questions**

1. In the smulation data, which variable fluctuates more, inflation (times 400) or output (times 100)?
2. Do the simulations suggest a positive or negative correlation between output and inflation?
3. Note that for the simulations, it was assumed that the variances on the inflation and demand shocks were about the same. Based on the answer to your previous question, which shocks appear to have a dominant effect on the dynamics of the model?

**Answers**

1.

2.

3.


# A TUTORIAL ON HAMINTONIAN MONTE CARLO

by Sebastian T. Glavind, January, 2022

```python
import jax # automatic differentiation
import jax.numpy as jnp
import jax.scipy.stats as jss

import numpy as np
import scipy
import scipy.stats as ss

from tqdm import tqdm # progress bars
from matplotlib import pyplot as plt
%matplotlib inline
```

```python
print('numpy: ', np.__version__,', scipy: ', scipy.__version__, ', jax: ', jax.__version__)
```

    numpy:  1.21.2 , scipy:  1.7.3 , jax:  0.2.26

# Hamintonian Monte Carlo

## Introduction

Hamiltonian Monte Carlo (HMC) is an effifient algorithm for performing Markov chain Monte Carlo (MCMC) inference, which uses gradients of the log-posterior distribution to guide the generation of new proposal states - the gradient at a given state provides information on the geometry of the posterior density function. Based on the gradient information, HMC can propose new states far from the current state with high acceptance probability by solving the so-called Hamiltonian equations, and it can thereby avoid the typical MCMC (local) random walk behaiviour. This enables HMC to scale to higher dimensions.

The basis for HMC is to consider the sampling domain as a conservative dynamical system, where the total energy - in terms of kinetic energy (momentum) and potential energy (position) - is preserved, i.e., if the system gains kinetic energy then is loses the same amount of potential energy. In a physical setting, the potential energy is due to the pull of gravity, and the momentum is due to the motion of the particle. As an example, consider a particle with zero momentum on a declining surface; it will be pulled downwards by gravity, and its potential energy (from its initial height) will be transferred into kinetic energy (motion). Conversely, if the particle is moving along a flat plain, and it encounters a rise, it will slow down as it ascends, transferring its kinetic energy into potential energy.

We thus consider the so-called momentum-position phase space, which can be represended by the following, convenient joint probability distribution:

$$p(\mathbf{p}, \mathbf{q}) = p(\mathbf{p} | \mathbf{q}) p(\mathbf{q}).$$

In this representation, we can immediately recover the target distribution for the position (parameters) $p(\mathbf{q})$ by marginalizing out the momentum $\mathbf{p}$, and thus $\mathbf{p}$ is regarded as an auxilary variable.

The joint distribution, also called the canonical density, for such a system may also be represented in terms of the invariant Hamiltonian function:

$$ p(\mathbf{p}, \mathbf{q}) = \exp( -H(\mathbf{p}, \mathbf{q}) ), $$

where the value of the Hamiltonian function at any point in phase space represents the total energy at that point. Thus, as the joint density decomposes, the Hamiltonian function may be written:

$$ H(\mathbf{p}, \mathbf{q}) = - \log p(\mathbf{p}, \mathbf{q}) = \underbrace{- \log p(\mathbf{p} | \mathbf{q})}_{K(\mathbf{p}, \mathbf{q})} \underbrace{- \log p(\mathbf{q})}_{V(\mathbf{q})}  =  K(\mathbf{p}, \mathbf{q}) + V(\mathbf{q}), $$

where $K(\mathbf{p}, \mathbf{q})$ resembles the kinetic energy and $V(\mathbf{q})$ resembles the potential energy.

The trajectory of a particle within an energy level set can be obtained (simulated) by solving the following continuous time differential equations, known as Hamilton’s equations;

$$
\require{xcancel}
\begin{align}
\frac{d\mathbf{q}}{dt} &= \frac{ \partial H }{\partial \mathbf{p}} = \frac{ \partial K }{\partial \mathbf{p}} + \xcancel{ \frac{ \partial V }{\partial \mathbf{p}} } \\
\frac{d\mathbf{p}}{dt} &= - \frac{ \partial H }{\partial \mathbf{q}} = - \frac{ \partial K }{\partial \mathbf{q}} - \frac{ \partial V }{\partial \mathbf{q}},
\end{align}
$$

where $\partial V / \partial \mathbf{p}$ cancels, as $V$ does not depend on $\mathbf{p}$. This mapping is volume preserving, i.e., it has a Jacobian determinant of 1, a facts that will is important when turning then system into an MCMC algorithm.

If we now chose the kinetic energy to be a Gaussian and drop the normalization constant, we have:

$$
K(\mathbf{p}, \mathbf{q}) = \frac{1}{2} \mathbf{p}^T M^{-1} \mathbf{p} + \log |M|,
$$

where M is called the mass matrix in the HMC literature. We scose this to be the identity matrix, i.e., $M=I$, which leads to

$$
K(\mathbf{p}, \mathbf{q}) = \frac{1}{2} \mathbf{p}^T \mathbf{p},
$$

which simplifies the Hamiltonian equations to

$$
\begin{align}
\frac{d\mathbf{q}}{dt} &= \mathbf{p} \\
\frac{d\mathbf{p}}{dt} &= - \frac{ \partial V }{\partial \mathbf{q}},
\end{align}
$$

as $\partial K / \partial \mathbf{p}=\mathbf{p}$ and $\partial K / \partial \mathbf{q}=0$ for the Gaussian definition of $K(\mathbf{p}, \mathbf{q})$ above; see e.g., Hoffman and Gelman (2014), Betancourt (2017), Martin et al (2021) and Murphy (2023) for further details.

A simulation procedure is now taken shape; thus, we can sample $\mathbf{p}_0 \sim \mathcal{N}(0,I)$ and simulate a trajectory $\{ \mathbf{q}_t, \mathbf{p}_t \}_{t=1}^T$ for some amount of time $T$, where $\mathbf{q}_T$ is then our new sample. In the following section, we will study how we simulate the trajectory $\{ \mathbf{q}_t, \mathbf{p}_t \}_{t=1}^T$ by solving Hamilton’s equations.

***
Hoffman, Matthew D., and Andrew Gelman. "The No-U-Turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo." J. Mach. Learn. Res. 15.1 (2014): 1593-1623 ([link](https://www.jmlr.org/papers/volume15/hoffman14a/hoffman14a.pdf)).

Betancourt, Michael. "A conceptual introduction to Hamiltonian Monte Carlo." arXiv preprint arXiv:1701.02434, 2017 ([link](https://arxiv.org/pdf/1701.02434.pdf?ref=https://githubhelp.com))

Martin, Osvaldo A. and Kumar, Ravin and Lao, Junpeng, "Bayesian Modeling and Computation in Python", CRC, 2021.

Murphy, Kevin P., "Probabilistic Machine Learning: Advanced Topics", MIT Press, 2023 ([draft](https://probml.github.io/pml-book/book2.html) from feb. 2022).
***

## Solving Hamilton's equations

In this section, we discuss how to simulate trajectories from Hamilton’s equations in discrete time.

### Euler's method

One simple way to simulate these trajectories is to update the position and momentum simultaneously by a small amount based on the gradients:

$$
\begin{align}
\mathbf{p}_{t+1} &= \mathbf{p}_{t} + \eta \frac{d\mathbf{p}}{dt}(\mathbf{q}_t,\mathbf{p}_t) = \mathbf{p}_{t} - \eta \frac{ \partial V }{\partial \mathbf{q}}(\mathbf{q}_t) \\
\mathbf{q}_{t+1} &= \mathbf{q}_{t} + \eta \frac{d\mathbf{q}}{dt}(\mathbf{q}_t,\mathbf{p}_t) = \mathbf{q}_{t} + \eta \mathbf{p}_{t},
\end{align}
$$

where $\eta$ is the step size, and we have included the simple definition for kinetic energy above in the second expression. Unfortunally, though, Euler’s method does not preserve volume, and can lead to inaccurate approximations after only a few steps.

### Modified Euler's method

In Modified Euler's method, we make Euler's method symplectic (volume preserving) by first updating the momentum, and then we update the position using the new momentum:

$$
\begin{align}
\mathbf{p}_{t+1} &= \mathbf{p}_{t} + \eta \frac{d\mathbf{p}}{dt}(\mathbf{q}_t,\mathbf{p}_t) = \mathbf{p}_{t} - \eta \frac{ \partial V }{\partial \mathbf{q}}(\mathbf{q}_t) \\
\mathbf{q}_{t+1} &= \mathbf{q}_{t} + \eta \frac{d\mathbf{q}}{dt}(\mathbf{q}_t,\mathbf{p}_{t+1}) = \mathbf{q}_{t} + \eta \mathbf{p}_{t+1},
\end{align}
$$

Note that we can equivalently perform the updates in the opposite order. Unfortunally, this procedure is not reversible due to the asymmetry in the updates.

### Leapfrog integrator

Using the Leapfrog integrator, we can make Euler's method symplectic and reversible by first performing a half-update of the momentum, next we perform a full update of the position, and then we perform another
half-update of the momentum:

$$
\begin{align}
\mathbf{p}_{t+1/2} &= \mathbf{p}_{t} + \frac{\eta}{2} \frac{d\mathbf{p}}{dt}(\mathbf{q}_t,\mathbf{p}_t) = \mathbf{p}_{t} - \frac{\eta}{2} \frac{ \partial V }{\partial \mathbf{q}}(\mathbf{q}_t) \\
\mathbf{q}_{t+1} &= \mathbf{q}_{t} + \eta \frac{d\mathbf{q}}{dt}(\mathbf{q}_t,\mathbf{p}_{t+1/2}) = \mathbf{q}_{t} + \eta \mathbf{p}_{t+1/2} \\
\mathbf{p}_{t+1} &= \mathbf{p}_{t+1/2} + \frac{\eta}{2} \frac{d\mathbf{p}}{dt}(\mathbf{q}_{t+1},\mathbf{p}_{t+1/2}) = \mathbf{p}_{t+1/2} - \frac{\eta}{2} \frac{ \partial V }{\partial \mathbf{q}}(\mathbf{q}_{t+1}),
\end{align}
$$

It can be shown that the Leapfrog integrator is volumen preserving, and given that we reverse the momentum at the end of the iteration, i.e., replace $\mathbf{p}$ by $-\mathbf{p}$, it is also reversible. Note that if the kinetic energy satisfies $K(\mathbf{p})=K(-\mathbf{p})$, the reversal of momentum is not needed to make the method reversible.

How effectively HMC explores the target distribution is determined by the integration length $T$, or equavantly the number of leapfrog steps $L$, at each iteration. One the one hand, if we integrate for only a short time then we do not take full advantage of the coherent exploration of the Hamiltonian trajectories, and the algorithm will exhibit random walk-like behaviour. On the other hand, trajectories will eventually return to previously explored neighborhoods, i.e., integrating too long can suffer from diminishing returns.

Unfortunally, the Leapfrog integrator does not exactly perserve energy, due to the finite step size. This, however, can be fixed by treating the method as a proposal distribution, and then we can use the Metropolis acceptance criterion to ensure that we sample form the target distribution. The Metropolis acceptance criterion for this case is defined as

$$
\mathbf{q}_{k+1} =
\begin{cases}
  \mathbf{q}_{T} & \text{with probability}\ \min\left( \left[p(\mathbf{q}_{T}, \mathbf{p}_{T}) \big/ p(\mathbf{q}_{k}, \mathbf{p}_0) \right], \ 1 \right) \\
  \mathbf{q}_{k} & \text{otherwise},
\end{cases}
$$

where $(\mathbf{q}_{T}, \mathbf{p}_{T})$ is our porposal, and $(\mathbf{q}_{k}, \mathbf{p}_{0})$ is the previous sample of $\mathbf{q}$ and the random, initial momentum for sample $k+1$, respectively. For more information on solving Hamilton's equations, see Murphy (2023), as well as this nice [blog post](https://bayesianbrad.github.io/posts/2019_hmc.html).

```python
# Code inspired by the implementation in the "Bayesian Modeling and Computation in Python" book
# https://bayesiancomputationbook.com/markdown/chp_11.html

def leapfrog(q, p, dVdq, path_len, step_size):
    p -= step_size * dVdq(q) / 2  # half step
    for _ in range(int(path_len / step_size) - 1):
        q += step_size * p  # whole step
        p -= step_size * dVdq(q)  # whole step
    q += step_size * p  # whole step
    p -= step_size * dVdq(q) / 2  # half step

    return q, -p  # momentum flip at end
```

## The HMC algorithm

Now that we have the boilding blocks, we are ready to summarize the HMC algorithm:

At each iteration, we

1. sample a momentum $\mathbf{p}_0 \sim \mathcal{N}(0,I)$ and

2. simulate a trajectory $\{ \mathbf{q}_t, \mathbf{p}_t \}_{t=1}^T$ for some amount of time $T$, i.e., $L$ Leapfrog steps.

3. chose $\mathbf{q}_T$ is then our new sample.

4. use the Metropolis acceptance criterion to judge whether $\mathbf{q}_T$ should be accepted or rejected as a sample from the traget distribution.

Note that in order the make the HMC implementation below general, we will use the autodiff library `jax`, whereby the implementation does not rely on our ability to provide the analytical gradients for the Leapfrog integrator.

```python
def hamiltonian_monte_carlo(
    n_samples, negative_log_prob, initial_position, path_len=1, step_size=0.5):
    # counter of acceptance
    counter = 0
    # jax autodiff magic - initialize
    dVdq = jax.grad(negative_log_prob)
    # collect all our samples in a list
    samples = [initial_position]
    # Keep a single object for momentum resampling
    momentum = ss.norm(0, 1)
    # If initial_position is a 10d vector and n_samples is 100, we want
    # 100 x 10 momentum draws. We can do this in one call to momentum.rvs, and
    # iterate over rows
    size = (n_samples,) + initial_position.shape[:1]
    for p0 in momentum.rvs(size=size):
        # Integrate over our path to get a new position and momentum
        q_new, p_new = leapfrog(
            samples[-1], p0, dVdq, path_len=path_len, step_size=step_size)
        # Check Metropolis acceptance criterion
        start_log_p = negative_log_prob(samples[-1]) - np.sum(momentum.logpdf(p0))
        new_log_p = negative_log_prob(q_new) - np.sum(momentum.logpdf(p_new))
        if np.log(np.random.rand()) < start_log_p - new_log_p:
            samples.append(q_new)
            counter += 1
        else:
            samples.append(np.copy(samples[-1]))

    return np.array(samples[1:]), counter
```

## Setting the step size by dual avaging

To set the step size $\epsilon$ in HMC (and its extension the no-uturn sampler (NUTS)), Hoffman and Gelman (2014) propose to use stochastic optimization with vanishing adaptation, specifically an extension of the primal-dual algorithm of Nesterov (2009).

Generally, we consider a statistic $H_k$ that describes some aspect of the behavior of an MCMC algorithm at iteration $k$, which has expectation $h(x)$, where $x \in \mathbb{R}$ is a tunable parameter of the MCMC algorithm. For our case, we will target the average Metropolis acceptance probability (see above) based on the criterion

$$
H_k = \delta - \alpha_k; \quad h(x) = \mathbb{E}[H_k|x]
$$

where $\delta$ is the desired average acceptance probability, and $\alpha_k$ is the Metropolis acceptance probability for iteration $k$. The following updating scheme:

$$
x_{k+1} = \mu - \frac{\sqrt(k)}{\gamma} \frac{1}{k + k_0} \sum_{i=1}^k H_i; \quad \bar{x}_{k+1} \leftarrow \eta_k x_{k+1} + (1 - \eta_k) \bar{x}_k
$$

than guarantees that the sequence of averaged iterates $\bar{x}_k$ converges to a value, such that $h(\bar{x}_k)$ converges to $0$ for $\eta_k = k^{-\kappa}$ with $\kappa \in (0.5, 1]$. Further, $\mu$ is a freely chosen point that the iterates $x_k$ are shrunk towards, $\gamma > 0$ is a free parameter that controls the amount of shrinkage towards $\mu$, $k_0 \geq 0$ is a free parameter that stabilizes the initial iterations of the algorithm, and we define $\bar{x}_1 = x_1$.

Hoffman and Gelman (2014) propose to use $x = \log \epsilon$ with parameters $\gamma = 0.05$, $t_0 = 10$, $\kappa = 0.75$, and $\mu = \log(10 \epsilon_0)$, where $\epsilon_0$ is the initial step size. Moreover, it has been shown that the optimal value of $\epsilon$ for a given simulation length $\epsilon L$ is the one that produces an average Metropolis acceptance probability of approximately $0.65$, thus we set $\delta=0.65$. Please refer to Hoffman and Gelman (2014) for further details.

***
Hoffman, Matthew D., and Andrew Gelman. "The No-U-Turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo." J. Mach. Learn. Res. 15.1 (2014): 1593-1623 ([link](https://www.jmlr.org/papers/volume15/hoffman14a/hoffman14a.pdf)).

Nesterov, Yurii. "Primal-dual subgradient methods for convex problems." Mathematical programming 120.1 (2009): 221-259.
***

```python
# See Hoffman and Gelman (2013), section 3.2.1.
# https://colindcarroll.com/2019/04/21/step-size-adaptation-in-hamiltonian-monte-carlo/
# https://github.com/ColCarroll/minimc/blob/master/minimc/minimc.py

class DualAveragingStepSize:
    def __init__(self, initial_step_size, target_accept=0.65, gamma=0.05, t0=10.0, kappa=0.75):
        self.mu = np.log(10 * initial_step_size)  # proposals are biased upwards to stay away from 0.
        self.target_accept = target_accept
        self.gamma = gamma
        self.t = t0
        self.kappa = kappa
        self.error_sum = 0
        self.log_averaged_step = 0

    def update(self, p_accept):
        # Running tally of absolute error. Can be positive or negative, but optimally 0.
        self.error_sum += self.target_accept - p_accept
        # This is the next proposed (log) step size. Note it is biased towards mu.
        log_step = self.mu - self.error_sum / (np.sqrt(self.t) * self.gamma)
        # Forgetting rate. As 't' gets bigger, 'eta' gets smaller.
        eta = self.t ** -self.kappa
        # Smoothed average step size
        self.log_averaged_step = eta * log_step + (1 - eta) * self.log_averaged_step
        # State update, such taht 't' keeps updating
        self.t += 1
        # Return both the noisy step size (tuning phase), and the smoothed step size (after tuning)
        return np.exp(log_step), np.exp(self.log_averaged_step)

def hamiltonian_monte_carlo_dual(
    n_samples, negative_log_prob, initial_position, path_len=1, initial_step_size=0.5, tune=500):
    # counter of acceptance
    counter = 0
    # jax autodiff magic - initialize
    dVdq = jax.grad(negative_log_prob)
    # collect all our samples in a list
    samples = [initial_position]
    # Keep a single object for momentum resampling
    momentum = ss.norm(0, 1)
    # Step size adaption - initialization
    step_size = initial_step_size
    step_size_tuning = DualAveragingStepSize(step_size)
    # If initial_position is a 10d vector and n_samples is 100, we want
    # 100 x 10 momentum draws. We can do this in one call to momentum.rvs, and
    # iterate over rows
    size = (n_samples + tune,) + initial_position.shape[:1]
    for idx, p0 in tqdm(enumerate(momentum.rvs(size=size)), total=size[0]):
    # for idx, p0 in enumerate(momentum.rvs(size=size)):
        # Integrate over our path to get a new position and momentum
        # NB!!! Jitter the path length to stabilize algorithm - on average it equals the specified path length
        q_new, p_new = leapfrog(
            samples[-1], p0, dVdq, path_len=2*np.random.rand()*path_len, step_size=step_size)
        # Check Metropolis acceptance criterion
        start_log_p = np.sum(momentum.logpdf(p0)) - negative_log_prob(samples[-1])
        new_log_p = np.sum(momentum.logpdf(p_new)) - negative_log_prob(q_new)
        p_accept = min(1, np.exp(new_log_p - start_log_p))
        if np.random.rand() < p_accept:
            samples.append(q_new)
            if idx > tune - 1:
                counter += 1
        else:
            samples.append(samples[-1])
        # Step size adaption
        if idx < tune - 1:
            step_size, _ = step_size_tuning.update(p_accept)
        elif idx == tune - 1:
            _, step_size = step_size_tuning.update(p_accept)

    return np.array(samples[1 + tune :]), counter, step_size
```

# Numerical example

In this tutorial, we will study how HMC can be used to perform Bayesian inference in linear regression. For this simple case, we can ofcause derive the analytical gradients needed by the Leapfrog integrator, see e.g., my tutorial on linear regression ([link](https://nbviewer.org/github/SebastianGlavind/PhD-study/blob/master/Linear-regression/LinReg.ipynb)), but to showcase the generality of the implementation above, we will used the autodiff library `jax` for this case as well.

## The model

In the simplest case of linear regression, sometimes called ordinary linear regression, the scalar output $y$ is assumed to be a linear combination of the inputs $\mathbf{x}$, and the observation errors follow a Gaussian white noise distribution, thus

$$ y | \mathbf{w}, \sigma, \mathbf{x} \sim \mathcal{N}(w_0 + \sum_{m=1}^{M-1} w_m x_m, \sigma^2) = \mathcal{N}(\mathbf{w}^T \mathbf{x}, \sigma^2), $$

where we have augmented the input vector $\mathbf{x}$ with an additional first element, which is always 1, i.e, $\mathbf{x} = (1, x_1, x_2, ..., x_{M-1})$, and the corresponding weight vector $\mathbf{w} = (w_0, w_1, w_2, ..., w_{M-1})$.

If we now consider a training data set $\mathcal{D}=\{ \mathbf{x}[n], y[n] \}_{n=1}^N = \{ \mathbf{X}, \mathbf{y} \}$, where $\mathbf{X}$ is a $ N \times M $ design matrix and $\mathbf{y}$ is a column vector of the corresponding output observations, the joint likelihood of the training data may be written:

$$ \mathbf{y} | \mathbf{w}, \sigma, \mathbf{X} \sim \prod_{n=1}^N \mathcal{N}(\mathbf{w}^T \mathbf{x}[n], \sigma^2) = \mathcal{N}(\mathbf{X}\mathbf{w},\sigma^2\mathbf{I}), $$

where $\mathbf{I}$ is an $ N \times N $ identity matrix.

## Sample data

```python
# True, underlaying model
w0_true = 4
w1_true = 8
sigma_true = 2

def generate_training_data(x, w0, w1, sigma):
    n = len(x)
    error = np.random.normal(loc=0, scale=sigma, size=n)
    return( w0 + w1*x + error )

# Generate and plot data
x_tr = np.arange(start=-1, stop=1, step=0.1)
np.random.seed(10) # good seed
y_tr = generate_training_data(x=x_tr, w0 = w0_true, w1 = w1_true, sigma = sigma_true)

XX = np.vstack((np.ones(len(x_tr)), x_tr)).T # design matrix (convention Bishop(2006))
nX_tr, mX_tr = XX.shape # - np.array([0, 1]); # data dimensions are defined as actual features ;)

x_te = np.arange(start=-1.25, stop=1.25, step=0.1)
XX_te = np.vstack((np.ones(len(x_te)), x_te)).T # design matrix (convention Bishop(2006))
nX_te, mX_te = XX_te.shape # - np.array([0, 1]); # data dimensions are defined as actual features ;)

plt.plot(x_tr, y_tr,'ob');
plt.plot(x_tr,(w0_true + w1_true * x_tr),'-r');
plt.xlabel('x')
plt.ylabel('y')
plt.title('Training data and true linear model')
plt.grid()
```

## Bayesian inference

See e.g. Gelman et al. (2013; Sec.12.4 and App.C.4 (R implementation)) for a reference on the Hamiltonian Monte Carlo algorithm.

### The generative story

We sample a realization of the parameter vector $\theta$ as

$$ p(\boldsymbol\theta |\mathcal{D}) \propto p(\mathbf{y}|\mathbf{X},\boldsymbol\theta) p(\boldsymbol\theta), $$

where $$ \theta_{0:1} = w_{0:1} \sim \mathcal{N}(0,10), $$

and $$ \theta_{2} = \sigma^2 \sim \text{Gamma}(2,2). $$

### The unnormalized parameter posterior

Recall that the parameter posterior is defined through Bayes' rule as

$$p(\boldsymbol\theta|\mathcal{D}) = \frac{ p(\mathbf{y},\boldsymbol\theta|\mathbf{x}) }{ p(\mathbf{y})}, $$

but we will work with the unnormalized version, i.e.

$$p(\boldsymbol\theta|\mathcal{D}) \propto p(\mathbf{y},\boldsymbol\theta|\mathbf{x}), $$

to avoid calculating the always troubling normalizing constant $\mathbf{y}$.

***
Gelman, Andrew, et al. Bayesian data analysis. CRC press, 2013.
***

```python
# This function calculates the unnormalized posterior for theta, i.e. p(theta|D) = p(y,theta|x)/p(y)
def neg_log_pos_hmc(theta,pam1=np.array([0.,10.]),pam2=np.array([0.,10.]),pam3=np.array([2.,2.])):
    w0 = theta[0]
    w1 = theta[1]
    sig2 = theta[2]
    mu = w0 + w1 * x_tr
    logPos = (sum( jss.norm.logpdf(y_tr, loc = mu, scale = jnp.sqrt(sig2)) )
              + jss.norm.logpdf(w0, loc=pam1[0], scale=pam1[1])
              + jss.norm.logpdf(w1, loc=pam2[0], scale=pam2[1])
              + jss.gamma.logpdf(sig2, a = pam3[0], scale = pam3[1])
             )
    negLogPos = - logPos
    return(negLogPos)

# Test the implementation
print( neg_log_pos_hmc(np.array([4., 8., 4.])) )
print( jax.grad(neg_log_pos_hmc)(np.array([4., 8., 4.])) )
```

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

    48.206924
    [-1.1471946   0.42492092  0.96945906]

## Basic Hamiltonian Monte Carlo

This shows how to sample from éne chain; in a real application, we would initialize multiple chains overdispersed in the plausible range of the variables and check for mixing.

### Sampling from the posterior

```python
# %%time
# Call the HMC implementation
np.random.seed(42)
n_samp_hmc = int(2e2)
samp_hmc, counter_hmc = hamiltonian_monte_carlo(n_samp_hmc, neg_log_pos_hmc, np.array([4., 8., 2.]), path_len=2, step_size=0.65)
counter_hmc/n_samp_hmc
```

    0.645

### Analyze samples

We regard the first half of the samples as burn-in and plot the posterior samples (blue).

```python
samp_hmc0 = samp_hmc[0:round(n_samp_hmc/2),:] # burn-in
samp_hmc1 = samp_hmc[round(n_samp_hmc/2):n_samp_hmc,:] # without burn-in

# Marginal distribution approximations
plt.figure()
plt.hist(samp_hmc0[:,0], color='red' );
plt.hist(samp_hmc1[:,0] , color='blue' );
plt.title('Noise standard deviation')

plt.figure()
plt.hist(samp_hmc0[:,1], color='red' );
plt.hist(samp_hmc1[:,1] , color='blue' );
plt.title('Kernel standard deviation (signal)')

plt.figure()
plt.hist(samp_hmc0[:,2], color='red' );
plt.hist(samp_hmc1[:,2] , color='blue' );
plt.title('Kernel length scale (signal)');
```

```python
# Scatter plot of regression parameter
plt.plot(samp_hmc0[:,0],samp_hmc0[:,1],'*r');
plt.plot(samp_hmc1[:,0],samp_hmc1[:,1],'*b');
plt.xlabel('$w_0$');
plt.ylabel('$w_1$');
plt.title('Weights samples');
```

#### Define sample statistics for the latent function (f) and the output (y)

```python
f_pred_hmc = samp_hmc1[:,0:2].dot(XX_te.T)
y_pred_hmc = f_pred_hmc + np.random.normal(scale=np.sqrt(samp_hmc1[:,2])).reshape([-1,1])
mu_pred_hmc = np.mean(f_pred_hmc,axis=0)
cf_pred_f_hmc= np.quantile(f_pred_hmc,q=[0.025, 0.975],axis=0)
cf_pred_y_hmc= np.quantile(y_pred_hmc,q=[0.025, 0.975],axis=0)
```

#### Plot the resulting model

95\% Bayesian credible interval of the posterior model.

```python
plt.plot(x_tr, y_tr,'ob');
plt.plot(x_tr,(w0_true + w1_true * x_tr),'-r');
plt.plot(x_te,mu_pred_hmc,'--k');
plt.fill_between(x=x_te, y1=cf_pred_f_hmc[0,:], y2=cf_pred_f_hmc[1,:], alpha=0.5);
plt.fill_between(x=x_te, y1=cf_pred_y_hmc[0,:], y2=cf_pred_y_hmc[1,:], alpha=0.1, color='blue');
plt.xlabel('x');
plt.ylabel('y');
plt.title('Training data, true linear model and linear regression');
plt.grid();
```

## Hamiltonian Monte Carlo with step size adaption

This shows how to sample from éne chain; in a real application, we would initialize multiple chains overdispersed in the plausible range of the variables and check for mixing.

### Sampling from the posterior

```python
# %%time
# Call the HMC implementation
np.random.seed(42)
n_samp_hmc_dual = int(2e2)
samp_hmc_dual, counter_hmc_dual, step_size_hmc_dual = hamiltonian_monte_carlo_dual(n_samp_hmc_dual, neg_log_pos_hmc, np.array([2., 16., 4.]),
                                                                                path_len=3, initial_step_size=0.33, tune=100)
print('Acceptance rate: ', counter_hmc_dual/n_samp_hmc_dual)
print('Optimized step size: ', step_size_hmc_dual)
```

    100%|█████████████████████████████████████████| 300/300 [01:53<00:00,  2.65it/s]

    Acceptance rate:  0.85
    Optimized step size:  5.009077704356696e-01

### Analyze samples

We regard the first half of the samples as burn-in and plot the posterior samples (blue).

```python
samp_hmc_dual0 = samp_hmc_dual[0:round(n_samp_hmc_dual/2),:] # burn-in
samp_hmc_dual1 = samp_hmc_dual[round(n_samp_hmc_dual/2):n_samp_hmc_dual,:] # without burn-in

# Marginal distribution approximations
plt.figure()
plt.hist(samp_hmc_dual0[:,0], color='red' );
plt.hist(samp_hmc_dual1[:,0] , color='blue' );
plt.title('Noise standard deviation')

plt.figure()
plt.hist(samp_hmc_dual0[:,1], color='red' );
plt.hist(samp_hmc_dual1[:,1] , color='blue' );
plt.title('Kernel standard deviation (signal)')

plt.figure()
plt.hist(samp_hmc_dual0[:,2], color='red' );
plt.hist(samp_hmc_dual1[:,2] , color='blue' );
plt.title('Kernel length scale (signal)');
```

```python
# Scatter plot of regression parameter
plt.plot(samp_hmc_dual0[:,0],samp_hmc_dual0[:,1],'*r');
plt.plot(samp_hmc_dual1[:,0],samp_hmc_dual1[:,1],'*b');
plt.xlabel('$w_0$');
plt.ylabel('$w_1$');
plt.title('Weights samples');
```

#### Define sample statistics for the latent function (f) and the output (y)

```python
f_pred_hmc_dual = samp_hmc_dual1[:,0:2].dot(XX_te.T)
y_pred_hmc_dual = f_pred_hmc_dual + np.random.normal(scale=np.sqrt(samp_hmc_dual1[:,2])).reshape([-1,1])
mu_pred_hmc_dual = np.mean(f_pred_hmc_dual,axis=0)
cf_pred_f_hmc_dual = np.quantile(f_pred_hmc_dual,q=[0.025, 0.975],axis=0)
cf_pred_y_hmc_dual = np.quantile(y_pred_hmc_dual,q=[0.025, 0.975],axis=0)
```

#### Plot the resulting model

95\% Bayesian credible interval of the posterior model.

```python
plt.plot(x_tr, y_tr,'ob');
plt.plot(x_tr,(w0_true + w1_true * x_tr),'-r');
plt.plot(x_te,mu_pred_hmc_dual,'--k');
plt.fill_between(x=x_te, y1=cf_pred_f_hmc_dual[0,:], y2=cf_pred_f_hmc_dual[1,:], alpha=0.5);
plt.fill_between(x=x_te, y1=cf_pred_y_hmc_dual[0,:], y2=cf_pred_y_hmc_dual[1,:], alpha=0.1, color='blue');
plt.xlabel('x');
plt.ylabel('y');
plt.title('Training data, true linear model and linear regression');
plt.grid();
```


```python
from sympy import *
from sympy.plotting import plot
# init_printing(pretty_print=True,use_latex=True)
# %matplotlib inline
# import matplotlib.pyplot as plt
from numpy import arange
```

```python
f,t,T,G = symbols('f t T G')
```

```python
f = G*t-G/2
print ('f =', f)
```

    f = G*t - G/2

```python
p = plot((f.subs(G,25),(t,0,1)),((f-25).subs(G,25),(t,1,2)),((f-50).subs(G,25),(t,2,3)), \
         ylabel = '$f = 25 t$ [V]', xlabel = '$t$ [sec]', xlim=(0,3), figsize=(10,8))
```

```python
c0 = integrate(f,(t,0,T))*(1/T)
print(c0)
print (c0.subs([(G,25),(T,1)]))
```

    (G*T**2/2 - G*T/2)/T
    0

```python
a1 = (2/T)*integrate(f*sin(2*pi*1*t),(t,0,T))
```

```python
print(a1)
print( a1.subs([(T,1),(pi, 3.14),(G,25)]))
```

    2*(-G*T*cos(2*pi*T)/(2*pi) + G*sin(2*pi*T)/(4*pi**2) + G*cos(2*pi*T)/(4*pi) - G/(4*pi))/T
    -7.96178343949045

```python
a,b = [],[]
for n in range(10):
    a.append((2/T)*integrate(f*sin(n*pi*t),(t,0,T)))
    b.append((2/T)*integrate(f*cos(n*pi*t),(t,0,T)))
```

```python
a[9], b[9]
```

    (2*G*(-T*cos(9*pi*T)/(9*pi) + sin(9*pi*T)/(81*pi**2))/T,
     2*(G*(T*sin(9*pi*T)/(9*pi) + cos(9*pi*T)/(81*pi**2)) - G/(81*pi**2))/T)

```python

c = []
for n in range(1,10):
    print( a[n])
    expr = sqrt(a[n]**2 + b[n]**2)
    expr = expr.subs([(T,1),(pi, 3.14),(G,25)])
    # print expr
    # val = lambdify(t,expr,'numpy')
    c.append(expr)
```

    2*G*(-T*cos(pi*T)/pi + sin(pi*T)/pi**2)/T
    2*G*(-T*cos(2*pi*T)/(2*pi) + sin(2*pi*T)/(4*pi**2))/T
    2*G*(-T*cos(3*pi*T)/(3*pi) + sin(3*pi*T)/(9*pi**2))/T
    2*G*(-T*cos(4*pi*T)/(4*pi) + sin(4*pi*T)/(16*pi**2))/T
    2*G*(-T*cos(5*pi*T)/(5*pi) + sin(5*pi*T)/(25*pi**2))/T
    2*G*(-T*cos(6*pi*T)/(6*pi) + sin(6*pi*T)/(36*pi**2))/T
    2*G*(-T*cos(7*pi*T)/(7*pi) + sin(7*pi*T)/(49*pi**2))/T
    2*G*(-T*cos(8*pi*T)/(8*pi) + sin(8*pi*T)/(64*pi**2))/T
    2*G*(-T*cos(9*pi*T)/(9*pi) + sin(9*pi*T)/(81*pi**2))/T

```python
import matplotlib.pyplot as plt
fig,ax = plt.subplots(figsize=(10,8))
ax.bar(-.5+arange(1,10),c,alpha=.5)
ax.set_xticks([1,2,3,4,5,6,7,8,9,10])
ax.set_xlabel('$n$',fontsize=16)
ax.set_ylabel('Fourier coefficients $a_n^2 + b_n^2$',fontsize=16)
```

```python

```


# Well testing

Consider a simple groundwater system that comprises a **confining layer** and a **confined aquifer** beneath it. The **piezometric surface**1 sits just below the ground surface.

A **pumping well** has been drilled that is screened across the entire depth of the aquifer. An **observation well** has also been drilled with the water level in this well reflecting the piezometric surface at that location.

> *Run the* ***notebook cell*** *below that begins* **`from`**` well_test `**`import`** *by clicking inside it and hitting* **`Ctrl+Enter`.**

1 The piezometric surface for a confined aquifer is **not necessarily** the same as the water table. It reflects the relative pressure in the confined aquifer and could be higher or lower than the water table depending land surface elevation or past abstraction. In an **unconfined aquifer** the piezometric surface **does** coincide with the water table.

```python
from well_test import*
conceptual_model()
```

    VBox(children=(HBox(children=(FloatSlider(value=200.0, continuous_update=False, description='$r$ [m]', max=500…

***Complete the Getting Started activities in `Conceptual_Exercises.pdf`***

## 1. Pumping a confined aquifer

The purpose of well testing is to **measure** specific properties of the nearby rock that contains groundwater, principally **storage**, $S$, and **transmissivity**, $T$. These help us to later **predict** how an aquifer will respond to future abstraction.

For a **confined aquifer**, the Theis equation can be used to estimate $S$ and $T$ from observations of **drawdown**, $h$.

\begin{equation}
h = \frac{1}{4\pi} \frac{Q}{T} W\left(\frac{r^2S}{4Tt}\right)\quad\quad\approx 0.183\frac{Q}{T}\log\left(\frac{2.25 Tt}{r^2S}\right)\quad\text{when}\quad\frac{r^2S}{4Tt}<0.05
\end{equation}

where $W(\cdot)$ is the well function, $r$ is distance to the observation well, $t$ is time, $S$ is storage, $Q$ is pumping rate, and $T$ is transmissivity.

> *Run the Notebook Cell below, which uses the* **log** *approximation above to model drawdown.*

```python
confined_aquifer(analysis=False)
```

    VBox(children=(HBox(children=(FloatSlider(value=1000.0, continuous_update=False, description='$Q$ [m$^3$/day]'…

***Complete Problem Set 1 in `Conceptual_Exercises.pdf`***

## 2. Pumping a semi-confined aquifer

Some confining layers allow a little bit of water to pass through. These leaky aquifers have a different response to drawdown, because now recharge can be drawn from above, in addition to radially.

Leaky aquifers are characterized by flattening of the drawdown-time curve. The Hantush-Jacob analytical solution describes this curve.

\begin{equation}
h = \frac{1}{4\pi} \frac{Q}{T} W_h\left(u,\frac{r}{L}\right), \quad\quad L = \sqrt{Tc}
\end{equation}

where $W_h(\cdot)$ is the Hantush well function and $c$ is the hydraulic resistance of the aquitard.

> *Run the Notebook Cell below, which implements the Hantush-Jacob solution.*

```python
leaky_aquifer()
```

    VBox(children=(HBox(children=(FloatSlider(value=1000.0, continuous_update=False, description='$Q$ [m$^3$/day]'…

***Complete Problem Set 2 in `Conceptual_Exercises.pdf`***

## 3. Pumping a laterally constrained aquifer

### Flow boundaries

Aquifers can be horizontally truncated in the presence of low permeability faults or where fault movement has laterally juxtaposed a confining unit against the aquifer. In either case, flow barriers affect the groundwater near the well and therefore the response to pumping.

At a flow barrier, the hydraulic gradient must be zero (if it wasn't, there would be flow, which there can't...) This can be enforced in a model by adding an image well at the same distance on the otherside of the flow barrier - any attempt to steepen the hydraulic gradient towards the real well will then be offset by the effect of the image well.

For a well at $\mathbf{x}_w=[R,0]$ and a boundary on the $y$-axis, the drawdown is

\begin{equation}
h = \frac{1}{4\pi} \frac{Q}{T} W\left(\frac{|\mathbf{x}-\mathbf{x}_w|^2S}{4Tt}\right) + \frac{1}{4\pi} \frac{Q}{T} W\left(\frac{|\mathbf{x}+\mathbf{x}_w|^2S}{4Tt}\right)
\end{equation}

where the second term is the Theis solution due to extraction from an image well located at -$\mathbf{x}_w$.

Superposition in this manner requires that the constitutive solutions be derived from linear governing equations. This condition is satisifed for both Theis and Hantush-Jacob solutions.

> *Run the Notebook Cell below, which demonstrates the image well approximation to a flow barrier.*

```python
flow_barrier()
```

    VBox(children=(HBox(children=(FloatSlider(value=1000.0, continuous_update=False, description='$Q$ [m$^3$/day]'…

### Recharge source

A source of recharge can be thought of as a boundary of zero drawdown. This drawdown can be modelled by adding an image well on the other side of the boundary that is injecting at -$Q$. For a well at $\mathbf{x}_w=[R,0]$ and a line of recharge (e.g., the coastline) on the $y$-axis, the drawdown is

\begin{equation}
h = \frac{1}{4\pi} \frac{Q}{T} W\left(\frac{|\mathbf{x}-\mathbf{x}_w|^2S}{4Tt}\right) - \frac{1}{4\pi} \frac{Q}{T} W\left(\frac{|\mathbf{x}+\mathbf{x}_w|^2S}{4Tt}\right)
\end{equation}

where the second term is the Theis solution due to injection (hence negative) into an image well located at -$\mathbf{x}_w$.

For any location on the $y$-axis, $\mathbf{x}=[0,y]$, then the radii in both terms becomes $|\mathbf{x}-\mathbf{x}_w|^2 = R^2+y^2=|\mathbf{x}+\mathbf{x}_w|^2$ and therefore the two terms cancel so that $h=0$.

> *Run the Notebook Cell below, which demonstrates the image well approximation to a source of recharge.*

```python
recharge_source()
```

    VBox(children=(HBox(children=(FloatSlider(value=1000.0, continuous_update=False, description='$Q$ [m$^3$/day]'…

***Complete Problem Set 3 in `Conceptual_Exercises.pdf`***

## 4. Pumping an unconfined aquifer

In an unconfined aquifer, we must contend with both changes of the hydraulic gradient AND a reduction of the saturated thickness. Furthermore, there are two storage responses: early on, it is the elastic storage that yields water to the well, with water table drawdown occurring later. The solution assumes that vertical flow in the aquifer can be neglected and that drawdown is small compared to saturated thickness.

Neuman (1975) formulated a semi-analytical solution to this problem, which was later updated by a more computationally efficient solution from Moench (1996).

> *Run the Notebook Cell below, which demonstrates the Moench-Neuman flow solution in an unconfined aquifer.*

```python
unconfined_aquifer()
```

    VBox(children=(HBox(children=(FloatSlider(value=1000.0, continuous_update=False, description='$Q$ [m$^3$/day]'…

***Complete Problem Set 4 in `Conceptual_Exercises.pdf`***

## All pumping tests

```python
all_options()
```

    VBox(children=(HBox(children=(Dropdown(description='Aquifer type', options={'confined': 1, 'leaky aquifer': 2,…