1import math, random, time
 2import numpy as np
 3import torch
 4
 5SEED = 2259
 6np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
 7
 8def compositions(d, m):
 9    out=[]
10    def rec(i, left, a):
11        if i == d-1:
12            out.append(tuple(a+[left])); return
13        for k in range(left+1): rec(i+1, left-k, a+[k])
14    rec(0,m,[])
15    return out
16
17def mult(alpha, m):
18    z=math.factorial(m)
19    for a in alpha: z//=math.factorial(a)
20    return z
21
22def features(x, alphas, weighted=True):
23    # x: [batch,d]
24    vals=[]
25    for a in alphas:
26        v=torch.ones(x.shape[0], device=x.device, dtype=x.dtype)
27        for j,k in enumerate(a):
28            if k: v=v*x[:,j]**k
29        if weighted: v=v*math.sqrt(mult(a, sum(a)))
30        vals.append(v)
31    return torch.stack(vals,1)
32
33def math_checks():
34    print('MATH_CHECKS')
35    # Prediction 1: exact parameter count and ratio.
36    for d in (2,4,8,16):
37        for m in (2,4,6):
38            ordered=d**m; orbit=math.comb(d+m-1,m)
39            print(f'COUNT d={d} m={m} observed_ratio={ordered/orbit:.6f} predicted={ordered/orbit:.6f} ordered={ordered} orbit={orbit}')
40    # Prediction 2: diagonal identity, independently expanded ordered tensor.
41    for d,m in [(3,2),(4,4),(5,6)]:
42        aa=compositions(d,m); x=np.random.randn(d); vals=np.random.randn(len(aa))
43        p=sum(mult(a,m)*v*np.prod(x**np.array(a)) for a,v in zip(aa,vals))
44        ordered=0.0
45        # enumerate ordered tuples and use histogram lookup
46        from itertools import product
47        lookup={a:v for a,v in zip(aa,vals)}
48        for tup in product(range(d), repeat=m):
49            hist=tuple(tup.count(j) for j in range(d))
50            ordered += lookup[hist]*np.prod(x[list(tup)])
51        print(f'IDENTITY d={d} m={m} abs_error={abs(p-ordered):.3e}')
52    # Prediction 3: the explicit weighted feature has second moment exactly N times
53    # the unweighted feature under iid standard normal inputs.
54    for d,m in [(4,2),(4,4),(4,6),(8,4)]:
55        aa=compositions(d,m); x=torch.randn(180000,d)
56        f0=features(x,aa,False); f1=features(x,aa,True)
57        # Select low, middle, max multiplicity orbit and average empirical ratios.
58        ns=np.array([mult(a,m) for a in aa]); picks=[int(np.argmin(ns)), int(np.argsort(ns)[len(ns)//2]), int(np.argmax(ns))]
59        for i in picks:
60            ratio=(f1[:,i].square().mean()/f0[:,i].square().mean()).item()
61            print(f'VAR d={d} m={m} N={ns[i]} observed={ratio:.4f} predicted={ns[i]} alpha={aa[i]}')
62    # Prediction 4: repeating each orbit coefficient N times gives exactly the
63    # ordered BH q-sum, for every q (including q=2m/(m+1)).
64    for d,m in [(3,2),(4,4),(5,6)]:
65        aa=compositions(d,m); q=2*m/(m+1); coeff=np.random.randn(len(aa))
66        orbit=sum(mult(a,m)*abs(c)**q for a,c in zip(aa,coeff))
67        repeated=sum(abs(coeff[aa.index(tuple(t.count(j) for j in range(d)))])**q
68                     for t in __import__('itertools').product(range(d), repeat=m))
69        print(f'BH d={d} m={m} q={q:.6f} abs_error={abs(orbit-repeated):.3e}')
70
71def regression():
72    # Same polynomial target, same orbit basis: compare exact multiplicity basis to unweighted compression.
73    torch.manual_seed(SEED)
74    d,m=8,4; aa=compositions(d,m); p=len(aa)
75    ntr,nva=6000,1500
76    xtr=torch.randn(ntr,d); xva=torch.randn(nva,d)
77    target_theta=torch.randn(p)/math.sqrt(p)
78    with torch.no_grad():
79        ytr=features(xtr,aa,True)@target_theta; yva=features(xva,aa,True)@target_theta
80    results={}
81    for name,weighted in [('balanced',True),('unweighted',False)]:
82        torch.manual_seed(SEED+int(weighted)); w=torch.nn.Parameter(torch.zeros(p))
83        opt=torch.optim.Adam([w],lr=0.03)
84        t=time.perf_counter(); losses=[]
85        for step in range(500):
86            ix=torch.randint(0,ntr,(128,)); pred=features(xtr[ix],aa,weighted)@w
87            loss=((pred-ytr[ix])**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
88            if step in (0,49,199,499): losses.append(float(loss))
89        with torch.no_grad(): val=float(((features(xva,aa,weighted)@w-yva)**2).mean())
90        results[name]=(val,time.perf_counter()-t,losses)
91    print('REGRESSION d=8 m=4 params_each=%d' % p)
92    for k,v in results.items(): print(f'{k} val_mse={v[0]:.6g} seconds={v[1]:.3f} checkpoints={v[2]}')
93    print('DENSE_ORDERED_PARAMETER_COUNT',d**m,'ORBIT_PARAMETER_COUNT',p)
94
95if __name__=='__main__':
96    math_checks(); regression()