Improve peak memory usage in unstructured grid conventions

Author: mx-moth

src/emsarray/conventions/ugrid.py

@@ -8,6 +8,7 @@
 import enum
 import logging
 import pathlib
+import math
 import warnings
 from collections import defaultdict
 from collections.abc import Hashable, Iterable, Sequence
@@ -1092,7 +1093,7 @@ def grid_kinds(self) -> frozenset[UGridKind]:
             items.append(UGridKind.edge)
         return frozenset(items)
 
-    def _make_polygons(self) -> numpy.ndarray:
+    def _make_polygons(self):
         """Generate list of Polygons"""
         topology = self.topology
         # X,Y coords of each node
@@ -1115,10 +1116,15 @@ def _make_polygons(self) -> numpy.ndarray:
         for unique_size in unique_sizes:
             # Extract the face node data for every polygon of this size
             indices = numpy.flatnonzero(polygon_sizes == unique_size)
-            nodes = numpy.ma.getdata(face_node)[indices, :unique_size]
-            coords = numpy.stack([node_x[nodes], node_y[nodes]], axis=-1)
-            # Generate the polygons directly in to their correct locations
-            shapely.polygons(coords, indices=indices, out=polygons)
+            chunk_size = 1000
+            chunk_count = math.ceil(len(indices) / chunk_size)
+            for chunk_index in range(chunk_count):
+                chunk_slice = slice(chunk_index * chunk_size, (chunk_index + 1) * chunk_size)
+                chunk_indices = indices[chunk_slice]
+                nodes = numpy.ma.getdata(face_node)[chunk_indices, :unique_size]
+                coords = numpy.stack([node_x[nodes], node_y[nodes]], axis=-1)
+                # Generate the polygons directly in to their correct locations
+                shapely.polygons(coords, indices=chunk_indices, out=polygons)
 
         return polygons
 

tests/conventions/test_ugrid.py

@@ -988,13 +988,13 @@ def test_has_valid_face_edge_connectivity():
 
 @pytest.mark.memory_usage
 def test_make_polygons_memory_usage():
-    dataset = make_dataset(width=500, height=400)
+    dataset = make_dataset(width=600, height=600)
 
     with track_peak_memory_usage() as tracker:
         assert len(dataset.ems.polygons) == dataset.ems.topology.face_count
 
     logger.info(f"current memory usage: %d, peak memory usage: %d", tracker.current, tracker.peak)
 
-    target = 124_000_000
+    target = 78_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?"