Add Inverse Face Indices to Subsetted Grids (#1122)

* Initial Work

* Addressed review comments

* Updated doc string

* Added inverse_indices support for data arrays and cross sections

* Added ability to choose which inverse_indices to store, stores as ds

* updated grid, doc strings, api, test cases

* fixed leftover variables

* Fixed failing tests

* Update grid.py

* New naming convention, fixed spelling errors

* Updated subsetting notebook

* Fixed precommit

* Added is_subset property

* Update uxarray/grid/grid.py

Co-authored-by: Philip Chmielowiec <[email protected]>

* Added doc string, updated test case

* Update grid.py

---------

Co-authored-by: Philip Chmielowiec <[email protected]>

Author: aaronzedwick

docs/api.rst

@@ -54,6 +54,7 @@ Indexing
    :toctree: generated/
 
    Grid.isel
+   Grid.inverse_indices
 
 Dimensions
 ~~~~~~~~~~

docs/user-guide/subset.ipynb

@@ -553,6 +553,95 @@
     "print(\"Bounding Box Mean:      \", bbox_subset_nodes.values.mean())\n",
     "print(\"Bounding Circle Mean:   \", bcircle_subset.values.mean())"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Retrieving Orignal Grid Indices"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Sometimes having the original grids' indices is useful. These indices can be stored within the subset with the `inverse_indices` variable. This can be used to store the indices of the original face centers, edge centers, and node indices. This variable can be used within the subset as follows:\n",
+    "\n",
+    "* Passing in `True`, which will store the face center indices\n",
+    "* Passing in a list of which indices to store, along with `True`, to indicate what kind of original grid indices to store.\n",
+    "    * Options for which indices to select include: `face`, `node`, and `edge`\n",
+    "\n",
+    "This currently only works when the element is `face centers`. Elements `nodes` and `edge centers` will be supported in the future."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "subset_indices = uxds[\"relhum_200hPa\"][0].subset.bounding_circle(\n",
+    "    center_coord,\n",
+    "    r,\n",
+    "    element=\"face centers\",\n",
+    "    inverse_indices=([\"face\", \"node\", \"edge\"], True),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "These indices can be retrieve through the grid:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "subset_indices.uxgrid.inverse_indices"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Determining if a Grid is a Subset"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "To check if a Grid (or dataset using `.uxgrid`) is a subset, we can use `Grid.is_subset`, which will return either `True` or `False`, depending on whether the `Grid` is a subset. Since `subset_indices` is a subset, using this feature we will return `True`:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "subset_indices.uxgrid.is_subset"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The file we have been using to create these subsets, `uxds`, is not a subset, so using the same call we will return `False:`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "uxds.uxgrid.is_subset"
+   ]
   }
  ],
  "metadata": {
@@ -571,7 +660,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.9"
+   "version": "3.12.7"
   }
  },
  "nbformat": 4,

test/test_subset.py

@@ -104,19 +104,49 @@ 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_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_indices is not None
+
+    # Test bounding box subsetting
+    box = [(-10, 10), (-10, 10)]
+    subset = grid.subset.bounding_box(box[0], box[1], element="face centers", inverse_indices=True)
+
+    assert subset.inverse_indices is not None
+
+    # Test bounding circle subsetting
+    center_coord = [0, 0]
+    subset = grid.subset.bounding_circle(center_coord, r=10, element="face centers", inverse_indices=True)
+
+    assert subset.inverse_indices is not None
+
+    # Ensure code raises exceptions when the element is edges or nodes or inverse_indices is incorrect
+    assert pytest.raises(Exception, grid.subset.bounding_circle, center_coord, r=10, element="edge centers", inverse_indices=True)
+    assert pytest.raises(Exception, grid.subset.bounding_circle, center_coord, r=10, element="nodes", inverse_indices=True)
+    assert pytest.raises(ValueError, grid.subset.bounding_circle, center_coord, r=10, element="face center", inverse_indices=(['not right'], True))
+
+    # Test isel directly
+    subset = grid.isel(n_face=[1], inverse_indices=True)
+    assert subset.inverse_indices.face.values == 1

uxarray/core/dataarray.py

@@ -1035,7 +1035,7 @@ def _edge_centered(self) -> bool:
         "n_edge" dimension)"""
         return "n_edge" in self.dims
 
-    def isel(self, ignore_grid=False, *args, **kwargs):
+    def isel(self, ignore_grid=False, inverse_indices=False, *args, **kwargs):
         """Grid-informed implementation of xarray's ``isel`` method, which
         enables indexing across grid dimensions.
 
@@ -1069,11 +1069,17 @@ def isel(self, ignore_grid=False, *args, **kwargs):
                 raise ValueError("Only one grid dimension can be sliced at a time")
 
             if "n_node" in kwargs:
-                sliced_grid = self.uxgrid.isel(n_node=kwargs["n_node"])
+                sliced_grid = self.uxgrid.isel(
+                    n_node=kwargs["n_node"], inverse_indices=inverse_indices
+                )
             elif "n_edge" in kwargs:
-                sliced_grid = self.uxgrid.isel(n_edge=kwargs["n_edge"])
+                sliced_grid = self.uxgrid.isel(
+                    n_edge=kwargs["n_edge"], inverse_indices=inverse_indices
+                )
             else:
-                sliced_grid = self.uxgrid.isel(n_face=kwargs["n_face"])
+                sliced_grid = self.uxgrid.isel(
+                    n_face=kwargs["n_face"], inverse_indices=inverse_indices
+                )
 
             return self._slice_from_grid(sliced_grid)
 

uxarray/cross_sections/dataarray_accessor.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Union, List, Set
 
 if TYPE_CHECKING:
     pass
@@ -22,7 +22,9 @@ def __repr__(self):
 
         return prefix + methods_heading
 
-    def constant_latitude(self, lat: float):
+    def constant_latitude(
+        self, lat: float, inverse_indices: Union[List[str], Set[str], bool] = False
+    ):
         """Extracts a cross-section of the data array by selecting all faces that
         intersect with a specified line of constant latitude.
 
@@ -31,6 +33,9 @@ def constant_latitude(self, lat: float):
         lat : float
             The latitude at which to extract the cross-section, in degrees.
             Must be between -90.0 and 90.0
+        inverse_indices : Union[List[str], Set[str], bool], optional
+            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)
 
         Returns
         -------
@@ -60,9 +65,11 @@ def constant_latitude(self, lat: float):
 
         faces = self.uxda.uxgrid.get_faces_at_constant_latitude(lat)
 
-        return self.uxda.isel(n_face=faces)
+        return self.uxda.isel(n_face=faces, inverse_indices=inverse_indices)
 
-    def constant_longitude(self, lon: float):
+    def constant_longitude(
+        self, lon: float, inverse_indices: Union[List[str], Set[str], bool] = False
+    ):
         """Extracts a cross-section of the data array by selecting all faces that
         intersect with a specified line of constant longitude.
 
@@ -71,6 +78,9 @@ def constant_longitude(self, lon: float):
         lon : float
             The latitude at which to extract the cross-section, in degrees.
             Must be between -180.0 and 180.0
+        inverse_indices : Union[List[str], Set[str], bool], optional
+            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)
 
         Returns
         -------
@@ -102,7 +112,7 @@ def constant_longitude(self, lon: float):
             lon,
         )
 
-        return self.uxda.isel(n_face=faces)
+        return self.uxda.isel(n_face=faces, inverse_indices=inverse_indices)
 
     def gca(self, *args, **kwargs):
         raise NotImplementedError

uxarray/cross_sections/grid_accessor.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Union, List, Set
 
 if TYPE_CHECKING:
     from uxarray.grid import Grid
@@ -25,6 +25,7 @@ def constant_latitude(
         self,
         lat: float,
         return_face_indices: bool = False,
+        inverse_indices: Union[List[str], Set[str], bool] = False,
     ):
         """Extracts a cross-section of the grid by selecting all faces that
         intersect with a specified line of constant latitude.
@@ -36,6 +37,9 @@ def constant_latitude(
             Must be between -90.0 and 90.0
         return_face_indices : bool, optional
             If True, also returns the indices of the faces that intersect with the line of constant latitude.
+        inverse_indices : Union[List[str], Set[str], bool], optional
+            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)
 
         Returns
         -------
@@ -66,7 +70,9 @@ def constant_latitude(
         if len(faces) == 0:
             raise ValueError(f"No intersections found at lat={lat}.")
 
-        grid_at_constant_lat = self.uxgrid.isel(n_face=faces)
+        grid_at_constant_lat = self.uxgrid.isel(
+            n_face=faces, inverse_indices=inverse_indices
+        )
 
         if return_face_indices:
             return grid_at_constant_lat, faces
@@ -77,6 +83,7 @@ def constant_longitude(
         self,
         lon: float,
         return_face_indices: bool = False,
+        inverse_indices: Union[List[str], Set[str], bool] = False,
     ):
         """Extracts a cross-section of the grid by selecting all faces that
         intersect with a specified line of constant longitude.
@@ -88,6 +95,9 @@ def constant_longitude(
             Must be between -90.0 and 90.0
         return_face_indices : bool, optional
             If True, also returns the indices of the faces that intersect with the line of constant longitude.
+        inverse_indices : Union[List[str], Set[str], bool], optional
+            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)
 
         Returns
         -------
@@ -117,7 +127,9 @@ def constant_longitude(
         if len(faces) == 0:
             raise ValueError(f"No intersections found at lon={lon}")
 
-        grid_at_constant_lon = self.uxgrid.isel(n_face=faces)
+        grid_at_constant_lon = self.uxgrid.isel(
+            n_face=faces, inverse_indices=inverse_indices
+        )
 
         if return_face_indices:
             return grid_at_constant_lon, faces