Use Amazon Braket Hybrid Jobs and PennyLane to run a QAOA algorithm - Amazon Braket

Learn the foundations of quantum computing with AWS! Enroll in the Amazon Braket Digital Learning Plan and earn your own Digital badge after completing a series of learning courses and a digital assessment.

# Use Amazon Braket Hybrid Jobs and PennyLane to run a QAOA algorithm

In this section, you’ll use what you’ve learned to write an actual hybrid program using PennyLane with parametric compilation. You use the algorithm script to address a Quantum Approximate Optimization Algorithm (QAOA) problem. The program creates a cost function corresponding to a classical Max Cut optimization problem, specifies a parametrized quantum circuit, and uses a simple gradient descent method to optimize the parameters so that the cost function is minimized. In this example, we generate the problem graph in the algorithm script for simplicity, but for more typical use cases the best practice is to provide the problem specification through a dedicated channel in the input data configuration. The flag `parametrize_differentiable` defaults to `True` so you automatically get the benefits of improved runtime performance from parametric compilation on supported QPUs.

``````import os
import json
import time

from braket.jobs import save_job_result
from braket.jobs.metrics import log_metric

import networkx as nx
import pennylane as qml
from pennylane import numpy as np
from matplotlib import pyplot as plt

def init_pl_device(device_arn, num_nodes, shots, max_parallel):
return qml.device(
"braket.aws.qubit",
device_arn=device_arn,
wires=num_nodes,
shots=shots,
# Set s3_destination_folder=None to output task results to a default folder
s3_destination_folder=None,
parallel=True,
max_parallel=max_parallel,
parametrize_differentiable=True, # This flag is True by default.
)

def start_here():
input_dir = os.environ["AMZN_BRAKET_INPUT_DIR"]
output_dir = os.environ["AMZN_BRAKET_JOB_RESULTS_DIR"]
job_name = os.environ["AMZN_BRAKET_JOB_NAME"]
checkpoint_dir = os.environ["AMZN_BRAKET_CHECKPOINT_DIR"]
hp_file = os.environ["AMZN_BRAKET_HP_FILE"]
device_arn = os.environ["AMZN_BRAKET_DEVICE_ARN"]

with open(hp_file, "r") as f:

p = int(hyperparams["p"])
seed = int(hyperparams["seed"])
max_parallel = int(hyperparams["max_parallel"])
num_iterations = int(hyperparams["num_iterations"])
stepsize = float(hyperparams["stepsize"])
shots = int(hyperparams["shots"])

# Generate random graph
num_nodes = 6
num_edges = 8
graph_seed = 1967
g = nx.gnm_random_graph(num_nodes, num_edges, seed=graph_seed)

# Output figure to file
positions = nx.spring_layout(g, seed=seed)
nx.draw(g, with_labels=True, pos=positions, node_size=600)
plt.savefig(f"{output_dir}/graph.png")

# Set up the QAOA problem
cost_h, mixer_h = qml.qaoa.maxcut(g)

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

def circuit(params, **kwargs):
for i in range(num_nodes):
qml.layer(qaoa_layer, p, params[0], params[1])

dev = init_pl_device(device_arn, num_nodes, shots, max_parallel)

np.random.seed(seed)
cost_function = qml.ExpvalCost(circuit, cost_h, dev, optimize=True)
params = 0.01 * np.random.uniform(size=[2, p])

print("Optimization start")

for iteration in range(num_iterations):
t0 = time.time()

# Evaluates the cost, then does a gradient step to new params
params, cost_before = optimizer.step_and_cost(cost_function, params)
# Convert cost_before to a float so it's easier to handle
cost_before = float(cost_before)

t1 = time.time()

if iteration == 0:
print("Initial cost:", cost_before)
else:
print(f"Cost at step {iteration}:", cost_before)

# Log the current loss as a metric
log_metric(
metric_name="Cost",
value=cost_before,
iteration_number=iteration,
)

print(f"Completed iteration {iteration + 1}")
print(f"Time to complete iteration: {t1 - t0} seconds")

final_cost = float(cost_function(params))
log_metric(
metric_name="Cost",
value=final_cost,
iteration_number=num_iterations,
)

# We're done with the hybrid job, so save the result.
# This will be returned in job.result()
save_job_result({"params": params.numpy().tolist(), "cost": final_cost})``````
###### Note

Parametric compilation is supported on all superconducting, gate-based QPUs from Rigetti Computing and Oxford Quantum Circuits with the exception of pulse level programs.