Merge remote-tracking branch 'upstream/dev' into fix-additional-metrics

Author: vpratz

README.md

@@ -135,7 +135,8 @@ Many examples from [Bayesian Cognitive Modeling: A Practical Course](https://bay
 6. [Bayesian experimental design](examples/Bayesian_Experimental_Design.ipynb)
 7. [Simple model comparison example](examples/One_Sample_TTest.ipynb)
 8. [Likelihood estimation](examples/Likelihood_Estimation.ipynb)
-9. [Moving from BayesFlow v1.1 to v2.0](examples/From_BayesFlow_1.1_to_2.0.ipynb)
+9. [Multimodal data](examples/Multimodal_Data.ipynb)
+10. [Moving from BayesFlow v1.1 to v2.0](examples/From_BayesFlow_1.1_to_2.0.ipynb)
 
 More tutorials are always welcome! Please consider making a pull request if you have a cool application that you want to contribute.
 

bayesflow/__init__.py

@@ -1,26 +1,64 @@
-from . import (
-    approximators,
-    adapters,
-    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
 
@@ -59,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/adapters/adapter.py

@@ -1,4 +1,4 @@
-from collections.abc import Callable, MutableSequence, Sequence, Mapping
+from collections.abc import Callable, MutableSequence, Sequence
 
 import numpy as np
 
@@ -18,7 +18,6 @@
     Keep,
     Log,
     MapTransform,
-    NNPE,
     NumpyTransform,
     OneHot,
     Rename,
@@ -87,16 +86,14 @@ def get_config(self) -> dict:
         return serialize(config)
 
     def forward(
-        self, data: dict[str, any], *, stage: str = "inference", log_det_jac: bool = False, **kwargs
+        self, data: dict[str, any], *, log_det_jac: bool = False, **kwargs
     ) -> dict[str, np.ndarray] | tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
         """Apply the transforms in the forward direction.
 
         Parameters
         ----------
-        data : dict
+        data : dict[str, any]
             The data to be transformed.
-        stage : str, one of ["training", "validation", "inference"]
-            The stage the function is called in.
         log_det_jac: bool, optional
             Whether to return the log determinant of the Jacobian of the transforms.
         **kwargs : dict
@@ -110,28 +107,26 @@ def forward(
         data = data.copy()
         if not log_det_jac:
             for transform in self.transforms:
-                data = transform(data, stage=stage, **kwargs)
+                data = transform(data, **kwargs)
             return data
 
         log_det_jac = {}
         for transform in self.transforms:
-            transformed_data = transform(data, stage=stage, **kwargs)
+            transformed_data = transform(data, **kwargs)
             log_det_jac = transform.log_det_jac(data, log_det_jac, **kwargs)
             data = transformed_data
 
         return data, log_det_jac
 
     def inverse(
-        self, data: dict[str, np.ndarray], *, stage: str = "inference", log_det_jac: bool = False, **kwargs
+        self, data: dict[str, any], *, log_det_jac: bool = False, **kwargs
     ) -> dict[str, np.ndarray] | tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
         """Apply the transforms in the inverse direction.
 
         Parameters
         ----------
-        data : dict
+        data : dict[str, any]
             The data to be transformed.
-        stage : str, one of ["training", "validation", "inference"]
-            The stage the function is called in.
         log_det_jac: bool, optional
             Whether to return the log determinant of the Jacobian of the transforms.
         **kwargs : dict
@@ -145,18 +140,18 @@ def inverse(
         data = data.copy()
         if not log_det_jac:
             for transform in reversed(self.transforms):
-                data = transform(data, stage=stage, inverse=True, **kwargs)
+                data = transform(data, inverse=True, **kwargs)
             return data
 
         log_det_jac = {}
         for transform in reversed(self.transforms):
-            data = transform(data, stage=stage, inverse=True, **kwargs)
+            data = transform(data, inverse=True, **kwargs)
             log_det_jac = transform.log_det_jac(data, log_det_jac, inverse=True, **kwargs)
 
         return data, log_det_jac
 
     def __call__(
-        self, data: Mapping[str, any], *, inverse: bool = False, stage="inference", **kwargs
+        self, data: dict[str, any], *, inverse: bool = False, **kwargs
     ) -> dict[str, np.ndarray] | tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
         """Apply the transforms in the given direction.
 
@@ -166,8 +161,6 @@ def __call__(
             The data to be transformed.
         inverse : bool, optional
             If False, apply the forward transform, else apply the inverse transform (default False).
-        stage : str, one of ["training", "validation", "inference"]
-            The stage the function is called in.
         **kwargs
             Additional keyword arguments passed to each transform.
 
@@ -177,9 +170,9 @@ def __call__(
             The transformed data or tuple of transformed data and log determinant of the Jacobian.
         """
         if inverse:
-            return self.inverse(data, stage=stage, **kwargs)
+            return self.inverse(data, **kwargs)
 
-        return self.forward(data, stage=stage, **kwargs)
+        return self.forward(data, **kwargs)
 
     def __repr__(self):
         result = ""
@@ -701,43 +694,6 @@ def map_dtype(self, keys: str | Sequence[str], to_dtype: str):
         self.transforms.append(transform)
         return self
 
-    def nnpe(
-        self,
-        keys: str | Sequence[str],
-        *,
-        spike_scale: float | None = None,
-        slab_scale: float | None = None,
-        per_dimension: bool = True,
-        seed: int | None = None,
-    ):
-        """Append an :py:class:`~transforms.NNPE` transform to the adapter.
-
-        Parameters
-        ----------
-        keys : str or Sequence of str
-            The names of the variables to transform.
-        spike_scale : float or np.ndarray or None, default=None
-            The scale of the spike (Normal) distribution. Automatically determined if None.
-        slab_scale : float or np.ndarray or None, default=None
-            The scale of the slab (Cauchy) distribution. Automatically determined if None.
-        per_dimension : bool, default=True
-            If true, noise is applied per dimension of the last axis of the input data.
-            If false, noise is applied globally.
-        seed : int or None
-            The seed for the random number generator. If None, a random seed is used.
-        """
-        if isinstance(keys, str):
-            keys = [keys]
-
-        transform = MapTransform(
-            {
-                key: NNPE(spike_scale=spike_scale, slab_scale=slab_scale, per_dimension=per_dimension, seed=seed)
-                for key in keys
-            }
-        )
-        self.transforms.append(transform)
-        return self
-
     def one_hot(self, keys: str | Sequence[str], num_classes: int):
         """Append a :py:class:`~transforms.OneHot` transform to the adapter.
 
@@ -857,6 +813,8 @@ def standardize(
         self,
         include: str | Sequence[str] = None,
         *,
+        mean: int | float | np.ndarray,
+        std: int | float | np.ndarray,
         predicate: Predicate = None,
         exclude: str | Sequence[str] = None,
         **kwargs,
@@ -865,10 +823,14 @@ def standardize(
 
         Parameters
         ----------
-        predicate : Predicate, optional
-            Function that indicates which variables should be transformed.
         include : str or Sequence of str, optional
             Names of variables to include in the transform.
+        mean : int or float
+            Specifies the mean (location) of the transform.
+        std : int or float
+            Specifies the standard deviation (scale) of the transform.
+        predicate : Predicate, optional
+            Function that indicates which variables should be transformed.
         exclude : str or Sequence of str, optional
             Names of variables to exclude from the transform.
         **kwargs :
@@ -879,6 +841,8 @@ def standardize(
             predicate=predicate,
             include=include,
             exclude=exclude,
+            mean=mean,
+            std=std,
             **kwargs,
         )
         self.transforms.append(transform)

bayesflow/adapters/transforms/__init__.py

@@ -12,7 +12,6 @@
 from .keep import Keep
 from .log import Log
 from .map_transform import MapTransform
-from .nnpe import NNPE
 from .numpy_transform import NumpyTransform
 from .one_hot import OneHot
 from .rename import Rename

bayesflow/adapters/transforms/broadcast.py

@@ -117,7 +117,7 @@ def forward(self, data: dict[str, np.ndarray], **kwargs) -> dict[str, np.ndarray
                 data[k] = np.expand_dims(data[k], axis=tuple(np.arange(0, len_diff)))
             elif self.expand == "right":
                 data[k] = np.expand_dims(data[k], axis=tuple(-np.arange(1, len_diff + 1)))
-            elif isinstance(self.expand, tuple):
+            elif isinstance(self.expand, Sequence):
                 if len(self.expand) is not len_diff:
                     raise ValueError("Length of `expand` must match the length difference of the involed arrays.")
                 data[k] = np.expand_dims(data[k], axis=self.expand)