Rename prior_ibm to prior_wiener_integrated and simplify the initial state (#827)

* Proof of concept: new API for integrated Wiener process

* Update most of the src

* Fix remaining tests

* Update examples

* Update examples and benchmarks

* Remove unused code

* Update pre-commit hook

* Rename variables to be more similar to previous versions

* Remove Pandoc installation from linter

* Remove pandoc installations

Author: pnkraemer

.github/workflows/ci.yaml

@@ -26,7 +26,6 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
         run: |
-          sudo apt-get install pandoc
           pip install --upgrade pip
           pip install .[cpu,format-and-lint]
       - name: Apply linter
@@ -111,7 +110,6 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
         run: |
-          sudo apt-get install pandoc
           pip install --upgrade pip
           pip install .[cpu,doc]
       - name: Build the HTML docs

.github/workflows/doc-publish.yaml

@@ -26,7 +26,6 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
         run: |
-          sudo apt-get install pandoc
           pip install --upgrade pip
           pip install .[cpu,doc]
       - name: Build the HTML docs

.pre-commit-config.yaml

@@ -8,7 +8,7 @@ repos:
       - id: end-of-file-fixer
       - id: check-merge-conflict
   - repo: https://github.com/lyz-code/yamlfix/
-    rev: 1.17.0
+    rev: 1.18.0
     hooks:
       - id: yamlfix
   - repo: https://github.com/astral-sh/ruff-pre-commit

docs/benchmarks/hires/run_hires.py

@@ -89,7 +89,7 @@ def param_to_solution(tol):
         # Build a solver
         vf_auto = functools.partial(vf_probdiffeq, t=t0)
         tcoeffs = taylor.odejet_padded_scan(vf_auto, (u0,), num=num_derivatives)
-        ibm, ssm = ivpsolvers.prior_ibm(tcoeffs, ssm_fact="dense")
+        init, ibm, ssm = ivpsolvers.prior_wiener_integrated(tcoeffs, ssm_fact="dense")
         ts1 = ivpsolvers.correction_ts1(ssm=ssm)
         strategy = ivpsolvers.strategy_filter(ssm=ssm)
         solver = ivpsolvers.solver_dynamic(strategy, prior=ibm, correction=ts1, ssm=ssm)
@@ -98,9 +98,6 @@ def param_to_solution(tol):
             solver, atol=1e-2 * tol, rtol=tol, control=control, ssm=ssm, clip_dt=True
         )
 
-        # Initial state
-        init = solver.initial_condition()
-
         # Solve
         dt0 = ivpsolve.dt0(vf_auto, (u0,))
         solution = ivpsolve.solve_adaptive_terminal_values(

docs/benchmarks/lotkavolterra/run_lotkavolterra.py

@@ -80,7 +80,9 @@ def param_to_solution(tol):
         vf_auto = functools.partial(vf_probdiffeq, t=t0)
         tcoeffs = taylor.odejet_padded_scan(vf_auto, (u0,), num=num_derivatives)
 
-        ibm, ssm = ivpsolvers.prior_ibm(tcoeffs, ssm_fact=implementation)
+        init, ibm, ssm = ivpsolvers.prior_wiener_integrated(
+            tcoeffs, ssm_fact=implementation
+        )
         strategy = ivpsolvers.strategy_filter(ssm=ssm)
         corr = correction(ssm=ssm)
         solver = ivpsolvers.solver_mle(strategy, prior=ibm, correction=corr, ssm=ssm)
@@ -89,9 +91,6 @@ def param_to_solution(tol):
             solver, atol=1e-2 * tol, rtol=tol, control=control, ssm=ssm
         )
 
-        # Initial state
-        init = solver.initial_condition()
-
         # Solve
         dt0 = ivpsolve.dt0(vf_auto, (u0,))
         solution = ivpsolve.solve_adaptive_terminal_values(

docs/benchmarks/pleiades/run_pleiades.py

@@ -99,7 +99,9 @@ def param_to_solution(tol):
         vf_auto = functools.partial(vf_probdiffeq, t=t0)
         tcoeffs = taylor.odejet_padded_scan(vf_auto, (u0, du0), num=num_derivatives - 1)
 
-        ibm, ssm = ivpsolvers.prior_ibm(tcoeffs, ssm_fact="isotropic")
+        init, ibm, ssm = ivpsolvers.prior_wiener_integrated(
+            tcoeffs, ssm_fact="isotropic"
+        )
         ts0_or_ts1 = correction_fun(ssm=ssm, ode_order=2)
         strategy = ivpsolvers.strategy_filter(ssm=ssm)
         solver = ivpsolvers.solver_dynamic(
@@ -110,9 +112,6 @@ def param_to_solution(tol):
             solver, atol=1e-3 * tol, rtol=tol, control=control, ssm=ssm
         )
 
-        # Initial state
-        init = solver.initial_condition()
-
         # Solve
         dt0 = ivpsolve.dt0(vf_auto, (u0, du0))
         solution = ivpsolve.solve_adaptive_terminal_values(

docs/benchmarks/vanderpol/run_vanderpol.py

@@ -81,7 +81,7 @@ def param_to_solution(tol):
         vf_auto = functools.partial(vf_probdiffeq, t=t0)
         tcoeffs = taylor.odejet_padded_scan(vf_auto, (u0, du0), num=num_derivatives - 1)
 
-        ibm, ssm = ivpsolvers.prior_ibm(tcoeffs, ssm_fact="dense")
+        init, ibm, ssm = ivpsolvers.prior_wiener_integrated(tcoeffs, ssm_fact="dense")
         ts0_or_ts1 = ivpsolvers.correction_ts1(ode_order=2, ssm=ssm)
         strategy = ivpsolvers.strategy_filter(ssm=ssm)
 
@@ -93,9 +93,6 @@ def param_to_solution(tol):
             solver, atol=1e-3 * tol, rtol=tol, control=control, ssm=ssm, clip_dt=True
         )
 
-        # Initial state
-        init = solver.initial_condition()
-
         # Solve
         dt0 = ivpsolve.dt0(vf_auto, (u0, du0))
         solution = ivpsolve.solve_adaptive_terminal_values(

docs/examples_advanced/equinox_while_loop.py

@@ -61,13 +61,12 @@ def vf(y, *, t):  # noqa: ARG001
     u0 = jnp.asarray([0.1])
 
     tcoeffs = taylor.odejet_padded_scan(lambda y: vf(y, t=t0), (u0,), num=1)
-    ibm, ssm = ivpsolvers.prior_ibm(tcoeffs, ssm_fact="isotropic")
+    init, ibm, ssm = ivpsolvers.prior_wiener_integrated(tcoeffs, ssm_fact="isotropic")
     ts0 = ivpsolvers.correction_ts0(ode_order=1, ssm=ssm)
 
     strategy = ivpsolvers.strategy_fixedpoint(ssm=ssm)
     solver = ivpsolvers.solver(strategy, prior=ibm, correction=ts0, ssm=ssm)
     adaptive_solver = ivpsolvers.adaptive(solver, ssm=ssm)
-    init = solver.initial_condition()
 
     def simulate(init_val):
         """Evaluate the parameter-to-solution function."""

docs/examples_advanced/neural_ode.py

@@ -149,15 +149,14 @@ def loss(
         """Loss function: log-marginal likelihood of the data."""
         # Build a solver
         tcoeffs = (*u0, vf(*u0, t=t0, p=p))
-        ibm, ssm = ivpsolvers.prior_ibm(
+        init, ibm, ssm = ivpsolvers.prior_wiener_integrated(
             tcoeffs, output_scale=output_scale, ssm_fact="isotropic"
         )
         ts0 = ivpsolvers.correction_ts0(ssm=ssm)
         strategy = ivpsolvers.strategy_smoother(ssm=ssm)
         solver_ts0 = ivpsolvers.solver(strategy, prior=ibm, correction=ts0, ssm=ssm)
 
         # Solve
-        init = solver_ts0.initial_condition()
         sol = ivpsolve.solve_fixed_grid(
             lambda *a, **kw: vf(*a, **kw, p=p),
             init,

docs/examples_advanced/parameter_estimation_blackjax.py

@@ -183,13 +183,12 @@ def solve_fixed(theta, *, ts):
     # Create a probabilistic solver
     tcoeffs = taylor.odejet_padded_scan(lambda y: vf(y, t=t0), (theta,), num=2)
     output_scale = 10.0
-    ibm, ssm = ivpsolvers.prior_ibm(
+    init, ibm, ssm = ivpsolvers.prior_wiener_integrated(
         tcoeffs, output_scale=output_scale, ssm_fact="isotropic"
     )
     ts0 = ivpsolvers.correction_ts0(ssm=ssm)
     strategy = ivpsolvers.strategy_filter(ssm=ssm)
     solver = ivpsolvers.solver(strategy, prior=ibm, correction=ts0, ssm=ssm)
-    init = solver.initial_condition()
     return ivpsolve.solve_fixed_grid(vf, init, grid=ts, solver=solver, ssm=ssm)
 
 
@@ -199,15 +198,13 @@ def solve_adaptive(theta, *, save_at):
     # Create a probabilistic solver
     tcoeffs = taylor.odejet_padded_scan(lambda y: vf(y, t=t0), (theta,), num=2)
     output_scale = 10.0
-    ibm, ssm = ivpsolvers.prior_ibm(
+    init, ibm, ssm = ivpsolvers.prior_wiener_integrated(
         tcoeffs, output_scale=output_scale, ssm_fact="isotropic"
     )
     ts0 = ivpsolvers.correction_ts0(ssm=ssm)
     strategy = ivpsolvers.strategy_filter(ssm=ssm)
     solver = ivpsolvers.solver(strategy, prior=ibm, correction=ts0, ssm=ssm)
     adaptive_solver = ivpsolvers.adaptive(solver, ssm=ssm)
-
-    init = solver.initial_condition()
     return ivpsolve.solve_adaptive_save_at(
         vf, init, save_at=save_at, adaptive_solver=adaptive_solver, dt0=0.1, ssm=ssm
     )