Merge remote-tracking branch 'upstream/dev' into standardize_in_approx

Author: vpratz

README.md

@@ -49,39 +49,6 @@ neural networks for parameter estimation, model comparison, and model validation
 when working with intractable simulators whose behavior as a whole is too
 complex to be described analytically.
 
-## Getting Started
-
-Using the high-level interface is easy, as demonstrated by the minimal working example below:
-
-```python
-import bayesflow as bf
-
-workflow = bf.BasicWorkflow(
-    inference_network=bf.networks.CouplingFlow(),
-    summary_network=bf.networks.TimeSeriesNetwork(),
-    inference_variables=["parameters"],
-    summary_variables=["observables"],
-    simulator=bf.simulators.SIR()
-)
-
-history = workflow.fit_online(epochs=15, batch_size=32, num_batches_per_epoch=200)
-
-diagnostics = workflow.plot_default_diagnostics(test_data=300)
-```
-
-For an in-depth exposition, check out our walkthrough notebooks below.
-
-1. [Linear regression starter example](examples/Linear_Regression_Starter.ipynb)
-2. [From ABC to BayesFlow](examples/From_ABC_to_BayesFlow.ipynb)
-3. [Two moons starter example](examples/Two_Moons_Starter.ipynb)
-4. [Rapid iteration with point estimators](examples/Lotka_Volterra_Point_Estimation_and_Expert_Stats.ipynb)
-5. [SIR model with custom summary network](examples/SIR_Posterior_Estimation.ipynb)
-6. [Bayesian experimental design](examples/Bayesian_Experimental_Design.ipynb)
-7. [Simple model comparison example](examples/One_Sample_TTest.ipynb)
-8. [Moving from BayesFlow v1.1 to v2.0](examples/From_BayesFlow_1.1_to_2.0.ipynb)
-
-More tutorials are always welcome! Please consider making a pull request if you have a cool application that you want to contribute.
-
 ## Install
 
 You can install the latest stable version from PyPI using:
@@ -132,9 +99,46 @@ export KERAS_BACKEND=jax
 
 This way, you also don't have to manually set the backend every time you are starting Python to use BayesFlow.
 
-**Caution:** Some development environments (e.g., VSCode or PyCharm) can silently overwrite environment variables. If you have set your backend as an environment variable and you still get keras-related import errors when loading BayesFlow, these IDE shenanigans might be the culprit. Try setting the keras backend in your Python script via `import os; os.environ["KERAS_BACKEND"] = "<YOUR-BACKEND>"`.
+## Getting Started
+
+Using the high-level interface is easy, as demonstrated by the minimal working example below:
+
+```python
+import bayesflow as bf
+
+workflow = bf.BasicWorkflow(
+    inference_network=bf.networks.CouplingFlow(),
+    summary_network=bf.networks.TimeSeriesNetwork(),
+    inference_variables=["parameters"],
+    summary_variables=["observables"],
+    simulator=bf.simulators.SIR()
+)
+
+history = workflow.fit_online(epochs=15, batch_size=32, num_batches_per_epoch=200)
+
+diagnostics = workflow.plot_default_diagnostics(test_data=300)
+```
+
+For an in-depth exposition, check out our expanding list of resources below.
+
+### Books
+
+Many examples from [Bayesian Cognitive Modeling: A Practical Course](https://bayesmodels.com/) by Lee & Wagenmakers (2013) in [BayesFlow](https://kucharssim.github.io/bayesflow-cognitive-modeling-book/).
+
+### Tutorial notebooks
+
+1. [Linear regression starter example](examples/Linear_Regression_Starter.ipynb)
+2. [From ABC to BayesFlow](examples/From_ABC_to_BayesFlow.ipynb)
+3. [Two moons starter example](examples/Two_Moons_Starter.ipynb)
+4. [Rapid iteration with point estimators](examples/Lotka_Volterra_Point_Estimation_and_Expert_Stats.ipynb)
+5. [SIR model with custom summary network](examples/SIR_Posterior_Estimation.ipynb)
+6. [Bayesian experimental design](examples/Bayesian_Experimental_Design.ipynb)
+7. [Simple model comparison example](examples/One_Sample_TTest.ipynb)
+8. [Moving from BayesFlow v1.1 to v2.0](examples/From_BayesFlow_1.1_to_2.0.ipynb)
+
+More tutorials are always welcome! Please consider making a pull request if you have a cool application that you want to contribute.
 
-### From Source
+## Contributing
 
 If you want to contribute to BayesFlow, we recommend installing it from source, see [CONTRIBUTING.md](CONTRIBUTING.md) for more details.
 

bayesflow/adapters/adapter.py

@@ -18,6 +18,7 @@
     Keep,
     Log,
     MapTransform,
+    NNPE,
     NumpyTransform,
     OneHot,
     Rename,
@@ -30,6 +31,7 @@
     Ungroup,
     RandomSubsample,
     Take,
+    NanToNum,
 )
 from .transforms.filter_transform import Predicate
 
@@ -699,6 +701,43 @@ 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.
 
@@ -918,3 +957,34 @@ def to_dict(self):
         transform = ToDict()
         self.transforms.append(transform)
         return self
+
+    def nan_to_num(
+        self,
+        keys: str | Sequence[str],
+        default_value: float = 0.0,
+        return_mask: bool = False,
+        mask_prefix: str = "mask",
+    ):
+        """
+        Append :py:class:`~bf.adapters.transforms.NanToNum` transform to the adapter.
+
+        Parameters
+        ----------
+        keys : str or sequence of str
+            The names of the variables to clean / mask.
+        default_value : float
+            Value to substitute wherever data is NaN. Defaults to 0.0.
+        return_mask : bool
+            If True, encode a binary missingness mask alongside the data. Defaults to False.
+        mask_prefix : str
+            Prefix for the mask key in the output dictionary. Defaults to 'mask_'. If the mask key already exists,
+            a ValueError is raised to avoid overwriting existing masks.
+        """
+        if isinstance(keys, str):
+            keys = [keys]
+
+        for key in keys:
+            self.transforms.append(
+                NanToNum(key=key, default_value=default_value, return_mask=return_mask, mask_prefix=mask_prefix)
+            )
+        return self

bayesflow/adapters/transforms/__init__.py

@@ -12,6 +12,7 @@
 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
@@ -28,6 +29,7 @@
 from .random_subsample import RandomSubsample
 from .take import Take
 from .ungroup import Ungroup
+from .nan_to_num import NanToNum
 
 from ...utils._docs import _add_imports_to_all
 

bayesflow/adapters/transforms/nan_to_num.py

@@ -0,0 +1,91 @@
+import numpy as np
+
+from bayesflow.utils.serialization import serializable, serialize
+from .transform import Transform
+
+
+@serializable("bayesflow.adapters")
+class NanToNum(Transform):
+    """
+    Replace NaNs with a default value, and optionally encode a missing-data mask as a separate output key.
+
+    This is based on "Missing data in amortized simulation-based neural posterior estimation" by Wang et al. (2024).
+
+    Parameters
+    ----------
+    default_value : float
+        Value to substitute wherever data is NaN.
+    return_mask : bool, default=False
+        If True, a mask array will be returned under a new key.
+    mask_prefix : str, default='mask_'
+        Prefix for the mask key in the output dictionary.
+    """
+
+    def __init__(self, key: str, default_value: float = 0.0, return_mask: bool = False, mask_prefix: str = "mask"):
+        super().__init__()
+        self.key = key
+        self.default_value = default_value
+        self.return_mask = return_mask
+        self.mask_prefix = mask_prefix
+
+    def get_config(self) -> dict:
+        return serialize(
+            {
+                "key": self.key,
+                "default_value": self.default_value,
+                "return_mask": self.return_mask,
+                "mask_prefix": self.mask_prefix,
+            }
+        )
+
+    @property
+    def mask_key(self) -> str:
+        """
+        Key under which the mask will be stored in the output dictionary.
+        """
+        return f"{self.mask_prefix}_{self.key}"
+
+    def forward(self, data: dict[str, any], **kwargs) -> dict[str, any]:
+        """
+        Forward transform: fill NaNs and optionally output mask under 'mask_<key>'.
+        """
+        data = data.copy()
+
+        # Check if the mask key already exists in the data
+        if self.mask_key in data.keys():
+            raise ValueError(
+                f"Mask key '{self.mask_key}' already exists in the data. Please choose a different mask_prefix."
+            )
+
+        # Identify NaNs and fill with default value
+        mask = np.isnan(data[self.key])
+        data[self.key] = np.nan_to_num(data[self.key], copy=False, nan=self.default_value)
+
+        if not self.return_mask:
+            return data
+
+        # Prepare mask array (1 for valid, 0 for NaN)
+        mask_array = (~mask).astype(np.int8)
+
+        # Return both the filled data and the mask under separate keys
+        data[self.mask_key] = mask_array
+        return data
+
+    def inverse(self, data: dict[str, any], **kwargs) -> dict[str, any]:
+        """
+        Inverse transform: restore NaNs using the mask under 'mask_<key>'.
+        """
+        data = data.copy()
+
+        # Retrieve mask and values to reconstruct NaNs
+        values = data[self.key]
+
+        if not self.return_mask:
+            values[values == self.default_value] = np.nan  # we assume default_value is not in data
+        else:
+            mask_array = data[self.mask_key].astype(bool)
+            # Put NaNs where mask is 0
+            values[~mask_array] = np.nan
+
+        data[self.key] = values
+        return data