Introduction to Jet Images and Computer Vision

Michela Paganini - Yale University

High Energy Phenomenology, Experiment and Cosmology Seminar Series

Material:
 1. Ben Nachman's plenary talk at ACAT 2017 (slides)
 1. My talk on Generative Adversarial Networks for jet images at the 2017 IML Workshop (video)
 1. Jannicke Pearkes's talk on boosted top tagging with jet images at the 2017 IML Workshop (video)
 1. Michael Kagan's overview talk at LHCP 2017 (slides)
 1. ATLAS PUB Note on quark vs gluon tagging with jet images(note)
 1. Lynn Huynh's summer report on jet image pre-processing (write-up)
 1. Ben Nachman's talk at DataScience@LHC 2015 (video)

# Dataset

[](https://doi.org/10.17632/4r4v785rgx.1)

Although the dataset was released in conjunction with the arXiv publication of our [work on Generative Adversarial Networks for jet images](https://arxiv.org/pdf/1701.05927.pdf), it was previously used in the original "[Jet Images -- Deep Learning Edition](https://link.springer.com/article/10.1007%2FJHEP07%282016%29069)" work on jet image classification. Feel free to explore the dataset and use it for any project you have in mind (please cite the dataset and relevant publications explaining its generation!)

### Download dataset from Zenodo

```python
import os
from keras.utils.data_utils import get_file

# Info for downloading the dataset from Zenodo
MD5_HASH = 'f9b11c46b6a0ff928bec2eccf865ecf0'
DATAFILE = 'jet-images_Mass60-100_pT250-300_R1.25_Pix25.hdf5'
URL_TEMPLATE = 'https://zenodo.org/record/{record}/files/{filename}'

print('[INFO] MD5 verification: {}'.format(MD5_HASH))

datafile = get_file(
    fname=DATAFILE,
    origin=URL_TEMPLATE.format(record=269622, filename=DATAFILE),
    md5_hash=MD5_HASH
)
```

    Using TensorFlow backend.

    [INFO] MD5 verification: f9b11c46b6a0ff928bec2eccf865ecf0

# Jets at the LHC

Jets are the observable result of quarks and gluons scattering at high energy. A collimated stream of
protons and other hadrons forms in the direction of the initiating quark or gluon. Clusters of such
particles are called jets.

# Jet Images

Mature field of research! (image courtesy of B.P.Nachman)

### What is a jet image?

> A jet image is a two-dimensional representation of the radiation pattern within a jet: the distribution of the locations and energies of the jet’s constituent particles. The jet image consists of a regular grid of pixels in η×φ.

Advantages of this data format include: easy visual inspection, fixed-length representation, suitable for application of computer vision techniques.

### Pre-processing

In the dataset we will be using today:

> The finite granularity of a calorimeter is simulated with a regular 0.1×0.1 grid in η and φ. The energy of each calorimeter cell is given by the sum of the energies of all particles incident on the cell. Cells with positive energy are assigned to jets using the anti-kt clustering algorithm with a radius parameter of R = 1.0 via the software package FastJet 3.2.1.

> To mitigate the contribution from the underlying event, jets are are trimmed by re-clustering the constituents into R = 0.3 kt subjets and dropping those which have less than 5% of the transverse momentum of the parent jet. Trimming also reduces the impact of pileup: multiple proton-proton collisions occurring in the same event as the hard-scatter process. Jet images are formed by translating the η and φ of all constituents of a given jet so that its highest pT subjet is centered at the origin.

> A rectangular grid of η × φ ∈ [−1.25, 1.25] × [−1.25, 1.25] with 0.1 × 0.1 pixels centered at the origin
forms the basis of the jet image. The intensity of each pixel is the pT corresponding to the energy
and pseudorapditiy of the constituent calorimeter cell, pT = E_cell/ cosh(η_cell). The radiation pattern
is symmetric about the origin of the jet image and so the images are rotated. The subjet with the
second highest pT (or, in its absence, the direction of the first principle component) is placed at an
angle of −π/2 with respect to the η − φ axes. Finally, a parity transform about the vertical axis is
applied if the left side of the image has more energy than the right side.

  Learning Particle Physics by Example: Location-Aware Generative Adversarial Networks for Physics Synthesis

  [arXiv:1701.05927](https://arxiv.org/pdf/1701.05927.pdf)

References:
* Section 3 of [arXiv:1511.05190](https://link.springer.com/content/pdf/10.1007%2FJHEP07%282016%29069.pdf)
* Appendix B of [arXiv:1701.05927](https://arxiv.org/pdf/1701.05927.pdf)

### Uniqueness with respect to natural images in ML literature

* Sparse (low occupancy)
* High dynamic range (pixel intensity represents pT of particles and spans several orders of magnitude)
* Pixel activations and positions are physically meaningful
* Small variations can drastically modify physical properties of a jet

# Hands-on tutorial

```python
import h5py
import numpy as np
import os
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from matplotlib.colors import LogNorm, Normalize
%matplotlib inline
```

```python
# number of images to load
nb_points = 800000
```

```python
# open hdf5 data file
d = h5py.File(datafile, 'r')
```

```python
# content of the dataset
d.items()
```

    [(u'image', <HDF5 dataset "image": shape (872666, 25, 25), type "<f4">),
     (u'jet_delta_R', <HDF5 dataset "jet_delta_R": shape (872666,), type "<f4">),
     (u'jet_eta', <HDF5 dataset "jet_eta": shape (872666,), type "<f4">),
     (u'jet_mass', <HDF5 dataset "jet_mass": shape (872666,), type "<f4">),
     (u'jet_phi', <HDF5 dataset "jet_phi": shape (872666,), type "<f4">),
     (u'jet_pt', <HDF5 dataset "jet_pt": shape (872666,), type "<f4">),
     (u'signal', <HDF5 dataset "signal": shape (872666,), type "<f4">),
     (u'tau_1', <HDF5 dataset "tau_1": shape (872666,), type "<f4">),
     (u'tau_2', <HDF5 dataset "tau_2": shape (872666,), type "<f4">),
     (u'tau_21', <HDF5 dataset "tau_21": shape (872666,), type "<f4">),
     (u'tau_3', <HDF5 dataset "tau_3": shape (872666,), type "<f4">),
     (u'tau_32', <HDF5 dataset "tau_32": shape (872666,), type "<f4">)]

```python
# extract a random subset of samples
ix = range(d['image'].shape[0]) # get indices
np.random.shuffle(ix) # shuffle them
ix = ix[:nb_points] # select out nb_points

# extract data from dataset
images, labels = d['image'][:][ix], d['signal'][:][ix]
mass = d['jet_mass'][:][ix]
delta_R = d['jet_delta_R'][:][ix]
tau_21 = d['tau_21'][:][ix]
```

You can compute the quantities above directly from the images:

\begin{align}
&p_\text{T}^2(I) =\left(\sum_{i=0}^{N} I_i\cos(\phi_i)\right)^2+\left(\sum_{i=0}^{N} I_i\sin(\phi_i)\right)^2
\label{eq:pt}
\\
&m^2(I) = \left(\sum_{i=0}^{N} I_i\right)^2-p_\text{T}^2(I)-\left(\sum_{i=0}^{N} I_i\sinh(\eta_i)\right)^2
\label{eq:m}
\\
&\tau_{21}(I)=\frac{\tau_2(I)}{\tau_1(I)},
\label{eq:tau21}
\end{align}
where:
\begin{equation*}
\tau_{n}(I)\propto\sum_{i=0}^{N} I_i \min_{1\leq a\leq n}\left\{\sqrt{\left(\eta_i-\eta_a\right)^2+\left(\phi_i-\phi_a\right)^2}\right\}
\end{equation*}

```python
def plot_jet_image(content, output_fname=None, vmin=1e-6, vmax=300, title=''):
    '''
    Function to help you visualize a jet image on a log scale
    Args:
    -----
        content : numpy array of dimensions 25x25, first arg to imshow,
            content of the image
            e.g.: images.mean(axis=0) --> the average image

        output_fname : string, name of the output file where the plot will be
            saved.

        vmin : (default = 1e-6) float, lower bound of the pixel intensity
            scale before saturation

        vmax : (default = 300) float, upper bound of the pixel intensity
            scale before saturation

        title : (default = '') string, title of the plot, to be displayed
            on top of the image
    '''
    fig, ax = plt.subplots(figsize=(7, 6))
    extent = [-1.25, 1.25, -1.25, 1.25]

    im = ax.imshow(content, interpolation='nearest',
                   norm=LogNorm(vmin=vmin, vmax=vmax), extent=extent)

    cbar = plt.colorbar(im, fraction=0.05, pad=0.05)
    cbar.set_label(r'Pixel $p_T$ (GeV)', y=0.85)
    plt.xlabel(r'[Transformed] Pseudorapidity $(\eta)$')
    plt.ylabel(r'[Transformed] Azimuthal Angle $(\phi)$')
    plt.title(title)

    if output_fname is None:
        plt.savefig(output_fname)
```

```python
def plot_diff_jet_image(content, output_fname=None, extr=None, title='',
                        cmap=matplotlib.cm.seismic):
    '''
    Function to help you visualize the difference between two sets of jet
        images on a linear scale
    Args:
    -----
        content : numpy array of dimensions 25x25, first arg to imshow,
            content of the image
            e.g.: sig_images.mean(axis=0) - bkg_images.mean(axis=0)
        output_fname : string, name of the output file where the plot will be
            saved.
        extr : (default = None) float, magnitude of the upper and lower
            bounds of the pixel intensity scale before saturation (symmetric
            around 0)
        title : (default = '') string, title of the plot, to be displayed on
            top of the image
        cmap : (default = matplotlib.cm.PRGn_r) matplotlib colormap, ideally
            white in the middle
    '''
    fig, ax = plt.subplots(figsize=(6, 6))
    extent = [-1.25, 1.25, -1.25, 1.25]
    if extr == None:
        extr = max(abs(content.min()), abs(content.max()))
    im = ax.imshow(
        content,
        interpolation='nearest',
        norm=Normalize(vmin=-extr, vmax=+extr), extent=extent,
        cmap=cmap
    )

    plt.colorbar(im, fraction=0.05, pad=0.05)
    plt.xlabel(r'[Transformed] Pseudorapidity $(\eta)$')
    plt.ylabel(r'[Transformed] Azimuthal Angle $(\phi)$')
    plt.title(title)

    if output_fname:
        plt.savefig(output_fname)

```

```python
# visualize a jet image
plot_jet_image(images[0])
```

```python
# visualize the average jet image
plot_jet_image(images.mean(axis=0))
```

```python
# visualize the difference between the average signal and the average background image
plot_diff_jet_image(
    images[labels == 1].mean(axis=0) - images[labels == 0].mean(axis=0)
)
```

## Looking at Physics features

```python
def plot_physics_feature(feature_name, feature, labels, bins=None, output_fname=None):

    # if bins are not defined when function is called, define them here
    if not bins:
        bins = np.linspace(feature.min(), feature.max(), 50)

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

    _ = plt.hist(feature[labels == 1], bins=bins, histtype='step',
                 label=r"Signal ($W' \rightarrow WZ$)",
                 normed=True, color='red')

    _ = plt.hist(feature[labels == 0], bins=bins, histtype='step',
                 label=r'Background (QCD dijets)', normed=True, color='blue')

    plt.xlabel(r'Discretized {} of Jet Image'.format(feature_name))
    plt.ylabel(r'Units normalized to unit area')
    plt.legend()

    if output_fname:
        plt.savefig(output_fname)

```

```python
plot_physics_feature(r'$m$', mass, labels)
```

```python
plot_physics_feature(r'$\Delta R$', delta_R, labels)
```

```python
plot_physics_feature(r'$\tau_{2,1}$', tau_21, labels)
```

# Jet Image Classification

We can now try to use various techniques to classify the jet images into signal (i.e. originating from boosted W bosons) and background (QCD).

We will start with a classic feature-based classifier, which will use properties of the jet such as mass, tau_21, and delta_R (known to have good discriminative power) to separate the two classes.

Then, we will construct different networks that operate directly at the pixel level and compare them all.

## Simple feature-based classifier

#### Data processing

Follow the procedure from yesterday to create your matrix of features X. Shuffle its entries, split them into train, test, and validation set, and scale them to zero mean and unit standard deviation.

```python
features = np.stack((mass, tau_21, delta_R)).T # What we called X yesterday
```

```python
features
```

    array([[  8.32666473e+01,   2.04822451e-01,   5.39974034e-01],
           [  8.12588806e+01,   1.50861815e-01,   7.13163793e-01],
           [  9.46128769e+01,   4.19182539e-01,   4.54179347e-01],
           ...,
           [  6.23480110e+01,   6.62943900e-01,   5.78597903e-01],
           [  8.06691437e+01,   8.02331641e-02,   6.24256849e-01],
           [  6.91418457e+01,   4.10301745e-01,   7.06015885e-01]], dtype=float32)

```python
from sklearn.model_selection import train_test_split
```

```python
# 80% train+validate, 20% test
images_train, images_test, \
    labels_train, labels_test, \
    features_train, features_test \
    = train_test_split(images, labels, features,
                       test_size=0.2)

# 64% train, 16% validate
images_train, images_val, \
    labels_train, labels_val, \
    features_train, features_val \
    = train_test_split(images_train, labels_train, features_train,
                       test_size=0.2)
```

```python
print ('{} training samples\n{} validation samples\n{} testing samples'
       .format(images_train.shape[0], images_val.shape[0], images_test.shape[0]))
```

    512000 training samples
    128000 validation samples
    160000 testing samples

```python
from sklearn.preprocessing import StandardScaler
```

```python
scaler = StandardScaler()
features_train = scaler.fit_transform(features_train)
features_val = scaler.transform(features_val)
features_test = scaler.transform(features_test)
```

#### Model

Build a simple keras model made of fully-connected (Dense) layers. Remember the steps:
1. Define the symbolic graph by connecting layers
1. Define an optimizer and a loss function to minimize
1. Train ('fit') the model to the training dataset, monitoring whether the validation loss continues to decrease
1. Stop the training automatically when the validation loss stops going down
1. Evaluate performance on test set

Recall activation functions: Rectified Linear Unit (relu) vs. Leaky Rectified Linear Unit

```python
from keras.layers import Input, Dense, Dropout
from keras.models import Model
from keras.callbacks import EarlyStopping, ModelCheckpoint
```

    Using TensorFlow backend.

```python
x = Input(shape=(features_train.shape[1], ))
h = Dense(64, activation='relu')(x)
h = Dense(64, activation='relu')(h)
h = Dense(64, activation='relu')(h)
y = Dense(1, activation='sigmoid')(h)
```

```python
baseline_model = Model(x, y)
baseline_model.compile('adam', 'binary_crossentropy', metrics=['acc'])
```

The command below trains the model. However, in the interest of time, I will not train the network on the spot. I will instead load in pre-trained weights from a training I performed last night.

```python
# baseline_model.fit(
#     features_train, labels_train, # X and y
#     epochs=200,
#     batch_size=128,
#     validation_data=(features_val, labels_val), # validation X and y
#     callbacks=[
#         EarlyStopping(verbose=True, patience=15, monitor='val_loss'),
#         ModelCheckpoint('./models/baseline-model.h5', monitor='val_loss',
#                         verbose=True, save_best_only=True)
#     ]
# )
```

    Train on 512000 samples, validate on 128000 samples
    Epoch 1/200
    511616/512000 [============================>.] - ETA: 0s - loss: 0.4542 - acc: 0.7907Epoch 00001: val_loss improved from inf to 0.45350, saving model to ./models/baseline-model.h5
    512000/512000 [==============================] - 20s 39us/step - loss: 0.4542 - acc: 0.7907 - val_loss: 0.4535 - val_acc: 0.7908
    Epoch 2/200
    511360/512000 [============================>.] - ETA: 0s - loss: 0.4515 - acc: 0.7923Epoch 00002: val_loss improved from 0.45350 to 0.45273, saving model to ./models/baseline-model.h5
    512000/512000 [==============================] - 20s 39us/step - loss: 0.4515 - acc: 0.7923 - val_loss: 0.4527 - val_acc: 0.7908
    Epoch 3/200
    510720/512000 [============================>.] - ETA: 0s - loss: 0.4510 - acc: 0.7925Epoch 00003: val_loss improved from 0.45273 to 0.45222, saving model to ./models/baseline-model.h5
    512000/512000 [==============================] - 19s 38us/step - loss: 0.4510 - acc: 0.7925 - val_loss: 0.4522 - val_acc: 0.7916
    Epoch 4/200
    511104/512000 [============================>.] - ETA: 0s - loss: 0.4508 - acc: 0.7929Epoch 00004: val_loss improved from 0.45222 to 0.45193, saving model to ./models/baseline-model.h5
    512000/512000 [==============================] - 20s 38us/step - loss: 0.4508 - acc: 0.7929 - val_loss: 0.4519 - val_acc: 0.7913
    Epoch 5/200
    511744/512000 [============================>.] - ETA: 0s - loss: 0.4507 - acc: 0.7928Epoch 00005: val_loss improved from 0.45193 to 0.45179, saving model to ./models/baseline-model.h5
    512000/512000 [==============================] - 18s 35us/step - loss: 0.4506 - acc: 0.7928 - val_loss: 0.4518 - val_acc: 0.7916
    Epoch 6/200
    511488/512000 [============================>.] - ETA: 0s - loss: 0.4504 - acc: 0.7928Epoch 00006: val_loss did not improve
    512000/512000 [==============================] - 20s 38us/step - loss: 0.4504 - acc: 0.7928 - val_loss: 0.4527 - val_acc: 0.7908
    Epoch 7/200
    511616/512000 [============================>.] - ETA: 0s - loss: 0.4504 - acc: 0.7927Epoch 00007: val_loss did not improve
    512000/512000 [==============================] - 19s 37us/step - loss: 0.4504 - acc: 0.7927 - val_loss: 0.4523 - val_acc: 0.7918
    Epoch 8/200
    510592/512000 [============================>.] - ETA: 0s - loss: 0.4503 - acc: 0.7929Epoch 00008: val_loss did not improve
    512000/512000 [==============================] - 19s 37us/step - loss: 0.4503 - acc: 0.7930 - val_loss: 0.4526 - val_acc: 0.7910
    Epoch 9/200
    511616/512000 [============================>.] - ETA: 0s - loss: 0.4502 - acc: 0.7929Epoch 00009: val_loss did not improve
    512000/512000 [==============================] - 18s 36us/step - loss: 0.4502 - acc: 0.7928 - val_loss: 0.4522 - val_acc: 0.7913
    Epoch 10/200
    510592/512000 [============================>.] - ETA: 0s - loss: 0.4502 - acc: 0.7929Epoch 00010: val_loss improved from 0.45179 to 0.45176, saving model to ./models/baseline-model.h5
    512000/512000 [==============================] - 19s 38us/step - loss: 0.4502 - acc: 0.7929 - val_loss: 0.4518 - val_acc: 0.7914
    Epoch 11/200
    511104/512000 [============================>.] - ETA: 0s - loss: 0.4502 - acc: 0.7927Epoch 00011: val_loss did not improve
    512000/512000 [==============================] - 19s 37us/step - loss: 0.4501 - acc: 0.7927 - val_loss: 0.4532 - val_acc: 0.7906
    Epoch 12/200
    510720/512000 [============================>.] - ETA: 0s - loss: 0.4501 - acc: 0.7929Epoch 00012: val_loss did not improve
    512000/512000 [==============================] - 19s 37us/step - loss: 0.4501 - acc: 0.7929 - val_loss: 0.4523 - val_acc: 0.7914
    Epoch 13/200
    511616/512000 [============================>.] - ETA: 0s - loss: 0.4501 - acc: 0.7928Epoch 00013: val_loss did not improve
    512000/512000 [==============================] - 19s 38us/step - loss: 0.4501 - acc: 0.7928 - val_loss: 0.4524 - val_acc: 0.7916
    Epoch 14/200
    511232/512000 [============================>.] - ETA: 0s - loss: 0.4499 - acc: 0.7928Epoch 00014: val_loss did not improve
    512000/512000 [==============================] - 19s 37us/step - loss: 0.4499 - acc: 0.7929 - val_loss: 0.4522 - val_acc: 0.7913
    Epoch 15/200
    511488/512000 [============================>.] - ETA: 0s - loss: 0.4500 - acc: 0.7929Epoch 00015: val_loss did not improve
    512000/512000 [==============================] - 19s 37us/step - loss: 0.4500 - acc: 0.7929 - val_loss: 0.4519 - val_acc: 0.7914
    Epoch 16/200
    511360/512000 [============================>.] - ETA: 0s - loss: 0.4498 - acc: 0.7929Epoch 00016: val_loss did not improve
    512000/512000 [==============================] - 19s 37us/step - loss: 0.4499 - acc: 0.7929 - val_loss: 0.4519 - val_acc: 0.7914
    Epoch 17/200
    511232/512000 [============================>.] - ETA: 0s - loss: 0.4498 - acc: 0.7932Epoch 00017: val_loss did not improve
    512000/512000 [==============================] - 19s 37us/step - loss: 0.4498 - acc: 0.7932 - val_loss: 0.4529 - val_acc: 0.7903
    Epoch 18/200
    511616/512000 [============================>.] - ETA: 0s - loss: 0.4498 - acc: 0.7928Epoch 00018: val_loss improved from 0.45176 to 0.45162, saving model to ./models/baseline-model.h5
    512000/512000 [==============================] - 19s 36us/step - loss: 0.4498 - acc: 0.7928 - val_loss: 0.4516 - val_acc: 0.7916
    Epoch 19/200
    510848/512000 [============================>.] - ETA: 0s - loss: 0.4497 - acc: 0.7930Epoch 00019: val_loss did not improve
    512000/512000 [==============================] - 19s 37us/step - loss: 0.4498 - acc: 0.7930 - val_loss: 0.4523 - val_acc: 0.7913
    Epoch 20/200
    511744/512000 [============================>.] - ETA: 0s - loss: 0.4497 - acc: 0.7932Epoch 00020: val_loss did not improve
    512000/512000 [==============================] - 19s 36us/step - loss: 0.4497 - acc: 0.7932 - val_loss: 0.4522 - val_acc: 0.7912
    Epoch 21/200
    511232/512000 [============================>.] - ETA: 0s - loss: 0.4498 - acc: 0.7928Epoch 00021: val_loss did not improve
    512000/512000 [==============================] - 19s 37us/step - loss: 0.4498 - acc: 0.7929 - val_loss: 0.4523 - val_acc: 0.7918
    Epoch 22/200
    511104/512000 [============================>.] - ETA: 0s - loss: 0.4497 - acc: 0.7932Epoch 00022: val_loss did not improve
    512000/512000 [==============================] - 19s 37us/step - loss: 0.4497 - acc: 0.7931 - val_loss: 0.4522 - val_acc: 0.7920
    Epoch 23/200
    511104/512000 [============================>.] - ETA: 0s - loss: 0.4498 - acc: 0.7932Epoch 00023: val_loss did not improve
    512000/512000 [==============================] - 19s 37us/step - loss: 0.4498 - acc: 0.7932 - val_loss: 0.4517 - val_acc: 0.7917
    Epoch 24/200
    510720/512000 [============================>.] - ETA: 0s - loss: 0.4497 - acc: 0.7934Epoch 00024: val_loss did not improve
    512000/512000 [==============================] - 19s 37us/step - loss: 0.4497 - acc: 0.7933 - val_loss: 0.4531 - val_acc: 0.7914
    Epoch 25/200
    511232/512000 [============================>.] - ETA: 0s - loss: 0.4496 - acc: 0.7932Epoch 00025: val_loss did not improve
    512000/512000 [==============================] - 19s 38us/step - loss: 0.4497 - acc: 0.7931 - val_loss: 0.4520 - val_acc: 0.7917
    Epoch 26/200
    511744/512000 [============================>.] - ETA: 0s - loss: 0.4496 - acc: 0.7932Epoch 00026: val_loss did not improve
    512000/512000 [==============================] - 19s 37us/step - loss: 0.4496 - acc: 0.7932 - val_loss: 0.4525 - val_acc: 0.7914
    Epoch 27/200
    511360/512000 [============================>.] - ETA: 0s - loss: 0.4496 - acc: 0.7932Epoch 00027: val_loss did not improve
    512000/512000 [==============================] - 19s 37us/step - loss: 0.4496 - acc: 0.7932 - val_loss: 0.4518 - val_acc: 0.7922
    Epoch 28/200
    511360/512000 [============================>.] - ETA: 0s - loss: 0.4495 - acc: 0.7931Epoch 00028: val_loss did not improve
    512000/512000 [==============================] - 19s 37us/step - loss: 0.4496 - acc: 0.7931 - val_loss: 0.4522 - val_acc: 0.7915
    Epoch 29/200
    511232/512000 [============================>.] - ETA: 0s - loss: 0.4496 - acc: 0.7932Epoch 00029: val_loss did not improve
    512000/512000 [==============================] - 19s 36us/step - loss: 0.4496 - acc: 0.7932 - val_loss: 0.4519 - val_acc: 0.7913
    Epoch 30/200
    510720/512000 [============================>.] - ETA: 0s - loss: 0.4496 - acc: 0.7933Epoch 00030: val_loss did not improve
    512000/512000 [==============================] - 19s 38us/step - loss: 0.4495 - acc: 0.7933 - val_loss: 0.4522 - val_acc: 0.7913
    Epoch 31/200
    511360/512000 [============================>.] - ETA: 0s - loss: 0.4496 - acc: 0.7930Epoch 00031: val_loss did not improve
    512000/512000 [==============================] - 19s 38us/step - loss: 0.4496 - acc: 0.7930 - val_loss: 0.4518 - val_acc: 0.7916
    Epoch 32/200
    510080/512000 [============================>.] - ETA: 0s - loss: 0.4495 - acc: 0.7932Epoch 00032: val_loss did not improve
    512000/512000 [==============================] - 19s 37us/step - loss: 0.4495 - acc: 0.7932 - val_loss: 0.4521 - val_acc: 0.7919
    Epoch 33/200
    510336/512000 [============================>.] - ETA: 0s - loss: 0.4495 - acc: 0.7934Epoch 00033: val_loss did not improve
    512000/512000 [==============================] - 18s 36us/step - loss: 0.4495 - acc: 0.7933 - val_loss: 0.4523 - val_acc: 0.7914
    Epoch 00033: early stopping

    <keras.callbacks.History at 0x7f68c898ab50>

If you were to actually run the training, you would be able to visualize its history. Keras saves the entire training history, keeping track of whatever metric you specify (here accuracy and loss).

```python
history = baseline_model.history.history
history.keys()
```

    ['acc', 'loss', 'val_acc', 'val_loss']

```python
# accuracy plot
plt.plot(100 * np.array(history['acc']), label='training')
plt.plot(100 * np.array(history['val_acc']), label='validation')
plt.xlim(0)
plt.xlabel('epoch')
plt.ylabel('accuracy %')
plt.legend(loc='lower right', fontsize=20)
plt.show()
```

```python
# loss plot
plt.plot(100 * np.array(history['loss']), label='training')
plt.plot(100 * np.array(history['val_loss']), label='validation')
plt.xlim(0)
plt.xlabel('epoch')
plt.ylabel('loss')
plt.legend(loc='upper right', fontsize=20)
# the line indicate the epoch corresponding to the best performance on the validation set
# plt.vlines(np.argmin(history['val_loss']), 45, 56, linestyle='dashed', linewidth=0.5)
plt.show()
```

```python
print 'Loss estimate on unseen examples (from validation set) = {0:.3f}'.format(np.min(history['val_loss']))
```

    Loss estimate on unseen examples (from validation set) = 0.452

#### Evaluate on test set

```python
baseline_model.load_weights('./models/baseline-model.h5')
```

```python
yhat_baseline = baseline_model.predict(features_test, batch_size=512)
```

```python
bins = np.linspace(0, 1, 20)
_ = plt.hist(yhat_baseline[labels_test==1],
         histtype='stepfilled', alpha=0.5, color='red', label=r"Signal ($W' \rightarrow WZ$)", bins=bins)
_ = plt.hist(yhat_baseline[labels_test==0],
         histtype='stepfilled', alpha=0.5, color='blue', label=r'Background (QCD dijets)', bins=bins)
plt.legend(loc='upper center')
plt.xlabel('P(signal) assigned by the baseline model')
```

## Convolutional Neural Network

We can now instead try to learn a model directly on the pixel space, instead of summarizing the information into engineered features such as mass, tau_21 and delta_R.

```python
from keras.layers import Conv2D, Flatten, LeakyReLU
```

```python
# add channel dimension (1 for grayscale)
images_train = np.expand_dims(images_train, -1)
images_test = np.expand_dims(images_test, -1)
images_val = np.expand_dims(images_val, -1)
```

```python
x = Input(shape=(images_train.shape[1:]))
h = Conv2D(32, kernel_size=7, strides=1)(x)
h = LeakyReLU()(h)
h = Dropout(0.2)(h)

h = Conv2D(64, kernel_size=7, strides=1)(h)
h = LeakyReLU()(h)
h = Dropout(0.2)(h)

h = Conv2D(128, kernel_size=5, strides=1)(h)
h = LeakyReLU()(h)
h = Dropout(0.2)(h)

h = Conv2D(256, kernel_size=5, strides=1)(h)
h = LeakyReLU()(h)
h = Flatten()(h)
h = Dropout(0.2)(h)
y = Dense(1, activation='sigmoid')(h)
```

```python
cnn_model = Model(x, y)
cnn_model.compile('adam', 'binary_crossentropy', metrics=['acc'])
```

```python
cnn_model.summary()
```

    _________________________________________________________________
    Layer (type)                 Output Shape              Param #
    =================================================================
    input_3 (InputLayer)         (None, 25, 25, 1)         0
    _________________________________________________________________
    conv2d_2 (Conv2D)            (None, 19, 19, 32)        1600
    _________________________________________________________________
    leaky_re_lu_1 (LeakyReLU)    (None, 19, 19, 32)        0
    _________________________________________________________________
    dropout_1 (Dropout)          (None, 19, 19, 32)        0
    _________________________________________________________________
    conv2d_3 (Conv2D)            (None, 13, 13, 64)        100416
    _________________________________________________________________
    leaky_re_lu_2 (LeakyReLU)    (None, 13, 13, 64)        0
    _________________________________________________________________
    dropout_2 (Dropout)          (None, 13, 13, 64)        0
    _________________________________________________________________
    conv2d_4 (Conv2D)            (None, 9, 9, 128)         204928
    _________________________________________________________________
    leaky_re_lu_3 (LeakyReLU)    (None, 9, 9, 128)         0
    _________________________________________________________________
    dropout_3 (Dropout)          (None, 9, 9, 128)         0
    _________________________________________________________________
    conv2d_5 (Conv2D)            (None, 5, 5, 256)         819456
    _________________________________________________________________
    leaky_re_lu_4 (LeakyReLU)    (None, 5, 5, 256)         0
    _________________________________________________________________
    flatten_1 (Flatten)          (None, 6400)              0
    _________________________________________________________________
    dropout_4 (Dropout)          (None, 6400)              0
    _________________________________________________________________
    dense_5 (Dense)              (None, 1)                 6401
    =================================================================
    Total params: 1,132,801
    Trainable params: 1,132,801
    Non-trainable params: 0
    _________________________________________________________________

Feel free to try to train it at home! For this tutorial, we will just load in pre-trained weights.

```python
# cnn_model.fit(
#     images_train, labels_train,
#     epochs=100,
#     batch_size=512,
#     validation_data=(images_val, labels_val),
#     callbacks=[
#         EarlyStopping(verbose=True, patience=30, monitor='val_loss'),
#         ModelCheckpoint('./models/cnn-model.h5', monitor='val_loss',
#                         verbose=True, save_best_only=True)
#     ]
# )
```

```python
cnn_history = cnn_model.history.history
```

```python
# accuracy plot
plt.plot(100 * np.array(cnn_history['acc']), label='training')
plt.plot(100 * np.array(cnn_history['val_acc']), label='validation')
plt.xlim(0)
plt.xlabel('epoch')
plt.ylabel('accuracy %')
plt.legend(loc='lower right', fontsize=20)
plt.show()
```

```python
# loss plot
plt.plot(100 * np.array(cnn_history['loss']), label='training')
plt.plot(100 * np.array(cnn_history['val_loss']), label='validation')
plt.xlim(0)
plt.xlabel('epoch')
plt.ylabel('loss')
plt.legend(loc='upper right', fontsize=20)
# the line indicate the epoch corresponding to the best performance on the validation set
plt.vlines(np.argmin(cnn_history['val_loss']), 43, 56, linestyle='dashed', linewidth=0.5)
plt.show()
```

```python
print 'Loss estimate on unseen examples (from validation set) = {0:.3f}'.format(np.min(cnn_history['val_loss']))
```

    Loss estimate on unseen examples (from validation set) = 0.398

```python
cnn_model.load_weights('models/cnn-model.h5')
yhat_cnn = cnn_model.predict(images_test, batch_size=512, verbose=True)
```

    160000/160000 [==============================] - 585s 4ms/step

## Locally-Connected Neural Network

```python
from keras.layers import LocallyConnected2D , MaxPool2D, Flatten
```

```python
x = Input(shape=(images_train.shape[1:]))
h = LocallyConnected2D(32, kernel_size=9, strides=2)(x)
h = LeakyReLU()(h)

h = Dropout(0.2)(h)
h = LocallyConnected2D(32, kernel_size=5, strides=1)(h)
h = LeakyReLU()(h)

h = Dropout(0.2)(h)
h = LocallyConnected2D(64, kernel_size=3, strides=1)(h)
h = LeakyReLU()(h)

h = Dropout(0.2)(h)
h = LocallyConnected2D(64, kernel_size=3, strides=1)(h)
h = LeakyReLU()(h)

h = Flatten()(h)
h = Dropout(0.2)(h)
y = Dense(1, activation='sigmoid')(h)
```

```python
lcn_model = Model(x, y)
lcn_model.compile('adam', 'binary_crossentropy', metrics=['acc'])
```

```python
lcn_model.summary()
```

    _________________________________________________________________
    Layer (type)                 Output Shape              Param #
    =================================================================
    input_4 (InputLayer)         (None, 25, 25, 1)         0
    _________________________________________________________________
    locally_connected2d_1 (Local (None, 9, 9, 32)          212544
    _________________________________________________________________
    leaky_re_lu_5 (LeakyReLU)    (None, 9, 9, 32)          0
    _________________________________________________________________
    dropout_5 (Dropout)          (None, 9, 9, 32)          0
    _________________________________________________________________
    locally_connected2d_2 (Local (None, 5, 5, 32)          640800
    _________________________________________________________________
    leaky_re_lu_6 (LeakyReLU)    (None, 5, 5, 32)          0
    _________________________________________________________________
    dropout_6 (Dropout)          (None, 5, 5, 32)          0
    _________________________________________________________________
    locally_connected2d_3 (Local (None, 3, 3, 64)          166464
    _________________________________________________________________
    leaky_re_lu_7 (LeakyReLU)    (None, 3, 3, 64)          0
    _________________________________________________________________
    dropout_7 (Dropout)          (None, 3, 3, 64)          0
    _________________________________________________________________
    locally_connected2d_4 (Local (None, 1, 1, 64)          36928
    _________________________________________________________________
    leaky_re_lu_8 (LeakyReLU)    (None, 1, 1, 64)          0
    _________________________________________________________________
    flatten_2 (Flatten)          (None, 64)                0
    _________________________________________________________________
    dropout_8 (Dropout)          (None, 64)                0
    _________________________________________________________________
    dense_6 (Dense)              (None, 1)                 65
    =================================================================
    Total params: 1,056,801
    Trainable params: 1,056,801
    Non-trainable params: 0
    _________________________________________________________________

Feel free to try to train it at home! For this tutorial, we will just load in pre-trained weights.

```python
# lcn_model.fit(
#     images_train, labels_train,
#     epochs=100,
#     batch_size=256,
#     validation_data=(images_val, labels_val),
#     callbacks=[
#         EarlyStopping(verbose=True, patience=30, monitor='val_loss'),
#         ModelCheckpoint('./models/lcn-model.h5', monitor='val_loss',
#                         verbose=True, save_best_only=True)
#     ]
# )
```

```python
lcn_history = lcn_model.history.history
```

```python
# accuracy plot
plt.plot(100 * np.array(lcn_history['acc']), label='training')
plt.plot(100 * np.array(lcn_history['val_acc']), label='validation')
plt.xlim(0)
plt.xlabel('epoch')
plt.ylabel('accuracy %')
plt.legend(loc='lower right', fontsize=20)
plt.show()
```

```python
# loss plot
plt.plot(100 * np.array(lcn_history['loss']), label='training')
plt.plot(100 * np.array(lcn_history['val_loss']), label='validation')
plt.xlim(0)
plt.xlabel('epoch')
plt.ylabel('loss')
plt.legend(loc='upper right', fontsize=20)
# the line indicate the epoch corresponding to the best performance on the validation set
plt.vlines(np.argmin(lcn_history['val_loss']), 43, 56, linestyle='dashed', linewidth=0.5)
plt.show()
```

```python
print 'Loss estimate on unseen examples (from validation set) = {0:.3f}'.format(np.min(lcn_history['val_loss']))
```

    Loss estimate on unseen examples (from validation set) = 0.390

```python
lcn_model.load_weights('models/lcn-model.h5')
yhat_lcn = lcn_model.predict(images_test, batch_size=512)
```

## Fully-Connected network

```python
x = Input(shape=(images_train.shape[1:]))
h = Flatten()(x)
h = Dense(25 ** 2, kernel_initializer='he_normal')(h)
h = LeakyReLU()(h)

h = Dropout(0.2)(h)
h = Dense(512, kernel_initializer='he_normal')(h)
h = LeakyReLU()(h)

h = Dropout(0.2)(h)
h = Dense(256, kernel_initializer='he_normal')(h)
h = LeakyReLU()(h)

h = Dropout(0.2)(h)
h = Dense(128, kernel_initializer='he_normal')(h)
h = LeakyReLU()(h)

h = Dropout(0.2)(h)
y = Dense(1, activation='sigmoid')(h)
```

```python
dense_model = Model(x, y)
dense_model.compile('adam', 'binary_crossentropy', metrics=['acc'])
```

```python
# dense_model.fit(
#     images_train, labels_train,
#     epochs=100,
#     batch_size=256,
#     validation_data=(images_val, labels_val),
#     callbacks=[
#         EarlyStopping(verbose=True, patience=30, monitor='val_loss'),
#         ModelCheckpoint('./models/dense-model.h5', monitor='val_loss',
#                         verbose=True, save_best_only=True)
#     ]
# )
```

```python
dense_history = dense_model.history.history
```

```python
# accuracy plot
plt.plot(100 * np.array(dense_history['acc']), label='training')
plt.plot(100 * np.array(dense_history['val_acc']), label='validation')
plt.xlim(0)
plt.xlabel('epoch')
plt.ylabel('accuracy %')
plt.legend(loc='lower right', fontsize=20)
plt.show()
```

```python
# loss plot
plt.plot(100 * np.array(dense_history['loss']), label='training')
plt.plot(100 * np.array(dense_history['val_loss']), label='validation')
plt.xlim(0)
plt.xlabel('epoch')
plt.ylabel('loss')
plt.legend(loc='upper right', fontsize=20)
# the line indicate the epoch corresponding to the best performance on the validation set
plt.vlines(np.argmin(dense_history['val_loss']), 43, 56, linestyle='dashed', linewidth=0.5)
plt.show()
```

```python
print 'Loss estimate on unseen examples (from validation set) = {0:.3f}'.format(np.min(dense_history['val_loss']))
```

    Loss estimate on unseen examples (from validation set) = 0.396

```python
dense_model.load_weights('models/dense-model.h5')
yhat_dense = dense_model.predict(images_test, batch_size=512, verbose=True)
```

    160000/160000 [==============================] - 11s 69us/step

## Plot ROC Curves

A standard way to visualize the tradeoff between low false positive rate (FPR) and high true positive rate (TPR) is by plotting them on a ROC curve. In Physics, people like to plot the TPR on the x-axis, and 1/FPR on the y-axis.

```python
from sklearn.metrics import roc_curve
```

```python
fpr_cnn, tpr_cnn, _ = roc_curve(labels_test, yhat_cnn)
fpr_lcn, tpr_lcn, _ = roc_curve(labels_test, yhat_lcn)
fpr_dense, tpr_dense, _ = roc_curve(labels_test, yhat_dense)
fpr_baseline, tpr_baseline, _ = roc_curve(labels_test, yhat_baseline)
```

```python
plt.figure(figsize=(10,10))
plt.grid(b = True, which = 'minor')
plt.grid(b = True, which = 'major')
_ = plt.plot(tpr_cnn, 1./fpr_cnn, label='CNN')
_ = plt.plot(tpr_lcn, 1./fpr_lcn, label='LCN')
_ = plt.plot(tpr_dense, 1./fpr_dense, label='FCN')
_ = plt.plot(tpr_baseline, 1./fpr_baseline, label='Baseline')
plt.legend()
plt.xlim((0.1, 0.9))
plt.ylim((1, 1000))
plt.yscale('log')
```

## What is the network learning?

```python
def get_correlations(images, disc_output):
    '''
    calculate linear correlation between each pixel and the output of the classifier
    to see what pixels are more indicative of a specific class.
    '''
    import pandas as pd
    # -- find the total number of pixels per image, here 25 x 25
    n_pixels = np.prod(images.shape[1:3])

    # -- add the pixels as columns to a dataframe
    df = pd.DataFrame(
        {i : np.squeeze(images).reshape(-1, n_pixels)[:, i] for i in range(n_pixels)}
    )
    # -- add a column to the end of the dataframe for the discriminator's output
    df['disc_output'] = disc_output
    # -- pandas offers an easy solution to calculate correlations
    # (even though it's slow because it also calculates the correlation between each pixel and every other pixel)
    correlations = df.corr().values[:-1, -1]
    return correlations

def plot_correlations(correlations, extent, title='', img_dim=(25, 25), cmap=plt.cm.seismic):
    '''
    call the function about and then plot the correlations in image format
    '''
    max_mag = max(
        abs(np.min(correlations[np.isfinite(correlations)])),
        abs(np.max(correlations[np.isfinite(correlations)])),
    ) # highest correlation value (abs value), to make the plot look nice and on a reasonable scale

    f, ax = plt.subplots(figsize=(6, 6))
    im = ax.imshow(
        correlations.reshape(img_dim),
        interpolation='nearest',
        norm=Normalize(vmin=-max_mag, vmax=max_mag),
        extent=extent,
        cmap=cmap
    )

    plt.colorbar(im, fraction=0.05, pad=0.05)
    plt.xlabel(r'[Transformed] Pseudorapidity $(\eta)$')
    plt.ylabel(r'[Transformed] Azimuthal Angle $(\phi)$')
    plt.title(title)
#     plt.savefig(os.path.join('..', outdir, outname))
```

You can now visualize what each network picks up on, at least to first order. These correlation plots tells us whether a specific pixel being strongly activated is a good indicator of that jet image belonging to a class or the other. Red represents the signal (boosted W from W'-->WZ), blue represents the background (QCD).

```python
plot_correlations(
    get_correlations(images_test[:10000], yhat_cnn[:10000]),
    extent=[-1.25, 1.25, -1.25, 1.25],
    title='Correlation between pixels \n and the CNN prediction'
)
```

```python
plot_correlations(
    get_correlations(images_test[:10000], yhat_lcn[:10000]),
    extent=[-1.25, 1.25, -1.25, 1.25],
    title='Correlation between pixels \n and the LCN prediction'
)
```

```python
plot_correlations(
    get_correlations(images_test[:10000], yhat_dense[:10000]),
    extent=[-1.25, 1.25, -1.25, 1.25],
    title='Correlation between pixels \n and the FCN prediction'
)
```

```python
plot_correlations(
    get_correlations(images_test[:10000], yhat_baseline[:10000]),
    extent=[-1.25, 1.25, -1.25, 1.25],
    title='Correlation between pixels \n and the Baseline prediction'
)
```

You can also look at how the output of each classifier is correlated with known quantities that are known to be discriminative, such as the ones used in the baseline classifier above (mass, tau_21, delta_R). This will inform us as to whether the network has 'learned' to internally calculate a representation that is close to these variables, thus eliminating our need to come up with these features ourselves.

```python
def plot_output_vs_kin(kin, output, xlabel, ylabel, nbins=30):
    '''
    Plot one output of the discriminator network vs. one of the 1D physics variables that describe jets
    Args:
    -----
        kin : numpy array, kinematic property (such as mass or pT) associated with each image. I.e.: discrete_mass(np.squeeze(generated_images))
        output : numpy array, one of the 2 outputs of the discriminator, evaluated on the same images that `kin` refers to
        xlabel : string, x-axis label that describes the meaning of `kin`
        ylabel : string, y-axis label that describes the meaning og `output`
        outname : name of the output file, to be placed in ../plots/
        nbins : (default = 30) number of bins to use to represent the distributions in a discretized way

    '''
    # set the colormap
    plt.set_cmap('jet')
    # draw a 2d histogram of the discriminator's output versus the kinematic variable of choice (mass, pT, etc.)
    h, binx, biny, _ = plt.hist2d(kin, output.reshape(-1,), bins=nbins)
    plt.clf() # we don't want to plot this 2D histogram, we want to normalize it per bin first

    # normalize the histogram such that the entries in each column add up to 1, such that the intensity
    # of each corresponds to the percentage of the jets in a given mass (or pT) bin that get assigned a p
    for i in range(nbins):
        h[i, :] = h[i, :] / float(np.sum(h[i, :]))
    # plot the normalized histogram as an image
    f, ax2 = plt.subplots(figsize=(6, 6))
    im = ax2.imshow(
        np.flipud(h.T),
        interpolation='nearest',
        norm=LogNorm(),
        extent=[binx.min(), binx.max(), biny.min(), biny.max()],
        aspect="auto"
    )
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    # add a custom colorbar
    cax = f.add_axes([0.93, 0.1, 0.03, 0.8])
    plt.colorbar(im, cax = cax)
    plt.set_cmap('viridis')
```

### Mass

```python
plot_output_vs_kin(
    scaler.inverse_transform(features_test)[:, 0], # mass
    yhat_cnn,
    xlabel='Discrete jet image mass (GeV)',
    ylabel='P(signal)',
)
```

```python
plot_output_vs_kin(
    scaler.inverse_transform(features_test)[:, 0], # mass
    yhat_lcn,
    xlabel='Discrete jet image mass (GeV)',
    ylabel='P(signal)',
)
```

```python
plot_output_vs_kin(
    scaler.inverse_transform(features_test)[:, 0], # mass
    yhat_dense,
    xlabel='Discrete jet image mass (GeV)',
    ylabel='P(signal)',
)
```

```python
plot_output_vs_kin(
    scaler.inverse_transform(features_test)[:, 0], # mass
    yhat_baseline,
    xlabel='Discrete jet image mass (GeV)',
    ylabel='P(signal)',
)
```

### Tau_21

```python
plot_output_vs_kin(
    scaler.inverse_transform(features_test)[:, 1], # tau21
    yhat_cnn,
    xlabel=r'Discrete jet $\tau_{21}$',
    ylabel='P(signal)',
)
```

```python
plot_output_vs_kin(
    scaler.inverse_transform(features_test)[:, 1], # tau21
    yhat_lcn,
    xlabel=r'Discrete jet $\tau_{21}$',
    ylabel='P(signal)',
)
```

```python
plot_output_vs_kin(
    scaler.inverse_transform(features_test)[:, 1], # tau21
    yhat_dense,
    xlabel=r'Discrete jet $\tau_{21}$',
    ylabel='P(signal)',
)
```

```python
plot_output_vs_kin(
    scaler.inverse_transform(features_test)[:, 1], # tau21
    yhat_baseline,
    xlabel=r'Discrete jet $\tau_{21}$',
    ylabel='P(signal)',
)
```

### Delta R

```python
plot_output_vs_kin(
    scaler.inverse_transform(features_test)[:, 2], # deltaR
    yhat_cnn,
    xlabel=r'Discrete jet $\Delta R$',
    ylabel='P(signal)',
)
```

```python
plot_output_vs_kin(
    scaler.inverse_transform(features_test)[:, 2], # deltaR
    yhat_lcn,
    xlabel=r'Discrete jet $\Delta R$',
    ylabel='P(signal)',
)
```

```python
plot_output_vs_kin(
    scaler.inverse_transform(features_test)[:, 2], # deltaR
    yhat_dense,
    xlabel=r'Discrete jet $\Delta R$',
    ylabel='P(signal)',
)
```

```python
plot_output_vs_kin(
    scaler.inverse_transform(features_test)[:, 2], # deltaR
    yhat_baseline,
    xlabel=r'Discrete jet $\Delta R$',
    ylabel='P(signal)',
)
```

Sometimes, however, it is not a good idea to have your network learn the mass of a jet and use that to classify jets. In that case, in fact, the network will successfully suppress all jets outside of the signal-like mass window and sculpt the mass of the background to look like a signal peak.

What we would like to be able to do, instead, is to have a classifier that is insensitive to mass, and that reduced the background across the entire mass spectrum.

For reference, see: C. Shimmin et al., Decorrelated Jet Substructure Tagging using Adversarial Neural Networks.

# Training a GAN (WGAN-GP) on jet images

WGAN-GP = a type of GAN that minimizes the Wasserstein distance between the target and generated distributions and enforces the Lipschitz contraint by penalizing the norm of the gradient instead of clipping weights.

```python
from keras.layers import Input, Dense, Reshape, Flatten
from keras.layers.merge import _Merge
from keras.layers.convolutional import Convolution2D, Conv2DTranspose
from keras.layers.normalization import BatchNormalization
from keras.layers.advanced_activations import LeakyReLU
from keras.optimizers import Adam
from keras import backend as K
```

```python
BATCH_SIZE = 100

# The training ratio is the number of discriminator updates per generator
# update. The paper uses 5.
TRAINING_RATIO = 5
GRADIENT_PENALTY_WEIGHT = 10  # As per the paper
```

Define new custom loss functions not included among the standard keras ones:

```python
def wasserstein_loss(y_true, y_pred):
    return K.mean(y_true * y_pred)

def gradient_penalty_loss(y_true, y_pred, averaged_samples,
                          gradient_penalty_weight):
    gradients = K.gradients(K.sum(y_pred), averaged_samples)
    gradient_l2_norm = K.sqrt(K.sum(K.square(gradients)))
    gradient_penalty = gradient_penalty_weight * K.square(1 - gradient_l2_norm)
    return gradient_penalty
```

```python
def make_generator():
    """Creates a generator model that takes a 100-dimensional latent prior and
    converts to size 25 x 25 x 1"""

    z = Input(shape=(100, ))

    x = Dense(1024, input_dim=100)(z)
    x = LeakyReLU()(x)

    x = Dense(128 * 7 * 7)(x)
    x = BatchNormalization()(x)
    x = LeakyReLU()(x)

    x = Reshape((7, 7, 128))(x)

    x = Conv2DTranspose(128, (5, 5), strides=2, padding='same')(x)
    x = BatchNormalization(axis=-1)(x)
    x = LeakyReLU()(x)

    x = Convolution2D(64, (5, 5), padding='same')(x)
    x = BatchNormalization(axis=-1)(x)
    x = LeakyReLU()(x)

    x = Conv2DTranspose(64, (5, 5), strides=2, padding='same')(x)
    x = BatchNormalization(axis=-1)(x)
    x = LeakyReLU()(x)

    y = Convolution2D(1, (4, 4), padding='valid', activation='relu')(x)

    return Model(z, y)
```

```python
def make_discriminator():

    x = Input(shape=(25, 25, 1))

    h = Convolution2D(64, (5, 5), padding='same')(x)
    h = LeakyReLU()(h)

    h = Convolution2D(128, (5, 5), kernel_initializer='he_normal',
                      strides=2)(h)
    h = LeakyReLU()(h)

    h = Convolution2D(256, (5, 5), kernel_initializer='he_normal',
                      padding='same', strides=2)(h)
    h = LeakyReLU()(h)

    h = Flatten()(h)
    h = Dense(1024, kernel_initializer='he_normal')(h)
    h = LeakyReLU()(h)

    y = Dense(1, kernel_initializer='he_normal')(h)

    return Model(x, y)
```

Build individual models:

```python
generator = make_generator()
discriminator = make_discriminator()
```

Build connected model with only generator trainable:

```python
discriminator.trainable = False
z = Input(shape=(100, ))
generator_model = Model(z, discriminator(generator(z)))
# We use the Adam paramaters from Gulrajani et al.
generator_model.compile(optimizer=Adam(0.0001, beta_1=0.5, beta_2=0.9),
                        loss=wasserstein_loss)
```

Build connected model with only discriminator trainable:

```python
discriminator.trainable = True
generator.trainable = False
```

```python
class RandomWeightedAverage(_Merge): # used for gradient norm penalty

    def _merge_function(self, inputs):
        weights = K.random_uniform((K.shape(inputs[0])[0], 1, 1, 1))
        return (weights * inputs[0]) + ((1 - weights) * inputs[1])
```

```python
real_samples = Input(shape=(25, 25, 1))
z = Input(shape=(100,))
fake_samples = generator(z)

critic_out_fake = discriminator(fake_samples)
critic_out_real = discriminator(real_samples)

# generate weighted-averages of real and generated
# samples, to use for the gradient norm penalty.
averaged_samples = RandomWeightedAverage()([real_samples, fake_samples])

# running them thru critic to get the gradient norm for the GP loss.
averaged_samples_out = discriminator(averaged_samples)

# The gradient penalty loss function requires the input averaged samples
def gp_loss(y_true, y_pred):
    return gradient_penalty_loss(
        y_true, y_pred,
        averaged_samples=averaged_samples,
        gradient_penalty_weight=GRADIENT_PENALTY_WEIGHT
    )

discriminator_model = Model(
    inputs=[real_samples, z],
    outputs=[critic_out_real, critic_out_fake, averaged_samples_out]
)

# We use the Adam paramaters from Gulrajani et al.
discriminator_model.compile(
    optimizer=Adam(0.0001, beta_1=0.5, beta_2=0.9),
    loss=[wasserstein_loss, wasserstein_loss, gp_loss]
)
```

```python
# positive_y is the label vector for real samples, with value 1.
# negative_y is the label vector for generated samples, with value -1.
# dummy_y vector is passed to the gradient_penalty loss function and is
# not used.

positive_y = np.ones((BATCH_SIZE, 1), dtype=np.float32)
negative_y = -positive_y
dummy_y = np.zeros((BATCH_SIZE, 1), dtype=np.float32)
```

```python
# do a little bit of scaling for stability
X_train = np.expand_dims(np.squeeze(images_train[:30000]) / 100, -1)
```

```python
overall_disc_loss = []

for epoch in range(200): # train for 200 iterations

    # at each epoch, shuffle the training set to get new samples
    np.random.shuffle(X_train)
    print "Epoch: ", epoch
    print "Number of batches: ", int(X_train.shape[0] // BATCH_SIZE)
    discriminator_loss = []
    generator_loss = []

    # we'll need this many samples per critic update
    critic_nb_samples = BATCH_SIZE * TRAINING_RATIO

    # loop through batches
    for i in range(int(X_train.shape[0] // (BATCH_SIZE * TRAINING_RATIO))):
        X_critic = X_train[i * critic_nb_samples:(i + 1) * critic_nb_samples]

        # critic gets trained 5 times more per iteration than the generator
        for j in range(TRAINING_RATIO):
            X_minibatch = X_critic[j * BATCH_SIZE:(j + 1) * BATCH_SIZE]

            # generate new input noise
            noise = np.random.rand(BATCH_SIZE, 100).astype(np.float32)

            # train the discriminator (or critic)
            disc_loss = discriminator_model.train_on_batch(
                [X_minibatch, noise],
                [positive_y, negative_y, dummy_y]
            )
            discriminator_loss.append(disc_loss)

        critic_score = np.array(discriminator_loss)[:, 0]
        if i % 10 == 0:
            print 'critic score =', critic_score.mean()
        overall_disc_loss.extend(critic_score.tolist())
        # train the generator
        gen_loss = generator_model.train_on_batch(
            np.random.rand(BATCH_SIZE, 100),
            positive_y
        )
        generator_loss.append(gen_loss)
```

```python
# discriminator.save_weights('./models/wgan-discriminator.h5')
# generator.save_weights('./models/wgan-generator.h5')
```

Load pre-trained wights:

```python
discriminator.load_weights('./models/wgan-discriminator.h5')
generator.load_weights('./models/wgan-generator.h5')
```

Now that we have a trained GAN, we can see if it actually works and what it produces. We can now get rid of the critic (discriminator) and focus only on the part that we really care about, the generator. Let's execute a forward pass on the generator.

```python
# input noise that will be transformed into jet images
noise = np.random.rand(1000, 100).astype(np.float32)
```

```python
# produce some jet images from the generator!
fake_jets = generator.predict(noise, batch_size=BATCH_SIZE, verbose=True)

# rescale energies and remove redundant dimension for grayscale channel
fake_jets = np.squeeze(fake_jets * 100)
```

    1000/1000 [==============================] - 10s 10ms/step

Let's look at some GAN-generated jet images!

```python
plot_jet_image(fake_jets.mean(0))
```

We can also look at the difference between the average generated image and the average real image to identify parts of the image that are not well-captured by the GAN and might need improvement. Green pixels are more strongly activated, on average, in fake images, while purple pixels are more strongly activated in background images.

```python
plot_diff_jet_image(fake_jets.mean(0) - images.mean(0), cmap='PRGn')
```

There are multiple ways of imporving this training, such as conditioning on class or attributes, using minibatch discrimination, enforcing physical constraints, etc.

```python

```


# Research
The title of the notebook should be coherent with file name. Namely, file name should be:
author's initials_progressive number_title.ipynb

# Purpose
State the purpose of the notebook.

# Methodology
Quickly describe assumptions and processing steps.

# WIP - improvements
(WORK IN PROGRESS)
Use this section only if the notebook is not final.

Notable TODOs:
* todo 1
* todo 2
* todo 3

## 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
#plt.style.use('paper')

#import data
import copy
from mdldb.tables 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 sklearn.metrics import r2_score
from src.data import database
```