First version of a measurement builder

shumpohl · shumpohl · commit cfcee75c0e28 · 2025-01-23T16:10:54.000+01:00
diff --git a/qupulse/program/measurement.py b/qupulse/program/measurement.py
@@ -0,0 +1,256 @@
+from typing import Sequence, Mapping, Iterable, Optional, Union, ContextManager
+from dataclasses import dataclass
+
+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
+
+
+@dataclass
+class MeasurementNode:
+    windows: Sequence[MeasurementWindow]
+    duration: HardwareTime
+
+
+@dataclass
+class MeasurementRepetition(MeasurementNode):
+    body: MeasurementNode
+    count: RepetitionCount
+
+@dataclass
+class MeasurementSequence(MeasurementNode):
+    nodes: Sequence[tuple[HardwareTime, MeasurementNode]]
+
+
+@dataclass
+class MeasurementFrame:
+    commands: list['Command']
+    has_duration: bool
+
+MeasurementID = str | int
+
+
+class MeasurementBuilder(ProgramBuilder):
+    def __init__(self):
+        super().__init__()
+
+        self._frames = [MeasurementFrame([], False)]
+        self._ranges: list[tuple[str, range]] = []
+        self._repetitions = []
+        self._measurements = []
+        self._label_counter = 0
+
+    def _with_new_frame(self, measurements):
+        self._frames.append(MeasurementFrame([], False))
+        yield self
+        frame = self._frames.pop()
+        if not frame.has_duration:
+            return
+        parent = self._frames[-1]
+        parent.has_duration = True
+        if measurements:
+            parent.commands.extend(map(Measure, measurements))
+        return frame.commands
+
+    def inner_scope(self, scope: Scope) -> Scope:
+        """This function is necessary to inject program builder specific parameter implementations into the build
+        process."""
+        if self._ranges:
+            name, rng = self._ranges[-1]
+            return scope.overwrite({name: SimpleExpression(base=rng.start, offsets={name: rng.step})})
+        else:
+            return 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
+
+    def play_arbitrary_waveform(self, waveform: Waveform):
+        """"""
+        self._frames[-1].commands.append(Wait(waveform.duration))
+        self._frames[-1].has_duration = 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
+
+    def with_repetition(self, repetition_count: RepetitionCount,
+                        measurements: Optional[Sequence[MeasurementWindow]] = None) -> Iterable['ProgramBuilder']:
+        """Measurements that are added to the new builder are dropped if the builder is empty upon exit"""
+        new_commands = yield from self._with_new_frame(measurements)
+        if new_commands is None:
+            return
+        parent = self._frames[-1]
+
+        self._label_counter += 1
+        label_idx = self._label_counter
+        parent.commands.append(LoopLabel(idx=label_idx, runtime_name=None, count=RepetitionCount))
+        parent.commands.extend(new_commands)
+        parent.commands.append(LoopJmp(idx=label_idx))
+
+    def with_sequence(self,
+                      measurements: Optional[Sequence[MeasurementWindow]] = None) -> ContextManager['ProgramBuilder']:
+        """
+
+        Measurements that are added in to the returned program builder are discarded if the sequence is empty on exit.
+
+        Args:
+            measurements: Measurements to attach to the potential child.
+        Returns:
+        """
+        new_commands = yield from self._with_new_frame(measurements)
+        if new_commands is None:
+            return
+        parent = self._frames[-1]
+        parent.commands.extend(new_commands)
+
+    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."""
+        yield self
+
+    def with_iteration(self, index_name: str, rng: range,
+                       measurements: Optional[Sequence[MeasurementWindow]] = None) -> Iterable['ProgramBuilder']:
+        self._ranges.append((index_name, rng))
+        new_commands = yield from self._with_new_frame(measurements)
+        self._ranges.pop()
+        if new_commands is None:
+            return
+        parent = self._frames[-1]
+
+        self._label_counter += 1
+        label_idx = self._label_counter
+        parent.commands.append(LoopLabel(idx=label_idx, runtime_name=index_name, count=len(rng)))
+        parent.commands.extend(new_commands)
+        parent.commands.append(LoopJmp(idx=label_idx))
+
+    def time_reversed(self) -> ContextManager['ProgramBuilder']:
+        self._frames.append(MeasurementFrame([], False))
+        yield self
+        frame = self._frames.pop()
+        if not frame.has_duration:
+            return
+
+        self._frames[-1].has_duration = 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]
+
+
+def _reversed_commands(cmds: Sequence[Command]) -> Sequence[Command]:
+    reversed_cmds = []
+    jumps = {}
+    for cmd in reversed(cmds):
+        if isinstance(cmd, LoopJmp):
+            jumps[cmd.idx] = len(reversed_cmds)
+            reversed_cmds.append(cmd)
+        elif isinstance(cmd, LoopLabel):
+            jump_idx = jumps[cmd.idx]
+            jump = reversed_cmds[jump_idx]
+            reversed_cmds[jump_idx] = cmd
+            reversed_cmds.append(jump)
+
+        elif isinstance(cmd, Measure):
+            if isinstance(cmd.delay, SimpleExpression) or isinstance(cmd.delay, SimpleExpression):
+                raise NotImplementedError("TODO")
+            reversed_cmds.append(Measure(meas_id=cmd.meas_id,
+                                         delay=-(cmd.delay + cmd.length),
+                                         length=cmd.length,))
+        elif isinstance(cmd, Wait):
+            reversed_cmds.append(cmd)
+        else:
+            raise ValueError("Not a command", cmd)
+
+    return reversed_cmds
+
+
+def to_table(commands: Sequence[Command]) -> dict[str, numpy.ndarray]:
+    time = TimeType(0)
+
+    memory = {}
+    counts = [None]
+
+    tables = {}
+
+    def eval_hardware_time(t: HardwareTime):
+        if isinstance(t, SimpleExpression):
+            value = t.base
+            for (factor_name, factor_val) in t.offsets.items():
+                count = counts[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
+
+        skip = 0
+        for idx, cmd in enumerate(sequence):
+            if idx < skip:
+                continue
+            if isinstance(cmd, LoopJmp):
+                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)
+
+                for iter_val in range(cmd.count):
+                    counts[cmd.idx] = iter_val
+                    pos = execute(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))
+            elif isinstance(cmd, Wait):
+                time += eval_hardware_time(cmd.duration)
+
+    execute(commands)
+    return {
+        name: numpy.array(values) for name, values in tables.items()
+    }
diff --git a/qupulse/pulses/measurement.py b/qupulse/pulses/measurement.py
@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: 2014-2024 Quantum Technology Group and Chair of Software Engineering, RWTH Aachen University
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
-
+import warnings
 from typing import Optional, List, Tuple, Union, Dict, Set, Mapping, AbstractSet
 from numbers import Real
 import itertools
@@ -55,8 +55,11 @@ def get_measurement_windows(self,
 
             begin_val = begin.evaluate_in_scope(parameters)
             length_val = length.evaluate_in_scope(parameters)
-            if begin_val < 0 or length_val < 0:
-                raise ValueError('Measurement window with negative begin or length: {}, {}'.format(begin, length))
+            try:
+                if begin_val < 0 or length_val < 0:
+                    warnings.warn('Measurement window with negative begin or length: {}, {}'.format(begin, length))
+            except TypeError:
+                pass
 
             resulting_windows.append(
                 (name,
diff --git a/tests/program/measurement_test.py b/tests/program/measurement_test.py
@@ -0,0 +1,39 @@
+import copy
+import unittest
+from unittest import TestCase
+
+import numpy as np
+
+from qupulse.pulses import *
+from qupulse.program.measurement import *
+
+
+class SingleRampTest(TestCase):
+    def setUp(self):
+        hold = ConstantPT(10 ** 6, {'a': '-1. + idx * 0.01'}, measurements=[('A', 10, 100), ('B', '1 + idx * 2', 200)])
+        self.pulse_template = hold.with_iteration('idx', 200)
+
+        self.commands = [
+            LoopLabel(1, 'idx', 200),
+            Measure('A', 10, 100),
+            Measure('B', SimpleExpression(base=1, offsets={'idx': 2}), 200),
+            Wait(TimeType(10 ** 6)),
+            LoopJmp(1)
+        ]
+
+        self.table_a = np.array([(10 + 10**6 * idx, 100) for idx in range(200)])
+        self.table_b = np.array([(1 + idx * 2 + 10**6 * idx, 200) for idx in range(200)])
+
+    def test_commands(self):
+        builder = MeasurementBuilder()
+        commands = self.pulse_template.create_program(program_builder=builder)
+        self.assertEqual(self.commands, commands)
+
+    def test_table(self):
+        table = to_table(self.commands)
+        tab_a = table['A']
+        tab_b = table['B']
+        np.testing.assert_array_equal(self.table_a, tab_a)
+        np.testing.assert_array_equal(self.table_b, tab_b)
+
+
