Feature/geometry overlay

backend/ibex/core/ibex_service.py

@@ -119,3 +119,10 @@ def get_multiple_node_data(uri: str) -> dict:
 
 def get_plot_data(plot_data_query: PlotDataRequestModel) -> dict:
     return data_source.get_plot_data(plot_data_query)
+
+
+def get_geometry_overlay_nodes(uri: str, show_empty_nodes: bool = False, show_error_bars: bool = False) -> dict:
+    uri_obj = IMAS_URI(uri)
+    return data_source.get_geometry_overlay_nodes(
+        uri=uri_obj.uri_entry_identifiers, show_empty_nodes=show_empty_nodes, show_error_bars=show_error_bars
+    )

backend/ibex/data_source/data_source_interface.py

@@ -133,3 +133,25 @@ def get_plot_data(self, plot_data_query: PlotDataRequestModel) -> dict:
         :return: Dictionary containing data values, metadata and coordinates.
         """
         ...
+
+    @abstractmethod
+    def get_geometry_overlay_nodes(
+        self,
+        uri: str,
+        show_empty_nodes: bool = False,
+        show_error_bars: bool = False,
+    ) -> dict:
+        """
+        Returns paths to metadata nodes that describe geometry overlays.
+
+        A node is included when:
+        - its type is ``outline_2d_geometry_static``, or
+        - its name contains ``outline`` and its type is ``rz1d_static`` or ``rz1d_dynamic_aos``.
+        Error bar nodes are filtered out by default and can be included with ``show_error_bars=True``.
+
+        :param uri: imas URI
+        :param show_empty_nodes: whether empty nodes should be returned, or not
+        :param show_error_bars: whether error bar nodes should be returned, or not
+        :return: dictionary {'outline_nodes': [{'geometry_node': '...', 'parameters': [...]}, ...]}
+        """
+        ...

backend/ibex/data_source/imas_python_source.py

@@ -165,18 +165,42 @@ def _jsonify_metadata(self, metadata: IDSMetadata, recursive: bool = False, show
         result["type"] = metadata.data_type or "structure"
         result["ndim"] = metadata.ndim
         result["shape"] = []  # empty for 0D data
+        result["is_geometry_node"] = self._is_geometry_node(metadata)
 
         if recursive:
             result["children"] = [self._jsonify_metadata(child, recursive) for child in metadata]
         else:
             result["children"] = [
-                {"name": child.name, "type": child.data_type, "ndim": child.ndim}
+                {
+                    "name": child.name,
+                    "type": child.data_type,
+                    "ndim": child.ndim,
+                    "is_geometry_node": self._is_geometry_node(child),
+                }
                 for child in metadata
                 if show_error_bars or not any(x in child.name for x in ["_error_upper", "_error_lower", "_error_index"])
             ]
 
         return result
 
+    def _is_geometry_node(self, metadata: IDSMetadata):
+        """
+        Checks if node lies inside geometry structure
+        :param metadata: metadata of ids node
+        :return:
+        """
+        if metadata is None:
+            return False
+
+        node_type = getattr(metadata, "structure_reference", None)
+        is_outline_static = node_type == "outline_2d_geometry_static"
+        is_outline_rz = "outline" in metadata.name and node_type in {"rz1d_static", "rz1d_dynamic_aos"}
+
+        if is_outline_rz or is_outline_static:
+            return True
+
+        return self._is_geometry_node(metadata._parent)
+
     def get_node_info(
         self,
         uri: str,
@@ -517,12 +541,17 @@ def find_paths(self, uri: str, searched_node: str, show_error_bars: bool = False
                         node_data_type = ids_obj.metadata[path].data_type
                         if node_data_type.value != "structure" and node_data_type.value != "struct_array":
                             path_name = f"#{ids}/{self._add_index_to_aos_in_path(ids_obj.metadata, path)}"
+                            is_geometry_node = self._is_geometry_node(ids_obj.metadata[path])
                             if not filled_paths:
                                 # every ids has at least one filled path. If not, it means functionality is not available.
-                                found_paths.append({"path": path_name, "has_data": None})
+                                found_paths.append(
+                                    {"path": path_name, "has_data": None, "is_geometry_node": is_geometry_node}
+                                )
                             else:
                                 path_has_data = path_in_filled_paths(path, filled_paths)
-                                found_paths.append({"path": path_name, "has_data": path_has_data})
+                                found_paths.append(
+                                    {"path": path_name, "has_data": path_has_data, "is_geometry_node": is_geometry_node}
+                                )
 
                 except imas.exception.DataEntryException:
                     continue
@@ -649,6 +678,147 @@ def _leaf_node_coordinates_contain_time(self, leaf_node_path: str, coordinates_t
 
         return False
 
+    def get_geometry_overlay_nodes(
+        self,
+        uri: str,
+        show_empty_nodes: bool = False,
+        show_error_bars: bool = False,
+    ) -> dict:
+        """
+        Returns paths to metadata nodes that describe geometry overlays.
+
+        A node is included when:
+        - its type is ``outline_2d_geometry_static``, or
+        - its name contains ``outline`` and its type is ``rz1d_static`` or ``rz1d_dynamic_aos``.
+        Error bar nodes are filtered out by default and can be included with ``show_error_bars=True``.
+
+        :param uri: imas URI
+        :param show_empty_nodes: whether empty nodes should be returned, or not
+        :param show_error_bars: whether error bar nodes should be returned, or not
+        :return: dictionary {'outline_nodes': [{'geometry_node': '...', 'parameters': [...]}, ...]}
+        """
+
+        # ============ HELPER FUNCTION ============
+        def _get_descendant_node_names(metadata: IDSMetadata):
+
+            res = []
+            if metadata.data_type == IDSDataType.STRUCTURE:
+                for child in metadata:
+                    res.extend([f"{metadata.name}/{x}" for x in _get_descendant_node_names(child)])
+            else:
+                res.append(f"{metadata.name}")
+            return res
+
+        def _walk_outline_nodes(
+            uri: str,
+            ids: str,
+            occurrence: int,
+            root_metadata: IDSMetadata,
+            metadata: IDSMetadata,
+            results: list[dict[str, list[str]]],
+            show_error_bars: bool = False,
+            filled_paths: list[str] | None = None,
+        ) -> None:
+            """
+            Recursively traverses IDS metadata tree and collects nodes describing geometry overlays.
+
+            A node is collected when its type is ``outline_2d_geometry_static``
+            or when its name contains ``outline`` and its type is ``rz1d_static`` or ``rz1d_dynamic_aos``.
+
+            :param uri: imas URI
+            :param ids: name of IDS (e.g. core_profiles)
+            :param occurrence: IDS occurrence number
+            :param root_metadata: root metadata of the IDS (used to resolve tensorized paths)
+            :param metadata: current metadata node to inspect
+            :param results: list to which collected geometry overlay entries are appended
+            :param show_error_bars: whether to include error bar parameter names (e.g. ``_error_upper``)
+            :param filled_paths: optional list of filled paths; when given, only nodes with filled parameters are collected
+            """
+            node_name = metadata.name
+            node_type = getattr(metadata, "structure_reference", None)
+
+            is_outline_static = node_type == "outline_2d_geometry_static"
+            is_outline_rz = "outline" in node_name and node_type in {"rz1d_static", "rz1d_dynamic_aos"}
+
+            if is_outline_static or is_outline_rz:
+                tensorized_path = self._add_index_to_aos_in_path(root_metadata, metadata.path_string)
+                full_uri_with_path = f"{uri}#{ids}:{occurrence}/{tensorized_path}"
+                parameters_entry = {"geometry_node": full_uri_with_path, "parameters": []}
+
+                params = []
+                for child in metadata:
+                    params.extend(_get_descendant_node_names(child))
+
+                for param in params:
+                    is_error_node = any(
+                        error_node in param for error_node in ["_error_upper", "_error_lower", "_error_index"]
+                    )
+                    if show_error_bars or not is_error_node:
+                        parameters_entry["parameters"].append(param)
+
+                if filled_paths is not None:
+                    node_filled = any(
+                        f"{metadata.path_string}/{parameter}" in filled_paths
+                        for parameter in parameters_entry["parameters"]
+                    )
+
+                    if node_filled and parameters_entry["parameters"]:  # don't put structures with empty "parameters"
+                        results.append(parameters_entry)
+                elif parameters_entry["parameters"]:  # don't put structures with empty "parameters"
+                    results.append(parameters_entry)
+
+            else:
+                for child in metadata:
+                    _walk_outline_nodes(
+                        uri=uri,
+                        ids=ids,
+                        occurrence=occurrence,
+                        root_metadata=root_metadata,
+                        metadata=child,
+                        results=results,
+                        show_error_bars=show_error_bars,
+                        filled_paths=filled_paths,
+                    )
+
+        # ============ END HELPER FUNCTION ============
+
+        # Iterate over all filled IDSes and their occurrences to collect geometry overlay nodes
+        filled_idses = self.list_idses(uri)["idses"]
+
+        with self._open_entry(uri) as entry:
+            result = []
+
+            for ids_dict in filled_idses:
+                # ids_dict = {'name': < name >, 'occurrences': [ < 0 >, < 1 >, ...]}
+                for occurrence in ids_dict["occurrences"]:
+                    ids_obj = self._get_ids_from_entry(entry, ids_dict["name"], occurrence)
+                    outline_nodes = []
+
+                    filled_paths = None
+                    if not show_empty_nodes:
+                        try:
+                            filled_paths = entry.list_filled_paths(ids_dict["name"], int(occurrence))
+                        except (AttributeError, imas.backends.imas_core.imas_interface.LLInterfaceError):
+                            # AttributeError - current version of IMAS-Python doesn't support list_filled paths
+                            # LLInterfaceError - current version of IMAS-Core doesn't support list_filled paths
+                            # proceed without filtering empty nodes
+                            ...
+
+                    _walk_outline_nodes(
+                        uri=uri,
+                        ids=ids_dict["name"],
+                        occurrence=occurrence,
+                        root_metadata=ids_obj.metadata,
+                        metadata=ids_obj.metadata,
+                        results=outline_nodes,
+                        show_error_bars=show_error_bars,
+                        filled_paths=filled_paths,
+                    )
+
+                    result.extend(outline_nodes)
+
+        return {"outline_nodes": result}
+
     def get_plot_data(self, plot_data_query: PlotDataRequestModel) -> dict:
         """
         Returns all data used to plot selected quantity. Result contains data values, metadata and coordinates.

backend/ibex/endpoints/ids_info.py

@@ -3,7 +3,12 @@
 from fastapi import APIRouter  # type: ignore
 
 from ibex.core import ibex_service
-from ibex.endpoints.schemas.response_ids_info_schemas import NodeInfoResponse, FindPathsResponse, ArraySummaryResponse
+from ibex.endpoints.schemas.response_ids_info_schemas import (
+    NodeInfoResponse,
+    FindPathsResponse,
+    ArraySummaryResponse,
+    GeometryOverlayNodesResponse,
+)
 
 router = APIRouter()
 
@@ -108,3 +113,34 @@ def array_summary(uri: str) -> dict:
 
     """
     return ibex_service.array_summary(uri.strip())
+
+
+@router.get(
+    "/ids_info/geometry_overlay_nodes",
+    status_code=200,
+    response_model=GeometryOverlayNodesResponse,
+    responses={
+        200: {"description": "Geometry overlay nodes returned successfully"},
+    },
+    description="Returns geometry overlay nodes metadata",
+)
+@ibex_service.measure_execution_time
+def geometry_overlay_nodes(uri: str, show_empty_nodes: bool = False, show_error_bars: bool = False) -> dict:
+    """
+    IBEX endpoint. Returns paths to geometry overlay nodes.
+
+    | Response JSON is constructed as follows:
+    | {
+    | "outline_nodes" : [<node1>, <node2>, ...]
+    | }
+
+    :param uri: IMAS URI
+    :param show_empty_nodes: switch used to hide empty nodes
+    :param show_error_bars: switch used to hide _error* nodes
+    :rtype: dict (automatically converted to JSON by FastAPI)
+    :return: JSON response
+
+    """
+    return ibex_service.get_geometry_overlay_nodes(
+        uri.strip(), show_empty_nodes=show_empty_nodes, show_error_bars=show_error_bars
+    )

backend/ibex/endpoints/schemas/response_ids_info_schemas.py

@@ -10,6 +10,7 @@ class NodeInfoChildModel(BaseModel):
     name: str = Field(description="Node name", examples=["t_i_average", "psi"])
     type: str = Field(description="Node type", examples=["FLT", "STR"])
     ndim: int = Field(description="Number of data dimensions stored in node", examples=[1, 5, 7])
+    is_geometry_node: bool = Field(description="True if node is inside geometry structure", examples=[True, False])
     has_data: Optional[bool] = Field(default=None, description="True if node contains data", examples=[True, False])
 
 
@@ -20,6 +21,7 @@ class NodeInfoResponse(BaseModel):
     type: str = Field(description="Node type", examples=["struct_array", "FLT", "STR"])
     ndim: int = Field(description="Number of data dimensions stored in node", examples=[1, 5, 7])
     shape: list[int] = Field(description="Shape of the data", examples=[[3], [5], [2, 5, 10]])
+    is_geometry_node: bool = Field(description="True if node is inside geometry structure", examples=[True, False])
     has_data: Optional[bool] = Field(default=None, description="True if node contains data", examples=[True, False])
     children: list[NodeInfoChildModel] = Field(description="Node info about node's children")
     coordinates: list[str] = Field(
@@ -35,6 +37,7 @@ class FoundPathModel(BaseModel):
 
     path: str = Field(description="", examples=["t_i_average", "path"])
     has_data: Optional[bool] = Field(default=None, description="True if node contains data", examples=[True, False])
+    is_geometry_node: bool = Field(description="True if node contains geometry data", examples=[True, False])
 
 
 class FindPathsResponse(BaseModel):
@@ -56,3 +59,24 @@ class ArraySummaryResponse(BaseModel):
     max: float = Field(description="Maximum value from the array", examples=[1.2, 3.4])
     mean: float = Field(description="Mean value from the array", examples=[1.2, 3.4])
     standard_deviation: float = Field(description="Standard deviation value of the array", examples=[1.2, 3.4])
+
+
+# ========== GEOMETRY OVERLAY NODES ==========
+
+
+class GeometryOverlayNodesSingleEntry(BaseModel):
+    """Response for /ids_info/geometry_overlay_nodes endpoint"""
+
+    geometry_node: str = Field(
+        description="Full uri pointing to a specific geometry node structure",
+        examples=["imas:hdf5?path=<entry_path>#equilibrium/time_slic[:]/profiles_2d[:]/psi"],
+    )
+    parameters: list[str] = Field(description="Name of child node of geometry_node", examples=["r", "z", "width"])
+
+
+class GeometryOverlayNodesResponse(BaseModel):
+    """Response for /ids_info/geometry_overlay_nodes endpoint"""
+
+    outline_nodes: list[GeometryOverlayNodesSingleEntry] = Field(
+        description="List of dicts describing geometry overlay node paths",
+    )

backend/pyproject.toml

@@ -66,6 +66,7 @@ test = [
     "pytest",
     "pytest-cov",
     "pytest-xdist",
+    "pytest-unordered",
 ]
 
 linting = [

backend/tests/conftest.py

@@ -97,6 +97,7 @@ def entry_path(tmp_path_factory):
 
     entry.put(core_profiles)
 
+    # ===== for geometry overlay =====
     # ===== for data smoothing 2D (one of coordinates is time) =====
 
     wall = entry.factory.wall()
@@ -110,7 +111,20 @@ def entry_path(tmp_path_factory):
             [1, 1, 1, 1, 1],
         ]
     )
+
+    wall.description_2d.resize(1)
+    wall.description_2d[0].limiter.unit.resize(1)
+    wall.description_2d[0].limiter.unit[0].outline.r = [1.0]
+    wall.description_2d[0].limiter.unit[0].outline.z = [1.0]
     entry.put(wall)
 
+    equilibrium = entry.factory.equilibrium()
+    equilibrium.ids_properties.homogeneous_time = 1
+    equilibrium.time = np.array([1.0], dtype=float)
+    equilibrium.time_slice.resize(1)
+    equilibrium.time_slice[0].boundary.outline.r = np.array([1.0], dtype=float)
+    equilibrium.time_slice[0].boundary.outline.z = np.array([2.0], dtype=float)
+    entry.put(equilibrium)
+
     entry.close()
     return tmp_path