utils.py

"""Shared helpers for the error-mitigation teaching-notebook series.

These utilities cover the pieces every notebook needs: the fixed 4-qubit
mirror circuit and its ⟨ZZZZ⟩ observable, Aer noise-model builders (one per
noise category), the readout-mitigation and zero-noise-extrapolation
routines developed by hand in notebooks 2 and 3, counts caching to JSON,
and a consistent plot style.

The running experiment (used everywhere in this series)
-------------------------------------------------------
`mirror_circuit(k)` applies a fixed entangling block U followed by U†, k
times, so the whole circuit composes to the identity. Started from |0000⟩
the ideal output is exactly |0000⟩, which pins the observable

    ⟨Z ⊗ Z ⊗ Z ⊗ Z⟩ = +1   (exactly, with zero shot ambiguity about the target)

Every notebook measures this same observable, so "how close to +1 did we
get?" is directly comparable across noise models, mitigation methods, and
real hardware. Counts keys follow Qiskit's little-endian convention
(rightmost character = qubit 0); every helper here strips register spaces
first.
"""

from __future__ import annotations

import json
from datetime import datetime, timezone
from functools import reduce
from pathlib import Path

import numpy as np
from qiskit import QuantumCircuit
from qiskit_aer import AerSimulator
from qiskit_aer.noise import (
    NoiseModel,
    ReadoutError,
    depolarizing_error,
    thermal_relaxation_error,
)

NUM_QUBITS = 4
IDEAL_EXPECTATION = 1.0  # ⟨ZZZZ⟩ of the mirror circuit, by construction

# Validated categorical palette (light surface) + chart chrome.
COLORS = {
    "blue": "#2a78d6",
    "aqua": "#1baf7a",
    "yellow": "#eda100",
    "red": "#e34948",
    "violet": "#4a3aa7",
    "ink": "#0b0b0b",
    "muted": "#898781",
    "grid": "#e1e0d9",
    "baseline": "#c3c2b7",
}


def set_plot_style() -> None:
    """Apply a consistent, recessive matplotlib style for the series."""
    import matplotlib as mpl

    mpl.rcParams.update(
        {
            "axes.grid": True,
            "grid.color": COLORS["grid"],
            "grid.linewidth": 0.8,
            "axes.edgecolor": COLORS["baseline"],
            "axes.labelcolor": COLORS["ink"],
            "axes.titlecolor": COLORS["ink"],
            "xtick.color": COLORS["muted"],
            "ytick.color": COLORS["muted"],
            "axes.spines.top": False,
            "axes.spines.right": False,
            "font.family": "sans-serif",
            "figure.figsize": (7.0, 4.0),
            "figure.dpi": 110,
            "lines.linewidth": 2.0,
        }
    )


# ---------------------------------------------------------------------------
# The fixed experiment: mirror circuit + ⟨ZZZZ⟩
# ---------------------------------------------------------------------------

MIRROR_ANGLES = (0.63, -1.18, 0.41, 0.97)  # fixed, arbitrary, hard-coded


def mirror_half(num_qubits: int = NUM_QUBITS) -> QuantumCircuit:
    """The forward half U of the mirror circuit: an Ry layer + a CX ladder.

    U on its own creates genuine 4-qubit entanglement; it is only the
    U·U† composition that returns to |0000⟩.
    """
    qc = QuantumCircuit(num_qubits)
    for q in range(num_qubits):
        qc.ry(MIRROR_ANGLES[q % len(MIRROR_ANGLES)], q)
    for q in range(num_qubits - 1):
        qc.cx(q, q + 1)
    return qc


def mirror_circuit(
    num_layers: int = 1, num_qubits: int = NUM_QUBITS, measure: bool = False
) -> QuantumCircuit:
    """(U · U†)^num_layers on |0000⟩ — logically the identity at any depth.

    Barriers separate U from U† so no optimizer can cancel the mirror away;
    the simulator (and real hardware) must execute every gate, which is the
    whole point: depth grows, the ideal answer does not change.
    """
    u = mirror_half(num_qubits)
    qc = QuantumCircuit(num_qubits)
    for _ in range(num_layers):
        qc.compose(u, inplace=True)
        qc.barrier()
        qc.compose(u.inverse(), inplace=True)
        qc.barrier()
    if measure:
        qc.measure_all()
    return qc


def zzzz_expectation(counts: dict[str, int]) -> float:
    """⟨Z⊗Z⊗...⊗Z⟩ from counts: +1 for even-parity bitstrings, -1 for odd."""
    total = sum(counts.values())
    signed = sum(
        (-1) ** key.replace(" ", "").count("1") * n for key, n in counts.items()
    )
    return signed / total


def sample_counts(
    circuit: QuantumCircuit,
    noise_model: NoiseModel | None = None,
    shots: int = 8192,
    seed: int = 1234,
) -> dict[str, int]:
    """Run a measured circuit on Aer (optionally noisy); return cleaned counts.

    The series builds circuits from {ry, x, cx} only — exactly the
    instructions the noise models attach errors to — so no transpilation is
    needed and no optimizer gets a chance to simplify the mirror away.
    """
    sim = AerSimulator(noise_model=noise_model)
    result = sim.run(circuit, shots=shots, seed_simulator=seed).result()
    return {k.replace(" ", ""): v for k, v in result.get_counts().items()}


# ---------------------------------------------------------------------------
# Noise models, one per category (built by hand in notebook 1)
# ---------------------------------------------------------------------------

READOUT_P01 = 0.02  # P(read 1 | prepared 0)
READOUT_P10 = 0.06  # P(read 0 | prepared 1) — asymmetric on purpose
GATE_P1 = 0.002  # depolarizing probability, 1-qubit gates
GATE_P2 = 0.015  # depolarizing probability, 2-qubit gates
T1_NS = 120_000.0  # relaxation time
T2_NS = 90_000.0  # dephasing time (T2 <= 2*T1 required)
TIME_1Q_NS = 60.0  # 1-qubit gate duration
TIME_2Q_NS = 400.0  # 2-qubit gate duration


def readout_noise_model(
    p01: float = READOUT_P01, p10: float = READOUT_P10
) -> NoiseModel:
    """Classical bit-flips at measurement only; gates stay perfect."""
    nm = NoiseModel()
    err = ReadoutError([[1 - p01, p01], [p10, 1 - p10]])
    nm.add_all_qubit_readout_error(err)
    return nm


def gate_noise_model(p1: float = GATE_P1, p2: float = GATE_P2) -> NoiseModel:
    """Depolarizing error after every gate; measurement stays perfect."""
    nm = NoiseModel()
    nm.add_all_qubit_quantum_error(depolarizing_error(p1, 1), ["ry", "x"])
    nm.add_all_qubit_quantum_error(depolarizing_error(p2, 2), ["cx"])
    return nm


def thermal_noise_model(
    t1: float = T1_NS,
    t2: float = T2_NS,
    time_1q: float = TIME_1Q_NS,
    time_2q: float = TIME_2Q_NS,
) -> NoiseModel:
    """T1/T2 decay during every gate's duration; measurement stays perfect."""
    nm = NoiseModel()
    err_1q = thermal_relaxation_error(t1, t2, time_1q)
    err_2q = thermal_relaxation_error(t1, t2, time_2q).tensor(
        thermal_relaxation_error(t1, t2, time_2q)
    )
    nm.add_all_qubit_quantum_error(err_1q, ["ry", "x"])
    nm.add_all_qubit_quantum_error(err_2q, ["cx"])
    return nm


def full_noise_model() -> NoiseModel:
    """All three categories together — the series' stand-in for a real device."""
    nm = NoiseModel()
    nm.add_all_qubit_readout_error(
        ReadoutError([[1 - READOUT_P01, READOUT_P01], [READOUT_P10, 1 - READOUT_P10]])
    )
    err_1q = depolarizing_error(GATE_P1, 1).compose(
        thermal_relaxation_error(T1_NS, T2_NS, TIME_1Q_NS)
    )
    err_2q = depolarizing_error(GATE_P2, 2).compose(
        thermal_relaxation_error(T1_NS, T2_NS, TIME_2Q_NS).tensor(
            thermal_relaxation_error(T1_NS, T2_NS, TIME_2Q_NS)
        )
    )
    nm.add_all_qubit_quantum_error(err_1q, ["ry", "x"])
    nm.add_all_qubit_quantum_error(err_2q, ["cx"])
    return nm


# ---------------------------------------------------------------------------
# Readout mitigation (developed by hand in notebook 2)
# ---------------------------------------------------------------------------


def counts_to_probs(counts: dict[str, int], num_qubits: int = NUM_QUBITS) -> np.ndarray:
    """Counts dict -> probability vector indexed by int(bitstring, 2)."""
    probs = np.zeros(2**num_qubits)
    total = sum(counts.values())
    for key, n in counts.items():
        probs[int(key.replace(" ", ""), 2)] = n / total
    return probs


def confusions_from_calibration(
    counts_all0: dict[str, int],
    counts_all1: dict[str, int],
    num_qubits: int = NUM_QUBITS,
) -> list[np.ndarray]:
    """Per-qubit 2x2 confusion matrices from just TWO calibration circuits.

    Prepare |00...0⟩ and |11...1⟩, then marginalize each qubit: assuming
    readout errors are uncorrelated between qubits, those two circuits
    calibrate every qubit at once. Column convention: column j = prepared
    |j⟩, row i = measured i, so each column sums to 1.
    """

    def flip_fraction(counts: dict[str, int], qubit: int, prepared: str) -> float:
        total = sum(counts.values())
        wrong = sum(
            n
            for key, n in counts.items()
            if key.replace(" ", "")[-(qubit + 1)] != prepared
        )
        return wrong / total

    mats = []
    for q in range(num_qubits):
        p01 = flip_fraction(counts_all0, q, "0")  # P(read 1 | prepared 0)
        p10 = flip_fraction(counts_all1, q, "1")  # P(read 0 | prepared 1)
        mats.append(np.array([[1 - p01, p10], [p01, 1 - p10]]))
    return mats


def tensored_confusion_matrix(single_qubit_mats: list[np.ndarray]) -> np.ndarray:
    """Kron per-qubit 2x2 matrices into the full 2^n confusion matrix.

    Index convention matches `counts_to_probs`: state index int(bits, 2),
    little-endian bits, so the highest qubit is the leftmost kron factor.
    """
    return reduce(np.kron, single_qubit_mats[::-1])


def mitigate_probs_inverse(probs: np.ndarray, confusion: np.ndarray) -> np.ndarray:
    """Undo readout error by matrix inversion: solve A x = p_noisy.

    The result is a QUASI-probability vector — entries can be negative.
    Fine for expectation values, unusable for sampling (notebook 2).
    """
    return np.linalg.solve(confusion, probs)


def mitigate_probs_lstsq(probs: np.ndarray, confusion: np.ndarray) -> np.ndarray:
    """Constrained least squares: min ||A x - p_noisy|| with x >= 0, sum(x) = 1."""
    from scipy.optimize import minimize

    dim = len(probs)
    result = minimize(
        lambda x: np.sum((confusion @ x - probs) ** 2),
        x0=np.full(dim, 1.0 / dim),
        method="SLSQP",
        bounds=[(0.0, 1.0)] * dim,
        constraints={"type": "eq", "fun": lambda x: np.sum(x) - 1.0},
        tol=1e-12,
    )
    return result.x


def expectation_from_probs(probs: np.ndarray) -> float:
    """⟨Z⊗Z⊗...⊗Z⟩ from a (quasi-)probability vector: parity-signed sum."""
    num_qubits = int(np.log2(len(probs)))
    signs = np.array(
        [(-1) ** bin(i).count("1") for i in range(2**num_qubits)], dtype=float
    )
    return float(signs @ probs)


def mitigate_expectation(
    counts: dict[str, int],
    single_qubit_mats: list[np.ndarray],
    method: str = "lstsq",
) -> float:
    """Readout-mitigated ⟨ZZZZ⟩ from raw counts, in one call (used in nb 4)."""
    probs = counts_to_probs(counts, len(single_qubit_mats))
    confusion = tensored_confusion_matrix(single_qubit_mats)
    if method == "inverse":
        quasi = mitigate_probs_inverse(probs, confusion)
    elif method == "lstsq":
        quasi = mitigate_probs_lstsq(probs, confusion)
    else:
        raise ValueError(f"unknown method: {method!r}")
    return expectation_from_probs(quasi)


# ---------------------------------------------------------------------------
# Zero-noise extrapolation (developed by hand in notebook 3)
# ---------------------------------------------------------------------------


def fold_circuit(circuit: QuantumCircuit, factor: int) -> QuantumCircuit:
    """Global gate folding: G -> G (G† G)^((factor-1)/2), factor odd.

    Logically identical to `circuit`, but with `factor` times the gates —
    turning the noise level itself into an experimental knob. Barriers
    separate the folds so no compiler cancels G†G. The input must not
    contain measurements (add them after folding).
    """
    if factor % 2 != 1 or factor < 1:
        raise ValueError(f"fold factor must be an odd positive integer, got {factor}")
    if circuit.num_clbits:
        raise ValueError("fold the unmeasured circuit; add measurements after")
    folded = circuit.copy()
    for _ in range((factor - 1) // 2):
        folded.barrier()
        folded.compose(circuit.inverse(), inplace=True)
        folded.barrier()
        folded.compose(circuit, inplace=True)
    return folded


def extrapolate_linear(factors, values) -> float:
    """Straight-line fit through (noise factor, value); return intercept at 0."""
    slope, intercept = np.polyfit(factors, values, 1)
    return float(intercept)


def extrapolate_exponential(factors, values) -> float:
    """Fit a·exp(-b·factor) and return the zero-noise limit a."""
    from scipy.optimize import curve_fit

    (a, b), _ = curve_fit(
        lambda x, a, b: a * np.exp(-b * np.asarray(x)),
        factors,
        values,
        p0=(values[0], 0.1),
        maxfev=10_000,
    )
    return float(a)


# ---------------------------------------------------------------------------
# Counts caching
# ---------------------------------------------------------------------------


def save_counts(counts: dict[str, int], path: str | Path, **metadata) -> None:
    """Cache counts as JSON with provenance metadata (backend, shots, ...)."""
    payload = {
        "metadata": {
            **metadata,
            "total_shots": int(sum(counts.values())),
            "saved_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        },
        "counts": {k: int(v) for k, v in counts.items()},
    }
    Path(path).write_text(json.dumps(payload, indent=2))


def load_counts(path: str | Path) -> tuple[dict[str, int], dict]:
    """Load cached counts; returns (counts, metadata)."""
    payload = json.loads(Path(path).read_text())
    return payload["counts"], payload["metadata"]