CLI for exporting waveforms

pyproject.toml

@@ -17,11 +17,12 @@ dependencies = [
     "click",
     "matplotlib",
     "scipy",
-    # "imas-python",
+    "imas-python[netcdf] @ git+https://github.com/iterorganization/IMAS-Python.git",
     "panel",
     "pygments",
     "plotly",
     "holoviews",
+    "kaleido", # required for exporting PNG to disk
 ]
 dynamic = ["version"]
 
@@ -54,6 +55,9 @@ homepage = "https://github.com/iterorganization/Waveform-Editor"
 where = ["."]
 include = ["waveform_editor*"]
 
+[project.scripts]
+waveform-editor = "waveform_editor.cli:cli"
+
 [tool.setuptools_scm]
 # version_file = "waveform_editor/_version.py"
 

tests/test_waveform.py

@@ -273,3 +273,9 @@ def test_overlap_derivatives():
     expected = [2, 2, -1.5, -1.5, -1.5, -1.5, -1.5]
     values = waveform.get_derivative(np.linspace(0, 3, 7))
     assert np.allclose(values, expected)
+
+
+def test_get_start_end(waveform):
+    """Test if waveform returns correct start and end value."""
+    assert waveform.get_start() == 0
+    assert waveform.get_end() == 14

tests/test_waveform_exporter.py

@@ -0,0 +1,217 @@
+import csv
+from pathlib import Path
+
+import imas
+import numpy as np
+import pytest
+
+from waveform_editor.waveform import Waveform
+from waveform_editor.waveform_exporter import WaveformExporter
+
+
+@pytest.fixture
+def waveform_list():
+    return [
+        {
+            "user_type": "linear",
+            "user_from": 0,
+            "user_to": 8,
+            "user_duration": 5,
+            "line_number": 1,
+        },
+        {
+            "user_type": "sine-wave",
+            "user_base": 8,
+            "user_amplitude": 2,
+            "user_frequency": 1,
+            "user_duration": 4,
+            "line_number": 2,
+        },
+        {
+            "user_type": "constant",
+            "user_value": 8,
+            "user_duration": 3,
+            "line_number": 3,
+        },
+        {
+            "user_type": "smooth",
+            "user_from": 8,
+            "user_to": 0,
+            "user_duration": 2,
+            "line_number": 4,
+        },
+    ]
+
+
+@pytest.fixture
+def waveform(waveform_list):
+    return Waveform(waveform=waveform_list)
+
+
+@pytest.fixture
+def times():
+    return np.array([0, 1, 2, 3, 4, 5])
+
+
+@pytest.fixture
+def exporter(waveform, times):
+    return WaveformExporter(waveform, times=times)
+
+
+def test_parse_uri(exporter):
+    """Test if URIs are parsed correctly."""
+    uri = "imas:hdf5?path=./testdb#ec_launchers/beam(1)/power_launched"
+    parsed_uri = exporter.parse_uri(uri)
+    assert parsed_uri == (
+        "imas:hdf5?path=./testdb",
+        "ec_launchers",
+        0,
+        "/beam(1)/power_launched",
+    )
+    uri2 = "imas:hdf5?path=./testdb/test_loc#ec_launchers:1/beam(1)/power_launched"
+    parsed_uri = exporter.parse_uri(uri2)
+    assert parsed_uri == (
+        "imas:hdf5?path=./testdb/test_loc",
+        "ec_launchers",
+        1,
+        "/beam(1)/power_launched",
+    )
+
+    uri3 = "imas:hdf5?path=./test_db#equilibrium/time_slice()/boundary/elongation"
+    parsed_uri = exporter.parse_uri(uri3)
+    assert parsed_uri == (
+        "imas:hdf5?path=./test_db",
+        "equilibrium",
+        0,
+        "/time_slice()/boundary/elongation",
+    )
+
+    uri4 = "test_db.nc#equilibrium/time_slice()/boundary/elongation"
+    parsed_uri = exporter.parse_uri(uri4)
+    assert parsed_uri == (
+        "test_db.nc",
+        "equilibrium",
+        0,
+        "/time_slice()/boundary/elongation",
+    )
+
+
+def test_to_csv(exporter, tmp_path):
+    """Test if exported CSV exists and is filled with correct times and values."""
+    file_path = tmp_path / "waveform.csv"
+    exporter.to_csv(file_path)
+
+    with open(file_path, newline="") as csvfile:
+        reader = csv.reader(csvfile)
+        rows = list(reader)
+
+    assert Path(file_path).exists()
+    assert len(rows) == len(exporter.times) + 1  # Header + data rows
+    for i, (time, value) in enumerate(zip(exporter.times, exporter.values), start=1):
+        assert float(rows[i][0]) == time
+        assert float(rows[i][1]) == value
+
+
+def test_to_png(exporter, tmp_path):
+    """Test if exported PNG exists."""
+    file_path = tmp_path / "waveform.png"
+    exporter.to_png(file_path)
+    assert Path(file_path).exists()
+
+
+def test_to_ids_flt1d(exporter, tmp_path):
+    """Check if FLT_1D is filled correctly for a homogeneous time IDS."""
+    file_path = f"{tmp_path}/test_ec_launchers.nc"
+
+    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):
+    """Check if FLT_1D is filled correctly for a heterogeneous time IDS."""
+    file_path = f"{tmp_path}/test_ec_launchers.nc"
+
+    uri = f"{file_path}#ec_launchers/beam(1)/power_launched"
+    exporter.to_ids(uri, is_homogeneous=False)
+
+    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):
+    """Check if FLT_1D is filled correctly when providing an occurrence number."""
+    file_path = f"{tmp_path}/test_ec_launchers.nc"
+
+    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
+        with pytest.raises(IndexError):
+            ids = dbentry.get("ec_launchers", occurrence=0, autoconvert=False)
+
+        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):
+    """Check if FLT_0D is filled correctly for a homogeneous time IDS."""
+    file_path = f"{tmp_path}/test_equilibrium.nc"
+
+    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):
+    """Check if FLT_0D is filled correctly for a heterogeneous time IDS."""
+    file_path = f"{tmp_path}/test_equilibrium.nc"
+
+    uri = f"{file_path}#equilibrium/time_slice()/boundary/elongation"
+    exporter.to_ids(uri, is_homogeneous=False)
+
+    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]
+
+
+def test_to_ids_no_times(tmp_path):
+    """Check if FLT_1D is filled correctly when not providing a time array."""
+    waveform_list = [
+        {
+            "user_type": "linear",
+            "user_from": 0,
+            "user_to": 8,
+            "user_start": 5,
+            "user_duration": 5,
+            "line_number": 1,
+        }
+    ]
+    waveform = Waveform(waveform=waveform_list)
+    exporter = WaveformExporter(waveform)
+
+    file_path = f"{tmp_path}/test_ec_launchers.nc"
+
+    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(np.isclose(ids.time, [5, 10]))
+        assert np.all(np.isclose(ids.beam[0].power_launched.data, [0, 8]))