Repository: NCAR/geocat-comp

Issue:
Allow user to define seasonal boundaries for climatology_average
**Describe the feature you'd like added to this project**
Currently `climatology_average` does not allow the user to select the seasonal boundaries. Seasonal climatologies are only calculated for the meteorological seasons, which are defined as follows:
- Winter: December, January, February
- Spring: March, April, May
- Summer: June, July, August
- Fall: September, October, November

This functionality should be added to the climatology module since it is available in NCL's `month_to_season` family of functions.

**Link to original NCL functionality, if applicable**
[month_to_season](https://www.ncl.ucar.edu/Document/Functions/Contributed/month_to_season.shtml)
[month_to_season12](https://www.ncl.ucar.edu/Document/Functions/Contributed/month_to_season12.shtml)
[month_to_seasonN](https://www.ncl.ucar.edu/Document/Functions/Contributed/month_to_seasonN.shtml)

**Describe the solution you'd like**
This could be done either by modifying `climatology_average` or by implementing one-to-one copies of NCL's `month_to_season12` and `month_to_seasonN` in Python.

`climatology_average` could be modified to take a parameter defining the first season of interest. This can be done using the NCL convention of three letter strings to represent the months (i.e. JFM for January, February, March as the first season). The function can take this and infer the other three seasons. Then climatologies for all four seasons can be calculated. This would not exactly replicate functionality in NCL's `month_to_season` functions and may be too specialized of a solution. Input from users is needed.

Discussion:

Patch:
diff --git a/docs/release-notes.rst b/docs/release-notes.rst
index 572c233..47116cb 100644
--- a/docs/release-notes.rst
+++ b/docs/release-notes.rst
@@ -5,6 +5,15 @@
 Release Notes
 =============
 
+v2023.07.0 (Jul X, 2023)
+-----------------------
+[Plain text summary here]
+
+New Features
+^^^^^^^^^^^^
+* User-defined seasonal boundaries, `custom_seasons`, enabled for `climatology_average` by `Julia Kent`_ in (:pr:441)
+
+
 v2023.06.1 (June 23, 2023)
 --------------------------
 This releases fixes the unintentional limitation of the 2023.06.0 release to python 3.11.0
diff --git a/geocat/comp/climatologies.py b/geocat/comp/climatologies.py
index f15336d..224b0a1 100644
--- a/geocat/comp/climatologies.py
+++ b/geocat/comp/climatologies.py
@@ -3,6 +3,7 @@ import cftime
 import numpy as np
 import typing
 import xarray as xr
+import warnings
 
 _FREQUENCIES = {"day", "month", "year", "season"}
 
@@ -182,7 +183,7 @@ def climate_anomaly(
     }
 
     if freq not in freq_dict:
-        raise KeyError(
+        raise ValueError(
             f"Received bad period {freq!r}. Expected one of {list(freq_dict.keys())!r}"
         )
     format, frequency = freq_dict[freq]
@@ -190,7 +191,10 @@ def climate_anomaly(
     if freq == 'year':
         clim = calendar_average(dset, freq, time_dim, keep_attrs)
     else:
-        clim = climatology_average(dset, freq, time_dim, keep_attrs)
+        clim = climatology_average(dset,
+                                   freq=freq,
+                                   time_dim=time_dim,
+                                   keep_attrs=keep_attrs)
     if freq == 'season':
         anom = dset.groupby(f"{time_dim}.season") - clim
         return anom.assign_attrs(attrs)
@@ -265,7 +269,7 @@ def month_to_season(
         raise ValueError(
             f"The {time_coord_name} axis length must be a multiple of {mod}.")
 
-    seasons_pd = {
+    seasons_dict = {
         "DJF": ([12, 1, 2], 'QS-DEC'),
         "JFM": ([1, 2, 3], 'QS-JAN'),
         "FMA": ([2, 3, 4], 'QS-FEB'),
@@ -280,10 +284,10 @@ def month_to_season(
         "NDJ": ([11, 12, 1], 'QS-NOV'),
     }
     try:
-        (months, quarter) = seasons_pd[season]
-    except KeyError:
-        raise KeyError(
-            f"contributed: month_to_season: bad season: SEASON = {season}. Valid seasons include: {list(seasons_pd.keys())}"
+        (months, quarter) = seasons_dict[season]
+    except ValueError:
+        raise ValueError(
+            f"contributed: month_to_season: bad season: SEASON = {season}. Valid seasons include: {list(seasons_dict.keys())}"
         )
 
     # Filter data to only contain the months of interest
@@ -370,7 +374,7 @@ def calendar_average(
     }
 
     if freq not in freq_dict:
-        raise KeyError(
+        raise ValueError(
             f"Received bad period {freq!r}. Expected one of {list(freq_dict.keys())!r}"
         )
 
@@ -430,6 +434,7 @@ def calendar_average(
 def climatology_average(
         dset: typing.Union[xr.DataArray, xr.Dataset],
         freq: str,
+        custom_seasons: typing.Union[list, str] = None,
         time_dim: str = None,
         keep_attrs: bool = None) -> typing.Union[xr.DataArray, xr.Dataset]:
     """This function calculates long term hourly, daily, monthly, or seasonal
@@ -447,7 +452,26 @@ def climatology_average(
         - `hour`: for hourly averages
         - `day`: for daily averages
         - `month`: for monthly averages
-        - `season`: for meteorological seasonal averages (DJF, MAM, JJA, and SON)
+        - `season`: for meteorological seasonal averages (default: DJF, JJA, MAM, and SON)
+
+    custom_seasons : list[str], str, optional
+        The list of 3-months season aliases or a single seaonal alias string.
+        Analysis is done on the provided seasons.
+        This parameter will be ignored if the `freq` is not set to `season`.
+        Accepted alias:
+
+        - `DJF` : for a season of December, January, and February
+        - `JFM` : for a season of January, February, and March
+        - `FMA` : for a season of February, March, and April
+        - `MAM` : for a season of March, April, and May
+        - `AMJ` : for a season of April, May, ad June
+        - `MJJ` : for a season of May, June, and July
+        - `JJA` : for a season of June, July, and August
+        - `JAS` : for a season of July, August, and September
+        - `ASO` : for a season of August, September, and October
+        - `SON` : for a season of September, October, and November
+        - `OND` : for a season of October, November, and December
+        - `NDJ` : for a season of November, December, and January
 
     time_dim : str, optional
         Name of the time coordinate for `xarray` objects. Defaults to ``None`` and
@@ -488,7 +512,6 @@ def climatology_average(
     `clmMonTLLL <https://www.ncl.ucar.edu/Document/Functions/Contributed/clmMonTLLL.shtml>`__,
     `month_to_season <https://www.ncl.ucar.edu/Document/Functions/Contributed/month_to_season.shtml>`__
     """
-    # TODO: add functionality for users to select specific seasons or hours for climatologies
     freq_dict = {
         'hour': ('%m-%d %H', 'H'),
         'day': ('%m-%d', 'D'),
@@ -496,11 +519,31 @@ def climatology_average(
         'season': (None, 'QS-DEC')
     }
 
+    seasons_dict = {
+        "DJF": [12, 1, 2],
+        "JFM": [1, 2, 3],
+        "FMA": [2, 3, 4],
+        "MAM": [3, 4, 5],
+        "AMJ": [4, 5, 6],
+        "MJJ": [5, 6, 7],
+        "JJA": [6, 7, 8],
+        "JAS": [7, 8, 9],
+        "ASO": [8, 9, 10],
+        "SON": [9, 10, 11],
+        "OND": [10, 11, 12],
+        "NDJ": [11, 12, 1],
+    }
+
     if freq not in freq_dict:
-        raise KeyError(
+        raise ValueError(
             f"Received bad period {freq!r}. Expected one of {list(freq_dict.keys())!r}"
         )
 
+    if custom_seasons and freq != 'season':
+        warnings.warn(
+            "You have specified `custom_seasons` without specifying seasonal climatology. Your `custom_seasons` will be ignored."
+        )
+
     # If freq is 'season', key is set to monthly in order to calculate monthly
     # averages which are then used to calculate seasonal averages
     key = 'month' if freq == 'season' else freq
@@ -518,24 +561,49 @@ def climatology_average(
     # Retrieve calendar name
     calendar = _infer_calendar_name(dset[time_dim])
 
+    attrs = {}
+    if keep_attrs in {True, None}:
+        attrs = dset.attrs
+
     if freq == 'season':
-        attrs = {}
-        if keep_attrs or keep_attrs is None:
-            attrs = dset.attrs
-        if xr.infer_freq(dset[time_dim]) != 'MS':
+        seasons = ['DJF', 'JJA', 'MAM', 'SON']
+        if custom_seasons:
+            seasons = custom_seasons
+            if isinstance(seasons, str):
+                seasons = [seasons]
+
+            invalid_seasons = [s for s in seasons if s not in seasons_dict]
+            if invalid_seasons:
+                raise ValueError(
+                    f"contributed: climatology_average: bad season(s): {invalid_seasons}. Valid seasons include: {list(seasons_dict.keys())}"
+                )
+
+        seasonal_climates = []
+        for season in seasons:
+
+            # Grab the months for each season
+            months = seasons_dict[season]
+
+            # Filter data to only contain the months of interest
+            dset_filter = dset.sel(
+                {time_dim: dset[time_dim].dt.month.isin(months)})
+
             # Calculate monthly average before calculating seasonal climatologies
-            dset = dset.resample({
+            dset_filter = dset_filter.resample({
                 time_dim: frequency
             }).mean(keep_attrs=keep_attrs).dropna(time_dim)
 
-        # Compute the weights for the months in each season so that the
-        # seasonal averages account for months being of different lengths
-        month_length = dset[time_dim].dt.days_in_month.groupby(
-            f"{time_dim}.season")
-        weights = month_length / month_length.sum(keep_attrs=keep_attrs)
-        dset = (dset * weights).groupby(f"{time_dim}.season")
-        dset = dset.sum(dim=time_dim, keep_attrs=keep_attrs)
+            # Compute the weights for the months in each season so that the
+            # seasonal averages account for months being of different lengths
+            month_length = dset_filter[time_dim].dt.days_in_month
+            weights = month_length / month_length.sum()
+            climatology = (dset_filter * weights).sum(dim=time_dim)
+
+            seasonal_climates.append(climatology)
+        dset = xr.concat(seasonal_climates, dim='season')
+        dset.coords['season'] = np.array(seasons).astype(object)
         return dset.assign_attrs(attrs)
+
     else:
         # Retrieve floor of median year
         median_yr = np.median(dset[time_dim].dt.year.values)


Test patch:
diff --git a/test/test_climatologies.py b/test/test_climatologies.py
index bf40c1b..d749899 100644
--- a/test/test_climatologies.py
+++ b/test/test_climatologies.py
@@ -612,7 +612,7 @@ class test_calendar_average(unittest.TestCase):
 
     @parameterized.expand([('freq=TEST', 'TEST'), ('freq=None', None)])
     def test_invalid_freq_calendar_average(self, name, freq):
-        with self.assertRaises(KeyError):
+        with self.assertRaises(ValueError):
             calendar_average(self.monthly, freq=freq)
 
     def test_custom_time_coord_calendar_average(self):
@@ -711,7 +711,7 @@ class test_climatology_average(unittest.TestCase):
     day_2_season_clim = xr.Dataset(
         data_vars={'data': (('season', 'lat', 'lon'), season_clim)},
         coords={
-            'season': season_clim_time,
+            'season': np.array(season_clim_time).astype(object),
             'lat': [-90.0],
             'lon': [-180.0]
         })
@@ -720,7 +720,7 @@ class test_climatology_average(unittest.TestCase):
     month_2_season_clim = xr.Dataset(
         data_vars={'data': (('season', 'lat', 'lon'), season_clim)},
         coords={
-            'season': season_clim_time,
+            'season': np.array(season_clim_time).astype(object),
             'lat': [-90.0],
             'lon': [-180.0]
         })
@@ -828,15 +828,26 @@ class test_climatology_average(unittest.TestCase):
             'lon': [-180.0]
         })
 
-    @parameterized.expand([('daily, "month", None', daily, 'month', None),
-                           ('daily, "month", True', daily, 'month', True),
-                           ('daily, "month", False', daily, 'month', False),
-                           ('monthly, "season", None', monthly, 'season', None),
-                           ('monthly, "season", True', monthly, 'season', True),
-                           ('monthly, "season", False', monthly, 'season',
-                            False)])
-    def test_climatology_average_keep_attrs(self, name, dset, freq, keep_attrs):
-        result = climatology_average(dset, freq, keep_attrs=keep_attrs)
+    @parameterized.expand([
+        ('daily, "month", None', daily, 'month', [], None),
+        ('daily, "month", True', daily, 'month', [], True),
+        ('daily, "month", False', daily, 'month', [], False),
+        ('monthly, "season", None', monthly, 'season', [], None),
+        ('monthly, "season", True', monthly, 'season', [], True),
+        ('monthly, "season", False', monthly, 'season', [], False),
+        ('monthly, "season", None', monthly, 'season',
+         ['DJF', 'MAM', 'JJA', 'SON'], None),
+        ('monthly, "season", True', monthly, 'season',
+         ['DJF', 'MAM', 'JJA', 'SON'], True),
+        ('monthly, "season", False', monthly, 'season',
+         ['DJF', 'MAM', 'JJA', 'SON'], False)
+    ])
+    def test_climatology_average_keep_attrs(self, name, dset, freq,
+                                            custom_seasons, keep_attrs):
+        result = climatology_average(dset,
+                                     freq=freq,
+                                     custom_seasons=custom_seasons,
+                                     keep_attrs=keep_attrs)
         if keep_attrs or keep_attrs == None:
             assert result.attrs == dset.attrs
         elif not keep_attrs:
@@ -854,6 +865,14 @@ class test_climatology_average(unittest.TestCase):
         result = climatology_average(self.daily, freq='month')
         xr.testing.assert_allclose(result, self.day_2_month_clim)
 
+    def test_custom_season_climatology_average(self):
+        result = climatology_average(
+            self.monthly,
+            freq='season',
+            custom_seasons=['DJF', 'JJA', 'MAM', 'SON'])
+        expected = climatology_average(self.monthly, freq='season')
+        xr.testing.assert_equal(result, expected)
+
     @parameterized.expand([('daily to seasonal', daily, day_2_season_clim),
                            ('monthly to seasonal', monthly, month_2_season_clim)
                           ])
@@ -864,7 +883,7 @@ class test_climatology_average(unittest.TestCase):
 
     @parameterized.expand([('freq=TEST', 'TEST'), ('freq=None', None)])
     def test_invalid_freq_climatology_average(self, name, freq):
-        with self.assertRaises(KeyError):
+        with self.assertRaises(ValueError):
             climatology_average(self.monthly, freq=freq)
 
     def test_custom_time_coord_climatology_average(self):
@@ -889,7 +908,7 @@ class test_climatology_average(unittest.TestCase):
         dset_encoded = xr.tutorial.open_dataset("air_temperature",
                                                 decode_cf=False)
         with self.assertRaises(ValueError):
-            climatology_average(dset_encoded, 'month')
+            climatology_average(dset_encoded, freq='month')
 
     def test_non_uniformly_spaced_data_climatology_average(self):
         time = pd.to_datetime(['2020-01-01', '2020-01-02', '2020-01-04'])