```python
import numpy as np
import scipy.misc
from scipy.fftpack import dct, idct
import sys
from PIL import Image
import matplotlib
import matplotlib.pyplot as plt
import random
from tqdm._tqdm_notebook import tqdm_notebook
from scipy.fftpack import dct, idct
import seaborn as sns
from skimage.metrics import structural_similarity as ssim
import pandas as pd
import sympy

%matplotlib inline

class ImageLoader:
    def __init__(self, FILE_PATH):
        self.img = np.array(Image.open(FILE_PATH))
        # 行数
        self.row_blocks_count = self.img.shape[0] // 8
        # 列数
        self.col_blocks_count = self.img.shape[1] // 8

    def get_points(self, POINT):
        Row = random.randint(0, len(self.img) - POINT - 1)
        Col = random.randint(0, len(self.img) - 1)
        return self.img[Row : Row + POINT, Col]

    def get_block(self, col, row):
        return self.img[col * 8 : (col + 1) * 8, row * 8 : (row + 1) * 8]

# plt.rcParams['font.family'] ='sans-serif'#使用するフォント
# plt.rcParams["font.sans-serif"] = "Source Han Sans"
plt.rcParams["font.family"] = "Source Han Sans JP"  # 使用するフォント
plt.rcParams["xtick.direction"] = "in"  # x軸の目盛線が内向き('in')か外向き('out')か双方向か('inout')
plt.rcParams["ytick.direction"] = "in"  # y軸の目盛線が内向き('in')か外向き('out')か双方向か('inout')
plt.rcParams["xtick.major.width"] = 1.0  # x軸主目盛り線の線幅
plt.rcParams["ytick.major.width"] = 1.0  # y軸主目盛り線の線幅
plt.rcParams["font.size"] = 12  # フォントの大きさ
plt.rcParams["axes.linewidth"] = 1.0  # 軸の線幅edge linewidth。囲みの太さ
matplotlib.font_manager._rebuild()

MONO_DIR_PATH = "../../Mono/"
AIRPLANE = ImageLoader(MONO_DIR_PATH + "airplane512.bmp")
BARBARA = ImageLoader(MONO_DIR_PATH + "barbara512.bmp")
BOAT = ImageLoader(MONO_DIR_PATH + "boat512.bmp")
GOLDHILL = ImageLoader(MONO_DIR_PATH + "goldhill512.bmp")
LENNA = ImageLoader(MONO_DIR_PATH + "lenna512.bmp")
MANDRILL = ImageLoader(MONO_DIR_PATH + "mandrill512.bmp")
MILKDROP = ImageLoader(MONO_DIR_PATH + "milkdrop512.bmp")
SAILBOAT = ImageLoader(MONO_DIR_PATH + "sailboat512.bmp")

IMAGES = [
    AIRPLANE,
    BARBARA,
    BOAT,
    GOLDHILL,
    LENNA,
    MANDRILL,
    MILKDROP,
    SAILBOAT
]
```

```python
# n_bar = 4
N = 8
```

```python
IMG = LENNA
```

```python
Fk = np.zeros(IMG.img.shape)
```

```python
sampling_x = (0.5 + np.arange(N)) / N
u_1 = (1 - sampling_x) ** 2 / 2
u_2 = 1 / 24 - sampling_x ** 2 / 6 + sampling_x ** 3 / 6 - sampling_x ** 4 / 24
u_3 = sampling_x ** 2 / 2
u_4 = -sampling_x ** 2 / 12 + sampling_x ** 4 / 24
```

```python
ak = - scipy.fftpack.dct(u_1,norm="ortho")
bk = scipy.fftpack.dct(u_2,norm="ortho")
ck = scipy.fftpack.dct(u_3,norm="ortho")
dk = scipy.fftpack.dct(u_4,norm="ortho")

alpha = ak
beta = ck

Ak = (2 * ak - 16 * bk) / np.sqrt(N)
Bk = (2 * ck - 16 * dk) / np.sqrt(N)
Ck = (2 * ak - 32 * bk) / np.sqrt(N)
Dk = (2 * ck - 32 * dk) / np.sqrt(N)
```

```python
def TPHLCT(IMG,N):
    Fk = np.zeros(IMG.img.shape)
    # DCT
    for row in range(IMG.img.shape[0] // N):
        for col in range(IMG.img.shape[1]):
            eight_points = IMG.img[N * row : N * (row + 1), col]
            c = scipy.fftpack.dct(eight_points,norm="ortho")
            Fk[N * row : N * (row + 1), col] = c

    # 3階
    for row in range(Fk.shape[0] // N):
        for col in range(Fk.shape[1]):
            F = Fk[N * row : N * (row + 1), col]

            F_0_r = 0
            F_1_r = 0
            if row is not Fk.shape[0] // N - 1:
                F_0_r = Fk[N * (row + 1), col]
                F_1_r = Fk[N * (row + 1) + 1, col]

            F_0_l = 0
            F_1_l = 1
            if row is not 0:
                F_0_l = Fk[N * (row - 1), col]
                F_1_l = Fk[N * (row - 1) + 1, col]

            # 残差
            F_0 = F[0]
            F_1 = F[1]

            F = (
                F
                - Ak * (F_0 - F_0_l)
                - Bk * (F_0_r - F_0)
                - Ck * (F_1 + F_1_l)
                - Dk * (F_1_r + F_1)
            )
            # F_0, F_1は残す
            F[0] = F_0
            F[1] = F_1
            # F_0 V_1 V_2 V_3 V_4 V_5 V_6 V_7
            Fk[N * row : N * (row + 1), col] = F

    # 1階
    for row in range(Fk.shape[0] // N):
        for col in range(Fk.shape[1]):
            F = Fk[N * row : N * (row + 1), col]

            F_0_r = 0
            if row is not Fk.shape[0] // N - 1:
                F_0_r = Fk[N * (row + 1), col]

            F_0_l = 0
            if row is not 0:
                F_0_l = Fk[N * (row - 1), col]

            # 残差
            F_0 = F[0]
            F_temp = F - alpha * (F_0_r - F_0) / np.sqrt(N) - beta * (F_0 - F_0_l) / np.sqrt(N)
            # F_0は残す
            F[1] = F_temp[1]
            # F_0 V_1 V_2 V_3 V_4 V_5 V_6 V_7
            Fk[N * row : N * (row + 1), col] = F

    # DCT
    for row in range(Fk.shape[0]):
        for col in range(Fk.shape[1] // N):
            eight_points = Fk[row, N * col : N * (col + 1)]
            c = scipy.fftpack.dct(eight_points,norm="ortho")
            Fk[row, N * col : N * (col + 1)] = c

    # 3階
    for row in range(Fk.shape[0]):
        for col in range(Fk.shape[1] // N):
            F = Fk[row, N * col : N * (col + 1)]

            F_0_r = 0
            F_1_r = 0
            if col is not Fk.shape[1] // N - 1:
                F_0_r = Fk[row, N * (col + 1)]
                F_1_r = Fk[row, N * (col + 1) + 1]

            F_0_l = 0
            F_1_l = 0
            if col is not 0:
                F_0_l = Fk[row, N * (col - 1)]
                F_1_l = Fk[row, N * (col - 1) + 1]

            # 残差
            F_0 = F[0]
            F_1 = F[1]

            F = (
                F
                - Ak * (F_0 - F_0_l)
                - Bk * (F_0_r - F_0)
                - Ck * (F_1 + F_1_l)
                - Dk * (F_1_r + F_1)
            )
            # F_0は残す
            F[0] = F_0
            F[1] = F_1
            # F_0 V_1 V_2 V_3 V_4 V_5 V_6 V_7
            Fk[row, N * col : N * (col + 1)] = F

    for row in range(Fk.shape[0]):
        for col in range(Fk.shape[1] // N):
            F = Fk[row, N * col : N * (col + 1)]

            F_0_r = 0
            if col is not Fk.shape[1] // N - 1:
                F_0_r = Fk[row, N * (col + 1)]

            F_0_l = 0
            if col is not 0:
                F_0_l = Fk[row, N * (col - 1)]

            # 残差
            F_0 = F[0]
            F_temp = F - alpha * (F_0_r - F_0) / np.sqrt(N) - beta * (F_0 - F_0_l) / np.sqrt(N)
            # F_0は残す
            F[1] = F_temp[1]
            # F_0 V_1 V_2 V_3 V_4 V_5 V_6 V_7
            Fk[row, N * col : N * (col + 1)] = F

    return Fk
```

```python
# DCT係数の平均を求める
```

```python
Fk_values = np.zeros((512,512))
```

```python
for IMG in tqdm_notebook(IMAGES):
    values = np.zeros((25,4))
    Fk = np.zeros(IMG.img.shape)
    for row in range(IMG.img.shape[0] // N):
        for col in range(IMG.img.shape[1] // N):
            block = IMG.img[row * N : (row + 1) * N, col * N : (col + 1) * N]
            c = scipy.fftpack.dct(
                scipy.fftpack.dct(block, axis=0, norm="ortho"), axis=1, norm="ortho"
            )
            Fk[row * N : (row + 1) * N, col * N : (col + 1) * N] = c

    Fk_values += np.abs(Fk)
```

    HBox(children=(IntProgress(value=0, max=8), HTML(value='')))

```python
Fk_values /= len(IMAGES)
pd.DataFrame(Fk_values).to_csv("DCT_coef_ave.csv",header=False,index=False)
```

```python
# 各n_barの残差係数の平均を求める
```

```python
Vk_values = np.zeros((512,512))
```

```python
Vk_values = np.zeros((512,512))
for IMG in tqdm_notebook(IMAGES):
    Fk = TPHLCT(IMG,N)
    Vk_values += np.abs(Fk)

pd.DataFrame(Vk_values / len(IMAGES)).to_csv("TPHLCT_coef_ave.csv",header=False,index=False)
```

    HBox(children=(IntProgress(value=0, max=8), HTML(value='')))

```python
# DCT係数を読み込む
```

```python
Fk_values = pd.read_csv("DCT_coef_ave.csv",header=None).values
```

```python
# 残差係数を読み込む
```

```python
Vk_values = pd.read_csv("TPHLCT_coef_ave.csv",header=None).values
```

```python
# NxNブロック1個当たりの係数の平均を求める
```

```python
Fk_block_ave_values = np.zeros((N,N))
for row in range(IMG.img.shape[0] // N):
    for col in range(IMG.img.shape[1] // N):
        if col == 0:
            continue
        if col == IMG.img.shape[1] // N -1:
            continue
        if row == 0:
            continue
        if row == IMG.img.shape[0] // N -1:
            continue
        block = Fk_values[row * N : (row + 1) * N, col * N : (col + 1) * N]
        Fk_block_ave_values += np.abs(block)

Fk_block_ave_values /= (IMG.img.shape[0]//N)**2
```

```python
Vk_block_ave_values = np.zeros((N,N))
for row in range(IMG.img.shape[0] // N):
    for col in range(IMG.img.shape[1] // N):
        if col == 0:
            continue
        if col == IMG.img.shape[1] // N -1:
            continue
        if row == 0:
            continue
        if row == IMG.img.shape[0] // N -1:
            continue
        block = Vk_values[row * N : (row + 1) * N, col * N : (col + 1) * N]
        Vk_block_ave_values += np.abs(block)

Vk_block_ave_values /= (IMG.img.shape[0]//N)**2
```

```python
df = pd.DataFrame(columns=["G(k_1,k_2)"])
```

```python
for index in range(1,k_max,1):
    for i in range(index):
        Gk1k2_1 = []
        Gk1k2_2 = []
        Vk = Vk_block_ave_values[i,index]
        Fk = Fk_block_ave_values[i,index]
        Gk1k2_1.append(100 * (1 - Vk/Fk))

        Vk = Vk_block_ave_values[index,i]
        Fk = Fk_block_ave_values[index,i]
        Gk1k2_2.append(100 * (1 - Vk/Fk))

        df.loc["(" + str(i) + "," + str(index) + ")"] = Gk1k2_1
        df.loc["(" + str(index) + "," + str(i) + ")"] = Gk1k2_2

    Gk1k2 = []

    Vk = Vk_block_ave_values[index,index]
    Fk = Fk_block_ave_values[index,index]
    Gk1k2.append(100 * (1 - Vk/Fk))

    df.loc["(" + str(index) + "," + str(index) + ")"] = Gk1k2

Gk1k2 = []

Vk_sum = 0
Fk_sum = 0
for row in range(N):
    for col in range(N):
        if row > k_max-1 or col > k_max-1:
            Vk_sum += Vk_block_ave_values[row,col]
            Fk_sum += Fk_block_ave_values[row,col]
Gk1k2.append(100 * (1 - Vk_sum/Fk_sum))
df.loc["others"] = Gk1k2
```

```python
df.to_csv("TPHLCT_high_freq_comp.csv")
df
```

      G(k_1,k_2)

      (0,1)
      20.717614

      (1,0)
      22.109009

      (1,1)
      13.555398

      (0,2)
      18.142591

      (2,0)
      17.046191

      (1,2)
      12.608437

      (2,1)
      12.401644

      (2,2)
      11.519603

      (0,3)
      8.359163

      (3,0)
      8.110313

      (1,3)
      4.960569

      (3,1)
      6.425333

      (2,3)
      6.951200

      (3,2)
      6.964452

      (3,3)
      5.062636

      (0,4)
      5.381535

      (4,0)
      5.198777

      (1,4)
      2.939087

      (4,1)
      3.619396

      (2,4)
      3.901762

      (4,2)
      4.711877

      (3,4)
      4.129811

      (4,3)
      4.055053

      (4,4)
      2.734539

      others
      1.490762

```python

```


###### Content under Creative Commons Attribution license CC-BY 4.0, code under MIT license (c)2014 L.A. Barba, C.D. Cooper, G.F. Forsyth.

# Spreading out

You've reached the fourth module of the open course [**"Practical Numerical Methods with Python"**](http://openedx.seas.gwu.edu/courses/GW/MAE6286/2014_fall/about), titled *Spreading out: Parabolic PDEs*. We hope that you are enjoying the ride of [#numericalmooc](https://twitter.com/hashtag/numericalmooc) so far!

We introduced finite-difference methods for partial differential equations (PDEs) in the [second module](https://github.com/numerical-mooc/numerical-mooc/tree/master/lessons/02_spacetime), and looked at convection problems in more depth in [module 3](https://github.com/numerical-mooc/numerical-mooc/tree/master/lessons/03_wave). Now we'll look at solving problems dominated by diffusion.

Why do we separate the discussion of how to solve convection-dominated and diffusion-dominated problems, you might ask? It's all about the harmony between mathematical model and numerical method. Convection and diffusion are inherently different physical processes.

The titles of the course modules are meant to spark your imagination:

* _Riding the wave_—imagine a surfer on a tall wave, moving fast towards the beach ... convection implies transport, speed, direction. The physics has a directional bias, and we discovered that numerical methods should be compatible with that. That's why we use _upwind_ methods for convection, and we pay attention to problems where waves move in opposite directions, needing special schemes.

* _Spreading out_—now imagine a drop of food dye in a cup of water, slowly spreading in all directions until all the liquid takes a uniform color. [Diffusion](http://en.wikipedia.org/wiki/Diffusion) spreads the concentration of something around (atoms, people, ideas, dirt, anything!). Since it is not a directional process, we need numerical methods that are isotropic (like central differences).

```
from IPython.display import Image
Image(url='http://upload.wikimedia.org/wikipedia/commons/f/f9/Blausen_0315_Diffusion.png')
```

## Parabolic PDEs

You already met the simplest parabolic PDE—the [1-D diffusion equation](http://nbviewer.ipython.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/02_spacetime/02_03_1DDiffusion.ipynb)—in module 2. Its main feature is that it has a second-order derivative in space. Here it is again:

$$
\frac{\partial u}{\partial t} = \alpha \frac{\partial^2 u}{\partial x^2}
$$

Check out the article on [parabolic PDEs](http://en.wikipedia.org/wiki/Parabolic_partial_differential_equation) in Wikipedia. Now compare with the diffusion equation above, with the two independent variables here being $x, t$. You'll see that with no mixed derivatives, and only one second-order derivative (in the spatial variable $x$), it satisfies the condition of a parabolic PDE. Work it out on a piece of paper if you need to.

In the previous module, discussing hyperbolic conservation laws, we learned that solutions have characteristics: information travels along certain paths on space-time phase space. In contrast, parabolic equations don't have characteristics, because any local change in the initial condition will eventually affect the entire domain, although its effect will be felt at smaller intensity with longer distances. This is typical of diffusion processes.

In this first lesson of the module, we first review the 1D diffusion equation and then take a deeper look into the issue of boundary conditions. In the next notebook, we'll introduce _implicit discretizations_ for the first time, which will take us to the land of linear solvers. In the third lesson we'll graduate to two dimensions—more boundary condition and stability issues will come up. We'll then study 2D implicit methods, and go into Crank-Nicolson method: perhaps the most popular of them all. _Enjoy!_

## Heat conduction

Heat conduction is a diffusive process. Let's remind ourselves of the heat equation in one spatial dimension:

\begin{equation}
\frac{\partial T}{\partial t} = \alpha \frac{\partial^2 T}{\partial x^2}
\end{equation}

Here, $\alpha$ is the thermal diffusivity, a property of the material, and $T$ is the temperature.

In the [third lesson of module 2](http://nbviewer.ipython.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/02_spacetime/02_03_1DDiffusion.ipynb), we discretized the diffusion equation  with a forward-time, centered-space scheme, subject to the following stability constraint:

\begin{equation}
\alpha \frac{\Delta t}{(\Delta x)^2} \leq \frac{1}{2}
\end{equation}

Let's look into it more deeply now, using a 1D temperature-evolution problem.

### Problem set up

Say we have a graphite rod, with [thermal diffusivity](http://en.wikipedia.org/wiki/Thermal_diffusivity) of $\alpha=1.22\times10^{-3} {\rm m}^2/{\rm s}$, length $L=1{\rm m}$, and temperature $T=0{\rm C}$ everywhere. At time $t=0$, we raise the temperature on the left-side end, $x=0$, to $T=100{\rm C}$, and hold it there. *How will the temperature evolve in the rod?*

As usual, start by importing some libraries and setting up the discretization. We'll begin by using a spatial grid with 51 points and advance for 100 time steps, using a forward-time/centered scheme.

.
#### Figure 1. Graphite rod, temperature fixed on ends.

```
import numpy
from matplotlib import pyplot
%matplotlib inline
from matplotlib import rcParams
rcParams['font.family'] = 'serif'
rcParams['font.size'] = 16
```

```
L = 1.
nt = 100
nx = 51.
alpha = 1.22e-3

dx = L/(nx-1)

Ti = numpy.zeros(nx)
Ti[0] = 100
```

Remember the forward-time, centered-space discretization? You should work it out on a piece of paper yourself (if you can't do it without looking it up, it means you need to do this more!).

\begin{equation}\frac{T_{i}^{n+1}-T_{i}^{n}}{\Delta t}=\alpha\frac{T_{i+1}^{n}-2T_{i}^{n}+T_{i-1}^{n}}{\Delta x^2}\end{equation}

To obtain the temperature at the next time step, $T^{n+1}_i$, from the known information at the current time step, we compute

\begin{equation}T_{i}^{n+1}=T_{i}^{n}+\frac{\alpha\Delta t}{\Delta x^2}(T_{i+1}^{n}-2T_{i}^{n}+T_{i-1}^{n})\end{equation}

Check the [third notebook of module 2](http://nbviewer.ipython.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/02_spacetime/02_03_1DDiffusion.ipynb), if you need to refresh your memory!

The following function implements this numerical scheme:

```
def ftcs(T, nt, dt, dx, alpha):
    """Solves the diffusion equation with forward-time, centered scheme

    Parameters:
    ----------
    T: array of float
        Initial temperature profile
    nt: int
        Number of time steps
    dt: float
        Time step size
    dx: float
        Mesh size
    alpha: float
        Diffusion coefficient (thermal diffusivity)

    Returns:
    -------
    T: array of float
        Temperature profile after nt time steps with forward in time scheme

    """
    for n in range(nt):
        Tn = T.copy()
        T[1:-1] = Tn[1:-1] + alpha*dt/dx**2*(Tn[2:] -2*Tn[1:-1] + Tn[0:-2])

    return T
```

We are all set to run! First, let's use a time step `dt` that satisfies the stability constraint

```
sigma = 1/2.0
dt = sigma * dx*dx/alpha
T = Ti.copy()
T = ftcs(T, nt, dt, dx, alpha)
```

And plot the solution!

```
pyplot.plot(numpy.linspace(0,1,nx), T, color='#003366', ls='-', lw=3)
pyplot.ylim(0,100)
pyplot.xlabel('Length of Rod')
pyplot.ylabel('Temperature');
```

## Boundary Conditions

In the rod problem above, we stated that the left-hand side of the rod is held to a temperature $T = 100{\rm C}$.  This is an example of a *boundary condition* (BC): a rule that defines how the problem interacts with the borders of the domain.  In this case, the domain spans the length of the rod, from $0 \leq x \leq 1$.

There are many types of boundary conditions, and they have an important effect on the solution of the problem. For example, in [module 2, lesson 4](http://nbviewer.ipython.org/urls/raw.github.com/numerical-mooc/numerical-mooc/master/lessons/02_spacetime/02_04_1DBurgers.ipynb) on Burgers' Equation, we used a periodic boundary condition to "connect" the right-hand and left-hand borders of the domain.

We need to discuss boundary conditions in a little more detail, and this is good place to do it. Read on!

### Dirichlet boundary

The rod example above uses a *Dirichlet* BC on the left-hand side.  A Dirichlet boundary is one in which the border is held to a specific value of the solution variable.

What about the right-hand end of the rod $(x=1)$?  In the discretization we set up above, the problem hasn't evolved for long enough time for the heat to travel the full length of the rod.  Let's increase the number of timesteps, `nt`, and see what happens.

```
nt = 1000
T = Ti.copy()
T = ftcs(T, nt, dt, dx, alpha)
```

```
pyplot.plot(numpy.linspace(0,1,nx), T, color='#003366', ls='-', lw=3)
pyplot.ylim(0,100)
pyplot.xlabel('Length of Rod')
pyplot.ylabel('Temperature');
```

Hmmm ... it looks like we're pinning the right BC to the value $T=0{\rm C}$, you agree?
Now study the code for the `ftbs` function, above, and try to figure out what happens at the right-most point of the spatial grid.

Did you figure it out?

It never updates the value at `T[-1]`!  That value is set to zero in the initial condition and stays that way throughout the entire simulation.  We effectively have a Dirichlet boundary at *both* ends.  The left-hand side of the rod is held to a temperature of $T = 100{\rm C}$ and the right-hand side is held to $T = 0{\rm C}$  Because both end tempera- tures are "pinned" to set values, the temperature distribution within the rod will relax (after enough time) to a linear temperature gradient across the length of the rod.

Once the solution is relaxed, no number of extra timesteps will make it change.  This is an example of a *steady-state solution* and we'll learn more about those in Module 5 of the course.

Dirichlet BCs show up in many engineering applications, among them thermo- dynamics (e.g., setting a surface temperature) and fluid dynamics (e.g., the no-slip conditions at walls in viscous fluids).

### Neumann boundary

Another commonly used BC is the *Neumann* boundary: rather than specifying the value of the solution at the border, they specify the value of the *derivative* of the solution at the border.

In our example, if we apply a Neumann boundary to the right-hand end of the rod, we can represent that mathematically as

\begin{equation}
\left. \frac{\partial T}{\partial x} \right|_{x = 1} = q(t)
\end{equation}

How can we enforce that in code?  One easy way is using a finite-difference discretization of the derivative at the boundary.  At time step $n$, for $N$ points, that would be

\begin{equation}
\left. \frac{\partial T}{\partial x} \right|_{x=1} \approx \frac{T^n_N - T^n_{N-1}}{\Delta x} = q(t)
\end{equation}

In the context of heat conduction, the space derivative of temperature is the heat flux density $q$, the amount of heat per unit time, per unit area.

For example, if the rod at $x=1$ has some insulating material, no heat is going to be able to get out through that end, and the Neumann boundary condition is:

\begin{equation}
\left. \frac{\partial T}{\partial x} \right|_{x=1} = 0
\end{equation}

Think about what the equation is saying: the change in temperature when moving in the $x$ direction is zero at the right-most edge of the rod. That means that the temperature should be equal on the last two spatial grid points.

To enforce this Neumann boundary condition at the right-hand end of the rod, we add the following line of code:

```Python
T[-1] = T[-2]
```

That is, the temperature at the boundary (`T[-1]`) is equal to the temperature directly to its left (`T[-2]`).  The spatial gradient is zero and a Neumann condition is satisfied there.

The Dirichlet conditions in the rod example were enforced without any extra input from us, but the Neumann condition will require an update at each iteration of the loop.

Finally, Dirichlet and Neumann BCs can sometimes be mixed to better represent the physics of a problem.

Let's revisit the heated-rod problem.  This time, the temperature at $x = 0$ will remain fixed at $T = 100{\rm C}$—that's the Dirichlet condition— and a Neumann BC applies at $x = 1$.

```
def ftcs_mixed(T, nt, dt, dx, alpha):
    """Solves the diffusion equation with forward-time, centered scheme using
    Dirichlet b.c. at left boundary and Neumann b.c. at right boundary

    Parameters:
    ----------
    u: array of float
        Initial temperature profile
    nt: int
        Number of time steps
    dt: float
        Time step size
    dx: float
        Mesh size
    alpha: float
        Diffusion coefficient (thermal diffusivity)

    Returns:
    -------
    u: array of float
        Temperature profile after nt time steps with forward in time scheme

    """
    for n in range(nt):
        Tn = T.copy()
        T[1:-1] = Tn[1:-1] + alpha*dt/dx**2*(Tn[2:] -2*Tn[1:-1] + Tn[0:-2])
        T[-1] = T[-2]

    return T
```

```
nt = 1000
T = Ti.copy()
T = ftcs_mixed(T, nt, dt, dx, alpha)
```

```
pyplot.plot(numpy.linspace(0,1,nx), T, color='#003366', ls='-', lw=3)
pyplot.ylim(0,100)
pyplot.title('Rod with LHS Dirichlet B.C. and RHS Neumann B.C.\n')
pyplot.xlabel('Length of Rod')
pyplot.ylabel('Temperature');
```

Now heat is accumulating in the domain. Our insulator works!  If you increase the number of timesteps `nt` further, you will see that the outflow temperature at $x = 1$ continues to increase.

### Explicit schemes and BCs

The numerical schemes we've learned so far are called _explicit_, because we compute the updated solution at $t^{n+1}$ using only known information at time $t^n$. This is simple but has limitations, particularly with the small sizes of time step necessary to obtain stability.

Here's another issue with explicit schemes. Figure 2 shows the superposed stencil of several grid points over three time steps computed using the forward-time, centered scheme. You know that to calculate $T_i^{n+1}$ you use the information from the grid points $i-1, i, i+1$ at the previous time step. Think about what happens at the boundary: any change in the boundary condition will feed into the solution only at the next time step, not immediately. But this contradicts the physics of the problem, as any change on the boundary should be felt right away under the diffusion equation. To include boundary effects on the same time level, we can use an _implicit_ update (subject of the next lesson).

.
#### Figure 2. Stencil with explicit scheme over 3 time steps.

### Explicit schemes and time step

Let's replace the parameter $\sigma=\alpha \frac{\Delta t}{(\Delta x)^2}$ into Equation 4 and rearrange it a little bit, grouping together the terms at grid point $i$ on the right-hand side:

\begin{equation}
T_{i}^{n+1} = \sigma T_{i-1}^{n}+(1- 2 \sigma) T_{i}^{n} + \sigma T_{i+1}^{n}
\end{equation}

It's helpful to look at a sketch of the stencil for the heat equation with the weights added in for the contribution of each grid point. See Figure 3, which shows the weights for two values of $\sigma$: $\frac{1}{2}$ and $1$. Notice that with $\sigma=\frac{1}{2}$, the solution variable at the next time step is _independent_ from its own value at the present step, which seems a bit weird. And with $\sigma=1$, the current value affects negatively the future value—now _that's_ a red flag!

In fact, the solution will develop growing errors with $\sigma>\frac{1}{2}$, i.e., become unstable. This limits the time step that can be used quite significantly, because $\Delta t \propto (\Delta x)^2$. This is a serious cost of _explicit_ methods and an incentive to consider _implicit_ alternatives (next lesson).

.
#### Figure 3. Stencils for the heat equation with $\sigma= \frac{1}{2}, 1$.

---
###### The cell below loads the style of the notebook

```
from IPython.core.display import HTML
css_file = '../../styles/numericalmoocstyle.css'
HTML(open(css_file, "r").read())
```

<link href='http://fonts.googleapis.com/css?family=Alegreya+Sans:100,300,400,500,700,800,900,100italic,300italic,400italic,500italic,700italic,800italic,900italic' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Arvo:400,700,400italic' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=PT+Mono' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Shadows+Into+Light' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Nixie+One' rel='stylesheet' type='text/css'>
<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: 750px;
}

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.8em;
}

/* 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: 'Alegreya Sans' sans-serif;
    line-height: 140%;
    font-size: 125%;
    font-weight: 400;
    width:600px;
    margin-left:auto;
    margin-right:auto;
}

/* Formatting for header cells */
.text_cell_render h1 {
    font-family: 'Nixie One', serif;
    font-style:regular;
    font-weight: 400;
    font-size: 45pt;
    line-height: 100%;
    color: rgb(0,51,102);
    margin-bottom: 0.5em;
    margin-top: 0.5em;
    display: block;
}
.text_cell_render h2 {
    font-family: 'Nixie One', serif;
    font-weight: 400;
    font-size: 30pt;
    line-height: 100%;
    color: rgb(0,51,102);
    margin-bottom: 0.1em;
    margin-top: 0.3em;
    display: block;
}

.text_cell_render h3 {
    font-family: 'Nixie One', serif;
    margin-top:16px;
	font-size: 22pt;
    font-weight: 600;
    margin-bottom: 3px;
    font-style: regular;
    color: rgb(102,102,0);
}

.text_cell_render h4 {    /*Use this for captions*/
    font-family: 'Nixie One', serif;
    font-size: 14pt;
    text-align: center;
    margin-top: 0em;
    margin-bottom: 2em;
    font-style: regular;
}

.text_cell_render h5 {  /*Use this for small titles*/
    font-family: 'Nixie One', sans-serif;
    font-weight: 400;
    font-size: 16pt;
    color: rgb(163,0,0);
    font-style: italic;
    margin-bottom: .1em;
    margin-top: 0.8em;
    display: block;
}

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

.CodeMirror{
        font-family: "PT Mono";
        font-size: 90%;
}

</style>


# Fault detection

We'll consider a problem of identifying faults that have occurred in a system based on sensor measurements of system performance.

# Topic references
- [Samar, Sikandar, Dimitry Gorinevsky, and Stephen Boyd. "Likelihood Bounds for Constrained Estimation with Uncertainty." Decision and Control, 2005 and 2005 European Control Conference. CDC-ECC'05. 44th IEEE Conference on. IEEE, 2005.](http://web.stanford.edu/~boyd/papers/pdf/map_bounds.pdf)

# Problem statement

Each of $n$ possible faults occurs independently with probability $p$.
The vector $x \in \lbrace 0,1 \rbrace^{n}$ encodes the fault occurrences, with $x_i = 1$ indicating that fault $i$ has occurred.
System performance is measured by $m$ sensors.
The sensor output is
\begin{equation}
y = Ax + v = \sum_{i=1}^n a_i x_i + v,
\end{equation}
where $A \in \mathbf{R}^{m \times n}$ is the sensing matrix with column $a_i$ being the **fault signature** of fault $i$,
and $v \in \mathbf{R}^m$ is a noise vector where $v_j$ is Gaussian with mean 0 and variance $\sigma^2$.

The objective is to guess $x$ (which faults have occurred) given $y$ (sensor measurements).

We are interested in the setting where $n > m$, that is, when we have more possible faults than measurements.
In this setting, we can expect a good recovery when the vector $x$ is sparse.
This is the subject of compressed sensing.

# Solution approach
To identify the faults, one reasonable approach is to choose $x \in \lbrace 0,1 \rbrace^{n}$ to minimize the negative log-likelihood function

\begin{equation}
\ell(x) = \frac{1}{2 \sigma^2} \|Ax-y\|_2^2 +  \log(1/p-1)\mathbf{1}^T x + c.
\end{equation}

However, this problem is nonconvex and NP-hard, due to the constraint that $x$ must be Boolean.

To make this problem tractable, we can relax the Boolean constraints and instead constrain $x_i \in [0,1]$.

The optimization problem

\begin{array}{ll}
\mbox{minimize} &  \|Ax-y\|_2^2 + 2 \sigma^2 \log(1/p-1)\mathbf{1}^T x\\
\mbox{subject to} &  0 \leq x_i \leq 1, \quad i=1, \ldots n
\end{array}

is convex.
We'll refer to the solution of the convex problem as the **relaxed ML** estimate.

By taking the relaxed ML estimate of $x$ and rounding the entries to the nearest of 0 or 1, we recover a Boolean estimate of the fault occurrences.

# Example

We'll generate an example with $n = 2000$ possible faults, $m = 200$ measurements, and fault probability $p = 0.01$.
We'll choose $\sigma^2$ so that the signal-to-noise ratio is 5.
That is,
\begin{equation}
\sqrt{\frac{\mathbf{E}\|Ax \|^2_2}{\mathbf{E} \|v\|_2^2}} = 5.
\end{equation}

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

np.random.seed(1)

n = 2000
m = 200
p = 0.01
snr = 5

sigma = np.sqrt(p*n/(snr**2))
A = np.random.randn(m,n)

x_true = (np.random.rand(n) <= p).astype(np.int)
v = sigma*np.random.randn(m)

y = A.dot(x_true) + v
```

Below, we show $x$, $Ax$ and the noise $v$.

```python
plt.plot(range(n),x_true)
```

```python
plt.plot(range(m), A.dot(x_true),range(m),v)
plt.legend(('Ax','v'))
```

# Recovery

We solve the relaxed maximum likelihood problem with CVXPY and then round the result to get a Boolean solution.

```python
%%time
import cvxpy as cp
x = cp.Variable(shape=n)
tau = 2*cp.log(1/p - 1)*sigma**2
obj = cp.Minimize(cp.sum_squares(A*x - y) + tau*cp.sum(x))
const = [0 <= x, x <= 1]
cp.Problem(obj,const).solve(verbose=True)
print("final objective value: {}".format(obj.value))

# relaxed ML estimate
x_rml = np.array(x.value).flatten()

# rounded solution
x_rnd = (x_rml >= .5).astype(int)
```

    ECOS 2.0.4 - (C) embotech GmbH, Zurich Switzerland, 2012-15. Web: www.embotech.com/ECOS

    It     pcost       dcost      gap   pres   dres    k/t    mu     step   sigma     IR    |   BT
     0  +7.343e+03  -3.862e+03  +5e+04  5e-01  5e-04  1e+00  1e+01    ---    ---    1  1  - |  -  -
     1  +4.814e+02  -9.580e+02  +8e+03  1e-01  6e-05  2e-01  2e+00  0.8500  1e-02   1  2  2 |  0  0
     2  -2.079e+02  -1.428e+03  +6e+03  1e-01  4e-05  8e-01  2e+00  0.7544  7e-01   2  2  2 |  0  0
     3  -1.321e+02  -1.030e+03  +5e+03  8e-02  3e-05  7e-01  1e+00  0.3122  2e-01   2  2  2 |  0  0
     4  -2.074e+02  -8.580e+02  +4e+03  6e-02  2e-05  6e-01  9e-01  0.7839  7e-01   2  2  2 |  0  0
     5  -1.121e+02  -6.072e+02  +3e+03  5e-02  1e-05  5e-01  7e-01  0.3859  4e-01   2  3  3 |  0  0
     6  -4.898e+01  -4.060e+02  +2e+03  3e-02  8e-06  3e-01  5e-01  0.5780  5e-01   2  2  2 |  0  0
     7  +7.778e+01  -5.711e+01  +8e+02  1e-02  3e-06  1e-01  2e-01  0.9890  4e-01   2  3  2 |  0  0
     8  +1.307e+02  +6.143e+01  +4e+02  6e-03  1e-06  6e-02  1e-01  0.5528  1e-01   3  3  3 |  0  0
     9  +1.607e+02  +1.286e+02  +2e+02  3e-03  4e-07  3e-02  5e-02  0.8303  3e-01   3  3  3 |  0  0
    10  +1.741e+02  +1.557e+02  +1e+02  2e-03  2e-07  2e-02  3e-02  0.6242  3e-01   3  3  3 |  0  0
    11  +1.834e+02  +1.749e+02  +5e+01  8e-04  9e-08  8e-03  1e-02  0.8043  3e-01   3  3  3 |  0  0
    12  +1.888e+02  +1.861e+02  +2e+01  3e-04  3e-08  2e-03  4e-03  0.9175  3e-01   3  3  2 |  0  0
    13  +1.909e+02  +1.902e+02  +4e+00  7e-05  7e-09  6e-04  1e-03  0.8198  1e-01   3  3  3 |  0  0
    14  +1.914e+02  +1.912e+02  +1e+00  2e-05  2e-09  2e-04  3e-04  0.8581  2e-01   3  2  3 |  0  0
    15  +1.916e+02  +1.916e+02  +1e-01  2e-06  3e-10  2e-05  4e-05  0.9004  3e-02   3  3  3 |  0  0
    16  +1.916e+02  +1.916e+02  +4e-02  7e-07  8e-11  7e-06  1e-05  0.8174  1e-01   3  3  3 |  0  0
    17  +1.916e+02  +1.916e+02  +8e-03  1e-07  1e-11  1e-06  2e-06  0.8917  9e-02   3  2  2 |  0  0
    18  +1.916e+02  +1.916e+02  +2e-03  4e-08  4e-12  4e-07  5e-07  0.8588  2e-01   3  3  3 |  0  0
    19  +1.916e+02  +1.916e+02  +2e-04  3e-09  3e-13  3e-08  5e-08  0.9309  2e-02   3  2  2 |  0  0
    20  +1.916e+02  +1.916e+02  +2e-05  4e-10  4e-14  4e-09  6e-09  0.8768  1e-02   4  2  2 |  0  0
    21  +1.916e+02  +1.916e+02  +4e-06  6e-11  6e-15  6e-10  9e-10  0.9089  6e-02   4  2  2 |  0  0
    22  +1.916e+02  +1.916e+02  +1e-06  2e-11  2e-15  2e-10  2e-10  0.8362  1e-01   2  1  1 |  0  0

    OPTIMAL (within feastol=1.8e-11, reltol=5.1e-09, abstol=9.8e-07).
    Runtime: 6.538894 seconds.

    final objective value: 191.6347201927456
    CPU times: user 6.51 s, sys: 291 ms, total: 6.8 s
    Wall time: 7.5 s

# Evaluation

We define a function for computing the estimation errors, and a function for plotting $x$, the relaxed ML estimate, and the rounded solutions.

```python
import matplotlib

def errors(x_true, x, threshold=.5):
    '''Return estimation errors.

    Return the true number of faults, the number of false positives, and the number of false negatives.
    '''
    n = len(x_true)
    k = sum(x_true)
    false_pos = sum(np.logical_and(x_true < threshold, x >= threshold))
    false_neg = sum(np.logical_and(x_true >= threshold, x < threshold))
    return (k, false_pos, false_neg)

def plotXs(x_true, x_rml, x_rnd, filename=None):
    '''Plot true, relaxed ML, and rounded solutions.'''
    matplotlib.rcParams.update({'font.size': 14})
    xs = [x_true, x_rml, x_rnd]
    titles = ['x_true', 'x_rml', 'x_rnd']

    n = len(x_true)
    k = sum(x_true)

    fig, ax = plt.subplots(1, 3, sharex=True, sharey=True, figsize=(12, 3))

    for i,x in enumerate(xs):
            ax[i].plot(range(n), x)
            ax[i].set_title(titles[i])
            ax[i].set_ylim([0,1])

    if filename:
        fig.savefig(filename, bbox_inches='tight')

    return errors(x_true, x_rml,.5)
```

We see that out of 20 actual faults, the rounded solution gives perfect recovery with 0 false negatives and 0 false positives.

```python
plotXs(x_true, x_rml, x_rnd, 'fault.pdf')
```


# Computing the Lithostatic Pressure

This is about the last paragraph of chapter 7, after introducing how to make numerical integration

We define the lithostatic pressure as the vertical pressure at a specific depth due to the weight of the overlying rocks. The pressure applied by a resting rock mass (assuming adiabatic compression, hydrostatic equilibrium, and spherical symmetry) under the acceleration of gravity (Poirier, 2000) is related to the rock density by

$\begin{equation}
p(z) = p_{0} + \int_{z_1=0}^{z}\rho(z_1)g(z_1)dz_1,
\label{eq:lithopress}
\end{equation}$

where $p(z)$ is the pressure at depth $z$, $p_0$ is the pressure at the surface, $\rho(z_1)$ is the bulk density for the rock mass as a function of depth, and $g(z_1)$ is the acceleration due to gravity.

As a zeroth-order approximation, we assume $p_0=0$ and that both $\rho(z)$ and $g(z)$ are constant, which reduces the above equation to

$\begin{equation}
p(z) = \rho g z.
\label{eq:simple_lithopress}
\end{equation}$

The following code listing shows its implementation in Python.

```python
def simple_lithopress(z,ro=2900, g=9.8):
    pressure = z*g*ro/1e6 # return the pressure in MPa
    return pressure

my_pressure = simple_lithopress(z=2000)
print('The pressure at 2000 meters is {0:.0f} MPa'.format(my_pressure))
```

    The pressure at 2000 meters is 57 MPa

| Layer      | $r$ from
 $km$ | $r$ to
 $km$ | Thickness
$km$ | Bottom density
$kg/m^3$   | Top density
$kg/m^3$   |
|:----------:|:-------------:|:-----------:|:--------------:|:-------------------------:|:----------------------:|
|Inner core  | 1             | 1220        | 1220           | 13100                     | 12800                  |
|Outer core  | 1221          | 3479        | 2259           | 12200                     | 9900                   |
|Lower mantle| 3480          | 5650        | 2171           | 5600                      | 4400                   |
|Upper mantle| 5651          | 6370        | 720            | 4130                      | 3400                   |
|Crust       | 6371          | 6400        | 30             | 3100                      | 2700                   |
|            |               |             |                |                           |                        |

We now take on the full implementation of the equation describing the lithostatic pressure. Using the data from Dziewonski and Anderson (1981) and Anderson (1989) and assuming that $\rho$ varies linearly between the upper and lower limits of each shell (i.e., crust, upper mantle, lower mantle, outer core, and inner core), we create an array that constitutes a first-order approximation of $\rho(z)$ (see the above Table). To simplify the presentation of the data, we define a new variable  $r$ (the radial distance from the Earth's center) as $r = R - z$, where $R\approx 6400~km$ is the Earth's radius and the Earth is approximated as a sphere.

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

r = np.linspace(1,6400,6400)

def density():
    ro_inner_core = np.linspace(13100, 12800, 1220)
    ro_outer_core = np.linspace(12200, 9900, 2259)
    ro_lower_mantle = np.linspace(5600,4400,2171)
    ro_upper_mantle = np.linspace(4130,3400,720)
    ro_crust = np.linspace(3100,2700,30)

    ro_final = np.concatenate((ro_inner_core, ro_outer_core, ro_lower_mantle, ro_upper_mantle, ro_crust))

    return ro_final

ro = density()

fig, ax = plt.subplots(figsize=(15,10))
ax.plot(r,ro, label="Densities on the Earth's Interior")
ax.set_ylabel(r"Density [Kg/m$^3$]")
ax.set_xlabel("Distance from the Earth center r [km]")
ax.legend()
ax.grid()
```

The acceleration $g(r)$ at a distance $r$ from Earth's center (Poirier,2000) is estimated  by using

$\begin{equation}
g(r) = \frac {4 \pi G}{r^2} \int_{r_1=0}^{r}\rho (r_1) r_1^2dr_1,
\end{equation}$

where $G=(6.67408 \pm 0.0031) \times 10^{-11}\ \text{m}^3\,\text{kg}^{-1}\, \text{s}^{-2}$ is the ''universal gravitational constant.''

```python
def gravity(r):

    g = np.zeros(len(r))
    Gr = 6.67408e-11
    r = r * 1000 # from Km to m

    for i in range(1,len(r)):

        r1 = r[0:i]
        ro1  = ro[0:i]
        r2 = r1[i-1]

        y = ro1*r1**2
        y_int = trapz(y,r1)

        g1 = ((4 * np.pi*Gr)/(r2**2)) * y_int
        g[i] = g1

    return g

g = gravity(r)

fig, ax = plt.subplots(figsize=(15,10))
ax.plot(r,g)
ax.grid()
ax.set_ylabel(r'g [m/s$^2]$')
ax.set_xlabel('Distance from the Earth center r [km]')
```

```python
def pressure(r, ro, g):

    p = np.zeros(len(r))
    r = r *1000

    for i in range(0,len(r)):
        r1 = r[i:len(r)]
        ro1 = ro[i:len(r)]
        g1 = g[i:len(r)]
        y = ro1*g1
        p1 = trapz(y,r1)
        p[i] = p1
    return p

p = pressure(r,ro,g)/1e9 # expressed in GPa
z = np.linspace(6400, 1, 6400)

fig, ax = plt.subplots(figsize=(15,10))
ax.plot(z,p)
ax.grid()
ax.set_ylabel('P [GPa]')
ax.set_xlabel('Depth z from the Earth Surface [km]')
```

### References:

Anderson, D. L. (1989).Theory of the Earth. Blackwell Scientific Publications.

Dziewonski, A. M., & Anderson, D. L. (1981). Preliminary reference Earth model.Physics of the Earth and Planetary Interiors,25(4), 297–356. https://doi.org/10.1016/0031-9201(81)90046-7

Poirier, J. (2000).Introduction to the physics of the earth’s interior(2ed.). CambridgeUniversity Press.

```python

```


# Dropout and maxout
In this lab we will explore the methods of [dropout](https://www.cs.toronto.edu/~hinton/absps/JMLRdropout.pdf), a regularisation method which stochastically drops out activations from the model during training, and [maxout](http://www.jmlr.org/proceedings/papers/v28/goodfellow13.pdf), another non-linear transformation that can be used in multiple layer models. This is based on material covered in the [fifth lecture slides](http://www.inf.ed.ac.uk/teaching/courses/mlp/2016/mlp05-hid.pdf).

## Exercise 1: Implementing a dropout layer

During training the forward propagation through a dropout layer produces outputs where a subset of the input dimensions are set to zero ('dropped out'). The dimensions to be dropped out are randomly sampled for each new batch, with each dimension having a probability $p$ of being included and the inclusion (or not) of each dimension independent of all the others. If the inputs to a dropout layer are $D$ dimensional vectors then we can represent the dropout operation by an elementwise multiplication by a $D$ dimensional *binary mask* vector $\boldsymbol{m} = \left[m_1 ~ m_2 ~\dots~ m_D\right]^{\rm T}$ where $m_d \sim \text{Bernoulli}(p) ~~\forall d \in \lbrace 1 \dots D\rbrace$.

As a first step implement a `random_binary_mask` function in the cell below to generate a binary mask array of a specified shape, where each value in the outputted array is either a one with probablity `prob_1` or zero with probability `1 - prob_1` and all values are sampled independently.

```python
def random_binary_mask(prob_1, shape, rng):
    """Generates a random binary mask array of a given shape.

    Each value in the outputted array should be an indepedently sampled
    binary value i.e in {0, 1} with the probability of each value
    being 1 being equal to `prob_1`.

    Args:
        prob_1: Scalar value in [0, 1] specifying probability each
            entry in output array is equal to one.
        shape: Shape of returned mask array.
        rng (RandomState): Seeded random number generator object.

    Returns:
        Random binary mask array of specified shape.
    """
    return rng.uniform(size=shape) < prob_1
```

Test your `random_binary_mask` function using the cell below (if your implementation is incorrect you will get an `AssertionError` - look at what the assert statement is checking for a clue as to what is wrong).

```python
import numpy as np
test_shapes = [(1, 1000), (10, 10, 10)]
test_probs = [0.1, 0.5, 0.7]
for i in range(10):
    for shape in test_shapes:
        for prob in test_probs:
            output = random_binary_mask(prob, shape, np.random)
            # Check generating correct shape output
            assert output.shape == shape
            # Check all outputs are binary values
            assert np.all((output == 1.) | (output == 0.))
            # Check proportion equal to one plausible
            # This will be noisy so there is a chance this will error
            # even for a correct implementation
            assert np.abs(output.mean() - prob) < 0.1
```

Given a randomly sampled binary mask $\boldsymbol{m}$, the outputs $\lbrace \boldsymbol{y}^{(b)} \rbrace_{b=1}^B$ of the stochastic forward propagation through a dropout layer given a batch of inputs $\lbrace \boldsymbol{x}^{(b)} \rbrace_{b=1}^B$ can be calculated by simply performing an elementwise multiplication of the inputs with the mask

\begin{equation}
  y^{(b)}_d = m_k x^{(b)}_d \qquad \forall d \in \lbrace 1 \dots D \rbrace
\end{equation}

The corresponding partial derivatives required for implementing back-propagation through a dropout layer are

\begin{equation}
  \frac{\partial y^{(b)}_k}{\partial x^{(b)}_d} =
  \begin{cases}
      m_k & \quad k = d \\
      0   & \quad k \neq d
  \end{cases}
  \qquad \forall k,\,d \in \lbrace 1 \dots D \rbrace
\end{equation}

As discussed in the lecture slides, when using a model trained with dropout at test time dimensions are no longer stochastically dropped out and instead all activations are deterministically fed forward through the model. So that the expected (mean) outputs of each layer are the same at test and training we scale the forward propagated inputs during testing by $p$ the probability of each dimension being included in the output. If we denote the deterministically forward-propagated batch of outputs of a dropout layer at test time as $\lbrace \boldsymbol{z}^{(b)} \rbrace_{b=1}^B$ then we have

\begin{equation}
  z^{(b)}_d =
  \mathbb{E}\left[ y^{(b)}_d \right] =
  \sum_{m_d \in \lbrace 0,1 \rbrace} \left( \mathbb{P}\left[\mathrm{m}_d = m_d\right] m_d x^{(b)}_d \right) =
  (p) (1) x^{(b)}_d  + (1-p) (0) x^{(b)}_d =
  p x^{(b)}_d  \qquad \forall d \in \lbrace 1 \dots D \rbrace
\end{equation}

To allow switching between this stochastic training time behaviour and deterministic test time behaviour, a new abstract `StochasticLayer` class has been defined in the `mlp.layers` module. This acts similarly to the layer objects we have already encountered other than adding an extra boolean argument `stochastic` to the `fprop` method interface. When `stochastic = True` (the default) a stochastic forward propagation should be caculated, for dropout this corresponding to $\boldsymbol{x}^{(b)} \to \boldsymbol{y}^{(b)}$ above. When `stochastic = False` a deterministic forward-propagation corresponding to the expected output of the stochastic forward-propagation should be calculated, for dropout this corresponding to $\boldsymbol{x}^{(b)} \to \boldsymbol{z}^{(b)}$ above.

Using the skeleton `DropoutLayer` class definition below, implement the `fprop` and `bprop` methods. You may wish to store the binary mask used in the forward propagation as an attribute of the class for use in back-propagation - it is fine to assume that the `fprop` and `bprop` will always be called in sync.

```python
from mlp.layers import StochasticLayer

class DropoutLayer(StochasticLayer):
    """Layer which stochastically drops input dimensions in its output."""

    def __init__(self, rng=None, incl_prob=0.5, share_across_batch=True):
        """Construct a new dropout layer.

        Args:
            rng (RandomState): Seeded random number generator.
            incl_prob: Scalar value in (0, 1] specifying the probability of
                each input dimension being included in the output.
            share_across_batch: Whether to use same dropout mask across
                all inputs in a batch or use per input masks.
        """
        super(DropoutLayer, self).__init__(rng)
        assert incl_prob > 0. and incl_prob <= 1.
        self.incl_prob = incl_prob
        self.share_across_batch = share_across_batch

    def fprop(self, inputs, stochastic=True):
        """Forward propagates activations through the layer transformation.

        Args:
            inputs: Array of layer inputs of shape (batch_size, input_dim).
            stochastic: Flag allowing different deterministic
                forward-propagation mode in addition to default stochastic
                forward-propagation e.g. for use at test time. If False
                a deterministic forward-propagation transformation
                corresponding to the expected output of the stochastic
                forward-propagation is applied.

        Returns:
            outputs: Array of layer outputs of shape (batch_size, output_dim).
        """
        if stochastic:
            mask_shape = (1,) + inputs.shape[1:] if self.share_across_batch else inputs.shape
            self._mask = (rng.uniform(size=mask_shape) < self.incl_prob)
            return inputs * self._mask
        else:
            return inputs * self.incl_prob

    def bprop(self, inputs, outputs, grads_wrt_outputs):
        """Back propagates gradients through a layer.

        Given gradients with respect to the outputs of the layer calculates the
        gradients with respect to the layer inputs. This should correspond to
        default stochastic forward-propagation.

        Args:
            inputs: Array of layer inputs of shape (batch_size, input_dim).
            outputs: Array of layer outputs calculated in forward pass of
                shape (batch_size, output_dim).
            grads_wrt_outputs: Array of gradients with respect to the layer
                outputs of shape (batch_size, output_dim).

        Returns:
            Array of gradients with respect to the layer inputs of shape
            (batch_size, input_dim).
        """
        return grads_wrt_outputs * self._mask

    def __repr__(self):
        return 'DropoutLayer(incl_prob={0:.1f})'.format(self.incl_prob)
```

Test your implementation by running the cell below (if your implementation is incorrect you will get an `AssertionError` - look at what the assert statement is checking for a clue as to what is wrong).

```python
seed = 31102016
rng = np.random.RandomState(seed)
test_incl_probs = [0.1, 0.5, 0.7]
input_shape = (5, 10)
for incl_prob in test_incl_probs:
    layer = DropoutLayer(rng, incl_prob)
    inputs = rng.normal(size=input_shape)
    grads_wrt_outputs = rng.normal(size=input_shape)
    for t in range(100):
        outputs = layer.fprop(inputs, stochastic=True)
        # Check outputted array correct shape
        assert outputs.shape == inputs.shape
        # Check all outputs are either equal to inputs or zero
        assert np.all((outputs == inputs) | (outputs == 0))
        grads_wrt_inputs = layer.bprop(inputs, outputs, grads_wrt_outputs)
        # Check back-propagated gradients only non-zero for non-zero outputs
        assert np.all((outputs != 0) == (grads_wrt_inputs != 0))
        assert np.all(grads_wrt_outputs[outputs != 0] == grads_wrt_inputs[outputs != 0])
    det_outputs = layer.fprop(inputs, stochastic=False)
    # Check deterministic fprop outputs are correct shape
    assert det_outputs.shape == inputs.shape
    # Check deterministic fprop outputs scaled correctly
    assert np.allclose(det_outputs, incl_prob * inputs)
```

### Optional extension

Above we assumed the same dropout mask was applied to each input in a batch, as specified in the lecture slides. In practice sometimes a different mask is sampled for each input. As an extension you could try implementing this per-input form of dropout either by defining a new layer or adding an extra argument to the constructor of the above layer which allows you to switch between the two forms.

## Exercise 2: Training with dropout

Experiment with training models with dropout layers to classify MNIST digits. Code has been provided below as a starting point for setting up the model objects though feel free to use any additional adaptive learning rules or learning rule schedulers you wrote during the coursework instead. You may also wish to change the model architecture to use a larger model with more parameters in which the regularisation provided by dropout is likely to have a more pronounced effect. You will probably also find that models with dropout generally need to be trained over more epochs than those without (can you suggest why this might be?).

You should training with a few different `incl_prob` settings for the dropout layers and try to establish how the values chosen affect the training performance. You may wish to experiment with using a different dropout probability at the input than for the intermediate layers (why?).

You may wish to start reading through and implementing exercise 3 while waiting for training runs to complete.

```python
import numpy as np
import logging
from mlp.data_providers import MNISTDataProvider
from mlp.models import MultipleLayerModel
from mlp.layers import ReluLayer, AffineLayer
from mlp.errors import CrossEntropySoftmaxError
from mlp.initialisers import GlorotUniformInit, ConstantInit
from mlp.learning_rules import MomentumLearningRule
from mlp.optimisers import Optimiser
import matplotlib.pyplot as plt
%matplotlib inline

# Seed a random number generator
seed = 31102016
rng = np.random.RandomState(seed)

# Set up a logger object to print info about the training run to stdout
logger = logging.getLogger()
logger.setLevel(logging.INFO)
logger.handlers = [logging.StreamHandler()]

# Create data provider objects for the MNIST data set
train_data = MNISTDataProvider('train', batch_size=50, rng=rng)
valid_data = MNISTDataProvider('valid', batch_size=50, rng=rng)
```

```python
# Probability of input being included in output in dropout layer
incl_prob = 0.5

input_dim, output_dim, hidden_dim = 784, 10, 125

# Use Glorot initialisation scheme for weights and zero biases
weights_init = GlorotUniformInit(rng=rng, gain=2.**0.5)
biases_init = ConstantInit(0.)

# Create three affine layer model with rectified linear non-linearities
# and dropout layers before every affine layer
model = MultipleLayerModel([
    DropoutLayer(rng, incl_prob),
    AffineLayer(input_dim, hidden_dim, weights_init, biases_init),
    ReluLayer(),
    DropoutLayer(rng, incl_prob),
    AffineLayer(hidden_dim, hidden_dim, weights_init, biases_init),
    ReluLayer(),
    DropoutLayer(rng, incl_prob),
    AffineLayer(hidden_dim, output_dim, weights_init, biases_init)
])

# Multiclass classification therefore use cross-entropy + softmax error
error = CrossEntropySoftmaxError()

# Use a momentum learning rule - you could use an adaptive learning rule
# implemented for the coursework here instead
learning_rule = MomentumLearningRule(0.02, 0.9)

# Monitor classification accuracy during training
data_monitors={'acc': lambda y, t: (y.argmax(-1) == t.argmax(-1)).mean()}

optimiser = Optimiser(
    model, error, learning_rule, train_data, valid_data, data_monitors)

num_epochs = 100
stats_interval = 5

stats, keys, run_time = optimiser.train(num_epochs=num_epochs, stats_interval=stats_interval)

# Plot the change in the validation and training set error over training.
fig_1 = plt.figure(figsize=(8, 4))
ax_1 = fig_1.add_subplot(111)
for k in ['error(train)', 'error(valid)']:
    ax_1.plot(np.arange(1, stats.shape[0]) * stats_interval,
              stats[1:, keys[k]], label=k)
ax_1.legend(loc=0)
ax_1.set_xlabel('Epoch number')

# Plot the change in the validation and training set accuracy over training.
fig_2 = plt.figure(figsize=(8, 4))
ax_2 = fig_2.add_subplot(111)
for k in ['acc(train)', 'acc(valid)']:
    ax_2.plot(np.arange(1, stats.shape[0]) * stats_interval,
              stats[1:, keys[k]], label=k)
ax_2.legend(loc=0)
ax_2.set_xlabel('Epoch number')
```

## Exercise 3: Implementing maxout

[Maxout](http://www.jmlr.org/proceedings/papers/v28/goodfellow13.pdf) can be considered a generalisation of the rectified linear transformation implemented in the previous lab.

For a rectified linear (`Relu`) layer the forward propagation corresponds to

\begin{equation}
    y^{(b)}_k =
    \max\left\lbrace 0,\,x^{(b)}_k \right\rbrace
\end{equation}

i.e. each output corresponds to an pairwise maximum of a constant (0) and the input.

Instead of taking the maximum of the input and a constant, we could instead consider taking the maximum over sets of inputs of a fixed size $s$.

\begin{equation}
    y^{(b)}_k =
    \max\left\lbrace x^{(b)}_{(k-1)s + 1},\, x^{(b)}_{(k-1)s + 2},\, \dots ,\, x^{(b)}_{ks} \right\rbrace
\end{equation}

If these inputs $x^{(b)}_d$ are themselves the outputs of an affine layer, then this corresponds to taking the maximum of a series of affine functions of the previous layer outputs. Like a rectified linear layer this leads to piecewise linear input-output relationships (which have well-behaved gradients which do not suffer from the saturation problems of logistic sigmoid / hyperbolic tangent transformations) but unlike the rectified linear case we do not end force a portion of the outputs to be zero.

Experimentally this form of transformation has been found to give good performance, with the name *maxout* chosen because the *out*put is the *max*imum of a set of inputs. Maxout is also commonly used with dropout layers however note they are not directly related - maxout defines a deterministic non-linear transformation which can help improve the representational capacity and trainability of models; dropout defines a stochastic transformation which is mainly aimed at regularising a model to reduce overfitting.

Using layers which take the maximum of fixed sized sets of inputs is also a common technique in models with convolutional layers which we will cover later in the course, with here the layer commonly being termed a *max-pooling* layer (with there being natural generalisation to other choices of reduction functions over pools such as the mean). We will adopt this terminology here for a layer implementing the transformation described above and we will be able to reuse our code implementing this maximum operation when experimenting with convolutional models.

The partial derivatives of this max-pooling transformation are sparse (lots of values are zero), with only the partial derivative of the output of a pool with respect to the maximum input in the pool non-zero. This can be expressed as

\begin{equation}
    \frac{\partial y^{(b)}_k}{\partial x^{(b)}_d} =
    \begin{cases}
      1 & \quad (k-1)s + 1 \leq d  \leq ks \quad\textrm{and} &x^{(b)}_d = \max\left\lbrace x^{(b)}_{(k-1)s + 1},\, x^{(b)}_{(k-1)s + 2},\, \dots ,\, x^{(b)}_{ks} \right\rbrace \\
      0 & \quad \textrm{otherwise}
    \end{cases}.
\end{equation}

Using these definitions implement the `fprop` and `bprop` methods of the skeleton `MaxPoolingLayer` class below.

Some hints

  * One way of organising the inputs into non-overlapping pools is using the `numpy.reshape` function.
  * The `numpy.max` function has an `axis` argument which allows you specify the axis (dimension) of the input array to take the maximum over.
  * It may help to construct a binary mask corresponding to the definitions of the partial derivatives above to allow you to implement the `bprop` method.
  * As with the `DropoutLayer` it is fine to temporarily store values calculated in the `fprop` method as attributes of the object (e.g. `self.val = val`) to use in the `bprop` method (although you don't necessarily need to do this).

```python
from mlp.layers import Layer

class MaxPoolingLayer(Layer):

    def __init__(self, pool_size=2):
        """Construct a new max-pooling layer.

        Args:
            pool_size: Positive integer specifying size of pools over
               which to take maximum value. The outputs of the layer
               feeding in to this layer must have a dimension which
               is a multiple of this pool size such that the outputs
               can be split in to pools with no dimensions left over.
        """
        self.pool_size = pool_size

    def fprop(self, inputs):
        """Forward propagates activations through the layer transformation.

        This corresponds to taking the maximum over non-overlapping pools of
        inputs of a fixed size `pool_size`.

        Args:
            inputs: Array of layer inputs of shape (batch_size, input_dim).

        Returns:
            outputs: Array of layer outputs of shape (batch_size, output_dim).
        """
        assert inputs.shape[-1] % self.pool_size == 0, (
            'Last dimension of inputs must be multiple of pool size')
        pooled_inputs = inputs.reshape(
            inputs.shape[:-1] +
            (inputs.shape[-1] // self.pool_size, self.pool_size))
        pool_maxes = pooled_inputs.max(-1)
        self._mask = pooled_inputs == pool_maxes[..., None]
        return pool_maxes

    def bprop(self, inputs, outputs, grads_wrt_outputs):
        """Back propagates gradients through a layer.

        Given gradients with respect to the outputs of the layer calculates the
        gradients with respect to the layer inputs.

        Args:
            inputs: Array of layer inputs of shape (batch_size, input_dim).
            outputs: Array of layer outputs calculated in forward pass of
                shape (batch_size, output_dim).
            grads_wrt_outputs: Array of gradients with respect to the layer
                outputs of shape (batch_size, output_dim).

        Returns:
            Array of gradients with respect to the layer inputs of shape
            (batch_size, input_dim).
        """
        return (self._mask * grads_wrt_outputs[..., None]).reshape(inputs.shape)

    def __repr__(self):
        return 'MaxPoolingLayer(pool_size={0})'.format(self.pool_size)
```

Test your implementation by running the cell below.

```python
test_inputs = np.array([[-3, -4, 5, 8], [0, -2, 3, -8], [1, 5, 3, 2]])
test_outputs_1 = np.array([[8], [3], [5]])
test_grads_wrt_outputs_1 = np.array([[10], [5], [-3]])
test_grads_wrt_inputs_1 = np.array([[0, 0, 0, 10], [0, 0, 5, 0], [0, -3, 0, 0]])
test_outputs_2 = np.array([[-3, 8], [0, 3], [5, 3]])
test_grads_wrt_outputs_2 = np.array([[3, -1], [2, 5], [5, 3]])
test_grads_wrt_inputs_2 = np.array([[3, 0, 0, -1], [2, 0, 5, 0], [0, 5, 3, 0]])
layer_1 = MaxPoolingLayer(4)
layer_2 = MaxPoolingLayer(2)
# Check fprop with pool_size = 4
assert np.allclose(layer_1.fprop(test_inputs), test_outputs_1)
# Check bprop with pool_size = 4
assert np.allclose(
    layer_1.bprop(test_inputs, test_outputs_1, test_grads_wrt_outputs_1),
    test_grads_wrt_inputs_1
)
# Check fprop with pool_size = 2
assert np.allclose(layer_2.fprop(test_inputs), test_outputs_2)
# Check bprop with pool_size = 2
assert np.allclose(
    layer_2.bprop(test_inputs, test_outputs_2, test_grads_wrt_outputs_2),
    test_grads_wrt_inputs_2
)
```

## Exercise 4: Training with maxout

Use your `MaxPoolingLayer` implementation in a multiple layer models to experiment with how well maxout networks are able to classify MNIST digits. As with the dropout training exercise, code has been provided below as a starting point for setting up the model objects, but again feel free to substitute any components.

If you have time you may wish to experiment with training a model using a combination of maxout and dropout or another regularisation method covered in the last lab notebook.

```python
import numpy as np
import logging
from mlp.data_providers import MNISTDataProvider
from mlp.models import MultipleLayerModel
from mlp.layers import AffineLayer
from mlp.errors import CrossEntropySoftmaxError
from mlp.initialisers import GlorotUniformInit, ConstantInit
from mlp.learning_rules import MomentumLearningRule
from mlp.optimisers import Optimiser
import matplotlib.pyplot as plt
%matplotlib inline

# Seed a random number generator
seed = 31102016
rng = np.random.RandomState(seed)

# Set up a logger object to print info about the training run to stdout
logger = logging.getLogger()
logger.setLevel(logging.INFO)
logger.handlers = [logging.StreamHandler()]

# Create data provider objects for the MNIST data set
train_data = MNISTDataProvider('train', batch_size=50, rng=rng)
valid_data = MNISTDataProvider('valid', batch_size=50, rng=rng)
```

```python
# Size of pools to take maximum over
pool_size = 2

input_dim, output_dim, hidden_dim = 784, 10, 100

# Use Glorot initialisation scheme for weights and zero biases
weights_init = GlorotUniformInit(rng=rng)
biases_init = ConstantInit(0.)

# Create three affine layer model interleaved with max-pooling layers
model = MultipleLayerModel([
    AffineLayer(input_dim, hidden_dim * pool_size, weights_init, biases_init),
    MaxPoolingLayer(pool_size),
    AffineLayer(hidden_dim, hidden_dim * pool_size, weights_init, biases_init),
    MaxPoolingLayer(pool_size),
    AffineLayer(hidden_dim, output_dim, weights_init, biases_init)
])

# Multiclass classification therefore use cross-entropy + softmax error
error = CrossEntropySoftmaxError()

# Use a momentum learning rule - you could use an adaptive learning rule
# implemented for the coursework here instead
learning_rule = MomentumLearningRule(0.02, 0.9)

# Monitor classification accuracy during training
data_monitors={'acc': lambda y, t: (y.argmax(-1) == t.argmax(-1)).mean()}

optimiser = Optimiser(
    model, error, learning_rule, train_data, valid_data, data_monitors)

num_epochs = 100
stats_interval = 5

stats, keys, run_time = optimiser.train(num_epochs=num_epochs, stats_interval=stats_interval)

# Plot the change in the validation and training set error over training.
fig_1 = plt.figure(figsize=(8, 4))
ax_1 = fig_1.add_subplot(111)
for k in ['error(train)', 'error(valid)']:
    ax_1.plot(np.arange(1, stats.shape[0]) * stats_interval,
              stats[1:, keys[k]], label=k)
ax_1.legend(loc=0)
ax_1.set_xlabel('Epoch number')

# Plot the change in the validation and training set accuracy over training.
fig_2 = plt.figure(figsize=(8, 4))
ax_2 = fig_2.add_subplot(111)
for k in ['acc(train)', 'acc(valid)']:
    ax_2.plot(np.arange(1, stats.shape[0]) * stats_interval,
              stats[1:, keys[k]], label=k)
ax_2.legend(loc=0)
ax_2.set_xlabel('Epoch number')
```


# Coherent dark states in a 3-level system
Investigating coherent dark states in 3-level system using the `toy_systems` package.

The 3-level system will consist of two ground states which are both coupled to the same excited state, which in turn can decay back to both of the ground states. A similar situation can be found in TlF when driving optical transitions from the electronic ground state X to the excited state B. The spacing of hyperfine levels in X is much smaller than the natural linewidth of the transitions, so that multiple groundstate hyperfine states end up coupling to one excited state hyperfine state. The number of coherent dark states in this scenario is N-1 where N is the number of ground states coupling to the same excited state. The dark states are typically not eigenstates of the internal Hamiltonian (e.g. TlF internal Hamiltonian) and so they evolve into bright states over time - we'll investigate the rate at which this occurs as the coupling strength of the ground states to the excited state and their energy splitting is varied.
## Imports
Start by importing the necessary packages

```python
%load_ext autoreload
%autoreload 2

import joblib
import matplotlib.pyplot as plt
plt.style.use("ggplot")
import numpy as np
import scipy
import qutip
from sympy import Symbol

from toy_systems.couplings import ToyCoupling, ToyEnergy
from toy_systems.dark_states import get_dark_states
from toy_systems.decays import ToyDecay
from toy_systems.hamiltonian import Hamiltonian
from toy_systems.quantum_system import QuantumSystem
from toy_systems.states import Basis, BasisState, ToyQuantumNumbers
from toy_systems.visualization import Visualizer
```

## Set up states and basis
We start by defining the three states of the system: the ground states $|g0\rangle$ and $|g1\rangle$, and the excited state $|e\rangle$>

```python
g0 = BasisState(qn=ToyQuantumNumbers(label="g0"))
g1 = BasisState(qn=ToyQuantumNumbers(label="g1"))
e = BasisState(qn=ToyQuantumNumbers(label="e"))

# Define basis
basis = Basis((g0,g1,e))
basis.print()
```

    |0> = |g0>
    |1> = |g1>
    |2> = |e>

## Define energies, couplings, decays and quantum system
I'm going to define the system in the rotating frame as usual.

### Energies

```python
δ = Symbol('delta') # Energy splitting between |g0> and |g1
Δ = Symbol('Delta') # Detuning of drive field from 0
E0 = ToyEnergy([g0], -δ/2)
E1 = ToyEnergy([g1], +δ/2)
Ee = ToyEnergy([e], Δ)
```

### Couplings

```python
Ω0 = Symbol('Omega_0') # Drive field Rabi rate for |g0> <-> |e>
Ω1 = Symbol('Omega_1') # Drive field Rabi rate for |g1> <-> |e>

coupling_0 = ToyCoupling(g0,e,Ω0)
coupling_1 = ToyCoupling(g1,e,Ω1)
```

### Decays
Defining a decay from $|e\rangle$ to $|g0\rangle$ and $|g1\rangle$:

```python
decay_0 = ToyDecay(e, ground = g0, gamma = Symbol("Gamma_0"))
decay_1 = ToyDecay(e, ground = g1, gamma = Symbol("Gamma_1"))
```

### Define a QuantumSystem
The QuantumSystem object combines the basis, Hamiltonian and decays to make setting parameters for time evolution using QuTiP more convenient.

```python
# Define the system
system = QuantumSystem(
    basis=basis,
    couplings=[E0, E1, Ee, coupling_0, coupling_1],
    decays=[decay_0, decay_1],
)
visualizer = Visualizer(system, vertical={"label":10}, horizontal={"label":50})

# Get representations of the Hamiltonian and the decays that will be accepted by qutip
Hqobj, c_qobj = system.get_qobjs()
```

## Find dark states
There is some rudimentary functionality in `toy_systems` for finding the bright and dark states when a number of ground states are coupled to a single excited state. Let's test that here:

```python
bright_state, dark_states, pol_dark_states = get_dark_states([g0, g1], e, [coupling_0, coupling_1])

print(f"|B> =\n{bright_state[0]}")

print(f"\n|D> =\n{dark_states[0]}")

print(f"\n<B|D> = {dark_states[0]@bright_state[0]}")
```

    |B> =
    0.71+0.00j x |g0>
    0.71+0.00j x |g1>

    |D> =
    0.71+0.00j x |g0>
    -0.71+0.00j x |g1>

    <B|D> = 0j

## Time evolution
No matter what state the system starts in, it should always end up in the dark state, from which it will slowly evolve out since the dark state is not an eigenstate of the Hamiltonian.

```python
# Get a pointer to the time-evolution arguments
args = Hqobj.args
print("Keys for setting arguments:")
print(f"args = {args}")
```

    Keys for setting arguments:
    args = {'delta': 1, 'Delta': 1, 'Omega_0': 1, 'Omega_1': 1, 'Gamma_0': 1, 'Gamma_1': 1}

```python
dark_states
```

    [0.71+0.00j x |g0>
     -0.71+0.00j x |g1>]

```python
bright_state
```

    [0.71+0.00j x |g0>
     0.71+0.00j x |g1>]

```python
# Set the parameters for the system
args['delta'] = 0.05
args['Delta'] = 0
args['Omega_0'] = 1
args['Omega_1'] = 1
args['Gamma_0'] = 1/2
args['Gamma_1'] = 1/2
Gamma = args['Gamma_0'] + args['Gamma_1']

# Get different dark and bright states for different values of Omegas so recalculate them
coupling_0.mag = args["Omega_0"]
coupling_1.mag = args["Omega_1"]
bright_states, dark_states, pol_dark_states = get_dark_states([g0, g1], e, [coupling_0, coupling_1])

# Generate a Qobj representing the initial state
psi0 = (1*dark_states[0]).qobj(basis)

# Operators for getting probability of being in each state as a function of time
P_g0 = qutip.Qobj((1*g0).density_matrix(basis), type = "oper")
P_g1 = qutip.Qobj((1*g1).density_matrix(basis), type = "oper")
P_e = qutip.Qobj((1*e).density_matrix(basis), type = "oper")
P_B = qutip.Qobj(bright_states[0].density_matrix(basis), type = "oper")
P_D = qutip.Qobj(dark_states[0].density_matrix(basis), type = "oper")
P_ops = [P_g0, P_g1, P_e, P_B, P_D]

# Times at which result is requested
times = np.linspace(0,100/args['delta'],1001)

# Setting the max_step is sometimes necessary
options = qutip.solver.Options(method = 'adams', nsteps=10000, max_step=1e0, rhs_reuse=True)

# Setup a progress bar
pb = qutip.ui.progressbar.EnhancedTextProgressBar()

# Run the time-evolution
result = qutip.mesolve(Hqobj, psi0, times, c_ops = c_qobj, e_ops = P_ops,
                      progress_bar=pb, options = options)
```

     Total run time:   0.79s*] Elapsed 0.79s / Remaining 00:00:00:00

```python
fig, ax = plt.subplots(2, 1, figsize = (16,9))
ax[0].plot(times, result.expect[0], label = "P_g0")
ax[0].plot(times, result.expect[1], label = "P_g1")
ax[0].plot(times, result.expect[2], label = "P_e")
ax[0].legend()
ax[0].set_ylabel("Population in each state")
ax[0].set_title("Energy eigenstate basis")

ax[1].plot(times, result.expect[3], label = "P_B")
ax[1].plot(times, result.expect[4], label = "P_D")
ax[1].plot(times, result.expect[2], label = "P_e")
ax[1].legend()
ax[1].set_ylabel("Population in each state")
ax[1].set_title("Dark and bright state basis")

print(f"|B> =\n{bright_state}")

print(f"\n|D> =\n{dark_states[0]}")

print(f"\nPopulation in excited state at the end: {result.expect[2][-1]*100:.1e} %")
print(f"Photons per unit time: {scipy.integrate.trapezoid(result.expect[2], x = times)/times[-1]:.2e}")
```

## Parameter scans

```python
default_args = {}
default_args['delta'] = 0.1
default_args['Delta'] = 0
default_args['Omega_0'] = 1
default_args['Omega_1'] = 1
default_args['Gamma_0'] = 1/2
default_args['Gamma_1'] = 1/2

# Define function that will be executed in parallel
def parallel_func(args):

    # Calculate total decay rate
    Gamma = args['Gamma_0'] + args['Gamma_1']

    # Get different dark and bright states for different values of Omegas so recalculate them
    coupling_0.mag = args["Omega_0"]
    coupling_1.mag = args["Omega_1"]
    bright_state, dark_states, pol_dark_states = get_dark_states([g0, g1], e, [coupling_0, coupling_1])

    # Generate a Qobj representing the initial state
    psi0 = (1*dark_states[0]).qobj(basis)

    # Operators for getting probability of being in each state as a function of time
    P_g0 = qutip.Qobj((1*g0).density_matrix(basis), type = "oper")
    P_g1 = qutip.Qobj((1*g1).density_matrix(basis), type = "oper")
    P_e = qutip.Qobj((1*e).density_matrix(basis), type = "oper")
    P_B = qutip.Qobj(bright_states[0].density_matrix(basis), type = "oper")
    P_D = qutip.Qobj(dark_states[0].density_matrix(basis), type = "oper")
    P_ops = [P_g0, P_g1, P_e, P_B, P_D]

    # Times at which result is requested
    times = np.linspace(0,100/args['delta'],1001)

    # Setting the max_step is sometimes necessary
    options = qutip.solver.Options(method = 'adams', nsteps=1000000, max_step=1e2/args["delta"], rhs_reuse=True)

    # Run the time-evolution
    result = qutip.mesolve(Hqobj, psi0, times, args = args, c_ops = c_qobj, e_ops = P_ops, options = options)

    # Calculate rate at which photons are emitted
    rate = scipy.integrate.trapezoid(result.expect[2], x = times)/times[-1]

    return rate
```

### Effect of energy splitting
Study the effect of the energy splitting of the ground states on the rate at which photons are emitted. Simple perturbation theory says that the rate at which photons are emitted is proportional to $\delta$.

For small values of $\delta$ (compared to what?) it seems that the dependence is more like quadratic.

```python
# Define values of energy splitting to use
δs = np.linspace(0.01, 10, 100)

# Generate list of arguments to pass to time evaluaation
args_list = [default_args.copy() | {'delta': δ} for δ in δs]

rates = np.array(joblib.Parallel(n_jobs = 10, backend='loky', verbose=True)(joblib.delayed(parallel_func)(args) for args in args_list))
```

    [Parallel(n_jobs=10)]: Using backend LokyBackend with 10 concurrent workers.
    [Parallel(n_jobs=10)]: Done  41 tasks      | elapsed:    0.4s
    [Parallel(n_jobs=10)]: Done 100 out of 100 | elapsed:    0.6s finished

```python
fig, ax = plt.subplots(figsize = (16,9))
ax.plot(δs, rates)
ax.set_xlabel("δ", fontsize = 16)
ax.set_ylabel("Photon emission rate", fontsize = 16)
```

### Quantum Zeno effect
For very large Rabi rates, the system should stay in the dark state, since we are effectively constantly measuring the state of the system.

```python
# Define values of energy splitting to use
Ωs = 10**np.linspace(-3, 1, 100)

# Generate list of arguments to pass to time evaluaation
args_list = [default_args.copy() | {'Omega_0': Ω, 'Omega_1': Ω} for Ω in Ωs]

# qutip.rhs_generate(Hqobj, c_qobj, args = default_args)
rates = np.array(joblib.Parallel(n_jobs = 10, backend='loky', verbose=True)(joblib.delayed(parallel_func)(args) for args in args_list))
```

    [Parallel(n_jobs=10)]: Using backend LokyBackend with 10 concurrent workers.
    [Parallel(n_jobs=10)]: Done  40 tasks      | elapsed:    0.3s
    [Parallel(n_jobs=10)]: Done 100 out of 100 | elapsed:    1.7s finished

```python
fig, ax = plt.subplots(figsize = (16,9))
ax.plot(Ωs, rates)
ax.set_xlabel("Ω", fontsize = 16)
ax.set_ylabel("Photon emission rate", fontsize = 16)
```

### Effect of different Rabi rates
In real systems the coupling strengths of the ground states to the excited state are typically different. Perturbation theory says that this should have a large effect on the rate at which the dark state evolves into the bright state (effect $\propto \left(\frac{\Omega_0\Omega_1}{\Omega_0^2+\Omega_1^2}\right)^2$)

```python
# Define values of energy splitting to use
Ωratios = 10**np.linspace(-3, 0, 100) # Defined as Omega_1/Omega_0

# Generate list of arguments to pass to time evaluaation
args_list = [default_args.copy() | {'Omega_1': Ω} for Ω in Ωratios]

qutip.rhs_clear()
rates = np.array(joblib.Parallel(n_jobs = 10, backend='loky', verbose=True)(joblib.delayed(parallel_func)(args) for args in args_list))
```

    [Parallel(n_jobs=10)]: Using backend LokyBackend with 10 concurrent workers.
    [Parallel(n_jobs=10)]: Done  40 tasks      | elapsed:    0.4s
    [Parallel(n_jobs=10)]: Done 100 out of 100 | elapsed:    1.0s finished

```python
fig, ax = plt.subplots(figsize = (16,9))
ax.plot(Ωratios, rates)
ax.set_xlabel("Ω", fontsize = 16)
ax.set_ylabel("Photon emission rate", fontsize = 16)
```

```python

```

examples/Coherent dark states in a 3-level system.ipynb


```python
import numpy as np
import numba
import matplotlib.pyplot as plt
import sympy as sym
plt.style.use('presentation')
%matplotlib notebook

colors_cycle=plt.rcParams.get('axes.prop_cycle')
colors = [item['color'] for item in colors_cycle]
def d2np(d):

    names = []
    numbers = ()
    dtypes = []
    for item in d:
        names += item
        if type(d[item]) == float:
            numbers += (d[item],)
            dtypes += [(item,float)]
        if type(d[item]) == int:
            numbers += (d[item],)
            dtypes += [(item,int)]
        if type(d[item]) == np.ndarray:
            numbers += (d[item],)
            dtypes += [(item,np.float64,d[item].shape)]
    return np.array([numbers],dtype=dtypes)
```

```python
psi_ds,psi_qs,psi_dr,psi_qr = sym.symbols('psi_ds,psi_qs,psi_dr,psi_qr')
i_ds,i_qs,i_dr,i_qr = sym.symbols('i_ds,i_qs,i_dr,i_qr')
di_ds,di_qs,di_dr,di_qr = sym.symbols('di_ds,di_qs,di_dr,di_qr')
L_s,L_r,L_m = sym.symbols('L_s,L_r,L_m')
R_s,R_r = sym.symbols('R_s,R_r')
omega_s,omega_r,sigma = sym.symbols('omega_s,omega_r,sigma')
v_ds,v_qs,v_dr,v_qr = sym.symbols('v_ds,v_qs,v_dr,v_qr')

eq_ds = (L_s+L_m)*i_ds + L_m*i_dr - psi_ds
eq_qs = (L_s+L_m)*i_qs + L_m*i_qr - psi_qs
eq_dr = (L_r+L_m)*i_dr + L_m*i_ds - psi_dr
eq_qr = (L_r+L_m)*i_qr + L_m*i_qs - psi_qr
dpsi_ds = v_ds - R_s*i_ds + omega_s*psi_qs
dpsi_qs = v_qs - R_s*i_qs - omega_s*psi_ds
dpsi_dr = v_dr - R_r*i_dr + sigma*omega_s*psi_qr
dpsi_qr = v_qr - R_r*i_qr - sigma*omega_s*psi_dr

'''
s = sym.solve([  eq_dr,  eq_qr, dpsi_ds, dpsi_qs, dpsi_dr, dpsi_qr],
              [   i_ds,   i_qs,  psi_ds, psi_qs,     v_dr, v_qr])
s = sym.solve([dpsi_ds,dpsi_qs,dpsi_dr,dpsi_qr],
              [   i_ds,   i_qs,   i_dr,   i_qr,
                psi_ds, psi_qs,   i_dr, psi_qr])
s = sym.solve([  eq_ds,  eq_qs,  eq_dr,  eq_qr,
               dpsi_ds,dpsi_qs,dpsi_dr,dpsi_qr],
              [   i_ds,   i_qs,   v_dr,   v_qr,
                psi_ds, psi_qs, psi_dr, psi_qr])

'''

s = sym.solve([    eq_dr,  eq_qr,
               dpsi_ds,dpsi_qs,dpsi_dr,dpsi_qr],
              [   i_ds,   i_qs,   v_dr,   v_qr,
                  psi_dr, psi_qr])

s = sym.solve([  eq_ds,  eq_qs,  eq_dr,  eq_qr,
               dpsi_ds,dpsi_qs,dpsi_dr,dpsi_qr],
              [   i_ds,   i_qs,   v_dr,   v_qr,
                psi_ds, psi_qs, psi_dr, psi_qr])

for item in s:
    print(item, '=', sym.simplify(s[item]))
```

    psi_dr = (L_m**2*R_s*i_qr*omega_s - L_m**2*i_dr*omega_s**2*(L_m + L_s) + L_m*R_s*v_ds + L_m*omega_s*v_qs*(L_m + L_s) + i_dr*(L_m + L_r)*(R_s**2 + omega_s**2*(L_m + L_s)**2))/(R_s**2 + omega_s**2*(L_m + L_s)**2)
    v_dr = (L_m**2*R_s*i_dr*omega_s**2*sigma + L_m**2*i_qr*omega_s**3*sigma*(L_m + L_s) - L_m*R_s*omega_s*sigma*v_qs + L_m*omega_s**2*sigma*v_ds*(L_m + L_s) + R_r*i_dr*(R_s**2 + omega_s**2*(L_m + L_s)**2) - i_qr*omega_s*sigma*(L_m + L_r)*(R_s**2 + omega_s**2*(L_m + L_s)**2))/(R_s**2 + omega_s**2*(L_m + L_s)**2)
    i_qs = (-L_m*R_s*i_dr*omega_s - L_m*i_qr*omega_s**2*(L_m + L_s) + R_s*v_qs - omega_s*v_ds*(L_m + L_s))/(R_s**2 + omega_s**2*(L_m + L_s)**2)
    i_ds = (L_m*R_s*i_qr*omega_s - L_m*i_dr*omega_s**2*(L_m + L_s) + R_s*v_ds + omega_s*v_qs*(L_m + L_s))/(R_s**2 + omega_s**2*(L_m + L_s)**2)
    psi_qs = (L_m*R_s**2*i_qr - L_m*R_s*i_dr*omega_s*(L_m + L_s) + R_s*v_qs*(L_m + L_s) - omega_s*v_ds*(L_m + L_s)**2)/(R_s**2 + omega_s**2*(L_m + L_s)**2)
    psi_qr = (-L_m**2*R_s*i_dr*omega_s - L_m**2*i_qr*omega_s**2*(L_m + L_s) + L_m*R_s*v_qs - L_m*omega_s*v_ds*(L_m + L_s) + i_qr*(L_m + L_r)*(R_s**2 + omega_s**2*(L_m + L_s)**2))/(R_s**2 + omega_s**2*(L_m + L_s)**2)
    psi_ds = (L_m*R_s**2*i_dr + L_m*R_s*i_qr*omega_s*(L_m + L_s) + R_s*v_ds*(L_m + L_s) + omega_s*v_qs*(L_m + L_s)**2)/(R_s**2 + omega_s**2*(L_m + L_s)**2)
    v_qr = (L_m**2*R_s*i_qr*omega_s**2*sigma - L_m**2*i_dr*omega_s**3*sigma*(L_m + L_s) + L_m*R_s*omega_s*sigma*v_ds + L_m*omega_s**2*sigma*v_qs*(L_m + L_s) + R_r*i_qr*(R_s**2 + omega_s**2*(L_m + L_s)**2) + i_dr*omega_s*sigma*(L_m + L_r)*(R_s**2 + omega_s**2*(L_m + L_s)**2))/(R_s**2 + omega_s**2*(L_m + L_s)**2)

```python
# [1] T. Demiray, F. Milano, and G. Andersson,
# “Dynamic phasor modeling of the doubly-fed induction generator under unbalanced conditions,” 2007 IEEE Lausanne POWERTECH, Proc., no. 2, pp. 1049–1054, 2007.

@numba.jit(nopython=True, cache=True)
def dfim_wecs(struct,i,m):
    '''

    '''

@numba.jit(nopython=True, cache=True)
def dfim_wecs_ctrl(struct,i,m):
    '''

    '''

    x_idx = struct[i]['dfim_wecs_ctrl_idx']

    L_m = struct[i]['L_m']
    L_r = struct[i]['L_r']
    L_s = struct[i]['L_s']
    R_r = struct[i]['R_r']
    R_s = struct[i]['R_s']
    N_pp = struct[i]['N_pp']

    Dt  = struct[i]['Dt']

    i_dr_ref = struct[i]['i_dr_ref']
    i_qr_ref = struct[i]['i_qr_ref']

    i_dr = i_dr_ref
    i_qr = i_qr_ref

    v_ds = struct[i]['v_ds']
    v_qs = struct[i]['v_qs']

    omega_r = struct[i]['omega_r']
    omega_s = struct[i]['omega_s']

    sigma = (omega_s - omega_r)/omega_s

    den = R_s**2 + omega_s**2*(L_m + L_s)**2
    i_qs = (-L_m*R_s*i_dr*omega_s - L_m*i_qr*omega_s**2*(L_m + L_s) + R_s*v_qs - omega_s*v_ds*(L_m + L_s))/den
    i_ds = ( L_m*R_s*i_qr*omega_s - L_m*i_dr*omega_s**2*(L_m + L_s) + R_s*v_ds + omega_s*v_qs*(L_m + L_s))/den

    v_qr = R_r*i_qr + omega_s*sigma*(L_m*i_dr + L_m*i_ds + L_r*i_dr)
    v_dr = R_r*i_dr - omega_s*sigma*(L_m*i_qr + L_m*i_qs + L_r*i_qr)
    psi_dr = L_m*i_dr + L_m*i_ds + L_r*i_dr
    psi_qs = (R_s*i_ds - v_ds)/omega_s
    psi_ds = (-R_s*i_qs + v_qs)/omega_s
    psi_qr = L_m*i_qr + L_m*i_qs + L_r*i_qr

    tau_e  = 3.0/2.0*N_pp*(psi_qr*i_dr - psi_dr*i_qr)

    struct[i]['v_dr'] = v_dr
    struct[i]['v_qr'] = v_qr

    struct[i]['i_ds'] = i_ds
    struct[i]['i_qs'] = i_qs
    struct[i]['i_dr'] = i_dr
    struct[i]['i_qr'] = i_qr

    struct[i]['psi_ds'] = psi_ds
    struct[i]['psi_qs'] = psi_qs
    struct[i]['psi_dr'] = psi_dr
    struct[i]['psi_qr'] = psi_qr

    struct[i]['tau_e'] = tau_e
    struct[i]['sigma'] = sigma

    struct[i]['p_s'] = 3.0/2.0*(v_ds*i_ds + v_qs*i_qs)
    struct[i]['q_s'] = 3.0/2.0*(v_ds*i_qs - v_qs*i_ds)

    struct[i]['p_r'] = 3.0/2.0*(v_dr*i_dr + v_qr*i_qr)
    struct[i]['q_r'] = 3.0/2.0*(v_dr*i_qr - v_qr*i_dr)

    return tau_e

@numba.jit(nopython=True, cache=True)
def vsc_grid_ctrl(struct,i,m):

    x_idx = struct[i]['mech_idx']
    omega_t = struct[i]['x'][x_idx,0]  # rad/s
    tau_t   = struct[i]['tau_t']
    tau_r   = struct[i]['tau_r']

    J_t  = struct[i]['J_t']
    N_tr = struct[i]['N_tr']
    Dt   = struct[i]['Dt']

    domega_t = 1.0/J_t*(tau_t - N_tr*tau_r)

    omega_r = N_tr*omega_t

    struct[i]['f'][x_idx,0] = domega_t

    struct[i]['omega_r'] = omega_r
    struct[i]['omega_t'] = omega_t

    return omega_t

@numba.jit(nopython=True, cache=True)
def vsc_rotor_ctrl(struct,i,m):

    x_idx = struct[i]['mech_idx']
    omega_t = struct[i]['x'][x_idx,0]  # rad/s
    tau_t   = struct[i]['tau_t']
    tau_r   = struct[i]['tau_r']

    J_t  = struct[i]['J_t']
    N_tr = struct[i]['N_tr']
    Dt   = struct[i]['Dt']

    domega_t = 1.0/J_t*(tau_t - N_tr*tau_r)

    omega_r = N_tr*omega_t

    struct[i]['f'][x_idx,0] = domega_t

    struct[i]['omega_r'] = omega_r
    struct[i]['omega_t'] = omega_t

    return omega_t

@numba.jit(nopython=True, cache=True)
def vsc_pitch_ctrl(struct,i,m):

    x_idx = struct[i]['mech_idx']
    omega_t = struct[i]['x'][x_idx,0]  # rad/s
    tau_t   = struct[i]['tau_t']
    tau_r   = struct[i]['tau_r']

    J_t  = struct[i]['J_t']
    N_tr = struct[i]['N_tr']
    Dt   = struct[i]['Dt']

    domega_t = 1.0/J_t*(tau_t - N_tr*tau_r)

    omega_r = N_tr*omega_t

    struct[i]['f'][x_idx,0] = domega_t

    struct[i]['omega_r'] = omega_r
    struct[i]['omega_t'] = omega_t

    return omega_t
```

```python
@numba.jit(nopython=True, cache=True)
def dfim_ctrl2(struct,i,m):
    '''
    Control level 2 for DFIM for stator active and reactive power.

    '''

    x_idx = struct[i]['ctrl2r_idx']
    xi_p_s = float(struct[i]['x'][x_idx+0,0])
    xi_q_s = float(struct[i]['x'][x_idx+1,0])

    K_r_p = struct[i]['K_r_p']
    K_r_i = struct[i]['K_r_i']

    p_s_ref = struct[i]['p_s_ref']
    q_s_ref = struct[i]['q_s_ref']

    p_s = struct[i]['p_s']
    q_s = struct[i]['q_s']

    S_b = struct[i]['S_b']

    omega_r = struct[i]['omega_r']
    omega_s = struct[i]['omega_s']
    R_r = struct[i]['R_r']

    I_b = S_b/(np.sqrt(3)*690.0)
    sigma = (omega_s - omega_r)/omega_s

    error_p_s = (p_s_ref - p_s)/S_b
    error_q_s = (q_s_ref - q_s)/S_b
    dxi_p_s = error_p_s
    dxi_q_s = error_q_s

    struct[i]['f'][x_idx+0,0] = dxi_p_s
    struct[i]['f'][x_idx+1,0] = dxi_q_s

    struct[i]['i_dr_ref'] = -I_b*(K_r_p*error_p_s + K_r_i*xi_p_s)
    struct[i]['i_qr_ref'] = -I_b*(K_r_p*error_q_s + K_r_i*xi_q_s)

    return struct[0]['i_dr_ref'],struct[0]['i_qr_ref']

```

```python
Omega_b = 2.0*np.pi*50.0
S_b = 2.0e6
U_b = 690.0
Z_b = U_b**2/S_b
#nu_w =np.linspace(0.1,15,N)
H = 2.0
N_pp = 2
N_tr = 20
# H = 0.5*J*Omega_t_n**2/S_b
S_b = 2.0e6
Omega_t_n = Omega_b/N_pp/N_tr
J_t = 2*H*S_b/Omega_t_n**2

#Z_b = 1.0
#Omega_b = 1.0
d =dict(S_b = S_b,
        Omega_b = Omega_b,
        R_r = 0.01*Z_b,
        R_s = 0.01*Z_b,
        L_r = 0.08*Z_b/Omega_b,
        L_s = 0.1*Z_b/Omega_b,
        L_m = 3.0*Z_b/Omega_b,
        N_pp = N_pp,
        psi_ds = 0.0,
        psi_qs = 0.0,
        p_s = 0.0,
        q_s = 0.0,
        p_r = 0.0,
        q_r = 0.0,
        psi_dr = 0.0,
        psi_qr = 0.0,
        p_s_ref = 0.0,
        q_s_ref = 0.0,
        i_ds   = 0.0,
        i_qs   = 0.0,
        i_dr   = 0.0,
        i_qr   = 0.0,
        i_dr_ref   = 0.0,
        i_qr_ref   = 0.0,
        v_ds   = 0.0,
        v_qs   = 0.0,
        v_dr   = 0.0,
        v_qr   = 0.0,
        omega_r   = Omega_b/N_pp,
        omega_s   = Omega_b/N_pp,
        sigma   = 0.0,
        tau_e   = 0.0,
        x = np.zeros((3,1)),
        f = np.zeros((3,1)),
        Dt = 0.0,
        J_t = J_t,
        omega_t = 0.0,
        tau_t = 0.0,
        tau_r = 0.0,
        N_tr = N_tr,
        K_r_p = 0.02,
        K_r_i = 20.0,
        dfim_idx = 0,
        mech_idx = 0,
        ctrl2r_idx = 1
        )

struct = d2np(d)
struct = np.hstack((struct[0],np.copy(struct[0])))

#wecs_mech_1(struct,0)

dfim_alg_ctrl1(struct,0,0)
dfim_ctrl2(struct,0,0)
dfim_alg_ctrl1(struct,1,0)
dfim_ctrl2(struct,1,0)
print(struct[0]['p_s']/1e6,struct[0]['q_s']/1e6,struct[0]['tau_e'])
print(struct[1]['p_s']/1e6,struct[0]['q_s']/1e6,struct[0]['tau_e'])
```

    0.0 0.0 0.0
    0.0 0.0 0.0

```python
struct = d2np(d)
struct = np.hstack((struct[0],np.copy(struct[0])))
sys_d = dict(x = np.zeros((6,1)),
             f = np.zeros((6,1)))

sys_struct = d2np(sys_d)

@numba.jit(nopython=True, cache=True)
def f_eval(sys_struct,struct):

    for i in range(2):
        struct[i]['x'][:,0] = sys_struct[0]['x'][3*i:3*(i+1),0]
        wecs_mech_1(struct,i,2)
        dfim_ctrl2(struct,i,2)
        dfim_alg_ctrl1(struct,i,2)
        sys_struct[0]['f'][3*i:3*(i+1),:] = struct[i]['f']
    return 0

```

```python
@numba.jit(nopython=True, cache=True)
def run(sys_struct,struct):
    N_steps = 1000
    N_states = 6

    Dt = 10.0e-3
    Omega_r = np.zeros((N_steps,1))
    Omega_t = np.zeros((N_steps,1))
    P_s_1 = np.zeros((N_steps,1))
    Q_s_1 = np.zeros((N_steps,1))
    P_r_1 = np.zeros((N_steps,1))
    Q_r_1 = np.zeros((N_steps,1))
    P_s_2 = np.zeros((N_steps,1))
    Q_s_2 = np.zeros((N_steps,1))
    P_r_2 = np.zeros((N_steps,1))
    Q_r_2 = np.zeros((N_steps,1))
    V_dr = np.zeros((N_steps,1))
    V_qr = np.zeros((N_steps,1))
    I_dr = np.zeros((N_steps,1))
    I_qr = np.zeros((N_steps,1))
    I_ds = np.zeros((N_steps,1))
    I_qs = np.zeros((N_steps,1))
    Tau_e = np.zeros((N_steps,1))
    Tau_t = np.zeros((N_steps,1))
    T     = np.zeros((N_steps,1))
    X = np.zeros((N_steps,N_states))
    p_ref = 0.0
    q_ref = 0.0
    xi_p = 0.0
    xi_q = 0.0

    struct[0]['x'][:,0] = np.copy(sys_struct[0]['x'][0:3,0])
    struct[1]['x'][:,0] = np.copy(sys_struct[0]['x'][3:6,0])

    for it in range(N_steps):
        t = Dt*float(it)

        # perturbations and references

        struct[0]['p_s_ref'] = 0.0
        struct[1]['p_s_ref'] = 0.0
        struct[0]['q_s_ref'] = 0.0
        struct[1]['q_s_ref'] = 0.0

        if t>1.0:
            Omega_t_b = struct[0]['Omega_b']/struct[0]['N_tr']/struct[0]['N_pp']
            struct[0]['tau_t'] = 1.5e6/Omega_t_b+np.random.normal(500e3,100e3)/Omega_t_b
        if t>1.5:
            struct[0]['p_s_ref'] = 1.0e6
        if t>3.0:
            struct[1]['p_s_ref'] = 1.50e6
        if t>4.0:
            struct[0]['q_s_ref'] = 0.5e6
        if t>5.0:
            struct[1]['q_s_ref'] = -0.7e6
        ## solver

        f_eval(sys_struct,struct)
        f1 = np.copy(sys_struct[0]['f'])
        x1 = np.copy(sys_struct[0]['x'])

        sys_struct[0]['x'][:]= np.copy(x1 + Dt*f1)
        f_eval(sys_struct,struct)
        f2 = np.copy(sys_struct[0]['f'])

        sys_struct[0]['x'][:]= np.copy(x1 + 0.5*Dt*(f1 + f2))

        for i in range(2):
            struct[i]['x'][:,0] = sys_struct[0]['x'][3*i:3*(i+1),0]

        struct[0]['tau_r'] = struct[0]['tau_e']
        struct[1]['tau_r'] = struct[1]['tau_e']

        T[it,0] = t
        P_s_1[it,0] = float(struct[0]['p_s'])
        Q_s_1[it,0] = float(struct[0]['q_s'])
        P_r_1[it,0] = float(struct[0]['p_r'])
        Q_r_1[it,0] = float(struct[0]['q_r'])
        P_s_2[it,0] = float(struct[1]['p_s'])
        Q_s_2[it,0] = float(struct[1]['q_s'])
        P_r_2[it,0] = float(struct[1]['p_r'])
        Q_r_2[it,0] = float(struct[1]['q_r'])

        I_dr[it,0] = float(struct[0]['i_dr'])
        I_qr[it,0] = float(struct[0]['i_qr'])
        I_ds[it,0] = float(struct[0]['i_ds'])
        I_qs[it,0] = float(struct[0]['i_qs'])
        Omega_r[it,0] = float(struct[0]['omega_r'])
        Omega_t[it,0] = float(struct[0]['omega_t'])

        V_dr[it,0] = float(struct[0]['v_dr'])
        V_qr[it,0] = float(struct[0]['v_qr'])
        Tau_e[it,0] = float(struct[0]['tau_e'])
        Tau_t[it,0] = float(struct[0]['tau_t'])
        X[it,:] = sys_struct[0]['x'][:].T
    return T,X,Tau_e,P_s_1,Q_s_1,P_r_1,Q_r_1,P_s_2,Q_s_2,P_r_2,Q_r_2,V_dr,V_qr,Omega_r,Omega_t,I_dr,I_qr,I_ds,I_qs,Tau_t
%timeit run(sys_struct, struct)
```

    The slowest run took 5.81 times longer than the fastest. This could mean that an intermediate result is being cached.
    1000 loops, best of 3: 1.33 ms per loop

```python
sys_struct['x'][:]= np.zeros((6,1))
struct['v_qs'] = 0.0
struct['v_ds'] = 690.0*np.sqrt(2.0/3.0)
struct['tau_t'] = 0.0
sys_struct[0]['x'][0,0] = Omega_b*0.9/struct[0]['N_tr']/struct[0]['N_pp']
sys_struct[0]['x'][3,0] = Omega_b*1.1/struct[1]['N_tr']/struct[0]['N_pp']
T,X,Tau_e,P_s_1,Q_s_1,P_r_1,Q_r_1,P_s_2,Q_s_2,P_r_2,Q_r_2,V_dr,V_qr,Omega_r,Omega_t,I_dr,I_qr,I_ds,I_qs,Tau_t = run(sys_struct, struct)
```

```python
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(8, 5), sharex = True)
axes.plot(T,Tau_e)
fig.savefig('dfim_tau_e.svg', bbox_inches='tight')
```

    <IPython.core.display.Javascript object>

```python
fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(9, 8), sharex = True)
axes[0].plot(T,P_s_1/1e6, label='$\sf p_{s1}$')
axes[0].plot(T,Q_s_1/1e6, label='$\sf q_{s1}$')
#axes[0].plot(T,P_s_2/1e6, label='$\sf p_{s2}$')
#axes[0].plot(T,Q_s_2/1e6, label='$\sf q_{s2}$')

axes[1].plot(T,P_r_1/1e6, label='$\sf p_{r1}$')
axes[1].plot(T,Q_r_1/1e6, label='$\sf q_{r1}$')
#axes[1].plot(T,P_r_2/1e6, label='$\sf p_{r2}$')
#axes[1].plot(T,Q_r_2/1e6, label='$\sf q_{r2}$')

axes[0].legend(loc='best')
axes[1].legend(loc='best')

axes[0].set_ylabel('Stator powers (MVA)')
axes[1].set_ylabel('Rotor powers (MVA)')
axes[1].set_xlabel('Time (s)')
axes[0].set_ylim([-0.1,1.1])
#axes[0].set_xlim([0,3.0])

fig.savefig('dfim_pq_s_pq_r.svg', bbox_inches='tight')

```

    <IPython.core.display.Javascript object>

```python
fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(8, 8), sharex = True)
axes[0].plot(T,Omega_t)

axes[1].plot(T,Tau_t/20.0/1000, label='$\sf \\tau_t$')
axes[1].plot(T,Tau_e/1000, label='$\sf \\tau_e$')
#axes[1].plot(T,P_r_2/1e6, label='$\sf p_{r2}$')
#axes[1].plot(T,Q_r_2/1e6, label='$\sf q_{r2}$')

axes[0].legend(loc='best')
axes[1].legend(loc='best')

axes[0].set_ylabel('Rotor speed')
axes[1].set_ylabel('Torques (kNm)')
axes[1].set_xlabel('Time (s)')
#axes[0].set_ylim([0,2.5])
#axes[0].set_xlim([0,3.0])

fig.savefig('dfim_omega_taus.svg', bbox_inches='tight')

```

    <IPython.core.display.Javascript object>

    /home/jmmauricio/bin/anaconda3/lib/python3.5/site-packages/matplotlib/axes/_axes.py:531: UserWarning: No labelled objects found. Use label='...' kwarg on individual plots.
      warnings.warn("No labelled objects found. "

```python
fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(9, 8), sharex = True)
axes[0].plot(T,V_dr, label='$\sf v_{dr}$')
axes[0].plot(T,V_qr, label='$\sf v_{qr}$')

axes[1].plot(T,P_r_1/1e6, label='$\sf p_{r}$', color = colors[2])
axes[1].plot(T,Q_r_1/1e6, label='$\sf q_{r}$', color = colors[3])
axes[1].plot(T,(P_r_1**2+Q_r_1**2)**0.5/1e6, label='$\sf s_{r}$', color = colors[4])

axes[0].legend(loc='best')
axes[1].legend(loc='best')

axes[0].set_ylabel('Rotor voltages (V)')
axes[1].set_ylabel('Rotor powers (MVA)')
axes[1].set_xlabel('Time (s)')
#axes[0].set_ylim([0,2.5])
#axes[0].set_xlim([0,3.0])

fig.savefig('dfim_rotor_v_powers.svg', bbox_inches='tight')

```

    <IPython.core.display.Javascript object>

```python
fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(9, 8), sharex = True)
axes[0].plot(T,I_dr/1000, label='$\sf i_{dr}$')
axes[0].plot(T,I_qr/1000, label='$\sf i_{qr}$')
axes[1].plot(T,I_ds/1000, label='$\sf i_{ds}$')
axes[1].plot(T,I_qs/1000, label='$\sf i_{qs}$')

axes[0].legend()
axes[1].legend()

axes[0].set_ylabel('Stator currents (kA)')
axes[1].set_ylabel('Rotor currents (kA)')
axes[1].set_xlabel('Time (s)')
#axes[0].set_ylim([0,2.5])
#axes[0].set_xlim([0,3.0])

fig.savefig('dfim_i_s_i_r.svg', bbox_inches='tight')

```

    <IPython.core.display.Javascript object>

```python

```

```python

```

```python

```

```python

```