fix stochastic sampler

Author: arrjon

bayesflow/experimental/diffusion_model.py

@@ -374,8 +374,8 @@ def __init__(
         subnet: str | type = "mlp",
         integrate_kwargs: dict[str, any] = None,
         subnet_kwargs: dict[str, any] = None,
-        noise_schedule: str | NoiseSchedule = "cosine",
-        prediction_type: str = "velocity",
+        noise_schedule: str | NoiseSchedule = "edm",
+        prediction_type: str = "F",
         **kwargs,
     ):
         """
@@ -398,10 +398,10 @@ def __init__(
             Keyword arguments passed to the subnet constructor or used to update the default MLP settings.
         noise_schedule : str or NoiseSchedule, optional
             The noise schedule used for the diffusion process. Can be "linear", "cosine", or "edm".
-            Default is "cosine".
+            Default is "edm".
         prediction_type: str, optional
             The type of prediction used in the diffusion model. Can be "velocity", "noise" or "F" (EDM).
-             Default is "velocity".
+             Default is "F".
         **kwargs
             Additional keyword arguments passed to the subnet and other components.
         """
@@ -425,10 +425,6 @@ def __init__(
         if prediction_type not in ["noise", "velocity", "F"]:  # F is EDM
             raise ValueError(f"Unknown prediction type: {prediction_type}")
         self._prediction_type = prediction_type
-        if noise_schedule.name == "edm_noise_schedule" and prediction_type != "F":
-            warnings.warn(
-                "EDM noise schedule is build for F-prediction. Consider using F-prediction instead.",
-            )
         self._loss_type = kwargs.get("loss_type", "noise")
         if self._loss_type not in ["noise", "velocity", "F"]:
             raise ValueError(f"Unknown loss type: {self._loss_type}")

bayesflow/utils/integrate.py

@@ -391,15 +391,18 @@ def integrate_stochastic(
 
     # Prepare step function with partial application
     step_fn = partial(step_fn, drift_fn=drift_fn, diffusion_fn=diffusion_fn, seed=seed, **kwargs)
-    step_size = (stop_time - start_time) / steps
 
+    step_size = (stop_time - start_time) / steps
     time = start_time
+    current_state = state.copy()
+
+    # keras.ops.fori_loop does not support keras seed generator in jax
+    for i in range(steps):
+        # Execute the step with the specific seed for this step
+        current_state, time = step_fn(
+            state=current_state,
+            time=time,
+            step_size=step_size,
+        )
 
-    def body(_loop_var, _loop_state):
-        _state, _time = _loop_state
-        _state, _time = step_fn(state=_state, time=_time, step_size=step_size)
-
-        return _state, _time
-
-    state, time = keras.ops.fori_loop(0, steps, body, (state, time))
-    return state
+    return current_state