Ans_4.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c6blGlxZRXV6"
      },
      "outputs": [],
      "source": [
        "import pymc3 as pm\n",
        "import numpy as np\n",
        "import jax.numpy as jnp\n",
        "\n",
        "from blackjax import nuts\n",
        "# import blackjax.stan_warmup as stan_warmup\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VCWDNaZsXLLV",
        "outputId": "d6b7f1b8-4ee7-48b4-c512-105426638396"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collecting blackjax\n",
            "  Downloading blackjax-0.4.0-py3-none-any.whl (74 kB)\n",
            "\u001b[?25l\r\u001b[K     |████▍                           | 10 kB 34.4 MB/s eta 0:00:01\r\u001b[K     |████████▉                       | 20 kB 35.0 MB/s eta 0:00:01\r\u001b[K     |█████████████▎                  | 30 kB 27.5 MB/s eta 0:00:01\r\u001b[K     |█████████████████▋              | 40 kB 19.7 MB/s eta 0:00:01\r\u001b[K     |██████████████████████          | 51 kB 7.7 MB/s eta 0:00:01\r\u001b[K     |██████████████████████████▌     | 61 kB 9.1 MB/s eta 0:00:01\r\u001b[K     |██████████████████████████████▉ | 71 kB 8.8 MB/s eta 0:00:01\r\u001b[K     |████████████████████████████████| 74 kB 3.1 MB/s \n",
            "\u001b[?25hInstalling collected packages: blackjax\n",
            "Successfully installed blackjax-0.4.0\n"
          ]
        }
      ],
      "source": [
        "pip install blackjax\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ohho3IzHXSx0"
      },
      "outputs": [],
      "source": [
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import jax.scipy.stats as stats\n",
        "import numpy as np\n",
        "\n",
        "import blackjax\n",
        "\n",
        "observed = np.random.normal(10, 20, size=1_000)\n",
        "def logprob_fn(x):\n",
        "  logpdf = stats.norm.logpdf(observed, x[\"loc\"], x[\"scale\"])\n",
        "  return jnp.sum(logpdf)\n",
        "\n",
        "# Build the kernel\n",
        "step_size = 1e-3\n",
        "inverse_mass_matrix = jnp.array([1., 1.])\n",
        "nuts = blackjax.nuts(logprob_fn, step_size, inverse_mass_matrix)\n",
        "\n",
        "# Initialize the state\n",
        "initial_position = {\"loc\": 1., \"scale\": 2.}\n",
        "state = nuts.init(initial_position)\n",
        "\n",
        "# Iterate\n",
        "rng_key = jax.random.PRNGKey(0)\n",
        "for _ in range(100):\n",
        "    _, rng_key = jax.random.split(rng_key)\n",
        "    state, _ = nuts.step(rng_key, state)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hJbBa7j18JA9"
      },
      "source": [
        "Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "haM-fjtY8G7o"
      },
      "outputs": [],
      "source": [
        "J = 8\n",
        "y = np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0])\n",
        "sigma = np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b_FttvSc1PAh"
      },
      "source": [
        "Modeling"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dtCl0tls1R9Z"
      },
      "outputs": [],
      "source": [
        "with pm.Model() as model:\n",
        "\n",
        "    mu = pm.Normal(\"mu\", mu=0.0, sigma=10.0)\n",
        "    tau = pm.HalfCauchy(\"tau\", 5.0)\n",
        "\n",
        "    theta = pm.Normal(\"theta\", mu=0, sigma=1, shape=J)\n",
        "    theta_1 = mu + tau * theta\n",
        "    obs = pm.Normal(\"obs\", mu=theta_1, sigma=sigma, shape=J, observed=y)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6xv1kokr0oue"
      },
      "source": [
        "Configuring the model for BlackJAX"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JGYL-lPs0zbs"
      },
      "outputs": [],
      "source": [
        "from theano.graph.fg import FunctionGraph\n",
        "from theano.link.jax.jax_dispatch import jax_funcify\n",
        "\n",
        "seed = jax.random.PRNGKey(1234)\n",
        "chains = 1\n",
        "\n",
        "# Get the FunctionGraph of the model.\n",
        "fgraph = FunctionGraph(model.free_RVs, [model.logpt])\n",
        "\n",
        "# Jax funcify builds Jax variant of the FunctionGraph.\n",
        "fns = jax_funcify(fgraph)\n",
        "logp_fn_jax = fns[0]\n",
        "\n",
        "# Now we build a Jax variant of the initial state/inputs to the model.\n",
        "rv_names = [rv.name for rv in model.free_RVs]\n",
        "init_state = [model.test_point[rv_name] for rv_name in rv_names]\n",
        "init_state_batched = jax.tree_map(\n",
        "    lambda x: np.repeat(x[None, ...], chains, axis=0), init_state\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PSu8Yys202nA"
      },
      "outputs": [],
      "source": [
        "# Then we transform the Jaxified input and FunctionGraph to a BlackJax NUTS sampler\n",
        "potential = lambda x: -logp_fn_jax(*x)\n",
        "initial_position = init_state\n",
        "initial_state = nuts.new_state(initial_position, potential)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8UtrZ5ji1GH9"
      },
      "source": [
        "Sampling"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kptFFWO01An1"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "\n",
        "kernel_factory = lambda step_size, inverse_mass_matrix: nuts.kernel(\n",
        "    potential, step_size, inverse_mass_matrix\n",
        ")\n",
        "\n",
        "last_state, (step_size, inverse_mass_matrix), _ = stan_warmup.run(\n",
        "    seed, kernel_factory, initial_state, 1000\n",
        ")\n",
        "\n",
        "\n",
        "def inference_loop(rng_key, kernel, initial_state, num_samples):\n",
        "    def one_step(state, rng_key):\n",
        "        state, info = kernel(rng_key, state)\n",
        "        return state, (state, info)\n",
        "\n",
        "    keys = jax.random.split(rng_key, num_samples)\n",
        "    _, (states, infos) = jax.lax.scan(one_step, initial_state, keys)\n",
        "\n",
        "    return states, infos\n",
        "\n",
        "# Build the kernel using the step size and inverse mass matrix returned from the window adaptation\n",
        "kernel = kernel_factory(step_size, inverse_mass_matrix)\n",
        "\n",
        "# Sample from the posterior distribution\n",
        "states, infos = inference_loop(seed, kernel, last_state, 50_000)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ne7Z-_pD_NZu"
      },
      "source": [
        "Resources\n",
        "\n",
        "\n",
        "*   https://blackjax-devs.github.io/blackjax/sampling.html#nuts\n",
        "  \n",
        "*   https://www.kaggle.com/code/s903124/numpyro-speed-benchmark-gpu\n",
        "\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Ans_4.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}