Fix _keep_weights() for departures and average modes

xCDAT · Sep 3, 2024 · b97e8ae · b97e8ae
1 parent 0f4a1af
commit b97e8ae
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/xcdat/temporal.py b/xcdat/temporal.py
@@ -783,6 +783,10 @@ def departures(
         # ----------------------------------------------------------------------
         ds_departs = self._calculate_departures(ds_obs, ds_climo, data_var)
 
+        if keep_weights:
+            self._weights = ds_climo[f"{self.dim}_wts"]
+            ds_departs = self._keep_weights(ds_departs)
+
         return ds_departs
 
     def _averager(
@@ -1644,6 +1648,8 @@ def _keep_weights(self, ds: xr.Dataset) -> xr.Dataset:
         if self._mode in ["group_average", "climatology"]:
             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[weights.name] = weights