update test

Author: arrjon

bayesflow/adapters/transforms/nan_to_num.py

@@ -45,6 +45,8 @@ 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()
+
         # 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)
@@ -63,6 +65,8 @@ 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]
 

tests/test_adapters/test_adapters.py

@@ -291,23 +291,22 @@ def test_log_det_jac_exceptions(random_data):
     assert np.allclose(forward_log_det_jac["p"], -inverse_log_det_jac)
 
 
-def test_replace_nan():
-    arr = {"test": np.array([1.0, np.nan, 3.0]), "test-2d": np.array([[1.0, np.nan], [np.nan, 4.0]])}
+def test_nan_to_num():
+    arr = {"test": np.array([1.0, np.nan, 3.0])}
     # test without mask
-    transform = bf.Adapter().replace_nan(keys="test", default_value=-1.0, encode_mask=False)
+    transform = bf.Adapter().nan_to_num(keys="test", default_value=-1.0, encode_mask=False)
     out = transform.forward(arr)["test"]
     np.testing.assert_array_equal(out, np.array([1.0, -1.0, 3.0]))
 
     # test with mask
-    transform = bf.Adapter().replace_nan(keys="test", default_value=0.0, encode_mask=True)
-    out = transform.forward(arr)["test"]
-    np.testing.assert_array_equal(out, np.array([[1.0, 1.0], [0.0, 0.0], [3.0, 1.0]]))
+    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)
+    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().replace_nan(keys="test-2d", default_value=0.5, encode_mask=True, axis=0)
-    out = transform.forward(arr)["test-2d"]
-    # Original shape (2,2) -> new shape (2,2,2) when expanding at axis=0
-    # Channel 0 along axis 0 should be the filled values
-    np.testing.assert_array_equal(out[0], np.array([[1.0, 0.5], [0.5, 4.0]]))
-    # Channel 1 along axis 0 should be the mask
-    np.testing.assert_array_equal(out[1], np.array([[1, 0], [0, 1]]))
+    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]]))