{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-odes-solvers-md-1 fw=all -->\n",
    "\n",
    "# Ordinary Differential Equations and Numerical Solvers\n",
    "\n",
    "Every continuous-time generative model, from Neural ODEs and continuous\n",
    "normalizing flows to the deterministic samplers of diffusion and flow\n",
    "matching, is an **ordinary differential equation that you integrate**. An ODE\n",
    "is a *velocity rule*: at every point of space (and time) it tells you which\n",
    "way to move and how fast. A solution is the path you trace by always\n",
    "following the local arrow; a numerical solver is a recipe for following\n",
    "arrows in finite steps; a deep residual network *is* such a recipe, with its\n",
    "layers as the steps; and a density transported by the flow changes at a rate\n",
    "read off the field's Jacobian. This section builds the mathematics behind\n",
    "that picture: what a vector field is, when a trajectory exists and is unique,\n",
    "how eigenvalues decide stability, and how numerical solvers trade step size\n",
    "for accuracy. The stochastic counterpart of the Euler method\n",
    "arrives in that section, the probability-flow ODE that\n",
    "turns a diffusion model into a deterministic flow in\n",
    "that section, and the solvers studied here\n",
    "set the step-count/quality tradeoff of every sampler in\n",
    "that section.\n",
    "\n",
    "We lean on that section (eigenvalues and\n",
    "eigenvectors; that section already defined the matrix exponential, and here\n",
    "we prove its properties), that section (Jacobians,\n",
    "Taylor expansion),\n",
    "that section (reverse-mode AD as vector--Jacobian\n",
    "products), that section (the integrals being\n",
    "approximated), and that section (the change-of-variables\n",
    "formula for densities). The numerical demonstrations are plain NumPy, because\n",
    "every solver is a handful of lines, until we *train* a Neural ODE below."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "odes-solvers-imports"
   },
   "source": [
    "%matplotlib inline\n",
    "from d2l import torch as d2l\n",
    "import numpy as np\n",
    "import torch"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-odes-solvers-md-2 fw=all -->\n",
    "\n",
    "## Vector Fields, Trajectories, and Well-Posedness\n",
    "\n",
    "### Velocity Fields and Integral Curves\n",
    "\n",
    "An **initial-value problem** specifies a velocity rule and a starting point:\n",
    "\n",
    "$$\n",
    "\\dot{\\mathbf{x}}(t) = \\mathbf{f}(\\mathbf{x}(t), t),\n",
    "\\qquad \\mathbf{x}(0) = \\mathbf{x}_0,\n",
    "$$\n",
    "\n",
    "where $\\mathbf{x}(t) \\in \\mathbb{R}^d$ is the state, the dot is the time\n",
    "derivative, and $\\mathbf{f} : \\mathbb{R}^d \\times \\mathbb{R} \\to \\mathbb{R}^d$\n",
    "is a **vector field**: a function that attaches an arrow to every point. A\n",
    "**trajectory** (or *integral curve*) is a differentiable path whose tangent\n",
    "matches the field everywhere: at each instant, the particle moves with\n",
    "exactly the velocity the field prescribes at its current location. When\n",
    "$\\mathbf{f}$ does not depend on $t$ we call the system **autonomous**; the\n",
    "arrows are frozen and only the particle moves. the figure\n",
    "shows the picture to keep in mind: a grid of arrows, and curves threading\n",
    "through them, everywhere tangent.\n",
    "\n",
    "![An ODE as a velocity field. The faint grid arrows are the field $\\mathbf{f}(\\mathbf{x})=A\\mathbf{x}$ with $A=\\left(\\begin{smallmatrix}-0.5&-1\\\\1&-0.5\\end{smallmatrix}\\right)$; the two coloured integral curves are trajectories that always follow the local arrow, spiralling into the stable fixed point at the origin (eigenvalues $-0.5\\pm i$).](../img/mdl-dyn-ode-field.svg)\n",
    "\n",
    "Integrating both sides of the equation from $0$ to $t$\n",
    "(that section) gives the equivalent **integral form**\n",
    "\n",
    "$$\n",
    "\\mathbf{x}(t) = \\mathbf{x}_0 + \\int_0^t \\mathbf{f}(\\mathbf{x}(s), s)\\,ds,\n",
    "$$\n",
    "\n",
    "which trades the derivative for an accumulation of velocities. This form does\n",
    "real work twice over: it is the fixed-point equation behind the existence\n",
    "theorem below, and it is what every numerical solver approximates; Euler's\n",
    "method is the left-endpoint Riemann sum of this very integral.\n",
    "\n",
    "Two scalar warm-ups anchor the notation. For $\\dot{x} = -x$, the rule \"move\n",
    "toward zero at a speed proportional to your distance\" is solved by\n",
    "$x(t) = e^{-t} x_0$: exponential decay, the prototype of a *contracting* mode.\n",
    "For the rotational field\n",
    "$\\dot{\\mathbf{x}} = \\left(\\begin{smallmatrix}0&-1\\\\1&0\\end{smallmatrix}\\right)\\mathbf{x}$,\n",
    "the velocity is always perpendicular to the position, so\n",
    "$\\tfrac{d}{dt}\\|\\mathbf{x}\\|^2 = 2\\,\\mathbf{x}^\\top\\dot{\\mathbf{x}} = 0$ and\n",
    "trajectories are circles traversed at constant angular speed: pure rotation,\n",
    "the prototype of an *oscillating* mode. The field of\n",
    "the figure is exactly the sum of these two behaviors, and\n",
    "we will see shortly that its eigenvalues $-0.5 \\pm i$ say so at a glance.\n",
    "\n",
    "### The Flow Map\n",
    "\n",
    "Fix the field and vary the starting point: for an autonomous system, write\n",
    "$\\Phi_t(\\mathbf{x}_0) = \\mathbf{x}(t)$ for the solution at time $t$ started at\n",
    "$\\mathbf{x}_0$. This **flow map** $\\Phi_t : \\mathbb{R}^d \\to \\mathbb{R}^d$\n",
    "moves *every* point of space forward $t$ seconds along its arrows, and it\n",
    "composes the way time does:\n",
    "\n",
    "$$\n",
    "\\Phi_0 = \\mathrm{id},\n",
    "\\qquad\n",
    "\\Phi_{t+s} = \\Phi_t \\circ \\Phi_s .\n",
    "$$\n",
    "\n",
    "Flowing for $s$ seconds and then $t$ more is the same as flowing for $t+s$,\n",
    "whenever all three solutions exist. This holds because the field is frozen in\n",
    "time: the second leg starts from $\\Phi_s(\\mathbf{x}_0)$ and follows the *same*\n",
    "rule. If solutions exist uniquely both forward and backward over the relevant\n",
    "interval, taking $s=-t$ gives\n",
    "$\\Phi_t \\circ \\Phi_{-t}=\\mathrm{id}$, so $\\Phi_t$ is a bijection with inverse\n",
    "$\\Phi_{-t}$. A field whose solutions exist for every positive and negative\n",
    "time therefore generates a one-parameter group of bijections. This is the\n",
    "seed of an invertible generative flow: push noise forward through a learned\n",
    "field to get data, and integrate backward to recover the noise (and, as we\n",
    "will see in that section, its likelihood).\n",
    "The qualifiers about existence and uniqueness matter, which is where we turn\n",
    "next.\n",
    "\n",
    "### Existence and Uniqueness\n",
    "\n",
    "Before building on this, we need to know that exactly one trajectory passes\n",
    "through each starting point. The right hypothesis is that the field does not\n",
    "change too abruptly from point to point.\n",
    "\n",
    "**Proposition (Picard--Lindelöf).** *Let\n",
    "$\\mathbf{f}(\\mathbf{x}, t)$ be continuous in $t$ and Lipschitz in\n",
    "$\\mathbf{x}$: there is an $L$ with*\n",
    "\n",
    "$$\n",
    "\\|\\mathbf{f}(\\mathbf{x}, t) - \\mathbf{f}(\\mathbf{y}, t)\\| \\le L\\,\\|\\mathbf{x} - \\mathbf{y}\\|\n",
    "\\qquad \\textrm{for all } \\mathbf{x}, \\mathbf{y}, t.\n",
    "$$\n",
    "\n",
    "*Then the initial-value problem the equation has exactly one\n",
    "solution on $[0, T]$, for every $\\mathbf{x}_0$ and every $T$*\n",
    "[@Coddington.Levinson.1955].\n",
    "\n",
    "**Proof sketch (the integral operator is a contraction).** A path\n",
    "$\\boldsymbol{\\varphi}$ solves the equation iff it is a fixed point\n",
    "of the **Picard operator** built from the integral form\n",
    "the equation,\n",
    "\n",
    "$$\n",
    "(P\\boldsymbol{\\varphi})(t) = \\mathbf{x}_0 + \\int_0^t \\mathbf{f}(\\boldsymbol{\\varphi}(s), s)\\,ds .\n",
    "$$\n",
    "\n",
    "For two candidate paths on a short interval $[0, \\delta]$,\n",
    "\n",
    "$$\n",
    "\\|(P\\boldsymbol{\\varphi})(t) - (P\\boldsymbol{\\psi})(t)\\|\n",
    "\\le \\int_0^t \\|\\mathbf{f}(\\boldsymbol{\\varphi}(s), s) - \\mathbf{f}(\\boldsymbol{\\psi}(s), s)\\|\\,ds\n",
    "\\le L\\,\\delta \\cdot \\max_{s \\le \\delta}\\|\\boldsymbol{\\varphi}(s) - \\boldsymbol{\\psi}(s)\\| ,\n",
    "$$\n",
    "\n",
    "so once $\\delta < 1/L$ the operator $P$ shrinks the distance between any two\n",
    "paths by a fixed factor: it is a *contraction*, and by the Banach fixed-point\n",
    "theorem a contraction on a complete space has exactly one fixed point\n",
    "(iterate $P$ from anywhere; the iterates form a Cauchy sequence, and two\n",
    "distinct fixed points would have to be strictly closer to each other than\n",
    "they are). Completeness is what turns the Cauchy sequence into a solution:\n",
    "the space $C([0, \\delta])$ of continuous paths under the sup norm is\n",
    "complete, so the iterates converge uniformly to a continuous fixed point.\n",
    "That settles $[0, \\delta]$; restart at $t = \\delta$ and repeat to cover all\n",
    "of $[0, T]$.\n",
    "\n",
    "Iterating $P$ from the constant path is **Picard iteration**; here it is once\n",
    "by hand. For $\\dot{x} = x$, $x_0 = 1$: starting from\n",
    "$\\varphi_0 \\equiv 1$,\n",
    "\n",
    "$$\n",
    "\\varphi_1(t) = 1 + t,\n",
    "\\qquad\n",
    "\\varphi_2(t) = 1 + t + \\tfrac{t^2}{2},\n",
    "\\qquad\n",
    "\\varphi_3(t) = 1 + t + \\tfrac{t^2}{2} + \\tfrac{t^3}{6},\n",
    "\\;\\;\\ldots\n",
    "$$\n",
    "\n",
    "These are the partial sums of $e^t$: the contraction *constructs* the\n",
    "exponential, one Taylor term per sweep through the integral.\n",
    "\n",
    "These hypotheses are a convenient sufficient package, not individually\n",
    "necessary conditions. Two standard counterexamples show the failures that\n",
    "they are designed to rule out. **Uniqueness can fail without Lipschitz.** The\n",
    "field $\\dot{x} = \\sqrt{|x|}$ has slope $\\to \\infty$ near $x = 0$ (no finite\n",
    "$L$ works there), and through $x_0 = 0$ it threads *infinitely many*\n",
    "solutions: $x(t) \\equiv 0$ is one, and for every waiting time $c \\ge 0$,\n",
    "\n",
    "$$\n",
    "x_c(t) =\n",
    "\\begin{cases}\n",
    "0, & t \\le c, \\\\\n",
    "(t - c)^2/4, & t > c\n",
    "\\end{cases}\n",
    "$$\n",
    "\n",
    "is another: sit at the origin for $c$ seconds, then leak away (check:\n",
    "$\\dot{x}_c = (t-c)/2 = \\sqrt{x_c}$ for $t > c$). The solutions form a *fan*\n",
    "leaving the origin, one blade per departure time, and the initial-value\n",
    "problem is genuinely ambiguous (the figure).\n",
    "\n",
    "![Solutions of $\\dot{x} = \\sqrt{|x|}$ through $x(0) = 0$: the rest solution $x \\equiv 0$ plus one parabolic blade $x_c(t) = (t-c)^2/4$ for each departure time $c$. The slope of $\\sqrt{|x|}$ is unbounded at the origin, so the Lipschitz hypothesis fails there, and with it uniqueness: every blade is a legitimate solution.](../img/mdl-dyn-uniqueness-fan.svg)\n",
    "\n",
    "**Global existence fails without a growth bound.** The\n",
    "field $\\dot{x} = x^2$ is Lipschitz on every bounded set but not globally, and\n",
    "from $x_0 > 0$, separation of variables gives\n",
    "\n",
    "$$\n",
    "x(t) = \\frac{x_0}{1 - x_0 t},\n",
    "$$\n",
    "\n",
    "which escapes to infinity at the finite time $t = 1/x_0$: the faster you grow,\n",
    "the faster you grow, and superlinear feedback compounds to a singularity. The\n",
    "solution exists only *locally*, on $[0, 1/x_0)$, which is all that\n",
    "local Lipschitz continuity can promise.\n",
    "\n",
    "For deep learning the theorem is a license. A neural field\n",
    "$\\mathbf{f}_\\theta$ built from linear maps and Lipschitz activations\n",
    "($\\tanh$, ReLU, GELU) is *globally* Lipschitz (a composition of globally\n",
    "Lipschitz maps is globally Lipschitz, with constant at most the product of\n",
    "the weight norms and the activations' Lipschitz constants), which is\n",
    "exactly the global-$L$ hypothesis of the proposition above. So a Neural ODE\n",
    "(that section) is\n",
    "well-posed: through every input there is exactly one trajectory, distinct\n",
    "inputs can never collide (two trajectories meeting at time $t$ would be two\n",
    "solutions of the same final-value problem run backward), and the learned map\n",
    "$\\Phi_T$ is automatically a bijection. Invertibility comes free, as a\n",
    "theorem.\n",
    "\n",
    "## Linear ODEs and Stability\n",
    "\n",
    "### The Matrix Exponential\n",
    "\n",
    "One family of ODEs can be solved in closed form, and it is the family that\n",
    "explains all the others locally: the linear system\n",
    "\n",
    "$$\n",
    "\\dot{\\mathbf{x}} = A\\mathbf{x},\n",
    "\\qquad\n",
    "\\mathbf{x}(0) = \\mathbf{x}_0,\n",
    "\\qquad A \\in \\mathbb{R}^{d \\times d}.\n",
    "$$\n",
    "\n",
    "In one dimension, $\\dot{x} = ax$ is solved by $e^{at} x_0$, so we engineer a\n",
    "matrix version of the exponential to make the same formula true. Define, for\n",
    "any square matrix $B$, the **matrix exponential** by the same power series\n",
    "that defines $e^z$:\n",
    "\n",
    "$$\n",
    "e^{B} = \\sum_{k=0}^{\\infty} \\frac{B^k}{k!}\n",
    "= I + B + \\frac{B^2}{2} + \\frac{B^3}{6} + \\cdots .\n",
    "$$\n",
    "\n",
    "The series converges for *every* matrix: term by term,\n",
    "$\\|B^k\\| \\le \\|B\\|^k$ (the operator norm is submultiplicative), so the series\n",
    "is dominated by the scalar series for $e^{\\|B\\|}$, which converges absolutely.\n",
    "\n",
    "**Proposition (the matrix exponential solves linear ODEs).**\n",
    "*$\\mathbf{x}(t) = e^{At}\\mathbf{x}_0$ is the unique solution of\n",
    "$\\dot{\\mathbf{x}} = A\\mathbf{x}$, $\\mathbf{x}(0) = \\mathbf{x}_0$.*\n",
    "\n",
    "**Proof.** Differentiate the series the equation with\n",
    "$B = At$ term by term (legitimate: on any bounded $t$-interval the\n",
    "differentiated series still converges uniformly, by the same domination):\n",
    "\n",
    "$$\n",
    "\\frac{d}{dt} e^{At}\n",
    "= \\sum_{k=1}^{\\infty} \\frac{A^k\\, t^{k-1}}{(k-1)!}\n",
    "= A \\sum_{j=0}^{\\infty} \\frac{(At)^j}{j!}\n",
    "= A\\, e^{At}.\n",
    "$$\n",
    "\n",
    "So $\\mathbf{x}(t) = e^{At}\\mathbf{x}_0$ satisfies the ODE, and at $t = 0$ the\n",
    "series collapses to $I$, so $\\mathbf{x}(0) = \\mathbf{x}_0$. Uniqueness is\n",
    "Picard--Lindelöf: the field $\\mathbf{x} \\mapsto A\\mathbf{x}$ is Lipschitz with\n",
    "$L = \\|A\\|$. $\\blacksquare$\n",
    "\n",
    "An infinite series of matrix powers sounds unwieldy, but the\n",
    "eigendecomposition of that section collapses it to\n",
    "$d$ scalar exponentials.\n",
    "\n",
    "**Proposition (matrix exponential by eigendecomposition).** *If\n",
    "$A = V \\Lambda V^{-1}$ with $\\Lambda = \\mathrm{diag}(\\lambda_1, \\ldots, \\lambda_d)$,\n",
    "then*\n",
    "\n",
    "$$\n",
    "e^{At} = V e^{\\Lambda t} V^{-1},\n",
    "\\qquad\n",
    "e^{\\Lambda t} = \\mathrm{diag}\\!\\left(e^{\\lambda_1 t}, \\ldots, e^{\\lambda_d t}\\right).\n",
    "$$\n",
    "\n",
    "**Proof.** Powers telescope: $(At)^k = V (\\Lambda t)^k V^{-1}$, because each\n",
    "adjacent $V^{-1} V$ cancels. Summing the series\n",
    "the equation and pulling $V, V^{-1}$ out of the sum\n",
    "gives $e^{At} = V\\left(\\sum_k (\\Lambda t)^k / k!\\right) V^{-1}$. The inner sum\n",
    "is the exponential of a diagonal matrix, which is diagonal with entries\n",
    "$\\sum_k (\\lambda_i t)^k / k! = e^{\\lambda_i t}$. $\\blacksquare$\n",
    "\n",
    "The formula is a *decoupling* statement. Expand the start point in the\n",
    "eigenbasis, $\\mathbf{x}_0 = \\sum_i c_i \\mathbf{v}_i$; then\n",
    "\n",
    "$$\n",
    "\\mathbf{x}(t) = \\sum_{i=1}^{d} c_i\\, e^{\\lambda_i t}\\, \\mathbf{v}_i :\n",
    "$$\n",
    "\n",
    "along each eigendirection the system is the scalar equation\n",
    "$\\dot{x} = \\lambda_i x$ we already solved, evolving independently of the\n",
    "others. Eigenvectors are the directions in which a linear ODE is\n",
    "one-dimensional; the $e^{\\lambda_i t}$ are its **modes**.\n",
    "\n",
    "### The Stability Dictionary\n",
    "\n",
    "Each mode $e^{\\lambda_i t}$ with $\\lambda_i = a_i + i b_i$ has magnitude\n",
    "$|e^{\\lambda_i t}| = e^{a_i t}$: *the real part sets growth or decay, the\n",
    "imaginary part sets rotation*. Reading the equation\n",
    "mode by mode gives the dictionary every dynamical argument in this chapter\n",
    "uses:\n",
    "\n",
    "* $\\operatorname{Re}\\lambda_i < 0$ for all $i$: every mode decays, and\n",
    "  $\\mathbf{x}(t) \\to \\mathbf{0}$ from every start: the origin is\n",
    "  **asymptotically stable**, with decay rate set by the slowest mode,\n",
    "  $\\max_i \\operatorname{Re}\\lambda_i$.\n",
    "* $\\operatorname{Re}\\lambda_i > 0$ for some $i$: that mode grows\n",
    "  exponentially, and almost every trajectory escapes: **unstable**.\n",
    "* $\\operatorname{Im}\\lambda_i \\neq 0$: the conjugate pair contributes\n",
    "  $e^{a t}(\\cos bt, \\sin bt)$ terms: **rotation**, decaying or growing\n",
    "  with the sign of $a$.\n",
    "\n",
    "The rotation entry is the eigenvalue--rotation dictionary of\n",
    "that section run in continuous time. On the invariant\n",
    "plane of a conjugate pair $a \\pm ib$, the matrix $A$ acts as the block\n",
    "$\\left(\\begin{smallmatrix}a & -b\\\\ b & \\phantom{-}a\\end{smallmatrix}\\right)$\n",
    "of the equation. Write it as $aI + bJ$ with\n",
    "$J = \\left(\\begin{smallmatrix}0 & -1\\\\ 1 & \\phantom{-}0\\end{smallmatrix}\\right)$;\n",
    "since $J^2 = -I$, the exponential series splits into cosine and sine parts,\n",
    "$e^{(aI + bJ)t} = e^{at}(\\cos bt \\, I + \\sin bt \\, J)$, which is scaling by\n",
    "$e^{at}$ composed with rotation by $bt$. Where the *iterated* matrix of\n",
    "that section turned by a fixed angle $\\theta$ per\n",
    "step, the *flow* turns at angular velocity $b$: this is the spiral, and it is\n",
    "exactly where the $e^{at}(\\cos bt, \\sin bt)$ terms above come from.\n",
    "\n",
    "In two dimensions the dictionary is a portrait gallery\n",
    "(the figure). The saddle deserves a\n",
    "note: with real eigenvalues of opposite signs, trajectories approach along the\n",
    "stable eigendirection and escape along the unstable one, the generic fate\n",
    "of an \"unstable equilibrium.\" The stable spiral is the field of\n",
    "the figure ($\\lambda = -0.5 \\pm i$: rotate while\n",
    "decaying); the center is our rotation warm-up.\n",
    "\n",
    "| eigenvalues of $A$ ($2\\times 2$)              | portrait        | behavior                    |\n",
    "|-----------------------------------------------|-----------------|-----------------------------|\n",
    "| real, both negative                           | stable node     | monotone decay              |\n",
    "| real, both positive                           | unstable node   | monotone growth             |\n",
    "| real, opposite signs                          | saddle          | decay in, escape out        |\n",
    "| complex pair, $\\operatorname{Re}\\lambda < 0$  | stable spiral   | decaying oscillation        |\n",
    "| complex pair, $\\operatorname{Re}\\lambda > 0$  | unstable spiral | growing oscillation         |\n",
    "| purely imaginary pair                         | center          | closed orbits               |\n",
    "\n",
    "![The portrait gallery of planar linear systems. Top row: stable node (real eigenvalues, both negative), unstable node (both positive), saddle (opposite signs). Bottom row: stable spiral (complex pair, $a < 0$), unstable spiral ($a > 0$), center (purely imaginary). Arrows show the field, curves a few trajectories; the real part decides growth or decay, the imaginary part rotation.](../img/mdl-dyn-phase-portraits.svg)\n",
    "\n",
    "(When $A$ is defective, i.e. not diagonalizable\n",
    "(that section), the modes pick up polynomial factors\n",
    "$t^k e^{\\lambda t}$, but polynomials never beat exponentials, so the\n",
    "stability verdict still depends only on the signs of\n",
    "$\\operatorname{Re}\\lambda_i$.)\n",
    "\n",
    "### Linearization at Fixed Points\n",
    "\n",
    "Nonlinear fields inherit all of this *locally*. A **fixed point** of an\n",
    "autonomous system is a state $\\mathbf{x}^\\star$ with\n",
    "$\\mathbf{f}(\\mathbf{x}^\\star) = \\mathbf{0}$: the particle parked there\n",
    "never moves. For a small displacement\n",
    "$\\boldsymbol{\\delta}(t) = \\mathbf{x}(t) - \\mathbf{x}^\\star$, Taylor expansion\n",
    "(that section) gives\n",
    "\n",
    "$$\n",
    "\\dot{\\boldsymbol{\\delta}}\n",
    "= \\mathbf{f}(\\mathbf{x}^\\star + \\boldsymbol{\\delta})\n",
    "= J\\,\\boldsymbol{\\delta} + O(\\|\\boldsymbol{\\delta}\\|^2),\n",
    "\\qquad\n",
    "J = \\frac{\\partial \\mathbf{f}}{\\partial \\mathbf{x}}(\\mathbf{x}^\\star),\n",
    "$$\n",
    "\n",
    "so near the fixed point the dynamics are the linear system\n",
    "$\\dot{\\boldsymbol{\\delta}} = J \\boldsymbol{\\delta}$, and the eigenvalues of\n",
    "the *Jacobian at the fixed point* decide local stability by the same\n",
    "dictionary. (This is rigorous whenever no eigenvalue sits exactly on the\n",
    "imaginary axis, by the Hartman--Grobman theorem\n",
    "[@Hartman.1960; @Grobman.1959]; on the axis, the neglected\n",
    "quadratic terms get a vote.) The damped pendulum\n",
    "$\\ddot{\\theta} = -\\sin\\theta - \\gamma\\dot{\\theta}$, written as a first-order\n",
    "system in $(\\theta, \\dot{\\theta})$, has Jacobian eigenvalues with negative\n",
    "real part at the hanging rest point $(0, 0)$ (a stable spiral, the\n",
    "ring-down of a released pendulum) and a saddle at the inverted balance\n",
    "$(\\pi, 0)$, which is why you can stand a pencil on its tip only along a\n",
    "measure-zero set of initial conditions. The same computation, applied to a\n",
    "trained network's dynamics or to the mean of the Ornstein--Uhlenbeck process\n",
    "(that section, whose deterministic skeleton is exactly\n",
    "the contracting mode $\\dot{x} = -\\theta x$), is the working stability tool of\n",
    "this whole chapter.\n",
    "\n",
    "Let us compute all of this. The cell builds $e^{At}$ for the spiral field of\n",
    "the figure three independent ways (the\n",
    "eigendecomposition formula the equation, the truncated\n",
    "series the equation, and the compounding limit\n",
    "$(I + tA/n)^n$, forward Euler in disguise and a preview of the next\n",
    "section) and checks the stability dictionary's prediction that\n",
    "$\\|\\mathbf{x}(t)\\|$ decays exactly like $e^{-t/2}$ (the rotation part of this\n",
    "particular $A$ is norm-preserving)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "odes-solvers-matrix-exponential"
   },
   "source": [
    "A = np.array([[-0.5, -1.0], [1.0, -0.5]])      # spiral sink: eigenvalues -0.5 +- i\n",
    "t = 2.0\n",
    "lam, V = np.linalg.eig(A)                      # A = V diag(lam) V^{-1}\n",
    "expAt_eig = (V @ np.diag(np.exp(lam * t)) @ np.linalg.inv(V)).real\n",
    "term, expAt_series = np.eye(2), np.eye(2)      # partial sums of sum_k (At)^k / k!\n",
    "for k in range(1, 30):\n",
    "    term = term @ (A * t) / k\n",
    "    expAt_series += term\n",
    "n = 2 ** 24                                    # Euler compounding (I + tA/n)^n\n",
    "expAt_euler = np.linalg.matrix_power(np.eye(2) + (t / n) * A, n)\n",
    "x0 = np.array([2.0, 0.0])\n",
    "print('eigenvalues of A:', lam.round(4))\n",
    "print('|series - eigen formula| :', f'{np.abs(expAt_series - expAt_eig).max():.1e}')\n",
    "print('|(I+tA/n)^n - eigen|     :', f'{np.abs(expAt_euler - expAt_eig).max():.1e}')\n",
    "print('||e^{At} x0|| =', f'{np.linalg.norm(expAt_eig @ x0):.6f}',\n",
    "      ' vs  e^{-t/2}||x0|| =', f'{np.exp(-0.5 * t) * np.linalg.norm(x0):.6f}')"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "eigenvalues of A: [-0.5+1.j -0.5-1.j]\n|series - eigen formula| : 5.6e-17\n|(I+tA/n)^n - eigen|     : 5.4e-08\n||e^{At} x0|| = 0.735759  vs  e^{-t/2}||x0|| = 0.735759\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-odes-solvers-md-3 fw=all -->\n",
    "\n",
    "Twenty-nine series terms already agree with the eigendecomposition formula to\n",
    "machine precision, the compounded Euler product lands within $10^{-7}$, and\n",
    "the trajectory norm matches $e^{-t/2}\\|\\mathbf{x}_0\\|$ to six digits: the\n",
    "spectrum ($-0.5 \\pm i$) told us the decay rate before we integrated anything.\n",
    "\n",
    "## Numerical Solvers: From Euler to Runge--Kutta\n",
    "\n",
    "### Forward Euler and Its Global Error\n",
    "\n",
    "Almost no ODE beyond the linear family has a closed-form solution, so we\n",
    "*march*: replace the integral form the equation over one\n",
    "short step by its left-endpoint approximation. With step size $h$ and\n",
    "$t_n = nh$, **forward Euler** is\n",
    "\n",
    "$$\n",
    "\\mathbf{x}_{n+1} = \\mathbf{x}_n + h\\, \\mathbf{f}(\\mathbf{x}_n, t_n)\n",
    "$$\n",
    "\n",
    "Take the arrow under your feet, follow it for time $h$, look again.\n",
    "Equivalently, the update is the Taylor expansion of the true solution truncated after\n",
    "the linear term, and the truncation error of one step is the first term\n",
    "dropped: if $\\|\\ddot{\\mathbf{x}}(t)\\| \\le M$ along the solution, then\n",
    "\n",
    "$$\n",
    "\\mathbf{x}(t_{n+1}) = \\mathbf{x}(t_n) + h\\,\\mathbf{f}(\\mathbf{x}(t_n), t_n) + \\mathbf{r}_n,\n",
    "\\qquad\n",
    "\\|\\mathbf{r}_n\\| \\le \\tfrac{M h^2}{2}.\n",
    "$$\n",
    "\n",
    "A per-step error of $O(h^2)$ does not mean a final error of $O(h^2)$: to\n",
    "reach a fixed horizon $T$ you take $n = T/h$ steps, and the per-step errors\n",
    "*accumulate*; worse, each step also inherits and possibly amplifies the\n",
    "error already made. One order of $h$ is lost to the accumulation, and the\n",
    "bookkeeping is short enough to do in full.\n",
    "\n",
    "**Proposition (Euler converges at order 1).** *Let $\\mathbf{f}$ be\n",
    "$L$-Lipschitz in $\\mathbf{x}$, and let the solution of\n",
    "the equation on $[0, T]$ satisfy\n",
    "$\\|\\ddot{\\mathbf{x}}(t)\\| \\le M$. Then the Euler iterates\n",
    "the equation obey*\n",
    "\n",
    "$$\n",
    "\\max_{n \\le T/h} \\|\\mathbf{x}_n - \\mathbf{x}(t_n)\\|\n",
    "\\;\\le\\; \\frac{M h}{2L}\\left(e^{LT} - 1\\right)\n",
    "\\;=\\; O(h).\n",
    "$$\n",
    "\n",
    "**Proof.** Let $e_n = \\|\\mathbf{x}_n - \\mathbf{x}(t_n)\\|$, and subtract the\n",
    "Taylor identity above from the update the equation:\n",
    "\n",
    "$$\n",
    "e_{n+1}\n",
    "\\le \\underbrace{\\|\\mathbf{x}_n - \\mathbf{x}(t_n)\\| + h\\,\\|\\mathbf{f}(\\mathbf{x}_n, t_n) - \\mathbf{f}(\\mathbf{x}(t_n), t_n)\\|}_{\\le\\, (1 + hL)\\, e_n \\textrm{ by Lipschitz}}\n",
    "\\;+\\; \\underbrace{\\|\\mathbf{r}_n\\|}_{\\le\\, Mh^2/2} .\n",
    "$$\n",
    "\n",
    "So each step multiplies the existing error by at most $(1 + hL)$ and adds at\n",
    "most $Mh^2/2$. Unrolling the recursion from $e_0 = 0$ sums a geometric\n",
    "series:\n",
    "\n",
    "$$\n",
    "e_n \\le \\frac{Mh^2}{2} \\sum_{j=0}^{n-1} (1 + hL)^j\n",
    "= \\frac{Mh^2}{2} \\cdot \\frac{(1 + hL)^n - 1}{hL}\n",
    "\\le \\frac{Mh}{2L}\\left(e^{LT} - 1\\right),\n",
    "$$\n",
    "\n",
    "using $(1 + hL)^n \\le e^{nhL} \\le e^{LT}$. $\\blacksquare$\n",
    "\n",
    "The structure of the bound recurs for\n",
    "every solver: a *local* error of order $h^{p+1}$ per step becomes a *global*\n",
    "error of order $h^p$ after $T/h$ steps, inflated by a stability factor\n",
    "$e^{LT}$ that prices how strongly the dynamics can amplify old mistakes. We\n",
    "say Euler is a method of **order** $p = 1$: halve the step, halve the error.\n",
    "\n",
    "### Runge--Kutta Methods\n",
    "\n",
    "Euler's weakness is that it commits to the slope at the *start* of the step,\n",
    "while the trajectory is already curving away. The fix is to spend a few extra\n",
    "field evaluations probing the slope *inside* the step and average them. The\n",
    "simplest upgrade, the **midpoint method**, takes half an Euler step only to\n",
    "*measure* the slope there, then takes the real step with that better slope:\n",
    "$\\mathbf{x}_{n+1} = \\mathbf{x}_n + h\\,\\mathbf{f}\\!\\left(\\mathbf{x}_n + \\tfrac{h}{2}\\mathbf{f}(\\mathbf{x}_n)\\right)$\n",
    "(autonomous notation). One Taylor line shows why it helps: in one dimension,\n",
    "\n",
    "$$\n",
    "x + h f\\!\\left(x + \\tfrac{h}{2} f(x)\\right)\n",
    "= x + h f + \\tfrac{h^2}{2} f' f + O(h^3),\n",
    "$$\n",
    "\n",
    "which matches the true expansion\n",
    "$x(t+h) = x + h\\dot{x} + \\tfrac{h^2}{2}\\ddot{x} + O(h^3)$ because\n",
    "$\\ddot{x} = \\tfrac{d}{dt} f(x(t)) = f' f$. The $h^2$ term is now *correct*\n",
    "instead of *absent*: local error $O(h^3)$, global order $2$, at the price of\n",
    "two field evaluations per step. Its trapezoid twin, **Heun's method**, takes\n",
    "a full trial Euler step and averages the slopes at its two ends,\n",
    "$\\mathbf{x}_{n+1} = \\mathbf{x}_n + \\tfrac{h}{2}\\left[\\mathbf{f}(\\mathbf{x}_n) + \\mathbf{f}\\!\\left(\\mathbf{x}_n + h\\,\\mathbf{f}(\\mathbf{x}_n)\\right)\\right]$\n",
    "(same order $2$, same two evaluations), and is the two-stage sampler of\n",
    "choice for the learned dynamics of\n",
    "that section.\n",
    "\n",
    "Pushing the same idea to four probe slopes gives the standard general-purpose\n",
    "method of scientific computing, **classical Runge--Kutta (RK4)**:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\mathbf{k}_1 &= \\mathbf{f}(\\mathbf{x}_n,\\, t_n), \\\\\n",
    "\\mathbf{k}_2 &= \\mathbf{f}(\\mathbf{x}_n + \\tfrac{h}{2}\\mathbf{k}_1,\\, t_n + \\tfrac{h}{2}), \\\\\n",
    "\\mathbf{k}_3 &= \\mathbf{f}(\\mathbf{x}_n + \\tfrac{h}{2}\\mathbf{k}_2,\\, t_n + \\tfrac{h}{2}), \\\\\n",
    "\\mathbf{k}_4 &= \\mathbf{f}(\\mathbf{x}_n + h\\,\\mathbf{k}_3,\\, t_n + h), \\\\\n",
    "\\mathbf{x}_{n+1} &= \\mathbf{x}_n + \\tfrac{h}{6}\\left(\\mathbf{k}_1 + 2\\mathbf{k}_2 + 2\\mathbf{k}_3 + \\mathbf{k}_4\\right).\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "The weights $\\tfrac16(1, 2, 2, 1)$ are Simpson's rule (the quadrature rule\n",
    "$\\int_a^b g \\approx \\tfrac{b-a}{6}\\bigl(g(a) + 4g(\\tfrac{a+b}{2}) + g(b)\\bigr)$)\n",
    "applied to the probe slopes at the beginning, middle (twice), and end of the\n",
    "step, and they are chosen so that the\n",
    "update's Taylor expansion matches the true solution's through the $h^4$ term.\n",
    "\n",
    "**Proposition (RK4 converges at order 4).** *For $\\mathbf{f}$ four times\n",
    "continuously differentiable ($C^4$), the\n",
    "RK4 iterates the equation have local truncation error\n",
    "$O(h^5)$ and global error $O(h^4)$ on a fixed interval $[0, T]$.*\n",
    "\n",
    "The proof is the same accumulation argument as for Euler, fed with a long but\n",
    "mechanical five-term Taylor match;\n",
    "see @Hairer.Norsett.Wanner.1993, the standard reference, for the\n",
    "general order theory. What matters for practice is the scaling: *halving $h$\n",
    "cuts RK4's error by a factor of 16*, so a method of order $p$ delivers error\n",
    "$\\varepsilon$ at cost proportional to $\\varepsilon^{-1/p}$ field evaluations,\n",
    "the difference between $10^6$ steps and $30$ steps for the same accuracy.\n",
    "This is why solver *order* is the headline spec, and why few-step samplers\n",
    "for diffusion models (that section) rely\n",
    "on higher-order integrators. Production solvers add one more trick:\n",
    "*embedded* pairs such as Dormand--Prince `RK45` compute two orders at once\n",
    "and use the gap between them as a free error estimate\n",
    "[@Dormand.Prince.1980]. That estimate drives an adaptive step size,\n",
    "shrinking $h$ where the field bends fast and stretching it where nothing\n",
    "happens.\n",
    "\n",
    "Both claims, slope $1$ and slope $4$, are measurable. We integrate the\n",
    "spiral field to $T = 2$, where we know the exact answer from\n",
    "`#odes-solvers-matrix-exponential`, and sweep the step size across five\n",
    "octaves. On log--log axes, $\\textrm{error} \\approx C h^p$ is a line of slope\n",
    "$p$."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "odes-solvers-euler-rk4-order"
   },
   "source": [
    "def euler(f, x0, T, n):\n",
    "    \"\"\"Forward Euler: n steps of size h = T/n.\"\"\"\n",
    "    x, h = np.asarray(x0, dtype=float), T / n\n",
    "    for k in range(n):\n",
    "        x = x + h * f(x)\n",
    "    return x\n",
    "\n",
    "def rk4(f, x0, T, n):\n",
    "    \"\"\"Classical Runge-Kutta: four probe slopes per step.\"\"\"\n",
    "    x, h = np.asarray(x0, dtype=float), T / n\n",
    "    for k in range(n):\n",
    "        k1 = f(x)\n",
    "        k2 = f(x + 0.5 * h * k1)\n",
    "        k3 = f(x + 0.5 * h * k2)\n",
    "        k4 = f(x + h * k3)\n",
    "        x = x + (h / 6) * (k1 + 2 * k2 + 2 * k3 + k4)\n",
    "    return x\n",
    "\n",
    "f = lambda x: A @ x                            # the spiral field from above\n",
    "T = 2.0\n",
    "x_true = expAt_eig @ x0                        # exact endpoint via e^{At}\n",
    "ns = 2 ** np.arange(3, 10)                     # 8, 16, ..., 512 steps\n",
    "hs = T / ns\n",
    "err_eu = np.array([np.linalg.norm(euler(f, x0, T, n) - x_true) for n in ns])\n",
    "err_rk = np.array([np.linalg.norm(rk4(f, x0, T, n) - x_true) for n in ns])\n",
    "print('measured slope, Euler:', f'{np.polyfit(np.log(hs), np.log(err_eu), 1)[0]:.3f}')\n",
    "print('measured slope, RK4  :', f'{np.polyfit(np.log(hs), np.log(err_rk), 1)[0]:.3f}')\n",
    "d2l.plot(hs, [err_eu, err_rk, err_eu[-1] * (hs / hs[-1]),\n",
    "              err_rk[-1] * (hs / hs[-1]) ** 4],\n",
    "         'step size h', 'global error at T', xscale='log', yscale='log',\n",
    "         legend=['Euler', 'RK4', 'slope 1', 'slope 4'])"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "measured slope, Euler: 1.040\nmeasured slope, RK4  : 4.022\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"utf-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<svg xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"258.278125pt\" height=\"183.35625pt\" viewBox=\"0 0 258.278125 183.35625\" xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\">\n <metadata>\n  <rdf:RDF xmlns:dc=\"http://purl.org/dc/elements/1.1/\" xmlns:cc=\"http://creativecommons.org/ns#\" xmlns:rdf=\"http://www.w3.org/1999/02/22-rdf-syntax-ns#\">\n   <cc:Work>\n    <dc:type rdf:resource=\"http://purl.org/dc/dcmitype/StillImage\"/>\n    <dc:date>2026-07-19T13:41:21.706997</dc:date>\n    <dc:format>image/svg+xml</dc:format>\n    <dc:creator>\n     <cc:Agent>\n      <dc:title>Matplotlib v3.10.8, https://matplotlib.org/</dc:title>\n     </cc:Agent>\n    </dc:creator>\n   </cc:Work>\n  </rdf:RDF>\n </metadata>\n <defs>\n  <style type=\"text/css\">*{stroke-linejoin: round; stroke-linecap: butt}</style>\n </defs>\n <g id=\"figure_1\">\n  <g id=\"patch_1\">\n   <path d=\"M 0 183.35625 \nL 258.278125 183.35625 \nL 258.278125 0 \nL 0 0 \nz\n\" style=\"fill: #ffffff\"/>\n  </g>\n  <g id=\"axes_1\">\n   <g id=\"patch_2\">\n    <path d=\"M 55.778125 145.8 \nL 251.078125 145.8 \nL 251.078125 7.2 \nL 55.778125 7.2 \nz\n\" style=\"fill: #ffffff\"/>\n   </g>\n   <g id=\"matplotlib.axis_1\">\n    <g id=\"xtick_1\">\n     <g id=\"line2d_1\">\n      <path d=\"M 104.784926 145.8 \nL 104.784926 7.2 \n\" clip-path=\"url(#p6f149f3ed3)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_2\">\n      <defs>\n       <path id=\"m7cc8d5938b\" d=\"M 0 0 \nL 0 3.5 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m7cc8d5938b\" x=\"104.784926\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_1\">\n      <!-- $\\mathdefault{10^{-2}}$ -->\n      <g transform=\"translate(93.034926 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-31\" d=\"M 794 531 \nL 1825 531 \nL 1825 4091 \nL 703 3866 \nL 703 4441 \nL 1819 4666 \nL 2450 4666 \nL 2450 531 \nL 3481 531 \nL 3481 0 \nL 794 0 \nL 794 531 \nz\n\" transform=\"scale(0.015625)\"/>\n        <path id=\"DejaVuSans-30\" d=\"M 2034 4250 \nQ 1547 4250 1301 3770 \nQ 1056 3291 1056 2328 \nQ 1056 1369 1301 889 \nQ 1547 409 2034 409 \nQ 2525 409 2770 889 \nQ 3016 1369 3016 2328 \nQ 3016 3291 2770 3770 \nQ 2525 4250 2034 4250 \nz\nM 2034 4750 \nQ 2819 4750 3233 4129 \nQ 3647 3509 3647 2328 \nQ 3647 1150 3233 529 \nQ 2819 -91 2034 -91 \nQ 1250 -91 836 529 \nQ 422 1150 422 2328 \nQ 422 3509 836 4129 \nQ 1250 4750 2034 4750 \nz\n\" transform=\"scale(0.015625)\"/>\n        <path id=\"DejaVuSans-2212\" d=\"M 678 2272 \nL 4684 2272 \nL 4684 1741 \nL 678 1741 \nL 678 2272 \nz\n\" transform=\"scale(0.015625)\"/>\n        <path id=\"DejaVuSans-32\" d=\"M 1228 531 \nL 3431 531 \nL 3431 0 \nL 469 0 \nL 469 531 \nQ 828 903 1448 1529 \nQ 2069 2156 2228 2338 \nQ 2531 2678 2651 2914 \nQ 2772 3150 2772 3378 \nQ 2772 3750 2511 3984 \nQ 2250 4219 1831 4219 \nQ 1534 4219 1204 4116 \nQ 875 4013 500 3803 \nL 500 4441 \nQ 881 4594 1212 4672 \nQ 1544 4750 1819 4750 \nQ 2544 4750 2975 4387 \nQ 3406 4025 3406 3419 \nQ 3406 3131 3298 2873 \nQ 3191 2616 2906 2266 \nQ 2828 2175 2409 1742 \nQ 1991 1309 1228 531 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(0 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-2212\" transform=\"translate(128.203125 39.046875) scale(0.7)\"/>\n       <use xlink:href=\"#DejaVuSans-32\" transform=\"translate(186.855469 39.046875) scale(0.7)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_2\">\n     <g id=\"line2d_3\">\n      <path d=\"M 203.083798 145.8 \nL 203.083798 7.2 \n\" clip-path=\"url(#p6f149f3ed3)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_4\">\n      <g>\n       <use xlink:href=\"#m7cc8d5938b\" x=\"203.083798\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_2\">\n      <!-- $\\mathdefault{10^{-1}}$ -->\n      <g transform=\"translate(191.333798 160.398438) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(0 0.684375)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0.684375)\"/>\n       <use xlink:href=\"#DejaVuSans-2212\" transform=\"translate(128.203125 38.965625) scale(0.7)\"/>\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(186.855469 38.965625) scale(0.7)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_3\">\n     <g id=\"line2d_5\">\n      <defs>\n       <path id=\"ma5265fd382\" d=\"M 0 0 \nL 0 2 \n\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#ma5265fd382\" x=\"65.667872\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_4\">\n     <g id=\"line2d_6\">\n      <g>\n       <use xlink:href=\"#ma5265fd382\" x=\"75.194017\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_5\">\n     <g id=\"line2d_7\">\n      <g>\n       <use xlink:href=\"#ma5265fd382\" x=\"82.977444\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_6\">\n     <g id=\"line2d_8\">\n      <g>\n       <use xlink:href=\"#ma5265fd382\" x=\"89.558238\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_7\">\n     <g id=\"line2d_9\">\n      <g>\n       <use xlink:href=\"#ma5265fd382\" x=\"95.258781\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_8\">\n     <g id=\"line2d_10\">\n      <g>\n       <use xlink:href=\"#ma5265fd382\" x=\"100.287016\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_9\">\n     <g id=\"line2d_11\">\n      <g>\n       <use xlink:href=\"#ma5265fd382\" x=\"134.375835\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_10\">\n     <g id=\"line2d_12\">\n      <g>\n       <use xlink:href=\"#ma5265fd382\" x=\"151.685407\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_11\">\n     <g id=\"line2d_13\">\n      <g>\n       <use xlink:href=\"#ma5265fd382\" x=\"163.966744\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_12\">\n     <g id=\"line2d_14\">\n      <g>\n       <use xlink:href=\"#ma5265fd382\" x=\"173.492889\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_13\">\n     <g id=\"line2d_15\">\n      <g>\n       <use xlink:href=\"#ma5265fd382\" x=\"181.276316\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_14\">\n     <g id=\"line2d_16\">\n      <g>\n       <use xlink:href=\"#ma5265fd382\" x=\"187.85711\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_15\">\n     <g id=\"line2d_17\">\n      <g>\n       <use xlink:href=\"#ma5265fd382\" x=\"193.557653\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_16\">\n     <g id=\"line2d_18\">\n      <g>\n       <use xlink:href=\"#ma5265fd382\" x=\"198.585888\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_17\">\n     <g id=\"line2d_19\">\n      <g>\n       <use xlink:href=\"#ma5265fd382\" x=\"232.674707\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_18\">\n     <g id=\"line2d_20\">\n      <g>\n       <use xlink:href=\"#ma5265fd382\" x=\"249.984279\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.6\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_3\">\n     <!-- step size h -->\n     <g transform=\"translate(126.571094 174.076563) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-73\" d=\"M 2834 3397 \nL 2834 2853 \nQ 2591 2978 2328 3040 \nQ 2066 3103 1784 3103 \nQ 1356 3103 1142 2972 \nQ 928 2841 928 2578 \nQ 928 2378 1081 2264 \nQ 1234 2150 1697 2047 \nL 1894 2003 \nQ 2506 1872 2764 1633 \nQ 3022 1394 3022 966 \nQ 3022 478 2636 193 \nQ 2250 -91 1575 -91 \nQ 1294 -91 989 -36 \nQ 684 19 347 128 \nL 347 722 \nQ 666 556 975 473 \nQ 1284 391 1588 391 \nQ 1994 391 2212 530 \nQ 2431 669 2431 922 \nQ 2431 1156 2273 1281 \nQ 2116 1406 1581 1522 \nL 1381 1569 \nQ 847 1681 609 1914 \nQ 372 2147 372 2553 \nQ 372 3047 722 3315 \nQ 1072 3584 1716 3584 \nQ 2034 3584 2315 3537 \nQ 2597 3491 2834 3397 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-74\" d=\"M 1172 4494 \nL 1172 3500 \nL 2356 3500 \nL 2356 3053 \nL 1172 3053 \nL 1172 1153 \nQ 1172 725 1289 603 \nQ 1406 481 1766 481 \nL 2356 481 \nL 2356 0 \nL 1766 0 \nQ 1100 0 847 248 \nQ 594 497 594 1153 \nL 594 3053 \nL 172 3053 \nL 172 3500 \nL 594 3500 \nL 594 4494 \nL 1172 4494 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-65\" d=\"M 3597 1894 \nL 3597 1613 \nL 953 1613 \nQ 991 1019 1311 708 \nQ 1631 397 2203 397 \nQ 2534 397 2845 478 \nQ 3156 559 3463 722 \nL 3463 178 \nQ 3153 47 2828 -22 \nQ 2503 -91 2169 -91 \nQ 1331 -91 842 396 \nQ 353 884 353 1716 \nQ 353 2575 817 3079 \nQ 1281 3584 2069 3584 \nQ 2775 3584 3186 3129 \nQ 3597 2675 3597 1894 \nz\nM 3022 2063 \nQ 3016 2534 2758 2815 \nQ 2500 3097 2075 3097 \nQ 1594 3097 1305 2825 \nQ 1016 2553 972 2059 \nL 3022 2063 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-70\" d=\"M 1159 525 \nL 1159 -1331 \nL 581 -1331 \nL 581 3500 \nL 1159 3500 \nL 1159 2969 \nQ 1341 3281 1617 3432 \nQ 1894 3584 2278 3584 \nQ 2916 3584 3314 3078 \nQ 3713 2572 3713 1747 \nQ 3713 922 3314 415 \nQ 2916 -91 2278 -91 \nQ 1894 -91 1617 61 \nQ 1341 213 1159 525 \nz\nM 3116 1747 \nQ 3116 2381 2855 2742 \nQ 2594 3103 2138 3103 \nQ 1681 3103 1420 2742 \nQ 1159 2381 1159 1747 \nQ 1159 1113 1420 752 \nQ 1681 391 2138 391 \nQ 2594 391 2855 752 \nQ 3116 1113 3116 1747 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-20\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-69\" d=\"M 603 3500 \nL 1178 3500 \nL 1178 0 \nL 603 0 \nL 603 3500 \nz\nM 603 4863 \nL 1178 4863 \nL 1178 4134 \nL 603 4134 \nL 603 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-7a\" d=\"M 353 3500 \nL 3084 3500 \nL 3084 2975 \nL 922 459 \nL 3084 459 \nL 3084 0 \nL 275 0 \nL 275 525 \nL 2438 3041 \nL 353 3041 \nL 353 3500 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-68\" d=\"M 3513 2113 \nL 3513 0 \nL 2938 0 \nL 2938 2094 \nQ 2938 2591 2744 2837 \nQ 2550 3084 2163 3084 \nQ 1697 3084 1428 2787 \nQ 1159 2491 1159 1978 \nL 1159 0 \nL 581 0 \nL 581 4863 \nL 1159 4863 \nL 1159 2956 \nQ 1366 3272 1645 3428 \nQ 1925 3584 2291 3584 \nQ 2894 3584 3203 3211 \nQ 3513 2838 3513 2113 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-73\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(52.099609 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(91.308594 0)\"/>\n      <use xlink:href=\"#DejaVuSans-70\" transform=\"translate(152.832031 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(216.308594 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(248.095703 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(300.195312 0)\"/>\n      <use xlink:href=\"#DejaVuSans-7a\" transform=\"translate(327.978516 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(380.46875 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(441.992188 0)\"/>\n      <use xlink:href=\"#DejaVuSans-68\" transform=\"translate(473.779297 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"matplotlib.axis_2\">\n    <g id=\"ytick_1\">\n     <g id=\"line2d_21\">\n      <path d=\"M 55.778125 124.239803 \nL 251.078125 124.239803 \n\" clip-path=\"url(#p6f149f3ed3)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_22\">\n      <defs>\n       <path id=\"ma2d0baa489\" d=\"M 0 0 \nL -3.5 0 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#ma2d0baa489\" x=\"55.778125\" y=\"124.239803\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_4\">\n      <!-- $\\mathdefault{10^{-10}}$ -->\n      <g transform=\"translate(20.878125 128.039021) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(0 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-2212\" transform=\"translate(128.203125 39.046875) scale(0.7)\"/>\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(186.855469 39.046875) scale(0.7)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(231.391602 39.046875) scale(0.7)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_2\">\n     <g id=\"line2d_23\">\n      <path d=\"M 55.778125 100.788062 \nL 251.078125 100.788062 \n\" clip-path=\"url(#p6f149f3ed3)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_24\">\n      <g>\n       <use xlink:href=\"#ma2d0baa489\" x=\"55.778125\" y=\"100.788062\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_5\">\n      <!-- $\\mathdefault{10^{-8}}$ -->\n      <g transform=\"translate(25.278125 104.58728) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-38\" d=\"M 2034 2216 \nQ 1584 2216 1326 1975 \nQ 1069 1734 1069 1313 \nQ 1069 891 1326 650 \nQ 1584 409 2034 409 \nQ 2484 409 2743 651 \nQ 3003 894 3003 1313 \nQ 3003 1734 2745 1975 \nQ 2488 2216 2034 2216 \nz\nM 1403 2484 \nQ 997 2584 770 2862 \nQ 544 3141 544 3541 \nQ 544 4100 942 4425 \nQ 1341 4750 2034 4750 \nQ 2731 4750 3128 4425 \nQ 3525 4100 3525 3541 \nQ 3525 3141 3298 2862 \nQ 3072 2584 2669 2484 \nQ 3125 2378 3379 2068 \nQ 3634 1759 3634 1313 \nQ 3634 634 3220 271 \nQ 2806 -91 2034 -91 \nQ 1263 -91 848 271 \nQ 434 634 434 1313 \nQ 434 1759 690 2068 \nQ 947 2378 1403 2484 \nz\nM 1172 3481 \nQ 1172 3119 1398 2916 \nQ 1625 2713 2034 2713 \nQ 2441 2713 2670 2916 \nQ 2900 3119 2900 3481 \nQ 2900 3844 2670 4047 \nQ 2441 4250 2034 4250 \nQ 1625 4250 1398 4047 \nQ 1172 3844 1172 3481 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(0 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-2212\" transform=\"translate(128.203125 39.046875) scale(0.7)\"/>\n       <use xlink:href=\"#DejaVuSans-38\" transform=\"translate(186.855469 39.046875) scale(0.7)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_3\">\n     <g id=\"line2d_25\">\n      <path d=\"M 55.778125 77.33632 \nL 251.078125 77.33632 \n\" clip-path=\"url(#p6f149f3ed3)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_26\">\n      <g>\n       <use xlink:href=\"#ma2d0baa489\" x=\"55.778125\" y=\"77.33632\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_6\">\n      <!-- $\\mathdefault{10^{-6}}$ -->\n      <g transform=\"translate(25.278125 81.135539) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-36\" d=\"M 2113 2584 \nQ 1688 2584 1439 2293 \nQ 1191 2003 1191 1497 \nQ 1191 994 1439 701 \nQ 1688 409 2113 409 \nQ 2538 409 2786 701 \nQ 3034 994 3034 1497 \nQ 3034 2003 2786 2293 \nQ 2538 2584 2113 2584 \nz\nM 3366 4563 \nL 3366 3988 \nQ 3128 4100 2886 4159 \nQ 2644 4219 2406 4219 \nQ 1781 4219 1451 3797 \nQ 1122 3375 1075 2522 \nQ 1259 2794 1537 2939 \nQ 1816 3084 2150 3084 \nQ 2853 3084 3261 2657 \nQ 3669 2231 3669 1497 \nQ 3669 778 3244 343 \nQ 2819 -91 2113 -91 \nQ 1303 -91 875 529 \nQ 447 1150 447 2328 \nQ 447 3434 972 4092 \nQ 1497 4750 2381 4750 \nQ 2619 4750 2861 4703 \nQ 3103 4656 3366 4563 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(0 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-2212\" transform=\"translate(128.203125 39.046875) scale(0.7)\"/>\n       <use xlink:href=\"#DejaVuSans-36\" transform=\"translate(186.855469 39.046875) scale(0.7)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_4\">\n     <g id=\"line2d_27\">\n      <path d=\"M 55.778125 53.884579 \nL 251.078125 53.884579 \n\" clip-path=\"url(#p6f149f3ed3)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_28\">\n      <g>\n       <use xlink:href=\"#ma2d0baa489\" x=\"55.778125\" y=\"53.884579\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_7\">\n      <!-- $\\mathdefault{10^{-4}}$ -->\n      <g transform=\"translate(25.278125 57.683798) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-34\" d=\"M 2419 4116 \nL 825 1625 \nL 2419 1625 \nL 2419 4116 \nz\nM 2253 4666 \nL 3047 4666 \nL 3047 1625 \nL 3713 1625 \nL 3713 1100 \nL 3047 1100 \nL 3047 0 \nL 2419 0 \nL 2419 1100 \nL 313 1100 \nL 313 1709 \nL 2253 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(0 0.684375)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0.684375)\"/>\n       <use xlink:href=\"#DejaVuSans-2212\" transform=\"translate(128.203125 38.965625) scale(0.7)\"/>\n       <use xlink:href=\"#DejaVuSans-34\" transform=\"translate(186.855469 38.965625) scale(0.7)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_5\">\n     <g id=\"line2d_29\">\n      <path d=\"M 55.778125 30.432838 \nL 251.078125 30.432838 \n\" clip-path=\"url(#p6f149f3ed3)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_30\">\n      <g>\n       <use xlink:href=\"#ma2d0baa489\" x=\"55.778125\" y=\"30.432838\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_8\">\n      <!-- $\\mathdefault{10^{-2}}$ -->\n      <g transform=\"translate(25.278125 34.232057) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(0 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-2212\" transform=\"translate(128.203125 39.046875) scale(0.7)\"/>\n       <use xlink:href=\"#DejaVuSans-32\" transform=\"translate(186.855469 39.046875) scale(0.7)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_9\">\n     <!-- global error at T -->\n     <g transform=\"translate(14.798437 116.697656) rotate(-90) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-67\" d=\"M 2906 1791 \nQ 2906 2416 2648 2759 \nQ 2391 3103 1925 3103 \nQ 1463 3103 1205 2759 \nQ 947 2416 947 1791 \nQ 947 1169 1205 825 \nQ 1463 481 1925 481 \nQ 2391 481 2648 825 \nQ 2906 1169 2906 1791 \nz\nM 3481 434 \nQ 3481 -459 3084 -895 \nQ 2688 -1331 1869 -1331 \nQ 1566 -1331 1297 -1286 \nQ 1028 -1241 775 -1147 \nL 775 -588 \nQ 1028 -725 1275 -790 \nQ 1522 -856 1778 -856 \nQ 2344 -856 2625 -561 \nQ 2906 -266 2906 331 \nL 2906 616 \nQ 2728 306 2450 153 \nQ 2172 0 1784 0 \nQ 1141 0 747 490 \nQ 353 981 353 1791 \nQ 353 2603 747 3093 \nQ 1141 3584 1784 3584 \nQ 2172 3584 2450 3431 \nQ 2728 3278 2906 2969 \nL 2906 3500 \nL 3481 3500 \nL 3481 434 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6c\" d=\"M 603 4863 \nL 1178 4863 \nL 1178 0 \nL 603 0 \nL 603 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6f\" d=\"M 1959 3097 \nQ 1497 3097 1228 2736 \nQ 959 2375 959 1747 \nQ 959 1119 1226 758 \nQ 1494 397 1959 397 \nQ 2419 397 2687 759 \nQ 2956 1122 2956 1747 \nQ 2956 2369 2687 2733 \nQ 2419 3097 1959 3097 \nz\nM 1959 3584 \nQ 2709 3584 3137 3096 \nQ 3566 2609 3566 1747 \nQ 3566 888 3137 398 \nQ 2709 -91 1959 -91 \nQ 1206 -91 779 398 \nQ 353 888 353 1747 \nQ 353 2609 779 3096 \nQ 1206 3584 1959 3584 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-62\" d=\"M 3116 1747 \nQ 3116 2381 2855 2742 \nQ 2594 3103 2138 3103 \nQ 1681 3103 1420 2742 \nQ 1159 2381 1159 1747 \nQ 1159 1113 1420 752 \nQ 1681 391 2138 391 \nQ 2594 391 2855 752 \nQ 3116 1113 3116 1747 \nz\nM 1159 2969 \nQ 1341 3281 1617 3432 \nQ 1894 3584 2278 3584 \nQ 2916 3584 3314 3078 \nQ 3713 2572 3713 1747 \nQ 3713 922 3314 415 \nQ 2916 -91 2278 -91 \nQ 1894 -91 1617 61 \nQ 1341 213 1159 525 \nL 1159 0 \nL 581 0 \nL 581 4863 \nL 1159 4863 \nL 1159 2969 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-61\" d=\"M 2194 1759 \nQ 1497 1759 1228 1600 \nQ 959 1441 959 1056 \nQ 959 750 1161 570 \nQ 1363 391 1709 391 \nQ 2188 391 2477 730 \nQ 2766 1069 2766 1631 \nL 2766 1759 \nL 2194 1759 \nz\nM 3341 1997 \nL 3341 0 \nL 2766 0 \nL 2766 531 \nQ 2569 213 2275 61 \nQ 1981 -91 1556 -91 \nQ 1019 -91 701 211 \nQ 384 513 384 1019 \nQ 384 1609 779 1909 \nQ 1175 2209 1959 2209 \nL 2766 2209 \nL 2766 2266 \nQ 2766 2663 2505 2880 \nQ 2244 3097 1772 3097 \nQ 1472 3097 1187 3025 \nQ 903 2953 641 2809 \nL 641 3341 \nQ 956 3463 1253 3523 \nQ 1550 3584 1831 3584 \nQ 2591 3584 2966 3190 \nQ 3341 2797 3341 1997 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-72\" d=\"M 2631 2963 \nQ 2534 3019 2420 3045 \nQ 2306 3072 2169 3072 \nQ 1681 3072 1420 2755 \nQ 1159 2438 1159 1844 \nL 1159 0 \nL 581 0 \nL 581 3500 \nL 1159 3500 \nL 1159 2956 \nQ 1341 3275 1631 3429 \nQ 1922 3584 2338 3584 \nQ 2397 3584 2469 3576 \nQ 2541 3569 2628 3553 \nL 2631 2963 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-54\" d=\"M -19 4666 \nL 3928 4666 \nL 3928 4134 \nL 2272 4134 \nL 2272 0 \nL 1638 0 \nL 1638 4134 \nL -19 4134 \nL -19 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-67\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(63.476562 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(91.259766 0)\"/>\n      <use xlink:href=\"#DejaVuSans-62\" transform=\"translate(152.441406 0)\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(215.917969 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(277.197266 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(304.980469 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(336.767578 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(398.291016 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(437.654297 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(476.517578 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(537.699219 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(578.8125 0)\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(610.599609 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(671.878906 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(711.087891 0)\"/>\n      <use xlink:href=\"#DejaVuSans-54\" transform=\"translate(742.875 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"line2d_31\">\n    <path d=\"M 242.200852 13.5 \nL 212.609943 17.5283 \nL 183.019034 21.296756 \nL 153.428125 24.94271 \nL 123.837216 28.529768 \nL 94.246307 32.088003 \nL 64.655398 35.631987 \n\" clip-path=\"url(#p6f149f3ed3)\" style=\"fill: none; stroke: #1f77b4; stroke-width: 1.5; stroke-linecap: square\"/>\n   </g>\n   <g id=\"line2d_32\">\n    <path d=\"M 242.200852 54.263535 \nL 212.609943 68.646627 \nL 183.019034 82.898175 \nL 153.428125 97.083729 \nL 123.837216 111.23621 \nL 94.246307 125.372139 \nL 64.655398 139.5 \n\" clip-path=\"url(#p6f149f3ed3)\" style=\"fill: none; stroke-dasharray: 5.55,2.4; stroke-dashoffset: 0; stroke: #bf00bf; stroke-width: 1.5\"/>\n   </g>\n   <g id=\"line2d_33\">\n    <path d=\"M 242.200852 14.452954 \nL 212.609943 17.982793 \nL 183.019034 21.512632 \nL 153.428125 25.04247 \nL 123.837216 28.572309 \nL 94.246307 32.102148 \nL 64.655398 35.631987 \n\" clip-path=\"url(#p6f149f3ed3)\" style=\"fill: none; stroke-dasharray: 9.6,2.4,1.5,2.4; stroke-dashoffset: 0; stroke: #008000; stroke-width: 1.5\"/>\n   </g>\n   <g id=\"line2d_34\">\n    <path d=\"M 242.200852 54.78387 \nL 212.609943 68.903225 \nL 183.019034 83.02258 \nL 153.428125 97.141935 \nL 123.837216 111.26129 \nL 94.246307 125.380645 \nL 64.655398 139.5 \n\" clip-path=\"url(#p6f149f3ed3)\" style=\"fill: none; stroke-dasharray: 1.5,2.475; stroke-dashoffset: 0; stroke: #ff0000; stroke-width: 1.5\"/>\n   </g>\n   <g id=\"patch_3\">\n    <path d=\"M 55.778125 145.8 \nL 55.778125 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_4\">\n    <path d=\"M 251.078125 145.8 \nL 251.078125 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_5\">\n    <path d=\"M 55.778125 145.8 \nL 251.078125 145.8 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_6\">\n    <path d=\"M 55.778125 7.2 \nL 251.078125 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"legend_1\">\n    <g id=\"patch_7\">\n     <path d=\"M 62.778125 107.35625 \nL 130.926562 107.35625 \nQ 132.926562 107.35625 132.926562 105.35625 \nL 132.926562 47.64375 \nQ 132.926562 45.64375 130.926562 45.64375 \nL 62.778125 45.64375 \nQ 60.778125 45.64375 60.778125 47.64375 \nL 60.778125 105.35625 \nQ 60.778125 107.35625 62.778125 107.35625 \nz\n\" style=\"fill: #ffffff; opacity: 0.8; stroke: #cccccc; stroke-linejoin: miter\"/>\n    </g>\n    <g id=\"line2d_35\">\n     <path d=\"M 64.778125 53.742188 \nL 74.778125 53.742188 \nL 84.778125 53.742188 \n\" style=\"fill: none; stroke: #1f77b4; stroke-width: 1.5; stroke-linecap: square\"/>\n    </g>\n    <g id=\"text_10\">\n     <!-- Euler -->\n     <g transform=\"translate(92.778125 57.242188) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-45\" d=\"M 628 4666 \nL 3578 4666 \nL 3578 4134 \nL 1259 4134 \nL 1259 2753 \nL 3481 2753 \nL 3481 2222 \nL 1259 2222 \nL 1259 531 \nL 3634 531 \nL 3634 0 \nL 628 0 \nL 628 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-75\" d=\"M 544 1381 \nL 544 3500 \nL 1119 3500 \nL 1119 1403 \nQ 1119 906 1312 657 \nQ 1506 409 1894 409 \nQ 2359 409 2629 706 \nQ 2900 1003 2900 1516 \nL 2900 3500 \nL 3475 3500 \nL 3475 0 \nL 2900 0 \nL 2900 538 \nQ 2691 219 2414 64 \nQ 2138 -91 1772 -91 \nQ 1169 -91 856 284 \nQ 544 659 544 1381 \nz\nM 1991 3584 \nL 1991 3584 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-45\"/>\n      <use xlink:href=\"#DejaVuSans-75\" transform=\"translate(63.183594 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(126.5625 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(154.345703 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(215.869141 0)\"/>\n     </g>\n    </g>\n    <g id=\"line2d_36\">\n     <path d=\"M 64.778125 68.420313 \nL 74.778125 68.420313 \nL 84.778125 68.420313 \n\" style=\"fill: none; stroke-dasharray: 5.55,2.4; stroke-dashoffset: 0; stroke: #bf00bf; stroke-width: 1.5\"/>\n    </g>\n    <g id=\"text_11\">\n     <!-- RK4 -->\n     <g transform=\"translate(92.778125 71.920313) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-52\" d=\"M 2841 2188 \nQ 3044 2119 3236 1894 \nQ 3428 1669 3622 1275 \nL 4263 0 \nL 3584 0 \nL 2988 1197 \nQ 2756 1666 2539 1819 \nQ 2322 1972 1947 1972 \nL 1259 1972 \nL 1259 0 \nL 628 0 \nL 628 4666 \nL 2053 4666 \nQ 2853 4666 3247 4331 \nQ 3641 3997 3641 3322 \nQ 3641 2881 3436 2590 \nQ 3231 2300 2841 2188 \nz\nM 1259 4147 \nL 1259 2491 \nL 2053 2491 \nQ 2509 2491 2742 2702 \nQ 2975 2913 2975 3322 \nQ 2975 3731 2742 3939 \nQ 2509 4147 2053 4147 \nL 1259 4147 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-4b\" d=\"M 628 4666 \nL 1259 4666 \nL 1259 2694 \nL 3353 4666 \nL 4166 4666 \nL 1850 2491 \nL 4331 0 \nL 3500 0 \nL 1259 2247 \nL 1259 0 \nL 628 0 \nL 628 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-52\"/>\n      <use xlink:href=\"#DejaVuSans-4b\" transform=\"translate(69.482422 0)\"/>\n      <use xlink:href=\"#DejaVuSans-34\" transform=\"translate(135.058594 0)\"/>\n     </g>\n    </g>\n    <g id=\"line2d_37\">\n     <path d=\"M 64.778125 83.098438 \nL 74.778125 83.098438 \nL 84.778125 83.098438 \n\" style=\"fill: none; stroke-dasharray: 9.6,2.4,1.5,2.4; stroke-dashoffset: 0; stroke: #008000; stroke-width: 1.5\"/>\n    </g>\n    <g id=\"text_12\">\n     <!-- slope 1 -->\n     <g transform=\"translate(92.778125 86.598438) scale(0.1 -0.1)\">\n      <use xlink:href=\"#DejaVuSans-73\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(52.099609 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(79.882812 0)\"/>\n      <use xlink:href=\"#DejaVuSans-70\" transform=\"translate(141.064453 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(204.541016 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(266.064453 0)\"/>\n      <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(297.851562 0)\"/>\n     </g>\n    </g>\n    <g id=\"line2d_38\">\n     <path d=\"M 64.778125 97.776563 \nL 74.778125 97.776563 \nL 84.778125 97.776563 \n\" style=\"fill: none; stroke-dasharray: 1.5,2.475; stroke-dashoffset: 0; stroke: #ff0000; stroke-width: 1.5\"/>\n    </g>\n    <g id=\"text_13\">\n     <!-- slope 4 -->\n     <g transform=\"translate(92.778125 101.276563) scale(0.1 -0.1)\">\n      <use xlink:href=\"#DejaVuSans-73\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(52.099609 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(79.882812 0)\"/>\n      <use xlink:href=\"#DejaVuSans-70\" transform=\"translate(141.064453 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(204.541016 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(266.064453 0)\"/>\n      <use xlink:href=\"#DejaVuSans-34\" transform=\"translate(297.851562 0)\"/>\n     </g>\n    </g>\n   </g>\n  </g>\n </g>\n <defs>\n  <clipPath id=\"p6f149f3ed3\">\n   <rect x=\"55.778125\" y=\"7.2\" width=\"195.3\" height=\"138.6\"/>\n  </clipPath>\n </defs>\n</svg>\n"
     },
     "metadata": {}
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-odes-solvers-md-4 fw=all -->\n",
    "\n",
    "The fitted slopes land within a few percent of the theoretical orders $1$ and\n",
    "$4$, and the RK4 curve sits *nearly nine decades* below Euler at the smallest\n",
    "step ($3.6 \\times 10^{-3}$ versus $5.0 \\times 10^{-12}$): same trajectory,\n",
    "same horizon, four times the work per step, nearly a billion times the\n",
    "accuracy.\n",
    "\n",
    "### Stiffness and Implicit Methods\n",
    "\n",
    "Order does not decide everything. Accuracy says how well you track the\n",
    "solution; **stability** asks whether your errors quietly die out or compound\n",
    "into an explosion, and for *explicit* methods like Euler and RK4, stability\n",
    "puts a hard ceiling on the step size. The phenomenon is already visible in one\n",
    "dimension.\n",
    "\n",
    "**Proposition (stability of Euler on the test equation).** *On\n",
    "$\\dot{x} = -\\lambda x$ with $\\lambda > 0$ (true solution: decay to zero):\n",
    "forward Euler gives $x_n = (1 - h\\lambda)^n x_0$, which decays iff\n",
    "$|1 - h\\lambda| < 1$, i.e.*\n",
    "\n",
    "$$\n",
    "h < \\frac{2}{\\lambda} .\n",
    "$$\n",
    "\n",
    "*Backward (implicit) Euler,\n",
    "$\\mathbf{x}_{n+1} = \\mathbf{x}_n + h\\,\\mathbf{f}(\\mathbf{x}_{n+1}, t_{n+1})$,\n",
    "gives $x_n = (1 + h\\lambda)^{-n} x_0$, which decays for* **every** *$h > 0$.*\n",
    "\n",
    "**Proof.** Forward Euler: $x_{n+1} = x_n - h\\lambda x_n = (1 - h\\lambda)x_n$;\n",
    "iterate. The factor has magnitude below $1$ iff $-1 < 1 - h\\lambda < 1$, and\n",
    "the left inequality is the binding one: $h\\lambda < 2$. Backward Euler:\n",
    "$x_{n+1} = x_n - h\\lambda x_{n+1}$, so $x_{n+1} = x_n / (1 + h\\lambda)$, and\n",
    "$0 < (1 + h\\lambda)^{-1} < 1$ for all $h > 0$. $\\blacksquare$\n",
    "\n",
    "Past the threshold, forward Euler *amplifies*: each step overshoots the\n",
    "origin and lands farther away on the\n",
    "other side, an oscillating divergence with growth factor $|1 - h\\lambda|$. A\n",
    "method that decays on the test equation for every $h > 0$, as backward Euler\n",
    "does, is called **A-stable** [@Dahlquist.1963]. (Our $\\lambda$ is real,\n",
    "but the full picture lives in the complex plane. Write the test equation as\n",
    "$\\dot{y} = \\lambda y$ with complex $\\lambda$, so that a decaying, rotating\n",
    "mode has $\\operatorname{Re}\\lambda < 0$; a method's\n",
    "**stability region** is the set of $z = h\\lambda$ for which its update\n",
    "decays, and A-stability proper demands that this region contain the entire\n",
    "left half-plane. the figure draws the regions;\n",
    "the computation above, $z = -h\\lambda$ with $\\lambda > 0$, is their slice\n",
    "along the negative real axis.) The price of\n",
    "implicitness is that each step\n",
    "*defines* $\\mathbf{x}_{n+1}$ only implicitly: you must solve an equation per\n",
    "step, a linear solve when $\\mathbf{f}$ is linear and a few Newton iterations\n",
    "(that section) otherwise.\n",
    "\n",
    "![Stability regions in the complex plane of $z = h\\lambda$ for the test equation $\\dot{y} = \\lambda y$ (a decaying mode sits in the left half-plane; our real computation is the slice along the negative real axis). Forward Euler decays only inside the disc $\\lvert 1 + z \\rvert < 1$, whose real slice is $-2 < h\\lambda < 0$; RK4 enlarges it to a lobed region reaching $\\approx -2.79$ on the real axis; backward Euler is stable everywhere *except* the disc around $+1$, in particular on the entire left half-plane, which is A-stability.](../img/mdl-dyn-stability-regions.svg)\n",
    "\n",
    "Why tolerate that price? **Stiffness.** A system is *stiff* when the step\n",
    "size that *stability* forces on an explicit method is far smaller than the\n",
    "one *accuracy* would need. The standard linear source is eigenvalues spread\n",
    "over wildly different scales, say eigenvalues $-50$ and $-1$, i.e. decay\n",
    "rates $\\lambda_{\\textrm{fast}} = 50$ and $\\lambda_{\\textrm{slow}} = 1$. The\n",
    "fast mode dies almost immediately, and what is left to track is the slow,\n",
    "smooth mode, for which a large step would be perfectly *accurate*. But\n",
    "forward Euler's *stability* constraint $h < 2/\\lambda_{\\textrm{fast}}$ is set\n",
    "by the fastest decay rate, including modes that died long ago and\n",
    "contribute nothing to the answer. The dead mode governs your budget: that is\n",
    "stiffness. An implicit method deletes the constraint and lets the step size\n",
    "follow the physics you actually care about. The sweep below shows the\n",
    "forward-Euler threshold appear exactly at $h = 2/\\lambda = 0.04$, backward\n",
    "Euler decaying at every step size, and the two-scale system blowing\n",
    "up in the fast component that had already decayed to $10^{-22}$."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "odes-solvers-stiffness-sweep"
   },
   "source": [
    "lam_f = 50.0                                   # the test equation dx/dt = -50 x\n",
    "T = 1.0\n",
    "print('forward Euler is stable iff h < 2/lambda =', 2 / lam_f)\n",
    "print(f'{\"h\":>8} {\"|1-h*lam|\":>10} {\"forward x(T)\":>14} {\"backward x(T)\":>14}')\n",
    "for n in [10, 20, 26, 40, 100]:\n",
    "    h = T / n\n",
    "    fwd = (1 - h * lam_f) ** n                 # forward Euler after n steps\n",
    "    bwd = (1 + h * lam_f) ** (-n)              # backward Euler after n steps\n",
    "    print(f'{h:8.3f} {abs(1 - h * lam_f):10.2f} {fwd:14.2e} {bwd:14.2e}')\n",
    "print('exact x(T) = e^{-50} =', f'{np.exp(-lam_f * T):.2e}')\n",
    "A_stiff = np.diag([-50.0, -1.0])               # one fast mode, one slow mode\n",
    "x_stiff = euler(lambda x: A_stiff @ x, np.array([1.0, 1.0]), T, 20)\n",
    "print('stiff system, forward Euler with h=0.05: x(T) =', x_stiff.round(2),\n",
    "      '(the long-dead fast mode explodes)')"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "forward Euler is stable iff h < 2/lambda = 0.04\n       h  |1-h*lam|   forward x(T)  backward x(T)\n   0.100       4.00       1.05e+06       1.65e-08\n   0.050       1.50       3.33e+03       1.31e-11\n   0.038       0.92       1.25e-01       7.73e-13\n   0.025       0.25       8.27e-25       8.18e-15\n   0.010       0.50       7.89e-31       2.46e-18\nexact x(T) = e^{-50} = 1.93e-22\nstiff system, forward Euler with h=0.05: x(T) = [3.32526e+03 3.60000e-01] (the long-dead fast mode explodes)\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-odes-solvers-md-5 fw=all -->\n",
    "\n",
    "Note what backward Euler's unconditional stability does and does not buy: at\n",
    "$h = 0.1$ it returns $1.65 \\times 10^{-8}$ where the truth is\n",
    "$1.93 \\times 10^{-22}$: *stable* (it decays, and further steps decay\n",
    "further) but not *accurate*. Stability prevents the explosion; only a small\n",
    "step or a higher order makes you right. The practical doctrine, which carries\n",
    "to every learned ODE in that section: use\n",
    "explicit adaptive solvers by default, and reach for implicit methods when the\n",
    "dynamics are stiff. For a nonlinear or learned field, stiff means that the\n",
    "Jacobian's eigenvalues along the trajectory are spread over widely different\n",
    "scales.\n",
    "\n",
    "### Gradient Descent Is a Solver\n",
    "\n",
    "Everything above already governs the oldest dynamical system in deep\n",
    "learning: training itself. Descending a loss $L$ by infinitesimal steps is\n",
    "the **gradient flow**\n",
    "\n",
    "$$\n",
    "\\dot{\\mathbf{x}} = -\\nabla L(\\mathbf{x}),\n",
    "$$\n",
    "\n",
    "a vector field whose fixed points are exactly the critical points of $L$ and\n",
    "along whose trajectories the loss can only fall:\n",
    "$\\tfrac{d}{dt} L(\\mathbf{x}(t)) = \\nabla L^\\top \\dot{\\mathbf{x}} = -\\|\\nabla L\\|^2 \\le 0$.\n",
    "Gradient descent with learning rate $\\eta$ is *forward Euler* on this flow\n",
    "($\\mathbf{x}_{k+1} = \\mathbf{x}_k - \\eta\\, \\nabla L(\\mathbf{x}_k)$ is\n",
    "the equation with $h = \\eta$), so the solver theory of\n",
    "this section *is* optimization theory.\n",
    "that section derived the learning-rate\n",
    "ceiling $\\eta < 2/\\lambda_{\\max}(H)$ from the quadratic analysis and ran\n",
    "essentially the experiment below; what is new here is the reading. Near a\n",
    "minimum with Hessian $H$ the dynamics linearize, mode by eigenmode, onto the\n",
    "test equation:\n",
    "\n",
    "$$\n",
    "\\dot{\\mathbf{x}} = -\\nabla L(\\mathbf{x})\n",
    "\\;\\;\\xrightarrow{\\textrm{linearize}}\\;\\;\n",
    "\\dot{\\boldsymbol{\\delta}} = -H\\boldsymbol{\\delta}\n",
    "\\;\\;\\xrightarrow{\\textrm{Euler}, \\, h = \\eta}\\;\\;\n",
    "\\boldsymbol{\\delta}_{k+1} = (I - \\eta H)\\,\\boldsymbol{\\delta}_k\n",
    "\\;\\;\\xrightarrow{\\textrm{stability}}\\;\\;\n",
    "\\eta < \\frac{2}{\\lambda_{\\max}(H)} .\n",
    "$$\n",
    "\n",
    "The solver bound $h < 2/\\lambda$ *is* the learning-rate ceiling: a diverging\n",
    "learning rate is a solver instability, and an ill-conditioned Hessian is a\n",
    "*stiff* gradient flow, with the steep, already-converged curvature directions\n",
    "setting the stability limit under which the shallow directions must then\n",
    "crawl. Momentum, Polyak's heavy-ball update, discretizes the second-order ODE\n",
    "$m\\, \\ddot{\\mathbf{x}} + \\gamma\\, \\dot{\\mathbf{x}} = -\\nabla L(\\mathbf{x})$\n",
    "analyzed in that section\n",
    "[@Polyak.1964]. We rerun that section's experiment, read now as the\n",
    "twin of the stiffness sweep, on a\n",
    "quadratic with $\\lambda_{\\max} = 10$: predicted flip at\n",
    "$\\eta = 2/10 = 0.2$."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "mdl-odes-solvers-gradient-descent-is-a-solver"
   },
   "source": [
    "H = np.diag([10.0, 1.0])                       # curvatures: lambda_max = 10\n",
    "print('GD on this quadratic is stable iff eta < 2/lambda_max =', 2 / H[0, 0])\n",
    "print(f'{\"eta\":>8} {\"|1-eta*lam_max|\":>16} {\"||x|| after 100 GD steps\":>25}')\n",
    "for eta in [0.05, 0.15, 0.19, 0.201, 0.21, 0.25]:\n",
    "    x = np.array([1.0, 1.0])\n",
    "    for _ in range(100):\n",
    "        x = x - eta * (H @ x)                  # gradient descent = forward Euler\n",
    "    print(f'{eta:8.3f} {abs(1 - eta * H[0, 0]):16.2f} {np.linalg.norm(x):25.2e}')"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "GD on this quadratic is stable iff eta < 2/lambda_max = 0.2\n     eta  |1-eta*lam_max|  ||x|| after 100 GD steps\n   0.050             0.50                  5.92e-03\n   0.150             0.50                  8.75e-08\n   0.190             0.90                  2.66e-05\n   0.201             1.01                  2.70e+00\n   0.210             1.10                  1.38e+04\n   0.250             1.50                  4.07e+17\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-odes-solvers-md-6 fw=all -->\n",
    "\n",
    "The verdict flips exactly at the predicted threshold. At $\\eta = 0.19$ the\n",
    "fast coordinate contracts by $|1 - 1.9| = 0.9$ per step and a hundred steps\n",
    "crush the iterate to $10^{-5}$; at $\\eta = 0.201$ the factor is $-1.01$ and\n",
    "the same hundred steps have *grown* it, oscillating in sign, past its\n",
    "starting point; at $\\eta = 0.25$ it has exploded seventeen orders of\n",
    "magnitude. Meanwhile the slow direction ($\\lambda = 1$) was converging\n",
    "at every one of these rates: the divergence is manufactured\n",
    "entirely by the loss surface's stiffest mode, its own \"fast dead mode.\"\n",
    "\n",
    "## Neural ODEs and the Adjoint Method\n",
    "\n",
    "### Residual Networks Are Euler Steps\n",
    "\n",
    "Now the bridge to deep learning, and it is one line long.\n",
    "\n",
    "**Proposition (a residual block is an Euler step).** *The residual update\n",
    "[@He.Zhang.Ren.ea.2016]*\n",
    "\n",
    "$$\n",
    "\\mathbf{x}_{l+1} = \\mathbf{x}_l + \\mathbf{f}_\\theta(\\mathbf{x}_l)\n",
    "$$\n",
    "\n",
    "*is exactly one forward-Euler step the equation of the\n",
    "ODE $\\dot{\\mathbf{x}} = \\mathbf{f}_\\theta(\\mathbf{x})$ with step size\n",
    "$h = 1$.*\n",
    "\n",
    "**Proof.** Set $h = 1$ and $\\mathbf{f} = \\mathbf{f}_\\theta$ in\n",
    "the equation; the two updates are the same formula.\n",
    "$\\blacksquare$\n",
    "\n",
    "![Left: a stack of residual blocks, each adding the learned branch $\\mathbf{f}_\\theta(\\mathbf{x}_l)$ to an identity skip. Right: the same update read as numerical integration, forward-Euler steps of size $h = 1$ tracking the smooth flow of the field $\\mathbf{f}_\\theta$. Depth is integration time.](../img/mdl-dyn-resnet-as-euler.svg)\n",
    "\n",
    "A trivial identification (the figure) with non-trivial consequences. Read in one\n",
    "direction: a ResNet with $N$ blocks is a *solver*; it integrates a vector\n",
    "field for $N$ unit steps of time, and \"depth\" is a discretization of a\n",
    "continuous deformation of the representation [@E.2017]. Read in the\n",
    "other direction:\n",
    "shrink the step while adding blocks, $\\mathbf{x}_{l+1} = \\mathbf{x}_l +\n",
    "\\tfrac{1}{N}\\mathbf{f}_\\theta(\\mathbf{x}_l)$, and by the Euler convergence\n",
    "proposition (for smooth activations, so that the field is $C^1$) the\n",
    "network's output approaches the *exact* flow map\n",
    "$\\Phi_1(\\mathbf{x}_0)$ of the field; the error of the \"infinitely deep\"\n",
    "limit is the $O(h)$ of that section with $h = 1/N$.\n",
    "That limit object is the **Neural ODE**\n",
    "[@Chen.Rubanova.Bettencourt.ea.2018]:\n",
    "\n",
    "$$\n",
    "\\dot{\\mathbf{x}} = \\mathbf{f}_\\theta(\\mathbf{x}, t),\n",
    "\\qquad\n",
    "\\textrm{output} = \\mathbf{x}(T) = \\mathbf{x}_0 + \\int_0^T \\mathbf{f}_\\theta(\\mathbf{x}(t), t)\\,dt .\n",
    "$$\n",
    "\n",
    "Layers became integration time; the architecture became a single learned\n",
    "vector field plus a choice of solver; and everything this section proved now\n",
    "applies to the model itself: a Lipschitz $\\mathbf{f}_\\theta$ makes the\n",
    "input--output map well-posed and invertible\n",
    "(that section), and a fancier solver (RK4, an\n",
    "adaptive method) is a drop-in upgrade that changes compute, not parameters.\n",
    "\n",
    "### Training Through the Solver\n",
    "\n",
    "How do you fit $\\theta$? The direct route is **discretize, then\n",
    "differentiate**: unroll a fixed solver, say $N$ Euler steps, into the\n",
    "computation graph, compute the loss on $\\mathbf{x}_N$, and let ordinary\n",
    "backpropagation flow through the unrolled steps. The \"network\" this builds is\n",
    "literally a ResNet with $N$ blocks that share the same weights $\\theta$, so\n",
    "ordinary automatic differentiation already trains it.\n",
    "\n",
    "Let us do exactly that, on a task small enough to\n",
    "watch: learn a planar field $\\mathbf{f}_\\theta$ (one hidden layer, $32$ tanh\n",
    "units) whose time-$1$ flow carries the unit circle onto a shifted, squashed\n",
    "ellipse. We use $64$ paired points, mean-squared loss on the endpoints, $10$\n",
    "unrolled Euler steps of size $h = 0.1$, and Adam\n",
    "(that section) as the optimizer. The flow\n",
    "deforms *all* of\n",
    "$\\mathbb{R}^2$ smoothly and invertibly; we supervise it only at the $64$\n",
    "points."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "odes-solvers-neural-ode-train"
   },
   "source": [
    "theta = np.linspace(0, 2 * np.pi, 64, endpoint=False)\n",
    "X0_ring = np.stack([np.cos(theta), np.sin(theta)], axis=1)\n",
    "X1_ring = np.stack([1.5 + 1.2 * np.cos(theta), 0.4 * np.sin(theta)], axis=1)\n",
    "torch.manual_seed(0)\n",
    "X, Y = torch.tensor(X0_ring).float(), torch.tensor(X1_ring).float()\n",
    "net = torch.nn.Sequential(torch.nn.Linear(2, 32), torch.nn.Tanh(),\n",
    "                          torch.nn.Linear(32, 2))\n",
    "N, h = 10, 0.1                                 # unrolled Euler: 10 steps to T = 1\n",
    "\n",
    "def flow(X):\n",
    "    for _ in range(N):                         # one Euler step = one residual block\n",
    "        X = X + h * net(X)\n",
    "    return X\n",
    "\n",
    "opt = torch.optim.Adam(net.parameters(), lr=0.05)\n",
    "for it in range(401):\n",
    "    loss = ((flow(X) - Y) ** 2).mean()\n",
    "    opt.zero_grad()\n",
    "    loss.backward()\n",
    "    opt.step()\n",
    "    if it % 100 == 0:\n",
    "        print(f'iter {it:3d}  loss {loss.item():.5f}')\n",
    "print(f'mean endpoint error: {(flow(X) - Y).norm(dim=1).mean().item():.4f}')"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "iter   0  loss 1.16972\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "iter 100  loss 0.00004\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "iter 200  loss 0.00000\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "iter 300  loss 0.00000\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "iter 400  loss 0.00000\nmean endpoint error: 0.0011\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-odes-solvers-md-7 fw=all -->\n",
    "\n",
    "A few hundred full-batch Adam iterations drive the loss to $\\sim 10^{-5}$ and\n",
    "the mean endpoint error to $\\sim 10^{-3}$: one tiny vector\n",
    "field, integrated by ten shared-weight \"residual blocks,\" learns to carry a\n",
    "circle onto a displaced ellipse. And because the model is a flow, you get for\n",
    "free what no plain MLP gives you: integrate the *same* field backward\n",
    "(negate it, or run time from $1$ to $0$) and the ellipse returns to the\n",
    "circle: invertibility by construction, courtesy of Picard--Lindelöf.\n",
    "\n",
    "### The Adjoint Method: Backpropagation in Continuous Time\n",
    "\n",
    "Differentiating through the unrolled solver works, but it stores every\n",
    "intermediate state; for an adaptive solver taking thousands of internal\n",
    "steps, that is a lot of tape. The **adjoint method** computes the same\n",
    "gradients by integrating a *second* ODE backward in time, storing essentially\n",
    "nothing. It is the continuous-time limit of backpropagation, and its central\n",
    "object, the **adjoint state**\n",
    "$\\mathbf{a}(t) = \\partial L / \\partial \\mathbf{x}(t)$ (\"how would the loss\n",
    "change if the trajectory were nudged at time $t$?\"), is the continuous\n",
    "analogue of the backprop delta of\n",
    "that section.\n",
    "\n",
    "**Proposition (the adjoint equations).** *Let $\\mathbf{f}$ be continuously\n",
    "differentiable in $\\mathbf{x}$ and $\\theta$, let $\\mathbf{x}(t)$ solve\n",
    "$\\dot{\\mathbf{x}} = \\mathbf{f}(\\mathbf{x}, \\theta, t)$ on $[0, T]$ from fixed\n",
    "$\\mathbf{x}(0)$, and let $L = \\ell(\\mathbf{x}(T))$ be a differentiable loss\n",
    "on the endpoint.\n",
    "Then the adjoint state satisfies the linear ODE*\n",
    "\n",
    "$$\n",
    "\\dot{\\mathbf{a}}(t) = -\\left(\\frac{\\partial \\mathbf{f}}{\\partial \\mathbf{x}}\\right)^{\\!\\top}\\! \\mathbf{a}(t),\n",
    "\\qquad\n",
    "\\mathbf{a}(T) = \\nabla \\ell(\\mathbf{x}(T)),\n",
    "$$\n",
    "\n",
    "*integrated backward from $T$, and the parameter gradient is the integral*\n",
    "\n",
    "$$\n",
    "\\frac{\\partial L}{\\partial \\theta}\n",
    "= \\int_0^T \\left(\\frac{\\partial \\mathbf{f}}{\\partial \\theta}\\right)^{\\!\\top}\\! \\mathbf{a}(t)\\, dt,\n",
    "$$\n",
    "\n",
    "*with both partial derivatives evaluated along the trajectory\n",
    "$(\\mathbf{x}(t), \\theta, t)$* [@Chen.Rubanova.Bettencourt.ea.2018].\n",
    "\n",
    "**Proof (variational).** Perturb the parameters by an infinitesimal\n",
    "$\\delta\\theta$ and ask how the trajectory responds. Because $\\mathbf{f}$ is\n",
    "$C^1$, the flow is differentiable in $\\theta$ and the first-order expansion\n",
    "below is rigorous: the perturbation\n",
    "$\\delta\\mathbf{x}(t)$ obeys the *linearization* of the dynamics (differentiate\n",
    "the ODE in $\\theta$ and use the chain rule):\n",
    "\n",
    "$$\n",
    "\\delta\\dot{\\mathbf{x}}\n",
    "= \\frac{\\partial \\mathbf{f}}{\\partial \\mathbf{x}}\\,\\delta\\mathbf{x}\n",
    "+ \\frac{\\partial \\mathbf{f}}{\\partial \\theta}\\,\\delta\\theta,\n",
    "\\qquad\n",
    "\\delta\\mathbf{x}(0) = \\mathbf{0}.\n",
    "$$\n",
    "\n",
    "Now let $\\mathbf{a}(t)$ be defined by the backward ODE\n",
    "the equation and watch the pairing\n",
    "$\\mathbf{a}^\\top \\delta\\mathbf{x}$ evolve; the derivative telescopes by\n",
    "design:\n",
    "\n",
    "$$\n",
    "\\frac{d}{dt}\\left(\\mathbf{a}^\\top \\delta\\mathbf{x}\\right)\n",
    "= \\dot{\\mathbf{a}}^\\top \\delta\\mathbf{x} + \\mathbf{a}^\\top \\delta\\dot{\\mathbf{x}}\n",
    "= -\\mathbf{a}^\\top \\frac{\\partial \\mathbf{f}}{\\partial \\mathbf{x}}\\,\\delta\\mathbf{x}\n",
    "+ \\mathbf{a}^\\top \\frac{\\partial \\mathbf{f}}{\\partial \\mathbf{x}}\\,\\delta\\mathbf{x}\n",
    "+ \\mathbf{a}^\\top \\frac{\\partial \\mathbf{f}}{\\partial \\theta}\\,\\delta\\theta\n",
    "= \\mathbf{a}^\\top \\frac{\\partial \\mathbf{f}}{\\partial \\theta}\\,\\delta\\theta .\n",
    "$$\n",
    "\n",
    "The state-coupling terms cancel exactly; that cancellation is *why* the\n",
    "adjoint ODE has the form it has. Integrate from $0$ to $T$, using\n",
    "$\\delta\\mathbf{x}(0) = \\mathbf{0}$ on the left endpoint:\n",
    "\n",
    "$$\n",
    "\\mathbf{a}(T)^\\top \\delta\\mathbf{x}(T)\n",
    "= \\int_0^T \\mathbf{a}(t)^\\top \\frac{\\partial \\mathbf{f}}{\\partial \\theta}\\, dt \\;\\delta\\theta .\n",
    "$$\n",
    "\n",
    "The left side is $\\nabla\\ell(\\mathbf{x}(T))^\\top \\delta\\mathbf{x}(T) = \\delta L$,\n",
    "the first-order change of the loss; the right side is linear in\n",
    "$\\delta\\theta$ with coefficient the equation. (The same\n",
    "telescoping, stopped at an intermediate time $s$, shows\n",
    "$\\delta L = \\mathbf{a}(s)^\\top \\delta\\mathbf{x}(s)$ for a nudge injected at\n",
    "time $s$, so the solution of the equation really is\n",
    "$\\partial L/\\partial \\mathbf{x}(s)$, earning its name.) $\\blacksquare$\n",
    "\n",
    "One caveat connects back to that section: the\n",
    "existence theorem tolerates ReLU fields, but the adjoint equations as stated\n",
    "need a $C^1$ field. A ReLU field is differentiable only piecewise; in\n",
    "practice autograd differentiates the unrolled solver piece by piece, and\n",
    "smooth activations (tanh, GELU) remove the issue entirely.\n",
    "\n",
    "Look at what the backward ODE *computes per step*:\n",
    "$-\\mathbf{a}^\\top \\partial\\mathbf{f}/\\partial\\mathbf{x}$ is a\n",
    "**vector--Jacobian product**, the very primitive that reverse-mode AD is made\n",
    "of (that section). Discretize\n",
    "the equation with the same Euler scheme as the forward pass\n",
    "and you get *literally* the backpropagation recursion through the unrolled\n",
    "network: backprop **is** the discrete adjoint method, a lineage that runs\n",
    "from optimal control [@Pontryagin.Boltyanskii.Gamkrelidze.ea.1962]\n",
    "straight to the backward pass of every deep-learning framework. The continuous formulation adds one practical\n",
    "twist: instead of storing the forward states for the VJPs, you may\n",
    "*re-integrate* $\\mathbf{x}(t)$ backward alongside $\\mathbf{a}(t)$, making\n",
    "memory $O(1)$ in the number of solver steps, at the cost of extra compute\n",
    "and of numerical drift when the reversed dynamics are unstable (a strongly\n",
    "contracting forward flow is, run backward, strongly expanding; in that regime\n",
    "checkpointing or plain unrolling is the sturdier choice). The literature\n",
    "names the two routes **discretize-then-optimize** (differentiate the\n",
    "unrolled solver, exact for the program you actually ran) versus\n",
    "**optimize-then-discretize** (discretize the continuous adjoint,\n",
    "memory-free but only approximately the gradient of anything)\n",
    "[@Kidger.2022]; Exercise 7 probes exactly this fault line.\n",
    "\n",
    "Everything above is checkable on the linear ODE\n",
    "$\\dot{\\mathbf{x}} = A\\mathbf{x}$, where every ingredient has a closed form:\n",
    "the trajectory is $e^{At}\\mathbf{x}_0$, the adjoint of\n",
    "$L = \\tfrac12\\|\\mathbf{x}(T)\\|^2$ is\n",
    "$\\mathbf{a}(t) = e^{A^\\top (T-t)}\\,\\mathbf{x}(T)$ (the adjoint ODE\n",
    "$\\dot{\\mathbf{a}} = -A^\\top\\mathbf{a}$ run backward), and with $\\theta$ the\n",
    "entries of $A$ itself, $\\partial\\mathbf{f}/\\partial A$ turns\n",
    "the equation into\n",
    "$\\partial L/\\partial A = \\int_0^T \\mathbf{a}(t)\\,\\mathbf{x}(t)^\\top dt$. The\n",
    "cell computes the gradient three ways: hand-written reverse mode through the\n",
    "unrolled Euler solver (the discrete adjoint: exactly what reverse-mode\n",
    "autograd builds, written out in six lines), central finite differences on the\n",
    "same unrolled program, and the continuous adjoint integral by quadrature."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "odes-solvers-adjoint-check"
   },
   "source": [
    "def expm_eig(B, s):\n",
    "    \"\"\"e^{Bs} via the eigendecomposition formula.\"\"\"\n",
    "    mu, W = np.linalg.eig(B)\n",
    "    return (W @ np.diag(np.exp(mu * s)) @ np.linalg.inv(W)).real\n",
    "\n",
    "def loss_and_grad_unrolled(A, x0, T, n):\n",
    "    \"\"\"L = ||x_n||^2 / 2 through n unrolled Euler steps; dL/dA by hand-written\n",
    "    reverse mode -- the discrete adjoint.\"\"\"\n",
    "    h, xs = T / n, [x0]\n",
    "    for k in range(n):\n",
    "        xs.append(xs[-1] + h * (A @ xs[-1]))   # forward sweep, tape the states\n",
    "    a, gA = xs[-1].copy(), np.zeros_like(A)    # a_n = dL/dx_n = x_n\n",
    "    for k in range(n - 1, -1, -1):             # backward sweep: one VJP per step\n",
    "        gA += h * np.outer(a, xs[k])           # dL/dA  +=  h a_{k+1} x_k^T\n",
    "        a = a + h * (A.T @ a)                  # a_k = (I + hA)^T a_{k+1}\n",
    "    return 0.5 * xs[-1] @ xs[-1], gA\n",
    "\n",
    "T, n = 2.0, 1000\n",
    "L, g_disc = loss_and_grad_unrolled(A, x0, T, n)\n",
    "g_fd = np.zeros_like(A)                        # central finite differences\n",
    "for i in range(2):\n",
    "    for j in range(2):\n",
    "        E = np.zeros_like(A); E[i, j] = 1e-5\n",
    "        g_fd[i, j] = (loss_and_grad_unrolled(A + E, x0, T, n)[0]\n",
    "                      - loss_and_grad_unrolled(A - E, x0, T, n)[0]) / 2e-5\n",
    "print('discrete adjoint vs finite differences:', f'{np.abs(g_disc - g_fd).max():.1e}')\n",
    "\n",
    "m = 2000                                       # continuous adjoint, trapezoid rule\n",
    "ts = np.linspace(0.0, T, m + 1)\n",
    "xT = expm_eig(A, T) @ x0\n",
    "F = np.stack([np.outer(expm_eig(A.T, T - s) @ xT, expm_eig(A, s) @ x0)\n",
    "              for s in ts])\n",
    "g_cont = (T / m) * (F[1:-1].sum(axis=0) + 0.5 * (F[0] + F[-1]))\n",
    "for n_steps in [100, 1000, 10000]:\n",
    "    gap = np.abs(loss_and_grad_unrolled(A, x0, T, n_steps)[1] - g_cont).max()\n",
    "    print(f'n = {n_steps:6d} Euler steps: |discrete - continuous adjoint| = {gap:.1e}')"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "discrete adjoint vs finite differences: 1.2e-10\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "n =    100 Euler steps: |discrete - continuous adjoint| = 2.9e-02\nn =   1000 Euler steps: |discrete - continuous adjoint| = 2.8e-03\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "n =  10000 Euler steps: |discrete - continuous adjoint| = 2.8e-04\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-odes-solvers-md-8 fw=all -->\n",
    "\n",
    "The discrete adjoint matches finite differences to $10^{-10}$ (it *is* the\n",
    "exact gradient of the unrolled program), and as the solver is refined the\n",
    "discrete gradient converges to the continuous adjoint integral at exactly the\n",
    "solver's order, $O(h)$: ten times the steps, one-tenth the gap. Backprop\n",
    "through a solver and the adjoint method are one algorithm, seen\n",
    "at two resolutions.\n",
    "\n",
    "## Continuous Normalizing Flows\n",
    "\n",
    "### The Instantaneous Change of Variables\n",
    "\n",
    "Push *samples* through a Neural ODE and you have a generative model; to train\n",
    "it by maximum likelihood you must know how the *density* changes along the\n",
    "flow. The discrete answer we know from that section:\n",
    "for an invertible map $\\mathbf{g}$ with Jacobian $J_{\\mathbf{g}}$,\n",
    "\n",
    "$$\n",
    "\\log p_{\\textrm{out}}(\\mathbf{g}(\\mathbf{x}))\n",
    "= \\log p_{\\textrm{in}}(\\mathbf{x}) - \\log \\left|\\det J_{\\mathbf{g}}(\\mathbf{x})\\right| ,\n",
    "$$\n",
    "\n",
    "and a normalizing flow built from $K$ layers pays one log-det-Jacobian per\n",
    "layer; the $O(d^3)$ determinant is the reason discrete flows constrain\n",
    "their layers to triangular or low-rank Jacobians. In continuous time the\n",
    "formula *simplifies*: over one infinitesimal step the flow map is\n",
    "near-identity, and the determinant of a near-identity matrix is governed by\n",
    "the trace.\n",
    "\n",
    "**Lemma.** *For any square matrix $J$,\n",
    "$\\det(I + hJ) = 1 + h \\operatorname{tr} J + O(h^2)$.*\n",
    "\n",
    "**Proof.** In the Leibniz expansion of the determinant, any permutation other\n",
    "than the identity uses at least two off-diagonal entries, each of size\n",
    "$O(h)$, so contributes $O(h^2)$. The identity permutation contributes\n",
    "$\\prod_i (1 + h J_{ii}) = 1 + h \\sum_i J_{ii} + O(h^2)$. $\\blacksquare$\n",
    "\n",
    "**Proposition (instantaneous change of variables).** *Let\n",
    "$\\dot{\\mathbf{x}} = \\mathbf{f}(\\mathbf{x}, t)$ with $\\mathbf{f}$ continuously\n",
    "differentiable and $\\mathbf{x}(0) \\sim p_0$ for a differentiable density\n",
    "$p_0$, and let $p_t$ denote the density of $\\mathbf{x}(t)$, which then exists\n",
    "and is differentiable along trajectories. Then along each\n",
    "trajectory* [@Chen.Rubanova.Bettencourt.ea.2018]\n",
    "\n",
    "$$\n",
    "\\frac{d}{dt} \\log p_t(\\mathbf{x}(t))\n",
    "= -\\operatorname{tr}\\!\\left(\\frac{\\partial \\mathbf{f}}{\\partial \\mathbf{x}}\\right)\\!(\\mathbf{x}(t), t).\n",
    "$$\n",
    "\n",
    "**Proof.** Over a short interval the flow advances by the near-identity map\n",
    "$\\Phi_h(\\mathbf{x}) = \\mathbf{x} + h\\,\\mathbf{f}(\\mathbf{x}, t) + O(h^2)$,\n",
    "whose Jacobian is $I + h J + O(h^2)$ with\n",
    "$J = \\partial\\mathbf{f}/\\partial\\mathbf{x}$. Apply the discrete\n",
    "change-of-variables formula to this one map and expand with the Lemma, which\n",
    "tolerates the perturbed argument: replacing $J$ by $J + O(h)$ changes\n",
    "$\\det(I + hJ)$ only at order $h^2$. Thus:\n",
    "\n",
    "$$\n",
    "\\log p_{t+h}(\\mathbf{x}(t+h))\n",
    "= \\log p_t(\\mathbf{x}(t)) - \\log\\det\\!\\left(I + hJ + O(h^2)\\right)\n",
    "= \\log p_t(\\mathbf{x}(t)) - h \\operatorname{tr} J + O(h^2),\n",
    "$$\n",
    "\n",
    "using $\\log(1 + u) = u + O(u^2)$ (for small $h$ the determinant is positive,\n",
    "so the absolute value is moot). Subtract, divide by $h$, and let\n",
    "$h \\to 0$. $\\blacksquare$\n",
    "\n",
    "The trace of the Jacobian is the **divergence** of the field: the local\n",
    "expansion rate of volume. Where the field spreads ($\\operatorname{tr} J > 0$),\n",
    "the cloud of samples dilutes and the log-density along each trajectory falls;\n",
    "where it compresses, probability concentrates. Integrating\n",
    "the equation along the trajectory turns the per-layer\n",
    "log-det *sum* of a discrete flow into a *trace integral*:\n",
    "\n",
    "$$\n",
    "\\log p_T(\\mathbf{x}(T))\n",
    "= \\log p_0(\\mathbf{x}(0)) - \\int_0^T \\operatorname{tr}\\!\\left(\\frac{\\partial \\mathbf{f}}{\\partial \\mathbf{x}}\\right) dt .\n",
    "$$\n",
    "\n",
    "This is the engine of the **continuous normalizing flow** (CNF): augment the\n",
    "state with a running log-density, integrate $(\\mathbf{x}, \\log p)$ together\n",
    "with one ODE solver call, and you have an *exact* likelihood for an\n",
    "unconstrained architecture, with no triangular Jacobians required.\n",
    "\n",
    "### The Hutchinson Trace Estimator\n",
    "\n",
    "One cost remains: $\\operatorname{tr}(\\partial\\mathbf{f}/\\partial\\mathbf{x})$\n",
    "naively needs all $d$ diagonal entries of the Jacobian, i.e. $d$ derivative\n",
    "passes. The fix is a one-line identity from randomized numerical linear\n",
    "algebra [@Hutchinson.1989].\n",
    "\n",
    "**Proposition (Hutchinson estimator).** *For any matrix\n",
    "$M \\in \\mathbb{R}^{d \\times d}$ and any random\n",
    "$\\boldsymbol{\\epsilon} \\in \\mathbb{R}^d$ with\n",
    "$\\mathbb{E}[\\boldsymbol{\\epsilon}] = \\mathbf{0}$ and\n",
    "$\\mathbb{E}[\\boldsymbol{\\epsilon}\\boldsymbol{\\epsilon}^\\top] = I$ (Gaussian\n",
    "or Rademacher, i.e. uniform $\\pm 1$ entries),*\n",
    "\n",
    "$$\n",
    "\\operatorname{tr}(M) = \\mathbb{E}\\!\\left[\\boldsymbol{\\epsilon}^\\top M \\boldsymbol{\\epsilon}\\right].\n",
    "$$\n",
    "\n",
    "**Proof.** $\\mathbb{E}[\\boldsymbol{\\epsilon}^\\top M \\boldsymbol{\\epsilon}]\n",
    "= \\mathbb{E}[\\operatorname{tr}(M \\boldsymbol{\\epsilon}\\boldsymbol{\\epsilon}^\\top)]\n",
    "= \\operatorname{tr}(M\\, \\mathbb{E}[\\boldsymbol{\\epsilon}\\boldsymbol{\\epsilon}^\\top])\n",
    "= \\operatorname{tr}(M)$, using the cyclic property of the trace and its\n",
    "linearity (which lets it commute with the expectation). $\\blacksquare$\n",
    "\n",
    "The point is *what the estimator touches*: $\\boldsymbol{\\epsilon}^\\top J$ is\n",
    "a single vector--Jacobian product, one reverse-mode pass through\n",
    "$\\mathbf{f}$ with cost comparable to one evaluation/backward pass, and no\n",
    "full Jacobian is ever materialized\n",
    "(that section), followed by a dot product.\n",
    "An unbiased stochastic log-likelihood at the price of one extra backward\n",
    "pass is what makes CNFs scale; this is precisely the FFJORD recipe\n",
    "[@Grathwohl.Chen.Bettencourt.ea.2018].\n",
    "\n",
    "For a *linear* field everything is checkable by hand:\n",
    "$\\dot{\\mathbf{x}} = A\\mathbf{x}$ has constant Jacobian $A$, so\n",
    "the equation says the log-density along every\n",
    "trajectory falls at the constant rate $\\operatorname{tr}(A)$:\n",
    "$\\log p_t(\\mathbf{x}(t)) = \\log p_0(\\mathbf{x}_0) - t \\operatorname{tr}(A)$.\n",
    "The cell flows standard-normal samples through the spiral field, where the\n",
    "time-$t$ density is the Gaussian\n",
    "$\\mathcal{N}(\\mathbf{0},\\, e^{At} e^{A^\\top t})$, computable in closed form,\n",
    "and compares; then it verifies the Hutchinson estimator on a random\n",
    "$6 \\times 6$ matrix."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "odes-solvers-cnf-trace"
   },
   "source": [
    "rng = np.random.default_rng(0)\n",
    "t = 1.5\n",
    "X0s = rng.standard_normal((4, 2))              # four samples from p_0 = N(0, I)\n",
    "Et = expm_eig(A, t)\n",
    "Xts = X0s @ Et.T                               # flow each sample to time t\n",
    "Sig = Et @ Et.T                                # x(t) ~ N(0, E_t E_t^T)\n",
    "logp0 = -0.5 * (X0s ** 2).sum(axis=1) - np.log(2 * np.pi)\n",
    "logpt = (-0.5 * np.einsum('ki,ij,kj->k', Xts, np.linalg.inv(Sig), Xts)\n",
    "         - 0.5 * np.log(np.linalg.det(Sig)) - np.log(2 * np.pi))\n",
    "print('log p_t(x(t)), Gaussian formula  :', logpt.round(6))\n",
    "print('log p_0(x_0) - t tr(A), our rule :', (logp0 - t * np.trace(A)).round(6))\n",
    "\n",
    "J = rng.standard_normal((6, 6))                # Hutchinson on a bigger matrix\n",
    "eps = rng.choice([-1.0, 1.0], size=(100000, 6))\n",
    "est = np.einsum('ki,ij,kj->k', eps, J, eps)\n",
    "print(f'Hutchinson: tr(J) = {np.trace(J):.4f},  estimate = {est.mean():.4f}'\n",
    "      f' +- {est.std() / np.sqrt(len(est)):.4f}')"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "log p_t(x(t)), Gaussian formula  : [-0.354507 -0.54845  -0.546723 -1.636566]\nlog p_0(x_0) - t tr(A), our rule : [-0.354507 -0.54845  -0.546723 -1.636566]\nHutchinson: tr(J) = -1.7535,  estimate = -1.7554 +- 0.0173\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-odes-solvers-md-9 fw=all -->\n",
    "\n",
    "The two log-density computations agree digit for digit (the trace integral\n",
    "*is* the log-det-Jacobian, evaluated the cheap way), and the Hutchinson\n",
    "estimate brackets the true trace within its standard error. Finally,\n",
    "the equation is the\n",
    "trajectory-wise form of the *continuity equation* governing how a whole\n",
    "density field is transported;\n",
    "that section starts there and turns\n",
    "a diffusion's noisy paths into a deterministic probability-flow ODE\n",
    "integrated by exactly the solvers of this section.\n",
    "\n",
    "## Summary\n",
    "\n",
    "* An ODE $\\dot{\\mathbf{x}} = \\mathbf{f}(\\mathbf{x}, t)$ is a velocity field;\n",
    "  a solution is everywhere tangent to it, and the flow map $\\Phi_t$ moves all\n",
    "  of space at once, composing as $\\Phi_{t+s} = \\Phi_t \\circ \\Phi_s$.\n",
    "* **Picard--Lindelöf**: a globally Lipschitz field has exactly one trajectory\n",
    "  through each point for every finite time (the Picard integral operator is a\n",
    "  contraction). Existence and uniqueness in both time directions make its\n",
    "  flow maps bijections, the well-posedness behind invertible generative flows.\n",
    "  Without Lipschitz, uniqueness fails ($\\dot{x} = \\sqrt{|x|}$, a fan of\n",
    "  solutions); without a growth bound, solutions can blow up in finite time\n",
    "  ($\\dot{x} = x^2$).\n",
    "* Linear systems are solved by the **matrix exponential**\n",
    "  $e^{At} = \\sum_k (At)^k / k!\\, = V e^{\\Lambda t} V^{-1}$: independent modes\n",
    "  $e^{\\lambda_i t}$ along eigendirections. **Stability dictionary**: real\n",
    "  parts decide decay vs. growth, imaginary parts rotation; nonlinear fixed\n",
    "  points inherit the verdict from the Jacobian's eigenvalues.\n",
    "* Explicit solvers march with global error $O(h^p)$: forward **Euler** has\n",
    "  order $1$ (local $O(h^2)$ errors, $T/h$ of them, times an $e^{LT}$\n",
    "  stability factor), **RK4** order $4$; both slopes are measurable on a\n",
    "  log--log plot. **Stiffness**: forward Euler is stable on\n",
    "  $\\dot{x} = -\\lambda x$ only for $h < 2/\\lambda$, so a fast dead mode can\n",
    "  dictate the step; **implicit (backward) Euler** is stable for every $h$, at\n",
    "  the price of solving an equation per step.\n",
    "* **Gradient descent is forward Euler on the gradient flow**\n",
    "  $\\dot{\\mathbf{x}} = -\\nabla L$: the solver stability bound becomes the\n",
    "  learning-rate rule $\\eta < 2/\\lambda_{\\max}(H)$, ill-conditioning is\n",
    "  stiffness, and momentum discretizes the heavy-ball ODE\n",
    "  $m\\ddot{\\mathbf{x}} + \\gamma\\dot{\\mathbf{x}} = -\\nabla L$.\n",
    "* A **residual block is one Euler step**, a ResNet is a solver, and the\n",
    "  continuous limit is a **Neural ODE**: a learned vector field whose flow is\n",
    "  fit by differentiating through the solver. The **adjoint method**\n",
    "  integrates $\\dot{\\mathbf{a}} = -(\\partial\\mathbf{f}/\\partial\\mathbf{x})^\\top\\mathbf{a}$\n",
    "  backward (each step a VJP) and is backpropagation in continuous time,\n",
    "  with $O(1)$-memory and numerical-drift tradeoffs.\n",
    "* Densities transported by a flow obey the **instantaneous change of\n",
    "  variables** $\\tfrac{d}{dt}\\log p_t = -\\operatorname{tr}(\\partial\\mathbf{f}/\\partial\\mathbf{x})$:\n",
    "  the log-det-Jacobian becomes a trace integral, estimated unbiasedly by\n",
    "  **Hutchinson's** $\\boldsymbol{\\epsilon}^\\top J \\boldsymbol{\\epsilon}$ trick\n",
    "  at the cost of one VJP: the mathematics of continuous normalizing flows.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Verify by differentiation that $\\mathbf{x}(t) = e^{-t}\\mathbf{x}_0$ solves\n",
    "   $\\dot{\\mathbf{x}} = -\\mathbf{x}$, and that trajectories of the rotational\n",
    "   field $\\dot{\\mathbf{x}} = \\left(\\begin{smallmatrix}0&-1\\\\1&0\\end{smallmatrix}\\right)\\mathbf{x}$\n",
    "   keep $\\|\\mathbf{x}(t)\\|$ constant. Write the integral form\n",
    "   the equation for $\\dot{x} = t$ and evaluate it.\n",
    "2. Show that a linear field $\\mathbf{f}(\\mathbf{x}) = A\\mathbf{x}$ is\n",
    "   Lipschitz with constant $\\|A\\|$ (the operator norm). Then construct a\n",
    "   solution of $\\dot{x} = \\sqrt{|x|}$, $x(0) = 0$ that waits at the origin\n",
    "   until time $c$ before leaving, verify it solves the ODE at every $t$, and\n",
    "   identify exactly where the Lipschitz hypothesis fails. Finally, compute\n",
    "   the blow-up time of $\\dot{x} = x^2$, $x(0) = x_0 > 0$, and explain why\n",
    "   Picard--Lindelöf only promises a *local* solution here.\n",
    "3. From the series the equation, show that $e^{At}$\n",
    "   commutes with $A$ and that $e^{A(t+s)} = e^{At}e^{As}$. Then classify the\n",
    "   phase portraits of\n",
    "   $\\left(\\begin{smallmatrix}-2&0\\\\0&-1\\end{smallmatrix}\\right)$,\n",
    "   $\\left(\\begin{smallmatrix}1&0\\\\0&-1\\end{smallmatrix}\\right)$, and\n",
    "   $\\left(\\begin{smallmatrix}0&-2\\\\2&0\\end{smallmatrix}\\right)$ using the\n",
    "   stability dictionary.\n",
    "4. Derive Euler's local truncation error $O(h^2)$ from Taylor's theorem with\n",
    "   remainder, and rehearse the accumulation argument that turns it into a\n",
    "   global $O(h)$. Then show that the midpoint method's update matches the\n",
    "   true Taylor expansion through $h^2$ in the vector case, where the scalar\n",
    "   product $f'f$ becomes the Jacobian--vector product\n",
    "   $(\\partial\\mathbf{f}/\\partial\\mathbf{x})\\mathbf{f}$.\n",
    "5. Re-derive the forward-Euler stability bound $h < 2/\\lambda$ and backward\n",
    "   Euler's unconditional stability. For the stiff matrix\n",
    "   $A = \\mathrm{diag}(-100, -1)$: what step size does forward Euler need for\n",
    "   stability, how many steps to reach $T = 5$, and how do both change if the\n",
    "   fast eigenvalue moves to $-10^4$? For nonlinear $\\mathbf{f}$, write the\n",
    "   equation one backward-Euler step must solve and the Newton iteration you\n",
    "   would use.\n",
    "6. A residual block computes\n",
    "   $\\mathbf{x}_{l+1} = \\mathbf{x}_l + \\mathbf{f}_\\theta(\\mathbf{x}_l)$.\n",
    "   Identify the implied step size; explain what halving the solver step while\n",
    "   doubling the step count corresponds to architecturally; and use\n",
    "   that section to argue that the time-$T$ flow\n",
    "   of a Lipschitz $\\mathbf{f}_\\theta$ is invertible; then explain why a\n",
    "   *plain* (non-residual) layer $\\mathbf{x} \\mapsto \\sigma(W\\mathbf{x})$ has\n",
    "   no such guarantee.\n",
    "7. Re-derive the adjoint equations the equation and\n",
    "   the equation from the variational argument without\n",
    "   looking. Show that one Euler step of the adjoint ODE is exactly one\n",
    "   vector--Jacobian product, and that discretizing the adjoint ODE with\n",
    "   Euler reproduces the backpropagation recursion of the\n",
    "   `#odes-solvers-adjoint-check` cell. When the forward dynamics contract\n",
    "   strongly (all $\\operatorname{Re}\\lambda \\ll 0$), what goes wrong with\n",
    "   reconstructing $\\mathbf{x}(t)$ by integrating backward, and why does the\n",
    "   $O(1)$-memory adjoint suffer where unrolled backprop does not?\n",
    "8. Derive the instantaneous change of variables\n",
    "   the equation from the discrete log-det formula as\n",
    "   $h \\to 0$, and check that for a single linear \"layer\" $\\Phi_t = e^{At}$ it\n",
    "   reproduces $-\\log|\\det e^{At}| = -t\\operatorname{tr}(A)$ (use\n",
    "   $\\det e^{At} = e^{t \\operatorname{tr} A}$, e.g. via\n",
    "   the equation). Prove the Hutchinson estimator\n",
    "   the equation is unbiased, compute its variance for\n",
    "   Rademacher probes (express it in terms of the off-diagonal entries of\n",
    "   $M + M^\\top$), and explain why estimating the *trace* costs one VJP while\n",
    "   the *determinant* has no comparably cheap unbiased estimator."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-odes-solvers-md-10 fw=pytorch -->\n",
    "\n",
    "[Discussions](https://d2l.discourse.group/)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-pytorch",
   "display_name": "d2l (PyTorch)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
