Lab_3_Optimization_PtI.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyN+9Dh/qsCaBlszqM1dDDQa",
      "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_3_Optimization_PtI.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Move on Up, or: Maximum likelihood Estimation & Optimization Pt I\n",
        "\n",
        "TBD: move notes from slides to here.\n"
      ],
      "metadata": {
        "id": "q1bg914lpQpN"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VAK5-ADoNq_I"
      },
      "outputs": [],
      "source": [
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import jax.random as rdm"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## MLE for iid Normal data\n",
        "Let $x_1, \\dotsc, x_n \\overset{\\mathrm{iid}}{\\sim} \\mathcal{N}(\\mu, \\sigma^2)$ where $\\mathcal{N}(\\mu, \\sigma^2)$ refers to the [Normal distribution](https://en.wikipedia.org/wiki/Normal_distribution) with mean parameter $\\mu$ and variance parameter $\\sigma^2$. The likelihood of our data is given by,\n",
        "$$\\begin{align*}\n",
        "\\mathcal{L}(\\mu, \\sigma^2 | x_1, \\dots, x_n) &=\n",
        "  \\prod_{i=1}^n \\mathcal{N}(x_i | \\mu, \\sigma^2) \\\\\n",
        "  &= \\prod_{i=1}^n \\frac{1}{\\sqrt{2\\pi \\sigma^2}} \\exp\\left(-\\frac{(x_i - \\mu)^2}{2\\sigma^2}\\right)\\\\\n",
        "  &= \\left(\\frac{1}{\\sqrt{2\\pi\\sigma^2}}\\right)^n \\exp\\left(-\\frac{1}{2\\sigma^2} \\sum_{i=1}^n (x_i - \\mu)^2\\right).\n",
        "\\end{align*}\\\\\n",
        "$$\n",
        "Thus, our _log_-likelihood is given by,\n",
        "The likelihood of our data is given by,\n",
        "$$\\begin{align*}\n",
        "\\ell(\\mu, \\sigma^2 | x_1, \\dots, x_n) &=\n",
        "  \\log \\left[\\left(\\frac{1}{\\sqrt{2\\pi\\sigma^2}}\\right)^n \\exp\\left(-\\frac{1}{2\\sigma^2} \\sum_{i=1}^n (x_i - \\mu)^2\\right)\\right]\\\\\n",
        "  &= -\\frac{n}{2}\\log(2\\pi\\sigma^2) - \\frac{1}{2\\sigma^2} \\sum_{i=1}^n (x_i - \\mu)^2.\n",
        "\\end{align*}\\\\\n",
        "$$"
      ],
      "metadata": {
        "id": "MbaYn27uwi28"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def norm_rv(key, n: int, mu: float, sigma_sq: float):\n",
        "  \"\"\"\n",
        "  Samples $n$ observations from $x_i \\sim N(\\mu, \\sigma^2)$\n",
        "\n",
        "  n: the number of observations\n",
        "  mu: the mean parameter\n",
        "  sigma_sq: the variance parameter\n",
        "\n",
        "  returns: x, Array of observations\n",
        "  \"\"\"\n",
        "  x = mu + jnp.sqrt(sigma_sq) * rdm.normal(key, shape=(n,))\n",
        "  return x\n",
        "\n",
        "\n",
        "def norm_mle(x):\n",
        "  \"\"\"\n",
        "  Computes $\\hat{\\mu}_{MLE}$ and $\\hat{\\sigma^2}_{MLE}$.\n",
        "\n",
        "  x: Array of observations\n",
        "\n",
        "  returns: Tuple of $\\hat{\\mu}_{MLE}$ and $\\hat{\\sigma^2}_{MLE}$.\n",
        "  \"\"\"\n",
        "  mu_hat = pass\n",
        "  ssq_hat = pass\n",
        "\n",
        "  return mu_hat, ssq_hat\n",
        "\n",
        "seed = 0\n",
        "key = rdm.PRNGKey(seed)\n",
        "key, x_key = rdm.split(key)\n",
        "\n",
        "N = 500\n",
        "\n",
        "mu = 58.\n",
        "sigma_sq = 100.\n",
        "x = norm_rv(x_key, N, mu, sigma_sq)\n",
        "#print(f\"x = {x}\")\n",
        "mu_hat, ssq_hat = norm_mle(x)\n",
        "print(f\"MLE[\\mu, \\sigma^2] = {mu_hat}, {ssq_hat}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wWoIiwnVwn6O",
        "outputId": "400bde50-5cf3-4de8-8640-7bd73379f6b1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "MLE[\\mu, \\sigma^2] = 58.59890365600586, 98.35617065429688\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def sq_diff(param, estimate):\n",
        "  return (param - estimate) ** 2\n",
        "\n",
        "mu = 58.\n",
        "sigma_sq = 10.\n",
        "for N in [50, 100, 1000, 10000]:\n",
        "  key, x_key = rdm.split(key)\n",
        "  # generate N observations\n",
        "  x_n = norm_rv(x_key, N, mu, sigma_sq)\n",
        "  # estimate mu, and sigma_sq\n",
        "  mu_hat, ssq_hat = norm_mle(x_n)\n",
        "  # compute the sq-diff for both and report\n",
        "  mu_err = sq_diff(mu, mu_hat)\n",
        "  ssq_err = sq_diff(sigma_sq, ssq_hat)\n",
        "  print(f\"MSE[{N} | mu, sigma^2] = {mu_err}, {ssq_err}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5Z8awgX6GXXv",
        "outputId": "67b111ee-2e05-415a-a097-c278e622ea9d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "MSE[50 | mu, sigma^2] = 1.0336028337478638, 6.931415557861328\n",
            "MSE[100 | mu, sigma^2] = 0.11436349898576736, 3.418867826461792\n",
            "MSE[1000 | mu, sigma^2] = 1.0024814400821924e-05, 0.01468171738088131\n",
            "MSE[10000 | mu, sigma^2] = 0.004498897586017847, 0.015031831339001656\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## MLE for iid Exponential data\n",
        "TBD: Add notes for Exponential PDF and MLE estimator"
      ],
      "metadata": {
        "id": "RTHf96slw__S"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def exp_rv(key, n: int, rate: float):\n",
        "  \"\"\"\n",
        "  Samples $n$ observations from $x_i \\sim Exp(\\lambda)$\n",
        "\n",
        "  n: the number of observations\n",
        "  rate: the $\\lambda$ parameter\n",
        "\n",
        "  returns: x, Array of observations\n",
        "  \"\"\"\n",
        "  mean = 1 / rate\n",
        "  x = mean * rdm.exponential(key, shape=(n,))\n",
        "  return x\n",
        "\n",
        "\n",
        "def exp_mle(x):\n",
        "  \"\"\"\n",
        "  Computes $\\hat{\\lambda}_{MLE}$.\n",
        "\n",
        "  x: Array of observations\n",
        "\n",
        "  returns: $\\hat{\\lambda}_{MLE}$.\n",
        "  \"\"\"\n",
        "  rate_hat = pass\n",
        "  return rate_hat\n",
        "\n",
        "key, x_key = rdm.split(key)\n",
        "N = 100\n",
        "rate = 1 / 500.\n",
        "x = exp_rv(x_key, N, rate)\n",
        "print(f\"x = {x}\")\n",
        "rate_hat = exp_mle(x)\n",
        "print(f\"MLE[\\lambda = {rate}] = {rate_hat}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "y3kQz53WA8YS",
        "outputId": "08a31e4d-46c3-480c-b52d-d19730b40eac"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "x = [4.35397430e+02 4.72187927e+02 1.38809436e+03 2.51057831e+02\n",
            " 7.86665115e+01 5.10070496e+02 6.66986755e+02 5.30637817e+02\n",
            " 6.95734985e+02 3.18214294e+02 5.19751167e+01 2.74846916e+01\n",
            " 3.22565138e-01 4.53188721e+02 6.49053101e+02 6.87411118e+01\n",
            " 3.77136902e+02 4.52345795e+02 8.62563843e+02 4.71622925e+02\n",
            " 1.00533356e+03 6.46890991e+02 2.95587494e+02 1.70733911e+03\n",
            " 8.86031799e+01 3.67736015e+01 1.70877844e+03 1.35494446e+02\n",
            " 6.13494530e+01 6.64146118e+02 2.19395728e+03 1.05528516e+03\n",
            " 9.19410034e+02 3.70226440e+01 2.18483673e+02 5.09372223e+02\n",
            " 2.02266724e+02 4.67761505e+02 1.08360504e+02 2.41796265e+03\n",
            " 4.80211212e+02 9.36368164e+02 5.30847473e+02 2.47510696e+02\n",
            " 3.98499237e+02 2.79448181e+02 1.04075806e+03 9.33896729e+02\n",
            " 3.96697357e+02 4.65119263e+02 7.25767761e+02 9.85023422e+01\n",
            " 5.21523010e+02 3.69823975e+02 1.16978165e+02 2.81428162e+02\n",
            " 1.20678604e+02 1.87231949e+02 1.82771103e+02 9.60759338e+02\n",
            " 7.37772095e+02 6.62937164e+01 3.19520782e+02 2.58933685e+02\n",
            " 1.22909531e+02 1.70570886e+03 3.33904846e+02 1.63073151e+02\n",
            " 7.69578674e+02 1.76380432e+03 8.24010742e+02 3.06480988e+02\n",
            " 2.07993018e+03 2.22385681e+02 2.11218094e+02 1.65863098e+02\n",
            " 1.19975000e+03 1.49656693e+02 9.81190918e+02 1.03857812e+03\n",
            " 2.94098389e+02 8.39220886e+01 1.84563950e+02 2.52515234e+03\n",
            " 3.35116913e+02 8.14175781e+02 3.83683411e+02 1.13784680e+03\n",
            " 3.77938690e+02 5.28888306e+02 2.28814484e+02 1.19466541e+03\n",
            " 4.82957916e+02 1.34612595e+02 3.37373169e+03 2.85590897e+01\n",
            " 4.15220612e+02 2.25339874e+02 5.76369019e+02 6.54426819e+02]\n",
            "MLE[\\lambda = 0.002] = 0.0016689717303961515\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "rate = 1 / 50.\n",
        "for N in [50, 100, 1000, 10000]:\n",
        "  key, x_key = rdm.split(key)\n",
        "  # generate N observations\n",
        "  x_n = exp_rv(x_key, N, rate)\n",
        "  # estimate rate\n",
        "  rate_hat = exp_mle(x_n)\n",
        "  # compute the sq-diff for rate\n",
        "  rate_err = sq_diff(rate, rate_hat)\n",
        "  print(f\"MSE[{N} | \\lambda = {rate}] = {rate_err}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-amDd6YTLuNI",
        "outputId": "5c17897c-4ba8-42e9-8831-424244458156"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "MSE[50 | \\lambda = 0.02] = 9.133363164437469e-06\n",
            "MSE[100 | \\lambda = 0.02] = 6.241853043320589e-07\n",
            "MSE[1000 | \\lambda = 0.02] = 2.641813523496239e-07\n",
            "MSE[10000 | \\lambda = 0.02] = 4.547181120528876e-08\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Gradient descent\n",
        "[Gradient descent](https://en.wikipedia.org/wiki/Gradient_descent) seeks to iteratively optimize a function $f(\\beta)$ by taking steps in the steepest direction,\n",
        "$$ \\hat{\\beta} = \\beta_t - \\rho_t \\nabla f(\\beta_t),$$\n",
        "where that direction is provided by the [gradient](https://en.wikipedia.org/wiki/Gradient) of (f).\n",
        "\n",
        "A helpful way to recast gradient descent is that we seek to perform a series of _local_ optimizations,\n",
        "\n",
        "$$\\hat{\\beta} = \\min_\\beta \\nabla f(\\beta_t)^T \\beta + \\frac{1}{2\\rho_t}\\|\\beta - \\beta_t\\|_2^2.$$\n",
        "\n",
        "To see how these are equivalent let's solve the local problem. but using inner product notation,\n",
        "$$m(\\beta) = \\nabla f(\\beta_t)^T \\beta + \\frac{1}{2\\rho_t} (\\beta - \\beta_t)^T(\\beta - \\beta_t).$$\n",
        "Now, using calculus again,\n",
        "$$\\begin{align*}\n",
        "\\nabla m(\\beta) &= \\nabla [ \\nabla f(\\beta_t)^T \\beta + \\frac{1}{2\\rho_t} (\\beta - \\beta_t)^T(\\beta - \\beta_t)] \\\\\n",
        "&= \\nabla [\\nabla f(\\beta_t)^T \\beta] + \\frac{1}{2\\rho_t} \\nabla [(\\beta - \\beta_t)^T(\\beta - \\beta_t)] \\\\\n",
        "&= \\nabla f(\\beta_t) + \\frac{1}{\\rho_t}(\\beta - \\beta_t) \\Rightarrow \\\\\n",
        "\\hat{\\beta} &= \\beta_t - \\rho_t \\nabla f(\\beta_t).\n",
        "\\end{align*}\n",
        "$$\n",
        "\n",
        "Neat! However, notice that the original local objective can be thought of as minimizing the directional derivative, but with a distance penalty, where that distance is defined by the geometry of the parameter space.\n",
        "\n",
        "$$\\hat{\\beta} = \\min_\\beta \\nabla f(\\beta_t)^T \\beta + \\frac{1}{2\\rho_t}\\text{dist}(\\beta, \\beta_t).$$\n",
        "\n",
        "When the natural geometry is $\\mathbb{R}^p$ then $\\text{dist}(\\cdot) = \\| \\cdot \\|_2^2$, however there are many  geometries that can describe the natural parameter space (for future class 😉)"
      ],
      "metadata": {
        "id": "I5mFAyAINs-B"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def sim_linear_reg(key, N, P, r2=0.5):\n",
        "  key, x_key = rdm.split(key)\n",
        "  X = rdm.normal(x_key, shape=(N, P))\n",
        "\n",
        "  key, b_key = rdm.split(key)\n",
        "  beta = rdm.normal(b_key, shape=(P,))\n",
        "\n",
        "  # g = jnp.dot(X, beta)\n",
        "  g = X @ beta\n",
        "  s2g = jnp.var(g)\n",
        "\n",
        "  # back out what s2e is, such that s2g / (s2g + s2e) == h2\n",
        "  s2e = (1 - r2) / r2 * s2g\n",
        "  key, y_key = rdm.split(key)\n",
        "\n",
        "  # add env noise to g, but scale such that var(e) == s2e\n",
        "  y = g + jnp.sqrt(s2e) * rdm.normal(y_key, shape=(N,))\n",
        "  return y, X, beta\n",
        "\n",
        "key, sim_key = rdm.split(key)\n",
        "\n",
        "N = 1000\n",
        "P = 5\n",
        "y, X, beta = sim_linear_reg(sim_key, N, P)\n",
        "\n",
        "def linreg_loss(beta_hat, y, X):\n",
        "  pass\n",
        "\n",
        "def gradient(beta_hat, y, X):\n",
        "  pass\n",
        "\n",
        "step_size = 1 / N\n",
        "diff = 10.\n",
        "last_loss = 1000.\n",
        "idx = 0\n",
        "beta_hat = jnp.zeros((P,))\n",
        "# while delta in loss is large, continue\n",
        "print(f\"true beta = {beta}\")\n",
        "while jnp.fabs(diff) > 1e-3:\n",
        "\n",
        "  # take a step in the direction of the gradient using step_size\n",
        "  beta_hat = beta_hat - step_size * gradient(beta_hat, y, X)\n",
        "\n",
        "  # update our current loss and compute delta\n",
        "  cur_loss = linreg_loss(beta_hat, y, X)\n",
        "  diff = last_loss - cur_loss\n",
        "  last_loss = cur_loss\n",
        "\n",
        "  # wave to the crowd\n",
        "  print(f\"Loss[{idx}]({beta_hat}) = {last_loss}\")\n",
        "  idx += 1\n",
        "\n",
        "# OLS solution\n",
        "beta_hat_ols = jnp.linalg.solve(X.T @ X, X.T @ y)\n",
        "print(f\"ols beta = {beta_hat_ols}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "g6LpP-pxNy8y",
        "outputId": "f8bf784f-e72a-4cc9-d58e-61a1abd328db"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "true beta = [-0.7478904   0.07329974  0.05228964 -1.0152605  -0.6137889 ]\n",
            "Loss[0]([-0.6955502   0.12941222 -0.00297071 -0.99138546 -0.57778835]) = 1002.5853271484375\n",
            "Loss[1]([-0.7247923   0.04623848  0.05397529 -0.9596198  -0.64419246]) = 994.5390014648438\n",
            "Loss[2]([-0.7290263   0.05385752  0.0567775  -0.96942943 -0.64901537]) = 994.4423217773438\n",
            "Loss[3]([-0.7293095   0.05279243  0.05733057 -0.96860015 -0.6497474 ]) = 994.4410400390625\n",
            "Loss[4]([-0.7293439   0.05291931  0.05735086 -0.96873957 -0.64979035]) = 994.4409790039062\n",
            "ols beta = [-0.72934663  0.05290483  0.05735677 -0.9687273  -0.64979887]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "key, sim_key = rdm.split(key)\n",
        "\n",
        "N = 1000\n",
        "P = 5\n",
        "y, X, beta = sim_linear_reg(sim_key, N, P)\n",
        "\n",
        "step_size = 1 / N\n",
        "diff = 10.\n",
        "last_loss = 1000.\n",
        "idx = 0\n",
        "beta_hat = jnp.zeros((P,))\n",
        "# while delta in loss is large, continue\n",
        "print(\"Using JAX to compute gradient\")\n",
        "print(f\"true beta = {beta}\")\n",
        "while jnp.fabs(diff) > 1e-3:\n",
        "  # take a step in the direction of the gradient using step_size\n",
        "\n",
        "  beta_hat = beta_hat - step_size * jax.grad(linreg_loss)(beta_hat, y, X)\n",
        "\n",
        "  # update our current loss and compute delta\n",
        "  cur_loss = linreg_loss(beta_hat, y, X)\n",
        "  diff = last_loss - cur_loss\n",
        "  last_loss = cur_loss\n",
        "\n",
        "  # wave to the crowd\n",
        "  print(f\"Loss[{idx}]({beta_hat}) = {last_loss}\")\n",
        "  idx += 1\n",
        "\n",
        "# OLS solution\n",
        "beta_hat_ols = jnp.linalg.solve(X.T @ X, X.T @ y)\n",
        "print(f\"ols beta = {beta_hat_ols}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pyZh3Msjuncp",
        "outputId": "d84627f8-b200-4f56-dbea-97d4fa45e621"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Using JAX to compute gradient\n",
            "true beta = [-0.49276644  0.36893874  0.8708915  -0.22680327  0.31571558]\n",
            "Loss[0]([-0.5512781   0.26551786  0.9118643  -0.22698347  0.35736504]) = 699.7918701171875\n",
            "Loss[1]([-0.5188158   0.29075423  0.8899959  -0.22931792  0.3277857 ]) = 698.3126220703125\n",
            "Loss[2]([-0.5198399   0.29171753  0.8909349  -0.23179689  0.33005244]) = 698.305419921875\n",
            "Loss[3]([-0.5196871   0.2918358   0.89090306 -0.23202236  0.32988426]) = 698.305419921875\n",
            "ols beta = [-0.5196884   0.29184714  0.8909022  -0.23206532  0.3298884 ]\n"
          ]
        }
      ]
    }
  ]
}