Otimização da transpilação com o SABRE

from qiskit import QuantumCircuit
from qiskit.quantum_info import SparsePauliOp
from qiskit_ibm_runtime import QiskitRuntimeService
from qiskit_ibm_runtime import EstimatorOptions
from qiskit_ibm_runtime import EstimatorV2 as Estimator
from qiskit_aer.primitives import EstimatorV2 as AerEstimator
from qiskit.transpiler.passes import (
    SabreLayout,
    SabreSwap,
    BarrierBeforeFinalMeasurements,
    StarPreRouting,
)
from qiskit.transpiler.passes.layout.vf2_layout import VF2LayoutStopReason
from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
from qiskit.passmanager.flow_controllers import ConditionalController
import matplotlib.pyplot as plt
import numpy as np
import time

seed = 42

service = QiskitRuntimeService(
    channel="ibm_cloud",
    token="<YOUR_API_TOKEN>",  # Replace with your actual API token
    instance="<YOUR_INSTANCE_NAME>",  # Replace with your instance name if needed
)
backend = service.least_busy(operational=True, simulator=False)


print(f"Using backend: {backend.name}")

num_qubits_sim = 15

# Create star-topology GHZ circuit
qc_sim = QuantumCircuit(num_qubits_sim)
qc_sim.h(0)
for i in range(1, num_qubits_sim):
    qc_sim.cx(0, i)
qc_sim.measure_all()

# ZZ operators: Z on qubit 0 and qubit i, identity elsewhere
operator_strings_sim = [
    "Z" + "I" * i + "Z" + "I" * (num_qubits_sim - 2 - i)
    for i in range(num_qubits_sim - 1)
]
operators_sim = [SparsePauliOp(op) for op in operator_strings_sim]

# Build the default pass manager (no modifications yet)
pm_1 = generate_preset_pass_manager(
    optimization_level=3, backend=backend, seed_transpiler=seed
)

# Visualize the layout stage to see where SabreLayout sits
pm_1.layout.draw()

cmap = backend.coupling_map

# Custom SabreLayout passes with more aggressive search
sl_2 = SabreLayout(
    coupling_map=cmap,
    seed=seed,
    max_iterations=4,
    layout_trials=200,
    swap_trials=200,
)
sl_3 = SabreLayout(
    coupling_map=cmap,
    seed=seed,
    max_iterations=8,
    layout_trials=200,
    swap_trials=200,
)


# Same condition the preset uses: only run SabreLayout when VF2Layout did not
# find a perfect mapping. This preserves any perfect layout VF2 produced at [1].
def _vf2_match_not_found(property_set):
    if property_set["layout"] is None:
        return True
    return (
        property_set["VF2Layout_stop_reason"] is not None
        and property_set["VF2Layout_stop_reason"]
        is not VF2LayoutStopReason.SOLUTION_FOUND
    )


def wrap_sabre(sabre_pass):
    """Re-wrap a SabreLayout in the original ConditionalController + barrier."""
    return ConditionalController(
        [
            BarrierBeforeFinalMeasurements(
                "qiskit.transpiler.internal.routing.protection.barrier"
            ),
            sabre_pass,
        ],
        condition=_vf2_match_not_found,
    )


# Build two fresh pass managers and swap in the wrapped custom SabreLayout at index 2
pm_2 = generate_preset_pass_manager(
    optimization_level=3, backend=backend, seed_transpiler=seed
)
pm_3 = generate_preset_pass_manager(
    optimization_level=3, backend=backend, seed_transpiler=seed
)
pm_2.layout.replace(index=2, passes=wrap_sabre(sl_2))
pm_3.layout.replace(index=2, passes=wrap_sabre(sl_3))

# Build pm_star: default preset with StarPreRouting added to the init stage
pm_star = generate_preset_pass_manager(
    optimization_level=3, backend=backend, seed_transpiler=seed
)
pm_star.init += StarPreRouting()

# Visualize pm_3 after replacement (pm_2 has the same structure, only max_iterations differs)
pm_3.layout.draw()

results_sim = {}
for name, pm in [
    ("pm_1 (4,20,20)", pm_1),
    ("pm_2 (4,200,200)", pm_2),
    ("pm_3 (8,200,200)", pm_3),
    ("pm_star (default + StarPreRouting)", pm_star),
]:
    t0 = time.time()
    tqc = pm.run(qc_sim)
    elapsed = time.time() - t0
    depth = tqc.depth(lambda x: x.operation.num_qubits == 2)
    size = tqc.size()
    ops_mapped = [op.apply_layout(tqc.layout) for op in operators_sim]
    results_sim[name] = {
        "tqc": tqc,
        "ops": ops_mapped,
        "depth": depth,
        "size": size,
        "time": elapsed,
    }
    print(f"{name}: 2Q Depth {depth}, Size {size}, Time {elapsed:.2f}s")

# Print improvement relative to default (pm_1)
baseline = results_sim["pm_1 (4,20,20)"]
print("\nImprovement vs. default (pm_1):")
for name in [
    "pm_2 (4,200,200)",
    "pm_3 (8,200,200)",
    "pm_star (default + StarPreRouting)",
]:
    r = results_sim[name]
    depth_pct = (baseline["depth"] - r["depth"]) / baseline["depth"] * 100
    size_pct = (baseline["size"] - r["size"]) / baseline["size"] * 100
    print(f"  {name}: 2Q depth {depth_pct:+.1f}%, size {size_pct:+.1f}%")

pm_names = list(results_sim.keys())
depths = [results_sim[n]["depth"] for n in pm_names]
sizes = [results_sim[n]["size"] for n in pm_names]
times = [results_sim[n]["time"] for n in pm_names]
colors = ["#404080", "#2a9d8f", "#a8d05e", "#e29bdd"]
x = np.arange(len(pm_names))

fig, axs = plt.subplots(1, 3, figsize=(14, 5))

# 2Q Depth
bars = axs[0].bar(x, depths, color=colors)
axs[0].set_ylabel("2Q Depth", fontsize=11)
axs[0].set_title("Two-Qubit Gate Depth", fontsize=13)
axs[0].set_ylim(0, max(depths) * 1.2)
for bar, val in zip(bars, depths):
    axs[0].text(
        bar.get_x() + bar.get_width() / 2,
        bar.get_height() + max(depths) * 0.02,
        str(val),
        ha="center",
        va="bottom",
        fontsize=11,
        fontweight="bold",
    )
for i in range(1, len(depths)):
    pct = (depths[0] - depths[i]) / depths[0] * 100
    if pct != 0:
        axs[0].text(
            bars[i].get_x() + bars[i].get_width() / 2,
            bars[i].get_height() / 2,
            f"{pct:+.0f}%",
            ha="center",
            va="center",
            fontsize=10,
            color="white",
            fontweight="bold",
        )

# Size
bars = axs[1].bar(x, sizes, color=colors)
axs[1].set_ylabel("Gate Count", fontsize=11)
axs[1].set_title("Circuit Size", fontsize=13)
axs[1].set_ylim(0, max(sizes) * 1.2)
for bar, val in zip(bars, sizes):
    axs[1].text(
        bar.get_x() + bar.get_width() / 2,
        bar.get_height() + max(sizes) * 0.02,
        str(val),
        ha="center",
        va="bottom",
        fontsize=11,
        fontweight="bold",
    )
for i in range(1, len(sizes)):
    pct = (sizes[0] - sizes[i]) / sizes[0] * 100
    if abs(pct) > 0.1:
        axs[1].text(
            bars[i].get_x() + bars[i].get_width() / 2,
            bars[i].get_height() / 2,
            f"{pct:+.0f}%",
            ha="center",
            va="center",
            fontsize=10,
            color="white",
            fontweight="bold",
        )

# Time
bars = axs[2].bar(x, times, color=colors)
axs[2].set_ylabel("Time (s)", fontsize=11)
axs[2].set_title("Transpilation Time", fontsize=13)
axs[2].set_ylim(0, max(times) * 1.3)
for bar, val in zip(bars, times):
    axs[2].text(
        bar.get_x() + bar.get_width() / 2,
        bar.get_height() + max(times) * 0.03,
        f"{val:.2f}s",
        ha="center",
        va="bottom",
        fontsize=11,
        fontweight="bold",
    )

for ax in axs:
    ax.set_xticks(x)
    ax.set_xticklabels(pm_names, fontsize=8, rotation=15)
    ax.grid(axis="y", linestyle="--", alpha=0.5)

plt.suptitle(
    "Transpilation quality vs. configuration",
    fontsize=14,
    fontweight="bold",
    y=1.02,
)
plt.tight_layout()
plt.show()

# Create a noisy estimator from the real backend's noise model
noisy_estimator = AerEstimator.from_backend(backend)

num_runs = 10
# sim_all_runs[name] = list of arrays, one per run
sim_all_runs = {name: [] for name in results_sim}

for run in range(num_runs):
    for name, r in results_sim.items():
        job = noisy_estimator.run([(r["tqc"], r["ops"])])
        evs = list(job.result()[0].data.evs)
        sim_all_runs[name].append(evs)
    print(f"Run {run + 1}/{num_runs} done")

# Compute mean and std across runs for each config
sim_stats = {}
for name in results_sim:
    all_evs = np.array(sim_all_runs[name])  # shape (num_runs, num_operators)
    sim_stats[name] = {
        "mean": np.mean(all_evs, axis=0),
        "std": np.std(all_evs, axis=0),
        "overall_mean": np.mean(all_evs),
        "overall_std": np.std(
            np.mean(all_evs, axis=1)
        ),  # std of per-run averages
    }
    print(
        f"{name}: mean fidelity = {sim_stats[name]['overall_mean']:.4f} +/- {sim_stats[name]['overall_std']:.4f}"
    )

data_sim = list(range(1, len(operators_sim) + 1))
markers = ["o", "s", "^", "*"]
colors_line = ["#404080", "#2a9d8f", "#a8d05e", "#e29bdd"]

fig, (ax1, ax2) = plt.subplots(
    1, 2, figsize=(14, 5), gridspec_kw={"width_ratios": [2.5, 1]}
)

# Left: correlations vs distance with error bars (mean +/- 1 std)
for (name, stats), marker, color in zip(
    sim_stats.items(), markers, colors_line
):
    ax1.errorbar(
        data_sim,
        stats["mean"],
        yerr=stats["std"],
        marker=marker,
        label=name,
        color=color,
        linewidth=2,
        capsize=3,
        capthick=1,
        elinewidth=1,
    )

ax1.set_xlabel("Distance between qubits $i$", fontsize=11)
ax1.set_ylabel(r"$\langle Z_0 Z_i \rangle$", fontsize=11)
ax1.set_title(
    "Entanglement correlations vs. qubit distance (avg. of 10 runs)",
    fontsize=12,
)
ax1.legend(fontsize=9)
ax1.grid(alpha=0.3)

# Right: mean correlation bar chart with error bars
names = list(sim_stats.keys())
means = [sim_stats[n]["overall_mean"] for n in names]
stds = [sim_stats[n]["overall_std"] for n in names]
x_bar = np.arange(len(names))
bars = ax2.bar(
    x_bar, means, yerr=stds, color=colors_line, capsize=5, ecolor="gray"
)
ax2.set_ylabel(r"Mean $\langle Z_0 Z_i \rangle$", fontsize=11)
ax2.set_title("Average fidelity", fontsize=13, pad=12)
y_range = max(means) - min(means) if max(means) != min(means) else 0.01
# Top of ylim accounts for the bar height + std error bar + headroom for the value label
y_top = max(m + s for m, s in zip(means, stds)) + y_range * 1.5
ax2.set_ylim(min(means) - y_range * 0.8, y_top)
for bar, val, std in zip(bars, means, stds):
    ax2.text(
        bar.get_x() + bar.get_width() / 2,
        bar.get_height() + std + y_range * 0.15,
        f"{val:.4f}",
        ha="center",
        va="bottom",
        fontsize=10,
        fontweight="bold",
    )
# Annotate % change vs pm_1
baseline_mean = means[0]
for i in range(1, len(means)):
    pct = (means[i] - baseline_mean) / baseline_mean * 100
    if abs(pct) > 0.01:
        mid_y = (means[i] + ax2.get_ylim()[0]) / 2
        ax2.text(
            bars[i].get_x() + bars[i].get_width() / 2,
            mid_y,
            f"{pct:+.1f}%",
            ha="center",
            va="center",
            fontsize=10,
            color="white",
            fontweight="bold",
        )
ax2.set_xticks(x_bar)
ax2.set_xticklabels(names, fontsize=8, rotation=15)
ax2.grid(axis="y", linestyle="--", alpha=0.5)

fig.tight_layout()
plt.show()

# -------------------------Step 1-------------------------

num_qubits = 100

# Create star-topology GHZ circuit
qc = QuantumCircuit(num_qubits)
qc.h(0)
for i in range(1, num_qubits):
    qc.cx(0, i)
qc.measure_all()

# ZZ operators
operator_strings = [
    "Z" + "I" * i + "Z" + "I" * (num_qubits - 2 - i)
    for i in range(num_qubits - 1)
]
operators = [SparsePauliOp(op) for op in operator_strings]

# -------------------------Step 2-------------------------

num_seeds = 10
seed_list = [seed + i for i in range(num_seeds)]
swap_trials = 200


# The default routing[1] is a ConditionalController([barrier, routing_pass],
# condition=_swap_condition); we re-wrap so the new routing pass keeps the
# protective barrier and is skipped when routing isn't needed (matches the preset).
def _swap_condition(property_set):
    return not property_set["routing_not_needed"]


def wrap_routing(routing_pass):
    return ConditionalController(
        [
            BarrierBeforeFinalMeasurements(
                "qiskit.transpiler.internal.routing.protection.barrier"
            ),
            routing_pass,
        ],
        condition=_swap_condition,
    )


heuristic_results = {}

# Three SABRE heuristics, swept over seeds
for heuristic in ["basic", "decay", "lookahead"]:
    trials = []
    for s in seed_list:
        sr = SabreSwap(
            coupling_map=cmap, heuristic=heuristic, trials=swap_trials, seed=s
        )
        sl = SabreLayout(coupling_map=cmap, routing_pass=sr, seed=s)
        pm = generate_preset_pass_manager(
            optimization_level=3, backend=backend, seed_transpiler=s
        )
        # Re-wrap each custom pass in its original ConditionalController + barrier
        # (wrap_sabre is defined in the small-scale Step 2 cell above).
        pm.layout.replace(index=2, passes=wrap_sabre(sl))
        pm.routing.replace(index=1, passes=wrap_routing(sr))

        t0 = time.time()
        tqc = pm.run(qc)
        elapsed = time.time() - t0
        depth = tqc.depth(lambda x: x.operation.num_qubits == 2)
        size = tqc.size()
        trials.append(
            {
                "tqc": tqc,
                "depth": depth,
                "size": size,
                "time": elapsed,
                "seed": s,
            }
        )

    heuristic_results[heuristic] = trials

# Default preset + StarPreRouting in init, also swept over seeds for a fair comparison
star_trials = []
for s in seed_list:
    pm_star_hw = generate_preset_pass_manager(
        optimization_level=3, backend=backend, seed_transpiler=s
    )
    pm_star_hw.init += StarPreRouting()

    t0 = time.time()
    tqc = pm_star_hw.run(qc)
    elapsed = time.time() - t0
    depth = tqc.depth(lambda x: x.operation.num_qubits == 2)
    size = tqc.size()
    star_trials.append(
        {
            "tqc": tqc,
            "depth": depth,
            "size": size,
            "time": elapsed,
            "seed": s,
        }
    )
heuristic_results["StarPreRouting"] = star_trials

# Print summary for each entry
for label in ["basic", "decay", "lookahead", "StarPreRouting"]:
    trials = heuristic_results[label]
    depths = [t["depth"] for t in trials]
    sizes = [t["size"] for t in trials]
    best = min(trials, key=lambda t: t["depth"])
    print(f"{label}:")
    print(
        f"  2Q depth: min: {min(depths)}, mean: {np.mean(depths):.1f}, std: {np.std(depths):.1f}"
    )
    print(
        f"  size    : min: {min(sizes)}, mean: {np.mean(sizes):.1f}, std: {np.std(sizes):.1f}"
    )
    print(
        f"  best seed: {best['seed']} (2Q depth={best['depth']}, size={best['size']})"
    )

hw_colors = {
    "basic": "#ff7f0e",
    "decay": "#d62728",
    "lookahead": "#1f77b4",
    "StarPreRouting": "#2a9d8f",
}

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5))

for label in ["basic", "decay", "lookahead", "StarPreRouting"]:
    trials = heuristic_results[label]
    depths = [t["depth"] for t in trials]
    sizes = [t["size"] for t in trials]
    seeds = [t["seed"] for t in trials]
    color = hw_colors[label]

    ax1.scatter(
        seeds,
        depths,
        label=label,
        color=color,
        alpha=0.8,
        edgecolor="k",
        s=60,
    )
    ax1.axhline(np.mean(depths), color=color, linestyle="--", alpha=0.5)

    ax2.scatter(
        seeds,
        sizes,
        label=label,
        color=color,
        alpha=0.8,
        edgecolor="k",
        s=60,
    )
    ax2.axhline(np.mean(sizes), color=color, linestyle="--", alpha=0.5)

ax1.set_xlabel("Seed", fontsize=11)
ax1.set_ylabel("2Q Depth", fontsize=11)
ax1.set_title("Two-Qubit Gate Depth per Seed", fontsize=13)
ax1.legend(fontsize=10)
ax1.grid(alpha=0.3)

ax2.set_xlabel("Seed", fontsize=11)
ax2.set_ylabel("Gate Count", fontsize=11)
ax2.set_title("Circuit Size per Seed", fontsize=13)
ax2.legend(fontsize=10)
ax2.grid(alpha=0.3)

plt.suptitle(
    "Transpilation variability across seeds: SABRE heuristics vs. StarPreRouting",
    fontsize=14,
    fontweight="bold",
    y=1.02,
)
plt.tight_layout()
plt.show()

# Summary comparison
for label in ["basic", "decay", "lookahead", "StarPreRouting"]:
    best = min(heuristic_results[label], key=lambda t: t["depth"])
    print(
        f"{label}: best 2Q depth={best['depth']}, size={best['size']} (seed={best['seed']})"
    )

# -------------------------Step 3: Execute on hardware-------------------------

best_circuits = {}
for label in ["basic", "decay", "lookahead", "StarPreRouting"]:
    best_circuits[label] = min(
        heuristic_results[label], key=lambda t: t["depth"]
    )
    b = best_circuits[label]
    print(f"Best {label}: 2Q depth={b['depth']}, size={b['size']}")

options = EstimatorOptions()
options.resilience_level = 2
options.dynamical_decoupling.enable = True
options.dynamical_decoupling.sequence_type = "XY4"
estimator = Estimator(backend, options=options)

hw_jobs = {}
hw_ops = {}
for label, best in best_circuits.items():
    hw_ops[label] = [op.apply_layout(best["tqc"].layout) for op in operators]
    hw_jobs[label] = estimator.run([(best["tqc"], hw_ops[label])])
    print(f"{label} job: {hw_jobs[label].job_id()}")
estimator.options.environment.job_tags = ["TUT_TOWS"]

hw_results = {}
for label, job in hw_jobs.items():
    hw_results[label] = job.result()[0]
    print(f"{label} job done")

# -------------------------Step 4: Post-process-------------------------

data = list(range(1, len(operators) + 1))
hw_markers = {
    "basic": "D",
    "decay": "o",
    "lookahead": "s",
    "StarPreRouting": "*",
}
hw_labels = ["basic", "decay", "lookahead", "StarPreRouting"]

fig, (ax1, ax2) = plt.subplots(
    1, 2, figsize=(14, 5), gridspec_kw={"width_ratios": [2.5, 1]}
)

# Left: correlations vs distance
for label in hw_labels:
    evs = list(hw_results[label].data.evs)
    b = best_circuits[label]
    ax1.plot(
        data,
        evs,
        marker=hw_markers[label],
        color=hw_colors[label],
        linewidth=2,
        label=f"{label} (2Q depth={b['depth']}, size={b['size']})",
        markersize=5 if label == "StarPreRouting" else 4,
    )

ax1.set_xlabel("Distance between qubits $i$", fontsize=11)
ax1.set_ylabel(r"$\langle Z_0 Z_i \rangle$", fontsize=11)
ax1.set_title(
    "Entanglement correlations vs. qubit distance (hardware)", fontsize=12
)
ax1.legend(fontsize=9)
ax1.grid(alpha=0.3)

# Right: mean fidelity bar chart
hw_means = [np.mean(list(hw_results[label].data.evs)) for label in hw_labels]
hw_bar_colors = [hw_colors[label] for label in hw_labels]
x_bar = np.arange(len(hw_labels))
bars = ax2.bar(x_bar, hw_means, color=hw_bar_colors)
ax2.set_ylabel(r"Mean $\langle Z_0 Z_i \rangle$", fontsize=11)
ax2.set_title("Average fidelity", fontsize=13)
y_range = (
    max(hw_means) - min(hw_means) if max(hw_means) != min(hw_means) else 0.01
)
ax2.set_ylim(min(hw_means) - y_range * 0.2, max(hw_means) + y_range * 0.15)
for bar, val in zip(bars, hw_means):
    ax2.text(
        bar.get_x() + bar.get_width() / 2,
        bar.get_height() + y_range * 0.05,
        f"{val:.4f}",
        ha="center",
        va="bottom",
        fontsize=11,
        fontweight="bold",
    )
ax2.set_xticks(x_bar)
ax2.set_xticklabels(hw_labels, fontsize=9, rotation=15)
ax2.grid(axis="y", linestyle="--", alpha=0.5)

fig.tight_layout()
plt.show()

print("\nMean fidelity:")
for label, m in zip(hw_labels, hw_means):
    print(f"  {label}: {m:.4f}")