text:
Metropolis and Gibbs Sampling
====

```python
%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
from functools import partial
```

```python
sns.set_context('notebook', font_scale=1.5)
```

## Introduction to MCMC

In regular Markov chain models, we are usually interested in finding the equilibrium  distribution $\pi$ at whihc $\pi^T T = \pi^T$ for a given transition kernel $T$.

MCMC inverts this thinking - we fix the equilibrium distribution to be the posterior distribution

$$
p(\theta \mid X) = \frac{p(X \mid \theta) \, p(\theta)}{\int{p(X \mid \theta) \, p(\theta) d\theta}}
$$

and look for a transition kernel $T$ that will converge to this equilibrium distribution.

### Island hopping

We first provide an example to show the mechanics of the Metropolis algorithm concretely, then explore why it works.

[Kruschke's book](https://sites.google.com/site/doingbayesiandataanalysis/) begins with a fun example of a politician visiting a chain of islands to canvas support - being callow, the politician uses a simple rule to determine which island to visit next. Each day, the politician chooses a neighboring island and compares the populations there with the population of the current island. If the neighboring island has a larger population, the politician goes over. If the neighboring island has a smaller population, then the politician visits with probability $p = N_\text{neighbor} / N_\text{current}$ where $N$ is the island population; otherwise the politician stays on the same island. After doing this for many days, the politician will end up spending time on each island proportional to the population of each island - in other words, estimating the distribution of island populations correctly. How a simple comparison of only two states at a time can lead to accurate estimation of a probability density is the topic of the next few lectures.

```python
def make_islands(n, low=10, high=101):
    islands = np.random.randint(low, high, n+2)
    islands[0] = 0
    islands[-1] = 0
    return islands
```

```python
def hop(islands, start=1, niter=1000):
    pos = start
    pop = islands[pos]
    thetas = np.zeros(niter+1, dtype='int')
    thetas[0] = pos
    for i in range(niter):
        # generate sample from proposal distribution
        k = np.random.choice([-1, 1], 1)
        next_pos = pos + k
        # evaluate unnormalized target distribution at proposed position
        next_pop = islands[next_pos]
        # calculate acceptance probability
        p = min(1, next_pop/pop)
        # use uniform random to decide accept/reject proposal
        if np.random.random() < p:
            pos = next_pos
            pop = next_pop
        thetas[i+1] = pos
    return thetas
```

```python
islands = make_islands(10)
thetas = hop(islands, start=1, niter=10000)
```

### True population proportions

```python
data = islands[1:-1]
data = data/data.sum()
sns.barplot(x=np.arange(len(data)), y=data)
pass
```

### Estimated population proportions

```python
data = np.bincount(thetas)[1:]
data = data/data.sum()
sns.barplot(x=np.arange(len(data)), y=data)
pass
```

### Generic Metropolis scheme

```python
def metroplis(start, target, proposal, niter, nburn=0):
    current = start
    post = [current]
    for i in range(niter):
        proposed = proposal(current)
        p = min(target(proposed)/target(current), 1)
        if np.random.random() < p:
            current = proposed
        post.append(current)
    return post[nburn:]
```

### Apply to island hooper

```python
target = lambda x: islands[x]
proposal = lambda x: x + np.random.choice([-1, 1])
post = metroplis(1, target, proposal, 2000)
data = np.bincount(post)[1:]
data = data/data.sum()
sns.barplot(x=np.arange(len(data)), y=data)
pass
```

Bayesian Data Analysis
----

The fundamental objective of Bayesian data analysis is to determine the posterior distribution

$$
p(\theta \ | \ X) = \frac{p(X \ | \ \theta) p(\theta)}{p(X)}
$$

where the denominator is

$$
p(X) = \int d\theta^* p(X \ | \ \theta^*) p(\theta^*)
$$

Here,

- $p(X \ | \ \theta)$ is the likelihood,
- $p(\theta)$ is the prior and
- $p(X)$ is a normalizing constant also known as the evidence or marginal likelihood

The computational issue is the difficulty of evaluating the integral in the denominator. There are many ways to address this difficulty, including:

- In cases with conjugate priors (with conjugate priors, the posterior has the same distribution family as the prior), we can get closed form solutions
- We can use numerical integration
- We can approximate the functions used to calculate the posterior with simpler functions and show that the resulting  approximate posterior is "close" to true posterior (variational Bayes)
- We can use Monte Carlo methods, of which the most important is Markov Chain Monte Carlo (MCMC). In simple Monte Carlo integration, we want to estimate the integral $\int f(x) \, p(x) dx$. With Bayesian models, the distribution $p(x)$ in the integral is the posterior

$$
p(x) = p(\theta \ | \ X) = \frac{p(X \ | \ \theta) p(\theta)}{\int d\theta^* p(X \ | \ \theta^*) p(\theta^*) }
$$
- MCMC allows to sample from the posterior distribution - the samples will not be independent unlike simple Monte Carlo integration, but this is OK as we can compensate for the auto-correlation by drawing a larger number  of samples.

### Motivating example

We will use the toy example of estimating the bias of a coin given a sample consisting of $n$ tosses to illustrate a few of the approaches.

#### Analytical solution

If we use a beta distribution as the prior, then the posterior distribution has a closed form solution. This is shown in the example below. Some general points:

- We need to choose a prior distribution family (i.e. the beta here) as well as its parameters (here a=10, b=10)
    - The prior distribution may be relatively uninformative (i.e. more flat) or informative (i.e. more peaked)
- The posterior depends on both the prior and the data
    - As the amount of data becomes large, the posterior approximates the MLE
    - An informative prior takes more data to shift than an uninformative one
- Of course, it is also important the model used (i.e. the likelihood) is appropriate for the fitting the data
- The mode of the posterior distribution is known as the maximum a posteriori (MAP) estimate (cf MLE which is the mode of the likelihood)

```python
import scipy.stats as stats
```

```python
n = 100
h = 61
p = h/n
rv = stats.binom(n, p)
mu = rv.mean()

a, b = 10, 10
prior = stats.beta(a, b)
post = stats.beta(h+a, n-h+b)
ci = post.interval(0.95)

thetas = np.linspace(0, 1, 200)
plt.plot(thetas, prior.pdf(thetas), label='Prior', c='blue')
plt.plot(thetas, post.pdf(thetas), label='Posterior', c='red')
plt.plot(thetas, n*stats.binom(n, thetas).pmf(h), label='Likelihood', c='green')
plt.axvline((h+a-1)/(n+a+b-2), c='red', linestyle='dashed', alpha=0.4, label='MAP')
plt.axvline(mu/n, c='green', linestyle='dashed', alpha=0.4, label='MLE')
plt.xlim([0, 1])
plt.axhline(0.3, ci[0], ci[1], c='black', linewidth=2, label='95% CI');
plt.xlabel(r'$\theta$', fontsize=14)
plt.ylabel('Density', fontsize=16)
plt.legend(loc='upper left')
pass
```

#### Numerical integration

One simple way of numerical integration is to estimate the values on a grid of values for $\theta$. To calculate the posterior, we find the prior and the likelihood for each value of $\theta$, and for the marginal likelihood, we replace the integral with the equivalent sum

$$
p(X) = \sum_{\theta^*} p(X | \theta^*) p(\theta^*)
$$

One advantage of this is that the prior does not have to be conjugate (although the example below uses the same beta prior for ease of comparison), and so we are not restricted in our choice of an appropriate prior distribution. For example, the prior can be a mixture distribution or estimated empirically from data. The disadvantage, of course, is that this is computationally very expensive when we need to estimate multiple parameters, since the number of grid points grows as $\mathcal{O}(n^d)$, where $n$ defines the grid resolution and $d$ is the size of $\theta$.

```python
thetas = np.linspace(0, 1, 200)
prior = stats.beta(a, b)

post = prior.pdf(thetas) * stats.binom(n, thetas).pmf(h)
# Normalzie so volume is 1
post /= (post.sum() / len(thetas))

plt.plot(thetas, prior.pdf(thetas), label='Prior', c='blue')
plt.plot(thetas, n*stats.binom(n, thetas).pmf(h), label='Likelihood', c='green')
plt.plot(thetas, post, label='Posterior', c='red')
plt.xlim([0, 1])
plt.xlabel(r'$\theta$', fontsize=14)
plt.ylabel('Density', fontsize=16)
plt.legend()
pass
```

### Markov Chain Monte Carlo (MCMC)

This lecture will only cover the basic ideas of MCMC and the 3 common variants - Metroplis, Metropolis-Hastings and Gibbs sampling. All code will be built from the ground up to illustrate what is involved in fitting an MCMC model, but only toy examples will be shown since the goal is conceptual understanding. More realistic computational examples will be shown in coming lectures using the `pymc3` and `pystan` packages.

In Bayesian statistics, we want to estimate the posterior distribution, but this is often intractable due to the high-dimensional integral in the denominator (marginal likelihood). A few other ideas we have encountered that are also relevant here are Monte Carlo integration with independent samples and the use of proposal distributions (e.g. rejection and importance sampling). As we have seen from the Monte Carlo integration lectures, we can approximate the posterior $p(\theta | X)$ if we can somehow draw many samples that come from the posterior distribution. With vanilla Monte Carlo integration, we need the samples to be independent draws from the posterior distribution, which is a problem if we do not actually know what the posterior distribution is (because we cannot integrate the marginal likelihood).

With MCMC, we draw samples from a (simple) proposal distribution so that each draw depends only on the state of the previous draw (i.e. the samples form a Markov chain). Under certain conditions, the Markov chain will have a unique stationary distribution. In addition, not all proposal draws are used - instead we set up acceptance criteria for each draw based on comparing successive states with respect to a target distribution that ensure that the stationary distribution is the posterior distribution of interest. The nice thing is that this target distribution only needs to be proportional to the posterior distribution, which means we don't need to evaluate the potentially intractable marginal likelihood, which is just a normalizing constant. We can find such a target distribution easily, since `posterior` $\propto$ `likelihood` $\times$ `prior`. After some time, the Markov chain of accepted draws will converge to the stationary distribution, and we can use those samples as (correlated) draws from the posterior distribution, and find functions of the posterior distribution in the same way as for vanilla Monte Carlo integration.

There are several flavors of MCMC, but the simplest to understand is the Metropolis-Hastings random walk algorithm, and we will start there.

#### Metropolis-Hastings random walk algorithm for estimating the bias of a coin

To carry out the Metropolis-Hastings algorithm, we need to draw random samples from the following distributions

- the standard uniform distribution
- a proposal distribution $p(x)$ that we choose to be $\mathcal{N}(0, \sigma)$
- the target function $g(x)$ which is proportional to the posterior probability (the target function is essentially an unnormalized distribution)

Given an initial guess for $\theta$ with positive probability of being drawn, the Metropolis-Hastings algorithm proceeds as follows

- Choose a new proposed value ($\theta_p$) such that  $\theta_p = \theta + \Delta\theta$ where $\Delta \theta \sim \mathcal{N}(0, \sigma)$
- Caluculate the ratio

$$
\rho = \frac{g(\theta_p \ | \ X)}{g(\theta \ | \ X)}
$$

where $g$ is the posterior probability.

- If the proposal distribution is not symmetrical, we need to weight the acceptance probability to maintain detailed balance (reversibility) of the stationary distribution, and instead calculate

$$
\rho = \frac{g(\theta_p \ | \ X) p(\theta \ | \ \theta_p)}{g(\theta \ | \ X) p(\theta_p \ | \ \theta)}
$$

Note: The Metropolis algorithm refers to symmetrical proposals, and Metropolis-Hastings refers to this correction for asymmetrical proposals.

Since we are taking ratios, the denominator cancels any distribution proportional to $g$ will also work - so we can use

$$
\rho = \frac{p(X | \theta_p ) p(\theta_p)}{p(X | \theta ) p(\theta)}
$$

- If $\rho \ge 1$, then set $\theta = \theta_p$
- If $\rho \lt 1$, then set $\theta = \theta_p$ with probability $\rho$, otherwise set $\theta = \theta$ (this is where we use the standard uniform distribution)
- Repeat the earlier steps

After some number of iterations $k$, the samples $\theta_{k+1}, \theta_{k+2}, \dots$ will be samples from the posterior distributions. Here are initial concepts to help your intuition about why this is so:

- We accept a proposed move to $\theta_{k+1}$ whenever the density of the (unnormalized) target function at $\theta_{k+1}$ is larger than the value of $\theta_k$ - so $\theta$ will more often be found in places where the target function is denser
- If this was all we accepted, $\theta$ would get stuck at a local mode of the target function, so we also accept occasional moves to lower density regions - it turns out that the correct probability of doing so is given by the ratio $\rho$
- The acceptance criteria only looks at ratios of the target function, so the denominator cancels out and does not matter - that is why we only need the target function to be proportional to the posterior distribution
- So, $\theta$ will be expected to bounce around in such a way that its spends its time in places proportional to the density of the posterior distribution - that is, $\theta$ is a draw from the posterior distribution.

Additional notes:

Different proposal distributions can be used for Metropolis-Hastings:

- The independence sampler uses a proposal distribution that is independent of the current value of $\theta$. In this case the proposal distribution needs to be similar to the posterior distribution for efficiency, while ensuring that the acceptance ratio is bounded in the tail region of the posterior.
- The random walk sampler (used in this example) takes a random step centered at the current value of $\theta$ - efficiency is a trade-off between small step size with high probability of acceptance and large step sizes with low probability of acceptance. Note (picture will be sketched in class) that the random walk may take a long time to traverse narrow regions of the probability distribution. Changing the step size (e.g. scaling $\Sigma$ for a multivariate normal proposal distribution) so that a target proportion of proposals are accepted is known as *tuning*.
- Much research is being conducted on different proposal distributions for efficient sampling of the posterior distribution.

We will first see a numerical example and then try to understand why it works.

```python
def target(lik, prior, n, h, theta):
    if theta < 0 or theta > 1:
        return 0
    else:
        return lik(n, theta).pmf(h)*prior.pdf(theta)

n = 100
h = 61
a = 10
b = 10
lik = stats.binom
prior = stats.beta(a, b)
sigma = 0.3

naccept = 0
theta = 0.1
niters = 10000
samples = np.zeros(niters+1)
samples[0] = theta
for i in range(niters):
    theta_p = theta + stats.norm(0, sigma).rvs()
    rho = min(1, target(lik, prior, n, h, theta_p)/target(lik, prior, n, h, theta ))
    u = np.random.uniform()
    if u < rho:
        naccept += 1
        theta = theta_p
    samples[i+1] = theta
nmcmc = len(samples)//2
print("Efficiency = ", naccept/niters)
```

    Efficiency =  0.1861

```python
post = stats.beta(h+a, n-h+b)

plt.hist(samples[nmcmc:], 40, histtype='step', normed=True, linewidth=1, label='Prior');
plt.hist(prior.rvs(nmcmc), 40, histtype='step', normed=True, linewidth=1, label='Posterior');
plt.plot(thetas, post.pdf(thetas), c='red', linestyle='--', alpha=0.5, label='True posterior')
plt.xlim([0,1]);
plt.legend(loc='upper left')
pass
```

#### Assessing for convergence

Trace plots are often used to informally assess for stochastic convergence. Rigorous demonstration of convergence is an unsolved problem, but simple ideas such as running multiple chains and checking that they are converging to similar distributions are often employed in practice.

```python
def mh_coin(niters, n, h, theta, lik, prior, sigma):
    samples = [theta]
    while len(samples) < niters:
        theta_p = theta + stats.norm(0, sigma).rvs()
        rho = min(1, target(lik, prior, n, h, theta_p)/target(lik, prior, n, h, theta ))
        u = np.random.uniform()
        if u < rho:
            theta = theta_p
        samples.append(theta)
    return samples
```

```python
n = 100
h = 61
lik = stats.binom
prior = stats.beta(a, b)
sigma = 0.05
niters = 100

sampless = [mh_coin(niters, n, h, theta, lik, prior, sigma) for theta in np.arange(0.1, 1, 0.2)]
```

```python
# Convergence of multiple chains

for samples in sampless:
    plt.plot(samples, '-o')
plt.xlim([0, niters])
plt.ylim([0, 1]);
```

#### Why does Metropolis-Hastings work?

There are two main ideas - first that the samples generated by MCMC constitute a Markov chain, and that this Markov chain has a unique stationary distribution that is always reached if we generate a very large number of samples. The second idea is to show that this stationary distribution is exactly the posterior distribution that we are looking for. We will only give the intuition here as a refresher.

#### One: There is a unique stationary state

Since possible transitions depend only on the current and the proposed values of $\theta$, the successive values of $\theta$ in a Metropolis-Hastings sample constitute a Markov chain. Recall that for a Markov chain with a transition matrix $A$

$$
\pi^T = \pi^T A
$$

means that $\pi$ is a stationary distribution. If it is possible to go from any state to any other state, then the matrix is irreducible. If in addition, it is not possible to get stuck in an oscillation, then the matrix is also aperiodic or mixing. For finite state spaces, irreducibility and aperiodicity guarantee the existence of a unique stationary state. For continuous state space, we need an additional property of positive recurrence - starting from any state, the expected time to come back to the original state must be finite. If we have all 3 properties of irreducibility, aperiodicity and positive recurrence, then there is a unique stationary distribution. The term ergodic is a little confusing - most standard definitions take ergodicity to be equivalent to irreducibility, but often Bayesian texts take ergodicity to mean irreducibility, aperiodicity and positive recurrence, and we will follow the latter convention. For another intuitive perspective, the random walk Metropolis-Hasting algorithm is analogous to a diffusion process. Since all states are communicating (by design), eventually the system will settle into an equilibrium state. This is analogous to converging on the stationary state.

#### Two: The stationary state is the posterior probability distribution

We will consider the simplest possible scenario for an explicit calculation. Suppose we have a two-state system where the posterior probabilities are $\theta$ and $1 - \theta$. Suppose $\theta \lt 0.5$. So we have the following picture with the Metropolis-Hastings algorithm:

and we find the stationary distribution $\pi = \left( \begin{array}{cc} p & 1-p \end{array} \right)$ by solving

$$
\begin{align}
\left( \begin{array}{cc} p & 1-p \end{array} \right) &=
\left( \begin{array}{cc} p & 1-p \end{array} \right) \left(
\begin{array}{cc}
0 & 1  \\
\frac{\theta}{1-\theta} & 1-\frac{\theta}{1-\theta}
\end{array}
\right)
\end{align}
$$

to be $\pi = \left( \begin{array}{cc} \theta & 1-\theta \end{array} \right)$, which is the posterior distribution.

The final point is that we can find a stationary distribution using the detailed balance (reversibility) criterion that says that the probability of being in state $x$ and moving to state $y$ must be the same as the probability of being in state $y$ and moving to state $x$. Note that detailed balance is a sufficient but not necessary condition ofr $\pi$ to be a steady state distribution (assuming ergodicity). Or, more briefly,

$$
\pi(x)T(x \to y) = \pi(y)T(y \to x)
$$

and the need to make sure that this condition is true accounts for the strange looking acceptance criterion

$$
\min \left(1, \frac{g(\theta_p \ | \ X) p(\theta \ | \ \theta_p)}{g(\theta \ | \ X) p(\theta_p \ | \ \theta)} \right)
$$

### Intuition

We want the stationary distribution $\pi(x)$ to be the posterior distribution $P(x)$. So we set

$$
P(x)T(x \to y) = P(y)T(y \to x)
$$

Rearranging, we get

$$
\frac{T(x \to y)}{T(y \to x)} = \frac{P(y)}{P(x)}
$$

We split the transition probability into separate proposal $q$ and acceptance $A$ parts, and after a little algebraic rearrangement get

$$
\frac{A(x \to y)}{A(y \to x)} = \frac{P(y) \, q(y \to x)}{P(x) \, q(x \to y)}
$$

An acceptance probability that meets this condition is
$$
A(x \to y) = \min \left(1, \frac{P(y) \, q(y \to x)}{P(x) \, q(x \to y)} \right)
$$

since $A$ in the numerator and denominator are both bounded above by 1.

See [Chib and Greenberg](https://eml.berkeley.edu/reprints/misc/understanding.pdf) for algebraic details.

### The Gibbs sampler

Suppose we have a vector of parameters $\theta = (\theta_1, \theta_2, \dots, \theta_k)$, and we want to estimate the joint posterior distribution $p(\theta | X)$. Suppose we can find and draw random samples from all the conditional distributions

$$
p(\theta_1 | \theta_2, \dots \theta_k, X) \\
p(\theta_2 | \theta_1, \dots \theta_k, X) \\
\dots \\
p(\theta_k | \theta_1, \theta_2, \dots, X)
$$

With Gibbs sampling, the Markov chain is constructed by sampling from the conditional distribution for each parameter $\theta_i$ in turn, treating all other parameters as observed. When we have finished iterating over all parameters, we are said to have completed one cycle of the Gibbs sampler. Since hierarchical models are typically set up as products of conditional distributions, the Gibbs sampler is ubiquitous in Bayesian modeling. Where it is difficult to sample from a conditional distribution, we can sample using a Metropolis-Hastings algorithm instead - this is known as Metropolis within Gibbs.

Gibbs sampling is a type of random walk through parameter space, and hence can be thought of as a Metropolis-Hastings algorithm with a special proposal distribution. At each iteration in the cycle, we are drawing a proposal for a new value of a particular parameter, where the proposal distribution *is* the conditional posterior probability of that parameter. This means that the proposal move is *always* accepted. Hence, if we can draw samples from the conditional distributions, Gibbs sampling can be much more efficient than regular Metropolis-Hastings.
More formally, we want to show that

$$
\frac{P(y) \, q(y \to x)}{P(x) \, q(x \to y)} = 1
$$

We start by noting that $P(x_{-i})$ is the same as $P(y_{-i})$ since apart from the component $i$, the old state and the proposed new state are identical in Gibbs sampling. We also recall that

$$P(x_i \mid x_{-i}) \, P(x_{-i}) = P(x_i, x_{-i}) = P(x)$$

by definition of conditional probability. So we have

$$
\begin{align}
\frac{P(y) \, q(y \to x)}{P(x) \, q(x \to y)} &= \frac{P(y_i \mid y_{-i}) \, P(y_{-i})\, P(x_i \mid x_{-i}) }{P(x_i \mid x_{-i}) \, P(x_{-i})\, P(y_i \mid y_{-1})} &= 1
\end{align}
$$

**Advantages of Gibbs sampling**

- No need to tune proposal distribution
- Proposals are always accepted

**Disadvantages of Gibbs sampling**

- Need to be able to derive conditional probability distributions
- Need to be able to (cheaply) draw random samples from conditional probability distributions
- Can be very slow if parameters are correlated because you cannot take "diagonal" steps (draw picture to illustrate)

### Motivating example

We will use the toy example of estimating the bias of two coins given sample pairs $(z_1, n_1)$ and $(z_2, n_2)$ where $z_i$ is the number of heads in $n_i$ tosses for coin $i$.

#### Setup

```python
def bern(theta, z, N):
    """Bernoulli likelihood with N trials and z successes."""
    return np.clip(theta**z * (1-theta)**(N-z), 0, 1)
```

```python
def bern2(theta1, theta2, z1, z2, N1, N2):
    """Bernoulli likelihood with N trials and z successes."""
    return bern(theta1, z1, N1) * bern(theta2, z2, N2)
```

```python
def make_thetas(xmin, xmax, n):
    xs = np.linspace(xmin, xmax, n)
    widths =(xs[1:] - xs[:-1])/2.0
    thetas = xs[:-1]+ widths
    return thetas
```

```python
from mpl_toolkits.mplot3d import Axes3D

def make_plots(X, Y, prior, likelihood, posterior, projection=None):
    fig, ax = plt.subplots(1,3, subplot_kw=dict(projection=projection, aspect='equal'), figsize=(12,3))
    if projection == '3d':
        ax[0].plot_surface(X, Y, prior, alpha=0.3, cmap=plt.cm.jet)
        ax[1].plot_surface(X, Y, likelihood, alpha=0.3, cmap=plt.cm.jet)
        ax[2].plot_surface(X, Y, posterior, alpha=0.3, cmap=plt.cm.jet)
        for ax_ in ax: ax_._axis3don = False
    else:
        ax[0].contour(X, Y, prior, cmap=plt.cm.jet)
        ax[1].contour(X, Y, likelihood, cmap=plt.cm.jet)
        ax[2].contour(X, Y, posterior, cmap=plt.cm.jet)
    ax[0].set_title('Prior')
    ax[1].set_title('Likelihood')
    ax[2].set_title('Posteior')
    plt.tight_layout()
```

```python
thetas1 = make_thetas(0, 1, 101)
thetas2 = make_thetas(0, 1, 101)
X, Y = np.meshgrid(thetas1, thetas2)
```

#### Analytic solution

```python
a = 2
b = 3

z1 = 11
N1 = 14
z2 = 7
N2 = 14

prior = stats.beta(a, b).pdf(X) * stats.beta(a, b).pdf(Y)
likelihood = bern2(X, Y, z1, z2, N1, N2)
posterior = stats.beta(a + z1, b + N1 - z1).pdf(X) * stats.beta(a + z2, b + N2 - z2).pdf(Y)
make_plots(X, Y, prior, likelihood, posterior)
make_plots(X, Y, prior, likelihood, posterior, projection='3d')
```

#### Grid approximation

```python
def c2d(thetas1, thetas2, pdf):
    width1 = thetas1[1] - thetas1[0]
    width2 = thetas2[1] - thetas2[0]
    area = width1 * width2
    pmf = pdf * area
    pmf /= pmf.sum()
    return pmf
```

```python
_prior = bern2(X, Y, 2, 8, 10, 10) + bern2(X, Y, 8, 2, 10, 10)
prior_grid = c2d(thetas1, thetas2, _prior)
_likelihood = bern2(X, Y, 1, 1, 2, 3)
posterior_grid = _likelihood * prior_grid
posterior_grid /= posterior_grid.sum()
make_plots(X, Y, prior_grid, likelihood, posterior_grid)
make_plots(X, Y, prior_grid, likelihood, posterior_grid, projection='3d')
```

#### Metropolis

```python
a = 2
b = 3

z1 = 11
N1 = 14
z2 = 7
N2 = 14

prior = lambda theta1, theta2: stats.beta(a, b).pdf(theta1) * stats.beta(a, b).pdf(theta2)
lik = partial(bern2, z1=z1, z2=z2, N1=N1, N2=N2)
target = lambda theta1, theta2: prior(theta1, theta2) * lik(theta1, theta2)

theta = np.array([0.5, 0.5])
niters = 10000
burnin = 500
sigma = np.diag([0.2,0.2])

thetas = np.zeros((niters-burnin, 2), np.float)
for i in range(niters):
    new_theta = stats.multivariate_normal(theta, sigma).rvs()
    p = min(target(*new_theta)/target(*theta), 1)
    if np.random.rand() < p:
        theta = new_theta
    if i >= burnin:
        thetas[i-burnin] = theta
```

```python
kde = stats.gaussian_kde(thetas.T)
XY = np.vstack([X.ravel(), Y.ravel()])
posterior_metroplis = kde(XY).reshape(X.shape)
make_plots(X, Y, prior(X, Y), lik(X, Y), posterior_metroplis)
make_plots(X, Y, prior(X, Y), lik(X, Y), posterior_metroplis, projection='3d')
```

#### Gibbs

```python
a = 2
b = 3

z1 = 11
N1 = 14
z2 = 7
N2 = 14

prior = lambda theta1, theta2: stats.beta(a, b).pdf(theta1) * stats.beta(a, b).pdf(theta2)
lik = partial(bern2, z1=z1, z2=z2, N1=N1, N2=N2)
target = lambda theta1, theta2: prior(theta1, theta2) * lik(theta1, theta2)

theta = np.array([0.5, 0.5])
niters = 10000
burnin = 500
sigma = np.diag([0.2,0.2])

thetas = np.zeros((niters-burnin,2), np.float)
for i in range(niters):
    theta = [stats.beta(a + z1, b + N1 - z1).rvs(), theta[1]]
    theta = [theta[0], stats.beta(a + z2, b + N2 - z2).rvs()]

    if i >= burnin:
        thetas[i-burnin] = theta
```

```python
kde = stats.gaussian_kde(thetas.T)
XY = np.vstack([X.ravel(), Y.ravel()])
posterior_gibbs = kde(XY).reshape(X.shape)
make_plots(X, Y, prior(X, Y), lik(X, Y), posterior_gibbs)
make_plots(X, Y, prior(X, Y), lik(X, Y), posterior_gibbs, projection='3d')
```

Hierarchical models
---

Hierarchical models have the following structure - first we specify that the data come from a distribution with parameters $\theta$

$$
X \sim f(X\ | \ \theta)
$$

and that the parameters themselves come from another distribution with hyperparameters $\lambda$

$$
\theta \sim g(\theta \ | \ \lambda)
$$

and finally that $\lambda$ comes from a prior distribution

$$
\lambda \sim h(\lambda)
$$

More levels of hierarchy are possible - i.e you can specify hyper-hyperparameters for the distribution of $\lambda$ and so on.

The essential idea of the hierarchical model is because the $\theta$s are not independent but rather are drawn from a common distribution with parameter $\lambda$, we can share information across the $\theta$s by also estimating $\lambda$ at the same time.

As an example, suppose we have data about the proportion of heads after some number of tosses from several coins, and we want to estimate the bias of each coin. We also know that the coins come from the same mint and so might share some common manufacturing defect. There are two extreme approaches - we could estimate the bias of each coin from its coin toss data independently of all the others, or we could pool the results together and estimate the same bias for all coins. Hierarchical models provide a compromise where we shrink individual estimates towards a common estimate.

Note that because of the conditionally independent structure of hierarchical models, Gibbs sampling is often a natural choice for the MCMC sampling strategy.

#### Gibbs sampler example from [Robert and Casella, 10.17](http://www.springer.com/statistics/statistical+theory+and+methods/book/978-0-387-21239-5)

Suppose we have data of the number of failures ($y_i$) for each of 10 pumps in a nuclear plant. We also have the times ($_i$) at which each pump was observed. We want to model the number of failures with a Poisson likelihood, where the expected number of failure $\lambda_i$ differs for each pump. Since the time which we observed each pump is different, we need to scale each $\lambda_i$ by its observed time $t_i$.

We now specify the hierarchical model - note change of notation from the overview above - that $\theta$ is $\lambda$ (parameter) and $\lambda$ is $\beta$ (hyperparameter) simply because $\lambda$ is traditional for the Poisson distribution parameter.

The likelihood $f$ is
$$
\prod_{i=1}^{10} \text{Poisson}(\lambda_i t_i)
$$

We let the prior $g$ for $\lambda$ be

$$
\lambda \sim \text{Gamma}(\alpha, \beta)
$$
with $\alpha = 1$

and let $\beta$ to be a random variable to be estimated from the data

$$
\beta \sim \text{Gamma}(\gamma, \delta)
$$

with $\gamma = 0.01$ and $\delta = 1$.

There are 11 unknown parameters (10 $\lambda$s and $\beta$) in this hierarchical model.

The posterior is
$$
p(\lambda, \beta \ | \ y, t) = \prod_{i=1}^{10} \text{Poisson}(\lambda_i t_i) \times \text{Gamma}(\alpha, \beta) \times \text{Gamma}(\gamma, \delta)
$$

with the conditional distributions needed for Gibbs sampling given by

$$
p(\lambda_i \ | \ \lambda_{-i}, \beta, y, t) = \text{Gamma}(y_i + \alpha, t_i + \beta)
$$

and

$$
p(\beta \ | \ \lambda, y, t) = \text{Gamma}(10\alpha + \gamma, \delta + \sum_{i=1}^10 \lambda_i)
$$

```python
from numpy.random import gamma as rgamma # rename so we can use gamma for parameter name
```

```python
def lambda_update(alpha, beta, y, t):
    return rgamma(size=len(y), shape=y+alpha, scale=1.0/(t+beta))

def beta_update(alpha, gamma, delta, lambd, y):
    return rgamma(size=1, shape=len(y) * alpha + gamma, scale=1.0/(delta + lambd.sum()))

def gibbs(niter, y, t, alpha, gamma, delta):
    lambdas_ = np.zeros((niter, len(y)), np.float)
    betas_ = np.zeros(niter, np.float)

    lambda_ = y/t

    for i in range(niter):
        beta_ = beta_update(alpha, gamma, delta, lambda_, y)
        lambda_ = lambda_update(alpha, beta_, y, t)

        betas_[i] = beta_
        lambdas_[i,:] = lambda_

    return betas_, lambdas_
```

#### Setup

```python
alpha = 1.8
gamma = 0.01
delta = 1.0
beta0 = 1
y = np.array([5, 1, 5, 14, 3, 19, 1, 1, 4, 22], np.int)
t = np.array([94.32, 15.72, 62.88, 125.76, 5.24, 31.44, 1.05, 1.05, 2.10, 10.48], np.float)
niter = 1000
```

```python
betas, lambdas = gibbs(niter, y, t, alpha, gamma, delta)
print('%.3f' % betas.mean())
print('%.3f' % betas.std(ddof=1))
print(lambdas.mean(axis=0))
print(lambdas.std(ddof=1, axis=0))
```

    2.450
    0.701
    [0.06905699 0.15597648 0.10703469 0.12354692 0.61892855 0.61668945
     0.8360853  0.86031655 1.28183177 1.83725868]
    [0.02517981 0.09228058 0.04070685 0.03129867 0.29249547 0.13439944
     0.54271836 0.5482083  0.55513921 0.38773635]

```python
plt.figure(figsize=(8, 16))
for i in range(len(lambdas.T)):
    plt.subplot(5,2,i+1)
    plt.plot(lambdas[::10, i]);
    plt.title('Trace for $\lambda$%d' % i)
plt.tight_layout()
```


text:
# Output feedback control for the harmonic oscillator
The harmonic oscillator is a useful model for systems that have a dominating resonance frequency with no, or very little damping. An example of such systems is the sideway movement of a container hanging from a gantry crane moving containers on and off ships.

by Tosaka, from Wikimedia.org

Consider a container of mass $m=10^4$ kg, hanging from a wire of length $l=40$ m. We can control the system by applying an acceleration $u$ of the cart on top of the crane. The ODE describing the dynamics of the container is obtained by considering a reference frame fixed in the cart
$$ ml^2 \ddot{\theta} = -lmg\sin\theta + lm\cos\theta u + l\cos\theta w,$$
where $\theta$ is the angle of the wires to the vertical, and $w$ is a disturbance force from wind blowing on the container. The small-angle approximation $\sin\theta \approx \theta$ works well in this case, giving the model
$$\ddot{\theta} = -\frac{g}{l}\theta + \frac{1}{l}u + \frac{1}{lm}w.$$
Write $y=\theta$ and the model in the s-domain becomes
$$ Y(s) = \frac{1}{s^2 + \omega_0^2}\Big(bU(s) + kW(s)\big),$$
where $\omega_0^2 = \frac{g}{l}= \frac{9.8}{40} \approx 0.25$ and $b=1/l= 0.025$ and $k=1/(lm) = 2.5\times 10^{-5}$

## Discrete-time state-space model
Zero-order-hold sampling of the continuous-time harmonic oscillator gives the pulse-transfer function
$$ H(z) = \frac{(1-\beta)(z+1)b}{z^2 - 2\beta z + 1}, \quad \beta=\cos(\omega_0 h)$$
The system can be written on state-space form (observable canonical form) as
\begin{align}
x(k+1) &= \begin{bmatrix} 2\beta & 1\\-1 & 0\end{bmatrix}x(k) = \begin{bmatrix} 1-\beta\\1-\beta\end{bmatrix}b u(k) \\
y(k) &= \begin{bmatrix}1 & 0\end{bmatrix} x(k)
\end{align}
where
### Verification by symbolic computation

```python
import numpy as np
import sympy as sy
sy.init_printing(use_latex='mathjax', order='lex')

h,omega0 = sy.symbols('h,omega0', real=True, positive=True)
z = sy.symbols('z')
beta = sy.cos(omega0*h)
Phi = sy.Matrix([[2*beta, 1],[-1, 0]])
Gamma = sy.Matrix([[1-beta],[1-beta]])
C = sy.Matrix([[1, 0]])
H = sy.simplify(C*(z*sy.eye(2)-Phi).inv()*Gamma)
H
```

$$\left[\begin{matrix}\frac{- z \cos{\left (h \omega_{0} \right )} + z - \cos{\left (h \omega_{0} \right )} + 1}{z^{2} - 2 z \cos{\left (h \omega_{0} \right )} + 1}\end{matrix}\right]$$

### Choosing the sampling ratio $h$
We may use the rule-of-thumb $\omega_0 h \approx 0.2\, \text{to} \, 0.6$ for choosing the sampling period. For our specific case we have $\omega_0^2 = 0.25$. Let's choose $\omega_0 h = \pi/6 \approx 0.53$, so that $\cos(h\omega_0) = \frac{\sqrt{3}}{2} \approx 0.866$ This gives the discrete-time system (ignoring the disturbance for now)
\begin{align}
    x(k+1) &= \begin{bmatrix} \sqrt{3} & 1 \\ -1 & 0 \end{bmatrix}x(k) + \begin{bmatrix} 1-\frac{\sqrt{3}}{2}\\1-\frac{\sqrt{3}}{2}\end{bmatrix}bu(k)\\
    y(k) &= \begin{bmatrix} 1 & 0\end{bmatrix} x
 \end{align}

```python
omegaval = 0.5
hval = np.pi/6/omegaval
Phi_np = np.array(Phi.subs({h:hval, omega0:omegaval})).astype(np.float64)
Phi_np
```

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

```python
Gamma_np = np.array(Gamma.subs({h:hval, omega0:omegaval})).astype(np.float64)
Gamma_np
```

    array([[0.1339746],
           [0.1339746]])

## Obervability
The observability matrix for this second order system becomes
$$ W_o = \begin{bmatrix} C\\C\Phi\end{bmatrix} = \begin{bmatrix} 1 & 0\\ \sqrt{3} & 1\end{bmatrix}, $$
with determinant
$$\det W_0 =  1-0 = 1 \neq 0.$$

```python
C_n = np.array(C).astype(np.float64)
Wo_n = np.vstack((C_n, np.dot(C_n,Phi_np)))
Wo_n
```

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

```python
np.linalg.det(Wo_n)
```

$$1.0$$

## Observer design
The observer is given by
\begin{align}
\hat{x}(k+1) &= \Phi \hat{x}(k) +\Gamma u(k) + K\big(y(k) - C\hat{x}(k)\big)\\
&= \big(\Phi - KC\big)\hat{x}(k) + \Gamma u(k) + Ky(k)
\end{align}
where
$$ KC = \begin{bmatrix} k_1\\ k_2\end{bmatrix}\begin{bmatrix}1 & 0\end{bmatrix} = \begin{bmatrix} k_1 & 0\\ k_2 & 0\end{bmatrix} $$
and
$$ \Phi - KC = \begin{bmatrix} \sqrt{3} & 1\\-1 & 0\end{bmatrix}-\begin{bmatrix} k_1 & 0\\ k_2 & 0\end{bmatrix}=\begin{bmatrix} \sqrt{3}-k_1 & 1\\-1-k_2 & 0\end{bmatrix}$$

```python
k1,k2 = sy.symbols('k1,k2')
K = sy.Matrix([[k1], [k2]])
Phi_o=Phi.subs({h:hval, omega0:omegaval}) - K*C
Phi_o
```

$$\left[\begin{matrix}- k_{1} + 1.73205080756888 & 1\\- k_{2} - 1 & 0\end{matrix}\right]$$

The characteristic polynomial of the observer is given by
\begin{align}
\det \left( zI - (\Phi-KC) \right) &= \det \left( \begin{bmatrix} z & 0\\0 & z \end{bmatrix} - \begin{bmatrix} \sqrt{3}-k_1 & 1\\-1-k_2\end{bmatrix} \right)\\
&= \det \begin{bmatrix} z-\sqrt{3}+k_1 & -1 \\ 1+k_2 & z\end{bmatrix}\\
&= z^2 + (-\sqrt{3}+k_1) z + (1 + k_2)
\end{align}

```python
z = sy.symbols('z')
zIminPhio = z*sy.eye(2) - Phi_o
ch_poly = sy.Poly(zIminPhio.det(), z)
ch_poly.as_expr()
```

$$1.0 k_{2} + 1.0 z^{2} + z \left(1.0 k_{1} - 1.73205080756888\right) + 1.0$$

### Desired observer poles
Assume we want the observer to have a double pole in $z=p_o$, which gives the following  desired characteristic polynomial for the observer
$$ (z-p_o)^2 = z^2 - 2p_oz + p_o^2$$
#### Determining the observer gain
As usual, we determine the unknown parameters by setting the coefficients equal in the two characteristic polynomials above. This gives the system of linear equations
\begin{align}
z^1: \qquad k_1 - \sqrt{3} &= -2p_o\\
z^0: \qquad 1+k_2 &= p_o^2
\end{align}
which has the straight-forward solution
\begin{align}
k_1 &= \sqrt{3}-2p_o\\
k_2 &= p_o^2 - 1.
\end{align}
For a deadbeat observer we get
\begin{align}
k_1 &= \sqrt{3}\\
k_2 &= - 1.
\end{align}

The solution is so simple because the state-space system was written on observable canonical form.

```python
po = sy.symbols('p_o')
K = sy.Matrix([[Phi[0,0]-2*po],[Phi[1,0]+po**2]])
K
```

$$\left[\begin{matrix}- 2 p_{o} + 2 \cos{\left (h \omega_{0} \right )}\\p_{o}^{2} - 1\end{matrix}\right]$$

```python
Phi_o = Phi - K*C
Phi_o
```

$$\left[\begin{matrix}2 p_{o} & 1\\- p_{o}^{2} & 0\end{matrix}\right]$$

```python
Phi_o.eigenvals()
```

$$\left \{ p_{o} : 2\right \}$$

## Reachability

```python
Wc = sy.BlockMatrix([[Gamma, Phi*Gamma]]).as_explicit()
Wc
```

$$\left[\begin{matrix}- \cos{\left (h \omega_{0} \right )} + 1 & 2 \left(- \cos{\left (h \omega_{0} \right )} + 1\right) \cos{\left (h \omega_{0} \right )} - \cos{\left (h \omega_{0} \right )} + 1\\- \cos{\left (h \omega_{0} \right )} + 1 & \cos{\left (h \omega_{0} \right )} - 1\end{matrix}\right]$$

```python
Wc.det()
```

$$- 2 \cos^{3}{\left (h \omega_{0} \right )} + 2 \cos^{2}{\left (h \omega_{0} \right )} + 2 \cos{\left (h \omega_{0} \right )} - 2$$

```python
Wc.det().subs({h:hval, omega0:omegaval})
```

## State feedback
Introducing the state-feedback control law
$$ u = -l_1x_1 - l_2 x_2 + l_0y_{ref} = -Lx + l_0y_{ref}$$
gives the closed-loop state-space system
\begin{align}
x(k+1) &= \Phi x(k) +\Gamma\big(-Lx(k) + l_0y_{ref}(k)\big) + \Gamma v(k) = \left( \Phi - \Gamma L \right) x(k) + l_0\Gamma y_{ref}(k) + \Gamma v(k)\\
y(k) &= C x(k),
\end{align}

```python
l1,l2 = sy.symbols('l1,l2')
L = sy.Matrix([[l1, l2]])
Phi_c=Phi.subs({h:hval, omega0:omegaval}) - Gamma.subs({h:hval, omega0:omegaval})*L
Phi_c
```

$$\left[\begin{matrix}- 0.133974596215561 l_{1} + 1.73205080756888 & - 0.133974596215561 l_{2} + 1\\- 0.133974596215561 l_{1} - 1 & - 0.133974596215561 l_{2}\end{matrix}\right]$$

The characteristic polynomial of the closed-loop system is given by
\begin{align}
\det \left( zI - (\Phi-\Gamma L) \right) \end{align}

```python
zIminPhic = z*sy.eye(2) - Phi_c
zIminPhic
```

$$\left[\begin{matrix}0.133974596215561 l_{1} + z - 1.73205080756888 & 0.133974596215561 l_{2} - 1\\0.133974596215561 l_{1} + 1 & 0.133974596215561 l_{2} + z\end{matrix}\right]$$

```python
ch_poly = sy.Poly(zIminPhic.det(), z)
ch_poly.as_expr()
```

$$0.133974596215561 l_{1} - 0.366025403784438 l_{2} + 1.0 z^{2} + z \left(0.133974596215561 l_{1} + 0.133974596215561 l_{2} - 1.73205080756888\right) + 1.0$$

### Desired closed-loop characteristic polynomial
Here we will aim for a pair of closed-loop poles that are critically damped, and about as fast as the plant poles.
$$ p_1 = p_2 = \sqrt{(1-\cos(\omega_0 h))^2 + \sin(\omega_0 h)^2} = p$$
$$ A_c(z) = (z-p)^2 = z^2 - 2pz + p^2 $$
In the same spirit as when designing an RST controller using the polynomial approach, we set the calculated characteristic polynomial - obtained when introducing the linear state feedback- equal to the desired characteristic polynomial.

```python
p = sy.symbols('p')
des_ch_poly = sy.Poly((z-p)**2, z)
dioph_eqn = ch_poly - des_ch_poly
sol = sy.solve(dioph_eqn.coeffs(), (l1,l2))
sol
```

$$\left \{ l_{1} : 2.0 p^{2} - 10.9282032302755 p + 7.46410161513779, \quad l_{2} : - 2.0 p^{2} - 4.00000000000001 p + 5.46410161513777\right \}$$

### Determining $l_0$

State feedback does not change the numerator of the open-loop system, which is
$$B(z) = (1-\beta)(z+1)b$$
The characteristic polynomial of the closed-loop system is already designed to be
$$A_c(z) = (z-p)^2$$.
This means that the closed-loop system has the pulse-transfer function
$$H_c(z) = \frac{(1-\beta)(z+1)b}{(z-p)^2} l_0$$
We want this pulse-transfer function to have unit DC-gain which gives
$$ H_c(1) = \frac{(1-\beta)2b}{(1-p)^2}l_0$$
$$ l_0 = \frac{A_c(1)}{B(z)} = \frac{(1-p)^2}{2(1-\beta)b}.$$

```python
l0 = (1-p)**2/(2*(1-beta)) # We have assumed b=1 all along
```

ext: ipynb

max_stars_repo_path: state-space/notebooks/Output feedback controller for the harmonic oscillator-observer canonical form.ipynb

max_issues_repo_name: alfkjartan/control-computarizado

max_forks_repo_name: alfkjartan/control-computarizado


text:
# Reference implementation of the QMC algorithm on a uniform Cartesian grid (1D version) #

### Abstract ###

This notebook presents a reference implementation of a quantum Monte-Carlo algorithm designed for sampling the ground state of the molecular Hamiltonian represented on a uniform Cartesian grid. Details of the theory as well as the aplications to model systems in 1D and 3D can be found in the accompaning paper. Here, we provide a simplified version of the code for 1D Hamiltonians and demonstrate its perfomance for several test cases.

## Implementation of the quantum Monte-Carlo method ##

**Random_walk** class provided below implements a grid propagator defined by the following expression for the transition probability:

$$ p_d = \frac{1}{2\pi} \int_{ -\pi}^{\pi} \cos(k d) e^{-\frac{2\tau}{\delta^{2}}\sin^2 (k/2)} dk, $$

where d is the length of the step on the 1-D integer lattice ($d = 0, \pm 1, \pm 2, \pm 3,...$), $\tau$ is the timestep of the diffusion random walk, and $\delta$ is a uniform grid spacing. The integral is calculated numerically using a general 1-D quadrature subroutine *quad* from *NumPy* for a series of d such that $p_d < 10^{-8}$ (the default value of thresh). The absolute integration accuracy is set to thresh/100 and the relative accuracy is $10^{-3}$.

```python
# Functor class generating grid displacements

class Random_walk:

    def __init__(self, delta, tau, thresh = 1e-8):

        # The names of the input parameters should be self-explanatory

        import numpy as np
        from scipy.integrate import quad

        assert delta > 0. and tau > 0.
        self.delta = delta
        self.tau = tau

        # Calculate the cumulative probabilites
        disp = [0]
        pd = []

        def kernel(k, d, t):
            return 0.5 * np.exp(-2. * t * np.sin(k/2.)**2 / d**2) / np.pi

        terminate = False
        while not terminate:
            d = disp[-1]
            r = quad(kernel, -np.pi, np.pi, (self.delta, self.tau), weight='cos', wvar=d, epsabs=thresh/100.,epsrel=1e-3)
            if r[1] > max(thresh/100., r[0] * 1e-3):
                print("WARNING! Insufficient integration accuracy for displacement {0}!\n".format(r[1]))

            pd.append(r[0])
            if abs(r[0]) < thresh:
                terminate = True
            else:
                disp.append(d + 1)

        self.pd = np.array(pd[-1:0:-1] + pd)
        self.cumm_pd = np.cumsum(self.pd)
        self.disp = np.array([-disp[j] for j in range(len(disp) - 1, 0, -1)] + disp)

    def __call__(self, num_samples):

        import numpy as np

        r = np.random.choice(self.disp, size=num_samples, p=self.pd)

        return r

    def __call_old__(self, num_samples):

        import numpy as np

        r = np.random.rand(num_samples)
        #print len(r)
        compare_mask = self.cumm_pd[np.newaxis,:] > r[:,np.newaxis]
        #print compare_mask.shape
        return np.array(self.disp[np.argmax(compare_mask, axis=1)])

```

It can be shown that $p_d$ resembles Gaussian distribution in the limit $\delta \to 0$ ($\delta << \sqrt(\tau)$). However, the two may differ significantly if $\delta$ is comparable to $\sqrt(\tau)$ as shown by the following example

```python
import numpy as np
from matplotlib import pyplot as plt
%matplotlib inline

delta = 0.1
tau = 0.015
rw = Random_walk(delta=delta, tau=tau)
disp = rw.disp
pd = rw.pd
#Discretized Gaussian
pd_gau = delta / np.sqrt(2. * np.pi * tau)  * np.exp(-delta**2 * disp**2 / (2. * tau))

plt.plot(disp, pd, 'bo', label='Numerical propagator')
plt.plot(disp, pd_gau,'r+', label='Discretized gaussian')

plt.xlabel("Displacement")
plt.ylabel("Transition probablity")
plt.legend()
```

The instances of the **Random_walk** class can be used to generate random walks on integer lattices by invoking the '__call__' method. In order to test the validity of our implementation we generate a series of *n_samples* and demonstrate that their distribution agrees with the analytic result.

```python
# Check whether the class produces correct displacement statistics
n_samples = 100000
displacements = rw(n_samples)
print("Number of sampled displacements {0}".format(displacements.shape[0]))
assert displacements.shape[0] == n_samples
#print(displacements)

hist = {}
for i in displacements:
    if i in hist:
        hist[i] += 1
    else:
        hist[i] = 1

# Normalize and compare with the exact distribution
plt.plot(hist.keys(), np.array(hist.values())/float(n_samples), 'go', label='Generated distribution')
plt.plot(disp, pd, 'b*', label='Exact propagator')

plt.xlabel("Displacement")
plt.ylabel("fraction of samples")
plt.legend()
```

**Walker_ensemble** class implements the core of the grid-based quantum Monte-Carlo algorithm. It uses instances of **Random_walk** class to generate diffusion random walks on spatial grids. The imaginary time propagation is executed by invoking a *one_time_step* method which in turn calls *_diffuse_and_branch, _update_phases* and *_annihilate_walkers* methods. The energy is then 'measured' by the growth estimator. The class also provides a functionality to implement constraints which can for example be used to perform fixed-node simulations and impose spatial symmetry. Present code is restricted to 1-D systems.

```python
class Walker_ensemble:

    def __init__(self, random_walk, potential, n_walkers, n_alpha, n_beta, constraint=None, pop_dyn=False, S=0.0):

        import numpy as np
        from sympy.combinatorics.permutations import Permutation
        import itertools as it

        L = 5. # Initial walker population will be (uniformly) distributed in the interval [-L, L]

        # Are we performing a regular run or recording population dynamics?
        self.pop_dyn = pop_dyn

        # System specification
        self.disp_generator = random_walk # An instance of random walk class
        self.V = potential # A callable python object capable of operating on a walker (=set of particle coordinates)
        self.S = S # initial guess for the growth estimator
        self.apply_constraint = False
        self.constraint = None

        if constraint != None:
            self.constraint = constraint # A callable python object capable of operating on a walker and returning True/False
            self.apply_constraint = True

        # Walker positions and weights;
        # The former are expressed by the coordinates on the integer lattice.
        # The latter should ALWAYS be \pm 1!
        self.n_alpha = n_alpha # number of alpha electrons
        self.n_beta = n_beta   # number of beta electrons
        self.lookup_a = {} # all permutations of alpha electrons with corresponding parities
        self.lookup_b = {} # same for beta

        pa = it.permutations(np.arange(n_alpha))
        pb = it.permutations(np.arange(n_beta))

        for i in pa:
            parity = 1 if Permutation(i).is_even else -1
            self.lookup_a[i] = parity
        for i in pb:
            parity = 1 if Permutation(i).is_even else -1
            self.lookup_b[i] = parity

        self.walker_weights = np.ones(n_walkers, dtype=int)
        self.walker_coords = np.zeros((n_walkers, n_alpha + n_beta), dtype=int)

        # Timing info for reference purposes
        self.time = {'total': 0.0, 'branching': 0.0, 'phase_update': 0.0, 'annihilation' : 0.0, 'n_steps' : 0}

        # Initialize walker ensemble
        # and make sure that the walkers conform to the Pauli principle and don't violate constraints
        d = self.disp_generator.delta
        grid_init = np.arange(-L/d, L/d + 1, dtype=int)
        kwargs = {'size' : n_alpha, 'replace' : False} # replace = False makes sure that the Pauli principle is not violated
        alpha_coord = np.apply_along_axis(np.random.choice, 1, np.broadcast_to(grid_init, (n_walkers, grid_init.shape[0])), **kwargs)
        alpha_coord = np.apply_along_axis(np.sort, 1, alpha_coord)
        kwargs = {'size' : n_beta, 'replace' : False}
        beta_coord = np.apply_along_axis(np.random.choice, 1, np.broadcast_to(grid_init, (n_walkers, grid_init.shape[0])), **kwargs)
        beta_coord = np.apply_along_axis(np.sort, 1, beta_coord)

        # Check if the coordinates of alpha and beta particles indeed satisfy the Pauli principle
        all_distinct = lambda ar : len(ar) == len(np.unique(ar))
        a = np.apply_along_axis(all_distinct, 1, alpha_coord)
        b = np.apply_along_axis(all_distinct, 1, beta_coord)

        if self.apply_constraint:
            constr_a = np.apply_along_axis(self.constraint, 1, alpha_coord)
            constr_b = np.apply_along_axis(self.constraint, 1, beta_coord)
            a = np.logical_and(a, constr_a)
            b = np.logical_and(b, constr_b)

        alpha_ok = True if n_alpha == 0 else np.all(a)
        beta_ok = True if n_beta == 0 else np.all(b)

        assert alpha_ok and beta_ok

        # Combine the coordinates of alpha and beta electrons (the resulting list may contain duplicates)
        self.walker_coords = np.hstack((alpha_coord, beta_coord))

    def one_time_step(self):

        # High level grid QMC driver

        import numpy as np
        import time

        start = time.clock()

        n_walkers_old = self.get_number_of_walkers()

        self._diffuse_and_branch()

        self.time['branching'] += time.clock() - start

        t1 = time.clock()

        self._update_phases()

        self.time['phase_update'] += time.clock() - t1

        t2 = time.clock()

        self._annihilate_walkers()

        self.time['annihilation'] += time.clock() - t2
        self.time['total'] += time.clock() - start
        self.time['n_steps'] += 1

        n_walkers_new = self.get_number_of_walkers()
        #print("Number of walkers at the end of the step is {0}".format(n_walkers_new))

        if not self.pop_dyn:
            self.S = self.S + 1./ self.disp_generator.tau * np.log(n_walkers_old/float(n_walkers_new))

    def _diffuse_and_branch(self):

        import numpy as np

        # Perform diffusion displacement

        random_walk =  self.disp_generator
        d = self.disp_generator.delta # grid spacing
        t = self.disp_generator.tau   # imaginary time step

        # the code snippet below will work correctly only if all weights are \pm 1
        n_walkers = self.get_number_of_walkers()
        assert n_walkers == self.walker_weights.shape[0]
        walker_size = self.n_alpha + self.n_beta
        disp = random_walk(n_walkers*walker_size)
        disp = disp.reshape(n_walkers, walker_size)

        # Some displacements give rise to the particles with
        # identical spin labels sitting on top of each other;
        # Those violate the Pauli principle and should be removed

        new_walker_coords = self.walker_coords + disp
        new_coords_a, new_coords_b = np.hsplit(new_walker_coords, [self.n_alpha])
        all_distinct = lambda ar : len(ar) == len(np.unique(ar))
        a = np.apply_along_axis(all_distinct, 1, new_coords_a)
        b = np.apply_along_axis(all_distinct, 1, new_coords_b)
        #print(a)
        #print(b)
        alive = np.logical_and(a,b) # positions of alive walkers

        # Apply constraints if any

        if self.apply_constraint:
            constr_a = np.apply_along_axis(self.constraint, 1, new_coords_a)
            constr_b = np.apply_along_axis(self.constraint, 1, new_coords_b)
            c = np.logical_and(constr_a, constr_b)
            alive = np.logical_and(c, alive)

        self.walker_coords = np.array(self.walker_coords[alive,:], dtype=int)
        self.walker_weights = np.array(self.walker_weights[alive], dtype=int)
        new_walker_coords = np.array(new_walker_coords[alive,:], dtype=int)

        # After this step one may still have all-electron walkers sitting atop of each other
        # but the algorithm is blind to that
        # Perform branching
        # NOTE: only alive walkers are participating in branching

        m = np.exp(-0.5 * t * ( np.apply_along_axis(self.V, 1, d*self.walker_coords) + np.apply_along_axis(self.V, 1, d*new_walker_coords) - 2. * self.S))
        r = np.random.rand(m.shape[0])

        # Assign integer weights and remove dead walkers

        new_walker_weights = np.asarray(m + r, dtype=int) * self.walker_weights
        alive = (new_walker_weights != 0)
        self.walker_coords = np.array(new_walker_coords[alive,:], dtype=int)
        self.walker_weights = np.array(new_walker_weights[alive], dtype=int)

        # Some walkers can acquire abs. weights > 1
        # Those clusters need to be separated
        # That will be done after the annihilation step

    def _update_phases(self):

        import numpy as np

        # Update walker signs based on whether the coordinates got interchanged
        # Split alpha and beta subsets first

        alpha_coord, beta_coord = np.hsplit(self.walker_coords, [self.n_alpha])
        p_alpha = np.argsort(alpha_coord, axis=1)
        p_beta = np.argsort(beta_coord, axis=1)

        parity = lambda p: self.lookup_a[tuple(p)]
        parity_a = np.apply_along_axis(parity, 1, p_alpha)
        parity = lambda p: self.lookup_b[tuple(p)]
        parity_b = np.apply_along_axis(parity, 1, p_beta)

        self.walker_weights = self.walker_weights * parity_a * parity_b
        alpha_coord = np.sort(alpha_coord, axis=1)
        beta_coord = np.sort(beta_coord, axis=1)

        self.walker_coords = np.hstack((alpha_coord, beta_coord))

    def _annihilate_walkers(self, catch_collisions=False):

        # Find and collect duplicates (perform annihilation)
        # A "quick and dirty" way to accomplish that is to hash
        # the walker coordinates & perform sorting;

        converter = lambda a: hash(a.tostring())
        walker_coords_str = np.apply_along_axis(converter, 1, self.walker_coords)

        if catch_collisions:

            # Check if we have any collisions, i.e. all the hashes are distinct for distinct walkers

            unique_hashes = set(walker_coords_str)
            walker_index = np.arange(walker_coords_str.shape[0], dtype=int)

            if len(unique_hashes) != len(walker_index):

                # Check if there is a collision
                for h in unique_hashes:
                    walkers_with_hash_h = np.array(self.walker_coords[(h == walker_coords_str), :], dtype=int)
                    if len(walkers_with_hash_h) == 1:
                        continue
                    else:
                        assert np.all(walkers_with_hash_h[0][:] == walkers_with_hash_h) # broadcasted!

        unique_walkers_str_, idx = np.unique(walker_coords_str, return_index = True)
        combined_weights = []

        for i in idx:
            idx_duplicates = (walker_coords_str[i] == walker_coords_str)
            combined_weights.append(np.sum(self.walker_weights[idx_duplicates]))

        self.walker_coords = np.array(self.walker_coords[idx, :], dtype=int)
        self.walker_weights = np.array(combined_weights, dtype=int)

        self._expand()

    def _expand(self):

        # Simplified version of expand function
        # implemented without using fancy NumPy features

        n = self.walker_weights.shape[0]
        split_weights = []
        split_coords = []

        for i in range(n):
            num_of_instances = abs(self.walker_weights[i])
            sign = 1 if self.walker_weights[i] > 0 else -1
            for j in range(num_of_instances):
                split_weights.append(sign)
                split_coords.append(self.walker_coords[i, :])

        self.walker_coords = np.vstack(tuple(split_coords))
        self.walker_weights = np.array(split_weights, dtype=int)

    def get_number_of_walkers(self):
        import numpy as np
        return np.sum(np.abs(self.walker_weights))

    def get_energy(self):
        return self.S

```

## Applications ##

### Fermions in harmonic trap ###

Non-interacting fermions in harmonic trap provide a convenient test case for the algorithm of the previous section for the following reasons:

1. Numerically exact energies for finite (non-zero) grid spacings can be readily calculated as sums of one-particle eigenvalues.
2. The fermions are confined by soft non-divergent potential so that the branching factors in QMC are always finite.

In the following we will assume that the potential energy oprator is $V = \frac{1}{2} \sum_{i=1}^{N_e} x_i^{2}$ and consider the cases where $N_e = 1, 2, 3$ and $4$

```python
def V_harm(x):
    return 0.5 * np.sum(x**2)
```

### Case $N_e = 1$ ###

As a first step, we determine the reference energies of one-particle states via direct diagonalization of the Hamiltonian obtained by approximating the kinetic energy operator using a three point finite difference  formula on a **uniform** grid with sapcing $\delta = 0.1$ a.u. spanning the interval $[-L, L]$:

$$ \hat{T}\psi(x_i) = \frac{2\psi(x_i) - \psi(x_{i + 1}) - \psi(x_{i - 1})}{2\delta^{2}}, $$

where $x_{i\pm1} = x_i \pm \delta$. Scince our goal is to approximate the solutions of the stationary Shrodinger equation for $L \to +\infty$ we perform a series of diagonalizations for increasing $L$ requiring that the energies of a selected subset of low-lying states have **at least 6 digits of accuracy**. The calculations are performed assuming the boundary condition $\psi(L) = \psi(-L) = 0$.

```python
from numpy import linalg as LA

def construct_H(L, delta):
    # Input:
    # L     - grid span
    # delta - grid spacing
    # Output:
    # $H_{\delta}$ - matrix of the discretized Hamiltonian
    x = np.linspace(-L, L, int(2. * L / delta) + 1) # grid points
    V = np.diag(0.5 * x**2)                        # potential energy
    T = -1./ (2. * delta**2) * ( np.eye(len(x), k = 1) - 2. * np.eye(len(x)) + np.eye(len(x), k =-1) ) # kinetic energy
    return T + V

delta = 0.1
L = 4. # Initial grid span
eps = 1e-6
converged = False
e5 = np.zeros(5)

while not converged:
    H = construct_H(L, delta)
    e = LA.eigvalsh(H)
    e5_new = np.array(e[0:5])
    if max(abs(e5_new - e5)) < eps:
        converged = True
    else:
        L *= 2.

    e5 = np.copy(e5_new)

print("Energies of harmonic oscillator states on a uniform grid with spacing {0}".format(delta))
print(e5)
```

    Energies of harmonic oscillator states on a uniform grid with spacing 0.1
    [ 0.4996873   1.49843574  2.49593063  3.49216962  4.48715031]

Alternatively, the ground state of the 1-D harmonic oscillator can be sampled using quantum Monte-Carlo scheme:

```python
rw0101 = Random_walk(delta=0.1, tau=0.1) # define a new random walk on a grid with spacing delta = 0.1 and tau = 0.1
D1 = Walker_ensemble(random_walk=rw0101, potential=V_harm, n_walkers=2000, n_alpha=1, n_beta=0, S=0.5)

N_equil = 500
N_prod = 2048
E = np.zeros(N_prod)
N = np.zeros(N_prod)

#Perform N_equil equilibration steps
for i in range(N_equil):
    D1.one_time_step()

print("Number of walkers at the end of equilibration run is {0}".format(D1.get_number_of_walkers()))
print("Starting production run...")

for i in range(N_prod):
    D1.one_time_step()
    E[i] = D1.get_energy()
    N[i] = D1.get_number_of_walkers()

print("Average energy at the end of the production run: {0:13.6f}".format(np.mean(E)))

```

    Number of walkers at the end of equilibration run is 2004
    Starting production run...
    Average energy at the end of the production run:      0.500943

In order to calculate statistical errors of the energy estimators we used blocking algorithm as described in [H.Flyvbjerg et al. *J. Chem. Phys.* 91, 461 (1989)](https://www.google.com/url?sa=t&rct=j&q=&esrc=s&source=web&cd=1&cad=rja&uact=8&ved=2ahUKEwiBjaaf68ndAhXH6YMKHSuyCdgQFjAAegQIBBAB&url=https%3A%2F%2Faip.scitation.org%2Fdoi%2F10.1063%2F1.457480&usg=AOvVaw0DEeYW79PHKKZfPb61xsKZ). The notebook relies on the python implementation of the bloking analysis due to J.Spencer which is freely available from the Github page [pyblock](https://github.com/jsspencer/pyblock)

```python
pyblock_is_installed = True
try:
    import pyblock as pb
except ImportError as e:
    print("PyBlock was not found! The error estimates will not be available.")
    pyblock_is_installed = False

```

```python
# If PyBlock is found we can proceed to calculate statistical errors
if pyblock_is_installed:
    reblock_data = pb.blocking.reblock(E)
    opt = pb.blocking.find_optimal_block(len(E), reblock_data)
    print("Optimal block: {}".format(opt))
    print(reblock_data[opt[0]])
    print("< E > = {0:10.6f} +- {1:10.6f}".format(float(reblock_data[opt[0]][2]), float(reblock_data[opt[0]][4])))
```

    Optimal block: [5]
    BlockTuple(block=5, ndata=64, mean=array(0.5009431527081165), cov=array(0.00014861497030582347), std_err=array(0.0015238467478813254), std_err_err=array(0.00013575505810765735))
    < E > =   0.500943 +-   0.001524

Error analysis demonstrates that the energy estimate agrees with the exact diagonlaization result within the statistical error margins.

### Case $N_e = 2$ ###

Admissible values of the total spin for 2 electrons in a harmonic trap are 0 (singlet) and 1 (triplet). Singlet state corresponds to the doubly occupied ground state of one particle problem and is nodeless. Triplet state exhibits a node at $x_1 = x_2$. The reference energies of the states are calculated below.

```python
# Singlet state occupation numbers
singlet_occ = np.array([2, 0, 0, 0, 0])
# Triplet state occupation numbers
triplet_occ = np.array([1, 1, 0, 0, 0])

print("E(singlet) = {0:13.6f} a.u.".format(np.dot(singlet_occ, e5)))
print("E(triplet) = {0:13.6f} a.u.".format(np.dot(triplet_occ, e5)))
```

    E(singlet) =      0.999375 a.u.
    E(triplet) =      1.998123 a.u.

```python
S2 = Walker_ensemble(random_walk=rw0101, potential=V_harm, n_walkers=2000, n_alpha=1, n_beta=1, S = 1.0)
T2 = Walker_ensemble(random_walk=rw0101, potential=V_harm, n_walkers=2000, n_alpha=2, n_beta=0, S = 2.0)

E_S = np.zeros(N_prod)
N_S = np.zeros(N_prod)

E_T = np.zeros(N_prod)
N_T = np.zeros(N_prod)

for i in range(N_equil):
    S2.one_time_step()
    T2.one_time_step()

print("Number of walkers at the end of equilibration run is {0} (singlet)".format(S2.get_number_of_walkers()))
print("Number of walkers at the end of equilibration run is {0} (triplet)".format(T2.get_number_of_walkers()))
print("Starting production runs...")

for i in range(N_prod):
    S2.one_time_step()
    T2.one_time_step()

    E_S[i] = S2.get_energy()
    N_S[i] = S2.get_number_of_walkers()
    E_T[i] = T2.get_energy()
    N_T[i] = T2.get_number_of_walkers()

print("Average energy at the end of the production run (S = 0): {0:13.6f}".format(np.mean(E_S)))
print("Average energy at the end of the production run (S = 1): {0:13.6f}".format(np.mean(E_T)))

```

    Number of walkers at the end of equilibration run is 2002 (singlet)
    Number of walkers at the end of equilibration run is 2004 (triplet)
    Starting production runs...
    Average energy at the end of the production run (S = 0):      1.000598
    Average energy at the end of the production run (S = 1):      1.998883

```python
# A helper function
def reblock(a):
    if pyblock_is_installed:
        reblock_data = pb.blocking.reblock(a)
        opt = pb.blocking.find_optimal_block(len(a), reblock_data)
        print("Optimal block: {}".format(opt))
        #print(reblock_data[opt[0]])
        return "< E > = {0:10.6f} +- {1:10.6f}".format(float(reblock_data[opt[0]][2]), float(reblock_data[opt[0]][4]))
    else:
        print("PyBlock is not installed!")
        return None

print("Singlet: " + reblock(E_S))
print("Triplet: " + reblock(E_T))

```

    Optimal block: [5]
    Singlet: < E > =   1.000598 +-   0.001908
    Optimal block: [5]
    Triplet: < E > =   1.998883 +-   0.002471

### Case $N_e = 3$ ###

Similar to $N_e = 2$ the system can have two values of total spin: $S = 1/2$ and $S = 3/2$. However, in contrast to the previous case none of the states is nodeless. The reference energies of the states are computed below.

```python
# Doublet state occupation numbers
doublet_occ = np.array([2, 1, 0, 0, 0])
# Quartet state occupation numbers
quartet_occ = np.array([1, 1, 1, 0, 0])

print("E(doublet) = {0:13.6f} a.u.".format(np.dot(doublet_occ, e5)))
print("E(quartet) = {0:13.6f} a.u.".format(np.dot(quartet_occ, e5)))
```

    E(doublet) =      2.497810 a.u.
    E(quartet) =      4.494054 a.u.

Quantum Monte-Carlo calculations performed for S=1/2 and S=3/2 produce energy estimates that are too low compared to the exact results if the number of random walkers is around 2000. This example demonstrates that **similar to FCIQMC present method guarantees convergence to the target state with correct nodal structure only provided the total number of walkers exceeds the critical value ($N_c$).** The $N_c$ values can be estimated from the calculations with the fixed energy onset.

```python
D3 = Walker_ensemble(random_walk=rw0101, potential=V_harm, n_walkers=2000, n_alpha=2, n_beta=1, S = 2.5)
Q3 = Walker_ensemble(random_walk=rw0101, potential=V_harm, n_walkers=2000, n_alpha=3, n_beta=0, S = 4.5)

E_D = np.zeros(N_prod)
N_D = np.zeros(N_prod)

E_Q = np.zeros(N_prod)
N_Q = np.zeros(N_prod)

for i in range(N_equil):
    D3.one_time_step()
    Q3.one_time_step()

print("Number of walkers at the end of equilibration run is {0} (doublet)".format(D3.get_number_of_walkers()))
print("Number of walkers at the end of equilibration run is {0} (quartet)".format(Q3.get_number_of_walkers()))
print("Starting production runs...")

for i in range(N_prod):
    D3.one_time_step()
    Q3.one_time_step()

    E_D[i] = D3.get_energy()
    N_D[i] = D3.get_number_of_walkers()
    E_Q[i] = Q3.get_energy()
    N_Q[i] = Q3.get_number_of_walkers()

print("Average energy at the end of the production run (S = 1/2): {0:13.6f}".format(np.mean(E_D)))
print("Average energy at the end of the production run (S = 3/2): {0:13.6f}".format(np.mean(E_Q)))
```

    Number of walkers at the end of equilibration run is 2065 (doublet)
    Number of walkers at the end of equilibration run is 2222 (quartet)
    Starting production runs...
    Average energy at the end of the production run (S = 1/2):      2.255693
    Average energy at the end of the production run (S = 3/2):      3.378452

```python
# Population dynamics runs
D3 = Walker_ensemble(random_walk=rw0101, potential=V_harm, n_walkers=500, n_alpha=2, n_beta=1, S = 2.51, pop_dyn=True)
Q3 = Walker_ensemble(random_walk=rw0101, potential=V_harm, n_walkers=500, n_alpha=3, n_beta=0, S = 4.505, pop_dyn=True)
N_max_Q = 9000
N_max_D = 6000
N_D = []
N_Q = []

print("Starting population dynamics run for the doublet state...")
while D3.get_number_of_walkers() < N_max_D:
    D3.one_time_step()
    N_D.append(D3.get_number_of_walkers())

print("Starting population dynamics run for the quartet state...")
while Q3.get_number_of_walkers() < N_max_Q:
    Q3.one_time_step()
    N_Q.append(Q3.get_number_of_walkers())

data = [('1/2', N_D), ('3/2', N_Q)]
fig, axes = plt.subplots(nrows=1, ncols=2)
for d, ax in zip(data, axes):
    spin = d[0]
    y = d[1]
    x = np.arange(len(y))
    ax.plot(x, y, 'r')
    ax.set_xlabel('QMC step')
    ax.set_ylabel('Number of walkers')
    ax.set_title('Population dynamics (S = {0})'.format(spin))

fig.tight_layout()
```

Based on the results of the simulations with $S = const$ one can conclude that the critical number of walkers for the states with $S = 1/2$ and $S = 3/2$ are around 4000 and 8000, respectively. If we rerun quantum Monte-Carlo simulations with the walker ensembles of the corresponding sizes, the correct energies can be obtained.

```python
D3 = Walker_ensemble(random_walk=rw0101, potential=V_harm, n_walkers=5000, n_alpha=2, n_beta=1, S = 2.5)
Q3 = Walker_ensemble(random_walk=rw0101, potential=V_harm, n_walkers=8000, n_alpha=3, n_beta=0, S = 4.5)

E_D = np.zeros(N_prod)
N_D = np.zeros(N_prod)

E_Q = np.zeros(N_prod)
N_Q = np.zeros(N_prod)

for i in range(N_equil):
    D3.one_time_step()
    Q3.one_time_step()

print("Number of walkers at the end of equilibration run is {0} (doublet)".format(D3.get_number_of_walkers()))
print("Number of walkers at the end of equilibration run is {0} (quartet)".format(Q3.get_number_of_walkers()))
print("Starting production runs...")

for i in range(N_prod):
    D3.one_time_step()
    Q3.one_time_step()

    E_D[i] = D3.get_energy()
    N_D[i] = D3.get_number_of_walkers()
    E_Q[i] = Q3.get_energy()
    N_Q[i] = Q3.get_number_of_walkers()

print("Average energy at the end of the production run (S = 1/2): {0:13.6f}".format(np.mean(E_D)))
print("Average energy at the end of the production run (S = 3/2): {0:13.6f}".format(np.mean(E_Q)))
```

    Number of walkers at the end of equilibration run is 4945 (doublet)
    Number of walkers at the end of equilibration run is 7917 (quartet)
    Starting production runs...
    Average energy at the end of the production run (S = 1/2):      2.499397
    Average energy at the end of the production run (S = 3/2):      4.488633

```python
# Timing information
print("The run took {0} s for doublet".format(D3.time['total']))
print("The run took {0} s for quartet".format(Q3.time['total']))
```

    The run took 953.185147 s for doublet
    The run took 1433.397774 s for quartet

```python
print("Doublet (QMC): " + reblock(E_D))
print("Quartet (QMC): " + reblock(E_Q))

print("Doublet (Ref): {}".format(np.dot(doublet_occ, e5)))
print("Quartet (Ref): {}".format(np.dot(quartet_occ, e5)))
```

    Optimal block: [6]
    Doublet (QMC): < E > =   2.499397 +-   0.003835
    Optimal block: [6]
    Quartet (QMC): < E > =   4.488633 +-   0.004065
    Doublet (Ref): 2.49781034531
    Quartet (Ref): 4.49405367447

The doublet energy is in slightly better agreement with the exact result compared to the quartet energy. As shown below, this can be attributed to the timestep error.

```python
rw01005 = Random_walk(delta=0.1, tau=0.05)
Q3_1 = Walker_ensemble(random_walk=rw01005, potential=V_harm, n_walkers=8000, n_alpha=3, n_beta=0, S = 4.5)

E_Q_1 = np.zeros(N_prod)
N_Q_1 = np.zeros(N_prod)

for i in range(N_equil):
    Q3_1.one_time_step()

print("Number of walkers at the end of equilibration run is {0} (quartet)".format(Q3_1.get_number_of_walkers()))
print("Starting production run...")

for i in range(N_prod):
    Q3_1.one_time_step()
    E_Q_1[i] = Q3_1.get_energy()
    N_Q_1[i] = Q3_1.get_number_of_walkers()

print("Average energy at the end of the production run (S = 3/2): {0:13.6f}".format(np.mean(E_Q_1)))
```

    Number of walkers at the end of equilibration run is 8064 (quartet)
    Starting production run...
    Average energy at the end of the production run (S = 3/2):      4.491615

```python
print("Quartet (QMC), tau = 0.05: " + reblock(E_Q_1))
print("Quartet (Ref): {0:9.6f}".format(np.dot(quartet_occ, e5)))
```

    Optimal block: [5]
    Quartet (QMC), tau = 0.05: < E > =   4.491615 +-   0.003630
    Quartet (Ref):  4.494054


text:
# Stability and Stiffness

In the past when you've talked about **stability**, it has likely been regarding the stability of a *system*. Stable systems are those will well-behaved exact solutions, meaning they do not grow unbounded.
In engineering we mostly focus (or want!) stable systems, although there are some interesting unstable systems such as those involving resonance, nonlinear dynamics, or chaos—generally we want to know when that happens so we can prevent it.

We can also define the stability of a *numerical scheme*, which is when the numerical solution exhibits unphysical behavior. In other words, it blows up.

For example, let's consider the relatively simple 1st-order ODE
\begin{equation}
\frac{dy}{dt} = -3 y
\end{equation}
with the initial condition $y(0) = 1$. As we will see, this ODE can cause explicit numerical schemes to become unstable, and thus it is a **stiff** ODE. (Note that we can easily obtain the exact solution for this problem, which is $y(t) = e^{-3 t}$.)

Let's try solving this with the Forward Euler method, integrating over $0 \leq t \leq 10$, for a range of time-step size values: $\Delta t = 0.1, 0.25, 0.5, 0.75$:

```matlab
clear
f = @(t,y) -3*y;

dt = 0.1;
t = 0 : dt : 20;
y = zeros(length(t), 1);
y(1) = 1;
for i = 1 : length(t) - 1
    y(i+1) = y(i) + dt*f(t(i), y(i));
end
subplot(4,1,1);
plot(t, y); title(sprintf('dt = %4.2f', dt));

dt = 0.25;
t = 0 : dt : 20;
y = zeros(length(t), 1);
y(1) = 1;
for i = 1 : length(t) - 1
    y(i+1) = y(i) + dt*f(t(i), y(i));
end
subplot(4,1,2);
plot(t, y); title(sprintf('dt = %4.2f', dt));

dt = 0.5;
t = 0 : dt : 20;
y = zeros(length(t), 1);
y(1) = 1;
for i = 1 : length(t) - 1
    y(i+1) = y(i) + dt*f(t(i), y(i));
end
subplot(4,1,3);
plot(t, y); title(sprintf('dt = %4.2f', dt));

dt = 0.75;
t = 0 : dt : 20;
y = zeros(length(t), 1);
y(1) = 1;
for i = 1 : length(t) - 1
    y(i+1) = y(i) + dt*f(t(i), y(i));
end
subplot(4,1,4);
plot(t, y); title(sprintf('dt = %4.2f', dt));
```

At the smaller step sizes, $\Delta t = 0.1$ and $\Delta t = 0.25$, we see that the solution is well-behaved. But, when we increase $\Delta t$ to 0.5, we see some instability that goes away with time. Then, when we increase $\Delta t$ to 0.75, the solution eventually blows up, leading to error **much** larger than what we should expect based on the method's order of accuracy (first) and the step size value.

Compare this behavior to that for the ODE
\begin{equation}
\frac{dy}{dt} = e^{-t}
\end{equation}
which is **non-stiff**:

```matlab
clear
f = @(t,y) exp(-t);

dt = 0.1;
t = 0 : dt : 10;
y = zeros(length(t), 1);
y(1) = 1;
for i = 1 : length(t) - 1
    y(i+1) = y(i) + dt*f(t(i), y(i));
end
subplot(4,1,1);
plot(t, y); title(sprintf('dt = %4.2f', dt));

dt = 0.25;
t = 0 : dt : 10;
y = zeros(length(t), 1);
y(1) = 1;
for i = 1 : length(t) - 1
    y(i+1) = y(i) + dt*f(t(i), y(i));
end
subplot(4,1,2);
plot(t, y); title(sprintf('dt = %4.2f', dt));

dt = 0.5;
t = 0 : dt : 10;
y = zeros(length(t), 1);
y(1) = 1;
for i = 1 : length(t) - 1
    y(i+1) = y(i) + dt*f(t(i), y(i));
end
subplot(4,1,3);
plot(t, y); title(sprintf('dt = %4.2f', dt));

dt = 0.75;
t = 0 : dt : 10;
y = zeros(length(t), 1);
y(1) = 1;
for i = 1 : length(t) - 1
    y(i+1) = y(i) + dt*f(t(i), y(i));
end
subplot(4,1,4);
plot(t, y); title(sprintf('dt = %4.2f', dt));
```

In this case, we see that the solution remains well-behaved even for larger time-step sizes, and the error matches the expected order based on the method and step-size value.

In general numerical schemes can be:

- **unstable**: the scheme blows up for any choice of parameters
- **conditionally stable**: the scheme is stable for a particular choice of parameters (for example, $\Delta t$ is less than some threshold
- **unconditionally stable**: the scheme is always stable

Schemes may be stable for some problem/system and not for another, and vice versa.

Stability is related to robustness of a method, which is generally a tradeoff between complexity and computational cost. The choice of method and solution strategy depends on what you want, and how long you can wait for it. In general, we almost always want to use the largest $\Delta t$ allowable.

Rather than reducing $\Delta t$ to avoid stability problems, we can also use a method that is unconditionally stable, such as the **Backward Euler** method.

## Backward Euler method

The Backward Euler method is very similar to the Forward Euler method, except in one way: it uses the slope at the *next* time step:
\begin{equation}
    \left(\frac{dy}{dx}\right)_{i+1} \approx \frac{y_{i+1} - y_i}{\Delta x}
\end{equation}
Then, the resulting recursion formula is
\begin{equation}
y_{i+1} = y_i + \Delta x \left(\frac{dy}{dx}\right)_{i+1}, \text{ or} \\
y_{i+1} = y_i + \Delta x \, f(x_{i+1}, y_{i+1})
\end{equation}
where $f(x,y) = dy/dx$.

Notice that this recursion formula cannot be directly solved, because $y_{i+1}$ shows up on both sides. This is an **implicit** method, where all the other methods we have covered (Forward Euler, Heun's, Midpoint, and 4th-order Runge-Kutta) are **explicit**. Implicit methods require more work to actually implement.

### Backward Euler example
For example, consider the problem
\begin{equation}
\frac{dy}{dx} = f(x,y) = 8 e^{-x} (1+x) - 2y
\end{equation}
To actually solve this problem with the Backward Euler method, we need to incorporate the derivative function $f(x,y)$ into the recursion formula and solve for $y_{i+1}$:
\begin{align}
y_{i+1} &= y_i + \Delta x \, f(x_{i+1}, y_{i+1}) \\
y_{i+1} &= y_i + \Delta x \left[ 8 e^{-x_{i+1}} (1 + x_{i+1}) - 2 y_{i+1} \right] \\
y_{i+1} &= y_i + 8 e^{-x_{i+1}} (1 + x_{i+1}) \Delta x - 2 \Delta x \, y_{i+1} \\
y_{i+1} + 2 \Delta x \, y_{i+1} &= y_i + 8 e^{-x_{i+1}} (1 + x_{i+1}) \Delta x \\
y_{i+1} &= \frac{ y_i + 8 e^{-x_{i+1}} (1 + x_{i+1}) \Delta x }{ 1 + 2 \Delta x }
\end{align}
Now we have a useable recursion formula that we can use to solve this problem. Let's use the initial condition $y(0) = 1$, the domain $0 \leq x \leq 7$, and $\Delta x = 0.2$.

```matlab
clear

dx = 0.2;
x = 0 : dx : 7;
y = zeros(length(x), 1);
y(1) = 1;

% Backward Euler loop
for i = 1 : length(x) - 1
    y(i+1) = (y(i) + 8*exp(-x(i+1))*(1 + x(i+1))*dx) / (1 + 2*dx);
end

x_exact = linspace(0, 7);
y_exact = exp(-2.*x_exact).*(8*x_exact.*exp(x_exact) + 1);
plot(x_exact, y_exact); hold on
plot(x, y, 'o--')
legend('Exact solution', 'Backward Euler solution')
```

This matches nearly what we saw with the Forward Euler method before—Backward Euler is also a **first-order** method, so the global error should be proportional to $\Delta x$.

Let's now return to the stiff ODE $y^{\prime} = -3 y$, and see how the Backward Euler method does. First, we need to obtain our useable recursion formula:
\begin{align}
y_{i+1} &= y_i + \Delta t \, f(t_{i+1}, y_{i+1}) \\
y_{i+1} &= y_i + \Delta t \, \left( -3 y_{i+1} \right) \\
y_{i+1} + 3 y_{i+1} \Delta t &= y_i \\
y_{i+1} &= \frac{y_i}{1 + 3 \Delta t}
\end{align}

```matlab
clear

dt = 0.1;
t = 0 : dt : 10;
y = zeros(length(t), 1);
y(1) = 1;
for i = 1 : length(t) - 1
    y(i+1) = y(i) / (1 + 3*dt);
end
subplot(4,1,1);
plot(t, y); title(sprintf('dt = %4.2f', dt));

dt = 0.25;
t = 0 : dt : 10;
y = zeros(length(t), 1);
y(1) = 1;
for i = 1 : length(t) - 1
    y(i+1) = y(i) / (1 + 3*dt);
end
subplot(4,1,2);
plot(t, y); title(sprintf('dt = %4.2f', dt));

dt = 0.5;
t = 0 : dt : 10;
y = zeros(length(t), 1);
y(1) = 1;
for i = 1 : length(t) - 1
    y(i+1) = y(i) / (1 + 3*dt);
end
subplot(4,1,3);
plot(t, y); title(sprintf('dt = %4.2f', dt));

dt = 0.75;
t = 0 : dt : 10;
y = zeros(length(t), 1);
y(1) = 1;
for i = 1 : length(t) - 1
    y(i+1) = y(i) / (1 + 3*dt);
end
subplot(4,1,4);
plot(t, y); title(sprintf('dt = %4.2f', dt));
```

In this case, we see that the solution remains well-behaved for all the step sizes, not showing any of the instability we saw with the Forward Euler method. This is because the Backward Euler method is **unconditionally stable**.

## Stability analysis

### Stability analysis of Forward Euler

We can perform a stability analysis of the stiff problem to identify when the Forward Euler method becomes unstable. Let's apply the method to the ODE at hand:
\begin{align}
\frac{dy}{dt} &= -3 y \\
y_{i+1} &= y_i + \Delta t f(t_i, y_i) \\
y_{i+1} &= y_i + \Delta t (-3 y_i) \\
 &= y_i (1 - 3 \Delta t) \\
\frac{y_{i+1}}{y_i} &= \sigma = 1 - 3 \Delta t
\end{align}
where $\sigma$ is the **amplification factor**. This defines whether the solution grows or decays each step—for a stable physical system, we expect the solution to get smaller or remain contant with each step.

Therefore, for the method to remain stable, we must have $\sigma | \leq 1$. We can use this stability criterion to find conditions on $\Delta t$ for stability:
\begin{gather}
| \sigma | = | 1 - 3 \Delta t | \leq 1 \\
-1 \leq 1 - 3 \Delta t \leq 1 \\
-1 \leq 1 - 3 \Delta t \quad \text{or} \quad 1 - 3 \Delta t \leq 1 \\
\frac{-2}{3} \leq -\Delta t \quad \quad -\Delta t \leq 0 \\
\rightarrow \Delta t \leq \frac{2}{3} \quad \text{and} \quad \Delta t \geq 0 \\
\therefore 0 \leq \Delta t \leq \frac{2}{3}
\end{gather}
for stability. (For safety, we might use $\Delta t < 1/2$ for safety, to stay away from the absolute stability limit.)

The Forward Euler method is then *conditionally stable*.

As a general rule of thumb, all **explicit** methods are conditionally stable; these are methods where the recursion formula for $y_{i+1}$ can be written and calculated explicitly in terms of known quantities.

### Stability analysis of Backward Euler

We can also perform a stability analysis on the Backward Euler method to show that its stability does not depend on the step size:
\begin{align}
\frac{dy}{dt} &= -3 y \\
y_{i+1} &= y_i + \Delta t f(t_{i+1}, y_{i+1}) \\
y_{i+1} &= y_i + \Delta t (-3 y_{i+1}) \\
y_{i+1} (1 + 3 \Delta t) &= y_i \\
\sigma &= \frac{y_{i+1}}{y_i} = \frac{1}{1 + 3 \Delta t}
\end{align}
For stability, we need $| \sigma | \leq 1$:
\begin{align}
| \sigma | &= \left| \frac{1}{1 + 3 \Delta t} \right| \leq 1 \\
\rightarrow \Delta t &> 0
\end{align}
Therefore the Backward Euler method is *unconditionally stable*.

## Backward Euler for 2nd-order ODEs

We saw how to implement the Backward Euler method for a 1st-order ODE, but what about a 2nd-order ODE? (Or in general a system of 1st-order ODEs?)

The recursion formula is the same, except now our dependent variable is an array/vector:
\begin{equation}
\mathbf{y}_{i+1} = \mathbf{y}_i + \Delta t \, \mathbf{f} \left( t_{i+1} , \mathbf{y}_{i+1} \right)
\end{equation}
where the bolded $\mathbf{y}$ and $\mathbf{f}$ indicate array quantities (in other words, they hold more than one value).

In general, we can use Backward Euler to solve 2nd-order ODEs in a similar fashion as our other numerical methods:

1. Convert the 2nd-order ODE into a system of two 1st-order ODEs
2. Insert the ODEs into the Backward Euler recursion formula and solve for $\mathbf{y}_{i+1}$

The main difference is that we will now have a system of two equations and two unknowns: $y_{1, i+1}$ and $y_{2, i+1}$.

Let's demonstrate with an example:
\begin{equation}
y^{\prime\prime} + 6 y^{\prime} + 5y = 10 \quad y(0) = 0 \quad y^{\prime}(0) = 5
\end{equation}
where the exact solution is
\begin{equation}
y(t) = -\frac{3}{4} e^{-5t} - \frac{5}{4} e^{-t} + 2
\end{equation}

To solve numerically,

1. Convert the 2nd-order ODE into a system of two 1st-order ODEs:
\begin{gather}
y_1 = y \quad y_1(t=0) = 0 \\
y_2 = y^{\prime} \quad y_2 (t=0) = 5
\end{gather}
Then, for the derivatives of these variables:
\begin{align}
y_1^{\prime} &= y_2 \\
y_2^{\prime} &= 10 - 6 y_2 - 5 y_1
\end{align}

2. Then plug these derivatives into the Backward Euler recursion formulas and solve for $y_{1,i+1}$ and $y_{2,i+1}$:
\begin{align}
y_{1, i+1} &= y_{1, i} + \Delta t \, y_{2,i+1} \\
y_{2, i+1} &= y_{2, i} + \Delta t \left( 10 - 6 y_{2, i+1} - 5 y_{1,i+1} \right) \\
\\
y_{1, i+1} - \Delta t \, y_{2, i+1} &= y_{1,i} \\
5 \Delta t \, y_{1, i+1} + (1 + 6 \Delta t) y_{2, i+1} &= y_{2,i} + 10 \Delta t \\
\text{or} \quad
\begin{bmatrix} 1 & -\Delta t \\ 5 \Delta t & (1+6\Delta t)\end{bmatrix}
\begin{bmatrix} y_{1, i+1} \\ y_{2, i+1} \end{bmatrix} &=
\begin{bmatrix} y_{1,i} \\ y_{2,i} + 10 \Delta t \\ \end{bmatrix} \\
\mathbf{A} \mathbf{y}_{i+1} &= \mathbf{b}
\end{align}
To isolate $\mathbf{y}_{i+1}$ and get a usable recursion formula, we need to solve this system of two equations. We could solve this by hand using the substitution method, or we can use Cramer's rule:
\begin{align}
y_{1, i+1} &= \frac{ y_{1,i} (1 + 6 \Delta t) + \Delta t \left( y_{2,i} + 10 \Delta t \right)}{1 + 6 \Delta t + 5 \Delta t^2} \\
y_{2, i+1} &= \frac{ y_{2,i} + 10 \Delta t - 5 \Delta t y_{1,i}}{1 + 6 \Delta t + 5 \Delta t^2}
\end{align}

Let's confirm that this gives us a good, well-behaved numerical solution and compare with the Forward Euler method:

```matlab
clear

% Exact solution
t = linspace(0, 5);
y_exact = @(t) -(3/4)*exp(-5*t) - (5/4)*exp(-t) + 2;
plot(t, y_exact(t)); hold on

dt = 0.1;
t = 0 : dt : 5;

% Forward Euler
f = @(t, y1, y2) 10 - 6*y2 - 5*y1;
y1 = zeros(length(t), 1); y2 = zeros(length(t), 1);
y1(1) = 0; y2(1) = 5;
for i = 1 : length(t) - 1
    y1(i+1) = y1(i) + dt*y2(i);
    y2(i+1) = y2(i) + dt*f(t(i), y1(i), y2(i));
end
plot(t, y1, '+')

Y = zeros(length(t), 2);
Y(1,1) = 0;
Y(1,2) = 5;
for i = 1 : length(t) - 1
    D = 1 + 6*dt + 5*dt^2;
    Y(i+1, 1) = (Y(i,1)*(1 + 6*dt) + dt*(Y(i,2) + 10*dt)) / D;
    Y(i+1, 2) = (Y(i,2) + 10*dt - Y(i,1)*5*dt) / D;
end
plot(t, Y(:,1), 'o')
legend('Exact', 'Forward Euler', 'Backward Euler', 'location', 'southeast')

fprintf('Maximum error of Forward Euler: %5.3f\n', max(abs(y1(:) - y_exact(t)')));
fprintf('Maximum error of Backward Euler: %5.3f', max(abs(Y(:,1) - y_exact(t)')));
```

So, for $\Delta t = 0.1$, we see that the Forward and Backward Euler methods give an error $\mathcal{O}(\Delta t)$, as expected since both methods are *first-order* accurate.

Let's see how they compare for a larger step size:

```matlab
clear

% Exact solution
t = linspace(0, 5);
y_exact = @(t) -(3/4)*exp(-5*t) - (5/4)*exp(-t) + 2;
plot(t, y_exact(t)); hold on

dt = 0.5;
t = 0 : dt : 5;

% Forward Euler
f = @(t, y1, y2) 10 - 6*y2 - 5*y1;
y1 = zeros(length(t), 1); y2 = zeros(length(t), 1);
y1(1) = 0; y2(1) = 5;
for i = 1 : length(t) - 1
    y1(i+1) = y1(i) + dt*y2(i);
    y2(i+1) = y2(i) + dt*f(t(i), y1(i), y2(i));
end
plot(t, y1, 'o')

% Backward Euler

Y = zeros(length(t), 2);
Y(1,1) = 0;
Y(1,2) = 5;
for i = 1 : length(t) - 1
    D = 1 + 6*dt + 5*dt^2;
    Y(i+1, 1) = (Y(i,1)*(1 + 6*dt) + dt*(Y(i,2) + 10*dt)) / D;
    Y(i+1, 2) = (Y(i,2) + 10*dt - Y(i,1)*5*dt) / D;
end
plot(t, Y(:,1), 'o')
legend('Exact', 'Backward Euler', 'location', 'southeast')

%fprintf('Maximum error of Forward Euler: %5.3f\n', max(abs(y1(:) - y_exact(t)')));
%fprintf('Maximum error of Backward Euler: %5.3f', max(abs(Y(:,1) - y_exact(t)')));

fprintf('Maximum error of Forward Euler: %5.3f\n', max(abs(y1(:) - y_exact(t)')));
fprintf('Maximum error of Backward Euler: %5.3f', max(abs(Y(:,1) - y_exact(t)')));
```

Backward Euler, since it is unconditionally stable, remains well-behaved at this larger step size, while the Forward Euler method blows up.

One other thing: instead of using Cramer's rule to get expressions for $y_{1,i+1}$ and $y_{2,i+1}$, we could instead use Matlab to solve the linear system of equations at each time step. To do that, we could replace
```OCTAVE
A = [1  -dt; 5*dt  (1+6*dt)];
b = [Y(i,1); Y(i,2)+10*dt];
Y(i+1,:) = (A\b)';
```
where `A\b` is equivalent to `inv(A)*b`, but faster. Let's confirm that this gives the same answer:

```matlab
clear

% Exact solution
t = linspace(0, 5);
y_exact = @(t) -(3/4)*exp(-5*t) - (5/4)*exp(-t) + 2;
plot(t, y_exact(t)); hold on

dt = 0.1;
t = 0 : dt : 5;

Y = zeros(length(t), 2);
Y(1,1) = 0;
Y(1,2) = 5;
for i = 1 : length(t) - 1
    A = [1  -dt; 5*dt  (1+6*dt)];
    b = [Y(i,1); Y(i,2)+10*dt];
    Y(i+1,:) = (A\b)';
end
plot(t, Y(:,1), 'o')
legend('Exact', 'Backward Euler', 'location', 'southeast')
```

## Cramer's Rule

Cramer's Rule provides a solution method for a system of linear equations, where the number of equations equals the number of unknowns. It works for any number of equations/unknowns, but isn't really practical for more than two or three. We'll focus on using it for a system of two equations, with two unknowns $x_1$ and $x_2$:
\begin{gather}
a_{11} + x_1 + a_{12} x_2 = b_1 \\
a_{21} + x_1 + a_{22} x_2 = b_2 \\
\text{or } \mathbf{A} \mathbf{x} = \mathbf{b}
\end{gather}
where
\begin{gather}
\mathbf{A} = \begin{bmatrix} a_{11} & a_{12} \\ a_{21} & a_{22} \end{bmatrix} \\
\mathbf{x} = \begin{bmatrix} x_1 \\ x_2 \end{bmatrix} \\
\mathbf{b} = \begin{bmatrix} b_1 \\ b_2 \end{bmatrix}
\end{gather}

The solutions for the unknowns are then
\begin{align}
x_1 &= \frac{ \begin{vmatrix} b_1 & a_{12} \\ b_2 & a_{22} \end{vmatrix} }{D} = \frac{b_1 a_{22} - a_{12} b_2}{D} \\
x_2 &= \frac{ \begin{vmatrix} a_{11} & b_1 \\ a_{21} & b_2 \end{vmatrix} }{D} = \frac{a_{11} b_2 - b_1 a_{21}}{D}
\end{align}
where $D$ is the determinant of $\mathbf{A}$:
\begin{equation}
D = \det(\mathbf{A}) = \begin{vmatrix} a_{11} & a_{12} \\ a_{21} & a_{22} \end{vmatrix} = a_{11} a_{22} - a_{12} a_{21}
\end{equation}
This works as long as the determinant does not equal zero.


text:
```python
# standard
import os
from datetime import datetime

# third party
import covidcast
import cvxpy as cp
import matplotlib.pyplot as plt
import numpy as np
import numpy.linalg as la
import pandas as pd
import scipy as sp
from sklearn.model_selection import LeaveOneOut

# first party
os.chdir("/Users/mjahja/Delphi/covidcast-nowcast/cases/deconv/code")
from weekday import Weekday
```

# Adjusting for day-of-week effects

As a preprocessing step (could be absorbed under the "sensorization" step, though it's different in spirit), we adjust each counties' case curve using a Poisson regression model. This is the same model used for adjusting the HSP signals, which was done by Aaron.

Let $C_{it}$ be the observed cases for location $i$ at time $t$. We assume

\begin{align}
\mathbb{E}[C_{it}] &= \mu_t\\
\log \mu_t &= \alpha_{wd(t)} + \phi_t
\end{align}

where $wd(t) \in \{0, ..., 6\}$ is the weekday corresponding to time $t$, and $\alpha_{wd(t)}$ is the correction for the weekday, and $\phi_t$ is the corrected case value. Notice we are fitting a single multiplier for each weekday, across all locations.

Since we will refit this model for each sequential timepoint, we assume that the weekday parameters do not change over time. We minimize the objective

\begin{align}
f(\alpha, \phi | \mu) &= - \log \ell(\alpha, \phi | \mu) + \lambda \lVert \Delta^3 \phi \rVert_1\\
\text{s.t.}\,\, \sum_{i=0}^6 \alpha_{wd(i)} &= 0 \Leftrightarrow \alpha_{wd(0)} = - \sum_{i=1}^6 \alpha_{wd(i)}
\end{align}

where $\ell$ represents the Poisson likelihood and $\Delta^3$ is the third-difference operator applied to $\phi$. The estimated counts are given by

\begin{equation}
\tilde C_{it} = C_{it} / \hat\alpha_{wd(t)}.
\end{equation}

## Fitting
Rewrite the problem as
\begin{equation}
\log \mu_t = X\beta
\end{equation}
where the first 6 elements of $\beta \in \mathbb{R}^{6+T}$ correspond to $\alpha$, and the remaining $T$ elements (corresponding to all the training timepoints) correspond to $\phi$. The design matrix $X\in\mathbb{R}^{T\times (6+T)}$ has the form
\begin{equation}
X =
\begin{bmatrix}
\vdots & & \vdots \\
\mathbf{1}_{i\ mod\ 7 = 1} & \cdots & \mathbf{1}_{i\ mod\ 7 = 6} & I_{T} \\
\vdots & &\vdots
\end{bmatrix}.
\end{equation}
In words, the rows of $X$ correspond to each timepoint, the first six columns are indicators for day-of-week, and the final $T$ columns are an identity matrix corresponding to unique $\phi_t$.

```python
# get all florida fips codes
geo_map = pd.read_csv(
            "~/Delphi/covid-19/geographical_scope/02_20_uszips.csv",
            usecols=["fips", "state_id", "population"],
            dtype={"state_id": str},
            converters={"fips": lambda x: str(x).zfill(5)},
        )

geo_map = geo_map.groupby(['fips', 'state_id']).sum().reset_index()
florida_geo = geo_map[geo_map.state_id.eq("FL")]
florida_population = florida_geo.groupby("fips").population.sum().reset_index()
florida_fips = florida_geo.fips.unique()
```

```python
# training dates
start_date = datetime(2020, 5, 15)
end_date = datetime(2020, 10, 15)

smoothed_cases_df = covidcast.signal(
        'indicator-combination',
        'confirmed_7dav_incidence_num',
        start_date, end_date,
        geo_type='county')

raw_cases_df = covidcast.signal(
        'indicator-combination',
        'confirmed_incidence_num',
        start_date, end_date,
        geo_type='county')
```

    /usr/local/lib/python3.7/site-packages/covidcast/covidcast.py:321: UserWarning: Problem obtaining data on 20201012: no results
      message=day_data["message"]))
    /usr/local/lib/python3.7/site-packages/covidcast/covidcast.py:321: UserWarning: Problem obtaining data on 20201013: no results
      message=day_data["message"]))
    /usr/local/lib/python3.7/site-packages/covidcast/covidcast.py:321: UserWarning: Problem obtaining data on 20201014: no results
      message=day_data["message"]))
    /usr/local/lib/python3.7/site-packages/covidcast/covidcast.py:321: UserWarning: Problem obtaining data on 20201015: no results
      message=day_data["message"]))

```python
raw_fl_cases = raw_cases_df[raw_cases_df.geo_value.isin(florida_fips)]
raw_fl_cases = raw_fl_cases.merge(florida_population, how="left", left_on="geo_value", right_on="fips")

raw_cases_df = raw_cases_df.groupby(['geo_value', 'time_value']).sum().reset_index()
raw_cases_df = raw_cases_df.merge(geo_map, how="left", left_on="geo_value", right_on="fips").dropna()

state_locs = raw_cases_df.state_id.unique()
state_cases_df = raw_cases_df.groupby(['state_id', 'time_value']).sum()
```

```python
# weekday adjustment on US states
lam = 50
for loc in ["PA", "FL"]:
    loc_cases = state_cases_df.loc[loc].groupby("time_value").sum().value
    state_wd_params = Weekday.get_params(loc_cases, lam)
    adj_value = np.exp(state_wd_params[6:])

    dates = loc_cases.index
    plt.plot(dates, loc_cases.values, label="raw", color="black", alpha=0.88)
    plt.plot(dates, loc_cases.rolling(7).mean(), label="7-day roll-avg")
    plt.plot(dates, adj_value, label="dow-adj")
    plt.xticks(rotation=45, size="small")
    plt.title(loc)
    plt.legend()
    plt.show()
```

```python
# weekday adjustment on florida counties

# sample a few counties
np.random.seed(2020)
locs = np.random.choice(raw_fl_cases.geo_value.unique(), 3)

# search over grid of lambdas
lam_grid = [1, 10, 25, 75, 100]

for loc in locs:
    loc_df = raw_fl_cases[raw_fl_cases.geo_value.eq(loc)]
    case_curve = loc_df.value.values.reshape(-1, 1)
    N = case_curve.size

    loo = LeaveOneOut()
    lam_scores = []
    for lam in lam_grid:
        score = []
        for train_i, test_i in loo.split(case_curve):
            test_i = test_i[0]
            if 1 < test_i < (N-2):
                raw_vals = loc_df.iloc[train_i].set_index("time_value")
                loo_params = Weekday.get_params(raw_vals.value, lam)
                fit = np.exp(loo_params[6:])
                preds_test = (fit[test_i-1] + fit[test_i+1])/2
                score.append(float(preds_test) - float(case_curve[test_i]))

        lam_scores.append(np.mean(np.square(score)))

    best_lam = lam_grid[np.argmin(lam_scores)]
    fips_wd_params = Weekday.get_params(loc_df.set_index("time_value").value, best_lam)
    adj_value = np.exp(fips_wd_params[6:])

    dates = loc_df.time_value
    plt.plot(dates, loc_df.value, label="raw", color="black", alpha=0.88)
    plt.plot(dates, loc_df.value.rolling(7).mean(), label="7-day roll-avg")
    plt.plot(dates, adj_value, label="dow-adj")
    plt.xticks(rotation=45, size="small")
    plt.title(loc)
    plt.legend()
    plt.show()
```

```python

```