```python
from sympy import *
from galgebra.printer import Format, latex, Fmt, GaLatexPrinter
Format()
from galgebra.ga import Ga
from galgebra.mv import ONE, ZERO, HALF
```

```python
def dot_basis_r_basis(ga):
    return [ga.dot(ga.basis[i], ga.r_basis[i]) for i in ga.n_range]
```

```python
def gg(ga):
    return simplify(ga.g * ga.g_inv)
```

```python
def conv_christoffel_symbols(cf):
    return permutedims(Array(cf), (2, 0, 1))
```

```python
def show_christoffel_symbols(ga):
    if ga.connect_flg:
        display(conv_christoffel_symbols(ga.Christoffel_symbols(mode=1)))
        display(conv_christoffel_symbols(ga.Christoffel_symbols(mode=2)))
```

## Minkowski Spacetime Metric

```python
coord = symbols('t x y z')
metric = Matrix([
                [ 1, 0, 0, 0 ],
                [ 0, -1, 0, 0 ],
                [ 0, 0, -1, 0 ],
                [ 0, 0, 0, -1 ]
            ])
minkowski = Ga('e', g=metric, coords=coord, norm=False)
```

```python
# NBVAL_IGNORE_OUTPUT
minkowski.basis
```

$\displaystyle \left[ \mathbf{e_{t}}, \  \mathbf{e_{x}}, \  \mathbf{e_{y}}, \  \mathbf{e_{z}}\right]$

```python
minkowski.g
```

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

```python
dot_basis_r_basis(minkowski)
```

$\displaystyle \left[ 1, \  1, \  1, \  1\right]$

```python
gg(minkowski)
```

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

## Kerr-Debney Metric

```python
g4coords = (u, x, y, z) = symbols("u x y z")
g = Matrix([
    [0, 0, -exp(-z), 0],
    [0, HALF * u ** 2 * exp(4 * z), 0, 0],
    [-exp(-z), 0, 12 * exp(-2 * z), u * exp(-z)],
    [0, 0, u * exp(-z), HALF * u ** 2],
])
g4 = Ga('e', g=g, coords=g4coords, norm=False)
```

```python
g4.basis
```

$\displaystyle \left[ \mathbf{e_{u}}, \  \mathbf{e_{x}}, \  \mathbf{e_{y}}, \  \mathbf{e_{z}}\right]$

```python
g4.g
```

$\displaystyle \left[\begin{matrix}0 & 0 & - e^{- z} & 0\\0 & \frac{u^{2} e^{4 z}}{2} & 0 & 0\\- e^{- z} & 0 & 12 e^{- 2 z} & u e^{- z}\\0 & 0 & u e^{- z} & \frac{u^{2}}{2}\end{matrix}\right]$

```python
g4.e_sq
```

$\displaystyle - \frac{u^{4} e^{2 z}}{4}$

```python
dot_basis_r_basis(g4)
```

$\displaystyle \left[ - \frac{u^{4} e^{2 z}}{4}, \  - \frac{u^{4} e^{2 z}}{4}, \  - \frac{u^{4} e^{2 z}}{4}, \  - \frac{u^{4} e^{2 z}}{4}\right]$

```python
gg(g4)
```

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

```python
show_christoffel_symbols(g4)
```

$\displaystyle \left[\begin{matrix}\left[\begin{matrix}0 & 0 & 0 & 0\\0 & - \frac{u e^{4 z}}{2} & 0 & 0\\0 & 0 & 0 & 0\\0 & 0 & 0 & - \frac{u}{2}\end{matrix}\right] & \left[\begin{matrix}0 & \frac{u e^{4 z}}{2} & 0 & 0\\\frac{u e^{4 z}}{2} & 0 & 0 & u^{2} e^{4 z}\\0 & 0 & 0 & 0\\0 & u^{2} e^{4 z} & 0 & 0\end{matrix}\right] & \left[\begin{matrix}0 & 0 & 0 & e^{- z}\\0 & 0 & 0 & 0\\0 & 0 & 0 & - 12 e^{- 2 z}\\e^{- z} & 0 & - 12 e^{- 2 z} & - u e^{- z}\end{matrix}\right] & \left[\begin{matrix}0 & 0 & 0 & \frac{u}{2}\\0 & - u^{2} e^{4 z} & 0 & 0\\0 & 0 & 12 e^{- 2 z} & 0\\\frac{u}{2} & 0 & 0 & 0\end{matrix}\right]\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}\left[\begin{matrix}0 & 0 & 0 & 0\\0 & 3 u e^{4 z} & 0 & 0\\0 & 0 & \frac{24 e^{- 2 z}}{u} & 12 e^{- z}\\0 & 0 & 12 e^{- z} & 6 u\end{matrix}\right] & \left[\begin{matrix}0 & \frac{1}{u} & 0 & 0\\\frac{1}{u} & 0 & 0 & 2\\0 & 0 & 0 & 0\\0 & 2 & 0 & 0\end{matrix}\right] & \left[\begin{matrix}0 & 0 & 0 & 0\\0 & \frac{u e^{5 z}}{2} & 0 & 0\\0 & 0 & 0 & 0\\0 & 0 & 0 & \frac{u e^{z}}{2}\end{matrix}\right] & \left[\begin{matrix}0 & 0 & 0 & \frac{1}{u}\\0 & - 3 e^{4 z} & 0 & 0\\0 & 0 & \frac{24 e^{- 2 z}}{u^{2}} & 0\\\frac{1}{u} & 0 & 0 & -1\end{matrix}\right]\end{matrix}\right]$

## Schwarzschild Metric

```python
G, M, c = symbols('G M c')
coords = (x0, x1, x2, x3) = symbols("t r theta phi")
g = Matrix([
        [ (1-(2*G*M)/(x1*c**2)), 0, 0, 0 ],
        [ 0, - (1-(2*G*M)/(x1*c**2))**(-1), 0, 0 ],
        [ 0, 0, - x1**2, 0 ],
        [ 0, 0, 0, - x1**2*sin(x2)**2 ]
    ])
schwarzschild = Ga('e', g=g, coords=coords, norm=False)
```

```python
# NBVAL_IGNORE_OUTPUT
schwarzschild.basis
```

$\displaystyle \left[ \mathbf{e_{t}}, \  \mathbf{e_{r}}, \  \mathbf{e_{\theta}}, \  \mathbf{e_{\phi}}\right]$

```python
schwarzschild.g
```

$\displaystyle \left[\begin{matrix}- \frac{2 G M}{c^{2} r} + 1 & 0 & 0 & 0\\0 & - \frac{1}{- \frac{2 G M}{c^{2} r} + 1} & 0 & 0\\0 & 0 & - r^{2} & 0\\0 & 0 & 0 & - r^{2} \sin^{2}{\left(\theta \right)}\end{matrix}\right]$

```python
schwarzschild.e_sq
```

$\displaystyle - r^{4} \sin^{2}{\left(\theta \right)}$

```python
dot_basis_r_basis(schwarzschild)
```

$\displaystyle \left[ 1, \  - \frac{\frac{2 G M}{c^{2} r} - 1}{- \frac{2 G M}{c^{2} r} + 1}, \  1, \  1\right]$

```python
gg(schwarzschild)
```

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

```python
show_christoffel_symbols(schwarzschild)
```

$\displaystyle \left[\begin{matrix}\left[\begin{matrix}0 & \frac{G M}{c^{2} r^{2}} & 0 & 0\\\frac{G M}{c^{2} r^{2}} & 0 & 0 & 0\\0 & 0 & 0 & 0\\0 & 0 & 0 & 0\end{matrix}\right] & \left[\begin{matrix}- \frac{G M}{c^{2} r^{2}} & 0 & 0 & 0\\0 & \frac{G M c^{2}}{4 G^{2} M^{2} + c^{2} r \left(- 4 G M + c^{2} r\right)} & 0 & 0\\0 & 0 & r & 0\\0 & 0 & 0 & r \sin^{2}{\left(\theta \right)}\end{matrix}\right] & \left[\begin{matrix}0 & 0 & 0 & 0\\0 & 0 & - r & 0\\0 & - r & 0 & 0\\0 & 0 & 0 & \frac{r^{2} \sin{\left(2 \theta \right)}}{2}\end{matrix}\right] & \left[\begin{matrix}0 & 0 & 0 & 0\\0 & 0 & 0 & - r \sin^{2}{\left(\theta \right)}\\0 & 0 & 0 & - \frac{r^{2} \sin{\left(2 \theta \right)}}{2}\\0 & - r \sin^{2}{\left(\theta \right)} & - \frac{r^{2} \sin{\left(2 \theta \right)}}{2} & 0\end{matrix}\right]\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}\left[\begin{matrix}0 & \frac{G M}{r \left(- 2 G M + c^{2} r\right)} & 0 & 0\\\frac{G M}{r \left(- 2 G M + c^{2} r\right)} & 0 & 0 & 0\\0 & 0 & 0 & 0\\0 & 0 & 0 & 0\end{matrix}\right] & \left[\begin{matrix}\frac{G M \left(- 2 G M + c^{2} r\right)}{c^{4} r^{3}} & 0 & 0 & 0\\0 & \frac{G M}{r \left(2 G M - c^{2} r\right)} & 0 & 0\\0 & 0 & \frac{2 G M}{c^{2}} - r & 0\\0 & 0 & 0 & \frac{\left(2 G M - c^{2} r\right) \sin^{2}{\left(\theta \right)}}{c^{2}}\end{matrix}\right] & \left[\begin{matrix}0 & 0 & 0 & 0\\0 & 0 & \frac{1}{r} & 0\\0 & \frac{1}{r} & 0 & 0\\0 & 0 & 0 & - \frac{\sin{\left(2 \theta \right)}}{2}\end{matrix}\right] & \left[\begin{matrix}0 & 0 & 0 & 0\\0 & 0 & 0 & \frac{1}{r}\\0 & 0 & 0 & \frac{1}{\tan{\left(\theta \right)}}\\0 & \frac{1}{r} & \frac{1}{\tan{\left(\theta \right)}} & 0\end{matrix}\right]\end{matrix}\right]$

```python

```


# Erlang_1,k
This notebook computes the CDF of an Erlang_1,k distribution

```
from sympy import *
from mpmath import *
%matplotlib inline
init_printing()
p, mu, t, k  = symbols('p, mu, t, k')

```

```
pdf = p * mu * exp(-mu*t)  + (1-p)* (mu**k) * ((t**(k-1))/fac(k-1)) *exp(-mu*t)
```

```

```


This is a demo of [`homogint`](https://github.com/olivierverdier/homogint), a simple Python library for integration on homogeneous spaces. The theoretical background is explained in the paper [Integrators on homogeneous spaces][paper], by [Oivier Verdier](http://www.olivierverdier.com/) and Hans Munthe-Kaas.
|
[paper]: http://arxiv.org/abs/1402.6981

## General imports

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

Plotting routine:

```python
from mpl_toolkits.mplot3d import Axes3D
def plot_sphere(ax=None):
    if ax is None:
        ax = plt.gcf().add_subplot(111, projection='3d')
    ax.autoscale(tight=True)
    ax.set_axis_off()
    ax.set_aspect("equal")

    u = np.linspace(0, 2 * np.pi, 100)
    v = np.linspace(0, np.pi, 100).reshape(-1,1)

    x = np.cos(u) * np.sin(v)
    y = np.sin(u) * np.sin(v)
    z = np.cos(v)
    ax.plot_wireframe(x, y, z,  rstride=4, cstride=4, color='k', alpha=.1)

    return ax
```

## Set up the solver

We will use the fourth order Rungke–Kutta–Munthe-Kaas method.

It can be described by the following *transition functions* (see [§3 in the paper][paper]):
\\begin{align}
    \theta\_{1,0} &= \frac{1}{2} F\_{0} \\\\
	\theta\_{2,0} &= \frac{1}{2}F\_1 - \frac{1}{8}[F\_{0},F\_1]\\\\
    \theta\_{3,0} &= F\_2 \\\\
    \theta\_{4,0} &= \frac{1}{6} (F\_0 + 2(F\_1+F\_2) + F\_3) - \frac{1}{12} [F\_0, F\_3]
\\end{align}

[paper]: http://arxiv.org/abs/1402.6981

```python
rkmk4 = RungeKutta(RKMK4())
```

We define a simple solver function that we use to solve the examples.

```python
def solve(vf,xs,stopping,action=None, maxit=10000):
    "Simple solver with stopping condition. The list xs is modified **in place**."
    for i in range(maxit):
        if stopping(i,xs[-1]):
            break
        xs.append(rkmk4.step(vf, xs[-1], action=action))
```

## Sphere: quadrature

*Example from [DiffMan](https://github.com/kenthe/DiffMan/blob/master/auxiliary/examples/vfex5.m)*.

We study the solution of the equation $x'(t) = ξ(t)x(t)$, where $x$ is on the sphere, and
\\[
ξ(t) = \begin{bmatrix}0 & t & -0.4\cos(t) \\ -t & 0 & 0.1t \\ 0.4 \cos(t) & -0.1 t & 0\end{bmatrix} \in \mathsf{so}(3)
\\]

This is equivalent to consider the problem
\\[
x'(t) = ω(t) \times x(t)
\\]
with $ω(t) = -(0.1t,0.4\cos(t),t)$.

We use a bit of a trickery here, and use instead the autonomous vector (in block notation):
\\[
\zeta(x) = \begin{bmatrix} \xi(t) & 0 &0 \\ 0 & 0 & 1 \\ 0 & 0 & 0 \end{bmatrix}
\\]
This amounts to work with the group $\mathsf{SO(3) \times \mathbf{R}}$ instead.

```python
def timedep_field(x):
    """
    Example from Diffman manual.
    """
    J = np.zeros([5,5])
    t = x[-2]
    J[0,1] = t
    J[0,2] = -.4*np.cos(t)
    J[1,2] = .1*t
    J -= J.T
    J[-2,-1] = 1.
    return J
```

```python
xs = [np.array([0.,0,1,0,1])]
dt = .02
solve(time_step(dt)(timedep_field),xs,lambda i,x:x[-2]>10)
axs = np.array(xs)
```

```python
def plot2(axs):
    plt.plot(axs[0,0],axs[0,1],'o')
    plt.plot(axs[:,0], axs[:,1],marker='.')
```

```python
plt.axis('equal')
plot2(axs)
```

```python
fig = plt.figure(figsize=(15,10))
ax = plot_sphere()
tot = len(xs)
for i,s in enumerate([slice(0,tot//2,None), slice(tot//2,None,None)]):
    ax.plot(axs[s,0],axs[s,1],axs[s,2],lw=2,marker='.',color=['black','blue'][i], alpha=[1.,0.2][i])
ax.view_init(50,-130)
# savefig('quad.pdf')
```

## $\mathsf{SO}(3)$: Quadrature

*Example from [DiffMan](https://github.com/kenthe/DiffMan/blob/master/auxiliary/examples/vfex1.m)*

The field is
\\[
\xi(t) = \begin{bmatrix} 0 & t & 1\\ -t & 0 & -t^2 \\ -1 & t^2 & 0\end{bmatrix}
\\]

```python
def so31_field(x):
    t = x[3,3]
    xi = np.zeros([5,5])
    xi[0,1] = t
    xi[0,2] = 1.
    xi[1,2] = -t*t
    xi -= xi.T
    xi[-2,-1] = 1.
    return xi
```

```python
x0 = np.zeros([5,4])
x0[:3,:3] = np.identity(3)
x0[-1,-1] = 1.
xs = [x0]
dt = .01
solve(time_step(dt)(so31_field), xs, lambda i,x: x[-2,-1] > 5)
```

```python
axs = np.array(xs)
```

Plot the three unit vectors of the rotation matrix:

```python
fig = plt.figure(figsize=(15,10))
ax = plot_sphere()
for i in range(3):
    ax.plot(axs[:,0,i],axs[:,1,i],axs[:,2,i],lw=2,marker='.')
ax.view_init(45,80)
plt.savefig('so3quad.svg', bbox_inches='tight', pad_inches=-1.5)
```

## Flat space: Lorenz equation

*Example from [DiffMan](https://github.com/kenthe/DiffMan/blob/master/auxiliary/examples/vfex2.m)*

This is the equation
\\[
(x,y,z)' = (-βx + yz, -σy + σz, -xy + ρy - z)
\\]
with the values
\\[
σ = 10 \qquad ρ = 28 \qquad β = 8/3
\\]

The idea here is to use the translation group, so the infinitesimal vector field is
\\[
\xi(x) = \begin{bmatrix} 0 & v(x)\\ 0 & 0\end{bmatrix}
\\]
where $v(x)\in\mathbf{R}^3$ is the Lorenz vector field above.

```python
def lorenz(x, sigma=10, beta=8./3., rho=28):
    y = x[1]
    A = np.array([[-beta, 0, y],
         [0, -sigma, sigma],
        [-y, rho, -1]])
    vf = np.dot(A,x[:3])
    xi = np.zeros([4,4])
    xi[:-1,-1] = vf # translation only
    return xi
```

```python
xs = [np.array([25.,0,-20,1])]
solve(time_step(0.02)(lorenz), xs, lambda i,x: i > 20/.02)
```

```python
axs = np.array(xs)
```

```python
fig = plt.figure(figsize=(10,8))
ax = fig.add_subplot(111, projection='3d')
ax.plot(axs[:,0],axs[:,1],axs[:,2],marker='.')
ax.plot([axs[0,0]],[axs[0,1]],[axs[0,2]],marker='o')
#ax.view_init(90,120)
```

## Isospectral Manifold: Toda flow

An isospectral flow is an equation of the form
\\[
P' = ξ(P)P - Pξ(P)
\\]
where $P$ is symmetric and $ξ(P)$ is skew-symmetric.

We implement what is known as the *Toda flow*:

```python
def iso_field(P):
    sk = np.tril(P) - np.triu(P) # skew symmetric
    return sk
```

Random symmetric matrix as initial condition.

```python
#init = np.array([[-1.,1,0],[1,.5,1],[0,1,.5]])
rmat = np.random.randn(20,20)
init = rmat + rmat.T
plt.matshow(init)
plt.savefig('matinit.png', bbox_inches='tight', pad_inches=0)
```

```python
from homogint.homogint import trans_adjoint
Ps = [init]
dt = .25
solve(time_step(dt)(iso_field), Ps, lambda i,x: i>30/dt, action=trans_adjoint)
```

The flow does not change the eigenvalues (hence the name *isospectral flow*)

```python
import numpy.linalg as nl
eigenvalues = [np.sort(nl.eigvals(P)) for P in Ps]
aeigenvalues = np.array(eigenvalues)
deig = aeigenvalues - aeigenvalues[0]
plt.plot(deig)
plt.title("eigenvalue drift")
```

```python
from IPython.html.widgets import interact
def view_matrix(i):
    plt.matshow(Ps[i])
interact(view_matrix, i=(0,len(Ps)-1,1))
```

The flow is a continuous version of the [QR algorithm](http://www.wikiwand.com/en/QR_algorithm), so it almost converges towards a diagonal matrix (*almost* because there is no shift and deflations).

```python
plt.matshow(Ps[-1])
plt.savefig('matfinal.png', bbox_inches='tight', pad_inches=0)
```

## Airy Equation

*Example from [Lie Group Method](http://www.damtp.cam.ac.uk/user/na/NA_papers/NA2000_03.pdf)*

We solve the equation
\\[
x'' + tx = 0
\\]

It is reformulated as
\\[
(x,v)' = \begin{bmatrix} 0 & 1\\ -t & 0\end{bmatrix} (x,v)
\\]

```python
def airy_field(x):
    mat = np.zeros([4,4])
    mat[0,1] = 1.
    mat[1,0] = -x[-2]
    mat[-2,-1] = 1 # time
    return mat
```

We solve it with initial condition $x(0) = 1.$, $x'(0) = 0$.

```python
x0 = np.array([1.,0,0,1])
xs = [x0]
dt=.05
solve(time_step(dt)(airy_field), xs, stopping=lambda i,x: i>100/dt)
```

```python
len(xs)
```

    2002

```python
fig = plt.figure(figsize=(15,5))
axs = np.array(xs)
plt.plot(axs[:,-2],axs[:,0])
#plot(axs[1900:,-2],axs[1900:,0],marker='.')
```

We can compute the exact solution using the `airy` function in `scipy.special`.

```python
from scipy.special import airy
from scipy.linalg import solve as linsolve
```

```python
# taken from the DiffMan examples
tstart = 0
m =np.array([airy(tstart)[0], airy(tstart)[2], -airy(tstart)[1], -airy(tstart)[3]]).reshape(2,-1)
c = linsolve(m, np.array([1.,0]))

# Computes the 'true' solution:
ts = np.linspace(90,100,1000)
def exact_airy(ts, c):
    return c[0]*np.real(airy(-ts)[0]) + c[1]*np.real(airy(-ts)[2])
```

```python
plt.figure(figsize=(15,5))
skip=1800
plt.plot(ts,exact_airy(ts,c),label="exact")
plt.plot(axs[skip:,-2],axs[skip:,0],marker='o',linestyle='',label="computed")
plt.legend()
```

```python
plt.figure(figsize=(15,5))
error = axs[:,0] - exact_airy(axs[:,-2],c)
plt.plot(axs[1:,-2],np.log10(np.abs(error[1:])))
plt.title("$\log_{10}(error)$")
```

## Stiefel manifold: Oja Flow

*Example from [Geometric Numerical Integration, § IV.9.2](http://link.springer.com/book/10.1007%2F3-540-30666-8)*

The *Oja flow* is given by
\\[
Q' = (I - QQ^T)A Q
\\]
for a given positive definite matrix $A$.

Using the connection formula
\\[
\langle ω,δQ \rangle_Q = δQ Q^T - QδQ^T -QδQ^T Q Q^T
\\]
we obtain the following vector field on the Lie algebra:
\\[
ξ(Q) = AQQ^T-QQ^TA
\\]

```python
D = np.diag([16.,8.,4.])
A = D
def oja_field(x):
    proj = np.dot(x,x.T)
    xi = commutator(A,proj)
    return xi
```

We choose a random starting point. It amounts to choose two orthogonal vectors of length one.

```python
def normalize(x):
    nx = np.sqrt(np.sum(np.square(x)))
    return x/nx

def rand_sphere_point():
    u,v = np.random.rand(2)
    phi = u*2*np.pi
    theta = np.arccos(2*v-1)
    sth = np.sin(theta)
    return np.array([sth*np.cos(phi), sth*np.sin(phi), np.cos(theta)])
r1 = rand_sphere_point()
r1_ = rand_sphere_point()
r2 = normalize(np.cross(r1,r1_))
x0 = np.array([r1,r2]).T
```

Check that the chosen vectors are orthogonal:

```python
print(np.allclose(np.dot(x0.T,x0), np.identity(2)))
```

    True

Starging value:

```python
print(x0)
```

    [[-0.386402   -0.82594557]
     [-0.87228066  0.47186382]
     [ 0.29969974  0.30847763]]

```python
xs = [x0]
dt = .1
solve(time_step(dt)(oja_field), xs, lambda i,x: np.allclose(oja_field(x),0,atol=1e-7))
```

```python
len(xs)
```

    40

The flow converges towards an invariant subspace. Here it converges towards the subspace containing the two largest eigenvalues:

```python
xs[-1]
```

    array([[ -4.37460593e-01,  -8.99237582e-01],
           [ -8.99237582e-01,   4.37460593e-01],
           [  2.15687191e-08,  -1.04927383e-08]])

```python
axs = np.array(xs)
fig = plt.figure(figsize=(15,10))
ax = plot_sphere()
ths = np.linspace(0,2*np.pi,200)
plt.plot(np.cos(ths), np.sin(ths), np.zeros_like(ths))
for i in range(2):
    for j in [0,-1]:
        ax.plot([axs[j,0,i]],[axs[j,1,i]],[axs[j,2,i]],lw=2,marker=['o','D'][j])
        ax.plot([0.,axs[j,0,i]],[0,axs[j,1,i]],[0,axs[j,2,i]],color=['black','red'][j])
    ax.plot(axs[:,0,i],axs[:,1,i],axs[:,2,i],marker='.')
ax.view_init(30,0)
plt.savefig('oja.pdf')
```

Check the convergence towards the plane with largest eigenvalues.

```python
for i in range(2):
    plt.plot(np.log10(np.abs(axs[:,-1,i])),marker='.')
plt.title("log10 of the z coordinate")
```


# Periodic Signals

*This Jupyter notebook is part of a [collection of notebooks](../index.ipynb) in the bachelors module Signals and Systems, Comunications Engineering, Universität Rostock. Please direct questions and suggestions to [Sascha.Spors@uni-rostock.de](mailto:Sascha.Spors@uni-rostock.de).*

## Relation between Spectrum and Fourier Series

The Fourier transform $X(j \omega) = \mathcal{F} \{ x(t) \}$ of a periodic signal $x(t)$, [as derived before](spectrum.ipynb#Fourier-Transform), is a line spectrum. It consists of a weighted series of Dirac impulses. Periodic functions can be represented alternatively by a [Fourier series](https://en.wikipedia.org/wiki/Fourier_series). The relation between the spectrum $X(j \omega)$ of a periodic signal and its Fourier series coefficients is derived in the following.

The complex Fourier series of a periodic signal $x(t)$ is defined as

\begin{equation}
x(t) = \sum_{n = - \infty}^{\infty} X_n \, e^{j n \frac{2 \pi}{T_\text{p}} t}
\end{equation}

where $T_\text{p} > 0$ denotes the period of the signal and $X_n$ the Fourier series coefficients of $x(t)$. The Fourier series represents the signal as weighted superposition of complex exponential signals. The weights (expansion coefficients) $X_n$ are given as

\begin{equation}
X_n = \frac{1}{T_\text{p}} \int_{0}^{T_\text{p}} x(t) \, e^{- j n \frac{2 \pi}{T_\text{p}} t} \; dt
\end{equation}

Introducing the [Fourier transform $X(j \omega)$ of a periodic signal](spectrum.ipynb#Fourier-Transform) into the [inverse Fourier transform](../fourier_transform/definition.ipynb#Definition) yields

\begin{align}
x(t) &= \frac{1}{2 \pi} \int_{-\infty}^{\infty} X_0(j \omega) \cdot {\bot \!\! \bot \!\! \bot} \left( \frac{\omega T_\text{p}}{2 \pi} \right) \, e^{j \omega t} \; d \omega \\
&= \frac{1}{T_\text{p}} \sum_{\mu = -\infty}^{\infty} X_0 \left( j \, \mu \frac{2 \pi}{T_\text{p}} \right) \, e^{j \, \mu \frac{2 \pi}{T_\text{p}} t}
\end{align}

where $X_0(j \omega) = \mathcal{F} \{ x_0(t) \}$ denotes the Fourier transform of one period $x_0(t)$ of the periodic signal. Note, the [definition of the Dirac comb](spectrum.ipynb#The-Dirac-Comb) and the multiplication property of the Dirac impulse was used to derive the last equality. Comparing this result with the definition of the Fourier series reveals that both are equal for

\begin{equation}
X_n = X_0 \left( j \, n \frac{2 \pi}{T_\text{p}} \right)
\end{equation}

The Fourier series coefficients $X_n$ of a periodic signal are equal to the Fourier transform $X_0(j \omega)$ of one period of the signal at the frequencies $\omega = n \frac{2 \pi}{T_\text{p}}$.

**Example**

The Fourier series coefficients of the pulse train can be derived from the [Fourier transform the pulse train](spectrum.ipynb#Fourier-Transform-of-the-Pulse-Train) as

\begin{equation}
X_n = T \, e^{-j \omega \frac{T}{2}} \cdot \text{sinc} \left( \frac{\omega T}{2} \right) \bigg\vert_{\omega = n \frac{2 \pi}{T_\text{p}}} = T \, e^{-j n \pi \frac{T}{T_\text{p}}} \cdot \text{sinc} \left( n \pi \frac{T}{T_\text{p}} \right)
\end{equation}

With these coefficients the pulse train can be represented by the Fourier series

\begin{equation}
x(t) = \sum_{n = -\infty}^{\infty} T \, e^{-j n \pi \frac{T}{T_\text{p}}} \cdot \text{sinc} \left( n \pi \frac{T}{T_\text{p}} \right) \, e^{j n \frac{2 \pi}{T_\text{p}} t}
\end{equation}

This series cannot be evaluated numerically due to its infinite limits. The series has to be truncated to a finite number of summands in a practical implementation. The consequences of truncating the series are illustrated in the following.

```python
%matplotlib inline
import sympy as sym
sym.init_printing()

n = sym.symbols('n', integer=True)
t = sym.symbols('t', real=True)

T = 2
Tp = 5

Xn = T * sym.exp(-sym.I * n * sym.pi * T/Tp) * sym.sinc(n * sym.pi * T/Tp)
Xn
```

Now the Fourier series is evaluated for a finite upper and lower limit $N$

\begin{equation}
x_N(t) = \sum_{n = -N}^{N} X_n \, e^{j n \frac{2 \pi}{T_\text{p}} t}
\end{equation}

```python
N = 15

x = sym.Sum(Xn * sym.exp(sym.I*n*2*sym.pi/Tp*t), (n, -N, N)).doit()
sym.plot(x, (t,0,10), xlabel='$t$', ylabel='$x_N(t)$');
```

Overshoots can be observed at the discontinuities of the pulse train. The relative magnitude of these overshoots remains at constantly 9% even when increasing the limits of the truncated Fourier series expansion. This effect is known as [*Gibbs phenomenon*](https://en.wikipedia.org/wiki/Gibbs_phenomenon). Truncated Fourier series are therefore not very well suited for the approximation of signals with discontinuities.

**Exercise**

* Examine the properties of the truncated Fourier series when you increase the limit $N$ in above example. Note: The evaluation of the Fourier series may take a while due to involved numerical complexity.

**Copyright**

The notebooks are provided as [Open Educational Resource](https://de.wikipedia.org/wiki/Open_Educational_Resources). Feel free to use the notebooks for your own educational purposes. The text is licensed under [Creative Commons Attribution 4.0](https://creativecommons.org/licenses/by/4.0/), the code of the IPython examples under the [MIT license](https://opensource.org/licenses/MIT). Please attribute the work as follows: *Lecture Notes on Signals and Systems* by Sascha Spors.


# Energy groups

To capture the variation of neutron energies in the reactor, we can **discretize the flux into energy groups.**

\begin{align}
\phi &= \sum_{g=1}^{g=G}\phi_g\\
\phi_g &= \int_{E_g}^{E_{g-1}}\phi(E)dE\rvert_{g = 1,G}
\end{align}

The various parameters (such as cross sections)  also need to be individually evaluated for each energy group such that we have a $\Sigma_{ag}$ for each group g.

The various cross sections and coefficients also need to be individually evaluated for each energy group such that we have a $\Sigma_{a,g}$ for each group  $g\in[1,G] \rightarrow E_g\in[E_1,E_G]$. Also, it is important to consider possible paths of demise for potential fission neutrons.

\begin{align}
\chi_g &= \int_{E_g}^{E_{g-1}}\chi(E)dE\\
\sum_{g=1}^{g=G}\chi_g &= \int_0^\infty\chi(E) dE\\
\end{align}

And the source of fission neutrons is:

\begin{align}
S &= \sum_{g'=1}^{g'=G}(\nu\Sigma_f)_{g'}\phi_{g'}
\end{align}

### Group-wise Scattering Cross Sections

This energy group notation is used to indicate scattering from one group into another as well.
Most of the scattering is from other groups into our energy group of interest, $g$, but some of the scattering is from this group into itself.

Scattering from group $g$ to group $g'$ is denoted as $\Sigma_{s,g\rightarrow g'}$.

### Group ordering

Let's define just two groups, $E_1$ and $E_{2}$. One of these groups will represent the fast neutrons and the other will represent the thermal neutrons.

### Think-pair share:

Which group should be **$g=1$** and which should be **$g=2$**?

For simple problems, one can typically also assume that prompt neutrons are born fast.
Since they are born fast, energy group 1 is always the fastest group.

### Upscattering and Downscattering

Downscattering (with cross section $\Sigma_{s,g\rightarrow g'}$) is when energy decreases because of the scatter ( so $g < g'$ ).

Upscattering (with cross section $\Sigma_{s,g\rightarrow g'}$) is when energy increases because of the scatter ( so $g > g'$ ).

### Think Pair Share

Given what you know about reactors and elastic scattering, which is likely more common in a reactor, upscattering or downscattering?

# $\infty$ Medium, Two-Group Neutron Balance

Balance equation for group g:

\begin{align*}
        \mbox{Loss} &= \mbox{Gain}\\
        \mbox{Absorption} + \mbox{Outscattering} &= \mbox{Fission} + \mbox{Inscattering}\\
\Sigma_{ag} \phi_g + \Sigma_{g\rightarrow g'}\phi_g &= \frac{\chi_g}{k}\left(\nu \Sigma_{fg} \phi_g + \nu\Sigma_{fg'}\phi_{g'}\right) + \Sigma_{g'\rightarrow g}\phi_{g'}
\end{align*}

Group 1 (fast):

\begin{align*}
\Sigma_{a1} \phi_1 + \Sigma_{1\rightarrow 2}\phi_1 &= \frac{\chi_1}{k}\left(\nu \Sigma_{f1} \phi_1 + \nu\Sigma_{f2}\phi_2\right) + \Sigma_{2\rightarrow 1}\phi_2
\end{align*}

Group 2 (thermal):

\begin{align*}
\Sigma_{a2} \phi_2 + \Sigma_{2\rightarrow 1}\phi_2 &= \frac{\chi_2}{k}\left(\nu \Sigma_{f2} \phi_2 + \nu\Sigma_{f1}\phi_1\right) + \Sigma_{1\rightarrow 2}\phi_1
\end{align*}

We can rewrite these equations in matrix form:

\begin{align*}
  \left[ {\begin{array}{cc}
   \Sigma_{a1} & 0 \\
   0 & \Sigma_{a2} \\
  \end{array} } \right]
  \left[ {\begin{array}{c}
   \phi_1 \\
   \phi_2 \\
  \end{array} } \right]
+
  \left[ {\begin{array}{cc}
   \Sigma_{1\rightarrow 2} & 0 \\
   0 & \Sigma_{2\rightarrow 1} \\
  \end{array} } \right]
  \left[ {\begin{array}{c}
   \phi_1 \\
   \phi_2 \\
  \end{array} } \right]
&=
\frac{1}{k}
  \left[ {\begin{array}{c}
   \chi_1 \\
   \chi_2 \\
  \end{array} } \right]
  \left[ {\begin{array}{cc}
    \nu \Sigma_{f1} & \nu \Sigma_{f2}
  \end{array}}\right]
  \left[ {\begin{array}{c}
   \phi_1 \\
   \phi_2 \\
  \end{array} } \right]
  +
  \left[ {\begin{array}{cc}
   0 & \Sigma_{2\rightarrow 1} \\
   \Sigma_{1\rightarrow 2} & 0 \\
  \end{array} } \right]
  \left[ {\begin{array}{c}
   \phi_1 \\
   \phi_2 \\
  \end{array} } \right]
\end{align*}

Moving the inscattering to the left side of the equation, the equation becomes:

\begin{align*}
  \left[ {\begin{array}{cc}
   \Sigma_{a1} & 0 \\
   0 & \Sigma_{a2} \\
  \end{array} } \right]
  \left[ {\begin{array}{c}
   \phi_1 \\
   \phi_2 \\
  \end{array} } \right]
+
  \left[ {\begin{array}{cc}
   \Sigma_{1\rightarrow 2} & 0 \\
   0 & \Sigma_{2\rightarrow 1} \\
  \end{array} } \right]
  \left[ {\begin{array}{c}
   \phi_1 \\
   \phi_2 \\
  \end{array} } \right]
-
  \left[ {\begin{array}{cc}
   0 & \Sigma_{2\rightarrow 1} \\
   \Sigma_{1\rightarrow 2} & 0 \\
  \end{array} } \right]
  \left[ {\begin{array}{c}
   \phi_1 \\
   \phi_2 \\
  \end{array} } \right]
&=
\frac{1}{k}
  \left[ {\begin{array}{cc}
    \chi_1\nu \Sigma_{f1} & \chi_1\nu \Sigma_{f2}\\
    \chi_2\nu \Sigma_{f1} & \chi_2\nu \Sigma_{f2}\\
  \end{array}}\right]
  \left[ {\begin{array}{c}
   \phi_1 \\
   \phi_2 \\
  \end{array} } \right]
\end{align*}

Here, we can define the macroscopic absorption, inscattering, outscattering,
and fission cross-section matrices:

\begin{align*}
\mbox{Absorption Matrix } (\mathbf{A}): &=
  \left[ {\begin{array}{cc}
    \Sigma_{a1} & 0\\
    0 & \Sigma_{a2}\\
  \end{array}}\right]\\
\mbox{Inscattering Matrix } (\mathbf{S_{in}}): &=
  \left[ {\begin{array}{cc}
    0 & \Sigma_{2\rightarrow 1}\\
    \Sigma_{1\rightarrow 2} & 0\\
  \end{array}}\right]\\
\mbox{Outscattering Matrix } (\mathbf{S_{out}}): &=
  \left[ {\begin{array}{cc}
    \Sigma_{1\rightarrow 2} & 0\\
    0 & \Sigma_{2\rightarrow 1}\\
  \end{array}}\right]\\
\mbox{Fission Matrix } (\mathbf{F}): &=
  \left[ {\begin{array}{cc}
    \chi_1\nu \Sigma_{f1} & \chi_1\nu \Sigma_{f2}\\
    \chi_2\nu \Sigma_{f1} & \chi_2\nu \Sigma_{f2}\\
  \end{array}}\right]\\
\end{align*}

And now our equation is reduced to:

\begin{align*}
\left[A + S_{out} - S_{in}\right]
  \left[ {\begin{array}{c}
   \phi_1 \\
   \phi_2 \\
\end{array} } \right]
&=
\frac{1}{k}
  \left[F\right]
  \left[ {\begin{array}{c}
   \phi_1 \\
   \phi_2 \\
  \end{array} } \right]
\end{align*}

So, the final version is:

\begin{align*}
k \left[ {\begin{array}{c}
   \phi_1 \\
   \phi_2 \\
  \end{array} } \right]
  &=
  \left[A + S_{out} - S_{in}\right]^{-1}
\left[F\right]
  \left[ {\begin{array}{c}
   \phi_1 \\
   \phi_2 \\
\end{array} } \right]
\end{align*}

This gives us the final definition to note, the migration matrix:

\begin{align*}
        \mbox{Migration Matrix:} &= \left[A + S_{out} - S_{in}\right]\\
\end{align*}

For the two group problem, let the eigenvector be
$\left(\begin{array}{c}\phi_1\\\phi_2\end{array}\right)$,
let the eigenvalue be $k$, and
calculate the eigenvalue and eigenvector of
$\left[A + S_{out} - S_{in}\right]^{-1}[F]$.

# Power Iteration

[1] Richard von Mises and H. Pollaczek-Geiringer, Praktische Verfahren der Gleichungsauflösung, ZAMM - Zeitschrift für Angewandte Mathematik und Mechanik 9, 152-164 (1929)

> Given a diagonalizable matrix $A$, the algorithm will produce a number $\lambda$, which is the greatest (in absolute value) eigenvalue of $A$, and a nonzero vector $v$, the corresponding eigenvector of $\lambda$, such that $Av=\lambda v$. The algorithm is also known as the Von Mises iteration.[1]

Facts about Power Iteration:
- It is simple to implement
- It can sometimes be quite slow to converge (reach the answer)
- It doesn't compute a matrix decomposition, so it can be used on very large sparse matrices

We'd like to solve for $\phi$. If we assume:
- B has an eigenvalue strictly greater in magnitude than its others
- the starting vector ($\phi_0$) has a nonzero component in the direction of an eigenvector associated with the dominant eigenvalue, the a sequence $\phi_k$ converges to an eigenvector associated with the dominant eigenvalue.

Iterate with the following definitions:
\begin{align*}
        \phi_{i+1} &= \frac{B\phi_i}{\lVert B\phi_i \rVert_2}\\
\end{align*}

Note the definition of the euclidean norm (a.k.a the $L^2$ norm) for a vector
$\vec{x} = (x_1, x_2, \cdots, x_n)$ is :

\begin{align*}
        \lVert \vec{x}\rVert_2 \equiv \sqrt{x_1^2 + x_2^2 + \cdots + x_n^2}.
\end{align*}

```python
import numpy as np

def power_iteration(B, num_simulations):
    """
    Returns the normalized
    """
    # By chosing a random solution vector we decrease the chance
    # that it is orthogonal to the eigenvector
    phi_k = np.random.rand(B.shape[1])

    for _ in range(num_simulations):
        # calculate the matrix-by-vector product Bphi
        phi_k1 = np.dot(B, phi_k)

        # calculate the norm
        phi_k1_norm = np.linalg.norm(phi_k1)

        # re normalize the vector
        phi_k = phi_k1 / phi_k1_norm

    return phi_k

B = np.array([[0.5, 0.5], [0.2, 0.8]])
power_iteration(B, 100)
```

    array([ 0.70710678,  0.70710678])

We can do the same with k using the Raleigh quotient and our solution for $\phi$. That is left for the reader. To find k, iterate with the following definition:

\begin{align*}
        k_{i+1} &= \frac{\left(B\phi_{i+1}\right)^T \phi_{i+1}}{\left(\phi_{i+1}^T\phi_{i+1}\right)}\\
\end{align*}


# Homework 6

**Submit the MP4 files you make along with a PDF of this notebook**

This week we will learn about how to turn the plots we have been making into animations! This will be really helpful for any groups that want to make simulations or other animations for their final project. This will build off of the differential equations we learned last week as well as the plotting skills we have learned so far.

First let us import the neccessary libraries,

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

from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FFMpegWriter

%config InlineBackend.figure_format='retina' # makes animation display better

%matplotlib osx
# ^ UNCOMMENT THIS LINE IF USING MAC

# %matplotlib qt
# ^ UNCOMMENT THIS LINE IF USING WINDOWS
```

    In /Users/jamessunseri/anaconda3/lib/python3.7/site-packages/matplotlib/mpl-data/stylelib/_classic_test.mplstyle:
    The text.latex.preview rcparam was deprecated in Matplotlib 3.3 and will be removed two minor releases later.
    In /Users/jamessunseri/anaconda3/lib/python3.7/site-packages/matplotlib/mpl-data/stylelib/_classic_test.mplstyle:
    The mathtext.fallback_to_cm rcparam was deprecated in Matplotlib 3.3 and will be removed two minor releases later.
    In /Users/jamessunseri/anaconda3/lib/python3.7/site-packages/matplotlib/mpl-data/stylelib/_classic_test.mplstyle: Support for setting the 'mathtext.fallback_to_cm' rcParam is deprecated since 3.3 and will be removed two minor releases later; use 'mathtext.fallback : 'cm' instead.
    In /Users/jamessunseri/anaconda3/lib/python3.7/site-packages/matplotlib/mpl-data/stylelib/_classic_test.mplstyle:
    The validate_bool_maybe_none function was deprecated in Matplotlib 3.3 and will be removed two minor releases later.
    In /Users/jamessunseri/anaconda3/lib/python3.7/site-packages/matplotlib/mpl-data/stylelib/_classic_test.mplstyle:
    The savefig.jpeg_quality rcparam was deprecated in Matplotlib 3.3 and will be removed two minor releases later.
    In /Users/jamessunseri/anaconda3/lib/python3.7/site-packages/matplotlib/mpl-data/stylelib/_classic_test.mplstyle:
    The keymap.all_axes rcparam was deprecated in Matplotlib 3.3 and will be removed two minor releases later.
    In /Users/jamessunseri/anaconda3/lib/python3.7/site-packages/matplotlib/mpl-data/stylelib/_classic_test.mplstyle:
    The animation.avconv_path rcparam was deprecated in Matplotlib 3.3 and will be removed two minor releases later.
    In /Users/jamessunseri/anaconda3/lib/python3.7/site-packages/matplotlib/mpl-data/stylelib/_classic_test.mplstyle:
    The animation.avconv_args rcparam was deprecated in Matplotlib 3.3 and will be removed two minor releases later.

# Problem 1 (10 points)

**Animating a Cool Function**

In homework 5, we plotted the following function:

\begin{equation}
f = \frac{\sin(10(x^2 + y^2))}{10}
\end{equation}

```python
x, y = np.meshgrid(np.linspace(0, 2*np.pi, 100),
                   np.linspace(0, 2*np.pi, 100))

f = np.sin(10*(x**2 + y**2))/10

plt.figure(figsize=(12,12))
plt.imshow(f)          # Plot the 2D array
plt.set_cmap('hot')    # Set the color scheme ("jet" is matplotlib's default)
plt.colorbar()         # Display a legend for the color values
plt.show()
```

Now let's animate it!

For each iteration, add 1 more to x and y. For example, add 1 to x and y for the first iteration, and add 10 to both x and y for the tenth iteration.

Rember to use plt.imshow( ) so that your function is shown.

```python
## SAVE AS MP4 ##
fig, ax = plt.subplots(figsize=(5,5))

num_iterations = 20

with writer.saving(fig, "3Dfunction.mp4", dpi=200):
    for i in range(num_iterations):

        ax.clear() # first clear the figure

        ### YOUR CODE HERE ###
        f = np.sin(10*((x+i)**2 + (y+i)**2))/10
        plt.imshow(f)

        plt.draw()
        plt.pause(0.01)
        writer.grab_frame() # save the current frame to mp4
```

# Problem 2 (20 points)
**Binary Star System**

In lecture we animated a single planet orbiting a star, but now let's animate a binary star system — two stars orbiting eachother. Because these two objects are of similar mass, we need to update the movements of both stars instead of just one.

First we'll define our initial conditions. For the sake of simplicity, don't worry about units for this problem.
- Star 1: mass $=1$, $x_i=-1$, $y_i=0$, $v_{xi}=1$, $v_{yx}=1$

- Star 2: mass $=1$, $x_i=1$, $y_i=0$, $v_{xi}=-1$, $v_{yx}=1$

```python
G = 1 #gravitational constant

# Define masses (in units of solar mass)
m1 = 1.0
m2 = 1.0

# Define initial position vectors
r1 = np.array([-1, 0])
r2 = np.array([1, 0])

# Define initial velocities
v1 = np.array([1, 1])
v2 = np.array([-1, 1])
```

The stars motions are characterized by the Law of Universal Gravitation:

\begin{equation}
 F = G \frac{m_1 m_2}{r^2}
\end{equation}

Ultimately, all we need to plot the stars motion is the position coordinates of the two stars. But to correctly update the stars positions we need the stars velocity and the change in velocity (acceleration). Instead of working with forces, let us work with accelerations. We can then rewrite the equations for star 1 and star 2 like this:

- $ a_1 = \frac{m_2(r_2 - r_1)}{r_{12}^3} $

- $ a_2 = \frac{m_1(r_1 - r_2)}{r_{12}^3} $

Fill out the following cell to define a function that we will integrate over to obtain our star positions.

Refer to the lecture demo notebook for help! OrbitEquation( ) that was used in the demo is set up the same way, but now we want to change position and velocity for another object as well.

```python
def TwoBodyEquations(w, t, m1, m2): # w is an array containing positions and velocities
    r1 = w[:2]
    r2 = w[2:4]
    v1 = w[4:6]
    v2 = w[6:8]

    r12 = np.linalg.norm(r2-r1)

    dv1bydt = m2*(r2-r1)/r12**3  # derivative of velocity
    dv2bydt = m1*(r1-r2)/r12**3

    dr1bydt = v1 # derivative of position
    dr2bydt = v2

    r_derivs = np.concatenate((dr1bydt,dr2bydt)) # joining arrays together
    v_derivs = np.concatenate((dv1bydt,dv2bydt))
    derivs = np.concatenate((r_derivs,v_derivs))

    return derivs
```

Now we want to run the ordinary differential equation solver. Then set r1_sol equal to the parts of two_body_sol that correspond to the first star. Do the same for r2_sol but for the second star.

```python
# Package initial parameters into one array (just easier to work with this way)
init_params = np.array([r1,r2,v1,v2])
init_params = init_params.flatten()
time_span = np.linspace(0, 20, 5000)  # run for t=20 (5000 points)

# Run the ordinary differential equation solver
two_body_sol = scipy.integrate.odeint(TwoBodyEquations, init_params, time_span, args=(m1,m2)) # use scipy.integrate.odeint()

r1_sol=two_body_sol[:,:2]
r2_sol=two_body_sol[:,2:4]
```

Now we have all of the data we want to plot. Using FFMpegWriter, let us loop through all of our data and save each frame to our MP4 file. Set your x and y axes to range from -2 to 2.

Again, this code is very similar to what was shown in the lecture demo!

```python
# Initilize writer
metadata = dict(title='2D animation', artist='Matplotlib')
writer = FFMpegWriter(fps=50, metadata=metadata, bitrate=200000)
fig = plt.figure(dpi=200)

## SAVE AS MP4 ##
fig, ax = plt.subplots()

with writer.saving(fig, "binary_stars.mp4", dpi=200):
    for i in range(len(time_span)):

        ax.clear()

        ### YOUR CODE HERE ###
        ax.plot(r1_sol[:i,0],r1_sol[:i,1],color="#00C9C8", alpha=0.5)
        ax.plot(r2_sol[:i,0],r2_sol[:i,1],color="#9296F0", alpha=0.5)

        ax.scatter(r1_sol[i,0],r1_sol[i,1],color="#00C9C8",marker="o",s= m1*30)
        ax.scatter(r2_sol[i,0],r2_sol[i,1],color="#9296F0",marker="o",s=m2*30)

        ax.set_xlim(-3, 3)
        ax.set_ylim(-3, 3)
        ###

        plt.draw()
        plt.pause(0.01)
        writer.grab_frame()
```

**Optional**: try changing the star masses, initial postions, and/or initial velocities and show us an animation that you think looks cool!

```python
# can be whatever other initial conditions
```

# Problem 3 (20 points)
## Swinging Pendulum
Now we will animate a basic swinging pendulum that you have probably seen in your Physics classes! The steps for animating this are similar to how we animated the binary star system. The main difference is how we define our ode_func and the differential equations.
For our pendulum system, we will have a point mass $m$ connected to a string of length $l$ with slight damping:
* $m = 1 kg$ is the mass,
* $g = 9.81 m/s^2$ is gravity,
* $l = 1 m$ is the length of the massless string
* $b = 0.05$ is the damping coefficient, and
* $t = 0$ to $10 s$ is the time duration with 100 intervals.
\
\
We can derive the change in the pendulum system as:
\
\
$$\frac{d^2\theta}{dt} = - \frac{b}{m} \frac{d\theta}{dt} - \frac{g  *  sin(\theta)}{l}$$
\
To further simplify this second order differential equation, we can rewrite this as two linear equations by defining

$$\theta_{1} = \theta  ,  \theta_{2} = \frac{d\theta}{dt}$$

$$\frac{d\theta_{1}}{dt} = \theta_{2}$$

$$\frac{d\theta_{2}}{dt} = -\frac{b  *  \theta_{2}}{m} - \frac{g  *  sin(\theta_{1})}{l}$$

Don't worry if you do not understand the physics equations and math for this derivation! We will only be needing the last two equations for this problem. Now use ode_func as an argument to scipy.integrate.odeint to solve for the theta1 and theta2 values for all time steps. Use the initial conditions of
* $\theta_{1} = \frac{\pi}{10} $ as the initial angle and
* $\theta_{2} = 0$ as the initial angular velocity.

```python
def ode_func(theta, t, g, l, m, b):
    '''
    Given a list theta that is [theta1, theta2] and the conditions of the system
    Computes the differential change for theta1 and theta2 following the last two eqs.
    Returns dtheta1dt and dtheta2dt as a list in that order
    '''
    # YOUR CODE HERE
    theta1, theta2 = theta
    dtheta1 = theta2
    dtheta2 = (- b * theta2 / m - g * np.sin(theta1) / l)
    return [dtheta1, dtheta2]

# Define system variables l, m, g, b, t and initial conditions
# YOUR CODE HERE
l = 1
m = 1
g = 9.81
b = 0.05
t = np.linspace(0, 10, 100)
initial_conditions = [np.pi/10,1]

# Calculate theta1 and theta2 at all time steps
# HINT: scipy.integrate.odeint returns a 2D matrix where
# the first column is the theta1 values and the second column is the theta2 values
# YOUR CODE HERE
theta = scipy.integrate.odeint(ode_func, initial_conditions, t, args=(g, l, m, b))
```

```python
print(theta)
```

    [[ 0.31415927  1.        ]
     [ 0.398018    0.64771035]
     [ 0.44306772  0.23782479]
     [ 0.44541221 -0.19139118]
     [ 0.4050558  -0.60124506]
     [ 0.32587182 -0.95407744]
     [ 0.21540056 -1.21543582]
     [ 0.08434552 -1.35790867]
     [-0.05437793 -1.36587406]
     [-0.18702947 -1.23896065]
     [-0.30062355 -0.99218707]
     [-0.38434018 -0.6526558 ]
     [-0.4305026  -0.25469205]
     [-0.43508437  0.16430654]
     [-0.39788089  0.56636046]
     [-0.32248769  0.91446998]
     [-0.21610397  1.17475525]
     [-0.08904618  1.32010269]
     [ 0.04616388  1.3346923 ]
     [ 0.17613041  1.21738761]
     [ 0.28811269  0.98204388]
     [ 0.37141091  0.6545109 ]
     [ 0.41834859  0.26802143]
     [ 0.42480095 -0.14089232]
     [ 0.39038601 -0.53496928]
     [ 0.31845253 -0.87787375]
     [ 0.21588882 -1.13631559]
     [ 0.09265104 -1.28353388]
     [-0.03911396 -1.30361605]
     [-0.1663489  -1.19480033]
     [-0.27656608 -0.96991274]
     [-0.35920571 -0.65364313]
     [-0.40661823 -0.27817347]
     [-0.41460943  0.12082177]
     [-0.38264889  0.50680309]
     [-0.31386723  0.84410259]
     [-0.21486975  1.10003385]
     [-0.09527732  1.24823432]
     [ 0.03311995  1.27278942]
     [ 0.15759615  1.17143583]
     [ 0.26592229  0.95609385]
     [ 0.34769527  0.65038302]
     [ 0.39531587  0.28547918]
     [ 0.4045465  -0.10378911]
     [ 0.37473486 -0.48160467]
     [ 0.30881971 -0.81297023]
     [ 0.21314903 -1.06581477]
     [ 0.09703176 -1.21421233]
     [-0.0280812  -1.24232403]
     [-0.14978745 -1.14749441]
     [-0.2561203  -0.9408506 ]
     [-0.33684714 -0.64502755]
     [-0.38443955 -0.29024138]
     [-0.39464008  0.08950979]
     [-0.36669871  0.45912963]
     [-0.30338652  0.78429348]
     [-0.21081781  1.03355532]
     [-0.09801129  1.18145801]
     [ 0.02390466  1.21230494]
     [ 0.14284272  1.12314448]
     [ 0.24710037  0.92441299]
     [ 0.3266268   0.63784241]
     [ 0.37398233  0.29273627]
     [ 0.38491068 -0.07771987]
     [ 0.35858614 -0.43914702]
     [ 0.29763419 -0.75789414]
     [ 0.20795719 -1.00314801]
     [ 0.09830372 -1.14994717]
     [-0.02050443 -1.1827944 ]
     [-0.13668639 -1.09852661]
     [-0.23880433 -0.90698197]
     [-0.31699849 -0.6290658 ]
     [-0.3639334  -0.29321584]
     [-0.37537275  0.06817479]
     [-0.35043508  0.42143914]
     [-0.2916205   0.73359985]
     [-0.20463935  0.9744828 ]
     [-0.09798863  1.11964476]
     [ 0.01780143  1.15383682]
     [ 0.13124769  1.07375836]
     [ 0.23117637  0.8887328 ]
     [ 0.30792618  0.61891031]
     [ 0.35427918  0.29190893]
     [ 0.36603578 -0.06064914]
     [ 0.34227687 -0.40580225]
     [ 0.28539555 -0.71124587]
     [ 0.20092836 -0.94744985]
     [ 0.09713789 -1.09050795]
     [-0.01572309 -1.12546133]
     [-0.12646046 -1.0489367 ]
     [-0.22416306 -0.86981778]
     [-0.29937384 -0.60756581]
     [-0.34500393 -0.28902375]
     [-0.35690524  0.05493486]
     [-0.33413724  0.39204552]
     [-0.27900277  0.69067486]
     [-0.19688124  0.92194029]
     [-0.09581655  1.06248806]
     [ 0.01420275  1.09768504]
     [ 0.12226303  1.02414186]]

Now that we have all of the $\theta_{1}$ and $\theta_{2}$ values from scipy.integrate.odeint, we can find the position of the point mass (x, y) values as
\
\
$$ x = l*sin(\theta_{1})$$
$$ y = - l*cos(\theta_{1})$$
\
For every time step, plot and save the (x, y) position using the $\theta$ values calculated above as an animation.
* Save the animation as "pendulum.mp4"
* Set y-axis limits as -1.1, 0
* Set x-axis limits as -0.5, 0.5
* If you would like to animate the string, include a line at each frame:
    "plt.plot( [ 0, x[i] ], [ 0, y[i] ] )" that will plot a line between the origin and the pendulum in the current frame i

```python
# Calculate x and y from theta1
# YOUR CODE HERE
theta1 = theta[:,0]
x = [l * np.sin(th) for th in theta1]
y = [-l * np.cos(th) for th in theta1]

# Initialize writer
# YOUR CODE HERE
metadata = dict(title='2D animation', artist='Matplotlib')
writer = FFMpegWriter(fps=100, metadata=metadata, bitrate=200000)
fig = plt.figure(dpi=200)

## SAVE AS MP4 ##
# YOUR CODE HERE
fig = plt.figure()

with writer.saving(fig, "pendulum.mp4", dpi=200):
    for i in range(len(t)):

        fig.clear()

        plt.plot([0, x[i]], [0, y[i]]) # plots the string
        plt.plot(x[i], y[i], marker="o", markersize=10) # plots the point mass
        plt.xlim([-0.5,0.5])
        plt.ylim([-1.1,0])

        plt.draw()
        plt.pause(0.01)
        writer.grab_frame()
```

```python

```