Repository: NCAR/geocat-comp

Issue:
month_to_season
Port the NCL **month_to_season** family of functions, including: Month_to_season, month_to_seasonN, month_to_season12. The **month_to_season** function computes a user-specified three-month seasonal mean (DJF, JFM, FMA, MAM, AMJ, MJJ, JJA, JAS, ASO, SON, OND, NDJ).

Comments:

Patch:
diff --git a/src/geocat/comp/month_to_season.py b/src/geocat/comp/month_to_season.py
new file mode 100644
index 0000000..c5b0e94
--- /dev/null
+++ b/src/geocat/comp/month_to_season.py
@@ -0,0 +1,77 @@
+
+
+def month_to_season(xMon, season, time_coord_name='time'):
+    """ This function takes an xarray dataset containing monthly data spanning years and
+        returns a dataset with one sample per year, for a specified three-month season.
+
+        This function requires the number of months to be a multiple of 12, i.e. full years must be provided.
+
+        Time stamps are centered on the season. For example, seasons='DJF' returns January timestamps.
+
+        If a calculated season's timestamp falls outside the original range of monthly values, then the calculated mean
+        is dropped.  For example, if the monthly data's time range is [Jan-2000, Dec-2003] and the season is "DJF", the
+        seasonal mean computed from the single month of Dec-2003 is dropped.
+    """
+    mod_check(xMon[time_coord_name].size, 12)
+
+    startDate = xMon[time_coord_name][0]
+    endDate = xMon[time_coord_name][-1]
+    seasons_pd = {
+        'DJF': (
+            'QS-DEC',
+            1),
+        'JFM': (
+            'QS-JAN',
+            2),
+        'FMA': (
+            'QS-FEB',
+            3),
+        'MAM': (
+            'QS-MAR',
+            4),
+        'AMJ': (
+            'QS-APR',
+            5),
+        'MJJ': (
+            'QS-MAY',
+            6),
+        'JJA': (
+            'QS-JUN',
+            7),
+        'JAS': (
+            'QS-JUL',
+            8),
+        'ASO': (
+            'QS-AUG',
+            9),
+        'SON': (
+            'QS-SEP',
+            10),
+        'OND': (
+            'QS-OCT',
+            11),
+        'NDJ': (
+            'QS-NOV',
+            12)}
+    try:
+        (season_pd, season_sel) = seasons_pd[season]
+    except KeyError:
+        raise KeyError(
+            f"contributed: month_to_season: bad season: SEASON = {season}. Valid seasons include: {list(seasons_pd.keys())}")
+
+    # Compute the three-month means, moving time labels ahead to the middle
+    # month.
+    month_offset = 'MS'
+    xSeasons = xMon.resample(
+        {time_coord_name: season_pd}, loffset=month_offset).mean()
+
+    # Filter just the desired season, and trim to the desired time range.
+    xSea = xSeasons.sel(
+        {time_coord_name: xSeasons[time_coord_name].dt.month == season_sel})
+    xSea = xSea.sel({time_coord_name: slice(startDate, endDate)})
+    return xSea
+
+
+def mod_check(value, mod):
+    if value % mod != 0:
+        raise ValueError(f'Expected a multiple of {mod} values')


Test patch:
diff --git a/test/test_month_to_season.py b/test/test_month_to_season.py
new file mode 100644
index 0000000..48afb06
--- /dev/null
+++ b/test/test_month_to_season.py
@@ -0,0 +1,131 @@
+import unittest
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+from geocat.comp.month_to_season import month_to_season
+
+
+def get_fake_dataset(start_month, nmonths, nlats, nlons):
+    """ Returns a very simple xarray dataset for testing.
+        Data values are equal to "month of year" for monthly time steps.
+    """
+    # Create coordinates
+    months = pd.date_range(
+        start=pd.to_datetime(start_month),
+        periods=nmonths,
+        freq='MS')
+    lats = np.linspace(start=-90, stop=90, num=nlats, dtype='float32')
+    lons = np.linspace(start=-180, stop=180, num=nlons, dtype='float32')
+
+    # Create data variable. Construct a 3D array with time as the first
+    # dimension.
+    month_values = np.expand_dims(
+        np.arange(
+            start=1,
+            stop=nmonths + 1),
+        axis=(
+            1,
+            2))
+    var_values = np.tile(month_values, (1, nlats, nlons))
+
+    ds = xr.Dataset(
+        data_vars={
+            'my_var': (('time', 'lat', 'lon'), var_values.astype('float32')),
+        },
+        coords={'time': months, 'lat': lats, 'lon': lons},
+    )
+    return ds
+
+
+class Test_month_to_season(unittest.TestCase):
+
+    def setUp(self):
+        # Create a dataset for the year 2000.
+        self.ds1 = get_fake_dataset(
+            start_month='2000-01', nmonths=12, nlats=1, nlons=1)
+
+        # Create another dataset for the year 2001.
+        self.ds2 = get_fake_dataset(
+            start_month='2001-01', nmonths=12, nlats=1, nlons=1)
+
+        # Create a dataset that combines the two previous datasets, for two
+        # years of data.
+        self.ds3 = xr.concat([self.ds1, self.ds2], dim='time')
+
+        # Create a dataset with the wrong number of months.
+        self.partial_year_dataset = get_fake_dataset(
+            start_month='2000-01', nmonths=13, nlats=1, nlons=1)
+
+        # Create a dataset with a custom time coordinate.
+        custom_time_dataset = get_fake_dataset(
+            start_month='2000-01', nmonths=12, nlats=1, nlons=1)
+        self.custom_time_dataset = custom_time_dataset.rename(
+            {'time': 'my_time'})
+
+        # Create a more complex dataset just to verify that get_fake_dataset()
+        # is generally working.
+        self.complex_dataset = get_fake_dataset(
+            start_month='2001-01', nmonths=12, nlats=10, nlons=10)
+
+        # Check all possible season choices for some tests.
+        self.all_seasons = [
+            'DJF',
+            'JFM',
+            'FMA',
+            'MAM',
+            'AMJ',
+            'MJJ',
+            'JJA',
+            'JAS',
+            'ASO',
+            'SON',
+            'OND',
+            'NDJ']
+
+    def test_m2s_returns_middle_month_value(self):
+        season_ds = month_to_season(self.ds1, 'JFM')
+        season_value_array = season_ds['my_var'].data
+
+        # Should equal the average of [1.0, 2.0, 3.0]
+        self.assertEqual(season_value_array[0, 0, 0], 2.0)
+
+        season_ds = month_to_season(self.ds1, 'JJA')
+        season_value_array = season_ds['my_var'].data
+
+        # Should equal the average of [6.0, 7.0, 8.0]
+        self.assertEqual(season_value_array[0, 0, 0], 7.0)
+
+    def test_bad_season_returns_exception(self):
+        with self.assertRaises(KeyError):
+            season_ds = month_to_season(self.ds1, 'XXX')
+
+    def test_partial_years_returns_exception(self):
+        with self.assertRaises(ValueError):
+            season_ds = month_to_season(self.partial_year_dataset, 'JFM')
+
+    def test_final_season_returns_2month_average(self):
+        season_ds = month_to_season(self.ds1, 'NDJ')
+        season_value_array = season_ds['my_var'].data
+        self.assertEqual(season_value_array[0, 0, 0], 11.5)
+
+    def test_each_season_returns_one_point_per_year(self):
+        nyears_of_data = self.ds3.sizes['time'] / 12
+        for season in self.all_seasons:
+            season_ds = month_to_season(self.ds3, season)
+            season_value_array = season_ds['my_var'].data
+            self.assertEqual(season_value_array.size, nyears_of_data)
+
+    def test_custom_time_coordinate(self):
+        season_ds = month_to_season(
+            self.custom_time_dataset,
+            'JFM',
+            time_coord_name='my_time')
+        season_value_array = season_ds['my_var'].data
+
+        # Should equal the average of [1.0, 2.0, 3.0]
+        self.assertEqual(season_value_array[0, 0, 0], 2.0)
+
+
+if __name__ == '__main__':
+    unittest.main()