```
from __future__ import division
import numpy as np
import sympy as sym
import matplotlib as mpl
import matplotlib.pyplot as plt
```

```
## Setting the parameters of the model (calibrated values)
R, rho, gamma, eta, L, A, kapa = sym.var('R, rho, gamma, eta, L, A, kapa')
```

```
### Defining variables ###
k, x, y, l, r = sym.var('k, x, y, l, r')

### Defining production function ###
gamma = 0.5
rho = 0.25
f = A*kapa*(eta*(k*x)**rho + (1- eta)*(y*(l/r))**rho)**(gamma/rho) - k*R
```

## Part 1: The first order condition.
Differentiating this function leads us to the first order condition:
$$\frac{(0.5A\eta\kappa(kx)^{1/4}(\eta(kx)^{1/4} - (ly/r)^{1/4}(1 -\eta)))}{k} - R = 0 $$

Solving it directly doesn't work:

```
sym.solve(0.5*A*eta*kapa*(k**(1/4))*(x**(1/4))*(eta*(k*x)**0.25 + (l*y/r)**0.25*(-eta + 1)) - R, k, set=True)
# Double click to enlarge output
```

# Part 2: A closer llok at the polynomial for the FOC

We can rename $K = k^{1/4}$ and obtain a thrid degree polynomial in $K$, and try to solve it by the usual equation solver:

```
## Defining auxiliary functions a, c, d
a = 2*R
c = -A*kapa*(eta**2)*x**0.5
d = -A*kapa*eta*(1-eta)*(x*y*(l/r))**0.25
K = sym.var('K')
pol_K = K**3*a + K*c + d
sym.solve(pol_K, K, warn=True, set=True, rational=True, manual=True)
```

    ([], set())

...which doesn't seem to work. How to solve this cubic equation?

There are several ways. But before proceeding, it is good to visualise the equation for the parameters of the model and reasonable values of the variables.

Next cell plots, for some range of values from 0 to 10, the value of the polynomial we are after. When the line crosses the vertical zero, we have a root. It is easy to see that there are, in fact, 3 real roots, two negatives and one positive.

```
K_range = np.arange(-3.0,3.0,0.01)
polik = sym.lambdify((K,x,y,l,r,A,kapa,eta,R),pol_K)
for i in np.arange(0,10,0.1):
    plt.plot(K_range, polik(K_range,i,i,1,1,1,1,0.89,0.13099))
plt.axvline()
plt.axhline()
plt.show()
```

The graph above shows how for most parameter values, this polynomial has three roots, of which one is positive.

Below is a zoomed in graph for values of $x,y,l$ very close to zero. There are still three roots, only one positive.

```
K_range = np.arange(-1.0,1.0,0.001)
plt.plot(K_range, polik(K_range,0.01,0.01,0.5,1,1,1,0.89,0.13099))
plt.axvline()
plt.axhline()
plt.show()
```

We can also vary $l$ (as we are going to normalise $r$ anyway) and see that it displaces the line to the right or to the left:

```
K_range = np.arange(-2.0,2.0,0.01)
for i in np.arange(0.01,10,0.1):
    plt.plot(K_range, polik(K_range,1,1,i,1,1,1,0.89,0.13099))
plt.axvline()
plt.axhline()
plt.show()
```

The extreme case is when $l=0$, in which case the second root is exactly zero. This is true for any *reasobale* combination of $x$ and $y$:

```
K_range = np.arange(-2.0,2.0,0.01)
for i in np.arange(0.01,10,0.5):
    plt.plot(K_range, polik(K_range,i,i,0,1,1,1,0.89,0.13099))
plt.axvline()
plt.axhline()
plt.show()
```

In the case when $l$ is also close to zero, the thried root gets closer to the zero as well. This makes sense! Zero workers, zero capital.

# Part 3: Solving for K
The first method I was using was to check for the condition below:

if delta < 0 --> only one real root

if delta > 0 --> three real roots

if delta = 0 --> combination of both

```
delta = 4*A*kapa*eta**4*x - 54*R*((1-eta)**2)*((y*(l/r))**(1/2))

def delta_condition(x,y,l,r,A,kappa,eta,R):
    '''
    Condition for one real root

    Input: A series of model parameters and values (x,y,l,r,A,kappa,eta,R)

    Output: the numerical value of the Delta equation followed to True if Delta<0.
    '''
    delta = 4*A*kappa*(eta**4)*x**(1/2) - 54*R*((1-eta)**2)*((y*(l/r))**(1/2))
    return delta, delta<0

# Example for (1,10,10):
delta_condition(1,1,1,1,1,1,0.89,0.13099)
```

    (2.4241007740000002, False)

As the condition fails, and as we saw in teh grpah above, there are likely to be three real roots.

A way of finding them is the *Trigonometric Method*: If the condition below applies,

```
def trig_condition(x,y,l,r,A,kappa,eta,R):
    '''
    Condition for three real roots

    Input: A series of model parameters and values (x,y,l,r,A,kappa,eta,R)

    Output: the numerical value of the trigonometric condition equation followed by the result.
    '''
    arcosine = (27/4)*(1-eta)*(y*l/r)**(1/2) - ((eta**4)*A*kappa*(x**(1/2)))/(2*R)
    if arcosine<=0:
        message = "The equation has three real roots"
    else:
        message = "The equation has at least one imaginary root"
    return arcosine, message
trig_condition(1,1,1,1,1,1,0.89,0.13099)
```

    (-1.6524248415909615, 'The equation has three real roots')

The condition is met! Let's apply the trigonometric solution found on wikipedia for the three roots!

```
p = -(eta**2*A*kapa*(x**0.5))/(2*R)
q = -(eta*(1-eta)*kapa*A*(x*y*(l/r))**(1/4))/(2*R)
qmin = (eta*(1-eta)*kapa*A*(x*y*(l/r))**(1/4))/(2*R)
root_zero = 2*(-p/3)**(0.5)*sym.cos(((1/3)*sym.acos(((3*q)/(2*p))*(-3/p)**(1/2))))
root_two = -(2*(-p/3)**(0.5)*sym.cos(((1/3)*sym.acos(((3*qmin)/(2*p))*(-3/p)**(1/2)))))
root_one =  - root_zero - root_two
print "First root =", root_zero
print "Second root =", root_one
print "Thrird root =", root_two
```

    First root = 2*6**(-0.5)*(A*eta**2*kapa*x**0.5/R)**0.5*cos(0.333333333333333*acos(3*6**0.5*x**(-0.5)*(l*x*y/r)**0.25*(R*x**(-0.5)/(A*eta**2*kapa))**0.5*(-eta + 1)/(2*eta)))
    Second root = 2*6**(-0.5)*(A*eta**2*kapa*x**0.5/R)**0.5*cos(0.333333333333333*acos(-3*6**0.5*x**(-0.5)*(l*x*y/r)**0.25*(R*x**(-0.5)/(A*eta**2*kapa))**0.5*(-eta + 1)/(2*eta))) - 2*6**(-0.5)*(A*eta**2*kapa*x**0.5/R)**0.5*cos(0.333333333333333*acos(3*6**0.5*x**(-0.5)*(l*x*y/r)**0.25*(R*x**(-0.5)/(A*eta**2*kapa))**0.5*(-eta + 1)/(2*eta)))
    Thrird root = -2*6**(-0.5)*(A*eta**2*kapa*x**0.5/R)**0.5*cos(0.333333333333333*acos(-3*6**0.5*x**(-0.5)*(l*x*y/r)**0.25*(R*x**(-0.5)/(A*eta**2*kapa))**0.5*(-eta + 1)/(2*eta)))

Now we neeed to know which one of this is the poisitive root. Let's evaluate it at the values given above.

```
Rt0 = sym.lambdify((x,y,l,r,A,kapa,eta,R), root_zero)
Rt1 = sym.lambdify((x,y,l,r,A,kapa,eta,R), root_one)
Rt2 = sym.lambdify((x,y,l,r,A,kapa,eta,R), root_two)
Rt0(1,1,1,1,1,1,0.89,0.13099), Rt1(1,1,1,1,1,1,0.89,0.13099), Rt2(1,1,1,1,1,1,0.89,0.13099)
```

    (1.7976085400993507, -0.12422961405626354, -1.6733789260430871)

Let's see if we have indeed solve the polynomial:

```
K_range = np.arange(-3.0,3.0,0.01)
plt.plot(K_range, polik(K_range,1,1,1,1,1,1,0.89,0.13099))
plt.plot((Rt0(1,1,1,1,1,1,0.89,0.13099)), 0, 'ro')
plt.plot(Rt1(1,1,1,1,1,1,0.89,0.13099), 0, 'bo')
plt.plot(Rt2(1,1,1,1,1,1,0.89,0.13099), 0, 'go')
plt.axvline()
plt.axhline()
plt.show()
```

Because $k = K^4$, all of the roots will deliver positive capital. However, I have tried all three of them in the solver (see *Examples with k* notebook) and the first (positive) root is the only one that produces strictly positive choices of capital.

That means we have a winner! Below is the expression for the final root:

```
k_star = (2*6**(-0.5)*(A*eta**2*kapa*x**0.5/R)**0.5*sym.cos(0.333333333333333*sym.acos(3*6**0.5*x**(-0.5)*(l*x*y/r)**0.25*(R*x**(-0.5)/(A*eta**2*kapa))**0.5*(-eta + 1)/(2*eta))))**(4)
```

# Appendix: Some comparative statics

Here I adapt the function we found to become plottable:

```
def kfuc2(x,y,l,r,A,kapa,eta,R):
    return (2*6**(-0.5)*(A*eta**2*kapa*x**0.5/R)**0.5*np.cos(0.333333333333333*np.arccos(3*6**0.5*x**(-0.5)*(l*x*y/r)**0.25*(R*x**(-0.5)/(A*eta**2*kapa))**0.5*(-eta + 1)/(2*eta))))**(4)
```

I'm assuming positive assortive matching, that is, high types match with high types.

### Case 1: High types hire more workers

- Labour and capital would be complements:

```
x_range = arange(0.01,1.01,0.01)
plt.plot(x_range, kfuc2(x_range,x_range,x_range*1.8,1,0.5105,1,0.89,0.13099), label="Rich country")
plt.plot(x_range, kfuc2(x_range,x_range,x_range*1.8,1,1,1,0.89,0.83684), label="Poor country")
plt.legend(loc=2)
plt.xlim(0.1,1)
plt.show()
```

### Case 2: High types hire less workers

```
x_range = arange(0.01,1.01,0.01)
plt.plot(x_range, kfuc2(x_range,x_range,(1/x_range)**0.5,1,0.5105,1,0.89,0.13099), label="Rich country")
plt.plot(x_range, kfuc2(x_range,x_range,(1/x_range)**0.5,1,1,1,0.89,0.83684), label="Poor country")
plt.legend(loc=2)
plt.xlim(0.1,1)
plt.show()
```

### Local effects of changing $x,y$ and $l$

- Increasing x,y and l from (1,1,1), rich country

```
x_range = arange(1,1.6,0.01)
plt.plot(x_range, kfuc2(x_range,1,1,1,1,1,0.89,0.13099), label="Increasing $x$")
plt.plot(x_range, kfuc2(1,x_range,1,1,1,1,0.89,0.13099), label="Increasing $y$")
plt.plot(x_range, kfuc2(1,1,x_range,1,1,1,0.89,0.13099), label="Increasing $l$")
plt.legend(loc=2)
plt.xlim(1,1.5)
plt.title("From (1,1,1), the effects on $k^{*}$ of:")
plt.show()
```

- Increasing x,y and l from (1,1,1), poor country (similar change, but less magnitude)

```
x_range = arange(1,1.6,0.01)
plt.plot(x_range, kfuc2(x_range,1,1,1,0.5105,1,0.89,0.83684), label="Increasing $x$")
plt.plot(x_range, kfuc2(1,x_range,1,1,0.5105,1,0.89,0.83684), label="Increasing $y$")
plt.plot(x_range, kfuc2(1,1,x_range,1,0.5105,1,0.89,0.83684), label="Increasing $l$")
plt.legend(loc=2)
plt.xlim(1,1.5)
plt.title("From (1,1,1), the effects on $k^{*}$ of:")
plt.show()
```

- Increasing y and l from 0 to 2, keeping the rest =1 ,rich country

```
x_range = arange(0.01,2,0.01)
plt.plot(x_range, kfuc2(1,x_range,1,1,1,1,0.89,0.13099), label="$y$"+str(i))
plt.plot(x_range, kfuc2(1,1,x_range,1,1,1,0.89,0.13099), label="$l$"+str(i))
plt.show()
```

- Increasing y and l from 0 to 2, keeping the rest =1 ,poor country

```
x_range = arange(0.01,2,0.01)
plt.plot(x_range, kfuc2(1,x_range,1,1,0.5105,1,0.89,0.83684), label="$y$"+str(i))
plt.plot(x_range, kfuc2(1,1,x_range,1,0.5105,1,0.89,0.83684), label="$l$"+str(i))
plt.show()
```

- Increasing $x$ from 0 to 2 and keeping the rest constant at 1 (much bigger, linear effects)

```
plt.plot(x_range, kfuc2(x_range,1,1,1,1,1,0.89,0.13099), label="Rich country")
plt.plot(x_range, kfuc2(x_range,1,1,1,0.5105,1,0.89,0.83684), label="Poor country")
plt.xlim(0.01,2)
plt.legend(loc=2)
plt.show()
```

```

```


A very wide range of physical processes lead to wave motion, where
signals are propagated through a medium in space and time, normally
with little or no permanent movement of the medium itself.
The shape of the signals may undergo changes as they travel through
matter, but usually not so much that the signals cannot be recognized
at some later point in space and time.
Many types of wave motion can be described by the equation
$u_{tt}=\nabla\cdot (c^2\nabla u) + f$, which we will solve
in the forthcoming text by finite difference methods.

# Simulation of waves on a string

We begin our study of wave equations by simulating one-dimensional
waves on a string, say on a guitar or violin.
Let the string in the undeformed state
coincide with the interval
$[0,L]$ on the $x$ axis, and let $u(x,t)$ be the displacement at
time $t$ in the $y$ direction of a point initially at $x$.
The displacement function $u$ is governed by the mathematical model

$$
\begin{equation}
\frac{\partial^2 u}{\partial t^2} =
c^2 \frac{\partial^2 u}{\partial x^2}, \quad x\in (0,L),\ t\in (0,T]
\label{wave:pde1} \tag{1}
\end{equation}
$$

$$
\begin{equation}
u(x,0) = I(x), \quad x\in [0,L]
\label{wave:pde1:ic:u} \tag{2}
\end{equation}
$$

$$
\begin{equation}
\frac{\partial}{\partial t}u(x,0) = 0, \quad x\in [0,L]
\label{wave:pde1:ic:ut} \tag{3}
\end{equation}
$$

$$
\begin{equation}
u(0,t)  = 0, \quad  t\in (0,T]
\label{wave:pde1:bc:0} \tag{4}
\end{equation}
$$

$$
\begin{equation}
u(L,t)  = 0, \quad  t\in (0,T]
\label{wave:pde1:bc:L} \tag{5}
\end{equation}
$$

The constant $c$ and the function $I(x)$ must be prescribed.

Equation ([1](#wave:pde1)) is known as the one-dimensional
*wave equation*. Since this PDE contains a second-order derivative
in time, we need *two initial conditions*. The condition
([2](#wave:pde1:ic:u)) specifies
the initial shape of the string, $I(x)$, and
([3](#wave:pde1:ic:ut)) expresses that the initial velocity of the
string is zero. In addition, PDEs need *boundary conditions*, given here as
([4](#wave:pde1:bc:0)) and ([5](#wave:pde1:bc:L)). These two
conditions specify that
the string is fixed at the ends, i.e., that the displacement $u$ is zero.

The solution $u(x,t)$ varies in space and time and describes waves that
move with velocity $c$ to the left and right.

Sometimes we will use a more compact notation for the partial derivatives
to save space:

$$
\begin{equation}
u_t = \frac{\partial u}{\partial t}, \quad
u_{tt} = \frac{\partial^2 u}{\partial t^2},
\label{_auto1} \tag{6}
\end{equation}
$$

and similar expressions
for derivatives with respect to other variables. Then the
wave equation can be written compactly as $u_{tt} = c^2u_{xx}$.

The PDE problem ([1](#wave:pde1))-([5](#wave:pde1:bc:L)) will now be
discretized in space and time by a finite difference method.

## Discretizing the domain

The temporal domain $[0,T]$ is represented by a finite number of mesh points

$$
\begin{equation}
0 = t_0 < t_1 < t_2 < \cdots < t_{N_t-1} < t_{N_t} = T    \label{_auto2} \tag{7}
\end{equation}
$$

Similarly, the spatial domain $[0,L]$ is replaced by a set of mesh points

$$
\begin{equation}
0 = x_0 < x_1 < x_2 < \cdots < x_{N_x-1} < x_{N_x} = L   \label{_auto3} \tag{8}
\end{equation}
$$

One may view the mesh as two-dimensional in the $x,t$ plane, consisting
of points $(x_i, t_n)$, with $i=0,\ldots,N_x$ and $n=0,\ldots,N_t$.

### Uniform meshes

For uniformly distributed mesh points we can introduce the constant
mesh spacings $\Delta t$ and $\Delta x$. We have that

$$
\begin{equation}
x_i = i\Delta x,\ i=0,\ldots,N_x,\quad
t_n = n\Delta t,\ n=0,\ldots,N_t
\label{_auto4} \tag{9}
\end{equation}
$$

We also have that $\Delta x = x_i-x_{i-1}$, $i=1,\ldots,N_x$, and
$\Delta t = t_n - t_{n-1}$, $n=1,\ldots,N_t$. [Figure](#wave:pde1:fig:mesh)
displays a mesh in the $x,t$ plane with $N_t=5$, $N_x=5$, and constant
mesh spacings.

## The discrete solution

The solution $u(x,t)$ is sought at the mesh points. We introduce
the mesh function $u_i^n$, which approximates the exact
solution at the
mesh point $(x_i,t_n)$ for $i=0,\ldots,N_x$ and $n=0,\ldots,N_t$.
Using the finite difference method, we shall
develop algebraic equations for computing the mesh function.

## Fulfilling the equation at the mesh points

In the finite difference method, we relax
the condition that ([1](#wave:pde1)) holds at all points in
the space-time domain $(0,L)\times (0,T]$ to the requirement that the PDE is
fulfilled at the *interior* mesh points only:

$$
\begin{equation}
\frac{\partial^2}{\partial t^2} u(x_i, t_n) =
c^2\frac{\partial^2}{\partial x^2} u(x_i, t_n),
\label{wave:pde1:step2} \tag{10}
\end{equation}
$$

for $i=1,\ldots,N_x-1$ and $n=1,\ldots,N_t-1$. For $n=0$ we have
the initial conditions $u=I(x)$ and $u_t=0$,
and at the boundaries $i=0,N_x$ we
have the boundary condition $u=0$.

## Replacing derivatives by finite differences

The second-order derivatives can be replaced by central
differences. The most widely used difference approximation of
the second-order derivative is

$$
\frac{\partial^2}{\partial t^2}u(x_i,t_n)\approx
\frac{u_i^{n+1} - 2u_i^n + u^{n-1}_i}{\Delta t^2}\
$$

mathcal{I}_t is convenient to introduce the finite difference operator notation

$$
[D_tD_t u]^n_i = \frac{u_i^{n+1} - 2u_i^n + u^{n-1}_i}{\Delta t^2}
$$

A similar approximation of the second-order derivative in the $x$
direction reads

$$
\frac{\partial^2}{\partial x^2}u(x_i,t_n)\approx
\frac{u_{i+1}^{n} - 2u_i^n + u^{n}_{i-1}}{\Delta x^2} = [D_xD_x u]^n_i
$$

### Algebraic version of the PDE

We can now replace the derivatives in ([10](#wave:pde1:step2))
and get

$$
\begin{equation}
\frac{u_i^{n+1} - 2u_i^n + u^{n-1}_i}{\Delta t^2} =
c^2\frac{u_{i+1}^{n} - 2u_i^n + u^{n}_{i-1}}{\Delta x^2},
\label{wave:pde1:step3b} \tag{11}
\end{equation}
$$

or written more compactly using the operator notation:

$$
\begin{equation}
[D_tD_t u = c^2 D_xD_x]^{n}_i
\label{wave:pde1:step3a} \tag{12}
\end{equation}
$$

### Interpretation of the equation as a stencil

A characteristic feature of ([11](#wave:pde1:step3b)) is that it
involves $u$ values from neighboring points only: $u_i^{n+1}$,
$u^n_{i\pm 1}$, $u^n_i$, and $u^{n-1}_i$.  The circles in [Figure](#wave:pde1:fig:mesh) illustrate such neighboring mesh points that
contribute to an algebraic equation. In this particular case, we have
sampled the PDE at the point $(2,2)$ and constructed
([11](#wave:pde1:step3b)), which then involves a coupling of $u_1^2$,
$u_2^3$, $u_2^2$, $u_2^1$, and $u_3^2$.  The term *stencil* is often
used about the algebraic equation at a mesh point, and the geometry of
a typical stencil is illustrated in [Figure](#wave:pde1:fig:mesh). One also often refers to the algebraic
equations as *discrete equations*, *(finite) difference equations* or
a *finite difference scheme*.

Mesh in space and time. The circles show points connected in a finite difference equation.

### Algebraic version of the initial conditions

We also need to replace the derivative in the initial condition
([3](#wave:pde1:ic:ut)) by a finite difference approximation.
A centered difference of the type

$$
\frac{\partial}{\partial t} u(x_i,t_0)\approx
\frac{u^1_i - u^{-1}_i}{2\Delta t} = [D_{2t} u]^0_i,
$$

seems appropriate. Writing out this equation and ordering the terms give

$$
\begin{equation}
u^{-1}_i=u^{1}_i,\quad i=0,\ldots,N_x
\label{wave:pde1:step3c} \tag{13}
\end{equation}
$$

The other initial condition can be computed by

$$
u_i^0 = I(x_i),\quad i=0,\ldots,N_x
$$

## Formulating a recursive algorithm

We assume that $u^n_i$ and $u^{n-1}_i$ are available for
$i=0,\ldots,N_x$.  The only unknown quantity in
([11](#wave:pde1:step3b)) is therefore $u^{n+1}_i$, which we now can
solve for:

$$
\begin{equation}
u^{n+1}_i = -u^{n-1}_i + 2u^n_i + C^2
\left(u^{n}_{i+1}-2u^{n}_{i} + u^{n}_{i-1}\right)
\label{wave:pde1:step4} \tag{14}
\end{equation}
$$

We have here introduced the parameter

$$
\begin{equation}
C = c\frac{\Delta t}{\Delta x},
\label{_auto5} \tag{15}
\end{equation}
$$

known as the *Courant number*.

**$C$ is the key parameter in the discrete wave equation.**

We see that the discrete version of the PDE features only one
parameter, $C$, which is therefore the key parameter, together with
$N_x$, that governs the quality of the numerical solution (see the section [Analysis of the difference equations](wave_analysis.ipynb) for details). Both the primary physical
parameter $c$ and the numerical parameters $\Delta x$ and $\Delta t$
are lumped together in $C$. Note that $C$ is a dimensionless
parameter.

Given that $u^{n-1}_i$ and $u^n_i$ are known for $i=0,\ldots,N_x$,
we find new values at the next time level by applying the formula
([14](#wave:pde1:step4)) for $i=1,\ldots,N_x-1$. [Figure](#wave:pde1:fig:mesh) illustrates the points that are used to
compute $u^3_2$. For the boundary points, $i=0$ and $i=N_x$, we apply
the boundary conditions $u_i^{n+1}=0$.

Even though sound reasoning leads up to
([14](#wave:pde1:step4)), there is still a minor challenge with it that needs
to be resolved. Think of the very first computational step to be made.
The scheme ([14](#wave:pde1:step4)) is supposed to start at $n=1$, which means
that we compute $u^2$ from $u^1$ and $u^0$. Unfortunately, we do not know the
value of $u^1$, so how to proceed? A standard procedure in such cases is to
apply ([14](#wave:pde1:step4)) also for $n=0$. This immediately seems strange,
since it involves $u^{-1}_i$, which is an undefined
quantity outside the time mesh (and the time domain). However, we can
use the initial condition ([13](#wave:pde1:step3c)) in combination with
([14](#wave:pde1:step4)) when $n=0$ to eliminate $u^{-1}_i$ and
arrive at a special formula for $u_i^1$:

$$
\begin{equation}
u_i^1 = u^0_i - \frac{1}{2}
C^2\left(u^{0}_{i+1}-2u^{0}_{i} + u^{0}_{i-1}\right)
\label{wave:pde1:step4:1} \tag{16}
\end{equation}
$$

[Figure](#wave:pde1:fig:stencil:u1) illustrates how ([16](#wave:pde1:step4:1))
connects four instead of five points: $u^1_2$, $u_1^0$, $u_2^0$, and $u_3^0$.

Modified stencil for the first time step.

We can now summarize the computational algorithm:

1. Compute $u^0_i=I(x_i)$ for $i=0,\ldots,N_x$

2. Compute $u^1_i$ by ([16](#wave:pde1:step4:1)) for $i=1,2,\ldots,N_x-1$ and set $u_i^1=0$
   for the boundary points given by $i=0$ and $i=N_x$,

3. For each time level $n=1,2,\ldots,N_t-1$

    a. apply ([14](#wave:pde1:step4)) to find $u^{n+1}_i$ for $i=1,\ldots,N_x-1$

    b. set $u^{n+1}_i=0$ for the boundary points having $i=0$, $i=N_x$.

The algorithm essentially consists of moving a finite difference
stencil through all the mesh points, which can be seen as an animation
in a [web page](mov-wave/D_stencil_gpl/index.html)
or a [movie file](mov-wave/D_stencil_gpl/movie.ogg).

## Sketch of an implementation

We start by defining some constants that will be used throughout our Devito code.

```python
import numpy as np

# Given mesh points as arrays x and t (x[i], t[n]),
# constant c and function I for initial condition
x = np.linspace(0, 2, 101)
t = np.linspace(0, 2, 101)
c = 1
I = lambda x: np.sin(x)

dx = x[1] - x[0]
dt = t[1] - t[0]
C = c*dt/dx            # Courant number
Nx = len(x)-1
Nt = len(t)-1
C2 = C**2              # Help variable in the scheme
L = 2.
```

Next, we define our 1D computational grid and create a function `u` as a symbolic `devito.TimeFunction`. We need to specify the `space_order` as 2 since our wave equation involves second-order derivatives with respect to $x$. Similarly, we specify the `time_order` as 2, as our equation involves second-order derivatives with respect to $t$. Setting these parameters allows us to use `u.dx2` and `u.dt2`.

```python
from devito import Grid, TimeFunction

# Initialise `u` for space and time order 2, using initialisation function I
grid = Grid(shape=(Nx+1), extent=(L))
u = TimeFunction(name='u', grid=grid, time_order=2, space_order=2)
u.data[:,:] = I(x[:])
```

Now that we have initialised `u`, we can solve our wave equation for the unknown quantity $u^{n+1}_i$ using forward and backward differences in space and time.

```python
from devito import Constant, Eq, solve

# Set up wave equation and solve for forward stencil point in time
pde = (1/c**2)*u.dt2-u.dx2
stencil = Eq(u.forward, solve(pde, u.forward))

print("LHS: %s" % stencil.lhs)
print("RHS: %s" % stencil.rhs)
```

    LHS: u(t + dt, x)
    RHS: 1.0*dt**2*(-2.0*u(t, x)/h_x**2 + u(t, x - h_x)/h_x**2 + u(t, x + h_x)/h_x**2 + 2.0*u(t, x)/dt**2 - 1.0*u(t - dt, x)/dt**2)

Great! From these print statements, we can see that Devito has taken the wave equation in ([1](#wave:pde1)) and solved it for $u^{n+1}_i$, giving us equation ([14](#wave:pde1:step4)). Note that `dx` is denoted as `h_x`, while `u(t, x)`, `u(t, x - h_x)` and `u(t, x + h_x)` denote the equivalent of $u^{n}_{i}$, $u^{n}_{i-1}$ and $u^{n}_{i+1}$ respectively.

We also need to create a separate stencil for the first timestep, where we substitute $u^{1}_i$ for $u^{-1}_i$, as given in ([13](#wave:pde1:step3c)).

```python
stencil_init = stencil.subs(u.backward, u.forward)
```

Now we can create expressions for our boundary conditions and build the operator. The results are plotted below.

```python
#NBVAL_IGNORE_OUTPUT
from devito import Operator

t_s = grid.stepping_dim

# Boundary conditions
bc = [Eq(u[t_s+1, 0], 0)]
bc += [Eq(u[t_s+1, Nx], 0)]

# Defining one Operator for initial timestep and one for the rest
op_init = Operator([stencil_init]+bc)
op = Operator([stencil]+bc)

op_init.apply(time_M=1, dt=dt)
op.apply(time_m=1,time_M=Nt, dt=dt)

```

    Data type float64 of runtime value `dt` does not match the Constant data type <class 'numpy.float32'>
    Operator `Kernel` run in 0.01 s
    Data type float64 of runtime value `dt` does not match the Constant data type <class 'numpy.float32'>
    Operator `Kernel` run in 0.01 s

    PerformanceSummary([(PerfKey(name='section0', rank=None),
                         PerfEntry(time=5.899999999999998e-05, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])

We can plot our results using `matplotlib`:

```python
import matplotlib.pyplot as plt

plt.plot(x, u.data[-1])
plt.xlabel('x')
plt.ylabel('u')
plt.show()
```

# Verification

Before implementing the algorithm, it is convenient to add a source
term to the PDE ([1](#wave:pde1)), since that gives us more freedom in
finding test problems for verification. Physically, a source term acts
as a generator for waves in the interior of the domain.

## A slightly generalized model problem

We now address the following extended initial-boundary value problem
for one-dimensional wave phenomena:

$$
\begin{equation}
u_{tt} = c^2 u_{xx} + f(x,t), \quad x\in (0,L),\ t\in (0,T]
\label{wave:pde2} \tag{17}
\end{equation}
$$

$$
\begin{equation}
u(x,0) = I(x), \quad x\in [0,L]
\label{wave:pde2:ic:u} \tag{18}
\end{equation}
$$

$$
\begin{equation}
u_t(x,0) = V(x), \quad x\in [0,L]
\label{wave:pde2:ic:ut} \tag{19}
\end{equation}
$$

$$
\begin{equation}
u(0,t)  = 0, \quad  t>0
\label{wave:pde2:bc:0} \tag{20}
\end{equation}
$$

$$
\begin{equation}
u(L,t)  = 0, \quad  t>0
\label{wave:pde2:bc:L} \tag{21}
\end{equation}
$$

Sampling the PDE at $(x_i,t_n)$ and using the same finite difference
approximations as above, yields

$$
\begin{equation}
[D_tD_t u = c^2 D_xD_x u + f]^{n}_i
\label{wave:pde2:fdop} \tag{22}
\end{equation}
$$

Writing this out and solving for the unknown $u^{n+1}_i$ results in

$$
\begin{equation}
u^{n+1}_i = -u^{n-1}_i + 2u^n_i + C^2
(u^{n}_{i+1}-2u^{n}_{i} + u^{n}_{i-1}) + \Delta t^2 f^n_i
\label{wave:pde2:step3b} \tag{23}
\end{equation}
$$

The equation for the first time step must be rederived. The discretization
of the initial condition $u_t = V(x)$ at $t=0$
becomes

$$
[D_{2t}u = V]^0_i\quad\Rightarrow\quad u^{-1}_i = u^{1}_i - 2\Delta t V_i,
$$

which, when inserted in ([23](#wave:pde2:step3b)) for $n=0$, gives
the special formula

$$
\begin{equation}
u^{1}_i = u^0_i + \Delta t V_i + {\frac{1}{2}}
C^2
\left(u^{0}_{i+1}-2u^{0}_{i} + u^{0}_{i-1}\right) + \frac{1}{2}\Delta t^2 f^0_i
\label{wave:pde2:step3c} \tag{24}
\end{equation}
$$

## Using an analytical solution of physical significance

Many wave problems feature sinusoidal oscillations in time
and space. For example, the original PDE problem
([1](#wave:pde1))-([5](#wave:pde1:bc:L)) allows an exact solution

$$
\begin{equation}
u_e(x,t) = A\sin\left(\frac{\pi}{L}x\right)
\cos\left(\frac{\pi}{L}ct\right)
\label{wave:pde2:test:ue} \tag{25}
\end{equation}
$$

This $u_e$ fulfills the PDE with $f=0$, boundary conditions
$u_e(0,t)=u_e(L,t)=0$, as well as initial
conditions $I(x)=A\sin\left(\frac{\pi}{L}x\right)$ and $V=0$.

**How to use exact solutions for verification.**

It is common to use such exact solutions of physical interest
to verify implementations. However, the numerical
solution $u^n_i$ will only be an approximation to $u_e(x_i,t_n)$.
We have no knowledge of the precise size of the error in
this approximation, and therefore we can never know if discrepancies
between $u^n_i$ and $u_e(x_i,t_n)$ are caused
by mathematical approximations or programming errors.
In particular, if plots of the computed solution $u^n_i$ and
the exact one ([25](#wave:pde2:test:ue)) look similar, many
are tempted to claim that the implementation works. However,
even if color plots look nice and the accuracy is "deemed good",
there can still be serious programming errors present!

The only way to use exact physical solutions like
([25](#wave:pde2:test:ue)) for serious and thorough verification is to
run a series of simulations on finer and finer meshes, measure the
integrated error in each mesh, and from this information estimate the
empirical convergence rate of the method.

An introduction to the computing of convergence rates is given in Section 3.1.6
 in [[Langtangen_decay]](#Langtangen_decay).
There is also a detailed example on computing convergence rates in
the [verification section](../01_vib/vib_undamped.ipynb#vib:ode1:verify) of the Vibration ODEs chapter.

In the present problem, one expects the method to have a convergence rate
of 2 (see the section [Analysis of the difference equations](wave_analysis.ipynb)), so if the computed rates
are close to 2 on a sufficiently fine mesh, we have good evidence that
the implementation is free of programming mistakes.

## Manufactured solution and estimation of convergence rates

### Specifying the solution and computing corresponding data

One problem with the exact solution ([25](#wave:pde2:test:ue)) is
that it requires a simplification (${V}=0, f=0$) of the implemented problem
([17](#wave:pde2))-([21](#wave:pde2:bc:L)). An advantage of using
a *manufactured solution* is that we can test all terms in the
PDE problem. The idea of this approach is to set up some chosen
solution and fit the source term, boundary conditions, and initial
conditions to be compatible with the chosen solution.
Given that our boundary conditions in the implementation are
$u(0,t)=u(L,t)=0$, we must choose a solution that fulfills these
conditions. One example is

$$
u_e(x,t) = x(L-x)\sin t
$$

Inserted in the PDE $u_{tt}=c^2u_{xx}+f$ we get

$$
-x(L-x)\sin t = -c^2 2\sin t + f\quad\Rightarrow f = (2c^2 - x(L-x))\sin t
$$

The initial conditions become

$$
\begin{align*}
u(x,0) =& I(x) = 0,\\
u_t(x,0) &= V(x) = x(L-x)
\end{align*}
$$

### Defining a single discretization parameter

To verify the code, we compute the convergence rates in a series of
simulations, letting each simulation use a finer mesh than the
previous one. Such empirical estimation of convergence rates relies on
an assumption that some measure $E$ of the numerical error is related
to the discretization parameters through

$$
E = C_t\Delta t^r + C_x\Delta x^p,
$$

where $C_t$, $C_x$, $r$, and $p$ are constants. The constants $r$ and
$p$ are known as the *convergence rates* in time and space,
respectively.  From the accuracy in the finite difference
approximations, we expect $r=p=2$, since the error terms are of order
$\Delta t^2$ and $\Delta x^2$.  This is confirmed by truncation error
analysis and other types of analysis.

By using an exact solution of the PDE problem, we will next compute
the error measure $E$ on a sequence of refined meshes and see if
the rates $r=p=2$ are obtained. We will not be concerned with estimating
the constants $C_t$ and $C_x$, simply because we are not interested in
their values.

mathcal{I}_t is advantageous to introduce a single discretization parameter
$h=\Delta t=\hat c \Delta x$ for some constant $\hat c$.  Since
$\Delta t$ and $\Delta x$ are related through the Courant number,
$\Delta t = C\Delta x/c$, we set $h=\Delta t$, and then $\Delta x =
hc/C$.  Now the expression for the error measure is greatly
simplified:

$$
E = C_t\Delta t^r + C_x\Delta x^r =
C_t h^r + C_x\left(\frac{c}{C}\right)^r h^r
= Dh^r,\quad D = C_t+C_x\left(\frac{c}{C}\right)^r
$$

### Computing errors

We choose an initial discretization parameter $h_0$ and run
experiments with decreasing $h$: $h_i=2^{-i}h_0$, $i=1,2,\ldots,m$.
Halving $h$ in each experiment is not necessary, but it is a common
choice.  For each experiment we must record $E$ and $h$.  Standard
choices of error measure are the $\ell^2$ and $\ell^\infty$ norms of the
error mesh function $e^n_i$:

$$
\begin{equation}
E = ||e^n_i||_{\ell^2} = \left( \Delta t\Delta x
\sum_{n=0}^{N_t}\sum_{i=0}^{N_x}
(e^n_i)^2\right)^{\frac{1}{2}},\quad e^n_i = u_e(x_i,t_n)-u^n_i,
\label{wave:pde2:fd:MMS:E:l2} \tag{26}
\end{equation}
$$

$$
\begin{equation}
E = ||e^n_i||_{\ell^\infty} = \max_{i,n} |e^n_i|
\label{wave:pde2:fd:MMS:E:linf} \tag{27}
\end{equation}
$$

In Python, one can compute $\sum_{i}(e^{n}_i)^2$ at each time step
and accumulate the value in some sum variable, say `e2_sum`.  At the
final time step one can do `sqrt(dt*dx*e2_sum)`.  For the
$\ell^\infty$ norm one must compare the maximum error at a time level
(`e.max()`) with the global maximum over the time domain: `e_max = max(e_max, e.max())`.

An alternative error measure is to use a spatial norm at one time step
only, e.g., the end time $T$ ($n=N_t$):

$$
\begin{equation}
E = ||e^n_i||_{\ell^2} = \left( \Delta x\sum_{i=0}^{N_x}
(e^n_i)^2\right)^{\frac{1}{2}},\quad e^n_i = u_e(x_i,t_n)-u^n_i,
\label{_auto6} \tag{28}
\end{equation}
$$

$$
\begin{equation}
E = ||e^n_i||_{\ell^\infty} = \max_{0\leq i\leq N_x} |e^{n}_i|
\label{_auto7} \tag{29}
\end{equation}
$$

The important point is that the error measure ($E$) for the simulation is represented by a single number.

### Computing rates

Let $E_i$ be the error measure in experiment (mesh) number $i$
(not to be confused with the spatial index $i$) and
let $h_i$ be the corresponding discretization parameter ($h$).
With the error model $E_i = Dh_i^r$, we can
estimate $r$ by comparing two consecutive
experiments:

$$
\begin{align*}
E_{i+1}& =D h_{i+1}^{r},\\
E_{i}& =D h_{i}^{r}
\end{align*}
$$

Dividing the two equations eliminates the (uninteresting) constant $D$.
Thereafter, solving for $r$ yields

$$
r = \frac{\ln E_{i+1}/E_{i}}{\ln h_{i+1}/h_{i}}
$$

Since $r$ depends on $i$, i.e., which simulations we compare,
we add an index to $r$: $r_i$, where $i=0,\ldots,m-2$, if we
have $m$ experiments: $(h_0,E_0),\ldots,(h_{m-1}, E_{m-1})$.

In our present discretization of the wave equation we expect $r=2$, and
hence the $r_i$ values should converge to 2 as $i$ increases.

## Constructing an exact solution of the discrete equations

With a manufactured or known analytical solution, as outlined above,
we can estimate convergence rates and see if they have the correct
asymptotic behavior. Experience shows that this is a quite good
verification technique in that many common bugs will destroy the
convergence rates. A significantly better test though,
would be to check that the
numerical solution is exactly what it should be. This will in general
require exact knowledge of the numerical error, which we do not normally have
(although we in the section [Analysis of the difference equations](wave_analysis.ipynb) establish such knowledge
in simple cases).
However, it is possible to look for solutions where we can show that
the numerical error vanishes, i.e., the solution of the original continuous
PDE problem is
also a solution of the discrete equations. This property often arises
if the exact solution of the PDE
is a lower-order polynomial. (Truncation error
analysis leads to error measures that involve derivatives of the
exact solution. In the present problem, the truncation error involves
4th-order derivatives of $u$ in space and time. Choosing $u$
as a polynomial of degree three or less
will therefore lead to vanishing error.)

We shall now illustrate the construction of an exact solution to both the
PDE itself and the discrete equations.
Our chosen manufactured solution is quadratic in space
and linear in time. More specifically, we set

$$
\begin{equation}
u_e (x,t) = x(L-x)(1+{\frac{1}{2}}t),
\label{wave:pde2:fd:verify:quadratic:uex} \tag{30}
\end{equation}
$$

which by insertion in the PDE leads to $f(x,t)=2(1+t)c^2$. This $u_e$
fulfills the boundary conditions $u=0$ and demands $I(x)=x(L-x)$
and $V(x)={\frac{1}{2}}x(L-x)$.

To realize that the chosen $u_e$ is also an exact
solution of the discrete equations,
we first remind ourselves that $t_n=n\Delta t$ so that

$$
\begin{equation}
\lbrack D_tD_t t^2\rbrack^n = \frac{t_{n+1}^2 - 2t_n^2 + t_{n-1}^2}{\Delta t^2}
= (n+1)^2 -2n^2 + (n-1)^2 = 2,
\label{_auto8} \tag{31}
\end{equation}
$$

$$
\begin{equation}
\lbrack D_tD_t t\rbrack^n = \frac{t_{n+1} - 2t_n + t_{n-1}}{\Delta t^2}
= \frac{((n+1) -2n + (n-1))\Delta t}{\Delta t^2} = 0
\label{_auto9} \tag{32}
\end{equation}
$$

Hence,

$$
[D_tD_t u_e]^n_i = x_i(L-x_i)[D_tD_t (1+{\frac{1}{2}}t)]^n =
x_i(L-x_i){\frac{1}{2}}[D_tD_t t]^n = 0
$$

Similarly, we get that

$$
\begin{align*}
\lbrack D_xD_x u_e\rbrack^n_i &=
(1+{\frac{1}{2}}t_n)\lbrack D_xD_x (xL-x^2)\rbrack_i\\
& =
(1+{\frac{1}{2}}t_n)\lbrack LD_xD_x x - D_xD_x x^2\rbrack_i \\
&= -2(1+{\frac{1}{2}}t_n)
\end{align*}
$$

Now, $f^n_i = 2(1+{\frac{1}{2}}t_n)c^2$, which results in

$$
[D_tD_t u_e - c^2D_xD_xu_e - f]^n_i = 0 +
c^2 2(1 + {\frac{1}{2}}t_{n}) +
2(1+{\frac{1}{2}}t_n)c^2 = 0
$$

Moreover, $u_e(x_i,0)=I(x_i)$,
$\partial u_e/\partial t = V(x_i)$ at $t=0$, and
$u_e(x_0,t)=u_e(x_{N_x},0)=0$. Also the modified scheme for the
first time step is fulfilled by $u_e(x_i,t_n)$.

Therefore, the exact solution $u_e(x,t)=x(L-x)(1+t/2)$ of the PDE
problem is also an exact solution of the discrete problem.  This means
that we know beforehand what numbers the numerical algorithm should
produce.  We can use this fact to check that the computed $u^n_i$
values from an implementation equals $u_e(x_i,t_n)$, within machine
precision.  This result is valid *regardless of the mesh spacings*
$\Delta x$ and $\Delta t$!  Nevertheless, there might be stability
restrictions on $\Delta x$ and $\Delta t$, so the test can only be run
for a mesh that is compatible with the stability criterion (which in
the present case is $C\leq 1$, to be derived later).

**Notice.**

A product of quadratic or linear expressions in the various
independent variables, as shown above, will often fulfill both the
PDE problem and the discrete equations, and can therefore be very useful
solutions for verifying implementations.

However, for 1D wave
equations of the type $u_{tt}=c^2u_{xx}$ we shall see that there is always
another much more powerful way of generating exact
solutions (which consists in just setting $C=1$ (!), as shown in
the section [Analysis of the difference equations](wave_analysis.ipynb)).


```python
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
import matplotlib.cm as cm
from IPython.display import display, Math, clear_output

import sympy
from sympy import *
from sympy.physics.vector import ReferenceFrame, CoordinateSym
from sympy.vector import CoordSys3D, divergence, curl

import ipyvolume as ipv
import time
np.seterr(divide='ignore', invalid='ignore')
```

    {'divide': 'warn', 'over': 'warn', 'under': 'ignore', 'invalid': 'warn'}

## Sympy Basics:
1. Unit vectors are indicated by a . (For example, R.x is the unit vector in the x direction).
2. Scalars are indicated by [] (For example, R[0] is x)

## First let's import packages

```python
## Check versions
### Need at least Sympy Version 1.2 (from conda-forge)
```

```python
print('Ipyvolume version = {}'.format(ipv.__version__))
print('Sympy version = {}'.format(sympy.__version__))
```

    Ipyvolume version = 0.4.6
    Sympy version = 1.2

```python
## Now let's checking that plotting and the plotting widget work
```

```python
# init_printing()
plt.figure()
plt.plot([1,2,3])
plt.show()
```

## Installation notes
Sympy1.2 is from conda-forge:
1. sudo conda install -c conda-forge sympy


# Stiff ODEs and implicit methods
In this notebook we look at examples where Runge-Kutta methods require very small steps to be accurate. For these **stiff** ODEs implicit methods are better.

Let's define our explicit methods first.

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

```python
# The below commands make the font and image size bigger
plt.rcParams.update({'font.size': 22})
plt.rcParams["figure.figsize"] = (15,10)
```

```python
def EulerStep(f, dx, xi, yi):
    return yi + dx*f(xi, yi)
```

```python
def RK2Step(f, dx, xi, yi):
    k1 = dx*f(xi, yi)
    k2 = dx*f(xi + dx, yi + k1)

    return yi + 0.5*(k1 + k2)
```

```python
def RK4Step(f, dx, xi, yi):
        k1 = dx*f(xi,yi)
        k2 = dx*f(xi + 0.5*dx, yi + 0.5*k1)
        k3 = dx*f(xi + 0.5*dx, yi + 0.5*k2)
        k4 = dx*f(xi + dx, yi + k3)

        return yi + 1/6*(k1 + 2*k2 + 2*k3 + k4)
```

```python
def ODESolve(f, dx, x0, y0, imax, method='RK4', plotSteps=False):

    xi = x0
    yi = y0

    # Create arrays to store the steps in
    steps = np.zeros((imax+1,2))
    steps[0,0] = x0
    steps[0,1] = y0

    i = 0
    while i < imax:
        if(method == 'RK4'):
            yi = RK4Step(f, dx, xi, yi)
        elif(method == 'RK2'):
            yi = RK2Step(f, dx, xi, yi)
        elif(method == 'Euler'):
            yi = EulerStep(f, dx, xi, yi)

        xi += dx
        i  += 1

        # Store the steps for plotting
        steps[i, 0] = xi
        steps[i, 1] = yi

    if(plotSteps):
        plt.scatter(steps[:,0], steps[:,1], color='red', linewidth=10)

    return [xi, yi]
```

# Stiff ODEs

Let's look at an example. Consider the *first-order* ODE:

$$ y'(x) = \lambda[-y(x) + \sin(x)] $$

with *initial conditions* $y(0) = 0$. This has the general solution of

$$ y(x) = C e^{-\lambda x} + \frac{\lambda^2 \sin(x) - \lambda \cos(x)}{1 + \lambda^2}$$

The initial condition gives

$$ C = \frac{\lambda}{1 + \lambda^2}$$

Let's define these equations and see how our RK methods do at solving the ODE. Try using the 'Euler', 'RK2' and 'RK4' methods below.

```python
lam = 30

def dydxStiff(x,y):
    global lam
    return lam*(-y + np.sin(x))

def yStiff(x):
    global lam
    C = lam/(1+lam**2)
    return C*np.exp(-lam*x) + (lam**2*np.sin(x) -lam*np.cos(x))/(1+lam**2)
```

```python
x = np.linspace(0, 2.0, 100)
y = yStiff(x)
plt.xlabel('x')
plt.ylabel('y')
plt.grid(True)
plt.plot(x,y)

ODESolve(dydxStiff, 0.1, 0, 0, 10, 'RK4', True)
plt.legend(['Exact solution', 'Explicit RK4']);
```

We see that even the more accurate RK4 method does not do well. It can be made to perform better by decreasing the step size but we want to avoid this. Furthermore, if we increase $\lambda$ the step size required to accurately obtain the solution drops rapidly.

To overcome this, let's look at implicit methods. One of the simplest is the Backwards Euler method.

## Implicit method: backwards Euler

The backwards Euler method is given by:

$$ \begin{align}
    x_{i+1} &= x_i + \Delta x \\
    y_{i+1} &= y_i + \Delta x f(x_{i+1}, y_{i+1}) \\
    \end{align}
    $$

Notice how, unlike the *explicit* methods we've looked at until now this does not immediately give a recipe for the next step in terms of the previous steps. We must solve the second equation for $y_{i+1}$. In general we won't be able to find an analytic solution, so we might use *numerical root finding*.

Let's look at the stiff ODE example given above, where it turns out we can find an analytic solution. In this case we have:

$$ \begin{align}
    x_{i+1} &= x_i + \Delta x \\
    y_{i+1} &= y_i + \Delta x \lambda[-y_{i+1} + \sin(x_{i+1})]
\end{align}$$

In this case we can solve the second equation explicitly for $y_{i+1}$:

$$ y_{i+1} = \frac{y_i + \Delta x \lambda \sin(x_{i+1})}{1 + \Delta x \lambda}$$

Let's now implement this *specific* example in code and test it out.

```python
def BackwardsEulerSpecific(dx, x0, y0, imax, plotSteps=False):
    global lam

    xi = x0
    yi = y0

    # Create arrays to store the steps in
    steps = np.zeros((imax+1,2))
    steps[0,0] = x0
    steps[0,1] = y0

    i = 0
    while i < imax:
        yi = (yi + dx * lam*np.sin(xi+dx))/(1+ dx*lam)
        xi += dx
        i += 1

        # Store the steps for plotting
        steps[i, 0] = xi
        steps[i, 1] = yi

    if(plotSteps):
        plt.scatter(steps[:,0], steps[:,1], color='green', linewidth=10)

    return [xi, yi]
```

```python
x = np.linspace(0, 2.0, 100)
y = yStiff(x)
plt.xlabel('x')
plt.ylabel('y')
plt.grid(True)
plt.plot(x,y)

ODESolve(dydxStiff, 0.1, 0, 0, 10, 'RK4', True)

BackwardsEulerSpecific(0.2, 0, 0, 10, True)

plt.legend(['Exact solution', 'Explicit RK4', 'Implicit backwards Euler']);
```

We can thus see that the (implicit) backwards Euler method is much more successful than the (explicit) RK4 method, even when it takes larger steps. This is great because each step of the RK4 method is more expensive than a step of the Euler method. The downside to implicit methods can be seen from the code above: our implicit ODE solver code is not general, it only works for this specific example. In general, implicit methods are more complicated to implement as they often require numerical methods, e.g., root finding to find $y_{i+1}$.

Note, unlike the forward Euler method, the backwards Euler method is stable across a much wider class of ODEs.

```python

```


# The minimum jerk hypothesis

> Marcos Duarte
> Laboratory of Biomechanics and Motor Control ([http://demotu.org/](http://demotu.org/))
> Federal University of ABC, Brazil

Hogan and Flash (1984, 1985), based on observations of voluntary movements in primates, suggested that movements are performed (organized) with the smoothest trajectory possible. In this organizing principle, the endpoint trajectory is such that the mean squared-jerk across time of this movement is minimum.

Jerk is the derivative of acceleration and the observation of the minimum-jerk trajectory is for the endpoint in the extracorporal coordinates (not for joint angles) and according to Flash and Hogan (1985), the minimum-jerk trajectory of a planar movement is such that minimizes the following objective function:

$$ C=\frac{1}{2} \int\limits_{t_{i}}^{t_{f}}\;\left[\left(\frac{d^{3}x}{dt^{3}}\right)^2+\left(\frac{d^{3}y}{dt^{3}}\right)^2\right]\:\mathrm{d}t $$

Hogan (1984) found that the solution for this objective function is a fifth-order polynomial trajectory (see Shadmehr and Wise (2004) for a simpler proof):

$$ \begin{array}{l l}
x(t) = a_0+a_1t+a_2t^2+a_3t^3+a_4t^4+a_5t^5 \\
y(t) = b_0+b_1t+b_2t^2+b_3t^3+b_4t^4+b_5t^5
\end{array} $$

With the following boundary conditions for $ x(t) $ and $ y(t) $: initial and final positions are $ (x_i,y_i) $ and $ (x_f,y_f) $ and initial and final velocities and accelerations are zero.

Let's employ [Sympy](http://sympy.org/en/index.html) to find the solution for the minimum jerk trajectory using symbolic algebra.

```python
import numpy as np
%matplotlib inline
import matplotlib.pyplot as plt
from IPython.display import display, Math, Latex
from sympy import symbols, Matrix, latex, Eq, collect, solve, diff, simplify
from sympy.utilities.lambdify import lambdify
```

Using Sympy, the equation for minimum jerk trajectory for x is:

```python
# declare the symbolic variables
x, xi, xf, y, yi, yf, d, t = symbols('x, x_i, x_f, y, y_i, y_f, d, t')
a0, a1, a2, a3, a4, a5 = symbols('a_0:6')
x = a0 + a1*t + a2*t**2 + a3*t**3 + a4*t**4 + a5*t**5
display(Math(latex('x(t)=') + latex(x)))
```

$$x(t)=a_{0} + a_{1} t + a_{2} t^{2} + a_{3} t^{3} + a_{4} t^{4} + a_{5} t^{5}$$

Without loss of generality, consider $ t_i=0 $ and let's use $ d $ for movement duration ($ d=t_f $). The system of equations with the boundary conditions for $ x $ is:

```python
# define the system of equations
s = Matrix([Eq(x.subs(t,0)          , xi),
            Eq(diff(x,t,1).subs(t,0),  0),
            Eq(diff(x,t,2).subs(t,0),  0),
            Eq(x.subs(t,d)          , xf),
            Eq(diff(x,t,1).subs(t,d),  0),
            Eq(diff(x,t,2).subs(t,d),  0)])
display(Math(latex(s, mat_str='matrix', mat_delim='[')))
```

$$\left[\begin{matrix}a_{0} = x_{i}\\a_{1} = 0\\2 a_{2} = 0\\a_{0} + a_{1} d + a_{2} d^{2} + a_{3} d^{3} + a_{4} d^{4} + a_{5} d^{5} = x_{f}\\a_{1} + 2 a_{2} d + 3 a_{3} d^{2} + 4 a_{4} d^{3} + 5 a_{5} d^{4} = 0\\2 a_{2} + 6 a_{3} d + 12 a_{4} d^{2} + 20 a_{5} d^{3} = 0\end{matrix}\right]$$

Which gives the following solution:

```python
# algebraically solve the system of equations
sol = solve(s, [a0, a1, a2, a3, a4, a5])
display(Math(latex(sol)))
```

$$\left \{ a_{0} : x_{i}, \quad a_{1} : 0, \quad a_{2} : 0, \quad a_{3} : \frac{10}{d^{3}} \left(x_{f} - x_{i}\right), \quad a_{4} : \frac{15}{d^{4}} \left(- x_{f} + x_{i}\right), \quad a_{5} : \frac{6}{d^{5}} \left(x_{f} - x_{i}\right)\right \}$$

Substituting this solution in the fifth order polynomial trajectory equation, we have the actual displacement trajectories:

```python
# substitute the equation parameters by the solution
x2 = x.subs(sol)
x2 = collect(simplify(x2, ratio=1), xf-xi)
display(Math(latex('x(t)=') + latex(x2)))
y2 = x2.subs([(xi, yi), (xf, yf)])
display(Math(latex('y(t)=') + latex(y2)))
```

$$x(t)=x_{i} + \left(x_{f} - x_{i}\right) \left(\frac{10 t^{3}}{d^{3}} - \frac{15 t^{4}}{d^{4}} + \frac{6 t^{5}}{d^{5}}\right)$$

$$y(t)=y_{i} + \left(y_{f} - y_{i}\right) \left(\frac{10 t^{3}}{d^{3}} - \frac{15 t^{4}}{d^{4}} + \frac{6 t^{5}}{d^{5}}\right)$$

And for the velocity, acceleration, and jerk trajectories in x:

```python
# symbolic differentiation
vx = x2.diff(t, 1)
display(Math(latex('v_x(t)=') + latex(vx)))
ax = x2.diff(t, 2)
display(Math(latex('a_x(t)=') + latex(ax)))
jx = x2.diff(t, 3)
display(Math(latex('j_x(t)=') + latex(jx)))
```

$$v_x(t)=\left(x_{f} - x_{i}\right) \left(\frac{30 t^{2}}{d^{3}} - \frac{60 t^{3}}{d^{4}} + \frac{30 t^{4}}{d^{5}}\right)$$

$$a_x(t)=\frac{60 t}{d^{3}} \left(x_{f} - x_{i}\right) \left(1 - \frac{3 t}{d} + \frac{2 t^{2}}{d^{2}}\right)$$

$$j_x(t)=\frac{60}{d^{3}} \left(x_{f} - x_{i}\right) \left(1 - \frac{6 t}{d} + \frac{6 t^{2}}{d^{2}}\right)$$

Let's plot the minimum jerk trajectory for x and its velocity, acceleration, and jerk considering $x_i=0,x_f=1,d=1$:

```python
# substitute by the numerical values
x3  = x2.subs([(xi, 0), (xf, 1), (d, 1)])
#create functions for calculation of numerical values
xfu = lambdify(t, diff(x3, t, 0), 'numpy')
vfu = lambdify(t, diff(x3, t, 1), 'numpy')
afu = lambdify(t, diff(x3, t, 2), 'numpy')
jfu = lambdify(t, diff(x3, t, 3), 'numpy')
#plots using matplotlib
ts = np.arange(0, 1.01, .01)
fig, axs = plt.subplots(1, 4, figsize=(12, 5), sharex=True, squeeze=True)
axs[0].plot(ts, xfu(ts), linewidth=3)
axs[0].set_title('Displacement [$\mathrm{m}$]')
axs[1].plot(ts, vfu(ts), linewidth=3)
axs[1].set_title('Velocity [$\mathrm{m/s}$]')
axs[2].plot(ts, afu(ts), linewidth=3)
axs[2].set_title('Acceleration [$\mathrm{m/s^2}$]')
axs[3].plot(ts, jfu(ts), linewidth=3)
axs[3].set_title('Jerk [$\mathrm{m/s^3}$]')

for axi in axs:
    axi.set_xlabel('Time [s]', fontsize=14)
    axi.grid(True)

fig.suptitle('Minimum jerk trajectory kinematics', fontsize=20, y=1.03)
fig.tight_layout()
plt.show()
```

Note that for the minimum jerk trajectory, initial and final values of both velocity and acceleration are zero, but not for the jerk.

Read more about the minimum jerk trajectory hypothesis in the [Shadmehr and Wise's book companion site](http://www.shadmehrlab.org/book/minimum_jerk/minimumjerk.htm) and in [Paul Gribble's website](http://www.gribblelab.org/compneuro/4_Computational_Motor_Control_Kinematics.html#sec-5-1).

### The angular trajectory of a minimum jerk trajectory

Let's calculate the resulting angular trajectory given a minimum jerk linear trajectory, supposing it is from a circular motion of an elbow flexion. The length of the forearm is 0.5 m, the movement duration is 1 s, the elbow starts flexed at 90$^o$ and the flexes to 180$^o$.

First, the linear trajectories for this circular motion:

```python
# substitute by the numerical values
x3 = x2.subs([(xi, 0.5), (xf, 0), (d, 1)])
y3 = x2.subs([(xi, 0), (xf, 0.5), (d, 1)])
display(Math(latex('y(t)=') + latex(x3)))
display(Math(latex('x(t)=') + latex(y3)))
#create functions for calculation of numerical values
xfux = lambdify(t, diff(x3, t, 0), 'numpy')
vfux = lambdify(t, diff(x3, t, 1), 'numpy')
afux = lambdify(t, diff(x3, t, 2), 'numpy')
jfux = lambdify(t, diff(x3, t, 3), 'numpy')
xfuy = lambdify(t, diff(y3, t, 0), 'numpy')
vfuy = lambdify(t, diff(y3, t, 1), 'numpy')
afuy = lambdify(t, diff(y3, t, 2), 'numpy')
jfuy = lambdify(t, diff(y3, t, 3), 'numpy')
```

$$y(t)=- 3.0 t^{5} + 7.5 t^{4} - 5.0 t^{3} + 0.5$$

$$x(t)=3.0 t^{5} - 7.5 t^{4} + 5.0 t^{3}$$

```python
#plots using matplotlib
ts = np.arange(0, 1.01, .01)
fig, axs = plt.subplots(1, 4, figsize=(12, 5), sharex=True, squeeze=True)
axs[0].plot(ts, xfux(ts), 'b', linewidth=3)
axs[0].plot(ts, xfuy(ts), 'r', linewidth=3)
axs[0].set_title('Displacement [$\mathrm{m}$]')
axs[1].plot(ts, vfux(ts), 'b', linewidth=3)
axs[1].plot(ts, vfuy(ts), 'r', linewidth=3)
axs[1].set_title('Velocity [$\mathrm{m/s}$]')
axs[2].plot(ts, afux(ts), 'b', linewidth=3)
axs[2].plot(ts, afuy(ts), 'r', linewidth=3)
axs[2].set_title('Acceleration [$\mathrm{m/s^2}$]')
axs[3].plot(ts, jfux(ts), 'b', linewidth=3)
axs[3].plot(ts, jfuy(ts), 'r', linewidth=3)
axs[3].set_title('Jerk [$\mathrm{m/s^3}$]')

for axi in axs:
    axi.set_xlabel('Time [s]', fontsize=14)
    axi.grid(True)

fig.suptitle('Minimum jerk trajectory kinematics', fontsize=20, y=1.03)
fig.tight_layout()
plt.show()
```

Now, the angular trajectories for this circular motion:

```python
from sympy import atan2
ang = atan2(y3, x3)*180/np.pi
display(Math(latex('angle(t)=') + latex(ang)))
xang = lambdify(t, diff(ang, t, 0), 'numpy')
vang = lambdify(t, diff(ang, t, 1), 'numpy')
aang = lambdify(t, diff(ang, t, 2), 'numpy')
jang = lambdify(t, diff(ang, t, 3), 'numpy')
```

$$angle(t)=57.2957795130823 \operatorname{atan_{2}}{\left (3.0 t^{5} - 7.5 t^{4} + 5.0 t^{3},- 3.0 t^{5} + 7.5 t^{4} - 5.0 t^{3} + 0.5 \right )}$$

```python
ts = np.arange(0, 1.01, .01)
fig, axs = plt.subplots(1, 4, figsize=(12, 5), sharex=True, squeeze=True)
axs[0].plot(ts, xang(ts), linewidth=3)
axs[0].set_title('Displacement [$\mathrm{m}$]')
axs[1].plot(ts, vang(ts), linewidth=3)
axs[1].set_title('Velocity [$\mathrm{m/s}$]')
axs[2].plot(ts, aang(ts), linewidth=3)
axs[2].set_title('Acceleration [$\mathrm{m/s^2}$]')
axs[3].plot(ts, jang(ts), linewidth=3)
axs[3].set_title('Jerk [$\mathrm{m/s^3}$]')

for axi in axs:
    axi.set_xlabel('Time [s]', fontsize=14)
    axi.grid(True)

fig.suptitle('Minimum jerk trajectory angular kinematics', fontsize=20, y=1.03)
fig.tight_layout()
plt.show()
```

## Problems

1. What is your opinion on the the minimum jerk hypothesis? Do you think humans control movement based on this principle? (Think about what biomechanical and neurophysiological properties are not considered on this hypothesis.)
2. Calculate and plot the position, velocity, acceleration, and jerk trajectories for different movement speeds (for example, consider always a displacement of 1 m and movement durations of 0.5, 1, and 2 s).
3. For the data in the previous item, calculate the ratio peak speed to average speed. Shadmehr and  Wise (2004) argue that psychophysical experiments show that reaching movements with the hand have this ratio equals to 1.75. Compare with the calculated values.
4. Can you propose alternative hypotheses for the control of movement?

## References

- Flash T, Hogan N (1985) [The coordination of arm movements: an experimentally confirmed mathematical model](http://www.jneurosci.org/cgi/reprint/5/7/1688.pdf). Journal of Neuroscience, 5, 1688-1703.
- Hogan N (1984) [An organizing principle for a class of voluntary movements](http://www.jneurosci.org/content/4/11/2745.full.pdf). Journal of Neuroscience, 4, 2745-2754.
- Shadmehr R, Wise S (2004) [The Computational Neurobiology of Reaching and Pointing: A Foundation for Motor Learning](http://www.shadmehrlab.org/book/). A Bradford Book. [Companion site](http://www.shadmehrlab.org/book/).
- Zatsiorsky VM (1998) [Kinematics of Human Motion](http://books.google.com.br/books/about/Kinematics_of_Human_Motion.html?id=Pql_xXdbrMcC&redir_esc=y). Champaign, Human Kinetics.


```python
from sympy import *
from calc import Symbol
init_printing()
```

## Simple Addition

```python
a = Symbol('a', val=1, err=0.01)
b = Symbol('b', val=2, err=0.01)
r1 = a + b
r1.simp()
r1.display(name="r1")
```

$\displaystyle
        \begin{align*}
            \textbf{Numeric: } \\
                              r1 &= \ 3.00000000000000 \pm 0.0141421356237310 \\
            \textbf{Symbolic: }& \\
            \text{Value: } r1 &= a + b \\
                                  &= 3 \\
            \text{Error: } \sigma_{r1}&= \sqrt{\sigma_{a}^{2} + \sigma_{b}^{2}} \\
                                  &= 0.014142135623731 \\
        \end{align*}
            $

## Multiplication

```python
r1 = a * b
r1.simp()
r1.display(name="r1")
```

$\displaystyle
        \begin{align*}
            \textbf{Numeric: } \\
                              r1 &= \ 2.00000000000000 \pm 0.0223606797749979 \\
            \textbf{Symbolic: }& \\
            \text{Value: } r1 &= a b \\
                                  &= 2 \\
            \text{Error: } \sigma_{r1}&= a b \sqrt{\frac{\sigma_{b}^{2}}{b^{2}} + \frac{\sigma_{a}^{2}}{a^{2}}} \\
                                  &= 0.0223606797749979 \\
        \end{align*}
            $

## Exponentiation

```python
r1 = a ** (1/2)
r1.simp()
r1.display(name="r1")
```

$\displaystyle
        \begin{align*}
            \textbf{Numeric: } \\
                              r1 &= \ 1.00000000000000 \pm 0.00500000000000000 \\
            \textbf{Symbolic: }& \\
            \text{Value: } r1 &= \sqrt{a} \\
                                  &= 1 \\
            \text{Error: } \sigma_{r1}&= \frac{a \sigma_{a}}{2} \\
                                  &= 0.005 \\
        \end{align*}
            $

## Standard error of the mean

```python
mean, err = Symbol.std_err_of_mean(1.3, 1.5, 1.7)
mean, err
```

## Miscellaneous

```python
r1 = a / b + a
r1.simp()
r1.display(name="r1")
```

$\displaystyle
        \begin{align*}
            \textbf{Numeric: } \\
                              r1 &= \ 1.50000000000000 \pm 0.0114564392373896 \\
            \textbf{Symbolic: }& \\
            \text{Value: } r1 &= a + \frac{a}{b} \\
                                  &= \frac{3}{2} \\
            \text{Error: } \sigma_{r1}&= \sqrt{\frac{a^{2} \sigma_{b}^{2}}{b^{4}} + \sigma_{a}^{2} + \frac{\sigma_{a}^{2}}{b^{2}}} \\
                                  &= 0.0114564392373896 \\
        \end{align*}
            $

## A more complex example

```python
U1 = Symbol('U_1', val=0.888, err=0.007)
U2 = Symbol('U_2', val=0.203, err=0.002)
V = Symbol('V', val=5.637, err=0.001)

C1 = (U1 - (U1**2 - 4*U1*U2)**0.5) / (V**2)
C2 = (U1 + (U1**2 - 4*U1*U2)**0.5) / (V**2)

C1.simp()
C2.simp()
```

```python
C1.display(name="C_1")
C2.display(name="C_2")
```

$\displaystyle
        \begin{align*}
            \textbf{Numeric: } \\
                              C_1 &= \ 0.0197702681531304 \pm 0.000221011614823886 \\
            \textbf{Symbolic: }& \\
            \text{Value: } C_1 &= \frac{U_{1} - \sqrt{U_{1} \left(U_{1} - 4 U_{2}\right)}}{V^{2}} \\
                                  &= 0.0197702681531304 \\
            \text{Error: } \sigma_{C_1}&= \frac{\sqrt{\frac{V^{2} \left(U_{1}^{2} \left(U_{1} - 4 U_{2}\right)^{2} \left(U_{1}^{2} \sigma_{U 1}^{2} + 4 U_{1}^{2} \sigma_{U 2}^{2} + 4 U_{2}^{2} \sigma_{U 1}^{2}\right) + \sigma_{U 1}^{2}\right) + 4 \sigma_{V}^{2} \left(U_{1} - \sqrt{U_{1} \left(U_{1} - 4 U_{2}\right)}\right)^{2}}{V^{2} \left(U_{1} - \sqrt{U_{1} \left(U_{1} - 4 U_{2}\right)}\right)^{2}}} \left(U_{1} - \sqrt{U_{1} \left(U_{1} - 4 U_{2}\right)}\right)}{V^{2}} \\
                                  &= 0.000221011614823886 \\
        \end{align*}
            $

$\displaystyle
        \begin{align*}
            \textbf{Numeric: } \\
                              C_2 &= \ 0.0361213768295607 \pm 0.000221271725056850 \\
            \textbf{Symbolic: }& \\
            \text{Value: } C_2 &= \frac{U_{1} + \sqrt{U_{1} \left(U_{1} - 4 U_{2}\right)}}{V^{2}} \\
                                  &= 0.0361213768295607 \\
            \text{Error: } \sigma_{C_2}&= \frac{\sqrt{\frac{V^{2} \left(U_{1}^{2} \left(U_{1} - 4 U_{2}\right)^{2} \left(U_{1}^{2} \sigma_{U 1}^{2} + 4 U_{1}^{2} \sigma_{U 2}^{2} + 4 U_{2}^{2} \sigma_{U 1}^{2}\right) + \sigma_{U 1}^{2}\right) + 4 \sigma_{V}^{2} \left(U_{1} + \sqrt{U_{1} \left(U_{1} - 4 U_{2}\right)}\right)^{2}}{V^{2} \left(U_{1} + \sqrt{U_{1} \left(U_{1} - 4 U_{2}\right)}\right)^{2}}} \left(U_{1} + \sqrt{U_{1} \left(U_{1} - 4 U_{2}\right)}\right)}{V^{2}} \\
                                  &= 0.00022127172505685 \\
        \end{align*}
            $

```python

```


# Fourier Series

Fourier series are a method we can use to solve inhomogeneous 2nd-order ODEs of the form
\begin{equation}
y^{\prime\prime} + p(t) y^{\prime} + q(t) y = r(t) \;,
\end{equation}
where the forcing function $r(t)$ is periodic. This means looking like one of these examples:

```matlab
% periodic square wave
subplot(3,1,1);
squareWave = repmat([1,1,1,1,0,0,0,0], [1, 10]);
t = linspace(0, 1, length(squareWave));
plot(t, squareWave); ylim([-0.5, 1.5]);

% sin wave
subplot(3,1,2);
t = linspace(0, 1, 100);
plot(t, sin(t*10*pi)); ylim([-1.5, 1.5]);

% sawtooth
subplot(3,1,3);
t = linspace(0, 1, 100);
y = ((mod(t,2*pi/40)/(pi*2/40))*2)-1;
plot(t, y); ylim([-1.5, 1.5]);
```

(Actually, as we'll see later, we can use a Fourier series to represent generic forcing function!)

Fourier series have been around a while, ever since in 1790 Jean-Baptiste Joseph Fourier found that generic periodic functions could be represented by a sum of series of `sin()` and `cos()` functions, harmonically related by frequency.

In general, a Fourier series represents a function $f(t)$ by
\begin{equation}
f(t) = a_0 + \sum_{n=1}^{\infty} a_n \cos (n \omega t) + \sum_{n=1}^{\infty} b_n \sin (n \omega t)
\end{equation}
where $a_0$, $a_n$, and $b_n$ are the Fourier coefficients, $\omega = \frac{2\pi}{T}$ is the frequency of the function $f(t)$, and $T$ is the period. $n$ is an integer used as an index.

## Properties of Fourier Series

Considering that $n$ is an integer, and the sine and cosine components of a Fourier series share the same fundamental frequency $\omega$, Fourier series have some useful properties:

1. The integral of each component trigonometric function over the period is zero:
\begin{equation}
\int_0^T \sin (n \omega t) dt = 0 = \int_0^T cos (n \omega t) dt
\end{equation}

2. The component trigonometric functions are **orthogonal** over their period:
\begin{equation}
\int_0^T \cos(n \omega t) \sin (m \omega t) dt = 0
\end{equation}
for all values of $n, m = 1, 2, \ldots, \infty$.

3. The component trigonometric functions are also orthogonal with themselves over their period:
\begin{align}
\int_0^T \cos (n \omega t) \cos (m \omega t) dt &= \begin{cases}0 \quad m \neq n \\ \frac{T}{2} \quad m = n \end{cases} \\
%
\int_0^T \sin (n \omega t) \sin (m \omega t) dt &= \begin{cases}0 \quad m \neq n \\ \frac{T}{2} \quad m = n \end{cases}
\end{align}
for all values of $n, m = 1, 2, \ldots, \infty$.

## Fourier coefficients

We can use the above properties to calculate the Fourier coefficients, given a periodic function $f(t)$.
First, recall
\begin{equation}
f(t) = a_0 + \sum_{n=1}^{\infty} a_n \cos (n \omega t) + \sum_{n=1}^{\infty} b_n \sin (n \omega t)
\end{equation}

1. To calculate $a_0$, integrate both sides of the equation over the period:
\begin{equation}
\int_0^T f(t) dt = \int_0^T a_0 dt + \int_0^T \left( \sum_{n=1}^{\infty} a_n \cos (n \omega t) \right)dt + \int_0^T \left( \sum_{n=1}^{\infty} b_n \sin (n \omega t) \right) dt
\end{equation}
For the integrals of the infinite sums, recall that the integral of the sum of some functions is the same as the sum of the integrals of the functions: $\int (a + b + c) = \int a + \int b + \int c$. That means that
\begin{align}
\int_0^T \left( \sum_{n=1}^{\infty} a_n \cos (n \omega t) \right)dt = \int_0^T a_1 \cos (\omega t) dt + \int_0^T a_2 \cos (2 \omega t) dt + \ldots &= 0 \;, \text{ and} \\
\int_0^T \left( \sum_{n=1}^{\infty} b_n \sin (n \omega t) \right)dt = \int_0^T b_1 \sin (\omega t) dt + \int_0^T b_2 \sin (2 \omega t) dt + \ldots &= 0 \;,
\end{align}
since the integrals of the trigonometric functions are all zero over the period. Thus,
\begin{equation}
a_0 = \frac{1}{T} \int_0^T f(t) dt
\end{equation}

2. To calculate $a_n$, multiply both sides of the equation by $\cos(m \omega t)$ and integrate over the period:
\begin{equation}
\int_0^T f(t) \cos(m \omega t) dt = a_0 \int_0^T \cos(m \omega t) dt + \int_0^T \left( \sum_{n=1}^{\infty} a_n \cos (n \omega t) \cos(m \omega t) \right)dt + \int_0^T \left( \sum_{n=1}^{\infty} b_n \sin (n \omega t) \cos(m \omega t) \right) dt
\end{equation}
Let's take a look at each of the three integrals on the right-hand side. First,
\begin{equation}
a_0 \int_0^T \cos(m \omega t) dt = 0
\end{equation}
because it just integrates cosine over the period.
Skipping to the last term,
\begin{equation}
\int_0^T \left( \sum_{n=1}^{\infty} b_n \sin (n \omega t) \cos(m \omega t) \right) dt = b_1 \int_0^T sin(\omega t) \cos (m \omega t) dt + b_2 \int_0^T \sin (2 \omega t) \cos (m \omega t) dt + \ldots = 0
\end{equation}
due to orthogonality. We are just left with the middle integral; let's expand a few terms to see what that looks like:
\begin{equation}
\int_0^T \left( \sum_{n=1}^{\infty} a_n \cos (n \omega t) \cos(m \omega t) \right)dt = a_1 \int_0^T \cos (\omega t) \cos (m \omega t) dt + a_2 \int_0^T \cos (2 \omega t) \cos (m \omega t) dt + \ldots
\end{equation}
Again, due to orthogonality, all of the terms of this infinite sum of integrals will be zero, *except* for the term where $n = m$. As a result, we are left with
\begin{align}
\int_0^T f(t) \cos(m \omega t) dt &= \int_0^T a_m \cos^2 (m \omega t) dt = a_m \frac{T}{2} \\
a_n = a_m &= \frac{2}{T} \int_0^T f(t) \cos (n \omega t) dt
\end{align}

3. We can find $b_n$ in the same way, multiplying the equation by $\sin (m \omega t)$ and integrating over the period. The work is the same, so let's skip that:
\begin{equation}
b_n = \frac{2}{T} \int_0^T f(t) \sin (n \omega t) dt
\end{equation}

## Example: periodic rectangular wave

Let's find the Fourier series for representing this periodic function $f(t)$:

```matlab
x = [0 1 1 2 2 3 3 4]; y = [2 2 1 1 2 2 1 1];
plot(x,y); ylim([0 2.5]); ylabel('f(t)'); xlabel('t')
```

First, we need to identify the fundamental period and frequency: $T = 2$ and then $\omega = \frac{2\pi}{T} = \pi$. Our work is then to calculate the Fourier coefficients. Since our periodic function $f(t)$ is not easily expressed with a function–hence the need for a Fourier series—we'll use piecewise integration.

First, calculate $a_0$:
\begin{align}
a_0 =& \frac{1}{T} \int_0^T f(t) dt \\
&= \frac{1}{2} \int_0^2 f(t) dt = \frac{1}{2}\left( \int_0^1 2dt + \int_1^2 1dt \right) = \frac{1}{2} (2\times 1 + 1 \times 1) \\
a_0 &= \frac{3}{2}
\end{align}

Then, get $a_n$:
\begin{align}
a_n &= \frac{2}{T} \int_0^T f(t) \cos (n \omega t) dt \\
&= \frac{2}{2} \int_0^2 f(t) \cos (n \pi t) dt = \left( \int_0^1 2 \cos (n \pi t) dt + \int_1^2 1 \cos (n \pi t)dt \right) \\
&= \frac{2}{n \pi} \sin(n \pi t)\Big|_0^1 + \frac{1}{n\pi} \sin(n \pi t)\Big|_1^2 \\
&= \frac{2}{n \pi}\left(\sin(n\pi) - \sin(0)\right) + \frac{1}{n\pi}\left( sin(2n\pi) - \sin(n\pi)\right) \\
a_n &= 0
\end{align}

Finally, we can calculate $b_n$:
\begin{align}
b_n &= \frac{2}{T} \int_0^T f(t) \sin (n \omega t) dt \\
&= \frac{2}{2} \int_0^2 f(t) \sin (n \pi t) dt = \left( \int_0^1 2 \sin (n \pi t) dt + \int_1^2 1 \sin (n \pi t)dt \right) \\
&= -\frac{2}{n \pi} \cos(n \pi t)\Big|_0^1 - \frac{1}{n\pi} \cos(n \pi t)\Big|_1^2 \\
&= -\frac{2}{n \pi}\left(\cos(n\pi) - \cos(0)\right) - \frac{1}{n\pi}\left( cos(2n\pi) - \cos(n\pi)\right) \\
b_n &= -\frac{2}{n \pi}\left(\cos(n\pi) - 1\right) - \frac{1}{n\pi}\left( 1 - \cos(n\pi)\right) = -\frac{1}{n\pi}\left( \cos(n\pi) - 1\right)
\end{align}
but recall that $n = 1, 2, \ldots, \infty$. Then,
\begin{align}
b_n &= -\frac{1}{n\pi} \times \begin{cases} -2 \text{ if } n \text{ odd} \\0 \text{ if } n \text{ even}\end{cases} \\
\rightarrow b_n &= \frac{2}{n\pi} \quad n = \text{odd}
\end{align}

Then, our Fourier series representation for the function shown above is
\begin{equation}
f(t) = \frac{3}{2} + \sum_{\substack{n = 1 \\n = \text{odd}}}^{\infty} \frac{2}{n\pi} \sin (n \pi t)
\end{equation}

Now, let's see how whether this actually works! Let's start with one term of the infinite sum, then gradually increase.

```matlab
t = linspace(0, 4);

% maximum number of terms
n_max = [1, 2, 3, 5, 10, 25, 50, 250, 500];

for i = 1 : length(n_max)
    N = n_max(i);

    s = 3./2.;
    for n = 1 : 2 : 2*N
        s = s + (2. / (n*pi)) .* sin(n * pi * t);
    end
    subplot(3, 3, i)
    plot(t, s); axis([0 4 0 3]); title(sprintf('%d terms', N));
end
```

As we increase the number of terms, adding higher-frequency sine waves, we are better able to match the original rectangular wave. Notice the discrepancies that remain near the sharp corners even after the rest of the series closely resembles the function: these are known as **Gibbs phenomena**, caused by the Fourier series overshooting or undershooting (or "ringing") near discontinuities.

## Even and Odd Functions

We can simplify our work generating a Fourier series if we can identify the given periodic function $f(t)$ as an **even function** or an **odd function**.

Even functions are those where $f(-x) = f(x)$.

Odd functions are those where $f(-x) = -f(x)$.

```matlab
y = [-1 0 1 0 -1];
x = [-2 -1 0 1 2];
subplot(2,1,1); plot(x,y); title('Even function');
ax = gca; ax.XAxisLocation = 'origin'; ax.YAxisLocation = 'origin';

y = [0 -1 0 1 0];
subplot(2,1,2); plot(x,y); title('Odd function');
ax = gca; ax.XAxisLocation = 'origin'; ax.YAxisLocation = 'origin';
```

An even function's Fourier series simplifies to a Fourier cosine series:
\begin{align}
f(x) &= a_0 + \sum_{n=1}^{\infty} a_n \cos (n \omega x) dx \\
a_0 &= \frac{2}{T} \int_0^{T/2} f(x) dx \\
a_n &= \frac{4}{T} \int_0^{T/2} f(x) \cos(n \omega x) dx
\end{align}

An odd function's Fourier series simplifies to a Fourier sine series:
\begin{align}
f(x) &= \sum_{n=1}^{\infty} b_n \sin (n \omega x) dx \\
b_n &= \frac{4}{T} \int_0^{T/2} f(x) \sin(n \omega x) dx
\end{align}

Note: not all periodic functions can be considered an even or an odd function.

## Application: Inhomogeneous 2nd-order ODE

One way we might use a Fourier series is to solve an inhomogeneous 2nd-order ODE, where the forcing term is given by a periodic function not easily expressed using our typical functions.

### Undamped mass-spring system

For example, let's consider an undamped mass-spring system, where the forcing is given by a periodic rectangular wave:
\begin{equation}
y^{\prime\prime} + 4y = f(t)
\end{equation}
where the forcing function $f(t)$ is

```matlab
t = [0 0 1 1 2 2 3 3];
y = [0 1 1 -1 -1 1 1 0];
plot(t,y); ylim([-1.5 1.5]);
title('Periodic rectangular forcing function');
ax = gca; ax.XAxisLocation = 'origin';
```

Using recognizing this as an odd function, we could construct a Fourier sine series to represent the forcing function:
\begin{equation}
f(t) = \sum_{\substack{n=1\\n=\text{odd}}}^{\infty} \frac{4}{n\pi} \sin(n \pi t)
\end{equation}

Let's confirm this works:

```matlab
t = linspace(0, 3, 500);
s = 0;
for n = 1 : 2 : 1000
    s = s + (4/(n*pi)).*sin(n*pi*t);
end
plot(t, s)
```

Looks good!

To find the exact solution for our displacement $y(t)$, we can follow our usual analytical solution approach: find the homogeneous solution $y_H$, then find the inhomogeneous solution $y_{IH}$; the overall solution is then $y(t) = y_H + y_{IH}$. The homogeneous solution is
\begin{equation}
y_H = c_1 \sin (2t) + c_2 \cos (2t)
\end{equation}
We then find the inhomogeneous solution using
\begin{equation}
y^{\prime\prime} + 4y = \frac{4}{n\pi} \sin (n \pi t) \quad n = 1, 3, \ldots, \infty
\end{equation}
Solving this will give us a specific $y_{IH, n}$; the complete inhomogeneous solution is then
\begin{equation}
y_{IH} = \sum_{\substack{n=1\\n=\text{odd}}}^{\infty} y_{IH, n} \;.
\end{equation}

Recognizing that our forcing function is sinusoidal, we should use the method of undetermined coefficients:
\begin{equation}
y_{IH, n} = K_1 \sin (n \pi t) + K_2 \cos (n \pi t)
\end{equation}
Inserting this into the above ODE gives
\begin{align}
K_1 &= \frac{4}{n \pi (4 - n^2 \pi^2)} \\
K_2 &= 0
\end{align}

Thus, the overall solution is
\begin{equation}
y(t) = c_1 \sin(2t) + c_2 \cos(2t) + \sum_{\substack{n=1\\n=\text{odd}}}^{\infty} \frac{4}{n \pi (4 - n^2 \pi^2)} \sin (n \pi t)
\end{equation}

### Damped mass-spring system

What about a damped mass-spring system? Recall that the homogeneous solution could take one of these three forms:
\begin{align}
y_H &= c_1 e^{-\lambda_1 t} + c_2 e^{-\lambda_2 t} \\
y_H &= c_1 e^{-\lambda_1 t} + c_2 t e^{-\lambda_2 t} \text{ or} \\
y_H &= e^{-\alpha t} (c_1 \sin(\beta t) + c_2 \cos(\beta t))
\end{align}
while the inhomogeneous solution, given a Fourier series forcing function, will take the form
\begin{equation}
y_{IH} = K_1 \sin() + K_2 \cos()
\end{equation}

The overall solution combines the homogenenous and inhomogeneous solutions. But, the homogeneous solution in this case is **transient**, because it decays to zero. On the other hand, the inhomogeneous solution remains, and is the **steady-state solution**.


[Source 1](https://www.geeksforgeeks.org/contingency-table-in-python/)

[Source 2](https://www.statology.org/contingency-table-python/)

Our HELLO WORD example

```python
import numpy as np
import pandas as pd
```

```python
x = np.array([[1, 0],
              [1, 0],
              [0, 1],
              [0, 1],
              [1, 1],
              [1, 1],
              [0, 0],
              [0, 1]])
b = np.array([0, 1, 0, 1, 0, 1, 0, 0])
# b is a line precursor of y. We transpose it.
y = b.reshape(-1,1)
```

```python
# trandform arrays to dataframe
df1 = pd.DataFrame(x, columns = ['X1', 'X2'])
df2 = pd.DataFrame(y, columns = ['Y'])
df1
```

      X1
      X2

      0
      1
      0

      1
      1
      0

      2
      0
      1

      3
      0
      1

      4
      1
      1

      5
      1
      1

      6
      0
      0

      7
      0
      1

```python
df2
```

      Y

      0
      0

      1
      1

      2
      0

      3
      1

      4
      0

      5
      1

      6
      0

      7
      0

```python
# merge dataframe x and y
df1['Y'] = df2
```

```python
df1
```

      X1
      X2
      Y

      0
      1
      0
      0

      1
      1
      0
      1

      2
      0
      1
      0

      3
      0
      1
      1

      4
      1
      1
      0

      5
      1
      1
      1

      6
      0
      0
      0

      7
      0
      1
      0

```python
# unique line counter for dataframe
dz = df1.groupby(['X1', 'X2', 'Y']).size().reset_index(name='Count')
# Without reset index, something strange is happening
dz
```

      X1
      X2
      Y
      Count

      0
      0
      0
      0
      1

      1
      0
      1
      0
      2

      2
      0
      1
      1
      1

      3
      1
      0
      0
      1

      4
      1
      0
      1
      1

      5
      1
      1
      0
      1

      6
      1
      1
      1
      1

```python
# basic table
pd.crosstab(index=df1['Y'], columns=df1['X2'])
```

      X2
      0
      1

      Y

      0
      2
      3

      1
      1
      2

```python
# Contingency Table with margins
pd.crosstab(index=df1['Y'], columns=df1['X1'], margins=True)
```

      X1
      0
      1
      All

      Y

      0
      3
      2
      5

      1
      1
      2
      3

      All
      4
      4
      8

```python
# Contingency Table with margins and normalize all
pd.crosstab(index=df1['Y'], columns=df1['X1'],normalize="all", margins=True)
```

      X1
      0
      1
      All

      Y

      0
      0.375
      0.25
      0.625

      1
      0.125
      0.25
      0.375

      All
      0.500
      0.50
      1.000

```python
# Contingency Table with margins and normalize by index
pd.crosstab(index=df1['Y'], columns=df1['X1'],normalize="index", margins=True)
```

      X1
      0
      1

      Y

      0
      0.600000
      0.400000

      1
      0.333333
      0.666667

      All
      0.500000
      0.500000

```python
# Contingency Table with margins and normalize by column
pd.crosstab(index=df1['Y'], columns=df1['X1'],normalize="columns", margins=True)
```

      X1
      0
      1
      All

      Y

      0
      0.75
      0.5
      0.625

      1
      0.25
      0.5
      0.375

```python
# Contingency Table with margins and all colums in dataframe
pd.crosstab(index=df1['Y'], columns=[df1['X1'], df1['X2']], margins=True)
```

      X1
      0
      1
      All

      X2
      0
      1
      0
      1

      Y

      0
      1
      2
      1
      1
      5

      1
      0
      1
      1
      1
      3

      All
      1
      3
      2
      2
      8

```python
from sympy import *
from fractions import Fraction

print(Fraction(5, 6))
```

    5/6

```python
test = pd.crosstab(index=df1['Y'], columns=df1['X1'], margins=True)
test
```

      X1
      0
      1
      All

      Y

      0
      3
      2
      5

      1
      1
      2
      3

      All
      4
      4
      8

```python
test.applymap(lambda x: Fraction(x, 5))
```

      X1
      0
      1
      All

      Y

      0
      3/5
      2/5
      1

      1
      1/5
      2/5
      3/5

      All
      4/5
      4/5
      8/5

```python

```

```python
def crosstab_to_fraction(test:pd.DataFrame) -> pd.DataFrame:
    # convert crosstab result to numpy
    data_top = list(test.columns)
    input_data = test.to_numpy()
    # select subarray
    lines, rows = input_data.shape
    subarray = input_data[:(lines-1),:rows-1]
    # define fraction denominator
    denom = input_data[:,-1:]
    denom = denom[:-1]
    # save all digit as a fraction
    input_data = input_data.astype('object')

    for i in range(0, len(subarray)):
        for j in range(0, len(subarray[i])):
            input_data[i,j] = Fraction(subarray[i,j], denom[i][0])

    # convert numpy to dataframe
    df = pd.DataFrame(input_data, columns=data_top)
    return df

res = crosstab_to_fraction(test)
res
```

      0
      1
      All

      0
      3/5
      2/5
      5

      1
      1/3
      2/3
      3

      2
      4
      4
      8

```python
from itertools import product

x = np.array([[2, 0],
              [1, 1],
              [0, 2],
              [2, 0],
              [1, 2],
              [0, 0],
              [2, 0]])

def unique_counter(test:pd.DataFrame) -> pd.DataFrame:
    # unique line counter for dataframe
    dz = test.groupby(list(test.columns)).size().reset_index(name='Count')
    return dz

def expand_array(test:np.array) -> np.array:
    # get unique element
    rows, cols = test.shape
    test = np.reshape(test, (cols, rows))
    # calculate product of colums in array
    lines = product(test[0], test[1])
    exploist = np.array(list(set(lines)))
    return exploist

def set_diff2d(A, B):
    """
    from here https://cloudstack.ninja/msunij/hot-to-get-the-set-difference-of-two-2d-numpy-arrays-or-equivalent-of-np-setdiff1d-in-a-2d-array/
    """
    nrows, ncols = A.shape
    dtype={'names':['f{}'.format(i) for i in range(ncols)],
        'formats':ncols * [A.dtype]}
    C = np.setdiff1d(A.view(dtype), B.view(dtype))

    return C.view(A.dtype).reshape(-1, ncols)

print("Original data")
print(x)
print("==================================")
print("Data after expansion ")
data = expand_array(x)
print(data)
print("==================================")
print("Difference between original and expansion")
news = set_diff2d(data, x)
print(news)
print("==================================")

# calculate unique counter of ofiginal data
df1 = pd.DataFrame(x, columns = ['y', 'yHat'])
newdf = unique_counter(df1)

# calculate data to append
adddf = pd.DataFrame(news, columns = ['y', 'yHat'])
adddf['Count'] = 0

frames = [newdf, adddf]
result = pd.concat(frames)
result.reset_index(drop=True, inplace=True)

print("Original data with extension and counting")
result
```

    Original data
    [[2 0]
     [1 1]
     [0 2]
     [2 0]
     [1 2]
     [0 0]
     [2 0]]
    ==================================
    Data after expansion
    [[0 1]
     [1 2]
     [0 0]
     [2 1]
     [2 0]
     [1 1]
     [2 2]
     [1 0]
     [0 2]]
    ==================================
    Difference between original and expansion
    [[0 1]
     [1 0]
     [2 1]
     [2 2]]
    ==================================
    Original data with extension and counting

      y
      yHat
      Count

      0
      0
      0
      1

      1
      0
      2
      1

      2
      1
      1
      1

      3
      1
      2
      1

      4
      2
      0
      3

      5
      0
      1
      0

      6
      1
      0
      0

      7
      2
      1
      0

      8
      2
      2
      0

```python
x = np.array([[2, 0],
              [1, 1],
              [0, 2],
              [2, 0],
              [1, 2],
              [0, 0],
              [2, 0]])
x.sort(axis=0)
x
```

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

```python
x = np.array([[2, 0],
              [1, 1],
              [0, 2],
              [2, 0],
              [1, 2],
              [0, 0],
              [2, 0]])
fliprX = np.fliplr(x)
fliprX
```

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

```python
fliprX.T
```

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

```python
np.lexsort(fliprX.T)
```

    array([5, 2, 1, 4, 0, 3, 6])

```python

print(x[np.lexsort(fliprX.T)], "\nx =\n", x)
# x[np.lexsort(np.fliplr(x).T)]
# x
```

    [[0 0]
     [0 2]
     [1 1]
     [1 2]
     [2 0]
     [2 0]
     [2 0]]
    x =
     [[2 0]
     [1 1]
     [0 2]
     [2 0]
     [1 2]
     [0 0]
     [2 0]]

```python

```

```python

```

Chapter 2/python/classification metrics/Contingency Table more advanced.ipynb


# Normalizing Flow Example: Real NVP

[](https://colab.research.google.com/github/EmoryMLIP/DeepGenerativeModelingIntro/blob/main/examples/RealNVP.ipynb)

## Some References

- Original work by [Dinh, Laurent and Sohl-Dickstein, Jascha and Bengio, Samy (2016)](https://arxiv.org/abs/1605.08803)
- Similar tutorial by [Arsenii Senya Ashukha](https://github.com/senya-ashukha/real-nvp-pytorch)
- Section 3.1 of our paper  [Introduction to Deep Generative Modeling](https://arxiv.org/abs/2103.05180)

## Short Mathematical Description

We train the generator by minimizing the negative log-likelihood of the samples:

\begin{equation}
    \newcommand{\bftheta}{\boldsymbol{\theta}}
    \newcommand{\bfx}{\mathbf{x}}
    \newcommand{\bfz}{\mathbf{z}}
    \newcommand{\bfy}{\mathbf{y}}
\min_{\bftheta}  \frac{1}{s} \sum_{i=1}^s \left( {\frac{1}{2}}\left\| g_{\bftheta}^{-1}\left(\bfx^{(i)}\right)\right\|^2 - \log\det\nabla  g_{\bftheta}^{-1}\left(\bfx^{(i)}\right) \right),
\end{equation}
where $\bfx^{(1)}, \ldots$ are sampled i.i.d. from the data distribution. Here, the idea is to use a generator like
\begin{equation*}
    g_{\bftheta}(\bfz) = f_K \circ f_{K-1} \circ \cdots \circ f_1(\bfz)
\end{equation*}
to obtain a tractable inverse and log-determinant.

```python
# install requirements  (Colab only)
import sys,os
if 'google.colab' in sys.modules:
    dgm_dir = '/content/DeepGenerativeModelingIntro/'
    if not os.path.exists(dgm_dir):
        !git clone git://github.com/EmoryMLIP/DeepGenerativeModelingIntro.git

    sys.path.append(os.path.dirname(dgm_dir))
```

```python
import torch
from torch import nn
import numpy as np
from torch import distributions
from sklearn import datasets
import matplotlib.pyplot as plt
%matplotlib inline
plt.rcParams.update({'font.size': 16})
plt.rcParams.update({'image.interpolation' : None})
plt.rcParams['figure.figsize'] = [10, 7]
plt.rcParams['figure.dpi'] = 200
```

## Some suggested experiments

1. reduce the widths of the moon shapes so that the intrinsic dimensionality reduces to one
1. try the other dataset or create your own more challenging distribution
1. vary the hyperparameters (learning rate, number of layers, width, ...) and see how the performance changes
1. (for later as this takes more time) use a different layer with tractable inverse and log-determinant

## Setup the Dataset and Plotting

```python
dataset = 'moons'
# dataset = 'uniform'

if dataset == 'moons':
    # grid of the spatial domain (for plotting)
    domain = (-1.2, 2.1, -1.2, 2.1)

    def get_data(batch_size,noise=0.05):
        return torch.tensor(datasets.make_moons(n_samples=batch_size, noise=noise)[0], dtype=torch.float32)
elif dataset == 'uniform':
    domain = (-1.0, 3.0, -1.0, 3.0)
    def get_data(batch_size):
        return torch.rand((batch_size,2)) + 1.0
else:
    print("incorrect dataset")

# grid of the spatial domain (for plotting)
x1 = torch.linspace(domain[0],domain[1], 100)
x2 = torch.linspace(domain[2],domain[3], 100)
xg = torch.meshgrid(x1, x2)
xx = torch.cat((xg[0].reshape(-1, 1), xg[1].reshape(-1, 1)), 1)

from plotting import *

xs = get_data(1000)
plot_x(xs,domain)
plt.title("dataset: %s" % (dataset))
```

## Choose the Network Architecture

We write the generator as

\begin{equation*}
    \newcommand{\bftheta}{\boldsymbol{\theta}}
	\newcommand{\bfz}{\mathbf{z}}
	\newcommand{\bfy}{\mathbf{y}}
	\newcommand{\bfm}{\mathbf{m}}
	g_{\bftheta}(\bfz) = f_K \circ f_{K-1} \circ \cdots \circ f_1(\bfz)
\end{equation*}

Here we are using the Real NVP layer

\begin{equation*}
f_j\left(\bfy^{(j)}\right) = \bfm^{(j)} \odot \bfy^{(j)} + (1-\bfm^{(j)}) \odot \left(y^{(j)} \odot \exp\left(s( \bfy^{(j)} \odot \bfm^{(j)})\right) + t\left(\bfy^{(j)} \odot \bfm^{(j)}\right)\right),
\end{equation*}

where $\bfm^{j} \in \{0,1\}^n$ is used to mask some components of the inputs. Below, we alternate between the first and second coordinate.

```python
width = 128   # width of neural net layers
K = 6         # number of layers

from realNVP import NF, RealNVPLayer

layers = torch.nn.ModuleList()
for k in range(K):
    m = torch.tensor([1 - (k % 2), k % 2])
    t = nn.Sequential(nn.Linear(2, width), nn.LeakyReLU(), nn.Linear(width, width), nn.LeakyReLU(), nn.Linear(width, 2),
                      nn.Tanh())
    s = nn.Sequential(nn.Linear(2, width), nn.LeakyReLU(), nn.Linear(width, width), nn.LeakyReLU(), nn.Linear(width, 2),
                      nn.Tanh())
    layer = RealNVPLayer(s, t, m)
    layers.append(layer)

prior = distributions.MultivariateNormal(torch.zeros(2), torch.eye(2))
flow = NF(layers, prior)
```

## Train the Generator

Here, we use ADAM, a stochastic approximation scheme that operates on minibatches.

```python
batch_size = 256
num_steps = 10000      # number of training steps
out_file=None          # base filename saving trained model
plot_interval = 1000    # interval for visualizing intermediate results

optimizer = torch.optim.Adam(flow.parameters(), lr=1e-4)
his = np.zeros((0,1))

print((3*"--" + " K=%d, width=%d, batch_size=%d, num_steps=%d" + 3*"--") % (K, width, batch_size, num_steps, ))

if out_file is not None:
    import os
    out_dir, fname = os.path.split(out_file)
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    print((3*"--" + "out_file: %s" + 3*"--") % (out_file))

print((2*"%7s    ") % ("step","J_ML"))

train_JML = 0.0
num_step = 0

for step in range(num_steps):

    x = get_data(batch_size)
    loss = -flow.log_prob(x).mean()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    train_JML += loss.item()
    num_step += 1

    if (step + 1) % plot_interval == 0:
        train_JML /= num_step

        print(("%06d   " + "%1.4e  ") %
              (step + 1, train_JML))
        his = np.vstack([his, [train_JML]])
        train_JML = 0.0
        num_step = 0

        zs = flow.ginv(x)[0].detach()
        xs = flow.sample(200).detach()
        log_px = flow.log_prob(xx).detach()

        plt.Figure()
        plt.subplot(1,3,1)
        plot_x(xs,domain)
        plt.title("generated samples")

        plt.subplot(1,3,2)
        plot_z(zs)
        plt.title("latent space")

        plt.subplot(1,3,3)
        plot_px(log_px.reshape(len(x1), len(x2)),domain)
        plt.title("likelihood estimate")
        plt.margins(0, 0)
        if out_file is not None:
            plt.savefig("%s-step-%d.png" % (out_file,step+1), bbox_inches='tight', pad_inches=0)
        plt.show()

if out_file is not None:
    torch.save(flow.state_dict(), ("%s.pt") % (out_file))
    from scipy.io import savemat
    savemat(("%s.mat") % (out_file), {"his":his})
```

## Visualize the Final Result

Compare this to Figure 5 in [https://arxiv.org/abs/2103.05180](https://arxiv.org/abs/2103.05180)

```python
plt.Figure()
plt.subplot(1,3,1)
plot_px(log_px.reshape(len(x1), len(x2)),domain)
plot_x(xs[:64, :],domain)
plt.title("likelihood and samples")

xx = get_data(20000)
zz = flow.ginv(xx)[0].detach().numpy()

plt.subplot(1,3,2)
plot_pz(zz)
plot_z(zs[:64,:])
plt.title("inverse of generator")

zall = torch.cat(flow.ginv(x)[2],1)

plt.subplot(1,3,3)
plt.plot(x[:, 0], x[:, 1], "bs",alpha=0.05)
plt.plot(zs[:, 0], zs[:, 1], "or",alpha=0.05)

for k in range(6):
    zk = zall[k,:].reshape(-1,2)
    plt.plot(zk[:, 0], zk[:, 1], "-xk")
    plt.plot(zk[0, 0], zk[0, 1], "bs")
    plt.plot(zk[-1, 0], zk[-1, 1], "or")

plt.axis("square")
plt.axis((-3.5, 3.5, -3.5, 3.5))
plt.xticks((-3.5, 3.5))
plt.yticks((-3.5, 3.5))
plt.margins(0, 0)
plt.title("hidden layers")
```

```python

```


# Numerical Assignment 1 Dodos

Name and Student ID: John Butler 999999999

Adapt the python code to model the initial value problem, which describes the Dodo population (p) in the 17th century the year (t=0) that the Dodo are exposed to man and subsequent interactions modeled by the over 10 years using a stepsize of h=0.1.
To learn more about the dodo look here [2].

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

## Dodo Population Equation before man
The Dodo population prior to being found by man is described by the non-linear,
\begin{equation}
\frac{d p}{dt}=\alpha_{Birth} p-\alpha_{Death}p^2,
\end{equation}
where $p$ is the population at time $t$,
for $a\leq t \leq b$ years, where  $\alpha_{Birth}$ is the birth rate and $\alpha_{Death}$ is the death rate, with the initial condition
$$ p(a)=A. $$
TIn this simulation we set the parameters of the different equations:
*  as $\alpha_{Birth}=0.3$, and $\alpha_{Death}=0.01$ with the inital condition $p(0)=30$, over the time period $0\leq t \leq 10$.

```python
alpha_birth=0.3
alpha_death=0.01
```

## Discrete Time Domain
To numerically approximate the solution we discretise the continuous domain using a step size, $h=0.1$,
which gives
\begin{equation}
t_i=0+i h,
\end{equation}
for $i=0,...100$.

The Figure below illustrate the discrete domain.

```python
h=0.1
t_end=10
t=np.arange(0,t_end+h/2,h)
fig = plt.figure(1,figsize=(14,4))
ax = fig.add_subplot(1,1,1)
ax.plot(t,0*t,'o:',color='k')
ax.set_xlabel('Time (yrs)')
plt.title('Discrete Time')
plt.show()
```

```python
N=int(t_end/h) # Number of Steps
P=np.zeros(N+1) # Numerical Solution P
P[0]=30
```

We numerically approximate the differential equation using Eulers method to give the difference equation
\begin{equation}
P_{i}=P_{i-1}+h \big(\alpha_{Birth} P_{i-1}-\alpha_{Death}P_{i-1}^2 \big),
\end{equation}
where $h$ is the stepsize and where $P_i$ is the numerical approximation of $p(t_i)$ at distance $t_i$, for $i=0,...,10.$

```python
for i in range (1,N+1):
    P[i]=P[i-1]+h*(alpha_birth*P[i-1]-alpha_death*P[i-1]*P[i-1])

```

## Results
The plot below shows the numerical solution $P_i$ for the Dodo population as a function of time $t$.

```python
## Plotting Figure
fig = plt.figure(figsize=(14,4))
plt.plot(t,P,'o:',color='red',label='Population')
plt.legend()
plt.xlabel('Time (yrs)')
plt.ylabel('Population (thousands)')
#ax.legend(loc='best')
plt.title('Dodo Population, P')
plt.ylim(0,40)
plt.show()
```

## Dodo Population Equation after man
The Dodo population prior to being found by man is described by the non-linear,
\begin{equation}
\frac{d p}{dt}=\alpha_{Birth} p-\alpha_{Death}p^2-p(e^{0.1t}-1),
\end{equation}
where $p$ is the population at time $t$,
for $a\leq t \leq b$ years, where  $\alpha_{Birth}$ is the birth rate and $\alpha_{Death}$ is the death rate, with the initial condition
$$ p(a)=A. $$
The introduction of humans to the island is simulated with the inclusion of the terms
$$-p(e^{0.1t}-1).$$

Like above the parameters are set:
*  as $\alpha_{Birth}=0.3$, and $\alpha_{Death}=0.01$ with the inital condition $p(0)=30$, over the time period $0\leq t \leq 10$.

## Problem 1
Adapt the text in blue and the code below to model the impact of humans on the Dodo population.

We numerically approximate the differential equation using Eulers method to give the difference equation
\begin{equation}
P_{i}=P_{i-1}+h \big(\alpha_{Birth} P_{i-1}-\alpha_{Death}P_{i-1}^2+ \ \color{blue}{HUMANS} \big),
\end{equation}
where $h$ is the stepsize and where $P_i$ is the numerical approximation of $p(t_i)$ at distance $t_i$, for $i=0,...,10.$

HINT: use np.exp(0.1*t[i]) for $e^{0.1t_i}$.

```python
for i in range (1,N+1):
    ### UPDATED DIFFERENCE FOR DODOS TO INCLUDE HUMANS
    P[i]=P[i-1]+h*(alpha_birth*P[i-1]-alpha_death*P[i-1]*P[i-1]-)
    ###
```

```python
## Plotting Figure
fig = plt.figure(figsize=(14,4))
plt.plot(t,P,'o:',color='red',label='Population')
plt.legend()
plt.xlabel('Time (yrs)')
plt.ylabel('Population (thousands)')
#ax.legend(loc='best')
plt.title('Dodo Population, P')
plt.show()
```

## Problem 2.

Write a brief description of what the solution of the model tells you about the population of Dodos.

# References

[1] Stover, Christopher and Weisstein, Eric W. "Population Growth." From MathWorld--A Wolfram Web Resource. http://mathworld.wolfram.com/PopulationGrowth.html

[2] Wikipedia contributors. (2021, January 13). Dodo. In Wikipedia, The Free Encyclopedia. Retrieved 15:19, February 22, 2021, from https://en.wikipedia.org/w/index.php?title=Dodo&oldid=1000163941

```python

```


```python
%matplotlib inline
%config InlineBackend.figure_format = 'retina'

%load_ext autoreload
%autoreload 2

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.pylab as pylab
from simplified_monorotor import Monorotor
import plotting
import testing
import trajectories

pylab.rcParams['figure.figsize'] = 10,10
```

# PD controller

#### TODO 1 - Implement PD controller

Implement the PD Controller math in `thrust_control`:

$$
\begin{align}
e &= z_{\text{target}} - z_{\text{actual}} \\
\dot{e} &= \dot{z}_{\text{target}} - \dot{z}_{\text{actual}} \\
\bar{u}_1 &= k_p e + k_d \dot{e} \\
u_1 &= m(g - \bar{u}_1)
\end{align}
$$

```python
class PDController:

    def __init__(self, k_p, k_d, m):
        self.k_p = k_p
        self.k_d = k_d
        self.vehicle_mass = m
        self.g = 9.81

    def thrust_control(self,
                z_target,
                z_actual,
                z_dot_target,
                z_dot_actual,
                z_dot_dot_ff=0.0): # IGNORE this for now.
        #
        # TODO
        #   implement PD controller math shown above
        #   and return a commanded thrust.
        #
        #   You can ignore the z_dot_dot_ff parameter for now.

        pass

testing.pd_controller_test(PDController, feed_forward=False)
```

#### TODO 2 - Adjust parameters

Start by running the code below with K_D = 0.0. This will remind you what a P-controlled trajectory looks like. Then try slowly increasing K_D to 0.5, 1.0, 2.0, etc...

What value of K_D gives a reasonable trajectory?

Is there a problem with setting K_D too high?

```python
MASS_ERROR = 1.5
K_P = 20.0
K_D = 0.0

# preparation
drone = Monorotor()
perceived_mass = drone.m * MASS_ERROR
controller = PDController(K_P, K_D, perceived_mass)

# generate trajectory
total_time = 3.0
dt = 0.001
t=np.linspace(0.0,total_time,int(total_time/dt))
z_path= -np.ones(t.shape[0])
z_dot_path = np.zeros(t.shape[0])

# run simulation
history = []
for z_target, z_dot_target in zip(z_path, z_dot_path):
    z_actual = drone.z
    z_dot_actual = drone.z_dot
    u = controller.thrust_control(z_target, z_actual,
                                  z_dot_target, z_dot_actual)
    drone.thrust = u
    drone.advance_state(dt)
    history.append(drone.X)

# generate plots
z_actual = [h[0] for h in history]
plotting.compare_planned_to_actual(z_actual, z_path, t)
```

[Solution](/notebooks/PD%20Controller%20Solution.ipynb)

```python

```