# CF Part 7 - Explainable Matrix factorization method
> Collaborative Filtering on MovieLens Latest-small Part 7 - Finding recommendations using model based explainable matrix factorization method

- toc: true
- badges: true
- comments: true
- categories: [movie, collaborative]
- image:

### How to quantify explainability ?

- Use the rating distribution within the active user’s neighborhood.
- If many neighbors have rated the recommended item, then this can provide a basis upon which to explain the recommendations, using neighborhood style explanation mechanisms

According to [(Abdollahi and Nasraoui, 2016)](https://www.researchgate.net/publication/301616080_Explainable_Matrix_Factorization_for_Collaborative_Filtering), an item $i$ is consider to be explainable for user $u$ if a considerable number of its neighbors rated item $i$. The explainability score $E_{ui}$ is the percentage of user $u$'s neighbors who have rated item $i$.

\begin{equation}
E_{ui} = \frac{|N_k^{(i)}(u)|}{|N_k(u)|},
\end{equation}

where $N_k(u)$ is the set of $k$ nearest neighbors of user $u$ and $N_k^{(i)}(u)$ is the set of user $u$'s neighbors who have rated item $i$. However, only explainable scores above an optimal threshold $\theta$ are accepted.

\begin{equation}
W_{ui} = \begin{cases} E_{ui} \text{  } if \text{  } E_{ui} > \theta \\ 0 \text{ } otherwise \end{cases},
\end{equation}

By including explainability weight in the training algorithm, the new objective function, to be minimized over the set of known ratings, has been formulated by [(Abdollahi and Nasraoui, 2016)](https://www.researchgate.net/publication/301616080_Explainable_Matrix_Factorization_for_Collaborative_Filtering) as:

\begin{equation}
 J = \sum_{(u,i)\in \kappa} (R_{ui} - \hat{R}_{ui})^2 +\frac{\beta}{2}(||P_u||^2 + ||Q_i||^2) + \frac{\lambda}{2}(P_u-Q_i)^2W_{ui},
\end{equation}

here, $\frac{\beta}{2}(||P_u||^2 + ||Q_i||^2)$ is the $L_2$ regularization term weighted by the coefficient $\beta$, and $\lambda$ is an explainability regularization coefficient that controls the smoothness of the new representation and tradeoff between explainability and accuracy. The idea here is that if item $i$ is explainable for user $u$, then their representations in the latent space, $Q_i$ and $P_u$, should be close to each other. Stochastic Gradient descent can be used to optimize the objectve function.

\begin{equation}
P_u \leftarrow P_u + \alpha\left(2(R_{u,i}-P_uQ_i^{\top})Q_i - \beta P_u - \lambda(P_u-Q_i)W_{ui}\right)
\end{equation}
\begin{equation}
Q_i \leftarrow Q_i + \alpha\left(2(R_{u,i}-P_uQ_i^{\top})P_u - \beta Q_i + \lambda(P_u-Q_i)W_{ui}\right)
\end{equation}

## Explainable Matrix Factorization : Model definition

```python
import os

if not (os.path.exists("recsys.zip") or os.path.exists("recsys")):
    !wget https://github.com/nzhinusoftcm/review-on-collaborative-filtering/raw/master/recsys.zip
    !unzip recsys.zip
```

### requirements

```
matplotlib==3.2.2
numpy==1.18.1
pandas==1.0.5
python==3.6.10
scikit-learn==0.23.1
scipy==1.5.0
```

```python
from recsys.memories.UserToUser import UserToUser

from recsys.preprocessing import mean_ratings
from recsys.preprocessing import normalized_ratings
from recsys.preprocessing import ids_encoder
from recsys.preprocessing import train_test_split
from recsys.preprocessing import rating_matrix
from recsys.preprocessing import get_examples

from recsys.datasets import ml100k, ml1m

import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

import sys
import os
```

### Compute Explainable Scores

Explainable score are computed using neighborhood based similarities. Here, we are using the user based algorithme to compute similarities

```python
def explainable_score(user2user, users, items, theta=0):

    def _progress(count):
        sys.stdout.write('\rCompute Explainable score. Progress status : %.1f%%'%(float(count/len(users))*100.0))
        sys.stdout.flush()
    # initialize explainable score to zeros
    W = np.zeros((len(users), len(items)))

    for count, u in enumerate(users):
        candidate_items = user2user.find_user_candidate_items(u)
        for i in candidate_items:
            user_who_rated_i, similar_user_who_rated_i = \
                user2user.similar_users_who_rated_this_item(u, i)
            if user_who_rated_i.shape[0] == 0:
                w = 0.0
            else:
                w = similar_user_who_rated_i.shape[0] / user_who_rated_i.shape[0]
            W[u,i] =  w  if w > theta else 0.0
        _progress(count)
    return W
```

### Explainable Matrix Factorization Model

```python
class ExplainableMatrixFactorization:

    def __init__(self, m, n, W, alpha=0.001, beta=0.01, lamb=0.1, k=10):
        """
            - R : Rating matrix of shape (m,n)
            - W : Explainability Weights of shape (m,n)
            - k : number of latent factors
            - beta : L2 regularization parameter
            - lamb : explainability regularization coefficient
            - theta : threshold above which an item is explainable for a user
        """
        self.W = W
        self.m = m
        self.n = n

        np.random.seed(64)

        # initialize the latent factor matrices P and Q (of shapes (m,k) and (n,k) respectively) that will be learnt
        self.k = k
        self.P = np.random.normal(size=(self.m,k))
        self.Q = np.random.normal(size=(self.n,k))

        # hyperparameter initialization
        self.alpha = alpha
        self.beta = beta
        self.lamb = lamb

        # training history
        self.history = {
            "epochs":[],
            "loss":[],
            "val_loss":[],
        }

    def print_training_parameters(self):
        print('Training EMF')
        print(f'k={self.k} \t alpha={self.alpha} \t beta={self.beta} \t lambda={self.lamb}')

    def update_rule(self, u, i, error):
        self.P[u] = self.P[u] + \
            self.alpha*(2 * error*self.Q[i] - self.beta*self.P[u] - self.lamb*(self.P[u] - self.Q[i]) * self.W[u,i])

        self.Q[i] = self.Q[i] + \
            self.alpha*(2 * error*self.P[u] - self.beta*self.Q[i] + self.lamb*(self.P[u] - self.Q[i]) * self.W[u,i])

    def mae(self,  x_train, y_train):
        """
        returns the Mean Absolute Error
        """
        # number of training exemples
        M = x_train.shape[0]
        error = 0
        for pair, r in zip(x_train, y_train):
            u, i = pair
            error += np.absolute(r - np.dot(self.P[u], self.Q[i]))
        return error/M

    def print_training_progress(self, epoch, epochs, error, val_error, steps=5):
        if epoch == 1 or epoch % steps == 0 :
                print(f"epoch {epoch}/{epochs} - loss : {round(error,3)} - val_loss : {round(val_error,3)}")

    def learning_rate_schedule(self, epoch, target_epochs = 20):
        if (epoch >= target_epochs) and (epoch % target_epochs == 0):
                factor = epoch // target_epochs
                self.alpha = self.alpha * (1 / (factor * 20))
                print("\nLearning Rate : {}\n".format(self.alpha))

    def fit(self, x_train, y_train, validation_data, epochs=10):
        """
        Train latent factors P and Q according to the training set

        :param
            - x_train : training pairs (u,i) for which rating r_ui is known
            - y_train : set of ratings r_ui for all training pairs (u,i)
            - validation_data : tuple (x_test, y_test)
            - epochs : number of time to loop over the entire training set.
            10 epochs by default

        Note that u and i are encoded values of userid and itemid
        """
        self.print_training_parameters()

        # get validation data
        x_test, y_test = validation_data

        for epoch in range(1, epochs+1):
            for pair, r in zip(x_train, y_train):
                u,i = pair
                r_hat = np.dot(self.P[u], self.Q[i])
                e = r - r_hat
                self.update_rule(u, i, error=e)

            # training and validation error  after this epochs
            error = self.mae(x_train, y_train)
            val_error = self.mae(x_test, y_test)
            self.update_history(epoch, error, val_error)
            self.print_training_progress(epoch, epochs, error, val_error, steps=1)

        return self.history

    def update_history(self, epoch, error, val_error):
        self.history['epochs'].append(epoch)
        self.history['loss'].append(error)
        self.history['val_loss'].append(val_error)

    def evaluate(self, x_test, y_test):
        """
        compute the global error on the test set

        :param
            - x_test : test pairs (u,i) for which rating r_ui is known
            - y_test : set of ratings r_ui for all test pairs (u,i)
        """
        error = self.mae(x_test, y_test)
        print(f"validation error : {round(error,3)}")

    def predict(self, userid, itemid):
        """
        Make rating prediction for a user on an item

        :param
        - userid
        - itemid

        :return
        - r : predicted rating
        """
        # encode user and item ids to be able to access their latent factors in
        # matrices P and Q
        u = uencoder.transform([userid])[0]
        i = iencoder.transform([itemid])[0]

        # rating prediction using encoded ids. Dot product between P_u and Q_i
        r = np.dot(self.P[u], self.Q[i])

        return r

    def recommend(self, userid, N=30):
        """
        make to N recommendations for a given user

        :return
        - (top_items,preds) : top N items with the highest predictions
        """
        # encode the userid
        u = uencoder.transform([userid])[0]

        # predictions for this user on all product
        predictions = np.dot(self.P[u], self.Q.T)

        # get the indices of the top N predictions
        top_idx = np.flip(np.argsort(predictions))[:N]

        # decode indices to get their corresponding itemids
        top_items = iencoder.inverse_transform(top_idx)

        # take corresponding predictions for top N indices
        preds = predictions[top_idx]

        return top_items, preds
```

```python
epochs = 10
```

# Model Evaluation

## 1. MovieLens 100K

### 1.1. Evaluation on raw data

```python
# load data
ratings, movies = ml100k.load()

# encode users and items ids
ratings, uencoder, iencoder = ids_encoder(ratings)

users = sorted(ratings.userid.unique())
items = sorted(ratings.itemid.unique())

m = len(users)
n = len(items)

# get examples as tuples of userids and itemids and labels from normalize ratings
raw_examples, raw_labels = get_examples(ratings)

# train test split
(x_train, x_test), (y_train, y_test) = train_test_split(examples=raw_examples, labels=raw_labels)
```

```python
# create the user to user model for similarity measure
usertouser = UserToUser(ratings, movies)

# compute explainable score
W = explainable_score(usertouser, users, items)
```

    Normalize users ratings ...
    Initialize the similarity model ...
    Compute nearest neighbors ...
    User to user recommendation model created with success ...
    Compute Explainable score. Progress status : 99.9%

```python
# initialize the model
EMF = ExplainableMatrixFactorization(m, n, W, alpha=0.01, beta=0.4, lamb=0.01, k=10)

history = EMF.fit(x_train, y_train, epochs=epochs, validation_data=(x_test, y_test))
```

    Training EMF
    k=10 	 alpha=0.01 	 beta=0.4 	 lambda=0.01
    epoch 1/10 - loss : 0.922 - val_loss : 1.036
    epoch 2/10 - loss : 0.79 - val_loss : 0.873
    epoch 3/10 - loss : 0.766 - val_loss : 0.837
    epoch 4/10 - loss : 0.757 - val_loss : 0.822
    epoch 5/10 - loss : 0.753 - val_loss : 0.814
    epoch 6/10 - loss : 0.751 - val_loss : 0.808
    epoch 7/10 - loss : 0.749 - val_loss : 0.805
    epoch 8/10 - loss : 0.748 - val_loss : 0.802
    epoch 9/10 - loss : 0.746 - val_loss : 0.799
    epoch 10/10 - loss : 0.745 - val_loss : 0.797

```python
EMF.evaluate(x_test, y_test)
```

    validation error : 0.797

### 1.2. Evaluation on normalized data

```python
# load data
ratings, movies = ml100k.load()

# encode users and items ids
ratings, uencoder, iencoder = ids_encoder(ratings)

users = sorted(ratings.userid.unique())
items = sorted(ratings.itemid.unique())

m = len(users)
n = len(items)

# normalize ratings by substracting means
normalized_column_name = "norm_rating"
ratings = normalized_ratings(ratings, norm_column=normalized_column_name)

# get examples as tuples of userids and itemids and labels from normalize ratings
raw_examples, raw_labels = get_examples(ratings, labels_column=normalized_column_name)

# train test split
(x_train, x_test), (y_train, y_test) = train_test_split(examples=raw_examples, labels=raw_labels)
```

```python
# initialize the model
EMF = ExplainableMatrixFactorization(m, n, W, alpha=0.022, beta=0.65, lamb=0.01, k=10)

history = EMF.fit(x_train, y_train, epochs=epochs, validation_data=(x_test, y_test))
```

    Training EMF
    k=10 	 alpha=0.022 	 beta=0.65 	 lambda=0.01
    epoch 1/10 - loss : 0.809 - val_loss : 0.842
    epoch 2/10 - loss : 0.809 - val_loss : 0.829
    epoch 3/10 - loss : 0.807 - val_loss : 0.821
    epoch 4/10 - loss : 0.799 - val_loss : 0.811
    epoch 5/10 - loss : 0.789 - val_loss : 0.8
    epoch 6/10 - loss : 0.782 - val_loss : 0.793
    epoch 7/10 - loss : 0.778 - val_loss : 0.789
    epoch 8/10 - loss : 0.776 - val_loss : 0.786
    epoch 9/10 - loss : 0.774 - val_loss : 0.784
    epoch 10/10 - loss : 0.773 - val_loss : 0.783

## 2. MovieLens 1M

### 2.1. Evaluation on raw data

```python
# load data
ratings, movies = ml1m.load()

# encode users and items ids
ratings, uencoder, iencoder = ids_encoder(ratings)

users = sorted(ratings.userid.unique())
items = sorted(ratings.itemid.unique())

m = len(users)
n = len(items)

# get examples as tuples of userids and itemids and labels from normalize ratings
raw_examples, raw_labels = get_examples(ratings)

# train test split
(x_train, x_test), (y_train, y_test) = train_test_split(examples=raw_examples, labels=raw_labels)
```

```python
# create the user to user model for similarity measure
usertouser = UserToUser(ratings, movies)

# compute explainable score
W = explainable_score(usertouser, users, items)
```

    Normalize users ratings ...
    Initialize the similarity model ...
    Compute nearest neighbors ...
    User to user recommendation model created with success ...
    Compute Explainable score. Progress status : 100.0%

```python
# initialize the model
EMF = ExplainableMatrixFactorization(m, n, W, alpha=0.01, beta=0.4, lamb=0.01, k=10)

history = EMF.fit(x_train, y_train, epochs=epochs, validation_data=(x_test, y_test))
```

    Training EMF
    k=10 	 alpha=0.01 	 beta=0.4 	 lambda=0.01
    epoch 1/10 - loss : 0.782 - val_loss : 0.807
    epoch 2/10 - loss : 0.762 - val_loss : 0.781
    epoch 3/10 - loss : 0.76 - val_loss : 0.775
    epoch 4/10 - loss : 0.758 - val_loss : 0.771
    epoch 5/10 - loss : 0.757 - val_loss : 0.769
    epoch 6/10 - loss : 0.756 - val_loss : 0.767
    epoch 7/10 - loss : 0.754 - val_loss : 0.764
    epoch 8/10 - loss : 0.752 - val_loss : 0.762
    epoch 9/10 - loss : 0.751 - val_loss : 0.761
    epoch 10/10 - loss : 0.75 - val_loss : 0.76

### 2.2. Evaluation on normalized data

```python
# load data
ratings, movies = ml1m.load()

# encode users and items ids
ratings, uencoder, iencoder = ids_encoder(ratings)

# normalize ratings by substracting means
normalized_column_name = "norm_rating"
ratings = normalized_ratings(ratings, norm_column=normalized_column_name)

# get examples as tuples of userids and itemids and labels from normalize ratings
raw_examples, raw_labels = get_examples(ratings, labels_column=normalized_column_name)

# train test split
(x_train, x_test), (y_train, y_test) = train_test_split(examples=raw_examples, labels=raw_labels)
```

```python
# initialize the model
EMF = ExplainableMatrixFactorization(m, n, W, alpha=0.023, beta=0.59, lamb=0.01, k=10)

history = EMF.fit(x_train, y_train, epochs=epochs, validation_data=(x_test, y_test))
```

    Training EMF
    k=10 	 alpha=0.023 	 beta=0.59 	 lambda=0.01
    epoch 1/10 - loss : 0.805 - val_loss : 0.814
    epoch 2/10 - loss : 0.764 - val_loss : 0.77
    epoch 3/10 - loss : 0.756 - val_loss : 0.762
    epoch 4/10 - loss : 0.755 - val_loss : 0.759
    epoch 5/10 - loss : 0.754 - val_loss : 0.759
    epoch 6/10 - loss : 0.754 - val_loss : 0.758
    epoch 7/10 - loss : 0.754 - val_loss : 0.758
    epoch 8/10 - loss : 0.753 - val_loss : 0.758
    epoch 9/10 - loss : 0.753 - val_loss : 0.758
    epoch 10/10 - loss : 0.753 - val_loss : 0.758

# Ratings prediction

```python
# get list of top N items with their corresponding predicted ratings
userid = 42
recommended_items, predictions = EMF.recommend(userid=userid)

# find corresponding movie titles
top_N = list(zip(recommended_items,predictions))
top_N = pd.DataFrame(top_N, columns=['itemid','predictions'])
top_N.predictions = top_N.predictions + ratings.loc[ratings.userid==userid].rating_mean.values[0]
List = pd.merge(top_N, movies, on='itemid', how='inner')

# show the list
List
```

      itemid
      predictions
      title
      genres

      0
      3460
      4.364036
      Hillbillys in a Haunted House (1967)
      Comedy

      1
      701
      4.324177
      Daens (1992)
      Drama

      2
      3057
      4.307404
      Where's Marlowe? (1999)
      Comedy

      3
      2214
      4.304979
      Number Seventeen (1932)
      Thriller

      4
      1145
      4.299559
      Snowriders (1996)
      Documentary

      5
      2258
      4.292125
      Master Ninja I (1984)
      Action

      6
      3353
      4.281912
      Closer You Get, The (2000)
      Comedy|Romance

      7
      868
      4.278937
      Death in Brunswick (1991)
      Comedy

      8
      826
      4.269901
      Diebinnen (1995)
      Drama

      9
      3305
      4.266769
      Bluebeard (1944)
      Film-Noir|Horror

      10
      2619
      4.265997
      Mascara (1999)
      Drama

      11
      763
      4.264092
      Last of the High Kings, The (a.k.a. Summer Fli...
      Drama

      12
      1852
      4.262517
      Love Walked In (1998)
      Drama|Thriller

      13
      642
      4.260353
      Roula (1995)
      Drama

      14
      682
      4.258829
      Tigrero: A Film That Was Never Made (1994)
      Documentary|Drama

      15
      792
      4.253339
      Hungarian Fairy Tale, A (1987)
      Fantasy

      16
      1316
      4.252915
      Anna (1996)
      Drama

      17
      3228
      4.245526
      Wirey Spindell (1999)
      Comedy

      18
      853
      4.240745
      Dingo (1992)
      Drama

      19
      3172
      4.238188
      Ulysses (Ulisse) (1954)
      Adventure

      20
      2254
      4.238008
      Choices (1981)
      Drama

      21
      2503
      4.234547
      Apple, The (Sib) (1998)
      Drama

      22
      2905
      4.224974
      Sanjuro (1962)
      Action|Adventure

      23
      744
      4.224278
      Brothers in Trouble (1995)
      Drama

      24
      757
      4.224226
      Ashes of Time (1994)
      Drama

      25
      858
      4.223665
      Godfather, The (1972)
      Action|Crime|Drama

      26
      789
      4.220788
      I, Worst of All (Yo, la peor de todas) (1990)
      Drama

      27
      3748
      4.216508
      Match, The (1999)
      Comedy|Romance

      28
      790
      4.216455
      An Unforgettable Summer (1994)
      Drama

      29
      745
      4.215986
      Close Shave, A (1995)
      Animation|Comedy|Thriller

**Note**: The recommendation list may content items already purchased by the user. This is just an illustration of how to implement matrix factorization recommender system. You can optimize the recommended list and return the top rated items that the user has not already purchased.

## Putting all together

This repository covers the following algorithms for collaborative filtering

1. User-based CF
2. Item-based CF
3. Singular Value Decomposition (SVD)
4. Matrix Factorization (MF)
5. Non-negative Matrix Factorization (NMF)
6. Explainable Matrix Factorization (EMF)

In the [next notebook](https://github.com/nzhinusoftcm/review-on-collaborative-filtering/blob/master/8.Performances_measure.ipynb), we present performances comparison of all these models.

## Reference

1. Yehuda Koren et al. (2009). Matrix Factorization Techniques for Recommender Systems
2. Abdollahi and Nasraoui (2016). [Explainable Matrix Factorization for Collaborative Filtering](https://www.researchgate.net/publication/301616080_Explainable_Matrix_Factorization_for_Collaborative_Filtering)
3. Abdollahi and Nasraoui (2017). [Using Explainability for Constrained Matrix Factorization](https://dl.acm.org/doi/abs/10.1145/3109859.3109913)
4. Shuo Wang et al, (2018). [Explainable Matrix Factorization with Constraints on Neighborhood in the Latent Space](https://dl.acm.org/doi/abs/10.1145/3109859.3109913)

## Author

[Carmel WENGA](https://www.linkedin.com/in/carmel-wenga-871876178/),

PhD student at Université de la Polynésie Française,

Applied Machine Learning Research Engineer,

[ShoppingList](https://shoppinglist.cm), NzhinuSoft.


# Computational Astrophysics
## Ordinary Differential Equations. 01

---
## Eduard Larrañaga

Observatorio Astronómico Nacional\
Facultad de Ciencias\
Universidad Nacional de Colombia

---

### About this notebook

In this notebook we present some of the techniques to solve a system of Ordinary Differential Equations.

`A. Garcia. Numerical Methods for Physics. (1999). Chapter 3 `

---

## Ordinary Differential Equations

A system of first-order ordinary differential equations (ODEs)
is a relationship between an unknown (vectorial) function
$y(x)$ and its derivative $y^\prime(x)$. The
general system of first-order ODEs has the form

\begin{equation}
y^\prime (x) =f(x,y(x))\,.
\end{equation}

A solution to the differential equation is,
obviously, any function $y(x)$ that satisfies it.

There are two general classes of first-order ODE problems:

1. Initial value problems: $y(x_i)$ is given at some starting point $x_i$.

2. Two-point boundary value problems: $y$ is known at two ends (boundaries) of the domain and these **boundary conditions** must be satisfied simultaneously.

---
### Reduction to First-Order ODE

Any ODE can be reduced to first-order form by introducing additional
variables.

#### Example

\begin{equation}
y''(x) + q(x) y'(x) = r(x)\,\,.
\end{equation}

Introducing a new function $z(x)$ this can be written as a system of first-order ODE,

\begin{align}
\text{(1)}\,\, y'(x) &= z(x)\\
\text{(2)}\,\, z'(x) &= r(x) - q(x) z(x).
\end{align}

---
## Explicit and Implicit Methods of Solution of ODE

Broadly speaking, the methods for integrating ODEs are clasiffied into *explicit* and *implicit* methods.

Explicit methods use differences of the system to construct and update of the solution using the current (and previous) values of the dependent variables. Some examples of explicit methods are Euler's and Runge-Kutta's methods.

Implicit methods evaluate the right-hand side using the new value of the function that we are solving for.

---
## Errors and ODE

All procedures to solve numerically an ODE consist of transforming
a continuous differential equation into a discrete iteration procedure
that starts from the initial conditions and returns the values of the
dependent variable $y(x)$ at points $x_m = x_0 + m * h$, where $h$
is the discretization step size  assumed to be constant here).

Two kinds of errors can arise in this procedure:
1. **Round-off error.** Due to limited float point accuracy. The global round-off is the sum of the local float point errors.

2. **Truncation error.**

* Local: The error made in one step when we replace a continuous process
  (e.g. a derivative) with a discrete procedure (e.g., a forward difference).

* Global: If the local truncation error is $\mathcal{O}(h^{n+1})$, then the global truncation error must be $\mathcal{O}(h^n)$, since the number of steps used in evaluating the derivatives to reach an arbitrary point $x_f$, having started at $x_0$, is $\frac{x_f - x_0}{h}$.

---
---

## Euler's Method

We want to solve the ODE

\begin{equation}
y' = f(x,y)
\end{equation}

with the initial condition $y(x_0) = y_0$. We introduce
a fixed stepsize $h$ and obtain an estimate of $y(x)$ at
$x_1 = x_0 + h$ using Taylor's theorem,

\begin{equation}
\begin{aligned}
y(x_1) &= y(x_0+h) = y(x_0) + y'(x_0) h + \mathcal{O}(h^2)\,,\\
       &= y(x_0) + h f(x_0,y(x_0)) +  \mathcal{O}(h^2)\,.
\end{aligned}
\end{equation}

By analogy, we obtain that the value $y_{n+1}$ of the function at the
point $x_{n+1} = x_0 + (n+1) h$ is given by

\begin{equation}
y_{n+1} = y(x_{n+1}) = y_n + h f(x_n,y(x_n)) + \mathcal{O}(h^2)\,.
\end{equation}

This is called the **forward Euler Method**.

---
Euler's method is extremely simple, but rather inaccurate and potentially unstable.

The error scales $\propto h^2$ locally. However, if $L$ is
the length of the domain, then $h = L / N$, where $N$ is the number of points
used to cover it. Since we are taking $N$ integration steps, the global
error is $\propto N h^2 = N L^2 / N^2 = L L/N \propto h$.

Hence, forward Euler is a first-order accurate method globally.

### Stability

Forward Euler is an **explicit** method. This means that
$y_{n+1}$ is given explicitly in terms of known quantities
$y_n$ and $f(x_n,y_n)$.

Explicit methods are simple and efficient, but the drawback is that
the step size must be small for stability.

### Example

\begin{equation}
\frac{dy}{dt} = -a y\,\,, \hspace{1cm}\text{with} \hspace{1cm}
y(0) = 1\,,\,\, a > 0\,\,.
\end{equation}

The exact solution to this problem is $y = e^{-at}$,
which is stable and  smooth with $y(0) = 1$ and
$y(\infty) = 0$.

Applying forward Euler,

\begin{equation}
y_{n+1} = y_n - a\, h\, y_n = (1-ah) y_{n}
\end{equation}

\begin{equation}
y_{n+1}  = (1-ah)^2 y_{n-1} = \cdots = (1-ah)^{n+1} y_0\,\,.
\end{equation}

This implies that in order to prevent any potential amplification of
errors, we must require that $|1-ah|<1$.

In fact, there are  3 cases,

| | | |
|:-: | :-: | :--|
| $0 < 1-ah < 1$  | : | $(1-ah)^{n+1}$ decays (good!) |
|$-1 < 1-ah < 0$ | : | $(1-ah)^{n+1}$ oscillates (not so good!)|
|$1-ah < -1$     | : | $(1-ah)^{n+1}$ oscillates and diverges (bad!)|

This gives the stability criterion of $0< h < \frac{2}{a}$.

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

def ODE(y):
    a=2.
    rhs = -a*y
    return rhs

def FEuler(h, t0, q0):
    '''
    ------------------------------------------
    FEuler(h, t0, q0)
    ------------------------------------------
    Forward Euler's method for solving a ODEs
    system.
    Arguments:
    h: stepsize for the iteration
    t0: independent parameter initial value
    q0: numpy array with the initial values of
        the functions in the ODEs system
    ------------------------------------------
    '''
    f = ODE(q0)
    q1 = q0 + h*f
    return q1

# Creation of the Grid
t_0 = 0.
t_f = 10.
n = 100
h = 0.2
#h = (t_f - t_0)/n

# Initial value
y_0 = 1.

# Variables setup
t = np.linspace(t_0, t_f, n)
y = np.zeros(n)
y[0] = y_0

# Main Loop
for i in range(1,n):
    y[i] = FEuler(h, t[i-1], y[i-1])

# Plot
fig, ax = plt.subplots()
ax.plot(t, y, label=f'$h=$ {h:.1f}')
ax.set_title('Forward Euler Method')
ax.set_xlabel(r'$t$')
ax.set_ylabel(r'$y(t)$')
plt.legend()
plt.show()
```

Lets check the stability of the method.

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

def ODE(y):
    a=2.
    rhs = -a*y
    return rhs

def FEuler(h, t0, q0):
    '''
    ------------------------------------------
    FEuler(h, t0, q0)
    ------------------------------------------
    Forward Euler's method for solving a ODEs
    system.
    Arguments:
    h: stepsize for the iteration
    t0: independent parameter initial value
    q0: numpy array with the initial values of
        the functions in the ODEs system
    ------------------------------------------
    '''
    f = ODE(q0)
    q1 = q0 + h*f
    return q1

# Initial Condition
y_0 = 1.

# Creation of the grid
t_0 = 0.
t_f = 10.

# Three stepsizes to compare
n1 = 100
h1 = (t_f - t_0)/n1

n2 = 15
h2 = (t_f - t_0)/n2

n3 = 10
h3 = (t_f - t_0)/n3

# Arrays to store the solution
t1 = np.linspace(t_0, t_f, n1)
y1 = np.zeros(n1)
y1[0] = y_0

t2 = np.linspace(t_0, t_f, n2)
y2 = np.zeros(n2)
y2[0] = y_0

t3 = np.linspace(t_0, t_f, n3)
y3 = np.zeros(n3)
y3[0] = y_0

for i in range(1,n1):
    y1[i] = FEuler(h1, t1[i-1], y1[i-1])

for i in range(1,n2):
    y2[i] = FEuler(h2, t2[i-1], y2[i-1])

for i in range(1,n3):
    y3[i] = FEuler(h3, t3[i-1], y3[i-1])

fig, ax = plt.subplots()
ax.plot(t3, y3, color='crimson', label=f'$h_3=$ {h3:.1f}')
ax.plot(t2, y2, color='orange', label=f'$h_2=$ {h2:.1f}')
ax.plot(t1, y1, color='cornflowerblue', label=f'$h_1=$ {h1:.1f}')

ax.set_title('Forward Euler Method')
ax.set_xlabel(r'$t$')
ax.set_ylabel(r'$y(t)$')
plt.legend()
plt.show()
```

---
## Predictor-Corrector Method

Consider the modification

\begin{equation}
y_{n+1} = y_{n} + h\, \frac{f(x_n, y_n) + f(x_{n+1},y_{n+1})}{2}\,\,.
\end{equation}

This may be a better estimate as it is using the *average slope* of $y$.

However, we don't know $y_{n+1}$ yet. We can get around this problem by using forward Euler to estimate $y_{n+1}$ and then use the above modification for a better estimate:

\begin{equation}
\begin{aligned}
y_{n+1}^{(\mathrm{P})} &= y_n + h f(x_n,y_n)\,\,, & \text{(predictor)}\\
y_{n+1} &= y_n + \frac{h}{2} \left[f(x_n,y_n) + f(x_{n+1}, y^{(\mathrm{P})}_{n+1}) \right]\,\,. & \text{(corrector)}
\end{aligned}
\end{equation}

One can show that the error of the **predictor-corrector method**
decreases locally with $h^3$, but globally with $h^2$.  One says it is
**second-order accurate** as opposed to the Euler method, which is
first-order accurate.

---
---

## Runge-Kutta Methods

The idea behind Runge-Kutta (RK) methods is to match the Taylor expansion
of $y(x)$ at $x=x_n$ up to the highest possible order.

---
### RK2. Second Order RK Method

For

\begin{equation}
\frac{dy}{dx} = f(x,y)\,\,,
\end{equation}

we have

\begin{equation}
y_{n+1} = y_n + a k_1 + b k_2\,\,,
\end{equation}

with

\begin{align}
k_1 &= h\,f(x_n,y_n)\,\,,\nonumber\\
k_2 &= h\,f(x_n+\alpha h,y_n + \beta k_1)\,\,.
\end{align}

The four parameters $a,b,\alpha,\beta$ will be fixed so that the equation for $y_{n+1}$ agrees as well as possible with the Taylor series expansion
of $y' = f(x,y)$:

\begin{align}
y_{n+1} &= y_n + h y_n' + \frac{h^2}{2} y_n'' + \mathcal{O}(h^3)\,\,,\nonumber\\
&= y_n + h f(x_n,y_n) + \frac{h^2}{2} \frac{d}{dx} f(x_n,y_n) + \mathcal{O}(h^3)\,\,,\nonumber\\
&= y_n + h f_n + h^2 \frac{1}{2}\left( \frac{\partial f_n}{\partial x} + \frac{\partial f_n}{\partial y} f_n \right) + \mathcal{O}(h^3)\,\,,
\end{align}
where $f_n = f(x_n,y_n)$.

Now, we will consider our supposition

\begin{equation}
y_{n+1} = y_n + a h f_n + b h f(x_n + \alpha h, y_n + \beta h f_n)\,\,.
\label{eq:rk2c}
\end{equation}

Expanding the last term in a Taylor series to first order in terms of $(x_n,y_n)$, we get

\begin{equation}
y_{n+1} = y_n + a h f_n + bh\left[ f_n + \frac{\partial f}{\partial x}(x_n,y_n) \alpha h +
\frac{\partial f}{\partial y}(x_n,y_n) \beta h f_n \right]\,\,,
\end{equation}

and compare to read off the constants

\begin{equation}
a+b = 1\,\,,\hspace{2em} \alpha b = \frac{1}{2}\,\,\hspace{2em} \beta b = \frac{1}{2}\,\,.
\end{equation}

So there are only 3 equations for 4 unknowns and we can assign an arbitrary value to one of the unknowns. Typical choices are:

\begin{equation}
\alpha = \beta = \frac{1}{2}\,\,,\hspace{2em} a=0\,\,, \hspace{2em} b = 1\,\,.
\end{equation}

With this, we have for RK2:

\begin{align}
k_1 & = h f(x_n,y_n)\,\,,\\
k_2 & = h f(x_n + \frac{1}{2}h, y_n + \frac{1}{2}k_1)\,\,,\\
y_{n+1} & = y_n + k_2 + \mathcal{O}(h^3)\,\,.
\end{align}

This method is locally $\mathcal{O}(h^3)$, but globally
only $\mathcal{O}(h^2)$.

**Note:** Using $a=b=1/2$ and $\alpha=\beta=1$, we recover the predictor-corrector method!

---

### RK3

\begin{align}
k_1 &=h f(x_n,y_n)\,\,\nonumber\\
k_2 &= h f(x_n + \frac{h}{2},y_n + \frac{1}{2} k_1)\,\,,\nonumber\\
k_3 &= h f(x_n + h, y_n - k_1 + 2 k_2)\,\,,\nonumber\\
y_{n+1} &= y_n + \frac{1}{6}(k_1 + 4 k_2 + k_3) + \mathcal{O}(h^4)\,\,.
\end{align}

---

### RK4

\begin{align}
k_1 &=h f(x_n,y_n)\,\,,\\
k_2 &=h f(x_n + \frac{h}{2},y_n + \frac{1}{2}k_1)\,\,,\nonumber\\
k_3 &=h f(x_n + \frac{h}{2},y_n + \frac{1}{2}k_2)\,\,,\nonumber\\
k_4 &=h f(x_n+h,y_n + k_3)\,\,,\nonumber\\
y_{n+1} &= y_n + \frac{1}{6}(k_1 + 2 k_2 + 2 k_3 + k_4) + \mathcal{O}(h^5)\,\,.
\end{align}

---

### Example

Consider again the ODE

\begin{equation}
\frac{dy}{dt} = -a y\,\,, \hspace{1cm}\text{with} \hspace{1cm}
y(0) = 1\,,\,\, a > 0
\end{equation}

with $y(0) = 1$.

```python
import numpy as np
import matplotlib.pyplot as plt
from RK4 import *

def ODE(x, y, a=1.):
    rhs = -a*y
    return rhs

# Initial Condition
y_0 = 1.

# Creation of the grid
t_0 = 0.
t_f = 20.

# Two stepsizes to compare
n1 = 100
h1 = (t_f - t_0)/n1

n2 = 15
h2 = (t_f - t_0)/n2

n3 = 10
h3 = (t_f - t_0)/n3

# Arrays to store the solution
t1 = np.linspace(t_0, t_f, n1)
y1 = np.zeros(n1)
y1[0] = y_0

t2 = np.linspace(t_0, t_f, n2)
y2 = np.zeros(n2)
y2[0] = y_0

t3 = np.linspace(t_0, t_f, n3)
y3 = np.zeros(n3)
y3[0] = y_0

for i in range(1,n1):
    y1[i] = RK4(ODE, h1, 0, y1[i-1])

for i in range(1,n2):
    y2[i] = RK4(ODE, h2, 0, y2[i-1])

for i in range(1,n3):
    y3[i] = RK4(ODE, h3, 0, y3[i-1])

fig, ax = plt.subplots()
ax.plot(t1, y1, color='cornflowerblue', label=f'$h_1=$ {h1:.1f}')
ax.plot(t2, y2, color='orange', label=f'$h_2=$ {h2:.1f}')
ax.plot(t3, y3, color='crimson', label=f'$h_3=$ {h3:.1f}')

ax.set_title('RK4 Method')
ax.set_xlabel(r'$t$')
ax.set_ylabel(r'$y(t)$')
plt.legend()
plt.show()

```

## Runge-Kutta Methods with Adaptive Step Size

The RK methods above require choosing a fixed step
size $h$ but, how should one choose this parameter?

It would be better to choose an **error tolerance** and let $h$ be chosen
automatically to satisfy this error tolerance.

This implies that  we need

1. A method for estimating the error.
2. A way to adjust the stepsize $h$, if the error is too large (or too small).

---
### Embedded RK Formulae

Embedded RK formulae provide an error estimator.  Now we will present the scheme for  3rd/4th order embedded RK (Bogaki and Shampine)

\begin{equation}
\begin{aligned}
k_1 &= h f(x_n, y_n)\,\,,\\
k_2 &= h f(x_n + \frac{1}{2} h, y_n + \frac{1}{2} k_1)\,\,,\\
k_3 &= h f(x_n + \frac{3}{4} h, y_n + \frac{3}{4} k_2)\,\,,\\
y_{n+1} &= y_n + \frac{2}{9} k_1 + \frac{1}{3} k_2 + \frac{4}{9} k_3 + \mathcal{O}(h^4)\,\,\\
k_4 &= h f(x_n + h, y_{n+1})\,\,\\
y^*_{n+1} &= y_n + \frac{7}{24} k_1 + \frac{1}{4}k_2 + \frac{1}{3}k_3 + \frac{1}{8} k_4 + \mathcal{O}(h^3)\,\,.
\end{aligned}
\end{equation}

The error is

\begin{equation}
\delta y_{n+1} = y_{n+1} - y^*_{n+1}\,\,.
\end{equation}

In this scheme, $k_4$ of step $n$ is the same as $k_1$ of step $n+1$. Therefore, $k_1$ does not need to be recomputed on step $n+1$; simply
save $k_4$ and re-use it on the next step.
This trick is called **FSAL (First Same As Last)**.

Given the error estimate $\delta y_{n+1} = y_{n+1} - y^*_{n+1}$
we want that it says smaller than some tolerance, $|\delta y_{n+1}| \le
\epsilon$ by adjusting $h$.

Usually, one sets

\begin{equation}
\epsilon = \epsilon_a + |y_{n+1}|\epsilon_r
\end{equation}

where

$\epsilon_a$ : absolute error tolerance.

$\epsilon_r$ : relative error tolerance.

### Adjusting the Step Size $h$

We define

\begin{equation}
\Delta = \frac{|\delta y_{n+1}|}{\epsilon}\,\,,
\end{equation}

and we want $\Delta \approx 1$.

Note that for a $p$-th-order formula, $\Delta \sim \mathcal{O}(h^p)$.
So if you took a step $h$ and got a value $\Delta$, then you change the step to $h_\text{desired}$,

\begin{equation}
h_\text{desired} = h \left|\frac{\Delta_\text{desired}}{\Delta} \right| ^\frac{1}{p}\,\,,
\end{equation}

to get the new  $\Delta_\text{desired} = 1$.

---

The algorithm to adjust $h$ can be written as follows:
1. Take step $h$, measure $\Delta$.
2. If $\Delta > 1$ (error too large), then

  - set $h_\text{new} = h \left| \frac{1}{\Delta} \right|^{\frac{1}{p}} S$, where $S$ is a fudge factor ($\sim$ $0.9$ or so).

  - **reject** the old step, redo with $h_\mathrm{new}$.

3. If $\Delta < 1$ (error too small), then

  - set $h_\text{new} = h \left| \frac{1}{\Delta} \right|^{\frac{1}{p}} S$.

  - **accept** old step, take next step with $h_\text{new}$.

```python

```


```python
%load_ext autoreload
%autoreload 2
import matplotlib.pyplot as plt
import numpy as np
import sys
import sympy as sp

sys.path.insert(0, '../')
from AircraftIden import FreqIdenSIMO, TransferFunctionFit, TransferFunctionParamModel
from AircraftIden.TransferFunctionFit import plot_fitter
```

    The autoreload extension is already loaded. To reload it, use:
      %reload_ext autoreload

```python
arr = np.load("../data/sweep_data_2017_11_16_11_47.npy")
time_seq_source = arr[:, 0] # Time series
ele_seq_source = arr[:, 1] # elevator input
q_seq_source = arr[:, 4] # angular velocity output
vvi_seq_source = arr[:,6] # verticle velocity output
theta_seq_source = arr[:,2] # pitch angle output
airspeed_seq = arr[:, 3] # airspeed output
simo_iden = FreqIdenSIMO(time_seq_source, 0.1, 50, ele_seq_source, q_seq_source,theta_seq_source,airspeed_seq, win_num=None)
```

    Sample rate is 46.70436437941652
    Sample rate is 46.70436437941652
    Setting-up composite win, sample rate 46.70436437941652 from omg 0.1 to 50
    Composite using [5, 61, 117, 173, 229]
    Setting-up composite win, sample rate 46.70436437941652 from omg 0.1 to 50
    Composite using [5, 61, 117, 173, 229]
    Setting-up composite win, sample rate 46.70436437941652 from omg 0.1 to 50
    Composite using [5, 61, 117, 173, 229]

```python
plt.rc("figure", figsize=(15,10))
plt.figure("ele->q")
simo_iden.plt_bode_plot(0)
plt.show()

plt.figure("ele->airspeed")
simo_iden.plt_bode_plot(2)
# Airspeed response is mainly at low frequency domain
plt.show()
```

## Fitting transfer function
Use sympy to present your transfer function and fit it!

A good estimation usually gives cost less than 15~20

Here we fit the elevator to pitch rate as a second order system with a latency

$$\frac{d}{a s^2 + b s +c} e^{-\tau s}$$

```python
freq, H, gamma2, gxx, gxy, gyy = simo_iden.get_freq_iden(0)
a, b, c, d, e, f, g, tau, s = sp.symbols("a b c d e f g tau s")
num = d*s+e
den = a*s*s + b*s + c
tfpm = TransferFunctionParamModel(num, den, tau)
fitter = TransferFunctionFit(freq, H, gamma2, tfpm, nw=20, iter_times=50, reg = 0.1)
tf = fitter.estimate(0.1, 10, accept_J=15)
plot_fitter(fitter, "$ele -> q$")
```

```python

```


###### Content under Creative Commons Attriubtion license CC-BY 4.0, code under MIT license (c)2015 Ye Qin.

# The Fisher-Kolmogorov equation

The Fisher-Kolmogorov equation (also known as KPP equation or Fisher's equation) has a form of
$$
\frac{\partial u}{\partial t}=ru(1-u)+D\frac{\partial u^2 }{\partial x^2}
$$
It can be used in many fields such as economy, chemistry, and so on.

## Discretizing our model equation

In our model, the Fisher's equation is used to describe the diffusion of an epidemic in a population in a 1-D space. Denote $n(x,t)$ the density of infected individuals, the model is formulated as:
$$
\frac{\partial n}{\partial t}=r(1-\frac{n}{K})n+\chi\frac{\partial n^2 }{\partial x^2}
\tag{1}
$$
where $t$ is time, $x$ is space, $r$ is rate of growth, $K$ is carrying capacity, and $\chi$ is diffusivity.
Eq. (1) seems complex. We can simplify it by using dimensionless variable. Let $u=\frac{n} {K}$, $\tilde{t}=rt$, and $\tilde{x}=\sqrt{\frac{r}{\chi}}x$. Then, Eq. (1) can be rewritten as
$$
\frac{\partial u}{\partial \tilde{t}}=(1-u)u+\frac{\partial u^2 }{\partial \tilde{x}^2}
\tag{2}
$$
For convenience, we drop the notation $\tilde{}$ in the following.

We establish the discrete system as:
$$
x_{i} = i\Delta x \quad  \quad t^{n}= n\Delta t
$$
$$
u_{i}^{n} = u(i\Delta x, n\Delta t)
$$

Choose Patankar scheme for the non-linear term, which means $\frac{\partial u}{\partial t}=u(1-u)$ is approximated as
$$
\frac{u_{i}^{n+1}-u_{i}^{n}}{\Delta t}=u_{i}^{n}(1-u_{i}^{n+1})
$$
Use forward-difference in time and central-difference in space, the partial differential part is approximated as
$$
\frac{u_{i}^{n+1}-u_{i}^{n}}{\Delta t}=\frac{u_{i-1}^{n}-2u_{i}^{n}+u_{i+1}^{n}}{\Delta x^2}
$$
Combine them, the discrete form of Eq. 2 is
$$
\frac{u_{i}^{n+1}-u_{i}^{n}}{\Delta t}=u_{i}^{n}(1-u_{i}^{n+1})+\frac{u_{i-1}^{n}-2u_{i}^{n}+u_{i+1}^{n}}{\Delta x^2}
$$
Put $u_{i}^{n+1}$ on the left-hand side ofthe equation and put the others on the right-hand side, we get
$$
u_{i}^{n+1}=(1+\Delta tu_{i}^{n})(u_{i}^{n}+\Delta t u_{i}^{n}+\frac{u_{i-1}^{n}-2u_{i}^{n}+u_{i+1}^{n}}{\Delta x^2})
$$
The stability criterion is $\sigma=\frac{\Delta t}{\Delta x^2}<0.5$.

## Compute!

First, define the initial and boundary conditions. They are defined as
\begin{equation}
u(x,0)=\begin{cases}1 & x=0\\
0 & 0<x<L
\end{cases}
\end{equation}
The length of the space $L$ is 150.

```python
import numpy
from matplotlib import pyplot
%matplotlib inline
from matplotlib import rcParams

rcParams['font.family'] = 'serif'
rcParams['font.size'] = 16
```

```python
# Parameters

L = 150             # Length
nx = 1200            # Number of grid points in space
dx = L/nx           # Space step
T = 56              # Period of time
sigma = 0.49        # To meet the stability criterion
dt = sigma*dx**2    # Time step

x = numpy.linspace(0,L,nx) # Define the space grid
```

Define the initial condition.

```python
def initial():
    u = numpy.zeros(nx)
    u[0]=1
    return u
```

```python
# How does it look like?
u = initial()
pyplot.figure(figsize=(10,5));
pyplot.plot(x, u, color='#003366', ls='-', lw=3)
pyplot.ylim(-1,1.5);
pyplot.xlim(-10,160);
pyplot.xlabel(r't', fontsize=18);
pyplot.ylabel(r'u', fontsize=18);
```

```python
def fisher(T,dt):
    """
    Solve the Fisher's equation

    Parameters
    ----------
    T : The period of time

    dt : time step
    ----------
    Produce a plot of the result
    """

    u = initial()       # Input the initial condition
    nt = int(T/dt)      # Number of grid points in time

    for n in range (nt):
        un = u.copy()
        u[1:-1] = (un[1:-1] + dt*un[1:-1]+dt/dx**2*(un[:-2]-2*un[1:-1]+un[2:]))/(1+dt*un[1:-1])

    # Plot the result
    pyplot.figure(figsize=(10,5));
    pyplot.plot(x, u, color='#003366', ls='-', lw=3)
    pyplot.ylim(-1,1.5);
    pyplot.xlim(-10,160);
    pyplot.xlabel(r'x', fontsize=18);
    pyplot.ylabel(r'u', fontsize=18);

```

Then, we can try some different $T$.

```python
fisher(10,dt)
```

```python
fisher(20,dt)
```

```python
fisher(30,dt)
```

```python
fisher(40,dt)
```

```python
fisher(60,dt)
```

Observe the figures above, we can find that location of the wave is always at around $x=2T$. It seems that the solution is a wave solution and the wave speed is $2$. However, we will talk about this in the next lesson.

Also, note that the solution on the figure aboves is for the dimensionalized system. When we want to get the solution of the initial variables, we just need to calculate $n(x,t)=Ku(x,t)$, $x=\tilde{x}\sqrt{\frac{\chi}{r}}$, and $t=\frac{\tilde{t}}{r}$.

### Dig deeper

For the non-linear term of the Fisher's equation, we use Patankar scheme and the term is approximated as
$$
\frac{u_{i}^{n+1}-u_{i}^{n}}{\Delta t}=u_{i}^{n}(1-u_{i}^{n+1})
$$
What would happened if this term is approximated as
$$
\frac{u_{i}^{n+1}-u_{i}^{n}}{\Delta t}=u_{i}^{n}(1-u_{i}^{n})
$$

## References:

+ Fisher's Equation, Wikipidia, https://en.wikipedia.org/wiki/Fisher%27s_equation
+ Schumacher, Eva.  "Ordinary and Fractional Diffusion in Simple Biological Models." https://www.uclouvain.be/cps/ucl/doc/commission-map/documents/Schumacher_TFE.pdf


# Demo - Stokes equations

**Mikael Mortensen** (email: `mikaem@math.uio.no`), Department of Mathematics, University of Oslo.

Date: **Jul 6, 2020**

Copyright 2020, Mikael Mortensen. Released under CC Attribution 4.0 license

**Summary.** The Stokes equations describe the flow of highly viscous fluids.
This is a demonstration of how the Python module [shenfun](https://github.com/spectralDNS/shenfun) can be used to solve Stokes
equations using a  mixed (coupled) basis in a 3D tensor product domain.
We assume homogeneous Dirichlet boundary conditions in one direction
and periodicity in the remaining two. The solver described runs with MPI
without any further considerations required from the user.
The solver assembles a block matrix with sparsity pattern as shown below
for the Legendre basis.

Coupled block matrix for Stokes equations.

# Model problem

## Stokes equations

The Stokes equations are given in strong form as

$$
\begin{align*}
\nabla^2 \mathbf{u} - \nabla p &= \mathbf{f} \quad \text{in }  \Omega, \\
\nabla \cdot \mathbf{u} &= h \quad \text{in } \Omega  \\
\int_{\Omega} p dx &= 0
\end{align*}
$$

where $\mathbf{u}$ and $p$ are, respectively, the
fluid velocity vector and pressure, and the domain
$\Omega = [0, 2\pi]^2 \times [-1, 1]$. The flow is assumed periodic
in $x$ and $y$-directions, whereas there is a no-slip homogeneous Dirichlet
boundary condition on $\mathbf{u}$ on the boundaries of the $z$-direction, i.e.,
$\mathbf{u}(x, y, \pm 1) = (0, 0, 0)$. (Note that we can configure shenfun with
non-periodicity in any of the three directions. However, since we are to
solve linear algebraic systems in the non-periodic direction, there is a speed
benefit from having the nonperiodic direction last. This has to do with Numpy
using a C-style row-major storage of arrays by default.)
The right hand side vector $\mathbf{f}(\mathbf{x})$ is an external applied body force.
The right hand side $h$ is usually zero in the regular Stokes equations. Here
we include it because it will be nonzero in the verification, which is using the
method of manufactured solutions. Note that the final $\int_{\Omega} p dx = 0$
is there because there is no Dirichlet boundary condition on the pressure
and the system of equations would otherwise be ill conditioned.

To solve Stokes equations with the Galerkin method we need basis
functions for both velocity and pressure. A
Dirichlet basis will be used for velocity, whereas there is no boundary restriction
on the pressure basis. For both three-dimensional bases we will use one basis
function for the $x$-direction,
$\mathcal{X}(x)$, one for the $y$-direction, $\mathcal{Y}(y)$, and one for the
$z$-direction, $\mathcal{Z}(z)$. And
then we create three-dimensional basis functions like

$$
\begin{equation}
v(x, y, z) = \mathcal{X}(x) \mathcal{Y}(y) \mathcal{Z} (z).
\label{_auto1} \tag{1}
\end{equation}
$$

The basis functions $\mathcal{X}(x)$ and $\mathcal{Y}(y)$ are chosen as Fourier
exponentials, since these functions are periodic:

$$
\begin{equation}
\mathcal{X}_l(x) = e^{\imath l x}, \forall \, l \in \mathbf{l}^{N_0},
\label{_auto2} \tag{2}
\end{equation}
$$

$$
\begin{equation}
\mathcal{Y}_m(y) =  e^{\imath m y}, \forall \, m \in \mathbf{m}^{N_1},
\label{_auto3} \tag{3}
\end{equation}
$$

where $\mathbf{l}^{N_0} = (-N_0/2, -N_0/2+1, \ldots, N_0/2-1)$ and
$\mathbf{m}^{N_1} = (-N_1/2, -N_1/2+1, \ldots, N_1/2-1)$.
The size of the discretized problem in real physical space is
$\mathbf{N} = (N_0, N_1, N_2)$, i.e., there are $N_0 \cdot N_1 \cdot N_2$ quadrature points
in total.

The basis functions for $\mathcal{Z}(z)$ remain to be decided.
For the velocity we need homogeneous Dirichlet boundary conditions, and for this
we use composite Legendre or Chebyshev polynomials

$$
\begin{equation}
\mathcal{Z}^0_n(z) = \phi_n(z) - \phi_{n+2}(z), \forall \, n \in \mathbf{n}^{N_2-2},
\label{_auto4} \tag{4}
\end{equation}
$$

where $\phi_n$ is the n'th Legendre or Chebyshev polynomial of the first kind.
$\mathbf{n}^{N_2-2} = (0, 1, \ldots, N_2-3)$, and the zero on $\mathcal{Z}^0$
is there to indicate the zero value on the boundary.

The pressure basis that comes with no restrictions for the boundary is a
little trickier. The reason for this has to do with
inf-sup stability. The obvious choice of basis is the regular Legendre or
Chebyshev basis, which is denoted as

$$
\begin{equation}
\mathcal{Z}_n(z) = \phi_n(z),  \forall \, n \in \mathbf{n}^{N_2}. \label{eq:Zn} \tag{5}
\end{equation}
$$

The problem is that for the natural choice of $n \in (0, 1, \ldots, N_2-1)$
there is a nullspace and one degree of freedom remains unresolved. It turns out
that the proper choice for the pressure basis is simply ([5](#eq:Zn)) for
$n \in \mathbf{n}^{N_2-2}$. (Also remember that we have to fix $\int_{\Omega} p dx = 0$.)

With given basis functions we obtain the spaces

$$
\begin{equation}
V^{N_0} = \text{span}\{ \mathcal{X}_l \}_{l\in\mathbf{l}^{N_0}},
\label{_auto5} \tag{6}
\end{equation}
$$

$$
\begin{equation}
V^{N_1} = \text{span}\{ \mathcal{Y}_m \}_{m\in\mathbf{m}^{N_1}},
\label{_auto6} \tag{7}
\end{equation}
$$

$$
\begin{equation}
V^{N_2} = \text{span}\{ \mathcal{Z}_n \}_{n\in\mathbf{n}^{N_2-2}},
\label{_auto7} \tag{8}
\end{equation}
$$

$$
\begin{equation}
V_0^{N_2} = \text{span}\{ \mathcal{Z}^0_n \}_{n\in\mathbf{n}^{N_2-2}},
\label{_auto8} \tag{9}
\end{equation}
$$

and from these we create two different tensor product spaces

$$
\begin{equation}
W_0^{\mathbf{N}}(\mathbf{x}) = V^{N_0}(x) \otimes V^{N_1}(y) \otimes V_0^{N_2}(z),
\label{_auto9} \tag{10}
\end{equation}
$$

$$
\begin{equation}
W^{\mathbf{N}}(\mathbf{x}) = V^{N_0}(x) \otimes V^{N_1}(y) \otimes V^{N_2}(z).
\label{_auto10} \tag{11}
\end{equation}
$$

The velocity vector is using a mixed basis, such that we will look for
solutions $\mathbf{u} \in [W_0^{\mathbf{N}}]^3 \, (=W_0^{\mathbf{N}} \times W_0^{\mathbf{N}} \times W_0^{\mathbf{N}})$,
whereas we look for the pressure
$p \in W^{\mathbf{N}}$. We now formulate a variational problem using the Galerkin method: Find
$\mathbf{u} \in [W_0^{\mathbf{N}}]^3$ and $p \in W^{\mathbf{N}}$ such that

$$
\begin{equation}
\int_{\Omega} (\nabla^2 \mathbf{u} - \nabla p ) \cdot \overline{\mathbf{v}} \, dx_w = \int_{\Omega} \mathbf{f} \cdot \overline{\mathbf{v}}\, dx_w \quad\forall \mathbf{v} \, \in \, [W_0^{\mathbf{N}}]^3, \label{eq:varform} \tag{12}
\end{equation}
$$

$$
\begin{equation}
\int_{\Omega} \nabla \cdot \mathbf{u} \, \overline{q} \, dx_w = \int_{\Omega} h \overline{q} \, dx_w \quad\forall q \, \in \, W^{\mathbf{N}}.
\label{_auto11} \tag{13}
\end{equation}
$$

Here $dx_w=w_xdxw_ydyw_zdz$ represents a weighted measure, with weights $w_x(x), w_y(y), w_z(z)$.
Note that it is only Chebyshev polynomials that
make use of a non-constant weight $w_x=1/\sqrt{1-x^2}$. The Fourier weights are $w_y=w_z=1/(2\pi)$
and the Legendre weight is $w_x=1$.
The overline in $\mathbf{\overline{v}}$ and $\overline{q}$ represents a complex conjugate, which is needed here because
the Fourier exponentials are complex functions.

## Mixed variational form

Since we are to solve for $\mathbf{u}$ and $p$ at the same time, we formulate a
mixed (coupled) problem: find $(\mathbf{u}, p) \in [W_0^{\mathbf{N}}]^3 \times W^{\mathbf{N}}$
such that

$$
\begin{equation}
a((\mathbf{u}, p), (\mathbf{v}, q)) = L((\mathbf{v}, q)) \quad \forall (\mathbf{v}, q) \in [W_0^{\mathbf{N}}]^3 \times W^{\mathbf{N}},
\label{_auto12} \tag{14}
\end{equation}
$$

where bilinear ($a$) and linear ($L$) forms are given as

$$
\begin{equation}
    a((\mathbf{u}, p), (\mathbf{v}, q)) = \int_{\Omega} (\nabla^2 \mathbf{u} - \nabla p) \cdot \overline{\mathbf{v}} \, dx_w + \int_{\Omega} \nabla \cdot \mathbf{u} \, \overline{q} \, dx_w,
\label{_auto13} \tag{15}
\end{equation}
$$

$$
\begin{equation}
    L((\mathbf{v}, q)) = \int_{\Omega} \mathbf{f} \cdot \overline{\mathbf{v}}\, dx_w + \int_{\Omega} h \overline{q} \, dx_w.
\label{_auto14} \tag{16}
\end{equation}
$$

Note that the bilinear form will assemble to block matrices, whereas the right hand side
linear form will assemble to block vectors.

# Implementation

## Preamble

We will solve the Stokes equations using the [shenfun](https://github.com/spectralDNS/shenfun) Python module. The first thing needed
is then to import some of this module's functionality
plus some other helper modules, like [Numpy](https://numpy.org) and [Sympy](https://sympy.org):

```
import os
import sys
import numpy as np
from mpi4py import MPI
from sympy import symbols, sin, cos
from shenfun import *
```

We use `Sympy` for the manufactured solution and `Numpy` for testing. MPI for
Python (`mpi4py`) is required for running the solver with MPI.

## Manufactured solution

The exact solutions $\mathbf{u}_e(\mathbf{x})$ and $p(\mathbf{x})$ are chosen to satisfy boundary
conditions, and the right hand sides $\mathbf{f}(\mathbf{x})$ and $h(\mathbf{x})$ are then
computed exactly using `Sympy`. These exact right hand sides will then be used to
compute a numerical solution that can be verified against the manufactured
solution. The chosen solution with computed right hand sides are:

```
x, y, z = symbols('x,y,z')
uex = sin(2*y)*(1-z**2)
uey = sin(2*x)*(1-z**2)
uez = sin(2*z)*(1-z**2)
pe = -0.1*sin(2*x)*cos(4*y)
fx = uex.diff(x, 2) + uex.diff(y, 2) + uex.diff(z, 2) - pe.diff(x, 1)
fy = uey.diff(x, 2) + uey.diff(y, 2) + uey.diff(z, 2) - pe.diff(y, 1)
fz = uez.diff(x, 2) + uez.diff(y, 2) + uez.diff(z, 2) - pe.diff(z, 1)
h = uex.diff(x, 1) + uey.diff(y, 1) + uez.diff(z, 1)
```

## Bases and tensor product spaces

Bases are created using the `Basis` function. A choice of
polynomials between Legendre or Chebyshev can be made, and the size
of the domain is given

```
N = (40, 40, 40)
family = 'Legendre'
K0 = FunctionSpace(N[0], 'Fourier', dtype='D', domain=(0, 2*np.pi))
K1 = FunctionSpace(N[1], 'Fourier', dtype='d', domain=(0, 2*np.pi))
SD = FunctionSpace(N[2], family, bc=(0, 0))
ST = FunctionSpace(N[2], family)
ST.slice = lambda: slice(0, ST.N-2)
```

Note that the last line of code is there to ensure that only the first
$N_2-2$ coefficients are used, see discussion around Eq. ([5](#eq:Zn)).
At the same time, we ensure that we are still using $N_2$
quadrature points, the same as for the Dirichlet basis.

Next the one-dimensional spaces are used to create two tensor product spaces Q = $W^{\mathbf{N}}$
and TD = $W_0^{\mathbf{N}}$, one vector V = $[W_0^{\mathbf{N}}]^3$ and one mixed
space  VQ = V $\times$ Q.

```
TD = TensorProductSpace(comm, (K0, K1, SD), axes=(2, 0, 1))
Q = TensorProductSpace(comm, (K0, K1, ST), axes=(2, 0, 1))
V = VectorSpace(TD)
VQ = CompositeSpace([V, Q])
```

Note that we choose to transform axes in the order $1, 0, 2$. This is to ensure
that the fully transformed arrays are aligned in the non-periodic direction 2.
And we need the arrays aligned in this direction, because this is the only
direction where there are tensor product matrices that are non-diagonal. All
Fourier matrices are, naturally, diagonal.

Test- and trialfunctions are created much like in a regular, non-mixed,
formulation. However, one has to create one test- and trialfunction for
the mixed space, and then split them up afterwards

```
up = TrialFunction(VQ)
vq = TestFunction(VQ)
u, p = up
v, q = vq
```

With the basisfunctions in place we may assemble the different blocks of the
final coefficient matrix. Since Legendre is using a constant weight function,
the equations may also be integrated by parts to obtain a symmetric system:

```
if family.lower() == 'chebyshev':
    A = inner(v, div(grad(u)))
    G = inner(v, -grad(p))
else:
    A = inner(grad(v), -grad(u))
    G = inner(div(v), p)
D = inner(q, div(u))
```

**Notice.**

The inner products may also be assembled with one single line, as

        AA = inner(v, div(grad(u))) + inner(v, -grad(u)) + inner(q, div(u))

However, this requires addition, not subtraction, of inner products and it is not
possible to move the negation to -inner(v, grad(u))

The assembled subsystems `A, G` and `D` are lists containg the different blocks of
the complete, coupled matrix. `A` actually contains 6
tensor product matrices of type `TPMatrix`. The first two
matrices are for vector component zero of the test function `v[0]` and
trial function `u[0]`, the
matrices 2 and 3 are for components 1 and the last two are for components
2. The first two matrices are as such for

          A[0:2] = inner(v[0], div(grad(u[0])))

Breaking it down the inner product is mathematically

$$
\begin{equation}
\label{eq:partialeq1} \tag{17}
\int_{\Omega} \mathbf{v}[0] \left(\frac{\partial^2 \mathbf{u}[0]}{\partial x^2} + \frac{\partial^2 \mathbf{u}[0]}{\partial y^2} + \frac{\partial^2 \mathbf{u}[0]}{\partial z^2}\right) w_x dx w_y dy w_z dz.
\end{equation}
$$

If we now use test function $\mathbf{v}[0]$

$$
\begin{equation}
\mathbf{v}[0]_{lmn} = \mathcal{X}_l \mathcal{Y}_m \mathcal{Z}_n,
\label{_auto15} \tag{18}
\end{equation}
$$

and trialfunction

$$
\begin{equation}
\mathbf{u}[0]_{pqr} = \sum_{p} \sum_{q} \sum_{r} \hat{\mathbf{u}}[0]_{pqr} \mathcal{X}_p \mathcal{Y}_q \mathcal{Z}_r,
\label{_auto16} \tag{19}
\end{equation}
$$

where $\hat{\mathbf{u}}$ are the unknown degrees of freedom, and then insert these functions
into ([17](#eq:partialeq1)), then we obtain after
performing some exact evaluations over the periodic directions

$$
\begin{equation}
 \Big( \underbrace{-\left(l^2 \delta_{lp} + m^2 \delta_{mq} \right) \int_{-1}^{1} \mathcal{Z}_r(z) \mathcal{Z}_n(z) w_z dz}_{A[0]} + \underbrace{\delta_{lp} \delta_{mq} \int_{-1}^{1} \frac{\partial^2 \mathcal{Z}_r(z)}{\partial z^2} \mathcal{Z}_n(z) w_z dz}_{A[1]} \Big) \hat{\mathbf{u}}[0]_{pqr},
\label{_auto17} \tag{20}
\end{equation}
$$

Similarly for components 1 and 2 of the test and trial vectors, leading to 6 tensor
product matrices in total for `A`. Similarly, we get three components of `G`
and  three of `D`.

Eliminating the Fourier diagonal matrices, we are left with block matrices like

$$
H(l, m) =
  \begin{bmatrix}
      A[0]+A[1] & 0 & 0 & G[0] \\ \label{eq:bmatrix} \tag{21}
      0 & A[2]+A[3] & 0 & G[1] \\
      0 & 0 &  A[4]+A[5] & G[2] \\
      D[0] & D[1] & D[2] & 0
  \end{bmatrix}
$$

Note that there will be one large block matrix $H(l, m)$ for each Fourier
wavenumber combination $(l, m)$. To solve the problem in the end we will need to
loop over these wavenumbers and solve the assembled linear systems one by one.
An example of the block matrix, for $l=m=5$ and $\mathbf{N}=(20, 20, 20)$ is given
in Fig. [fig:BlockMat](#fig:BlockMat).

In the end we create a block matrix through

```
M = BlockMatrix(A+G+D)
```

The right hand side can easily be assembled since we have already
defined the functions $\mathbf{f}$ and $h$, see Sec. [Manufactured solution](#sec:mansol)

```
# Get mesh (quadrature points)
X = TD.local_mesh(True)

# Get f and h on quad points
fh = Array(VQ, buffer=(fx, fy, fz, h))
f_, h_ = fh

# Compute inner products
fh_hat = Function(VQ)
f_hat, h_hat = fh_hat
f_hat = inner(v, f_, output_array=f_hat)
h_hat = inner(q, h_, output_array=h_hat)
```

In the end all that is left is to solve and compare with
the exact solution.

```
# Solve problem
up_hat = M.solve(fh_hat, constraints=((3, 0, 0),))
up = up_hat.backward()
u_, p_ = up

# Exact solution
ux, uy, uz = Array(V, buffer=(uex, uey, uez))
pe = Array(Q, buffer=pe)

error = [comm.reduce(np.linalg.norm(ux-u_[0])),
         comm.reduce(np.linalg.norm(uy-u_[1])),
         comm.reduce(np.linalg.norm(uz-u_[2])),
         comm.reduce(np.linalg.norm(pe-p_))]
print(error)
```

Note that solve has a keyword argument
`constraints=((3, 0, 0),)` that takes care of the restriction
$\int_{\Omega} p dx = 0$ by indenting the row in M corresponding to the
first degree of freedom for the pressure. The value $(3, 0, 0)$
indicates that pressure is
in block 3 of the block vector solution (the velocity vector holds
positions 0, 1 and 2), whereas the two zeros ensures that the first dof
(dof 0) should obtain value 0.

## Complete solver

A complete solver can be found in demo [Stokes3D.py](https://github.com/spectralDNS/shenfun/blob/master/demo/Stokes3D.py).


# Spots: Short Timescale Approximation

Simple relations for changes in stellar fundamental parameters resulting from the presence of spots under the short timescale approximation. Influence on the stellar bolometric luminosity relative to that of an unspotted star is
\begin{equation}
    \zeta = 1 - \varrho (1 - \varpi^4).
\end{equation}
This leads to the relation between short timescale spots and changes in the average photospheric temperature of a star,
\begin{equation}
    \frac{T_{\rm avg}}{T_{\rm eff}} = \zeta^{1/4}.
\end{equation}
We can now explore how luminosity variations change as a function of $\varrho$ and $\varpi$.

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

```python
# initialize some spot properties
varrho = np.arange(0.0, 1.0, 0.01)
varpi  = np.arange(0.0, 1.0, 0.01)
zetas  = np.arange(0.1, 1.0, 0.1)
```

First we can look at combinations of spot areal coverage and spot temperature contrast.

```python
def arealCoverage(zeta, varpis):
    return (1. - zeta)/(1. - varpis**4)
```

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

for zeta in zetas:
    ax.plot(varpi, arealCoverage(zeta, varpi), '-', lw=3, c='#0094b2')

ax.grid(True)
ax.set_xlim(0.0, 1.0)
ax.set_ylim(0.0, 1.0)
ax.set_xlabel('Temperature Contrast $\\varpi$', fontsize=16., family='serif')
ax.set_ylabel('Areal Coverage $\\varrho$', fontsize=16., family='serif')
ax.tick_params(which='major', axis='both', length=15., labelsize=16.)
```

Moving on to the average surface temperature,

```python
zetas  = np.arange(0.0, 1.0, 0.01)

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

ax.plot(zetas, zetas**0.25, '-', lw=3, c='#0094b2')

ax.grid(True)
ax.set_xlim(0.0, 1.0)
ax.set_ylim(0.0, 1.0)
ax.set_xlabel('Luminosity Ratio $\\zeta$', fontsize=16., family='serif')
ax.set_ylabel('Average Surface Temperature Ratio', fontsize=16., family='serif')
ax.tick_params(which='major', axis='both', length=15., labelsize=16.)
```

```python
zetas = np.arange(0.1, 1.0, 0.1)
def arealCoverage(zeta, varpis):
    return (1. - zeta**0.25)/(1. - varpis**4)
```

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

for zeta in zetas:
    ax.plot(varpi, arealCoverage(zeta, varpi), '-', lw=3, c='#0094b2')

ax.grid(True)
ax.set_xlim(0.0, 1.0)
ax.set_ylim(0.0, 1.0)
ax.set_xlabel('Temperature Contrast $\\varpi$', fontsize=16., family='serif')
ax.set_ylabel('Areal Coverage $\\varrho$', fontsize=16., family='serif')
ax.tick_params(which='major', axis='both', length=15., labelsize=16.)
```

For a range of stellar effective temperatures, this results in an average surface temperature anomaly (in K).

```python
Teffs = np.arange(3000., 7000.)

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

for zeta in zetas:
    ax.plot(Teffs, Teffs*(zeta**0.25 - 1.), '-', lw=3, c='#0094b2')

ax.grid(True)
#ax.set_xlim(0.0, 1.0)
#ax.set_ylim(0.0, 1.0)
ax.set_xlabel('Original Effective Temperature (K)', fontsize=16., family='serif')
ax.set_ylabel('Mean Temperature Anomaly (K)', fontsize=16., family='serif')
ax.tick_params(which='major', axis='both', length=15., labelsize=16.)
```

```python

```


## VQE - Variational Quantum Eigensolver

- What is Optimization?

   Search for maxima or minima. (Machine Learning: from data to pattern recognition)

------------------

- What is Simulation?

  A sampling technique. (Physics and Engineering : from pattern(mathematical equation) to data generation))

### VQE problem statement:

$~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~$ Given a Hamiltonian H, solve the following optimization problem

$$
\min_{\theta}~~~ \langle \psi(\theta) |H| \psi (\theta) \rangle
$$

### Simple Example

```python
from qiskit import *
```

* Lets us consider very simple $|\psi\rangle $ as probability vector (i.e. H is I matrix)

```python
import numpy as np
np.random.seed(999999)
target_distr = np.random.rand(2)
# We now convert the random vector into a valid probability vector
target_distr /= sum(target_distr)
```

* Variational Form can be prepared with three paramentrs $\theta, \phi, \lambda $

$$
\begin{align}
    U3(\theta, \phi, \lambda) = \begin{pmatrix}\cos(\frac{\theta}{2}) & -e^{i\lambda}\sin(\frac{\theta}{2}) \\ e^{i\phi}\sin(\frac{\theta}{2}) & e^{i\lambda + i\phi}\cos(\frac{\theta}{2}) \end{pmatrix}
\end{align}
$$

```python
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister
def get_var_form(params):
    qr = QuantumRegister(1, name="q")
    cr = ClassicalRegister(1, name='c')
    qc = QuantumCircuit(qr, cr)
    qc.u3(params[0], params[1], params[2], qr[0])
    qc.measure(qr, cr[0])
    return qc
```

* Classical Part: Cost function (Objective function)

```python
from qiskit import Aer
from qiskit.compiler import transpile,assemble
backend = Aer.get_backend("qasm_simulator")
NUM_SHOTS = 10000

def get_probability_distribution(counts):
    output_distr = [v / NUM_SHOTS for v in counts.values()]
    if len(output_distr) == 1:
        output_distr.append(1 - output_distr[0])
    return output_distr

def objective_function(params):
    # Obtain a quantum circuit instance from the paramters
    qc = get_var_form(params)
    # Execute the quantum circuit to obtain the probability distribution associated with the current parameters
    t_qc = transpile(qc, backend)
    qobj = assemble(t_qc, shots=NUM_SHOTS)
    result = backend.run(qobj).result()
    # Obtain the counts for each measured state, and convert those counts into a probability vector
    output_distr = get_probability_distribution(result.get_counts(qc))
    # Calculate the cost as the distance between the output distribution and the target distribution
    cost = sum([np.abs(output_distr[i] - target_distr[i]) for i in range(2)])
    return cost
```

* Classical Part: Optimization

```python
from qiskit.aqua.components.optimizers import COBYLA

# Initialize the COBYLA optimizer
optimizer = COBYLA(maxiter=500, tol=0.0001)

# Create the initial parameters (noting that our single qubit variational form has 3 parameters)
params = np.random.rand(3)
ret = optimizer.optimize(num_vars=3, objective_function=objective_function, initial_point=params)

# Obtain the output distribution using the final parameters
qc = get_var_form(ret[0])
t_qc = transpile(qc, backend)
qobj = assemble(t_qc, shots=NUM_SHOTS)
counts = backend.run(qobj).result().get_counts(qc)
output_distr = get_probability_distribution(counts)

print("Target Distribution:", target_distr)
print("Obtained Distribution:", output_distr)
print("Output Error (Manhattan Distance):", ret[1])
print("Parameters Found:", ret[0])
```

    Target Distribution: [0.51357006 0.48642994]
    Obtained Distribution: [0.5248, 0.4752]
    Output Error (Manhattan Distance): 0.030059881261160826
    Parameters Found: [1.53560768 1.2072113  0.63106461]

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```


# The shallow water equations

In this chapter we study a model for shallow water waves in one dimension:

\begin{align}
    h_t + (hu)_x & = 0 \label{SW_mass} \\
    (hu)_t + \left(hu^2 + \frac{1}{2}gh^2\right)_x & = 0. \label{SW_mom}
\end{align}

Here $h$ is the depth, $u$ is the velocity, and $g$ is a constant representing the force of gravity.  These equations are a relatively simple but surprisingly effective model for water waves; they are based on several important physical assumptions.  For instance, viscosity and compressibility are neglected, and the pressure is assumed to be hydrostatic.  Nevertheless, this is a surprisingly effective model for many applications.

Previously we have looked at a *linear hyperbolic system* (acoustics) and *scalar hyperbolic equations* (advection, Burgers' equation, traffic flow).  The shallow water system is our first example of a *nonlinear hyperbolic system*; solutions of the Riemann problem will consist of multiple waves, each of which may be a shock or rarefaction.

## Hyperbolic structure
We can write (\ref{SW_mass})-(\ref{SW_mom}) in the canonical form $q_t + f(q)_x = 0$ if we define

\begin{align}
q & = \begin{pmatrix} h \\ hu \end{pmatrix} & f & = \begin{pmatrix} hu \\ hu^2 + \frac{1}{2}gh^2 \end{pmatrix}.
\end{align}

In terms of the conserved quantities, the flux is

\begin{align}
f(q) & = \begin{pmatrix} q_2 \\ q_2^2/q_1 + \frac{1}{2}g q_1^2 \end{pmatrix}
\end{align}

Thus the flux jacobian is

\begin{align}
f'(q) & = \begin{pmatrix} 0 & 1 \\ -(q_2/q_1)^2 + g q_1 & 2 q_2/q_1 \end{pmatrix}
        = \begin{pmatrix} 0 & 1 \\ -u^2 + g h & 2 u \end{pmatrix}.
\end{align}

Its eigenvalues are

\begin{align} \label{SW:char-speeds}
    \lambda_1 & = u - \sqrt{gh} & \lambda_2 & = u + \sqrt{gh},
\end{align}

with corresponding eigenvectors

\begin{align} \label{SW:fjac-evecs}
    r_1 & = \begin{bmatrix} 1 \\ u-\sqrt{gh} \end{bmatrix} &
    r_2 & = \begin{bmatrix} 1 \\ u+\sqrt{gh} \end{bmatrix}
\end{align}

Notice that -- unlike for acoustics, but similar to the LWR traffic model -- the eigenvalues and eigenvectors depend on $q$.  Because of this, the waves appearing in the Riemann problem move at different speeds and may be either jump discontinuities (shocks) or smoothly-varying regions (rarefactions).

```python
%matplotlib inline
%config InlineBackend.figure_format = 'svg'
import matplotlib as mpl
mpl.rcParams['font.size'] = 8
figsize =(8,4)
mpl.rcParams['figure.figsize'] = figsize

import numpy as np
import matplotlib.pyplot as plt
from exact_solvers import shallow_water
from collections import namedtuple
from utils import riemann_tools
from ipywidgets import interact
from ipywidgets import widgets, FloatSlider, Checkbox, RadioButtons
State = namedtuple('State', shallow_water.conserved_variables)
Primitive_State = namedtuple('PrimState', shallow_water.primitive_variables)
g = 1.
```

Here we define a function for each wave speed, for use later.

```python
def lambda_1(q, xi, g=1.):
    "Characteristic speed for shallow water 1-waves."
    h, hu = q
    if h > 0:
        u = hu/h
        return u - np.sqrt(g*h)
    else:
        return 0

def lambda_2(q, xi, g=1.):
    "Characteristic speed for shallow water 2-waves."
    h, hu = q
    if h > 0:
        u = hu/h
        return u + np.sqrt(g*h)
    else:
        return 0
```

## The Riemann problem
Consider the Riemann problem with left and right states

\begin{align*}
q_l & = \begin{bmatrix} h_l \\ h_r u_l \end{bmatrix}, &
q_r & = \begin{bmatrix} h_r \\ h_r u_r \end{bmatrix}.
\end{align*}

Typically the Riemann solution will consist of two waves, one related to each of the eigenvectors in \eqref{SW:fjac-evecs}.  We refer to the wave related to $r_1$ (the first characteristic family) as a 1-wave and the wave related to $r_2$ (the second characteristic family) as a 2-wave.  Each wave may be a shock or rarefaction.  There will be an intermediate state $q_m$ between them.  In the figure below, we illustrate a typical situation in which the 1-wave happens to be a rarefaction and the 2-wave is a shock:

In the figure we have one wave going in each direction, but since the wave speeds depend on $q$ and can have either sign, it is possible to have both waves going left, or both going right.

To solve the Riemann problem, we must find $q_m$.  To do so we must find a state that can be connected to $q_l$ by a 1-shock or 1-rarefaction and to $q_r$ by a 2-shock or 2-rarefaction.  We must also ensure that the resulting waves satisfy the entropy condition.

## Shock waves
Let us now consider a shock wave connecting a known state $q_*=(h_*, h_* u_*)$ to an unknown state $q=(h,hu)$.  In the context of the Riemann problem, $q_*$ will be the initial left or right state and $q$ will be the middle state.
The Rankine-Hugoniot jump conditions $s(q_r - q_l) = f(q_r) - f(q_l)$ for a shock wave moving
at speed $s$ read

\begin{align}
s (h_* - h) & = h_* u_* - h u \\
s (h_* u_* - h u) & = h_* u_*^2 - h u^2 + \frac{g}{2}(h_*^2 - h^2).
\end{align}

It is convenient to define $\alpha = h - h_*$.  Then we can combine the two equations above to show that the momenta are related by

\begin{align} \label{SW:hugoniot-locus}
    h u & = h_* u_* + \alpha \left[u_* \pm \sqrt{gh_* \left(1+\frac{\alpha}{h_*}\right)\left(1+\frac{\alpha}{2h_*}\right)}\right]
\end{align}

for a shock moving with speed

\begin{align} \label{SW:shock-speed}
    s & = \frac{h_* u_* - h u}{h_* - h} = u_* \mp \sqrt{gh_* \left(1+\frac{\alpha}{h_*}\right)\left(1+\frac{\alpha}{2h_*}\right)} = u_* \mp \sqrt{gh \frac{h_* + h}{2h_*} }
\end{align}

Choosing the plus sign in \eqref{SW:hugoniot-locus} (which yields the minus sign in \eqref{SW:shock-speed}) gives a 1-shock.  Choosing the minus sign in \eqref{SW:hugoniot-locus} (which yields the plus sign in \eqref{SW:shock-speed}) gives a 2-shock.  Notice from the last expression in \eqref{SW:shock-speed} that as $h \to h^*$ the shock speed approaches the corresponding characteristic speed \eqref{SW:char-speeds}, while the ratio of the jumps in momentum and depth approaches that suggested by the eigenvectors \eqref{SW:fjac-evecs}.

We can now consider the set of all possible states $(h, h u)$ satisfying \eqref{SW:hugoniot-locus}.  This curve in the $h - hu$ plane is known as the hugoniot locus; there is a family of hugoniot loci for 1-shocks and another for 2-shocks.  Below we plot some members of these two families of curves.

```python
interact(shallow_water.plot_hugoniot_loci,
         y_axis=widgets.fixed('hu'));
```

The plot above shows the Hugoniot loci in the $h$-$hu$ plane, the natural phase plane in terms of the two conserved quantities.  Note that they all approach the origin as $h \rightarrow 0$. Alternatively, we can plot these same curves in the $h$-$u$ plane:

```python
interact(shallow_water.plot_hugoniot_loci,
         y_axis=widgets.fixed('u'));
```

Note that in this plane, the curves in each family are simply vertical translates of one another, and all curves asymptote to $\pm \infty$ as $h \rightarrow 0$.

### All-shock Riemann solutions
If we knew *a priori* that both waves in the Riemann solution were shock waves, we could find $q_m$ by finding the intersection of the 1-Hugoniot locus passing through $q_l$ with the 2-Hugoniot locus passing through $q_r$.  The Riemann solution would then consist simply of two discontinuities propagating at the speeds given by \eqref{SW:shock-speed}.

Here is an example for which this is the correct solution.

```python
def make_plot_function(q_l,q_r,g=1.,force_waves=None,extra_lines=None):
    states, speeds, reval, wave_types = \
        shallow_water.exact_riemann_solution(q_l,q_r,g,
                      force_waves=force_waves)
    def plot_function(t,plot_1_chars=False,plot_2_chars=False):
        ax = riemann_tools.plot_riemann(states,speeds,reval,wave_types,
                           t=t,t_pointer=0,extra_axes=True,
                           variable_names=shallow_water.conserved_variables,
                           fill=[0]);
        if plot_1_chars:
            riemann_tools.plot_characteristics(reval,lambda_1,axes=ax[0],
                                               extra_lines=extra_lines)
        if plot_2_chars:
            riemann_tools.plot_characteristics(reval,lambda_2,axes=ax[0],
                                               extra_lines=extra_lines)
        shallow_water.phase_plane_plot(q_l,q_r,g,ax=ax[3],
                                       force_waves=force_waves,y_axis='u')
        plt.show()
    return plot_function

def plot_riemann_SW(q_l,q_r,g=1.,force_waves=None,extra_lines=None):
    plot_function = make_plot_function(q_l,q_r,g,
                                       force_waves,extra_lines)
    interact(plot_function,
             t=widgets.FloatSlider(value=0.,min=0,max=.9),
             plot_1_chars=Checkbox(description='1-characteristics',
                                   value=False),
             plot_2_chars=Checkbox(description='2-characteristics'))
```

```python
q_l  = State(Depth = 1.,
             Momentum = 1.)
q_r = State(Depth = 1.,
            Momentum = -1.)

plot_riemann_SW(q_l,q_r)
```

*Check the boxes to plot the characteristic families, and notice that the 1-characteristics impinge on the 1-shock; the 2-characteristics impinge on the 2-shock.  Thus these shocks satisfy the entropy condition.*

#### Dam-break problem: all-shock solution
Here's another example, which you might think of as modeling a dam that suddenly breaks.  The water is initially at rest on both sides, but much deeper on the left.

```python
q_l  = State(Depth = 4.,
             Momentum = 0.)
q_r = State(Depth = 1.,
            Momentum = 0.)
plot_riemann_SW(q_l,q_r,force_waves='shock',
                extra_lines=[[[0,0],[-1.317,1.]]]);
```

*Check the box to plot the 1-characteristics.*  Notice that they don't impinge on the 1-shock; instead, characteristics are diverging away from it.  This shock does not satisfy the entropy condition and should be replaced with a rarefaction.  The corresponding part of the Hugoniot locus is plotted with a dashed line to remind us that it is unphysical.

*Now check the box to plot the 2-characteristics.*  Notice that these *do* impinge on the 2-shock; this is an entropy-satisfying shock, and the Hugoniot locus is plotted as a solid line.

### The entropy condition
To be physically correct, each of these waves must satisfy the Lax entropy condition.  Notice that the shock speed \eqref{SW:shock-speed} can be written as

\begin{align}
    s & = u_* \pm \sqrt{gh \frac{h_* + h}{2h_*} }
\end{align}
Thus the 1-shock entropy condition reads

\begin{align*}
u_l - \sqrt{gh_l} > u_l - \sqrt{gh_l \frac{h_l+h_m}{2h_l}} = u_m - \sqrt{gh_m \frac{h_m+h_l}{2h_m}} > u_m - \sqrt{gh_m}.
\end{align*}
The second expression for the shock speed is obtained by noticing
that the shock speed is invariant under transposition of the left and right states.
The first and second inequalities both simplify to

\begin{align} \label{SW:1-entropy}
    h_m > h_l.
\end{align}
Similarly, it can be shown that the entropy condition for a 2-shock connecting $q_m$ and $q_r$ is satisfied if and only if

\begin{align} \label{SW:2-entropy}
    h_m > h_r.
\end{align}
If \eqref{SW:1-entropy} or \eqref{SW:2-entropy} is violated, then the
corresponding wave must be a rarefaction rather than a shock.

Go back and look at the phase plane plots above.  You'll notice that the 1-Hugoniot locus passing through $q_\ell$ is plotted with a solid line for $h>h_\ell$ and with a dashed line for $h<h_\ell$ (the same is done for the 2-Hugoniot locus passing through $q_r$).  This is to remind us that states connected to $q_\ell$ by the dashed part are unphysical.

## Rarefaction waves

### Integral curves
In the LWR traffic flow model, we saw that a rarefaction wave
consisted of a smoothly-varying density.  For the shallow water
system, both the depth $h$ and the momentum $hu$ will vary smoothly
through a rarefaction wave.  A given rarefaction wave is associated
with just one characteristic family, and the variation within the
wave is at all points proportional to the corresponding eigenvector
$r_p$.  In other words, all values $\tilde{q}=(h,hu)$ in the rarefaction
lie on a curve defined by

\begin{align} \label{SW:gen_ode}
    \tilde{q}'(\xi) & = r_p(\tilde{q}(\xi)).
\end{align}

Here $\tilde{q}(\xi)$ is a parameterization of the solution.
For the shallow water system, \eqref{SW:gen_ode} reads

\begin{align}
    h'(\xi)    = q_1'(\xi) & = 1 \label{SW:dh} \\
    (hu)'(\xi) = q_2'(\xi) & = u \pm \sqrt{gh} = \tilde{q}_2/\tilde{q}_1 - \sqrt{g\tilde{q}_1}, \label{SW:dhu}
\end{align}
where we take the minus sign for 1-waves and the plus sign for 2-waves.
We can think of \eqref{SW:dh}-\eqref{SW:dhu} as an initial value ODE;
fixing the value of $\tilde{q}$ at one point in the rarefaction wave
determines the whole solution of \eqref{SW:dh}-\eqref{SW:dhu}.  We
refer to each of these solutions as an *integral curve*.

For the shallow water system, these equations can be integrated exactly.  If we fix one point $(h_*, h_*u_*)$ on the curve, the whole integral curve is defined by

\begin{align}
    hu & = hu_* \pm 2h(\sqrt{gh_*} - \sqrt{gh}).
\end{align}
where now the plus sign is for 1-waves and the minus sign for 2-waves.
This can be rewritten as

\begin{align} \label{SW:riemann-invariant}
u \pm 2 \sqrt{gh} = u_* \pm 2 \sqrt{gh_*}.
\end{align}
From \eqref{SW:riemann-invariant} we see that the value $u + 2 \sqrt{gh}$
is constant along any 1-integral curve, while the value $u-2\sqrt{gh}$ is constant along any 2-integral curve.  We refer to these quantities as *Riemann invariants* for the shallow water system:

\begin{align}
w_1(q) & = u + 2 \sqrt{gh} \\
w_2(q) & = u - 2 \sqrt{gh}.
\end{align}
In other words, the trajectories plotted above are just the level curves of $w_1$ and $w_2$.  This allows us to easily plot the integral curves as a contour plot of $w_1$ or $w_2$:

```python
def plot_int_curves(plot_1=True,plot_2=False,y_axis='hu'):
    N = 400
    g = 1.
    if y_axis == 'hu':
        h, hu = np.meshgrid(np.linspace(0.01,3,N),
                            np.linspace(-3,3,N))
        u = hu/h
        yvar = hu
        ylabel = 'momentum (hu)'
    else:
        h, u = np.meshgrid(np.linspace(0.01,3,N),
                           np.linspace(-3,3,N))
        yvar = u
        ylabel = 'velocity (u)'
    w1 = u + 2 * np.sqrt(g*h)
    w2 = u - 2 * np.sqrt(g*h)
    if plot_1:
        clines = np.linspace(-4,6,21)
        plt.contour(h,yvar,w1,clines,colors='coral',
                    linestyles='solid')
    if plot_2:
        clines = np.linspace(-6,4,21)
        plt.contour(h,yvar,w2,clines,colors='maroon',
                    linestyles='solid')

    plt.axis((0,3,-3,3))
    plt.xlabel('depth h')
    plt.ylabel(ylabel)
    plt.title('Integral curves')
    plt.show()

interact(plot_int_curves, y_axis=widgets.fixed('hu'),
         plot_1=widgets.Checkbox(description='1-wave curves',
                                 value=True),
         plot_2=widgets.Checkbox(description='2-wave curves'));
```

We can also plot the integral curves in the $h$--$u$ plane:

```python
interact(plot_int_curves, y_axis=widgets.fixed('u'),
         plot_1=widgets.Checkbox(description='1-wave curves',
                                 value=True),
         plot_2=widgets.Checkbox(description='2-wave curves'));
```

Note that in this plane the integral curves of each family are simply vertical translates of one another due to the form of the functions $w_1$ and $w_2$.  Unlike the Hugoniot loci, the integral curves do not asymptote to $\pm\infty$ as $h \rightarrow 0$ and instead approach a finite value.

### Comparison of integral curves and Hugoniot loci
You may have noticed that the integral curves look very similar to the Hugoniot loci, especially when plotted in the $h-hu$ plane.  Below we emphasize the differences by plotting the curve from each of these families that passes through a specified point.  The difference is less noticeable in the $h-hu$ plane since all of the curves approach the origin.

```python
def compare_curves(wave_family=1, h0=1., u0=0., y_axis='u'):
    h = np.linspace(1.e-2,3,400)

    u = shallow_water.integral_curve(h,h0,h0*u0,
                                     wave_family,y_axis=y_axis)
    plt.plot(h,u,'b')
    u = shallow_water.hugoniot_locus(h,h0,h0*u0,
                                     wave_family,y_axis=y_axis)
    plt.plot(h,u,'r')
    plt.legend(['Integral curve','Hugoniot locus'])
    if y_axis == 'u':
        plt.plot(h0,u0,'ok')
        plt.ylabel('velocity (u)')
    else:
        plt.plot(h0,h0*u0,'ok')
        plt.ylabel('momentum (hu)')
    plt.xlabel('depth h')
    plt.show()

interact(compare_curves,
         wave_family=RadioButtons(options=[1,2],
                                  description='Wave family:'),
         y_axis=RadioButtons(options=['u','hu'],
                             description='Vertical axis:'),
         h0=FloatSlider(min=1.e-1,max=3.,value=1.,
                        description='$h_*$'),
         u0=FloatSlider(min=-3,max=3,description='$u_*$'));
```

Near the point of intersection, the curves are very close; indeed, they must be tangent at this point since their direction is parallel to the corresponding eigenvector there.  Far from this point they diverge; for small depths they must diverge greatly, since the Hugoniot locus never reaches $h=0$ at any finite velocity.

### The structure of centered rarefaction waves
Rarefactions appearing in the Riemann solution are also similarity solutions; that is they are constant along any ray $\xi=x/t$.  For a $p$-rarefaction connecting states $q_\ell$ and $q_r$, the rarefaction thus has the form

\begin{align}
    q(x,t) & = \begin{cases} q_\ell & \text{if } x/t \le \lambda_p(q_\ell) \\
                            \tilde{q}(x/t) & \text{if } \lambda_p(q_\ell) \le \lambda_p(q_r) \\
                            q_r & \text{if } x/t \ge \lambda_p(q_r). \end{cases}
\end{align}

It remains to be determined how $\tilde{q}$ varies as a function of $x/t$.  This can be determined by noting that each solution value $\tilde{q}$ propagates at the characteristic speed $\lambda_p(\tilde{q})$.  Letting $\xi=x/t$, this gives us the condition

\begin{align} \label{SW:char-sim}
    \lambda_p(\tilde{q}(\xi)) = \xi.
\end{align}

By combining \eqref{SW:char-sim} with the Riemann invariant condition, we can find $\tilde{h}$ and $\tilde{u}$ within the rarefaction.  For concreteness, let us consider a 1-rarefaction connecting $q_l$ and $q_m$.  Then \eqref{SW:char-sim} reads $\tilde{u} - \sqrt{g\tilde{h}} = \xi$, which simplifies to

\begin{align} \label{SW:char-sim1}
\tilde{h} = \frac{(\tilde{u}-\xi)^2}{g}.
\end{align}
Meanwhile, the fact that $w_1(\tilde{q})$ is constant means that

\begin{align} \label{SW:w1-const}
    \tilde{u} + 2 \sqrt{g\tilde{h}} = u_\ell + 2 \sqrt{gh_\ell}.
\end{align}
Combining \eqref{SW:char-sim1} with \eqref{SW:w1-const}, we find

\begin{align} \label{SW:1-raref-struct-h}
    \tilde{h} & = \frac{(u_\ell + 2\sqrt{gh_\ell} - \xi)^2}{9g} = \frac{(w_1(q_l) - \xi)^2}{9g} \\
    \tilde{u} & = \frac{1}{3}w_1(q_l) + \frac{2}{3}\xi. \label{SW:1-raref-struct-u}
\end{align}
We see that $h$ varies quadratically through the rarefaction, while $u$ varies linearly.

### Connecting states
Suppose we know that the middle state $q_m$ is connected to the left and right states by rarefactions.  Then we can find $q_m$ by finding the intersection of the corresponding integral curves.  Here is an example for which this is the physically correct solution:

\begin{align}
    q_l & = \begin{bmatrix} 1 \\ -1 \end{bmatrix} &
    q_r & = \begin{bmatrix} 1 \\ 1 \end{bmatrix}
\end{align}

```python
q_l  = State(Depth = 1.,
             Momentum = -1.)
q_r = State(Depth = 1.,
            Momentum = 1.)
plot_riemann_SW(q_l,q_r)
```

Notice that the segment of each integral curve that connects to states with a smaller depth is plotted as a solid line, while the segment connecting to states with greater depth is plotted with a dashed line.  This again is to remind us that states connected by a rarefaction through the solid part are physical (entropy-satisfying), while states connected by the dashed part would be unphysical (entropy-violating).

*Try changing the velocities to $(2, -2)$, and notice that the Riemann solution has a state with almost zero depth in the middle.  What happens if you use even larger outflow velocities?  We'll discuss this later, in the section on dry states.*

### Dam-break problem: all-rarefaction solution
Next let us revisit the problem from [the section above](Shallow_water.ipynb#Dam-break-problem:-all-shock-solution) where we found that the all-shock solution was incorrect.  This time we'll try to find an all-rarefaction solution.

```python
q_l  = State(Depth = 4.,
             Momentum = 0.)
q_r = State(Depth = 1.,
            Momentum = 0.)
plot_riemann_SW(q_l,q_r,force_waves='raref')
```

Plot the 2-characteristics and notice that they cross each other within the 2-rarefaction. This means that the solution we constructed is triple-valued and nonsensical as a solution to this one-dimensional conservation law, and so this portion of the solution is omitted in the plots of depth and momentum.  In this case a rarefaction wave is not physical and should be replaced with a shock; the corresponding part of the integral curve is hence shown as a dashed line.

## The Riemann solution
To determine $q_m$, we need to know whether each of the two waves is a shock or rarefaction, so that we can use the appopriate Hugoniot locus or integral curve.  But to determine the nature of each wave, we need to check \eqref{SW:1-entropy} or \eqref{SW:2-entropy}, which requires knowledge of $h_m$.  To deal with this we define a piecewise function $\phi_l(h)$ that gives:

- for $h>h_l$, the momentum connected to $q_l$ by a 1-shock;
- for $h<h_l$, the momentum connected to $q_l$ by a 1-integral curve.

Thus $\phi$ is partly a 1-Hugoniot locus and partly a 1-integral curve; it gives the value of each precisely in the regime where they are physically correct.  We can similarly define a function $\phi_r(h)$ related to the 2-Hugoniot locus and 2-integral curve.  The middle state depth $h_m$ is the value for which $\phi_l(h) = \phi_r(h)$.  This value can be found computationally via a rootfinding algorithm.

### Dam-break problem: correct solution
Below we plot (at last) the correct solution to the dam-break problem, involving a 1-rarefaction and a 2-shock.  Plot the characteristics and notice how each family behaves correctly with respect to the corresponding wave in the Riemann solution.

```python
q_l  = State(Depth = 4.,
             Momentum = 0.)
q_r = State(Depth = 1.,
            Momentum = 0.)

plot_riemann_SW(q_l,q_r)
```

[This webpage](phase_plane_shallow_water_small.html) provides an interactive view that allows you to see how the solution changes as you move the left and right states around in the phase plane.  Note how the waves switch between shocks and rarefactions depending on the relative position of the left and right states.

## Dry states
The approach just described will yield the correct Riemann solution (provided that the rootfinding algorithm converges) in most situations.  However, it will fail if the solution involves dry ($h=0$) states.

### Dry middle state
Sometimes it seems that there is no way to connect a given pair of left and right states via entropy-satisfying shocks or rarefactions.  In the cell below we plot only the physically relevant parts of the integral curves and hugoniot loci passing through the left and right states.  Both the $h-u$ and $h-hu$ planes are shown.  *Try to adjust those states to find a combination for which none of the curves intersect.*

```python
def connect_states(h_l=1.,u_l=-1.,h_r=1.,u_r=1.):
    q_l = np.array([h_l,h_l*u_l])
    q_r = np.array([h_r,h_r*u_r])
    fig, ax = plt.subplots(1,2,figsize=figsize)
    shallow_water.phase_plane_curves(q_l[0], q_l[1], 'qleft',
                  wave_family=1,ax=ax[0])
    shallow_water.phase_plane_curves(q_r[0], q_r[1], 'qright',
                  wave_family=2,ax=ax[0])
    shallow_water.phase_plane_curves(q_l[0], q_l[1], 'qleft',
                  wave_family=1,y_axis='hu',ax=ax[1])
    shallow_water.phase_plane_curves(q_r[0], q_r[1], 'qright',
                  wave_family=2,y_axis='hu',ax=ax[1])
    ax[0].set_title('h-u plane'); ax[1].set_title('h-hu plane');
    ax[0].set_xlim(0,3); ax[1].set_xlim(0,3);
    ax[0].set_ylim(-10,10); ax[1].set_ylim(-10,10);
    plt.tight_layout(); plt.show()

interact(connect_states,
         h_l=widgets.FloatSlider(min=0.001,max=2,value=1),
         u_l=widgets.FloatSlider(min=-5,max=5,value=-1),
         h_r=widgets.FloatSlider(min=0.001,max=2,value=1),
         u_r=widgets.FloatSlider(min=-5,max=5,value=1));
```

You should find that by making the initial states flow sufficiently fast away from each other, there is no intersection in the $h-u$ plane.  In the $h-hu$ plane, the integral curves always intersect at the origin.  The reason for this ambiguity is that for zero depth it isn't meaningful to assign a velocity.  Thus in the $h-u$ plane we could think of the entire $u$-axis as being part of every integral curve.  That means that we can always connect the left and right states via an intermediate state with depth $h=0$ (a dry state).

Since the 1-integral curve through $q_l$ reaches $h=0$ at $u = u_l + 2 \sqrt{gh_l}$ and the 2-integral curve reaches $h=0$ at $u=u_r-2\sqrt{gh_r}$, we see that a dry middle state occurs if and only if $u_l + 2 \sqrt{gh_l} < u_r-2\sqrt{gh_r}$.

It is clear in this case that the solution must consist of two centered rarefaction waves, connecting the left and right states to the middle dry state.  These centered rarefactions still have the structure derived above; to complete the solution we only need to know the range of centered characteristics included in each rarefaction.  This is most conveniently determined by considering the $h-u$ plane.  The 1-integral curve passing through $q_l$ is given by $u = u_l + 2(\sqrt{gh_l}-\sqrt{gh_r})$; it reaches a dry state $h=0$ at $u=u_l + 2\sqrt{gh_l}$.  Similarly, the 2-integral curve passing through $q_r$ reaches a dry state $h=0$ at $u=u_r - 2\sqrt{gh_r}$.  Thus the dry state solution is given by

\begin{align}
    q(x,t) & = \begin{cases}
                        q_l &         \text{if } x/t \le u_l - \sqrt{gh_l} \\
                        \tilde{q}_\ell & \text{if } u_l -  \sqrt{gh_\ell} \le x/t \le u_l+2\sqrt{gh_l} \\
                        0 &           \text{if } u_l + 2\sqrt{gh_l} \le x/t \le u_r - 2\sqrt{gh_r} \\
                        \tilde{q}_r & \text{if } u_r - 2\sqrt{gh_r} \le x/t \le u_r+\sqrt{gh_r} \\
                        q_r &         \text{if } x/t \ge u_r + \sqrt{gh_r},
                       \end{cases}
\end{align}
where $\tilde{q}_\ell$ is given by \eqref{SW:1-raref-struct-h}-\eqref{SW:1-raref-struct-u} and $\tilde{q}_r$ is given by a similar expression.

Here's an example of a Riemann problem whose solution involves a dry middle state:

```python
q_l  = State(Depth = 1.,
             Momentum = -3.)
q_r = State(Depth = 1.,
            Momentum = 4.)
plot_riemann_SW(q_l,q_r)
```

### Dry initial states
What if the left or right state in the Riemann problem is dry?  In this case, the solution consists only of a rarefaction wave connecting the non-dry state to $h=0$.

```python
q_l  = State(Depth = 1.,
             Momentum = 0.)
q_r = State(Depth = 0.,
            Momentum = 0.)
plot_riemann_SW(q_l,q_r)
```

## Interactive phase plane and solver
In the phase plane image below, the hugoniot loci are plotted in red and the integral curves in blue.  The portion of each curve that satisfies the entropy condition is plotted as a solid line, while the entropy-violating parts are plotted as dashed lines.  The curves are plotted for both the left and right states.  Thus the function $\phi_l$ just described consists of the solid curve (blue and red parts) passing through the left state $q_\ell$; the function $\phi_r$ is the solid curve (blue and red parts) passing through $q_r$.  The intersection of these curves is labeled $q_m$.  [Click here to go to the interactive phase plane](phase_plane_shallow_water_small.html).  While you're there, be sure to click and drag the dark circles indicating the values of $q_\ell$ and $q_r$ in the top left plot.  The time can also be adjusted by dragging the circle in the top right plot.

Finally, the code below generates an interactive widget that shows the solution to the Riemann problem for arbitrary left and right states, including solutions with dry states.

```python
# Unlike javascript version, this works with dry states
def plot_exact_riemann_solution(h_l=3.,u_l=0.,h_r=1.,u_r=0.,t=0.2):
    q_l = shallow_water.primitive_to_conservative(h_l,u_l)
    q_r = shallow_water.primitive_to_conservative(h_r,u_r)

    x = np.linspace(-1.,1.,1000)
    states, speeds, reval, wave_types = \
        shallow_water.exact_riemann_solution(q_l, q_r, grav=g)
    q = reval(x/t)
    primitive = shallow_water.conservative_to_primitive(q[0],q[1])

    fig = plt.figure(figsize=figsize)
    names = shallow_water.primitive_variables
    axes = [0]*len(q_l)
    for i in range(len(q_l)):
        axes[i] = fig.add_subplot(1,len(q_l),i+1)
        q = primitive[i]
        plt.plot(x,q,'-k',linewidth=3)
        if i==0:
            plt.fill_between(x,q,color='b')
        plt.title(names[i])
        qmax = max(q)
        qmin = min(q)
        qdiff = qmax - qmin
        axes[i].set_ylim((qmin-0.1*qdiff,qmax+0.1*qdiff))
        axes[i].set_xlim(-1,1)
    plt.show()

interact(plot_exact_riemann_solution,
         h_l=FloatSlider(min=0.,max=10.,step=0.1,value=3.,
                                 description=r'$h_l$'),
         u_l=FloatSlider(min=-5.,max=5.,step=0.1,value=0.,
                                 description=r'$u_l$'),
         h_r=FloatSlider(min=0.,max=10.,step=0.1,value=1.,
                                 description=r'$h_r$'),
         u_r=FloatSlider(min=-5.,max=5.,step=0.1,value=0.,
                                 description=r'$u_r$'),
         t=(0.1,1.,0.1));
```