Avoid converting chunked data to Numpy arrays (i.e. .values calls)

test/core/test_topological_agg.py

@@ -1,5 +1,7 @@
 import uxarray as ux
 
+import numpy as np
+import numpy.testing as nt
 import pytest
 
 
@@ -32,3 +34,67 @@ def test_node_to_edge_aggs(gridpath):
         grid_reduction = getattr(uxds['areaTriangle'], agg_func)(destination='edge')
 
         assert 'n_edge' in grid_reduction.dims
+
+
+def test_node_to_face_dask_reproduces_numpy(gridpath):
+    # the numpy (eager) and dask (chunked) branches must agree
+    da = pytest.importorskip("dask")  # dask-backed branch requires dask
+    uxds = ux.open_dataset(gridpath("mpas", "QU", "oQU480.231010.nc"), gridpath("mpas", "QU", "oQU480.231010.nc"))
+    uxda = uxds['areaTriangle']
+
+    for agg_func in AGGS:
+        numpy_result = getattr(uxda, agg_func)(destination='face')
+        dask_result = getattr(uxda.chunk(), agg_func)(destination='face')
+
+        assert isinstance(dask_result.data, da.array.Array)
+        assert numpy_result.dims == dask_result.dims
+        assert numpy_result.dtype == dask_result.dtype
+        # both paths run the same kernel over the same partitions, so they must
+        # agree exactly -- a mere allclose would hide a reordering regression
+        nt.assert_array_equal(numpy_result.values, dask_result.values)
+
+
+def test_node_to_edge_dask_reproduces_numpy(gridpath):
+    # the numpy (eager) and dask (chunked) branches must agree
+    da = pytest.importorskip("dask")  # dask-backed branch requires dask
+    uxds = ux.open_dataset(gridpath("mpas", "QU", "oQU480.231010.nc"), gridpath("mpas", "QU", "oQU480.231010.nc"))
+    uxda = uxds['areaTriangle']
+
+    for agg_func in AGGS:
+        numpy_result = getattr(uxda, agg_func)(destination='edge')
+        dask_result = getattr(uxda.chunk(), agg_func)(destination='edge')
+
+        assert isinstance(dask_result.data, da.array.Array)
+        assert numpy_result.dims == dask_result.dims
+        assert numpy_result.dtype == dask_result.dtype
+        # both paths run the same kernel over the same partitions, so they must
+        # agree exactly -- a mere allclose would hide a reordering regression
+        nt.assert_array_equal(numpy_result.values, dask_result.values)
+
+
+@pytest.mark.parametrize("destination", ["face", "edge"])
+def test_node_aggs_dask_reproduces_numpy_blockwise(gridpath, destination):
+    # 'areaTriangle' is 1D, so chunking it leaves a single block and the
+    # blockwise machinery is never exercised. Add a leading dimension so the
+    # dask path really runs the kernel once per chunk.
+    da = pytest.importorskip("dask")  # dask-backed branch requires dask
+    uxds = ux.open_dataset(gridpath("mpas", "QU", "oQU480.231010.nc"), gridpath("mpas", "QU", "oQU480.231010.nc"))
+    uxgrid = uxds['areaTriangle'].uxgrid
+
+    rng = np.random.default_rng(0)
+    uxda = ux.UxDataArray(
+        rng.random((6, uxgrid.n_node)), dims=("lev", "n_node"), uxgrid=uxgrid, name="var"
+    )
+
+    for agg_func in AGGS:
+        numpy_result = getattr(uxda, agg_func)(destination=destination)
+        chunked = uxda.chunk({"lev": 2})
+        dask_result = getattr(chunked, agg_func)(destination=destination)
+
+        assert isinstance(dask_result.data, da.array.Array)
+        # three chunks along 'lev', so the kernel is applied three times
+        assert len(dask_result.chunks[0]) == 3
+
+        assert numpy_result.dims == dask_result.dims
+        assert numpy_result.dtype == dask_result.dtype
+        nt.assert_array_equal(numpy_result.values, dask_result.values)

test/grid/integrate/test_basic.py

@@ -28,6 +28,7 @@ def test_multi_dim(gridpath):
     assert integral.ndim == len(dims) - 1
     nt.assert_almost_equal(integral, np.ones((5, 5)) * 4 * np.pi)
 
+
 def test_integrate_crashes_when_nnode_equals_nface():
     """Ensure UxDataArray.integrate() crashes for non-face_centered data, even if n_node==n_face.
     regression test for issue #1616.
@@ -79,3 +80,92 @@ def test_integrate_crashes_when_nnode_equals_nface():
     # ensure integrate() crashes for non-face_centered data, even if n_node==n_face
     with pytest.raises(ValueError):
         uxarr.integrate()
+
+
+def _random_uxda(gridpath, shape, dims):
+    """Non-uniform test data, so a reordered reduction actually shows up."""
+    uxgrid = ux.open_grid(gridpath("ugrid", "outCSne30", "outCSne30.ug"))
+    rng = np.random.default_rng(0)
+    return ux.UxDataArray(
+        data=rng.random(shape), dims=dims, uxgrid=uxgrid, name='var2'
+    )
+
+
+@pytest.mark.parametrize(
+    "shape_dims",
+    [((5400,), ("n_face",)), ((4, 5, 5400), ("a", "b", "n_face"))],
+    ids=["single_dim", "multi_dim"],
+)
+def test_integrate_returns_uxdataarray_matching_numpy(gridpath, shape_dims):
+    """Both branches of integrate() return a UxDataArray matching a plain numpy dot product."""
+    dask_array = pytest.importorskip("dask.array")
+    shape, dims = shape_dims
+    uxda = _random_uxda(gridpath, shape, dims)
+
+    # independent reference, computed without going through integrate()
+    expected = uxda.values @ uxda.uxgrid.face_areas.values
+
+    numpy_integral = uxda.integrate()
+    dask_integral = uxda.chunk({"n_face": 100}).integrate()
+
+    # backing type is preserved: eager stays eager, chunked stays lazy
+    assert isinstance(numpy_integral.data, np.ndarray)
+    assert isinstance(dask_integral.data, dask_array.Array)
+
+    for integral in (numpy_integral, dask_integral):
+        assert isinstance(integral, ux.UxDataArray)
+        assert integral.uxgrid is uxda.uxgrid
+        assert integral.dims == dims[:-1]
+        nt.assert_allclose(integral.values, expected, rtol=1e-12, atol=0.0)
+
+
+@pytest.mark.parametrize(
+    "shape_dims",
+    [((5400,), ("n_face",)), ((4, 5, 5400), ("a", "b", "n_face"))],
+    ids=["single_dim", "multi_dim"],
+)
+def test_integrate_dask_reproduces_numpy_whole_face_dim(gridpath, shape_dims):
+    # 'n_face' in a single chunk: xr.dot performs one reduction per block, in the
+    # same order as the numpy path's einsum, so agreement must be exact.
+    dask_array = pytest.importorskip("dask.array")  # dask-backed branch requires dask
+    shape, dims = shape_dims
+    uxda = _random_uxda(gridpath, shape, dims)
+
+    numpy_integral = uxda.integrate()
+
+    for chunks in ({"n_face": -1}, {"a": 2, "n_face": -1}):
+        chunks = {d: s for d, s in chunks.items() if d in dims}
+        dask_integral = uxda.chunk(chunks).integrate()
+
+        assert isinstance(dask_integral.data, dask_array.Array)
+        assert numpy_integral.dims == dask_integral.dims
+        assert numpy_integral.dtype == dask_integral.dtype
+        nt.assert_array_equal(numpy_integral.values, dask_integral.values)
+
+
+@pytest.mark.parametrize(
+    "shape_dims",
+    [((5400,), ("n_face",)), ((4, 5, 5400), ("a", "b", "n_face"))],
+    ids=["single_dim", "multi_dim"],
+)
+def test_integrate_dask_reproduces_numpy_chunked_face_dim(gridpath, shape_dims):
+    # Splitting 'n_face' splits the reduction: xr.dot sums per-chunk partials, so
+    # the result differs from the numpy path's single einsum in the last few ULP.
+    # This is the one changed routine that is close-but-not-bitwise, hence a
+    # tolerance rather than assert_array_equal. Observed worst case is ~2e-15
+    # relative on this mesh, independent of chunk size.
+    dask_array = pytest.importorskip("dask.array")  # dask-backed branch requires dask
+    shape, dims = shape_dims
+    uxda = _random_uxda(gridpath, shape, dims)
+
+    numpy_integral = uxda.integrate()
+
+    for chunks in ({"n_face": 100}, {"n_face": 1350}, {"n_face": 7}):
+        dask_integral = uxda.chunk(chunks).integrate()
+
+        assert isinstance(dask_integral.data, dask_array.Array)
+        assert numpy_integral.dims == dask_integral.dims
+        assert numpy_integral.dtype == dask_integral.dtype
+        nt.assert_allclose(
+            numpy_integral.values, dask_integral.values, rtol=1e-12, atol=0.0
+        )

test/precomputed_weights_test.py

@@ -31,6 +31,38 @@ def _write_sparse_map(path: Path, source_size: int, destination_size: int) -> Pa
     return path
 
 
+def _write_dense_map(
+    path: Path, source_size: int, destination_size: int, nnz_per_row: int = 3
+) -> Path:
+    """A map with several non-trivial weights per destination row.
+
+    ``_write_sparse_map`` builds a pure permutation with unit weights, so the
+    sparse multiply never sums anything. Use this when the summation itself
+    matters.
+    """
+    rng = np.random.default_rng(0)
+    rows = np.repeat(np.arange(1, destination_size + 1), nnz_per_row).astype(np.int32)
+    cols = (rng.integers(0, source_size, size=destination_size * nnz_per_row) + 1).astype(
+        np.int32
+    )
+    values = rng.random(destination_size * nnz_per_row)
+
+    ds = xr.Dataset(
+        data_vars={
+            "row": (("n_s",), rows),
+            "col": (("n_s",), cols),
+            "S": (("n_s",), values),
+        },
+        coords={"n_s": np.arange(rows.size, dtype=np.int32)},
+    )
+    ds = ds.assign_coords(
+        n_a=np.arange(source_size, dtype=np.int32),
+        n_b=np.arange(destination_size, dtype=np.int32),
+    )
+    ds.to_netcdf(path)
+    return path
+
+
 def test_load_remap_weights_and_apply_vector(tmp_path, gridpath):
     grid = ux.open_grid(gridpath("ugrid", "quad-hexagon", "grid.nc"))
     weight_file = _write_sparse_map(
@@ -191,3 +223,37 @@ def test_remap_weights_cache_is_lru_bounded(tmp_path, gridpath):
         _write_sparse_map(path, grid.n_face, grid.n_face)
         load_remap_weights(path)
     assert len(_WEIGHTS_CACHE) == _WEIGHTS_CACHE_MAXSIZE
+
+
+def test_apply_weights_dask_reproduces_numpy(tmp_path, gridpath):
+    # the numpy (eager) and dask (blockwise) branches must agree. The source dim
+    # is forced into a single chunk, so the sparse multiply is never split and
+    # agreement must be exact.
+    dask_array = pytest.importorskip("dask.array")  # dask-backed branch requires dask
+    grid = ux.open_grid(gridpath("mpas", "QU", "mesh.QU.1920km.151026.nc"))
+    weight_file = _write_dense_map(tmp_path / "dense_map.nc", grid.n_face, grid.n_face)
+
+    rng = np.random.default_rng(1)
+    source = ux.UxDataArray(
+        xr.DataArray(
+            rng.random((6, grid.n_face)),
+            dims=["time", "n_face"],
+            name="temperature",
+            attrs={"units": "K"},
+        ),
+        uxgrid=grid,
+    )
+
+    numpy_result = source.remap.apply_weights(grid, weight_file)
+    assert isinstance(numpy_result.data, np.ndarray)
+
+    # including chunkings of the source dim, which _apply_weights must rechunk away
+    for chunks in ({"time": 2}, {"n_face": 50}, {"time": 2, "n_face": 50}):
+        dask_result = source.chunk(chunks).remap.apply_weights(grid, weight_file)
+
+        assert isinstance(dask_result.data, dask_array.Array)
+
+        assert numpy_result.dims == dask_result.dims
+        assert numpy_result.dtype == dask_result.dtype
+        nt.assert_array_equal(numpy_result.values, dask_result.values)
+        nt.assert_equal(dask_result.attrs["units"], "K")

test/test_cross_sections.py

@@ -226,3 +226,70 @@ def test_cross_section(gridpath, datasetpath):
         _ = uxds['RELHUM'].cross_section(start=(45, 45))
         _ = uxds['RELHUM'].cross_section(lon=45, end=(45, 45))
         _ = uxds['RELHUM'].cross_section()
+
+
+CROSS_SECTION_MODES = [
+    dict(start=(-45, -45), end=(45, 45)),
+    dict(lat=45),
+    dict(lon=45),
+    dict(lon=45, steps=3),
+]
+
+
+def test_cross_section_dask_reproduces_numpy(gridpath, datasetpath):
+    # the numpy (eager) and dask (lazy gather) branches must agree
+    dask_array = pytest.importorskip("dask.array")  # dask-backed branch requires dask
+    uxds = ux.open_dataset(gridpath("scrip", "ne30pg2", "grid.nc"), datasetpath("scrip", "ne30pg2", "data.nc"))
+    uxda = uxds['RELHUM']  # ('lev', 'n_face'), so the face dim is not the leading one
+
+    for kwargs in CROSS_SECTION_MODES:
+        numpy_result = uxda.cross_section(**kwargs)
+
+        # chunk the face dim as well: the gather must not rely on n_face being whole
+        for chunks in ({"lev": 8}, {"n_face": 4000}, {"lev": 8, "n_face": 4000}):
+            dask_result = uxda.chunk(chunks).cross_section(**kwargs)
+
+            # the accessor no longer calls .compute(), so the result stays lazy
+            assert isinstance(dask_result.data, dask_array.Array)
+
+            assert numpy_result.dims == dask_result.dims
+            assert numpy_result.dtype == dask_result.dtype
+            nt.assert_array_equal(numpy_result.values, dask_result.values)
+            nt.assert_array_equal(numpy_result.lat.values, dask_result.lat.values)
+            nt.assert_array_equal(numpy_result.lon.values, dask_result.lon.values)
+
+
+def test_cross_section_dask_reproduces_numpy_partial_coverage(gridpath, datasetpath):
+    # steps with no containing face become NaN; that fill path must match too
+    dask_array = pytest.importorskip("dask.array")  # dask-backed branch requires dask
+    uxds = ux.open_dataset(gridpath("ugrid", "quad-hexagon", "grid.nc"),
+                           datasetpath("ugrid", "quad-hexagon", "multi_dim_data.nc"))
+    uxda = uxds['multi_dim_data']  # ('time', 'lev', 'n_face') over a 4-face regional patch
+
+    # arc starts inside the patch and leaves it, so the result mixes data with NaN
+    kwargs = dict(start=(-0.043, -0.112), end=(5, 5))
+    numpy_result = uxda.cross_section(**kwargs)
+    nan_mask = np.isnan(numpy_result.values)
+    assert nan_mask.any() and not nan_mask.all()
+
+    for chunks in ({"time": 2}, {"time": 2, "lev": 3}, {"time": 2, "n_face": 1}):
+        dask_result = uxda.chunk(chunks).cross_section(**kwargs)
+
+        assert isinstance(dask_result.data, dask_array.Array)
+
+        assert numpy_result.dims == dask_result.dims
+        assert numpy_result.dtype == dask_result.dtype
+        nt.assert_array_equal(numpy_result.values, dask_result.values)
+
+    # 'n_face' in a leading position: 'steps' must land in its place. This is the
+    # case the replaced numpy path handled with an explicit moveaxis round-trip.
+    leading = uxda.transpose("n_face", "time", "lev")
+    leading_numpy = leading.cross_section(**kwargs)
+    leading_dask = leading.chunk({"time": 2}).cross_section(**kwargs)
+
+    assert leading_numpy.dims == ("steps", "time", "lev")
+    assert leading_numpy.dims == leading_dask.dims
+    nt.assert_array_equal(leading_numpy.values, leading_dask.values)
+    nt.assert_array_equal(
+        leading_numpy.transpose(*numpy_result.dims).values, numpy_result.values
+    )