Add skipna arg to all temporal APIs

- Add unit tests

Author: tomvothecoder

tests/test_spatial.py

@@ -180,6 +180,28 @@ def test_spatial_average_for_lat_region(self):
 
         assert result.identical(expected)
 
+    def test_spatial_average_for_lat_region_and_skipna(self):
+        ds = self.ds.copy(deep=True)
+        ds.ts[0] = np.nan
+
+        # Specifying axis as a str instead of list of str.
+        result = ds.spatial.average("ts", axis=["Y"], lat_bounds=(-5.0, 5), skipna=True)
+
+        expected = self.ds.copy()
+        expected["ts"] = xr.DataArray(
+            data=np.array(
+                [
+                    [np.nan, np.nan, np.nan, np.nan],
+                    [1.0, 1.0, 1.0, 1.0],
+                    [1.0, 1.0, 1.0, 1.0],
+                ]
+            ),
+            coords={"time": expected.time, "lon": expected.lon},
+            dims=["time", "lon"],
+        )
+
+        assert result.identical(expected)
+
     def test_spatial_average_for_domain_wrapping_p_meridian_non_cf_conventions(
         self,
     ):

tests/test_temporal.py

@@ -520,6 +520,57 @@ def test_weighted_annual_averages(self):
         assert result.ts.attrs == expected.ts.attrs
         assert result.time.attrs == expected.time.attrs
 
+    def test_weighted_annual_averages_and_skipna(self):
+        ds = self.ds.copy(deep=True)
+        ds.ts[0] = np.nan
+
+        result = ds.temporal.group_average("ts", "year", skipna=True)
+        expected = ds.copy()
+        expected = expected.drop_dims("time")
+        expected["ts"] = xr.DataArray(
+            name="ts",
+            data=np.array([[[1]], [[2.0]]]),
+            coords={
+                "lat": expected.lat,
+                "lon": expected.lon,
+                "time": xr.DataArray(
+                    data=np.array(
+                        [
+                            cftime.DatetimeGregorian(2000, 1, 1),
+                            cftime.DatetimeGregorian(2001, 1, 1),
+                        ],
+                    ),
+                    coords={
+                        "time": np.array(
+                            [
+                                cftime.DatetimeGregorian(2000, 1, 1),
+                                cftime.DatetimeGregorian(2001, 1, 1),
+                            ],
+                        )
+                    },
+                    dims=["time"],
+                    attrs={
+                        "axis": "T",
+                        "long_name": "time",
+                        "standard_name": "time",
+                        "bounds": "time_bnds",
+                    },
+                ),
+            },
+            dims=["time", "lat", "lon"],
+            attrs={
+                "test_attr": "test",
+                "operation": "temporal_avg",
+                "mode": "group_average",
+                "freq": "year",
+                "weighted": "True",
+            },
+        )
+
+        xr.testing.assert_allclose(result, expected)
+        assert result.ts.attrs == expected.ts.attrs
+        assert result.time.attrs == expected.time.attrs
+
     @requires_dask
     def test_weighted_annual_averages_with_chunking(self):
         ds = self.ds.copy().chunk({"time": 2})
@@ -1161,6 +1212,68 @@ def test_weighted_seasonal_climatology_with_DJF(self):
 
         xr.testing.assert_identical(result, expected)
 
+    def test_weighted_seasonal_climatology_with_DJF_and_skipna(self):
+        ds = self.ds.copy(deep=True)
+
+        # Replace all MAM values with np.nan.
+        djf_months = [3, 4, 5]
+        for mon in djf_months:
+            ds["ts"] = ds.ts.where(ds.ts.time.dt.month != mon, np.nan)
+
+        result = ds.temporal.climatology(
+            "ts",
+            "season",
+            season_config={"dec_mode": "DJF", "drop_incomplete_djf": True},
+            skipna=True,
+        )
+
+        expected = ds.copy()
+        expected = expected.drop_dims("time")
+        expected_time = xr.DataArray(
+            data=np.array(
+                [
+                    cftime.DatetimeGregorian(1, 1, 1),
+                    cftime.DatetimeGregorian(1, 4, 1),
+                    cftime.DatetimeGregorian(1, 7, 1),
+                    cftime.DatetimeGregorian(1, 10, 1),
+                ],
+            ),
+            coords={
+                "time": np.array(
+                    [
+                        cftime.DatetimeGregorian(1, 1, 1),
+                        cftime.DatetimeGregorian(1, 4, 1),
+                        cftime.DatetimeGregorian(1, 7, 1),
+                        cftime.DatetimeGregorian(1, 10, 1),
+                    ],
+                ),
+            },
+            attrs={
+                "axis": "T",
+                "long_name": "time",
+                "standard_name": "time",
+                "bounds": "time_bnds",
+            },
+        )
+        expected["ts"] = xr.DataArray(
+            name="ts",
+            data=np.ones((4, 4, 4)),
+            coords={"lat": expected.lat, "lon": expected.lon, "time": expected_time},
+            dims=["time", "lat", "lon"],
+            attrs={
+                "operation": "temporal_avg",
+                "mode": "climatology",
+                "freq": "season",
+                "weighted": "True",
+                "dec_mode": "DJF",
+                "drop_incomplete_djf": "True",
+            },
+        )
+        expected.ts[1] = np.nan
+
+        # MAM should be np.nan
+        assert result.identical(expected)
+
     @requires_dask
     def test_chunked_weighted_seasonal_climatology_with_DJF(self):
         ds = self.ds.copy().chunk({"time": 2})
@@ -1947,6 +2060,62 @@ def test_weighted_seasonal_departures_with_DJF(self):
 
         xr.testing.assert_identical(result, expected)
 
+    def test_weighted_seasonal_departures_with_DJF_and_skipna(self):
+        ds = self.ds.copy(deep=True)
+
+        # Replace all MAM values with np.nan.
+        djf_months = [3, 4, 5]
+        for mon in djf_months:
+            ds["ts"] = ds.ts.where(ds.ts.time.dt.month != mon, np.nan)
+
+        result = ds.temporal.departures(
+            "ts",
+            "season",
+            weighted=True,
+            season_config={"dec_mode": "DJF", "drop_incomplete_djf": True},
+            skipna=True,
+        )
+
+        expected = ds.copy()
+        expected = expected.drop_dims("time")
+        expected["ts"] = xr.DataArray(
+            name="ts",
+            data=np.array([[[np.nan]], [[0.0]], [[0.0]], [[0.0]]]),
+            coords={
+                "lat": expected.lat,
+                "lon": expected.lon,
+                "time": xr.DataArray(
+                    data=np.array(
+                        [
+                            cftime.DatetimeGregorian(2000, 4, 1),
+                            cftime.DatetimeGregorian(2000, 7, 1),
+                            cftime.DatetimeGregorian(2000, 10, 1),
+                            cftime.DatetimeGregorian(2001, 1, 1),
+                        ],
+                    ),
+                    dims=["time"],
+                    attrs={
+                        "axis": "T",
+                        "long_name": "time",
+                        "standard_name": "time",
+                        "bounds": "time_bnds",
+                    },
+                ),
+            },
+            dims=["time", "lat", "lon"],
+            attrs={
+                "test_attr": "test",
+                "operation": "temporal_avg",
+                "mode": "departures",
+                "freq": "season",
+                "weighted": "True",
+                "dec_mode": "DJF",
+                "drop_incomplete_djf": "True",
+            },
+        )
+
+        assert result.identical(expected)
+
     def test_weighted_seasonal_departures_with_DJF_and_keep_weights(self):
         ds = self.ds.copy()
 

xcdat/temporal.py

@@ -160,7 +160,7 @@ def average(
         data_var: str,
         weighted: bool = True,
         keep_weights: bool = False,
-        skipna: Union[bool, None] = None,
+        skipna: bool | None = None,
     ):
         """
         Returns a Dataset with the average of a data variable and the time
@@ -202,7 +202,7 @@ def average(
         keep_weights : bool, optional
             If calculating averages using weights, keep the weights in the
             final dataset output, by default False.
-        skipna : bool or None, optional
+        skipna : bool | None, optional
             If True, skip missing values (as marked by NaN). By default, only
             skips missing values for float dtypes; other dtypes either do not
             have a sentinel missing value (int) or ``skipna=True`` has not been
@@ -257,6 +257,7 @@ def group_average(
         weighted: bool = True,
         keep_weights: bool = False,
         season_config: SeasonConfigInput = DEFAULT_SEASON_CONFIG,
+        skipna: bool | None = None,
     ):
         """Returns a Dataset with average of a data variable by time group.
 
@@ -335,6 +336,11 @@ def group_average(
                 >>>     ["Jul", "Aug", "Sep"],  # "JulAugSep"
                 >>>     ["Oct", "Nov", "Dec"],  # "OctNovDec"
                 >>> ]
+        skipna : bool | None, optional
+            If True, skip missing values (as marked by NaN). By default, only
+            skips missing values for float dtypes; other dtypes either do not
+            have a sentinel missing value (int) or ``skipna=True`` has not been
+            implemented (object, datetime64 or timedelta64).
 
         Returns
         -------
@@ -413,6 +419,7 @@ def group_average(
             weighted=weighted,
             keep_weights=keep_weights,
             season_config=season_config,
+            skipna=skipna,
         )
 
     def climatology(
@@ -423,6 +430,7 @@ def climatology(
         keep_weights: bool = False,
         reference_period: Optional[Tuple[str, str]] = None,
         season_config: SeasonConfigInput = DEFAULT_SEASON_CONFIG,
+        skipna: bool | None = None,
     ):
         """Returns a Dataset with the climatology of a data variable.
 
@@ -510,6 +518,11 @@ def climatology(
                 >>>     ["Jul", "Aug", "Sep"],  # "JulAugSep"
                 >>>     ["Oct", "Nov", "Dec"],  # "OctNovDec"
                 >>> ]
+        skipna : bool | None, optional
+            If True, skip missing values (as marked by NaN). By default, only
+            skips missing values for float dtypes; other dtypes either do not
+            have a sentinel missing value (int) or ``skipna=True`` has not been
+            implemented (object, datetime64 or timedelta64).
 
         Returns
         -------
@@ -593,6 +606,7 @@ def climatology(
             keep_weights,
             reference_period,
             season_config,
+            skipna,
         )
 
     def departures(
@@ -603,6 +617,7 @@ def departures(
         keep_weights: bool = False,
         reference_period: Optional[Tuple[str, str]] = None,
         season_config: SeasonConfigInput = DEFAULT_SEASON_CONFIG,
+        skipna: bool | None = None,
     ) -> xr.Dataset:
         """
         Returns a Dataset with the climatological departures (anomalies) for a
@@ -697,6 +712,11 @@ def departures(
                 >>>     ["Jul", "Aug", "Sep"],  # "JulAugSep"
                 >>>     ["Oct", "Nov", "Dec"],  # "OctNovDec"
                 >>> ]
+        skipna : bool | None, optional
+            If True, skip missing values (as marked by NaN). By default, only
+            skips missing values for float dtypes; other dtypes either do not
+            have a sentinel missing value (int) or ``skipna=True`` has not been
+            implemented (object, datetime64 or timedelta64).
 
         Returns
         -------
@@ -777,11 +797,7 @@ def departures(
         inferred_freq = _infer_freq(ds[self.dim])
         if inferred_freq != freq:
             ds_obs = ds_obs.temporal.group_average(
-                data_var,
-                freq,
-                weighted,
-                keep_weights,
-                season_config,
+                data_var, freq, weighted, keep_weights, season_config, skipna
             )
 
         # 4. Calculate the climatology of the data variable.
@@ -794,6 +810,7 @@ def departures(
             keep_weights,
             reference_period,
             season_config,
+            skipna,
         )
 
         # 5. Calculate the departures for the data variable.
@@ -815,7 +832,7 @@ def _averager(
         keep_weights: bool = False,
         reference_period: Optional[Tuple[str, str]] = None,
         season_config: SeasonConfigInput = DEFAULT_SEASON_CONFIG,
-        skipna: Union[bool, None] = None,
+        skipna: bool | None = None,
     ) -> xr.Dataset:
         """Averages a data variable based on the averaging mode and frequency."""
         ds = self._dataset.copy()
@@ -1141,7 +1158,7 @@ def _drop_leap_days(self, ds: xr.Dataset):
         return ds
 
     def _average(
-        self, ds: xr.Dataset, data_var: str, skipna: Union[bool, None] = None
+        self, ds: xr.Dataset, data_var: str, skipna: bool | None = None
     ) -> xr.DataArray:
         """Averages a data variable with the time dimension removed.
 
@@ -1151,7 +1168,7 @@ def _average(
             The dataset.
         data_var : str
             The key of the data variable.
-        skipna : bool or None, optional
+        skipna : bool | None, optional
             If True, skip missing values (as marked by NaN). By default, only
             skips missing values for float dtypes; other dtypes either do not
             have a sentinel missing value (int) or ``skipna=True`` has not been
@@ -1178,7 +1195,7 @@ def _average(
         return dv
 
     def _group_average(
-        self, ds: xr.Dataset, data_var: str, skipna: Union[bool, None] = None
+        self, ds: xr.Dataset, data_var: str, skipna: bool | None = None
     ) -> xr.DataArray:
         """Averages a data variable by time group.
 
@@ -1188,7 +1205,7 @@ def _group_average(
             The dataset.
         data_var : str
             The key of the data variable.
-        skipna : bool or None, optional
+        skipna : bool | None, optional
             If True, skip missing values (as marked by NaN). By default, only
             skips missing values for float dtypes; other dtypes either do not
             have a sentinel missing value (int) or ``skipna=True`` has not been