# Fractal drum

Nori Parelius

This project was a part of an exam in Computational Physics that I took in May 2015. At the time I used Fortran to solve it, but I have since rewritten it in Matlab and now Python.

The project is about finding the eigenvalues and eigenvectors of a "fractal drum" - a thin membrane stretched on a frame with a fractal shaped border. The fractal is a square Koch fractal.

Oscillations of this drum's membrane follow the wave equation
\begin{equation}
\nabla^2 u=\frac{1}{v^2}\frac{\partial^2 u}{\partial t^2},
\end{equation}
with $u$ being the displacement, $v$ velocity and $t$ time. The value of $u$ at the boundary (where the membrane is attached to the frame) is always 0. Fourier transforming the wave equation over time gives the Helmholtz equation:
\begin{equation}
-\nabla^2U(\vec{x},\omega)=\frac{\omega^2}{v^2}U(\vec{x},\omega), \quad \rm{in}\ \Omega,
\end{equation}
where $U$ is the Fourier transform of $u$ and $\omega$ is angular frequency. And for the Dirichlet boundary condition $u=0$
\begin{equation}
U(\vec{x},\omega)=0, \quad \rm{on}\ \partial \Omega
\end{equation}

Using a finite difference approximation for the Laplacian operator, the Helmholtz equation is transformed into a form
\begin{equation}
LU=\frac{\omega^2}{v^2}U
\end{equation}
where $L$ is a matrix generated from the Laplacian operator and taking into account the boundary conditions and $U$ is a vector containing the displacements for each point within the boundary.
This equation is an eigenvalue problem which means that matrix $L$ can be used to find eigenvalues $\omega^2/v^2$ and corresponding eigenvectors $U$ for which the equation is true.

To find the eigenvalues, several steps have to be taken. First the shape of the fractal boundary has to be generated and positioned onto a spatial grid so that all the corners in the boundary fall onto grid points. Then it is necessary to find out which grid points are within the boundary and which are outside, as only the points inside will be a part of the vector $U$ from the eigenvalue problem. The matrix $L$ is then generated from the finite difference scheme for the Laplace operator and using the boundary conditions. The last step is to solve the eigenvalue problem.

## 1. Creating the fractal border on a grid

### Importing libraries

```python
import numpy as np
import scipy as sc
import matplotlib.pyplot as plt
%matplotlib inline
```

### Making the fractal shape
The fractal border is created by first defining a square, given by the coordinates of its corners. Function FRACTAL takes the coordinates of two points and creates a Koch square fractal from the line. It returnes the coordinates of the corners of this fractal. Function FRACTALIZE then uses function FRACTAL to do this for several connected lines, such as the starting square. The function FRACTALIZE applies the "fractalization" f_level times.

```python
# length of the side
L_side=4.0

# position of top left corner
LT=[5.0,9.0]

# number of times the initial square will be "fractalized"
f_level=2

```

```python
# Initializing the square
corners=np.zeros((5,2))
corners[0,:]=LT
corners[1,:]=[corners[0,0]+L_side,corners[0,1]]
corners[2,:]=[corners[0,0]+L_side,corners[0,1]-L_side]
corners[3,:]=[corners[0,0],corners[0,1]-L_side]
corners[4,:]=corners[0,:]
plt.plot(corners[:,0],corners[:,1])
```

```python
def fractal(A,B):
    '''
    FUNCTION FRACTAL takes as input the coordinates of two points A and B, 2D
    only, and they should either have the same x or y coordinate. And returns
    an array of points that give the Koch fractal shape
    '''

    f=np.zeros((9,2));
    f[0,:]=A
    f[8,:]=B

    if abs(A[0]-B[0]) > abs(A[1]-B[1]): # horizontal
        c=1.0
    else: # vertical
        c=-1.0

    # Points on the line connecting A and B
    step=1; # how far from A we are stepping, for first point 1 step, then 2
    for i in [1,4,7]: # which point in f it is
        for coor in range(2): # which coordinate, x or y
            f[i,coor]=A[coor]+step*(B[coor]-A[coor])/4.0
        step=step+1

    # Points on the line to one direction
    step=1;
    for i in range(2,4,1):
        f[i,0]=A[0]+step*(B[0]-A[0])/4.0+c*(B[1]-A[1])/4.0
        f[i,1]=A[1]+step*(B[1]-A[1])/4.0+c*(B[0]-A[0])/4.0
        step=step+1

    # Points on the line to the other direction
    step=2;
    for i in range(5,7,1):
        f[i,0]=A[0]+step*(B[0]-A[0])/4.0-c*(B[1]-A[1])/4.0
        f[i,1]=A[1]+step*(B[1]-A[1])/4.0-c*(B[0]-A[0])/4.0
        step=step+1
    return f
```

#### Example of a Koch fractal

```python
new_fractal=fractal(corners[0,:],corners[1,:])
plt.plot(new_fractal[:,0],new_fractal[:,1])
```

```python
def fractalize(corners,f_level):
    '''
    Function that takes an array containing coordinates of corners
    of a closed shape (first and last corner the same) and fractalizes each
    segment f_level number of times
    '''
    for k in range(f_level):
        n_sides=corners.shape[0]-1;
        corn=np.zeros((8*(n_sides-1)+9,2))
        corn[0,:]=corners[0,:]
        for i in range(1,n_sides+1):
            corn_new=fractal(corners[i-1,:],corners[i,:])
            corn[(8*(i-1)):(8*i+1),:]=corn_new
        corners=corn
    return corners
```

#### Fractalizing

```python
f=fractalize(corners,f_level)
fig, axes=plt.subplots(nrows=1,ncols=3,figsize=(12,3))
axes[0].plot(corners[:,0],corners[:,1])
axes[1].plot(fractalize(corners,1)[:,0],fractalize(corners,1)[:,1])
axes[2].plot(f[:,0],f[:,1])
for ax in axes:
    ax.set_aspect('equal')
```

### Grid to put the drum on
I created a square lattice that the border will be positioned on. This lattice is the minimum size necessary to cover the whole border structure. This means that the edge of the lattice is further out from the initial square box. This grid offset depends on $l$. For each $l$ the fractal border is moved outside by $L\_side/4$ where length is the length of the side at previous $l$. This adds up to $L\_side/4+L\_side/4^2+L\_side/4^3+...L\_side/4^f\_level$, where $L\_side$ is the length of the side of the initial box. One grid constant $\delta$ was added to the offset.

Grid constant $\delta$ was chosen depending on the desired $f\_level$. If each corner of the fractal border is supposed to fall onto a grid point for any $f\_level<f\_level_{max}$ then
\begin{equation}
\delta=\frac{L\_side}{4^{f\_level_{max}}}.
\end{equation}
In addition, for $f\_level=1$ with this grid spacing (meaning corners on the grid but no additional grid points between them), there would be only one point within the boundaries. For $f\_level>1$ this is not the case any more, but I have anyway defined a variable which allows to add grid points so that the grid is more dense than the corners of fractal shape.
\begin{equation}
\delta=\frac{L\_side}{4^{f\_level_{max}}(points\_between+1)},
\end{equation}
where $n_{points\_between}$ is the number of added points between, not including, the corners.
As a result, I had to make a routine that would add these points to my array of corners and create an array containing all the edge points.

```python
# Grid parameters
# number of grid points between corner points
points_between=1

# With each iteration, the border grows by 1/4 of L_side
grid_offset=np.sum(L_side/4**(np.arange(1,f_level+1)))

# Step of the grid, so that points between fall on it too
delta= L_side/(4**f_level)/(points_between+1)

# Origin of the grid, with the offset and with adding one point at the edge
grid_origin=[LT[0]-grid_offset-delta,LT[1]+grid_offset+delta]

# Number of grid points in both directions
n_grid=int((L_side+2*grid_offset+2*delta)/delta+1)
```

#### Adding more points between the corners, on the grid

```python
def fill_edges(corners,MP):
    '''
    Function that takes the corners and adds a given number (MP) of points
    between them, to fill out all the grid points that have an edge
    '''
    n_points=corners.shape[0]
    n_sides=n_points-1
    edges=np.zeros((n_points+n_sides*MP,2))

    for i in range(n_sides):
        for coor in range(2): # x y coordinates
            edges[((MP+1)*i):((MP+1)*(i+1)+1),coor]=corners[i,coor]+(corners[i+1,coor]-corners[i,coor])/(MP+1)*range(MP+2)
    return edges
```

```python
# Filling in the edges
edges=fill_edges(f,points_between)
plt.plot(edges[:,0],edges[:,1],marker='o')
```

## Identifying border, inside and outside grid points
To know which points are inside and which points are outside, I have created a mask, a two dimensional array of the same size as the grid, that for each point given its position in the array $mask(i,j)$ holds a value defining whether the point is in, out or on the border of the drum. The $mask$ for boundary points is -1, for points outside it is -2 and for points inside it holds the index of the point - from 0 to # points inside -1.

```python
# Grid
x=grid_origin[0]+delta*np.arange(n_grid)
y=grid_origin[1]-delta*np.arange(n_grid)
```

First function MASK_EDGES sets in the first step all the mask to -2. Next, using the coordinates(x, y) of the edge points the $mask$ is set to -1 for these edge points. The positions (i, j) of these points in the $mas$k were found by rearranging  the equation that generated the physical coordinates x and y,
\begin{equation}
x=A_x+(i-1)\delta \quad \rm{and}\ y=A_y-(j-1)\delta.
\end{equation}
Where $A_x$ and $A_y$ are the x and y coordinates of the top left corner of the grid.

```python
def mask_edges(n_grid,delta,LeftTop,edges):
    '''
    FUNCTION MASK_EDGES creates an array the size of the grid (n_grid) and
    puts its value to -1 everywhere, except for the points that belong to the
    edges, which are 0, this is calculated by using the (x,y) coordinates of the edges
    give in edges
    x=LTx+(i-1)delta --> i=(x-LTx)/delta +1
    y=LTy-(j-1)delta --> j=-(y-LTy)/delta +1
    '''
    me=-2*np.ones((n_grid,n_grid))
    for point in range(edges.shape[0]-1):
        me[int((edges[point,0]-LeftTop[0])/delta),
           int((LeftTop[1]-edges[point,1])/delta)]=0
    return me
```

Function INSIDE cycles through each of the points of the grid checking whether the point is inside the polygon.The algorithm consists of deciding how many times a horizontal line passing through the point in question crosses a side of the polygon on the left and on the right of the point. If the number of intersections to the left and to the right is odd, than the point is inside the polygon.

This procedure does not give clear results for the points that are right on the boundary, but these points are already known from the initial mask returned by MASK_EDGES.

```python
def point_in_polygon(e,P):
    '''
    FUNCTION POINT_IN_POLYGON takes as input the coordinates of the edge
    points and the coordinates of the point in question P and returns a
    boolean p depending on whether the point is inside the polygon or not.
    Using the number of times a horizontal line passing through the point
    crosses some segment
    '''
    x=P[0]
    y=P[1]

    p=False
    n_corners=e.shape[0]-1

    j=n_corners-1
    for i in range(n_corners):
        #it is horizontally between the end points
        if e[i,1]<y and e[j,1] >= y or e[j,1]<y and e[i,1]>=y:
            #whether x is right of the segment
            if e[i,0]+(y-e[i,1])/(e[j,1]-e[i,1])*(e[j,0]-e[i,0])<x:
                p=not p
        j=i
    return p
```

```python
def inside(edges,x,y,mask):
    '''
    FUNCTION INSIDE takes the edge points, the x and y points of the grid and
    the mask with the edges being 0, rest being -1. It uses point_in_polygon
    to find out if the point is in or out and then sets insides to numbers 1
    and upwards (counting them), edge to 0 and outside to -1
    '''
    n_grid=x.shape[0]
    me=np.zeros((n_grid,n_grid))
    n_inside=0
    for i in range(n_grid):
        for j in range(n_grid):
            # Find if it's in (1) or out (0)
            me[i,j]=point_in_polygon(edges,[x[i],y[j]])
            # If it's out, set it to -2,
            # if it's in, number it and count it
            if me[i,j]==0: # out
                me[i,j]=-2
                # If it corresponds to an edge point, set it to -1
                if mask[i,j]==-1:
                    me[i,j]=-1
            else: # in
                # If it corresponds to an edge point, set it to -1
                if mask[i,j]==-1:
                    me[i,j]=-1
                else:
                    me[i,j]=n_inside;
                    n_inside=n_inside+1;
    return me,n_inside
```

```python
from matplotlib import cm
# First all -1, edges 0
mask1=mask_edges(n_grid,delta,grid_origin,edges)
X, Y = np.meshgrid(x, y)
fig, ax = plt.subplots()
ax.pcolormesh(X,Y,mask1,cmap=cm.coolwarm)

```

```python
mask,n_in=inside(edges,x,y,mask1)
fig, ax = plt.subplots()
ax.contourf(X,Y,mask)
```

```python

```

### The Laplacian matrix
By using a Taylor expansion on the Laplace operator, we get
\begin{equation}
-\nabla^2u(x,y) \approx \frac{1}{\delta^2}(4u(x,y)-u(x-\delta,y)-u(x+\delta,y)-u(x,y-\delta)-u(x,y+\delta)).
\end{equation}
Applying this to the discretized x and y, with step $\delta$ leads to the 5-point central finite difference approximation of the Laplace, and this gives for our eigenvalue matrix equation ($LU=\frac{\omega^2}{v^2}U$):
\begin{equation}
4u_{i,j}-u_{i-1,j}-u_{i+1,j}-u_{i,j-1}-u_{i,j+1}=\omega^2/v^2\ \delta^2 u_{i,j}.
\end{equation}
This equation can be turned into a matrix equation
\begin{equation}
\frac{1}{\delta^2}L\ V=\frac{\omega^2}{v^2}V,
\end{equation}
where $V=[V_k]$ is a vector that holds values $u_{i,j}$ for i,j such that the corresponding point is within the fractal boundary, and $L$ is a matrix that has the value 4 on its diagonal and up to four times the value -1 in each row i, depending on whether the neigbouring points of $u_{i,i}$ are inside the boundary, or not.

This means that the matrix $L$ has to be constructed with the shape of the boundary taken into account.

To construct the matrix, I have created an array $n2grid$ that allows to turn the index of the drum point (the number held by the $mask$) into its position on the grid.

```python
# Vector containing the grid coordinates of each point on the inside,
n2grid=np.zeros((n_in,2))
for i in range(n_grid):
    for j in range(n_grid):
        if mask[i,j]>=0:
            n2grid[int(mask[i,j]),:]=[i,j]

```

```python
def matrix_laplace(mask,n2grid,n_in,delta):
    '''
    FUNCTION LAP_MAT returns the Laplacian matrix for the points inside the
    drum
    '''
    mat_lap=np.zeros((n_in,n_in))
    for i in range(n_in): # one row of the laplacian at a time
        mat_lap[i,i]=4.0/delta**2 # the diagonal is 4
        for ii in [-1,1]:
            if mask[int(n2grid[i,0]+ii),int(n2grid[i,1])]>=0: # if the neighbour at x+-1 is inside the drum
                mat_lap[i,int(mask[int(n2grid[i,0]+ii),int(n2grid[i,1])])]=-1.0/delta**2
            if mask[int(n2grid[i,0]),int(n2grid[i,1]+ii)]>=0: # same for y dir
                mat_lap[i,int(mask[int(n2grid[i,0]),int(n2grid[i,1]+ii)])]=-1.0/delta**2
    return mat_lap
```

```python
# Creating the Laplacian
laplace=matrix_laplace(mask,n2grid,n_in,delta)
```

### Solving the eigenvalue problem

```python
# Getting the eigenvalues and eigenvectors
eigenval,evec=np.linalg.eig(laplace)

# Sorting them from lowest to largest
idx = eigenval.argsort()
eigenval = eigenval[idx]
evec = evec[:,idx]
```

```python
# First five eigenvalues
print(eigenval[0:4])
```

    [4.27736565 8.42218947 8.42218947 8.74559167]

```python
U=np.zeros((n_grid,n_grid,5))
for n in range(5):
    for i in range(n_in):
        U[int(n2grid[i,0]),int(n2grid[i,1]),n]=evec[i,n]
```

### The plots of the first 5 vibrational modes
The following plots show the first five eigenvectors plotted onto the grid. These represent the displacement of each drum points, so we are basically looking at how the drum membrane would physically move. The real life vibration would be a combination of these and other modes.

```python
from mpl_toolkits.mplot3d import Axes3D

fig=plt.figure(figsize=(5,25))

X, Y = np.meshgrid(x, y)

# Plot the surface.
for i in range(5):
    ax[i]= fig.add_subplot(5,1,i+1, projection='3d')
    ax[i].plot_surface(X, Y, U[:,:,i],rstride=1,cstride=1, cmap=cm.YlOrRd,
                       linewidth=0, antialiased=False,alpha=0.5)
    ax[i].plot_surface(X,Y,0.001*mask1,rstride=1,cstride=1,cmap=cm.Blues,linewidth=0,
        antialiased=False,alpha=0.5)
    ax[i].set_xlabel("X")
    ax[i].set_ylabel("Y")
    ax[i].set_zlabel("U")

```

```python

```

```python

```


# Radar tracking EKF {-}

An aircraft is flying in a 2D coordinate frame with constant speed and altitude. The range to the arcraft is observed by radar.

Both speed and altitude are unknown.

State vector:
\begin{equation}
  x=
  \begin{bmatrix}
    pos\\
    vel\\
    alt\\
  \end{bmatrix}
  =
  \begin{bmatrix}
    x_1\\
    x_2\\
    x_3\\
  \end{bmatrix}
\end{equation}

Dynamic equation
\begin{equation}
  \begin{bmatrix}
    \dot{x_1}\\
    \dot{x_2}\\
    \dot{x_3}\\
  \end{bmatrix}
  =
  \begin{bmatrix}
    0 &1 &0\\
    0 &0 &0\\
    0 &0 &0\\
  \end{bmatrix}
  \begin{bmatrix}
    x_1\\
    x_2\\
    x_3\\
  \end{bmatrix}
  +
  \begin{bmatrix}
    0 &0\\
    \sqrt{q_v} &0\\
    0 &\sqrt{q_h}\\
  \end{bmatrix}
  u(t)
\end{equation}

Observation equation
\begin{equation}
  r=\sqrt{x_1^2 + x_3^2} + v
\end{equation}

```python
from sympy import Matrix, diff, symbols, sqrt

x1, x2, x3, r = symbols('x1 x2 x3 r')

# Measurement equation
r = sqrt(x1**2 + x3**2)

# Design matrix (linearized)
H = Matrix([[r]])
dH = H.jacobian([x1, x2, x3])
dH
```

$\displaystyle \left[\begin{matrix}\frac{x_{1}}{\sqrt{x_{1}^{2} + x_{3}^{2}}} & 0 & \frac{x_{3}}{\sqrt{x_{1}^{2} + x_{3}^{2}}}\end{matrix}\right]$

```python
from numpy import arange, array, eye, sqrt, zeros, random, diag
from numpy.linalg import inv, norm, cholesky
import matplotlib.pyplot as plt
from lib.vanloan import numeval

# System values
dt = 0.05       # [second]
samples = 400   # number of samples
num = 3         # number of iterations (design matrix)
qv = 0.01       # [meter^2/second^3]
qh = 0.01       # [meter^2/second]

# Dynamics matrix
F = array([[0, 1, 0],
           [0, 0, 0],
           [0, 0, 0]])

# White noise coefficients
G = array([[0, 0],
           [sqrt(qv), 0],
           [0, sqrt(qh)]])

# Observation covariance matrix
R = array([[10]])    # [meter^2]

# Initial state
x = array([[-1000],  # position
           [100],    # velocity
           [1000]])  # height
xt = x

# Inital state covariance matrix
P = diag([1, 0.1, 1])

# Numerical evaluation (van Loan)
[phi, Q] = numeval(F, G, dt)

# Cholesky decomposition of process noise covariance
C = cholesky(Q)

# Linearized design matrix
def dH(x_nom):

    x1 = x_nom[0, 0]
    x3 = x_nom[2, 0]

    dH = array([[x1/sqrt(x1**2 + x3**2), 0, x3/sqrt(x1**2 + x3**2)]])

    return dH

# Computed observations
def hx(x_nom):

    x1 = x_nom[0, 0]
    x3 = x_nom[2, 0]

    z = array([[sqrt(x1**2 + x3**2)]])

    return z

# Plot vectors
x_all = []; P_all = []; xt_all =[]
res = []; err = []

# Main loop
for k in range(0, samples):

    # Process noise vector
    w = array([[0],
               [random.normal(0, sqrt(qv))],
               [random.normal(0, sqrt(qh))]])

    # Compute true trajectory
    xt = phi@xt + C@w

    # Generate noisy measurements
    z = hx(xt) + random.normal(0, sqrt(R[0, 0]))

    # Time update
    xp = phi@x
    Pp = phi@P@phi.T + Q

    # Iterations
    for k in range(0, num):

        # Design matrix
        H = dH(x)

        # System size
        (n, e) = H.shape

        # Computed observations
        zp = hx(xp)

        # Kalman gain
        K = Pp@H.T@inv(H@Pp@H.T + R)

        # Measurement update
        x = xp + K@(z - zp)
        P = (eye(e) - K@H)@Pp@(eye(e) - K@H).T + K@R@K.T

    # Accumulate plot vectors
    x_all.append(x)
    P_all.append(P)
    xt_all.append(xt)

    res.append(norm(z - zp))
    err.append(norm(xt[0] - x[0]))

# Extract plot vectors
pos_est = [x[0] for x in x_all]
vel_est = [x[1] for x in x_all]
height_est = [x[2] for x in x_all]

pos_true = [xt[0] for xt in xt_all]
vel_true = [xt[1] for xt in xt_all]
height_true = [xt[2] for xt in xt_all]

std = [sqrt(P[0, 0]) for P in P_all]

# Time
time = arange(0, samples)*dt

plt.figure(1)
plt.plot(time, pos_est, 'g', label = 'Estimated trajectory')
plt.plot(time, pos_true, 'b', label = 'True trajectory')
plt.title('System State')
plt.xlabel('Time (second)')
plt.ylabel('Position (meter)')
plt.legend(loc='lower right')
plt.grid(True, which='both')
plt.show()

plt.figure(2)
plt.plot(time, pos_est, 'g.', label='Estimated trajectory')
plt.plot(time, pos_true, 'b', label='True trajectory')
plt.title('System State (0-500m)')
plt.xlabel('Time (second)')
plt.ylabel('Position (meter)')
plt.xlim(9, 11)
plt.ylim(-75, 75)
plt.legend(loc='lower right')
plt.grid(True, which='both')
plt.show()

plt.figure(3)
plt.plot(time, vel_est, 'g.', label='Estimated velocity')
plt.plot(time, vel_true, 'b', label='True velocity')
plt.title('System State')
plt.xlabel('Time (second)')
plt.ylabel('Velocity (meter/second)')
plt.ylim(90, 110)
plt.legend(loc='lower right')
plt.grid(True, which='both')
plt.show()

plt.figure(4)
plt.plot(time, height_est, 'g.', label='Estimated altitude')
plt.plot(time, height_true, 'b', label='True altitude')
plt.title('System State')
plt.xlabel('Time (second)')
plt.ylabel('Altitude (meter)')
plt.ylim(990, 1010)
plt.legend(loc='lower right')
plt.grid(True, which='both')
plt.show()

plt.figure(5)
plt.plot(time, res, 'r', label='Residual')
plt.plot(time, err, 'b', label='True error')
plt.plot(time, std, 'g', label='Estimated Standard Deviation')
plt.title('Kalman filter performance')
plt.xlabel('Time (second)')
plt.ylabel('Position (meter)')
plt.legend(loc='upper right')
plt.grid(True, which='both')
plt.show()
```


# Description of multiple equilibria in solution

Solution chemistry is still widely teaching in the first year post-bac. Typically in pH problems, you are looking for an approaching solutions using the dominant reaction approach. However, you must be aware than today, it's relatively easy to implement numerical algorithm to find the exact solution.

These numerical approaches are largely used to discuss the validity of the common approximations. Pedagogical software like ChimGéné able someone to simulate E-pH diagrams, i-E curves (<https://chimsoft.com/chimsoft/chimgene>), ...

Your goal is to understand how to describe a chemical equilibrium, without any approximation, using simple numerical recipes.

# A first approach : find the equilibrium of a system

## Weak acid equilibrium in water

$$
\begin{align} HA (aq) &= H^+ (aq) + A^-(aq) & K_A\\ H_2O(l) &= H^+(aq) + HO^- (aq) & K_e\end{align}
$$
**Q. 1** The system is describe by 4 degrees of liberty. Identify them. Write two equilibria and two conservation laws.

**R.1**

**Q.2** Combine the equation in order to get a third order polynomial which describe the problem.

**R.2**

In order to solve numerical the previous polynomial, called $\mathit{P}$, you might use robust numerical methods. One of the solvers is the dichotomy algorithm.

This algorithm consists in defining an interval $[x_1, x_2]$ where you want to find a solution :

$$
\mathit{P} (x_s) = 0 \quad \text{with } x_s \in [x_1, x_2] = [x_1^{(0)}, x_2^{(0)}]
$$
This n-step solution $x_s^{(n)}$ of the iterative method might converge to the final solution $x_s \pm \epsilon$ . $\epsilon$ is a user-controlled convergence criteria. The iterative transformation is defined as :

$$
\begin{align}
\forall n > 1, \quad &\begin{cases}
x_1^{(n)} &= \frac{x_1^{(n-1)} + x_2^{(n-1)}}{2} \\
x_2^{(n)} &= x_2^{(n-1)}
\end{cases} \text{ if } P(x_1^{(n-1)})\cdot P(x_1^{(n)})>0 \\
\quad &\begin{cases}
x_2^{(n)} &= \frac{x_1^{(n-1)} + x_2^{(n-1)}}{2} \\
x_1^{(n)} &= x_1^{(n-1)}
\end{cases} \text{ else }
\end{align}
$$
where $n$ is a given iteration step.

**Q.3** Interpret the transformation. Write a code using the dichotomy algorithm. Benchmark the function on a polynomial of your choice and plot the iterative process.

**R.3**

```
#R.3
```

*Note :* Implement conditions on the maximum number of iteration and on the value of $x_2$ compared to $x_1$.

Now, let's implement a second iterative method called *Regula falsi* method*.* The process is exactly the same than the one before, but the iterative transformation is different :

$$
\begin{align}

\forall n > 1, \quad &\begin{cases}
x_1^{(n)} &= x_1^{(n-1)} - \frac{x_1^{(n-1)} - x_2^{(n-1)}}{P(x_1^{(n-1)}) - P(x_2^{(n-1)})} P(x_1^{(n-1)})\\
x_2^{(n)} &= x_2^{(n-1)}
\end{cases} \text{ if } P(x_1^{(n-1)})\cdot P(x_1^{(n)})>0 \\
\quad &\begin{cases}
x_2^{(n)} &= x_2^{(n-1)} - \frac{x_2^{(n-1)} - x_1^{(n-1)}}{P(x_2^{(n-1)}) - P(x_1^{(n-1)})} P(x_2^{(n-1)}) \\
x_1^{(n)} &= x_1^{(n-1)}
\end{cases} \text{ else }
\end{align}
$$
**Q.4** Interpret the transformation. Add this transformation in the iterative solver presented in **Q.3**. Benchmark the function on the polynomial used in **Q.3** and plot the iterative process. Compare the convergence time of the two methods.

```
#R.4
```

**R.4**

**Q.5** Consider the acetic acid in water $pK_A (CH_3COOH/CH_3COO^-, 25 ^\circ C) = 4.8$ . Find the solution for three concentrations :

$$
c = 10^{-2}, 10^{-4}, 10^{-6} mol \cdot L^{-1}
$$
**R.5** Compare the results obtained from **Q.3** algorithm with the one obtain with the dominant reaction approximation.

## Another way to solve the problem

The first step is to implement a Newton-Raphson algorithm. This robust method consists in finding the roots of an $N$-variable real function $f$ :

$$
f(X_s) = 0
$$
In 1D, the transformation is easy to visualize. From the knowledge of the tangent at a given point, the cross-point between $y = 0$ and the tangent, allow you to approach the root of $f$.

**Q.6** Using the first order development of f around $x_n$ write a single step transformation $x_n \to x_{n+1}$

**R.6**

As all iterative method, convergence criteria might be defined :

1. $\lvert f(x_n) \rvert < \epsilon_1$, the residue in between your function and 0;
2. $\lvert x_{n+1} - x_n\rvert < \epsilon_2$, the resolution of your solution, estimating the accuracy of your estimation compared to the exact solution $\lvert x_1^{(\infty)} - x_s\rvert < \epsilon_2$ ;
3. a maximum number of iteration

**Q.7** Implement this transformation for 1D function as in **Q.3**. Benchmark the function on the polynomial used in **Q.3** and plot the iterative process. Compare the convergence time of the previous methods.

**R.7**

Chemical media are complex system described by :

* multiple equilibria ;
* multiple charge/mass conservation law .

In our simple weak acid example, system state belongs to the value of a single vector :

$$
X = \begin{pmatrix}
[AH] \\
[A^-] \\
[H^+] \\
[HO^-]
\end{pmatrix}
$$
The chemical system are always multicomponents system (at least, one solvent and one solute). So mathematically, any transformation $f$ of a chemical system is multidimensional. You must generalized the Newton-Raphson method in $N$ dimensions. The extension to multiple dimension can be summarize by this block diagram :

**Q.8** What is the name of the matrix $[\nabla f]$ ?

**R.8**

**Q.9** Using the system defined in **Q.1**, write the matrix $[\nabla f]$ for a system state described by $X$.

**R.9**

**Q.10** Develop a code using a N-dimension Newton-Raphson algorithm. Find the equilibrium state of acetic acid **Q.5** in the three cases.

```
#R.10
```

*Note :*  You can write the equilibria relations in log scale.

Before starting with more complex systems, you shall be aware of some limits of the previous approach. Classically, Newton-Raphson might be unstable if your system present too many degrees of freedom. Each result obtained must be interpreted as we say : "avec l'aide de votre $6^e$ sens chimique. Le plus important de tous !" .

# Treat a complex system with different kind of equilibria

## Single equilibrium state for one composition

Using the numerical recipes we study in the previous section, combined with a well-written object-oriented code, we are able to describe most of the single phase chemical equilibria.

Consider a well-known chemical system $Fe(II)/Ce(IV)$ composed of :

* 1L water solution,
* 0.1 mol of $FeCl_2$, 0.05 mol of $Ce(SO_4 )_2$ and 0.1 mol of $H_2SO_4$,
* restrict the system state to these species :

$$
H^+, HO^-, Cl^-, Fe^{3+}, Fe^{2+}, Ce^{4+}, Ce^{3+}, SO_4^{2-}, FeSO_4^+
$$
**Q.11** Write the 11 equations which describe the system.

**R.11**

**Q.12** Modify the code **Q.10** for an undermined N number of chemical species described by N (non-)linear equations. Find the equilibrium state of the previous system.

```
#R.12
```

*Data : Thermodynamic information for the Fe/Ce system in water at 25°C*

$$
\begin{align}
pK_e &= 14 \\
\log \beta (FeSO_4^+/Fe^{3+}) &= 2.03\\
E^\circ (Fe^{3+}/Fe^{2+}) &= 0.77 V/ESH \\
E^\circ (Ce^{4+}/Ce^{3+}) &= 1.71 V/ESH
\end{align}
$$
## Multiple equilibrium state for successive composition : titration

Titration reactions are necessarily process controlled by the thermodynamic (the reaction is fast) so they can be describe by the previous formalism. An experimenter prepares this experimental setup.

*Material :*

* $0.1 mol \cdot L^{-1}$ Mohr's salt solution , $0.1 mol \cdot L^{-1}$ titration solution of ceric disulfate,
* Volumetric Pipette, buret, magnetic stirrer, platin electrode, saturated calomel electrode,
* voltmeter.

*Protocol :*

Put $20 cm^{3}$ of ferric solution in a beaker and add a magnetic stirrer. Stire the solution and introduce both electrode on the solution. Connect it to a voltmeter. Fill the burette with the ceric disulfate reaction. Start the burette fall and acquire the potential.

To treat this problem numerically, one might consider it as succession of N equilibria. The complete burette fall is considered in the $[0; 2V_e]$ interval volume.

The volume increment between two equilibria corresponds to equal volume burette fall, defined as :

$$
\delta V = \frac{2V_{eq}}{N-1}
$$
where $V_{eq}$ is equivalence volume, $N$the number of point in the curve. The iterative method can be simplified by the block diagram below :

**Q.13** Calling the equilibrium solver developed in **Q.12**, developed an iterative method to obtain the curve :

$$
E = f(V) \quad [\emptyset]
$$
Save the volume, the electrochemical potential and the concentration of each chemical species in a npz file.

```
#R.13
```

**Q.14** Write a post analysis script which import the outputs. Check the electric potential profile to what you are expecting and benchmark your code.

Implement four methods to determine the volume of end of the titration :

* tangent method;
* first derivative method;
* second derivative method;
* Gram method.

Compare the value obtained with the one expected. Discuss the interest and precision of each method.

```
#R.14
```

*Note :* You might have data with different values of $N$ to get an interesting discussion.

**Q.15**  How the pH changing during the titration process ? Add an acid buffer and compare the evolution.

```
#R.15
```

**R.15**

*Data :* Complementary thermodynamic information in water at 25°C

$$
\begin{align}
pK_A (H_2SO_4 / HSO_4^-) & < 0 \\
pK_A (HSO_4^- / SO_4^{2-}) & = 2.0
\end{align}
$$
# An original approach : Use semantic tableau

Now, you get some ideas about solution model. Your goal in this section is to implement a more elegant methods. One interesting construction is to create a logic process to get a semantic tableau.

Let's imagine you are studying a fix concentration of iron in a solution with one degree of oxidation : $III$. We choose to describe the solution as composed by :

$$
H^+, Fe^{3+}, HO^-, Fe(OH)_1^{2+}, Fe(OH)_2^{+}, Fe(OH)_4^{-} \quad (\star)
$$
with these equilibria constants in water at 25°C :

$$
\begin{align}
K_e &= 10^{-14} \\
\log \beta_1 &= 11.81 \\
\log \beta_2 &= 22.33\\
\log \beta_4 &= 34.4
\end{align}
$$
**Q.16** How many concentrations an experimenter should know to determine the concentration of all the species using the mass action law. Why ?

**R.16**

Now, you must simplify the problem !  In order to construct an easiest mathematical problem, the number of variables will be reduce from 6 to 2. We choose the pH (so $[H^+]$) and $[Fe^{3+}]$ as our basis set to generate all the configurations.

**Q.17** Using the equilibrium constants, rewrite the concentration of the 6 species $(\star)$ in such a way :

$$
\forall i \in [1, 6], C_i = [H^+]^{n_H} \cdot[Fe^{3+}]^{n_{Fe}} \cdot 10^{-pK}
$$
where $n_H$, $n_{Fe}$ and $pK$ are three parameters to determine.

**R.17**

**Q.18** Using **Q.17**, complete this tableau :

What is the meaning of a row ? What is the meaning of the first two columns $H_T$ and $Fe_T$.

**R.18**

**Q.19** Now we want to convert the problem in a simple matrix problem where the unknown component is **X** :

$$
X = \left( \log [H^+] \log[Fe^{3+}] \right)
$$
From the tableau in **Q.18**, write :

* a vector (6,) containing the log species concentrations **(C);**
* a vector (6,) containing the reactions constants **(pK);**
* a matrix (6,2) containing the stoichiometric coefficients **(A)**;
* a vector(2,) containing the conservation law constrain **(T)**.

Write **C** as a function of **pK**, **A** and **X**.

**R.19**

To solve the problem, we want to minimize a relation. The mass balance is an ideal quantity to minimize in this problem :

$$
R = {^t A} \cdot 10^C - T
$$
where **T** is the mass contrain, **$^t A$** the transposition of stoichiometric coefficient, **T** the conservation law constrain.

**Q.20** What does mass balance represent ? Using the previously Newton-Raphson method, minimize the residuals in the mass balance. For a total E-5 M of iron, find the reconstruct the distribution of iron species as a function of pH.

 *Note :* To minimize **R** as a function of **X**, you need to find the jacobian and fix an origin.

**R.20**

```
#R.20
```

# Open Question

**R.21** Generalize the semantic tableau method for multiphasic system. You can add $Fe(OH)_3(s)$ to extend the previous problem for higher concentration in iron.

```
#R.21
```

*Data :* $pK_s(Fe(OH)_3 (s), 25 ^\circ C) = 15.09$

*Note :* There is as many tableau as there is phases ...

This project is derived from the architecture of *Chimie des solutions : De l'élémentaire aux calculs numériques, Marc Blétry & Marc Presset, Chapter 12 : Modélisation numérique des équilibres en solution idéale.*

To complete, these resources are used :

*L'Oxydoréduction : Concepts et expériences*, Jean Sarrazin & Michel Verdaguer.

*Chemical Reaction Equilibrium Analysis Theory and Algorithms*,  William R. Smith and Ronald W. Missen

*Solution of simultaneous chemical equilibria in heterogeneous systems : implementation in Matlab*, Scott Smith, 23 September 2019


# Multiresolution Chemistry made easy (through Vampyr)

## 1 Vampyr prologue

import the packages like so

```python
from vampyr import vampyr3d as vp
import matplotlib.pyplot as plt
import numpy as np
```

If you want to find out how to use a function, or what we have in the package use the function below

```python
help(vp.ScalingProjector)
```

In case the help was not enough, ask **us** for help. For this you need only your voice.

## 2 Math in Vampyr

The following exercises are introductions in working with functions, application of operators and some equation solving.

### 2.1 Projecting a function

In this task you will create a function, project it at different precisions and basis orders, and plot it against its projection.
This is meant to give you understanding in how one can approach completeness in multiwavelets.

First we start by initializing the multiresolution analysis (mra) by using the following command

```python
mra = vp.MultiResolutionAnalysis(box=[-20,20], order=7)
```

Here you will change the `order` value later. This defines the polynomial order of the basis.

An allowed function for projection in vampyr needs to have a single coordinate argument and it has to output a float value. You will construct a gaussian function defined by
\begin{align}
f(\mathbf{r}) = \beta \exp(-\alpha[r_x^2 + r_y^2 + r_z^2])
\end{align}
the alpha value is a measure of how sharp the gaussian function is, and the beta is just a normalization factor, for this exercise you can set it to anything.
Fill out the following python function for this

```python
def f(r):
    # implement your function here
    return

```

You should then be able to plot this function by running the next cell

```python
r_x = np.linspace(-0.99, 0.99, 1000) # create an evenly spaced set of points between -0.99 and 0.99
r_y = np.zeros(1000)
r_z = np.zeros(1000)
r = [r_x, r_y, r_z]

plt.plot(r_x, f(r))

```

Now we want to project this function, the way we do that is by creating a projection operator `vp.ScalingProjector` and apply it on the function. We apply this operator by using the function as a function argument of it, `P(f)` for example.

Try creating one and project the function you created earlier.

hint: use the `help` function to fill out the blanks)

```python
### write your code here
P_eps =

f_tree =    # <-- here is where you will put the projected function
```

when you are done you can run the following cell to plot your projected function.

```python
f_plt = [ f_tree([x, 0.0, 0.0]) for x in r_x ]
plt.plot(r_x, f(r), "tab:blue") # your function
plt.plot(r_x, f_plt, "tab:red") # your projected function
plt.show()
```

Now that you have projected the function, try changing three values, the `alpha` value of your original function, the `order` of the basis and the precision (`prec`) at which you project the function. What you should see is that the sharper the function, the higher order and precision you need to represent it properly.

One last thing I need to point out before we go to the next exercise is that it is convention to call our projected functions trees, as they all are `vp.FunctionTree` type.

### 2.2 Arithmethic

Arithmethic in vampyr is mostly the same as in normal python.
Lets work with the previous function you created.

First you can normalize your projected function by calling the function `normalize()` of `f_tree`.
Try it in the next cell

```python
# normalize your function here

```

Before we go on into the rest of the exercise, it is important to note that whenever a function is shown, and it doesn't start with `vp.`, it is assumed that it must be called as a member function of the tree, that is `f_tree.normalize()`, otherwise all functions must be called by using the `vp.` prefix.

Start by multiplying it by a factor of $2$ and assigning it to a new tree `g_tree`

```python
g_tree =  # <--- assign your rescaled function to this
```

You should be able to compare it against `f_tree` by runnning the following cell

```python
g_plt = [ g_tree([x, 0.0, 0.0]) for x in r_x ]
plt.plot(r_x, g_plt, "tab:red")
plt.plot(r_x, f_plt, "tab:blue")
plt.show()
```

try adding `f_tree` and `g_tree` together

```python
h_tree =  # <--- assign your added functions to this
```

again, you can test this by plotting it against the two other functions by running the following cell

```python
h_plt = [ h_tree([x, 0.0, 0.0]) for x in r_x ]

plt.plot(r_x, h_plt, "tab:orange")
plt.plot(r_x, g_plt, "tab:red")
plt.plot(r_x, f_plt, "tab:blue")
plt.show()
```

You can also check that the norm of the different trees makes sense. This is done by calling the `norm()` value of your tree, try printing it below using the print function. The `f_tree` should be normalized, so its norm should be about $1.0$

```python
print() # print the norm of your functions
```

Try multiplying your trees with the `*` operator, it should return a new tree which is the product of both, and you can check its norm again.

```python
i_tree =  # <----- assign you multiplied functions to this
```

Finally, you can take the dot product, also called inner product, between two trees by using the `vp.dot` function. You can try this in the cell below. print out the norm of the product as shown above, is it what you would expect?

```python
# take the dot product of f_tree and g_tree
```

## 2.3 Operators

You have already worked with the `vp.ScalingProjector` and will work with the `vp.HelmholtzOperator`. In this part we will work with the Derivative operator, defined as `vp.ABGVDerivative`.
Try constructing (with the help of `help()`) and applying the operator to your `f_tree`

```python
D =  # <--- initialize the ABGVDerivative operator here
df_tree =  #  <-- set your differentiated f_tree here
```

Run the cell below to plot your differentiated function

```python
df_plt = [ df_tree([x, 0.0, 0.0]) for x in r_x ]
plt.plot(r_x, df_plt, "tab:blue") # your differentiated function
plt.plot(r_x, f_plt, "tab:red") # your function tree
plt.show()
```

```python

```

## 3 Quantum chemistry in Vampyr - The Hydrogen atom

In order to solve the KS equations in MWs we reformulate them in an integral form [1].

\begin{align}
    &\phi_i = -2\hat{G}_{\mu_i}\hat{V}\phi_i \\
\end{align}

Where $\hat{V}$ is the potential acting on the system, $\phi_i$ is the i-th KS orbital,   $\hat{G}$ is the Helmholtz integral operator, where its kernel is defined as
\begin{align}
G_\mu(r - r') = \frac{\exp(-\mu |r - r'|)}{4\pi |r - r'|}
\end{align}
and $\mu$ is a parameter defined above through the orbital energies. The kinetic and orbital energies are included in $\hat{G}$.

The Helmholtz operator is already implemented in vampyr, therefore the only things you need are the integral KS equation and the definition of $\mu$

\begin{equation}
     \mu_i = \sqrt{-2\epsilon_i}
\end{equation}

The way you initialize the Helmholtz operator is as follows
```
H = vp.HelmholtzOperator( mra, exp=mu, prec=eps )
```
where `mu` is the $\mu$ is the parameter defined above, mra you have seen before, and `eps` is the desired threshold precision. This operator is applied the same way you applied the `vp.ScalingProjector` earlier.

In this exercise you will be solving the KS equation iteratively for a simple system, the Hydrogen atom. This means that you only have the nuclear potential to take into account for the potential term in the KS equation.

\begin{align}
    V_{nuc}(\mathbf{r}) = -\frac{1}{|\mathbf{r}|}
\end{align}

We will also be working with a single orbital, of which the initial guess is

\begin{align}
    \phi_0(\mathbf{r}) = e^{-|\mathbf{r}|^2}
\end{align}
where
\begin{align}
    |\mathbf{r}| = \sqrt{x^2 + y^2 + z^2}
\end{align}

The orbital update is defined as follows
\begin{align}
    \Delta\tilde{\phi}^n &= -2\hat{G}[V_{nuc}\phi^n] - \phi^n \\
    \Delta\tilde{\phi}^n &= \tilde{\phi}^{n+1} - \phi^n
\end{align}
where we use \~ to denote a function that is **not** normalized, and $n$ is the iteration index.

The energy update is then defined by
\begin{align}
    \Delta \epsilon^n = \frac{<\Delta\tilde{\phi}^n|V\phi_n>}{||\tilde{\phi}^{n+1}||}
\end{align}

What you need to do in this exercise is
1. Make a nuclear potential as a python function
2. Make an initial guess for the orbital as a python function
3. Project both nuclear potential and orbital through vampyr
4. Normalize orbital
5. Create a Helmholtz operator with $\mu$ as shown above, use an initial guess for $\epsilon$ (-0.5 is a good starting guess)
6. Compute the orbital update through application of the Helmholtz operator
7. Compute the energy update
8. Update both orbital and energy
9. Repeat steps 4-8 until your orbital energies have converged

The convergence criterion is the norm of $\Delta \phi^n$, but you should start by looping a set amount of times before trying the threshold.

```python
from vampyr import vampyr3d as vp
import numpy as np
import matplotlib.pyplot as plt

r_x = np.linspace(-0.99, 0.99, 1000) # create an evenly spaced set of points between -0.99 and 0.99
r_y = np.zeros(1000)
r_z = np.zeros(1000)
r = [r_x, r_y, r_z]

# Analytic nuclear potential
def f_nuc(r):
    #implement the nuclear potential
    return

# Analytic guess for solution
def f_phi(r):
    # implement the initial guess for the orbital
    return

#implement the starting guess

i = 0
while (i < 10):
    # implement the SCF

    # this will plot each of your orbitals, ucomment the lines below to do so
    #phi_n_plt = [phi_n([x, 0.0, 0.0]) for x in r_x]
    #plt.plot(r_x, phi_n_plt)

    # you should also print your energies among other things, uncomment the line below to do so
    #print(E_n)
    i += 1

plt.show()
```

## 4 Sources

[1] Stig Rune Jensen, Santanu Saha, José A. Flores-Livas, William Huhn, Volker Blum, Stefan Goedecker, and Luca Frediani The Elephant in the Room of Density Functional Theory Calculations. The Journal of Physical Chemistry Letters 2017 8 (7), 1449-1457
DOI: 10.1021/acs.jpclett.7b00255


# Smooth FIR
Smooth FIR model (see [Goutte et al., 2000](https://ieeexplore.ieee.org/document/897811)).

```python
import numpy as np
import matplotlib.pyplot as plt
```

Smooth FIR is like Tikhonov regression with a following multivariate normal prior on the parameters:

\begin{equation}
\Sigma_{ij} = \nu \exp(- \frac{h}{2} (i - j) ^ 2)
\end{equation}

The smoothness parameters, $h$, is often set as:

\begin{equation}
h = \frac{1}{\sqrt(\frac{7}{\mathrm{sp}})}
\end{equation}

where $\mathrm{sp}$ is the sampling period (TR in fMRI lingo).

In code:

```python
sp = 2  # sampling period
h = (1 / np.sqrt(7 / sp))  # smoothness factor
λ = 1   # regularization strength
p = 20   # number of parameters

sigma = np.zeros((p, p))
for i in range(p):
    for j in range(p):
        sigma[i, j] = np.exp(- (h / 2) * (i - j) ** 2)

plt.imshow(sigma)
plt.colorbar()
```

```python
n = 100
X = np.random.normal(0, 1, size=(n, p))
b = np.random.multivariate_normal(np.zeros(p), cov=sigma)
y = X @ b + np.random.normal(0, 7, n)
b_ols = np.linalg.inv(X.T @ X) @ X.T @ y

λ = 100
b_sf = np.linalg.inv(X.T @ X + λ * np.linalg.inv(sigma)) @ X.T @ y

plt.figure(figsize=(15, 5))
plt.plot(b)
plt.plot(b_ols)
plt.plot(b_sf)
plt.xlim(-1, p)
plt.legend(['True β', 'OLS β', 'sFIR β']);
```


# Section 2

This notebook replicates Figure 2.1 in your textbook. The goal is to understand the Figure and the parameters. Then, you can play with them and analyze what happends with the equilibrium when you move them.

```python
import collections
from pylab import *
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import sympy as sy
```

## Description:

$p_{Ch}$ and $q_{Ch}$ are the price and quantity of tea consumed in China, respectively. The inverse demand function for tea in China is

$p_{Ch} = 20 - q_{Ch}$, where 20 is the choke price for China

$p_{US}$ and $q_{US}$ are the price and quantity of tea consumed in the U.S., respectively. The total quantity of tea, $q_T$, is 10. Then, the quantity in the U.S., $q_{US}$, is total quantity minus the quantity in China, $q_{Ch}$.

$q_{US} = 10 - q_{Ch}$

The inverse demand function for tea in the U.S. is

$p_{US} = 18 - q_{US} = 18 - (10 - q_{Ch})$, where 18 is the choke price for the U.S.

## Parameters:

You can change the parameters below to analize what happens with the equilibrium.

```python
aCh = 20 # choke price for China
aUS = 18 # choke price for the U.S.
qT = 10 # total quantity
b = 0 # transportation cost
```

We consider iceberg transportation cost. For instance, $b=0.3$ implies tranportation costs equals to 30\% of the purchase price. How do transportation costs affect the equilibrium price and quantity in the two markets?

## Equilibrium:

First, we calculate the price and equilibrium quantities. If you change the parameters above the equilibrium will change. Don't change the code below!

```python
# Equilibrium
EqCh = (qT + ((1+b)*aCh) - aUS)/(2+b) # quantity in China
EqUS = qT - EqCh # quantity in the U.S.
EpCh = aCh - EqCh # price
EpUS = (1+b)*EpCh

```

```python
print("The equilibrium price in China is", "{:.2f}".format(EpCh),
      "The equilibrium price in the U.S. is", "{:.2f}".format(EpUS),
      ",the equilibrium quantity in China is", "{:.2f}".format(EqCh),
      " and the equilibrium quantity in the U.S.(transportation-inclusive) is", "{:.2f}".format(EqUS),
      sep=" ")
```

    The equilibrium price in China is 14.00 The equilibrium price in the U.S. is 14.00 ,the equilibrium quantity in China is 6.00  and the equilibrium quantity in the U.S.(transportation-inclusive) is 4.00

## Figure 2.1

The code below replicates Figure 2.1 in your texbook. Don't change it!

```python
qCh = np.linspace(0, qT, 10) # quantity in China
pCh = aCh - qCh # price in China
qUS = np.linspace(0, qT, 10)  # quantity in the U.S.
pUS = (1/(1+b))*(aUS - (qT - qCh)) # price in the U.S.

fig, ax1 = plt.subplots()
ax2 = ax1.twinx()
ax1.plot(qCh, pCh, 'g-')
ax2.plot(qUS, pUS, 'b-')
ax1.plot(EqCh, EpCh, 'o', markersize = 10)

ax1.set_xlabel('tea')
ax1.set_ylabel('price')
ax2.set_ylabel('price')

ax1.set_ylim([0,25])
ax2.set_ylim([0,25])

title('Arbitrage')
show()

```


# Introduction to Programming and Numerical Analysis: Model Project

```python
# We import all the necessary packets at the beginning of our code:

import numpy as np
from scipy import linalg
from scipy import optimize
import sympy as sm
sm.init_printing(use_unicode=True)

%matplotlib inline
import matplotlib.pyplot as plt
plt.style.use('seaborn-whitegrid')
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D

import ipywidgets as widgets
```

## **Basic OLG**

### **Introduction & household optimization**

The Overlapping Generations Model (OLG) is a class of models in order to study the development of an economy, mostly economic growth. OLG models are characterized by discrete time periods such that in every period at least two different generations live. In the case we study here, there are two generations per period; a young generation that works and an old generation that is retired (meaning that the individuals themselves also only live for two periods, however they overlap with each others life cycles). They all consume in both periods but they only save in their life-time period 1 as they die at the end of period 2 and use all their resources before doing so. Hence, they have a labor income in their 1st period and a capital income in their 2nd period.

- **Beginning of period t:** generation t born
- **Period t:** generation t works
- **End of period t:** genertaion t consumes and saves
- **Period t+1:** generation t lends
- **End of period t+1:** generation t consumes and dies

We will first solve the general case of the model and then complement the model by including the government, which raises taxes or issues debt-bonds. By that, we show that the model is applicable in a much broader setting and therefore also of use to answer more sophisticated questions.

First of all, we have the following CRRA life-time utility function:

\begin{eqnarray*}
U_t=\frac{C_{1t}^{1-\theta}}{1-\theta} + \frac{1}{(1+\rho)}\frac{C_{2t+1}^{1-\theta}}{1-\theta}
\end{eqnarray*}

However, due to simpliticity, but also as theta is not significantly different from 1 in many studies, we take

\begin{eqnarray*}
\theta=1
\end{eqnarray*}

Consequently, we need then to logarithmize and get:

\begin{eqnarray*}
U_t=ln(C_{1t})+ \frac{1}{(1+\rho)}ln(C_{2t+1})
\end{eqnarray*}

Furthermore, households have the following budget constraint:

\begin{eqnarray*}
A_tw_t=C_{1t}+ \frac{C_{2t+1}}{1+r_{t+1}}
\end{eqnarray*}

which comes from the single period constraints:

\begin{eqnarray*}
C_{1t}=A_tw_t-S_t
\end{eqnarray*}

and

\begin{eqnarray*}
C_{2t+1}=(1+r_{t+1} )S_t
\end{eqnarray*}

- $U_t$ = life-time utility
- $C_1$ = consumption as young (in period t)
- $C_2$ = consumption as old (in period t+1)
- $\rho$ = discount rate on consumption (>-1)
- $\theta$ = risk aversion parameter
- $A_t$ = total factor productivity in period t
- $w_t$ = wage rate
- $S_t$ = savings in period t
- $r_{t+1}$ = interest rate on savings (technicially paid at the end of period t+1)

```python
# Define the first sympy variables and parameters:

Ut = sm.symbols('U_t')
C1t = sm.symbols('C_1t')
C2t1 = sm.symbols('C_2t+1')

alpha = sm.symbols('alpha')
rho = sm.symbols('rho')

rt1 = sm.symbols('r_t+1')
wt = sm.symbols('w_t')

At = sm.symbols('A_t')
St = sm.symbols('S_t') # total savings
s = sm.symbols('s(r_t+1)') # saving rate
```

```python
# Define the utiliy function as sympy:

utility_function = sm.ln(C1t)+(1/(1+rho))*sm.ln(C2t1)
print('Utility Function:')
display(sm.Eq(Ut,utility_function))
```

```python
# Define the budget constraints as sympy:

budget_constraint_p1 = sm.Eq(C1t,wt*At-St)
budget_constraint_p2 = sm.Eq(C2t1,(1+rt1)*St)

St_from_bcp2 = sm.solve(budget_constraint_p2,St)
budget_constraint_by_sympy = budget_constraint_p1.subs(St,St_from_bcp2[0])

# Define budget constraint as in the outline:

budget_constraint = sm.Eq(C1t+(C2t1/(1+rt1)),wt*At)

# Showing the results:
print('Budget Constraints 1 & 2 as defined above:')
display(budget_constraint_p1)
display(budget_constraint_p2)
print('Combine the constraints:')
display(budget_constraint_by_sympy)
print('Reformulate:')
display(budget_constraint)
```

```python
# Optimum value for C1t by inserting budget constraint into the utility function and differentiating:

C2t1_from_con = sm.solve(budget_constraint,C2t1)
objective_subs = utility_function.subs(C2t1,C2t1_from_con[0])
foc1 = sm.diff(objective_subs,C1t)
sol1 = sm.solve(sm.Eq(foc1,0),C1t)
opt_C1t = sm.Eq(C1t,sol1[0])

# Optimum value for C2t1 by inserting budget constraint into the utility function and differentiating:

C1t_from_con = sm.solve(budget_constraint,C1t)
objective_subs = utility_function.subs(C1t,C1t_from_con[0])
foc2 = sm.diff(objective_subs,C2t1)
sol2 = sm.solve(sm.Eq(foc2,0),C2t1)
opt_C2t1 = sm.Eq(C2t1,sol2[0])

# We can then straightforward derive the Euler-equation (C2t1/C1t) by taking together the two latter results.
# The Euler-equation has the advantage that it shows the relative change in consumption and not the level.

euler = sol2[0]/sol1[0]
euler_equation=sm.Eq(C2t1/C1t,euler)

# Showing the results:
print('optimal Consumption in period 1:')
display(opt_C1t)
print('optimal Consumption in period 2:')
display(opt_C2t1)
print('Euler Equation:')
display(euler_equation)
```

```python
# In order to get total savings, we only have to rearrange the budget constraint from period 1 and insert optimal consumption we just derived:

savings_period1 = sm.solve(budget_constraint_p1,St)
savings_period1_fullequation = sm.Eq(St,savings_period1[0])
savings_formula = savings_period1_fullequation.subs(C1t,sol1[0])
savings_formula_simplified=sm.simplify(savings_formula)

# We print all the steps and the result:
print('Saving Equation:')
display(savings_period1_fullequation)
print('Insert optimal first period consumption:')
display(savings_formula)
print('Simplify:')
display(savings_formula_simplified)
```

```python
# Having calculated total savings, we can also calculate the saving rate:

total_savings=sm.Eq(St ,s*wt*At)
saving_r = sm.solve(total_savings,s)
saving_r2 = sm.Eq(s ,saving_r[0])
saving_rate_equation = saving_r2.subs(St,sm.simplify(savings_period1[0].subs(C1t,sol1[0])))

print('Total savings:')
display(total_savings)
print('Reformulate:')
display(saving_r2)
print('Insert Saving Equation and simplify:')
display(saving_rate_equation)
```

```python
# For later, we define the saving rate
# as saving_rate_equation shows the whole equation for illustration reasons, but we only need the right-hand term).

saving_rate = 1/(2+rho)
print('Saving Rate: s=')
display(saving_rate)
```

### **Production**

So far we have just introduced the households. We consider the case where firms have a standart neoclssical Cobb-Douglas production function:

\begin{eqnarray*}
Y_t=F(K_t,A_tL_t)=(K_t)^\alpha(A_tL_t)^{1-\alpha}
\end{eqnarray*}

where total factor productivity as well as population grow with a certain rate:

\begin{eqnarray*}
A_t=(1+g)A_{t-1}
\end{eqnarray*}

\begin{eqnarray*}
L_t=(1+n)L_{t-1}
\end{eqnarray*}

- $K_t$ = capital stock in period t
- $A_t$ = total factor productivity in t
- $L_t$ = amount of labour in t
- $\alpha$ = capital share in production
- $g$ = growth rate of total factor productivity
- $n$ = growth rate of population

As we use the terms of effectice labor, it is however much easier to exert the intensive form of production. Hence we define:

\begin{eqnarray*}
\frac{1}{A_tL_t}F(K_t,A_tL_t)=F(\frac{K_t}{A_tL_t},1)
\end{eqnarray*}

and in this particular case

\begin{eqnarray*}
y=f(k_t)=k_t^\alpha
\end{eqnarray*}

with

\begin{eqnarray*}
k_t=\frac{K_t}{A_tL_t}, y_t=\frac{Y_t}{A_tL_t}, f(k_t)=F(k_t,1)
\end{eqnarray*}

We further assume that factor markets are perfektly competitive, i.e.:

\begin{eqnarray*}
r_t=f'(k_t)
\end{eqnarray*}

\begin{eqnarray*}
w_t=f(k_t)-f'(k_t)k_t
\end{eqnarray*}

```python
# We define the missing corresponding sympy variables:

yt = sm.symbols('y_t')
kt = sm.symbols('k_t')
rt = sm.symbols('r_t')
```

```python
# Define the production function:

f = kt**alpha
print('Production Function:')
display(sm.Eq(yt,f))
```

```python
# Calculate the real interest rate:

real_interest_rate = sm.diff(f,kt)
real_interest_rate = sm.simplify(real_interest_rate)
print('Real Interest Rate:')
display(sm.Eq(rt,real_interest_rate))
```

```python
# Calculate the wage rate per unit of effective labor:

wage_rate = f-sm.diff(f,kt)*kt
wage_rate = sm.simplify(wage_rate)
print('Wage Rate per Unit of Effective Labor:')
display(sm.Eq(wt,wage_rate))
```

### **Balanced growth path & steady state**

We are now able to set up the transition to the balanced growth path:

\begin{eqnarray*}
K_{t+1}=s(r_{t+1})w_tA_tL_t
\end{eqnarray*}

and therefore

\begin{eqnarray*}
k_{t+1}=s(r_{t+1})w_t\frac{A_tL_t}{A_{t+1}L_{t+1}}=s(r_{t+1})w_t\frac{A_tL_t}{A_t(1+g)L_t(1+n)}=s(r_{t+1})w_t\frac{1}{(1+g)(1+n)}
\end{eqnarray*}

```python
# Define again the necessary sympy variables:

kt1 = sm.symbols('k_t+1')

g = sm.symbols('g')
n = sm.symbols('n')
```

```python
# We also define the transition function as sympy:

transition_int = sm.Eq(kt1,s*wt*(1/((1+g)*(1+n))))
print('Transition Equation:')
display(transition_int)
```

```python
# And we can plug in the saving rate s as well as the wage rate wt:

transition_int2 = transition_int.subs(s,saving_rate)
transition_int2 = transition_int2.subs(wt,wage_rate)
print('Transition Equation with Saving and Wage rates plugged-in:')
display(transition_int2)
```

```python
# We then simply replace k_t+1 with k_t:

ss = transition_int2.subs(kt1,kt)
print('In SS k is constant:')
display(ss)
```

```python
# And are the able to calculate the steady state. The numerator looks complicated is however only the product of (1+g)(1+n)(2+p):

kss = sm.solve(ss,kt)[0]
print('Steady-State k:')
display(kss)
```

### **Steady-state calculation**

In this part we will now set up functions in order to calculate the steady-state.

```python
# It is then straightforward to set up a steady-sate function:

ss_func = sm.lambdify((alpha,g,n,rho),kss)

# we define a function with which we can calcluate the steady state with flexible parameteres:
def ss_calc(alpha,g,n,rho):
    result = ss_func(alpha,g,n,rho)
    print(f'The steady state solved by the standard python function is: {result}')
```

```python
# with this widget one can choose the parameter and it yields the steady state capital stock
widgets.interact(
        ss_calc,
        alpha = widgets.FloatSlider(description="$\\alpha$", min=0.01, max=0.99, step=0.01, continuous_update=False,),
        rho = widgets.FloatSlider(description="$\\rho$", min=0, max=0.99, step=0.01, continuous_update=False,),
        g = widgets.FloatSlider(description="g", min=-0.99, max=0.99, step=0.01, continuous_update=False,),
        n = widgets.FloatSlider(description="n", min=-0.99, max=0.99, step=0.01, continuous_update=False,),
)
```

    interactive(children=(FloatSlider(value=0.01, continuous_update=False, description='$\\alpha$', max=0.99, min=…

    <function __main__.ss_calc(alpha, g, n, rho)>

```python
# we call the steady-state function with certain values as input as comparison:
# we choose parameters, which we think are reasonable today

result = ss_func(1/3,0.03,0.02,0.05)
print(f'The steady state solved by a standart python function is: {result}')
```

    The steady state solved by a standart python function is: 0.17221708150191095

```python
# We can now solve this the problem numerically with optimize.root_scalar.
# However we have to make sure that we do not overwrite the already defined variables:

alpha_2 = 1/3
g_2 = 0.03
n_2 = 0.02
rho_2 =0.05

obj_kss = lambda kss: kss - ((kss**alpha_2*(1-alpha_2))/((1+g_2)*(1+n_2)*(2+rho_2)))

result2 = optimize.root_scalar(obj_kss,bracket=[0.1,100],method='brentq')
print(f'The steady state solved numerically with optimize.root_scalar is: {result2.root}')
print(f'Details are as followed:')
result2
```

    The steady state solved numerically with optimize.root_scalar is: 0.1722170815014273
    Details are as followed:

          converged: True
               flag: 'converged'
     function_calls: 8
         iterations: 7
               root: 0.1722170815014273

```python
# We test if the two solutions are the same:

np.allclose(result, result2.root)
```

    True

### **Steady-state illustration**

We now set up a plot, where the balanced growth path as well as the steady-state are shown. With several sliders one then can choose values for the different variabels and parameters.

```python
# We have to define first an equation corresponding to transition_int2 i.e. the corresponding path:

def kt1_func(kt, alpha, rho, n, g):
    return ((1-alpha)*kt**alpha)/((2+rho)*(1+n)*(1+g))

kt = np.linspace(0,1,500)
```

```python
# And the plot function:

def plot_OLG(alpha,rho,g,n):
    plt.figure(dpi=150)
    plt.plot(kt,kt1_func(kt, alpha, rho, n, g), color = 'green')
    plt.plot(range(10),range(10), color = 'blue')
    plt.plot(ss_func(alpha,g,n,rho),ss_func(alpha,g,n,rho), marker='*',color='red')
    plt.title('OLG model: Basic version')
    plt.xlabel('$ k_t $')
    plt.ylabel( '$ k_ {t + 1} $')
    plt.xlim([0,1])
    plt.ylim([0,1])
    plt.legend(('BGP', '45\N{DEGREE SIGN}-line', 'steady state'), loc='upper left')
    plt.grid()
    sol = np.round(ss_func(alpha,g,n,rho),3)
    text = 'Steady state at: ' + str(sol)
    plt.text(0.64,0.95,text, fontsize='small', c='red', bbox=dict(facecolor='none', edgecolor='red'))
    text = 'Here you can choose all variables and parameters of the basic OLG model!'
    plt.text(0, 1.1, text, fontsize='x-small', c='black')

# Plotting plots, with interactive slider
widgets.interact(
        plot_OLG,
        alpha = widgets.FloatSlider(description="$\\alpha$", min=0.01, max=0.99, step=0.01, value=alpha_2, continuous_update=False,),
        rho = widgets.FloatSlider(description="$\\rho$", min=-0.99, max=0.99, step=0.01, value=rho_2, continuous_update=False,),
        g = widgets.FloatSlider(description="g", min=-0.99, max=0.99, step=0.01, value=g_2, continuous_update=False,),
        n = widgets.FloatSlider(description="n", min=-0.99, max=0.99, step=0.01, value=n_2, continuous_update=False,),
)
```

    interactive(children=(FloatSlider(value=0.3333333333333333, continuous_update=False, description='$\\alpha$', …

    <function __main__.plot_OLG(alpha, rho, g, n)>

### Comparison of the steeady state depending on the parameters

In the steady state the capital per unit of effective labor is constant and depends only on four parameters. In the this section, we compare some output variables like the consumption, utility and output depending on these input parameters.

```python
def parameter_comparison(alpha, rho, n, g, alpha2, rho2, n2, g2):
    # steady-state k
    result = ss_func(alpha, rho, n, g)
    result2 = ss_func(alpha2, rho2, n2, g2)

    fig, axs = plt.subplots(4, 2, figsize=(15,15))
    fig.suptitle('Steady State Comparison Depending on Different Parameters', fontsize=20)

    # Technology
    At = A*(1+g)**t
    At2 = A*(1+g2)**t
    axs[0, 0].plot(t, At, c='blue')
    axs[0, 0].plot(t, At2, c='green')
    axs[0, 0].set_title('Technology A')
    axs[0, 0].grid()
    axs[0, 0].legend(('Parameters 1', 'Parameters 2'), loc='upper left')

    # Labor
    Lt = L*(1+n)**t
    Lt2 = L*(1+n2)**t
    axs[0, 1].plot(t, Lt, c='blue')
    axs[0, 1].plot(t, Lt2, c='green')
    axs[0, 1].set_title('Labor L')
    axs[0, 1].grid()
    axs[0, 1].legend(('Parameters 1', 'Parameters 2'), loc='upper left')

    # wage as defined above
    w = np.empty(t.size)
    w2 = np.empty(t.size)
    for i in range(t.size):
        w[i] = result**alpha *(1-alpha)
        w2[i] = result2**alpha2 *(1-alpha2)

    axs[1, 0].plot(t, w, c='blue')
    axs[1, 0].plot(t, w2, c='green')
    axs[1, 0].set_title('Wage Rate per Unit of Effective Labor')
    axs[1, 0].grid()
    axs[1, 0].legend(('Parameters 1', 'Parameters 2'), loc='center left')

    # real interest rate as defined above
    r = np.empty(t.size)
    r2 = np.empty(t.size)
    for i in range(t.size):
        r[i] = result**(alpha-1) *(alpha)
        r2[i] = result2**(alpha2-1) *(alpha2)

    axs[1, 1].plot(t, r, c='blue')
    axs[1, 1].plot(t, r2, c='green')
    axs[1, 1].set_title('Real Interest Rate')
    axs[1, 1].grid()
    axs[1, 1].legend(('Parameters 1', 'Parameters 2'), loc='center left')

    # consumption period 1 as defined above
    C_t1 = A*(1+g)**t*w*(rho+1)/(rho+2)
    C2_t1 = A*(1+g2)**t*w2*(rho2+1)/(rho2+2)

    axs[2, 0].plot(t, C_t1, c='blue')
    axs[2, 0].plot(t, C2_t1, c='green')
    axs[2, 0].set_title('HH Consumption in Period 1')
    axs[2, 0].grid()
    axs[2, 0].legend(('Parameters 1', 'Parameters 2'), loc='upper left')

    # consumption period 2 as defined above
    C_t2 = A*(1+g)**t*w*(1+r)/(rho+2)
    C2_t2 = A*(1+g2)**t*w2*(1+r2)/(rho2+2)

    axs[2, 1].plot(t, C_t2, c='blue')
    axs[2, 1].plot(t, C2_t2, c='green')
    axs[2, 1].set_title('HH Consumption in Period 2')
    axs[2, 1].grid()
    axs[2, 1].legend(('Parameters 1', 'Parameters 2'), loc='upper left')

    # utility
    U = np.log(C_t1) + np.log(C_t2)/(rho+1)
    U2 = np.log(C2_t1) + np.log(C2_t2)/(rho2+1)
    axs[3, 0].plot(t, U, c='blue')
    axs[3, 0].plot(t, U2, c='green')
    axs[3, 0].set_title('Lifetime Utility of a HH')
    axs[3, 0].grid()
    axs[3, 0].legend(('Parameters 1', 'Parameters 2'), loc='upper left')

    # production
    Y = result**alpha*(A*(1+g)**t *(L*(1+n)**t))
    Y2 = result2**alpha2*(A*(1+g2)**t *(L*(1+n2)**t))
    axs[3, 1].plot(t, Y, c='blue')
    axs[3, 1].plot(t, Y2, c='green')
    axs[3, 1].set_title('Production of the Whole Economy Y')
    axs[3, 1].grid()
    axs[3, 1].legend(('Parameters 1', 'Parameters 2'), loc='upper left')

    fig.text(0.5, 0.10, 'Time', ha='center', va='center',  fontsize=16)
    fig.text(0.08, 0.5, 'Variable of Interest', ha='center', va='center', rotation='vertical',  fontsize=16)

```

```python
A = 1
L = 1
t = np.linspace(0,100,100)
widgets.interact(
        parameter_comparison,
        alpha = widgets.FloatSlider(description="$\\alpha_1$", min=0.01, max=0.99, step=0.01, value=alpha_2, continuous_update=False,),
        rho = widgets.FloatSlider(description="$\\rho_1$", min=-0.99, max=0.99, step=0.01, value= rho_2, continuous_update=False,),
        g = widgets.FloatSlider(description="$g_1$", min=-0.99, max=0.99, step=0.01, continuous_update=False,),
        n = widgets.FloatSlider(description="$n_1$", min=-0.99, max=0.99, step=0.01, continuous_update=False,),
        alpha2 = widgets.FloatSlider(description="$\\alpha_2$", min=0.01, max=0.99, step=0.01, value=alpha_2, continuous_update=False,),
        rho2 = widgets.FloatSlider(description="$\\rho_2$", min=-0.99, max=0.99, step=0.01, value=rho_2, continuous_update=False,),
        g2 = widgets.FloatSlider(description="$g_2$", min=-0.99, max=0.99, step=0.01, value=g_2, continuous_update=False,),
        n2 = widgets.FloatSlider(description="$n_2$", min=-0.99, max=0.99, step=0.01, value=n_2, continuous_update=False,),
)
```

    interactive(children=(FloatSlider(value=0.3333333333333333, continuous_update=False, description='$\\alpha_1$'…

    <function __main__.parameter_comparison(alpha, rho, n, g, alpha2, rho2, n2, g2)>

The figure shows that:
- in the steady state the wage per unit of effective labor is constant, this statement is true for all variables measured in effective labor units
- the real interest rate is constant, because it depends on the steady state capital stock per unit of effective labor, which is constant
- the houshold consumption and lifetime utility depends crucially on the technology growth, g
- the whole production of the economy grows with g and n

## **OLG with a government**

The basic OLG from the first part is not that realistic. What is missing for example is goverment spending, which of course has to be financed somehow. The two basic means to finance government spending is either through taxes or otherwise by issuing debt-bonds. Both of them reduce the available income of some generations, and it is therefore interesting to study what effect such an implementation of a goverment has for the balanced-growth-path as well as the steady-state. As before, we do not change the setting of log-utility and Cobb-Douglas productivity. Furthermore, goverment spending G_t is also measured as spending per unit of effective labor (as well as taxes and bonds). What is obvious is that the government spending must be balanced, i.e.

\begin{eqnarray*}
G_t=T_t\\ G_t=b_t
\end{eqnarray*}

We first consider a labor tax, i.e.

\begin{eqnarray*}
C_{1t}=A_t(w_t-T_t)-S_t
\end{eqnarray*}

Savings are then:

\begin{eqnarray*}
S_t=\frac{1}{2+\rho}A_t(w_t-T_t)
\end{eqnarray*}

Ant the balanced growth path becomes:

\begin{eqnarray*}
k_{t+1}=\frac{(1-\alpha)k_t^\alpha-T_t}{(1+g)(1+n)(2+\rho)}
\end{eqnarray*}

With debt issuing, it is much more straightforward, as individuals buy government bonds instead of saving and the balanced growth path becomes therefore:

\begin{eqnarray*}
k_{t+1}=\frac{(1-\alpha)k_t^\alpha}{(1+g)(1+n)(2+\rho)}-b_{t+1}
\end{eqnarray*}

Let's analyse the effect with a nice plot:

```python
# We redefine the corresponding path steady-state functions:

kt = sm.symbols('k_t')
kt1 = sm.symbols('k_t+1')
alpha = sm.symbols('alpha')
T = sm.symbols('T_t')
bt1 = sm.symbols('b_t+1')
g = sm.symbols('g')
n = sm.symbols('n')
rho = sm.symbols('rho')

tax_transition = sm.Eq(kt1,((1-alpha)*kt**alpha-T)/((1+g)*(1+n)*(2+rho)))
bond_transition = sm.Eq(kt1,(((1-alpha)*kt**alpha)/((1+g)*(1+n)*(2+rho)))-bt1)

ss_tax = tax_transition.subs(kt1,kt)
# kss = sm.solve(ss_tax,kt)[0]          # Try this command; we are not able to calculate the kss-formula as before!!!

                                        # This means that we are not able to plot stars for the new steady states, which is no that bad as
                                        # steady states are obvious when plotting a 45°-line.
```

```python
# We also redefine the corresponding path equations:

def kt1_no_gov_func(kt, alpha, rho, n, g):
    return ((1-alpha)*kt**alpha)/((2+rho)*(1+n)*(1+g))

def kt1_with_tax_func(kt, alpha, rho, n, g, G):            # With G as variable instead of T_t
    return ((1-alpha)*kt**alpha-G)/((2+rho)*(1+n)*(1+g))

def kt1_with_bonds_func(kt, alpha, rho, n, g, G):          # With G as variable instead of b_t+1
    return ((1-alpha)*kt**alpha)/((2+rho)*(1+n)*(1+g))-G

kt = np.linspace(0,1,500)
```

```python
# And again the plot function:

def plot_OLG_with_G(alpha,rho,g,n,G):
    plt.figure(dpi=150)
    plt.plot(kt,kt1_no_gov_func(kt, alpha, rho, n, g))
    plt.plot(kt,kt1_with_tax_func(kt, alpha, rho, n, g, G))
    plt.plot(kt,kt1_with_bonds_func(kt, alpha, rho, n, g, G))
    plt.plot(range(10),range(10),color='black')
    plt.title('OLG model with government G = ' + str(G))
    plt.xlabel('$ k_t $')
    plt.ylabel('$ k_ {t + 1} $')
    plt.xlim([0,1])
    plt.ylim([0,1])
    plt.legend(('G=0','Tax','Bonds','45\N{DEGREE SIGN}-line'), loc='upper left')
    plt.grid()
    text = 'Here you can also choose the size of government!'
    plt.text(0, 1.1, text, fontsize='x-small', c='black')
```

```python
widgets.interact(
        plot_OLG_with_G,
        G = widgets.FloatSlider(description="G_size", min=0, max=0.25, step=0.01, value=0.1, continuous_update=False),
        alpha = widgets.FloatSlider(description="$\\alpha$", min=0.01, max=0.99, step=0.01, value=alpha_2, continuous_update=False,),
        rho = widgets.FloatSlider(description="$\\rho$", min=-0.99, max=0.99, step=0.01, value=rho_2, continuous_update=False,),
        g = widgets.FloatSlider(description="g", min=-0.99, max=0.99, step=0.01, value=g_2, continuous_update=False,),
        n = widgets.FloatSlider(description="n", min=-0.99, max=0.99, step=0.01, value=n_2, continuous_update=False,),
)
```

    interactive(children=(FloatSlider(value=0.3333333333333333, continuous_update=False, description='$\\alpha$', …

    <function __main__.plot_OLG_with_G(alpha, rho, g, n, G)>

As a result, we see that goverment spending i.e. its financing reduces the steady-state capital stock, however the effect of taxes is smaller than the effect of debt-issuing.

```python
# As above we can now solve this the problem numerically with optimize.root_scalar.

alpha = 1/3
g = 0.03
n = 0.02
rho =0.05

# As can be seen in the figure, the government spending can not be really big otherwise the model does not converge to a steady state.
G= 0.04

obj_kss_ = lambda kss: kss - ((1-alpha)*kss**alpha)/((2+rho)*(1+n)*(1+g))

obj_kss_tax = lambda kss: kss - ((1-alpha)*kss**alpha-G)/((2+rho)*(1+n)*(1+g))  # With G as variable instead of T_t

obj_kss_bond = lambda kss: kss - (((1-alpha)*kss**alpha)/((2+rho)*(1+n)*(1+g))-G)  # With G as variable instead of b_t+1

result_ = optimize.root_scalar(obj_kss_,bracket=[0.1,100],method='brentq')
result_t = optimize.root_scalar(obj_kss_tax,bracket=[0.1,100],method='brentq')
result_b = optimize.root_scalar(obj_kss_bond,bracket=[0.1,100],method='brentq')

print(f'The steady state solved numerically with optimize.root_scalar is without government: {result_.root}')

print(f'The steady state solved numerically with optimize.root_scalar is with taxes: {result_t.root}')

print(f'The steady state solved numerically with optimize.root_scalar is with bonds: {result_b.root}')

```

    The steady state solved numerically with optimize.root_scalar is without government: 0.1722170815014273
    The steady state solved numerically with optimize.root_scalar is with taxes: 0.1434754466865211
    The steady state solved numerically with optimize.root_scalar is with bonds: 0.106913475354846

## Conclusion

In our model project we analyzed a standard overlapping generations model. First, we solved the household optimization using sympy. Similarly, we solved the production side. Combining these two optimal behaviors, we are able to calculate the steady state. We then visualize the steady state and the evolution of the key variables in this steady state. Moreover, we make this visualization interactive, so we can directly observe the effect of the model parameter. In the last step, we introduce a government and analyze the steady state in this case. Unfortunately, we are unable to calculate a closed form solution for the steady state. But we can show that the effect of taxes is smaller than the effect of bonds.


```python
%matplotlib inline

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import sklearn.linear_model as linear_model

np.random.seed(1337)
from timeit import timeit

import matplotlib
import matplotlib.pyplot as plt

kwargs = {'linewidth' : 3.5}
font = {'weight' : 'normal', 'size'   : 24}

def error_plot(ys, yscale='log'):
    plt.figure(figsize=(8, 8))
    plt.xlabel('Step')
    plt.ylabel('Error')
    plt.yscale(yscale)
    plt.plot(range(len(ys)), ys, **kwargs)

def error_plot3(ys1, ys2, ys3, yscale='log'):
    plt.figure(figsize=(8, 8))
    plt.xlabel('Step')
    plt.ylabel('Error')
    plt.yscale(yscale)
    plt.plot(range(len(ys1)), ys1, **kwargs, label='GD')
    plt.plot(range(len(ys2)), ys2, **kwargs, label='Newton')
    plt.plot(range(len(ys3)), ys3, **kwargs, label='BFGS')
    plt.legend()

def error_plot2(ys1, ys2, yscale='log'):
    plt.figure(figsize=(8, 8))
    plt.xlabel('Step')
    plt.ylabel('Error')
    plt.yscale(yscale)
    plt.plot(range(len(ys1)), ys1, **kwargs, label='GD')
    plt.plot(range(len(ys2)), ys2, **kwargs, label='Newton')
    plt.legend()

def convergence_plot(fs, gs):
    plt.figure(figsize=(14,6))
    plt.subplot(121)
    plt.title('Convergence in objective')
    plt.xlabel('Step')
    plt.ylabel('Error')
    plt.yscale('log')
    plt.plot(range(len(fs)), fs, **kwargs)
    plt.subplot(122)
    plt.title('Convergence in domain')
    plt.xlabel('Step')
    plt.yscale('log')
    plt.plot(range(len(gs)), gs, **kwargs)
    plt.tight_layout()

def setup_layout():
    display(HTML(open('custom.html').read()))
    matplotlib.rc('font', **font)

```

### Solving linear regression with Newton's method

**Purpose of this demo**: Motivate alternatives of gradient descent

+ Disclaimer: I'm not expert in Python - I use Python/Matlab as tools to validate algorithms and theorems.
+ Thus, my implementations are not the most efficient ones + there might be bugs

**Problem definition: Linear regression**.

\begin{align}
  y = A x^\star
\end{align}

+ $A \in \mathbb{R}^{n \times n}$
+ $x^\star \in \mathbb{R}^n$

Assume $A$ is in general position.
Given $y$ and $A$:

\begin{equation*}
\begin{aligned}
& \underset{x \in \mathbb{R}^p}{\text{min}}
& & f(x) \triangleq \|y - A x\|_2^2 \\
\end{aligned}
\end{equation*}

+ We want to solve this problem with Newton's method

\begin{align}
x_{t+1} = x_t - \eta \left(\nabla^2 f(x_t) \right)^{-1} \nabla f(x_t)
\end{align}

+ What is the gradient and Hessian of this function?

\begin{align}
\nabla f(x) = -A^\top \left(y - Ax\right) \in \mathbb{R}^n
\end{align}

and
\begin{align}
\nabla^2 f(x) = A^\top A \in \mathbb{R}^{n \times n}
\end{align}

+ Assuming that $A$ is in general position, $A^\top A$ is full rank

+ Let us use the above in Newton's iteration
+ Assuming $x_0$ is zero (initial point), and the step size is $\eta = 1$, then

\begin{align}
x_1 = x_0 - \left(\nabla^2 f(x_0) \right)^{-1} \nabla f(x_0) = - (A^\top A)^{-1} (- A^\top y) = (A^\top A)^{-1} A^\top y
\end{align}

+ Does this solution ring a bell?

#### Questions on Newton's method for linear regression

+ What about underdetermined linear systems?
+ What is the complexity of running Newton's method in well-conditioned linear systems?

\begin{align}
O\left( \log \frac{1}{\varepsilon} \cdot n^2 \right) \quad \quad \text{vs} \quad \quad O \left(n^3\right)
\end{align}

### Logistic regression with quasi-Newton methods

We will solve the logistic regression problem using second-order and quasi-second-order methods, and compare them with first order methods, such as plain gradient descent.

Let's first redefine the sigmoid function:

```python
def sigmoid(x):
    """Sigmoid function."""
    return 1./(1.+np.exp(-x))
```

```python
xs = np.linspace(-5, 5, 100)
plt.plot(xs, sigmoid(xs), **kwargs);
```

```python
def logistic_model(weights, data):
    return sigmoid(data.dot(weights))

def log_likelihood(weights, data, labels):
    scores = np.dot(data, weights)
    return -np.mean(labels*scores - np.log(1 + np.exp(scores)))

def log_likelihood_gradient(weights, data, labels):
    predictions = logistic_model(weights, data)
    return -data.T.dot(labels - predictions)/len(data)

def log_likelihood_hessian(weights, data, labels):
    predictions = logistic_model(weights, data)
    diag = np.diag(predictions * (1 - predictions))
    return data.T.dot(diag.dot(data))/len(data)
```

```python
m, n = 1000, 100
data = np.random.normal(0, 1, (m, n))
wopt = np.random.normal(0, 1, n)
labels = (data.dot(wopt) > 0).astype(np.float)
gradient = lambda w: log_likelihood_gradient(w, data, labels)
hessian = lambda w: log_likelihood_hessian(w, data, labels)
init = np.zeros(n)
```

```python
import sklearn.linear_model as linear_model
```

```python
cls = linear_model.LogisticRegression()
cls.fit(data, labels)
weights = np.reshape(cls.coef_, cls.coef_.shape[1])
baseline = log_likelihood(weights, data, labels)
baseline
```

    0.05166262561432568

```python
def gradient_descent(init, step_sizes, grad):
    X = [init]
    for step in step_sizes:
        X.append(X[-1] - step * grad(X[-1]))
    return X
```

```python
def newton_descent(init, step_sizes, grad, hessian):
    X = [init]
    for step in step_sizes:
        Hinv = np.linalg.pinv(hessian(X[-1]))
        X.append(X[-1] - step * Hinv.dot(grad(X[-1])))
    return X
```

```python
ws_gd = gradient_descent(init, np.ones(200), gradient)
ws_newton = newton_descent(init, np.ones(200), gradient, hessian)
```

```python
error_plot2([log_likelihood(w, data, labels) for w in ws_gd],
            [log_likelihood(w, data, labels) for w in ws_newton])
```

### Using BFGS in logistic regression

```python
def bfgs(init, grad, hessian, line_search, num_steps, num_to_keep=None):
    """BFGS algorithm (toy implementation).

    Parameters
    ----------
        initial : array
            starting point
        grad : function
            mapping arrays to arrays of same shape
        hessian : function
            mapping 1d arrays to 2d arrays
        line_search : function
            method to determine step size
        num_steps : number
            number of iterations
        num_to_keep : integer, optional
            number of points to keep

    Returns
    -------
        List of points computed by Newton iteration. Length of the
        list is determined by `num_to_keep`.
    """
    xs = [init]
    I = np.eye(len(init))
    Hs = [hessian(init)]
    Hinvs = [I]

    for _ in range(num_steps):
        # update point
        direction = Hinvs[-1].dot(grad(xs[-1]))
        step = line_search(xs[-1], direction)
        xs.append(xs[-1] + step * direction)

        s = xs[-1] - xs[-2]
        y = grad(xs[-1]) - grad(xs[-2])
        ys = np.dot(y, s)
        ssT = np.outer(s, s)

        # udpate H
        H = Hs[-1] + np.outer(y, y)/ys
        H -= Hs[-1].dot(ssT.dot(Hs[-1]))/np.dot(s, Hs[-1].dot(s))
        Hs.append(H)

        # update Hinv
        left = I - np.outer(s, y)/ys
        right = I - np.outer(y, s)/ys
        Hinvs.append(left.dot(Hinvs[-1]).dot(right) + ssT/ys)

        if num_to_keep:
          xs = xs[-num_to_keep:]

    return xs
```

```python
def line_search(x, d):
    """Constant step size."""
    return -0.1

ws_bfgs = bfgs(init, gradient, hessian, line_search, 200)

error_plot3([log_likelihood(w, data, labels) for w in ws_gd],
            [log_likelihood(w, data, labels) for w in ws_newton],
            [log_likelihood(w, data, labels) for w in ws_bfgs])

# error_plot([log_likelihood(w, data, labels) for w in ws_bfgs])
# plt.plot(range(len(ws_bfgs)), [baseline]*len(ws_bfgs),
#         label='sklearn', **kwargs)
# _ = plt.legend()
```

```python

```

```python

```


The Maxwell's equations are

\begin{align}
\frac{\partial \vec{B}}{\partial t} &= - \nabla \times \vec{E}            \hspace{2cm} \text{Faraday equations}\\
\frac{\partial \vec{D}}{\partial t} &= - \nabla \times \vec{H} - \vec{J} - \vec{s} \hspace{2cm} \text{Ampere's law.}
\end{align}

where $\vec{E}$ is the electric field, $\vec{H}$ is the magnetic field, $\vec{J}$ is the electric flux, $\vec{D}$ is the displacement flux, $\vec{B}$ is the magnetic flux and $\vec{s}$ is the source.

The constitutive realtations between the fields and flux are as follows.

\begin{align}
\vec{J} &= \sigma \vec{E}\\
\vec{B} &= \mu \vec{H}\\
\vec{D} &= \epsilon \vec{E}
\end{align}

Using the contitutive relations, the Maxwell's equations can written in terms of the electric field and the magnetic flux
\begin{align}
\mu \frac{\partial \vec{B}}{\partial t} &= - \nabla \times \vec{E}  \\
\epsilon \frac{\partial \vec{E}}{\partial t} &= -  \nabla \times \frac{1}{\mu} \vec{B} - \sigma \vec{E} - \vec{s}
\end{align}

Futher by droping the $\epsilon$ term and using a $e^{-i\omega t}$ Fourier time relation, Maxwell's equations can be written in the frequency domain as

\begin{align}
-i \omega \mu  \hat{\vec{B}} &= - \nabla \times \hat{\vec{E}}  \\
- \vec{s} &= -  \nabla \times \frac{1}{\mu} \hat{\vec{B}} - \sigma \hat{\vec{E}}
\end{align}

The weak form is
\begin{align}
-i \omega \mu (\hat{\vec{B}},\hat{\vec{F}}) + (\nabla \times \hat{\vec{E}},\hat{\vec{F}}) &= 0 \\
(\nabla \times \mu^{-1} \hat{\vec{B}},\hat{\vec{W}}) + (\sigma \hat{\vec{E}},\hat{\vec{W}}) &= (\hat{\vec{s}},\hat{\vec{W}})
\end{align}

$ (\mu^{-1} B, \nabla \times W) + (\sigma E, W) = (\mu^{-1}B W )|_0^{end} $

$ (i \omega B,F) + (\nabla \times E , f) = 0 $

```

```

```

```


# Basic Symbolic Quantum Mechanics

```python
from sympy import init_printing
init_printing(use_latex=True)
```

```python
from sympy import sqrt, symbols, Rational, srepr
from sympy import expand, Eq, Symbol, simplify, exp, sin
from sympy.physics.quantum import *
from sympy.physics.quantum.qubit import *
from sympy.physics.quantum.gate import *
from sympy.physics.quantum.grover import *
from sympy.physics.quantum.qft import QFT, IQFT, Fourier
from sympy.physics.quantum.circuitplot import circuit_plot
```

    /Users/bgranger/anaconda/lib/python3.5/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.')
    /Users/bgranger/anaconda/lib/python3.5/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.')

## Bras and kets

Symbolic kets can be created using the `Ket` class as seen here:

```python
phi, psi = Ket('phi'), Ket('psi')
```

These ket instances are fully symbolic and behave exactly like the corresponding mathematical entities.
For example, one can form a linear combination using addition and scalar multiplication:

```python
alpha = Symbol('alpha', complex=True)
beta = Symbol('beta', complex=True)
state = alpha*psi + beta*phi; state
```

Bras can be created using the `Bra` class directly or by using the `Dagger` class
on an expression involving kets:

```python
ip = Dagger(state)*state; ip
```

Because this is a standard SymPy expression, we can use standard SymPy functions and methods
for manipulating expression. Here we use expand to multiply this expression out, followed
by `qapply` which identifies inner and outer products in an expression.

```python
qapply(expand(ip))
```

## Operators

SymPy also has a full set of classes for handling symbolic operators.  Here we create three operators,
one of which is hermitian:

```python
A = Operator('A')
B = Operator('B')
C = HermitianOperator('C')
```

When used in arithmetic expressions SymPy knows that operators do not commute under
multiplication/composition as is seen by expanding a polynomial of operators:

```python
expand((A+B)**2)
```

Commutators of operators can also be created:

```python
comm = Commutator(A*B,B+C); comm
```

The `expand` function has custom logic for expanding commutators using standard commutator
relations:

```python
comm.expand(commutator=True)
```

Any commutator can be performed ($[A,B]\rightarrow AB-BA$) using the `doit` method:

```python
_.doit().expand()
```

The `Dagger` class also works with operators and is aware of the properties of unitary
and hermitian operators:

```python
Dagger(_)
```

## Tensor products

Symbolic tensor products of operators and states can also be created and manipulated:

```python
op = TensorProduct(A,B+C)
state = TensorProduct(psi,phi)
op*state
```

Once a tensor product has been created, it can be simplified,

```python
tensor_product_simp(_)
```

and expanded:

```python
expand(_)
```