changed name to return_mask

Author: arrjon

bayesflow/adapters/adapter.py

@@ -924,7 +924,7 @@ def nan_to_num(
         self,
         keys: str | Sequence[str],
         default_value: float = 0.0,
-        encode_mask: bool = False,
+        return_mask: bool = False,
     ):
         """
         Append :py:class:`~bf.adapters.transforms.NanToNum` transform to the adapter.
@@ -935,12 +935,12 @@ def nan_to_num(
             The names of the variables to clean / mask.
         default_value : float
             Value to substitute wherever data is NaN.
-        encode_mask : bool
+        return_mask : bool
             If True, encode a binary missingness mask alongside the data.
         """
         if isinstance(keys, str):
             keys = [keys]
 
         for key in keys:
-            self.transforms.append(NanToNum(key=key, default_value=default_value, encode_mask=encode_mask))
+            self.transforms.append(NanToNum(key=key, default_value=default_value, return_mask=return_mask))
         return self

bayesflow/adapters/transforms/nan_to_num.py

@@ -15,22 +15,22 @@ class NanToNum(Transform):
     ----------
     default_value : float
         Value to substitute wherever data is NaN.
-    encode_mask : bool, default=False
+    return_mask : bool, default=False
         If True, a mask array will be returned under a new key.
     """
 
-    def __init__(self, key: str, default_value: float = 0.0, encode_mask: bool = False):
+    def __init__(self, key: str, default_value: float = 0.0, return_mask: bool = False):
         super().__init__()
         self.key = key
         self.default_value = default_value
-        self.encode_mask = encode_mask
+        self.return_mask = return_mask
 
     def get_config(self) -> dict:
         return serialize(
             {
                 "key": self.key,
                 "default_value": self.default_value,
-                "encode_mask": self.encode_mask,
+                "return_mask": self.return_mask,
             }
         )
 
@@ -51,7 +51,7 @@ def forward(self, data: dict[str, any], **kwargs) -> dict[str, any]:
         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.encode_mask:
+        if not self.return_mask:
             return data
 
         # Prepare mask array (1 for valid, 0 for NaN)
@@ -70,7 +70,7 @@ def inverse(self, data: dict[str, any], **kwargs) -> dict[str, any]:
         # Retrieve mask and values to reconstruct NaNs
         values = data[self.key]
 
-        if not self.encode_mask:
+        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)

tests/test_adapters/test_adapters.py

@@ -301,19 +301,19 @@ def test_log_det_jac_exceptions(random_data):
 def test_nan_to_num():
     arr = {"test": np.array([1.0, np.nan, 3.0])}
     # test without mask
-    transform = bf.Adapter().nan_to_num(keys="test", default_value=-1.0, encode_mask=False)
+    transform = bf.Adapter().nan_to_num(keys="test", default_value=-1.0, return_mask=False)
     out = transform.forward(arr)["test"]
     np.testing.assert_array_equal(out, np.array([1.0, -1.0, 3.0]))
 
     # test with mask
     arr = {"test": np.array([1.0, np.nan, 3.0]), "test-2d": np.array([[1.0, np.nan], [np.nan, 4.0]])}
-    transform = bf.Adapter().nan_to_num(keys="test", default_value=0.0, encode_mask=True)
+    transform = bf.Adapter().nan_to_num(keys="test", default_value=0.0, return_mask=True)
     out = transform.forward(arr)
     np.testing.assert_array_equal(out["test"], np.array([1.0, 0.0, 3.0]))
     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, encode_mask=True)
+    transform = bf.Adapter().nan_to_num(keys="test-2d", default_value=0.5, return_mask=True)
     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]]))