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')