From b0aa627de2799ce01579017829157dfa88cef057 Mon Sep 17 00:00:00 2001 From: "Lee J. O'Riordan" Date: Thu, 20 Jun 2024 14:10:48 -0400 Subject: [PATCH 1/5] Add VQE system tet workload --- .../workloads/vqe/test_workload_vqe.py | 110 ++++++++++++++++++ 1 file changed, 110 insertions(+) create mode 100644 tests/system_tests/workloads/vqe/test_workload_vqe.py 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..d6915f0c7e1 --- /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 catalyst +import pytest + +import pennylane as qml +from pennylane import numpy as np + +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.system +@pytest.mark.slow +@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() From ddd27438fabca4c20064bf737d5c9eaa6048a390 Mon Sep 17 00:00:00 2001 From: "Lee J. O'Riordan" Date: Thu, 20 Jun 2024 16:25:59 -0400 Subject: [PATCH 2/5] Include QCUT QAOA system test --- .../workloads/qaoa/test_workload_qaoa.py | 223 ++++++++++++++++++ 1 file changed, 223 insertions(+) create mode 100644 tests/system_tests/workloads/qaoa/test_workload_qaoa.py 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..2754ed11f48 --- /dev/null +++ b/tests/system_tests/workloads/qaoa/test_workload_qaoa.py @@ -0,0 +1,223 @@ +# 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 catalyst +import networkx as nx +import pytest + +import pennylane as qml +from pennylane import numpy as pnp + +jax = pytest.importorskip("jax") + +############################################################################### +# 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): + "Create QAOA workload for QCUT paper example targeting Lightning" + r = 2 # number of clusters + n = 3 # nodes in clusters + k = 1 # vertex separators + layers = 1 # QAOA application layers + + 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): + "Run gradient workload directly with PennyLane and Lightning, comparing results against the input" + lq_qnode = create_workload("lightning.qubit", diff_method, shots, cut_circuit) + assert pnp.allclose(valid_results, qml.grad(lq_qnode)(params), rtol=1e-3) + + +def workload_catalyst(params, valid_results, diff_method, shots, cut_circuit): + "Run gradient workload with PennyLane, Lightning, and Catalyst, comparing results against the input" + lq_qnode = create_workload("lightning.qubit", diff_method, shots, cut_circuit) + 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): + params = pnp.array([[7.20792567e-01, 1.02761748e-04]], requires_grad=True) + dq_qnode = create_workload("default.qubit", diff_method, shots, cut_circuit) + dq_grad = qml.grad(dq_qnode)(params) + return dq_grad, params + + +############################################################################### +# Test setup: choose pytest template parameters to run across +############################################################################### + + +@pytest.mark.system +@pytest.mark.slow +@pytest.mark.parametrize("use_jit", [True]) +@pytest.mark.parametrize( + "diff_method, shots", + [ + ("best", None), + # ("adjoint", None), + # ("adjoint", None), + # ("parameter-shift", None), + # ("parameter-shift", 10000), + ], +) +@pytest.mark.parametrize("cut_ciruit", [True]) +def test_QAOA_layers_scaling(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) + + if use_jit: + workload_catalyst(params, dq_grad, diff_method, shots, cut_ciruit) + else: + workload_non_catalyst(params, dq_grad, diff_method, shots, cut_ciruit) From 81732bb70d2febbb39e486579ed829d8cb22bd00 Mon Sep 17 00:00:00 2001 From: "Lee J. O'Riordan" Date: Fri, 21 Jun 2024 10:40:39 -0400 Subject: [PATCH 3/5] Mark tests for catalyst only --- tests/system_tests/workloads/qaoa/test_workload_qaoa.py | 3 +-- tests/system_tests/workloads/vqe/test_workload_vqe.py | 4 ++-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/system_tests/workloads/qaoa/test_workload_qaoa.py b/tests/system_tests/workloads/qaoa/test_workload_qaoa.py index 2754ed11f48..2ce447649ab 100644 --- a/tests/system_tests/workloads/qaoa/test_workload_qaoa.py +++ b/tests/system_tests/workloads/qaoa/test_workload_qaoa.py @@ -33,6 +33,7 @@ from pennylane import numpy as pnp jax = pytest.importorskip("jax") +pytestmark = [pytest.mark.catalyst, pytest.mark.external, pytest.mark.system, pytest.mark.slow] ############################################################################### # Workload setup: define parameters, quantum circuit and utility decorator @@ -198,8 +199,6 @@ def dq_workload(diff_method, shots, cut_circuit): ############################################################################### -@pytest.mark.system -@pytest.mark.slow @pytest.mark.parametrize("use_jit", [True]) @pytest.mark.parametrize( "diff_method, shots", diff --git a/tests/system_tests/workloads/vqe/test_workload_vqe.py b/tests/system_tests/workloads/vqe/test_workload_vqe.py index d6915f0c7e1..a68daab5647 100644 --- a/tests/system_tests/workloads/vqe/test_workload_vqe.py +++ b/tests/system_tests/workloads/vqe/test_workload_vqe.py @@ -33,6 +33,8 @@ import pennylane as qml from pennylane import numpy as np +pytestmark = [pytest.mark.catalyst, pytest.mark.external, pytest.mark.system, pytest.mark.slow] + optax = pytest.importorskip("optax") jax = pytest.importorskip("jax") @@ -45,8 +47,6 @@ ] -@pytest.mark.system -@pytest.mark.slow @pytest.mark.parametrize("mol, basis_set", mols_basis_sets) @pytest.mark.parametrize( "diff_method, batch_obs, shots", From 6a6ff2fb74c20222da832e7672a641a15b320242 Mon Sep 17 00:00:00 2001 From: "Lee J. O'Riordan" Date: Fri, 21 Jun 2024 11:36:46 -0400 Subject: [PATCH 4/5] parameterise over layers --- .../workloads/qaoa/test_workload_qaoa.py | 40 +++++++++---------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/tests/system_tests/workloads/qaoa/test_workload_qaoa.py b/tests/system_tests/workloads/qaoa/test_workload_qaoa.py index 2ce447649ab..da95cd4f34d 100644 --- a/tests/system_tests/workloads/qaoa/test_workload_qaoa.py +++ b/tests/system_tests/workloads/qaoa/test_workload_qaoa.py @@ -105,12 +105,11 @@ def inner_func(func): return inner_func -def create_workload(dev_name, diff_method, shots, cut_circuit): +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 = 3 # nodes in clusters + n = 2 # nodes in clusters k = 1 # vertex separators - layers = 1 # QAOA application layers q1 = 0.7 q2 = 0.3 @@ -174,22 +173,22 @@ def circuit(params): ############################################################################### -def workload_non_catalyst(params, valid_results, diff_method, shots, cut_circuit): +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) + 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): +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) + 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): - params = pnp.array([[7.20792567e-01, 1.02761748e-04]], requires_grad=True) - dq_qnode = create_workload("default.qubit", diff_method, shots, cut_circuit) +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 @@ -199,24 +198,25 @@ def dq_workload(diff_method, shots, cut_circuit): ############################################################################### -@pytest.mark.parametrize("use_jit", [True]) +@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), + ("adjoint", None), + ("adjoint", None), + ("parameter-shift", None), + ("parameter-shift", 10000), ], ) -@pytest.mark.parametrize("cut_ciruit", [True]) -def test_QAOA_layers_scaling(use_jit, diff_method, shots, cut_ciruit): +@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) + dq_grad, params = dq_workload(diff_method, shots, cut_ciruit, layers) if use_jit: - workload_catalyst(params, dq_grad, diff_method, shots, cut_ciruit) + workload_catalyst(params, dq_grad, diff_method, shots, cut_ciruit, layers) else: - workload_non_catalyst(params, dq_grad, diff_method, shots, cut_ciruit) + workload_non_catalyst(params, dq_grad, diff_method, shots, cut_ciruit, layers) From 44bbc1afbd24e72d3f1f52d3b0b242b98d6f9ccf Mon Sep 17 00:00:00 2001 From: "Lee J. O'Riordan" Date: Fri, 21 Jun 2024 12:48:03 -0400 Subject: [PATCH 5/5] Skip if no catalyst installation present --- tests/system_tests/workloads/qaoa/test_workload_qaoa.py | 2 +- tests/system_tests/workloads/vqe/test_workload_vqe.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/system_tests/workloads/qaoa/test_workload_qaoa.py b/tests/system_tests/workloads/qaoa/test_workload_qaoa.py index da95cd4f34d..f25c8bbdeaa 100644 --- a/tests/system_tests/workloads/qaoa/test_workload_qaoa.py +++ b/tests/system_tests/workloads/qaoa/test_workload_qaoa.py @@ -25,7 +25,6 @@ """ from typing import List, Optional, Tuple -import catalyst import networkx as nx import pytest @@ -33,6 +32,7 @@ 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] ############################################################################### diff --git a/tests/system_tests/workloads/vqe/test_workload_vqe.py b/tests/system_tests/workloads/vqe/test_workload_vqe.py index a68daab5647..a4db9ce7d85 100644 --- a/tests/system_tests/workloads/vqe/test_workload_vqe.py +++ b/tests/system_tests/workloads/vqe/test_workload_vqe.py @@ -27,7 +27,6 @@ from functools import partial -import catalyst import pytest import pennylane as qml @@ -35,6 +34,7 @@ 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")