1import os, sys, json, time
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, train_model, sweep_baseline, make_report
  8
  9SEEDS = tuple(range(8))
 10# The sequence track is structurally matched: it is a temporal forecasting task.
 11# Quiet-clock threshold is fixed before running; no tuning on test data.
 12QUIET_THRESHOLD = 0.035
 13HIDDEN = 64
 14
 15
 16def seed_all(seed):
 17    np.random.seed(seed)
 18    torch.manual_seed(seed)
 19    if torch.cuda.is_available():
 20        torch.cuda.manual_seed_all(seed)
 21
 22
 23class ResidualState(nn.Module):
 24    """Shared neural vector field, with explicit or Stieltjes implicit updates."""
 25    def __init__(self, implicit, clock_scale=1.0, hidden=HIDDEN):
 26        super().__init__()
 27        self.implicit = bool(implicit)
 28        self.clock_scale = float(clock_scale)
 29        self.inp = nn.Linear(1, hidden)
 30        self.rec = nn.Linear(hidden, hidden, bias=False)
 31        self.bias = nn.Parameter(torch.zeros(hidden))
 32        self.head = nn.Linear(hidden, 1)
 33
 34    def vector_field(self, z, x):
 35        return torch.tanh(self.inp(x.unsqueeze(-1)) + torch.matmul(z, self.rec.weight.t()) + self.bias) - z
 36
 37    def clock(self, x):
 38        # Effective clock pauses on small chronological changes and jumps at events.
 39        change = torch.zeros_like(x)
 40        change[:, 1:] = (x[:, 1:] - x[:, :-1]).abs()
 41        return self.clock_scale * (change > QUIET_THRESHOLD).to(x.dtype)
 42
 43    def forward(self, x):
 44        b, t = x.shape
 45        z = torch.zeros(b, self.inp.out_features, device=x.device, dtype=x.dtype)
 46        ds = self.clock(x) if self.implicit else torch.ones_like(x)
 47        for k in range(t):
 48            d = ds[:, k:k+1]
 49            if self.implicit:
 50                # Damped fixed-point solve of z_next=z+d f(z_next,x).
 51                zn = z
 52                for _ in range(10):
 53                    target = z + d * self.vector_field(zn, x[:, k])
 54                    zn = 0.5 * zn + 0.5 * target
 55                z = torch.where(d > 0, zn, z)
 56            else:
 57                z = z + d * self.vector_field(z, x[:, k])
 58        return self.head(z)
 59
 60
 61def train_one(seed, implicit, lr, clock_scale=1.0, n_train=400, n_test=200, epochs=15):
 62    seed_all(seed)
 63    ds = get_dataset('sequence', seed, n_train=n_train, n_test=n_test)
 64    model = ResidualState(implicit=implicit, clock_scale=clock_scale)
 65    net, metric, history = train_model(model, ds, epochs=epochs, lr=lr, batch=128, log=lambda _: None)
 66    return float(metric) if metric is not None else float('nan')
 67
 68
 69def baseline_factory(cfg):
 70    return lambda seed: train_one(seed, False, cfg['lr'], 1.0)
 71
 72
 73def idea_factory(cfg):
 74    return lambda seed: train_one(seed, True, cfg['lr'], cfg['clock_scale'])
 75
 76
 77def mechanism_signature(seed=0):
 78    """Measure behavior of trained systems, rather than asserting an identity."""
 79    seed_all(seed)
 80    ds = get_dataset('sequence', seed, n_train=400, n_test=200)
 81    base = ResidualState(False); idea = ResidualState(True, 1.0)
 82    # Train both systems independently on the same data.
 83    base, _, _ = train_model(base, ds, epochs=15, lr=0.003, batch=128, log=lambda _: None)
 84    idea, _, _ = train_model(idea, ds, epochs=15, lr=0.003, batch=128, log=lambda _: None)
 85    x = ds['xte']
 86    with torch.no_grad():
 87        def states(m, implicit):
 88            dev = next(m.parameters()).device
 89            xd = x.to(dev)
 90            z = torch.zeros(x.shape[0], HIDDEN, device=dev)
 91            clock = m.clock(xd)
 92            ds0 = clock if implicit else torch.ones_like(clock)
 93            changes=[]; active=[]
 94            for k in range(x.shape[1]):
 95                old=z.clone(); d=ds0[:, k:k+1]
 96                if implicit:
 97                    zn=z
 98                    for _ in range(10):
 99                        zn=0.5*zn+0.5*(z+d*m.vector_field(zn,xd[:,k]))
100                    z=torch.where(d>0,zn,z)
101                else: z=z+d*m.vector_field(z,xd[:,k])
102                changes.append((z-old).norm(dim=1)); active.append((d[:,0]>0))
103            c=torch.stack(changes); a=torch.stack(active)
104            return float(c[~a].mean()) if (~a).any() else 0., float(c[a].mean()) if a.any() else 0., float(a.float().mean())
105        bi, be, ba=states(base,False); ii, ie, ia=states(idea,True)
106    # Predicted inactive change is zero; prediction is confirmed only with a measured small value.
107    return {'predicted_inactive_state_change': 0.0,
108            'observed_baseline_inactive_change': bi,
109            'observed_idea_inactive_change': ii,
110            'observed_baseline_event_change': be,
111            'observed_idea_event_change': ie,
112            'idea_active_fraction': ia,
113            'predicted_active_fraction_is_data_dependent': True,
114            'confirmed': bool(ii < 1e-7 and ia > 0.0)}
115
116
117def main():
118    # Search-space parity: every lr used by idea is also evaluated by baseline.
119    lrs = [0.001, 0.003, 0.01]
120    baseline_grid = [{'lr': lr, 'clock_scale': 1.0} for lr in lrs]
121    idea_grid = [{'lr': lr, 'clock_scale': 1.0} for lr in lrs]
122    base = sweep_baseline(baseline_factory, baseline_grid, seeds=(0,1,2,3))
123    # Evaluate each idea setting on all paired seeds; select by the same sweep seeds.
124    idea_sweep=[]
125    for cfg in idea_grid:
126        vals=[idea_factory(cfg)(s) for s in (0,1,2,3)]
127        idea_sweep.append({'cfg':cfg, 'mean':float(np.mean(vals))})
128    best_cfg=min(idea_sweep, key=lambda q:q['mean'])['cfg']
129    idea_vals=[idea_factory(best_cfg)(s) for s in SEEDS]
130    idea_res={'mean':float(np.mean(idea_vals)), 'std':float(np.std(idea_vals)), 'per_seed':idea_vals, 'n':len(idea_vals), 'best_cfg':best_cfg, 'sweep':idea_sweep}
131    sig=mechanism_signature(0)
132    report=make_report('sequence','custom_residual_state',base,idea_res,sig)
133    report['protocol_notes']={'structural_match':'sequence-level temporal forecasting; recurrent state transition replacement', 'baseline_method':'explicit residual Euler with one chronological transition per token', 'idea_method':'implicit residual solve with event-clock deltas', 'paired_seeds':list(SEEDS), 'epochs':15, 'n_train':400, 'n_test':200}
134    os.makedirs('artifacts',exist_ok=True)
135    with open('artifacts/bench_report.json','w') as f: json.dump(report,f,indent=2)
136    print(json.dumps(report,indent=2))
137
138if __name__=='__main__': main()