# Sparse-Group Lasso Inductive Matrix Completion via ADMM

```python
import numpy as np

%matplotlib inline
import matplotlib.pyplot as plt
```

Fix the random state

```python
random_state = np.random.RandomState(0x0BADCAFE)
```

## Problem?

```python
PROBLEM = "classification" if True else "regression"
```

### Synthetic data

```python
assert PROBLEM in ("classification", "regression")
```

Produce a low rank matrix

```python
n_samples, n_objects = 19990, 201
n_rank, n_features = 5, 20
```

```python
n_samples, n_objects = 1990, 2010
n_rank, n_features = 5, 20
```

```python
n_samples, n_objects = 199, 2010
n_rank, n_features = 5, 20
```

```python
n_samples, n_objects = 199, 201
n_rank, n_features = 5, 200
```

```python
n_samples, n_objects = 75550, 40
n_rank, n_features = 5, 20
```

```python
n_samples, n_objects = 199, 201
n_rank, n_features = 5, 20
```

```python
n_samples, n_objects = 1990, 2010
n_rank, n_features = 5, 100
```

```python
n_samples, n_objects = 550, 550
n_rank, n_features = 5, 25
```

Transform the problem

```python
from sgimc.utils import make_imc_data, sparsify

X, W_ideal, Y, H_ideal, R_full = make_imc_data(
    n_samples, n_features, n_objects, n_features,
    n_rank, scale=(0.05, 0.05), noise=0,
    binarize=(PROBLEM == "classification"),
    random_state=random_state)
```

Drop the bulk of the values from $R$

```python
R, mask = sparsify(R_full, 0.10, random_state=random_state)
```

Plot the matrix

```python
fig = plt.figure(figsize=(12, 12))
ax = fig.add_subplot(111, title="The synthetic matrix")
ax.imshow(R.todense(), cmap=plt.cm.RdBu, origin="upper")

print("Observed entries: %d / %d" % (R.nnz, np.prod(R.shape)))
plt.show()
```

# The IMC problem

```python
from sgimc import IMCProblem
```

The IMC problem is:
$$\begin{aligned}
    & \underset{W, H}{\text{miminize}}
      & & \sum_{(i,j)\in \Omega} l(p_{ij}, R_{ij})
          + \nu_W \sum_{m=1}^{d_1} \bigl\| W' e_m \bigr\|_2
          + \nu_H \sum_{m=1}^{d_2} \bigl\| H' e_m \bigr\|_2
          + \mu_W \bigl\| W \bigr\|_1
          + \mu_H \bigl\| H \bigr\|_1
            \,, \\
    & \text{with}
      & & p_{ij} = e_i'\, X W \, H' Y'\, e_j
      \,,
\end{aligned}$$
where $X \in \mathbb{R}^{n_1 \times d_1}$, $Y \in \mathbb{R}^{n_2 \times d_2}$,
$W \in \mathbb{R}^{d_1\times k}$ and $H \in \mathbb{R}^{d_2\times k}$.

### Quadratic Approximation
The target objective without regularization (holding $H$ fixed) is
$$ F(W; H)
    = \sum_{(i,j)\in \Omega}
        l(p_{ij}, R_{ij})
    \,, $$
in which $p = p(W) = (e_i' X W H' Y' e_j)_{(i,j)\in \Omega}$ are the current
predictions.

The Quadratic Approximation to $F$ around $W_0$ is
$$ Q(W; W_0)
    = F(W_0)
        + \nabla F(W_0)' \delta
        + \frac12 \delta' \nabla^2 F(W_0) \delta
    \,, $$
for $\delta = \mathtt{vec}(W - W_0)$. Now the gradient of $F$ w.r.t. vec-form of $W$ is
$$ \nabla F(W_0)
    = \mathtt{vec}\bigl(
        X' g Y H
    \bigr)
    \,, $$
with $g = g(W_0) = (l{'}_{p}(p(W_0)_{ij}, R_{ij}))_{(i,j)\in \Omega}$ is $\Omega$-sparse
matrix of first-order (gradient) data. For a matrix $D \in \mathbb{R}^{d_1 \times k}$
$$ \nabla^2F(W_0)\, \mathtt{vec}(D)
    = \mathtt{vec}\Bigl(
        X' \underbrace{\bigl\{h \odot (X D H'Y')\bigr\}}_{\Omega-\text{sparse}} YH
    \Bigr)
    \,, $$
where $h = h(W_0) = (l{''}_{pp}(p(W_0)_{ij}, R_{ij}))_{(i,j)\in \Omega}$ is
the $\Omega$-sparse matrix of the second order (hessian) values and $\odot$
is the element-wise matrix product.

The quadratic approximation with respect to $H$ around $H_0$ holding $W$ and $\Sigma$ fixed
is similar up to transposing $R$ and swapping $X \leftrightarrow Y$ and $W \leftrightarrow H$
in the above formulae.

**Note** that although the expressions for the gradient and the hessian-vector product presented above are identical to the fast operations in *section 3.1* of [H. Yu et al. (2014)](http://bigdata.ices.utexas.edu/publication/993/), the fomulae here have been derived independently. In fact, they are obvious products of simple block-matrix and **vech** algebra.

The implementation below is, however, completely original (although, nothing special).

#### Implementation details

To compute the gradient and the hessian-vector product we need the following
"elementary" operations:
* $\mathtt{Op}_d: D \mapsto (e_i' X D H'Y' e_j)_{(i, j)\in \Omega}$ -- a map
of some $\mathbb{R}^{d_1\times k}$ dense $D$ to a $\mathbb{R}^{n_1\times n_2}$
$\Omega$-sparse matrix $S$;
* $\mathtt{Op}_s: S \mapsto X'S YH$ mapping an $\mathbb{R}^{n_1\times n_2}$
$\Omega$-sparse $S$ to a $\mathbb{R}^{d_1\times k}$ dense matrix $D$.

The gradient becomes
$$ \nabla F(W_0)
    = \mathtt{vec}(\mathtt{Op}_s(g)) \,, $$
and the hessian-vector product transforms into
$$ \nabla^2 F(W_0)\,\mathtt{vec}(D)
    = \mathtt{vec}\bigl(\mathtt{Op}_s\bigl(h\odot \mathtt{Op}_d(D)\bigr)\bigr) \,. $$

In fact the predictions $p = p(W_0)$ also form an $\mathbb{R}^{n_1\times n_2}$
$\Omega$-matrix, that cam be computed by $p(W_0) = \mathtt{Op}_d(W_0)$. The
gradient $g(W_0)$ and hessian $h(W_0)$ statistics are also $\Omega$-sparse,
and can be computed by element-wise application of $l'_p$ and $l''_{pp}$
to $p$.

Similar formulae hold for $H$ with appropriate re-labellings and transpositions.

**Note** that the sparsity structure remains unchanged and the thin matrix $YH$
can be cached, since both $H$ and $Y$ fit in memory and $k < d_2 \ll n_2$.

```python
# from sgimc import op_s, op_d
```

Define the objectives

The $l_2$ loss $l(p, t) = \frac12 (p-t)^2$.

```python
from sgimc.qa_objective import QAObjectiveL2Loss
```

The log-loss $l(p, t) = \log \bigl(1 + e^{-t p}\bigr)$ for $t\in \{-1, +1\}$
and $p\in \mathbb{R}$.

\begin{align}
    \sigma(x)
        &= \frac1{1+e^{-x}}
        \,, \\
    \sigma'(x)
        &= -\frac{- e^{-x}}{(1+e^{-x})^2}
        = \frac{e^{-x}}{1+e^{-x}} \frac1{1+e^{-x}}
        = (1-\sigma(x))\,\sigma(x)
        \,, \\
    l(p, t)
        &= \log \bigl(1 + e^{-t p}\bigr)
         = \log \bigl(1 + e^{- \lvert p \rvert}\bigr)
         - \min\bigl\{t p, 0\bigr\}
        \,, \\
    l_p'(p, t)
        &= \frac{-t e^{-t p}}{1 + e^{-t p}}
         = -t (1 - \sigma(t p))
        \,, \\
      l_p''(p, t)
        &= (1 - \sigma(t p))\sigma(t p)
        = (1 - \sigma(p))\sigma(p)
        \,.
\end{align}

```python
from sgimc.qa_objective import QAObjectiveLogLoss
```

Huber loss:
$$ l(x; \epsilon)
    = \begin{cases}
        \frac12 x^2
            & \text{if } \lvert x \rvert \leq \epsilon\,, \\
        \epsilon \bigl(\lvert x \rvert - \frac\epsilon2\bigr)
            & \text{otherwise}
    \end{cases}
    \,. $$
Therefore
$$ l_p'
    = \begin{cases}
        x   & \text{if } \lvert x \rvert \leq \epsilon\,, \\
        \epsilon \frac{x}{\lvert x \rvert}
            & \text{otherwise}
    \end{cases}
    \,, $$
and
$$ l_{pp}''
    = \begin{cases}
        1   & \text{if } \lvert x \rvert \leq \epsilon\,, \\
        0   & \text{otherwise}
    \end{cases}
    \,. $$

```python
from sgimc.qa_objective import QAObjectiveHuberLoss
```

Choose the objective

```python
if PROBLEM == "classification":
    QAObjectiveLoss = QAObjectiveLogLoss
else:
    QAObjectiveLoss = QAObjectiveL2Loss  # QAObjectiveHuberLoss

problem = IMCProblem(QAObjectiveLoss, X, Y, R, n_threads=4)
```

### Optimisation

Fix $H$ and consider the problem with respect to $W$:
$$\begin{aligned}
    & \underset{W \in \mathbb{R}^{d_1\times k}}{\text{miminize}}
      & & Q(W; W_0)
          + \sum_{m=1}^{d_1}
              \nu_m \bigl\| W' e_m \bigr\|_2
              + \mu_m \bigl\| W' e_m \bigr\|_1
              + \frac{\kappa_m}2 \bigl\| W' e_m \bigr\|_2^2
            \,.
\end{aligned}$$
Let's move to an equivalent problem by splitting the variables in
the objective, introducing linear consensus constraints and adding
$d_1$ ridge-like regularizers (augmenation)
$$\begin{aligned}
    & \underset{Z_m, \delta_m \in \mathbb{R}^{k\times 1}}{\text{miminize}}
      & & Q(\delta; W_0)
          + \sum_{m=1}^{d_1}
              \nu_m \bigl\| Z_m \bigr\|_2
              + \mu_m \bigl\| Z_m \bigr\|_1
              + \frac{\kappa_m}2 \bigl\| Z_m \bigr\|_2^2
          + \frac1{2\eta}
              \sum_{m=1}^{d_1} \bigl\| \delta_m - (Z_m - W_0'e_m) \bigr\|_2^2
          \,, \\
    & \text{subject to}
      & & Z_m - \delta_m = W_0' e_m\,, m=1 \ldots d_1
          \,,
\end{aligned}$$
with $\sum_{m=1}^{d_1} e_m \delta_m' = \delta$.

The objective is convex and the constraints are linear, which means that
Strong Duality holds for this problem. The lagrangian is
\begin{align}
    \mathcal{L}(Z_m, \delta_m; \lambda_m)
        &= F(W_0)
            + \nabla F(W_0)' \mathtt{vec}(\delta)
            + \frac12 \mathtt{vec}(\delta)' \nabla^2 F(W_0) \mathtt{vec}(\delta)
        \\
        & + \sum_{m=1}^{d_1}
              \nu_m \bigl\| Z_m \bigr\|_2
               + \mu_m \bigl\| Z_m \bigr\|_1
               + \frac{\kappa_m}2 \bigl\| Z_m \bigr\|_2^2
        \\
        & + \frac1\eta
              \sum_{m=1}^{d_1} \lambda_m'\bigl(\delta_m - (Z_m - W_0'e_m)\bigr)
          + \frac1{2\eta}
              \sum_{m=1}^{d_1} \bigl\| \delta_m - (Z_m - W_0'e_m) \bigr\|_2^2
        \,.
\end{align}
Note the following expressions
\begin{align}
    \sum_{m=1}^{d_1} \lambda_m'\bigl(\delta_m - (Z_m - W_0'e_m)\bigr)
        &= \mathtt{tr}\bigl((\delta - (Z - W_0))\Lambda'\bigr) \,,
        \\
    \sum_{m=1}^{d_1} \bigl\| \delta_m - (Z_m - W_0'e_m) \bigr\|_2^2
        &= \Bigl\| \delta - (Z - W_0) \Bigr\|_\text{F}^2 \,,
        \\
\end{align}
where $\Lambda = \sum_{m=1}^{d_1}e_m \lambda_m'$ and $Z = \sum_{m=1}^{d_1}e_m Z_m'$.

#### Sub-0

Consider the following subproblem ($\mathtt{Sub}_0^\text{QA}$):
$$\begin{aligned}
    & \underset{\delta \in \mathbb{R}^{d_1\times k}}{\text{miminize}}
      & & \nabla F(W_0)' \mathtt{vec}(\delta)
            + \frac12 \mathtt{vec}(\delta)' \nabla^2 F(W_0) \mathtt{vec}(\delta)
          \\
%     & & & + \frac1\eta
%               \mathtt{tr}\bigl((\delta - (Z - W_0))\Lambda'\bigr)
%           + \frac1{2\eta}
%               \Bigl\| \delta - (Z - W_0) \Bigr\|_\text{F}^2
    & & & + \frac1{2\eta}
              \Bigl\| \delta + W_0 - Z + \Lambda \Bigr\|_\text{F}^2
          - \frac1{2\eta} \| \Lambda \|_\text{F}^2
            \,.
\end{aligned}$$
The first-order-conditions for this convex problem w.r.t. $\mathtt{vec}(\delta)$
are
$$ \nabla F(W_0) + \nabla^2 F(W_0) \mathtt{vec}(\delta)
    + \frac1\eta
        \mathtt{vec}\bigl( \delta - (\underbrace{Z - W_0 - \Lambda}_{D}) \bigr)
    = 0 \,. $$
Since computing the inverse of the hessian is out of the quiestion, we use Conjugate
Gradient method to solve for $\delta$, because it queries the hessian
only through matrix-vector priducts, which are efficicnetly computable.

The map $\mathtt{Sub}_0^\text{QA}(D; \eta)$ returns the $\delta$ which satisfies
$$\Bigl( \nabla^2 F(W_0) + \frac1\eta I\Bigr)\mathtt{vec}(\delta)
    = \frac1\eta \mathtt{vec}\bigl(D \bigr) - \nabla F(W_0) \,. $$

```python
from sgimc.algorithm.admm import sub_0_cg
```

Using a more comprehensive solver, like `L-BFGS` we can tackle the original
objective, instead of its Quadratic approximation.

Consider the subproblem ($\mathtt{Sub}_0^\text{Orig}$):
$$\begin{aligned}
    & \underset{W \in \mathbb{R}^{d_1\times k}}{\text{miminize}}
      & & F(W; H) + \frac1{2\eta}
              \Bigl\|W - Z + \Lambda \Bigr\|_\text{F}^2
          - \frac1{2\eta} \| \Lambda \|_\text{F}^2
            \,.
\end{aligned}$$
The L-BFGS requires the gradient of the final objective:
$$ \nabla F(W)
    + \frac1\eta
        \mathtt{vec}\bigl( W - (Z - \Lambda) \bigr) \,. $$

```python
from sgimc.algorithm.admm import sub_0_lbfgs
```

#### Sub-m

The next set of subproblems is represented by the following problem ($\mathtt{Sub}_m$):
$$\begin{aligned}
    & \underset{Z_m \in \mathbb{R}^{k\times 1}}{\text{miminize}}
      & & \mu_m \bigl\| Z_m \bigr\|_1 + \nu_m \bigl\| Z_m \bigr\|_2
          + \frac{\kappa_m}2 \bigl\| Z_m \bigr\|_2^2
          \\
%     & & & + \frac1\eta \lambda_m'\bigl(\delta_m - (Z_m - W_0'e_m)\bigr)
%           + \frac1{2\eta} \bigl\| \delta_m - (Z_m - W_0'e_m) \bigr\|_2^2
    & & & + \frac1{2\eta} \bigl\| (\delta_m + W_0'e_m + \lambda_m) - Z_m\bigr\|_2^2
          - \frac1{2\eta} \| \lambda_m \|_2^2
            \,.
\end{aligned}$$
After a **lot of math** this problem admits a closed form solution:
$$ Z_m
    = \frac1{1 + \kappa_m \eta}
      \biggl(1 - \frac{\nu_m \eta}{\|S(V_m; \mu_m \eta)\|_2}\biggr)_+
        S(V_m; \mu_m \eta)
    \,, $$
where $V_m = \delta_m + W_0'e_m + \lambda_m$ and
$$ S(u; \mu_m \eta)
    = \Bigl(\Bigl(1 - \frac{\mu_m \eta}{\lvert u_i \rvert}\Bigr)_+ u_i\Bigr)_{i=1}^k\,, $$
is the **soft_thresholding** operator.

The map $\mathtt{Sub}_m(D; \eta)$ returns $Z_m$ defined above.

```python
from sgimc.algorithm.admm import sub_m
```

#### ADMM

Thus the QA-ADMM for $W$ around $W_0$ with $H$ fixed is the follwing
iterative procedure:
\begin{align}
    Z^{t+1}_m &= \mathtt{Sub}_m(W^t_m + \lambda^t_m) \,,\, m = 1\ldots d_1 \,,\\
    W^{t+1} &= \mathtt{Sub}_0(Z^{t+1} - W_0 - \Lambda^t) + W_0  \,,\\
    % W^{t+1} &= \mathtt{Sub}_0(Z^t - W_0 - \Lambda^t) + W_0  \,,\\
    % Z^{t+1}_m &= \mathtt{Sub}_m(W^{t+1}_m + \lambda^t_m) \,,\, m = 1\ldots d_1 \,,\\
    \Lambda^{t+1} &= \Lambda^t + (W^{t+1} - Z^{t+1})\,,\\
\end{align}
where $W^{t+1}_m$ is the $m$-th row of $W^{t+1}$, $Z^{t+1}_m$ is the
$m$-th row of $Z^{t+1}$ and $\lambda_m$ is the $m$-th row of $\Lambda$.
These iterations necessarily converge to a fixed point, which is the
solution of the original optimisation problem. If stopped early, the
current values of $W^t$ and $Z^t$ would be close to each other, however
$Z^t$ would be sparse and $W^t$ -- dense.

Note that we can also consider ADMM with a linear approximation of $F$
w.r.t. $W$ at $W_0$, instead of the quadratic (LA-ADMM). This way the algorithm
reduces to prox-gradient descent with step $\eta$. Although it does not utilize
the second order infromation, it can be fused with Nesterov's Accelerated
gradient.

```python
from sgimc.algorithm import admm_step

def step_qaadmm(problem, W, H, C, eta, method="l-bfgs", sparse=True,
                n_iterations=50, rtol=1e-5, atol=1e-8):

    approx_type = "quadratic" if method in ("cg",) else "linear"
    Obj = problem.objective(W, H, approx_type=approx_type)

    return admm_step(Obj, W, C, eta, sparse=sparse, method=method,
                     n_iterations=n_iterations, rtol=rtol, atol=atol)
```

```python
from sgimc.algorithm.decoupled import step as decoupled_step

def step_decoupled(problem, W, H, C, eta, rtol=1e-5, atol=1e-8):

    Obj = problem.objective(W, H, approx_type="linear")

    return decoupled_step(Obj, W, C, eta, rtol=rtol, atol=atol)
```

Ad-hoc procedure. No guarantees for convergence.

```python
# def step_adhoc(problem, W, H, C, eta, rtol=1e-5, atol=1e-8):
#     Obj = problem.objective(W, H, approx_type="quadratic")
#     delta = sub_0_cg(np.zeros_like(W), Obj, eta=eta, tol=1e-8)
#     return sub_m(delta + W, *C, eta=eta)
```

```python
# def QA_argmin(D, Obj, tol=1e-8):

#     # set up the CG arguments
#     x = D.reshape(-1).copy()
#     b = - Obj.grad().reshape(-1)
#     Ax = lambda x: Obj.hess_v(x.reshape(D.shape)).reshape(-1)

#     n_iter = simple_cg(Ax, b, x, tol=tol)
#     return x.reshape(D.shape)
```

Thus Sparse Group IMC via QA-ADMM is the follwing iterative procedure:
* $W^{t+1} = \mathtt{ADMM}\bigl(W^t; H^t\bigr)$,
* $H^{t+1} = \mathtt{ADMM}\bigl(H^t; W^{t+1}\bigr)$,

until convergence.

```python
from sgimc import imc_descent
```

The loss information: value and regularization on the train data and value of the full matrix.

```python
from sgimc.utils import performance
```

### Illustration

```python
step_fn = step_qaadmm
# step_fn = step_decoupled
```

$$\bigl(C_\mathtt{lasso}, C_\mathtt{group}, C_\mathtt{ridge}\bigr) = C \,.$$
It seems that it must hold $C_\mathtt{lasso} > C_\mathtt{group}$ so that
individual sparsity preceeds group sparsity.

```python
if PROBLEM == "classification":
    C = 1e0, 1e-1, 1e-3
    eta = 1e0
else:
    # C = 2e-5, 2e-3, 0
    C = 2e-3, 2e-4, 1e-4  # 1e-2
    eta = 1e1
```

```python
if step_fn == step_decoupled:
    eta = 1e-3
```

Let's see how the feature coefficients look like.

```python
from sgimc.utils import plot_WH, plot_loss
```

Initialization

```python
K = 10  # n_rank
# K = n_rank

W_0 = random_state.normal(size=(X.shape[1], K))
H_0 = random_state.normal(size=(Y.shape[1], K))

# W_0 = W_ideal.copy() # + random_state.normal(scale=0.1, size=(X.shape[1], K))
# H_0 = H_ideal.copy() # + random_state.normal(scale=0.1, size=(Y.shape[1], K))
```

Now in this experiment the ideal solution is a unit matrix stacked atop a zero martix.

```python
plot_WH(W_ideal, H_ideal)

loss_arr, exp_type, norm_type = performance(
    problem, W_ideal, H_ideal, C, R_full)

print("The loss on the initial guess is:")
print("%.3e + %.3e -- partial matrix" % (loss_arr[0, -1], loss_arr[1, -1]))
print("%.3e -- full matrix" % loss_arr[3, -1])
print("score %.4f" % loss_arr[2, -1])
```

```python
plt.imshow(np.dot(W_ideal, H_ideal.T))
```

The initial guess is:

```python
plot_WH(W_0, H_0)

loss_arr, exp_type, norm_type = performance(
    problem, W_0, H_0, C, R_full)

print("The loss on the initial guess is:")
print("%.3e + %.3e -- partial matrix" % (loss_arr[0, -1], loss_arr[1, -1]))
print("%.3e -- full matrix" % loss_arr[3, -1])
print("score %.4f" % loss_arr[2, -1])
```

```python
plt.imshow(np.dot(W_0, H_0.T))
```

Run!

```python
W, H = W_0.copy(), H_0.copy()
```

Setup the parameters of the subalgorithm.

```python
step_kwargs = {
    "C": C,                 # the regularizr constants (C_lasso, C_group, C_ridge)
    "eta": eta,             # the eta of the ADMM (larger - faster but more unstable)
    "rtol": 1e-5,           # the relative tolerance for stopping the ADMM
    "atol": 1e-8,           # the absolute tolerance
    "method": "l-bfgs",     # the method to use in Sub_0
    "n_iterations": 2,      # the number of iterations of the inner ADMM
}

# n_iterations = 2, 10, 25, 50
```

Run the alternating minimization with ADMM as the sub-algorithm.

```python
W, H = imc_descent(problem, W, H,
                   step_fn,                  # the inner optimization
                   step_kwargs=step_kwargs,  # asrtguments for the inner optimizer
                   n_iterations=1000,        # the number of outer iterations (Gauss-Siedel)
                   return_history=True,      # Record the evolution of the matrices (W, H)
                   rtol=1e-5,                # relative stopping tolerance for the outer iterations
                   atol=1e-7,                # absolute tolerance
                   verbose=True,             # show the progress bar
                   check_product=True,       # use the product W H' for stopping
                  )
```

     17%|█▋        | 166/1000 [01:03<04:34,  3.04it/s]

Inspect

```python
plot_WH(abs(W[..., -1]), abs(H[..., -1]))

loss_arr, exp_type, norm_type = performance(problem, W, H, C, R_full)

print("The loss on the final estimates is:")
print("%.3e + %.3e -- partial matrix" % (loss_arr[0, -1], loss_arr[1, -1]))
print("%.3e -- full matrix" % loss_arr[3, -1])
print("score %.4f" % loss_arr[2, -1])
```

```python
plt.imshow(np.dot(W[..., -1], H[..., -1].T))
```

```python
fig = plt.figure(figsize=(12, 12))
ax = fig.add_subplot(111, title="Elementwise loss value")

R_hat = problem.prediction(W[..., -1], H[..., -1])
ax.imshow(problem.loss(R_hat, R_full), cmap=plt.cm.hot, origin="upper")

plt.show()
```

```python
print(str(np.array(["#", "."])[np.isclose(W[..., -1], 0)*1]).replace("' '", ""))
```

    [['#.........']
     ['....#.....']
     ['..#.......']
     ['......#...']
     ['.........#']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']]

```python
print(str(np.array(["#", "."])[np.isclose(H[..., -1], 0)*1]).replace("' '", ""))
```

    [['#.........']
     ['....#.....']
     ['..#.......']
     ['......#...']
     ['.........#']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..#.......']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']
     ['..........']]

```python
np.linalg.norm(W[..., -1], 2, axis=-1)
```

    array([  84.33452596,  116.27803303,   68.39865667,   94.56376613,
             84.2719859 ,    0.        ,    0.        ,    0.        ,
              0.        ,    0.        ,    0.        ,    0.        ,
              0.        ,    0.        ,    0.        ,    0.        ,
              0.        ,    0.        ,    0.        ,    0.        ,
              0.        ,    0.        ,    0.        ,    0.        ,    0.        ])

```python
np.linalg.norm(H[..., -1], 2, axis=-1)
```

    array([  92.46341611,   66.942329  ,  114.10945471,   82.50728563,
             92.46151265,    0.        ,    0.        ,    0.        ,
              0.        ,    0.        ,    0.19861052,    0.        ,
              0.        ,    0.        ,    0.        ,    0.        ,
              0.        ,    0.        ,    0.        ,    0.        ,
              0.        ,    0.        ,    0.        ,    0.        ,    0.        ])

```python
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111)
ax.imshow(~ np.isclose(np.dot(W[..., -1], H[..., -1].T),
                       np.dot(W_ideal, H_ideal.T)),
          cmap=plt.cm.binary_r)
plt.show()
```

```python
plt.hist(abs(np.dot(W[..., -1], H[..., -1].T)).reshape(-1), bins=20) ;
```

```python
plt.hist(abs(R_hat).reshape(-1), bins=200) ;
```

```python
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, projection='3d', xlabel="col", ylabel="row")

ZZ = problem.loss(R_hat, R_full)
mesh_ = np.meshgrid(*[np.linspace(0, 1, num=n) for n in ZZ.shape[::-1]])

surf = ax.plot_surface(*mesh_, ZZ, alpha=0.5, lw=0, antialiased=True,
                       cmap=plt.cm.coolwarm)

fig.colorbar(surf, shrink=0.5, aspect=10)

ax.view_init(37, 15)

```

```python
ZZ[~mask].std()
```

    0.079106809138822079

```python
ZZ[mask].std()
```

    0.080047102742687851

```python
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, projection='3d', xlabel="col", ylabel="row")

ZZ = np.dot(W[..., -1], H[..., -1].T)
mesh_ = np.meshgrid(*[np.linspace(0, 1, num=n) for n in ZZ.shape[::-1]])

surf = ax.plot_surface(*mesh_, ZZ, alpha=0.5, lw=0, antialiased=True,
                       cmap=plt.cm.coolwarm)

fig.colorbar(surf, shrink=0.5, aspect=10)

ax.view_init(37, 15)

```

```python
plot_loss(loss_arr, exp_type, norm_type,
          fig_size=4, max_cols=4, yscale="log")
```

<hr/>


# 1D Plug Flow Reactor Model with Surface Chemistry

In this model, we will illustrate the derivation of the governing differential equations and algebraic constraints, calculation of the initial conditions of the variables and their spatial derivatives and use the [scikits.odes.dae](http://scikits-odes.readthedocs.io/en/latest/guide.html#object-oriented-interface-ode-and-dae) IDA solver to solve this system of differential algebraic equations (DAE).

```python
from __future__ import print_function, division
import numpy as np
from scikits.odes import dae
import cantera as ct
import matplotlib.pyplot as plt
%matplotlib inline
print('Runnning Cantera version: ' + ct.__version__)
```

    Runnning Cantera version: 2.4.0

## Define gas species, bulk species, surface species and the interface

Here, we use a kinetic mechanism involving the chemical vapor deposition of silicon nitride (Si3N4) from SiF4 and NH3. 25 gas species, 6 surface species and 2 bulk species mechanism is applied by [Richard S. Larson et al. 1996, SAND96-8211](https://github.com/yuj056/yuj056.github.io/blob/master/_posts/Sandia.pdf).

```python
#import the SiF4 + NH3 reaction mechanism
mech = 'data/SiF4_NH3_mec.cti'
#import the models for gas and bulk
gas, bulk_Si, bulk_N = ct.import_phases(mech, ['gas', 'SiBulk', 'NBulk'])
#import the model for gas-Si-N interface
gas_Si_N_interface = ct.Interface(mech, 'SI3N4', [gas,bulk_Si,bulk_N])
```

# Case 1: isothermal reactor
## Define reactor conditions : temperature, pressure, fuel, and some important parameters

```python
T0 = 1713  # Kelvin
p0 = 2 * ct.one_atm / 760.0  # Pa ~2 Torr
gas.TPX = T0, p0, "NH3:6, SiF4:1"
bulk_Si.TP = T0, p0
bulk_N.TP = T0, p0
gas_Si_N_interface.TP = T0, p0
D = 5.08e-2  # diameter of the tube [m]
Ac = np.pi * D**2/4  # cross section of the tube [m^2]
mu = 5.7e-5  # kg/(m-s) dynamic viscosity
perim = np.pi * D  # perimeter of the tube
# calculate the site fractions of surface species at the entrance of the tube at steady state
gas_Si_N_interface.advance_coverages(100.0)  # Here we assume after 100s, the system reaches the steady state
Zk_0 = gas_Si_N_interface.coverages
N = gas.n_species  # number of gas species
M = gas_Si_N_interface.n_species  # number of surface species
```

## Define a residual function for IDA solver

For the isothermal tube with laminar flow, since the temperature of the flow and tube is constant, the energy conservation equation can be ignored. The governing equations include conservation of mass and species, momentum equation, equation of state, and the algebraic constraints that the net production rate of surface species by heterogeneous reactions are zero and that the sum of site fractions equals 1.

Here we define a residual function, an equation which should always evaluate to the zero vector, as the input of IDA solver, which listed as follows:

```latex
%%latex
\begin{align}
    R[0] &= u\frac{d\rho}{dz} + \rho\frac{du}{dz} - \frac{p'}{A_c}\sum^{K_g}\dot{s}_{k,g}W_{k,g} \\
    R[1] &= \rho u A_c\frac{dY_k}{dz} + Y_k p'\sum^{K_g}\dot{s}_{k,g}W_{k,g} - \dot{\omega_k}W_kA_c - \dot{s}_{k,g}W_{k,g} p' \\
    R[2] &= 2\rho u \frac{du}{dz} + u^2\frac{d\rho}{dz} + \frac{dP}{dz} + \frac{32u\mu}{D^2}\\
    R[3] &= P\bar{W} - \rho RT\\
    R[4] &= \dot{s}_{k,s} \\
    R[5] &= \sum_{phase}{Z_{k,s}} - 1
\end{align}
```

\begin{align}
    R[0] &= u\frac{d\rho}{dz} + \rho\frac{du}{dz} - \frac{p'}{A_c}\sum^{K_g}\dot{s}_{k,g}W_{k,g} \\
    R[1] &= \rho u A_c\frac{dY_k}{dz} + Y_k p'\sum^{K_g}\dot{s}_{k,g}W_{k,g} - \dot{\omega_k}W_kA_c - \dot{s}_{k,g}W_{k,g} p' \\
    R[2] &= 2\rho u \frac{du}{dz} + u^2\frac{d\rho}{dz} + \frac{dP}{dz} + \frac{32u\mu}{D^2}\\
    R[3] &= P\bar{W} - \rho RT\\
    R[4] &= \dot{s}_{k,s} \\
    R[5] &= \sum_{phase}{Z_{k,s}} - 1
\end{align}

The detailed derivation of the DAE system can be found in [my report](https://github.com/yuj056/yuj056.github.io/blob/master/Week1/yuj056_github_io.pdf).

```python
def residual(z, vec, vecp, result):
    """ we create the residual equations for the problem
        vec = [u, rho, Yk, p, Zk]
        vecp = [dudz, drhodz, dYkdz, dpdz, dZkdz]
    """
    # temporary variables
    u = vec[0]  # velocity
    rho = vec[1]  # density
    Y = vec[2:2+N]  # vector of mass fractions of all gas species
    p = vec[2+N]  # pressure
    Z = vec[3+N:]  # vector of site fractions of all surface species

    dudz = vecp[0]  # velocity spatial derivative
    drhodz = vecp[1]  # density spatial derivative
    dYdz = vecp[2:2+N]  # mass fraction spatial derivative
    dpdz = vecp[2+N]  # pressure spatial derivative

    # Use unnormalized mass fractions to avoid over-constraining the system
    gas.set_unnormalized_mass_fractions(Y)
    gas.TP = T0,p

    bulk_Si.TP = T0,p
    bulk_N.TP = T0,p

    # Use unnormalized site fractions (coverages) to avoid over-constraining the system
    gas_Si_N_interface.set_unnormalized_coverages(Z)
    gas_Si_N_interface.TP = T0,p

    # temporary variables (based on the given state)
    coverages = gas_Si_N_interface.coverages  # site fraction vector
    sdot_g = gas_Si_N_interface.net_production_rates[:N]  # heterogeneous production rate of gas species
    sdot_s = gas_Si_N_interface.net_production_rates[-M:]
    wdot_g = gas.net_production_rates  # homogeneous production rate of gas species
    W_g = gas.molecular_weights  # vector of molecular weight of gas species

    # mass continuity equation
    result[0] = u*drhodz + rho*dudz - perim*np.sum(sdot_g*W_g)/Ac

    # conservation of species
    for k in range(N):
        result[1+k] = (rho*u*Ac*dYdz[k] + Y[k]*perim*np.sum(sdot_g*W_g)
                       - wdot_g[k]*W_g[k]*Ac
                       - sdot_g[k]*W_g[k]*perim)
    # conservation of momentum
    result[1+N] = 2*rho*u*dudz + np.power(u,2)*drhodz + dpdz + 32*u*mu/D**2

    # equation of state
    result[2+N] = gas.density - rho

    # algebraic constraints
    for j in range(M):
        result[3+N+j] = sdot_s[j]

    # replace the constraint with the condition sum(Zk) = 1 for the largest site fraction species
    index = np.argmax(coverages)
    result[3+N+index] = np.sum(coverages) - 1
```

## Determine the initial values of the spatial derivatives of the unknowns which need to be used as the initial conditions for the IDA solver

The following linear equation system has been solved by [np.linalg.solve](https://docs.scipy.org/doc/numpy-1.14.0/reference/generated/numpy.linalg.solve.html), a linear solver, to calculate the initial values of the spatial derivatives of the unknowns.

```latex
%%latex
\begin{align}
    u_0\rho_0' + \rho_0 u_0' - \frac{p'}{A_c}\sum^{K_g}\dot{s}_{k,g}W_{k,g} &= 0\\
    \rho_0 u_0 A_c Y_{k,0}' + Y_{k,0} p'\sum^{K_g}\dot{s}_{k,g}W_{k,g} - \dot{\omega_k}W_kA_c - \dot{s}_{k,g}W_{k,g} p' &=0 \\
    2\rho_0 u_0 u_0' + u_0^2\rho_0' + P_0' + \frac{32u_0 \mu}{D^2} &=0\\
    -RT\rho_0' + \bar{W_0}P_0' - P_0\frac{\sum^{K_g}Y_{k}'/W_{k,g}}{(\sum^{K_g}Y_{k}/W_{k,g})^2} &= 0
\end{align}
```

\begin{align}
    u_0\rho_0' + \rho_0 u_0' - \frac{p'}{A_c}\sum^{K_g}\dot{s}_{k,g}W_{k,g} &= 0\\
    \rho_0 u_0 A_c Y_{k,0}' + Y_{k,0} p'\sum^{K_g}\dot{s}_{k,g}W_{k,g} - \dot{\omega_k}W_kA_c - \dot{s}_{k,g}W_{k,g} p' &=0 \\
    2\rho_0 u_0 u_0' + u_0^2\rho_0' + P_0' + \frac{32u_0 \mu}{D^2} &=0\\
    -RT\rho_0' + \bar{W_0}P_0' - P_0\frac{\sum^{K_g}Y_{k}'/W_{k,g}}{(\sum^{K_g}Y_{k}/W_{k,g})^2} &= 0
\end{align}

We assume the derivatives of the site fractions are equal to zero, although it is trivial for the IDA solver.

```python
######## Solve linear system for the initial vecp ###########
"""
   a = coefficient of [u', rho', Yk', P']
   b = RHS constant of each conservation equations
"""
rho0 = gas.density  # initial density of the flow
u0 = 11.53  # m/s initial velocity of the flow
W = gas.molecular_weights
W_avg = gas.mean_molecular_weight
sdot = gas_Si_N_interface.net_production_rates  # heterogeneous molar production rate
wdot = gas.net_production_rates  # homogeneous molar production rate
################### a #########################
a = np.zeros((3+N,3+N))
a[0,:] = np.hstack((rho0, u0, np.zeros(1+N)))
for i in range(N):
    a[1+i,2+i] = rho0*u0*Ac
a[1+N,:] = np.hstack((2*rho0*u0, u0**2, np.zeros(N), 1))
coef = np.zeros(N)
for j in range(N):
    coef[j] = gas.P/W[j]/np.power(np.sum(gas.Y/W),2)
a[2+N,:] = np.hstack((0, ct.gas_constant*T0, coef, -W_avg))
################### b ###########################
b = np.zeros(3+gas.n_species)
b[0] = perim*np.sum(sdot[:N]*W)/Ac
for i in range(gas.n_species):
    b[1+i] = (wdot[i]*W[i]*Ac
              + sdot[i]*W[i]*perim
              - gas.Y[i]*perim*np.sum(sdot[:N]*W))
b[1+gas.n_species] = -32*u0*mu/D**2
b[2+gas.n_species] = 0
part_vecp0 = np.linalg.solve(a,b)

vecp0 = np.hstack((part_vecp0, np.zeros(M)))
vec0 = np.hstack((11.53, gas.density, gas.Y, gas.P, Zk_0))
```

## Run the IDA solver to calculate the unknowns varying in the flow direction

```python
solver = dae(
    'ida',
    residual,
    first_step_size=1e-16,
    atol=1e-8,  # absolute tolerance for solution
    rtol=1e-8,  # relative tolerance for solution
    # If the given problem is of type DAE, some items of the residual vector
    # returned by the 'resfn' have to be treated as algebraic equations, and
    # algebraic variables must be defined. These algebraic variables are
    # denoted by the position (index) in the state vector y. All these
    # indexes have to be specified in the 'algebraic_vars_idx' array.
    algebraic_vars_idx=[np.arange(3+N,3+N+M,1)],
    max_steps=5000,
    old_api=False  # Forces use of new api (namedtuple)
)

times = np.arange(0,0.7,0.01)
solution = solver.solve(times, vec0, vecp0)
```

## Plot the results

```python
# plot velocity of gas along the flow direction
f, ax = plt.subplots(3,2, figsize=(9,9), dpi=96)
ax[0,0].plot(times, solution.values.y[:,0], color='C0')
ax[0,0].set_xlabel('Distance (m)')
ax[0,0].set_ylabel('Velocity (m/s)')

# plot gas density along the flow direction
ax[0,1].plot(times, solution.values.y[:,1], color='C1')
ax[0,1].set_xlabel('Distance (m)')
ax[0,1].set_ylabel('Density ($\mathregular{kg/m^3}$)')
ax[0,1].ticklabel_format(axis='y', style='sci', scilimits=(-2,2))  # scientific notation

# plot major and minor gas species separately
minor_idx = []
major_idx = []
for i,name in enumerate(gas.species_names):
    mean = np.mean(solution.values.y[:,2+i])
    if mean <= 0.01:
        minor_idx.append(i)
    else:
        major_idx.append(i)

# plot minor species
for i in minor_idx:
    style = '-' if i < 10 else '--'
    ax[1,0].plot(times, solution.values.y[:,2+i], label=gas.species_names[i], linestyle=style)
ax[1,0].legend(fontsize=7, loc='upper right')
ax[1,0].set_xlabel('Distance (m)')
ax[1,0].set_ylabel('Mass Fraction')
ax[1,0].ticklabel_format(axis='y', style='sci', scilimits=(-2,2))  # scientific notation

# plot major species
for j in major_idx:
    ax[1,1].plot(times,solution.values.y[:,2+j], label=gas.species_names[j])
ax[1,1].legend(loc='best')
ax[1,1].set_xlabel('Distance (m)')
ax[1,1].set_ylabel('Mass Fraction')

# plot the pressure of the gas along the flow direction
ax[2,0].plot(times, solution.values.y[:,2+N], color='C2')
ax[2,0].set_xlabel('Distance (m)')
ax[2,0].set_ylabel('Pressure (Pa)')

# plot the site fraction of the surface species along the flow direction
for i,name in enumerate(gas_Si_N_interface.species_names):
    ax[2,1].plot(times, solution.values.y[:,3+N+i], label=name)
ax[2,1].legend()
ax[2,1].set_xlabel('Distance (m)')
ax[2,1].set_ylabel('Site Fraction')
f.tight_layout(pad=0.5)
```

# Case 2: Adiabatic reactor

Since the application of isothermal reactor is not prevalent, to improve the model for real use, the adiabatic reator is considered. Here, the energy balance equation is also considered.

The heat flow rate into the system has two components. One is due to the heat flux $q_e$ from the surroundings to the outer tube wall (whose surface area per unit length is $a_e$) and accumulation of enthalpy in the bulk solid. The other is due to $q_i$, the heat flux to the gas from the inner tube wall, and accumulation of enthalpy in the surface species. The expression of energy balance equation for this problem is as follows:

```latex
%%latex
\begin{align}
    \rho u A_c c_p \frac{dT}{dz} +A_c \sum_{K_g}\dot{\omega}_k W_k h_k + p'\sum_{K_g}h_k\dot{s}_k W_k &= a_eq_e - p'\sum^{K_b}_{bulk}\dot{\omega}_kh_k\\&=p'q_i + p'\sum^{K_g}_{gas}\dot{s_k}W_kh_k
\end{align}
```

\begin{align}
    \rho u A_c c_p \frac{dT}{dz} +A_c \sum_{K_g}\dot{\omega}_k W_k h_k + p'\sum_{K_g}h_k\dot{s}_k W_k &= a_eq_e - p'\sum^{K_b}_{bulk}\dot{\omega}_kh_k\\&=p'q_i + p'\sum^{K_g}_{gas}\dot{s_k}W_kh_k
\end{align}

Since the adiabatic reactor is considered, $q_e = 0$. Similar to the procedure for the isothermal reactor model, add the energy equation into the residual function and calculate the initial value of the spatial derivative of the temperature.

```python
############################### initial conditions ##################################################################
# import the SiF4 + NH3 reaction mechanism
mech = 'data/SiF4_NH3_mec.cti'
# import the models for gas and bulk
gas, bulk_Si, bulk_N = ct.import_phases(mech,['gas','SiBulk','NBulk'])

# import the model for gas-Si-N interface
gas_Si_N_interface = ct.Interface(mech, 'SI3N4',[gas,bulk_Si,bulk_N])
T0 = 1713  # K
p0 = 2 * ct.one_atm / 760.0  # Pa ~2Torr
gas.TPX = T0, p0, "NH3:6, SiF4:1"
bulk_Si.TP = T0, p0
bulk_N.TP = T0, p0
gas_Si_N_interface.TP = T0, p0
D = 5.08e-2  # diameter of the tube [m]
Ac = np.pi * D**2/4  # cross section of the tube [m]
mu = 5.7e-5  # kg/(m-s) dynamic viscosity
perim = np.pi * D  # perimeter of the tube
# calculate the site fractions of surface species at the entrance of the tube at steady state
gas_Si_N_interface.advance_coverages(100.0)
Zk_0 = gas_Si_N_interface.coverages
######################################## IDA solver ###################################################################
def residual(z, vec, vecp, result):
    """ we create the residual equations for the problem
        vec = [u, rho, Yk, p, Zk, T]
        vecp = [dudz, drhodz, dYkdz, dpdz, dZkdz, dTdz]
    """
    # temporary variables
    u = vec[0]  # velocity
    rho = vec[1]  # density
    Y = vec[2:2+N]  # vector of mass fractions of all gas species
    p = vec[2+N]  # pressure
    Z = vec[3+N:-1]  # vector of site fractions of all surface species
    T = vec[-1]  # temperature

    dudz = vecp[0]  # velocity spatial derivative
    drhodz = vecp[1]  # density spatial derivative
    dYdz = vecp[2:2+N]  # mass fraction spatial derivative
    dpdz = vecp[2+N]  # pressure spatial derivative
    dTdz = vecp[-1]  # temperature spatial derivative

    h = gas.enthalpy_mass  # enthalpy of gas species per mass
    h_Si = bulk_Si.enthalpy_mass  # enthalpy of Si per mass
    h_N = bulk_N.enthalpy_mass  # enthalpy of N per mass

    # initial conditions
    gas.set_unnormalized_mass_fractions(Y)
    gas.TP = T, p

    bulk_Si.TP = T, p
    bulk_N.TP = T, p
    gas_Si_N_interface.set_unnormalized_coverages(Z)
    gas_Si_N_interface.TP = T, p

    # temporary variables (based on the current system state)
    coverages = gas_Si_N_interface.coverages  # site fraction vector
    sdot_g = gas_Si_N_interface.net_production_rates[:N]  # heterogeneous production rate of gas species
    sdot_s = gas_Si_N_interface.net_production_rates[-M:]  # molar production rate of surface speceis
    wdot_g = gas.net_production_rates  # homogeneous production rate of gas species
    W_g = gas.molecular_weights  # vector of molecular weight of gas species
    W_Si_b = bulk_Si.molecular_weights
    W_N_b = bulk_N.molecular_weights
    bdot = gas_Si_N_interface.net_production_rates[gas.n_species:gas.n_species+2]  # bulk production rate

    # mass continuity equation
    result[0] = u*drhodz+rho*dudz-perim*np.sum(sdot_g*W_g)/Ac
    # conservation of species
    for k in range(gas.n_species):
        result[1+k] = (rho*u*Ac*dYdz[k] + Y[k]*perim*np.sum(sdot_g*W_g)
                       - wdot_g[k]*W_g[k]*Ac
                       - sdot_g[k]*W_g[k]*perim)
    # conservation of momentum
    result[1+gas.n_species] = 2*rho*u*dudz + np.power(u,2)*drhodz + dpdz + 32*u*mu/D**2

    # equation of state
    result[2+gas.n_species] = gas.density - rho

    # algebraic constraints
    for j in range(M):
        result[3+N+j] = sdot_s[j]

    # replace the constraints with the condition sum(Zk) = 1 for the largest site fraction species
    index = np.argmax(coverages)
    result[3+N+index] = np.sum(coverages) - 1

    # energy equation
    result[3+N+M] = (rho*u*Ac*gas.cp*dTdz
        + Ac*np.sum(wdot_g*W_g*h)
        + perim*np.sum(h*sdot_g*W_g)
        + perim*(bdot[0]*W_Si_b*h_Si + bdot[1]*W_N_b*h_N))
```

```python
######## Solve linear system for the initial values of vecp ###########
"""
   a = coefficient of [u', rho', Yk', P',T]
   b = RHS constant of each conservation equations
"""
rho0 = gas.density  # initial density of the flow
u0 = 11.53  # m/s initial velocity of the flow
W = gas.molecular_weights
W_avg = gas.mean_molecular_weight
sdot = gas_Si_N_interface.net_production_rates  # heterogeneous molar production rate
wdot = gas.net_production_rates  # homogeneours molar production rate
h = gas.enthalpy_mass
h_Si = bulk_Si.enthalpy_mass
h_N = bulk_N.enthalpy_mass
W_Si = bulk_Si.molecular_weights
W_N = bulk_N.molecular_weights
################### a #########################
a = np.zeros((4+N,4+N))
a[0,:] = np.hstack((rho0, u0, np.zeros(2+N)))
for i in range(N):
    a[1+i,2+i] = rho0*u0*Ac
a[1+N,:] = np.hstack((2*rho0*u0, u0**2, np.zeros(N), 1,0))
coef = np.zeros(N)
for j in range(N):
    coef[j] = gas.P/W[j]/np.power(np.sum(gas.Y/W),2)
a[2+N,:] = np.hstack((0, ct.gas_constant*T0, coef, -W_avg, 0))
a[3+N,:] = np.hstack((np.zeros(3+N), rho0*u0*Ac*gas.cp))
################### b ###########################
b = np.zeros(4+gas.n_species)
b[0] = perim*np.sum(sdot[:N]*W)/Ac
for i in range(N):
    b[1+i] = (wdot[i]*W[i]*Ac
              + sdot[i]*W[i]*perim
              - gas.Y[i]*perim*np.sum(sdot[:N]*W))
b[1+gas.n_species] = -32*u0*mu/D**2
b[2+gas.n_species] = 0
b[3+gas.n_species] = (- Ac*np.sum(wdot*W*h)
                      - perim*np.sum(h*sdot[:N]*W)
                      - perim*np.sum(sdot[N]*W_Si*h_Si + sdot[N+1]*W_N*h_N))
part_vecp0 = np.linalg.solve(a,b)

vecp0 = np.hstack((part_vecp0[:-1], np.zeros(M), part_vecp0[-1]))
vec0 = np.hstack((11.53, gas.density, gas.Y, gas.P, Zk_0, T0))
```

```python
solver = dae(
    'ida',
    residual,
    atol=1e-8,  # absolute tolerance for solution
    rtol=1e-8,  # relative tolerance for solution
    algebraic_vars_idx=[np.arange(3+N,3+N+M,1)],
    max_steps=5000,
    one_step_compute=True,
    old_api=False
)

time = []
solution = []
state = solver.init_step(0.0, vec0, vecp0)
while state.values.t < 0.7:
    time.append(state.values.t)
    solution.append(state.values.y)
    state = solver.step(0.7)

time = np.array(time)
solution = np.array(solution)
```

```python
f, ax = plt.subplots(4,2, figsize=(9,12), dpi=96)

# plot gas velocity along the flow direction
ax[0,0].plot(time, solution[:,0], color='C0')
ax[0,0].set_xlabel('Distance (m)')
ax[0,0].set_ylabel('Velocity (m/s)')

# plot gas density along the flow direction
ax[0,1].plot(time, solution[:,1], color='C1')
ax[0,1].set_xlabel('Distance (m)')
ax[0,1].set_ylabel('Density ($\mathregular{kg/m^3}$)')
ax[0,1].ticklabel_format(axis='y', style='sci', scilimits=(-2,2))  # scientific notation

# plot major and minor gas species separately
minor_idx = []
major_idx = []
for i,name in enumerate(gas.species_names):
    mean = np.mean(solution[:,2+i])
    if mean <= 0.01:
        minor_idx.append(i)
    else:
        major_idx.append(i)

# plot minor gas species along the flow direction
for i in minor_idx:
    style = '-' if i < 10 else '--'
    ax[1,0].plot(time, solution[:,2+i], label=gas.species_names[i], linestyle=style)
ax[1,0].legend(fontsize=7.5, loc='best')
ax[1,0].set_xlabel('Distance (m)')
ax[1,0].set_ylabel('Mass Fraction')
ax[1,0].ticklabel_format(axis='y', style='sci', scilimits=(-2,2))  # scientific notation

# plot major gas species along the flow direction
for j in major_idx:
    ax[1,1].plot(time, solution[:,2+j], label=gas.species_names[j])
ax[1,1].legend(fontsize=8, loc='best')
ax[1,1].set_xlabel('Distance (m)')
ax[1,1].set_ylabel('Mass Fraction')

# plot the pressure of the gas along the flow direction
ax[2,0].plot(time, solution[:,2+N], color='C2')
ax[2,0].set_xlabel('Distance (m)')
ax[2,0].set_ylabel('Pressure (Pa)')

# plot the site fraction of the surface species along the flow direction
for i,name in enumerate(gas_Si_N_interface.species_names):
    ax[2,1].plot(time, solution[:,3+N+i], label=name)
ax[2,1].legend(fontsize=8)
ax[2,1].set_xlabel('Distance (m)')
ax[2,1].set_ylabel('Site Fraction')

# plot the temperature profile along the flow direction
ax[3,0].plot(time, solution[:,-1], color='C3')
ax[3,0].set_xlabel('Distance (m)')
ax[3,0].set_ylabel('Temperature (K)')
f.tight_layout(pad=0.5)
```


These are the packages we need:

```python
import sympy as sp
import numpy as np
from itertools import combinations_with_replacement as itTuples
import os.path
from multiprocessing import Pool
```

# Below you will find *all* functions defined in the module.

Generates al possible combination --of length k-- from a List of elements

```python
def Tuples(List,k):
    return list(itTuples(List,k))
```

MatrixProd([list1,list2,list3,...]) returns np.dot(list1,np.dot(list2,np.dot(list3,...)).
Notice that it works recursively.

```python
def MatrixProd(a):
    n=len(a)-1
    if n!=0:
        return np.dot(MatrixProd(a[:n]),a[n])
    else:
        return a[0]
```

Calculates the derivative of $\mathcal{L}$ with respect to $\phi$: $\frac{d\mathcal{L}}{d\phi}$ (with all fields sent to 0).

If $\phi$ is a list, it calculates the derivative $\frac{d^{n}\mathcal{L}}{d\phi_{1}d\phi_{2}...d\phi_{n}}$.

This is used to get the feynmann rules.

```python
def Deriv(L,a):
    try:

        n=len(a)-1
        if n>=0:
            return sp.diff(Deriv(L,a[:n]),a[n])
        else:
            return L
    except:
        return sp.diff(L,a)

```

Get specific assumtions--given in list(assL)--for a Symbol--Sym.

It is used for ParameterSymbols, which is a list of all parameters and their assumtions.

```python
def GetAssumptions(Sym,assL):
    tmpA=[]
    for i in assL:
        try:
            tmpA.append(Sym.assumptions0[i] )
        except:
            tmpA.append(None )
    return tmpA
```

Defines the paricles, parameters of the model for SU(DimN).

If Gauge='un', the G^{0}, G^{+} and G^{-} are not defined.

Also, this function defines the various useful rules for substitution needed, such as subs0, which sets all fields to 0 (needed for idendifying the vertices and minimizing the potential).

```python
def Definitions(DimN, Gauge):
    global gauge, dimN
    global dimRange, indexRange, mPhi2, mPhip2, v, vPhi, muH, lamH, lamHPhi, lamPhi
    global Gp, H0, Gm, H0t, h, G0, H, Ht, Phi, Phit, chi, rho, phi, s
    global  sqrt2, subsvev, subsexpand

    '''gauge, dimN, dimRange, indexRange, mPhi2, mPhip2, v, vPhi, muH, lamH, lamHPhi, lamPhi,\
            Gp, H0, Gm, H0t, h, G0, H, Ht, Phi, Phit, chi, rho, phi, s,\
            sqrt2, subsvev, subsexpand=0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\
            0, 0, 0'''

    dimN=DimN
    gauge=Gauge

    dimRange=np.arange(1,dimN+1);
    dimRangeM=range(dimN-1)
    indexRange=range(0,dimN);

    sqrt2=sp.sqrt(2);

    mPhi2=np.array( sp.symbols('mPhi2(1:{})(1:{})'.format(str(dimN+1),str(dimN+1)),complex=True,real=False      )         ).reshape(dimN,dimN)
    mPhi2[dimN-1][dimN-1]=sp.Symbol('mPhi2{}{}'.format(dimN,dimN),real=True )#this is real, due to the minimization conditions
    mPhip2=np.array( sp.symbols('mPhip2(1:{})(1:{})'.format(str(dimN+1),str(dimN+1)),complex=True,real=False      )         ).reshape(dimN,dimN)

    #make mPhi symmetric (faster than np.triu(mPhi,+1).T+np.triu(mPhi))
    for i in range(dimN):
        for j in range(i+1,dimN):
            mPhi2[j][i]=mPhi2[i][j]

    #make mPhip hermitian (faster than np.conjugate(np.triu(mPhi,+1).T)+np.triu(mPhi))
    for i in range(dimN):
        for j in range(i+1,dimN):
            mPhip2[j][i]=sp.conjugate(mPhip2[i][j])
    #make the diagonal  real. keep in mind that the squared elements of the diagonal are real.
    #So the elements can be either real or imaginary
    for i in range(dimN):
        exec( 'mPhip2[{}][{}]=sp.Symbol(  \'mPhip2{}{}\'  ,real=True)'.format(str(i),str(i),str(i+1),str(i+1)) )

    tmpMPHI=(np.triu(mPhi2)).reshape(dimN**2)

    ParameterSymbols= np.array( [ (tmpMPHI[i], GetAssumptions(tmpMPHI[i],['complex','real','positive'] ) ) \
                                 for i in np.nonzero(tmpMPHI)[0]] )

    tmpMPHI=(np.triu(mPhip2)).reshape(dimN**2)
    ParameterSymbols=np.append(ParameterSymbols, np.array( [ (tmpMPHI[i], GetAssumptions(tmpMPHI[i],['complex','real','positive'] ) )\
                                                            for i in np.nonzero(tmpMPHI)[0]] ) )
    del tmpMPHI
    #print EverySymbol

    Phi = sp.symbols('Phi1:{}'.format(str(dimN+1)))
    Phit = sp.symbols('Phi1:{}t'.format(str(dimN+1)))

    if gauge=='un':
        H0, H0t=sp.symbols('H0, H0t')
        H = [0,H0];
        Ht = [0, H0t];

    else:
        H0,H0t,Gp,Gm,G0=sp.symbols('H0,H0t,Gp,Gm,G0')
        H = [Gp,H0];
        Ht = [Gm, H0t];

    ##################--Declare symbols for expaned scalars
    phi = list(sp.symbols('phi1:{}'.format(str(dimN))))
    s = list(sp.symbols('s1:{}'.format(str(dimN))))
    h , chi, rho=sp.symbols('h chi rho')

    v=sp.Symbol('v',positive=True);
    vPhi=sp.Symbol('vPhi',positive=True);
    muH=sp.Symbol('muH');
    lamH=sp.Symbol('lamH',real=True,positive=True);
    lamHPhi=sp.Symbol('lamHPhi',real=True,positive=None);
    lamPhi=sp.Symbol('lamPhi',real=True,positive=True);

    ParameterSymbols=np.append(ParameterSymbols, np.array( [\
                                                            (v,GetAssumptions(v,['complex','real','positive'] )),\
                                                            (vPhi,GetAssumptions(vPhi,['complex','real','positive'] )),\
                                                            (lamH,GetAssumptions(lamH,['complex','real','positive'] )),\
                                                            (lamHPhi,GetAssumptions(lamHPhi,['complex','real','positive'] )),\
                                                            (lamPhi,GetAssumptions(lamPhi,['complex','real','positive'] ))]))

    #Expand the fields at their vevs
    if gauge=='un':
        subsexpand =np.array(\
                     [(H0,(h+v)/sqrt2 ),(H0t,(h+v)/sqrt2 ),\
                      (Phi[dimN-1],(rho+ sp.I*chi+vPhi)/sqrt2 ),\
                      (Phit[dimN-1],(rho-sp.I*chi+vPhi)/sqrt2 )]+ \
                     [(Phi[i], (phi[i]+sp.I*s[i])/sqrt2   ) for i in dimRangeM]+\
                     [(Phit[i],(phi[i]-sp.I*s[i])/sqrt2)  for i in dimRangeM])

        Fields=np.array(sp.flatten([h,rho,s,chi,phi]))
        subsvev = np.array(\
                      [(H0,v/sqrt2 ),(H0t,v/sqrt2 ),\
                      (Phi[dimN-1], vPhi/sqrt2 ),\
                      (Phit[dimN-1],vPhi/sqrt2 )]+ \
                     [(Phi[i], 0) for i in dimRangeM]+\
                     [(Phit[i],0) for i in dimRangeM])

    else:
        subsexpand = np.array(\
                      [(H0,(h+sp.I*G0+v)/sqrt2 ),(H0t,(h-sp.I*G0+v)/sqrt2 ),\
                      (Phi[dimN-1], (rho+sp.I*chi+vPhi)/sqrt2 ),\
                      (Phit[dimN-1],(rho-sp.I*chi+vPhi)/sqrt2 )]+ \
                     [(Phi[i], (phi[i]+sp.I*s[i])/sqrt2) for i in dimRangeM]+\
                     [(Phit[i],(phi[i]-sp.I*s[i])/sqrt2) for i in dimRangeM])

        Fields=np.array(sp.flatten([h,rho,s,chi,phi,G0,Gp,Gm]))

        subsvev = np.array(\
                          [(H0,v/sqrt2 ),(H0t,v/sqrt2 ),\
                           (G0,0),(Gm,0),(Gp,0),\
                          (Phi[dimN-1], vPhi/sqrt2 ),\
                          (Phit[dimN-1],vPhi/sqrt2 )]+ \
                         [(Phi[i], 0) for i in dimRangeM]+\
                         [(Phit[i],0) for i in dimRangeM])

    return list(Fields),ParameterSymbols
```

Should run after Definitions(DimN,Gauge)! Since all parameters, fields and rules are global parameters in Definitions, GetLagrangian(AllFields) takes them and calculates the Potential and returns the Lagrangian. Here we define the substitution rules for the minimization of the potential.

AllFields is needed in order to run CheckMinimizations, which check the vanishing of the first derivatives of the potential.

```python

def GetLagrangian(AllFields=False):
    #global V, constV, subsmin#these are for internal checks. Not really useful

    mPhi2C=[[sp.conjugate(i) for i in x] for x in mPhi2]

    V0=-muH**2/2*MatrixProd([H,Ht])+lamH/2*MatrixProd([H,Ht])**2+lamPhi/2*MatrixProd([Phi,Phit])**2\
    +lamHPhi*MatrixProd([H,Ht])*MatrixProd([Phi,Phit] );

    Vsoft=MatrixProd([Phi,mPhi2,Phi])+MatrixProd([Phit,mPhi2C,Phit])+MatrixProd([Phit,mPhip2,Phi])

    V=(V0+Vsoft)#.subs(subsexpand)

    subsmin= [ (mPhi2[i][dimN-1], -mPhip2[dimN-1][i]/2 ) for i in range(0,dimN-1)]+ \
    [(muH, sp.sqrt(v**2*lamH + vPhi**2*lamHPhi)),\
    (lamPhi,-(lamHPhi*v**2 + 2*mPhi2[dimN-1][dimN-1] + 2*mPhip2[dimN-1][dimN-1] + 2*sp.conjugate(mPhi2[dimN-1][dimN-1]))/vPhi**2),\
    (sp.conjugate(mPhi2[dimN-1][dimN-1]),mPhi2[dimN-1][dimN-1] )]

    constV=sp.simplify((V.subs(subsmin).subs(subsvev)) )

    if AllFields!=False:
        try:
            CheckMinimizations(AllFields,V, constV, subsmin)
        except:
            print 'Something went wrong while checking the minimization. \nHave you passed the fields correctly? '

    LMassInt = -( (V.subs(subsmin)).subs(subsexpand) -constV );
    return LMassInt

def CheckMinimizations(AllFields,V, constV, subsmin):#uses only global
    subs0=[ (i,0) for i in AllFields]

    print 'Checking vanishing of the first derivatives of the potential...'

    minV=np.unique(map(lambda i: \
             sp.simplify(Deriv(V.subs(subsexpand),i ).subs(subs0).subs(subsmin) ),AllFields))
    if (minV==0).all():
        print 'The conditions are correct!'
    else:
        print 'The potential is not minimized correctlly...'

```

IdentifyInteractions(Langrangian,All_Fields,Parallel=True/False) idendifies the 2-,3-,4-point interactions
of the Fields given a Langrangian. It returns a dictionary of the form:

$$\rm{ \{2:[2-point interactions], 3:[3-point interactions],4:[4-point interactions]\} }$$

DEF_TMP is needed so that Pool does not complain:
Pool needs defined functions at the top level. So we need to define a functions which defines TMP_int (called in
IdentifyInteractions)

TMP_int calculates the derivative of the Lagrangian with respect to particles (a list of particles) and returns  the particles, tmpval=the interaction term in the Langrangian and SymF=the symmetry factor (the factorial of the number of identical particles )

```python
def DEF_TMP(Langrangian,Fields):
    set_fields_to_0=[(i,0) for i in Fields  ]
    global TMP_int
    def TMP_int(particles):

        SymF=np.product([ sp.factorial(particles.count(j)) for j in set(particles)])
        tmpval=1/SymF*sp.simplify(Deriv(Langrangian,particles).subs(set_fields_to_0))
        if tmpval!=0:
            return [particles, tmpval,SymF]
        else:
            return 0

OPTIONS_Int=['Parallel']
DEF_OPT_Int={'Parallel':True}
def IdentifyInteractions(Langrangian,All_Fields,**opts):

    #----------------Begin check opts
    if len(opts) == 0:
        print 'Using default options...'
        opts=DEF_OPT_Int

    for i in opts:
        if not (i in OPTIONS_Int):
            print 'invalid option '+i
            print 'availabe options: '
            print OPTIONS_Int
            return 'ERR:: invalid option. Abort!'

    xtmp=opts.copy()
    for i in OPTIONS_Int:
        if not (i in opts):
            xtmp.update({i:DEF_OPT_Int[i]})

    Parallel=xtmp['Parallel']

    if Parallel!=True:
        Parallel=False

    #----------------End check opts

    #extract all interactions  involving from Min_in to Max_int particles
    Min_int=2
    Max_int=4

    Point_N={}
    DEF_TMP(Langrangian,All_Fields)

    ###########################################################
    for i in range(Min_int,Max_int+1):
        tmpTuples=Tuples(All_Fields,i)
        print 'calculating {}-point interactions'.format(i)

        if Parallel:
            p=Pool()
            FR=np.array(p.map(TMP_int,tmpTuples))
            Point_N[i]= [FR[TMPI] for TMPI  in np.nonzero(FR)[0] ]
            p.close()
            del p,FR
        else:
            FR=np.array(map(TMP_int,tmpTuples))
            Point_N[i]= [FR[TMPI] for TMPI  in np.nonzero(FR)[0] ]
            del FR
    return Point_N

```

FRules takes a list  with the n-point interactions:

$$\text{ [(particles_1, interaction_term_1, symmetry_factor_1 ), (particles_2, interaction_term_2, symmetry_factor_2 ),...  ] }$$

and returns a list with the  feynamn rules and mass matrix entries.

It multiplies each 2-point interactions with -1*symmetry_factor (mass matrix entries).

It multiplies each n-point (n>2) interactions with -I*symmetry_factor (Feynman rules).

Make_Feynman_Rules calls FRules for a dictionary of the form:

$$\text{ \{2:[2-point interactions], 3:[3-point interactions],4:[4-point interactions]\} }$$

and calculates a dictionary --which is globaly available-- of the form:

$$\text{ \{2:[mass matrix entries], 3:[3-point Fenman rules],4:[4-point Fenman rules]\} }$$

```python

#The function prtcls, gets a list of particles (prts) and returns a sorted list of them
#(needed for use in FRules and VertexValue)
#Example:
#prtcls(['z','b'])---->('b', 'z')

def prtcls(prts):
    return  tuple(sorted( prts ) )
#---------------------------------------------------------------------

def FRules(N_Point):
    N=len(N_Point[0][0])
    NPoint_dict={}
    if N==2:
        for i in N_Point:
            NPoint_dict[prtcls( map( str, i[0] ) ) ]=i[1]*(-i[2])
    else:
        for i in N_Point:
            NPoint_dict[prtcls( map( str, i[0] ) ) ]=i[1]*(-sp.I*i[2])

    return NPoint_dict

def Make_Feynman_Rules(NPoint_dict):
    global DictP
    DictP={}
    for k in NPoint_dict.keys():

        DictP[k] = FRules(NPoint_dict[k])
```

VertexValue(particle_1, particle_2, ...) gets a number of particles, and returns the  corresponding Feynman rule (or
mass matrix entry, if the input consists of two particles).

```python

def VertexValue(*particles):
    lp=len(particles)

    try:
        return DictP[lp][ prtcls(   map(str, particles)  ) ]
        #return eval('DictP'+str(lp)+'[ prtcls(   map(str, particles)  ) ]'   )
    except:
        return 0

```

CheckInteractions  takes the output of IdentifyInteractions, the Initial_Lagrangian (subtracted by its constant term) and the Fields, and compares them.

```python
def CheckInteractions(N_Point_dict, Initial_Lagrangian,AllFields):

    if N_Point_dict!=False and Initial_Lagrangian!=False and AllFields!=False:
        testL=True
    else:
        testL=False

    if testL:
        global    LMassIntfinal, L_in
        print 'Checking Vertices...'
        LMassIntfinal=0
        SUBS0=[ (i,0) for i in AllFields]

        for TypeOfVert in N_Point_dict.keys():
            TypeV=N_Point_dict[TypeOfVert]

            LMassIntfinal+=np.sum([ np.product(tmpi[0])*tmpi[1] for tmpi in TypeV])
        L_in=Initial_Lagrangian-sp.simplify(Initial_Lagrangian.subs(SUBS0))

        if (sp.simplify(LMassIntfinal-L_in))==0:
            print 'The interactions have been identified correctly!!'
        else:
            print 'The final Lagrangian is not the same as the initial one... (check it!)'
```

StoreVert takes a dictionary of the form--the output of IdentifyInteractions-- :

$\rm{ \{2:[2-point interactions], 3:[3-point interactions],4:[4-point interactions]\} },$
All Fields, All Parameter Symbols, and  prints files with the Feynman rules, mass matrix entries, fields and parameetrs.

Change 'Directory', to store them in anoter directory.

```python
def StoreVert(N_Points,AllFields,AllParameterSymbols,Directory='Frules'):
    print 'Writing Vertices (Feynman Rules and mass matrix entries)...'
    dirV=Directory

    if not os.path.exists(dirV):
        os.makedirs(dirV)

    if not os.path.exists(dirV+"/SU" + str(dimN)):
        os.makedirs(dirV+"/SU" + str(dimN))

    files=N_Points.keys()

    tmp =open(dirV+"/SU" + str(dimN)+ "/SU" + str(dimN) +'_'+gauge+ ".fields","w")
    [tmp.write(str(ff)+'\n') for ff in AllFields]

    tmp =open(dirV+"/SU" + str(dimN)+ "/SU" + str(dimN)+'_'+gauge+".parameters","w")
    [tmp.write(str(ff)+'\n') for ff in AllParameterSymbols]

    for file in files:
        tmp = open(dirV+"/SU" + str(dimN)+ "/SU" + str(dimN)+"_" +str(file)+"-point_"+gauge + ".vrt","w")
        if file==2:
            factorI=-1
        else:
            factorI=-sp.I
        for i in N_Points[file]:
            particles=str(i[0])

            vertex=str(factorI*i[1]*i[2])

            line='{:<40} {:<40} {:<0}'.format(particles, '|' , vertex)
            #tmp.write( particles +"|\t|"+ vertex + "\n" )
            tmp.write( line +'\n')
        tmp.close()
    print 'All Done!'
```

# Putting everything together:

```python
#Run first this. This example defines SU(2) in the Feynman gauge.
Fields ,ParameterSymbol =Definitions(2,'feyn')

#The definitions can be used to construct the Lagrangian.
LMassInt=GetLagrangian(Fields)

#The Lagrangian can be used to find all interaction terms.
Point_N=IdentifyInteractions(LMassInt,Fields ,Parallel=True)

#Once the interactions are known, Make_Feynman_Rules makes the Feynamn rules,
#and defines a global dictionary DictP used in VertexValue.
Make_Feynman_Rules(Point_N)

#This checks that the interactions have been identified correctly.
CheckInteractions(Point_N,LMassInt,Fields )

#Saves the Feynamn rules, and parameters in a directory (./test in this case)
StoreVert(Point_N,Fields ,ParameterSymbol,'test' )
```

    Checking vanishing of the first derivatives of the potential...
    The conditions are correct!
    calculating 2-point interactions
    calculating 3-point interactions
    calculating 4-point interactions
    Checking Vertices...
    The interactions have been identified correctly!!
    Writing Vertices (Feynman Rules and mass matrix entries)...
    All Done!

```python

```


```python
%matplotlib inline
```

# Introduction to Optimal Transport with Python

This example gives an introduction on how to use Optimal Transport in Python.

```python
# Author: Remi Flamary, Nicolas Courty, Aurelie Boisbunon
#
# License: MIT License
# sphinx_gallery_thumbnail_number = 1
```

## POT Python Optimal Transport Toolbox

### POT installation

* Install with pip::

    pip install pot
* Install with conda::

    conda install -c conda-forge pot

### Import the toolbox

```python
import numpy as np  # always need it
import pylab as pl  # do the plots

import ot  # ot

import time
```

### Getting help

Online  documentation : `<https://pythonot.github.io/all.html>`_

Or inline help:

```python
help(ot.dist)
```

## First OT Problem

We will solve the Bakery/Cafés problem of transporting croissants from a
number of Bakeries to Cafés in a City (in this case Manhattan). We did a
quick google map search in Manhattan for bakeries and Cafés:

We extracted from this search their positions and generated fictional
production and sale number (that both sum to the same value).

We have acess to the position of Bakeries ``bakery_pos`` and their
respective production ``bakery_prod`` which describe the source
distribution. The Cafés where the croissants are sold are defined also by
their position ``cafe_pos`` and ``cafe_prod``, and describe the target
distribution. For fun we also provide a
map ``Imap`` that will illustrate the position of these shops in the city.

Now we load the data

```python
data = np.load('../data/manhattan.npz')

bakery_pos = data['bakery_pos']
bakery_prod = data['bakery_prod']
cafe_pos = data['cafe_pos']
cafe_prod = data['cafe_prod']
Imap = data['Imap']

print('Bakery production: {}'.format(bakery_prod))
print('Cafe sale: {}'.format(cafe_prod))
print('Total croissants : {}'.format(cafe_prod.sum()))
```

## Plotting bakeries in the city

Next we plot the position of the bakeries and cafés on the map. The size of
the circle is proportional to their production.

```python
pl.figure(1, (7, 6))
pl.clf()
pl.imshow(Imap, interpolation='bilinear')  # plot the map
pl.scatter(bakery_pos[:, 0], bakery_pos[:, 1], s=bakery_prod, c='r', ec='k', label='Bakeries')
pl.scatter(cafe_pos[:, 0], cafe_pos[:, 1], s=cafe_prod, c='b', ec='k', label='Cafés')
pl.legend()
pl.title('Manhattan Bakeries and Cafés')
```

## Cost matrix

We can now compute the cost matrix between the bakeries and the cafés, which
will be the transport cost matrix. This can be done using the
`ot.dist <https://pythonot.github.io/all.html#ot.dist>`_ function that
defaults to squared Euclidean distance but can return other things such as
cityblock (or Manhattan distance).

```python
C = ot.dist(bakery_pos, cafe_pos)

labels = [str(i) for i in range(len(bakery_prod))]
f = pl.figure(2, (14, 7))
pl.clf()
pl.subplot(121)
pl.imshow(Imap, interpolation='bilinear')  # plot the map
for i in range(len(cafe_pos)):
    pl.text(cafe_pos[i, 0], cafe_pos[i, 1], labels[i], color='b',
            fontsize=14, fontweight='bold', ha='center', va='center')
for i in range(len(bakery_pos)):
    pl.text(bakery_pos[i, 0], bakery_pos[i, 1], labels[i], color='r',
            fontsize=14, fontweight='bold', ha='center', va='center')
pl.title('Manhattan Bakeries and Cafés')

ax = pl.subplot(122)
im = pl.imshow(C, cmap="coolwarm")
pl.title('Cost matrix')
cbar = pl.colorbar(im, ax=ax, shrink=0.5, use_gridspec=True)
cbar.ax.set_ylabel("cost", rotation=-90, va="bottom")

pl.xlabel('Cafés')
pl.ylabel('Bakeries')
pl.tight_layout()
```

The red cells in the matrix image show the bakeries and cafés that are
further away, and thus more costly to transport from one to the other, while
the blue ones show those that are very close to each other, with respect to
the squared Euclidean distance.

## Solving the OT problem with `ot.emd <https://pythonot.github.io/all.html#ot.emd>`_

```python
start = time.time()
ot_emd = ot.emd(bakery_prod, cafe_prod, C)
time_emd = time.time() - start
```

The function returns the transport matrix, which we can then visualize (next section).

### Transportation plan vizualization

A good vizualization of the OT matrix in the 2D plane is to denote the
transportation of mass between a Bakery and a Café by a line. This can easily
be done with a double ``for`` loop.

In order to make it more interpretable one can also use the ``alpha``
parameter of plot and set it to ``alpha=G[i,j]/G.max()``.

```python
# Plot the matrix and the map
f = pl.figure(3, (14, 7))
pl.clf()
pl.subplot(121)
pl.imshow(Imap, interpolation='bilinear')  # plot the map
for i in range(len(bakery_pos)):
    for j in range(len(cafe_pos)):
        pl.plot([bakery_pos[i, 0], cafe_pos[j, 0]], [bakery_pos[i, 1], cafe_pos[j, 1]],
                '-k', lw=3. * ot_emd[i, j] / ot_emd.max())
for i in range(len(cafe_pos)):
    pl.text(cafe_pos[i, 0], cafe_pos[i, 1], labels[i], color='b', fontsize=14,
            fontweight='bold', ha='center', va='center')
for i in range(len(bakery_pos)):
    pl.text(bakery_pos[i, 0], bakery_pos[i, 1], labels[i], color='r', fontsize=14,
            fontweight='bold', ha='center', va='center')
pl.title('Manhattan Bakeries and Cafés')

ax = pl.subplot(122)
im = pl.imshow(ot_emd)
for i in range(len(bakery_prod)):
    for j in range(len(cafe_prod)):
        text = ax.text(j, i, '{0:g}'.format(ot_emd[i, j]),
                       ha="center", va="center", color="w")
pl.title('Transport matrix')

pl.xlabel('Cafés')
pl.ylabel('Bakeries')
pl.tight_layout()
```

The transport matrix gives the number of croissants that can be transported
from each bakery to each café. We can see that the bakeries only need to
transport croissants to one or two cafés, the transport matrix being very
sparse.

## OT loss and dual variables

The resulting wasserstein loss loss is of the form:

\begin{align}W=\sum_{i,j}\gamma_{i,j}C_{i,j}\end{align}

where $\gamma$ is the optimal transport matrix.

```python
W = np.sum(ot_emd * C)
print('Wasserstein loss (EMD) = {0:.2f}'.format(W))
```

## Regularized OT with Sinkhorn

The Sinkhorn algorithm is very simple to code. You can implement it directly
using the following pseudo-code

In this algorithm, $\oslash$ corresponds to the element-wise division.

An alternative is to use the POT toolbox with
`ot.sinkhorn <https://pythonot.github.io/all.html#ot.sinkhorn>`_

Be careful of numerical problems. A good pre-processing for Sinkhorn is to
divide the cost matrix ``C`` by its maximum value.

### Algorithm

```python
# Compute Sinkhorn transport matrix from algorithm
reg = 0.1
K = np.exp(-C / C.max() / reg)
nit = 100
u = np.ones((len(bakery_prod), ))
for i in range(1, nit):
    v = cafe_prod / np.dot(K.T, u)
    u = bakery_prod / (np.dot(K, v))
ot_sink_algo = np.atleast_2d(u).T * (K * v.T)  # Equivalent to np.dot(np.diag(u), np.dot(K, np.diag(v)))

# Compute Sinkhorn transport matrix with POT
ot_sinkhorn = ot.sinkhorn(bakery_prod, cafe_prod, reg=reg, M=C / C.max())

# Difference between the 2
print('Difference between algo and ot.sinkhorn = {0:.2g}'.format(np.sum(np.power(ot_sink_algo - ot_sinkhorn, 2))))
```

### Plot the matrix and the map

```python
print('Min. of Sinkhorn\'s transport matrix = {0:.2g}'.format(np.min(ot_sinkhorn)))

f = pl.figure(4, (13, 6))
pl.clf()
pl.subplot(121)
pl.imshow(Imap, interpolation='bilinear')  # plot the map
for i in range(len(bakery_pos)):
    for j in range(len(cafe_pos)):
        pl.plot([bakery_pos[i, 0], cafe_pos[j, 0]],
                [bakery_pos[i, 1], cafe_pos[j, 1]],
                '-k', lw=3. * ot_sinkhorn[i, j] / ot_sinkhorn.max())
for i in range(len(cafe_pos)):
    pl.text(cafe_pos[i, 0], cafe_pos[i, 1], labels[i], color='b',
            fontsize=14, fontweight='bold', ha='center', va='center')
for i in range(len(bakery_pos)):
    pl.text(bakery_pos[i, 0], bakery_pos[i, 1], labels[i], color='r',
            fontsize=14, fontweight='bold', ha='center', va='center')
pl.title('Manhattan Bakeries and Cafés')

ax = pl.subplot(122)
im = pl.imshow(ot_sinkhorn)
for i in range(len(bakery_prod)):
    for j in range(len(cafe_prod)):
        text = ax.text(j, i, np.round(ot_sinkhorn[i, j], 1),
                       ha="center", va="center", color="w")
pl.title('Transport matrix')

pl.xlabel('Cafés')
pl.ylabel('Bakeries')
pl.tight_layout()
```

We notice right away that the matrix is not sparse at all with Sinkhorn,
each bakery delivering croissants to all 5 cafés with that solution. Also,
this solution gives a transport with fractions, which does not make sense
in the case of croissants. This was not the case with EMD.

### Varying the regularization parameter in Sinkhorn

```python
reg_parameter = np.logspace(-3, 0, 20)
W_sinkhorn_reg = np.zeros((len(reg_parameter), ))
time_sinkhorn_reg = np.zeros((len(reg_parameter), ))

f = pl.figure(5, (14, 5))
pl.clf()
max_ot = 100  # plot matrices with the same colorbar
for k in range(len(reg_parameter)):
    start = time.time()
    ot_sinkhorn = ot.sinkhorn(bakery_prod, cafe_prod, reg=reg_parameter[k], M=C / C.max())
    time_sinkhorn_reg[k] = time.time() - start

    if k % 4 == 0 and k > 0:  # we only plot a few
        ax = pl.subplot(1, 5, k // 4)
        im = pl.imshow(ot_sinkhorn, vmin=0, vmax=max_ot)
        pl.title('reg={0:.2g}'.format(reg_parameter[k]))
        pl.xlabel('Cafés')
        pl.ylabel('Bakeries')

    # Compute the Wasserstein loss for Sinkhorn, and compare with EMD
    W_sinkhorn_reg[k] = np.sum(ot_sinkhorn * C)
pl.tight_layout()
```

This series of graph shows that the solution of Sinkhorn starts with something
very similar to EMD (although not sparse) for very small values of the
regularization parameter, and tends to a more uniform solution as the
regularization parameter increases.

### Wasserstein loss and computational time

```python
# Plot the matrix and the map
f = pl.figure(6, (4, 4))
pl.clf()
pl.title("Comparison between Sinkhorn and EMD")

pl.plot(reg_parameter, W_sinkhorn_reg, 'o', label="Sinkhorn")
XLim = pl.xlim()
pl.plot(XLim, [W, W], '--k', label="EMD")
pl.legend()
pl.xlabel("reg")
pl.ylabel("Wasserstein loss")
```

In this last graph, we show the impact of the regularization parameter on
the Wasserstein loss. We can see that higher
values of ``reg`` leads to a much higher Wasserstein loss.

The Wasserstein loss of EMD is displayed for
comparison. The Wasserstein loss of Sinkhorn can be a little lower than that
of EMD for low values of ``reg``, but it quickly gets much higher.