Tutorial 3.3 Normalizing_Flows_2D_Scalar_Field_Theory.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "**Normalizing Flows for Lattice Field Theory**\n",
        "----\n",
        "\n",
        ">Code courtesy: 100% of the code is taken from the paper [Introduction to Normalizing Flows for Lattice Field Theory](https://arxiv.org/abs/2101.08176).\n",
        "\n",
        "Given below is a snippet from the same paper.\n",
        "![picture](https://drive.google.com/uc?export=view&id=1axWOeSJC7j4WOwnBLIyNIiCYQySmyVAi)\n",
        "In this tutorial, we will consider the map $f: z \\to \\phi$, contrary to the convention above.\n",
        "\n",
        "\n",
        "Consider a Neural Network that we want to train, to obtain a target output  $\\phi$ with a probability distribution function\n",
        "$$P[\\phi] = \\frac{1}{Z}e^{-S[\\phi]} = \\frac{1}{Z} \\exp\\big( - \\int d^dx (\\partial_\\mu \\phi(x) )^2 + m^2 \\phi(x)^2 + \\lambda \\phi(x)^4 \\big).$$\n",
        "$Z$ is the partition function or the overall normalization of this theory. Here, $d$ is the dimension of a lattice of inputs for this Neural Network. In scalar field theory language, this lattice is discretized (Euclidean) space-time. For simplicity, let us consider the case $d=2$, then the target output density is given as\n",
        "$$P_{\\text{lattice}}[\\phi] = \\frac{1}{Z} e^{-S_{\\text{lattice}}[\\phi]} \\\\\n",
        "= \\frac{1}{Z} \\exp \\Big( - \\sum_{\\vec{n}} \\phi(\\vec{n}) \\Big[\\sum_{\\mu \\in \\{1,2 \\} } 2\\phi(\\vec{n}) - \\phi(\\vec{n} + {\\mu} ) - \\phi(\\vec{n}- {\\mu}) \\Big] + m^2 \\phi(\\vec{n})^2 + \\lambda \\phi(\\vec{n})^4 \\Big). $$\n",
        "The $2$ dimensional lattice of inputs is given as $L_x \\times L_y$, with number of lattice sites $\\vec{n} = (n_x,n_y)$. The field $\\phi(\\vec{n})$ is defined only on the lattice sites, and $n_x,n_y$ are integers. Lattice spacing is taken as $1$, in appropriate units.\n",
        "\n",
        "The Neural Network takes in lattice points $x ,y \\in \\{L_x, L_y\\}$, and trainable priors (a.k.a. stochastic parameters) $z$. Let $r(z)$ denote the prior density. The actual output density of this Network is denoted by $Q[\\phi]$. *Normalizing Flows* $f: z \\to \\phi$ are differentiable and invertible maps from the space of priors to the space of outputs.\n",
        "\n",
        "* Let us choose the convention that $L_x = L_y = L$, and $n_x, n_y \\in [0, L-1]$.\n",
        "* Let us also introduce alternate notions $P[\\phi]:= P(\\alpha)$ and $Q[\\phi] := Q(\\alpha)$ for the target density, and actual output density, respectively.\n",
        "* Neural Network output $\\alpha$ is a random variable, which in this case, corresponds to the value of the scalar field $\\phi(\\vec{n})$ evaluated on a single lattice site in $2$ dimensions.\n",
        "* Notations $\\phi(\\vec{n})$ and $\\phi_{\\vec{n}}$ respectively denote the field configuration that lives on the sites of the lattice, and the unraveled one-dimensional vector of lattice degrees of freedom.\n",
        "* Field $\\phi(\\vec{n}):= \\phi_1(\\vec{n}_x) \\times \\phi_2(\\vec{n}_y)$ can be thought of as some $(L \\times L)$ dimensional tensor.\n",
        "* Field action $S_{\\text{lattice}}[\\phi] \\in \\mathbb{R}$, and the path integral measure is a product over discretized ones $D\\phi(\\vec{n}) = \\prod_{\\vec{n}} d \\phi_{\\vec{n}}$.  \n",
        "* The partition function $Z = \\int \\prod_{\\vec{n}} d\\phi(\\vec{n}) \\, e^{- S_{\\text{lattice}}[\\phi]}$ is difficult to calculate, and obtaining the target density $P_{\\text{lattice}}[\\phi]$ up to some arbitrary constant is fine, for this tutorial.\n",
        "* *There is a curse of dimensionality here*. The dimensionality of the field density depends on the number of lattice sites, and grows with the volume of the lattice.\n",
        "\n",
        "*Reguralizing the scalar field theory on such a $2$ dimensional lattice, with one real degree of freedom per each lattice site, lets us access the non-perturbative physics in the strong coupling regime of the same theory.*\n",
        "\n",
        "**Let us introduce Normalizing Flows**\n",
        "----\n",
        "*A normalizing flow describes the transformation of a probability density through a sequence of invertible, differentiable mappings.*\n",
        "\n",
        "Any invertible, differentiable function $f(z)$ will transform the prior density $r(z)$ on random variable $z$ to a output density $Q(x):= Q(\\alpha)$ on variable $\\alpha = f(z)$; both $z$ and $\\alpha$ can be multi-dimensional random variables.\n",
        "* Let $f^{-1} = g$, then $z = g \\circ f(z)$.\n",
        "* The normalizing flow $f$ induces a change of variable(s) from $z$ to $\\alpha$.\n",
        "* The output density $Q(\\alpha)$ can be obtained from $r(z)$, if we can compute the Jacobian factor $J(z) = \\big|\\det_{kl} \\partial f(z_k)/\\partial z_l \\big|$ for this transformation. $$Q(\\alpha) = r(z)\\big[ J(z)\\big]^{-1} = r(z) \\Bigg|\\det_{kl} \\frac{\\partial f_k(z)}{\\partial z_l} \\Bigg|^{-1} $$\n",
        "\n",
        "**Let us introduce Coupling Layers**\n",
        "----\n",
        "\n",
        "An invertible, differentiable expressive function $f$ can be constructed as a composition of multiple simpler invertible, differentiable functions $f_i$,\n",
        "$$f = f_k \\circ f_{k-1} \\circ \\cdots f_1.$$\n",
        "Coupling layers are one of the approaches to define the simpler mappings $f_i$, in a way such that only a subset of the input variables are updated through each $f_i$, conditioned on the complimentary (\"frozen\") subset.\n",
        "* As a simple example, if the inputs to a set of coupling layers is a lattice with one real number per each lattice site, then the first two coupling layers $f_1$ and $f_2$ can be designed to update only the even and odd sites, respectively.\n",
        "* The transformation applied by each coupling layer is manifestly invertible, e.g. a scaling transformation $x \\to e^s x$, or an affine transformation $x \\to e^s x + t$.\n",
        "* Next, a coupling layer generating outputs $x_1', x'_2 = f_i(x_1,x_2)$ via a scaling transformation is defined below\n",
        "$$x'_1 = e^{s(x_2)} x_1, \\\\\n",
        "x'_2 = x_2, $$ where $x_1, x_2$ are subsets of the layer input $x$.\n",
        "* The coupling layer in the above example keeps $x_2$ frozen, and $x_1$ is updated based on frozen $x_2$.\n",
        "* The product in $e^{s(x_2)}x_1$ is elementwise, and $e^{s(x_2)}$ is defined elementwise on $x_2$, and therefore, has the same shape as $x_1$.\n",
        "* The transformation inducing parameters in $s(x_2)$ can be complicated non-invertible functions of frozen $x_2$.\n",
        "* Such choices are not a problem as the inverse transformation can easily be defined, using the same functions $s(x_2)$, as\n",
        "$$x_1 = e^{-s(x'_2)} x'_1, \\\\\n",
        "x_2 = x'_2. $$\n",
        "* The Jacobian of the forward transformation\n",
        "$$J_{\\text{forward}}(x) = \\Bigg|\\det_{kl} \\frac{\\partial \\big[ f_i(x_1, x_2) \\big]_k }{\\partial x_l} \\Bigg| = \\prod_k e^{[s(x_2)]_k} ,$$\n",
        "is the inverse of the Jacobian of the reverse transformation associated with $f^{-1}_{i}$, i.e. $J_{\\text{reverse}}(x') = 1/J_{\\text{forward}}(x)$.\n",
        "\n",
        "*The Jacobian for the normalizing flow $f$ is obtained by composing all Jacobians of all coupling layers together, $$Q(\\alpha) = r(z)\\prod_{i=1}^{k} J_{i}^{-1}(z) = r(z) \\Bigg|\\det \\frac{\\partial f(z)}{\\partial z} \\Bigg|^{-1} .$$*\n",
        "\n",
        "\n",
        "\n",
        "**Coupling Layers for Scalar Field Theory on $2$D Lattice**\n",
        "---\n",
        "\n",
        "Previously, we defined scalar field theory on a $2$ dimensional lattice of size $L \\times L$. We will now define coupling layers for this problem, such that the real degrees of freedom on each lattice site $\\phi$ can be divided into two subsets $\\phi_1, \\phi_2$, and the coupling layers will induce affine and / or scaling transformations on one subset, while keeping the other frozen.\n",
        "\n",
        "We choose the same set of affine transformations as before. Then, each coupling layer of the form $f_i: (\\phi_1, \\phi_2) \\mapsto (\\phi_1' , \\phi_2') $ induces the following transformation\n",
        "$$ \\phi'_1 = e^{s(\\phi_2)}\\phi_1 + t(\\phi_2), \\\\\n",
        "\\phi'_2 = \\phi_2,$$\n",
        "on its inputs. The inverse of this layer is given by $f_i^{-1}(\\phi'_1, \\phi'_2) = \\big( (\\phi'_1 - t(\\phi'_2) )\\, e^{-s(\\phi'_2) } , \\phi'_2 \\big) .$\n",
        "* Here, $s(\\phi_2), t(\\phi_2)$ are vectors of the same size as $\\phi_1$, and all operations are elementwise.\n",
        "* Later, we will define functions $s,t$ using Neural Networks.\n",
        "* The logarithm of Jacobians for $f_i$ (forward) and its inverse $f^{-1}_{i}$ (reverse) transformations are respectively\n",
        "$$\\log J_{\\text{forward}}(\\phi) = \\sum_k \\big[s(\\phi_2) \\big]_k, \\\\\n",
        "\\log J_{\\text{reverse}}(\\phi') = \\sum_k - \\big[s(\\phi'_2) \\big]_k. $$\n",
        "* For implimentations in Neural Networks, we will define subsets $\\phi_1, \\phi_2$ using a mask $m(\\vec{n}) \\in \\{0,1 \\}$. When $m(\\vec{n})=1$, $\\phi_1$ is updated using functions $s,t$, while $\\phi_2$ is kept frozen.\n",
        "\n",
        "\n",
        "**KL Divergence Loss for Normalizing Flow**\n",
        "----\n",
        "Once we implement the coupling layers $f_i$ for scalar field theory on a $2$ dimensional lattice using a Neural Network, it will result in output distribution $Q[\\phi']$ on transformed variable $\\phi'$. A priori, there is no reason to believe that $Q[\\phi']$ would be close to target scalar field density $P_{\\text{lattice}}[\\phi']$ on $\\phi'$. We need a good loss function to minimize the differences between the approximate and true distributions $Q[\\phi']$ and $P_{\\text{lattice}}[\\phi']$, respectively.\n",
        "\n",
        "Recall [Tutorial 3.1 KL_divergence.ipynb](https://github.com/iaifi/summer-school-2023/blob/main/Tutorial%203.1%20KL_divergence.ipynb) from earlier today. Target output data $\\phi'$ from target distribution $P_{\\text{lattice}}[\\phi']$ can be scarce, which may pose a problem to training. To avoid this, we define the reverse Kullback-Leibler diveregnce between $Q[\\phi']$ and $P_{\\text{lattice}}[\\phi']$, under the assumption that output $\\phi'$ is a continuous random variable. Then, the reverse KL divergence loss is\n",
        "$$D_{KL}(Q||P_{\\text{lattice}}) = \\int d\\phi \\, Q(\\phi') \\big( \\log Q(\\phi') - \\log P_{\\text{lattice}}(\\phi')  \\big) .$$\n",
        "Actual outputs $\\phi'$ from the coupling layers would be discrete multidimensional random variables. Then, a trainable KL divergence loss would require some modifications of the above equation.\n",
        "\n",
        "Let's redefine the reverse KL divergence for approximate and true distributions $Q[\\phi']$ and $P_{\\text{lattice}}[\\phi']$ as\n",
        "$$\\tilde{D}_{KL}(Q||P_{\\text{lattice}}) = \\frac{1}{N} \\sum_{i=1}^{N} \\big( \\log Q[\\phi'_i]  - \\log P_{\\text{lattice}}[\\phi'_i] \\big) ,$$\n",
        "for $N$ samples $\\phi'_i$ drawn from target distribution $P_{\\text{lattice}}$. A problem with this choice may be the mode-seeking or zero-forcing behavior of reverse KL diveregnce, we will not discuss it here.\n",
        "\n",
        "\n",
        "**Effective Sample Size**\n",
        "---\n",
        "\n",
        "Defined below is Effective Sample Size or ESS for a batch of samples $\\phi'_i$. It's a measure of the model quality that does not require overall normalization of $P_{\\text{lattice}}[\\phi']$,\n",
        "$$ESS = \\frac{\\Big(\\frac{1}{N}\\sum_i \\frac{ P_{\\text{lattice}}[\\phi'_i] }{ Q[\\phi'_i]} \\Big)^2 }{ \\frac{1}{N}\\sum_i \\Big( \\frac{ P_{\\text{lattice}}[\\phi'_i] }{ Q[\\phi'_i]} \\Big)^2 } .$$\n",
        "It can be seen that ESS lies in the range $[0,1]$; larger values indicate better effective sampling, and a perfect value of $ESS = 1$ indicates perfect independent draw from target distribution $P_{\\text{lattice}}$ for each sample.\n",
        "\n",
        "ESS is a noisier quantity than KL divergence, and therefore, is not used as a loss function.\n"
      ],
      "metadata": {
        "id": "kQsePaHfUkyh"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "**It's time to implement everything in codes.**\n",
        "----\n",
        "\n",
        "At this point, we will import all requires libraries, and define functions that we may need to call at a later point."
      ],
      "metadata": {
        "id": "mxk4LNiLBDoK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# import libraries\n",
        "import base64\n",
        "import io\n",
        "import pickle\n",
        "import numpy as np\n",
        "import torch\n",
        "print(f'TORCH VERSION: {torch.__version__}')\n",
        "import packaging.version\n",
        "if packaging.version.parse(torch.__version__) < packaging.version.parse('1.5.0'):\n",
        "    raise RuntimeError('Torch versions lower than 1.5.0 not supported')\n",
        "\n",
        "%matplotlib inline\n",
        "import matplotlib.pyplot as plt\n",
        "import seaborn as sns\n",
        "sns.set_style('whitegrid')\n",
        "\n",
        "# set up environmet for pytorch\n",
        "if torch.cuda.is_available():\n",
        "    torch_device = 'cuda'\n",
        "    float_dtype = np.float32 # single\n",
        "    torch.set_default_tensor_type(torch.cuda.FloatTensor)\n",
        "else:\n",
        "    torch_device = 'cpu'\n",
        "    float_dtype = np.float64 # double\n",
        "    torch.set_default_tensor_type(torch.DoubleTensor)\n",
        "print(f\"TORCH DEVICE: {torch_device}\")\n",
        "\n",
        "# define two functions for mod and remainder\n",
        "def torch_mod(x):\n",
        "    return torch.remainder(x, 2*np.pi)\n",
        "def torch_wrap(x):\n",
        "    return torch_mod(x+np.pi) - np.pi\n",
        "\n",
        "# a function to detach tensors from the computational graph and pull them to CPU as a numpy array\n",
        "def grab(var):\n",
        "    return var.detach().cpu().numpy()\n",
        "\n",
        "#####\n",
        "# the set of functions defined below make a live-updating plot during training\n",
        "from IPython.display import display\n",
        "\n",
        "def init_live_plot(dpi=125, figsize=(8,4)):\n",
        "    fig, ax_ess = plt.subplots(1,1, dpi=dpi, figsize=figsize)\n",
        "    plt.xlim(0, N_era*N_epoch)\n",
        "    plt.ylim(0, 1)\n",
        "\n",
        "    ess_line = plt.plot([0],[0], alpha=0.5) # dummy\n",
        "    plt.grid(False)\n",
        "    plt.ylabel('ESS')\n",
        "\n",
        "    ax_loss = ax_ess.twinx()\n",
        "    loss_line = plt.plot([0],[0], alpha=0.5, c='orange') # dummy\n",
        "    plt.grid(False)\n",
        "    plt.ylabel('Loss')\n",
        "\n",
        "    plt.xlabel('Epoch')\n",
        "\n",
        "    display_id = display(fig, display_id=True)\n",
        "\n",
        "    return dict(\n",
        "        fig=fig, ax_ess=ax_ess, ax_loss=ax_loss,\n",
        "        ess_line=ess_line, loss_line=loss_line,\n",
        "        display_id=display_id\n",
        "    )\n",
        "\n",
        "def moving_average(x, window=10):\n",
        "    if len(x) < window:\n",
        "        return np.mean(x, keepdims=True)\n",
        "    else:\n",
        "        return np.convolve(x, np.ones(window), 'valid') / window\n",
        "\n",
        "def update_plots(history, fig, ax_ess, ax_loss, ess_line, loss_line, display_id):\n",
        "    Y = np.array(history['ess'])\n",
        "    Y = moving_average(Y, window=15)\n",
        "    ess_line[0].set_ydata(Y)\n",
        "    ess_line[0].set_xdata(np.arange(len(Y)))\n",
        "    Y = history['loss']\n",
        "    Y = moving_average(Y, window=15)\n",
        "    loss_line[0].set_ydata(np.array(Y))\n",
        "    loss_line[0].set_xdata(np.arange(len(Y)))\n",
        "    ax_loss.relim()\n",
        "    ax_loss.autoscale_view()\n",
        "    fig.canvas.draw()\n",
        "    display_id.update(fig) # need to force colab to update plot\n",
        "#####"
      ],
      "metadata": {
        "id": "iy57N9X2e8h6",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2405cab5-0a54-4bde-f16c-1a2bfdbe5942"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "TORCH VERSION: 2.0.1+cu118\n",
            "TORCH DEVICE: cpu\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Exercise 0: Compute the Jacobian for the transformation $f:(Z_1, Z_2) \\to (U_1, U_2)$ where $$Z_1 = \\sqrt{-2 \\ln U_1}\\cos(\\pi U_2), \\\\ Z_2 = \\sqrt{-2 \\ln U_1}\\sin(\\pi U_2).$$**\n",
        "---\n"
      ],
      "metadata": {
        "id": "k7_jddgSclso"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "**Exercise 1: Sample $U_1, U_2 \\sim \\text{unif}(0,1)$ for a batch size of $2^{12}$. Write a code in python to plot transformed variables $Z_2$ against $Z_1$, and original variables $U_2$ against $U_1$, separately.**\n",
        "---"
      ],
      "metadata": {
        "id": "L9iTSIa4dNMt"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Next, let us implement a simple Coupling Layer using a feedforward Neural Network.**\n",
        "---\n",
        "\n",
        "Consider two random variables $x_1, x_2 \\sim \\text{unif}[0,1)$, and a map $f: (x_1, x_2) \\to (x'_1, x'_2)$ with\n",
        "$$x'_1 = e^{s(x_2)} x_1, \\\\\n",
        "x'_2 = x_2.$$\n",
        "\n",
        "* First, we define the coupling layer as a class called ```SimpleCouplingLayer```. There are two different functions in this class, one of them is $f: (x_1, x_2) \\to (x'_1, x'_2)$, called ```forward(self, x)```, and the other one is its inverse ```reverse(self, fx)```. Here, ```fx``` is the transformed output $f(x_1, x_2)$.\n",
        "* Then, priors are initialized in a way to give interesting behaviour without training.\n",
        "* Next, we evaluate $f(x_1, x_2)$, and $f^{-1}\\big(f(x_1, x_2) \\big)$.\n",
        "* Lastly, we plot (a)  $(x_1, x_2)$, (b) $f(x_1, x_2)$, and (c) $f^{-1}\\big(f(x_1, x_2) \\big)$ in three different plots.\n",
        "\n",
        "\n",
        "**Exercise 2: Write a function ```reverse(self, fx)``` in the class ```SimpleCouplingLayer``` to return $f^{-1}\\big(f(x_1, x_2) \\big)$. If it is properly defined, the plots (a) and (c) in the last step of this example would be identical.**\n",
        "---"
      ],
      "metadata": {
        "id": "-sKBXYcDKkL6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# let's first define the class for the Coupling Layer\n",
        "\n",
        "class SimpleCouplingLayer(torch.nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.s = torch.nn.Sequential(\n",
        "            torch.nn.Linear(1, 8),\n",
        "            torch.nn.ReLU(),\n",
        "            torch.nn.Linear(8, 8),\n",
        "            torch.nn.ReLU(),\n",
        "            torch.nn.Linear(8, 1),\n",
        "            torch.nn.Tanh()\n",
        "        )\n",
        "\n",
        "    # the map f: x --> x' is defined below\n",
        "    def forward(self, x):\n",
        "        x1, x2 = x[:,0], x[:,1]\n",
        "        s = self.s(x2.unsqueeze(-1)).squeeze(-1)\n",
        "        fx1 = torch.exp(s) * x1\n",
        "        fx2 = x2\n",
        "        logJ = s\n",
        "        return torch.stack((fx1, fx2), dim=-1), logJ\n",
        "\n",
        "\n",
        "    # write a function reverse(self, fx) as solution to exercise 2.\n",
        "    ####\n",
        "    ####\n",
        "\n",
        "coupling_layer = SimpleCouplingLayer()\n",
        "\n",
        "# init weights in a way that gives interesting behavior without training\n",
        "def set_weights(m):\n",
        "    if hasattr(m, 'weight') and m.weight is not None:\n",
        "        torch.nn.init.normal_(m.weight, mean=1, std=2)\n",
        "    if hasattr(m, 'bias') and m.bias is not None:\n",
        "        m.bias.data.fill_(-1)\n",
        "torch.manual_seed(1234)\n",
        "coupling_layer.s.apply(set_weights);\n",
        "\n",
        "# we generate x, f(x), and f^{-1} (f(x))\n",
        "batch_size = 750\n",
        "np_x = (2*np.random.random(size=(batch_size, 2)) - 1).astype(float_dtype)\n",
        "x = torch.from_numpy(np_x).to(torch_device)\n",
        "\n",
        "gx, fwd_logJ = coupling_layer.forward(x)\n",
        "xp, bwd_logJ = coupling_layer.reverse(gx)\n",
        "\n",
        "# lastly, we plot x, f(x), and f^{-1} (f(x)). Expect first and last plots to be identical.\n",
        "fig, ax = plt.subplots(1,3, dpi=125, figsize=(6,2.3), sharex=True, sharey=True)\n",
        "np_gx, np_xp = grab(gx), grab(xp)\n",
        "for a in ax:\n",
        "    a.set_xlim(-1.1,1.1)\n",
        "    a.set_ylim(-1.1,1.1)\n",
        "ax[0].scatter(np_x[:,0], np_x[:,1], marker='.')\n",
        "ax[0].set_title(r'$x$')\n",
        "ax[1].scatter(np_gx[:,0], np_gx[:,1], marker='.', color='tab:orange')\n",
        "ax[1].set_title(r'$f(x)$')\n",
        "ax[2].scatter(np_xp[:,0], np_xp[:,1], marker='.')\n",
        "ax[2].set_title(r\"$f^{-1}(f(x))$\")\n",
        "fig.set_tight_layout(True)\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "3WrX3bEvmFYX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Next, let us encode the $2$D scalar field theory on lattice.**\n",
        "---\n",
        "\n",
        "Let us go step by step.\n",
        "* First, we will define a $8 \\times 8$  lattice, and randomly initialize $\\phi(\\vec{n})$ on each lattice site, as a real number drawn from Gaussian distribution $\\text{Normal}(0,1)$.\n",
        "* There are two independent field configurations $\\phi_1, \\phi_2$ in two orthogonal directions. We will call this scalar field $\\vec{\\phi} = \\phi_1 \\times \\phi_2$, defined over all lattice sites, as ```cfgs```.\n",
        "* Then, we will define a class called ```ScalarPhi4Action``` to compute and return scalar field theory action $S_{\\text{lattice}}[\\phi]$, given each $\\phi(\\vec{n})$. This class takes in the values for scalar field mass $m$ and quartic coupling constant $\\lambda$.\n",
        "\n",
        "**Exercise 3: Write a callable object in class ```ScalarPhi4Action``` that will evaluate actions  $S_{\\text{lattice}}[\\phi_1], S_{\\text{lattice}}[\\phi_2]$ for each independent field configuration, and return them in an array format. Note that each $S_{\\text{lattice}}[\\phi_i]$ is a scalar.**\n",
        "----"
      ],
      "metadata": {
        "id": "zIjSabdCrM5O"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# first, we define a 8x8 lattice\n",
        "L = 8\n",
        "lattice_shape = (L,L)\n",
        "\n",
        "# then, we initialize a real values d.o.f. on each lattice site. Overall, these are two independent field configurations.\n",
        "phi_ex1 = np.random.normal(size=lattice_shape).astype(float_dtype)\n",
        "phi_ex2 = np.random.normal(size=lattice_shape).astype(float_dtype)\n",
        "cfgs = torch.from_numpy(np.stack((phi_ex1, phi_ex2), axis=0)).to(torch_device)\n",
        "\n",
        "# here, we define a class to return the action for each real d.o.f. on each lattice site.\n",
        "class ScalarPhi4Action:\n",
        "    def __init__(self, M2, lam):\n",
        "        self.M2 = M2\n",
        "        self.lam = lam\n",
        "\n",
        "    ### write a callable object here __call__(self, cfgs) as solution to exercise 3.\n",
        "    ###\n",
        "    ###\n",
        "\n",
        "\n",
        "# We can choose different values for square of mass and the quartic coupling.\n",
        "M2 = -4.0\n",
        "lam = 8.0\n",
        "phi4_action = ScalarPhi4Action(M2=M2, lam=lam)\n",
        "print(phi4_action(cfgs))"
      ],
      "metadata": {
        "id": "MLzUZz2SvVVc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* Next, we will define a class ```SimpleNormal``` that initializes priors drawn from uncorrelated Gaussian distributions.\n",
        "* Using this, we will generate priors for this $8 \\times 8$ lattice.\n",
        "* Lastly, we will use ```draw``` function, defined in the beginning, to acquire samples from this prior and plot them.  "
      ],
      "metadata": {
        "id": "H2dy_jnQ2RKE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# first, let's define a class for uncorrelated Gaussians for prior distributions\n",
        "class SimpleNormal:\n",
        "    def __init__(self, loc, var):  # attributes 'loc' and 'var' determine the shape of the samples drawn\n",
        "        self.dist = torch.distributions.normal.Normal(\n",
        "            torch.flatten(loc), torch.flatten(var))\n",
        "        self.shape = loc.shape\n",
        "    def log_prob(self, x):\n",
        "        logp = self.dist.log_prob(x.reshape(x.shape[0], -1))\n",
        "        return torch.sum(logp, dim=1)\n",
        "    def sample_n(self, batch_size):\n",
        "        x = self.dist.sample((batch_size,))\n",
        "        return x.reshape(batch_size, *self.shape)\n",
        "\n",
        "# we use this to initialize the priors on the lattice\n",
        "prior = SimpleNormal(torch.zeros(lattice_shape), torch.ones(lattice_shape))\n",
        "\n",
        "# draw samples from this prior and plot them\n",
        "torch_z = prior.sample_n(1024)\n",
        "z = grab(torch_z)\n",
        "\n",
        "fig, ax = plt.subplots(4,4, dpi=125, figsize=(4,4))\n",
        "for i in range(4):\n",
        "    for j in range(4):\n",
        "        ind = i*4 + j\n",
        "        ax[i,j].imshow(np.tanh(z[ind]), vmin=-1, vmax=1, cmap='viridis')\n",
        "        ax[i,j].axes.xaxis.set_visible(False)\n",
        "        ax[i,j].axes.yaxis.set_visible(False)\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "nRDonw0SlrP0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* Given below is a code to plot two-way histograms between lattice sites. The plots show that the priors are uncorrelated."
      ],
      "metadata": {
        "id": "V4CjiwzQ4FsG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "fig, ax = plt.subplots(4,4, dpi=125, figsize=(4,4))\n",
        "for x1 in range(2):\n",
        "    for y1 in range(2):\n",
        "        i1 = x1*2 + y1\n",
        "        for x2 in range(2):\n",
        "            for y2 in range(2):\n",
        "                i2 = x2*2 + y2\n",
        "                ax[i1,i2].hist2d(z[:,x1,y1], z[:,x2,y2], range=[[-3,3],[-3,3]], bins=20)\n",
        "                ax[i1,i2].set_xticks([])\n",
        "                ax[i1,i2].set_yticks([])\n",
        "                if i1 == 3:\n",
        "                    ax[i1,i2].set_xlabel(rf'$\\phi({x2},{y2})$')\n",
        "                if i2 == 0:\n",
        "                    ax[i1,i2].set_ylabel(rf'$\\phi({x1},{y1})$')\n",
        "fig.suptitle(\"Correlations in Various Lattice Sites\")\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "vlZTKG4W4Npe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "* Finally, we evaluate scalar field theory action $S_{\\text{lattice}}[z]$ on these uncorrelated priors as\n",
        "$$S_{\\text{lattice}}[z] = \\sum_{\\vec{n}} z(\\vec{n}) \\Big[\\sum_{\\mu \\in \\{1,2 \\} } 2z(\\vec{n}) - z(\\vec{n} + {\\mu} ) - z(\\vec{n}- {\\mu}) \\Big] + m^2 z(\\vec{n})^2 + \\lambda z(\\vec{n})^4,$$\n",
        "and compare it against the actual action associated with prior distribution $r(z)$. The actual action is given as $S_{\\text{eff}} = - \\log r(z)$. The plot shows differences between the actual and scalar field theory actions on $z$.\n"
      ],
      "metadata": {
        "id": "ycwmKqoJx6kr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "S_eff = -grab(prior.log_prob(torch_z))\n",
        "S = grab(phi4_action(torch_z))\n",
        "fit_b = np.mean(S) - np.mean(S_eff)\n",
        "print(f'slope 1 linear regression S = -logr + {fit_b:.4f}')\n",
        "fig, ax = plt.subplots(1,1, dpi=125, figsize=(4,4))\n",
        "ax.hist2d(S_eff, S, bins=20, range=[[-800, 800], [200,1800]])\n",
        "xs = np.linspace(-800, 800, num=4, endpoint=True)\n",
        "ax.plot(xs, xs + fit_b, ':', color='w', label='slope 1 fit')\n",
        "ax.set_xlabel(r'$S_{\\mathrm{eff}} \\equiv -\\log~r(z)$')\n",
        "ax.set_ylabel(r'$S(z)$')\n",
        "ax.set_aspect('equal')\n",
        "plt.legend(prop={'size': 6})\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "YGfQTx0APtEA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## The prior distribution is expected to change due to training, such that the slope of this plot would become $\\sim 1$.\n"
      ],
      "metadata": {
        "id": "zI6yRSW35wtA"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Next, let us encode affine Coupling Layers for this lattice field theory problem.**\n",
        "---\n",
        "\n",
        "The steps are given below.\n",
        "* We start by designing a checkerboard masking for $m(\\vec{n}) \\in \\{0,1 \\}$, defined earlier in the tutorial.\n",
        "* Then, we define a class ```AffineCoupling``` for transformations given by $f: \\phi \\to \\phi^{-1}$ and its inverse transformation $f^{-1}$.\n",
        "* Later, when we will set up affine transformations on priors on the $2$D lattice using Neural Networks, we will let functions $s(\\phi_2), t(\\phi_2)$ transform both $\\phi_2$ and $\\phi_1$, and then use the masking to get\n",
        "$$ \\phi'_1 = e^{s(\\phi_2)}\\phi_1 + t(\\phi_2), \\\\\n",
        "\\phi'_2 = \\phi_2.$$\n",
        "* A single Neural Network will be used to define both $s,t$. The shared parameters in $s$ amd $t$ do not add to formal complexities.\n",
        "* We will use Convolutional Neural Nets for this $2$D lattice field theory problem. The implementation of Pytorch's $2$D CNNs require inputs in the shape ```(batch_size, n_input_channels, L_x, L_y)```.\n",
        "* To ensure that, we use ```.unsqueeze(1)``` in the class ```AffineCoupling```, to add a fake ```n_input_channels``` dimension of length $1$.\n",
        "* The outputs of the CNNs are shaped as ```(batch_size,n_output_channels,L_x, L_y)```. We will choose ```n_output_channels=2```, one for $s(x_2)$ and $t(x_2)$ each.\n",
        "\n",
        "**Exercise 4: Define a function called ```reverse(self, fx)``` in class ```AffineCoupling``` that induces inverse affine transformation $f^{-1}$. Get outputs as tuples of the form $\\big(f^{-1}\\big( f(\\phi)\\big), \\log J\\big) $, where $J$ denotes the Jacobian for $f^{-1}: \\phi' \\to \\phi$.**\n",
        "----"
      ],
      "metadata": {
        "id": "JaczZ4WZ6X3q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# define checkerboard maskig first\n",
        "def make_checker_mask(shape, parity):\n",
        "    checker = torch.ones(shape, dtype=torch.uint8) - parity\n",
        "    checker[::2, ::2] = parity\n",
        "    checker[1::2, 1::2] = parity\n",
        "    return checker.to(torch_device)\n",
        "\n",
        "#print(\"For example this is the mask for an 8x8 configuration:\\n\", make_checker_mask(lattice_shape, 0))\n",
        "\n",
        "# define the class for affine coupling layer\n",
        "class AffineCoupling(torch.nn.Module):\n",
        "    def __init__(self, net, *, mask_shape, mask_parity):\n",
        "        super().__init__()\n",
        "        self.mask = make_checker_mask(mask_shape, mask_parity)\n",
        "        self.net = net\n",
        "\n",
        "    def forward(self, x):\n",
        "        x_frozen = self.mask * x\n",
        "        x_active = (1 - self.mask) * x\n",
        "        net_out = self.net(x_frozen.unsqueeze(1))\n",
        "        s, t = net_out[:,0], net_out[:,1]\n",
        "        fx = (1 - self.mask) * t + x_active * torch.exp(s) + x_frozen\n",
        "        axes = range(1,len(s.size()))\n",
        "        logJ = torch.sum((1 - self.mask) * s, dim=tuple(axes))\n",
        "        return fx, logJ\n",
        "\n",
        "    # define s function reverse(self, fx) for the inverse affine transformation.\n",
        "    ####\n",
        "    ####\n",
        "\n",
        "# both function s,t associated with each coupling layer are defined using a CNN\n",
        "def make_conv_net(*, hidden_sizes, kernel_size, in_channels, out_channels, use_final_tanh):\n",
        "    sizes = [in_channels] + hidden_sizes + [out_channels]\n",
        "    assert packaging.version.parse(torch.__version__) >= packaging.version.parse('1.5.0')\n",
        "    assert kernel_size % 2 == 1, 'kernel size must be odd for PyTorch >= 1.5.0'\n",
        "    padding_size = (kernel_size // 2)\n",
        "    net = []\n",
        "    for i in range(len(sizes) - 1):\n",
        "        net.append(torch.nn.Conv2d(\n",
        "            sizes[i], sizes[i+1], kernel_size, padding=padding_size,\n",
        "            stride=1, padding_mode='circular'))\n",
        "        if i != len(sizes) - 2:\n",
        "            net.append(torch.nn.LeakyReLU())\n",
        "        else:\n",
        "            if use_final_tanh:\n",
        "                net.append(torch.nn.Tanh())\n",
        "    return torch.nn.Sequential(*net)"
      ],
      "metadata": {
        "id": "Eg2LWLuE_-3n"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## The goal here is to obtain output $\\phi'$ from target density $P_{\\text{lattice}}[\\phi']$. Variable $\\phi'$ is obtained from lattice degrees of freedom $\\phi$, after affine transformation by $f$.\n",
        "* We hope to achieve this by composing a sequence of affine coupling layers.\n",
        "* We will use $16$ affine coupling layers, while embedding alternate lattice sites with checkerboard parity.\n",
        "* We will define a function ```make_phi4_affine_layers```, that imposes class ```AffineCoupling``` on the CNN defined above.\n",
        "* The function ```make_phi4_affine_layers``` will ensure that this CNN is transformed according to all $16$ affine coupling layers, in appropriate order of sequence.\n",
        "* CNN kernel size is kept at a minimum $3 \\times 3$ to induce correlations between neighboring sites. Larger kernel size will require more compute power.\n",
        "\n",
        "**Exercise 5: Fill in the missing parts in function ```make_phi4_affine_layers``` which ensures that this function returns appropriately transformed CNN layers.**\n",
        "---"
      ],
      "metadata": {
        "id": "VJxYeXvdD8FM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# define the function to induce affine transformations on CNN layers\n",
        "def make_phi4_affine_layers(*, n_layers, lattice_shape, hidden_sizes, kernel_size):\n",
        "    layers = []\n",
        "    for i in range(n_layers):\n",
        "        parity = i % 2\n",
        "        net = make_conv_net(\n",
        "            in_channels=1, out_channels=2, hidden_sizes=hidden_sizes,\n",
        "            kernel_size=kernel_size, use_final_tanh=True)\n",
        "        # write down necessary lines to ensure each layer undergoes affine transformation and gets appended to the transformed CNN\n",
        "        ####\n",
        "        ####\n",
        "\n",
        "    return torch.nn.ModuleList(layers)\n",
        "\n",
        "# get the transformed CNN\n",
        "n_layers = 8\n",
        "hidden_sizes = [8,8]\n",
        "kernel_size = 3\n",
        "layers = make_phi4_affine_layers(\n",
        "    lattice_shape=lattice_shape, n_layers=n_layers,\n",
        "    hidden_sizes=hidden_sizes, kernel_size=kernel_size)\n",
        "model = {'layers': layers, 'prior': prior}"
      ],
      "metadata": {
        "id": "Ehxe8fCFHZim"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Next, we define a function called ```apply_flow_to_prior``` that returns the transformed prior density $Q[\\phi']$.\n",
        "* This is used to define a function for the the training step, called ```train_step```.\n",
        "* The function ```train_step``` calls in reverse Kl divergence loss.\n",
        "\n",
        "**Exercise 6: Define a function called ```calc_dkl```, that takes in $\\log Q$ and $\\log P_{\\text{lattice}}$ as arguments, computes reverse KL divergence, and returns a scalar valued loss.**\n",
        "---\n",
        "\n",
        "**Exercise 7: Write a function called ```compute_ess``` that takes in $\\log Q$ and $\\log P_{\\text{lattice}}$ as arguments, and returns the ESS for each batch of samples $\\phi'_i$.**\n",
        "---\n"
      ],
      "metadata": {
        "id": "Ts59OQCoKvHH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# given below is the function that computes the transformed priors\n",
        "def apply_flow_to_prior(prior, coupling_layers, *, batch_size):\n",
        "    x = prior.sample_n(batch_size)\n",
        "    logq = prior.log_prob(x)\n",
        "    for layer in coupling_layers:\n",
        "        x, logJ = layer.forward(x)\n",
        "        logq = logq - logJ\n",
        "    return x, logq\n",
        "\n",
        "\n",
        "# write a function calc_dkl(logp, logq) to return reverse KL divergence loss\n",
        "####\n",
        "####\n",
        "\n",
        "# the function below defines the training step\n",
        "def train_step(model, action, loss_fn, optimizer, metrics):\n",
        "    layers, prior = model['layers'], model['prior']\n",
        "    optimizer.zero_grad()\n",
        "\n",
        "    x, logq = apply_flow_to_prior(prior, layers, batch_size=batch_size)\n",
        "    logp = -action(x)\n",
        "    loss = calc_dkl(logp, logq)\n",
        "    loss.backward()\n",
        "\n",
        "    optimizer.step()\n",
        "\n",
        "    metrics['loss'].append(grab(loss))\n",
        "    metrics['logp'].append(grab(logp))\n",
        "    metrics['logq'].append(grab(logq))\n",
        "    metrics['ess'].append(grab( compute_ess(logp, logq) ))\n",
        "\n",
        "# write a function compute_ess(logp, logq) to get ESS for each batch of samples\n",
        "###\n",
        "###\n",
        "\n",
        "# a function to print training history\n",
        "def print_metrics(history, avg_last_N_epochs):\n",
        "    print(f'== Era {era} | Epoch {epoch} metrics ==')\n",
        "    for key, val in history.items():\n",
        "        avgd = np.mean(val[-avg_last_N_epochs:])\n",
        "        print(f'\\t{key} {avgd:g}')"
      ],
      "metadata": {
        "id": "TEAEKcAU4fAU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Next, let's train this model.\n"
      ],
      "metadata": {
        "id": "R1BkliUiQMkn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Training specifications\n",
        "base_lr = .001\n",
        "optimizer = torch.optim.Adam(model['layers'].parameters(), lr=base_lr)\n",
        "use_pretrained = False\n",
        "\n",
        "# if not pretrained, then use the code below. If you have pretrained weights for similar lattice field configurations, that can be loaded here too.\n",
        "N_era = 25\n",
        "N_epoch = 100\n",
        "batch_size = 64\n",
        "print_freq = N_epoch\n",
        "plot_freq = 1\n",
        "\n",
        "history = {\n",
        "    'loss' : [],\n",
        "    'logp' : [],\n",
        "    'logq' : [],\n",
        "    'ess' : []\n",
        "}\n",
        "\n",
        "if not use_pretrained:\n",
        "    [plt.close(plt.figure(fignum)) for fignum in plt.get_fignums()] # close all existing figures\n",
        "    live_plot = init_live_plot()\n",
        "\n",
        "    for era in range(N_era):\n",
        "        for epoch in range(N_epoch):\n",
        "            train_step(model, phi4_action, calc_dkl, optimizer, history)\n",
        "\n",
        "            if epoch % print_freq == 0:\n",
        "                print_metrics(history, avg_last_N_epochs=print_freq)\n",
        "\n",
        "            if epoch % plot_freq == 0:\n",
        "                update_plots(history, **live_plot)\n",
        "else:\n",
        "    print('Skipping training')\n",
        "\n",
        "\n",
        "# next, the weights are serialized to distribute the model in this state\n",
        "print('Model weights blob:\\n===')\n",
        "serialized_model = io.BytesIO()\n",
        "torch.save(model['layers'].state_dict(), serialized_model)\n",
        "print(base64.b64encode(serialized_model.getbuffer()).decode('utf-8'))\n",
        "print('===')"
      ],
      "metadata": {
        "id": "O7S91AvWQmeT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Now that we have a trained model at hand, let us draw samples from this, and check the quality.\n",
        "* We plot the drawn samples below. These have regions with smoother, largely correlated fluctuations, in comparison to untrained priors.\n",
        "* Output distribution $Q[\\phi']$ is an approximate one, and can be used for reweighting or resampling.\n",
        "* As we do not know what $1/Z$ is, there is a finite offset between approximate and true actions $- \\log Q[\\phi']$ and $S_{\\text{lattice}}[\\phi']$.\n",
        "\n",
        "**Exercise 8: Find an expression for the offset between approximate and target actions $- \\log Q[\\phi']$ and $S_{\\text{lattice}}[\\phi']$. Use it to correct the approximate action, call it $-\\log Q'[\\phi']$. Plot and compare $-\\log Q'[\\phi']$ against target action $S_{\\text{lattice}}[\\phi']$.**\n",
        "---"
      ],
      "metadata": {
        "id": "RZzOuTk3bRxK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# we draw samples from this trained weight distribution and plot below\n",
        "torch_x, torch_logq = apply_flow_to_prior(prior, layers, batch_size=1024)\n",
        "x = grab(torch_x)\n",
        "\n",
        "fig, ax = plt.subplots(4,4, dpi=125, figsize=(4,4))\n",
        "for i in range(4):\n",
        "    for j in range(4):\n",
        "        ind = i*4 + j\n",
        "        ax[i,j].imshow(np.tanh(x[ind]), vmin=-1, vmax=1, cmap='viridis')\n",
        "        ax[i,j].axes.xaxis.set_visible(False)\n",
        "        ax[i,j].axes.yaxis.set_visible(False)\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "GZd3WmEwSGqr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "References\n",
        "* [Introduction to Normalizing Flows for Lattice Field Theory](https://arxiv.org/abs/2101.08176) by Boyda et. al.\n",
        "* [Learning Lattice Quantum Field Theories with Equivariant Continuous Flows](https://arxiv.org/abs/2207.00283) by Gerdes et. al.\n",
        "* [Quantum Field Theory](https://www.cambridge.org/highereducation/books/quantum-field-theory/718DD037728FB3745F48A40A6D9A8A1C#overview), Mark Srednicki, Cembridge University Press, Chapter 9.\n",
        "* [Quantum Field on a Lattice](https://www.cambridge.org/core/books/quantum-fields-on-a-lattice/4401A88CD232B0AEF1409BF6B260883A), Istvan Montvay and Gernot Munster, Cambridge University Press, Chapters 1, 2."
      ],
      "metadata": {
        "id": "1INQfCC8d9Wm"
      }
    }
  ]
}