DRAFT: Weighted Average

benchmarks/mpas_ocean.py

@@ -98,11 +98,28 @@ def teardown(self, resolution):
 
     def time_n_nodes_per_face(self, resolution):
         self.uxds.uxgrid.n_nodes_per_face
-
+        
     def time_face_face_connectivity(self, resolution):
         ux.grid.connectivity._populate_face_face_connectivity(self.uxds.uxgrid)
 
 
+class WeightedMean:
+
+    param_names = ['resolution']
+    params = ['480km', '120km']
+
+    def setup(self, resolution):
+        self.uxds = ux.open_dataset(file_path_dict[resolution][0], file_path_dict[resolution][1])
+        _ = self.uxds.uxgrid.face_areas
+        _ = self.uxds.uxgrid.edge_node_distances
+
+    def teardown(self, resolution):
+        del self.uxds
+
+
+    def time_weighted_mean_face_centered(self, resolution):
+        self.uxds['bottomDepth'].mean(weighted=True)
+
 class MatplotlibConversion:
     param_names = ['resolution', 'periodic_elements']
     params = (['480km', '120km'], ['include', 'exclude', 'split'])
@@ -116,7 +133,6 @@ def teardown(self, resolution, periodic_elements):
     def time_dataarray_to_polycollection(self, resolution, periodic_elements):
         self.uxds[data_var].to_polycollection()
 
-
 class ConstructTreeStructures:
     param_names = ['resolution']
     params = ['480km', '120km']

docs/user-guide/weighted-average.ipynb

@@ -0,0 +1,45 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# Weighted Average"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "4850ba8becab8b0d"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "ee8c405203e6faf"
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/userguide.rst

@@ -49,6 +49,9 @@ These user guides provide detailed explanations of the core functionality in UXa
 `Topological Aggregations <user-guide/topological-aggregations.ipynb>`_
  Aggregate data across grid dimensions
 
+`Weighted Average <user-guide/weighted-average.ipynb>`_
+ Compute the weighted average
+
 `Calculus Operators <user-guide/calculus.ipynb>`_
  Apply calculus operators (gradient, integral) on unstructured grid data
 

test/test_weighted_mean.py

@@ -0,0 +1,67 @@
+import uxarray as ux
+import os
+
+import numpy.testing as nt
+import numpy as np
+
+import pytest
+
+from pathlib import Path
+
+current_path = Path(os.path.dirname(os.path.realpath(__file__)))
+
+
+csne30_grid_path = current_path / 'meshfiles' / "ugrid" / "outCSne30" / "outCSne30.ug"
+csne30_data_path = current_path / 'meshfiles' / "ugrid" / "outCSne30" / "outCSne30_vortex.nc"
+
+quad_hex_grid_path = current_path / 'meshfiles' / "ugrid" / "quad-hexagon" / "grid.nc"
+quad_hex_data_path_face_centered = current_path / 'meshfiles' / "ugrid" / "quad-hexagon" / "data.nc"
+quad_hex_data_path_edge_centered = current_path / 'meshfiles' / "ugrid" / "quad-hexagon" / "random-edge-data.nc"
+
+
+def test_quad_hex_face_centered():
+    """Compares the weighted average computation for the quad hexagon grid
+    using a face centered data variable to the expected value computed by
+    hand."""
+    uxds = ux.open_dataset(quad_hex_grid_path, quad_hex_data_path_face_centered)
+
+    # expected weighted average computed by hand
+    expected_weighted_mean = 297.55
+
+    # compute the weighted mean
+    result = uxds['t2m'].mean(weighted=True)
+
+    # ensure values are within 3 decimal points of each other
+    nt.assert_almost_equal(result.values, expected_weighted_mean, decimal=3)
+
+def test_quad_hex_edge_centered():
+    """Compares the weighted average computation for the quad hexagon grid
+    using an edge centered data variable to the expected value computed by
+    hand."""
+    uxds = ux.open_dataset(quad_hex_grid_path, quad_hex_data_path_edge_centered)
+
+    # expected weighted average computed by hand
+    # expected_weighted_mean = 297.55
+    expected_weighted_mean = (uxds['random_data_edge'].values * uxds.uxgrid.edge_node_distances).sum() / uxds.uxgrid.edge_node_distances.sum()
+
+    # compute the weighted mean
+    result = uxds['random_data_edge'].mean(weighted=True)
+
+    nt.assert_equal(result, expected_weighted_mean)
+
+
+def test_csne30_equal_area():
+    """Compute the weighted average with a grid that has equal-area faces and
+    compare the result to the regular mean."""
+    uxds = ux.open_dataset(csne30_grid_path, csne30_data_path)
+    face_areas = uxds.uxgrid.face_areas
+
+    # set the area of each face to be one
+    uxds.uxgrid._ds['face_areas'].data = np.ones(uxds.uxgrid.n_face)
+
+
+    weighted_mean = uxds['psi'].mean(weighted=True)
+    unweighted_mean = uxds['psi'].mean(weighted=False)
+
+    # with equal area, both should be equal
+    nt.assert_equal(weighted_mean, unweighted_mean)

uxarray/core/dataarray.py

@@ -140,6 +140,12 @@ def uxgrid(self):
     def uxgrid(self, ugrid_obj):
         self._uxgrid = ugrid_obj
 
+    def from_xarray(self, data_array, uxgrid):
+        """Take a ``xarray.DataArray`` and convert it to a
+        ``uxarray.UxDataArray``"""
+
+        pass
+
     def to_geodataframe(self, override=False, cache=True, exclude_antimeridian=False):
         """Constructs a ``spatialpandas.GeoDataFrame`` with a "geometry"
         column, containing a collection of Shapely Polygons or MultiPolygons
@@ -411,6 +417,42 @@ def nodal_average(self):
 
         return self.topological_mean(destination="face")
 
+    # weighted_mean()
+    # .weighted.mean()
+    # weights are only on grid - not supporting custom weights for now 
+    def mean(self, dim=None, *, skipna=None, keep_attrs=None, weighted=False, **kwargs):
+        if weighted:
+            if self._face_centered():
+                grid_dim = "n_face"
+                # use face areas as weight
+                weights = da.from_array(self.uxgrid.face_areas.values)
+            elif self._edge_centered():
+                grid_dim = "n_edge"
+                # use edge magnitude as weight
+                weights = da.from_array(self.uxgrid.edge_node_distances.values)
+            else:
+                # apply regular Xarray mean
+                warnings.warn(
+                    "Attempting to perform a weighted mean calculation on a variable that does not have"
+                    "associated weights. Weighted mean is only supported for face or edge centered "
+                    "variables. Performing an unweighted mean."
+                )
+
+                return super().mean(dim=dim, skipna=skipna, keep_attrs=keep_attrs, **kwargs)
+
+            # compute the total weight
+            total_weight = weights.sum()
+
+            # compute weighted mean #assumption on index of dimension (last one is geometry)
+            weighted_mean = (self * weights).sum(dim=grid_dim) / total_weight
+
+            # create a UxDataArray and return it
+            return UxDataArray(weighted_mean, uxgrid=self.uxgrid)
+
+        else:
+            # apply regular Xarray mean
+            return super().mean(dim=None, skipna=None, keep_attrs=None, **kwargs)
+
     def topological_mean(
         self,
         destination: Literal["node", "edge", "face"],