Merge pull request #1068 from qiboteam/opx1000-0.2

Support for QM OPX1000 in 0.2

src/qibolab/_core/instruments/qm/components/configs.py

@@ -2,9 +2,18 @@
 
 from pydantic import Field
 
-from qibolab._core.components import AcquisitionConfig, DcConfig
+from qibolab._core.components import AcquisitionConfig, DcConfig, OscillatorConfig
 
-__all__ = ["OpxOutputConfig", "QmAcquisitionConfig", "QmConfigs"]
+__all__ = [
+    "OpxOutputConfig",
+    "QmAcquisitionConfig",
+    "QmConfigs",
+    "OctaveOscillatorConfig",
+]
+
+OctaveOutputModes = Literal[
+    "always_on", "always_off", "triggered", "triggered_reversed"
+]
 
 
 class OpxOutputConfig(DcConfig):
@@ -25,6 +34,15 @@ class OpxOutputConfig(DcConfig):
     for more details.
     Changing the filters affects the calibration of single shot discrimination (threshold and angle).
     """
+    output_mode: Literal["direct", "amplified"] = "direct"
+
+
+class OctaveOscillatorConfig(OscillatorConfig):
+    """Oscillator confing that allows switching the output mode."""
+
+    kind: Literal["octave-oscillator"] = "octave-oscillator"
+
+    output_mode: OctaveOutputModes = "triggered"
 
 
 class QmAcquisitionConfig(AcquisitionConfig):
@@ -41,4 +59,4 @@ class QmAcquisitionConfig(AcquisitionConfig):
     """Constant voltage to be applied on the input."""
 
 
-QmConfigs = Union[OpxOutputConfig, QmAcquisitionConfig]
+QmConfigs = Union[OpxOutputConfig, OctaveOscillatorConfig, QmAcquisitionConfig]

src/qibolab/_core/instruments/qm/config/__init__.py

@@ -1,2 +1,3 @@
 from .config import Configuration
+from .devices import ControllerId, ModuleTypes
 from .pulses import SAMPLING_RATE, operation

src/qibolab/_core/instruments/qm/config/config.py

@@ -12,7 +12,17 @@
 from qibolab._core.pulses import Pulse, Readout
 
 from ..components import OpxOutputConfig, QmAcquisitionConfig
-from .devices import AnalogOutput, Controller, Octave, OctaveInput, OctaveOutput
+from .devices import (
+    AnalogOutput,
+    Controller,
+    ControllerId,
+    Controllers,
+    FemAnalogOutput,
+    ModuleTypes,
+    Octave,
+    OctaveInput,
+    OctaveOutput,
+)
 from .elements import AcquireOctaveElement, DcElement, Element, RfOctaveElement
 from .pulses import (
     QmAcquisition,
@@ -42,7 +52,7 @@ class Configuration:
     """
 
     version: int = 1
-    controllers: dict[str, Controller] = field(default_factory=dict)
+    controllers: Controllers = field(default_factory=Controllers)
     octaves: dict[str, Octave] = field(default_factory=dict)
     elements: dict[str, Element] = field(default_factory=dict)
     pulses: dict[str, Union[QmPulse, QmAcquisition]] = field(default_factory=dict)
@@ -53,20 +63,27 @@ class Configuration:
     integration_weights: dict = field(default_factory=dict)
     mixers: dict = field(default_factory=dict)
 
-    def add_controller(self, device: str):
+    def add_controller(self, device: ControllerId, modules: dict[str, ModuleTypes]):
         if device not in self.controllers:
-            self.controllers[device] = Controller()
+            self.controllers[device] = Controller(type=modules[device])
 
-    def add_octave(self, device: str, connectivity: str):
+    def add_octave(
+        self, device: str, connectivity: ControllerId, modules: dict[str, ModuleTypes]
+    ):
         if device not in self.octaves:
-            self.add_controller(connectivity)
+            self.add_controller(connectivity, modules)
             self.octaves[device] = Octave(connectivity)
 
     def configure_dc_line(
         self, id: ChannelId, channel: DcChannel, config: OpxOutputConfig
     ):
         controller = self.controllers[channel.device]
-        controller.analog_outputs[channel.port] = AnalogOutput.from_config(config)
+        if controller.type == "opx1":
+            controller.analog_outputs[channel.port] = AnalogOutput.from_config(config)
+        else:
+            controller.analog_outputs[channel.port] = FemAnalogOutput.from_config(
+                config
+            )
         self.elements[id] = DcElement.from_channel(channel)
 
     def configure_iq_line(
@@ -116,7 +133,11 @@ def configure_acquire_line(
         )
 
     def register_waveforms(
-        self, pulse: Pulse, element: Optional[str] = None, dc: bool = False
+        self,
+        pulse: Pulse,
+        max_voltage: float,
+        element: Optional[str] = None,
+        dc: bool = False,
     ):
         if dc:
             qmpulse = QmPulse.from_dc_pulse(pulse)
@@ -125,34 +146,38 @@ def register_waveforms(
                 qmpulse = QmPulse.from_pulse(pulse)
             else:
                 qmpulse = QmAcquisition.from_pulse(pulse, element)
-        waveforms = waveforms_from_pulse(pulse)
+        waveforms = waveforms_from_pulse(pulse, max_voltage)
         if dc:
             self.waveforms[qmpulse.waveforms["single"]] = waveforms["I"]
         else:
             for mode in ["I", "Q"]:
                 self.waveforms[getattr(qmpulse.waveforms, mode)] = waveforms[mode]
         return qmpulse
 
-    def register_iq_pulse(self, element: str, pulse: Pulse):
+    def register_iq_pulse(self, element: str, pulse: Pulse, max_voltage: float):
         op = operation(pulse)
         if op not in self.pulses:
-            self.pulses[op] = self.register_waveforms(pulse)
+            self.pulses[op] = self.register_waveforms(pulse, max_voltage)
         self.elements[element].operations[op] = op
         return op
 
-    def register_dc_pulse(self, element: str, pulse: Pulse):
+    def register_dc_pulse(self, element: str, pulse: Pulse, max_voltage: float):
         op = operation(pulse)
         if op not in self.pulses:
-            self.pulses[op] = self.register_waveforms(pulse, dc=True)
+            self.pulses[op] = self.register_waveforms(pulse, max_voltage, dc=True)
         self.elements[element].operations[op] = op
         return op
 
-    def register_acquisition_pulse(self, element: str, readout: Readout):
+    def register_acquisition_pulse(
+        self, element: str, readout: Readout, max_voltage: float
+    ):
         """Registers pulse, waveforms and integration weights in QM config."""
         op = operation(readout)
         acquisition = f"{op}_{element}"
         if acquisition not in self.pulses:
-            self.pulses[acquisition] = self.register_waveforms(readout.probe, element)
+            self.pulses[acquisition] = self.register_waveforms(
+                readout.probe, max_voltage, element
+            )
         self.elements[element].operations[op] = acquisition
         return op
 

src/qibolab/_core/instruments/qm/config/devices.py

@@ -1,14 +1,25 @@
 from dataclasses import dataclass, field
-from typing import Any, Generic, TypeVar
+from typing import Generic, Literal, TypeVar, Union
 
 from qibolab._core.components import OscillatorConfig
 
-from ..components import OpxOutputConfig, QmAcquisitionConfig
+from ..components import OctaveOscillatorConfig, OpxOutputConfig, QmAcquisitionConfig
+from ..components.configs import OctaveOutputModes
 
-__all__ = ["AnalogOutput", "OctaveOutput", "OctaveInput", "Controller", "Octave"]
+__all__ = [
+    "AnalogOutput",
+    "FemAnalogOutput",
+    "ModuleTypes",
+    "OctaveOutput",
+    "OctaveInput",
+    "Controller",
+    "Octave",
+    "ControllerId",
+    "Controllers",
+]
 
 
-DEFAULT_INPUTS = {"1": {}, "2": {}}
+DEFAULT_INPUTS = {"1": {"offset": 0}, "2": {"offset": 0}}
 """Default controller config section.
 
 Inputs are always registered to avoid issues with automatic mixer
@@ -25,7 +36,7 @@ class PortDict(Generic[V], dict[str, V]):
     in the QUA config.
     """
 
-    def __setitem__(self, key: Any, value: V):
+    def __setitem__(self, key: Union[str, int], value: V):
         super().__setitem__(str(key), value)
 
 
@@ -39,6 +50,17 @@ def from_config(cls, config: OpxOutputConfig):
         return cls(offset=config.offset, filter=config.filter)
 
 
+@dataclass(frozen=True)
+class FemAnalogOutput(AnalogOutput):
+    output_mode: Literal["direct", "amplified"] = "direct"
+
+    @classmethod
+    def from_config(cls, config: OpxOutputConfig):
+        return cls(
+            offset=config.offset, filter=config.filter, output_mode=config.output_mode
+        )
+
+
 @dataclass(frozen=True)
 class AnalogInput:
     offset: float = 0.0
@@ -53,24 +75,32 @@ def from_config(cls, config: QmAcquisitionConfig):
 class OctaveOutput:
     LO_frequency: int
     gain: int = 0
-    LO_source: str = "internal"
-    output_mode: str = "triggered"
+    LO_source: Literal["internal", "external"] = "internal"
+    output_mode: OctaveOutputModes = "triggered"
 
     @classmethod
-    def from_config(cls, config: OscillatorConfig):
-        return cls(LO_frequency=config.frequency, gain=config.power)
+    def from_config(cls, config: Union[OscillatorConfig, OctaveOscillatorConfig]):
+        kwargs = dict(LO_frequency=config.frequency, gain=config.power)
+        if isinstance(config, OctaveOscillatorConfig):
+            kwargs["output_mode"] = config.output_mode
+        return cls(**kwargs)
 
 
 @dataclass(frozen=True)
 class OctaveInput:
     LO_frequency: int
-    LO_source: str = "internal"
-    IF_mode_I: str = "direct"
-    IF_mode_Q: str = "direct"
+    LO_source: Literal["internal", "external"] = "internal"
+    IF_mode_I: Literal["direct", "envelop", "mixer"] = "direct"
+    IF_mode_Q: Literal["direct", "envelop", "mixer"] = "direct"
+
+
+ModuleTypes = Literal["opx1", "LF", "MW"]
 
 
 @dataclass
 class Controller:
+    type: ModuleTypes = "opx1"
+    """https://docs.quantum-machines.co/latest/docs/Introduction/config/?h=opx10#controllers"""
     analog_outputs: PortDict[dict[str, AnalogOutput]] = field(default_factory=PortDict)
     digital_outputs: PortDict[dict[str, dict]] = field(default_factory=PortDict)
     analog_inputs: PortDict[dict[str, AnalogInput]] = field(
@@ -90,8 +120,66 @@ def add_octave_input(self, port: int, config: QmAcquisitionConfig):
         )
 
 
+@dataclass
+class Opx1000:
+    type: Literal["opx1000"] = "opx1000"
+    fems: dict[str, Controller] = field(default_factory=PortDict)
+
+
 @dataclass
 class Octave:
-    connectivity: str
+    connectivity: Union[str, tuple[str, int]]
     RF_outputs: PortDict[dict[str, OctaveOutput]] = field(default_factory=PortDict)
     RF_inputs: PortDict[dict[str, OctaveInput]] = field(default_factory=PortDict)
+
+    def __post_init__(self):
+        if "/" in self.connectivity:
+            con, fem = self.connectivity.split("/")
+            self.connectivity = (con, int(fem))
+
+
+ControllerId = Union[str, tuple[str, int]]
+
+
+def process_controller_id(id: ControllerId):
+    """Convert controller identifier depending on cluster type.
+
+    For OPX+ clusters ``id`` is just the controller name (eg. 'con1').
+    For OPX1000 clusters ``id`` has the format
+    '{controller_name}/{fem_number}' (eg. 'con1/4').
+    In that case ``id`` may also be a ``tuple``
+    `(controller_name, fem_number)`
+    """
+    if isinstance(id, tuple):
+        con, fem = id
+        return con, str(fem)
+    if "/" in id:
+        return id.split("/")
+    return id, None
+
+
+class Controllers(dict[str, Union[Controller, Opx1000]]):
+    """Dictionary of controllers compatible with OPX+ and OPX1000."""
+
+    def __contains__(self, key: ControllerId) -> bool:
+        con, fem = process_controller_id(key)
+        contains = super().__contains__(con)
+        if fem is None:
+            return contains
+        return contains and fem in self[con].fems
+
+    def __getitem__(self, key: ControllerId) -> Controller:
+        con, fem = process_controller_id(key)
+        value = super().__getitem__(con)
+        if fem is None:
+            return value
+        return value.fems[fem]
+
+    def __setitem__(self, key: ControllerId, value: Controller):
+        con, fem = process_controller_id(key)
+        if fem is None:
+            super().__setitem__(key, value)
+        else:
+            if con not in self:
+                super().__setitem__(con, Opx1000())
+            self[con].fems[fem] = value