from __future__ import annotations
import logging
import tomllib
from pathlib import Path
from typing import TYPE_CHECKING, TypeVar
from semi_cr.core.helpers import deprecated
from semi_cr.core.station.graph import Connector, GraphEdge, GraphNode, StationGraph
from semi_cr.core.station.hydration import UnloadedSpec, hydrate_station_config
from semi_cr.core.station.instruments import LabInstrument
from semi_cr.core.station.manifest import NQCPStationManifest
StaticInstrumentT = TypeVar("StaticInstrumentT", bound=LabInstrument)
if TYPE_CHECKING:
from collections.abc import Iterable
from semi_cr.core.lab.station.base import ConfigBundle
from semi_cr.core.station.chip import Chip
from semi_cr.core.station.device import Gate, QuantumDevice
SETUP_DIR = Path(__file__).parent.parent.parent / "setup"
[docs]
class NQCPStation:
config: ConfigBundle
manifest: NQCPStationManifest
graph: StationGraph
chips: list[Chip]
connectors: list[Connector]
instruments: list[LabInstrument]
unloaded_specs: list[UnloadedSpec]
def __init__(self, manifest: str | Path | NQCPStationManifest) -> None:
"""Sets up a station handler, given a manifest. In the constructor we establish an object-oriented model of
the station constituents, meaning instruments, wiring and the quantum device model as well."""
logging.info("Instantiating NQCPStation, manifest path: %s", manifest)
# If `manifest` is a string or even Path, we may assume it's a file path to a station manifest.
if isinstance(manifest, str):
manifest = NQCPStationManifest(Path(manifest))
if isinstance(manifest, Path):
manifest = NQCPStationManifest(manifest)
self.manifest = manifest
self.config = manifest.get_config_bundle()
# Turning our configuration into "live" objects is something we call "hydration". Think of it as adding water
# to a dry configuration. This is what happens implicitly in QCoDeS when loading a configuration; it looks at
# the config and creates QCoDeS objects for each "instrument" in the configuration. We kind of do the same;
# declare one hydration handler per type (in hydration.py), call the hydration function and then pull it back
# into the station here.
hydrated = hydrate_station_config(self.config, self)
self.chips = hydrated.chips
self.connectors = hydrated.connectors
self.instruments = hydrated.instruments
self.unloaded_specs = hydrated.unloaded_specs
logging.info("Hydrated %d chips, %d connectors and %d instruments", len(self.chips), len(self.connectors), len(self.instruments))
if self.unloaded_specs:
logging.info("Skipped %d specs without static loaders", len(self.unloaded_specs))
logging.info("Extracting routable information from loaded chips and instruments")
logging.info("Building a station graph")
# Once we know about our chips and instruments, we can start to wire up the station.
self.graph = StationGraph.from_routable(self)
@property
def available_chips(self) -> list[Chip]:
return self.chips
@property
def chip(self) -> Chip:
"""Returns the configuration of the currently loaded chip"""
return self.get_chip()
[docs]
def get_chip(self, chip_id: str | None = None) -> Chip:
"""Gets a chip by name.
:returns: The chip matching ``chip_id``, or the loaded chip when omitted.
:raises KeyError: If no matching chip exists.
"""
chip_id = chip_id or self.manifest.loaded_chip_id
for chip in self.chips:
if chip.chip_id == chip_id:
return chip
raise KeyError(f"Chip not found: {chip_id}")
@property
def device(self) -> QuantumDevice:
return self.get_device()
[docs]
def get_device(self, device_id: str | None = None) -> QuantumDevice:
"""Gets a device on the loaded chip.
:returns: The named device, or the only device on the loaded chip when omitted.
:raises ValueError: If ``device_id`` is omitted and the loaded chip has multiple devices.
"""
devices = self.chip.devices
if device_id is not None:
return self.chip.device(device_id)
if len(devices) != 1:
raise ValueError("device_id is required when the loaded chip has multiple devices")
return devices[0]
[docs]
def get_gates(self, device_id: str | None = None) -> list[Gate]:
"""Gets gates on the station.
:returns: All gates, or gates on ``device_id`` when provided.
"""
if device_id is not None:
return self.get_device(device_id).gates
gates = []
for chip in self.chips:
for device in chip.devices:
for gate in device.gates:
gates.append(gate)
return gates
[docs]
def get_instruments_of_type(self, instrument_type: type[StaticInstrumentT], *, name: str | None = None) -> list[StaticInstrumentT]:
"""Gets static instruments matching ``instrument_type`` and optional ``name``."""
return [instrument for instrument in self.instruments if isinstance(instrument, instrument_type) and (name is None or instrument.name == name)]
[docs]
def get_instrument_of_type(self, instrument_type: type[StaticInstrumentT], *, name: str | None = None) -> StaticInstrumentT:
"""Gets exactly one static instrument matching ``instrument_type`` and optional ``name``.
:raises LookupError: If zero or multiple instruments match.
"""
matches = self.get_instruments_of_type(instrument_type, name=name)
if len(matches) != 1:
names = [instrument.name for instrument in matches]
raise LookupError(f"Expected one {instrument_type.__name__} instrument, found {names}")
return matches[0]
[docs]
@deprecated("We need to figure out if we can do this supporting multiple chips, devices etc.")
def gate(self, gate_name: str) -> Gate:
"""Gets a gate by name. This is kind of a placeholder method because it kind of glosses over the fact that a
station may contain more than one chip with more than one device, and so a gate name can't be considered a
unique property.
:return: The gate matching the requested name
:raises: KeyError if the gate doesn't exist
:raises: ValueError if there is more than one gate with that name"""
resulting_set = list(filter(lambda gate: gate.name == gate_name, self.get_gates()))
if len(resulting_set) == 0:
raise KeyError(f"No gate found with the name {gate_name}.")
if len(resulting_set) > 1:
raise ValueError(f"Gate name {gate_name} not unique for the station.")
return resulting_set[0]
def __rich_repr__(self): # noqa: PLW3201
yield "Name", self.manifest.name
[docs]
@classmethod
def available_stations(cls) -> dict[str, NQCPStationManifest]:
"""
Finds setup manifests under `semi_cr.setup`.
"""
setups: dict[str, NQCPStationManifest] = {}
# We find `station.toml` files, extract the names of them in order to let the user pick a station name.
for manifest_path in SETUP_DIR.glob("*/station.toml"):
# Quickly peek in the manifest to extract the name of the station.
with manifest_path.open("rb") as file:
manifest = tomllib.load(file)
station_name: str = manifest["name"]
setups[station_name] = NQCPStationManifest(manifest_path)
# `setups` now contains a name -> manifest dict
return setups
[docs]
@classmethod
def get(cls, name: str) -> NQCPStation:
"""
Load a station description without connecting to hardware.
"""
stations = NQCPStation.available_stations()
try:
return NQCPStation(stations[name])
except KeyError:
raise KeyError(f"Station {name} not found")
[docs]
def graph_nodes(self) -> Iterable[GraphNode]:
"""Iterates over loaded objects to return graph nodes for this station."""
for chip in self.chips:
yield from chip.graph_nodes()
for connector in self.connectors:
yield from connector.graph_nodes()
for instrument in self.instruments:
yield from instrument.graph_nodes()
[docs]
def graph_edges(self) -> Iterable[GraphEdge]:
"""Iterates over loaded objects to return graph edges for this station."""
for chip in self.chips:
yield from chip.graph_edges()
for connector in self.connectors:
yield from connector.graph_edges()
for instrument in self.instruments:
yield from instrument.graph_edges()