1import json
 2import numpy as np
 3import torch
 4from torch import nn
 5from volatility_clock import simulate, params
 6
 7
 8def make_batch(kind, rng, batch=96, length=32):
 9    # A deliberately small synthetic diffusion-like task. x0 is an AR(1)
10    # signal; the denoiser is blind to realized A and sees only xt.
11    x0 = np.zeros((batch, length), dtype=np.float32)
12    for j in range(1, length):
13        x0[:, j] = .85*x0[:, j-1] + rng.standard_normal(batch).astype(np.float32)
14    if kind == 'clock':
15        _, eps = simulate(.65, .8, .7, batch, length, rng, burn=120)
16    else:
17        eps = rng.standard_normal((batch, length))
18    ab = .55
19    xt = (np.sqrt(ab)*x0 + np.sqrt(1-ab)*eps).astype(np.float32)
20    return torch.tensor(xt[:, None, :], dtype=torch.float32), torch.tensor(eps[:, None, :], dtype=torch.float32)
21
22
23class TinyDenoiser(nn.Module):
24    def __init__(self):
25        super().__init__()
26        self.net = nn.Sequential(nn.Conv1d(1, 16, 5, padding=2), nn.GELU(),
27                                 nn.Conv1d(16, 16, 5, padding=2), nn.GELU(),
28                                 nn.Conv1d(16, 1, 5, padding=2))
29    def forward(self, x): return self.net(x)
30
31
32def run_one(kind):
33    torch.manual_seed(1142)
34    rng = np.random.default_rng(9000 + (kind == 'clock'))
35    model = TinyDenoiser()
36    opt = torch.optim.Adam(model.parameters(), lr=2e-3)
37    model.train()
38    for _ in range(250):
39        x, y = make_batch(kind, rng)
40        loss = ((model(x)-y)**2).mean()
41        opt.zero_grad(); loss.backward(); opt.step()
42    model.eval(); vals=[]
43    with torch.no_grad():
44        for _ in range(12):
45            x,y=make_batch(kind,rng)
46            vals.append(float(((model(x)-y)**2).mean()))
47    return float(np.mean(vals)), float(np.std(vals)/np.sqrt(len(vals)))
48
49if __name__ == '__main__':
50    out={'iid_gaussian':run_one('iid'), 'tempered_stable_clock_blind':run_one('clock')}
51    print(json.dumps(out, indent=2))
52    with open('denoiser_results.json','w') as f: json.dump(out,f,indent=2)