address remaining comments

Author: LarsKue

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":

bayesflow/approximators/backend_approximators/numpy_approximator.py

@@ -448,7 +448,7 @@ def sample(
         conditions = self._prepare_data(conditions, **kwargs)
 
         # Remove any superfluous keys, just retain actual conditions
-        conditions = {k: v for k, v in conditions.items() if k in ContinuousApproximator.CONDITION_KEYS}
+        conditions = {k: v for k, v in conditions.items() if k in self.CONDITION_KEYS}
 
         # Sample and undo optional standardization
         samples = self._sample(num_samples=num_samples, **conditions, **kwargs)
@@ -485,7 +485,7 @@ def _prepare_data(
             ldj_inference = None
 
         # Standardize conditions
-        for key in ContinuousApproximator.CONDITION_KEYS:
+        for key in self.CONDITION_KEYS:
             if key in self.standardize and key in data:
                 data[key] = self.standardize_layers[key](data[key])
 
@@ -514,8 +514,12 @@ def _sample(
         summary_variables: Tensor = None,
         **kwargs,
     ) -> Tensor:
-        if (self.summary_network is None) != (summary_variables is None):
-            raise ValueError("Summary variables and summary network must be used together.")
+        if self.summary_network is None:
+            if summary_variables is not None:
+                raise ValueError("Cannot use summary variables without a summary network.")
+        else:
+            if summary_variables is None:
+                raise ValueError("Summary variables are required when a summary network is present.")
 
         if self.summary_network is not None:
             summary_outputs = self.summary_network(
@@ -606,8 +610,12 @@ def _log_prob(
         summary_variables: Tensor = None,
         **kwargs,
     ) -> Tensor:
-        if (self.summary_network is None) != (summary_variables is None):
-            raise ValueError("Summary variables and summary network must be used together.")
+        if self.summary_network is None:
+            if summary_variables is not None:
+                raise ValueError("Cannot use summary variables without a summary network.")
+        else:
+            if summary_variables is None:
+                raise ValueError("Summary variables are required when a summary network is present.")
 
         if self.summary_network is not None:
             summary_outputs = self.summary_network(

bayesflow/approximators/continuous_approximator.py

@@ -92,7 +92,7 @@ def build(self, data_shapes: dict[str, tuple[int] | dict[str, dict]]) -> None:
 
         # Set up standardization layers if requested
         if self.standardize == "all":
-            self.standardize = [var for var in ModelComparisonApproximator.CONDITION_KEYS if var in data_shapes]
+            self.standardize = [var for var in self.CONDITION_KEYS if var in data_shapes]
             self.standardize_layers = {var: Standardization(trainable=False) for var in self.standardize}
 
         # Build all standardization layers
@@ -242,7 +242,7 @@ def compute_metrics(
     def fit(
         self,
         *,
-        adapter: Adapter | str = "auto",
+        adapter: Adapter = "auto",
         dataset: keras.utils.PyDataset = None,
         simulator: ModelComparisonSimulator = None,
         simulators: Sequence[Simulator] = None,
@@ -256,7 +256,7 @@ def fit(
 
         Parameters
         ----------
-        adapter : Adapter or str, optional
+        adapter : Adapter or 'auto', optional
             The data adapter that will make the simulated / real outputs neural-network friendly.
         dataset : keras.utils.PyDataset, optional
             A dataset containing simulations for training. If provided, `simulator` must be None.
@@ -392,17 +392,20 @@ def predict(
         conditions = self.adapter(conditions, strict=False, stage="inference", **kwargs)
 
         # Ensure only keys relevant for sampling are present in the conditions dictionary
-        conditions = {k: v for k, v in conditions.items() if k in ModelComparisonApproximator.CONDITION_KEYS}
+        conditions = {k: v for k, v in conditions.items() if k in self.CONDITION_KEYS}
         conditions = keras.tree.map_structure(keras.ops.convert_to_tensor, conditions)
 
         # Optionally standardize conditions
-        for key in ModelComparisonApproximator.CONDITION_KEYS:
+        for key in self.CONDITION_KEYS:
             if key in conditions and key in self.standardize:
                 conditions[key] = self.standardize_layers[key](conditions[key])
 
         output = self._predict(**conditions, **kwargs)
 
-        return keras.ops.convert_to_numpy(keras.ops.softmax(output) if probs else output)
+        if probs:
+            output = keras.ops.softmax(output)
+
+        return keras.ops.convert_to_numpy(output)
 
     def summaries(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray:
         """

bayesflow/approximators/model_comparison_approximator.py

@@ -58,7 +58,7 @@ def estimate(
         # Adapt, optionally standardize and convert conditions to tensor.
         conditions = self._prepare_data(conditions, **kwargs)
         # Remove any superfluous keys, just retain actual conditions.  # TODO: is this necessary?
-        conditions = {k: v for k, v in conditions.items() if k in ContinuousApproximator.CONDITION_KEYS}
+        conditions = {k: v for k, v in conditions.items() if k in self.CONDITION_KEYS}
 
         estimates = self._estimate(**conditions, **kwargs)
 
@@ -77,9 +77,9 @@ def estimate(
             estimates = split_arrays(estimates, axis=-1)
 
         # Reorder the nested dictionary so that original variable names are at the top.
-        estimates = PointApproximator._reorder_estimates(estimates)
+        estimates = self._reorder_estimates(estimates)
         # Remove unnecessary nesting.
-        estimates = PointApproximator._squeeze_estimates(estimates)
+        estimates = self._squeeze_estimates(estimates)
 
         return estimates
 
@@ -124,7 +124,7 @@ def sample(
         # Adapt, optionally standardize and convert conditions to tensor.
         conditions = self._prepare_data(conditions, **kwargs)
         # Remove any superfluous keys, just retain actual conditions.  # TODO: is this necessary?
-        conditions = {k: v for k, v in conditions.items() if k in ContinuousApproximator.CONDITION_KEYS}
+        conditions = {k: v for k, v in conditions.items() if k in self.CONDITION_KEYS}
 
         # Sample and undo optional standardization
         samples = self._sample(num_samples, **conditions, **kwargs)
@@ -183,7 +183,7 @@ def log_prob(
         if log_det_jac is not None:
             log_prob = keras.tree.map_structure(lambda x: x + log_det_jac, log_prob)
 
-        log_prob = PointApproximator._squeeze_parametric_score_major_dict(log_prob)
+        log_prob = self._squeeze_parametric_score_major_dict(log_prob)
 
         return log_prob
 

bayesflow/approximators/point_approximator.py

@@ -1,4 +1,4 @@
-from collections.abc import Mapping, Callable
+from collections.abc import Callable, Mapping, Sequence
 
 import os
 import pathlib as pl
@@ -36,7 +36,7 @@ def __init__(
         load_fn: Callable = None,
         adapter: Adapter | None,
         stage: str = "training",
-        augmentations: Mapping[str, Callable] | Callable = None,
+        augmentations: Callable | Mapping[str, Callable] | Sequence[Callable] = None,
         shuffle: bool = True,
         **kwargs,
     ):
@@ -58,13 +58,14 @@ def __init__(
             Optional adapter to transform the loaded batch.
         stage : str, default="training"
             Current stage (e.g., "training", "validation", etc.) used by the adapter.
-        augmentations : dict of str to Callable or Callable, optional
-            Dictionary of augmentation functions to apply to each corresponding key in the batch
-            or a function to apply to the entire batch (possibly adding new keys).
+        augmentations : Callable or Mapping[str, Callable] or Sequence[Callable], optional
+            A single augmentation function, dictionary of augmentation functions, or sequence of augmentation functions
+            to apply to the batch.
 
             If you provide a dictionary of functions, each function should accept one element
-            of your output batch and return the corresponding transformed element. Otherwise,
-            your function should accept the entire dictionary output and return a dictionary.
+            of your output batch and return the corresponding transformed element.
+
+            Otherwise, your function should accept the entire dictionary output and return a dictionary.
 
             Note - augmentations are applied before the adapter is called and are generally
             transforms that you only want to apply during training.
@@ -81,7 +82,7 @@ def __init__(
         self.files = list(map(str, self.root.glob(pattern)))
         self.stage = stage
 
-        self.augmentations = augmentations
+        self.augmentations = augmentations or []
         self._shuffle = shuffle
         if self._shuffle:
             self.shuffle()
@@ -101,6 +102,9 @@ def __getitem__(self, item) -> dict[str, np.ndarray]:
         elif isinstance(self.augmentations, Mapping):
             for key, fn in self.augmentations.items():
                 batch[key] = fn(batch[key])
+        elif isinstance(self.augmentations, Sequence):
+            for fn in self.augmentations:
+                batch = fn(batch)
         elif isinstance(self.augmentations, Callable):
             batch = self.augmentations(batch)
         else:

bayesflow/datasets/disk_dataset.py

@@ -1,4 +1,4 @@
-from collections.abc import Mapping, Callable
+from collections.abc import Callable, Mapping, Sequence
 
 import numpy as np
 
@@ -23,7 +23,7 @@ def __init__(
         num_samples: int = None,
         *,
         stage: str = "training",
-        augmentations: Mapping[str, Callable] | Callable = None,
+        augmentations: Callable | Mapping[str, Callable] | Sequence[Callable] = None,
         shuffle: bool = True,
         **kwargs,
     ):
@@ -42,13 +42,14 @@ def __init__(
             Number of samples in the dataset. If None, it will be inferred from the data.
         stage : str, default="training"
             Current stage (e.g., "training", "validation", etc.) used by the adapter.
-        augmentations : dict of str to Callable or Callable, optional
-            Dictionary of augmentation functions to apply to each corresponding key in the batch
-            or a function to apply to the entire batch (possibly adding new keys).
+        augmentations : Callable or Mapping[str, Callable] or Sequence[Callable], optional
+            A single augmentation function, dictionary of augmentation functions, or sequence of augmentation functions
+            to apply to the batch.
 
             If you provide a dictionary of functions, each function should accept one element
-            of your output batch and return the corresponding transformed element. Otherwise,
-            your function should accept the entire dictionary output and return a dictionary.
+            of your output batch and return the corresponding transformed element.
+
+            Otherwise, your function should accept the entire dictionary output and return a dictionary.
 
             Note - augmentations are applied before the adapter is called and are generally
             transforms that you only want to apply during training.
@@ -71,7 +72,7 @@ def __init__(
 
         self.indices = np.arange(self.num_samples, dtype="int64")
 
-        self.augmentations = augmentations
+        self.augmentations = augmentations or []
         self._shuffle = shuffle
         if self._shuffle:
             self.shuffle()
@@ -111,6 +112,9 @@ def __getitem__(self, item: int) -> dict[str, np.ndarray]:
         elif isinstance(self.augmentations, Mapping):
             for key, fn in self.augmentations.items():
                 batch[key] = fn(batch[key])
+        elif isinstance(self.augmentations, Sequence):
+            for fn in self.augmentations:
+                batch = fn(batch)
         elif isinstance(self.augmentations, Callable):
             batch = self.augmentations(batch)
         else:

bayesflow/datasets/offline_dataset.py

@@ -1,4 +1,4 @@
-from collections.abc import Mapping, Callable
+from collections.abc import Callable, Mapping, Sequence
 
 import keras
 import numpy as np
@@ -20,7 +20,7 @@ def __init__(
         adapter: Adapter | None,
         *,
         stage: str = "training",
-        augmentations: Mapping[str, Callable] | Callable = None,
+        augmentations: Callable | Mapping[str, Callable] | Sequence[Callable] = None,
         **kwargs,
     ):
         """
@@ -38,13 +38,14 @@ def __init__(
             Optional adapter to transform the simulated batch.
         stage : str, default="training"
             Current stage (e.g., "training", "validation", etc.) used by the adapter.
-        augmentations : dict of str to Callable or Callable, optional
-            Dictionary of augmentation functions to apply to each corresponding key in the batch
-            or a function to apply to the entire batch (possibly adding new keys).
+        augmentations : Callable or Mapping[str, Callable] or Sequence[Callable], optional
+            A single augmentation function, dictionary of augmentation functions, or sequence of augmentation functions
+            to apply to the batch.
 
             If you provide a dictionary of functions, each function should accept one element
-            of your output batch and return the corresponding transformed element. Otherwise,
-            your function should accept the entire dictionary output and return a dictionary.
+            of your output batch and return the corresponding transformed element.
+
+            Otherwise, your function should accept the entire dictionary output and return a dictionary.
 
             Note - augmentations are applied before the adapter is called and are generally
             transforms that you only want to apply during training.
@@ -58,7 +59,7 @@ def __init__(
         self.adapter = adapter
         self.simulator = simulator
         self.stage = stage
-        self.augmentations = augmentations
+        self.augmentations = augmentations or []
 
     def __getitem__(self, item: int) -> dict[str, np.ndarray]:
         """
@@ -81,6 +82,9 @@ def __getitem__(self, item: int) -> dict[str, np.ndarray]:
         elif isinstance(self.augmentations, Mapping):
             for key, fn in self.augmentations.items():
                 batch[key] = fn(batch[key])
+        elif isinstance(self.augmentations, Sequence):
+            for fn in self.augmentations:
+                batch = fn(batch)
         elif isinstance(self.augmentations, Callable):
             batch = self.augmentations(batch)
         else: