From a12cf30423b7cf4d2ced9b57a8b9813a99e62ed7 Mon Sep 17 00:00:00 2001 From: tomvothecoder Date: Thu, 4 Apr 2024 14:44:25 -0700 Subject: [PATCH] Fix tests --- tests/test_temporal.py | 120 ++++++++++++++++++++++++++++++++--------- xcdat/temporal.py | 24 +++++---- 2 files changed, 108 insertions(+), 36 deletions(-) diff --git a/tests/test_temporal.py b/tests/test_temporal.py index e9ead770..9ba0763e 100644 --- a/tests/test_temporal.py +++ b/tests/test_temporal.py @@ -1,4 +1,5 @@ import logging +import warnings import cftime import numpy as np @@ -620,7 +621,7 @@ def test_weighted_seasonal_averages_with_DJF_without_dropping_incomplete_seasons }, ) - assert result.identical(expected) + xr.testing.assert_identical(result, expected) def test_weighted_seasonal_averages_with_DJF_and_drop_incomplete_seasons(self): ds = generate_dataset(decode_times=True, cf_compliant=True, has_bounds=True) @@ -872,7 +873,7 @@ def test_weighted_custom_seasonal_averages_drops_incomplete_seasons(self): }, ) - assert result.identical(expected) + xr.testing.assert_identical(result, expected) def test_weighted_custom_seasonal_averages_with_seasons_spanning_calendar_years( self, @@ -1151,7 +1152,7 @@ def test_raises_error_if_reference_period_arg_is_incorrect(self): ds.temporal.climatology( "ts", "season", - season_config={"dec_mode": "DJF", "drop_incomplete_djf": True}, + season_config={"dec_mode": "DJF", "drop_incomplete_seasons": True}, reference_period=("01-01-2000", "01-01-2000"), ) @@ -1159,7 +1160,7 @@ def test_raises_error_if_reference_period_arg_is_incorrect(self): ds.temporal.climatology( "ts", "season", - season_config={"dec_mode": "DJF", "drop_incomplete_djf": True}, + season_config={"dec_mode": "DJF", "drop_incomplete_seasons": True}, reference_period=("01-01-2000"), ) @@ -1169,7 +1170,7 @@ def test_subsets_climatology_based_on_reference_period(self): result = ds.temporal.climatology( "ts", "season", - season_config={"dec_mode": "DJF", "drop_incomplete_djf": True}, + season_config={"dec_mode": "DJF", "drop_incomplete_seasons": True}, reference_period=("2000-01-01", "2000-06-01"), ) @@ -1201,7 +1202,7 @@ def test_subsets_climatology_based_on_reference_period(self): "freq": "season", "weighted": "True", "dec_mode": "DJF", - "drop_incomplete_djf": "True", + "drop_incomplete_seasons": "True", }, ) @@ -1261,6 +1262,70 @@ def test_weighted_seasonal_climatology_with_DJF(self): xr.testing.assert_identical(result, expected) + 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. + ds = self.ds.copy() + + with warnings.catch_warnings(record=True) as w: + result = ds.temporal.climatology( + "ts", + "season", + season_config={"dec_mode": "DJF", "drop_incomplete_djf": True}, + ) + + assert len(w) == 1 + assert issubclass(w[0].category, DeprecationWarning) + assert str(w[0].message) == ( + "The `season_config` argument 'drop_incomplete_djf' is being deprecated. " + "Please use 'drop_incomplete_seasons' instead." + ) + + 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", + "drop_incomplete_seasons": "True", + "dec_mode": "DJF", + }, + ) + + xr.testing.assert_identical(result, expected) + @requires_dask def test_chunked_weighted_seasonal_climatology_with_DJF(self): ds = self.ds.copy().chunk({"time": 2}) @@ -1468,7 +1533,7 @@ def test_weighted_custom_seasonal_climatology_with_seasons_spanning_calendar_yea expected["ts"] = xr.DataArray( name="ts", - data=np.ones((4, 4, 4)), + data=np.ones((1, 4, 4)), coords={"lat": expected.lat, "lon": expected.lon, "time": expected_time}, dims=["time", "lat", "lon"], attrs={ @@ -1481,7 +1546,7 @@ def test_weighted_custom_seasonal_climatology_with_seasons_spanning_calendar_yea }, ) - assert result.identical(expected) + xr.testing.assert_identical(result, expected) def test_weighted_monthly_climatology(self): result = self.ds.temporal.climatology("ts", "month") @@ -1902,7 +1967,7 @@ def test_raises_error_if_reference_period_arg_is_incorrect(self): ds.temporal.departures( "ts", "season", - season_config={"dec_mode": "DJF", "drop_incomplete_djf": True}, + season_config={"dec_mode": "DJF", "drop_incomplete_seasons": True}, reference_period=("01-01-2000", "01-01-2000"), ) @@ -1910,7 +1975,7 @@ def test_raises_error_if_reference_period_arg_is_incorrect(self): ds.temporal.departures( "ts", "season", - season_config={"dec_mode": "DJF", "drop_incomplete_djf": True}, + season_config={"dec_mode": "DJF", "drop_incomplete_seasons": True}, reference_period=("01-01-2000"), ) @@ -1921,7 +1986,7 @@ def test_seasonal_departures_relative_to_climatology_reference_period(self): "ts", "season", weighted=True, - season_config={"dec_mode": "DJF", "drop_incomplete_djf": True}, + season_config={"dec_mode": "DJF", "drop_incomplete_seasons": False}, reference_period=("2000-01-01", "2000-06-01"), ) @@ -1929,13 +1994,14 @@ def test_seasonal_departures_relative_to_climatology_reference_period(self): expected = expected.drop_dims("time") expected["ts"] = xr.DataArray( name="ts", - data=np.array([[[0.0]], [[np.nan]], [[np.nan]], [[np.nan]]]), + data=np.array([[[0.0]], [[0.0]], [[np.nan]], [[np.nan]], [[0.0]]]), 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(2000, 10, 1), @@ -1959,7 +2025,7 @@ def test_seasonal_departures_relative_to_climatology_reference_period(self): "freq": "season", "weighted": "True", "dec_mode": "DJF", - "drop_incomplete_djf": "True", + "drop_incomplete_seasons": "False", }, ) @@ -1974,7 +2040,7 @@ def test_monthly_departures_relative_to_climatology_reference_period_with_same_o "ts", "month", weighted=True, - season_config={"dec_mode": "DJF", "drop_incomplete_djf": True}, + season_config={"dec_mode": "DJF", "drop_incomplete_seasons": True}, reference_period=("2000-01-01", "2000-06-01"), ) @@ -2057,20 +2123,21 @@ def test_weighted_seasonal_departures_with_DJF(self): "ts", "season", weighted=True, - season_config={"dec_mode": "DJF", "drop_incomplete_seasons": True}, + season_config={"dec_mode": "DJF", "drop_incomplete_seasons": False}, ) expected = ds.copy() expected = expected.drop_dims("time") expected["ts"] = xr.DataArray( name="ts", - data=np.array([[[0.0]], [[0.0]], [[0.0]], [[0.0]]]), + data=np.array([[[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]]]), 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(2000, 10, 1), @@ -2094,7 +2161,7 @@ def test_weighted_seasonal_departures_with_DJF(self): "freq": "season", "weighted": "True", "dec_mode": "DJF", - "drop_incomplete_seasons": "True", + "drop_incomplete_seasons": "False", }, ) @@ -2108,20 +2175,21 @@ def test_weighted_seasonal_departures_with_DJF_and_keep_weights(self): "season", weighted=True, keep_weights=True, - season_config={"dec_mode": "DJF", "drop_incomplete_seasons": True}, + season_config={"dec_mode": "DJF", "drop_incomplete_seasons": False}, ) expected = ds.copy() expected = expected.drop_dims("time") expected["ts"] = xr.DataArray( name="ts", - data=np.array([[[0.0]], [[0.0]], [[0.0]], [[0.0]]]), + data=np.array([[[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]]]), 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(2000, 10, 1), @@ -2145,16 +2213,17 @@ def test_weighted_seasonal_departures_with_DJF_and_keep_weights(self): "freq": "season", "weighted": "True", "dec_mode": "DJF", - "drop_incomplete_seasons": "True", + "drop_incomplete_seasons": "False", }, ) expected["time_wts"] = xr.DataArray( name="ts", - data=np.array([1.0, 1.0, 1.0, 1.0]), + data=np.array([0.52542373, 1.0, 1.0, 1.0, 0.47457627]), coords={ "time_original": xr.DataArray( data=np.array( [ + "2000-01-16T12:00:00.000000000", "2000-03-16T12:00:00.000000000", "2000-06-16T00:00:00.000000000", "2000-09-16T00:00:00.000000000", @@ -2174,7 +2243,7 @@ def test_weighted_seasonal_departures_with_DJF_and_keep_weights(self): dims=["time_original"], ) - xr.testing.assert_identical(result, expected) + xr.testing.assert_allclose(result, expected) def test_unweighted_seasonal_departures_with_DJF(self): ds = self.ds.copy() @@ -2183,20 +2252,21 @@ def test_unweighted_seasonal_departures_with_DJF(self): "ts", "season", weighted=False, - season_config={"dec_mode": "DJF", "drop_incomplete_seasons": True}, + season_config={"dec_mode": "DJF", "drop_incomplete_seasons": False}, ) expected = ds.copy() expected = expected.drop_dims("time") expected["ts"] = xr.DataArray( name="ts", - data=np.array([[[0.0]], [[0.0]], [[0.0]], [[0.0]]]), + data=np.array([[[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]]]), 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(2000, 10, 1), @@ -2219,7 +2289,7 @@ def test_unweighted_seasonal_departures_with_DJF(self): "mode": "departures", "freq": "season", "weighted": "False", - "drop_incomplete_seasons": "True", + "drop_incomplete_seasons": "False", "dec_mode": "DJF", }, ) diff --git a/xcdat/temporal.py b/xcdat/temporal.py index 6275e4b1..f7efcaaa 100644 --- a/xcdat/temporal.py +++ b/xcdat/temporal.py @@ -67,7 +67,6 @@ SeasonConfigInput = TypedDict( "SeasonConfigInput", { - "drop_incomplete_djf": bool, "drop_incomplete_seasons": bool, "dec_mode": Literal["DJF", "JFD"], "custom_seasons": Optional[List[List[str]]], @@ -78,7 +77,6 @@ SeasonConfigAttr = TypedDict( "SeasonConfigAttr", { - "drop_incomplete_djf": bool, "drop_incomplete_seasons": bool, "dec_mode": Literal["DJF", "JFD"], "custom_seasons": Optional[Dict[str, List[str]]], @@ -319,7 +317,7 @@ def group_average( * "custom_seasons" ([List[List[str]]], by default None) List of sublists containing month strings, with each sublist - representing a custom season. This config overrides the `decod + representing a custom season. * Month strings must be in the three letter format (e.g., 'Jan') * Each month must be included once in a custom season @@ -504,7 +502,7 @@ def climatology( * "custom_seasons" ([List[List[str]]], by default None) List of sublists containing month strings, with each sublist - representing a custom season. This config overrides the `decod + representing a custom season. * Month strings must be in the three letter format (e.g., 'Jan') * Each month must be included once in a custom season @@ -971,7 +969,9 @@ def _set_arg_attrs( # "season" frequency specific configuration attributes. for key in season_config.keys(): - if key not in DEFAULT_SEASON_CONFIG.keys(): + # TODO: Deprecate `drop_incomplete_djf`. + valid_keys = list(DEFAULT_SEASON_CONFIG.keys()) + ["drop_incomplete_djf"] + if key not in valid_keys: raise KeyError( f"'{key}' is not a supported season config. Supported " f"configs include: {DEFAULT_SEASON_CONFIG.keys()}." @@ -989,7 +989,7 @@ def _set_arg_attrs( "deprecated. Please use 'drop_incomplete_seasons' instead.", DeprecationWarning, ) - self._season_config["drop_incomplete_seasons"] = drop_incomplete_djf + self._season_config["drop_incomplete_seasons"] = drop_incomplete_djf # type: ignore else: self._season_config["drop_incomplete_seasons"] = season_config.get( "drop_incomplete_seasons", False @@ -1089,7 +1089,7 @@ def _preprocess_dataset(self, ds: xr.Dataset) -> xr.Dataset: # "NDJFM", we should subset the dataset for time coordinates # belonging to those months. months = self._season_config["custom_seasons"].values() # type: ignore - months = list(chain.from_iterable(months.values())) # type: ignore + months = list(chain.from_iterable(months)) if len(months) != 12: ds = self._subset_coords_for_custom_seasons(ds, months) @@ -1133,12 +1133,14 @@ def _subset_coords_for_custom_seasons( The dataset with time coordinate subsetted to months used in custom seasons. """ - months_ints = sorted([MONTH_STR_TO_INT[month] for month in months]) + month_ints = sorted([MONTH_STR_TO_INT[month] for month in months]) coords_by_month = ds.time.groupby(f"{self.dim}.month").groups - months_idxs = {k: coords_by_month[k] for k in months_ints} - months_idxs = sorted(list(chain.from_iterable(months_idxs.values()))) # type: ignore - ds_new = ds.isel({f"{self.dim}": months_idxs}) + month_to_time_idx = { + k: coords_by_month[k] for k in month_ints if k in coords_by_month + } + month_to_time_idx = sorted(list(chain.from_iterable(month_to_time_idx.values()))) # type: ignore + ds_new = ds.isel({f"{self.dim}": month_to_time_idx}) return ds_new