Update slicing to only update edges when possible (#1256)

* created docstring for azimuthal_mean

* wrote azimuthal mean computation in public function

* draft azimuthal mean ready to test

* fixed typos, return hit count for radial bins

* made azimuthal mean more robust to axis ordering

* added central coord as attribute in output UxDataArray of azimuthal_mean

* run pre-commit

* update parameters and set default values to nan

* fix edge slicing issue

* remove code from other brahnch

* fix slicing on edge arrays

* use the correct dict

---------

Co-authored-by: Joshua Pan <[email protected]>

Author: philipc2

test/test_cross_sections.py

@@ -198,3 +198,25 @@ def test_latitude_along_arc(self):
         out2 = uxgrid.get_faces_at_constant_latitude(lat=25.41)
 
         nt.assert_array_equal(out1, out2)
+
+
+
+def test_double_cross_section():
+    uxgrid = ux.open_grid(quad_hex_grid_path)
+
+    # construct edges
+    sub_lat = uxgrid.cross_section.constant_latitude(0.0)
+
+    sub_lat_lon = sub_lat.cross_section.constant_longitude(0.0)
+
+    assert "n_edge" not in sub_lat_lon._ds.dims
+
+    _ = uxgrid.face_edge_connectivity
+    _ = uxgrid.edge_node_connectivity
+    _ = uxgrid.edge_lon
+
+    sub_lat = uxgrid.cross_section.constant_latitude(0.0)
+
+    sub_lat_lon = sub_lat.cross_section.constant_longitude(0.0)
+
+    assert "n_edge" in sub_lat_lon._ds.dims

uxarray/grid/slice.py

@@ -94,64 +94,85 @@ def _slice_face_indices(
         Indicates whether to store the original grids indices. Passing `True` stores the original face centers,
         other reverse indices can be stored by passing any or all of the following: (["face", "edge", "node"], True)
     """
+    from uxarray.grid import Grid
+
     if inclusive is False:
         raise ValueError("Exclusive slicing is not yet supported.")
 
-    from uxarray.grid import Grid
-
     ds = grid._ds
-
-    indices = np.asarray(indices, dtype=INT_DTYPE)
-
-    if indices.ndim == 0:
-        indices = np.expand_dims(indices, axis=0)
-
-    face_indices = indices
+    face_indices = np.atleast_1d(np.asarray(indices, dtype=INT_DTYPE))
 
     # nodes of each face (inclusive)
     node_indices = np.unique(grid.face_node_connectivity.values[face_indices].ravel())
     node_indices = node_indices[node_indices != INT_FILL_VALUE]
 
-    # index original dataset to obtain a 'subgrid'
+    # Index Node and Face variables
     ds = ds.isel(n_node=node_indices)
     ds = ds.isel(n_face=face_indices)
 
-    # Only slice edge dimension if we already have the connectivity
-    if "face_edge_connectivity" in grid._ds:
+    # Only slice edge dimension if we have the face edge connectivity
+    if "face_edge_connectivity" in ds:
         edge_indices = np.unique(
             grid.face_edge_connectivity.values[face_indices].ravel()
         )
         edge_indices = edge_indices[edge_indices != INT_FILL_VALUE]
         ds = ds.isel(n_edge=edge_indices)
         ds["subgrid_edge_indices"] = xr.DataArray(edge_indices, dims=["n_edge"])
+    # Otherwise, drop any edge variables
     else:
+        if "n_edge" in ds.dims:
+            ds = ds.drop_dims(["n_edge"])
         edge_indices = None
 
     ds["subgrid_node_indices"] = xr.DataArray(node_indices, dims=["n_node"])
     ds["subgrid_face_indices"] = xr.DataArray(face_indices, dims=["n_face"])
 
-    # mapping to update existing connectivity
-    node_indices_dict = {
-        key: val for key, val in zip(node_indices, np.arange(0, len(node_indices)))
-    }
+    # Construct updated Node Index Map
+    node_indices_dict = {orig: new for new, orig in enumerate(node_indices)}
     node_indices_dict[INT_FILL_VALUE] = INT_FILL_VALUE
 
-    for conn_name in grid._ds.data_vars:
-        # update or drop connectivity variables to correctly point to the new index of each element
+    # Construct updated Edge Index Map
+    if edge_indices is not None:
+        edge_indices_dict = {orig: new for new, orig in enumerate(edge_indices)}
+        edge_indices_dict[INT_FILL_VALUE] = INT_FILL_VALUE
+    else:
+        edge_indices_dict = None
+
+    def map_node_indices(i):
+        return node_indices_dict.get(i, INT_FILL_VALUE)
+
+    if edge_indices is not None:
+
+        def map_edge_indices(i):
+            return edge_indices_dict.get(i, INT_FILL_VALUE)
+    else:
+        map_edge_indices = None
+
+    for conn_name in list(ds.data_vars):
+        if conn_name.endswith("_node_connectivity"):
+            map_fn = map_node_indices
 
-        if "_node_connectivity" in conn_name:
-            # update connectivity vars that index into nodes
-            ds[conn_name] = xr.DataArray(
-                np.vectorize(node_indices_dict.__getitem__, otypes=[INT_DTYPE])(
-                    ds[conn_name].values
-                ),
-                dims=ds[conn_name].dims,
-                attrs=ds[conn_name].attrs,
-            )
+        elif conn_name.endswith("_edge_connectivity"):
+            if edge_indices_dict is None:
+                ds = ds.drop_vars(conn_name)
+                continue
+            map_fn = map_edge_indices
 
         elif "_connectivity" in conn_name:
-            # drop any conn that would require re-computation
+            # anything else we can't remap
             ds = ds.drop_vars(conn_name)
+            continue
+
+        else:
+            # not a connectivity var, skip
+            continue
+
+        # Apply Remapping
+        ds[conn_name] = xr.DataArray(
+            np.vectorize(map_fn, otypes=[INT_DTYPE])(ds[conn_name].values),
+            dims=ds[conn_name].dims,
+            attrs=ds[conn_name].attrs,
+        )
 
     if inverse_indices:
         inverse_indices_ds = xr.Dataset()