Bounding-box subset allocates memory proportional to the whole mesh, not the crop

[rajeeja]

Version

2026.9.0 (and main)

How did you install UXarray?

Source

What happened?

Two separate costs make a regional crop of a large SCRIP mesh scale with the whole mesh rather than the crop. The first is fixable in the subsetting code; the second is a reader/representation question and is the harder one.

1. Grid.bounds is computed for every face before any filtering.

subset.bounding_box filters on face_bounds_lon/lat, which come from _populate_face_bounds (uxarray/grid/bounds.py:84-96). That materializes the node coordinate and connectivity arrays with .values regardless of how the grid was opened. Measured, 12-corner SCRIP, crop fixed at ~0.006% of the mesh:

n_face	crop	memory
400,000	26 faces	+0.076 GiB
1,600,000	116 faces	+0.851 GiB

Asking for a smaller box does not help — the index is built before the box is consulted.

2. face_node_connectivity is a lazy graph over the dedup, so any slice replays it.

On the 63 GB CONUS-RRM np4 grid (299,999,162 faces), after open_grid leaves ~119 GiB resident, even a contiguous conn[:30M] kills a 236 GiB worker: computing the connectivity re-reads the 53.6 GiB corner arrays.

The decisive measurement is crop-size independence:

crop	candidates	outcome
Texas	30,914,134	killed on slice
Key West	15,866	killed on slice

A 2000× smaller crop fails identically, so the cost is not the crop.

For contrast, the pg2 twin — same 200 m geometry, 4 corners instead of 12 — subsets in 113 s.

What did you expect to happen?

That a regional crop costs something proportional to the region, and that a mesh which can be opened can also be subset.

Can you provide a MCVE to reproduce the bug?

"""Memory cost of a bbox subset scales with the mesh, not the crop."""
import gc, os, sys, tempfile, time
from pathlib import Path
import numpy as np, psutil, uxarray as ux, xarray as xr

P = psutil.Process(os.getpid())
rss = lambda: P.memory_info().rss / 1024**3
n_face, n_corner = int(sys.argv[1]), 12

tmp = Path(tempfile.mkdtemp()) / "mesh.nc"
rng = np.random.default_rng(0)
clon, clat = rng.uniform(0, 360, n_face), rng.uniform(-85, 85, n_face)
off = np.linspace(0, 2 * np.pi, n_corner, endpoint=False)
xr.Dataset({
    "grid_corner_lon": (("grid_size", "grid_corners"),
                        (clon[:, None] + 0.02 * np.cos(off)[None, :]) % 360.0),
    "grid_corner_lat": (("grid_size", "grid_corners"),
                        np.clip(clat[:, None] + 0.02 * np.sin(off)[None, :], -90, 90)),
    "grid_center_lon": (("grid_size",), clon),
    "grid_center_lat": (("grid_size",), clat),
    "grid_imask": (("grid_size",), np.ones(n_face, "i4")),
    "grid_area": (("grid_size",), np.zeros(n_face)),
}).to_netcdf(tmp)
del clon, clat; gc.collect()

grid = ux.open_grid(str(tmp), chunks="auto")
base = rss(); t0 = time.perf_counter()
sub = grid.subset.bounding_box(lon_bounds=[-1.0, 1.0], lat_bounds=[-1.0, 1.0])
print(f"n_face={n_face:,} crop={int(sub.n_face):,} +{rss()-base:.3f} GiB "
      f"{time.perf_counter()-t0:.1f}s")

Run at 400000 and 1600000: per-face cost stays flat while the crop stays at ~0.006%.

Possible directions

Not prescribing a fix; each is a design call.

1. Prefilter on face centres before computing exact bounds — addresses (1). Filter a bounding-box subset on face latitudes before computing bounds #1779 does this and measures 0.84 GiB → 0.00 GiB at 1.6M faces.
2. Persist connectivity after dedup, so slicing does not replay the graph — addresses (2).
3. Subset SCRIP at the file level, bypassing global connectivity entirely. grid_center_lat/lon gives candidates and grid_corner_lat/lon settles them exactly; neither needs node dedup. Prototyped outside uxarray against the 63 GB file: Key West returns 8,858 faces in 39 s at 0.47 GiB, verified to return the identical face set to subset.bounding_box on every SCRIP mesh small enough for both to run. The file is chunking=contiguous, so candidate rows are direct seeks.
4. Do not hold both coordinate representations: node_x/y/z are derivable from node_lon/lat and are ~46% of what survives a subset.

A narrower INT_DTYPE for connectivity would halve another array, but #1658 shows INT_FILL_VALUE does not survive that narrowing — mentioned only to rule it out.