Add PulseTemplate.pad_to

changes.d/801.feature

@@ -0,0 +1,2 @@
+Add ``PulseTemplate.pad_to`` method to help padding to minimal lengths or multiples of given durations.
+

qupulse/pulses/pulse_template.py

@@ -383,6 +383,47 @@ def with_appended(self, *appended: 'PulseTemplate'):
         else:
             return self
 
+    def pad_to(self, to_new_duration: Union[ExpressionLike, Callable[[Expression], ExpressionLike]],
+               pt_kwargs: Mapping[str, Any] = None) -> 'PulseTemplate':
+        """Pad this pulse template to the given duration.
+        The target duration can be numeric, symbolic or a callable that returns a new duration from the current
+        duration.
+
+        Examples:
+            # pad to a fixed duration
+            >>> padded_1 = my_pt.pad_to(1000)
+
+            # pad to a fixed sample coun
+            >>> padded_2 = my_pt.pad_to('sample_rate * 1000')
+
+            # pad to the next muliple of 16 samples with a symbolic sample rate
+            >>> padded_3 = my_pt.pad_to(to_next_multiple('sample_rate', 16))
+
+            # pad to the next muliple of 16 samples with a fixed sample rate of 1 GHz
+            >>> padded_4 = my_pt.pad_to(to_next_multiple(1, 16))
+        Args:
+            to_new_duration: Duration or callable that maps the current duration to the new duration
+            pt_kwargs: Keyword arguments for the newly created sequence pulse template.
+
+        Returns:
+            A pulse template that has the duration given by ``to_new_duration``. It can be ``self`` if the duration is
+            already as required. It is never ``self`` if ``pt_kwargs`` is non-empty.
+        """
+        from qupulse.pulses import ConstantPT, SequencePT
+        current_duration = self.duration
+        if callable(to_new_duration):
+            new_duration = to_new_duration(current_duration)
+        else:
+            new_duration = ExpressionScalar(to_new_duration)
+        pad_duration = new_duration - current_duration
+        if not pt_kwargs and pad_duration == 0:
+            return self
+        pad_pt = ConstantPT(pad_duration, self.final_values)
+        if pt_kwargs:
+            return SequencePT(self, pad_pt, **pt_kwargs)
+        else:
+            return self @ pad_pt
+
     def __format__(self, format_spec: str):
         if format_spec == '':
             format_spec = self._DEFAULT_FORMAT_SPEC

tests/pulses/pulse_template_tests.py

@@ -3,6 +3,8 @@
 from unittest import mock
 
 from typing import Optional, Dict, Set, Any, Union
+
+import frozendict
 import sympy
 
 from qupulse.parameter_scope import Scope, DictScope
@@ -23,19 +25,22 @@
 
 class PulseTemplateStub(PulseTemplate):
     """All abstract methods are stubs that raise NotImplementedError to catch unexpected calls. If a method is needed in
-    a test one should use mock.patch or mock.patch.object"""
+    a test one should use mock.patch or mock.patch.object.
+    Properties can be passed as init argument because mocking them is a pita."""
     def __init__(self, identifier=None,
                  defined_channels=None,
                  duration=None,
                  parameter_names=None,
                  measurement_names=None,
+                 final_values=None,
                  registry=None):
         super().__init__(identifier=identifier)
 
         self._defined_channels = defined_channels
         self._duration = duration
         self._parameter_names = parameter_names
         self._measurement_names = set() if measurement_names is None else measurement_names
+        self._final_values = final_values
         self.internal_create_program_args = []
         self._register(registry=registry)
 
@@ -89,7 +94,10 @@ def initial_values(self) -> Dict[ChannelID, ExpressionScalar]:
 
     @property
     def final_values(self) -> Dict[ChannelID, ExpressionScalar]:
-        raise NotImplementedError()
+        if self._final_values is None:
+            raise NotImplementedError()
+        else:
+            return self._final_values
 
 
 def get_appending_internal_create_program(waveform=DummyWaveform(),
@@ -343,6 +351,73 @@ def test_create_program_volatile(self):
 
         _internal_create_program.assert_called_once_with(**expected_internal_kwargs, parent_loop=Loop())
 
+    def test_pad_to(self):
+        from qupulse.pulses import SequencePT
+
+        def to_multiple_of_192(x: Expression) -> Expression:
+            return (x + 191) // 192 * 192
+
+        final_values = frozendict.frozendict({'A': ExpressionScalar(0.1), 'B': ExpressionScalar('a')})
+        measurements = [('M', 0, 'y')]
+
+        pt = PulseTemplateStub(duration=ExpressionScalar(10))
+        padded = pt.pad_to(10)
+        self.assertIs(pt, padded)
+
+        pt = PulseTemplateStub(duration=ExpressionScalar('duration'))
+        padded = pt.pad_to('duration')
+        self.assertIs(pt, padded)
+
+        # padding with numeric durations
+
+        pt = PulseTemplateStub(duration=ExpressionScalar(10),
+                               final_values=final_values,
+                               defined_channels=final_values.keys())
+        padded = pt.pad_to(20)
+        self.assertEqual(padded.duration, 20)
+        self.assertEqual(padded.final_values, final_values)
+        self.assertIsInstance(padded, SequencePT)
+        self.assertIs(padded.subtemplates[0], pt)
+
+        padded = pt.pad_to(20, pt_kwargs=dict(measurements=measurements))
+        self.assertEqual(padded.duration, 20)
+        self.assertEqual(padded.final_values, final_values)
+        self.assertIsInstance(padded, SequencePT)
+        self.assertIs(padded.subtemplates[0], pt)
+        self.assertEqual(measurements, padded.measurement_declarations)
+
+        padded = pt.pad_to(10, pt_kwargs=dict(measurements=measurements))
+        self.assertEqual(padded.duration, 10)
+        self.assertEqual(padded.final_values, final_values)
+        self.assertIsInstance(padded, SequencePT)
+        self.assertIs(padded.subtemplates[0], pt)
+        self.assertEqual(measurements, padded.measurement_declarations)
+
+        # padding with numeric duation and callable
+        padded = pt.pad_to(to_multiple_of_192)
+        self.assertEqual(padded.duration, 192)
+        self.assertEqual(padded.final_values, final_values)
+        self.assertIsInstance(padded, SequencePT)
+        self.assertIs(padded.subtemplates[0], pt)
+
+        # padding with symbolic durations
+
+        pt = PulseTemplateStub(duration=ExpressionScalar('duration'),
+                               final_values=final_values,
+                               defined_channels=final_values.keys())
+        padded = pt.pad_to('new_duration')
+        self.assertEqual(padded.duration, 'new_duration')
+        self.assertEqual(padded.final_values, final_values)
+        self.assertIsInstance(padded, SequencePT)
+        self.assertIs(padded.subtemplates[0], pt)
+
+        # padding symbolic durations with callable
+
+        padded = pt.pad_to(to_multiple_of_192)
+        self.assertEqual(padded.duration, '(duration + 191) // 192 * 192')
+        self.assertEqual(padded.final_values, final_values)
+        self.assertIsInstance(padded, SequencePT)
+        self.assertIs(padded.subtemplates[0], pt)
 
     def test_create_program_none(self) -> None:
         template = PulseTemplateStub(defined_channels={'A'}, parameter_names={'foo'})