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

Author: vpratz

README.md

@@ -130,7 +130,7 @@ Many examples from [Bayesian Cognitive Modeling: A Practical Course](https://bay
 1. [Linear regression starter example](examples/Linear_Regression_Starter.ipynb)
 2. [From ABC to BayesFlow](examples/From_ABC_to_BayesFlow.ipynb)
 3. [Two moons starter example](examples/Two_Moons_Starter.ipynb)
-4. [Rapid iteration with point estimators](examples/Lotka_Volterra_Point_Estimation_and_Expert_Stats.ipynb)
+4. [Rapid iteration with point estimators](examples/Lotka_Volterra_Point_Estimation.ipynb)
 5. [SIR model with custom summary network](examples/SIR_Posterior_Estimation.ipynb)
 6. [Bayesian experimental design](examples/Bayesian_Experimental_Design.ipynb)
 7. [Simple model comparison example](examples/One_Sample_TTest.ipynb)

bayesflow/adapters/adapter.py

@@ -820,7 +820,7 @@ def split(self, key: str, *, into: Sequence[str], indices_or_sections: int | Seq
 
         return self
 
-    def squeeze(self, keys: str | Sequence[str], *, axis: int | tuple):
+    def squeeze(self, keys: str | Sequence[str], *, axis: int | Sequence[int]):
         """Append a :py:class:`~transforms.Squeeze` transform to the adapter.
 
         Parameters

bayesflow/adapters/transforms/nan_to_num.py

@@ -13,12 +13,14 @@ class NanToNum(Transform):
 
     Parameters
     ----------
-    default_value : float
-        Value to substitute wherever data is NaN.
-    return_mask : bool, default=False
-        If True, a mask array will be returned under a new key.
-    mask_prefix : str, default='mask_'
-        Prefix for the mask key in the output dictionary.
+    key : str
+        The variable key to look for in the simulation data dict.
+    default_value : float, optional
+        Value to substitute wherever data is NaN. Default is 0.0.
+    return_mask : bool, optional
+        If True, a mask array will be returned under a new key. Default is False.
+    mask_prefix : str, optional
+        Prefix for the mask key in the output dictionary. Default is 'mask_'.
     """
 
     def __init__(self, key: str, default_value: float = 0.0, return_mask: bool = False, mask_prefix: str = "mask"):
@@ -81,10 +83,10 @@ def inverse(self, data: dict[str, any], **kwargs) -> dict[str, any]:
         values = data[self.key]
 
         if not self.return_mask:
-            values[values == self.default_value] = np.nan  # we assume default_value is not in data
+            # assumes default_value is not in nan
+            values[values == self.default_value] = np.nan
         else:
             mask_array = data[self.mask_key].astype(bool)
-            # Put NaNs where mask is 0
             values[~mask_array] = np.nan
 
         data[self.key] = values

bayesflow/adapters/transforms/nnpe.py

@@ -65,8 +65,8 @@ class NNPE(ElementwiseTransform):
     def __init__(
         self,
         *,
-        spike_scale: float | np.ndarray | None = None,
-        slab_scale: float | np.ndarray | None = None,
+        spike_scale: np.typing.ArrayLike | None = None,
+        slab_scale: np.typing.ArrayLike | None = None,
         per_dimension: bool = True,
         seed: int | None = None,
     ):
@@ -80,14 +80,14 @@ def __init__(
     def _resolve_scale(
         self,
         name: str,
-        passed: float | np.ndarray | None,
+        passed: np.typing.ArrayLike | None,
         default: float,
         data: np.ndarray,
     ) -> np.ndarray | float:
         """
         Determine spike/slab scale:
-         - If passed is None: Automatic determination via default * std(data) (per‐dimension or global).
-         - Else: validate & cast passed to the correct shape/type.
+         - If `passed` is None: Automatic determination via default * std(data) (per‐dimension or global).
+         - Else: Validate & cast `passed` to the correct shape/type.
 
         Parameters
         ----------
@@ -103,8 +103,8 @@ def _resolve_scale(
 
         Returns
         -------
-        float or np.ndarray
-            The resolved scale, either as a scalar (if per_dimension=False) or an 1D array of length data.shape[-1]
+        np.ndarray
+            The resolved scale, either as a 0D array (if per_dimension=False) or an 1D array of length data.shape[-1]
             (if per_dimension=True).
         """
 
@@ -119,22 +119,22 @@ def _resolve_scale(
 
         # If no scale is passed, determine scale automatically given the dimensionwise or global std
         if passed is None:
-            return default * std
+            return np.array(default * std)
         # If a scale is passed, check if the passed shape matches the expected shape
         else:
-            if self.per_dimension:
+            try:
                 arr = np.asarray(passed, dtype=float)
-                if arr.shape != expected_shape or arr.ndim != 1:
+            except Exception as e:
+                raise TypeError(f"{name}: expected values convertible to float, got {type(passed).__name__}") from e
+
+            if self.per_dimension:
+                if arr.ndim != 1 or arr.shape != expected_shape:
                     raise ValueError(f"{name}: expected array of shape {expected_shape}, got {arr.shape}")
                 return arr
             else:
-                try:
-                    scalar = float(passed)
-                except TypeError:
-                    raise TypeError(f"{name}: expected a scalar convertible to float, got type {type(passed).__name__}")
-                except ValueError:
-                    raise ValueError(f"{name}: expected a scalar convertible to float, got value {passed!r}")
-                return scalar
+                if arr.ndim != 0:
+                    raise ValueError(f"{name}: expected scalar, got array of shape {arr.shape}")
+                return arr
 
     def forward(self, data: np.ndarray, stage: str = "inference", **kwargs) -> np.ndarray:
         """
@@ -173,7 +173,7 @@ def forward(self, data: np.ndarray, stage: str = "inference", **kwargs) -> np.nd
         return data + noise
 
     def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
-        """Non-invertible transform."""
+        # Non-invertible transform
         return data
 
     def get_config(self) -> dict:

bayesflow/adapters/transforms/random_subsample.py

@@ -8,8 +8,10 @@ class RandomSubsample(ElementwiseTransform):
     """
     A transform that takes a random subsample of the data within an axis.
 
-    Example: adapter.random_subsample("x", sample_size = 3, axis = -1)
+    Examples
+    --------
 
+    >>> adapter = bf.Adapter().random_subsample("x", sample_size=3, axis=-1)
     """
 
     def __init__(
@@ -20,23 +22,22 @@ def __init__(
         super().__init__()
         if isinstance(sample_size, float):
             if sample_size <= 0 or sample_size >= 1:
-                ValueError("Sample size as a percentage must be a float between 0 and 1 exclusive. ")
+                raise ValueError("Sample size as a percentage must be a float between 0 and 1 exclusive. ")
         self.sample_size = sample_size
         self.axis = axis
 
     def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
-        axis = self.axis
-        max_sample_size = data.shape[axis]
+        max_sample_size = data.shape[self.axis]
 
         if isinstance(self.sample_size, int):
             sample_size = self.sample_size
         else:
             sample_size = np.round(self.sample_size * max_sample_size)
 
         # random sample without replacement
-        sample_indices = np.random.permutation(max_sample_size)[0 : sample_size - 1]
+        sample_indices = np.random.permutation(max_sample_size)[:sample_size]
 
-        return np.take(data, sample_indices, axis)
+        return np.take(data, sample_indices, self.axis)
 
     def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
         # non invertible transform

bayesflow/adapters/transforms/squeeze.py

@@ -1,5 +1,6 @@
 import numpy as np
 
+from collections.abc import Sequence
 from bayesflow.utils.serialization import serializable, serialize
 
 from .elementwise_transform import ElementwiseTransform
@@ -29,8 +30,10 @@ class Squeeze(ElementwiseTransform):
     It is recommended to precede this transform with a :class:`~bayesflow.adapters.transforms.ToArray` transform.
     """
 
-    def __init__(self, *, axis: int | tuple):
+    def __init__(self, *, axis: int | Sequence[int]):
         super().__init__()
+        if isinstance(axis, Sequence):
+            axis = tuple(axis)
         self.axis = axis
 
     def get_config(self) -> dict:

bayesflow/adapters/transforms/standardize.py

@@ -1,4 +1,5 @@
 from collections.abc import Sequence
+import warnings
 
 import numpy as np
 
@@ -69,6 +70,14 @@ def __init__(
     ):
         super().__init__()
 
+        if mean is None or std is None:
+            warnings.warn(
+                "Dynamic standardization is deprecated and will be removed in later versions."
+                "Instead, use the standardize argument of the approximator / workflow instance or provide "
+                "fixed mean and std arguments. You may incur some redundant computations if you keep this transform.",
+                FutureWarning,
+            )
+
         self.mean = mean
         self.std = std
 

bayesflow/adapters/transforms/transform.py

@@ -22,6 +22,7 @@ def __repr__(self):
 
     @classmethod
     def from_config(cls, config: dict, custom_objects=None):
+        # noinspection PyArgumentList
         return cls(**deserialize(config, custom_objects=custom_objects))
 
     def get_config(self) -> dict:

bayesflow/approximators/approximator.py

@@ -11,18 +11,16 @@
 
 
 class Approximator(BackendApproximator):
-    def build(self, data_shapes: any) -> None:
-        mock_data = keras.tree.map_structure(keras.ops.zeros, data_shapes)
-        self.build_from_data(mock_data)
+    def build(self, data_shapes: dict[str, tuple[int] | dict[str, dict]]) -> None:
+        raise NotImplementedError
 
     @classmethod
     def build_adapter(cls, **kwargs) -> Adapter:
         # implemented by each respective architecture
         raise NotImplementedError
 
-    def build_from_data(self, data: dict[str, any]) -> None:
-        self.compute_metrics(**filter_kwargs(data, self.compute_metrics), stage="training")
-        self.built = True
+    def build_from_data(self, adapted_data: dict[str, any]) -> None:
+        raise NotImplementedError
 
     @classmethod
     def build_dataset(
@@ -61,6 +59,9 @@ def build_dataset(
             max_queue_size=max_queue_size,
         )
 
+    def call(self, *args, **kwargs):
+        return self.compute_metrics(*args, **kwargs)
+
     def fit(self, *, dataset: keras.utils.PyDataset = None, simulator: Simulator = None, **kwargs):
         """
         Trains the approximator on the provided dataset or on-demand data generated from the given simulator.
@@ -132,6 +133,7 @@ def fit(self, *, dataset: keras.utils.PyDataset = None, simulator: Simulator = N
             logging.info("Building on a test batch.")
             mock_data = dataset[0]
             mock_data = keras.tree.map_structure(keras.ops.convert_to_tensor, mock_data)
-            self.build_from_data(mock_data)
+            mock_data_shapes = keras.tree.map_structure(keras.ops.shape, mock_data)
+            self.build(mock_data_shapes)
 
         return super().fit(dataset=dataset, **kwargs)

bayesflow/approximators/backend_approximators/backend_approximator.py

@@ -6,8 +6,6 @@
 match keras.backend.backend():
     case "jax":
         from .jax_approximator import JAXApproximator as BaseBackendApproximator
-    case "numpy":
-        from .numpy_approximator import NumpyApproximator as BaseBackendApproximator
     case "tensorflow":
         from .tensorflow_approximator import TensorFlowApproximator as BaseBackendApproximator
     case "torch":