diff --git a/CHANGELOG.md b/CHANGELOG.md index 4fc114c6..2e1555d4 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/quam_builder/architecture/superconducting/components/twpa.py b/quam_builder/architecture/superconducting/components/twpa.py index 6fff7fd7..c08dbf02 100644 --- a/quam_builder/architecture/superconducting/components/twpa.py +++ b/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) diff --git a/tests/test_twpa_initialize.py b/tests/test_twpa_initialize.py new file mode 100644 index 00000000..dce1220b --- /dev/null +++ b/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 diff --git a/tests/test_twpa_keepalive.py b/tests/test_twpa_keepalive.py index bf086e3b..a8fc84cc 100644 --- a/tests/test_twpa_keepalive.py +++ b/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