Repository: qutech/qupulse
Issue: Measurements defined on subtemplates of MultiChannelPTs are ignored

Only `build_waveform` is called by the AtomicPTs `build_sequence` method.

Changed files:
tests/pulses/sequencing_dummies.py, tests/pulses/multi_channel_pulse_template_tests.py, qctoolkit/pulses/multi_channel_pulse_template.py

diff --git a/qctoolkit/pulses/multi_channel_pulse_template.py b/qctoolkit/pulses/multi_channel_pulse_template.py
index d25014e4..2665754b 100644
--- a/qctoolkit/pulses/multi_channel_pulse_template.py
+++ b/qctoolkit/pulses/multi_channel_pulse_template.py
@@ -21,7 +21,7 @@ from qctoolkit._program.waveforms import MultiChannelWaveform
 from qctoolkit.pulses.pulse_template import PulseTemplate, AtomicPulseTemplate
 from qctoolkit.pulses.mapping_pulse_template import MappingPulseTemplate, MappingTuple
 from qctoolkit.pulses.parameters import Parameter, ParameterConstrainer
-from qctoolkit.pulses.measurement import MeasurementDeclaration
+from qctoolkit.pulses.measurement import MeasurementDeclaration, MeasurementWindow
 from qctoolkit.expressions import Expression, ExpressionScalar
 
 __all__ = ["AtomicMultiChannelPulseTemplate"]
@@ -117,6 +117,16 @@ class AtomicMultiChannelPulseTemplate(AtomicPulseTemplate, ParameterConstrainer)
         else:
             return MultiChannelWaveform(sub_waveforms)
 
+    def get_measurement_windows(self,
+                                parameters: Dict[str, numbers.Real],
+                                measurement_mapping: Dict[str, Optional[str]]) -> List[MeasurementWindow]:
+        measurements = super().get_measurement_windows(parameters=parameters,
+                                                       measurement_mapping=measurement_mapping)
+        for st in self.subtemplates:
+            measurements.extend(st.get_measurement_windows(parameters=parameters,
+                                                           measurement_mapping=measurement_mapping))
+        return measurements
+
     def requires_stop(self,
                       parameters: Dict[str, Parameter],
                       conditions: Dict[str, 'Condition']) -> bool:
diff --git a/tests/pulses/multi_channel_pulse_template_tests.py b/tests/pulses/multi_channel_pulse_template_tests.py
index a17d5e81..0ef37990 100644
--- a/tests/pulses/multi_channel_pulse_template_tests.py
+++ b/tests/pulses/multi_channel_pulse_template_tests.py
@@ -5,8 +5,9 @@ import numpy
 
 from qctoolkit.utils.types import time_from_float
 from qctoolkit.pulses.multi_channel_pulse_template import MultiChannelWaveform, MappingPulseTemplate, ChannelMappingException, AtomicMultiChannelPulseTemplate
-from qctoolkit.pulses.parameters import ParameterConstraint, ParameterConstraintViolation
+from qctoolkit.pulses.parameters import ParameterConstraint, ParameterConstraintViolation, ConstantParameter
 from qctoolkit.expressions import ExpressionScalar, Expression
+from qctoolkit._program.instructions import InstructionBlock
 
 from tests.pulses.sequencing_dummies import DummyPulseTemplate, DummyWaveform
 from tests.serialization_dummies import DummySerializer
@@ -198,6 +199,46 @@ class MultiChannelPulseTemplateSequencingTests(unittest.TestCase):
         wf = pt.build_waveform(parameters, channel_mapping=channel_mapping)
         self.assertIsNone(wf)
 
+    def test_build_sequence(self):
+        wfs = [DummyWaveform(duration=1.1, defined_channels={'A'}), DummyWaveform(duration=1.1, defined_channels={'B'})]
+        sts = [DummyPulseTemplate(duration='t1', defined_channels={'A'}, waveform=wfs[0], measurements=[('m', 0, 1)]),
+               DummyPulseTemplate(duration='t1', defined_channels={'B'}, waveform=wfs[1]),
+               DummyPulseTemplate(duration='t1', defined_channels={'C'}, waveform=None)]
+
+        pt = AtomicMultiChannelPulseTemplate(*sts, parameter_constraints=['a < b'], measurements=[('n', .1, .2)])
+
+        params = dict(a=ConstantParameter(1.0), b=ConstantParameter(1.1))
+        measurement_mapping = dict(m='foo', n='bar')
+        channel_mapping = {'A': 'A', 'B': 'B', 'C': None}
+
+        block = InstructionBlock()
+        pt.build_sequence(None, parameters=params, conditions={}, measurement_mapping=measurement_mapping,
+                          channel_mapping=channel_mapping, instruction_block=block)
+
+        expected_waveform = MultiChannelWaveform(wfs)
+
+        expected_block = InstructionBlock()
+        measurements = [('bar', .1, .2), ('foo', 0, 1)]
+        expected_block.add_instruction_meas(measurements)
+        expected_block.add_instruction_exec(waveform=expected_waveform)
+
+        self.assertEqual(len(block.instructions), len(expected_block.instructions))
+        self.assertEqual(block.instructions[0].compare_key, expected_block.instructions[0].compare_key)
+        self.assertEqual(block.instructions[1].compare_key, expected_block.instructions[1].compare_key)
+
+    def test_get_measurement_windows(self):
+        wfs = [DummyWaveform(duration=1.1, defined_channels={'A'}), DummyWaveform(duration=1.1, defined_channels={'B'})]
+        sts = [DummyPulseTemplate(duration='t1', defined_channels={'A'}, waveform=wfs[0], measurements=[('m', 0, 1),
+                                                                                                        ('n', 0.3, 0.4)]),
+               DummyPulseTemplate(duration='t1', defined_channels={'B'}, waveform=wfs[1], measurements=[('m', 0.1, .2)])]
+
+        pt = AtomicMultiChannelPulseTemplate(*sts, parameter_constraints=['a < b'], measurements=[('n', .1, .2)])
+
+        measurement_mapping = dict(m='foo', n='bar')
+        expected = [('bar', .1, .2), ('foo', 0, 1), ('bar', .3, .4), ('foo', .1, .2)]
+        meas_windows = pt.get_measurement_windows({}, measurement_mapping)
+        self.assertEqual(expected, meas_windows)
+
 
 class AtomicMultiChannelPulseTemplateSerializationTests(SerializableTests, unittest.TestCase):
 
diff --git a/tests/pulses/sequencing_dummies.py b/tests/pulses/sequencing_dummies.py
index 7f5dd355..f0ddeb40 100644
--- a/tests/pulses/sequencing_dummies.py
+++ b/tests/pulses/sequencing_dummies.py
@@ -327,10 +327,6 @@ class DummyPulseTemplate(AtomicPulseTemplate):
     def parameter_names(self) -> Set[str]:
         return set(self.parameter_names_)
 
-    def get_measurement_windows(self, parameters: Dict[str, Parameter] = None) -> List[MeasurementWindow]:
-        """Return all measurement windows defined in this PulseTemplate."""
-        raise NotImplementedError()
-
     @property
     def build_sequence_calls(self):
         return len(self.build_sequence_arguments)