cf_notebooks/3.PartIII-VariationalInference.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/PartIII-VariationalInference.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hq9pSZfBrT1x"
      },
      "source": [
        "# Guided Data Challenge Part III: Variational Posterior Inference\n",
        "\n",
        "Author:\n",
        " - [@EiffL](https://github.com/EiffL) (Fancois Lanusse)\n",
        "\n",
        "## Overview\n",
        "\n",
        "In this last notebook, we will use everything we have seen so far, and try to perform posterior using Variational Inference.\n",
        "\n",
        "\n",
        "### Learning objectives:\n",
        "\n",
        "In this notebook we will put into practice:\n",
        "  - Perform MAP inference\n",
        "  - Variational inference"
      ]
    },
    {
      "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": "CrCEhAmmrT1y"
      },
      "outputs": [],
      "source": [
        "%pylab inline\n",
        "import jax\n",
        "import jax.numpy as jnp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jpbKCpkkrT1z"
      },
      "source": [
        "## Step I: Load your generative model\n",
        "\n",
        "\n",
        "Here I'm going to load an existing pretrained model, you should feel free to replace this by a model you might have  trained yourself :-)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "joDnjqHmrT10"
      },
      "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\n",
        "\n",
        "# Specific models built by EiffL\n",
        "from quarks2cosmos.models.vae import Decoder\n",
        "from quarks2cosmos.models.flow import AffineFlow"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cKBNDxIPrT10"
      },
      "outputs": [],
      "source": [
        "# Create a random sequence\n",
        "rng_seq = hk.PRNGSequence(42)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IUuGfYTzrT10"
      },
      "outputs": [],
      "source": [
        "# Restore model parameters\n",
        "import pickle\n",
        "with open('galsim/model-50000.pckl', 'rb') as file:\n",
        "    params, state, _ = pickle.load(file)\n",
        "with open('galsim/model-20000.pckl', 'rb') as file:\n",
        "    params_flow, _ = pickle.load(file)\n",
        "\n",
        "params = hk.data_structures.merge(params, params_flow)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HHxgnAVsrT10"
      },
      "source": [
        "#### Create a forward model combining latent flow with VAE"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Uaz_xUR4rT11"
      },
      "outputs": [],
      "source": [
        "def generative_model_fn(z):\n",
        "    # Transform from Gaussian space to VAE latent space\n",
        "    z1 = AffineFlow()().bijector.forward(z)\n",
        "\n",
        "    # Decode sample with decoder\n",
        "    likelihood = Decoder()(z1, is_training=False)\n",
        "\n",
        "    return likelihood.mean()\n",
        "\n",
        "generative_model = hk.without_apply_rng(hk.transform_with_state(generative_model_fn))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5OdkYNj1rT11"
      },
      "outputs": [],
      "source": [
        "# To sample from the model, we draw from a Gaussian...\n",
        "z = tfd.MultivariateNormalDiag(jnp.zeros(32)).sample(16, seed=next(rng_seq))\n",
        "# And we run it through the forward model\n",
        "x, _ = generative_model.apply(params, state, z)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5hcbS8TIrT11"
      },
      "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(x[i+4*j],cmap='gray')\n",
        "        axis('off')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gi-wiGMzrT12"
      },
      "source": [
        "Not too bad :-)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fgfGjVkDrT12"
      },
      "source": [
        "## Step II: Back to our inverse problems"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E9hDrdLwrT12"
      },
      "outputs": [],
      "source": [
        "import quarks2cosmos.datasets\n",
        "import tensorflow_datasets as tfds\n",
        "from quarks2cosmos import galjax as gj"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zVDI_aJPrT12"
      },
      "outputs": [],
      "source": [
        "dset_cosmos = tfds.load(\"Cosmos/23.5\", split=tfds.Split.TRAIN,\n",
        "                        data_dir='galsim/tensorflow_datasets') # Load the TRAIN split\n",
        "dset_cosmos = dset_cosmos.as_numpy_iterator()                  # Convert the dataset to numpy iterator\n",
        "\n",
        "dset_hsc = tfds.load(\"HSC\", split=tfds.Split.TRAIN,\n",
        "                     data_dir='galsim/tensorflow_datasets')\n",
        "dset_hsc = dset_hsc.as_numpy_iterator()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-u-J_rjQrT12"
      },
      "outputs": [],
      "source": [
        "# Extract a new example from the dataset\n",
        "cosmos = next(dset_cosmos)\n",
        "\n",
        "figure(figsize=[10,5])\n",
        "subplot(121)\n",
        "imshow(cosmos['image'],cmap='gray')\n",
        "title('Galaxy')\n",
        "subplot(122)\n",
        "imshow(cosmos['psf'],cmap='gray')\n",
        "title('PSF');"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cxxPvLzsrT12"
      },
      "outputs": [],
      "source": [
        "# Extract a new example from the dataset\n",
        "hsc = next(dset_hsc)\n",
        "\n",
        "figure(figsize=[20,5])\n",
        "subplot(141)\n",
        "imshow(hsc['image'],cmap='gray')\n",
        "title('Galaxy')\n",
        "subplot(142)\n",
        "imshow(hsc['psf'],cmap='gray')\n",
        "title('PSF')\n",
        "subplot(143)\n",
        "imshow(hsc['mask'] == 44,cmap='gray')\n",
        "title('Interpolated pixels')\n",
        "subplot(144)\n",
        "imshow(hsc['variance'],cmap='gray')\n",
        "title('Variance plane');"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4t3A0b0rrT12"
      },
      "outputs": [],
      "source": [
        "def simulate_hsc(x, in_psf, out_psf):\n",
        "    \"\"\" This function will simulate an image at HSC resolution given an image at HST resolution,\n",
        "    accounting for input PSF and convolving by output PSF\n",
        "    Args:\n",
        "        x: HST resolution image (MUST BE ODD SIZE!!!!)\n",
        "        in_psf: HST PSF\n",
        "        out_psf: HSC PSF\n",
        "    Returns:\n",
        "        y: HSC simulated image of size [41,41]\n",
        "    \"\"\"\n",
        "    y = gj.deconvolve(x, in_psf)         # Deconvolve by input PSF\n",
        "    y = gj.kresample(y, 0.03, 0.168, 41) # Resample image to HSC grid\n",
        "    y = gj.convolve(y,  out_psf)         # Reconvolve by HSC PSF\n",
        "    return 2.587*y                       # Conversion factor for the flux"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8tl8G06OrT12"
      },
      "outputs": [],
      "source": [
        "likelihood = tfd.Independent(tfd.Normal(loc=simulate_hsc(cosmos['image'], cosmos['psf'], hsc['psf']),\n",
        "                                        scale=jnp.sqrt(hsc['variance'])),\n",
        "                             reinterpreted_batch_ndims=2) # This is to make sure TFP understand we have a 2d image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jn7wUzQfrT12"
      },
      "outputs": [],
      "source": [
        "im_noise = likelihood.sample(seed=jax.random.PRNGKey(1))\n",
        "x_true = cosmos['image']\n",
        "cr_mask = 1.*(hsc['mask'] == 44)\n",
        "y_obs = im_noise * (1 - cr_mask)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o7Dv52l5rT13"
      },
      "outputs": [],
      "source": [
        "figure(figsize=[15,5])\n",
        "subplot(131)\n",
        "imshow(x_true)\n",
        "title('Hubble image to recover')\n",
        "subplot(132)\n",
        "imshow(y_obs)\n",
        "title('Observed image')\n",
        "subplot(133)\n",
        "imshow(cr_mask)\n",
        "title('Cosmic Ray mask');"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0CtujABIrT13"
      },
      "source": [
        "## Step III: MAP Inference\n",
        "\n",
        "We now have all the tools for trying to perform Maximum A Posterior inference for our inverse problem, i.e.:\n",
        "\n",
        "$$\\hat{z} = \\arg \\max_{z} \\log p(y | z) + \\log p(z) $$\n",
        "\n",
        "In order to achieve this, you will need to put together the following elements:\n",
        "\n",
        "- Combine the physical forward model with generative model for an end-to-end forward model going from latent variable $z$ to HSC image.\n",
        "- Write a function that computes the log posterior for a given $z$\n",
        "- Use the tools from day I to do the optmization and recover a solution\n",
        "\n",
        "Your turn :-)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kka4OISkrT13"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PpaMoWIlrT13"
      },
      "source": [
        "## Step IV: Variational Inference\n",
        "\n",
        "In the previous section, we only recover a single point estimate of the solution, but ideally we want to access the full posterior. In this section, we will try to use VI.\n",
        "\n",
        "\n",
        "The idea of VI, is to use a parametric model $q_\\theta$ to approximate the posterior distribution $p(z | x)$. You need two things:\n",
        "- a tractable and flexible parametric model $q_\\theta$, we can use a Normalizing Flow for instance ;-)\n",
        "- a loss function that minimizes the distance between $p$ and $q_\\theta$\n",
        "\n",
        "\n",
        "The loss function typically used for VI is the Evidence Lower-Bound (ELBO) (the same one as we used in the VAE ;-) ). The ELBO is the right hand side part of this expression:\n",
        "\n",
        "$$ p_\\theta(y) \\geq \\mathbb{E}_{z \\sim q_\\theta}\\left[ p(y | z) \\right] - KL(q_\\theta || p) $$\n",
        "where $p$ in the KL divergence term is the latent space prior.\n",
        "\n",
        "In other words, maximizing the ELBO tries to maximize the likelihood of the data under the model.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eQf0UuDwrT13"
      },
      "outputs": [],
      "source": [
        "# We are going to need a normalizing flow to model the posterior then\n",
        "def sample_and_logp(N=1):\n",
        "    flow = AffineFlow()()\n",
        "    z = flow.sample(N, seed=hk.next_rng_key())\n",
        "    log_p = flow.log_prob(z)\n",
        "    return z, log_p"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j3TDWHYorT13"
      },
      "outputs": [],
      "source": [
        "q_sample_logp = hk.transform(sample_and_logp)\n",
        "\n",
        "# We initialize the parameters for the variational distribution\n",
        "q_params = q_sample_logp.init(next(rng_seq), 1)\n",
        "\n",
        "# And here is our prior distribution\n",
        "p = tfd.MultivariateNormalDiag(jnp.zeros(32),\n",
        "                               scale_identity_multiplier=1.)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b35SCz_9rT13"
      },
      "outputs": [],
      "source": [
        "# Let's write a concrete ELBO\n",
        "def elbo(params, rng_key):\n",
        "\n",
        "    # Sample from the log posterior\n",
        "    z, log_q = q_sample_logp.apply(params, rng_key, N=100)\n",
        "\n",
        "    # KL term\n",
        "    kl = log_q - p.log_prob(z)\n",
        "\n",
        "    # You need to plug your forward model producing a likelihood object here\n",
        "    likelihood = # .....\n",
        "\n",
        "    log_likelihood.log_prob(y_obs)\n",
        "\n",
        "    # Form the ELBO\n",
        "    elbo = jnp.mean(log_likelihood - kl)\n",
        "\n",
        "    return -elbo"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LlaRjtK0rT13"
      },
      "source": [
        "The rest is now up to you :-) Use this ELBO to optimize the parameters of the posterior variational distribution $q_\\theta$. Once you have achieved a good solution, try to sample from that posterior."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SSQOeYfLrT13"
      },
      "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
}