Add diffrax explainer

examples/diffrax_comparison.ipynb

@@ -21,7 +21,14 @@
     "$$\n",
     "dx = -A(x-b)dt + \\mathcal{N}(0, Ddt)\n",
     "$$\n",
-    "and compare `thermox` and `diffrax` for a varying number of timesteps."
+    "and compare `thermox` and `diffrax` for a varying number of timesteps.\n",
+    "\n",
+    "\n",
+    "`thermox` simulates OU processes *exactly*, whilst [`diffrax`](https://github.com/patrick-kidger/diffrax)\n",
+    "is a differential equation library that uses discretization to\n",
+    "*approximately* simulate general ODEs and SDEs (thus `diffrax` has a broader scope but\n",
+    "still represents the most natural open source comparison for simulating OU processes\n",
+    "with JAX)."
    ]
   },
   {
@@ -123,7 +130,6 @@
     "    brownian_motion = diffrax.VirtualBrownianTree(ts[0], ts[-1], tol=1e-3, shape=(dim,), key=rk)\n",
     "    terms = diffrax.MultiTerm(diffrax.ODETerm(drift), diffrax.ControlTerm(diffusion, brownian_motion))\n",
     "    solver = diffrax.ItoMilstein()\n",
-    "    # saveat = diffrax.SaveAt(dense=True)\n",
     "    saveat = diffrax.SaveAt(ts=ts)\n",
     "    sol = diffrax.diffeqsolve(terms, solver, ts[0], ts[-1], y0=x0, saveat=saveat,\n",
     "                               dt0=0.01, max_steps=10000000\n",
@@ -216,6 +222,13 @@
     "    "
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let's visualise the runtimes"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 43,
@@ -259,7 +272,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "And there you have it! `thermox` is roughly constant with respect to the number of time points, because the computation is dominated by the diagonalisation that is performed at initialisation. For a small number of time points, performance is comparable, but `thermox` really starts to shine for $\\sim 10000$ time steps, where we get $\\sim 200\\times$ speedup with respect to `diffrax`, as well as being exact!"
+    "And there you have it! `thermox` is roughly constant with respect to the number of time points, because the computation is dominated by the diagonalisation that is performed at initialisation. For a small number of time points, performance is comparable, but `thermox` really starts to shine for $\\sim 10000$ time steps, where we get $\\sim 800\\times$ speedup with respect to `diffrax`, as well as being exact!"
    ]
   },
   {