Merge pull request #1408 from UXARRAY/rajeeja/fix_issue1406

Fix .sel() wrapper when  a slice indexer and method are used together

Author: rajeeja

test/core/test_api.py

@@ -36,17 +36,21 @@ def test_open_dataset(gridpath, datasetpath, mesh_constants):
     nt.assert_equal(len(uxds_var2_ne30.uxgrid._ds.data_vars), mesh_constants['DATAVARS_outCSne30'])
     nt.assert_equal(uxds_var2_ne30.source_datasets, str(data_path))
 
-def test_open_mf_dataset(gridpath, test_data_dir, mesh_constants):
+def test_open_mf_dataset(gridpath, datasetpath, mesh_constants):
     """Loads multiple datasets with their grid topology file using
     uxarray's open_dataset call."""
 
     grid_path = gridpath("ugrid", "outCSne30", "outCSne30.ug")
-    dsfiles_mf_ne30 = str(test_data_dir) + "/ugrid/outCSne30/outCSne30_*.nc"
+    dsfiles_mf_ne30 = datasetpath(
+        "ugrid",
+        "outCSne30",
+        ["outCSne30_var2.nc", "outCSne30_vortex.nc"],
+    )
     uxds_mf_ne30 = ux.open_mfdataset(grid_path, dsfiles_mf_ne30)
 
     nt.assert_equal(uxds_mf_ne30.uxgrid.node_lon.size, mesh_constants['NNODES_outCSne30'])
     nt.assert_equal(len(uxds_mf_ne30.uxgrid._ds.data_vars), mesh_constants['DATAVARS_outCSne30'])
-    nt.assert_equal(uxds_mf_ne30.source_datasets, dsfiles_mf_ne30)
+    nt.assert_equal(uxds_mf_ne30.source_datasets, str(dsfiles_mf_ne30))
 
 def test_open_grid(gridpath, mesh_constants):
     """Loads only a grid topology file using uxarray's open_grid call."""

test/core/test_dataset.py

@@ -67,3 +67,35 @@ def test_get_dual(gridpath, datasetpath):
 
     assert isinstance(dual, UxDataset)
     assert len(uxds.data_vars) == len(dual.data_vars)
+
+
+def _load_sel_subset(gridpath, datasetpath):
+    grid_file = gridpath("ugrid", "outCSne30", "outCSne30.ug")
+    data_file = datasetpath("ugrid", "outCSne30", "outCSne30_sel_timeseries.nc")
+    uxds = ux.open_dataset(grid_file, data_file)
+    base_time = np.datetime64("2018-04-28T00:00:00")
+    offsets = np.arange(uxds.sizes["time"], dtype="timedelta64[h]")
+    uxds = uxds.assign_coords(time=(base_time + offsets).astype("datetime64[ns]"))
+    return uxds
+
+
+def test_sel_time_slice(gridpath, datasetpath):
+    uxds = _load_sel_subset(gridpath, datasetpath)
+
+    times = uxds["time"].values
+    sliced = uxds.sel(time=slice(times[0], times[2]))
+
+    assert sliced.dims["time"] == 3
+    np.testing.assert_array_equal(sliced["time"].values, times[:3])
+
+
+def test_sel_method_forwarded(gridpath, datasetpath):
+    uxds = _load_sel_subset(gridpath, datasetpath)
+
+    target = np.datetime64("2018-04-28T02:20:00")
+    nearest = uxds.sel(time=target, method="nearest")
+
+    np.testing.assert_array_equal(
+        nearest["time"].values,
+        np.array(uxds["time"].values[2], dtype="datetime64[ns]"),
+    )

test/meshfiles/ugrid/outCSne30/README.md

@@ -0,0 +1,11 @@
+# outCSne30 UGRID assets
+
+This folder houses the NE30 cubed-sphere mesh (5400 faces) and a few companion datasets used throughout the test suite.
+
+## Files
+
+- `outCSne30.ug` – canonical NE30 UGRID mesh with 5400 faces (nMesh2_face), max 4 vertices per face, and 5402 nodes. Provides `Mesh2_face_nodes`, `Mesh2_node_x`, and `Mesh2_node_y`.
+- `outCSne30_var2.nc` – single-variable dataset (`var2(ncol=5400)`) used for basic integration/zonal workflows.
+- `outCSne30_vortex.nc` – single-variable dataset (`psi(ncol=5400)`) containing the barotropic vortex test case.
+- `outCSne30_sel_timeseries.nc` – synthetic 6-step hourly time series (`psi(time=6, ncol=5400)`) where each time slice ramps with longitude; input for the `.sel()` regression tests. Regenerate via `python generate_sel_timeseries.py`.
+- `generate_sel_timeseries.py` – helper script that rewrites the synthetic time series using `outCSne30_var2.nc` as a template for face count.

test/meshfiles/ugrid/outCSne30/generate_sel_timeseries.py

@@ -0,0 +1,36 @@
+"""Generate the synthetic NE30 time series used in sel() regression tests.
+
+Run from repo root: ``python test/meshfiles/ugrid/outCSne30/generate_sel_timeseries.py``
+"""
+
+from pathlib import Path
+
+import numpy as np
+import xarray as xr
+
+TIME_STEPS = 6
+
+def main() -> None:
+    base = Path(__file__).parent
+
+    data_src = base / "outCSne30_var2.nc"
+    template = xr.open_dataset(data_src)
+    n_face = template.dims["ncol"]
+    template.close()
+
+    times = np.arange(TIME_STEPS, dtype=np.int64)
+    time_vals = (np.datetime64("2018-04-28T00:00:00") + times * np.timedelta64(1, "h"))
+    faces = np.arange(n_face, dtype=np.float32)
+    field = (times[:, None].astype(np.float32) + faces[None, :] / 100.0)
+
+    ds = xr.Dataset(
+        {"psi": (("time", "ncol"), field)},
+        coords={"time": time_vals},
+        attrs={"source": "Synthetic field for sel regression"},
+    )
+    data_path = base / "outCSne30_sel_timeseries.nc"
+    ds.to_netcdf(data_path)
+
+
+if __name__ == "__main__":
+    main()

test/meshfiles/ugrid/outCSne30/outCSne30_sel_timeseries.nc

@@ -690,7 +690,13 @@ def sel(
         self, indexers=None, method=None, tolerance=None, drop=False, **indexers_kwargs
     ):
         return UxDataset(
-            self.to_xarray().sel(indexers, tolerance, drop, **indexers_kwargs),
+            self.to_xarray().sel(
+                indexers=indexers,
+                method=method,
+                tolerance=tolerance,
+                drop=drop,
+                **indexers_kwargs,
+            ),
             uxgrid=self.uxgrid,
         )