ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | """Compact spectral/toroidal superposition field. | |
| Atoms are projected into a fixed periodic set of modes. The field is state, | |
| not a Transformer activation stack: atoms remain the structural memory while | |
| ``alpha`` is their compact superposition used by the continuous dynamics. | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| class FractalSuperpositionState(nn.Module): | |
| def __init__(self, n_modes: int = 256, d_model: int = 256) -> None: | |
| super().__init__() | |
| self.n_modes = n_modes | |
| self.d_model = d_model | |
| self.alpha = nn.Parameter(torch.zeros(n_modes, d_model)) | |
| self.t = nn.Parameter(torch.tensor(0.0), requires_grad=False) | |
| self.register_buffer( | |
| "mode_phase", | |
| torch.arange(n_modes, dtype=torch.float32) * (2.0 * torch.pi / n_modes), | |
| ) | |
| def _projection( | |
| self, | |
| r: torch.Tensor, | |
| phi: torch.Tensor, | |
| omega: torch.Tensor, | |
| E: torch.Tensor, | |
| kappa: torch.Tensor, | |
| ) -> torch.Tensor: | |
| """Project atoms continuously across toroidal modes. | |
| Frequency selects the centre of the mode distribution, phase controls | |
| its periodic modulation, and ``r`` carries the vector-valued content. | |
| No hard ``.item()`` indexing is used, so the projection remains smooth. | |
| """ | |
| if r.ndim == 1: | |
| r, phi, omega, E, kappa = [x.unsqueeze(0) for x in (r, phi, omega, E, kappa)] | |
| mode_phase = self.mode_phase.to(device=r.device, dtype=r.dtype).view(1, self.n_modes, 1) | |
| phase = phi.mean(dim=-1, keepdim=True).unsqueeze(1) | |
| freq = omega.mean(dim=-1, keepdim=True).unsqueeze(1).abs() | |
| freq_angle = 2.0 * torch.pi * freq / (1.0 + freq) | |
| centre = (freq_angle / (2.0 * torch.pi)) * self.n_modes | |
| mode_index = torch.arange(self.n_modes, device=r.device, dtype=r.dtype).view(1, self.n_modes, 1) | |
| # Wrapped circular distance keeps the projection toroidal. | |
| delta = torch.remainder(mode_index - centre + self.n_modes / 2, self.n_modes) - self.n_modes / 2 | |
| width = 1.0 + 0.25 * torch.sigmoid(freq) | |
| kernel = torch.exp(-(delta / width) ** 2) | |
| modulation = 0.5 + 0.5 * torch.cos(mode_phase - phase) | |
| strength = (E * kappa).mean(dim=-1, keepdim=True).unsqueeze(1) | |
| vectors = r.unsqueeze(1) * strength | |
| contribution = vectors * kernel * modulation | |
| return contribution.sum(dim=0) | |
| def add_atom_contribution(self, r, phi, omega, E, kappa) -> None: | |
| self.alpha.add_(self._projection(r, phi, omega, E, kappa)) | |
| def remove_atom_contribution(self, r, phi, omega, E, kappa) -> None: | |
| self.alpha.sub_(self._projection(r, phi, omega, E, kappa)) | |
| def get_field(self) -> torch.Tensor: | |
| return self.alpha | |
| def get_state_dict(self) -> dict: | |
| return {"alpha": self.alpha.detach().cpu().clone(), "t": self.t.detach().cpu().clone()} | |
| def load_state_dict(self, state: dict, strict: bool = True): | |
| with torch.no_grad(): | |
| self.alpha.copy_(state["alpha"].to(self.alpha.device)) | |
| self.t.copy_(state["t"].to(self.t.device)) | |
| def to(self, device): | |
| super().to(device) | |
| return self | |