Stateless adapters

bayesflow/__init__.py

@@ -1,6 +1,7 @@
 from . import (
     approximators,
     adapters,
+    augmentations,
     datasets,
     diagnostics,
     distributions,

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/standardize.py

@@ -1,6 +1,3 @@
-from collections.abc import Sequence
-import warnings
-
 import numpy as np
 
 from bayesflow.utils.serialization import serializable, serialize
@@ -11,120 +8,50 @@
 @serializable("bayesflow.adapters")
 class Standardize(ElementwiseTransform):
     """
-    Transform that when applied standardizes data using typical z-score standardization
-    i.e. for some unstandardized data x the standardized version z would be
+    Transform that when applied standardizes data using typical z-score standardization with
+    fixed means and std, i.e. for some unstandardized data x the standardized version z would be
 
     >>> z = (x - mean(x)) / std(x)
 
+    Important: Ensure dynamic standardization (employed by BayesFlow approximators) has been
+    turned off when using this transform.
+
     Parameters
     ----------
-    mean : int or float, optional
-        Specify a mean if known but will be estimated from data when not provided
-    std : int or float, optional
-        Specify a standard devation if known but will be estimated from data when not provided
-    axis : int, optional
-        A specific axis along which standardization should take place. By default
-        standardization happens individually for each dimension
-    momentum : float in (0,1)
-        The momentum during training
+    mean : int or float
+        Specifies the mean (location) of the transform.
+    std : int or float
+        Specifies the standard deviation (scale) of the transform.
 
     Examples
     --------
-    1) Standardize all variables using their individually estimated mean and stds.
-
-    >>> adapter = (
-            bf.adapters.Adapter()
-                .standardize()
-        )
-
-
-    2) Standardize all with same known mean and std.
-
-    >>> adapter = (
-            bf.adapters.Adapter()
-                .standardize(mean = 5, sd = 10)
-        )
-
-
-    3) Mix of fixed and estimated means/stds. Suppose we have priors for "beta" and "sigma" where we
-    know the means and stds. However for all other variables, the means and stds are unknown.
-    Then standardize should be used in several stages specifying which variables to include or exclude.
-
-    >>> adapter = (
-            bf.adapters.Adapter()
-                # mean fixed, std estimated
-                .standardize(include = "beta", mean = 1)
-                # both mean and SD fixed
-                .standardize(include = "sigma", mean = 0.6, sd = 3)
-                # both means and stds estimated for all other variables
-                .standardize(exclude = ["beta", "sigma"])
-        )
+    >>> adapter = bf.Adapter().standardize(include="beta", mean=5, std=10)
     """
 
     def __init__(
         self,
-        mean: int | float | np.ndarray = None,
-        std: int | float | np.ndarray = None,
-        axis: int | Sequence[int] = None,
-        momentum: float | None = 0.99,
+        mean: int | float | np.ndarray,
+        std: int | float | np.ndarray,
     ):
         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
 
-        if isinstance(axis, Sequence):
-            # numpy hates lists
-            axis = tuple(axis)
-        self.axis = axis
-        self.momentum = momentum
-
     def get_config(self) -> dict:
         config = {
             "mean": self.mean,
             "std": self.std,
-            "axis": self.axis,
-            "momentum": self.momentum,
         }
         return serialize(config)
 
-    def forward(self, data: np.ndarray, stage: str = "inference", **kwargs) -> np.ndarray:
-        if self.axis is None:
-            self.axis = tuple(range(data.ndim - 1))
-
-        if self.mean is None:
-            self.mean = np.mean(data, axis=self.axis, keepdims=True)
-        else:
-            if self.momentum is not None and stage == "training":
-                self.mean = self.momentum * self.mean + (1.0 - self.momentum) * np.mean(
-                    data, axis=self.axis, keepdims=True
-                )
-
-        if self.std is None:
-            self.std = np.std(data, axis=self.axis, keepdims=True, ddof=1)
-        else:
-            if self.momentum is not None and stage == "training":
-                self.std = self.momentum * self.std + (1.0 - self.momentum) * np.std(
-                    data, axis=self.axis, keepdims=True, ddof=1
-                )
-
+    def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
         mean = np.broadcast_to(self.mean, data.shape)
         std = np.broadcast_to(self.std, data.shape)
 
         return (data - mean) / std
 
     def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
-        if self.mean is None or self.std is None:
-            raise RuntimeError("Cannot call `inverse` before calling `forward` at least once.")
-
         mean = np.broadcast_to(self.mean, data.shape)
         std = np.broadcast_to(self.std, data.shape)
 

bayesflow/approximators/continuous_approximator.py

@@ -476,7 +476,7 @@ def _prepare_data(
         Handles inputs containing only conditions, only inference_variables, or both.
         Optionally tracks log-determinant Jacobian (ldj) of transformations.
         """
-        adapted = self.adapter(data, strict=False, stage="inference", log_det_jac=log_det_jac, **kwargs)
+        adapted = self.adapter(data, strict=False, log_det_jac=log_det_jac, **kwargs)
 
         if log_det_jac:
             data, ldj = adapted
@@ -565,7 +565,7 @@ def summarize(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray:
         if self.summary_network is None:
             raise ValueError("A summary network is required to compute summaries.")
 
-        data_adapted = self.adapter(data, strict=False, stage="inference", **kwargs)
+        data_adapted = self.adapter(data, strict=False, **kwargs)
         if "summary_variables" not in data_adapted or data_adapted["summary_variables"] is None:
             raise ValueError("Summary variables are required to compute summaries.")
 

bayesflow/approximators/model_comparison_approximator.py

@@ -390,7 +390,7 @@ def predict(
             probs = not logits
 
         # Apply adapter transforms to raw simulated / real quantities
-        conditions = self.adapter(conditions, strict=False, stage="inference", **kwargs)
+        conditions = self.adapter(conditions, strict=False, **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 self.CONDITION_KEYS}
@@ -429,7 +429,7 @@ def summarize(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray:
         if self.summary_network is None:
             raise ValueError("A summary network is required to compute summaries.")
 
-        data_adapted = self.adapter(data, strict=False, stage="inference", **kwargs)
+        data_adapted = self.adapter(data, strict=False, **kwargs)
         if "summary_variables" not in data_adapted or data_adapted["summary_variables"] is None:
             raise ValueError("Summary variables are required to compute summaries.")
 

bayesflow/augmentations/__init__.py

@@ -0,0 +1,5 @@
+r"""
+A collection of augmentations to modify the training data on the fly.
+"""
+
+from .nnpe import NNPE