From 3d6acfefbc7842078c56572d0c698427f5a7e4f4 Mon Sep 17 00:00:00 2001 From: Matthew Iannucci Date: Wed, 24 Jul 2024 16:57:03 -0400 Subject: [PATCH] WIP --- tests/test_grids/test_ugrid.py | 17 +++++++++++++++++ xarray_subset_grid/accessor.py | 22 ++++++++++++++++++++++ xarray_subset_grid/grid.py | 25 +++++++++++++++---------- 3 files changed, 54 insertions(+), 10 deletions(-) diff --git a/tests/test_grids/test_ugrid.py b/tests/test_grids/test_ugrid.py index d97fb54..f8cebba 100644 --- a/tests/test_grids/test_ugrid.py +++ b/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",) diff --git a/xarray_subset_grid/accessor.py b/xarray_subset_grid/accessor.py index 53c8f7c..9011ec6 100644 --- a/xarray_subset_grid/accessor.py +++ b/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 diff --git a/xarray_subset_grid/grid.py b/xarray_subset_grid/grid.py index 199fa7c..aa082d4 100644 --- a/xarray_subset_grid/grid.py +++ b/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