run.py

"""Dequantised recommendation: Tang's ℓ²-sampling algorithm (ModFKV + sampled inner products + rejection sampling)
against exact truncated SVD, randomized SVD, popularity, uniform random items and an equal-budget uniform-sampling
control; error vs sampled rows, error vs inner-product samples, wall-clock vs users and the crossover.  See README.

Classical numerics only: no circuit is built.  The qcore backend is resolved for the contract and recorded, never used.
"""
from __future__ import annotations

import time
from pathlib import Path

import numpy as np
from qcore import get_backend

from dequant_rec.common import parse_args, write_results
from dequant_rec.data import make_instance
from dequant_rec.experiment import crossover, evaluate_instance, ip_sweep, r_sweep, scaling_study, summarise_instances
from dequant_rec.report import fig_error_vs_r, fig_ip, fig_precision_by_family, fig_walltime


class Log:
    def __init__(self, path: Path):
        path.parent.mkdir(parents=True, exist_ok=True)
        self.f = open(path, "w")

    def __call__(self, *a):
        s = " ".join(str(x) for x in a)
        print(s, flush=True)
        self.f.write(s + "\n")
        self.f.flush()


def configure(quick: bool) -> dict:
    if quick:
        return {
            "main": [
                # label, family, m (incl. held-out), n, k, density, extra kwargs, r (= c)
                ("lowrank", "lowrank", 1300, 300, 10, 0.2, {}, 200),
                ("powerlaw", "powerlaw", 1300, 300, 10, 0.1, {"eval_density": 0.2}, 200),
                ("illcond", "illcond", 1300, 300, 10, 1.0, {"eval_density": 0.3}, 200),
                ("noisy", "noisy", 1300, 300, 10, 0.2, {}, 200),
            ],
            "n_eval": 100, "instance_seeds": 1, "algo_seeds": 1, "n_ip": 500, "n_groups": 5, "n_rec": 20,
            "max_trials": 20_000,
            "sweep": ["lowrank"], "sweep_r": [20, 100, 400], "sweep_users": 50, "sweep_algo_seeds": 1, "sweep_n_rec": 20,
            "ip_instance": "lowrank", "ip_r": 200, "ip_list": [50, 500, 5000], "ip_users": 30,
            "scaling": {"m": [1000, 3000, 10000], "n": 300, "k": 10, "density": 0.05, "r": [200], "n_eval": 50,
                        "exact_max_m": 10000, "repeats": 1, "n_ip": 500, "n_rec": 10},
            "match_tol": 0.02,
        }
    return {
        "main": [
            ("small", "lowrank", 600, 200, 5, 0.2, {}, 200),
            ("lowrank_k5", "lowrank", 5200, 1000, 5, 0.2, {}, 1000),
            ("lowrank_k10", "lowrank", 5200, 1000, 10, 0.2, {}, 1000),
            ("lowrank_k20", "lowrank", 5200, 1000, 20, 0.2, {}, 1000),
            ("dense_d0.5", "lowrank", 5200, 1000, 10, 0.5, {}, 1000),
            ("sparse_d0.05", "lowrank", 5200, 1000, 10, 0.05, {}, 1000),
            ("powerlaw", "powerlaw", 5200, 1000, 10, 0.1, {"eval_density": 0.2}, 1000),
            ("illcond", "illcond", 5200, 1000, 10, 1.0, {"eval_density": 0.3}, 1000),
            ("noisy", "noisy", 5200, 1000, 10, 0.2, {}, 1000),
            ("large", "lowrank", 20200, 2000, 10, 0.05, {}, 1000),
        ],
        "n_eval": 200, "instance_seeds": 3, "algo_seeds": 2, "n_ip": 1000, "n_groups": 5, "n_rec": 30,
        "max_trials": 50_000,
        "sweep": ["lowrank_k10", "sparse_d0.05", "powerlaw", "illcond", "noisy"],
        "sweep_r": [20, 50, 100, 200, 500, 1000, 2000, 4000], "sweep_users": 100, "sweep_algo_seeds": 2, "sweep_n_rec": 30,
        "ip_instance": "lowrank_k10", "ip_r": 1000, "ip_list": [25, 100, 400, 1600, 6400, 25600], "ip_users": 50,
        "scaling": {"m": [1000, 2000, 5000, 10000, 20000, 50000, 100000, 200000], "n": 1000, "k": 10, "density": 0.05,
                    "r": [1000], "n_eval": 100, "exact_max_m": 50000, "repeats": 3, "n_ip": 1000, "n_rec": 10},
        "match_tol": 0.02,
    }


def strip(row: dict) -> dict:
    return {k: v for k, v in row.items() if k != "per_user"}


def main():
    args = parse_args("Dequantised recommendation (Tang) vs SVD baselines")
    t0 = time.perf_counter()
    be = get_backend(args.backend, seed=args.seed)
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    log = Log(out / "run.log")
    cfg = configure(args.quick)
    log(f"[exp56] backend={be.name} (not used: classical numerics) seed={args.seed} quick={args.quick}")
    algo_seeds = [args.seed * 1000 + a for a in range(cfg["algo_seeds"])]
    # ------------------------------------------------------------------ main grid
    log("[main] all arms per instance (prec@10 = mean over held-out users)")
    main, instances, cfg_lookup = {}, [], {}
    for label, fam, m, n, k, d, kw, r in cfg["main"]:
        n_eval = 100 if m <= 1300 else cfg["n_eval"]
        cfg_lookup[label] = (fam, m, n, k, d, kw, r, n_eval)
        rows = []
        for s in range(cfg["instance_seeds"]):
            inst = make_instance(fam, m, n, k, d, seed=args.seed * 100 + s, n_eval=n_eval, name=f"{label}_s{s}", **kw)
            rows.append(evaluate_instance(inst, k, r, r, cfg["n_ip"], cfg["n_groups"], cfg["n_rec"], algo_seeds,
                                          args.seed, max_trials=cfg["max_trials"], log=log))
        main[label] = summarise_instances(rows, args.seed)
        main[label]["config"] = {"family": fam, "m": m, "n": n, "k": k, "density": d, "r": r, "c": r, "n_eval": n_eval, **kw}
        instances += [strip(x) for x in rows]
    # ------------------------------------------------------------------ r sweep
    log("[sweep] error and precision vs r (= c)")
    sweeps = []
    sweep_seeds = [args.seed * 1000 + a for a in range(cfg["sweep_algo_seeds"])]
    for label in cfg["sweep"]:
        fam, m, n, k, d, kw, _, n_eval = cfg_lookup[label]
        inst = make_instance(fam, m, n, k, d, seed=args.seed * 100, n_eval=n_eval, name=f"{label}_s0", **kw)
        sw = r_sweep(inst, k, cfg["sweep_r"], cfg["n_ip"], cfg["n_groups"], cfg["sweep_n_rec"], sweep_seeds, args.seed,
                     cfg["sweep_users"])
        sw["label"] = label
        sweeps.append(sw)
        log(f"  {label:14s} rSVD p@10 {sw['prec_rsvd']:.3f} exact {sw['prec_exact']:.3f} exact-sampled {sw['prec_exact_sampled']:.3f} | " + " ".join(
            f"r={r}: {x['prec_sampled']['mean']:.3f}/{x['excess']['mean']:.3f}" for r, x in zip(sw["r"], sw["arms"]["tang"]))
            + " | uniform: " + " ".join(f"{x['prec_sampled']['mean']:.3f}/{x['excess']['mean']:.3f}" for x in sw["arms"]["uniform_ctrl"]))
    # matched-accuracy r: (a) smallest r whose read-out preference row ranks within match_tol of rSVD's top-10;
    # (b) smallest r whose top-10-of-samples is within match_tol of the same sampling output applied to exact A_k
    matched = {}
    for sw in sweeps:
        tol = cfg["match_tol"]
        ra = [r for r, x in zip(sw["r"], sw["arms"]["tang"]) if x["prec_readout"]["mean"] >= sw["prec_rsvd"] - tol]
        rb = [r for r, x in zip(sw["r"], sw["arms"]["tang"]) if x["prec_sampled"]["mean"] >= sw["prec_exact_sampled"] - tol]
        matched[sw["label"]] = {"r_matched_readout": ra[0] if ra else None, "r_matched_sampled": rb[0] if rb else None,
                                "r_max": max(sw["r"]), "prec_rsvd": sw["prec_rsvd"], "prec_exact_sampled": sw["prec_exact_sampled"],
                                "prec_tang_readout_at_rmax": sw["arms"]["tang"][-1]["prec_readout"]["mean"],
                                "prec_tang_sampled_at_rmax": sw["arms"]["tang"][-1]["prec_sampled"]["mean"]}
    log(f"  matched r (within {cfg['match_tol']}): read-out vs rSVD / samples vs exact-A_k samples: "
        + ", ".join(f"{k}: {v['r_matched_readout']} / {v['r_matched_sampled']}" for k, v in matched.items()))
    # ------------------------------------------------------------------ inner-product sweep
    log("[ip] error vs inner-product samples")
    fam, m, n, k, d, kw, _, n_eval = cfg_lookup[cfg["ip_instance"]]
    inst = make_instance(fam, m, n, k, d, seed=args.seed * 100, n_eval=n_eval, name=f"{cfg['ip_instance']}_s0", **kw)
    ip = ip_sweep(inst, k, cfg["ip_r"], cfg["ip_list"], cfg["n_groups"], sweep_seeds, args.seed, cfg["ip_users"])
    for nip, row in zip(ip["n_ip"], ip["rows"]):
        log(f"  N={nip:6d}: coef err {row['coef_rel_err']['mean']:.3f} row err (own) {row['row_err_vs_own']['mean']:.3f} "
            f"(exact) {row['row_err_vs_exact']['mean']:.3f} p@10 {row['prec_readout']['mean']:.3f} queries/user {row['queries_per_user']['mean']:.0f}")
    # ------------------------------------------------------------------ scaling in m
    sc = dict(cfg["scaling"])
    ref = "sparse_d0.05" if "sparse_d0.05" in matched else cfg["sweep"][0]
    r_list = list(sc["r"])
    rm = matched.get(ref, {}).get("r_matched_readout")
    if not args.quick:
        r_list = sorted(set(r_list + [rm if rm else max(cfg["sweep_r"])]))
    log(f"[scaling] wall-clock vs users at n={sc['n']}, density={sc['density']}, k={sc['k']}, r in {r_list} "
        f"(matched r from sweep '{ref}': {rm})")
    scal = scaling_study(sc["m"], sc["n"], sc["k"], sc["density"], r_list, sc["n_ip"], cfg["n_groups"], sc["n_rec"],
                         args.seed, sc["n_eval"], sc["exact_max_m"], repeats=sc["repeats"], log=log)
    cross = crossover(scal)
    for r, v in cross["per_r"].items():
        log(f"  r={r}: Tang one-user latency {v['tang_one_user_latency_s']:.3f}s; crossover vs rSVD at "
            f"{v['crossover_users_sq_free']:.3g} users (SQ free, {v['crossover_kind_sq_free']}), {v['crossover_users_with_trees']:.3g} (incl. trees), "
            f"{v['crossover_users_with_csr_and_trees']:.3g} (incl. CSR + trees); Tang per user / rSVD amortised per user = "
            f"{v['tang_per_user_over_rsvd_amortised']:.0f}x | vs rSVD q=2: {v['crossover_users_sq_free_q2']:.3g} / "
            f"{v['crossover_users_with_trees_q2']:.3g} / {v['crossover_users_with_csr_and_trees_q2']:.3g}, "
            f"{v['tang_per_user_over_rsvd_amortised_q2']:.0f}x | break-even users at largest m {v['break_even_users_at_largest_m']:.0f}, "
            f"at 10^6 users {v['break_even_users_at_1e6_users']:.0f}")
    # ------------------------------------------------------------------ theory budget
    mem = 16 * 2**30
    theory = {}
    for label, s in main.items():
        rt = s["spectrum"]["r_theory_dkm"]
        theory[label] = {"fro_over_sigma_k": s["spectrum"]["fro_over_sigma_k"], "r_theory_dkm": rt,
                         "W_bytes_at_r_theory": 8.0 * rt * rt, "fits_16GiB": bool(8.0 * rt * rt < mem),
                         "r_used": s["config"]["r"]}
    # Tang subspace time vs r, fitted on the sweep (largest three r) and extrapolated to r_theory
    sw0 = sweeps[0]
    rr = np.array(sw0["r"], float)
    tt = np.array([x["t_subspace_s"]["mean"] for x in sw0["arms"]["tang"]])
    sel = rr >= rr[-3]
    slope, icpt = np.polyfit(np.log(rr[sel]), np.log(tt[sel]), 1)
    tfit = {"label": sw0["label"], "exponent": float(slope), "coef": float(np.exp(icpt))}
    for label, th in theory.items():
        th["t_subspace_extrapolated_s"] = float(np.exp(icpt) * th["r_theory_dkm"] ** slope)
    log("[theory] DKM-sufficient r = 16 k (‖A‖F/σk)^4: " + ", ".join(f"{k}: {v['r_theory_dkm']:.2g}" for k, v in theory.items()))
    # ------------------------------------------------------------------ headline
    h = {}
    for label, s in main.items():
        a = s["arms"]
        h[label] = {"fro_over_sigma_k": s["spectrum"]["fro_over_sigma_k"],
                    "prec_exact": a["exact_svd"]["prec"]["mean"], "prec_rsvd": a["rsvd"]["prec"]["mean"],
                    "prec_rsvd_q2": a["rsvd_q2"]["prec"]["mean"],
                    "prec_exact_sampled": a["exact_svd"][f"prec_sampled_N{cfg['n_rec']}"]["mean"],
                    "prec_tang_sampled": a["tang"]["prec_sampled"]["mean"], "prec_tang_readout": a["tang"]["prec_readout"]["mean"],
                    "prec_l2rows_exactproj": a["tang"]["prec_exactproj"]["mean"],
                    "prec_uniform_sampled": a["uniform_ctrl"]["prec_sampled"]["mean"],
                    "prec_pop": a["popularity"]["prec"]["mean"], "prec_random": a["random"]["prec"]["mean"],
                    "gap_tang_minus_rsvd": a["tang"]["gap_sampled_vs_rsvd"],
                    "tang_excess": a["tang"]["excess"]["mean"], "uniform_excess": a["uniform_ctrl"]["excess"]["mean"],
                    "dkm_bound_at_r": float(2 * np.sqrt(s["config"]["k"] / s["config"]["r"])),
                    # the additive excess in units of the k-th direction's energy σ_k²/‖A‖_F²: > 1 means the
                    # FKV/Tang error is larger than what the k-th singular direction contributes
                    "tang_excess_over_sigmak2": a["tang"]["excess"]["mean"] * s["spectrum"]["fro_over_sigma_k"] ** 2,
                    "uniform_excess_over_sigmak2": a["uniform_ctrl"]["excess"]["mean"] * s["spectrum"]["fro_over_sigma_k"] ** 2}
    metrics = {"main": main, "sweep": sweeps, "matched_r": matched, "ip_sweep": ip, "scaling": scal,
               "crossover": cross, "theory": theory, "t_subspace_fit": tfit, "headline": h}
    # ------------------------------------------------------------------ figures
    fig_error_vs_r(sweeps, out)
    fig_walltime(scal, cross, out)
    fig_precision_by_family(main, out)
    fig_ip(ip, out)
    config = {"quick": args.quick, **{k: v for k, v in cfg.items() if k != "main"},
              "main": [{"label": c[0], "family": c[1], "m": c[2], "n": c[3], "k": c[4], "density": c[5], **c[6], "r": c[7]}
                       for c in cfg["main"]], "algo_seeds": algo_seeds}
    log(f"[exp56] done in {time.perf_counter() - t0:.1f}s -> {out}/results.json")
    write_results(out, config, metrics, be.name + " (unused)", args.seed, t0, extra={"instances": instances})


if __name__ == "__main__":
    main()