From 9d05bfc93ae9088b3b5f5f6e0fd32c68461d481f Mon Sep 17 00:00:00 2001 From: Nicholas Delli Carpini Date: Tue, 7 Nov 2023 15:23:53 -0500 Subject: [PATCH 1/2] add selfe grid support --- xpublish_wms/grid.py | 168 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 165 insertions(+), 3 deletions(-) diff --git a/xpublish_wms/grid.py b/xpublish_wms/grid.py index e0bc6ae..a260d46 100644 --- a/xpublish_wms/grid.py +++ b/xpublish_wms/grid.py @@ -273,7 +273,7 @@ def __init__(self, ds: xr.Dataset): @staticmethod def recognize(ds: xr.Dataset) -> bool: - return ds.attrs.get("title", "").startswith("HYCOM") + return ds.attrs.get("title", "").lower().startswith("hycom") @property def name(self) -> str: @@ -336,7 +336,7 @@ def __init__(self, ds: xr.Dataset): @staticmethod def recognize(ds: xr.Dataset) -> bool: - return ds.attrs.get("source", "").startswith("FVCOM") + return ds.attrs.get("source", "").lower().startswith("fvcom") @property def name(self) -> str: @@ -528,7 +528,169 @@ def tessellate(self, da: xr.DataArray) -> np.ndarray: ).triangles -_grid_impls = [HYCOMGrid, FVCOMGrid, ROMSGrid, RegularGrid] +class SELFEGrid(Grid): + def __init__(self, ds: xr.Dataset): + self.ds = ds + + @staticmethod + def recognize(ds: xr.Dataset) -> bool: + return ds.attrs.get("source", "").lower().startswith("selfe") + + @property + def name(self) -> str: + return "selfe" + + @property + def render_method(self) -> RenderMethod: + return RenderMethod.Triangle + + @property + def crs(self) -> str: + return "EPSG:4326" + + def has_elevation(self, da: xr.DataArray) -> bool: + return "nv" in da.dims + + def elevation_units(self, da: xr.DataArray) -> Optional[str]: + if self.has_elevation(da): + return "sigma" + else: + return None + + def elevation_positive_direction(self, da: xr.DataArray) -> Optional[str]: + if self.has_elevation(da): + return self.ds.cf["vertical"].attrs.get("positive", "up") + else: + return None + + def elevations(self, da: xr.DataArray) -> Optional[xr.DataArray]: + if self.has_elevation(da): + # clean up elevation values using nv as index array + vertical = self.ds.cf["vertical"].values + elevations = [] + for index in da.nv.values: + if index < len(vertical): + elevations.append(vertical[index]) + + return xr.DataArray( + data=elevations, + dims=da.nv.dims, + coords=da.nv.coords, + name=self.ds.cf["vertical"].name, + attrs=self.ds.cf["vertical"].attrs, + ) + + return None + + def sel_lat_lng( + self, + subset: xr.Dataset, + lng, + lat, + parameters, + ) -> Tuple[xr.Dataset, list, list]: + """Select the given dataset by the given lon/lat and optional elevation""" + + lng_rad = np.deg2rad(subset.cf["longitude"]) + lat_rad = np.deg2rad(subset.cf["latitude"]) + + stacked = np.stack([lng_rad, lat_rad], axis=-1) + tree = BallTree(stacked, leaf_size=2, metric="haversine") + + idx = tree.query( + [[np.deg2rad((360 + lng) if lng < 0 else lng), np.deg2rad(lat)]], + return_distance=False, + ) + + if "nele" in subset.dims: + subset = subset.isel(nele=idx[0][0]) + else: + subset = subset.isel(node=idx[0][0]) + + x_axis = [strip_float(subset.cf["longitude"])] + y_axis = [strip_float(subset.cf["latitude"])] + return subset, x_axis, y_axis + + def select_by_elevation( + self, + da: xr.DataArray, + elevations: Optional[Sequence[float]], + ) -> xr.DataArray: + """Select the given data array by elevation""" + if not self.has_elevation(da): + return da + + if ( + elevations is None + or len(elevations) == 0 + or all(v is None for v in elevations) + ): + elevations = [0.0] + + da_elevations = self.elevations(da) + elevation_index = [ + int(np.absolute(da_elevations - elevation).argmin().values) + for elevation in elevations + ] + if len(elevation_index) == 1: + elevation_index = elevation_index[0] + + if "vertical" not in da.cf: + if da.nv.shape[0] > da_elevations.shape[0]: + # need to fill the nv array w/ nan to match dimensions of the var's nv + new_nv_data = da_elevations.values.tolist() + for i in range(da.nv.shape[0] - da_elevations.shape[0]): + new_nv_data.append(np.nan) + + da_elevations = xr.DataArray( + data=new_nv_data, + dims=da_elevations.dims, + coords=da_elevations.coords, + name=da_elevations.name, + attrs=da_elevations.attrs, + ) + + da.__setitem__("nv", da_elevations) + + if "vertical" in da.cf: + da = da.cf.isel(vertical=elevation_index) + + return da + + def project(self, da: xr.DataArray, crs: str) -> Any: + if crs == "EPSG:4326": + da = da.assign_coords({"x": da.cf["longitude"], "y": da.cf["latitude"]}) + elif crs == "EPSG:3857": + x, y = to_mercator.transform(da.cf["longitude"], da.cf["latitude"]) + x_chunks = ( + da.cf["longitude"].chunks if da.cf["longitude"].chunks else x.shape + ) + y_chunks = da.cf["latitude"].chunks if da.cf["latitude"].chunks else y.shape + + da = da.assign_coords( + { + "x": ( + da.cf["longitude"].dims, + dask_array.from_array(x, chunks=x_chunks), + ), + "y": ( + da.cf["latitude"].dims, + dask_array.from_array(y, chunks=y_chunks), + ), + }, + ) + + da = da.unify_chunks() + return da + + def tessellate(self, da: xr.DataArray) -> np.ndarray: + return tri.Triangulation( + da.cf["longitude"], + da.cf["latitude"], + self.ds.ele[0].T - 1, + ).triangles + +_grid_impls = [HYCOMGrid, FVCOMGrid, SELFEGrid, ROMSGrid, RegularGrid] def register_grid_impl(grid_impl: Grid, priority: int = 0): From 4c44e9a87d7127e10276dba2393feef6b8c7f72b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 7 Nov 2023 20:25:50 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xpublish_wms/grid.py | 1 + 1 file changed, 1 insertion(+) diff --git a/xpublish_wms/grid.py b/xpublish_wms/grid.py index a260d46..63ad7d7 100644 --- a/xpublish_wms/grid.py +++ b/xpublish_wms/grid.py @@ -690,6 +690,7 @@ def tessellate(self, da: xr.DataArray) -> np.ndarray: self.ds.ele[0].T - 1, ).triangles + _grid_impls = [HYCOMGrid, FVCOMGrid, SELFEGrid, ROMSGrid, RegularGrid]