Lab_2_Sampling.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "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_2_Sampling.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import jax.random as rdm"
      ],
      "metadata": {
        "id": "dGIjVYw8-as8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# random variables in JAX\n",
        "key = rdm.PRNGKey(0)\n",
        "\n",
        "key, y_key = rdm.split(key)\n",
        "N = 500\n",
        "mu_y = 50\n",
        "std_y = 100\n",
        "y = mu_y + std_y * rdm.normal(y_key, shape=(N,)) # y_i ~ N(mu_y, std_y)\n",
        "print(f\"mean(y) = {jnp.mean(y)}\")\n",
        "\n",
        "P = 100\n",
        "key, x_key = rdm.split(key)\n",
        "X = rdm.normal(x_key, shape=(N,P))\n",
        "print(f\"shape(x) = {X.shape}\")"
      ],
      "metadata": {
        "id": "oUiwVMCfJAWw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Simulations under a linear model\n",
        "Goal: given $n=500$, $p=100$, simulate a random normal matrix of size $n \\times p$, simulate $p$ effect sizes $\\beta$ under a standard normal distribution, compute $g = X \\beta$, and $s^2_g :=$ var($g$) and lastly compute\n",
        "$y = g + \\epsilon$ where $\\epsilon_i \\sim N(0, \\sigma^2e)$ where $\\sigma^2_e := (1 - r^2)/r^2 * s^2_g$. where $r^2 = 0.5$."
      ],
      "metadata": {
        "id": "9_4Sn8YRA1DQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# our PRNG seed to initialize state for sampling and replication\n",
        "seed = 0\n",
        "\n",
        "# defined quantities\n",
        "N = 500\n",
        "P = 100\n",
        "r2 = 0.5\n",
        "\n",
        "# create initial key\n",
        "key = rdm.PRNGKey(seed)\n",
        "\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",
        "\n",
        "# compute basic stats and sanity check against specified h2\n",
        "print(f\"mean(y) = {jnp.mean(y)} | var(y) = {jnp.var(y)}\")\n",
        "print(f\"hat(h2) = {s2g / jnp.var(y)} | true h2 = {h2}\")"
      ],
      "metadata": {
        "id": "WVpjOl6ECPFn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's wrap the above functionality into a function called `sim_data` that takes as its arguments,\n",
        "  1. an PRNGKey\n",
        "  2. N\n",
        "  3. P\n",
        "  4. h2\n",
        "and returns `y`, `X`, `beta`"
      ],
      "metadata": {
        "id": "dVGhJajXGFpx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def sim_data(key, N, P, h2):\n",
        "  \"\"\"\n",
        "  Simulate a trait under a linear model;\n",
        "  \"\"\"\n",
        "  pass"
      ],
      "metadata": {
        "id": "-44xpVZyGYAW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def compute_stat(y, x):\n",
        "  \"\"\"\n",
        "  Compute OLS between y and x;\n",
        "  Return hat{beta}_x and se(hat{beta}_x)\n",
        "  \"\"\"\n",
        "  pass"
      ],
      "metadata": {
        "id": "WpCxclEnH6Nn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Compute stats!\n",
        "Let's loop over each column in `X` and compute its _marginal_ statistics using the above function."
      ],
      "metadata": {
        "id": "FzLbWW9aKXet"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# import our stats library\n",
        "import jax.scipy.stats as stats\n",
        "\n",
        "beta_hats = []\n",
        "ses = []\n",
        "\n",
        "for i in range(P):\n",
        "  beta_i, se_i = compute_stat(y, X[:, i])\n",
        "  beta_hats.append(beta_i)\n",
        "  ses.append(se_i)\n",
        "\n",
        "beta_hats = jnp.array(beta_hats)\n",
        "ses = jnp.array(ses)\n",
        "zscores = beta_hats / ses\n",
        "pvalues = 2 * stats.norm.sf(jnp.abs(zscores))"
      ],
      "metadata": {
        "id": "5PmVRT94KXE8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# permutation tests wow!\n",
        "let's write a function `perm` that takes as input,\n",
        "  1. PRNGKey,\n",
        "  2. y,\n",
        "  3. X,\n",
        "  4. beta_hat\n",
        "  5. B, the number of permutations to perform\n",
        "\n",
        "and returns empirical/permutation pvalues for each effect"
      ],
      "metadata": {
        "id": "ws-AwCeLPAOL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def compute_marginals(y, X):\n",
        "  \"\"\"\n",
        "  Compute the marginal statistics for reach column in X against y;\n",
        "  Return tuple of (beta_hat, ses) where each are jax arrays\n",
        "  \"\"\"\n",
        "  pass\n",
        "\n",
        "def perm(key, y, X, zscores, B):\n",
        "  \"\"\"\n",
        "  Compute a permutation test for each marginal effect over X;\n",
        "  Returns the permutation/empirical p-value for each marginal effect\n",
        "  as a single jax array.\n",
        "  \"\"\"\n",
        "  pass\n",
        "\n",
        "# split key for fun\n",
        "key, p_key = rdm.split(key)\n",
        "\n",
        "# compute marginals under observed\n",
        "beta_hat, ses = compute_marginals(y, X)\n",
        "\n",
        "# compute zscores and pvalues under normality\n",
        "zscores = beta_hats / ses\n",
        "pvalues = 2 * stats.norm.sf(jnp.abs(zscores))\n",
        "\n",
        "# compute empirical pvalues\n",
        "B = 100\n",
        "perm(p_key, y, X, zscores, B)"
      ],
      "metadata": {
        "id": "JQwa15dEPT2j"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "s-9DLBaFDb0e"
      }
    }
  ]
}