Repository: qutech/qupulse
Issue: `roll_constant_waveforms` does not handle measurement windows in leaves correctly

Long term: Replace with different functionality

Short term either
  1. emit warning if measurement windows in leaves are detected
  2. do not roll waveforms with measurement windows
  3. move measurement windows to parent loop (computationally more expensive)

Accompany 1 with a function `extract_measurement_windows` to `Loop` that is `get_measurement_windows` but removes them from the `Loop`.

Changed files:
qupulse/hardware/setup.py, qupulse/_program/_loop.py, tests/_program/loop_tests.py

diff --git a/qupulse/_program/_loop.py b/qupulse/_program/_loop.py
index d87aa0d7..92d3f7eb 100644
--- a/qupulse/_program/_loop.py
+++ b/qupulse/_program/_loop.py
@@ -231,9 +231,13 @@ class Loop(Node):
                           measurements=None if self._measurements is None else list(self._measurements),
                           children=(child.copy_tree_structure() for child in self))
 
-    def _get_measurement_windows(self) -> Mapping[str, np.ndarray]:
+    def _get_measurement_windows(self, drop: bool) -> Mapping[str, np.ndarray]:
         """Private implementation of get_measurement_windows with a slightly different data format for easier tiling.
 
+        Args:
+            drop: Drops the measurements from the Loop i.e. the Loop will no longer have measurements attached after
+            collecting them
+
         Returns:
              A dictionary (measurement_name -> array) with begin == array[:, 0] and length == array[:, 1]
         """
@@ -245,46 +249,43 @@ class Loop(Node):
             for mw_name, begin_length_list in temp_meas_windows.items():
                 temp_meas_windows[mw_name] = [np.asarray(begin_length_list, dtype=float)]
 
+            if drop:
+                self._measurements = None
+
         # calculate duration together with meas windows in the same iteration
         if self.is_leaf():
             body_duration = float(self.body_duration)
         else:
             offset = TimeType(0)
             for child in self:
-                for mw_name, begins_length_array in child._get_measurement_windows().items():
+                for mw_name, begins_length_array in child._get_measurement_windows(drop).items():
                     begins_length_array[:, 0] += float(offset)
                     temp_meas_windows[mw_name].append(begins_length_array)
                 offset += child.duration
 
             body_duration = float(offset)
 
-        # this gives us regular dict behaviour of the returned object
-        temp_meas_windows.default_factory = None
+        # formatting like this for easier debugging
+        result = {}
 
         # repeat and add repetition based offset
         for mw_name, begin_length_list in temp_meas_windows.items():
-            temp_begin_length_array = np.concatenate(begin_length_list)
-
-            begin_length_array = np.tile(temp_begin_length_array, (self.repetition_count, 1))
-
-            shaped_begin_length_array = np.reshape(begin_length_array, (self.repetition_count, -1, 2))
+            result[mw_name] = _repeat_loop_measurements(begin_length_list, self.repetition_count, body_duration)
 
-            shaped_begin_length_array[:, :, 0] += (np.arange(self.repetition_count) * body_duration)[:, np.newaxis]
+        return result
 
-            temp_meas_windows[mw_name] = begin_length_array
-
-        # the cast is here because static type analysis struggles to detect that we replace _all_ values by ndarray in
-        # the previous loop
-        return cast(Mapping[str, np.ndarray], temp_meas_windows)
-
-    def get_measurement_windows(self) -> Dict[str, Tuple[np.ndarray, np.ndarray]]:
+    def get_measurement_windows(self, drop=False) -> Dict[str, Tuple[np.ndarray, np.ndarray]]:
         """Iterates over all children and collect the begin and length arrays of each measurement window.
 
+        Args:
+            drop: Drops the measurements from the Loop i.e. the Loop will no longer have measurements attached after
+            collecting them.
+
         Returns:
             A dictionary (measurement_name -> (begin, length)) with begin and length being :class:`numpy.ndarray`
         """
         return {mw_name: (begin_length_list[:, 0], begin_length_list[:, 1])
-                for mw_name, begin_length_list in self._get_measurement_windows().items()}
+                for mw_name, begin_length_list in self._get_measurement_windows(drop=drop).items()}
 
     def split_one_child(self, child_index=None) -> None:
         """Take the last child that has a repetition count larger one, decrease it's repetition count and insert a copy
@@ -597,7 +598,8 @@ def make_compatible(program: Loop, minimal_waveform_length: int, waveform_quantu
 
 def roll_constant_waveforms(program: Loop, minimal_waveform_quanta: int, waveform_quantum: int, sample_rate: TimeType):
     """This function finds waveforms in program that can be replaced with repetitions of shorter waveforms and replaces
-    them. Complexity O(N_waveforms)
+    them. Complexity O(N_waveforms). Drops measurements because they are not correctly handled here for performance
+    reasons.
 
     This is possible if:
      - The waveform is constant on all channels
@@ -608,7 +610,15 @@ def roll_constant_waveforms(program: Loop, minimal_waveform_quanta: int, wavefor
         minimal_waveform_quanta:
         waveform_quantum:
         sample_rate:
+
+    Warnings:
+        DroppedMeasurementWarning: This warning is raised if a measurement is dropped.
     """
+    if program._measurements:
+        warnings.warn("Dropping measurements. Remove measurements before calling roll_constant_waveforms by calling"
+                      " get_measurement_windows(drop=True).", category=DroppedMeasurementWarning)
+        program._measurements = None
+
     waveform = program.waveform
 
     if waveform is None:
@@ -647,6 +657,21 @@ def roll_constant_waveforms(program: Loop, minimal_waveform_quanta: int, wavefor
         program._waveform = new_waveform
 
 
+def _repeat_loop_measurements(begin_length_list: List[np.ndarray],
+                              repetition_count: int,
+                              body_duration: float
+                              ) -> np.ndarray:
+    temp_begin_length_array = np.concatenate(begin_length_list)
+
+    begin_length_array = np.tile(temp_begin_length_array, (repetition_count, 1))
+
+    shaped_begin_length_array = np.reshape(begin_length_array, (repetition_count, -1, 2))
+
+    shaped_begin_length_array[:, :, 0] += (np.arange(repetition_count) * body_duration)[:, np.newaxis]
+
+    return begin_length_array
+
+
 class MakeCompatibleWarning(ResourceWarning):
     pass
 
diff --git a/qupulse/hardware/setup.py b/qupulse/hardware/setup.py
index bce71903..64681574 100644
--- a/qupulse/hardware/setup.py
+++ b/qupulse/hardware/setup.py
@@ -70,7 +70,7 @@ class MarkerChannel(_SingleChannel):
 
 
 RegisteredProgram = NamedTuple('RegisteredProgram', [('program', Loop),
-                                                     ('measurement_windows', Dict[str, Tuple[float, float]]),
+                                                     ('measurement_windows', Dict[str, Tuple[np.ndarray, np.ndarray]]),
                                                      ('run_callback', Callable),
                                                      ('awgs_to_upload_to', Set[AWG]),
                                                      ('dacs_to_arm', Set[DAC])])
@@ -102,7 +102,7 @@ class HardwareSetup:
                 channels - set(self._channel_map.keys())))
 
         temp_measurement_windows = defaultdict(list)
-        for mw_name, begins_lengths in program.get_measurement_windows().items():
+        for mw_name, begins_lengths in program.get_measurement_windows(drop=True).items():
             temp_measurement_windows[mw_name].append(begins_lengths)
 
         if set(temp_measurement_windows.keys()) - set(self._measurement_map.keys()):
diff --git a/tests/_program/loop_tests.py b/tests/_program/loop_tests.py
index 40d1c266..3bd811cd 100644
--- a/tests/_program/loop_tests.py
+++ b/tests/_program/loop_tests.py
@@ -120,6 +120,28 @@ LOOP 1 times:
         self.assertNotEqual(tree1, tree4)
         self.assertEqual(tree1, tree5)
 
+    def test_get_measurement_windows(self):
+        wf_1 = ConstantWaveform(channel='A', duration=32, amplitude=1)
+        wf_2 = ConstantWaveform(channel='A', duration=10, amplitude=2)
+
+        prog = Loop(children=[
+            Loop(waveform=wf_1, measurements=[('x', 0, 16)], repetition_count=3),
+            Loop(waveform=wf_2, measurements=[('y', 5, 5)], repetition_count=5),
+        ], repetition_count=2)
+
+        expected_measurements = {
+            'x': (np.array([0, 32, 64, 146, 178, 210]), np.array([16]*6)),
+            'y': (np.array([101, 111, 121, 131, 141, 247, 257, 267, 277, 287]), np.array([5]*10))
+        }
+        measurements_no_drop = prog.get_measurement_windows()
+        np.testing.assert_equal(expected_measurements, measurements_no_drop)
+
+        measurements_drop = prog.get_measurement_windows(drop=True)
+        np.testing.assert_equal(expected_measurements, measurements_drop)
+
+        # no measurements left
+        self.assertEqual({}, prog.get_measurement_windows())
+
     def test_repr(self):
         wf_gen = WaveformGenerator(num_channels=1)
         wfs = [wf_gen() for _ in range(11)]