Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 49 additions & 2 deletions backend/ibex/data_source/imas_python_source.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion backend/ibex/endpoints/schemas/response_data_schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down
5 changes: 5 additions & 0 deletions backend/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 =====
Expand Down
45 changes: 42 additions & 3 deletions backend/tests/test_data_endpoints.py
Original file line number Diff line number Diff line change
@@ -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):
Expand Down Expand Up @@ -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",
}
Expand All @@ -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"
Expand Down Expand Up @@ -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]
Loading