Skip to content

Commit

Permalink
Fix _drop_incomplete_seasons() adding dims to variables
Browse files Browse the repository at this point in the history 
- Add conditional that determines whether subsetting time coordinates is necessary with custom seasons
- Update docstrings for `season_config`
- Add tests
  • Loading branch information
@tomvothecoder
tomvothecoder committed Mar 7, 2023
1 parent fa087b7 commit c11e505
Show file tree
Hide file tree
Showing 2 changed files with 110 additions and 91 deletions.
77 changes: 31 additions & 46 deletions tests/test_temporal.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -575,27 +575,28 @@ 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_and_drop_incomplete_seasons(self):
def test_weighted_seasonal_averages_with_DJF_without_dropping_incomplete_seasons(
self,
):
ds = self.ds.copy()

result = ds.temporal.group_average(
"ts",
"season",
season_config={"dec_mode": "DJF", "drop_incomplete_seasons": True},
season_config={"dec_mode": "DJF", "drop_incomplete_seasons": False},
)
expected = ds.copy()
# Drop the incomplete DJF seasons
expected = expected.isel(time=slice(2, -1))
expected = expected.drop_dims("time")
expected["ts"] = xr.DataArray(
name="ts",
data=np.array([[[1]], [[1]], [[1]], [[2.0]]]),
data=np.array([[[2.0]], [[1.0]], [[1.0]], [[1.0]], [[2.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),
Expand All @@ -618,35 +619,33 @@ def test_weighted_seasonal_averages_with_DJF_and_drop_incomplete_seasons(self):
"mode": "group_average",
"freq": "season",
"weighted": "True",
"drop_incomplete_seasons": "False",
"dec_mode": "DJF",
"drop_incomplete_seasons": "True",
},
)

assert result.identical(expected)

def test_weighted_seasonal_averages_with_DJF_without_dropping_incomplete_seasons(
self,
):
ds = self.ds.copy()
def test_weighted_seasonal_averages_with_DJF_and_drop_incomplete_seasons(self):
ds = generate_dataset(decode_times=True, cf_compliant=True, has_bounds=True)

result = ds.temporal.group_average(
"ts",
"season",
season_config={"dec_mode": "DJF", "drop_incomplete_seasons": False},
season_config={"dec_mode": "DJF", "drop_incomplete_seasons": True},
)

expected = ds.copy()
expected = expected.drop_dims("time")
expected["ts"] = xr.DataArray(
name="ts",
data=np.array([[[2.0]], [[1.0]], [[1.0]], [[1.0]], [[2.0]]]),
data=np.ones((4, 4, 4)),
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),
Expand All @@ -664,13 +663,12 @@ def test_weighted_seasonal_averages_with_DJF_without_dropping_incomplete_seasons
},
dims=["time", "lat", "lon"],
attrs={
"test_attr": "test",
"operation": "temporal_avg",
"mode": "group_average",
"freq": "season",
"weighted": "True",
"drop_incomplete_seasons": "True",
"dec_mode": "DJF",
"drop_incomplete_seasons": "False",
},
)

Expand Down Expand Up @@ -729,6 +727,7 @@ def test_weighted_seasonal_averages_with_JFD(self):
"mode": "group_average",
"freq": "season",
"weighted": "True",
"drop_incomplete_seasons": "False",
"dec_mode": "JFD",
},
)
Expand Down Expand Up @@ -810,13 +809,14 @@ def test_weighted_custom_seasonal_averages(self):
"operation": "temporal_avg",
"mode": "group_average",
"freq": "season",
"weighted": "True",
"drop_incomplete_seasons": "False",
"custom_seasons": [
"JanFebMar",
"AprMayJun",
"JulAugSep",
"OctNovDec",
],
"weighted": "True",
},
)

Expand Down Expand Up @@ -870,8 +870,9 @@ def test_weighted_custom_seasonal_averages_drops_incomplete_seasons(self):
"operation": "temporal_avg",
"mode": "group_average",
"freq": "season",
"custom_seasons": ["NovDec", "FebMarApr"],
"weighted": "True",
"drop_incomplete_seasons": "True",
"custom_seasons": ["NovDec", "FebMarApr"],
},
)

Expand Down Expand Up @@ -926,8 +927,9 @@ def test_weighted_custom_seasonal_averages_with_seasons_spanning_calendar_years(
"operation": "temporal_avg",
"mode": "group_average",
"freq": "season",
"custom_seasons": ["NovDecJanFebMar"],
"weighted": "True",
"drop_incomplete_seasons": "False",
"custom_seasons": ["NovDecJanFebMar"],
},
)

Expand Down Expand Up @@ -1198,8 +1200,8 @@ def test_weighted_seasonal_climatology_with_DJF(self):
"mode": "climatology",
"freq": "season",
"weighted": "True",
"dec_mode": "DJF",
"drop_incomplete_seasons": "True",
"dec_mode": "DJF",
},
)

Expand Down Expand Up @@ -1305,6 +1307,7 @@ def test_weighted_seasonal_climatology_with_JFD(self):
"mode": "climatology",
"freq": "season",
"weighted": "True",
"drop_incomplete_seasons": "False",
"dec_mode": "JFD",
},
)
Expand Down Expand Up @@ -1363,6 +1366,7 @@ def test_weighted_custom_seasonal_climatology(self):
"mode": "climatology",
"freq": "season",
"weighted": "True",
"drop_incomplete_seasons": "False",
"custom_seasons": [
"JanFebMar",
"AprMayJun",
Expand All @@ -1374,46 +1378,30 @@ def test_weighted_custom_seasonal_climatology(self):

assert result.identical(expected)

@pytest.mark.xfail
def test_weighted_custom_seasonal_climatology_with_seasons_spanning_calendar_years(
self,
):
ds = self.ds.copy()

custom_seasons = [
["Jan", "Feb", "Mar"],
["Apr", "May", "Jun"],
["Jul", "Aug", "Sep"],
["Oct", "Nov", "Dec"],
]
custom_seasons = [["Nov", "Dec", "Jan", "Feb", "Mar"]]
result = ds.temporal.climatology(
"ts",
"season",
season_config={
"drop_incomplete_seasons": False,
"custom_seasons": custom_seasons,
"drop_incomplete_seasons": True,
},
)

expected = ds.copy()
expected = expected.drop_dims("time")
expected_time = xr.DataArray(
data=np.array(
[
cftime.DatetimeGregorian(1, 2, 1),
cftime.DatetimeGregorian(1, 5, 1),
cftime.DatetimeGregorian(1, 8, 1),
cftime.DatetimeGregorian(1, 11, 1),
],
[cftime.DatetimeGregorian(1, 1, 1)],
),
coords={
"time": np.array(
[
cftime.DatetimeGregorian(1, 2, 1),
cftime.DatetimeGregorian(1, 5, 1),
cftime.DatetimeGregorian(1, 8, 1),
cftime.DatetimeGregorian(1, 11, 1),
],
[cftime.DatetimeGregorian(1, 1, 1)],
),
},
attrs={
Expand All @@ -1434,12 +1422,8 @@ def test_weighted_custom_seasonal_climatology_with_seasons_spanning_calendar_yea
"mode": "climatology",
"freq": "season",
"weighted": "True",
"custom_seasons": [
"JanFebMar",
"AprMayJun",
"JulAugSep",
"OctNovDec",
],
"drop_incomplete_seasons": "False",
"custom_seasons": ["NovDecJanFebMar"],
},
)

Expand Down Expand Up @@ -1938,8 +1922,8 @@ def test_unweighted_seasonal_departures_with_DJF(self):
"mode": "departures",
"freq": "season",
"weighted": "False",
"dec_mode": "DJF",
"drop_incomplete_seasons": "True",
"dec_mode": "DJF",
},
)

Expand Down Expand Up @@ -1969,6 +1953,7 @@ def test_unweighted_seasonal_departures_with_JFD(self):
"mode": "departures",
"freq": "season",
"weighted": "False",
"drop_incomplete_seasons": "False",
"dec_mode": "JFD",
},
)
Expand Down
Loading

0 comments on commit c11e505

Please sign in to comment.