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

test/core/test_topological_agg.py

@@ -1,5 +1,6 @@
 import uxarray as ux
 
+import numpy as np
 import pytest
 
 
@@ -32,3 +33,31 @@ 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
+    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 numpy_result.dims == dask_result.dims
+        assert np.allclose(numpy_result.values, dask_result.values, equal_nan=True)
+
+
+def test_node_to_edge_dask_reproduces_numpy(gridpath):
+    # the numpy (eager) and dask (chunked) branches must agree
+    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 numpy_result.dims == dask_result.dims
+        assert np.allclose(numpy_result.values, dask_result.values, equal_nan=True)

uxarray/core/aggregation.py

@@ -1,4 +1,5 @@
 import numpy as np
+import xarray as xr
 
 import uxarray.core.dataarray
 from uxarray.grid.connectivity import get_face_node_partitions
@@ -80,41 +81,40 @@ def _node_to_face_aggregation(uxda, aggregation, aggregation_func_kwargs):
 
     if isinstance(uxda.data, np.ndarray):
         # apply aggregation using numpy
-        aggregated_var = _apply_node_to_face_aggregation_numpy(
+        return _apply_node_to_face_aggregation_numpy(
             uxda, NUMPY_AGGREGATIONS[aggregation], aggregation_func_kwargs
         )
     elif isinstance(uxda.data, da.Array):
-        # apply aggregation on dask array, TODO:
-        aggregated_var = _apply_node_to_face_aggregation_numpy(
+        # apply aggregation lazily on a dask array
+        return _apply_node_to_face_aggregation_dask(
             uxda, NUMPY_AGGREGATIONS[aggregation], aggregation_func_kwargs
         )
     else:
         raise ValueError
 
-    return uxarray.core.dataarray.UxDataArray(
-        uxgrid=uxda.uxgrid,
-        data=aggregated_var,
-        dims=uxda.dims,
-        name=uxda.name,
-    ).rename({"n_node": "n_face"})
-
 
-def _apply_node_to_face_aggregation_numpy(
-    uxda, aggregation_func, aggregation_func_kwargs
+def _node_to_face_kernel(
+    data,
+    face_node_conn,
+    n_nodes_per_face,
+    n_face,
+    aggregation_func,
+    aggregation_func_kwargs,
 ):
-    """Applies a Node to Face Topological aggregation on a Numpy array."""
-    data = uxda.values
-    face_node_conn = uxda.uxgrid.face_node_connectivity.values
-    n_nodes_per_face = uxda.uxgrid.n_nodes_per_face.values
+    """Node-to-face topological aggregation on a single numpy block.
 
+    ``data`` has the node dimension as its last axis, shape ``(..., n_node)``;
+    the result has the face dimension as its last axis, shape ``(..., n_face)``.
+    Shared verbatim by the numpy path and the (blockwise) dask path.
+    """
     (
         change_ind,
         n_nodes_per_face_sorted_ind,
         element_sizes,
         size_counts,
     ) = get_face_node_partitions(n_nodes_per_face)
 
-    result = np.empty(shape=(data.shape[:-1]) + (uxda.uxgrid.n_face,))
+    result = np.empty(shape=(data.shape[:-1]) + (n_face,))
 
     for e, start, end in zip(element_sizes, change_ind[:-1], change_ind[1:]):
         face_inds = n_nodes_per_face_sorted_ind[start:end]
@@ -131,9 +131,60 @@ def _apply_node_to_face_aggregation_numpy(
     return result
 
 
-def _apply_node_to_face_aggregation_dask(*args, **kwargs):
-    """Applies a Node to Face Topological aggregation on a Dask array."""
-    pass
+def _apply_node_to_face_aggregation_numpy(
+    uxda, aggregation_func, aggregation_func_kwargs
+):
+    """Applies a Node to Face Topological aggregation on a Numpy array."""
+    aggregated_var = _node_to_face_kernel(
+        uxda.values,
+        uxda.uxgrid.face_node_connectivity.values,
+        uxda.uxgrid.n_nodes_per_face.values,
+        uxda.uxgrid.n_face,
+        aggregation_func,
+        aggregation_func_kwargs,
+    )
+    return uxarray.core.dataarray.UxDataArray(
+        uxgrid=uxda.uxgrid,
+        data=aggregated_var,
+        dims=uxda.dims,
+        name=uxda.name,
+    ).rename({"n_node": "n_face"})
+
+
+def _apply_node_to_face_aggregation_dask(
+    uxda, aggregation_func, aggregation_func_kwargs
+):
+    """Applies a Node to Face Topological aggregation on a Dask array, lazily"""
+
+    uxgrid = uxda.uxgrid
+    n_face = uxgrid.n_face
+
+    # n_node must live in a single chunk: any node can feed any face. The grid
+    # metadata are core-dim inputs, so their core dims must be single-chunk too.
+    da_in = uxda.chunk({"n_node": -1})
+    face_node_conn = uxgrid.face_node_connectivity.chunk(
+        {"n_face": -1, "n_max_face_nodes": -1}
+    )
+    n_nodes_per_face = uxgrid.n_nodes_per_face.chunk({"n_face": -1})
+
+    result = xr.apply_ufunc(
+        _node_to_face_kernel,
+        da_in,
+        face_node_conn,
+        n_nodes_per_face,
+        input_core_dims=[["n_node"], ["n_face", "n_max_face_nodes"], ["n_face"]],
+        output_core_dims=[["n_face"]],
+        dask="parallelized",
+        output_dtypes=[np.float64],
+        dask_gufunc_kwargs={"output_sizes": {"n_face": n_face}},
+        kwargs={
+            "n_face": n_face,
+            "aggregation_func": aggregation_func,
+            "aggregation_func_kwargs": aggregation_func_kwargs,
+        },
+    )
+
+    return uxarray.core.dataarray.UxDataArray(result, uxgrid=uxgrid, name=uxda.name)
 
 
 def _node_to_edge_aggregation(uxda, aggregation, aggregation_func_kwargs):
@@ -146,39 +197,95 @@ def _node_to_edge_aggregation(uxda, aggregation, aggregation_func_kwargs):
             f"{uxda.uxgrid.n_face}."
         )
 
+    aggregation_func = NUMPY_AGGREGATIONS[aggregation]
+
     if isinstance(uxda.data, np.ndarray):
         # apply aggregation using numpy
-        aggregation_var = _apply_node_to_edge_aggregation_numpy(
-            uxda, NUMPY_AGGREGATIONS[aggregation], aggregation_func_kwargs
+        return _apply_node_to_edge_aggregation_numpy(
+            uxda, aggregation_func, aggregation_func_kwargs
         )
     elif isinstance(uxda.data, da.Array):
-        # apply aggregation on dask array, TODO:
-        aggregation_var = _apply_node_to_edge_aggregation_numpy(
-            uxda, NUMPY_AGGREGATIONS[aggregation], aggregation_func_kwargs
+        # apply aggregation lazily on a dask array
+        return _apply_node_to_edge_aggregation_dask(
+            uxda, aggregation_func, aggregation_func_kwargs
         )
     else:
         raise ValueError
 
+
+def _node_to_edge_kernel(
+    data, edge_node_conn, aggregation_func, aggregation_func_kwargs
+):
+    """Node-to-edge topological aggregation on a single numpy block.
+
+    ``data`` has the node dimension as its last axis, shape ``(..., n_node)``;
+    the result has the edge dimension as its last axis, shape ``(..., n_edge)``.
+    Each edge reduces over its two endpoint nodes. Shared verbatim by the numpy
+    path and the (blockwise) dask path.
+    """
+    return aggregation_func(
+        data[..., edge_node_conn], axis=-1, **aggregation_func_kwargs
+    )
+
+
+def _apply_node_to_edge_aggregation_numpy(
+    uxda, aggregation_func, aggregation_func_kwargs
+):
+    """Applies a Node to Edge topological aggregation on a numpy array."""
+    aggregated_var = _node_to_edge_kernel(
+        uxda.values,
+        uxda.uxgrid.edge_node_connectivity.values,
+        aggregation_func,
+        aggregation_func_kwargs,
+    )
     return uxarray.core.dataarray.UxDataArray(
         uxgrid=uxda.uxgrid,
-        data=aggregation_var,
+        data=aggregated_var,
         dims=uxda.dims,
         name=uxda.name,
     ).rename({"n_node": "n_edge"})
 
 
-def _apply_node_to_edge_aggregation_numpy(
+def _apply_node_to_edge_aggregation_dask(
     uxda, aggregation_func, aggregation_func_kwargs
 ):
-    """Applies a Node to Edge topological aggregation on a numpy array."""
-    data = uxda.values
-    edge_node_conn = uxda.uxgrid.edge_node_connectivity.values
-    result = aggregation_func(
-        data[..., edge_node_conn], axis=-1, **aggregation_func_kwargs
+    """Applies a Node to Edge topological aggregation on a Dask array, lazily.
+
+    Mirrors the node-to-face dask path: the data field and the edge-node
+    connectivity enter ``apply_ufunc`` as dask inputs (nothing computed
+    eagerly), ``n_node`` is a single (uncrunked) core dimension, and the leading
+    dimensions are processed blockwise. Each edge reduces over its two endpoint
+    nodes.
+    """
+
+    uxgrid = uxda.uxgrid
+    n_edge = uxgrid.n_edge
+
+    # n_node must live in a single chunk: any node can feed any edge. The grid
+    # metadata is a core-dim input, so its core dims must be single-chunk too.
+    da_in = uxda.chunk({"n_node": -1})
+    edge_node_conn = uxgrid.edge_node_connectivity.chunk({"n_edge": -1, "two": -1})
+
+    # match the numpy path's dtype (float for numeric aggs, bool for all/any)
+    out_dtype = aggregation_func(
+        np.empty((1, edge_node_conn.shape[-1]), dtype=uxda.dtype),
+        axis=-1,
+        **aggregation_func_kwargs,
+    ).dtype
+
+    result = xr.apply_ufunc(
+        _node_to_edge_kernel,
+        da_in,
+        edge_node_conn,
+        input_core_dims=[["n_node"], ["n_edge", "two"]],
+        output_core_dims=[["n_edge"]],
+        dask="parallelized",
+        output_dtypes=[out_dtype],
+        dask_gufunc_kwargs={"output_sizes": {"n_edge": n_edge}},
+        kwargs={
+            "aggregation_func": aggregation_func,
+            "aggregation_func_kwargs": aggregation_func_kwargs,
+        },
     )
-    return result
-
 
-def _apply_node_to_edge_aggregation_dask(*args, **kwargs):
-    """Applies a Node to Edge topological aggregation on a dask array."""
-    pass
+    return uxarray.core.dataarray.UxDataArray(result, uxgrid=uxgrid, name=uxda.name)

uxarray/core/dataarray.py

@@ -246,14 +246,14 @@ def to_geodataframe(
             same name as the ``UxDataArray`` (or named ``var`` if no name exists)
         """
 
-        if self.values.ndim > 1:
+        if self.ndim > 1:
             # data is multidimensional, must be a 1D slice
             raise ValueError(
                 f"Data Variable must be 1-dimensional, with shape {self.uxgrid.n_face} "
                 f"for face-centered data."
             )
 
-        if self.values.size == self.uxgrid.n_face:
+        if self.size == self.uxgrid.n_face:
             gdf, non_nan_polygon_indices = self.uxgrid.to_geodataframe(
                 periodic_elements=periodic_elements,
                 projection=projection,
@@ -290,21 +290,21 @@ def to_geodataframe(
 
             gdf[var_name] = _data
 
-        elif self.values.size == self.uxgrid.n_node:
+        elif self.size == self.uxgrid.n_node:
             raise ValueError(
-                f"Data Variable with size {self.values.size} does not match the number of faces "
+                f"Data Variable with size {self.size} does not match the number of faces "
                 f"({self.uxgrid.n_face}. Current size matches the number of nodes. Consider running "
                 f"``UxDataArray.topological_mean(destination='face') to aggregate the data onto the faces."
             )
-        elif self.values.size == self.uxgrid.n_edge:
+        elif self.size == self.uxgrid.n_edge:
             raise ValueError(
-                f"Data Variable with size {self.values.size} does not match the number of faces "
+                f"Data Variable with size {self.size} does not match the number of faces "
                 f"({self.uxgrid.n_face}. Current size matches the number of edges."
             )
         else:
             # data is not mapped to
             raise ValueError(
-                f"Data Variable with size {self.values.size} does not match the number of faces "
+                f"Data Variable with size {self.size} does not match the number of faces "
                 f"({self.uxgrid.n_face}."
             )
 
@@ -340,7 +340,7 @@ def to_polycollection(
             Flag to indicate whether to override a cached PolyCollection, if it exists
         """
         # data is multidimensional, must be a 1D slice
-        if self.values.ndim > 1:
+        if self.ndim > 1:
             raise ValueError(
                 f"Data Variable must be 1-dimensional, with shape {self.uxgrid.n_face} "
                 f"for face-centered data."
@@ -612,24 +612,29 @@ def integrate(
         >>> uxds = ux.open_dataset("grid.ug", "centroid_pressure_data_ug")
         >>> integral = uxds["psi"].integrate()
         """
-        if self.values.shape[-1] == self.uxgrid.n_face:
-            face_areas = self.uxgrid.face_areas.values
-
-            # perform dot product between face areas and last dimension of data
-            integral = np.einsum("i,...i", face_areas, self.values)
+        if self.shape[-1] == self.uxgrid.n_face:
+            # dot product between face areas and the face dimension of the data
+            if isinstance(self.data, np.ndarray):
+                # eager data: a direct einsum avoids xr.dot's per-call overhead
+                integral = np.einsum(
+                    "i,...i", self.uxgrid.face_areas.values, self.values
+                )
+            else:
+                # dask-backed data: xr.dot keeps the reduction lazy
+                integral = xr.dot(self, self.uxgrid.face_areas, dim="n_face")
 
-        elif self.values.shape[-1] == self.uxgrid.n_node:
+        elif self.shape[-1] == self.uxgrid.n_node:
             raise ValueError("Integrating data mapped to each node not yet supported.")
 
-        elif self.values.shape[-1] == self.uxgrid.n_edge:
+        elif self.shape[-1] == self.uxgrid.n_edge:
             raise ValueError("Integrating data mapped to each edge not yet supported.")
 
         else:
             raise ValueError(
                 f"The final dimension of the data variable does not match the number of nodes, edges, "
                 f"or faces. Expected one of "
                 f"{self.uxgrid.n_node}, {self.uxgrid.n_edge}, or {self.uxgrid.n_face}, "
-                f"but received {self.values.shape[-1]}"
+                f"but received {self.shape[-1]}"
             )
 
         # construct a uxda with integrated quantity
@@ -1653,7 +1658,7 @@ def curl(
         )
 
         # Compute curl = ∂v/∂x - ∂u/∂y
-        curl_values = grad_v_zonal.values - grad_u_meridional.values
+        curl_values = grad_v_zonal.data - grad_u_meridional.data
 
         u_units = self.attrs.get("units", "")
         has_sphere_radius = "sphere_radius" in self.uxgrid._ds.attrs
@@ -2180,11 +2185,10 @@ def get_dual(self):
         # Get correct dimensions for the dual
         dims = [dim_map.get(dim, dim) for dim in self.dims]
 
-        # Get the values from the data array
-        data = np.array(self.values)
-
         # Construct the new data array
-        uxda = uxarray.UxDataArray(uxgrid=dual, data=data, dims=dims, name=self.name)
+        uxda = uxarray.UxDataArray(
+            uxgrid=dual, data=self.data, dims=dims, name=self.name
+        )
 
         return uxda