SELFE Support

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,170 @@ 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):