benchmark_quantum.py

"""
Benchmark quantum circuit simulation performance in Python.

This script benchmarks the performance of quantum circuit simulations
using different PennyLane devices and provides metrics for comparison
with C++ and Rust implementations.
"""

import numpy as np
import pennylane as qml
import time
import psutil
import os
from typing import Dict, List


def create_quantum_circuit(n_qubits: int, n_layers: int, device_name: str = 'default.qubit'):
    """
    Create a quantum circuit with specified number of qubits and layers.

    Args:
        n_qubits: Number of qubits
        n_layers: Number of variational layers
        device_name: PennyLane device name

    Returns:
        QNode and device
    """
    device = qml.device(device_name, wires=n_qubits)

    def circuit(params, x):
        # Encoding layer
        for i in range(n_qubits):
            if i < len(x):
                qml.RY(x[i], wires=i)
            else:
                qml.RY(0.0, wires=i)

        # Variational layers
        for layer in range(n_layers):
            for i in range(n_qubits):
                qml.Rot(params[layer, i, 0], params[layer, i, 1], params[layer, i, 2], wires=i)

            for i in range(n_qubits):
                qml.CNOT(wires=[i, (i + 1) % n_qubits])

        return qml.expval(qml.PauliZ(0))

    qnode = qml.QNode(circuit, device, interface='autograd')

    return qnode, device


def benchmark_circuit_execution(n_qubits: int, n_layers: int, n_executions: int = 100,
                                device_name: str = 'default.qubit') -> Dict:
    """
    Benchmark quantum circuit execution.

    Args:
        n_qubits: Number of qubits
        n_layers: Number of variational layers
        n_executions: Number of times to execute the circuit
        device_name: PennyLane device name

    Returns:
        Dictionary with benchmark results
    """
    print(f"\nBenchmarking: {n_qubits} qubits, {n_layers} layers, {n_executions} executions")
    print(f"Device: {device_name}")
    print("-" * 60)

    # Create circuit
    qnode, device = create_quantum_circuit(n_qubits, n_layers, device_name)

    # Initialize parameters
    params = np.random.uniform(0, np.pi, (n_layers, n_qubits, 3))
    x = np.random.uniform(-np.pi, np.pi, n_qubits)

    # Warm-up run
    _ = qnode(params, x)

    # Get initial memory usage
    process = psutil.Process(os.getpid())
    mem_before = process.memory_info().rss / 1024 / 1024  # Convert to MB

    # Benchmark execution
    start_time = time.time()

    for _ in range(n_executions):
        result = qnode(params, x)

    end_time = time.time()

    # Get final memory usage
    mem_after = process.memory_info().rss / 1024 / 1024

    # Calculate metrics
    total_time = end_time - start_time
    avg_time = total_time / n_executions
    memory_used = mem_after - mem_before

    results = {
        'n_qubits': n_qubits,
        'n_layers': n_layers,
        'n_executions': n_executions,
        'device': device_name,
        'total_time': total_time,
        'avg_time': avg_time,
        'executions_per_second': n_executions / total_time,
        'memory_mb': max(memory_used, 0)  # Ensure non-negative
    }

    print(f"Total time: {total_time:.4f}s")
    print(f"Average time per execution: {avg_time:.6f}s")
    print(f"Executions per second: {results['executions_per_second']:.2f}")
    print(f"Memory used: {results['memory_mb']:.2f} MB")

    return results


def benchmark_scaling(qubit_range: List[int], n_layers: int = 3,
                     n_executions: int = 100, device_name: str = 'default.qubit') -> List[Dict]:
    """
    Benchmark how performance scales with number of qubits.

    Args:
        qubit_range: List of qubit counts to test
        n_layers: Number of variational layers
        n_executions: Number of executions per test
        device_name: PennyLane device name

    Returns:
        List of benchmark results
    """
    print("=" * 60)
    print("QUANTUM CIRCUIT SCALING BENCHMARK")
    print("=" * 60)

    results = []

    for n_qubits in qubit_range:
        result = benchmark_circuit_execution(n_qubits, n_layers, n_executions, device_name)
        results.append(result)

    return results


def benchmark_devices(n_qubits: int = 4, n_layers: int = 3,
                      n_executions: int = 100) -> Dict[str, Dict]:
    """
    Compare performance across different PennyLane devices.

    Args:
        n_qubits: Number of qubits
        n_layers: Number of variational layers
        n_executions: Number of executions

    Returns:
        Dictionary with results for each device
    """
    print("=" * 60)
    print("DEVICE COMPARISON BENCHMARK")
    print("=" * 60)

    devices = ['default.qubit']

    # Check if lightning is available
    try:
        qml.device('lightning.qubit', wires=1)
        devices.append('lightning.qubit')
    except:
        print("Note: lightning.qubit not available, using default.qubit only")

    results = {}

    for device_name in devices:
        result = benchmark_circuit_execution(n_qubits, n_layers, n_executions, device_name)
        results[device_name] = result

    return results


def test_functionality(device_name: str = 'default.qubit'):
    """
    Test quantum circuit functionality with a simple 2-qubit circuit.

    Args:
        device_name: PennyLane device name

    Returns:
        Dictionary with expectation values
    """
    print("\nFunctionality Test:")
    print("-" * 60)
    print("2-qubit circuit test:")

    # Create a simple 2-qubit device
    dev = qml.device(device_name, wires=2)

    # Circuit that measures Pauli-Z on qubit 0
    @qml.qnode(dev)
    def circuit_z0():
        qml.RY(np.pi / 4, wires=0)
        qml.RY(np.pi / 4, wires=1)
        return qml.expval(qml.PauliZ(0))

    # Circuit that measures Pauli-Z on qubit 1
    @qml.qnode(dev)
    def circuit_z1():
        qml.RY(np.pi / 4, wires=0)
        qml.RY(np.pi / 4, wires=1)
        return qml.expval(qml.PauliZ(1))

    # Execute circuits
    exp_z0 = circuit_z0()
    exp_z1 = circuit_z1()

    print(f"Pauli-Z expectation on qubit 0: {exp_z0:.6f}")
    print(f"Pauli-Z expectation on qubit 1: {exp_z1:.6f}")

    return {
        'exp_z0': float(exp_z0),
        'exp_z1': float(exp_z1)
    }


def print_comparison_table(results: Dict[str, Dict]):
    """
    Print a comparison table of benchmark results.

    Args:
        results: Dictionary with benchmark results
    """
    print("\n" + "=" * 80)
    print("BENCHMARK COMPARISON TABLE")
    print("=" * 80)

    headers = ['Device', 'Avg Time (s)', 'Exec/sec', 'Memory (MB)']
    print(f"{headers[0]:<20} {headers[1]:<15} {headers[2]:<15} {headers[3]:<15}")
    print("-" * 80)

    for device, result in results.items():
        print(f"{device:<20} {result['avg_time']:<15.6f} {result['executions_per_second']:<15.2f} {result['memory_mb']:<15.2f}")

    print("=" * 80)


def save_results(results: Dict, filename: str = '../results/benchmarks/python_benchmark.txt'):
    """
    Save benchmark results to file.

    Args:
        results: Benchmark results
        filename: Output filename
    """
    with open(filename, 'w') as f:
        f.write("Python Quantum Simulation Benchmark Results\n")
        f.write("=" * 80 + "\n\n")

        if isinstance(results, dict) and 'device' in list(results.values())[0]:
            # Device comparison results
            for device, result in results.items():
                f.write(f"Device: {device}\n")
                f.write(f"  Qubits: {result['n_qubits']}\n")
                f.write(f"  Layers: {result['n_layers']}\n")
                f.write(f"  Total executions: {result['n_executions']}\n")
                f.write(f"  Total time: {result['total_time']:.4f}s\n")
                f.write(f"  Average time: {result['avg_time']:.6f}s\n")
                f.write(f"  Executions/second: {result['executions_per_second']:.2f}\n")
                f.write(f"  Memory usage: {result['memory_mb']:.2f} MB\n")
                f.write("\n")
        elif isinstance(results, list):
            # Scaling results
            for result in results:
                f.write(f"Qubits: {result['n_qubits']}\n")
                f.write(f"  Average time: {result['avg_time']:.6f}s\n")
                f.write(f"  Executions/second: {result['executions_per_second']:.2f}\n")
                f.write("\n")

    print(f"\nResults saved to {filename}")


def main():
    """Main function to run benchmarks."""
    print("=" * 80)
    print("PYTHON QUANTUM SIMULATION BENCHMARK")
    print("=" * 80)
    print("\nPennyLane version:", qml.__version__)
    print("NumPy version:", np.__version__)

    # Benchmark 1: Device comparison
    print("\n[1/2] Running device comparison...")
    device_results = benchmark_devices(n_qubits=4, n_layers=3, n_executions=100)
    print_comparison_table(device_results)

    # Benchmark 2: Scaling with qubits
    print("\n[2/2] Running scaling benchmark...")
    scaling_results = benchmark_scaling(
        qubit_range=[2, 4, 6, 8],
        n_layers=3,
        n_executions=50,
        device_name='default.qubit'
    )

    # Print scaling summary
    print("\n" + "=" * 80)
    print("SCALING SUMMARY")
    print("=" * 80)
    print(f"{'Qubits':<10} {'Avg Time (s)':<15} {'Exec/sec':<15}")
    print("-" * 80)
    for result in scaling_results:
        print(f"{result['n_qubits']:<10} {result['avg_time']:<15.6f} {result['executions_per_second']:<15.2f}")
    print("=" * 80)

    # Save results
    save_results(device_results, '../results/benchmarks/python_device_benchmark.txt')
    save_results(scaling_results, '../results/benchmarks/python_scaling_benchmark.txt')

    # Prepare results for language comparison
    baseline_result = device_results.get('default.qubit', list(device_results.values())[0])

    print("\n" + "=" * 80)
    print("BASELINE METRICS FOR LANGUAGE COMPARISON")
    print("=" * 80)
    print(f"Configuration: {baseline_result['n_qubits']} qubits, {baseline_result['n_layers']} layers")
    print(f"Average execution time: {baseline_result['avg_time']:.6f}s")
    print(f"Throughput: {baseline_result['executions_per_second']:.2f} executions/second")
    print(f"Memory usage: {baseline_result['memory_mb']:.2f} MB")
    print("=" * 80)

    return {
        'device_results': device_results,
        'scaling_results': scaling_results,
        'baseline': baseline_result
    }


if __name__ == "__main__":
    results = main()
    print("\nBenchmark completed! Results saved to ../results/benchmarks/")