diff --git a/pyproject.toml b/pyproject.toml index 992d635..996147a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,13 @@ dependencies = ["pyyaml>=6.0"] [project.optional-dependencies] # Hunter learning research only; the core install stays PyYAML. -hunt = ["dspy>=3.4,<3.5", "pandas>=2.2,<3"] +hunt = [ + "dspy>=3.4,<3.5", + "pandas>=2.2,<3", + "scipy>=1.13,<2", + "scikit-learn>=1.6,<2", + "statsmodels>=0.14,<1", +] # The documentation site (mkdocs.yml, docs/). Mermaid renders client-side through Material's superfences fence. docs = [ "mkdocs>=1.6,<2", diff --git a/src/openultrasast/hunt/features.py b/src/openultrasast/hunt/features.py new file mode 100644 index 0000000..5e577d5 --- /dev/null +++ b/src/openultrasast/hunt/features.py @@ -0,0 +1,123 @@ +"""Feature engineering over the feature table, and the tie between engineered features and the DSPy annotations. + +The table's raw columns (`feature..`) are what the producers measured. This module turns them +into the numeric frame the statistics read (distances, counts, flags), and keeps the map from every engineered +feature back to the annotation the model is shown. An ablation that masks an annotation therefore drops both the +rendered input and the engineered features that came from the same column, and the separation screen can say +which annotation carries signal before any model call is spent on it. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping, Sequence +from typing import Any + +from .table import FeatureTable + +# engineered feature -> (raw column, annotation input name in program.INPUTS) +ENGINEERED: dict[str, tuple[str, str]] = { + "diff.distance": ("feature.diff.distance", "hunks"), + "diff.on_changed_line": ("feature.diff.on_changed_line", "hunks"), + "diff.hunk_count": ("feature.diff.hunks", "hunks"), + "diff.hunk_lines": ("feature.diff.hunks", "hunks"), + "source.span_lines": ("feature.source.span", "function_text"), + "source.is_module": ("feature.source.function", "change_summary"), + "source.caller_count": ("feature.source.callers", "callers"), + "source.caller_files": ("feature.source.callers", "callers"), + "source.entry_point_count": ("feature.source.entry_points", "entry_points"), + "semantic.sources": ("feature.semantic.roles", "roles"), + "semantic.sinks": ("feature.semantic.roles", "roles"), + "semantic.sanitizers": ("feature.semantic.roles", "roles"), + "semantic.binds": ("feature.semantic.def_use", "def_use"), + "semantic.calls": ("feature.semantic.def_use", "def_use"), + "semantic.obligations": ("feature.semantic.obligations", "obligations"), + "semantic.obligations_undischarged": ("feature.semantic.obligations", "obligations"), +} +KEYS = ("change_id", "repo", "head", "path", "line", "kind", "fold", "label", "family") + + +def _count(value: object) -> int | None: + return None if value is None else len(value) if isinstance(value, (list, dict, tuple)) else None + + +def _engineer(row: Mapping[str, Any]) -> dict[str, Any]: + hunks = row.get("feature.diff.hunks") + span = row.get("feature.source.span") + callers = row.get("feature.source.callers") + roles = row.get("feature.semantic.roles") or None + def_use = row.get("feature.semantic.def_use") or None + obligations = row.get("feature.semantic.obligations") + function = row.get("feature.source.function") + out: dict[str, Any] = {key: row.get(key) for key in KEYS} + out.update( + { + "diff.distance": row.get("feature.diff.distance"), + "diff.on_changed_line": None + if row.get("feature.diff.on_changed_line") is None + else int(bool(row["feature.diff.on_changed_line"])), + "diff.hunk_count": _count(hunks), + "diff.hunk_lines": None if hunks is None else sum(end - start + 1 for start, end in hunks), + "source.span_lines": None if not span else span[1] - span[0] + 1, + "source.is_module": None if function is None else int(function == ""), + "source.caller_count": _count(callers), + "source.caller_files": None if callers is None else len({path for path, _ in callers}), + "source.entry_point_count": _count(row.get("feature.source.entry_points")), + "semantic.sources": None if roles is None else sum(1 for role in roles.values() if role == "source"), + "semantic.sinks": None if roles is None else sum(1 for role in roles.values() if role == "sink"), + "semantic.sanitizers": None if roles is None else sum(1 for role in roles.values() if role == "sanitizer"), + "semantic.binds": None if def_use is None else len(def_use.get("binds", [])), + "semantic.calls": None if def_use is None else len(def_use.get("calls", [])), + "semantic.obligations": _count(obligations), + "semantic.obligations_undischarged": None if obligations is None else sum(1 for o in obligations if not o.get("discharged")), + } + ) + for producer in ("diff", "source", "semantic", "taxonomy"): + out[f"ok.{producer}"] = row.get(f"provenance.{producer}.ok") + return out + + +def feature_frame(table: FeatureTable | Sequence[Mapping[str, Any]]) -> Any: + """The engineered frame: one row per candidate, key and label columns, numeric features, provenance flags.""" + import pandas as pd + + rows = table.rows if isinstance(table, FeatureTable) else list(table) + return pd.DataFrame([_engineer(row) for row in rows]) + + +def features_of(annotation: str) -> list[str]: + """The engineered features that come from one annotation (what an ablation mask on it also drops).""" + return [name for name, (_, input_name) in ENGINEERED.items() if input_name == annotation] + + +def annotations_of(mask: Iterable[str]) -> set[str]: + """Annotations an ablation mask removes, whether it names a producer, a column or an input.""" + dropped = set(mask) + out = set() + for name, (column, input_name) in ENGINEERED.items(): + producer = column.split(".")[1] + if input_name in dropped or column in dropped or producer in dropped or name in dropped: + out.add(input_name) + return out + + +def screen(frame: Any, *, label: str = "label") -> Any: + """Separation of every engineered feature against the label, grouped by the annotation it renders into. + + The caller reads it two ways: an annotation with no feature above 0.5 carries no signal the model could use + from this column alone, and an annotation near 1.0 on training data is a leak to audit before it is trusted.""" + import pandas as pd + + from ..stats import separation + + columns = [name for name in ENGINEERED if name in frame.columns] + table = separation(frame, label, columns) + table["annotation"] = [ENGINEERED[name][1] for name in table["column"]] + best = ( + table.dropna(subset=["auc"]) + .assign(distance=lambda t: (t["auc"] - 0.5).abs()) + .sort_values("distance", ascending=False) + .groupby("annotation", as_index=False) + .first()[["annotation", "column", "auc", "n"]] + .rename(columns={"column": "strongest_feature"}) + ) + return pd.DataFrame(best).sort_values("annotation").reset_index(drop=True) diff --git a/src/openultrasast/stats.py b/src/openultrasast/stats.py new file mode 100644 index 0000000..5543e3f --- /dev/null +++ b/src/openultrasast/stats.py @@ -0,0 +1,186 @@ +"""Statistics the research side reports with: intervals, paired comparisons, grouped folds, feature separation. + +One place, one convention. Wilson intervals come from statsmodels, exact McNemar from statsmodels, the bootstrap +is a repository-clustered percentile bootstrap with linear-interpolation quantiles, folds come from scikit-learn's +grouped splitters, and separation is the area under the ROC curve of one column against the label. The libraries +are imported inside the functions so the core install (PyYAML only) never pays for them; the `[hunt]` extra does. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable, Sequence +from dataclasses import dataclass +from typing import Any + +RESAMPLES = 2000 +ALPHA = 0.05 + + +def wilson(successes: int, n: int, *, alpha: float = ALPHA) -> tuple[float, float] | None: + """The Wilson score interval of a proportion; None when there is no denominator.""" + if n <= 0: + return None + from statsmodels.stats.proportion import proportion_confint + + low, high = proportion_confint(successes, n, alpha=alpha, method="wilson") + return float(low), float(high) + + +@dataclass(frozen=True) +class McNemar: + """Exact McNemar on paired binary outcomes: `b` units where only A succeeds, `c` where only B does.""" + + b: int + c: int + concordant: int + p_value: float + + +def mcnemar_exact(a: Sequence[bool], b: Sequence[bool]) -> McNemar: + """Two-sided exact McNemar (binomial on the discordant pairs); p = 1 when nothing is discordant.""" + if len(a) != len(b): + raise ValueError("paired outcomes must have the same length") + only_a = sum(1 for x, y in zip(a, b, strict=True) if x and not y) + only_b = sum(1 for x, y in zip(a, b, strict=True) if y and not x) + concordant = len(a) - only_a - only_b + if only_a + only_b == 0: + return McNemar(only_a, only_b, concordant, 1.0) + from statsmodels.stats.contingency_tables import mcnemar + + result = mcnemar([[concordant, only_a], [only_b, 0]], exact=True) + return McNemar(only_a, only_b, concordant, float(min(1.0, result.pvalue))) + + +def cluster_resample(groups: Sequence[str], rng: Any) -> list[tuple[int, int]]: + """Draw groups with replacement; return (draw, row index) pairs so a repeated group keeps its own identity.""" + names = sorted(set(groups)) + members: dict[str, list[int]] = {} + for index, group in enumerate(groups): + members.setdefault(group, []).append(index) + out: list[tuple[int, int]] = [] + for draw, name in enumerate(rng.choice(names, size=len(names), replace=True)): + out.extend((draw, index) for index in members[name]) + return out + + +def _quantiles(values: Sequence[float], alpha: float) -> tuple[float, float]: + import numpy as np + + low, high = np.quantile(np.asarray(values, dtype=float), [alpha / 2, 1 - alpha / 2]) + return float(low), float(high) + + +def cluster_bootstrap( + frame: Any, group: str, stat: Callable[[Any], float | None], *, resamples: int = RESAMPLES, seed: int = 0, alpha: float = ALPHA +) -> tuple[float, float] | None: + """Percentile interval of `stat(frame)` under a bootstrap that resamples whole groups (repositories).""" + import numpy as np + + if len(frame) == 0: + return None + rng = np.random.default_rng(seed) + groups = list(frame[group].astype(str)) + values = [] + for _ in range(resamples): + rows = cluster_resample(groups, rng) + sample = frame.iloc[[index for _, index in rows]].copy() + sample[group] = [f"{draw}:{groups[index]}" for draw, index in rows] + value = stat(sample) + if value is not None: + values.append(float(value)) + return _quantiles(values, alpha) if len(values) >= 2 else None + + +@dataclass(frozen=True) +class Estimate: + """A paired comparison B minus A over units, with its clustered-bootstrap interval.""" + + a: float | None + b: float | None + diff: float | None + ci: tuple[float, float] | None + se: float | None + units: int + groups: int + resamples: int + + +def paired_bootstrap( + frame: Any, + group: str, + a: str, + b: str, + *, + stat: Callable[[Any], float | None] | None = None, + resamples: int = RESAMPLES, + seed: int = 0, + alpha: float = ALPHA, +) -> Estimate: + """B minus A where columns `a` and `b` hold each unit's outcome in both arms; `stat` reduces a column + (mean by default) and may return None when a resample cannot be scored.""" + import numpy as np + + def mean(column: Any) -> float | None: + return None if len(column) == 0 else float(column.mean()) + + reduce = stat or mean + a_all, b_all = reduce(frame[a]), reduce(frame[b]) + diff = None if a_all is None or b_all is None else b_all - a_all + groups = list(frame[group].astype(str)) + rng = np.random.default_rng(seed) + values = [] + for _ in range(resamples if groups else 0): + rows = [index for _, index in cluster_resample(groups, rng)] + sample = frame.iloc[rows] + x, y = reduce(sample[a]), reduce(sample[b]) + if x is not None and y is not None: + values.append(y - x) + ci = se = None + if len(values) >= 2: + ci = _quantiles(values, alpha) + se = float(np.std(values, ddof=1)) + return Estimate(a_all, b_all, diff, ci, se, len(frame), len(set(groups)), len(values)) + + +def group_folds(frame: Any, group: str, *, k: int = 5, seed: int = 0) -> list[tuple[list[int], list[int]]]: + """K folds that never split a group across train and test (scikit-learn's GroupKFold, shuffled).""" + from sklearn.model_selection import GroupKFold + + groups = list(frame[group].astype(str)) + if len(set(groups)) < k: + raise ValueError(f"{len(set(groups))} groups cannot make {k} folds") + splitter = GroupKFold(n_splits=k, shuffle=True, random_state=seed) + return [(list(map(int, train)), list(map(int, test))) for train, test in splitter.split(frame, groups=groups)] + + +def leave_one_group_out(frame: Any, group: str) -> list[tuple[list[int], list[int]]]: + from sklearn.model_selection import LeaveOneGroupOut + + groups = list(frame[group].astype(str)) + return [(list(map(int, train)), list(map(int, test))) for train, test in LeaveOneGroupOut().split(frame, groups=groups)] + + +def separation(frame: Any, label: str, columns: Iterable[str]) -> Any: + """How well each numeric column alone separates label 1 from label 0: area under the ROC curve, the number + of rows it was computed on, and a note when it could not be. 0.5 is no signal; near 0 or 1 is a leak or a + strong feature, which the caller decides.""" + import pandas as pd + from sklearn.metrics import roc_auc_score + + rows = [] + labelled = frame[frame[label].isin([0, 1])] + for column in columns: + series = pd.to_numeric(labelled[column], errors="coerce") if column in labelled else None + if series is None: + rows.append({"column": column, "auc": None, "n": 0, "note": "missing"}) + continue + mask = series.notna() + y = labelled.loc[mask, label].astype(int) + if mask.sum() == 0 or y.nunique() < 2: + rows.append({"column": column, "auc": None, "n": int(mask.sum()), "note": "one class or empty"}) + continue + if series[mask].nunique() < 2: + rows.append({"column": column, "auc": 0.5, "n": int(mask.sum()), "note": "constant"}) + continue + rows.append({"column": column, "auc": float(roc_auc_score(y, series[mask])), "n": int(mask.sum()), "note": ""}) + return pd.DataFrame(rows, columns=["column", "auc", "n", "note"]) diff --git a/tests/test_hunt_features.py b/tests/test_hunt_features.py new file mode 100644 index 0000000..cd463a3 --- /dev/null +++ b/tests/test_hunt_features.py @@ -0,0 +1,101 @@ +"""Engineered features from the feature table, their tie to the annotations, and the separation screen.""" + +from __future__ import annotations + +import pytest + +pd = pytest.importorskip("pandas") +pytest.importorskip("sklearn") + +from openultrasast.hunt import features # noqa: E402 + +ROW = { + "change_id": "c1", + "repo": "r", + "head": "h", + "path": "app.py", + "line": 5, + "kind": "vulnerability", + "fold": 1, + "label": 1, + "family": "injection", + "feature.diff.hunks": [[3, 3], [10, 12]], + "feature.diff.on_changed_line": False, + "feature.diff.distance": 2, + "feature.source.function": "second", + "feature.source.span": [5, 7], + "feature.source.callers": [["web.py", 2], ["web.py", 9], ["cli.py", 1]], + "feature.source.entry_points": [{"path": "web.py", "line": 1, "kind": "route"}], + "feature.source.language": "python", + "feature.semantic.roles": {"input": "source", "eval": "sink", "escape": "sanitizer", "x": "sink"}, + "feature.semantic.def_use": {"binds": [{"line": 6}], "calls": [{"line": 6}, {"line": 7}]}, + "feature.semantic.obligations": [ + {"operation": "auth", "line": 6, "discharged": False}, + {"operation": "csrf", "line": 7, "discharged": True}, + ], + "feature.taxonomy.families": [{"id": "injection", "definition": "x"}], + "provenance.diff.ok": True, + "provenance.source.ok": True, + "provenance.semantic.ok": True, + "provenance.taxonomy.ok": True, +} + + +def test_feature_frame_engineers_counts_distances_and_flags(): + absent = { + **ROW, + "change_id": "c2", + "label": 0, + "feature.source.function": "", + "feature.source.span": None, + "feature.semantic.roles": None, + "feature.semantic.def_use": None, + "feature.semantic.obligations": None, + "provenance.semantic.ok": False, + } + frame = features.feature_frame([ROW, absent]) + first = frame.iloc[0] + assert first["diff.distance"] == 2 and first["diff.on_changed_line"] == 0 + assert first["diff.hunk_count"] == 2 and first["diff.hunk_lines"] == 4 + assert first["source.span_lines"] == 3 and first["source.is_module"] == 0 + assert first["source.caller_count"] == 3 and first["source.caller_files"] == 2 and first["source.entry_point_count"] == 1 + assert (first["semantic.sources"], first["semantic.sinks"], first["semantic.sanitizers"]) == (1, 2, 1) + assert first["semantic.binds"] == 1 and first["semantic.calls"] == 2 + assert first["semantic.obligations"] == 2 and first["semantic.obligations_undischarged"] == 1 + second = frame.iloc[1] + assert second["source.is_module"] == 1 and pd.isna(second["source.span_lines"]) + assert pd.isna(second["semantic.sinks"]) and second["ok.semantic"] is False or second["ok.semantic"] == False # noqa: E712 + assert list(frame["label"]) == [1, 0] + + +def test_every_engineered_feature_maps_to_an_annotation_and_masks_follow(): + from openultrasast.hunt.program import INPUTS + + inputs = {name for name, _, _ in INPUTS} + assert all(annotation in inputs for _, annotation in features.ENGINEERED.values()) + assert set(features.features_of("roles")) == {"semantic.sources", "semantic.sinks", "semantic.sanitizers"} + assert features.annotations_of(["semantic"]) == {"roles", "def_use", "obligations"} + assert features.annotations_of(["feature.source.callers"]) == {"callers"} + assert features.annotations_of(["hunks"]) == {"hunks"} + assert features.annotations_of([]) == set() + + +def test_screen_reports_the_strongest_feature_per_annotation(): + rows = [] + for i in range(12): + label = i % 2 + rows.append( + { + **ROW, + "change_id": f"c{i}", + "label": label, + "feature.diff.distance": 0 if label else 30, + "feature.source.callers": [["a.py", 1]] * (3 if label else 1), + "feature.semantic.roles": {"eval": "sink"}, + } + ) + frame = features.feature_frame(rows) + out = features.screen(frame).set_index("annotation") + assert out.loc["hunks", "strongest_feature"] == "diff.distance" and out.loc["hunks", "auc"] == 0.0 + assert out.loc["callers", "strongest_feature"] == "source.caller_count" and out.loc["callers", "auc"] == 1.0 + assert out.loc["roles", "auc"] == 0.5 # constant: no signal diff --git a/tests/test_stats.py b/tests/test_stats.py new file mode 100644 index 0000000..3d9a338 --- /dev/null +++ b/tests/test_stats.py @@ -0,0 +1,104 @@ +"""The shared statistics against the formulas the benchmark and the learn harness used.""" + +from __future__ import annotations + +import math + +import pytest + +pd = pytest.importorskip("pandas") +pytest.importorskip("statsmodels") +pytest.importorskip("sklearn") + +from openultrasast import stats # noqa: E402 + + +def test_wilson_matches_the_benchmark_scorer_to_three_decimals(): + from benchmarks.unseen.score import wilson as scorer_wilson + + for successes, n in ((6, 12), (5, 12), (130, 239), (1, 149), (0, 12), (221, 239)): + low, high = stats.wilson(successes, n) + old_low, old_high = scorer_wilson(successes, n) + assert abs(low - old_low) < 1e-3 and abs(high - old_high) < 1e-3, (successes, n) + assert stats.wilson(0, 0) is None + low, high = stats.wilson(6, 12) + assert round(low, 3) == 0.254 and round(high, 3) == 0.746 # the pilot's hunter recall interval + + +def test_exact_mcnemar_matches_the_learn_harness(): + import random + + from openultrasast.learn.experiments import exact_mcnemar + + rng = random.Random(1) + for _ in range(20): + n = rng.randint(1, 40) + a = [rng.random() < 0.5 for _ in range(n)] + b = [rng.random() < 0.5 for _ in range(n)] + ours = stats.mcnemar_exact(a, b) + theirs = exact_mcnemar(list(zip(a, b, strict=True))) + assert (ours.b, ours.c, ours.concordant) == (theirs.b, theirs.c, theirs.concordant) + assert math.isclose(ours.p_value, theirs.p_value, abs_tol=1e-9), (a, b) + assert stats.mcnemar_exact([True, False], [True, False]).p_value == 1.0 + with pytest.raises(ValueError): + stats.mcnemar_exact([True], [True, False]) + + +def _frame(seed=0, groups=8, per=6): + import numpy as np + + rng = np.random.default_rng(seed) + rows = [] + for g in range(groups): + for _ in range(per): + a = rng.random() < 0.4 + rows.append({"repo": f"r{g}", "a": int(a), "b": int(a or rng.random() < 0.3), "label": int(rng.random() < 0.5)}) + return pd.DataFrame(rows) + + +def test_cluster_bootstrap_is_seeded_and_brackets_the_estimate(): + frame = _frame() + stat = lambda f: float(f["a"].mean()) # noqa: E731 + first = stats.cluster_bootstrap(frame, "repo", stat, resamples=300, seed=3) + second = stats.cluster_bootstrap(frame, "repo", stat, resamples=300, seed=3) + assert first == second and first[0] <= stat(frame) <= first[1] + assert stats.cluster_bootstrap(frame.iloc[0:0], "repo", stat) is None + + +def test_paired_bootstrap_reports_b_minus_a_with_an_interval(): + frame = _frame() + estimate = stats.paired_bootstrap(frame, "repo", "a", "b", resamples=300, seed=1) + assert estimate.diff == pytest.approx(frame["b"].mean() - frame["a"].mean()) + assert estimate.ci is not None and estimate.ci[0] <= estimate.diff <= estimate.ci[1] + assert estimate.units == len(frame) and estimate.groups == 8 and estimate.resamples == 300 + assert estimate.se is not None and estimate.se > 0 + + +def test_group_folds_never_split_a_repository(): + frame = _frame(groups=10) + folds = stats.group_folds(frame, "repo", k=5, seed=0) + assert len(folds) == 5 + for train, test in folds: + assert not set(frame.iloc[train]["repo"]) & set(frame.iloc[test]["repo"]) + assert sorted(train + test) == list(range(len(frame))) + assert len(stats.leave_one_group_out(frame, "repo")) == 10 + with pytest.raises(ValueError): + stats.group_folds(_frame(groups=3), "repo", k=5) + + +def test_separation_scores_each_column_alone(): + frame = pd.DataFrame( + { + "label": [1, 1, 1, 0, 0, 0, None], + "perfect": [5, 6, 7, 1, 2, 3, 9], + "noise": [1, 1, 1, 1, 1, 1, 1], + "inverse": [1, 2, 3, 5, 6, 7, 0], + "sparse": [None, None, None, None, None, None, 1], + } + ) + out = stats.separation(frame, "label", ["perfect", "noise", "inverse", "sparse", "missing"]).set_index("column") + assert out.loc["perfect", "auc"] == 1.0 and out.loc["perfect", "n"] == 6 + assert out.loc["inverse", "auc"] == 0.0 + assert out.loc["noise", "auc"] == 0.5 and out.loc["noise", "note"] == "constant" + assert out.loc["sparse", "auc"] is None or pd.isna(out.loc["sparse", "auc"]) + assert out.loc["missing", "note"] == "missing"