1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 233
  7np.random.seed(SEED)
  8random.seed(SEED)
  9torch.manual_seed(SEED)
 10
 11
 12def spd_np(d=3):
 13    x = np.random.randn(d, d)
 14    l = np.tril(x)
 15    l[np.diag_indices(d)] += 2.0
 16    return l @ l.T + 1e-6 * np.eye(d)
 17
 18
 19def sqrt_spd_np(x):
 20    x = (x + x.T) / 2
 21    w, v = np.linalg.eigh(x)
 22    return (v * np.sqrt(np.maximum(w, 1e-12))) @ v.T
 23
 24
 25def mutate_np(states, B, k, eps=0.0):
 26    d = states[0].shape[0]
 27    p = np.eye(d)
 28    n = np.eye(d)
 29    # Fixed ascending order is the arrow order in the construction.
 30    for i in range(len(states)):
 31        b = int(B[i, k])
 32        if b > 0:
 33            p = p @ np.linalg.matrix_power(states[i], b)
 34        elif b < 0:
 35            n = n @ np.linalg.matrix_power(states[i], -b)
 36    m = p + n
 37    s = sqrt_spd_np(m)
 38    out = list(states)
 39    out[k] = s @ np.linalg.inv(states[k]) @ s + eps * np.eye(d)
 40    return out, m, p, n
 41
 42
 43def math_check():
 44    B = np.array([[0, 1, 0], [-1, 0, 3], [0, -1, 0]], dtype=int)
 45    inverse_errors, conditions = [], []
 46    for k in [0, 1, 2]:
 47        states = [spd_np() for _ in range(3)]
 48        after, _, _, _ = mutate_np(states, B, k)
 49        back, _, _, _ = mutate_np(after, B, k)
 50        inverse_errors.append(float(np.linalg.norm(back[k] - states[k]) / np.linalg.norm(states[k])))
 51        conditions.append(float(np.linalg.cond(after[k])))
 52
 53    # A branch containing ordered products exposes the noncommutative issue:
 54    # products of SPD matrices are generally not symmetric, so M need not be SPD.
 55    B4 = np.array([[0, 1, 1, 0], [-1, 0, 1, 0], [-1, -1, 0, 1], [0, 0, -1, 0]], dtype=int)
 56    states4 = [spd_np() for _ in range(4)]
 57    after4, m4, p4, n4 = mutate_np(states4, B4, 2)
 58    back4, _, _, _ = mutate_np(after4, B4, 2)
 59    nonsym = float(np.linalg.norm(m4 - m4.T) / np.linalg.norm(m4))
 60    inverse4 = float(np.linalg.norm(back4[2] - states4[2]) / np.linalg.norm(states4[2]))
 61    return {
 62        'inverse_errors': inverse_errors,
 63        'max_inverse_error': max(inverse_errors),
 64        'updated_condition_numbers': conditions,
 65        'ordered_product_relative_nonsymmetry': nonsym,
 66        'four_node_inverse_error': inverse4,
 67        'four_node_product_norms': [float(np.linalg.norm(p4)), float(np.linalg.norm(n4))]
 68    }
 69
 70
 71class ResidualNet(nn.Module):
 72    def __init__(self, d, hidden=32):
 73        super().__init__()
 74        self.net = nn.Sequential(nn.Linear(d, hidden), nn.Tanh(), nn.Linear(hidden, 1))
 75    def forward(self, x):
 76        return self.net(x).squeeze(-1)
 77
 78
 79class MatrixClusterNet(nn.Module):
 80    def __init__(self, d=2, nodes=3, hidden=32):
 81        super().__init__()
 82        self.d, self.nodes = d, nodes
 83        self.readout = nn.Sequential(nn.Linear(nodes*d*d, hidden), nn.Tanh(), nn.Linear(hidden, 1))
 84    def forward(self, x):
 85        # x: batch,nodes,d,d; mutate node 1 using the weighted three-node quiver
 86        B = torch.tensor([[0.,1.,0.],[-1.,0.,3.],[0.,-1.,0.]], device=x.device)
 87        a = x[:, 1]
 88        eye = torch.eye(self.d, device=x.device).expand_as(a)
 89        p = eye.clone(); n = eye.clone()
 90        # For k=1: P=A0, N=A2^1 (ordered products are trivial here).
 91        p = x[:, 0]
 92        n = x[:, 2]
 93        m = (p+n + (p+n).transpose(-1,-2))/2
 94        w,v = torch.linalg.eigh(m)
 95        s = (v * torch.sqrt(torch.clamp(w, min=1e-5)).unsqueeze(-2)) @ v.transpose(-1,-2)
 96        updated = s @ torch.linalg.inv(a) @ s
 97        z = x.clone(); z[:,1] = updated
 98        return self.readout(z.reshape(x.shape[0], -1)).squeeze(-1)
 99
100
101def training_check(device):
102    torch.manual_seed(SEED)
103    n, nodes, d = 512, 3, 2
104    raw = torch.randn(n, nodes, d, d, device=device)
105    raw = raw @ raw.transpose(-1,-2) + 0.4*torch.eye(d, device=device)
106    # Target is a relational matrix quantity, with a mild nonlinear scalar map.
107    target = torch.logdet(raw[:,0] + raw[:,2]) - 0.25*torch.logdet(raw[:,1])
108    models = {'baseline': ResidualNet(nodes*d*d), 'matrix_cluster': MatrixClusterNet(d, nodes)}
109    # Match the baseline input/output role; both are small and trained identically.
110    xflat = raw.reshape(n, -1)
111    results = {}
112    for name, model in models.items():
113        model.to(device); opt = torch.optim.Adam(model.parameters(), lr=3e-3)
114        losses=[]
115        for step in range(180):
116            opt.zero_grad()
117            pred = model(xflat) if name == 'baseline' else model(raw)
118            loss = ((pred-target)**2).mean()
119            loss.backward(); opt.step()
120            if step in (0, 59, 119, 179): losses.append(float(loss.detach().cpu()))
121        with torch.no_grad():
122            pred = model(xflat) if name == 'baseline' else model(raw)
123            final = float(((pred-target)**2).mean().cpu())
124        results[name] = {'loss_checkpoints': losses, 'final_mse': final, 'parameters': sum(p.numel() for p in model.parameters())}
125    return results
126
127
128def main():
129    try:
130        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
131        math_result = math_check()
132        train_result = training_check(device)
133        used = str(device)
134    except Exception as exc:
135        device = torch.device('cpu')
136        math_result = math_check()
137        train_result = training_check(device)
138        used = 'cpu_fallback:' + repr(exc)
139    out = {'device': used, 'math': math_result, 'training': train_result}
140    print(json.dumps(out, indent=2))
141    with open('results.json', 'w') as f:
142        json.dump(out, f, indent=2)
143
144if __name__ == '__main__':
145    main()