Refactoring of integration routines

[arrjon]

This pull request introduces significant improvements to the integration utilities, focusing on broader support for stochastic integration methods. The main changes include adding new samplers, refactoring the diffusion model to cleanly separate the score and velocity functions, updating integration method defaults and handling, and adding comprehensive tests for deterministic and stochastic solvers.

• RK45 was the classical Runge-Kutta 4 method. Now it is really a Runge-Kutta method with error controlled step size selection. I updated the coefficients to be the one from Dormand-Prince, which is the standard implementation in lots of packages and has higher accuracy. I also added TSIT5, a newer Runge-Kutta method, which is the default in DifferentialEquations.jl and Diffrax. In the benchmarking, both methods were equally accurate and fast, but TSIT5 was marginally better.

• I added several new SDE solvers (inspired again by Diffrax): SEA has same order as Euler with better properties, ShARK is a higher order method and the fast adaptive solver from Jolicoeur-Martineau et al. The latter outperformed all other samplers in speed and accuracy. I also added Langevin-sampling and predictor-corrector methods. However, these do not perform nearly as well as the ODE or SDE solvers.

• Refactored the DiffusionModel class to separate the score and velocity functions. The new score method computes the score function directly, and velocity now calls score and composes the velocity using the noise schedule, clarifying their roles and simplifying downstream usage.

• Improved the handling of stochastic integration in the diffusion model by supporting additional stochastic methods and corrector steps, and by passing the score function to integrate_stochastic where needed (e.g., for Langevin dynamics).

• Added extensive parameterized tests for both deterministic and stochastic integration methods, including forward and backward SDEs, zero-noise (ODE) reduction, and Langevin dynamics.

• Changed the default integration methods and step settings in both the diffusion model (two_step_adaptive, "adaptive") and flow matching model (tsit5, "adaptive"). This was a result of benchmarking both models on the Lueckmann low-dimensional problems.

[codecov]

Codecov Report

❌ Patch coverage is 92.88793% with 33 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
bayesflow/utils/integrate.py	93.39%	29 Missing ⚠️
...esflow/networks/diffusion_model/diffusion_model.py	77.77%	4 Missing ⚠️

Files with missing lines	Coverage Δ
bayesflow/networks/flow_matching/flow_matching.py	92.96% <100.00%> (+0.11%)	⬆️
bayesflow/utils/__init__.py	100.00% <100.00%> (ø)
...esflow/networks/diffusion_model/diffusion_model.py	83.74% <77.77%> (-1.24%)	⬇️
bayesflow/utils/integrate.py	91.49% <93.39%> (+26.42%)	⬆️

[paul-buerkner]

Thank you! Will you let us know once we shall start reviewing the PR?

[arrjon]

Thank you! Will you let us know once we shall start reviewing the PR?

All tests are passing now so it is ready :)

[arrjon]

@stefanradev93 the sampling methods are now ready for review, for jax we need to pre-allocate random numbers in the stochastic sampling, but not for the other backends.