Type aware constants

docs/source/derived.rst

@@ -105,6 +105,12 @@ In the example below, waveform ``test/3`` is the sum of the waveforms ``test/1``
    :width: 600px
    :align: center
 
+.. note::
+    A derived waveform's value type mirrors that of the waveform(s) it depends on.
+    When an expression references multiple waveforms, they must all share the same
+    value type (see :ref:`Value Types <constant-value-types>`). Derived waveforms
+    cannot depend on string-typed waveforms.
+
 
 Using NumPy Functions
 ---------------------

docs/source/tendencies.rst

@@ -51,6 +51,45 @@ If the ``value`` is not specified, it will be set to the last value of the previ
     - {type: linear, to: 3, duration: 10}
     - {type: constant, duration: 10}
 
+.. _constant-value-types:
+
+Value Types
+-----------
+
+The ``value`` of a constant tendency may be a number (integer or floating-point) or a string.
+
+An integer value:
+
+.. code-block:: yaml
+
+    - {type: constant, value: 3, duration: 2}
+
+A floating-point value:
+
+.. code-block:: yaml
+
+    - {type: constant, value: 3.5, duration: 2}
+
+A string value:
+
+.. code-block:: yaml
+
+    - {type: constant, value: ohmic, duration: 2}
+    - {type: constant, value: nbi, duration: 2}
+
+It is also possible to use ``True`` and ``False`` as the ``value``, in which case the 
+waveform will be considered an integer type (where ``True = 1``, and ``False = 0``):
+
+.. code-block:: yaml
+
+    - {type: constant, value: True, duration: 2}
+    - {type: constant, value: False, duration: 2}
+
+.. warning::
+    Integers and floats may be freely combined within a single waveform (the
+    resulting waveform will be float-typed), but it's not allowed to combine string 
+    values with numbers in a waveform.
+
 Linear Tendency
 ===============
 

tests/tendencies/test_constant.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pytest
 
 from waveform_editor.tendencies.constant import ConstantTendency
 
@@ -69,6 +70,59 @@ def test_generate():
     assert not tendency.annotations
 
 
+@pytest.mark.parametrize("value", ["ec", 3, 3.5], ids=["str", "int", "float"])
+def test_categorical_value(value):
+    tendency = ConstantTendency(user_start=0, user_duration=1, user_value=value)
+    assert tendency.value == value
+    assert tendency.value_type is type(value)
+    assert not tendency.annotations
+
+    time, values = tendency.get_value()
+    assert np.all(time == np.array([0, 1]))
+    assert list(values) == [value, value]
+
+
+@pytest.mark.parametrize("value,expected", [(True, 1), (False, 0)])
+def test_bool_value_treated_as_int(value, expected):
+    tendency = ConstantTendency(user_start=0, user_duration=1, user_value=value)
+    assert tendency.value == expected
+    assert type(tendency.value) is int
+    assert tendency.value_type is int
+    assert not tendency.annotations
+
+
+def test_unsupported_value_type():
+    tendency = ConstantTendency(user_start=0, user_duration=1, user_value=[1, 2, 3])
+    assert tendency.annotations
+    assert tendency.value == 0.0
+
+
+@pytest.mark.parametrize("value", [5, 5.5, "ec"], ids=["int", "float", "str"])
+def test_inherited_value(value):
+    t1 = ConstantTendency(user_value=value, user_start=0, user_duration=1)
+    t2 = ConstantTendency(user_duration=1)
+    t2.set_previous_tendency(t1)
+
+    assert t2.value_type is type(value)
+    assert type(t2.value) is type(value)
+
+
+@pytest.mark.parametrize("value", [5, 5.5, "ec"], ids=["int", "float", "str"])
+def test_inherited_value_stays_plain_type_across_chain(value):
+    """start_value/end_value/value must be plain Python types"""
+    t1 = ConstantTendency(user_value=value, user_start=0, user_duration=1)
+    t2 = ConstantTendency(user_duration=1)
+    t2.set_previous_tendency(t1)
+    t3 = ConstantTendency(user_duration=1)
+    t3.set_previous_tendency(t2)
+
+    for tendency in (t1, t2, t3):
+        assert tendency.value_type is type(value)
+        assert type(tendency.value) is type(value)
+        assert type(tendency.start_value) is type(value)
+        assert type(tendency.end_value) is type(value)
+
+
 def test_declarative_assignments():
     t1 = ConstantTendency(user_duration=1)
     t2 = ConstantTendency(user_duration=1)

tests/tendencies/test_repeat.py

@@ -177,6 +177,19 @@ def test_too_short(repeat_waveform):
     assert repeat_tendency.annotations[0]["type"] == "warning"
 
 
+def test_string_value_not_supported():
+    """String values inside a repeat tendency are not allowed"""
+    repeat_tendency = RepeatTendency(
+        user_duration=4,
+        user_waveform=[
+            {"user_type": "constant", "user_value": "ec", "user_duration": 1},
+        ],
+    )
+    assert repeat_tendency.annotations
+    times = np.linspace(0, 4, 9)
+    repeat_tendency.get_value(times)
+
+
 def test_period(repeat_waveform):
     """Check values when period is provided."""
     repeat_waveform["user_period"] = 1

tests/test_configuration.py

@@ -3,6 +3,7 @@
 import pytest
 
 from waveform_editor.configuration import WaveformConfiguration
+from waveform_editor.derived_waveform import DerivedWaveform
 from waveform_editor.tendencies.constant import ConstantTendency
 from waveform_editor.tendencies.linear import LinearTendency
 from waveform_editor.tendencies.periodic.sine_wave import SineWaveTendency
@@ -109,6 +110,54 @@ def test_replace_waveform(config):
         config.replace_waveform(waveform3)
 
 
+def test_replace_waveform_revalidates_dependents(config):
+    """Test if replacing a waveform re-evaluates the value_type of derived waveforms"""
+    path = ["ec_launchers"]
+
+    int_waveform = Waveform(
+        waveform=[{"user_type": "constant", "user_value": 3, "user_duration": 1}],
+        name="A",
+    )
+    config.add_waveform(int_waveform, path)
+
+    derived_b = DerivedWaveform("B: \"'A'\"", "B", config)
+    config.add_waveform(derived_b, path)
+    derived_c = DerivedWaveform("C: \"'B'\"", "C", config)
+    config.add_waveform(derived_c, path)
+
+    assert derived_b.value_type is int
+    assert derived_c.value_type is int
+
+    flt_waveform = Waveform(
+        waveform=[{"user_type": "constant", "user_value": 3.5, "user_duration": 1}],
+        name="A",
+    )
+    config.replace_waveform(flt_waveform)
+
+    assert derived_b.value_type is float
+    assert derived_c.value_type is float
+
+
+def test_add_waveform_revalidates_previously_missing_dependency(config):
+    """Test if a derived waveform's error is cleared once its dependency is added to
+    the configuration."""
+    path = ["ec_launchers"]
+
+    derived = DerivedWaveform("B: \"'A'\"", "B", config)
+    config.add_waveform(derived, path)
+    assert derived.annotations
+    assert derived.value_type is float
+
+    waveform = Waveform(
+        waveform=[{"user_type": "constant", "user_value": 3, "user_duration": 1}],
+        name="A",
+    )
+    config.add_waveform(waveform, path)
+
+    assert not derived.annotations
+    assert derived.value_type is int
+
+
 def test_remove_waveform(config):
     """Test if waveforms are removed correctly from configuration."""
 

tests/test_dependency_graph.py

@@ -87,3 +87,20 @@ def test_detect_cycles_with_start_node():
     dg.graph["C"] = {"A"}
     with pytest.raises(RuntimeError):
         dg.detect_cycles("A")
+
+
+def test_topological_order():
+    dg = DependencyGraph()
+    dg.add_node("A", ["B"])
+    dg.add_node("B", ["C"])
+    dg.add_node("C", [])
+
+    assert dg.topological_order() == ["C", "B", "A"]
+
+
+def test_topological_order_ignores_leaf_dependencies():
+    """Dependencies that are not nodes themselves should not show up."""
+    dg = DependencyGraph()
+    dg.add_node("A", ["leaf"])
+
+    assert dg.topological_order() == ["A"]

tests/test_derived_waveform.py

@@ -120,9 +120,10 @@ def test_rename_waveform(filled_config):
     name = "waveform/2"
     yaml_str = f"{name}: |\n  'waveform/1'"
     waveform = DerivedWaveform(yaml_str, name, filled_config)
+    filled_config.add_waveform(waveform, ["root_group"])
     assert waveform.dependencies == {"waveform/1"}
     assert waveform.get_yaml_string() == "'waveform/1'"
-    waveform.rename_dependency("waveform/1", "waveform/3")
+    filled_config.rename_waveform("waveform/1", "waveform/3")
     assert waveform.dependencies == {"waveform/3"}
     assert waveform.get_yaml_string() == "'waveform/3'"
 
@@ -151,3 +152,92 @@ def test_function_access_control(filled_config):
         else:
             with pytest.raises(NameError):
                 waveform.get_value(time_ret)
+
+
+def test_derived_waveform_type_matches_original():
+    yaml_str = """
+    root_group:
+      wf1:
+        - {type: constant, value: 3, duration: 2}
+      wf2: |
+        'wf1'
+    """
+    config = WaveformConfiguration()
+    config.load_yaml(yaml_str)
+
+    assert not config["wf1"].annotations
+    assert config["wf1"].value_type is int
+    assert config["wf2"].dependencies == {"wf1"}
+    assert config["wf2"].value_type == config["wf1"].value_type
+
+
+def test_derived_waveform_chain_type_order_independent():
+    yaml_str = """
+    root_group:
+      wf1: |
+        'wf2' + 'wf3'
+      wf2: |
+        'wf3'
+      wf3: |
+        'wf4'
+      wf4:
+        - {type: constant, value: 3, duration: 2}
+    """
+    config = WaveformConfiguration()
+    config.load_yaml(yaml_str)
+
+    for wf in ["wf1", "wf2", "wf3", "wf4"]:
+        assert config[wf].value_type is int
+        assert not config[wf].annotations
+
+
+def test_derived_waveform_type_mixing():
+    yaml_str = """
+    root_group:
+      wf1:
+        - {type: constant, value: 3, duration: 2}
+      wf2:
+        - {type: constant, value: test, duration: 2}
+      derived_waveform: |
+        'wf1' + 'wf2'
+    """
+    config = WaveformConfiguration()
+    config.load_yaml(yaml_str)
+
+    assert config["wf1"].value_type is int
+    assert config["wf2"].value_type is str
+    assert config["derived_waveform"].annotations  # not allowed to mix str and int
+
+
+def test_derived_waveform_disallows_string_dependency():
+    yaml_str = """
+    root_group:
+      wf1:
+        - {type: constant, value: ohmic, duration: 2}
+      derived_waveform: |
+        'wf1'
+    """
+    config = WaveformConfiguration()
+    config.load_yaml(yaml_str)
+
+    assert config["wf1"].value_type is str
+    assert config["derived_waveform"].annotations
+
+
+def test_derived_waveform_int_float_mixing():
+    yaml_str = """
+    root_group:
+      wf1:
+        - {type: constant, value: 3, duration: 2}
+      wf2:
+        - {type: constant, value: 3.5, duration: 2}
+      derived_waveform: |
+        'wf1' + 'wf2'
+    """
+    config = WaveformConfiguration()
+    config.load_yaml(yaml_str)
+
+    assert config["wf1"].value_type is int
+    assert config["wf2"].value_type is float
+    assert not config["derived_waveform"].annotations
+    assert config["derived_waveform"].value_type is float

tests/test_exporter.py

@@ -647,6 +647,43 @@ def test_export_constant(tmp_path):
         assert np.array_equal(ids.beam[2].phase.angle, [3.3e3] * 3)
 
 
+def test_export_typed_waveforms(tmp_path):
+    """Check that constant waveforms of each supported value type"""
+
+    yaml_str = """
+    globals:
+      dd_version: 4.0.0
+    core_profiles:
+      core_profiles/profiles_1d/electrons/temperature_validity:
+      - {type: constant, value: 0, duration: 2}
+      - {type: constant, value: 1, duration: 2}
+      core_profiles/profiles_1d/grid/psi_magnetic_axis:
+      - {type: constant, value: 1.5, duration: 2}
+      - {type: constant, value: 3.0, duration: 2}
+      core_profiles/profiles_1d/ion(1)/name:
+      - {type: constant, value: D, duration: 2}
+      - {type: constant, value: He, duration: 2}
+      core_profiles/profiles_1d/ion(1)/multiple_states_flag:
+      - {type: constant, value: 1, duration: 2}
+      - {type: constant, value: 0, duration: 2}
+    """
+    file_path = f"{tmp_path}/test.nc"
+    times = np.array([0, 2.0])
+    _export_ids(file_path, yaml_str, times)
+
+    with imas.DBEntry(file_path, "r", dd_version="4.0.0") as dbentry:
+        core_profiles = dbentry.get("core_profiles", autoconvert=False)
+        assert core_profiles.profiles_1d[0].grid.psi_magnetic_axis == 1.5
+        assert core_profiles.profiles_1d[0].electrons.temperature_validity == 0
+        assert core_profiles.profiles_1d[0].ion[0].multiple_states_flag == 1
+        assert core_profiles.profiles_1d[0].ion[0].name == "D"
+
+        assert core_profiles.profiles_1d[1].grid.psi_magnetic_axis == 3.0
+        assert core_profiles.profiles_1d[1].electrons.temperature_validity == 1
+        assert core_profiles.profiles_1d[1].ion[0].multiple_states_flag == 0
+        assert core_profiles.profiles_1d[1].ion[0].name == "He"
+
+
 def test_example_yaml(tmp_path):
     """Test for an example YAML file if all IDSs are correctly filled."""