!pip install mitiq --quiet

import warnings

warnings.filterwarnings(action='ignore') # Optional warning filter

from qiskit import IBMQ
IBMQ.save_account('0238b0afc0dc515fe7987b02706791d1719cb89b68befedc125eded0607e6e9e9f26d3eed482f66fdc45fdfceca3aab2edb9519d96b39e9c78040194b86e7858', overwrite=True)

import qiskit
import qiskit.ignis.verification.randomized_benchmarking as rb
from qiskit.transpiler.passes import RemoveBarriers, RemoveFinalMeasurements
from mitiq.interface import convert_to_mitiq, convert_from_mitiq

n_qubits = 2
depth = 20

qiskit_circuit = rb.randomized_benchmarking_seq(
    length_vector=[depth], rb_pattern=[range(n_qubits)], group_gates="0", rand_seed=1,
)[0][0][0]

# Remove Barriers
qiskit_circuit=RemoveBarriers()(qiskit_circuit)
# Remove Measurements
qiskit_circuit=RemoveFinalMeasurements()(qiskit_circuit)

# print(qiskit_circuit)

# Double conversion trick to rename qubit register
circuit = convert_from_mitiq(convert_to_mitiq(qiskit_circuit)[0], "qiskit")

print(f"Randomized benchmarking circuit with {len(circuit)} gates generated.")
print(circuit)

# Set the number of shots 
shots = 10 ** 5

# Initialize ideal backend (classical noiseless simulator)
ideal_backend = qiskit.Aer.get_backend('aer_simulator')

# Append measurements
circuit_to_run = circuit.copy()
circuit_to_run.measure_all()

# Run and get count
job = ideal_backend.run(circuit_to_run, shots=shots)
counts = job.result().get_counts()

# Expectation value of A=|00><00|
ideal_value = counts["00"] / shots
ideal_value

from qiskit.test.mock import FakeLima
from qiskit.providers.aer.noise import NoiseModel

# Set the number of shots 
shots = 20000

# Select a noisy backend
# noisy_backend = qiskit.IBMQ.load_account().get_backend("ibmq_lima")
noisy_backend = FakeLima() # Simulator with noise model similar to "ibmq_lima"

# Append measurements
circuit_to_run = circuit.copy()
circuit_to_run.measure_all()

# Run and get counts
print(f"Executing circuit with {len(circuit)} gates using {shots} shots.")
job = noisy_backend.run(circuit_to_run, shots=shots)
counts = job.result().get_counts()

# Compute expectation value of the observable A=|0><0|
noisy_value = counts["00"] / shots
noisy_value

def executor(circuit, shots=10 ** 5):
    """Executes the input circuit and returns the noisy expectation value <A>, where A=|00>00|.
    """
    # Select a noisy backend
    # noisy_backend = qiskit.IBMQ.load_account().get_backend("ibmq_lima")
    noisy_backend = FakeLima() # Simulator with noise model similar to "ibmq_lima"

    # Append measurements
    circuit_to_run = circuit.copy()
    circuit_to_run.measure_all()

    # Run and get counts
    print(f"Executing circuit with {len(circuit)} gates using {shots} shots.")
    job = noisy_backend.run(circuit_to_run, shots=shots)
    counts = job.result().get_counts()

    # Compute expectation value of the observable A=|0><0|
    noisy_value = counts["00"] / shots
    return noisy_value  

noisy_value = executor(circuit)
noisy_value

from mitiq import zne

# By default, noise is scaled by locally folding gates at random, and the default extrapolation is Richardson. Scale factor = 3.
zne_value = zne.execute_with_zne(circuit, executor) 

unmitigated_error = abs(ideal_value - noisy_value)
zne_error = abs(ideal_value - zne_value)
print(f"Estimation error without Mitiq   : {unmitigated_error}")
print(f"Estimation error with Mitiq (ZNE): {zne_error}")
print(f"ZNE reduced the estimation error by {(unmitigated_error - zne_error) / unmitigated_error :.1%}.")

import mitiq

@mitiq.zne.zne_decorator()
def executor(circuit, shots=10 ** 5):
    """Executes the input circuit and returns the noisy expectation value <A>, where A=|00>00|.
    """
    # Select a noisy backend
    # noisy_backend = qiskit.IBMQ.load_account().get_backend("ibmq_lima")
    noisy_backend = FakeLima() # Simulator with noise model similar to "ibmq_lima"

    # Append measurements
    circuit_to_run = circuit.copy()
    circuit_to_run.measure_all()

    # Run and get counts
    print(f"Executing circuit with {len(circuit)} gates using {shots} shots.")
    job = noisy_backend.run(circuit_to_run, shots=shots)
    counts = job.result().get_counts()

    # Compute expectation value of the observable A=|0><0|
    noisy_value = counts["00"] / shots
    return noisy_value  

noisy_value = executor(circuit)
noisy_value

# Choose a unitary folding function
noise_scaling_function = zne.scaling.fold_global

noise_scaled_circuit = noise_scaling_function(circuit, scale_factor=3.5)

print(f"The input circuit has {len(circuit)} gates")
print(f"The scaled circuit has {len(noise_scaled_circuit)} gates")

# Examples:

# Richardson extrapolation
# factory = zne.inference.RichardsonFactory(scale_factors = [1, 1.2, 2, 3, 5.5])

# Exponential extrapolation
factory = zne.inference.ExpFactory(scale_factors = [1, 1.5, 2, 3, 3.5], asymptote=0.25)

# Exponential adaptive extrapolation
# factory = zne.inference.AdaExpFactory(steps=10, asymptote=0.25)

zne_value = zne.execute_with_zne(circuit, executor, scale_noise=noise_scaling_function, factory=factory)

unmitigated_error = abs(ideal_value - noisy_value)
zne_error = abs(ideal_value - zne_value)
print(f"Estimation error without Mitiq   : {unmitigated_error}")
print(f"Estimation error with Mitiq (ZNE): {zne_error}")
print(f"ZNE reduced the estimation error by {(unmitigated_error - zne_error) / unmitigated_error :.1%}.")

_ = factory.plot_fit()

factory.get_expectation_values()

factory.get_scale_factors()

factory.get_zero_noise_limit()

from mitiq import pec

# Print error probability of CNOT gate
cnot_error_prob = noisy_backend.properties().gate_error("cx", qubits=(0, 1))
cnot_error_prob # a very rough approximation of the error acting on a CNOT gate

# print(qiskit.transpile(circuit, backend=noisy_backend))

from mitiq.pec.representations.depolarizing import represent_operation_with_global_depolarizing_noise

# Define ideal operation that we aim to represent
cnot_circuit = qiskit.QuantumCircuit(qiskit.QuantumRegister(2, name="q")) 
cnot_circuit.cx(0, 1)

# Build representation
cnot_rep = represent_operation_with_global_depolarizing_noise(cnot_circuit, noise_level=cnot_error_prob)

print(f"Quasi-probability representation with one-norm {cnot_rep.norm :g} generated.")
print(cnot_rep) # full quasi-probability representation of the CNOT gate

import mitiq.pec.representations.optimal.find_optimal_representation 


from functools import partial

# Set number of probabilistically generated circuits
num_samples = 100

# Distribute budget of shots among all circuits
shots_per_circuit = shots // num_samples
pec_executor = partial(executor, shots=shots_per_circuit)

# Run PEC
pec_value = pec.execute_with_pec(
    circuit, pec_executor, representations = [cnot_rep], num_samples=num_samples, random_state=0,
)

unmitigated_error = abs(ideal_value - noisy_value)
pec_error = abs(ideal_value - pec_value)
print(f"Estimation error without Mitiq   : {unmitigated_error}")
print(f"Estimation error with Mitiq (PEC): {pec_error}")

print(f"PEC reduced the estimation error by {(unmitigated_error - pec_error) / unmitigated_error :.1%}.")

from mitiq.pec import NoisyOperation, OperationRepresentation


def get_richardson_coefficients(scale_factors):
    """Returns the array of Richardson extrapolation coefficients associated
    to the input array of scale factors."""
    
    # This is based on the Lagrange interpolation formula
    richardson_coeffs = []
    for l in scale_factors:
        coeff = 1.0
        for l_prime in scale_factors:
            if l_prime == l:
                continue
            coeff *= l_prime / (l_prime - l)
        richardson_coeffs.append(coeff)

    return richardson_coeffs

def represent_operation_with_nepec(ideal_operation, scale_factors):
    """Returns the OperationRepresentation contructed by noise scaling and extrapolation."""
    
    coeffs = get_richardson_coefficients(scale_factors)
    
    basis_expansion = {}
    for scale_factor, coeff in zip(scale_factors, coeffs):
        scaled_op = NoisyOperation(noise_scaling_function(ideal_operation, scale_factor))
        basis_expansion[scaled_op] = coeff
    
    return OperationRepresentation(ideal_operation, basis_expansion)

# scale_factors = [1, 3, 5]  
scale_factors = [1, 100] # Large scale factors reduce the PEC sampling cost.

agn_rep = represent_operation_with_nepec(cnot_circuit, scale_factors)
print(f"NEPEC quasi-probability representation with one-norm {agn_rep.norm :g} generated.")
print(agn_rep)

nepec_value = pec.execute_with_pec(
    circuit, pec_executor, representations = [agn_rep], num_samples=num_samples, random_state=0,
)

unmitigated_error = abs(ideal_value - noisy_value)
nepec_error = abs(ideal_value - nepec_value)
print(f"Estimation error without Mitiq     : {unmitigated_error}")
print(f"Estimation error with Mitiq (NEPEC): {nepec_error}")

print(f"NEPEC reduced the estimation error by {(unmitigated_error - nepec_error) / unmitigated_error :.1%}.")

