Make special case for fvcom sigma layer subsetting

xarray_subset_grid/accessor.py

@@ -5,9 +5,9 @@
 import xarray as xr
 
 from xarray_subset_grid.grid import Grid
-from xarray_subset_grid.grids import RegularGrid, RegularGrid2d, SGrid, UGrid
+from xarray_subset_grid.grids import FVCOMGrid, RegularGrid, RegularGrid2d, SGrid, UGrid
 
-_grid_impls = [UGrid, SGrid, RegularGrid2d, RegularGrid]
+_grid_impls = [FVCOMGrid, UGrid, SGrid, RegularGrid2d, RegularGrid]
 
 
 def register_grid_impl(grid_impl: Grid, priority: int = 0):

xarray_subset_grid/grid.py

@@ -73,7 +73,7 @@ def subset_bottom_level(self, ds: xr.Dataset) -> xr.Dataset:
         if positive_direction == "down":
             return self.subset_vertical_level(ds, FLOAT_MAX, method="nearest")
         else:
-            return self.subset_vertical_level(ds, -.975, method="nearest")
+            return self.subset_vertical_level(ds, -0.975, method="nearest")
 
     def subset_top_level(self, ds: xr.Dataset) -> xr.Dataset:
         """Subset the dataset to the top level according to the datasets CF metadata
@@ -101,17 +101,8 @@ def subset_vertical_level(
             return ds
 
         vertical_coords = ds.cf.coordinates["vertical"]
-        if all([i in ds.indexes for i in vertical_coords]):
-            selection = {coord: level for coord in vertical_coords}
-            return ds.sel(selection, method=method)
-        else:
-            # Otherwise, the vertical coordinates are not indexible, so we have to find the
-            # closest levels manually, then slice the dataset using the found level
-            selections = {}
-            for coord in vertical_coords:
-                elevation_index = int(np.absolute(ds[coord] - level).argmin().values)
-                selections[coord] = elevation_index
-            return ds.isel(selections)
+        selection = {coord: level for coord in vertical_coords}
+        return ds.sel(selection, method=method)
 
     def subset_vertical_levels(
         self, ds: xr.Dataset, levels: tuple[float, float], method: str | None = None
@@ -132,20 +123,8 @@ def subset_vertical_levels(
             raise ValueError("The minimum level must be smaller than the maximum level")
 
         vertical_coords = ds.cf.coordinates["vertical"]
-        if all([i in ds.indexes for i in vertical_coords]):
-            selection = {coord: slice(levels[0], levels[1]) for coord in vertical_coords}
-            return ds.sel(selection, method=method)
-        else:
-            # Otherwise, the vertical coordinates are not indexible, so we have to find the
-            # closest levels manually, then slice the dataset using the found levels
-            selections = {}
-            for coord in vertical_coords:
-                da_elevations = ds[coord]
-                elevation_indexes = [
-                    int(np.absolute(da_elevations - level).argmin().values) for level in levels
-                ]
-                selections[coord] = slice(*elevation_indexes)
-            return ds.isel(selections)
+        selection = {coord: slice(levels[0], levels[1]) for coord in vertical_coords}
+        return ds.sel(selection, method=method)
 
     @abstractmethod
     def subset_polygon(

xarray_subset_grid/grids/__init__.py

@@ -1,3 +1,4 @@
+from .fvcom_grid import FVCOMGrid #noqa
 from .regular_grid import RegularGrid #noqa
 from .regular_grid_2d import RegularGrid2d #noqa
 from .sgrid import SGrid #noqa

xarray_subset_grid/grids/fvcom_grid.py

@@ -0,0 +1,56 @@
+import xarray as xr
+
+from xarray_subset_grid.grids.ugrid import UGrid
+
+
+class FVCOMGrid(UGrid):
+    """Grid implementation for FVCOM datasets, extending the UGrid implementation"""
+
+    @staticmethod
+    def recognize(ds: xr.Dataset) -> bool:
+        """Recognize if the dataset matches the given grid"""
+        is_ugrid = UGrid.recognize(ds)
+        if not is_ugrid:
+            return False
+
+        # For now, the only reason for this subclass is to add support
+        # for subsetting FVCOMs non standard vertical levels
+        is_recognized = False
+        coords = ["siglay", "siglev"]
+        for coord in coords:
+            if coord in ds:
+                if len(ds[coord].shape) > 1 and "node" in ds[coord].dims:
+                    is_recognized = True
+                    break
+
+        return is_recognized
+
+    @property
+    def name(self) -> str:
+        """Name of the grid type"""
+        return "fvcom"
+
+    def subset_vertical_level(
+        self, ds: xr.Dataset, level: float, method: str | None = None
+    ) -> xr.Dataset:
+        """Subset the dataset to the vertical level
+
+        :param ds: The dataset to subset
+        :param level: The vertical level to subset to
+        :param method: The method to use for the selection, this is the
+            same as the method in xarray.Dataset.sel
+        """
+        if not self.has_vertical_levels(ds):
+            return ds
+
+        selections = {}
+        if "siglay" in ds.dims:
+            siglay = ds["siglay"].isel(node=0)
+            elevation_index = int((siglay - level).argmin().values)
+            selections["siglay"] = elevation_index
+        if "siglev" in ds.dims:
+            siglev = ds["siglev"].isel(node=0)
+            elevation_index = int((siglev - level).argmin().values)
+            selections["siglev"] = elevation_index
+
+        return ds.isel(selections)