Lab_3_Optimization_PtI.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyOKg/m6xBhuKcTfU57+dgpo",
      "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": 1,
      "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",
        "TBD: Add notes for Normal PDF and MLE estimator"
      ],
      "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": 5,
      "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": 7,
      "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": 8,
      "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": 9,
      "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",
        "TBD: Add notes for 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",
        "  pass\n",
        "\n",
        "def gradient(y, X, beta_hat):\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",
        "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"
          ]
        }
      ]
    },
    {
      "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",
        "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 * # COMPUTE GRADIENT USING JAX\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": {
        "id": "pyZh3Msjuncp"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}