Type aware constants

docs/source/derived.rst

@@ -105,6 +105,11 @@ 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>`).
+
 
 Using NumPy Functions
 ---------------------

docs/source/tendencies.rst

@@ -51,6 +51,22 @@ 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, a string, or a boolean:
+
+.. code-block:: yaml
+
+    - {type: constant, value: ohmic, duration: 2}
+    - {type: constant, value: nbi, duration: 2}
+
+.. warning::
+    Integers and floats may be freely combined within a single waveform, but other
+    value types may not be combined with each other.
+
 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,38 @@ def test_generate():
     assert not tendency.annotations
 
 
+@pytest.mark.parametrize(
+    "value", ["ec", True, 3, 3.5], ids=["str", "bool", "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]
+
+
+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", True], ids=["int", "float", "str", "bool"]
+)
+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)
+
+
 def test_declarative_assignments():
     t1 = ConstantTendency(user_duration=1)
     t2 = ConstantTendency(user_duration=1)

tests/tendencies/test_repeat.py

@@ -177,6 +177,20 @@ def test_too_short(repeat_waveform):
     assert repeat_tendency.annotations[0]["type"] == "warning"
 
 
+@pytest.mark.parametrize("value", ["ec", True], ids=["str", "bool"])
+def test_categorical_value_not_supported(value):
+    """Categorical values inside a repeat tendency are not allowed"""
+    repeat_tendency = RepeatTendency(
+        user_duration=4,
+        user_waveform=[
+            {"user_type": "constant", "user_value": value, "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_dependency_graph.py

@@ -87,3 +87,22 @@ 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", [])
+
+    order = dg.topological_order()
+    assert set(order) == {"A", "B", "C"}
+    assert order.index("C") < order.index("B") < order.index("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,74 @@ def test_function_access_control(filled_config):
         else:
             with pytest.raises(NameError):
                 waveform.get_value(time_ret)
+
+
+def test_derived_waveform_type_matches_original(config):
+    original_name = "wf1"
+    original = Waveform(
+        waveform=[{"user_type": "constant", "user_value": 3, "line_number": 1}],
+        name=original_name,
+    )
+    assert not original.annotations
+    assert original.value_type is int
+    config.add_waveform(original, ["root_group"])
+
+    derived_name = "wf2"
+    yaml_str = f"{derived_name}: |\n  '{original_name}'"
+    derived = DerivedWaveform(yaml_str, derived_name, config)
+    assert derived.dependencies == {original_name}
+    assert derived.value_type == original.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: hello, duration: 2}
+    """
+    config = WaveformConfiguration()
+    config.load_yaml(yaml_str)
+
+    wf4 = config["wf4"]
+    wf3 = config["wf3"]
+    wf2 = config["wf2"]
+    wf1 = config["wf1"]
+    assert wf4.value_type is str
+    assert wf3.value_type is str
+    assert wf2.value_type is str
+    assert wf1.value_type is str
+    assert not wf4.annotations
+    assert not wf3.annotations
+    assert not wf2.annotations
+    assert not wf1.annotations
+
+
+def test_derived_waveform_type_mixing(config):
+    wf1_name = "wf1"
+    wf1 = Waveform(
+        waveform=[{"user_type": "constant", "user_value": 3, "line_number": 1}],
+        name=wf1_name,
+    )
+    assert not wf1.annotations
+    assert wf1.value_type is int
+    config.add_waveform(wf1, ["root_group"])
+
+    wf2_name = "wf2"
+    wf2 = Waveform(
+        waveform=[{"user_type": "constant", "user_value": "test", "line_number": 2}],
+        name=wf2_name,
+    )
+    assert not wf2.annotations
+    assert wf2.value_type is str
+    config.add_waveform(wf2, ["root_group"])
+
+    derived_name = "derived_waveform"
+    yaml_str = f"{derived_name}: |\n  '{wf1_name}' + '{wf2_name}'"
+    derived = DerivedWaveform(yaml_str, derived_name, config)
+    assert derived.annotations  # not allowed to mix str and int type waveforms

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: true, duration: 2}
+      - {type: constant, value: false, 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."""