Sync mirror with GitHub main (v8.1.62): subpackage refactor, mitigation promotions, MCP server refactor

2a87ede verified 2 months ago

9.78 kB

"""
Real Zero-Noise Extrapolation (ZNE) error mitigation, wired to the same
engine and real noise channels (dense_evolution.NoiseModel) the rest of
the Composer uses. Noise is scaled by running the real channel at
noise_p, 2*noise_p, 3*noise_p and Richardson-extrapolating the measured
Pauli expectation back to zero noise via dense_evolution's own
zero_noise_extrapolation -- not a fabricated mitigated curve.
"""

from dataclasses import dataclass
from typing import Optional

import numpy as np

import dense_evolution as de

__all__ = ['MitigationResult', 'run_zne_mitigation', 'DensityMatrixZNEResult', 'run_density_matrix_zne']

_DEFAULT_NOISE_FACTORS = (1.0, 2.0, 3.0)
# polynomial_extrapolate's own docstring: with exactly degree+1 points the
# fit is the unique interpolating polynomial, IDENTICAL to
# richardson_extrapolate -- the extra 2 points here (4x, 5x) are what
# actually changes anything, trading a bit of interpolation bias for
# averaging down statistical noise across more measured scales (verified
# in that docstring against a real 5-seed sweep, not asserted here).
_POLYNOMIAL_NOISE_FACTORS = (1.0, 2.0, 3.0, 4.0, 5.0)
_POLYNOMIAL_DEGREE = 2


@dataclass
class MitigationResult:
    n_qubits: int
    pauli_string: str
    ideal_expectation: float
    noise_factors: list
    noisy_expectations: list
    zne_extrapolated: float
    extrapolation_method: str = "richardson"


def run_zne_mitigation(
    qasm_text: str,
    pauli_string: str,
    noise_model: str,
    noise_p: float,
    seed: Optional[int] = None,
    noise_factors=None,
    n_trials: int = 200,
    extrapolation_method: str = "richardson",
) -> MitigationResult:
    """Real ZNE: <P> is measured at the real ideal state and at the real
    channel applied at noise_p * each factor, then extrapolated to zero
    noise -- either Richardson (dense_evolution.zero_noise_extrapolation,
    the exact interpolating polynomial through 3 points, the default) or
    a degree-2 least-squares polynomial fit through 5 points
    (dense_evolution.polynomial_extrapolate) -- caller's choice, not
    silently picked: noise_factors defaults to the 3- or 5-point set that
    matches whichever method was requested, unless the caller overrides
    it explicitly.

    NoiseModel.apply_to_sv is a *stochastic single-shot* Kraus draw (one
    random outcome per call, not the channel's averaged/ensemble
    behavior) -- feeding a single draw straight into ZNE gives a
    discontinuous, meaningless curve (verified: bitflip/depolarizing
    jumped 1.0 -> 0.0 -> 0.0 across noise scales instead of decaying
    smoothly). <P> under a Kraus channel is Tr(rho P) = mean over the
    channel's trajectories, so each scale here averages n_trials
    independent stochastic draws -- the actual expectation value ZNE is
    defined against, not one random sample of it.

    pauli_string uses dense_evolution's own qubit-0-is-position-0
    convention (same as pauli_expectation), independent of Qiskit's
    little-endian display convention used elsewhere on this page --
    this function never touches a Qiskit-ordered array.
    """
    if extrapolation_method not in ("richardson", "polynomial"):
        raise ValueError(
            f"unknown extrapolation_method {extrapolation_method!r}, must be 'richardson' or 'polynomial'"
        )
    if noise_factors is None:
        noise_factors = (
            _POLYNOMIAL_NOISE_FACTORS if extrapolation_method == "polynomial" else _DEFAULT_NOISE_FACTORS
        )

    # Parsed with dense_evolution's own QASMParser, never Qiskit's --
    # QuantumCircuit.from_qasm_str itself segfaults on macOS (see
    # dashboard_core/engine.py's module docstring for the full story);
    # this function never needs a Qiskit circuit object at all, only the
    # ideal statevector, so there's no reason to build one here.
    parsed = de.QASMParser().parse(qasm_text)
    n_qubits = parsed.n_qubits
    if len(pauli_string) != n_qubits:
        raise ValueError(f"pauli_string length {len(pauli_string)} != n_qubits {n_qubits}")
    if noise_model not in de.NoiseModel.MODELS:
        raise ValueError(f"unknown noise model {noise_model!r}, must be one of {de.NoiseModel.MODELS}")

    # Same real anti-OOM guard as dashboard_core.engine.run_circuit_from_qasm
    # -- this kernel now runs on whatever machine a Composer visitor has,
    # not just a dev laptop, and this function was the one real gap that
    # never checked before allocating (engine.py's own functions always did).
    required_mb = (2 ** n_qubits) * 16 / 1e6
    de.chunk.SafeMemoryGuard().check_allocation(required_mb, context=f"{n_qubits}-qubit statevector")

    sim = de.DenseSVSimulator(n_qubits, use_float32=False)
    sim.run_circuit(parsed.to_tuples())
    sv_ideal = np.asarray(sim.sv)
    ideal_expectation = float(np.real(de.pauli_expectation(sv_ideal, pauli_string)))

    rng = np.random.default_rng(seed)
    noisy_expectations = []
    for factor in noise_factors:
        scaled_p = min(noise_p * factor, 1.0)
        trial_values = np.empty(n_trials)
        for i in range(n_trials):
            sv_noisy = de.NoiseModel.apply_to_sv(
                sv_ideal.copy(), n_qubits, noise_model, scaled_p, rng=rng,
            )
            trial_values[i] = np.real(de.pauli_expectation(sv_noisy, pauli_string))
        noisy_expectations.append(float(trial_values.mean()))

    if extrapolation_method == "polynomial":
        zne_value = float(de.polynomial_extrapolate(noisy_expectations, list(noise_factors), degree=_POLYNOMIAL_DEGREE))
    else:
        zne_value = float(de.zero_noise_extrapolation(noisy_expectations, list(noise_factors)))

    return MitigationResult(
        n_qubits=n_qubits,
        pauli_string=pauli_string,
        ideal_expectation=ideal_expectation,
        noise_factors=list(noise_factors),
        noisy_expectations=noisy_expectations,
        zne_extrapolated=zne_value,
        extrapolation_method=extrapolation_method,
    )


@dataclass
class DensityMatrixZNEResult:
    n_qubits: int
    noise_factors: list
    fidelity_raw: float          # Uhlmann fidelity: ideal vs raw noisy (base scale) rho
    fidelity_corrected: float    # Uhlmann fidelity: ideal vs ZNE-corrected rho


def run_density_matrix_zne(
    qasm_text: str,
    noise_model: str,
    noise_p: float,
    seed: Optional[int] = None,
    noise_factors=_DEFAULT_NOISE_FACTORS,
    n_trials: int = 200,
) -> DensityMatrixZNEResult:
    """Density-matrix ZNE (dense_evolution.zne_density_matrix): builds a
    real Monte-Carlo density-matrix estimate at each noise scale (the
    mean of n_trials |psi_k><psi_k| projectors from independent
    NoiseModel draws -- a real ensemble reconstruction, not a single
    trajectory), extrapolates to zero noise, and projects onto the
    nearest physical (PSD, trace-1) density matrix internally.

    Graded (never fed back into the extrapolation) against the true
    ideal density matrix via dense_evolution.uhlmann_fidelity, so the
    reported improvement is an honest measurement of whether the
    correction actually helped -- matches the pattern in
    zne_density_matrix's own docstring (experiments/matrix_healing_zne.py:
    raw ~0.865, corrected ~0.947 on a 2-qubit Bell state).
    """
    # Same QASMParser-only parsing as run_zne_mitigation above -- see its
    # comment for why a Qiskit circuit object is never built here.
    parsed = de.QASMParser().parse(qasm_text)
    n_qubits = parsed.n_qubits
    if noise_model not in de.NoiseModel.MODELS:
        raise ValueError(f"unknown noise model {noise_model!r}, must be one of {de.NoiseModel.MODELS}")

    # Density matrices are dim x dim (dim = 2**n_qubits), not just dim --
    # this holds rho_ideal, one rho per noise factor (rhos_at_scales, kept
    # alive simultaneously so zne_density_matrix can extrapolate across all
    # of them at once) and rho_corrected: (len(noise_factors) + 2) separate
    # dim*dim complex128 arrays, quadratically worse than a plain
    # statevector at the same qubit count. Same real anti-OOM guard as
    # dashboard_core.engine.run_circuit_from_qasm, sized for what this
    # function actually allocates -- this was the one real gap that never
    # checked before allocating.
    dim = 2 ** n_qubits
    required_mb = dim * dim * 16 / 1e6 * (len(noise_factors) + 2)
    de.chunk.SafeMemoryGuard().check_allocation(required_mb, context=f"{n_qubits}-qubit density matrix ZNE")

    sim = de.DenseSVSimulator(n_qubits, use_float32=False)
    sim.run_circuit(parsed.to_tuples())
    sv_ideal = np.asarray(sim.sv)
    rho_ideal = np.outer(sv_ideal, sv_ideal.conj())

    rng = np.random.default_rng(seed)
    rhos_at_scales = []
    for factor in noise_factors:
        scaled_p = min(noise_p * factor, 1.0)
        rho_acc = np.zeros((dim, dim), dtype=complex)
        for _ in range(n_trials):
            sv_noisy = de.NoiseModel.apply_to_sv(
                sv_ideal.copy(), n_qubits, noise_model, scaled_p, rng=rng,
            )
            rho_acc += np.outer(sv_noisy, sv_noisy.conj())
        rho_acc /= n_trials
        rhos_at_scales.append(rho_acc)

    rho_corrected = np.asarray(de.zne_density_matrix(rhos_at_scales, list(noise_factors)))

    fidelity_raw = float(de.uhlmann_fidelity(rho_ideal, rhos_at_scales[0]))
    fidelity_corrected = float(de.uhlmann_fidelity(rho_ideal, rho_corrected))

    return DensityMatrixZNEResult(
        n_qubits=n_qubits,
        noise_factors=list(noise_factors),
        fidelity_raw=fidelity_raw,
        fidelity_corrected=fidelity_corrected,
    )

	"""
	Real Zero-Noise Extrapolation (ZNE) error mitigation, wired to the same
	engine and real noise channels (dense_evolution.NoiseModel) the rest of
	the Composer uses. Noise is scaled by running the real channel at
	noise_p, 2noise_p, 3noise_p and Richardson-extrapolating the measured
	Pauli expectation back to zero noise via dense_evolution's own
	zero_noise_extrapolation -- not a fabricated mitigated curve.
	"""

	from dataclasses import dataclass
	from typing import Optional

	import numpy as np

	import dense_evolution as de

	__all__ = ['MitigationResult', 'run_zne_mitigation', 'DensityMatrixZNEResult', 'run_density_matrix_zne']

	_DEFAULT_NOISE_FACTORS = (1.0, 2.0, 3.0)
	# polynomial_extrapolate's own docstring: with exactly degree+1 points the
	# fit is the unique interpolating polynomial, IDENTICAL to
	# richardson_extrapolate -- the extra 2 points here (4x, 5x) are what
	# actually changes anything, trading a bit of interpolation bias for
	# averaging down statistical noise across more measured scales (verified
	# in that docstring against a real 5-seed sweep, not asserted here).
	_POLYNOMIAL_NOISE_FACTORS = (1.0, 2.0, 3.0, 4.0, 5.0)
	_POLYNOMIAL_DEGREE = 2


	@dataclass
	class MitigationResult:
	n_qubits: int
	pauli_string: str
	ideal_expectation: float
	noise_factors: list
	noisy_expectations: list
	zne_extrapolated: float
	extrapolation_method: str = "richardson"


	def run_zne_mitigation(
	qasm_text: str,
	pauli_string: str,
	noise_model: str,
	noise_p: float,
	seed: Optional[int] = None,
	noise_factors=None,
	n_trials: int = 200,
	extrapolation_method: str = "richardson",
	) -> MitigationResult:
	"""Real ZNE: <P> is measured at the real ideal state and at the real
	channel applied at noise_p * each factor, then extrapolated to zero
	noise -- either Richardson (dense_evolution.zero_noise_extrapolation,
	the exact interpolating polynomial through 3 points, the default) or
	a degree-2 least-squares polynomial fit through 5 points
	(dense_evolution.polynomial_extrapolate) -- caller's choice, not
	silently picked: noise_factors defaults to the 3- or 5-point set that
	matches whichever method was requested, unless the caller overrides
	it explicitly.

	NoiseModel.apply_to_sv is a stochastic single-shot Kraus draw (one
	random outcome per call, not the channel's averaged/ensemble
	behavior) -- feeding a single draw straight into ZNE gives a
	discontinuous, meaningless curve (verified: bitflip/depolarizing
	jumped 1.0 -> 0.0 -> 0.0 across noise scales instead of decaying
	smoothly). <P> under a Kraus channel is Tr(rho P) = mean over the
	channel's trajectories, so each scale here averages n_trials
	independent stochastic draws -- the actual expectation value ZNE is
	defined against, not one random sample of it.

	pauli_string uses dense_evolution's own qubit-0-is-position-0
	convention (same as pauli_expectation), independent of Qiskit's
	little-endian display convention used elsewhere on this page --
	this function never touches a Qiskit-ordered array.
	"""
	if extrapolation_method not in ("richardson", "polynomial"):
	raise ValueError(
	f"unknown extrapolation_method {extrapolation_method!r}, must be 'richardson' or 'polynomial'"
	)
	if noise_factors is None:
	noise_factors = (
	_POLYNOMIAL_NOISE_FACTORS if extrapolation_method == "polynomial" else _DEFAULT_NOISE_FACTORS
	)

	# Parsed with dense_evolution's own QASMParser, never Qiskit's --
	# QuantumCircuit.from_qasm_str itself segfaults on macOS (see
	# dashboard_core/engine.py's module docstring for the full story);
	# this function never needs a Qiskit circuit object at all, only the
	# ideal statevector, so there's no reason to build one here.
	parsed = de.QASMParser().parse(qasm_text)
	n_qubits = parsed.n_qubits
	if len(pauli_string) != n_qubits:
	raise ValueError(f"pauli_string length {len(pauli_string)} != n_qubits {n_qubits}")
	if noise_model not in de.NoiseModel.MODELS:
	raise ValueError(f"unknown noise model {noise_model!r}, must be one of {de.NoiseModel.MODELS}")

	# Same real anti-OOM guard as dashboard_core.engine.run_circuit_from_qasm
	# -- this kernel now runs on whatever machine a Composer visitor has,
	# not just a dev laptop, and this function was the one real gap that
	# never checked before allocating (engine.py's own functions always did).
	required_mb = (2 ** n_qubits) * 16 / 1e6
	de.chunk.SafeMemoryGuard().check_allocation(required_mb, context=f"{n_qubits}-qubit statevector")

	sim = de.DenseSVSimulator(n_qubits, use_float32=False)
	sim.run_circuit(parsed.to_tuples())
	sv_ideal = np.asarray(sim.sv)
	ideal_expectation = float(np.real(de.pauli_expectation(sv_ideal, pauli_string)))

	rng = np.random.default_rng(seed)
	noisy_expectations = []
	for factor in noise_factors:
	scaled_p = min(noise_p * factor, 1.0)
	trial_values = np.empty(n_trials)
	for i in range(n_trials):
	sv_noisy = de.NoiseModel.apply_to_sv(
	sv_ideal.copy(), n_qubits, noise_model, scaled_p, rng=rng,
	)
	trial_values[i] = np.real(de.pauli_expectation(sv_noisy, pauli_string))
	noisy_expectations.append(float(trial_values.mean()))

	if extrapolation_method == "polynomial":
	zne_value = float(de.polynomial_extrapolate(noisy_expectations, list(noise_factors), degree=_POLYNOMIAL_DEGREE))
	else:
	zne_value = float(de.zero_noise_extrapolation(noisy_expectations, list(noise_factors)))

	return MitigationResult(
	n_qubits=n_qubits,
	pauli_string=pauli_string,
	ideal_expectation=ideal_expectation,
	noise_factors=list(noise_factors),
	noisy_expectations=noisy_expectations,
	zne_extrapolated=zne_value,
	extrapolation_method=extrapolation_method,
	)


	@dataclass
	class DensityMatrixZNEResult:
	n_qubits: int
	noise_factors: list
	fidelity_raw: float # Uhlmann fidelity: ideal vs raw noisy (base scale) rho
	fidelity_corrected: float # Uhlmann fidelity: ideal vs ZNE-corrected rho


	def run_density_matrix_zne(
	qasm_text: str,
	noise_model: str,
	noise_p: float,
	seed: Optional[int] = None,
	noise_factors=_DEFAULT_NOISE_FACTORS,
	n_trials: int = 200,
	) -> DensityMatrixZNEResult:
	"""Density-matrix ZNE (dense_evolution.zne_density_matrix): builds a
	real Monte-Carlo density-matrix estimate at each noise scale (the
	mean of n_trials \|psi_k><psi_k\| projectors from independent
	NoiseModel draws -- a real ensemble reconstruction, not a single
	trajectory), extrapolates to zero noise, and projects onto the
	nearest physical (PSD, trace-1) density matrix internally.

	Graded (never fed back into the extrapolation) against the true
	ideal density matrix via dense_evolution.uhlmann_fidelity, so the
	reported improvement is an honest measurement of whether the
	correction actually helped -- matches the pattern in
	zne_density_matrix's own docstring (experiments/matrix_healing_zne.py:
	raw ~0.865, corrected ~0.947 on a 2-qubit Bell state).
	"""
	# Same QASMParser-only parsing as run_zne_mitigation above -- see its
	# comment for why a Qiskit circuit object is never built here.
	parsed = de.QASMParser().parse(qasm_text)
	n_qubits = parsed.n_qubits
	if noise_model not in de.NoiseModel.MODELS:
	raise ValueError(f"unknown noise model {noise_model!r}, must be one of {de.NoiseModel.MODELS}")

	# Density matrices are dim x dim (dim = 2**n_qubits), not just dim --
	# this holds rho_ideal, one rho per noise factor (rhos_at_scales, kept
	# alive simultaneously so zne_density_matrix can extrapolate across all
	# of them at once) and rho_corrected: (len(noise_factors) + 2) separate
	# dim*dim complex128 arrays, quadratically worse than a plain
	# statevector at the same qubit count. Same real anti-OOM guard as
	# dashboard_core.engine.run_circuit_from_qasm, sized for what this
	# function actually allocates -- this was the one real gap that never
	# checked before allocating.
	dim = 2 ** n_qubits
	required_mb = dim * dim * 16 / 1e6 * (len(noise_factors) + 2)
	de.chunk.SafeMemoryGuard().check_allocation(required_mb, context=f"{n_qubits}-qubit density matrix ZNE")

	sim = de.DenseSVSimulator(n_qubits, use_float32=False)
	sim.run_circuit(parsed.to_tuples())
	sv_ideal = np.asarray(sim.sv)
	rho_ideal = np.outer(sv_ideal, sv_ideal.conj())

	rng = np.random.default_rng(seed)
	rhos_at_scales = []
	for factor in noise_factors:
	scaled_p = min(noise_p * factor, 1.0)
	rho_acc = np.zeros((dim, dim), dtype=complex)
	for _ in range(n_trials):
	sv_noisy = de.NoiseModel.apply_to_sv(
	sv_ideal.copy(), n_qubits, noise_model, scaled_p, rng=rng,
	)
	rho_acc += np.outer(sv_noisy, sv_noisy.conj())
	rho_acc /= n_trials
	rhos_at_scales.append(rho_acc)

	rho_corrected = np.asarray(de.zne_density_matrix(rhos_at_scales, list(noise_factors)))

	fidelity_raw = float(de.uhlmann_fidelity(rho_ideal, rhos_at_scales[0]))
	fidelity_corrected = float(de.uhlmann_fidelity(rho_ideal, rho_corrected))

	return DensityMatrixZNEResult(
	n_qubits=n_qubits,
	noise_factors=list(noise_factors),
	fidelity_raw=fidelity_raw,
	fidelity_corrected=fidelity_corrected,
	)