ensure AoS in structure is long enough to fill IDS quantity

Author: SBlokhuizen

pyproject.toml

@@ -17,7 +17,7 @@ dependencies = [
     "click",
     "matplotlib",
     "scipy",
-    # "imas-python",
+    "imas-python[netcdf] @ git+https://github.com/iterorganization/IMAS-Python.git",
     "panel",
     "pygments",
     "plotly",

tests/test_waveform_exporter.py

@@ -1,6 +1,8 @@
 import csv
 from pathlib import Path
 
+# TODO: use netcdf and imas instead of imaspy
+import imaspy as imas
 import numpy as np
 import pytest
 
@@ -109,6 +111,95 @@ def test_to_png(exporter, tmp_path):
     assert Path(file_path).exists()
 
 
-# TODO: Write tests for exporting to IDS, this requires IMASPy as a dependency
-# def test_to_ids(exporter):
-#     pass
+def test_to_ids_flt1d(exporter, tmp_path):
+    ids = imas.IDSFactory().ec_launchers()
+    ids.time = [0]
+    ids.ids_properties.homogeneous_time = imas.ids_defs.IDS_TIME_MODE_HOMOGENEOUS
+    file_path = f"imas:hdf5?path={tmp_path}/test_ec_launchers"
+    with imas.DBEntry(file_path, "w") as dbentry:
+        dbentry.put(ids)
+
+    uri = f"{file_path}#ec_launchers/beam(1)/power_launched"
+    exporter.to_ids(uri)
+
+    with imas.DBEntry(file_path, "r") as dbentry:
+        ids = dbentry.get("ec_launchers", autoconvert=False)
+        assert np.all(ids.time == exporter.times)
+        assert np.all(ids.beam[0].power_launched.data == exporter.values)
+
+
+def test_to_ids_flt1d_heterogeneous(exporter, tmp_path):
+    ids = imas.IDSFactory().ec_launchers()
+    ids.time = [0]
+    ids.ids_properties.homogeneous_time = imas.ids_defs.IDS_TIME_MODE_HETEROGENEOUS
+    file_path = f"imas:hdf5?path={tmp_path}/test_ec_launchers"
+    with imas.DBEntry(file_path, "w") as dbentry:
+        dbentry.put(ids)
+
+    uri = f"{file_path}#ec_launchers/beam(1)/power_launched"
+    exporter.to_ids(uri)
+
+    with imas.DBEntry(file_path, "r") as dbentry:
+        ids = dbentry.get("ec_launchers", autoconvert=False)
+        assert np.all(ids.beam[0].power_launched.time == exporter.times)
+        assert np.all(ids.beam[0].power_launched.data == exporter.values)
+
+
+def test_to_ids_flt1d_occurrence(exporter, tmp_path):
+    ids = imas.IDSFactory().ec_launchers()
+    ids.time = [0]
+    ids.ids_properties.homogeneous_time = imas.ids_defs.IDS_TIME_MODE_HOMOGENEOUS
+    file_path = f"imas:hdf5?path={tmp_path}/test_ec_launchers"
+    with imas.DBEntry(file_path, "w") as dbentry:
+        dbentry.put(ids)
+        dbentry.put(ids, occurrence=1)
+
+    uri = f"{file_path}#ec_launchers:1/beam(1)/power_launched"
+    exporter.to_ids(uri)
+
+    with imas.DBEntry(file_path, "r") as dbentry:
+        # Check if only occurrence 1 is filled
+        ids = dbentry.get("ec_launchers", occurrence=0, autoconvert=False)
+        assert np.all(ids.time == [0])
+        assert not ids.beam
+
+        ids = dbentry.get("ec_launchers", occurrence=1, autoconvert=False)
+        assert np.all(ids.time == exporter.times)
+        assert np.all(ids.beam[0].power_launched.data == exporter.values)
+
+
+def test_to_ids_flt0d(exporter, tmp_path):
+    ids = imas.IDSFactory().equilibrium()
+    ids.time = [0]
+    ids.ids_properties.homogeneous_time = imas.ids_defs.IDS_TIME_MODE_HOMOGENEOUS
+    file_path = f"imas:hdf5?path={tmp_path}/test_equilibrium"
+    with imas.DBEntry(file_path, "w") as dbentry:
+        dbentry.put(ids)
+
+    uri = f"{file_path}#equilibrium/time_slice()/boundary/elongation"
+    exporter.to_ids(uri)
+
+    with imas.DBEntry(file_path, "r") as dbentry:
+        ids = dbentry.get("equilibrium", autoconvert=False)
+        assert np.all(ids.time == exporter.times)
+        for i, time_slice in enumerate(ids.time_slice):
+            assert time_slice.boundary.elongation == exporter.values[i]
+
+
+def test_to_ids_flt0d_heterogeneous(exporter, tmp_path):
+    ids = imas.IDSFactory().equilibrium()
+    ids.time = [0]
+    ids.ids_properties.homogeneous_time = imas.ids_defs.IDS_TIME_MODE_HETEROGENEOUS
+    file_path = f"imas:hdf5?path={tmp_path}/test_equilibrium"
+    with imas.DBEntry(file_path, "w") as dbentry:
+        dbentry.put(ids)
+
+    uri = f"{file_path}#equilibrium/time_slice()/boundary/elongation"
+    exporter.to_ids(uri)
+
+    with imas.DBEntry(file_path, "r") as dbentry:
+        ids = dbentry.get("equilibrium", autoconvert=False)
+        imas.util.print_tree(ids)
+        for i, time_slice in enumerate(ids.time_slice):
+            assert time_slice.boundary.elongation == exporter.values[i]
+            assert time_slice.time == exporter.times[i]

waveform_editor/waveform_exporter.py

@@ -1,7 +1,8 @@
 import csv
 import logging
+import re
 
-import imaspy
+import imas
 import plotly.graph_objects as go
 
 logger = logging.getLogger(__name__)
@@ -53,17 +54,17 @@ def to_ids(self, uri, dd_version=None):
                 default version will be used.
         """
         uri_entry, uri_ids, occurrence, path = self.parse_uri(uri)
-        entry = imaspy.DBEntry(uri_entry, "r", dd_version=dd_version)
+        entry = imas.DBEntry(uri_entry, "r", dd_version=dd_version)
         ids = entry.get(uri_ids, occurrence, autoconvert=False)
 
         if (
             ids.ids_properties.homogeneous_time
-            == imaspy.ids_defs.IDS_TIME_MODE_HETEROGENEOUS
+            == imas.ids_defs.IDS_TIME_MODE_HETEROGENEOUS
         ):
             is_homogeneous = False
         elif (
             ids.ids_properties.homogeneous_time
-            == imaspy.ids_defs.IDS_TIME_MODE_HOMOGENEOUS
+            == imas.ids_defs.IDS_TIME_MODE_HOMOGENEOUS
         ):
             is_homogeneous = True
             ids.time = self.times
@@ -72,12 +73,14 @@ def to_ids(self, uri, dd_version=None):
                 "The time mode must be homogeneous or heterogeneous."
             )
 
+        self._ensure_path_exists(ids, path)
+
         if "()" in path:
             self._fill_flt_0d(ids, path, is_homogeneous)
         else:
             self._fill_flt_1d(ids, path, is_homogeneous)
 
-        entry.put(ids)
+        entry.put(ids, occurrence)
         entry.close()
 
     def to_csv(self, file_path):
@@ -132,7 +135,6 @@ def _fill_flt_0d(self, ids, path, is_homogeneous):
         aos_path = aos_path.strip("/")
         remaining_path = remaining_path.strip("/")
         aos = ids[aos_path]
-        aos.resize(len(self.times))
 
         for i, time in enumerate(self.times):
             if aos[i][remaining_path].data_type == "FLT_0D":
@@ -158,3 +160,48 @@ def _fill_flt_1d(self, ids, path, is_homogeneous):
                 quantity.time = self.times
         else:
             raise NotImplementedError("Invalid data")
+
+    def _ensure_path_exists(self, ids, path):
+        """
+        Traverses a given path and modifies the AoS in the IDS to ensure the IDS
+        quantity to be filled exists.
+
+        Examples:
+
+        - imas:hdf5?path=./testdb#ec_launchers/beam(123)/power_launched
+          This will ensure ec_launchers.beam has a length of at least 123. Note, 1-based
+          indexing is used in the URI.
+
+        - imas:hdf5?path=./testdb#equilibrium/time_slice()/boundary/elongation
+          When '()' is encountered, it is assumed that this AoS should be the length of
+          the exported time array, i.e. len(equilibrium.time_slice) == len(self.times)
+
+        Args:
+            ids: The IDS to export to.
+            path: The path of the IDS quantity to export to.
+
+        Returns:
+            None. Modifies the `ids` structure in place.
+        """
+        path = path.strip("/")
+        path_parts = path.split("/")
+        current = ids
+        for part in path_parts:
+            if "()" in part:
+                current = current[part.split("(")[0]]
+                current.resize(len(self.times))
+                current = current[0]
+
+            elif "(" in part:
+                match = re.search(r"\((\d+)\)", part)
+                index = int(match.group(1))
+                current = current[part.split("(")[0]]
+
+                # We use 1-based indexing in the URI
+                if len(current) < index:
+                    current.resize(index)
+
+                # Revert to 0-based indexing
+                current = current[index - 1]
+            else:
+                current = current[part]