Update the docs (#821)

pnkraemer · probnumschool · web-flow · commit 0c20563b76c1 · 2025-09-04T10:42:24.000+02:00
* Avoid linebreaks in docs

* Update the transitioning guide

* Format and lint

---------

Co-authored-by: pnkraemer &lt;pnkraemer@duck.com&gt;
diff --git a/docs/benchmarks/taylor_fitzhughnagumo/plot.md b/docs/benchmarks/taylor_fitzhughnagumo/plot.md
@@ -13,7 +13,7 @@ jupyter:
     name: python3
 ---
 
-# Taylor-series: FitzHugh-Nagumo
+# FitzHugh-Nagumo
 
 The FHN problem is a common non-stiff differential equation.
 
diff --git a/docs/benchmarks/taylor_node/plot.md b/docs/benchmarks/taylor_node/plot.md
@@ -13,7 +13,7 @@ jupyter:
     name: python3
 ---
 
-# Taylor-series: Neural ODE problem
+# Neural ODE
 
 ```python
 """Benchmark all Taylor-series estimators on a Neural ODE."""
diff --git a/docs/benchmarks/taylor_pleiades/plot.md b/docs/benchmarks/taylor_pleiades/plot.md
@@ -13,7 +13,7 @@ jupyter:
     name: python3
 ---
 
-# Taylor-series: Pleiades
+# Pleiades
 
 The Pleiades problem is a common non-stiff differential equation.
 
diff --git a/docs/examples_advanced/neural_ode.py b/docs/examples_advanced/neural_ode.py
@@ -12,7 +12,7 @@
 #     name: python3
 # ---
 
-# # Training a Neural ODE with Optax
+# # Neural ODEs
 #
 # We can use the parameter estimation functionality
 # to fit a neural ODE to a time series data set.
diff --git a/docs/examples_advanced/physics_enhanced_regression_1.py b/docs/examples_advanced/physics_enhanced_regression_1.py
@@ -12,7 +12,7 @@
 #     name: python3
 # ---
 
-# # Estimating ODE parameters with Optax
+# # Parameter estimation (Optax)
 #
 # **Time-series data and optimization with ``optax``**
 #
diff --git a/docs/examples_advanced/physics_enhanced_regression_2.py b/docs/examples_advanced/physics_enhanced_regression_2.py
@@ -12,7 +12,7 @@
 #     name: python3
 # ---
 
-# # Estimating ODE parameters with BlackJAX
+# # Parameter estimation (BlackJAX)
 #
 #
 # This tutorial explains how to estimate unknown parameters of
diff --git a/docs/examples_advanced/use_equinox_bounded_while_loop.py b/docs/examples_advanced/use_equinox_bounded_while_loop.py
@@ -12,7 +12,7 @@
 #     name: python3
 # ---
 
-# # Using Equinox's reverse-mode differentiable while-loops
+# # Equinox's while-loops
 #
 # Use [Equinox's](https://docs.kidger.site/equinox/)
 # bounded while loop to enable reverse-mode differentiation of adaptive IVP solvers.
diff --git a/docs/examples_basic/conditioning-on-zero-residual.py b/docs/examples_basic/conditioning-on-zero-residual.py
@@ -12,7 +12,7 @@
 #     name: python3
 # ---
 
-# # Understanding probabilistic solvers as collocation methods
+# # How probabilistic solvers work
 #
 # Probabilistic solvers condition a prior distribution
 # on satisfying a zero-ODE-residual on a specified grid.
diff --git a/docs/examples_basic/dynamic_output_scales.py b/docs/examples_basic/dynamic_output_scales.py
@@ -12,7 +12,7 @@
 #     name: python3
 # ---
 
-# # Choosing between dynamic and non-dynamic solvers
+# # Solver types
 #
 # You can choose between a `adaptive.solver_calibrationfree()`
 # (which does not calibrate the output-scale),
diff --git a/docs/examples_basic/posterior_uncertainties.py b/docs/examples_basic/posterior_uncertainties.py
@@ -12,7 +12,7 @@
 #     name: python3
 # ---
 
-# # Exploring posterior uncertainties
+# # Posterior uncertainties
 
 # +
 """Display the marginal uncertainties of filters and smoothers."""
diff --git a/docs/examples_basic/second_order_problems.py b/docs/examples_basic/second_order_problems.py
@@ -12,7 +12,7 @@
 #     name: python3
 # ---
 
-# # Simulating second-order systems
+# # Second-order systems
 
 # +
 """Demonstrate how to solve second-order IVPs without transforming them first."""
diff --git a/docs/examples_basic/taylor_coefficients.py b/docs/examples_basic/taylor_coefficients.py
@@ -12,7 +12,7 @@
 #     name: python3
 # ---
 
-# # Understanding the role of Taylor coefficients
+# # Taylor coefficients
 #
 # To build a probabilistic solver, we need to build a specific state-space model.
 # To build this specific state-space model, we interact with Taylor coefficients.
diff --git a/docs/examples_quickstart/easy_example.py b/docs/examples_quickstart/easy_example.py
@@ -12,7 +12,7 @@
 #     name: python3
 # ---
 
-# # Getting started with ProbDiffEq
+# # Quickstart
 #
 # Let's have a look at an easy example.
 
diff --git a/docs/getting_started/transitioning_from_other_packages.md b/docs/getting_started/transitioning_from_other_packages.md
@@ -1,4 +1,4 @@
-# Transitioning from other packages
+# Transitioning
 
 Here is how you get started with ProbDiffEq for solving ordinary differential equations (ODEs) 
 if you already have experience with other (probabilistic) ODE solver packages in Python and Julia.
@@ -26,22 +26,6 @@ It has been used, for instance, to solve [million-dimensional](https://arxiv.org
 ProbDiffEq is (more or less) a successor of Tornadox: it can do almost everything that Tornadox can do, but is generally faster (compiling entire solver loops instead of only single steps), offers more solvers, and provides more features built ``around'' IVP solutions: e.g. dense output or posterior sampling.
 ProbDiffEq is also more thoroughly tested and documented, and has a few features that are not yet (that is, at the time of writing this document) implemented in Tornadox.
 
-ProbDiffEq can reproduce most of the implementations in Tornadox:
-
-| In Tornadox:                        | In ProbDiffEq:                                    | Comments                                                                                                                                                         |
-|---------------------------------------|---------------------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| `ek0.KroneckerEK0()`                  | `dynamic(strategy_filter(ibm_adaptive(), ts0()))` | Combine with `impl.select("isotropic", ...)`                                                                                                                     |
-| `ek0.DiagonalEK0()`                   | `dynamic(strategy_filter(ibm_adaptive(), ts0()))` | Combine with `impl.select("blockdiag", ...)`                                                                                                                     |
-| `ek0.ReferenceEK0()`                  | `dynamic(strategy_filter(ibm_adaptive(), ts0()))` | Combine with `impl.select("dense", ...)`                                                                                                                         |
-| `ek1.ReferenceEK1()`                  | `dynamic(strategy_filter(ibm_adaptive(), ts1()))` | Combine with `impl.select("dense", ...)`                                                                                                                         |
-| `ek1.ReferenceEK1ConstantDiffusion()` | `mle(strategy_filter(ibm_adaptive(), ts1()))`     | Combine with `impl.select("dense", ...)`.                                                                                                                        |
-| `ek1.DiagonalEK1()`                   | Work in progress.                                 |                                                                                                                                                                  |
-| `solver.solve()`                      | `solve_adaptive_save_every_step()`                     | Try `solve_adaptive_save_at()` instead.                                                                                                                               |
-| `solver.simulate_final_state()`       | `simulate_terminal_values()`                      | ProbDiffEq compiles the whole loop; it will be much faster.                                                                                                      |
-| `solver.solution_generator()`         | Work in progress.                                 |                                                                                                                                                                  |
-| `init.TaylorMode()`                   | `taylor.taylor.taylor_mode`                     | Consider `taylor.taylor.odejet_via_jvp()` for low numbers of derivatives and `taylor.taylor.odejet_doubling_unroll()` for (absurdly) high numbers of derivatives |
-| `init.RungeKutta()`                   | `taylor.taylor.runge_kutta_starter()`         |                                                                                                                                                                  |
-
 
 Recently, the development of Tornadox has not been very active, so its API is relatively stable.
 In contrast, APIs in ProbDiffEq may change.
@@ -70,49 +54,15 @@ To translate between the two packages, consider the following:
 * Everything termed `EK0` or `EK1` in ProbNumDiffEq.jl is `ts0` or `ts1` in ProbDiffEq ("ts" stands for "Taylor series" linearisation and stands in contrast to "slr", i.e. "statistical linear regression").
 * ProbNumDiffEq.jl calibrates output scales via `DynamicDiffusion`, `FixedDiffusion`, `DynamicMVDiffusion`, or `FixedMVDiffusion`. 
  Their equivalents in ProbDiffEq are `ivpsolvers.solver_dynamic()` or `ivpsolvers.solver_mle()`. Feed them with any strategies (Filters/Smoothers) and any state-space model implementations. Use a block-diagonal implementation (e.g. `impl.choose("blockdiag")`) for multivariate output scales ("`MVDiffusion`"). Try the `ivpsolvers.solver()` with a manual (gradient-based?) calibration if the other routines are unsatisfactory.
-* ProbNumDiffEq.jl refers to `ibm_adaptive(output_scale=x)` as `IWP(diffusion=x^2)`. They are the same processes. 
+* ProbNumDiffEq.jl refers to `prior_ibm(output_scale=x)` as `IWP(diffusion=x^2)`. They are the same processes. 
 * ProbNumDiffEq.jl switches between filtering and smoothing with a `smooth=true/false` flag. ProbDiffEq uses different strategies to distinguish these strategies and offers a third one (fixedpoint-smoothing). 
-* Initialisation schemes like those in `ProbNumDiffEq` are in `probdiffeq/taylor/*.py`. ProbDiffEq offers some rules for high-order differential equations and some unique methods (e.g. doubling). But the feature lists are relatively similar.
-
-Should I replace ProbNumDiffEq.jl with ProbDiffEq?
-Short answer: No. 
-Long answer: Use ProbNumDiffEq.jl in Julia and ProbDiffEq in JAX. 
-
-
+* Initialisation schemes like those in `ProbNumDiffEq` are in `probdiffeq.taylor`. ProbDiffEq offers some rules for high-order differential equations and some unique methods (e.g. doubling). But the feature lists are relatively similar.
 
 ## Transitioning from ProbNum
 
 [ProbNum](https://probnum.readthedocs.io/en/latest/) is a probabilistic numerics library in Python, just like ProbDiffEq.
 ProbNum collects probabilistic solvers for many problems, not just ordinary differential equations.
-Its API and documentation are more mature than the API and documentation in ProbDiffEq.
-That said, ProbDiffEq specialises in pure JAX and state-space-model-based IVP solvers, which leads to significant efficiency gains.
-
-
-The features (ignoring the non-IVP-related features in ProbNum, such as linear solvers or numerical integration) 
-differ between ProbDiffEq and ProbNum as follows:
-
-* ProbNum implements IVP solvers, and general continuous/discrete filters and smoothers. 
-  At the moment, ProbDiffEq's support for discrete-time filtering and smoothing is very limited. 
-  The ability to seamlessly switch between ODE problems and filtering problems in ProbNum is helpful, for example, to build [latent force models](https://arxiv.org/abs/2103.10153).
-* The filtering and smoothing options in ProbNum are broader than what ProbDiffEq needs;
-  e.g., ProbNum offers particle filtering or nonlinear, continuous-time state-space models.
-* ProbNum offers perturbation-based solvers, which ProbDiffEq does not provide.
-* ProbNum offers callbacks, which are not yet supported in ProbDiffEq
-
-In contrast:
-* ProbDiffEq offers state-space model factorisations, different modes of linearisation, 
-  more output-scale calibration routines, more estimation strategies, 
-  and different modes of solving.
-* ProbDiffeq computes IVP solutions faster.
-* ProbDiffEq is compatible with everything that builds on JAX.
-
-
-At the time of writing, ProbNum is NumPy-based (but a [JAX-backend has been in development](https://github.com/probabilistic-numerics/probnum/pull/581)), and ProbDiffEq is pure JAX.
-Using JAX implies we can take a few specialisations (mostly vmap- or  PyTree-centric ones) that lead to drastic efficiency gains, but ProbDiffEq will generally not be compatible with non-JAX based code.
-
-In other words, the solvers in ProbNum are excellent for their breadth and for didactic purposes, but ProbDiffEq is more efficient in what it provides (check the benchmarks!).
-Our solvers also naturally work with function transformations such as `vmap`, `jacfwd`, or `jit`, which is quite helpful for combining ProbDiffEq with, e.g. `optax` or `blackjax` (check the examples).
-
+That said, ProbDiffEq specialises in pure JAX and state-space-model-based IVP solvers, which leads to significant efficiency gains. Probdiffeq implements more solvers than ProbNum.
 
 
 ## Transitioning from Diffrax
@@ -139,7 +89,7 @@ To roughly translate the Diffrax IVP solvers to ProbDiffEq solvers, consider the
 
 | If you use the following solvers in Diffrax:              | You might like the following solvers in ProbDiffEq:                                                                     | Comments                                                                               | 
 |-----------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------|
-| `Heun()`, `Midpoint()`, `Ralston()`, `LeapfrogMidpoint()` | e.g. `ibm_adaptive(num_derivatives=1)`, `ts0()` with an `isotropic` or `blockdiag` implementation                       | Use a block-diagonal factorisation if the ODE dimensions have greatly different scales |
+| `Heun()`, `Midpoint()`, `Ralston()`, `LeapfrogMidpoint()` | e.g. `prior_ibm(num_derivatives=1)`, `ts0()` with an `isotropic` or `blockdiag` state-space model                       | Use a block-diagonal state-space model if the ODE dimensions have greatly different scales |
 | `Bosh3()`                                                 | increase `num_derivatives` to `num_derivatives=2` in the above                                                          | See above.                                                                             |
 | `Tsit5()`, `Dopri5()`                                     | increase `num_derivatives` to `num_derivatives=4` in the above                                                          | See above.                                                                             |
 | `Dopri8()`                                                | increase `num_derivatives` to `num_derivatives={5,6,7}` in the above. If this is inefficient, try a `ts1()` correction. | See above.                                                                             |
diff --git a/docs/getting_started/troubleshooting.md b/docs/getting_started/troubleshooting.md
@@ -15,15 +15,13 @@ This is because, during solver initialization, it computes the Cholesky factor o
 
 If a solution routine takes an unexpectedly long time to compile but runs quickly afterward, the issue might be related to how Taylor coefficients are computed. 
 Some functions in `probdiffeq.taylor` unroll a small loop, which can slow down compilation.  
-
-To avoid this, try using `probdiffeq.taylor.taylor.odejet_padded_scan()`, which replaces loop unrolling with a scan.  
-
+To avoid this, try using the padded scan, which replaces loop unrolling with a scan.  
 If the problem persists, consider:  
 
 - Reducing the number of derivatives (if appropriate for your problem).  
-- Switching to a different Taylor-coefficient routine, such as a Runge-Kutta starter with `probdiffeq.taylor.taylor.runge_kutta_starter()`.  
+- Switching to a different Taylor-coefficient routine, such as a Runge-Kutta starter.
 
-For $\nu < 5$, using a Runge-Kutta starter should maintain solver performance. However, for higher-order methods (e.g., \(\nu = 9\)), `taylor_fn=taylor.odejet_fn` appears to be the best choice.  
+For $\nu < 5$, using a Runge-Kutta starter should maintain solver performance. However, for higher-order methods (e.g., \(\nu = 9\)), Taylor-mode ("jets") appears to be the best choice.  
 
 
 ## Taylor-derivative routines yield NaNs
diff --git a/mkdocs.yml b/mkdocs.yml
@@ -75,22 +75,20 @@ extra:
 nav:
   - Probabilistic solvers for differential equations in JAX: index.md
   - An easy example: examples_quickstart/easy_example.ipynb
-  - 'FAQ: PROBDIFFEQ 101':
-      - Transitioning from other packages: getting_started/transitioning_from_other_packages.md
-      - Choosing a solver: getting_started/choosing_a_solver.md
-      - Troubleshooting: getting_started/troubleshooting.md
-  - SOLVER CONFIGURATION:
-      - examples_solver_config/conditioning-on-zero-residual.ipynb
-      - examples_solver_config/posterior_uncertainties.ipynb
-      - examples_solver_config/dynamic_output_scales.ipynb
-      - examples_solver_config/second_order_problems.ipynb
-      - examples_solver_config/taylor_coefficients.ipynb
-  - PARAMETER ESTIMATION:
-      - Parameter estimation with Optax: examples_parameter_estimation/physics_enhanced_regression_1.ipynb
-      - Parameter estimation with BlackJAX: examples_parameter_estimation/physics_enhanced_regression_2.ipynb
-      - Neural ODEs with Optax: examples_parameter_estimation/neural_ode.ipynb
-  - MISCELLANEOUS EXAMPLES:
-      - examples_misc/use_equinox_bounded_while_loop.ipynb
+  - Transitioning from other packages: getting_started/transitioning_from_other_packages.md
+  - Choosing a solver: getting_started/choosing_a_solver.md
+  - Troubleshooting: getting_started/troubleshooting.md
+  - EXAMPLES | BASIC:
+      - examples_basic/conditioning-on-zero-residual.ipynb
+      - examples_basic/posterior_uncertainties.ipynb
+      - examples_basic/dynamic_output_scales.ipynb
+      - examples_basic/second_order_problems.ipynb
+      - examples_basic/taylor_coefficients.ipynb
+  - EXAMPLES | ADVANCED:
+      - examples_advanced/physics_enhanced_regression_1.ipynb
+      - examples_advanced/physics_enhanced_regression_2.ipynb
+      - examples_advanced/neural_ode.ipynb
+      - examples_advanced/use_equinox_bounded_while_loop.ipynb
   - API DOCUMENTATION:
       - ivpsolve: api_docs/ivpsolve.md
       - ivpsolvers: api_docs/ivpsolvers.md
@@ -106,6 +104,7 @@ nav:
       - benchmarks/pleiades/plot.ipynb
       - benchmarks/hires/plot.ipynb
       - benchmarks/vanderpol/plot.ipynb
-      - benchmarks/taylor_pleiades/plot.ipynb
-      - benchmarks/taylor_fitzhughnagumo/plot.ipynb
-      - benchmarks/taylor_node/plot.ipynb
+      - Taylor coefficients:
+          - benchmarks/taylor_pleiades/plot.ipynb
+          - benchmarks/taylor_fitzhughnagumo/plot.ipynb
+          - benchmarks/taylor_node/plot.ipynb

Original file line number	Diff line number	Diff line change
`@@ -12,7 +12,7 @@`
`12`	`12`	`# name: python3`
`13`	`13`	`# ---`
`14`	`14`
`15`		`-# # Training a Neural ODE with Optax`
	`15`	`+# # Neural ODEs`
`16`	`16`	`#`
`17`	`17`	`# We can use the parameter estimation functionality`
`18`	`18`	`# to fit a neural ODE to a time series data set.`