## Testing look-elsewhere effect by creating 2d chi-square random fields with a Gaussian Process

by Kyle Cranmer, Dec 7, 2015

The correction for 2d look-elsewhere effect presented in
*Estimating the significance of a signal in a multi-dimensional search* by  Ofer Vitells and Eilam Gross http://arxiv.org/pdf/1105.4355v1.pdf

is based on the fact that the test statistic

\begin{equation}
q(\nu_1, \nu_2) = -2 \log \frac{ \max_{\theta} L(\mu=0, \nu_1, \nu_2, \theta)}{ \max_{\mu, \theta} L(\mu, \nu_1, \nu_2, \theta)}
\end{equation}

is a chi-square random field (with 1 degree of freedom). That means that, for any point in $\nu_1, \nu_2$, the quantity $q(\nu_1, \nu_2)$ would have a chi-square distribution if you repeated the experiment many times.

That is what you expect if you have a background model $p_b(x|\theta)$ and you look for a signal on top of it with signal strength $\mu$. Creating that scan is somewhat time consuming, so here we make realizations of a chi-square random field by using a Gaussian Process.
The main trick we will use is that a chi-square distribution for one degree of freedom is the same as the distribution of $x^2$ if $x$ is normally distributed. As you might have guessed, a Gaussian Process (GP) is like a chi-square random field, but it is Gaussian-distributed at each point.

Note, the distributions are not independent at each point, there is some covaraince. So if the $q(\nu_1, \nu_2)$ is high at one point, you can expect it to be high near by. We can control this behavior via the GP's kernel.

For more on the theory of Gaussian Processes, the best resource is available for free online: [Rasmussen & Williams (2006)](http://www.gaussianprocess.org/gpml/). We will [`george`](http://dan.iel.fm/george/current/) -- a nice python package for Gaussian Processes (GP).

```python
%pylab inline --no-import-all
```

    Populating the interactive namespace from numpy and matplotlib

    //anaconda/lib/python2.7/site-packages/matplotlib/font_manager.py:273: UserWarning: Matplotlib is building the font cache using fc-list. This may take a moment.
      warnings.warn('Matplotlib is building the font cache using fc-list. This may take a moment.')

The main trick we will use is that a chi-square distribution for one degree of freedom is the same as the distribution of $x^2$ if $x$ is normally distributed. Here's a quick demonstration of that:

```python
from scipy.stats import chi2, norm
chi2_array = chi2.rvs(1, size=10000)
norm_array = norm.rvs(size=10000)
_ = plt.hist(chi2_array, bins=100, alpha=.5, label='chi-square')
_ = plt.hist(norm_array**2, bins=100, alpha=.5, color='r', label='x^2')
plt.yscale('log', nonposy='clip')
plt.legend(('chi-square', 'x^2'))
#plt.semilogy()
```

### Ok, now to the Gaussian processes.

```python
import george
from george.kernels import ExpSquaredKernel
```

```python
length_scale_of_correaltion=0.1
kernel = ExpSquaredKernel(length_scale_of_correaltion, ndim=2)
```

```python
# Create the Gaussian process
# gp = george.GP(kernel)
gp = george.GP(kernel, solver=george.HODLRSolver) #faster
```

```python
n_scan_points=50
aspect_ratio =  10. # make excesses look like stripes
x_scan = np.arange(0,aspect_ratio,aspect_ratio/n_scan_points)
y_scan = np.arange(0,1,1./n_scan_points)
xx, yy = np.meshgrid(x_scan, y_scan)
```

```python
# reformat the independent coordinates where we evaluate the GP
indep = np.vstack((np.hstack(xx),np.hstack(yy))).T
```

```python
# illustration of what is being done here
np.vstack([[1,2],[3,4]]).T
```

    array([[1, 3],
           [2, 4]])

```python
# slow part: pre-compute internal stuff for the GP
gp.compute(indep)
```

```python
# evaluate one realization of the GP
z = gp.sample(indep)
```

```python
# reformat output for plotting
zz = z.reshape((n_scan_points,n_scan_points))
```

```python
# plot the chi-square random field
plt.imshow(zz**2, cmap='gray')
plt.colorbar()
```

Now lets histogram the values of the random field.
Don't get confused here... if you pick a single point and histogram the value of over many instances, you expect a Gaussian. However, for a single instance, you don't expect the histogram for the value of the field to be Gaussian (because of the correlations). Thought experiments: if you make `length_scale_of_correaltion` very small, then each point is essentially independent and you do expect to see a Gaussian; however, if `length_scale_of_correaltion` is very large then you expect the field to be nearly constant and the histogram below would be a delta function.

```python
# plot the gaussian distributed x and chi-square distributed x**2
plt.subplot(1,2,1)
count, edges, patches = plt.hist(np.hstack(zz), bins=100)
plt.xlabel('z')
plt.subplot(1,2,2)
count, edges, patches = plt.hist(np.hstack(zz)**2, bins=100)
plt.xlabel('q=z**2')
plt.yscale('log', nonposy='clip')
```

## Ok, now let's repeat that several times and test lee2d

```python
from lee2d import *
```

```python
from scipy.ndimage import grey_closing, binary_closing

def fill_holes(array):
    zero_array = array==0.
    temp = grey_closing(array, size=2)*zero_array
    return temp+array
```

Generate 25 realizations of the GP, calculate the Euler characteristic for two thresholds, and use the mean of those Euler characteristics to estimate $N_1$ and $N_2$

```python
n_samples = 100
z_array = gp.sample(indep,n_samples)

q_max = np.zeros(n_samples)
phis = np.zeros((n_samples,2))

u1,u2 = 0.5, 1.

n_plots = 3
plt.figure(figsize=(9,n_plots*3))
for scan_no, z in enumerate(z_array):

    scan = z.reshape((n_scan_points,n_scan_points))**2

    q_max[scan_no] = np.max(scan)

    # fill holes from failures in original likelihood
    scan = fill_holes(scan)

    #get excursion sets above those two levels
    exc1 = (scan>u1) + 0. #add 0. to convert from bool to double
    exc2 = (scan>u2) + 0.
    #print '\nu1,u2 = ', u1, u2

    #print 'diff = ', np.sum(exc1), np.sum(exc2)

    if scan_no < n_plots:
        aspect = 1.
        plt.subplot(n_plots,3,3*scan_no+1)
        aspect = 1.*scan.shape[0]/scan.shape[1]
        plt.imshow(scan.T, cmap='gray', aspect=aspect)
        plt.subplot(n_plots,3,3*scan_no+2)
        plt.imshow(exc1.T, cmap='gray', aspect=aspect, interpolation='none')
        plt.subplot(n_plots,3,3*scan_no+3)
        plt.imshow(exc2.T, cmap='gray', aspect=aspect, interpolation='none')

    phi1 = calculate_euler_characteristic(exc1)
    phi2 = calculate_euler_characteristic(exc2)
    #print 'phi1, phi2 = ', phi1, phi2
    #print 'q_max = ', np.max(scan)

    phis[scan_no] = [phi1, phi2]

plt.savefig('chi-square-random-fields.png')
```

```python
exp_phi_1, exp_phi_2 = np.mean(phis[:,0]), np.mean(phis[:,1])
exp_phi_1, exp_phi_2
```

    (14.359999999999999, 13.76)

```python
n1, n2 = get_coefficients(u1=u1, u2=u2, exp_phi_1=exp_phi_1, exp_phi_2=exp_phi_2)
print n1, n2
```

    7.34442273108 14.81882537

With estimates of $N_1$ and $N_2$ predict the global p-value vs. u

```python
u = np.linspace(5,25,100)
global_p = global_pvalue(u,n1,n2)
```

### Generate 5000 instances of the Gaussian Process, find maximum local significance for each, and check the prediction for the LEE-corrected global p-value

```python
n_samples = 5000
z_array = gp.sample(indep,n_samples)

q_max = np.zeros(n_samples)

for scan_no, z in enumerate(z_array):
    scan = z.reshape((n_scan_points,n_scan_points))**2
    q_max[scan_no] = np.max(scan)
```

```python
bins, edges, patches = plt.hist(q_max, bins=30)
icdf = 1.-np.cumsum(bins/n_samples)
icdf = np.hstack((1.,icdf))
icdf_error = np.sqrt(np.cumsum(bins))/n_samples
icdf_error = np.hstack((0.,icdf_error))
plt.xlabel('q_max')
plt.ylabel('counts / bin')
```

```python
# plot the p-value
plt.subplot(121)
plt.plot(edges,icdf, c='r')
plt.errorbar(edges,icdf,yerr=icdf_error)
plt.plot(u, global_p)
plt.xlabel('u')
plt.ylabel('P(q_max >u)')
plt.xlim(0,25)
plt.subplot(122)
plt.plot(edges,icdf, c='r', label='toys')
plt.errorbar(edges,icdf,yerr=icdf_error)
plt.plot(u, global_p, label='prediction')
plt.xlabel('u')
plt.legend(('toys', 'prediction'))
#plt.ylabel('P(q>u)')
plt.ylim(1E-3,10)
plt.xlim(0,25)
plt.semilogy()
```

### Study statistical uncertainty

Outline:
 1. generate `n_samples` likelihood scans using the GP
 1. make exclusion sets, calculate phi1, phi2 for levels u1, u2
 1. look at histogram of phi1, phi2 (notice that they are narrower than Poisson)
 1. look at 2-d scatter of phi1, phi2 (notice that they are positively correlated)
 1. look at 2-d scatter of coefficients n1, n2 (notice tha they are negatively correlated)
 1. Compare three ways of propagating error to global p-value
     1. Poisson, no correlations: estimate uncertainty on Exp[phi1] as sqrt(exp_phi_1)/sqrt(n_samples)
     1. Gaus approx of observed, no correlations: estimate uncertainty on Exp[phi1] as std(exp_phi_1)/sqrt(n_samples)
     1. Gaus approx of observed, with correlations: estimate covariance of (Exp[phi1], Exp[phi2]) with cov(phi1, phi2)/n_samples  -- note since it's covariance we divide by n_samples not sqrt(n_samples)

Conclusions:

The number of islands (as quantified by the Euler characteristic) is not Poisson distributed.
Deviation from the Poisson distribution will depend on the properties of the underlying 2-d fit (equivalently, the Gaussian Process kernel). In this example, the deviation isn't that big. It is probably generic that the uncertainty in phi is smaller than Poisson because one can only fit in so many islands into the scan... so it's probably more like a Binomial.

Unsurpringly there is also a positive correlation between the number of islands at levels u1 and u2.
This turns into an anti-correlation on the coefficients n1 and n2.

The two effects lead to the Poisson approximation over estimating the uncertainty on the global p-value.

```python
from scipy.stats import poisson
```

```python
n_samples = 1000
z_array = gp.sample(indep,n_samples)
phis = np.zeros((n_samples,2))

for scan_no, z in enumerate(z_array):
    scan = z.reshape((n_scan_points,n_scan_points))**2

    #get excursion sets above those two levels
    exc1 = (scan>u1) + 0. #add 0. to convert from bool to double
    exc2 = (scan>u2) + 0.

    phi1 = calculate_euler_characteristic(exc1)
    phi2 = calculate_euler_characteristic(exc2)
    phis[scan_no] = [phi1, phi2]
```

```python
bins = np.arange(0,25)
counts, bins, patches = plt.hist(phis[:,0], bins=bins, normed=True, alpha=.3, color='b')
_ = plt.hist(phis[:,1], bins=bins, normed=True,alpha=.3, color='r')
plt.plot(bins,poisson.pmf(bins,np.mean(phis[:,0])), c='b')
plt.plot(bins,poisson.pmf(bins,np.mean(phis[:,1])), c='r')
plt.xlabel('phi_i')
plt.legend(('obs phi1', 'obs phi2', 'poisson(mean(phi1)', 'poisson(mean(phi2))'), loc='upper left')
```

```python
print 'Check Poisson phi1', np.mean(phis[:,0]), np.std(phis[:,0]), np.sqrt(np.mean(phis[:,0]))
print 'Check Poisson phi1', np.mean(phis[:,1]), np.std(phis[:,1]), np.sqrt(np.mean(phis[:,1]))

print 'correlation coefficients:'
print np.corrcoef(phis[:,0], phis[:,1])
print 'covariance:'
print np.cov(phis[:,0], phis[:,1])
```

    Check Poisson phi1 14.628 2.87256261899 3.82465684735
    Check Poisson phi1 13.88 2.38528824254 3.72558720204
    correlation coefficients:
    [[ 1.          0.43701228]
     [ 0.43701228  1.        ]]
    covariance:
    [[ 8.25987588  2.99735736]
     [ 2.99735736  5.6952953 ]]

```python
x, y = np.random.multivariate_normal([np.mean(phis[:,0]),np.mean(phis[:,0])], np.cov(phis[:,0], phis[:,1]), 5000).T
_ = plt.scatter(phis[:,0], phis[:,1], alpha=0.1)
plt.plot(x, y, 'x', alpha=0.1)
plt.axis('equal')
plt.xlabel('phi_0')
plt.ylabel('phi_1')
```

```python
toy_n1, toy_n2 = np.zeros(x.size),np.zeros(x.size)
for i, (toy_exp_phi_1, toy_exp_phi_2) in enumerate(zip(x,y)):
    n1, n2 = get_coefficients(u1=u1, u2=u2, exp_phi_1=toy_exp_phi_1, exp_phi_2=toy_exp_phi_2)
    toy_n1[i] = n1
    toy_n2[i] = n2
plt.scatter(toy_n1, toy_n2, alpha=.1)
plt.xlabel('n1')
plt.ylabel('n2')
```

```python
# now propagate error exp_phi_1 and exp_phi_2 (by dividing cov matrix by n_samples) including correlations
x, y = np.random.multivariate_normal([np.mean(phis[:,0]),np.mean(phis[:,1])],
                                     np.cov(phis[:,0], phis[:,1])/n_samples,
                                     5000).T

'''
# check consistency with next cell by using diagonal covariance
dummy_cov = np.cov(phis[:,0], phis[:,1])/n_samples
dummy_cov[0,1]=0
dummy_cov[1,0]=0
print dummy_cov
x, y = np.random.multivariate_normal([np.mean(phis[:,0]),np.mean(phis[:,1])],
                                     dummy_cov,
                                     5000).T
'''

toy_global_p = np.zeros(x.size)
for i, (toy_exp_phi_1, toy_exp_phi_2) in enumerate(zip(x,y)):
    n1, n2 = get_coefficients(u1=u1, u2=u2, exp_phi_1=toy_exp_phi_1, exp_phi_2=toy_exp_phi_2)
    u = 16
    #global_p = global_pvalue(u,n1,n2)
    toy_global_p[i] = global_pvalue(u,n1,n2)
```

```python
# now propagate error assuming uncorrelated but observed std. on phi_1 and phi_2 / sqrt(n_samples)
x = np.random.normal(np.mean(phis[:,0]), np.std(phis[:,0])/np.sqrt(n_samples), 5000)
y = np.random.normal(np.mean(phis[:,1]), np.std(phis[:,1])/np.sqrt(n_samples), 5000)

toy_global_p_uncor = np.zeros(x.size)
for i, (toy_exp_phi_1, toy_exp_phi_2) in enumerate(zip(x,y)):
    n1, n2 = get_coefficients(u1=u1, u2=u2, exp_phi_1=toy_exp_phi_1, exp_phi_2=toy_exp_phi_2)
    u = 16
    #global_p = global_pvalue(u,n1,n2)
    toy_global_p_uncor[i] = global_pvalue(u,n1,n2)
```

```python
# now propagate error assuming uncorrelated Poisson stats on phi_1 and phi_2
x = np.random.normal(np.mean(phis[:,0]), np.sqrt(np.mean(phis[:,0]))/np.sqrt(n_samples), 5000)
y = np.random.normal(np.mean(phis[:,1]), np.sqrt(np.mean(phis[:,1]))/np.sqrt(n_samples), 5000)

toy_global_p_uncor_pois = np.zeros(x.size)
for i, (toy_exp_phi_1, toy_exp_phi_2) in enumerate(zip(x,y)):
    n1, n2 = get_coefficients(u1=u1, u2=u2, exp_phi_1=toy_exp_phi_1, exp_phi_2=toy_exp_phi_2)
    u = 16
    #global_p = global_pvalue(u,n1,n2)
    toy_global_p_uncor_pois[i] = global_pvalue(u,n1,n2)
```

```python
counts, bins, patches = plt.hist(toy_global_p_uncor_pois, bins=50, normed=True, color='g', alpha=.3)
counts, bins, patches = plt.hist(toy_global_p_uncor, bins=bins, normed=True, color='r', alpha=.3)
counts, bins, patches = plt.hist(toy_global_p, bins=bins, normed=True, color='b', alpha=.3)

plt.xlabel('global p-value')
#plt.ylim(0,1.4*np.max(counts))
plt.legend(('uncorrelated Poisson approx from mean',
            'uncorrelated Gaus. approx of observed dist',
            'correlated Gaus. approx of observed dist'),
           bbox_to_anchor=(1., 1.3))
```

Conclusion: The two effects lead to the Poisson approximation over estimating the uncertainty on the global p-value.


```python
# Make SymPy available to this program:
import sympy
from sympy import *

# Make GAlgebra available to this program:
from galgebra.ga import *
from galgebra.mv import *
from galgebra.printer import Fmt, GaPrinter, Format
    # Fmt:       sets the way that a multivector's basis expansion is output.
    # GaPrinter: makes GA output a little more readable.
    # Format:    turns on latex printer.
from galgebra.gprinter import gFormat, gprint
gFormat()
```

```python
# Set up standard G^3 geometric algebra
g3coords = (x,y,z) = symbols('x y z', real=True) # Without real=True, symbols are complex
g3 = Ga('\mathbf{e}', g=[1,1,1], coords=g3coords)
(ex, ey, ez) = g3.mv()     # Program names of basis vectors.
(exr, eyr, ezr) = g3.mvr() # Program names of reciprocal basis vectors.
```

```python
gprint(r'word word\ word \cdot ex<ey', r'\\ \text{word word\ word \cdot ex<ey)}')
    # \\ gives a new line. The second string encloses the first in \text{}.
```

```python
B = g3.mv('B', 'bivector')
Fmt(1)  # Set Fmt globally
gprint(r'\mathbf{B} =', B)         # B will be bold.
gprint(r'\mathbf{B} =', B.Fmt(3))  # Fmt(3) here only.
gprint(r'\mathbf{B} =', B)         # Global Fmt remembered.
```

```python
gprint(r'\mathbf{B}^2 =', B*B)
```

```python
M = g3.mv('M', 'mv')
gprint(r'\langle \mathbf{M} \rangle_2 =', M.grade(2))
   # grade(2) could be replaced by, e.g., odd(), or omitted altogether.
```

```python
gprint(r'\alpha_1\mathbf{X}/\gamma_r^3')
```

```python
# Program name and output are different
theta = symbols('theta', real = True)
th = symbols('theta', real = True)  # This will save typing if theta is used a lot.
gprint(theta, ', ', th)
```

```python
grad = g3.grad
```

```python
grad
```

```python
gprint(r'{\nabla}')
```

```python

```


Traçar um esboço do gráfico e obter uma equação da parábola que satisfaça as condições dadas.

13. Vértice: $V(0,0)$; Foco: $F(0,-3)$

Se a distância do vértice ao o foco é igual a distância do vértice até a diretriz, temos que $||F|| = ||D||$, logo $3 = d$

Equação da diretriz
$D : y = 3$

Descobrindo o valor de $p$

$F = -3$

$-3 = \frac{p}{2}$

$-6 = p$

Montando a equação

$x^2 = 2py$

$x^2 = 2\cdot -6 \cdot x$

$x^2 = -12y$

Gráfico da função

```python
from sympy import *
from sympy.plotting import plot_implicit
x, y = symbols("x y")
plot_implicit(Eq((x-0)**2, -12*(y+0)), (x,-10,10), (y,-10,10),
title=u'Gráfico da parábola', xlabel='x', ylabel='y');
```

Problemas Propostos. Pag. 172 - 175/13.ipynb


# The Helmholtz equation on the Torus

Helmholtz equation is given as

$$
\nabla^2 u + \alpha u = f.
$$

In this notebook we will solve this equation (without boundary conditions) on the surface of a torus, using curvilinear coordinates. The surface of the torus is parametrized by

$$
\begin{align}
x(\theta, \phi) &= (R + r \cos \theta) \cos \phi \\
y(\theta, \phi) &= (R + r \cos \theta) \sin \phi \\
z(\theta, \phi) &= r \sin \theta
\end{align}
$$

where $\theta, \phi$ are angles which make a full circle, so that their values start and end at the same point, $R$ is the distance from the center of the tube to the center of the torus,
$r$ is the radius of the tube. Note that $\theta$ is the angle in the small circle (around its center), whereas $\phi$ is the angle of the large circle, around origo.

We start the implementation by importing necessary functionality from shenfun and sympy and then defining the coordinates of the surface of the torus. Note that `rv` represents the position vector spanning the surface.

```python
from shenfun import *
from shenfun.la import SolverGeneric1ND
import sympy as sp
from IPython.display import Math

N = 96
R = 3
r = 1
theta, phi = psi = sp.symbols('x,y', real=True, positive=True)
rv = ((R + r*sp.cos(theta))*sp.cos(phi), (R + r*sp.cos(theta))*sp.sin(phi), r*sp.sin(theta))
```

Now create necessary bases and function spaces. Due to the geometry of the problem, the solution will be periodic in both $\theta$ and $\phi$ directions, and we choose Fourier basis functions. The basis for the $\phi$-direction can be either real to complex or complex to complex, depending on the type of the solution. Here we assume a real solution

```python
B1 = FunctionSpace(N, 'F', dtype='D', domain=(0, 2*np.pi))
B2 = FunctionSpace(N, 'F', dtype='d', domain=(0, 2*np.pi))
T = TensorProductSpace(comm, (B1, B2), coordinates=(psi, rv, sp.Q.positive(r*sp.cos(theta)+R)))
V = VectorSpace(T)
u = TrialFunction(T)
v = TestFunction(T)
T.coors.sg
```

Note the assumption `sp.Q.positive(r*sp.cos(theta)+R))`, which is there to help sympy when computing basis vectors and scaling factors. It is not completely necessary, but try to omit it and look at what happens to the expanded Helmholtz equation below.

```python
alpha = 0
du = div(grad(u))+alpha*u
g = sp.Symbol('g', real=True, positive=True) # The Jacobian**2 (T.coors.sg**2)
#replace = [(3*sp.cos(theta)+5, g/2), (sp.cos(theta)**2+6*sp.cos(theta)+10, sp.cos(theta)**2+g), (3*sp.cos(theta)+10, g/2+5)] # to simplify the look
replace = [(sp.cos(theta)+3, g)]
Math(du.tolatex(symbol_names={r: 'r', theta: '\\theta', phi: '\\phi'}, replace=replace))
```

We now create a manufactured solution that satisfies periodicity and compute the right hand side $f$.

```python
ue = sp.sin(sp.cos(theta*4))*sp.cos(4*phi)
f = (div(grad(u))+alpha*u).tosympy(basis=ue, psi=psi)
fj = Array(T, buffer=f*T.sg)
f_hat = Function(T)
f_hat = inner(v, fj, output_array=f_hat)
```

Assemble coefficient matrix and solve problem. Note that the tensorproduct matrices along axis 0 can be non-diagonal due to the measure $\cos \theta + 3$. The matrices along the second axis will all be diagonal, so we can choose to use `SolverGeneric1ND`

```python
mats = inner(v*T.sg, (div(grad(u))+alpha*u))
#mats = inner(grad(v*T.sg), -grad(u)) # + inner(v, alpha*u)
u_hat = Function(T)
sol = SolverGeneric1ND(mats)
u_hat = sol(f_hat, u_hat, constraints=((0, 0),))
uj = u_hat.backward()
uq = Array(T, buffer=ue)
print('Error =', np.sqrt(inner(1, (uj-uq)**2)))
```

```python
mats = inner(v*T.sg, u)
```

Finally, just plot the solution using plotly.

```python
surf3D(u_hat, wrapaxes=[0, 1])
```

Inspect sparsity pattern at a given Fourier wavenumber

```python
import matplotlib.pyplot as plt
%matplotlib notebook
plt.figure(figsize=(6,4))
plt.spy(sol.MM[6].mat, markersize=0.1)
```

Inspect the vector Laplacian

```python
uv = TrialFunction(V)
vv = TestFunction(V)
du = div(grad(uv))
Math(du.tolatex(symbol_names={r: 'r', theta: '\\theta', phi: '\\phi'}, replace=replace))
```

```python
Math(T.coors.latex_basis_vectors(symbol_names={r: 'r', theta: '\\theta', phi: '\\phi'}, replace=replace))
```

```python
%matplotlib notebook
M = inner(vv*T.sg, du)
B = BlockMatrix(M)
plt.figure(figsize=(6,4))
plt.spy(B.diags(it=(4,)), markersize=0.2)
```


```python
from sympy import MatrixSymbol, Matrix, Identity

import sympy as sym

sym.init_printing()
%matplotlib inline

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import control
```

#  Espacio de estados

Hasta el momento se han analizado sistemas lineales e invariantes en el tiempo, de una entrada y una salida en el dominio de la frecuencia. La representación en espacio de estados permite modelar y analizar sistemas con múltiples entradas y salidas, tanto para sistemas lineales como no lineales.

Un modelo en espacio de estados de un sistema físico describe relaciones entre entradas, salidas y variables de estado relacionadas por ecuaciones diferenciales de orden $n$ en el dominio del tiempo. Estas ecuaciones se compactan en una ecuación diferencial matricial de primer orden. Las variables de entradas, salidas y estados son expresadas como vectores, y si es un sistema LTI, las ecuaciones se escriben matricialmente.

La función de transferencia de un sistema cuya entrada es una señal $u$ y cuya salida es $y$ de un sistema corresponde a $\frac{U(s)}{y(s)}$.

Su representación en espacio de estados se muestra en el siguiente diagrama.

- La entrada $u$ puede tener un efecto directo sobre la salida a través de $D$.
- $X$ es un vector de señales internas del sistema.
- $\dot{X}$ es el vector de derivadas de $X$.
- $u$ altera la derivada de $X$ a través de $B$.
- $X$ realimenta a sus derivadas a través de $A$.
- $X$ envía información desde el estado hacia la salida $y$ a través de $C$.

$$\dot{X} = A \cdot X + B \cdot u $$
$$y =  C \cdot X +  D \cdot u $$

- **Estado**: Los valores que toma, en un instante determinado, un conjunto de variables a partir de cuál se define el comportamiento de un sistema.
- **Variables de estado**: Las variables que determinan el estado de un sistema dinámico. Las variables de estado no necesitan ser físicamente medibles (observables) pero es recomendable que lo sean.
- **Vector de estado**: El compuesto por todas las variables de estado.
- **Espacio de estados**: Un espacio compuesto por las variables de estado.
- **Ecuaciones de estado**: Las relaciones entre los estados y las entradas.

**Ejemplo**

Considere un sistema físico de entrada $u(t)$ y salida $y(t)$ que es modelado por la siguiente ecuación diferencial:

$$m \frac{d^2y}{dt^2} + b \frac{dy}{dt} + k y(t) = u(t)$$

$$s^2mY(s) + bsY(s) + kY(s) = U(s)$$

$$(s^2m + bs + k )Y(s) = U(s)$$

$$\frac{Y(s)}{U(s)} = \frac{1}{s^2m + bs + k}$$

Como el sistema es de orden 2, se define un vector de estados en 2 dimensiones.

$$x = \left [ \matrix{x_1 \\ x_2} \right ]  $$

Se escogen las variables de estado:

\begin{align}
x_1(t) &= y(t) \\
x_2(t) &= \frac{dy}{dt} = \dot{y}(t)
\end{align}

De donde se desprende que:
\begin{align}
\dot{x}_1(t) &= \frac{dy}{dt} = x_2(t) \\
\dot{x}_2(t) &= \frac{d^2y}{dt^2} \\
&=  \frac{1}{m}\left ( -b\frac{dy}{dt} - ky(t) + u(t)\right ) \\
&=  \frac{1}{m}\left ( -b x_2(t) - k x_1(t) + u(t)\right ) \\
&=  - \frac{k}{m} x_1(t) - \frac{b}{m} x_2(t) + \frac{1}{m}u(t)
\end{align}

Además, la ecuación de salida es:

$$y(t) = x_1(t)$$

Estas ecuaciones se pueden representar vectorialmente como:

\begin{align}
\dot{x}_1(t) &= \left [ \matrix{0 & 1} \right ] \cdot \left [ \matrix{x_1 \\ x_2} \right ]  \\
\dot{x}_2(t) &=  \left [ \matrix{-\frac{k}{m} & -\frac{b}{m}} \right ] \cdot \left [ \matrix{x_1 \\ x_2} \right ]  + \frac{1}{m}u\\
y(t) &= x_1(t)
\end{align}

O matricialmente como:

$$ \left [ \matrix{\dot{x}_1 \\ \dot{x}_2} \right ] = \left [ \matrix{ 0 & 1 \\ -\frac{k}{m} & -\frac{b}{m} } \right ]  \left [ \matrix{x_1 \\ x_2} \right ] + \left [ \matrix{0 \\ \frac{1}{m} } \right ]u$$

$$y = \left [ \matrix{1 & 0} \right ] \left [ \matrix{x_1 \\ x_2} \right ]$$

Dada la forma estándar de los modelos en espacio de estados,

$$\dot{X} = A \cdot X + B \cdot u $$
$$y =  C \cdot X +  D \cdot u $$

Se tiene que:

$$A = \left [ \matrix{ 0 & 1 \\ -\frac{k}{m} & -\frac{b}{m} } \right ], \;
B = \left [ \matrix{0 \\ \frac{1}{m} } \right ], \;
C = \left [ \matrix{1 & 0} \right ], \;
D = 0$$

## El modelo analizado como función de transferencia

Teniendo en cuenta condiciones iniciales iguales a cero, se aplica la transformada de Laplace.

$$sX(s) = A \cdot X(s) + B \cdot U(s) $$
$$Y(S) =  C \cdot X(s) +  D \cdot U(s) $$

Operando la ecuación de estados:

$$sX(s) - A \cdot X(s) = B \cdot U(s) $$

o

$$\left ( sI - A \right ) X(s) = B \cdot U(s) $$

$$ X(s) = \left ( sI - A \right )^{-1} B \cdot U(s) $$

Incorporando la ecuación de estados en la de salida,

$$Y(S) =  C \cdot \left ( sI - A \right )^{-1} B \cdot U(s)  +  D \cdot U(s)$$

$$Y(S) =  \left ( C \cdot \left ( sI - A \right )^{-1} B  +  D \right ) \cdot U(s)$$

La función de transferencia en términos de las matrices del modelo en espacio de estados es:

$$\frac{Y(S)}{U(s)} = G(s)=  C \cdot \left ( sI - A \right )^{-1} B  +  D $$

Cuando el sistema tiene múltiples entradas y múltiples salidas, se obtiene una matriz de transferencia.

```python
k,m,b,t = sym.symbols(['k','m','b','t'],real = True,positive = True)
s = sym.symbols('s',complex = True)
```

```python
A = Matrix([[0,1],[-k/m,-b/m]])
A
```

$\displaystyle \left[\begin{matrix}0 & 1\\- \frac{k}{m} & - \frac{b}{m}\end{matrix}\right]$

```python
B = Matrix([[0],[1/m]])
B
```

$\displaystyle \left[\begin{matrix}0\\\frac{1}{m}\end{matrix}\right]$

```python
C = Matrix([[1,0]])
C
```

$\displaystyle \left[\begin{matrix}1 & 0\end{matrix}\right]$

```python
D = Matrix([[0]])
D
```

$\displaystyle \left[\begin{matrix}0\end{matrix}\right]$

```python
I =  Matrix([[1,0],[0,1]])
I
```

$\displaystyle \left[\begin{matrix}1 & 0\\0 & 1\end{matrix}\right]$

```python
s*I-A
```

$\displaystyle \left[\begin{matrix}s & -1\\\frac{k}{m} & \frac{b}{m} + s\end{matrix}\right]$

```python
INV=(s*I-A).inv()
INV
```

$\displaystyle \left[\begin{matrix}- \frac{\frac{b}{m} + s}{- \frac{k}{m} - s \left(\frac{b}{m} + s\right)} & - \frac{1}{- \frac{k}{m} - s \left(\frac{b}{m} + s\right)}\\\frac{k}{m \left(- \frac{k}{m} - s \left(\frac{b}{m} + s\right)\right)} & - \frac{s}{- \frac{k}{m} - s \left(\frac{b}{m} + s\right)}\end{matrix}\right]$

Encuentre la función de transferencia $G(s)$

$$\frac{Y(S)}{U(s)} = G(s)=  C \cdot \left ( sI - A \right )^{-1} B  +  D $$

```python
C*INV
```

$\displaystyle \left[\begin{matrix}- \frac{\frac{b}{m} + s}{- \frac{k}{m} - s \left(\frac{b}{m} + s\right)} & - \frac{1}{- \frac{k}{m} - s \left(\frac{b}{m} + s\right)}\end{matrix}\right]$

```python
C*INV*B
```

$\displaystyle \left[\begin{matrix}- \frac{1}{m \left(- \frac{k}{m} - s \left(\frac{b}{m} + s\right)\right)}\end{matrix}\right]$

```python
GS = C*INV*B+D
GS[0].simplify()
```

$$\frac{Y(s)}{U(s)} = \frac{1}{s^2m + bs + k}$$

## Representación de un modelo dinámico

Con la notación matricial analizada puede llevarse una ecuación diferencial de orden $n$ mediante una ecuación matricial de orden 1.

Considere un sistema de entrada $u(t)$ y salida $y(t)$ que se modela según la siguiente ecuación diferencial.

$$a_0y(t)+a_1\frac{dy}{dt}+a_2\frac{d^2y}{dt^2}+\cdots+a_n\frac{d^ny}{dt^n} = u(t)$$

$$a_0y(t) = -a_1\frac{dy}{dt}-a_2\frac{d^2y}{dt^2}-\cdots-a_{n-1}\frac{d^{n-1}y}{dt^{n-1}}-a_n\frac{d^ny}{dt^n} + u(t)$$

$$a_n\frac{d^ny}{dt^n} = -a_0y(t) -a_1\frac{dy}{dt} -a_2\frac{d^2y}{dt^2}-\cdots-a_{n-1}\frac{d^{n-1}y}{dt^{n-1}} +u(t)$$

Se definen las variables de estado:

\begin{align}
x_1 &= y \\
x_2 &= \frac{dy}{dt} \\
x_3 &= \frac{d^2y}{dt^2} \\
\cdots &= \cdots \\
x_n &= \frac{d^{n-1}y}{dt^{n-1}} \\
\end{align}

La ecuación diferencial puede reescribirse como:

\begin{align}
\dot{x}_1 &= \frac{dy}{dt} = x_2 \\
\dot{x}_2 &= \frac{d^2y}{dt^2} = x_3 \\
\dot{x}_3 &= \frac{d^3y}{dt^3} = x_4 \\
&\vdots\\
\dot{x}_{n-1} &= \frac{d^{n-1}y}{dt^{n-1}} =  x_{n}  \\
\dot{x}_n &= -\frac{a_0}{a_n}x_1 -\frac{a_1}{a_n}x_2 -\frac{a_2}{a_n}x_3 -\dots -\frac{a_{n-2}}{a_{n}}x_{n-1}  -\frac{a_{n-1}}{a_{n}}x_{n}  + \frac{1}{a_{n}}u
\end{align}

Matricialmente se representa la ecuación de estados:

$$ \left[\matrix{
\dot{x}_1 \\ \dot{x}_2 \\ \dot{x}_3 \\ \vdots \\ \dot{x}_{n-1} \\ \dot{x}_n }
\right ] =
\left [ \matrix{
0 & 1 & 0 & 0 & \dots & 0 & 0 \\
0 & 0 & 1 & 0 & \dots & 0 & 0 \\
0 & 0 & 0 & 1 & \dots & 0 & 0 \\
\vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\
0 & 0 & 0 & 0 & \dots & 0 & 1 \\
-\frac{a_{0}}{a_{n}} & -\frac{a_{1}}{a_{n}} & -\frac{a_{2}}{a_{n}} & -\frac{a_{3}}{a_{n}} & \dots & -\frac{a_{n-2}}{a_{n}} & -\frac{a_{n-1}}{a_{n}} \\
} \right ]
\left [ \matrix{x_1 \\ x_2 \\ x_3 \\ \vdots \\ x_{n-1} \\x_n}
\right ]
+
\left [ \matrix{0 \\ 0 \\ 0 \\ \vdots \\ 0 \\ \frac{1}{a_{n}} } \right ]u$$

Y la ecuación de salida:

$$y = \left [ \matrix{1 & 0 & 0 \dots & 0 & 0} \right ]
\left [ \matrix{x_1 \\ x_2 \\ x_3 \\ \vdots \\ x_{n-1} \\x_n}
\right ]$$

El modelo en espacio de estados es:

$$\dot{X} = A \cdot X + B \cdot u $$
$$y =  C \cdot X +  D \cdot u $$

donde

$$A = \left [ \matrix{
0 & 1 & 0 & 0 & \dots & 0 & 0 \\
0 & 0 & 1 & 0 & \dots & 0 & 0 \\
0 & 0 & 0 & 1 & \dots & 0 & 0 \\
\vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\
0 & 0 & 0 & 0 & \dots & 0 & 1 \\
-\frac{a_{0}}{a_{n}} & -\frac{a_{1}}{a_{n}} & -\frac{a_{2}}{a_{n}} & -\frac{a_{3}}{a_{n}} & \dots & -\frac{a_{n-2}}{a_{n}} & -\frac{a_{n-1}}{a_{n}} \\
} \right ]  $$

$$ B = \left [ \matrix{0 \\ 0 \\ 0 \\ \vdots \\ 0 \\ \frac{1}{a_{n}} } \right ] $$

$$ C = \left [ \matrix{1 & 0 & 0 \dots & 0 & 0} \right ] $$

```python

```

## Valores propios de la matriz A

La función de transferencia en términos de las matrices del modelo en espacio de estados es:

$$\frac{Y(S)}{U(s)} = G(s)=  C \cdot \left ( sI - A \right )^{-1} B  +  D $$

Observe el papel que tiene el término $sI - A$. El operador de inversión aporta el polinomio denominador, es decir, define los polos del sistema. Así, los valores propios de la matriz $A$ son las raíces de la ecuación característica del sistema.

Considere un sistema cuya matriz de transición A es:

$$A = \left [ \matrix{0&1&0\\0&0&1\\-6&-11&-6} \right ]$$

```python
A = Matrix([[0,1,0],[0,0,1],[-6,-11,-6]])
A
```

$\displaystyle \left[\begin{matrix}0 & 1 & 0\\0 & 0 & 1\\-6 & -11 & -6\end{matrix}\right]$

```python
A.eigenvals()
```

## Transformaciones de matrices

Dada una matriz de transición de estados $A$, puede transformarse para que tenga otras propiedades mediante una matriz de transformación $P$:

$$P^{-1}AP$$

### Diagonalización de A

Considere que $A$ es una matriz de transición de $n \times n$ con valores propios distintos como la que se muestra a continuación:

$$A = \left [ \matrix{
0 & 1 & 0 & 0 & \dots & 0 & 0 \\
0 & 0 & 1 & 0 & \dots & 0 & 0 \\
0 & 0 & 0 & 1 & \dots & 0 & 0 \\
\vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\
0 & 0 & 0 & 0 & \dots & 0 & 1 \\
-{a_{n}} & -{a_{n-1}} & -{a_{n-2}} & -{a_{n-3}} & \dots & -{a_{2}} & -{a_{1}} \\
} \right ]  $$

Cuyos $n$ valores propios distintos son:

$$\lambda_1,\lambda_2,\dots,\lambda_n,$$

Considere la matriz de transformación

$$P = \left [ \matrix{
1 & 1 & 1 & 1 & \dots & 1 & 1 \\
\lambda_1 & \lambda_2 & \lambda_3 & \lambda_4 & \dots & \lambda_{n-1} & \lambda_n \\
\lambda_1^2 & \lambda_2^2 & \lambda_3^2 & \lambda_4^2 & \dots & \lambda_{n-1}^2 & \lambda_n^2 \\
\vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\
\lambda_1^{n-2} & \lambda_2^{n-2} & \lambda_3^{n-2} & \lambda_4^{n-2} & \dots & \lambda_{n-1}^{n-2} & \lambda_n^{n-2} \\
\lambda_1^{n-1} & \lambda_2^{n-1} & \lambda_3^{n-1} & \lambda_4^{n-1} & \dots & \lambda_{n-1}^{n-1} & \lambda_n^{n-1} \\
} \right ]  $$

La matríz de transformación $P^{-1}AP$ es la matriz diagonal

$$P^{-1}AP= \left [ \matrix{
\lambda_1 & 0 & 0 & 0 & \dots & 0 & 0 \\
0 & \lambda_2 & 0 & 0 & \dots & 0 & 0 \\
0 & 0 & \lambda_3 & 0 & \dots & 0 & 0 \\
\vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\
0 & 0 & 0 & 0 & \dots & \lambda_{n-1} & 0 \\
0 & 0 & 0 & 0 & \dots & 0 & \lambda_{n} \\
} \right ]  $$

**Ejemplo**

Considere el sistema descrito por:

$$\dot{X} = A \cdot X + B \cdot u $$
$$y =  C \cdot X +  D \cdot u $$

con

$$A = \left [ \matrix{0&1&0\\0&0&1\\-6&-11&-6} \right ]$$
$$B = \left [ \matrix{0\\0\\6} \right ]$$
$$C = \left [ \matrix{1&0&0} \right ]$$
$$D = 0$$

```python
A = Matrix([[0,1,0],[0,0,1],[-6,-11,-6]])
B = Matrix([[0],[0],[6]])
C = Matrix([[1,0,0]])
D = Matrix([[0]])
```

Los valores propios de $A$ son:

```python
A.eigenvals()
```

Se define la matriz de transformación

```python
P = Matrix([[1,1,1],[-3,-2,-1],[9,4,1]])
P
```

$\displaystyle \left[\begin{matrix}1 & 1 & 1\\-3 & -2 & -1\\9 & 4 & 1\end{matrix}\right]$

A partir de la matriz de transformación se define un nuevo conjunto de variables de estado

$$x=P \cdot z$$

Al introducir la transformación a la ecuación del sistema se obtiene:

$$P \cdot \dot{z} = A \cdot P \cdot z + B \cdot u $$
$$y =  C \cdot P \cdot z +  D \cdot u $$

Como la ecuación de transición de estados debe definir las derivadas de las variables de estado, se cambia a:

$$\dot{z} = \overbrace{P^{-1} \cdot A \cdot P}^{A_1} \cdot z + \overbrace{P^{-1} \cdot B}^{B_1} \cdot u $$
$$y =  \underbrace{C \cdot P}_{C_1} \cdot z +  D \cdot u $$

```python
A1 = P.inv()*A*P
A1
```

$\displaystyle \left[\begin{matrix}-3 & 0 & 0\\0 & -2 & 0\\0 & 0 & -1\end{matrix}\right]$

```python
B1 = P.inv()*B
B1
```

$\displaystyle \left[\begin{matrix}3\\-6\\3\end{matrix}\right]$

```python
C1 = C*P
C1
```

$\displaystyle \left[\begin{matrix}1 & 1 & 1\end{matrix}\right]$

Así, el modelo original

$$\dot{X} = A \cdot X + B \cdot u $$
$$y =  C \cdot X +  D \cdot u $$

con

$$A = \left [ \matrix{0&1&0\\0&0&1\\-6&-11&-6} \right ]$$
$$B = \left [ \matrix{0\\0\\6} \right ]$$
$$C = \left [ \matrix{1&0&0} \right ]$$
$$D = 0$$

También puede representarse como:

$$\dot{Z} = A_1 \cdot Z + B_1 \cdot u $$
$$y =  C_1 \cdot Z +  D_1 \cdot u $$

con

$$A_1 = \left [ \matrix{-3&0&0\\0&-2&0\\0&0&-1} \right ]$$
$$B_1 = \left [ \matrix{3\\-6\\3} \right ]$$
$$C_1 = \left [ \matrix{1&1&1} \right ]$$
$$D_1 = 0$$

Compare la función de transferencia que se obtiene con cada modelo en espacio de estados

$$ C \cdot \left ( sI - A \right )^{-1} B  +  D$$

```python
I = Matrix([[1,0,0],[0,1,0],[0,0,1]])
(C*(s*I-A).inv()*B + D)[0].simplify()
```

```python
(C1*(s*I-A1).inv()*B1 + D)[0].simplify()
```

```python
Gss = control.ss(A,B,C,D)
Gss
```

\[
\left(
\begin{array}{rllrllrll|rll}
0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&1\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}\\
0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&1\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}\\
-6\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&-11\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&-6\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&6\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}\\
\hline
1\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}\\
\end{array}\right)
\]

```python
Gss1 = control.ss(A1,B1,C1,D)
Gss1
```

\[
\left(
\begin{array}{rllrllrll|rll}
-3\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&3\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}\\
0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&-2\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&-6\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}\\
0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&-1\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&3\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}\\
\hline
1\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&1\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&1\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}\\
\end{array}\right)
\]

```python
t,y = control.step_response(Gss)
plt.plot(t,y)
```

```python
t1,y1 = control.step_response(Gss1)
t,y = control.step_response(Gss)
plt.plot(t,y,t1,y1)
```

```python

```


# COMP90051 Workshop 2
## Part B: Bayesian inference

***

In this part of the workshop, we'll develop some intuition for priors and posteriors, which are crucial to Bayesian inference.

```python
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import bernoulli, beta
```

### 1. A lucky find

On the way to class, you discover an unusual coin on the ground.

As a dedicated student in statistical ML, you're interested in determining whether the coin is biased.
More specifically, you want to estimate the probability $\theta$ that the coin will land heads-up when you toss it.

You can use the function below to simulate a coin toss: it returns `1` for heads and `0` for tails.

```python
def toss_coin():
    if bernoulli.rvs(p = (int.from_bytes("coin".encode(), 'little') % 10000)/10000):
        return 1
    return 0
```

### 2. Prior belief
Before you even toss the coin, you notice that the heads side appears to have more mass.
Thus, your _prior belief_ is that $\theta$ is slightly biased away from 0.5 towards 0—i.e. you expect tails are more likely.

To quantify this prior belief, we assume that the prior distribution for $\theta$ is $\mathrm{Beta}(a,b)$, for some choice of the hyperparameters $a, b > 0$.
(See [link](https://en.wikipedia.org/wiki/Beta_distribution) for info about the Beta distribution.)
The prior probability density function for $\theta$ is therefore given by:

$$ p(\theta) = \frac{1}{B(a,b)} \theta^{a-1} (1 - \theta)^{b-1} $$
where $B(a,b)$ is a special function called the _Beta function_.

Select appropriate values for $a$ and $b$ by looking at the plot of $p(\theta)$ below.

```python
a = 2
b = 1
theta = np.linspace(0, 1, 1000)
plt.plot(theta, beta.pdf(theta, a, b), 'b-', lw=1)
plt.xlabel(r'$\theta$')
plt.ylabel(r'$p(\theta)$')
plt.show()
```

### 3. Posterior updates
Now toss the coin once and denote the outcome by $x_1$.

```python
x1 = toss_coin()
print(x1)
```

    1

We can update our belief about $\theta$, based on this new evidence $x_1$.
To do this we apply Bayes' rule to compute the posterior for $\theta$:
$$ p(\theta | x_1) = \frac{p(x_1 | \theta) \, p(\theta)}{p(x_1)} \propto p(x_1 | \theta) \, p(\theta)$$
where $p(\theta)$ is the prior given above and
$$ p(x_1 | \theta) = \theta^{x_1} (1 - \theta)^{1 - x_1} $$
is the likelihood.

***
**Exercise:** Show (on paper) that
$$ p(\theta | x_1) \propto \theta^{x_1 + a - 1} (1 - \theta)^{(1 - x_1) + b - 1} $$
which implies that $\theta | x_1 \sim \mathrm{Beta}[x_1 + a, (1 - x_1) + b]$.
***

Toss the coin a second time, denoting the outcome by $x_2$.

```python
x2 = toss_coin()
print(x2)
```

    1

Again, we want to update our belief about $\theta$ based on the new information $x_2$.
We take the previous posterior $p(\theta|x_1)$ as the new prior and apply Bayes' rule:
$$ p(\theta | x_1, x_2) \propto p(x_2 | \theta) p(\theta | x_1)$$
\[Note: We assume the tosses are independent, otherwise the likelihood for $x_2$ would depend on $x_1$.\]
This gives $\theta | x_1, x_2 \sim \mathrm{Beta}[x_1 + x_2 + a, (2 - x_1 - x_2) + b]$.

***
**Exercise:** Show that for $n$ coin tosses, the posterior is $\theta | x_1, \ldots, x_n \sim \operatorname{Beta}[n_H + a, n - n_H + b]$ where $n_H = \sum_{i = 1}^{n} x_i$ is the number of heads observed.
***

### 4. MAP estimator and MLE estimator

The posterior $\theta|x_1, \ldots, x_n$ contains all the information we know about $\theta$ after observing $n$ coin tosses.
One way of obtaining a point estimate of $\theta$ from the posterior, is to take the value with the maximum a posteriori probability (MAP):
$$
\begin{align}
    \hat{\theta}_{MAP} &= \arg \max_{\theta} p(\theta|x_1, \ldots, x_n) \\
        & = \frac{n_H + a - 1}{n + a + b - 2}
\end{align}
$$

In general, the MAP estimator gives a different result to the maximum likelihood estimator (MLE) for $\theta$:
$$
\begin{align}
    \hat{\theta}_{MLE} &=\arg \max_{\theta} p(x_1, \ldots, x_n|\theta) \\
        & = \frac{n_H}{n}
\end{align}
$$

***
**Exercise:** Derive the above results for $\hat{\theta}_{MAP}$ and  $\hat{\theta}_{MLE}$.
***

### 5. Convergence of the estimates

Let's now toss the coin an additional 48 times (so that $n = 50$), recording $\hat{\theta}_{MLE}$ and $\hat{\theta}_{MAP}$ after each toss.

```python
extra_tosses = 48
num_tosses = 2 + extra_tosses
num_heads = 0
theta_map = np.zeros(num_tosses)
theta_mle = np.zeros(num_tosses)
for i in range(0, num_tosses):
    if i == 0:
        num_heads += x1
    elif i == 1:
        num_heads += x2
    else:
        num_heads += toss_coin()
    theta_map[i] = ... # fill in
    theta_mle[i] = ... # fill in
```

We plot the results below.

```python
plt.plot(theta_map, label = "MAP")
plt.plot(theta_mle, label = "MLE")
plt.xlabel('Number of draws')
plt.ylabel(r'$\hat{\theta}$')
plt.legend()
plt.show()
```

**Questions:**

1. Is the coin biased?
1. Do the MAP and MLE estimates converge to the same value for $\theta$?
1. What happens if you set $a = 1; b = 1$?
1. How does the posterior distribution for $\theta$ compare to the prior plotted above? (Use the code block below to plot the posterior.)

```python
plt.plot(theta, beta.pdf(theta, a + num_heads, b + num_tosses - num_heads), 'b-', lw=1)
plt.xlabel(r'$\theta$')
plt.ylabel(r'$p(\theta|x_1, \ldots, x_n)$')
plt.show()
```

```python

```


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

# Discussion: Week 6

## Exercise: Two-Period Cake-Eating Problem

Recall the two-period cake-eating problem. A person has an initial quantity of cake $K_0>0$ and receives utility from consuming cake in period 0 and period 1:

\begin{align}
U(C_0,C_1) & = \log C_0 + \beta \log C_1.
\end{align}

The person chooses how much to consume in each period, $C_0$ and $C_1$, and a quantity of cake $K_1$ to carry forward from period 0 to maximize her utility subject to the following two budget constraints:

\begin{align}
C_0 + K_1 & = K_0\\
C_1 & = K_1,
\end{align}

where the boundary condition $K_2 = 0$ has been imposed. Use the budget constraints to substitute $C_0$ and $C_1$ out of the utility function. Then the maximization problem can be written concisely as a choice of just $K_1$:

\begin{align}
\max_{K_1} \; \log (K_0-K_1) + \beta \log K_1
\end{align}

The first-order condition with respect to $K_1$ is:

\begin{align}
-\frac{1}{K_0-K_1} + \frac{\beta}{K_1} & = 0,
\end{align}

Given $K_1$, values for $C_0$ and $C_1$ are implied by the budget constraints

The function in the next cell computes the solution for the two-period cake-eating problem and provides an example of how to use the function.

```python
# Define a function to compute the solution for the two-period cake-eating problem. CELL PROVIDED.
def two_period_cake(beta,cake0):

    '''Function for computing solution to two-period cake-eating problem

            Args:
            beta (float):   subjective discount factor
            cake0 (float):  initial quantity of cake

        Returns:
            numpy array, numpy array
    '''

    # Initialize arrays for cake and consumption
    cake = np.zeros(3)
    consumption = np.zeros(2)

    # Assign values to cake array elements
    cake[0] = cake0
    cake[1] = beta/(1+beta)*cake[0]

    # Assign values to consumption array elements as implied by the budget constraints
    consumption[0] = cake[0]-cake[1]
    consumption[1] = cake[1]

    return cake,consumption

# Example of function use: Plot cake path for beta=0.99 and cake0=100
beta = 0.99
cake0 = 100

cake,consumption = two_period_cake(beta=beta,cake0=cake0)

plt.plot(cake)
plt.title('Cake Quantity for $\\beta='+str(beta)+'$')
plt.grid()
```

Suppose that the initial quantity of cake is $K_0 = 100$. For $\beta = 0, 0.25, 0.5, 0.75, 0.99, 1$, use the function `two_period_cake()` to plot the optimal path for the quantity of cake for periods 0 through 2. Add labels to the plot command so you know which line corresponds to which value of $\beta$.

```python
# Create a variable called 'cake0' that stores the initial quantity of cake

# Plot optimal path of cake for the different values of beta.

# Put a legend to the right of the current axis. PROVIDED
plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))

```

**Questions**

1. How does increasing $\beta$ affect the otpimal path of $K$?
2. For $\beta=0$, $K_1 = 0$. Why does this make sense?

**Answers**

1.

2.

## Exercise: Three-Period Cake-Eating Problem

Now consider a the three-period cake-eating problem. A person has an initial quantity of cake $K_0>0$ and receives utility from consuming cake in period 0, period 1, and period 2:

\begin{align}
U(C_0,C_1,C_2) & = \log C_0 + \beta \log C_1+ \beta^2 \log C_2.
\end{align}

The person chooses how much to consume in each period, $C_0$, $C_1$, and $C_2$, and quantities of cake $K_1$ and $K_2$ to carry foward from period 0 and from period 1 to maximize her utility subject to the following two budget constraints:

\begin{align}
C_0 + K_1 & = K_0\\
C_1 + K_2 & = K_1\\
C_2 & = K_2,
\end{align}

where the boundary condition $K_3 = 0$ has been imposed. Use the budget constraints to substitute $C_0$, $C_1$, and $C_2$ out of the utility function. Then the maximiation problem can be written concisely as a choice of $K_1$ and $K_2$:

\begin{align}
\max_{K_1,K_2} \; \log (K_0-K_1) + \beta\log (K_1-K_2) + \beta^2 \log K_2
\end{align}

The first-order conditions with respect to $K_1$ and $K_2$ are:

\begin{align}
-\frac{1}{K_0-K_1} + \frac{\beta}{K_1-K_2} & = 0\\
-\frac{1}{K_1-K_2} + \frac{\beta}{K_2} & = 0,
\end{align}

### Part (a)

Use the first-order conditions for the three-period cake-eating problem to solve by hand for:

1. $K_2$ as a function of $K_1$
2. $K_1$ as a function of $K_0$

Use what you derive in this part to do **Part (b)** and **Part (c)**.

### Part (b)

Copy and paste the full definition of the function `two_period_cake()` into the cell below and do the following:

1. Rename the function `three_period_cake`
2. Edit the function so that it computes the optimal path for the three-period problem

```python
# Define a function to compute the solution for the three-period cake-eating problem.

```

### Part (c)

Suppose that the initial quantity of cake is $K_0 = 100$. For $\beta = 0, 0.25, 0.5, 0.75, 0.99, 1$, plot the optimal path for the quantity of cake for periods 0 through 3. Add labels to the plot command so you know which line corresponds to which value of $\beta$.

```python
# Create a variable called 'cake0' that stores the initial quantity of cake

# Plot optimal path of cake for the different values of beta.

# Put a legend to the right of the current axis. PROVIDED
plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))

```


Computational Seismology

Numerical derivatives based on the Fourier Transform

Seismo-Live: http://seismo-live.org

##### Authors:
* Fabian Linder ([@fablindner](https://github.com/fablindner))
* Heiner Igel ([@heinerigel](https://github.com/heinerigel))
* David Vargas ([@dvargas](https://github.com/davofis))
---

## Basic Equations
The derivative of function $f(x)$ with respect to the spatial coordinate $x$ is calculated using the differentiation theorem of the Fourier transform:

\begin{equation}
\frac{d}{dx} f(x) = \frac{1}{\sqrt{2\pi}} \int_{-\infty}^{\infty} ik F(k) e^{ikx} dk
\end{equation}

In general, this formulation can be extended to compute the n−th derivative of $f(x)$ by considering that $F^{(n)}(k) = D(k)^{n}F(k) = (ik)^{n}F(k)$. Next, the inverse Fourier transform is taken to return to physical space.

\begin{equation}
f^{(n)}(x) = \mathscr{F}^{-1}[(ik)^{n}F(k)] = \frac{1}{\sqrt{2\pi}} \int_{-\infty}^{\infty} (ik)^{n} F(k) e^{ikx} dk
\end{equation}

```python
# Import all necessary libraries, this is a configuration step for the exercise.
# Please run it before the simulation code!
import numpy as np
import matplotlib.pyplot as plt

# Show the plots in the Notebook.
plt.switch_backend("nbagg")
```

#### Exercise 1

Define a python function call "fourier_derivative(f, dx)" that compute the first derivative of a function $f$ using the Fourier transform properties.

```python
def fourier_derivative(f, dx):
    # Length of vector f
    nx = np.size(f)
    # Initialize k vector up to Nyquist wavenumber
    kmax = np.pi/dx
    dk = kmax/(nx/2)
    k = np.arange(float(nx))
    k[: int(nx/2)] = k[: int(nx/2)] * dk
    k[int(nx/2) :] = k[: int(nx/2)] - kmax
    # Fourier derivative
    ff = np.fft.fft(f); ff = 1j*k*ff
    df_num = np.real(np.fft.ifft(ff))
    return df_num
```

#### Exercise 2

Calculate the numerical derivative based on the Fourier transform to show that the derivative is exact. Define an arbitrary function (e.g. a Gaussian) and initialize its analytical derivative on the same spatial grid. Calculate the numerical derivative and the difference to the analytical solution. Vary the wavenumber content of the analytical function. Does it make a difference? Why is the numerical result not entirely exact?

```python
# Basic parameters
# ---------------------------------------------------------------
nx = 128
x, dx = np.linspace(2*np.pi/nx, 2*np.pi, nx, retstep=True)
sigma = 0.5
xo = np.pi

# Initialize Gauss function
f = np.exp(-1/sigma**2 * (x - xo)**2)

# Numerical derivative
df_num = fourier_derivative(f, dx)

# Analytical derivative
df_ana = -2*(x-xo)/sigma**2 * np.exp(-1/sigma**2 * (x-xo)**2)

# To make the error visible, it is multiply by 10^13
df_err = 1e13*(df_ana - df_num)

# Error between analytical and numerical solution
err = np.sum((df_num - df_ana)**2) / np.sum(df_ana**2) * 100
print('Error: %s' %err)

```

#### Exercise 3

Now that the numerical derivative is available, we can visually inspect our results. Make a plot of both, the analytical and numerical derivatives together with the difference error.

```python
# Plot analytical and numerical derivatives
# ---------------------------------------------------------------

plt.subplot(2,1,1)
plt.plot(x, f, "g", lw = 1.5, label='Gaussian')
plt.legend(loc='upper right', shadow=True)
plt.xlabel('$x$')
plt.ylabel('$f(x)$')
plt.axis([2*np.pi/nx, 2*np.pi, 0, 1])

plt.subplot(2,1,2)
plt.plot(x, df_ana, "b", lw = 1.5, label='Analytical')
plt.plot(x, df_num, 'k--', lw = 1.5, label='Numerical')
plt.plot(x, df_err, "r", lw = 1.5, label='Difference')
plt.legend(loc='upper right', shadow=True)
plt.xlabel('$x$')
plt.ylabel('$\partial_x f(x)$')
plt.axis([2*np.pi/nx, 2*np.pi, -2, 2])

plt.show()
#plt.savefig('Fig_5.9.png')
```

```python

```


# Hamiltonian Monte Carlo (HMC)

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

## Use of auxiliary variables

Slice sampling is a simple MCMC algorithm that introduces the idea of auxiliary variables. The motivation for slice sampling is that if we can sample uniformly from the region under the graph of the target distribution, we will have random samples from the target distribution. In the univariate case, the algorithm is as follows

- start with some $x$ where $p(x) \ne 0$
- repeat
    - sample $y$ (auxiliary variable) uniformly from 0 to $f(x)$
    - draw a horizontal line at $y$ within $p(x)$ (this may consist of multiple intervals)
    - sample $x$ from the horizontal segments

The auxiliary $y$ variable allows us to sample $(x, y)$ points that are in the region under the graph of the target distribution. Only the $x$ variable is used for the Monte Carlo samples - the $y$ variables are simply discarded. This works because the joint distribution is constructed so that it factors $p(x, y) = p(y \mid x) p(x)$ and so projecting leaves just $p(x)$. The slice sampler is a Markov Chain Monte Carlo method since the next $(x, y)$ position depends only on the current position. Like Gibbs sampling, there is no tuning process and all proposals are accepted. For slice sampling, you either need the inverse distribution function or some way to estimate it. Later we will see that Hamiltonian Monte Carlo also uses auxiliary variables to generate a new proposal in an analogous way.

A toy example illustrates the process - Suppose we want to draw random samples from the posterior distribution $\mathcal{N}(0, 1)$ using slice sampling

Start with some value $x$
- sample $y$ from $\mathcal{U}(0, f(x))$ - this is the horizontal "slice" that gives the method its name
- sample the next $x$ from $f^{-1}(y)$ - this is typically done numerically
- repeat

Will sketch picture in class to show what is going on.

### A simple slice sampler example

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

```python
dist = stats.norm(5, 3)
w = 0.5
x = dist.rvs()

niters = 1000
xs = []
while len(xs) < niters:
    y = np.random.uniform(0, dist.pdf(x))
    lb = x
    rb = x
    while y < dist.pdf(lb):
        lb -= w
    while y < dist.pdf(rb):
        rb += w
    x = np.random.uniform(lb, rb)
    if y > dist.pdf(x):
        if np.abs(x-lb) < np.abs(x-rb):
            lb = x
        else:
            lb = y
    else:
        xs.append(x)
```

```python
plt.hist(xs, 20)
pass
```

Notes on the slice sampler:

- the slice may consist of disjoint pieces for multimodal distributions
- the slice can be a rectangular hyperslab for multivariable posterior distributions
- sampling from the slice (i.e. finding the boundaries at level $y$) is non-trivial and may involve iterative rejection steps - see figure below (from Wikimedia) for a typical approach - the blue bars represent disjoint pieces of the true slice through a bimodal distribution and the black lines are the proposal distribution approximating the true slice

## Hamiltonian Monte Carlo (HMC)

HMC uses an auxiliary variable corresponding to the momentum of particles in a potential energy well to generate proposal distributions that can make use of gradient information in the posterior distribution. For reversibility to be maintained, the total energy of the particle has to be conserved - hence we are interested in Hamiltonian systems. The main attraction of HMC is that it works much better than other methods when variables of interest are highly correlated. Because we have to solve problems involving momentum, we need to understand how to numerically solve differential equations in a way that is both accurate (i.e. second order) and preserves total energy (necessary for a Hamiltonian system).

Example adapted from [MCMC: Hamiltonian Monte Carlo (a.k.a. Hybrid Monte Carlo)](https://theclevermachine.wordpress.com/2012/11/18/mcmc-hamiltonian-monte-carlo-a-k-a-hybrid-monte-carlo/)

### Hamiltonian systems

In a Hamiltonian system, we consider particles with position $x$ and momentum (or velocity if we assume unit mass) $v$. The total energy of the system $H(x, v) = K(v) + U(x)$, where $K$ is the kinetic energy and $U$ is the potential energy, is conserved. Such a system satisfies the following Hamiltonian equations

$$
\begin{align}
\frac{dx}{dt} &= & \frac{\delta H}{dv} \\
\frac{dv}{dt} &= & -\frac{\delta H}{dx}
\end{align}
$$

Since $K$ depends only on $v$ and $U$ depends only on $x$, we have
$$
\begin{align}
\frac{dx}{dt} &= & \frac{\delta K}{dv} \\
\frac{dv}{dt} &= & -\frac{\delta U}{dx}
\end{align}
$$

#### Harmonic oscillator

We will consider solving a classical Hamiltonian system - that of a undamped spring governed by the second order differential equation

$$
x'' + x = 0
$$

We convert this to two first order ODEs by using a dummy variable $x' = v$ to get

$$
\begin{align}
x' &= v \\
v' &= -x
\end{align}
$$

From the Hamiltonian equations above, this is equivalent to a system with kinetic energy $K(v) = \frac{1}{2}v^2$ and potential energy $U(x) = \frac{1}{2}x^2$.

Writing in matrix form,

$$
A = \pmatrix{ x' \\ v' } = \pmatrix{0 & 1 \\ -1 & 0} \pmatrix{x \\ v}
$$

and in general, for the state vector $x$,

$$
x' = Ax
$$

We note that $A$ is anti- or skew-symmetric ($A^T = -A$), and hence has purely imaginary eigenvalues. Solving $|A - \lambda I = 0$, we see that the eigenvalues and eigenvectors are $i, \pmatrix{1\\i}$ and $-i, \pmatrix{1\\-i}$. Since the eigenvalues are pure imaginary, we see that the solution for the initial conditions $(x,v) = (1, 0)$ is $x(t) = e^{it}$ and the orbit just goes around a circle with a period of $2\pi$, neither growing nor decaying. Another weay of seeing this is that the Hamiltonian $H(u, v)$ or sum of potential ($U(x)) = \frac{1}{2}x^2$) and kinetic energy ($K(v) = \frac{1}{2}v^2$) is constant, i.e. in vector form, $(x^T x) = \text{constant}$.

### Finite difference methods

We want to find a finite difference approximation to $u' = Au$ that is **accurate** and **preserves total energy**. If total energy is not preserved, the orbit will either spiral in towards zero or outwards away from the unit circle. If the accuracy is poor, the orbit will not be close to its starting value after $t = 2\pi$. This gives us an easy way to visualize how good our numerical scheme is. We can also compare the numerical scheme to the Taylor series to evaluate its accuracy.

#### Forward Euler

The simplest finite difference scheme for integrating ODEs is the forward Euler

$$
\frac{u_{n+1} - u_n}{\Delta t} = A u_n
$$

Rearranging terms, we get

$$
u_{n+1} = u_n + \Delta t A u_n = \left( I + \Delta t A \right) u_n
$$

Since the eigenvalues of $A$ are $\pm i$, we see that the eigenvalues of the forward Euler matrix are $1 \pm i$. Since the absolute value of the eigenvalues is greater than 1, we expect **growing** solutions - i.e. the solution will spiral away from the unit circle.

```python
import scipy.linalg as la
```

```python
def f_euler(A, u, N):
    orbit = np.zeros((N,2))

    dt = 2*np.pi/N
    for i in range(N):
        u = u + dt * A @ u
        orbit[i] = u
    return orbit
```

```python
A = np.array([[0,1],[-1,0]])
u = np.array([1.0,0.0])
N = 64
orbit = f_euler(A, u, N)
```

##### Accuracy

```python
la.norm(np.array([1.0,0.0]) - orbit[-1])
```

    0.3600318484671193

##### Conservation of energy

```python
plt.plot([p @ p for p in orbit])
pass
```

```python
ax = plt.subplot(111)
plt.plot(orbit[:, 0], orbit[:,1], 'o')
ax.axis('square')
plt.axis([-1.5, 1.5, -1.5, 1.5])
pass
```

##### Accuracy and conservation of energy

We can see that forward Euler is not very accurate and also does not preserve energy since the orbit spirals away from the unit circle.

#### The trapezoidal method

The trapezoidal method uses the following scheme

$$
\frac{u_{n+1} - u_n}{\Delta t} = \frac{1}{2}  ( A u_{n+1} + A u_{n})
$$

This is an implicit scheme (because $u_{n+1}$ appears on the RHS) whose solution is

$$
u_{n+1} = \left(I - \frac{\Delta t}{2} A \right)^{-1} \left(I + \frac{\Delta t}{2} A \right) u_{n} = B u_n
$$

By inspection, we see that the eigenvalues are the complex conjugates of

$$
\frac{1 + \frac{\Delta t}{2} i}{1 - \frac{\Delta t}{2} i}
$$

whose absolute value is 1 - hence, energy is conserved. If we expand the matrix $B$ using the geometric series and compare with the Taylor expansion, we see that the trapezoidal method has local truncation error $O(h^3)$ and hence accuracy $O(h^2)$, where $h$ is the time step.

```python
def trapezoidal(A, u, N):
    p = len(u)
    orbit = np.zeros((N,p))

    dt = 2*np.pi/N
    for i in range(N):
        u = la.inv(np.eye(p) - dt/2 * A) @ (np.eye(p) + dt/2 * A) @ u
        orbit[i] = u
    return orbit
```

```python
A = np.array([[0,1],[-1,0]])
u = np.array([1.0,0.0])
N = 64
orbit = trapezoidal(A, u, N)
```

##### Accuracy

```python
la.norm(np.array([1.0,0.0]) - orbit[-1])
```

    0.005039305635733781

##### Conservation of energy

```python
plt.plot([p @ p for p in orbit])
pass
```

```python
ax = plt.subplot(111)
plt.plot(orbit[:, 0], orbit[:,1], 'o')
ax.axis('square')
plt.axis([-1.5, 1.5, -1.5, 1.5])
pass
```

#### The leapfrog method

The leapfrog method uses a second order difference to update $u_n$. The algorithm simplifies to the following explicit scheme:

- First take one half-step for v
- Then take a full step for u
- Then take one final half step for v

It performs almost as well as the trapezoidal method, with the advantage of being an explicit scheme and cheaper to calculate, so the leapfrog method is used in HMC.

```python
def leapfrog(A, u, N):
    orbit = np.zeros((N,2))

    dt = 2*np.pi/N
    for i in range(N):
        u[1] = u[1] + dt/2 * A[1] @ u
        u[0] = u[0] + dt * A[0] @ u
        u[1] = u[1] + dt/2 * A[1] @ u
        orbit[i] = u
    return orbit
```

##### If we don't care about the intermediate steps, it is more efficient to just take 1/2 steps at the beginning and end

```python
def leapfrog2(A, u, N):
    dt = 2*np.pi/N

    u[1] = u[1] + dt/2 * A[1] @ u
    for i in range(N-1):
        u[0] = u[0] + dt * A[0] @ u
        u[1] = u[1] + dt * A[1] @ u

    u[0] = u[0] + dt * A[0] @ u
    u[1] = u[1] + dt/2 * A[1] @ u
    return u
```

```python
A = np.array([[0,1],[-1,0]])
u = np.array([1.0,0.0])
N = 64
```

```python
orbit = leapfrog(A, u, N)
```

##### Accuracy

```python
la.norm(np.array([1.0,0.0]) - orbit[-1])
```

    0.0025229913808033464

##### Conservation of energy

Note that unlike the trapezoidal scheme, energy is not perfectly conserved.

```python
plt.plot([p @ p for p in orbit])
pass
```

```python
ax = plt.subplot(111)
plt.plot(orbit[:, 0], orbit[:,1], 'o')
ax.axis('square')
plt.axis([-1.5, 1.5, -1.5, 1.5])
pass
```

### From Hamiltonians to probability distributions

The physical analogy considers the negative log likelihood of the target distribution $p(x)$ to correspond to a potential energy well, with a collection of particles moving on the surface of the well. The state of each particle is given only by its position and momentum (or velocity if we assume unit mass for each particle). In a Hamiltonian system, the total energy $H(x,, v) = U(x) + K(v)$ is conserved. From statistical mechanics, the probability of each state is related to the total energy of the system

$$
\begin{align}
p(x, v) & \propto e^{-H(x, v)} \\
&= e^{-U(x) - K(v)} \\
&= e^{-P(x)}e^{-K(v)} \\
& \propto p(x) \, p(v)
\end{align}
$$

Since the joint distribution factorizes $p(x, v) = p(x)\, p(v)$, we can select an initial random $v$ for a particle, numerically integrate using a finite difference method such as the leapfrog and then use the updated $x^*$ as the new proposal. The acceptance ratio for the new $x^*$ is

$$
\frac{ e^{ -U(x^*)-K(v^*) }} { e^{-U(x)-K(v)} } = e^{U(x)-U(x^*)+K(x)-K(x^*)}
$$

If our finite difference scheme was exact, the acceptance ration would be 1 since energy is conserved with Hamiltonian dynamics. However, as we have seen, the leapfrog method does not conserve energy perfectly and an accept/reject step is still needed.

#### Example of HMC

We will explore how HMC works when the target distribution is bivariate normal centered at zero

$$
x \sim N(0, \Sigma)
$$

In practice of course, the target distribution will be the posterior distribution and depend on both data and distributional parameters.

The potential energy or negative log likelihood is proportional to
$$
U(x) = \frac{x^T\Sigma^{-1} x}{2}
$$

The kinetic energy is given by
$$
K(v) = \frac{v^T v}{2}
$$

where the initial $v_0$ is chosen at random from the unit normal at each step.

To find the time updates, we use the Hamiltonian equations and find the first derivatives of total energy with respect to $x$ and $v$

$$
\begin{align}
x' &= \frac{\delta K}{\delta v} &= v \\
v' &= -\frac{\delta U}{\delta x} &= -\Sigma^{-1} x \\
\end{align}
$$

giving us the block matrix

$$
A = \pmatrix{0 & 1 \\ -\Sigma^{-1} & 0}
$$

By using the first derivatives, we are making use of the gradient information on the log posterior to guide the proposal distribution.

##### This is what the target distribution should look like

```python
sigma = np.array([[1,0.8],[0.8,1]])
mu = np.zeros(2)
ys = np.random.multivariate_normal(mu, sigma, 1000)
sns.kdeplot(ys[:,0], ys[:,1])
plt.axis([-3.5,3.5,-3.5,3.5])
pass
```

##### This is the HMC posterior

```python
def E(A, u0, v0, u, v):
    """Total energy."""
    return (u0 @ tau @ u0 + v0 @ v0) - (u @ tau@u + v @ v)
```

```python
def leapfrog(A, u, v, h, N):
    """Leapfrog finite difference scheme."""
    v = v - h/2 * A @ u
    for i in range(N-1):
        u = u + h * v
        v = v - h * A @ u

    u = u + h * v
    v = v - h/2 * A @ u

    return u, v
```

```python
niter = 100
h = 0.01
N = 100

tau = la.inv(sigma)

orbit = np.zeros((niter+1, 2))
u = np.array([-3,3])
orbit[0] = u
for k in range(niter):
    v0 = np.random.normal(0,1,2)
    u, v = leapfrog(tau, u, v0, h, N)

    # accept-reject
    u0 = orbit[k]
    a = np.exp(E(A, u0, v0, u, v))
    r = np.random.rand()

    if r < a:
        orbit[k+1] = u
    else:
        orbit[k+1] = u0
```

```python
sns.kdeplot(orbit[:, 0], orbit[:, 1])
plt.plot(orbit[:,0], orbit[:,1], alpha=0.2)
plt.scatter(orbit[:1,0], orbit[:1,1],  c='red', s=30)
plt.scatter(orbit[1:,0], orbit[1:,1],  c=np.arange(niter)[::-1], cmap='Reds')
plt.axis([-3.5,3.5,-3.5,3.5])
pass
```

```python

```


# Simple pendulum using Lagrange's equation

Defines a LagrangianPendulum class that is used to generate basic pendulum plots from solving Lagrange's equations.

* Last revised 17-Mar-2019 by Dick Furnstahl (furnstahl.1@osu.edu).

## Euler-Lagrange equation

For a simple pendulum, the Lagrangian with generalized coordinate $\phi$ is

$\begin{align}
  \mathcal{L} = \frac12 m L^2 \dot\phi^2 - mgL(1 - \cos\phi)
\end{align}$

The Euler-Lagrange equation is

$\begin{align}
 \frac{d}{dt}\frac{\partial\mathcal{L}}{\partial \dot\phi} = \frac{\partial\mathcal L}{\partial\phi}
 \quad\Longrightarrow\quad
 m L^2 \ddot \phi = -mgL\sin\phi
  \ \mbox{or}\ \ddot\phi = - \omega_0^2\sin\phi = 0
  \;.
\end{align}$

## Hamilton's equations

The generalized momentum corresponding to $\phi$ is

$\begin{align}
   \frac{\partial\mathcal{L}}{\partial \dot\phi} = m L^2 \dot\phi \equiv p_\phi
   \;.
\end{align}$

We can invert this equation to find $\dot\phi = p_\phi / m L^2$.
Constructing the Hamiltonian by Legendre transformation we find

$\begin{align}
  \mathcal{H} &= \dot\phi p_\phi  - \mathcal{L} \\
              &= \frac{p_\phi^2}{m L^2} - \frac12 m L^2 \dot\phi^2 + mgL(1 - \cos\phi) \\
              &= \frac{p_\phi^2}{2 m L^2} + mgL(1 - \cos\phi)
              \;.
\end{align}$

Thus $\mathcal{H}$ is simply $T + V$.  Hamilton's equations are

$\begin{align}
  \dot\phi &= \frac{\partial\mathcal{H}}{\partial p_\phi} = \frac{p_\phi}{m L^2} \\
  \dot p_\phi &= -\frac{\partial\mathcal{H}}{\partial \phi} = -mgL \sin\phi
  \;.
\end{align}$

```python
%matplotlib inline
```

```python
import numpy as np
from scipy.integrate import odeint, solve_ivp

import matplotlib.pyplot as plt
```

```python
# The dpi (dots-per-inch) setting will affect the resolution and how large
#  the plots appear on screen and printed.  So you may want/need to adjust
#  the figsize when creating the figure.
plt.rcParams['figure.dpi'] = 100.    # this is the default for notebook

# Change the common font size (smaller when higher dpi)
font_size = 12
plt.rcParams.update({'font.size': font_size})
```

## Pendulum class and utility functions

```python
class LagrangianPendulum():
    """
    Pendulum class implements the parameters and Lagrange's equations for
     a simple pendulum (no driving or damping).

    Parameters
    ----------
    L : float
        length of the simple pendulum
    g : float
        gravitational acceleration at the earth's surface
    omega_0 : float
        natural frequency of the pendulum (\sqrt{g/l} where l is the
        pendulum length)
    mass : float
        mass of pendulum

    Methods
    -------
    dy_dt(t, y)
        Returns the right side of the differential equation in vector y,
        given time t and the corresponding value of y.
    """
    def __init__(self, L=1., mass=1., g=1.
                ):
        self.L = L
        self.g = g
        self.omega_0 = np.sqrt(g/L)
        self.mass = mass

    def dy_dt(self, t, y):
        """
        This function returns the right-hand side of the diffeq:
        [dphi/dt d^2phi/dt^2]

        Parameters
        ----------
        t : float
            time
        y : float
            A 2-component vector with y[0] = phi(t) and y[1] = dphi/dt

        Returns
        -------

        """
        return [y[1], -self.omega_0**2 * np.sin(y[0]) ]

    def solve_ode(self, t_pts, phi_0, phi_dot_0,
                  abserr=1.0e-9, relerr=1.0e-9):
        """
        Solve the ODE given initial conditions.
        Specify smaller abserr and relerr to get more precision.
        """
        y = [phi_0, phi_dot_0]
        solution = solve_ivp(self.dy_dt, (t_pts[0], t_pts[-1]),
                             y, t_eval=t_pts,
                             atol=abserr, rtol=relerr)
        phi, phi_dot = solution.y

        return phi, phi_dot
```

```python
def plot_y_vs_x(x, y, axis_labels=None, label=None, title=None,
                color=None, linestyle=None, semilogy=False, loglog=False,
                ax=None):
    """
    Generic plotting function: return a figure axis with a plot of y vs. x,
    with line color and style, title, axis labels, and line label
    """
    if ax is None:        # if the axis object doesn't exist, make one
        ax = plt.gca()

    if (semilogy):
        line, = ax.semilogy(x, y, label=label,
                            color=color, linestyle=linestyle)
    elif (loglog):
        line, = ax.loglog(x, y, label=label,
                          color=color, linestyle=linestyle)
    else:
        line, = ax.plot(x, y, label=label,
                    color=color, linestyle=linestyle)

    if label is not None:    # if a label if passed, show the legend
        ax.legend()
    if title is not None:    # set a title if one if passed
        ax.set_title(title)
    if axis_labels is not None:  # set x-axis and y-axis labels if passed
        ax.set_xlabel(axis_labels[0])
        ax.set_ylabel(axis_labels[1])

    return ax, line
```

```python
def start_stop_indices(t_pts, plot_start, plot_stop):
    start_index = (np.fabs(t_pts-plot_start)).argmin()  # index in t_pts array
    stop_index = (np.fabs(t_pts-plot_stop)).argmin()  # index in t_pts array
    return start_index, stop_index
```

```python

```

## Make simple pendulum plots

```python
# Labels for individual plot axes
phi_vs_time_labels = (r'$t$', r'$\phi(t)$')
phi_dot_vs_time_labels = (r'$t$', r'$d\phi/dt(t)$')
state_space_labels = (r'$\phi$', r'$d\phi/dt$')

# Common plotting time (generate the full time then use slices)
t_start = 0.
t_end = 50.
delta_t = 0.001

t_pts = np.arange(t_start, t_end+delta_t, delta_t)

L = 1.
g = 1.
mass = 1.

# Instantiate a pendulum
p1 = LagrangianPendulum(L=L, g=g, mass=mass)

```

```python
# both plots: same initial conditions
phi_0 = (3./4.)*np.pi
phi_dot_0 = 0.
phi, phi_dot = p1.solve_ode(t_pts, phi_0, phi_dot_0)

# start the plot!
fig = plt.figure(figsize=(15,5))
overall_title = 'Simple pendulum from Lagrangian:  ' + \
                rf' $\omega_0 = {p1.omega_0:.2f},$' + \
                rf'  $\phi_0 = {phi_0:.2f},$' + \
                rf' $\dot\phi_0 = {phi_dot_0:.2f}$' + \
                '\n'     # \n means a new line (adds some space here)
fig.suptitle(overall_title, va='baseline')

# first plot: phi plot
ax_a = fig.add_subplot(1,3,1)

start, stop = start_stop_indices(t_pts, t_start, t_end)
plot_y_vs_x(t_pts[start : stop], phi[start : stop],
            axis_labels=phi_vs_time_labels,
            color='blue',
            label=None,
            title=r'$\phi(t)$',
            ax=ax_a)

# second plot: phi_dot plot
ax_b = fig.add_subplot(1,3,2)

start, stop = start_stop_indices(t_pts, t_start, t_end)
plot_y_vs_x(t_pts[start : stop], phi_dot[start : stop],
            axis_labels=phi_dot_vs_time_labels,
            color='blue',
            label=None,
            title=r'$\dot\phi(t)$',
            ax=ax_b)

# third plot: state space plot from t=30 to t=50
ax_c = fig.add_subplot(1,3,3)

start, stop = start_stop_indices(t_pts, t_start, t_end)
plot_y_vs_x(phi[start : stop], phi_dot[start : stop],
            axis_labels=state_space_labels,
            color='blue',
            label=None,
            title='State space',
            ax=ax_c)

fig.tight_layout()
fig.savefig('simple_pendulum_Lagrange.png', bbox_inches='tight')
```

Now trying the power spectrum, plotting only positive frequencies and cutting off the lower peaks:

```python
start, stop = start_stop_indices(t_pts, t_start, t_end)
signal = phi[start:stop]
power_spectrum = np.abs(np.fft.fft(signal))**2
freqs = np.fft.fftfreq(signal.size, delta_t)
idx = np.argsort(freqs)

fig_ps = plt.figure(figsize=(5,5))
ax_ps = fig_ps.add_subplot(1,1,1)
ax_ps.semilogy(freqs[idx], power_spectrum[idx], color='blue')
ax_ps.set_xlim(0, 1.)
ax_ps.set_ylim(1.e5, 1.e11)
ax_ps.set_xlabel('frequency')
ax_ps.set_title('Power Spectrum')

fig_ps.tight_layout()
```

```python

```


# Linear First Order System

This notebook demonstrates simulation of a linear first-order system in Pyomo using two distinct approaches. The first uses the `Simulator` class from Pyomo which can employ the

```
%%capture
!pip install -q pyomo
```

## First-Order Differential Equation with Initial Condition

The following cell implements a solution to a first-order linear model in the form

\begin{align}
\tau\frac{dy}{dt} + y & = K u(t) \\
\end{align}

where $\tau$ and $K$ are model parameters, and $u(t)$ is an external process input.

```
% matplotlib inline
from pyomo.environ import *
from pyomo.dae import *
import matplotlib.pyplot as plt

tf = 10
tau = 1
K = 5

# define u(t)
u = lambda t: 1

# create a model object
model = ConcreteModel()

# define the independent variable
model.t = ContinuousSet(bounds=(0, tf))

# define the dependent variables
model.y = Var(model.t)
model.dydt = DerivativeVar(model.y)

# fix the initial value of y
model.y[0].fix(0)

# define the differential equation as a constraint
model.ode = Constraint(model.t, rule=lambda model, t: tau*model.dydt[t] + model.y[t] == K*u(t))

# transform dae model to discrete optimization problem
#TransformationFactory('dae.finite_difference').apply_to(model, nfe=50, method='BACKWARD')

# solve the model
#SolverFactory('ipopt').solve(model).write()

tsim, profiles = Simulator(model, package='scipy').simulate(numpoints=100)

plt.plot(tsim, profiles)

# access elements of a ContinuousSet object
t = [t for t in model.t]

# access elements of a Var object
y = [model.y[t]() for t in model.y]

plt.plot(t,y)
plt.xlabel('time / sec')
plt.ylabel('response')
plt.title('Response of a linear first-order ODE')
```

```

```


# Theory

### TODO

Explain why it is that we use likelihood (you can't just optimize the probability of values, we need to optimize some relationship between our model output and the training output)

Logistic regression is a technique that was developed for approximating the probability of a random variable Y taking a certain value given some experimental data.  It uses the logistic (a.k.a. sigmoid) function, $\theta(t) =  \frac{1}{1+e^{-t}}$, to map the results of a linear combination of the features to the range [0,1].  Summarily, we can express the probability of a random variable Y taking a value of $y_i$ given feature data $\vec{x}$ , as:  $$ p(Y=y_i|\vec{x}) = \frac{1}{1+e^{-(\vec{w}\cdot\vec{x} + b)}}$$

In order to generate this model we note that our hypothesis fits nicely with the idea of likelihood, which is a measure of how likely a model's parameters are given some sample data.  The likelihood of the parameters of our model, $\vec{w}$, given our training examples $\vec{x}$ can be written as $L(\vec{w} | \vec{x})$.  We note that the likelihood is then the probability of our training labels given our model's output for each training example.  Now, let us derive some things more formally for the case of binary classification.

Let us define the following terms:

- X : A set of training examples (assuming a 1 is added to each example as the bias)
- Y : The labels for our training examples
- $\vec{w}$ : The parameters for our logistic regression model
- $ g_{\vec{w}}(x) = \frac{1}{1+e^{-(\vec{w}\cdot x)}}$

As stated earlier, we can express $p(Y=1|X;\vec{w}) = g_{\vec{w}}(X)$.  Since this is a binary classification (A or not A), we can write $p(Y=0|X;\vec{w}) = 1 - g_{\vec{w}}(X)$.

Noting that this is a Bernoulli distribution, we can use the differentiable form, $*^1$, of its the probability mass function, $f(k;p) = p^k(1-p)^{1-k} $ , to estimate $P(y|X;\vec{w})$.  This yields $$P(y|X;\vec{w}) = g_{\vec{w}}(X)^y(1-g_{\vec{w}}(X))^{1-y}$$

With this equation we can express our likelihood as a function of our training data:

$\begin{align}
 L(\vec{w} | X) & = P(Y | X;\vec{w}) \\
 & = \prod_{i=1}^{N} P(Y = y_i | x_i;\vec{w}) \\
 & = \prod_{i=1}^{N} g_{\vec{w}}(x_i)^{y_i}(1-g_{\vec{w}}(x_i))^{1-y_i}
\end{align}$

Since we want to maximize likelihood and this equation would be a monster to differentiate, one often uses the log likelihood instead.  Using the properties of logarithms, we can express our log likelihood as:

$\begin{align}
 log(L(\vec{w} | X)) & = log(\prod_{i=1}^{N} P(Y = y_i | x_i;\vec{w})) \\
 & = \sum_{i=1}^{N} log(P(Y = y_i | x_i;\vec{w})) \\
 & = \sum_{i=1}^{N} log(g_{\vec{w}}(x_i)^{y_i}(1-g_{\vec{w}}(x_i))^{1-y_i}) \\
 & = \sum_{i=1}^{N} log(g_{\vec{w}}(x_i)^{y_i}) + log((1-g_{\vec{w}}(x_i))^{1-y_i}) \\
 & = \sum_{i=1}^{N} y_ilog(g_{\vec{w}}(x_i)) + (1-y_i)log(1-g_{\vec{w}}(x_i))
\end{align}$

Before determining the gradient (derivative with respect to the weights) of the entire equation, let us derive some of the embedded gradients :

$$\begin{align}
\frac{d}{d\vec{w}}g_{\vec{w}}(x_i) & = g_{\vec{w}}(x_i)(1-g_{\vec{w}}(x_i))\frac{d}{d\vec{w}}(\vec{w}\cdot x) \\
& = g_{\vec{w}}(x_i)(1-g_{\vec{w}}(x_i))x
\end{align}$$

And the derivative for a logarithm of some function $u(\vec{w})$:
$$\frac{d}{d\vec{w}}log(u) = \frac{u'}{u} $$

Using these two building blocks, and $\theta = g_{\vec{w}}(x_i)$ for the sake of brevity, we can easily solve for the derivate of our log likelihood with respect to the weights $\vec{w}$ yields:

$\begin{align}
\frac{d}{d\vec{w}}log(L(\vec{w} | X)) & = \\
& = \frac{d}{d\vec{w}}\sum_{i=1}^{N} y_ilog(\theta) + (1-y_i)log(1-\theta) \\
& = \sum_{i=1}^{N} \frac{d}{d\vec{w}}y_ilog(\theta) + \frac{d}{d\vec{w}}(1-y_i)log(1-\theta) \\
& = \sum_{i=1}^{N} y_i\frac{\theta'}{g_{\vec{w}}(x_i)} + (1-y_i)\frac{-\theta'}{1 - \theta}  \\
& = \sum_{i=1}^{N} y_i\frac{\theta(1-\theta)x_i}{\theta} + (1-y_i)\frac{-\theta(1-\theta)x_i}{1 - \theta}  \\
& = \sum_{i=1}^{N} y_i(1-\theta)x_i + (1-y_i)(-\theta)x_i  \\
& = \sum_{i=1}^{N} x_i(y_i-\theta)
\end{align}$

#### Footnotes

- $*^1$ : The bernoulli probability mass function is defined as a stepwise function that take a value of $p$ for one outcome and the complementary value, $1-p$, for the other outcome.  We note that $p^k(1-p)^{1-k}$ is a clever way of writing this same thing as a continuous function, since one of the exponents will be zero for each of the 2 possible outcomes, thus reducing the probability to $p$ or $1-p$.

## Implementation

## Binary Classifier

```python
import numpy as np
import math

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

class logistic_regression:

    def __init__(self, max_iterations=100, learning_factor=0.01):
        self._max_iterations = max_iterations
        self._learning_rate = learning_factor
        self._weights = []

    def predict(self, x):
        if len(self._weights) > 0:
            return 0 if sigmoid(np.dot(self._weights,np.append(x,1))) < 0.5 else 1
        else:
            raise Exception('The model has not been fitted, prediction is impossible')

    def fit(self, raw_features,labels):
        features = np.append(raw_features,np.ones((len(raw_features),1)),axis=1) #add bias
        self._weights = np.random.rand(features.shape[1])

        #calculate the gradient over all examples
        for epoch in range(self._max_iterations):
            predictions = sigmoid(np.dot(self._weights,features.T))
            error = labels - predictions
            gradient = (features.T * error).T.sum(axis=0)

            self._weights += gradient*self._learning_rate

#             print(np.sum(np.abs(error)))

```

```python
#This is mostly for my machine learning homework assignment
import numpy as np
from sklearn import datasets

iris_ds = datasets.load_iris()

features = iris_ds['data']
labels = iris_ds['target']

labels[labels > 1] = 1   #binary classification

LR = logistic_regression(500,0.001)

LR.fit(features,labels)

[[LR.predict(features[i]), labels[i]] for i in range(len(features)) if LR.predict(features[i]) != labels[i]]

```

    []


# Quickly Simulate an RBC Model

Here I demonstrate how how relatively straightforward it is to appoximate, solve, and simulate a DSGE model using `linearsolve`. In the example that follows, I describe the procedure more carefully.

\begin{align}
C_t^{-\sigma} & = \beta E_t \left[C_{t+1}^{-\sigma}(\alpha A_{t+1} K_{t+1}^{\alpha-1} + 1 - \delta)\right]\\
C_t + K_{t+1} & = A_t K_t^{\alpha} + (1-\delta)K_t\\
\log A_{t+1} & = \rho_a \log A_{t} + \epsilon_{t+1}
\end{align}

In the block of code that immediately follows, I input the model, solve for the steady state, compute the log-linear approximation of the equilibirum conditions, and compute some impulse responses following a shock to technology $A_t$.

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

# Input model parameters
parameters = pd.Series()
parameters['alpha']  = .35
parameters['beta']  = 0.99
parameters['delta']   = 0.025
parameters['rhoa'] = .9
parameters['sigma'] = 1.5
parameters['A'] = 1

# Funtion that evaluates the equilibrium conditions
def equilibrium_equations(variables_forward,variables_current,parameters):

    # Parameters
    p = parameters

    # Variables
    fwd = variables_forward
    cur = variables_current

    # Household Euler equation
    euler_eqn = p.beta*fwd.c**-p.sigma*(p.alpha*cur.a*fwd.k**(p.alpha-1)+1-p.delta) - cur.c**-p.sigma

    # Goods market clearing
    market_clearing = cur.c + fwd.k - (1-p.delta)*cur.k - cur.a*cur.k**p.alpha

    # Exogenous technology
    technology_proc = p.rhoa*np.log(cur.a) - np.log(fwd.a)

    # Stack equilibrium conditions into a numpy array
    return np.array([
            euler_eqn,
            market_clearing,
            technology_proc
        ])

# Initialize the model
model = ls.model(equations = equilibrium_equations,
                 n_states=2,
                 var_names=['a','k','c'],
                 shock_names=['e_a','e_k'],
                 parameters = parameters)

# Compute the steady state numerically
guess = [1,1,1]
model.compute_ss(guess)

# Find the log-linear approximation around the non-stochastic steady state and solve
model.approximate_and_solve()

# Compute impulse responses and plot
model.impulse(T=41,t0=5,shocks=None)

fig = plt.figure(figsize=(12,4))
ax1 =fig.add_subplot(1,2,1)
model.irs['e_a'][['a','k','c']].plot(lw='5',alpha=0.5,grid=True,ax=ax1).legend(loc='upper right',ncol=3)
ax2 =fig.add_subplot(1,2,2)
model.irs['e_a'][['e_a','a']].plot(lw='5',alpha=0.5,grid=True,ax=ax2).legend(loc='upper right',ncol=2)
```