Lab_2_LinearSolve.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_LinearSolve.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "J01SPGCfMEma"
      },
      "outputs": [],
      "source": [
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import jax.numpy.linalg as jnpla\n",
        "import jax.random as rdm\n",
        "import jax.scipy.linalg as jspla"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Sum(mer) Madness, or Solving Linear Equations\n",
        "Or how I learned not to take the inverse and perform [matrix-vector products](https://en.wikipedia.org/wiki/Operator_(mathematics)).\n",
        "\n",
        "Given $n \\times n$ non-singular (i.e. \"nice\") matrix $A$, $n \\times 1$ column-vector $b$, we can describe a [linear system of equations](https://en.wikipedia.org/wiki/System_of_linear_equations#General_form) related $A$, $b$ as $$Ax=b.$$\n",
        "Algebraically, we can solve for $x$ as, $x = A^{-1}b$, but how should we go about this [_numerically_](https://en.wikipedia.org/wiki/Numerical_linear_algebra)?"
      ],
      "metadata": {
        "id": "AO0OrFtLOxoA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def sim_linear_system(key, n: int):\n",
        "  key, x_key = rdm.split(key)\n",
        "  A = rdm.normal(key, shape=(n,n))\n",
        "  x = rdm.normal(x_key, shape=(n,))\n",
        "  b = A @ x\n",
        "  return A, x, b\n",
        "\n",
        "seed = 0\n",
        "N = 100\n",
        "\n",
        "key = rdm.PRNGKey(seed)\n",
        "key, sim_key = rdm.split(key)\n",
        "A, x, b = sim_linear_system(sim_key, N)\n",
        "\n",
        "# solve using algebraic approach\n",
        "Ainv = jnpla.inv(A)\n",
        "x_hat_direct = Ainv @ b\n",
        "\n",
        "# solve using 'blackbox(!?)' solver\n",
        "x_hat = jnpla.solve(A, b)\n",
        "\n",
        "# measure distance from truth using 2norm\n",
        "direct_dist = jnpla.norm(x - x_hat_direct)\n",
        "solve_dist = jnpla.norm(x - x_hat)\n",
        "\n",
        "print(f\"direct dist = {direct_dist} | solve_dist = {solve_dist}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xp93hPA4O-QL",
        "outputId": "e5cd49db-38d2-4ae7-c6cd-a6b0fa9caf1f"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "direct dist = 0.0003477725840639323 | solve_dist = 0.0002090586203848943\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We see less [numerical error](https://en.wikipedia.org/wiki/Numerical_analysis#Numerical_stability_and_well-posed_problems) in the solution using the 'solve' approach compared to our attempt at using the 'algebraic' solution. What happens as `N` increases?"
      ],
      "metadata": {
        "id": "oasiBAY7R3Tj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for N in [100, 1000, 5000]:\n",
        "\n",
        "  key, sim_key = rdm.split(key)\n",
        "  A, x, b = sim_linear_system(sim_key, N)\n",
        "\n",
        "  # solve using algebraic approach\n",
        "  Ainv = jnpla.inv(A)\n",
        "  x_hat_direct = Ainv @ b\n",
        "\n",
        "  # solve using 'blackbox(!?)' solver\n",
        "  x_hat = jnpla.solve(A, b)\n",
        "\n",
        "  # measure distance from truth using 2norm\n",
        "  direct_dist = jnpla.norm(x - x_hat_direct)\n",
        "  solve_dist = jnpla.norm(x - x_hat)\n",
        "\n",
        "  print(f\"N = {N} | direct dist = {direct_dist} | solve_dist = {solve_dist}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EoUXCHaiSH0S",
        "outputId": "d3ffe3b6-5c37-4a42-9310-60efb466d183"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "N = 100 | direct dist = 0.0002750760759226978 | solve_dist = 4.570868986775167e-05\n",
            "N = 1000 | direct dist = 0.02133580483496189 | solve_dist = 0.003495120210573077\n",
            "N = 5000 | direct dist = 0.3149307668209076 | solve_dist = 0.04819103330373764\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### How is `jnpla.solve` _solving_ the system?\n",
        "\n",
        "Numerical solvers typically perform some _decomopsition_ of $A$ into a product of structured matrices (e.g., orthogonal, lower/upper diagonal, permutations, etc.) which permit direct solutions through forward/backward solvers. A couple prominant examples are given below.\n",
        "\n",
        "[QR decomposition](https://en.wikipedia.org/wiki/QR_decomposition#Using_for_solution_to_linear_inverse_problems)\n",
        "\n",
        "\n",
        "$A = QR$ where $Q$ is $n \\times n$ [orthonormal matrix](https://en.wikipedia.org/wiki/Orthogonal_matrix) (i.e. $Q Q' = Q'Q = I$, thus $Q^{-1} = Q'$) and $R$ is $n \\times n$ upper [triangular matrix](https://en.wikipedia.org/wiki/Triangular_matrix) (i.e. $R_{ij} = 0, i < j$).\n",
        "\n",
        "Solving $Ax = b$ using QR amounts to noticing,\n",
        "$$QRx = b ⇒ Rx = Q^{-1}b = Q'b,$$ which can then be solved using a backwards solve. See, [jax.numpy.linalg.qr](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.qr.html), [jax.scipy.linalg.solve_triangular](https://jax.readthedocs.io/en/latest/_autosummary/jax.scipy.linalg.solve_triangular.html) for reference.\n",
        "\n",
        "[LU decomposition](https://en.wikipedia.org/wiki/LU_decomposition#Solving_linear_equations)\n",
        "\n",
        "$A = PLU$ where $P$ is an $n \\times n$ [permutation matrix](https://en.wikipedia.org/wiki/Permutation_matrix), $L$ is $n \\times n$ lower triangular matrix and $U$ is an $n \\times n$ upper triangular matrix. Solving $Ax = b$ using LU amounts to, $$PLU = b ⇒ LU = P^{-1}b = P'b,$$ which can be solved with one forwards solve $Ly = P'b$ and then a backwards solve $Ux = y$. See, [jax.scipy.linalg.lu_factor](https://jax.readthedocs.io/en/latest/_autosummary/jax.scipy.linalg.lu_factor.html), [jax.scipy.linalg.lu_solve](https://jax.readthedocs.io/en/latest/_autosummary/jax.scipy.linalg.lu_solve.html) for reference.\n",
        "\n",
        "Let's try to assess which is used by `jnpla.solve` by comparing error of the solutions (not ideal, but not worst idea possible; see [jax.numpy.linalg.solve](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.solve.html) for reference)."
      ],
      "metadata": {
        "id": "_n7nGU4VnoCc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "N = 25\n",
        "key, sim_key = rdm.split(key)\n",
        "A, x, b = sim_linear_system(sim_key, N)\n",
        "\n",
        "# solve using algebraic approach\n",
        "Ainv = jnpla.inv(A)\n",
        "x_direct = Ainv @ b\n",
        "\n",
        "# QR\n",
        "Q, R = jnpla.qr(A)\n",
        "x_qr = jspla.solve_triangular(R, Q.T @ b)\n",
        "\n",
        "# LU\n",
        "LU, P = jspla.lu_factor(A)\n",
        "x_lu = jspla.lu_solve((LU, P), b)\n",
        "\n",
        "# direct solve for baseline\n",
        "x_solve = jnpla.solve(A, b)\n",
        "\n",
        "direct_dist = jnpla.norm(x - x_direct)\n",
        "qr_dist = jnpla.norm(x - x_qr)\n",
        "lu_dist = jnpla.norm(x - x_lu)\n",
        "solve_dist = jnpla.norm(x - x_solve)\n",
        "\n",
        "\n",
        "print(f\"direct dist = {direct_dist} | qr dist = {qr_dist} | lu dist = {lu_dist} | solve_dist = {solve_dist}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dr1IIFXjS44u",
        "outputId": "d40f02c8-2c48-46b6-bb27-bb435ed399a5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "direct dist = 2.836631210811902e-05 | qr dist = 2.3936559955473058e-05 | lu dist = 4.185290890745819e-06 | solve_dist = 4.185290890745819e-06\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### How much time does it take to solve a system of linear equations?\n",
        "Complexity is $O(n^3)$ when $A$ is $n \\times n$ (under exact numerical precision, i.e. ignoring bit complexity of numerical results), but can we do better in terms of constants? Pls no [mention](https://en.wikipedia.org/wiki/Computational_complexity_of_matrix_multiplication#Matrix_multiplication_exponent) of $O(n^\\omega)$.\n",
        "\n",
        "What if we know that our matrix $A$ has additional structure from the start?\n",
        "\n",
        "[Cholesky decomposition](https://en.wikipedia.org/wiki/Cholesky_decomposition)\n",
        "\n",
        "When $A$ is an $n \\times n$ [_positive definite_](https://en.wikipedia.org/wiki/Definite_matrix) (i.e. $x'Ax \\geq 0, ∀ x \\in R^n$), the Cholesky decomposition gives $A = LL'$ where $L$ is an $n \\times n$ lower triangular matrix.\n",
        "\n",
        "Solving $Ax=b$ using Cholesky decomposition amounts to, $$Ax=b ⇒LL'x = b,$$\n",
        "which  can be solved with one forwards solve $Ly=b$ and then a backwards solve $L'x=y$. This will be roughly twice as fast compared with LU decomposition approaches. See [jax.scipy.linalg.cho_factor](https://jax.readthedocs.io/en/latest/_autosummary/jax.scipy.linalg.cho_factor.html) and [jax.scipy.linalg.cho_solve](https://jax.readthedocs.io/en/latest/_autosummary/jax.scipy.linalg.cho_solve.html) for reference.\n",
        "\n",
        "Let's compare the error of the solve solutions."
      ],
      "metadata": {
        "id": "uCM1v4ATZlC8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def sim_sym_linear_system(key, n: int):\n",
        "  key, x_key = rdm.split(key, 2)\n",
        "  # sample larger row to improve conditioning\n",
        "  A = rdm.normal(key, shape=(2*n,n))\n",
        "  # rescale back to n x n\n",
        "  A = A.T @ A\n",
        "  x = rdm.normal(x_key, shape=(n,))\n",
        "  b = A @ x\n",
        "  return A, x, b\n",
        "\n",
        "key, sim_key = rdm.split(key)\n",
        "N = 500\n",
        "\n",
        "A, x, b = sim_sym_linear_system(sim_key, N)\n",
        "\n",
        "# solve using algebraic approach\n",
        "Ainv = jnpla.inv(A)\n",
        "x_direct = Ainv @ b\n",
        "\n",
        "# QR\n",
        "Q, R = jnpla.qr(A)\n",
        "x_qr = jspla.solve_triangular(R, Q.T @ b)\n",
        "\n",
        "# Cholesky\n",
        "L, lower = jspla.cho_factor(A)\n",
        "x_cho = jspla.cho_solve((L, lower), b)\n",
        "\n",
        "# direct solve for baseline\n",
        "x_solve = jnpla.solve(A, b)\n",
        "\n",
        "direct_dist = jnpla.norm(x - x_direct)\n",
        "qr_dist = jnpla.norm(x - x_qr)\n",
        "cho_dist = jnpla.norm(x - x_cho)\n",
        "solve_dist = jnpla.norm(x - x_solve)\n",
        "\n",
        "print(f\"direct dist = {direct_dist} | qr dist = {qr_dist} | cho dist = {cho_dist} | solve_dist = {solve_dist}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "68jTTKDDCt7O",
        "outputId": "9530755d-868e-4e4e-e8fb-63724e577a93"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "direct dist = 2.7083056920673698e-05 | qr dist = 2.6027397325378843e-05 | cho dist = 1.7303784261457622e-05 | solve_dist = 1.732515011099167e-05\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### LAB portion\n",
        "\n",
        "What is the average runtime for each of these operations?\n",
        "Hint: use the `%timeit` magic command and [`block_until_ready`](https://jax.readthedocs.io/en/latest/_autosummary/jax.block_until_ready.html) function."
      ],
      "metadata": {
        "id": "XDtGoVs_GyeP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "N = 1000\n",
        "key, sim_key = rdm.split(key)\n",
        "A, x, b = sim_sym_linear_system(sim_key, N)\n",
        "\n",
        "# solve using algebraic approach\n",
        "def _direct(A, b):\n",
        "  pass\n",
        "\n",
        "# QR\n",
        "def _qr(A, b):\n",
        "  pass\n",
        "\n",
        "# Cholesky\n",
        "def _cholesky(A, b):\n",
        "  pass\n",
        "\n",
        "# general solver (uses LU)\n",
        "def _solve(A, b):\n",
        "  return jnpla.solve(A, b)\n",
        "\n",
        "%timeit x_solve = _solve(A, b).block_until_ready()\n",
        "%timeit x_direct = _direct(A, b).block_until_ready()\n",
        "%timeit x_qr = _qr(A, b).block_until_ready()\n",
        "%timeit x_cho = _cholesky(A, b).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "v65o0jsHG1Rt",
        "outputId": "de588426-3821-4cdd-e46c-6187ef9bd632"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "48 ms ± 15.1 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Great! Cholesky is the fastest and has error on par with the general `solve` command. What can we gain in terms of speed by using [`jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html)?"
      ],
      "metadata": {
        "id": "_ORvGCONKJV8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "N = 1000\n",
        "key, sim_key = rdm.split(key)\n",
        "A, x, b = sim_sym_linear_system(sim_key, N)\n",
        "\n",
        "jit_direct = jax.jit(_direct)\n",
        "jit_qr = jax.jit(_qr)\n",
        "jit_cholesky = jax.jit(_cholesky)\n",
        "jit_solve = jax.jit(jnpla.solve)\n",
        "\n",
        "%timeit x_direct = jit_direct(A, b).block_until_ready()\n",
        "%timeit x_qr = jit_qr(A, b).block_until_ready()\n",
        "%timeit x_cho = jit_cholesky(A, b).block_until_ready()\n",
        "%timeit x_solve = jit_solve(A, b).block_until_ready()"
      ],
      "metadata": {
        "id": "zcf4ynyAKSYY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Not much improvement! What gives?\n",
        "\n",
        "Turns out, JAX is already JIT compiling these lower level functions (e.g., solve, decompositions) and our code is not much larger in scope, thus benefit not clearly demonstrated. The takeaway here is not that JIT doesn't help, but rather, to get the most out of JIT'd code, we should apply it to larger programs/functions, which can better optimize the computational graph.\n",
        "\n"
      ],
      "metadata": {
        "id": "E1ZgjnttNJqv"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Conjugate Gradient (CG)\n",
        "TBD"
      ],
      "metadata": {
        "id": "rgRH20aEqdAt"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## And now for something completely different... [Lineax](https://docs.kidger.site/lineax/)\n",
        "Library for solving linear systems and related linear computational tasks."
      ],
      "metadata": {
        "id": "aB7C0NUkqb7j"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install lineax"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OvHjwTYmNE-F",
        "outputId": "082ac1a8-3f0a-4165-b2ec-1f6609b24394"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: lineax in /usr/local/lib/python3.11/dist-packages (0.0.7)\n",
            "Requirement already satisfied: equinox>=0.11.5 in /usr/local/lib/python3.11/dist-packages (from lineax) (0.11.11)\n",
            "Requirement already satisfied: jax>=0.4.26 in /usr/local/lib/python3.11/dist-packages (from lineax) (0.5.0)\n",
            "Requirement already satisfied: jaxtyping>=0.2.20 in /usr/local/lib/python3.11/dist-packages (from lineax) (0.2.36)\n",
            "Requirement already satisfied: typing-extensions>=4.5.0 in /usr/local/lib/python3.11/dist-packages (from lineax) (4.12.2)\n",
            "Requirement already satisfied: jaxlib<=0.5.0,>=0.5.0 in /usr/local/lib/python3.11/dist-packages (from jax>=0.4.26->lineax) (0.5.0)\n",
            "Requirement already satisfied: ml_dtypes>=0.4.0 in /usr/local/lib/python3.11/dist-packages (from jax>=0.4.26->lineax) (0.4.1)\n",
            "Requirement already satisfied: numpy>=1.25 in /usr/local/lib/python3.11/dist-packages (from jax>=0.4.26->lineax) (1.26.4)\n",
            "Requirement already satisfied: opt_einsum in /usr/local/lib/python3.11/dist-packages (from jax>=0.4.26->lineax) (3.4.0)\n",
            "Requirement already satisfied: scipy>=1.11.1 in /usr/local/lib/python3.11/dist-packages (from jax>=0.4.26->lineax) (1.13.1)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import lineax as lx\n",
        "\n",
        "N = 1000\n",
        "key, sim_key = rdm.split(key)\n",
        "A, x, b = sim_sym_linear_system(sim_key, N)\n",
        "\n",
        "A_op = lx.MatrixLinearOperator(A)\n",
        "sol = lx.linear_solve(A_op, b)\n",
        "x_sol = sol.value\n",
        "\n",
        "A_psd_op = lx.MatrixLinearOperator(A, lx.positive_semidefinite_tag)\n",
        "psd_sol = lx.linear_solve(A_psd_op, b)\n",
        "x_psd_sol = psd_sol.value\n",
        "\n",
        "# initialize our CG solver\n",
        "cg_solver = lx.CG(atol=1e-5, rtol=1e-4)\n",
        "y0 = b / jnp.diag(A)\n",
        "cg_sol = lx.linear_solve(A_psd_op, b, solver=cg_solver, options={\"y0\":y0})\n",
        "x_cg_sol = cg_sol.value\n",
        "\n",
        "x_solve = jnpla.solve(A, b)\n",
        "\n",
        "dist_sol = jnpla.norm(x - x_sol)\n",
        "dist_psd = jnpla.norm(x - x_psd_sol)\n",
        "dist_cg = jnpla.norm(x - x_cg_sol)\n",
        "dist_solve = jnpla.norm(x - x_solve)\n",
        "\n",
        "print(f\"dist lineax solve = {dist_sol}\")\n",
        "print(f\"psd lineax solve = {dist_psd}\")\n",
        "print(f\"cg lineax solve = {dist_cg}\")\n",
        "print(f\"jax solve = {dist_solve}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0amWH5fNOYqA",
        "outputId": "28ec0727-6dc4-4ddd-a299-a106aee42892"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "dist lineax solve = 2.793022031255532e-05\n",
            "psd lineax solve = 3.014422873093281e-05\n",
            "cg lineax solve = 3.1980012863641605e-05\n",
            "jax solve = 2.793022031255532e-05\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's look at slightly more sophisticated linear operator. Let $A = B B' + I_n$ we want to solve $Ax = b$, where $B$ is shape $n \\times k$ for $k < n$."
      ],
      "metadata": {
        "id": "ZbVsbI7Bx0Dr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "N, K = 1000, 50\n",
        "key, b_key, x_key = rdm.split(key, 3)\n",
        "\n",
        "B = rdm.normal(b_key, shape=(N, K))\n",
        "x = rdm.normal(x_key, shape=(N,))\n",
        "b = B @ (B.T @ x) + x\n",
        "\n",
        "B_op = lx.MatrixLinearOperator(B)\n",
        "I_op = lx.IdentityLinearOperator(jax.eval_shape(lambda: x))\n",
        "A_op = lx.TaggedLinearOperator(B_op @ B_op.T + I_op, lx.positive_semidefinite_tag)\n",
        "\n",
        "sol = lx.linear_solve(A_op, b)\n",
        "cg_sol = lx.linear_solve(A_op, b, solver=cg_solver)\n",
        "\n",
        "dist_sol = jnpla.norm(x - sol.value)\n",
        "dist_cg = jnpla.norm(x - cg_sol.value)\n",
        "print(f\"dist lineax solve = {dist_sol}\")\n",
        "print(f\"cg lineax solve = {dist_cg}\")"
      ],
      "metadata": {
        "id": "ChXNObUUyQlF",
        "outputId": "10eea4cf-5da0-410c-e423-e3e76c783bc9",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "dist lineax solve = 0.003705778392031789\n",
            "cg lineax solve = 0.0006713042384944856\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit lx.linear_solve(A_op, b).value.block_until_ready()\n",
        "%timeit lx.linear_solve(A_op, b, solver=cg_solver).value.block_until_ready()"
      ],
      "metadata": {
        "id": "ntO14u_70qoi",
        "outputId": "f9d40c4a-f0cb-497e-940b-4e4abb6e34ba",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "22.2 ms ± 10.2 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
            "1.99 ms ± 48.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def sim_linear_reg(key, N, P, r2=0.5):\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",
        "  return y, X, beta\n",
        "\n",
        "\n",
        "N, P = 1000, 150\n",
        "key, sim_key = rdm.split(key)\n",
        "y, X, beta = sim_linear_reg(sim_key, N, P)\n",
        "\n",
        "XtX = X.T @ X\n",
        "Xty = X.T @ y\n",
        "\n",
        "# dist func\n",
        "def _dist(sol):\n",
        "  return jnpla.norm(beta - sol.value)\n",
        "\n",
        "X_op = lx.MatrixLinearOperator(X)\n",
        "XtX_op = lx.MatrixLinearOperator(XtX, lx.positive_semidefinite_tag)\n",
        "\n",
        "# 1st: linear operator over `XtX` and solve against Xty\n",
        "sol_1 = lx.linear_solve(XtX_op, Xty)\n",
        "sol_1_dist = _dist(sol_1)\n",
        "\n",
        "# 2nd: linear operator over `X` and solve against `y`\n",
        "sol_2 = lx.linear_solve(X_op, y, solver=lx.AutoLinearSolver(well_posed=False))\n",
        "sol_2_dist = _dist(sol_2)\n",
        "\n",
        "# 3rd: linear operator over `X`, then compose with X.T (ie Xop.T @ Xop) solve against Xty\n",
        "sol_3 = lx.linear_solve(X_op.T @ X_op, Xty)\n",
        "sol_3_dist = _dist(sol_3)\n",
        "\n",
        "# 4th: linear operator over `X` then solve against `y` using NormalCG\n",
        "solver = lx.NormalCG(atol=1e-4, rtol=1e-4)\n",
        "sol_4 = lx.linear_solve(X_op, y, solver=solver)\n",
        "sol_4_dist = _dist(sol_4)\n",
        "print(f\" XtX Xty = {sol_1_dist} | X y = {sol_2_dist} | XtX Xty #2 = {sol_3_dist} | XtX Xty CG = {sol_4_dist}\")"
      ],
      "metadata": {
        "id": "lK54mdzzSKpH",
        "outputId": "2a914c6d-e4e5-4254-ed2d-1b7c3a95dd43",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            " XtX Xty = 4.768477916717529 | X y = 4.768476486206055 | XtX Xty #2 = 4.768477916717529 | XtX Xty CG = 4.768479347229004\n"
          ]
        }
      ]
    }
  ]
}