From ce10f8b1662b9fe9b71e1b38b4554dcbbd17c819 Mon Sep 17 00:00:00 2001 From: Fred Bunt Date: Tue, 18 Apr 2023 13:34:47 -0600 Subject: [PATCH] Fix issue where ufunc ops had side effects on the input masks --- raster_tools/utils.py | 2 +- tests/test_raster.py | 38 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/raster_tools/utils.py b/raster_tools/utils.py index 00a3005..235830d 100644 --- a/raster_tools/utils.py +++ b/raster_tools/utils.py @@ -49,7 +49,7 @@ def merge_masks(masks): mask = None for m in masks: if mask is None: - mask = m + mask = m.copy() else: mask |= m return mask diff --git a/tests/test_raster.py b/tests/test_raster.py index 70b290c..6f3034f 100644 --- a/tests/test_raster.py +++ b/tests/test_raster.py @@ -1097,6 +1097,44 @@ def test_ufuncs_multiple_input_against_raster(ufunc): assert np.allclose(r._ds.mask, mask, equal_nan=True) +def test_ufunc_different_masks(): + r1 = arange_raster((3, 3)).set_null_value(0) + r1.mask[..., :2, :] = True + r2 = arange_raster((3, 3)).set_null_value(0) + r2.mask[..., :, :2] = True + r1_mask = np.array( + [ + [ + [1, 1, 1], + [1, 1, 1], + [0, 0, 0], + ] + ] + ) + r2_mask = np.array( + [ + [ + [1, 1, 0], + [1, 1, 0], + [1, 1, 0], + ] + ] + ) + result_mask = r1_mask | r2_mask + assert np.allclose(r1.mask.compute(), r1_mask) + assert np.allclose(r2.mask.compute(), r2_mask) + assert np.allclose( + result_mask, np.array([[1, 1, 1], [1, 1, 1], [1, 1, 0]]) + ) + + result = r1 * r2 + # Make sure that there where no side effects on the input raster masks + assert np.allclose(r1.mask.compute(), r1_mask) + assert np.allclose(r2.mask.compute(), r2_mask) + result_data = np.where(result_mask, result.null_value, 64) + assert np.allclose(result, result_data) + + def test_invert(): x = arange_nd((4, 5, 5)) rs = Raster(x)