Improve peak memory usage in Arakawa C grid conventions

Author: mx-moth

src/emsarray/conventions/arakawa_c.py

@@ -261,18 +261,31 @@ def pack_index(self, grid_kind: ArakawaCGridKind, indexes: Sequence[int]) -> Ara
         return cast(ArakawaCIndex, (grid_kind, *indexes))
 
     def _make_polygons(self) -> numpy.ndarray:
-        # Make an array of shape (j, i, 2) of all the nodes
-        grid = numpy.stack([self.node.longitude.values, self.node.latitude.values], axis=-1)
-
-        # Transform this in to an array of shape (topology.size, 4, 2)
-        points = numpy.stack([
-            grid[:-1, :-1],
-            grid[:-1, +1:],
-            grid[+1:, +1:],
-            grid[+1:, :-1],
-        ], axis=2).reshape((-1, 4, 2))
-
-        return utils.make_polygons_with_holes(points)
+        j_size, i_size = self.face.shape
+        longitude = self.node.longitude.values
+        latitude = self.node.latitude.values
+
+        # Preallocate the points array. We will copy data straight in to this
+        # to save repeated memory allocations.
+        points = numpy.empty(shape=(i_size, 4, 2), dtype=self.node.longitude.dtype)
+        # Preallocate the output array so we can fill it in batches
+        out = numpy.full(shape=self.face.size, fill_value=None, dtype=object)
+        # Construct polygons row by row
+        for j in range(self.face.shape[0]):
+            points[:, 0, 0] = longitude[j + 0, :-1]
+            points[:, 1, 0] = longitude[j + 0, +1:]
+            points[:, 2, 0] = longitude[j + 1, +1:]
+            points[:, 3, 0] = longitude[j + 1, :-1]
+
+            points[:, 0, 1] = latitude[j + 0, :-1]
+            points[:, 1, 1] = latitude[j + 0, +1:]
+            points[:, 2, 1] = latitude[j + 1, +1:]
+            points[:, 3, 1] = latitude[j + 1, :-1]
+
+            j_slice = slice(j * i_size, (j + 1) * i_size)
+            utils.make_polygons_with_holes(points, out=out[j_slice])
+
+        return out
 
     @cached_property
     def face_centres(self) -> numpy.ndarray:

tests/conventions/test_shoc_standard.py

@@ -1,7 +1,7 @@
 import itertools
 import json
-import pathlib
 import logging
+import pathlib
 
 import numpy
 import pandas
@@ -18,7 +18,8 @@
 from emsarray.conventions.shoc import ShocStandard
 from emsarray.operations import geometry
 from tests.utils import (
-    DiagonalShocGrid, ShocGridGenerator, ShocLayerGenerator, mask_from_strings, track_peak_memory_usage
+    DiagonalShocGrid, ShocGridGenerator, ShocLayerGenerator, mask_from_strings,
+    track_peak_memory_usage
 )
 
 logger = logging.getLogger(__name__)
@@ -670,8 +671,8 @@ def test_make_polygons_memory_usage():
     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 = 570_000_000
+    target = 133_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?"