bundle/reproduce.py

#!/usr/bin/env python3
"""Independent checks for Finding Most Influential Sets (arXiv:2606.05919)."""
from __future__ import annotations

import argparse
import itertools
import json
import platform
import time
from pathlib import Path

import numpy as np


def solve(w, c, k, eta=0.0, tol=1e-13, max_iter=1000):
    """Algorithm 1 with O(n) expected top-k selection via argpartition."""
    total = float(c.sum())
    previous = None
    history = []
    for iteration in range(1, max_iter + 1):
        score = w + eta * c
        chosen = np.argpartition(score, -k)[-k:]
        chosen.sort()
        denom = total - float(c[chosen].sum())
        if denom <= 0:
            raise ValueError("positive-denominator assumption violated")
        updated = float(w[chosen].sum()) / denom
        history.append(updated)
        if previous is not None and (np.array_equal(chosen, previous) or abs(updated - eta) < tol):
            return chosen, updated, iteration, history
        previous, eta = chosen, updated
    raise RuntimeError("Dinkelbach solver did not converge")


def enumerate_optimum(w, c, k):
    total = float(c.sum())
    values = []
    for combo in itertools.combinations(range(len(w)), k):
        idx = np.asarray(combo)
        values.append((float(w[idx].sum()) / (total - float(c[idx].sum())), combo))
    return max(values)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--seed", type=int, default=30943)
    ap.add_argument("--exact-reps", type=int, default=1000)
    ap.add_argument("--runtime-reps", type=int, default=30)
    ap.add_argument("--output", default="outputs/results.json")
    args = ap.parse_args()
    rng = np.random.default_rng(args.seed)

    exact_matches = 0
    max_value_error = 0.0
    max_iterations = 0
    for _ in range(args.exact_reps):
        n, k = 18, 3
        c = rng.lognormal(0, 0.5, n)
        w = rng.normal(size=n)
        idx, value, iterations, _ = solve(w, c, k)
        oracle_value, oracle_idx = enumerate_optimum(w, c, k)
        max_value_error = max(max_value_error, abs(value - oracle_value))
        exact_matches += set(idx) == set(oracle_idx)
        max_iterations = max(max_iterations, iterations)

    # Claim 4 scale. Keep generated arrays fixed so the timer isolates selection.
    n, k = 1_000_000, 100_000
    x = rng.normal(size=n)
    residual = rng.normal(size=n)
    w, c = x * residual, x * x
    times_ms, iterations_seen = [], []
    for _ in range(args.runtime_reps):
        start = time.perf_counter()
        _, _, its, _ = solve(w, c, k)
        times_ms.append((time.perf_counter() - start) * 1000)
        iterations_seen.append(its)

    # Algorithm 2: compare a 1..K warm path with zero-started independent runs.
    n_path, K = 20_000, 100
    x = rng.normal(size=n_path)
    residual = rng.normal(size=n_path)
    wp, cp = x * residual, x * x
    eta = 0.0
    warm_iterations, cold_iterations, same_values = [], [], []
    for path_k in range(1, K + 1):
        _, warm_value, warm_it, _ = solve(wp, cp, path_k, eta=eta)
        _, cold_value, cold_it, _ = solve(wp, cp, path_k, eta=0.0)
        eta = warm_value
        warm_iterations.append(warm_it)
        cold_iterations.append(cold_it)
        same_values.append(abs(warm_value - cold_value) < 1e-12)

    # A finite-sample stability proxy for Theorem 2: bounded perturbations of
    # oracle scores should uniformly shrink and recover a separated maximizer.
    n_stab, k_stab = 80, 4
    c_stab = rng.uniform(0.5, 1.5, n_stab)
    oracle_w = rng.normal(size=n_stab)
    oracle_idx, oracle_value, _, _ = solve(oracle_w, c_stab, k_stab)
    stability = []
    for noise in [0.5, 0.2, 0.1, 0.05, 0.02, 0.01]:
        errs, recovered = [], []
        for _ in range(200):
            empirical_w = oracle_w + rng.normal(scale=noise, size=n_stab)
            idx, value, _, _ = solve(empirical_w, c_stab, k_stab)
            errs.append(abs(value - oracle_value))
            recovered.append(set(idx) == set(oracle_idx))
        stability.append({"noise_sd": noise, "median_value_error": float(np.median(errs)),
                          "set_recovery_rate": float(np.mean(recovered))})

    result = {
        "scope": "independent NumPy implementation; not the authors' R/Rcpp hardware",
        "environment": {"python": platform.python_version(), "numpy": np.__version__,
                        "machine": platform.machine(), "processor": platform.processor()},
        "claim_1_2_exactness": {"trials": args.exact_reps, "exact_set_matches": exact_matches,
                                "max_value_abs_error": max_value_error,
                                "max_iterations": max_iterations},
        "claim_4_runtime": {"n": n, "k": k, "repetitions": args.runtime_reps,
                            "median_ms": float(np.median(times_ms)),
                            "p95_ms": float(np.percentile(times_ms, 95)),
                            "median_iterations": float(np.median(iterations_seen)),
                            "max_iterations": int(max(iterations_seen))},
        "claim_5_warm_start": {"K": K, "all_objectives_match_cold": bool(all(same_values)),
                               "warm_total_iterations": int(sum(warm_iterations)),
                               "cold_total_iterations": int(sum(cold_iterations)),
                               "warm_median_iterations": float(np.median(warm_iterations)),
                               "cold_median_iterations": float(np.median(cold_iterations))},
        "claim_3_stability_proxy": stability,
    }
    out = Path(args.output)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()