Differentiable Euler-density morphology loss / run_experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
 1import json
 2import random
 3import numpy as np
 4import torch
 5import torch.nn as nn
 6import torch.nn.functional as F
 7
 8
 9def euler_soft(x, thresholds, temperature=0.1):
10    if x.ndim == 3: x = x[:, None]
11    t = torch.as_tensor(thresholds, device=x.device, dtype=x.dtype)
12    s = torch.sigmoid((x[:, None] - t[None, :, None, None, None]) / temperature)
13    p1 = s.mean((-1, -2, -3))
14    eh = (s[..., :, 1:] * s[..., :, :-1]).mean((-1, -2, -3))
15    ev = (s[..., 1:, :] * s[..., :-1, :]).mean((-1, -2, -3))
16    f = (s[..., :-1, :-1] * s[..., 1:, :-1] * s[..., :-1, 1:] * s[..., 1:, 1:]).mean((-1, -2, -3))
17    return p1 - eh - ev + f
18
19
20def euler_hard(x, thresholds):
21    if x.ndim == 3: x = x[:, None]
22    vals = []
23    for v in thresholds:
24        s = (x > v).float()
25        p1 = s.mean((-1, -2, -3))
26        eh = (s[..., :, 1:] * s[..., :, :-1]).mean((-1, -2, -3))
27        ev = (s[..., 1:, :] * s[..., :-1, :]).mean((-1, -2, -3))
28        f = (s[..., :-1, :-1] * s[..., 1:, :-1] * s[..., :-1, 1:] * s[..., 1:, 1:]).mean((-1, -2, -3))
29        vals.append(p1 - eh - ev + f)
30    return torch.stack(vals, 1)
31
32
33def shapes(n, h=16, seed=0):
34    rng = np.random.default_rng(seed); yy, xx = np.mgrid[:h, :h]; out=[]
35    for _ in range(n):
36        im = np.zeros((h,h), np.float32)
37        if rng.random() < .5:
38            cx,cy=rng.uniform(4,h-4,2); r=rng.uniform(2.5,5.5)
39            im[((xx-cx)**2+(yy-cy)**2)<r*r]=1
40        else:
41            x0,y0=rng.integers(2,h-8,2); w,hh=rng.integers(3,8,2)
42            im[y0:y0+hh,x0:x0+w]=1
43        im += rng.normal(0,.08,(h,h)).astype(np.float32)
44        out.append(np.clip(im,0,1))
45    return torch.tensor(np.array(out)[:,None])
46
47class Net(nn.Module):
48    def __init__(self):
49        super().__init__(); self.a=nn.Conv2d(1,16,3,padding=1); self.b=nn.Conv2d(16,16,3,padding=1); self.c=nn.Conv2d(16,1,3,padding=1)
50    def forward(self,x): return torch.sigmoid(self.c(F.relu(self.b(F.relu(self.a(x))))))
51
52def main():
53    torch.set_num_threads(2)
54    random.seed(7); np.random.seed(7); torch.manual_seed(7)
55    thresholds=[.2,.4,.6,.8]; results={}
56    # Check 1: T -> 0 agrees with hard excursion away from threshold ties.
57    x=torch.rand(32,1,12,12)*.8+.1; hard=euler_hard(x,thresholds)
58    errs=[]
59    for T in [.2,.1,.05,.02,.01]: errs.append(float((euler_soft(x,thresholds,T)-hard).abs().mean()))
60    results['temperature_errors']={'T':[.2,.1,.05,.02,.01],'error':errs}
61    # Check 2: constant field has chi density 1 if occupied, 0 otherwise.
62    const=torch.full((3,1,10,11),.9); constv=float(euler_soft(const,[.5],.01).mean())
63    emptyv=float(euler_soft(torch.zeros_like(const),[.5],.01).mean())
64    results['constant_field']={'occupied':constv,'empty':emptyv,'predicted':[0.0,0.0]}
65    # Check 3: iid Bernoulli occupancy predicts q-2q^2+q^4.
66    q=.35; g=torch.Generator().manual_seed(11); z=(torch.rand((512,1,48,48),generator=g)<q).float()
67    empirical=float(euler_hard(z,[.5]).mean()); predicted=q-2*q*q+q**4
68    results['iid_bernoulli']={'q':q,'empirical':empirical,'predicted':predicted,'abs_error':abs(empirical-predicted),'note':'finite-grid variance and overlapping windows make exact agreement approximate'}
69    # Small equal-step denoising test; targets are clean morphology statistics.
70    clean=shapes(128,16,3); noisy=(clean+.22*torch.randn(clean.shape)).clamp(0,1)
71    real_target=euler_soft(clean,thresholds,.05).mean(0).detach()
72    def train(lam):
73        torch.manual_seed(19); net=Net(); opt=torch.optim.Adam(net.parameters(),lr=.004)
74        for step in range(80):
75            ix=torch.randint(0,len(clean),(32,)); pred=net(noisy[ix]); loss=F.mse_loss(pred,clean[ix])
76            if lam: loss=loss+lam*((euler_soft(pred,thresholds,.1).mean(0)-real_target)**2).mean()
77            opt.zero_grad(); loss.backward(); opt.step()
78        with torch.no_grad():
79            pred=net(noisy); mse=float(F.mse_loss(pred,clean)); ee=float((euler_hard(pred,thresholds)-euler_hard(clean,thresholds)).abs().mean()); soft=float(((euler_soft(pred,thresholds,.1).mean(0)-real_target)**2).mean())
80        return {'mse':mse,'hard_euler_mae':ee,'soft_target_mse':soft}
81    results['denoising']={'baseline':train(0.0),'euler_lambda_0.2':train(.2)}
82    with open('results.json','w') as f: json.dump(results,f,indent=2)
83    print(json.dumps(results,indent=2))
84
85if __name__=='__main__': main()