merge new samplers

Author: arrjon

bayesflow/networks/diffusion_model/compositional_diffusion_model.py

@@ -5,11 +5,7 @@
 from keras import ops
 
 from bayesflow.types import Tensor
-from bayesflow.utils import (
-    expand_right_as,
-    integrate,
-    integrate_stochastic,
-)
+from bayesflow.utils import expand_right_as, integrate, integrate_stochastic, STOCHASTIC_METHODS
 from bayesflow.utils.serialization import serializable
 from .diffusion_model import DiffusionModel
 from .schedules.noise_schedule import NoiseSchedule
@@ -318,7 +314,7 @@ def _inverse_compositional(
         z = z / ops.sqrt(ops.cast(scale_latent, dtype=ops.dtype(z)))
 
         if density:
-            if integrate_kwargs["method"] == "euler_maruyama":
+            if integrate_kwargs["method"] in STOCHASTIC_METHODS:
                 raise ValueError("Stochastic methods are not supported for density computation.")
 
             def deltas(time, xz):
@@ -346,7 +342,7 @@ def deltas(time, xz):
 
         state = {"xz": z}
 
-        if integrate_kwargs["method"] == "euler_maruyama":
+        if integrate_kwargs["method"] in STOCHASTIC_METHODS:
 
             def deltas(time, xz):
                 return {
@@ -365,20 +361,19 @@ def diffusion(time, xz):
                 return {"xz": self.diffusion_term(xz, time=time, training=training)}
 
             score_fn = None
-            if "corrector_steps" in integrate_kwargs:
-                if integrate_kwargs["corrector_steps"] > 0:
-
-                    def score_fn(time, xz):
-                        return {
-                            "xz": self.compositional_score(
-                                xz,
-                                time=time,
-                                conditions=conditions,
-                                compute_prior_score=compute_prior_score,
-                                mini_batch_size=mini_batch_size,
-                                training=training,
-                            )
-                        }
+            if "corrector_steps" in integrate_kwargs or integrate_kwargs.get("method") == "langevin":
+
+                def score_fn(time, xz):
+                    return {
+                        "xz": self.compositional_score(
+                            xz,
+                            time=time,
+                            conditions=conditions,
+                            compute_prior_score=compute_prior_score,
+                            mini_batch_size=mini_batch_size,
+                            training=training,
+                        )
+                    }
 
             state = integrate_stochastic(
                 drift_fn=deltas,
@@ -390,7 +385,6 @@ def score_fn(time, xz):
                 **integrate_kwargs,
             )
         else:
-            integrate_kwargs.pop("corrector_steps", None)
 
             def deltas(time, xz):
                 return {

bayesflow/workflows/basic_workflow.py

@@ -291,7 +291,7 @@ def compositional_sample(
         *,
         num_samples: int,
         conditions: Mapping[str, np.ndarray],
-        compute_prior_score: Callable[[Mapping[str, np.ndarray]], np.ndarray],
+        compute_prior_score: Callable[[Mapping[str, np.ndarray]], Mapping[str, np.ndarray]],
         **kwargs,
     ) -> dict[str, np.ndarray]:
         """