Source code for semi_cr.core.quam_components.cryo_switch

from __future__ import annotations

from quam.components.channels import InOutSingleChannel
from quam.components.pulses import SquareReadoutPulse
from quam.core import QuamComponent, quam_dataclass
from quam.utils.qua_types import ScalarInt

__all__ = ["CryoSwitch"]




@quam_dataclass
class CryoSwitch(QuamComponent):
    """QUAM component for a cryogenic RF switch / multiplexer.

    Groups multiple :class:`~quam.components.channels.InOutSingleChannel` instances
    and exposes helpers to drive them together inside QUA programs (frequency sweeps,
    alignment, simultaneous measurements, etc.).

    Attributes:
        channels: Mapping from logical channel names to in/out reflectometry channels.

    """

    in_channel: InOutSingleChannel
    channels: dict[str, InOutSingleChannel]
    amplitude: float = 0.1
    intermediate_frequency: int = 100_000_000
    readout_length: int = 1000

    @property
    def valid_channel_names(self) -> list[str]:
        return list(self.channels.keys())



    def all_channels(self) -> list[InOutSingleChannel]:
        if self.in_channel is None:
            raise ValueError("In channel is not set")
        return [self.in_channel] + list(self.channels.values())




    def setup_switch(self) -> None:
        if self.in_channel is None:
            raise ValueError("In channel is not set")
        # Update intermediate frequency of all channels to match self.intermediate_frequency
        self.in_channel.intermediate_frequency = self.intermediate_frequency
        for channel in self.channels.values():
            channel.intermediate_frequency = self.intermediate_frequency

        inop = SquareReadoutPulse(
            id="cryo_switch_readout",
            length=self.readout_length,
            amplitude=self.amplitude,
        )
        self.in_channel.operations["cryo_switch_readout"] = inop
        for channel in self.channels.values():
            outop = SquareReadoutPulse(
                id="cryo_switch_readout",
                length=self.readout_length,
                amplitude=0,
            )
            channel.operations["cryo_switch_readout"] = outop




    def validate_switch(self) -> None:
        # Validate all channels' intermediate frequency matches
        for channel in self.all_channels():
            if channel.intermediate_frequency != self.intermediate_frequency:
                raise ValueError(
                    f"Channel '{channel.name}' has intermediate_frequency={channel.intermediate_frequency}, "
                    f"expected {self.intermediate_frequency}"
                )

        # Validate 'cryo_switch_readout' operations
        for channel in self.all_channels():
            if "cryo_switch_readout" not in channel.operations:
                raise ValueError(
                    f"Channel {channel.name} does not have a 'cryo_switch_readout' operation"
                )

            op = channel.operations["cryo_switch_readout"]
            # Ensure op meets specs
            if op.length != self.readout_length:
                raise ValueError(
                    f"Channel {channel.name} has a 'cryo_switch_readout' operation with a length of {op.length}, "
                    f"but the required readout length is {self.readout_length}"
                )
            expected_amplitude = self.amplitude if channel is self.in_channel else 0
            if op.amplitude != expected_amplitude:
                raise ValueError(
                    f"Channel {channel.name} has a 'cryo_switch_readout' amplitude {op.amplitude}, "
                    f"expected {expected_amplitude} (amplitude only on in_channel)"
                )




    def align(self) -> None:
        """Align timing across the selected channels (default: all)."""
        self.in_channel.align(*self.channels)




    def wait(
        self,
        duration: ScalarInt,
    ) -> None:
        """Wait on all selected channels without outputting a pulse."""
        self.in_channel.wait(duration, *self.channels)




    def update_frequency(
        self,
        new_frequency: ScalarInt,
        units: str = "Hz",
        keep_phase: bool = False,
    ) -> None:
        """Update the intermediate frequency on all selected channels."""
        for channel in self.all_channels():
            channel.update_frequency(new_frequency, units=units, keep_phase=keep_phase)




    def reset_phase(self) -> None:
        """Reset phase on all channels"""
        for channel in self.all_channels():
            channel.reset_if_phase()




    def measure(
        self,
    ):
        self.validate_switch()
        self.align()
        results = {channel.name: {} for channel in self.all_channels()}
        for channel in self.all_channels():
            results[channel.name]["I"], results[channel.name]["Q"] = channel.measure(
                pulse_name="cryo_switch_readout"
            )
        return results