from __future__ import annotations
import enum
from collections.abc import Mapping
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Protocol, overload
import networkx as nx
from semi_cr.core.lab.station.graphing.ids import terminal_id
from semi_cr.core.lab.station.graphing.parsing import parse_endpoint
from semi_cr.core.station.drivers import DCDriver
from semi_cr.core.station.instruments.qcodes_proxy import QCoDeSInstrumentSpec
if TYPE_CHECKING:
from collections.abc import Callable, Iterable, Iterator
from semi_cr.core.station.device import DevicePin
[docs]
def graph_attributes(config: Mapping[str, Any]) -> dict[str, Any]:
"""Return the explicit graph metadata declared in a YAML config section.
Configuration keys are intentionally not copied wholesale: only values
under ``attributes`` become graph metadata. This keeps constructor
settings such as ``type`` and ``slot`` out of the graph namespace.
"""
raw = config.get("attributes", {})
if raw is None:
return {}
if not isinstance(raw, Mapping):
raise TypeError("Graph attributes must be a mapping")
return dict(raw)
[docs]
@dataclass(frozen=True)
class GraphNode:
id: str
obj: object
kind: str
attrs: Mapping[str, Any] = field(default_factory=dict)
[docs]
class Conduction(enum.Flag):
DC = enum.auto()
AC = enum.auto()
[docs]
@dataclass(frozen=True)
class GraphEdge:
# TODO: Validate that this is a good data model.
id: str
source: str
target: str
kind: str
attrs: Mapping[str, Any] = field(default_factory=dict)
conducts: Conduction = Conduction.AC | Conduction.DC # by default we path both DC and AC over an edge
[docs]
class Routable(Protocol):
"""Use this protocol for anything that can add nodes or edges to the station graph. The idea is that devices
export themselves as nodes at least. If they have internal edges that need to exist for the routing to work,
they can supply those, too."""
[docs]
def graph_nodes(self) -> Iterable[GraphNode]: ...
[docs]
def graph_edges(self) -> Iterable[GraphEdge]: ...
[docs]
@dataclass
class Connector:
"""Implements a non-QCoDeS variant of semi_cr.core.lab.instrument.connector.Connector, a class representing a
junction or set of junctions between two instruments in our setup.
The class is thought to be instantiated using the `from_instrument_spec` class method, and exposes the
:class:`Routable` protocol."""
name: str
# TODO: I don't really like this property of our current configuration schema; that one connection can have two
# endpoints. I think it'd be a better model to just have two edges out of the same origin.
connections: Mapping[str, list[str]]
connection_attributes: Mapping[str, Mapping[str, Any]] = field(default_factory=dict)
[docs]
def graph_nodes(self) -> Iterable[GraphNode]:
"""Returns graph nodes for this connector's terminals."""
for connection_name in self.connections:
yield GraphNode(
id=terminal_id(self.name, connection_name),
obj=self,
kind="terminal",
attrs={
"instrument": self.name,
"terminal": connection_name,
**self.connection_attributes.get(connection_name, {}),
},
)
[docs]
def graph_edges(self) -> Iterable[GraphEdge]:
"""Emits graph edges from the connections stored on the Connector element."""
for connection_name, endpoints in self.connections.items():
local = terminal_id(self.name, connection_name)
for endpoint in endpoints:
yield GraphEdge(
id=f"{endpoint}->{local}",
source=endpoint,
target=local,
kind="physical",
)
[docs]
@classmethod
def from_instrument_spec(cls, spec: QCoDeSInstrumentSpec) -> Connector:
name, config = spec
# Walks through the connections declared in the YaML format and creates as-yet unresolved connections
connections: dict[str, list[str]] = {}
connection_attributes: dict[str, Mapping[str, Any]] = {}
for connection in config.get("init", {}).get("connections", []):
connection_name = connection["name"]
endpoints = list(connection.get("endpoints", []))
connections[connection_name] = endpoints
connection_attributes[connection_name] = graph_attributes(connection)
return cls(name=name, connections=connections, connection_attributes=connection_attributes)
[docs]
def build_node_aliases(graph: nx.Graph) -> dict[str, str]:
"""In some config references there seems to be a convention that e.g. tank inductors can be referred to by
concatenating the chip ID with the device ID, such as B083004D010033 where B083004D01 is the chip part and 0033
is the device part. When configurations refer to endpoints like this the graph resolver gets confused. So for
each device pin we add an alias that we can use when looking up the node to make edges later on."""
aliases = {}
for node_id, data in graph.nodes(data=True):
if data.get("kind") != "device_pin":
continue
obj = data["obj"]
chip_id = obj.parent_ref.parent_ref.chip_id
device_id = obj.parent_ref.device_id
aliases[f"{chip_id}{device_id}[{obj.name}]"] = node_id
aliases[f"{chip_id}/{device_id}[{obj.name}]"] = node_id
return aliases
[docs]
def resolve_endpoint_reference(reference: str, aliases: Mapping[str, str]) -> str:
"""Resolve an edge endpoint reference to a canonical graph node ID.
Config strings like ``foo[bar]`` are ambiguous: they may refer to a modeled device pin, or to a wiring terminal.
Device-pin aliases are tried first. If no alias matches, the string is treated as a terminal reference. Existing
``lab://`` node IDs are already canonical and are returned unchanged.
"""
if reference in aliases:
return aliases[reference]
if reference.startswith("lab://"):
return reference
instrument, channel = parse_endpoint(reference)
return terminal_id(instrument, channel)
[docs]
@dataclass
class StationGraph:
graph: nx.Graph = field(default_factory=nx.Graph)
aliases: dict[str, str] = field(default_factory=dict)
[docs]
def opx_source_for_pin(self, pin: DevicePin) -> object | None:
"""Return the OPX source terminal reachable from a device pin.
:param pin: Device pin object with a ``node_id``.
:returns: The first reachable OPX source terminal, or ``None``.
"""
from semi_cr.core.station.instruments.opx1000 import OPX1000LFTerminal # noqa: PLC0415
return self.opx_terminal_for_pin(pin, OPX1000LFTerminal.Role.Source)
[docs]
def opx_receiver_for_pin(self, pin: DevicePin) -> object | None:
"""Return the OPX receiver terminal reachable from a device pin.
:param pin: Device pin object with a ``node_id``.
:returns: The first reachable OPX receiver terminal, or ``None``.
"""
from semi_cr.core.station.instruments.opx1000 import OPX1000LFTerminal # noqa: PLC0415
return self.opx_terminal_for_pin(pin, OPX1000LFTerminal.Role.Receiver)
[docs]
def opx_terminal_for_pin(self, pin: DevicePin, role: object) -> object | None:
"""Return the first OPX terminal with ``role`` reachable from a device pin.
Driver lookup is rooted at the pin and walks the station topology outward. Matching is based on the graph
node's ``obj`` relation, not graph metadata.
:param pin: Device pin object with a ``node_id``.
:param role: :class:`OPX1000LFTerminal.Role` member to match.
:returns: The first matching OPX terminal object, or ``None``.
"""
from semi_cr.core.station.instruments.opx1000 import OPX1000LFTerminal # noqa: PLC0415
return next(
self.conducting(Conduction.AC).reachable_from_pin(
pin,
OPX1000LFTerminal,
lambda terminal: terminal.role is role,
),
None,
)
[docs]
def dc_driver_for_pin(self, pin: DevicePin) -> DCDriver | None:
"""Gets the DC driver for a given pin.
It does this by extracting the subgraph that only contains DC-carrying edges, and find a path in that one from
the pin to any DCDriver."""
drivers = list(self.conducting(Conduction.DC).reachable_from_pin(pin, DCDriver))
if len(drivers) > 1:
raise ValueError(f"{pin.node_id} has {len(drivers)} DC-drivers: {[d.node_id for d in drivers]}")
return drivers[0] if drivers else None
@overload
def reachable_from_pin(self, pin: DevicePin) -> Iterator[object]: ...
@overload
def reachable_from_pin[T](
self,
pin: DevicePin,
object_type: type[T],
where: Callable[[T], bool] | None = None,
) -> Iterator[T]: ...
[docs]
def reachable_from_pin(
self,
pin: DevicePin,
object_type: type[Any] = object,
where: Callable[[Any], bool] | None = None,
) -> Iterator[Any]:
"""Iterate over matching hydrated objects reachable from a device pin.
This returns values stored in graph-node ``obj`` metadata, not NetworkX node IDs or node metadata. The
starting pin and nodes without an ``obj`` value are excluded. Omitting ``object_type`` returns all other
reachable hydrated objects.
:param pin: Device pin from which to start the graph traversal.
:param object_type: Type of object to return. Defaults to :class:`object`.
:param where: Optional predicate applied to objects of ``object_type``.
:returns: Matching objects in breadth-first traversal order.
"""
objects = (self.graph.nodes[node_id].get("obj") for node_id in nx.bfs_tree(self.graph, pin.node_id) if node_id != pin.node_id)
matches = (obj for obj in objects if obj is not None and isinstance(obj, object_type))
return matches if where is None else filter(where, matches)
[docs]
@classmethod
def from_routable(cls, routable: Routable) -> StationGraph:
"""Creates a station graph from a Routable object, usually :class:`NQCPStation`.
Ordering matters: all nodes are created before edges. Edge endpoints may use
different names for the same node, so we need the full node set before we can
resolve those names to canonical node IDs.
For example, a device pin may be referred to as "B083004D010033[RF_IN]" in
one place and "B083004D01/0033[RF_IN]" in another. Before adding the edges
exported by the Routable, we build aliases for all device pins and use them
to resolve edge sources and targets.
"""
graph: nx.Graph = nx.Graph()
for node in routable.graph_nodes():
graph.add_node(
node.id,
kind=node.kind,
obj=node.obj,
**node.attrs,
)
aliases = build_node_aliases(graph)
for edge in routable.graph_edges():
graph.add_edge(
resolve_endpoint_reference(edge.source, aliases),
resolve_endpoint_reference(edge.target, aliases),
kind=edge.kind,
conducts=edge.conducts,
**edge.attrs,
)
return cls(graph=graph, aliases=aliases)
[docs]
def conducting(self, signal: Conduction) -> StationGraph:
"""Returns a subgraph of the station graph containing only the edges that conduct the specified signal type."""
view = nx.subgraph_view(self.graph, filter_edge=lambda u, v: signal in self.graph.edges[u, v]["conducts"])
return StationGraph(graph=view, aliases=self.aliases)