1import json
  2import math
  3import numpy as np
  4
  5
  6def tc_gaussian(cov, eps=1e-10):
  7    cov = np.asarray(cov, float)
  8    d = np.diag(cov)
  9    sign, ld = np.linalg.slogdet(cov + eps * np.eye(len(d)))
 10    if sign <= 0:
 11        return float("nan")
 12    return float(0.5 * (np.log(d + eps).sum() - ld))
 13
 14
 15def shrink_cov(x, alpha=0.0):
 16    s = np.cov(x, rowvar=False, ddof=1)
 17    return (1-alpha) * s + alpha * np.diag(np.diag(s))
 18
 19
 20def merge_score(cov, a, b, eps=1e-10):
 21    ab = list(a) + list(b)
 22    return tc_gaussian(cov[np.ix_(ab, ab)], eps) - tc_gaussian(cov[np.ix_(list(a), list(a))], eps) - tc_gaussian(cov[np.ix_(list(b), list(b))], eps)
 23
 24
 25def agglomerative(cov, k, eps=1e-10):
 26    groups = [[i] for i in range(cov.shape[0])]
 27    while len(groups) > k:
 28        best = None
 29        for i in range(len(groups)):
 30            for j in range(i+1, len(groups)):
 31                val = merge_score(cov, groups[i], groups[j], eps)
 32                if best is None or val > best[0]:
 33                    best = (val, i, j)
 34        _, i, j = best
 35        groups[i] = groups[i] + groups[j]
 36        del groups[j]
 37    return sorted([sorted(g) for g in groups])
 38
 39
 40def equicorr(d, rho):
 41    return (1-rho)*np.eye(d) + rho*np.ones((d,d))
 42
 43
 44def exact_equicorr_tc(d, rho):
 45    # det(Sigma)=(1-rho)^(d-1)*(1+(d-1)rho), unit variances.
 46    return -0.5*((d-1)*math.log(1-rho) + math.log(1+(d-1)*rho))
 47
 48
 49
 50def fit_shared_lowrank(x, y, q=2, ridge=1e-3, iters=80):
 51    # Shared adapter basis U with task-specific heads V: Y ~= X U V.
 52    n, p = x.shape; t = y.shape[1]
 53    U = np.random.default_rng(777).normal(size=(p, q)) * .1
 54    V = np.random.default_rng(778).normal(size=(q, t)) * .1
 55    for _ in range(iters):
 56        xu = x @ U
 57        V = np.linalg.solve(xu.T @ xu + ridge*np.eye(q), xu.T @ y)
 58        xv = y @ V.T
 59        U = np.linalg.solve(x.T @ x + ridge*np.eye(p), x.T @ xv)
 60    return U, V
 61
 62
 63def predict_shared_lowrank(x, model):
 64    U, V = model
 65    return x @ U @ V
 66
 67def fit_cluster_models(xtr, ytr, groups, ridge=1e-3):
 68    # Frozen base is zero; each routed adapter is a linear weight vector.
 69    p = xtr.shape[1]
 70    ws = []
 71    for g in groups:
 72        yy = ytr[:, g].mean(axis=1)
 73        a = xtr.T @ xtr + ridge*np.eye(p)
 74        w = np.linalg.solve(a, xtr.T @ yy)
 75        ws.append(w)
 76    return ws
 77
 78
 79def predict_cluster(x, ws, groups):
 80    out = np.zeros((len(x), sum(len(g) for g in groups)))
 81    # groups are contiguous in this experiment after canonical sorting.
 82    for w, g in zip(ws, groups):
 83        out[:, g] = x @ w[:, None]
 84    return out
 85
 86
 87def regression_demo(rng):
 88    t, p, ntr, nv = 6, 12, 700, 1000
 89    true_groups = [[0,1], [2,3], [4,5]]
 90    xtr, xv = rng.normal(size=(ntr,p)), rng.normal(size=(nv,p))
 91    base_w = rng.normal(size=(3,p))
 92    w = np.zeros((t,p))
 93    for j,g in enumerate(true_groups):
 94        for task in g: w[task] = base_w[j] + .10*rng.normal(size=p)
 95    # Correlated task noise makes residual dependence operationally visible.
 96    noise_cov = np.eye(t)*.35
 97    for g in true_groups:
 98        for i in g:
 99            for j in g:
100                if i != j: noise_cov[i,j] = .28
101    etr = rng.multivariate_normal(np.zeros(t), noise_cov, ntr)
102    ev = rng.multivariate_normal(np.zeros(t), noise_cov, nv)
103    ytr, yv = xtr@w.T + etr, xv@w.T + ev
104    # Warm-up predictor used only to collect held-out residuals.
105    warm_model = fit_shared_lowrank(xtr, ytr, q=2)
106    warm = predict_shared_lowrank(xv, warm_model)
107    residual = yv - warm
108    cov = shrink_cov(residual, alpha=.10)
109    discovered = agglomerative(cov, 3)
110    random_groups = [[0,2], [1,5], [3,4]]
111    results = {}
112    # Equal nominal budget: rank-2 shared factor has p*q + q*t = 36 parameters;
113    # each of three rank-1 routed adapters has 3*p = 36 parameters.
114    shared_model = fit_shared_lowrank(xtr, ytr, q=2)
115    results['shared'] = {'mse': float(np.mean((predict_shared_lowrank(xv, shared_model)-yv)**2)),
116                         'groups': [list(range(t))], 'params': p*2+2*t}
117    for name, groups in [('random', random_groups), ('residual_tc', discovered), ('oracle', true_groups)]:
118        ws = fit_cluster_models(xtr, ytr, groups)
119        pred = predict_cluster(xv, ws, groups)
120        results[name] = {'mse': float(np.mean((pred-yv)**2)), 'groups': groups, 'params': len(groups)*p}
121    return results
122
123
124def main():
125    rng = np.random.default_rng(1015)
126    # Prediction 1: exact TC rises according to closed form with rho.
127    d = 4
128    rhos = [0.0, .1, .2, .4, .6, .8]
129    tc_curve = []
130    for rho in rhos:
131        x = rng.multivariate_normal(np.zeros(d), equicorr(d,rho), 30000)
132        observed = tc_gaussian(np.cov(x, rowvar=False))
133        predicted = exact_equicorr_tc(d,rho)
134        tc_curve.append({'rho':rho, 'predicted':predicted, 'observed':observed, 'abs_error':abs(observed-predicted)})
135    # Prediction 2: near independence TC is quadratic in rho, not linear.
136    small = []
137    for rho in [.02,.04,.08,.12,.16]:
138        pred = exact_equicorr_tc(d,rho)
139        quadratic = d*(d-1)/4 * rho*rho
140        small.append({'rho':rho, 'observed_exact':pred, 'quadratic_prediction':quadratic, 'ratio':pred/quadratic})
141    # Prediction 3: pairwise Gaussian TC is -1/2 log(1-rho^2), and
142    # block clustering should recover the population pairs once signal exceeds sampling noise.
143    pair_tc = []
144    for rho in [.0,.2,.4,.6,.8]:
145        x = rng.multivariate_normal(np.zeros(2), equicorr(2,rho), 12000)
146        pair_tc.append({'rho': rho, 'predicted': -0.5*math.log(max(1-rho*rho, 1e-12)),
147                        'observed': tc_gaussian(np.cov(x, rowvar=False))})
148    recovery = []
149    for rho in [.0,.2,.4,.6,.8]:
150        cov = np.eye(6)
151        for g in [[0,1],[2,3],[4,5]]:
152            i,j=g; cov[i,j]=cov[j,i]=rho
153        x = rng.multivariate_normal(np.zeros(6), cov, 5000)
154        est = shrink_cov(x, .05)
155        groups = agglomerative(est, 3)
156        recovery.append({'rho':rho, 'groups':groups, 'exact_recovery':groups==[[0,1],[2,3],[4,5]], 'within_tc':float(np.mean([tc_gaussian(est[np.ix_(g,g)]) for g in [[0,1],[2,3],[4,5]]]))})
157    regression = regression_demo(rng)
158    output = {'tc_rho_sweep':tc_curve, 'small_rho_quadratic_sweep':small, 'pair_tc_sweep': pair_tc, 'block_clustering_sweep':recovery, 'regression_demo':regression}
159    with open('results.json','w') as f: json.dump(output,f,indent=2)
160    print(json.dumps(output, indent=2))
161
162if __name__ == '__main__': main()