Lab_5_ExpFam.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyP7CW5ra3pS12nBUmbBQ9WV",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/USCbiostats/PM520/blob/main/Lab_5_ExpFam.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# It's a Family Affair, or: Exponential Families\n",
        "\n",
        "[Exponential Families](https://en.wikipedia.org/wiki/Exponential_family) (sometimes abbreviated as ExpFam) provide a [succinct characterization](https://en.wikipedia.org/wiki/Exponential_family#Table_of_distributions) of many distributions (e.g., [Normal](https://en.wikipedia.org/wiki/Normal_distribution), [Gamma](https://en.wikipedia.org/wiki/Gamma_distribution), [Poisson](https://en.wikipedia.org/wiki/Poisson_distribution), [Bernoulli](https://en.wikipedia.org/wiki/Bernoulli_distribution), [Wishart](https://en.wikipedia.org/wiki/Wishart_distribution), etc.). We'll take an informal look at their properties and how to perform inference."
      ],
      "metadata": {
        "id": "4R7VI43gT937"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Exponential Families\n",
        "Let $\\eta = [\\eta_1, \\dotsc, \\eta_k]$ be a $k$-vector of parameters, and $x$ be an observation such that $x \\sim f(\\eta)$. We can define its [PDF](https://en.wikipedia.org/wiki/Probability_density_function) (or [PMF](https://en.wikipedia.org/wiki/Probability_mass_function) in case of discrete $x$) as $$f(x | \\eta) = h(x)\\exp(\\eta \\cdot T(x) - A(\\eta)),$$ where $h(x)$ is a *base measure*, $\\eta$ are the *natural parameters*, $T(x)$ are the [*sufficient statistics*](https://en.wikipedia.org/wiki/Sufficient_statistic), and $A(\\eta)$ is the [*log-partition function*](https://en.wikipedia.org/wiki/Partition_function_%28mathematics%29). If $\\eta$ is *finite*, and the [*support*](https://en.wikipedia.org/wiki/Support_%28mathematics%29%23In_probability_and_measure_theory) of $f$ does not depend on the value of $\\eta$, then $f$ can be said to be a member of the [Exponential Families](https://en.wikipedia.org/wiki/Exponential_family).\n",
        "\n",
        "### Example: Normal Distribution\n",
        "Recall if $x \\sim N(\\mu, \\sigma^2)$, then the PDF of $x$ is given by,\n",
        "$$f(x | \\mu, \\sigma^2) = \\frac{1}{\\sqrt{2 \\pi \\sigma^2}} \\exp\\left(-\\frac{(x - \\mu)^2}{2 \\sigma^2}\\right).$$ To see that the two-parameter Normal distribution is a member of the Exponential Families, define $\\eta = [\\frac{\\mu}{\\sigma^2}, -\\frac{1}{2\\sigma^2}]$, $h(x) = \\frac{1}{\\sqrt{2\\pi}}$, $T(x) = [x, x^2]^T$, and $A(\\eta) = \\frac{\\mu^2}{2\\sigma^2} + \\log |\\sigma| = -\\frac{\\eta_1^2}{4\\eta_2} + \\frac{1}{2}\\log|\\frac{1}{2\\eta_2}|$. Placing this all together we have,\n",
        "$$\\begin{align*}\n",
        "f(x | \\eta) &= h(x)\\exp(\\eta \\cdot T(x) - A(\\eta)) \\\\\n",
        "  &= \\frac{1}{\\sqrt{2\\pi}} \\exp(\\eta \\cdot T(x) - A(\\eta)) \\\\\n",
        "  &= \\frac{1}{\\sqrt{2\\pi}} \\exp\\left([\\frac{\\mu}{\\sigma^2}, -\\frac{1}{2\\sigma^2}] \\cdot T(x) - A(\\eta)\\right) \\\\\n",
        "  &= \\frac{1}{\\sqrt{2\\pi}} \\exp\\left([\\frac{\\mu}{\\sigma^2}, -\\frac{1}{2\\sigma^2}] \\cdot [x, x^2]^T - A(\\eta)\\right) \\\\\n",
        "  &= \\frac{1}{\\sqrt{2\\pi}} \\exp\\left(\\frac{\\mu x}{\\sigma^2} -\\frac{x^2}{2\\sigma^2} - A(\\eta)\\right) \\\\\n",
        "  &= \\frac{1}{\\sqrt{2\\pi}} \\exp\\left(\\frac{\\mu x}{\\sigma^2} -\\frac{x^2}{2\\sigma^2} + \\frac{\\eta_1^2}{4\\eta_2} - \\frac{1}{2}\\log\\left|\\frac{1}{2\\eta_2}\\right|\\right) \\\\\n",
        "  &= \\frac{1}{\\sqrt{2\\pi}} \\exp\\left(\\frac{\\mu x}{\\sigma^2} -\\frac{x^2}{2\\sigma^2} - \\frac{\\mu^2}{2\\sigma^2} - \\log |\\sigma|\\right) \\\\\n",
        "  &= \\frac{1}{\\sqrt{2\\pi}} \\exp\\left(-\\frac{(x - \\mu)^2}{2\\sigma^2} - \\log |\\sigma|\\right) \\\\\n",
        "  &= \\frac{1}{\\sqrt{2\\pi}\\sigma} \\exp\\left(-\\frac{(x - \\mu)^2}{2\\sigma^2}\\right) \\\\\n",
        "  &= \\frac{1}{\\sqrt{2\\pi\\sigma^2}} \\exp\\left(-\\frac{(x - \\mu)^2}{2\\sigma^2}\\right) \\\\\n",
        "\\end{align*}$$\n",
        "\n",
        "## Expectations\n",
        "A key property of ExpFam distributions is that one can define the [moments](https://en.wikipedia.org/wiki/Moment_%28mathematics%29) of $T(x)$ from the log-partition function $A(\\eta)$. Two key moments to remember are, $\\mathbb{E}[T(x)] = \\frac{\\partial}{\\partial \\eta} A(\\eta)$ and $\\mathbb{V}[T(x)] = \\frac{\\partial^2}{\\partial \\eta \\partial \\eta^T} A(\\eta)$.\n",
        "\n",
        "In the case of scalar Normal variables, we have $\\mathbb{E}[T(x)] = \\mathbb{E}\\left[[x, x^2]^T\\right] = [\\mu, \\sigma^2 + \\mu^2]^T$ and $\\mathbb{V}[T(x)] = \\begin{bmatrix} \\sigma^2 & 0 \\\\ 0 & 2\\sigma^4 \\end{bmatrix}$.\n",
        "$$\\begin{align*}\n",
        "\\mathbb{E}[T(x)] &= \\frac{\\partial}{\\partial \\eta} A(\\eta) \\\\\n",
        "  &= \\frac{\\partial}{\\partial \\eta} \\left[-\\frac{\\eta_1^2}{4\\eta_2} + \\frac{1}{2}\\log|\\frac{1}{2\\eta_2}|\\right] \\\\\n",
        "  &= -\\frac{\\partial}{\\partial \\eta} \\frac{\\eta_1^2}{4\\eta_2} + \\frac{\\partial}{\\partial \\eta}\\frac{1}{2}\\log|\\frac{1}{2\\eta_2}| \\\\\n",
        "  &= -\\frac{\\partial}{\\partial \\eta} \\frac{\\eta_1^2}{4\\eta_2} - \\frac{\\partial}{\\partial \\eta}\\frac{1}{2}\\log|2\\eta_2| \\\\\n",
        "\\mathbb{E}[T(x)]_1 &= -\\frac{\\partial}{\\partial \\eta_1} \\frac{\\eta_1^2}{4\\eta_2} - \\frac{\\partial}{\\partial \\eta_1}\\frac{1}{2}\\log|2\\eta_2| \\\\\n",
        "  &= - \\frac{2 \\eta_1}{4 \\eta_2} = -\\frac{\\eta_1}{2 \\eta_2} = -\\frac{\\mu / \\sigma^2}{-2 / (2\\sigma^2)} = \\mu \\\\\n",
        "\\mathbb{E}[T(x)]_2 &= -\\frac{\\partial}{\\partial \\eta_2} \\frac{\\eta_1^2}{4\\eta_2} - \\frac{\\partial}{\\partial \\eta_2}\\frac{1}{2}\\log|2\\eta_2| \\\\\n",
        "  &= \\frac{\\eta^2_1}{4\\eta_2^2} - \\frac{1}{2\\eta_2} = \\frac{\\mu^2/\\sigma^4}{1/\\sigma^4} + \\frac{1}{1 / \\sigma^2} = \\mu^2 + \\sigma^2.\n",
        "\\end{align*}$$\n",
        "\n",
        "## MLE Inference for ExpFam\n",
        "Given $x_1, \\dotsc, x_n$ we assume that $x_i \\sim f(\\eta)$ independently and identically distributed under some ExpFam dist $f$ with natural parameters $\\eta$. Our log-likelihood is then given by,\n",
        "$$ \\ell(\\eta) = \\sum_{i=1}^n \\log f(x_i | \\eta).$$ To identify the\n",
        "maximum likelihood estimates we first compute the gradient of $\\ell$ wrt $\\eta$, which is given by\n",
        "$$\\begin{align*}\n",
        "\\nabla \\ell(\\eta) &= \\sum_{i=1}^n \\nabla \\log f(x_i | \\eta) \\\\\n",
        "  &= \\sum_{i=1}^n \\nabla \\log \\left[ h(x_i)\\exp(\\eta \\cdot T(x_i) - A(\\eta))\\right] \\\\\n",
        "  &= \\sum_{i=1}^n \\nabla \\left[\\log h(x_i) + \\log \\exp(\\eta \\cdot T(x_i) - A(\\eta))\\right] \\\\\n",
        "  &= \\sum_{i=1}^n \\nabla \\log h(x_i) + \\nabla \\log \\exp(\\eta \\cdot T(x_i) - A(\\eta)) \\\\\n",
        "  &= \\sum_{i=1}^n \\nabla [\\eta \\cdot T(x_i)] - \\nabla A(\\eta) \\\\\n",
        "  &= \\sum_{i=1}^n T(x_i) - \\mathbb{E}[T(x_i)],\n",
        "\\end{align*}$$\n",
        "where we used the fact that  $\\nabla A(\\eta) = \\mathbb{E}[T(x)]$. Setting this to zero and solving implies, we would like to find values $\\eta$ such that $\\sum_i T(x_i) = \\sum_i \\mathbb{E}[T(x_i)] = n \\mathbb{E}[T(X)] ⇒ \\frac{1}{n} \\sum_i T(x_i) = \\mathbb{E}[T(X)]$! In other words, MLE under ExpFam seeks values of $\\eta$ that *match the empirical mean of the observed sufficient statistics to their expectation*!"
      ],
      "metadata": {
        "id": "tz3bA7fUcApH"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## It's an Equinox for all seasons\n",
        "TBD: Information on `equinox`, (functional) objective oriented programming, abstract base classes. More documentation on `Normal` class."
      ],
      "metadata": {
        "id": "4cm1S74tRxqA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install equinox"
      ],
      "metadata": {
        "id": "ykn-OzCv3xjQ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2cf5ee4e-3e91-44fa-ce38-d9d4163cc56d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting equinox\n",
            "  Downloading equinox-0.11.11-py3-none-any.whl.metadata (18 kB)\n",
            "Collecting jax>=0.4.38 (from equinox)\n",
            "  Downloading jax-0.5.0-py3-none-any.whl.metadata (22 kB)\n",
            "Collecting jaxtyping>=0.2.20 (from equinox)\n",
            "  Downloading jaxtyping-0.2.38-py3-none-any.whl.metadata (6.6 kB)\n",
            "Requirement already satisfied: typing-extensions>=4.5.0 in /usr/local/lib/python3.11/dist-packages (from equinox) (4.12.2)\n",
            "Collecting jaxlib<=0.5.0,>=0.5.0 (from jax>=0.4.38->equinox)\n",
            "  Downloading jaxlib-0.5.0-cp311-cp311-manylinux2014_x86_64.whl.metadata (978 bytes)\n",
            "Requirement already satisfied: ml_dtypes>=0.4.0 in /usr/local/lib/python3.11/dist-packages (from jax>=0.4.38->equinox) (0.4.1)\n",
            "Requirement already satisfied: numpy>=1.25 in /usr/local/lib/python3.11/dist-packages (from jax>=0.4.38->equinox) (1.26.4)\n",
            "Requirement already satisfied: opt_einsum in /usr/local/lib/python3.11/dist-packages (from jax>=0.4.38->equinox) (3.4.0)\n",
            "Requirement already satisfied: scipy>=1.11.1 in /usr/local/lib/python3.11/dist-packages (from jax>=0.4.38->equinox) (1.13.1)\n",
            "Collecting wadler-lindig>=0.1.3 (from jaxtyping>=0.2.20->equinox)\n",
            "  Downloading wadler_lindig-0.1.3-py3-none-any.whl.metadata (17 kB)\n",
            "Downloading equinox-0.11.11-py3-none-any.whl (179 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m179.2/179.2 kB\u001b[0m \u001b[31m5.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading jax-0.5.0-py3-none-any.whl (2.3 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.3/2.3 MB\u001b[0m \u001b[31m29.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading jaxtyping-0.2.38-py3-none-any.whl (56 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m56.4/56.4 kB\u001b[0m \u001b[31m3.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading jaxlib-0.5.0-cp311-cp311-manylinux2014_x86_64.whl (102.0 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m102.0/102.0 MB\u001b[0m \u001b[31m5.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading wadler_lindig-0.1.3-py3-none-any.whl (20 kB)\n",
            "Installing collected packages: wadler-lindig, jaxtyping, jaxlib, jax, equinox\n",
            "  Attempting uninstall: jaxlib\n",
            "    Found existing installation: jaxlib 0.4.33\n",
            "    Uninstalling jaxlib-0.4.33:\n",
            "      Successfully uninstalled jaxlib-0.4.33\n",
            "  Attempting uninstall: jax\n",
            "    Found existing installation: jax 0.4.33\n",
            "    Uninstalling jax-0.4.33:\n",
            "      Successfully uninstalled jax-0.4.33\n",
            "Successfully installed equinox-0.11.11 jax-0.5.0 jaxlib-0.5.0 jaxtyping-0.2.38 wadler-lindig-0.1.3\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from abc import abstractmethod\n",
        "\n",
        "import equinox as eqx\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "from jaxtyping import Array, ArrayLike\n",
        "\n",
        "\n",
        "class AbstractExpFam(eqx.Module):\n",
        "  \"\"\"\n",
        "  Simple base class for Exponential Families. Will provide means to compute\n",
        "  sufficient statistics and evaluat the loglikelihood of downstream\n",
        "  implementations.\n",
        "  \"\"\"\n",
        "\n",
        "\n",
        "  @abstractmethod\n",
        "  def base_measure(self, x: ArrayLike) -> Array:\n",
        "    \"\"\"\n",
        "    Computes the base measure for an implementation of ExpFam.\n",
        "\n",
        "    x: ArrayLike, the observations.\n",
        "\n",
        "    Returns:\n",
        "      The base measure for an implementation of ExpFam.\n",
        "    \"\"\"\n",
        "    ...\n",
        "\n",
        "  @abstractmethod\n",
        "  def sufficient_statistics(self, x: ArrayLike) -> Array:\n",
        "    \"\"\"\n",
        "    Computes the sufficient statistics (i.e. $T(x)$) for an implementation\n",
        "    of ExpFam.\n",
        "\n",
        "    x: ArrayLike, the observations.\n",
        "\n",
        "    Returns:\n",
        "      The sufficient statistics for each observation under an implementation\n",
        "      of ExpFam.\n",
        "    \"\"\"\n",
        "    ...\n",
        "\n",
        "  @abstractmethod\n",
        "  def log_partition(self, eta: ArrayLike) -> Array:\n",
        "    \"\"\"\n",
        "    Computes the log partition function (i.e. $A(\\eta)$) for an implementation\n",
        "    of ExpFam with natural parameters $\\eta$.\n",
        "\n",
        "    eta: ArrayLike, the natural parameters to evaluate under the log partition.\n",
        "\n",
        "    Returns:\n",
        "      The value of the log partition function for each observation under an\n",
        "      implementation of ExpFam.\n",
        "    \"\"\"\n",
        "    ...\n",
        "\n",
        "  def loglikelihood(self, eta: ArrayLike, x: ArrayLike) -> Array:\n",
        "    \"\"\"\n",
        "    Computes the log likelihood for each observation $x$ under an implementation\n",
        "    of ExpFam with natural parameters $\\eta$.\n",
        "\n",
        "    eta: ArrayLike, the natural parameters to evaluate under the log likelihood.\n",
        "    x: ArrayLike, the observations.\n",
        "\n",
        "    Returns:\n",
        "      The log likelihood for each observation under an implementation of ExpFam.\n",
        "    \"\"\"\n",
        "    t_x = self.sufficient_statistics(x)\n",
        "    log_h_x = jnp.log(self.base_measure(x))\n",
        "    log_eta = self.log_partition(eta)\n",
        "    if eta.ndim == 0:\n",
        "      inner = t_x * eta\n",
        "    else:\n",
        "      inner = t_x @ eta\n",
        "\n",
        "    return inner - log_eta + log_h_x\n",
        "\n",
        "  def __call__(self, eta: ArrayLike, x: ArrayLike) -> Array:\n",
        "    return self.loglikelihood(eta, x)\n",
        "\n",
        "  def fit_mle(self, x: ArrayLike) -> Array:\n",
        "    \"\"\"\n",
        "    Maximizes the log likelihood for each observation under an implementation\n",
        "    of ExpFam with natural parameters.\n",
        "\n",
        "    x: ArrayLike, the observations.\n",
        "\n",
        "    Returns:\n",
        "      The moment parameters that maximize log likelihood.\n",
        "    \"\"\"\n",
        "    ...\n",
        "\n",
        "\n",
        "class Normal(AbstractExpFam):\n",
        "\n",
        "  def base_measure(self, x: ArrayLike) -> Array:\n",
        "    return 1. / jnp.sqrt(2 * jnp.pi)\n",
        "\n",
        "  def sufficient_statistics(self, x: ArrayLike) -> Array:\n",
        "    x_sq = x ** 2\n",
        "    return jnp.concatenate((x[:, jnp.newaxis], x_sq[:, jnp.newaxis]), axis=1)\n",
        "\n",
        "  def log_partition(self, eta: ArrayLike) -> Array:\n",
        "    eta_1, eta_2 = eta\n",
        "    term1 = -(eta_1**2 / (4 * eta_2))\n",
        "    term2 = -0.5 * jnp.log(- 2 * eta_2)\n",
        "    return term1 + term2"
      ],
      "metadata": {
        "id": "ppu0j4ydwTnn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Great! Having implemented our objects to perform MLE in a general setting, with a specific example of normal distributions, let's perform some sanity checks before proceeding with inference."
      ],
      "metadata": {
        "id": "If1EgDTJSFjA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import jax.random as rdm\n",
        "\n",
        "# initialize our PRNG state\n",
        "seed = 0\n",
        "key = rdm.PRNGKey(seed)\n",
        "\n",
        "# define some parameter values for N(mu, sigma2)\n",
        "mu = 5.0\n",
        "sigma_sq = 10.0\n",
        "\n",
        "# generate random data\n",
        "N = 10\n",
        "key, x_key = rdm.split(key)\n",
        "obs = mu + jnp.sqrt(sigma_sq) * rdm.normal(x_key, shape=(N,))\n",
        "\n",
        "# transform parameters to natural parameters and create some other guess at eta\n",
        "eta_true = jnp.array([mu / sigma_sq, - 1. / (2 * sigma_sq)])\n",
        "eta_guess = jnp.array([-9, -100.])\n",
        "\n",
        "# create an instance of Normal distribution using our implementation above\n",
        "model = Normal()\n",
        "\n",
        "# calculate the sum of log likelihood\n",
        "# equivalent to `jnp.sum(model.loglikelihood(eta_true))`\n",
        "sum_ll_true = jnp.sum(model(eta_true, obs))\n",
        "sum_ll_guess = jnp.sum(model(eta_guess, obs))\n",
        "print(f\"true logl $\\ell$({eta_true}) = {sum_ll_true} | guess logl $\\ell$({eta_guess}) = {sum_ll_guess}\")\n",
        "\n",
        "# sanity check against jax loglikelihood\n",
        "import jax.scipy.stats as stats\n",
        "jax_ll_true = jnp.sum(stats.norm.logpdf(obs, mu, jnp.sqrt(sigma_sq)))\n",
        "print(f\"Our $\\ell$({eta_true}) = {sum_ll_true} | JAX logl $\\ell$({eta_true}) = {jax_ll_true}\")\n",
        "\n",
        "fitted = model.fit_mle(obs)\n",
        "print(f\"Our MLE moment-based parameters are {fitted}\")\n",
        "fitted_con = jnp.array([fitted[0], fitted[1] - fitted[0]**2])\n",
        "print(f\"Our MLE canonical parameters are {fitted_con}\")\n",
        "print(f\"Our true canonical parameters are {jnp.array([mu, sigma_sq])}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uOmELPydTaY-",
        "outputId": "2b6abf52-524e-4f2c-f33e-3172a03c2f7a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "true logl $\\ell$([ 0.5  -0.05]) = -2621.6142578125 | guess logl $\\ell$([  -9. -100.]) = -3551774.5\n",
            "Our $\\ell$([ 0.5  -0.05]) = -2621.6142578125 | JAX logl $\\ell$([ 0.5  -0.05]) = -2621.6142578125\n",
            "Our MLE moment-based parameters are [ 4.906379 35.091446]\n",
            "Our MLE canonical parameters are [ 4.906379 11.018888]\n",
            "Our true canonical parameters are [ 5. 10.]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from jax.scipy.special import gamma\n",
        "\n",
        "class Poisson(AbstractExpFam):\n",
        "\n",
        "  def base_measure(self, x: ArrayLike) -> Array:\n",
        "    ...\n",
        "\n",
        "  def sufficient_statistics(self, x: ArrayLike) -> Array:\n",
        "    ...\n",
        "\n",
        "  def log_partition(self, eta: ArrayLike) -> Array:\n",
        "    ..."
      ],
      "metadata": {
        "id": "k4gT7TI-zKF4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# initialize our PRNG state\n",
        "seed = 0\n",
        "key = rdm.PRNGKey(seed)\n",
        "\n",
        "# define some parameter values for N(mu, sigma2)\n",
        "rate = 5.0\n",
        "\n",
        "# generate random data\n",
        "N = 10\n",
        "key, x_key = rdm.split(key)\n",
        "obs = rdm.poisson(x_key, rate, shape=(N,))\n",
        "\n",
        "# transform parameters to natural parameters and create some other guess at eta\n",
        "eta_true = jnp.log(rate)\n",
        "eta_guess = jnp.log(100.)\n",
        "\n",
        "# create an instance of Poisson distribution using our implementation above\n",
        "model = Poisson()\n",
        "\n",
        "# calculate the sum of log likelihood\n",
        "# equivalent to `jnp.sum(model.loglikelihood(eta_true))`\n",
        "sum_ll_true = jnp.sum(model(eta_true, obs))\n",
        "sum_ll_guess = jnp.sum(model(eta_guess, obs))\n",
        "print(f\"true logl $\\ell$({eta_true}) = {sum_ll_true} | guess logl $\\ell$({eta_guess}) = {sum_ll_guess}\")\n",
        "\n",
        "# sanity check against jax loglikelihood\n",
        "import jax.scipy.stats as stats\n",
        "jax_ll_true = jnp.sum(stats.poisson.logpmf(obs, rate))\n",
        "print(f\"Our $\\ell$({eta_true}) = {sum_ll_true} | JAX logl $\\ell$({eta_true}) = {jax_ll_true}\")\n",
        "\n",
        "fitted = model.fit_mle(obs)\n",
        "print(f\"Our MLE moment-based parameters are {fitted}\")\n",
        "print(f\"Our true canonical parameters are {rate}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Unr2ullt3pbe",
        "outputId": "207f6ac2-d2c0-4d0a-8bd1-705c8b6ea110"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "true logl $\\ell$(1.6094379425048828) = -21.40349578857422 | guess logl $\\ell$(4.605170249938965) = -806.6383056640625\n",
            "Our $\\ell$(1.6094379425048828) = -21.40349578857422 | JAX logl $\\ell$(1.6094379425048828) = -21.40349578857422\n",
            "Our MLE moment-based parameters are 5.5\n",
            "Our true canonical parameters are 5.0\n"
          ]
        }
      ]
    }
  ]
}