flow_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I-DStRWBSe_3"
   },
   "source": [
    "# Normalizing flows exercise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KHlmT4cuSe_-"
   },
   "source": [
    "## 1.) Motivation\n",
    "\n",
    "Just as Variational Autoencoders (VAEs) or Generative Adversarial Networks (GANs), Normalizing Flows are **Probabilistic Generative Models (PGMs)**, which describe a probability distribution that we attempt to learn from a set of observed data. \n",
    "\n",
    "PGMs are useful for generating new samples from the learned distribution, evaluating the likelihood of new data points etc. \n",
    "\n",
    "Normalizing flows are PGMs built on **invertible transformations**. Their advantages are that it is typically possible to efficiently **sample** and **evaluate** the learned distributions. Normalizing flows also are **highly expressive** and come with a **useful latent space representation**, since we have a one-to-one mapping between our input and the latent space. Finally, they are also **easy to train**, since we just need to conduct a simple maximum likelihood training.\n",
    "\n",
    "In this hands-on exercise we will show examples for using normalizing flows to estimate (conditional) probability densities and to generate new samples. While some instructive mock datasets will be used for getting to know Normalizing Flows, we will also have one dedicated exercise that will use a Particle Physics dataset in the end. As a framework for actually implementing Normalizing Flows quickly, we will use the **TensorFlow probability** package. We hope this tutorial will give you at least a rough idea of how Normalizing Flows work in practice and - most importantly - that you find it instructive and fun :)\n",
    "\n",
    "## 2.) Normalizing Flows - The general idea (recap)\n",
    "\n",
    "At the heart of normalizing flows stands the change of random variables formula. Let's say we had a random variable $U$ and now, we apply a simple transformation, defining the transformed random variable as $X$. Then, the change of random variable formula is given by:\n",
    "\n",
    "$$ p(x) = p(u)\\left| \\frac{df(u)}{du}\\right|^{-1} $$\n",
    "\n",
    "As an example, let's have a look at the figure below:\n",
    "\n",
    "<br />\n",
    "<img src=\"https://drive.google.com/uc?id=1SoaFiiAtpYMSOepjtZkgZACQeV4Zxwbr\" width=\"250\">\n",
    "<br />\n",
    "\n",
    "We start out with a uniform distribution $U$ between 0 and 1, and then do a transformation $f(U)=2\\cdot U + 1$. This, however leads to a distribution that is not normalized, so in order to get a proper probability distribution, we still need to scale it with $ \\left| \\left(\\frac{df(u)}{du}\\right)\\right|^{-1} $.\n",
    "\n",
    "The derivative of $f$ with respect to $u$ is clearly 2, and the inverse is $\\frac{1}{2}$, which is exactly the scaling factor that we need for getting a properly normalized distribution.\n",
    "\n",
    "This, however, is only a simple univariate example. For the multivariate case, instead of the derivative, we need to scale our transformed distribution with the Jacobian determinant:\n",
    "\n",
    "$$ \\left|\\det \\left( \\frac{\\partial \\boldsymbol{f}}{\\partial \\boldsymbol{u}} \\right)  \\right|^{-1} = \\left|\\det \\left( \\frac{\\partial \\boldsymbol{f}^{-1}}{\\partial \\boldsymbol{x}} \\right)  \\right| $$\n",
    "\n",
    "The main idea of normalizing flows is to learn an invertible mapping between a very complex distribution (for example a distribution of a physics variable) and a very simple distribution. Let us consider the \"generative\" part first, where we attempt to generate new samples from a Normalizing Flow. In this case, we start out with a simple distribution, for example a standard Gaussian, and then repeatedly apply the random variable transformations to acquire a complex distribution similar to the input distribution that we want to approximate. Let's say we apply $k$ transformations $f_{1}...f_{k}$, starting out from a simple distribution $u_{0}$ and going to a complex distribution $u_{k}$, then the change of random variables formula becomes:\n",
    "\n",
    "$$ p(\\boldsymbol{u}_{k}) = p(\\boldsymbol{u}_{0})\\prod_{i} \\left|\\det \\left( \\frac{\\partial \\boldsymbol{f}_{i}}{\\partial \\boldsymbol{u}_{i-1}} \\right)  \\right|^{-1} $$\n",
    "\n",
    "\n",
    "From this formula, we also see the reason why this method is called \"normalizing flows\". The random variable \"flows\" through a series of transformations, while staying normalized by the scaling with the Jacobian determinant.\n",
    "\n",
    "<br />\n",
    "<img src=\"https://drive.google.com/uc?id=1fHZeW9lKGOPl7lZehy8dY0wxLVzjnQFZ\" width=\"1000\">\n",
    "<br />\n",
    "\n",
    "In a practical case, our \"complex\" random variables $\\boldsymbol{u}_{k}$ would then approximate our input variables.\n",
    "\n",
    "Since we choose our $\\boldsymbol{f}_{i}$ to be invertible, we can also run the Normalizing Flow in the other direction, starting out with our complex input distribution and mapping this distribution tho a standard Gaussian. This is the direction used for density/likelihood estimation (i.e. if you want to evaluate the likelihood of new data under the learned distribution) and also the one used for training.\n",
    "\n",
    "For training a normalizing flow, remember that our $\\boldsymbol{f}_{i}$ are actually neural networks with parameters $\\theta$. We train our flow by optimizing these parameters in a negative logarithmic likelihood minimization (which is our loss function then). So we basically take our input samples $\\boldsymbol{u}_{k}$, take it \"backwards\" through the flow (i.e. in the direction of the blue arrows in the picture above) and then tune the parameters of our $\\boldsymbol{f}_{i}$ such that the likelihood of the so transformed distribution under a standard Gaussian gets maximized.\n",
    "\n",
    "To be able to use normalizing flows efficiently, we therefore need functions $f$ that are both **invertible** and **have a tractable Jacobian** that is **easy to compute**. Several methods have been proposed so far and today, you will work mainly with so called \"autoregressive flows\", where the functions $f$ are chosen in a way such that the Jacobian is guaranteed to be an upper triangular matrix, so we can compute the determinant simply by multiplying its diagonal elements.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PjLFuT_3Se__"
   },
   "source": [
    "## 3.) The TensorFlow Probability package \n",
    "\n",
    "Today, we are going to use the TensorFlow Probability library, that is easy to set up and already contains implementations of many of the different flow models. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_UZnAEM-SfAA",
    "outputId": "35d6a035-c9fe-4a42-9b1c-064767faba44"
   },
   "outputs": [],
   "source": [
    "# Let's import some TensorFlow probability as well as some modules first\n",
    "import tensorflow as tf\n",
    "import tensorflow_probability as tfp\n",
    "from tensorflow_probability import distributions as tfd\n",
    "from tensorflow_probability import bijectors as tfb\n",
    "import tensorflow.keras as tfk\n",
    "import tensorflow.keras.layers as tfkl\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "print(f\"CUDA available: {len(tf.config.list_physical_devices('GPU')) > 0}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "me9llvR_SfAB"
   },
   "source": [
    "Some background about the different modules:\n",
    "\n",
    "The **bijectors module** contains the implementation of the actual flow models. Each class in this model is already containing the calculation of the Jacobian determinant, so you do not need to do this yourself. The models contain bijective (flow-) equivalents of \"standard\" layers such as BatchNormalization as well as models used for specific flows (such as the \"AutoregressiveNetwork\").\n",
    "\n",
    "The **distributions module** contains many \"base distributions\" that we can build our flow on (for example a normal distribution as with our example above). The most important class in this module is the `TransformedDistribution`, which takes as parameter a base distribution and an invertible, differentiable transformation (typically one or more transformations from the bijectors module) and models the distribution resulting from applying said tranformation to the base distribution.\n",
    "\n",
    "As a first step. let's do a very simple example (adapted from TensorFlow probability guide (<a href=\"https://www.tensorflow.org/probability/api_docs/python/tfp/bijectors/AutoregressiveNetwork\">link</a>)):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "m2Krshw4SfAB"
   },
   "outputs": [],
   "source": [
    "# Generate a dataset, the density of which we would like to approximate\n",
    "num_samples = 2000\n",
    "x2 = np.random.randn(num_samples).astype(dtype=np.float32) * 2.\n",
    "x1 = np.random.randn(num_samples).astype(dtype=np.float32) + (x2 * x2 / 4.)\n",
    "data = np.stack([x1, x2], axis=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5g2UYcxtSfAB"
   },
   "source": [
    "Let's plot the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 279
    },
    "id": "2T_DJT2-SfAB",
    "outputId": "3d3a2ab6-6275-4410-de5c-a6c148d27be0"
   },
   "outputs": [],
   "source": [
    "plt.scatter(data[:, 0], data[:, 1], color='black', marker='.', linewidth=0)\n",
    "plt.xlabel(\"x1\")\n",
    "plt.ylabel(\"x2\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EJCJe_sMSfAC"
   },
   "source": [
    "The autoregressive models that we will use for this exercise are built with several blocks of MADE networks (Masked Autoencoder for Distribution Estimation, <a href=\"https://arxiv.org/abs/1502.03509\">arxiv</a>). For this very first test, we will not directly build the flow, but just use a single MADE network for estimation of the above density."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "W5Bx-UG7SfAC"
   },
   "outputs": [],
   "source": [
    "# Density estimation with MADE.\n",
    "# Since we later want to use a Gaussian as our autoregressive distribution, we need to set\n",
    "# \"params\" to 2, such that the MADE network can parameterize its mean and logarithmic standard deviation.\n",
    "# The \"hidden units\" parameter tells us that the MADE network has two layers with 10 nodes each.\n",
    "made = tfb.AutoregressiveNetwork(params=2, hidden_units=[10, 10])\n",
    "\n",
    "# Create the TransformedDistribution. We will use a normal distribution with mean 0 and standard\n",
    "# deviation 1 as base distribution.\n",
    "# The tfd.Sample allows us to define normal distributions in multiple dimensions.\n",
    "# Since we have two dimensions (x1 and x2), we use tfd.Sample and set sample_shape\n",
    "# to [2].\n",
    "# We also define the bijector that transforms this distribution. In this case, it is\n",
    "# a masked autoregressive flow, just consisting of our single MADE block.\n",
    "bijector = tfb.MaskedAutoregressiveFlow(made)\n",
    "distribution = tfd.TransformedDistribution(\n",
    "    distribution=tfd.Sample(tfd.Normal(loc=0., scale=1.), sample_shape=[2]),\n",
    "    bijector=bijector)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "k8xhHJ6ESfAD"
   },
   "outputs": [],
   "source": [
    "# Construct the model:\n",
    "# We use as input a standard input layer.\n",
    "# The output is the logarithmic probability (i.e. log likelihood) of our transformed\n",
    "# distribution:\n",
    "x_ = tfkl.Input(shape=(2,), dtype=tf.float32)\n",
    "log_prob_ = distribution.log_prob(x_)\n",
    "model = tfk.Model(x_, log_prob_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "2MI-gB99SfAE",
    "outputId": "390b1d2c-e5b4-4502-d24e-91f9f477b0c8"
   },
   "outputs": [],
   "source": [
    "# Let's compile the model, using the Adam optimizer.\n",
    "# The loss is simply the negative log likelihood,\n",
    "# here simply wrapped in a lambda function.\n",
    "model.compile(optimizer=tf.optimizers.Adam(),\n",
    "              loss=lambda _, log_prob: -log_prob)\n",
    "\n",
    "# let's train the model\n",
    "batch_size = 25\n",
    "\n",
    "# since normalizing flows are an unsupervised method,\n",
    "# we simply provide an array of zeros for the \"y\"\n",
    "# parameter here, which has the same number of entries\n",
    "# as our input data.\n",
    "model.fit(x=data,\n",
    "          y=np.zeros((num_samples, 0), dtype=np.float32),\n",
    "          batch_size=batch_size,\n",
    "          epochs=20,\n",
    "          steps_per_epoch=num_samples // batch_size,\n",
    "          shuffle=True,\n",
    "          verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 279
    },
    "id": "rwoxoJXQSfAE",
    "outputId": "df1d3404-d9fc-41ee-ba53-5ce898ea7803"
   },
   "outputs": [],
   "source": [
    "# Let's use the fitted distribution and sample from it\n",
    "samples = distribution.sample(1000)\n",
    "plt.scatter(samples[:, 0], samples[:, 1], color='darkblue', marker='o', linewidth=0)\n",
    "plt.xlabel(\"x1\")\n",
    "plt.ylabel(\"x2\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 283
    },
    "id": "5Fy5VaGgSfAE",
    "outputId": "6c8c166e-0835-481a-fee9-ea40c9131f18"
   },
   "outputs": [],
   "source": [
    "# Let's also plot how well our data is mapped to the Gaussian base distribution\n",
    "\n",
    "inv = bijector.inverse(data)\n",
    "plt.scatter(inv[:, 0], inv[:, 1], color='black', marker='.', linewidth=0)\n",
    "plt.xlabel(\"Transformed x1\")\n",
    "plt.ylabel(\"Transformed x2\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UbtlEySDSfAF"
   },
   "source": [
    "As we can see, our initial density is estimated rather poorly. We also see that our mapped data is far from being Gaussian distributed in the latent space. So what did we do wrong? We actually didn't use a proper flow! For getting to know the TensorFlow probability module, a single MADE network was used here. This would be equivalent of using only a single transformation $f$ instead of using many of them. With the knowledge of how to use the module, we can now start building a powerful flow, consisting of multiple MADE blocks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "odmPNp-hSfAF"
   },
   "source": [
    "# Self-defined MADE block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "K23FrzYXSfAF"
   },
   "outputs": [],
   "source": [
    "# This is simply a copy of the original \"AutoregressiveNetwork\" class in tfp.bijectors. The only reason we need to do this is that we want to apply a tanh\n",
    "# on the output log-scale when the Network is called. This allows for better regularization and helps with \"inf\" and \"nan\" values that otherwise would\n",
    "# frequently occur during training.\n",
    "class Made(tfb.AutoregressiveNetwork):\n",
    "    def __init__(self, params, event_shape=None, conditional=False, conditional_event_shape=None, conditional_input_layers='first_layer', hidden_units=None,\n",
    "                 input_order='left-to-right', hidden_degrees='equal', activation=None, use_bias=True,kernel_initializer='glorot_uniform', bias_initializer='zeros',\n",
    "                 kernel_regularizer=None, bias_regularizer=None, kernel_constraint=None, bias_constraint=None, validate_args=False, **kwargs):\n",
    "        \n",
    "        super().__init__(params=params, event_shape=event_shape, conditional=conditional, conditional_event_shape=conditional_event_shape,\n",
    "                         conditional_input_layers=conditional_input_layers, hidden_units=hidden_units, input_order=input_order, hidden_degrees=hidden_degrees,\n",
    "                         activation=activation, use_bias=use_bias, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer,\n",
    "                         kernel_regularizer=kernel_regularizer, bias_regularizer=bias_regularizer, kernel_constraint=kernel_constraint, bias_constraint=bias_constraint,\n",
    "                         validate_args=validate_args, **kwargs)\n",
    "    \n",
    "    def call(self, x, conditional_input=None):\n",
    "        \n",
    "\n",
    "        result = super().call(x, conditional_input=conditional_input)\n",
    "        \n",
    "        shift, log_scale = tf.unstack(result, num=2, axis=-1)\n",
    "\n",
    "        return shift, tf.math.tanh(log_scale)\n",
    "    \n",
    "    def get_config(self):\n",
    "        \n",
    "        config = super().get_config().copy()\n",
    "        \n",
    "        return config"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u5mFe-45SfAG"
   },
   "source": [
    "# Define model and compile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Und8zFbjSfAG"
   },
   "outputs": [],
   "source": [
    "def compile_MAF_model(num_made, num_inputs, num_cond_inputs=None, made_layers=[128], base_lr=1.0e-3, end_lr=1.0e-4, return_layer_list=False):\n",
    "\n",
    "  if num_cond_inputs is not None:\n",
    "    conditional = True\n",
    "    cond_event_shape = (num_cond_inputs,)\n",
    "  else:\n",
    "    conditional = False\n",
    "    cond_event_shape = None\n",
    "\n",
    "  made_list = []\n",
    "  for i in range(num_made):\n",
    "    made_list.append(tfb.MaskedAutoregressiveFlow(\n",
    "        shift_and_log_scale_fn=Made(params=2, hidden_units=made_layers, event_shape=(num_inputs,), conditional=conditional,\n",
    "                                    conditional_event_shape=cond_event_shape, activation='relu', name=f\"made_{i}\"), name=f\"maf_{i}\"))\n",
    "    \n",
    "    #made_list.append(tfb.BatchNormalization())\n",
    "    made_list.append(tfb.Permute(permutation=np.array(np.arange(0, num_inputs)[::-1])))\n",
    "                     \n",
    "  # remove final permute layer\n",
    "  made_chain = tfb.Chain(list(reversed(made_list[:-1])))\n",
    "\n",
    "  # we want to transform to gaussian distribution with mean 0 and std 1 in latent space\n",
    "  distribution = tfd.TransformedDistribution(\n",
    "    distribution=tfd.Sample(tfd.Normal(loc=0., scale=1.), sample_shape=[num_inputs]),\n",
    "    bijector=made_chain)\n",
    "\n",
    "  x_ = tfk.layers.Input(shape=(num_inputs,), name=\"aux_input\")\n",
    "  input_list = [x_]\n",
    "\n",
    "  if conditional:\n",
    "    c_ = tfk.layers.Input(shape=(num_cond_inputs,), name=\"cond_input\")\n",
    "    input_list.append(c_)\n",
    "\n",
    "    current_kwargs = {}\n",
    "    for i in range(num_made):\n",
    "      current_kwargs[f\"maf_{i}\"] = {'conditional_input' : c_}\n",
    "    \n",
    "  else:\n",
    "    current_kwargs = {}\n",
    "  \n",
    "  log_prob_ = distribution.log_prob(x_, bijector_kwargs=current_kwargs)\n",
    "  \n",
    "  model = tfk.Model(input_list, log_prob_)\n",
    "  max_epochs = 100  # maximum number of epochs of the training\n",
    "  learning_rate_fn = tfk.optimizers.schedules.PolynomialDecay(base_lr, max_epochs, end_lr, power=0.5)\n",
    "  model.compile(optimizer=tfk.optimizers.Adam(learning_rate=learning_rate_fn),\n",
    "              loss=lambda _, log_prob: -log_prob)\n",
    "  \n",
    "  if return_layer_list:\n",
    "    return model, distribution, made_list\n",
    "  else:\n",
    "    return model, distribution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6J_LURVNjNlF"
   },
   "source": [
    "# Data Generation using Flows\n",
    "\n",
    "One major advantage of Normalizing Flow Networks is that they are inherently invertible. This means that we can not only use them for density estimation by mapping our data set to a well known distribution like a Normal Gaussian, but we can also do the reverse, mapping the Normal Gaussian to our data distribution. This essentially allows us to generate new data by feeding Gaussian samples backwards through our model. \n",
    "\n",
    "One common example used for benchmarking is the so-called 'two moon' dataset, consisting of two interlocking half circles. On the one hand it tests the generative model's ability to replicate a complex structure, on the other hand this data set also features two separate subsets, which our generative model will have to keep separate as well. \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ag-jHVUJSfAH"
   },
   "source": [
    "# Load two moon data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 265
    },
    "id": "tJyPtwUJSfAH",
    "outputId": "ee39c73d-7594-4dfe-ec16-a1497d44f352"
   },
   "outputs": [],
   "source": [
    "from sklearn import cluster, datasets, mixture\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "# for two moons dataset\n",
    "n = 10000\n",
    "X_moons, _ = datasets.make_moons(n_samples=n, noise=.05)\n",
    "X_moons = StandardScaler().fit_transform(X_moons)\n",
    "\n",
    "plt.scatter(X_moons[:, 0], X_moons[:, 1], color='black', marker='.', linewidth=0)\n",
    "plt.savefig(\"./original_two_moons.pdf\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TklSlEUvldqe"
   },
   "source": [
    "As a first test, let us take our simple Flow and train it on this dataset:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "68w2L47LSfAH"
   },
   "source": [
    "# Train Simple Flow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6iZm_lqYXlD8"
   },
   "outputs": [],
   "source": [
    "#Define two-block made network\n",
    "num_made = 2\n",
    "epochs = 10\n",
    "\n",
    "model, distribution, list_made = compile_MAF_model(num_made, num_inputs=2, return_layer_list=True)\n",
    "\n",
    "# definie intermediate outputs for later\n",
    "feat_extraction_dists = []\n",
    "listR = list(reversed(list_made[:-1]))\n",
    "\n",
    "made_chain = tfb.Chain([])\n",
    "dist = tfd.TransformedDistribution(\n",
    "  distribution=tfd.Sample(tfd.Normal(loc=0., scale=1.), sample_shape=[2]),\n",
    "  bijector=made_chain)\n",
    "feat_extraction_dists.append(dist)\n",
    "\n",
    "for i in range(1, len(listR), 2):\n",
    "  made_chain = tfb.Chain(listR[0:i])\n",
    "  dist = tfd.TransformedDistribution(\n",
    "    distribution=tfd.Sample(tfd.Normal(loc=0., scale=1.), sample_shape=[2]),\n",
    "    bijector=made_chain)\n",
    "  feat_extraction_dists.append(dist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-FY5WTjvSfAH",
    "outputId": "d41ea23f-1a14-4504-f97e-303eed9adbbd"
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "#as this is once again a unsupervised task, the target vector y is again zeros\n",
    "batch_size = 32\n",
    "model.fit(x=X_moons,\n",
    "          y=np.zeros((X_moons.shape[0], 0), dtype=np.float32),\n",
    "          batch_size= batch_size,\n",
    "          epochs=epochs,\n",
    "          steps_per_epoch=X_moons.shape[0] // batch_size,\n",
    "          verbose=1,\n",
    "          shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2gfuIzAUmxOh"
   },
   "source": [
    "Now that the model is trained we can use it to try and generate new samples. \n",
    "\n",
    "Thankfully TensorFlow Probability provides an easy way of doing this, we simply need to call the sample() function on the returned distribution variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 265
    },
    "id": "bxVawO9mmv-2",
    "outputId": "25287158-4bbe-4285-bc8c-c860553e1710"
   },
   "outputs": [],
   "source": [
    "#generate 1000 new samples\n",
    "samples = distribution.sample(1000)\n",
    "\n",
    "#transform them like we did when plotting the original dataset\n",
    "samples = StandardScaler().fit_transform(samples)\n",
    "\n",
    "#plot the results\n",
    "plt.scatter(samples[:, 0], samples[:, 1], color='darkblue', marker='o', linewidth=0)\n",
    "plt.scatter(X_moons[:, 0], X_moons[:, 1], color='black', marker='.', linewidth=0)\n",
    "plt.savefig(\"./sampled_two_moons.pdf\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GvxYLZYWuRbO"
   },
   "source": [
    "Additionally we can use the earlier prepared code to show intermediate steps of the transformations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 761
    },
    "id": "S0BjIpxntCjE",
    "outputId": "5955dc75-10a5-4709-d5fc-34e7e332f6b1"
   },
   "outputs": [],
   "source": [
    "for d in feat_extraction_dists:\n",
    "  out = d.sample(1000)\n",
    "  out = StandardScaler().fit_transform(out)\n",
    "\n",
    "  plt.scatter(out[:, 0], out[:, 1], color='darkblue', marker='o', linewidth=0)\n",
    "  plt.show()\n",
    "  plt.close()\n",
    "\n",
    "out = distribution.sample(1000)\n",
    "out = StandardScaler().fit_transform(out)\n",
    "\n",
    "plt.scatter(out[:, 0], out[:, 1], color='darkblue', marker='o', linewidth=0)\n",
    "plt.show()\n",
    "plt.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XiJE_Bljodqs"
   },
   "source": [
    "As we can see, the result is somewhat underwhelming. This is because we once again used a very simple setup, that is unable to correctly model the dataset. \n",
    "\n",
    "**Task: Modify the Code below to improve the generation result**\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 595
    },
    "id": "6Of8xenvSfAH",
    "outputId": "172492ae-0fdf-4b6d-d809-98fcf09674bb"
   },
   "outputs": [],
   "source": [
    "#Define two-block made network\n",
    "num_made = 2\n",
    "epochs = 10\n",
    "\n",
    "model, distribution, list_made = compile_MAF_model(num_made, num_inputs=2, return_layer_list=True)\n",
    "\n",
    "# definie intermediate outputs for later\n",
    "feat_extraction_dists = []\n",
    "listR = list(reversed(list_made[:-1]))\n",
    "\n",
    "made_chain = tfb.Chain([])\n",
    "dist = tfd.TransformedDistribution(\n",
    "  distribution=tfd.Sample(tfd.Normal(loc=0., scale=1.), sample_shape=[2]),\n",
    "  bijector=made_chain)\n",
    "feat_extraction_dists.append(dist)\n",
    "\n",
    "for i in range(1, len(listR), 2):\n",
    "  made_chain = tfb.Chain(listR[0:i])\n",
    "  dist = tfd.TransformedDistribution(\n",
    "    distribution=tfd.Sample(tfd.Normal(loc=0., scale=1.), sample_shape=[2]),\n",
    "    bijector=made_chain)\n",
    "  feat_extraction_dists.append(dist)\n",
    "\n",
    "\n",
    "#as this is once again a unsupervised task, the target vector y is again zeros\n",
    "batch_size = 32\n",
    "model.fit(x=X_moons,\n",
    "          y=np.zeros((X_moons.shape[0], 0), dtype=np.float32),\n",
    "          batch_size= batch_size,\n",
    "          epochs=epochs,\n",
    "          steps_per_epoch=X_moons.shape[0] // batch_size,\n",
    "          verbose=1,\n",
    "          shuffle=True)\n",
    "\n",
    "\n",
    "#generate 1000 new samples\n",
    "samples = distribution.sample(1000)\n",
    "\n",
    "#transform them like we did when plotting the original dataset\n",
    "samples = StandardScaler().fit_transform(samples)\n",
    "\n",
    "#plot the results\n",
    "plt.scatter(samples[:, 0], samples[:, 1], color='darkblue', marker='o', linewidth=0)\n",
    "plt.scatter(X_moons[:, 0], X_moons[:, 1], color='black', marker='.', linewidth=0)\n",
    "plt.savefig(\"./sampled_two_moons.pdf\")\n",
    "plt.show()\n",
    "plt.close()\n",
    "\n",
    "#plot the intermediate results\n",
    "for d in feat_extraction_dists:\n",
    "  out = d.sample(1000)\n",
    "  out = StandardScaler().fit_transform(out)\n",
    "\n",
    "  plt.scatter(out[:, 0], out[:, 1], color='darkblue', marker='o', linewidth=0)\n",
    "  plt.show()\n",
    "  plt.close()\n",
    "\n",
    "out = distribution.sample(1000)\n",
    "out = StandardScaler().fit_transform(out)\n",
    "\n",
    "plt.scatter(out[:, 0], out[:, 1], color='darkblue', marker='o', linewidth=0)\n",
    "plt.show()\n",
    "plt.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CXtsCPlXSfAI"
   },
   "source": [
    "Now that we have a nicely working model, we can try and adapt this to other problems\n",
    "\n",
    "**Task: Use what you learned previously to train a generative flow on the first example dataset**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 279
    },
    "id": "aOjtzjMMt7nv",
    "outputId": "615330f4-841c-47dd-ff5d-e139cc078363"
   },
   "outputs": [],
   "source": [
    "# Generate a dataset, the density of which we would like to approximate\n",
    "num_samples = 2000\n",
    "x2 = np.random.randn(num_samples).astype(dtype=np.float32) * 2.\n",
    "x1 = np.random.randn(num_samples).astype(dtype=np.float32) + (x2 * x2 / 4.)\n",
    "data = np.stack([x1, x2], axis=-1)\n",
    "\n",
    "plt.scatter(data[:, 0], data[:, 1], color='black', marker='.', linewidth=0)\n",
    "plt.xlabel(\"x1\")\n",
    "plt.ylabel(\"x2\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 265
    },
    "id": "zXHKXLwDYtFe",
    "outputId": "53905e54-b0c7-4cea-91de-80790074ffef"
   },
   "outputs": [],
   "source": [
    "######################\n",
    "#Add code here\n",
    "######################\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fOkCS4rTZBKK"
   },
   "source": [
    "# Conditional Generative Flows\n",
    "\n",
    "Many generative tasks in particle physics require the generative model to not only generate a random point from the data set, but instead to generate a point with a specific property. One Example is shower simulation, where a generative shower simulator should be able to produce showers for a given particle energy, rather than just for a random one. \n",
    "\n",
    "This is where conditional generative models become important. In the previous example the energy value would be the condition given to the model. \n",
    "\n",
    "Let us try this principle on flows with our previous two moons data set. As one can see, it is made up of two distinct subsets, (the eponymous moons). Our conditioning will be to which moon a datapoint belongs, e.g. if we tell the model to generate points with the label 0 they should be from the top moon, while points with the label 1 should be from the bottom."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 265
    },
    "id": "7BHN2hoKaDNC",
    "outputId": "96aefd26-3d92-465a-e563-b40651ea3480"
   },
   "outputs": [],
   "source": [
    "from sklearn import cluster, datasets, mixture\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "# for two moons dataset\n",
    "n = 10000\n",
    "# the two moons sampling function already returns labels to which moon a data point belongs\n",
    "X_moons, label_moon = datasets.make_moons(n_samples=n, noise=.05)\n",
    "X_moons = StandardScaler().fit_transform(X_moons)\n",
    "\n",
    "\n",
    "\n",
    "plt.scatter(X_moons[:, 0][label_moon==0], X_moons[:, 1][label_moon==0], color='orangered', marker='.', linewidth=0)\n",
    "plt.scatter(X_moons[:, 0][label_moon==1], X_moons[:, 1][label_moon==1], color='green', marker='.', linewidth=0)\n",
    "\n",
    "\n",
    "plt.savefig(\"./original_two_moons.pdf\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "y5rYwUtonxFc",
    "outputId": "082a108b-827f-4b21-8f01-5898300cb7f0"
   },
   "outputs": [],
   "source": [
    "#Define one-block made network\n",
    "num_made = 15\n",
    "epochs = 30\n",
    "\n",
    "model, distribution = compile_MAF_model(num_made, num_inputs=2, num_cond_inputs=1)\n",
    "\n",
    "#as this is once again an unsupervised task, the target vector y is again zeros during training\n",
    "batch_size = 32\n",
    "model.fit(x=[X_moons, label_moon],\n",
    "          y=np.zeros((X_moons.shape[0], 0), dtype=np.float32),\n",
    "          batch_size= batch_size,\n",
    "          epochs=epochs,\n",
    "          steps_per_epoch=X_moons.shape[0] // batch_size,\n",
    "          verbose=1,\n",
    "          shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "c_QgmwT5zGI9"
   },
   "source": [
    "Now lets see how well our model performs, first by using random labels. This should look like the default two moons set. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 265
    },
    "id": "lNcs7DP-p4ON",
    "outputId": "6c256dea-8efd-4067-8422-9c19da5e78fd"
   },
   "outputs": [],
   "source": [
    "n_samples = 10000\n",
    "\n",
    "cond = np.random.randint(2, size=(n_samples,1))\n",
    "current_kwargs ={}\n",
    "for i in range(num_made):\n",
    "    current_kwargs[f\"maf_{i}\"] = {'conditional_input' : cond}\n",
    "\n",
    "\n",
    "samples = distribution.sample ( (n_samples, ), bijector_kwargs=current_kwargs)\n",
    "\n",
    "plt.scatter(samples[:, 0], samples[:, 1], color='darkblue', marker='o', linewidth=0)\n",
    "plt.scatter(X_moons[:, 0], X_moons[:, 1], color='black', marker='.', linewidth=0)\n",
    "plt.savefig(\"./sampled_two_moons.pdf\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jO8tZDAYbYj3"
   },
   "source": [
    "We can also test how well our conditioning performs by passing only the label 0 or only the label 1 to our network.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 265
    },
    "id": "-1mTh2tCroPU",
    "outputId": "82cd3fc9-a2e3-45c4-84f5-793b1f9eea7e"
   },
   "outputs": [],
   "source": [
    "n_samples = 10000\n",
    "\n",
    "cond1 = np.ones((n_samples,1))\n",
    "cond0 = np.zeros((n_samples,1))\n",
    "\n",
    "\n",
    "current_kwargs0 ={}\n",
    "for i in range(num_made):\n",
    "    current_kwargs0[f\"maf_{i}\"] = {'conditional_input' : cond0}\n",
    "samples0 = distribution.sample ( (n_samples, ), bijector_kwargs=current_kwargs0)\n",
    "\n",
    "current_kwargs1 ={}\n",
    "for i in range(num_made):\n",
    "    current_kwargs1[f\"maf_{i}\"] = {'conditional_input' : cond1}\n",
    "samples1 = distribution.sample ( (n_samples, ), bijector_kwargs=current_kwargs1)\n",
    "\n",
    "plt.scatter(samples0[:, 0], samples0[:, 1], color='orangered', marker='.', linewidth=0)\n",
    "plt.scatter(samples1[:, 0], samples1[:, 1], color='green', marker='.', linewidth=0)\n",
    "plt.savefig(\"./sampled_two_moons.pdf\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9ALnN9Vewh3Y"
   },
   "source": [
    "# High energy physics use case"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "49n56ljzwqAi"
   },
   "source": [
    "To illustrate the use of normalizing flows, we now want to a apply it to a high energy physics example. For this aim, we use the LHCO 2020 challenge R&D dataset (https://lhco2020.github.io/homepage/). It consists of simulated QCD multijet events in addition to a signal process of a W' boson with a mass of $m_{W'} = 3.5$ TeV decaying into an X boson ($m_{X} = 500$ GeV) and a Y boson ($m_{Y} = 100$ GeV). Specifically, we use the high-level features set, which contains the 4-momenta of the two leading jets, their (1,2,3) subjettiness and a label denoting whether the event is a signal (1) or background (0) process.\n",
    "\n",
    "While the main purpose of this dataset was an anomaly detection competition, we will continue looking at normalizing flows from the perspective of artificial data generation. Nevertheless, if you manage to solve the advanced exercise in the end, you will also be well prepared to understand and use the ANODE method (Anomaly Detection with Density Estimation, <a href=\"https://arxiv.org/abs/2001.04990\">arxiv</a>). More on this might follow in the anomaly detection workshop.\n",
    "\n",
    "So let's first download the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Mi9qyAMfwlni",
    "outputId": "ff37a079-b586-4f47-ac91-96b35c5faca4"
   },
   "outputs": [],
   "source": [
    "!wget https://zenodo.org/record/4536377/files/events_anomalydetection_v2.features.h5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ltEQiOQtwvWR"
   },
   "source": [
    "The data is stored as a pandas dataframe. We load the data and pick a random 10% subset for the sake of computation speed. Feel free to use the full data in your own time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_E6vLfBywxPH",
    "outputId": "78ba3497-a2b5-43db-9be7-9f721e361974"
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "dataset = pd.read_hdf(\"events_anomalydetection_v2.features.h5\")\n",
    "dataset = dataset.sample(frac=0.1)\n",
    "print(\"The variables are:\", [col for col in dataset.columns])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5so7o1UBw0of"
   },
   "source": [
    "We now order the two jets by mass, such that $m_{j1} \\leq m_{j2}$. Then we extract four interesting features, which could be used here to distinguish signal from background: $m_{j1}$, $m_{j2}-m_{j1}$, $\\tau_{21,1}$, $\\tau_{21,2}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JlJYsGU8w2Wm",
    "outputId": "5cdf4787-71fe-4bbc-cfea-2ef75a0d7060"
   },
   "outputs": [],
   "source": [
    "mj1mj2=np.array(dataset[['mj1','mj2']])\n",
    "tau21=np.array(dataset[['tau2j1','tau2j2']])/(1e-5+np.array(dataset[['tau1j1','tau1j2']])) # add small number to denominator to avoid division by zero\n",
    "mjmin=mj1mj2[range(len(mj1mj2)),np.argmin(mj1mj2,axis=1)]\n",
    "mjmax=mj1mj2[range(len(mj1mj2)),np.argmax(mj1mj2,axis=1)]\n",
    "tau21min=tau21[range(len(mj1mj2)),np.argmin(mj1mj2,axis=1)]\n",
    "tau21max=tau21[range(len(mj1mj2)),np.argmax(mj1mj2,axis=1)]\n",
    "data = np.dstack((mjmin/1000, (mjmax-mjmin)/1000, tau21min, tau21max))[0] # put data together and convert GeV to TeV\n",
    "labels = np.array(dataset['label'])\n",
    "print(\"data.shape =\", data.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4UJ8WEKvw5Hu"
   },
   "source": [
    "Let's plot these four features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 591
    },
    "id": "i9xbcZsXw65-",
    "outputId": "b44b376f-ca72-4cd4-fd91-4097dbf5207a"
   },
   "outputs": [],
   "source": [
    "## visualize the 4 extracted variables\n",
    "fig, ax = plt.subplots(2,2, figsize=(12,8))\n",
    "_, common_bins, _ = ax[0,0].hist(data[:,0], 100, alpha=0.6, label=\"data\")\n",
    "ax[0,0].hist(data[:,0][labels==0], common_bins, label=\"background\", edgecolor=\"red\", linewidth=1.3, histtype=\"step\")\n",
    "ax[0,0].hist(data[:,0][labels==1], common_bins, label=\"signal\", edgecolor=\"green\", linewidth=1.3, histtype=\"step\")\n",
    "ax[0,0].set_xlabel(\"lower jet mass (TeV)\")\n",
    "ax[0,0].set_ylabel(\"events\")\n",
    "ax[0,0].set_yscale(\"log\")\n",
    "_, common_bins, _ = ax[0,1].hist(data[:,1], 100, alpha=0.6, label=\"data\")\n",
    "ax[0,1].hist(data[:,1][labels==0], common_bins, label=\"background\", edgecolor=\"red\", linewidth=1.3, histtype=\"step\")\n",
    "ax[0,1].hist(data[:,1][labels==1], common_bins, label=\"signal\", edgecolor=\"green\", linewidth=1.3, histtype=\"step\")\n",
    "ax[0,1].set_xlabel(\"jet mass difference (TeV)\")\n",
    "ax[0,1].set_ylabel(\"events\")\n",
    "ax[0,1].set_yscale(\"log\")\n",
    "_, common_bins, _ = ax[1,0].hist(data[:,2], 100, alpha=0.6, label=\"data\")\n",
    "ax[1,0].hist(data[:,2][labels==0], common_bins, label=\"background\", edgecolor=\"red\", linewidth=1.3, histtype=\"step\")\n",
    "ax[1,0].hist(data[:,2][labels==1], common_bins, label=\"signal\", edgecolor=\"green\", linewidth=1.3, histtype=\"step\")\n",
    "ax[1,0].set_xlabel(r\"$\\tau_{21,1}$\")\n",
    "ax[1,0].set_ylabel(\"events\")\n",
    "_, common_bins, _ = ax[1,1].hist(data[:,3], 100, alpha=0.6, label=\"data\")\n",
    "ax[1,1].hist(data[:,3][labels==0], common_bins, label=\"background\", edgecolor=\"red\", linewidth=1.3, histtype=\"step\")\n",
    "ax[1,1].hist(data[:,3][labels==1], common_bins, label=\"signal\", edgecolor=\"green\", linewidth=1.3, histtype=\"step\")\n",
    "ax[1,1].set_xlabel(r\"$\\tau_{21,2}$\")\n",
    "ax[1,1].set_ylabel(\"events\")\n",
    "leg_handles, leg_labels = ax[0,0].get_legend_handles_labels()\n",
    "fig.legend(leg_handles, leg_labels, loc='upper right')\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7N0H01gTw9t2"
   },
   "source": [
    "We want train a density estimator to generate such data ourselves. We use the same MAF model as before, just reinstantiate it. This time with 4 input variables. In addition, we increase the complexity of the MADE layers in order to keep up with the more complex task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bAUKIc2Tw_8f"
   },
   "outputs": [],
   "source": [
    "num_made = 15\n",
    "num_inputs = 4\n",
    "\n",
    "model, distribution = compile_MAF_model(num_made, num_inputs, made_layers=[128, 128])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mRuMFBwDxBq7"
   },
   "source": [
    "Before plugging our data into the MAF, we should preprocess the variables, so the model can learn them more easily. Density estimators tend to have difficulties with sharp  edges and boundaries. We first normalize each variable to lie between 0 and 1, then we apply a logit transformation $\\log(\\frac{x}{1-x})$ mapping the variables to ($-\\infty$, $+\\infty$) and then we shift the mean to 0 and divide by the standard deviation. Here we define these three transformations, as well as their inverse functions. Be aware that for computing likelihoods, we would also need to take into account the Jacobian of each transformation. In a purely generative approach, we don't need to worry about this now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "G1gwGkwkxDs2"
   },
   "outputs": [],
   "source": [
    "## defining variable transformations\n",
    "\n",
    "def normalize_data(in_data, max_val, min_val):\n",
    "  new_data = (in_data-min_val)/(max_val-min_val)\n",
    "  mask = np.prod(((new_data < 1) & (new_data > 0 )), axis=1, dtype=bool)\n",
    "  new_data = new_data[mask]\n",
    "  return new_data, mask\n",
    "\n",
    "def logit_data(in_data):\n",
    "  new_data = np.log(in_data/(1-in_data))\n",
    "  return new_data\n",
    "\n",
    "def standardize_data(in_data, mean_val, std_val):\n",
    "  new_data = (in_data - mean_val)/std_val\n",
    "  return new_data\n",
    "\n",
    "\n",
    "## defining their inverse transformations\n",
    "\n",
    "def normalize_inverse(in_data, max_val, min_val):\n",
    "  new_data = in_data*(max_val-min_val) + min_val\n",
    "  return new_data\n",
    "\n",
    "def logit_inverse(in_data):\n",
    "  new_data = (1+np.exp(-in_data))**(-1)\n",
    "  return new_data\n",
    "\n",
    "def standardize_inverse(in_data, mean_val, std_val):\n",
    "  new_data = std_val*in_data + mean_val\n",
    "  return new_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "x-STT0vaxGf9"
   },
   "outputs": [],
   "source": [
    "## transform data and save max, min, mean, std values for the backtransformation later\n",
    "\n",
    "max_values = np.max(data, keepdims=True, axis=0)\n",
    "min_values = np.min(data, keepdims=True, axis=0)\n",
    "\n",
    "## normalize\n",
    "transformed_data, mask = normalize_data(data, max_values, min_values)\n",
    "\n",
    "## logit\n",
    "transformed_data = logit_data(transformed_data)\n",
    "\n",
    "## standardize\n",
    "mean_values = np.mean(transformed_data, keepdims=True, axis=0)\n",
    "std_values = np.std(transformed_data, keepdims=True, axis=0)\n",
    "transformed_data = standardize_data(transformed_data, mean_values, std_values)\n",
    "\n",
    "## apply mask also to labels\n",
    "transformed_labels = labels[mask]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "88V4XWXrxIsr"
   },
   "source": [
    "Now we are ready to train the model again. Here we just train for 10 epochs (maybe even decrease it to 5 if you couldn't get a GPU). The resulting density estimate might not be perfectly optimized, but it will illustrate its power on a HEP example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "dZTm4MOxxKxd",
    "outputId": "e5bd3c96-1a80-47fb-a36c-336238352500"
   },
   "outputs": [],
   "source": [
    "batch_size = 256\n",
    "model.fit(x=transformed_data,\n",
    "          y=np.zeros((transformed_data.shape[0], 0), dtype=np.float32),\n",
    "          batch_size= batch_size,\n",
    "          epochs=10,\n",
    "          steps_per_epoch=transformed_data.shape[0] // batch_size,\n",
    "          verbose=1,\n",
    "          shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_2UeTA_ExNK4"
   },
   "source": [
    "Then we sample artificial data points from this learned density. We produce the same number of datapoints that we have in our training data. Then we have to perform all the variable transformations backwards."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "rRwgPbxmxOp-",
    "outputId": "d71e3f12-0f7a-447e-cb6d-0d798f999880"
   },
   "outputs": [],
   "source": [
    "## sample artificial data points\n",
    "samples = distribution.sample(data.shape[0])\n",
    "\n",
    "## inverse standardize\n",
    "retransformed_samples = standardize_inverse(samples, mean_values, std_values)\n",
    "\n",
    "## inverse logit\n",
    "retransformed_samples = logit_inverse(retransformed_samples)\n",
    "\n",
    "## inverse normalize\n",
    "retransformed_samples = normalize_inverse(retransformed_samples, max_values, min_values)\n",
    "\n",
    "print(\"sampled data shape =\", retransformed_samples.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SrwBpUoixRvf"
   },
   "source": [
    "Let's plot the original data again and in addition the sampled data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 591
    },
    "id": "CV5cCEkKxUxu",
    "outputId": "e47854e5-9da3-4e73-a209-b19457f23bd4"
   },
   "outputs": [],
   "source": [
    "## plot the sampled against the original data\n",
    "fig, ax = plt.subplots(2,2, figsize=(12,8))\n",
    "_, common_bins, _ = ax[0,0].hist(data[:,0], 100, alpha=0.6, label=\"data\")\n",
    "ax[0,0].hist(retransformed_samples[:,0], common_bins, label=\"sampled\", edgecolor=\"orange\", linewidth=1.3, histtype=\"step\")\n",
    "ax[0,0].set_xlabel(\"lower jet mass (TeV)\")\n",
    "ax[0,0].set_ylabel(\"events\")\n",
    "ax[0,0].set_yscale(\"log\")\n",
    "_, common_bins, _ = ax[0,1].hist(data[:,1], 100, alpha=0.6, label=\"data\")\n",
    "ax[0,1].hist(retransformed_samples[:,1], common_bins, label=\"sampled\", edgecolor=\"orange\", linewidth=1.3, histtype=\"step\")\n",
    "ax[0,1].set_xlabel(\"jet mass difference (TeV)\")\n",
    "ax[0,1].set_ylabel(\"events\")\n",
    "ax[0,1].set_yscale(\"log\")\n",
    "_, common_bins, _ = ax[1,0].hist(data[:,2], 100, alpha=0.6, label=\"data\")\n",
    "ax[1,0].hist(retransformed_samples[:,2], common_bins, label=\"sampled\", edgecolor=\"orange\", linewidth=1.3, histtype=\"step\")\n",
    "ax[1,0].set_xlabel(r\"$\\tau_{21,1}$\")\n",
    "ax[1,0].set_ylabel(\"events\")\n",
    "_, common_bins, _ = ax[1,1].hist(data[:,3], 100, alpha=0.6, label=\"data\")\n",
    "ax[1,1].hist(retransformed_samples[:,3], common_bins, label=\"sampled\", edgecolor=\"orange\", linewidth=1.3, histtype=\"step\")\n",
    "ax[1,1].set_xlabel(r\"$\\tau_{21,2}$\")\n",
    "ax[1,1].set_ylabel(\"events\")\n",
    "leg_handles, leg_labels = ax[0,0].get_legend_handles_labels()\n",
    "fig.legend(leg_handles, leg_labels, loc='upper right')\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kjU5HJhrxXH1"
   },
   "source": [
    "Even though we haven't fully trained the model, there is a striking agreement already between the original and the sampled distribution. We want to make sure now that the model has also learned the correct correlations between the four features. Let's compute the Pearson correlation coefficients and compare the two correlation matrices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 313
    },
    "id": "Q8WPgelExadO",
    "outputId": "4d0fffc9-5170-4d66-f62c-7dbf537b0dd3"
   },
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "\n",
    "variable_labels = [r'$m_{j1}$', r'$m_{j2}-m_{j1}$', r'$\\tau_{21,1}$', r'$\\tau_{21,2}$']\n",
    "corr_matrix_data = np.corrcoef(np.transpose(data))\n",
    "\n",
    "corr_matrix_data = pd.DataFrame(np.corrcoef(np.transpose(data)), columns=variable_labels, index=variable_labels)\n",
    "corr_matrix_sample = pd.DataFrame(np.corrcoef(np.transpose(retransformed_samples)), columns=variable_labels, index=variable_labels)\n",
    "\n",
    "plt.figure(figsize=(12,5))\n",
    "plt.subplot(1,2,1)\n",
    "sns.heatmap(corr_matrix_data, vmax=1., vmin=-1., cmap='coolwarm', annot=True, square=True)\n",
    "plt.title(\"data correlations\")\n",
    "\n",
    "plt.subplot(1,2,2)\n",
    "sns.heatmap(corr_matrix_sample, vmax=1., vmin=-1., cmap='coolwarm', annot=True, square=True)\n",
    "plt.title(\"sample correlations\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2q2vQplKxczH"
   },
   "source": [
    "While the comparison might still yield some room for improvement, it's clear that the normalizing flow has indeed learned the correlations between the features, rather than just sampling four independent distributions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t1Q4bKNt_SY_"
   },
   "source": [
    "**Task: Use your knowledge from the previous section to train a conditional flow that can separately generate signal and background events. Sample from these separately and plot them in comparison to the actual data as before. This works best if you sample the same number of signal and background events as present in the data respectively.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "uuXd9LrEbk5N",
    "outputId": "e474577a-eea2-4ca3-fe1b-f5b632fa1a4d",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#Add code here\n",
    "######################\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hFe2BWVBTAPn"
   },
   "source": [
    "## Advanced Task"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lKF0ZBzvQnJW"
   },
   "source": [
    "If the exercises before were too easy for you or if you would like to have an additional challenge in your own time, consider the following task. It's not rocket science but it's very likely to take more time than what we have in this exercise class.\n",
    "\n",
    "So far, we have for the sake of simplicity only used a binary conditional variable. However, this would work just as well on a general continuous conditional. If you extracted also the dijet mass (mjj) from the dataset and plotted it. You would see that the signal peaks quite sharply at the resonance mass of 3.5 TeV. We can divide the mjj spectrum into a signal region (SR), ranging from 3.3 to 3.7 TeV, and the remaining sideband region (SB). Drawing events from the SB, you should only see very background-like distributions in the other variables, whereas the SR events should yield a very distinct signal peak in addition to the background. \n",
    "\n",
    "**Task: Extract and plot mjj. Train a conditional flow model with mjj as the conditional variable, then sample data from the two regions separately in the same ratio as found in data. Plot the variables and compare with the original data. Since this task is also more difficult for the density estimator, you might need to increase the model complexity, the number of training epochs and/or the fraction of used data to achieve a satisfying result.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-ade-E5inizh"
   },
   "outputs": [],
   "source": [
    "#Add code here\n",
    "######################\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "flow_tutorial.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python (flow_env)",
   "language": "python",
   "name": "flow_env"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}