gauge/techniques/t_diff_means.py

"""
Diff-of-means concept vectors. Attempts R (recovery), F (faithfulness), X (control).

For a binary target, the concept vector is the difference of class means:
    v = mean(z | y=1) - mean(z | y=0).
- R: classify by projection onto v (threshold at the midpoint).
- F: name v as the causal direction for y_causal. NOTE this is the interesting
     case — the diff-of-means picks up BOTH the causal axis (v_causal, which Rd
     reads) AND any *correlated* readable-not-causal contamination (v_ro). Whether
     steering v actually moves Rd is exactly what F1 tests causally.
- X: to ADD concept m, estimate u_m as the diff-of-means between items with m
     active vs not (needs labeled probe items); to FLIP the label, push along the
     causal diff-of-means. The challenge verifies via Rd.

We make X *causally honest*: we estimate the edit budget so Rd actually flips,
using the technique's own probe data and the synth readout is NOT consulted (we
calibrate budget from the projection statistics of the labeled train set).
"""

from __future__ import annotations

import numpy as np

from .base import Technique


def _diff_means(Z, y):
    y = np.asarray(y)
    pos = Z[y == 1].mean(0) if (y == 1).any() else np.zeros(Z.shape[1])
    neg = Z[y == 0].mean(0) if (y == 0).any() else np.zeros(Z.shape[1])
    return pos - neg


class DiffMeans(Technique):
    NAME = "diff_means"
    AXES = "RFX"

    # ---- R ----
    def recover(self, encoder, train, test, target):
        Ztr = encoder.encode(train["items"]); ytr = np.asarray(train["labels"])
        Zte = encoder.encode(test["items"]); yte = np.asarray(test["labels"])
        if target["kind"] == "regression":
            # diff-of-means is for binary; regress on projection onto top corr dir
            # (correlate z dims with y, use as direction) — a linear surrogate.
            zc = Ztr - Ztr.mean(0)
            w = zc.T @ (ytr - ytr.mean())
            w = w / (np.linalg.norm(w) + 1e-9)
            a = Ztr @ w
            # least-squares scale/offset
            A = np.vstack([a, np.ones_like(a)]).T
            coef, *_ = np.linalg.lstsq(A, ytr, rcond=None)
            pred = (Zte @ w) * coef[0] + coef[1]
            tol = target.get("tol", 0.5)
            acc = float(np.mean(np.abs(pred - yte) <= tol))
            r2 = max(0.0, 1 - np.var(pred - yte) / (np.var(yte) + 1e-9))
            return r2 > 0.2, {"pred": pred.tolist()}, float(np.clip(r2, 0, 1))
        v = _diff_means(Ztr, ytr); v = v / (np.linalg.norm(v) + 1e-9)
        atr = Ztr @ v
        thr = 0.5 * (atr[ytr == 1].mean() + atr[ytr == 0].mean())
        pred = (Zte @ v > thr).astype(int)
        cv = float(np.mean((atr > thr).astype(int) == ytr))
        return cv > 0.6, {"pred": pred.tolist()}, float(np.clip(cv, 0, 1))

    # ---- F: name the causal direction for y_causal ----
    def faithfulness(self, encoder, train, test, target):
        Ztr = encoder.encode(train["items"]); ytr = np.asarray(train["labels"])
        v = _diff_means(Ztr, ytr)
        v = v / (np.linalg.norm(v) + 1e-9)
        # confidence from how separable the classes are along v
        a = Ztr @ v
        sep = abs(a[ytr == 1].mean() - a[ytr == 0].mean()) / (a.std() + 1e-9)
        conf = float(np.clip(sep / 3.0, 0, 1))
        return True, {"direction": v}, conf

    # ---- X (SONAR inject variant): diff-of-means concept injection ----
    def _control_inject(self, encoder, train, test, target):
        """Inject a target concept by adding the diff-of-means direction between
        target-concept sentences and the source sentences (constructed probe)."""
        Zsrc_probe = encoder.encode(train["items"])
        Ztgt = encoder.encode(target["target_sents"])
        direction = Ztgt.mean(0) - Zsrc_probe.mean(0)
        direction = direction / (np.linalg.norm(direction) + 1e-9)
        items = test["items"]
        Z = encoder.encode(items)
        # scale to a fraction of the typical embedding norm (known budget)
        scale = float(np.median(np.linalg.norm(Z, axis=1))) * 0.6 * target.get("budget", 1.0)
        deltas = [(scale * direction).tolist() for _ in Z]
        return True, {"deltas": deltas}, 0.7

    # ---- X: targeted edit (add concept + flip label) ----
    def control(self, encoder, train, test, target):
        if target.get("kind") == "inject":
            return self._control_inject(encoder, train, test, target)
        # Build concept + causal direction estimates from labeled PROBE items the
        # technique constructs itself (it may query the encoder).
        synth_K = encoder._synth.K
        M = encoder._synth.M
        rng = np.random.default_rng(0)

        # estimate each concept direction u_m via diff-of-means on probe items.
        n_probe = 120
        # build probe items: half with concept m, half without
        u_hat = {}
        for m in range(M):
            on, off = [], []
            for _ in range(n_probe):
                L = 4
                tok = list(rng.integers(0, synth_K, size=L))
                pos = list(range(L))
                base = {"tokens": tok, "positions": pos, "number": None,
                        "entities": {}, "ro_g": 0.0, "causal_g": 0.0}
                on.append({**base, "concepts": [m]})
                off.append({**base, "concepts": []})
            Zon = encoder.encode(on); Zoff = encoder.encode(off)
            u_hat[m] = (Zon - Zoff).mean(0)

        # estimate causal flip direction via diff-of-means on causal_g sign.
        pos_items, neg_items = [], []
        for _ in range(n_probe):
            L = 4
            tok = list(rng.integers(0, synth_K, size=L)); p = list(range(L))
            b = {"tokens": tok, "positions": p, "number": None, "concepts": [],
                 "entities": {}, "ro_g": 0.0}
            pos_items.append({**b, "causal_g": 1.2})
            neg_items.append({**b, "causal_g": -1.2})
        Zp = encoder.encode(pos_items); Zn = encoder.encode(neg_items)
        v_causal_hat = (Zp - Zn).mean(0)
        v_causal_hat = v_causal_hat / (np.linalg.norm(v_causal_hat) + 1e-9)

        # build deltas per test item
        items = test["items"]; targets = test["targets"]
        Z = encoder.encode(items)
        deltas = []
        for z, tgt in zip(Z, targets):
            m_add = tgt["add_concept"]
            d = np.zeros(z.shape[0])
            # add the concept (unit u_m direction, budget to clear threshold ~1)
            um = u_hat[m_add]; um = um / (np.linalg.norm(um) + 1e-9)
            d = d + 1.2 * um
            # flip the label: push causal direction to opposite side, budget large
            proj = float(z @ v_causal_hat)
            d = d + (-np.sign(proj) if proj != 0 else 1.0) * (abs(proj) + 1.5) * v_causal_hat
            deltas.append(d)
        return True, {"deltas": [dd.tolist() for dd in deltas]}, 0.7