Move vector field from solution routines to corrections (#829)

* Move vector_field to corrections

* Update tests

* Update example notebooks

* Update benchmarks

* Rerun benchmarks

Author: pnkraemer

docs/benchmarks/hires/plot_ts.npy

@@ -90,7 +90,7 @@ def param_to_solution(tol):
         vf_auto = functools.partial(vf_probdiffeq, t=t0)
         tcoeffs = taylor.odejet_padded_scan(vf_auto, (u0,), num=num_derivatives)
         init, ibm, ssm = ivpsolvers.prior_wiener_integrated(tcoeffs, ssm_fact="dense")
-        ts1 = ivpsolvers.correction_ts1(ssm=ssm)
+        ts1 = ivpsolvers.correction_ts1(vf_probdiffeq, ssm=ssm)
         strategy = ivpsolvers.strategy_filter(ssm=ssm)
         solver = ivpsolvers.solver_dynamic(strategy, prior=ibm, correction=ts1, ssm=ssm)
         control = ivpsolvers.control_proportional_integral()
@@ -101,13 +101,7 @@ def param_to_solution(tol):
         # Solve
         dt0 = ivpsolve.dt0(vf_auto, (u0,))
         solution = ivpsolve.solve_adaptive_terminal_values(
-            vf_probdiffeq,
-            init,
-            t0=t0,
-            t1=t1,
-            dt0=dt0,
-            adaptive_solver=adaptive_solver,
-            ssm=ssm,
+            init, t0=t0, t1=t1, dt0=dt0, adaptive_solver=adaptive_solver, ssm=ssm
         )
 
         # Return the terminal value

docs/benchmarks/hires/plot_ys.npy

@@ -84,7 +84,7 @@ def param_to_solution(tol):
             tcoeffs, ssm_fact=implementation
         )
         strategy = ivpsolvers.strategy_filter(ssm=ssm)
-        corr = correction(ssm=ssm)
+        corr = correction(vf_probdiffeq, ssm=ssm)
         solver = ivpsolvers.solver_mle(strategy, prior=ibm, correction=corr, ssm=ssm)
         control = ivpsolvers.control_proportional_integral()
         adaptive_solver = ivpsolvers.adaptive(
@@ -94,13 +94,7 @@ def param_to_solution(tol):
         # Solve
         dt0 = ivpsolve.dt0(vf_auto, (u0,))
         solution = ivpsolve.solve_adaptive_terminal_values(
-            vf_probdiffeq,
-            init,
-            t0=t0,
-            t1=t1,
-            dt0=dt0,
-            adaptive_solver=adaptive_solver,
-            ssm=ssm,
+            init, t0=t0, t1=t1, dt0=dt0, adaptive_solver=adaptive_solver, ssm=ssm
         )
 
         # Return the terminal value

docs/benchmarks/hires/results.npy

@@ -102,7 +102,7 @@ def param_to_solution(tol):
         init, ibm, ssm = ivpsolvers.prior_wiener_integrated(
             tcoeffs, ssm_fact="isotropic"
         )
-        ts0_or_ts1 = correction_fun(ssm=ssm, ode_order=2)
+        ts0_or_ts1 = correction_fun(vf_probdiffeq, ssm=ssm, ode_order=2)
         strategy = ivpsolvers.strategy_filter(ssm=ssm)
         solver = ivpsolvers.solver_dynamic(
             strategy, prior=ibm, correction=ts0_or_ts1, ssm=ssm
@@ -115,13 +115,7 @@ def param_to_solution(tol):
         # Solve
         dt0 = ivpsolve.dt0(vf_auto, (u0, du0))
         solution = ivpsolve.solve_adaptive_terminal_values(
-            vf_probdiffeq,
-            init,
-            t0=t0,
-            t1=t1,
-            dt0=dt0,
-            adaptive_solver=adaptive_solver,
-            ssm=ssm,
+            init, t0=t0, t1=t1, dt0=dt0, adaptive_solver=adaptive_solver, ssm=ssm
         )
 
         # Return the terminal value