"""Shared helpers for pipeline stages."""
from __future__ import annotations
from pathlib import Path
from typing import Any
import numpy as np
FILTER_CLASSIFIERS = frozenset(
{"threshold", "leakage", "multilevel", "herqules", "lda", "qda", "mlp", "gmm"},
)
[docs]
def resolve_cache_paths(
dcfg: Any,
*,
cache_dir: str | Path | None,
demod_cache_path: str | Path | None,
) -> dict[str, Path | None]:
"""Resolve demod / transition / spectral cache locations.
Args:
dcfg: Data config block (``demodulation``, ``transitions``).
cache_dir: If set, all caches live under this directory
(``demod_traces.pkl``, ``spectral_cache``, ``transitions``).
demod_cache_path: Explicit demod-cache file; overrides
``dcfg.demodulation.cache_path`` when ``cache_dir`` is unset.
Returns:
Dict with keys ``demod``, ``spectral``, and ``transition``
(each a :class:`~pathlib.Path` or ``None``).
"""
dm = dcfg.demodulation
tcfg = dcfg.transitions
if cache_dir is not None:
root = Path(cache_dir)
root.mkdir(parents=True, exist_ok=True)
return {
"demod": root / "demod_traces.pkl",
"spectral": root / "spectral_cache",
"transition": root / "transitions",
}
demod: Path | None = None
if demod_cache_path is not None:
demod = Path(demod_cache_path)
elif dm.cache_path:
demod = Path(dm.cache_path)
spectral = Path(tcfg.spectral_cache_path) if tcfg.spectral_cache_path else None
transition = Path(tcfg.cache_path) if tcfg.cache_path else None
return {"demod": demod, "spectral": spectral, "transition": transition}
def _compute_split_indices(labels: np.ndarray, sc: Any) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Compute train/val/test indices from labels before demodulation."""
from arcade.data.processors.splitter import split
n = len(labels)
dummy = np.arange(n)
split_data = split(
dummy,
labels,
mode=sc.mode,
train_ratio=sc.train_ratio,
val_ratio=sc.val_ratio,
test_ratio=sc.test_ratio,
train_per_class=sc.train_per_class,
val_per_class=sc.val_per_class,
test_per_class=sc.test_per_class,
stratified=sc.stratified,
shuffle=sc.shuffle,
seed=sc.seed,
)
return split_data.train_indices, split_data.val_indices, split_data.test_indices
[docs]
def transition_plan(
tcfg: Any,
classifiers_run: list[str] | None,
) -> tuple[bool, bool, str]:
"""Decide which transition detectors to run for this pipeline.
Args:
tcfg: Transitions config (``enabled``, ``auto``, ``method``, …).
classifiers_run: Classifier names from ``classifiers.run``. Used when
``tcfg.auto`` is true to enable leakage or filter-only detection.
Returns:
Tuple ``(run_relaxations, run_leakage, leakage_method)``.
"""
if not tcfg.enabled:
return False, False, tcfg.method
if tcfg.auto and classifiers_run is not None:
needs_leakage = bool({"leakage", "multilevel"} & set(classifiers_run))
needs_filters = any(c in FILTER_CLASSIFIERS for c in classifiers_run)
if needs_leakage:
method = tcfg.method if tcfg.method in ("spectral", "gmm") else "spectral"
return True, True, method
if needs_filters:
return True, False, "centroid_radius"
return False, False, tcfg.method
if tcfg.method in ("spectral", "gmm"):
return True, True, tcfg.method
return True, False, tcfg.method
[docs]
def hydrate_data_out(data_out: dict[str, Any]) -> dict[str, Any]:
"""Ensure flat train/val/test keys exist (e.g. after loading old pickles).
Args:
data_out: Data-stage dict that may only expose nested ``splits``.
Returns:
The same dict if already flat, otherwise a shallow copy with
``train_traces`` / ``train_labels`` / … copied from ``splits``.
"""
splits = data_out.get("splits")
if splits is None or "train_traces" in data_out:
return data_out
return {
**data_out,
"train_traces": splits.train_traces,
"train_labels": splits.train_labels,
"val_traces": splits.val_traces,
"val_labels": splits.val_labels,
"test_traces": splits.test_traces,
"test_labels": splits.test_labels,
"train_indices": splits.train_indices,
"val_indices": splits.val_indices,
"test_indices": splits.test_indices,
}
[docs]
def data_out_data_fraction(data_out: dict[str, Any], frac: float) -> dict[str, Any]:
"""Subsample train/val/test splits; keep ``transition_result`` unchanged.
Args:
data_out: Hydrated or nest-only data-stage dict.
frac: Fraction of each split to keep, in ``(0, 1]``.
Returns:
Shallow copy with subsampled traces, labels, and indices.
Raises:
ValueError: If ``frac`` is not in ``(0, 1]``.
"""
if not 0.0 < frac <= 1.0:
raise ValueError("frac must be in (0, 1]")
data = hydrate_data_out(data_out)
def _n(tr: Any) -> int:
if isinstance(tr, dict):
return int(next(iter(tr.values())).shape[0])
return int(np.asarray(tr).shape[0])
def _take(tr: Any, labels: Any) -> tuple[Any, np.ndarray, np.ndarray]:
n = _n(tr)
k = max(1, int(round(n * frac)))
idx = np.arange(min(k, n))
if isinstance(tr, dict):
tr_sub = {int(q): np.asarray(t)[idx] for q, t in tr.items()}
else:
tr_sub = np.asarray(tr)[idx]
lab = np.asarray(labels)[idx]
return tr_sub, lab, idx
out = dict(data)
tt, tl, ti = _take(data["train_traces"], data["train_labels"])
vt, vl, vi = _take(data["val_traces"], data["val_labels"])
xt, xl, xe = _take(data["test_traces"], data["test_labels"])
out["train_traces"], out["train_labels"] = tt, tl
out["val_traces"], out["val_labels"] = vt, vl
out["test_traces"], out["test_labels"] = xt, xl
out["train_indices"] = np.asarray(data["train_indices"])[ti]
out["val_indices"] = np.asarray(data["val_indices"])[vi]
out["test_indices"] = np.asarray(data["test_indices"])[xe]
return out
def trace_shape_descr(traces: Any) -> str:
"""Format a short shape description for logging.
Args:
traces: Per-qubit dict of arrays or a single array-like.
Returns:
Human-readable shape string (e.g. ``{q0(100, 200), …}``).
"""
if isinstance(traces, dict):
parts = [f"q{k}{tuple(v.shape)}" for k, v in sorted(traces.items())]
return "{" + ", ".join(parts) + "}"
return str(np.asarray(traces).shape)
[docs]
def data_out_labels_placeholder(features: Any) -> Any:
"""Zero labels placeholder for DSE resource-only sweeps.
Args:
features: Feature matrix whose leading dimension sets label length.
Returns:
``int64`` zero vector of length ``len(features)``.
"""
return np.zeros(len(features), dtype=np.int64)