from dataclasses import dataclass, field
from io import StringIO
from typing import Any, Literal
from collections.abc import Iterable
import networkx as nx
from qcodes.instrument import Instrument
from ruamel.yaml import YAML
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.routing_v1.models import ConnectionRoutable
from semi_cr.core.lab.station.graphing.ids import qcodes_node_id
from semi_cr.core.lab.station.graphing.context import ConfigContext
from semi_cr.core.lab.station.graphing.resolver import GraphNodeResolver
from semi_cr.core.lab.station.graphing.runtime_builder import ContainmentGraphBuilder
from semi_cr.core.lab.station.graphing.overlay import DependencyOverlayBuilder
from semi_cr.core.lab.station.graphing.static_builder import build_wiring_graph
from semi_cr.core.lab.station.graphing.view import graph_view
from semi_cr.core.lab.station.metadata import LabStation
from semi_cr.core.lab.station.runtime import iter_routables
from semi_cr.core.lab.station.base import (
GraphRelation,
InstrumentSpec,
static_instrument_node_id,
ConfigBundle,
)
from copy import deepcopy
# %%
[docs]
@dataclass(frozen=True)
class StaticTerminalRef:
instrument: str
module: str | None
terminal: str
# def resolve_instrument_name(
# self,
# name: str,
# ) -> str:
# return self.static.instruments[
# name
# ].runtime_name
[docs]
@dataclass
class InfrastructureContext:
wiring_config: ConfigBundle
electronics_config: ConfigBundle
config: ConfigBundle = field(init=False)
graph: nx.MultiDiGraph = field(init=False)
def __post_init__(self) -> None:
self.config = ConfigBundle.combine(
self.wiring_config,
self.electronics_config,
)
self.graph = self.build_graph()
[docs]
def build_graph(self) -> nx.MultiDiGraph:
return build_wiring_graph(self.config.config)
[docs]
def view_graph(
self,
what: Literal["nodes", "edges", "both"] = "both",
include_attrs: bool = True,
):
return graph_view(self.graph, what=what, include_attrs=include_attrs)
[docs]
def build_instrument_specs(
config: dict[str, Any],
) -> dict[str, InstrumentSpec]:
return {
name: InstrumentSpec(
name=name,
type=cfg["type"],
init=dict(cfg.get("init", {})),
)
for name, cfg
in config.get("instruments", {}).items()
}
[docs]
@dataclass
class StaticStationContext:
"""
Description of the complete station without talking to hardware.
"""
device: ConfigContext
infrastructure: InfrastructureContext
config: ConfigBundle = field(init=False)
instruments: dict[str, InstrumentSpec] = field(init=False)
graph: nx.MultiDiGraph = field(init=False)
def __post_init__(self) -> None:
self.config = ConfigBundle.combine(
self.device.config,
self.infrastructure.config,
)
self.instruments = build_instrument_specs(
self.config.config
)
self.graph = build_wiring_graph(
self.config.config
)
[docs]
@dataclass(frozen=True)
class RuntimeBinding:
static_id: str
runtime_id: str
runtime_object: Any
[docs]
class RuntimeBindingRegistry:
def __init__(self) -> None:
self._by_static_id: dict[str, RuntimeBinding] = {}
[docs]
def bind(
self,
static_id: str,
runtime_id: str,
runtime_object: Any,
) -> None:
self._by_static_id[static_id] = RuntimeBinding(
static_id=static_id,
runtime_id=runtime_id,
runtime_object=runtime_object,
)
[docs]
def resolve(self, static_id: str) -> Any:
try:
return self._by_static_id[static_id].runtime_object
except KeyError as exc:
raise KeyError(
f"No runtime binding exists for {static_id!r}."
) from exc
[docs]
def get_binding(
self,
static_id: str,
) -> RuntimeBinding:
return self._by_static_id[static_id]
[docs]
def is_bound(self, static_id: str) -> bool:
return static_id in self._by_static_id
[docs]
def items(self):
return self._by_static_id.items()
[docs]
def values(self):
return self._by_static_id.values()
def __len__(self) -> int:
return len(self._by_static_id)
[docs]
def normalize_qcodes_config(
config: dict[str, Any],
) -> dict[str, Any]:
config = deepcopy(config)
instruments = config.get("instruments", {})
for instrument_cfg in instruments.values():
init = instrument_cfg.get("init", {})
devices = init.get("devices")
if isinstance(devices, dict):
init["devices"] = {
str(getattr(device_id, "value", device_id)): device_cfg
for device_id, device_cfg in devices.items()
}
return config
# %%
from enum import Enum, auto
import logging
logger = logging.getLogger(__name__)
[docs]
class ContextState(Enum):
CREATED = auto()
LOADED = auto()
PARTIALLY_CONNECTED = auto()
CONNECTED = auto()
[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)
# _instruments_by_config_name: 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_components_by_config_name: dict[str, Any] = field(init=False, default_factory=dict, 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 instruments_by_config_name(
# self,
# ) -> dict[str, Instrument]:
# return self._instruments_by_config_name
@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._runtime_components_by_config_name
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._runtime_components_by_config_name
self._instruments
)
@property
def has_connections(self) -> bool:
"""
True when at least one instrument owned by this context
is currently connected.
"""
return bool(
# self._runtime_components_by_config_name
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 _load_runtime_instruments(self) -> None:
# """
# Refresh runtime instrument views from components owned by this context.
# Only instruments that have actually been instantiated/connected
# are included.
# """
# self._instruments_by_config_name = {
# config_name: component
# for config_name, component
# in self._runtime_components_by_config_name.items()
# if (
# isinstance(component, Instrument)
# and not isinstance(component, Chip)
# and not isinstance(component, PCB)
# )
# }
# self._instruments = {}
# for config_name, instrument in (
# self._instruments_by_config_name.items()
# ):
# station_name = self._station_component_name(
# instrument
# )
# # Prefer the station registration name.
# # Fall back to the QCoDeS instrument name.
# name = (
# station_name
# if station_name is not None
# else instrument.name
# )
# self._instruments[name] = instrument
# self._connectors = tuple(
# instrument
# for instrument in self._instruments.values()
# if isinstance(instrument, Connector)
# )
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,
)
# setattr(
# instrument,
# "config_name",
# config_name,
# )
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._instruments_by_config_name.clear()
self._connectors = ()
self._pcbs = ()
self._chips = ()
self._devices = ()
def _clear_runtime_state(self) -> None:
# self._runtime_components_by_config_name.clear()
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._load_runtime_instruments()
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
# 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,
# )