diff --git a/backend/ibex/data_source/imas_python_source.py b/backend/ibex/data_source/imas_python_source.py index 91d57203..61110056 100644 --- a/backend/ibex/data_source/imas_python_source.py +++ b/backend/ibex/data_source/imas_python_source.py @@ -628,6 +628,43 @@ def _check_data_is_leaf_node(self, data) -> None: elif isinstance(data, IDSStructure): raise NotALeafNodeException("Cannot serialize non-leaf node") + 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 = {"alias": None, "unit_alias": 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/ 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_alias"] = units[dim_index] + except (ValueError, KeyError, AttributeError): + result["unit_alias"] = None + + # Extract axis labels + try: + axis_labels = imas.identifiers.poloidal_plane_coordinates_identifier(grid_type_index).axis_labels.split(",") + result["alias"] = 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["alias"] = axis_labels.get(grid_node.metadata.name, None) + + return result + def get_plot_data( self, uri: str, @@ -803,10 +840,20 @@ def get_plot_data( except ValueError: coord_data_shape = "irregular" + coord_name = coord.split("/")[-1] + alias = None + alias_unit = None + if re.search(r"dim[1-9]", coord_name): + alias_dict = self._generate_grid_quantity_alias(first_value) + alias = alias_dict["alias"] + alias_unit = alias_dict["unit_alias"] + + coord_name = alias if alias else coord_name + units = alias_unit if alias_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, diff --git a/backend/ibex/endpoints/schemas/data_schemas.py b/backend/ibex/endpoints/schemas/data_schemas.py index 68052c77..1fbb7e3a 100644 --- a/backend/ibex/endpoints/schemas/data_schemas.py +++ b/backend/ibex/endpoints/schemas/data_schemas.py @@ -17,10 +17,12 @@ class PlotDataCoordinateModel(BaseModel): """Intermediate model for /data/plot_data endpoint""" name: str = Field(description="Node name", examples=["dim1"]) + alias: str = Field(description="Node name alias", examples=["Height"]) 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"]) + unit_alias: str = Field(description="Data units alias", examples=["m"]) 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]) diff --git a/backend/tests/conftest.py b/backend/tests/conftest.py index 220783e1..a7d06e3a 100644 --- a/backend/tests/conftest.py +++ b/backend/tests/conftest.py @@ -88,8 +88,13 @@ def entry_path(tmp_path_factory): ion.temperature = np.array([[i, +1, i + 2], [i + 10, i + 11, i + 12], [i + 20, i + 21, i + 32]]) profiles_2d.grid.dim1 = np.array([0, 1, 2]) profiles_2d.grid.dim2 = np.array([0, 1, 2]) + 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 coordinate aliases/units + core_profiles.profiles_2d[0].grid_type = 1 + core_profiles.profiles_2d[1].grid_type = 2 + entry.put(core_profiles) entry.close() diff --git a/backend/tests/test_data_endpoints.py b/backend/tests/test_data_endpoints.py index 5ddd8f51..723c5d0f 100644 --- a/backend/tests/test_data_endpoints.py +++ b/backend/tests/test_data_endpoints.py @@ -93,13 +93,13 @@ 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["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["target"] == "#core_profiles/profiles_2d[:]/ion[:]/temperature" assert dim2_coordinate["shape"] == [5, 3] assert dim2_coordinate["path"] == "#core_profiles/profiles_2d[:]/grid/dim2" @@ -122,3 +122,28 @@ def test_plot_data_1_N_coord(entry_path): assert numeric_coordinate["ndim"] == 1 assert numeric_coordinate["path"] == "" assert numeric_coordinate["description"] == "1...N" + + +def test_plot_data_coordinate_aliases(entry_path): + + 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" + + 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] + assert dim1_coordinate["name"].lower() == "rho"