from dataclasses import dataclass
from typing import Any
from collections.abc import Callable, Sequence
import networkx as nx
from qcodes.instrument import Instrument, InstrumentModule, InstrumentChannel
from qcodes.station import Station
from semi_cr.core.lab.station.graphing.models import Node, Edge
from semi_cr.core.lab.station.graphing.ids import (
instrument_node_id,
instrument_module_node_id,
instrument_connection_node_id,
terminal_node_id,
)
from semi_cr.core.lab.station.graphing.parsing import (
parse_endpoint_ref,
parse_channel_number,
)
from semi_cr.core.lab.station.graphing.operations import build_nx_graph
[docs]
def build_connection_graph_for_instrument(
instrument: Instrument,
connections: Sequence[dict[str, Any]],
parent_name: str | None = None,
) -> nx.MultiDiGraph:
nodes: list[Node] = []
edges: list[Edge] = []
owner_name = getattr(instrument, "short_name", instrument.name)
if parent_name is None:
root_id = instrument_node_id(owner_name)
root_kind = "instrument"
else:
root_id = instrument_module_node_id(parent_name,owner_name)
root_kind = "instrument_module"
nodes.append(
Node(
id=root_id,
kind=root_kind,
name=owner_name,
obj=instrument,
)
)
for connection in connections:
connection_name = str(connection["name"])
connection_id = instrument_connection_node_id(
owner_name,
connection_name,
)
terminal_id = terminal_node_id(
owner_name,
connection_name,
)
nodes.extend(
[
Node(
id=connection_id,
kind=f"{root_kind}_connection",
name=connection_name,
obj=None,
),
Node(
id=terminal_id,
kind="terminal",
name=f"{owner_name}[{connection_name}]",
obj=None,
),
]
)
# Root module/instrument contains its named connection.
edges.append(
Edge(
source=root_id,
target=connection_id,
kind="contains",
)
)
# The named connection represents the electrical terminal.
edges.append(
Edge(
source=connection_id,
target=terminal_id,
kind="represents",
)
)
# The root object also represents the terminal, so route lookups
# can find module-level attributes like gain_v_per_a.
edges.append(
Edge(
source=root_id,
target=terminal_id,
kind="represents",
)
)
for endpoint in connection.get("endpoints", []):
endpoint_owner, endpoint_name = parse_endpoint_ref(endpoint)
endpoint_terminal_id = terminal_node_id(
endpoint_owner,
endpoint_name,
)
nodes.append(
Node(
id=endpoint_terminal_id,
kind="terminal",
name=f"{endpoint_owner}[{endpoint_name}]",
obj=None,
)
)
edges.append(
Edge(
source=endpoint_terminal_id,
target=terminal_id,
kind="dependency",
)
)
return build_nx_graph(nodes, edges)
[docs]
def build_connection_graph_for_instrument_module(
module: InstrumentModule,
connections: Sequence[dict[str, Any]],
parent_name: str,
owner_name: str | None = None,
) -> nx.MultiDiGraph:
nodes: list[Node] = []
edges: list[Edge] = []
owner_name = owner_name or getattr(module, "short_name", module.name)
module_id = instrument_module_node_id(
parent_name,
owner_name,
)
nodes.append(
Node(
id=module_id,
kind="instrument_module",
name=owner_name,
obj=module,
attrs={
"module_type": getattr(module, "module_type", None),
"slot": getattr(module, "slot", None),
},
)
)
for connection in connections:
connection_name = str(connection["name"])
role = connection.get("role")
# connection_id = instrument_connection_node_id(
# owner_name,
# connection_name,
# )
terminal_id = terminal_node_id(
owner_name,
connection_name,
)
common_attrs = {
"role": role,
"modality": "rf",
"controller": getattr(module, "controller", None),
"module_type": getattr(module, "module_type", None),
"slot": getattr(module, "slot", None),
"number": parse_channel_number(connection_name),
}
# nodes.append(
# Node(
# id=connection_id,
# kind="instrument_connection",
# name=connection_name,
# attrs=common_attrs,
# )
# )
nodes.append(
Node(
id=terminal_id,
kind="terminal",
name=f"{owner_name}[{connection_name}]",
attrs=common_attrs,
)
)
# edges.append(
# Edge(
# source=module_id,
# target=connection_id,
# kind="contains",
# )
# )
# edges.append(
# Edge(
# source=connection_id,
# target=terminal_id,
# kind="represents",
# )
# )
edges.append(
Edge(
source=module_id,
target=terminal_id,
kind="contains",
)
)
for endpoint in connection.get("endpoints", []):
endpoint_owner, endpoint_name = parse_endpoint_ref(endpoint)
endpoint_terminal_id = terminal_node_id(
endpoint_owner,
endpoint_name,
)
nodes.append(
Node(
id=endpoint_terminal_id,
kind="terminal",
name=f"{endpoint_owner}[{endpoint_name}]",
)
)
edges.append(
Edge(
source=endpoint_terminal_id,
target=terminal_id,
kind="dependency",
)
)
return build_nx_graph(nodes, edges)
[docs]
@dataclass(frozen=True)
class ContainmentGraphBuilder:
"""
Builds a containment graph from a QCoDeS Station.
Stores build policy + ID strategy.
"""
# station_id: str = "qcodes://station/noma"
attach_objects: bool = True
include_instrument_parameters: bool = True
include_submodules: bool = True
include_submodule_parameters: bool = True
# dependency-injected ID resolver
node_id: Callable[[Any], str] = lambda obj: f"lab://obj/{type(obj).__name__}/{getattr(obj, 'name', str(id(obj)))}"
[docs]
def station_node_id(self, station: Station) -> str:
station_name = getattr(station, "name", None) or "station"
return f"qcodes://station/{station_name}"
[docs]
def build(self, station: Station) -> nx.MultiDiGraph:
G = nx.MultiDiGraph()
station_id = self.station_node_id(station)
G.add_node(
station_id,
kind="station",
name=getattr(station, "name", "station"),
obj=station if self.attach_objects else None,
)
for _, comp in station.components.items():
if isinstance(comp, Instrument):
inst_id = self.node_id(comp)
G.add_node(inst_id, kind="instrument", name=comp.name, obj=comp if self.attach_objects else None)
G.add_edge(station_id, inst_id, kind="contains")
if self.include_instrument_parameters:
for p_name, p in getattr(comp, "parameters", {}).items():
pid = self.node_id(p)
G.add_node(pid, kind="parameter", name=p_name, obj=p if self.attach_objects else None)
G.add_edge(inst_id, pid, kind="contains")
if self.include_submodules:
for sm_name, sm in getattr(comp, "submodules", {}).items():
if not isinstance(sm, InstrumentChannel):
continue
sm_id = self.node_id(sm)
G.add_node(sm_id, kind="submodule", name=sm_name, obj=sm if self.attach_objects else None)
G.add_edge(inst_id, sm_id, kind="contains")
if self.include_submodule_parameters:
for p_name, p in getattr(sm, "parameters", {}).items():
pid = self.node_id(p)
G.add_node(pid, kind="parameter", name=p_name, obj=p if self.attach_objects else None)
G.add_edge(sm_id, pid, kind="contains")
else:
cid = self.node_id(comp)
G.add_node(cid, kind="component", name=getattr(comp, "name", str(comp)), obj=comp if self.attach_objects else None)
G.add_edge(station_id, cid, kind="contains")
return G