# Using Fairness Dashboard with Loan Allocation Data

This notebook shows how to use [`Fairlearn`](https://fairlearn.github.io/) and the Responsible AI Widget's Fairness dashboard to generate predictors for the Census dataset. This dataset is a classification problem - given a range of data about 32,000 individuals, predict whether their annual income is above or below fifty thousand dollars per year.

For the purposes of this notebook, we shall treat this as a loan decision problem. We will pretend that the label indicates whether or not each individual repaid a loan in the past. We will use the data to train a predictor to predict whether previously unseen individuals will repay a loan or not. The assumption is that the model predictions are used to decide whether an individual should be offered a loan.

We will first train a fairness-unaware predictor and use the fairness dashboard to demonstrate how this model leads to unfair decisions (under a specific notion of fairness called *demographic parity*). We then mitigate unfairness by applying the `GridSearch` algorithm from `Fairlearn` package.

## Install Required Packages

```python
# %pip install --upgrade fairlearn
# %pip install --upgrade raiwidgets

# %pip install -e "/mnt/c/Users/t-alexquach/Desktop/fairlearn"
# %pip install raiwidgets

# %pip install -e "/mnt/c/Users/t-alexquach/Music/fairlearn"
%pip install -e "/mnt/c/Users/t-alexquach/Music/responsible-ai-widgets/raiwidgets"
# %pip install -e "/mnt/c/Users/t-alexquach/Music/responsible-ai-widgets/responsibleai"
```

    Obtaining file:///mnt/c/Users/t-alexquach/Music/responsible-ai-widgets/raiwidgets
    Requirement already satisfied: erroranalysis>=0.1.8 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from raiwidgets==0.0.0) (0.1.8)
    Requirement already satisfied: ipython==7.16.1 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from raiwidgets==0.0.0) (7.16.1)
    Requirement already satisfied: jinja2==2.11.3 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from raiwidgets==0.0.0) (2.11.3)
    Requirement already satisfied: lightgbm>=2.0.11 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from raiwidgets==0.0.0) (3.2.1)
    Requirement already satisfied: numpy>=1.17.2 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from raiwidgets==0.0.0) (1.20.3)
    Requirement already satisfied: pandas>=0.25.1 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from raiwidgets==0.0.0) (1.2.4)
    Requirement already satisfied: rai-core-flask==0.2.2 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from raiwidgets==0.0.0) (0.2.2)
    Requirement already satisfied: responsibleai==0.0.0 in /mnt/c/Users/t-alexquach/Music/responsible-ai-widgets/responsibleai (from raiwidgets==0.0.0) (0.0.0)
    Requirement already satisfied: scikit-learn>=0.22.1 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from raiwidgets==0.0.0) (0.24.2)
    Requirement already satisfied: scipy>=1.4.1 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from raiwidgets==0.0.0) (1.6.3)
    Requirement already satisfied: pygments in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from ipython==7.16.1->raiwidgets==0.0.0) (2.9.0)
    Requirement already satisfied: prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from ipython==7.16.1->raiwidgets==0.0.0) (3.0.18)
    Requirement already satisfied: setuptools>=18.5 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from ipython==7.16.1->raiwidgets==0.0.0) (44.0.0)
    Requirement already satisfied: backcall in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from ipython==7.16.1->raiwidgets==0.0.0) (0.2.0)
    Requirement already satisfied: traitlets>=4.2 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from ipython==7.16.1->raiwidgets==0.0.0) (5.0.5)
    Requirement already satisfied: pexpect in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from ipython==7.16.1->raiwidgets==0.0.0) (4.8.0)
    Requirement already satisfied: pickleshare in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from ipython==7.16.1->raiwidgets==0.0.0) (0.7.5)
    Requirement already satisfied: jedi>=0.10 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from ipython==7.16.1->raiwidgets==0.0.0) (0.18.0)
    Requirement already satisfied: decorator in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from ipython==7.16.1->raiwidgets==0.0.0) (4.4.2)
    Requirement already satisfied: MarkupSafe>=0.23 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from jinja2==2.11.3->raiwidgets==0.0.0) (2.0.1)
    Requirement already satisfied: Flask==1.1.2 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from rai-core-flask==0.2.2->raiwidgets==0.0.0) (1.1.2)
    Requirement already satisfied: gevent==20.9.0 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from rai-core-flask==0.2.2->raiwidgets==0.0.0) (20.9.0)
    Requirement already satisfied: Flask-Cors==3.0.9 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from rai-core-flask==0.2.2->raiwidgets==0.0.0) (3.0.9)
    Requirement already satisfied: greenlet==0.4.17 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from rai-core-flask==0.2.2->raiwidgets==0.0.0) (0.4.17)
    Requirement already satisfied: itsdangerous>=0.24 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from Flask==1.1.2->rai-core-flask==0.2.2->raiwidgets==0.0.0) (2.0.1)
    Requirement already satisfied: Werkzeug>=0.15 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from Flask==1.1.2->rai-core-flask==0.2.2->raiwidgets==0.0.0) (2.0.1)
    Requirement already satisfied: click>=5.1 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from Flask==1.1.2->rai-core-flask==0.2.2->raiwidgets==0.0.0) (8.0.1)
    Requirement already satisfied: Six in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from Flask-Cors==3.0.9->rai-core-flask==0.2.2->raiwidgets==0.0.0) (1.16.0)
    Requirement already satisfied: zope.event in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from gevent==20.9.0->rai-core-flask==0.2.2->raiwidgets==0.0.0) (4.5.0)
    Requirement already satisfied: zope.interface in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from gevent==20.9.0->rai-core-flask==0.2.2->raiwidgets==0.0.0) (5.4.0)
    Requirement already satisfied: dice-ml<0.7,>=0.6.1 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from responsibleai==0.0.0->raiwidgets==0.0.0) (0.6.1)
    Requirement already satisfied: econml>=0.12.0b2 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from responsibleai==0.0.0->raiwidgets==0.0.0) (0.12.0b2)
    Requirement already satisfied: interpret-community>=0.18.1 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from responsibleai==0.0.0->raiwidgets==0.0.0) (0.18.1)
    Requirement already satisfied: h5py in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from dice-ml<0.7,>=0.6.1->responsibleai==0.0.0->raiwidgets==0.0.0) (3.2.1)
    Requirement already satisfied: tqdm in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from dice-ml<0.7,>=0.6.1->responsibleai==0.0.0->raiwidgets==0.0.0) (4.61.0)
    Requirement already satisfied: statsmodels>=0.10 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from econml>=0.12.0b2->responsibleai==0.0.0->raiwidgets==0.0.0) (0.12.2)
    Requirement already satisfied: shap~=0.38.1 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from econml>=0.12.0b2->responsibleai==0.0.0->raiwidgets==0.0.0) (0.38.1)
    Requirement already satisfied: dowhy in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from econml>=0.12.0b2->responsibleai==0.0.0->raiwidgets==0.0.0) (0.6)
    Requirement already satisfied: numba!=0.42.1 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from econml>=0.12.0b2->responsibleai==0.0.0->raiwidgets==0.0.0) (0.53.1)
    Requirement already satisfied: sparse in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from econml>=0.12.0b2->responsibleai==0.0.0->raiwidgets==0.0.0) (0.12.0)
    Requirement already satisfied: graphviz in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from econml>=0.12.0b2->responsibleai==0.0.0->raiwidgets==0.0.0) (0.16)
    Requirement already satisfied: joblib>=0.13.0 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from econml>=0.12.0b2->responsibleai==0.0.0->raiwidgets==0.0.0) (1.0.1)
    Requirement already satisfied: interpret-core[required]<=0.2.4,>=0.1.20 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from interpret-community>=0.18.1->responsibleai==0.0.0->raiwidgets==0.0.0) (0.2.4)
    Requirement already satisfied: packaging in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from interpret-community>=0.18.1->responsibleai==0.0.0->raiwidgets==0.0.0) (20.9)
    Requirement already satisfied: parso<0.9.0,>=0.8.0 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from jedi>=0.10->ipython==7.16.1->raiwidgets==0.0.0) (0.8.2)
    Requirement already satisfied: wheel in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from lightgbm>=2.0.11->raiwidgets==0.0.0) (0.36.2)
    Requirement already satisfied: llvmlite<0.37,>=0.36.0rc1 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from numba!=0.42.1->econml>=0.12.0b2->responsibleai==0.0.0->raiwidgets==0.0.0) (0.36.0)
    Requirement already satisfied: python-dateutil>=2.7.3 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from pandas>=0.25.1->raiwidgets==0.0.0) (2.8.1)
    Requirement already satisfied: pytz>=2017.3 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from pandas>=0.25.1->raiwidgets==0.0.0) (2021.1)
    Requirement already satisfied: wcwidth in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0->ipython==7.16.1->raiwidgets==0.0.0) (0.2.5)
    Requirement already satisfied: threadpoolctl>=2.0.0 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from scikit-learn>=0.22.1->raiwidgets==0.0.0) (2.1.0)
    Requirement already satisfied: cloudpickle in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from shap~=0.38.1->econml>=0.12.0b2->responsibleai==0.0.0->raiwidgets==0.0.0) (1.6.0)
    Requirement already satisfied: slicer==0.0.7 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from shap~=0.38.1->econml>=0.12.0b2->responsibleai==0.0.0->raiwidgets==0.0.0) (0.0.7)
    Requirement already satisfied: patsy>=0.5 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from statsmodels>=0.10->econml>=0.12.0b2->responsibleai==0.0.0->raiwidgets==0.0.0) (0.5.1)
    Requirement already satisfied: ipython-genutils in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from traitlets>=4.2->ipython==7.16.1->raiwidgets==0.0.0) (0.2.0)
    Requirement already satisfied: sympy>=1.4 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from dowhy->econml>=0.12.0b2->responsibleai==0.0.0->raiwidgets==0.0.0) (1.8)
    Requirement already satisfied: pydot>=1.4 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from dowhy->econml>=0.12.0b2->responsibleai==0.0.0->raiwidgets==0.0.0) (1.4.2)
    Requirement already satisfied: networkx>=2.0 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from dowhy->econml>=0.12.0b2->responsibleai==0.0.0->raiwidgets==0.0.0) (2.5.1)
    Requirement already satisfied: pyparsing>=2.1.4 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from pydot>=1.4->dowhy->econml>=0.12.0b2->responsibleai==0.0.0->raiwidgets==0.0.0) (2.4.7)
    Requirement already satisfied: mpmath>=0.19 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from sympy>=1.4->dowhy->econml>=0.12.0b2->responsibleai==0.0.0->raiwidgets==0.0.0) (1.2.1)
    Requirement already satisfied: ptyprocess>=0.5 in /home/alex/.local/share/virtualenvs/responsible-ai-widgets-ySgAgbNB/lib/python3.8/site-packages (from pexpect->ipython==7.16.1->raiwidgets==0.0.0) (0.7.0)
    Installing collected packages: raiwidgets
      Running setup.py develop for raiwidgets
    Successfully installed raiwidgets-0.0.0
    Note: you may need to restart the kernel to use updated packages.

## Load and preprocess the data set

For simplicity, we import the data set from the `shap` package, which contains the data in a cleaned format. We start by importing the various modules we're going to use:

```python
from raiwidgets import FairnessDashboard

from fairlearn.reductions import GridSearch
from fairlearn.reductions import DemographicParity
from fairlearn.metrics import MetricFrame, selection_rate

from sklearn import svm, neighbors, tree
from sklearn.metrics import accuracy_score
from sklearn.preprocessing import LabelEncoder,StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import fetch_openml

import pandas as pd

import numpy as np
```

We can now load and inspect the data from `openml`:

```python
data = fetch_openml(data_id=1590, as_frame=True)
X_raw = data.data
y_true = (data.target == '>50K') * 1

X_raw["race"].value_counts().to_dict()
```

We are going to treat the sex of each individual as a protected attribute (where 0 indicates female and 1 indicates male), and in this particular case we are going separate this attribute out and drop it from the main data. We then perform some standard data preprocessing steps to convert the data into a format suitable for the ML algorithms

```python
sensitive_features = X_raw[['sex','race']]
X = X_raw.drop(labels=['sex', 'race'],axis = 1)
X = pd.get_dummies(X)

sc = StandardScaler()
X_scaled = sc.fit_transform(X)
X_scaled = pd.DataFrame(X_scaled, columns=X.columns)

le = LabelEncoder()
y_true = le.fit_transform(y_true)
```

Finally, we split the data into training and test sets:

```python
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test, sensitive_features_train, sensitive_features_test = \
    train_test_split(X_scaled, y_true, sensitive_features,
                     test_size = 0.2, random_state=0, stratify=y_true)

# Work around indexing bug
X_train = X_train.reset_index(drop=True)
sensitive_features_train = sensitive_features_train.reset_index(drop=True)
X_test = X_test.reset_index(drop=True)
sensitive_features_test = sensitive_features_test.reset_index(drop=True)
```

## Train a fairness-unaware predictor

To show the effect of `Fairlearn` we will first train a standard ML predictor that does not incorporate fairness For speed of demonstration, we use a simple logistic regression estimator from `sklearn`:

```python
unmitigated_predictor = LogisticRegression(solver='liblinear', fit_intercept=True)

unmitigated_predictor.fit(X_train, y_train)
```

We can load this predictor into the Fairness dashboard, and examine how it is unfair (there is a warning about AzureML since we are not yet integrated with that product):

```python
y_pred = unmitigated_predictor.predict(X_test)

FairnessDashboard(sensitive_features=sensitive_features_test,
                  y_true=y_test,
                  y_pred=y_pred)
```

Looking at the disparity in accuracy, we see that males have an error rate about three times greater than the females. More interesting is the disparity in opportunitiy - males are offered loans at three times the rate of females.

Despite the fact that we removed the feature from the training data, our predictor still discriminates based on sex. This demonstrates that simply ignoring a protected attribute when fitting a predictor rarely eliminates unfairness. There will generally be enough other features correlated with the removed attribute to lead to disparate impact.

## Mitigation with GridSearch

The `GridSearch` class in `Fairlearn` implements a simplified version of the exponentiated gradient reduction of [Agarwal et al. 2018](https://arxiv.org/abs/1803.02453). The user supplies a standard ML estimator, which is treated as a blackbox. `GridSearch` works by generating a sequence of relabellings and reweightings, and trains a predictor for each.

For this example, we specify demographic parity (on the protected attribute of sex) as the fairness metric. Demographic parity requires that individuals are offered the opportunity (are approved for a loan in this example) independent of membership in the protected class (i.e., females and males should be offered loans at the same rate). We are using this metric for the sake of simplicity; in general, the appropriate fairness metric will not be obvious.

```python
sweep = GridSearch(LogisticRegression(solver='liblinear', fit_intercept=True),
                   constraints=DemographicParity(),
                   grid_size=70)
```

Our algorithms provide `fit()` and `predict()` methods, so they behave in a similar manner to other ML packages in Python. We do however have to specify two extra arguments to `fit()` - the column of protected attribute labels, and also the number of predictors to generate in our sweep.

After `fit()` completes, we extract the full set of predictors from the `GridSearch` object.

```python
sweep.fit(X_train, y_train,
          sensitive_features=sensitive_features_train.sex)

predictors = sweep.predictors_
```

We could load these predictors into the Fairness dashboard now. However, the plot would be somewhat confusing due to their number. In this case, we are going to remove the predictors which are dominated in the accuracy-disparity space by others from the sweep (note that the disparity will only be calculated for the protected attribute; other potentially protected attributes will not be mitigated). In general, one might not want to do this, since there may be other considerations beyond the strict optimization of accuracy and disparity (of the given protected attribute).

```python
accuracies, disparities = [], []

for predictor in predictors:
    accuracy_metric_frame = MetricFrame(accuracy_score, y_train, predictor.predict(X_train), sensitive_features=sensitive_features_train.sex)
    selection_rate_metric_frame = MetricFrame(selection_rate, y_train, predictor.predict(X_train), sensitive_features=sensitive_features_train.sex)
    accuracies.append(accuracy_metric_frame.overall)
    disparities.append(selection_rate_metric_frame.difference())

all_results = pd.DataFrame( {"predictor": predictors, "accuracy": accuracies, "disparity": disparities})

all_models_dict = {"unmitigated": unmitigated_predictor}
dominant_models_dict = {"unmitigated": unmitigated_predictor}
base_name_format = "grid_{0}"
row_id = 0
for row in all_results.itertuples():
    model_name = base_name_format.format(row_id)
    all_models_dict[model_name] = row.predictor
    accuracy_for_lower_or_eq_disparity = all_results["accuracy"][all_results["disparity"] <= row.disparity]
    if row.accuracy >= accuracy_for_lower_or_eq_disparity.max():
        dominant_models_dict[model_name] = row.predictor
    row_id = row_id + 1
```

We can construct predictions for all the models, and also for the dominant models:

```python
dashboard_all = {}
for name, predictor in all_models_dict.items():
    value = predictor.predict(X_test)
    dashboard_all[name] = value

dominant_all = {}
for name, predictor in dominant_models_dict.items():
    dominant_all[name] = predictor.predict(X_test)

FairnessDashboard(sensitive_features=sensitive_features_test,
                  y_true=y_test.tolist(),
                  y_pred=dominant_all)
```

We see a Pareto front forming - the set of predictors which represent optimal tradeoffs between accuracy and disparity in predictions. In the ideal case, we would have a predictor at (1,0) - perfectly accurate and without any unfairness under demographic parity (with respect to the protected attribute "sex"). The Pareto front represents the closest we can come to this ideal based on our data and choice of estimator. Note the range of the axes - the disparity axis covers more values than the accuracy, so we can reduce disparity substantially for a small loss in accuracy.

By clicking on individual models on the plot, we can inspect their metrics for disparity and accuracy in greater detail. In a real example, we would then pick the model which represented the best trade-off between accuracy and disparity given the relevant business constraints.

```python

```

```python

```


# KVLCC2 Ikeda at various frequencies

# Purpose
Will the viscous damping predicted with Ikeda change with frequency?

# Methodology
Make a variation around the natural frequency

## Results
Describe and comment the most important results.

# Suggested next steps
State suggested next steps, based on results obtained in this notebook.

# Setup

```python
# %load imports.py
"""
These is the standard setup for the notebooks.
"""

%matplotlib inline
%load_ext autoreload
%autoreload 2

from jupyterthemes import jtplot
jtplot.style(theme='onedork', context='notebook', ticks=True, grid=False)

import pandas as pd
pd.options.display.max_rows = 999
pd.options.display.max_columns = 999
pd.set_option("display.max_columns", None)
import numpy as np
import os
import matplotlib.pyplot as plt
from collections import OrderedDict
#plt.style.use('paper')

#import data
import copy
from mdldb.run import Run

from sklearn.pipeline import Pipeline
from rolldecayestimators.transformers import CutTransformer, LowpassFilterDerivatorTransformer, ScaleFactorTransformer, OffsetTransformer
from rolldecayestimators.direct_estimator_cubic import EstimatorQuadraticB, EstimatorCubic
from rolldecayestimators.ikeda_estimator import IkedaQuadraticEstimator
import rolldecayestimators.equations as equations
import rolldecayestimators.lambdas as lambdas
from rolldecayestimators.substitute_dynamic_symbols import lambdify
import rolldecayestimators.symbols as symbols
import sympy as sp

from sympy.physics.vector.printing import vpprint, vlatex
from IPython.display import display, Math, Latex

from sklearn.metrics import r2_score
from src.data import database
from mdldb import tables

```

```python
import pyscores2
import pyscores2.runScores2
import pyscores2.xml_hydrostatics
from pyscores2.output import OutputFile
from rolldecayestimators.ikeda import Ikeda, IkedaR

from rolldecayestimators.simplified_ikeda_class import SimplifiedIkeda, SimplifiedIkedaABS
from rolldecayestimators.simplified_ikeda import limits_kawahara
from pyscores2.runScores2 import Calculation
from pyscores2.indata import Indata
import joblib
from scipy.optimize import least_squares
from reports.paper_writing import save_fig
```

```python
db = database.get_db()
```

```python
sql = """
SELECT * from run
INNER JOIN loading_conditions
ON (run.loading_condition_id = loading_conditions.id)
INNER JOIN models
ON (run.model_number = models.model_number)
INNER JOIN ships
ON (run.ship_name = ships.name)
WHERE run.model_number='M5057-01-A' and run.test_type='roll decay' and run.project_number=40178362;
"""
df_rolldecays = pd.read_sql(sql=sql, con=db.engine)
df_rolldecays=df_rolldecays.loc[:,~df_rolldecays.columns.duplicated()]
df_rolldecays.set_index('id', inplace=True)

df_rolldecays['ship_speed'].fillna(0, inplace=True)

```

```python
df_rolldecays.head()
```

```python
def calculate_ikeda(ikeda):

    output = pd.DataFrame()
    output['B_44_hat']   = ikeda.calculate_B44()
    output['B_W0_hat']   = ikeda.calculate_B_W0()
    output['B_W_hat']    = ikeda.calculate_B_W()
    output['B_F_hat']    = ikeda.calculate_B_F()
    output['B_E_hat']    = ikeda.calculate_B_E()
    output['B_BK_hat']   = ikeda.calculate_B_BK()
    output['B_L_hat']    = ikeda.calculate_B_L()
    output['Bw_div_Bw0'] = ikeda.calculate_Bw_div_Bw0()
    return output
```

```python
#df_rolldecays=df_rolldecays.loc[[21337,21338,21340,]].copy()
df_rolldecays=df_rolldecays.loc[[21338,21340,]].copy()
```

```python
resources={
    21338 : {
        'scores_indata_path':'../models/KVLCC2_speed.IN',
        'scores_outdata_path':'../data/interim/KVLCC2_speed.out',
        'roll_decay_model':'../models/KVLCC2_0_speed.pkl',
    },
    21340 : {
        'scores_indata_path':'../models/KVLCC2_speed.IN',
        'scores_outdata_path':'../data/interim/KVLCC2_speed.out',
        'roll_decay_model':'../models/KVLCC2_speed.pkl',
    }
}
```

```python
def unhat(df_ikeda, volume, beam, g, rho):
    for key,data in df_ikeda.items():
        if '_hat' in key:
            new_key = key.replace('_hat','')
            df_ikeda[new_key] = lambdas.B_from_hat_lambda(B_44_hat=data, Disp=volume, beam=beam, g=g, rho=rho)

    return df_ikeda
```

```python
g=9.81
rho=1000

phi_a = np.deg2rad(10)
results = pd.DataFrame()

for id, row in df_rolldecays.iterrows():
    run = db.session.query(Run).get(int(row.name))
    run = database.load_run(run, save_as_example=False, prefer_hdf5=True)
    scale_factor = run.model.scale_factor

    resource = resources[id]

    ## Load ScoresII results
    indata = Indata()
    indata.open(indataPath='../models/KVLCC2_speed.IN')
    output_file = OutputFile(filePath='../data/interim/KVLCC2_speed.out')

    ## Compare with model test
    model = joblib.load(resource['roll_decay_model'])
    estimator = model['estimator']

    ## Non. Lin. linear equivalent damping
    GM = run.loading_condition.gm/scale_factor
    volume = run.loading_condition.Volume/(scale_factor**3)
    GM = run.loading_condition.gm/scale_factor
    beam = run.ship.beam/scale_factor

    meta_data = {
        'Volume':volume,
        'GM':GM,
        'rho':rho,
        'g':g,
    }
    parameters = estimator.result_for_database(meta_data = meta_data)

    ## Run Ikeda
    N=99
    ws_factor = np.linspace(0.5,1.5,N)
    ws = parameters['omega0']*ws_factor

    scale_factor=run.model.scale_factor
    V = row.ship_speed*1.852/3.6/np.sqrt(scale_factor)
    Vs = V*np.ones(len(ws))
    phi_as = phi_a*np.ones(len(ws))

    if not run.ship.BKL:
        BKL=0
    else:
        BKL=run.ship.BKL/scale_factor

    if not run.ship.BKB:
        BKB = 0
    else:
        BKB=run.ship.BKB/scale_factor

    BKL_ = BKL*np.ones(N)
    BKB_ = BKB*np.ones(N)

    ikeda = Ikeda.load_scoresII(V=Vs, w=ws, fi_a=phi_as, indata=indata, output_file=output_file,
                                scale_factor=scale_factor, BKL=BKL_, BKB=BKB_)

    R = 0.15*run.ship.beam/scale_factor  # Just guessing...
    ikeda.R = R

    df_ikeda = calculate_ikeda(ikeda=ikeda)
    df_ikeda['w']=ws
    df_ikeda['ws_factor']=ws_factor
    df_ikeda['phi_a']=phi_as
    df_ikeda.set_index('ws_factor',inplace=True)
    df_ikeda['id']=id
    df_ikeda = unhat(df_ikeda=df_ikeda, volume=volume, beam=beam, g=g, rho=rho)

    # assuming : phi = phi_a*sin(w*t)
    #phi1ds = phi_as*w*np.cos()

    #df_ikeda['B'] = (parameters['B_1'] + parameters['B_2']*np.abs(ws))*ws + parameters['B_3']*ws**3
    #df_ikeda['B']=df_ikeda['B'].astype(float)

    df_ikeda['B']=lambdas.B_e_lambda_cubic(B_1=parameters['B_1'], B_2=parameters['B_2'], B_3=parameters['B_3'], omega0=ws, phi_a=phi_as)
    df_ikeda['B']=df_ikeda['B'].astype(float)

    results=results.append(df_ikeda)

```

```python
parameters
```

```python
equations.B_e_equation_cubic
```

```python
for id,df_ikeda in results.groupby(by='id'):
    fig,ax=plt.subplots()
    interesting_ = ['B_W','B_L','B_F','B_E',]

    df_ikeda.plot.area(x='w', y=interesting_, ax=ax)

    #df_ikeda.plot(x='w', y='B', ax=ax)
    df_ikeda.plot(x='w', y='B', label=r'$B_e$ roll decay', ax=ax)

    ax.legend()
    ax.set_xlabel(r'$ \omega $ [rad/s] (Roll frequency)')
    ax.set_ylabel(r'$ B_e $ [Nm/s]');

    s = df_rolldecays.loc[id]
    phi_a = np.rad2deg(df_ikeda['phi_a'].iloc[0])
    title='Ship speed: %0.1f [kts] $\phi_a$: %i [deg]' % (s.ship_speed,phi_a)
    ax.set_title(title)

```

```python
for id,df_ikeda in results.groupby(by='id'):
    fig,ax=plt.subplots()
    interesting_ = ['B_L_hat','B_F_hat','B_E_hat',]
    df_ikeda_fraction = df_ikeda[interesting_].copy()
    df_ikeda_fraction/=df_ikeda_fraction.loc[1]

    df_ikeda_fraction.plot.area(y=interesting_, ax=ax)

    ax.legend()
    ax.set_xlabel(r'$ \frac{\omega}{\omega_0} $ [-] (Roll frequency)')
    ax.set_ylabel(r'$\frac{\hat{B_e}}{\hat{B_e^0}}$ [-]');

    s = df_rolldecays.loc[id]
    phi_a = np.rad2deg(df_ikeda['phi_a'].iloc[0])
    title='Ship speed: %0.1f [kts] $\phi_a$: %i [deg]' % (s.ship_speed,phi_a)
    ax.set_title(title)

    figure_name = 'KVLCC2_B_visc_frequency%0.1f' % s.ship_speed
    save_fig(fig, name=figure_name)

```

```python

```


# Saturating the Cosmological Field Information

by Lucas Makinen [](https://orcid.org/0000-0002-3795-6933 "") [](https://twitter.com/lucasmakinen?lang=en ""), Tom Charnock [](https://orcid.org/0000-0002-7416-3107 "Redirect to orcid") [](https://twitter.com/t_charnock?lang=en "")), Justin Alsing [](https://scholar.google.com/citations?user=ICPFL8AAAAAJ&hl=en "Redirect to orcid"), and Ben Wandelt [](https://twitter.com/bwandelt?lang=en "")

>read the paper: [accompanying arXiv paper](arxiv link)

>get the code:  [https://github.com/tlmakinen/FieldIMNNs](https://github.com/tlmakinen/FieldIMNNs)

In this short walkthrough we will train an optimal neural compression on cosmological fields generated from a fiducial model and demonstrate theoretical information saturation. This code can be run so long as the user has a Google Colab Pro account with increased computational resources.

First we'll install the relevant libraries and walk through the simulation implementation. Then we'll build a neural IMNN compressor to generate two optimal summaries for our simulations.

### Q: Wait a second -- how do we know this works ?
If you're not convinced by our method by the end of this tutorial, we invite you to take a look at our [benchmarking tutorial with Gaussian fields from power spectra](https://www.aquila-consortium.org/doc/imnn/pages/examples/2d_field_inference/2d_field_inference.html), which is also runnable in-browser on [this Colab notebook](https://colab.research.google.com/drive/1_y_Rgn3vrb2rlk9YUDUtfwDv9hx774ZF#scrollTo=EW4H-R8I0q6n).

# step 1: loading packages and setting up environment

1.  check that Colab is set to run on a GPU ! Go to `Runtime`>`change runtime type` and select `GPU` from the dropdown menu. Next, enable dark mode by going to `settings`>`Theme` and selecting `dark` (protect your eyes !)

2.   install packages. The current code relies on several libraries, namely `jax` and `tensorflow_probability`. However, we require both plain `tensorflow_probability` (`tfp`) and the experimental `tensorflow_probability.substrates.jax` (`tfpj`) packages for different parts of our inference
3. for some Colab sessions, you may need to run the second cell so that `!pip install jax-cosmo` gets the package imported properly.

```python
#@title set up environment <font color='lightgreen'>[RUN ME FIRST]</font>
%tensorflow_version 2.x
import tensorflow as tf
print('tf version', tf.__version__)

!pip install -q jax==0.2.11

!pip install -q tensorflow-probability

import tensorflow_probability as tfp
print('tfp version:', tfp.__version__)
tfd = tfp.distributions
tfb = tfp.bijectors

!pip install -q imnn

!python -m pip install -q jax-cosmo
```

    tf version 2.5.0
         |████████████████████████████████| 604kB 8.0MB/s
      Building wheel for jax (setup.py) ... done
    tfp version: 0.13.0
         |████████████████████████████████| 133kB 9.1MB/s
         |████████████████████████████████| 81kB 8.4MB/s
         |████████████████████████████████| 286kB 9.3MB/s
      Building wheel for jax-cosmo (setup.py) ... done

```python
# now import all the required libraries
import jax.numpy as np
from jax import grad, jit, vmap
from jax import random
import jax

print('jax version:', jax.__version__)

# for nn model stuff
import jax.experimental.optimizers as optimizers
import jax.experimental.stax as stax

# tensorflow-prob VANILLA
tfd = tfp.distributions
tfb = tfp.bijectors

# tensorflow-prob-JAX
import tensorflow_probability.substrates.jax as tfpj

tfdj = tfpj.distributions
tfbj = tfpj.bijectors

# for imnn
import imnn
import imnn.lfi

print('IMNN version:', imnn.__version__)

# jax-cosmo module
!python -m pip install -q jax-cosmo
import jax_cosmo as jc
print('jax-cosmo version:', jc.__version__)

# matplotlib stuff
import matplotlib.pyplot as plt
from scipy.linalg import toeplitz
import seaborn as sns
sns.set()

rng = random.PRNGKey(2)
```

    jax version: 0.2.11
    IMNN version: 0.3.0
    jax-cosmo version: 0.1rc7

note: if the above cell fails for installing jax-cosmo, just run it again: Colab will rearrange the headings needed.

```python
from jax.config import config
config.update('jax_enable_x64', True)
```

make sure we're using 64-bit precision and running on a GPU !

```python
from jax.lib import xla_bridge
print(xla_bridge.get_backend().platform)
```

    gpu

# Cosmological Fields from the Eisenstein-Hu linear matter power spectrum
We're interested in extracting the cosmological parameters $\Omega_c$ and $\sigma_8$ *directly* from cosmological field pixels. To generate our simulations we'll need to install the library `jax-cosmo` to generate our differentiable model power spectra.

## choose fiducial model
To train our neural compression, we need to choose a fiducial model to train the IMNN.

For example lets say that our fiducial cosmology has $\Omega_c=0.60$ and $\sigma_8=0.60$. This is *deliberately* far from, say, Planck parameters -- we want to investigate how our compression behaves if we don't know our universe's true parameters.

```python
cosmo_params = jc.Planck15(Omega_c=0.60, sigma8=0.60)
θ_fid = np.array(
    [cosmo_params.Omega_c,
     cosmo_params.sigma8],
    dtype=np.float32)

n_params = θ_fid.shape[0]
```

Our power spectrum $P_{\rm LN}(k)$ is the linear matter power spectrum defined as

```python
def P(k, A=0.60, B=0.60):
    cosmo_params = jc.Planck15(Omega_c=A, sigma8=B)
    return jc.power.linear_matter_power(cosmo_params, k)
```

and we can visualize it in $k$-space (small $k$ <=> big $r$, big $k$ <=> small $r$) :

```python
#@title plot the Eisenstein-Hu $P(k)$ <font color='lightgreen'>[run me]</font>

sns.set()
L = 250.
N = 128.
#kmax = 1.0
#kmin = 0.5 / (N)

kmax = N / L
kmin = 1. / L

kbin = np.linspace(kmin, kmax, num=100)

power_spec = P(kbin, A=cosmo_params.Omega_c, B=cosmo_params.sigma8)

plt.style.use('dark_background')
plt.grid(b=None)

plt.plot(kbin, power_spec, linewidth=2)
plt.xlabel(r'$k\ \rm [h\ Mpc^{-1}]$', fontsize=14)
plt.ylabel(r'$P(k)\ \rm$', fontsize=14)
plt.ylim((1e2, 1e4))
plt.xscale('log')
plt.yscale('log')
```

## Lognormal Fields from Power Spectra: how much information is embedded in the field ?
Cosmologists often use lognormal fields as "the poor man's large scale structure" since they're analytically interrogable and easy to obtain from Gaussian fields. We'll walk through how to obtain the *theoretical* information content of such fields using the Fisher formalism.

The likelihood for an $N_{\rm pix}\times N_{\rm pix}$ Gaussian field, $\boldsymbol{\delta}$, can be explicitly written down for the Fourier transformed data, $\boldsymbol{\Delta}$ as
$$\mathcal{L}(\boldsymbol{\Delta}|\boldsymbol{\theta}) = \frac{1}{(2\pi)^{N_{\rm pix}^2 / 2} |P_{\rm G}({\bf k}, \boldsymbol{\theta})|^{1/2}}\exp{\left(-\frac{1}{2}\boldsymbol{\Delta}\left(P_{\rm G}({\bf k}, \boldsymbol{\theta})\right)^{-1}\boldsymbol{\Delta}\right)}$$
Since the Fisher information can be calculated from the expectation value of the second derivative of the score, i.e. the log likelihood
$${\bf F}_{\alpha\beta} = - \left.\left\langle\frac{\partial^2\ln\mathcal{L}(\Delta|\boldsymbol{\theta})}{\partial\theta_\alpha\partial\theta_\beta}\right\rangle\right|_{\boldsymbol{\theta}=\boldsymbol{\theta}^\textrm{fid}}$$
then we know that analytically the Fisher information must be
$${\bf F}_{\alpha\beta} = \frac{1}{2} {\rm Tr} \left(\frac{\partial P_{\rm G}({\bf k}, \boldsymbol{\theta})}{\partial\theta_\alpha}\left(P_{\rm G}({\bf k}, \boldsymbol{\theta})\right)^{-1}\frac{\partial P_{\rm G}({\bf k}, \boldsymbol{\theta})}{\partial\theta_\beta}\left(P_{\rm G}({\bf k}, \boldsymbol{\theta})\right)^{-1}\right)$$
where $\alpha$ and $\beta$ label the parameters (for instance $ \Omega_c, \sigma_8$) in the power spectrum. As each $k$-mode is uncoupled for this power law form we require the derivatives
$$\begin{align}
\left(\frac{\partial P_{\rm G}({\bf k}, \boldsymbol{\theta})}{\partial \Omega_c},\
\frac{\partial P_{\rm G}({\bf k}, \boldsymbol{\theta})}{\partial \sigma_8}\right) \\
\end{align}$$
We can set up these derivative functions *so long as our code for $P(k)$ is differentiable*.

For *lognormal* fields, this likelihood changes somewhat. Formally, if a random variable $Y$ has a normal distribution, then the exponential function of $Y$, $X = \exp(Y)$, has a log-normal distribution. We will generate our log-normal fields with a power spectrum such that the *lognormal field has the specified $P_{\rm LN}(k)$*. This means that we need to employ the *backwards conversion formula* , presented by [M. Greiner? and T.A. Enßlin](https://arxiv.org/pdf/1312.1354.pdf), to obtain the correct form for $P_{\rm G}(k)$ needed for the above Fisher evaluation:
$$ P_{\rm G} = \int d^u x e^{i \textbf{k} \cdot \textbf{x}} \ln \left( \int \frac{d^u q}{(2\pi)^u} e^{i \textbf{q} \cdot \textbf{x}} P_{\rm LN}(\textbf{q}) \right) $$

which we can do numerically (and differentiably !) in `Jax`:

```python
def ln_fisher(θ, k, N):
    A, B = θ

    def Pgauss(k, A, B):
      return np.real(np.fft.fftn(np.log(np.fft.fftshift(np.fft.ifftn(P(k, A, B))))))

    C_A,C_B = jax.vmap(jax.grad(Pgauss, argnums=(1,2)))(k.flatten(),
                                        np.repeat(A, len(k.flatten())),
                                        np.repeat(B, len(k.flatten())))

    C_A = C_A.reshape(k.shape)
    C_B = C_B.reshape(k.shape)

    pgauss = Pgauss(k, A, B)[1:N // 2, 1:N // 2]
    k = k[1:N // 2, 1:N // 2]

    Cinv = np.diag(1. / pgauss.flatten())

    C_A = np.diag(C_A[1:N // 2, 1:N // 2].flatten())
    C_B = np.diag(C_B[1:N // 2, 1:N // 2].flatten())

    F_AA = 0.5 * np.trace((Cinv @ C_A @ Cinv @ C_A))
    F_AB = 0.5 * np.trace((Cinv @ C_B @ Cinv @ C_A ))
    F_BA = 0.5 * np.trace((Cinv @ C_A @ Cinv @ C_B))
    F_BB = 0.5 * np.trace((Cinv @ C_B @ Cinv @ C_B))

    return np.array([[F_AA, F_AB], [F_BA, F_BB]])
```

We can now set the simulator arguments, i.e. the $k$-modes to evaluate, the length of the side of a box, the shape of the box and whether to normalise via the volume and squeeze the output dimensions

## choose $k$-modes and define fiducial model
Next, we're going to set our $N$-side to 128 (the size of our data vector), $k$-vector, as well as the $L$-side (the physical dimensions of the universe-in-a-box:

```python
N = 128
shape = (N, N)

k = np.sqrt(
    np.sum(
        np.array(
            np.meshgrid(
                *((np.hstack(
                    (np.arange(0, _shape // 2 + 1),
                     np.arange(-_shape // 2 + 1, 0)))
                   * 2 * np.pi / _shape)**2.
                  for _shape in shape))),
        axis=0))
```

Now that we have the setup, we can compute the theoretical Fisher information for this cosmological power spectrum:

```python
F = ln_fisher(θ_fid, k, int(N))
```

```python
print('theoretical det F:', np.linalg.det(F))
print('theoretical Shannon information content in nats: ', 0.5*np.log(np.linalg.det(F)))
```

    theoretical det F: 245524.33983948187
    theoretical Shannon information content in nats:  6.205575682625643

### Next, we provide you our universe simulator in `jax`. This is how it works:

### 2D random field simulator in jax

To create a 2D lognormal random field we can follow these steps:

1. Generate a $(N_\textrm{pix}\times N_\textrm{pix})$ white noise field $\varphi$ such that $\langle \varphi_k \varphi_{-k} \rangle' = 1$

2. Fourier Transform $\varphi$ to real space: $R_{\rm white}({\bf x}) \rightarrow R_{\rm white}({\bf k})$
   Note that NumPy's DFT Fourier convention is:
   $$\phi_{ab}^{\bf k} = \sum_{c,d = 0}^{N-1} \exp{(-i x_c k_a - i x_d k_b) \phi^{\bf x}_{cd}}$$
   $$\phi_{ab}^{\bf x} = \frac{1}{N^2}\sum_{c,d = 0}^{N-1} \exp{(-i x_c k_a - i x_d k_b) \phi^{\bf k}_{cd}}$$

3. Evaluate the chosen power spectrum over a field of $k$ values and do the lognormal transformation:
   $$P_{\rm LN}(k) \gets \ln(1 + P(k)) $$
   Here we need to ensure that this array of amplitudes are Hermitian, e.g. $\phi^{* {\bf k}}_{a(N/2 + b)} = \phi^{{\bf k}}_{a(N/2 - b)}$. This is accomplished by choosing indices $k_a = k_b = \frac{2\pi}{N} (0, \dots, N/2, -N/2+1, \dots, -1)$ (as above) and then evaluating the square root of the outer product of the meshgrid between the two: $k = \sqrt{k^2_a + k^2_b}$. We can then evaluate $P_{\rm LN}^{1/2}(k)$.

4. Scale white noise $R_{\rm white}({\bf k})$ by the power spectrum:
   $$R_P({\bf k}) = P_{\rm LN}^{1/2}(k) R_{\rm white}({\bf k}) $$

5. Fourier Transform $R_{P}({\bf k})$ to real space: $R_P({\bf x}) = \int d^d \tilde{k} e^{i{\bf k} \cdot {\bf x}} R_p({\bf k})$
   $$R_{ab}^{\bf x} = \frac{1}{N^2}\sum_{c,d = 0}^{N-1} \exp{(-i x_c k_a - i x_d k_b) R^{\bf k}_{cd}}$$

We are going to use a broadcastable jax simultor which takes in a variety of different shaped parameter arrays and vmaps them until a single parameter pair are passed. This is very efficient for generating many simulations at once, for Approximate Bayesian Computation for example.

```python
#@title simulator code <font color='lightgreen'>[RUN ME]</font>

def simulator(rng, θ, simulator_args, foregrounds=None):
    def fn(rng, A, B):
        dim = len(simulator_args["shape"])
        L = simulator_args["L"]
        if np.isscalar(L):
            L = [L] * int(dim)
        Lk = ()
        shape = ()
        for i, _shape in enumerate(simulator_args["shape"]):
            Lk += (_shape / L[i],)
            if _shape % 2 == 0:
                shape += (_shape + 1,)
            else:
                shape += (_shape,)

        k = simulator_args["k"]
        k_shape = k.shape
        k = k.flatten()[1:]
        tpl = ()
        for _d in range(dim):
            tpl += (_d,)

        V = np.prod(np.array(L))
        scale = V**(1. / dim)
        fft_norm = np.prod(np.array(Lk))

        rng, key = jax.random.split(rng)

        mag = jax.random.normal(
            key, shape=shape)
        pha = 2. * np.pi * jax.random.uniform(
            key, shape=shape)

        # now make hermitian field (reality condition)
        revidx = (slice(None, None, -1),) * dim
        mag = (mag + mag[revidx]) / np.sqrt(2)
        pha = (pha - pha[revidx]) / 2 + np.pi
        dk = mag * (np.cos(pha) + 1j * np.sin(pha))
        cutidx = (slice(None, -1),) * dim
        dk = dk[cutidx]

        powers = np.concatenate(
            (np.zeros(1),
             np.sqrt(P(k, A=A, B=B)))).reshape(k_shape)

        if simulator_args['vol_norm']:
            powers /= V

        if simulator_args["log_normal"]:
            powers = np.real(
                np.fft.ifftshift(
                    np.fft.ifftn(
                        powers)
                    * fft_norm) * V)

            powers = np.log(1. + powers)
            powers = np.abs(np.fft.fftn(powers))

        fourier_field = powers * dk
        fourier_field = jax.ops.index_update(
            fourier_field,
            np.zeros(dim, dtype=int),
            np.zeros((1,)))

        if simulator_args["log_normal"]:
            field = np.real(np.fft.ifftn(fourier_field)) * fft_norm * np.sqrt(V)
            sg = np.var(field)
            field = np.exp(field - sg / 2.) - 1.

        else:
            field = np.real(np.fft.ifftn(fourier_field) * fft_norm * np.sqrt(V)**2)

        if simulator_args["N_scale"]:
            field *= scale

        if foregrounds is not None:
            rng, key = jax.random.split(key)
            foreground = foregrounds[
                jax.random.randint(
                    key,
                    minval=0,
                    maxval=foregrounds.shape[0],
                    shape=())]
            field = np.expand_dims(field + foreground, (0,))

        if not simulator_args["squeeze"]:
            field = np.expand_dims(field, (0, -1))

        return np.array(field, dtype='float32')

    if isinstance(θ, tuple):
        A, B = θ
    else:
        A = np.take(θ, 0, axis=-1)
        B = np.take(θ, 1, axis=-1)
    if A.shape == B.shape:
        if len(A.shape) == 0:
            return fn(rng, A, B)
        else:
            keys = jax.random.split(rng, num=A.shape[0] + 1)
            rng = keys[0]
            keys = keys[1:]
            return jax.vmap(
                lambda key, A, B: simulator(
                    key, (A, B), simulator_args=simulator_args))(
                keys, A, B)
    else:
        if len(A.shape) > 0:
            keys = jax.random.split(rng, num=A.shape[0] + 1)
            rng = keys[0]
            keys = keys[1:]
            return jax.vmap(
                lambda key, A: simulator(
                    key, (A, B), simulator_args=simulator_args))(
                keys, A)
        elif len(B.shape) > 0:
            keys = jax.random.split(rng, num=B.shape[0])
            return jax.vmap(
                lambda key, B: simulator(
                    key, (A, B), simulator_args=simulator_args))(
                keys, B)
```

```python
simulator_args = dict(
    k=k,               # k-vector (grid units)
    L=250,             # in Mpc h^-1
    shape=shape,
    vol_norm=True,     # whether to normalise P(k) by volume
    N_scale=False,     # scale field values up or down
    squeeze=True,
    log_normal=True)
```

By constructing our random field simulator *and* cosmological power spectrum in `Jax`, we have access to *exact numerical derivatives*, meaning we can simulate a *differentiable* universe. Let's visualize what our universe and derivatives look like at our fiducial model below:

```python
#@title visualize a fiducial universe and gradients <font color='lightgreen'>[run me]</font>
from imnn.utils import value_and_jacrev, value_and_jacfwd

def simulator_gradient(rng, θ, simulator_args=simulator_args):
    return value_and_jacrev(simulator, argnums=1, allow_int=True, holomorphic=True)(rng, θ, simulator_args=simulator_args)

simulation, simulation_gradient = value_and_jacfwd(simulator, argnums=1)(rng, θ_fid,
                                                        simulator_args=simulator_args)
cmap = 'viridis'
from mpl_toolkits.axes_grid1 import make_axes_locatable

fig,ax = plt.subplots(nrows=1, ncols=3, figsize=(12,15))

im1 = ax[0].imshow(np.squeeze(simulation),
                   extent=(0,1,0,1), cmap=cmap)
ax[0].title.set_text(r'example fiducial $\rm d$')
divider = make_axes_locatable(ax[0])
cax = divider.append_axes('right', size='5%', pad=0.05)
fig.colorbar(im1, cax=cax, orientation='vertical')

im1 = ax[1].imshow(np.squeeze(simulation_gradient).T[0].T,
                   extent=(0,1,0,1), cmap=cmap)
ax[1].title.set_text(r'$\nabla_{\Omega_m} \rm d$')
divider = make_axes_locatable(ax[1])
cax = divider.append_axes('right', size='5%', pad=0.05)
fig.colorbar(im1, cax=cax, orientation='vertical')

im1 = ax[2].imshow(np.squeeze(simulation_gradient).T[1].T,
                   extent=(0,1,0,1), cmap=cmap)
ax[2].title.set_text(r'$\nabla_{\sigma_8} \rm d$')
divider = make_axes_locatable(ax[2])
cax = divider.append_axes('right', size='5%', pad=0.05)
fig.colorbar(im1, cax=cax, orientation='vertical')

for a in ax:
    a.set_xticks([])
    a.set_yticks([])

plt.show()
```

Nice ! Since we can differentiate our universe and power spectrum, we can also compute the theoretical Fisher information embedded in the$128\times128$ field at the fiducial model:

---
## Training an IMNN

$\quad$

The details behind the IMNN algorithm [can be found here on arxiv](https://arxiv.org/abs/1802.03537), but we'll summarize the gist briefly:

1.   We want to maximise the Fisher information, $\textbf{F}$, of compressed summaries to satisfy the Cramer-Rao bound:
    $$ \langle (\vartheta_\alpha - \langle \vartheta_\alpha \rangle )      (\vartheta_\beta - \langle \vartheta_\beta
    \rangle) \rangle \geq \textbf{F}^{-1}_{\alpha \beta} $$  which means saturating the Fisher information minimizes the average variance of the parameter estimates.

2.   To do this, we compute a pseudolikelihood to compute our Fisher information:
$$     -2 \ln \mathcal{L}(\textbf{x} | \textbf{d}) = (\textbf{x} - \boldsymbol{\mu}_f(\vartheta))^T \textbf{C}_f^{-1}(\textbf{x} - \boldsymbol{\mu}_f(\vartheta)) $$ where $\boldsymbol{\mu}_f$ and $\textbf{C}$ are the mean and covariance of the network output (summaries). The Fisher is then $$ \textbf{F}_{\alpha \beta} =  {\rm tr} [\boldsymbol{\mu}_{f,\alpha}^T C^{-1}_f \boldsymbol{\mu}_{f, \beta}] $$

Since we can differentiate through our neural network *and* simulated universe, we have the exact derivatives with respect to the pipeline we need to compute the Fisher matrix of compressed summaries on-the-fly during compression training.
___

### Q: wait -- what if my simulator isn't differentiable ?
We don't *need* to have the exact derivatives for IMNN training ! Having the gradients accessible just means that we don't have to optimize finite-differencing for estimating derivatives by hand, however (as is done in the original IMNN paper).

___

Let's train an IMNN on cosmological fields to see how much information we can extract an what sort of constraints we can get. We will use 2000 simulations to estimate the covariance and use all of their derivatives and we'll summarise the whole cosmological field using 2 summaries.

```python
n_s = 200     # number of simulations used to estimate covariance of network outputs
n_d = n_s     # number of simulations used to estimate the numerical derivative of
              # the mean of the network outputs
n_summaries = 2
```

We're going to use a fully convolutional inception network built using stax with some custom designed blocks. The inception block itself is implemented in the following block:

```python
#@title nn model stuff <font color='lightgreen'>[RUN ME]</font>
def InceptBlock(filters, strides, do_5x5=True, do_3x3=True):
    """InceptNet convolutional striding block.
    filters: tuple: (f1,f2,f3)
    filters1: for conv1x1
    filters2: for conv1x1,conv3x3
    filters3L for conv1x1,conv5x5"""

    filters1, filters2, filters3 = filters
    conv1x1 = stax.serial(stax.Conv(filters1, (1, 1), strides, padding="SAME"))

    filters4 = filters2
    conv3x3 = stax.serial(stax.Conv(filters2, (1, 1), strides=None, padding="SAME"),
                        stax.Conv(filters4, (3, 3), strides, padding="SAME"))

    filters5 = filters3
    conv5x5 = stax.serial(stax.Conv(filters3, (1, 1), strides=None, padding="SAME"),
                         stax.Conv(filters5, (5, 5), strides, padding="SAME"))

    maxpool = stax.serial(stax.MaxPool((3, 3), padding="SAME"),
                         stax.Conv(filters4, (1, 1), strides, padding="SAME"))

    if do_3x3:
        if do_5x5:
            return stax.serial(
                  stax.FanOut(4),
                  stax.parallel(conv1x1, conv3x3, conv5x5, maxpool),
                  stax.FanInConcat(),
                  stax.LeakyRelu)
        else:
            return stax.serial(
                  stax.FanOut(3),
                  stax.parallel(conv1x1, conv3x3, maxpool),
                  stax.FanInConcat(),
                  stax.LeakyRelu)
    else:
        return stax.serial(
              stax.FanOut(2),
              stax.parallel(conv1x1, maxpool),
              stax.FanInConcat(),
              stax.LeakyRelu)
```

We'll also want to make sure that the output of the network is the correct shape, for which we'll introduce a Reshaping layer

```python
def Reshape(shape):
    """Layer function for a reshape layer."""
    init_fun = lambda rng, input_shape: (shape,())
    apply_fun = lambda params, inputs, **kwargs: np.reshape(inputs, shape)
    return init_fun, apply_fun
```

Now we can build the network, with 64 filters and strides of 4 in each direction in each layer

```python
fs = 55
layers = [
        InceptBlock((fs, fs, fs), strides=(4, 4)),
        InceptBlock((fs, fs, fs), strides=(4, 4)),
        InceptBlock((fs, fs, fs), strides=(4, 4)),
        InceptBlock((fs, fs, fs), strides=(2, 2), do_5x5=False, do_3x3=False),
        stax.Conv(n_summaries, (1, 1), strides=(1, 1), padding="SAME"),
        stax.Flatten,
        Reshape((n_summaries,))
]
model = stax.serial(*layers)
```

We'll also introduce a function to check our model output:

```python
def print_model(layers, input_shape, rng):
    print('input_shape: ', input_shape)
    for l in range(len(layers)):
        _m = stax.serial(*layers[:l+1])
        print('layer %d shape: '%(l+1), _m[0](rng, input_shape)[0])

# print model specs
key,rng = jax.random.split(rng)
input_shape = (1,) + shape + (1,)

print_model(layers, input_shape, rng)
```

    input_shape:  (1, 128, 128, 1)
    layer 1 shape:  (1, 32, 32, 220)
    layer 2 shape:  (1, 8, 8, 220)
    layer 3 shape:  (1, 2, 2, 220)
    layer 4 shape:  (1, 1, 1, 110)
    layer 5 shape:  (1, 1, 1, 2)
    layer 6 shape:  (1, 2)
    layer 7 shape:  (2,)

We'll also grab an adam optimiser from jax.experimental.optimizers

```python
optimiser = optimizers.adam(step_size=1e-3)
```

Note that due to the form of the network we'll want to have simulations that have a "channel" dimension, which we can set up by not allowing for squeezing in the simulator.

### Train an IMNN

Finally we can train an IMNN to maxmise the Fisher information of its summaries and compare its compression efficiency to the theoretical Fisher.

We've used a SimulatorIMNN trained on new simulations on-the-fly.

```python
# initialize IMNN with pre-trained state
rng, key = jax.random.split(rng)
IMNN = imnn.IMNN(
        n_s=n_s,
        n_d=n_d,
        n_params=n_params,
        n_summaries=n_summaries,
        input_shape=(1,) + shape + (1,),
        θ_fid=θ_fid,
        model=model,
        optimiser=optimiser,
        key_or_state=key,  # <---- initialize with state
        simulator=lambda rng, θ: simulator(
            rng, θ, simulator_args={
                **simulator_args,
                **{"squeeze": False}}))
```

    `simulator` provided, using SimulatorIMNN

```python
%%time
rng, key = jax.random.split(rng)
IMNN.fit(λ=10., ϵ=0.1, rng=key, print_rate=None, min_iterations=500, best=True, patience=500)
```

    CPU times: user 41min 24s, sys: 1min 53s, total: 43min 17s
    Wall time: 42min 49s

```python
sns.set()
ax = IMNN.plot(expected_detF=np.linalg.det(F));
ax[0].set_yscale('log')
```

```python
print('det F from IMNN:', np.linalg.det(IMNN.F))
```

    det F from IMNN: 333175.6

```python
print('% Fisher information captured by IMNN: ', np.linalg.det(IMNN.F) / np.linalg.det(F))
```

    % Fisher information captured by IMNN:  1.3569961901448244

```python
print('difference in Shannon information captured by IMNN: ', 0.5*(np.log(np.linalg.det(IMNN.F)) - np.log(np.linalg.det(F))))
```

    difference in Shannon information captured by IMNN:  0.15263678838265804

```python
np.save('cosmo-IMNN-history', IMNN.history['detF'])
np.save('cosmo-IMNN-F', IMNN.F)
np.save('cosmo-IMNN-fid', IMNN.θ_fid)
np.save('cosmo-IMNN-theory-F', F)
np.save('cosmo-IMNN-theory-detF', np.linalg.det(F))
```


# Matplotlib Random Walk project

### First we import the necessary libraries

```python
import matplotlib.pyplot as plt
from random import choice
%matplotlib inline
```

### Then we change the default display settings for our plots

```python
plt.rc('figure',figsize=(12,6))
```

### Now we create x and y values

```python
values = list(range(0,55,5))
values
```

    [0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50]

### And we plot them

```python
plt.plot(values)
plt.show()
```

### We can label the x and y axis and add a title

```python
x_axis = [0.0,0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1.0]
```

```python
plt.plot(x_axis,values)
plt.xlim(0,1.0)
plt.ylim(0,50)
plt.title('Our Plot')
plt.xlabel('This is the horizantal axis')
plt.ylabel('Vertical axis')
plt.show()
```

### We can also change the style of the plot

```python
plt.plot(x_axis,values,'ko--')
plt.xlim(0,1.0)
plt.ylim(0,50)
plt.title('Our Plot')
plt.xlabel('This is the horizantal axis')
plt.ylabel('Vertical axis')
plt.show()
```

### Now let's create a Random Walk

```python
step_option = [1,-1]
step_choice = choice(step_option)
# step_choice
```

```python
walk = []
walk.append(step_choice)
for step in range(1,1000):
    next_step = walk[step-1] + choice(step_option)
    walk.append(next_step)
print(walk)
```

    [1, 0, -1, 0, 1, 0, -1, 0, -1, 0, -1, -2, -3, -2, -1, -2, -3, -4, -5, -6, -5, -4, -3, -4, -3, -4, -3, -4, -5, -4, -3, -4, -3, -2, -3, -2, -3, -2, -1, -2, -1, 0, 1, 2, 3, 4, 3, 4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 5, 4, 5, 4, 3, 2, 3, 2, 1, 0, 1, 0, 1, 2, 1, 2, 3, 4, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 11, 12, 13, 14, 15, 14, 13, 12, 11, 10, 9, 10, 11, 12, 11, 12, 11, 12, 11, 10, 9, 8, 9, 8, 9, 10, 9, 10, 9, 10, 11, 10, 9, 8, 7, 8, 7, 8, 9, 10, 11, 10, 11, 10, 9, 8, 9, 10, 9, 10, 11, 12, 11, 10, 11, 10, 9, 8, 7, 8, 9, 8, 7, 6, 7, 8, 7, 8, 9, 10, 11, 12, 11, 12, 11, 10, 9, 8, 7, 8, 7, 8, 7, 8, 7, 8, 9, 10, 9, 10, 11, 10, 11, 10, 9, 8, 7, 8, 9, 8, 7, 6, 7, 8, 9, 8, 7, 6, 5, 6, 7, 6, 5, 6, 5, 4, 3, 4, 3, 4, 3, 2, 3, 2, 1, 2, 3, 2, 1, 2, 3, 4, 5, 6, 7, 6, 7, 8, 9, 8, 7, 6, 5, 4, 3, 2, 1, 2, 3, 2, 1, 0, -1, -2, -3, -2, -3, -2, -3, -4, -5, -4, -3, -4, -5, -6, -5, -4, -5, -4, -3, -4, -3, -2, -1, 0, -1, 0, 1, 2, 1, 2, 1, 2, 1, 0, 1, 2, 1, 2, 3, 2, 3, 4, 5, 6, 5, 4, 5, 6, 7, 6, 7, 6, 5, 4, 3, 4, 5, 6, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 8, 7, 6, 5, 6, 7, 8, 9, 8, 9, 10, 9, 8, 9, 10, 9, 8, 9, 8, 7, 8, 7, 8, 9, 10, 9, 10, 9, 10, 11, 12, 13, 12, 13, 12, 13, 14, 13, 12, 13, 14, 15, 14, 13, 14, 15, 16, 15, 16, 17, 16, 17, 16, 15, 14, 13, 14, 15, 16, 17, 16, 17, 16, 17, 18, 19, 20, 19, 18, 17, 18, 19, 18, 19, 18, 17, 18, 19, 18, 17, 18, 17, 16, 17, 18, 19, 18, 19, 20, 19, 20, 19, 18, 19, 20, 21, 20, 21, 22, 23, 24, 23, 24, 23, 24, 23, 24, 23, 24, 23, 24, 23, 24, 23, 24, 23, 24, 25, 26, 27, 26, 25, 24, 25, 26, 25, 26, 27, 26, 27, 26, 27, 28, 29, 28, 29, 30, 29, 28, 29, 28, 29, 28, 29, 30, 29, 30, 29, 30, 29, 30, 29, 28, 27, 26, 27, 28, 29, 28, 27, 28, 27, 26, 25, 24, 25, 24, 23, 24, 25, 24, 23, 22, 23, 24, 23, 24, 25, 24, 25, 26, 25, 24, 23, 22, 21, 22, 23, 22, 23, 24, 23, 22, 21, 22, 23, 24, 25, 24, 25, 24, 23, 24, 23, 22, 23, 22, 23, 22, 23, 24, 23, 24, 23, 22, 21, 22, 23, 22, 21, 22, 23, 22, 21, 22, 21, 20, 21, 22, 23, 24, 25, 26, 27, 28, 27, 26, 27, 28, 27, 28, 29, 28, 29, 28, 27, 28, 27, 26, 25, 26, 25, 24, 25, 26, 25, 24, 23, 24, 25, 24, 23, 22, 21, 22, 21, 22, 23, 24, 25, 24, 25, 24, 23, 24, 23, 22, 23, 24, 25, 26, 27, 26, 27, 28, 29, 30, 31, 32, 31, 32, 33, 32, 33, 32, 33, 32, 31, 30, 31, 32, 31, 32, 31, 32, 31, 32, 31, 32, 33, 32, 31, 30, 29, 30, 31, 30, 31, 32, 31, 32, 31, 30, 29, 28, 27, 26, 25, 24, 23, 24, 23, 24, 23, 22, 23, 24, 23, 22, 23, 22, 23, 24, 25, 26, 25, 24, 25, 26, 25, 24, 23, 24, 25, 26, 25, 26, 25, 24, 25, 24, 23, 24, 25, 26, 25, 24, 25, 24, 25, 24, 23, 24, 23, 22, 23, 22, 21, 20, 21, 20, 19, 20, 19, 18, 19, 18, 19, 20, 19, 18, 17, 18, 17, 18, 19, 18, 19, 18, 19, 18, 19, 20, 19, 18, 19, 20, 21, 20, 21, 20, 19, 20, 21, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 29, 30, 31, 30, 29, 28, 29, 30, 29, 28, 27, 28, 29, 28, 29, 28, 27, 26, 27, 26, 25, 24, 25, 24, 25, 26, 27, 26, 27, 26, 25, 24, 23, 22, 23, 24, 25, 26, 25, 26, 25, 24, 25, 26, 27, 26, 27, 26, 25, 24, 25, 26, 25, 24, 25, 24, 23, 22, 21, 20, 21, 22, 23, 22, 23, 24, 25, 24, 25, 24, 25, 24, 25, 24, 23, 22, 23, 24, 23, 24, 23, 22, 23, 24, 25, 24, 23, 24, 23, 24, 25, 24, 25, 24, 23, 24, 25, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 33, 34, 35, 34, 33, 34, 35, 34, 35, 34, 35, 34, 33, 32, 31, 30, 31, 32, 31, 32, 33, 32, 31, 30, 31, 30, 29, 30, 31, 32, 31, 32, 33, 32, 33, 32, 31, 30, 29, 30, 31, 32, 31, 32, 33, 34, 35, 36, 35, 34, 33, 34, 33, 32, 31, 32, 33, 34, 33, 32, 31, 32, 33, 34, 35, 34, 33, 34, 33, 34, 35, 34, 33, 34, 33, 32, 31, 30, 29, 30, 31, 32, 31, 32, 31, 32, 33, 32, 33, 34, 33, 32, 31, 30, 29, 30, 29, 28, 29, 28, 29, 28, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 35, 34, 35, 36, 37, 36, 37, 36, 35, 36, 35, 36, 37, 38, 37, 36, 35, 36, 37, 38, 37, 36, 35, 36, 35, 36, 35, 36, 35, 36, 35, 34, 33, 34, 33, 34, 33, 32, 31, 32, 31, 30, 29, 30]

```python
plt.plot(walk)
plt.show()
```

### Using functions would be better

```python
def rand_walk(step_num):
    walk = []
    step_choice = choice([1,-1])
    walk.append(step_choice)
    for i in range(1,step_num):
        step_choice_2 = choice([1,-1])
        next_step = walk[i-1] + step_choice_2
        walk.append(next_step)
    return walk
```

```python
random_w = rand_walk(100)
print(random_w)
```

    [-1, 0, 1, 0, -1, 0, 1, 2, 3, 4, 5, 4, 5, 6, 5, 4, 5, 4, 5, 4, 3, 2, 1, 2, 1, 0, -1, 0, 1, 0, 1, 0, -1, -2, -1, 0, 1, 2, 3, 2, 1, 2, 3, 2, 1, 0, 1, 2, 3, 4, 5, 4, 3, 4, 3, 2, 1, 0, 1, 2, 3, 2, 1, 0, -1, 0, 1, 0, 1, 2, 3, 4, 3, 2, 3, 4, 3, 2, 1, 2, 3, 2, 3, 4, 3, 2, 3, 4, 3, 2, 1, 2, 1, 2, 3, 2, 1, 2, 1, 2]

```python
def plot_walk(walk):

    plt.plot(walk)
    plt.xlabel('Number of Steps')
    plt.ylabel('Distance from Origin')
    plt.title('Our Random Walk')
    plt.show()
```

```python
plot_walk(rand_walk(1000))

```

### We can create a function to plot more than 1 random walk

```python
def plot_walks(num,walk_steps):

    lab_list = list(range(1,num+1))
    for i in range(0,num):

        x = list(range(1,walk_steps+1))

        plt.plot(x,rand_walk(walk_steps),label = 'Plot Number' + ' ' + str(lab_list[i]))
        plt.xlabel('Number of Steps')
        plt.ylabel('Distance from Origin')
        plt.title('Our Random Walk')
        plt.legend(loc='lower left')
    plt.show()
```

```python
plot_walks(5,100000)
```

# Sierpinski Triangle

First Transformation
\begin{equation}
x_{n+1} = 0.5x_n
\end{equation}
\begin{equation}
y_{n+1} = 0.5y_n
\end{equation}
Second Transformation
\begin{equation}
x_{n+1} = 0.5x_n + 0.5
\end{equation}
\begin{equation}
y_{n+1} = 0.5y_n + 0.5
\end{equation}
Third Transformation
\begin{equation}
x_{n+1} = 0.5x_n + 1
\end{equation}
\begin{equation}
y_{n+1} = 0.5y_n
\end{equation}

```python
def trans_1(p):
    x = p[0]
    y = p[1]
    x1 = 0.5 * x
    y1 = 0.5 * y
    return x1,y1

def trans_2(p):
    x = p[0]
    y = p[1]
    x1 = 0.5 * x + 0.5
    y1 = 0.5 * y + 0.5
    return x1,y1

def trans_3(p):
    x = p[0]
    y = p[1]
    x1 = 0.5 * x + 1
    y1 = 0.5 * y
    return x1,y1

transformations = [trans_1,trans_2,trans_3]
a1 = [0]
b1 = [0]
a,b = 0,0

for i in range(1000000):
    trans = choice(transformations)
    a,b = trans((a,b))
    a1.append(a)
    b1.append(b)
```

```python
plt.rc('figure',figsize=(16,16))
```

```python
plt.plot(a1,b1,'o')
plt.savefig('my_figure.png')
```

```python
%magic
```

```python

```

```python

```


# Two Degree-of-Freedom Caldera Model

## Introduction

Dynamical matching is an interesting chemical dynamical phenomenon that occurs in a variety of  organic chemical reactions. A caldera PES  arises  in many organic chemical reactions, such as the  vinylcyclopropane-cyclopentene rearrangement \cite{baldwin2003,gold1988}, the stereomutation of cyclopropane \cite{doubleday1997}, the degenerate rearrangement of bicyclo[3.1.0]hex-2-ene \cite{doubleday1999,doubleday2006} or that of 5-methylenebicyclo[2.1.0]pentane \cite{reyes2002}. It is characterized by a flat region or shallow minimum at its center  surrounded by potential walls and multiple symmetry related index one saddle points that allow entrance and exit from this intermediate region. This shape of the potential resembles  the collapsed  region within an errupted volcano (caldera), and this is the reason that Doering \cite{doering2002} and co-workers  refer to this type of  potential  as a caldera.

The manifestation of the dynamical matching phenomenon is essentially a statement of  momentum conservation and Newton's first law of motion. It is observed that a trajectory entering the Caldera from a channel corresponding to a high energy index-1 saddle (reactant) experiences little force in the caldera due to the ''flatness'' of the PES, and it exits through the diametrically opposing low energy index-1 saddle (product). Consequently, this mechanism determines to a considerable extent the outcome of the chemical reaction. However, not all trajectories entering the caldera experience dynamical matching. It is observed that some trajectories may interact with the shallow potential well and become temporarily trapped in the region. This can dramatically influence the manner in which they exit from the well.

A detailed study of the trajectory behavior in a two DoF caldera PES was given in \cite{collins2014}, where a more general discussion of caldera-like PESs in organic reactions is also presented. Further work elucidating the phenomena of dynamical matching and trapping  in this caldera model was carried out in \cite{ katsanikas2018,katsanikas2019}. We will describe the results in these papers in more detail when we describe the Hamiltonian model in the next section.

In this chapter we describe the phase space mechanism that controls dynamical matching. We show that  dynamical matching is controlled by a heteroclinic intersection between the unstable manifold of a periodic orbit controlling entrance to the caldera and the stable manifold of a periodic orbit  in the region of the shallow minimum.  When a heteroclinic connection exists, trajectories that enter the caldera are transported to the shallow minimum, and they experience temporary trapping in this region. When there is no heteroclinic connection, trajectories enter and exit the caldera without interacting with the  region of the central minimum.  Knowledge of this phase space mechanism is significant because it allows us to predict existence, and non-existence, of dynamical matching.

This chapter is outlined as follows. In Section \ref{sec:model} we describe the Caldera PES that we use in this work, its critical points and stability, and the resulting Hamiltonian model. Section \ref{sec:DM_mech} is devoted to analyzing the phase space structures that govern dynamical matching.

## Development of the Problem
\label{sec:model}

We give a brief description of the caldera potential energy surface (Fig. \ref{fig:equi1}) and Hamiltonian as described in \cite{collins2014}. The caldera potential has a stable equilibrium point at the center, referred to as the central minimum. This potential has an axis of symmetry, the y-axis. We have also the existence of potential walls around the central minimum. On these potential walls we encounter four 1-index saddles (two for lower values of energy, referred to as the lower saddles,  and two for  higher values of energy, referred to as the upper saddles).  In this chapter we consider the stretched version of the caldera potential:

\begin{equation*}
\label{eq1}
V(x,y)=c_1(y^2+(\lambda x)^2) + c_2y - c_3((\lambda x)^4 + y^4 - 6 (\lambda x)^2 y^2)
\end{equation*}

The potential parameters are  $c_1=5$, $c_2=3$,$c_3=-3/10$ and $0<\lambda \leq 1$. The classical symmetric caldera PES \cite{collins2014,katsanikas2018} corresponds to $\lambda = 1$ and is shown in the upper left hand panel of Fig. \ref{fig:equi}.

The Hamiltonian for the system with two DoF is the sum of kinetic plus potential energy:

\begin{equation}
H(x,y,p_x,p_y) = \frac{p_x^2}{2m_x} + \frac{p_y^2}{2m_y} + V(x,y)
\label{eq2}
\end{equation}

where $V(x,y)$ is the Caldera PES in Eq. \eqref{eq1}, and $m_x$, $m_y$ are the masses of the $x$ and $y$ DoF respectively. In this work, for simplicity, we take $m_x = m_y = 1$. Hamilton's equations of motion are given by:

\begin{equation}
\begin{cases}
\dot x = \dfrac{\partial H} {\partial p_x} = \dfrac{p_x}{m_x} \\[.4cm]
\dot y = \dfrac{\partial H} {\partial p_y} = \dfrac{p_y}{m_y} \\[.4cm]
\dot p_x = -\dfrac{\partial H} {\partial x} = 2 \lambda \, (\lambda x) \left[2c_3 \left((\lambda x)^2 - 3 y^2 \right) - c_1 \right] \\[.4cm]
\dot p_y = -\dfrac {\partial H} {\partial y} = 2 y \left[ 2 c_3 \left(y^2 - 3 (\lambda x)^2\right) - c_1 \right] - c_2
\end{cases}
\label{eq3}
\end{equation}

## Revealing the Phase Space Structures

In Fig. \ref{fig:equi} we show the contours and the equilibrium points of the potential for different values of $\lambda$, for example $\lambda=0.8$, $\lambda=0.6$ and $\lambda=0.2$. Table \ref{tab:ta08} summarises the positions and energies of the upper index-1 saddles for different values of $\lambda$. We observe that the positions of the index-1 saddles move away from the center of the Caldera as we decrease the parameter $\lambda$.  The position of the central minimum is $(x,y) = (0,-0.297)$ with energy $E = -0.448$ for all values of the stretching parameter $\lambda$.

\label{fig:equi1}
\caption{Caldera potential energy surface for $\lambda=1$.}

\label{fig:equi}
\caption{The stable equilibrium point in the center (depicted by a black point), the upper saddles (depicted by red points), the lower saddles (depicted by blue points) and the  contours of the potential for $\lambda = 1$ (upper left panel), $\lambda = 0.8$ (upper right panel), $\lambda = 0.6$ (lower left panel) and  $\lambda = 0.2$ (lower right panel).}

| Equilibrium point       | x        | y       | $\lambda$ |
|-------------------------|----------|---------|-----      |
| Upper LH index-1 saddle | -2.149   |  2.0778 | 1         |
| Upper RH index-1 saddle |  2.149   |  2.0778 | 1         |
| Upper LH index-1 saddle | -2.6862  |  2.0778 | 0.8       |
| Upper RH index-1 saddle | 2.6862   |  2.0778 | 0.8       |
| Upper LH index-1 saddle | -3.5815  |  2.0778 | 0.6       |
| Upper RH index-1 saddle | 3.5815   |  2.0778 | 0.6       |
| Upper LH index-1 saddle | -10.7446 |  2.0778 | 0.2       |
| Upper RH index-1 saddle | 10.7446  |  2.0778 | 0.2       |

\label{tab:ta08}
\caption{The upper index-1 saddles of the PES in Eq. \ref{eq1} ("RH" and "LH" are the abbreviations for right hand and left hand respectively) for different values of $\lambda$. The energy for each of the cases is $E = 27.0123$.}

### Lagrangian Descriptors for revealing Phase Space Structures

In order to reveal the phase space structures that are responsible for the mechanism that allows and prevents dynamical matching, we use in this work the method of Lagrangian descriptors (LDs), see e.g. \cite{mancho2013lagrangian,lopesino2017,naik2019a}. Lagrangian descriptors is a trajectory-based scalar diagnostic that has been developed in the nonlinear dynamics literature to explore the geometrical template of phase space structures that characterizes qualitatively distinct dynamical behavior. Details on how they are applied for revealing phase space structures in caldera-like PESs are described in \cite{KGW2019,KGW2019a}. In this chapter we focus on presenting the results relevant to dynamical matching.

## Implications for Reaction Dynamics
\label{sec:DM_mech}

As we have described in the introduction, the caldera gets its name from the shape of the PES. However, transport across the caldera is a dynamical phenomenon governed by the template of geometrical structures in phase space, and dynamical matching is just one particular type of  dynamical phenomenon that we are considering in this chapter. First, we describe the phase space structures that mediate transport into the caldera.

For a two DoF system, the fixed energy surface is three dimensional. For energies above that of the upper saddles  an unstable periodic orbit exists in the energy surface. This is a consequence of the Lyapunov subcenter manifold theorem \cite{moser1976, weinstein1973, rabinowitz1982}. In a fixed energy surface, these periodic orbits have two dimensional stable and unstable manifolds. Trajectories move away from the periodic orbits along the direction of the unstable manifold in forward time. In the upper left panel of Fig. \ref{fig:fig_panel} we show a segment of the unstable manifold of the upper right-hand saddle directed towards the interior of the caldera.

The region of the central minimum of the caldera may also contain unstable periodic orbits. The stable manifolds of these periodic orbits direct trajectories towards the central minimum. In the upper left panel of Fig. \ref{fig:fig_panel} we show a segment of the stable manifold of an unstable periodic orbit in the region of the central minimum directed away from the central minimum.

If the stable manifold of a periodic orbit in the central minimum intersects the unstable manifolds of one of the  upper  saddles we have a mechanism for trajectories to enter the caldera and be directed towards the region of the central minimum. In dynamical systems terminology this is referred to as a heteroclinic connection. This would inhibit dynamical matching, as trajectories entering the caldera would exhibit (temporary) trapping in the region of the central minimum. If the heteroclinic connection breaks, as might occur if a parameter is varied, the mechanism for directing trajectories towards the regions of the central minimum no longer exists, and dynamical matching is possible.  Hence, a heteroclinic bifurcation is the critical phase space structure that inhibits or allows dynamical matching, which we now present.

In order to explore the formation of a heteroclinic intersection between any stable manifold coming from an UPO of the central region of the Caldera and the unstable manifold of the UPO of the upper-right index-1 saddle, as the stretching parameter of the Caldera PES is varied, we probe the phase space structures in the following Poincare surface of section:

\begin{equation}
\mathcal{U}^{+}_{x,p_x} = \lbrace (x,y,p_x,p_y) \in \mathbb{R}^4 \;|\; y = 1.88409 \; ,\; p_y > 0 \;,\; E = 29 \rbrace
\label{psos}
\end{equation}

In the middle-left panel of Fig. \ref{fig:fig_panel}, we observe that there is a critical value of the stretching parameter ($\lambda=0.778$) for the formation of this heteroclinic connection. For values of the stretching parameter above the critical value there is no  heteroclinic connection between any stable manifold coming from an UPO of the central region of the Caldera and the unstable manifold of the UPO of the upper index-1 saddle (see the upper left   panel of Fig. \ref{fig:fig_panel}). The non-existence of these  heteroclinic connections results in the phenomenon of dynamical matching. In this case, if we integrate an initial condition inside the region of the unstable manifold of UPO of the upper-right index-1 saddle forward and backward in time, we see in the upper right panel of  Fig. \ref{fig:fig_panel} that the resulting trajectory comes from the region of the upper-right index-1 saddle and exits the caldera through the region of the opposite lower saddle without any interaction with the central area of the caldera.

Now, for values of the stretching parameter equal or above the critical value we have the formation of heteroclinic connections between  the stable manifold coming from an UPO of the central region of the Caldera and the unstable manifold of the UPO of the upper-right index-1 saddle, (see  middle and lower left panels of Fig. \ref{fig:fig_panel}). This heteroclinic connection destroys the dynamical matching mechanism because many trajectories become trapped inside the lobes between the two invariant manifolds. We can see this better if we choose an initial condition inside a lobe, as we illustrate in the middle and lower left panels of Fig. \ref{fig:fig_panel} and integrate it forward and backward. We observe that the resulting trajectory is temporarily trapped in the central area of the caldera before it exits from this area, see the middle and lower right panels of Fig. \ref{fig:fig_panel}.

\label{fig:fig_panel}
\caption{Phase space structures calculated on the Poincar\'e section described in Eq. \eqref{psos} located in the vicinity of the UPO of the upper-right index-1 saddle. We have similar structures in the vicinity of the UPO of the upper-left index-1 saddle because of the symmetry of the potential. We illustrate the formation of an heteroclinic connection, as the stretching parameter of the Caldera PES is varied, between a stable manifold (blue curve) of an UPO of the central region of the Caldera and the unstable manifold (red curve) of the UPO associated to the upper index-1 saddle. The first row corresponds to $\lambda = 0.8$, the second row is for the critical stretching value $\lambda = 0.778$, and for the third row we use $\lambda = 0.7$. In the second column, we capture the inhibition of dynamical matching by depicting the projection onto configuration space of the trajectory of two initial conditions marked as a yellow dot (outside the lobe - in all rows) and a yellow diamond (inside the lobe - in the second and third row). In the first row, forward and backward evolution of the trajectory  are represented in black and green respectively. In the second and third row, the red line indicate the part of the trajectories at backward integration that correspond to both of them, circle and diamond. In addition, the black and blue line indicate the part of the trajectories at forward integration that correspond to the circle and diamond respectively. The magenta curve represents the energy boundary in all rows.}

```python

```


Trusted Notebook" width="250 px" align="left">

## _*Quantum Process Tomography*_

The latest version of this notebook is available on https://github.com/QISKit/qiskit-tutorial.

***
### Contributors
Christopher J. Wood$^{1}$

1. IBM T.J. Watson Research Center, Yorktown Heights, NY, USA

## Introduction

In this notebook we demonstrate how to design and run experiments to perform quantum process tomography using QISKit, and demonstrate this using both simulators, and the IBM Quantum Experience. This is implemented similarly to the [quantum state tomography](state_tomography.ipynb) example notebook [here](state_tomography.ipynb).

The functions used in this notebook to implement state tomography may be imported from the **`tools.tomography`** module.

#### Reconstruction Methods

State reconstruction may be done using a variety of methods. In this notebook, we implement two simple cases:
* Least squares fitting
* Fast maximum likelihood (*see [J Smolin, JM Gambetta, G Smith, Phys. Rev. Lett. 108, 070502](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.108.070502), open access: [arXiv:1106.5458](https://arxiv.org/abs/1106.5458)*).

Note that while the fast maximum likelihood method constrains the resulting channel to be positive, it does not enfore complete positivity of the final estimate.

```python
# Checking the version of PYTHON; we only support > 3.5
import sys
if sys.version_info < (3,5):
    raise Exception('Please use Python version 3.5 or greater.')

import numpy as np

# importing the QISKit
from qiskit import QuantumCircuit, QuantumProgram
import Qconfig

# import tomography libary
import qiskit.tools.qcvv.tomography as tomo

# useful additional packages
from qiskit.tools.visualization import plot_state
from qiskit.tools.qi.qi import *
```

## Building process tomography measurement circuits

To perform quantum process tomography of a n-qubit quantum operation we must construct a family of circuits which implement a tomographically complete set of preparations and measurements of each of those qubits. The **`tools.tomography.py`** module contains functions to generate these measurement circuits for general n-qubit systems.

#### Specifying the qubits, preparation basis, and measurement basis

A process tomography experiment is specified by a **`process_tomography_set`** object:
```python
tomo_set = tomo.process_tomography_set(qubits, meas_basis, prep_basis)
```
where
- `qubits` is a list of the qubit indexes within a register (eg. [0, 1] for the first 2 qubits)
- `meas_basis` is the basis to measure each qubit in. The default value is `meas_basis='Pauli'`.
- `prep_basis` is the basis to measure each qubit in. The default value is `prep_basis='SIC'`.

This contains information about which qubits we are constructing the quantum operation for, and what bases to prepare and measure these qubits in. The default preparation basis corresponds to preparing each qubit in each of the projectors of a symmetric informationally complete positive operator valued measurement (SIC-POVM):
$$\begin{align}
\left|S_0\right\rangle & = \left|0\right\rangle \\
\left|S_1\right\rangle & = \frac{1}{\sqrt{3}}\left|0\right\rangle
                            + \sqrt{\frac{2}{3}}\left|1\right\rangle \\
\left|S_2\right\rangle & = \frac{1}{\sqrt{3}}\left|0\right\rangle
                            + e^{-2i \pi/3}\sqrt{\frac{2}{3}}\left|1\right\rangle \\
\left|S_3\right\rangle & = \frac{1}{\sqrt{3}}\left|0\right\rangle
                            + e^{2i \pi/3}\sqrt{\frac{2}{3}}\left|1\right\rangle.
\end{align}$$
The default measurement basis corresponds to performing quantum state tomography on each input state by measuring each qubit in the $X$, $Y$, and $Z$ Pauli bases.

For $n$-qubits this gives a total of $4^n3^{n}$ circuits that must be run. Hence it is not recommended to perform process tomography for more than 2 or 3 qubits.

#### Adding preparation and measurement circuits to a Quantum Program

Once a process tomography set has been defined the function to append the tomography circuits to a `QuantumProgram` is:
```python
tomo_labels = tomo.create_tomography_circuits(Q_program, 'name', qr, cr, process_tomo_set)
```
where
- `Q_program` is the quantum program containing the preparation circuit `'name'`.
- `'name'` is the name of the preparation circuit.
- `qr` is the quantum register in `Q_program` to be measured.
- `cr` is the classical register in `Q_program` to store the measurement results.
- `process_tomo_set` is a `process_tomography_set`.

The function output, `tomo_labels`, is a list of the strings for each of the tomography circuit names stored in the `Q_program`, which is used to execute the tomography experiments. The prepared and measured qubits will be `qr[j]`, with corresponding measurement outcomes stored in classical bits `cr[j]` where `j` corresponds to each qubit in the process tomography set.

## Extracting process tomography measurement data

After building the process tomography circuits, and executing them on an available backend we must process the output results. This is done by the function:
```python
tomo_data = tomo.tomography_data(results, 'name', process_tomo_set)
```
where
- `results` is the Result object returned from execution of `process_tomo_set` circuits.
- `name` is the name of the original circuit being measured via process tomography.
- `process_tomo_set` is the tomography set used for generation of tomography data.

This returns a dictionary which stores the preparation basis, measurement basis, and measured counts along with preparation and mesurement basis configuration corresponding to each outcome.

## Performing tomographic reconstruction

Finally we may perform the fitting of the tomographic data to reconstruct a description of the process matrix. This is done by the function
```python
tomo.fit_tomography_data(tomo_data, method)
```
where
- `tomo_data` is the data returned from `process_tomography_data`.
- `method` is the reconstruction method to use. The default value is `method='wizard'`.

The return process matrix is a Choi-matrix in the column vectorization convention. For a quantum process $\mathcal E$ this is defined as the matrix
$J(\mathcal E) = \sum_{i,j} |i\rangle\langle j | \otimes \mathcal{E}(|i\rangle\langle j |)$. By default the trace of the returned Choi-matrix will be 1, but this can be changed with optional arguments. We also note that the default reconstruction method constrains the process-matrix to be a completely-positive (CP) map (by fitting a positive-semidefinite Choi-matrix), however it does *does not* constrain it to be a trace-preserving (TP) map.

### Reconstruction Methods
Currently only two methods are available.
- `method='wizard'` (default) uses the fast maximum liklihood reconstruction from [1] to return a postive semidefinite process matrix.
- `method='leastsq'` does not constrain the returned matrix to be positive-semidefinite.

Both of these methods are the same as for performing state tomography, as process tomography is implemented as state tomographic reconstruction of the Choi-matrix.

# Process Tomography Examples

We now present two examples of 1 and 2-qubit quantum process tomography. We will consider reconstruction of a Hadamard gate and a CNOT gate.

### Setting up the Quantum Program

First we construct a quantum program with two circuits, one containing a single Hadamard gate on qubit-0, and one containing a single CNOT gate between qubits-0 and 1.

```python
Q_program = QuantumProgram()
Q_program.set_api(Qconfig.APItoken, Qconfig.config["url"]) # set the APIToken and API url

# Creating registers
qr = Q_program.create_quantum_register("qr", 2)
cr = Q_program.create_classical_register("cr", 2)

# hadamard on qubit-0 only
had = Q_program.create_circuit("had", [qr], [cr])
had.h(qr[0])

# CNOT gate with qubit-1 control, qubit-0 target (target for ibmqx2)
cnot = Q_program.create_circuit("cnot", [qr], [cr])
cnot.cx(qr[0], qr[1])
```

    <qiskit.extensions.standard.cx.CnotGate at 0x10b7c1b38>

## Process tomography of a Hadamard gate

We will first demonstrate a process tomography experiment for a single qubit H gate.

### Visualization of the ideal process

Plot choi-matrix, rauli

```python
U_had = np.array([[1,1],[1,-1]])/np.sqrt(2)
# compute Choi-matrix from unitary
had_choi = outer(vectorize(U_had))
plot_state(had_choi)
```

### Simulation of process tomography experiment

Next we build the process tomography circuits and execute them on the simulator backend

```python
# process tomography set for a quantum operation on qubit-0
had_tomo_set =  tomo.process_tomography_set([0])

# Generate process tomography preparation and measurement circuits
had_tomo_circuits = tomo.create_tomography_circuits(Q_program, 'had', qr, cr, had_tomo_set)
print('Tomography circuit labels for "had" circuit:')
for label in had_tomo_circuits:
    print(label)
```

    Tomography circuit labels for "had" circuit:
    had_prep_S0(0)_meas_X(0)
    had_prep_S0(0)_meas_Y(0)
    had_prep_S0(0)_meas_Z(0)
    had_prep_S1(0)_meas_X(0)
    had_prep_S1(0)_meas_Y(0)
    had_prep_S1(0)_meas_Z(0)
    had_prep_S2(0)_meas_X(0)
    had_prep_S2(0)_meas_Y(0)
    had_prep_S2(0)_meas_Z(0)
    had_prep_S3(0)_meas_X(0)
    had_prep_S3(0)_meas_Y(0)
    had_prep_S3(0)_meas_Z(0)

```python
backend = 'local_qasm_simulator'
shots = 1000
had_tomo_results = Q_program.execute(had_tomo_circuits, shots=shots, backend=backend)
```

Next we fit the tomographic data

```python
had_process_data = tomo.tomography_data(had_tomo_results, 'had', had_tomo_set)
had_choi_fit = tomo.fit_tomography_data(had_process_data, options={'trace':2})
print('Process Fidelity = ', state_fidelity(vectorize(U_had)/2, had_choi_fit))
plot_state(had_choi_fit)
```

### Running process tomography on a real device

Now we run tomography on the IBM Quantum Experience device.

```python
# Use the IBM Quantum Experience
backend = 'ibmqx2'
# Take 500 shots for each measurement basis
# Note: reduce this number for larger number of qubits
shots = 500
# set max credits
max_credits = 8
# Run the experiment
hadqx_tomo_results = Q_program.execute(had_tomo_circuits, backend=backend, shots=shots,
                                       max_credits=max_credits, timeout=240)
print(hadqx_tomo_results)
```

    COMPLETED

```python
hadqx_process_data = tomo.tomography_data(hadqx_tomo_results, 'had', had_tomo_set)
hadqx_choi_fit = tomo.fit_tomography_data(hadqx_process_data, options={'trace':2})
print('Process Fidelity = ', state_fidelity(vectorize(U_had)/2, hadqx_choi_fit))
plot_state(hadqx_choi_fit)
```

## Process tomography of a CNOT gate

We will demonstrate a process tomography experiment for a CNOT gate.

### Visualization of the ideal process

```python
U_cnot = np.array([[1,0,0,0],[0,0,0,1],[0,0,1,0],[0,1,0,0]])
# compute Choi-matrix from unitary
cnot_choi = outer(vectorize(U_cnot))
plot_state(cnot_choi)
```

### Simulation of process tomography experiment

Our test circuit simply consists of a quantum program of two qubits with a single gate, the CNOT gate. Note that we must include a two-qubit register to record measurement outcomes.

```python
# process tomography set for quantum operation on qubits 0 and 1
cnot_tomo_set = tomo.process_tomography_set([0, 1])

# Generate process tomography circuits
cnot_tomo_circuits = tomo.create_tomography_circuits(Q_program, 'cnot', qr, cr, cnot_tomo_set)
```

```python
backend = 'local_qasm_simulator'
shots = 1000
cnot_tomo_results = Q_program.execute(cnot_tomo_circuits, backend=backend, shots=shots, timeout=300)
cnot_process_data = tomo.tomography_data(cnot_tomo_results, 'cnot', cnot_tomo_set)
```

Now we fit the data and plot the results

```python
cnot_choi_fit = tomo.fit_tomography_data(cnot_process_data, options={'trace':4})
print('Process Fidelity = ', state_fidelity(vectorize(U_cnot)/2, cnot_choi_fit/4))
plot_state(cnot_choi_fit)
```

### Running process tomography on a real device

We now demonstate running tomography on a real device. Note that 2-qubit process tomography requires evaluation of 144 circuits, so circuits may need to be submitted to the device in batches and require a lot of credits for evalaution on the IBM Q Experience.

We will break out submission up into 8 batches of 18 circuits for submission

```python
# Use the IBM Quantum Experience
backend = 'ibmqx2'
# Take 500 shots for each measurement basis
# Note: reduce this number for larger number of qubits
shots = 500
# set max credits
max_credits = 8
# Set batch size and complete
batch_size = 18
for j in range(8):
    batch_circs = cnot_tomo_circuits[j*batch_size:(j+1)*batch_size]
    if j == 0:
        cnotqx_tomo_results = Q_program.execute(batch_circs, backend=backend, shots=shots,
                                                max_credits=max_credits, timeout=240)
    else:
        cnotqx_tomo_results += Q_program.execute(batch_circs, backend=backend, shots=shots,
                                                 max_credits=max_credits, timeout=240)
    msg = 'Batch %d/%d: %s' % (j+1, 8, cnotqx_tomo_results)
    print(msg)
```

    Batch 1/8: COMPLETED
    Batch 2/8: COMPLETED
    Batch 3/8: COMPLETED
    Batch 4/8: COMPLETED
    Batch 5/8: COMPLETED
    Batch 6/8: COMPLETED
    Batch 7/8: COMPLETED
    Batch 8/8: COMPLETED

Finally we process and fit the process tomography data

```python
cnotqx_process_data = tomo.tomography_data(cnotqx_tomo_results, 'cnot', cnot_tomo_set)
cnotqx_choi_fit = tomo.fit_tomography_data(cnotqx_process_data, options={'trace':4})
print('Process Fidelity = ', state_fidelity(vectorize(U_cnot)/2, cnotqx_choi_fit/4))
plot_state(cnotqx_choi_fit)
```

```python
%run "../version.ipynb"
```

Version information

Please note that this tutorial is targeted to the stable version of the QISKit SDK. The following versions of the packages are recommended:

Package	Version

QISKit	 0.4.8

IBMQuantumExperience	>= 1.8.26

numpy	>= 1.13, < 1.14

scipy	>= 0.19, < 0.20

matplotlib	>= 2.0, < 2.1