cleanup

Author: dcherian

properties/test_properties.py

@@ -5,11 +5,11 @@
 pytest.importorskip("hypothesis")
 
 import hypothesis.strategies as st
-from hypothesis import given
+from hypothesis import given, note
 
 import xarray as xr
 import xarray.testing.strategies as xrst
-from xarray.groupers import season_to_month_tuple
+from xarray.groupers import find_independent_seasons, season_to_month_tuple
 
 
 @given(attrs=xrst.simple_attrs)
@@ -46,3 +46,18 @@ def test_property_season_month_tuple(roll, breaks):
         rolled_months[start:stop] for start, stop in itertools.pairwise(breaks)
     )
     assert expected == actual
+
+
+@given(data=st.data(), nmonths=st.integers(min_value=1, max_value=11))
+def test_property_find_independent_seasons(data, nmonths):
+    chars = "JFMAMJJASOND"
+    # if stride > nmonths, then we can't infer season order
+    stride = data.draw(st.integers(min_value=1, max_value=nmonths))
+    chars = chars + chars[:nmonths]
+    seasons = [list(chars[i : i + nmonths]) for i in range(0, 12, stride)]
+    note(seasons)
+    groups = find_independent_seasons(seasons)
+    for group in groups:
+        inds = tuple(itertools.chain(*group.inds))
+        assert len(inds) == len(set(inds))
+        assert len(group.codes) == len(set(group.codes))

xarray/groupers.py

@@ -604,6 +604,14 @@ def unique_value_groups(
 
 
 def season_to_month_tuple(seasons: Sequence[str]) -> tuple[tuple[int, ...], ...]:
+    """
+    >>> season_to_month_tuple(["DJF", "MAM", "JJA", "SON"])
+    ((12, 1, 2), (3, 4, 5), (6, 7, 8), (9, 10, 11))
+    >>> season_to_month_tuple(["DJFM", "MAMJ", "JJAS", "SOND"])
+    ((12, 1, 2, 3), (3, 4, 5, 6), (6, 7, 8, 9), (9, 10, 11, 12))
+    >>> season_to_month_tuple(["DJFM", "SOND"])
+    ((12, 1, 2, 3), (9, 10, 11, 12))
+    """
     initials = "JFMAMJJASOND"
     starts = dict(
         ("".join(s), i + 1)
@@ -629,7 +637,7 @@ def season_to_month_tuple(seasons: Sequence[str]) -> tuple[tuple[int, ...], ...]
     return tuple(result)
 
 
-def inds_to_string(asints: tuple[tuple[int, ...], ...]) -> tuple[str, ...]:
+def inds_to_season_string(asints: tuple[tuple[int, ...], ...]) -> tuple[str, ...]:
     inits = "JFMAMJJASOND"
     return tuple("".join([inits[i_ - 1] for i_ in t]) for t in asints)
 
@@ -660,24 +668,36 @@ def is_sorted_periodic(lst):
     return lst[-1] <= lst[0]
 
 
-@dataclass
+@dataclass(kw_only=True, frozen=True)
 class SeasonsGroup:
     seasons: tuple[str, ...]
+    # tuple[integer months] corresponding to each season
     inds: tuple[tuple[int, ...], ...]
+    # integer code for each season, this is not simply range(len(seasons))
+    # when the seasons have overlaps
     codes: Sequence[int]
 
 
 def find_independent_seasons(seasons: Sequence[str]) -> Sequence[SeasonsGroup]:
     """
     Iterates though a list of seasons e.g. ["DJF", "FMA", ...],
     and splits that into multiple sequences of non-overlapping seasons.
+
+    >>> find_independent_seasons(
+    ...     ["DJF", "FMA", "AMJ", "JJA", "ASO", "OND"]
+    ... )  # doctest: +NORMALIZE_WHITESPACE
+    [SeasonsGroup(seasons=('DJF', 'AMJ', 'ASO'), inds=((12, 1, 2), (4, 5, 6), (8, 9, 10)), codes=[0, 2, 4]),
+    SeasonsGroup(seasons=('FMA', 'JJA', 'OND'), inds=((2, 3, 4), (6, 7, 8), (10, 11, 12)), codes=[1, 3, 5])]
+
+    >>> find_independent_seasons(["DJF", "MAM", "JJA", "SON"])
+    [SeasonsGroup(seasons=('DJF', 'MAM', 'JJA', 'SON'), inds=((12, 1, 2), (3, 4, 5), (6, 7, 8), (9, 10, 11)), codes=[0, 1, 2, 3])]
     """
     season_inds = season_to_month_tuple(seasons)
     grouped = defaultdict(list)
     codes = defaultdict(list)
     seen: set[tuple[int, ...]] = set()
     idx = 0
-    # This is quadratic, but the length of seasons is at most 12
+    # This is quadratic, but the number of seasons is at most 12
     for i, current in enumerate(season_inds):
         # Start with a group
         if current not in seen:
@@ -699,7 +719,7 @@ def find_independent_seasons(seasons: Sequence[str]) -> Sequence[SeasonsGroup]:
 
     grouped_ints = tuple(tuple(idx) for idx in grouped.values() if idx)
     return [
-        SeasonsGroup(seasons=inds_to_string(inds), inds=inds, codes=codes)
+        SeasonsGroup(seasons=inds_to_season_string(inds), inds=inds, codes=codes)
         for inds, codes in zip(grouped_ints, codes.values(), strict=False)
     ]
 

xarray/tests/test_groupby.py

@@ -3378,7 +3378,7 @@ def test_season_grouper_with_partial_years(self, calendar):
 
         assert_allclose(expected, actual)
 
-    @pytest.mark.parametrize("calendar", _ALL_CALENDARS)
+    @pytest.mark.parametrize("calendar", ["standard"])
     def test_season_grouper_with_single_month_seasons(self, calendar):
         time = date_range("2001-01-01", "2002-12-30", freq="MS", calendar=calendar)
         # fmt: off
@@ -3392,29 +3392,32 @@ def test_season_grouper_with_single_month_seasons(self, calendar):
         da = DataArray(data, dims="time", coords={"time": time})
         da["year"] = da.time.dt.year
 
-        actual = da.groupby(
-            year=UniqueGrouper(),
-            time=SeasonGrouper(
-                ["J", "F", "M", "A", "M", "J", "J", "A", "S", "O", "N", "D"]
-            ),
-        ).mean()
+        # TODO: Consider supporting this if needed
+        # It does not work without flox, because the group labels are not unique,
+        # and so the stack/unstack approach does not work.
+        with pytest.raises(ValueError):
+            da.groupby(
+                year=UniqueGrouper(),
+                time=SeasonGrouper(
+                    ["J", "F", "M", "A", "M", "J", "J", "A", "S", "O", "N", "D"]
+                ),
+            ).mean()
 
         # Expected if single month seasons are handled correctly
-        expected = xr.DataArray(
-            data=np.array(
-                [
-                    [1.0, 1.25, 1.5, 1.75, 2.0, 1.1, 1.35, 1.6, 1.85, 1.2, 1.45, 1.7],
-                    [1.95, 1.05, 1.3, 1.55, 1.8, 1.15, 1.4, 1.65, 1.9, 1.25, 1.5, 1.75],
-                ]
-            ),
-            dims=["year", "season"],
-            coords={
-                "year": [2001, 2002],
-                "season": ["J", "F", "M", "A", "M", "J", "J", "A", "S", "O", "N", "D"],
-            },
-        )
-
-        assert_allclose(expected, actual)
+        # expected = xr.DataArray(
+        #     data=np.array(
+        #         [
+        #             [1.0, 1.25, 1.5, 1.75, 2.0, 1.1, 1.35, 1.6, 1.85, 1.2, 1.45, 1.7],
+        #             [1.95, 1.05, 1.3, 1.55, 1.8, 1.15, 1.4, 1.65, 1.9, 1.25, 1.5, 1.75],
+        #         ]
+        #     ),
+        #     dims=["year", "season"],
+        #     coords={
+        #         "year": [2001, 2002],
+        #         "season": ["J", "F", "M", "A", "M", "J", "J", "A", "S", "O", "N", "D"],
+        #     },
+        # )
+        # assert_allclose(expected, actual)
 
     @pytest.mark.parametrize("calendar", _ALL_CALENDARS)
     def test_season_grouper_with_months_spanning_calendar_year_using_previous_year(
@@ -3481,13 +3484,7 @@ def test_season_resampling_raises_unsorted_seasons(self, seasons):
             da.resample(time=SeasonResampler(seasons))
 
     @pytest.mark.parametrize(
-        "use_cftime",
-        [
-            pytest.param(
-                True, marks=pytest.mark.skipif(not has_cftime, reason="no cftime")
-            ),
-            False,
-        ],
+        "use_cftime", [pytest.param(True, marks=requires_cftime), False]
     )
     @pytest.mark.parametrize("drop_incomplete", [True, False])
     @pytest.mark.parametrize(