1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 2870
  6rng = np.random.default_rng(SEED)
  7
  8
  9def make_incidence(n=7, edges=None):
 10    if edges is None:
 11        edges = [(0,1),(1,2),(2,3),(3,4),(4,5),(5,6),(0,6),(1,5),(2,6)]
 12    B = np.zeros((n, len(edges)), dtype=np.float64)
 13    for j, (tail, head) in enumerate(edges):
 14        B[tail, j] = -1.0
 15        B[head, j] = 1.0
 16    return B, edges
 17
 18
 19def core_sweeps():
 20    B, edges = make_incidence()
 21    n, m = B.shape
 22    ones = np.ones(n)
 23    # Prediction 1: arbitrary parallel channels remain exactly conservative.
 24    channel_rows = []
 25    for k in [1, 2, 4, 8, 16, 32]:
 26        Pk = rng.normal(size=(m, k))
 27        P = Pk.sum(axis=1)
 28        residual = float(abs(ones @ (B @ P)))
 29        channel_rows.append({"K": k, "abs_total_internal_residual": residual})
 30
 31    # Prediction 2: if an unconstrained construction adds leakage with amplitude
 32    # lambda, the mass residual is linear in lambda. Here r is fixed, so this is
 33    # a clean quantitative slope test rather than a noisy training result.
 34    P = rng.normal(size=m)
 35    r = rng.normal(size=n)
 36    lambdas = [0.0, 0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.0]
 37    leakage_rows = []
 38    for lam in lambdas:
 39        q_unconstrained = B @ P + lam * r
 40        observed = float(abs(ones @ q_unconstrained))
 41        predicted = float(lam * abs(ones @ r))
 42        leakage_rows.append({"lambda": lam, "observed_abs_residual": observed,
 43                             "predicted_abs_residual": predicted})
 44    nonzero = [(x["lambda"], x["observed_abs_residual"]) for x in leakage_rows if x["lambda"] > 0]
 45    slope = float(sum(a*b for a,b in nonzero) / sum(a*a for a,b in nonzero))
 46    expected_slope = float(abs(ones @ r))
 47
 48    # Prediction 3: with internal-only dynamics, conservative Euler updates
 49    # preserve total mass for every rollout length; constant leakage accumulates
 50    # linearly in the number of steps.
 51    dt, eps = 0.05, 0.02
 52    x0 = rng.normal(size=n)
 53    flow = rng.normal(size=m)
 54    q_internal = B @ flow
 55    steps_list = [1, 2, 5, 10, 20, 40, 80]
 56    drift_rows = []
 57    for steps in steps_list:
 58        xc = x0.copy()
 59        xu = x0.copy()
 60        for _ in range(steps):
 61            xc += dt * q_internal
 62            xu += dt * (q_internal + eps * np.ones(n))
 63        cdrift = float(abs(ones @ (xc - x0)))
 64        udrift = float(abs(ones @ (xu - x0)))
 65        predicted_u = float(steps * dt * eps * n)
 66        drift_rows.append({"steps": steps, "conservative_abs_mass_drift": cdrift,
 67                           "unconstrained_abs_mass_drift": udrift,
 68                           "predicted_unconstrained_drift": predicted_u})
 69    return {
 70        "incidence_shape": [n, m],
 71        "incidence_column_sum_max_abs": float(np.max(np.abs(ones @ B))),
 72        "channel_sweep": channel_rows,
 73        "leakage_sweep": leakage_rows,
 74        "leakage_fitted_slope": slope,
 75        "leakage_predicted_slope": expected_slope,
 76        "rollout_sweep": drift_rows,
 77    }
 78
 79
 80def toy_model_comparison():
 81    # Small supervised graph-flow task: true interaction is a sum of two
 82    # domain-specific edge channels. The conservative model has two scalar
 83    # channel regressors and aggregates only through B. The baseline predicts
 84    # both endpoint updates independently from the composite edge features.
 85    try:
 86        import torch
 87        torch.manual_seed(SEED)
 88        np.random.seed(SEED)
 89        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 90        B_np, edges = make_incidence(n=6, edges=[(0,1),(1,2),(2,3),(3,4),(4,5),(0,5),(1,4)])
 91        B = torch.tensor(B_np, dtype=torch.float32, device=device)
 92        n, m = B_np.shape
 93        N = 1200
 94        x = torch.randn(N, n, device=device)
 95        u = torch.randn(N, m, 2, device=device)
 96        tail = torch.tensor([e[0] for e in edges], device=device)
 97        head = torch.tensor([e[1] for e in edges], device=device)
 98        d = x[:, tail] - x[:, head]
 99        target_p = 0.8*torch.tanh(d) + 0.25*u[:,:,0] + 0.35*torch.sin(d*u[:,:,1])
100        target_q = torch.einsum('nm,bm->bn', B, target_p)
101        train = slice(0, 900); test = slice(900, 1200)
102
103        class Conservative(torch.nn.Module):
104            def __init__(self):
105                super().__init__()
106                self.c0 = torch.nn.Sequential(torch.nn.Linear(3, 16), torch.nn.Tanh(), torch.nn.Linear(16, 1))
107                self.c1 = torch.nn.Sequential(torch.nn.Linear(3, 16), torch.nn.Tanh(), torch.nn.Linear(16, 1))
108            def forward(self, xx, uu):
109                dd = xx[:,tail] - xx[:,head]
110                z = torch.stack((dd, uu[:,:,0], uu[:,:,1]), dim=-1)
111                p = self.c0(z).squeeze(-1) + self.c1(z).squeeze(-1)
112                return torch.einsum('nm,bm->bn', B, p)
113
114        class Unconstrained(torch.nn.Module):
115            def __init__(self):
116                super().__init__()
117                self.edge = torch.nn.Sequential(torch.nn.Linear(4, 24), torch.nn.Tanh(), torch.nn.Linear(24, 2))
118            def forward(self, xx, uu):
119                dd = xx[:,tail] - xx[:,head]
120                z = torch.stack((xx[:,tail], xx[:,head], uu[:,:,0], uu[:,:,1]), dim=-1)
121                msg = self.edge(z)
122                out = torch.zeros((xx.shape[0], n), device=xx.device)
123                out.index_add_(1, tail, msg[:,:,0])
124                out.index_add_(1, head, msg[:,:,1])
125                return out
126
127        models = {"conservative": Conservative().to(device), "unconstrained": Unconstrained().to(device)}
128        results = {}
129        for name, model in models.items():
130            opt = torch.optim.Adam(model.parameters(), lr=3e-3)
131            for step in range(500):
132                opt.zero_grad()
133                pred = model(x[train], u[train])
134                loss = ((pred-target_q[train])**2).mean()
135                loss.backward(); opt.step()
136            with torch.no_grad():
137                pred = model(x[test], u[test])
138                err = ((pred-target_q[test])**2).mean().sqrt().item()
139                residual = (pred.sum(dim=1)).abs().mean().item()
140                # internal conservation is exact for conservative architecture;
141                # baseline residual is measured independently of task error.
142                results[name] = {"test_q_rmse": err, "mean_abs_mass_residual": residual,
143                                 "parameters": sum(p.numel() for p in model.parameters())}
144        return {"device": str(device), "results": results}
145    except Exception as exc:
146        return {"device": "cpu-fallback", "error": repr(exc)}
147
148
149def main():
150    out = {"seed": SEED, "core_math": core_sweeps(), "toy_model": toy_model_comparison()}
151    Path("results.json").write_text(json.dumps(out, indent=2))
152    print(json.dumps(out, indent=2))
153
154if __name__ == "__main__":
155    main()