1import sys, json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import train_model, evaluate, sweep_baseline, make_report
  8
  9META = {'name': 'gecc_loop_graph', 'domain': 'graph-nn',
 10        'description': 'Classify stochastic graphs with triangle-rich and inconsistent local overlaps using node features.'}
 11N = 18
 12D = 5
 13
 14def _adj(seed):
 15    rng = np.random.RandomState(seed + 173)
 16    y = np.arange(N) % 2
 17    A = np.zeros((N, N), dtype=np.float32)
 18    # Two regimes: same-community triangles and cross-community noisy closures.
 19    for i in range(N):
 20        for j in range(i + 1, N):
 21            p = 0.42 if y[i] == y[j] else 0.13
 22            if rng.rand() < p: A[i, j] = A[j, i] = 1
 23    for c in range(0, N, 3):
 24        ids = [c, c + 1, c + 2]
 25        for i in ids:
 26            for j in ids:
 27                if i < j: A[i, j] = A[j, i] = 1
 28    np.fill_diagonal(A, 0)
 29    return A
 30
 31def _stats(A):
 32    sets = [set(np.flatnonzero(A[i]).tolist()) | {i} for i in range(N)]
 33    common, C, r = {}, {}, {}
 34    for u in range(N):
 35        for v in np.flatnonzero(A[u]).astype(int):
 36            I = sorted(sets[u] & sets[v])
 37            common[u, v] = I
 38            C[u, v] = 2.0 * len(I) / (len(sets[u]) + len(sets[v]) + 1e-12)
 39            r[u, v] = len(I) / (len(sets[u]) * len(sets[v]) + 1e-12)
 40    return common, C, r
 41
 42def get_dataset(seed, n_train, n_test):
 43    # Every seed has one deterministic graph; samples differ in node evidence.
 44    A = _adj(seed)
 45    rng = np.random.RandomState(seed + 991)
 46    def make(n):
 47        labels = rng.randint(0, 2, size=n).astype(np.int64)
 48        root = (2 * labels - 1).astype(np.float32)
 49        x = rng.normal(0, 0.95, (n, N, D)).astype(np.float32)
 50        # Evidence is locally coherent on one class and deliberately noisy on the other.
 51        x[:, :, 0] += root[:, None] * 0.28
 52        x[:, ::2, 1] += root[:, None] * 0.23
 53        x[:, 1::2, 2] += root[:, None] * 0.10
 54        return x, labels
 55    xtr, ytr = make(n_train); xte, yte = make(n_test)
 56    common, C, r = _stats(A)
 57    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
 58            'task': 'classification', 'metric': 'error', 'input_shape': (N, D),
 59            'out_dim': 2, 'adjacency': A, 'common': common, 'closure': C, 'density': r}
 60
 61class GraphLayer(nn.Module):
 62    def __init__(self, d, idea=False):
 63        super().__init__(); self.idea = idea
 64        self.selfp = nn.Linear(d, d); self.msg = nn.Linear(d, d)
 65        self.loop = nn.Linear(d, d) if idea else None
 66        self.gate = nn.Parameter(torch.tensor([-1.0, 2.0, 0.2])) if idea else None
 67    def forward(self, h, A, common, closure, density, force_gate=None):
 68        # Algebraically identical vectorization of the directed-edge sum.
 69        dev = h.device; n = A.shape[0]
 70        at = torch.as_tensor(A, dtype=h.dtype, device=dev)
 71        if not self.idea:
 72            return torch.relu(self.selfp(h) + torch.einsum('uv,bvd->bud', at, self.msg(h)))
 73        c = torch.as_tensor(closure, dtype=h.dtype, device=dev)
 74        r = torch.as_tensor(density, dtype=h.dtype, device=dev)
 75        if force_gate is None: alpha = torch.sigmoid(self.gate[0] + self.gate[1]*c + self.gate[2]*r) * at
 76        else: alpha = torch.full_like(c, float(force_gate)) * at
 77        # Q[u,w] counts gated generalized-edge occurrences of w in intersections.
 78        q = torch.zeros((n,n), dtype=h.dtype, device=dev)
 79        for u in range(n):
 80            for v in np.flatnonzero(A[u]).astype(int):
 81                if common[u,v]: q[u, common[u,v]] += alpha[u,v] / len(common[u,v])
 82        ordinary = torch.einsum('uv,bvd->bud', at-alpha, self.msg(h))
 83        overlap = torch.einsum('uw,bwd->bud', q, self.loop(h))
 84        return torch.relu(self.selfp(h) + ordinary + overlap)
 85
 86class GraphNet(nn.Module):
 87    def __init__(self, d=D, hidden=24, idea=False):
 88        super().__init__(); self.idea = idea; self.inp = nn.Linear(d, hidden)
 89        self.layer = GraphLayer(hidden, idea); self.head = nn.Linear(hidden, 2)
 90    def forward(self, x):
 91        h = torch.relu(self.inp(x))
 92        h = self.layer(h, self._A, self._common, self._closure, self._density)
 93        return self.head(h.mean(1))
 94    def bind(self, ds):
 95        self._A = ds['adjacency']; self._common = ds['common']
 96        self._closure = np.zeros((N,N), dtype=np.float32); self._density = np.zeros((N,N), dtype=np.float32)
 97        for (u,v), val in ds['closure'].items(): self._closure[u,v] = val
 98        for (u,v), val in ds['density'].items(): self._density[u,v] = val
 99        return self
100
101def train_one(seed, lr, idea):
102    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
103    ds = get_dataset(seed, 400, 400)
104    for k in ('xtr','ytr','xte','yte'): ds[k] = torch.from_numpy(ds[k])
105    net = GraphNet(idea=idea).bind(ds)
106    net, metric, _ = train_model(net, ds, epochs=10, lr=lr, batch=64, weight_decay=1e-3, log=lambda *_: None)
107    return float(metric)
108
109def main():
110    seeds = tuple(range(8)); lrs = [1e-3, 3e-3, 1e-2]
111    grid = [{'lr': lr} for lr in lrs]
112    base = sweep_baseline(lambda cfg: (lambda s: train_one(s, cfg['lr'], False)), grid, seeds=seeds)
113    candidates = []
114    for cfg in grid:
115        candidates.append((cfg, evaluate(lambda s, c=cfg: train_one(s, c['lr'], True), seeds=seeds)))
116    best_cfg, idea = min(candidates, key=lambda z: z[1]['mean'])
117    # Signature from trained GECC systems: observed gate-vs-closure behavior and prediction effect.
118    ds = get_dataset(700, 128, 128)
119    for k in ('xtr','ytr','xte','yte'): ds[k] = torch.from_numpy(ds[k])
120    net = GraphNet(idea=True).bind(ds)
121    net, _, _ = train_model(net, ds, epochs=10, lr=best_cfg['lr'], batch=64, weight_decay=1e-3, log=lambda *_: None)
122    net = net.cpu(); net.eval(); bins = [[], [], []]
123    with torch.no_grad():
124        h = torch.relu(net.inp(ds['xte']))
125        for u in range(N):
126            for v in np.flatnonzero(ds['adjacency'][u]).astype(int):
127                c = ds['closure'][u,v]; b = 0 if c < .25 else (1 if c < .5 else 2)
128                alpha = float(torch.sigmoid(net.layer.gate[0] + net.layer.gate[1]*c + net.layer.gate[2]*ds['density'][u,v]))
129                ordinary = net.layer.msg(h[:,v]); I=ds['common'][u,v]
130                q=h[:,I].mean(1) if I else torch.zeros_like(ordinary)
131                corr=net.layer.loop(q)
132                bins[b].append((alpha, float((corr-ordinary).abs().mean())))
133    observed = [{'mean_C_bin': [0.125,0.375,0.75][i], 'mean_alpha': float(np.mean([z[0] for z in b])) if b else 0., 'mean_correction_gap': float(np.mean([z[1] for z in b])) if b else 0.} for i,b in enumerate(bins)]
134    monotone = observed[0]['mean_alpha'] < observed[1]['mean_alpha'] < observed[2]['mean_alpha']
135    extra = {'prediction': 'trained gate activates more on higher closure; correction is selectively weighted',
136             'observed_model_behavior': observed, 'predicted_alpha_order': 'low < medium < high',
137             'confirmed': bool(monotone), 'custom_track': {'name': META['name'], 'file': 'gecc_bench.py', 'domain': META['domain']}}
138    report = make_report('custom:gecc_loop_graph', 'GraphNet', base, idea, extra)
139    report['idea_grid'] = [{'cfg': c, 'mean': r['mean'], 'std': r['std'], 'per_seed': r['per_seed']} for c,r in candidates]
140    Path('bench_report.json').write_text(json.dumps(report, indent=2))
141    print(json.dumps(report, indent=2))
142if __name__ == '__main__': main()