"""Control-equipment container: drive lines and signal chain.
This container, not the individual drives, owns the signal transforms.
"""
from __future__ import annotations
from typing import Any
from quchip.control.drive import BaseDrive, CouplingDrive
from quchip.declarative.parameters import parameter, setting
from quchip.control.signal import Crosstalk, SignalMap, SignalTransform
from quchip.utils.jax_utils import (
is_jax_array as _is_traced,
select_array_module as _select_array_module,
)
[docs]
class CrosstalkMatrix(SignalTransform, serializable=True):
"""Dense crosstalk transform and matrix view in control-line order.
Attributes
----------
labels : tuple[str, ...]
Drive labels in wiring order (the order in which drives appear in
:attr:`ControlEquipment.lines`). Row / column ``i`` corresponds to
``labels[i]``.
beta : Any
``[n, n]`` amplitude matrix. ``beta[i, j]`` is the leakage amplitude
from source ``labels[j]`` onto victim ``labels[i]`` (column = source,
row = victim). Diagonal entries are self-coupling, ``1.0`` by
convention.
theta : Any
``[n, n]`` phase matrix (radians), same indexing as ``beta``.
delay : Any
``[n, n]`` delay matrix (ns), same indexing as ``beta``.
Notes
-----
Every off-diagonal edge reads the same input signal map, so reciprocal
entries form one linear mixing stage and do not recursively leak each
other's output. Matrix entries keep end-to-end JAX traceability.
"""
# Matrix entries go directly into the signal-program IR
# (`PolarScale`/`Shift`), which is what keeps end-to-end JAX traceability.
labels: tuple[str, ...] = setting()
beta: Any = parameter()
theta: Any = parameter()
delay: Any = parameter()
def __init__(self, labels: tuple[str, ...], beta: Any, theta: Any, delay: Any) -> None:
super().__init__(labels=tuple(labels), beta=beta, theta=theta, delay=delay)
[docs]
def validate(self) -> None:
"""Check that all crosstalk matrices match the declared line order."""
shape = (len(self.labels), len(self.labels))
for name in ("beta", "theta", "delay"):
matrix = getattr(self, name)
if getattr(matrix, "shape", None) != shape:
raise ValueError(f"{name} matrix shape {getattr(matrix, 'shape', None)} does not match {shape}")
[docs]
def apply(self, signals: SignalMap) -> SignalMap:
"""Apply all directed leakage edges to one shared input snapshot.
Parameters
----------
signals : SignalMap
Signals keyed by ``(line, source_index)``.
"""
output = dict(signals)
line_index = {label: index for index, label in enumerate(self.labels)}
for key, signal in signals.items():
source_index = line_index.get(key[0])
if source_index is None:
continue
for victim_index, victim in enumerate(self.labels):
if victim_index == source_index:
continue
leaked = signal.shifted(
self.delay[victim_index, source_index]
).polar_scaled(
self.beta[victim_index, source_index],
self.theta[victim_index, source_index],
)
victim_key = (victim, key[1])
existing = output.get(victim_key)
output[victim_key] = leaked if existing is None else existing + leaked
return output
[docs]
def referenced_lines(self) -> tuple[str, ...]:
"""Return drive labels referenced by this transform."""
return self.labels
[docs]
def without_line(self, line: str) -> "CrosstalkMatrix | None":
"""Return a copy with ``line`` removed, or ``None`` below two lines.
Parameters
----------
line : str
Drive label to remove.
"""
if line not in self.labels:
return self
keep = [index for index, label in enumerate(self.labels) if label != line]
if len(keep) < 2:
return None
labels = tuple(self.labels[index] for index in keep)
return type(self)(
labels=labels,
beta=self.beta[keep][:, keep],
theta=self.theta[keep][:, keep],
delay=self.delay[keep][:, keep],
)
[docs]
def edges(self) -> list[Crosstalk]:
"""Return the directed-edge view used by topology and visualization."""
return [
Crosstalk(
source=self.labels[source],
victim=self.labels[victim],
beta=self.beta[victim, source],
theta=self.theta[victim, source],
delay=self.delay[victim, source],
)
for victim in range(len(self.labels))
for source in range(len(self.labels))
if victim != source
]
def _leaves(value: Any):
if isinstance(value, (list, tuple)):
for item in value:
yield from _leaves(item)
else:
yield value
def _as_matrix(matrix: Any) -> Any:
"""Return *matrix* as an array; nested sequences are stacked, tracers preserved."""
if matrix is None or hasattr(matrix, "shape"):
return matrix
traced = any(_is_traced(leaf) for leaf in _leaves(matrix) if hasattr(leaf, "shape"))
return _select_array_module(traced).asarray(matrix, dtype=float)
[docs]
class ControlEquipment:
"""Ordered drive lines plus a sequence of signal-chain transforms.
The equipment pipes complete analytic signals through ``signal_chain``
before destination drives author Hamiltonian terms.
Parameters
----------
lines : list of BaseDrive
Drive lines in wiring order.
signal_chain : list of SignalTransform, optional
Sequential transforms applied after scheduling.
"""
def __init__(
self,
lines: list[BaseDrive],
*,
signal_chain: list[SignalTransform] | None = None,
) -> None:
invalid = [line for line in lines if not isinstance(line, BaseDrive)]
if invalid:
raise TypeError(
"ControlEquipment lines must be BaseDrive instances; coherent field "
"inputs are scheduled directly and propagate through PortNetwork."
)
self._lines = list(lines)
self._signal_chain = list(signal_chain) if signal_chain else []
@property
def lines(self) -> list[BaseDrive]:
"""Ordered drive lines (defensive copy)."""
return list(self._lines)
@property
def signal_chain(self) -> list[SignalTransform]:
"""Signal-chain transforms (defensive copy)."""
return list(self._signal_chain)
[docs]
def apply_signal_chain(self, signals: SignalMap) -> SignalMap:
"""Apply every signal-chain transform to *signals*, in order.
Each transform receives the output of the previous one, so transforms
compose sequentially and the order of :attr:`signal_chain` affects the
result. For example, a :class:`Delay` applied before a :class:`Gain`
sees the undelayed signal.
Parameters
----------
signals : SignalMap
``{(line_label, source_index): AnalyticSignal}`` map.
``source_index`` distinguishes scheduled pulses through mixing.
Returns
-------
SignalMap
Transformed signal map, which can contain keys absent from
*signals*. For example, a :class:`Crosstalk` transform adds an
entry under the victim drive's label for every source entry it
leaks from.
"""
built = dict(signals)
for transform in self._signal_chain:
built = transform.apply(built)
return built
@property
def crosstalks(self) -> list[Crosstalk]:
"""Directed crosstalk edges represented by the signal chain."""
edges: list[Crosstalk] = []
for transform in self._signal_chain:
if isinstance(transform, Crosstalk):
edges.append(transform)
elif isinstance(transform, CrosstalkMatrix):
edges.extend(transform.edges())
return edges
[docs]
def crosstalk_matrix(self) -> CrosstalkMatrix:
"""Return a dense matrix view of the :class:`Crosstalk` transforms.
The matrix uses wiring order (``self.lines``) as the stable axis
ordering. Column index = source drive, row index = victim drive.
Diagonal entries are ``beta=1``, ``theta=0``, ``delay=0`` by convention
(self-coupling). Off-diagonal entries aggregate every
:class:`Crosstalk` transform in the signal chain. Lines with no
corresponding transform contribute zeros.
Non-:class:`Crosstalk` transforms (``Gain``, ``Delay``) are ignored
here, because this is strictly a view of the crosstalk edges.
Returns
-------
CrosstalkMatrix
``labels`` (wiring order), ``beta``, ``theta``, ``delay``
as ``[n, n]`` arrays. Arrays use ``jax.numpy`` when any
stored entry is a JAX tracer or array, otherwise
``numpy``.
"""
labels = tuple(line.label for line in self._lines)
index = {label: i for i, label in enumerate(labels)}
n = len(labels)
entries: dict[tuple[int, int], tuple[Any, Any, Any]] = {}
any_traced = False
for t in self.crosstalks:
if t.source not in index or t.victim not in index:
continue
i = index[t.victim]
j = index[t.source]
entries[i, j] = (t.beta, t.theta, t.delay)
for val in (t.beta, t.theta, t.delay):
if _is_traced(val):
any_traced = True
xp = _select_array_module(any_traced)
grid = [[entries.get((i, j), (float(i == j), 0.0, 0.0)) for j in range(n)] for i in range(n)]
beta, theta, delay = (
xp.asarray([[cell[field] for cell in row] for row in grid], dtype=float).reshape(n, n)
for field in range(3)
)
return CrosstalkMatrix(labels=labels, beta=beta, theta=theta, delay=delay)
[docs]
def set_crosstalk_matrix(
self,
beta: Any,
theta: Any | None = None,
delay: Any | None = None,
*,
labels: tuple[str, ...] | list[str] | None = None,
) -> None:
"""Rehydrate the crosstalk edges from dense matrices.
Removes every crosstalk transform currently in the signal chain and
replaces them with one :class:`CrosstalkMatrix`. Other signal-chain
transforms (``Gain``, ``Delay``, and user-defined subclasses) stay in
their order.
Parameters
----------
beta : Any
``[n, n]`` amplitude matrix. ``beta[i, j]`` is emitted as
the leakage amplitude from source ``labels[j]`` onto victim
``labels[i]``. Diagonal entries are ignored (self-coupling
belongs to the drive itself, not to a crosstalk edge).
theta : Any, optional
``[n, n]`` phase matrix (radians). Defaults to zeros.
delay : Any, optional
``[n, n]`` delay matrix (ns). Defaults to zeros.
labels : tuple[str, ...] | list[str] | None, optional
Axis ordering. Defaults to wiring order
(``self.lines``). Must match ``beta.shape[0]``.
Notes
-----
Traced JAX entries flow unchanged into :class:`CrosstalkMatrix` and
therefore into the signal-program IR without concretization.
"""
order = tuple(line.label for line in self._lines) if labels is None else tuple(labels)
n = len(order)
# Plain nested lists/tuples are accepted; arrays and tracers already carry
# a shape and pass through untouched. A list holding traced leaves is
# stacked with jax.numpy so no tracer is ever concretised.
beta, theta, delay = (_as_matrix(matrix) for matrix in (beta, theta, delay))
# beta is required, so it is always shape-checked; theta and delay are
# optional and validated only when provided.
for name, matrix in (("beta", beta), ("theta", theta), ("delay", delay)):
if matrix is None and name != "beta":
continue
if getattr(matrix, "shape", None) != (n, n):
raise ValueError(
f"{name} matrix shape {getattr(matrix, 'shape', None)} does not match {n} drive lines"
)
any_traced = any(_is_traced(matrix) for matrix in (beta, theta, delay) if matrix is not None)
xp = _select_array_module(any_traced)
matrix_beta = xp.asarray(beta, dtype=float)
matrix_theta = xp.zeros((n, n), dtype=float) if theta is None else xp.asarray(theta, dtype=float)
matrix_delay = xp.zeros((n, n), dtype=float) if delay is None else xp.asarray(delay, dtype=float)
diagonal = xp.eye(n, dtype=bool)
matrix_beta = xp.where(diagonal, 1.0, matrix_beta)
matrix_theta = xp.where(diagonal, 0.0, matrix_theta)
matrix_delay = xp.where(diagonal, 0.0, matrix_delay)
kept = [
transform for transform in self._signal_chain
if not isinstance(transform, (Crosstalk, CrosstalkMatrix))
]
self._signal_chain = kept + [
CrosstalkMatrix(order, matrix_beta, matrix_theta, matrix_delay)
]
[docs]
def copy(self, device_map: dict[str, Any], coupling_map: dict[str, Any] | None = None) -> "ControlEquipment":
"""Return a structural copy with drive lines rebound to target maps.
Parameters
----------
device_map : dict[str, object]
Device labels to copied devices.
coupling_map : dict[str, object], optional
Coupling labels to copied couplings.
Coupling-target lines rebind via *coupling_map*, keyed by coupling
label; device-target lines rebind via *device_map*.
"""
copied_lines = []
for line in self._lines:
if isinstance(line, CouplingDrive):
if coupling_map is None or line.target_label not in coupling_map:
raise KeyError(
f"No coupling '{line.target_label}' in the target map for edge line '{line.label}'."
)
copied_lines.append(line.copy(target=coupling_map[line.target_label]))
else:
copied_lines.append(
line.copy(target=None if line.device_label is None else device_map[line.device_label])
)
return type(self)(lines=copied_lines,
signal_chain=[transform.copy() for transform in self._signal_chain] or None)
[docs]
def to_dict(self) -> dict[str, Any]:
"""Serialize into a JSON-safe dictionary."""
data: dict[str, Any] = {"lines": [line.to_dict() for line in self._lines]}
if self._signal_chain:
data["signal_chain"] = [transform.to_dict() for transform in self._signal_chain]
return data
[docs]
@classmethod
def from_dict(
cls,
d: dict[str, Any],
dev_map: dict[str, Any],
coupling_map: dict[str, Any] | None = None,
) -> "ControlEquipment":
"""Reconstruct from :meth:`to_dict` output and target maps.
Parameters
----------
d : dict
Serialized equipment dictionary.
dev_map : dict[str, object]
Device labels to target objects.
coupling_map : dict[str, object], optional
Coupling labels to target objects.
Each line's ``target_label`` is resolved against *dev_map* first,
then *coupling_map* — device and coupling labels are disjoint by
Chip construction, so at most one map holds the label.
"""
lines: list[BaseDrive] = []
for line_dict in d.get("lines", []):
target_label = line_dict.get("target_label")
if target_label is None:
target = None
elif target_label in dev_map:
target = dev_map[target_label]
elif coupling_map is not None and target_label in coupling_map:
target = coupling_map[target_label]
else:
raise KeyError(f"No device or coupling named '{target_label}' in the target maps.")
lines.append(BaseDrive.from_dict(line_dict, target=target))
signal_chain = [SignalTransform.from_dict(td) for td in d.get("signal_chain", [])]
return cls(lines=lines, signal_chain=signal_chain or None)
