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

import models
import payoffs
import selection_functions
```

```python
x1, x2, x3 = sym.symbols("x1, x2, x3", real=True, nonnegative=True)
T, R, P, S = sym.symbols('T, R, P, S', real=True, positive=True)
UG = selection_functions.UGA
Ug = selection_functions.UgA
```

# Equilibria with $x_g=1$

Characterize set of equilibria where all females are carrying the $g$ allele of the $\gamma$ gene. In such equilibria we have $x_1=x_2=0$ and $x_4=1-x_3$

```python
# start with generic system
x4 = 1 - x1 - x2 - x3
x = np.array([[x1], [x2], [x3], [x4]])
payoff_kernel = np.array([[R, S], [T, P]])
W = models.generalized_sexual_selection(x, UG, Ug, payoff_kernel, 0, 0, 0)
_, _, (N3,), (N4,) = models._offspring_by_genotype(W, x)
N, = models.total_offspring(W, x)
```

```python
sym.factor(N3.subs({x1: 0, x2: 0}), Ug(x3))
```

    2.0*(0.5*P*x3 + (-1.0*P*x3 + 0.5*S*x3 + 0.5*S + 0.5*T*x3)*UgA(x3) + (0.5*P*x3 + 0.5*R*x3 + 0.5*R - 0.5*S*x3 - 0.5*S - 0.5*T*x3)*UgA(x3)**2)

```python
sym.simplify(N4.subs({x1: 0, x2: 0}))
```

    x3*(P*(UgA(x3) - 1) - T*UgA(x3))*(UgA(x3) - 1) - (x3 - 1)*(2.0*(P*(UgA(x3) - 1) - T*UgA(x3))*(UgA(x3) - 1) + 1.0*(R*UgA(x3) - S*(UgA(x3) - 1))*UgA(x3))

```python
sym.simplify(N.subs({x1: 0, x2: 0}))
```

    2.0*P*UgA(x3)**2 - 4.0*P*UgA(x3) + 2.0*P + 2.0*R*UgA(x3)**2 - 2.0*S*UgA(x3)**2 + 2.0*S*UgA(x3) - 2.0*T*UgA(x3)**2 + 2.0*T*UgA(x3)

$$ F\big(Ug(x_3)\big) = 2\big(P + U_g(x_3)(-2P + S + T) + U_g(x_3)^{2}(P + R - S - T)\big) $$

# Locus of potential equilibria

```python
(f1,), (f2,), (f3,), _ = models.offspring_genotypes_evolution(W, x)
```

```python
xA = sym.symbols("xA", real=True, nonnegative=True)
xG = sym.symbols("xG", real=True, nonnegative=True)

_, denominator = sym.fraction(sym.together(f1.subs({x1 + x3: xA}), x1))
```

```python
f1
```

    -x1 + (x1*(1.0*x1*(2*R*UGA(x1 + x3) + 2*S*(-UGA(x1 + x3) + 1))*UGA(x1 + x3)/(x1 + x3) + 0.5*x2*(2*P*(-UGA(x1 + x3) + 1) + 2*T*UGA(x1 + x3))*(-UGA(x1 + x3) + 1)/(-x1 - x3 + 1) + 0.5*x3*(2*R*UGA(x1 + x3) + 2*S*(-UGA(x1 + x3) + 1))*UGA(x1 + x3)/(x1 + x3) + 0.25*(2*P*(-UGA(x1 + x3) + 1) + 2*T*UGA(x1 + x3))*(-UGA(x1 + x3) + 1)*(-x1 - x2 - x3 + 1)/(-x1 - x3 + 1)) + 0.25*x1*(2*R*UgA(x1 + x3) + 2*S*(-UgA(x1 + x3) + 1))*(-x1 - x2 - x3 + 1)*UgA(x1 + x3)/(x1 + x3) + x2*(0.5*x1*(2*R*UGA(x1 + x3) + 2*S*(-UGA(x1 + x3) + 1))*UGA(x1 + x3)/(x1 + x3) + 0.25*x3*(2*R*UGA(x1 + x3) + 2*S*(-UGA(x1 + x3) + 1))*UGA(x1 + x3)/(x1 + x3)) + x3*(0.5*x1*(2*R*UgA(x1 + x3) + 2*S*(-UgA(x1 + x3) + 1))*UgA(x1 + x3)/(x1 + x3) + 0.25*x2*(2*P*(-UgA(x1 + x3) + 1) + 2*T*UgA(x1 + x3))*(-UgA(x1 + x3) + 1)/(-x1 - x3 + 1)))/(0.25*x1*(2*P*(-UGA(x1 + x3) + 1) + 2*T*UGA(x1 + x3))*(-UGA(x1 + x3) + 1)*(-x1 - x2 - x3 + 1)/(-x1 - x3 + 1) + x1*(0.5*x2*(2*P*(-UGA(x1 + x3) + 1) + 2*T*UGA(x1 + x3))*(-UGA(x1 + x3) + 1)/(-x1 - x3 + 1) + 0.25*(2*P*(-UGA(x1 + x3) + 1) + 2*T*UGA(x1 + x3))*(-UGA(x1 + x3) + 1)*(-x1 - x2 - x3 + 1)/(-x1 - x3 + 1)) + x1*(0.5*x3*(2*R*UGA(x1 + x3) + 2*S*(-UGA(x1 + x3) + 1))*UGA(x1 + x3)/(x1 + x3) + 0.25*(2*P*(-UGA(x1 + x3) + 1) + 2*T*UGA(x1 + x3))*(-UGA(x1 + x3) + 1)*(-x1 - x2 - x3 + 1)/(-x1 - x3 + 1)) + x1*(1.0*x1*(2*R*UGA(x1 + x3) + 2*S*(-UGA(x1 + x3) + 1))*UGA(x1 + x3)/(x1 + x3) + 0.5*x2*(2*P*(-UGA(x1 + x3) + 1) + 2*T*UGA(x1 + x3))*(-UGA(x1 + x3) + 1)/(-x1 - x3 + 1) + 0.5*x3*(2*R*UGA(x1 + x3) + 2*S*(-UGA(x1 + x3) + 1))*UGA(x1 + x3)/(x1 + x3) + 0.25*(2*P*(-UGA(x1 + x3) + 1) + 2*T*UGA(x1 + x3))*(-UGA(x1 + x3) + 1)*(-x1 - x2 - x3 + 1)/(-x1 - x3 + 1)) + 0.25*x1*(2*R*UgA(x1 + x3) + 2*S*(-UgA(x1 + x3) + 1))*(-x1 - x2 - x3 + 1)*UgA(x1 + x3)/(x1 + x3) + 0.25*x2*x3*(2*P*(-UgA(x1 + x3) + 1) + 2*T*UgA(x1 + x3))*(-UgA(x1 + x3) + 1)/(-x1 - x3 + 1) + 0.25*x2*x3*(2*R*UGA(x1 + x3) + 2*S*(-UGA(x1 + x3) + 1))*UGA(x1 + x3)/(x1 + x3) + x2*(0.5*x1*(2*R*UGA(x1 + x3) + 2*S*(-UGA(x1 + x3) + 1))*UGA(x1 + x3)/(x1 + x3) + 0.25*x3*(2*R*UGA(x1 + x3) + 2*S*(-UGA(x1 + x3) + 1))*UGA(x1 + x3)/(x1 + x3)) + x2*(0.25*x3*(2*R*UGA(x1 + x3) + 2*S*(-UGA(x1 + x3) + 1))*UGA(x1 + x3)/(x1 + x3) + 0.5*(2*P*(-UGA(x1 + x3) + 1) + 2*T*UGA(x1 + x3))*(-UGA(x1 + x3) + 1)*(-x1 - x2 - x3 + 1)/(-x1 - x3 + 1)) + x2*(0.5*x1*(2*R*UGA(x1 + x3) + 2*S*(-UGA(x1 + x3) + 1))*UGA(x1 + x3)/(x1 + x3) + 1.0*x2*(2*P*(-UGA(x1 + x3) + 1) + 2*T*UGA(x1 + x3))*(-UGA(x1 + x3) + 1)/(-x1 - x3 + 1) + 0.25*x3*(2*R*UGA(x1 + x3) + 2*S*(-UGA(x1 + x3) + 1))*UGA(x1 + x3)/(x1 + x3) + 0.5*(2*P*(-UGA(x1 + x3) + 1) + 2*T*UGA(x1 + x3))*(-UGA(x1 + x3) + 1)*(-x1 - x2 - x3 + 1)/(-x1 - x3 + 1)) + x3*(0.5*x1*(2*R*UgA(x1 + x3) + 2*S*(-UgA(x1 + x3) + 1))*UgA(x1 + x3)/(x1 + x3) + 0.25*x2*(2*P*(-UgA(x1 + x3) + 1) + 2*T*UgA(x1 + x3))*(-UgA(x1 + x3) + 1)/(-x1 - x3 + 1)) + x3*(0.25*x2*(2*P*(-UgA(x1 + x3) + 1) + 2*T*UgA(x1 + x3))*(-UgA(x1 + x3) + 1)/(-x1 - x3 + 1) + 0.5*(2*P*(-UgA(x1 + x3) + 1) + 2*T*UgA(x1 + x3))*(-UgA(x1 + x3) + 1)*(-x1 - x2 - x3 + 1)/(-x1 - x3 + 1)) + x3*(0.5*x1*(2*R*UgA(x1 + x3) + 2*S*(-UgA(x1 + x3) + 1))*UgA(x1 + x3)/(x1 + x3) + 0.25*x2*(2*P*(-UgA(x1 + x3) + 1) + 2*T*UgA(x1 + x3))*(-UgA(x1 + x3) + 1)/(-x1 - x3 + 1) + 1.0*x3*(2*R*UgA(x1 + x3) + 2*S*(-UgA(x1 + x3) + 1))*UgA(x1 + x3)/(x1 + x3) + 0.5*(2*P*(-UgA(x1 + x3) + 1) + 2*T*UgA(x1 + x3))*(-UgA(x1 + x3) + 1)*(-x1 - x2 - x3 + 1)/(-x1 - x3 + 1)) + (0.25*x1*(2*R*UgA(x1 + x3) + 2*S*(-UgA(x1 + x3) + 1))*UgA(x1 + x3)/(x1 + x3) + 0.5*x2*(2*P*(-UgA(x1 + x3) + 1) + 2*T*UgA(x1 + x3))*(-UgA(x1 + x3) + 1)/(-x1 - x3 + 1))*(-x1 - x2 - x3 + 1) + (0.25*x1*(2*R*UgA(x1 + x3) + 2*S*(-UgA(x1 + x3) + 1))*UgA(x1 + x3)/(x1 + x3) + 0.5*x3*(2*R*UgA(x1 + x3) + 2*S*(-UgA(x1 + x3) + 1))*UgA(x1 + x3)/(x1 + x3))*(-x1 - x2 - x3 + 1) + (-x1 - x2 - x3 + 1)*(0.25*x1*(2*R*UgA(x1 + x3) + 2*S*(-UgA(x1 + x3) + 1))*UgA(x1 + x3)/(x1 + x3) + 0.5*x2*(2*P*(-UgA(x1 + x3) + 1) + 2*T*UgA(x1 + x3))*(-UgA(x1 + x3) + 1)/(-x1 - x3 + 1) + 0.5*x3*(2*R*UgA(x1 + x3) + 2*S*(-UgA(x1 + x3) + 1))*UgA(x1 + x3)/(x1 + x3) + 1.0*(2*P*(-UgA(x1 + x3) + 1) + 2*T*UgA(x1 + x3))*(-UgA(x1 + x3) + 1)*(-x1 - x2 - x3 + 1)/(-x1 - x3 + 1)))

```python

```

```python
F = sym.Matrix([f1, f2, f3])
```

```python
_, _, Lg = F.subs({x1: 0, x2: 0})
```

```python
equilibrium_Ug = sym.symbols("equilibrium_Ug", nonnegative=True, real=True)
equilibrium_x3, = sym.solve(Lg.subs({Ug(x3): equilibrium_Ug}), x3)
```

```python
sym.factor(equilibrium_x3, equilibrium_Ug)
```

    0.5*equilibrium_Ug*(S + equilibrium_Ug*(R - S))/(0.5*P + equilibrium_Ug**2*(0.5*P + 0.5*R - 0.5*S - 0.5*T) + equilibrium_Ug*(-1.0*P + 0.5*S + 0.5*T))

$$ x_3^* = \frac {2U_g^{*2}R + 2SU_g^*(1 - U_g^{*})}{F\big(Ug^*\big)} $$

# Invadability

Equilibria with $x_g=1$ are invadable if and only if one of the eigenvalues of the Jacobian matrix of the dynamical system evaluated at the equilibrium is strictly positive.

```python
F_jac = F.jacobian([x1, x2, x3])
```

```python
equilibrium_UG = sym.symbols("equilibrium_UG", nonnegative=True, real=True)
equilibrium_Ug_prime = sym.symbols("equilibrium_Ug_prime", nonnegative=True, real=True)

equilibrium_F_jac = (F_jac.subs({x1: 0, x2: 0})
                          .doit()
                          .subs({UG(x3): equilibrium_UG, Ug(x3): equilibrium_Ug, sym.Derivative(Ug(x3), x3): equilibrium_Ug_prime}))
```

```python
sym.factor(equilibrium_F_jac[0,0], equilibrium_UG, equilibrium_Ug, x3)
```

    1.0*(-0.375*P*x3 + 0.125*S*equilibrium_Ug + equilibrium_UG**2*x3*(0.125*P + 0.25*R - 0.25*S - 0.125*T) + equilibrium_UG*x3*(-0.25*P + 0.25*S + 0.125*T) + equilibrium_Ug**2*x3*(-0.5*P - 0.375*R + 0.375*S + 0.5*T) + equilibrium_Ug**2*(0.125*R - 0.125*S) + equilibrium_Ug*x3*(1.0*P - 0.375*S - 0.5*T))/(x3*(0.5*P + equilibrium_Ug**2*(0.5*P + 0.5*R - 0.5*S - 0.5*T) + equilibrium_Ug*(-1.0*P + 0.5*S + 0.5*T)))

```python
sym.cancel(sym.factor(equilibrium_F_jac[0,0].subs({x3: equilibrium_x3}), equilibrium_UG, equilibrium_Ug))
```

    1.0*(0.125*P*equilibrium_UG**2 - 0.25*P*equilibrium_UG - 0.375*P*equilibrium_Ug**2 + 0.75*P*equilibrium_Ug - 0.25*P + 0.25*R*equilibrium_UG**2 - 0.25*R*equilibrium_Ug**2 - 0.25*S*equilibrium_UG**2 + 0.25*S*equilibrium_UG + 0.25*S*equilibrium_Ug**2 - 0.25*S*equilibrium_Ug - 0.125*T*equilibrium_UG**2 + 0.125*T*equilibrium_UG + 0.375*T*equilibrium_Ug**2 - 0.375*T*equilibrium_Ug)/(0.5*P*equilibrium_Ug**2 - 1.0*P*equilibrium_Ug + 0.5*P + 0.5*R*equilibrium_Ug**2 - 0.5*S*equilibrium_Ug**2 + 0.5*S*equilibrium_Ug - 0.5*T*equilibrium_Ug**2 + 0.5*T*equilibrium_Ug)

```python
1/6
```

    0.16666666666666666

$$ \frac{(-2P + U_G^{*2}(P + 2R - 2S - T) + U_G^*(-2P + 2S + T) + U_g^{*2}(-3P - 2R + 2S + 3T) + U_g^*(6P - 2S - 3T))}{2F(U_g^*)} $$

```python
n, d = sym.fraction(sym.simplify(sym.factor(equilibrium_F_jac[0,0].subs({x3: equilibrium_x3}), equilibrium_UG, equilibrium_Ug) + sym.factor(equilibrium_F_jac[1,1].subs({x3: equilibrium_x3}), equilibrium_UG, equilibrium_Ug)))
```

```python
sym.factor(n, equilibrium_UG, equilibrium_Ug)
```

    1.25*(-0.2*P + equilibrium_UG**2*(0.3*P + 0.3*R - 0.3*S - 0.3*T) + equilibrium_UG*(-0.6*P + 0.3*S + 0.3*T) + equilibrium_Ug**2*(-0.5*P - 0.5*R + 0.5*S + 0.5*T) + equilibrium_Ug*(1.0*P - 0.5*S - 0.5*T))

```python
0.3*1.25
```

    0.375

```python
sym.factor(d, equilibrium_UG, equilibrium_Ug)
```

    1.0*(0.5*P + equilibrium_Ug**2*(0.5*P + 0.5*R - 0.5*S - 0.5*T) + equilibrium_Ug*(-1.0*P + 0.5*S + 0.5*T))

```python

```

```python
sym.factor(equilibrium_F_jac[0,1], equilibrium_UG, equilibrium_Ug)
```

    -0.125*(-P*x3 + equilibrium_UG**2*(R*x3 - R - S*x3 + S) + equilibrium_UG*(S*x3 - S) + equilibrium_Ug**2*(-P*x3 + T*x3) + equilibrium_Ug*(2*P*x3 - T*x3))/((-x3 + 1)*(0.5*P + equilibrium_Ug**2*(0.5*P + 0.5*R - 0.5*S - 0.5*T) + equilibrium_Ug*(-1.0*P + 0.5*S + 0.5*T)))

```python
sym.factor(equilibrium_F_jac[0,2], equilibrium_UG, equilibrium_Ug)
```

    0

```python
sym.factor(equilibrium_F_jac[1,0], equilibrium_UG, equilibrium_Ug)
```

    0.125*(P*x3 + equilibrium_UG**2*(P*x3 - T*x3) + equilibrium_UG*(-2*P*x3 + T*x3) + equilibrium_Ug**2*(-R*x3 + R + S*x3 - S) + equilibrium_Ug*(-S*x3 + S))/(x3*(0.5*P + equilibrium_Ug**2*(0.5*P + 0.5*R - 0.5*S - 0.5*T) + equilibrium_Ug*(-1.0*P + 0.5*S + 0.5*T)))

```python
sym.factor(equilibrium_F_jac[1,1], equilibrium_UG, equilibrium_Ug)
```

    -0.75*(-0.166666666666667*P*x3 + equilibrium_UG**2*(0.333333333333333*P*x3 - 0.333333333333333*P + 0.166666666666667*R*x3 - 0.166666666666667*R - 0.166666666666667*S*x3 + 0.166666666666667*S - 0.333333333333333*T*x3 + 0.333333333333333*T) + equilibrium_UG*(-0.666666666666667*P*x3 + 0.666666666666667*P + 0.166666666666667*S*x3 - 0.166666666666667*S + 0.333333333333333*T*x3 - 0.333333333333333*T) + equilibrium_Ug**2*(-0.5*P*x3 + 0.333333333333333*P - 0.666666666666667*R*x3 + 0.666666666666667*R + 0.666666666666667*S*x3 - 0.666666666666667*S + 0.5*T*x3 - 0.333333333333333*T) + equilibrium_Ug*(1.0*P*x3 - 0.666666666666667*P - 0.666666666666667*S*x3 + 0.666666666666667*S - 0.5*T*x3 + 0.333333333333333*T))/((-x3 + 1)*(0.5*P + equilibrium_Ug**2*(0.5*P + 0.5*R - 0.5*S - 0.5*T) + equilibrium_Ug*(-1.0*P + 0.5*S + 0.5*T)))

```python
sym.factor(equilibrium_F_jac[1,1].subs({x3: equilibrium_x3}), equilibrium_UG, equilibrium_Ug)
```

    0.5*(equilibrium_UG**2*(0.5*P + 0.25*R - 0.25*S - 0.5*T) + equilibrium_UG*(-1.0*P + 0.25*S + 0.5*T) + equilibrium_Ug**2*(-0.5*P - 0.75*R + 0.75*S + 0.5*T) + equilibrium_Ug*(1.0*P - 0.75*S - 0.5*T))/(0.5*P + equilibrium_Ug**2*(0.5*P + 0.5*R - 0.5*S - 0.5*T) + equilibrium_Ug*(-1.0*P + 0.5*S + 0.5*T))

```python
numerator, denominator = sym.fraction(sym.simplify(sym.factor(equilibrium_F_jac[0,0].subs({x3: equilibrium_x3}), equilibrium_UG, equilibrium_Ug) * sym.factor(equilibrium_F_jac[1,1].subs({x3: equilibrium_x3}), equilibrium_UG, equilibrium_Ug) - sym.factor(equilibrium_F_jac[0,1].subs({x3: equilibrium_x3}), equilibrium_UG, equilibrium_Ug)*sym.factor(equilibrium_F_jac[1,0].subs({x3: equilibrium_x3}), equilibrium_UG, equilibrium_Ug)))
```

```python
denominator
```

    (0.5*P + 0.5*equilibrium_Ug**2*(P + R - S - T) + equilibrium_Ug*(-1.0*P + 0.5*S + 0.5*T))**2

```python
sym.factor(numerator)
```

    0.375*(1.0*equilibrium_UG - 1.0*equilibrium_Ug)*(0.5*P*equilibrium_UG + 0.5*P*equilibrium_Ug - 1.0*P + 0.5*R*equilibrium_UG + 0.5*R*equilibrium_Ug - 0.5*S*equilibrium_UG - 0.5*S*equilibrium_Ug + 0.5*S - 0.5*T*equilibrium_UG - 0.5*T*equilibrium_Ug + 0.5*T)*(0.166666666666667*P*equilibrium_UG**2 - 0.333333333333333*P*equilibrium_UG - 0.5*P*equilibrium_Ug**2 + 1.0*P*equilibrium_Ug - 0.333333333333333*P + 0.166666666666667*R*equilibrium_UG**2 - 0.5*R*equilibrium_Ug**2 - 0.166666666666667*S*equilibrium_UG**2 + 0.166666666666667*S*equilibrium_UG + 0.5*S*equilibrium_Ug**2 - 0.5*S*equilibrium_Ug - 0.166666666666667*T*equilibrium_UG**2 + 0.166666666666667*T*equilibrium_UG + 0.5*T*equilibrium_Ug**2 - 0.5*T*equilibrium_Ug)

```python
p = sym.poly(numerator, equilibrium_UG, equilibrium_Ug)
```

```python
p.coeffs()
```

    [0.03125*P**2 + 0.0625*P*R - 0.0625*P*S - 0.0625*P*T + 0.03125*R**2 - 0.0625*R*S - 0.0625*R*T + 0.03125*S**2 + 0.0625*S*T + 0.03125*T**2,
     -0.125*P**2 - 0.125*P*R + 0.1875*P*S + 0.1875*P*T + 0.0625*R*S + 0.0625*R*T - 0.0625*S**2 - 0.125*S*T - 0.0625*T**2,
     -0.125*P**2 - 0.25*P*R + 0.25*P*S + 0.25*P*T - 0.125*R**2 + 0.25*R*S + 0.25*R*T - 0.125*S**2 - 0.25*S*T - 0.125*T**2,
     0.25*P**2 + 0.25*P*R - 0.375*P*S - 0.375*P*T - 0.125*R*S - 0.125*R*T + 0.125*S**2 + 0.25*S*T + 0.125*T**2,
     0.0625*P**2 - 0.0625*P*R - 0.0625*P*S - 0.0625*P*T + 0.03125*S**2 + 0.0625*S*T + 0.03125*T**2,
     0.25*P**2 + 0.25*P*R - 0.375*P*S - 0.375*P*T - 0.125*R*S - 0.125*R*T + 0.125*S**2 + 0.25*S*T + 0.125*T**2,
     -0.5*P**2 + 0.5*P*S + 0.5*P*T - 0.125*S**2 - 0.25*S*T - 0.125*T**2,
     0.125*P**2 - 0.0625*P*S - 0.0625*P*T,
     0.09375*P**2 + 0.1875*P*R - 0.1875*P*S - 0.1875*P*T + 0.09375*R**2 - 0.1875*R*S - 0.1875*R*T + 0.09375*S**2 + 0.1875*S*T + 0.09375*T**2,
     -0.375*P**2 - 0.375*P*R + 0.5625*P*S + 0.5625*P*T + 0.1875*R*S + 0.1875*R*T - 0.1875*S**2 - 0.375*S*T - 0.1875*T**2,
     0.4375*P**2 + 0.0625*P*R - 0.4375*P*S - 0.4375*P*T + 0.09375*S**2 + 0.1875*S*T + 0.09375*T**2,
     -0.125*P**2 + 0.0625*P*S + 0.0625*P*T]

```python
sym.factor(equilibrium_F_jac[1,2], equilibrium_UG, equilibrium_Ug)
```

    0

```python
sym.factor(equilibrium_F_jac[2,0], equilibrium_UG, equilibrium_Ug)
```

    -0.75*(-0.0833333333333333*P**2*x3 - 0.166666666666667*P*S*equilibrium_Ug_prime*x3 + equilibrium_UG**2*equilibrium_Ug**2*(0.166666666666667*P**2*x3**2 - 0.0833333333333333*P**2*x3 + 0.333333333333333*P*R*x3**2 - 0.0833333333333333*P*R*x3 - 0.333333333333333*P*S*x3**2 + 0.0833333333333333*P*S*x3 - 0.333333333333333*P*T*x3**2 + 0.166666666666667*P*T*x3 + 0.166666666666667*R**2*x3**2 - 0.333333333333333*R*S*x3**2 - 0.333333333333333*R*T*x3**2 + 0.0833333333333333*R*T*x3 + 0.166666666666667*S**2*x3**2 + 0.333333333333333*S*T*x3**2 - 0.0833333333333333*S*T*x3 + 0.166666666666667*T**2*x3**2 - 0.0833333333333333*T**2*x3) + equilibrium_UG**2*equilibrium_Ug*(-0.333333333333333*P**2*x3**2 + 0.166666666666667*P**2*x3 - 0.333333333333333*P*R*x3**2 + 0.333333333333333*P*R*x3 + 0.5*P*S*x3**2 - 0.25*P*S*x3 + 0.5*P*T*x3**2 - 0.25*P*T*x3 + 0.166666666666667*R*S*x3**2 + 0.166666666666667*R*T*x3**2 - 0.166666666666667*R*T*x3 - 0.166666666666667*S**2*x3**2 - 0.333333333333333*S*T*x3**2 + 0.0833333333333333*S*T*x3 - 0.166666666666667*T**2*x3**2 + 0.0833333333333333*T**2*x3) + equilibrium_UG**2*(0.166666666666667*P**2*x3**2 - 0.0833333333333333*P**2*x3 + 0.166666666666667*P*R*x3**2 - 0.166666666666667*P*R*x3 - 0.166666666666667*P*S*x3**2 + 0.166666666666667*P*S*x3 - 0.166666666666667*P*T*x3**2 + 0.0833333333333333*P*T*x3) + equilibrium_UG*equilibrium_Ug**2*(-0.333333333333333*P**2*x3**2 + 0.166666666666667*P**2*x3 - 0.333333333333333*P*R*x3**2 - 0.166666666666667*P*R*x3 + 0.5*P*S*x3**2 + 0.5*P*T*x3**2 - 0.25*P*T*x3 + 0.166666666666667*R*S*x3**2 + 0.166666666666667*R*T*x3**2 + 0.0833333333333333*R*T*x3 - 0.166666666666667*S**2*x3**2 - 0.333333333333333*S*T*x3**2 + 0.0833333333333333*S*T*x3 - 0.166666666666667*T**2*x3**2 + 0.0833333333333333*T**2*x3) + equilibrium_UG*equilibrium_Ug*(0.666666666666667*P**2*x3**2 - 0.333333333333333*P**2*x3 - 0.666666666666667*P*S*x3**2 + 0.166666666666667*P*S*x3 - 0.666666666666667*P*T*x3**2 + 0.333333333333333*P*T*x3 + 0.166666666666667*S**2*x3**2 + 0.333333333333333*S*T*x3**2 - 0.0833333333333333*S*T*x3 + 0.166666666666667*T**2*x3**2 - 0.0833333333333333*T**2*x3) + equilibrium_UG*(-0.333333333333333*P**2*x3**2 + 0.166666666666667*P**2*x3 + 0.166666666666667*P*S*x3**2 - 0.166666666666667*P*S*x3 + 0.166666666666667*P*T*x3**2 - 0.0833333333333333*P*T*x3) + equilibrium_Ug**4*(-0.166666666666667*P**2*x3**2 - 0.333333333333333*P*R*x3**2 + 0.0833333333333333*P*R*x3 + 0.0833333333333333*P*R + 0.333333333333333*P*S*x3**2 - 0.0833333333333333*P*S*x3 - 0.0833333333333333*P*S + 0.333333333333333*P*T*x3**2 - 0.166666666666667*R**2*x3**2 + 0.0833333333333333*R**2*x3 + 0.0833333333333333*R**2 + 0.333333333333333*R*S*x3**2 - 0.166666666666667*R*S*x3 - 0.166666666666667*R*S + 0.333333333333333*R*T*x3**2 - 0.0833333333333333*R*T*x3 - 0.0833333333333333*R*T - 0.166666666666667*S**2*x3**2 + 0.0833333333333333*S**2*x3 + 0.0833333333333333*S**2 - 0.333333333333333*S*T*x3**2 + 0.0833333333333333*S*T*x3 + 0.0833333333333333*S*T - 0.166666666666667*T**2*x3**2) + equilibrium_Ug**3*(0.666666666666667*P**2*x3**2 + 0.666666666666667*P*R*x3**2 - 0.166666666666667*P*R*x3 - 0.166666666666667*P*R - 1.0*P*S*x3**2 + 0.25*P*S*x3 + 0.25*P*S - 1.0*P*T*x3**2 - 0.333333333333333*R*S*x3**2 + 0.166666666666667*R*S*x3 + 0.166666666666667*R*S - 0.333333333333333*R*T*x3**2 + 0.0833333333333333*R*T*x3 + 0.0833333333333333*R*T + 0.333333333333333*S**2*x3**2 - 0.166666666666667*S**2*x3 - 0.166666666666667*S**2 + 0.666666666666667*S*T*x3**2 - 0.166666666666667*S*T*x3 - 0.166666666666667*S*T + 0.333333333333333*T**2*x3**2) + equilibrium_Ug**2*(-0.833333333333333*P**2*x3**2 - 0.0833333333333333*P**2*x3 + 0.333333333333333*P*R*equilibrium_Ug_prime*x3 - 0.166666666666667*P*R*x3**2 + 0.166666666666667*P*R*x3 + 0.0833333333333333*P*R - 0.166666666666667*P*S*equilibrium_Ug_prime*x3 + 0.833333333333333*P*S*x3**2 - 0.333333333333333*P*S*x3 - 0.25*P*S + 0.833333333333333*P*T*x3**2 + 0.0833333333333333*P*T*x3 - 0.166666666666667*R*T*equilibrium_Ug_prime*x3 - 0.166666666666667*S**2*x3**2 + 0.0833333333333333*S**2*x3 + 0.0833333333333333*S**2 - 0.333333333333333*S*T*x3**2 + 0.0833333333333333*S*T*x3 + 0.0833333333333333*S*T - 0.166666666666667*T**2*x3**2) + equilibrium_Ug*(0.333333333333333*P**2*x3**2 + 0.166666666666667*P**2*x3 - 0.333333333333333*P*R*equilibrium_Ug_prime*x3 + 0.333333333333333*P*S*equilibrium_Ug_prime*x3 - 0.166666666666667*P*S*x3**2 + 0.166666666666667*P*S*x3 + 0.0833333333333333*P*S - 0.166666666666667*P*T*x3**2 - 0.0833333333333333*P*T*x3))/(x3*(0.5*P + equilibrium_Ug**2*(0.5*P + 0.5*R - 0.5*S - 0.5*T) + equilibrium_Ug*(-1.0*P + 0.5*S + 0.5*T))**2)

```python
sym.factor(equilibrium_F_jac[2,1], equilibrium_UG, equilibrium_Ug)
```

    0.75*(-0.0833333333333333*P**2*x3 + equilibrium_UG**2*equilibrium_Ug**2*(0.166666666666667*P**2*x3**2 - 0.166666666666667*P**2*x3 + 0.333333333333333*P*R*x3**2 - 0.25*P*R*x3 - 0.0833333333333333*P*R - 0.333333333333333*P*S*x3**2 + 0.25*P*S*x3 + 0.0833333333333333*P*S - 0.333333333333333*P*T*x3**2 + 0.333333333333333*P*T*x3 + 0.166666666666667*R**2*x3**2 - 0.0833333333333333*R**2*x3 - 0.0833333333333333*R**2 - 0.333333333333333*R*S*x3**2 + 0.166666666666667*R*S*x3 + 0.166666666666667*R*S - 0.333333333333333*R*T*x3**2 + 0.25*R*T*x3 + 0.0833333333333333*R*T + 0.166666666666667*S**2*x3**2 - 0.0833333333333333*S**2*x3 - 0.0833333333333333*S**2 + 0.333333333333333*S*T*x3**2 - 0.25*S*T*x3 - 0.0833333333333333*S*T + 0.166666666666667*T**2*x3**2 - 0.166666666666667*T**2*x3) + equilibrium_UG**2*equilibrium_Ug*(-0.333333333333333*P**2*x3**2 + 0.333333333333333*P**2*x3 - 0.333333333333333*P*R*x3**2 + 0.5*P*R*x3 - 0.166666666666667*P*R + 0.5*P*S*x3**2 - 0.5*P*S*x3 + 0.5*P*T*x3**2 - 0.5*P*T*x3 + 0.166666666666667*R*S*x3**2 - 0.0833333333333333*R*S*x3 - 0.0833333333333333*R*S + 0.166666666666667*R*T*x3**2 - 0.25*R*T*x3 + 0.0833333333333333*R*T - 0.166666666666667*S**2*x3**2 + 0.0833333333333333*S**2*x3 + 0.0833333333333333*S**2 - 0.333333333333333*S*T*x3**2 + 0.25*S*T*x3 + 0.0833333333333333*S*T - 0.166666666666667*T**2*x3**2 + 0.166666666666667*T**2*x3) + equilibrium_UG**2*(0.166666666666667*P**2*x3**2 - 0.166666666666667*P**2*x3 + 0.166666666666667*P*R*x3**2 - 0.25*P*R*x3 + 0.0833333333333333*P*R - 0.166666666666667*P*S*x3**2 + 0.25*P*S*x3 - 0.0833333333333333*P*S - 0.166666666666667*P*T*x3**2 + 0.166666666666667*P*T*x3) + equilibrium_UG*equilibrium_Ug**2*(-0.333333333333333*P**2*x3**2 + 0.333333333333333*P**2*x3 - 0.333333333333333*P*R*x3**2 + 0.333333333333333*P*R + 0.5*P*S*x3**2 - 0.25*P*S*x3 - 0.25*P*S + 0.5*P*T*x3**2 - 0.5*P*T*x3 + 0.166666666666667*R*S*x3**2 - 0.0833333333333333*R*S*x3 - 0.0833333333333333*R*S + 0.166666666666667*R*T*x3**2 - 0.166666666666667*R*T - 0.166666666666667*S**2*x3**2 + 0.0833333333333333*S**2*x3 + 0.0833333333333333*S**2 - 0.333333333333333*S*T*x3**2 + 0.25*S*T*x3 + 0.0833333333333333*S*T - 0.166666666666667*T**2*x3**2 + 0.166666666666667*T**2*x3) + equilibrium_UG*equilibrium_Ug*(0.666666666666667*P**2*x3**2 - 0.666666666666667*P**2*x3 - 0.666666666666667*P*S*x3**2 + 0.5*P*S*x3 + 0.166666666666667*P*S - 0.666666666666667*P*T*x3**2 + 0.666666666666667*P*T*x3 + 0.166666666666667*S**2*x3**2 - 0.0833333333333333*S**2*x3 - 0.0833333333333333*S**2 + 0.333333333333333*S*T*x3**2 - 0.25*S*T*x3 - 0.0833333333333333*S*T + 0.166666666666667*T**2*x3**2 - 0.166666666666667*T**2*x3) + equilibrium_UG*(-0.333333333333333*P**2*x3**2 + 0.333333333333333*P**2*x3 + 0.166666666666667*P*S*x3**2 - 0.25*P*S*x3 + 0.0833333333333333*P*S + 0.166666666666667*P*T*x3**2 - 0.166666666666667*P*T*x3) + equilibrium_Ug**4*(-0.166666666666667*P**2*x3**2 + 0.0833333333333333*P**2*x3 - 0.333333333333333*P*R*x3**2 + 0.25*P*R*x3 + 0.333333333333333*P*S*x3**2 - 0.25*P*S*x3 + 0.333333333333333*P*T*x3**2 - 0.166666666666667*P*T*x3 - 0.166666666666667*R**2*x3**2 + 0.166666666666667*R**2*x3 + 0.333333333333333*R*S*x3**2 - 0.333333333333333*R*S*x3 + 0.333333333333333*R*T*x3**2 - 0.25*R*T*x3 - 0.166666666666667*S**2*x3**2 + 0.166666666666667*S**2*x3 - 0.333333333333333*S*T*x3**2 + 0.25*S*T*x3 - 0.166666666666667*T**2*x3**2 + 0.0833333333333333*T**2*x3) + equilibrium_Ug**3*(0.666666666666667*P**2*x3**2 - 0.333333333333333*P**2*x3 + 0.666666666666667*P*R*x3**2 - 0.5*P*R*x3 - 1.0*P*S*x3**2 + 0.75*P*S*x3 - 1.0*P*T*x3**2 + 0.5*P*T*x3 - 0.333333333333333*R*S*x3**2 + 0.333333333333333*R*S*x3 - 0.333333333333333*R*T*x3**2 + 0.25*R*T*x3 + 0.333333333333333*S**2*x3**2 - 0.333333333333333*S**2*x3 + 0.666666666666667*S*T*x3**2 - 0.5*S*T*x3 + 0.333333333333333*T**2*x3**2 - 0.166666666666667*T**2*x3) + equilibrium_Ug**2*(-0.833333333333333*P**2*x3**2 + 0.333333333333333*P**2*x3 - 0.166666666666667*P*R*x3**2 + 0.25*P*R*x3 - 0.166666666666667*P*R + 0.833333333333333*P*S*x3**2 - 0.75*P*S*x3 + 0.166666666666667*P*S + 0.833333333333333*P*T*x3**2 - 0.333333333333333*P*T*x3 - 0.166666666666667*S**2*x3**2 + 0.166666666666667*S**2*x3 - 0.333333333333333*S*T*x3**2 + 0.25*S*T*x3 - 0.166666666666667*T**2*x3**2 + 0.0833333333333333*T**2*x3) + equilibrium_Ug*(0.333333333333333*P**2*x3**2 - 0.166666666666667*P*S*x3**2 + 0.25*P*S*x3 - 0.166666666666667*P*S - 0.166666666666667*P*T*x3**2))/((-x3 + 1)*(0.5*P + equilibrium_Ug**2*(0.5*P + 0.5*R - 0.5*S - 0.5*T) + equilibrium_Ug*(-1.0*P + 0.5*S + 0.5*T))**2)

```python
sym.factor(equilibrium_F_jac[2,2], equilibrium_UG, equilibrium_Ug)
```

    -0.75*(0.166666666666667*P**2 - 0.166666666666667*P*S*equilibrium_Ug_prime + equilibrium_Ug**4*(0.166666666666667*P**2 + 0.333333333333333*P*R - 0.333333333333333*P*S - 0.333333333333333*P*T + 0.166666666666667*R**2 - 0.333333333333333*R*S - 0.333333333333333*R*T + 0.166666666666667*S**2 + 0.333333333333333*S*T + 0.166666666666667*T**2) + equilibrium_Ug**3*(-0.666666666666667*P**2 - 0.666666666666667*P*R + 1.0*P*S + 1.0*P*T + 0.333333333333333*R*S + 0.333333333333333*R*T - 0.333333333333333*S**2 - 0.666666666666667*S*T - 0.333333333333333*T**2) + equilibrium_Ug**2*(1.0*P**2 + 0.333333333333333*P*R*equilibrium_Ug_prime + 0.333333333333333*P*R - 0.166666666666667*P*S*equilibrium_Ug_prime - 1.0*P*S - 1.0*P*T - 0.166666666666667*R*T*equilibrium_Ug_prime + 0.166666666666667*S**2 + 0.333333333333333*S*T + 0.166666666666667*T**2) + equilibrium_Ug*(-0.666666666666667*P**2 - 0.333333333333333*P*R*equilibrium_Ug_prime + 0.333333333333333*P*S*equilibrium_Ug_prime + 0.333333333333333*P*S + 0.333333333333333*P*T))/(0.5*P + equilibrium_Ug**2*(0.5*P + 0.5*R - 0.5*S - 0.5*T) + equilibrium_Ug*(-1.0*P + 0.5*S + 0.5*T))**2

```python
charpoly = equilibrium_F_jac.charpoly()
```

```python
eigenvals = equilibrium_F_jac.eigenvals()
```

```python
e1, e2, e3 = eigenvals.keys()
```

```python
numerator, denominator = sym.fraction(sym.factor(e1, equilibrium_UG, equilibrium_Ug))
```

```python
# always positive!
denominator
```

    2*(P + equilibrium_Ug**2*(P + R - S - T) + equilibrium_Ug*(-2*P + S + T))**2

```python
# yields condition on Ug'
sym.factor(numerator, equilibrium_Ug_prime)
```

    -P**2*equilibrium_Ug**4 + 4*P**2*equilibrium_Ug**3 - 6*P**2*equilibrium_Ug**2 + 4*P**2*equilibrium_Ug - P**2 - 2*P*R*equilibrium_Ug**4 + 4*P*R*equilibrium_Ug**3 - 2*P*R*equilibrium_Ug**2 + 2*P*S*equilibrium_Ug**4 - 6*P*S*equilibrium_Ug**3 + 6*P*S*equilibrium_Ug**2 - 2*P*S*equilibrium_Ug + 2*P*T*equilibrium_Ug**4 - 6*P*T*equilibrium_Ug**3 + 6*P*T*equilibrium_Ug**2 - 2*P*T*equilibrium_Ug - R**2*equilibrium_Ug**4 + 2*R*S*equilibrium_Ug**4 - 2*R*S*equilibrium_Ug**3 + 2*R*T*equilibrium_Ug**4 - 2*R*T*equilibrium_Ug**3 - S**2*equilibrium_Ug**4 + 2*S**2*equilibrium_Ug**3 - S**2*equilibrium_Ug**2 - 2*S*T*equilibrium_Ug**4 + 4*S*T*equilibrium_Ug**3 - 2*S*T*equilibrium_Ug**2 - T**2*equilibrium_Ug**4 + 2*T**2*equilibrium_Ug**3 - T**2*equilibrium_Ug**2 - equilibrium_Ug_prime*(2*P*R*equilibrium_Ug**2 - 2*P*R*equilibrium_Ug - P*S*equilibrium_Ug**2 + 2*P*S*equilibrium_Ug - P*S - R*T*equilibrium_Ug**2)

```python
sym.cancel(e2.subs({x3: equilibrium_x3}))
```

```python
numerator, denominator = sym.fraction(sym.simplify(e2.subs({x3: equilibrium_x3})))
```

```python
sym.factor(denominator, equilibrium_Ug)
```

    -16.0*(-1.0*P + equilibrium_Ug*(1.0*P - 1.0*T))*(S + equilibrium_Ug*(R - S))*(1.0*equilibrium_Ug - 1.0)*(0.5*P + equilibrium_Ug**2*(0.5*P + 0.5*R - 0.5*S - 0.5*T) + equilibrium_Ug*(-1.0*P + 0.5*S + 0.5*T))

$$ -4(P + U_g^*(T - P))(S + U_g^*(R - S))(1 - U_g^*)\big(2(P + U_g^{*2}(P + R - S - T) + U_g^*(-P + S + T))\big) $$

```python
sym.factor(sym.expand(numerator), equilibrium_Ug, equilibrium_UG)
```

```python
.subs({x3: equilibrium_x3})
```

# Emergence of cooperation

Start in an equilibrium with mixed $\gamma$ and pure $a$ (i.e., mix of $Ga$ and $ga$ genotypes). Necessary and sufficient conditions for stability of such as equilibrium is as follows.

$$ U'_{GA}(0) < \frac{P}{S} $$

Furthermore suppose that the necessary and sufficient condition for stability of a mixed $\gamma$ pure $A$ (i.e., a mix of $GA$ and $gA$ genotypes) is also satisfied.

$$ U'_{GA}(1) < \frac{R}{T} $$

Cooperative behavior is said to emerge spontaneously if, starting from a stable equilibrium with mixed $\gamma$ and pure $a$, a negative (positive) shock to $P$ ($S$ or $d$) results in the stability condition for the equilibrium to be violated and results in the system converging to a mixed $\gamma$ pure $A$ equilibrium.

### Example: Kirkpatrick Selection

Suppose that we take our selection function from Kirkpatrick paper.

$$ U_{GA} = \frac{d x_A}{1 + (d-1) x_A} $$

Derivative of this selection function is

$$ U'_{GA} = \frac{d}{\big(1 + (d-1)x_A\big)^2}$$

which yields $U'_{GA}(0)=1$ and $U'_{GA}(1)=\frac{1}{d}$.

Assuming payoffs are from a Prisoner's Dilemma we see that the stability condition for the mixed $\gamma$ and pure $a$ equilibrium will always hold and thus spontaneous emergence of cooperation is not possible.

### Example: Seger Selection

Suppose that we take our selection function from the Seger paper.

$$ U_{GA} = x_A\big(1 + d(1 - x_A)\big) $$

$$ U'_{GA} = 1 + d - 2dx_A$$

which yields $U'_{GA}(0)=1 + d$ and $U'_{GA}(1)=1-d$.  Substituting the results into the stability conditions yields

$$ 1 - \frac{R}{T} < d < \frac{P}{S}-1 $$

is sufficient to guarantee the stability of both equilibria.

```python
def seger_selection(x_A, d=0):
    return x_A * (1 + d * (1 - x_A))
```

```python
def plot_generalized_sexual_selection(x0=0.25, x1=0.25, x2=0.25, d0=2, d2=1, T=15, R=6, P=3, S=1, max_time=10):

    fig, ax = plt.subplots(1,1, figsize=(10,8))

    # prepare the axes
    ax.set_ylim((0, 1))
    ax.set_xlabel(r"Time, $t$", fontsize=15)
    ax.set_ylabel(r"Female genotype shares, $x_i$", fontsize=15)

    # create the initial condition
    x3 = 1 - (x0 + x1 + x2)
    y0=np.array([x0,x1,x2,x3])
    assert np.allclose(y0.sum(), 1)

    # create the payoff kernel
    assert T > R > P > S
    payoff_kernel = np.array([[R, S],
                              [T, P]])

    # create the selection functions
    UGA = lambda x_A: seger_selection(x_A, d0)
    UgA = lambda x_A: seger_selection(x_A, d2)

    # simulate the model starting from a random initial condition
    f = lambda t, y: generalized_sexual_selection(y, UGA, UgA, payoff_kernel)
    result = integrate.solve_ivp(f, t_span=(0, max_time), y0=y0, rtol=1e-7, atol=1e-9,
                                 dense_output=True, vectorized=True)

    ax.plot(result.t, result.y[0], label="GA")
    ax.plot(result.t, result.y[1], label="Ga")
    ax.plot(result.t, result.y[2], label="gA")
    ax.plot(result.t, result.y[3], label="ga")
    ax.legend()
    plt.show()

    return result
```


# Mixture Density Networks in GPflow

In this notebook we explain how to implement a Mixture Density Network (MDN) [1] using GPflow. In theory, this is similar to [this blog post](http://blog.otoro.net/2015/11/24/mixture-density-networks-with-tensorflow/) from 2015, but instead of using TensorFlow directly we'll use GPflow. GPflow is typically used for building Gaussian Process-based models, but the framework contains many useful methods and classes that can be used to quickly prototype a wide variety of ML algorithms. Excellent for doing research!

We start by explaining why MDNs can be useful. We then examine a GPflow implementation of the model and use it for a couple of toy experiments.

## Conditional Density Estimation models
Imagine we are interested in performing regression on the following dataset.

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

%matplotlib inline

np.random.seed(1)  # for reproducibility of this notebook

CMAP = plt.get_cmap("Blues")
```

```python
N = 200
NOISE_STD = 5.0e-2

def sinusoidal_data(N, noise):
    Y = np.linspace(-2, 2, N)[:, None]
    X = np.sin(4 * Y) * 2.0 + Y * 0.5
    X += np.random.randn(N, 1) * noise
    Y += np.random.randn(N, 1) * noise
    return X, Y

X, Y = data = sinusoidal_data(N, NOISE_STD)

plt.plot(X, Y, "ro", alpha=0.3)
plt.xlabel("$x$")
_ = plt.ylabel("$y$")
```

At first sight, this dataset doesn't seem overly complex. Both input and output have a single dimension, and the data has a clear sinusoidal pattern. However, notice that a single input $x$ can correspond to multiple output values $y$, so for example $x=0$ can yield any of the values $\{-1.5, -3/4, 0, 0.8, 1.5\}$. Typical regression algorithms such as Linear Regression, Gaussian Process regression and Multilayer Perceptrons (MLPs) struggle as they can only predict one output value for every input.

To model this dataset we can use a Conditional Density Estimation (CDE) model. CDE models infer $p(f(x)|x)$ instead of just calculating the expectation $E[f(x) | x]$. Modeling the complete distribution $p(f(x)|x)$ is typically harder but it reveals more interesting properties, such as the modes, outlier boundaries and samples. A real-world example might be modeling taxi drop-offs, conditioned on the pick-up location. We would expect a taxi drop-off location to be multi-modal as passengers need to go to different destinations (airport/city center/suburbs and so on) and the density depends on the starting location [2].

## Mixture Density Network models

Mixture Density Networks (MDNs) are a parametric class of models that allow for conditional density estimation. They consist of two parts: a neural net and a Mixture of Gaussians (MoG). The neural net is responsible for producing the characteristics of the MoG. In practice, given that the MoG consists of $M$ Gaussians, the neural net will output a collection of $M$ means, variances and weights $\{\mu_m, \sigma_m^2, \pi_m\}_{m=1}^M$. These means, variances and weights are used to define the conditional probability distribution function:
\begin{equation}
p(Y = y\,|\,X = x) = \sum_{m=1}^{M} \pi_{m}(x)\,\mathcal{N}\big(y\, \left|\,\mu_{m}(x), \sigma_{m}^2(x)\big)\right.
\end{equation}

Each of the parameters  $\pi_{m}(x), \mu_{m}(x), \sigma_{m}(x)$ of the distribution are determined by the neural net, as a function of the input $x$.

We train the MDN's neural net by optimizing the model's likelihood:
\begin{equation}
\mathcal{L} \triangleq \text{argmax}_{\Theta} \prod_{n=1}^N p(Y = y_n | X = x_n)
\end{equation}

where $\Theta$ collects the neural net's weights and biases and $\{x_n, y_n\}_{n=1}^N$ represents our training dataset.

## A GPflow MDN implementation

GPflow doesn't reinvent the wheel; most of what follows is just plain Python/TensorFlow code. We choose to use GPflow, however, because it provides us with functionality to easily define a model. Once we have a GPflow model, we can specify its objective function, parameters and dataset. This extra layer of abstraction makes interacting with the model much easier, for example optimizing or performing inference.

We begin by importing the required packages from TensorFlow and GPflow.

```python
import tensorflow as tf
```

```python
import gpflow
from gpflow.models import BayesianModel, ExternalDataTrainingLossMixin
from gpflow.base import Parameter
```

Next, we create a `MDN` class that inherits from GPflow's `Model` class. We need to do the following:
1. Store each of the feature and target matrices (X, Y) as a `DataHolder` object.
2. Define our model's parameters using GPflow's `Parameter` and `ParamList` objects.
3. Define the objective function using the `_build_likelihood` method. When we optimize the model the negative of this function will be minimized.

```python
from typing import Callable, Optional, Tuple

class MDN(BayesianModel, ExternalDataTrainingLossMixin):
    def __init__(
        self,
        num_mixtures: Optional[int] = 5,
        inner_dims: Optional[list] = [10, 10,],
        activation: Optional[Callable[[tf.Tensor], tf.Tensor]] = tf.keras.activations.relu,
    ):
        super().__init__()

        # `self.dims` collects the neural net's input, hidden and output dimensions.
        # The number of output dims `self.dims[-1]` equals `num_mixtures` means +
        # `num _mixtures` variances + `num_mixtures` weights, a total of
        # 3 times `num_mixtures` variables.
        self.dims = [1,] + list(inner_dims) + [3 * num_mixtures]
        self.activation = activation
        self._create_network()

    def _create_network(self):
        self.Ws, self.bs = [], []

        for dim_in, dim_out in zip(self.dims[:-1], self.dims[1:]):
            init_xavier_std = (2.0 / (dim_in + dim_out)) ** 0.5
            self.Ws.append(Parameter(np.random.randn(dim_in, dim_out) * init_xavier_std))
            self.bs.append(Parameter(np.zeros(dim_out)))

    def eval_network(self, X):
        for i, (W, b) in enumerate(zip(self.Ws, self.bs)):
            X = tf.matmul(X, W) + b
            if i < len(self.bs) - 1:
                X = self.activation(X)

        pis, mus, sigmas = tf.split(X, 3, axis=1)
        pis = tf.nn.softmax(pis)  # make sure they normalize to 1
        sigmas = tf.exp(sigmas)  # make sure std. dev. are positive

        return pis, mus, sigmas

    def maximum_log_likelihood_objective(self, data: Tuple[tf.Tensor, tf.Tensor]):
        x, y = data
        pis, mus, sigmas = self.eval_network(x)
        Z = (2 * np.pi) ** 0.5 * sigmas
        log_probs_mog = (-0.5 * (mus - y) ** 2 / sigmas ** 2) - tf.math.log(Z) + tf.math.log(pis)
        log_probs = tf.reduce_logsumexp(log_probs_mog, axis=1)
        return tf.reduce_sum(log_probs)
```

### Notes
- Given we are dealing with a MoG, the neural net output must comply with the following restrictions:
  \begin{equation}
  \sum_{m=1}^{M} \pi_{m}(x) = 1, \pi_m \ge 0\ \text{and}\ \sigma_m\ \forall\ m
  \end{equation}
  We achieve this by applying the `softmax` operator to the $\pi$'s and by taking the `exp` to the $\sigma$'s.

- We use the "Xavier" initialization for the neural net's weights. (Glorot and Bengio, 2010).

- Instead of calculating the pdf of the Gaussians, we work with the pdf `log` and use `tf.reduce_logsumexp`. This is mainly for numerical stability.

## Experiment 1: The sinusoidal dataset

Let's see how our model works in practice with the sinusoidal dataset presented earlier. We do this by initializing a new instance of our MDN model, and then specifying the dataset $(X, Y)$, the number of hidden units of the MDN's neural net, and the number of mixture components $M$.

```python
model = MDN(inner_dims=[100, 100], num_mixtures=5)

from gpflow.utilities import print_summary

print_summary(model)
```

    ╒═══════════╤═══════════╤═════════════╤═════════╤═════════════╤════════════╤═════════╤═══════════════════════════════════════════╕
    │ name      │ class     │ transform   │ prior   │ trainable   │ shape      │ dtype   │ value                                     │
    ╞═══════════╪═══════════╪═════════════╪═════════╪═════════════╪════════════╪═════════╪═══════════════════════════════════════════╡
    │ MDN.Ws[0] │ Parameter │             │         │ True        │ (1, 100)   │ float64 │ [[-0.18385483, 0.01074822, 0.05167668...  │
    ├───────────┼───────────┼─────────────┼─────────┼─────────────┼────────────┼─────────┼───────────────────────────────────────────┤
    │ MDN.Ws[1] │ Parameter │             │         │ True        │ (100, 100) │ float64 │ [[-0.17193945, 0.0057121, -0.07995475...  │
    ├───────────┼───────────┼─────────────┼─────────┼─────────────┼────────────┼─────────┼───────────────────────────────────────────┤
    │ MDN.Ws[2] │ Parameter │             │         │ True        │ (100, 15)  │ float64 │ [[-0.08246711, -0.37777674, 0.09595574... │
    ├───────────┼───────────┼─────────────┼─────────┼─────────────┼────────────┼─────────┼───────────────────────────────────────────┤
    │ MDN.bs[0] │ Parameter │             │         │ True        │ (100,)     │ float64 │ [0., 0., 0....                            │
    ├───────────┼───────────┼─────────────┼─────────┼─────────────┼────────────┼─────────┼───────────────────────────────────────────┤
    │ MDN.bs[1] │ Parameter │             │         │ True        │ (100,)     │ float64 │ [0., 0., 0....                            │
    ├───────────┼───────────┼─────────────┼─────────┼─────────────┼────────────┼─────────┼───────────────────────────────────────────┤
    │ MDN.bs[2] │ Parameter │             │         │ True        │ (15,)      │ float64 │ [0., 0., 0....                            │
    ╘═══════════╧═══════════╧═════════════╧═════════╧═════════════╧════════════╧═════════╧═══════════════════════════════════════════╛

The objective function for MDN instances is the `maximum_log_likelihood_objective`, which we use for optimization of the parameters. GPflow ensures that only the variables stored in `Parameter` objects are optimized. For the MDN, the only parameters are the weights and the biases of the neural net.

We use the `Scipy` optimizer, which is a wrapper around SciPy's L-BFGS optimization algorithm. Note that GPflow supports other TensorFlow optimizers such as `Adam`, `Adagrad`, and `Adadelta` as well.

```python
from gpflow.optimizers import Scipy
from gpflow.ci_utils import ci_niter

Scipy().minimize(
    model.training_loss_closure(data, compile=True),
    model.trainable_variables,
    options=dict(maxiter=ci_niter(1500)),
)

print("Final Likelihood", model.maximum_log_likelihood_objective(data).numpy())
```

    Final Likelihood -139.11678932818785

To evaluate the validity of our model, we draw the posterior density. We also plot $\mu(x)$ of the optimized neural net. Remember that for every $x$ the neural net outputs $M$ means $\mu_m(x)$. These determine the location of the Gaussians. We plot all $M$ means and use their corresponding mixture weight $\pi_m(X)$ to determine their size. Larger dots will have more impact in the Gaussian ensemble.

```python
try:
    from mdn_plotting import plot
except:
    # VS CODE's root directory is GPflow's top-level directory
    from doc.source.notebooks.tailor.mdn_plotting import plot

fig, axes = plt.subplots(1, 2, figsize=(12, 6))
for a in axes:
    a.set_xlim(-4, 4)
    a.set_ylim(-3, 3)
plot(model, X, Y, axes, cmap=CMAP)
```

## Experiment 2: The half moon dataset

The half moon dataset is available in the `scikit-learn` package.

```python
from sklearn.datasets import make_moons

def moon_data(N, noise):
    data, _ = make_moons(n_samples=N, shuffle=True, noise=noise)
    X, Y = data[:, 0].reshape(-1, 1), data[:, 1].reshape(-1, 1)
    return X, Y
```

```python
X, Y = data = moon_data(N, NOISE_STD)
plt.plot(X, Y, "ro", alpha=0.3)
plt.xlabel("$x$")
_ = plt.ylabel("$y$")
```

The only difference in the MDN's setup is that we lower the number of mixture components.

```python
model = MDN(inner_dims=[100, 100], num_mixtures=5)
```

```python
Scipy().minimize(
    model.training_loss_closure(data, compile=True),
    model.trainable_variables,
    options=dict(maxiter=ci_niter(int(10e3))),
)

print("Final Likelihood", model.maximum_log_likelihood_objective(data).numpy())
```

    Final Likelihood 203.3477853001398

```python
fig, axes = plt.subplots(1, 2, figsize=(12, 6))
for a in axes:
    a.set_xlim(-2, 3)
    a.set_ylim(-1.5, 2)
plot(model, X, Y, axes, cmap=CMAP)
```

## References

[1] Bishop, Christopher M. Mixture density networks. Technical Report NCRG/4288, Aston University, Birmingham, UK, 1994.

[2] Dutordoir, Vincent, et al. "Gaussian Process Conditional Density Estimation." Advances in Neural Information Processing Systems. 2018.


# 2. laboratorijska vježba

```python
# učitvanje potrebnih biblioteka

import numpy as np
import matplotlib.pyplot as plt
import scipy.signal as ss
```

```python
#@title pomoćna funkcija
# izvršite ovu ćeliju ali se ne opterećujte detaljima implementacije

def plot_frequency_response(f, Hm, fc=None, ylim_min=None):
    """Grafički prikaz prijenosne funkcije filtra.

    Args
        f (numpy.ndarray) : frekvencije
        Hm (numpy.ndarray) : apsolutne vrijednosti prijenosne funkcije
        fc (number) : cutoff frekvencija
        ylim_min (number): minimalna vrijednost na y-osi za dB skalu

    Returns
        (matplotlib.figure.Figure, matplotlib.axes._subplots.AxesSubplot)
    """
    Hc = 1 / np.sqrt(2)
    if fc is None:
        fc_idx = np.where(np.isclose(Hm, Hc, rtol=1e-03))[0][0]
        fc = f[fc_idx]
    H_db = 20 * np.log10(Hm)

    fig, ax = plt.subplots(nrows=2, ncols=2, figsize=(12, 7.5))

    ax[0, 0].plot(f, Hm, label='$H(f)$')
    ax[0, 0].plot(fc, Hc, 'o', label='$H(f_c)$')
    ax[0, 0].vlines(fc, Hm.min(), Hc, linestyle='--')
    ax[0, 0].annotate(f'$f_c = {fc:.3f}$ Hz\n$H(f_c)={Hc:.3f}$', (fc * 1.4, Hc))
    ax[0, 0].set_xscale('log')
    ax[0, 0].set_ylabel('$|V_{out}$ / $V_{in}$|')
    ax[0, 0].set_title('log scale')
    ax[0, 0].legend(loc='lower left')
    ax[0, 0].grid()

    ax[0, 1].plot(f, Hm, label='$H(f)$')
    ax[0, 1].plot(fc, Hc, 'o', label='$H(f_c)$')
    ax[0, 1].annotate(f'$f_c = {fc:.3f}$ Hz\n$H(f_c)={Hc:.3f}$', (fc * 1.4, Hc))
    ax[0, 1].set_title('linear scale')
    ax[0, 1].legend()
    ax[0, 1].grid()

    ax[1, 0].plot(f, H_db, label='$H_{dB}(f)$')
    ax[1, 0].plot(fc, H_db.max() - 3, 'o', label='$H_{dB}(f_c)$')
    ax[1, 0].vlines(fc, H_db.min(), H_db.max() - 3, linestyle='--')
    ax[1, 0].annotate(f'$f_c = {fc:.3f}$ Hz\n$H(f_c)={H_db.max() - 3:.3f} dB$',
                      (fc * 1.4, H_db.max() - 3))
    ax[1, 0].set_xscale('log')
    ax[1, 0].set_xlabel('$f$ [Hz]')
    ax[1, 0].set_ylabel('$20 \\cdot \\log$ |$V_{out}$ / $V_{in}$|')
    if ylim_min:
        ax[1, 0].set_ylim((ylim_min, 10))
    ax[1, 0].legend(loc='lower left')
    ax[1, 0].grid()

    ax[1, 1].plot(f, H_db, label='$H_{dB}(f)$')
    ax[1, 1].plot(fc, H_db.max() - 3, 'o', label='$H_{dB}(f_c)$')
    ax[1, 1].annotate(f'$f_c = {fc:.3f}$ Hz\n$H(f_c)={H_db.max() - 3:.3f} dB$',
                      (fc * 1.4, H_db.max() - 3))
    ax[1, 1].set_xlabel('$f$ [Hz]')
    if ylim_min:
        ax[1, 1].set_ylim((ylim_min, 10))
    ax[1, 1].legend()
    ax[1, 1].grid()

    fig.tight_layout
    return fig, ax
```

## Električni filtri

Općenito, filtar u domeni elektromagnetske kompatibilnosti, kao i kroz ovaj kolegij, promatramo kao model električnog kruga definiranog koristeći koncentrirane parametre. Signal s ulaza se oblikuje na način da njegove komponente pri određenim frekvencijama filtar propušta bez izobličenja (u teoriji), dok na drugim frekvencijama komponente uopće ne propušta (u teoriji, u stvarnosti dolazi do jače ili slabije atenuacije). Selektivno propuštanje signala se radi ponajprije kako bi se potisnule ili u potpunosti otklonile neželjene komponente signala.

S obzirom na frekvencijski pojas propusta, filtre dijelimo u osnovne kategorije:
1. nisko-propusni filtri - propuštaju se sve frekvencije signala ispod tzv. *cutoff* frekvencije;
2. visoko-propusni filtri - propuštaju se sve frekvencije signala iznad *cutoff* frekvencije;
3. pojasno-propusni filtri - propušta se definirani frekvencijski opseg;
4. pojasno-nepropusni filtri - ne propušta se određeni frekvencijski opseg.

S obzirom na tip koncentriranih parametara, odnosno komponenti koje koristimo kako bismo računalno modelirali rad filtra, filtre dijelimo na dodatne kategorije:
1. aktivni filtri - realizirani uz pomoć aktivnih komponenti (tranzistori i operacijska pojačala);
2. pasivni filtri - realizirani uz pomoć pasivnih komponenti (otpornici, kondenzatori i zavojnice).

### Pasivni nisko-propusni filtar

Realizacija nisko-propusnog filtra u ovim laboratorijskim vježbama se ostvaruje korištenjem otpornika i kondenzatora povezanih u seriju, pri čemu se izlaz promatra kao napon na kondenzatoru, $V_{out}$. Uz pretpostavku da je signal na ulazu, $V_{in}$, sinusoidalni naponski izvor, analizu možemo prebaciti u frekvencijsku domenu koristeći impedancijski model. Na ovaj način zaobilazimo potrebu korištenja diferencijalnog i čitav problem svodimo na jednostavni algebarski proračun.

Napon na kondenzatoru, $V_{out}$, definiramo kroz podjelu ulaznog napona:

$$
\begin{align}
    V_{out} &= \frac{Z_c}{Z_c + Z_r} \cdot V_{in} \\
    \frac{V_{out}}{V_{in}} &= \frac{Z_c}{Z_c + Z_r} = \frac{1/(j\omega C)}{1/(j\omega C) + R} = \frac{1}{1 + j\omega C R}
\end{align}
$$

Dobiveni odnos izlaznog i ulaznog napona se naziva **funkcija prijenosnog odziva**, $H$, u stacionarnom stanju:

$$
H(\omega) = \frac{V_{out}}{V_{in}} = \frac{1}{1 + j\omega C R}
$$

Kako je $H$ funkcija frekvencije, imamo dva rubna slučaja:
* za iznimno niske frekvencije kada je $\omega \sim 0$ slijedi da je $H(\omega) = 1$;
* za iznimno visoke frekvencije kada $\omega \rightarrow \infty$ slijedi da je $H(\omega) \rightarrow 0$.

Potrebno je dodatno definirati već spomenutu *cutoff* frekvenciju, $f_c$, za koju amplituda funkcije frekvencijskog odziva, $H$, pada za $\sqrt 2$ puta, odnosno za $3$ dB:

$$
\begin{align}
    \frac{H(\omega)}{\sqrt 2} &= \frac{1}{1 + j\omega_c C R} \\
    \omega_c &= \frac{1}{CR} \\
    f_c &= \frac{1}{2\pi CR}
\end{align}
$$

Link za interaktivni rad s pasivnim nisko-propusnim filtrom: http://sim.okawa-denshi.jp/en/CRtool.php

#### Zadatak 1

Prvi zadatak je implementirati funkciju `cutoff_frequency` koja na ulazu prima iznose otpora, `R`, i kapaciteta, `C`, a na izlazu daje *cutoff* frekvenciju nisko-propusnog filtra.

```python
def cutoff_frequency(R, C):
    """Cutoff frekvencija nisko-propusnog filtra.

    Args:
        R (number) : vrijednost otpora otpornika
        C (number) : kapacitet kondenzatora

    Returns:
        number
    """
    #######################################################
    ## TO-DO: implementiraj proračun cutoff frekvencije ##
    # Nakon toga zakomentiraj sljedeću liniju.
    raise NotImplementedError('Implementiraj proračun cutoff frekvencije.')
    #######################################################

    # definiraj cutoff frekvenciju
    fc = ...
    return fc
```

Kolika je *cutoff* frekvencija za otpor od $4.5 k \Omega$ i kapacitet kondenzatora od $7 \mu F$?

```python
R = ...  # otpor u Ohm
C = ...  # kapacitet u F

fc = ...  # cutoff frekvencija

print(f'R = {R / 1000:.2f} kΩ')
print(f'C = {C * 10**6:.2f} µF')
print(f'cutoff frekvencija iznosi {fc:.2f} Hz, '
      'očekivana vrijednost je 5.05 Hz')
```

#### Zadatak 2

Drugi zadatak je implementirati funkciju `rc_lowpass` koja na ulazu prima iznose otpora, `R`, kapaciteta, `C`, i frekvenciju, `f`, a na izlazu daje prijenosni odziv nisko-propusnog filtra.

```python
def rc_lowpass(R, C, f):
    """Funkcija prijenosnog odziva RC nisko-propusnog filtra.

    Args:
        R (number) : vrijednost otpora otpornika
        C (number) : kapacitet kondenzatora
        f (number or numpy.ndarray) : frekvencija/e

    Returns:
        float or numpy.ndarray
    """
    ######################################################
    ## TO-DO: implementiraj funkciju prijenosnog odziva ##
    # Nakon toga zakomentiraj sljedeću liniju.
    raise NotImplementedError('Implementiraj funckiju prijenosnog odziva.')
    ######################################################

    # definiraj funkciju prijenosnog pazeći da `f` može biti ili broj (int,
    # float) ili 1-D niz (`numpy.ndarray`)
    H = ...
    return H
```

Kolika je vrijednost prijenosne funkcije pri *cutoff* frekvencija za otpor od $4.5 k \Omega$ i kapacitet kondenzatora od $7 \mu F$?

```python
R = ...  # otpor u Ohm
C = ...  # kapacitet u F

Hc = ...  # prijenosna funkcija pri cutoff frekvenciji

print(f'R = {R / 1000:.2f} kΩ')
print(f'C = {C * 10**6:.2f} µF')
print(f'pojačanje pri cutoff frekvenciji iznosi {abs(Hc):.4f}, '
      'očekivana vrijednost je 1/√2\n\n'
      'provjerite ispravnost dobivenog rezutltata')
```

```python
# ćelija za provjeru rezultata

```

Pretvorite vrijednost prijenosne funkcije pri *cutoff* frekvenciju u decibele i uvjerite se u tvrdnju da amplituda funkcije frekvencijskog odziva, $H$, pada za $3$ dB pri *cutoff* frekvenciji.

```python
Hc_dB = ...  # pretvorba prijenosne funkcije pri cutoff frekvenciji u dB skalu
print(Hc_dB)
```

Za raspon od $10000$ vrijednosti frekvencija do $100 Hz$ te za otpor od $4.5 k\Omega$ i kapacitet kondenzatora od $7 \mu F$, izračunajte vrijednosti prijenosne funkcije.

```python
f = ...  # raspon frekvencija od 0 do 100 Hz

H = rc_lowpass(R, C, f)  # prijenosna funkcija
```

S obzirom da su vrijednosti prijenosne funkcije kompleksne veličine, razmilite što je potrebno napraviti s njima prije nego ih grafički prikažemo?

```python
Hm = ...  # konverzija u apsolutne vrijednosti
```

Vizualizirajte ovisnost prijenosne funkcije o frekvenciji koristeći `matplotlib` i funkciju `matplotlib.pyplot.plot`.

```python
plt.plot(...)  # prvi argument su vrijednosti na x-osi, drugi argument vrijednosti na y-osi
plt.xlabel('f')
plt.ylabel('H(f)')
plt.show()
```

Vizualizirajte sada rezultate koristeći već implementiranu funkciju `plot_frequency_response`.

Napomena: za provjeru načina korištenja prethodne funkcije koristite sljedeću naredbu:

```python
help(plot_frequency_response)
```

ili jednostavno

```python
plot_frequency_response?
```

```python
# provjerite način korištenja funkcije

```

```python
fig, ax = plot_frequency_response(...)  # grafički prikaz dobivenih rezultata
```

### Butterworthov filtar

Butterworthov filtar vrsta je filtra za obradu signala dizajnirana tako da je frekvencijski odziv što je moguće ravniji u propusnom opsegu.

Funkcija prijenosnog odziva, $H$, niskopropusnog Butterworth filtra trećeg reda (Cauer, ljestvasta topologija), prikazanog na slici poviše, ima sljedeći oblik:

$$H(s) = \frac{V_{out}(s)}{V_{in}(s)} = \frac{R_4}{s^3(L_1 C_2 L_3) + s^2(L_1 C_2 R_4) + s(L_1 + L_3) + R_4}$$

Ukoliko zadržimo istu *cutoff* frekvenciju kao i u prvom primjeru, funkcija Butterworthovog filtra će, kao glavnu prednost, imati puno strmiju granicu propusnog i nepropusnog frekvencijskog područja (ovisno o redu filtra).

```python
f = np.linspace(0, 100, 10000)  # raspon frekvencija od 0 do 100 Hz
fc = cutoff_frequency(R, C)  # cutoff frekvencija
fs = 1000  # frekvencija uzorkovanja
fs_nyq = 0.5 * fs  # Nyquistova frekvencija
fc_norm = fc / fs_nyq  # normalizirana cutoff frekvencija
order = 3  # red filtra

b, a = ss.butter(order, fc_norm, btype='low')  # brojnik i nazivnik IIR filtra
f_norm, H = ss.freqz(b, a, worN=f.size)  # normalizirane frekvencije i prijenosna funkcija filtra
f_denorm = f_norm / np.pi * fs_nyq  # rekonstruirane frekvencije
Hm = abs(H)  # apsolutna vrijednost prijenosne funkcije filtra
```

```python
fig, ax = plot_frequency_response(f_denorm, Hm, fc, ylim_min=-80)  # grafički prikaz dobivenih rezultata
```

#### Zadatak 3

Implementiraj funkciju generatora sinusnog signala `signal_generator` koja prima vremensku seriju trenutaka u kojima želimo generirati signal, amplitudu signala, željenu frekvenciju signala, i razinu šuma koji se nadodaje signalu.

```python
def signal_generator(t, A, f, noise_factor):
    """Generator sinusnog signala.

    Args:
        t (number or numpy.ndarray) : vrijeme u kojem generiramo signal
        A (number) : amplituda signala
        f (number) : frekvencija signala
        noise_factor (number) : razina šuma

    Returns:
        tuple
    """
    ######################################################
    ## TO-DO: implementiraj funkciju generatora signala ##
    # Nakon toga zakomentiraj sljedeću liniju.
    raise NotImplementedError('Implementiraj funkciju generatora signala.')
    ######################################################

    # definiraj sinusni signal amplitude `A` za svaki `t`
    signal = ...

    # definiraj bijeli šum skaliran za `noise_factor`
    noise = ...

    # konačni signal = suma signala i šuma
    noisy_signal = ...
    return signal, noisy_signal
```

Generiraj sinusni signal i signal uz prisustvo šuma frekvencije $1 Hz$ u trajanju od $5$ sekunda, koristeći frekvenciju uzorkovanja od $30 Hz$. Amplituda signala neka bude $2$ a amplituda šuma $1$.

```python
T = ...  # ukupno trajanje signala
fs = ...  # frekvencija uzorkovanja
A = ...  # amplituda signala
f = ...  # frekvencija signala
noise_factor = ...  # amplituda bijelog šuma
t = np.linspace(0, T, T * fs)  # trenuci u kojima generiramo signal

signal, noisy_signal = signal_generator(t, A, f, noise_factor)
```

Grafički prikaži ovisnost signala i signala uz prisustvo šuma u vremenu.

```python
plt.plot(..., label='originalni signal')
plt.plot(..., label='mjereni signal uz prisustvo šuma')
plt.xlabel('t')
plt.ylabel('signal')
plt.legend()
plt.grid()
plt.show()
```

```python
def butter_filter(noisy_signal, fc, fs, order=3):
    """Butterworthov nisko-propusni filtar.

    Args:
        noisy_signal (numpy.ndarray) : signal uz prisustvo šuma
        fc (number) : cutoff frekvencija
        fs (number) : frekvencija uzorkovanja signala
        order (int) : red filtra

    Returns:
        numpy.ndarray
    """
    fs_nyq = 0.5 * fs
    fc_norm = fc / fs_nyq
    b, a = ss.butter(order, fc_norm, btype='low')
    filtered_signal = ss.filtfilt(b, a, noisy_signal)
    return filtered_signal
```

Filtriraj signal uz prisustvo šuma koristeći funkciju `butter_filter`. Mijenjaj *cutoff* frekvenciju dok ne dobiješ zadovoljavajuću rekonstrukciju signala. Počni s vrijednošću dobivenu za RC nisko-propusni filtar iz prvog zadatka.

```python
fc = ...  # proizvoljna cutoff frekvencija

filtered_signal = butter_filter(noisy_signal, fc, fs)
```

Grafički prikaz ovisnost signala i signala uz prisustvo šuma u vremenu.

```python
plt.plot(t, signal, label='originalni signal')
plt.plot(t, noisy_signal, label='mjereni signal uz prisustvo šuma')
plt.plot(t, filtered_signal, label='filtrirani signal')
plt.xlabel('t')
plt.ylabel('signal')
plt.legend()
plt.grid()
plt.show()
```

Za koju *cutoff* frekvenciju dobivate optimalnu rekonstrukciju signala? Komentirajte rezultate.

Zaključak:


###### Content under Creative Commons Attribution license CC-BY 4.0, code under BSD 3-Clause License © 2018  by D. Koehn, heterogeneous models are from [this Jupyter notebook](https://krischer.github.io/seismo_live_build/html/Computational%20Seismology/The%20Finite-Difference%20Method/fd_ac2d_heterogeneous_wrapper.html) by Heiner Igel ([@heinerigel](https://github.com/heinerigel)), Florian Wölfl and Lion Krischer ([@krischer](https://github.com/krischer)) which is a supplemenatry material to the book [Computational Seismology: A Practical Introduction](http://www.computational-seismology.org/), notebook style sheet by L.A. Barba, N.C. Clementi

```python
# Execute this cell to load the notebook's style sheet, then ignore it
from IPython.core.display import HTML
css_file = '../style/custom.css'
HTML(open(css_file, "r").read())
```

<link href="https://fonts.googleapis.com/css?family=Merriweather:300,300i,400,400i,700,700i,900,900i" rel='stylesheet' >
<link href="https://fonts.googleapis.com/css?family=Source+Sans+Pro:300,300i,400,400i,700,700i" rel='stylesheet' >
<link href='http://fonts.googleapis.com/css?family=Source+Code+Pro:300,400' rel='stylesheet' >
<style>

@font-face {
    font-family: "Computer Modern";
    src: url('http://mirrors.ctan.org/fonts/cm-unicode/fonts/otf/cmunss.otf');
}

#notebook_panel { /* main background */
    background: rgb(245,245,245);
}

div.cell { /* set cell width */
    width: 800px;
}

div #notebook { /* centre the content */
    background: #fff; /* white background for content */
    width: 1000px;
    margin: auto;
    padding-left: 0em;
}

#notebook li { /* More space between bullet points */
margin-top:0.5em;
}

/* draw border around running cells */
div.cell.border-box-sizing.code_cell.running {
    border: 1px solid #111;
}

/* Put a solid color box around each cell and its output, visually linking them*/
div.cell.code_cell {
    background-color: rgb(256,256,256);
    border-radius: 0px;
    padding: 0.5em;
    margin-left:1em;
    margin-top: 1em;
}

div.text_cell_render{
    font-family: 'Source Sans Pro', sans-serif;
    line-height: 140%;
    font-size: 110%;
    width:680px;
    margin-left:auto;
    margin-right:auto;
}

/* Formatting for header cells */
.text_cell_render h1 {
    font-family: 'Merriweather', serif;
    font-style:regular;
    font-weight: bold;
    font-size: 250%;
    line-height: 100%;
    color: #004065;
    margin-bottom: 1em;
    margin-top: 0.5em;
    display: block;
}
.text_cell_render h2 {
    font-family: 'Merriweather', serif;
    font-weight: bold;
    font-size: 180%;
    line-height: 100%;
    color: #0096d6;
    margin-bottom: 0.5em;
    margin-top: 0.5em;
    display: block;
}

.text_cell_render h3 {
    font-family: 'Merriweather', serif;
	font-size: 150%;
    margin-top:12px;
    margin-bottom: 3px;
    font-style: regular;
    color: #008367;
}

.text_cell_render h4 {    /*Use this for captions*/
    font-family: 'Merriweather', serif;
    font-weight: 300;
    font-size: 100%;
    line-height: 120%;
    text-align: left;
    width:500px;
    margin-top: 1em;
    margin-bottom: 2em;
    margin-left: 80pt;
    font-style: regular;
}

.text_cell_render h5 {  /*Use this for small titles*/
    font-family: 'Source Sans Pro', sans-serif;
    font-weight: regular;
    font-size: 130%;
    color: #e31937;
    font-style: italic;
    margin-bottom: .5em;
    margin-top: 1em;
    display: block;
}

.text_cell_render h6 { /*use this for copyright note*/
    font-family: 'Source Code Pro', sans-serif;
    font-weight: 300;
    font-size: 9pt;
    line-height: 100%;
    color: grey;
    margin-bottom: 1px;
    margin-top: 1px;
}

    .CodeMirror{
            font-family: "Source Code Pro";
			font-size: 90%;
    }
/*    .prompt{
        display: None;
    }*/

    .warning{
        color: rgb( 240, 20, 20 )
        }
</style>

# 2D acoustic finite-difference modelling

One approximation of the elastic equations of motion is the 2D acoustic wave equation. In this notebook we will discuss the finite difference approximation and model wave propagation for problems of different complexity. From a simple homogeneous half-space to complex subsurface models with significant spatial variations of the acoustic (P-wave) velocity.

## Finite difference solution of 2D acoustic wave equation

As derived in [this](https://nbviewer.jupyter.org/github/daniel-koehn/Theory-of-seismic-waves-II/blob/master/01_Analytical_solutions/3_Acoustic_medium.ipynb) and [this lecture](https://nbviewer.jupyter.org/github/daniel-koehn/Theory-of-seismic-waves-II/blob/master/01_Analytical_solutions/4_2D_1D_elastic_acoustic_approx.ipynb), the acoustic wave equation in 2D with constant density is

\begin{equation}
\frac{\partial^2 p(x,z,t)}{\partial t^2} \ = \ vp(x,z)^2 \biggl(\frac{\partial^2 p(x,z,t)}{\partial x^2}+\frac{\partial^2 p(x,z,t)}{\partial z^2}\biggr) + f(x,z,t) \nonumber
\end{equation}

with pressure $p$, acoustic velocity $vp$ and source term $f$. We can split the source term into a spatial and temporal part. Spatially, we assume that the source is localized at one point ($x_s, z_s$). Therefore, the spatial source contribution consists of two Dirac $\delta$-functions $\delta(x-x_s)$ and $\delta(z-z_s)$. The temporal source part is an arbitrary source wavelet $s(t)$:

\begin{equation}
\frac{\partial^2 p(x,z,t)}{\partial t^2} \ = \ vp(x,z)^2 \biggl(\frac{\partial^2 p(x,z,t)}{\partial x^2}+\frac{\partial^2 p(x,z,t)}{\partial z^2}\biggr) + \delta(x-x_s)\delta(z-z_s)s(t) \nonumber
\end{equation}

Both second derivatives can be approximated by a 3-point difference formula. For example for the time derivative, we get:

\begin{equation}
\frac{\partial^2 p(x,z,t)}{\partial t^2} \ \approx \ \frac{p(x,z,t+dt) - 2 p(x,z,t) + p(x,z,t-dt)}{dt^2}, \nonumber
\end{equation}

and similar for the spatial derivatives:

\begin{equation}
\frac{\partial^2 p(x,z,t)}{\partial x^2} \ \approx \ \frac{p(x+dx,z,t) - 2 p(x,z,t) + p(x-dx,z,t)}{dx^2}, \nonumber
\end{equation}

\begin{equation}
\frac{\partial^2 p(x,z,t)}{\partial x^2} \ \approx \ \frac{p(x,z+dz,t) - 2 p(x,z,t) + p(x,z-dz,t)}{dz^2}, \nonumber
\end{equation}

Injecting these approximations into the wave equation allows us to formulate the pressure p(x) for the time step $t+dt$ (the future) as a function of the pressure at time $t$ (now) and $t-dt$ (the past). This is called an **explicit time integration scheme** allowing the $extrapolation$  of the space-dependent field into the future only looking at the nearest neighbourhood.

In the next step, we discretize the P-wave velocity and pressure wavefield at the discrete spatial grid points

\begin{align}
x &= i*dx\nonumber\\
z &= j*dz\nonumber\\
\end{align}

with $i = 0, 1, 2, ..., nx$, $j = 0, 1, 2, ..., nz$ on a 2D Cartesian grid.

Using the discrete time steps

\begin{align}
t &= n*dt\nonumber
\end{align}

with $n = 0, 1, 2, ..., nt$ and time step $dt$, we can replace the time-dependent part (upper index time, lower indices space) by

\begin{equation}
 \frac{p_{i,j}^{n+1} - 2 p_{i,j}^n + p_{i,j}^{n-1}}{\mathrm{d}t^2} \ = \ vp_{i,j}^2 \biggl( \frac{\partial^2 p}{\partial x^2} + \frac{\partial^2 p}{\partial z^2}\biggr) \ + \frac{s_{i,j}^n}{dx\;dz}. \nonumber
\end{equation}

The spatial $\delta$-functions $\delta(x-x_s)$ and $\delta(z-z_s)$ in the source term are approximated by the boxcar function:

$$
\delta_{bc}(x) = \left\{
\begin{array}{ll}
	1/dx &|x|\leq dx/2 \\
	0 &\text{elsewhere}
\end{array}
\right.
$$

Solving for $p_{i,j}^{n+1}$ leads to the extrapolation scheme:

\begin{equation}
p_{i,j}^{n+1} \ = \ vp_{i,j}^2 \mathrm{d}t^2 \left( \frac{\partial^2 p}{\partial x^2} + \frac{\partial^2 p}{\partial z^2} \right) + 2p_{i,j}^n - p_{i,j}^{n-1} + \frac{\mathrm{d}t^2}{dx\; dz} s_{i,j}^n.
\end{equation}

The spatial derivatives are determined by

\begin{equation}
\frac{\partial^2 p(x,z,t)}{\partial x^2} \ \approx \ \frac{p_{i+1,j}^{n} - 2 p_{i,j}^n + p_{i-1,j}^{n}}{\mathrm{d}x^2} \nonumber
\end{equation}

and

\begin{equation}
\frac{\partial^2 p(x,z,t)}{\partial z^2} \ \approx \ \frac{p_{i,j+1}^{n} - 2 p_{i,j}^n + p_{i,j-1}^{n}}{\mathrm{d}z^2}. \nonumber
\end{equation}

Eq. (1) is the essential core of the 2D FD modelling code. Because we derived analytical solutions for wave propagation in a homogeneous medium, we should test our first code implementation for a similar medium, by setting

\begin{equation}
vp_{i,j} = vp0\notag
\end{equation}

at each spatial grid point $i = 0, 1, 2, ..., nx$; $j = 0, 1, 2, ..., nz$, in order to compare the numerical with the analytical solution. For a complete description of the problem we also have to define initial and boundary conditions. The **initial condition** is

\begin{equation}
p_{i,j}^0 = 0, \nonumber
\end{equation}

so the modelling starts with zero pressure amplitude at each spatial grid point $i, j$. As **boundary conditions**, we assume

\begin{align}
p_{i,0}^n = 0, \nonumber\\
\end{align}

for all time steps n at the top boundary. This **Dirichlet boundary condition**, leads to a boundary reflection which represents a free-surface reflection, which would occur. At the other boundaries, we damp artifical reflections using an **Absorbing Boundary frame** in order to model a realistic half-space model. For further details, how to implement this boundary frame, I refer to [this notebook](https://nbviewer.jupyter.org/github/daniel-koehn/Theory-of-seismic-waves-II/blob/master/05_2D_acoustic_FD_modelling/4_fdac2d_absorbing_boundary.ipynb).

Let's implement the 2D acoustic FD modelling code. As usual, we first import the required Python libraries ...

```python
# Import Libraries
# ----------------
import numpy as np
from numba import jit
import matplotlib
import matplotlib.pyplot as plt
from pylab import rcParams

# Ignore Warning Messages
# -----------------------
import warnings
warnings.filterwarnings("ignore")
```

... define the modelling parameters ...

```python
# Definition of initial modelling parameters
# ------------------------------------------
xmax = 2000.0 # maximum spatial extension of the 1D model in x-direction (m)
zmax = xmax   # maximum spatial extension of the 1D model in z-direction (m)
dx   = 10.0   # grid point distance in x-direction (m)
dz   = dx     # grid point distance in z-direction (m)

tmax = 0.75    # maximum recording time of the seismogram (s)
dt   = 0.0010 # time step

vp0  = 3000.  # P-wave speed in medium (m/s)

# acquisition geometry
xsrc = 1000.0 # x-source position (m)
zsrc = xsrc   # z-source position (m)

f0   = 100.0 # dominant frequency of the source (Hz)
t0   = 0.1   # source time shift (s)

isnap = 1  # snapshot interval (timesteps)
```

... compute the spatial derivatives in a function `update_d2px_d2pz` ...

```python
@jit(nopython=True) # use JIT for C-performance
def update_d2px_d2pz(p, dx, dz, nx, nz, d2px, d2pz):

    for i in range(1, nx - 1):
        for j in range(1, nz - 1):

            d2px[i,j] = (p[i + 1,j] - 2 * p[i,j] + p[i - 1,j]) / dx**2
            d2pz[i,j] = (p[i,j + 1] - 2 * p[i,j] + p[i,j - 1]) / dz**2

    return d2px, d2pz
```

... define the absorbing boundary frame at the left, right and bottom boundary in the function `absorb`

```python
# Define simple absorbing boundary frame based on wavefield damping
# according to Cerjan et al., 1985, Geophysics, 50, 705-708
def absorb(nx,nz):

    FW = 60     # thickness of absorbing frame (gridpoints)
    a = 0.0053

    coeff = np.zeros(FW)

    # define coefficients in absorbing frame
    for i in range(FW):
        coeff[i] = np.exp(-(a**2 * (FW-i)**2))

    # initialize array of absorbing coefficients
    absorb_coeff = np.ones((nx,nz))

    # compute coefficients for left grid boundaries (x-direction)
    zb=0
    for i in range(FW):
        ze = nz - i - 1
        for j in range(zb,ze):
            absorb_coeff[i,j] = coeff[i]

    # compute coefficients for right grid boundaries (x-direction)
    zb=0
    for i in range(FW):
        ii = nx - i - 1
        ze = nz - i - 1
        for j in range(zb,ze):
            absorb_coeff[ii,j] = coeff[i]

    # compute coefficients for bottom grid boundaries (z-direction)
    xb=0
    for j in range(FW):
        jj = nz - j - 1
        xb = j
        xe = nx - j
        for i in range(xb,xe):
            absorb_coeff[i,jj] = coeff[j]

    return absorb_coeff
```

Finally, we assemble all parts in the 2D acoustic finite difference modelling code `FD_2D_acoustic_JIT`

```python
# FD_2D_acoustic code with JIT optimization
# -----------------------------------------
def FD_2D_acoustic_JIT(dt,dx,dz,f0,xsrc,zsrc):

    # define model discretization
    # ---------------------------

    nx = (int)(xmax/dx) # number of grid points in x-direction
    print('nx = ',nx)

    nz = (int)(zmax/dz) # number of grid points in x-direction
    print('nz = ',nz)

    nt = (int)(tmax/dt) # maximum number of time steps
    print('nt = ',nt)

    isrc = (int)(xsrc/dx)  # source location in grid in x-direction
    jsrc = (int)(zsrc/dz)  # source location in grid in x-direction

    # Source time function (Gaussian)
    # -------------------------------
    src  = np.zeros(nt + 1)
    time = np.linspace(0 * dt, nt * dt, nt)

    # 1st derivative of Gaussian
    src  = -2. * (time - t0) * (f0 ** 2) * (np.exp(- (f0 ** 2) * (time - t0) ** 2))

    # define clip value: 0.1 * absolute maximum value of source wavelet
    clip = 0.1 * max([np.abs(src.min()), np.abs(src.max())]) / (dx*dz) * dt**2

    # Define absorbing boundary frame
    # -------------------------------
    absorb_coeff = absorb(nx,nz)

    # Define model
    # ------------
    vp  = np.zeros((nx,nz))
    vp  = model(nx,nz,vp,dx,dz)
    vp2 = vp**2

    # Initialize empty pressure arrays
    # --------------------------------
    p    = np.zeros((nx,nz)) # p at time n (now)
    pold = np.zeros((nx,nz)) # p at time n-1 (past)
    pnew = np.zeros((nx,nz)) # p at time n+1 (present)
    d2px = np.zeros((nx,nz)) # 2nd spatial x-derivative of p
    d2pz = np.zeros((nx,nz)) # 2nd spatial z-derivative of p

    # Initalize animation of pressure wavefield
    # -----------------------------------------
    fig = plt.figure(figsize=(8,2.5))  # define figure size
    plt.tight_layout()
    extent = [0.0,xmax,zmax,0.0]     # define model extension

    # Plot pressure wavefield movie
    ax1 = plt.subplot(121)
    image = plt.imshow(p.T, animated=True, cmap="RdBu", extent=extent,
                          interpolation='spline36', vmin=-clip, vmax=clip)
    plt.title('Pressure wavefield')
    plt.xlabel('x [m]')
    plt.ylabel('z [m]')
    cbar = plt.colorbar(image)
    cbar.set_label(r'Pressure [Pa]')

    # Plot Vp-model
    ax2 = plt.subplot(122)
    image1 = plt.imshow((vp.T)/1000, cmap=plt.cm.viridis, interpolation='nearest',
                        extent=extent)

    cbar1 = plt.colorbar(image1)
    cbar1.set_label(r'$V_p$ [km/s]')
    plt.title('Vp-model')
    plt.xlabel('x [m]')
    plt.setp(ax2.get_yticklabels(), visible=False)

    plt.ion()
    plt.show(block=False)

    snap_count = 0

    # Calculate Partial Derivatives
    # -----------------------------
    for it in range(nt):

        # FD approximation of spatial derivative by 3 point operator
        d2px, d2pz = update_d2px_d2pz(p, dx, dz, nx, nz, d2px, d2pz)

        # Time Extrapolation
        # ------------------
        pnew = 2 * p - pold + vp2 * dt**2 * (d2px + d2pz)

        # Add Source Term at isrc
        # -----------------------
        # Absolute pressure w.r.t analytical solution
        pnew[isrc,jsrc] = pnew[isrc,jsrc] + src[it] / (dx * dz) * dt ** 2

        # Apply absorbing boundary frame
        # ------------------------------
        p *= absorb_coeff
        pnew *= absorb_coeff

        # Remap Time Levels
        # -----------------
        pold, p = p, pnew

        # display pressure snapshots
        if (it % isnap) == 0:
            image.set_data(p.T)
            fig.canvas.draw()

            # write snapshots to Tiff files
            name_snap = "image_out/acoustic_2D_" + "%0.*f" %(0,np.fix(snap_count+1000)) + ".tiff"
            plt.savefig(name_snap, format='tiff', bbox_inches='tight', dpi=125)
            snap_count += 1

    return vp
```

## Problem 1: Homogeneous Model

Let 's start with a very simple problem - a homogeneous half-space model. We can define, and later modify, the subsurface model in the function `model`:

```python
# Homogeneous model
def model(nx,nz,vp,dx,dz):

    vp += vp0

    return vp
```

Time to define some additional modelling parameters:

- We place an explosive source, e.g. representing a subsurface nuclear explosion test, in the center of the model at $(xsrc,zsrc)^T$ = $(1000\; m,1000\; m)^T$

- Use a spatial sampling interval $dx = dz = 5\; m$ of the finite-difference grid

- Define the centre frequency of the source wavelet

- Calculate the time sample interval $dt$ in order to satisfy the CFL-criterion.

before finally running the FD code ...

```python
%matplotlib notebook
xsrc = 1000.0 # x-source position (m)
zsrc = xsrc   # z-source position (m)

dx   = 5.0   # grid point distance in x-direction (m)
dz   = dx     # grid point distance in z-direction (m)
f0   = 100.0  # centre frequency of the source wavelet (Hz)

# calculate dt according to the CFL-criterion
dt = dx / (np.sqrt(2.0) * vp0)

vp_hom = FD_2D_acoustic_JIT(dt,dx,dz,f0,xsrc,zsrc)
```

    nx =  400
    nz =  400
    nt =  636

    <IPython.core.display.Javascript object>

```python
from IPython.display import YouTubeVideo
YouTubeVideo('97QY9vDaEbo')
```

Notice, how the P-wavefront propagates in all directions, is absorbed at the left, right and bottom boundary, while reflected at the free-surface boundary on top, representing the earth surface.

## Problem 2: Random Medium

Next, we add some random perturbations to the homogeneous Vp-model:

```python
# Random medium model
def model(nx,nz,vp,dx,dz):

    pert = 0.4
    r = 2.0 * (np.random.rand(nz, nx) - 0.5) * pert
    vp += vp0 * (1 + r)

    return vp
```

```python
xsrc = 1000.0 # x-source position (m)
zsrc = xsrc    # z-source position (m)

dx   = 5.0   # grid point distance in x-direction (m)
dz   = dx     # grid point distance in z-direction (m)
f0   = 100.0  # centre frequency of the source wavelet (Hz)
vpmax = 4200.0
# vpmax = np.max(vp0)

# calculate dt according to the CFL-criterion
dt = dx / (np.sqrt(2.0) * vpmax)

vp_rand = FD_2D_acoustic_JIT(dt,dx,dz,f0,xsrc,zsrc)
```

    nx =  400
    nz =  400
    nt =  890

    <IPython.core.display.Javascript object>

```python
from IPython.display import YouTubeVideo
YouTubeVideo('Kvbrw5oXSfY')
```

As in the first example, we get a circular P-wavefront. However, the small scale random model variations scatter the seismic energy, leading to a significant damping of the direct P-wave amplitude.

## Problem 3: Fault Zone

In this problem, we model acoustic wave propagation in a vertical fault zone

```python
# Vertical fault zone model
def model(nx,nz,vp,dx,dz):

    fault_ext = 25.0           # half-extension of the fault in x-direction (m)
    n_ft = (int)(fault_ext/dx) # half-extension of the fault in x-direction (gridpoints)
    vp += vp0
    vp[nx // 2 - n_ft: nx // 2 + n_ft,:] *= 0.8

    return vp
```

```python
xsrc = 1000.0 # x-source position (m)
zsrc = 1000.0 # z-source position (m)

dx   = 5.0   # grid point distance in x-direction (m)
dz   = dx     # grid point distance in z-direction (m)
f0   = 100.0  # centre frequency of the source wavelet (Hz)
vpmax = 4200.0
# vpmax = np.max(vp0)

# calculate dt according to the CFL-criterion
dt = dx / (np.sqrt(2.0) * vpmax)

vp_fault = FD_2D_acoustic_JIT(dt,dx,dz,f0,xsrc,zsrc)
```

    nx =  400
    nz =  400
    nt =  890

    <IPython.core.display.Javascript object>

```python
from IPython.display import YouTubeVideo
YouTubeVideo('1anNRXd1cpU')
```

The vertical low-velocity fault leads to the development of an acoustic interface wave, sticking to the fault zone, with a lower velocity compared to the body P-wave propagating in the homogeneous medium, surrounding the fault zone.

## Problem 4: Simplified Vulcano

How does the surface topography of a vulcano scatter the acoustic wavefield?

```python
# Simplified vulcano (Gaussian hill)
def model(nx,nz,vp,dx,dz):

    vp += vp0 # fill model with vp0

    for i in range(0,nx):
        for j in range(0,nz):

            depth = j * dz  # depth of grid point j (m)
            x = i * dx      # horizontal distance of grid point i (m)

            x_center = (int)(nx/2) * dx  # center of vulcano cone
            height_vul = 1500.0          # vulcano height
            sigma = 500.0                # half-width of vulcano cone

            # approximate vulcano topography by Gaussian hill
            topo = - height_vul * np.exp(-(x - x_center)**2/sigma**2) + dx*nz

            # add air above the topography
            if(depth < topo):
                vp[i,j] = 0.0

    return vp
```

```python
xsrc = 1000.0 # x-source position (m)
zsrc = 1250.0 # z-source position (m)

dx   = 5.0   # grid point distance in x-direction (m)
dz   = dx     # grid point distance in z-direction (m)
f0   = 4.5  # centre frequency of the source wavelet (Hz)
vpmax = vp0

# calculate dt according to the CFL-criterion
dt = dx / (np.sqrt(2.0) * vpmax)

vp_topo = FD_2D_acoustic_JIT(dt,dx,dz,f0,xsrc,zsrc)
```

    nx =  400
    nz =  400
    nt =  636

    <IPython.core.display.Javascript object>

```python
from IPython.display import YouTubeVideo
YouTubeVideo('52YBfHOvAUI')
```

Modelling example 3, represents the occurence of a microseismic event in a Gaussian shaped volcano. Such events can be related to the movement of magma in the volcanic cone. After a short time, the seismic event leads to natural oscillations of the volcano.

## Problem 4: Building velocity models based on images

It is a quite tedious approach to build complex velocity models via IF-statements in the FOR-loops over the Cartesian FD grid. Alternatively, you can read an image and match colors with seismic P-wave velocities. Be sure to use PNGs or other non-compression image formats in order to avoid problems with compression artifacts. Let's take a look at this subsurface model, based on the image file `subsurface.png` located in the `images` directory:

```python
# Build "realistic" subsurface model based on imported image

from scipy import misc # import image read from scipy

def model(nx,nz,vp,dx,dz):

    # read jpeg image - be sure that the number of pixels coincides
    # with the number of grid points nx and ny of the FD model
    image = misc.imread('images/subsurface.png')
    image_red = image[:,:,0]  # extract red channel

    vp += vp0 # fill model with vp0

    for i in range(0,nx):
        for j in range(0,nz):

            # assume vaccum in the violet caves
            vp[i,j] = 0.

            # if image_red[j,i]<50, set vp = 2000 m/s
            if(image_red[j,i]<112):
                vp[i,j] = 2000.

            # if image_red[j,i]>200, set vp = 3000 m/s
            if(image_red[j,i]>150):
                vp[i,j] = 2700.

    return vp
```

```python
xsrc = 1000.0 # x-source position (m)
zsrc = 100.0  # z-source position (m)

dx   = 5.0    # grid point distance in x-direction (m)
dz   = dx     # grid point distance in z-direction (m)
f0   = 100.   # centre frequency of the source wavelet (Hz)
vpmax = vp0

# calculate dt according to the CFL-criterion
dt = dx / (np.sqrt(2.0) * vpmax)

vp_image = FD_2D_acoustic_JIT(dt,dx,dz,f0,xsrc,zsrc)
```

    nx =  400
    nz =  400
    nt =  636

    <IPython.core.display.Javascript object>

```python
from IPython.display import YouTubeVideo
YouTubeVideo('cQoshYobv_0')
```

In this complex subsurface model, the interface between the top low-velocity layer with $vp=2000\;m/s$ and the underlying model with $vp=2700\;m/s$, as well as the cavities leads to a complex seismic wavefield, consisting of reflections and diffractions.

## Problem 6: Create your own problem here

##### Excercise

Now it's your turn, create a 2D P-wave velocity model and compute the pressure
wavefield by 2D acoustic FD modelling. Pay attention to satisfy the CFL and grid dispersion criteria by choosing appropriate spatial grid point distances $dx = dz$
and time step dt.

## What we learned:

- How to model P-wave propgation in simple heterogeneous 2D acoustic media


# MHD Equation with CentPy in 1D

### Import packages

```python
# Install the centpy package
!pip install centpy
```

    Requirement already satisfied: centpy in /usr/local/lib/python3.6/dist-packages (0.1)
    Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from centpy) (1.18.5)

```python
# Import numpy and centpy for the solution
import numpy as np
import centpy
```

```python
# Imports functions from matplotlib and setup for the animation
import matplotlib.pyplot as plt
from matplotlib import animation
from IPython.display import HTML
```

### Equation

We solve the equations of ideal magnetohydrodynamics in 1D

\begin{equation}
\partial_t
\begin{bmatrix}
\rho \\ \rho v_x \\ \rho v_y \\ \rho v_z \\ B_y \\ B_z \\ E
\end{bmatrix}
+
\partial_x
\begin{bmatrix}
\rho v_x \\ \rho v_x^2 + p^* - B_x^2 \\ \rho v_x v_y - B_x B_y \\
\rho v_x v_z - B_x B_z \\ B_y v_x - B_x v_y \\ B_z v_x - B_x v_z \\
(E+p^*) v_x - B_x (B_x v_x + B_y v_y + B_z v_Z)
\end{bmatrix}
= 0
\end{equation}

where the total pressure is given by

\begin{equation}
p^* = p + \frac{1}{2} (B_x^2 + B_y^2 + B_z^2)
\end{equation}

with the equation of state

\begin{equation}
p = (\gamma-1) \left(E-\frac{1}{2} \rho (v_x^2+v_y^2+v_z^2) -
\frac{1}{2}(B_x^2 + B_y^2 + B_z^2)\right), \qquad \gamma=2.0
\end{equation}

The solution is computed on the domain $(x,t)\in([-1,1]\times[0,0.2])$ with initial data for a *Brio-Wu shock tube*:

\begin{equation}
(\rho, v_x, v_y, v_z, B_y, B_z, p)_{t=0} =
\begin{cases}
(1,0,0,0,1,0,1) & \text{if} & -1<x\leq 0 \\
(0.125, 0, 0, 0, -1, 0, 0.1) & \text{if} & \ \ 0<x<1
\end{cases}
\end{equation}

and Dirichlet boundary data set by initial data on each boundary. The solution is computed using 400 cells and CFL number 0.475.

```python
pars = centpy.Pars1d(x_init=-1., x_final=1., t_final=0.2, dt_out=0.002, J=400, cfl=0.475, scheme="fd2")
pars.B1 = 0.75
```

```python
# MHD equation
class MHD1d(centpy.Equation1d):
    def pressure(self, u):
        return u[:, 6] - 0.5*((u[:, 1] ** 2 + u[:, 2] ** 2 + u[:, 3] ** 2)/u[:, 0])
        - 0.5 * (self.B1 ** 2 + u[:, 4] ** 2 + u[:, 5] ** 2)

    def initial_data(self):
        u = np.zeros((self.J + 4, 7))
        midpoint = int(self.J / 2) + 2

        # Left side
        u[:midpoint, 0] = 1.0
        u[:midpoint, 1] = 0.0
        u[:midpoint, 2] = 0.0
        u[:midpoint, 3] = 0.0
        u[:midpoint, 4] = 1.0
        u[:midpoint, 5] = 0.0
        u[:midpoint, 6] = 1.0 + 25.0 / 32.0

        # Right side
        u[midpoint:, 0] = 0.125
        u[midpoint:, 1] = 0.0
        u[midpoint:, 2] = 0.0
        u[midpoint:, 3] = 0.0
        u[midpoint:, 4] = -1.0
        u[midpoint:, 5] = 0.0
        u[midpoint:, 6] = 0.1 + 25.0 / 32.0

        return u

    def boundary_conditions(self, u):

        left_v = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0 + 25.0 / 32.0]
        right_v = [0.125, 0.0, 0.0, 0.0, -1.0, 0.0, 0.1 + 25.0 / 32]

        if self.odd:
            u[0] = left_v
            u[-1] = right_v
            u[-2] = right_v
        else:
            u[0] = left_v
            u[1] = left_v
            u[-1] = right_v

    def flux_x(self, u):
        f = np.zeros_like(u)

        B1 = self.B1

        p_star = self.pressure(u) + 0.5 * (B1 ** 2 + u[:, 4] ** 2 + u[:, 5] ** 2)

        f[:, 0] = u[:, 1]
        f[:, 1] = u[:, 1] ** 2 / u[:, 0] + p_star
        f[:, 2] = u[:, 1] * u[:, 2] / u[:, 0] - B1 * u[:, 4]
        f[:, 3] = u[:, 1] * u[:, 3] / u[:, 0] - B1 * u[:, 5]
        f[:, 4] = u[:, 1] * u[:, 4] / u[:, 0] - B1 * u[:, 2] / u[:, 0]
        f[:, 5] = u[:, 1] * u[:, 5] / u[:, 0] - B1 * u[:, 3] / u[:, 0]
        f[:, 6] = (u[:, 6] + p_star) * (u[:, 1] / u[:, 0]) - B1 * (
            B1 * u[:, 1] + u[:, 2] * u[:, 4] + u[:, 3] * u[:, 5]
        ) / u[:, 0]

        return f

    def spectral_radius_x(self, u):
        rho = u[:, 0]
        u_x = u[:, 1] / rho
        p = self.pressure(u)
        A = 2.0 * p / rho
        B = (self.B1 ** 2 + u[:, 4] ** 2 + u[:, 5] ** 2) / rho
        cf = np.sqrt(
            0.5 * (A + B + np.sqrt((A + B) ** 2 - 4.0 * A * self.B1 ** 2 / rho))
        )
        return np.abs(u_x) + cf
```

### Solution

```python
eqn = MHD1d(pars)
soln = centpy.Solver1d(eqn)
soln.solve()
```

### Animation

```python
# Animation
j0 = slice(2,-2)

fig = plt.figure(figsize=(12,6))
ax1=fig.add_subplot(1,3,1)
ax2=fig.add_subplot(2,3,2)
ax3=fig.add_subplot(2,3,3)
ax4=fig.add_subplot(2,3,5)
ax5=fig.add_subplot(2,3,6)

line_u=[]
for ax in [ax1,ax2,ax3,ax4,ax5]:
  ax.set_xlim(-1, 1)
  line_u.append(ax.plot([], [], linewidth=1, marker='o', markersize=2)[0])

ax1.set_xlabel('x')
ax4.set_xlabel('x')
ax5.set_xlabel('x')

ax2.set_xticks([])
ax3.set_xticks([])
ax2.set_yticks([])
ax3.set_yticks([])
ax4.set_yticks([])
ax5.set_yticks([])

ax1.set_ylabel(r'$\rho$', fontsize=12)
ax2.set_ylabel(r'$v_x$', fontsize=12)
ax3.set_ylabel(r'$v_y$', fontsize=12)
ax4.set_ylabel(r'$B_y$', fontsize=12)
ax5.set_ylabel(r'$p$', fontsize=12)

# Primitive variables
rho=soln.u_n[:,j0,0]
v_x = soln.u_n[:,j0,1]/soln.u_n[:,j0,0]
v_y = soln.u_n[:,j0,2]/soln.u_n[:,j0,0]
B_y = soln.u_n[:,j0,4]

ax1.set_ylim(np.min(rho), np.max(rho))
ax2.set_ylim(np.min(v_x), np.max(v_x))
ax3.set_ylim(np.min(v_y), np.max(v_y))
ax4.set_ylim(np.min(B_y), np.max(B_y))
ax5.set_ylim(0.5, 2.)

def animate(i):
    p = eqn.pressure(soln.u_n[i,j0,:])
    line_u[0].set_data(soln.x[j0], rho[i])
    line_u[1].set_data(soln.x[j0], v_x[i])
    line_u[2].set_data(soln.x[j0], v_y[i])
    line_u[3].set_data(soln.x[j0], B_y[i])
    line_u[4].set_data(soln.x[j0], p)

plt.close()
anim = animation.FuncAnimation(fig, animate, frames=soln.Nt, interval=100, blit=False);
HTML(anim.to_html5_video())

```


## SurfinPy

#### Tutorial 1 - Generating a phase diagram as a function of chemical potential and surface energy at 0 K

In this tutorial we will learn how to generate a basic phase diagram from DFT energies. This example will consider a series of surfaces that contain differing amounts of surface oxygen and adsorbed water species.
The physical quantity that is used to define the stability of a surface with a given composition is its surface energy $\gamma$ (J $m^{-2}$). Going forward in this tutorial we will use the example of water adsorbing on to defective Ti$O_2$ surfaces.

\begin{align}
\gamma_{Surf} & = \frac{1}{2A} \Bigg( E_{TiO_2}^{slab} - \frac{nTi_{slab}}{nTi_{Bulk}} E_{TiO_2}^{Bulk} \Bigg) - \Gamma_O \mu_O - \Gamma_{H_2O} \mu_{H_2O} ,
\end{align}

where A is the surface area, $E_{TiO_2}^{slab}$ is the DFT energy of the slab, $nTi_{Slab}$ is the number of cations in the slab, $nTi_{Bulk}$ is the number of cations in the bulk unit cell, $E_{TiO_2}^{Bulk}$ is the DFT energy of the bulk unit cell and

\begin{align}
\Gamma_O & = \frac{1}{2A} \Bigg( nO_{Slab} - \frac{nO_{Bulk}}{nTi_{Bulk}}nTi_{Slab}  \Bigg) ,
\end{align}

\begin{align}
\Gamma_{H_2O} & = \frac{nH_2O}{2A} ,
\end{align}

where $nO_{Slab}$ is the number of anions in the slab, $nO_{Bulk}$ is the number of anions in the bulk and $nH_2O$ is the number of adsorbing water molecules. $\Gamma_O$ / $\Gamma_{H_2O}$ is the excess oxygen / water at the surface and $\mu_O$ / $\mu_{H_2O}$ is the oxygen / water chemcial potential. Clearly $\Gamma$ and $\mu$ will only matter when the surface is non stoichiometric.

So now lets work through an example

```python
from surfinpy import mu_vs_mu
```

The first thing to do is input the data that we have generated from our DFT simulations. The input data needs to be contained within a dictionary. First we have created the dictionary for the bulk data, where 'Cation' is the number of cations, 'Anion' is the number of anions, 'Energy' is the DFT energy and 'F-Units' is the number of formula units.
 ```
bulk = {'Cation' : Cations in Bulk Unit Cell,
        'Anion' : Anions in Bulk Unit Cell,
        'Energy' :  Energy of Bulk Calculation,
        'F-Units' : Formula units in Bulk Calculation}
```

```python
bulk = {'Cation' : 1, 'Anion' : 2, 'Energy' : -780.0, 'F-Units' : 4}
```

Next we create the surface dictionaries - one for each surface or "phase". 'Cation' is the number of cations, 'X' is in this case the number of oxygen species (corresponding to the X axis of the phase diagram), 'Y' is the number of in this case water molecules (corresponding to the Y axis of our phase diagram), 'Area' is the surface area, 'Energy' is the DFT energy, 'Label' is the label for the surface (appears on the phase diagram) and finally nSpecies is the number of adsorbin species.

```
    surface = {'Cation': Cations in Slab,
               'X': Number of Species X in Slab,
               'Y': Number of Species Y in Slab,
               'Area': Surface area in the slab,
               'Energy': Energy of Slab,
               'Label': Label for phase,
               'nSpecies': How many species are non stoichiometric}
```

```python
pure =     {'Cation': 24, 'X': 48, 'Y': 0, 'Area': 60.0, 'Energy': -575.0,   'Label': 'Stoich',  'nSpecies': 1}
H2O =      {'Cation': 24, 'X': 48, 'Y': 2, 'Area': 60.0, 'Energy': -612.0,   'Label': '1 Water', 'nSpecies': 1}
H2O_2 =    {'Cation': 24, 'X': 48, 'Y': 4, 'Area': 60.0, 'Energy': -640.0,   'Label': '2 Water', 'nSpecies': 1}
H2O_3 =    {'Cation': 24, 'X': 48, 'Y': 8, 'Area': 60.0, 'Energy': -676.0,   'Label': '3 Water', 'nSpecies': 1}
Vo =       {'Cation': 24, 'X': 46, 'Y': 0, 'Area': 60.0, 'Energy': -558.0,   'Label': 'Vo', 'nSpecies': 1}
H2O_Vo =   {'Cation': 24, 'X': 46, 'Y': 2, 'Area': 60.0, 'Energy': -594.0,  'Label': 'Vo + 1 Water', 'nSpecies': 1}
H2O_Vo_2 = {'Cation': 24, 'X': 46, 'Y': 4, 'Area': 60.0, 'Energy': -624.0,  'Label': 'Vo + 2 Water', 'nSpecies': 1}
H2O_Vo_3 = {'Cation': 24, 'X': 46, 'Y': 6, 'Area': 60.0, 'Energy': -640.0, 'Label': 'Vo + 3 Water', 'nSpecies': 1}
H2O_Vo_4 = {'Cation': 24, 'X': 46, 'Y': 8, 'Area': 60.0, 'Energy': -670.0, 'Label': 'Vo + 4 Water', 'nSpecies': 1}
```

Next we need to create a list of our data. DOn't worry about the order, surfinpy will sort that out for you.

```python
data = [pure, H2O_2, H2O_Vo, H2O,  H2O_Vo_2, H2O_3, H2O_Vo_3,  H2O_Vo_4, Vo]
```

We now need to generate our X and Y axis, or more appropriately, our chemical potential values. Again these exist in a dictionary. 'Range' corresponds to the range of chemcial potential values to be considered and 'Label' is the axis label.
```
    deltaX = {'Range': Range of Chemical Potential,
              'Label': Species Label}

```

```python
deltaX = {'Range': [ -12, -6],  'Label': 'O'}
deltaY = {'Range': [ -19, -12], 'Label': 'H_2O'}
```

And finally we can generate our plot using these 4 variables of data.

```python
system = mu_vs_mu.calculate(data, bulk, deltaX, deltaY)
system.plot_phase()
```

This plot is a good start in that the relative stability of the surfaces has been evaluated. However the chemical potential values are essentially meaningless and can be dependent on the pseudopotentials. The first thing that we can do is add in the DFT energy of species X/Y, this will give us a 0K phase diagram.

```python
Zero_K = mu_vs_mu.calculate(data, bulk, deltaX, deltaY, x_energy=-4.54, y_energy=-14.84)
Zero_K.plot_phase(output="seaborn_rdybu.png", set_style="seaborn-dark-palette", colourmap="RdYlBu")
```