add mask naming

Author: arrjon

bayesflow/adapters/adapter.py

@@ -925,6 +925,7 @@ def nan_to_num(
         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.
@@ -934,13 +935,18 @@ def nan_to_num(
         keys : str or sequence of str
             The names of the variables to clean / mask.
         default_value : float
-            Value to substitute wherever data is NaN.
+            Value to substitute wherever data is NaN. Defaults to 0.0.
         return_mask : bool
-            If True, encode a binary missingness mask alongside the data.
+            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))
+            self.transforms.append(
+                NanToNum(key=key, default_value=default_value, return_mask=return_mask, mask_prefix=mask_prefix)
+            )
         return self

bayesflow/adapters/transforms/nan_to_num.py

@@ -17,20 +17,24 @@ class NanToNum(Transform):
         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):
+    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,
             }
         )
 
@@ -39,14 +43,20 @@ def mask_key(self) -> str:
         """
         Key under which the mask will be stored in the output dictionary.
         """
-        return f"mask_{self.key}" if self.key else "mask"
+        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)

tests/test_adapters/test_adapters.py

@@ -313,7 +313,7 @@ def test_nan_to_num():
     np.testing.assert_array_equal(out["mask_test"], np.array([1.0, 0.0, 1.0]))
 
     # test two-d array
-    transform = bf.Adapter().nan_to_num(keys="test-2d", default_value=0.5, return_mask=True)
+    transform = bf.Adapter().nan_to_num(keys="test-2d", default_value=0.5, return_mask=True, mask_prefix="new_mask")
     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["mask_test-2d"], np.array([[1, 0], [0, 1]]))
+    np.testing.assert_array_equal(out["new_mask_test-2d"], np.array([[1, 0], [0, 1]]))