diff --git a/tests/conftest.py b/tests/conftest.py index b9f8215..551bf45 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -72,4 +72,5 @@ def pytest_collection_modifyitems(config, items): EXAMPLE_DATA / "2D-rectangular_grid_wind.nc", EXAMPLE_DATA / "rectangular_grid_decreasing.nc", EXAMPLE_DATA / "AMSEAS-subset.nc", + EXAMPLE_DATA / "unknown_2D-rtofs_example.nc", ] diff --git a/tests/test_accessor.py b/tests/test_accessor.py index 858df65..e3356f8 100644 --- a/tests/test_accessor.py +++ b/tests/test_accessor.py @@ -7,9 +7,8 @@ def test_accessor_warns_when_no_grid_recognized(): ds = xr.Dataset() - with pytest.warns(UserWarning, match="no grid type"): - accessor = ds.xsg - assert accessor.grid is None + with pytest.raises(ValueError, match="Cannot find grid or coords for"): + ds.xsg def test_subset_polygon_and_bbox_return_none_without_grid(): @@ -22,19 +21,19 @@ def test_subset_polygon_and_bbox_return_none_without_grid(): [-72.0, 41.0], ] ) - with pytest.warns(UserWarning, match="no grid type"): - assert ds.xsg.subset_polygon(poly) is None - assert ds.xsg.subset_bbox((-72, 39, -70, 41)) is None + with pytest.raises(ValueError, match="Cannot find grid or coords for"): + ds.xsg.subset_polygon(poly) + with pytest.raises(ValueError, match="Cannot find grid or coords for"): + ds.xsg.subset_bbox((-72, 39, -70, 41)) def test_subset_vars_raises_without_grid(): ds = xr.Dataset({"a": (("x",), [1, 2, 3])}) - with pytest.warns(UserWarning, match="no grid type"): - with pytest.raises(ValueError, match="subset_vars requires a recognized grid"): - ds.xsg.subset_vars(["a"]) + with pytest.raises(ValueError, match="Cannot find grid or coords for"): + ds.xsg.subset_vars(["a"]) def test_has_vertical_levels_false_without_grid(): ds = xr.Dataset() - with pytest.warns(UserWarning, match="no grid type"): - assert ds.xsg.has_vertical_levels is False + with pytest.raises(ValueError, match="Cannot find grid or coords for"): + ds.xsg.has_vertical_levels diff --git a/tests/test_grids/test_regular_grid.py b/tests/test_grids/test_regular_grid.py index 42bd5cd..9274084 100644 --- a/tests/test_grids/test_regular_grid.py +++ b/tests/test_grids/test_regular_grid.py @@ -15,7 +15,7 @@ import xarray as xr from tests.conftest import RGRID_FILES, SGRID_FILES, UGRID_FILES -from xarray_subset_grid.grids.regular_grid import RegularGrid +from xarray_subset_grid.grids.unknown_grid import RegularGrid EXAMPLE_DATA = Path(__file__).parent.parent / "example_data" diff --git a/xarray_subset_grid/accessor.py b/xarray_subset_grid/accessor.py index 45f5929..8bc26fc 100644 --- a/xarray_subset_grid/accessor.py +++ b/xarray_subset_grid/accessor.py @@ -1,4 +1,3 @@ -import warnings from collections.abc import Iterable # from typing import Optional, Union @@ -21,7 +20,6 @@ SELFEGrid, UGrid, SGrid, - # RegularGrid2d, RegularGrid, ] @@ -46,8 +44,8 @@ def grid_factory(ds: xr.Dataset) -> Grid | None: for grid_impl in _grid_impls: if grid_impl.recognize(ds): return grid_impl() - warnings.warn("no grid type found in this dataset") - return None + msg = f"Cannot find grid or coords for\n{ds.cf}" + raise ValueError(msg) @xr.register_dataset_accessor("xsg") diff --git a/xarray_subset_grid/grids/__init__.py b/xarray_subset_grid/grids/__init__.py index 23017e9..103a46c 100644 --- a/xarray_subset_grid/grids/__init__.py +++ b/xarray_subset_grid/grids/__init__.py @@ -1,6 +1,13 @@ -from .fvcom_grid import FVCOMGrid # noqa -from .regular_grid import RegularGrid # noqa -from .regular_grid_2d import RegularGrid2d # noqa -from .selfe_grid import SELFEGrid # noqa -from .sgrid import SGrid # noqa -from .ugrid import UGrid # noqa +from .fvcom_grid import FVCOMGrid +from .selfe_grid import SELFEGrid +from .sgrid import SGrid +from .ugrid import UGrid +from .unknown_grid import RegularGrid + +__all__ = [ + "FVCOMGrid", + "RegularGrid", + "SELFEGrid", + "SGrid", + "UGrid", +] diff --git a/xarray_subset_grid/grids/regular_grid_2d.py b/xarray_subset_grid/grids/regular_grid_2d.py deleted file mode 100644 index 09ab9c2..0000000 --- a/xarray_subset_grid/grids/regular_grid_2d.py +++ /dev/null @@ -1,104 +0,0 @@ -# 2D and 3D should share a lot of code -# -# do we need a separate class for this? -# This doesn't appear, right now, to check for depth to dedermine if it's 2D - - -import numpy as np -import xarray as xr - -from xarray_subset_grid.grid import Grid -from xarray_subset_grid.selector import Selector -from xarray_subset_grid.utils import compute_2d_subset_mask - - -class RegularGrid2dSelector(Selector): - polygon: list[tuple[float, float]] | np.ndarray - _subset_mask: xr.DataArray - - def __init__( - self, polygon: list[tuple[float, float]] | np.ndarray, subset_mask: xr.DataArray, name: str - ): - super().__init__() - self.name = name - self.polygon = polygon - self._subset_mask = subset_mask - - def select(self, ds: xr.Dataset) -> xr.Dataset: - # First, we need to add the mask as a variable in the dataset - # so that we can use it to mask and drop via xr.where, which requires that - # the mask and data have the same shape and both are DataArrays with matching - # dimensions - ds_subset = ds.assign(subset_mask=self._subset_mask) - - # Now we can use the mask to subset the data - ds_subset = ds_subset.where(ds_subset.subset_mask, drop=True).drop_encoding() - ds_subset.drop_vars("subset_mask") - return ds_subset - - -class RegularGrid2d(Grid): - """Grid implementation for 2D regular grids.""" - - @staticmethod - def recognize(ds) -> bool: - """Recognize if the dataset matches the given grid.""" - lat = ds.cf.coordinates.get("latitude", None) - lon = ds.cf.coordinates.get("longitude", None) - if lat is None or lon is None: - return False - - # Make sure the coordinates are 1D and match - lat_dim = ds[lat[0]].dims - ndim = ds[lat[0]].ndim - lon_dim = ds[lon[0]].dims - return lat_dim == lon_dim and ndim == 2 - - @property - def name(self) -> str: - """Name of the grid type.""" - return "regular_grid_2d" - - def grid_vars(self, ds: xr.Dataset) -> set[str]: - """Set of grid variables. - - These variables are used to define the grid and thus should be - kept when subsetting the dataset - """ - lat = ds.cf.coordinates["latitude"][0] - lon = ds.cf.coordinates["longitude"][0] - return {lat, lon} - - def data_vars(self, ds: xr.Dataset) -> set[str]: - """Set of data variables. - - These variables exist on the grid and are available to used for - data analysis. These can be discarded when subsetting the - dataset when they are not needed. - """ - lat = ds.cf.coordinates["latitude"][0] - lon = ds.cf.coordinates["longitude"][0] - data_vars = { - var.name - for var in ds.data_vars.values() - if var.name not in {lat, lon} - and "latitude" in var.cf.coordinates - and "longitude" in var.cf.coordinates - } - return data_vars - - def compute_polygon_subset_selector( - self, - ds: xr.Dataset, - polygon: list[tuple[float, float]] | np.ndarray, - name: str | None = None, - ) -> Selector: - lat = ds.cf["latitude"] - lon = ds.cf["longitude"] - subset_mask = compute_2d_subset_mask(lat=lat, lon=lon, polygon=polygon) - - return RegularGrid2dSelector( - polygon=polygon, - subset_mask=subset_mask, - name=name or "selector", - ) diff --git a/xarray_subset_grid/grids/sgrid.py b/xarray_subset_grid/grids/sgrid.py index c6822f2..44355f8 100644 --- a/xarray_subset_grid/grids/sgrid.py +++ b/xarray_subset_grid/grids/sgrid.py @@ -105,7 +105,6 @@ def compute_polygon_subset_selector( grid_topology_key = ds.cf.cf_roles["grid_topology"][0] grid_topology = ds[grid_topology_key] subset_masks: list[tuple[list[str], xr.DataArray]] = [] - node_info = _get_location_info_from_topology(grid_topology, "node") node_dims = node_info["dims"] node_coords = node_info["coords"] diff --git a/xarray_subset_grid/grids/regular_grid.py b/xarray_subset_grid/grids/unknown_grid.py similarity index 61% rename from xarray_subset_grid/grids/regular_grid.py rename to xarray_subset_grid/grids/unknown_grid.py index 14432ba..ba0042f 100644 --- a/xarray_subset_grid/grids/regular_grid.py +++ b/xarray_subset_grid/grids/unknown_grid.py @@ -1,12 +1,6 @@ """ -Implementation for Rectangular grids +Implementation for any unknown 1D and 2D grids -NOTE: it's called "regular", but I think this will work for -any (grid-aligned) rectangular grid: - -The grid is defined by two 1-d arrays. - -* delta_lat and delta_lon do not have to be constant. """ import numpy as np @@ -19,29 +13,6 @@ normalize_polygon_x_coords, ) -# class RegularGridPolygonSelector(Selector): -# """Polygon Selector for regular lat/lon grids.""" -# # with a regular grid, you have to select the full bounding box anyway -# # this this simply computes the bounding box, and used that - -# polygon: list[tuple[float, float]] | np.ndarray -# _polygon_mask: xr.DataArray - -# def __init__(self, polygon: list[tuple[float, float]] | np.ndarray, mask: xr.DataArray, -# name: str): -# super().__init__() -# self.name = name -# self.polygon = polygon -# self.polygon_mask = mask - -# def select(self, ds: xr.Dataset) -> xr.Dataset: -# """Perform the selection on the dataset.""" -# ds_subset = ds.cf.isel( -# lon=self._polygon_mask, -# lat=self._polygon_mask, -# ) -# return ds_subset - class RegularGridBBoxSelector(Selector): """Selector for regular lat/lng grids.""" @@ -53,8 +24,6 @@ class RegularGridBBoxSelector(Selector): def __init__(self, bbox: tuple[float, float, float, float]): super().__init__() self.bbox = bbox - self._longitude_selection = slice(bbox[0], bbox[2]) - self._latitude_selection = slice(bbox[1], bbox[3]) def select(self, ds: xr.Dataset) -> xr.Dataset: """ @@ -62,19 +31,14 @@ def select(self, ds: xr.Dataset) -> xr.Dataset: """ lat = ds[ds.cf.coordinates.get("latitude")[0]] lon = ds[ds.cf.coordinates.get("longitude")[0]] - if np.all(np.diff(lat) < 0): - # swap the slice if the latitudes are descending - self._latitude_selection = slice( - self._latitude_selection.stop, self._latitude_selection.start - ) - # and np.all(np.diff(lon) > 0): - if np.all(np.diff(lon) < 0): - # swap the slice if the longitudes are descending - self._longitude_selection = slice( - self._longitude_selection.stop, self._longitude_selection.start - ) - return ds.cf.sel(lon=self._longitude_selection, lat=self._latitude_selection) + xmin, xmax = self.bbox[0], self.bbox[2] + ymin, ymax = self.bbox[1], self.bbox[3] + + return ds.where( + (xmin <= lon) & (lon <= xmax) & (ymin <= lat) & (lat <= ymax), + drop=True, + ) class RegularGridPolygonSelector(RegularGridBBoxSelector): @@ -102,24 +66,25 @@ def recognize(ds: xr.Dataset) -> bool: """ Recognize if the dataset matches the given grid. """ + # Short-circut to known grids. + grid = ds.variables.get("grid", None) + if grid is not None: + return False + + # Are coords available? lat = ds.cf.coordinates.get("latitude", None) lon = ds.cf.coordinates.get("longitude", None) if lat is None or lon is None: return False - # choose first one -- valid assumption?? - lat = lat[0] - lon = lon[0] - # Make sure the coordinates are 1D and match - if not (1 == ds[lat].ndim == ds[lon].ndim): + # Must have only one lon, lat! + if (len(lat) != len(lon)) or len(lat) > 1: return False - # make sure that at least one variable is using both the - # latitude and longitude dimensions - # (ugrids have both coordinates, but not both dimensions) - for var_name, var in ds.data_vars.items(): - if (lon in var.dims) and (lat in var.dims): - return True + # If lat, lon are consistent and not 3D, we have a grid! + lat, lon = lat[0], lon[0] + if (ds[lon].ndim == ds[lat].ndim) or ds[lon].ndim < 3: + return True return False @property