DRAFT: Weighted Average

benchmarks/mpas_ocean.py

@@ -4,8 +4,6 @@
 
 import uxarray as ux
 
-import numpy as np
-
 current_path = Path(os.path.dirname(os.path.realpath(__file__)))
 
 data_var = 'bottomDepth'
@@ -29,74 +27,120 @@
                   "120km": [current_path / grid_filename_120, current_path / data_filename_120]}
 
 
+class Gradient:
 
-class DatasetBenchmark:
-    """Class used as a template for benchmarks requiring a ``UxDataset`` in
-    this module across both resolutions."""
-    param_names = ['resolution',]
-    params = [['480km', '120km'],]
+    param_names = ['resolution']
+    params = ['480km', '120km']
 
-    def setup(self, resolution, *args, **kwargs):
+
+    def setup(self, resolution):
         self.uxds = ux.open_dataset(file_path_dict[resolution][0], file_path_dict[resolution][1])
 
-    def teardown(self, resolution, *args, **kwargs):
+    def teardown(self, resolution):
         del self.uxds
 
-class GridBenchmark:
-    """Class used as a template for benchmarks requiring a ``Grid`` in this
-    module across both resolutions."""
-    param_names = ['resolution', ]
-    params = [['480km', '120km'], ]
-
-    def setup(self, resolution, *args, **kwargs):
-        self.uxgrid = ux.open_grid(file_path_dict[resolution][0])
-
-    def teardown(self, resolution, *args, **kwargs):
-        del self.uxgrid
-
-
-class Gradient(DatasetBenchmark):
     def time_gradient(self, resolution):
         self.uxds[data_var].gradient()
 
     def peakmem_gradient(self, resolution):
         grad = self.uxds[data_var].gradient()
 
-class Integrate(DatasetBenchmark):
+class Integrate:
+
+    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])
+
+    def teardown(self, resolution):
+        del self.uxds
+
     def time_integrate(self, resolution):
         self.uxds[data_var].integrate()
 
     def peakmem_integrate(self, resolution):
         integral = self.uxds[data_var].integrate()
 
+class GeoDataFrame:
+
+    param_names = ['resolution', 'exclude_antimeridian']
+    params = [['480km', '120km'],
+              [True, False]]
 
-class GeoDataFrame(DatasetBenchmark):
-    param_names = DatasetBenchmark.param_names + ['exclude_antimeridian']
-    params = DatasetBenchmark.params + [[True, False]]
 
+    def setup(self, resolution, exclude_antimeridian):
+        self.uxds = ux.open_dataset(file_path_dict[resolution][0], file_path_dict[resolution][1])
+
+    def teardown(self, resolution, exclude_antimeridian):
+        del self.uxds
 
     def time_to_geodataframe(self, resolution, exclude_antimeridian):
         self.uxds[data_var].to_geodataframe(exclude_antimeridian=exclude_antimeridian)
 
+    def peakmem_to_geodataframe(self, resolution, exclude_antimeridian):
+        gdf = self.uxds[data_var].to_geodataframe(exclude_antimeridian=exclude_antimeridian)
+
+
+class ConnectivityConstruction:
 
+    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])
+
+
+    def teardown(self, resolution):
+        del self.uxds
 
-class ConnectivityConstruction(DatasetBenchmark):
     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 MatplotlibConversion(DatasetBenchmark):
-    param_names = DatasetBenchmark.param_names + ['periodic_elements']
-    params = DatasetBenchmark.params + [['include', 'exclude', 'split']]
+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'].weighted_mean()
+
+class MatplotlibConversion:
+    param_names = ['resolution', 'periodic_elements']
+    params = (['480km', '120km'], ['include', 'exclude', 'split'])
+
+    def setup(self, resolution, periodic_elements):
+        self.uxds = ux.open_dataset(file_path_dict[resolution][0], file_path_dict[resolution][1])
+
+    def teardown(self, resolution, periodic_elements):
+        del self.uxds
 
     def time_dataarray_to_polycollection(self, resolution, periodic_elements):
         self.uxds[data_var].to_polycollection()
 
+class ConstructTreeStructures:
+    param_names = ['resolution']
+    params = ['480km', '120km']
 
-class ConstructTreeStructures(DatasetBenchmark):
+    def setup(self, resolution):
+        self.uxds = ux.open_dataset(file_path_dict[resolution][0], file_path_dict[resolution][1])
+
+    def teardown(self, resolution):
+        del self.uxds
 
     def time_kd_tree(self, resolution):
         self.uxds.uxgrid.get_kd_tree()
@@ -135,42 +179,3 @@ def time_nearest_neighbor_remapping(self):
 
     def time_inverse_distance_weighted_remapping(self):
         self.uxds_480["bottomDepth"].remap.inverse_distance_weighted(self.uxds_120.uxgrid)
-
-class HoleEdgeIndices(DatasetBenchmark):
-    def time_construct_hole_edge_indices(self, resolution):
-        ux.grid.geometry._construct_hole_edge_indices(self.uxds.uxgrid.edge_face_connectivity)
-
-
-class DualMesh(DatasetBenchmark):
-    def time_dual_mesh_construction(self, resolution):
-        self.uxds.uxgrid.get_dual()
-
-class ConstructFaceLatLon(GridBenchmark):
-    def time_welzl(self, resolution):
-        self.uxgrid.construct_face_centers(method='welzl')
-
-    def time_cartesian_averaging(self, resolution):
-        self.uxgrid.construct_face_centers(method='cartesian average')
-
-
-class CheckNorm:
-    param_names = ['resolution']
-    params = ['480km', '120km']
-
-    def setup(self, resolution):
-        self.uxgrid = ux.open_grid(file_path_dict[resolution][0])
-
-    def teardown(self, resolution):
-        del self.uxgrid
-
-    def time_check_norm(self, resolution):
-        from uxarray.grid.validation import _check_normalization
-        _check_normalization(self.uxgrid)
-
-
-class CrossSections(DatasetBenchmark):
-    param_names = DatasetBenchmark.param_names + ['n_lat']
-    params = DatasetBenchmark.params + [[1, 2, 4, 8]]
-    def time_constant_lat_fast(self, resolution, n_lat):
-        for lat in np.linspace(-89, 89, n_lat):
-            self.uxds.uxgrid.constant_latitude_cross_section(lat, method='fast')

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

@@ -55,6 +55,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'].weighted_mean()
+
+    # 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'].weighted_mean()
+
+    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'].weighted_mean()
+    unweighted_mean = uxds['psi'].mean()
+
+    # with equal area, both should be equal
+    nt.assert_equal(weighted_mean, unweighted_mean)