diff --git a/tests/test_temporal.py b/tests/test_temporal.py index 0047c551..9c1b9d74 100644 --- a/tests/test_temporal.py +++ b/tests/test_temporal.py @@ -121,6 +121,7 @@ def test_averages_for_yearly_time_series(self): "mode": "average", "freq": "year", "weighted": "True", + "min_weight": 0.0, }, ) @@ -208,6 +209,7 @@ def test_averages_for_monthly_time_series(self): "mode": "average", "freq": "month", "weighted": "True", + "min_weight": 0.0, }, ) @@ -295,6 +297,7 @@ def test_averages_for_daily_time_series(self): "mode": "average", "freq": "day", "weighted": "True", + "min_weight": 0.0, }, ) @@ -380,6 +383,7 @@ def test_averages_for_hourly_time_series(self): "mode": "average", "freq": "hour", "weighted": "True", + "min_weight": 0.0, }, ) @@ -518,6 +522,7 @@ def test_weighted_annual_averages(self): "mode": "group_average", "freq": "year", "weighted": "True", + "min_weight": 0.0, }, ) @@ -569,6 +574,7 @@ def test_weighted_annual_averages_and_skipna(self): "mode": "group_average", "freq": "year", "weighted": "True", + "min_weight": 0.0, }, ) @@ -620,6 +626,7 @@ def test_weighted_annual_averages_with_chunking(self): "mode": "group_average", "freq": "year", "weighted": "True", + "min_weight": 0.0, }, ) @@ -627,7 +634,7 @@ def test_weighted_annual_averages_with_chunking(self): assert result.ts.attrs == expected.ts.attrs assert result.time.attrs == expected.time.attrs - def test_weighted_seasonal_averages_with_DJF_without_dropping_incomplete_seasons( + def test_weighted_seasonal_averages_with_DJF( self, ): ds = self.ds.copy() @@ -671,6 +678,7 @@ def test_weighted_seasonal_averages_with_DJF_without_dropping_incomplete_seasons "mode": "group_average", "freq": "season", "weighted": "True", + "min_weight": 0.0, "drop_incomplete_seasons": "False", "dec_mode": "DJF", }, @@ -678,6 +686,167 @@ def test_weighted_seasonal_averages_with_DJF_without_dropping_incomplete_seasons xr.testing.assert_identical(result, expected) + @pytest.mark.parametrize( + "min_weight, ts_data, expected_data", + [ + # min_weight=0.0: all missing, output is all np.nan + ( + 0.0, + np.array([[[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]]]), + np.array([[[np.nan]], [[np.nan]], [[np.nan]]]), + ), + # min_weight=0.0: all present, output is weighted mean + ( + 0.0, + np.array([[[1.0]], [[1.0]], [[2.0]], [[1.0]], [[1.0]]]), + np.array([[[1.318681]], [[1.0]], [[1.0]]]), + ), + # min_weight=0.25: (2000, DJF) meet threshold; + # (2000, MAM) below threshold (np.nan); + # (2000, JJA) meet threshold. + ( + 0.25, + np.array([[[np.nan]], [[1.0]], [[np.nan]], [[np.nan]], [[2.0]]]), + np.array([[[1.0]], [[np.nan]], [[2.0]]]), + ), + # min_weight=0.33: (2000, DJF), (2000, MAM), (2000, JJA) all meet threshold + ( + 0.33, + np.array([[[np.nan]], [[1.0]], [[np.nan]], [[1.0]], [[2.0]]]), + np.array([[[1.0]], [[1.0]], [[2.0]]]), + ), + # min_weight=0.33: (2000, DJF) below threshold (edge case, Feb has + # less weight); (2000, MAM), (2000, JJA) meet threshold + ( + 0.33, + np.array([[[np.nan]], [[np.nan]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[np.nan]], [[1.0]], [[2.0]]]), + ), + # min_weight=0.66: (2000, DJF), (2000, MAM), (2000, JJA) all meet threshold + ( + 0.66, + np.array([[[1.0]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[1.0]], [[1.0]], [[2.0]]]), + ), + # min_weight=1.0: all missing, output is all np.nan + ( + 1.0, + np.array([[[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]]]), + np.array([[[np.nan]], [[np.nan]], [[np.nan]]]), + ), + # min_weight=1.0: (2000, DJF) below threshold; + # (2000, MAM), (2000, JJA) all meet threshold + ( + 1.0, + np.array([[[np.nan]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[np.nan]], [[1.0]], [[2.0]]]), + ), + # min_weight=1.0: all meet threshold, output is weighted mean + ( + 1.0, + np.array([[[1.0]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[1]], [[1.0]], [[2.0]]]), + ), + ], + ) + def test_weighted_seasonal_averages_with_DJF_and_min_weight_threshold( + self, min_weight, ts_data, expected_data + ): + time = xr.DataArray( + data=np.array( + [ + "1999-12-16T00:00:00.000000000", # (2000, DJF) + "2000-01-16T12:00:00.000000000", # (2000, DJF) + "2000-02-15T12:00:00.000000000", # (2000, DJF) + "2000-03-16T12:00:00.000000000", # (2000, MAM) + "2000-06-16T00:00:00.000000000", # (2000, JJA) + ], + dtype="datetime64[ns]", + ), + dims=["time"], + attrs={ + "axis": "T", + "long_name": "time", + "standard_name": "time", + "bounds": "time_bnds", + }, + ) + time.encoding = {"calendar": "standard"} + time_bnds = xr.DataArray( + name="time_bnds", + data=np.array( + [ + ["1999-12-01T00:00:00.000000000", "2000-01-01T00:00:00.000000000"], + ["2000-01-01T00:00:00.000000000", "2000-02-01T00:00:00.000000000"], + ["2000-02-01T00:00:00.000000000", "2000-03-01T00:00:00.000000000"], + ["2000-03-01T00:00:00.000000000", "2000-04-01T00:00:00.000000000"], + ["2000-06-01T00:00:00.000000000", "2000-07-01T00:00:00.000000000"], + ], + dtype="datetime64[ns]", + ), + coords={"time": time}, + dims=["time", "bnds"], + attrs={"xcdat_bounds": "True"}, + ) + + ds = xr.Dataset( + data_vars={"time_bnds": time_bnds}, + coords={"lat": [-90], "lon": [0], "time": time}, + ) + + ds["ts"] = xr.DataArray( + data=ts_data, + coords={"time": time, "lat": ds.lat, "lon": ds.lon}, + dims=["time", "lat", "lon"], + attrs={"test_attr": "test"}, + ) + + result = ds.temporal.group_average( + "ts", + "season", + season_config={"dec_mode": "DJF"}, + min_weight=min_weight, + ) + expected = ds.copy() + expected = expected.drop_dims("time") + expected["ts"] = xr.DataArray( + name="ts", + data=expected_data, + coords={ + "lat": expected.lat, + "lon": expected.lon, + "time": xr.DataArray( + data=np.array( + [ + cftime.DatetimeGregorian(2000, 1, 1), + cftime.DatetimeGregorian(2000, 4, 1), + cftime.DatetimeGregorian(2000, 7, 1), + ], + ), + dims=["time"], + attrs={ + "axis": "T", + "long_name": "time", + "standard_name": "time", + "bounds": "time_bnds", + }, + ), + }, + dims=["time", "lat", "lon"], + attrs={ + "test_attr": "test", + "operation": "temporal_avg", + "mode": "group_average", + "freq": "season", + "weighted": "True", + "min_weight": 0.0, + "drop_incomplete_seasons": "False", + "dec_mode": "DJF", + }, + ) + + xr.testing.assert_allclose(result["ts"], expected["ts"]) + def test_weighted_seasonal_averages_with_DJF_and_drop_incomplete_seasons(self): ds = generate_dataset(decode_times=True, cf_compliant=True, has_bounds=True) @@ -719,6 +888,7 @@ def test_weighted_seasonal_averages_with_DJF_and_drop_incomplete_seasons(self): "mode": "group_average", "freq": "season", "weighted": "True", + "min_weight": 0.0, "drop_incomplete_seasons": "True", "dec_mode": "DJF", }, @@ -769,6 +939,7 @@ def test_weighted_seasonal_averages_with_DJF_and_drop_incomplete_djf(self): "mode": "group_average", "freq": "season", "weighted": "True", + "min_weight": 0.0, "dec_mode": "DJF", "drop_incomplete_djf": "True", }, @@ -802,17 +973,6 @@ def test_weighted_seasonal_averages_with_JFD(self): cftime.DatetimeGregorian(2001, 1, 1), ], ), - coords={ - "time": np.array( - [ - cftime.DatetimeGregorian(2000, 1, 1), - cftime.DatetimeGregorian(2000, 4, 1), - cftime.DatetimeGregorian(2000, 7, 1), - cftime.DatetimeGregorian(2000, 10, 1), - cftime.DatetimeGregorian(2001, 1, 1), - ], - ) - }, dims=["time"], attrs={ "axis": "T", @@ -829,6 +989,7 @@ def test_weighted_seasonal_averages_with_JFD(self): "mode": "group_average", "freq": "season", "weighted": "True", + "min_weight": 0.0, "drop_incomplete_seasons": "False", "dec_mode": "JFD", }, @@ -836,6 +997,169 @@ def test_weighted_seasonal_averages_with_JFD(self): xr.testing.assert_identical(result, expected) + @pytest.mark.parametrize( + "min_weight, ts_data, expected_data", + [ + # min_weight=0.0: all missing, output is all np.nan + ( + 0.0, + np.array([[[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]]]), + np.array([[[np.nan]], [[np.nan]], [[np.nan]]]), + ), + # min_weight=0.0: all present, output is weighted mean + ( + 0.0, + np.array([[[1.0]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[1.50413223]], [[1.0]], [[1.0]]]), + ), + # min_weight=0.25: (2000, JFD) meets threshold; + # (2000, MAM) below threshold (np.nan); + # (2000, JJA) meets threshold + ( + 0.25, + np.array([[[np.nan]], [[1.0]], [[np.nan]], [[1.0]], [[2.0]]]), + np.array([[[1.6777778]], [[np.nan]], [[1.0]]]), + ), + # min_weight=0.33: (2000, JFD), (2000, MAM), and (2000, JJA) all + # meet threshold + ( + 0.33, + np.array([[[np.nan]], [[np.nan]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[2.0]], [[1.0]], [[1.0]]]), + ), + # min_weight=0.66: (2000, JFD) below threshold (np.nan); + # (2000, MAM) and (2000, JJA) meet threshold + ( + 0.66, + np.array([[[np.nan]], [[np.nan]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[np.nan]], [[1.0]], [[1.0]]]), + ), + # min_weight=0.66: (2000, JFD), (2000, MAM), and (2000, JJA) all + # meet threshold + ( + 0.66, + np.array([[[np.nan]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[1.6777778]], [[1.0]], [[1.0]]]), + ), + # min_weight=1.0: all missing, output is all np.nan + ( + 1.0, + np.array([[[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]]]), + np.array([[[np.nan]], [[np.nan]], [[np.nan]]]), + ), + # min_weight=1.0: (2000, JFD) below threshold (np.nan); + # (2000, MAM) and (2000, JJA) meet threshold. + ( + 1.0, + np.array([[[np.nan]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[np.nan]], [[1.0]], [[1.0]]]), + ), + # min_weight=1.0: all present, output is weighted mean + ( + 1.0, + np.array([[[1.0]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[1.50413223]], [[1.0]], [[1.0]]]), + ), + ], + ) + def test_weighted_seasonal_averages_with_JFD_and_min_weight_threshold( + self, min_weight, ts_data, expected_data + ): + time = xr.DataArray( + data=np.array( + [ + "2000-01-16T12:00:00.000000000", # JFD + "2000-02-15T12:00:00.000000000", # JFD + "2000-03-16T12:00:00.000000000", # MAM + "2000-06-16T00:00:00.000000000", # JJA + "2000-12-16T00:00:00.000000000", # JFD + ], + dtype="datetime64[ns]", + ), + dims=["time"], + attrs={ + "axis": "T", + "long_name": "time", + "standard_name": "time", + "bounds": "time_bnds", + }, + ) + time.encoding = {"calendar": "standard"} + time_bnds = xr.DataArray( + name="time_bnds", + data=np.array( + [ + ["2000-01-01T00:00:00.000000000", "2000-02-01T00:00:00.000000000"], + ["2000-02-01T00:00:00.000000000", "2000-03-01T00:00:00.000000000"], + ["2000-03-01T00:00:00.000000000", "2000-04-01T00:00:00.000000000"], + ["2000-06-01T00:00:00.000000000", "2000-07-01T00:00:00.000000000"], + ["2000-11-01T00:00:00.000000000", "2001-01-01T00:00:00.000000000"], + ], + dtype="datetime64[ns]", + ), + coords={"time": time}, + dims=["time", "bnds"], + attrs={"xcdat_bounds": "True"}, + ) + + ds = xr.Dataset( + data_vars={"time_bnds": time_bnds}, + coords={"lat": [-90], "lon": [0], "time": time}, + ) + + ds["ts"] = xr.DataArray( + data=ts_data, + coords={"time": time, "lat": ds.lat, "lon": ds.lon}, + dims=["time", "lat", "lon"], + attrs={"test_attr": "test"}, + ) + + result = ds.temporal.group_average( + "ts", + "season", + season_config={"dec_mode": "JFD"}, + min_weight=min_weight, + ) + expected = ds.copy() + expected = expected.drop_dims("time") + expected["ts"] = xr.DataArray( + name="ts", + data=expected_data, + coords={ + "lat": expected.lat, + "lon": expected.lon, + "time": xr.DataArray( + data=np.array( + [ + cftime.DatetimeGregorian(2000, 1, 1), + cftime.DatetimeGregorian(2000, 4, 1), + cftime.DatetimeGregorian(2000, 7, 1), + ], + ), + dims=["time"], + attrs={ + "axis": "T", + "long_name": "time", + "standard_name": "time", + "bounds": "time_bnds", + }, + ), + }, + dims=["time", "lat", "lon"], + attrs={ + "test_attr": "test", + "operation": "temporal_avg", + "mode": "group_average", + "freq": "season", + "weighted": "True", + "min_weight": 0.0, + "drop_incomplete_seasons": "False", + "dec_mode": "JFD", + }, + ) + + xr.testing.assert_allclose(result["ts"], expected["ts"]) + def test_raises_error_with_incorrect_custom_seasons_argument(self): # Test raises error with non-3 letter strings with pytest.raises(ValueError): @@ -927,6 +1251,7 @@ def test_weighted_custom_seasonal_averages(self): "mode": "group_average", "freq": "season", "weighted": "True", + "min_weight": 0.0, "drop_incomplete_seasons": "False", "custom_seasons": [ "JanFebMar", @@ -993,6 +1318,7 @@ def test_weighted_seasonal_averages_with_custom_seasons_and_all_complete_seasons "mode": "group_average", "freq": "season", "weighted": "True", + "min_weight": 0.0, "drop_incomplete_seasons": "True", "custom_seasons": ["JanMarJun", "FebSep"], }, @@ -1050,6 +1376,7 @@ def test_weighted_custom_seasonal_averages_drops_incomplete_seasons(self): "mode": "group_average", "freq": "season", "weighted": "True", + "min_weight": 0.0, "drop_incomplete_seasons": "True", "custom_seasons": ["NovDec", "FebMarApr"], }, @@ -1107,6 +1434,7 @@ def test_weighted_custom_seasonal_averages_with_seasons_spanning_calendar_years( "mode": "group_average", "freq": "season", "weighted": "True", + "min_weight": 0.0, "drop_incomplete_seasons": "False", "custom_seasons": ["NovDecJanFebMar"], }, @@ -1153,6 +1481,7 @@ def test_weighted_monthly_averages(self): "mode": "group_average", "freq": "month", "weighted": "True", + "min_weight": 0.0, }, ) @@ -1204,11 +1533,159 @@ def test_weighted_monthly_averages_with_masked_data(self): "mode": "group_average", "freq": "month", "weighted": "True", + "min_weight": 0.0, }, ) xr.testing.assert_identical(result, expected) + @pytest.mark.parametrize( + "min_weight, ts_data, expected_data", + [ + # min_weight=0.0, all missing, output is all np.nan + ( + 0.0, + np.array( + [ + [[np.nan]], + [[np.nan]], + [[np.nan]], + [[np.nan]], + [[np.nan]], + [[np.nan]], + ] + ), + np.array([[[np.nan]], [[np.nan]], [[np.nan]]]), + ), + # min_weight=0.0, all months meet threshold, output is weighted mean. + ( + 0.0, + np.array([[[2.0]], [[1.0]], [[3.0]], [[4.0]], [[5.0]], [[6.0]]]), + np.array([[[1.5]], [[3.5]], [[5.5]]]), + ), + # min_weight=0.5, (2000, Jan) meets threshold; + # (2000, Feb) and (2000, Mar) below threshold (np.nan). + ( + 0.5, + np.array( + [ + [[1.0]], + [[np.nan]], + [[np.nan]], + [[np.nan]], + [[np.nan]], + [[np.nan]], + ] + ), + np.array([[[1.0]], [[np.nan]], [[np.nan]]]), + ), + # min_weight=0.5, (2000, Jan) and (2000, Feb) meet threshold; + # (2000, Mar) below threshold (np.nan). + ( + 0.5, + np.array( + [[[2.0]], [[np.nan]], [[3.0]], [[np.nan]], [[np.nan]], [[np.nan]]] + ), + np.array([[[2.0]], [[3.0]], [[np.nan]]]), + ), + # min_weight=1.0, (2000, Jan) meets threshold; + # (2000, Feb) and (2000, Mar) below threshold (np.nan). + ( + 1.0, + np.array( + [[[2.0]], [[1.0]], [[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]]] + ), + np.array([[[1.5]], [[np.nan]], [[np.nan]]]), + ), + # min_weight=1.0, all months meet threshold, output is weighted mean. + ( + 1.0, + np.array([[[2.0]], [[1.0]], [[3.0]], [[4.0]], [[5.0]], [[6.0]]]), + np.array([[[1.5]], [[3.5]], [[5.5]]]), + ), + ], + ) + def test_weighted_monthly_averages_min_weight_threshold_additional( + self, min_weight, ts_data, expected_data + ): + ds = xr.Dataset( + coords={ + "lat": [-90], + "lon": [0], + "time": xr.DataArray( + data=np.array( + [ + "2000-01-01T00:00:00.000000000", # (2000, Jan) + "2000-01-15T00:00:00.000000000", # (2000, Jan) + "2000-02-01T00:00:00.000000000", # (2000, Feb) + "2000-02-15T00:00:00.000000000", # (2000, Feb) + "2000-03-01T00:00:00.000000000", # (2000, Mar) + "2000-03-15T00:00:00.000000000", # (2000, Mar) + ], + dtype="datetime64[ns]", + ), + dims=["time"], + attrs={ + "axis": "T", + "long_name": "time", + "standard_name": "time", + "bounds": "time_bnds", + }, + ), + } + ) + ds.time.encoding = {"calendar": "standard"} + ds["time_bnds"] = xr.DataArray( + name="time_bnds", + data=np.array( + [ + ["2000-01-01T00:00:00.000000000", "2000-02-01T00:00:00.000000000"], + ["2000-01-01T00:00:00.000000000", "2000-02-01T00:00:00.000000000"], + ["2000-02-01T00:00:00.000000000", "2000-03-01T00:00:00.000000000"], + ["2000-02-01T00:00:00.000000000", "2000-03-01T00:00:00.000000000"], + ["2000-03-01T00:00:00.000000000", "2000-04-01T00:00:00.000000000"], + ["2000-03-01T00:00:00.000000000", "2000-04-01T00:00:00.000000000"], + ], + dtype="datetime64[ns]", + ), + coords={"time": ds.time}, + dims=["time", "bnds"], + attrs={"xcdat_bounds": "True"}, + ) + ds["ts"] = xr.DataArray( + data=ts_data, + coords={"lat": ds.lat, "lon": ds.lon, "time": ds.time}, + dims=["time", "lat", "lon"], + ) + result = ds.temporal.group_average("ts", "month", min_weight=min_weight) + expected = ds.copy().drop_dims("time") + expected["ts"] = xr.DataArray( + name="ts", + data=expected_data, + coords={ + "lat": expected.lat, + "lon": expected.lon, + "time": xr.DataArray( + data=np.array( + [ + cftime.DatetimeGregorian(2000, 1, 1), + cftime.DatetimeGregorian(2000, 2, 1), + cftime.DatetimeGregorian(2000, 3, 1), + ], + ), + dims=["time"], + attrs={ + "axis": "T", + "long_name": "time", + "standard_name": "time", + "bounds": "time_bnds", + }, + ), + }, + dims=["time", "lat", "lon"], + ) + xr.testing.assert_equal(result["ts"], expected["ts"]) + def test_weighted_daily_averages(self): ds = self.ds.copy() @@ -1247,6 +1724,7 @@ def test_weighted_daily_averages(self): "mode": "group_average", "freq": "day", "weighted": "True", + "min_weight": 0.0, }, ) @@ -1291,6 +1769,7 @@ def test_weighted_hourly_averages(self): "mode": "group_average", "freq": "hour", "weighted": "True", + "min_weight": 0.0, }, ) @@ -1383,6 +1862,7 @@ def test_subsets_climatology_based_on_reference_period(self): "mode": "climatology", "freq": "season", "weighted": "True", + "min_weight": 0.0, "dec_mode": "DJF", "drop_incomplete_seasons": "True", }, @@ -1437,6 +1917,7 @@ def test_weighted_seasonal_climatology_with_DJF(self): "mode": "climatology", "freq": "season", "weighted": "True", + "min_weight": 0.0, "drop_incomplete_seasons": "True", "dec_mode": "DJF", }, @@ -1444,6 +1925,169 @@ def test_weighted_seasonal_climatology_with_DJF(self): xr.testing.assert_identical(result, expected) + @pytest.mark.parametrize( + "min_weight, ts_data, expected_data", + [ + # min_weight=0.0: all missing (np.nan), output is all np.nan + ( + 0.0, + np.array([[[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]]]), + np.array([[[np.nan]], [[np.nan]], [[np.nan]]]), + ), + # min_weight=0.0: all present, output is weighted mean. + ( + 0.0, + np.array([[[1.0]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[1.504132]], [[1.0]], [[1.0]]]), + ), + # min_weight=0.25: DJF, MAM, and JJA all meet threshold. + ( + 0.25, + np.array([[[np.nan]], [[np.nan]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[2.0]], [[1.0]], [[1.0]]]), + ), + # min_weight=0.33: DJF, MAM, and JJA all meet threshold. + ( + 0.33, + np.array([[[np.nan]], [[np.nan]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[2.0]], [[1.0]], [[1.0]]]), + ), + # min_weight=0.66: DJF below threshold (np.nan); + # MAM, and JJA meet threshold. + ( + 0.66, + np.array([[[np.nan]], [[np.nan]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[np.nan]], [[1.0]], [[1.0]]]), + ), + # min_weight=0.66: JFD, MAM, and JJA all meet threshold. + ( + 0.66, + np.array([[[np.nan]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[1.677778]], [[1.0]], [[1.0]]]), + ), + # min_weight=1.0: all missing, output is all np.nan. + ( + 1.0, + np.array([[[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]]]), + np.array([[[np.nan]], [[np.nan]], [[np.nan]]]), + ), + # min_weight=1.0: DJF below threshold (np.nan); + # MAM, and JJA meet threshold. + ( + 1.0, + np.array([[[np.nan]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[np.nan]], [[1.0]], [[1.0]]]), + ), + # min_weight=1.0: All values present, output is weighted mean. + ( + 1.0, + np.array([[[1.0]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[1.504132]], [[1.0]], [[1.0]]]), + ), + ], + ) + def test_weighted_seasonal_climatology_with_DJF_and_min_weight_threshold( + self, min_weight, ts_data, expected_data + ): + time = xr.DataArray( + data=np.array( + [ + "2000-01-16T12:00:00.000000000", # DJF + "2000-02-15T12:00:00.000000000", # DJF + "2000-03-16T12:00:00.000000000", # MAM + "2000-06-16T00:00:00.000000000", # JJA + "2000-12-16T00:00:00.000000000", # DJF + ], + dtype="datetime64[ns]", + ), + dims=["time"], + attrs={ + "axis": "T", + "long_name": "time", + "standard_name": "time", + "bounds": "time_bnds", + }, + ) + time.encoding = {"calendar": "standard"} + time_bnds = xr.DataArray( + name="time_bnds", + data=np.array( + [ + ["2000-01-01T00:00:00.000000000", "2000-02-01T00:00:00.000000000"], + ["2000-02-01T00:00:00.000000000", "2000-03-01T00:00:00.000000000"], + ["2000-03-01T00:00:00.000000000", "2000-04-01T00:00:00.000000000"], + ["2000-06-01T00:00:00.000000000", "2000-07-01T00:00:00.000000000"], + ["2000-11-01T00:00:00.000000000", "2001-01-01T00:00:00.000000000"], + ], + dtype="datetime64[ns]", + ), + coords={"time": time}, + dims=["time", "bnds"], + attrs={"xcdat_bounds": "True"}, + ) + + ds = xr.Dataset( + data_vars={"time_bnds": time_bnds}, + coords={"lat": [-90], "lon": [0], "time": time}, + ) + + ds["ts"] = xr.DataArray( + data=ts_data, + coords={"time": time, "lat": ds.lat, "lon": ds.lon}, + dims=["time", "lat", "lon"], + attrs={"test_attr": "test"}, + ) + + result = ds.temporal.climatology( + "ts", + "season", + season_config={"dec_mode": "DJF"}, + min_weight=min_weight, + ) + expected = ds.copy() + expected = expected.drop_dims("time") + expected_time = xr.DataArray( + data=np.array( + [ + cftime.DatetimeGregorian(1, 1, 1), + cftime.DatetimeGregorian(1, 4, 1), + cftime.DatetimeGregorian(1, 7, 1), + ], + ), + coords={ + "time": np.array( + [ + cftime.DatetimeGregorian(1, 1, 1), + cftime.DatetimeGregorian(1, 4, 1), + cftime.DatetimeGregorian(1, 7, 1), + ], + ), + }, + attrs={ + "axis": "T", + "long_name": "time", + "standard_name": "time", + "bounds": "time_bnds", + }, + ) + expected["ts"] = xr.DataArray( + name="ts", + data=expected_data, + coords={"lat": expected.lat, "lon": expected.lon, "time": expected_time}, + dims=["time", "lat", "lon"], + attrs={ + "operation": "temporal_avg", + "mode": "climatology", + "freq": "season", + "weighted": "True", + "min_weight": min_weight, + "drop_incomplete_seasons": "False", + "dec_mode": "DJF", + }, + ) + + xr.testing.assert_allclose(result["ts"], expected["ts"]) + def test_raises_deprecation_warning_with_drop_incomplete_djf_season_config(self): # NOTE: This will test will also cover the other public APIs that # have drop_incomplete_djf as a season_config arg. @@ -1501,6 +2145,7 @@ def test_raises_deprecation_warning_with_drop_incomplete_djf_season_config(self) "mode": "climatology", "freq": "season", "weighted": "True", + "min_weight": 0.0, "drop_incomplete_djf": "True", "dec_mode": "DJF", }, @@ -1561,6 +2206,7 @@ def test_weighted_seasonal_climatology_with_DJF_and_skipna(self): "mode": "climatology", "freq": "season", "weighted": "True", + "min_weight": 0.0, "dec_mode": "DJF", "drop_incomplete_djf": "True", }, @@ -1575,9 +2221,62 @@ def test_chunked_weighted_seasonal_climatology_with_DJF(self): ds = self.ds.copy().chunk({"time": 2}) result = ds.temporal.climatology( - "ts", - "season", - season_config={"dec_mode": "DJF", "drop_incomplete_seasons": True}, + "ts", + "season", + season_config={"dec_mode": "DJF", "drop_incomplete_seasons": True}, + ) + + expected = ds.copy() + expected = expected.drop_dims("time") + expected_time = xr.DataArray( + data=np.array( + [ + cftime.DatetimeGregorian(1, 1, 1), + cftime.DatetimeGregorian(1, 4, 1), + cftime.DatetimeGregorian(1, 7, 1), + cftime.DatetimeGregorian(1, 10, 1), + ], + ), + coords={ + "time": np.array( + [ + cftime.DatetimeGregorian(1, 1, 1), + cftime.DatetimeGregorian(1, 4, 1), + cftime.DatetimeGregorian(1, 7, 1), + cftime.DatetimeGregorian(1, 10, 1), + ], + ), + }, + attrs={ + "axis": "T", + "long_name": "time", + "standard_name": "time", + "bounds": "time_bnds", + }, + ) + expected["ts"] = xr.DataArray( + name="ts", + data=np.ones((4, 4, 4)), + coords={"lat": expected.lat, "lon": expected.lon, "time": expected_time}, + dims=["time", "lat", "lon"], + attrs={ + "operation": "temporal_avg", + "mode": "climatology", + "freq": "season", + "weighted": "True", + "min_weight": 0.0, + "dec_mode": "DJF", + "drop_incomplete_seasons": "True", + }, + ) + + xr.testing.assert_identical(result, expected) + + def test_weighted_seasonal_climatology_with_JFD(self): + ds = self.ds.copy() + + result = ds.temporal.climatology( + "ts", "season", season_config={"dec_mode": "JFD"} ) expected = ds.copy() @@ -1618,20 +2317,133 @@ def test_chunked_weighted_seasonal_climatology_with_DJF(self): "mode": "climatology", "freq": "season", "weighted": "True", - "dec_mode": "DJF", - "drop_incomplete_seasons": "True", + "min_weight": 0.0, + "drop_incomplete_seasons": "False", + "dec_mode": "JFD", }, ) xr.testing.assert_identical(result, expected) - def test_weighted_seasonal_climatology_with_JFD(self): - ds = self.ds.copy() + @pytest.mark.parametrize( + "min_weight, ts_data, expected_data", + [ + # min_weight=0.0: all missing (np.nan), output is all np.nan + ( + 0.0, + np.array([[[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]]]), + np.array([[[np.nan]], [[np.nan]], [[np.nan]]]), + ), + # min_weight=0.0: all present, output is weighted mean. + ( + 0.0, + np.array([[[1.0]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[1.504132]], [[1.0]], [[1.0]]]), + ), + # min_weight=0.25: JFD, MAM, and JJA all meet threshold. + ( + 0.25, + np.array([[[np.nan]], [[np.nan]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[2.0]], [[1.0]], [[1.0]]]), + ), + # min_weight=0.33: JFD, MAM, and JJA all meet threshold. + ( + 0.33, + np.array([[[np.nan]], [[np.nan]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[2.0]], [[1.0]], [[1.0]]]), + ), + # min_weight=0.66: JFD below threshold (np.nan); + # MAM, and JJA meet threshold. + ( + 0.66, + np.array([[[np.nan]], [[np.nan]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[np.nan]], [[1.0]], [[1.0]]]), + ), + # min_weight=0.66: JFD, MAM, and JJA all meet threshold. + ( + 0.66, + np.array([[[np.nan]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[1.677778]], [[1.0]], [[1.0]]]), + ), + # min_weight=1.0: all missing, output is all np.nan. + ( + 1.0, + np.array([[[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]]]), + np.array([[[np.nan]], [[np.nan]], [[np.nan]]]), + ), + # min_weight=1.0: JFD below threshold (np.nan); + # MAM, and JJA meet threshold. + ( + 1.0, + np.array([[[np.nan]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[np.nan]], [[1.0]], [[1.0]]]), + ), + # min_weight=1.0: All values present, output is weighted mean. + ( + 1.0, + np.array([[[1.0]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[1.504132]], [[1.0]], [[1.0]]]), + ), + ], + ) + def test_weighted_seasonal_climatology_with_JFD_and_min_weight_threshold( + self, min_weight, ts_data, expected_data + ): + time = xr.DataArray( + data=np.array( + [ + "2000-01-16T12:00:00.000000000", # JFD + "2000-02-15T12:00:00.000000000", # JFD + "2000-03-16T12:00:00.000000000", # MAM + "2000-06-16T00:00:00.000000000", # JJA + "2000-12-16T00:00:00.000000000", # JFD + ], + dtype="datetime64[ns]", + ), + dims=["time"], + attrs={ + "axis": "T", + "long_name": "time", + "standard_name": "time", + "bounds": "time_bnds", + }, + ) + time.encoding = {"calendar": "standard"} + time_bnds = xr.DataArray( + name="time_bnds", + data=np.array( + [ + ["2000-01-01T00:00:00.000000000", "2000-02-01T00:00:00.000000000"], + ["2000-02-01T00:00:00.000000000", "2000-03-01T00:00:00.000000000"], + ["2000-03-01T00:00:00.000000000", "2000-04-01T00:00:00.000000000"], + ["2000-06-01T00:00:00.000000000", "2000-07-01T00:00:00.000000000"], + ["2000-11-01T00:00:00.000000000", "2001-01-01T00:00:00.000000000"], + ], + dtype="datetime64[ns]", + ), + coords={"time": time}, + dims=["time", "bnds"], + attrs={"xcdat_bounds": "True"}, + ) - result = ds.temporal.climatology( - "ts", "season", season_config={"dec_mode": "JFD"} + ds = xr.Dataset( + data_vars={"time_bnds": time_bnds}, + coords={"lat": [-90], "lon": [0], "time": time}, + ) + + ds["ts"] = xr.DataArray( + data=ts_data, + coords={"time": time, "lat": ds.lat, "lon": ds.lon}, + dims=["time", "lat", "lon"], + attrs={"test_attr": "test"}, ) + result = ds.temporal.climatology( + "ts", + "season", + season_config={"dec_mode": "JFD"}, + min_weight=min_weight, + ) expected = ds.copy() expected = expected.drop_dims("time") expected_time = xr.DataArray( @@ -1640,7 +2452,6 @@ def test_weighted_seasonal_climatology_with_JFD(self): cftime.DatetimeGregorian(1, 1, 1), cftime.DatetimeGregorian(1, 4, 1), cftime.DatetimeGregorian(1, 7, 1), - cftime.DatetimeGregorian(1, 10, 1), ], ), coords={ @@ -1649,7 +2460,6 @@ def test_weighted_seasonal_climatology_with_JFD(self): cftime.DatetimeGregorian(1, 1, 1), cftime.DatetimeGregorian(1, 4, 1), cftime.DatetimeGregorian(1, 7, 1), - cftime.DatetimeGregorian(1, 10, 1), ], ), }, @@ -1662,7 +2472,7 @@ def test_weighted_seasonal_climatology_with_JFD(self): ) expected["ts"] = xr.DataArray( name="ts", - data=np.ones((4, 4, 4)), + data=expected_data, coords={"lat": expected.lat, "lon": expected.lon, "time": expected_time}, dims=["time", "lat", "lon"], attrs={ @@ -1670,12 +2480,13 @@ def test_weighted_seasonal_climatology_with_JFD(self): "mode": "climatology", "freq": "season", "weighted": "True", + "min_weight": min_weight, "drop_incomplete_seasons": "False", "dec_mode": "JFD", }, ) - xr.testing.assert_identical(result, expected) + xr.testing.assert_allclose(result["ts"], expected["ts"]) def test_weighted_custom_seasonal_climatology(self): ds = self.ds.copy() @@ -1729,6 +2540,7 @@ def test_weighted_custom_seasonal_climatology(self): "mode": "climatology", "freq": "season", "weighted": "True", + "min_weight": 0.0, "drop_incomplete_seasons": "False", "custom_seasons": [ "JanFebMar", @@ -1785,6 +2597,7 @@ def test_weighted_custom_seasonal_climatology_with_seasons_spanning_calendar_yea "mode": "climatology", "freq": "season", "weighted": "True", + "min_weight": 0.0, "drop_incomplete_seasons": "False", "custom_seasons": ["NovDecJanFebMar"], }, @@ -1850,6 +2663,7 @@ def test_weighted_monthly_climatology(self): "mode": "climatology", "freq": "month", "weighted": "True", + "min_weight": 0.0, }, ) @@ -1974,6 +2788,7 @@ def test_weighted_daily_climatology(self): "mode": "climatology", "freq": "day", "weighted": "True", + "min_weight": 0.0, }, ) @@ -2065,6 +2880,7 @@ def test_weighted_daily_climatology_drops_leap_days_with_matching_calendar(self) "mode": "climatology", "freq": "day", "weighted": "True", + "min_weight": 0.0, }, ) @@ -2268,6 +3084,7 @@ def test_seasonal_departures_relative_to_climatology_reference_period(self): "mode": "departures", "freq": "season", "weighted": "True", + "min_weight": 0.0, "dec_mode": "DJF", "drop_incomplete_seasons": "False", }, @@ -2323,6 +3140,7 @@ def test_monthly_departures_relative_to_climatology_reference_period_with_same_o "mode": "departures", "freq": "month", "weighted": "True", + "min_weight": 0.0, }, ) expected["time_bnds"] = xr.DataArray( @@ -2404,6 +3222,7 @@ def test_weighted_seasonal_departures_with_DJF(self): "mode": "departures", "freq": "season", "weighted": "True", + "min_weight": 0.0, "dec_mode": "DJF", "drop_incomplete_seasons": "False", }, @@ -2411,6 +3230,162 @@ def test_weighted_seasonal_departures_with_DJF(self): xr.testing.assert_identical(result, expected) + @pytest.mark.parametrize( + "min_weight, ts_data, expected_data", + [ + # min_weight=0.0: all missing (np.nan), output is all np.nan + ( + 0.0, + np.array([[[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]]]), + np.array([[[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]]]), + ), + # min_weight=0.0: all present, output are departures from the mean. + ( + 0.0, + np.array([[[2.0]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[-0.340659]], [[0.0]], [[0.0]], [[0.659341]]]), + ), + # min_weight=0.25: (2000, DJF), (2000, MAM) and (2000, JJA) meet threshold; + # (2001, DJF) below threshold (np.nan); + ( + 0.25, + np.array([[[np.nan]], [[np.nan]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[0.0]], [[0.0]], [[0.0]], [[np.nan]]]), + ), + # min_weight=0.33: (2000, DJF) and (2001, DJF) below threshold (np.nan); + # (2000, MAM), (2000, JJA) meet threshold. + ( + 0.33, + np.array([[[np.nan]], [[np.nan]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[np.nan]], [[0.0]], [[0.0]], [[np.nan]]]), + ), + # min_weight=1.0: (2000, DJF) and (2001, DJF) below threshold (np.nan); + # (2000, MAM) and (2000, JJA) meet threshold. + ( + 0.66, + np.array([[[np.nan]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[np.nan]], [[0.0]], [[0.0]], [[np.nan]]]), + ), + # min_weight=1.0: all missing, output is all np.nan. + ( + 1.0, + np.array([[[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]]]), + np.array([[[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]]]), + ), + # min_weight=1.0: (2000, DJF) and (2001, DJF) below threshold (np.nan); + # (2000, MAM) and (2000, JJA) meet threshold. + ( + 1.0, + np.array([[[np.nan]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[np.nan]], [[0.0]], [[0.0]], [[np.nan]]]), + ), + # min_weight=1.0: All values present, output are departures from the mean. + ( + 1.0, + np.array([[[1.0]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[0.0]], [[0.0]], [[0.0]], [[0.0]]]), + ), + ], + ) + def test_weighted_seasonal_departures_with_DJF_and_min_weight_threshold( + self, min_weight, ts_data, expected_data + ): + time = xr.DataArray( + data=np.array( + [ + "1999-12-16T00:00:00.000000000", # DJF + "2000-01-16T12:00:00.000000000", # DJF + "2000-02-15T12:00:00.000000000", # DJF + "2000-03-16T12:00:00.000000000", # MAM + "2000-06-16T00:00:00.000000000", # JJA + ], + dtype="datetime64[ns]", + ), + dims=["time"], + attrs={ + "axis": "T", + "long_name": "time", + "standard_name": "time", + "bounds": "time_bnds", + }, + ) + time.encoding = {"calendar": "standard"} + time_bnds = xr.DataArray( + name="time_bnds", + data=np.array( + [ + ["1999-12-01T00:00:00.000000000", "2000-01-01T00:00:00.000000000"], + ["2000-01-01T00:00:00.000000000", "2000-02-01T00:00:00.000000000"], + ["2000-02-01T00:00:00.000000000", "2000-03-01T00:00:00.000000000"], + ["2000-03-01T00:00:00.000000000", "2000-04-01T00:00:00.000000000"], + ["2000-06-01T00:00:00.000000000", "2000-07-01T00:00:00.000000000"], + ], + dtype="datetime64[ns]", + ), + coords={"time": time}, + dims=["time", "bnds"], + attrs={"xcdat_bounds": "True"}, + ) + + ds = xr.Dataset( + data_vars={"time_bnds": time_bnds}, + coords={"lat": [-90], "lon": [0], "time": time}, + ) + + ds["ts"] = xr.DataArray( + data=ts_data, + coords={"time": time, "lat": ds.lat, "lon": ds.lon}, + dims=["time", "lat", "lon"], + attrs={"test_attr": "test"}, + ) + + result = ds.temporal.departures( + "ts", + "season", + season_config={"dec_mode": "DJF"}, + min_weight=min_weight, + ) + expected = ds.copy() + expected = expected.drop_dims("time") + expected["ts"] = xr.DataArray( + name="ts", + data=expected_data, + coords={ + "lat": expected.lat, + "lon": expected.lon, + "time": xr.DataArray( + data=np.array( + [ + cftime.DatetimeGregorian(2000, 1, 1), + cftime.DatetimeGregorian(2000, 4, 1), + cftime.DatetimeGregorian(2000, 7, 1), + cftime.DatetimeGregorian(2001, 1, 1), + ], + ), + dims=["time"], + attrs={ + "axis": "T", + "long_name": "time", + "standard_name": "time", + "bounds": "time_bnds", + }, + ), + }, + dims=["time", "lat", "lon"], + attrs={ + "test_attr": "test", + "operation": "temporal_avg", + "mode": "departures", + "freq": "season", + "weighted": "True", + "min_weight": min_weight, + "dec_mode": "DJF", + "drop_incomplete_seasons": "False", + }, + ) + + xr.testing.assert_allclose(result["ts"], expected["ts"]) + def test_weighted_seasonal_departures_with_DJF_and_skipna(self): ds = self.ds.copy(deep=True) @@ -2460,6 +3435,7 @@ def test_weighted_seasonal_departures_with_DJF_and_skipna(self): "mode": "departures", "freq": "season", "weighted": "True", + "min_weight": 0.0, "dec_mode": "DJF", "drop_incomplete_djf": "True", }, @@ -2512,6 +3488,7 @@ def test_weighted_seasonal_departures_with_DJF_and_keep_weights(self): "mode": "departures", "freq": "season", "weighted": "True", + "min_weight": 0.0, "dec_mode": "DJF", "drop_incomplete_seasons": "False", }, @@ -2648,6 +3625,159 @@ def test_unweighted_seasonal_departures_with_JFD(self): xr.testing.assert_identical(result, expected) + @pytest.mark.parametrize( + "min_weight, ts_data, expected_data", + [ + # min_weight=0.0: all missing (np.nan), output is all np.nan + ( + 0.0, + np.array([[[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]]]), + np.array([[[np.nan]], [[np.nan]], [[np.nan]]]), + ), + # min_weight=0.0: all present, output are departures from the mean. + ( + 0.0, + np.array([[[2.0]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[0.0]], [[0.0]], [[0.0]]]), + ), + # min_weight=0.25: (2000, DJF), (2000, MAM) and (2000, JJA) meet threshold; + ( + 0.25, + np.array([[[np.nan]], [[np.nan]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[0.0]], [[0.0]], [[0.0]]]), + ), + # min_weight=0.33: (2000, DJF), (2000, MAM), (2000, JJA) meet threshold. + ( + 0.33, + np.array([[[np.nan]], [[np.nan]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[0.0]], [[0.0]], [[0.0]]]), + ), + # min_weight=1.0: (2000, DJF) below threshold (np.nan); + # (2000, MAM) and (2000, JJA) meet threshold. + ( + 0.66, + np.array([[[np.nan]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[np.nan]], [[0.0]], [[0.0]]]), + ), + # min_weight=1.0: all missing, output is all np.nan. + ( + 1.0, + np.array([[[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]], [[np.nan]]]), + np.array([[[np.nan]], [[np.nan]], [[np.nan]]]), + ), + # min_weight=1.0: (2000, DJF) below threshold (np.nan); + # (2000, MAM) and (2000, JJA) meet threshold. + ( + 1.0, + np.array([[[np.nan]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[np.nan]], [[0.0]], [[0.0]]]), + ), + # min_weight=1.0: All values present, output are departures from the mean. + ( + 1.0, + np.array([[[1.0]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]), + np.array([[[0.0]], [[0.0]], [[0.0]]]), + ), + ], + ) + def test_weighted_seasonal_departures_with_JFD_and_min_weight_threshold( + self, min_weight, ts_data, expected_data + ): + time = xr.DataArray( + data=np.array( + [ + "2000-01-16T12:00:00.000000000", # DJF + "2000-02-15T12:00:00.000000000", # DJF + "2000-03-16T12:00:00.000000000", # MAM + "2000-06-16T00:00:00.000000000", # JJA + "2000-12-16T00:00:00.000000000", # DJF + ], + dtype="datetime64[ns]", + ), + dims=["time"], + attrs={ + "axis": "T", + "long_name": "time", + "standard_name": "time", + "bounds": "time_bnds", + }, + ) + time.encoding = {"calendar": "standard"} + time_bnds = xr.DataArray( + name="time_bnds", + data=np.array( + [ + ["2000-01-01T00:00:00.000000000", "2000-02-01T00:00:00.000000000"], + ["2000-02-01T00:00:00.000000000", "2000-03-01T00:00:00.000000000"], + ["2000-03-01T00:00:00.000000000", "2000-04-01T00:00:00.000000000"], + ["2000-06-01T00:00:00.000000000", "2000-07-01T00:00:00.000000000"], + ["2000-12-01T00:00:00.000000000", "2001-01-01T00:00:00.000000000"], + ], + dtype="datetime64[ns]", + ), + coords={"time": time}, + dims=["time", "bnds"], + attrs={"xcdat_bounds": "True"}, + ) + + ds = xr.Dataset( + data_vars={"time_bnds": time_bnds}, + coords={"lat": [-90], "lon": [0], "time": time}, + ) + + ds["ts"] = xr.DataArray( + data=ts_data, + coords={"time": time, "lat": ds.lat, "lon": ds.lon}, + dims=["time", "lat", "lon"], + attrs={"test_attr": "test"}, + ) + + result = ds.temporal.departures( + "ts", + "season", + season_config={"dec_mode": "JFD"}, + min_weight=min_weight, + ) + expected = ds.copy() + expected = expected.drop_dims("time") + expected["ts"] = xr.DataArray( + name="ts", + data=expected_data, + coords={ + "lat": expected.lat, + "lon": expected.lon, + "time": xr.DataArray( + data=np.array( + [ + cftime.DatetimeGregorian(2000, 1, 1), + cftime.DatetimeGregorian(2000, 4, 1), + cftime.DatetimeGregorian(2000, 7, 1), + ], + ), + dims=["time"], + attrs={ + "axis": "T", + "long_name": "time", + "standard_name": "time", + "bounds": "time_bnds", + }, + ), + }, + dims=["time", "lat", "lon"], + attrs={ + "test_attr": "test", + "operation": "temporal_avg", + "mode": "departures", + "freq": "season", + "weighted": "True", + "min_weight": min_weight, + "drop_incomplete_seasons": "False", + "dec_mode": "JFD", + }, + ) + + xr.testing.assert_allclose(result["ts"], expected["ts"]) + def test_weighted_daily_departures_drops_leap_days_with_matching_calendar(self): time = xr.DataArray( data=np.array( @@ -2734,6 +3864,7 @@ def test_weighted_daily_departures_drops_leap_days_with_matching_calendar(self): "mode": "departures", "freq": "day", "weighted": "True", + "min_weight": 0.0, }, ), }, diff --git a/xcdat/temporal.py b/xcdat/temporal.py index 14cbe2df..023e94b9 100644 --- a/xcdat/temporal.py +++ b/xcdat/temporal.py @@ -19,6 +19,7 @@ from xcdat._logger import _setup_custom_logger from xcdat.axis import get_dim_coords from xcdat.dataset import _get_data_var +from xcdat.utils import _get_masked_weights, _validate_min_weight logger = _setup_custom_logger(__name__) @@ -266,6 +267,7 @@ def group_average( keep_weights: bool = False, season_config: SeasonConfigInput = DEFAULT_SEASON_CONFIG, skipna: bool | None = None, + min_weight: float | None = None, ): """Returns a Dataset with average of a data variable by time group. @@ -367,6 +369,10 @@ def group_average( skips missing values for float dtypes; other dtypes either do not have a sentinel missing value (int) or ``skipna=True`` has not been implemented (object, datetime64 or timedelta64). + min_weight : float | None, optional + Fraction of data coverage (i.e., weight) needed to return a + temporal average value. Value must range from 0 to 1, by default + None (equivalent to ``min_weight=0.0``). Returns ------- @@ -446,6 +452,7 @@ def group_average( keep_weights=keep_weights, season_config=season_config, skipna=skipna, + min_weight=min_weight, ) def climatology( @@ -457,6 +464,7 @@ def climatology( reference_period: tuple[str, str] | None = None, season_config: SeasonConfigInput = DEFAULT_SEASON_CONFIG, skipna: bool | None = None, + min_weight: float | None = None, ): """Returns a Dataset with the climatology of a data variable. @@ -567,6 +575,10 @@ def climatology( skips missing values for float dtypes; other dtypes either do not have a sentinel missing value (int) or ``skipna=True`` has not been implemented (object, datetime64 or timedelta64). + min_weight : float | None, optional + Fraction of data coverage (i.e., weight) needed to return a + temporal average value. Value must range from 0 to 1, by default + None (equivalent to ``min_weight=0.0``). Returns ------- @@ -651,6 +663,7 @@ def climatology( reference_period, season_config, skipna, + min_weight=min_weight, ) def departures( @@ -662,6 +675,7 @@ def departures( reference_period: tuple[str, str] | None = None, season_config: SeasonConfigInput = DEFAULT_SEASON_CONFIG, skipna: bool | None = None, + min_weight: float | None = None, ) -> xr.Dataset: """ Returns a Dataset with the climatological departures (anomalies) for a @@ -783,6 +797,10 @@ def departures( skips missing values for float dtypes; other dtypes either do not have a sentinel missing value (int) or ``skipna=True`` has not been implemented (object, datetime64 or timedelta64). + min_weight : float | None, optional + Fraction of data coverage (i.e., weight) needed to return a + temporal average value. Value must range from 0 to 1, by default + None (equivalent to ``min_weight=0.0``). Returns ------- @@ -863,7 +881,13 @@ def departures( inferred_freq = _infer_freq(ds[self.dim]) if inferred_freq != freq: ds_obs = ds_obs.temporal.group_average( - data_var, freq, weighted, keep_weights, season_config, skipna + data_var, + freq, + weighted, + keep_weights, + season_config, + skipna, + min_weight, ) # 4. Calculate the climatology of the data variable. @@ -877,6 +901,7 @@ def departures( reference_period, season_config, skipna, + min_weight=min_weight, ) # 5. Calculate the departures for the data variable. @@ -899,10 +924,13 @@ def _averager( reference_period: tuple[str, str] | None = None, season_config: SeasonConfigInput = DEFAULT_SEASON_CONFIG, skipna: bool | None = None, + min_weight: float | None = None, ) -> xr.Dataset: """Averages a data variable based on the averaging mode and frequency.""" ds = self._dataset.copy() - self._set_arg_attrs(mode, freq, weighted, reference_period, season_config) + self._set_arg_attrs( + mode, freq, weighted, reference_period, season_config, min_weight + ) # Preprocess the dataset based on method argument values. ds = self._preprocess_dataset(ds) @@ -983,6 +1011,7 @@ def _set_arg_attrs( weighted: bool, reference_period: tuple[str, str] | None = None, season_config: SeasonConfigInput = DEFAULT_SEASON_CONFIG, + min_weight: float | None = None, ): """Validates method arguments and sets them as object attributes. @@ -998,6 +1027,10 @@ def _set_arg_attrs( A dictionary for "season" frequency configurations. If configs for predefined seasons are passed, configs for custom seasons are ignored and vice versa, by default DEFAULT_SEASON_CONFIG. + min_weight : float | None, optional + Fraction of data coverage (i.e., weight) needed to return a + temporal average value. Value must range from 0 to 1, by default + None (equivalent to ``min_weight=0.0``). Raises ------ @@ -1025,6 +1058,7 @@ def _set_arg_attrs( self._mode = mode self._freq = freq self._weighted = weighted + self._min_weight = _validate_min_weight(min_weight) self._reference_period = None if reference_period is not None: @@ -1541,53 +1575,119 @@ def _group_average( """ dv = _get_data_var(ds, data_var) - # Label the time coordinates for grouping weights and the data variable - # values. + # Label the time coordinates with groups for grouping data and weights. self._labeled_time = self._label_time_coords(dv[self.dim]) dv = dv.assign_coords({self.dim: self._labeled_time}) if self._weighted: - self._weights = self._get_weights(ds, data_var) + dv_avg = self._weighted_group_average(ds, dv, skipna) + else: + dv_avg = self._group_data(dv).mean(skipna=skipna) - # Weight the data variable. - dv *= self._weights + # After grouping and aggregating, xarray removes attributes from the + # grouped time coordinate. The `keep_attrs=True` option only preserves + # attributes for data variables, not coordinates. Therefore, we manually + # restore the time coordinate's attributes below. + dv_avg[self.dim].attrs = self._labeled_time.attrs + dv_avg[self.dim].encoding = self._labeled_time.encoding - # Ensure missing data (`np.nan`) receives no weight (zero). To - # achieve this, first broadcast the one-dimensional (temporal - # dimension) shape of the `weights` DataArray to the - # multi-dimensional shape of its corresponding data variable. - weights = self._weights + dv_avg = self._add_operation_attrs(dv_avg) + + return dv_avg + + def _weighted_group_average( + self, + ds: xr.Dataset, + dv: xr.DataArray, + skipna: bool | None, + ) -> xr.DataArray: + """Compute the weighted group average of a data variable. + + This method applies weights to the data variable, groups the weighted data, + and computes the average by dividing the sum of weighted data by the sum of + weights for non-missing data. It handles missing values according to the + `skipna` parameter and ensures that weights for missing data are excluded + from the denominator. Optionally, results are masked where the sum of weights + falls below a minimum threshold. + + Parameters + ---------- + ds : xr.Dataset + The input xarray Dataset containing the data variable and any coordinate information. + dv : xr.DataArray + The data variable to be averaged. + skipna : bool | None + If True, skip missing values (as marked by NaN). By default, only + skips missing values for float dtypes; other dtypes either do not + have a sentinel missing value (int) or ``skipna=True`` has not been + implemented (object, datetime64 or timedelta64). + + Returns + ------- + xr.DataArray + The weighted group average of the data variable, with the same name as `dv`. + Values are set to NaN where the sum of weights is below the minimum threshold. + + Notes + ----- + - Weights are masked to zero where data is missing. + - For Dask-backed data, weights are chunked to avoid eager evaluation. + - The minimum weight threshold is controlled by `self._min_weight`. + """ + with xr.set_options(keep_attrs=True): + # Keep the original weights for other operations and make a copy + # to avoid modifying the original weights. + self._weights = self._get_weights(ds, str(dv.name)) + weights = self._weights.copy() + + # For Dask-backed data variables, chunk the weights along the + # time dimension before broadcasting to avoid eager evaluation + # of the masking step. if dv.chunks: - # For Dask-backed data variables, chunk the weights along the - # time dimension before broadcasting to avoid eager evaluation - # of the masking step. weights = weights.chunk({self.dim: dv.chunksizes[self.dim]}) - weights, _ = xr.broadcast(self._weights, dv) - weights = xr.where(dv.copy().isnull(), 0.0, weights) - - # Perform weighted average using the formula - # WA = sum(data*weights) / sum(weights). The denominator must be - # included to take into account zero weight for missing data. - with xr.set_options(keep_attrs=True): - dv = self._group_data(dv).sum(skipna=skipna) / self._group_data( - weights - ).sum(skipna=skipna) - - # Restore the data variable's name. - dv.name = data_var - else: - dv = self._group_data(dv).mean(skipna=skipna) - # After grouping and aggregating, the grouped time dimension's - # attributes are removed. Xarray's `keep_attrs=True` option only keeps - # attributes for data variables and not their coordinates, so the - # coordinate attributes have to be restored manually. - dv[self.dim].attrs = self._labeled_time.attrs - dv[self.dim].encoding = self._labeled_time.encoding + # Apply the weights to data. + dv_weighted = dv * weights - dv = self._add_operation_attrs(dv) + # Group and sum weighted data, skipping NaNs if specified. + dv_group_sum = self._group_data(dv_weighted).sum(skipna=skipna) - return dv + # Mask weights where data is missing (set to zero), then + # group and sum the masked weights. This ensures that only weights + # corresponding to non-missing data are used in the denominator of + # the weighted average. + masked_weights = _get_masked_weights(dv, self._weights) + masked_weights_group_sum = self._group_data(masked_weights).sum( + skipna=skipna + ) + + # Compute weighted average using the formula: + # WA = sum(data * weights) / sum(weights for non-missing data) + dv_avg = dv_group_sum / masked_weights_group_sum + + # Mask averaged data where the fraction of weights in each group + # does not meet the minimum weight threshold (fractional). + if self._min_weight > 0.0: + # The sum of all weights in each group (i.e., full coverage) + weight_sum_all = self._group_data(self._weights).sum(skipna=skipna) + + # Fraction of weights present in each group. + weight_fraction = masked_weights_group_sum / weight_sum_all + + # Mask the averaged data where the weight fraction is below + # the minimum weight threshold. + dv_avg = xr.where( + weight_fraction >= self._min_weight, + dv_avg, + np.nan, + keep_attrs=True, + ) + + # Restore the data variables name which gets lost after arithmetic + # and masking operations. + dv_avg.name = dv.name + + return dv_avg def _get_weights(self, ds: xr.Dataset, data_var: str) -> xr.DataArray: """Calculates weights for a data variable using time bounds. @@ -2003,14 +2103,15 @@ def _add_operation_attrs(self, data_var: xr.DataArray) -> xr.DataArray: xr.DataArray The data variable with a temporal averaging attributes. """ - data_var.attrs.update( - { - "operation": "temporal_avg", - "mode": self._mode, - "freq": self._freq, - "weighted": str(self._weighted), - } - ) + attrs_to_set = { + "operation": "temporal_avg", + "mode": self._mode, + "freq": self._freq, + "weighted": str(self._weighted), + } + + if self._weighted: + attrs_to_set["min_weight"] = self._min_weight # type: ignore if self._freq == "season": drop_incomplete_seasons = self._season_config["drop_incomplete_seasons"] @@ -2019,16 +2120,18 @@ def _add_operation_attrs(self, data_var: xr.DataArray) -> xr.DataArray: # TODO: Deprecate drop_incomplete_djf. This attr is only set if the # user does not set drop_incomplete_seasons. if drop_incomplete_seasons is False and drop_incomplete_djf is not False: - data_var.attrs["drop_incomplete_djf"] = str(drop_incomplete_djf) + attrs_to_set["drop_incomplete_djf"] = str(drop_incomplete_djf) else: - data_var.attrs["drop_incomplete_seasons"] = str(drop_incomplete_seasons) + attrs_to_set["drop_incomplete_seasons"] = str(drop_incomplete_seasons) custom_seasons = self._season_config.get("custom_seasons") if custom_seasons is not None: - data_var.attrs["custom_seasons"] = list(custom_seasons.keys()) + attrs_to_set["custom_seasons"] = list(custom_seasons.keys()) # type: ignore else: dec_mode = self._season_config.get("dec_mode") - data_var.attrs["dec_mode"] = dec_mode + attrs_to_set["dec_mode"] = dec_mode # type: ignore + + data_var.attrs.update(attrs_to_set) return data_var