Lab_4_Optimization_PtI.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNgY/qVpFZKbwZsA3r3nC0q",
      "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_4_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"
      ],
      "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 = jnp.mean(x)\n",
        "  ssq_hat = jnp.mean((x - mu_hat) ** 2) # jnp.var(x)\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 = 100\n",
        "\n",
        "mu = 58.\n",
        "sigma_sq = 10.\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": "5fe02f0f-86e4-4b0e-cbce-14fde574e42a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "x = [54.748363 61.49207  58.068737 62.63578  63.038616 57.023346 62.89525\n",
            " 67.758804 53.49503  64.813095 54.75522  56.203518 58.778862 59.815628\n",
            " 59.685024 52.485527 54.44686  58.31798  59.458866 54.291496 62.230354\n",
            " 61.92895  57.596672 54.478622 62.95002  61.07032  58.497906 61.5559\n",
            " 57.91157  48.89393  58.041267 57.05834  54.320095 60.814167 58.863586\n",
            " 52.72084  57.37599  58.616085 61.504147 53.660152 53.529938 56.553493\n",
            " 59.48462  57.536884 58.604736 60.350792 59.59793  58.373173 56.29943\n",
            " 59.44     60.093395 58.018135 60.842957 58.10686  62.98405  56.63515\n",
            " 60.17766  62.55351  61.50157  63.707592 55.19841  63.116325 57.727917\n",
            " 57.98455  62.07337  59.92226  55.25571  55.92521  60.866608 58.606293\n",
            " 58.526638 62.975803 60.693626 63.814217 62.8752   57.88707  61.743687\n",
            " 62.84971  61.9028   60.084473 61.1948   56.62765  59.065964 56.75269\n",
            " 57.404243 53.602177 56.415016 54.36041  60.074253 60.310368 59.518055\n",
            " 62.36892  57.469944 62.59477  61.664345 62.104057 55.20062  62.079853\n",
            " 56.311634 54.813797]\n",
            "MLE[\\mu, \\sigma^2] = 58.88652038574219, 10.556812286376953\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": "8118f5b0-5db4-41c4-9523-1f5c8125a08d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "MSE[50 | mu, sigma^2] = 0.041574958711862564, 0.10048199445009232\n",
            "MSE[100 | mu, sigma^2] = 0.009534507989883423, 0.43518951535224915\n",
            "MSE[1000 | mu, sigma^2] = 0.0006926158675923944, 0.03188581019639969\n",
            "MSE[10000 | mu, sigma^2] = 0.0016879187896847725, 0.001828653272241354\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## MLE for iid Exponential data"
      ],
      "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 = 1 / jnp.mean(x)\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": "1f7d7ba3-44f9-437e-f1e7-18c36fec0ec9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "x = [3.88764771e+02 3.65960579e+01 7.46025085e+01 7.58861938e+02\n",
            " 3.56365692e+02 2.62509766e+02 1.04530891e+02 6.67349365e+02\n",
            " 1.46979480e+03 2.09354004e+02 1.62024744e+03 5.43406982e+01\n",
            " 1.95166306e+02 2.59302582e+02 9.02553589e+02 4.48792053e+02\n",
            " 1.42324921e+02 4.51267151e+02 3.21048798e+02 1.19272351e+03\n",
            " 4.99932159e+02 5.77076836e+01 5.10972626e+02 3.86068085e+02\n",
            " 2.08017014e+02 9.29669556e+02 1.19520203e+03 3.36191833e+02\n",
            " 1.69074356e+02 5.36690735e+02 8.64235901e+02 1.34896851e+03\n",
            " 3.47149231e+02 5.78509941e+01 1.26216560e+02 1.29906509e+02\n",
            " 1.21480568e+02 3.94892082e+01 2.79801056e+02 5.42540550e+01\n",
            " 3.20104187e+02 7.94236511e+02 2.21294403e+02 6.73425903e+02\n",
            " 1.19045776e+02 5.03411026e+01 4.75495392e+02 2.45755737e+02\n",
            " 4.47672882e+02 1.20910376e+03 1.23082751e+03 9.31572266e+01\n",
            " 9.14312134e+02 9.95121338e+02 5.24672974e+02 7.28671570e+02\n",
            " 9.60228699e+02 2.89802608e+01 4.31902161e+02 3.25366783e+01\n",
            " 6.16884277e+02 3.53967224e+02 1.63407410e+02 7.94963501e+02\n",
            " 9.10376282e+02 1.26778746e+00 3.08394745e+02 1.89620941e+02\n",
            " 5.72219604e+02 6.55903778e+01 1.97578461e+02 7.51722595e+02\n",
            " 9.56931381e+01 9.55107971e+02 3.99790527e+02 7.86239548e+01\n",
            " 6.38027039e+01 5.11758900e+00 7.73605713e+02 1.04103004e+02\n",
            " 1.00399670e+03 1.98847855e+02 6.43844681e+01 1.56428146e+01\n",
            " 4.07316467e+02 3.06054718e+02 1.96847366e+02 2.59816223e+02\n",
            " 2.15031787e+03 2.82647095e+02 4.67662781e+02 9.49611023e+02\n",
            " 4.08943420e+02 1.10794202e+03 1.66897995e+02 2.79430359e+02\n",
            " 1.09352661e+02 7.39537109e+02 2.04390405e+03 4.11656860e+02]\n",
            "MLE[\\lambda = 0.002] = 0.002101506572216749\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": "d1d5a7f5-c30f-4811-ad67-7d6a12185723"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "MSE[50 | \\lambda = 0.02] = 8.534788321412634e-06\n",
            "MSE[100 | \\lambda = 0.02] = 8.773293302510865e-06\n",
            "MSE[1000 | \\lambda = 0.02] = 2.0062362882100615e-09\n",
            "MSE[10000 | \\lambda = 0.02] = 1.0255300963990521e-07\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Gradient descent"
      ],
      "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(y, X, beta_hat):\n",
        "  pred_y = X @ beta_hat\n",
        "  return jnp.sum((y - pred_y)**2)\n",
        "\n",
        "def gradient(y, X, beta_hat):\n",
        "  Xty = X.T @ y\n",
        "  XtX = X.T @ X\n",
        "  return (XtX @ beta_hat) - Xty\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",
        "while jnp.fabs(diff) > 1e-3:\n",
        "  # take a step in the direction of the gradient using step_size\n",
        "  beta_hat = beta_hat - step_size * gradient(y, X, beta_hat)\n",
        "  # update our current loss and compute delta\n",
        "  cur_loss = linreg_loss(y, X, beta_hat)\n",
        "  diff = last_loss - cur_loss\n",
        "  last_loss = cur_loss\n",
        "  # wave to the crowd\n",
        "  print(f\"Loss[{idx} | {beta}] = {last_loss} @ {beta_hat}\")\n",
        "  idx += 1"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "g6LpP-pxNy8y",
        "outputId": "63101837-cf6a-48c9-a2a3-54b8f0b54417"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loss[0 | [-1.0257084   1.8594663   1.1235882   0.23114179 -0.34901348]] = 5639.55810546875 @ [-1.079702    1.8402483   1.1758611   0.23298487 -0.52076465]\n",
            "Loss[1 | [-1.0257084   1.8594663   1.1235882   0.23114179 -0.34901348]] = 5623.48291015625 @ [-1.1577635   1.8372455   1.081856    0.24313869 -0.47761846]\n",
            "Loss[2 | [-1.0257084   1.8594663   1.1235882   0.23114179 -0.34901348]] = 5623.23828125 @ [-1.1617126  1.8369826  1.0963572  0.2397221 -0.4716044]\n",
            "Loss[3 | [-1.0257084   1.8594663   1.1235882   0.23114179 -0.34901348]] = 5623.23193359375 @ [-1.1620235   1.8380398   1.094032    0.23982307 -0.47083738]\n",
            "Loss[4 | [-1.0257084   1.8594663   1.1235882   0.23114179 -0.34901348]] = 5623.2314453125 @ [-1.1620889   1.8379307   1.0944436   0.23976085 -0.4707122 ]\n"
          ]
        }
      ]
    }
  ]
}