Make longitude range lazy for lazy Grid constructors

benchmarks/lazy_grid_construction.py

@@ -0,0 +1,112 @@
+"""Cost of constructing a ``Grid`` from a chunked open.
+
+``LazyGridConstruction`` counts dask graph executions; ``OpenGridChunked`` times
+and peak-measures a chunked ``open_grid`` at two resolutions, which the rest of
+the suite, all eager, does not cover.
+"""
+
+import math
+import os
+import warnings
+from pathlib import Path
+
+from dask.callbacks import Callback
+
+import uxarray as ux
+
+from .helpers._fixtures import OQU_GRIDS, OQU_RESOLUTIONS
+from .helpers._peakmem import peak_allocated
+
+current_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[0]
+
+grid_path = current_path / "test" / "meshfiles" / "ugrid" / "outCSne30" / "outCSne30.ug"
+
+#: Small enough to leave several chunks per coordinate, which is what makes a
+#: reduction over one visible as a scheduled graph rather than a fused no-op.
+CHUNKS = {"n_node": 1000}
+
+
+class _CountComputes(Callback):
+    def __init__(self):
+        self.n = 0
+
+    def _start(self, dsk):
+        self.n += 1
+
+
+class LazyGridConstruction:
+    def setup(self):
+        # Opening a grid for the first time in a process pulls in xarray's
+        # backend machinery and the netCDF library
+        ux.open_grid(grid_path, chunks=CHUNKS)
+
+    def track_computes_open_grid_chunked(self):
+        """Dask graph executions during a chunked ``open_grid``.
+
+        Two remain after the lazy wrap, both from ``_standardize_connectivity``
+        checking connectivity for nulls.
+        """
+        counter = _CountComputes()
+        with counter:
+            ux.open_grid(grid_path, chunks=CHUNKS)
+        return counter.n
+
+    def track_computes_isel(self):
+        """Dask graph executions during ``isel`` on a chunked grid.
+
+        One remains, from ``_slice_face_indices`` loading the connectivity it slices by.
+        """
+        uxgrid = ux.open_grid(grid_path, chunks=CHUNKS)
+        counter = _CountComputes()
+        with counter:
+            uxgrid.isel(n_face=slice(0, 100))
+        return counter.n
+
+
+#: Chunks per grid dimension, held fixed across resolutions so the graph is
+#: the same shape at both and only the data under it grows.
+N_CHUNKS = 4
+
+
+class OpenGridChunked:
+    """``open_grid`` with ``chunks=`` across both oQU resolutions.
+
+    Compare each resolution to its own history: oQU480 is netCDF4/HDF5 and oQU120
+    is netCDF3, so 480km opens slower despite a sixteenth of the data.
+    """
+
+    param_names = ["resolution"]
+    params = [OQU_RESOLUTIONS]
+
+    def setup(self, resolution):
+        self.grid_path = OQU_GRIDS[resolution]
+        # An eager open for the sizes, which also pays the first-open cost of
+        # the netCDF backend here rather than in the first sample.
+        uxgrid = ux.open_grid(self.grid_path)
+        self.chunks = {
+            "n_node": math.ceil(uxgrid.n_node / N_CHUNKS),
+            "n_face": math.ceil(uxgrid.n_face / N_CHUNKS),
+        }
+        self._open()
+
+    def _open(self):
+        with warnings.catch_warnings():
+            # oQU480 stores layerThickness, ssh and zMid as one chunk of all
+            # 1791 cells, so any n_face chunking splits them and xarray warns
+            # once per open. They are data variables the grid reader drops.
+            warnings.filterwarnings(
+                "ignore", message="The specified chunks separate the stored chunks"
+            )
+            # A copy: open_grid adds the source-format dimension names to the
+            # dict it is handed (match_chunks_to_ugrid), so reusing one would
+            # have every sample after the first open with a different argument.
+            return ux.open_grid(self.grid_path, chunks=dict(self.chunks))
+
+    def time_open_grid(self, resolution):
+        self._open()
+
+    def track_peakmem_open_grid(self, resolution):
+        """Transient high-water allocation of a chunked ``open_grid``."""
+        return peak_allocated(self._open)
+
+    track_peakmem_open_grid.unit = "bytes"

test/grid/grid/test_io.py

@@ -6,6 +6,17 @@
 from uxarray.constants import ERROR_TOLERANCE
 
 
+def _assert_lon_close(actual, desired, err_msg, atol=ERROR_TOLERANCE):
+    """Compare longitudes modulo 360, to an absolute tolerance.
+
+    Exodus round-trips pass through ``_xyz_to_lonlat_deg``, which returns -180.0
+    where the original keeps 180.0, and near-zero longitudes come back ~1e-14 off,
+    which a relative tolerance alone rejects.
+    """
+    diff = (np.asarray(actual) - np.asarray(desired) + 180.0) % 360.0 - 180.0
+    np.testing.assert_allclose(diff, 0.0, atol=atol, err_msg=err_msg)
+
+
 def test_normalize_existing_coordinates_non_norm_initial(gridpath):
     from uxarray.grid.validation import _check_normalization
     uxgrid = ux.open_grid(gridpath("mpas", "QU", "mesh.QU.1920km.151026.nc"))
@@ -145,11 +156,10 @@ def test_grid_ugrid_exodus_roundtrip(gridpath):
             err_msg=f"UGRID longitude mismatch for {grid_name}",
             rtol=ERROR_TOLERANCE
         )
-        np.testing.assert_allclose(
+        _assert_lon_close(
             original_grid.node_lon.values,
             reloaded_exodus.node_lon.values,
             err_msg=f"Exodus longitude mismatch for {grid_name}",
-            rtol=ERROR_TOLERANCE
         )
         np.testing.assert_allclose(
             original_grid.node_lat.values,
@@ -185,11 +195,10 @@ def test_exodus_roundtrip_rll1deg_node_lonlat(gridpath, tmp_path):
 
     reloaded_exodus = ux.open_grid(exodus_filepath)
 
-    np.testing.assert_allclose(
+    _assert_lon_close(
         grid.node_lon.values,
         reloaded_exodus.node_lon.values,
         err_msg="Exodus longitude mismatch for RLL1deg",
-        rtol=ERROR_TOLERANCE,
     )
     np.testing.assert_allclose(
         grid.node_lat.values,

test/grid/grid/test_lon_range.py

@@ -0,0 +1,179 @@
+"""Tests for the elementwise longitude wrap in ``_set_desired_longitude_range``.
+
+It must stay lazy on dask-backed grids and must not grow the graph when called
+repeatedly.
+"""
+
+import numpy as np
+import numpy.testing as nt
+import pytest
+import xarray as xr
+from dask.callbacks import Callback
+
+import uxarray as ux
+from uxarray.grid.coordinates import (
+    _lon_within_range,
+    _set_desired_longitude_range,
+)
+
+
+class _CountComputes(Callback):
+    """Counts dask graph executions inside the ``with`` block."""
+
+    def __init__(self):
+        self.n = 0
+
+    def _start(self, dsk):
+        self.n += 1
+
+
+class _Shim:
+    """The only attribute ``_set_desired_longitude_range`` needs of a Grid."""
+
+    def __init__(self, ds):
+        self._ds = ds
+
+
+def _shim(values, name="node_lon", **attrs):
+    return _Shim(xr.Dataset({name: (f"n_{name.split('_')[0]}", values, attrs)}))
+
+
+def test_wrap_is_elementwise_and_leaves_in_range_values_exact():
+    """Outside [-180, 180] wraps; inside passes through bit-for-bit."""
+    rng = np.random.default_rng(0)
+    values = np.concatenate(
+        [rng.uniform(0.0, 180.0, 500), rng.uniform(180.0, 360.0, 500)]
+    )
+
+    grid = _shim(values.copy())
+    _set_desired_longitude_range(grid)
+    wrapped = grid._ds["node_lon"].values
+
+    in_range = values <= 180.0
+    nt.assert_array_equal(wrapped[in_range], values[in_range])
+    nt.assert_array_equal(
+        wrapped[~in_range], (values[~in_range] + 180.0) % 360.0 - 180.0
+    )
+    assert wrapped.max() <= 180.0 and wrapped.min() >= -180.0
+
+
+def test_both_endpoints_of_the_antimeridian_are_left_alone():
+    """180.0 and -180.0 are both kept, whatever else the array holds.
+
+    ``antimeridian_face_indices`` detects a crossing from a face's longitude span,
+    which folding 180 to -180 would collapse.
+    """
+    for companion in (10.0, 270.0, -170.0):
+        grid = _shim(np.array([180.0, companion]))
+        _set_desired_longitude_range(grid)
+        assert grid._ds["node_lon"].values[0] == 180.0
+
+    grid = _shim(np.array([-180.0, 0.0, 179.5]))
+    _set_desired_longitude_range(grid)
+    nt.assert_array_equal(grid._ds["node_lon"].values, [-180.0, 0.0, 179.5])
+
+
+def test_wrap_normalizes_both_tails_and_is_idempotent():
+    """One pass lands everything in [-180, 180], below -180 too; a second is a no-op."""
+    rng = np.random.default_rng(1)
+    grid = _shim(rng.uniform(-720.0, 720.0, 1000))
+
+    _set_desired_longitude_range(grid)
+    once = grid._ds["node_lon"].values.copy()
+    assert once.min() >= -180.0 and once.max() <= 180.0
+
+    grid._wrapped_lon_vars = {}
+    _set_desired_longitude_range(grid)
+    nt.assert_array_equal(grid._ds["node_lon"].values, once)
+
+
+@pytest.mark.parametrize("dtype", [np.float64, np.float32])
+def test_wrap_preserves_dtype_name_and_attrs(dtype):
+    grid = _shim(np.array([0.0, 180.0, 270.0], dtype=dtype), units="degrees_east")
+    _set_desired_longitude_range(grid)
+    out = grid._ds["node_lon"]
+
+    assert out.dtype == dtype
+    assert out.name == "node_lon"
+    assert out.attrs == {"units": "degrees_east"}
+
+
+def test_nan_longitudes_survive_the_wrap():
+    """``nan > 180`` is False, so NaN takes the pass-through branch."""
+    grid = _shim(np.array([np.nan, 270.0, 10.0]))
+    _set_desired_longitude_range(grid)
+
+    nt.assert_array_equal(grid._ds["node_lon"].values, [np.nan, -90.0, 10.0])
+
+
+def test_open_grid_does_not_compute_longitudes(gridpath):
+    """The wrap runs no dask compute on a chunked grid, the regression fixed here."""
+    path = gridpath("ugrid", "outCSne30", "outCSne30.ug")
+
+    grid = ux.open_grid(path, chunks={"n_node": 1000})
+    assert grid._ds["node_lon"].chunks is not None, "fixture is not chunked"
+
+    counter = _CountComputes()
+    with counter:
+        _set_desired_longitude_range(grid)
+    assert counter.n == 0
+
+
+def test_repeated_calls_do_not_grow_the_graph(gridpath):
+    """Repeated calls, as ``edge_lat`` makes, add no ``where`` layers to the graph."""
+    path = gridpath("ugrid", "outCSne30", "outCSne30.ug")
+    grid = ux.open_grid(path, chunks={"n_node": 1000})
+
+    before = len(grid._ds["node_lon"].data.dask)
+    for _ in range(20):
+        _set_desired_longitude_range(grid)
+    assert len(grid._ds["node_lon"].data.dask) == before
+
+
+def test_memo_reopens_when_the_variable_is_replaced(gridpath):
+    """The memo keys on the ``xr.Variable``, so assignment invalidates it."""
+    path = gridpath("ugrid", "outCSne30", "outCSne30.ug")
+    grid = ux.open_grid(path)
+
+    grid._ds["node_lon"] = grid._ds["node_lon"] + 200.0
+    assert float(grid._ds["node_lon"].max()) > 180.0
+
+    _set_desired_longitude_range(grid)
+    assert float(grid._ds["node_lon"].max()) <= 180.0
+
+
+def test_eager_fast_path_agrees_with_the_where_element_for_element():
+    """The in-memory guard must give exactly what the unguarded ``where`` gives."""
+    rng = np.random.default_rng(2)
+    cases = {
+        "strictly inside": rng.uniform(-179.0, 179.0, 500),
+        "on both endpoints": np.array([-180.0, 180.0, 0.0, 179.5, -179.5]),
+        "needs wrapping": rng.uniform(0.0, 360.0, 500),
+        "negative tail": rng.uniform(-540.0, -180.5, 500),
+        "with nan": np.array([np.nan, 10.0, 200.0]),
+    }
+
+    for label, values in cases.items():
+        da = xr.DataArray(values, dims="n_node", name="node_lon")
+        unguarded = xr.where(
+            (da > 180) | (da < -180), (da + 180) % 360 - 180, da
+        ).values
+
+        grid = _shim(values.copy())
+        _set_desired_longitude_range(grid)
+
+        nt.assert_array_equal(grid._ds["node_lon"].values, unguarded, err_msg=label)
+
+
+def test_guard_is_skipped_for_dask_backed_arrays(gridpath):
+    """``_lon_within_range`` computes on dask, so the wrap only runs it eagerly."""
+    grid = ux.open_grid(
+        gridpath("ugrid", "outCSne30", "outCSne30.ug"), chunks={"n_node": 1000}
+    )
+    da = grid._ds["node_lon"]
+    assert da.chunks is not None
+
+    counter = _CountComputes()
+    with counter:
+        _lon_within_range(da)
+    assert counter.n > 0, "guard is lazy here; skipping it would be pointless"

uxarray/grid/coordinates.py

@@ -724,9 +724,25 @@ def _is_projected_grid(uxgrid) -> bool:
     return False
 
 
+def _lon_within_range(da) -> bool:
+    """Whether all of ``da`` lies in [-180, 180], making the wrap the identity.
+
+    Lets in-memory arrays skip the elementwise wrap, which costs ~5x a reduction
+    plus ~18 bytes per node of temporaries. Only call it when ``da.chunks is
+    None``: on a dask array these reductions are the compute being avoided.
+    """
+    return bool(da.max() <= 180 and da.min() >= -180)
+
+
 def _set_desired_longitude_range(uxgrid):
     """Sets the longitude range to [-180, 180] for all longitude variables.
 
+    Wraps elementwise with ``xr.where``, so dask-backed coordinates stay lazy and
+    in-range values pass through bit-for-bit; both endpoints are kept because
+    ``antimeridian_face_indices`` relies on a vertex at 180. Each variable is
+    wrapped once, memoized on its ``xr.Variable``, since ``edge_lat`` calls this
+    on every access.
+
     Skipped entirely for projected grids: wrapping meter-scale coordinates
     as if they were degrees silently corrupts the geometry. A ``UserWarning``
     is issued once per Grid instance so users know which operations are invalid.
@@ -747,16 +763,26 @@ def _set_desired_longitude_range(uxgrid):
             uxgrid._projected_warning_issued = True
         return
 
+    memo = getattr(uxgrid, "_wrapped_lon_vars", None)
+    if memo is None:
+        memo = uxgrid._wrapped_lon_vars = {}
+
     with xr.set_options(keep_attrs=True):
         for lon_name in ["node_lon", "edge_lon", "face_lon"]:
             if lon_name in uxgrid._ds:
                 da = uxgrid._ds[lon_name]
                 if da.size == 0:
                     continue
-                if da.max() > 180:
-                    wrapped = (uxgrid._ds[lon_name] + 180) % 360 - 180
-                    wrapped.name = da.name
-                    uxgrid._ds[lon_name] = wrapped
+                if memo.get(lon_name) is da.variable:
+                    continue
+                if da.chunks is None and _lon_within_range(da):
+                    memo[lon_name] = da.variable
+                    continue
+                out_of_range = (da > 180) | (da < -180)
+                wrapped = xr.where(out_of_range, (da + 180) % 360 - 180, da)
+                wrapped.name = da.name
+                uxgrid._ds[lon_name] = wrapped
+                memo[lon_name] = uxgrid._ds[lon_name].variable
 
 
 def prepare_points(points, normalize):