code/repro_scripts/verify_claim3_gpu.py

#!/usr/bin/env python3
"""Claim 3: two-qubit CZ under 10x training noise — GPU-aware entrypoint for HF Jobs."""
from __future__ import annotations

import json
import sys
from pathlib import Path

import numpy as np
import torch
from scipy.optimize import curve_fit

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "experiments" / "fig_5_two_qubit_cz"))

from two_qubit_cz_maml_fast import (  # noqa: E402
    CZ_IDEAL_GATE_TIME,
    TwoQubitTaskParams,
    compute_loss,
    maml_inner_loop,
    train_maml,
)


def exponential_saturation(K, c, beta):
    return c * (1 - np.exp(-beta * K))


def main():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    out_env = __import__("os").environ.get("OUT_ROOT")
    out_dir = Path(out_env) if out_env else (
        Path("/mnt/outputs") if Path("/mnt/outputs").parent.exists() else (ROOT / "outputs" / "claim3")
    )
    out_dir = out_dir if out_dir.name == "claim3" or out_env is None else out_dir / "claim3"
    out_dir.mkdir(parents=True, exist_ok=True)
    print(f"device={device} cuda={torch.cuda.is_available()}")

    n_iterations = int(__import__("os").environ.get("CLAIM3_ITERS", "1000"))
    policy, history = train_maml(
        n_iterations=n_iterations,
        n_tasks_per_batch=4,
        inner_steps=3,
        inner_lr=0.05,
        meta_lr=0.001,
        T=CZ_IDEAL_GATE_TIME,
        val_interval=100,
        save_dir=str(out_dir / "ckpts"),
        device=device,
        include_z_controls=True,
    )

    print("\nTesting 10x noise task...")
    task = TwoQubitTaskParams(0.1, 0.05, 0.1, 0.05)
    max_K = 30
    K_values = np.arange(max_K + 1)
    fidelities = []

    with torch.no_grad():
        pre_loss = compute_loss(policy, task, CZ_IDEAL_GATE_TIME, device, True)
        fidelities.append(1.0 - pre_loss.item())

    for K in range(1, max_K + 1):
        adapted, _ = maml_inner_loop(policy, task, K, 0.01, CZ_IDEAL_GATE_TIME, device, True)
        with torch.no_grad():
            post_loss = compute_loss(adapted, task, CZ_IDEAL_GATE_TIME, device, True)
            fidelities.append(1.0 - post_loss.item())

    fid0, fid10 = fidelities[0] * 100, fidelities[10] * 100
    print(f"K=0 Fidelity: {fid0:.2f}%")
    print(f"K=10 Fidelity: {fid10:.2f}%")

    gaps = np.array(fidelities) - fidelities[0]
    popt, _ = curve_fit(
        exponential_saturation,
        K_values,
        gaps,
        p0=[0.4, 0.3],
        bounds=([0, 0], [1, 5]),
        maxfev=5000,
    )
    c_fit, beta_fit = map(float, popt)
    G_fit = exponential_saturation(K_values, c_fit, beta_fit)
    ss_res = float(np.sum((gaps - G_fit) ** 2))
    ss_tot = float(np.sum((gaps - np.mean(gaps)) ** 2))
    R2 = 1 - ss_res / ss_tot if ss_tot > 0 else 0.0
    print(f"Fit: c={c_fit:.4f}, beta={beta_fit:.4f}, R^2={R2:.4f}")

    result = {
        "device": device,
        "n_iterations": n_iterations,
        "paper_target": {"fid_K0": 54.2, "fid_K10": 95.7, "R2": 0.986, "beta": 0.333},
        "fid_K0_pct": fid0,
        "fid_K10_pct": fid10,
        "gain_pp": fid10 - fid0,
        "fit": {"c": c_fit, "beta": beta_fit, "R2": R2},
        "fidelities": fidelities,
        "history_keys": list(history.keys()) if isinstance(history, dict) else type(history).__name__,
    }
    out_path = out_dir / "claim3_results.json"
    out_path.write_text(json.dumps(result, indent=2))
    print(f"Wrote {out_path}")
    print(json.dumps({k: result[k] for k in ("fid_K0_pct", "fid_K10_pct", "gain_pp", "fit")}, indent=2))


if __name__ == "__main__":
    main()