1import json, math, os
  2import numpy as np
  3from scipy.optimize import minimize, brentq
  4from itertools import product
  5
  6rng = np.random.default_rng(7)
  7
  8def mons(n, deg):
  9    out=[]
 10    def rec(rem,k,p):
 11        if k==n-1: out.append(tuple(p+[rem])); return
 12        for a in range(rem+1): rec(rem-a,k+1,p+[a])
 13    rec(deg,0,[]); return out
 14
 15def poly_add(a,b,scale=1.):
 16    c=dict(a)
 17    for e,v in b.items(): c[e]=c.get(e,0.)+scale*v
 18    return {e:v for e,v in c.items() if abs(v)>1e-14}
 19
 20def poly_mul(a,b):
 21    c={}
 22    for x,v in a.items():
 23      for y,w in b.items():
 24        e=tuple(i+j for i,j in zip(x,y)); c[e]=c.get(e,0.)+v*w
 25    return c
 26
 27def power(p,k):
 28    q={(0,)*len(next(iter(p))):1.}
 29    for _ in range(k): q=poly_mul(q,p)
 30    return q
 31
 32def gram_system(poly,n,deg):
 33    mm=mons(n,deg); pairs=[(i,j) for i in range(len(mm)) for j in range(i,len(mm))]
 34    exps=sorted(set(tuple(mm[i][k]+mm[j][k] for k in range(n)) for i,j in pairs)|set(poly))
 35    rows=[]; rhs=[]
 36    for e in exps:
 37      row=np.zeros(len(pairs))
 38      for k,(i,j) in enumerate(pairs):
 39        ee=tuple(mm[i][z]+mm[j][z] for z in range(n))
 40        if ee==e: row[k]=1 if i==j else 2
 41      rows.append(row); rhs.append(poly.get(e,0.))
 42    return mm,pairs,np.asarray(rows),np.asarray(rhs)
 43
 44def solve_sos(poly,n,r,tries=1):
 45    # Fast approximate SDP: alternate projection onto coefficient affine space and PSD cone.
 46    norm2={tuple(2 if i==j else 0 for i in range(n)):1. for j in range(n)}
 47    lifted=poly_mul(power(norm2,r),poly)
 48    deg=max(sum(e) for e in lifted)//2
 49    mm,pairs,C,b=gram_system(lifted,n,deg); m=len(mm)
 50    # map symmetric matrix to the upper-triangular coefficient vector
 51    x=np.linalg.lstsq(C,b,rcond=None)[0]
 52    def unpack(v):
 53      Q=np.zeros((m,m))
 54      for a,(i,j) in zip(v,pairs): Q[i,j]=Q[j,i]=a
 55      return Q
 56    def pack(Q): return np.array([Q[i,j] for i,j in pairs])
 57    # affine projection, with a PSD projection between corrections
 58    G=np.linalg.pinv(C@C.T,rcond=1e-10)
 59    for _ in range(100):
 60      Q=unpack(x); w,U=np.linalg.eigh((Q+Q.T)/2); Q=(U*np.maximum(w,0.))@U.T
 61      x=x+C.T@(G@(b-C@pack(Q)))
 62    Q=unpack(x); ev=np.linalg.eigvalsh((Q+Q.T)/2)
 63    return {'success':bool(ev.min()>=-2e-5 and np.max(np.abs(C@x-b))<2e-5),
 64            'margin':float(ev.min()),'residual':float(np.max(np.abs(C@x-b))),
 65            'status':'alternating-projection'}
 66
 67def normpow(n,k):
 68    return power({tuple(2 if i==j else 0 for i in range(n)):1. for j in range(n)},k)
 69
 70def shifted_motzkin(eps):
 71    # Motzkin form plus eps*||x||_2^6; coefficients are accumulated explicitly.
 72    p={(4,2,0):1.,(2,4,0):1.,(0,0,6):1.,(2,2,2):-3.}
 73    return poly_add(p,normpow(3,3),scale=eps)
 74
 75def threshold(n,r):
 76    # Coarse threshold scan avoids expensive repeated nonlinear solves.
 77    grid=[0.,.05,.15,.35,1.0]
 78    vals=[]
 79    for e in grid:
 80      rec=solve_sos(shifted_motzkin(e),n,r,tries=1)
 81      vals.append((e,rec['margin']))
 82      if rec['margin']>=-2e-5: return e,rec['margin']
 83    return grid[-1],vals[-1][1]
 84
 85def rollout_boundary():
 86    # scalar z+=gamma*a*z; assess stability from 250-step bounded rollout.
 87    rows=[]
 88    for a in [0.5,0.8,1.0,1.2,1.5]:
 89      g=np.linspace(.1,2.2/a,401); empirical=[]
 90      for gamma in g:
 91        z=1.;
 92        for _ in range(250): z*=gamma*a
 93        empirical.append(np.isfinite(z) and abs(z)<=1.0001)
 94      stable=[x for x,ok in zip(g,empirical) if ok]
 95      last=max(stable,default=0.)
 96      rows.append({'a':a,'predicted_gamma':1/a,'observed_gamma':float(last),
 97                   'relative_error':abs(last-1/a)/(1/a)})
 98    return rows
 99
100def main():
101    # Mechanism test: Motzkin is nonnegative, but not level-0 SOS; multiplier levels improve margin.
102    levels=[]
103    for r in range(2):
104      rec=solve_sos(shifted_motzkin(0.),3,r,tries=1); levels.append({'r':r,**rec})
105    thresholds=[]
106    for r in range(2):
107      th,mg=threshold(3,r); thresholds.append({'r':r,'epsilon':th,'margin':mg})
108    boundary=rollout_boundary()
109    result={'levels':levels,'thresholds':thresholds,'boundary':boundary,
110      'predictions':{'threshold_monotone':all(thresholds[i]['epsilon']>=thresholds[i+1]['epsilon']-2e-3 for i in range(len(thresholds)-1)),
111                     'level_improves':levels[-1]['margin']>levels[0]['margin']+1e-4,
112                     'boundary_relative_error_max':max(abs(x['observed_gamma']-x['predicted_gamma'])/x['predicted_gamma'] for x in boundary)}}
113    open('results.json','w').write(json.dumps(result,indent=2))
114    print(json.dumps(result,indent=2))
115if __name__=='__main__': main()