Improve error estimation (API *and* implementation) (#852)

* Merge local-residual error estimates and add the suffix *std

* Avoid testing the product space of configs in test_adaptive_save_at to save minutes of test time

* Update the benchmarks to the new API

* Move PDE example to examples_basic because it does not have an external dependency

* Fix some tests

* Fix remaining tests

* Explain the improved normalisation

* Update the VdP benchmarK

Author: pnkraemer

docs/examples_advanced/equinox_while_loop.py

@@ -47,9 +47,9 @@ def vf(y, /, *, t):  # noqa: ARG001
 
     strategy = probdiffeq.strategy_smoother_fixedpoint(ssm=ssm)
     solver = probdiffeq.solver(strategy=strategy, prior=ibm, constraint=ts0, ssm=ssm)
-    errorest = probdiffeq.errorest_local_residual_cached(prior=ibm, ssm=ssm)
+    error = probdiffeq.error_residual_std(constraint=ts0, prior=ibm, ssm=ssm)
     solve_adaptive = ivpsolve.solve_adaptive_terminal_values(
-        solver=solver, errorest=errorest, while_loop=while_loop
+        solver=solver, error=error, while_loop=while_loop
     )
 
     def simulate(init_val):

docs/examples_advanced/parameter_estimation_blackjax.py

@@ -201,8 +201,8 @@ def solve_adaptive(theta, *, save_at):
     # Create a probabilistic solver
     tcoeffs = taylor.odejet_padded_scan(lambda y: vf(y, t=t0), (theta,), num=2)
     init, _ibm, _ssm = probdiffeq.prior_wiener_integrated(tcoeffs, ssm_fact="isotropic")
-    errorest = probdiffeq.errorest_local_residual_cached(prior=ibm, ssm=ssm)
-    solve = ivpsolve.solve_adaptive_save_at(solver=solver, errorest=errorest)
+    error = probdiffeq.error_residual_std(constraint=ts0, prior=ibm, ssm=ssm)
+    solve = ivpsolve.solve_adaptive_save_at(solver=solver, error=error)
     return solve(init, save_at=save_at, dt0=0.1, atol=1e-4, rtol=1e-2)
 
 

docs/examples_basic/conditioning_on_zero_residual.py

@@ -89,10 +89,10 @@ def vector_field(y, /, *, t):
 ts1 = probdiffeq.constraint_ode_ts1(vector_field, ssm=ssm)
 strategy = probdiffeq.strategy_smoother_fixedpoint(ssm=ssm)
 solver = probdiffeq.solver(strategy=strategy, prior=ibm, constraint=ts1, ssm=ssm)
-errorest = probdiffeq.errorest_local_residual_cached(prior=ibm, ssm=ssm)
+error = probdiffeq.error_residual_std(constraint=ts1, prior=ibm, ssm=ssm)
 
 dt0 = ivpsolve.dt0(lambda y: vector_field(y, t=t0), (u0,))
-solve = ivpsolve.solve_adaptive_save_at(solver=solver, errorest=errorest)
+solve = ivpsolve.solve_adaptive_save_at(solver=solver, error=error)
 sol = solve(init, save_at=ts, dt0=dt0, atol=1e-1, rtol=1e-1)
 markov_seq_posterior = sol.solution_full
 

docs/examples_basic/custom_information_operator.py

@@ -99,8 +99,8 @@ def hamiltonian_2nd(u, du):
 solver_1st = probdiffeq.solver_mle(
     strategy=strategy, prior=ibm, constraint=ts1, ssm=ssm
 )
-errorest = probdiffeq.errorest_local_residual_cached(prior=ibm, ssm=ssm)
-solve = ivpsolve.solve_adaptive_save_at(solver=solver_1st, errorest=errorest)
+error = probdiffeq.error_residual_std(constraint=ts1, prior=ibm, ssm=ssm)
+solve = ivpsolve.solve_adaptive_save_at(solver=solver_1st, error=error)
 
 # -
 
@@ -165,11 +165,11 @@ def root(u, du, ddu, /, *, t):
 
 # +
 
-error_norm = probdiffeq.errorest_error_norm_rms_then_scale()
-errorest = probdiffeq.errorest_local_residual_cached(
-    prior=ibm, ssm=ssm, error_norm=error_norm
+error_norm = probdiffeq.error_norm_rms_then_scale()
+error = probdiffeq.error_residual_std(
+    constraint=ts1, prior=ibm, ssm=ssm, error_norm=error_norm
 )
-solve = ivpsolve.solve_adaptive_save_at(solver=solver_2nd, errorest=errorest)
+solve = ivpsolve.solve_adaptive_save_at(solver=solver_2nd, error=error)
 
 sol_2 = jax.jit(solve)(init, save_at=save_at, atol=1e-3, rtol=1e-1)
 ham_2 = jax.vmap(hamiltonian_2nd)(sol_2.u.mean[0], sol_2.u.mean[1])

docs/examples_basic/posterior_uncertainties.py

@@ -65,8 +65,8 @@ def vf(y, /, *, t):  # noqa: ARG001
 strategy = probdiffeq.strategy_smoother_fixedpoint(ssm=ssm)
 
 solver = probdiffeq.solver_mle(strategy=strategy, prior=ibm, constraint=ts, ssm=ssm)
-errorest = probdiffeq.errorest_local_residual_cached(prior=ibm, ssm=ssm)
-solve = ivpsolve.solve_adaptive_save_at(solver=solver, errorest=errorest)
+error = probdiffeq.error_residual_std(constraint=ts, prior=ibm, ssm=ssm)
+solve = ivpsolve.solve_adaptive_save_at(solver=solver, error=error)
 
 
 # -

docs/examples_basic/second_order_problems.py

@@ -62,7 +62,7 @@ def vf_1(y, /, *, t):
 solver_1st = probdiffeq.solver_mle(
     strategy=strategy, prior=ibm, constraint=ts0, ssm=ssm
 )
-errorest_1st = probdiffeq.errorest_local_residual_cached(prior=ibm, ssm=ssm)
+error_1st = probdiffeq.error_residual_std(constraint=ts0, prior=ibm, ssm=ssm)
 
 
 # -
@@ -71,7 +71,7 @@ def vf_1(y, /, *, t):
 
 
 save_at = jnp.linspace(t0, t1, endpoint=True, num=250)
-solve = ivpsolve.solve_adaptive_save_at(solver=solver_1st, errorest=errorest_1st)
+solve = ivpsolve.solve_adaptive_save_at(solver=solver_1st, error=error_1st)
 solution = jax.jit(solve)(init, save_at=save_at, atol=1e-5, rtol=1e-5)
 plt.plot(solution.u.mean[0][:, 0], solution.u.mean[0][:, 1], marker=".")
 plt.show()
@@ -104,15 +104,15 @@ def vf_2(y, dy, /, *, t):
 solver_2nd = probdiffeq.solver_mle(
     strategy=strategy, prior=ibm, constraint=ts0, ssm=ssm
 )
-errorest_2nd = probdiffeq.errorest_local_residual_cached(prior=ibm, ssm=ssm)
+error_2nd = probdiffeq.error_residual_std(constraint=ts0, prior=ibm, ssm=ssm)
 
 # -
 
 
 # +
 
 
-solve = ivpsolve.solve_adaptive_save_at(solver=solver_2nd, errorest=errorest_2nd)
+solve = ivpsolve.solve_adaptive_save_at(solver=solver_2nd, error=error_2nd)
 solution = jax.jit(solve)(init, save_at=save_at, atol=1e-5, rtol=1e-5)
 
 plt.plot(solution.u.mean[0][:, 0], solution.u.mean[0][:, 1], marker=".")

docs/examples_advanced/solve_pde.py docs/examples_basic/solve_pde.pydocs/examples_advanced/solve_pde.py renamed to docs/examples_basic/solve_pde.py

@@ -52,11 +52,11 @@ def vf(y, /, *, t):  # noqa: ARG001
     solver = probdiffeq.solver_dynamic(
         ssm=ssm, strategy=strategy, prior=ibm, constraint=ts
     )
-    errorest = probdiffeq.errorest_local_residual_cached(prior=ibm, ssm=ssm)
+    error = probdiffeq.error_residual_std(constraint=ts, prior=ibm, ssm=ssm)
 
     # Solve the ODE
     save_at = jnp.linspace(t0, t1, num=5, endpoint=True)
-    simulate = simulator(save_at=save_at, errorest=errorest, solver=solver)
+    simulate = simulator(save_at=save_at, error=error, solver=solver)
     (u, u_std) = simulate(init)
 
     _fig, axes = plt.subplots(
@@ -86,12 +86,12 @@ def vf(y, /, *, t):  # noqa: ARG001
 # +
 
 
-def simulator(save_at, errorest, solver):
+def simulator(save_at, error, solver):
     """Simulate a PDE."""
 
     @jax.jit
     def solve(init):
-        solve = ivpsolve.solve_adaptive_save_at(errorest=errorest, solver=solver)
+        solve = ivpsolve.solve_adaptive_save_at(error=error, solver=solver)
         solution = solve(init, save_at=save_at, atol=1e-4, rtol=1e-2)
         return (solution.u.mean[0], solution.u.std[0])
 

docs/examples_basic/taylor_coefficients.py

@@ -68,8 +68,8 @@ def solve(tc):
         strategy=strategy, prior=prior, constraint=ts0, ssm=ssm
     )
     ts = jnp.linspace(t0, t1, endpoint=True, num=10)
-    errorest = probdiffeq.errorest_local_residual_cached(prior=prior, ssm=ssm)
-    solve = ivpsolve.solve_adaptive_save_at(solver=solver, errorest=errorest)
+    error = probdiffeq.error_residual_std(constraint=ts0, prior=prior, ssm=ssm)
+    solve = ivpsolve.solve_adaptive_save_at(solver=solver, error=error)
     return solve(init, save_at=ts, atol=1e-2, rtol=1e-2)
 
 

docs/examples_benchmarks/convergence-rates-lotka-volterra.py

@@ -181,15 +181,13 @@ def param_to_solution(tol):
         # Build a solver
         init, ibm, ssm = probdiffeq.prior_wiener_integrated(tcoeffs, ssm_fact="dense")
         strategy = probdiffeq.strategy_filter(ssm=ssm)
-        corr = probdiffeq.constraint_ode_ts1(vf_probdiffeq, ssm=ssm)
-        solver = probdiffeq.solver(
-            strategy=strategy, prior=ibm, constraint=corr, ssm=ssm
-        )
-        errorest = probdiffeq.errorest_local_residual_cached(prior=ibm, ssm=ssm)
+        ts = probdiffeq.constraint_ode_ts1(vf_probdiffeq, ssm=ssm)
+        solver = probdiffeq.solver(strategy=strategy, prior=ibm, constraint=ts, ssm=ssm)
+        error = probdiffeq.error_residual_std(constraint=ts, prior=ibm, ssm=ssm)
 
         control = ivpsolve.control_proportional_integral()
         solve = ivpsolve.solve_adaptive_terminal_values(
-            solver=solver, errorest=errorest, control=control
+            solver=solver, error=error, control=control
         )
 
         # Solve

docs/examples_benchmarks/work-precision-hires.py

@@ -170,11 +170,11 @@ def param_to_solution(tol):
         solver = probdiffeq.solver_dynamic(
             strategy=strategy, prior=ibm, constraint=ts1, ssm=ssm
         )
-        errorest = probdiffeq.errorest_local_residual_cached(prior=ibm, ssm=ssm)
+        error = probdiffeq.error_residual_std(constraint=ts1, prior=ibm, ssm=ssm)
 
         control = ivpsolve.control_proportional_integral()
         solve = ivpsolve.solve_adaptive_terminal_values(
-            solver=solver, clip_dt=True, control=control, errorest=errorest
+            solver=solver, clip_dt=True, control=control, error=error
         )
 
         # Solve