Function find_independent_seasons

xarray/groupers.py:758–797

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=('D

(seasons: Sequence[str])

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()
    # 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:
            grouped[i].append(current)
            codes[i].append(i)
            seen.add(current)

        # Loop through remaining groups, and look for overlaps
        for j, second in enumerate(season_inds[i:]):
            if not (set(chain(*grouped[i])) & set(second)) and second not in seen:
                grouped[i].append(second)
                codes[i].append(j + i)
                seen.add(second)
        if len(seen) == len(seasons):
            break
        # found all non-overlapping groups for this row start over

    grouped_ints = tuple(tuple(idx) for idx in grouped.values() if idx)
    return [
        SeasonsGroup(seasons=inds_to_season_string(inds), inds=inds, codes=codes)
        for inds, codes in zip(grouped_ints, codes.values(), strict=False)
    ]


@dataclass