Task2_5/Task2_5.py - master - timshk/Shalaev_Stepik

Task2_5/Task2_5.py
271 строка10 KB
timshk
Stepik задания
01 июн 2026, 18:58
01 июн 2026, 18:582e75834
"""
ZNE (Zero-Noise Extrapolation) Wrapper для квантовых схем
ИСПРАВЛЕННАЯ ВЕРСИЯ — удалена проблема с add_register
"""

import numpy as np
from qiskit import QuantumCircuit, transpile, ClassicalRegister
from qiskit_aer import AerSimulator
from qiskit_aer.noise import NoiseModel, depolarizing_error
from qiskit.quantum_info import SparsePauliOp
from qiskit_ibm_runtime.fake_provider import FakeSantiagoV2
from scipy.optimize import curve_fit
import warnings


def _fold_circuit_global(circuit, scale_factor):
    """Глобальный фолдинг схемы: повторяет схему целиком scale_factor раз."""
    if scale_factor <= 1.0:
        return circuit.copy()
    
    k = int(round(scale_factor))
    if k < 1:
        k = 1
    
    folded = circuit.copy()
    
    if k == 1:
        return folded
    
    for i in range(k - 1):
        if i % 2 == 0:
            folded = folded.compose(circuit.inverse(), qubits=range(circuit.num_qubits))
        else:
            folded = folded.compose(circuit, qubits=range(circuit.num_qubits))
    
    return folded


def _add_measurements_for_observable(circuit, observable):
    """
    Добавляет измерения для наблюдаемой.
    ИСПРАВЛЕНО: корректное создание классических регистров.
    """
    if not isinstance(observable, SparsePauliOp):
        try:
            observable = SparsePauliOp.from_list(observable)
        except Exception as e:
            raise ValueError(f"observable должен быть SparsePauliOp или списком кортежей: {e}")
    
    # Создаём копию схемы
    meas_circuit = circuit.copy()
    num_qubits = circuit.num_qubits
    
    # Берём первый терм для измерения (для ZNE обычно один терм)
    pauli_label = list(observable.paulis)[0].to_label()
    coeff = observable.coeffs[0]
    
    # Добавляем гейты для измерения в нужном базисе
    for i, p in enumerate(pauli_label):
        if p == 'X':
            meas_circuit.h(i)
        elif p == 'Y':
            meas_circuit.sdg(i)
            meas_circuit.h(i)
    
    # ИСПРАВЛЕНО: создаём новый классический регистр и добавляем его
    cr = ClassicalRegister(num_qubits, name="meas")
    meas_circuit.add_register(cr)
    
    # Добавляем измерения
    meas_circuit.measure(range(num_qubits), cr)
    
    return meas_circuit, coeff


def _run_and_get_expectation(circuit, backend, observable, shots=4096):
    """
    Запускает схему на бэкенде и вычисляет ожидаемое значение наблюдаемой.
    """
    # Добавляем измерения и получаем коэффициент
    meas_circuit, coeff = _add_measurements_for_observable(circuit, observable)
    
    # Транспиляция и запуск
    transpiled = transpile(meas_circuit, backend)
    job = backend.run(transpiled, shots=shots)
    result = job.result()
    counts = result.get_counts()
    
    # Получаем метку классического регистра для извлечения результатов
    # В Qiskit результаты сохраняются по имени регистра
    meas_creg = meas_circuit.cregs[0].name if meas_circuit.cregs else "meas"
    
    # Извлекаем counts (может быть вложенным словарём)
    if meas_creg in counts:
        counts_dict = counts[meas_creg]
    else:
        counts_dict = counts
    
    # Вычисляем ожидаемое значение
    pauli_label = list(observable.paulis)[0].to_label()
    expectation = 0.0
    total_shots = 0
    
    for bitstring, count in counts_dict.items():
        # Qiskit возвращает строки в little-endian
        bits = bitstring[::-1]
        
        term_val = 1.0
        for i, p in enumerate(pauli_label):
            if p == 'Z':
                bit_val = 1 if bits[i] == '0' else -1
                term_val *= bit_val
        
        expectation += term_val * count
        total_shots += count
    
    if total_shots > 0:
        expectation = expectation / total_shots
    else:
        expectation = 0.0
    
    expectation *= float(coeff)
    
    return expectation


def _extrapolate(scale_factors, raw_values, model='linear'):
    """Экстраполирует значения к scale_factor = 0."""
    x = np.array(scale_factors)
    y = np.array(raw_values)
    
    if model == 'linear':
        coeffs = np.polyfit(x, y, deg=1)
        return float(coeffs[1])
    
    elif model == 'poly':
        with warnings.catch_warnings():
            warnings.simplefilter('ignore', np.RankWarning)
            coeffs = np.polyfit(x, y, deg=2)
        return float(coeffs[2])
    
    elif model == 'exp':
        def exp_func(x, a, b, c):
            return a * np.exp(-b * x) + c
        
        p0 = [y[0] - y[-1], 0.5, y[-1]]
        
        try:
            popt, _ = curve_fit(exp_func, x, y, p0=p0, maxfev=1000)
            return float(popt[0] + popt[2])
        except Exception:
            warnings.warn("Экспоненциальная экстраполяция не сходится, используется линейная")
            coeffs = np.polyfit(x, y, deg=1)
            return float(coeffs[1])
    
    else:
        raise ValueError(f"Неизвестная модель экстраполяции: {model}")


def apply_zne(circuit, backend, observable, scale_factors=None, 
              extrapolation_model='linear', shots=4096, noise_amplification='global'):
    """
    Zero-Noise Extrapolation (ZNE) wrapper.
    """
    if scale_factors is None:
        scale_factors = [1, 2, 3]
    
    print("\n" + "="*60)
    print("ZERO-NOISE EXTRAPOLATION (ZNE)")
    print("="*60)
    print(f"Scale factors: {scale_factors}")
    print(f"Extrapolation model: {extrapolation_model}")
    print(f"Shots: {shots}")
    
    raw_values = []
    
    for sf in scale_factors:
        print(f"\n  Scale factor = {sf}...")
        
        if noise_amplification == 'global':
            folded_circuit = _fold_circuit_global(circuit, sf)
        else:
            # Для простоты используем глобальный фолдинг
            folded_circuit = _fold_circuit_global(circuit, sf)
        
        exp_val = _run_and_get_expectation(folded_circuit, backend, observable, shots)
        raw_values.append(exp_val)
        print(f"    ⟨observable⟩ = {exp_val:.6f}")
    
    mitigated = _extrapolate(scale_factors, raw_values, extrapolation_model)
    
    print("\n" + "="*60)
    print("РЕЗУЛЬТАТЫ ZNE")
    print("="*60)
    print(f"Сырые значения: {dict(zip(scale_factors, raw_values))}")
    print(f"Митигированное значение: {mitigated:.6f}")
    
    return {
        'mitigated_value': mitigated,
        'scale_factors': scale_factors,
        'raw_values': raw_values,
        'extrapolation_model': extrapolation_model
    }


# ============================================================================
# ДЕМОНСТРАЦИЯ
# ============================================================================

def demo():
    """Демонстрация работы ZNE на 2-кубитной схеме состояния Белла."""
    print("\n" + "="*60)
    print("ДЕМОНСТРАЦИЯ ZNE: СОСТОЯНИЕ БЕЛЛА |Φ⁺⟩")
    print("="*60)
    
    # Создание схемы Белла (без измерений)
    circuit = QuantumCircuit(2, name="Bell State")
    circuit.h(0)
    circuit.cx(0, 1)
    
    print(f"\nИсходная схема:")
    print(circuit.draw(output='text'))
    
    # Наблюдаемая: ⟨Z₀Z₁⟩
    observable = SparsePauliOp.from_list([("ZZ", 1.0)])
    
    print(f"\nИдеальное значение ⟨ZZ⟩: 1.0")
    
    # Шумовая модель (деполяризация)
    noise_model = NoiseModel()
    depol_1q = depolarizing_error(0.01, 1)
    depol_2q = depolarizing_error(0.02, 2)
    noise_model.add_all_qubit_quantum_error(depol_1q, ['h', 'x', 'rz'])
    noise_model.add_all_qubit_quantum_error(depol_2q, ['cx'])
    
    noisy_backend = AerSimulator(noise_model=noise_model)
    
    print(f"\nШумовая модель: деполяризация (1% на 1-куб., 2% на 2-куб.)")
    
    # Запуск ZNE
    result = apply_zne(
        circuit=circuit,
        backend=noisy_backend,
        observable=observable,
        scale_factors=[1, 2, 3],
        extrapolation_model='linear',
        shots=4096,
        noise_amplification='global'
    )
    
    # Сравнение
    print("\n" + "="*60)
    print("СРАВНЕНИЕ РЕЗУЛЬТАТОВ")
    print("="*60)
    print(f"Идеальное значение:         1.000000")
    print(f"Зашумленное (sf=1):         {result['raw_values'][0]:.6f}")
    print(f"ZNE-митигированное:         {result['mitigated_value']:.6f}")
    
    error_noisy = abs(1.0 - result['raw_values'][0])
    error_mitigated = abs(1.0 - result['mitigated_value'])
    improvement = error_noisy / error_mitigated if error_mitigated > 0 else float('inf')
    
    print(f"\nОшибка без митигации:        {error_noisy:.6f}")
    print(f"Ошибка с митигацией (ZNE):   {error_mitigated:.6f}")
    print(f"Улучшение:                   {improvement:.1f}x")
    
    return result


if __name__ == "__main__":
    demo()