from __future__ import annotations
import bisect
import json
import tempfile
import webbrowser
from dataclasses import dataclass
from typing import TYPE_CHECKING
from urllib.parse import unquote
import networkx as nx
from semi_cr.core.station.drivers import DCDriver
from semi_cr.core.station.graph import StationGraph
if TYPE_CHECKING:
from collections.abc import Mapping
import plotly.graph_objects as go
def _group_of(node_id: str) -> str:
"""Nodes sharing a parent (FEM, DAC module, connector, quantum device) share their node ID minus the last part."""
return node_id.rsplit("/", 1)[0]
def _short(node_id: str) -> str:
"""``lab://terminal/fbrd_1_ffc_J2/49`` -> ``fbrd_1_ffc_J2/49``."""
return unquote(node_id.split("://", 1)[-1].split("/", 1)[-1])
[docs]
@dataclass
class StationGraphPlotter:
"""Plots a :class:`StationGraph` in columns: device pins on the left, sources towards the right.
Usage::
StationGraphPlotter(station.graph).plot()
StationGraphPlotter(station.graph).chip_and_drivers().merge_boards().plot() # debug severed AC/DC paths
"""
station_graph: StationGraph
def _is_driver(self, node_id: str) -> bool:
from semi_cr.core.station.instruments.opx1000 import OPX1000LFTerminal # noqa: PLC0415
return isinstance(self.station_graph.graph.nodes[node_id].get("obj"), (DCDriver, OPX1000LFTerminal))
[docs]
def chip_and_drivers(self) -> StationGraphPlotter:
"""Keep only wiring connected to a device pin or to an AC/DC driver.
A driver whose wiring never reaches the chip ends up in the last column, so a severed path shows as a pin
marked missing AC/DC on the left and a dangling driver on the right.
"""
g = self.station_graph.graph
keep: set[str] = set()
for component in nx.connected_components(g):
if any(g.nodes[n].get("kind") == "device_pin" or self._is_driver(n) for n in component):
keep |= component
return StationGraphPlotter(StationGraph(graph=g.subgraph(keep), aliases=self.station_graph.aliases))
[docs]
def merge_boards(self) -> StationGraphPlotter:
"""Draw each board channel and the connector pins it is traced to as a single row.
A board's own connectors are recognised by name: ``fbrd_1_ffc_J2`` belongs to board ``fbrd_1``. The traces
between a board and its connectors reshuffle line order, which causes crossings no layout can remove, so every
set of nodes joined by board traces is contracted into one node. It keeps the board-side node's ID (and data)
and lists the IDs it absorbed under ``merged``.
"""
g = self.station_graph.graph
board_of: dict[str, str] = {}
names = sorted(self.groups(), key=len)
for board in names:
for other in names:
if other.startswith(board + "_"):
board_of.setdefault(other, board) # shortest prefix wins for nested names
def family(node_id: str) -> str:
return board_of.get(_group_of(node_id), _group_of(node_id))
boards = set(board_of.values())
traces: nx.Graph = nx.Graph()
traces.add_edges_from((u, v) for u, v in g.edges if family(u) == family(v) in boards)
keep: dict[str, str] = {}
for component in nx.connected_components(traces):
board = family(next(iter(component)))
row = min(component, key=lambda n: (_group_of(n) != board, n)) # prefer the board's own node
keep.update(dict.fromkeys(component, row))
merged: nx.Graph = nx.Graph()
merged.add_nodes_from((n, data) for n, data in g.nodes(data=True) if keep.get(n, n) == n)
for node_id, row in sorted(keep.items()):
if node_id != row:
merged.nodes[row].setdefault("merged", []).append(node_id)
for u, v, data in g.edges(data=True):
a, b = keep.get(u, u), keep.get(v, v)
if a != b:
merged.add_edge(a, b, **data)
return StationGraphPlotter(StationGraph(graph=merged, aliases=self.station_graph.aliases))
[docs]
def pin_status(self, node_id: str) -> tuple[bool, bool]:
"""Whether a device pin reaches an AC (OPX source) and a DC driver, by the same lookup the pins use."""
pin = self.station_graph.graph.nodes[node_id]["obj"]
return self.station_graph.opx_source_for_pin(pin) is not None, self.station_graph.dc_driver_for_pin(pin) is not None
def _discovery_rank(self) -> dict[str, int]:
"""Order in which a breadth-first walk out from the (sorted) device pins reaches each node.
Used as a deterministic starting order for the layout; the crossing heuristic depends a lot on where it starts.
"""
g = self.station_graph.graph
order = sorted(n for n, kind in g.nodes(data="kind") if kind == "device_pin")
seen = set(order)
for node_id in order: # grows while we walk it
for neighbour in sorted(g[node_id]):
if neighbour not in seen:
seen.add(neighbour)
order.append(neighbour)
order += sorted(n for n in g if n not in seen)
return {n: i for i, n in enumerate(order)}
[docs]
def groups(self) -> dict[str, list[str]]:
"""Node IDs grouped by parent, e.g. ``lab://terminal/opx1000_1/LF_1`` or ``lab://pin/TestChip/D1``.
Groups and their members are in :meth:`_discovery_rank` order.
"""
rank = self._discovery_rank()
groups: dict[str, list[str]] = {}
for node_id in sorted(self.station_graph.graph, key=rank.__getitem__):
groups.setdefault(_group_of(node_id), []).append(node_id)
return groups
[docs]
def columns(self) -> list[list[str]]:
"""Group names per column, laid out so every wire between groups points from the chip towards the electronics.
Groups are ordered by how close they get to a device pin (in hops), and every wire is directed from the
closer group to the farther one. Each group then goes one column right of the furthest group wired into it
(longest-path layering), so no wire runs vertically or backwards. Wiring not connected to the chip is
ordered by closeness to a driver instead and right-aligned, so a severed path shows as a stub hanging off the
electronics side.
"""
g = self.station_graph.graph
groups = self.groups()
rank = self._discovery_rank()
pins = [n for n, kind in g.nodes(data="kind") if kind == "device_pin"]
drivers = [n for n in g if self._is_driver(n)]
from_pin = nx.multi_source_dijkstra_path_length(g, pins) if pins else {}
to_driver = nx.multi_source_dijkstra_path_length(g, drivers) if drivers else {}
def key(members: list[str]) -> tuple[int, int, int]:
first = rank[members[0]]
if hops := [from_pin[n] for n in members if n in from_pin]:
return 0, min(hops), first
return 1, -min(to_driver.get(n, 0) for n in members), first
keys = {name: key(members) for name, members in groups.items()}
dag: nx.DiGraph = nx.DiGraph()
dag.add_nodes_from(groups)
for u, v in g.edges:
a, b = _group_of(u), _group_of(v)
if a != b:
dag.add_edge(*sorted((a, b), key=keys.__getitem__)) # acyclic: always from lower to higher key
layer: dict[str, int] = {}
for name in nx.topological_sort(dag):
layer[name] = max((layer[p] + 1 for p in dag.predecessors(name)), default=0)
last = max(layer.values(), default=0)
for component in nx.weakly_connected_components(dag):
if all(keys[name][0] == 1 for name in component):
shift = last - max(layer[name] for name in component)
for name in component:
layer[name] += shift
columns: dict[int, list[str]] = {}
for name in sorted(groups, key=keys.__getitem__):
columns.setdefault(layer[name], []).append(name)
return [columns[x] for x in sorted(columns)]
[docs]
def crossings(self, pos: Mapping[str, tuple[float, float]]) -> int:
"""Count edge crossings between adjacent columns for a layout. Edges skipping a column are not counted."""
by_column: dict[float, list[tuple[float, float]]] = {}
for u, v in self.station_graph.graph.edges:
(xu, yu), (xv, yv) = sorted((pos[u], pos[v]))
if xv - xu == 1:
by_column.setdefault(xu, []).append((-yu, -yv))
total = 0
for edges in by_column.values():
# Sorted by left end, so every earlier edge whose right end is lower crosses this one.
seen: list[float] = []
for _, right in sorted(edges):
total += len(seen) - bisect.bisect_right(seen, right)
bisect.insort(seen, right)
return total
[docs]
def layout(self, sweeps: int = 10) -> dict[str, tuple[float, float]]:
"""Compute node positions in :meth:`columns`, with each group drawn as a block of consecutive nodes.
Crossings are reduced with the layered (Sugiyama) barycenter heuristic, applied to groups as blocks: each
column is re-sorted by the mean height of its neighbours in the column before it (left-to-right sweep) and
then after it (right-to-left sweep). Groups move as a whole, and nodes are re-sorted only inside their group.
The layout with the fewest :meth:`crossings` is returned.
:param sweeps: Number of left-to-right plus right-to-left sweep pairs.
:returns: Mapping from node ID to ``(x, y)``.
"""
# ponytail: heuristic, not optimal (minimal crossings is NP-hard); more sweeps rarely help past ~10.
g = self.station_graph.graph
groups = self.groups()
columns = self.columns()
column_of = {n: x for x, names in enumerate(columns) for name in names for n in groups[name]}
pos: dict[str, tuple[float, float]] = {}
def place(x: int) -> None:
y = 0.0
for name in columns[x]:
for node_id in groups[name]:
pos[node_id] = (float(x), y)
y -= 1
y -= 1 # gap between groups
def reorder(x: int, step: int) -> None:
"""Sort column ``x`` by its neighbours on the side we came from (``step`` = +1 left, -1 right)."""
def barycenter(node_id: str) -> float:
ys = [pos[m][1] for m in g[node_id] if (x - column_of[m]) * step > 0]
return sum(ys) / len(ys) if ys else pos[node_id][1] # no neighbours there: stay put
heights = {n: barycenter(n) for name in columns[x] for n in groups[name]}
for name in columns[x]:
groups[name].sort(key=lambda n: -heights[n])
columns[x].sort(key=lambda name: -sum(heights[n] for n in groups[name]) / len(groups[name]))
place(x)
for x in range(len(columns)):
place(x)
best, best_crossings = dict(pos), self.crossings(pos)
for _ in range(sweeps):
for x in range(1, len(columns)):
reorder(x, +1)
for x in range(len(columns) - 2, -1, -1):
reorder(x, -1)
if (crossings := self.crossings(pos)) < best_crossings:
best, best_crossings = dict(pos), crossings
return best
[docs]
def plot(self) -> go.Figure:
"""Draw each parent group as a labelled box, with edges coloured by kind (physical/device).
Hover a node for its full ID and, for device pins, whether it reaches an AC/DC driver. Call ``.show()`` on the
result to open it in a browser (``config={"scrollZoom": True}`` enables mouse-wheel zoom).
"""
import plotly.graph_objects as go # noqa: PLC0415
g = self.station_graph.graph
pos = self.layout()
fig = go.Figure()
# Lines are batched into one trace per style (segments separated by None), which keeps big stations fast.
def lines(segments: list[list[tuple[float, float]]], name: str, color: str) -> None:
xs: list[float | None] = []
ys: list[float | None] = []
for segment in segments:
xs += [x for x, _ in segment] + [None]
ys += [y for _, y in segment] + [None]
fig.add_trace(go.Scattergl(x=xs, y=ys, mode="lines", name=name, line={"color": color, "width": 1}, hoverinfo="skip"))
group_labels: list[tuple[float, float, str]] = []
boxes = []
for name, members in self.groups().items():
x = pos[members[0]][0]
top, bottom = max(pos[n][1] for n in members) + 0.4, min(pos[n][1] for n in members) - 0.4
left, right = x - 0.1, x + 0.75
boxes.append([(left, top), (right, top), (right, bottom), (left, bottom), (left, top)])
group_labels.append((left, top + 0.1, _short(name)))
lines(boxes, "groups", "lightgray")
edge_colors = {"physical": "gray", "device": "royalblue"}
for kind, color in edge_colors.items():
segments = [[pos[u], pos[v]] for u, v, data in g.edges(data=True) if data.get("kind") == kind]
lines(segments, kind, color)
# Pins: green = AC+DC, orange = one of them, red = neither. Drivers: purple. Everything else: black.
node_ids = list(pos)
colors, labels, hovers = [], [], []
for node_id in node_ids:
label = unquote(node_id.rsplit("/", 1)[-1])
hover = unquote(node_id)
if merged := g.nodes[node_id].get("merged"):
board = _short(_group_of(node_id))
label += f" ({', '.join(_short(m).removeprefix(board + '_') for m in merged)})"
hover += "".join(f"<br>+ {unquote(m)}" for m in merged)
if g.nodes[node_id].get("kind") == "device_pin":
ac, dc = self.pin_status(node_id)
colors.append(("red", "orange", "green")[ac + dc])
status = f"[{'AC' if ac else '--'}|{'DC' if dc else '--'}]"
label += f" {status}"
hover += f"<br>AC driver: {'yes' if ac else 'NO'}<br>DC driver: {'yes' if dc else 'NO'}"
else:
colors.append("purple" if self._is_driver(node_id) else "black")
labels.append(label)
hovers.append(hover)
fig.add_trace(
go.Scattergl(
x=[pos[n][0] for n in node_ids],
y=[pos[n][1] for n in node_ids],
mode="markers+text",
name="nodes",
marker={"color": colors, "size": 7},
text=labels,
textposition="middle right",
textfont={"size": 10},
hovertext=hovers,
hoverinfo="text",
customdata=node_ids,
)
)
fig.add_trace(go.Scattergl(x=[], y=[], mode="lines", name="highlight", line={"color": "gold", "width": 4}, hoverinfo="skip"))
fig.add_trace(
go.Scattergl(
x=[x for x, _, _ in group_labels],
y=[y for _, y, _ in group_labels],
mode="text",
name="group names",
text=[t for *_, t in group_labels],
textfont={"weight": "bold"}, # WebGL text ignores HTML tags like <b>
textposition="top right",
hoverinfo="skip",
)
)
fig.update_layout(
plot_bgcolor="white",
dragmode="pan",
hovermode="closest",
xaxis={"visible": False},
yaxis={"visible": False},
margin={"l": 10, "r": 10, "t": 10, "b": 10},
)
return fig
[docs]
def show(self, path: str | None = None) -> str:
"""Write :meth:`plot` to an HTML file and open it in the browser.
Click a node to highlight all wiring connected to it (e.g. a pin's path to its drivers) and fade the rest;
double-click to reset. Scroll to zoom, drag to pan.
:param path: Where to write the HTML. Defaults to a temporary file.
:returns: The path written.
"""
fig = self.plot()
g = self.station_graph.graph
traces = {trace.name: i for i, trace in enumerate(fig.data)}
nodes = fig.data[traces["nodes"]]
index = {n: i for i, n in enumerate(nodes.customdata)}
data = {
"x": list(nodes.x),
"y": list(nodes.y),
"adjacency": [[index[m] for m in g[n]] for n in nodes.customdata],
"nodes": traces["nodes"],
"highlight": traces["highlight"],
"dim": [i for name, i in traces.items() if name not in {"nodes", "highlight", "group names"}],
}
script = f"const gd = document.getElementById('{{plot_id}}');\nconst d = {json.dumps(data)};\n" + _HIGHLIGHT_JS
if path is None:
with tempfile.NamedTemporaryFile(suffix=".html", prefix="station_graph_", delete=False) as f:
path = f.name
fig.write_html(path, config={"scrollZoom": True}, post_script=script)
webbrowser.open(f"file://{path}")
return path
# Breadth-first walk from the clicked node; the reached wiring is drawn into the highlight trace.
_HIGHLIGHT_JS = """
gd.on('plotly_click', (event) => {
const point = event.points.find((p) => p.curveNumber === d.nodes);
if (!point) return;
const seen = new Set([point.pointIndex]);
const queue = [point.pointIndex];
const hx = [], hy = [];
while (queue.length) {
const n = queue.shift();
for (const m of d.adjacency[n]) {
hx.push(d.x[n], d.x[m], null);
hy.push(d.y[n], d.y[m], null);
if (!seen.has(m)) { seen.add(m); queue.push(m); }
}
}
Plotly.restyle(gd, {x: [hx], y: [hy]}, [d.highlight]);
Plotly.restyle(gd, {'marker.opacity': [d.x.map((_, i) => (seen.has(i) ? 1 : 0.15))]}, [d.nodes]);
Plotly.restyle(gd, {opacity: 0.2}, d.dim);
});
gd.on('plotly_doubleclick', () => {
Plotly.restyle(gd, {x: [[]], y: [[]]}, [d.highlight]);
Plotly.restyle(gd, {'marker.opacity': 1}, [d.nodes]);
Plotly.restyle(gd, {opacity: 1}, d.dim);
});
"""