I've used probdiffeq regularly in recent times, and parts of the API made my life surprisingly hard. While working with it, I refactored several components to make the library easier to extend by users (including myself). The result is (hopefully) a more readable and maintainable codebase.

Warning: Big refactor. Description below is still a work in progress.

This PR addresses several long-standing API issues:

• The lack of separation between corrections, strategies, priors, and solvers made it difficult to implement new ideas (which is unfortunate for a research-oriented library).
• The adaptive time-stepping loops were hard to read and modify, especially when changing error estimation or controller logic.
• The lack of separation between dynamic and static arguments in ivpsolve.solve_adaptive* made jit/vmap/differentiation unnecessarily tedious.
• And more.

All of these are resolved now 🎉 (Is v1.0.0 getting closer?)

Notable changes

• Adaptive time-stepping is now independent of solvers. Instead of solve(ivp, adaptive(solver, control, ...)), one now calls solve(ivp, solver, errorest, control, ...).

  This flattened hierarchy makes each component smaller and thus easier to maintain. Controllers and adaptive logic are now completely solver-agnostic and live in ivpsolve.py.

• stats.py removed. Its functionality has been moved into the respective strategy implementations. Since they share SSM and solution/argument types, co-locating them simplifies the structure. Offgrid marginals are solver methods for the same reason.

• ivpsolvers.py renamed to probdiffeq.py. After the refactor, it contains the core probabilistic solver code (and essentially no non-probnum logic except quadrature rules). The new name makes the core contribution clearer. (ivpsolve.py and taylor.py could, in principle, be standalone libraries.)

• Renamed smoothers for clarity. smoother is now smoother_fixedinterval . fixedpoint is now smoother_fixedpoint . Their relationship is now more explicit.

• Error estimation decoupled from solvers. Error estimators are now constructed via probdiffeq.errorest_*() and are only required for adaptive runs. Corrections no longer implement error estimation. Currently, two error estimators are provided.

• Corrections renamed. Since they no longer “correct” in the previous sense, they are now named
  constraint_ode_ts0/ts1/....

• Builder pattern made explicit. Strategies and solvers already followed a builder-style design; this is now more visible (e.g., explicit finalize() methods). The code is closer to textbook OOP — still fully compatible with JAX — and easier to maintain and explain.

• solve_and_save_every_step deprecated The non-jittable adaptive solver has been deprecated, which in this case it's been moved to probdiffeq.util.test_util. It remains a valuable tool for unittests, but I don't recommend using it instead of jittable solve_adaptive_save_at or other methods.

• IVPSolution type removed: Now, the outputs of the ivpsolve routines depend on the solution type of the solver. since the probabilistic solvers operate on ProbabilisticSolution types (during the forward pass), these are returned. They are almost identical to the previous IVPSolution type.

• Documentation improved: Massive overhaul of docstrings and typing (clearer structure is easier to type). As a part of this effort, all base classes (like MarkovStrategy) have been made public to communicate the code structure more openly

• Calibration built-in: Removed stats.calibrate; calibration now happens in the finalize() methods of strategies.

• ivpsolve-methods refactored: The solution routines in probdiffeq.ivpsolve.py now cleanly separate static (strings, classes, callable, integers, etc.) from dynamic arguments (pytrees of arrays): instead of solve(dynamic1, dynamic2, static1, static2), the solution routines implement solve = solve_factory(static1, static2); solution = solve(dynamic1, dynamic2), which makes the code much easier to jit/vmap/differentiate.

• While loops more accessible: While-loop selection (eg equinox's bounded while loop) is now governed by an argument to the solve-routines. No more context managers.

• Controller types public: The controller types have been made public to clarify the API

• Backend modules renamed Some of the more verbose backend module names have been shortened to streamling the src a bit

What gives?

According to the benchmarks, the code is a little bit slower than before, presumably because some of the "user-friendliness" is now built in (e.g., computing marginals of the smoothing solution, or always returning standard deviations at each step during the forward pass instead of computing them once afterwards). However, the code still has the same "big-Oh" complexity, and the wall-time differences shouldn't really be noticeable outside of dedicated benchmarks.

To make up for that:

• Since now, it's a lot easier to build custom probabilistic ODE solver models, we might experience significant performance gains down the line by building better solvers.
• The early versions of probdiffeq focussed a lot on performance, sacrificing extendability at times, which just did not age well in my opinion. (I stopped using it for a bit because I couldn't extend it). Now, with a stronger focus on human-readable source code, I hope probdiffeq gains plenty of usability.