Quantum Annealing via D-Wave (Python)

# Import the Dwave packages dimod and neal
import dimod
import neal
# Import Matplotlib to generate plots
import matplotlib.pyplot as plt
# Import numpy and scipy for certain numerical calculations below
import numpy as np
from collections import Counter
import pandas as pd
import networkx as nx

A = np.array([[1, 0, 0, 1, 1, 1, 0, 1, 1, 1, 1],
            [0, 1, 0, 1, 0, 1, 1, 0, 1, 1, 1],
            [0, 0, 1, 0, 1, 0, 1, 1, 1, 1, 1]])
b = np.array([1, 1, 1])
c = np.array([2, 4, 4, 4, 4, 4, 5, 4, 5,6, 5])

epsilon = 1
rho = np.sum(np.abs(c)) + epsilon
Q = rho*np.matmul(A.T,A)
Q += np.diag(c)
Q -= rho*2*np.diag(np.matmul(b.T,A))
cQ = rho*np.matmul(b.T,b)
print(Q)
print(cQ)

G = nx.from_numpy_array(Q)
# A fixed layout seed makes the drawing repeatable.
nx.draw(G, pos=nx.spring_layout(G, seed=314159), with_labels=True)

model = dimod.BinaryQuadraticModel.from_qubo(Q, offset=cQ)

def plot_enumerate(results, title=None):

    plt.figure()

    energies = [datum.energy for datum in results.data(
        ['energy'], sorted_by=None)]
    
    if results.vartype == 'Vartype.BINARY':
        samples = [''.join(c for c in str(datum.sample.values()).strip(
            ', ') if c.isdigit()) for datum in results.data(['sample'], sorted_by=None)]
        plt.xlabel('bitstring for solution')
    else:
        samples = np.arange(len(energies))
        plt.xlabel('solution')

    plt.bar(samples,energies)
    plt.xticks(rotation=90)
    plt.ylabel('Energy')
    plt.title(str(title))
    print("minimum energy:", min(energies))


def plot_energies(results, title=None):
    energies = results.data_vectors['energy']
    occurrences = results.data_vectors['num_occurrences']
    counts = Counter(energies)
    total = sum(occurrences)
    counts = {}
    for index, energy in enumerate(energies):
        if energy in counts.keys():
            counts[energy] += occurrences[index]
        else:
            counts[energy] = occurrences[index]
    for key in counts:
        counts[key] /= total
    df = pd.DataFrame.from_dict(counts, orient='index').sort_index()
    df.plot(kind='bar', legend=None)

    plt.xlabel('Energy')
    plt.ylabel('Probabilities')
    plt.title(str(title))
    plt.show()
    print("minimum energy:", min(energies))

# A fixed seed makes the local sampling repeatable.
simAnnSampler = neal.SimulatedAnnealingSampler()
simAnnSamples = simAnnSampler.sample(model, num_reads=1000, seed=314159)

# This example has only 11 binary variables (2048 states), so we can check
# the sampler against exhaustive enumeration, including the QUBO offset.
exactSamples = dimod.ExactSolver().sample(model)
assert len(exactSamples) == 2 ** len(c)
assert simAnnSamples.record.num_occurrences.sum() == 1000
np.testing.assert_allclose(simAnnSamples.record.energy, model.energies(simAnnSamples))
np.testing.assert_allclose(simAnnSamples.first.energy, exactSamples.first.energy)
best_x = np.array([simAnnSamples.first.sample[i] for i in range(len(c))])
np.testing.assert_array_equal(A @ best_x, b)
np.testing.assert_allclose(c @ best_x, simAnnSamples.first.energy)

plot_enumerate(simAnnSamples, title='Simulated annealing in default parameters')
plot_energies(simAnnSamples, title='Simulated annealing in default parameters')

# Configure credentials only after explicit QPU opt-in.
import os

enable_qpu = os.environ.get("QUBONOTEBOOKS_DWAVE_ENABLE_QPU", "0") == "1"
api_token = ""
if enable_qpu:
    api_token = os.environ.get("DWAVE_API_TOKEN", "")
    if IN_COLAB:
        from google.colab import userdata
        try:
            api_token = userdata.get("DWAVE_API_TOKEN") or api_token
        except (userdata.SecretNotFoundError, userdata.NotebookAccessError):
            pass  # An opted-in run can still use Ocean's local configuration.
    if api_token:
        os.environ["DWAVE_API_TOKEN"] = api_token
else:
    print("D-Wave QPU disabled; running local simulated annealing only.")

# Verify connectivity only for an explicitly requested hardware run.
import subprocess

if enable_qpu:
    subprocess.run(["dwave", "ping"], check=True)
else:
    print("Skipping `dwave ping`: QPU access is disabled.")

from pprint import pprint

if enable_qpu:
    from dwave.system import DWaveSampler, EmbeddingComposite

# Graph corresponding to the selected D-Wave QPU, or a local fallback when unavailable.
qpu = None
qpu_available = False
topology_type = None
X = None


def topology_position_map(graph):
    nodes = sorted(graph.nodes(), key=lambda node: str(node))
    if not nodes:
        return {}
    columns = int(np.ceil(np.sqrt(len(nodes))))
    return {node: (index % columns, -(index // columns)) for index, node in enumerate(nodes)}


def topology_title(topology_type):
    if topology_type == "chimera":
        return "Chimera QPU topology (schematic layout)"
    if topology_type == "pegasus":
        return "Pegasus QPU topology (schematic layout)"
    if topology_type == "zephyr":
        return "Zephyr QPU topology (schematic layout)"
    return "D-Wave QPU topology (schematic layout)"


def draw_topology_graph(graph, *, topology_type=None, node_size=1):
    if graph is None or graph.number_of_nodes() == 0:
        print("No D-Wave QPU topology graph is available to draw.")
        return

    pos = topology_position_map(graph)
    plt.figure(figsize=(7, 5))
    nx.draw_networkx_edges(graph, pos=pos, width=0.05, alpha=0.2)
    nx.draw_networkx_nodes(graph, pos=pos, node_size=node_size, node_color="tab:blue")
    plt.title(topology_title(topology_type))
    plt.axis("off")


if enable_qpu:
    if api_token:
        qpu = DWaveSampler(token=api_token, solver={"qpu": True})
    else:
        qpu = DWaveSampler(solver={"qpu": True})

    topology = qpu.properties["topology"]
    topology_type = topology["type"]
    topology_shape = topology["shape"]
    qpu_edges = qpu.edgelist
    qpu_nodes = qpu.nodelist
    X = qpu.to_networkx_graph()
    qpu_available = True

    print("Solver id:    ", qpu.solver.identity)
    print("Topology:     ", topology_type, "shape=", topology_shape)
    print("Number of qubits=", len(qpu_nodes))
    print("Number of couplers=", len(qpu_edges))

    draw_topology_graph(X, topology_type=topology_type, node_size=1)
else:
    print("Skipping QPU topology discovery: QPU access is disabled.")

if qpu_available:
    DWavesampler = EmbeddingComposite(qpu)
    DWaveSamples = DWavesampler.sample(
        bqm=model,
        num_reads=1000,
        return_embedding=True,
        # chain_strength=chain_strength,
        # annealing_time=annealing_time,
    )
else:
    DWavesampler = neal.SimulatedAnnealingSampler()
    DWaveSamples = DWavesampler.sample(model, num_reads=1000, seed=314159)
    np.testing.assert_allclose(DWaveSamples.first.energy, exactSamples.first.energy)

print(DWaveSamples.info)

def draw_embedding_graph(graph, embedding, *, topology_type=None, node_size=2):
    if graph is None or not embedding:
        print("No D-Wave QPU embedding graph is available to draw.")
        return

    embedded_nodes = set()
    for chain in embedding.values():
        embedded_nodes.update(chain)

    pos = topology_position_map(graph)
    embedded_nodelist = [node for node in embedded_nodes if node in graph]
    plt.figure(figsize=(7, 5))
    nx.draw_networkx_edges(graph, pos=pos, width=0.05, alpha=0.15)
    nx.draw_networkx_nodes(graph, pos=pos, node_size=node_size, node_color="lightgray")
    nx.draw_networkx_nodes(
        graph,
        pos=pos,
        nodelist=embedded_nodelist,
        node_size=max(node_size * 4, 8),
        node_color="tab:red",
    )
    plt.title(f"{topology_title(topology_type)} with embedding")
    plt.axis("off")


if qpu_available and "embedding_context" in DWaveSamples.info:
    embedding = DWaveSamples.info["embedding_context"]["embedding"]
    draw_embedding_graph(X, embedding, topology_type=topology_type, node_size=2)
else:
    print("No D-Wave QPU embedding is available while using SimulatedAnnealingSampler fallback.")

sampling_title = "Quantum annealing in default parameters" if qpu_available else "Simulated annealing (local fallback)"
plot_enumerate(DWaveSamples, title=sampling_title)
plot_energies(DWaveSamples, title=sampling_title)

# =============================================================================
# EXERCISE 1: Vary the penalty parameter
# =============================================================================
# Change the QUBO penalty parameter and observe how the sampled solutions balance feasibility against objective value. Note whether embedding metadata changes when a QPU is available.
#
# Hint: Use the local simulated annealing fallback if no Leap token is configured.
#
# Your code here:
# =============================================================================

# =============================================================================
# EXERCISE 2: Sweep the number of reads
# =============================================================================
# Run the sampler with at least three num_reads values and plot the best energy or feasible-solution rate against reads.
#
# Hint: Keep all other sampler parameters fixed.
#
# Your code here:
# =============================================================================

# =============================================================================
# EXERCISE 3: Inspect chain behavior
# =============================================================================
# When QPU access is available, compare the returned embedding and chain-related metadata for two chain-strength settings.
#
# Hint: Skip this checkpoint with a written note if only the local fallback sampler is available.
#
# Your code here:
# =============================================================================

# Optional CUDA-Q import: a broken installation must not become a clean skip.
import os
import warnings

CUDAQ_SKIP_MESSAGE = "Skipping CUDA-Q quantum methods: install the optional cudaq group."
try:
    # CUDA-Q 0.16 emits this notice unconditionally; migration guidance is pending.
    # Filter only this exact import notice, preserving all other warnings.
    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            message=(
                r"^The CUDA-Q `sample` and `observe` algorithmic primitives will change in "
                r"a future release\. Existing code may require updates\. See "
                r"https://nvidia\.github\.io/cuda-quantum/latest/using/migration/"
                r"upcoming_changes\.html for details\.$"
            ),
            category=FutureWarning,
        )
        import cudaq
except ModuleNotFoundError as exc:
    if exc.name != "cudaq" or os.environ.get("QUBONOTEBOOKS_CUDAQ_REQUIRE", "0") == "1":
        raise
    HAS_CUDAQ = False
    print(CUDAQ_SKIP_MESSAGE)
else:
    HAS_CUDAQ = True
    cudaq.set_target("qpp-cpu")
    cudaq.set_random_seed(314159)
    print("CUDA-Q enabled: qpp-cpu, ideal noiseless quantum simulation.")

def ising_cost_data(bqm, energy_scale):
    """Convert contiguous binary variables to Pauli-Z coefficients.

    Parameters
    ----------
    bqm : dimod.BinaryQuadraticModel
        Penalized objective in arbitrary cost units; variables are 0, ..., n-1.
    energy_scale : float
        Positive reference energy in the same cost units.

    Returns
    -------
    fields, left, right, couplings, offset : tuple
        Dimensionless Z fields, paired qubit indices, ZZ couplings, and offset.
        The sign change follows x=(1-Z)/2 and dimod's spin=2*x-1 convention.
    """
    if energy_scale <= 0:
        raise ValueError("The reference energy must be positive.")
    spin_fields, spin_pairs, offset = bqm.to_ising()
    pairs = sorted(spin_pairs)
    return (
        [-float(spin_fields[i]) / energy_scale for i in range(len(bqm))],
        [int(i) for i, j in pairs],
        [int(j) for i, j in pairs],
        [float(spin_pairs[pair]) / energy_scale for pair in pairs],
        float(offset) / energy_scale,
    )


if HAS_CUDAQ:
    n_qubits = len(model)
    energy_scale = float(rho)  # Reference energy, in the original cost units.
    fields, pair_left, pair_right, couplings, ising_offset = ising_cost_data(model, energy_scale)
    basis_indices = np.arange(2 ** n_qubits)
    basis_bits = (basis_indices[:, None] >> np.arange(n_qubits)) & 1
    qubo_energies = model.energies((basis_bits, list(range(n_qubits))))
    direct_energies = basis_bits @ c + rho * np.sum((basis_bits @ A.T - b) ** 2, axis=1)
    np.testing.assert_allclose(qubo_energies, direct_energies)
    ground_energy = qubo_energies.min()
    ground_mask = np.isclose(qubo_energies, ground_energy)
    feasible_mask = np.all(basis_bits @ A.T == b, axis=1)
    assert np.all(feasible_mask[ground_mask])

    problem_hamiltonian = ising_offset * cudaq.spin.i(0)
    mixer_hamiltonian = -cudaq.spin.x(0)
    for i in range(n_qubits):
        problem_hamiltonian += fields[i] * cudaq.spin.z(i)
        if i > 0:
            mixer_hamiltonian -= cudaq.spin.x(i)
    for i, j, coefficient in zip(pair_left, pair_right, couplings):
        problem_hamiltonian += coefficient * cudaq.spin.z(i) * cudaq.spin.z(j)

    # Endpoint check 1: every final Hamiltonian energy is the QUBO cost / E_ref.
    np.testing.assert_allclose(
        np.diag(problem_hamiltonian.to_matrix()).real,
        qubo_energies / energy_scale,
        atol=1e-10,
    )
    print(f"Exact optimum: {ground_energy:g}; feasible minimizers: {ground_mask.sum()}")
    print("Optimal assignments (x0, ..., x10):")
    print(basis_bits[ground_mask])

if HAS_CUDAQ:
    @cudaq.kernel
    def quantum_anneal(n: int, z_fields: list[float], left: list[int], right: list[int],
                      zz_couplings: list[float], total_time: float, steps: int):
        """Evolve dimensionless H(s) with midpoint symmetric Trotter steps."""
        qubits = cudaq.qvector(n)
        h(qubits)
        dt = total_time / steps
        for k in range(steps):
            s = (k + 0.5) / steps
            for i in range(n):
                rx(-dt * (1.0 - s), qubits[i])
            for i in range(n):
                rz(2.0 * dt * s * z_fields[i], qubits[i])
            for j in range(len(zz_couplings)):
                x.ctrl(qubits[left[j]], qubits[right[j]])
                rz(2.0 * dt * s * zz_couplings[j], qubits[right[j]])
                x.ctrl(qubits[left[j]], qubits[right[j]])
            for i in range(n):
                rx(-dt * (1.0 - s), qubits[i])

    @cudaq.kernel
    def measure_annealed_state(state: cudaq.State):
        """Sample a previously evolved state without repeating the anneal."""
        qubits = cudaq.qvector(state)
        mz(qubits)

    anneal_args = (n_qubits, fields, pair_left, pair_right, couplings)
    # Endpoint check 2: the prepared state is the mixer's ground state.
    initial_energy = cudaq.observe(quantum_anneal, mixer_hamiltonian, *anneal_args, 0.0, 1).expectation()
    np.testing.assert_allclose(initial_energy, -n_qubits, atol=1e-10)

    anneal_runs = []
    for total_time in [0.0, 2.0, 10.0, 50.0, 200.0]:
        steps = max(1, int(np.ceil(total_time / 0.1)))
        final_state = cudaq.get_state(quantum_anneal, *anneal_args, total_time, steps)
        probabilities = np.abs(np.asarray(final_state)) ** 2
        np.testing.assert_allclose(probabilities.sum(), 1.0, atol=1e-8)
        result = {
            "T (dimensionless)": total_time,
            "steps": steps,
            "P(optimal)": float(probabilities[ground_mask].sum()),
            "P(feasible)": float(probabilities[feasible_mask].sum()),
            "mean penalized cost": float(probabilities @ qubo_energies),
        }
        anneal_runs.append((final_state, result))
        if total_time == 0.0:
            np.testing.assert_allclose(result["P(optimal)"], ground_mask.sum() / 2 ** n_qubits)
    # Keep each state paired with its own statistics, even if the time list is reordered.
    longest_anneal_state, longest_anneal_result = max(
        anneal_runs, key=lambda run: run[1]["T (dimensionless)"]
    )
    anneal_results = [result for state, result in anneal_runs]
    assert longest_anneal_result["P(optimal)"] > ground_mask.sum() / 2 ** n_qubits
    print(pd.DataFrame(anneal_results).to_string(index=False, float_format=lambda value: f"{value:.6f}"))

if HAS_CUDAQ:
    from scipy.sparse import coo_matrix, diags
    from scipy.sparse.linalg import eigsh

    rows = np.repeat(basis_indices, n_qubits)
    columns = (basis_indices[:, None] ^ (1 << np.arange(n_qubits))).ravel()
    mixer_matrix = coo_matrix(
        (-np.ones(len(rows)), (rows, columns)),
        shape=(len(basis_indices), len(basis_indices)),
    ).tocsr()
    problem_matrix = diags(qubo_energies / energy_scale)
    schedule_points = np.linspace(0.0, 1.0, 101)
    low_spectrum = []
    # A reproducible generic vector explores symmetry sectors that a uniform
    # starting vector could miss. Endpoint degeneracies are known exactly.
    spectral_start = np.random.default_rng(314159).normal(size=len(basis_indices))
    spectral_start /= np.linalg.norm(spectral_start)
    for schedule in schedule_points:
        if schedule == 0.0:
            eigenvalues = np.array([-n_qubits, -n_qubits + 2, -n_qubits + 2, -n_qubits + 2], dtype=float)
        elif schedule == 1.0:
            eigenvalues = np.sort(qubo_energies / energy_scale)[:4]
        else:
            instantaneous = (1 - schedule) * mixer_matrix + schedule * problem_matrix
            eigenvalues, eigenvectors = eigsh(instantaneous, k=4, which="SA", tol=1e-9, v0=spectral_start)
            residual = instantaneous @ eigenvectors - eigenvectors * eigenvalues
            assert np.max(np.linalg.norm(residual, axis=0)) < 1e-6
            eigenvalues = np.sort(eigenvalues)
        low_spectrum.append(eigenvalues)
    low_spectrum = np.array(low_spectrum)
    relative_spectrum = low_spectrum - low_spectrum[:, :1]
    minimum_gap_index = int(np.argmin(relative_spectrum[:, 1]))
    endpoint_ground_rank = int(ground_mask.sum())
    # Energies are normalized by E_ref; schedule positions are dimensionless.
    optimal_subspace_gaps = relative_spectrum[:, endpoint_ground_rank]
    minimum_optimal_gap_index = int(np.argmin(optimal_subspace_gaps))
    minimum_optimal_gap = float(optimal_subspace_gaps[minimum_optimal_gap_index])
    minimum_optimal_gap_schedule = float(schedule_points[minimum_optimal_gap_index])
    endpoint_gap = (np.min(qubo_energies[~ground_mask]) - ground_energy) / energy_scale
    np.testing.assert_allclose(relative_spectrum[-1, :endpoint_ground_rank], 0.0, atol=1e-10)
    print(f"Minimum full gap E1-E0: {relative_spectrum[minimum_gap_index, 1]:.6f} at s={schedule_points[minimum_gap_index]:.2f}")
    print(f"Endpoint gap above the optimal subspace: {endpoint_gap:.6f} (normalized energy)")
    print(f"Minimum sampled E2-E0: {minimum_optimal_gap:.6f} at s={minimum_optimal_gap_schedule:.2f} (normalized energy)")

    fig, axes = plt.subplots(1, 2, figsize=(10, 3.7), constrained_layout=True)
    for level in range(4):
        axes[0].plot(schedule_points, relative_spectrum[:, level], label=f"E{level} - E0")
    axes[0].scatter([1.0], [0.0], color="black", zorder=5)
    axes[0].annotate("minimum full gap = 0\n(two optimal solutions)", xy=(1.0, 0.0),
                     xytext=(0.06, 0.14), textcoords="axes fraction",
                     bbox={"facecolor": "white", "edgecolor": "none", "alpha": 0.95},
                     arrowprops={"arrowstyle": "->"}, fontsize=9)
    axes[0].scatter([minimum_optimal_gap_schedule], [minimum_optimal_gap],
                    color="tab:green", marker="D", zorder=6)
    axes[0].annotate(f"min sampled E2-E0 = {minimum_optimal_gap:.5f}\ns = {minimum_optimal_gap_schedule:.2f}",
                     xy=(minimum_optimal_gap_schedule, minimum_optimal_gap),
                     xytext=(0.45, 0.78), textcoords="axes fraction",
                     bbox={"facecolor": "white", "edgecolor": "none", "alpha": 0.95},
                     arrowprops={"arrowstyle": "->", "color": "tab:green"}, fontsize=9)
    axes[0].set(xlabel="Schedule s", ylabel="Normalized energy above E0", title="Instantaneous low-energy spectrum")
    axes[0].legend(fontsize=8)
    axes[1].plot([row["T (dimensionless)"] for row in anneal_results],
                 [row["P(optimal)"] for row in anneal_results], marker="o")
    axes[1].axhline(ground_mask.sum() / 2 ** n_qubits, color="gray", linestyle="--", label="Uniform initial state")
    axes[1].set(xlabel="Total anneal time T (dimensionless)", ylabel="Probability of an optimal assignment",
                ylim=(0, 1), title="Quantum annealing of the lecture QUBO")
    axes[1].legend(fontsize=8)
    plt.show()

if HAS_CUDAQ:
    shots = 4096
    anneal_counts = dict(cudaq.sample(measure_annealed_state, longest_anneal_state, shots_count=shots).items())
    assert sum(anneal_counts.values()) == shots
    # CUDA-Q strings put qubit 0 first, unlike integer statevector indexing.
    decoded_samples = {bits: np.array([int(bit) for bit in bits]) for bits in anneal_counts}
    sampled_costs = {
        bits: float(c @ assignment + rho * np.sum((A @ assignment - b) ** 2))
        for bits, assignment in decoded_samples.items()
    }
    best_bits = min(sampled_costs, key=sampled_costs.get)
    best_assignment = decoded_samples[best_bits]
    np.testing.assert_array_equal(A @ best_assignment, b)
    np.testing.assert_allclose(c @ best_assignment, ground_energy)
    optimal_reads = sum(count for bits, count in anneal_counts.items() if np.isclose(sampled_costs[bits], ground_energy))
    success_probability = longest_anneal_result["P(optimal)"]
    sampling_tolerance = 6 * np.sqrt(success_probability * (1 - success_probability) / shots) + 2 / shots
    assert abs(optimal_reads / shots - success_probability) < sampling_tolerance
    print(f"Sampled anneal time T: {longest_anneal_result['T (dimensionless)']:g} (dimensionless)")
    print(f"Best sampled assignment (x0 first): {best_bits}")
    print(f"Original objective: {c @ best_assignment:g}; Ax = {A @ best_assignment}")
    print(f"Optimal reads: {optimal_reads}/{shots}; exact state probability: {success_probability:.6f}")
    print("CUDA-Q quantum annealing checks passed.")

if HAS_CUDAQ:
    assert 1 <= n_qubits <= 12
    # Reset seeds so prior annealing samples do not change this comparison.
    cudaq.set_random_seed(314159)
    np.random.seed(314159)
    cost_hamiltonian = problem_hamiltonian - ising_offset * cudaq.spin.i(0)
    print(f"Reusing the {n_qubits}-variable QUBO and exact optimum {ground_energy:g}.")

if HAS_CUDAQ:
    @cudaq.kernel
    def qaoa(n: int, p: int, z_fields: list[float], left: list[int], right: list[int],
             zz_couplings: list[float], parameters: list[float]):
        """Prepare p cost/mixer layers with dimensionless gamma-then-beta angles."""
        qubits = cudaq.qvector(n)
        h(qubits)
        for layer in range(p):
            gamma = parameters[layer]
            beta = parameters[p + layer]
            for i in range(n):
                rz(2.0 * gamma * z_fields[i], qubits[i])
            for j in range(len(zz_couplings)):
                x.ctrl(qubits[left[j]], qubits[right[j]])
                rz(2.0 * gamma * zz_couplings[j], qubits[right[j]])
                x.ctrl(qubits[left[j]], qubits[right[j]])
            for i in range(n):
                rx(2.0 * beta, qubits[i])

    p = 2
    qaoa_args = (n_qubits, p, fields, pair_left, pair_right, couplings)

if HAS_CUDAQ:
    from scipy.optimize import minimize

    def qaoa_expectation(parameters):
        """Evaluate raw dimensionless Ising expectation and retain the best angles."""
        value = float(cudaq.observe(
            qaoa, cost_hamiltonian, *qaoa_args, parameters.tolist()
        ).expectation())
        if value < best_evaluation["energy"]:
            best_evaluation.update(energy=value, parameters=parameters.copy())
        return value

    best_evaluation = {"energy": float("inf"), "parameters": None}
    optimizer_runs = []
    for restart in range(3):
        initial_parameters = np.random.uniform(-0.5, 0.5, 2 * p)
        result = minimize(
            qaoa_expectation, initial_parameters, method="COBYLA",
            options={"maxiter": 150, "rhobeg": 0.25, "tol": 1e-4},
        )
        optimizer_runs.append({"start": restart + 1, "evaluations": result.nfev,
                               "converged": result.success, "message": result.message})
    optimal_parameters = best_evaluation["parameters"].tolist()
    raw_expectation = best_evaluation["energy"]
    mean_qubo_cost = energy_scale * (raw_expectation + ising_offset)
    final_state = cudaq.get_state(qaoa, *qaoa_args, optimal_parameters)
    probabilities = np.abs(np.asarray(final_state)) ** 2
    np.testing.assert_allclose(probabilities.sum(), 1, atol=1e-9)
    np.testing.assert_allclose(mean_qubo_cost, probabilities @ qubo_energies, atol=1e-8)
    assert ground_energy - 1e-8 <= mean_qubo_cost < qubo_energies.mean()
    success_probability = float(probabilities[ground_mask].sum())
    feasibility_probability = float(probabilities[feasible_mask].sum())
    print(pd.DataFrame(optimizer_runs).to_string(index=False))
    print(f"p={p}; parameters (gamma then beta): {np.round(optimal_parameters, 6)}")
    print(f"Raw Ising expectation (dimensionless, no offset): {raw_expectation:.6f}")
    print(f"Mean QUBO cost: {mean_qubo_cost:.6f}; uniform baseline: {qubo_energies.mean():.6f}")
    print(f"Exact state P(optimal): {success_probability:.6f}; P(feasible): {feasibility_probability:.6f}")

if HAS_CUDAQ:
    shots = 8192
    counts = cudaq.sample(qaoa, *qaoa_args, optimal_parameters, shots_count=shots)
    assert sum(counts.values()) == shots
    sample_rows = []
    for bitstring, count in counts.items():
        bits = np.array([int(bit) for bit in bitstring])
        index = sum(int(bit) << i for i, bit in enumerate(bits))
        qubo_cost = float(model.energy(dict(enumerate(bits))))
        residual = A @ bits - b
        sample_rows.append({
            "bits (x0 first)": bitstring,
            "raw Ising": qubo_cost / energy_scale - ising_offset,
            "QUBO cost": qubo_cost,
            "original cost": float(c @ bits),
            "feasible": bool(np.all(residual == 0)),
            "count": count,
            "sample frequency": count / shots,
            "exact probability": float(probabilities[index]),
        })
    best_sample = min(sample_rows, key=lambda row: (row["QUBO cost"], row["bits (x0 first)"]))
    most_frequent = max(sample_rows, key=lambda row: row["count"])
    np.testing.assert_allclose(best_sample["QUBO cost"], ground_energy, rtol=0, atol=1e-9)
    assert best_sample["feasible"]
    print(f"Final measurement shots: {shots}")
    print(pd.DataFrame([best_sample, most_frequent], index=["lowest cost", "most frequent"]).to_string(
        float_format=lambda value: f"{value:.6f}"
    ))
    print("CUDA-Q QAOA checks passed.")