Improve peak memory usage in CF grid conventions

Author: mx-moth

src/emsarray/conventions/grid.py

@@ -413,39 +413,32 @@ def _make_polygons(self) -> numpy.ndarray:
         lon_bounds = self.topology.longitude_bounds.values
         lat_bounds = self.topology.latitude_bounds.values
 
-        # Make a bounds array as if this dataset had 2D coordinates.
-        # 1D bounds are (lon, 2) and (lat, 2).
-        # 2D bounds are (lat, lon, 4)
-        # where the 4 points are (j-i, i-1), (j-1, i+1), (j+1, i+1), (j+1, i-1).
-        # The bounds values are repeated as required, are given a new dimension,
-        # then repeated along that new dimension.
-        # They will come out as array with shape (y_size, x_size, 4)
-
-        lon_bounds_2d = numpy.stack([
-            lon_bounds[:, 0],
-            lon_bounds[:, 1],
-            lon_bounds[:, 1],
-            lon_bounds[:, 0],
-        ], axis=-1)
-        lon_bounds_2d = numpy.broadcast_to(numpy.expand_dims(lon_bounds_2d, 0), (y_size, x_size, 4))
-
-        lat_bounds_2d = numpy.stack([
-            lat_bounds[:, 0],
-            lat_bounds[:, 0],
-            lat_bounds[:, 1],
-            lat_bounds[:, 1],
-        ], axis=-1)
-        lat_bounds_2d = numpy.broadcast_to(numpy.expand_dims(lat_bounds_2d, 0), (x_size, y_size, 4))
-        lat_bounds_2d = numpy.transpose(lat_bounds_2d, (1, 0, 2))
-
-        assert lon_bounds_2d.shape == lat_bounds_2d.shape == (y_size, x_size, 4)
-
-        # points is a (topology.size, 4, 2) array of the corners of each cell
-        points = numpy.stack([lon_bounds_2d, lat_bounds_2d], axis=-1).reshape((-1, 4, 2))
-
-        polygons = utils.make_polygons_with_holes(points)
-
-        return polygons
+        # Create the polygons batched by row.
+        # The point array is copied by shapely before being used,
+        # so this can accidentally use a whole bunch of memory for large datasets.
+        # Creating them one by one is very slow but very memory efficient.
+        # Creating the polygons in one batch is faster but uses up a huge amount of memory.
+        # Batching them row by row is a decent compromise.
+        out = numpy.full(shape=y_size * x_size, dtype=object, fill_value=None)
+
+        # By preallocating this array, we can copy data in to it to save on a number of allocations.
+        chunk_points = numpy.empty(shape=(x_size, 4, 2), dtype=lon_bounds.dtype)
+        # By chunking by row, the longitude bounds never change between loops
+        chunk_points[:, 0, 0] = lon_bounds[:, 0]
+        chunk_points[:, 1, 0] = lon_bounds[:, 1]
+        chunk_points[:, 2, 0] = lon_bounds[:, 1]
+        chunk_points[:, 3, 0] = lon_bounds[:, 0]
+
+        for row in range(y_size):
+            chunk_points[:, 0, 1] = lat_bounds[row, 0]
+            chunk_points[:, 1, 1] = lat_bounds[row, 0]
+            chunk_points[:, 2, 1] = lat_bounds[row, 1]
+            chunk_points[:, 3, 1] = lat_bounds[row, 1]
+
+            row_slice = slice(row * x_size, (row + 1) * x_size)
+            utils.make_polygons_with_holes(chunk_points, out=out[row_slice])
+
+        return out
 
     @cached_property
     def face_centres(self) -> numpy.ndarray:
@@ -597,13 +590,22 @@ def check_dataset(cls, dataset: xarray.Dataset) -> int | None:
 
     def _make_polygons(self) -> numpy.ndarray:
         # Construct polygons from the bounds of the cells
-        lon_bounds = self.topology.longitude_bounds.values
-        lat_bounds = self.topology.latitude_bounds.values
-
-        # points is a (topology.size, 4, 2) array of the corners of each cell
-        points = numpy.stack([lon_bounds, lat_bounds], axis=-1).reshape((-1, 4, 2))
-
-        return utils.make_polygons_with_holes(points)
+        j_size, i_size = self.topology.shape
+        lon_bounds = self.topology.longitude_bounds
+        lat_bounds = self.topology.latitude_bounds
+
+        assert lon_bounds.shape == (j_size, i_size, 4)
+        assert lat_bounds.shape == (j_size, i_size, 4)
+
+        chunk_points = numpy.empty(shape=(i_size, 4, 2), dtype=lon_bounds.dtype)
+        out = numpy.full(shape=j_size * i_size, dtype=object, fill_value=None)
+        for j in range(j_size):
+            chunk_points[:, :, 0] = lon_bounds[j, :, :]
+            chunk_points[:, :, 1] = lat_bounds[j, :, :]
+            chunk_slice = slice(j * i_size, (j + 1) * i_size)
+            utils.make_polygons_with_holes(chunk_points, out=out[chunk_slice])
+
+        return out
 
     @cached_property
     def face_centres(self) -> numpy.ndarray:

tests/conventions/test_cfgrid1d.py

@@ -1,6 +1,6 @@
 import json
-import pathlib
 import logging
+import pathlib
 
 import numpy
 import pandas
@@ -16,7 +16,9 @@
     CFGrid1D, CFGrid1DTopology, CFGridKind, CFGridTopology
 )
 from emsarray.operations import geometry
-from tests.utils import assert_property_not_cached, box, mask_from_strings, track_peak_memory_usage
+from tests.utils import (
+    assert_property_not_cached, box, mask_from_strings, track_peak_memory_usage
+)
 
 logger = logging.getLogger(__name__)
 
@@ -512,8 +514,8 @@ def test_make_polygon_memory_usage() -> None:
     with track_peak_memory_usage() as tracker:
         assert len(dataset.ems.polygons) == width * height
 
-    logger.info(f"current memory usage: %d, peak memory usage: %d", tracker.current, tracker.peak)
+    logger.info("current memory usage: %d, peak memory usage: %d", tracker.current, tracker.peak)
 
-    target = 537_000_000
+    target = 135_000_000
     assert tracker.peak < target, "Peak memory allocation is too large"
     assert tracker.peak > target * 0.9, "Peak memory allocation is suspiciously small - did you improve things?"

tests/conventions/test_cfgrid2d.py

@@ -7,8 +7,8 @@
 """
 import itertools
 import json
-import pathlib
 import logging
+import pathlib
 
 import numpy
 import pandas
@@ -26,7 +26,8 @@
 from emsarray.operations import geometry
 from tests.utils import (
     AxisAlignedShocGrid, DiagonalShocGrid, ShocGridGenerator,
-    ShocLayerGenerator, assert_property_not_cached, plot_geometry, track_peak_memory_usage,
+    ShocLayerGenerator, assert_property_not_cached, plot_geometry,
+    track_peak_memory_usage
 )
 
 logger = logging.getLogger(__name__)
@@ -509,8 +510,8 @@ def test_make_polygon_memory_usage() -> None:
     with track_peak_memory_usage() as tracker:
         assert len(dataset.ems.polygons) == j_size * i_size
 
-    logger.info(f"current memory usage: %d, peak memory usage: %d", tracker.current, tracker.peak)
+    logger.info("current memory usage: %d, peak memory usage: %d", tracker.current, tracker.peak)
 
-    target = 665_000_000
+    target = 300_000_000
     assert tracker.peak < target, "Peak memory allocation is too large"
     assert tracker.peak > target * 0.9, "Peak memory allocation is suspiciously small - did you improve things?"

tests/utils.py

@@ -1,6 +1,5 @@
 import abc
 import contextlib
-import dataclasses
 import importlib.metadata
 import itertools
 import tracemalloc
@@ -467,7 +466,6 @@ def plot_geometry(
     figure.savefig(out)
 
 
-
 class TracemallocTracker:
     _finished = False
     _usage = None