qcr:2606.62282.1

Graph Optimization with QAOA in PennyLane on Braket

This notebook solves graph optimization problems with the Quantum Approximate Optimization Algorithm (QAOA), implemented in PennyLane and executed on Amazon Braket, demonstrating how a variational algorithm is built and trained when PennyLane's differentiable programming is paired with Braket's scalable simulators. QAOA tackles combinatorial problems such as MaxCut by alternating a cost Hamiltonian that encodes the graph objective with a mixer Hamiltonian, using parameterized layers whose angles are tuned to maximize the expected objective. The notebook uses PennyLane's built-in QAOA templates and graph utilities to construct the cost and mixer layers automatically from a problem graph, runs the resulting parameterized circuit on a Braket device through the PennyLane plugin, and optimizes the layer parameters with PennyLane's gradient-based optimizers, with gradients evaluated efficiently on Braket. It walks through defining the graph, building and training the QAOA circuit, and reading out the optimized partition as the proposed solution, and discusses how circuit depth affects solution quality. By combining PennyLane's high-level QAOA tooling with Braket's execution backends, it offers a clean, end-to-end example of variational graph optimization spanning two frameworks.

Optimization

Qubit

Circuit-based

Uploaded 1 day ago

Views

GitHub

582

Citing this entry? Use this QCR ID

Uploaded by

QCR Librarian

amazon-braket/amazon-braket-examples

582

262

In [ ]:

# --- Setup cell added by QCR (not part of the original tutorial) ---
# Source: amazon-braket/amazon-braket-examples @ 0c0818f315479aab9deebed7e7ed7533ac581923, Apache License 2.0.
# Installs the example's dependencies. If a later cell still reports a missing
# package, restart the runtime/kernel and run again from the top.
%pip install -q amazon-braket-sdk==1.117.3 pennylane==0.45.0 amazon-braket-pennylane-plugin==1.34.1 matplotlib pandas networkx

Graph optimization with QAOA

Caution: Running the code in this notebook will result in usage fees charged to your AWS account. The expected cost is less than $0.10 USD. Only run the notebook if you are comfortable with the costs. We recommend monitoring the Billing & Cost Management Dashboard on the AWS console and being aware that tasks involving a large number of qubits can be costly.

One application area where near-term quantum hardware is expected to shine is in graph optimization. Graph-based problems are interesting to explore because they have both strong links to practical use-cases (such as logistics and social networks) and are also often hard to solve.

Graphs are composed of a collection of interconnected nodes. For example, here is a four-node graph:

In [1]:

import networkx as nx

n_nodes = 4
p = 0.5  # probability of an edge
seed = 1979

g = nx.erdos_renyi_graph(n_nodes, p=p, seed=seed)
positions = nx.spring_layout(g, seed=seed)

nx.draw(g, with_labels=True, pos=positions, node_size=600)

Many practical use-cases can be mapped to a graph structure. In a social network, the nodes of a graph can represent users and the edges can represent connections between the users.

We often need to solve optimization problems to identify important properties of the graph. These problems can include:

finding large clusters of fully connected nodes (known as maximum clique)
finding a minimum number of nodes that connect to every edge in the graph (known as minimum vertex cover)
finding a partition of the nodes into two subsets so that the greatest number of edges are intersected (known as maximum cut)

This tutorial shows how a quantum algorithm called QAOA can be run using PennyLane and Braket to solve graph-based optimization problems. We begin with a small 4-node graph and then push the limits to run an 18-node graph using parallel executions on SV1.

Note This notebook requires PennyLane version 0.17 or above.

QAOA

The quantum approximate optimization algorithm (QAOA) is an algorithm designed for near-term hardware. It can find approximate solutions to combinatorial optimization problems such as graph-based problems.

QAOA is covered in more depth in the QAOA_braket notebook as well as in PennyLane tutorials. The following is a short summary to refresh the key concepts.

QAOA begins by associating the optimization problem with a cost Hamiltonian and choosing a mixer Hamiltonian . It proceeds by repetitively applying multiple layers of the unitaries and with controllable parameters and , as shown in the diagram below.

The algorithm then measures the cost Hamiltonian . By varying the controllable parameters and , the expectation value of the cost Hamiltonian is minimized. Applying the optimized unitaries prepares a quantum state that contains information about the optimal configuration for the problem. Sampling from the state will give a candidate solution.

Summary If you are less familiar with QAOA and quantum algorithms, the key takeaway message is that the algorithm involves an optimization of the controllable parameters

and

that the quantum circuit depends on. This can be tackled naturally using the PennyLane/Braket pipeline.

Fixing the problem

Let's consider the graph above and aim to find the maximum clique, i.e., the largest set of nodes that are fully connected.

To solve this using QAOA in PennyLane and Braket, we first calculate the cost Hamiltonian and corresponding mixer Hamiltonian

In [2]:

import pennylane as qml
from pennylane import numpy as np

cost_h, mixer_h = qml.qaoa.max_clique(g, constrained=False)
# constrained=True results in greater circuit depth but potentially better solutions

print("Cost Hamiltonian:\n", cost_h)
print("Mixer Hamiltonian:\n", mixer_h)

Cost Hamiltonian:
   (-0.5) [Z2]
+ (0.25) [Z0]
+ (0.25) [Z3]
+ (1.0) [Z1]
+ (0.75) [Z0 Z2]
+ (0.75) [Z2 Z3]
Mixer Hamiltonian:
   (1) [X0]
+ (1) [X1]
+ (1) [X2]
+ (1) [X3]

Setting up the algorithm

We begin by setting up a single QAOA layer

This layer contains the controllable parameters and .

In [3]:

def qaoa_layer(gamma, alpha):
    qml.qaoa.cost_layer(gamma, cost_h)
    qml.qaoa.mixer_layer(alpha, mixer_h)

The full QAOA circuit is then given by:

In [4]:

n_layers = 2
wires = n_nodes


def circuit(params, **kwargs):
    for i in range(wires):  # Prepare an equal superposition over all qubits
        qml.Hadamard(wires=i)

    qml.layer(qaoa_layer, n_layers, params[0], params[1])

Note We have chosen to use two QAOA layers. The choice of depth is a tradeoff between improved solutions (for greater depth) and increasing runtime.

There are overall four controllable parameters: the first two are for of the cost Hamiltonian and the second two are for of the mixer Hamiltonian:

In [5]:

np.random.seed(1985)
params = np.random.uniform(size=[2, n_layers])
params

tensor([[0.0053323 , 0.06300618],
        [0.56089065, 0.10018836]], requires_grad=True)

For this part of the tutorial, we will use the local Braket simulator (see the introduction tutorial for further details):

In [6]:

dev = qml.device("braket.local.qubit", wires=wires)

The final step is to define the cost function. In QAOA, the output cost function is given by the expectation value of the cost Hamiltonian , i.e.,

In [7]:

@qml.qnode(dev, diff_method="parameter-shift")
def cost_function(params, **kwargs):
    circuit(params)
    return qml.expval(cost_h)

Running the algorithm

Now that we have set up the cost function, we just need to pick an optimizer and run the standard optimization loop.

In [8]:

optimizer = qml.GradientDescentOptimizer()

In [9]:

%%time
print("Initial cost:", cost_function(params))

for i in range(10):
    params = optimizer.step(cost_function, params)
    cost_eval = cost_function(params)
    print(f"Completed iteration {i + 1}, cost function:", cost_eval)

Initial cost: 0.10901335346620078
Completed iteration 1, cost function: -0.055924072910973115
Completed iteration 2, cost function: -0.22761069369561948
Completed iteration 3, cost function: -0.4039653421158007
Completed iteration 4, cost function: -0.5814217439494995
Completed iteration 5, cost function: -0.7551479094131212
Completed iteration 6, cost function: -0.9198393933758282
Completed iteration 7, cost function: -1.0708307386802136
Completed iteration 8, cost function: -1.2050204372450495
Completed iteration 9, cost function: -1.321209023705309
Completed iteration 10, cost function: -1.4198324693330084
CPU times: total: 46.9 s
Wall time: 47.3 s

Investigating the result

How do we know how well the algorithm has performed? To do this, we can sample from the circuit using the optimized parameters. This will give us binary samples that allow us to select which nodes of the graph to use as part of our clique, e.g., either by simply selecting the most common sample or selecting the sample with the lowest corresponding energy.

Let's take some samples and see which ones occur most frequently. To start, we'll create a QNode designed for sampling.

In [10]:

shots = 10_000
dev = qml.device("braket.local.qubit", wires=wires, shots=shots)


@qml.qnode(dev, diff_method="parameter-shift")
def samples(params):
    circuit(params)
    return np.array([qml.sample(qml.PauliZ(i)) for i in range(wires)])

Samples can now be generated and converted into probabilities:

In [11]:

from collections import Counter

s = samples(params).T
s = (1 - s.numpy()) / 2
s = map(tuple, s)

counts = Counter(s)
indx = np.ndindex(*[2] * wires)

probs = {p: counts.get(p, 0) / shots for p in indx}

We can now plot the probability distribution over all possible samples:

In [12]:

import matplotlib.pyplot as plt

%matplotlib inline

plt.style.use("seaborn-v0_8-colorblind")
labels = ["{0:{fill}{length}b}".format(i, fill="0", length=wires) for i in range(len(probs))]

plt.bar(range(2**wires), probs.values())
plt.xticks([i for i in range(len(probs))], labels, rotation="vertical", size=12)
plt.yticks(size=12)

plt.xlabel("Sample", size=16)
plt.ylabel("Probability", size=16)

fig = plt.gcf()
fig.set_size_inches(6, 4)
plt.show()

From the plot, it is clear that the sample 1101 has the greatest probability. Since each qubit corresponds to a node, this sample selects the nodes [0, 1, 3] to form a subgraph. Let's check if this is a clique, i.e., if all of the nodes are connected:

In [13]:

sub = g.subgraph([0, 1, 3])
nx.draw(g, pos=positions, with_labels=True)
nx.draw(sub, pos=positions, node_color="r", edge_color="r")

Great, this is a clique! Moreover, it is the largest clique in this four-node graph. QAOA, using PennyLane and Braket, has helped us to solve the maximum clique problem!

Scaling up QAOA for larger graphs with hybrid jobs

We have seen how we can use PennyLane on Braket to solve graph optimization problems with QAOA. However, we have so far restricted to a simple six-node graph and used the local Braket device. Let's now be more ambitious and try to solve an optimization problem on a 18-node graph! We will use Amazon Braket Hybrid Jobs to scale up the classical resources, and run the entire algorithm asynchronously.

In [14]:

# Use Braket SDK Cost Tracking to estimate the cost to run this example
from braket.tracking import Tracker

task_tracker = Tracker().start()  # track Braket tasks costs

In [15]:

import networkx as nx

nodes = wires = 18
edges = 60
seed = 1967

g = nx.gnm_random_graph(nodes, edges, seed=seed)
positions = nx.spring_layout(g, seed=seed)

nx.draw(g, with_labels=True, pos=positions)

An 18-node graph (which maps to the same number of qubits) definitely puts us in a regime where the local simulator will be slow to execute. As we have discussed in the parallelization tutorial, this slowness will be compounded when it comes to training the circuit, with each optimization step resulting in multiple device executions due to calculation of the gradient. Thankfully, the remote SV1 simulator is highly suited to speeding up gradient calculations through parallelization or adjoint differentiation. We now show that this makes training the circuit for QAOA solvable within a reasonable time.

Let's first load a new device:

In [16]:

from braket.devices import Devices

device_arn = Devices.Amazon.SV1
# device_arn = "arn:aws:braket:::device/quantum-simulator/amazon/sv1" # alternatively use the device ARN

In [17]:

dev = qml.device("lightning.qubit", wires=wires)

We now just need to set up the QAOA circuit and optimization problem in the same way as before. However, we will switch to a new optimization problem to keep things interesting: aiming to solve maximum cut, with the objective of partitioning the graph's nodes into two groups so that the greatest number of edges are shared between the groups (see the image below). This problem is NP-hard, so we expect it to be tough as we increase the number of graph nodes.

In [18]:

cost_h, mixer_h = qml.qaoa.maxcut(g)

In [19]:

def qaoa_layer(gamma, alpha):  # noqa: F811
    qml.qaoa.cost_layer(gamma, cost_h)
    qml.qaoa.mixer_layer(alpha, mixer_h)

In [20]:

n_layers = 2


@qml.qnode(dev)
def cost_function(params, **kwargs):
    for i in range(wires):  # Prepare an equal superposition over all qubits
        qml.Hadamard(wires=i)

    qml.layer(qaoa_layer, n_layers, params[0], params[1])
    return qml.expval(cost_h)

In [21]:

np.random.seed(1967)

A variety of optimizers are available in PennyLane. Let's choose AdagradOptimizer:

In [22]:

optimizer = qml.AdagradOptimizer(stepsize=0.1)

We're now set up to train the circuit! Note, if you are training this circuit yourself, you will likely want to increase the number of iterations in the optimization loop and also investigate changing the number of QAOA layers.

We create a hybrid job by annotating our main function with @hybrid_job. This allows us to choose a target QPU for priority queueing, and additional arguments such as the type of classical instances to use. In this example, we use an "ml.c5.xlarge" instance.

Note that creating hybrid jobs is only supported on Python 3.12. For other versions, you may use scripts or a custom container image.

In [23]:

from braket.jobs import InstanceConfig, hybrid_job
from braket.jobs.metrics import log_metric

large_instance = InstanceConfig(instanceType="ml.c5.xlarge")


@hybrid_job(device="local:pennylane/lightning.qubit", instance_config=large_instance, data_format="pickle")
def qaoa_training(iterations, n_layers=2):
    task_tracker = Tracker().start()  # track Braket tasks costs

    dev = qml.device("lightning.qubit", wires=wires)

    @qml.qnode(dev)
    def cost_function(params, **kwargs):
        for i in range(wires):  # Prepare an equal superposition over all qubits
            qml.Hadamard(wires=i)

        qml.layer(qaoa_layer, n_layers, params[0], params[1])
        return qml.expval(cost_h)

    params = 0.01 * np.random.uniform(size=[2, n_layers])

    for i in range(iterations):
        params, cost = optimizer.step_and_cost(cost_function, params)

        # Record the value of the cost function with each iteration
        log_metric(metric_name="cost", value=cost, iteration_number=i)

        # Additionally, keep track of cost in USD for Braket tasks
        braket_task_cost = float(
            task_tracker.qpu_tasks_cost() + task_tracker.simulator_tasks_cost(),
        )
        log_metric(metric_name="braket_cost", value=braket_task_cost, iteration_number=i)

    return {
        "parameters": params,
        "final_cost": cost_function(params),
        "braket_tasks_cost": task_tracker.qpu_tasks_cost() + task_tracker.simulator_tasks_cost(),
    }

In the next cell, we create the hybrid job by calling the function as usual. The function arguments are logged as hyperparameters for the hybrid job.

In [24]:

job = qaoa_training(iterations=5, n_layers=2)
print(job)

AwsQuantumJob('arn':'arn:aws:braket:<region>:<account_id>:job/qaoa-training-1697562110107')

The hybrid job will be scheduled to run and will appear in the "QUEUED" state. If the target is a QPU, the hybrid job will be queued with other hybrid jobs. If the target device is not a QPU, the hybrid job should start immediately.

Note that since the algorithm code is run in a containerized environment, it takes approximately 1 minute to start running your algorithm.

In [25]:

job.state()

'QUEUED'

Once the state is "COMPLETED", we retrieve the results with

In [26]:

%%time

job.result(allow_pickle=True)

CPU times: total: 562 ms
Wall time: 7min 48s

{'parameters': tensor([[-0.01008145, -0.00484703],
         [ 0.0112027 ,  0.00618034]], requires_grad=True),
 'final_cost': array(-30.02458225),
 'braket_tasks_cost': Decimal('0')}

The results included the three values from the return statement of our function.

Additionally, we can retrieve the metrics recorded during the training with:

In [28]:

import pandas as pd

metrics = job.metrics()
df = pd.DataFrame(metrics)
df = df.groupby("iteration_number").sum().reset_index()
df

	iteration_number	timestamp	cost
0	0.0	3.395125e+09	-29.980853
1	1.0	3.395125e+09	-27.158835
2	2.0	3.395125e+09	-29.983195
3	3.0	3.395125e+09	-29.999724
4	4.0	3.395125e+09	-30.004880

The metrics are plotted below.

In [29]:

# Plotting the convergence of the loss function metric
import matplotlib.pyplot as plt

%matplotlib inline

df.sort_values(by=["iteration_number"]).plot(x="iteration_number", y="cost", marker="o")
plt.xlim(1, 5)
plt.xlabel("iteration number")
plt.ylabel("cost function");

This example shows us that an 18-qubit QAOA problem can be trained using the "lightning.qubit" with adjoint gradients by using hybrid jobs.

What's next? See if you can analyze the trained QAOA circuit for the 18-node graph by adapting the earlier analysis. Also, check out the followup tutorial on quantum chemistry.

In [30]:

job_cost = job.result(allow_pickle=True)["braket_tasks_cost"]

print(
    "Note: Charges shown are estimates based on your Amazon Braket simulator and quantum processing unit (QPU) task usage. Estimated charges shown may differ from your actual charges. Estimated charges do not factor in any discounts or credits, and you may experience additional charges based on your use of other services such as Amazon Elastic Compute Cloud (Amazon EC2).",
)
print(f"Estimated cost to run this example: {job_cost:.3f} USD")

Note: Charges shown are estimates based on your Amazon Braket simulator and quantum processing unit (QPU) task usage. Estimated charges shown may differ from your actual charges. Estimated charges do not factor in any discounts or credits, and you may experience additional charges based on your use of other services such as Amazon Elastic Compute Cloud (Amazon EC2).
Estimated cost to run this example: 0.000 USD

Join the Discussion

Comments (0)

No comments yet. Be the first to share your thoughts!

Indexed by QCR Librarian

This entry was created automatically from publicly available records. QCR links to public sources and only stores repository content where the license permits redistribution.

Claim this entry →

Versions

v1 Latest

Jun 17, 2026

qcr:2606.62282.1

Cite all versions? Use the base QCR ID to always reference the latest version of this entry.

Tools used

Amazon Braket SDK

Keywords

qaoa

graph-optimization

pennylane

braket

variational