ASV better memory benchmarks

benchmarks/face_bounds.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 
 import uxarray as ux
+from .helpers._memsize import grid_nbytes
 
 current_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[0]
 
@@ -25,6 +26,31 @@ def time_face_bounds(self, grid_path):
         """Time to obtain ``Grid.face_bounds``"""
         self.uxgrid.bounds
 
-    def peakmem_face_bounds(self, grid_path):
-        """Peak memory usage obtain ``Grid.face_bounds."""
-        face_bounds = self.uxgrid.bounds
+    def track_nbytes_face_bounds(self, grid_path):
+        """Size of the materialized ``Grid.face_bounds`` array."""
+        return self.uxgrid.bounds.nbytes
+
+    track_nbytes_face_bounds.unit = "bytes"
+
+    def track_nbytes_grid_with_bounds(self, grid_path):
+        """Grid footprint after populating bounds -- catches cached arrays that
+        ``bounds`` adds to the ``Grid`` beyond the returned array itself."""
+        self.uxgrid.bounds
+        return grid_nbytes(self.uxgrid)
+
+    track_nbytes_grid_with_bounds.unit = "bytes"
+
+
+class FaceBoundsPeakMem:
+    """Peak memory of a cold start: import uxarray, open a grid, get its bounds."""
+
+    params = FaceBounds.params
+    param_names = ["grid_path"]
+
+    def setup_cache(self):
+        """Compile the njit kernels before anything is measured."""
+        for grid_path in self.params:
+            ux.open_grid(grid_path).bounds
+
+    def peakmem_open_and_bounds(self, grid_path):
+        ux.open_grid(grid_path).bounds

benchmarks/helpers/_memsize.py

@@ -0,0 +1,12 @@
+
+__all__ = ["grid_nbytes", "dataset_nbytes"]
+
+
+def grid_nbytes(uxgrid):
+    """Total size of the arrays a ``Grid`` currently holds, in bytes."""
+    return uxgrid._ds.nbytes
+
+
+def dataset_nbytes(uxds):
+    """Total size of a ``UxDataset``, in bytes, including its grid."""
+    return uxds.nbytes + grid_nbytes(uxds.uxgrid)

benchmarks/import.py

@@ -3,3 +3,7 @@ class Imports:
 
     def timeraw_import_uxarray(self):
         return "import uxarray"
+
+    def peakmem_import_uxarray(self):
+        """Peak memory of a process that has imported uxarray."""
+        import uxarray  # noqa: F401

benchmarks/mpas_ocean.py

@@ -5,6 +5,7 @@
 import numpy as np
 
 import uxarray as ux
+from .helpers._memsize import grid_nbytes
 
 current_path = Path(os.path.dirname(os.path.realpath(__file__)))
 
@@ -60,16 +61,44 @@ class Gradient(DatasetBenchmark):
     def time_gradient(self, resolution):
         self.uxds[data_var].gradient()
 
-    def peakmem_gradient(self, resolution):
-        grad = self.uxds[data_var].gradient()
+    def track_nbytes_gradient(self, resolution):
+        """Size of the gradient result."""
+        return self.uxds[data_var].gradient().nbytes
+
+    track_nbytes_gradient.unit = "bytes"
 
 
 class Integrate(DatasetBenchmark):
     def time_integrate(self, resolution):
         self.uxds[data_var].integrate()
 
-    def peakmem_integrate(self, resolution):
-        integral = self.uxds[data_var].integrate()
+    def track_nbytes_integrate(self, resolution):
+        """Grid footprint after integrating."""
+        self.uxds[data_var].integrate()
+        return grid_nbytes(self.uxds.uxgrid)
+
+    track_nbytes_integrate.unit = "bytes"
+
+
+class GradientPeakMem:
+    """Peak memory of a cold start: import uxarray, open a dataset, take a gradient.
+
+    Not a :class:`DatasetBenchmark` subclass -- that would open the dataset in
+    ``setup``, and asv counts setup memory towards ``peakmem_*``.
+    """
+
+    param_names = ["resolution"]
+    params = [["480km", "120km"]]
+
+    def setup_cache(self):
+        """Compile the njit kernels before anything is measured."""
+        for resolution in self.params[0]:
+            grid, data = file_path_dict[resolution]
+            ux.open_dataset(grid, data)[data_var].gradient()
+
+    def peakmem_gradient(self, resolution):
+        grid, data = file_path_dict[resolution]
+        ux.open_dataset(grid, data)[data_var].gradient()
 
 
 class GeoDataFrame(DatasetBenchmark):

benchmarks/quad_hexagon.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 
 import uxarray as ux
+from .helpers._memsize import dataset_nbytes, grid_nbytes
 
 current_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[0]
 
@@ -14,23 +15,18 @@ def time_open_grid(self):
         """Time to open a `Grid`"""
         ux.open_grid(grid_path)
 
-    # def mem_open_grid(self):
-    #     """Memory Occupied by a `Grid`"""
-    #     return ux.open_grid(grid_path)
-
-    def peakmem_open_grid(self):
-        """Peak memory usage of a `Grid`"""
-        uxgrid = ux.open_grid(grid_path)
+    def track_nbytes_open_grid(self):
+        """Memory occupied by a `Grid`"""
+        return grid_nbytes(ux.open_grid(grid_path))
 
+    track_nbytes_open_grid.unit = "bytes"
 
     def time_open_dataset(self):
         """Time to open a `UxDataset`"""
         ux.open_dataset(grid_path, data_path)
 
-    # def mem_open_dataset(self):
-    #     """Memory occupied by a `UxDataset`"""
-    #     return ux.open_dataset(grid_path, data_path)
+    def track_nbytes_open_dataset(self):
+        """Memory occupied by a `UxDataset`, including its grid"""
+        return dataset_nbytes(ux.open_dataset(grid_path, data_path))
 
-    def peakmem_open_dataset(self):
-        """Peak memory usage of a `UxDataset`"""
-        uxds = ux.open_dataset(grid_path, data_path)
+    track_nbytes_open_dataset.unit = "bytes"