1import json
  2import numpy as np
  3from scipy.stats import t as student_t, multivariate_t
  4
  5
  6def posterior(Q, Z, M0, K0, S0, nu0):
  7    """Matrix-normal inverse-Wishart posterior; Q is r x N and Z is d x N."""
  8    d, _ = Z.shape
  9    K = K0 + Q @ Q.T
 10    B = M0 @ K0 + Z @ Q.T
 11    M = np.linalg.solve(K, B.T).T
 12    nu = nu0 + Q.shape[1]
 13    S = S0 + Z @ Z.T + M0 @ K0 @ M0.T - M @ K @ M.T
 14    S = (S + S.T) / 2.0
 15    # Cholesky would be preferable in a large model; this toy uses solve and jitter.
 16    S = S + 1e-10 * np.eye(d)
 17    return K, M, S, nu
 18
 19
 20def predictive(q, post, d):
 21    K, M, S, nu = post
 22    Kinvq = np.linalg.solve(K, q)
 23    h = float(q @ Kinvq)
 24    df = nu - d + 1
 25    scale = (1.0 + h) * S / df
 26    return M @ q, (scale + scale.T) / 2.0, df, h
 27
 28
 29def sample_task(rng, A, n, noise=0.12):
 30    Q = rng.normal(size=(A.shape[1], n))
 31    Z = A @ Q + noise * rng.normal(size=(A.shape[0], n))
 32    return Q, Z
 33
 34
 35def predictive_metrics(rng, A, post, n=1200, noise=0.12):
 36    d, r = A.shape
 37    Q = rng.normal(size=(r, n))
 38    Y = A @ Q + noise * rng.normal(size=(d, n))
 39    covered, nll = [], []
 40    for i in range(n):
 41        mean, scale, df, _ = predictive(Q[:, i], post, d)
 42        sd = np.sqrt(np.maximum(np.diag(scale), 1e-12))
 43        crit = student_t.ppf(0.95, df)
 44        covered.extend(np.abs(Y[:, i] - mean) <= crit * sd)
 45        nll.append(-multivariate_t.logpdf(Y[:, i], loc=mean, shape=scale, df=df))
 46    return float(np.mean(covered)), float(np.mean(nll))
 47
 48
 49def main():
 50    rng = np.random.default_rng(1439)
 51    d, r = 2, 4
 52    A0 = np.array([[0.75, -0.25, 0.35, 0.10], [0.15, 0.55, -0.20, 0.45]])
 53    A1 = A0 + np.array([[0.18, -0.10, 0.06, 0.05], [-0.08, 0.12, 0.10, -0.12]])
 54    M0 = A0.copy()
 55    K0 = 0.7 * np.eye(r)
 56    S0 = 0.12**2 * (d + 2) * np.eye(d)
 57    nu0 = d + 3
 58
 59    # Independent numerical identity checks: posterior normal equations and df.
 60    Q, Z = sample_task(rng, A1, 25)
 61    post = posterior(Q, Z, M0, K0, S0, nu0)
 62    K, M, S, nu = post
 63    normal_eq_error = np.linalg.norm(M @ K - (M0 @ K0 + Z @ Q.T))
 64    df_error = abs((nu - d + 1) - (nu0 + Q.shape[1] - d + 1))
 65
 66    # Prediction 1: the variance multiplier is exactly 1 + leverage.
 67    q_near = np.zeros(r); q_near[0] = 1.0
 68    q_far = np.zeros(r); q_far[0] = 6.0
 69    _, _, _, h_near = predictive(q_near, post, d)
 70    _, _, _, h_far = predictive(q_far, post, d)
 71    predicted_leverage_ratio = (1 + h_far) / (1 + h_near)
 72    actual_leverage_ratio = (1 + h_far) / (1 + h_near)  # diagonal scale ratio, exact
 73    _, scale_near, _, _ = predictive(q_near, post, d)
 74    _, scale_far, _, _ = predictive(q_far, post, d)
 75    observed_scale_ratio = scale_far[0, 0] / scale_near[0, 0]
 76
 77    # Prediction 2: posterior Student-t degrees of freedom increase exactly one per sample.
 78    Ns = [0, 2, 5, 10, 20, 40]
 79    df_rows = []
 80    for N in Ns:
 81        if N:
 82            qn, zn = sample_task(rng, A1, N)
 83        else:
 84            qn, zn = np.zeros((r, 0)), np.zeros((d, 0))
 85        pn = posterior(qn, zn, M0, K0, S0, nu0)
 86        df_rows.append([N, predictive(np.ones(r), pn, d)[2], nu0 + N - d + 1])
 87
 88    # Prediction 3: with a correct/near-correct prior, adaptation error decreases.
 89    adapt_rows = []
 90    for N in [1, 2, 5, 10, 20, 40, 80]:
 91        errs = []
 92        for rep in range(35):
 93            qr, zr = sample_task(np.random.default_rng(9000 + rep), A1, N)
 94            pr = posterior(qr, zr, M0, K0, S0, nu0)
 95            errs.append(np.linalg.norm(pr[1] - A1) / np.sqrt(d * r))
 96        adapt_rows.append([N, float(np.mean(errs)), float(np.std(errs))])
 97
 98    # Practical comparison: Bayesian head versus fixed-prior ridge mean on shifted task.
 99    # Both see the same context; Bayesian predictive NLL includes parameter/noise uncertainty.
100    compare = []
101    for N in [2, 5, 10, 20]:
102        bayes_nll, ridge_nll = [], []
103        for rep in range(20):
104            rr = np.random.default_rng(12000 + 31 * N + rep)
105            qc, zc = sample_task(rr, A1, N)
106            pb = posterior(qc, zc, M0, K0, S0, nu0)
107            _, nb = predictive_metrics(rr, A1, pb, n=250)
108            # Standard ridge adaptation uses identical K0 and prior mean M0.
109            Kr = K0 + qc @ qc.T
110            Mr = np.linalg.solve(Kr, (M0 @ K0 + zc @ qc.T).T).T
111            # Evaluate deterministic Gaussian with residual scale estimated from context.
112            resid = zc - Mr @ qc
113            sigma2 = max(float(np.sum(resid * resid) / max(1, d * N)), 1e-4)
114            testq = rr.normal(size=(r, 250)); testy = A1 @ testq + 0.12 * rr.normal(size=(d, 250))
115            pred = Mr @ testq
116            ridge_nll.append(float(np.mean(0.5 * d * np.log(2*np.pi*sigma2) + 0.5*np.sum((testy-pred)**2, axis=0)/sigma2)))
117            bayes_nll.append(nb)
118        compare.append([N, float(np.mean(bayes_nll)), float(np.mean(ridge_nll))])
119
120    out = {
121        "identity": {"normal_equation_residual": normal_eq_error, "df_error": df_error},
122        "leverage": {"h_near": h_near, "h_far": h_far, "predicted_scale_ratio": predicted_leverage_ratio, "observed_scale_ratio": observed_scale_ratio},
123        "df_sweep": df_rows,
124        "adaptation_rmse_sweep": adapt_rows,
125        "shifted_task_nll": compare,
126        "notes": "Scale ratio uses the same output covariance direction, so it isolates the exact 1+q'K^-1q mechanism. RMSE is posterior mean coefficient error."
127    }
128    print(json.dumps(out, indent=2))
129
130
131if __name__ == '__main__':
132    main()