Merge remote-tracking branch 'upstream/dev' into docs-spatial-data

Author: vpratz

bayesflow/__init__.py

@@ -1,27 +1,64 @@
-from . import (
-    approximators,
-    adapters,
-    augmentations,
-    datasets,
-    diagnostics,
-    distributions,
-    experimental,
-    networks,
-    simulators,
-    utils,
-    workflows,
-    wrappers,
-)
-
-from .adapters import Adapter
-from .approximators import ContinuousApproximator, PointApproximator
-from .datasets import OfflineDataset, OnlineDataset, DiskDataset
-from .simulators import make_simulator
-from .workflows import BasicWorkflow
+# ruff: noqa: E402
+# disable E402 to allow for setup code before importing any internals (which could import keras)
 
 
 def setup():
     # perform any necessary setup without polluting the namespace
+    import os
+    from importlib.util import find_spec
+
+    issue_url = "https://github.com/bayesflow-org/bayesflow/issues/new?template=bug_report.md"
+
+    if "KERAS_BACKEND" not in os.environ:
+        # check for available backends and automatically set the KERAS_BACKEND env variable or raise an error
+        class Backend:
+            def __init__(self, display_name, package_name, env_name, install_url, priority):
+                self.display_name = display_name
+                self.package_name = package_name
+                self.env_name = env_name
+                self.install_url = install_url
+                self.priority = priority
+
+        backends = [
+            Backend("JAX", "jax", "jax", "https://docs.jax.dev/en/latest/quickstart.html#installation", 0),
+            Backend("PyTorch", "torch", "torch", "https://pytorch.org/get-started/locally/", 1),
+            Backend("TensorFlow", "tensorflow", "tensorflow", "https://www.tensorflow.org/install", 2),
+        ]
+
+        found_backends = []
+        for backend in backends:
+            if find_spec(backend.package_name) is not None:
+                found_backends.append(backend)
+
+        if not found_backends:
+            message = "No suitable backend found. Please install one of the following:\n"
+            for backend in backends:
+                message += f"{backend.display_name}\n"
+            message += "\n"
+
+            message += f"If you continue to see this error, please file a bug report at {issue_url}.\n"
+            message += (
+                "You can manually select a backend by setting the KERAS_BACKEND environment variable as shown below:\n"
+            )
+            message += "https://keras.io/getting_started/#configuring-your-backend"
+
+            raise ImportError(message)
+
+        if len(found_backends) > 1:
+            import warnings
+
+            found_backends.sort(key=lambda b: b.priority)
+            chosen_backend = found_backends[0]
+
+            warnings.warn(
+                f"Multiple Keras-compatible backends detected ({', '.join(b.display_name for b in found_backends)}).\n"
+                f"Defaulting to {chosen_backend.display_name}.\n"
+                "To override, set the KERAS_BACKEND environment variable before importing bayesflow.\n"
+                "See: https://keras.io/getting_started/#configuring-your-backend"
+            )
+        else:
+            os.environ["KERAS_BACKEND"] = found_backends[0].env_name
+
     import keras
     import logging
 
@@ -60,3 +97,24 @@ def setup():
 # call and clean up namespace
 setup()
 del setup
+
+from . import (
+    approximators,
+    adapters,
+    augmentations,
+    datasets,
+    diagnostics,
+    distributions,
+    experimental,
+    networks,
+    simulators,
+    utils,
+    workflows,
+    wrappers,
+)
+
+from .adapters import Adapter
+from .approximators import ContinuousApproximator, PointApproximator
+from .datasets import OfflineDataset, OnlineDataset, DiskDataset
+from .simulators import make_simulator
+from .workflows import BasicWorkflow

bayesflow/approximators/continuous_approximator.py

@@ -537,7 +537,7 @@ def _sample(
             )
             batch_shape = keras.ops.shape(inference_conditions)[:-1]
         else:
-            batch_shape = keras.ops.shape(inference_conditions)[1:-1]
+            batch_shape = (num_samples,)
 
         return self.inference_network.sample(
             batch_shape, conditions=inference_conditions, **filter_kwargs(kwargs, self.inference_network.sample)

bayesflow/approximators/point_approximator.py

@@ -143,12 +143,7 @@ def sample(
 
         return samples
 
-    def log_prob(
-        self,
-        *,
-        data: Mapping[str, np.ndarray],
-        **kwargs,
-    ) -> np.ndarray | dict[str, np.ndarray]:
+    def log_prob(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray | dict[str, np.ndarray]:
         """
         Computes the log-probability of given data under the parametric distribution(s) for given input conditions.
 

bayesflow/diagnostics/plots/calibration_ecdf.py

@@ -144,9 +144,11 @@ def calibration_ecdf(
             tq_targets = test_quantity_fn(data=targets)
             test_quantities_targets[key] = np.expand_dims(tq_targets, axis=1)
 
-            # # Flatten estimates for batch processing in test_quantity_fn, apply function, and restore shape
+            # Flatten estimates for batch processing in test_quantity_fn, apply function, and restore shape
             num_conditions, num_samples = next(iter(estimates.values())).shape[:2]
-            flattened_estimates = keras.tree.map_structure(lambda t: np.reshape(t, (-1, *t.shape[2:])), estimates)
+            flattened_estimates = keras.tree.map_structure(
+                lambda t: np.reshape(t, (num_conditions * num_samples, *t.shape[2:])), estimates
+            )
             flat_tq_estimates = test_quantity_fn(data=flattened_estimates)
             test_quantities_estimates[key] = np.reshape(flat_tq_estimates, (num_conditions, num_samples, 1))
 

bayesflow/links/positive_definite.py

@@ -43,8 +43,8 @@ def compute_input_shape(self, output_shape):
 
         There are m nonzero elements of a lower triangular nxn matrix with m = n * (n + 1) / 2.
 
-        Example
-        -------
+        Examples
+        --------
         >>> PositiveDefinite().compute_output_shape((None, 3, 3))
         6
         """

bayesflow/networks/consistency_models/consistency_model.py

@@ -4,7 +4,7 @@
 import numpy as np
 
 from bayesflow.types import Tensor
-from bayesflow.utils import find_network, layer_kwargs, weighted_mean
+from bayesflow.utils import find_network, layer_kwargs, weighted_mean, tensor_utils, expand_right_as
 from bayesflow.utils.serialization import deserialize, serializable, serialize
 
 from ..inference_network import InferenceNetwork
@@ -67,6 +67,11 @@ def __init__(
             Final number of discretization steps
         subnet_kwargs: dict[str, any], optional
             Keyword arguments passed to the subnet constructor or used to update the default MLP settings.
+        concatenate_subnet_input: bool, optional
+            Flag for advanced users to control whether all inputs to the subnet should be concatenated
+            into a single vector or passed as separate arguments. If set to False, the subnet
+            must accept three separate inputs: 'x' (noisy parameters), 't' (time),
+            and optional 'conditions'. Default is True.
         **kwargs    : dict, optional, default: {}
             Additional keyword arguments
         """
@@ -77,6 +82,7 @@ def __init__(
         subnet_kwargs = subnet_kwargs or {}
         if subnet == "mlp":
             subnet_kwargs = ConsistencyModel.MLP_DEFAULT_CONFIG | subnet_kwargs
+        self._concatenate_subnet_input = kwargs.get("concatenate_subnet_input", True)
 
         self.subnet = find_network(subnet, **subnet_kwargs)
         self.output_projector = keras.layers.Dense(
@@ -119,6 +125,7 @@ def get_config(self):
             "eps": self.eps,
             "s0": self.s0,
             "s1": self.s1,
+            "concatenate_subnet_input": self._concatenate_subnet_input,
             # we do not need to store subnet_kwargs
         }
 
@@ -161,18 +168,23 @@ def build(self, xz_shape, conditions_shape=None):
 
         input_shape = list(xz_shape)
 
-        # time vector
-        input_shape[-1] += 1
+        if self._concatenate_subnet_input:
+            # construct time vector
+            input_shape[-1] += 1
+            if conditions_shape is not None:
+                input_shape[-1] += conditions_shape[-1]
+            input_shape = tuple(input_shape)
 
-        if conditions_shape is not None:
-            input_shape[-1] += conditions_shape[-1]
-
-        input_shape = tuple(input_shape)
-
-        self.subnet.build(input_shape)
-
-        input_shape = self.subnet.compute_output_shape(input_shape)
-        self.output_projector.build(input_shape)
+            self.subnet.build(input_shape)
+            out_shape = self.subnet.compute_output_shape(input_shape)
+        else:
+            # Multiple separate inputs
+            time_shape = tuple(xz_shape[:-1]) + (1,)  # same batch/sequence dims, 1 feature
+            self.subnet.build(x_shape=xz_shape, t_shape=time_shape, conditions_shape=conditions_shape)
+            out_shape = self.subnet.compute_output_shape(
+                x_shape=xz_shape, t_shape=time_shape, conditions_shape=conditions_shape
+            )
+        self.output_projector.build(out_shape)
 
         # Choose coefficient according to [2] Section 3.3
         self.c_huber = 0.00054 * ops.sqrt(xz_shape[-1])
@@ -256,6 +268,35 @@ def _inverse(self, z: Tensor, conditions: Tensor = None, training: bool = False,
             x = self.consistency_function(x_n, t, conditions=conditions, training=training)
         return x
 
+    def _apply_subnet(
+        self, x: Tensor, t: Tensor, conditions: Tensor = None, training: bool = False
+    ) -> Tensor | tuple[Tensor, Tensor, Tensor]:
+        """
+        Prepares and passes the input to the subnet either by concatenating the latent variable `x`,
+        the time `t`, and optional conditions or by returning them separately.
+
+        Parameters
+        ----------
+        x : Tensor
+            The parameter tensor, typically of shape (..., D), but can vary.
+        t : Tensor
+            The time tensor, typically of shape (..., 1).
+        conditions : Tensor, optional
+            The optional conditioning tensor (e.g. parameters).
+        training : bool, optional
+            The training mode flag, which can be used to control behavior during training.
+
+        Returns
+        -------
+        Tensor
+            The output tensor from the subnet.
+        """
+        if self._concatenate_subnet_input:
+            xtc = tensor_utils.concatenate_valid([x, t, conditions], axis=-1)
+            return self.subnet(xtc, training=training)
+        else:
+            return self.subnet(x=x, t=t, conditions=conditions, training=training)
+
     def consistency_function(self, x: Tensor, t: Tensor, conditions: Tensor = None, training: bool = False) -> Tensor:
         """Compute consistency function.
 
@@ -271,12 +312,8 @@ def consistency_function(self, x: Tensor, t: Tensor, conditions: Tensor = None,
             Whether internal layers (e.g., dropout) should behave in train or inference mode.
         """
 
-        if conditions is not None:
-            xtc = ops.concatenate([x, t, conditions], axis=-1)
-        else:
-            xtc = ops.concatenate([x, t], axis=-1)
-
-        f = self.output_projector(self.subnet(xtc, training=training))
+        subnet_out = self._apply_subnet(x, t, conditions, training=training)
+        f = self.output_projector(subnet_out)
 
         # Compute skip and out parts (vectorized, since self.sigma2 is of shape (1, input_dim)
         # Thus, we can do a cross product with the time vector which is (batch_size, 1) for
@@ -316,8 +353,8 @@ def compute_metrics(
 
         log_p = ops.log(p)
         times = keras.random.categorical(ops.expand_dims(log_p, 0), ops.shape(x)[0], seed=self.seed_generator)[0]
-        t1 = ops.take(discretized_time, times)[..., None]
-        t2 = ops.take(discretized_time, times + 1)[..., None]
+        t1 = expand_right_as(ops.take(discretized_time, times), x)
+        t2 = expand_right_as(ops.take(discretized_time, times + 1), x)
 
         # generate noise vector
         noise = keras.random.normal(keras.ops.shape(x), dtype=keras.ops.dtype(x), seed=self.seed_generator)