Compositional sampling diffusion

[arrjon]

This pull request introduces compositional sampling support to the BayesFlow framework, enabling diffusion models to handle multiple compositional conditions efficiently. The main changes span the continuous approximator, diffusion model, and inference network modules, adding new methods and refactoring existing ones to support compositional structures in sampling, inference, and diffusion processes.

Larger changes include:

• Added a new compositional_sample method to ContinuousApproximator, which generates samples with compositional structure and handles flattening, reshaping, and prior score computation for multiple compositional conditions. Supporting internal method _compositional_sample was also introduced.
• In DiffusionModel, implemented compositional diffusion support including:
  • New compositional_bridge and compositional_velocity methods for compositional score calculation.
  • _compute_individual_scores helper for handling multiple compositional conditions.
  • _inverse_compositional method for inverse compositional diffusion sampling.

The idea is that the workflow now has the method compositional_sample, which expects conditions in the form (n_datasets, n_conditions, ...). Then we can perform compositional sampling with diffusion models.
compositional_sample allows to set a mini_batch_size for memory efficient computation of the compositional score, which does not work with jax backend however, as jax does not like stochasticity in its integrators which cannot be precomputed. We could support here only fixed step sizes though?

To compute the compositional score we need access to the score of the prior. Here we need to handle the adapter carefully so that we compute the correct score. In the current draft, I am not sure I computed the prior score correctly. Some ideas would be great, currently it fails for jax because the adpater is converting stuff to numpy back and forth, but for torch it is working.

[codecov]

Codecov Report

❌ Patch coverage is 56.47668% with 84 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
bayesflow/approximators/continuous_approximator.py	3.57%	54 Missing ⚠️
...esflow/networks/diffusion_model/diffusion_model.py	77.88%	23 Missing ⚠️
bayesflow/networks/inference_network.py	55.55%	4 Missing ⚠️
bayesflow/simulators/simulator.py	50.00%	1 Missing ⚠️
bayesflow/utils/integrate.py	95.00%	1 Missing ⚠️
bayesflow/workflows/basic_workflow.py	50.00%	1 Missing ⚠️

Files with missing lines	Coverage Δ
bayesflow/simulators/sequential_simulator.py	46.03% <ø> (ø)
bayesflow/simulators/simulator.py	82.92% <50.00%> (-1.69%)	⬇️
bayesflow/utils/integrate.py	67.68% <95.00%> (+16.01%)	⬆️
bayesflow/workflows/basic_workflow.py	69.71% <50.00%> (-0.23%)	⬇️
bayesflow/networks/inference_network.py	70.58% <55.55%> (-3.23%)	⬇️
...esflow/networks/diffusion_model/diffusion_model.py	80.41% <77.88%> (+0.61%)	⬆️
bayesflow/approximators/continuous_approximator.py	72.33% <3.57%> (-19.55%)	⬇️

... and 1 file with indirect coverage changes