Repository: NCAR/geocat-comp (python)

Issue:
Add `keep_attrs` kwarg to all `climatology.py` functions
Changes in xarray have resulted in an error when passing numpy objects into some `meteorology.py` functions. It has been determined in https://github.com/NCAR/geocat-comp/pull/312 that this arose from the global xarray option `keep_attrs` being set to `True` at the beginning of `climatology.py`. That line is being removed in #312, but to maintain the same the same behavior of `climatology.py` functions, a new kwarg must be added to each function in the file. This issue addresses that.

Pull request:
Add `keep_attrs` to functions
Closes #314

I cannot test these functions with `np.ndarray` inputs since they don't yet explicitly support that data structure. As far as I can tell, the aggregation functions for `xarray.core.groupby` and `xarray.core.resample` objects are the only ones used in this file that have `keep_attrs` as an argument.

TODO from our discussion below:
- [x] set the default to None to preserve xarray default behavior
- [x] add a test for the new kwarg

Patch:
diff --git a/src/geocat/comp/climatologies.py b/src/geocat/comp/climatologies.py
index 264351b..db7a6d8 100644
--- a/src/geocat/comp/climatologies.py
+++ b/src/geocat/comp/climatologies.py
@@ -60,12 +60,13 @@ def _setup_clim_anom_input(dset, freq, time_coord_name):
 
 
 def _calculate_center_of_time_bounds(
-    dset: typing.Union[xr.Dataset,
-                       xr.DataArray], time_dim: str, freq: str, calendar: str,
-    start: typing.Union[str,
-                        cftime.datetime], end: typing.Union[str,
-                                                            cftime.datetime]
-) -> typing.Union[xr.Dataset, xr.DataArray]:
+        dset: typing.Union[xr.Dataset, xr.DataArray],
+        time_dim: str,
+        freq: str,
+        calendar: str,
+        start: typing.Union[str, cftime.datetime],
+        end: typing.Union[str, cftime.datetime],
+        keep_attrs: bool = None) -> typing.Union[xr.Dataset, xr.DataArray]:
     """Helper function to determine the time bounds based on the given dataset
     and frequency and then calculate the averages of them.
 
@@ -93,6 +94,11 @@ def _calculate_center_of_time_bounds(
     end : str, :class:`cftime.datetime`
         The ending date of the data. The string representation must be in ISO format
 
+    keep_attrs : bool, optional
+        If True, attrs will be copied from the original object to the new one.
+        If False, the new object will be returned without attributes.
+        Defaults to None which means the attrs will only be kept in unambiguous circumstances.
+
     Returns
     -------
     computed_dset: :class:`xarray.Dataset`, :class:`xarray.DataArray`
@@ -107,7 +113,7 @@ def _calculate_center_of_time_bounds(
     time_bounds = time_bounds.append(time_bounds[-1:].shift(1, freq=freq))
     time =  xr.DataArray(np.vstack((time_bounds[:-1], time_bounds[1:])).T,
                          dims=[time_dim, 'nbd']) \
-        .mean(dim='nbd')
+        .mean(dim='nbd', keep_attrs=keep_attrs)
     return dset.assign_coords({time_dim: time})
 
 
@@ -128,7 +134,8 @@ def _infer_calendar_name(dates):
 def climatology(
         dset: typing.Union[xr.DataArray, xr.Dataset],
         freq: str,
-        time_coord_name: str = None) -> typing.Union[xr.DataArray, xr.Dataset]:
+        time_coord_name: str = None,
+        keep_attrs: bool = None) -> typing.Union[xr.DataArray, xr.Dataset]:
     """Compute climatologies for a specified time frequency.
 
     Parameters
@@ -148,6 +155,11 @@ def climatology(
          Name for time coordinate to use. Defaults to ``None`` and infers the name
          from the data.
 
+    keep_attrs : bool, optional
+        If True, attrs will be copied from the original object to the new one.
+        If False, the new object will be returned without attributes.
+        Defaults to None which means the attrs will only be kept in unambiguous circumstances.
+
     Returns
     -------
     computed_dset : :class:`xarray.Dataset`, :class:`xarray.DataArray`
@@ -215,9 +227,15 @@ def climatology(
 
     grouped = data.groupby(time_dot_freq)
     # TODO: Compute weighted climatologies when `time_bounds` are available
-    clim = grouped.mean(time_coord_name)
+    clim = grouped.mean(time_coord_name, keep_attrs=keep_attrs)
     if time_invariant_vars:
-        return xr.concat([dset[time_invariant_vars], clim], dim=time_coord_name)
+        if keep_attrs == False:
+            return xr.concat([dset[time_invariant_vars], clim],
+                             combine_attrs='drop',
+                             dim=time_coord_name)
+        else:
+            return xr.concat([dset[time_invariant_vars], clim],
+                             dim=time_coord_name)
     else:
         return clim
 
@@ -314,10 +332,10 @@ def anomaly(
 
 
 def month_to_season(
-    dset: typing.Union[xr.Dataset, xr.DataArray],
-    season: str,
-    time_coord_name: str = None,
-) -> typing.Union[xr.Dataset, xr.DataArray]:
+        dset: typing.Union[xr.Dataset, xr.DataArray],
+        season: str,
+        time_coord_name: str = None,
+        keep_attrs: bool = None) -> typing.Union[xr.Dataset, xr.DataArray]:
     """Computes a user-specified three-month seasonal mean.
 
     This function takes an xarray dataset containing monthly data spanning years and
@@ -347,6 +365,11 @@ def month_to_season(
         Name for time coordinate to use. Defaults to ``None`` and infers the name
         from the data.
 
+    keep_attrs : bool, optional
+        If True, attrs will be copied from the original object to the new one.
+        If False, the new object will be returned without attributes.
+        Defaults to None which means the attrs will only be kept in unambiguous circumstances.
+
     Returns
     -------
     computed_dset : :class:`xarray.Dataset`, :class:`xarray.DataArray`
@@ -406,7 +429,7 @@ def month_to_season(
     # Group the months into three and take the mean
     means = data_filter.resample({
         time_coord_name: quarter
-    }, loffset='MS').mean()
+    }, loffset='MS').mean(keep_attrs=keep_attrs)
 
     # The line above tries to take the mean for all quarters even if there is not data for some of them
     # Therefore, we must filter out the NaNs
@@ -417,7 +440,8 @@ def month_to_season(
 def calendar_average(
         dset: typing.Union[xr.DataArray, xr.Dataset],
         freq: str,
-        time_dim: str = None) -> typing.Union[xr.DataArray, xr.Dataset]:
+        time_dim: str = None,
+        keep_attrs: bool = 'default') -> typing.Union[xr.DataArray, xr.Dataset]:
     """This function divides the data into time periods (months, seasons, etc)
     and computes the average for the data in each one.
 
@@ -445,6 +469,11 @@ def calendar_average(
     computed_dset : :class:`xarray.Dataset`, :class:`xarray.DataArray`
         The computed data with the same type as `dset`
 
+    keep_attrs : bool, optional
+        If True, attrs will be copied from the original object to the new one.
+        If False, the new object will be returned without attributes.
+        Defaults to None which means the attrs will only be kept in unambiguous circumstances.
+
     Examples
     --------
     See this `example notebook <https://geocat-comp.readthedocs.io/en/latest/examples/calendar_average.html>`__.
@@ -495,10 +524,16 @@ def calendar_average(
     calendar = _infer_calendar_name(dset[time_dim])
 
     # Group data
-    dset = dset.resample({time_dim: frequency}).mean().dropna(time_dim)
+    dset = dset.resample({
+        time_dim: frequency
+    }).mean(keep_attrs=keep_attrs).dropna(time_dim)
 
     # Weight the data by the number of days in each month
     if freq in ['season', 'year']:
+        attrs = {}
+        if keep_attrs or keep_attrs is None:
+            attrs = dset.attrs
+
         key = freq
         format, frequency = freq_dict[key]
         # Compute the weights for the months in each season so that the
@@ -506,7 +541,9 @@ def calendar_average(
         month_length = dset[time_dim].dt.days_in_month.resample(
             {time_dim: frequency})
         weights = month_length.map(lambda group: group / group.sum())
-        dset = (dset * weights).resample({time_dim: frequency}).sum()
+
+        dset_weighted = dset * weights
+        dset = (dset_weighted).resample({time_dim: frequency}).sum()
 
     # Center the time coordinate by inferring and then averaging the time bounds
     dset = _calculate_center_of_time_bounds(dset,
@@ -514,14 +551,19 @@ def calendar_average(
                                             frequency,
                                             calendar,
                                             start=dset[time_dim].values[0],
-                                            end=dset[time_dim].values[-1])
-    return dset
+                                            end=dset[time_dim].values[-1],
+                                            keep_attrs=keep_attrs)
+    if freq in ['season', 'year']:
+        return dset.assign_attrs(attrs)
+    else:
+        return dset
 
 
 def climatology_average(
         dset: typing.Union[xr.DataArray, xr.Dataset],
         freq: str,
-        time_dim: str = None) -> typing.Union[xr.DataArray, xr.Dataset]:
+        time_dim: str = None,
+        keep_attrs: bool = None) -> typing.Union[xr.DataArray, xr.Dataset]:
     """This function calculates long term hourly, daily, monthly, or seasonal
     averages across all years in the given dataset.
 
@@ -543,6 +585,10 @@ def climatology_average(
         Name of the time coordinate for `xarray` objects. Defaults to ``None`` and
         infers the name from the data.
 
+    keep_attrs : bool, optional
+        If True, attrs will be copied from the original object to the new one.
+        If False, the new object will be returned without attributes.
+        Defaults to None which means the attrs will only be kept in unambiguous circumstances.
 
     Returns
     -------
@@ -605,24 +651,30 @@ def climatology_average(
     calendar = _infer_calendar_name(dset[time_dim])
 
     if freq == 'season':
+        attrs = {}
+        if keep_attrs or keep_attrs is None:
+            attrs = dset.attrs
         if xr.infer_freq(dset[time_dim]) != 'MS':
             # Calculate monthly average before calculating seasonal climatologies
-            dset = dset.resample({time_dim: frequency}).mean().dropna(time_dim)
+            dset = dset.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()
+        weights = month_length / month_length.sum(keep_attrs=keep_attrs)
         dset = (dset * weights).groupby(f"{time_dim}.season")
-        dset = dset.sum(dim=time_dim)
+        dset = dset.sum(dim=time_dim, keep_attrs=keep_attrs)
+        return dset.assign_attrs(attrs)
     else:
         # Retrieve floor of median year
         median_yr = np.median(dset[time_dim].dt.year.values)
 
         # Group data by format then calculate average of groups
-        dset = dset.groupby(dset[time_dim].dt.strftime(format)).mean().rename(
-            {'strftime': time_dim})
+        dset = dset.groupby(dset[time_dim].dt.strftime(format)).mean(
+            keep_attrs=keep_attrs).rename({'strftime': time_dim})
 
         # Center the time coordinate by inferring and then averaging the time bounds
         start_time = dset[time_dim].values[0]
@@ -633,6 +685,7 @@ def climatology_average(
             frequency,
             calendar,
             start=f'{median_yr:.0f}-{start_time}',
-            end=f'{median_yr:.0f}-{end_time}')
+            end=f'{median_yr:.0f}-{end_time}',
+            keep_attrs=keep_attrs)
 
-    return dset
+        return dset


Test patch:
diff --git a/test/test_climatologies.py b/test/test_climatologies.py
index 421787e..0d25dd1 100644
--- a/test/test_climatologies.py
+++ b/test/test_climatologies.py
@@ -37,16 +37,15 @@ def get_fake_dataset(start_month, nmonths, nlats, nlons):
                                   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
-        },
-    )
+    ds = xr.Dataset(data_vars={
+        "my_var": (("time", "lat", "lon"), var_values.astype("float32")),
+    },
+                    coords={
+                        "time": months,
+                        "lat": lats,
+                        "lon": lons
+                    },
+                    attrs={'Description': 'This is dummy data for testing.'})
     return ds
 
 
@@ -76,10 +75,22 @@ def _get_dummy_data(start_date,
                         'time': time,
                         'lat': lats,
                         'lon': lons
-                    })
+                    },
+                    attrs={'Description': 'This is dummy data for testing.'})
     return ds
 
 
+@pytest.mark.parametrize("dataset", [dset_a, dset_b, dset_c["Tair"]])
+@pytest.mark.parametrize("freq", ["day", "month", "year", "season"])
+@pytest.mark.parametrize("keep_attrs", [None, True, False])
+def test_climatology_keep_attrs(dataset, freq, keep_attrs):
+    computed_dset = climatology(dataset, freq, keep_attrs=keep_attrs)
+    if keep_attrs or keep_attrs == None:
+        assert computed_dset.attrs == dataset.attrs
+    elif not keep_attrs:
+        assert computed_dset.attrs == {}
+
+
 def test_climatology_invalid_freq():
     with pytest.raises(ValueError):
         climatology(dset_a, "hourly")
@@ -134,6 +145,15 @@ complex_dataset = get_fake_dataset(start_month="2001-01",
                                    nlons=10)
 
 
+@pytest.mark.parametrize("keep_attrs", [None, True, False])
+def test_month_to_season_keep_attrs(keep_attrs):
+    season_ds = month_to_season(ds1, 'JFM', keep_attrs=keep_attrs)
+    if keep_attrs or keep_attrs == None:
+        assert season_ds.attrs == ds1.attrs
+    elif not keep_attrs:
+        assert season_ds.attrs == {}
+
+
 @pytest.mark.parametrize("dataset, season, expected", [(ds1, "JFM", 2.0),
                                                        (ds1, "JJA", 7.0)])
 def test_month_to_season_returns_middle_month_value(dataset, season, expected):
@@ -205,7 +225,19 @@ daily = _get_dummy_data('2020-01-01', '2021-12-31', 'D', 1, 1)
 
 monthly = _get_dummy_data('2020-01-01', '2021-12-01', 'MS', 1, 1)
 
+
 # Computational Tests for calendar_average()
+@pytest.mark.parametrize('dset, freq', [(daily, 'month'), (monthly, 'season'),
+                                        (monthly, 'year')])
+@pytest.mark.parametrize('keep_attrs', [None, True, False])
+def test_calendar_average_keep_attrs(dset, freq, keep_attrs):
+    result = calendar_average(dset, freq, keep_attrs=keep_attrs)
+    if keep_attrs or keep_attrs == None:
+        assert result.attrs == dset.attrs
+    elif not keep_attrs:
+        assert result.attrs == {}
+
+
 hour_avg = np.arange(0.5, 35088.5, 2).reshape((365 + 366) * 24, 1, 1)
 hour_avg_time = xr.cftime_range('2020-01-01 00:30:00',
                                 '2021-12-31 23:30:00',
@@ -330,6 +362,15 @@ def test_daily_monthly_to_yearly_calendar_average(dset, expected):
 
 
 # Computational Tests for climatology_average()
+@pytest.mark.parametrize('dset, freq', [(daily, 'month'), (monthly, 'season')])
+@pytest.mark.parametrize('keep_attrs', [None, True, False])
+def test_climatology_average_keep_attrs(dset, freq, keep_attrs):
+    result = climatology_average(dset, freq, keep_attrs=keep_attrs)
+    if keep_attrs or keep_attrs == None:
+        assert result.attrs == dset.attrs
+    elif not keep_attrs:
+        assert result.attrs == {}
+
 hour_clim = np.concatenate([np.arange(8784.5, 11616.5, 2),
                             np.arange(2832.5, 2880.5, 2),
                             np.arange(11640.5, 26328.5, 2)])\