Remove dead code, improve readability and create dedicated vm.

Author: shumpohl

qupulse/program/measurement.py

@@ -1,37 +1,68 @@
-from typing import Sequence, Mapping, Iterable, Optional, Union, ContextManager
+import contextlib
+from typing import Sequence, Mapping, Iterable, Optional, Union, ContextManager, Callable
 from dataclasses import dataclass
+from functools import cached_property
 
 import numpy
-from rich.measure import Measurement
 
 from qupulse.utils.types import TimeType
 from qupulse.program import (ProgramBuilder, Program, HardwareVoltage, HardwareTime,
                              MeasurementWindow, Waveform, RepetitionCount, SimpleExpression)
 from qupulse.parameter_scope import Scope
 
 
+MeasurementID = str | int
+
+
+@dataclass
+class LoopLabel:
+    idx: int
+    runtime_name: str | None
+    count: RepetitionCount
+
+
 @dataclass
-class MeasurementNode:
-    windows: Sequence[MeasurementWindow]
+class Measure:
+    meas_id: MeasurementID
+    delay: HardwareTime
+    length: HardwareTime
+
+
+@dataclass
+class Wait:
     duration: HardwareTime
 
 
 @dataclass
-class MeasurementRepetition(MeasurementNode):
-    body: MeasurementNode
-    count: RepetitionCount
+class LoopJmp:
+    idx: int
+
+
+Command = Union[LoopLabel, LoopJmp, Wait, Measure]
+
 
 @dataclass
-class MeasurementSequence(MeasurementNode):
-    nodes: Sequence[tuple[HardwareTime, MeasurementNode]]
+class MeasurementInstructions(Program):
+    commands: Sequence[Command]
+
+    @cached_property
+    def duration(self) -> float:
+        latest = 0.
+
+        def process(_, begin, length):
+            nonlocal latest
+            end = begin + length
+            latest = max(latest, end)
+
+        vm = MeasurementVM(process)
+        vm.execute(commands=self.commands)
+        return latest
 
 
 @dataclass
 class MeasurementFrame:
     commands: list['Command']
-    has_duration: bool
-
-MeasurementID = str | int
+    keep: bool
 
 
 class MeasurementBuilder(ProgramBuilder):
@@ -48,12 +79,11 @@ def _with_new_frame(self, measurements):
         self._frames.append(MeasurementFrame([], False))
         yield self
         frame = self._frames.pop()
-        if not frame.has_duration:
+        if not frame.keep:
             return
-        parent = self._frames[-1]
-        parent.has_duration = True
-        if measurements:
-            parent.commands.extend(map(Measure, measurements))
+        self.measure(measurements)
+        # measure does not keep if there are no measurements
+        self._frames[-1].keep = True
         return frame.commands
 
     def inner_scope(self, scope: Scope) -> Scope:
@@ -68,19 +98,19 @@ def inner_scope(self, scope: Scope) -> Scope:
     def hold_voltage(self, duration: HardwareTime, voltages: Mapping[str, HardwareVoltage]):
         """Supports dynamic i.e. for loop generated offsets and duration"""
         self._frames[-1].commands.append(Wait(duration))
-        self._frames[-1].has_duration = True
+        self._frames[-1].keep = True
 
     def play_arbitrary_waveform(self, waveform: Waveform):
         """"""
         self._frames[-1].commands.append(Wait(waveform.duration))
-        self._frames[-1].has_duration = True
+        self._frames[-1].keep = True
 
     def measure(self, measurements: Optional[Sequence[MeasurementWindow]]):
         """Unconditionally add given measurements relative to the current position."""
         if measurements:
             commands = self._frames[-1].commands
             commands.extend(Measure(*meas) for meas in measurements)
-            self._frames[-1].has_duration = True
+            self._frames[-1].keep = True
 
     def with_repetition(self, repetition_count: RepetitionCount,
                         measurements: Optional[Sequence[MeasurementWindow]] = None) -> Iterable['ProgramBuilder']:
@@ -92,10 +122,11 @@ def with_repetition(self, repetition_count: RepetitionCount,
 
         self._label_counter += 1
         label_idx = self._label_counter
-        parent.commands.append(LoopLabel(idx=label_idx, runtime_name=None, count=RepetitionCount))
+        parent.commands.append(LoopLabel(idx=label_idx, runtime_name=None, count=repetition_count))
         parent.commands.extend(new_commands)
         parent.commands.append(LoopJmp(idx=label_idx))
 
+    @contextlib.contextmanager
     def with_sequence(self,
                       measurements: Optional[Sequence[MeasurementWindow]] = None) -> ContextManager['ProgramBuilder']:
         """
@@ -112,6 +143,7 @@ def with_sequence(self,
         parent = self._frames[-1]
         parent.commands.extend(new_commands)
 
+    @contextlib.contextmanager
     def new_subprogram(self, global_transformation: 'Transformation' = None) -> ContextManager['ProgramBuilder']:
         """Create a context managed program builder whose contents are translated into a single waveform upon exit if
         it is not empty."""
@@ -136,43 +168,16 @@ def time_reversed(self) -> ContextManager['ProgramBuilder']:
         self._frames.append(MeasurementFrame([], False))
         yield self
         frame = self._frames.pop()
-        if not frame.has_duration:
+        if not frame.keep:
             return
 
-        self._frames[-1].has_duration = True
+        self._frames[-1].keep = True
         self._frames[-1].commands.extend(_reversed_commands(frame.commands))
 
     def to_program(self) -> Optional[Program]:
         """Further addition of new elements might fail after finalizing the program."""
-        if self._frames[0].has_duration:
-            return self._frames[0].commands
-
-
-@dataclass
-class LoopLabel:
-    idx: int
-    runtime_name: str | None
-    count: RepetitionCount
-
-
-@dataclass
-class Measure:
-    meas_id: MeasurementID
-    delay: HardwareTime
-    length: HardwareTime
-
-
-@dataclass
-class Wait:
-    duration: HardwareTime
-
-
-@dataclass
-class LoopJmp:
-    idx: int
-
-
-Command = Union[LoopLabel, LoopJmp, Wait, Measure]
+        if self._frames[0].keep:
+            return MeasurementInstructions(self._frames[0].commands)
 
 
 def _reversed_commands(cmds: Sequence[Command]) -> Sequence[Command]:
@@ -202,30 +207,26 @@ def _reversed_commands(cmds: Sequence[Command]) -> Sequence[Command]:
     return reversed_cmds
 
 
-def to_table(commands: Sequence[Command]) -> dict[str, numpy.ndarray]:
-    time = TimeType(0)
-
-    memory = {}
-    counts = [None]
+class MeasurementVM:
+    """A VM that is capable of executing the measurement commands"""
 
-    tables = {}
+    def __init__(self, callback: Callable[[str, float, float], None]):
+        self._time = TimeType(0)
+        self._memory = {}
+        self._counts = {}
+        self._callback = callback
 
-    def eval_hardware_time(t: HardwareTime):
+    def _eval_hardware_time(self, t: HardwareTime):
         if isinstance(t, SimpleExpression):
             value = t.base
             for (factor_name, factor_val) in t.offsets.items():
-                count = counts[memory[factor_name]]
+                count = self._counts[self._memory[factor_name]]
                 value += factor_val * count
             return value
         else:
             return t
 
-    def execute(sequence: Sequence[Command]) -> int:
-        nonlocal time
-        nonlocal tables
-        nonlocal memory
-        nonlocal counts
-
+    def _execute_after_label(self, sequence: Sequence[Command]) -> int:
         skip = 0
         for idx, cmd in enumerate(sequence):
             if idx < skip:
@@ -234,23 +235,30 @@ def execute(sequence: Sequence[Command]) -> int:
                 return idx
             elif isinstance(cmd, LoopLabel):
                 if cmd.runtime_name:
-                    memory[cmd.runtime_name] = cmd.idx
-                if cmd.idx == len(counts):
-                    counts.append(0)
-                assert cmd.idx < len(counts)
+                    self._memory[cmd.runtime_name] = cmd.idx
 
                 for iter_val in range(cmd.count):
-                    counts[cmd.idx] = iter_val
-                    pos = execute(sequence[idx + 1:])
+                    self._counts[cmd.idx] = iter_val
+                    pos = self._execute_after_label(sequence[idx + 1:])
                 skip = idx + pos + 2
+
             elif isinstance(cmd, Measure):
-                meas_time = float(eval_hardware_time(cmd.delay) + time)
-                meas_len = float(eval_hardware_time(cmd.length))
-                tables.setdefault(cmd.meas_id, []).append((meas_time, meas_len))
+                meas_time = float(self._eval_hardware_time(cmd.delay) + self._time)
+                meas_len = float(self._eval_hardware_time(cmd.length))
+                self._callback(cmd.meas_id, meas_time, meas_len)
+
             elif isinstance(cmd, Wait):
-                time += eval_hardware_time(cmd.duration)
+                self._time += self._eval_hardware_time(cmd.duration)
+
+    def execute(self, commands: Sequence[Command]):
+        self._execute_after_label(commands)
+
+
+def to_table(commands: Sequence[Command]) -> dict[str, numpy.ndarray]:
+    tables = {}
 
-    execute(commands)
+    vm = MeasurementVM(lambda name, begin, length: tables.setdefault(name, []).append((begin, length)))
+    vm.execute(commands)
     return {
         name: numpy.array(values) for name, values in tables.items()
     }

tests/program/measurement_test.py

@@ -26,8 +26,8 @@ def setUp(self):
 
     def test_commands(self):
         builder = MeasurementBuilder()
-        commands = self.pulse_template.create_program(program_builder=builder)
-        self.assertEqual(self.commands, commands)
+        instructions = self.pulse_template.create_program(program_builder=builder)
+        self.assertEqual(self.commands, instructions.commands)
 
     def test_table(self):
         table = to_table(self.commands)
@@ -37,3 +37,24 @@ def test_table(self):
         np.testing.assert_array_equal(self.table_b, tab_b)
 
 
+class ComplexPulse(TestCase):
+    def setUp(self):
+        hold = ConstantPT(10 ** 6, {'a': 1}, measurements=[('A', 10, 100), ('B', '1 + ii * 2 + jj', '3 + ii + jj')])
+        dyn_hold = ConstantPT('10 ** 6 - 4 * ii', {'a': 1}, measurements=[('A', 10, 100), ('B', '1 + ii * 2 + jj', '3 + ii + jj')])
+
+        self.pulse_template = SequencePT(
+            hold.with_repetition(2).with_iteration('ii', 100).with_repetition(2).with_iteration('jj', 200),
+            measurements=[('A', 1, 100)]
+        ).with_repetition(2)
+
+        self.commands = []
+
+    def test_commands(self):
+        builder = MeasurementBuilder()
+        commands = self.pulse_template.create_program(program_builder=builder)
+        to_table(commands.commands)
+        raise NotImplementedError("TODO")
+
+    def test_table(self):
+
+        raise NotImplementedError("TODO")