"""Chip topology (pyvis) and dressed-spectrum (matplotlib) plots."""
from __future__ import annotations
from pathlib import Path
from typing import Any
import matplotlib.pyplot as plt
from matplotlib.figure import Figure
from quchip.chip.chip import Chip
from quchip.utils.jax_utils import maybe_concrete_scalar
from quchip.viz._common import _basis_label, _draw_energy_ladder
from quchip.viz._style import _GRAPH_PALETTE, _cyclic_colors, _quchip_style, _resolve_single_axes
_NODE_SHAPES = {"device": "dot", "drive": "diamond", "coupling": "square"}
_NODE_SIZES = {"device": 25, "drive": 15, "coupling": 12}
_EDGE_COLORS = {"coupling": "#4d4d4d", "wiring": "#ff7f0e", "crosstalk": "#b22222"}
_EDGE_LENGTHS = {"coupling": 150, "wiring": 80, "crosstalk": 300}
_VALUE_MODES = {"bare", "dressed", "both"}
def _device_node_id(label: str) -> str:
"""Namespace a device label as a topology node id (``"device:<label>"``)."""
return f"device:{label}"
def _coupling_node_id(label: str) -> str:
"""Namespace a coupling label as a topology node id (``"coupling:<label>"``)."""
return f"coupling:{label}"
def _drive_node_id(label: str) -> str:
"""Namespace a drive label as a topology node id (``"drive:<label>"``)."""
return f"drive:{label}"
def _frequency_text(value: Any, *, precision: int = 3) -> str:
"""Format a concrete frequency in GHz, or ``"n/a"`` when unavailable."""
concrete = maybe_concrete_scalar(value)
return "n/a" if concrete is None else f"{concrete:.{precision}f} GHz"
def _device_label(device: Any, *, dressed_frequency: Any = None, values: str = "bare") -> str:
"""Return a device label annotated with bare, dressed, or both frequencies."""
bare_text = _frequency_text(getattr(device, "freq", None))
dressed_text = _frequency_text(dressed_frequency, precision=6)
if values == "bare":
return f"{device.label}\n{bare_text}"
if values == "dressed":
return f"{device.label}\nf01={dressed_text}"
return f"{device.label}\nbare={bare_text}\nf01={dressed_text}"
def _coupling_label(coupling: Any, *, dressed_kerr: Any = None, values: str = "bare") -> tuple[str, str]:
"""Return junction and edge labels for bare strength and/or dressed cross-Kerr."""
strength = maybe_concrete_scalar(coupling.coupling_strength)
strength_text = "n/a" if strength is None else f"{strength:.3g} GHz"
kerr = maybe_concrete_scalar(dressed_kerr)
kerr_text = "n/a" if kerr is None else f"{kerr:.3g} GHz"
class_name = type(coupling).__name__
if values == "bare":
physics_lines = [f"g={strength_text}"]
elif values == "dressed":
physics_lines = [f"K={kerr_text}"]
else:
physics_lines = [f"g={strength_text}", f"K={kerr_text}"]
edge_label = "\n".join((class_name, *physics_lines))
return f"{coupling.label}\n{edge_label}", edge_label
def _control_label(drive: Any) -> str:
"""Return the multi-line node label for *drive* (edge pumps always name their pumped coupling)."""
from quchip.control.drive import CouplingDrive
class_name = type(drive).__name__
if isinstance(drive, CouplingDrive):
return f"{drive.label}\n({class_name})\n→ {drive.target_label}"
if drive.label:
return f"{drive.label}\n({class_name})"
return f"{class_name}\n→ {drive.device_label}"
def _collect_topology(
chip: Chip,
*,
exclude: set[str] | None = None,
values: str = "bare",
) -> tuple[dict[str, dict[str, Any]], list[tuple[str, str, dict[str, Any]]]]:
"""Collect node/edge dicts for devices, couplings, drives, and crosstalks.
Every coupling is its own junction node splitting the device-device
edge in two (``device -- junction -- device``), so an edge-pump
control (:class:`~quchip.control.drive.ParametricDrive`) has a node
to attach to — it modulates the coupling element itself, not either
endpoint device. Node ids are namespaced by kind
(``"device:q0"``, ``"coupling:tc"``, ``"drive:pump"``) so a device,
coupling, and drive that happen to share a user-chosen label can
never collide or silently overwrite one another in *nodes*.
When couplings are excluded from the render (``"coupling" in
exclude``), any edge-pump drives that would otherwise attach to a
now-absent junction node are omitted too, rather than left dangling.
Returns ``(nodes, edges)`` where *nodes* maps namespaced node id to
attributes (insertion-ordered: devices, then couplings, then
drives) and *edges* is a list of ``(start, end, attributes)``
triples of namespaced node ids.
"""
if values not in _VALUE_MODES:
raise ValueError(f"values must be one of {sorted(_VALUE_MODES)}, got {values!r}")
_exclude = exclude or set()
nodes: dict[str, dict[str, Any]] = {}
edges: list[tuple[str, str, dict[str, Any]]] = []
coupling_labels: set[str] = set()
dressed_frequencies: dict[str, Any] = {}
dressed_kerr: Any = None
if values != "bare":
dressed_frequencies = chip.freq()
dressed_kerr = chip.kerr_matrix()
for device in chip.devices:
nodes[_device_node_id(device.label)] = {
"kind": "device",
"label": _device_label(
device,
dressed_frequency=dressed_frequencies.get(device.label),
values=values,
),
"class_name": type(device).__name__,
"computational": device.computational,
}
if "coupling" not in _exclude:
for coupling in chip.couplings:
coupling_labels.add(coupling.label)
junction_id = _coupling_node_id(coupling.label)
junction_label, edge_label = _coupling_label(
coupling,
dressed_kerr=(
None
if dressed_kerr is None
else dressed_kerr[coupling.device_a_label, coupling.device_b_label]
),
values=values,
)
nodes[junction_id] = {
"kind": "coupling",
"label": junction_label,
"class_name": type(coupling).__name__,
"computational": False,
}
edges.append((
_device_node_id(coupling.device_a_label), junction_id, {"kind": "coupling", "label": edge_label},
))
edges.append((
junction_id, _device_node_id(coupling.device_b_label), {"kind": "coupling", "label": edge_label},
))
if "drive" not in _exclude and chip.control_equipment is not None:
from quchip.control.drive import CouplingDrive
for drive in chip.control_equipment.lines:
is_edge_pump = isinstance(drive, CouplingDrive)
if is_edge_pump and drive.target_label not in coupling_labels:
# No junction node to attach to — either couplings were
# excluded from this render, or the pump's coupling isn't
# on this chip. Either way, omit the dangling control.
continue
drive_id = _drive_node_id(drive.label)
nodes[drive_id] = {
"kind": "drive",
"label": _control_label(drive),
"class_name": type(drive).__name__,
"computational": False,
}
target_id = (
# target_label is non-None here: edge-pump implies a
# ParametricDrive, whose target_label override returns str,
# never None.
_coupling_node_id(drive.target_label) # type: ignore[arg-type]
if is_edge_pump
# device_label is non-None here: Chip.connect() rejects unconnected
# drives before they reach control_equipment.lines.
else _device_node_id(drive.device_label) # type: ignore[arg-type]
)
edges.append((
drive_id,
target_id,
{"kind": "wiring", "label": f"{type(drive).__name__} → {drive.target_label}"},
))
if "crosstalk" not in _exclude:
for xt in chip.crosstalks:
# Only connect channels the drive pass actually rendered: a
# crosstalk edge must never resurrect a drive that was
# deliberately omitted (an edge pump whose coupling is
# excluded from this render).
source_id, victim_id = _drive_node_id(xt.source), _drive_node_id(xt.victim)
if source_id not in nodes or victim_id not in nodes:
continue
beta = maybe_concrete_scalar(abs(xt.beta))
beta_text = "n/a" if beta is None else f"{beta:.3g}"
edges.append((
source_id, victim_id,
{"kind": "crosstalk", "label": f"|xi|={beta_text}"},
))
return nodes, edges
[docs]
def plot_graph(
chip: Chip,
path: str = "chip_topology.html",
*,
full: bool = True,
exclude: set[str] | None = None,
values: str = "bare",
layout: str = "force_atlas",
height: str = "800px",
width: str = "100%",
) -> str:
"""Write an interactive HTML chip topology and return its path.
The HTML file is standalone and works offline. Devices show as dots and
drives as diamonds. Edge colour and dash pattern identify couplings,
control wiring, and crosstalk. Node colours follow the *class*
(``DuffingTransmon``, ``Resonator``, ``ChargeDrive``, ...) automatically,
so new kinds of device and drive look different without user action.
Layouts:
- ``"force_atlas"`` (default): ``force_atlas_2based`` physics with larger
and heavier computational devices, so the register sits at the center and
the auxiliary couplers and drives orbit it.
- ``"hierarchical"``: pyvis ``hrepulsion`` on explicit levels
(computational devices → other devices/couplings → drives). It is useful
for chips with many auxiliary drives or cross-resonance lines.
*exclude* has priority over *full*. If you pass ``full=False`` and no
*exclude*, the plot hides drives and crosstalk (chip-only view). Each
coupling renders as its own junction node that splits the device-device
edge, so an edge-pump control has a node to attach to. Excluding couplings
also removes all edge-pump controls that would otherwise dangle.
``values="bare"`` keeps the declaration view: device frequencies and
coupling strengths come directly from component parameters, without
diagonalization. ``values="dressed"`` labels devices with dressed
:math:`f_{01}` and coupling junctions with the full-pull cross-Kerr
:math:`K_{ab}`. ``values="both"`` shows both annotations. These modes do
not change the topology.
Parameters
----------
chip : Chip
The chip to visualise.
path : str
Destination ``.html`` path, also returned.
full : bool
If ``False``, hide drives and crosstalk unless *exclude* is given.
exclude : set of str, optional
The items of ``{"coupling", "drive", "crosstalk"}`` to hide fully.
values : {"bare", "dressed", "both"}
Physical values that show on device and coupling nodes. The default is
the declared bare values.
layout : str
``"force_atlas"`` or ``"hierarchical"``.
height, width : str
CSS dimensions passed to ``pyvis.network.Network``.
Returns
-------
str
*path*, unchanged, so you can chain the call into code that opens or
serves the file.
Raises
------
ImportError
``pyvis`` is not installed.
ValueError
*layout* or *values* is not one of its supported choices.
"""
# The removal of dangling edge-pump controls is implemented in
# `_collect_topology`.
try:
from pyvis.network import Network
except ImportError as exc:
raise ImportError("pyvis is required for graph visualization: pip install 'quchip[viz]'") from exc
valid_layouts = {"force_atlas", "hierarchical"}
if layout not in valid_layouts:
raise ValueError(f"layout must be one of {sorted(valid_layouts)}, got {layout!r}")
if exclude is None and not full:
exclude = {"drive", "crosstalk"}
nodes, edges = _collect_topology(chip, exclude=exclude, values=values)
class_names = list(dict.fromkeys(data["class_name"] for data in nodes.values()))
class_colors = _cyclic_colors(class_names, _GRAPH_PALETTE)
net = Network(height=height, width=width, directed=False, notebook=False, cdn_resources="in_line")
if layout == "force_atlas":
net.force_atlas_2based(
gravity=-50, central_gravity=0.005, spring_length=150,
spring_strength=0.04, damping=0.5, overlap=1,
)
else: # layout == "hierarchical" (the only other option; validated above)
net.hrepulsion(
node_distance=200, central_gravity=0.1, spring_length=200,
spring_strength=0.05, damping=0.15,
)
net.neighborhood_highlight = True
net.select_menu = False
net.set_edge_smooth("continuous")
net.toggle_hide_edges_on_drag(True)
net.toggle_stabilization(True)
net.set_options('{"physics": {"stabilization": {"iterations": 200}}}')
use_levels = layout == "hierarchical"
for node, data in nodes.items():
kind = data["kind"]
is_computational = data.get("computational", False)
node_kwargs: dict[str, Any] = {
"label": data["label"],
"title": data["label"],
"color": class_colors.get(data["class_name"], "#999999"),
"shape": _NODE_SHAPES.get(kind, "dot"),
}
if use_levels:
node_kwargs["size"] = _NODE_SIZES.get(kind, 20)
if kind == "drive":
node_kwargs["level"] = 2
elif is_computational:
node_kwargs["level"] = 0
else:
node_kwargs["level"] = 1
elif is_computational:
node_kwargs["size"], node_kwargs["mass"] = 30, 3
elif kind == "device":
node_kwargs["size"], node_kwargs["mass"] = 20, 2
elif kind == "coupling":
node_kwargs["size"], node_kwargs["mass"] = 10, 1
else:
node_kwargs["size"], node_kwargs["mass"] = 12, 1
net.add_node(node, **node_kwargs)
for start, end, data in edges:
kind = data["kind"]
edge_kwargs: dict[str, Any] = {
"title": data["label"],
"color": _EDGE_COLORS.get(kind, "#4d4d4d"),
"dashes": kind == "crosstalk",
"width": 2.5 if kind == "coupling" else 1.5,
}
if not use_levels:
edge_kwargs["length"] = _EDGE_LENGTHS.get(kind, 150)
net.add_edge(start, end, **edge_kwargs)
net.write_html(path)
# Center tooltip text on multi-line labels.
out = Path(path)
out.write_text(out.read_text().replace(
"<head>", "<head>\n<style>.vis-tooltip { text-align: center; }</style>", 1,
))
return path
[docs]
def plot_energy_levels(
chip: Chip,
*,
ax: Any = None,
max_states: int | None = None,
) -> Figure:
"""Plot dressed chip eigenenergies relative to the ground state.
Each level shows as a horizontal bar with its dressed represented-basis
tuple label ``|n_1 n_2 ...>``. This label is the per-device bare-basis
assignment that the chip's dressing analysis selects (see
``chip.analysis``). The y-axis is energy in GHz, with the ground state
shifted to zero.
Parameters
----------
chip : Chip
Chip whose dressed spectrum is plotted.
ax : matplotlib.axes.Axes, optional
Existing axes to draw on. When ``None``, the function makes a new figure.
max_states : int, optional
Number of lowest levels to show. The default is
``min(12, total_levels)``.
Returns
-------
Figure
The figure that holds the energy-ladder axes (``ax.figure`` when *ax*
was given).
Examples
--------
>>> import quchip as qc
>>> from quchip.viz.chip import plot_energy_levels
>>> qb0 = qc.DuffingTransmon(freq=5.0, anharmonicity=-0.3, levels=3)
>>> qb1 = qc.DuffingTransmon(freq=5.2, anharmonicity=-0.3, levels=3)
>>> chip = qc.Chip([qb0, qb1], [qc.Capacitive(qb0, qb1, g=0.005)])
>>> fig = plot_energy_levels(chip, max_states=6)
"""
dressed = chip._ensure_dressed()
levels = sorted(dressed.dressed_eigenvalues.items(), key=lambda item: item[1])
if max_states is None:
max_states = min(12, len(levels))
displayed = levels[:max_states]
ground_energy = 0.0 if not displayed else float(displayed[0][1])
line_color = plt.get_cmap("tab10")(0)
entries = [(float(energy) - ground_energy, _basis_label(state)) for state, energy in displayed]
with _quchip_style():
fig, axis = _resolve_single_axes(ax)
_draw_energy_ladder(axis, entries, color=line_color, linewidth=2.0)
axis.set_ylabel("Energy relative to ground (GHz)")
axis.set_title("Dressed energy levels", fontfamily="sans-serif")
fig.tight_layout()
return fig
