Runge-Kutta method for Numbat

examples/tests/numerics.nbt

@@ -25,3 +25,31 @@ fn dist(t: Time) -> Length = 0.5 g0 t^2
 fn velocity(t: Time) -> Velocity = diff(dist, t)
 
 assert_eq(velocity(2.0 s), 2.0 s × g0, 1e-3 m/s)
+
+# Differential equations
+
+let t_min = 0 s
+let t_max = 2 s
+let n_points = 2_000
+let t_list = linspace(t_min, t_max, n_points)
+
+let μ = 0.7 / s
+let x0 = 2 m
+
+# Numerically solve x'(t) = μ x(t) with x(0) = x0.
+# The exact solution is x(t) = x0 e^(μ t).
+#
+# 'ode(t, x)' is the right-hand side of the ODE.
+fn ode(t, x) = μ x
+
+let x_list = dsolve_runge_kutta(ode, t_list, x0)
+
+fn numerical_solution(t) = element_at(idx, x_list)
+  where t_range = t_max - t_min
+    and idx = floor((t - t_min) / t_range * (n_points - 1))
+
+fn exact_solution(t) = x0 exp(μ t)
+
+assert_eq(numerical_solution(0 s), exact_solution(0 s))
+assert_eq(numerical_solution(1 s), exact_solution(1 s), 1e-3 x0)
+assert_eq(numerical_solution(2 s), exact_solution(2 s), 1e-3 x0)

numbat/modules/numerics/diff.nbt

@@ -1,4 +1,5 @@
 use core::quantities
+use core::lists
 
 @name("Numerical differentiation")
 @url("https://en.wikipedia.org/wiki/Numerical_differentiation")
@@ -7,3 +8,22 @@ fn diff<X: Dim, Y: Dim>(f: Fn[(X) -> Y], x: X) -> Y / X =
   (f(x + Δx) - f(x - Δx)) / 2 Δx
   where
     Δx = 1e-10 × unit_of(x)
+
+@name("Runga-Kutta method")
+@url("https://en.wikipedia.org/wiki/Runge-Kutta_methods")
+@description("Solve the ordinary differential equation $y' = f(x, y)$ with initial conditions $y(x_0) = y_0$ using the fourth-order Runge-Kutta method.")
+fn dsolve_runge_kutta<X: Dim, Y: Dim>(
+  f: Fn[(X, Y) -> Y / X],
+  xs: List<X>,
+  y0: Y) -> List<Y> =
+    if len(xs) == 2
+      then [y0, y1]
+      else cons(y0, dsolve_runge_kutta(f, tail(xs), y1))
+  where x0 = element_at(0, xs)
+    and x1 = element_at(1, xs)
+    and Δx = x1 - x0
+    and k1 = f(x0, y0)
+    and k2 = f(x0 + Δx / 2, y0 + Δx k1 / 2)
+    and k3 = f(x0 + Δx / 2, y0 + Δx k2 / 2)
+    and k4 = f(x0 + Δx, y0 + Δx k3)
+    and y1 = y0 + Δx / 6 × (k1 + 2 k2 + 2 k3 + k4)