Extra dimension selectors (#93)

* Selection working but no data rendering

* Cleanup debug print

* Get additional coords working with metadata

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* get caps working with additional dims

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* GFI

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove unused import

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

xpublish_wms/grids/__init__.py

@@ -116,6 +116,12 @@ def select_by_elevation(
         else:
             return self._grid.select_by_elevation(da, elevations)
 
+    def additional_coords(self, da: xr.DataArray) -> list[str]:
+        if self._grid is None:
+            return None
+        else:
+            return self._grid.additional_coords(da)
+
     def mask(
         self,
         da: Union[xr.DataArray, xr.Dataset],

xpublish_wms/grids/grid.py

@@ -100,6 +100,32 @@ def select_by_elevation(
 
         return da
 
+    def additional_coords(self, da: xr.DataArray) -> list[str]:
+        """Return the additional coordinate dimensions for the given data array
+
+        Given a data array with a shape of (time, latitude, longitude),
+        this function would return [].
+
+        Given a data array with a shape of (time, latitude, longitude, vertical),
+        this function would return [].
+
+        Given a data array with a shape of (band, latitude, longitude),
+        this function would return ["band"].
+        """
+        lat_dim = da.cf["latitude"].name
+        lng_dim = da.cf["longitude"].name
+        filter_dims = [lat_dim, lng_dim]
+
+        time_dim = da.cf.coords.get("time", None)
+        if time_dim is not None:
+            filter_dims.append(time_dim.name)
+
+        elevation_dim = da.cf.coords.get("vertical", None)
+        if elevation_dim is not None:
+            filter_dims.append(elevation_dim.name)
+
+        return [dim for dim in da.dims if dim not in filter_dims]
+
     def mask(
         self,
         da: Union[xr.DataArray, xr.Dataset],

xpublish_wms/wms/get_capabilities.py

@@ -176,12 +176,16 @@ def get_capabilities(ds: xr.Dataset, request: Request, query_params: dict) -> Re
         create_text_element(
             layer,
             "Title",
-            attrs.get("long_name", attrs.get("name", var)),
+            attrs.get("standard_name", attrs.get("name", var)),
         )
+        long_name = attrs.get("long_name", None)
+        if long_name is not None or type(long_name) is not str:
+            long_name = attrs.get("name", var)
+
         create_text_element(
             layer,
             "Abstract",
-            attrs.get("long_name", attrs.get("name", var)),
+            long_name,
         )
         create_text_element(layer, crs_tag, "EPSG:4326")
         create_text_element(layer, crs_tag, "EPSG:3857")
@@ -202,7 +206,7 @@ def get_capabilities(ds: xr.Dataset, request: Request, query_params: dict) -> Re
 
         ET.SubElement(layer, "BoundingBox", attrib=bounds)
 
-        if "time" in da.cf.coords:
+        if "time" in da.cf.coords and "time" in da.cf.indexes:
             times = format_timestamp(da.cf["time"])
             default_time = format_timestamp(
                 da.cf["time"].cf.sel(time=current_date, method="nearest"),
@@ -240,6 +244,21 @@ def get_capabilities(ds: xr.Dataset, request: Request, query_params: dict) -> Re
             )
             elevation_dimension_element.text = ",".join(elevations)
 
+        additonal_coords = ds.gridded.additional_coords(da)
+        for coord in additonal_coords:
+            values = da.cf.coords[coord].values
+            units = da.cf.coords[coord].attrs.get("units", "")
+            coord_element = ET.SubElement(
+                layer,
+                "Dimension",
+                attrib={
+                    "name": coord,
+                    "units": units,
+                    "default": str(values[0]),
+                },
+            )
+            coord_element.text = ",".join([str(v) for v in values])
+
         for style in styles:
             style_element = ET.SubElement(
                 layer,

xpublish_wms/wms/get_feature_info.py

@@ -21,6 +21,7 @@ def create_parameter_feature_data(
     z_axis,
     x_axis,
     y_axis,
+    extra_axes: dict = None,
     values=None,
     name=None,
     id=None,
@@ -56,6 +57,8 @@ def create_parameter_feature_data(
         axis_names.append("t")
     if z_axis is not None:
         axis_names.append("z")
+    if extra_axes is not None:
+        axis_names.extend(extra_axes.keys())
     axis_names.extend(["x", "y"])
 
     values = values if values is not None else ds[parameter]
@@ -65,6 +68,9 @@ def create_parameter_feature_data(
         shape.append(len(t_axis))
     if z_axis is not None:
         shape.append(len(z_axis))
+    if extra_axes is not None:
+        for _axis in extra_axes.values():
+            shape.append(1)
 
     if isinstance(values, float):
         shape.extend([1, 1])
@@ -215,6 +221,16 @@ def get_feature_info(ds: xr.Dataset, query: dict) -> Response:
         else:
             z_axis = selected_ds.cf["vertical"].values.tolist()
 
+    additional_queries = {}
+    queries = set(query.keys())
+    for param in parameters:
+        available_coords = selected_ds.gridded.additional_coords(selected_ds[param])
+        valid_quiries = list(set(available_coords) & queries)
+        for q in valid_quiries:
+            additional_queries[q] = query[q]
+
+    selected_ds = selected_ds.sel(additional_queries)
+
     parameter_info = {}
     ranges = {}
 
@@ -226,6 +242,7 @@ def get_feature_info(ds: xr.Dataset, query: dict) -> Response:
             z_axis,
             x_axis,
             y_axis,
+            extra_axes=additional_queries,
         )
         parameter_info[parameter] = info
         ranges[parameter] = range
@@ -243,6 +260,7 @@ def get_feature_info(ds: xr.Dataset, query: dict) -> Response:
             z_axis,
             x_axis,
             y_axis,
+            additional_queries,
             speed,
             "Magnitude of velocity",
             "magnitude_of_velocity",
@@ -258,6 +276,7 @@ def get_feature_info(ds: xr.Dataset, query: dict) -> Response:
             z_axis,
             x_axis,
             y_axis,
+            additional_queries,
             direction,
             "Direction of velocity",
             "direction_of_velocity",
@@ -266,15 +285,22 @@ def get_feature_info(ds: xr.Dataset, query: dict) -> Response:
         parameter_info[direction_parameter_name] = direction_info
         ranges[direction_parameter_name] = direction_range
 
-    axis = {
-        "t": {"values": t_axis} if any_has_time_axis else {},
-        "x": {"values": x_axis},
-        "y": {"values": y_axis},
-        "z": {"values": z_axis} if any_has_vertical_axis else {},
-    }
+    axis = {}
 
     referencing = []
+
+    if len(additional_queries) > 0:
+        for key, value in additional_queries.items():
+            axis[key] = {"values": [value]}
+            referencing.append(
+                {
+                    "coordinates": [key],
+                    "id": "custom",
+                },
+            )
+
     if any_has_time_axis:
+        axis["t"] = {"values": t_axis}
         referencing.append(
             {
                 "coordinates": ["t"],
@@ -285,6 +311,7 @@ def get_feature_info(ds: xr.Dataset, query: dict) -> Response:
             },
         )
     if any_has_vertical_axis:
+        axis["z"] = {"values": z_axis}
         referencing.append(
             {
                 "coordinates": ["z"],
@@ -303,6 +330,8 @@ def get_feature_info(ds: xr.Dataset, query: dict) -> Response:
             },
         )
 
+    axis["x"] = {"values": x_axis}
+    axis["y"] = {"values": y_axis}
     referencing.append(
         {
             "coordinates": ["x", "y"],

xpublish_wms/wms/get_map.py

@@ -42,6 +42,7 @@ class GetMap:
     has_time: bool
     elevation: float = None
     has_elevation: bool
+    dim_selectors: dict
 
     # Grid
     crs: str
@@ -68,7 +69,7 @@ def get_map(self, ds: xr.Dataset, query: dict) -> StreamingResponse:
         da = self.select_layer(ds)
         da = self.select_time(da)
         da = self.select_elevation(ds, da)
-        # da = self.select_custom_dim(da)
+        da = self.select_custom_dim(da)
 
         # Render the data using the render that matches the dataset type
         # The data selection and render are coupled because they are both driven by
@@ -101,6 +102,7 @@ def get_minmax(self, ds: xr.Dataset, query: dict) -> dict:
         da = self.select_layer(ds)
         da = self.select_time(da)
         da = self.select_elevation(ds, da)
+        da = self.select_custom_dim(da)
 
         # Prepare the data as if we are going to render it, but instead grab the min and max
         # values from the data to represent the range of values in the given area
@@ -162,6 +164,9 @@ def ensure_query_types(self, ds: xr.Dataset, query: dict):
         ]
         self.autoscale = query.get("autoscale", "false") == "true"
 
+        available_selectors = ds.gridded.additional_coords(ds[self.parameter])
+        self.dim_selectors = {k: query[k] for k in available_selectors}
+
     def select_layer(self, ds: xr.Dataset) -> xr.DataArray:
         """
         Select Dataset variable, according to WMS layers request
@@ -224,19 +229,27 @@ def select_custom_dim(self, da: xr.DataArray) -> xr.DataArray:
         :param da:
         :return:
         """
+        # Filter dimension from custom query, if any
+        for dim, value in self.dim_selectors.items():
+            if dim in da.coords:
+                dtype = da[dim].dtype
+                if np.issubdtype(dtype, np.integer):
+                    value = int(value)
+                elif np.issubdtype(dtype, np.floating):
+                    value = float(value)
+                da = da.sel({dim: value})
+
         # Squeeze single value dimensions
         da = da.squeeze()
 
-        # TODO: Filter dimension from custom query, if any
-
         # Squeeze multiple values dimensions, by selecting the last value
-        for key in da.cf.coordinates.keys():
-            if key in ("latitude", "longitude", "X", "Y"):
-                continue
+        # for key in da.cf.coordinates.keys():
+        #     if key in ("latitude", "longitude", "X", "Y"):
+        #         continue
 
-            coord = da.cf.coords[key]
-            if coord.size > 1:
-                da = da.cf.isel({key: -1})
+        #     coord = da.cf.coords[key]
+        #     if coord.size > 1:
+        #         da = da.cf.isel({key: -1})
 
         return da
 

xpublish_wms/wms/get_metadata.py

@@ -110,6 +110,11 @@ def get_layer_details(ds: xr.Dataset, layer_name: str) -> dict:
     else:
         timesteps = None
 
+    additional_coords = ds.gridded.additional_coords(da)
+    additional_coord_values = {
+        coord: da.cf.coords[coord].values.tolist() for coord in additional_coords
+    }
+
     return {
         "layerName": da.name,
         "standard_name": da.cf.attrs.get("standard_name", da.name),
@@ -121,6 +126,8 @@ def get_layer_details(ds: xr.Dataset, layer_name: str) -> dict:
         "elevation_positive": elevation_positive,
         "elevation_units": elevation_units,
         "timesteps": timesteps,
+        "additional_coords": additional_coords,
+        **additional_coord_values,
     }