Feature/replace dim x with identifier (#86)

Author: jwasikpsnc

backend/ibex/data_source/imas_python_source.py

@@ -678,6 +678,43 @@ def _leaf_node_coordinates_contain_time(self, leaf_node_path: str, coordinates_t
 
         return False
 
+    def _generate_grid_quantity_alias(self, grid_node: IDSNumericArray):
+        """
+        Generates alias and unit for selected grid node. Assumes grid_node.name == "dimX" X=(1...N)
+        :param grid_node: IDSNode (named dimX, X = [1...N])
+        :return:
+        """
+        result = {"axis_label": None, "unit": None}
+        if grid_node._parent is None or grid_node._parent._parent is None:
+            return result
+        if not re.search(r"dim[1-9]", grid_node.metadata.name):
+            return result
+
+        # assume grid_node is located inside XXX/grid/<node> and grid_type is located in XXX/grid_type
+        grid_type_index = grid_node._parent._parent.grid_type.index
+        if grid_type_index == imas.ids_defs.EMPTY_INT:
+            return result
+
+        dim_index = int(grid_node.metadata.name[-1]) - 1  # dim1->0, dim2->1 etc...
+        # Extract units
+        try:
+            units = imas.identifiers.poloidal_plane_coordinates_identifier(grid_type_index).units.split(",")
+            result["unit"] = units[dim_index]
+        except (ValueError, KeyError, AttributeError):
+            result["unit"] = None
+
+        # Extract axis labels
+        try:
+            axis_labels = imas.identifiers.poloidal_plane_coordinates_identifier(grid_type_index).axis_labels.split(",")
+            result["axis_label"] = axis_labels[dim_index]
+        except AttributeError:
+            description = imas.identifiers.poloidal_plane_coordinates_identifier(grid_type_index).description
+            match = re.findall(r"(\w+)=(dim[1-9])", description)
+            axis_labels = {v: k for k, v in match}
+            result["axis_label"] = axis_labels.get(grid_node.metadata.name, None)
+
+        return result
+
     def get_geometry_overlay_nodes(
         self,
         uri: str,
@@ -982,10 +1019,20 @@ def get_plot_data(self, plot_data_query: PlotDataRequestModel) -> dict:
                         except ValueError:
                             coord_data_shape = "irregular"
 
+                        coord_name = coord.split("/")[-1]
+                        axis_label = None
+                        unit = None
+                        if re.search(r"dim[1-9]", coord_name):
+                            labels_dict = self._generate_grid_quantity_alias(first_value)
+                            axis_label = labels_dict["axis_label"]
+                            unit = labels_dict["unit"]
+
+                        coord_name = axis_label if axis_label else coord_name
+                        units = unit if unit else first_value.metadata.units
                         c = {
-                            "name": coord.split("/")[-1],
+                            "name": coord_name,
                             "target": f"#{ids}/{target}",
-                            "unit": first_value.metadata.units,
+                            "unit": units,
                             "shape": coord_data_shape,  # coord_data could be np.ndarray or list[np.ndarray]
                             "downsampled_shape": coord_data_shape,
                             "ndim": first_value.metadata.ndim,

backend/ibex/endpoints/schemas/response_data_schemas.py

@@ -20,7 +20,7 @@ class PlotDataCoordinateModel(BaseModel):
     target: str = Field(
         description="Which node coordinate is it", examples=["#equilibrium/time_slice[0]/profiles_2d[0]/psi"]
     )
-    unit: str = Field(description="Data units", examples=[""])
+    unit: str = Field(description="Data units", examples=["m", "mixed"])
     shape: list[int] | str = Field(description="Shape of the data", examples=[[129]])
     downsampled_shape: list[int] | str = Field(description="Shape of the data after downsampling", examples=[[129]])
     ndim: int = Field(description="Number of data dimensions stored in node", examples=[1])

backend/tests/conftest.py

@@ -90,11 +90,16 @@ def entry_path(tmp_path_factory):
             )
         profiles_2d.grid.dim1 = np.array([0, 1, 2], dtype=float)
         profiles_2d.grid.dim2 = np.array([0, 1, 2], dtype=float)
+        profiles_2d.grid.volume_element = np.array([[1.0, 2.0, 3.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], dtype=float)
         i += 10
 
     # ===== for data smoothing (must be time-based) =====
     core_profiles.global_quantities.ip = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
 
+    # for coordinate aliases/units
+    core_profiles.profiles_2d[0].grid_type = 1
+    core_profiles.profiles_2d[1].grid_type = 2
+
     entry.put(core_profiles)
 
     # ===== for geometry overlay =====

backend/tests/test_data_endpoints.py

@@ -1,5 +1,9 @@
 import pytest
 import numpy as np
+from packaging.version import Version
+import imas
+
+_IMAS_GE_2_3 = Version(imas.__version__) >= Version("2.3.0")
 
 
 def test_status_codes(entry_path):
@@ -123,7 +127,8 @@ def test_plot_data_smoothing_with_wrong_target_node(entry_path):
     assert response.status_code == 466
 
 
-def test_plot_data_2d(entry_path):
+@pytest.mark.parametrize("expected_unit", ["m"] if _IMAS_GE_2_3 else ["mixed"])
+def test_plot_data_2d(entry_path, expected_unit):
     parameters = {
         "uri": f"imas:hdf5?path={entry_path}#core_profiles/profiles_2d[:]/ion[:]/temperature",
     }
@@ -147,13 +152,15 @@ def test_plot_data_2d(entry_path):
     assert time_coordinate["description"] == "Generic time"
 
     dim1_coordinate = response_body["data"]["coordinates"][0]
-    assert dim1_coordinate["name"] == "dim1"
+    assert dim1_coordinate["name"] == "R"  # alias for dim1
+    assert dim1_coordinate["unit"] == expected_unit
     assert dim1_coordinate["target"] == "#core_profiles/profiles_2d[:]/ion[:]/temperature"
     assert dim1_coordinate["shape"] == [5, 3]
     assert dim1_coordinate["path"] == "#core_profiles/profiles_2d[:]/grid/dim1"
 
     dim2_coordinate = response_body["data"]["coordinates"][1]
-    assert dim2_coordinate["name"] == "dim2"
+    assert dim2_coordinate["name"] == "Z"  # alias for dim2
+    assert dim2_coordinate["unit"] == expected_unit
     assert dim2_coordinate["target"] == "#core_profiles/profiles_2d[:]/ion[:]/temperature"
     assert dim2_coordinate["shape"] == [5, 3]
     assert dim2_coordinate["path"] == "#core_profiles/profiles_2d[:]/grid/dim2"
@@ -205,3 +212,35 @@ def test_plot_data_requires_savgol_window_length_and_polyorder(entry_path):
     )
     assert response.status_code == 422
     assert "savgol_smoothing_polyorder is required" in response.text
+
+
+@pytest.mark.parametrize("expected_unit", [("m", "rad")] if _IMAS_GE_2_3 else [("mixed", "mixed")])
+def test_plot_data_coordinate_aliases(entry_path, expected_unit):
+
+    parameters = {
+        "uri": f"imas:hdf5?path={entry_path}#core_profiles/profiles_2d[0]/grid/volume_element",
+    }
+    response = pytest.test_client.get("/data/plot_data", params=parameters)
+    response_body = response.json()
+
+    assert response.status_code == 200
+
+    dim1_coordinate = response_body["data"]["coordinates"][0]
+    assert dim1_coordinate["name"].lower() == "r"
+    assert dim1_coordinate["unit"].lower() == expected_unit[0]
+
+    parameters = {
+        "uri": f"imas:hdf5?path={entry_path}#core_profiles/profiles_2d[1]/grid/volume_element",
+    }
+    response = pytest.test_client.get("/data/plot_data", params=parameters)
+    response_body = response.json()
+
+    assert response.status_code == 200
+
+    dim1_coordinate = response_body["data"]["coordinates"][0]
+    dim2_coordinate = response_body["data"]["coordinates"][1]
+    assert dim1_coordinate["name"].lower() == "rho"
+    assert dim1_coordinate["unit"].lower() == expected_unit[0]
+
+    assert dim2_coordinate["name"].lower() == "theta"
+    assert dim2_coordinate["unit"].lower() == expected_unit[1]