1import json
  2import numpy as np
  3from scipy.linalg import eigh
  4from sklearn.linear_model import LogisticRegression
  5from sklearn.model_selection import train_test_split
  6from sklearn.metrics import accuracy_score
  7
  8EPS = 1e-10
  9
 10def sym(a): return (a + a.T) / 2
 11
 12def spectral(a, fn):
 13    v, u = eigh(sym(a)); v = np.maximum(v, EPS)
 14    return sym((u * fn(v)) @ u.T)
 15
 16def sqrt_spd(a): return spectral(a, np.sqrt)
 17def invsqrt_spd(a): return spectral(a, lambda x: 1 / np.sqrt(x))
 18def spd_log(a): return spectral(a, np.log)
 19def spd_exp(a): return spectral(a, np.exp)
 20def project(a, lo, hi): return spectral(a, lambda x: np.clip(x, lo, hi))
 21
 22def bw2(p, q):
 23    hp = sqrt_spd(p); r = sqrt_spd(hp @ q @ hp)
 24    return float(max(np.trace(p + q) - 2 * np.trace(r), 0))
 25
 26def objective(s, mats, weights):
 27    return .5 * sum(w * bw2(s, a) for w, a in zip(weights, mats))
 28
 29def eigvals(a): return eigh(sym(a), eigvals_only=True)
 30def condition(a):
 31    v = eigvals(a); return float(v[-1] / max(v[0], EPS))
 32
 33def bw_update(s, mats, weights, bounds=None):
 34    hs = sqrt_spd(s); z = np.zeros_like(s)
 35    for w, a in zip(weights, mats): z += w * sqrt_spd(hs @ a @ hs)
 36    his = invsqrt_spd(s)
 37    x = his @ (z @ z) @ his                 # article.md Eq. (3)
 38    return project(x, *bounds) if bounds else sym(x)
 39
 40def arithmetic(mats, w): return sum(x * a for x, a in zip(w, mats))
 41def log_euclidean(mats, w): return spd_exp(sum(x * spd_log(a) for x, a in zip(w, mats)))
 42
 43def random_spd(rng, d, lo, hi):
 44    q, _ = np.linalg.qr(rng.normal(size=(d, d)))
 45    v = np.exp(rng.uniform(np.log(lo), np.log(hi), d))
 46    return sym((q * v) @ q.T)
 47
 48def math_check(rng):
 49    lo, hi, d = .01, 3., 2
 50    # Direct numerical check of the claimed non-expansiveness.
 51    ratios = []
 52    for _ in range(120):
 53        x, y = random_spd(rng, 5, .01, 20), random_spd(rng, 5, .01, 20)
 54        ratios.append(np.sqrt(bw2(project(x,.1,10), project(y,.1,10))) / max(np.sqrt(bw2(x,y)),1e-12))
 55    # Search the paper's stated type of transient floor exit: feasible inputs and S0,
 56    # but an unprojected unit step below lo. This is a diagnostic, not cherry-picked output.
 57    found = None
 58    for trial in range(1200):
 59        mats = [random_spd(rng, d, lo, hi) for _ in range(3)]
 60        w = np.ones(3) / 3; s = arithmetic(mats, w)
 61        if eigvals(s)[0] < lo - 1e-8: continue
 62        raw, clipped = [s], [s]
 63        for _ in range(5):
 64            raw.append(bw_update(raw[-1], mats, w))
 65            clipped.append(bw_update(clipped[-1], mats, w, (lo, hi)))
 66        if min(eigvals(x)[0] for x in raw) < lo - 1e-6:
 67            found = (mats, w, raw, clipped, trial + 1); break
 68    if found is None:
 69        # Still provide a reproducible ordinary feasible run.
 70        mats = [random_spd(rng, d, lo, hi) for _ in range(3)]; w=np.ones(3)/3
 71        s=arithmetic(mats,w); raw=[s]; clipped=[s]
 72        for _ in range(5): raw.append(bw_update(raw[-1],mats,w)); clipped.append(bw_update(clipped[-1],mats,w,(lo,hi)))
 73        trials=None
 74    else: mats,w,raw,clipped,trials=found
 75    return {
 76      'projection_max_ratio':float(max(ratios)), 'projection_mean_ratio':float(np.mean(ratios)),
 77      'exit_search_trials':trials, 'raw_min_eigenvalues':[float(eigvals(x)[0]) for x in raw],
 78      'projected_min_eigenvalues':[float(eigvals(x)[0]) for x in clipped],
 79      'raw_objective':[objective(x,mats,w) for x in raw],
 80      'projected_objective':[objective(x,mats,w) for x in clipped],
 81      'raw_condition':[condition(x) for x in raw], 'projected_condition':[condition(x) for x in clipped]
 82    }
 83
 84def pooling_task(rng):
 85    n,d,m=500,5,6; q,_=np.linalg.qr(rng.normal(size=(d,d)))
 86    bases=[q@np.diag([4,2,1,.5,.25])@q.T,q@np.diag([.25,.5,1,2,4])@q.T]
 87    feats={k:[] for k in ('arith','log','bw','pbw')}; mats_by={k:[] for k in feats}; all_inputs=[]; labels=[]
 88    for j in range(n):
 89        base=bases[j%2]; mats=[]
 90        for _ in range(m):
 91            z=rng.normal(size=(d,d)); mats.append(sym(base+.15*(z@z.T/d)+.03*np.eye(d)))
 92        w=np.ones(m)/m; ar=arithmetic(mats,w); vals={'arith':ar,'log':log_euclidean(mats,w)}
 93        bw=ar; pbw=ar
 94        for _ in range(4): bw=bw_update(bw,mats,w); pbw=bw_update(pbw,mats,w,(.2,5))
 95        vals.update({'bw':bw,'pbw':pbw}); all_inputs.append(mats); labels.append(j%2)
 96        for k,s in vals.items(): mats_by[k].append(s); feats[k].append(s[np.triu_indices(d)])
 97    labels=np.asarray(labels); tr,te=train_test_split(np.arange(n),test_size=.3,random_state=7,stratify=labels); out={}
 98    for k in feats:
 99        x=np.asarray(feats[k]); clf=LogisticRegression(max_iter=1000,random_state=7).fit(x[tr],labels[tr])
100        out[k]={'accuracy':float(accuracy_score(labels[te],clf.predict(x[te]))), 'mean_condition':float(np.mean([condition(a) for a in mats_by[k]])), 'mean_objective':float(np.mean([objective(s,aa,np.ones(m)/m) for s,aa in zip(mats_by[k],all_inputs)]))}
101    return out
102
103def main():
104    result={'math_check':math_check(np.random.default_rng(3150)),'pooling_task':pooling_task(np.random.default_rng(3151))}
105    with open('results.json','w') as f: json.dump(result,f,indent=2)
106    print(json.dumps(result,indent=2))
107
108if __name__=='__main__': main()