Preserve Coordinates when Indexing UxDataArray (#1003)

* preserve coordinates in isel * consider dropping coordinates when remapping * add preserve coordinates util function * simpler handling of coordinates * update test * use del * fix isel failure --------- Co-authored-by: Orhan Eroglu <[email protected]>
UXARRAY · Oct 29, 2024 · e30586f · e30586f
1 parent 6061d2e
commit e30586f
Show file tree

Hide file tree

Showing 6 changed files with 70 additions and 108 deletions.
diff --git a/docs/user-guide/subset.ipynb b/docs/user-guide/subset.ipynb
@@ -16,10 +16,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:36:46.872906Z",
-     "start_time": "2024-06-07T22:36:34.761168Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
@@ -72,10 +68,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:36:46.917137Z",
-     "start_time": "2024-06-07T22:36:46.873648Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
@@ -95,10 +87,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:36:53.691169Z",
-     "start_time": "2024-06-07T22:36:46.918225Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
@@ -113,10 +101,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:36:53.701638Z",
-     "start_time": "2024-06-07T22:36:53.692509Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
@@ -131,10 +115,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:36:53.707801Z",
-     "start_time": "2024-06-07T22:36:53.706375Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
@@ -165,19 +145,15 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:36:59.545160Z",
-     "start_time": "2024-06-07T22:36:53.709243Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
     }
    },
    "outputs": [],
    "source": [
-    "uxds[\"relhum_200hPa\"][0].plot.rasterize(\n",
-    "    method=\"polygon\", exclude_antimeridian=True, title=\"Global Grid\", **plot_opts\n",
+    "uxds[\"relhum_200hPa\"][0].plot(\n",
+    "    rasterize=True, periodic_elements=\"exclude\", title=\"Global Grid\", **plot_opts\n",
     ") * features"
    ]
   },
@@ -227,10 +203,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:38:21.645685Z",
-     "start_time": "2024-06-07T22:38:21.641552Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
@@ -245,10 +217,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:38:23.218966Z",
-     "start_time": "2024-06-07T22:38:23.172443Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
@@ -277,10 +245,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:38:37.739473Z",
-     "start_time": "2024-06-07T22:38:37.697998Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
@@ -310,10 +274,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:39:14.680397Z",
-     "start_time": "2024-06-07T22:39:12.596281Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
@@ -325,9 +285,9 @@
     "    lon_bounds, lat_bounds, element=\"nodes\"\n",
     ")\n",
     "\n",
-    "bbox_subset_nodes.plot.rasterize(\n",
-    "    method=\"polygon\",\n",
-    "    exclude_antimeridian=True,\n",
+    "bbox_subset_nodes.plot(\n",
+    "    rasterize=True,\n",
+    "    periodic_elements=\"exclude\",\n",
     "    clim=clim,\n",
     "    title=\"Bounding Box Subset (Corner Node Query)\",\n",
     "    **plot_opts,\n",
@@ -350,10 +310,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:38:57.658080Z",
-     "start_time": "2024-06-07T22:38:55.725793Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
@@ -365,9 +321,9 @@
     "    lon_bounds, lat_bounds, element=\"face centers\"\n",
     ")\n",
     "\n",
-    "bbox_subset_faces.plot.rasterize(\n",
-    "    method=\"polygon\",\n",
-    "    exclude_antimeridian=True,\n",
+    "bbox_subset_faces.plot(\n",
+    "    rasterize=True,\n",
+    "    periodic_elements=\"exclude\",\n",
     "    clim=clim,\n",
     "    title=\"Bounding Box Subset (Face Center Query)\",\n",
     "    **plot_opts,\n",
@@ -408,10 +364,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:39:43.139200Z",
-     "start_time": "2024-06-07T22:39:43.026055Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
@@ -428,10 +380,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:39:46.520147Z",
-     "start_time": "2024-06-07T22:39:44.107409Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
@@ -441,9 +389,9 @@
    "source": [
     "bcircle_subset = uxds[\"relhum_200hPa\"][0].subset.bounding_circle(center_coord, r)\n",
     "\n",
-    "bcircle_subset.plot.rasterize(\n",
-    "    method=\"polygon\",\n",
-    "    exclude_antimeridian=True,\n",
+    "bcircle_subset.plot(\n",
+    "    rasterize=True,\n",
+    "    periodic_elements=\"exclude\",\n",
     "    clim=clim,\n",
     "    title=\"Bounding Circle Subset\",\n",
     "    **plot_opts,\n",
@@ -468,10 +416,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:39:53.170082Z",
-     "start_time": "2024-06-07T22:39:53.017349Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
@@ -486,10 +430,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:39:56.343421Z",
-     "start_time": "2024-06-07T22:39:54.415082Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
@@ -501,9 +441,9 @@
     "    center_coord, k=30, element=\"nodes\"\n",
     ")\n",
     "\n",
-    "nn_subset.plot.rasterize(\n",
-    "    method=\"polygon\",\n",
-    "    exclude_antimeridian=True,\n",
+    "nn_subset.plot(\n",
+    "    rasterize=True,\n",
+    "    periodic_elements=\"exclude\",\n",
     "    clim=clim,\n",
     "    title=\"Nearest Neighbor Subset\",\n",
     "    **plot_opts,\n",
@@ -526,10 +466,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:40:46.514866Z",
-     "start_time": "2024-06-07T22:40:44.152659Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
@@ -541,9 +477,9 @@
     "    center_coord, k=120, element=\"face centers\"\n",
     ")\n",
     "\n",
-    "nn_subset_120.plot.rasterize(\n",
-    "    method=\"polygon\",\n",
-    "    exclude_antimeridian=True,\n",
+    "nn_subset_120.plot(\n",
+    "    rasterize=True,\n",
+    "    periodic_elements=\"exclude\",\n",
     "    clim=clim,\n",
     "    title=\"Nearest Neighbor Subset (120 Faces)\",\n",
     "    **plot_opts,\n",
@@ -566,10 +502,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:41:27.554399Z",
-     "start_time": "2024-06-07T22:41:25.616381Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
@@ -581,9 +513,9 @@
     "    center_coord, k=1, element=\"face centers\"\n",
     ")\n",
     "\n",
-    "nn_subset_1.plot.rasterize(\n",
-    "    method=\"polygon\",\n",
-    "    exclude_antimeridian=True,\n",
+    "nn_subset_1.plot(\n",
+    "    rasterize=True,\n",
+    "    periodic_elements=\"exclude\",\n",
     "    clim=clim,\n",
     "    title=\"Nearest Neighbor Subset (Closest Face)\",\n",
     "    **plot_opts,\n",
@@ -610,10 +542,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-07T22:43:24.764492Z",
-     "start_time": "2024-06-07T22:43:24.759423Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false

diff --git a/test/test_remap.py b/test/test_remap.py
@@ -214,6 +214,16 @@ def test_value_errors(self):
                 remap_to="nodes"
             )
 
+    def test_preserve_coordinates(self):
+        # Dataset with 'time' coordinates
+        source_uxds = ux.open_dataset(gridfile_geoflow, dsfile_v1_geoflow)
+        destination_grid = ux.open_grid(mpasfile_QU)
+
+        res = source_uxds.remap.nearest_neighbor(destination_grid=destination_grid)
+
+        assert "time" in res.coords
+
+
 
 class TestInverseDistanceWeightedRemapping(TestCase):
     """Testing for inverse distance weighted remapping."""

diff --git a/test/test_subset.py b/test/test_subset.py
@@ -102,3 +102,21 @@ def test_grid_bounding_box_subset():
 
             grid_subset_antimeridian = grid.subset.bounding_box(
                 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
diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py
@@ -1085,20 +1085,18 @@ def _slice_from_grid(self, sliced_grid):
         """Slices a  ``UxDataArray`` from a sliced ``Grid``, using cached
         indices to correctly slice the data variable."""
 
-        from uxarray.core.dataarray import UxDataArray
-
         if self._face_centered():
-            d_var = self.isel(
+            da_sliced = self.isel(
                 n_face=sliced_grid._ds["subgrid_face_indices"], ignore_grid=True
             )
 
         elif self._edge_centered():
-            d_var = self.isel(
+            da_sliced = self.isel(
                 n_edge=sliced_grid._ds["subgrid_edge_indices"], ignore_grid=True
             )
 
         elif self._node_centered():
-            d_var = self.isel(
+            da_sliced = self.isel(
                 n_node=sliced_grid._ds["subgrid_node_indices"], ignore_grid=True
             )
 
@@ -1107,14 +1105,7 @@ def _slice_from_grid(self, sliced_grid):
                 "Data variable must be either node, edge, or face centered."
             )
 
-        return UxDataArray(
-            uxgrid=sliced_grid,
-            data=d_var,
-            name=self.name,
-            coords=self.coords,
-            dims=self.dims,
-            attrs=self.attrs,
-        )
+        return UxDataArray(da_sliced, uxgrid=sliced_grid)
 
     def get_dual(self):
         """Compute the dual mesh for a data array, returns a new data array

diff --git a/uxarray/remap/inverse_distance_weighted.py b/uxarray/remap/inverse_distance_weighted.py
@@ -12,6 +12,8 @@
 from uxarray.grid import Grid
 import warnings
 
+from copy import deepcopy
+
 from uxarray.remap.utils import _remap_grid_parse
 
 
@@ -144,13 +146,19 @@ def _inverse_distance_weighted_remap_uxda(
         power,
         k,
     )
+
+    # preserve only non-spatial coordinates
+    destination_coords = deepcopy(source_uxda.coords)
+    if destination_dim in destination_coords:
+        del destination_coords[destination_dim]
+
     # construct data array for remapping variable
     uxda_remap = uxarray.core.dataarray.UxDataArray(
         data=destination_data,
         name=source_uxda.name,
-        coords=source_uxda.coords,
         dims=destination_dims,
         uxgrid=destination_grid,
+        coords=destination_coords,
     )
 
     # return UxDataArray with remapped variable