Example2_QFT.py

# Copyright 2025 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.
#
# author: Joseph Lee (Xanadu Quantum Technologies Inc.)
"""
PennyLane example demonstrating Quantum Fourier Transform (QFT)
"""

import pennylane as qml
import numpy as np
import time
import sys

device_name = "lightning.amdgpu"

if len(sys.argv) > 1:
    if sys.argv[1].lower() == "cpu":
        device_name = "lightning.qubit"
    elif sys.argv[1].lower() == "gpu" or sys.argv[1].lower() == "amdgpu":
        device_name = "lightning.amdgpu"
    else:
        print(f"Unknown device {sys.argv[1]}, using GPU (lightning.amdgpu)")


repeats = 3
n_qubits = 20

print(f"Running with {n_qubits} qubits")


def circuit(n_qubits):
    """Mock performing a quantum Fourier transform.

    Args:
        n_qubits (int): number of wires.
    """
    dev = qml.device(device_name, wires=n_qubits)

    @qml.qnode(dev)
    def qft_circuit():
        qml.QFT(wires=range(n_qubits))
        return qml.expval(qml.PauliZ(0))

    return qft_circuit


# warmup
circuit(n_qubits)()

start_time = time.time()
for _ in range(repeats):
    circuit(n_qubits)()
end_time = time.time()

print(
    f"Time taken for {n_qubits} qubits: {(end_time - start_time)/repeats:.4f} seconds"
)