Add bucket metrics and contrast delta analysis

README.md

@@ -32,11 +32,16 @@ prompts/                    # your JSONL corpora
 
 1. Activation census extraction (`mlp`, `gate`, `up` by default)
 2. Finalize `.npz` chunks
-3. Per-layer analysis: taxonomy, heatmap, F-stat separation, co-activation, code cross-reference
+3. Per-layer analysis: taxonomy, heatmap, F-stat separation, bucket quality, contrast deltas, co-activation, code cross-reference
 4. OV-circuit spectral analysis: SVD over `W_V @ W_O` per head
 5. Compliance-behaviour axis extraction (positive vs negative corpus)
 6. Atlas build + SQLite index
 
+Per-component analysis now writes:
+
+- `l<N>_<component>_bucket_metrics.json`: dominant bucket, entropy, eta-squared, and bucket-quality score per feature.
+- `l<N>_<component>_contrast_delta.json`: feature deltas for explicit `contrast_pair_id` pairs when the corpus provides exactly two rows per pair.
+
 ## Dev-mode setup
 
 Inside the HF Space terminal:

qwip_atlas/analyze_layers.py

@@ -423,6 +423,183 @@ def compute_separation_scores(A: np.ndarray, buckets: list[str]) -> np.ndarray:
     return fstat.astype(np.float32)
 
 
+def compute_feature_bucket_metrics(A: np.ndarray, buckets: list[str]) -> list[dict]:
+    """
+    Per-feature bucket quality metrics.
+
+    These make bucket behavior visible instead of hiding it behind a single
+    F-statistic:
+      - dominant_bucket: bucket with highest mean activation
+      - dominance_margin: top bucket mean minus runner-up
+      - bucket_entropy: normalized entropy over positive bucket means
+      - eta_squared: fraction of activation variance explained by buckets
+      - bucket_quality: eta_squared weighted by low entropy
+    """
+    bucket_arr = np.asarray(buckets)
+    unique = sorted(set(buckets))
+    n_features, n_prompts = A.shape
+    if n_prompts != len(buckets):
+        raise ValueError(f"A has {n_prompts} prompts, but buckets has {len(buckets)} entries")
+
+    overall_mean = A.mean(axis=1, keepdims=True)
+    bucket_means = np.zeros((n_features, len(unique)), dtype=np.float64)
+    bucket_counts = np.zeros(len(unique), dtype=np.int64)
+
+    for j, bucket in enumerate(unique):
+        mask = bucket_arr == bucket
+        bucket_counts[j] = int(mask.sum())
+        if bucket_counts[j]:
+            bucket_means[:, j] = A[:, mask].mean(axis=1)
+
+    between_var = np.zeros(n_features, dtype=np.float64)
+    for j in range(len(unique)):
+        if bucket_counts[j] < 1:
+            continue
+        delta = bucket_means[:, j] - overall_mean.squeeze(-1)
+        between_var += bucket_counts[j] * delta**2
+
+    centered = A - overall_mean
+    total_var = np.sum(centered**2, axis=1)
+    eta_squared = np.divide(
+        between_var,
+        total_var,
+        out=np.zeros_like(between_var),
+        where=total_var > 1e-10,
+    )
+
+    # Entropy over positive bucket means. If all bucket means are non-positive,
+    # fall back to absolute means so inhibitory/signed features still get shape.
+    positive = np.clip(bucket_means, 0, None)
+    fallback = np.abs(bucket_means)
+    weights = np.where(positive.sum(axis=1, keepdims=True) > 1e-10, positive, fallback)
+    weight_sums = weights.sum(axis=1, keepdims=True)
+    probs = np.divide(weights, weight_sums, out=np.zeros_like(weights), where=weight_sums > 1e-10)
+    log_probs = np.zeros_like(probs)
+    nz = probs > 0
+    log_probs[nz] = np.log(probs[nz])
+    raw_entropy = -(probs * log_probs).sum(axis=1)
+    norm = np.log(len(unique)) if len(unique) > 1 else 1.0
+    bucket_entropy = raw_entropy / norm if norm > 0 else raw_entropy
+
+    top_idx = np.argmax(bucket_means, axis=1)
+    order_desc = np.argsort(-bucket_means, axis=1)
+    second_idx = order_desc[:, 1] if len(unique) > 1 else top_idx
+
+    metrics = []
+    for i in range(n_features):
+        top = int(top_idx[i])
+        second = int(second_idx[i])
+        margin = float(bucket_means[i, top] - bucket_means[i, second]) if len(unique) > 1 else 0.0
+        entropy = float(bucket_entropy[i])
+        eta = float(eta_squared[i])
+        metrics.append({
+            "feature": i,
+            "dominant_bucket": unique[top],
+            "dominant_bucket_mean": float(bucket_means[i, top]),
+            "runner_up_bucket": unique[second],
+            "runner_up_bucket_mean": float(bucket_means[i, second]),
+            "dominance_margin": margin,
+            "bucket_entropy": entropy,
+            "eta_squared": eta,
+            "bucket_quality": float(eta * (1.0 - entropy)),
+            "bucket_means": {bucket: float(bucket_means[i, j]) for j, bucket in enumerate(unique)},
+        })
+    return metrics
+
+
+def compute_contrast_delta_scores(A: np.ndarray, records: list[dict]) -> dict:
+    """
+    Rank features by consistent activation change across explicit contrast pairs.
+
+    Pairs are not inferred from prompt text. The corpus must provide
+    ``contrast_pair_id`` metadata; exactly two rows per pair are required.
+    Within each pair, rows are ordered by bucket name and the delta is:
+
+        activation(second bucket lexically) - activation(first bucket lexically)
+
+    The report includes the direction labels so downstream readers can see what
+    a positive delta means.
+    """
+    from collections import defaultdict
+
+    by_pair: dict[str, list[int]] = defaultdict(list)
+    for i, record in enumerate(records):
+        pair_id = record.get("contrast_pair_id")
+        if pair_id:
+            by_pair[str(pair_id)].append(i)
+
+    pair_rows = []
+    skipped = []
+    directions = []
+    for pair_id in sorted(by_pair):
+        idxs = by_pair[pair_id]
+        if len(idxs) != 2:
+            skipped.append(pair_id)
+            continue
+        ordered = sorted(idxs, key=lambda idx: (
+            str(records[idx].get("bucket", "")),
+            str(records[idx].get("id", idx)),
+        ))
+        left, right = ordered
+        left_bucket = str(records[left].get("bucket", ""))
+        right_bucket = str(records[right].get("bucket", ""))
+        directions.append((left_bucket, right_bucket))
+        pair_rows.append({
+            "pair_id": pair_id,
+            "left_index": int(left),
+            "right_index": int(right),
+            "left_bucket": left_bucket,
+            "right_bucket": right_bucket,
+            "left_id": records[left].get("id"),
+            "right_id": records[right].get("id"),
+        })
+
+    if not pair_rows:
+        return {
+            "n_pairs": 0,
+            "direction": "lexical_bucket_order",
+            "direction_counts": {},
+            "skipped_pair_ids": skipped,
+            "pairs": [],
+            "features": [],
+        }
+
+    left_idx = np.array([p["left_index"] for p in pair_rows], dtype=np.int64)
+    right_idx = np.array([p["right_index"] for p in pair_rows], dtype=np.int64)
+    deltas = A[:, right_idx] - A[:, left_idx]
+    mean_delta = deltas.mean(axis=1)
+    abs_mean_delta = np.abs(mean_delta)
+    std_delta = deltas.std(axis=1)
+    positive_fraction = (deltas > 0).mean(axis=1)
+    consistency = np.maximum(positive_fraction, 1.0 - positive_fraction)
+    effect_score = abs_mean_delta * consistency
+
+    from collections import Counter
+    direction_counts = Counter(f"{left}->{right}" for left, right in directions)
+
+    order = np.argsort(effect_score)[::-1]
+    features = []
+    for idx in order:
+        features.append({
+            "feature": int(idx),
+            "mean_delta": float(mean_delta[idx]),
+            "abs_mean_delta": float(abs_mean_delta[idx]),
+            "std_delta": float(std_delta[idx]),
+            "positive_fraction": float(positive_fraction[idx]),
+            "consistency": float(consistency[idx]),
+            "effect_score": float(effect_score[idx]),
+        })
+
+    return {
+        "n_pairs": len(pair_rows),
+        "direction": "lexical_bucket_order",
+        "direction_counts": dict(direction_counts),
+        "skipped_pair_ids": skipped,
+        "pairs": pair_rows,
+        "features": features,
+    }
+
+
 def compute_silhouette_scores(A: np.ndarray, buckets: list[str]) -> np.ndarray:
     """
     Per-feature silhouette score (kept for backward comparison).
@@ -461,6 +638,44 @@ def print_top_separators(scores: np.ndarray, n: int = 20,
         print(f"  {idx:>8}  {scores[idx]:>10.4f}  {cls:<30}  {act_rate:>9.3f}")
 
 
+def print_top_bucket_metrics(metrics: list[dict], n: int = 20) -> None:
+    top = sorted(metrics, key=lambda row: row["bucket_quality"], reverse=True)[:n]
+    print(f"\nTop {n} bucket-clean features (eta^2 * low entropy):")
+    print(
+        f"  {'feature':>8}  {'quality':>10}  {'eta^2':>8}  "
+        f"{'entropy':>8}  {'dominant':<24}  {'margin':>9}"
+    )
+    for row in top:
+        print(
+            f"  {row['feature']:>8}  {row['bucket_quality']:>10.4f}  "
+            f"{row['eta_squared']:>8.4f}  {row['bucket_entropy']:>8.4f}  "
+            f"{row['dominant_bucket']:<24}  {row['dominance_margin']:>9.4f}"
+        )
+
+
+def print_contrast_delta_summary(report: dict, n: int = 20) -> None:
+    print(f"\nContrast pairs: {report['n_pairs']} usable")
+    if report["skipped_pair_ids"]:
+        print(f"  skipped pair ids (not exactly 2 rows): {report['skipped_pair_ids'][:10]}")
+    if report["direction_counts"]:
+        print("  positive delta direction counts:")
+        for direction, count in sorted(report["direction_counts"].items()):
+            print(f"    {direction}: {count}")
+    if not report["features"]:
+        return
+    print(f"\nTop {n} contrast-shifting features:")
+    print(
+        f"  {'feature':>8}  {'effect':>10}  {'mean_delta':>11}  "
+        f"{'consistency':>11}  {'pos_frac':>9}"
+    )
+    for row in report["features"][:n]:
+        print(
+            f"  {row['feature']:>8}  {row['effect_score']:>10.4f}  "
+            f"{row['mean_delta']:>11.4f}  {row['consistency']:>11.4f}  "
+            f"{row['positive_fraction']:>9.4f}"
+        )
+
+
 # ---------------------------------------------------------------------------
 # Phase 6: Cross-reference with capability damage
 # ---------------------------------------------------------------------------
@@ -655,11 +870,27 @@ def analyze_one(name: str, A: np.ndarray, records: list[dict],
     np.save(out / f"{prefix}_separation_scores.npy", scores)
     print_top_separators(scores, n=20, classifications=classifications, A=A)
 
-    # Phase 6: Code cross-reference
-    print(f"\n--- {name}: Phase 6: Code feature cross-reference ---")
+    # Phase 6: Bucket quality metrics
+    print(f"\n--- {name}: Phase 6: Bucket quality metrics ---")
+    bucket_metrics = compute_feature_bucket_metrics(A, buckets)
+    write_json(out / f"{prefix}_bucket_metrics.json", bucket_metrics)
+    print_top_bucket_metrics(bucket_metrics)
+
+    # Phase 7: Contrast pair deltas
+    print(f"\n--- {name}: Phase 7: Contrast pair deltas ---")
+    contrast_report = compute_contrast_delta_scores(A, records)
+    if contrast_report["n_pairs"]:
+        write_json(out / f"{prefix}_contrast_delta.json", contrast_report)
+    print_contrast_delta_summary(contrast_report)
+
+    # Phase 8: Code cross-reference
+    print(f"\n--- {name}: Phase 8: Code feature cross-reference ---")
     code_analysis = analyze_code_neurons(A, buckets, scores, code_bucket=code_bucket)
     write_json(out / f"{prefix}_code_analysis.json", code_analysis)
 
+    best_bucket = max(bucket_metrics, key=lambda row: row["bucket_quality"]) if bucket_metrics else {}
+    best_contrast = contrast_report["features"][0] if contrast_report["features"] else {}
+
     # Per-component summary stats
     from collections import Counter
     def _norm_cls(c):
@@ -672,6 +903,14 @@ def analyze_one(name: str, A: np.ndarray, records: list[dict],
         "taxonomy":       dict(tax_counts),
         "top_sep_score":  float(np.max(scores)),
         "mean_sep_score": float(np.mean(scores)),
+        "top_bucket_quality": float(best_bucket.get("bucket_quality", 0.0)),
+        "top_bucket_feature": int(best_bucket.get("feature", 0)) if best_bucket else None,
+        "top_bucket": best_bucket.get("dominant_bucket"),
+        "bucket_entropy_at_top": float(best_bucket.get("bucket_entropy", 0.0)) if best_bucket else None,
+        "eta_squared_at_top": float(best_bucket.get("eta_squared", 0.0)) if best_bucket else None,
+        "n_contrast_pairs": int(contrast_report["n_pairs"]),
+        "top_contrast_feature": int(best_contrast.get("feature", 0)) if best_contrast else None,
+        "top_contrast_effect": float(best_contrast.get("effect_score", 0.0)) if best_contrast else None,
         "n_active":       int((A > ACTIVATION_THRESHOLD).any(axis=1).sum()),
         "mean":           float(A.mean()),
         "std":            float(A.std()),
@@ -719,6 +958,20 @@ def print_comparison(summaries: list[dict], layer: int):
     row("  best feature",    lambda s: f"{s['top_sep_score']:.4f}")
     row("  mean (all feats)",lambda s: f"{s['mean_sep_score']:.4f}")
 
+    print(f"  {'-'*28}" + "".join("-" * col_width for _ in summaries))
+    print("  Bucket quality (eta^2 * low entropy):")
+    row("  best quality",    lambda s: f"{s.get('top_bucket_quality', 0.0):.4f}")
+    row("  best bucket",     lambda s: str(s.get("top_bucket") or "-")[:10])
+    row("  best feature",    lambda s: s.get("top_bucket_feature") if s.get("top_bucket_feature") is not None else "-")
+    row("  eta^2 at top",    lambda s: f"{s.get('eta_squared_at_top'):.4f}" if s.get("eta_squared_at_top") is not None else "-")
+    row("  entropy at top",  lambda s: f"{s.get('bucket_entropy_at_top'):.4f}" if s.get("bucket_entropy_at_top") is not None else "-")
+
+    print(f"  {'-'*28}" + "".join("-" * col_width for _ in summaries))
+    print("  Contrast pairs (explicit contrast_pair_id only):")
+    row("  usable pairs",    lambda s: s.get("n_contrast_pairs", 0))
+    row("  best feature",    lambda s: s.get("top_contrast_feature") if s.get("top_contrast_feature") is not None else "-")
+    row("  best effect",     lambda s: f"{s.get('top_contrast_effect'):.4f}" if s.get("top_contrast_effect") is not None else "-")
+
     print(f"  {'-'*28}" + "".join("-" * col_width for _ in summaries))
     print("  Code cross-reference (top 30 code-preferring):")
     row("  entangled",       lambda s: s["code_entangled"])
@@ -733,6 +986,13 @@ def print_comparison(summaries: list[dict], layer: int):
           f"(n={best_specific['taxonomy'].get('specific', 0)})")
     print(f"    Highest single-feature separation: {best_sep['name']}  "
           f"(score={best_sep['top_sep_score']:.4f})")
+    best_bucket = max(summaries, key=lambda s: s.get("top_bucket_quality", 0.0))
+    print(f"    Cleanest bucket signal: {best_bucket['name']}  "
+          f"(bucket={best_bucket.get('top_bucket')}, quality={best_bucket.get('top_bucket_quality', 0.0):.4f})")
+    if any(s.get("n_contrast_pairs", 0) for s in summaries):
+        best_contrast = max(summaries, key=lambda s: s.get("top_contrast_effect", 0.0) or 0.0)
+        print(f"    Strongest contrast delta: {best_contrast['name']}  "
+              f"(effect={best_contrast.get('top_contrast_effect', 0.0):.4f})")
 
 
 def main():
@@ -846,6 +1106,7 @@ def main():
         prefix = f"l{args.layer}_{s['name']}"
         for suffix in ["neuron_taxonomy.json", "census_heatmap.png",
                        "coactivation_pairs.json", "separation_scores.npy",
+                       "bucket_metrics.json", "contrast_delta.json",
                        "code_analysis.json"]:
             p = out / f"{prefix}_{suffix}"
             if p.exists():

qwip_atlas/build_atlas.py

@@ -213,6 +213,8 @@ def cmd_merge_layer(args):
             "taxonomy.json":      analysis_dir / f"{prefix}neuron_taxonomy.json",
             "separation.npy":     analysis_dir / f"{prefix}separation_scores.npy",
             "coactivation.json":  analysis_dir / f"{prefix}coactivation_pairs.json",
+            "bucket_metrics.json":analysis_dir / f"{prefix}bucket_metrics.json",
+            "contrast_delta.json":analysis_dir / f"{prefix}contrast_delta.json",
             "code_analysis.json": analysis_dir / f"{prefix}code_analysis.json",
             "census_heatmap.png": analysis_dir / f"{prefix}census_heatmap.png",
         }
@@ -245,6 +247,29 @@ def cmd_merge_layer(args):
             "code_selective":code.get("selective_count", 0),
             "files":         copied,
         }
+        bucket_metrics_path = cdir / "bucket_metrics.json"
+        if bucket_metrics_path.exists():
+            bucket_metrics = read_json(bucket_metrics_path)
+            if bucket_metrics:
+                top_bucket = max(bucket_metrics, key=lambda row: row.get("bucket_quality", 0.0))
+                summary.update({
+                    "bucket_quality_top": top_bucket.get("bucket_quality"),
+                    "bucket_quality_top_idx": top_bucket.get("feature"),
+                    "bucket_quality_top_bucket": top_bucket.get("dominant_bucket"),
+                    "bucket_quality_top_entropy": top_bucket.get("bucket_entropy"),
+                    "bucket_quality_top_eta_squared": top_bucket.get("eta_squared"),
+                })
+        contrast_path = cdir / "contrast_delta.json"
+        if contrast_path.exists():
+            contrast = read_json(contrast_path)
+            features = contrast.get("features") or []
+            top_contrast = features[0] if features else {}
+            summary.update({
+                "contrast_pairs": contrast.get("n_pairs", 0),
+                "contrast_top_idx": top_contrast.get("feature"),
+                "contrast_top_effect": top_contrast.get("effect_score"),
+                "contrast_top_mean_delta": top_contrast.get("mean_delta"),
+            })
         write_json(cdir / "summary.json", summary)
 
     # 3. Per-head ingest

tests/test_analyze_layers_metrics.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import numpy as np
+
+from qwip_atlas.analyze_layers import (
+    compute_contrast_delta_scores,
+    compute_feature_bucket_metrics,
+)
+
+
+def test_feature_bucket_metrics_identifies_clean_dominant_bucket() -> None:
+    A = np.array(
+        [
+            [5.0, 4.0, 0.0, 0.0],
+            [1.0, 1.0, 1.0, 1.0],
+            [0.0, 0.0, 3.0, 4.0],
+        ],
+        dtype=np.float32,
+    )
+    buckets = ["auth", "auth", "corp", "corp"]
+
+    metrics = compute_feature_bucket_metrics(A, buckets)
+
+    assert metrics[0]["dominant_bucket"] == "auth"
+    assert metrics[0]["dominance_margin"] > 3.0
+    assert metrics[0]["bucket_entropy"] < 0.1
+    assert metrics[0]["eta_squared"] > 0.9
+
+    assert metrics[1]["dominant_bucket"] == "auth"
+    assert metrics[1]["dominance_margin"] == 0.0
+    assert metrics[1]["bucket_entropy"] > 0.9
+    assert metrics[1]["eta_squared"] == 0.0
+
+    assert metrics[2]["dominant_bucket"] == "corp"
+    assert metrics[2]["dominance_margin"] > 3.0
+    assert metrics[2]["bucket_entropy"] < 0.1
+    assert metrics[2]["eta_squared"] > 0.9
+
+
+def test_contrast_delta_scores_use_exact_two_row_pairs() -> None:
+    A = np.array(
+        [
+            [1.0, 2.0, 2.0, 3.0, 100.0],
+            [5.0, 2.0, 0.0, 1.0, 100.0],
+            [2.0, 0.0, 5.0, 2.0, 100.0],
+        ],
+        dtype=np.float32,
+    )
+    records = [
+        {"id": "a1", "contrast_pair_id": "p1", "bucket": "auth"},
+        {"id": "c1", "contrast_pair_id": "p1", "bucket": "corp"},
+        {"id": "a2", "contrast_pair_id": "p2", "bucket": "auth"},
+        {"id": "c2", "contrast_pair_id": "p2", "bucket": "corp"},
+        {"id": "orphan", "contrast_pair_id": "bad", "bucket": "other"},
+    ]
+
+    report = compute_contrast_delta_scores(A, records)
+
+    assert report["n_pairs"] == 2
+    assert report["skipped_pair_ids"] == ["bad"]
+    assert report["direction"] == "lexical_bucket_order"
+
+    by_feature = {row["feature"]: row for row in report["features"]}
+    assert by_feature[0]["mean_delta"] == 1.0
+    assert by_feature[0]["positive_fraction"] == 1.0
+    assert by_feature[1]["mean_delta"] == -1.0
+    assert by_feature[1]["positive_fraction"] == 0.5
+    assert by_feature[2]["mean_delta"] == -2.5
+    assert by_feature[2]["positive_fraction"] == 0.0