Fix to annual bug (#95)

* return nan when summing across nans

* add testing for temporal module

* add pytest lazy fixture testing

update test fixture

CHANGELOG.rst

@@ -2,6 +2,14 @@
 Changelog History
 =================
 
+esmtools v1.1.4 (2020-##-##)
+============================
+
+Bug Fixes
+---------
+- :py:func:`~esmtools.temporal.to_annual` no longer returns ``0.0`` where the original
+  dataset had NaNs. (:issue:`75`) (:pr:`95`) `Riley X. Brady`_.
+
 esmtools v1.1.3 (2020-07-17)
 ============================
 

ci/environment-dev-3.6.yml

@@ -29,6 +29,7 @@ dependencies:
   - pre-commit
   - pytest
   - pytest-cov
+  - pytest-lazy-fixture
   - pytest-sugar
   # Statistics
   - scipy

esmtools/temporal.py

@@ -6,11 +6,11 @@
 from .constants import CALENDARS
 from .timeutils import get_calendar, get_days_per_month
 
-GROUPBY_TIMES = {"annual": "time.year"}
+GROUPBY_TIMES = {'annual': 'time.year'}
 TIME_RESOLUTIONS = [k for k in GROUPBY_TIMES]
 
 
-def _weighted_resample(ds, calendar=None, dim="time", resample_resolution=None):
+def _weighted_resample(ds, calendar=None, dim='time', resample_resolution=None):
     """Generalized function for time-weighted resampling.
 
     Args:
@@ -28,12 +28,12 @@ def _weighted_resample(ds, calendar=None, dim="time", resample_resolution=None):
         calendar = get_calendar(ds[dim])
 
     if resample_resolution not in TIME_RESOLUTIONS:
-        raise ValueError(f"Please submit a temporal resolution from {TIME_RESOLUTIONS}")
+        raise ValueError(f'Please submit a temporal resolution from {TIME_RESOLUTIONS}')
 
     time_length = xr.DataArray(
         get_days_per_month(ds.time.to_index(), calendar=calendar),
         coords=[ds.time],
-        name="time_length",
+        name='time_length',
     )
 
     time_res = GROUPBY_TIMES[resample_resolution]
@@ -45,11 +45,12 @@ def _weighted_resample(ds, calendar=None, dim="time", resample_resolution=None):
     np.testing.assert_allclose(weights_sum, np.ones(len(weights_sum)))
 
     # Calculate the weighted average
-    ds_weighted = (ds * weights).groupby(time_res).sum(dim=dim)
+    # `skipna=False` ensures that masked locations like land remain nan.
+    ds_weighted = (ds * weights).groupby(time_res).sum(dim=dim, skipna=False)
     return ds_weighted
 
 
-def to_annual(ds, calendar=None, how="mean", dim="time"):
+def to_annual(ds, calendar=None, how='mean', dim='time'):
     """Resample sub-annual temporal resolution to annual resolution with weighting.
 
     .. note::
@@ -80,11 +81,11 @@ def to_annual(ds, calendar=None, how="mean", dim="time"):
     Returns:
         ds_weighted (xarray object): Dataset or DataArray resampled to annual resolution
     """
-    if how != "mean":
+    if how != 'mean':
         raise NotImplementedError(
-            "Only annual-weighted averaging is currently"
+            'Only annual-weighted averaging is currently'
             + "supported. Please change `how` keyword to 'mean'"
         )
     return _weighted_resample(
-        ds, calendar=calendar, dim=dim, resample_resolution="annual"
+        ds, calendar=calendar, dim=dim, resample_resolution='annual'
     )

esmtools/tests/test_temporal.py

@@ -0,0 +1,48 @@
+import numpy as np
+import pytest
+import xarray as xr
+
+from esmtools.temporal import to_annual
+
+
+@pytest.mark.parametrize(
+    'dataset',
+    (
+        pytest.lazy_fixture('gridded_da_datetime'),
+        pytest.lazy_fixture('gridded_da_cftime'),
+    ),
+)
+def test_to_annual(dataset):
+    """General checks that `to_annual` time conversion is working as expected."""
+    data = dataset()
+    result = to_annual(data)
+    assert result.notnull().all()
+    assert 'year' in result.dims
+
+
+def test_to_annual_accuracy(ts_monthly_da):
+    """Tests that weighted sum correctly takes the annual mean."""
+    data = ts_monthly_da().isel(time=slice(0, 12))
+    MONTH_LENGTHS = [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
+    manual_sum = []
+    for i in range(len(data)):
+        manual_sum.append(data[i].values * MONTH_LENGTHS[i] / 365)
+    expected = sum(manual_sum)
+    actual = to_annual(data)
+    assert np.abs(actual.values - expected) < 1e-5
+
+
+def test_to_annual_retains_nans(gridded_da_landmask):
+    """Tests that `to_annual` function retains nans where the original dataset had nans
+
+    .. note::
+        Previous versions of `esmtools` did not do this, since xarray automatically
+        skips nans with the grouped sum operator, returning zeroes where there used
+        to be nans.
+    """
+    data = gridded_da_landmask
+    data['time'] = xr.cftime_range(
+        start='1990-01', freq='MS', periods=data['time'].size
+    )
+    result = to_annual(data)
+    assert result.isel(lat=0, lon=0).isnull().all()