"""Check subgroup patterns using forward out-of-fold chronological TRAIN only.""" from __future__ import annotations import hashlib import json from pathlib import Path import numpy as np import pandas as pd from sklearn.metrics import roc_curve from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.linear_model import SGDClassifier ROOT = Path(__file__).resolve().parents[1] CORPUS = ROOT / "data/processed/corpus.parquet" TRAIN = ROOT / "data/splits/chronological/train.json" DATE_CACHE = ROOT / "data/processed/secvuleval_commit_dates.jsonl" OUT_JSON = ROOT / "results/phase7/char_forward_train_subgroups.json" OUT_MD = ROOT / "results/phase7/char_forward_train_subgroups.md" SEEDS = tuple(range(10)) FOLDS = (2019, 2020, 2021) def wilson(k: int, n: int) -> list[float] | None: if not n: return None z, p = 1.96, k / n d = 1 + z*z/n c = p + z*z/(2*n) h = z*np.sqrt(p*(1-p)/n + z*z/(4*n*n)) return [float((c-h)/d), float((c+h)/d)] def metric(mask: np.ndarray, y: np.ndarray, pred: np.ndarray) -> dict: safe = mask & (y == 0) vuln = mask & (y == 1) good = pred == y n, k = int(mask.sum()), int(np.sum(mask & good)) ks, kv = int(np.sum(safe & good)), int(np.sum(vuln & good)) return { "n_rows": n, "n_pairs": n // 2, "accuracy": k/n if n else None, "wilson_95": wilson(k, n), "safe_n": int(safe.sum()), "safe_accuracy": ks/int(safe.sum()) if safe.any() else None, "safe_wilson_95": wilson(ks, int(safe.sum())), "vulnerable_n": int(vuln.sum()), "vulnerable_accuracy": kv/int(vuln.sum()) if vuln.any() else None, "vulnerable_wilson_95": wilson(kv, int(vuln.sum())), } def fit_score(fit: pd.DataFrame, held: pd.DataFrame, sample_weight: np.ndarray | None = None) -> np.ndarray: vectorizer = TfidfVectorizer(analyzer="char", ngram_range=(3, 5), min_df=2, max_features=300_000, sublinear_tf=True, dtype=np.float32) xfit = vectorizer.fit_transform(fit.text.tolist()) xheld = vectorizer.transform(held.text.tolist()) yfit = fit.target.to_numpy(dtype=np.int8) all_scores = [] for seed in SEEDS: model = SGDClassifier(loss="log_loss", alpha=1e-5, max_iter=200, tol=1e-3, random_state=seed, average=True) model.fit(xfit, yfit, sample_weight=sample_weight) all_scores.append(model.decision_function(xheld)) return np.mean(np.stack(all_scores), axis=0) def length_label_weights(fit: pd.DataFrame) -> np.ndarray: """Equalize the six observed label × fixed-length cells using TRAIN only.""" labels = fit.target.to_numpy(dtype=np.int8) lines = fit.vuln_lines.to_numpy() bins = np.where(lines <= 30, 0, np.where(lines <= 100, 1, 2)) cell = labels * 3 + bins counts = np.bincount(cell, minlength=6) present = counts > 0 n_present = int(present.sum()) weights = np.asarray([len(fit) / (n_present * counts[c]) for c in cell], dtype=np.float64) if not np.isclose(weights.mean(), 1.0): raise AssertionError("sample weights must have mean one") return weights def select_threshold(scores: np.ndarray, labels: np.ndarray) -> float: """Select a balanced-accuracy threshold on calibration-only predictions.""" fpr, tpr, thresholds = roc_curve(labels, scores, pos_label=1) finite = np.isfinite(thresholds) thresholds, fpr, tpr = thresholds[finite], fpr[finite], tpr[finite] prevalence = float(np.mean(labels == 1)) ba = (tpr + 1 - fpr) / 2 accuracy = prevalence * tpr + (1 - prevalence) * (1 - fpr) order = np.lexsort((- np.abs(thresholds), accuracy, ba)) return float(thresholds[order[-1]]) def nested_length_thresholds(oof: pd.DataFrame) -> list[dict]: """Calibrate per fixed length band on earlier OOF years, score next year.""" reports = [] rng = np.random.default_rng(42) bands = (("short_le_30", 1, 30), ("medium_31_100", 31, 100), ("long_gt_100", 101, None)) for year in FOLDS: calibration = oof[oof.year < year] held = oof[oof.year == year] if calibration.empty: continue cy = calibration.target.to_numpy(dtype=np.int8) cs = calibration.score.to_numpy() global_threshold = select_threshold(cs, cy) thresholds = {} for name, low, high in bands: cmask = calibration.vuln_lines.to_numpy() >= low if high is not None: cmask &= calibration.vuln_lines.to_numpy() <= high counts = [int(np.sum(cmask & (cy == label))) for label in (0, 1)] if min(counts) >= 50: threshold = select_threshold(cs[cmask], cy[cmask]) fallback = False else: threshold = global_threshold fallback = True thresholds[name] = {"threshold": threshold, "calibration_safe_n": counts[0], "calibration_vulnerable_n": counts[1], "pooled_threshold_fallback": fallback} ey = held.target.to_numpy(dtype=np.int8) es = held.score.to_numpy() lens = held.vuln_lines.to_numpy() pred = np.empty(len(held), dtype=np.int8) for name, low, high in bands: mask = lens >= low if high is not None: mask &= lens <= high pred[mask] = (es[mask] >= thresholds[name]["threshold"]).astype(np.int8) fixed = (es >= 0).astype(np.int8) pair_effects = [] rows = pd.DataFrame({"pair_id": held.pair_id.to_numpy(), "y": ey, "fixed": fixed, "calibrated": pred}) for _, group in rows.groupby("pair_id", sort=True): if len(group) != 2 or set(group.y.astype(int)) != {0, 1}: raise ValueError("pair integrity failed in threshold comparison") pair_effects.append(float(np.mean(group.calibrated.to_numpy() == group.y.to_numpy()) - np.mean(group.fixed.to_numpy() == group.y.to_numpy()))) effects = np.asarray(pair_effects) draws = effects[rng.integers(0, len(effects), size=(10_000, len(effects)))].mean(axis=1) reports.append({ "held_year": int(year), "calibration_years": sorted(int(y) for y in calibration.year.unique()), "calibration_pairs": int(calibration.pair_id.nunique()), "thresholds_by_length": thresholds, "fixed_zero_threshold": metric(np.ones(len(held), dtype=bool), ey, fixed), "length_calibrated": metric(np.ones(len(held), dtype=bool), ey, pred), "balanced_accuracy_delta_pp": float(100 * effects.mean()), "pair_bootstrap_95_ci_delta_pp": [ float(100 * x) for x in np.quantile(draws, [0.025, 0.975])], }) return reports def main() -> None: frame = pd.read_parquet(CORPUS) date_map = {} for line in DATE_CACHE.read_text().splitlines(): obj = json.loads(line) if obj.get("date"): date_map[str(obj["commit_id"])] = str(obj["date"]) missing = frame.commit_date.isna() frame.loc[missing, "commit_date"] = frame.loc[missing, "commit_id"].astype(str).map(date_map) frame["year"] = pd.to_numeric(frame.commit_date.astype(str).str[:4], errors="coerce") train_ids = set(map(str, json.loads(TRAIN.read_text()))) sub = frame[frame.pair_id.astype(str).isin(train_ids)] records = [] for pid, grp in sub.groupby("pair_id", sort=True): if len(grp) != 2 or set(grp.target.astype(int)) != {0, 1}: raise ValueError(f"pair integrity failure: {pid}") if grp.func_vuln.nunique() != 1 or grp.func_patch.nunique() != 1: raise ValueError(f"paired code mismatch: {pid}") if grp.cwe.nunique(dropna=False) != 1 or grp.year.nunique(dropna=False) != 1: raise ValueError(f"paired metadata mismatch: {pid}") vrow = grp[grp.target.astype(int) == 1].iloc[0] for target in (0, 1): row = grp[grp.target.astype(int) == target].iloc[0] records.append({ "pair_id": str(pid), "target": target, "year": row.year, "text": str(row.func_vuln if target else row.func_patch), "cwe": str(vrow.cwe) if pd.notna(vrow.cwe) else "UNKNOWN", "vuln_lines": max(1, len(str(vrow.func_vuln).splitlines())), }) train = pd.DataFrame(records) oof_parts = [] weighted_parts = [] fold_info = [] for year in FOLDS: fit = train[train.year < year].reset_index(drop=True) held = train[train.year == year].reset_index(drop=True) if fit.empty or held.empty or held.target.nunique() != 2: raise ValueError(f"invalid forward fold {year}") scores = fit_score(fit, held) weighted_scores = fit_score(fit, held, length_label_weights(fit)) held = held.copy() held["score"] = scores held["prediction"] = (scores >= 0).astype(np.int8) oof_parts.append(held) weighted_held = held.drop(columns=["score", "prediction"]).copy() weighted_held["score"] = weighted_scores weighted_held["prediction"] = (weighted_scores >= 0).astype(np.int8) weighted_parts.append(weighted_held) fold_info.append({"held_year": year, "fit_pairs": int(fit.pair_id.nunique()), "held_pairs": int(held.pair_id.nunique()), "metrics": metric(np.ones(len(held), dtype=bool), held.target.to_numpy(), held.prediction.to_numpy())}) print(json.dumps({"held_year": year, "fit_pairs": int(fit.pair_id.nunique()), "held_pairs": int(held.pair_id.nunique())}), flush=True) oof = pd.concat(oof_parts, ignore_index=True) weighted_oof = pd.concat(weighted_parts, ignore_index=True) y = oof.target.to_numpy(dtype=np.int8) pred = oof.prediction.to_numpy(dtype=np.int8) all_mask = np.ones(len(oof), dtype=bool) nested_calibration = nested_length_thresholds(oof) # Paired out-of-fold comparison for TRAIN-only length × label weighting. wy = weighted_oof.target.to_numpy(dtype=np.int8) wp = weighted_oof.prediction.to_numpy(dtype=np.int8) base_rows = oof[["pair_id", "year", "target", "vuln_lines", "prediction"]].copy() base_rows = base_rows.rename(columns={"prediction": "base_prediction"}) weighted_rows = weighted_oof[["pair_id", "target", "prediction"]].rename( columns={"prediction": "weighted_prediction"}) paired_compare = base_rows.merge(weighted_rows, on=["pair_id", "target"], validate="one_to_one") if len(paired_compare) != len(oof): raise AssertionError("weighted and baseline OOF row alignment failed") pair_effects = [] for _, group in paired_compare.groupby("pair_id", sort=True): if len(group) != 2 or set(group.target.astype(int)) != {0, 1}: raise ValueError("pair integrity failure in weighted comparison") base_ok = group.base_prediction.to_numpy() == group.target.to_numpy() weighted_ok = group.weighted_prediction.to_numpy() == group.target.to_numpy() pair_effects.append(float(weighted_ok.mean() - base_ok.mean())) pair_effects = np.asarray(pair_effects) rng = np.random.default_rng(77) pair_draws = pair_effects[rng.integers(0, len(pair_effects), size=(10_000, len(pair_effects)))].mean(axis=1) weighted_comparison = { "weighting": "equal total sample weight for each observed target × fixed length band cell", "baseline": metric(np.ones(len(oof), dtype=bool), y, pred), "length_label_weighted": metric(np.ones(len(weighted_oof), dtype=bool), wy, wp), "paired_accuracy_delta_pp": float(100 * pair_effects.mean()), "pair_bootstrap_95_ci_delta_pp": [ float(100 * v) for v in np.quantile(pair_draws, [0.025, 0.975])], "by_held_year": {}, "by_fixed_length_band": {}, } for year in FOLDS: bm = oof.year.to_numpy() == year wm = weighted_oof.year.to_numpy() == year weighted_comparison["by_held_year"][str(year)] = { "baseline": metric(bm, y, pred), "length_label_weighted": metric(wm, wy, wp), } for name, low, high in (("short_le_30", 1, 30), ("medium_31_100", 31, 100), ("long_gt_100", 101, None)): lm = oof.vuln_lines.to_numpy() >= low if high is not None: lm &= oof.vuln_lines.to_numpy() <= high wlm = weighted_oof.vuln_lines.to_numpy() >= low if high is not None: wlm &= weighted_oof.vuln_lines.to_numpy() <= high weighted_comparison["by_fixed_length_band"][name] = { "baseline": metric(lm, y, pred), "length_label_weighted": metric(wlm, wy, wp), } cwe_results = [] for cwe, inds in oof.groupby("cwe", sort=True).groups.items(): mask = np.zeros(len(oof), dtype=bool) mask[np.asarray(list(inds), dtype=int)] = True by_year = {} for year in FOLDS: ym = mask & (oof.year.to_numpy() == year) by_year[str(year)] = metric(ym, y, pred) cwe_results.append({"cwe": cwe, "pooled": metric(mask, y, pred), "by_held_year": by_year}) length_results = [] bands = [("short_le_30", 1, 30), ("medium_31_100", 31, 100), ("long_gt_100", 101, None)] for name, low, high in bands: lens = oof.vuln_lines.to_numpy() mask = (lens >= low) & ((lens <= high) if high is not None else True) by_year = {str(year): metric(mask & (oof.year.to_numpy() == year), y, pred) for year in FOLDS} length_results.append({"band": name, "min_lines": low, "max_lines": high, "pooled": metric(mask, y, pred), "by_held_year": by_year}) result = { "method": "fixed ten-seed char TF-IDF/SGD, forward TRAIN OOF, zero threshold", "folds": fold_info, "pooled_oof": metric(all_mask, y, pred), "cwe_slices": cwe_results, "vulnerable_function_length_bands": length_results, "nested_forward_length_calibration": nested_calibration, "length_label_weighting_comparison": weighted_comparison, "coverage": { "train_pairs": len(train_ids), "oof_pairs": int(oof.pair_id.nunique()), "oof_rows": len(oof), "undated_or_other_year_pairs_not_in_folds": len(train_ids - set(oof.pair_id)), }, "protocol": { "held_years": list(FOLDS), "seeds": list(SEEDS), "threshold": 0.0, "subgroup_analysis_is_descriptive": True, "length_bands_prespecified_lines": [30, 100], "valid_used_for_selection": False, "test_read": False, }, "provenance": { "corpus_sha256": hashlib.sha256(CORPUS.read_bytes()).hexdigest(), "train_split_sha256": hashlib.sha256(TRAIN.read_bytes()).hexdigest(), "date_cache_sha256": hashlib.sha256(DATE_CACHE.read_bytes()).hexdigest(), }, "limitations": [ "Only TRAIN pairs dated 2019–2021 with earlier dated TRAIN history are scored.", "CWE labels may be incomplete/noisy; small subgroup estimates are unstable.", "Forward folds share overlapping fit histories and are not independent replications.", "Per-row Wilson intervals do not account for repository clustering.", ], } OUT_JSON.parent.mkdir(parents=True, exist_ok=True) OUT_JSON.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n") lines = ["# Forward TRAIN subgroup replication", "", "Fixed model and zero threshold; all rows are forward out-of-fold TRAIN predictions. " "No VALID selection or TEST access.", "", f"Pooled: {result['pooled_oof']['accuracy']:.1%} accuracy " f"(Wilson 95% {result['pooled_oof']['wilson_95'][0]:.1%}–" f"{result['pooled_oof']['wilson_95'][1]:.1%}); SAFE " f"{result['pooled_oof']['safe_accuracy']:.1%}, VULNERABLE " f"{result['pooled_oof']['vulnerable_accuracy']:.1%}; " f"{result['coverage']['oof_pairs']} pairs.", "", "## Prespecified function-length bands", "", "| Band | pairs | SAFE accuracy | VULNERABLE accuracy | overall |", "|---|---:|---:|---:|---:|"] for item in length_results: m = item["pooled"] lines.append(f"| {item['band']} | {m['n_pairs']} | {m['safe_accuracy']:.1%} | " f"{m['vulnerable_accuracy']:.1%} | {m['accuracy']:.1%} |") lines += ["", "## CWE slices with at least 30 pooled OOF pairs", "", "| CWE | pairs | SAFE accuracy | VULNERABLE accuracy | held-year vulnerable recalls |", "|---|---:|---:|---:|---|"] for item in cwe_results: m = item["pooled"] if m["n_pairs"] < 30: continue yearly = ", ".join(f"{yr}: {v['vulnerable_accuracy']:.1%}" for yr, v in item["by_held_year"].items() if v["vulnerable_n"]) lines.append(f"| {item['cwe']} | {m['n_pairs']} | {m['safe_accuracy']:.1%} | " f"{m['vulnerable_accuracy']:.1%} | {yearly} |") lines += ["", "## Nested prior-year length calibration", "", "For each held year, thresholds use only earlier-year OOF scores; year 2019 has no prior OOF calibration and is omitted.", "", "| Held year | fixed balanced acc. | calibrated balanced acc. | delta pp (pair-bootstrap 95% CI) |", "|---:|---:|---:|---:|"] for item in nested_calibration: fixed, calibrated = item["fixed_zero_threshold"], item["length_calibrated"] fixed_ba = (fixed["safe_accuracy"] + fixed["vulnerable_accuracy"]) / 2 cal_ba = (calibrated["safe_accuracy"] + calibrated["vulnerable_accuracy"]) / 2 ci = item["pair_bootstrap_95_ci_delta_pp"] lines.append(f"| {item['held_year']} | {fixed_ba:.2%} | {cal_ba:.2%} | " f"{item['balanced_accuracy_delta_pp']:+.2f} ({ci[0]:+.2f}, {ci[1]:+.2f}) |") lines += ["", "## TRAIN-only length × label sample weighting", "", f"Pooled paired accuracy delta: {weighted_comparison['paired_accuracy_delta_pp']:+.2f} pp " f"(pair-bootstrap 95% CI {weighted_comparison['pair_bootstrap_95_ci_delta_pp'][0]:+.2f}, " f"{weighted_comparison['pair_bootstrap_95_ci_delta_pp'][1]:+.2f}).", "", "| Held year | baseline accuracy | weighted accuracy |", "|---:|---:|---:|"] for year, values in weighted_comparison["by_held_year"].items(): lines.append(f"| {year} | {values['baseline']['accuracy']:.2%} | " f"{values['length_label_weighted']['accuracy']:.2%} |") lines += ["", "This is one forward TRAIN diagnostic; length × label weights use only the corresponding fit fold. " "CWE and length slices are descriptive; no VALID selection or TEST access.", ""] OUT_MD.write_text("\n".join(lines)) print(json.dumps({"pooled_oof": result["pooled_oof"], "coverage": result["coverage"], "artifacts": [str(OUT_JSON), str(OUT_MD)]}, indent=2), flush=True) if __name__ == "__main__": main()