Add weight threshold option for spatial averaging

tests/test_spatial.py

@@ -15,13 +15,13 @@ def setup(self):
         )
 
     def test__init__(self):
-        ds = self.ds.copy()
+        ds = self.ds.copy(deep=True)
         obj = SpatialAccessor(ds)
 
         assert obj._dataset.identical(ds)
 
     def test_decorator_call(self):
-        ds = self.ds.copy()
+        ds = self.ds.copy(deep=True)
         obj = ds.spatial
 
         assert obj._dataset.identical(ds)
@@ -50,7 +50,7 @@ def test_raises_error_if_axis_list_contains_unsupported_axis(self):
             self.ds.spatial.average("ts", axis=["Y", "incorrect_axis"])
 
     def test_raises_error_if_lat_axis_coords_cant_be_found(self):
-        ds = self.ds.copy()
+        ds = self.ds.copy(deep=True)
 
         # Update CF metadata to invalid values so cf_xarray can't interpret them.
         del ds.lat.attrs["axis"]
@@ -64,7 +64,7 @@ def test_raises_error_if_lat_axis_coords_cant_be_found(self):
             ds.spatial.average("ts", axis=["X", "Y"])
 
     def test_raises_error_if_lon_axis_coords_cant_be_found(self):
-        ds = self.ds.copy()
+        ds = self.ds.copy(deep=True)
 
         # Update CF metadata to invalid values so cf_xarray can't interpret them.
         del ds.lon.attrs["axis"]
@@ -141,13 +141,13 @@ def test_raises_error_if_weights_lat_and_lon_dims_dont_align_with_data_var_dims(
             self.ds.spatial.average("ts", axis=["X", "Y"], weights=weights)
 
     def test_spatial_average_for_lat_region_and_keep_weights(self):
-        ds = self.ds.copy()
+        ds = self.ds.copy(deep=True)
 
         result = ds.spatial.average(
             "ts", axis=["Y"], lat_bounds=(-5.0, 5), keep_weights=True
         )
 
-        expected = self.ds.copy()
+        expected = self.ds.copy(deep=True)
         expected["ts"] = xr.DataArray(
             data=np.array(
                 [[2.25, 2.25, 2.25, 2.25], [1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]]
@@ -164,12 +164,12 @@ def test_spatial_average_for_lat_region_and_keep_weights(self):
         xr.testing.assert_allclose(result, expected)
 
     def test_spatial_average_for_lat_region(self):
-        ds = self.ds.copy()
+        ds = self.ds.copy(deep=True)
 
         # Specifying axis as a str instead of list of str.
         result = ds.spatial.average("ts", axis=["Y"], lat_bounds=(-5.0, 5))
 
-        expected = self.ds.copy()
+        expected = self.ds.copy(deep=True)
         expected["ts"] = xr.DataArray(
             data=np.array(
                 [[2.25, 2.25, 2.25, 2.25], [1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]]
@@ -205,7 +205,7 @@ def test_spatial_average_for_lat_region_and_skipna(self):
     def test_spatial_average_for_domain_wrapping_p_meridian_non_cf_conventions(
         self,
     ):
-        ds = self.ds.copy()
+        ds = self.ds.copy(deep=True)
 
         # get spatial average for original dataset
         ref = ds.spatial.average("ts").ts
@@ -222,14 +222,14 @@ def test_spatial_average_for_domain_wrapping_p_meridian_non_cf_conventions(
 
     @requires_dask
     def test_spatial_average_for_lat_region_and_keep_weights_with_dask(self):
-        ds = self.ds.copy().chunk(2)
+        ds = self.ds.copy(deep=True).chunk(2)
 
         # Specifying axis as a str instead of list of str.
         result = ds.spatial.average(
             "ts", axis=["Y"], lat_bounds=(-5.0, 5), keep_weights=True
         )
 
-        expected = self.ds.copy()
+        expected = self.ds.copy(deep=True)
         expected["ts"] = xr.DataArray(
             data=np.array(
                 [[2.25, 2.25, 2.25, 2.25], [1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]]
@@ -246,7 +246,7 @@ def test_spatial_average_for_lat_region_and_keep_weights_with_dask(self):
         xr.testing.assert_allclose(result, expected)
 
     def test_spatial_average_for_lat_and_lon_region_and_keep_weights(self):
-        ds = self.ds.copy()
+        ds = self.ds.copy(deep=True)
         result = ds.spatial.average(
             "ts",
             axis=["X", "Y"],
@@ -255,7 +255,7 @@ def test_spatial_average_for_lat_and_lon_region_and_keep_weights(self):
             keep_weights=True,
         )
 
-        expected = self.ds.copy()
+        expected = self.ds.copy(deep=True)
         expected["ts"] = xr.DataArray(
             data=np.array([2.25, 1.0, 1.0]),
             coords={"time": expected.time},
@@ -277,7 +277,7 @@ def test_spatial_average_for_lat_and_lon_region_and_keep_weights(self):
         xr.testing.assert_allclose(result, expected)
 
     def test_spatial_average_for_lat_and_lon_region_with_custom_weights(self):
-        ds = self.ds.copy()
+        ds = self.ds.copy(deep=True)
 
         weights = xr.DataArray(
             data=np.array([[1, 2, 3, 4], [2, 4, 6, 8], [3, 6, 9, 12], [4, 8, 12, 16]]),
@@ -292,7 +292,7 @@ def test_spatial_average_for_lat_and_lon_region_with_custom_weights(self):
             data_var="ts",
         )
 
-        expected = self.ds.copy()
+        expected = self.ds.copy(deep=True)
         expected["ts"] = xr.DataArray(
             data=np.array([2.25, 1.0, 1.0]),
             coords={"time": expected.time},
@@ -302,6 +302,179 @@ def test_spatial_average_for_lat_and_lon_region_with_custom_weights(self):
         assert result.identical(expected)
 
 
+class TestAveragerMinWeight:
+    @pytest.fixture(autouse=True)
+    def setup(self):
+        self.ds = generate_dataset(
+            decode_times=True, cf_compliant=False, has_bounds=True
+        )
+
+        # Limit to just 3 data points to simplify testing.
+        self.ds = self.ds.isel(time=slice(None, 3))
+
+        # Change the value of the first element so that it is easier to identify
+        # changes in the output.
+        self.ds["ts"].data[0] = np.full((4, 4), 2.25)
+
+    def test_raises_error_if_min_weight_is_negative(self):
+        with pytest.raises(ValueError):
+            self.ds.spatial.average("ts", axis=["X", "Y"], min_weight=-0.1)
+
+    def test_raises_error_if_min_weight_is_greater_than_one(self):
+        with pytest.raises(ValueError):
+            self.ds.spatial.average("ts", axis=["X", "Y"], min_weight=1.1)
+
+    def test_spatial_average_with_min_weight_zero(self):
+        ds = self.ds.copy(deep=True)
+
+        # Insert NaN values into the dataset (no minimum required to compute value).
+        ds["ts"][0, :, 2] = np.nan
+
+        # min_weight=0.0 means no minimum weight threshold required to compute value
+        result = ds.spatial.average(
+            "ts",
+            axis=["X", "Y"],
+            min_weight=0.0,
+        )
+
+        expected = self.ds.copy(deep=True)
+        expected["ts"] = xr.DataArray(
+            data=np.array([2.25, 1.0, 1.0]),
+            coords={"time": expected.time},
+            dims="time",
+        )
+
+        xr.testing.assert_allclose(result, expected)
+
+    def test_spatial_average_with_min_weight_none_equivalent_to_zero(self):
+        ds = self.ds.copy(deep=True)
+
+        # Insert NaN values into the dataset.
+        ds["ts"][0, :, 2] = np.nan
+
+        # min_weight=None means no minimum weight threshold required to compute value
+        result = ds.spatial.average(
+            "ts",
+            axis=["X", "Y"],
+            min_weight=None,
+        )
+
+        expected = self.ds.copy(deep=True)
+        expected["ts"] = xr.DataArray(
+            data=np.array([2.25, 1.0, 1.0]),
+            coords={"time": expected.time},
+            dims="time",
+        )
+
+        xr.testing.assert_allclose(result, expected)
+
+    def test_spatial_average_with_min_weight_half(self):
+        ds = self.ds.copy(deep=True)
+
+        # Insert NaN values into the dataset > 50% at second time point.
+        ds["ts"][1, :, :] = np.nan
+
+        # At least 50% of the weights must be non-NaN to compute value.
+        result = ds.spatial.average(
+            "ts",
+            axis=["X", "Y"],
+            min_weight=0.5,
+        )
+
+        # The second grouping window will by NaN because the minimum weight
+        # threshold is not met (>50%).
+        expected = self.ds.copy(deep=True)
+        expected["ts"] = xr.DataArray(
+            data=np.array([2.25, np.nan, 1.0]),
+            coords={"time": expected.time},
+            dims="time",
+        )
+
+        xr.testing.assert_allclose(result, expected)
+
+    def test_spatial_average_with_min_weight_one(self):
+        ds = self.ds.copy(deep=True)
+
+        # Insert a single NaN value at the last time point.
+        ds["ts"][2, 0, 0] = np.nan
+
+        # 100% of the weights must be non-NaN to compute value.
+        result = ds.spatial.average(
+            "ts",
+            axis=["X", "Y"],
+            min_weight=1.0,
+        )
+
+        # The last grouping window will by NaN because the minimum weight
+        # threshold is not met (1 value is NaN out of 4).
+        expected = self.ds.copy(deep=True)
+        expected["ts"] = xr.DataArray(
+            data=np.array([2.25, 1.0, np.nan]),
+            coords={"time": expected.time},
+            dims="time",
+        )
+
+        xr.testing.assert_allclose(result, expected)
+
+    def test_spatial_average_with_min_weight_edge_case_zero_weights(self):
+        ds = self.ds.copy(deep=True)
+
+        # Set all weights to zero.
+        weights = xr.DataArray(
+            data=np.zeros((4, 4)),
+            coords={"lat": ds.lat, "lon": ds.lon},
+            dims=["lat", "lon"],
+        )
+
+        result = ds.spatial.average(
+            "ts",
+            axis=["X", "Y"],
+            weights=weights,
+            min_weight=0.5,
+        )
+
+        # With all weights set to zero, the results for all grouping windows
+        # should be NaN because the minimum weight threshold is not met (>50%).
+        expected = self.ds.copy(deep=True)
+        expected["ts"] = xr.DataArray(
+            data=np.array([np.nan, np.nan, np.nan]),
+            coords={"time": expected.time},
+            dims="time",
+        )
+
+        xr.testing.assert_allclose(result, expected)
+
+    def test_spatial_average_with_min_weight_edge_case_partial_nan_weights(self):
+        ds = self.ds.copy(deep=True)
+
+        # Insert NaN values into the weights.
+        weights = xr.DataArray(
+            data=np.array(
+                [[1, np.nan, 1, 1], [1, 1, np.nan, 1], [1, 1, 1, np.nan], [1, 1, 1, 1]]
+            ),
+            coords={"lat": ds.lat, "lon": ds.lon},
+            dims=["lat", "lon"],
+        )
+
+        result = ds.spatial.average(
+            "ts",
+            axis=["X", "Y"],
+            weights=weights,
+            min_weight=0.5,
+        )
+
+        # With partial NaN weights, the averages are still computed because
+        # at least 50% of the weights are non-NaN for each grouping window.
+        expected = self.ds.copy(deep=True)
+        expected["ts"] = xr.DataArray(
+            data=np.array([2.25, 1.0, 1.0]),
+            coords={"time": expected.time},
+            dims="time",
+        )
+
+        xr.testing.assert_allclose(result, expected)
+
+
 class TestGetWeights:
     @pytest.fixture(autouse=True)
     def setup(self):
@@ -322,7 +495,7 @@ def test_value_error_thrown_for_multiple_out_of_order_lon_bounds(self):
             self.ds.spatial._get_longitude_weights(domain_bounds, region_bounds=None)
 
     def test_raises_error_if_dataset_has_multiple_bounds_variables_for_an_axis(self):
-        ds = self.ds.copy()
+        ds = self.ds.copy(deep=True)
 
         # Create a second "Y" axis dimension and associated bounds
         ds["lat2"] = ds.lat.copy()
@@ -336,7 +509,7 @@ def test_raises_error_if_dataset_has_multiple_bounds_variables_for_an_axis(self)
             ds.spatial.get_weights(axis=["Y", "X"])
 
     def test_data_var_weights_for_region_in_lat_and_lon_domains(self):
-        ds = self.ds.copy()
+        ds = self.ds.copy(deep=True)
 
         result = ds.spatial.get_weights(
             axis=["Y", "X"], lat_bounds=(-5, 5), lon_bounds=(-170, -120), data_var="ts"
@@ -416,7 +589,7 @@ def test_weights_for_region_in_lon_domain(self):
     def test_dataset_weights_for_region_in_lon_domain_with_region_spanning_p_meridian(
         self,
     ):
-        ds = self.ds.copy()
+        ds = self.ds.copy(deep=True)
 
         result = ds.spatial._get_longitude_weights(
             domain_bounds=ds.lon_bnds,

tests/test_utils.py

@@ -1,7 +1,7 @@
 import pytest
 import xarray as xr
 
-from xcdat.utils import compare_datasets, str_to_bool
+from xcdat.utils import _validate_min_weight, compare_datasets, str_to_bool
 
 
 class TestCompareDatasets:
@@ -103,3 +103,23 @@ def test_raises_error_if_str_is_not_a_python_bool(self):
 
         with pytest.raises(ValueError):
             str_to_bool("1")
+
+
+class TestValidateMinWeight:
+    def test_pass_None_returns_0(self):
+        result = _validate_min_weight(None)
+
+        assert result == 0
+
+    def test_returns_error_if_less_than_0(self):
+        with pytest.raises(ValueError):
+            _validate_min_weight(-1)
+
+    def test_returns_error_if_greater_than_1(self):
+        with pytest.raises(ValueError):
+            _validate_min_weight(1.1)
+
+    def test_returns_valid_min_weight(self):
+        result = _validate_min_weight(1)
+
+        assert result == 1