Add stretch support

changelog.d/384.added.md

@@ -0,0 +1 @@
+Added support for `Delay` operations with `Stretch` durations.

samplomatic/builders/build.py

@@ -60,6 +60,7 @@ def _build(stream: DAGCircuit, builder: Builder):
         qubit_remapping = dict(zip(nested_instr.qargs, nested_instr.op.body.qubits))
 
         remapped_template_state = builder.template_state.remap(qubit_remapping, idx)
+        remapped_template_state.add_stretches(nested_instr.op.body.iter_stretches())
         remapped_pre_samplex = builder.samplex_state.remap(remapped_template_state.qubit_map)
         inner_builder = inner_builder.set_template_state(remapped_template_state).set_samplex_state(
             remapped_pre_samplex

samplomatic/builders/template_state.py

@@ -14,8 +14,9 @@
 
 from collections.abc import Iterable, Sequence
 
-from qiskit.circuit import ClassicalRegister, Clbit, QuantumCircuit, QuantumRegister, Qubit
+from qiskit.circuit import ClassicalRegister, Clbit, Delay, QuantumCircuit, QuantumRegister, Qubit
 from qiskit.circuit.classical import expr
+from qiskit.circuit.classical.expr import Stretch
 from qiskit.converters import dag_to_circuit
 from qiskit.dagcircuit import DAGCircuit
 
@@ -32,6 +33,8 @@ class TemplateState:
         qubit_map: A map from qubits in the source circuit to corresponding qubits in the template.
         param_iter: An iterator over parameters to use in the circuit being built.
         scope_idx: The nested index of the scope currently being built.
+        stretch_map: A map from stretches in the source circuit to corresponding stretches in the
+            template.
     """
 
     def __init__(
@@ -40,12 +43,14 @@ def __init__(
         qubit_map: dict[Qubit, QubitIndex],
         param_iter: ParamIter,
         scope_idx: list[int],
+        stretch_map: dict[Stretch, Stretch],
         debug: bool = False,
     ):
         self.template: DAGCircuit = template
         self.qubit_map = qubit_map
         self.param_iter = param_iter
         self.scope_idx = scope_idx
+        self.stretch_map = stretch_map
         self.debug = debug
 
     def remap(
@@ -66,7 +71,9 @@ def remap(
             for parent_scope_qubit, qubit in scoped_qubit_map.items()
         }
         scope_idx = self.scope_idx if last_scope_idx is None else self.scope_idx + [last_scope_idx]
-        return TemplateState(self.template, new_qubit_map, self.param_iter, scope_idx, self.debug)
+        return TemplateState(
+            self.template, new_qubit_map, self.param_iter, scope_idx, self.stretch_map, self.debug
+        )
 
     @classmethod
     def construct_for_circuit(cls, circuit: QuantumCircuit, debug: bool = False) -> Self:
@@ -83,6 +90,11 @@ def construct_for_circuit(cls, circuit: QuantumCircuit, debug: bool = False) ->
             template_circuit.add_creg(creg)
 
         qubit_map = {q: idx for idx, q in enumerate(circuit.qubits)}
+        stretch_map = {}
+        for stretch in circuit.iter_stretches():
+            new_stretch = Stretch.new(stretch.name)
+            stretch_map[stretch] = new_stretch
+            template_circuit.add_declared_stretch(new_stretch)
 
         # quick and dirty heuristic to get the max params roughly correct with a safety factor
         # TODO: This estimate might not hold for dynamic circuits, where the same qubit will be
@@ -93,7 +105,7 @@ def construct_for_circuit(cls, circuit: QuantumCircuit, debug: bool = False) ->
         max_params += circuit.num_parameters
         param_iter = ParamIter(5 * max_params)
 
-        return cls(template_circuit, qubit_map, param_iter, [], debug)
+        return cls(template_circuit, qubit_map, param_iter, [], stretch_map, debug)
 
     def qubits(self, idxs: Iterable[int] | None = None) -> Sequence[Qubit]:
         """Return the qubits in the template at the given indices.
@@ -107,6 +119,23 @@ def qubits(self, idxs: Iterable[int] | None = None) -> Sequence[Qubit]:
         idxs = self.qubit_map.values() if idxs is None else idxs
         return [self.template.qubits[i] for i in idxs]
 
+    def add_stretches(self, stretches: Iterable[Stretch]):
+        """Add stretches to the template circuit.
+
+        This method adds new stretches to the template circuit with the same name as those in
+        ``stretches`` with the scope index prepended if it has not already been added.
+
+        Args:
+            stretches: An iterable of stretches to add.
+        """
+        prefix = ".".join(str(i) for i in self.scope_idx)
+        for stretch in stretches:
+            if stretch in self.stretch_map:
+                continue
+            new_stretch = Stretch.new(f"{prefix}.{stretch.name}")
+            self.stretch_map[stretch] = new_stretch
+            self.template.add_declared_stretch(new_stretch)
+
     def append_remapped_gate(self, dag_op_node: DAGOpNode) -> ParamSpec:
         """Remap the parameters and qubits of a gate and append it to the circuit."""
         new_params = []
@@ -118,6 +147,10 @@ def append_remapped_gate(self, dag_op_node: DAGOpNode) -> ParamSpec:
                 param_mapping.append([self.param_iter.idx, param])
                 new_params.append(next(self.param_iter))
             new_operation = type(dag_op_node.op)(*new_params) if new_params else dag_op_node.op
+        elif isinstance(delay := dag_op_node.op, Delay) and isinstance(
+            duration := delay.duration, Stretch
+        ):
+            new_operation = Delay(self.stretch_map[duration])
         else:
             new_operation = dag_op_node.op
 

test/unit/test_builders/test_builders_pre_samplex.py

@@ -40,7 +40,7 @@ def get_builder(self, qreg, creg=None):
         circuit = DAGCircuit()
         circuit.add_qreg(qreg)
         circuit.add_creg(creg)
-        template_state = TemplateState(circuit, qubit_map, ParamIter(), [0])
+        template_state = TemplateState(circuit, qubit_map, ParamIter(), [0], {})
         pre_samplex = PreSamplex(qubit_map=qubit_map, cregs=[creg])
         qubits = QubitPartition.from_elements(qreg)
         builder = LeftBoxBuilder(

test/unit/test_builders/test_builders_template.py

@@ -13,6 +13,7 @@
 """Test BoxBuilder for Templates"""
 
 from qiskit.circuit import Parameter, QuantumCircuit
+from qiskit.circuit.classical.expr import Stretch
 
 from samplomatic.annotations import Twirl
 from samplomatic.builders import pre_build
@@ -33,6 +34,40 @@ def test_qubits(self):
         assert len(new_state.qubits()) == 1
         assert new_state.qubits() == [state.template.qubits[2]]
 
+    def test_stretch(self):
+        """Test that stretches are added appropriately."""
+        circuit = QuantumCircuit(2)
+        x = circuit.add_stretch("x")
+        circuit.delay(x, 0)
+        circuit.x(0)
+        circuit.delay(x, 0)
+        circuit.measure_all()
+
+        state = TemplateState.construct_for_circuit(circuit)
+
+        assert state.template.num_stretches == 1
+
+        stretch = next(state.template.iter_stretches())
+        assert stretch != circuit.get_stretch("x")
+
+    def test_rescoped_stretch(self):
+        """Test that stretches with the same name in different scopes do not overlap."""
+        circuit = QuantumCircuit(1)
+        state = TemplateState.construct_for_circuit(circuit)
+
+        new_state = state.remap({state.template.qubits[0]: 0}, 0)
+        new_new_state = new_state.remap({0: 0}, 1)
+        new_new_state.add_stretches([Stretch.new("x")])  # scope 0.1
+        new_state.add_stretches([Stretch.new("x")])  # scope 0
+
+        new_state = state.remap({state.template.qubits[0]: 0}, 2)  # scope 2
+        new_state.add_stretches([stretch := Stretch.new("x")])
+
+        new_state = new_state.remap({0: 0}, 0)
+        new_state.add_stretches([stretch])  # scope 2.0, but this stretch added in outer scope
+
+        assert {s.name for s in state.stretch_map.values()} == {"0.x", "0.1.x", "2.x"}
+
 
 class TestTemplateBuilder:
     """Test strictly the template aspects of build."""
@@ -140,6 +175,30 @@ def test_box_left_right(self):
         for idx, (instr, name) in enumerate(zip(template, expected_names)):
             assert instr.name == name, f"Instruction {idx}"
 
+    def test_box_stretches(self):
+        """Test boxes with stretches."""
+        circuit = QuantumCircuit(2)
+        with circuit.box([Twirl()]):
+            x = circuit.add_stretch("x")
+            circuit.delay(x)
+            circuit.measure_all()
+
+        with circuit.box([Twirl()]):
+            x = circuit.add_stretch("x")
+            circuit.delay(x)
+            circuit.measure_all()
+
+        x = circuit.add_stretch("x")
+        circuit.delay(x)
+        circuit.measure_all()
+
+        template_state, _ = pre_build(circuit)
+        template = template_state.finalize()
+
+        assert template.num_stretches == 3
+        assert {s.name for s in template.iter_stretches()} == {"0.x", "1.x", "x"}
+        assert template.get_stretch("x") != x
+
     def test_box_decomposition(self):
         """Test decomposition modes of a box."""
         circuit = QuantumCircuit(2)