Qblox

pyproject.toml

@@ -26,7 +26,7 @@ qibo = "^0.2.8"
 numpy = "^1.26.4"
 scipy = "^1.13.0"
 pydantic = "^2.6.4"
-qblox-instruments = { version = "0.12.0", optional = true }
+qblox-instruments = { version = "^0.14.2", optional = true }
 lark = { version = "^1.1.9", optional = true }
 qcodes = { version = "^0.37.0", optional = true }
 qcodes_contrib_drivers = { version = "0.18.0", optional = true }
@@ -59,7 +59,7 @@ nbsphinx = "^0.9.1"
 ipython = "^8.12.0"
 sphinx-copybutton = "^0.5.1"
 # extras
-qblox-instruments = "0.12.0"
+qblox-instruments = "^0.14.2"
 qcodes = "^0.37.0"
 qcodes_contrib_drivers = "0.18.0"
 qibosoq = { version = "^0.1.2", python = "<3.12" }

src/qibolab/_core/components/channels.py

@@ -46,21 +46,21 @@ class DcChannel(Channel):
 class IqChannel(Channel):
     """Channel that can be used to send IQ pulses."""
 
-    mixer: Optional[str]
+    mixer: Optional[str] = None
     """Name of the IQ mixer component corresponding to this channel.
 
     None, if the channel does not have a mixer, or it does not need
     configuration.
     """
-    lo: Optional[str]
+    lo: Optional[str] = None
     """Name of the local oscillator component corresponding to this channel.
 
     None, if the channel does not have an LO, or it is not configurable.
     """
 
 
 class AcquisitionChannel(Channel):
-    twpa_pump: Optional[str]
+    twpa_pump: Optional[str] = None
     """Name of the TWPA pump component.
 
     None, if there is no TWPA, or it is not configurable.

src/qibolab/_core/instruments/qblox/__init__.py

@@ -1 +1 @@
-from . import ast_
+__all__ = []

src/qibolab/_core/instruments/qblox/ast_.py

@@ -13,6 +13,8 @@
 
 from ...serialize import Model
 
+__all__ = []
+
 
 class Register(Model):
     number: int

src/qibolab/_core/instruments/qblox/cluster.py

@@ -0,0 +1,175 @@
+from collections import defaultdict
+from functools import cached_property
+from typing import Optional
+
+import qblox_instruments as qblox
+from qblox_instruments.qcodes_drivers.module import Module
+from qcodes.instrument import find_or_create_instrument
+
+from qibolab._core.components.configs import Config
+from qibolab._core.execution_parameters import ExecutionParameters
+from qibolab._core.identifier import ChannelId, Result
+from qibolab._core.instruments.abstract import Controller
+from qibolab._core.sequence import PulseSequence
+from qibolab._core.serialize import Model
+from qibolab._core.sweeper import ParallelSweepers
+
+from .sequence import Sequence
+
+__all__ = ["Cluster"]
+
+SAMPLING_RATE = 1
+
+
+SlotId = int
+SequencerId = int
+
+
+class PortAddress(Model):
+    slot: SlotId
+    ports: tuple[int, Optional[int]]
+    input: bool = False
+
+    @classmethod
+    def from_path(cls, path: str):
+        """Load address from :attr:`qibolab.Channel.path`."""
+        els = path.split("/")
+        assert len(els) == 2
+        ports = els[1][1:].split("_")
+        assert 1 <= len(ports) <= 2
+        return cls(
+            slot=int(els[0]),
+            ports=(int(ports[0]), int(ports[1]) if len(ports) == 2 else None),
+            input=els[1][0] == "i",
+        )
+
+    @property
+    def local_address(self):
+        """Physical address within the module.
+
+        It will generate a string in the format ``<direction><channel>`` or
+        ``<direction><I-channel>_<Q-channel>``.
+        ``<direction>`` is ``in`` for a connection between an input and the acquisition
+        path, ``out`` for a connection from the waveform generator to an output, or
+        ``io`` to do both.
+        The channels must be integer channel indices.
+        Only one channel is present for a real mode operating sequencer; two channels
+        are used for complex mode.
+
+        .. note::
+
+            Description adapted from
+            https://docs.qblox.com/en/main/api_reference/cluster.html#qblox_instruments.Cluster.connect_sequencer
+        """
+        direction = "in" if self.input else "out"
+        channels = (
+            str(self.ports[0])
+            if self.ports[1] is None
+            else f"{self.ports[0]}_{self.ports[1]}"
+        )
+        return f"{direction}{channels}"
+
+
+class Cluster(Controller):
+    name: str
+    """Device name.
+
+    As described in:
+    https://docs.qblox.com/en/main/getting_started/setup.html#connecting-to-multiple-instruments
+    """
+    bounds: str = "qblox/bounds"
+    _cluster: Optional[qblox.Cluster] = None
+
+    @cached_property
+    def _modules(self) -> dict[SlotId, Module]:
+        assert self._cluster is not None
+        return {mod.slot_idx: mod for mod in self._cluster.modules if mod.present()}
+
+    @property
+    def sampling_rate(self) -> int:
+        return SAMPLING_RATE
+
+    def connect(self):
+        if self.is_connected:
+            return
+
+        self._cluster = find_or_create_instrument(
+            qblox.Cluster, recreate=True, name=self.name, identifier=self.address
+        )
+        self._cluster.reset()
+
+    @property
+    def is_connected(self) -> bool:
+        return self._cluster is not None
+
+    def disconnect(self):
+        assert self._cluster is not None
+
+        for module in self._modules.values():
+            module.stop_sequencer()
+        self._cluster.reset()
+        self._cluster = None
+
+    def play(
+        self,
+        configs: dict[str, Config],
+        sequences: list[PulseSequence],
+        options: ExecutionParameters,
+        sweepers: list[ParallelSweepers],
+    ) -> dict[int, Result]:
+        results = {}
+        for ps in sequences:
+            sequences_ = _prepare(ps, sweepers, options)
+            sequencers = self._upload(sequences_)
+            results |= self._execute(sequencers)
+        return results
+
+    def _upload(
+        self, sequences: dict[ChannelId, Sequence]
+    ) -> dict[SlotId, dict[ChannelId, SequencerId]]:
+        channels_by_module = _channels_by_module()
+        sequencers = defaultdict(dict)
+        for mod, chs in channels_by_module.items():
+            module = self._modules[mod]
+            assert len(module.sequencers) > len(chs)
+            # Map sequencers to specific outputs (but first disable all sequencer connections)
+            module.disconnect_outputs()
+            for idx, (ch, sequencer) in enumerate(zip(chs, module.sequencers)):
+                sequencers[mod][ch] = idx
+                sequencer.sequence(sequences[ch].model_dump())
+                # Configure the sequencers to synchronize
+                sequencer.sync_en(True)
+                sequencer.connect_sequencer(
+                    PortAddress.from_path(self.channels[ch].path).local_address
+                )
+
+        return sequencers
+
+    def _execute(self, sequencers: dict[SlotId, dict[ChannelId, SequencerId]]) -> dict:
+        # TODO: implement
+        for mod, seqs in sequencers.items():
+            module = self._modules[mod]
+            for seq in seqs.values():
+                module.arm_sequencer(seq)
+            module.start_sequencer()
+
+        return {}
+
+
+def _channels_by_module() -> dict[SlotId, list[ChannelId]]:
+    # TODO: implement
+    return {}
+
+
+def _prepare(
+    sequence: PulseSequence,
+    sweepers: list[ParallelSweepers],
+    options: ExecutionParameters,
+) -> dict[ChannelId, Sequence]:
+    sequences = {}
+    for channel in sequence.channels:
+        filtered = PulseSequence((ch, pulse) for ch, pulse in sequence if ch == channel)
+        seq = Sequence.from_pulses(filtered, sweepers, options)
+        sequences[channel] = seq
+
+    return sequences

src/qibolab/_core/instruments/qblox/parse.py

@@ -4,6 +4,8 @@
 
 from .ast_ import INSTRUCTIONS, Comment, Line, Program
 
+__all__ = []
+
 GRAMMAR_FILE = Path(__file__).parent / "q1asm.lark"
 GRAMMAR = GRAMMAR_FILE.read_text()
 

src/qibolab/_core/instruments/qblox/platform.py

@@ -0,0 +1,97 @@
+from collections import defaultdict
+from typing import Optional, Union
+
+from qibolab._core.components.channels import (
+    AcquisitionChannel,
+    Channel,
+    DcChannel,
+    IqChannel,
+)
+
+__all__ = ["map_ports"]
+
+
+def _eltype(el: str):
+    return "qubits" if el[0] == "q" else "couplers"
+
+
+def _chtype(mod: str, input: bool) -> tuple[str, type[Channel]]:
+    if mod.startswith("qcm_rf"):
+        return "drive", IqChannel
+    if mod.startswith("qcm"):
+        return "flux", DcChannel
+    if mod.startswith("qrm_rf"):
+        if input:
+            return "acquisition", AcquisitionChannel
+        return "probe", IqChannel
+    raise ValueError
+
+
+def _port_channels(mod: str, port: Union[int, str], slot: int) -> dict:
+    if isinstance(port, str) and port.startswith("io"):
+        return {
+            "probe": IqChannel(path=f"{slot}/o{port[2:]}"),
+            "acquisition": AcquisitionChannel(path=f"{slot}/o{port[2:]}"),
+        }
+    port = f"o{port}" if isinstance(port, int) else port
+    name, cls = _chtype(mod, port[0] == "i")
+    return {name: cls(path=f"{slot}/{port}")}
+
+
+def _premap(cluster: dict):
+    d = {
+        "qubits": defaultdict(lambda: defaultdict(dict)),
+        "couplers": defaultdict(lambda: defaultdict(dict)),
+    }
+
+    for mod, props in cluster.items():
+        slot = props[0]
+        for port, els in props[1].items():
+            for el in els:
+                nel = el[1:]
+                d[_eltype(el)][nel] |= _port_channels(mod, port, slot)
+
+    return d
+
+
+def map_ports(cluster: dict, qubits: dict, couplers: Optional[dict] = None) -> dict:
+    """Extract channels from compact representation.
+
+    Conventions:
+    - each item is a module
+    - the first element of each value is the module's slot ID
+    - the second element is a map from ports to qubits
+        - ports
+            - they are `i{n}` or `o{n}` for the inputs and outputs respectively
+            - `io{n}` is also allowed, to signal that both are connected (cater for the specific
+              case of the QRM_RF where there are only one port of each type)
+            - if it's just an integer, it is intended to be an output (equivalent to `o{n}`)
+        - values
+            - list of element names
+            - they are `q{name}` or `c{name}` for qubits and couplers respectively
+            - multiple elements are allowed, for multiplexed ports
+
+    .. note::
+
+        integer qubit names are not allowed
+
+    .. todo::
+
+        Currently channel types are inferred from the module type, encoded in its name. At
+        least an override should be allowed (per module, per port, per qubit).
+    """
+    if couplers is None:
+        couplers = {}
+
+    premap = _premap(cluster)
+
+    channels = {}
+    for kind, elements in [("qubits", qubits), ("couplers", couplers)]:
+        for name, el in elements.items():
+            for chname, ch in premap[kind][name].items():
+                channels[getattr(el, chname)] = ch
+            if kind == "qubits":
+                channels[el.acquisition] = channels[el.acquisition].model_copy(
+                    update={"probe": el.probe}
+                )
+    return channels

src/qibolab/_core/instruments/qblox/sequence.py

@@ -0,0 +1,46 @@
+from typing import Annotated
+
+from pydantic import AfterValidator, PlainSerializer, PlainValidator
+
+from qibolab._core.execution_parameters import ExecutionParameters
+from qibolab._core.sequence import PulseSequence
+from qibolab._core.serialize import ArrayList, Model
+from qibolab._core.sweeper import ParallelSweepers
+
+from .ast_ import Program
+from .parse import parse
+
+__all__ = []
+
+
+class Waveform(Model):
+    data: Annotated[ArrayList, AfterValidator(lambda a: a.astype(float))]
+    index: int
+
+
+Weight = Waveform
+
+
+class Acquisition(Model):
+    num_bins: int
+    index: int
+
+
+class Sequence(Model):
+    waveforms: dict[str, Waveform]
+    weights: dict[str, Weight]
+    acquisitions: dict[str, Acquisition]
+    program: Annotated[
+        Program, PlainSerializer(lambda p: p.asm()), PlainValidator(parse)
+    ]
+
+    @classmethod
+    def from_pulses(
+        cls,
+        sequence: PulseSequence,
+        sweepers: list[ParallelSweepers],
+        options: ExecutionParameters,
+    ):
+        return cls(
+            waveforms={}, weights={}, acquisitions={}, program=Program(elements=[])
+        )