From 428675fef6998d47cf31892883446d46b7c10c8f Mon Sep 17 00:00:00 2001
From: Stephen Po-Chedley <pochedley@gmail.com>
Date: Wed, 17 Aug 2022 17:54:17 -0700
Subject: [PATCH] initial fix for #304

---
 tests/test_temporal.py |  2 +-
 xcdat/temporal.py      | 24 +++++++++++++++++++++---
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/tests/test_temporal.py b/tests/test_temporal.py
index 67a5fcea..a233355d 100644
--- a/tests/test_temporal.py
+++ b/tests/test_temporal.py
@@ -831,7 +831,7 @@ def test_weighted_monthly_averages_with_masked_data(self):
         expected = expected.drop_dims("time")
         expected["ts"] = xr.DataArray(
             name="ts",
-            data=np.array([[[2.0]], [[0.0]], [[1.0]], [[1.0]], [[2.0]]]),
+            data=np.array([[[2.0]], [[np.nan]], [[1.0]], [[1.0]], [[2.0]]]),
             coords={
                 "lat": expected.lat,
                 "lon": expected.lon,
diff --git a/xcdat/temporal.py b/xcdat/temporal.py
index 0fe8d905..40003c94 100644
--- a/xcdat/temporal.py
+++ b/xcdat/temporal.py
@@ -919,10 +919,28 @@ def _group_average(self, data_var: xr.DataArray) -> xr.DataArray:
 
         if self._weighted:
             self._weights = self._get_weights()
+            # grab metadata
+            dv_attrs = dv.attrs
+            dv_name = dv.name
+            # grab metadata (populated after calling ._group_data())
+            lt_attrs = self._labeled_time.attrs
+            lt_encoding = self._labeled_time.encoding
+            # weight the data variable
             dv *= self._weights
-            dv = self._group_data(dv).sum()
+            # cast the weights to match the dv.shape / dims
+            weights, x = xr.broadcast(self._weights, dv)
+            # ensure missing data receives no weight
+            weights = xr.where(np.isnan(dv), 0.0, weights)
+            # perform weighted average
+            dv = self._group_data(dv).sum() / self._group_data(weights).sum()
+            # add dv attributes
+            dv.attrs = dv_attrs
+            dv.name = dv_name
         else:
             dv = self._group_data(dv).mean()
+            # grab metadata (populated after calling ._group_data())
+            lt_attrs = self._labeled_time.attrs
+            lt_encoding = self._labeled_time.encoding
 
         # After grouping and aggregating the data variable values, the
         # original time dimension is replaced with the grouped time dimension.
@@ -936,8 +954,8 @@ def _group_average(self, data_var: xr.DataArray) -> xr.DataArray:
         # attributes are removed. Xarray's `keep_attrs=True` option only keeps
         # attributes for data variables and not their coordinates, so the
         # coordinate attributes have to be restored manually.
-        dv[self._dim].attrs = self._labeled_time.attrs
-        dv[self._dim].encoding = self._labeled_time.encoding
+        dv[self._dim].attrs = lt_attrs
+        dv[self._dim].encoding = lt_encoding
 
         dv = self._add_operation_attrs(dv)