Repository Public Research preview
1"""Donor→recipient representation maps fitted on disjoint alignment data.23All maps are affine ``y ≈ x W + b`` from donor space (d_D) to recipient space (d_R) and expose the same4interface; ``fit`` consumes paired activations ``X`` [N, d_D] and ``Y`` [N, d_R] collected on the *same*5inputs. Leakage is audited by the caller through item hashes (alignment items ∩ evaluation items = ∅).67Families:8 identity : requires d_D == d_R; W = I, b = 0 (no fitting)9 procrustes : orthogonal W (Schönemann) after centring; optional isotropic scale; requires d_D == d_R10 (for d_D != d_R the smaller space is zero-padded — reported as such)11 ridge : least squares with L2 penalty (general affine map)12 cca : canonical correlation analysis map through the top-k canonical directions13 random_orth : random orthogonal W (sham mapping; same norm behaviour as procrustes)14 unrelated : ridge map fitted on *unrelated* data (items without the target skill) — supplied by caller15Each map records n_fit, parameter count, fit time and in-sample / held-out R².16"""17from __future__ import annotations1819import time20from dataclasses import dataclass, field2122import numpy as np232425@dataclass26class AffineMap:27 family: str28 W: np.ndarray # [d_D, d_R]29 b: np.ndarray # [d_R]30 info: dict = field(default_factory=dict)3132 def __call__(self, X: np.ndarray) -> np.ndarray:33 return X @ self.W + self.b3435 @property36 def n_params(self) -> int:37 return int(self.W.size + self.b.size)3839 def r2(self, X: np.ndarray, Y: np.ndarray) -> float:40 pred = self(X)41 ss_res = float(((Y - pred) ** 2).sum())42 ss_tot = float(((Y - Y.mean(0)) ** 2).sum())43 return 1.0 - ss_res / max(ss_tot, 1e-12)444546def _pad(X: np.ndarray, d: int) -> np.ndarray:47 if X.shape[1] == d:48 return X49 if X.shape[1] > d:50 raise ValueError("cannot pad down")51 return np.concatenate([X, np.zeros((X.shape[0], d - X.shape[1]), X.dtype)], axis=1)525354def fit_identity(X: np.ndarray, Y: np.ndarray) -> AffineMap:55 if X.shape[1] != Y.shape[1]:56 raise ValueError("identity map requires equal widths")57 return AffineMap("identity", np.eye(X.shape[1], dtype=np.float64), np.zeros(Y.shape[1]), {"n_fit": 0})585960def fit_procrustes(X: np.ndarray, Y: np.ndarray, scale: bool = True) -> AffineMap:61 t0 = time.time()62 d = max(X.shape[1], Y.shape[1])63 Xp, Yp = _pad(X.astype(np.float64), d), _pad(Y.astype(np.float64), d)64 mx, my = Xp.mean(0), Yp.mean(0)65 Xc, Yc = Xp - mx, Yp - my66 U, S, Vt = np.linalg.svd(Xc.T @ Yc, full_matrices=False)67 R = U @ Vt68 s = float(S.sum() / max((Xc ** 2).sum(), 1e-12)) if scale else 1.069 W = s * R70 b = my - mx @ W71 # crop to the true output width (zero-padded rows/cols carry no information)72 W = W[: X.shape[1], : Y.shape[1]]73 b = b[: Y.shape[1]]74 return AffineMap("procrustes", W, b, {"n_fit": int(len(X)), "scale": s, "padded": X.shape[1] != Y.shape[1], "fit_seconds": time.time() - t0})757677def fit_ridge(X: np.ndarray, Y: np.ndarray, lam: float = 1e-2, family: str = "ridge") -> AffineMap:78 t0 = time.time()79 Xd, Yd = X.astype(np.float64), Y.astype(np.float64)80 mx, my = Xd.mean(0), Yd.mean(0)81 Xc, Yc = Xd - mx, Yd - my82 d = Xc.shape[1]83 W = np.linalg.solve(Xc.T @ Xc + lam * len(Xc) * np.eye(d), Xc.T @ Yc)84 b = my - mx @ W85 return AffineMap(family, W, b, {"n_fit": int(len(X)), "lambda": lam, "fit_seconds": time.time() - t0})868788def fit_cca(X: np.ndarray, Y: np.ndarray, k: int | None = None, eps: float = 1e-3) -> AffineMap:89 """Map through canonical directions: x -> (x-mx) Wx (top-k) -> Wy^+ -> y space (regression on canonical variates)."""90 t0 = time.time()91 Xd, Yd = X.astype(np.float64), Y.astype(np.float64)92 mx, my = Xd.mean(0), Yd.mean(0)93 Xc, Yc = Xd - mx, Yd - my94 n = len(Xc)95 Sxx = Xc.T @ Xc / n + eps * np.eye(Xc.shape[1])96 Syy = Yc.T @ Yc / n + eps * np.eye(Yc.shape[1])97 Sxy = Xc.T @ Yc / n9899 def inv_sqrt(S):100 w, V = np.linalg.eigh(S)101 return V @ np.diag(1 / np.sqrt(np.clip(w, 1e-10, None))) @ V.T102103 Kx, Ky = inv_sqrt(Sxx), inv_sqrt(Syy)104 U, S, Vt = np.linalg.svd(Kx @ Sxy @ Ky, full_matrices=False)105 k = k or len(S)106 A = Kx @ U[:, :k] # x canonical weights [d_D, k]107 Zx = Xc @ A # canonical variates108 # regress Y on canonical variates (least squares)109 B = np.linalg.lstsq(Zx, Yc, rcond=None)[0] # [k, d_R]110 W = A @ B111 b = my - mx @ W112 return AffineMap("cca", W, b, {"n_fit": int(n), "k": int(k), "top_canonical_corr": float(S[0]), "fit_seconds": time.time() - t0})113114115def fit_random_orthogonal(X: np.ndarray, Y: np.ndarray, seed: int = 0) -> AffineMap:116 """Sham mapping: random orthogonal rotation (padded/cropped) with the Procrustes scale and the mean offset."""117 rng = np.random.default_rng(seed)118 d = max(X.shape[1], Y.shape[1])119 Q, _ = np.linalg.qr(rng.standard_normal((d, d)))120 Xd, Yd = X.astype(np.float64), Y.astype(np.float64)121 mx, my = Xd.mean(0), Yd.mean(0)122 s = float(np.linalg.norm(Yd - my) / max(np.linalg.norm(Xd - mx), 1e-12))123 W = (s * Q)[: X.shape[1], : Y.shape[1]]124 b = my - mx @ W125 return AffineMap("random_orth", W, b, {"n_fit": 0, "seed": seed, "scale": s})126127128FITTERS = {"identity": fit_identity, "procrustes": fit_procrustes, "ridge": fit_ridge, "cca": fit_cca, "random_orth": fit_random_orthogonal}129