Fix remaining tests

Author: pnkraemer

tests/test_probdiffeq/test_calibration_mle.py

@@ -43,7 +43,7 @@ def case_simulate_terminal_values(fact):
     def solver_to_solution(solver_fun, strategy_fun):
         strategy = strategy_fun(ssm=ssm)
         solver = solver_fun(strategy=strategy, prior=ibm, constraint=ts0, ssm=ssm)
-        error = probdiffeq.error_residual_std(prior=ibm, ssm=ssm)
+        error = probdiffeq.error_residual_std(constraint=ts0, prior=ibm, ssm=ssm)
         solve = ivpsolve.solve_adaptive_terminal_values(error=error, solver=solver)
         return func.jit(solve)(init, t0=t0, t1=t1, dt0=dt0, atol=1e-2, rtol=1e-2)
 

tests/test_probdiffeq/test_constraints.py

@@ -52,7 +52,7 @@ def fixture_solution(constraint_ode_factory, fact):
     solver = probdiffeq.solver_mle(
         strategy=strategy, prior=ibm, constraint=constraint, ssm=ssm
     )
-    error = probdiffeq.error_residual_std(prior=ibm, ssm=ssm)
+    error = probdiffeq.error_residual_std(constraint=constraint, prior=ibm, ssm=ssm)
     solve = ivpsolve.solve_adaptive_terminal_values(solver=solver, error=error)
     return solve(init, t0=t0, t1=t1, atol=1e-2, rtol=1e-2, damp=1e-9)
 

tests/test_probdiffeq/test_log_marginal_likelihood.py

@@ -16,7 +16,7 @@ def fixture_solution(fact):
     strategy = probdiffeq.strategy_smoother_fixedpoint(ssm=ssm)
 
     solver = probdiffeq.solver(strategy=strategy, prior=ibm, constraint=ts0, ssm=ssm)
-    error = probdiffeq.error_residual_std(prior=ibm, ssm=ssm)
+    error = probdiffeq.error_residual_std(constraint=ts0, prior=ibm, ssm=ssm)
 
     save_at = np.linspace(t0, t1, endpoint=True, num=4)
     solve = ivpsolve.solve_adaptive_save_at(error=error, solver=solver)

tests/test_probdiffeq/test_log_marginal_likelihood_terminal_values.py

@@ -30,7 +30,7 @@ def fixture_solution(strategy_func, fact):
     ts0 = probdiffeq.constraint_ode_ts0(vf, ssm=ssm)
     strategy = strategy_func(ssm=ssm)
     solver = probdiffeq.solver(strategy=strategy, prior=ibm, constraint=ts0, ssm=ssm)
-    error = probdiffeq.error_residual_std(prior=ibm, ssm=ssm)
+    error = probdiffeq.error_residual_std(constraint=ts0, prior=ibm, ssm=ssm)
     solve = ivpsolve.solve_adaptive_terminal_values(solver=solver, error=error)
     sol = func.jit(solve)(init, t0=t0, t1=t1, atol=1e-2, rtol=1e-2)
     return sol, strategy

tests/test_probdiffeq/test_offgrid_marginals.py

@@ -16,7 +16,7 @@ def test_save_at_result_matches_interpolated_adaptive_result(fact):
     ts0 = probdiffeq.constraint_ode_ts0(vf, ssm=ssm)
     strategy = probdiffeq.strategy_filter(ssm=ssm)
     solver = probdiffeq.solver(strategy=strategy, prior=ibm, constraint=ts0, ssm=ssm)
-    error = probdiffeq.error_residual_std(prior=ibm, ssm=ssm)
+    error = probdiffeq.error_residual_std(constraint=ts0, prior=ibm, ssm=ssm)
 
     # Compute an adaptive solution and interpolate
     ts = np.linspace(t0, t1, num=15, endpoint=True)

tests/test_probdiffeq/test_sample.py

@@ -15,7 +15,7 @@ def fixture_approximation_and_strategy(fact):
     strategy = probdiffeq.strategy_smoother_fixedpoint(ssm=ssm)
     solver = probdiffeq.solver(strategy=strategy, prior=ibm, constraint=ts0, ssm=ssm)
 
-    error = probdiffeq.error_residual_std(prior=ibm, ssm=ssm)
+    error = probdiffeq.error_residual_std(constraint=ts0, prior=ibm, ssm=ssm)
     solve = ivpsolve.solve_adaptive_save_at(solver=solver, error=error)
     save_at = np.linspace(t0, t1, endpoint=True, num=7)
     sol = func.jit(solve)(init, save_at=save_at, atol=1e-2, rtol=1e-2)

tests/test_probdiffeq/test_strategy_smoother_fixedinterval_vs_fixedpoint.py

@@ -24,7 +24,7 @@ def fixture_solution_smoother(solver_setup):
     ts0 = probdiffeq.constraint_ode_ts0(solver_setup["vf"], ssm=ssm)
     strategy = probdiffeq.strategy_smoother_fixedinterval(ssm=ssm)
     solver = probdiffeq.solver(strategy=strategy, prior=ibm, constraint=ts0, ssm=ssm)
-    error = probdiffeq.error_residual_std(prior=ibm, ssm=ssm)
+    error = probdiffeq.error_residual_std(constraint=ts0, prior=ibm, ssm=ssm)
     solve = test_util.solve_adaptive_save_every_step(error=error, solver=solver)
     t0, t1 = solver_setup["t0"], solver_setup["t1"]
     return solve(init, t0=t0, t1=t1, dt0=0.1, atol=1e-3, rtol=1e-3)
@@ -37,7 +37,7 @@ def test_fixedpoint_smoother_equivalent_same_grid(solver_setup, solution_smoothe
     ts0 = probdiffeq.constraint_ode_ts0(solver_setup["vf"], ssm=ssm)
     strategy = probdiffeq.strategy_smoother_fixedpoint(ssm=ssm)
     solver = probdiffeq.solver(strategy=strategy, prior=ibm, constraint=ts0, ssm=ssm)
-    error = probdiffeq.error_residual_std(prior=ibm, ssm=ssm)
+    error = probdiffeq.error_residual_std(constraint=ts0, prior=ibm, ssm=ssm)
 
     save_at = solution_smoother.t
     solve = ivpsolve.solve_adaptive_save_at(error=error, solver=solver)
@@ -92,7 +92,7 @@ def test_fixedpoint_smoother_equivalent_different_grid(solver_setup, solution_sm
     ts0 = probdiffeq.constraint_ode_ts0(solver_setup["vf"], ssm=ssm)
     strategy_fp = probdiffeq.strategy_smoother_fixedpoint(ssm=ssm)
     solver = probdiffeq.solver(strategy=strategy_fp, prior=ibm, constraint=ts0, ssm=ssm)
-    error = probdiffeq.error_residual_std(prior=ibm, ssm=ssm)
+    error = probdiffeq.error_residual_std(constraint=ts0, prior=ibm, ssm=ssm)
     solve = ivpsolve.solve_adaptive_save_at(error=error, solver=solver)
     solution_fixedpoint = func.jit(solve)(
         init, save_at=ts, dt0=0.1, atol=1e-3, rtol=1e-3

tests/test_probdiffeq/test_strategy_warnings_for_wrong_strategies.py

@@ -15,7 +15,7 @@ def test_warning_for_fixedpoint_in_save_every_step_mode(fact):
     ts0 = probdiffeq.constraint_ode_ts0(vf, ssm=ssm)
     strategy = probdiffeq.strategy_smoother_fixedpoint(ssm=ssm)
     solver = probdiffeq.solver(strategy=strategy, prior=ibm, constraint=ts0, ssm=ssm)
-    error = probdiffeq.error_residual_std(prior=ibm, ssm=ssm)
+    error = probdiffeq.error_residual_std(constraint=ts0, prior=ibm, ssm=ssm)
 
     with testing.warns():
         _ = test_util.solve_adaptive_save_every_step(error=error, solver=solver)
@@ -30,6 +30,6 @@ def test_warning_for_smoother_in_save_at_mode(fact):
     ts0 = probdiffeq.constraint_ode_ts0(vf, ssm=ssm)
     strategy = probdiffeq.strategy_smoother_fixedinterval(ssm=ssm)
     solver = probdiffeq.solver(strategy=strategy, prior=ibm, constraint=ts0, ssm=ssm)
-    error = probdiffeq.error_residual_std(prior=ibm, ssm=ssm)
+    error = probdiffeq.error_residual_std(constraint=ts0, prior=ibm, ssm=ssm)
     with testing.warns():
         _ = ivpsolve.solve_adaptive_save_at(solver=solver, error=error)