Start writing notes on differentiable simulation

_code/notes/diffsim/derivations.py

+from sympy import *
+from sympy.assumptions.assume import global_assumptions
+
+def print_expr(name, expr):
+    print(name + ": ")
+    print(latex(expr))
+    print("")
+
+t = Symbol('t', real=True)
+dt = Symbol('\Delta t', real=True)
+x = Symbol('x', real=True)
+xn = Function('x_n', real=True)
+f = Function('f', real=True)
+
+def euler():
+    expr = xn(x) + dt * f(t, xn(x))
+    return simplify(Derivative(expr, x))
+
+def rk4():
+    k1 = dt * f(t, xn(x))
+    k2 = dt * f(t + dt/2, xn(x) + k1/2)
+    k3 = dt * f(t + dt/2, xn(x) + k2/2)
+    k4 = dt * f(t + dt, xn(x) + k3)
+    expr = xn(x) + k1/6 + k2/3 + k3/3 + k4/6
+    return simplify(Derivative(expr, x))
+
+print_expr("euler", euler())
+print_expr("rk4", rk4())

_notes/differentiable_simulation.md

+---
+title: Differentiable Simulation
+---
+
+## ODEs
+
+Let's consider ODEs of the form $$y'(x) = f(x, y)$$. Assuming a fixed step size $$\Delta x$$, the
+most common integration schemes are basically functions $$g$$ such that
+
+$$
+y(x + \Delta x) \approx g(x, y)
+$$
+
+For instance, Euler integration uses $$g(x, y) = y + (\Delta x) f(x, y)$$.
+
+The function $$g$$ is applied recursively to produce a series of points, starting from some initial
+condition $$y_0$$.
+
+$$
+\begin{aligned}
+y_1 &= g(0, y_0) \\
+y_2 &= g(\Delta t, y_1) \\
+y_3 &= g(2 \Delta t, y_2) \\
+\vdots
+\end{aligned}
+$$
+
+We can write this more compactly as the recurrence relation
+
+$$
+y_{n+1} = g(n \Delta t, y_n)
+$$
+
+By differentiating both sides of this relation and applying the chain rule, we find that
+
+$$
+\frac{\partial}{\partial y_0} y_{n+1} = \frac{\partial}{\partial y_n} g(n \Delta t, y_n) \frac{\partial}{\partial y_0} y_n
+$$
+
+Assuming that we can compute $$\partial g / \partial y$$ on demand, this gives a convenient
+recurrence relation for $$\partial y_{n+1} / \partial y_0$$ in terms of $$\partial y_{n} / \partial
+y_0$$. So to differentiate through such a simulation, all we need to do is figure out $$\partial g /
+\partial y$$, and compute the series of derivate values $$\partial y_{n} / \partial y_0$$ alongside
+the values $$y_n$$.
+
+### Euler
+
+As mentioned above, for Euler integration $$g(x, y) = y + (\Delta x) f(x, y)$$. Thus
+
+$$
+\frac{\partial}{\partial y} g(x, y) = 1 + (\Delta x) \frac{\partial}{\partial y} f(x, y)
+$$
+
+And
+
+$$
+\begin{aligned}
+\frac{\partial}{\partial y_0} y_{n+1}
+&= \frac{\partial}{\partial y_n} g(n \Delta t, y_n) \frac{\partial}{\partial y_0} y_n \\
+&= \left( 1 + (\Delta x) \frac{\partial}{\partial y} f(x, y) \right) \frac{\partial}{\partial y_0} y_n \\
+\end{aligned}
+$$