cf_notebooks/2.PartII-GenerativeModels.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/EiffL/Quarks2CosmosDataChallenge/blob/colab/notebooks/PartII-GenerativeModels.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1ZvzzJJPpYQb"
      },
      "source": [
        "# Guided Data Challenge Part II: Generative Modeling\n",
        "\n",
        "Author:\n",
        " - [@EiffL](https://github.com/EiffL) (Fancois Lanusse)\n",
        "\n",
        "## Overview\n",
        "\n",
        "In this notebook, we will go over several types of generative models which can be used to learn priors in order to solve the inverse problems illustrated in the previous notebook. Concretely, we will see how to fit a model $p_{\\theta}(x)$ to existing data.\n",
        "\n",
        "\n",
        "### Learning objectives:\n",
        "\n",
        "In this notebook we will put into practice:\n",
        "  - Write a Neural Network with DeepMind's Haiku\n",
        "  - Build a Variational Auto-Encoder\n",
        "  - Build a latent Normalizing Flow\n",
        "  - (Bonus) Build a score matching network!\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Installing dependencies and accessing data"
      ],
      "metadata": {
        "id": "0ktLeF7bjnbw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install git+https://github.com/EiffL/Quarks2CosmosDataChallenge.git\n",
        "!echo \"deb http://packages.cloud.google.com/apt gcsfuse-bionic main\" > /etc/apt/sources.list.d/gcsfuse.list\n",
        "!curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -\n",
        "!apt -qq update\n",
        "!apt -qq install gcsfuse\n",
        "!mkdir galsim\n",
        "\n",
        "import logging\n",
        "logger = logging.getLogger()\n",
        "class CheckTypesFilter(logging.Filter):\n",
        "    def filter(self, record):\n",
        "        return \"check_types\" not in record.getMessage()\n",
        "logger.addFilter(CheckTypesFilter())"
      ],
      "metadata": {
        "id": "p_56Uqv6h0QH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Authenticating and mounting cloud data storage\n",
        "from google.colab import auth\n",
        "auth.authenticate_user()\n",
        "!gcsfuse --implicit-dirs galsim galsim"
      ],
      "metadata": {
        "id": "rypc1fA8iK9p"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ywlFYhf3pYQd"
      },
      "outputs": [],
      "source": [
        "%pylab inline\n",
        "import jax\n",
        "import jax.numpy as jnp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3myZJSjKpYQe"
      },
      "source": [
        "## Step I: Preparing training data\n",
        "\n",
        "\n",
        "Here we will be using the space-based data and try to learn a generative model from it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tw0cFmippYQe"
      },
      "outputs": [],
      "source": [
        "import quarks2cosmos\n",
        "import tensorflow_datasets as tfds\n",
        "import tensorflow as tf\n",
        "\n",
        "def augment_data(example):\n",
        "    \"\"\" This applies random flipping to the input image.\n",
        "    \"\"\"\n",
        "    x = example['image']\n",
        "    x = x[...,tf.newaxis]\n",
        "    x = tf.image.random_flip_left_right(x)\n",
        "    x = tf.image.random_flip_up_down(x)\n",
        "    x = x[..., 0]\n",
        "    return x\n",
        "\n",
        "dset = tfds.load('Cosmos/23.5', split=tfds.Split.TRAIN,\n",
        "                 data_dir='galsim/tensorflow_datasets')\n",
        "dset = dset.cache()\n",
        "dset = dset.repeat()\n",
        "dset = dset.shuffle(50000)\n",
        "dset = dset.batch(64)\n",
        "dset = dset.map(augment_data)\n",
        "dset = iter(tfds.as_numpy(dset))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzYIhoxUpYQe"
      },
      "outputs": [],
      "source": [
        "batch = next(dset)\n",
        "\n",
        "figure(figsize=[10,10])\n",
        "for i in range(16):\n",
        "    subplot(4,4,i+1)\n",
        "    imshow(batch[i])\n",
        "    axis('off')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0AP0ZnuSpYQe"
      },
      "source": [
        "## Step II: Building a Variational Auto-Encoder\n",
        "\n",
        "We will start by a conventional Variational Auto-Encoder, this model here is not very complicated, probably the most simple one you could think of.\n",
        "\n",
        "It is an opportunity to see how to build a neural network in Haiku, and feel free to try to improve the model :-)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pP7Fdi4CpYQf"
      },
      "outputs": [],
      "source": [
        "# Let's start with the imports\n",
        "import haiku as hk     # NN library\n",
        "import optax           # Optimizer library\n",
        "import pickle\n",
        "\n",
        "# Utility function for tensoboard\n",
        "from flax.metrics import tensorboard\n",
        "\n",
        "# TensorFlow probability\n",
        "from tensorflow_probability.substrates import jax as tfp\n",
        "tfd = tfp.distributions\n",
        "tfb = tfp.bijectors"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OxSalG3ppYQf"
      },
      "outputs": [],
      "source": [
        "# Create a random sequence\n",
        "rng_seq = hk.PRNGSequence(42)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vTYd1Rd2pYQf"
      },
      "source": [
        "#### Building an encoder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SsLedM4cpYQf"
      },
      "outputs": [],
      "source": [
        "class Encoder(hk.Module):\n",
        "  \"\"\"Simple Convolutional encoder model.\"\"\"\n",
        "  def __init__(self, latent_size=64):\n",
        "    super().__init__()\n",
        "    self._latent_size = latent_size\n",
        "\n",
        "  def __call__(self, x):\n",
        "    # Add channel dimension to x\n",
        "    x = x[..., jnp.newaxis]\n",
        "\n",
        "    x = hk.Conv2D(16, kernel_shape=3)(x)\n",
        "    x = jax.nn.leaky_relu(x)\n",
        "    x = hk.avg_pool(x, window_shape=3, strides=2, padding='SAME')\n",
        "\n",
        "    x = hk.Conv2D(32, kernel_shape=3)(x)\n",
        "    x = jax.nn.leaky_relu(x)\n",
        "    x = hk.avg_pool(x,  window_shape=3, strides=2, padding='SAME')\n",
        "\n",
        "    x = hk.Conv2D(64, kernel_shape=3)(x)\n",
        "    x = jax.nn.leaky_relu(x)\n",
        "    x = hk.avg_pool(x, window_shape=3, strides=2, padding='SAME')\n",
        "\n",
        "    x = hk.Conv2D(128, kernel_shape=3)(x)\n",
        "    x = jax.nn.leaky_relu(x)\n",
        "    x = hk.avg_pool(x, window_shape=3, strides=2, padding='SAME')\n",
        "\n",
        "    x = hk.Conv2D(128, kernel_shape=3)(x)\n",
        "    x = jax.nn.leaky_relu(x)\n",
        "    x = hk.avg_pool(x, window_shape=3, strides=2, padding='SAME')\n",
        "\n",
        "    x = hk.Flatten()(x)\n",
        "\n",
        "    x = hk.Linear(256)(x)\n",
        "    x = jax.nn.leaky_relu(x)\n",
        "\n",
        "    # Returns the variational distribution encoding the input image\n",
        "    loc = hk.Linear(self._latent_size)(x)\n",
        "\n",
        "    scale = jax.nn.softplus(hk.Linear(self._latent_size)(x)) + 1e-3\n",
        "    return tfd.MultivariateNormalDiag(loc, scale)                    # Note that this returns a distribution !"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-mfBC5VKpYQg"
      },
      "outputs": [],
      "source": [
        "# Transform the model into pure functions\n",
        "encoder = hk.without_apply_rng(hk.transform(lambda x : Encoder()(x)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ueR-3W0dpYQg"
      },
      "source": [
        "What `hk.transform` returns is a structure that contains two functions:\n",
        "  - `model.init`: A function that initializes the parameters\n",
        "  - `model.apply`: A function that applies the neural network\n",
        "\n",
        "Let's start by creating params:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YsxUgPN3pYQg"
      },
      "outputs": [],
      "source": [
        "params_enc = encoder.init(next(rng_seq), jnp.zeros([1,101,101]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n3vyVQnSpYQg"
      },
      "source": [
        "And then you apply the NN like this:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_318eEZOpYQg"
      },
      "outputs": [],
      "source": [
        "z = encoder.apply(params_enc, jnp.zeros([1,101,101]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WPUA8HLvpYQg"
      },
      "outputs": [],
      "source": [
        "z"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R8ff2k-opYQg"
      },
      "source": [
        "and what is returned is an actual distribution object :-) So you can sample from it and evaluate probability."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QhnBmrXWpYQg"
      },
      "source": [
        "#### Building a decoder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tyvE1vAypYQg"
      },
      "outputs": [],
      "source": [
        "class Decoder(hk.Module):\n",
        "  \"\"\"Simple Convolutional decoder model.\"\"\"\n",
        "  def __call__(self, z, scale=1.0):\n",
        "\n",
        "    # Reshape latent variable to an image\n",
        "    x = hk.Linear(256)(z)\n",
        "    x = jax.nn.leaky_relu(x)\n",
        "\n",
        "    x = hk.Linear(3*3*128)(x)\n",
        "    x = x.reshape([-1,3,3,128])\n",
        "\n",
        "    x = hk.Conv2DTranspose(128, kernel_shape=3, stride=2)(x)\n",
        "    x = jax.nn.leaky_relu(x)\n",
        "\n",
        "    x = hk.Conv2DTranspose(64, kernel_shape=3, stride=2)(x)\n",
        "    x = jax.nn.leaky_relu(x)\n",
        "\n",
        "    x = hk.Conv2DTranspose(32, kernel_shape=3, stride=2)(x)\n",
        "    x = jax.nn.leaky_relu(x)\n",
        "\n",
        "    x = hk.Conv2DTranspose(16, kernel_shape=3, stride=2)(x)\n",
        "    x = jax.nn.leaky_relu(x)\n",
        "\n",
        "    x = hk.Conv2DTranspose(8, kernel_shape=3, stride=2)(x)\n",
        "\n",
        "    x = hk.Conv2D(1, kernel_shape=5)(x)\n",
        "\n",
        "    x = x[...,0]\n",
        "    x = jnp.pad(x, [[0,0],[3,2],[3,2]])  # This step is to pad the image for the 101x101 expected size\n",
        "\n",
        "    return tfd.Independent(tfd.Normal(loc=x, scale=scale), reinterpreted_batch_ndims=2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JfvMAqRTpYQg"
      },
      "outputs": [],
      "source": [
        "# We turn the model into pure functions\n",
        "decoder = hk.without_apply_rng(hk.transform(lambda z : Decoder()(z)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Mi_JQoWnpYQg"
      },
      "outputs": [],
      "source": [
        "# Create parameters\n",
        "params_dec = decoder.init(next(rng_seq), jnp.zeros([1,64]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xH40qrQApYQg"
      },
      "outputs": [],
      "source": [
        "decoder.apply(params_dec, jnp.zeros([1,64]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FpThMCthpYQh"
      },
      "source": [
        "#### Building the auto-encoder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9SlVxXDSpYQh"
      },
      "outputs": [],
      "source": [
        "# Combine both set of parameters\n",
        "params = hk.data_structures.merge(params_enc, params_dec)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yVSpTE1kpYQh"
      },
      "outputs": [],
      "source": [
        "@jax.jit\n",
        "def loss_fn(params, rng_key, batch):\n",
        "    x = batch\n",
        "\n",
        "    # Autoencode an example\n",
        "    q = encoder.apply(params, x)\n",
        "\n",
        "    # Sample from the posterior\n",
        "    z = q.sample(seed=rng_key)\n",
        "\n",
        "    # Decode the sample\n",
        "    p = decoder.apply(params, z)\n",
        "\n",
        "    # Compute loss\n",
        "    kl = tfd.kl_divergence(q, tfd.MultivariateNormalDiag(jnp.zeros(64),\n",
        "                                                         scale_identity_multiplier=1.))\n",
        "    log_likelihood = p.log_prob(x)\n",
        "\n",
        "    elbo = log_likelihood - 0.0001*kl      # Here we apply a factor on the KL term, you can try to change it!\n",
        "    return -jnp.mean(elbo)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ddnZLJjNpYQh"
      },
      "outputs": [],
      "source": [
        "def lr_schedule(step):\n",
        "  \"\"\"Linear scaling rule optimized for 90 epochs.\"\"\"\n",
        "  steps_per_epoch = 40000 // 64\n",
        "\n",
        "  current_epoch = step / steps_per_epoch  # type: float\n",
        "  boundaries = jnp.array((20, 40, 60)) * steps_per_epoch\n",
        "  values = jnp.array([1., 0.1, 0.01, 0.001])\n",
        "\n",
        "  index = jnp.sum(boundaries < step)\n",
        "  return jnp.take(values, index)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_Oq1Zyn4pYQh"
      },
      "outputs": [],
      "source": [
        "optimizer = optax.chain(\n",
        "      optax.adam(1e-3),\n",
        "      optax.scale_by_schedule(lr_schedule))\n",
        "\n",
        "opt_state = optimizer.init(params)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vHFPIMEgpYQh"
      },
      "outputs": [],
      "source": [
        "@jax.jit\n",
        "def update(params, rng_key, opt_state, batch):\n",
        "    \"\"\"Single SGD update step.\"\"\"\n",
        "    loss, grads = jax.value_and_grad(loss_fn)(params, rng_key, batch)\n",
        "    updates, new_opt_state = optimizer.update(grads, opt_state)\n",
        "    new_params = optax.apply_updates(params, updates)\n",
        "    return loss, new_params, new_opt_state"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vESc62YLpYQh"
      },
      "outputs": [],
      "source": [
        "# To keep a log of training loss\n",
        "summary_writer = tensorboard.SummaryWriter(\"models/vae\")\n",
        "step = 0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_OlED9Y2pYQh"
      },
      "outputs": [],
      "source": [
        "# Before starting the training, we can start tensorboard to check on the training [only works with colab pro]\n",
        "# %load_ext tensorboard\n",
        "# %tensorboard --logdir models"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "86Ua7Z2lpYQh"
      },
      "outputs": [],
      "source": [
        "for s in range(50000):\n",
        "    loss, params, opt_state = update(params, next(rng_seq), opt_state, next(dset))\n",
        "    step+=1\n",
        "\n",
        "    summary_writer.scalar('train_loss', loss, step)\n",
        "    summary_writer.scalar('learning_rate', lr_schedule(step)*1e-3, step)\n",
        "\n",
        "    if step%5000 ==0:\n",
        "        with open('models/vae/model-%d.pckl'%step, 'wb') as file:\n",
        "            pickle.dump([params, opt_state], file)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "teYyou1vpYQh"
      },
      "outputs": [],
      "source": [
        "# If I wanted to reload a specific outputed set of params I can use the following:\n",
        "# with open('models/vae/model-5000.pckl', 'rb') as file:\n",
        "#     params, opt_state = pickle.load(, file)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gc6sVVFRpYQh"
      },
      "source": [
        "Now let's quickly test the auto-encoder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6GPLeLU7pYQh"
      },
      "outputs": [],
      "source": [
        "# Retrieving a batch of data\n",
        "batch = next(dset)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mqtUUF0DpYQh"
      },
      "outputs": [],
      "source": [
        "# Use the encoder to get the variational posterior distribution\n",
        "latent_distribution = encoder.apply(params, batch)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mDftKIsVpYQh"
      },
      "outputs": [],
      "source": [
        "# Sample some variables from the distribution\n",
        "z = latent_distribution.sample(seed=next(rng_seq))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NYP5DGtbpYQh"
      },
      "outputs": [],
      "source": [
        "scatter(z[:,0], z[:,1])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K8lETq4tpYQh"
      },
      "outputs": [],
      "source": [
        "# And we can now get the data likelihood given latent variables\n",
        "likelihood = decoder.apply(params, z)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7MH3iF73pYQi"
      },
      "outputs": [],
      "source": [
        "# We will look at the mean of the data likelihood\n",
        "reconstruction = likelihood.mean()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EcSdXpdqpYQp"
      },
      "outputs": [],
      "source": [
        "imshow(reconstruction[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "scrolled": false,
        "id": "AF4e-TmgpYQp"
      },
      "outputs": [],
      "source": [
        "# Let's draw from a few examples\n",
        "for i in range(4):\n",
        "    figure(figsize=[10,5])\n",
        "    subplot(131)\n",
        "    imshow(batch[i,:,:])\n",
        "    subplot(132)\n",
        "    imshow(im[i,:,:],vmin=0)\n",
        "    subplot(133)\n",
        "    imshow(im[i,:,:] - batch[i,:,:],vmin=-.03,vmax=0.03)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bFMc6RBjpYQp"
      },
      "source": [
        "Now, it's your turn! Try to get the best possible (Variational)-Auto-Encoder."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-dcSeI8ZpYQp"
      },
      "source": [
        "But first let's check something, we see that the autoencoding is not perfect but can we at least sample from it!?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SCA1dg8CpYQq"
      },
      "outputs": [],
      "source": [
        "latent_prior = tfd.MultivariateNormalDiag(jnp.zeros(64))\n",
        "z_prior = latent_prior.sample(64, seed=next(rng_seq))\n",
        "\n",
        "likelihood = decoder.apply(params, z_prior)\n",
        "rec = likelihood.mean()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "q4m6jXIqpYQq"
      },
      "outputs": [],
      "source": [
        "figure(figsize=(10,10))\n",
        "for i in range(4):\n",
        "    for j in range(4):\n",
        "        subplot(4,4,i+4*j+1)\n",
        "        imshow(rec[i+4*j],cmap='gray')\n",
        "        axis('off')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gztpTWQqpYQq"
      },
      "source": [
        "## Step III: Latent Normalizing Flow\n",
        "\n",
        "Even with a very good VAE, the latent space of the VAE is rarely very Gaussian, which means that it can be hard to sample realistic images from it. In this part of the notebook, we'll see how you can create a latent model using a Normalizing Flow, which will learn how to sample from the effective latent distribution.\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bpDbnj80pYQq"
      },
      "outputs": [],
      "source": [
        "# Let's store the dimensionality of our latent space\n",
        "d = 64"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kpg_FQNTpYQq"
      },
      "source": [
        "The first part of a Normalizing Flow is to define some coupling layers, i.e. the layers of the Flow."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WQKYOQnWpYQq"
      },
      "outputs": [],
      "source": [
        "class AffineCoupling(hk.Module):\n",
        "  \"\"\"This is the coupling layer used in the Flow.\"\"\"\n",
        "  def __init__(self, scale_only=True, **kwargs):\n",
        "    super().__init__(**kwargs)\n",
        "    self.scale_only = scale_only\n",
        "\n",
        "  def __call__(self, x, output_units, **condition_kwargs):\n",
        "    net = hk.Linear(128)(x)\n",
        "    net = jax.nn.leaky_relu(net)\n",
        "    net = hk.Linear(128)(net)\n",
        "    net = jax.nn.leaky_relu(net)\n",
        "    shifter = tfb.Shift(hk.Linear(output_units)(net))\n",
        "    if self.scale_only:\n",
        "      return shifter\n",
        "    else:\n",
        "      scaler = tfb.Scale(jnp.clip(jax.nn.softplus(hk.Linear(output_units)(net)), 1e-2, 1e1))\n",
        "      return tfb.Chain([shifter, scaler])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sG8hOBOxpYQq"
      },
      "source": [
        "With these layers we can then compose the flow:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iuQkJzuKpYQq"
      },
      "outputs": [],
      "source": [
        "class AffineFlow(hk.Module):\n",
        "    \"\"\"This is a normalizing flow using the coupling layers defined\n",
        "    above.\"\"\"\n",
        "    def __call__(self):\n",
        "        chain = tfb.Chain([\n",
        "            tfb.RealNVP(d//2, bijector_fn=AffineCoupling(name='aff1')),\n",
        "            tfb.Permute(np.arange(d)[::-1]),\n",
        "            tfb.RealNVP(d//2, bijector_fn=AffineCoupling(name='aff2')),\n",
        "            tfb.Permute(np.arange(d)[::-1]),\n",
        "            tfb.RealNVP(d//2, bijector_fn=AffineCoupling(name='aff3',\n",
        "                                                         scale_only=False)),\n",
        "            tfb.Permute(np.arange(d)[::-1]),\n",
        "            tfb.RealNVP(d//2, bijector_fn=AffineCoupling(name='aff4',\n",
        "                                                         scale_only=False)),\n",
        "            tfb.Permute(np.arange(d)[::-1]),\n",
        "        ])\n",
        "\n",
        "        nvp = tfd.TransformedDistribution(\n",
        "            tfd.MultivariateNormalDiag(jnp.zeros(d),jnp.ones(d)),\n",
        "            bijector=chain)\n",
        "        return nvp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "64OS35xJpYQq"
      },
      "source": [
        "We can now define two types of functions, one that will evaluate the log likelihood, and one that will sample from the flow."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MTBbfm6lpYQq"
      },
      "outputs": [],
      "source": [
        "model_logp = hk.without_apply_rng(hk.transform(lambda x : AffineFlow()().log_prob(x)))\n",
        "model_sample = hk.transform(lambda n : AffineFlow()().sample(n, seed=hk.next_rng_key()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y6QIx4OlpYQq"
      },
      "outputs": [],
      "source": [
        "params_flow = model_logp.init(next(rng_seq), jnp.zeros([1, d]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RZv9KwDepYQq"
      },
      "outputs": [],
      "source": [
        "optimizer = optax.adam(1e-4)\n",
        "opt_state = optimizer.init(params_flow)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kJyF_BFxpYQq"
      },
      "outputs": [],
      "source": [
        "@jax.jit\n",
        "def loss_fn(params_flow, rng_key, batch):\n",
        "    x = batch\n",
        "\n",
        "    # Encode example using the trained encoder\n",
        "    # params here are the params of the VAE as trained in the\n",
        "    q = encoder.apply(params, x)\n",
        "\n",
        "    z = q.sample(seed=rng_key)\n",
        "\n",
        "    logp = model_logp.apply(params_flow, z)\n",
        "\n",
        "    return -jnp.mean(logp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tIGyVzHkpYQr"
      },
      "outputs": [],
      "source": [
        "@jax.jit\n",
        "def update(params, rng_key, opt_state, batch):\n",
        "    \"\"\"Single SGD update step.\"\"\"\n",
        "    loss, grads = jax.value_and_grad(loss_fn)(params, rng_key, batch)\n",
        "    updates, new_opt_state = optimizer.update(grads, opt_state)\n",
        "    new_params = optax.apply_updates(params, updates)\n",
        "    return loss, new_params, new_opt_state"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R8hdkjGSpYQr"
      },
      "outputs": [],
      "source": [
        "summary_writer = tensorboard.SummaryWriter(\"models/flow\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i5DVRhFPpYQr"
      },
      "outputs": [],
      "source": [
        "step = 0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yHG4w2AXpYQr"
      },
      "outputs": [],
      "source": [
        "for s in range(10000):\n",
        "    loss, params_flow, opt_state = update(params_flow, next(rng_seq), opt_state, next(dset))\n",
        "    step+=1\n",
        "    if step % 100 == 0:\n",
        "        summary_writer.scalar('train_loss', loss, step)\n",
        "        print(step, loss)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1X9vRupEpYQr"
      },
      "outputs": [],
      "source": [
        "# Once trained, we can apply the flow and check how we're doing\n",
        "batch = next(dset)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7rUc1mKipYQr"
      },
      "outputs": [],
      "source": [
        "latent_distribution = encoder.apply(params, batch)\n",
        "z = latent_distribution.sample(seed=next(rng_seq))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-qcmommWpYQr"
      },
      "outputs": [],
      "source": [
        "# And also sample from the Flow\n",
        "z2 = model_sample.apply(params_flow, next(rng_seq), n=256)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RlzubTPNpYQr"
      },
      "outputs": [],
      "source": [
        "scatter(z2[:,0],  z2[:,1], label='Flow latent')\n",
        "scatter(z[:,0],  z[:,1], label='VAE latent')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dxo5x6p8pYQr"
      },
      "source": [
        "Here we see that the 2 distributions don't match very well, more training is necessary"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p0lBsivEpYQr"
      },
      "source": [
        "#### Combining the VAE with the Flow\n",
        "\n",
        "We can now combine both models, the flow allowing us to sample from the latent space of the VAE, and then we decode these samples with the decoder of the VAE:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TxMk4Q7YpYQr"
      },
      "outputs": [],
      "source": [
        "# Sample latent space values with the flow\n",
        "z = model_sample.apply(params_flow, next(rng_seq), n=256)\n",
        "# Decode them with the Decoder\n",
        "likelihood = decoder.apply(params, z)\n",
        "rec = likelihood.mean()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UvzQIzwcpYQr"
      },
      "outputs": [],
      "source": [
        "figure(figsize=(10,10))\n",
        "for i in range(4):\n",
        "    for j in range(4):\n",
        "        subplot(4,4,i+4*j+1)\n",
        "        imshow(rec[i+4*j],cmap='gray')\n",
        "        axis('off')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gi0L_kknpYQr"
      },
      "source": [
        "## Bonus:  Denoising Score Matching\n",
        "\n",
        "In this section, we will now explore how to use Denoising Score Matching to learn a Generative Model.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jwvXTAZypYQr"
      },
      "outputs": [],
      "source": [
        "# Training set preparation\n",
        "def load_dataset(batch_size, noise_dist_std, train_split):\n",
        "\n",
        "  def pre_process(im):\n",
        "    \"\"\" Pre-processing function preparing data for denoising task\n",
        "    \"\"\"\n",
        "    # Cutout a portion of the map\n",
        "    x = tf.image.resize_with_crop_or_pad(im['image'][...,tf.newaxis], 128,128)\n",
        "    x = tf.image.random_flip_left_right(x)\n",
        "    x = tf.image.random_flip_up_down(x)\n",
        "    # Sample random Gaussian noise\n",
        "    u = tf.random.normal(tf.shape(x))\n",
        "    # Sample standard deviation of noise corruption\n",
        "    s = noise_dist_std * tf.random.normal((1, 1, 1))\n",
        "    # Create noisy image\n",
        "    y = x + s * u\n",
        "    return {'x':x, 'y':y, 'u':u,'s':s}\n",
        "\n",
        "  ds = tfds.load('Cosmos/23.5', split=train_split, shuffle_files=True,\n",
        "                 data_dir='galsim/tensorflow_datasets')\n",
        "  ds = ds.shuffle(buffer_size=10*batch_size)\n",
        "  ds = ds.repeat()\n",
        "  ds = ds.map(pre_process)\n",
        "  ds = ds.batch(batch_size)\n",
        "  ds = ds.prefetch(tf.data.experimental.AUTOTUNE)\n",
        "  return ds.as_numpy_iterator()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hoyQbDjVpYQs"
      },
      "outputs": [],
      "source": [
        "cosmos_dset = load_dataset(32, 0.025, tfds.Split.TRAIN)\n",
        "batch = next(cosmos_dset)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4JqdG3DZpYQs"
      },
      "outputs": [],
      "source": [
        "for i in range(4):\n",
        "    figure()\n",
        "    subplot(131)\n",
        "    imshow(batch['x'][i])\n",
        "    axis('off')\n",
        "    subplot(132)\n",
        "    imshow(batch['u'][i])\n",
        "    axis('off')\n",
        "    subplot(133)\n",
        "    imshow(batch['y'][i])\n",
        "    axis('off')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tgE0tOXFpYQs"
      },
      "outputs": [],
      "source": [
        "# Ok, so now we just need to train\n",
        "from quarks2cosmos.models import SmallUResNet, SNParamsTree\n",
        "import haiku as hk\n",
        "import optax"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UMKLNRGIpYQs"
      },
      "outputs": [],
      "source": [
        "model = hk.transform_with_state(lambda x, sigma, is_training=False: SmallUResNet()(x, sigma, is_training))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vW3qLqQLpYQs"
      },
      "outputs": [],
      "source": [
        "params, state = model.init(jax.random.PRNGKey(0), jnp.zeros([1,128,128,1]),jnp.zeros([1,1,1,1]), is_training=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aHYlubJUpYQs"
      },
      "outputs": [],
      "source": [
        "def lr_schedule(step):\n",
        "  \"\"\"Linear scaling rule optimized for 90 epochs.\"\"\"\n",
        "  steps_per_epoch = 40000 // 32\n",
        "\n",
        "  current_epoch = step / steps_per_epoch  # type: float\n",
        "  lr = (1.0 * 32) / 32\n",
        "  boundaries = jnp.array((20, 40, 60)) * steps_per_epoch\n",
        "  values = jnp.array([1., 0.1, 0.01, 0.001]) * lr\n",
        "\n",
        "  index = jnp.sum(boundaries < step)\n",
        "  return jnp.take(values, index)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HUvpjuI5pYQs"
      },
      "outputs": [],
      "source": [
        "optimizer = optax.chain(\n",
        "  optax.adam(learning_rate=1e-3),\n",
        "  optax.scale_by_schedule(lr_schedule)\n",
        ")\n",
        "rng_seq = hk.PRNGSequence(42)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5V41dF7TpYQs"
      },
      "outputs": [],
      "source": [
        "opt_state = optimizer.init(params)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "buKxjsYPpYQs"
      },
      "outputs": [],
      "source": [
        "sn_fn = hk.transform_with_state(lambda x: SNParamsTree(ignore_regex='[^?!.]*b$|[^?!.]*offset$',\n",
        "                                                          val=1.)(x))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1cpVfXtmpYQs"
      },
      "outputs": [],
      "source": [
        "_, sn_state = sn_fn.init(next(rng_seq), params)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QjfucgYTpYQs"
      },
      "outputs": [],
      "source": [
        "params, sn_state = sn_fn.apply(None, sn_state, None, params)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VZ96HgI1pYQs"
      },
      "outputs": [],
      "source": [
        "def loss_fn(params, state, rng_key, batch):\n",
        "    score, state = model.apply(params, state, rng_key, batch['y'], batch['s'], is_training=True)\n",
        "    loss = jnp.mean((batch['u'] + batch['s'] * (score))**2)\n",
        "    return loss, state"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "31FGcCzHpYQs"
      },
      "outputs": [],
      "source": [
        "@jax.jit\n",
        "def update(params, state, sn_state, rng_key, opt_state, batch):\n",
        "    (loss, state), grads = jax.value_and_grad(loss_fn, has_aux=True)(params, state, rng_key, batch)\n",
        "    updates, new_opt_state = optimizer.update(grads, opt_state)\n",
        "    new_params = optax.apply_updates(params, updates)\n",
        "    new_params, new_sn_state = sn_fn.apply(None, sn_state, None, new_params)\n",
        "    return loss, new_params, state, new_sn_state, new_opt_state"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vjft4Y0XpYQs"
      },
      "outputs": [],
      "source": [
        "from flax.metrics import tensorboard\n",
        "import pickle"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Yr48l5QdpYQt"
      },
      "outputs": [],
      "source": [
        "summary_writer = tensorboard.SummaryWriter('score_model_0.025')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P55FtpxrpYQt"
      },
      "outputs": [],
      "source": [
        "for step in range(100000):\n",
        "    loss, params, state, sn_state, opt_state = update(params, state, sn_state,\n",
        "                                                      next(rng_seq), opt_state,\n",
        "                                                      next(cosmos_dset))\n",
        "    summary_writer.scalar('train_loss', loss, step)\n",
        "    summary_writer.scalar('learning_rate', lr_schedule(step)*1e-3, step)\n",
        "\n",
        "    if step%5000 ==0:\n",
        "        with open('score_model_0.025/model-%d.pckl'%step, 'wb') as file:\n",
        "            pickle.dump([params, state, sn_state], file)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QvI9SRlwpYQt"
      },
      "outputs": [],
      "source": [
        "def score_fn(y, s):\n",
        "    score, _ = model.apply(params, state, None, y, s.reshape((-1,1,1,1)), is_training=False)\n",
        "    return score"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g6m9dmYBpYQt"
      },
      "outputs": [],
      "source": [
        "score = score_fn(batch['y'], batch['s'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "scrolled": false,
        "id": "Rl8k2SGTpYQt"
      },
      "outputs": [],
      "source": [
        "for i in range(16):\n",
        "    figure(figsize=[20,5])\n",
        "    subplot(141)\n",
        "    imshow(batch['x'][i])\n",
        "    axis('off')\n",
        "    subplot(142)\n",
        "    imshow(batch['y'][i])\n",
        "    axis('off')\n",
        "    subplot(143)\n",
        "    imshow(score[i])\n",
        "    axis('off')\n",
        "    subplot(144)\n",
        "    imshow(batch['y'][i] + (batch['s']**2 * score)[i])\n",
        "    axis('off')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Et-C1h9epYQt"
      },
      "outputs": [],
      "source": [
        "# with open('score_model_0.05/model-%d.pckl'%step, 'rb') as file:\n",
        "#     params, state, sn_state = pickle.load(, file)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ztfQV4kWpYQt"
      },
      "outputs": [],
      "source": [
        "from scipy import integrate\n",
        "\n",
        "@jax.jit\n",
        "def score_fn(y, s):\n",
        "    score, _ = model.apply(params, state, None, y, s.reshape((-1,1,1,1)), is_training=False)\n",
        "    return score\n",
        "\n",
        "@jax.jit\n",
        "def dynamics(t, x):\n",
        "  x = x.reshape([-1,128,128,1])\n",
        "  return - 0.5*score_fn(x, s=jnp.ones([1])*jnp.sqrt(t)).reshape([-1])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6QYqDGqIpYQt"
      },
      "outputs": [],
      "source": [
        "subplot(121)\n",
        "imshow(batch['x'][0] )\n",
        "subplot(122)\n",
        "imshow(batch['y'][0] )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nr_ep4QwpYQt"
      },
      "outputs": [],
      "source": [
        "noise = 0.2\n",
        "init_image = batch['y'][0] + noise*randn(1,128,128,1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7ajxHBxxpYQt"
      },
      "outputs": [],
      "source": [
        "start_and_end_times = jnp.logspace(log10(0.99*noise**2),-7)\n",
        "\n",
        "solution = integrate.solve_ivp(dynamics,\n",
        "                               [noise**2,(0.0)],\n",
        "                               init_image.flatten(),\n",
        "                               t_eval=start_and_end_times)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3fGzuV4_pYQt"
      },
      "outputs": [],
      "source": [
        "imshow(solution.y[:,-1].reshape(128,128))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yq9vmTyQpYQt"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 - AI",
      "language": "python",
      "name": "python3-ai"
    },
    "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.7.9"
    },
    "colab": {
      "provenance": [],
      "include_colab_link": true
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}