iqm-pulse 9.21.0__py3-none-any.whl → 10.1.0__py3-none-any.whl

iqm/pulse/gate_implementation.py

@@ -25,10 +25,10 @@
 from __future__ import annotations
 
 import abc
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping
 import dataclasses
 import re
-from typing import TYPE_CHECKING, Any, TypeAlias, get_type_hints
+from typing import TYPE_CHECKING, Any, TypeAlias, final, get_type_hints
 from uuid import uuid4
 
 import numpy as np
@@ -36,7 +36,7 @@ import numpy as np
 from exa.common.data.parameter import Parameter, Setting
 from exa.common.data.setting_node import SettingNode
 from exa.common.qcm_data.chip_topology import ChipTopology
-from iqm.pulse.playlist.instructions import Instruction
+from iqm.pulse.playlist.instructions import IQPulse
 from iqm.pulse.playlist.schedule import Schedule
 from iqm.pulse.timebox import TimeBox
 from iqm.pulse.utils import map_waveform_param_types
@@ -81,23 +81,31 @@ PROBE_LINES_LOCUS_MAPPING = "probe_lines"
 class GateImplementation(abc.ABC):
     """ABC for implementing quantum gates and other quantum operations using instruction schedules.
 
-    There is a separate GateImplementation subclass for every implementation of every operation type.
+    Every implementation of every operation type can have its own GateImplementation subclass.
     Each GateImplementation instance represents a particular locus for that implementation, and encapsulates
     the calibration data it requires.
 
-    All GateImplementation subclasses ``__init__`` must have exactly the below arguments in order to be
+    All GateImplementation subclasses :meth:`__init__` must have exactly the below arguments in order to be
     usable via :meth:`.ScheduleBuilder.get_implementation`.
 
-    GateImplementations also have the :meth:`~GateImplementation.__call__` method, which takes the operation parameters
-    (e.g. rotation angles) as input, and returns a :class:`.TimeBox` instance which implements
-    an instance of the operation at that locus.
+    GateImplementation also has the :meth:`__call__` method, which takes the operation parameters
+    (e.g. rotation angles) as input, and builds, caches and then returns a :class:`.TimeBox` instance which
+    implements an instance of the operation at that locus. :meth:`__call__` normally should not be
+    reimplemented by the subclasses, instead they should define :meth:`_call` which contains the actual
+    implementation without the caching logic.
+
+    Even though it is possible, GateImplementations *should not* use other gates to implement themselves,
+    unless they are subclasses of :class:`CompositeGate`. This is to encapsulate the calibration data better,
+    and avoid unpredictable dependencies between gates.
 
     Args:
         parent: Quantum operation this instance implements.
         name: Name of the implementation provided by this instance.
         locus: Locus the operation acts on.
-        calibration_data: (Raw) calibration data for the (operation, implementation, locus) represented by this instance
-        builder: Schedule builder.
+        calibration_data: (Raw) calibration data for the (operation, implementation, locus) represented
+            by this instance.
+        builder: Schedule builder instance that can be used to access system properties (and in the
+            case of :class:`CompositeGate`, other gates).
 
     """
 
@@ -120,8 +128,13 @@ class GateImplementation(abc.ABC):
         self.locus = locus
         self.calibration_data = calibration_data
         self.builder = builder
-        self.id = str(uuid4())  # unique str identifier, needed for certain caching properties
+        self.id = str(uuid4())
+        """Unique str identifier, needed for certain caching properties."""
+        self.sub_implementations: dict[str, GateImplementation] = {}
+        """Single-component sub-implementations for factorizable gates with len(locus) > 1, otherwise empty.
+        At least one of self.sub_implementations and self.calibration_data is always empty."""
         self._timebox_cache: dict[tuple[Any, ...], TimeBox | list[TimeBox]] = {}
+        """Cache for :meth:`__call__` results."""
 
     @property
     def qualified_name(self) -> str:
@@ -130,13 +143,38 @@ class GateImplementation(abc.ABC):
 
     @classmethod
     def needs_calibration(cls) -> bool:
-        """Whether the implementation needs calibration data
+        """True iff the implementation needs calibration data.
+
+        Returns:
+            True iff :attr:`OpCalibrationDataTree` must contain a node ``f"{self.parent.name}.{self.name}.{self.locus}``
+            for the calibration data this instance needs.
 
-        Returns True if the calibration dict must contain a node with keyed with
-        <operation name>: <implementation name>: <appropriate locus> in order to use this implementation.
         """
         return bool(cls.parameters)
 
+    @classmethod
+    def optional_calibration_keys(cls) -> tuple[str, ...]:
+        """Optional calibration data keys for this class, in addition to the required items in :attr:`parameters`.
+
+        These keys are not required to be present in :attr:`OILCalibrationData` when validating it.
+
+        Returns:
+            Optional top-level calibration data keys.
+
+        """
+        return ()
+
+    @classmethod
+    def may_have_calibration(cls) -> bool:
+        """True iff the implementation may have calibration data.
+
+        Returns:
+            True iff :attr:`OpCalibrationDataTree` may contain a node ``"{self.parent.name}.{self.name}.{self.locus}``
+            for the calibration data of this instance.
+
+        """
+        return cls.needs_calibration() or bool(cls.optional_calibration_keys())
+
     def __call__(self, *args, **kwargs) -> TimeBox | list[TimeBox]:
         """TimeBox that implements a quantum operation.
 
@@ -171,26 +209,32 @@ class GateImplementation(abc.ABC):
         """
         return NotImplementedError  # type: ignore[return-value]
 
-    def build(
-        self, op_name: str, locus: Locus, impl_name: str | None = None, strict_locus: bool = False
+    @final
+    @classmethod
+    def construct_factorizable(
+        cls,
+        parent: QuantumOp,
+        name: str,
+        locus: Locus,
+        builder: ScheduleBuilder,
+        sub_implementations: dict[str, GateImplementation],
     ) -> GateImplementation:
-        """Utility method for constructing a GateImplementation with ``self.builder``.
-
-        Inheriting classes may override this in order to add additional logic.
-
-        Args:
-            op_name: operation name
-            locus: locus the operation acts on
-            impl_name: implementation name. Uses the assigned default implementation if not specified.
-            strict_locus: iff False, for non-symmetric implementations of symmetric ops the locus order may
-                be changed if no calibration data is available for the requested locus order
-
-        Returns:
-            Calibrated gate implementation.
+        """Construct an implementation for a factorizable operation.
 
+        Instead of calibration data this method is given ``sub_implementations``, which contains single-qubit
+        implementations for all the components in ``locus``.
         """
-        return self.builder.get_implementation(op_name, locus, impl_name, strict_locus=strict_locus)
+        impl = cls(
+            parent=parent,
+            name=name,
+            locus=locus,
+            calibration_data={},
+            builder=builder,
+        )
+        impl.sub_implementations = sub_implementations
+        return impl
 
+    @final
     def to_timebox(self, schedule: Schedule) -> TimeBox:
         """Wraps the given instruction schedule into an atomic/resolved timebox."""
         return TimeBox.atomic(
@@ -207,22 +251,30 @@ class GateImplementation(abc.ABC):
         """
         raise NotImplementedError
 
+    @final
     @classmethod
     def convert_calibration_data(
         cls,
         calibration_data: OILCalibrationData,
         params: NestedParams,
         channel_props: ChannelProperties,
+        *,
         duration: float | None = None,
+        _top_level: bool = True,
     ) -> OILCalibrationData:
-        """Convert time-like items in the calibration data to fractions of the time duration of the gate.
+        """Convert time- and frequency-like items in the calibration data to fractions of the time duration of the gate.
 
         This is a convenience method for converting calibration data items involving time
-        durations measured in seconds into fractions of the duration of the gate.
+        durations measured in seconds and frequencies measured in Hz into fractions of the duration
+        of the gate, e.g. to be used to parameterize :class:`.Waveform` classes.
 
         * Values of items that are not measured in seconds or Hz are returned as is.
-        * Additionally, converts ``duration`` to channel samples and adds it in the converted
-          calibration data under the key ``"n_samples"``, while the original ``"duration"`` key is removed.
+        * Items named ``duration`` get special treatment for convenience.
+          They are not included in the converted data.
+          If the ``duration`` parameter is None, there must be a ``"duration"`` item at the top level in
+          ``calibration_data`` whose value will be used instead.
+        * The ``duration`` parameter is converted to channel samples and included in the converted
+          data under the top-level key ``"n_samples"``.
 
         Args:
             calibration_data: (subset of) calibration data for the gate/implementation/locus
@@ -243,7 +295,7 @@ class GateImplementation(abc.ABC):
             """
             if value is None:
                 return None
-            if (unit not in {"s", "Hz"}) or name == "duration":
+            if unit not in {"s", "Hz"}:
                 return value
             if unit == "s":
 
@@ -261,27 +313,30 @@ class GateImplementation(abc.ABC):
             return conversion(value)
 
         if duration is None:
-            # duration should be found in the outermost dict
+            # if not given, duration should be found in the outermost dict
             duration = calibration_data["duration"]
             if duration is None:
                 raise ValueError(f"Duration for {cls.__name__} has not been set.")
-            converted = {"n_samples": channel_props.duration_to_int_samples(duration) if duration > 0 else 0}
-        else:
-            converted = {}
+
+        # n_samples will only be included on the top level
+        converted = (
+            {"n_samples": channel_props.duration_to_int_samples(duration) if duration > 0 else 0} if _top_level else {}
+        )
 
         for p_name, p in params.items():
-            if p_name in calibration_data:
+            if p_name in calibration_data and p_name != "duration":
+                # duration is not included in the converted data
                 data = calibration_data[p_name]
                 if isinstance(p, Setting | Parameter):
                     value = convert(p_name, p.unit, data, duration)
                 else:
                     # recursion for nested parameter dicts
-                    value = cls.convert_calibration_data(data, p, channel_props, duration=duration)
-                if p_name != "duration":  # duration is converted to n_samples so we don't want it in seconds here
-                    converted[p_name] = value
+                    value = cls.convert_calibration_data(data, p, channel_props, duration=duration, _top_level=False)
+                converted[p_name] = value
 
         return converted
 
+    @final
     @classmethod
     def get_parameters(cls, locus: Iterable[str], path: Iterable[str] = ()) -> SettingNode:
         """Calibration data tree the GateImplementation subclass expects for each locus.
@@ -342,7 +397,7 @@ class GateImplementation(abc.ABC):
 
     @classmethod
     def get_custom_locus_mapping(
-        cls, chip_topology: ChipTopology, component_to_channels: dict[str, Iterable[str]]
+        cls, chip_topology: ChipTopology, component_to_channels: Mapping[str, Iterable[str]]
     ) -> dict[tuple[str, ...] | frozenset[str], tuple[str, ...]] | None:
         """Get custom locus mapping for this GateImplementation.
 
@@ -393,6 +448,7 @@ class CustomIQWaveforms(GateImplementation):
         self, parent: QuantumOp, name: str, locus: Locus, calibration_data: OILCalibrationData, builder: ScheduleBuilder
     ) -> None:
         super().__init__(parent, name, locus, calibration_data, builder)
+        # TODO why does not this check happen in __init_subclass__?
         if getattr(self, "wave_i", None) is None or getattr(self, "wave_q", None) is None:
             raise ValueError(
                 "You must provide valid Waveforms for both of the arguments `wave_i` and `wave_q` when inheriting"
@@ -532,7 +588,8 @@ class SinglePulseGate(GateImplementation):
         """
         return self.builder.get_drive_channel(self.locus[0])
 
-    def _get_pulse(self, **kwargs) -> Instruction:
+    @classmethod
+    def _get_pulse(cls, **kwargs) -> IQPulse:
         """Return pulse based on the provided calibration data."""
         raise NotImplementedError
 
@@ -543,32 +600,70 @@ class SinglePulseGate(GateImplementation):
 
 
 class CompositeGate(GateImplementation):
-    """Utility base class for creating gate implementations that are defined in terms of other gate implementations.
-
-    Gates can be implemented using other pre-existing gate implementations by just utilizing the ScheduleBuilder in
-    :attr:`~GateImplementation.builder` in the :meth:`__call__` method (e.g. by calling
-    ``self.builder.get_implementation(<some gate>, <some locus>)``. In this way, any such "member gates" will use
-    the common calibration that exists in :attr:`~GateImplementation.builder`. In order for a composite gate
-    implementation to be able to calibrate its member gates with different calibration values from the common
-    calibration, it needs to know what gates it considers as its "members". This is what the CompositeGate ABC is for.
-
-    Inheriting from this class and defining e.g. ``registered_gates = ["prx", "cz"]`` allows one to calibrate the
-    member operations (i.e. ``"prx"`` and ``"cz"`` in this example) inside this composite gate differently from the
-    common calibration. However, if no specific calibration data is provided, the gate implementation will be calibrated
-    with the common calibration.
+    """Base class for gate implementations that are defined in terms of other gate implementations.
+
+    Composite gates can be implemented using other pre-existing gate implementations (called its
+    *member gates*) by using :meth:`build` in the :meth:`_call` method. You *should not* call
+    :meth:`ScheduleBuilder.get_implementation` directly in composite gate code.
+
+    A CompositeGate subclass needs to declare what its member gates are, e.g. to be able to
+    verify that they are calibrated, using the :attr:`registered_gates` class attribute.
+
+    It is possibe to calibrate (some of) the member gates separately from the common calibration,
+    by listing their names in :attr:`customizable_gates` class attribute.
+    However, if no custom calibration data is provided, the composite gate will use
+    the common calibration for the member operations.
+
+    .. example::
+
+       Inheriting this class and defining ``registered_gates = ("prx", "cz")``, ``customizable_gates = ("prx",)``
+       allows one to use ``prx`` and ``cz`` gates as member operations, and calibrate ``prx`` independently of
+       the common calibration.
+
+    .. note::
+
+       :meth:`CompositeGate.needs_calibration` only tells whether the implementation class itself needs
+       calibration data, not whether the member gates need some.
+
+    """
+
+    registered_gates: tuple[str, ...] = ()
+    """Names of the member operations used by the composite gate.
+    There must be corresponding keys in :attr:`builder.op_table`.
+    """
+
+    customizable_gates: tuple[str, ...] | None = None
+    """These member operations can be calibrated separately from their common
+    calibration by adding :attr:`OCalibrationData` nodes for them under the
+    :attr:`OILCalibrationData` node of the composite gate.
+    Must be a subset of :attr:`registered_gates`.
+    By default all member operations are customizable.
     """
 
-    registered_gates: list[str] = []
-    """Gates that can be calibrated separately from their common calibration existing in ``self.builder``. The gate
-    names should correspond to the keys in ``self.builder.op_table``. Other gates besides the ones given here can
-    also be constructed via ``self.builder``, but these will always use the common calibration."""
     default_implementations: dict[str, str] = {}
-    """Mapping from operation names to the designated default implementation of that operation. Filling this attribute
-    allows one to define a different default implementation from the common default in ``self.builder.op_table`` to
-    be used in he context of this composite gate. If an operation is not found in this dict as a key, this
-    CompositeGate will use the common default as the default implementation for it.
+    """Mapping from member operation names to the designated default implementation of that
+    operation. Filling this attribute allows one to define a different default implementation from
+    the common default in :attr:`builder.op_table` to be used in the context of this composite
+    gate. If a member operation is not found in this dict as a key, the CompositeGate will use the
+    common default as its default implementation.
     """
 
+    def __init_subclass__(cls):
+        if not cls.registered_gates:
+            # this would be pointless
+            raise ValueError(f"CompositeGate {cls.__name__} has no registered gates.")
+        # TODO we should also check that customizable_gates may_have_calibration (otherwise it's pointless
+        # to call them customizable), but we don't currently have access to their implementation classes here...
+        if cls.customizable_gates is None:
+            cls.customizable_gates = cls.registered_gates
+        elif not set(cls.customizable_gates) <= set(cls.registered_gates):
+            raise ValueError(f"CompositeGate {cls.__name__}: customizable_gates must be a subset of registered_gates.")
+
+    @classmethod
+    def optional_calibration_keys(cls) -> tuple[str, ...]:
+        # Custom calibration data for member gates is optional. Affects may_have_calibration.
+        return cls.customizable_gates or ()
+
     def __init__(
         self,
         parent: QuantumOp,
@@ -577,14 +672,18 @@ class CompositeGate(GateImplementation):
         calibration_data: OILCalibrationData,
         builder: ScheduleBuilder,
     ) -> None:
+        # validate the registered gates
+        for op_name in self.registered_gates:
+            if (op := builder.op_table.get(op_name)) is None:
+                raise ValueError(f"Unknown registered gate '{op_name}'.")
+            if parent.factorizable:
+                if not (op.factorizable or op.arity == 1):
+                    raise ValueError(
+                        f"'{parent.name}' is factorizable, but registered gate '{op_name}'"
+                        " is neither factorizable nor arity-1."
+                    )
+
         super().__init__(parent, name, locus, calibration_data, builder)
-        custom_defaults = {
-            op: data.get("default_implementation") for op, data in calibration_data.items() if isinstance(data, dict)
-        }
-        self._default_implementations = self.default_implementations.copy()
-        for op_name, impl in custom_defaults.items():
-            if op_name in self.registered_gates and impl is not None:
-                self._default_implementations[op_name] = impl
 
     def __call__(self, *args, **kwargs):
         default_cache_key = tuple(args) + tuple(kwargs.items())
@@ -601,46 +700,85 @@ class CompositeGate(GateImplementation):
         return box
 
     def build(
-        self, op_name: str, locus: Locus, impl_name: str | None = None, strict_locus: bool = False
+        self,
+        op_name: str,
+        locus: Locus,
+        impl_name: str | None = None,
+        *,
+        strict_locus: bool = False,
+        priority_calibration: OILCalibrationData | None = None,
     ) -> GateImplementation:
-        """Construct a member gate implementation.
+        """Construct an implementation for a member (registered) gate.
 
-        If the gate ``op_name`` is registered, a specific calibration for it in the context of this CompositeGate
-        will be sought for from ``self.builder.calibration``. If any (non-empty) calibration values are found in
-        ``self.builder.calibration[self.name][op_name][<impl_name>]`` they will be merged to the common calibration
-        (only non-empty values will be merged). If there are no values found, the
-        common calibration will be used.
+        A custom calibration for ``op_name`` will be sought in :attr:`calibration_data`.
+        If any calibration parameters are found, they override the corresponding parameters
+        in the common calibration data.
 
         Args:
-            op_name: operation name
+            op_name: member operation name
             locus: locus the operation acts on
             impl_name: Implementation name. If not given, uses the default implementation defined in the class instance
-                if any, and otherwise the common default in ``self.builder.op_table``
+                if any, and otherwise the common default in :attr:`builder.op_table`.
             strict_locus: iff False, for non-symmetric implementations of symmetric ops the locus order may
                 be changed if no calibration data is available for the requested locus order
+            priority_calibration: If given, overrides the custom calibration for the member gate. Deprecated,
+                should not be used.
 
         Returns:
-            Calibrated gate implementation.
+            Calibrated member gate implementation.
 
         """
-        impl_name = impl_name or self._default_implementations.get(op_name, None)
-        composite_calibration = None
-        if op_name in self.registered_gates:
-            # if op is registered and needs calibration, find the cal data under the CompositeGate
-            impl_class = self.builder.get_implementation_class(op_name, impl_name)
-            if impl_class.needs_calibration() or (
-                isinstance(impl_class, CompositeGate) and impl_class.registered_gates
-            ):
-                try:
-                    composite_calibration = self.builder.get_calibration(self.parent.name, self.name, self.locus)
-                except ValueError:
-                    pass
+        # FIXME remove priority_calibration, it's a HACK to make a single current use case work
+        # (MOVE_NCZ_MOVE_Composite in exa-core).
+        if op_name not in self.registered_gates:
+            raise ValueError(f"'{op_name}' not found in registered_gates.")
+
+        op = self.builder.op_table[op_name]
+
+        # implementation to use: given or class default
+        impl_name = impl_name or self.default_implementations.get(op_name, None)
+        # or, finally, the global default
+        impl_name, locus = self.builder._find_implementation_and_locus(
+            op,
+            impl_name=impl_name,
+            locus=locus,
+            strict_locus=strict_locus,
+        )
+
+        def get_custom_oi(cal_impl: GateImplementation) -> OICalibrationData:
+            """Return the custom calibration data node for the requested member op/implementation in
+            the calibration data tree of the given implementation, or an empty dict if it does not exist.
+            """
+            return cal_impl.calibration_data.get(op_name, {}).get(impl_name, {})
+
+        if priority_calibration is None:
+            # Find the custom cal data for the member op (if allowed and present).
+            if op_name in (self.customizable_gates or ()):
+                impl_class = self.builder.get_implementation_class(op_name, impl_name)
+                if impl_class.may_have_calibration():
+                    if self.sub_implementations:
+                        # self is a len(locus) > 1 factorizable CompositeGate (e.g. reset).
+                        # It has 1-qubit subimplementations that have their own cal data.
+                        # It may only have factorizable or arity-1 member ops.
+                        if op.factorizable:
+                            # Combine the custom cal datas scattered in the sub_implementations.
+                            # priority calibration for factorizable ops is OICalibrationData
+                            priority_calibration = {}
+                            for c in locus:
+                                priority_calibration |= get_custom_oi(self.sub_implementations[c])  # type: ignore[arg-type]
+                        else:
+                            priority_calibration = get_custom_oi(self.sub_implementations[locus[0]]).get(locus)
+                    else:
+                        # self has normal cal data
+                        oi = get_custom_oi(self)
+                        priority_calibration = oi if op.factorizable else oi.get(locus)  # type: ignore
+
         return self.builder.get_implementation(
             op_name,
             locus,
             impl_name=impl_name,
             strict_locus=strict_locus,
-            priority_calibration=composite_calibration,
+            priority_calibration=priority_calibration,
             use_priority_order=True,
         )
 
@@ -652,17 +790,21 @@ class CompositeCache:
     cache ``GateImplementation._timebox_cache`` as composites can include any gates in their calls, and we cannot trust
     that the cache is flushed correctly just based on if the composite itself has its own calibration data changed
     (we would have to flush also when any of the composite's members get new calibration, and this cannot consistently
-    be deduced). For this reason, CompositeCache will be flushed whenever ANY gate implementation gets new calibration
-    data.
+    be deduced). For this reason, the CompositeCache is located in the ScheduleBuilder class, and will be flushed
+    whenever *any* gate implementation gets new calibration data.
     """
 
     def __init__(self) -> None:
         self._cache: dict[tuple[Any, ...], TimeBox] = {}
 
     def set(
-        self, gate_implementation: GateImplementation, cache_key: tuple[Any, ...], timebox: TimeBox, extra_id: str = ""
+        self,
+        gate_implementation: GateImplementation,
+        cache_key: tuple[Any, ...],
+        timebox: TimeBox,
+        extra_id: str = "",
     ) -> None:
-        """Set a TimeBox into the cache.
+        """Store a TimeBox in the cache.
 
         Args:
             gate_implementation: gate implementation that created the TimeBox.