Adding stay alive method for sticky elements (TWPA pumps) within shot loops

CHANGELOG.md

@@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ### Added
 
+- Added `BaseQuam.twpa_keepalive()` to keep sticky TWPA pumps on across a real-time loop (#123).
 - Added `custom_gates` section to `architecture/superconducting/README.md` documenting single-qubit macros (`MeasureMacro`, `ResetMacro`, `VirtualZMacro`, `DelayMacro`, `IdMacro`) and the `CZGate` two-qubit macro, including pulse-naming conventions and usage examples. Gate macros are currently specific to the superconducting architecture; `CZGate` requires `FluxTunableTransmonPair`.
 - Added incremental add/remove helpers for qubits, channels, and ports, including typed port helpers ``add_mw_port`` (MW-FEM) and ``add_lf_port`` (LF-FEM / OPX+ baseband) with a required ``type="input"`` or ``type="output"`` argument.
 

quam_builder/architecture/superconducting/qpu/base_quam.py

@@ -6,9 +6,10 @@
 from qm import QuantumMachinesManager, QuantumMachine
 from qm.octave import QmOctaveConfig
 from qm.qua.type_hints import QuaVariable, StreamType
-from qm.qua import declare_stream, declare, fixed
+from qm.qua import declare_stream, declare, fixed, align
 
 from quam.components import FrequencyConverter
+from quam.components.quantum_components import Qubit
 from quam.core import QuamRoot, quam_dataclass
 from quam.components.octave import Octave
 from quam.components.ports import FEMPortsContainer, OPXPlusPortsContainer
@@ -51,6 +52,7 @@ class BaseQuam(QuamRoot):
         thermalization_time: Return longest thermalization time.
         declare_qua_variables: Declare necessary QUA variables for qubits.
         initialize_qpu: Initialize the QPU with specified settings.
+        twpa_keepalive: Align the TWPA pumps with the given qubits to keep them on.
     """
 
     octaves: Dict[str, Octave] = field(default_factory=dict)
@@ -360,10 +362,47 @@ def declare_qua_variables(
     def initialize_qpu(self, isolation: bool = False, **kwargs):
         """Initialize the QPU with the calibrated TWPA pumping points.
 
+        The TWPA pumps require calling :meth:`twpa_keepalive` once per
+        inner-most-loop iteration to keep them on.
+
         Args:
         isolation : bool, optional
             If True, also configure and play the isolation tone. Use when the TWPA
             has isolation and you want it active. Default False.
         """
         for twpa in self.twpas.values():
             twpa.initialize(isolation=isolation)
+
+    def twpa_keepalive(self, qubits=None, isolation: bool = False) -> None:
+        """Align the TWPA pumps with the given qubits to keep them on across a loop.
+
+        Call once per iteration of the inner-most loop around your pulse operations,
+        so the sticky pumps stay in the program timeline.
+
+        Args:
+            qubits (Union[Qubit, Iterable[Qubit]], optional): Qubit(s) whose channels the
+                pumps are aligned with. Defaults to ``self.active_qubits``.
+            isolation (bool, optional): If True, also keep the isolation tone alive.
+                Default False.
+        """
+        if qubits is None:
+            qubits = self.active_qubits
+        elif isinstance(qubits, Qubit):
+            qubits = [qubits]
+        else:
+            qubits = list(qubits)
+
+        twpa_elements = []
+        for twpa in self.twpas.values():
+            if not twpa.initialization:
+                continue
+            if twpa.pump is not None:
+                twpa_elements.append(twpa.pump.name)
+            if isolation and twpa.isolation is not None:
+                twpa_elements.append(twpa.isolation.name)
+
+        if not twpa_elements:
+            return
+
+        qubit_channels = [ch.name for q in qubits for ch in q.channels.values()]
+        align(*twpa_elements, *qubit_channels)

quam_builder/architecture/superconducting/qpu/flux_tunable_quam.py

@@ -130,6 +130,9 @@ def initialize_qpu(self, isolation: bool = False, **kwargs):
         """Initialize the QPU with the calibrated TWPA pumping points and
            with the specified flux point and target
 
+        The TWPA pumps require calling :meth:`twpa_keepalive` once per
+        inner-most-loop iteration to keep them on.
+
         Args:
             flux_point (str): The flux point to set. Default is 'joint'.
             target: The qubit under study.

tests/test_twpa_keepalive.py

@@ -0,0 +1,103 @@
+"""Tests for ``BaseQuam.twpa_keepalive``.
+"""
+
+import pytest
+
+from quam.config.models.quam import QuamConfig
+from quam.config.resolvers import get_quam_config
+from quam.config.vars import CONFIG_PATH_ENV_NAME
+from quam.components.ports import FEMPortsContainer
+
+import quam_builder.architecture.superconducting.qpu.base_quam as base_quam_mod
+from quam_builder.architecture.superconducting.qpu import FixedFrequencyQuam
+from quam_builder.builder.superconducting.modify_quam import add_qubit
+from quam_builder.architecture.superconducting.components.twpa import TWPA
+from quam_builder.architecture.superconducting.components.xy_drive import XYDriveMW
+
+
+@pytest.fixture(autouse=True)
+def compatible_quam_config(tmp_path, monkeypatch):
+    """Use a qualibrate config that matches the installed quam package version."""
+    config_file = tmp_path / "config.toml"
+    config_file.write_text(f"[quam]\nversion = {QuamConfig.version}\n")
+    monkeypatch.setenv(CONFIG_PATH_ENV_NAME, str(config_file))
+    get_quam_config.cache_clear()
+
+
+@pytest.fixture
+def machine() -> FixedFrequencyQuam:
+    """A real machine with two qubits (q0 active, q1 inactive) and one TWPA."""
+    machine = FixedFrequencyQuam(ports=FEMPortsContainer())
+    for qubit_id, base in (("q0", 1), ("q1", 3)):
+        add_qubit(
+            machine,
+            qubit_id,
+            {
+                "xy": {"opx_output": f"#/ports/mw_outputs/con1/{base}/1"},
+                "rr": {
+                    "opx_output": f"#/ports/mw_outputs/con1/{base + 1}/1",
+                    "opx_input": f"#/ports/mw_inputs/con1/{base + 1}/1",
+                },
+            },
+        )
+    machine.active_qubit_names = ["q0"]
+    machine.twpas["twpa1"] = TWPA(
+        id="twpa1", pump=XYDriveMW(opx_output="#/ports/mw_outputs/con1/7/1")
+    )
+    return machine
+
+
+@pytest.fixture
+def record_align(monkeypatch):
+    """Capture every call to ``align`` (as element names) instead of emitting QUA."""
+    calls = []
+    monkeypatch.setattr(base_quam_mod, "align", lambda *names: calls.append(set(names)))
+    return calls
+
+
+def _channel_names(qubit):
+    return {ch.name for ch in qubit.channels.values()}
+
+
+def test_default_aligns_pump_with_active_qubits(machine, record_align):
+    machine.twpa_keepalive()
+
+    pump = machine.twpas["twpa1"].pump.name
+    assert record_align == [{pump} | _channel_names(machine.qubits["q0"])]
+    # inactive qubit must not be dragged into the align
+    assert record_align[0].isdisjoint(_channel_names(machine.qubits["q1"]))
+
+
+def test_explicit_qubits_include_all_given(machine, record_align):
+    q0, q1 = machine.qubits["q0"], machine.qubits["q1"]
+    machine.twpa_keepalive([q0, q1])
+
+    pump = machine.twpas["twpa1"].pump.name
+    assert record_align[0] == {pump} | _channel_names(q0) | _channel_names(q1)
+
+
+def test_single_qubit_object(machine, record_align):
+    q0 = machine.qubits["q0"]
+    machine.twpa_keepalive(q0)  # a bare Qubit, not a list
+
+    pump = machine.twpas["twpa1"].pump.name
+    assert record_align[0] == {pump} | _channel_names(q0)
+
+
+def test_skips_uninitialized_twpa(machine, record_align):
+    machine.twpas["twpa1"].initialization = False
+    machine.twpa_keepalive()
+
+    assert record_align == []  # no initialized pumps -> no align emitted
+
+
+def test_isolation_flag(machine, record_align):
+    twpa = machine.twpas["twpa1"]
+    twpa.isolation = XYDriveMW(opx_output="#/ports/mw_outputs/con1/8/1")
+
+    machine.twpa_keepalive(isolation=False)
+    assert twpa.isolation.name not in record_align[0]
+
+    record_align.clear()
+    machine.twpa_keepalive(isolation=True)
+    assert twpa.isolation.name in record_align[0]