tutorial_solution.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Practical Session: Probabilistic Linear Solvers and Their Applications\n",
    "Author: Marvin Pförtner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial we will implement a **matrix-free** **probabilistic linear solver** (PLS) from scratch and apply it to linear systems arising in two different applications: Gaussian process regression and the numerical solution of a linear PDE."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Clone the GitHub repository\n",
    "\n",
    "`https://github.com/marvinpfoertner/2024-pn-spring-school`\n",
    "\n",
    "To participate in this tutorial, you will need a fresh Python 3.11 virtual environment (or a conda/miniconda installation).\n",
    "Activate the environment, and run `pip install -r requirements.txt` in the repository to install all necessary requirements.\n",
    "\n",
    "If you use conda, you can simply run `conda env create -f environment.yml` in the repository.\n",
    "This will create a new conda env named `pn-school-linalg` with the correct Python version and all necessary requirements installed automatically.\n",
    "\n",
    "In either case, run `jupyter lab` and navigate to the tutorial notebook.\n",
    "\n",
    "If everything is installed correctly, the following two cells should run without errors (the first one might produce a `[KeOps] Warning : ...`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "import scipy.linalg\n",
    "\n",
    "import probnum as pn\n",
    "import linpde_gp\n",
    "\n",
    "from numpy.typing import ArrayLike\n",
    "from probnum.typing import LinearOperatorLike"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%config InlineBackend.figure_formats = ['svg']\n",
    "%matplotlib inline\n",
    "\n",
    "FIGSIZE = (8, 4)\n",
    "plt.rcParams['figure.figsize'] = FIGSIZE\n",
    "plt.rcParams[\"figure.autolayout\"] = True\n",
    "plt.rcParams[\"text.usetex\"] = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Matrix-Free Linear Algebra"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is well-known that solving a linear system $A x = b$ with $A \\in \\mathbb{R}^{n \\times n}$ has a worst-case time complexity of $\\mathcal{O}(n^3)$.\n",
    "This is often used to argue that the cost of solving such problems scales prohibitively with $n$.\n",
    "However, in practice, the quadratic memory cost of storing a (dense) $n \\times n$ matrix in memory is arguably much worse.\n",
    "\n",
    "Luckily, for many problems, so-called **matrix-free** methods tackle both of these problems at the same time.\n",
    "In contrast to classical linear solvers, they only require access to matrix-vector products with the system matrix, i.e. $v \\mapsto A v$.\n",
    "Many system matrices $A$ arising in practical applications admit efficient computation of the maps $v \\mapsto A v$ without incurring $\\mathcal{O}(n^2)$ memory cost.\n",
    "We call the such maps $A[\\cdot] : \\mathbb{R}^n \\to \\mathbb{R}^n, v \\to A v$ **matrix-free linear operators**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Examples of Matrix-Free Linear Operators**:\n",
    "- sparse matrices (e.g. banded matrices) in $\\mathcal{O}(\\operatorname{nnz}(A))$ time and space\n",
    "- diagonal matrices $A = \\operatorname{diag}(a)$, since $$A[v] = a \\odot v$$ in $\\mathcal{O}(n)$ time and space\n",
    "- symbolic matrices (e.g. kernel Gram matrices) $A_{ij} = F(x_i, x_j)$, since $$(A[v])_i = \\sum_{j = 1}^n F(x_i, x_j) v_j$$ in $\\mathcal{O}(n^2)$ time and $\\mathcal{O}(n)$ space\n",
    "- low-rank matrices $A = U V^T$ with $U, V \\in \\mathbb{R}^{n \\times r}$ and $r \\ll d$, since $$A[v] = U[\\underbrace{(V^\\top[v])}_{\\in \\mathbb{R}^n}]$$\n",
    "- Toeplitz matrices $A = \\operatorname{toep}(a)$ with $a \\in \\mathbb{R}^n$, since $$A[v] = \\operatorname{rfft}^{-1}(\\operatorname{rfft}(a) \\odot \\operatorname{rfft}(v))$$\n",
    "in $\\mathcal{O}(n \\log n)$ time and $\\mathcal{O}(n)$ space\n",
    "- recursive algebraic combinations of matrix-free linear operators\n",
    "    - linear combinations $A = \\beta B + \\gamma C$, since $A[v] = \\beta B[v] + \\gamma C[v]$\n",
    "    - composition $A = B C$, since $A[v] = B[C[v]]$\n",
    "    - Kronecker (and Khatri–Rao) products $A = B \\otimes C$ with $B \\in \\mathbb{R}^{n_B \\times n_B}, C \\in \\mathbb{R}^{n_C \\times n_C}$ and $n_B, n_C \\in \\mathcal{O}(\\sqrt{n})$, since $$A[\\operatorname{vec}(M)] = \\operatorname{vec}(C[B[M^\\top]^\\top])$$ in $\\mathcal{O}(n \\sqrt{n})$ time and $\\mathcal{O}(n)$ space.\n",
    "    - block matrices $$A = \\begin{pmatrix} B & C \\\\ D & E \\end{pmatrix},$$ since $$A \\begin{pmatrix} u \\\\ v\\end{pmatrix} = \\begin{pmatrix} B[u] + C[v] \\\\ D[u] + E[v] \\end{pmatrix}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The subpackage [`probnum.linops`](https://probnum.readthedocs.io/en/latest/api/linops.html) of the `probnum` Python library implements a small computer algebra system for working with matrix-free linear operators."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "\n",
    "**Task 1:** Implement a [`probnum.linops.LinearOperator`](https://probnum.readthedocs.io/en/latest/api/automod/probnum.linops.LinearOperator.html#probnum.linops.LinearOperator) for representing a symmetric tridiagonal matrix\n",
    "\\begin{equation*}\n",
    "A =\n",
    "\\begin{pmatrix}\n",
    "    a_1    & b_1    & 0      & 0      & \\cdots & 0      \\\\\n",
    "    b_1    & a_2    & b_2    & 0      & \\ddots & \\vdots \\\\\n",
    "    0      & b_2    & a_3    & b_3    & \\ddots & 0 \\\\\n",
    "    0      & 0      & b_3    & a_4    & \\ddots & 0 \\\\\n",
    "    \\vdots & \\ddots & \\ddots & \\ddots & \\ddots   & b_{n - 1} \\\\\n",
    "    0      & \\cdots & 0      & 0      &  b_{n - 1} & a_n\n",
    "\\end{pmatrix}\n",
    "\\end{equation*}\n",
    "in $\\mathcal{O}(n)$ space.\n",
    "The scaffolding is already given below, you just have to implement the method `_matmul` realizing $x \\mapsto A x$ and the method `_solve` realizing $x \\mapsto A^{-1} v$.\n",
    "\n",
    "*Hint:* For the `_solve` method, the function `scipy.linalg.solve_banded` could come in handy.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SymmetricTridiagonalMatrix(pn.linops.LinearOperator):\n",
    "    def __init__(self, diag: ArrayLike, offdiag: ArrayLike):\n",
    "        self._diag = np.asarray(diag)\n",
    "        self._offdiag = np.asarray(offdiag)\n",
    "\n",
    "        if self._diag.ndim != 1:\n",
    "            raise ValueError(\"`diag` must be a 1D array\")\n",
    "\n",
    "        if self._offdiag.shape != (self._diag.size - 1,):\n",
    "            raise ValueError(\"`offdiag` must have shape `(diag.size - 1,)`.\")\n",
    "\n",
    "        super().__init__(\n",
    "            shape=(self._diag.size, self._diag.size),\n",
    "            dtype=np.promote_types(self._diag.dtype, self._offdiag.dtype),\n",
    "        )\n",
    "\n",
    "        self.is_symmetric = True\n",
    "\n",
    "        # SOLUTION\n",
    "        self._diags_padded = np.zeros((3, self._diag.size), dtype=self.dtype)\n",
    "        self._diags_padded[0, 1:] = self._offdiag\n",
    "        self._diags_padded[1, :] = self._diag\n",
    "        self._diags_padded[2, :-1] = self._offdiag\n",
    "        # SOLUTION END\n",
    "\n",
    "    def _matmul(self, x: np.ndarray) -> np.ndarray:\n",
    "        \"\"\"(Vectorized) matrix-matrix product `y = self @ x`.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        x\n",
    "            Stack of matrices of shape `(D1, ..., Dk, N, L)`, where `self.shape == (N, N)`\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        Stack `y` of matrices of shape `(D1, ..., Dk, N, L)`, where\n",
    "\n",
    "        ..code::\n",
    "            y[i1, ..., ik, :, :] == self @ x[i1, ..., ik, :, :]\n",
    "        \"\"\"\n",
    "        # SOLUTION\n",
    "        res = self._diag[:, None] * x\n",
    "        res[1:] += self._offdiag[:, None] * x[..., :-1, :]\n",
    "        res[:-1] += self._offdiag[:, None] * x[..., 1:, :]\n",
    "        return res\n",
    "        # SOLUTION END\n",
    "\n",
    "    @pn.linops.LinearOperator.broadcast_matmat(method=True)\n",
    "    def _solve(self, b: np.ndarray) -> np.ndarray:\n",
    "        \"\"\"Solves the linear system `self @ x = b`.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        b\n",
    "            Matrix of shape `(N, L)`, where `self.shape == (N, N)`.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        Solution `x` to the linear system of shape `(N, L)`, i.e.\n",
    "\n",
    "        ..code::\n",
    "            y[i1, ..., ik, :, :] == A @ x[i1, ..., ik, :, :]\n",
    "        \"\"\"\n",
    "        # SOLUTION\n",
    "        return scipy.linalg.solve_banded((1, 1), self._diags_padded, b)\n",
    "        # SOLUTION END"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below you will find some code to test your implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 5\n",
    "\n",
    "diag = np.arange(1, n + 1, dtype=np.double)\n",
    "offdiag = 2 * np.arange(1, n, dtype=np.double)\n",
    "\n",
    "A = SymmetricTridiagonalMatrix(diag, offdiag)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To test your implementation of `_matmul`, we can call `LinearOperator.todense`.\n",
    "Its default implementation simply applies the linear operator to a unit matrix, i.e. $A[I_n]$, which produces a dense representation of the matrix.\n",
    "If the following cell produces the matrix\n",
    "\\begin{equation*}\n",
    "\\begin{pmatrix}\n",
    "1 & 2 & 0 & 0 & 0 \\\\\n",
    "2 & 2 & 4 & 0 & 0 \\\\\n",
    "0 & 4 & 3 & 6 & 0 \\\\\n",
    "0 & 0 & 6 & 4 & 8 \\\\\n",
    "0 & 0 & 0 & 8 & 5 \\\\\n",
    "\\end{pmatrix}\n",
    "\\end{equation*}\n",
    "for $d = 5$, then your code is likely correct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A.todense()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The implementation of `_solve` is likely correct if `A.inv() @ A` is approximately equal to the identity matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with np.printoptions(precision=5):\n",
    "    print((A.inv() @ A).todense())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Probabilistic Linear Solvers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now use our matrix-free linear operators to implement a matrix-free probabilistic linear solver (PLS).\n",
    "In the lectures you have seen that a (solution-based) PLS posits a prior $\\mu_0(x) = \\mathcal{N}(x; x_0, \\Sigma_0)$ over the unknown solution of the linear system $A x = b$, and then proceeds to iteratively compute the conditional distributions\n",
    "$$\\mu_i(x) := \\mu_{i - 1}(x \\mid s_i^\\top A x = s_i^\\top b) = \\mu_0(x \\mid s_j^\\top A x = s_j^\\top b\\quad \\forall j \\in \\{1, \\dotsc, i\\})$$\n",
    "for $i < n$.\n",
    "The so-called actions $s_i \\in \\mathbb{R}^n$ are generated by a policy, which depends on $\\mu_{i - 1}$.\n",
    "Due to the closure of Gaussian probability measures under conditioning on affine constraints, the conditional distributions $\\mu_i$ are also Gaussian\n",
    "$$\\mu_i(x) = \\mathcal{N}(x; x_i, \\Sigma_i)$$\n",
    "whose moments admit the recursion\n",
    "\\begin{align*}\n",
    "x_i    & = x_{i - 1} + \\frac{\\alpha_i}{\\eta_i} d_i \\\\\n",
    "\\Sigma_i & = \\Sigma_{i - 1} - \\frac{1}{\\eta_i} d_i d_i^\\top\n",
    "\\end{align*}\n",
    "with\n",
    "\\begin{align*}\n",
    "r_{i - 1} & = b - A[x_{i - 1}] \\\\\n",
    "\\alpha_i  & = s_i^\\top r_{i - 1} \\\\\n",
    "d_i       & = \\Sigma_{i - 1} A^\\top [s_i] \\\\\n",
    "\\eta_i    & = s_i^\\top A d_i.\n",
    "\\end{align*}\n",
    "Partially resolving the recursion for $\\Sigma_i$, we find that\n",
    "$$\\Sigma_i = \\Sigma_0 - D_i D_i^\\top,$$\n",
    "where $D_0 = (\\ ) \\in \\mathbb{R}^{n \\times 0}$ and $D_i = \\begin{pmatrix} D_{i - 1} & \\frac{1}{\\sqrt{\\eta_i}} d_i \\end{pmatrix} \\in \\mathbb{R}^{n \\times i}$.\n",
    "In practice, it is often helpful to implement the recursion in terms of the downdate matrices $D_i$, instead of the full covariance matrix.\n",
    "In this case, we also need that\n",
    "$$d_i = (\\Sigma_0 - D_i D_i^\\top) A^\\top[s_i].$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial, we will assume $A$ to be symmetric, positive-definite and we will write specialized code for a specific choice of prior covariance, namely $\\Sigma_0 = A^{-1}$.\n",
    "Combined with the policy $\\operatorname{Policy}(\\mu_{i - 1}) = b - A x_{i - 1} = r_{i - 1}$, the estimates $x_i$ of the solution coincide with the iterates of the method of conjugate gradients (CG)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-box alert-info\">\n",
    "\n",
    "**Task 2:**\n",
    "Implement a probabilistic linear solver with hard-coded prior $\\mathcal{N}(x_0, A^{-1})$ for the linear system $A x = b$.\n",
    "The scaffolding is already given in the cell below.\n",
    "You just have to implement the inner loop and the CG policy.\n",
    "The solver should return $(x_i, D_i)$ instead of $(x_i, \\Sigma_i)$.\n",
    "Note that, under the prior with $\\Sigma_0 = A^{-1}$, the expression for $d_i$ simplifies to\n",
    "\\begin{equation*}\n",
    "    d_i = (I - D_i D_i^\\top A) s_i.\n",
    "\\end{equation*}\n",
    "Pay close attention to keeping the algorithm matrix-free and use the associativity of matrix multiplication to your advantage.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CGPolicy:\n",
    "    def __call__(self, *, i: int, x: np.ndarray, D: np.ndarray, r: np.ndarray, **kwargs):\n",
    "        # SOLUTION\n",
    "        return r\n",
    "        # SOLUTION END\n",
    "\n",
    "class RandomPolicy:\n",
    "    def __init__(self, rng: np.random.Generator) -> None:\n",
    "        self._rng = rng\n",
    "\n",
    "    def __call__(self, *, r: np.ndarray, **kwargs) -> None:\n",
    "        return self._rng.normal(size=r.shape)\n",
    "\n",
    "class CoordinatePolicy:\n",
    "    def __call__(self, *, i: int, r: np.ndarray, **kwargs):\n",
    "        s = np.zeros_like(r)\n",
    "        s[i - 1] = 1\n",
    "        return s\n",
    "\n",
    "class EigenvectorPolicy:\n",
    "    def __init__(self, A: ArrayLike, desc: bool = True):\n",
    "        _, eigvecs = np.linalg.eigh(A)\n",
    "\n",
    "        self._eigvecs = eigvecs[:, ::-1] if desc else eigvecs\n",
    "\n",
    "    def __call__(self, *, i: int, **kwargs):\n",
    "        return self._eigvecs[:, i]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def problinsolve(\n",
    "    A: pn.linops.LinearOperatorLike,\n",
    "    b: ArrayLike,\n",
    "    x0: ArrayLike | None = None,\n",
    "    policy: callable = CGPolicy(),\n",
    "    abstol: float = 1e-10,\n",
    "    reltol: float = 1e-12,\n",
    "    maxiter: int | None = None,\n",
    "):\n",
    "    A = pn.linops.aslinop(A)\n",
    "    b = np.asarray(b)\n",
    "\n",
    "    if x0 is None:\n",
    "        x0 = np.zeros(A.shape[1])\n",
    "\n",
    "    x0 = np.asarray(x0, dtype=np.promote_types(A.dtype, b.dtype))\n",
    "\n",
    "    # Initial Residual\n",
    "    r0 = b - A @ x0\n",
    "\n",
    "    # Initial Belief\n",
    "    x = x0\n",
    "    D = np.zeros((A.shape[1], 0), dtype=x.dtype, order=\"F\")\n",
    "\n",
    "    r = r0\n",
    "\n",
    "    # Stopping Criteria\n",
    "    tol = max(abstol, reltol * np.linalg.norm(b, 2))\n",
    "\n",
    "    if maxiter is None:\n",
    "        maxiter = x.size\n",
    "\n",
    "    i = 0\n",
    "\n",
    "    while i < maxiter and np.linalg.norm(r, 2) > tol:\n",
    "        s = policy(i=i, x=x, D=D, r=r)\n",
    "\n",
    "        # SOLUTION\n",
    "        alpha = s @ r\n",
    "        As = A @ s\n",
    "        d = s - D @ (D.T @ As)\n",
    "        eta = As @ d\n",
    "\n",
    "        x = x + (alpha / eta) * d\n",
    "        D = np.concatenate((D, np.sqrt(1 / eta) * d[:, None]), axis=1)\n",
    "        # END SOLUTION\n",
    "\n",
    "        i = i + 1\n",
    "        r = b - A @ x\n",
    "\n",
    "    return x, D"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To test your implementation, we use `probnum` to generate random linear systems with spd system matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from probnum.problems.zoo.linalg import random_linear_system, random_spd_matrix\n",
    "\n",
    "n = 2\n",
    "\n",
    "linsys = random_linear_system(\n",
    "    np.random.default_rng(243),\n",
    "    random_spd_matrix,\n",
    "    dim=n,\n",
    ")\n",
    "\n",
    "linsys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, D = problinsolve(\n",
    "    linsys.A,\n",
    "    linsys.b,\n",
    "    # maxiter=n / 2,\n",
    ")\n",
    "\n",
    "np.testing.assert_allclose(x, linsys.solution)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also take a look at the full covariance matrix.\n",
    "For $i \\approx d$, it should be close to zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma = pn.linops.aslinop(linsys.A).inv() - D @ D.T\n",
    "sigma.todense()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The matrix $D_i D_i^\\top$ is a low-rank approximation of $A^{-1}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(D @ D.T) @ linsys.A"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Application: IterGP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have a working implementation of a probabilistic linear solver, we will apply it to the linear system arising in Gaussian process (GP) regression.\n",
    "This strategy is the core idea behind IterGP and was already introduced in the lecture."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In GP regression, we posit a Gaussian process prior $f \\sim \\mathcal{GP}(m, k)$ over an unknown function and we assume that we measure the function at a set $X = (x_i)_{i = 1}^n \\in \\mathbb{X}^n$ of inputs according to the measurement model $f(X) + \\epsilon$, where $\\epsilon \\sim \\mathcal{N}(0, \\lambda^2 I_n)$ with $\\epsilon \\perp\\!\\!\\perp f$.\n",
    "We then measure $y \\in \\mathbb{R}^n$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean = pn.functions.Zero(())\n",
    "cov = pn.randprocs.covfuncs.Matern((), 1.5)\n",
    "\n",
    "prior = pn.randprocs.GaussianProcess(mean, cov)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domain = (-1, 1)\n",
    "\n",
    "n = 10\n",
    "\n",
    "# Observation model\n",
    "noise_std = 0.2\n",
    "epsilon = pn.randvars.Normal(\n",
    "    np.zeros(n),\n",
    "    noise_std ** 2 * pn.linops.Identity(n),\n",
    ")\n",
    "\n",
    "# Generate data by sampling from the prior and the observation model\n",
    "rng = np.random.default_rng(324)\n",
    "X = rng.uniform(*domain, n)\n",
    "y = prior(X).sample(rng) + epsilon.sample(rng)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xs_plot = np.linspace(*domain, 100)\n",
    "\n",
    "prior.plot(\n",
    "    plt.gca(),\n",
    "    xs_plot,\n",
    "    rng=np.random.default_rng(42),\n",
    "    num_samples=10,\n",
    "    label=r\"Prior $f$\",\n",
    ")\n",
    "\n",
    "plt.errorbar(\n",
    "    X,\n",
    "    y,\n",
    "    yerr=1.96 * noise_std,\n",
    "    fmt=\"+\",\n",
    "    capsize=5,\n",
    "    label=r\"Data $(X, y)$\",\n",
    ")\n",
    "\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To incorporate the measurements into our model, we compute the conditional process $f \\mid f(X) + \\epsilon = y \\sim \\mathcal{GP}(m^y, k^y)$, where\n",
    "\\begin{align*}\n",
    "    m^y(x)        & = m(x) + k(x, X) w\\\\\n",
    "    k^y(x_1, x_2) & = k(x_1, x_2) - k(x_1, X) \\hat{K}^{-1} k(X, x_2)\n",
    "\\end{align*}\n",
    "with $w = \\hat{K}^{-1} (y - m(X))$ and $\\hat{K} = k(X, X) + \\lambda^2 I_n$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "html"
    }
   },
   "source": [
    "<div class=\"alert alert-box alert-info\">\n",
    "\n",
    "**Task 3:**\n",
    "Implement the expressions for the conditional mean and the conditional covariance function in the `ConditionalGaussianProcess` class below.\n",
    "Use a Cholesky decomposition to multiply with $\\hat{K}^{-1}$.\n",
    "\n",
    "*Hints:*\n",
    "- the methods `self._Khat.solve(..)` and `self._Khat.inv()` use the Cholesky decomposition to multiply with the inverse and to construct the inverse in a lazy fashion, respectively.\n",
    "- the method `Covariance._kxX` might come in handy for implementing the broadcasting in $k(x_1, X) \\hat{K}^{-1} k(X, x_2)$ correctly\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConditionalGaussianProcess(pn.randprocs.GaussianProcess):\n",
    "    def __init__(\n",
    "        self,\n",
    "        prior: pn.randprocs.GaussianProcess,\n",
    "        X: ArrayLike,\n",
    "        y: ArrayLike,\n",
    "        noise: pn.randvars.Normal,\n",
    "    ) -> None:\n",
    "        self._prior = prior\n",
    "\n",
    "        self._X = np.asarray(X)\n",
    "        self._y = np.asarray(y)\n",
    "        self._noise = noise\n",
    "\n",
    "        self._Khat = self._prior.cov.linop(self._X) + pn.linops.aslinop(self._noise.cov)\n",
    "        self._Khat.is_symmetric = True\n",
    "\n",
    "        # SOLUTION\n",
    "        self._w = self._Khat.solve(self._y - (self._prior.mean(X) + self._noise.mean))\n",
    "        # END SOLUTION\n",
    "\n",
    "        super().__init__(\n",
    "            mean=ConditionalGaussianProcess.Mean(prior.mean, prior.cov, self._X, self._w),\n",
    "            cov=ConditionalGaussianProcess.CovarianceFunction(prior.cov, self._X, self._Khat),\n",
    "        )\n",
    "\n",
    "    class Mean(pn.functions.Function):\n",
    "        def __init__(\n",
    "            self,\n",
    "            prior_mean: pn.functions.Function,\n",
    "            prior_cov: pn.randprocs.covfuncs.CovarianceFunction,\n",
    "            X: np.ndarray,\n",
    "            w: np.ndarray,\n",
    "        ) -> None:\n",
    "            self._prior_mean = prior_mean\n",
    "            self._prior_cov = prior_cov\n",
    "\n",
    "            self._X = X\n",
    "\n",
    "            self._w = w\n",
    "\n",
    "            super().__init__(\n",
    "                input_shape=prior_mean.input_shape,\n",
    "                output_shape=prior_mean.output_shape,\n",
    "            )\n",
    "\n",
    "        def _evaluate(self, x: np.ndarray) -> np.ndarray:\n",
    "            \"\"\"Evaluate the conditional mean function at location `x`.\n",
    "\n",
    "            Parameters\n",
    "            ----------\n",
    "            x\n",
    "                Single input of shape `()` or a batch of inputs of shape `(N,)`.\n",
    "\n",
    "            Returns\n",
    "            -------\n",
    "            Posterior mean evaluated at the given input(s).\n",
    "            Output shape should be `()` for a single input or `(N,)` for a batch of inputs.\n",
    "            \"\"\"\n",
    "            # SOLUTION\n",
    "            return self._prior_mean(x) + self._prior_cov.linop(x, self._X) @ self._w\n",
    "            # END SOLUTION\n",
    "\n",
    "    class CovarianceFunction(pn.randprocs.covfuncs.CovarianceFunction):\n",
    "        def __init__(\n",
    "            self,\n",
    "            prior_cov: pn.randprocs.covfuncs.CovarianceFunction,\n",
    "            X: np.ndarray,\n",
    "            Khat: pn.linops.LinearOperator,\n",
    "        ):\n",
    "            self._prior_cov = prior_cov\n",
    "\n",
    "            self._X = X\n",
    "\n",
    "            self._Khat = Khat\n",
    "\n",
    "            super().__init__(\n",
    "                input_shape=prior_cov.input_shape,\n",
    "                output_shape_0=prior_cov.output_shape_0,\n",
    "                output_shape_1=prior_cov.output_shape_1,\n",
    "            )\n",
    "\n",
    "        def _evaluate(self, x0: np.ndarray, x1: np.ndarray | None) -> np.ndarray:\n",
    "            \"\"\"Evaluates the conditional covariance function between locations `x0` and `x1`.\n",
    "\n",
    "            Parameters\n",
    "            ----------\n",
    "            x0\n",
    "                (Batch of) input(s) with shape `batch_shape_0`.\n",
    "            x1\n",
    "                (Batch of) input(s) with shape `batch_shape_1`.\n",
    "                If this is `None` it is assumed that `x1 == x0`.\n",
    "\n",
    "            Returns\n",
    "            -------\n",
    "            Covariance between function values at `x0` and `x1` as an array with shape\n",
    "            `np.broadcast_shapes(batch_shape_0, batch_shape_1)`.\n",
    "            This function has to broadcast correctly according to the rules of NumPy broadcasting.\n",
    "            For example, if `x0.shape == (100, 1)` and `x1.shape == (1, 100)`, then\n",
    "            `self._evaluate(x0, x1).shape == (100, 100)` and\n",
    "            `self._evaluate(x0, x1)[i, j] = self._evaluate(x0[i, 0], x1[j, 0])`.\n",
    "            \"\"\"\n",
    "            # SOLUTION\n",
    "            k_x0_X = self._kxX(x0)\n",
    "            k_x1_X = self._kxX(x1) if x1 is not None else k_x0_X\n",
    "\n",
    "            return self._prior_cov(x0, x1) - (k_x0_X[..., None, :] @ self._Khat.inv() @ k_x1_X[..., :, None])[..., 0, 0]\n",
    "            # END SOLUTION\n",
    "\n",
    "        def _kxX(self, x: np.ndarray) -> np.ndarray:\n",
    "            \"\"\"Computes test-train cross-covariances according to the prior covariance function.\n",
    "\n",
    "            Parameters\n",
    "            ----------\n",
    "            x\n",
    "                (Batch of) input(s) with shape `batch_shape`.\n",
    "\n",
    "            Returns\n",
    "            -------\n",
    "            Prior test-train cross-covariances with shape `batch_shape + (n,)`, where `n == X.shape[0]`.\n",
    "            \"\"\"\n",
    "            return self._prior_cov(np.expand_dims(x, axis=x.ndim - self.input_ndim_0), self._X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "posterior_cho = ConditionalGaussianProcess(\n",
    "    prior,\n",
    "    X,\n",
    "    y,\n",
    "    epsilon,\n",
    ")\n",
    "\n",
    "xs_plot = np.linspace(*domain, 100)\n",
    "\n",
    "posterior_cho.plot(\n",
    "    plt.gca(),\n",
    "    xs_plot,\n",
    "    rng=np.random.default_rng(42),\n",
    "    num_samples=10,\n",
    "    label=r\"Posterior $f \\mid f(X) + \\epsilon = y$\",\n",
    ")\n",
    "\n",
    "plt.errorbar(\n",
    "    X,\n",
    "    y,\n",
    "    yerr=1.96 * noise_std,\n",
    "    fmt=\"+\",\n",
    "    capsize=5,\n",
    "    label=r\"Data $(X, y)$\",\n",
    ")\n",
    "\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Cholesky decomposition used above has a time complexity of $\\frac{1}{3} n^3$ (to leading order) and a memory complexity of $\\mathcal{O}(n^2)$.\n",
    "This quickly becomes prohibitive for large datasets.\n",
    "\n",
    "To address this, the lecture introduced IterGP, a GP approximation method based on a matrix-free probabilistic linear solver.\n",
    "In essence, in the following, we will use our PLS to solve the system $\\hat{K} w = y - m(X)$.\n",
    "This yields an estimate $w_i$ of the solution and a downdate matrix $D_i$.\n",
    "As seen in the lecture, we can use these two quantities to implement the approximate conditional mean and covariance functions:\n",
    "\\begin{align*}\n",
    "    m^y_i(x) & = m(x) + k(x, X) w_i \\\\\n",
    "    k^y_i(x_1, x_2) & = k(x_1, x_2) - k(x_1, X) D_i D_i^\\top k(X, x_2).\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-box alert-info\">\n",
    "\n",
    "**Task 4:**\n",
    "Implement the expressions for the conditional mean $m^y_i$ and the conditional covariance function $k^y_i$ in the `IterGPConditionalGaussianProcess` class below.\n",
    "Again, use the associativity of matrix products to your advantage.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class IterGPConditionalGaussianProcess(pn.randprocs.GaussianProcess):\n",
    "    def __init__(\n",
    "        self,\n",
    "        prior: pn.randprocs.GaussianProcess,\n",
    "        X: ArrayLike,\n",
    "        y: ArrayLike,\n",
    "        noise: pn.randvars.Normal,\n",
    "        **problinsolve_kwargs,\n",
    "    ) -> None:\n",
    "        self._prior = prior\n",
    "\n",
    "        self._X = np.asarray(X)\n",
    "        self._y = np.asarray(y)\n",
    "        self._noise = noise\n",
    "\n",
    "        self._gram = self._prior.cov.linop(self._X) + pn.linops.aslinop(self._noise.cov)\n",
    "\n",
    "        self._w, self._D = problinsolve(\n",
    "            self._gram,\n",
    "            self._y - (self._prior.mean(X) + self._noise.mean),\n",
    "            **problinsolve_kwargs,\n",
    "        )\n",
    "\n",
    "        super().__init__(\n",
    "            mean=IterGPConditionalGaussianProcess.Mean(prior.mean, prior.cov, self._X, self._w),\n",
    "            cov=IterGPConditionalGaussianProcess.CovarianceFunction(prior.cov, self._X, self._D),\n",
    "        )\n",
    "\n",
    "    class Mean(pn.functions.Function):\n",
    "        def __init__(\n",
    "            self,\n",
    "            prior_mean: pn.functions.Function,\n",
    "            prior_cov: pn.randprocs.covfuncs.CovarianceFunction,\n",
    "            X: np.ndarray,\n",
    "            w: np.ndarray,\n",
    "        ) -> None:\n",
    "            self._prior_mean = prior_mean\n",
    "            self._prior_cov = prior_cov\n",
    "\n",
    "            self._X = X\n",
    "\n",
    "            self._w = w\n",
    "\n",
    "            super().__init__(\n",
    "                input_shape=prior_mean.input_shape,\n",
    "                output_shape=prior_mean.output_shape,\n",
    "            )\n",
    "\n",
    "        def _evaluate(self, x: np.ndarray) -> np.ndarray:\n",
    "            \"\"\"Evaluate the IterGP conditional mean function at location `x`.\n",
    "\n",
    "            Parameters\n",
    "            ----------\n",
    "            x\n",
    "                Single input of shape `()` or a batch of inputs of shape `(N,)`.\n",
    "\n",
    "            Returns\n",
    "            -------\n",
    "            Posterior mean evaluated at the given input(s).\n",
    "            Output shape should be `()` for a single input or `(N,)` for a batch of inputs.\n",
    "            \"\"\"\n",
    "            # SOLUTION\n",
    "            return self._prior_mean(x) + self._prior_cov.linop(x, self._X) @ self._w\n",
    "            # END SOLUTION\n",
    "\n",
    "    class CovarianceFunction(pn.randprocs.covfuncs.CovarianceFunction):\n",
    "        def __init__(\n",
    "            self,\n",
    "            prior_cov: pn.randprocs.covfuncs.CovarianceFunction,\n",
    "            X: np.ndarray,\n",
    "            D: np.ndarray,\n",
    "        ):\n",
    "            self._prior_cov = prior_cov\n",
    "\n",
    "            self._X = X\n",
    "\n",
    "            self._D = D\n",
    "\n",
    "            super().__init__(\n",
    "                input_shape=prior_cov.input_shape,\n",
    "                output_shape_0=prior_cov.output_shape_0,\n",
    "                output_shape_1=prior_cov.output_shape_1,\n",
    "            )\n",
    "\n",
    "        def _evaluate(self, x0: np.ndarray, x1: np.ndarray | None) -> np.ndarray:\n",
    "            \"\"\"Evaluates the IterGP conditional covariance function between locations `x0` and `x1`.\n",
    "\n",
    "            Parameters\n",
    "            ----------\n",
    "            x0\n",
    "                (Batch of) input(s) with shape `batch_shape_0`.\n",
    "            x1\n",
    "                (Batch of) input(s) with shape `batch_shape_1`.\n",
    "                If this is `None` it is assumed that `x1 == x0`.\n",
    "\n",
    "            Returns\n",
    "            -------\n",
    "            Covariance between function values at `x0` and `x1` as an array with shape\n",
    "            `np.broadcast_shapes(batch_shape_0, batch_shape_1)`.\n",
    "            This function has to broadcast correctly according to the rules of NumPy broadcasting.\n",
    "            For example, if `x0.shape == (100, 1)` and `x1.shape == (1, 100)`, then\n",
    "            `self._evaluate(x0, x1).shape == (100, 100)` and\n",
    "            `self._evaluate(x0, x1)[i, j] = self._evaluate(x0[i, 0], x1[j, 0])`.\n",
    "            \"\"\"\n",
    "            # SOLUTION\n",
    "            k_x0_X_D = self._kxX(x0) @ self._D\n",
    "            k_x1_X_D = self._kxX(x1) @ self._D if x1 is not None else k_x0_X_D\n",
    "\n",
    "            return self._prior_cov(x0, x1) - (k_x0_X_D[..., None, :] @ k_x1_X_D[..., :, None])[..., 0, 0]\n",
    "            # END SOLUTION\n",
    "\n",
    "        def _kxX(self, x: np.ndarray) -> np.ndarray:\n",
    "            \"\"\"Computes test-train cross-covariances according to the prior covariance function.\n",
    "\n",
    "            Parameters\n",
    "            ----------\n",
    "            x\n",
    "                (Batch of) input(s) with shape `batch_shape`.\n",
    "\n",
    "            Returns\n",
    "            -------\n",
    "            Prior test-train cross-covariances with shape `batch_shape + (n,)`, where `n == X.shape[0]`.\n",
    "            \"\"\"\n",
    "            return self._prior_cov(np.expand_dims(x, axis=x.ndim - self.input_ndim_0), self._X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "posterior_itergp = IterGPConditionalGaussianProcess(\n",
    "    prior,\n",
    "    X,\n",
    "    y,\n",
    "    epsilon,\n",
    "    maxiter=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "posterior_itergp.plot(\n",
    "    ax,\n",
    "    xs_plot,\n",
    "    rng=np.random.default_rng(42),\n",
    "    num_samples=10,\n",
    "    label=\"IterGP Posterior\",\n",
    ")\n",
    "\n",
    "posterior_cho.plot(\n",
    "    ax,\n",
    "    xs_plot,\n",
    "    label=\"Exact Posterior\",\n",
    "    color=\"C3\",\n",
    ")\n",
    "\n",
    "plt.errorbar(\n",
    "    X,\n",
    "    y,\n",
    "    yerr=1.96 * noise_std,\n",
    "    fmt=\"+\",\n",
    "    capsize=5,\n",
    "    label=r\"Data $(X, y)$\",\n",
    ")\n",
    "\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets\n",
    "\n",
    "%matplotlib widget\n",
    "\n",
    "fig, ax = plt.subplots(num=\"IterGP\")\n",
    "\n",
    "def interact(policy: str, maxiter: int):\n",
    "    if policy == \"CG\":\n",
    "        policy = CGPolicy()\n",
    "    elif policy == \"Coordinate\":\n",
    "        policy = CoordinatePolicy()\n",
    "    elif policy == \"Random\":\n",
    "        policy = RandomPolicy(np.random.default_rng(89745))\n",
    "    else:\n",
    "        raise IndexError(\"Unknown policy\")\n",
    "\n",
    "    posterior_itergp = IterGPConditionalGaussianProcess(\n",
    "        prior,\n",
    "        X,\n",
    "        y,\n",
    "        epsilon,\n",
    "        policy=policy,\n",
    "        maxiter=maxiter,\n",
    "    )\n",
    "\n",
    "    ax.cla()\n",
    "    posterior_itergp.plot(\n",
    "        ax,\n",
    "        xs_plot,\n",
    "        rng=np.random.default_rng(42),\n",
    "        num_samples=10,\n",
    "        label=\"IterGP Posterior\",\n",
    "    )\n",
    "\n",
    "    posterior_cho.plot(\n",
    "        ax,\n",
    "        xs_plot,\n",
    "        label=\"Exact Posterior\",\n",
    "        color=\"C3\",\n",
    "    )\n",
    "\n",
    "    plt.errorbar(\n",
    "        X,\n",
    "        y,\n",
    "        yerr=1.96 * noise_std,\n",
    "        fmt=\"+\",\n",
    "        capsize=5,\n",
    "        label=r\"Data $(X, y)$\",\n",
    "    )\n",
    "\n",
    "    ax.legend()\n",
    "\n",
    "ipywidgets.interactive(\n",
    "    interact,\n",
    "    policy=ipywidgets.ToggleButtons(\n",
    "        options=['CG', 'Coordinate', 'Random'],\n",
    "        description='Policy',\n",
    "    ),\n",
    "    maxiter=ipywidgets.IntSlider(\n",
    "        value=int(0.5 * n),\n",
    "        min=0,\n",
    "        max=n,\n",
    "        description=\"maxiter (%)\",\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Application: Galerkin Methods for Linear PDEs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Consider the linear partial differential equation (PDE)\n",
    "$$-\\Delta u = f$$\n",
    "with $u \\colon \\overline{\\mathbb{D}} \\to \\mathbb{R}$ and $\\mathbb{D} \\subset \\mathbb{R}^d$ bounded and open, subject to so-called Dirichlet boundary conditions\n",
    "$$u(x) = g(x)$$\n",
    "on $x \\in \\partial \\mathbb{D}$, where $f$ and $g$ are known.\n",
    "\n",
    "In general, this boundary value problem can not be solved in closed-form, so we have to use a numerical method to approximate it's solution.\n",
    "To this end, we will focus on the **Galerkin methods**, a particularly popular family of numerical methods for solving PDEs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<hr/>\n",
    "\n",
    "In the following, I will give a very terse, intuitive overview over Galerkin methods.\n",
    "\n",
    "<div class=\"alert alert-box alert-warning\">\n",
    "\n",
    "**Disclaimer:**\n",
    "Since the main focus of this tutorial is on linear algebra and not on PDEs, the following overview is very crude.\n",
    "For the sake of brevity, I'm glossing over many important theoretical concepts (e.g. Sobolev spaces).\n",
    "More complete introductions to these methods and the related theoretical concepts can be found in\n",
    "\n",
    "- L.C. Evans. *Partial Differential Equations: Second Edition*. AMS, 2010.\n",
    "- C.A.J. Fletcher. *Computational Galerkin Methods*. Springer, 1984.\n",
    "- H.P. Langtangen and K.-A. Mardal. *Introduction to Numerical Methods for Variational Problems*. Springer, 2019.\n",
    "\n",
    "Probabilistic numerical versions of (Petrov-)Galerkin methods and more general methods of weighted residuals are introduced in\n",
    "\n",
    "M. Pförtner, I. Steinwart, P. Hennig, and J. Wenger. Physics-Informed Gaussian Process Regression Generalizes Linear PDE Solvers. 2022.\n",
    "\n",
    "</div>\n",
    "\n",
    "The first step in applying Galerkin methods is to convert the PDE into its so-called **weak** or **variational formulation**.\n",
    "To this end, let $v: \\mathbb{D} \\to \\mathbb{R}$ be a so-called **test function** with $v \\vert_{\\partial \\mathbb{D}} = 0$.\n",
    "Integrating $v$ against both sides of the PDE yields\n",
    "\\begin{equation*}\n",
    "    - \\int_{\\mathbb{D}} v(x) \\Delta u(x)\\;\\mathrm{d}x\n",
    "    = \\int_{\\mathbb{D}} v(x) f(x)\\;\\mathrm{d}x.\n",
    "\\end{equation*}\n",
    "If this equation holds for a sufficiently large set of test functions $v$, then it is equivalent to the original PDE.\n",
    "Applying Green's first identity (think: multivariate integration by parts) to the left-hand site of the equation, we obtain\n",
    "\\begin{equation*}\n",
    "    \\underbrace{\\int_{\\mathbb{D}} \\langle \\nabla v(x), \\nabla u(x) \\rangle\\;\\mathrm{d}x}_{=: B[u, v]}\n",
    "    = \\underbrace{\\int_{\\mathbb{D}} v(x) f(x)\\;\\mathrm{d}x}_{=: l[v]},\n",
    "\\end{equation*}\n",
    "where the boundary integral vanishes due to the fact that $v \\vert_{\\partial \\mathbb{D}} = 0$.\n",
    "Note that the expression on the left-hand side makes sense even if $u$ is once differentiable almost everywhere.\n",
    "This equation, i.e.\n",
    "\\begin{equation*}\n",
    "    B[u, v] = l[v] \\qquad \\qquad \\forall v\n",
    "\\end{equation*}\n",
    "is the so-called weak formulation of the original PDE.\n",
    "\n",
    "The weak formulation constitutes a system of infinitely many equations in infinitely many unknowns.\n",
    "To render the problem tractable, Galerkin methods proceed by only requiring it to hold for finitely many test functions $\\{ \\psi_i \\}_{i = 1}^n$ and by making the \"ansatz\"\n",
    "\\begin{equation*}\n",
    "  u(x) = \\sum_{j = 1}^n w_j \\phi_j(x)\n",
    "\\end{equation*}\n",
    "where the so-called **trial functions** $\\{\\phi_i\\}_{j = 1}^n$ are linearly independent and fulfill the boundary conditions.\n",
    "In other words, we approximate the solution of the PDE in an $n$-dimensional subspace of the infinite-dimensional solution space.\n",
    "If $\\psi_i$ and $\\phi_j$ are locally supported, then the resuting approximation scheme is referred to as a **finite element method** (FEM).\n",
    "Since $B$ is bilinear, this amounts to\n",
    "\\begin{equation*}\n",
    "    \\sum_{j = 1}^n \\underbrace{B[\\phi_j, \\psi_i]}_{=: A_{ij}} w_j = \\underbrace{l[\\psi_i]}_{=: b_i} \\qquad \\qquad i = 1, \\dotsc, n,\n",
    "\\end{equation*}\n",
    "or, equivalently, we need to solve the linear system $A w = b$.\n",
    "\n",
    "<hr/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For simplicity, we will set $d = 1$ and $\\mathbb{D} = [-1, 1]$, as well as $f(x) = c$ and $g(x) = 0$.\n",
    "Then the BVP amounts to solving\n",
    "\\begin{align*}\n",
    "    -u''(x) & = c \\\\\n",
    "    u(-1)   & = 0 \\\\\n",
    "    u(1)    & = 0\n",
    "\\end{align*}\n",
    "for $u$.\n",
    "In this case, we have $u(x) = \\frac{c}{2}(1 - x^2)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domain = (-1, 1)\n",
    "xs_plot = np.linspace(*domain, 500)\n",
    "\n",
    "c = 2\n",
    "\n",
    "def u_star(x: ArrayLike):\n",
    "    return c / 2 * (1 - x ** 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For solving the PDE with a Galerkin method, we choose so-called linear Lagrange elements as test and trial functions, i.e.\n",
    "$$\n",
    "\\phi_i(x) \n",
    "= \\psi_i(x)\n",
    "= \\begin{cases}\n",
    "    \\frac{x - x_{i - 1}}{x_i - x_{i - 1}} & \\text{if}\\ x_{i - 1} \\le x \\le x_i, \\\\\n",
    "    \\frac{x_{i + 1} - x}{x_{i + 1} - x_i} & \\text{if}\\ x_i < x \\le x_{i + 1}, \\\\\n",
    "    0 & \\text{otherwise}.\n",
    "\\end{cases}\n",
    "$$\n",
    "on a grid $-1 = x_0 < x_1 < \\dotsb < x_n < x_{n + 1} = 1$ which we choose to be regular, i.e. $x_i = (-1) + 2 \\frac{i}{n + 1}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class UnivariateLinearLagrangeElements(pn.functions.Function):\n",
    "    def __init__(self, grid: ArrayLike):\n",
    "        self._grid = np.asarray(grid)\n",
    "\n",
    "        if self._grid.ndim != 1 or self._grid.size < 3:\n",
    "            raise ValueError(\"`grid` must be a 1D array with at least 3 elements\")\n",
    "\n",
    "        self._num_elements = self._grid.size - 2\n",
    "\n",
    "        super().__init__(\n",
    "            input_shape=(),\n",
    "            output_shape=(self._num_elements,),\n",
    "        )\n",
    "\n",
    "    def _evaluate(self, x: np.ndarray) -> np.ndarray:\n",
    "        return np.maximum(\n",
    "            0,\n",
    "            np.where(\n",
    "                x[..., None] < self._grid[1:-1],\n",
    "                (x[..., None] - self._grid[:-2]) / (self._grid[1:-1] - self._grid[:-2]),\n",
    "                (self._grid[2:] - x[..., None]) / (self._grid[2:] - self._grid[1:-1]),\n",
    "            )\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 5\n",
    "xs_grid = np.linspace(*domain, n + 2)\n",
    "phis = UnivariateLinearLagrangeElements(xs_grid)\n",
    "\n",
    "plt.plot(\n",
    "    xs_plot,\n",
    "    phis(xs_plot),\n",
    "    label=[rf\"$\\phi_{{{i}}}$\" for i in range(n)],\n",
    ")\n",
    "\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The $\\phi_j$ are piecewise linear and locally supported.\n",
    "Moreover, their linear span amounts to the set of piecewise linear functions that interpolate the coefficients $w_i$ on the grid nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(\n",
    "    xs_plot,\n",
    "    phis(xs_plot) * u_star(xs_grid[1:-1]),\n",
    "    label=[rf\"$w_{{{i}}} \\phi_{{{i}}}$\" for i in range(n)],\n",
    ")\n",
    "\n",
    "plt.plot(\n",
    "    xs_plot,\n",
    "    phis(xs_plot) @ u_star(xs_grid[1:-1]),\n",
    "    label=r\"$\\sum_{j = 1}^n w_j \\phi_j$\",\n",
    "    color=\"black\",\n",
    ")\n",
    "\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Computing the entries of the stiffness matrix $A = (B[\\phi_j, \\psi_i])_{i,j = 1}^n$ constructed by the Galerkin method yields\n",
    "$$\n",
    "A =\n",
    "\\frac{1}{\\delta x} \\cdot\n",
    "\\begin{pmatrix}\n",
    " 2     & -1     &  0     & \\cdots &  0     \\\\\n",
    "-1     &  2     & -1     & \\ddots & \\vdots \\\\\n",
    " 0     & -1     &  2     & \\ddots &  0     \\\\\n",
    "\\vdots & \\ddots & \\ddots & \\ddots & -1     \\\\\n",
    " 0     & \\cdots &  0     & -1     &  2\n",
    "\\end{pmatrix}\n",
    "$$\n",
    "i.e. a tridiagonal Toeplitz matrix.\n",
    "For the force vector $b = (l[\\psi_i])_{i = 1}^n$, we obtain $b_i = c \\delta_x$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta_x = (domain[1] - domain[0]) / (n + 1)\n",
    "\n",
    "A_fem = SymmetricTridiagonalMatrix(\n",
    "    np.full(n, 2 / delta_x),\n",
    "    np.full(n - 1, -1 / delta_x),\n",
    ")\n",
    "\n",
    "b_fem = np.full(n, c * delta_x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now solve the resulting linear system using our probabilistic linear solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w_fem, D_fem = problinsolve(\n",
    "    A_fem,\n",
    "    b_fem,\n",
    ")\n",
    "\n",
    "w_fem"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-box alert-info\">\n",
    "\n",
    "**Task 5:**\n",
    "The output of the PLS for the linear system $A x = b$ can be interpreted as the random variable $\\mathrm{x} \\sim \\mathcal{N}(x_i, \\Sigma_i)$.\n",
    "Implement a function `pls_randvar` that transforms the current return value of the PLS function into a [`probnum.randvars.Normal`](https://probnum.readthedocs.io/en/v0.1.6/automod/probnum.randvars.Normal.html) object.\n",
    "The scaffolding is given below.\n",
    "\n",
    "*Hint:* `LinearOperator.inv()` could come in handy for performance and numerical stability.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pls_randvar(x: np.ndarray, D: np.ndarray, A: LinearOperatorLike) -> pn.randvars.Normal:\n",
    "    A = pn.linops.aslinop(A)\n",
    "    D = pn.linops.aslinop(D)\n",
    "\n",
    "    # SOLUTION\n",
    "    Sigma = A.inv() - D @ D.T\n",
    "    Sigma.is_symmetric = True\n",
    "\n",
    "\n",
    "    return pn.randvars.Normal(x, Sigma)\n",
    "    # END SOLUTION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w_fem_randvar = pls_randvar(\n",
    "    w_fem,\n",
    "    D_fem,\n",
    "    A_fem,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall that $$u(x) = \\sum_{j = 1}^n w_j \\phi_j(x).$$\n",
    "Since our PLS gives us acces to $w \\in \\mathbb{R}^n$ in the form of a random variable $\\mathrm{w} \\sim \\mathcal{N}(w_i, \\Sigma_i)$, we have\n",
    "$$u = \\phi(\\cdot)^\\top \\mathrm{w} \\sim \\mathcal{GP}(\\phi(\\cdot)^\\top w_i, \\phi(\\cdot)^\\top \\Sigma_i \\phi(\\bullet)),$$\n",
    "i.e. the pushforward of our belief about the solution $w$ through the trial functions is a parametric GP that encodes our belief about the PDE solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_fem = linpde_gp.randprocs.ParametricGaussianProcess(\n",
    "    weights=w_fem_randvar,\n",
    "    feature_fn=phis,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_fem.plot(\n",
    "    plt.gca(),\n",
    "    xs_plot,\n",
    "    rng=np.random.default_rng(2345),\n",
    "    num_samples=10,\n",
    "    label=r\"$\\phi(\\cdot)^{\\top} \\mathrm{w}$\",\n",
    ")\n",
    "\n",
    "plt.plot(\n",
    "    xs_plot,\n",
    "    u_star(xs_plot),\n",
    "    label=r\"$u^\\star$\",\n",
    ")\n",
    "\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets\n",
    "\n",
    "%matplotlib widget\n",
    "\n",
    "fig, ax = plt.subplots(num=\"FEM\")\n",
    "\n",
    "def interact(policy: str, maxiter: int):\n",
    "    if policy == \"CG\":\n",
    "        policy = CGPolicy()\n",
    "    elif policy == \"Coordinate\":\n",
    "        policy = CoordinatePolicy()\n",
    "    elif policy == \"Random\":\n",
    "        policy = RandomPolicy(np.random.default_rng(89745))\n",
    "    elif policy == \"Eigenvector (desc)\":\n",
    "        policy = EigenvectorPolicy(A_fem.todense(), desc=True)\n",
    "    elif policy == \"Eigenvector (asc)\":\n",
    "        policy = EigenvectorPolicy(A_fem.todense(), desc=False)\n",
    "    else:\n",
    "        raise IndexError(\"Unknown policy\")\n",
    "\n",
    "    u_fem = linpde_gp.randprocs.ParametricGaussianProcess(\n",
    "        weights=pls_randvar(\n",
    "            *problinsolve(\n",
    "                A_fem,\n",
    "                b_fem,\n",
    "                policy=policy,\n",
    "                maxiter=maxiter,\n",
    "            ),\n",
    "            A_fem,\n",
    "        ),\n",
    "        feature_fn=phis,\n",
    "    )\n",
    "\n",
    "    ax.cla()\n",
    "    u_fem.plot(\n",
    "        ax,\n",
    "        xs_plot,\n",
    "        rng=np.random.default_rng(2345),\n",
    "        num_samples=10,\n",
    "        label=r\"$\\phi(\\cdot)^{\\top} \\mathrm{w}$\",\n",
    "    )\n",
    "\n",
    "    ax.plot(\n",
    "        xs_plot,\n",
    "        u_star(xs_plot),\n",
    "        label=r\"$u^\\star$\",\n",
    "    )\n",
    "\n",
    "    plt.legend()\n",
    "\n",
    "    ax.legend()\n",
    "\n",
    "ipywidgets.interactive(\n",
    "    interact,\n",
    "    policy=ipywidgets.ToggleButtons(\n",
    "        options=['CG', 'Coordinate', 'Random', 'Eigenvector (desc)', 'Eigenvector (asc)'],\n",
    "        description='Policy',\n",
    "    ),\n",
    "    maxiter=ipywidgets.IntSlider(\n",
    "        value=int(0.5 * n),\n",
    "        min=0,\n",
    "        max=n,\n",
    "        description=\"maxiter\",\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(np.linalg.eigvalsh(A_fem.todense()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}