Initial Work

Author: aaronzedwick

test/test_subset.py

@@ -104,19 +104,40 @@ def test_grid_bounding_box_subset():
                 bbox_antimeridian[0], bbox_antimeridian[1], element=element)
 
 
-
-
 def test_uxda_isel():
     uxds = ux.open_dataset(GRID_PATHS[0], DATA_PATHS[0])
 
     sub = uxds['bottomDepth'].isel(n_face=[1, 2, 3])
 
     assert len(sub) == 3
 
+
 def test_uxda_isel_with_coords():
     uxds = ux.open_dataset(GRID_PATHS[0], DATA_PATHS[0])
     uxds = uxds.assign_coords({"lon_face": uxds.uxgrid.face_lon})
     sub = uxds['bottomDepth'].isel(n_face=[1, 2, 3])
 
     assert "lon_face" in sub.coords
     assert len(sub.coords['lon_face']) == 3
+
+
+def test_inverse_face_indices():
+    grid = ux.open_grid(GRID_PATHS[0])
+
+    # Test nearest neighbor subsetting
+    coord = [0, 0]
+    subset = grid.subset.nearest_neighbor(coord, k=1, element="face centers", inverse_indices=True)
+
+    assert subset.inverse_face_indices is not None
+
+    # Test bounding box subsetting
+    box = [(-10, 10), (-10, 10)]
+    subset = grid.subset.bounding_box(box[0], box[1], element="edge centers", inverse_indices=True)
+
+    assert subset.inverse_face_indices is not None
+
+    # Test bounding circle subsetting
+    center_coord = [0, 0]
+    subset = grid.subset.bounding_circle(center_coord, r=10, element="nodes", inverse_indices=True)
+
+    assert subset.inverse_face_indices is not None

uxarray/grid/grid.py

@@ -1506,6 +1506,18 @@ def global_sphere_coverage(self):
         (i.e. contains no holes)"""
         return not self.partial_sphere_coverage
 
+    @property
+    def inverse_face_indices(self):
+        if self._is_subset:
+            return self._ds["inverse_face_indices"]
+
+    @property
+    def _is_subset(self):
+        """Boolean indicator for whether the Grid is from a subset or not."""
+        if "_is_subset" not in self._ds:
+            self._ds["_is_subset"] = False
+        return self._ds["_is_subset"]
+
     def chunk(self, n_node="auto", n_edge="auto", n_face="auto"):
         """Converts all arrays to dask arrays with given chunks across grid
         dimensions in-place.
@@ -2201,7 +2213,7 @@ def get_dual(self):
 
         return dual
 
-    def isel(self, **dim_kwargs):
+    def isel(self, inverse_indices=False, **dim_kwargs):
         """Indexes an unstructured grid along a given dimension (``n_node``,
         ``n_edge``, or ``n_face``) and returns a new grid.
 
@@ -2226,13 +2238,19 @@ def isel(self, **dim_kwargs):
             raise ValueError("Indexing must be along a single dimension.")
 
         if "n_node" in dim_kwargs:
-            return _slice_node_indices(self, dim_kwargs["n_node"])
+            return _slice_node_indices(
+                self, dim_kwargs["n_node"], inverse_indices=inverse_indices
+            )
 
         elif "n_edge" in dim_kwargs:
-            return _slice_edge_indices(self, dim_kwargs["n_edge"])
+            return _slice_edge_indices(
+                self, dim_kwargs["n_edge"], inverse_indices=inverse_indices
+            )
 
         elif "n_face" in dim_kwargs:
-            return _slice_face_indices(self, dim_kwargs["n_face"])
+            return _slice_face_indices(
+                self, dim_kwargs["n_face"], inverse_indices=inverse_indices
+            )
 
         else:
             raise ValueError(

uxarray/grid/slice.py

@@ -10,7 +10,7 @@
     pass
 
 
-def _slice_node_indices(grid, indices, inclusive=True):
+def _slice_node_indices(grid, indices, inclusive=True, inverse_indices=False):
     """Slices (indexes) an unstructured grid given a list/array of node
     indices, returning a new Grid composed of elements that contain the nodes
     specified in the indices.
@@ -33,10 +33,10 @@ def _slice_node_indices(grid, indices, inclusive=True):
     face_indices = np.unique(grid.node_face_connectivity.values[indices].ravel())
     face_indices = face_indices[face_indices != INT_FILL_VALUE]
 
-    return _slice_face_indices(grid, face_indices)
+    return _slice_face_indices(grid, face_indices, inverse_indices=inverse_indices)
 
 
-def _slice_edge_indices(grid, indices, inclusive=True):
+def _slice_edge_indices(grid, indices, inclusive=True, inverse_indices=False):
     """Slices (indexes) an unstructured grid given a list/array of edge
     indices, returning a new Grid composed of elements that contain the edges
     specified in the indices.
@@ -59,10 +59,10 @@ def _slice_edge_indices(grid, indices, inclusive=True):
     face_indices = np.unique(grid.edge_face_connectivity.values[indices].ravel())
     face_indices = face_indices[face_indices != INT_FILL_VALUE]
 
-    return _slice_face_indices(grid, face_indices)
+    return _slice_face_indices(grid, face_indices, inverse_indices=inverse_indices)
 
 
-def _slice_face_indices(grid, indices, inclusive=True):
+def _slice_face_indices(grid, indices, inclusive=True, inverse_indices=False):
     """Slices (indexes) an unstructured grid given a list/array of face
     indices, returning a new Grid composed of elements that contain the faces
     specified in the indices.
@@ -77,7 +77,6 @@ def _slice_face_indices(grid, indices, inclusive=True):
         Whether to perform inclusive (i.e. elements must contain at least one desired feature from a slice) as opposed
         to exclusive (i.e elements be made up all desired features from a slice)
     """
-
     if inclusive is False:
         raise ValueError("Exclusive slicing is not yet supported.")
 
@@ -132,4 +131,9 @@ def _slice_face_indices(grid, indices, inclusive=True):
             # drop any conn that would require re-computation
             ds = ds.drop_vars(conn_name)
 
+    if inverse_indices:
+        ds["inverse_face_indices"] = indices
+
+    ds["_is_subset"] = True
+
     return Grid.from_dataset(ds, source_grid_spec=grid.source_grid_spec)

uxarray/subset/grid_accessor.py

@@ -33,6 +33,7 @@ def bounding_box(
         lat_bounds: Union[Tuple, List, np.ndarray],
         element: Optional[str] = "nodes",
         method: Optional[str] = "coords",
+        inverse_indices=False,
         **kwargs,
     ):
         """Subsets an unstructured grid between two latitude and longitude
@@ -53,6 +54,8 @@ def bounding_box(
             face centers, or edge centers lie within the bounds.
         element: str
             Element for use with `coords` comparison, one of `nodes`, `face centers`, or `edge centers`
+        inverse_indices : bool
+            Flag to indicate whether to store the original grids face indices for later use
         """
 
         if method == "coords":
@@ -101,11 +104,11 @@ def bounding_box(
                 )
 
             if element == "nodes":
-                return self.uxgrid.isel(n_node=indices)
+                return self.uxgrid.isel(inverse_indices, n_node=indices)
             elif element == "face centers":
-                return self.uxgrid.isel(n_face=indices)
+                return self.uxgrid.isel(inverse_indices, n_face=indices)
             elif element == "edge centers":
-                return self.uxgrid.isel(n_edge=indices)
+                return self.uxgrid.isel(inverse_indices, n_edge=indices)
 
         else:
             raise ValueError(f"Method '{method}' not supported.")
@@ -115,6 +118,7 @@ def bounding_circle(
         center_coord: Union[Tuple, List, np.ndarray],
         r: Union[float, int],
         element: Optional[str] = "nodes",
+        inverse_indices=False,
         **kwargs,
     ):
         """Subsets an unstructured grid by returning all elements within some
@@ -128,6 +132,8 @@ def bounding_circle(
             Radius of bounding circle (in degrees)
         element: str
             Element for use with `coords` comparison, one of `nodes`, `face centers`, or `edge centers`
+        inverse_indices : bool
+            Flag to indicate whether to store the original grids face indices for later use
         """
 
         coords = np.asarray(center_coord)
@@ -141,13 +147,14 @@ def bounding_circle(
                 f"No elements founding within the bounding circle with radius {r} when querying {element}"
             )
 
-        return self._index_grid(ind, element)
+        return self._index_grid(ind, element, inverse_indices)
 
     def nearest_neighbor(
         self,
         center_coord: Union[Tuple, List, np.ndarray],
         k: int,
         element: Optional[str] = "nodes",
+        inverse_indices=False,
         **kwargs,
     ):
         """Subsets an unstructured grid by returning the ``k`` closest
@@ -161,6 +168,8 @@ def nearest_neighbor(
             Number of neighbors to query
         element: str
             Element for use with `coords` comparison, one of `nodes`, `face centers`, or `edge centers`
+        inverse_indices : bool
+            Flag to indicate whether to store the original grids face indices for later use
         """
 
         coords = np.asarray(center_coord)
@@ -169,7 +178,7 @@ def nearest_neighbor(
 
         _, ind = tree.query(coords, k)
 
-        return self._index_grid(ind, element)
+        return self._index_grid(ind, element, inverse_indices=inverse_indices)
 
     def _get_tree(self, coords, tree_type):
         """Internal helper for obtaining the desired KDTree or BallTree."""
@@ -187,12 +196,12 @@ def _get_tree(self, coords, tree_type):
 
         return tree
 
-    def _index_grid(self, ind, tree_type):
+    def _index_grid(self, ind, tree_type, inverse_indices=False):
         """Internal helper for indexing a grid with indices based off the
         provided tree type."""
         if tree_type == "nodes":
-            return self.uxgrid.isel(n_node=ind)
+            return self.uxgrid.isel(inverse_indices, n_node=ind)
         elif tree_type == "edge centers":
-            return self.uxgrid.isel(n_edge=ind)
+            return self.uxgrid.isel(inverse_indices, n_edge=ind)
         else:
-            return self.uxgrid.isel(n_face=ind)
+            return self.uxgrid.isel(inverse_indices, n_face=ind)