Improve peak memory usage in unstructured grid conventions

Author: mx-moth

src/emsarray/conventions/ugrid.py

@@ -7,6 +7,7 @@
 """
 import enum
 import logging
+import math
 import pathlib
 import warnings
 from collections import defaultdict
@@ -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

@@ -1,7 +1,7 @@
 import json
+import logging
 import pathlib
 import warnings
-import logging
 
 import geojson
 import numpy
@@ -21,7 +21,9 @@
     ConventionViolationError, ConventionViolationWarning
 )
 from emsarray.operations import geometry
-from tests.utils import assert_property_not_cached, filter_warning, track_peak_memory_usage
+from tests.utils import (
+    assert_property_not_cached, filter_warning, track_peak_memory_usage
+)
 
 logger = logging.getLogger(__name__)
 
@@ -988,13 +990,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)
+    logger.info("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?"