fix(twpa): scope TWPA init dedup to active QUA program

CHANGELOG.md

@@ -14,6 +14,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ### Fixed
 
+- Fixed TWPA pump initialization being skipped in subsequent QUA programs due to process-global deduplication state ([#122](https://github.com/qua-platform/quam-builder/issues/122)).
 - `add_default_transmon_pair_macros` now passes `flux_pulse_qubit="const"` (was `flux_pulse_control="const"`) to `CZGate`, matching the field rename introduced in v0.4.0 and aligning with the `"const"` pulse added to every `FluxTunableTransmon` Z line by `add_default_transmon_pulses`.
 - Applied black formatting across the entire repo.
 

quam_builder/architecture/superconducting/components/twpa.py

@@ -1,6 +1,7 @@
 from quam.core import quam_dataclass
 from quam import QuamComponent
-from typing import Union, ClassVar
+from typing import Union, Optional, Any
+from qm.qua._scope_management.scopes_manager import scopes_manager
 from .xy_drive import XYDriveMW, XYDriveIQ
 
 __all__ = ["TWPA"]
@@ -49,9 +50,9 @@ class TWPA(QuamComponent):
         Qubits whose readout signals are amplified by this TWPA.
     initialization : bool
         If True, use this TWPA in the QUA program (e.g. call initialize). Default True.
-    _initialized_ids : ClassVar[set]
-        Class-level set of initialized instance ids to ensure pump is initialized
-        only once per program run. Not serialized.
+    _initialized_prog_scope : Any, optional
+        Runtime-only reference to the active QUA program scope in which this TWPA
+        was initialized. Not serialized.
     """
 
     id: Union[int, str]
@@ -74,7 +75,8 @@ class TWPA(QuamComponent):
     qubits: list = None
 
     initialization: bool = True
-    _initialized_ids: ClassVar[set] = set()
+    _initialized_prog_scope: Optional[Any] = None
+    _skip_attrs = ["_initialized_prog_scope"]
 
     @property
     def name(self):
@@ -86,11 +88,11 @@ def initialize(self, isolation: bool = False):
         Set the TWPA pump (and optionally isolation) to the calibrated tones for the QUA program.
 
         Has no effect if :attr:`initialization` is False. Each instance is initialized
-        at most once per program run; further calls return immediately. Sets the pump
-        frequency from :attr:`pump_frequency` and :attr:`pump_amplitude`, then plays
-        the pump operation. If ``isolation`` is True and this TWPA has isolation
-        channels, also sets and plays the isolation tone from :attr:`isolation_frequency`
-        and :attr:`isolation_amplitude`.
+        at most once per active ``program()`` block; further calls in the same block
+        return immediately. Sets the pump frequency from :attr:`pump_frequency` and
+        :attr:`pump_amplitude`, then plays the pump operation. If ``isolation`` is
+        True and this TWPA has isolation channels, also sets and plays the isolation
+        tone from :attr:`isolation_frequency` and :attr:`isolation_amplitude`.
 
         Parameters
         ----------
@@ -100,18 +102,18 @@ def initialize(self, isolation: bool = False):
 
         Notes
         -----
-        Initialization state is tracked in :attr:`_initialized_ids`, so the pump is
-        turned on onl
+        Initialization state is tracked in :attr:`_initialized_prog_scope` keyed to the
+        active QUA program scope object, so the pump is turned on only once per program
+        compilation while each new ``program()`` block re-initializes.
         """
         # dont use twpa for the QUA program if initialization is set to False
         if not self.initialization:
             return
-        # Initialize TWPA pump only when it hasn't been initialized yet
-        # This won't be serialized since it's stored in a class-level set
-        obj_id = id(self)
-        # Check initialization state using object ID (memory address)
-        if obj_id in self._initialized_ids:
+
+        current_program_scope = scopes_manager.program_scope
+        if self._initialized_prog_scope is current_program_scope:
             return
+        self._initialized_prog_scope = current_program_scope
 
         if self.pump_frequency is not None:
             self.pump.update_frequency(int(self.pump_frequency - self.pump.LO_frequency))
@@ -129,7 +131,3 @@ def initialize(self, isolation: bool = False):
                 self.isolation.play("pump", amplitude_scale=self.isolation_amplitude)
             else:
                 self.isolation.play("pump")
-
-        # Store object ID externally (won't be serialized)
-        # guarantee initializing twpa pump only once per QUA program execution
-        self._initialized_ids.add(obj_id)

tests/test_twpa_initialize.py

@@ -0,0 +1,80 @@
+"""Tests for ``TWPA.initialize`` program scoping."""
+
+import pytest
+from qm.qua import program
+
+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
+
+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))
+    if hasattr(get_quam_config, "cache_clear"):
+        get_quam_config.cache_clear()
+
+
+@pytest.fixture
+def machine() -> FixedFrequencyQuam:
+    """A machine with one TWPA configured for pump initialization."""
+    machine = FixedFrequencyQuam(ports=FEMPortsContainer())
+    add_qubit(
+        machine,
+        "q0",
+        {
+            "xy": {"opx_output": "#/ports/mw_outputs/con1/1/1"},
+            "rr": {
+                "opx_output": "#/ports/mw_outputs/con1/2/1",
+                "opx_input": "#/ports/mw_inputs/con1/2/1",
+            },
+        },
+    )
+    machine.active_qubit_names = ["q0"]
+    machine.twpas["twpa1"] = TWPA(
+        id="twpa1",
+        pump=XYDriveMW(opx_output="#/ports/mw_outputs/con1/7/1"),
+        pump_amplitude=0.5,
+    )
+    return machine
+
+
+@pytest.fixture
+def record_pump_play(machine, monkeypatch):
+    """Record every call to the TWPA pump ``play`` method."""
+    plays = []
+
+    def record_play(*args, **kwargs):
+        plays.append(1)
+
+    monkeypatch.setattr(machine.twpas["twpa1"].pump, "play", record_play)
+    return plays
+
+
+def test_twpa_initialize_once_per_program(machine, record_pump_play):
+    """Two initialize_qpu calls in one program block emit TWPA init only once."""
+    with program():
+        machine.initialize_qpu()
+        machine.initialize_qpu()
+
+    assert len(record_pump_play) == 1
+
+
+def test_twpa_initialize_in_each_program(machine, record_pump_play):
+    """Each new program block re-initializes the TWPA pump."""
+    with program():
+        machine.initialize_qpu()
+    assert len(record_pump_play) == 1
+
+    with program():
+        machine.initialize_qpu()
+    assert len(record_pump_play) == 2

tests/test_twpa_keepalive.py

@@ -21,7 +21,8 @@ def compatible_quam_config(tmp_path, monkeypatch):
     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()
+    if hasattr(get_quam_config, "cache_clear"):
+        get_quam_config.cache_clear()
 
 
 @pytest.fixture