DRAFT: Reductions

test/test_topological_agg.py

@@ -7,7 +7,6 @@
 
 current_path = Path(os.path.dirname(os.path.realpath(__file__)))
 
-
 ds_path = current_path / 'meshfiles' / "mpas" / "QU" / 'oQU480.231010.nc'
 
 AGGS = ["topological_mean",
@@ -32,11 +31,22 @@ def test_node_to_face_aggs():
         assert 'n_face' in grid_reduction.dims
 
 
-
 def test_node_to_edge_aggs():
     uxds = ux.open_dataset(ds_path, ds_path)
 
     for agg_func in AGGS:
         grid_reduction = getattr(uxds['areaTriangle'], agg_func)(destination='edge')
 
         assert 'n_edge' in grid_reduction.dims
+
+
+def test_edge_to_face_aggs():
+    grid_path = '/Users/aaronzedwick/uxarray/test/meshfiles/mpas/QU/mesh.QU.1920km.151026.nc'
+
+    uxds = ux.open_dataset(grid_path, grid_path)
+
+    uxds = uxds['cellsOnEdge'].subset.nearest_neighbor(k=3, center_coord=[0, 0])
+
+    uxda_edge_face_agg = uxds.topological_mean(destination="node")
+
+    print(uxda_edge_face_agg)

uxarray/conventions/ugrid.py

@@ -224,6 +224,33 @@
 
 N_NODES_PER_FACE_DIMS = ["n_face"]
 
+N_FACES_PER_NODE_ATTRS = {
+    "cf_role": "n_faces_per_node",
+    "long name": "Number of faces per node",
+}
+
+N_FACES_PER_NODE_DIMS = ["n_node"]
+
+N_FACES_PER_EDGE_ATTRS = {
+    "cf_role": "n_faces_per_edge",
+    "long name": "Number of faces per edge",
+}
+
+N_FACES_PER_EDGE_DIMS = ["n_edge"]
+
+N_EDGES_PER_FACE_ATTRS = {
+    "cf_role": "n_edges_per_face",
+    "long name": "Number of edges per face",
+}
+
+N_EDGES_PER_FACE_DIMS = ["n_face"]
+
+N_EDGES_PER_NODE_ATTRS = {
+    "cf_role": "n_edges_per_node",
+    "long name": "Number of edges per node",
+}
+
+N_EDGES_PER_NODE_DIMS = ["n_node"]
 
 CONNECTIVITY_NAMES = [
     "face_node_connectivity",

uxarray/core/aggregation.py

@@ -37,32 +37,32 @@ def _uxda_grid_aggregate(uxda, destination, aggregation, **kwargs):
         else:
             raise ValueError(
                 f"Invalid destination for a node-centered data variable. Expected"
-                f"one of ['face', 'edge' but received {destination}"
+                f"one of ['face', 'edge'] but received {destination}"
             )
 
     elif uxda._edge_centered():
         # aggregation of an edge-centered data variable
-        raise NotImplementedError(
-            "Aggregation of edge-centered data variables is not yet supported."
-        )
-        # if destination == "node":
-        #     pass
-        # elif destination == "face":
-        #     pass
-        # else:
-        #     raise ValueError("TODO: )
+        if destination == "face":
+            return _edge_to_face_aggregation(uxda, aggregation, kwargs)
+        elif destination == "node":
+            return _edge_to_node_aggregation(uxda, aggregation, kwargs)
+        else:
+            raise ValueError(
+                f"Invalid destination for an edge-centered data variable. Expected"
+                f"one of ['face', 'node'] but received {destination}"
+            )
 
     elif uxda._face_centered():
         # aggregation of a face-centered data variable
-        raise NotImplementedError(
-            "Aggregation of face-centered data variables is not yet supported."
-        )
-        # if destination == "node":
-        #     pass
-        # elif destination == "edge":
-        #     pass
-        # else:
-        #     raise ValueError("TODO: ")
+        if destination == "node":
+            return _face_to_node_aggregation(uxda, aggregation, kwargs)
+        elif destination == "edge":
+            return _face_to_edge_aggregation(uxda, aggregation, kwargs)
+        else:
+            raise ValueError(
+                f"Invalid destination for a face-centered data variable. Expected"
+                f"one of ['node', 'edge'] but received {destination}"
+            )
 
     else:
         raise ValueError(
@@ -181,3 +181,247 @@ def _apply_node_to_edge_aggregation_numpy(
 def _apply_node_to_edge_aggregation_dask(*args, **kwargs):
     """Applies a Node to Edge topological aggregation on a dask array."""
     pass
+
+
+def _apply_edge_to_face_aggregation_numpy(
+    uxda, aggregation_func, aggregation_func_kwargs
+):
+    """Applies an Edge to Face Topological aggregation on a Numpy array."""
+    data = uxda.values
+    face_edge_conn = uxda.uxgrid.face_edge_connectivity.values
+    n_nodes_per_face = uxda.uxgrid.n_nodes_per_face.values
+
+    (
+        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,))
+
+    for e, start, end in zip(element_sizes, change_ind[:-1], change_ind[1:]):
+        face_inds = n_nodes_per_face_sorted_ind[start:end]
+        face_nodes_par = face_edge_conn[face_inds, 0:e]
+
+        # apply aggregation function to current face edge partition
+        aggregation_par = aggregation_func(
+            data[..., face_nodes_par], axis=-1, **aggregation_func_kwargs
+        )
+
+        # store current aggregation
+        result[..., face_inds] = aggregation_par
+
+    return result
+
+
+def _edge_to_face_aggregation(uxda, aggregation, aggregation_func_kwargs):
+    """Applies an Edge to Face Topological aggregation."""
+    if not uxda._edge_centered():
+        raise ValueError(
+            f"Data Variable must be mapped to the edge centers of each face, with dimension "
+            f"{uxda.uxgrid.n_edge}."
+        )
+
+    if isinstance(uxda.data, np.ndarray):
+        # apply aggregation using numpy
+        aggregated_var = _apply_edge_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_edge_to_face_aggregation_numpy(
+            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_edge": "n_face"})
+
+
+def _apply_edge_to_node_aggregation_numpy(
+    uxda, aggregation_func, aggregation_func_kwargs
+):
+    """Applies an Edge to Node Topological aggregation on a Numpy array."""
+    data = uxda.values
+    node_edge_conn = uxda.uxgrid.node_edge_connectivity.values
+    n_edges_per_node = uxda.uxgrid.n_edges_per_node.values
+
+    (
+        change_ind,
+        n_edges_per_node_sorted_ind,
+        element_sizes,
+        size_counts,
+    ) = get_face_node_partitions(n_edges_per_node)
+
+    result = np.empty(shape=(data.shape[:-1]) + (uxda.uxgrid.n_node,))
+
+    for e, start, end in zip(element_sizes, change_ind[:-1], change_ind[1:]):
+        node_inds = n_edges_per_node_sorted_ind[start:end]
+        node_edges_par = node_edge_conn[node_inds, 0:e]
+
+        # apply aggregation function to current node edge partition
+        aggregation_par = aggregation_func(
+            data[..., node_edges_par], axis=-1, **aggregation_func_kwargs
+        )
+
+        # store current aggregation
+        result[..., node_inds] = aggregation_par
+
+    return result
+
+
+def _edge_to_node_aggregation(uxda, aggregation, aggregation_func_kwargs):
+    """Applies an Edge to Node Topological aggregation."""
+    if not uxda._edge_centered():
+        raise ValueError(
+            f"Data Variable must be mapped to the edge centers of each face, with dimension "
+            f"{uxda.uxgrid.n_edge}."
+        )
+
+    if isinstance(uxda.data, np.ndarray):
+        # apply aggregation using numpy
+        aggregated_var = _apply_edge_to_node_aggregation_numpy(
+            uxda, NUMPY_AGGREGATIONS[aggregation], aggregation_func_kwargs
+        )
+    elif isinstance(uxda.data, da.Array):
+        # apply aggregation on dask array, TODO:
+        aggregated_var = _apply_edge_to_node_aggregation_numpy(
+            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_edge": "n_node"})
+
+
+def _apply_face_to_node_aggregation_numpy(
+    uxda, aggregation_func, aggregation_func_kwargs
+):
+    """Applies a Face to Node Topological aggregation on a Numpy array."""
+    data = uxda.values
+    node_face_conn = uxda.uxgrid.node_face_connectivity.values
+    n_faces_per_node = uxda.uxgrid.n_faces_per_node.values
+
+    (
+        change_ind,
+        n_faces_per_node_sorted_ind,
+        element_sizes,
+        size_counts,
+    ) = get_face_node_partitions(n_faces_per_node)
+
+    result = np.empty(shape=(data.shape[:-1]) + (uxda.uxgrid.n_node,))
+
+    for e, start, end in zip(element_sizes, change_ind[:-1], change_ind[1:]):
+        face_inds = n_faces_per_node_sorted_ind[start:end]
+        face_nodes_par = node_face_conn[face_inds, 0:e]
+
+        # apply aggregation function to current node face partition
+        aggregation_par = aggregation_func(
+            data[..., face_nodes_par], axis=-1, **aggregation_func_kwargs
+        )
+
+        # store current aggregation
+        result[..., face_inds] = aggregation_par
+
+    return result
+
+
+def _face_to_node_aggregation(uxda, aggregation, aggregation_func_kwargs):
+    """Applies a Face to Node Topological aggregation."""
+    if not uxda._face_centered():
+        raise ValueError(
+            f"Data Variable must be mapped to the face centers of each face, with dimension "
+            f"{uxda.uxgrid.n_face}."
+        )
+
+    if isinstance(uxda.data, np.ndarray):
+        # apply aggregation using numpy
+        aggregated_var = _apply_face_to_node_aggregation_numpy(
+            uxda, NUMPY_AGGREGATIONS[aggregation], aggregation_func_kwargs
+        )
+    elif isinstance(uxda.data, da.Array):
+        # apply aggregation on dask array, TODO:
+        aggregated_var = _apply_face_to_node_aggregation_numpy(
+            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_face": "n_node"})
+
+
+def _apply_face_to_edge_aggregation_numpy(
+    uxda, aggregation_func, aggregation_func_kwargs
+):
+    """Applies a Face to Edge Topological aggregation on a Numpy array."""
+    data = uxda.values
+    edge_face_conn = uxda.uxgrid.edge_face_connectivity.values
+    n_faces_per_edge = uxda.uxgrid.n_faces_per_edge.values
+
+    (
+        change_ind,
+        n_faces_per_edge_sorted_ind,
+        element_sizes,
+        size_counts,
+    ) = get_face_node_partitions(n_faces_per_edge)
+
+    result = np.empty(shape=(data.shape[:-1]) + (uxda.uxgrid.n_edge,))
+
+    for e, start, end in zip(element_sizes, change_ind[:-1], change_ind[1:]):
+        face_inds = n_faces_per_edge_sorted_ind[start:end]
+        face_edge_par = edge_face_conn[face_inds, 0:e]
+
+        # apply aggregation function to current node face partition
+        aggregation_par = aggregation_func(
+            data[..., face_edge_par], axis=-1, **aggregation_func_kwargs
+        )
+
+        # store current aggregation
+        result[..., face_inds] = aggregation_par
+
+    return result
+
+
+def _face_to_edge_aggregation(uxda, aggregation, aggregation_func_kwargs):
+    """Applies a Face to Edge Topological aggregation."""
+    if not uxda._face_centered():
+        raise ValueError(
+            f"Data Variable must be mapped to the face centers of each face, with dimension "
+            f"{uxda.uxgrid.n_face}."
+        )
+
+    if isinstance(uxda.data, np.ndarray):
+        # apply aggregation using numpy
+        aggregated_var = _apply_face_to_edge_aggregation_numpy(
+            uxda, NUMPY_AGGREGATIONS[aggregation], aggregation_func_kwargs
+        )
+    elif isinstance(uxda.data, da.Array):
+        # apply aggregation on dask array, TODO:
+        aggregated_var = _apply_face_to_edge_aggregation_numpy(
+            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_face": "n_edge"})