Differentiable Euler-density morphology loss / run_experiment.py
Mechanism confirmed, baseline not beaten
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()