Repository: xCDAT/xcdat
Issue category: Performance Issue

[Enhancement]: Temporal averaging performance
### Is your feature request related to a problem?

This may not be a high priority issue, but I think it is worthwhile to document here:
When refactoring e3sm_diags with xcdat, I was working on using temporal.climatology operation to get annual cycle of a data stream which has 1hourly data for 30 years (ntime = 262800) from 1 grid point. Using xcdat is slower than the native function from e3sm_diags. It is probably expected, because the xcdat code is more sophisticated in reading in, writing out the data. I just noticed that it seems any operations on xarray.dataarray is much slower than numpy arrays, so the native e3sm_diags climo function needs to be optimized accordingly to get better performance. 

### Describe the solution you'd like

We can check if there are places where it is possible to get performance gain

### Describe alternatives you've considered

_No response_

### Additional context

Attaching example code and performance data.

```
### 1. Using temporal.climatology from xcdat
import xcdat as xc

file_path = '/global/cfs/cdirs/e3sm/e3sm_diags/postprocessed_e3sm_v2_data_for_e3sm_diags/20221103.v2.LR.amip.NGD_v3atm.chrysalis/arm-diags-data/PRECT_sgpc1_198501_201412.nc'
ds = xc.open_dataset(file_path)
ds_annual_cycle = ds.temporal.climatology("PRECT", "month")
#real    0m48.785s
#user    0m46.604s
#sys     0m12.709s

### 2. With external climo function from E3SM Diags
import xcdat as xc
from e3sm_diags.driver.utils.climo_xr import climo
file_path = '/global/cfs/cdirs/e3sm/e3sm_diags/postprocessed_e3sm_v2_data_for_e3sm_diags/20221103.v2.LR.amip.NGD_v3atm.chrysalis/arm-diags-data/PRECT_sgpc1_198501_201412.nc'
ds = xc.open_dataset(file_path)
ds_annual_cycle = climo(ds, "PRECT", "ANNUALCYCLE")

#real    0m23.995s
#user    0m19.630s
#sys     0m16.825s

### 3. Similar implementation but with cdat
import cdms2
from e3sm_diags.driver.utils.climo import climo
file_path = '/global/cfs/cdirs/e3sm/e3sm_diags/postprocessed_e3sm_v2_data_for_e3sm_diags/20221103.v2.LR.amip.NGD_v3atm.chrysalis/arm-diags-data/PRECT_sgpc1_198501_201412.nc'
ds = cdms2.open(file_path)
ds_annual_cycle = climo(ds("PRECT"), "ANNUALCYCLE")

#real    0m8.332s
#user    0m5.654s
#sys     0m11.613s

###4. Use cdat and cdutil for climatology
import cdms2
import cdutil

from e3sm_diags.driver.utils.climo import climo
file_path = '/global/cfs/cdirs/e3sm/e3sm_diags/postprocessed_e3sm_v2_data_for_e3sm_diags/20221103.v2.LR.amip.NGD_v3atm.chrysalis/arm-diags-data/PRECT_sgpc1_198501_201412.nc'ds = cdms2.open(file_path)
ds_annual_cycle = cdutil.ANNUALCYCLE.climatology(ds("PRECT"))

#real    3m40.007s
#user    3m34.030s
#sys     0m12.117s

```


Functions to edit:
xcdat/temporal.py:TemporalAccessor.departures
xcdat/temporal.py:TemporalAccessor._group_average
xcdat/temporal.py:TemporalAccessor._get_weights
xcdat/temporal.py:TemporalAccessor._group_data
xcdat/temporal.py:TemporalAccessor._label_time_coords
xcdat/temporal.py:TemporalAccessor._keep_weights
Functions added: xcdat/temporal.py:TemporalAccessor._calculate_departures

Patch:
diff --git a/.vscode/xcdat.code-workspace b/.vscode/xcdat.code-workspace
index 906eb68e..8eaa4e47 100644
--- a/.vscode/xcdat.code-workspace
+++ b/.vscode/xcdat.code-workspace
@@ -61,7 +61,7 @@
         "configurations": [
             {
                 "name": "Python: Current File",
-                "type": "python",
+                "type": "debugpy",
                 "request": "launch",
                 "program": "${file}",
                 "console": "integratedTerminal",
diff --git a/xcdat/temporal.py b/xcdat/temporal.py
index 834eb394..5fb0d736 100644
--- a/xcdat/temporal.py
+++ b/xcdat/temporal.py
@@ -747,16 +747,16 @@ def departures(
         # calling the group average or climatology APIs in step #3.
         ds = self._preprocess_dataset(ds)
 
-        # 3. Calculate the grouped average and climatology of the data variable.
+        # 3. Get the observational data variable.
         # ----------------------------------------------------------------------
-        # The climatology and grouped average APIs are called on the copied
-        # dataset to create separate instances of the `TemporalAccessor` class.
-        # This is done to avoid overriding the attributes of the current
-        # instance of `TemporalAccessor` (set in step #1 above).
+        # NOTE: The xCDAT APIs are called on copies of the original dataset to
+        # create separate instances of the `TemporalAccessor` class. This is
+        # done to avoid overriding the attributes of the current instance of
+        # `TemporalAccessor`, which is set by step #1.
+        ds_obs = ds.copy()
 
         # Group averaging is only required if the dataset's frequency (input)
         # differs from the `freq` arg (output).
-        ds_obs = ds.copy()
         inferred_freq = _infer_freq(ds[self.dim])
         if inferred_freq != freq:
             ds_obs = ds_obs.temporal.group_average(
@@ -767,7 +767,10 @@ def departures(
                 season_config,
             )
 
-        ds_climo = ds.temporal.climatology(
+        # 4. Calculate the climatology of the data variable.
+        # ----------------------------------------------------------------------
+        ds_climo = ds.copy()
+        ds_climo = ds_climo.temporal.climatology(
             data_var,
             freq,
             weighted,
@@ -776,44 +779,15 @@ def departures(
             season_config,
         )
 
-        # 4. Group the averaged data variable values by the time `freq`.
+        # 5. Calculate the departures for the data variable.
         # ----------------------------------------------------------------------
-        # This step allows us to perform xarray's grouped arithmetic to
-        # calculate departures.
-        dv_obs = ds_obs[data_var].copy()
-        self._labeled_time = self._label_time_coords(dv_obs[self.dim])
-        dv_obs_grouped = self._group_data(dv_obs)
-
-        # 5. Align time dimension names using the labeled time dimension name.
-        # ----------------------------------------------------------------------
-        # The climatology's time dimension is renamed to the labeled time
-        # dimension in step #4 above (e.g., "time" -> "season"). xarray requires
-        # dimension names to be aligned to perform grouped arithmetic, which we
-        # use for calculating departures in step #5. Otherwise, this error is
-        # raised: "`ValueError: incompatible dimensions for a grouped binary
-        # operation: the group variable '<FREQ ARG>' is not a dimension on the
-        # other argument`".
-        dv_climo = ds_climo[data_var]
-        dv_climo = dv_climo.rename({self.dim: self._labeled_time.name})
-
-        # 6. Calculate the departures for the data variable.
-        # ----------------------------------------------------------------------
-        # departures = observation - climatology
-        with xr.set_options(keep_attrs=True):
-            dv_departs = dv_obs_grouped - dv_climo
-            dv_departs = self._add_operation_attrs(dv_departs)
-            ds_obs[data_var] = dv_departs
-
-            # The original time dimension name is restored after grouped
-            # arithmetic, so the labeled time dimension name is no longer needed
-            # and therefore dropped.
-            ds_obs = ds_obs.drop_vars(str(self._labeled_time.name))
+        ds_departs = self._calculate_departures(ds_obs, ds_climo, data_var)
 
         if weighted and keep_weights:
-            self._weights = ds_climo.time_wts
-            ds_obs = self._keep_weights(ds_obs)
+            self._weights = ds_climo[f"{self.dim}_wts"]
+            ds_departs = self._keep_weights(ds_departs)
 
-        return ds_obs
+        return ds_departs
 
     def _averager(
         self,
@@ -1196,6 +1170,7 @@ def _group_average(self, ds: xr.Dataset, data_var: str) -> xr.DataArray:
         # Label the time coordinates for grouping weights and the data variable
         # values.
         self._labeled_time = self._label_time_coords(dv[self.dim])
+        dv = dv.assign_coords({self.dim: self._labeled_time})
 
         if self._weighted:
             time_bounds = ds.bounds.get_bounds("T", var_key=data_var)
@@ -1222,14 +1197,6 @@ def _group_average(self, ds: xr.Dataset, data_var: str) -> xr.DataArray:
         else:
             dv = self._group_data(dv).mean()
 
-        # After grouping and aggregating the data variable values, the
-        # original time dimension is replaced with the grouped time dimension.
-        # For example, grouping on "year_season" replaces the time dimension
-        # with "year_season". This dimension needs to be renamed back to
-        # the original time dimension name before the data variable is added
-        # back to the dataset so that the original name is preserved.
-        dv = dv.rename({self._labeled_time.name: self.dim})
-
         # After grouping and aggregating, the grouped time dimension's
         # attributes are removed. Xarray's `keep_attrs=True` option only keeps
         # attributes for data variables and not their coordinates, so the
@@ -1254,9 +1221,6 @@ def _get_weights(self, time_bounds: xr.DataArray) -> xr.DataArray:
         divided by the total sum of the time lengths in its group to get its
         corresponding weight.
 
-        The sum of the weights for each group is validated to ensure it equals
-        1.0.
-
         Parameters
         ----------
         time_bounds : xr.DataArray
@@ -1285,7 +1249,7 @@ def _get_weights(self, time_bounds: xr.DataArray) -> xr.DataArray:
         # or time unit (with rare exceptions see release notes). To avoid this
         # warning please use the scalar types `np.float64`, or string notation.`
         if isinstance(time_lengths.data, Array):
-            time_lengths.load()
+            time_lengths = time_lengths.astype("timedelta64[ns]")
 
         time_lengths = time_lengths.astype(np.float64)
 
@@ -1293,11 +1257,6 @@ def _get_weights(self, time_bounds: xr.DataArray) -> xr.DataArray:
         weights: xr.DataArray = grouped_time_lengths / grouped_time_lengths.sum()
         weights.name = f"{self.dim}_wts"
 
-        # Validate the sum of weights for each group is 1.0.
-        actual_sum = self._group_data(weights).sum().values
-        expected_sum = np.ones(len(grouped_time_lengths.groups))
-        np.testing.assert_allclose(actual_sum, expected_sum)
-
         return weights
 
     def _group_data(self, data_var: xr.DataArray) -> DataArrayGroupBy:
@@ -1321,8 +1280,8 @@ def _group_data(self, data_var: xr.DataArray) -> DataArrayGroupBy:
         if self._mode == "average":
             dv_gb = dv.groupby(f"{self.dim}.{self._freq}")
         else:
-            dv.coords[self._labeled_time.name] = self._labeled_time
-            dv_gb = dv.groupby(self._labeled_time.name)
+            dv = dv.assign_coords({self.dim: self._labeled_time})
+            dv_gb = dv.groupby(self.dim)
 
         return dv_gb
 
@@ -1374,9 +1333,9 @@ def _label_time_coords(self, time_coords: xr.DataArray) -> xr.DataArray:
         dt_objects = self._convert_df_to_dt(df_dt_components)
 
         time_grouped = xr.DataArray(
-            name="_".join(df_dt_components.columns),
+            name=self.dim,
             data=dt_objects,
-            coords={self.dim: time_coords[self.dim]},
+            coords={self.dim: dt_objects},
             dims=[self.dim],
             attrs=time_coords[self.dim].attrs,
         )
@@ -1669,6 +1628,10 @@ def _convert_df_to_dt(self, df: pd.DataFrame) -> np.ndarray:
     def _keep_weights(self, ds: xr.Dataset) -> xr.Dataset:
         """Keep the weights in the dataset.
 
+        The labeled time coordinates for the weights are replaced with the
+        original time coordinates and the dimension name is appended with
+        "_original".
+
         Parameters
         ----------
         ds : xr.Dataset
@@ -1679,16 +1642,13 @@ def _keep_weights(self, ds: xr.Dataset) -> xr.Dataset:
         xr.Dataset
             The dataset with the weights used for averaging.
         """
-        # Append "_original" to the name of the weights` time coordinates to
-        # avoid conflict with the grouped time coordinates in the Dataset (can
-        # have a different shape).
         if self._mode in ["group_average", "climatology"]:
-            self._weights = self._weights.rename({self.dim: f"{self.dim}_original"})
-            # Only keep the original time coordinates, not the ones labeled
-            # by group.
-            self._weights = self._weights.drop_vars(self._labeled_time.name)
+            weights = self._weights.assign_coords({self.dim: self._dataset[self.dim]})
+            weights = weights.rename({self.dim: f"{self.dim}_original"})
+        else:
+            weights = self._weights
 
-        ds[self._weights.name] = self._weights
+        ds[weights.name] = weights
 
         return ds
 
@@ -1733,6 +1693,59 @@ def _add_operation_attrs(self, data_var: xr.DataArray) -> xr.DataArray:
 
         return data_var
 
+    def _calculate_departures(
+        self,
+        ds_obs: xr.Dataset,
+        ds_climo: xr.Dataset,
+        data_var: str,
+    ) -> xr.Dataset:
+        """Calculate the departures for a data variable.
+
+        How this methods works:
+
+            1. Label the observational data variable's time coordinates by their
+               appropriate time group. For example, the first two time
+               coordinates 2000-01-01 and 2000-02-01 are replaced with the
+               "01-01-01" and "01-02-01" monthly groups.
+            2. Calculate departures by subtracting the climatology from the
+               labeled observational data using Xarray's grouped arithmetic with
+               automatic broadcasting (departures = obs - climo).
+            3. Restore the original time coordinates to the departures variable
+               to preserve the "year" of the time coordinates. For example,
+               the first two time coordinates 01-01-01 and 01-02-01 are reverted
+               back to 2000-01-01 and 2000-02-01.
+
+        Parameters
+        ----------
+        ds_obs : xr.Dataset
+            The observational dataset.
+        dv_climo : xr.Dataset
+            The climatology dataset.
+        data_var : str
+            The key of the data variable for calculating departures.
+
+        Returns
+        -------
+        xr.Dataset
+            The dataset containing the departures for a data variable.
+        """
+        ds_departs = ds_obs.copy()
+
+        dv_obs = ds_obs[data_var].copy()
+        self._labeled_time = self._label_time_coords(dv_obs[self.dim])
+        dv_obs_grouped = self._group_data(dv_obs)
+
+        dv_climo = ds_climo[data_var].copy()
+
+        with xr.set_options(keep_attrs=True):
+            dv_departs = dv_obs_grouped - dv_climo
+
+        dv_departs = self._add_operation_attrs(dv_departs)
+        dv_departs = dv_departs.assign_coords({self.dim: ds_obs[self.dim]})
+        ds_departs[data_var] = dv_departs
+
+        return ds_departs
+
 
 def _infer_freq(time_coords: xr.DataArray) -> Frequency:
     """Infers the time frequency from the coordinates.

Test patch: diff --git a/tests/test_temporal.py b/tests/test_temporal.py
index 6e2d6049..e5489b1b 100644
--- a/tests/test_temporal.py
+++ b/tests/test_temporal.py
@@ -121,7 +121,7 @@ def test_averages_for_yearly_time_series(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
         # Test unweighted averages
         result = ds.temporal.average("ts", weighted=False)
@@ -139,7 +139,7 @@ def test_averages_for_yearly_time_series(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_averages_for_monthly_time_series(self):
         # Set up dataset
@@ -293,7 +293,7 @@ def test_averages_for_daily_time_series(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
         # Test unweighted averages
         result = ds.temporal.average("ts", weighted=False)
@@ -310,7 +310,7 @@ def test_averages_for_daily_time_series(self):
                 "weighted": "False",
             },
         )
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_averages_for_hourly_time_series(self):
         ds = xr.Dataset(
@@ -378,7 +378,7 @@ def test_averages_for_hourly_time_series(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
         # Test unweighted averages
         result = ds.temporal.average("ts", weighted=False)
@@ -396,7 +396,7 @@ def test_averages_for_hourly_time_series(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
 
 class TestGroupAverage:
@@ -619,7 +619,7 @@ def test_weighted_seasonal_averages_with_DJF_and_drop_incomplete_seasons(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_weighted_seasonal_averages_with_DJF_without_dropping_incomplete_seasons(
         self,
@@ -670,7 +670,7 @@ def test_weighted_seasonal_averages_with_DJF_without_dropping_incomplete_seasons
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_weighted_seasonal_averages_with_JFD(self):
         ds = self.ds.copy()
@@ -729,7 +729,7 @@ def test_weighted_seasonal_averages_with_JFD(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_weighted_custom_seasonal_averages(self):
         ds = self.ds.copy()
@@ -787,7 +787,7 @@ def test_weighted_custom_seasonal_averages(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_raises_error_with_incorrect_custom_seasons_argument(self):
         # Test raises error with non-3 letter strings
@@ -873,7 +873,7 @@ def test_weighted_monthly_averages(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_weighted_monthly_averages_with_masked_data(self):
         ds = self.ds.copy()
@@ -924,7 +924,7 @@ def test_weighted_monthly_averages_with_masked_data(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_weighted_daily_averages(self):
         ds = self.ds.copy()
@@ -967,7 +967,7 @@ def test_weighted_daily_averages(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_weighted_hourly_averages(self):
         ds = self.ds.copy()
@@ -1011,7 +1011,7 @@ def test_weighted_hourly_averages(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
 
 class TestClimatology:
@@ -1105,7 +1105,7 @@ def test_subsets_climatology_based_on_reference_period(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_weighted_seasonal_climatology_with_DJF(self):
         ds = self.ds.copy()
@@ -1159,7 +1159,7 @@ def test_weighted_seasonal_climatology_with_DJF(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     @requires_dask
     def test_chunked_weighted_seasonal_climatology_with_DJF(self):
@@ -1214,7 +1214,7 @@ def test_chunked_weighted_seasonal_climatology_with_DJF(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_weighted_seasonal_climatology_with_JFD(self):
         ds = self.ds.copy()
@@ -1265,7 +1265,7 @@ def test_weighted_seasonal_climatology_with_JFD(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_weighted_custom_seasonal_climatology(self):
         ds = self.ds.copy()
@@ -1328,7 +1328,7 @@ def test_weighted_custom_seasonal_climatology(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_weighted_monthly_climatology(self):
         result = self.ds.temporal.climatology("ts", "month")
@@ -1391,7 +1391,7 @@ def test_weighted_monthly_climatology(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_unweighted_monthly_climatology(self):
         result = self.ds.temporal.climatology("ts", "month", weighted=False)
@@ -1453,7 +1453,7 @@ def test_unweighted_monthly_climatology(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_weighted_daily_climatology(self):
         result = self.ds.temporal.climatology("ts", "day", weighted=True)
@@ -1515,7 +1515,7 @@ def test_weighted_daily_climatology(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_weighted_daily_climatology_drops_leap_days_with_matching_calendar(self):
         time = xr.DataArray(
@@ -1606,7 +1606,7 @@ def test_weighted_daily_climatology_drops_leap_days_with_matching_calendar(self)
                 },
             )
 
-            assert result.identical(expected)
+            xr.testing.assert_identical(result, expected)
 
     def test_unweighted_daily_climatology(self):
         result = self.ds.temporal.climatology("ts", "day", weighted=False)
@@ -1668,7 +1668,7 @@ def test_unweighted_daily_climatology(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
 
 class TestDepartures:
@@ -1810,7 +1810,7 @@ def test_seasonal_departures_relative_to_climatology_reference_period(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_monthly_departures_relative_to_climatology_reference_period_with_same_output_freq(
         self,
@@ -1895,7 +1895,7 @@ def test_monthly_departures_relative_to_climatology_reference_period_with_same_o
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_weighted_seasonal_departures_with_DJF(self):
         ds = self.ds.copy()
@@ -1945,7 +1945,7 @@ def test_weighted_seasonal_departures_with_DJF(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_weighted_seasonal_departures_with_DJF_and_keep_weights(self):
         ds = self.ds.copy()
@@ -2021,7 +2021,7 @@ def test_weighted_seasonal_departures_with_DJF_and_keep_weights(self):
             dims=["time_original"],
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_unweighted_seasonal_departures_with_DJF(self):
         ds = self.ds.copy()
@@ -2071,7 +2071,7 @@ def test_unweighted_seasonal_departures_with_DJF(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_unweighted_seasonal_departures_with_JFD(self):
         ds = self.ds.copy()
@@ -2121,7 +2121,7 @@ def test_unweighted_seasonal_departures_with_JFD(self):
             },
         )
 
-        assert result.identical(expected)
+        xr.testing.assert_identical(result, expected)
 
     def test_weighted_daily_departures_drops_leap_days_with_matching_calendar(self):
         time = xr.DataArray(
@@ -2214,7 +2214,27 @@ def test_weighted_daily_departures_drops_leap_days_with_matching_calendar(self):
                 },
             )
 
-            assert result.identical(expected)
+            xr.testing.assert_identical(result, expected)
+
+
+def _check_each_weight_group_adds_up_to_1(ds: xr.Dataset, weights: xr.DataArray):
+    """Check that the sum of the weights in each group adds up to 1.0 (or 100%).
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The dataset with the temporal accessor class attached.
+    weights : xr.DataArray
+        The weights to check, produced by the `_get_weights` method.
+    """
+    time_lengths = ds.time_bnds[:, 1] - ds.time_bnds[:, 0]
+    time_lengths = time_lengths.astype(np.float64)
+
+    grouped_time_lengths = ds.temporal._group_data(time_lengths)
+
+    actual_sum = ds.temporal._group_data(weights).sum().values
+    expected_sum = np.ones(len(grouped_time_lengths.groups))
+    np.testing.assert_allclose(actual_sum, expected_sum)
 
 
 class Test_GetWeights:
@@ -2289,6 +2309,8 @@ def test_weights_for_yearly_averages(self):
             )
             assert np.allclose(result, expected)
 
+            _check_each_weight_group_adds_up_to_1(ds, result)
+
         def test_weights_for_monthly_averages(self):
             ds = self.ds.copy()
 
@@ -2352,6 +2374,8 @@ def test_weights_for_monthly_averages(self):
             )
             assert np.allclose(result, expected)
 
+            _check_each_weight_group_adds_up_to_1(ds, result)
+
     class TestWeightsForGroupAverageMode:
         @pytest.fixture(autouse=True)
         def setup(self):
@@ -2423,6 +2447,8 @@ def test_weights_for_yearly_averages(self):
             )
             assert np.allclose(result, expected)
 
+            _check_each_weight_group_adds_up_to_1(ds, result)
+
         def test_weights_for_monthly_averages(self):
             ds = self.ds.copy()
 
@@ -2469,6 +2495,8 @@ def test_weights_for_monthly_averages(self):
             expected = np.ones(15)
             assert np.allclose(result, expected)
 
+            _check_each_weight_group_adds_up_to_1(ds, result)
+
         def test_weights_for_seasonal_averages_with_DJF_and_drop_incomplete_seasons(
             self,
         ):
@@ -2631,6 +2659,8 @@ def test_weights_for_seasonal_averages_with_DJF_and_drop_incomplete_seasons(
             )
             assert np.allclose(result, expected, equal_nan=True)
 
+            _check_each_weight_group_adds_up_to_1(ds, result)
+
         def test_weights_for_seasonal_averages_with_JFD(self):
             ds = self.ds.copy()
 
@@ -2702,6 +2732,8 @@ def test_weights_for_seasonal_averages_with_JFD(self):
             )
             assert np.allclose(result, expected)
 
+            _check_each_weight_group_adds_up_to_1(ds, result)
+
         def test_custom_season_time_series_weights(self):
             ds = self.ds.copy()
 
@@ -2780,6 +2812,8 @@ def test_custom_season_time_series_weights(self):
             )
             assert np.allclose(result, expected)
 
+            _check_each_weight_group_adds_up_to_1(ds, result)
+
         def test_weights_for_daily_averages(self):
             ds = self.ds.copy()
 
@@ -2873,6 +2907,8 @@ def test_weights_for_hourly_averages(self):
             expected = np.ones(15)
             assert np.allclose(result, expected)
 
+            _check_each_weight_group_adds_up_to_1(ds, result)
+
     class TestWeightsForClimatologyMode:
         @pytest.fixture(autouse=True)
         def setup(self):
@@ -3035,6 +3071,8 @@ def test_weights_for_seasonal_climatology_with_DJF(self):
 
             assert np.allclose(result, expected, equal_nan=True)
 
+            _check_each_weight_group_adds_up_to_1(ds, result)
+
         def test_weights_for_seasonal_climatology_with_JFD(self):
             ds = self.ds.copy()
 
@@ -3101,6 +3139,8 @@ def test_weights_for_seasonal_climatology_with_JFD(self):
             )
             assert np.allclose(result, expected, equal_nan=True)
 
+            _check_each_weight_group_adds_up_to_1(ds, result)
+
         def test_weights_for_annual_climatology(self):
             ds = self.ds.copy()
 
@@ -3165,6 +3205,8 @@ def test_weights_for_annual_climatology(self):
             )
             assert np.allclose(result, expected)
 
+            _check_each_weight_group_adds_up_to_1(ds, result)
+
         def test_weights_for_daily_climatology(self):
             ds = self.ds.copy()
 
@@ -3227,6 +3269,8 @@ def test_weights_for_daily_climatology(self):
             )
             assert np.allclose(result, expected)
 
+            _check_each_weight_group_adds_up_to_1(ds, result)
+
 
 class Test_Averager:
     # NOTE: This private method is tested because it is more redundant to