# Variational Quantum Eigensolver

<em> Copyright (c) 2021 Institute for Quantum Computing, Baidu Inc. All Rights Reserved. </em>

## Overview

It is widely believed that one of the most promising applications of quantum computing in the near future is solving quantum chemistry problems [1-2]. **Variational Quantum Eigensolver** (VQE) is a strong proof to this possibility of studying quantum chemistry with **Noisy Intermediate-Scale Quantum** (NISQ) devices [1-4]. The core task is to solve the ground state of any molecular Hamiltonian $\hat{H}$ by preparing a parametrized wave function ansatz $|\Psi(\boldsymbol\theta)\rangle$ on a quantum computer and adopt classical optimization methods (e.g. gradient descent) to adjust the parameters $\boldsymbol\theta$ to minimize the expectation value $\langle \Psi(\boldsymbol\theta)|\hat{H}|\Psi(\boldsymbol\theta)\rangle$. This approach is based on the **Rayleigh-Ritz variational principle**.

$$
E_0 = \min_{\boldsymbol\theta} \langle \Psi(\boldsymbol\theta)|\hat{H}|\Psi(\boldsymbol\theta)\rangle.
\tag{1}
$$

where $E_0$ denotes the ground state energy. Numerically, it can be understood as finding the smallest eigenvalue $\lambda_{\min}$ of a **discretized** Hamiltonian $H$ (hermitian matrix) and its corresponding eigenvector $|\Psi_0\rangle$. How such a discretization can be done on a classical computer belongs to the art of quantum chemistry and is far beyond the scope of this tutorial. We will discuss this part with a few words in the background section. In general, such a Hamiltonian $H$ is expressed as a weighted sum of Pauli spin operators $\{X,Y,Z\}$ (native to quantum devices) such that this information can be processed on a quantum computer.

$$
H = \sum_k c_k ~ \bigg( \bigotimes_{j=0}^{M-1} \sigma_j^{(k)} \bigg),
\tag{2}
$$

where $\sigma_j^{(k)} \in \{I,X,Y,Z\}$ and $M$ stands for qubit number. We refer this form of Hamiltonian as **Pauli strings**. For example,

$$
H= 0.12~Y_0 \otimes I_1-0.04~X_0\otimes Z_1.
\tag{3}
$$

In the next section, we will provide a brief review on the electronic structure problem which essentially tells us how to calculate the Hamiltonian $H$. For those who are already familiar with this topic or only interested in how to implement VQE on Paddle Quantum, please skip this part and jump into the illustrative example of hydrogen molecule $H_2$.

## Background: the electronic structure problem

In this section, we focus on one of the fundamental problems in quantum chemistry --  **the electronic structure problem**. To be more specific, we are interested in the low lying energy eigenstates of any given molecule. These knowledge could help predict reaction rates and location of stable structures [5]. Suppose a molecule consists of $N_n$ nuclei and $N_e$ electrons, the first quantized (canonical quantization) Hamiltonian operator $\hat{H}_{mol}$ describing the total energy of this molecular system can be written as

$$
\begin{align}
\hat{H}_{\text{mol}} & = -\sum_{i}\frac{\nabla_{R_i}^2}{2M_i} - \sum_{i} \frac{\nabla_{r_i}^2}{2} -\sum_{i,j}\frac{Z_i}{\lvert R_i - r_j\lvert} + \sum_{i,j>i}\frac{Z_iZ_j}{\lvert R_i - R_j\lvert} + \sum_{i, j>i}\frac{1}{\lvert r_i - r_j\lvert},
\tag{4}
\end{align}
$$

where $R_i, M_i,$ and $Z_i$ denote the position, mass and atomic number (the number of protons) of the $i^{th}$ nucleus, and the positions of electrons are $r_i$. The first two sums describe the kinetic energy of nuclei and electrons, respectively. The third sum describes the attractive Coulomb interaction between the positively charged nuclei and the negatively charged electrons. The last two terms represent the repulsive nuclei-nuclei and electron-electron interactions.  Here, the molecular Hamiltonian $\hat{H}_\text{mol}$ is already in atomic units of energy, **Hartree**. $1$ Hartree is $[\hbar^2/(m_ee^2a_0^2)] = 27.2$ eV or $630$ kcal/mol, where $m_e, e,$ and $a_0$ stand for the mass of electron, charge of electron, and Bohr radius.

**Note:** The spin-orbit interaction and hyperfine interaction are not considered in this picture. They can be treated as perturbations if necessary.

### Born-Oppenheimer approximation

Since the nuclei are much heavier than electrons, the electrons will move much faster than the nuclei. It is reasonable to treat the positions of nuclei as fixed, $R_i =$ constants. This is known as the Born-Oppenheimer approximation by decoupling the behavior of nuclei and electrons in time scale. Consequently, the kinetic energy term of nuclei will disappear and the nuclei-nuclei repulsive interaction term can be viewed as an energy shift (independent of electron positions $r_i$). We could derive the electronic Hamiltonian $\hat{H}_{\text{electron}}$ as

$$
\begin{align}
\hat{H}_{\text{electron}} & =  - \sum_{i} \frac{\nabla_{r_i}^2}{2} -\sum_{i,j}\frac{Z_i}{\lvert R_i - r_j\lvert} + \sum_{i, j>i}\frac{1}{\lvert r_i - r_j\lvert}
\tag{5},
\end{align}
$$

The energy levels of the electrons in the molecule can be found by solving the time independent Schrödinger equation

$$
\hat{H}_{\text{electron}} |\Psi_n \rangle = E_n |\Psi_n \rangle,
\tag{6}
$$

where $n$ stands for the energy level. Notice the electron repulsion terms scale as $N_e(N_e-1)/2$ which means for the Oxygen molecule $O_2$ carrying 16 electrons there will be 120 electron repulsion terms in total! In general, this problem cannot be solved analytically. As Dirac concluded in [Quantum mechanics of many-electron systems](https://royalsocietypublishing.org/doi/10.1098/rspa.1929.0094) [6],

> *The underlying physical laws necessary for the mathematical theory of a large part of physics and the whole of chemistry are thus completely known, and the difficulty is only that the exact application of these laws leads to equations much too complicated to be soluble.*
>
> 																			-- Paul Dirac (1929)

A straightforward numerical approach is discretizing the infinite-dimensional Hilbert space into equidistant grid points where linear algebra guides the whole calculation. Suppose each axis of space is discretized into $k$ points, the $N$-electron (drop the subscript e for simplicity) wave function can be written as [2]

$$
|\Psi \rangle = \sum_{\mathbf{x_1}, \ldots, \mathbf{x_N}} \psi(\mathbf{x_1}, \ldots, \mathbf{x_N}) \mathcal{A}(|\mathbf{x_1}, \ldots, \mathbf{x_N}\rangle).
\tag{7}
$$

where coordinate $|\mathbf{x_j}\rangle = |r_j\rangle |\sigma_j\rangle$ records the spatial location and spin of the $j^{th}$ electron, $|r_j\rangle  = |x_j,y_j,z_j\rangle$ for $j\in \{1,2,\cdots,N\}$, $x_j,y_j,z_j \in \{0,1,\cdots,k-1\}$ and  $\sigma_j \in \{\downarrow,\uparrow\}$ for spin down or up.  There will be $k^{3N}\times 2^{N}$ complex amplitudes in total. Here, $\mathcal{A}$ denotes anti-symmetrization and a consequence of the Pauli exclusion principle (electrons are fermion), and $\psi(\mathbf{x_1}, \mathbf{x_2}, \ldots, \mathbf{x_N})=\langle\mathbf{x_1}, \mathbf{x_2}, \ldots, \mathbf{x_N}|\Psi\rangle$. One can see that storing such a wave function already requires **exponentially growing memory** with respect to the number of electrons $N$. This would make classical simulation methods based on this naive numerical approach intractable for systems size larger than few tens of electrons. Now, the question becomes can we prepare such a wave function $|\Psi\rangle$ directly on a quantum computer and measure the expectation value $E_0$? In the next section, let's take the simplest molecular system -- hydrogen molecule $H_2$ as a concrete example.

**Note:** A detailed review on quantum chemistry and existing classical computational methods are far beyond the scope of this tutorial, we refer the enthusiastic readers to the standard textbooks *'Molecular Electronic-Structure Theory'* [5] by Helgaker and *'Modern Quantum Chemistry: Introduction to Advanced Electronic Structure Theory'* [7] by Szabo & Ostlund. To bridge to knowledge gap between quantum chemistry and quantum computing, please check the following review papers [Quantum computational chemistry](https://journals.aps.org/rmp/abstract/10.1103/RevModPhys.92.015003) [2] and [Quantum Chemistry in the Age of Quantum Computing](https://pubs.acs.org/doi/10.1021/acs.chemrev.8b00803) [1].

**Note:** For energy calculation, it is desired to reach the **chemical accuracy** of $1.6\times10^{-3}$ Hartree or 1 kcal/mol .

## Ground state of the hydrogen molecule $H_2$

### Building electronic Hamiltonian

First of all, let us import the necessary libraries and packages.`qchem` module in Paddle Quantum is developed basing on `psi4` and `openfermion`, so you need to install these two packages before executing the following codes. We strongly encourage you to read the tutorial [Building molecular Hamiltonian](./BuildingMolecule_EN.ipynb) first, which introduces how to utilize our quantum chemistry toolkit `qchem`.

**Note: As to the environment setting, please refer to [README.md](https://github.com/PaddlePaddle/Quantum/blob/master/README.md).**

```python
import paddle
import paddle_quantum.qchem as qchem
from paddle_quantum.utils import Hamiltonian
from paddle_quantum.circuit import UAnsatz

import os
import matplotlib.pyplot as plt

import numpy
from numpy import pi as PI
from numpy import savez, zeros

# Eliminate noisy python warnings
import warnings
warnings.filterwarnings("ignore")
```

To analyze specific molecules, we need several key information such as **geometry**, **basis set** (such as STO-3G), **multiplicity**, and **charge** to model the molecule and achieve more information about the molecule, such as one-body integrations, two-body integrations, and others. Next, through our built-in quantum chemistry toolkit, we could set a `Hamiltonian` class to carry molecule Hamiltonian information, which may simplify the following calculations.

```python
geo = qchem.geometry(structure=[['H', [-0., 0., 0.0]], ['H', [-0., 0., 0.74]]])
# geo = qchem.geometry(file='h2.xyz')

# Save molecule information in to variable molecule, including one-body integrations, one-body integrations, molecular and Hamiltonian
molecule = qchem.get_molecular_data(
    geometry=geo,
    basis='sto-3g',
    charge=0,
    multiplicity=1,
    method="fci",
    if_save=True,
    if_print=True
)
# Recall Hamiltonian
molecular_hamiltonian = qchem.spin_hamiltonian(molecule=molecule,
                                               filename=None,
                                               multiplicity=1,
                                               mapping_method = 'jordan_wigner',)
# Print results
print("\nThe generated h2 Hamiltonian is \n", molecular_hamiltonian)
```

    FCI energy for H2_sto-3g_singlet (2 electrons) is -1.137283834485513.

    The generated h2 Hamiltonian is
     -0.09706626861762556 I
    -0.04530261550868938 X0, X1, Y2, Y3
    0.04530261550868938 X0, Y1, Y2, X3
    0.04530261550868938 Y0, X1, X2, Y3
    -0.04530261550868938 Y0, Y1, X2, X3
    0.1714128263940239 Z0
    0.16868898168693286 Z0, Z1
    0.12062523481381837 Z0, Z2
    0.16592785032250773 Z0, Z3
    0.17141282639402394 Z1
    0.16592785032250773 Z1, Z2
    0.12062523481381837 Z1, Z3
    -0.2234315367466399 Z2
    0.17441287610651626 Z2, Z3
    -0.2234315367466399 Z3

**Note:** This Hamiltonian is generated with an interatomic distance of $d = 74$ pm.

In addition to inputting molecular geometry directly, inputting molecular geometry file (.xyz file) is also allowed. For more information about quantum chemistry toolkit, please refer to the tutorial [Building molecular Hamiltonian](./BuildingMolecule_EN.ipynb). If you need to test the geometric configuration of more molecules, please check out this [database](http://smart.sns.it/molecules/index.html).

### Building QNN and trial wave function

To implement VQE, we firstly need to design a quantum neural network QNN to prepare the wave function ansatz $|\Psi(\boldsymbol\theta)\rangle$. Here, we provide a 4-qubit quantum circuit template with a depth of $D$ blocks. The dotted frame in the figure below denotes a single block:

Next, we use the `UAnsatz` class and the built-in `real_entangled_layer(theta, D)` circuit template in Paddle Quantum to realize this QNN.

```python
def U_theta(theta, Hamiltonian, N, D):
    """
    Quantum Neural Network
    """

    # Initialize the quantum neural network according to the number of qubits N
    cir = UAnsatz(N)

    # Built-in {R_y + CNOT} circuit template
    cir.real_entangled_layer(theta[:D], D)

    # Lay R_y gates in the last row
    for i in range(N):
        cir.ry(theta=theta[D][i][0], which_qubit=i)

    # The quantum neural network acts on the default initial state |0000>
    fin_state = cir.run_state_vector()

    # Calculate the expectated value of the given Hamiltonian
    expectation_val = cir.expecval(Hamiltonian)

    return expectation_val, cir, fin_state
```

### Setting up the loss function and model

Now that we have the target Hamiltonian and QNN, we will further define the training model and loss function. By applying the QNN $U(\theta)$ on the initial state $|0..0\rangle$, we get the output state $|\psi(\boldsymbol\theta)\rangle $. Then, the loss function to be minimized is the expectation value,

$$
\min_{\boldsymbol\theta}  \mathcal{L}(\boldsymbol \theta) = \min_{\boldsymbol\theta} \langle \Psi(\boldsymbol\theta)|H |\Psi(\boldsymbol\theta)\rangle
= \min_{\boldsymbol\theta} \sum_k c_k~\langle \Psi(\boldsymbol\theta)| \bigotimes_j \sigma_j^{(k)}|\Psi(\boldsymbol\theta)\rangle.
\tag{8}
$$

```python
class StateNet(paddle.nn.Layer):
    """
    Construct the model net
    """

    def __init__(self, shape, dtype="float64"):
        super(StateNet, self).__init__()

        # Initialize the list of theta parameters, filling the initial values with a uniform distribution of [0, 2* PI]
        self.theta = self.create_parameter(shape=shape,
                                           default_initializer=paddle.nn.initializer.Uniform(low=0.0, high=2*PI),
                                           dtype=dtype, is_bias=False)

    # Define loss function and forward propagation mechanism
    def forward(self, N, D):

        # Calculate loss function/expected value
        loss, cir, fin_state = U_theta(self.theta, molecular_hamiltonian.pauli_str, N, D)

        return loss, cir, fin_state
```

### Hyper-parameters

Before training the QNN, we also need to set some training hyper-parameters, mainly the learning rate (LR), the number of iterations (ITR), and the depth (D) of repeated blocks.

```python
ITR = 80  # Set the number of optimization iterations
LR = 0.4   # Set the learning rate
D = 2      # Set the depth of the repetitive calculation module in QNN
N = molecular_hamiltonian.n_qubits # Set number of qubits
```

### Training

After all the training model parameters are set, we convert the data into Tensor in the Paddle, and then train the quantum neural network. The results of the training process is stored in the summary_data file.

```python
# Determine the parameter dimensions of the network
net = StateNet(shape=[D + 1, N, 1])

# In general, we use the Adam optimizer to obtain relatively good convergence,
# Of course you can change it to SGD or RMS prop.
opt = paddle.optimizer.Adam(learning_rate=LR, parameters=net.parameters())

# Record optimization results
summary_iter, summary_loss = [], []

# Optimize iterations
for itr in range(1, ITR + 1):

    # Forward propagation calculates the loss function
    loss, cir, fin_state = net(N, D)

    # Back propagation minimizes the loss function
    loss.backward()
    opt.minimize(loss)
    opt.clear_grad()

    # Update optimization results
    summary_loss.append(loss.numpy())
    summary_iter.append(itr)

    # Print results
    if itr % 20 == 0:
        print("iter:", itr, "loss:", "%.4f" % loss.numpy())
        print("iter:", itr, "Ground state energy:", "%.4f Ha"
                                            % loss.numpy())
    if itr == ITR:
        print("\nCircuit after training:")
        print(cir)

# Save the training results in the Output folder
os.makedirs("output", exist_ok=True)
savez("./output/summary_data", iter = summary_iter,
                               energy=summary_loss)
```

    iter: 20 loss: -1.0510
    iter: 20 Ground state energy: -1.0510 Ha
    iter: 40 loss: -1.1298
    iter: 40 Ground state energy: -1.1298 Ha
    iter: 60 loss: -1.1361
    iter: 60 Ground state energy: -1.1361 Ha
    iter: 80 loss: -1.1371
    iter: 80 Ground state energy: -1.1371 Ha

    Circuit after training:
    --Ry(4.710)----*--------------x----Ry(1.569)----*--------------x----Ry(-0.01)--
                   |              |                 |              |
    --Ry(1.507)----x----*---------|----Ry(1.566)----x----*---------|----Ry(4.392)--
                        |         |                      |         |
    --Ry(5.984)---------x----*----|----Ry(4.476)---------x----*----|----Ry(1.692)--
                             |    |                           |    |
    --Ry(0.132)--------------x----*----Ry(7.773)--------------x----*----Ry(3.194)--

### Benchmarking
We have now completed the training of the quantum neural network, and the estimated value of the ground state energy obtained is $E_0 \approx -1.137 $ Hartree, we compare it with the theoretical value $E_0 = -1.1371$ to benchmark our model. The estimation obtained with VQE agree with a full configuration-interaction (FCI) calculation within chemical accuracy $\varepsilon = 1.6 \times 10^{-3}$ Hartree.

```python
result = numpy.load('./output/summary_data.npz')

eig_val, eig_state = numpy.linalg.eig(
                     molecular_hamiltonian.construct_h_matrix())
min_eig_H = numpy.min(eig_val.real)
min_loss = numpy.ones([len(result['iter'])]) * min_eig_H

plt.figure(1)
func1, = plt.plot(result['iter'], result['energy'],
                  alpha=0.7, marker='', linestyle="-", color='r')
func_min, = plt.plot(result['iter'], min_loss,
                  alpha=0.7, marker='', linestyle=":", color='b')
plt.xlabel('Number of iteration')
plt.ylabel('Energy (Ha)')

plt.legend(handles=[
    func1,
    func_min
],
    labels=[
        r'$\left\langle {\psi \left( {\theta } \right)} '
        r'\right|H\left| {\psi \left( {\theta } \right)} \right\rangle $',
        'Ground-state energy',
    ], loc='best')

#plt.savefig("vqe.png", bbox_inches='tight', dpi=300)
plt.show()
```

## Determining the interatomic distance

Recall the above calculation is done with an interatomic distance $d = 74$ pm between two hydrogen atoms. Another interesting aspect we can try with VQE is determining the true interatomic distance by modifying the `h2.xyz` file. The results are summarize in figure below,

The lowest value is found around $d = 74$ pm (1 pm = $1\times 10^{-12}$m), which is consistent with the [experimental data](https://cccbdb.nist.gov/exp2x.asp?casno=1333740&charge=0) $d_{exp} (H_2) = 74.14$ pm.

_______

## References

[1] Cao, Yudong, et al. Quantum Chemistry in the Age of Quantum Computing. [Chemical reviews 119.19 (2019): 10856-10915.](https://pubs.acs.org/doi/10.1021/acs.chemrev.8b00803)

[2] McArdle, Sam, et al. Quantum computational chemistry. [Reviews of Modern Physics 92.1 (2020): 015003.](https://journals.aps.org/rmp/abstract/10.1103/RevModPhys.92.015003)

[3] Peruzzo, A. et al. A variational eigenvalue solver on a photonic quantum processor. [Nat. Commun. 5, 4213 (2014).](https://www.nature.com/articles/ncomms5213)

[4] Moll, Nikolaj, et al. Quantum optimization using variational algorithms on near-term quantum devices. [Quantum Science and Technology 3.3 (2018): 030503.](https://iopscience.iop.org/article/10.1088/2058-9565/aab822)

[5] Helgaker, Trygve, Poul Jorgensen, and Jeppe Olsen. Molecular electronic-structure theory. John Wiley & Sons, 2014.

[6] Dirac, Paul Adrien Maurice. Quantum mechanics of many-electron systems. [Proceedings of the Royal Society of London. Series A, Containing Papers of a Mathematical and Physical Character 123.792 (1929): 714-733.](https://royalsocietypublishing.org/doi/10.1098/rspa.1929.0094)

[7] Szabo, Attila, and Neil S. Ostlund. Modern quantum chemistry: introduction to advanced electronic structure theory. Courier Corporation, 2012.


# Assessing the accuracy of ML-based uncertainties in the context of galaxy photometric redshifts

A data challenge for [From Quarks to Cosmos with AI](https://events.mcs.cmu.edu/qtc2021/)

_Alex I. Malz (German Centre for Cosmological Lensing)_

## Motivation: Uncertainty quantification

More and more ML/AI methods, particularly in Bayesian deep learning, yield uncertainties $\hat{p}(y | x_{i})$ of target parameters $y$ given observed random variables $x_{i}$, rather than just point estimates $\hat{y}_{i}$.
Though rarely framed as such, these uncertainties are _posteriors_, and there's more than meets the eye to what lies on the righthand side of the conditional.

Really, an estimated posterior should be writen as $p(y | x_{i}, \{y_{n}, x_{n}\}_{N}, \pi)$ for training data $\{y_{n}, x_{n}\}_{N}$ and algorithm $\pi$.
Though the dependence on training data is straightforward, the dependence on the algorithm, meaning the estimation model and its implementation, is subtle.
However, if it were not there, then every method trained on the same data would yield identical estimated uncertainties, which is not observed.
While the training set $\{y_{n}, x_{n}\}_{N}$ is equivalent to an _explicit prior_, the algorithm $\pi$ must be an _implicit prior_.

Another way of looking at estimated posteriors is in terms of the type of uncertainty encompassed by each term on the righthand side of the conditional.
For a noisy measurement or a stochastic generative process, the random variable $x_{i} \sim p(x | y_{i})$ represents the _aleatoric uncertainty_, that inherent to the data.
However, the training set $\{y_{n}, x_{n}\}_{N}$ and implemented algorithm $\pi$ could potentially be improved to yield a better estimate and thus constitute sources of _epistemic uncertainty_.
In physics, we want to learn the aleatoric uncertainty $p(x | y_{i})$, which we can't get from the $p(y | x_{i}, \{y_{n}, x_{n}\}_{N}, \pi)$ in hand without knowing $p(y | \{y_{n}, x_{n}\}_{N}, \pi)$.

Meanwhile, assessments of the performance of estimated $p(y | x_{i}, \{y_{n}, x_{n}\}_{N}, \pi)$ are almost always made by comparison to known $y_{i}$, leaving unanswered the question of how well the estimator approximates $p(y | x_{i})$.
Why?
The problem is that $p(y | x_{i})$ is not necessarily known, certainly not for observed $y_{i}$ measured in nature, but also generally for simulated $y_{i}$, at least in astrophysics.

## tl;dr Goals

This data challenge invites participants to address any of the following open problems.
- How can we use AI to estimate posteriors?
- How should we assess posterior estimates?
- How may we quantify the implicit prior(s)?

This tutorial/demo provides some starting points for answering these questions in the context of photometric redshifts (photo-$z$s).

## Context: Photometric redshifts

[Photo-$z$s](https://en.wikipedia.org/wiki/Photometric_redshift) provide an excellent testbed for addressing these issues but require some introduction.
There isn't a definitive primer, but here are several overviews that are informative, if a bit dry.
- [basic intro from Rubin Observatory](https://www.lsst.org/science/dark-energy/photometric-redshift)
- [old overview covering classic concepts](https://ned.ipac.caltech.edu/level5/Glossary/Essay_photredshifts.html)
- [recent review of estimation methods](https://arxiv.org/abs/1805.12574)
If these don't answer your questions, Chapter 0 of [my thesis](https://zenodo.org/record/3973536) might, or, better yet, [the slides from my defense](https://github.com/aimalz/ship-of-theses/tree/master/presentation) make for a better tl;dr.

One reason photo-$z$s are perfect for a data challenge is that they are simple enough to obtain $p(y | x_{i})$ along the way to generating a sample of $(y_{i}, x_{i})$ pairs.
Here, the target variable $y \to z$ is redshift, a scalar, and the data $x \to \vec{d} = (u, g, r, i, z, y)$, or some trivial function thereof, is a vector of length $<10$ observed [photometric magnitudes](https://en.wikipedia.org/wiki/Magnitude_(astronomy)) of galaxies through [broadband optical filters](https://en.wikipedia.org/wiki/Photometric_system) (which I somewhat arbitrarily choose to be those of [the Vera C. Rubin Observatory](https://www.lsst.org/)).
Because of the extremely low dimensionality of the problem, we can forward model not just individual pairs $(z_{i}, \vec{d}_{i})$ but the entire joint probability space $p(z, \vec{d})$, thereby obtaining $p(y | x_{i})$ for every $(z_{i}, \vec{d}_{i})$.
That doesn't mean it's trivial to do so -- most of this tutorial concerns that forward-modeling procedure -- but it is possible.

Another reason photo-$z$s are perfect for this data challenge is that comprehensive uncertainty quantification for galaxy redshifts in the absence of spectroscopy are crucial for the Legacy Survey of Space and Time (LSST), an upcoming photometric survey on the Rubin Observatory.
This data challenge has the ulterior motive of strengthening the cosmology analysis pipeline of LSST's Dark Energy Science Collaboration (DESC).
This tutorial makes use of two pieces of software DESC members are publicly developing right now, [RAIL](https://github.com/LSSTDESC/RAIL) and [qp](https://github.com/LSSTDESC/qp), whose functionality will be introduced where relevant, along with other code dependencies.
Participants need not use either in their responses to the challenge questions, but the development team welcomes feedback from potential users, contributions to the codebase that could result from this data challenge, and new team members, DESC membership not required.

# Tutorial and challenge prompts

Finally, without more ado, we set the stage with a tutorial that:
1. creates realistically complex mock photo-$z$ posteriors, redshifts, and photometry to define training and test sets;
2. estimates photo-$z$ posteriors of the test set given the training set using ML;
3. quantifies how closely the estimated posteriors approximate the true ones.

_The three-pronged structure of this tutorial is inspired by that of RAIL, which has subpackages corresponding to each of the enumerated parts of the tutorial: `rail.creation`, `rail.estimation`, and `rail.evaluation`._

The challenge comes from building on the tutorial content to conduct a self-guided investigation of the open questions.
To get the most out of this opportunity, think of this tutorial as a lab manual that presents an experimental procedure to answer each question and includes an example of a possible solution, then invites participants to "riff" on those solutions to devise and implement their own.
The greatest opportunities for such investigations can be found under the headings **"Your turn!"** throughout the tutorial, but there are certainly more interesting directions to go in.

## 0. Setup

This tutorial uses several packages that may not be installed in the conference's kernel.
The following commands installed them successfully on Google Colab, but there might be a better way to install it on PSC machines.
On the PSC machines, first follow [these instructions](https://gist.github.com/EiffL/a8cfa39c635334d063f0d3c07a85d1ec).

```python
# !pip install cde-diagnostics
# !pip install cdetools
# !pip install corner
# !pip install FlexCode
# !pip install pzflow
# !pip install sklearn
# !pip install xgboost==0.90
```

```python
import jax.numpy as jnp
import numpy as np
import pandas as pd
import sklearn
```

```python
# !git clone https://github.com/COINtoolbox/photoz_catalogues.git ..
# !pip install git+https://github.com/LSSTDESC/RAIL.git#egg=rail[Full]
# !pip install git+https://github.com/LSSTDESC/qp.git#egg=qp[Full]
```

```python
import cdetools
import flexcode
import pzflow
import qp
import rail
```

```python
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.patches as mpl_patches
mpl.rc('text', usetex=False)

import corner
```

## 1. Generating test and training data, including true photo-$z$ posteriors

This is the most complex of the parts of the tutorial, which should leave readers with a sense of why this kind of experiment has not been done for higher-dimensional data.
There are three main steps:
0. Prepare and explore data to use as the basis for a realistically complex generative model.
1. Emulate a realistically complex $p(z, \vec{d})$ probability space using input data.
2. Emulate a distinct $p'(z, \vec{d})$ probability space from an existing model $p(z, \vec{d})$ of that space.
3. Sample $p(z, \vec{d})$ and $p'(z, \vec{d})$ to produce test and training sets, respectively, with known $z$, $\vec{d}$, and $p(z | \vec{d})$.

### Data

Though we're going to forward-model mock data to experiment on, we want it to be _realistically complex_, meaning it shares the physical degeneracies and systematic errors that would be present in a real data set.
The [Happy/Teddy data sets](https://github.com/COINtoolbox/photoz_catalogues) (see [Beck, et al 2017](https://arxiv.org/abs/1701.08748) for full release notes) are curated subsamples of the [Sloan Digital Sky Survey (SDSS) Data Release (DR) 12](https://www.sdss.org/dr12/), a spectroscopic survey with high-fidelity redshift measurements, and were created by the [Cosmostatistics Initiative (COIN)](https://cosmostatistics-initiative.org/).
The data sets are defined to emulate the kinds of differences in observational properties of galaxies with measured spectroscopic redshifts and those for which only photometry is available, and they were created with the goal of determining the impact of imbalance between training, validation, and test sets for photo-$z$ point estimation.
They are thus an appropriate starting point for creating a realistically complex model of the joint probability space $p(z, \vec{d})$.

Of course there are plenty of other potential data sets available, including those that are simulated, which may be advantageous for the potential to estimate other galaxy properties such as stellar mass and star formation rate that are known in a simulation but not directly measurable with real data.

```python
happy_path = '../photoz_catalogues/Happy/happy_'
header = pd.read_csv(happy_path+'A', delim_whitespace=True, nrows=0).columns[1:]
teddy_path = '../photoz_catalogues/Teddy/teddy_'

happy, teddy = {}, {}
for lett in ['A', 'B', 'C', 'D']:
    happy[lett] = pd.read_csv(happy_path+lett, delim_whitespace=True, header=None, skiprows=1, names=header)
    happy[lett] = happy[lett]
    teddy[lett] = pd.read_csv(teddy_path+lett, delim_whitespace=True, header=None, skiprows=7, names=header)
    teddy[lett] = teddy[lett]
```

```python
# TODO: separate Happy and Teddy

fig, ax = plt.subplots(2, 2, figsize=(10, 10))
for j, col in enumerate(['z_spec', 'mag_r']):
    for i, lett in enumerate(['A', 'B', 'C', 'D']):
        ax[j][0].hist(happy[lett][col], alpha=0.25, bins=100, density=False, label='happy_'+lett)
        ax[j][1].hist(teddy[lett][col], alpha=0.25, bins=100, density=False, label='teddy_'+lett)
        ax[j][0].set_xlabel(col)
        ax[j][0].legend()
        ax[j][1].set_xlabel(col)
        ax[j][1].legend()
```

For convenience later on, we can safely cut off the tiny fraction of outliers in redshift and $r$-band magnitude.

```python
z_min, z_max = 0., 1.5
r_min, r_max = 10., 25.
```

```python
colorcycle = 'rbgcmy'
fig = corner.corner(happy['A'][['u-g', 'g-r', 'r-i', 'i-z']], color='k', alpha=0.25)
for i, lett in enumerate(['B', 'C', 'D']):
    corner.corner(happy[lett][['u-g', 'g-r', 'r-i', 'i-z']], fig=fig, color=colorcycle[i], alpha=0.25)

fig = corner.corner(teddy['A'][['u-g', 'g-r', 'r-i', 'i-z']], color='k', alpha=0.25)
for i, lett in enumerate(['B', 'C', 'D']):
    corner.corner(teddy[lett][['u-g', 'g-r', 'r-i', 'i-z']], fig=fig, color=colorcycle[i], alpha=0.25)
```

### Emulating $p(z, \vec{d})$ from input data with `pzflow`

[`pzflow`](https://github.com/jfcrenshaw/pzflow) is a package for making normalizing flows from sets of redshifts and photometry in order to estimate or otherwise model photo-$z$ posteriors.
We'll use it to make a normalizing flow that will serve as the model for $p(z, \vec{d})$.
This is the only part of the tutorial that explicitly takes advantage of GPU capabilities.

_This content is adapted from pzflow's [conditional redshifts demo](https://github.com/jfcrenshaw/pzflow/blob/main/examples/conditional_redshift_example.ipynb), written by John Franklin Crenshaw (UW)._
_The same functionality can also be accessed through a `rail.Creator()` object._

```python
from pzflow import Flow
from pzflow.bijectors import Chain, ColorTransform, InvSoftplus, StandardScaler, RollingSplineCoupling
from pzflow.examples import galaxy_data
```

Let's start with a demonstration of how pzflow makes a model of $p(z, \vec{d}$ from its own demo data.

```python
data = galaxy_data()

# restrict to Happy/Teddy range for coverage in demo
data = data[(data['redshift'] > z_min) & (data['redshift'] < z_max) & (data['r'] > r_min) & (data['r'] < r_max)]

# normalize
data = data.values

# use fewer bands to be able to compare with Happy/Teddy
data = data[:, :-1]

# convert magnitudes to a reference magnitude and colors
data[:, 1], data[:, 2:] = data[:, 3], np.diff(data[:,1:])

# standard scale the reference magnitude and colors
data[:, 1:] = (data[:, 1:] - data[:, 1:].mean(axis=0))/data[:, 1:].std(axis=0)

# calculate the mean and standard deviations of the dimensions
# note that on this side of the color transform, we are working
# in color space, so I calculate the mean and std dev of these
# variables instead of the raw magnitudes
means = data[:, 1:].mean(axis=1)
stds = data[:, 1:].std(axis=1)

# save the new set
data = pd.DataFrame(data, columns=('redshift', 'r', 'u-g', 'g-r', 'r-i', 'i-z'))
# trim_data[i][lett] = data
```

```python
# set the inverse softplus parameters, estimated
# to ensure that sampled redshifts are positive
column_idx = 0
sharpness = 5
n_epoch = 125

# construct our bijector
# by chaining all these layers
bijector = Chain(
    InvSoftplus(column_idx, sharpness),
    StandardScaler(means[column_idx], stds[column_idx]),
    RollingSplineCoupling(nlayers=1, n_conditions=len(data.columns)-1),
)

# To create the conditional flow, we have to provide
# 1. The names of the data columns
# 2. The bijector
# 3. The names of the conditional columns
flow = Flow(data.columns[:1], bijector, conditional_columns=data.columns[1:])

losses = flow.train(data, epochs=n_epoch, verbose=True)

plt.plot(losses)
plt.xlabel("Epoch")
plt.ylabel("Training loss")
plt.show()
```

We need to choose a redshift grid upon which to evaluate the posterior PDFs.

```python
granularity = 100
grid = np.linspace(z_min, z_max, granularity)
```

Let's check one of them before making more.

```python
chosen = 999

galaxy = data.loc[[chosen]]
pdf = flow.posterior(galaxy, column="redshift", grid=grid)

plt.plot(grid, pdf[0], label='Posterior')
plt.axvline(galaxy['redshift'].values[0], 0, 1, c='C3', label='True redshift')
plt.legend()
plt.xlabel("redshift")
plt.show()
```

The code below limits the number of sample posteriors to generate only because I ran out of memory when attempting to save posteriors for the whole set of demo data on PSC (but not on Google Colab).

```python
samples = flow.sample(1, conditions=data[:5000], seed=0)
plt.hist(data['redshift'], range=(0, 2.5), bins=40, histtype='step', label='data', density=True)
plt.hist(samples['redshift'], range=(0, 2.5), bins=40, histtype='step', label='samples', density=True)
plt.xlabel('z')
plt.legend()
plt.show()
```

```python
z = samples['redshift']
z.to_csv('test_set_redshifts.csv')

phot = samples[['r', 'u-g', 'g-r', 'r-i', 'i-z']]
phot.to_csv('test_set_photometry.csv')

posteriors = flow.posterior(samples, column="redshift", grid=grid)
with open('test_set_posteriors.csv', 'wb') as fn:
    jnp.save(fn, posteriors)
```

We can do this procedure for all the Happy/Teddy samples so we can experiment with them later.
It's kind of slow, at least without using GPU, and it only needs to be done once, so we won't do it in realtime.

```python
# full_data = {'happy': happy, 'teddy': teddy}
# n_out = 1000

# for name, dat in full_data.items():
#     for lett in ['A', 'B', 'C', 'D']:
#         print(name+lett)
#         print(dat[lett].columns)
#         dat[lett] = dat[lett][['z_spec', 'mag_r', 'u-g', 'g-r', 'r-i', 'i-z']]
#         dat[lett] = dat[lett][(dat[lett]['z_spec'] > z_min) & (dat[lett]['z_spec'] < z_max) & (dat[lett]['mag_r'] > r_min) & (dat[lett]['mag_r'] < r_max)]
#         dat[lett] = dat[lett].rename(columns={'z_spec': 'redshift', 'mag_r': 'r'})

#         data = dat[lett].values

#     # standard scale the reference magnitude and colors
#         data[:, 1:] = (data[:, 1:] - data[:, 1:].mean(axis=0))/data[:, 1:].std(axis=0)

#     # calculate the mean and standard deviations of thefrom sklearn.cross_validation import train_test_split dimensions
#     # note that on this side of the color transform, we are working
#     # in color space, so I calculate the mean and std dev of these
#     # variables instead of the raw magnitudes
#         means = data[:, 1:].mean(axis=0)
#         stds = data[:, 1:].std(axis=0)

#     # save the new set
#         data = pd.DataFrame(data, columns=('redshift', 'r', 'u-g', 'g-r', 'r-i', 'i-z'))

#     # construct our bijector
#   # by chaining all these layers
#         bijector = Chain(
#           InvSoftplus(column_idx, sharpness),
#           StandardScaler(means[column_idx], stds[column_idx]),
#           RollingSplineCoupling(nlayers=1, n_conditions=len(data.columns)-1),
#         )

#   # To create the conditional flow, we have to provide
#   # 1. The names of the data columns
#   # 2. The bijector
#   # 3. The names of the conditional columns
#         flow = Flow(data.columns[:1], bijector, conditional_columns=data.columns[1:])

#         losses = flow.train(data, epochs=n_epoch, verbose=False)

#         samples = flow.sample(1, conditions=data.sample(n_out), seed=0)

#         z = samples['redshift']
#         z.to_csv(name+lett+'redshifts.csv', index=False)

#         phot = samples[['r', 'u-g', 'g-r', 'r-i', 'i-z']]
#         phot.to_csv(name+lett+'photometry.csv', index=False)

#         posteriors = flow.posterior(samples, column="redshift", grid=grid)
#         with open(name+lett+'posteriors.csv', 'wb') as fn:
#             jnp.save(fn, posteriors)
```

### Emulating $p'(z, \vec{d})$ from $p(z, \vec{d})$ using `rail.creation`

While one can build $p(z, \vec{d})$ from an input data set, more systematic investigation can be conducted by creating many training, validation, and test set combinations that are incremental variations on one another.
As one of RAIL's purposes is to conduct one such experiment, the `rail.creation` subpackage not only wraps `pzflow` (the above example can be executed just as easily through the RAIL interface) but also includes physically motivated forms of systematic discrepancy between training and test sets that can be applied to a tunable degree to create "degraded" versions of an original $p(z, \vec{d})$ space.
So, for completeness, let's demonstrate how to create a new $p'(z, \vec{d})$ from an existing $p'(z, \vec{d})$ using `rail.creation`.

_This content is taken directly from RAIL's [degradation demo](), written by John Franklin Crenshaw (UW)._

```python
from pzflow.examples import example_flow
from rail.creation import Creator, engines
from rail.creation.degradation import InvRedshiftIncompleteness, LineConfusion
```

We can create a new normalizing flow model of the $p(z, \vec{d})$ space based on an existing one.
Let's start with one that emulates redshift incompleteness.

```python
flow = engines.FlowEngine(example_flow())
creator = Creator(flow)
degraded_creator = Creator(flow, degrader=InvRedshiftIncompleteness(0.8))
```

```python
samples = creator.sample(100000)
degraded_samples = degraded_creator.sample(100000)
```

```python
fig, ax = plt.subplots(figsize=(4.5,4.5), dpi=100)
ax.hist(samples['redshift'], bins=20, range=(0,2.3), histtype='step', label="Unbiased samples")
ax.hist(degraded_samples['redshift'], bins=20, range=(0,2.3), histtype='step', label="Degraded samples")
ax.legend()
ax.set(xlabel="Redshift", ylabel="Number of galaxies", xlim=(0,2.3))
plt.show()
```

As expected, the "degraded" model has a different redshift distribution than the original one.
Next, we can try to emulate errors in the spectroscopic catalog that can arise from misidentified emission lines.

```python
def OxygenLineConfusion(data, seed=None):
    OII = 3727
    OIII = 5007

    data = LineConfusion(true_wavelen=OII, wrong_wavelen=OIII, frac_wrong=0.02)(data, seed)
    data = LineConfusion(true_wavelen=OIII, wrong_wavelen=OII, frac_wrong=0.01)(data, seed)
    return data
```

```python
flow = engines.FlowEngine(example_flow())
creator = Creator(flow)
degraded_creator = Creator(flow, degrader=OxygenLineConfusion)
```

```python
samples = creator.sample(100000, seed=0)
degraded_samples = degraded_creator.sample(100000, seed=0)
```

```python
fig, ax = plt.subplots(figsize=(4.5,4.5), dpi=100)
ax.scatter(samples["redshift"], degraded_samples["redshift"], s=0.1)
ax.set(xlabel="True spec-z", ylabel="Erroneous spec-z")
plt.show()
```

### Your turn!

These are just two ways to create a probability space $p(z, \vec{d})$ that can forward model redshifts, photometry, and photo-$z$ posteriors.
Next, we'll use the samples drawn from these systematically different models as training, validation, and test sets for estimators of photo-$z$ posteriors.

**CHALLENGE**: Try modeling $p(z, \vec{d})$ with different starting data sets.

**CHALLENGE**: How would you incorporate the notion of measurement errors ($\vec{d}_{meas} \sim p(\vec{d}_{meas} | \vec{d}_{true})$) into the normalizing flow model for the data?

**CHALLENGE**: Make systematically different sample data sets to use as training, validation, and test sets, varying them in controlled ways.

**CHALLENGE**: Devise and implement another creator of a joint probability space that can encompass realistic complexity of photometric data.

**CHALLENGE**: Extend the probability space to $p(z, SED)$ to draw redshifts, full spectra, and redshift posteriors.

## 2. Estimating photo-z posterior PDFs

There are many estimators of photo-$z$ posteriors, and many of those are compared to one another in [Schmidt & Malz, et al. 2020](https://arxiv.org/abs/2001.03621).
Of the tested estimators, including ML and non-ML methods, the most promising was [`FlexCode`](https://github.com/tpospisi/FlexCode) ([Izbicki & Lee, 2017](https://arxiv.org/abs/1704.08095)), which also happens to be one of the easiest to install and apply, so we'll demonstrate it as an example of an estimator of photo-$z$ posteriors.

### Estimating photo-$z$ posterior PDFs with FlexCode

A demonstration of `FlexCode` in the context of photo-$z$s can be found in [Dalmasso, et al 2019](https://arxiv.org/abs/1908.11523), with demos in `R` published on [GitHub](https://github.com/Mr8ND/cdetools_applications).
We'll demonstrate it on the pzflow-based samples generated from the Happy/Teddy data sets.*italicized text*

_This content is adapted from FlexCode's [Teddy tutorial](https://github.com/tpospisi/FlexCode/blob/master/tutorial/Flexcode-tutorial-teddy.ipynb), written by Nic Dalmasso (CMU)._
_The same functionality can also be accessed through a `rail.Estimator()` object; see [tutorial](https://github.com/LSSTDESC/RAIL/blob/master/examples/estimation/RAIL_estimation_demo.ipynb) written by Sam Schmidt (UC Davis)._

```python
n_grid = granularity

# Select regression method
from flexcode.regression_models import NN

# Parameters
basis_system = "cosine"  # Basis system
max_basis = 31           # Maximum number of basis. If the model is not tuned,
                         # max_basis is set as number of basis

# Regression Parameters
# If a list is passed for any parameter automatic 5-fold CV is used to
# determine the best parameter combination.
params = {"k": 5}#[5, 10, 15, 20]}       # A dictionary with method-specific regression parameters.
```

Let's try first with a representative training/validation set.

```python
x_orig = pd.read_csv('test_set_photometry.csv')[['r', 'u-g', 'g-r', 'r-i', 'i-z']].to_numpy()
y_orig = pd.read_csv('test_set_redshifts.csv')[['redshift']].to_numpy()
posteriors_orig = pd.DataFrame(np.load('test_set_posteriors.csv')).to_numpy()

# n_samp = 10000
from sklearn.model_selection import train_test_split
x_train, x_test, y_train, y_test, posteriors_train, posteriors_test = train_test_split(x_orig, y_orig, posteriors_orig,
                                                                                train_size=2000, random_state=42)
x_train, x_validation, y_train, y_validation, posteriors_train, posteriors_validation = train_test_split(x_train, y_train, posteriors_train,
                                                                                train_size=1000, random_state=42)
```

```python
# Parameterize model
model = flexcode.FlexCodeModel(NN, max_basis, basis_system, regression_params=params)

# Fit model - this will also choose the optimal number of neighbors `k`
model.fit(x_train, y_train)

# Tune model - Select the best number of basis
model.tune(x_validation, y_validation)

# Predict new densities on grid
cde_test, y_grid = model.predict(x_test, n_grid=n_grid)
```

We can examine one of the photo-$z$ posteriors estimated with the perfectly representative training/validation set.

```python
chosen = 9

plt.plot(y_grid, cde_test[chosen], label='Estimated posterior')
plt.plot(grid, posteriors_test[chosen], label='True posterior')
plt.axvline(y_test[chosen], 0, 1, c='C3', label='True redshift')
plt.legend()
plt.xlabel("redshift")
plt.show()
```

It looks pretty good!
Now let's try training and validating with some Happy/Teddy data but estimating posteriors on the test set from the pzflow demo.

```python
y_train = pd.read_csv('teddyAredshifts.csv')['redshift'].to_numpy()
x_train = pd.read_csv('teddyAphotometry.csv')[['r', 'u-g', 'g-r', 'r-i', 'i-z']].to_numpy()

y_validation = pd.read_csv('teddyBredshifts.csv')['redshift'].to_numpy()
x_validation = pd.read_csv('teddyBphotometry.csv')[['r', 'u-g', 'g-r', 'r-i', 'i-z']].to_numpy()
```

```python
# Parameterize model
model = flexcode.FlexCodeModel(NN, max_basis, basis_system, regression_params=params)

# Fit model - this will also choose the optimal number of neighbors `k`
model.fit(x_train, y_train)

# # # Tune model - Select the best number of basis
model.tune(x_validation, y_validation)

# # Predict new densities on grid
cde_test_bias, y_grid_bias = model.predict(x_test, n_grid=n_grid)
```

```python
plt.plot(y_grid_bias, cde_test_bias[chosen], label='Estimated posterior (biased training set)')
plt.plot(y_grid, cde_test[chosen], label='Estimated posterior (unbiased training set)')
plt.plot(grid, posteriors_test[chosen], label='True posterior')
plt.axvline(y_test[chosen], 0, 1, c='C3', label='True redshift')
plt.legend()
plt.xlabel("redshift")
plt.show()
```

As expected, the biased training and validation sets worsen the estimated posterior PDFs.

### Your turn!

AI is underutilized in estimating photo-$z$ posteriors, so the most exciting aspect of this challenge is to improve upon existing estimators.

**CHALLENGE**: Implement another estimator of photo-$z$ posteriors; consider those used in [Brian Nord's data challenge](https://github.com/deepskies/DeeplyUncertain-Public) for inspiration.

**CHALLENGE**: Devise and implement a way to obtain likelihoods $p(\vec{d} | z, \{z_{n}, \vec{d}_{n}\}_{N}, \pi)$ rather than posteriors $p(z | \vec{d}, \{z_{n}, \vec{d}_{n}\}_{N}, \pi)$ .

## 3. Evaluating the performance of estimated photo-$z$ posterior PDFs

Once we have estimated photo-$z$ posterior PDFs, we need a way to determine if they're actually any good.
Since the tutorial has only one method but multiple training/validation sets, that's all we can compare for now.

### Metrics of estimated photo-$z$ posteriors and true redshifts

First, let's try out a couple metrics of estimated photo-$z$ posteriors that do not require knowledge of the true photo-$z$ posteriors.
There's additional functionality for the case of having true redshifts but not true posteriors in [cdetools](https://github.com/tpospisi/cdetools) and [cde-diagnostics](https://github.com/zhao-david/CDE-diagnostics), but this should give a general idea.

```python
from cdetools import cde_loss, cdf_coverage, hpd_coverage
```

The Probability Integral Transform (PIT) is defined as
\begin{equation}
PIT = \int_{-\infty}^{z_{true}} p(z | \vec{d}, \{z_{n}, \vec{d}_{n}\}_{N}, \pi) dz .
\end{equation}
A histogram of PIT values is commonly used to assess how consistent a population of photo-$z$ PDFs are with the true redshifts.
Ideally, it would be a uniform distribution, meaning N% of galaxies have their true redshift within the Nth percentile of their estimated photo-$z$ posterior PDF.

```python
pit_values = cdf_coverage.cdf_coverage(cde_test, y_grid, y_test)
pit_values_bias = cdf_coverage.cdf_coverage(cde_test_bias, y_grid_bias, y_test)

plt.hist(pit_values, alpha=0.5, bins=100, label='representative')
plt.hist(pit_values_bias, alpha=0.5, bins=100, label='biased')
plt.ylim(0, 100)
plt.legend()
plt.xlabel('PIT')
```

The Highest Predictive Density (HPD)
\begin{equation}
HPD = \int_{z': p(z' | \vec{d}_{i}, \{z_{n}, \vec{d}_{n}\}_{N}, \pi) \geq p(z | \vec{d}_{i}, \{z_{n}, \vec{d}_{n}\}_{N}, \pi)} p(z' | \vec{d}_{i}, \{z_{n}, \vec{d}_{n}\}_{N}, \pi) dz
\end{equation}
is like the area of the PDF where it exceeds a given value.
Over a population, it would ideally be flat, like the PIT.
[A talk by David Zhao (CMU)](https://drive.google.com/file/d/1uvPtK_RcTUHEwt0ZYld41VKEPHnehWbN/view) has a lovely visualization of the HPD.

```python
hpd_cov = hpd_coverage.hpd_coverage(cde_test, y_grid, y_test)
hpd_cov_bias = hpd_coverage.hpd_coverage(cde_test_bias, y_grid_bias, y_test)

plt.hist(hpd_cov, alpha=0.5, bins=100, label='representative')
plt.hist(hpd_cov_bias, alpha=0.5, bins=100, label='biased')
plt.ylim(0, 100)
plt.legend()
plt.xlabel('HPD')
```

The CDE loss
\begin{equation}
\hat{L} = \frac{1}{K} \sum_{i=1}^{K} \int \left(p(z | \vec{d}_{i}, \{z_{n}, \vec{d}_{n}\}_{N}, \pi)\right)^{2} dz - \frac{2}{K} \sum_{i=1}^{K} p(z_{i} | \vec{d}_{i}, \{z_{n}, \vec{d}_{n}\}_{N}, \pi)
\end{equation}
approximates the true posterior from the estimated posterior evaluated at the true redshift.
It's explained quite well in [a talk by Nic Dalmasso (CMU)](https://www.dropbox.com/s/2r4tl4qv0iyqo9b/STAMPS_LSST_CDE_Tools_Presentation.pdf?dl=0).
A lower value indicates a better estimator.

```python
print(cde_loss.cde_loss(cde_test, y_grid, y_test))
print(cde_loss.cde_loss(cde_test_bias, y_grid_bias, y_test))
```

### Comparison of estimated and true photo-$z$ posterior PDFs

There are two categories of metrics of approximated and true PDF:
those that either rely upon or force the normalization condition $\int p(z) dz = 1$and those that evaluate differences between arbitrary functions.
`qp` [(Malz, et al 2018)](https://arxiv.org/abs/1806.00014) is a package for manipulating univariate PDFs under many parameterizations and includes a few comparison metrics.

The [original version](https://github.com/aimalz/qp) consistently enforced normalization but had limited functionality, whereas the [new version](https://github.com/LSSTDESC/qp) includes many more parameterizations whose usage is "at your own risk" in terms of possibly violating normalization.
We'll use the new version for the sake of speed but evaluating metrics using simplified functions ported from the old version due to a (hopefully transient) bug in the handling of large sets of posteriors.
The first step is to get both the true posteriors and the approximations evaluated on the same grid of redshifts.

_This content is adapted from the [qp demo](https://github.com/LSSTDESC/qp/blob/master/docs/notebooks/demo.ipynb), written by Alex Malz (GCCL@RUB), Phil Marshall (SLAC), and Eric Charles (SLAC), and [qp metrics demo](https://github.com/LSSTDESC/qp/blob/master/docs/notebooks/kld.ipynb), written by Alex Malz (GCCL@RUB) and Phil Marshall (SLAC)._

```python
# P = qp.Ensemble(qp.interp, data=dict(xvals=grid.reshape(grid.shape[0]), yvals=posteriors_test))
Q = qp.Ensemble(qp.interp, data=dict(xvals=y_grid.reshape(y_grid.shape[0]), yvals=cde_test))
Q_bias = qp.Ensemble(qp.interp, data=dict(xvals=y_grid.reshape(y_grid_bias.shape[0]), yvals=cde_test_bias))
grid, approx_pdf_on_grid = Q.gridded(grid)
grid, approx_pdf_on_grid_bias = Q_bias.gridded(grid)
```

The Kullback Leibler Divergence (KLD)
\begin{equation}
KLD = \int_{-\infty}^{\infty} p(z | \vec{d}) \log\left[\frac{p(z | \vec{d}, \{z_{n}, \vec{d}_{n}\}_{N}, \pi)}{p(z | \vec{d})}\right] dz
\end{equation}
is a directional measure of how much information is lost by using the estimated $p(z | \vec{d}, \{z_{n}, \vec{d}_{n}\}_{N}, \pi)$ instead of the true $p(z | \vec{d})$.
We want the KLD for each galaxy to be low.

```python
KLDs = np.array([qp.metrics.quick_kld(p, q, dx=np.mean(grid[1:] - grid[:-1])) for p, q in zip(posteriors_test, approx_pdf_on_grid)])
KLDs_bias = np.array([qp.metrics.quick_kld(p, q, dx=np.mean(grid[1:] - grid[:-1])) for p, q in zip(posteriors_test, approx_pdf_on_grid_bias)])

plt.hist(np.log(KLDs), alpha=0.5, bins=100, label='representative', density=True)
plt.hist(np.log(KLDs_bias), alpha=0.5, bins=100, label='biased', density=True)
plt.xlabel('KLD')
plt.legend()
plt.semilogy()
```

The root-mean-square-error (RMSE) is a symmetric measure commonly used to compare 1D functions.
**TODO: write it out?** Similarly, a lower value corresponds to a more closely approximating posterior PDF.

```python
RMSEs = np.array([qp.metrics.quick_rmse(p, q, N=granularity) for p, q in zip(posteriors_test, approx_pdf_on_grid)])
RMSEs_bias = np.array([qp.metrics.quick_rmse(p, q, N=granularity) for p, q in zip(posteriors_test, approx_pdf_on_grid_bias)])

plt.hist(np.log(RMSEs), alpha=0.5, bins=100, label='representative', density=True)
plt.hist(np.log(RMSEs_bias), alpha=0.5, bins=100, label='biased', density=True)
plt.xlabel('RMSE')
plt.legend()
```

### Your turn!

How else can we quantify the performance of estimators of aleatoric uncertainty?

**CHALLENGE**: Apply and visualize the local metrics of [Zhao, Dalmasso, Izbicki & Lee, 2021](https://arxiv.org/abs/2102.10473), or any other metrics not included in this demo; the [cde-diagnostics tutorial](https://github.com/zhao-david/CDE-diagnostics/blob/main/tutorial/tutorial-cde-diagnostics.ipynb), written by David Zhao (CMU), may be a good starting point.

**CHALLENGE**: Use the differences between estimated and true photo-$z$ posterior PDFs as a function of photometry to isolate the implicit prior of epistemic uncertainty from the aleatoric uncertainty.

# Your turn: How to participate in this challenge tl;dr

The organizers have identified three hack-able aspects to this data challenge:
- Quantify sensitivity of ML photo-z posterior estimators to training set non-representativity
- Implement and test additional AI methods for photo-z posterior estimation
- Implement/interpret additional metrics of posterior precision

Of course, there are many other possibilities for what to investigate based on this starting material!

To participate, clone this tutorial's [GitHub repo](https://github.com/aimalz/qtc2021), make or comment on [an issue](https://github.com/aimalz/qtc2021/issues) saying what you'd like to hack on, and work in a branch corresponding to that issue.
The only rule is that if your hack leads to a publication, you cite the repo (and, of course, invite those who contributed to your solution to be authors).

```python

```


```python
# This cell is added by sphinx-gallery
!pip install mrsimulator --quiet

%matplotlib inline

import mrsimulator
print(f'You are using mrsimulator v{mrsimulator.__version__}')
```

# Writing Custom methods (HahnEcho)

Writing custom methods using the Event objects.

Import the relevant modules

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

from mrsimulator import Simulator, SpinSystem, Site, Coupling
from mrsimulator.methods import Method1D
from mrsimulator.method.event import MixingEvent, SpectralEvent
```

For demonstration, we will create two spin systems, one with a single site and other
with two spin 1/2 sites.

```python
S1 = Site(
    isotope="1H",
    isotropic_chemical_shift=10,  # in ppm
    shielding_symmetric={"zeta": -80, "eta": 0.25},  # zeta in ppm
)
S2 = Site(isotope="1H", isotropic_chemical_shift=-10)
S12 = Coupling(
    site_index=[0, 1], isotropic_j=100, dipolar={"D": 2000, "eta": 0, "alpha": 0}
)

spin_system_1 = SpinSystem(sites=[S1], label="Uncoupled system")
spin_system_2 = SpinSystem(sites=[S1, S2], couplings=[S12], label="Coupled system")
```

**Create a custom method**

Writing a custom method is simply specifying an appropriate list of event objects per
spectral dimension. In this example, we are interested in a one-dimensional Hahnecho
method, and we use the generic `Method1D` class as a template. For a Hahnecho, we will
use two types of Event objects---SpectralEvent and MixingEvent.

A SpectralEvent object is where we sample the frequency contributions. The net
frequency along a given spectral dimension is a weighted average of the frequencies
from all SpectralEvent objects within a given SpectralDimension, i.e.,

\begin{align}f = \sum_j w_j \times \nu_j,\end{align}

where $f$ is the net averaged frequency along a spectral dimension, $w_j$
is the weight (the attribute ``fraction``), and $\nu_j$ is the frequency from
the $j^\text{th}$ SpectralEvent. The index $j$ runs over all spectral
events within a spectral dimension.

In the case of a one-dimensional Hahnecho method, the frequency is equally averaged
over two spectral events, corresponding to the symmetry pathway,

\begin{align}p = 1 \rightarrow -1.\end{align}

In the following code, we define the two SpectralEvent objects with fraction 0.5 and
the transition_query on channel-1 of P=[1] and P=[-1], respectively. Notice, the value
for the ``P`` attribute is a list. Here, it is a list with a single integer. The list
notation, ``[1]``, implies that the query selects all transitions where exactly one
spin is undergoing a $p=+1$ transition with the remaining spin at $p=0$.
A similar argument holds for ``[-1]`` query. By implementing query objects, we
decouple the method from the spin system, i.e., once a method is defined, it can be
used to simulate spectra from any given spin system. We will demonstrate this
momentarily by simulating a Hahnecho spectrum from single and two-site spin systems.

Besides the SpectralEvent, you may also notice a MixingEvent sandwitched in-between
the two SpectralEvent. A MixingEvent does not directly contribute to the frequencies.
As the name suggests, a mixing event is used for the mixing of transitions in a
multi-event method such as HahnEcho. In the following code, we define a mixing query
on channel-1 by setting the attributes ``tip_angle`` and ``phase`` to $\pi$ and
0, respectively. There two parameters are analogous to the pulse angle and phase.

```python
hahn_echo = Method1D(
    channels=["1H"],
    magnetic_flux_density=9.4,  # in T
    spectral_dimensions=[
        {
            "count": 512,
            "spectral_width": 2e4,  # in Hz
            "events": [
                SpectralEvent(fraction=0.5, transition_query=[{"ch1": {"P": [1]}}]),
                MixingEvent(mixing_query={"ch1": {"tip_angle": np.pi, "phase": 0}}),
                SpectralEvent(fraction=0.5, transition_query=[{"ch1": {"P": [-1]}}]),
            ],
        }
    ],
)
```

You may also visualize the method using the `plot` function.
hahn_echo.plot()

As mentioned before, a method object is decoupled from the spin system object. Notice,
when we get the transition pathways from this method for a single-site spin system, we
get a single transition pathway.

```python
print(hahn_echo.get_transition_pathways(spin_system_1))
```

In the case of a homonuclear two-site spin 1/2 spin system, the same method returns
four transition pathways.

```python
print(hahn_echo.get_transition_pathways(spin_system_2))
```

Create the Simulator object, add the method and spin system objects, and run the
simulation.

```python
sim = Simulator()
sim.spin_systems = [spin_system_1, spin_system_2]  # add the spin systems
sim.methods = [hahn_echo]  # add the method
sim.config.decompose_spectrum = "spin_system"

sim.run()
```

The simulation from each spin system is stored as a dependent variable within the
CSDM object. Use the `split` function to split the list of the dependent variables
into a list of CSDM objects.

```python
simulation_results = sim.methods[0].simulation.split()

# The plot of the two simulations.
fig, ax = plt.subplots(1, 2, figsize=(8.0, 3.0), subplot_kw={"projection": "csdm"})
for i in range(2):
    ax[i].plot(simulation_results[i].real, color="black", linewidth=1)
    ax[i].invert_xaxis()
plt.tight_layout()
plt.show()
```

Notice, in the single-site spin system, the hahn echo refocuses the isotropic chemical
shits and chemical shift anisotropies. The end result is a resonance at zero
frequency. In the case of the two homonuclear spin 1/2 coupled spin system, the Hahn
echo refocuses the isotropic chemical shits and chemical shift anisotropies, but not
the dipolar and `J` couplings.


```python
from IPython.display import Image
```

# CNTK 203: Reinforcement Learning Basics

Reinforcement learning (RL) is an area of machine learning inspired by behaviorist psychology, concerned with how [software agents](https://en.wikipedia.org/wiki/Software_agent) ought to take [actions](https://en.wikipedia.org/wiki/Action_selection) in an environment so as to maximize some notion of cumulative reward. In machine learning, the environment is typically formulated as a [Markov decision process](https://en.wikipedia.org/wiki/Markov_decision_process) (MDP) as many reinforcement learning algorithms for this context utilize [dynamic programming](https://en.wikipedia.org/wiki/Dynamic_programming) techniques.

In some machine learning settings, we do not have immediate access to labels, so we cannot rely on supervised learning techniques. If, however, there is something we can interact with and thereby get some feedback that tells us occasionally, whether our previous behavior was good or not, we can use RL to learn how to improve our behavior.

Unlike in supervised learning, in RL, labeled correct input/output pairs are never presented and sub-optimal actions are never explicitly corrected. This mimics many of the online learning paradigms which involves finding a balance between exploration (of conditions or actions never learnt before) and exploitation (of already learnt conditions or actions from previous encounters). Multi-arm bandit problems is one of the category of RL algorithms where exploration vs. exploitation trade-off have been thoroughly studied. See figure below for [reference](http://www.simongrant.org/pubs/thesis/3/2.html).

```python
# Figure 1
Image(url="https://cntk.ai/jup/polecart.gif", width=300, height=300)
```

**Problem**

We will use the [CartPole](https://gym.openai.com/envs/CartPole-v0) environment from OpenAI's [gym](https://github.com/openai/gym) simulator to teach a cart to balance a pole. As described in the link above, in the CartPole example, a pole is attached by an un-actuated joint to a cart, which moves along a frictionless track. The system is controlled by applying a force of +1 or -1 to the cart. A reward of +1 is provided for every timestep that the pole remains upright. The episode ends when the pole is more than 15 degrees from vertical, or the cart moves more than 2.4 units from the center. See figure below for reference.

**Goal**
Our goal is to prevent the pole from falling over as the cart moves with the pole in upright position (perpendicular to the cart) as the starting state. More specifically if the pole is less than 15 degrees from vertical while the cart is within 2.4 units of the center we will collect reward. In this tutorial, we will train till we learn a set of actions (policies) that lead to an average reward of 200 or more over last 50 batches.

In, RL terminology, the goal is to find _policies_ $a$, that maximize the _reward_ $r$ (feedback) through interaction with some environment (in this case the pole being balanced on the cart). So given a series of experiences $$s \xrightarrow{a} r, s'$$ we then can learn how to choose action $a$ in a given state $s$ to maximize the accumulated reward $r$ over time:
\begin{align}
Q(s,a) &= r_0 + \gamma r_1 + \gamma^2 r_2 + \ldots \newline
&= r_0 + \gamma \max_a Q^*(s',a)
\end{align}
where $\gamma \in [0,1)$ is the discount factor that controls how much we should value reward that is further away. This is called the [*Bellmann*-equation](https://en.wikipedia.org/wiki/Bellman_equation).

In this tutorial we will show how to model the state space, how to use the received reward to figure out which action yields the highest future reward.

We present two different popular approaches here:

**Deep Q-Networks**: DQNs have become famous in 2015 when they were successfully used to train how to play Atari just form raw pixels. We train neural network to learn the $Q(s,a)$ values (thus _Q-Network _). From these $Q$ functions values we choose the best action.

**Policy gradient**: This method directly estimates the policy (set of actions) in the network. The outcome is a learning of an ordered set of actions which leads to maximize reward by probabilistically choosing a subset of actions. In this tutorial, we learn the actions using a gradient descent approach to learn the policies.

In this tutorial, we focus how to implement RL in CNTK. We choose a straight forward shallow network. One can extend the approaches by replacing our shallow model with deeper networks that are introduced in other CNTK tutorials.

Additionally, this tutorial is in its early stages and will be evolving in future updates.

## Before we start...
Please run the following cell from the menu above or select the cell below and hit `Shift + Enter` to ensure the environment is ready. Verify that the following imports work in your notebook.

```python
from __future__ import print_function
from __future__ import division
import matplotlib.pyplot as plt
from matplotlib import style
import numpy as np
import pandas as pd
import seaborn as sns

style.use('ggplot')
%matplotlib inline
```

We use the following construct to install the OpenAI gym package if it is not installed. For users new to Jupyter environment, this construct can be used to install any python package.

```python
try:
    import gym
except:
    !pip install gym
    import gym
```

### Select the notebook run mode

There are two run modes:
- *Fast mode*: `isFast` is set to `True`. This is the default mode for the notebooks, which means we train for fewer iterations or train / test on limited data. This ensures functional correctness of the notebook though the models produced are far from what a completed training would produce.

- *Slow mode*: We recommend the user to set this flag to `False` once the user has gained familiarity with the notebook content and wants to gain insight from running the notebooks for a longer period with different parameters for training.

```python
isFast = True
```

# CartPole: Data and Environment

We will use the [CartPole](https://gym.openai.com/envs/CartPole-v0) environment from OpenAI's [gym](https://github.com/openai/gym) simulator to teach a cart to balance a pole. Please follow the links to get more details.

In every time step, the agent
 * gets an observation $(x, \dot{x}, \theta, \dot{\theta})$, corresponding to *cart position*, *cart velocity*, *pole angle with the vertical*, *pole angular velocity*,
 * performs an action `LEFT` or `RIGHT`, and
 * receives
  * a reward of +1 for having survived another time step, and
  * a new state $(x', \dot{x}', \theta', \dot{\theta}')$

The episode ends, if
 * the pole is more than 15 degrees from vertical and/or
 * the cart is moving more than 2.4 units from center.

The task is considered done, if
 * the agent achieved and averaged reward of 200 over the last 50 episodes (if you manage to get a reward of 200 averaged over the last 100 episode you can consider submitting it to OpenAI).

In fast mode these targets are relaxed.

# Part 1: DQN

After a transition $(s,a,r,s')$, we are trying to move our value function $Q(s,a)$ closer to our target $r+\gamma \max_{a'}Q(s',a')$, where $\gamma$ is a discount factor for future rewards and ranges in value between 0 and 1.

DQNs
 * learn the _Q-function_ that maps observation (state, action) to a `score`
 * use memory replay (previously recorded $Q$ values corresponding to different $(s,a)$ to decorrelate experiences (sequence state transitions)
 * use a second network to stabilize learning (*not* part of this tutorial)

### Setting up the model
\begin{equation}
l_1 = relu( x W_1 + b_1) \\
Q(s,a) = l_1 W_2 + b_2 \\
\end{equation}

We will start with a slightly modified version for Keras, https://github.com/jaara/AI-blog/blob/master/CartPole-basic.py, published by Jaromír Janisch in his [AI blog](https://jaromiru.com/2016/09/27/lets-make-a-dqn-theory/), and will then incrementally convert it to use CNTK.

We use a simple two-layer densely connected network, for simpler illustrations. More advance networks can be substituted.

**CNTK** concepts: The commented out code is meant to be an illustration of the similarity of concepts between CNTK API/abstractions against Keras.

```python
import numpy
import math
import os
import random

import cntk as C
```

In the block below, we check if we are running this notebook in the CNTK internal test machines by looking for environment variables defined there. We then select the right target device (GPU vs CPU) to test this notebook. In other cases, we use CNTK's default policy to use the best available device (GPU, if available, else CPU).

```python
# Select the right target device when this notebook is being tested:
if 'TEST_DEVICE' in os.environ:
    if os.environ['TEST_DEVICE'] == 'cpu':
        C.device.try_set_default_device(C.device.cpu())
    else:
        C.device.try_set_default_device(C.device.gpu(0))
```

STATE_COUNT = 4 (corresponding to $(x, \dot{x}, \theta, \dot{\theta})$),

ACTION_COUNT = 2 (corresponding to `LEFT` or `RIGHT`)

```python
env = gym.make('CartPole-v0')

STATE_COUNT  = env.observation_space.shape[0]
ACTION_COUNT = env.action_space.n

STATE_COUNT, ACTION_COUNT
```

    [2017-05-16 20:19:32,618] Making new env: CartPole-v0

    (4, 2)

Note: in the cell below we highlight how one would do it in Keras. And a marked similarity with CNTK. While CNTK allows for more compact representation, we present a slightly verbose illustration for ease of learning.

Additionally, you will note that, CNTK model doesn't need to be compiled explicitly and is implicitly done when data is processed during training.

CNTK effectively uses available memory on the system between minibatch execution. Thus the learning rates are stated as **rates per sample** instead of **rates per minibatch** (as with other toolkits).

```python
# Targetted reward
REWARD_TARGET = 30 if isFast else 200
# Averaged over these these many episodes
BATCH_SIZE_BASELINE = 20 if isFast else 50

H = 64 # hidden layer size

class Brain:
    def __init__(self):
        self.params = {}
        self.model, self.trainer, self.loss = self._create()
        # self.model.load_weights("cartpole-basic.h5")

    def _create(self):
        observation = C.sequence.input_variable(STATE_COUNT, np.float32, name="s")
        q_target = C.sequence.input_variable(ACTION_COUNT, np.float32, name="q")

        # Following a style similar to Keras
        l1 = C.layers.Dense(H, activation=C.relu)
        l2 = C.layers.Dense(ACTION_COUNT)
        unbound_model = C.layers.Sequential([l1, l2])
        model = unbound_model(observation)

        self.params = dict(W1=l1.W, b1=l1.b, W2=l2.W, b2=l2.b)

        # loss='mse'
        loss = C.reduce_mean(C.square(model - q_target), axis=0)
        meas = C.reduce_mean(C.square(model - q_target), axis=0)

        # optimizer
        lr = 0.00025
        lr_schedule = C.learning_rate_schedule(lr, C.UnitType.minibatch)
        learner = C.sgd(model.parameters, lr_schedule, gradient_clipping_threshold_per_sample=10)
        trainer = C.Trainer(model, (loss, meas), learner)

        # CNTK: return trainer and loss as well
        return model, trainer, loss

    def train(self, x, y, epoch=1, verbose=0):
        #self.model.fit(x, y, batch_size=64, nb_epoch=epoch, verbose=verbose)
        arguments = dict(zip(self.loss.arguments, [x,y]))
        updated, results =self.trainer.train_minibatch(arguments, outputs=[self.loss.output])

    def predict(self, s):
        return self.model.eval([s])
```

The `Memory` class stores the different states, actions and rewards.

```python
class Memory:   # stored as ( s, a, r, s_ )
    samples = []

    def __init__(self, capacity):
        self.capacity = capacity

    def add(self, sample):
        self.samples.append(sample)

        if len(self.samples) > self.capacity:
            self.samples.pop(0)

    def sample(self, n):
        n = min(n, len(self.samples))
        return random.sample(self.samples, n)
```

The `Agent` uses the `Brain` and `Memory` to replay the past actions to choose optimal set of actions that maximize the rewards.

```python
MEMORY_CAPACITY = 100000
BATCH_SIZE = 64

GAMMA = 0.99 # discount factor

MAX_EPSILON = 1
MIN_EPSILON = 0.01 # stay a bit curious even when getting old
LAMBDA = 0.0001    # speed of decay

class Agent:
    steps = 0
    epsilon = MAX_EPSILON

    def __init__(self):
        self.brain = Brain()
        self.memory = Memory(MEMORY_CAPACITY)

    def act(self, s):
        if random.random() < self.epsilon:
            return random.randint(0, ACTION_COUNT-1)
        else:
            return numpy.argmax(self.brain.predict(s))

    def observe(self, sample):  # in (s, a, r, s_) format
        self.memory.add(sample)

        # slowly decrease Epsilon based on our eperience
        self.steps += 1
        self.epsilon = MIN_EPSILON + (MAX_EPSILON - MIN_EPSILON) * math.exp(-LAMBDA * self.steps)

    def replay(self):
        batch = self.memory.sample(BATCH_SIZE)
        batchLen = len(batch)

        no_state = numpy.zeros(STATE_COUNT)

        # CNTK: explicitly setting to float32
        states = numpy.array([ o[0] for o in batch ], dtype=np.float32)
        states_ = numpy.array([(no_state if o[3] is None else o[3]) for o in batch ], dtype=np.float32)

        p = agent.brain.predict(states)
        p_ = agent.brain.predict(states_)

        # CNTK: explicitly setting to float32
        x = numpy.zeros((batchLen, STATE_COUNT)).astype(np.float32)
        y = numpy.zeros((batchLen, ACTION_COUNT)).astype(np.float32)

        for i in range(batchLen):
            s, a, r, s_ = batch[i]

            # CNTK: [0] because of sequence dimension
            t = p[0][i]
            if s_ is None:
                t[a] = r
            else:
                t[a] = r + GAMMA * numpy.amax(p_[0][i])

            x[i] = s
            y[i] = t

        self.brain.train(x, y)
```

### Brain surgery

As any learning experiences, we expect to see the initial state of actions to be wild exploratory and over the iterations the system learns the range of actions that yield longer runs and collect more rewards. The tutorial below implements the [$\epsilon$-greedy](https://en.wikipedia.org/wiki/Reinforcement_learning) approach.

```python
def plot_weights(weights, figsize=(7,5)):
    '''Heat map of weights to see which neurons play which role'''
    sns.set(style="white")
    f, ax = plt.subplots(len(weights), figsize=figsize)
    cmap = sns.diverging_palette(220, 10, as_cmap=True)

    for i, data in enumerate(weights):
        axi = ax if len(weights)==1 else ax[i]
        if isinstance(data, tuple):
            w, title = data
            axi.set_title(title)
        else:
            w = data

        sns.heatmap(w.asarray(), cmap=cmap, square=True, center=True, #annot=True,
                    linewidths=.5, cbar_kws={"shrink": .25}, ax=axi)
```

### Exploration - exploitation trade-off

Note that the initial $\epsilon$ is set to 1 which implies we are entirely exploring but as steps increase we reduce exploration and start leveraging the learnt space to collect rewards (a.k.a. exploitation) as well.

```python
def epsilon(steps):
    return MIN_EPSILON + (MAX_EPSILON - MIN_EPSILON) * np.exp(-LAMBDA * steps)
plt.plot(range(10000), [epsilon(x) for x in range(10000)], 'r')
plt.xlabel('step');plt.ylabel('$\epsilon$')
```

We are now ready to train our agent using **DQN**. Note this would take anywhere between 2-10 min and we stop whenever the learner hits the average reward of 200 over past 50 batches. One would get better results if they could train the learner until say one hits a reward of 200 or higher for say larger number of runs. This is left as an exercise.

```python
TOTAL_EPISODES = 2000 if isFast else 3000

def run(agent):
    s = env.reset()
    R = 0

    while True:
        # Uncomment the line below to visualize the cartpole
        # env.render()

        # CNTK: explicitly setting to float32
        a = agent.act(s.astype(np.float32))

        s_, r, done, info = env.step(a)

        if done: # terminal state
            s_ = None

        agent.observe((s, a, r, s_))
        agent.replay()

        s = s_
        R += r

        if done:
            return R

agent = Agent()

episode_number = 0
reward_sum = 0
while episode_number < TOTAL_EPISODES:
    reward_sum += run(agent)
    episode_number += 1
    if episode_number % BATCH_SIZE_BASELINE == 0:
        print('Episode: %d, Average reward for episode %f.' % (episode_number,
                                                               reward_sum / BATCH_SIZE_BASELINE))
        if episode_number%200==0:
            plot_weights([(agent.brain.params['W1'], 'Episode %i $W_1$'%episode_number)], figsize=(14,5))
        if reward_sum / BATCH_SIZE_BASELINE > REWARD_TARGET:
            print('Task solved in %d episodes' % episode_number)
            plot_weights([(agent.brain.params['W1'], 'Episode %i $W_1$'%episode_number)], figsize=(14,5))
            break
        reward_sum = 0
agent.brain.model.save('dqn.mod')
```

If you run it, you should see something like
```
Episode: 20, Average reward for episode 20.700000.
Episode: 40, Average reward for episode 20.150000.
Episode: 60, Average reward for episode 21.100000.
...
Episode: 960, Average reward for episode 20.150000.
Episode: 980, Average reward for episode 26.700000.
Episode: 1000, Average reward for episode 32.900000.
Task solved in 1000 episodes
```

#### Task 1.1
Rewrite the model without using the layer lib.
#### Task 1.2
Play with different [learners](https://cntk.ai/pythondocs/cntk.learner.html#module-cntk.learner). Which one works better? Worse? Think about which parameters you would need to adapt when switching from one learner to the other.

### Running the DQN model

```python
env = gym.make('CartPole-v0')

num_episodes = 10  # number of episodes to run

modelPath = 'dqn.mod'
root = C.load_model(modelPath)

for i_episode in range(num_episodes):
    print(i_episode)
    observation = env.reset()  # reset environment for new episode
    done = False
    while not done:
        if not 'TEST_DEVICE' in os.environ:
            env.render()
        action = np.argmax(root.eval([observation.astype(np.float32)]))
        observation, reward, done, info  = env.step(action)
```

    [2017-05-16 20:21:24,188] Making new env: CartPole-v0

    0
    1
    2
    3
    4
    5
    6
    7
    8
    9

# Part 2: Policy gradient
**Goal:**
\begin{equation}\text{maximize } E [R | \pi_\theta]
\end{equation}

**Approach:**
1. Collect experience (sample a bunch of trajectories through $(s,a)$ space)
2. Update the policy so that _good_ experiences become more probable

**Difference to DQN: **
 * we don't consider single $(s,a,r,s')$ transitions, but rather use whole episodes for the gradient updates
 * our parameters directly model the policy (output is an action probability), whereas in DQN they model the value function (output is raw score)

#### Rewards:
Remember, we get +1 reward for every time step, in which we still were in the game.

The problem: we normally do not know, which action led to a continuation of the game, and which was actually a bad one. Our simple heuristic: actions in the beginning of the episode are good, and those towards the end are likely bad (they led to losing the game after all).

```python
def discount_rewards(r, gamma=0.999):
    """Take 1D float array of rewards and compute discounted reward """
    discounted_r = np.zeros_like(r)
    running_add = 0
    for t in reversed(range(0, r.size)):
        running_add = running_add * gamma + r[t]
        discounted_r[t] = running_add
    return discounted_r
```

```python
discounted_epr = discount_rewards(np.ones(10))
f, ax = plt.subplots(1, figsize=(5,2))
sns.barplot(list(range(10)), discounted_epr, color="steelblue")
```

We normalize the rewards so that they tank below zero towards the end. gamma controls how late the rewards tank.

```python
discounted_epr_cent = discounted_epr - np.mean(discounted_epr)
discounted_epr_norm = discounted_epr_cent/np.std(discounted_epr_cent)
f, ax = plt.subplots(1, figsize=(5,2))
sns.barplot(list(range(10)), discounted_epr_norm, color="steelblue")
```

```python
discounted_epr = discount_rewards(np.ones(10), gamma=0.5)
discounted_epr_cent = discounted_epr - np.mean(discounted_epr)
discounted_epr_norm = discounted_epr_cent/np.std(discounted_epr_cent)
f, ax = plt.subplots(2, figsize=(5,3))
sns.barplot(list(range(10)), discounted_epr, color="steelblue", ax=ax[0])
sns.barplot(list(range(10)), discounted_epr_norm, color="steelblue", ax=ax[1])
```

### Setting up the model
\begin{equation}
l_1 = relu( x W_1 + b_1) \\
l_2 = l_1 W_2 + b_2 \\
\pi(a|s) = sigmoid(l_2)
\end{equation}

Note: in policy gradient approach, the output of the dense layer is mapped into to a 0-1 range via the sigmoid function.

```python
TOTAL_EPISODES = 2000 if isFast else 10000

D = 4  # input dimensionality
H = 10 # number of hidden layer neurons

observations = C.sequence.input_variable(STATE_COUNT, np.float32, name="obs")

W1 = C.parameter(shape=(STATE_COUNT, H), init=C.glorot_uniform(), name="W1")
b1 = C.parameter(shape=H, name="b1")
layer1 = C.relu(C.times(observations, W1) + b1)

W2 = C.parameter(shape=(H, ACTION_COUNT), init=C.glorot_uniform(), name="W2")
b2 = C.parameter(shape=ACTION_COUNT, name="b2")
score = C.times(layer1, W2) + b2
# Until here it was similar to DQN

probability = C.sigmoid(score, name="prob")
```

**Policy Search**: The optimal policy search can be carried out with either gradient free approaches or by computing gradients over the policy space ($\pi_\theta$) which is parameterized by $\theta$. In this tutorial, we use the classic forward (`loss.forward`) and back (`loss.backward`) propagation of errors over the parameterized space $\theta$. In this case, $\theta = \{W_1, b_1, W_2, b_2\}$, our model parameters.

```python
input_y = C.sequence.input_variable(1, np.float32, name="input_y")
advantages = C.sequence.input_variable(1, np.float32, name="advt")

loss = -C.reduce_mean(C.log(C.square(input_y - probability) + 1e-4) * advantages, axis=0, name='loss')

lr = 0.001
lr_schedule = C.learning_rate_schedule(lr, C.UnitType.sample)
sgd = C.sgd([W1, W2], lr_schedule)

gradBuffer = dict((var.name, np.zeros(shape=var.shape)) for var in loss.parameters if var.name in ['W1', 'W2', 'b1', 'b2'])

xs, hs, label, drs = [], [], [], []
running_reward = None
reward_sum = 0
episode_number = 1

observation = env.reset()

while episode_number <= TOTAL_EPISODES:
    x = np.reshape(observation, [1, STATE_COUNT]).astype(np.float32)

    # Run the policy network and get an action to take.
    prob = probability.eval(arguments={observations: x})[0][0][0]
    action = 1 if np.random.uniform() < prob else 0

    xs.append(x)  # observation
    # grad that encourages the action that was taken to be taken

    y = 1 if action == 0 else 0  # a "fake label"
    label.append(y)

    # step the environment and get new measurements
    observation, reward, done, info = env.step(action)
    reward_sum += float(reward)

    # Record reward (has to be done after we call step() to get reward for previous action)
    drs.append(float(reward))

    if done:
        # Stack together all inputs, hidden states, action gradients, and rewards for this episode
        epx = np.vstack(xs)
        epl = np.vstack(label).astype(np.float32)
        epr = np.vstack(drs).astype(np.float32)
        xs, label, drs = [], [], []  # reset array memory

        # Compute the discounted reward backwards through time.
        discounted_epr = discount_rewards(epr)
        # Size the rewards to be unit normal (helps control the gradient estimator variance)
        discounted_epr -= np.mean(discounted_epr)
        discounted_epr /= np.std(discounted_epr)

        # Forward pass
        arguments = {observations: epx, input_y: epl, advantages: discounted_epr}
        state, outputs_map = loss.forward(arguments, outputs=loss.outputs,
                                          keep_for_backward=loss.outputs)

        # Backward psas
        root_gradients = {v: np.ones_like(o) for v, o in outputs_map.items()}
        vargrads_map = loss.backward(state, root_gradients, variables=set([W1, W2]))

        for var, grad in vargrads_map.items():
            gradBuffer[var.name] += grad

        # Wait for some batches to finish to reduce noise
        if episode_number % BATCH_SIZE_BASELINE == 0:
            grads = {W1: gradBuffer['W1'].astype(np.float32),
                     W2: gradBuffer['W2'].astype(np.float32)}
            updated = sgd.update(grads, BATCH_SIZE_BASELINE)

            # reset the gradBuffer
            gradBuffer = dict((var.name, np.zeros(shape=var.shape))
                              for var in loss.parameters if var.name in ['W1', 'W2', 'b1', 'b2'])

            print('Episode: %d. Average reward for episode %f.' % (episode_number, reward_sum / BATCH_SIZE_BASELINE))

            if reward_sum / BATCH_SIZE_BASELINE > REWARD_TARGET:
                print('Task solved in: %d ' % episode_number)
                break

            reward_sum = 0

        observation = env.reset()  # reset env
        episode_number += 1
probability.save('pg.mod')
```

    Episode: 20. Average reward for episode 21.800000.
    Episode: 40. Average reward for episode 23.450000.
    Episode: 60. Average reward for episode 20.650000.
    Episode: 80. Average reward for episode 23.900000.
    Episode: 100. Average reward for episode 17.750000.
    Episode: 120. Average reward for episode 27.550000.
    Episode: 140. Average reward for episode 20.300000.
    Episode: 160. Average reward for episode 20.450000.
    Episode: 180. Average reward for episode 14.700000.
    Episode: 200. Average reward for episode 18.900000.
    Episode: 220. Average reward for episode 18.350000.
    Episode: 240. Average reward for episode 20.050000.
    Episode: 260. Average reward for episode 19.600000.
    Episode: 280. Average reward for episode 22.150000.
    Episode: 300. Average reward for episode 16.000000.
    Episode: 320. Average reward for episode 21.300000.
    Episode: 340. Average reward for episode 19.250000.
    Episode: 360. Average reward for episode 17.550000.
    Episode: 380. Average reward for episode 20.050000.
    Episode: 400. Average reward for episode 20.250000.
    Episode: 420. Average reward for episode 20.200000.
    Episode: 440. Average reward for episode 21.100000.
    Episode: 460. Average reward for episode 23.400000.
    Episode: 480. Average reward for episode 20.100000.
    Episode: 500. Average reward for episode 18.400000.
    Episode: 520. Average reward for episode 15.550000.
    Episode: 540. Average reward for episode 16.850000.
    Episode: 560. Average reward for episode 18.000000.
    Episode: 580. Average reward for episode 23.650000.
    Episode: 600. Average reward for episode 16.550000.
    Episode: 620. Average reward for episode 21.200000.
    Episode: 640. Average reward for episode 19.100000.
    Episode: 660. Average reward for episode 25.150000.
    Episode: 680. Average reward for episode 22.400000.
    Episode: 700. Average reward for episode 19.150000.
    Episode: 720. Average reward for episode 19.900000.
    Episode: 740. Average reward for episode 26.500000.
    Episode: 760. Average reward for episode 20.500000.
    Episode: 780. Average reward for episode 19.300000.
    Episode: 800. Average reward for episode 21.250000.
    Episode: 820. Average reward for episode 27.900000.
    Episode: 840. Average reward for episode 25.600000.
    Episode: 860. Average reward for episode 27.250000.
    Episode: 880. Average reward for episode 27.000000.
    Episode: 900. Average reward for episode 29.900000.
    Episode: 920. Average reward for episode 28.900000.
    Episode: 940. Average reward for episode 30.200000.
    Task solved in: 940

# Solutions
#### Solution 1.1

```python
observation = C.sequence.input_variable(STATE_COUNT, np.float32, name="s")

W1 = C.parameter(shape=(STATE_COUNT, H), init=C.glorot_uniform(), name="W1")
b1 = C.parameter(shape=H, name="b1")
layer1 = C.relu(C.times(observation, W1) + b1)

W2 = C.parameter(shape=(H, ACTION_COUNT), init=C.glorot_uniform(), name="W2")
b2 = C.parameter(shape=ACTION_COUNT, name="b2")
model = C.times(layer1, W2) + b2
W1.shape, b1.shape, W2.shape, b2.shape, model.shape
```

    ((4, 10), (10,), (10, 2), (2,), (2,))

```python
# Correspoding layers implementation - Preferred solution
def create_model(input):
    with C.layers.default_options(init=C.glorot_uniform()):
        z = C.layers.Sequential([C.layers.Dense(H, name="layer1"),
                                 C.layers.Dense(ACTION_COUNT, name="layer2")])
        return z(input)

model = create_model(observation)

model.layer1.W.shape, model.layer1.b.shape, model.layer2.W.shape, model.layer2.b.shape, model.shape
```

    ((4, 10), (10,), (10, 2), (2,), (2,))

```python

```