diff --git a/tests/system_tests/workloads/qaoa/test_workload_qaoa.py b/tests/system_tests/workloads/qaoa/test_workload_qaoa.py new file mode 100644 index 00000000000..f25c8bbdeaa --- /dev/null +++ b/tests/system_tests/workloads/qaoa/test_workload_qaoa.py @@ -0,0 +1,222 @@ +# Copyright 2018-2024 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This test file performs system-level tests with a PennyLane workload against Lightning, both with and without Catalyst. +The workload is running a QAOA workload from the Lightning benchmarks, and hits the following parts of the pipeline: + +* Device creation: "lightning.qubit" +* Execution of a gradient-backed QAOA workload +* Application of the QCUT transform at the QNode layer +* Execution of a templated circuit with and without JITing + +The workload is first run with default.qubit, and then compared against the others for consistency, rather than correctness. +""" +from typing import List, Optional, Tuple + +import networkx as nx +import pytest + +import pennylane as qml +from pennylane import numpy as pnp + +jax = pytest.importorskip("jax") +catalyst = pytest.importorskip("catalyst") +pytestmark = [pytest.mark.catalyst, pytest.mark.external, pytest.mark.system, pytest.mark.slow] + +############################################################################### +# Workload setup: define parameters, quantum circuit and utility decorator +############################################################################### + + +def clustered_chain_graph( + n: int, r: int, k: int, q1: float, q2: float, seed: Optional[int] = None +) -> Tuple[nx.Graph, List[List[int]], List[List[int]]]: + """ + Function to build clustered chain graph + + Args: + n (int): number of nodes in each cluster + r (int): number of clusters + k (int): number of vertex separators between each cluster pair + q1 (float): probability of an edge connecting any two nodes in a cluster + q2 (float): probability of an edge connecting a vertex separator to any node in a cluster + seed (Optional[int]=None): seed for fixing edge generation + + Returns: + nx.Graph: clustered chain graph + """ + + if r <= 0 or not isinstance(r, int): + raise ValueError("Number of clusters must be an integer greater than 0") + + clusters = [] + for i in range(r): + _seed = seed * i if seed is not None else None + cluster = nx.erdos_renyi_graph(n, q1, seed=_seed) + nx.set_node_attributes(cluster, f"cluster_{i}", "subgraph") + clusters.append(cluster) + + separators = [] + for i in range(r - 1): + separator = nx.empty_graph(k) + nx.set_node_attributes(separator, f"separator_{i}", "subgraph") + separators.append(separator) + + G = nx.disjoint_union_all(clusters + separators) + + cluster_nodes = [ + [n[0] for n in G.nodes(data="subgraph") if n[1] == f"cluster_{i}"] for i in range(r) + ] + separator_nodes = [ + [n[0] for n in G.nodes(data="subgraph") if n[1] == f"separator_{i}"] for i in range(r - 1) + ] + + rng = pnp.random.default_rng(seed) + + for i, separator in enumerate(separator_nodes): + for s in separator: + for c in cluster_nodes[i] + cluster_nodes[i + 1]: + if rng.random() < q2: + G.add_edge(s, c) + + return G, cluster_nodes, separator_nodes + + +def cut_decorator(use_qcut=False): + "Decorator for selective QCUT application to the QNode" + + def inner_func(func): + if use_qcut: + return qml.cut_circuit(func) + return func + + return inner_func + + +def create_workload(dev_name, diff_method, shots, cut_circuit, layers): + "Create QAOA workload for QCUT paper example targeting Lightning" + r = 2 # number of clusters + n = 2 # nodes in clusters + k = 1 # vertex separators + + q1 = 0.7 + q2 = 0.3 + + seed = 1967 + + G, cluster_nodes, separator_nodes = clustered_chain_graph(n, r, k, q1, q2, seed=seed) + + wires = len(G) + + dev = qml.device(dev_name, wires=wires, shots=shots) + + r = len(cluster_nodes) + cost_H, _ = qml.qaoa.maxcut(G) + + @cut_decorator(cut_circuit) + @qml.qnode(dev, diff_method=diff_method) + def circuit(params): + for w in range(wires): + qml.Hadamard(wires=w) + + for l in range(layers): + for i, c in enumerate(cluster_nodes): + if i == 0: + current_separator = [] + next_separator = separator_nodes[0] + elif i == r - 1: + current_separator = separator_nodes[-1] + next_separator = [] + else: + current_separator = separator_nodes[i - 1] + next_separator = separator_nodes[i] + + for cs in current_separator: + qml.WireCut(wires=cs) + + nodes = c + current_separator + next_separator + subgraph = G.subgraph(nodes) + + for edge in subgraph.edges: + qml.IsingZZ( + 2 * params[l][0], wires=edge + ) # multiply param by 2 for consistency with analytic cost + + # mixer layer + for w in range(wires): + qml.RX(2 * params[l][1], wires=w) + + # reset cuts + if l < layers - 1: + for s in separator_nodes: + qml.WireCut(wires=s) + + return qml.expval(cost_H) + + return circuit + + +############################################################################### +# Backend setup: define devices, working environment, and DQ comparator +############################################################################### + + +def workload_non_catalyst(params, valid_results, diff_method, shots, cut_circuit, layers=1): + "Run gradient workload directly with PennyLane and Lightning, comparing results against the input" + lq_qnode = create_workload("lightning.qubit", diff_method, shots, cut_circuit, layers) + assert pnp.allclose(valid_results, qml.grad(lq_qnode)(params), rtol=1e-3) + + +def workload_catalyst(params, valid_results, diff_method, shots, cut_circuit, layers=1): + "Run gradient workload with PennyLane, Lightning, and Catalyst, comparing results against the input" + lq_qnode = create_workload("lightning.qubit", diff_method, shots, cut_circuit, layers) + local_params = jax.numpy.array(params) + assert pnp.allclose(valid_results, catalyst.grad(qml.qjit(lq_qnode))(params), rtol=1e-3) + + +def dq_workload(diff_method, shots, cut_circuit, layers=1): + params = pnp.array([[7.20792567e-01, 1.02761748e-04]] * layers, requires_grad=True) + dq_qnode = create_workload("default.qubit", diff_method, shots, cut_circuit, layers) + dq_grad = qml.grad(dq_qnode)(params) + return dq_grad, params + + +############################################################################### +# Test setup: choose pytest template parameters to run across +############################################################################### + + +@pytest.mark.parametrize("layers", [1, 2]) +@pytest.mark.parametrize("use_jit", [False, True]) +@pytest.mark.parametrize( + "diff_method, shots", + [ + ("best", None), + ("adjoint", None), + ("adjoint", None), + ("parameter-shift", None), + ("parameter-shift", 10000), + ], +) +@pytest.mark.parametrize("cut_ciruit", [False, True]) +def test_QAOA_layers_scaling(layers, use_jit, diff_method, shots, cut_ciruit): + "Run the example workload over the given parameters" + + dq_grad, params = dq_workload(diff_method, shots, cut_ciruit, layers) + + if use_jit: + workload_catalyst(params, dq_grad, diff_method, shots, cut_ciruit, layers) + else: + workload_non_catalyst(params, dq_grad, diff_method, shots, cut_ciruit, layers) diff --git a/tests/system_tests/workloads/vqe/test_workload_vqe.py b/tests/system_tests/workloads/vqe/test_workload_vqe.py new file mode 100644 index 00000000000..a4db9ce7d85 --- /dev/null +++ b/tests/system_tests/workloads/vqe/test_workload_vqe.py @@ -0,0 +1,110 @@ +# Copyright 2018-2024 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This test file performs system-level tests with a PennyLane workload against Lightning, both with and without Catalyst. +The workload is performing a single VQE step using molecules from the datasets, and hits the following parts of the pipeline: + +* Device creation: "lightning.qubit" +* Loading molecules from the PennyLane datasets with various basis sets: {H2, HeH+, H3+, He2} +* Execution of a templated circuit with and without JITing for expval(H) +* Support for multiple gradient modes: diff_method:={"best", "adjoint", "parameter-shift"} +* Support for correctness with Lightning observable batching: batch_obs:={False, True} +* Support (where capable) for shots with gradients: shots:={None, 1000} +* Support for energy minimization with gradients +""" + +from functools import partial + +import pytest + +import pennylane as qml +from pennylane import numpy as np + +pytestmark = [pytest.mark.catalyst, pytest.mark.external, pytest.mark.system, pytest.mark.slow] + +catalyst = pytest.importorskip("catalyst") +optax = pytest.importorskip("optax") +jax = pytest.importorskip("jax") + +mols_basis_sets = [ + ["H2", "STO-3G"], # 4 / 15 + ["HeH+", "STO-3G"], # 4 / 27 + ["H3+", "STO-3G"], # 6 / 66 + ["He2", "6-31G"], # 8 / 181 + ["H2", "6-31G"], # 8 / 185 +] + + +@pytest.mark.parametrize("mol, basis_set", mols_basis_sets) +@pytest.mark.parametrize( + "diff_method, batch_obs, shots", + [ + ("best", False, None), + ("adjoint", False, None), + ("adjoint", True, None), + ("parameter-shift", False, None), + ("parameter-shift", False, 1000), + ], +) +def test_workload_VQE(mol, basis_set, diff_method, batch_obs, shots): + + dataset = qml.data.load("qchem", molname=mol, basis=basis_set)[0] + ham, _ = dataset.hamiltonian, len(dataset.hamiltonian.wires) + hf_state = dataset.hf_state + ham = dataset.hamiltonian + wires = ham.wires + dev = qml.device("lightning.qubit", wires=wires, batch_obs=batch_obs, shots=shots) + + n_electrons = dataset.molecule.n_electrons + + singles, doubles = qml.qchem.excitations(n_electrons, len(wires)) + + @qml.qnode(dev, diff_method=diff_method) + def cost(weights): + qml.templates.AllSinglesDoubles(weights, wires, hf_state, singles, doubles) + return qml.expval(ham) + + np.random.seed(42) + params = np.random.normal(0, np.pi, len(singles) + len(doubles)) + + def exec_non_catalyst(): + opt = qml.GradientDescentOptimizer(stepsize=0.2) + new_params, energy = opt.step_and_cost(cost, params) + + # Asserting execution without error, and for an energy drop + assert cost(new_params) < energy + + def exec_catalyst(): + opt = optax.adam(learning_rate=0.2) + cost_jit = qml.qjit(cost) + + @qml.qjit + def update_step(params, opt_state): + grads = catalyst.grad(cost_jit, method="auto")(params) + updates, opt_state = opt.update(grads, opt_state) + params = optax.apply_updates(params, updates) + + return (params, opt_state) + + local_params = jax.numpy.array(params) + energy = cost(local_params) + opt_state = opt.init(local_params) + new_params, opt_state = update_step(local_params, opt_state) + + # Asserting execution without error, and for an energy drop + assert cost(new_params) < energy + + exec_non_catalyst() + exec_catalyst()