Add memory benchmarks (#5960)

* add memory benchmarks

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix memory benchmarks

* generalise benchmark repeats

* fix indent

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Patrick Peglar <[email protected]>

Author: 3 people

benchmarks/benchmarks/__init__.py

@@ -114,6 +114,33 @@ def _wrapper(*args, **kwargs):
         decorated_func.unit = "Mb"
         return _wrapper
 
+    @staticmethod
+    def decorator_repeating(repeats=3):
+        """Benchmark to track growth in resident memory during execution.
+
+        Tracks memory for repeated calls of decorated function.
+
+        Intended for use on ASV ``track_`` benchmarks. Applies the
+        :class:`TrackAddedMemoryAllocation` context manager to the benchmark
+        code, sets the benchmark ``unit`` attribute to ``Mb``.
+
+        """
+
+        def decorator(decorated_func):
+            def _wrapper(*args, **kwargs):
+                assert decorated_func.__name__[:6] == "track_"
+                # Run the decorated benchmark within the added memory context
+                # manager.
+                with TrackAddedMemoryAllocation() as mb:
+                    for _ in range(repeats):
+                        decorated_func(*args, **kwargs)
+                return mb.addedmem_mb()
+
+            decorated_func.unit = "Mb"
+            return _wrapper
+
+        return decorator
+
 
 def on_demand_benchmark(benchmark_object):
     """Disable these benchmark(s) unless ON_DEMAND_BENCHARKS env var is set.

benchmarks/benchmarks/merge_concat.py

@@ -8,6 +8,7 @@
 
 from iris.cube import CubeList
 
+from . import TrackAddedMemoryAllocation
 from .generate_data.stock import realistic_4d_w_everything
 
 
@@ -33,6 +34,10 @@ def setup(self):
     def time_merge(self):
         _ = self.cube_list.merge_cube()
 
+    @TrackAddedMemoryAllocation.decorator_repeating()
+    def track_mem_merge(self):
+        _ = self.cube_list.merge_cube()
+
 
 class Concatenate:
     # TODO: Improve coverage.
@@ -50,3 +55,7 @@ def setup(self):
 
     def time_concatenate(self):
         _ = self.cube_list.concatenate_cube()
+
+    @TrackAddedMemoryAllocation.decorator_repeating()
+    def track_mem_merge(self):
+        _ = self.cube_list.concatenate_cube()

benchmarks/benchmarks/regridding.py

@@ -14,6 +14,8 @@
 from iris.analysis import AreaWeighted, PointInCell
 from iris.coords import AuxCoord
 
+from . import TrackAddedMemoryAllocation
+
 
 class HorizontalChunkedRegridding:
     def setup(self) -> None:
@@ -51,6 +53,20 @@ def time_regrid_area_w_new_grid(self) -> None:
         # Realise data
         out.data
 
+    @TrackAddedMemoryAllocation.decorator_repeating()
+    def track_mem_regrid_area_w(self) -> None:
+        # Regrid the chunked cube
+        out = self.cube.regrid(self.template_cube, self.scheme_area_w)
+        # Realise data
+        out.data
+
+    @TrackAddedMemoryAllocation.decorator_repeating()
+    def track_mem_regrid_area_w_new_grid(self) -> None:
+        # Regrid the chunked cube
+        out = self.chunked_cube.regrid(self.template_cube, self.scheme_area_w)
+        # Realise data
+        out.data
+
 
 class CurvilinearRegridding:
     def setup(self) -> None:
@@ -93,3 +109,10 @@ def time_regrid_pic(self) -> None:
         out = self.cube.regrid(self.template_cube, self.scheme_pic)
         # Realise the data
         out.data
+
+    @TrackAddedMemoryAllocation.decorator_repeating()
+    def track_mem_regrid_pic(self) -> None:
+        # Regrid the cube onto the template.
+        out = self.cube.regrid(self.template_cube, self.scheme_pic)
+        # Realise the data
+        out.data

benchmarks/benchmarks/stats.py

@@ -8,6 +8,8 @@
 from iris.analysis.stats import pearsonr
 import iris.tests
 
+from . import TrackAddedMemoryAllocation
+
 
 class PearsonR:
     def setup(self):
@@ -30,9 +32,21 @@ def setup(self):
     def time_real(self):
         pearsonr(self.cube_a, self.cube_b, weights=self.weights)
 
+    @TrackAddedMemoryAllocation.decorator_repeating()
+    def track_real(self):
+        pearsonr(self.cube_a, self.cube_b, weights=self.weights)
+
     def time_lazy(self):
         for cube in self.cube_a, self.cube_b:
             cube.data = cube.lazy_data()
 
         result = pearsonr(self.cube_a, self.cube_b, weights=self.weights)
         result.data
+
+    @TrackAddedMemoryAllocation.decorator_repeating()
+    def track_lazy(self):
+        for cube in self.cube_a, self.cube_b:
+            cube.data = cube.lazy_data()
+
+        result = pearsonr(self.cube_a, self.cube_b, weights=self.weights)
+        result.data

benchmarks/benchmarks/trajectory.py

@@ -13,6 +13,8 @@
 import iris
 from iris.analysis.trajectory import interpolate
 
+from . import TrackAddedMemoryAllocation
+
 
 class TrajectoryInterpolation:
     def setup(self) -> None:
@@ -33,8 +35,22 @@ def time_trajectory_linear(self) -> None:
         # Realise the data
         out_cube.data
 
+    @TrackAddedMemoryAllocation.decorator_repeating()
+    def track_trajectory_linear(self) -> None:
+        # Regrid the cube onto the template.
+        out_cube = interpolate(self.cube, self.sample_points, method="linear")
+        # Realise the data
+        out_cube.data
+
     def time_trajectory_nearest(self) -> None:
         # Regrid the cube onto the template.
         out_cube = interpolate(self.cube, self.sample_points, method="nearest")
         # Realise the data
         out_cube.data
+
+    @TrackAddedMemoryAllocation.decorator_repeating()
+    def track_trajectory_nearest(self) -> None:
+        # Regrid the cube onto the template.
+        out_cube = interpolate(self.cube, self.sample_points, method="nearest")
+        # Realise the data
+        out_cube.data