ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | """Explicit toroidal interactions between computational atoms and field modes.""" | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| class ToroidalInteraction(nn.Module): | |
| """Interaction engine based on phase/frequency/coupling, not Transformer attention.""" | |
| def __init__(self, d_model: int = 256, n_modes: int = 256) -> None: | |
| super().__init__() | |
| self.d_model = d_model | |
| self.n_modes = n_modes | |
| def field_interaction(self, alpha: torch.Tensor) -> torch.Tensor: | |
| """Toroidal nearest-neighbour interaction energy over the field.""" | |
| if alpha.numel() == 0: | |
| return alpha.new_zeros(0) | |
| prev = torch.roll(alpha, 1, dims=0) | |
| nxt = torch.roll(alpha, -1, dims=0) | |
| coupling = torch.cosine_similarity(alpha, prev, dim=-1) | |
| forward = torch.cosine_similarity(alpha, nxt, dim=-1) | |
| local_gradient = (alpha - prev).pow(2).mean(dim=-1).sqrt() | |
| return 0.5 * (coupling + forward) - local_gradient | |
| def pairwise_interaction( | |
| self, | |
| atoms_r: torch.Tensor, | |
| atoms_phi: torch.Tensor, | |
| atoms_kappa: torch.Tensor, | |
| ) -> torch.Tensor: | |
| """K_ij = <kappa_i*kappa_j*cos(phi_i-phi_j)> with optional radial affinity.""" | |
| n = atoms_r.shape[0] | |
| if n == 0: | |
| return atoms_r.new_zeros((0, 0)) | |
| dphi = atoms_phi[:, None, :] - atoms_phi[None, :, :] | |
| phase = torch.cos(dphi) | |
| coupling = atoms_kappa[:, None, :] * atoms_kappa[None, :, :] | |
| radial = torch.exp(-0.5 * (atoms_r[:, None, :] - atoms_r[None, :, :]).pow(2).mean(dim=-1)) | |
| return (phase * coupling).mean(dim=-1) * radial | |
| def get_interaction_matrix(self, atoms_phi: torch.Tensor, atoms_kappa: torch.Tensor) -> torch.Tensor: | |
| return self.pairwise_interaction(torch.zeros_like(atoms_phi), atoms_phi, atoms_kappa) | |