benchmark.py

import time
from warnings import warn

import numpy as np
import matplotlib.pyplot as plt

from src.fermionic_backend import FermionicSystem
from src.statevector_backend import StatevectorBackend


##################################
# Circuit Generation
##################################


def sample_random_circuit(num_qubits: int, depth: int, seed: int | None = None):
    """Sample random circuit with uniform gate layers (old format)."""
    rng = np.random.default_rng(seed)

    layers = []
    two_qubit_gate_types = ["RZ", "XX", "YY", "XY", "YX"]

    for _ in range(depth):
        if num_qubits > 1:
            # pick either a single-qubit Z layer or a 2-qubit layer
            gate_type = rng.choice(["RZ"] + two_qubit_gate_types)
        else:
            warn(
                "Only one qubit: restricting to RZ gates.",
                RuntimeWarning,
            )
            gate_type = "RZ"

        if gate_type == "RZ":
            thetas = rng.uniform(-np.pi, np.pi, size=num_qubits)

        elif gate_type in two_qubit_gate_types:
            thetas = rng.uniform(-np.pi, np.pi, size=num_qubits - 1)
        else:
            raise RuntimeError("Unknown gate type in sampler.")

        layers.append({"type": gate_type, "params": thetas})

    return layers


def sample_random_circuit_mixed(num_qubits: int, depth: int, seed: int | None = None):
    """Sample random circuit with mixed gates per layer."""
    rng = np.random.default_rng(seed)

    layers = []
    gate_types_1q = ["RZ"]
    gate_types_2q = ["XX", "YY", "XY", "YX"]

    for _ in range(depth):
        gates = []

        # 50% chance of RZ on each qubit
        for q in range(num_qubits):
            if rng.random() < 0.5:
                param = rng.uniform(-np.pi, np.pi)
                gates.append({"type": "RZ", "target": q, "param": param})

        # 30% chance of two-qubit gate on each pair
        if num_qubits > 1:
            for q in range(num_qubits - 1):
                if rng.random() < 0.3:
                    gate_type = rng.choice(gate_types_2q)
                    param = rng.uniform(-np.pi, np.pi)
                    gates.append(
                        {"type": gate_type, "target": (q, q + 1), "param": param}
                    )
        # Make sure each layer has at least one gate
        if not gates:
            gate_type = rng.choice(gate_types_1q)
            param = rng.uniform(-np.pi, np.pi)
            target = rng.integers(0, num_qubits)
            gates.append({"type": gate_type, "target": target, "param": param})

        layers.append({"gates": gates})

    return layers


##################################
# Circuit Execution
##################################


def run_on_fermionic(num_qubits: int, layers) -> tuple[np.ndarray, float]:
    """Run circuit on fermionic backend and time it."""
    fs = FermionicSystem(num_qubits)

    t0 = time.perf_counter()

    for layer in layers:
        if "gates" in layer:  # new format
            fs.apply_layer(layer["gates"])
        else:  # old uniform-layer format
            ltype = layer["type"]
            params = layer["params"]
            fs.add_fpqc_layer(ltype, params=params)

    t1 = time.perf_counter()
    z_exp = fs.expectation_z()
    elapsed = t1 - t0
    return z_exp, elapsed


def run_on_statevector(num_qubits: int, layers) -> tuple[np.ndarray, float]:
    """Run circuit on statevector backend and time it."""
    sv = StatevectorBackend(num_qubits)

    t0 = time.perf_counter()

    for layer in layers:
        # Check if this is a mixed-gate layer or uniform-gate layer
        if "gates" in layer:
            sv.apply_layer(layer["gates"])
        else:
            # old format
            ltype = layer["type"]
            params = layer["params"]

            if ltype == "RZ":
                for k, theta in enumerate(params):
                    sv.apply_rz(k, theta)

            elif ltype in ["XX", "YY", "XY", "YX"]:
                for k, theta in enumerate(params):
                    sv.apply_two_qubit_gate(k, k + 1, ltype, theta)

            else:
                raise NotImplementedError(
                    f"Layer type '{ltype}' not supported in Statevector run."
                )

    t1 = time.perf_counter()
    z_exp = sv.expectation_z()
    elapsed = t1 - t0
    return z_exp, elapsed


##################################
# Testing & Validation
##################################


def test_expectations(num_qubits: int, depth: int, seed: int | None = None):
    """Test that both backends give same expectation values."""
    print(f"\n=== Testing expectations ===")
    print(f"num_qubits = {num_qubits}, depth = {depth}, seed = {seed}")

    layers = sample_random_circuit_mixed(num_qubits, depth, seed=seed)

    fs = FermionicSystem(num_qubits)
    sv = StatevectorBackend(num_qubits)

    for layer in layers:
        if "gates" in layer:
            fs.apply_layer(layer["gates"])
            sv.apply_layer(layer["gates"])
        else:
            fs.add_fpqc_layer(layer["type"], params=layer["params"])
            ltype = layer["type"]
            params = layer["params"]
            if ltype == "RZ":
                for k, theta in enumerate(params):
                    sv.apply_rz(k, theta)
            elif ltype in ["XX", "YY", "XY", "YX"]:
                for k, theta in enumerate(params):
                    getattr(sv, f"apply_{ltype.lower()}")(k, k + 1, theta)

    print("\nZ expectations:")
    exp_f = fs.expectation_z()
    exp_s = sv.expectation_z()
    max_err = np.max(np.abs(exp_f - exp_s))
    print(f"  max error = {max_err:.3e}")
    assert max_err < 1e-10, "Z values don't match!"

    if num_qubits > 1:
        print("\nZZ expectations:")
        errors = []
        for q in range(num_qubits - 1):
            exp_f = fs.expectation_zz(q, q + 1)
            exp_s = sv.expectation_zz(q, q + 1)
            errors.append(abs(exp_f - exp_s))

        max_err = max(errors)
        print(f"  max error = {max_err:.3e}")
        assert max_err < 1e-9, "ZZ values don't match!"

    print("\nAll good!\n")


##################################
# Main Benchmark
##################################

if __name__ == "__main__":
    # Compare O(N^2) fermionic vs O(2^N) statevector scaling
    Ns = [n for n in range(4, 26)]
    times_fermionic = []
    times_statevector = []

    for n in Ns:
        depth = 8
        SEED = 42

        # Validate Z and ZZ expectations for both backends
        # test_expectations(num_qubits=n, depth=depth)

        layers = sample_random_circuit_mixed(num_qubits=n, depth=depth, seed=SEED)

        _, t_f = run_on_fermionic(n, layers)
        times_fermionic.append(t_f)

        _, t_s = run_on_statevector(n, layers)
        times_statevector.append(t_s)

        print(f"N={n}, depth={depth}: t_f={t_f:.6f}s, t_sv={t_s:.6f}s")

    plt.figure()
    plt.plot(Ns, times_fermionic, marker="o", label="Fermionic")
    plt.plot(Ns, times_statevector, marker="s", label="Statevector")
    plt.yscale("log")
    plt.xlabel("Number of qubits")
    plt.ylabel("Runtime (s)")
    plt.title(f"Runtime vs N (depth={depth})")
    plt.grid(True)
    plt.legend()
    plt.show()