from dataclasses import dataclass, field
import networkx as nx
from typing import Any
from collections.abc import Iterable
from ruamel.yaml import YAML
from io import StringIO
from qcodes.instrument import Instrument
from semi_cr.core.lab.devices.base import BaseDevice
from semi_cr.core.lab.instrument.chip import Chip
from semi_cr.core.lab.instrument.connector import Connector
from semi_cr.core.lab.instrument.pcb import PCB
from semi_cr.core.lab.station.metadata import LabStation
from semi_cr.core.lab.station.graphing.ids import qcodes_node_id
from semi_cr.core.lab.station.base import (
static_instrument_node_id,
)
from .models import ContextState
from .static import StaticStationContext
from .graph import StationGraphBuilder
from .bindings import RuntimeBindingRegistry
from .config import normalize_qcodes_config
import logging
logger = logging.getLogger(__name__)
[docs]
@dataclass
class StationContext:
"""
Runtime realization of a StaticStationContext.
Construction is intentionally side-effect free.
Lifecycle:
context = StationContext(station, static)
context.load()
context.connect()
...
context.disconnect()
"""
station: LabStation
static: StaticStationContext
_state: ContextState = field(init=False, default=ContextState.CREATED)
_bindings: RuntimeBindingRegistry = field(init=False, default_factory=RuntimeBindingRegistry)
_instruments: dict[str, Instrument] = field(init=False, default_factory=dict, repr=False)
_pcbs: tuple[PCB, ...] = field(init=False, default_factory=tuple, repr=False)
_chips: tuple[Chip, ...] = field(init=False, default_factory=tuple, repr=False)
_devices: tuple[BaseDevice, ...] = field(init=False, default_factory=tuple, repr=False)
_connectors: tuple[Connector, ...] = field(init=False, default_factory=tuple, repr=False)
_graph: nx.MultiDiGraph = field(init=False, default_factory=nx.MultiDiGraph, repr=False)
# ---------------------------------------------------------
# Runtime binding
# ---------------------------------------------------------
@property
def bindings(self) -> RuntimeBindingRegistry:
return self._bindings
[docs]
def resolve_runtime(self, static_id: str) -> Any:
"""
Resolve a static object ID to its runtime realization.
Runtime resolution is available once bindings have been created,
including while the context is being connected.
"""
try:
return self.bindings.resolve(static_id)
except KeyError as exc:
raise LookupError(
f"No runtime binding exists for {static_id!r}."
) from exc
# ---------------------------------------------------------
# Runtime collections
# ---------------------------------------------------------
@property
def instruments(self) -> dict[str, Instrument]:
return self._instruments
@property
def devices(self) -> tuple[BaseDevice, ...]:
return self._devices
@property
def chips(self) -> tuple[Chip, ...]:
return self._chips
@property
def pcbs(self) -> tuple[PCB, ...]:
return self._pcbs
@property
def connectors(self) -> tuple[Connector, ...]:
return self._connectors
@property
def graph(self) -> nx.MultiDiGraph:
return self._graph
@property
def modules(self) -> dict[str, Any]:
modules: dict[str, Any] = {}
for instrument in self._instruments.values():
inst_modules = getattr(
instrument,
"modules",
None,
)
if inst_modules is None:
continue
if hasattr(inst_modules, "items"):
modules.update(inst_modules.items())
return modules
@property
def has_runtime(self) -> bool:
return bool(self._instruments)
@property
def state(self) -> ContextState:
return self._state
def _update_state(self) -> None:
if self._state is ContextState.CREATED:
return
n_configured = len(self.static.instruments)
n_connected = len(
self._instruments
)
if n_connected == 0:
self._state = ContextState.LOADED
elif n_connected == n_configured:
self._state = ContextState.CONNECTED
else:
self._state = ContextState.PARTIALLY_CONNECTED
@property
def loaded(self) -> bool:
return self._state in {
ContextState.LOADED,
ContextState.PARTIALLY_CONNECTED,
ContextState.CONNECTED,
}
@property
def connected(self) -> bool:
"""
True when all configured instruments are connected.
"""
return self._state is ContextState.CONNECTED
@property
def partially_connected(self) -> bool:
return self._state is ContextState.PARTIALLY_CONNECTED
@property
def connected_instrument_names(self) -> tuple[str, ...]:
return tuple(
self._instruments
)
@property
def has_connections(self) -> bool:
"""
True when at least one instrument owned by this context
is currently connected.
"""
return bool(
self._instruments
)
def _resolve_instrument_names(
self,
instruments: str | Iterable[str] | None,
) -> tuple[str, ...]:
"""
Normalize an instrument selector to QCoDeS config names.
Parameters
----------
instruments:
None
Select every configured instrument.
str
Select one instrument by config name.
Iterable[str]
Select multiple instruments by config name.
"""
if instruments is None:
names = tuple(self.static.instruments)
elif isinstance(instruments, str):
names = (instruments,)
else:
names = tuple(instruments)
unknown = [
name
for name in names
if name not in self.static.instruments
]
if unknown:
available = ", ".join(self.static.instruments)
raise KeyError(
f"Unknown instrument name(s): {unknown}. "
f"Available instruments: {available}"
)
# Preserve order while removing duplicates.
return tuple(dict.fromkeys(names))
def _load_qcodes_config(self) -> None:
config = normalize_qcodes_config(
self.static.config.config
)
yaml_text = StringIO()
yaml_rt = YAML()
yaml_rt.default_flow_style = False
yaml_rt.dump(config, yaml_text)
self.station.load_config(
yaml_text.getvalue()
)
def _rebuild_graph(self) -> None:
self._graph = self._create_graph()
def _station_component_name(
self,
component: Any,
) -> str | None:
for name, registered_component in self.station.components.items():
if registered_component is component:
return name
return None
def _refresh_station_objects(self) -> None:
self._pcbs = tuple(
obj
for obj in self.station.components.values()
if isinstance(obj, PCB)
)
self._chips = tuple(
obj
for obj in self.station.components.values()
if isinstance(obj, Chip)
)
self._devices = tuple(
device
for chip in self._chips
for device in chip.devices
)
def _refresh_runtime_views(self) -> None:
self._connectors = tuple(
instrument
for instrument in self._instruments.values()
if isinstance(instrument, Connector)
)
def _connect_instruments(
self,
config_names: Iterable[str],
) -> None:
for config_name in config_names:
# Idempotent: already connected by this context.
if config_name in self._instruments:
continue
spec = self.static.instruments[config_name]
instrument = self.station.load_instrument(
config_name,
**spec.init,
)
self._instruments[
config_name
] = instrument
static_id = static_instrument_node_id(spec)
runtime_id = qcodes_node_id(instrument)
self._bindings.bind(
static_id=static_id,
runtime_id=runtime_id,
runtime_object=instrument,
)
def _clear_runtime_views(self) -> None:
self._instruments.clear()
self._connectors = ()
self._pcbs = ()
self._chips = ()
self._devices = ()
def _clear_runtime_state(self) -> None:
self._clear_runtime_views()
self._bindings = RuntimeBindingRegistry()
def _rebuild_runtime_bindings(self) -> None:
bindings = RuntimeBindingRegistry()
for name, component in (
self._instruments.items()
):
spec = self.static.instruments[name]
bindings.bind(
static_id=static_instrument_node_id(spec),
runtime_id=qcodes_node_id(component),
runtime_object=component,
)
self._bindings = bindings
# ---------------------------------------------------------
# Lifecycle functions
# ---------------------------------------------------------
[docs]
def load(self) -> None:
"""
Load the static configuration into the QCoDeS station.
This does not connect to physical instruments.
"""
if self.loaded:
return
self._load_qcodes_config()
self._rebuild_graph()
self._state = ContextState.LOADED
[docs]
def connect(
self,
instruments: str | Iterable[str] | None = None,
) -> None:
"""
Connect one, several, or all configured instruments.
Parameters
----------
instruments:
QCoDeS instrument config name, iterable of config names,
or None to connect all configured instruments.
Examples
--------
context.connect()
context.connect("qdac_1")
context.connect(["qdac_1", "smu"])
"""
if not self.loaded:
self.load()
names = self._resolve_instrument_names(
instruments
)
# Only connect instruments not already connected.
missing = tuple(
name
for name in names
if name not in self._instruments
)
if not missing:
return
self._connect_instruments(missing)
# Components such as PCB/Chip may now exist in
# station.components.
self._refresh_station_objects()
self._refresh_runtime_views()
self._rebuild_graph()
self._update_state()
def _disconnect_instrument(
self,
name: str,
) -> None:
component = self._instruments.get(name)
if component is None:
# Already disconnected.
return
try:
if isinstance(component, Instrument):
component.close()
except Exception:
logger.exception("Failed to close component %s",name)
finally:
# Important:
# Instrument.close() may itself modify registration state.
#
# Therefore locate the component AFTER close(), rather than
# relying on the name found before close().
station_name = self._station_component_name(
component
)
if station_name is not None:
self.station.remove_component(
station_name
)
self._instruments.pop(
name,
None,
)
[docs]
def disconnect(
self,
instruments: str | Iterable[str] | None = None,
) -> None:
"""
Disconnect one, several, or all runtime instruments
instantiated by this context.
QCoDeS configuration remains loaded.
Parameters
----------
instruments:
Config name, iterable of config names, or None to
disconnect all instruments owned by this context.
Examples
--------
context.disconnect()
context.disconnect("qdac_1")
context.disconnect(["qdac_1", "smu"])
"""
if not self.loaded:
return
if instruments is None:
# Reverse connection order when shutting everything down.
names = tuple(reversed(tuple(self._instruments)))
else:
names = self._resolve_instrument_names(instruments)
for name in names:
self._disconnect_instrument(name)
# Bindings corresponding to removed instruments must disappear.
self._rebuild_runtime_bindings()
if self._instruments:
self._refresh_station_objects()
self._refresh_runtime_views()
else:
self._clear_runtime_views()
self._rebuild_graph()
self._update_state()
# def _add_runtime_bindings_to_graph(
# self,
# graph: nx.MultiDiGraph,
# ) -> None:
# for static_id, binding in self._bindings.items():
# if static_id not in graph:
# raise KeyError(
# f"Static node {static_id!r} is missing "
# "from the station graph."
# )
# if binding.runtime_id not in graph:
# raise KeyError(
# f"Runtime node {binding.runtime_id!r} "
# "is missing from the station graph."
# )
# graph.add_edge(
# static_id,
# binding.runtime_id,
# relation=GraphRelation.REPRESENTS,
# )
# def _create_graph(
# self,
# ) -> nx.MultiDiGraph:
# # -----------------------------------------------------
# # 1. Canonical static topology
# # -----------------------------------------------------
# static_graphs = [
# self.static.graph,
# *(
# device.graph
# for device in self.devices
# ),
# ]
# canonical_graph = nx.compose_all(
# static_graphs
# )
# if not self.has_runtime:
# return canonical_graph
# # -----------------------------------------------------
# # 2. Runtime QCoDeS containment
# # -----------------------------------------------------
# containment_builder = (
# ContainmentGraphBuilder(
# node_id=qcodes_node_id,
# attach_objects=True,
# )
# )
# runtime_containment = (
# containment_builder.build(
# self.station
# )
# )
# base_graph = nx.compose(
# canonical_graph,
# runtime_containment,
# )
# # -----------------------------------------------------
# # 3. Collect runtime routable graphs
# # -----------------------------------------------------
# runtime_graphs: list[
# nx.MultiDiGraph
# ] = []
# seen_routables: set[int] = set()
# for instrument in (
# self._instruments.values()
# ):
# candidates = [
# instrument,
# *iter_routables(instrument),
# ]
# for routable in candidates:
# routable_id = id(routable)
# if routable_id in seen_routables:
# continue
# seen_routables.add(
# routable_id
# )
# if not isinstance(
# routable,
# ConnectionRoutable,
# ):
# continue
# runtime_graphs.append(
# routable.graph
# )
# # -----------------------------------------------------
# # 4. Merge runtime metadata into canonical terminals
# # -----------------------------------------------------
# resolver = GraphNodeResolver(
# qcodes_id=qcodes_node_id,
# )
# overlay_builder = (
# DependencyOverlayBuilder(
# resolver=resolver,
# )
# )
# graph = overlay_builder.overlay(
# base_graph,
# runtime_graphs,
# )
# # -----------------------------------------------------
# # 5. Bind static instruments to runtime instruments
# # -----------------------------------------------------
# self._add_runtime_bindings_to_graph(
# graph
# )
# return graph
def _create_graph(self) -> nx.MultiDiGraph:
return StationGraphBuilder(
station=self.station
).build(
static_graph=self.static.graph,
devices=self.devices,
instruments=self.instruments,
bindings=self.bindings,
)
# COULD BE IMPLEMENTED HERE LATER
# def _attach_routed_parameters(self) -> None:
# from semi_cr.core.lab.station.routing_v1.attachment import (
# # attach_routed_pin_parameters,
# attach_routed_pin_quantities,
# )
# from semi_cr.core.lab.devices.pinned import PinnedDevice
# # for instrument in self.instruments.values():
# # devices = getattr(instrument, "devices", None)
# # if devices is None:
# # continue
# for device in self.devices:
# if isinstance(device, PinnedDevice):
# # attach_routed_pin_parameters(
# # graph=self.graph,
# # device=device,
# # )
# attach_routed_pin_quantities(
# graph=self.graph,
# device=device,
# )