1import json, math, os, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 2016
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8
  9
 10def tents(diagram, grid):
 11    d = np.asarray(diagram, dtype=float).reshape(-1, 2)
 12    b, death = d[:, 0:1], d[:, 1:2]
 13    return np.maximum(0.0, np.minimum(grid[None, :] - b, death - grid[None, :]))
 14
 15
 16def landscape(diagram, grid, K):
 17    v = tents(diagram, grid)
 18    if len(v) == 0:
 19        return np.zeros((K, len(grid)))
 20    v = np.sort(v, axis=0)[::-1]
 21    out = np.zeros((K, len(grid)))
 22    out[:min(K, len(v))] = v[:K]
 23    return out
 24
 25
 26def lp_grid(a, b, p=2, dt=0.002):
 27    # Common grid over the supports; sufficiently fine for this sanity check.
 28    lo = min([x[0] for x in a] + [x[0] for x in b]) - .01
 29    hi = max([x[1] for x in a] + [x[1] for x in b]) + .01
 30    g = np.arange(lo, hi + dt/2, dt)
 31    return (np.sum(np.abs(landscape(a, g, max(len(a),len(b))) -
 32                         landscape(b, g, max(len(a),len(b))))**p) * dt) ** (1/p)
 33
 34
 35def tent_distance(x, y, dt=0.001):
 36    lo = min(x[0], y[0]) - .01; hi = max(x[1], y[1]) + .01
 37    g = np.arange(lo, hi + dt/2, dt)
 38    return np.sqrt(np.sum(np.abs(tents([x], g)[0] - tents([y], g)[0])**2) * dt)
 39
 40
 41def matching_bound(a, b, dt=0.001):
 42    # Equal-cardinality matched-pair W_2^triangle upper bound (the chosen matching).
 43    return math.sqrt(sum(tent_distance(x, y, dt)**2 for x, y in zip(a, b)))
 44
 45
 46class PersistenceLandscape(nn.Module):
 47    def __init__(self, t_min=0., t_max=1., n_grid=32, K=3):
 48        super().__init__()
 49        self.K, self.n_grid = K, n_grid
 50        self.register_buffer('grid', torch.linspace(t_min, t_max, n_grid))
 51    def forward(self, diagrams, mask=None):
 52        # diagrams [B,N,2], mask [B,N], padded rows are ignored.
 53        b = diagrams[..., 0:1]; d = diagrams[..., 1:2]
 54        t = self.grid.view(1, 1, -1)
 55        v = torch.relu(torch.minimum(t-b, d-t))
 56        if mask is not None:
 57            v = v.masked_fill(~mask[..., None], -1e9)
 58        vals, _ = torch.topk(v, k=min(self.K, v.shape[1]), dim=1)
 59        if vals.shape[1] < self.K:
 60            vals = torch.cat([vals, torch.zeros(*vals.shape[:1], self.K-vals.shape[1], vals.shape[2], device=vals.device)], 1)
 61        return vals.clamp_min(0.)
 62
 63
 64def make_data(n, N=6, noise=0., seed=0):
 65    rng=np.random.default_rng(seed); ds=[]; ys=[]
 66    for i in range(n):
 67        y=int(i % 2); center=.34 if y==0 else .66
 68        births=np.clip(rng.normal(center, .09, N), .03, .82)
 69        pers=np.clip(rng.normal(.22, .05, N), .04, .38)
 70        deaths=np.minimum(births+pers, .97)
 71        ds.append(np.stack([births,deaths],1)); ys.append(y)
 72    ds=np.asarray(ds)
 73    if noise: ds=np.stack([np.clip(ds[:,:,0]+rng.normal(0,noise,ds.shape[:2]),0,.9), np.clip(ds[:,:,1]+rng.normal(0,noise,ds.shape[:2]),.05,1)],2); ds[:,:,1]=np.maximum(ds[:,:,1],ds[:,:,0]+.01)
 74    return ds.astype('float32'), np.asarray(ys, dtype='int64')
 75
 76
 77def train_eval(kind, train, ytr, test, yte, epochs=80):
 78    device='cuda' if torch.cuda.is_available() else 'cpu'
 79    try:
 80        if kind=='landscape':
 81            layer=PersistenceLandscape(.0,1.,32,3); Xtr=layer(torch.tensor(train)).flatten(1); Xte=layer(torch.tensor(test)).flatten(1)
 82        else:
 83            Xtr=torch.tensor(train).flatten(1); Xte=torch.tensor(test).flatten(1)
 84        model=nn.Sequential(nn.Linear(Xtr.shape[1],32),nn.ReLU(),nn.Linear(32,2)).to(device)
 85        Xtr,Xte=Xtr.to(device),Xte.to(device); yt=torch.tensor(ytr,device=device); yv=torch.tensor(yte,device=device)
 86        opt=torch.optim.Adam(model.parameters(),lr=.01)
 87        for _ in range(epochs):
 88            opt.zero_grad(); loss=nn.functional.cross_entropy(model(Xtr),yt); loss.backward(); opt.step()
 89        with torch.no_grad(): acc=(model(Xte).argmax(1)==yv).float().mean().item()
 90        return acc
 91    except Exception:
 92        # CUDA OOM or driver issues: rerun on CPU.
 93        torch.cuda.empty_cache() if torch.cuda.is_available() else None
 94        old=torch.cuda.is_available; torch.cuda.is_available=lambda:False
 95        try: return train_eval(kind,train,ytr,test,yte,epochs)
 96        finally: torch.cuda.is_available=old
 97
 98
 99def main():
100    # Mechanism verification: predictions are (P1) one point saturates the bound,
101    # (P2) every multi-point ratio is <=1, and (P3) larger noise cannot amplify
102    # landscape perturbation beyond the chosen matching bound.
103    grid=np.linspace(0,1,2001); rows=[]
104    a=[(.15,.55)];
105    for shift in [.002,.01,.03,.08]:
106        b=[(a[0][0]+shift,a[0][1]+shift)]
107        lhs=np.linalg.norm(landscape(a,grid,3)-landscape(b,grid,3))*math.sqrt(1/2000)
108        rhs=tent_distance(a[0],b[0]); rows.append({'single_shift':shift,'landscape_norm':lhs,'bound':rhs,'ratio':lhs/rhs})
109    rng=np.random.default_rng(4); ratios=[]
110    for _ in range(100):
111        A=[]; B=[]
112        for j in range(5):
113            x=float(rng.uniform(.05,.75)); p=float(rng.uniform(.08,.3)); A.append((x,x+p))
114            dx=float(rng.normal(0,.04)); dp=float(rng.normal(0,.03)); z=max(.01,min(.9,x+dx)); B.append((z,min(.99,z+max(.02,p+dp))))
115        ratios.append(lp_grid(A,B)/matching_bound(A,B))
116    # P3: piecewise-linear tents integrated on a uniform grid should converge
117    # to the continuous L2 norm as dt shrinks (first-order rectangle estimate).
118    ca=[(.13,.57),(.31,.76),(.61,.91)]; cb=[(.16,.55),(.28,.79),(.65,.88)]
119    ref=lp_grid(ca, cb, dt=0.0001)
120    discretization=[]
121    for dt in [.02,.01,.005,.0025,.00125]:
122        val=lp_grid(ca, cb, dt=dt)
123        discretization.append({'dt':dt,'grid_norm':val,'abs_error':abs(val-ref),'error_over_dt':abs(val-ref)/dt})
124    # Autograd prediction: away from max/tent kinks, the layer has finite gradients.
125    q=torch.tensor([[[.18,.52],[.43,.81]]],dtype=torch.float32,requires_grad=True)
126    qout=PersistenceLandscape(.0,1.,64,2)(q)
127    qout.sum().backward()
128    grad_ok=bool(torch.isfinite(q.grad).all() and q.grad.abs().sum()>0)
129    prediction={'single_point_equality':rows,'multi_point_max_ratio':float(max(ratios)),
130                'multi_point_mean_ratio':float(np.mean(ratios)),'predicted_bound':1.0,
131                'grid_convergence':{'reference_dt':0.0001,'sweep':discretization,
132                                    'prediction':'error decreases as dt decreases'},
133                'autograd_finite_nonzero':grad_ok}
134    tr,ytr=make_data(160,seed=10); te,yte=make_data(80,seed=11)
135    clean={'padded_birth_death':train_eval('raw',tr,ytr,te,yte),'landscape':train_eval('landscape',tr,ytr,te,yte)}
136    noisy=[]
137    for sigma in [.01,.03,.07]:
138        nt,ny=make_data(80,noise=sigma,seed=20); noisy.append({'sigma':sigma,'padded_birth_death':train_eval('raw',tr,ytr,nt,ny),'landscape':train_eval('landscape',tr,ytr,nt,ny)})
139    result={'prediction_checks':prediction,'clean_accuracy':clean,'noise_accuracy':noisy}
140    with open('results.json','w') as f: json.dump(result,f,indent=2)
141    print(json.dumps(result,indent=2))
142if __name__=='__main__': main()