#!/usr/bin/env python3 """Analyze budgeted teacher dataset: per-budget stats, error improvement, saturation. Usage: python outlook/analyze_budget_teacher.py python outlook/analyze_budget_teacher.py --data-dir outlook/data/budget_density_dataset_check python outlook/analyze_budget_teacher.py --sat-threshold 0.05 """ import argparse from pathlib import Path import numpy as np def load_dataset(data_dir: Path): """Load all sample .npz files, return list of dicts.""" data_dir = Path(data_dir) params_path = data_dir / "params_list.npz" if params_path.exists(): n_samples = len(np.load(params_path)["params"]) else: # Fall back to scanning for numbered .npz files n_samples = len(sorted(data_dir.glob("[0-9]*.npz"))) samples = [] for sid in range(n_samples): npz_path = data_dir / f"{sid:04d}.npz" if not npz_path.exists(): continue d = np.load(npz_path) samples.append({ "sid": sid, "budgets": d["budgets"], "actual_elements": d["actual_elements"], "teacher_aw_rel": d["teacher_aw_rel"], "teacher_max_err": d["teacher_max_err"], "valid_budgets": d["valid_budgets"] if "valid_budgets" in d else np.ones(len(d["budgets"]), dtype=bool), "match_ratios": d["match_ratios"] if "match_ratios" in d else d["actual_elements"] / d["budgets"], "k": float(d["k"]), "eps_r": float(d["eps_r"]), }) return samples def per_budget_summary(samples): """Aggregate stats grouped by budget value.""" groups = {} # budget -> list of sample dicts for s in samples: for bi, b in enumerate(s["budgets"]): b = int(b) groups.setdefault(b, {"aw_rel": [], "max_err": [], "actual": [], "valid": []}) groups[b]["aw_rel"].append(s["teacher_aw_rel"][bi]) groups[b]["max_err"].append(s["teacher_max_err"][bi]) groups[b]["actual"].append(s["actual_elements"][bi]) groups[b]["valid"].append(bool(s["valid_budgets"][bi])) rows = [] for b in sorted(groups.keys()): g = groups[b] aw = np.array(g["aw_rel"]) me = np.array(g["max_err"]) rows.append({ "budget": b, "n_samples": len(aw), "n_valid": sum(g["valid"]), "mean_actual": np.mean(g["actual"]), "mean_aw_rel": np.mean(aw), "std_aw_rel": np.std(aw), "mean_max_err": np.mean(me), }) return rows def compute_improvements(rows): """Compute Δaw_rel and relative improvement between adjacent budgets.""" for i in range(len(rows)): rows[i]["delta_aw_rel"] = None rows[i]["rel_improvement"] = None for i in range(1, len(rows)): prev = rows[i - 1]["mean_aw_rel"] curr = rows[i]["mean_aw_rel"] delta = prev - curr # positive = improvement rows[i]["delta_aw_rel"] = delta rows[i]["rel_improvement"] = delta / prev if prev > 0 else 0.0 return rows def estimate_saturation(samples, threshold=0.02): """For each sample, find the first budget where further improvement < threshold. Saturation budget = first budget b such that (aw_rel[b-1] - aw_rel[b]) / aw_rel[b-1] < threshold i.e. relative improvement from previous budget is below threshold. Returns list of (sid, saturation_budget, sat_reason). """ results = [] for s in samples: budgets = s["budgets"] aw = s["teacher_aw_rel"] sat_budget = None for bi in range(1, len(budgets)): if np.isnan(aw[bi - 1]) or np.isnan(aw[bi]): continue if aw[bi - 1] <= 0: continue rel_imp = (aw[bi - 1] - aw[bi]) / aw[bi - 1] if rel_imp < threshold: sat_budget = int(budgets[bi]) break # If never saturated, mark as the largest budget if sat_budget is None: sat_budget = int(budgets[-1]) results.append({ "sid": s["sid"], "saturation_budget": sat_budget, }) return results def print_table(rows): """Print the per-budget summary table.""" hdr = (f"{'budget':>8s} {'n':>4s} {'valid':>5s} " f"{'mean_act':>9s} {'mean_aw':>9s} {'std_aw':>9s} {'mean_me':>9s} " f"{'Δaw_rel':>9s} {'rel_imp':>9s}") print(hdr) print("-" * len(hdr)) for r in rows: delta = f"{r['delta_aw_rel']:.6f}" if r["delta_aw_rel"] is not None else "—" rel = f"{r['rel_improvement']:.4f}" if r["rel_improvement"] is not None else "—" print(f"{r['budget']:8d} {r['n_samples']:4d} {r['n_valid']:5d} " f"{r['mean_actual']:9.1f} {r['mean_aw_rel']:.6f} {r['std_aw_rel']:.6f} " f"{r['mean_max_err']:.6f} {delta:>9s} {rel:>9s}") def print_saturation(sat_results, threshold): """Print saturation budget distribution.""" budgets = [r["saturation_budget"] for r in sat_results] unique = sorted(set(budgets)) print(f"\nSaturation budget distribution (threshold={threshold:.0%} rel. improvement):") print(f" {'sat_budget':>12s} {'count':>6s} {'pct':>6s}") for b in unique: count = budgets.count(b) pct = 100.0 * count / len(budgets) print(f" {b:12d} {count:6d} {pct:5.1f}%") print(f" {'median':>12s} {np.median(budgets):.0f}") print(f" {'mean':>12s} {np.mean(budgets):.1f}") def main(): parser = argparse.ArgumentParser(description="Analyze budgeted teacher dataset") parser.add_argument("--data-dir", type=str, default="outlook/data", help="Path to budget dataset directory (default: outlook/data)") parser.add_argument("--sat-threshold", type=float, default=0.02, help="Relative improvement threshold for saturation (default: 0.02)") parser.add_argument("--sat-threshold-2", type=float, default=0.05, help="Second saturation threshold to report (default: 0.05)") args = parser.parse_args() data_dir = Path(args.data_dir) print("=" * 70) print("Budget Teacher Dataset Analysis") print("=" * 70) print(f" Data dir: {data_dir}") samples = load_dataset(data_dir) print(f" Samples: {len(samples)}") budgets_all = samples[0]["budgets"] print(f" Budgets: {[int(b) for b in budgets_all]}") print() # 1. Per-budget summary rows = per_budget_summary(samples) rows = compute_improvements(rows) print("Per-budget summary:") print_table(rows) # 2. Saturation analysis for thresh, label in [(args.sat_threshold, "primary"), (args.sat_threshold_2, "secondary")]: sat = estimate_saturation(samples, threshold=thresh) print_saturation(sat, thresh) # 3. Per-sample detail (compact) print(f"\nPer-sample detail:") print(f" {'sid':>4s} {'k':>6s} {'eps_r':>6s} " + " ".join(f"b{int(b):>5d}" for b in budgets_all)) print(f" {'':>4s} {'':>6s} {'':>6s} " + " ".join(f"{'aw_rel':>5s}" for _ in budgets_all)) print(" " + "-" * (20 + 8 * len(budgets_all))) for s in samples: vals = " ".join(f"{v:.4f}" for v in s["teacher_aw_rel"]) print(f" {s['sid']:4d} {s['k']:6.2f} {s['eps_r']:6.2f} {vals}") print("\nDone.") if __name__ == "__main__": main()