1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 2967
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_num_threads(4)
  9
 10
 11def projectors(normals):
 12    I = torch.eye(3, dtype=normals.dtype, device=normals.device)
 13    return I[None] - normals[..., :, None] * normals[..., None, :]
 14
 15
 16def estimate_normals(x, k=12):
 17    # Small, explicit kNN covariance estimator; normal orientation is irrelevant to P.
 18    d = torch.cdist(x, x)
 19    nn_idx = d.topk(k + 1, largest=False).indices[:, 1:]
 20    neigh = x[nn_idx] - x[:, None, :]
 21    cov = torch.einsum('nki,nkj->nij', neigh, neigh) / float(k)
 22    vals, vecs = torch.linalg.eigh(cov)
 23    n = vecs[:, :, 0]
 24    return n / (n.norm(dim=-1, keepdim=True) + 1e-12), nn_idx
 25
 26
 27def rotation(seed):
 28    g = torch.Generator().manual_seed(seed)
 29    z = torch.randn(3, 3, generator=g)
 30    q, r = torch.linalg.qr(z)
 31    q = q @ torch.diag(torch.sign(torch.diag(r)))
 32    if torch.linalg.det(q) < 0: q[:, 0] *= -1
 33    return q
 34
 35
 36def make_cloud(seed, n=240, noise=0.8):
 37    g = torch.Generator().manual_seed(seed)
 38    x = torch.randn(n, 3, generator=g)
 39    x = x / x.norm(dim=1, keepdim=True)
 40    R = rotation(10000 + seed)
 41    x = x @ R.T
 42    # A smooth scalar field and its exact tangent gradient on the unit sphere.
 43    a = torch.tensor([0.8, -0.3, 0.5])
 44    f = x @ a
 45    normal = x
 46    clean_v = a[None] - f[:, None] * normal
 47    # Deliberate ambient normal corruption, which intrinsic projection should discard.
 48    eps = noise * torch.randn(n, 1, generator=g)
 49    v = clean_v + eps * normal
 50    target = clean_v.norm(dim=1)
 51    return x, f[:, None], v, target
 52
 53
 54class LocalLayer(nn.Module):
 55    def __init__(self, intrinsic=False, hidden=48):
 56        super().__init__(); self.intrinsic = intrinsic
 57        # [s_i,s_j,v_i,v_j,r,d] = 12 numbers
 58        self.mlp = nn.Sequential(nn.Linear(12, hidden), nn.SiLU(), nn.Linear(hidden, hidden), nn.SiLU(), nn.Linear(hidden, 1))
 59
 60    def forward(self, x, s, v):
 61        n, idx = estimate_normals(x)
 62        P = projectors(n)
 63        if self.intrinsic: v = torch.einsum('nij,nj->ni', P, v)
 64        xj = x[idx]; sj = s[idx]; vj = v[idx]
 65        delta = xj - x[:, None, :]
 66        if self.intrinsic: r = torch.einsum('nij,nkj->nki', P, delta)
 67        else: r = delta
 68        d = r.norm(dim=-1, keepdim=True)
 69        vi = v[:, None, :].expand_as(vj)
 70        inp = torch.cat([s[:, None, :].expand_as(sj), sj, vi, vj, r, d], dim=-1)
 71        msg = self.mlp(inp).mean(dim=1)
 72        return msg[:, 0], n, P
 73
 74
 75def math_check():
 76    # Verify P is symmetric/idempotent, removes normal output, and PJP removes
 77    # both normal output and normal input derivative directions.
 78    torch.manual_seed(11)
 79    n = torch.randn(100, 3); n = n / n.norm(dim=1, keepdim=True); P = projectors(n)
 80    sym = (P - P.transpose(1, 2)).abs().max().item()
 81    idem = (P @ P - P).abs().max().item()
 82    normal_removed = (torch.einsum('ni,nij->nj', n, P)).abs().max().item()
 83    J = torch.randn(100, 3, 3)
 84    K = P @ J @ P
 85    left = torch.einsum('ni,nij->nj', n, K).abs().max().item()
 86    right = torch.einsum('nij,nj->ni', K, n).abs().max().item()
 87    return dict(projector_symmetry=sym, projector_idempotence=idem,
 88                normal_removed=normal_removed, jacobian_left_normal=left,
 89                jacobian_right_normal=right)
 90
 91
 92def train_eval(intrinsic, epochs=65):
 93    model = LocalLayer(intrinsic=intrinsic)
 94    opt = torch.optim.Adam(model.parameters(), lr=3e-3, weight_decay=1e-5)
 95    train = [make_cloud(i, noise=0.8) for i in range(8)]
 96    test = [make_cloud(100+i, noise=1.5) for i in range(4)]
 97    for ep in range(epochs):
 98        random.shuffle(train); model.train()
 99        for x,s,v,y in train:
100            pred,_,_ = model(x,s,v)
101            loss = ((pred-y)**2).mean()
102            opt.zero_grad(); loss.backward(); opt.step()
103    model.eval(); errs=[]
104    with torch.no_grad():
105        for x,s,v,y in test:
106            p,n,P = model(x,s,v); errs.append(torch.mean((p-y)**2).sqrt().item() / (y.pow(2).mean().sqrt().item()+1e-8))
107    return float(np.mean(errs)), model
108
109
110def main():
111    check = math_check()
112    b, bm = train_eval(False); i, im = train_eval(True)
113    # Mechanism observable: projection gives machine-zero normal velocity after each update.
114    x,s,v,y = make_cloud(999, noise=2.0)
115    with torch.no_grad():
116        _, n, P = im(x,s,v); vp = torch.einsum('nij,nj->ni', P, v)
117        leakage = (torch.abs((n*vp).sum(1))).max().item()
118        raw_leakage = (torch.abs((n*v).sum(1))).mean().item()
119    out = {'math_check': check, 'relative_rmse_baseline': b,
120           'relative_rmse_intrinsic': i, 'relative_rmse_improvement_pct': 100*(b-i)/b,
121           'raw_mean_normal_velocity': raw_leakage, 'projected_max_normal_velocity': leakage,
122           'seed': SEED}
123    print(json.dumps(out, indent=2))
124
125if __name__ == '__main__': main()