improve

Author: arrjon

bayesflow/adapters/transforms/nan_to_num.py

@@ -39,11 +39,11 @@ 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"mask_{self.key}" if self.key else "mask"
 
     def forward(self, data: dict[str, any], **kwargs) -> dict[str, any]:
         """
-        Forward transform: fill NaNs and optionally output mask under '_mask_<key>'.
+        Forward transform: fill NaNs and optionally output mask under 'mask_<key>'.
         """
         data = data.copy()
 
@@ -63,7 +63,7 @@ def forward(self, data: dict[str, any], **kwargs) -> dict[str, any]:
 
     def inverse(self, data: dict[str, any], **kwargs) -> dict[str, any]:
         """
-        Inverse transform: restore NaNs using the mask under '_mask_<key>'.
+        Inverse transform: restore NaNs using the mask under 'mask_<key>'.
         """
         data = data.copy()
 

tests/test_adapters/test_adapters.py

@@ -303,10 +303,10 @@ def test_nan_to_num():
     transform = bf.Adapter().nan_to_num(keys="test", default_value=0.0, encode_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]))
+    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)
     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["mask_test-2d"], np.array([[1, 0], [0, 1]]))