Remove nnpe from adapter

Author: stefanradev93

bayesflow/__init__.py

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

bayesflow/adapters/adapter.py

@@ -18,7 +18,6 @@
     Keep,
     Log,
     MapTransform,
-    NNPE,
     NumpyTransform,
     OneHot,
     Rename,
@@ -695,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.
 

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/augmentations/__init__.py

@@ -0,0 +1 @@
+from .nnpe import NNPE

bayesflow/adapters/transforms/nnpe.py renamed to bayesflow/augmentations/nnpe.py

@@ -1,12 +1,7 @@
 import numpy as np
 
-from bayesflow.utils.serialization import serializable, serialize
 
-from .elementwise_transform import ElementwiseTransform
-
-
-@serializable("bayesflow.adapters")
-class NNPE(ElementwiseTransform):
+class NNPE:
     """Implements noisy neural posterior estimation (NNPE) as described in [1], which adds noise following a
     spike-and-slab distribution to the training data as a mild form of data augmentation to robustify against noisy
     real-world data (see [1, 2] for benchmarks). Adds the options of automatic noise scale determination and
@@ -34,12 +29,18 @@ class NNPE(ElementwiseTransform):
         automatic scale determination occurs globally. The original implementation in [1] uses global application
         (i.e., per_dimension=False), whereas dimensionwise is recommended if the data dimensions are heterogeneous.
     seed : int or None
-        The seed for the random number generator. If None, a random seed is used. Used instead of np.random.Generator
-        here to enable easy serialization.
+        The seed for the random number generator. If None, a random seed is used.
 
     Examples
     --------
-    >>> adapter = bf.Adapter().nnpe(["x"])
+    >>> nnpe_aug = bf.augmentations.NNPE(spike_scale=0.01, slab_scale=0.2, per_dimension=True, seed=42)
+    >>> dataset = bf.datasets.OnlineDataset(
+    ...     simulator=my_sim,
+    ...     batch_size=64,
+    ...     num_batches=100,
+    ...     adapter=None,
+    ...     augmentations={"data": nnpe_aug},
+    ... )
     """
 
     DEFAULT_SPIKE = 0.01
@@ -54,12 +55,40 @@ def __init__(
         seed: int | None = None,
     ):
         super().__init__()
+
         self.spike_scale = spike_scale
         self.slab_scale = slab_scale
         self.per_dimension = per_dimension
         self.seed = seed
         self.rng = np.random.default_rng(seed)
 
+    def __call__(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        """
+        Add spike‐and‐slab noise to `data` using automatic scale determination if not provided.
+        See “Notes” section of the class docstring for details).
+
+        Parameters
+        ----------
+        data : np.ndarray
+            Input array to be perturbed.
+        **kwargs
+            Unused keyword arguments.
+        """
+
+        # Check data validity
+        if not np.all(np.isfinite(data)):
+            raise ValueError("NNPE.forward: `data` contains NaN or infinite values.")
+
+        spike_scale = self._resolve_scale("spike_scale", self.spike_scale, self.DEFAULT_SPIKE, data)
+        slab_scale = self._resolve_scale("slab_scale", self.slab_scale, self.DEFAULT_SLAB, data)
+
+        # Apply spike-and-slab noise
+        mixture_mask = self.rng.binomial(n=1, p=0.5, size=data.shape).astype(bool)
+        noise_spike = self.rng.standard_normal(size=data.shape) * spike_scale
+        noise_slab = self.rng.standard_cauchy(size=data.shape) * slab_scale
+        noise = np.where(mixture_mask, noise_slab, noise_spike)
+        return data + noise
+
     def _resolve_scale(
         self,
         name: str,
@@ -118,44 +147,3 @@ def _resolve_scale(
                 if arr.ndim != 0:
                     raise ValueError(f"{name}: expected scalar, got array of shape {arr.shape}")
                 return arr
-
-    def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
-        """
-        Add spike‐and‐slab noise to `data` using automatic scale determination if not provided.
-        See “Notes” section of the class docstring for details).
-
-        Parameters
-        ----------
-        data : np.ndarray
-            Input array to be perturbed.
-        **kwargs
-            Unused keyword arguments.
-        """
-
-        # Check data validity
-        if not np.all(np.isfinite(data)):
-            raise ValueError("NNPE.forward: `data` contains NaN or infinite values.")
-
-        spike_scale = self._resolve_scale("spike_scale", self.spike_scale, self.DEFAULT_SPIKE, data)
-        slab_scale = self._resolve_scale("slab_scale", self.slab_scale, self.DEFAULT_SLAB, data)
-
-        # Apply spike-and-slab noise
-        mixture_mask = self.rng.binomial(n=1, p=0.5, size=data.shape).astype(bool)
-        noise_spike = self.rng.standard_normal(size=data.shape) * spike_scale
-        noise_slab = self.rng.standard_cauchy(size=data.shape) * slab_scale
-        noise = np.where(mixture_mask, noise_slab, noise_spike)
-        return data + noise
-
-    def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
-        # Non-invertible transform
-        return data
-
-    def get_config(self) -> dict:
-        return serialize(
-            {
-                "spike_scale": self.spike_scale,
-                "slab_scale": self.slab_scale,
-                "per_dimension": self.per_dimension,
-                "seed": self.seed,
-            }
-        )

tests/test_adapters/test_adapters.py

@@ -317,75 +317,3 @@ def test_nan_to_num():
     out = transform.forward(arr)
     np.testing.assert_array_equal(out["test-2d"], np.array([[1.0, 0.5], [0.5, 4.0]]))
     np.testing.assert_array_equal(out["new_mask_test-2d"], np.array([[1, 0], [0, 1]]))
-
-
-def test_nnpe(random_data):
-    # NNPE cannot be integrated into the adapter fixture and its tests since it modifies the input data
-    # and therefore breaks existing allclose checks
-    import numpy as np
-    from bayesflow.adapters import Adapter
-
-    # Test basic case with global noise application
-    ad = Adapter().nnpe("x1", spike_scale=1.0, slab_scale=1.0, per_dimension=False, seed=42)
-    result_training = ad(random_data)
-    result_inversed = ad(random_data, inverse=True)
-    serialized = serialize(ad)
-    deserialized = deserialize(serialized)
-    reserialized = serialize(deserialized)
-
-    assert keras.tree.lists_to_tuples(serialized) == keras.tree.lists_to_tuples(reserialized)
-
-    # check that only x1 is changed
-    assert "x1" in result_training
-    assert not np.allclose(result_training["x1"], random_data["x1"])
-
-    # all other keys are untouched
-    for k, v in random_data.items():
-        if k == "x1":
-            continue
-        assert np.allclose(result_training[k], v)
-
-    # check that the validation and inference data as well as inversed results are unchanged
-    for k, v in random_data.items():
-        assert np.allclose(result_inversed[k], v)
-
-    # Test both scales and seed are None case (automatic scale determination) with dimensionwise noise application
-    ad_auto = Adapter().nnpe("y1", slab_scale=None, spike_scale=None, per_dimension=True, seed=None)
-    result_training_auto = ad_auto(random_data)
-    assert not np.allclose(result_training_auto["y1"], random_data["y1"])
-    for k, v in random_data.items():
-        if k == "y1":
-            continue
-        assert np.allclose(result_training_auto[k], v)
-
-    serialized_auto = serialize(ad_auto)
-    deserialized_auto = deserialize(serialized_auto)
-    reserialized_auto = serialize(deserialized_auto)
-    assert keras.tree.lists_to_tuples(serialized_auto) == keras.tree.lists_to_tuples(serialize(reserialized_auto))
-
-    # Test dimensionwise versus global noise application (per_dimension=True vs per_dimension=False)
-    # Create data with second dimension having higher variance
-    data_shape = (32, 16, 1)
-    rng = np.random.default_rng(42)
-    zero = np.ones(shape=data_shape)
-    high = rng.normal(0, 100.0, size=data_shape)
-    var_data = {"x": np.concatenate([zero, high], axis=-1)}
-
-    # Apply dimensionwise and global adapters with automatic slab_scale scale determination
-    ad_partial_global = Adapter().nnpe("x", spike_scale=0, slab_scale=None, per_dimension=False, seed=42)
-    ad_partial_dim = Adapter().nnpe("x", spike_scale=[0, 1], slab_scale=None, per_dimension=True, seed=42)
-    res_dim = ad_partial_dim(var_data)
-    res_glob = ad_partial_global(var_data)
-
-    # Compute standard deviations of noise per last axis dimension
-    noise_dim = res_dim["x"] - var_data["x"]
-    noise_glob = res_glob["x"] - var_data["x"]
-    std_dim = np.std(noise_dim, axis=(0, 1))
-    std_glob = np.std(noise_glob, axis=(0, 1))
-
-    # Dimensionwise should assign zero noise, global some noise to zero-variance dimension
-    assert std_dim[0] == 0
-    assert std_glob[0] > 0
-    # Both should assign noise to high-variance dimension
-    assert std_dim[1] > 0
-    assert std_glob[1] > 0

tests/test_augmentations/__init__.py

@@ -0,0 +1,10 @@
+import numpy as np
+import pytest
+
+
+@pytest.fixture()
+def random_data():
+    return {
+        "x1": np.random.standard_normal(size=(4, 1)),
+        "x2": np.random.standard_normal(size=(8, 10, 1)),
+    }

tests/test_augmentations/conftest.py

@@ -0,0 +1,44 @@
+def test_nnpe(random_data):
+    import numpy as np
+    from bayesflow.augmentations import NNPE
+
+    # Test basic case with global noise application
+    aug = NNPE(spike_scale=1.0, slab_scale=1.0, per_dimension=False, seed=42)
+    result = aug(random_data["x1"])
+
+    # Check that result is the same shape as x1, but changed
+    assert result.shape == random_data["x1"].shape
+    assert not np.allclose(result, random_data["x1"])
+
+    # Test both scales and seed are None case (automatic scale determination) with dimensionwise noise application
+    aug_auto = NNPE(slab_scale=None, spike_scale=None, per_dimension=True, seed=None)
+    result_auto = aug_auto(random_data["x2"])
+    assert result_auto.shape == random_data["x2"].shape
+    assert not np.allclose(result_auto, random_data["x2"])
+
+    # Test dimensionwise versus global noise application (per_dimension=True vs per_dimension=False)
+    # Create data with second dimension having higher variance
+    data_shape = (32, 16, 1)
+    rng = np.random.default_rng(42)
+    zero = np.ones(shape=data_shape)
+    high = rng.normal(0, 100.0, size=data_shape)
+    var_data = np.concatenate([zero, high], axis=-1)
+
+    # Apply dimensionwise and global adapters with automatic slab_scale scale determination
+    aug_partial_global = NNPE(spike_scale=0, slab_scale=None, per_dimension=False, seed=42)
+    aug_partial_dim = NNPE(spike_scale=[0, 1], slab_scale=None, per_dimension=True, seed=42)
+    result_dim = aug_partial_dim(var_data)
+    result_glob = aug_partial_global(var_data)
+
+    # Compute standard deviations of noise per last axis dimension
+    noise_dim = result_dim - var_data
+    noise_glob = result_glob - var_data
+    std_dim = np.std(noise_dim, axis=(0, 1))
+    std_glob = np.std(noise_glob, axis=(0, 1))
+
+    # Dimensionwise should assign zero noise, global some noise to zero-variance dimension
+    assert std_dim[0] == 0
+    assert std_glob[0] > 0
+    # Both should assign noise to high-variance dimension
+    assert std_dim[1] > 0
+    assert std_glob[1] > 0