Adjust class name and add docstring to forward method

Author: elseml

bayesflow/adapters/adapter.py

@@ -18,7 +18,7 @@
     Keep,
     Log,
     MapTransform,
-    Nnpe,
+    NNPE,
     NumpyTransform,
     OneHot,
     Rename,
@@ -708,7 +708,7 @@ def nnpe(
         spike_scale: float = 0.01,
         seed: int | None = None,
     ):
-        """Append an :py:class:`~transforms.Nnpe` transform to the adapter.
+        """Append an :py:class:`~transforms.NNPE` transform to the adapter.
 
         Parameters
         ----------
@@ -724,7 +724,7 @@ def nnpe(
         if isinstance(keys, str):
             keys = [keys]
 
-        transform = MapTransform({key: Nnpe(slab_scale=slab_scale, spike_scale=spike_scale, seed=seed) for key in keys})
+        transform = MapTransform({key: NNPE(slab_scale=slab_scale, spike_scale=spike_scale, seed=seed) for key in keys})
         self.transforms.append(transform)
         return self
 

bayesflow/adapters/transforms/__init__.py

@@ -12,7 +12,7 @@
 from .keep import Keep
 from .log import Log
 from .map_transform import MapTransform
-from .nnpe import Nnpe
+from .nnpe import NNPE
 from .numpy_transform import NumpyTransform
 from .one_hot import OneHot
 from .rename import Rename

bayesflow/adapters/transforms/nnpe.py

@@ -6,7 +6,7 @@
 
 
 @serializable("bayesflow.adapters")
-class Nnpe(ElementwiseTransform):
+class NNPE(ElementwiseTransform):
     """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).
@@ -48,6 +48,23 @@ def __init__(self, *, slab_scale: float = 0.25, spike_scale: float = 0.01, seed:
         self.rng = np.random.default_rng(seed)
 
     def forward(self, data: np.ndarray, stage: str = "inference", **kwargs) -> np.ndarray:
+        """
+        Add spike‐and‐slab noise (see “Notes” section of the class docstring for details) to `data` during training.
+
+        Parameters
+        ----------
+        data : np.ndarray
+            Input array to be perturbed.
+        stage : str, default='inference'
+            If 'training', noise is added; else data is returned unchanged.
+        **kwargs
+            Unused keyword arguments.
+
+        Returns
+        -------
+        np.ndarray
+            Noisy data when `stage` is 'training', otherwise the original input.
+        """
         if stage != "training":
             return data
         mixture_mask = self.rng.binomial(n=1, p=0.5, size=data.shape).astype(bool)
@@ -57,6 +74,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."""
         return data
 
     def get_config(self) -> dict: