Improve some documentation and make minor usability adjustments (#816)

* Add an error message for non-vector-valued initial conditions

* Improve the visuals of the ODE in the quickstart example

* Update the ruff-pre-commit hook

* Treat warnings as errors (and ditch Diffrax from tests because it raises warnings)

* Allow passing strategies to solvers as keyword arguments

Author: pnkraemer

.pre-commit-config.yaml

@@ -13,10 +13,10 @@ repos:
       - id: yamlfix
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.12.10
+    rev: v0.12.11
     hooks:
       # Run the linter.
-      - id: ruff
+      - id: ruff-check
         args: [--fix]
       # Run the formatter.
       - id: ruff-format

docs/examples_quickstart/easy_example.py

@@ -36,11 +36,11 @@
 @jax.jit
 def vf(y, *, t):  # noqa: ARG001
     """Evaluate the vector field."""
-    return 0.5 * y * (1 - y)
+    return 2.0 * y * (1 - y)
 
 
 u0 = jnp.asarray([0.1])
-t0, t1 = 0.0, 1.0
+t0, t1 = 0.0, 5.0
 # -
 
 # Configuring a probabilistic IVP solver is a little more

makefile

@@ -2,7 +2,7 @@ format-and-lint:
 	pre-commit run --all-files
 
 test:
-	pytest -n auto -v 		# parallelise, verbose output
+	pytest -n auto -v -Werror 		# parallelise, verbose output, warnings as errors
 
 quickstart:
 	# Run some code without installing any of the optional dependencies

probdiffeq/backend/ode.py

@@ -2,7 +2,6 @@
 
 import jax
 import jax.experimental.ode
-import jax.numpy as jnp
 
 
 def odeint_and_save_at(vf, y0: tuple, /, save_at, *, atol, rtol):
@@ -15,40 +14,6 @@ def vf_wrapped(y, t):
     return jax.experimental.ode.odeint(vf_wrapped, *y0, save_at, atol=atol, rtol=rtol)
 
 
-def odeint_dense(vf, y0: tuple, /, t0, t1, *, atol, rtol):
-    # Local import because diffrax is not an official dependency
-    import diffrax
-
-    assert isinstance(y0, tuple | list)
-    assert len(y0) == 1
-
-    @diffrax.ODETerm
-    @jax.jit
-    def vf_wrapped(t, y, _args):
-        return vf(y, t=t)
-
-    solution_object = diffrax.diffeqsolve(
-        vf_wrapped,
-        diffrax.Dopri5(),
-        t0=t0,
-        t1=t1,
-        dt0=0.1,
-        y0=y0[0],
-        saveat=diffrax.SaveAt(dense=True),
-        stepsize_controller=diffrax.PIDController(atol=atol, rtol=rtol),
-    )
-
-    def solution(t):
-        # Automatic batching
-        if jnp.ndim(t) > 0:
-            return jax.vmap(solution)(t)
-
-        # Interpolate
-        return solution_object.evaluate(t)
-
-    return solution
-
-
 def ivp_lotka_volterra():
     # Local imports because diffeqzoo is not an official dependency
     from diffeqzoo import backend, ivps

probdiffeq/impl/impl.py

@@ -1,6 +1,7 @@
 """State-space model implementations."""
 
 from probdiffeq.backend import containers, functools, tree_util
+from probdiffeq.backend import numpy as np
 from probdiffeq.impl import _conditional, _linearise, _normal, _prototypes, _stats
 
 
@@ -26,6 +27,12 @@ def __eq__(self, other):
 
 def choose(which: str, /, *, tcoeffs_like) -> FactImpl:
     """Choose a state-space model implementation."""
+    u0 = np.asarray(tcoeffs_like[0])
+    if u0.ndim != 1:
+        msg = "'tcoeffs' expected to have shape=(d,), "
+        msg += f"but shape={u0.shape} received."
+        raise ValueError(msg)
+
     if which == "dense":
         return _select_dense(tcoeffs_like=tcoeffs_like)
     if which == "isotropic":

probdiffeq/ivpsolvers.py

@@ -766,7 +766,7 @@ def initial_condition(self):
         return posterior, self.prior.output_scale
 
 
-def solver_mle(extrapolation, /, *, correction, prior, ssm):
+def solver_mle(strategy, *, correction, prior, ssm):
     """Create a solver that calibrates the output scale via maximum-likelihood.
 
     Warning: needs to be combined with a call to stats.calibrate()
@@ -777,15 +777,15 @@ def step_mle(state, /, *, dt, vector_field, calibration):
         output_scale_prior, _calibrated = calibration.extract(state.output_scale)
 
         prior_discretized = prior.discretize(dt)
-        hidden, extra = extrapolation.begin(
+        hidden, extra = strategy.begin(
             state.hidden, state.aux_extra, prior_discretized=prior_discretized
         )
         t = state.t + dt
         error, _, corr = correction.estimate_error(
             hidden, vector_field=vector_field, t=t
         )
 
-        hidden, extra = extrapolation.complete(
+        hidden, extra = strategy.complete(
             hidden, extra, output_scale=output_scale_prior
         )
         hidden, observed = correction.complete(hidden, corr)
@@ -800,7 +800,7 @@ def step_mle(state, /, *, dt, vector_field, calibration):
         prior=prior,
         calibration=_calibration_running_mean(ssm=ssm),
         step_implementation=step_mle,
-        extrapolation=extrapolation,
+        extrapolation=strategy,
         correction=correction,
         requires_rescaling=True,
     )
@@ -829,12 +829,12 @@ def extract(state, /):
     return _Calibration(init=init, update=update, extract=extract)
 
 
-def solver_dynamic(extrapolation, *, correction, prior, ssm):
+def solver_dynamic(strategy, *, correction, prior, ssm):
     """Create a solver that calibrates the output scale dynamically."""
 
     def step_dynamic(state, /, *, dt, vector_field, calibration):
         prior_discretized = prior.discretize(dt)
-        hidden, extra = extrapolation.begin(
+        hidden, extra = strategy.begin(
             state.hidden, state.aux_extra, prior_discretized=prior_discretized
         )
         t = state.t + dt
@@ -845,7 +845,7 @@ def step_dynamic(state, /, *, dt, vector_field, calibration):
         output_scale = calibration.update(state.output_scale, observed=observed)
 
         prior_, _calibrated = calibration.extract(output_scale)
-        hidden, extra = extrapolation.complete(hidden, extra, output_scale=prior_)
+        hidden, extra = strategy.complete(hidden, extra, output_scale=prior_)
         hidden, corr = correction.complete(hidden, corr)
 
         # Return solution
@@ -855,7 +855,7 @@ def step_dynamic(state, /, *, dt, vector_field, calibration):
     return _ProbabilisticSolver(
         prior=prior,
         ssm=ssm,
-        extrapolation=extrapolation,
+        extrapolation=strategy,
         correction=correction,
         calibration=_calibration_most_recent(ssm=ssm),
         name="Dynamic probabilistic solver",
@@ -877,22 +877,22 @@ def extract(state, /):
     return _Calibration(init=init, update=update, extract=extract)
 
 
-def solver(extrapolation, /, *, correction, prior, ssm):
+def solver(strategy, *, correction, prior, ssm):
     """Create a solver that does not calibrate the output scale automatically."""
 
     def step(state: _State, *, vector_field, dt, calibration):
         del calibration  # unused
 
         prior_discretized = prior.discretize(dt)
-        hidden, extra = extrapolation.begin(
+        hidden, extra = strategy.begin(
             state.hidden, state.aux_extra, prior_discretized=prior_discretized
         )
         t = state.t + dt
         error, _, corr = correction.estimate_error(
             hidden, vector_field=vector_field, t=t
         )
 
-        hidden, extra = extrapolation.complete(
+        hidden, extra = strategy.complete(
             hidden, extra, output_scale=state.output_scale
         )
         hidden, corr = correction.complete(hidden, corr)
@@ -906,7 +906,7 @@ def step(state: _State, *, vector_field, dt, calibration):
     return _ProbabilisticSolver(
         ssm=ssm,
         prior=prior,
-        extrapolation=extrapolation,
+        extrapolation=strategy,
         correction=correction,
         calibration=_calibration_none(),
         step_implementation=step,

tests/test_ivpsolve/test_save_every_step.py

@@ -59,5 +59,4 @@ def python_loop_solution(ivp, *, fact, strategy_fun):
 
 def reference_solution(ivp, ts):
     vf, u0, (t0, t1) = ivp
-    sol = ode.odeint_dense(vf, u0, t0=t0, t1=t1, atol=1e-10, rtol=1e-10)
-    return sol(ts)
+    return ode.odeint_and_save_at(vf, u0, save_at=ts, atol=1e-10, rtol=1e-10)

tests/test_ivpsolvers/test_calibration_dynamic_vs_mle.py

@@ -26,8 +26,10 @@ def test_exponential_approximated_well(fact):
     solver_kwargs = {"grid": grid, "solver": solver, "ssm": ssm}
     approximation = ivpsolve.solve_fixed_grid(*problem_args, **solver_kwargs)
 
-    solution = ode.odeint_dense(vf, u0, t0=t0, t1=t1, atol=1e-5, rtol=1e-5)
-    rmse = _rmse(approximation.u[0][-1], solution(t1))
+    solution = ode.odeint_and_save_at(
+        vf, u0, save_at=np.asarray([t0, t1]), atol=1e-5, rtol=1e-5
+    )
+    rmse = _rmse(approximation.u[0][-1], solution[-1])
     assert rmse < 0.1
 
 

tests/test_ivpsolvers/test_corrections.py

@@ -58,14 +58,15 @@ def fixture_solution(correction_impl, fact):
     )
 
 
-@testing.fixture(name="reference_solution")
-def fixture_reference_solution():
-    vf, (u0,), (t0, t1) = ode.ivp_lotka_volterra()
-    return ode.odeint_dense(vf, (u0,), t0=t0, t1=t1, atol=1e-10, rtol=1e-10)
-
-
-def test_terminal_value_simulation_matches_reference(solution, reference_solution):
+def test_terminal_value_simulation_matches_reference(solution):
     expected = reference_solution(solution.t)
     received = solution.u[0]
-
     assert np.allclose(received, expected, rtol=1e-2)
+
+
+@functools.jit
+def reference_solution(t1):
+    vf, (u0,), (t0, t1) = ode.ivp_lotka_volterra()
+    ts = np.asarray([t0, t1])
+    sol = ode.odeint_and_save_at(vf, (u0,), save_at=ts, atol=1e-10, rtol=1e-10)
+    return sol[-1]

tests/test_ivpsolvers/test_strategy_smoother_vs_filter.py

@@ -43,18 +43,10 @@ def fixture_smoother_solution(solver_setup):
     )
 
 
-@testing.fixture(name="reference_solution")
-def fixture_reference_solution():
-    vf, (u0,), (t0, t1) = ode.ivp_lotka_volterra()
-    return ode.odeint_dense(vf, (u0,), t0=t0, t1=t1, atol=1e-10, rtol=1e-10)
-
-
-def test_compare_filter_smoother_rmse(
-    filter_solution, smoother_solution, reference_solution
-):
+def test_compare_filter_smoother_rmse(filter_solution, smoother_solution):
     assert np.allclose(filter_solution.t, smoother_solution.t)  # sanity check
 
-    reference = reference_solution(filter_solution.t)
+    reference = _reference_solution(filter_solution.t)
     filter_rmse = _rmse(filter_solution.u[0], reference)
     smoother_rmse = _rmse(smoother_solution.u[0], reference)
 
@@ -66,5 +58,10 @@ def test_compare_filter_smoother_rmse(
     assert filter_rmse < 0.01
 
 
+def _reference_solution(ts):
+    vf, (u0,), (t0, t1) = ode.ivp_lotka_volterra()
+    return ode.odeint_and_save_at(vf, (u0,), save_at=ts, atol=1e-10, rtol=1e-10)
+
+
 def _rmse(a, b):
     return linalg.vector_norm((a - b) / b) / np.sqrt(b.size)