```python
%matplotlib inline
```

Autograd: Automatic Differentiation
===================================

Central to all neural networks in PyTorch is the ``autograd`` package.
Let’s first briefly visit this, and we will then go to training our
first neural network.

The ``autograd`` package  provides automatic differentiation for all operations
on Tensors. It is a define-by-run framework, which means that your backprop is defined by how your code is run, and that every single iteration can be different.

Let us see this in more simple terms with some examples.

Tensor
--------

``torch.Tensor`` is the central class of the package. If you set its attribute
``.requires_grad`` as ``True``, it starts to track all operations on it. When
you finish your computation you can call ``.backward()`` and have all the
gradients computed automatically. The gradient for this tensor will be
accumulated into ``.grad`` attribute.

To stop a tensor from tracking history, you can call ``.detach()`` to detach
it from the computation history, and to prevent future computation from being
tracked.

To prevent tracking history (and using memory), you can also wrap the code block
in ``with torch.no_grad():``. This can be particularly helpful when evaluating a
model because the model may have trainable parameters with
``requires_grad=True``, but for which we don't need the gradients.

There’s one more class which is very important for autograd
implementation - a ``Function``.

``Tensor`` and ``Function`` are interconnected and build up an **acyclic
graph**, that encodes a complete history of computation. Each tensor has
a ``.grad_fn`` attribute that references a ``Function`` that has created
the ``Tensor`` (except for Tensors created by the user - their
``grad_fn is None``).

If you want to compute the derivatives, you can call ``.backward()`` on
a ``Tensor``. If ``Tensor`` is a scalar (i.e. it holds a one element
data), you don’t need to specify any arguments to ``backward()``,
however if it has more elements, you need to specify a ``gradient``
argument that is a tensor of matching shape.

```python
import torch
torch.__version__
```

    '0.4.1.post2'

Create a tensor and set ``requires_grad=True`` to track computation with it

```python
x = torch.ones(2, 2, requires_grad=True)
print(x)
```

    tensor([[1., 1.],
            [1., 1.]], requires_grad=True)

Do a tensor operation:

```python
y = x + 2
print(y)
```

    tensor([[3., 3.],
            [3., 3.]], grad_fn=<AddBackward>)

``y`` was created as a result of an operation, so it has a ``grad_fn``.

```python
print(y.grad_fn)
```

    <AddBackward object at 0x7f347d2c96d8>

Do more operations on ``y``

```python
z = y * y * 3
out = z.mean()

print(z, out)
```

    tensor([[27., 27.],
            [27., 27.]], grad_fn=<MulBackward>) tensor(27., grad_fn=<MeanBackward1>)

``.requires_grad_( ... )`` changes an existing Tensor's ``requires_grad``
flag in-place. The input flag defaults to ``False`` if not given.

```python
a = torch.randn(2, 2)
print('[1] a=', a)

a = ((a * 3) / (a - 1))
print('[2] a=', a)
print(a.requires_grad)

a.requires_grad_(True)
print(a.requires_grad)

b = (a * a).sum()
print('[3] b=', b)
print(b.grad_fn)
```

    [1] a= tensor([[ 0.3127,  0.6842],
            [-0.0578, -1.8010]])
    [2] a= tensor([[-1.3650, -6.4986],
            [ 0.1641,  1.9289]])
    False
    True
    [3] b= tensor(47.8424, grad_fn=<SumBackward0>)
    <SumBackward0 object at 0x7f347d2614a8>

Gradients
---------
Let's backprop now.
Because ``out`` contains a single scalar, ``out.backward()`` is
equivalent to ``out.backward(torch.tensor(1.))``.

```python
out.backward()
```

Print gradients d(out)/dx

```python
print(x.grad)
```

    tensor([[4.5000, 4.5000],
            [4.5000, 4.5000]])

You should have got a matrix of ``4.5``. Let’s call the ``out``
*Tensor* “$o$”.

We have that $o = \frac{1}{4}\sum_i z_i$,
$z_i = 3(x_i+2)^2$ and $z_i\bigr\rvert_{x_i=1} = 27$.

Therefore,
$\frac{\partial o}{\partial x_i} = \frac{3}{2}(x_i+2)$, hence
$\frac{\partial o}{\partial x_i}\bigr\rvert_{x_i=1} = \frac{9}{2} = 4.5$.

Mathematically, if you have a vector valued function $\vec{y}=f(\vec{x})$,
then the gradient of $\vec{y}$ with respect to $\vec{x}$
is a **Jacobian matrix**:

\begin{align}J=\left(\begin{array}{ccc}
   \frac{\partial y_{1}}{\partial x_{1}} & \cdots & \frac{\partial y_{1}}{\partial x_{n}}\\
   \vdots & \ddots & \vdots\\
   \frac{\partial y_{m}}{\partial x_{1}} & \cdots & \frac{\partial y_{m}}{\partial x_{n}}
   \end{array}\right)\end{align}

Generally speaking, ``torch.autograd`` is an engine for computing
**vector-Jacobian product**.

That is, given any vector
$v=\left(\begin{array}{cccc} v_{1} & v_{2} & \cdots & v_{m}\end{array}\right)^{T}$,
compute the product $v^{T}\cdot J$.
 If $v$ happens to be
the gradient of a scalar function $l=g\left(\vec{y}\right)$,
that is,
$v=\left(\begin{array}{ccc}\frac{\partial l}{\partial y_{1}} & \cdots & \frac{\partial l}{\partial y_{m}}\end{array}\right)^{T}$,
then by the chain rule, the vector-Jacobian product would be the
gradient of $l$ with respect to $\vec{x}$:

\begin{align}J^{T}\cdot v=\left(\begin{array}{ccc}
   \frac{\partial y_{1}}{\partial x_{1}} & \cdots & \frac{\partial y_{m}}{\partial x_{1}}\\
   \vdots & \ddots & \vdots\\
   \frac{\partial y_{1}}{\partial x_{n}} & \cdots & \frac{\partial y_{m}}{\partial x_{n}}
   \end{array}\right)\left(\begin{array}{c}
   \frac{\partial l}{\partial y_{1}}\\
   \vdots\\
   \frac{\partial l}{\partial y_{m}}
   \end{array}\right)=\left(\begin{array}{c}
   \frac{\partial l}{\partial x_{1}}\\
   \vdots\\
   \frac{\partial l}{\partial x_{n}}
   \end{array}\right)\end{align}

(Note that $v^{T}\cdot J$ gives a row vector which can be
treated as a column vector by taking $J^{T}\cdot v$.)

This characteristic of vector-Jacobian product makes it very
convenient to feed external gradients into a model that has
non-scalar output.

Now let's take a look at an example of vector-Jacobian product, where `.norm()` is **L2-norm**:

```python
x = torch.randn(3, requires_grad=True)

y = x * 2
count = 0
while y.data.norm() < 1000:
    count += 1
    y = y * 2

print(count, '=> ', y)
```

    tensor([-856.5516,  499.1843,  456.1978], grad_fn=<MulBackward>)

Now in this case ``y`` is no longer a scalar. ``torch.autograd``
could not compute the full Jacobian directly, but if we just
want the vector-Jacobian product, simply pass the vector to
``backward`` as argument:

```python
v = torch.tensor([0.1, 1.0, 0.0001], dtype=torch.float)
y.backward(v)

print(x.grad)
```

    tensor([ 204.8000, 2048.0000,    0.2048])

You can also stop autograd from tracking history on Tensors
with ``.requires_grad=True`` by wrapping the code block in
``with torch.no_grad():``

```python
print(x.requires_grad)
print((x ** 2).requires_grad)

with torch.no_grad():
	print((x ** 2).requires_grad)
```

    True
    True
    False

**Read Later:**

Documentation of ``autograd`` and ``Function`` is at
https://pytorch.org/docs/autograd


[Sebastian Raschka](http://sebastianraschka.com)

- [Link to the containing GitHub Repository](https://github.com/rasbt/algorithms_in_ipython_notebooks)
- [Link to this IPython Notebook on GitHub](https://github.com/rasbt/algorithms_in_ipython_notebooks/blob/master/ipython_nbs/geometry/eucldist_numpy_vectorization.ipynb)

```
%load_ext watermark
```

```
%watermark -a 'Sebastian Raschka' -v -d
```

    Sebastian Raschka 03/07/2014

    CPython 3.4.1
    IPython 2.0.0

[More information](http://nbviewer.ipython.org/github/rasbt/python_reference/blob/master/ipython_magic/watermark.ipynb) about the `watermark` magic command extension.

# Vectorizing a classic for-loop in NumPy for calculating Euclidean distances

One of the biggest advantages of NumPy, besides its convenient usage, is the speed gain over classic Python loop structures via vectorized arithmetic operations for its `ndarray` objects.

In general, "vectorizing" means that artithmetic operations on elments  in a vector can be done in parallel as an 1-step process.
In theory, if we forget about additional overheads, a vectorized arithmetic addition, e.g.,

$\begin{pmatrix} 1 \\ 2 \\ 3 \\ 4 \end{pmatrix} + 1 = \begin{pmatrix} 1+1 \\ 2+1 \\ 3+1 \\ 4+1 \end{pmatrix} = \begin{pmatrix} 2 \\ 3 \\ 4 \\ 5\end{pmatrix}$

could be 4 times faster than adding the constant 1 to every number in sequential order.

In NumPy, some operations that are using the architecture's "BLAS" (Basic Linear Algebra Subroutines) can also take advantage of CPUs with multiple cores and run multiple processes in parallel, for example, the matrix dot product (`numpy.np(A,B)` or `A.dot(B)`).
Other array operations.

Other arithmetic operations in NumPy, e.g, arithmetic addition, don't make use of BLAS in the current implementation of NumPy. However, those are still overcoming Python's GIL (Global Interpreter Lock) to allow multi-threading (not to be confused with multi-processing), which can still result in a significant speed boost.
(Source: http://wiki.scipy.org/ParallelProgramming)

### Euclidean Distance

For the following benchmarks, we will use a simple Euclidean distance calculation with the following equation

\begin{equation} d = \sqrt{(X_1 - Y_1)^2 + (X_2 - Y_2)^2 + (X_3 - Y_3)^2 + ... (X_d - Y_d)}\end{equation}

But before we skip to the actual benchmark, let us visualize it for a set of 2 3D coordinates:

```
%matplotlib inline
```

```
import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d import proj3d

coords1 = [1, 2, 3]
coords2 = [4, 5, 6]

fig = plt.figure(figsize=(7,7))
ax = fig.add_subplot(111, projection='3d')

ax.scatter((coords1[0], coords2[0]),
        (coords1[1], coords2[1]),
        (coords1[2], coords2[2]),
         color="k", s=150)

ax.plot((coords1[0], coords2[0]),
        (coords1[1], coords2[1]),
        (coords1[2], coords2[2]),
         color="r")

ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')

ax.text(x=2.5, y=3.5, z=4.0, s='d = 5.19')

plt.title('Euclidean distance between 2 3D-coordinates')

plt.show()
```

### Python and NumPy Implementations

```
# Small sample data of 3D coordinates

coords1 = [1, 2, 3]
coords2 = [4, 5, 6]
np_c1 = np.array(coords1)
np_c2 = np.array(coords2)
```

```
# Classic For-loop

def eucldist_forloop(coords1, coords2):
    """ Calculates the euclidean distance between 2 lists of coordinates. """
    dist = 0
    for (x, y) in zip(coords1, coords2):
        dist += (x - y)**2
    return dist**0.5
```

```
# Generator expression

def eucldist_generator(coords1, coords2):
    """ Calculates the euclidean distance between 2 lists of coordinates. """
    return sum((x - y)**2 for x, y in zip(coords1, coords2))**0.5
```

```
# Vectorized version using NumPy

def eucldist_vectorized(coords1, coords2):
    """ Calculates the euclidean distance between 2 lists of coordinates. """
    return np.sqrt(np.sum((coords1 - coords2)**2))
```

```
# Using an in-built NumPy function

np.linalg.norm(np_c1 - np_c2)
```

    5.196152422706632

```
print(eucldist_forloop(coords1, coords2))
print(eucldist_generator(coords1, coords2))
print(eucldist_vectorized(np_c1, np_c2))
print(np.linalg.norm(np_c1 - np_c2))
```

    5.196152422706632
    5.196152422706632
    5.19615242271
    5.19615242271

# `timeit` benchmarks

```
import timeit
import random
random.seed(123)

from numpy.linalg import norm as np_linalg_norm

funcs = ('eucldist_forloop', 'eucldist_generator', 'eucldist_vectorized', 'np_linalg_norm')
times = {f:[] for f in funcs}
orders_n = [10**i for i in range(1, 8)]
for n in orders_n:

    c1 = [random.randint(0,100) for _ in range(n)]
    c2 = [random.randint(0,100) for _ in range(n)]
    np_c1 = np.array(c1)
    np_c2 = np.array(c2)

    assert(eucldist_forloop(c1, c2)
           == eucldist_generator(c1, c2)
           == eucldist_vectorized(np_c1, np_c2)
           == np_linalg_norm(np_c1 - np_c2)
           )

    times['eucldist_forloop'].append(min(timeit.Timer('eucldist_forloop(c1, c2)',
            'from __main__ import c1, c2, eucldist_forloop').repeat(repeat=50, number=1)))
    times['eucldist_generator'].append(min(timeit.Timer('eucldist_generator(c1, c2)',
            'from __main__ import c1, c2, eucldist_generator').repeat(repeat=50, number=1)))
    times['eucldist_vectorized'].append(min(timeit.Timer('eucldist_vectorized(np_c1, np_c2)',
            'from __main__ import np_c1, np_c2, eucldist_vectorized').repeat(repeat=50, number=1)))
    times['np_linalg_norm'].append(min(timeit.Timer('np_linalg_norm(np_c1 - np_c2)',
            'from __main__ import np_c1, np_c2, np_linalg_norm').repeat(repeat=50, number=1)))
```

```
labels = {'eucldist_forloop': 'for-loop',
          'eucldist_generator': 'generator expression (comprehension equiv.)',
          'eucldist_vectorized': 'NumPy vectorization',
          'np_linalg_norm': 'numpy.linalg.norm'
          }

def plot(times, orders_n, labels):

    colors = ('cyan', '#7DE786', 'black', 'blue')
    linestyles = ('-', '-', '--', '--')
    fig = plt.figure(figsize=(11,10))
    for lb,c,l in zip(labels.keys(), colors, linestyles):
        plt.plot(orders_n, times[lb], alpha=1, label=labels[lb],
                 lw=3, color=c, linestyle=l)
    plt.xlabel('sample size n (items in the list)', fontsize=14)
    plt.ylabel('time per computation in seconds', fontsize=14)
    plt.xlim([min(orders_n) / 10, max(orders_n)* 10])
    plt.legend(loc=2, fontsize=14)
    plt.grid()
    plt.xticks(fontsize=16)
    plt.yticks(fontsize=16)
    plt.xscale('log')
    plt.yscale('log')
    plt.title('Python for-loop/generator expr. vs. NumPy vectorized code', fontsize=18)
    plt.show()
```

```
plot(times, orders_n, labels)
```

```
%watermark -m -v -d -t -z -p numpy
```

    03/07/2014 14:00:21 EDT

    CPython 3.4.1
    IPython 2.0.0

    numpy 1.8.1

    compiler   : GCC 4.2.1 (Apple Inc. build 5577)
    system     : Darwin
    release    : 13.2.0
    machine    : x86_64
    processor  : i386
    CPU cores  : 4
    interpreter: 64bit

As we can see in the plot, the vectorized code runs a lot faster than our classic Python `for`-loop (note that this graph is plotted on logarithmic scale).


```python
%matplotlib inline
```

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

import solowpy
```

# 4. Solving the model

## 4.1 Solow model as an initial value problem

The Solow model with can be formulated as an initial value problem (IVP) as follows.

$$ \dot{k}(t) = sf(k(t)) - (g + n + \delta)k(t),\ t\ge t_0,\ k(t_0) = k_0 \tag{4.1.0} $$

The solution to this IVP is a function $k(t)$ describing the time-path of capital stock (per unit effective labor). Our objective in this section will be to explore methods for approximating the function $k(t)$. The methods we will learn are completely general and can be used to solve any IVP. Those interested in learning more about these methods should start by reading Chapter 10 of [*Numerical Methods for Economists*](http://mitpress.mit.edu/books/numerical-methods-economics) by Ken Judd before proceeding to John Butcher's excellent book entitled [*Numerical Methods for solving Ordinary Differential Equations*](http://www.shivanian.com/teaching%20avtivities/Butcher.pdf).

Before discussing numerical methods we should stop and consider whether or not there are any special cases (i.e., combintions of parameters) for which the the initial value problem defined in 4.2.1 might have an analytic solution. Analytic results can be very useful in building intuition about the economic mechanisms at play in a model and are invaluable for debugging code.

## 4.2 Analytic methods

### 4.2.1 Analytic solution for a model with Cobb-Douglas production

The Solow model with Cobb-Douglas production happens to have a completely [general analytic solution](https://github.com/davidrpugh/numerical-methods/raw/master/labs/lab-1/solow-analytic-solution.pdf):

$$ k(t) = \left[\left(\frac{s}{n+g+\delta}\right)\left(1 -  e^{-(n + g + \delta) (1-\alpha) t}\right) + k_0^{1-\alpha}e^{-(n + g + \delta) (1-\alpha) t}\right]^{\frac{1}{1-\alpha}} \tag{4.2.0}$$

This analytic result is available via the `analytic_solution` method of the `solow.CobbDouglasModel` class.

```python
solowpy.CobbDouglasModel.analytic_solution?
```

### Example: Computing the analytic trajectory
We can compute an analytic solution for our Solow model like so...

```python
# define model parameters
cobb_douglas_params = {'A0': 1.0, 'L0': 1.0, 'g': 0.02, 'n': 0.03, 's': 0.15,
                      'delta': 0.05, 'alpha': 0.33}

# create an instance of the solow.Model class
cobb_douglas_model = solowpy.CobbDouglasModel(params=cobb_douglas_params)
```

```python
# specify some initial condition
k0 = 0.5 * cobb_douglas_model.steady_state

# grid of t values for which we want the value of k(t)
ti = np.linspace(0, 100, 10)

# generate a trajectory!
cobb_douglas_model.analytic_solution(ti, k0)
```

    array([[   0.        ,    0.91578343],
           [  11.11111111,    1.37081903],
           [  22.22222222,    1.60723533],
           [  33.33333333,    1.72380474],
           [  44.44444444,    1.78011125],
           [  55.55555556,    1.80706521],
           [  66.66666667,    1.81991504],
           [  77.77777778,    1.8260292 ],
           [  88.88888889,    1.82893579],
           [ 100.        ,    1.83031695]])

...and we can make a plot of this solution like so...

```python
fig, ax = plt.subplots(1, 1, figsize=(8,6))

# compute the solution
ti = np.linspace(0, 100, 1000)
analytic_traj = cobb_douglas_model.analytic_solution(ti, k0)

# plot this trajectory
ax.plot(ti, analytic_traj[:,1], 'r-')

# equilibrium value of capital stock (per unit effective labor)
ax.axhline(cobb_douglas_model.steady_state, linestyle='dashed',
           color='k', label='$k^*$')

# axes, labels, title, etc
ax.set_xlabel('Time, $t$', fontsize=20, family='serif')
ax.set_ylabel('$k(t)$', rotation='horizontal', fontsize=20, family='serif')
ax.set_title('Analytic solution to a Solow model\nwith Cobb-Douglas production',
             fontsize=25, family='serif')
ax.legend(loc=0, frameon=False, bbox_to_anchor=(1.0, 1.0))
ax.grid('on')

plt.show()
```

### 4.2.2 Linearized solution to general model

In general there will not be closed-form solutions for the Solow model. The standard approach to obtaining general analytical results for the Solow model is to linearize the equation of motion for capital stock (per unit effective labor). Linearizing the equation of motion of capital (per unit effective labor) amounts to taking a first-order [Taylor approximation](http://en.wikipedia.org/wiki/Taylor_series) of equation 4.1.0 around its long-run equilibrium $k^*$:

$$ \dot{k}(t) \approx -\lambda (k(t) - k^*),\ t \ge t_0,\ k(t_0)=k_0 \tag{4.2.1}$$

where the *speed of convergence*, $\lambda$, is defined as

$$ \lambda = -\frac{\partial \dot{k}(k(t))}{\partial k(t)}\bigg|_{k(t)=k^*} \tag{4.2.2} $$

The solution the the linear differential equation 4.2.1 is

$$ k(t) = k^* + e^{-\lambda t}(k_0 - k^*). \tag{4.2.3} $$

To complete the solution it remains to find an expression for the speed of convergence $\lambda$:

\begin{align}
    \lambda \equiv -\frac{\partial \dot{k}(k(t))}{\partial k(t)}\bigg|_{k(t)=k^*} =& -[sf'(k^*) - (g + n+ \delta)] \\
    =& (g + n+ \delta) - sf'(k^*) \\
    =& (g + n + \delta) - (g + n + \delta)\frac{k^*f'(k^*)}{f(k^*)} \\
    =& (1 - \alpha_K(k^*))(g + n + \delta) \tag{4.2.4}
\end{align}
where the elasticity of output with respect to capital, $\alpha_K(k)$, is

$$\alpha_K(k) = \frac{k^*f'(k^*)}{f(k^*)}. \tag{4.2.5}$$

### Example: Computing the linearized trajectory
One can compute a linear approximation of the model solution using the `linearized_solution` method of the `solow.Model` class as follows.

```python
# specify some initial condition
k0 = 0.5 * cobb_douglas_model.steady_state

# grid of t values for which we want the value of k(t)
ti = np.linspace(0, 100, 10)

# generate a trajectory!
cobb_douglas_model.linearized_solution(ti, k0)
```

    array([[   0.        ,    0.91578343],
           [  11.11111111,    1.39657145],
           [  22.22222222,    1.62494486],
           [  33.33333333,    1.73342179],
           [  44.44444444,    1.78494813],
           [  55.55555556,    1.80942305],
           [  66.66666667,    1.82104859],
           [  77.77777778,    1.82657069],
           [  88.88888889,    1.82919369],
           [ 100.        ,    1.8304396 ]])

### 4.2.3 Accuracy of the linear approximation

```python
# initial condition
t0, k0 = 0.0, 0.5 * cobb_douglas_model.steady_state

# grid of t values for which we want the value of k(t)
ti = np.linspace(t0, 100, 1000)

# generate the trajectories
analytic = cobb_douglas_model.analytic_solution(ti, k0)
linearized = cobb_douglas_model.linearized_solution(ti, k0)
```

```python
fig, ax = plt.subplots(1, 1, figsize=(8,6))

ax.plot(ti, analytic[:,1], 'r-', label='Analytic')
ax.plot(ti, linearized[:,1], 'b-', label='Linearized')

# demarcate k*
ax.axhline(cobb_douglas_model.steady_state, linestyle='dashed',
            color='k', label='$k^*$')

# axes, labels, title, etc
ax.set_xlabel('Time, $t$', fontsize=20, family='serif')
ax.set_ylabel('$k(t)$', rotation='horizontal', fontsize=20, family='serif')
ax.set_title('Analytic vs. linearized solutions', fontsize=25, family='serif')
ax.legend(loc='best', frameon=False, prop={'family':'serif'},
          bbox_to_anchor=(1.0, 1.0))
ax.grid('on')

fig.show()
```

## 4.3 Finite-difference methods

Four of the best, most widely used ODE integrators have been implemented in the `scipy.integrate` module (they are called `dopri5`, `dop85`, `lsoda`, and `vode`). Each of these integrators uses some type of adaptive step-size control: the integrator adaptively adjusts the step size $h$ in order to keep the approximation error below some user-specified threshold). The cells below contain code which compares the approximation error of the forward Euler with those of [`lsoda`](http://computation.llnl.gov/casc/nsde/pubs/u113855.pdf) and [`vode`](http://jeffreydk.site50.net/papers/BDFmethodpaper.pdf). Instead of simple linear interpolation (i.e., `k=`1), I set `k=5` for 5th order [B-spline](http://en.wikipedia.org/wiki/B-spline) interpolation.

...finally, we can plot trajectories for different initial conditions. Note that the analytic solutions converge to the long-run equilibrium no matter the initial condition of capital stock (per unit of effective labor) providing a nice graphical demonstration that the Solow model is globally stable.

```python
fig, ax = plt.subplots(1, 1, figsize=(8,6))

# lower and upper bounds for initial conditions
k_star = solow.cobb_douglas.analytic_steady_state(cobb_douglas_model)
k_l = 0.5 * k_star
k_u = 2.0 * k_star

for k0 in np.linspace(k_l, k_u, 5):

    # compute the solution
    ti = np.linspace(0, 100, 1000)
    analytic_traj = solow.cobb_douglas.analytic_solution(cobb_douglas_model, ti, k0)

    # plot this trajectory
    ax.plot(ti, analytic_traj[:,1])

# equilibrium value of capital stock (per unit effective labor)
ax.axhline(k_star, linestyle='dashed', color='k', label='$k^*$')

# axes, labels, title, etc
ax.set_xlabel('Time, $t$', fontsize=15, family='serif')
ax.set_ylabel('$k(t)$', rotation='horizontal', fontsize=20, family='serif')
ax.set_title('Analytic solution to a Solow model\nwith Cobb-Douglas production',
             fontsize=20, family='serif')
ax.legend(loc=0, frameon=False, bbox_to_anchor=(1.0, 1.0))
ax.grid('on')

plt.show()
```

```python
k0 = 0.5 * ces_model.steady_state
numeric_trajectory = ces_model.ivp.solve(t0=0, y0=k0, h=0.5, T=100, integrator='dopri5')

ti = numeric_trajectory[:,0]
linearized_trajectory = ces_model.linearized_solution(ti, k0)

```

### 4.3.2 Accuracy of finite-difference methods

```python
t0, k0 = 0.0, 0.5
numeric_soln = cobb_douglas_model.ivp.solve(t0, k0, T=100, integrator='lsoda')
```

```python
fig, ax = plt.subplots(1, 1, figsize=(8,6))

# compute and plot the numeric approximation
t0, k0 = 0.0, 0.5
numeric_soln = cobb_douglas_model.ivp.solve(t0, k0, T=100, integrator='lsoda')
ax.plot(numeric_soln[:,0], numeric_soln[:,1], 'bo', markersize=3.0)

# compute and plot the analytic solution
ti = np.linspace(0, 100, 1000)
analytic_soln = solow.cobb_douglas.analytic_solution(cobb_douglas_model, ti, k0)
ax.plot(ti, analytic_soln[:,1], 'r-')

# equilibrium value of capital stock (per unit effective labor)
k_star = solow.cobb_douglas.analytic_steady_state(cobb_douglas_model)
ax.axhline(k_star, linestyle='dashed', color='k', label='$k^*$')

# axes, labels, title, etc
ax.set_xlabel('Time, $t$', fontsize=15, family='serif')
ax.set_ylabel('$k(t)$', rotation='horizontal', fontsize=20, family='serif')
ax.set_title('Numerical approximation of the solution',
             fontsize=20, family='serif')
ax.legend(loc=0, frameon=False, bbox_to_anchor=(1.0, 1.0))
ax.grid('on')

plt.show()
```

```python
ti = np.linspace(0, 100, 1000)
interpolated_soln = cobb_douglas_model.ivp.interpolate(numeric_soln, ti, k=3)
```

```python
fig, ax = plt.subplots(1, 1, figsize=(8,6))

# compute and plot the numeric approximation
ti = np.linspace(0, 100, 1000)
interpolated_soln = cobb_douglas_model.ivp.interpolate(numeric_soln, ti, k=3)
ax.plot(ti, interpolated_soln[:,1], 'b-')

# compute and plot the analytic solution
analytic_soln = solow.cobb_douglas.analytic_solution(cobb_douglas_model, ti, k0)
ax.plot(ti, analytic_soln[:,1], 'r-')

# equilibrium value of capital stock (per unit effective labor)
k_star = solow.cobb_douglas.analytic_steady_state(cobb_douglas_model)
ax.axhline(k_star, linestyle='dashed', color='k', label='$k^*$')

# axes, labels, title, etc
ax.set_xlabel('Time, $t$', fontsize=15, family='serif')
ax.set_ylabel('$k(t)$', rotation='horizontal', fontsize=20, family='serif')
ax.set_title('Numerical approximation of the solution',
             fontsize=20, family='serif')
ax.legend(loc=0, frameon=False, bbox_to_anchor=(1.0, 1.0))
ax.grid('on')

plt.show()
```

```python

```

```python
ti = np.linspace(0, 100, 1000)
residual = cobb_douglas_model.ivp.compute_residual(numeric_soln, ti, k=3)
```

```python
# extract the raw residuals
capital_residual = residual[:, 1]

# typically, normalize residual by the level of the variable
norm_capital_residual = np.abs(capital_residual) / interpolated_soln[:,1]

# create the plot
fig = plt.figure(figsize=(8, 6))
plt.plot(ti, norm_capital_residual, 'b-', label='$k(t)$')
plt.axhline(np.finfo('float').eps, linestyle='dashed', color='k', label='Machine eps')
plt.xlabel('Time', fontsize=15)
plt.ylim(1e-16, 1)
plt.ylabel('Residuals (normalized)', fontsize=15, family='serif')
plt.yscale('log')
plt.title('Residual', fontsize=20, family='serif')
plt.grid()
plt.legend(loc=0, frameon=False, bbox_to_anchor=(1.0,1.0))
plt.show()
```

```python

```


# Job Shop Scheduling

## Imports

The following cell specifies the solver to used in the subsequent calculations. Some of these problems can become quite larger, and therefore the `gurobi` solver has been set as a default. If you don't have the `gurobi` solver then adjust the code to use the `glpk` solver, but know the calculations may take longer (and the benchmark problem will not solve at all). If you do have the `gurobi` solver, edit the location of the executable to match the location on your computer.

```python
# Import Pyomo and solvers for Google Colab
import sys
if "google.colab" in sys.modules:
    !wget -N -q https://raw.githubusercontent.com/jckantor/MO-book/main/tools/install_on_colab.py
    %run install_on_colab.py
```

## Job Shop Scheduling

Job shop scheduling is one of the classic problems in Operations Research.

A *job shop* refers to a collection of *machines* that process *jobs*. The jobs may be simple or complex, such as the production of a custom part, a print job at printing house, a patient treated in a hospital, or a meal produced in a fast food restaurant. The term "machines" can refer any device, person, or resource involved in the processing the jobs. In what follows in this notebook, a job comprises a series of tasks that requiring use of particular machines for known durations, and which must be completed in specified order.

The *job shop scheduling problem* is to schedule a set of jobs on the available machines to optimize a metric of productivity. A typical metric is the *makespan* which refers to the time needed to process all jobs.

At a basic level, the data required to specify a job shop schedulig problem consists of two tables. The first table is a decomposition of the jobs into a series of tasks, each task listing the job name, name of the required machine, and task duration. A second table list task pairs where the first task must be completed before the second task can be started.

Real applications of job shop sceduling oten include additional considerations such as time that may be required to set up machines prior to the start of the next job, or limits on the time a job can wait between processing steps.

## Job shop example

The following example of a job shop is from from Christelle Gueret, Christian Prins, Marc Sevaux, "Applications of Optimization with Xpress-MP," Dash Optimization, 2000.

In this example, there are three printed paper products that must pass through color printing presses in a particular order. The given data consists of a flowsheet showing the order in which each job passes through the color presses

and a table of data showing, in minutes, the amount of time each job requires on each machine.

| Machine | Color | Paper 1 | Paper 2 | Paper 3 |
| :-----: | :---: | :-----: | :-----: | :-----: |
| 1       | Blue  | 45      | 20      | 12      |
| 2       | Green | -       | 10      | 17      |
| 3       | Yellow| 10      | 34      | 28      |

What is the minimum amount of time (i.e, what is the makespan) for this set of jobs?

## Task decomposition

The first step in the analysis is to decompose the process into a series of tasks. Each task c(job, machine) pair. Some tasks cannot start until a prerequisite task is completed.

| Task (Job,Machine) | Duration | Prerequisite Task |
| :----------------: | :------: | :---------------: |
| (Paper 1, Blue)    | 45       | -                 |
| (Paper 1, Yellow)  | 10       | (Paper 1,Blue)    |
| (Paper 2, Blue)    | 20       | (Paper 2, Green)  |
| (Paper 2, Green)   | 10       | -                 |
| (Paper 2, Yellow)  | 34       | (Paper 2, Blue)   |
| (Paper 3, Blue)    | 12       | (Paper 3, Yellow) |
| (Paper 3, Green)   | 17       | (Paper 3, Blue)   |
| (Paper 3, Yellow)  | 28       | -                 |

We convert this to a JSON style representation where tasks are denoted by (Job,Machine) tuples in Python. The task data is stored in a Python dictionary indexed by (Job,Machine) tuples. The task data conists of a dictionary with duration ('dur') and (Job,Machine) pair for any prerequisite task.

```python
TASKS = {
    ('Paper_1','Blue')   : {'dur': 45, 'prec': None},
    ('Paper_1','Yellow') : {'dur': 10, 'prec': ('Paper_1','Blue')},
    ('Paper_2','Blue')   : {'dur': 20, 'prec': ('Paper_2','Green')},
    ('Paper_2','Green')  : {'dur': 10, 'prec': None},
    ('Paper_2','Yellow') : {'dur': 34, 'prec': ('Paper_2','Blue')},
    ('Paper_3','Blue')   : {'dur': 12, 'prec': ('Paper_3','Yellow')},
    ('Paper_3','Green')  : {'dur': 17, 'prec': ('Paper_3','Blue')},
    ('Paper_3','Yellow') : {'dur': 28, 'prec': None},
}
```

## Model formulation

Each task is indexed by an ordered pair $(j, m)$ where $j$ refers to a job, and $m$ refers to a machine. Associated with each task is data describing the time needed to perform the task, and a preceeding task that must be completed before the index task can start.

| Parameter | Description |
| :-------- | :-----------|
| $\text{duration}_{j,m}$ | Duration of task $(j,m)$ |
| $\text{prec}_{j,m}$ | A task $(k,n) = \text{prec}_{j,m}$ that must be completed before task $(j,m)$|

The choice of decision variables for this problem are key to modeling. We introduce $makespan$ as the time needed to complete all tasks. $makespan$ is a candidate objective function. Variable $start_{j,m}$ denotes the time when task $(j,m)$ begins.

| Decision Variables | Description |
| :-------- | :-----------|
| $\text{makespan}$  | Completion of all jobs |
| $\text{start}_{j,m}$ | Start time for task $(j,m)$ |

The constraints include lower bounda on the start and an upper bound on the completion of each task $(j,m)$

44
\begin{align}
\text{start}_{j,m} & \geq 0\\
\text{start}_{j,m}+\text{dur}_{j,m} & \leq \text{makespan}
\end{align}
$$

Any preceeding tasks must be completed before task $(j,m)$ can start.

$$
\begin{align}
\text{start}_{k,n}+\text{dur}_{k,n}\leq\text{start}_{j,m}\ \ \ \ \text{for } (k,n) =\text{prec}_{j,m}
\end{align}
$$

Finally, for every task performed on machine $m$, there can be no overlap among those tasks. This leads to a set of pair-wise disjunctive constraints for each machine.

$$
\begin{align}
\left[\text{start}_{j,m}+\text{dur}_{j,m} \leq \text{start}_{k,m}\right] \vee \left[\text{start}_{k,m}+\text{dur}_{k,m} \leq \text{start}_{j,m}\right]
\end{align}
$$

avoids conflicts for use of the same machine.

## Pyomo implementation

The job shop scheduling problem is implemented below in Pyomo. The implementation consists of of a function JobShopModel(TASKS) that accepts a dictionary of tasks and returns a Pyomo model.

```python
import pyomo.environ as pyo

def jobshop_model(TASKS):

    model = pyo.ConcreteModel()

    # tasks is a two dimensional set of (j,m) constructed from the dictionary keys
    model.TASKS = pyo.Set(initialize = TASKS.keys(), dimen=2)

    # the set of jobs is constructed from a python set
    @model.Set()
    def JOBS(model):
        return list(set([j for (j, m) in model.TASKS]))
    #model.JOBS = pyo.Set(initialize = )))

    # set of machines is constructed from a python set
    model.MACHINES = pyo.Set(initialize = list(set([m for (j, m) in model.TASKS])))

    # the order of tasks is constructed as a cross-product of tasks and filtering
    model.TASKORDER = pyo.Set(initialize = model.TASKS * model.TASKS, dimen=4,
        filter = lambda model, j, m, k, n: (k, n) == TASKS[(j, m)]['prec'])

    # the set of disjunctions is cross-product of jobs, jobs, and machines
    model.DISJUNCTIONS = pyo.Set(initialize = model.JOBS * model.JOBS * model.MACHINES, dimen=3,
        filter = lambda model, j, k, m: j < k and (j, m) in model.TASKS and (k, m) in model.TASKS)

    # load duration data into a model parameter for later access
    @model.Param(model.TASKS)
    def dur(model, j, m):
        return TASKS[(j, m)]['dur']

    # establish an upper bound on makespan
    ub = sum([model.dur[j, m] for (j,m) in model.TASKS])

    # create decision variables
    model.makespan = pyo.Var(bounds=(0, ub))
    model.start = pyo.Var(model.TASKS, bounds=(0, ub))

    @model.Objective(sense=pyo.minimize)
    def minimize_makespan(model):
        return model.makespan

    @model.Constraint(model.TASKS)
    def finish_tasks(model, j, m):
        return model.start[j, m] + model.dur[j, m] <= model.makespan

    @model.Constraint(model.TASKORDER)
    def preceding(model, j, m, k, n):
        return model.start[k, n] + model.dur[k, n] <= model.start[j, m]

    @model.Disjunction(model.DISJUNCTIONS)
    def no_overlap(model, j, k, m):
        return [model.start[j, m] + model.dur[j, m] <= model.start[k, m],
                model.start[k, m] + model.dur[k, m] <= model.start[j, m]]

    pyo.TransformationFactory('gdp.bigm').apply_to(model)
    return model

jobshop_model(TASKS)
```

```python
def jobshop_solve(model):
    pyo.SolverFactory('cbc').solve(model)
    results = [{'Job': j,
                'Machine': m,
                'Start': model.start[j, m](),
                'Duration': model.dur[j,m],
                'Finish': model.start[(j, m)]() + model.dur[j,m]}
               for j,m in model.TASKS]
    return results

def jobshop(TASKS):
    return jobshop_solve(jobshop_model(TASKS))

results = jobshop(TASKS)
results
```

    [{'Job': 'Paper_1',
      'Machine': 'Blue',
      'Start': 42.0,
      'Duration': 45,
      'Finish': 87.0},
     {'Job': 'Paper_1',
      'Machine': 'Yellow',
      'Start': 87.0,
      'Duration': 10,
      'Finish': 97.0},
     {'Job': 'Paper_2',
      'Machine': 'Blue',
      'Start': 10.0,
      'Duration': 20,
      'Finish': 30.0},
     {'Job': 'Paper_2',
      'Machine': 'Green',
      'Start': 0.0,
      'Duration': 10,
      'Finish': 10.0},
     {'Job': 'Paper_2',
      'Machine': 'Yellow',
      'Start': 30.0,
      'Duration': 34,
      'Finish': 64.0},
     {'Job': 'Paper_3',
      'Machine': 'Blue',
      'Start': 30.0,
      'Duration': 12,
      'Finish': 42.0},
     {'Job': 'Paper_3',
      'Machine': 'Green',
      'Start': 42.0,
      'Duration': 17,
      'Finish': 59.0},
     {'Job': 'Paper_3',
      'Machine': 'Yellow',
      'Start': 0.0,
      'Duration': 28,
      'Finish': 28.0}]

## Printing schedules

```python
import pandas as pd

schedule = pd.DataFrame(results)

print('\nSchedule by Job')
print(schedule.sort_values(by=['Job','Start']).set_index(['Job', 'Machine']))

print('\nSchedule by Machine')
print(schedule.sort_values(by=['Machine','Start']).set_index(['Machine', 'Job']))
```

    Schedule by Job
                     Start  Duration  Finish
    Job     Machine
    Paper_1 Blue      42.0        45    87.0
            Yellow    87.0        10    97.0
    Paper_2 Green      0.0        10    10.0
            Blue      10.0        20    30.0
            Yellow    30.0        34    64.0
    Paper_3 Yellow     0.0        28    28.0
            Blue      30.0        12    42.0
            Green     42.0        17    59.0

    Schedule by Machine
                     Start  Duration  Finish
    Machine Job
    Blue    Paper_2   10.0        20    30.0
            Paper_3   30.0        12    42.0
            Paper_1   42.0        45    87.0
    Green   Paper_2    0.0        10    10.0
            Paper_3   42.0        17    59.0
    Yellow  Paper_3    0.0        28    28.0
            Paper_2   30.0        34    64.0
            Paper_1   87.0        10    97.0

## Visualizing Results with Gantt Charts

```python
import matplotlib.pyplot as plt
import matplotlib as mpl

def visualize(results):

    schedule = pd.DataFrame(results)
    JOBS = sorted(list(schedule['Job'].unique()))
    MACHINES = sorted(list(schedule['Machine'].unique()))
    makespan = schedule['Finish'].max()

    bar_style = {'alpha':1.0, 'lw':25, 'solid_capstyle':'butt'}
    text_style = {'color':'white', 'weight':'bold', 'ha':'center', 'va':'center'}
    colors = mpl.cm.Dark2.colors

    schedule.sort_values(by=['Job', 'Start'])
    schedule.set_index(['Job', 'Machine'], inplace=True)

    fig, ax = plt.subplots(2,1, figsize=(12, 5+(len(JOBS)+len(MACHINES))/4))

    for jdx, j in enumerate(JOBS, 1):
        for mdx, m in enumerate(MACHINES, 1):
            if (j,m) in schedule.index:
                xs = schedule.loc[(j,m), 'Start']
                xf = schedule.loc[(j,m), 'Finish']
                ax[0].plot([xs, xf], [jdx]*2, c=colors[mdx%7], **bar_style)
                ax[0].text((xs + xf)/2, jdx, m, **text_style)
                ax[1].plot([xs, xf], [mdx]*2, c=colors[jdx%7], **bar_style)
                ax[1].text((xs + xf)/2, mdx, j, **text_style)

    ax[0].set_title('Job Schedule')
    ax[0].set_ylabel('Job')
    ax[1].set_title('Machine Schedule')
    ax[1].set_ylabel('Machine')

    for idx, s in enumerate([JOBS, MACHINES]):
        ax[idx].set_ylim(0.5, len(s) + 0.5)
        ax[idx].set_yticks(range(1, 1 + len(s)))
        ax[idx].set_yticklabels(s)
        ax[idx].text(makespan, ax[idx].get_ylim()[0]-0.2, "{0:0.1f}".format(makespan), ha='center', va='top')
        ax[idx].plot([makespan]*2, ax[idx].get_ylim(), 'r--')
        ax[idx].set_xlabel('Time')
        ax[idx].grid(True)

    fig.tight_layout()

visualize(results)
```

## Application to the scheduling of batch processes

We will now turn our attention to the application of the job shop scheduling problem to the short term scheduling of batch processes. We illustrate these techniques using Example II from Dunn (2013).

| Recipe  | Mixer | Reactor | Separator | Packaging |
| :-----: | :---: | :-----: | :-------: | :-------: |
| A       | 1.0   | 5.0     | 4.0       | 1.5       |
| B       | -     | -       | 4.5       | 1.0       |
| C       | -     | 3.0     | 5.0       | 1.5       |

### Single product strategies

Before going further, we create a function to streamline the generation of the TASKS dictionary.

```python
def recipe_to_tasks(jobs, machines, durations):
    TASKS = {}
    for j in jobs:
        prec = (None,None)
        for m,d in zip(machines,durations):
            task = (j,m)
            if prec == (None,None):
                TASKS.update({(j,m): {'dur': d, 'prec': None}})
            else:
                TASKS.update({(j,m): {'dur': d, 'prec': prec}})
            prec = task
    return TASKS

recipeA = recipe_to_tasks('A', ['Mixer', 'Reactor', 'Separator', 'Packaging'], [1, 5, 4, 1.5])
visualize(jobshop(recipeA))
```

```python
recipeB = recipe_to_tasks('B', ['Separator', 'Packaging'], [4.5, 1])
visualize(jobshop(recipeB))
```

```python
recipeC = recipe_to_tasks('C', ['Separator', 'Reactor', 'Packaging'], [5, 3, 1.5])
visualize(jobshop(recipeC))
```

### Multiple Overlapping tasks

Let's now consider an optimal scheduling problem where we are wish to make two batches of Product A.

```python
TASKS = recipe_to_tasks(['A1','A2','A3', 'A4'],['Mixer','Reactor','Separator','Packaging'],[1,5,4,1.5])
results = jobshop(TASKS)
visualize(results)
print("Makespan =", max([task['Finish'] for task in results]))
```

Earlier we found it tood 11.5 hours to produce one batch of product A.  As we see here, we can produce a second batch with only 5.0 additional hours because some of the tasks overlap.  The overlapping of tasks is the key to gaining efficiency in batch processing facilities.

Let's next consider production of a single batch each of products A, B, and C.

```python
# update is used to append dictionaries
TASKS = recipeA
TASKS.update(recipeB)
TASKS.update(recipeC)

for k, v in TASKS.items():
    print(k, v)

results = jobshop(TASKS)
visualize(results)
print("Makespan =", max([task['Finish'] for task in results]))
```

The individual production of A, B, and C required 11.5, 5.5, and 9.5 hours, respectively, for a total of 25.5 hours. As we see here, by scheduling the production simultaneously, we can get all three batches done in just 15 hours.

As we see below, each additional set of three products takes an additionl 13 hours.  So there is considerable efficiency gained by scheduling over longer intervals whenever possible.

```python
TASKS = recipe_to_tasks(['A1','A2'],['Mixer','Reactor','Separator','Packaging'],[1,5,4,1.5])
TASKS.update(recipe_to_tasks(['B1','B2'],['Separator','Packaging'],[4.5,1]))
TASKS.update(recipe_to_tasks(['C1','C2'],['Separator','Reactor','Packaging'],[5,3,1.5]))

results = jobshop(TASKS)
visualize(results)
print("Makespan =", max([task['Finish'] for task in results]))
```

### Adding time for unit clean out

A common feature of batch unit operations is a requirement that equipment be cleaned prior to reuse.

In most cases the time needed for clean out would be specific to the equipment and product. But for the purposes this notebook, we implement can implement a simple clean out policy with a single non-negative parameter $t_{clean} \geq 0$ which, if specified, requires a period no less than $t_{clean}$ between the finish of one task and the start of another on every piece of equipment.

This policy is implemented by modifying the usual disjunctive constraints to avoid machine conflicts to read

$$
\begin{align}
\left[\text{start}_{j,m}+\text{dur}_{j,m} + t_{clean} \leq \text{start}_{k,m}\right] \vee \left[\text{start}_{k,m}+\text{dur}_{k,m} + t_{clean} \leq \text{start}_{j,m}\right]
\end{align}
$$

For this purpose, we write a new JobShopModel_Clean

```python
def jobshop_model_clean(TASKS, tclean=0):

    model = pyo.ConcreteModel()

    # tasks is a two dimensional set of (j,m) constructed from the dictionary keys
    model.TASKS = pyo.Set(initialize = TASKS.keys(), dimen=2)

    # the set of jobs is constructed from a python set
    model.JOBS = pyo.Set(initialize = list(set([j for (j, m) in model.TASKS])))

    # set of machines is constructed from a python set
    model.MACHINES = pyo.Set(initialize = list(set([m for (j, m) in model.TASKS])))

    # the order of tasks is constructed as a cross-product of tasks and filtering
    model.TASKORDER = pyo.Set(initialize = model.TASKS * model.TASKS, dimen=4,
        filter = lambda model, j, m, k, n: (k, n) == TASKS[(j, m)]['prec'])

    # the set of disjunctions is cross-product of jobs, jobs, and machines
    model.DISJUNCTIONS = pyo.Set(initialize = model.JOBS * model.JOBS * model.MACHINES, dimen=3,
        filter = lambda model, j, k, m: j < k and (j, m) in model.TASKS and (k, m) in model.TASKS)

    # load duration data into a model parameter for later access
    @model.Param(model.TASKS)
    def dur(model, j, m):
        return TASKS[(j, m)]['dur']

    # establish an upper bound on makespan
    ub = sum([model.dur[j, m] for (j, m) in model.TASKS])

    model.makespan = pyo.Var(bounds=(0, ub))

    model.start = pyo.Var(model.TASKS, bounds=(0, ub))

    @model.Objective(sense=pyo.minimize)
    def objective(m):
        return model.makespan

    @model.Constraint(model.TASKS)
    def finish(model, j, m):
        return model.start[j, m] + model.dur[j, m] <= model.makespan

    @model.Constraint(model.TASKORDER)
    def preceding(model, j, m, k, n):
        return model.start[k, n] + model.dur[k, n] <= model.start[j, m]

    @model.Disjunction(model.DISJUNCTIONS)
    def disjunctions(model, j, k, m):
        return [model.start[j, m] + model.dur[j, m] + tclean <= model.start[k, m],
                model.start[k, m] + model.dur[k, m] + tclean <= model.start[j, m]]

    pyo.TransformationFactory('gdp.hull').apply_to(model)
    return model

model = jobshop_model_clean(TASKS, tclean=0.5)
results = jobshop_solve(model)
visualize(results)
print("Makespan =", max([task['Finish'] for task in results]))
```

### Adding a zero wait policy

One of the issues in the use of job shop scheduling for chemical process operations are situations where there it is not possible to store intermediate materials. If there is no way to store intermediates, either in the processing equipment or in external vessels, then a **zero-wait** policy may be required.

A zero-wait policy requires subsequent processing machines to be available immediately upon completion of any task. To implement this policy, the usual precident sequencing constraint of a job shop scheduling problem, i.e.,

$$
\begin{align*}
\text{start}_{k,n}+\text{Dur}_{k,n} \leq \text{start}_{j,m}\ \ \ \ \text{for } (k,n) =\text{Prec}_{j,m}
\end{align*}
$$

is changed to

$$
\begin{align*}
\text{start}_{k,n}+\text{Dur}_{k,n} = \text{start}_{j,m}\ \ \ \ \text{for } (k,n) =\text{Prec}_{j,m}\text{ and ZW is True}
\end{align*}
$$

if the zero-wait policy is in effect.

While this could be implemented on an equipment or product specific basis, here we add an optional ZW flag to the JobShop function that, by default, is set to False.

```python
def jobshop_model_clean_zw(TASKS, tclean=0, ZW=False):

    model = pyo.ConcreteModel()

    # tasks is a two dimensional set of (j,m) constructed from the dictionary keys
    model.TASKS = pyo.Set(initialize = TASKS.keys(), dimen=2)

    # the set of jobs is constructed from a python set
    model.JOBS = pyo.Set(initialize = list(set([j for (j, m) in model.TASKS])))

    # set of machines is constructed from a python set
    model.MACHINES = pyo.Set(initialize = list(set([m for (j, m) in model.TASKS])))

    # the order of tasks is constructed as a cross-product of tasks and filtering
    model.TASKORDER = pyo.Set(initialize = model.TASKS * model.TASKS, dimen=4,
        filter = lambda model, j, m, k, n: (k, n) == TASKS[(j, m)]['prec'])

    # the set of disjunctions is cross-product of jobs, jobs, and machines
    model.DISJUNCTIONS = pyo.Set(initialize = model.JOBS * model.JOBS * model.MACHINES, dimen=3,
        filter = lambda model, j, k, m: j < k and (j, m) in model.TASKS and (k, m) in model.TASKS)

    # load duration data into a model parameter for later access
    @model.Param(model.TASKS)
    def dur(model, j, m):
        return TASKS[(j, m)]['dur']

    # establish an upper bound on makespan
    ub = sum([model.dur[j, m] for (j, m) in model.TASKS])

    model.makespan = pyo.Var(bounds=(0, ub))

    model.start = pyo.Var(model.TASKS, bounds=(0, ub))

    @model.Objective(sense=pyo.minimize)
    def objective(model):
        return model.makespan

    @model.Constraint(model.TASKS)
    def finish(mode, j, m):
        return model.start[j, m] + model.dur[j, m] <= model.makespan

    @model.Constraint(model.TASKORDER)
    def preceding(model, j, m, k, n):
        if ZW:
            return  model.start[k, n] + model.dur[k, n] == model.start[j, m]
        else:
            return model.start[k, n] + model.dur[k, n] <= model.start[j, m]

    @model.Disjunction(model.DISJUNCTIONS)
    def disjunctions(model, j, k, m):
        return [model.start[j, m] + model.dur[j, m] + tclean <= model.start[k, m],
                model.start[k, m] + model.dur[k, m] + tclean <= model.start[j, m]]

    pyo.TransformationFactory('gdp.hull').apply_to(model)
    return model

model = jobshop_model_clean_zw(TASKS, tclean=0.5, ZW=True)
results = jobshop_solve(model)
visualize(results)
print("Makespan =", max([task['Finish'] for task in results]))
```

## References

* Applegate, David, and William Cook. ["A computational study of the job-shop scheduling problem."](https://doi.org/10.1287/ijoc.3.2.149) ORSA Journal on computing 3, no. 2 (1991): 149-156. [pdf available](http://www.math.uwaterloo.ca/~bico/papers/jobshop.pdf)
* Beasley, John E. ["OR-Library: distributing test problems by electronic mail."](https://www.jstor.org/stable/pdf/2582903.pdf?casa_token=RWUXQ6e2VngAAAAA:g-dwMXjJfCUWqoLJk5y-MpDDhLmJKSGE36YU5kpkXZ71dmos25vntN0RlDFLxDme_QsOw8Oc3oi6juXCkVCfm2gcycCK0SYNoOnOZqRYDp3EqSrXP6nDeQ) Journal of the operational research society 41, no. 11 (1990): 1069-1072. [OR-Library](http://people.brunel.ac.uk/~mastjjb/jeb/info.html)
* Guéret, Christelle, Christian Prins, and Marc Sevaux. ["Applications of optimization with Xpress-MP."](http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.69.9634&rep=rep1&type=pdf) contract (1999): 00034.
* Manne, Alan S. ["On the job-shop scheduling problem."](https://doi.org/10.1287/opre.8.2.219) Operations Research 8, no. 2 (1960): 219-223.

## Exercises

### Task specific cleanout

Clean out operations are often slow and time consuming. Further more, the time required to perform a clean out frequently depends on the type of machine, and the task performed by the machine. For this exercise, create a data format to include task-specific clean out times, and model the job shop model to accomodate this additional informaton.

### Computational impact of a zero-wait policy

Repeat the benchmark problem calculation, but with a zero-wait policy. Does the execution time increase or descrease as a consequence of specifying zero-wait?


```python
import sympy as sym
from functools import reduce
from sympy.matrices import Matrix, MatrixSymbol
```

```python
sym.init_printing()
```

```python
y = MatrixSymbol('y', 1, 1)
x = MatrixSymbol('x', 1, 3)
#x = Matrix([[1,2,3]])
theta = Matrix([[1],[2],[3]])
prediction = x * theta
loss = (prediction - y)**2
loss
```

```python
sym.diff(loss, x)
```

```python

```


# Analysis of Wilhelm - "The Smallest Chemical Reaction with Bistability"

# Preliminaries

## Imports

```python
import tellurium as te
import sympy
import matplotlib.pyplot as plt
import numpy as np
from common_python.sympy import sympyUtil as su
from common_python.ODEModel.ODEModel import ODEModel
```

## Constants

```python
su.addSymbols("k_1 x k_2 k_3 y k_4")
```

# Equations

```python
stateDct = {
    x: 2 * k_1 * y - k_2 * x**2 - k_3 * x * y - k_4 * x,
    y: k_2 * x **2 - k_1 * y
}
```

# Analysis of Fixed Points

```python
stateDct[x]
```

$\displaystyle 2 k_{1} y - k_{2} x^{2} - k_{3} x y - k_{4} x$

```python
stateDct[y]
```

$\displaystyle - k_{1} y + k_{2} x^{2}$

```python
model = ODEModel(stateDct)
```

```python
valueDcts = [f.valueDct for f in model.fixedPoints]
valueDcts
```

    [{x: 0.0, y: 0.0},
     {x: -(-1.0*sqrt(k_1)*sqrt(k_2)*sqrt(0.25*k_1*k_2 - k_3*k_4)/k_3**2 + 0.5*(k_1*k_2 - 2.0*k_3*k_4)/k_3**2)*(-k_1*k_2 + k_3**2*(-1.0*sqrt(k_1)*sqrt(k_2)*sqrt(0.25*k_1*k_2 - k_3*k_4)/k_3**2 + 0.5*(k_1*k_2 - 2.0*k_3*k_4)/k_3**2) + k_3*k_4)/(k_2*k_4),
      y: -1.0*sqrt(k_1)*sqrt(k_2)*sqrt(0.25*k_1*k_2 - k_3*k_4)/k_3**2 + 0.5*(k_1*k_2 - 2.0*k_3*k_4)/k_3**2},
     {x: -(1.0*sqrt(k_1)*sqrt(k_2)*sqrt(0.25*k_1*k_2 - k_3*k_4)/k_3**2 + 0.5*(k_1*k_2 - 2.0*k_3*k_4)/k_3**2)*(-k_1*k_2 + k_3**2*(1.0*sqrt(k_1)*sqrt(k_2)*sqrt(0.25*k_1*k_2 - k_3*k_4)/k_3**2 + 0.5*(k_1*k_2 - 2.0*k_3*k_4)/k_3**2) + k_3*k_4)/(k_2*k_4),
      y: 1.0*sqrt(k_1)*sqrt(k_2)*sqrt(0.25*k_1*k_2 - k_3*k_4)/k_3**2 + 0.5*(k_1*k_2 - 2.0*k_3*k_4)/k_3**2}]

```python
sympy.simplify(valueDcts[1][x] - valueDcts[2][x])
```

$\displaystyle - \frac{2.0 \sqrt{k_{1}} \sqrt{0.25 k_{1} k_{2} - k_{3} k_{4}}}{\sqrt{k_{2}} k_{3}}$

So, require that $0.25 k_1 k_2 \geq k_3 k_4$ for real fixed point.

## Non-zero fixed points

```python
fp = valueDcts[1]
fp[x].expand()
```

$\displaystyle - \frac{1.0 \sqrt{k_{1}} \sqrt{0.25 k_{1} k_{2} - k_{3} k_{4}}}{\sqrt{k_{2}} k_{3}} + \frac{0.5 k_{1}}{k_{3}}$

```python
fp = valueDcts[2]
fp[x].expand()
```

$\displaystyle \frac{1.0 \sqrt{k_{1}} \sqrt{0.25 k_{1} k_{2} - k_{3} k_{4}}}{\sqrt{k_{2}} k_{3}} + \frac{0.5 k_{1}}{k_{3}}$

```python
fp = valueDcts[1]
fp[y].expand()
```

$\displaystyle - \frac{1.0 \sqrt{k_{1}} \sqrt{k_{2}} \sqrt{0.25 k_{1} k_{2} - k_{3} k_{4}}}{k_{3}^{2}} + \frac{0.5 k_{1} k_{2}}{k_{3}^{2}} - \frac{1.0 k_{4}}{k_{3}}$

```python
fp = valueDcts[2]
fp[y].expand()
```

$\displaystyle \frac{1.0 \sqrt{k_{1}} \sqrt{k_{2}} \sqrt{0.25 k_{1} k_{2} - k_{3} k_{4}}}{k_{3}^{2}} + \frac{0.5 k_{1} k_{2}}{k_{3}^{2}} - \frac{1.0 k_{4}}{k_{3}}$

## When are fixed points > 0?

### Value of $x$

```python
# Evaluation of x^* for fixed point 1
fp = valueDcts[1]
expandedX = fp[x].expand()
expandedX
```

$\displaystyle - \frac{1.0 \sqrt{k_{1}} \sqrt{0.25 k_{1} k_{2} - k_{3} k_{4}}}{\sqrt{k_{2}} k_{3}} + \frac{0.5 k_{1}}{k_{3}}$

This is equivalent to the following.

```python
inequality = (k_2 * k_3 * expandedX.args[0]) **2 / (k_1 * k_2) > (k_2 * k_3 * expandedX.args[1]) ** 2 / (k_1 * k_2)
inequality
```

$\displaystyle 0.25 k_{1} k_{2} > 0.25 k_{1} k_{2} - 1.0 k_{3} k_{4}$

Since all $k > 0 $, $x^{\star} > 0$.

### Value of $y$

```python
# Evaluation of y^* for fixed point 1. This should be greater than 0.
fp = valueDcts[1]
expandedY = fp[y].expand()
expandedY > 0
```

$\displaystyle - \frac{1.0 \sqrt{k_{1}} \sqrt{k_{2}} \sqrt{0.25 k_{1} k_{2} - k_{3} k_{4}}}{k_{3}^{2}} + \frac{0.5 k_{1} k_{2}}{k_{3}^{2}} - \frac{1.0 k_{4}}{k_{3}} > 0$

```python
su.addSymbols("k_5")  # k_5 = k_3 * k_4 /(k_1 * k_2)
term = k_3 **2 / (k_1 * k_2)
# lhs = term * expandedY.args[1]
rhs = term * expandedY.args[0] - term * expandedY.args[2]
rhs = rhs.subs(k_3 * k_4 /(k_1 * k_2), k_5)
term * expandedY.args[1] > - sympy.simplify(term * expandedY.args[0] - term * expandedY.args[2])
```

$\displaystyle 0.5 > \frac{1.0 k_{3} k_{4}}{k_{1} k_{2}} - \frac{1.0 \sqrt{0.25 k_{1} k_{2} - k_{3} k_{4}}}{\sqrt{k_{1}} \sqrt{k_{2}}}$

```python
# Re-written
0.5 > k_5 - sympy.sqrt(0.25  - k_5)
```

$\displaystyle k_{5} - \sqrt{0.25 - k_{5}} < 0.5$

```python
-sympy.sqrt(0.25  - k_5) < 0.5 - k_5
```

$\displaystyle - \sqrt{0.25 - k_{5}} < 0.5 - k_{5}$

```python
k_5 - 0.5 < sympy.sqrt(0.25  - k_5)
```

$\displaystyle k_{5} - 0.5 < \sqrt{0.25 - k_{5}}$

This is true for $k_5 \in [0, 0.25]$. If $k_5 > 0.25$, there is no real solution.

### Evaluation of fixed points

Plot for values of $k_5$ outside the operating region.


# Assignment 3: Combustor Design

## Introduction

The global desire to reduce greenhouse gas emissions is the main reason for the interest in the use of hydrogen for power generation.
Although hydrogen shows to be a promising solution, there are many challenges that need to be solved.
One of the challenges focuses on the use of hydrogen as a fuel in gas turbines.

In gas turbines hydrogen could replace natural gas as a fuel in the combustor. Unfortunately, this is accompanied with a technical challenge which deals with an important property in premixed combustion: the flame speed. The flame speed of hydrogen is an order of magnitude higher than natural gas due to the highly reactive nature of hydrogen. As a result a hydrogen flame is more prone to flashback than a natural gas flame.

Flame flashback is the undesired upstream propagation of a flame into the premix section of a combustor. Flashback occurs when the flame speed is higher than the velocity of the incoming fresh mixture. This could cause severe equipment damage and turbine shutdown. Adjustments to traditional combustors are required in order to guarantee safe operation when using hydrogen as fuel.

To this end the students are asked to investigate the use of hydrogen, natural gas and a blend thereof in gas turbines. The first part will focus on the impact of the different fuels on the combustor geometry. Finally, we will have a closer look at the influence of different fuels on the $CO_2$ and $NO_x$ emissions. For simplicty, it is assumed that natural gas consists purely of methane ($CH_4$).

## Tasks

### Diameter of the combustor
A gas turbine has a power output of 100 MW. The combustion section consists of 8 can combustors. Each can combustor is, for the sake of simplicty, represented by a tube with a diameter $D$.

The inlet temperature $T_2$ of the compressor is 293 K and the inlet pressure $p_2$ is 101325 Pa. To prevent damage of the turbine blades a turbine inlet temperature (TIT) of 1800 K is desired. Furthermore, assume that the specific heat of the fluid is constant through the compressor, i.e. specific heat capacity $c_{p,c}$=1.4 and a heat capacity ratio $\gamma_c$=1.4. The polytropic efficiency of the compressor and turbine are 0.90 and 0.85, respectively.

The pressure ratio over the compressor will depend on your studentID:

PR = 10 if (numpy.mod(studentID, 2) + 1) == 1

PR = 20 if (numpy.mod(studentID, 2) + 1) == 2

Assume the TIT to be equal to the temperature of the flame inside the combustor. The flame temperature depends on the equivalence ratio ($\phi$), the hydrogen volume percentage of the fuel ($H_2\%$) and the combustor inlet temperature and pressure. For now consider the fuel to consist of pure natural gas ($H_2\%=0$). Note that the equivalence ratio is given by:

\begin{align}
\phi = \frac{\frac{m_{fuel}}{m_{air}}}{(\frac{m_{fuel}}{m_{air}})_{stoich}}
\end{align}

**1. Calculate the inlet temperature $T_3$ and inlet pressure $p_3$ of the combustor and determine the required equivalence ratio (adjust PART A and PART B and run the code), so that the TIT specification is met.**

Inside the combustor the flow is turbulent. Turbulence causes an increase in the flame speed, so that the turbulent flame speed $S_T \approx 10 S_L$.

**2. With the equivalence ratio determined in the previous question, calculate the total mass flow rate ($\dot{m}_{tot}$) through the gas turbine and the maximal diameter $D$ of a single combustor tube, so that flashback is prevented. Adjust PART A, PART B, PART C and PART D in the code and run it again. Report the steps you have taken.

Is there also a minimum diameter? If so, no calculation required, discuss what could be the reason for the necessity of a minimal diameter of the combustor tube.**

The combustion of methane is represented by the reaction: $CH_4 + 2 (O_2 + 3.76 N_2) \rightarrow CO_2 + 2 H_2O + 7.52 N_2$

**3. Use the above reaction equation and the definition of $\phi$ to find the mass flow rate of the fuel $\dot{m}_{fuel}$.**

**4. Calculate the total heat input using $\dot{m}_{fuel}$ and calculate the efficiency of the complete cycle.**

**5. Repeat tasks 1-4 for a fuel consisting of $50\%H_2$/$50\%CH_4$ and $100\%H_2$. Discuss the effect of the addition of hydrogen to the fuel on the combustor geometry and cycle performance.**

### $CO_2$ and $NO_x$ emissions

**6. A gas turbine manufacturer claims that their gas turbines can be fired with a hydrogen content of 30%. Discuss wheter this could be regarded an achievement (use the top plot in Figure 5).**

**7. Consider an equivalence ratio $\phi=0.5$. Regarding emissions, discuss the advantages and disadvantages of increasing the hydrogen content of the fuel. Adjust PART A and use Figure 5.**

### Bonus assignment
For simplicty, it was assumed that natural gas does consist of pure methane. In reality, it could be a mix of methane, higher hydrocarbons and nitrogen.
An example is Dutch Natural Gas (DNG), which consists of $80\%CH_4$, $5\%C_2H_6$ and $15\%N_2$.

**Repeat tasks 1-4 for a fuel consisting of $50\%H_2$/$50\%DNG$.
 Hint1: Nitrogen does not participate in the reaction.
 Hint2: This requires more adjustment of the code than just PARTS A, B, C, D.**

## Code
Two properties of importance in this assignment are the laminar flame speed $S_L$ and the adiabtic flame temperature $T_{ad}$ of a mixture. These properties can be determined by solving the equations for continuity, momentum, species and energy in one dimension. Fortunetaly, we do not need to solve these equations by hand, instead a chemical kinetics software (Cantera) is used to solve these equations by running a simulation. The simulation is illustrated in the sketch below. Keep in mind that the simulation can take some time to complete.

For more information about Cantera visit: https://cantera.org/.

For more background information regarding the 1D flame simulation visti: https://cantera.org/science/flames.html

```python
#%% Load required packages
import sys
import cantera as ct
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import cm

#%% Constants
R_gas_mol = 8314 # Universal gas constant [units: J*K^-1*kmol^-1]
R_gas_mass = 287 # universal gas constant [units: J*K^-1*kg^-1]

#%% Start
# Power output of turbine
power_output = 100 # units: MW
power_output*=1e6

# Compressor and turbine polytropic efficiencies
etap_c = 1
etap_t = 1

# Pressure ratio
PR = 10

# Compressor inlet temperature and pressure
T2 = 293.15 # units: K
p2 = 101325 # units: Pa

# Heat capacity ratio of air at T=293.15 K
gam_c = 1.4

# Compressor stage
# Specific heat capacity (heat capacity per unit mass) of mixture in compressor
cp_c = R_gas_mass*gam_c/(gam_c-1)
# cp_c = 1006 # units: J.kg^-1.K^-1
# cv_c = 717 # units: J.kg^-1.K^-1

# Molar mass of species [units: kg*kmol^-1]
M_H = 1.008
M_C = 12.011
M_N = 14.007
M_O = 15.999
M_H2 = M_H*2
M_CH4 = M_C + M_H*4
M_CO2 = M_C + M_O*4
M_O2 = M_O*2
M_N2 = M_N*2

# Define volume fractions of species in air [units: -]
f_O2 = 0.21
f_N2 = 0.79

########## PART A: ADJUST CODE HERE ##########

# Equivalence ratios
phis = [None, None, None] # Set equivalence ratios ranging from 0.4 to 0.8

# Hydrogen percentages
H2_percentages = [None, None, None] # Set hydrogen volume percentages of the fuel ranging from 0 to 100

################# END PART A ##################

# Define colors to make distinction between different mixtures based on hydrogen percentage
colors = cm.rainbow(np.linspace(0, 1, len(H2_percentages)))

#%% Premixed flame object
class mixture_class:

    def __init__(self, phi, H2_percentage, T_u=293.15, p_u=101325):

        # Color and label for plots
        self.color = colors[H2_percentages.index(H2_percentage)]
        self.label = str(int(H2_percentage)) + r'$\% H_2$'

        # Temperature and pressure of the unburnt mixture
        self.T_u = T_u # units: K
        self.p_u = p_u # units: Pa

        # Equivalence ratio
        self.phi = phi

        # Hydrogen percentage of fuel
        self.H2_percentage = H2_percentage

        # DNG percentage of fuel
        self.CH4_percentage = 100 - self.H2_percentage

        # Volume fractions of fuel
        self.f_H2 = self.H2_percentage/100
        self.f_CH4 = self.CH4_percentage/100

        # Mass densities of fuel species
        rho_H2 = M_H2*self.p_u/(self.T_u*R_gas_mol)
        rho_CH4 = M_CH4*self.p_u/(self.T_u*R_gas_mol)

        # Check if volume fractions of fuel and air are correct
        check_air = f_O2 + f_N2
        check_fuel = self.f_H2 + self.f_CH4

        if check_air == 1.0 and round(check_fuel,3) == 1.0:
            pass
        else:
            sys.exit("fuel or air composition is incorrect!")

        if round(check_fuel,3) == 1.0:
            pass
        else:
            sys.exit("fuel composition is incorrect!")

        # Definition of the mixture
        # 1. Set the reaction mechanism
        self.gas = ct.Solution('gri30.cti')

        # 2. Define the fuel and air composition
        fuel = {'H2':self.f_H2, 'CH4':self.f_CH4}
        air = {'N2':f_N2/f_O2, 'O2':1.0}

        # 3. Set the equivalence ratio
        self.gas.set_equivalence_ratio(phi, fuel, air)

        # 4. Set the transport model
        self.gas.transport_model= 'Multi'

        # 5. Set the unburnt mixture temperature and pressure
        self.gas.TP = T_u, p_u

        # Unburnt mixture properties
        self.h_u = self.gas.enthalpy_mass # units: J.kg^-1
        self.cp_u = self.gas.cp_mass # units: J*K^-1*kg^-1
        self.cv_u = self.gas.cv_mass # units: J*K^-1*kg^-1
        self.rho_u = self.gas.density_mass # units: kg.m^-3
        self.rho_u_H2 = rho_H2 # units: kg.m^-3
        self.rho_u_CH4 = rho_CH4 # units: kg.m^-3
        self.mu_u = self.gas.viscosity # Pa.s
        self.nu_u = self.mu_u/self.rho_u # units: m^2.s^-1
        self.lambda_u= self.gas.thermal_conductivity # units: W.m^-1.K^-1
        self.alpha_u = self.lambda_u/(self.rho_u*self.cp_u) # units: m^2.s^-1

    def solve_equations(self):

        # Unburnt molar fractions
        self.X_H2 = self.gas["H2"].X[0]
        self.X_CH4 = self.gas["CH4"].X[0]
        self.X_O2 = self.gas["O2"].X[0]
        self.X_N2 = self.gas["N2"].X[0]

        # Set domain size (1D)
        width = 0.05  # units: m

        # Create object for freely-propagating premixed flames
        flame = ct.FreeFlame(self.gas, width=width)

        # Set the criteria used to refine one domain
        flame.set_refine_criteria(ratio=3, slope=0.1, curve=0.1)

        # Solve the equations
        flame.solve(loglevel=0, auto=True)

        # Result 1: Laminar flame speed
        self.S_L0 = flame.velocity[0]*100 # units: cm.s^-1
        self.S_T = 10*self.S_L0/100 # units: m.s^-1 Rough estimation of the turbulent flame speed

        # Result 2: Adiabtaic flame temperature
        self.T_ad = self.gas.T

        # Burnt mixture properties
        self.h_b = self.gas.enthalpy_mass # units: J.kg^-1
        self.cp_b = self.gas.cp_mass # units: J*K^-1*kg^-1
        self.cv_b = self.gas.cv_mass # units: J*K^-1*kg^-1
        self.rho_b = self.gas.density_mass # units: kg.m^-3
        self.mu_b = self.gas.viscosity # Pa.s
        self.nu_b = self.mu_b/self.rho_b # units: m^2.s^-1
        self.lambda_b = self.gas.thermal_conductivity # units: W.m^-1.K^-1
        self.alpha_b = self.lambda_b/(self.rho_b*self.cp_b) # units: m^2.s^-1

        # Burnt mixture molar fractions
        self.X_CO2 = self.gas["CO2"].X[0]
        self.X_NO = self.gas["NO"].X[0]
        self.X_NO2 = self.gas["NO2"].X[0]

#%% Function to retrieve the LHV of different kind of fuels
def heating_value(fuel):
    """ Returns the LHV and HHV for the specified fuel """
    T_u = 293.15
    p_u = 101325
    gas1 = ct.Solution('gri30.cti')
    gas1.TP = T_u, p_u
    gas1.set_equivalence_ratio(1.0, fuel, 'O2:1.0')
    h1 = gas1.enthalpy_mass
    Y_fuel = gas1[fuel].Y[0]

    # complete combustion products
    Y_products = {'CO2': gas1.elemental_mole_fraction('C'),
                  'H2O': 0.5 * gas1.elemental_mole_fraction('H'),
                  'N2': 0.5 * gas1.elemental_mole_fraction('N')}

    gas1.TPX = None, None, Y_products
    h2 = gas1.enthalpy_mass
    LHV = -(h2-h1)/Y_fuel
    return LHV

# Lower Heating Values of well-known combustion fuels
LHV_H2 = heating_value('H2')
LHV_CH4 = heating_value('CH4')
LHV_C2H6 = heating_value('C2H6')

#%% Create list of flame objects for multiple mixtures depending on the equivalence ratio
# and the percentage of hydrogen in the fuel (volume based)

# Initialize list for flame objects
mixtures = []

# Create flame objects and start simulations
for phi in phis:
    for H2_percentage in H2_percentages:

        ########## PART B: ADJUST CODE HERE ##########

        # Compressor stage
        # Temperature after compressor stage
        T3 = None # units: K
        p3 = None # units: Pa

        # Combustor inlet temperature in K and pressure in Pa
        T_u = T3 # units: K
        p_u = p3 # units: Pa

        ################# END PART B ##################

        # Combustor stage
        # Define unburnt mixture that goes into the combustor
        mixture = mixture_class(phi, H2_percentage, T_u, p_u)

        # Solve equations and obtain burnt mixture properties
        mixture.solve_equations()

        # Append the mixture (with unburnt and burnt properties) to list of mixtures
        mixtures.append(mixture)

        # Turbine stage
        # Heat capacity ratio of mixture in turbine
        gam_t = mixture.cp_b/mixture.cv_b

        # Turbine inlet temperature
        T4 = mixture.T_ad

        ########## PART C: ADJUST CODE HERE ##########

        # Turbine outlet temperature
        T5 = None

        ################# END PART C ##################

        print('mixture solved: phi=' + str(phi) + ', H2%=' + str(H2_percentage))

#%% Plots A: Laminar flame speed/adiabatic flame temperture vs equivalence ratio
plt.close('all')

# Plot parameters
fontsize = 12
marker = 'o'
markersize = 8
linewidth = 1
linestyle = 'None'

# Figure 1: Laminar flame speed vs equivalence ratio
fig1, ax1 = plt.subplots()
ax1.set_xlabel(r'$\phi$ [-]', fontsize=fontsize)
ax1.set_ylabel(r'$S_L$ [cm.s$^{-1}$]', fontsize=fontsize)
ax1.set_xlim(0.3, 1.1)
ax1.set_ylim(0, 250)
ax1.set_title('Laminar flame speed vs. equivalence ratio \n $T_u=$' + str(round(T_u,2)) + ' K, $p_u$=' + str(p_u*1e-5) + ' bar')
ax1.grid()

# Figure 2: Adiabatic flame temperature vs equivalence ratio
fig2, ax2 = plt.subplots()
ax2.set_xlabel(r'$\phi$ [-]', fontsize=fontsize)
ax2.set_ylabel(r'$T_{ad}$ [K]', fontsize=fontsize)
ax2.set_xlim(0.3, 1.1)
ax2.set_ylim(1200, 2800)
ax2.grid()
ax2.set_title('Adiabtic flame temperature vs. equivalence ratio \n $T_u=$' + str(round(T_u,2)) + ' K, $p_u$=' + str(p_u*1e-5) + ' bar')

# Initialize list for laminar flame speeds
S_L0_lists = [[] for i in range(len(H2_percentages))]
# Initialize list for adiabatic flame temperatures
T_ad_lists = [[] for i in range(len(H2_percentages))]

# Fill Figure 1 and 2
for mixture in mixtures:

    index = H2_percentages.index(mixture.H2_percentage)

    ax1.plot(mixture.phi, mixture.S_L0, ls=linestyle,  marker=marker, ms=markersize, c=mixture.color, label=mixture.label if mixture.phi == phis[0] else "")
    ax2.plot(mixture.phi, mixture.T_ad, ls=linestyle,  marker=marker, ms=markersize, c=mixture.color, label=mixture.label if mixture.phi == phis[0] else "")

    S_L0_lists[index] = np.append(S_L0_lists[index], mixture.S_L0)
    T_ad_lists[index] = np.append(T_ad_lists[index], mixture.T_ad)

# Plot polynomial fits to show trends for laminar flame speed and adiabatic flame temperature as a function of the equivalence ratio
if len(phis) == 1:
    pass
else:
    # Create zipped lists for polynomial fits
    lists_zipped = zip(S_L0_lists, T_ad_lists, colors)

    # Order of polynomial
    poly_order = 3

    for (S_L0, T_ad, color) in lists_zipped:
        # Create new array for phi
        phis_fit = np.linspace(phis[0], phis[-1])

        # Plot 4th order polynomial fit for laminar flame speed
        coeff_S_L0 = np.polyfit(phis, S_L0, poly_order)
        poly_S_L0 = np.poly1d(coeff_S_L0)
        S_L0_fit = poly_S_L0(phis_fit)
        ax1.plot(phis_fit, S_L0_fit, ls="--", c=color)

        # Plot 4th order polynomial fit for adiabatic flame temperature
        coeff_T_ad = np.polyfit(phis, T_ad, poly_order)
        poly_T_ad = np.poly1d(coeff_T_ad)
        T_ad_fit = poly_T_ad(phis_fit)
        ax2.plot(phis_fit, T_ad_fit, ls="--", c=color)

#% Plots B: Fuel blend properties and emissions
# Assume constant power (or heat input): heat_input = m_H2_dot*LHV_H2 + m_CH4_dot*LHV_CH4 = 1 (constant)

# Plot parameters
x_ticks = np.linspace(0, 100, 11)
y_ticks = x_ticks
bin_width = 5

# Initialize lists
H2_fraction_heat_input, H2_fraction_mass, CH4_fraction_heat_input, CH4_fraction_mass, CO2_fraction, fuel_energy_mass = ([] for i in range(6))

# Densities of hydrogen and methane of unburnt mixture
rho_H2 = mixture.rho_u_H2
rho_CH4 = mixture.rho_u_CH4

# Reference: Amount of CO2 when H2%=0 (1 mol of CH4 == 1 mol CO2)
Q_CO2_ref = 1 / (rho_CH4*LHV_CH4)

# Hydrogen fraction in the fuel
H2_fraction_volume = np.linspace(0, 1, 21)

# Mixture calculations
for x in H2_fraction_volume:

    # Fractions of H2 and CH4 by heat input
    H2_part = rho_H2*LHV_H2*x
    CH4_part = rho_CH4*LHV_CH4*(1-x)
    H2_fraction_heat_input_i = H2_part / (H2_part + CH4_part)
    CH4_fraction_heat_input_i = 1 - H2_fraction_heat_input_i
    H2_fraction_heat_input = np.append(H2_fraction_heat_input, H2_fraction_heat_input_i)
    CH4_fraction_heat_input = np.append(CH4_fraction_heat_input, CH4_fraction_heat_input_i)

    # Fraction of CO2 reduction
    Q_u_i = 1 / (rho_H2*LHV_H2*x + rho_CH4*LHV_CH4*(1-x))
    Q_CH4_i = Q_u_i*(1-x)
    Q_CO2_i = Q_CH4_i
    CO2_fraction_i = Q_CO2_i/Q_CO2_ref
    CO2_fraction = np.append(CO2_fraction, CO2_fraction_i)

    # Fractions of H2 and CH4 by mass
    H2_part = rho_H2*x
    CH4_part = rho_CH4*(1-x)
    H2_fraction_mass_i = H2_part / (H2_part + CH4_part)
    CH4_fraction_mass_i = 1- H2_fraction_mass_i
    H2_fraction_mass = np.append(H2_fraction_mass, H2_fraction_mass_i)
    CH4_fraction_mass = np.append(CH4_fraction_mass, CH4_fraction_mass_i)

    # Fuel energy content
    fuel_energy_mass_i = (H2_fraction_mass_i*LHV_H2 + CH4_fraction_mass_i*LHV_CH4)/1e6 # units: MJ.kg^-1
    fuel_energy_mass = np.append(fuel_energy_mass, fuel_energy_mass_i)

# Convert fractions to percentages
CO2_percentage = 100*CO2_fraction
CO2_reduction_percentage = 100 - CO2_percentage
H2_percentage_volume = H2_fraction_volume*100
CH4_percentage_heat_input = CH4_fraction_heat_input*100
H2_percentage_heat_input = H2_fraction_heat_input*100
H2_percentage_mass = 100*H2_fraction_mass
CH4_percentage_mass = 100*CH4_fraction_mass

# Plots

fig3, ax3 = plt.subplots()
line3_0 = ax3.plot(H2_percentage_volume, CH4_percentage_heat_input, marker=marker, color='tab:blue', label=r'$CH_{4}$')
ax3.set_xticks(x_ticks)
ax3.set_yticks(y_ticks)
ax3.set_xlabel(r'$H_{2}$% (by volume)', fontsize=fontsize)
ax3.set_ylabel(r'$CH_{4}$% (by heat input)', fontsize=fontsize, color='tab:blue')
ax3.set_title('Heat input vs. volume percentage for a methane/hydrogen fuel blend')
ax3.grid()
ax3_1 = ax3.twinx()
line3_1 = ax3_1.plot(H2_percentage_volume, H2_percentage_heat_input, marker=marker, color='tab:orange', label=r'$H_{2}$')
ax3_1.set_ylabel(r'$H_{2}$% (by heat input)', fontsize=fontsize, color='tab:orange')
lines3 = line3_0 + line3_1
labels3 = [l.get_label() for l in lines3]

fig4, ax4 = plt.subplots()
ax4.plot(H2_percentage_heat_input, CO2_reduction_percentage, marker=marker)
ax4.set_xticks(x_ticks)
ax4.set_yticks(y_ticks)
ax4.set_xlabel(r'$H_{2}$% (by heat input)', fontsize=fontsize)
ax4.set_ylabel(r'$CO_2$ reduction [%]', fontsize=fontsize)
ax4.set_title(r'$CO_2$ emissions vs. hydrogen/methane fuel blends (heat input %)')
ax4.grid()

fig5, (ax5, ax5_1) = plt.subplots(2)
ax5.plot(H2_percentage_volume, CO2_percentage, marker=marker)
ax5.set_xticks(x_ticks)
ax5.set_yticks(y_ticks)
ax5.set_xlabel(r'$H_{2}$% (by volume)', fontsize=fontsize)
ax5.set_ylabel(r'$CO_2$ emissions [%]', fontsize=fontsize)
ax5.set_title(r'$CO_2$ emissions vs. hydrogen/methane fuel blends (volume %)')
ax5.grid()

for i, mixture in enumerate(mixtures):
    if mixture.phi == phis[-1]:
        NO_percentage_volume = mixture.X_NO*100
        NO2_percentage_volume = mixture.X_NO2*100
        ax5_1.bar(mixture.H2_percentage, NO_percentage_volume, bin_width, color='tab:red', label=r'$NO$' if i == 0 else "")
        ax5_1.bar(mixture.H2_percentage, NO2_percentage_volume, bin_width, bottom=NO_percentage_volume, color='tab:blue', label=r'$NO_2$' if i == 0 else "")
        ax5_1.set_title(r'$NO_x$ emissions for $\phi=$' + str(mixture.phi))

ax5_1.set_xlim(-5, 105)
ax5_1.set_xticks(x_ticks)
ax5_1.set_xlabel(r'$H_{2}$% (by volume)', fontsize=fontsize)
ax5_1.set_ylabel(r'$NO_x$ [%]', fontsize=fontsize)
ax5_1.grid()

fig6, ax6 = plt.subplots()
ax6.plot(H2_percentage_volume, H2_percentage_mass, marker=marker, color='tab:blue', label=r'$H_{2}$')
ax6.plot(H2_percentage_volume, CH4_percentage_mass, marker=marker, color='tab:orange', label=r'$CH_{4}$')
ax6.set_xticks(x_ticks)
ax6.set_yticks(y_ticks)
ax6.set_xlabel(r'$H_{2}$% (by volume)', fontsize=fontsize)
ax6.set_ylabel(r'wt.% (by mass)', fontsize=fontsize)
ax6.set_title(r'Weight vs. volume percentage for hydrogen/methane fuel blends')
ax6.grid()

fig7, ax7 = plt.subplots()
ax7.plot(H2_percentage_volume, fuel_energy_mass, lw=2, marker=marker, color='tab:red')
ax7.set_xticks(x_ticks)
ax7.set_xlabel(r'$H_{2}$% (by volume)', fontsize=fontsize)
ax7.set_ylabel(r'Fuel energy content [MJ.kg$^{-1}$]', fontsize=fontsize)
ax7.set_title(r'Energy content vs. volume percentage for hydrogen/methane fuel blends')
ax7.grid()

# Turn on legends
ax1.legend()
ax2.legend()
ax3.legend(lines3, labels3, loc='center left')
ax5.legend(bbox_to_anchor=(1, 1))
ax5_1.legend(bbox_to_anchor=(1, 1))
ax6.legend(bbox_to_anchor=(1, 1))

# Fix figures layout
fig1.tight_layout()
fig2.tight_layout()
fig3.tight_layout()
fig4.tight_layout()
fig5.tight_layout()
fig6.tight_layout()
fig7.tight_layout()

# Uncomment to save figures as .svg
# fig1.savefig('turbo3_1.svg')
# fig2.savefig('turbo3_2.svg')
# fig3.savefig('turbo3_3.svg')
# fig4.savefig('turbo3_4.svg')
# fig5.savefig('turbo3_5.svg')
# fig6.savefig('turbo3_6.svg')
# fig7.savefig('turbo3_7.svg')

########## PART D: ADJUST CODE HERE ##########

for mixture in mixtures:

    # Equivalence of the mixture
    phi = mixture.phi

    # Density of the unburnt mixture (before entering the combustor)
    rho_u = mixture.rho_u # units: kg.s^-1

    # Volumtric fractions of hydrogen and methane
    f_H2 = mixture.f_H2
    f_CH4 = mixture.f_CH4

    # Specific heat capacity (heat capacity per unit mass)
    cp_t = mixture.cp_b

    # Total mass flow rate
    m_dot = None # units: kg.s^-1

    # Velocity in the combustor
    V = mixture.S_T # units: m.s^-1

    # Area and diameter of the combustor
    A = None # units: m^2
    D = None # units: m

    # ratio of mass fuel and mass air at stoichiometric conditions
    m_f_over_m_a_stoich = None

    # Mass flow of the fuel
    m_f_dot = None

    # Heat input or Heat of combustion
    Q = None # units: W

    # thermal cycle efficiency
    eta_cycle = None

################# END PART D ##################

```

```python

```


```python
from __future__ import division
import math
import numpy as np
import numpy.random
import scipy as sp
import scipy.stats
from hypergeometric import *

%matplotlib inline
import matplotlib
import matplotlib.pyplot as plt

```

# Tri-hypergeometric test using the diluted margin as the test statistic

We will sample individual ballots without replacement from a population of $N$ ballots, of which
$N_w$ ballots have a vote for $w$ but not for $\ell$, $N_\ell$ ballots have a vote for $\ell$ but not for $w$, and $N_u = N - N_w - N_\ell$ ballots have votes for both $w$ and $\ell$ or neither $w$ nor $\ell$, including undervotes and invalid ballots.
We might draw a simple random sample of $n$ ballots ($n$ fixed ahead of time), or we might draw
sequentially, so the sample size $B$ could be random, including stopping rules that depend on the data.

Regardless, we assume that, conditional on the attained sample size $n$, the ballots are a simple random sample
of size $n$ from the $N$ ballots in the population.
Then, conditional on $B=n$, the joint distribution of
$(B_w, B_\ell, B_u)$ is tri-hypergeometric.

That is,

\begin{equation}
    \mathbb{P}_{N_w, N_\ell, N} \{ B_w = i, B_\ell = j \vert B=n \} =
     \frac{ {N_w \choose i}{N_\ell \choose j}{N - N_w - N_\ell \choose n-i-j}}{{N \choose n}}.
\end{equation}

To test, we will condition on the event $B=n$. (In contrast, the BRAVO ballot-polling
method conditions on $B_w+B_\ell = m$;
we might instead consider conditioning on $B=n$ _and_ $B_w + B_\ell = m$.)
The test statistic will be the diluted sample margin, $D \equiv (B_w - B_\ell)/B$.
This is the sample difference in the number of ballots for the winner and for the loser, divided by the
total number of ballots in the sample.
We want to test the compound hypothesis $N_w - N_\ell \le c$ against the compound alternative
$N_w - N_\ell > c$.
(To use Wald's Sequential Probability Ratio Test, we might pick a simple alternative instead, e.g.,
$N_w = V_w$ and $N_\ell = V_\ell$, the reported values, assuming $V_w - V_\ell > c$.)
Hence, we will reject for large values of $D$.
Conditional on $B=n$, the event $D = (B_w - B_\ell)/B = d$ is the event $B_w - B_\ell = nd$.

Suppose we observe $D=d$.
The $p$-value of the simple hypothesis that there are $N_w$ ballots with
a vote for $w$ but not for $\ell$, $N_\ell$ ballots with a vote for $\ell$ but not for $w$, and $N - N_w - N_\ell$ ballots with votes for both $w$ and $\ell$ or neither $w$ nor $\ell$ (including undervotes and
invalid ballots) is

\begin{equation}
   \mathbb{P}_{N_w, N_\ell, N} \left \{ D \geq d \vert B = n \right \} =
   \mathbb{P}_{N_w, N_\ell, N}\{ B_w - B_\ell \geq nd \vert B = n \} =
   \sum_{(i, j) : i - j \geq nd; i+j \leq n; i, j\geq 0}
         \frac{ {N_w \choose i}{N_\ell \choose j}{N - N_w - N_\ell \choose n-i-j}}{{N \choose n}}.
\end{equation}
Below we show empirically (but do not prove) that this tail probability is monotone increasing in $N_w$.

This expression can be simplified for numerical stability:

\begin{equation}
\frac{ {N_w \choose i}{N_\ell \choose j}{N - N_w - N_\ell \choose n-i-j}}{{N \choose n}} =
  \frac{ N_w! N_\ell! (N - N_w - N_\ell)! n! (N-n)!}
       {i! (N_w-i)! j! (N_\ell - j)! (n-i-j)! (N - N_w - N_\ell - n +i+j)! N!}
\end{equation}

The composite null hypothesis does not specify $N_w$ or $N_\ell$ separately, only that $N_w - N_\ell \le c$ for
some fixed, known $c$.
The (conditional) $p$-value of this composite hypothesis for $D=d$ is
\begin{equation}
  \max_{N_w, N_\ell \in \{0, 1, \ldots, N \}: N_w - N_\ell \le c, N_w + N_\ell \le N}
   \sum_{(i, j) : i - j \geq nd; i+j \leq n; i, j\geq 0}
         \frac{ {N_w \choose i}{N_\ell \choose j}{N - N_w - N_\ell \choose n-i-j}}{{N \choose n}},
\end{equation}
wherever the summand is defined.
(Equivalently, define ${N \choose M} \equiv 0$ if $M > N$, $M < 0$, or $N \le 0$.)
If the empirical result mentioned above is true, then finding the maximum is trivial; if not,
it is still at most a one-dimensional optimization problem.

## Experiment 1 - few invalids

Testing the hypothesis $\pi_w - \pi_\ell \leq 0$ -- the same as testing whether the winner and the loser got exactly the same vote share. This is in blue.
Then we also test the hypotheses $N_w - N_\ell \leq 1$ (in green) and $N_w - N_\ell \leq 2$ (in red).

Let $N = 100$, $n=30$, $n_w = 15$, $n_\ell=10$, and $n_u = 5$.

Below, we plot the $p$-values for varying values of $N_w$.

The $p$-value is maximized for each contour at the largest possible value of $N_w$, and the contours have increasing $p$-values as the constant $N_w - N_\ell$ decreases.

```python
pvalues = np.zeros(47)
for B in range(47):
    pvalues[B] = diluted_margin_trihypergeometric(w=15, l=10, n=30, N_w=B, N_l=B, N=100)

pvalues2 = np.zeros(47)
for B in range(47):
    pvalues2[B] = diluted_margin_trihypergeometric(w=15, l=10, n=30, N_w=B+1, N_l=B, N=100)

pvalues3 = np.zeros(47)
for B in range(47):
    pvalues3[B] = diluted_margin_trihypergeometric(w=15, l=10, n=30, N_w=B+2, N_l=B, N=100)

plt.plot(range(47), pvalues, 'blue',
         range(1,48), pvalues2, 'green',
         range(2,49), pvalues3, 'red')
plt.xlabel("N_w")
plt.ylabel("p-value")
```

## Experiment 2 - many invalids

Testing the hypothesis $\pi_w - \pi_\ell \leq 0$ -- the same as testing whether the winner and the loser got exactly the same vote share. This is in blue.
Then we also test the hypotheses $N_w - N_\ell \leq 1$ (in green) and $N_w - N_\ell \leq 2$ (in red).

Let $N = 100$, $n=30$, $n_w = 8$, $n_\ell=7$, and $n_u = 15$.

Below, we plot the $p$-values for varying values of $N_w$.

The $p$-value is maximized for each contour at the largest possible value of $N_w$, and the contours have increasing $p$-values as the constant $N_w - N_\ell$ decreases.

```python
pvalues = np.zeros(47)
for B in range(47):
    pvalues[B] = diluted_margin_trihypergeometric(w=8, l=7, n=30, N_w=B, N_l=B, N=100)

pvalues2 = np.zeros(47)
for B in range(47):
    pvalues2[B] = diluted_margin_trihypergeometric(w=8, l=7, n=30, N_w=B+1, N_l=B, N=100)

pvalues3 = np.zeros(47)
for B in range(47):
    pvalues3[B] = diluted_margin_trihypergeometric(w=8, l=7, n=30, N_w=B+2, N_l=B, N=100)

plt.plot(range(47), pvalues, 'blue',
         range(47), pvalues2, 'green',
         range(47), pvalues3, 'red')
plt.xlabel("N_w")
plt.ylabel("p-value")
```

# Hypergeometric test

Conditional on $B_{w\ell} = n_w+n_\ell$, the $p$-value is

\begin{equation}
\mathbb{P} \{ d \geq \frac{n_w - n_\ell}{n} | B_{w\ell} = n_w+n_\ell \} = \mathbb{P}\{ B_w - B_\ell \geq n_w - n_\ell | B_{w\ell} = n_w+n_\ell \}
   = \sum_{(i, j) : i - j \geq n_w-n_\ell; i+j = n_w+n_\ell; i, j\geq 0} \frac{ {N_w \choose i}{N_\ell \choose j}}{{N_{w\ell} \choose n_w+n_\ell}}.
\end{equation}

Under the null hypothesis, $N_w - N_\ell \leq c$ for specified $c$, but neither parameter value is known.

We show empirically, but not analytically, below that the tail probability is monotone increasing in $N_w$.

Actually compute the $p$-value as

\begin{equation}
\mathbb{P} \{ d \geq \frac{n_w - n_\ell}{n} | B_{w\ell} = n_w+n_\ell \} = \mathbb{P}\{ B_w - B_\ell \geq n_w - n_\ell | B_{w\ell} = n_w+n_\ell \}
   = \sum_{(i, j) : i - j \geq n_w-n_\ell; i+j = n_w+n_\ell; i, j\geq 0} \frac{ {N_\ell+c \choose i}{N_\ell \choose j}}{{2N_\ell+c \choose n_w+n_\ell}}.
\end{equation}

## Experiment 1 - few invalids

Testing the hypothesis $\pi_w - \pi_\ell \leq 0$ -- the same as testing whether the winner and the loser got exactly the same vote share. This is in blue.
Then we also test the hypotheses $N_w - N_\ell \leq 1$ (in green) and $N_w - N_\ell \leq 2$ (in red).

Let $N = 100$, $n=30$, $n_w = 15$, $n_\ell=10$, and $n_u = 5$.

Below, we plot the $p$-values for varying values of $N_w$.

The $p$-value is maximized for each contour at the largest possible value of $N_w$, and the contours have increasing $p$-values as the constant $N_w - N_\ell$ decreases.

```python
pvalues = np.zeros(47)
for B in range(1,47):
    pvalues[B] = diluted_margin_hypergeometric(w=15, l=10, N_w=B, N_l=B)

pvalues2 = np.zeros(47)
for B in range(1,47):
    pvalues2[B] = diluted_margin_hypergeometric(w=15, l=10, N_w=B+1, N_l=B)

pvalues3 = np.zeros(47)
for B in range(1,47):
    pvalues3[B] = diluted_margin_hypergeometric(w=15, l=10, N_w=B+2, N_l=B)

plt.plot(range(47), pvalues, 'blue',
         range(1, 47+1), pvalues2, 'green',
         range(2, 47+2), pvalues3, 'red')
plt.xlabel("N_w")
plt.ylabel("p-value")
```

## Experiment 1.5 - w < l

Testing the hypothesis $\pi_w - \pi_\ell \leq 0$ -- the same as testing whether the winner and the loser got exactly the same vote share. This is in blue.
Then we also test the hypotheses $N_w - N_\ell \leq 1$ (in green) and $N_w - N_\ell \leq 2$ (in red).

Let $N = 100$, $n=30$, $n_w = 10$, $n_\ell=15$, and $n_u = 5$.

Below, we plot the $p$-values for varying values of $N_w$.

The $p$-value is maximized for each contour at the largest possible value of $N_w$, and the contours have increasing $p$-values as the constant $N_w - N_\ell$ decreases.

```python
pvalues = np.zeros(47)
for B in range(1,47):
    pvalues[B] = diluted_margin_hypergeometric(w=10, l=15, N_w=B, N_l=B)

pvalues2 = np.zeros(47)
for B in range(1,47):
    pvalues2[B] = diluted_margin_hypergeometric(w=10, l=15, N_w=B+1, N_l=B)

pvalues3 = np.zeros(47)
for B in range(1,47):
    pvalues3[B] = diluted_margin_hypergeometric(w=10, l=15, N_w=B+2, N_l=B)

plt.plot(range(47), pvalues, 'blue',
         range(1, 47+1), pvalues2, 'green',
         range(2, 47+2), pvalues3, 'red')
plt.xlabel("N_w")
plt.ylabel("p-value")
```

## Experiment 2 - c < 0

Testing the hypothesis $\pi_w - \pi_\ell \leq 0$ -- the same as testing whether the winner and the loser got exactly the same vote share. This is in blue.
Then we also test the hypotheses $N_w - N_\ell \leq -1$ (in green) and $N_w - N_\ell \leq -2$ (in red).

Let $N = 100$, $n=30$, $n_w = 15$, $n_\ell=10$, and $n_u = 5$.

Below, we plot the $p$-values for varying values of $N_w$.

The $p$-value is maximized for each contour at the largest possible value of $N_w$, and the contours have increasing $p$-values as the constant $N_w - N_\ell$ decreases.

```python
pvalues = np.zeros(47)
for B in range(1,47):
    pvalues[B] = diluted_margin_hypergeometric(w=15, l=10, N_w=B, N_l=B)

pvalues2 = np.zeros(47)
for B in range(1,47):
    pvalues2[B] = diluted_margin_hypergeometric(w=15, l=10, N_w=B-1, N_l=B)

pvalues3 = np.zeros(47)
for B in range(1,47):
    pvalues3[B] = diluted_margin_hypergeometric(w=15, l=10, N_w=B-2, N_l=B)

plt.plot(range(47), pvalues, 'blue',
         range(1, 47+1), pvalues2, 'green',
         range(2, 47+2), pvalues3, 'red')
plt.xlabel("N_w")
plt.ylabel("p-value")
```

# Experiment 3 - varying the margin in the sample

Testing the hypothesis $N_w - N_\ell \leq 2$.

Let $N = 100$, $n=30$, and vary $n_w - n_\ell$ with $n_u = 15$.

Below, we plot the $p$-values for varying values of $N_w$. We see that the $p$-value is *not* always monotonically increasing in $N_w$.

```python
pvalues = np.zeros(47)
for B in range(1,47):
    pvalues[B] = diluted_margin_hypergeometric(w=8, l=7, N_w=B+2, N_l=B)

pvalues2 = np.zeros(47)
for B in range(1,47):
    pvalues2[B] = diluted_margin_hypergeometric(w=9, l=6, N_w=B+2, N_l=B)

pvalues3 = np.zeros(47)
for B in range(1,47):
    pvalues3[B] = diluted_margin_hypergeometric(w=10, l=5, N_w=B+2, N_l=B)

plt.plot(range(47), pvalues, 'blue',
         range(47), pvalues2, 'green',
         range(47), pvalues3, 'red')
plt.xlabel("N_w")
plt.ylabel("p-value")
```

# Comparing tri-hypergeometric and plain hypergeometric

## Experiment 1

Testing the hypothesis $\pi_w - \pi_\ell \leq 2/N$.
The tri-hypergeometric $p$-values are in blue and the hypergeometric $p$-values are in green, for increasing values of the parameter $N_w$.

Let $N = 100$, $n=30$, $n_w = 15$, $n_\ell=10$, and $n_u = 15$.

The $p$-value is maximized for each contour at the largest possible value of $N_w$, and the contours have increasing $p$-values as the constant $N_w - N_\ell$ decreases.

```python
pvalues = np.zeros(42)
for B in range(42):
    pvalues[B] = diluted_margin_trihypergeometric(w=15, l=10, n=30, N_w=B+2, N_l=B, N=100)

pvalues2 = np.zeros(42)
for B in range(42):
    pvalues2[B] = diluted_margin_hypergeometric(w=15, l=10, N_w=B+2, N_l=B)

plt.plot(range(2, 42+2), pvalues, 'blue',
         range(2, 42+2), pvalues2, 'green')
plt.xlabel("N_w")
plt.ylabel("p-value")
```

# Comparing tri-hypergeometric and plain hypergeometric

## Experiment 2

Testing the hypothesis $N_w - N_\ell \leq 2$.
The tri-hypergeometric $p$-values are in blue and the hypergeometric $p$-values are in green, for increasing values of the parameter $N_w$.

Let $N = 100$, $n=30$, $n_w = 8$, $n_\ell=7$, and $n_u = 15$.

```python
pvalues = np.zeros(42)
for B in range(42):
    pvalues[B] = diluted_margin_trihypergeometric(w=8, l=7, n=30, N_w=B+2, N_l=B, N=100)

pvalues2 = np.zeros(42)
for B in range(42):
    pvalues2[B] = diluted_margin_hypergeometric(w=8, l=7, N_w=B+2, N_l=B)

plt.plot(range(2, 42+2), pvalues, 'blue',
         range(2, 42+2), pvalues2, 'green')
plt.xlabel("N_w")
plt.ylabel("p-value")
```

## Takeaways

* Not consistent whether the hypergeometric or tri-hypergeometric has more power. Perhaps hypergeometric actually is better!
* The $p$-value of each test is maximized at the smallest or largest value of $N_w$. Thus, to construct a conservative test, take $N_w + N_\ell = N - n_u$ or $N_w+N_\ell = w+\ell$, where $n_u$ is the number of invalid ballots or ballots for other candidates *in the sample*. These are crude upper and lower bounds on how many valid ballots for the winner or loser exist in the population.
* The p-values seem to be **increasing in c, the null hypothesis margin**. Thus, when testing a null hypothesis, the most conservative test lies on the boundary $N_w - N_\ell = c$, with $N_w$ as large as possible.