Vincent quantum final code

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
This script simulates an EAQEC code with a single-parameter recovery gate.
It uses a depolarizing noise model and implements a grid search over rotation angles
to find the optimal recovery angle that maximizes state fidelity.



"""

import numpy as np
import matplotlib.pyplot as plt
from qiskit import QuantumCircuit, transpile
from qiskit.quantum_info import Statevector, state_fidelity
from qiskit_aer import AerSimulator
from qiskit_aer.noise import NoiseModel, depolarizing_error
from qiskit.visualization import plot_histogram
from scipy.signal import find_peaks # Import find_peaks

# Set random seed for reproducibility (only affects NumPy, not AerSimulator unless seeded)
np.random.seed(42)

class EAQECSimulator:
    """
    A class to simulate Entanglement-Assisted Quantum Error Correction (EAQEC) codes
    with a single-parameter recovery gate and depolarizing noise.
    """

    def __init__(self, data_qubits=1, ancilla_qubits=2, noise_strength=0.05):
        """
        Initialize the EAQEC simulator.

        Args:
            data_qubits (int): Number of data qubits
            ancilla_qubits (int): Number of ancilla qubits (including entangled qubits)
            noise_strength (float): Strength of the depolarizing noise (0 to 1)
        """
        self.data_qubits = data_qubits
        self.ancilla_qubits = ancilla_qubits
        self.total_qubits = data_qubits + ancilla_qubits
        self.noise_strength = noise_strength
        # Use unseeded simulator for fluctuations
        self.simulator = AerSimulator()
        self.noise_model = self._create_noise_model()

    def _create_noise_model(self):
        """
        Create a depolarizing noise model for the simulation.

        Returns:
            NoiseModel: A Qiskit noise model with depolarizing errors
        """
        noise_model = NoiseModel()

        # Add depolarizing error to all single-qubit gates
        error = depolarizing_error(self.noise_strength, 1)
        noise_model.add_all_qubit_quantum_error(error, ["rx", "ry", "rz", "h", "x", "y", "z"])

        # Add depolarizing error to all two-qubit gates
        error = depolarizing_error(self.noise_strength, 2)
        noise_model.add_all_qubit_quantum_error(error, ["cx", "cz", "swap"])

        return noise_model

    def create_encoding_circuit(self, initial_state=None):
        """
        Create the encoding circuit for the EAQEC code.

        Args:
            initial_state (list, optional): Initial state amplitudes for data qubits

        Returns:
            QuantumCircuit: The encoding circuit
        """
        # Create a quantum circuit with data and ancilla qubits
        qc = QuantumCircuit(self.total_qubits)

        # Prepare the initial state if provided
        if initial_state is not None:
            # Initialize the data qubit to the specified state
            qc.initialize(initial_state, 0)
        else:
            # Default to |+⟩ state
            qc.h(0)

        # Create entanglement between ancilla qubits (Bell pair)
        qc.h(1)
        qc.cx(1, 2)

        # Entangle data qubit with ancilla qubits
        qc.cx(0, 1)
        qc.cx(0, 2)

        return qc

    def apply_error_correction(self, circuit, recovery_angle):
        """
        Apply error correction to the circuit with a parameterized recovery gate.

        Args:
            circuit (QuantumCircuit): The quantum circuit to apply error correction to
            recovery_angle (float): The angle for the recovery rotation gate

        Returns:
            QuantumCircuit: The circuit with error correction applied
        """
        # Create a new circuit with the same qubits
        qc = circuit.copy()

        # Decode the EAQEC code
        qc.cx(0, 2)
        qc.cx(0, 1)

        # Apply the recovery gate (Y-axis rotation) to the data qubit
        qc.ry(recovery_angle, 0)

        return qc

    def simulate_ideal_circuit(self, initial_state=None):
        """
        Simulate the ideal circuit without noise.

        Args:
            initial_state (list, optional): Initial state amplitudes for data qubits

        Returns:
            Statevector: The final state vector
        """
        # Create the encoding circuit
        qc = self.create_encoding_circuit(initial_state)

        # Apply error correction with no recovery (angle=0)
        qc = self.apply_error_correction(qc, 0)

        # Simulate the circuit
        qc_transpiled = transpile(qc, self.simulator)
        # Save statevector is needed for Statevector.from_instruction
        qc_transpiled.save_statevector(label="ideal_statevector")
        result = self.simulator.run(qc_transpiled).result()

        # Get the final state vector from result
        state = result.data()["ideal_statevector"]
        state = Statevector(state)

        return state

    def simulate_noisy_circuit(self, recovery_angle, initial_state=None):
        """
        Simulate the circuit with noise and recovery.

        Args:
            recovery_angle (float): The angle for the recovery rotation gate
            initial_state (list, optional): Initial state amplitudes for data qubits

        Returns:
            Statevector: The final state vector
        """
        # Create the encoding circuit
        qc = self.create_encoding_circuit(initial_state)

        # Apply error correction with the specified recovery angle
        qc = self.apply_error_correction(qc, recovery_angle)

        # Transpile the circuit first
        qc_transpiled = transpile(qc, self.simulator)

        # Save statevector before running
        qc_transpiled.save_statevector(label="final_statevector")

        # Simulate the circuit with noise
        result = self.simulator.run(
            qc_transpiled,
            noise_model=self.noise_model,
            # Ensure the simulator is not seeded by default to keep fluctuations
            # seed_simulator=None # Explicitly None or omit
        ).result()

        # Get the final state vector from the result data
        state = result.data()["final_statevector"]

        # Convert result to Statevector object for fidelity calculation
        state = Statevector(state)

        return state

    def grid_search(self, angle_range=(0, 2*np.pi), num_steps=40, initial_state=None):
        """
        Perform a grid search over recovery angles to find the optimal angle.

        Args:
            angle_range (tuple): Range of angles to search (min, max)
            num_steps (int): Number of steps in the grid search
            initial_state (list, optional): Initial state amplitudes for data qubits

        Returns:
            tuple: (angles, fidelities, optimal_angle, max_fidelity)
        """
        # Generate angles for grid search
        angles = np.linspace(angle_range[0], angle_range[1], num_steps)
        fidelities = []

        # Get the ideal state
        ideal_state = self.simulate_ideal_circuit(initial_state)

        # Perform grid search
        for angle in angles:
            # Simulate the noisy circuit with the current recovery angle
            noisy_state = self.simulate_noisy_circuit(angle, initial_state)

            # Calculate the fidelity between the noisy and ideal states
            fidelity = state_fidelity(noisy_state, ideal_state)
            fidelities.append(fidelity)

        # Find the optimal angle (global max)
        max_idx = np.argmax(fidelities)
        optimal_angle = angles[max_idx]
        max_fidelity = fidelities[max_idx]

        # Return global max info along with angles/fidelities for plotting
        return angles, np.array(fidelities), optimal_angle, max_fidelity

    def plot_fidelity_vs_angle(self, angles, fidelities, optimal_angle, max_fidelity):
        """
        Plot the fidelity vs recovery angle, marking local peaks.

        Args:
            angles (array): Array of angles
            fidelities (array): Array of fidelities
            optimal_angle (float): The optimal recovery angle (global max, not plotted)
            max_fidelity (float): The maximum fidelity achieved (global max, not plotted)

        Returns:
            matplotlib.figure.Figure: The figure object
        """
        fig, ax = plt.subplots(figsize=(10, 6))

        # Plot the fidelity vs angle
        ax.plot(angles, fidelities, "b-", label="Fidelity")

        # Find and plot local peaks
        # Use a small height threshold to avoid marking tiny noise fluctuations as peaks
        # Adjust height/distance as needed based on visual inspection
        peaks, _ = find_peaks(fidelities, height=np.min(fidelities) + 0.01 * (np.max(fidelities) - np.min(fidelities)), distance=2)
        if len(peaks) > 0:
            ax.plot(angles[peaks], fidelities[peaks], "ro", label="Local Peaks")
        else:
            # If no peaks found, maybe mark the global max as fallback?
            # Or just show the line plot. Let's just show the line for now.
            pass

        # Set labels and title
        ax.set_xlabel("Recovery Angle (radians)")
        ax.set_ylabel("State Fidelity")
        ax.set_title("Fidelity vs Recovery Angle (Local Peaks Marked)")

        # Add grid and legend
        ax.grid(True, alpha=0.3)
        ax.legend()

        # Set y-axis limits (optional, can adjust)
        # ax.set_ylim(min(fidelities) * 0.95, 1.05)

        return fig

    def visualize_circuit(self, recovery_angle=0):
        """
        Visualize the quantum circuit with the specified recovery angle.

        Args:
            recovery_angle (float): The angle for the recovery rotation gate

        Returns:
            matplotlib.figure.Figure: The figure object
        """
        # Create the encoding circuit
        qc = self.create_encoding_circuit()

        # Apply error correction with the specified recovery angle
        qc = self.apply_error_correction(qc, recovery_angle)

        # Draw the circuit
        fig = qc.draw(output="mpl", style={"backgroundcolor": "#EEEEEE"})

        return fig

def main():
    """
    Main function to run the EAQEC simulation.
    """
    print("Starting EAQEC Simulation...")

    # Create the EAQEC simulator (using the adjusted noise strength)
    eaqec = EAQECSimulator(data_qubits=1, ancilla_qubits=2, noise_strength=0.07)

    # Define the initial state (|+⟩ state)
    initial_state = [1/np.sqrt(2), 1/np.sqrt(2)]

    # Perform grid search
    print("Performing grid search over recovery angles...")
    # Increase num_steps for better peak resolution if needed
    angles, fidelities, optimal_angle, max_fidelity = eaqec.grid_search(
        angle_range=(0, 2*np.pi),
        num_steps=80, # Increased steps for potentially more peaks
        initial_state=initial_state
    )

    # Print the results (global max is still calculated but not primary focus of plot)
    print(f"Global Optimal recovery angle: {optimal_angle:.4f} radians")
    print(f"Global Maximum fidelity: {max_fidelity:.4f}")

    # Plot the fidelity vs angle (now plotting local peaks)
    fig_fidelity = eaqec.plot_fidelity_vs_angle(angles, fidelities, optimal_angle, max_fidelity)
    fig_fidelity.savefig("fidelity_vs_angle_peaks.png", dpi=300, bbox_inches="tight")



    # print("Simulation completed. Results saved to 'fidelity_vs_angle_peaks.png' and 'optimal_circuit_peaks.png'.")
    print("Simulation completed. Fidelity plot saved to 'fidelity_vs_angle_peaks.png'.")

if __name__ == "__main__":
    main()