Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Comment thread
ocefpaf marked this conversation as resolved.
]
21 changes: 10 additions & 11 deletions tests/test_accessor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand All @@ -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
2 changes: 1 addition & 1 deletion tests/test_grids/test_regular_grid.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should probably rename all the classes and functions that refer to "regular grid" to be unknown (1D/2D) grid b/c the logic to subset them is the same. However, that is a breaking change and I want to be sure all devs are onboard before we rename everything.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't like "unknown" -- but "regular grid" is a misnomer -- so I'm all for renaming it "rectangular grid".

And it it's a 2D lat-lon, then it's not a rectangular grid, it's curvilinear, which is what SGRID handles, though it's the simplest case of SGRID. It might be hard to have an SGRID-with-a-proper grid spec and simple curvilinear grid be the same class -- maybe a mix-in for the curvilinear grid logic?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It might be hard to have an SGRID-with-a-proper grid spec and simple curvilinear grid be the same class -- maybe a mix-in for the curvilinear grid logic?

I like the "special case inside SGRID". It is SGRID after all and, even if we crack the fix-the-metadata route, the code to subset would be the same b/c it is likely the most accurate there (TBD).


EXAMPLE_DATA = Path(__file__).parent.parent / "example_data"

Expand Down
7 changes: 2 additions & 5 deletions xarray_subset_grid/accessor.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
import warnings

# from typing import Optional, Union
import numpy as np
import xarray as xr
Expand All @@ -20,7 +18,6 @@
SELFEGrid,
UGrid,
SGrid,
# RegularGrid2d,
RegularGrid,
]

Expand All @@ -45,8 +42,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")
Expand Down
19 changes: 13 additions & 6 deletions xarray_subset_grid/grids/__init__.py
Original file line number Diff line number Diff line change
@@ -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",
]
101 changes: 0 additions & 101 deletions xarray_subset_grid/grids/regular_grid_2d.py

This file was deleted.

1 change: 0 additions & 1 deletion xarray_subset_grid/grids/sgrid.py
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,6 @@ def compute_polygon_subset_selector(
grid_topology = ds[grid_topology_key]
dims = _get_sgrid_dim_coord_names(grid_topology)
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"]
Expand Down
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -19,29 +13,6 @@
normalize_polygon_x_coords,
)

# class RegularGridPolygonSelector(Selector):

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should look at how rioxarray does their polygon subsets. Maybe we can use that instead of custom code.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Worth a look, but oddly the docs don't seem to describe what rioxarray is at a very high level, but it looks like it's for rasters, which is not what we need here :-(

The big distinction is that a raster is a truly regular grid -- it can be defined by a corner, number of pixels and a delta x and delta y.

whereas most of our use cases have delta x and/or delta y that vary.

And, oddly, CF doesn't have a way to support a "proper" raster at all.

And we need the custom polygon code for curvilinear and unstructured grids anyway.

Though now that I think about it -- if you are subsetting a north-aligned rectangular grid, and returning a rectangular grid, then all you can do is make the subset the bounding box of the polygon anyway.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Worth a look, but oddly the docs don't seem to describe what rioxarray is at a very high level, but it looks like it's for rasters, which is not what we need here :-(

Sorry, should've been more specific. It is focused on sub-setting raster images loaded as xarray datasets. The subset method (clip in their method) is generic enough to work with CF datasets that have a project. Ideally this method should be a stand alone library that both rioxarray and xarray-subet-grid could use. However, it is likely that we will re-create their methods centered at CF logic.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ChrisBarker-NOAA here is an example on how to use a polygon to subset with rioxarray. If it wasn't for the rasterio dependency, b/c t brings in gdal, we could just depend on rioxarray. That is why I prefer to re-write their sub-setting mechanism here.

https://gist.github.com/ocefpaf/e0b56132d3ad1d81f020dd18062d4d21

PS: I'm not sure I used xarray-subset-grid polygon feature correctly there. I cannot get the same triangle as rioxarray's clip.

# """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."""
Expand All @@ -53,28 +24,21 @@ 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:
"""
Perform the selection on the 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):
Expand Down Expand Up @@ -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
Expand Down
Loading