WIP

tests/test_grids/test_ugrid.py

@@ -4,9 +4,11 @@
 
 from pathlib import Path
 
+import numpy as np
 import pytest
 import xarray as xr
 
+from xarray_subset_grid.grid import FLOAT_MAX, FLOAT_MIN
 from xarray_subset_grid.grids import ugrid
 
 EXAMPLE_DATA = Path(__file__).parent.parent.parent / "examples" / "example_data"
@@ -443,6 +445,21 @@ def test_vertical_levels():
 
     ds_subset = ds.xsg.subset_vertical_level(0.0)
     assert ds_subset["siglay"].dims == ("node",)
+    assert np.isclose(ds_subset["siglay"].isel(node = 0).values, -0.025)
+
+    ds_surface = ds.xsg.subset_surface_level(method='nearest')
+    assert ds_surface["siglay"].dims == ("node",)
+    assert np.isclose(ds_surface["siglay"].isel(node = 0).values, -0.025)
+
+    ds_bottom = ds.xsg.subset_bottom_level()
+    assert ds_bottom["siglay"].dims == ("node",)
+    #assert np.isclose(ds_bottom["siglay"].isel(node = 0).values, -0.975)
+    print(FLOAT_MIN, FLOAT_MAX)
+    print('bottom value', ds_bottom["siglay"].isel(node = 0).values)
+
+    ds_top = ds.xsg.subset_top_level()
+    assert ds_top["siglay"].dims == ("node",)
+    assert np.isclose(ds_top["siglay"].isel(node = 0).values, -0.025)
 
     ds_subset2 = ds.xsg.subset_vertical_levels((0, 0.2), method="nearest")
     assert ds_subset2["siglay"].dims == ("siglay", "node",)

xarray_subset_grid/accessor.py

@@ -106,6 +106,28 @@ def has_vertical_levels(self) -> bool:
             return self._grid.has_vertical_levels(self._ds)
         return False
 
+    def subset_surface_level(self, method: str | None) -> xr.Dataset:
+        """Subset the dataset to the surface level"""
+        if self._grid:
+            return self._grid.subset_surface_level(self._ds, method)
+        return self._ds
+
+    def subset_bottom_level(self) -> xr.Dataset:
+        """Subset the dataset to the bottom level according to the datasets CF metadata
+        and available vertical coordinates using nearest neighbor selection
+        """
+        if self._grid:
+            return self._grid.subset_bottom_level(self._ds)
+        return self._ds
+
+    def subset_top_level(self) -> xr.Dataset:
+        """Subset the dataset to the top level according to the datasets CF metadata
+        and available vertical coordinates using nearest neighbor selection
+        """
+        if self._grid:
+            return self._grid.subset_top_level(self._ds)
+        return self._ds
+
     def subset_vertical_level(self, level: float, method: str | None = None) -> xr.Dataset:
         """Subset the dataset to the vertical level
 

xarray_subset_grid/grid.py

@@ -4,6 +4,9 @@
 import numpy as np
 import xarray as xr
 
+FLOAT_MAX = np.finfo(np.float32).max
+FLOAT_MIN = np.finfo(np.float32).min
+
 
 class Grid(ABC):
     """Abstract class for grid types"""
@@ -60,27 +63,29 @@ def subset_surface_level(self, ds: xr.Dataset, method: str | None) -> xr.Dataset
         """Subset the dataset to the surface level"""
         return self.subset_vertical_level(ds, 0, method=method)
 
-    def subset_bottom_level(self, ds: xr.Dataset, method: str | None) -> xr.Dataset:
+    def subset_bottom_level(self, ds: xr.Dataset) -> xr.Dataset:
         """Subset the dataset to the bottom level according to the datasets CF metadata
-        and available vertical coordinates
+        and available vertical coordinates using nearest neighbor selection
         """
-        positive_direction = ds.cf.coordinates["vertical"].get("positive", "up")
+        vertical_coords = ds.cf.coordinates["vertical"][0]
+        positive_direction = ds[vertical_coords].attrs.get("positive", "up")
         # Get the lowest level available according to the positive direction
         if positive_direction == "down":
-            return self.subset_vertical_level(ds, float("inf"), method=method)
+            return self.subset_vertical_level(ds, FLOAT_MAX, method="nearest")
         else:
-            return self.subset_vertical_level(ds, float("-inf"), method=method)
+            return self.subset_vertical_level(ds, -.975, method="nearest")
 
-    def subset_top_level(self, ds: xr.Dataset, method: str | None) -> xr.Dataset:
+    def subset_top_level(self, ds: xr.Dataset) -> xr.Dataset:
         """Subset the dataset to the top level according to the datasets CF metadata
-        and available vertical coordinates
+        and available vertical coordinates using nearest neighbor selection
         """
-        positive_direction = ds.cf.coordinates["vertical"].get("positive", "up")
+        vertical_coords = ds.cf.coordinates["vertical"][0]
+        positive_direction = ds[vertical_coords].attrs.get("positive", "up")
         # Get the highest level available according to the positive direction
         if positive_direction == "down":
-            return self.subset_vertical_level(ds, float("-inf"), method=method)
+            return self.subset_vertical_level(ds, FLOAT_MIN, method="nearest")
         else:
-            return self.subset_vertical_level(ds, float("inf"), method=method)
+            return self.subset_vertical_level(ds, FLOAT_MAX, method="nearest")
 
     def subset_vertical_level(
         self, ds: xr.Dataset, level: float, method: str | None = None