Remove the need to pass inference variables to workflow

Author: stefanradev93

bayesflow/adapters/adapter.py

@@ -1,6 +1,7 @@
-from collections.abc import MutableSequence, Sequence
+from collections.abc import MutableSequence, Sequence, Mapping
 
 import numpy as np
+
 from keras.saving import (
     deserialize_keras_object as deserialize,
     register_keras_serializable as serializable,
@@ -121,16 +122,16 @@ def inverse(self, data: dict[str, np.ndarray], **kwargs) -> dict[str, any]:
 
         return data
 
-    def __call__(self, data: dict[str, any], *, inverse: bool = False, **kwargs) -> dict[str, np.ndarray]:
+    def __call__(self, data: Mapping[str, any], *, inverse: bool = False, **kwargs) -> dict[str, np.ndarray]:
         """Apply the transforms in the given direction.
 
         Parameters
         ----------
-        data : dict
+        data : Mapping[str, any]
             The data to be transformed.
         inverse : bool, optional
             If False, apply the forward transform, else apply the inverse transform (default False).
-        **kwargs : dict
+        **kwargs
             Additional keyword arguments passed to each transform.
 
         Returns

bayesflow/approximators/approximator.py

@@ -1,6 +1,9 @@
-import keras
+from collections.abc import Mapping
+
 import multiprocessing as mp
 
+import keras
+
 from bayesflow.adapters import Adapter
 from bayesflow.datasets import OnlineDataset
 from bayesflow.simulators import Simulator
@@ -19,7 +22,7 @@ def build_adapter(cls, **kwargs) -> Adapter:
         # implemented by each respective architecture
         raise NotImplementedError
 
-    def build_from_data(self, data: dict[str, any]) -> None:
+    def build_from_data(self, data: Mapping[str, any]) -> None:
         self.compute_metrics(**data, stage="training")
         self.built = True
 
@@ -72,7 +75,7 @@ def fit(self, *, dataset: keras.utils.PyDataset = None, simulator: Simulator = N
             A dataset containing simulations for training. If provided, `simulator` must be None.
         simulator : Simulator, optional
             A simulator used to generate a dataset. If provided, `dataset` must be None.
-        **kwargs : dict
+        **kwargs
             Additional keyword arguments passed to `keras.Model.fit()`, including (see also `build_dataset`):
 
             batch_size : int or None, default='auto'

bayesflow/approximators/continuous_approximator.py

@@ -1,7 +1,8 @@
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence, Callable
 
-import keras
 import numpy as np
+
+import keras
 from keras.saving import (
     deserialize_keras_object as deserialize,
     register_keras_serializable as serializable,
@@ -23,7 +24,7 @@ class ContinuousApproximator(Approximator):
 
     Parameters
     ----------
-    adapter : Adapter
+    adapter : bayesflow.adapters.Adapter
         Adapter for data processing. You can use :py:meth:`build_adapter`
         to create it.
     inference_network : InferenceNetwork
@@ -53,7 +54,7 @@ def build_adapter(
         inference_variables: Sequence[str],
         inference_conditions: Sequence[str] = None,
         summary_variables: Sequence[str] = None,
-        sample_weight: Sequence[str] = None,
+        sample_weight: str = None,
     ) -> Adapter:
         """Create an :py:class:`~bayesflow.adapters.Adapter` suited for the approximator.
 
@@ -159,7 +160,7 @@ def fit(self, *args, **kwargs):
             A dataset containing simulations for training. If provided, `simulator` must be None.
         simulator : Simulator, optional
             A simulator used to generate a dataset. If provided, `dataset` must be None.
-        **kwargs : dict
+        **kwargs
             Additional keyword arguments passed to `keras.Model.fit()`, including (see also `build_dataset`):
 
             batch_size : int or None, default='auto'
@@ -221,12 +222,50 @@ def get_config(self):
 
     def estimate(
         self,
-        conditions: dict[str, np.ndarray],
+        conditions: Mapping[str, np.ndarray],
         split: bool = False,
-        estimators: dict[str, callable] = None,
+        estimators: Mapping[str, Callable] = None,
         num_samples: int = 1000,
         **kwargs,
     ) -> dict[str, dict[str, np.ndarray]]:
+        """
+        Estimate summary statistics for variables based on given conditions.
+
+        This function samples data using the object's ``sample`` method according to the provided
+        conditions and then computes summary statistics for each variable using a set of estimator
+        functions. By default, it calculates the mean, median, and selected quantiles (10th, 50th,
+        and 90th percentiles). Users can also supply custom estimators that override or extend the
+        default ones.
+
+        Parameters
+        ----------
+        conditions : Mapping[str, np.ndarray]
+            A mapping from variable names to numpy arrays representing the conditions under which
+            samples should be generated.
+        split : bool, optional
+            If True, indicates that the data sampling process should split the samples based on an
+            internal logic. The default is False.
+        estimators : Mapping[str, Callable], optional
+            A dictionary where keys are estimator names and values are callables. Each callable must
+            accept an array and an axis parameter, and return a dictionary with the computed statistic.
+            If not provided, a default set of estimators is used:
+                - 'mean': Computes the mean along the specified axis.
+                - 'median': Computes the median along the specified axis.
+                - 'quantiles': Computes the 10th, 50th, and 90th percentiles along the specified axis,
+                  then rearranges the axes for convenience.
+        num_samples : int, optional
+            The number of samples to generate for each variable. The default is 1000.
+        **kwargs
+            Additional keyword arguments passed to the ``sample`` method.
+
+        Returns
+        -------
+        dict[str, dict[str, np.ndarray]]
+            A nested dictionary where the outer keys correspond to variable names and the inner keys
+            correspond to estimator names. Each inner dictionary contains the computed statistic(s) for
+            the variable, potentially with reduced nesting via ``squeeze_inner_estimates_dict``.
+        """
+
         estimators = estimators or {}
         estimators = (
             dict(
@@ -261,7 +300,7 @@ def sample(
         self,
         *,
         num_samples: int,
-        conditions: dict[str, np.ndarray],
+        conditions: Mapping[str, np.ndarray],
         split: bool = False,
         **kwargs,
     ) -> dict[str, np.ndarray]:
@@ -338,14 +377,14 @@ def _sample(
             **filter_kwargs(kwargs, self.inference_network.sample),
         )
 
-    def log_prob(self, data: dict[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 model. The `data` dictionary is preprocessed using the
         `adapter`. Log-probabilities are returned as NumPy arrays.
 
         Parameters
         ----------
-        data : dict[str, np.ndarray]
+        data : Mapping[str, np.ndarray]
             Dictionary of observed data as NumPy arrays.
         **kwargs : dict
             Additional keyword arguments for the adapter and log-probability computation.

bayesflow/approximators/point_approximator.py

@@ -1,5 +1,8 @@
-import keras
+from collections.abc import Mapping
+
 import numpy as np
+
+import keras
 from keras.saving import (
     register_keras_serializable as serializable,
 )
@@ -21,7 +24,7 @@ class PointApproximator(ContinuousApproximator):
 
     def estimate(
         self,
-        conditions: dict[str, np.ndarray],
+        conditions: Mapping[str, np.ndarray],
         split: bool = False,
         **kwargs,
     ) -> dict[str, dict[str, np.ndarray | dict[str, np.ndarray]]]:
@@ -33,7 +36,7 @@ def estimate(
 
         Parameters
         ----------
-        conditions : dict[str, np.ndarray]
+        conditions : Mapping[str, np.ndarray]
             A dictionary mapping variable names to arrays representing the conditions
             for the estimation process.
         split : bool, optional
@@ -71,7 +74,7 @@ def sample(
         self,
         *,
         num_samples: int,
-        conditions: dict[str, np.ndarray],
+        conditions: Mapping[str, np.ndarray],
         split: bool = False,
         **kwargs,
     ) -> dict[str, dict[str, np.ndarray]]:
@@ -111,7 +114,7 @@ def sample(
         # Optionally split the arrays along the last axis.
         if split:
             raise NotImplementedError("split=True is currently not supported for `PointApproximator`.")
-            samples = split_arrays(samples, axis=-1)
+
         # Squeeze sample dictionary if there's only one key-value pair.
         samples = self._squeeze_parametric_score_major_dict(samples)
 
@@ -120,7 +123,7 @@ def sample(
     def log_prob(
         self,
         *,
-        data: dict[str, np.ndarray],
+        data: Mapping[str, np.ndarray],
         **kwargs,
     ) -> np.ndarray | dict[str, np.ndarray]:
         """
@@ -152,14 +155,14 @@ def log_prob(
 
         return log_prob
 
-    def _prepare_conditions(self, conditions: dict[str, np.ndarray], **kwargs) -> dict[str, Tensor]:
+    def _prepare_conditions(self, conditions: Mapping[str, np.ndarray], **kwargs) -> dict[str, Tensor]:
         """Adapts and converts the conditions to tensors."""
         conditions = self.adapter(conditions, strict=False, stage="inference", **kwargs)
         conditions.pop("inference_variables", None)
         return keras.tree.map_structure(keras.ops.convert_to_tensor, conditions)
 
     def _apply_inverse_adapter_to_estimates(
-        self, estimates: dict[str, dict[str, Tensor]], **kwargs
+        self, estimates: Mapping[str, Mapping[str, Tensor]], **kwargs
     ) -> dict[str, dict[str, dict[str, np.ndarray]]]:
         """Applies the inverse adapter on each inner element of the _estimate output dictionary."""
         estimates = keras.tree.map_structure(keras.ops.convert_to_numpy, estimates)
@@ -183,7 +186,7 @@ def _apply_inverse_adapter_to_estimates(
         return processed
 
     def _apply_inverse_adapter_to_samples(
-        self, samples: dict[str, Tensor], **kwargs
+        self, samples: Mapping[str, Tensor], **kwargs
     ) -> dict[str, dict[str, np.ndarray]]:
         """Applies the inverse adapter to a dictionary of samples."""
         samples = keras.tree.map_structure(keras.ops.convert_to_numpy, samples)
@@ -198,7 +201,7 @@ def _apply_inverse_adapter_to_samples(
         return processed
 
     def _reorder_estimates(
-        self, estimates: dict[str, dict[str, dict[str, np.ndarray]]]
+        self, estimates: Mapping[str, Mapping[str, Mapping[str, np.ndarray]]]
     ) -> dict[str, dict[str, dict[str, np.ndarray]]]:
         """Reorders the nested dictionary so that the inference variable names become the top-level keys."""
         # Grab the variable names from one sample inner dictionary.
@@ -212,7 +215,7 @@ def _reorder_estimates(
         return reordered
 
     def _squeeze_estimates(
-        self, estimates: dict[str, dict[str, dict[str, np.ndarray]]]
+        self, estimates: Mapping[str, Mapping[str, Mapping[str, np.ndarray]]]
     ) -> dict[str, dict[str, np.ndarray]]:
         """Squeezes each inner estimate dictionary to remove unnecessary nesting."""
         squeezed = {}
@@ -224,7 +227,7 @@ def _squeeze_estimates(
         return squeezed
 
     def _squeeze_parametric_score_major_dict(
-        self, samples: dict[str, np.ndarray]
+        self, samples: Mapping[str, np.ndarray]
     ) -> np.ndarray or dict[str, np.ndarray]:
         """Squeezes the dictionary to just the value if there is only one key-value pair."""
         if len(samples) == 1:

bayesflow/datasets/offline_dataset.py

@@ -1,6 +1,9 @@
-import keras
+from collections.abc import Mapping
+
 import numpy as np
 
+import keras
+
 from bayesflow.adapters import Adapter
 from bayesflow.utils import logging
 
@@ -13,7 +16,12 @@ class OfflineDataset(keras.utils.PyDataset):
     """
 
     def __init__(
-        self, data: dict[str, np.ndarray], batch_size: int, adapter: Adapter | None, num_samples: int = None, **kwargs
+        self,
+        data: Mapping[str, np.ndarray],
+        batch_size: int,
+        adapter: Adapter | None,
+        num_samples: int = None,
+        **kwargs,
     ):
         super().__init__(**kwargs)
         self.batch_size = batch_size
@@ -60,7 +68,7 @@ def shuffle(self) -> None:
         np.random.shuffle(self.indices)
 
     @staticmethod
-    def _get_num_samples_from_data(data: dict) -> int:
+    def _get_num_samples_from_data(data: Mapping) -> int:
         for key, value in data.items():
             if hasattr(value, "shape"):
                 ndim = len(value.shape)

bayesflow/diagnostics/metrics/calibration_error.py

@@ -1,4 +1,4 @@
-from typing import Sequence, Any, Mapping, Callable
+from collections.abc import Sequence, Mapping, Callable
 
 import numpy as np
 
@@ -14,7 +14,7 @@ def calibration_error(
     aggregation: Callable = np.median,
     min_quantile: float = 0.005,
     max_quantile: float = 0.995,
-) -> Mapping[str, Any]:
+) -> dict[str, any]:
     """
     Computes an aggregate score for the marginal calibration error over an ensemble of approximate
     posteriors. The calibration error is given as the aggregate (e.g., median) of the absolute deviation

bayesflow/diagnostics/metrics/expected_calibration_error.py

@@ -1,6 +1,7 @@
+from collections.abc import Sequence
+
 import numpy as np
 from keras import ops
-from typing import Sequence, Any, Mapping
 
 from ...utils.exceptions import ShapeError
 from sklearn.calibration import calibration_curve
@@ -12,7 +13,7 @@ def expected_calibration_error(
     model_names: Sequence[str] = None,
     n_bins: int = 10,
     return_probs: bool = False,
-) -> Mapping[str, Any]:
+) -> dict[str, any]:
     """
     Estimates the expected calibration error (ECE) of a model comparison network according to [1].
 

bayesflow/diagnostics/metrics/posterior_contraction.py

@@ -1,4 +1,4 @@
-from typing import Sequence, Any, Mapping, Callable
+from collections.abc import Sequence, Mapping, Callable
 
 import numpy as np
 
@@ -11,7 +11,7 @@ def posterior_contraction(
     variable_keys: Sequence[str] = None,
     variable_names: Sequence[str] = None,
     aggregation: Callable = np.median,
-) -> Mapping[str, Any]:
+) -> dict[str, any]:
     """
     Computes the posterior contraction (PC) from prior to posterior for the given samples.
 

bayesflow/diagnostics/metrics/root_mean_squared_error.py

@@ -1,4 +1,4 @@
-from typing import Sequence, Any, Mapping, Callable
+from collections.abc import Sequence, Mapping, Callable
 
 import numpy as np
 
@@ -12,7 +12,7 @@ def root_mean_squared_error(
     variable_names: Sequence[str] = None,
     normalize: bool = True,
     aggregation: Callable = np.median,
-) -> Mapping[str, Any]:
+) -> dict[str, any]:
     """
     Computes the (Normalized) Root Mean Squared Error (RMSE/NRMSE) for the given posterior and prior samples.
 

bayesflow/diagnostics/plots/calibration_ecdf.py

@@ -1,15 +1,16 @@
+from collections.abc import Mapping, Sequence
+
 import numpy as np
 import matplotlib.pyplot as plt
 
-from typing import Sequence
 from ...utils.plot_utils import prepare_plot_data, add_titles_and_labels, prettify_subplots
 from ...utils.ecdf import simultaneous_ecdf_bands
 from ...utils.ecdf.ranks import fractional_ranks, distance_ranks
 
 
 def calibration_ecdf(
-    estimates: dict[str, np.ndarray] | np.ndarray,
-    targets: dict[str, np.ndarray] | np.ndarray,
+    estimates: Mapping[str, np.ndarray] | np.ndarray,
+    targets: Mapping[str, np.ndarray] | np.ndarray,
     variable_keys: Sequence[str] = None,
     variable_names: Sequence[str] = None,
     difference: bool = False,