Add diffusion model implementation

bayesflow/experimental/__init__.py

@@ -4,6 +4,7 @@
 
 from .cif import CIF
 from .continuous_time_consistency_model import ContinuousTimeConsistencyModel
+from .diffusion_model import DiffusionModel
 from .free_form_flow import FreeFormFlow
 
 from ..utils._docs import _add_imports_to_all

bayesflow/utils/__init__.py

@@ -29,9 +29,7 @@
     repo_url,
 )
 from .hparam_utils import find_batch_size, find_memory_budget
-from .integrate import (
-    integrate,
-)
+from .integrate import integrate, integrate_stochastic
 from .io import (
     pickle_load,
     format_bytes,

bayesflow/utils/integrate.py

@@ -4,10 +4,11 @@
 import keras
 
 import numpy as np
-from typing import Literal
+from typing import Literal, Union, List
 
 from bayesflow.types import Tensor
 from bayesflow.utils import filter_kwargs
+
 from . import logging
 
 ArrayLike = int | float | Tensor
@@ -293,3 +294,112 @@ def integrate(
         return integrate_scheduled(fn, state, steps, method, **kwargs)
     else:
         raise RuntimeError(f"Type or value of `steps` not understood (steps={steps})")
+
+
+def euler_maruyama_step(
+    drift_fn: Callable,
+    diffusion_fn: Callable,
+    state: dict[str, ArrayLike],
+    time: ArrayLike,
+    step_size: ArrayLike,
+    seed: keras.random.SeedGenerator,
+) -> (dict[str, ArrayLike], ArrayLike, ArrayLike):
+    """
+    Performs a single Euler-Maruyama step for stochastic differential equations.
+
+    Args:
+        drift_fn: Function that computes the drift term.
+        diffusion_fn: Function that computes the diffusion term.
+        state: Dictionary containing the current state.
+        time: Current time.
+        step_size: Size of the integration step.
+        seed: Random seed for noise generation.
+
+    Returns:
+        Tuple of (new_state, new_time, new_step_size).
+    """
+    # Compute drift term
+    drift = drift_fn(time, **filter_kwargs(state, drift_fn))
+
+    # Compute diffusion term
+    diffusion = diffusion_fn(time, **filter_kwargs(state, diffusion_fn))
+
+    # Generate noise for this step
+    noise = {}
+    for key in state.keys():
+        eps = keras.random.normal(keras.ops.shape(state[key]), dtype=keras.ops.dtype(state[key]), seed=seed)
+        noise[key] = eps * keras.ops.sqrt(keras.ops.abs(step_size))
+
+    # Check if diffusion and noise have the same keys
+    if set(diffusion.keys()) != set(noise.keys()):
+        raise ValueError("Keys of diffusion terms and noise do not match.")
+
+    # Apply updates using Euler-Maruyama formula: dx = f(x)dt + g(x)dW
+    new_state = state.copy()
+    for key in drift.keys():
+        if key in diffusion:
+            new_state[key] = state[key] + (step_size * drift[key]) + (diffusion[key] * noise[key])
+        else:
+            # If no diffusion term for this variable, apply deterministic update
+            new_state[key] = state[key] + step_size * drift[key]
+
+    new_time = time + step_size
+
+    return new_state, new_time
+
+
+def integrate_stochastic(
+    drift_fn: Callable,
+    diffusion_fn: Callable,
+    state: dict[str, ArrayLike],
+    start_time: ArrayLike,
+    stop_time: ArrayLike,
+    steps: int,
+    seed: keras.random.SeedGenerator,
+    method: str = "euler_maruyama",
+    **kwargs,
+) -> Union[dict[str, ArrayLike], tuple[dict[str, ArrayLike], dict[str, List[ArrayLike]]]]:
+    """
+    Integrates a stochastic differential equation from start_time to stop_time.
+
+    Args:
+        drift_fn: Function that computes the drift term.
+        diffusion_fn: Function that computes the diffusion term.
+        state: Dictionary containing the initial state.
+        start_time: Starting time for integration.
+        stop_time: Ending time for integration.
+        steps: Number of integration steps.
+        seed: Random seed for noise generation.
+        method: Integration method to use ('euler_maruyama').
+        **kwargs: Additional arguments to pass to the step function.
+
+    Returns:
+        If return_noise is False, returns the final state dictionary.
+        If return_noise is True, returns a tuple of (final_state, noise_history).
+    """
+    if steps <= 0:
+        raise ValueError("Number of steps must be positive.")
+
+    # Select step function based on method
+    match method:
+        case "euler_maruyama":
+            step_fn = euler_maruyama_step
+        case str() as name:
+            raise ValueError(f"Unknown integration method name: {name!r}")
+        case other:
+            raise TypeError(f"Invalid integration method: {other!r}")
+
+    # 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
+
+    time = start_time
+
+    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

bayesflow/utils/optimal_transport/log_sinkhorn.py

@@ -1,7 +1,6 @@
 import keras
 
 from .. import logging
-from ..tensor_utils import is_symbolic_tensor
 
 from .euclidean import euclidean
 
@@ -27,9 +26,6 @@ def log_sinkhorn_plan(x1, x2, regularization: float = 1.0, rtol=1e-5, atol=1e-8,
 
     log_plan = cost / -(regularization * keras.ops.mean(cost) + 1e-16)
 
-    if is_symbolic_tensor(log_plan):
-        return log_plan
-
     def contains_nans(plan):
         return keras.ops.any(keras.ops.isnan(plan))
 

bayesflow/utils/tensor_utils.py

@@ -97,9 +97,6 @@ def is_symbolic_tensor(x: Tensor) -> bool:
     if keras.utils.is_keras_tensor(x):
         return True
 
-    if not keras.ops.is_tensor(x):
-        return False
-
     match keras.backend.backend():
         case "jax":
             import jax

tests/test_networks/conftest.py

@@ -3,6 +3,23 @@
 from bayesflow.networks import MLP
 
 
+@pytest.fixture()
+def diffusion_model():
+    from bayesflow.experimental import DiffusionModel
+
+    return DiffusionModel(
+        subnet_kwargs={"widths": [64, 64]},
+        integrate_kwargs={"method": "rk45", "steps": 100},
+    )
+
+
+@pytest.fixture()
+def diffusion_model_subnet(subnet):
+    from bayesflow.experimental import DiffusionModel
+
+    return DiffusionModel(subnet=subnet)
+
+
 @pytest.fixture()
 def flow_matching():
     from bayesflow.networks import FlowMatching
@@ -84,6 +101,7 @@ def typical_point_inference_network_subnet():
         "affine_coupling_flow",
         "spline_coupling_flow",
         "flow_matching",
+        "diffusion_model",
         "free_form_flow",
         "consistency_model",
     ],
@@ -98,6 +116,7 @@ def inference_network(request):
         "typical_point_inference_network_subnet",
         "coupling_flow_subnet",
         "flow_matching_subnet",
+        "diffusion_model_subnet",
         "free_form_flow_subnet",
     ],
     scope="function",
@@ -107,7 +126,14 @@ def inference_network_subnet(request):
 
 
 @pytest.fixture(
-    params=["affine_coupling_flow", "spline_coupling_flow", "flow_matching", "free_form_flow", "consistency_model"],
+    params=[
+        "affine_coupling_flow",
+        "spline_coupling_flow",
+        "flow_matching",
+        "diffusion_model",
+        "free_form_flow",
+        "consistency_model",
+    ],
     scope="function",
 )
 def generative_inference_network(request):