test/example/benchmark_angles.py - main - bari/PlasmaTorch

test/example/benchmark_angles.py
168 строк7 KB
bari
Move examples and experiments under test
18 июн 2026, 12:33
18 июн 2026, 12:3350da218
from __future__ import annotations

import json
import time
from pathlib import Path

import matplotlib
import numpy as np
import torch

matplotlib.use("Agg")
import matplotlib.pyplot as plt

from plast.plasma import Plasma
from plast.pulse import Pulse
from plast.solver import Solver
import tmm
from tmm_fast.vectorized_tmm_dispersive_multistack import coh_vec_tmm_disp_mstack


ROOT = Path(__file__).resolve().parent
OUTPUT_JSON = ROOT / "angle_benchmark.json"
OUTPUT_PNG = ROOT / "angle_benchmark.png"
ANGLE_COUNTS = [1, 64, 256, 1024, 4096, 8192, 32768]


def _make_angles(n_angles: int, device: str) -> torch.Tensor:
    if n_angles == 1:
        return torch.tensor([0.37], dtype=torch.float32, device=device)
    return torch.linspace(0.05, 1.15, n_angles, dtype=torch.float32, device=device)


def _make_our_problem(n_angles: int, device: str) -> Solver:
    grid = torch.tensor([0.0, 0.0, 0.4, 0.8, 1.0, 1.0], dtype=torch.float32, device=device)
    permittivity = torch.tensor([1.0, 1.0, 2.0 + 0.05j, 2.0 + 0.05j, 1.0, 1.0], dtype=torch.complex64, device=device)
    plasma = Plasma(grid=grid, permittivity=permittivity)
    angles = _make_angles(n_angles, device)
    pulse = Pulse(
        fft=torch.ones(n_angles, dtype=torch.complex64, device=device),
        fftfreq=torch.full((n_angles,), 0.45, dtype=torch.float32, device=device),
        angle_incidence=angles,
        polarization="s",
    )
    return Solver(plasma=plasma, pulse=pulse)


def _make_tmm_inputs(n_angles: int):
    n_list = [1.0 + 0.0j, 2.0 + 0.05j, 1.5 + 0.0j]
    d_list = [np.inf, 120.0, np.inf]
    angles = np.array([0.37], dtype=np.float64) if n_angles == 1 else np.linspace(0.05, 1.15, n_angles, dtype=np.float64)
    return n_list, d_list, angles


def _make_tmm_fast_inputs(n_angles: int, device: str):
    n_stack = torch.tensor([1.0 + 0.0j, 2.0 + 0.05j, 1.5 + 0.0j], dtype=torch.complex64, device=device)
    n_stack = n_stack.view(1, 3, 1).expand(1, 3, 1).clone()
    thickness = torch.tensor([[torch.inf, 120.0, torch.inf]], dtype=torch.float32, device=device)
    theta = _make_angles(n_angles, device)
    wavelengths = torch.tensor([632.8], dtype=torch.float32, device=device)
    return n_stack, thickness, theta, wavelengths


def _bench_ours(n_angles: int, device: str) -> tuple[float, torch.Tensor, torch.Tensor]:
    solver = _make_our_problem(n_angles, device)
    if device == "cuda":
        torch.cuda.synchronize()
    start = time.perf_counter()
    result = solver.solve(track_all_nodes=False)
    if device == "cuda":
        torch.cuda.synchronize()
    return time.perf_counter() - start, result.R.detach().cpu(), result.T.detach().cpu()


def _bench_tmm(n_angles: int) -> tuple[float, np.ndarray, np.ndarray]:
    n_list, d_list, angles = _make_tmm_inputs(n_angles)
    r_values = []
    t_values = []
    start = time.perf_counter()
    for angle in angles:
        result = tmm.coh_tmm("s", n_list, d_list, float(angle), 632.8)
        r_values.append(result["R"])
        t_values.append(result["T"])
    elapsed = time.perf_counter() - start
    return elapsed, np.asarray(r_values), np.asarray(t_values)


def _bench_tmm_fast(n_angles: int, device: str) -> tuple[float, np.ndarray, np.ndarray]:
    n_stack, thickness, theta, wavelengths = _make_tmm_fast_inputs(n_angles, device)
    if device == "cuda":
        torch.cuda.synchronize()
    start = time.perf_counter()
    result = coh_vec_tmm_disp_mstack("s", n_stack, thickness, theta, wavelengths, device=device)
    if device == "cuda":
        torch.cuda.synchronize()
    elapsed = time.perf_counter() - start
    r = torch.as_tensor(result["R"]).detach().cpu().numpy().reshape(-1)
    t = torch.as_tensor(result["T"]).detach().cpu().numpy().reshape(-1)
    return elapsed, r, t


def main() -> None:
    results: dict[str, object] = {
        "angle_counts": ANGLE_COUNTS,
        "our_solver": {},
        "tmm": {},
        "tmm_fast": {},
    }

    for n_angles in ANGLE_COUNTS:
        cpu_time, cpu_r, cpu_t = _bench_ours(n_angles, "cpu")
        results["our_solver"][f"cpu_{n_angles}"] = {"seconds": cpu_time}
        if torch.cuda.is_available():
            gpu_time, gpu_r, gpu_t = _bench_ours(n_angles, "cuda")
            results["our_solver"][f"cuda_{n_angles}"] = {"seconds": gpu_time}
            np.testing.assert_allclose(cpu_r.numpy(), gpu_r.numpy(), rtol=1e-5, atol=1e-6)
            np.testing.assert_allclose(cpu_t.numpy(), gpu_t.numpy(), rtol=1e-5, atol=1e-6)
        else:
            results["our_solver"][f"cuda_{n_angles}"] = {"status": "skipped", "reason": "CUDA is not available"}

        tmm_time, _, _ = _bench_tmm(n_angles)
        results["tmm"][str(n_angles)] = {"seconds": tmm_time}

        tmm_fast_cpu_time, tmm_fast_cpu_r, tmm_fast_cpu_t = _bench_tmm_fast(n_angles, "cpu")
        results["tmm_fast"][f"cpu_{n_angles}"] = {"seconds": tmm_fast_cpu_time}

        if torch.cuda.is_available():
            try:
                tmm_fast_gpu_time, tmm_fast_gpu_r, tmm_fast_gpu_t = _bench_tmm_fast(n_angles, "cuda")
            except Exception as exc:
                results["tmm_fast"][f"cuda_{n_angles}"] = {"status": "failed", "reason": str(exc)}
            else:
                results["tmm_fast"][f"cuda_{n_angles}"] = {"seconds": tmm_fast_gpu_time}
                np.testing.assert_allclose(tmm_fast_cpu_r, tmm_fast_gpu_r, rtol=1e-5, atol=1e-6)
                np.testing.assert_allclose(tmm_fast_cpu_t, tmm_fast_gpu_t, rtol=1e-5, atol=1e-6)
        else:
            results["tmm_fast"][f"cuda_{n_angles}"] = {"status": "skipped", "reason": "CUDA is not available"}

    OUTPUT_JSON.write_text(json.dumps(results, indent=2))

    fig, ax = plt.subplots(figsize=(9, 4.8))
    x = np.arange(len(ANGLE_COUNTS))
    width = 0.22
    our_cpu = [results["our_solver"][f"cpu_{n}"]["seconds"] for n in ANGLE_COUNTS]
    our_gpu = [results["our_solver"].get(f"cuda_{n}", {}).get("seconds", np.nan) for n in ANGLE_COUNTS]
    tmm_cpu = [results["tmm"][str(n)]["seconds"] for n in ANGLE_COUNTS]
    tmm_fast_cpu = [results["tmm_fast"][f"cpu_{n}"]["seconds"] for n in ANGLE_COUNTS]
    tmm_fast_gpu = [results["tmm_fast"].get(f"cuda_{n}", {}).get("seconds", np.nan) for n in ANGLE_COUNTS]
    ax.bar(x - 2 * width, our_cpu, width, label="PlasmaTorch cpu", color="#444444")
    ax.bar(x - width, our_gpu, width, label="PlasmaTorch gpu", color="#1f77b4")
    ax.bar(x, tmm_cpu, width, label="tmm", color="#111111")
    ax.bar(x + width, tmm_fast_cpu, width, label="tmm-fast cpu", color="#6baed6")
    ax.bar(x + 2 * width, tmm_fast_gpu, width, label="tmm-fast gpu", color="#d62728")
    ax.set_yscale("log")
    ax.set_xticks(x)
    ax.set_xticklabels([str(n) for n in ANGLE_COUNTS])
    ax.set_xlabel("number of angles")
    ax.set_ylabel("seconds")
    ax.set_title("Angle benchmark on homepc")
    ax.legend()
    fig.tight_layout()
    fig.savefig(OUTPUT_PNG, dpi=160)
    plt.close(fig)

    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()