mirror of
https://github.com/velocitatem/PHANTOM.git
synced 2026-07-15 17:43:36 +00:00
455 lines
15 KiB
Python
455 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import subprocess
|
|
from pathlib import Path
|
|
from typing import Iterable
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
from scipy import stats
|
|
|
|
|
|
def _project_root() -> Path:
|
|
return Path(__file__).resolve().parents[5]
|
|
|
|
|
|
def _default_bundle_dir() -> Path:
|
|
base = _project_root() / "engine" / "studies" / "results" / "wandb_sweep_bundles"
|
|
bundles = sorted(
|
|
[path for path in base.glob("bundle_*") if path.is_dir()],
|
|
key=lambda path: path.stat().st_mtime,
|
|
reverse=True,
|
|
)
|
|
if not bundles:
|
|
raise FileNotFoundError(f"No sweep bundle directories found in {base}")
|
|
return bundles[0]
|
|
|
|
|
|
def _bundle_dir_from_id(bundle_id: str) -> Path:
|
|
token = str(bundle_id).strip()
|
|
name = token if token.startswith("bundle_") else f"bundle_{token}"
|
|
path = (
|
|
_project_root()
|
|
/ "engine"
|
|
/ "studies"
|
|
/ "results"
|
|
/ "wandb_sweep_bundles"
|
|
/ name
|
|
)
|
|
if not path.exists():
|
|
raise FileNotFoundError(f"Bundle not found: {path}")
|
|
return path
|
|
|
|
|
|
def _default_output_dir() -> Path:
|
|
return Path(__file__).resolve().parent / "generated" / "final"
|
|
|
|
|
|
def _truthy(value: object) -> bool:
|
|
if isinstance(value, bool):
|
|
return value
|
|
if value is None:
|
|
return False
|
|
return str(value).strip().lower() in {"1", "true", "yes", "on"}
|
|
|
|
|
|
def _mode_of(row: pd.Series) -> str:
|
|
mode_hint = str(row.get("study_mode", "")).strip().lower()
|
|
if mode_hint in {"baseline", "no_robust"}:
|
|
return "baseline"
|
|
if mode_hint in {"defended", "robust"}:
|
|
return "defended"
|
|
if _truthy(row.get("baseline_mode")) or _truthy(row.get("no_robust")):
|
|
return "baseline"
|
|
return "defended"
|
|
|
|
|
|
def _coerce_numeric(frame: pd.DataFrame, columns: Iterable[str]) -> None:
|
|
for column in columns:
|
|
if column in frame.columns:
|
|
frame[column] = pd.to_numeric(frame[column], errors="coerce")
|
|
|
|
|
|
def _load_runs(bundle_dir: Path) -> pd.DataFrame:
|
|
path = bundle_dir / "runs_finished.csv"
|
|
if not path.exists():
|
|
raise FileNotFoundError(f"Missing required file: {path}")
|
|
frame = pd.read_csv(path)
|
|
frame["mode"] = frame.apply(_mode_of, axis=1)
|
|
_coerce_numeric(
|
|
frame,
|
|
[
|
|
"alpha",
|
|
"n_products",
|
|
"eta_ux",
|
|
"lambda_coi",
|
|
"eval_revenue_mean",
|
|
],
|
|
)
|
|
frame = frame[frame["mode"].isin({"baseline", "defended"})].copy()
|
|
return frame
|
|
|
|
|
|
def _get_git_commit() -> str:
|
|
try:
|
|
result = subprocess.run(
|
|
["git", "rev-parse", "HEAD"],
|
|
check=True,
|
|
text=True,
|
|
capture_output=True,
|
|
cwd=_project_root(),
|
|
)
|
|
except Exception:
|
|
return "unknown"
|
|
return result.stdout.strip()
|
|
|
|
|
|
def _apply_filters(frame: pd.DataFrame, args: argparse.Namespace) -> pd.DataFrame:
|
|
data = frame.copy()
|
|
if args.sweep_id:
|
|
allowed = {str(value) for value in args.sweep_id}
|
|
data = data[data["sweep_id"].astype(str).isin(allowed)]
|
|
if args.mode != "all":
|
|
data = data[data["mode"] == args.mode]
|
|
if args.n_products is not None:
|
|
data = data[data["n_products"] == float(args.n_products)]
|
|
if args.eta_ux is not None:
|
|
data = data[data["eta_ux"] == float(args.eta_ux)]
|
|
if args.lambda_coi is not None:
|
|
data = data[data["lambda_coi"] == float(args.lambda_coi)]
|
|
data = data[data["alpha"].notna() & data["eval_revenue_mean"].notna()]
|
|
data = data[data["alpha"] >= float(args.alpha_min)]
|
|
data = data[data["alpha"] <= float(args.alpha_max)]
|
|
return data.reset_index(drop=True)
|
|
|
|
|
|
def _design_matrix(
|
|
frame: pd.DataFrame,
|
|
*,
|
|
include_sweep_fixed_effects: bool,
|
|
) -> tuple[np.ndarray, np.ndarray, list[str]]:
|
|
y = frame["eval_revenue_mean"].to_numpy(dtype=float)
|
|
x_alpha = frame["alpha"].to_numpy(dtype=float)
|
|
columns = ["intercept", "alpha"]
|
|
blocks = [np.ones_like(x_alpha), x_alpha]
|
|
if include_sweep_fixed_effects:
|
|
dummies = pd.get_dummies(
|
|
frame["sweep_id"].astype(str), prefix="sweep", drop_first=True
|
|
)
|
|
if not dummies.empty:
|
|
blocks.append(dummies.to_numpy(dtype=float).T)
|
|
columns.extend(dummies.columns.tolist())
|
|
X = np.vstack(blocks).T
|
|
return X, y, columns
|
|
|
|
|
|
def _covariance_hc1(X: np.ndarray, residuals: np.ndarray) -> np.ndarray:
|
|
n, k = X.shape
|
|
xtx_inv = np.linalg.pinv(X.T @ X)
|
|
xr = X * residuals[:, None]
|
|
meat = xr.T @ xr
|
|
scale = float(n) / max(n - k, 1)
|
|
return scale * (xtx_inv @ meat @ xtx_inv)
|
|
|
|
|
|
def _covariance_cluster(
|
|
X: np.ndarray, residuals: np.ndarray, groups: pd.Series
|
|
) -> tuple[np.ndarray, int]:
|
|
xtx_inv = np.linalg.pinv(X.T @ X)
|
|
unique = pd.Series(groups).astype(str).dropna().unique().tolist()
|
|
g = len(unique)
|
|
n, k = X.shape
|
|
if g <= 1:
|
|
return _covariance_hc1(X, residuals), g
|
|
meat = np.zeros((k, k), dtype=float)
|
|
for value in unique:
|
|
mask = pd.Series(groups).astype(str).to_numpy() == value
|
|
Xg = X[mask]
|
|
ug = residuals[mask]
|
|
xu = Xg.T @ ug
|
|
meat += np.outer(xu, xu)
|
|
c = (g / (g - 1.0)) * ((n - 1.0) / max(n - k, 1.0))
|
|
return c * (xtx_inv @ meat @ xtx_inv), g
|
|
|
|
|
|
def _fit_ols(
|
|
X: np.ndarray,
|
|
y: np.ndarray,
|
|
columns: list[str],
|
|
*,
|
|
cov_type: str,
|
|
groups: pd.Series | None = None,
|
|
) -> dict[str, object]:
|
|
n, k = X.shape
|
|
beta, _, _, _ = np.linalg.lstsq(X, y, rcond=None)
|
|
fitted = X @ beta
|
|
residuals = y - fitted
|
|
dof = max(n - k, 1)
|
|
sse = float(np.sum(residuals**2))
|
|
y_centered = y - float(np.mean(y))
|
|
sst = float(np.sum(y_centered**2))
|
|
r2 = float(1.0 - sse / sst) if sst > 0 else 0.0
|
|
adj_r2 = float(1.0 - (1.0 - r2) * ((n - 1.0) / max(n - k, 1.0)))
|
|
|
|
if cov_type == "iid":
|
|
sigma2 = sse / dof
|
|
cov = sigma2 * np.linalg.pinv(X.T @ X)
|
|
df_t = dof
|
|
clusters = None
|
|
elif cov_type == "hc1":
|
|
cov = _covariance_hc1(X, residuals)
|
|
df_t = dof
|
|
clusters = None
|
|
elif cov_type == "cluster":
|
|
if groups is None:
|
|
raise ValueError("groups are required when cov_type='cluster'")
|
|
cov, clusters = _covariance_cluster(X, residuals, groups)
|
|
df_t = max(clusters - 1, 1)
|
|
else:
|
|
raise ValueError(f"Unsupported cov_type: {cov_type}")
|
|
|
|
se = np.sqrt(np.clip(np.diag(cov), 0.0, np.inf))
|
|
t_stats = np.divide(beta, se, out=np.zeros_like(beta), where=se > 0)
|
|
p_values = 2.0 * (1.0 - stats.t.cdf(np.abs(t_stats), df=df_t))
|
|
t_crit = float(stats.t.ppf(0.975, df=df_t))
|
|
ci_low = beta - t_crit * se
|
|
ci_high = beta + t_crit * se
|
|
|
|
coef_rows = []
|
|
for idx, name in enumerate(columns):
|
|
coef_rows.append(
|
|
{
|
|
"name": name,
|
|
"coef": float(beta[idx]),
|
|
"std_error": float(se[idx]),
|
|
"t_stat": float(t_stats[idx]),
|
|
"p_value": float(p_values[idx]),
|
|
"ci95_low": float(ci_low[idx]),
|
|
"ci95_high": float(ci_high[idx]),
|
|
}
|
|
)
|
|
|
|
return {
|
|
"n": int(n),
|
|
"k": int(k),
|
|
"dof": int(dof),
|
|
"df_t": int(df_t),
|
|
"cov_type": cov_type,
|
|
"clusters": int(clusters) if clusters is not None else None,
|
|
"r2": r2,
|
|
"adj_r2": adj_r2,
|
|
"sse": sse,
|
|
"coefficients": coef_rows,
|
|
"residuals": residuals,
|
|
"fitted": fitted,
|
|
"beta": beta,
|
|
}
|
|
|
|
|
|
def _diagnostics(
|
|
X: np.ndarray, y: np.ndarray, fit: dict[str, object]
|
|
) -> dict[str, object]:
|
|
residuals = np.asarray(fit["residuals"], dtype=float)
|
|
n, k = X.shape
|
|
if residuals.size < 8:
|
|
normality = {"test": "jarque_bera", "available": False}
|
|
else:
|
|
jb_stat, jb_p = stats.jarque_bera(residuals)
|
|
normality = {
|
|
"test": "jarque_bera",
|
|
"available": True,
|
|
"statistic": float(jb_stat),
|
|
"p_value": float(jb_p),
|
|
}
|
|
|
|
if k <= 1:
|
|
hetero = {"test": "breusch_pagan", "available": False}
|
|
else:
|
|
u2 = residuals**2
|
|
aux = _fit_ols(X, u2, [f"x{i}" for i in range(k)], cov_type="iid")
|
|
lm = float(len(u2) * float(aux["r2"]))
|
|
df_bp = k - 1
|
|
p_bp = float(1.0 - stats.chi2.cdf(lm, df_bp))
|
|
hetero = {
|
|
"test": "breusch_pagan",
|
|
"available": True,
|
|
"lm_stat": lm,
|
|
"df": int(df_bp),
|
|
"p_value": p_bp,
|
|
}
|
|
|
|
xtx_inv = np.linalg.pinv(X.T @ X)
|
|
leverages = np.sum((X @ xtx_inv) * X, axis=1)
|
|
mse = float(np.sum(residuals**2) / max(n - k, 1))
|
|
if mse <= 0:
|
|
cooks = np.zeros(n, dtype=float)
|
|
else:
|
|
denom = np.clip((1.0 - leverages) ** 2, 1e-10, np.inf)
|
|
cooks = ((residuals**2) / (k * mse)) * (leverages / denom)
|
|
|
|
return {
|
|
"normality": normality,
|
|
"heteroskedasticity": hetero,
|
|
"influence": {
|
|
"max_leverage": float(np.max(leverages)) if leverages.size else 0.0,
|
|
"mean_leverage": float(np.mean(leverages)) if leverages.size else 0.0,
|
|
"high_leverage_threshold": float(2.0 * k / max(n, 1)),
|
|
"high_leverage_count": int(np.sum(leverages > (2.0 * k / max(n, 1)))),
|
|
"max_cooks_distance": float(np.max(cooks)) if cooks.size else 0.0,
|
|
"high_cooks_threshold": float(4.0 / max(n, 1)),
|
|
"high_cooks_count": int(np.sum(cooks > (4.0 / max(n, 1)))),
|
|
},
|
|
}
|
|
|
|
|
|
def run(args: argparse.Namespace) -> list[Path]:
|
|
output_dir = Path(args.output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
runs = _load_runs(Path(args.bundle_dir))
|
|
filtered = _apply_filters(runs, args)
|
|
if len(filtered) < 3:
|
|
raise ValueError("Filtered cohort must contain at least 3 rows")
|
|
if filtered["alpha"].nunique() < 2:
|
|
raise ValueError("Filtered cohort must contain at least 2 unique alpha values")
|
|
|
|
filtered_csv = output_dir / "revenue_alpha_filtered.csv"
|
|
filtered.to_csv(filtered_csv, index=False)
|
|
|
|
sample_accounting = {
|
|
"bundle_dir": str(Path(args.bundle_dir)),
|
|
"git_commit": _get_git_commit(),
|
|
"cohort_name": str(args.cohort_name),
|
|
"filters": {
|
|
"sweep_id": args.sweep_id,
|
|
"mode": args.mode,
|
|
"n_products": args.n_products,
|
|
"eta_ux": args.eta_ux,
|
|
"lambda_coi": args.lambda_coi,
|
|
"alpha_min": args.alpha_min,
|
|
"alpha_max": args.alpha_max,
|
|
},
|
|
"n_rows": int(len(filtered)),
|
|
"n_sweeps": int(filtered["sweep_id"].nunique()),
|
|
"alpha_unique": sorted(
|
|
float(v) for v in filtered["alpha"].dropna().unique().tolist()
|
|
),
|
|
"rows_by_sweep": filtered.groupby("sweep_id").size().astype(int).to_dict(),
|
|
"rows_by_mode": filtered.groupby("mode").size().astype(int).to_dict(),
|
|
}
|
|
sample_path = output_dir / "revenue_alpha_sample_accounting.json"
|
|
sample_path.write_text(json.dumps(sample_accounting, indent=2) + "\n")
|
|
|
|
X_simple, y, cols_simple = _design_matrix(
|
|
filtered, include_sweep_fixed_effects=False
|
|
)
|
|
fit_simple = _fit_ols(X_simple, y, cols_simple, cov_type="iid")
|
|
simple_path = output_dir / "revenue_alpha_simple_ols.json"
|
|
simple_path.write_text(
|
|
json.dumps(
|
|
{
|
|
k: v
|
|
for k, v in fit_simple.items()
|
|
if k not in {"residuals", "fitted", "beta"}
|
|
},
|
|
indent=2,
|
|
)
|
|
+ "\n"
|
|
)
|
|
|
|
X_fe, y_fe, cols_fe = _design_matrix(filtered, include_sweep_fixed_effects=True)
|
|
cov_type = "cluster" if filtered["sweep_id"].nunique() > 1 else "hc1"
|
|
fit_fe = _fit_ols(
|
|
X_fe, y_fe, cols_fe, cov_type=cov_type, groups=filtered["sweep_id"]
|
|
)
|
|
fe_path = output_dir / "revenue_alpha_fixed_effects.json"
|
|
fe_path.write_text(
|
|
json.dumps(
|
|
{
|
|
k: v
|
|
for k, v in fit_fe.items()
|
|
if k not in {"residuals", "fitted", "beta"}
|
|
},
|
|
indent=2,
|
|
)
|
|
+ "\n"
|
|
)
|
|
|
|
per_sweep_rows: list[dict[str, float | str | int]] = []
|
|
for sweep_id, group in filtered.groupby("sweep_id"):
|
|
if len(group) < 3 or group["alpha"].nunique() < 2:
|
|
continue
|
|
X_sw, y_sw, cols_sw = _design_matrix(group, include_sweep_fixed_effects=False)
|
|
fit_sw = _fit_ols(X_sw, y_sw, cols_sw, cov_type="hc1")
|
|
alpha_row = next(
|
|
row for row in fit_sw["coefficients"] if row["name"] == "alpha"
|
|
)
|
|
per_sweep_rows.append(
|
|
{
|
|
"sweep_id": str(sweep_id),
|
|
"n": int(fit_sw["n"]),
|
|
"alpha_coef": float(alpha_row["coef"]),
|
|
"alpha_std_error": float(alpha_row["std_error"]),
|
|
"alpha_t_stat": float(alpha_row["t_stat"]),
|
|
"alpha_p_value": float(alpha_row["p_value"]),
|
|
"alpha_ci95_low": float(alpha_row["ci95_low"]),
|
|
"alpha_ci95_high": float(alpha_row["ci95_high"]),
|
|
"r2": float(fit_sw["r2"]),
|
|
}
|
|
)
|
|
per_sweep_frame = pd.DataFrame(per_sweep_rows)
|
|
if not per_sweep_frame.empty:
|
|
per_sweep_frame = per_sweep_frame.sort_values("sweep_id").reset_index(drop=True)
|
|
per_sweep_path = output_dir / "revenue_alpha_per_sweep.csv"
|
|
per_sweep_frame.to_csv(per_sweep_path, index=False)
|
|
|
|
fit_for_diagnostics = fit_fe if cov_type == "cluster" else fit_simple
|
|
X_for_diagnostics = X_fe if cov_type == "cluster" else X_simple
|
|
diagnostics = _diagnostics(X_for_diagnostics, y, fit_for_diagnostics)
|
|
diagnostics_path = output_dir / "revenue_alpha_diagnostics.json"
|
|
diagnostics_path.write_text(json.dumps(diagnostics, indent=2) + "\n")
|
|
|
|
return [
|
|
filtered_csv,
|
|
sample_path,
|
|
simple_path,
|
|
fe_path,
|
|
per_sweep_path,
|
|
diagnostics_path,
|
|
]
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description="Reproducible contamination-vs-revenue analysis from a sweep bundle"
|
|
)
|
|
parser.add_argument("--bundle-dir", type=Path, default=None)
|
|
parser.add_argument("--bundle-id", type=str, default=None)
|
|
parser.add_argument("--output-dir", type=Path, default=_default_output_dir())
|
|
parser.add_argument("--cohort-name", type=str, default="custom")
|
|
parser.add_argument("--sweep-id", action="append", default=[])
|
|
parser.add_argument(
|
|
"--mode", choices=["all", "baseline", "defended"], default="all"
|
|
)
|
|
parser.add_argument("--n-products", type=float, default=None)
|
|
parser.add_argument("--eta-ux", type=float, default=None)
|
|
parser.add_argument("--lambda-coi", type=float, default=None)
|
|
parser.add_argument("--alpha-min", type=float, default=0.0)
|
|
parser.add_argument("--alpha-max", type=float, default=1.0)
|
|
args = parser.parse_args()
|
|
|
|
if args.bundle_id:
|
|
args.bundle_dir = _bundle_dir_from_id(args.bundle_id)
|
|
elif args.bundle_dir is None:
|
|
args.bundle_dir = _default_bundle_dir()
|
|
|
|
outputs = run(args)
|
|
for path in outputs:
|
|
print(path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|