Draft: rust extension interface

qupulse/_program/__init__.py

@@ -1 +1,79 @@
 """This is a private package meaning there are no stability guarantees."""
+from abc import ABC, abstractmethod
+from typing import Optional, Union, Sequence, ContextManager, Mapping
+
+import numpy as np
+
+from qupulse._program.waveforms import Waveform
+from qupulse.utils.types import MeasurementWindow, TimeType
+from qupulse._program.volatile import VolatileRepetitionCount
+
+try:
+    import qupulse_rs
+except ImportError:
+    qupulse_rs = None
+    RsProgramBuilder = None
+else:
+    from qupulse_rs.replacements import ProgramBuilder as RsProgramBuilder
+
+try:
+    from typing import Protocol, runtime_checkable
+except ImportError:
+    Protocol = object
+
+    def runtime_checkable(cls):
+        return cls
+
+
+RepetitionCount = Union[int, VolatileRepetitionCount]
+
+
+@runtime_checkable
+class Program(Protocol):
+    """This protocol is used to inspect and or manipulate programs"""
+
+    def to_single_waveform(self) -> Waveform:
+        pass
+
+    def get_measurement_windows(self) -> Mapping[str, np.ndarray]:
+        pass
+
+    @property
+    def duration(self) -> TimeType:
+        raise NotImplementedError()
+
+    def make_compatible_inplace(self):
+        # TODO: rename?
+        pass
+
+
+class ProgramBuilder(Protocol):
+    """This protocol is used by PulseTemplate to build the program."""
+
+    def append_leaf(self, waveform: Waveform,
+                    measurements: Optional[Sequence[MeasurementWindow]] = None,
+                    repetition_count: int = 1):
+        pass
+
+    def potential_child(self, measurements: Optional[Sequence[MeasurementWindow]],
+                        repetition_count: Union[VolatileRepetitionCount, int] = 1) -> ContextManager['ProgramBuilder']:
+        """
+
+        Args:
+            measurements: Measurements to attach to the potential child. Is not repeated with repetition_count.
+            repetition_count:
+
+        Returns:
+
+        """
+
+    def to_program(self) -> Optional[Program]:
+        pass
+
+
+def default_program_builder() -> ProgramBuilder:
+    if RsProgramBuilder is None:
+        from qupulse._program._loop import Loop
+        return Loop()
+    else:
+        return RsProgramBuilder()

qupulse/_program/_loop.py

@@ -1,4 +1,5 @@
-from typing import Union, Dict, Iterable, Tuple, cast, List, Optional, Generator, Mapping
+import contextlib
+from typing import Union, Dict, Iterable, Tuple, cast, List, Optional, Generator, Mapping, ContextManager, Sequence
 from collections import defaultdict
 from enum import Enum
 import warnings
@@ -15,6 +16,7 @@
 from qupulse.utils.tree import Node, is_tree_circular
 from qupulse.utils.numeric import smallest_factor_ge
 
+from qupulse._program import ProgramBuilder, Program
 from qupulse._program.waveforms import SequenceWaveform, RepetitionWaveform
 
 __all__ = ['Loop', 'make_compatible', 'MakeCompatibleWarning']
@@ -101,6 +103,9 @@ def add_measurements(self, measurements: Iterable[MeasurementWindow]):
         Args:
             measurements: Measurements to add
         """
+        warnings.warn("Loop.add_measurements is deprecated since qupulse 0.7 and will be removed in a future version.",
+                      DeprecationWarning,
+                      stacklevel=2)
         body_duration = float(self.body_duration)
         if body_duration == 0:
             measurements = measurements
@@ -198,23 +203,47 @@ def encapsulate(self) -> None:
         self._measurements = None
         self.assert_tree_integrity()
 
-    def _get_repr(self, first_prefix, other_prefixes) -> Generator[str, None, None]:
+    def __repr__(self):
+        kwargs = []
+
+        repetition_count = self._repetition_definition
+        if repetition_count != 1:
+            kwargs.append(f"repetition_count={repetition_count!r}")
+
+        waveform = self._waveform
+        if waveform:
+            kwargs.append(f"waveform={waveform!r}")
+
+        children = self.children
+        if children:
+            try:
+                kwargs.append(f"children={self._children_repr()}")
+            except RecursionError:
+                kwargs.append("children=[...]")
+
+        measurements = self._measurements
+        if measurements:
+            kwargs.append(f"measurements={measurements!r}")
+
+        return f"Loop({','.join(kwargs)})"
+
+    def _get_str(self, first_prefix, other_prefixes) -> Generator[str, None, None]:
         if self.is_leaf():
             yield '%sEXEC %r %d times' % (first_prefix, self._waveform, self.repetition_count)
         else:
             yield '%sLOOP %d times:' % (first_prefix, self.repetition_count)
 
             for elem in self:
-                yield from cast(Loop, elem)._get_repr(other_prefixes + '  ->', other_prefixes + '    ')
+                yield from cast(Loop, elem)._get_str(other_prefixes + '  ->', other_prefixes + '    ')
 
-    def __repr__(self) -> str:
+    def __str__(self) -> str:
         is_circular = is_tree_circular(self)
         if is_circular:
             return '{}: Circ {}'.format(id(self), is_circular)
 
         str_len = 0
         repr_list = []
-        for sub_repr in self._get_repr('', ''):
+        for sub_repr in self._get_str('', ''):
             str_len += len(sub_repr)
 
             if self.MAX_REPR_SIZE and str_len > self.MAX_REPR_SIZE:
@@ -404,6 +433,21 @@ def _merge_single_child(self):
         self._invalidate_duration()
         return True
 
+    @contextlib.contextmanager
+    def potential_child(self,
+                        measurements: Optional[List[MeasurementWindow]],
+                        repetition_count: Union[VolatileRepetitionCount, int] = 1) -> ContextManager['Loop']:
+        if repetition_count != 1 and measurements:
+            # current design requires an extra level of nesting here because the measurements are NOT to be repeated
+            # with the repetition count
+            inner_child = Loop(repetition_count=repetition_count)
+            child = Loop(measurements=measurements, children=[inner_child])
+        else:
+            inner_child = child = Loop(measurements=measurements, repetition_count=repetition_count)
+        yield inner_child
+        if inner_child.waveform or len(inner_child):
+            self.append_child(child)
+
     def cleanup(self, actions=('remove_empty_loops', 'merge_single_child')):
         """Apply the specified actions to cleanup the Loop.
 
@@ -451,6 +495,32 @@ def get_duration_structure(self) -> Tuple[int, Union[TimeType, tuple]]:
         else:
             return self.repetition_count, tuple(child.get_duration_structure() for child in self)
 
+    def to_single_waveform(self) -> Waveform:
+        if self.is_leaf():
+            if self.repetition_count == 1:
+                return self.waveform
+            else:
+                return RepetitionWaveform.from_repetition_count(self.waveform, self.repetition_count)
+        else:
+            if len(self) == 1:
+                sequenced_waveform = to_waveform(cast(Loop, self[0]))
+            else:
+                sequenced_waveform = SequenceWaveform.from_sequence([to_waveform(cast(Loop, sub_program))
+                                                                     for sub_program in self])
+            if self.repetition_count > 1:
+                return RepetitionWaveform.from_repetition_count(sequenced_waveform, self.repetition_count)
+            else:
+                return sequenced_waveform
+
+    def append_leaf(self, waveform: Waveform,
+                    measurements: Optional[Sequence[MeasurementWindow]] = None,
+                    repetition_count: int = 1):
+        self.append_child(waveform=waveform, measurements=measurements, repetition_count=repetition_count)
+
+    def to_program(self) -> Optional['Loop']:
+        if self.waveform or self.children:
+            return self
+
     def reverse_inplace(self):
         if self.is_leaf():
             self._waveform = self._waveform.reversed()
@@ -465,29 +535,45 @@ def reverse_inplace(self):
                 for name, begin, length in self._measurements
             ]
 
+    def make_compatible_inplace(self, minimal_waveform_length: int, waveform_quantum: int, sample_rate: TimeType):
+        program = self
+        comp_level = _is_compatible(program,
+                                    min_len=minimal_waveform_length,
+                                    quantum=waveform_quantum,
+                                    sample_rate=sample_rate)
+        if comp_level == _CompatibilityLevel.incompatible_fraction:
+            raise ValueError(
+                'The program duration in samples {} is not an integer'.format(program.duration * sample_rate))
+        if comp_level == _CompatibilityLevel.incompatible_too_short:
+            raise ValueError('The program is too short to be a valid waveform. \n'
+                             ' program duration in samples: {} \n'
+                             ' minimal length: {}'.format(program.duration * sample_rate, minimal_waveform_length))
+        if comp_level == _CompatibilityLevel.incompatible_quantum:
+            raise ValueError('The program duration in samples {} '
+                             'is not a multiple of quantum {}'.format(program.duration * sample_rate, waveform_quantum))
+
+        elif comp_level == _CompatibilityLevel.action_required:
+            warnings.warn(
+                "qupulse will now concatenate waveforms to make the pulse/program compatible with the chosen AWG."
+                " This might take some time. If you need this pulse more often it makes sense to write it in a "
+                "way which is more AWG friendly.", MakeCompatibleWarning)
+
+            _make_compatible(program,
+                             min_len=minimal_waveform_length,
+                             quantum=waveform_quantum,
+                             sample_rate=sample_rate)
+
+        else:
+            assert comp_level == _CompatibilityLevel.compatible
+
 
 class ChannelSplit(Exception):
     def __init__(self, channel_sets):
         self.channel_sets = channel_sets
 
 
 def to_waveform(program: Loop) -> Waveform:
-    if program.is_leaf():
-        if program.repetition_count == 1:
-            return program.waveform
-        else:
-            return RepetitionWaveform.from_repetition_count(program.waveform, program.repetition_count)
-    else:
-        if len(program) == 1:
-            sequenced_waveform = to_waveform(cast(Loop, program[0]))
-        else:
-            sequenced_waveform = SequenceWaveform.from_sequence(
-                [to_waveform(cast(Loop, sub_program))
-                 for sub_program in program])
-        if program.repetition_count > 1:
-            return RepetitionWaveform.from_repetition_count(sequenced_waveform, program.repetition_count)
-        else:
-            return sequenced_waveform
+    return program.to_single_waveform()
 
 
 class _CompatibilityLevel(Enum):
@@ -568,32 +654,7 @@ def _make_compatible(program: Loop, min_len: int, quantum: int, sample_rate: Tim
 
 def make_compatible(program: Loop, minimal_waveform_length: int, waveform_quantum: int, sample_rate: TimeType):
     """ check program for compatibility to AWG requirements, make it compatible if necessary and  possible"""
-    comp_level = _is_compatible(program,
-                                min_len=minimal_waveform_length,
-                                quantum=waveform_quantum,
-                                sample_rate=sample_rate)
-    if comp_level == _CompatibilityLevel.incompatible_fraction:
-        raise ValueError('The program duration in samples {} is not an integer'.format(program.duration * sample_rate))
-    if comp_level == _CompatibilityLevel.incompatible_too_short:
-        raise ValueError('The program is too short to be a valid waveform. \n'
-                         ' program duration in samples: {} \n'
-                         ' minimal length: {}'.format(program.duration * sample_rate, minimal_waveform_length))
-    if comp_level == _CompatibilityLevel.incompatible_quantum:
-        raise ValueError('The program duration in samples {} '
-                         'is not a multiple of quantum {}'.format(program.duration * sample_rate, waveform_quantum))
-
-    elif comp_level == _CompatibilityLevel.action_required:
-        warnings.warn("qupulse will now concatenate waveforms to make the pulse/program compatible with the chosen AWG."
-                      " This might take some time. If you need this pulse more often it makes sense to write it in a "
-                      "way which is more AWG friendly.", MakeCompatibleWarning)
-
-        _make_compatible(program,
-                         min_len=minimal_waveform_length,
-                         quantum=waveform_quantum,
-                         sample_rate=sample_rate)
-
-    else:
-        assert comp_level == _CompatibilityLevel.compatible
+    program.make_compatible_inplace(minimal_waveform_length, waveform_quantum, sample_rate)
 
 
 def roll_constant_waveforms(program: Loop, minimal_waveform_quanta: int, waveform_quantum: int, sample_rate: TimeType):

qupulse/_program/transformation.py

@@ -6,7 +6,15 @@
 
 from qupulse import ChannelID
 from qupulse.comparable import Comparable
-from qupulse.utils.types import SingletonABCMeta
+from qupulse.utils.types import SingletonABCMeta, use_rs_replacements
+
+try:
+    import qupulse_rs
+except ImportError:
+    qupulse_rs = None
+    transformation_rs = None
+else:
+    from qupulse_rs.replacements import transformation as transformation_rs
 
 
 class Transformation(Comparable):
@@ -325,4 +333,13 @@ def chain_transformations(*transformations: Transformation) -> Transformation:
     elif len(parsed_transformations) == 1:
         return parsed_transformations[0]
     else:
-        return ChainedTransformation(*parsed_transformations)
+        return ChainedTransformation(*parsed_transformations)
+
+
+
+if transformation_rs:
+    use_rs_replacements(globals(), transformation_rs, Transformation)
+
+    py_chain_transformations = chain_transformations
+    rs_chain_transformations = transformation_rs.chain_transformations
+    chain_transformations = rs_chain_transformations

qupulse/_program/waveforms.py

@@ -25,10 +25,18 @@
 from qupulse.expressions import ExpressionScalar
 from qupulse.pulses.interpolation import InterpolationStrategy
 from qupulse.utils import checked_int_cast, isclose
-from qupulse.utils.types import TimeType, time_from_float, FrozenDict
+from qupulse.utils.types import TimeType, time_from_float, FrozenDict, use_rs_replacements
 from qupulse._program.transformation import Transformation
 from qupulse.utils import pairwise
 
+try:
+    import qupulse_rs
+except ImportError:
+    qupulse_rs = None
+    waveforms_rs = None
+else:
+    from qupulse_rs.replacements import waveforms as waveforms_rs
+
 class ConstantFunctionPulseTemplateWarning(UserWarning):
     """  This warning indicates a constant waveform is constructed from a FunctionPulseTemplate """
     pass
@@ -850,8 +858,8 @@ def unsafe_sample(self,
         return output_array
 
     @property
-    def compare_key(self) -> Tuple[Any, int]:
-        return self._body.compare_key, self._repetition_count
+    def compare_key(self) -> Tuple[int, Any]:
+        return self._repetition_count, self._body
 
     def unsafe_get_subset_for_channels(self, channels: AbstractSet[ChannelID]) -> Waveform:
         return RepetitionWaveform.from_repetition_count(
@@ -987,9 +995,14 @@ def unsafe_sample(self,
         return self.inner_waveform.unsafe_sample(channel, sample_times, output_array)
 
     def constant_value_dict(self) -> Optional[Mapping[ChannelID, float]]:
-        d = self._inner_waveform.constant_value_dict()
-        if d is not None:
-            return {ch: d[ch] for ch in self._channel_subset}
+        constant_values = {}
+        for ch in self.defined_channels:
+            value = self._inner_waveform.constant_value(ch)
+            if value is None:
+                return
+            else:
+                constant_values[ch] = value
+        return constant_values
 
     def constant_value(self, channel: ChannelID) -> Optional[float]:
         if channel not in self._channel_subset:
@@ -1231,3 +1244,9 @@ def compare_key(self) -> Hashable:
 
     def reversed(self) -> 'Waveform':
         return self._inner
+
+
+
+
+if waveforms_rs:
+    use_rs_replacements(globals(), waveforms_rs, Waveform)