ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | """ | |
| ToroidalAtom: a single computational primitive with 8 properties. | |
| Design choice: atoms are represented as *tensors* (not Python objects) | |
| so that the entire collection can be batched and differentiable. | |
| Each property is a torch.Tensor of shape [d_model] (or scalar where appropriate). | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| class ToroidalAtom(nn.Module): | |
| """ | |
| A single toroidal atom. | |
| Properties (all differentiable tensors): | |
| r [d_model] — position / radius direction in latent space | |
| phi [d_model] — phase (angle latent) | |
| omega [d_model] — natural frequency (positive) | |
| E [d_model] — energy / activation / importance (0..1) | |
| kappa [d_model] — coupling strength (positive) | |
| M [d_model] — local memory direction | |
| tau [d_model] — temporal scale (positive) | |
| rho [1] — hierarchical level (integer) | |
| """ | |
| def __init__( | |
| self, | |
| r: torch.Tensor, | |
| phi: torch.Tensor, | |
| omega: torch.Tensor, | |
| E: torch.Tensor, | |
| kappa: torch.Tensor, | |
| M: torch.Tensor, | |
| tau: torch.Tensor, | |
| rho: torch.Tensor, | |
| payload: bytes = b"", | |
| ) -> None: | |
| super().__init__() | |
| self.register_buffer("r", r) | |
| self.register_buffer("phi", phi) | |
| self.register_buffer("omega", omega) | |
| self.register_buffer("E", E) | |
| self.register_buffer("kappa", kappa) | |
| self.register_buffer("M", M) | |
| self.register_buffer("tau", tau) | |
| self.register_buffer("rho", rho) | |
| # Surface bytes carried by this living atom (not a buffer; episode memory). | |
| self.payload = bytes(payload) if payload else b"" | |
| def energy(self) -> torch.Tensor: | |
| """Scalar energy of this atom.""" | |
| return self.E.sum() | |
| def phase_distance(self, other: "ToroidalAtom") -> torch.Tensor: | |
| """Phase difference (wrapped) with another atom.""" | |
| return torch.remainder(self.phi - other.phi, 2 * torch.pi) | |
| def coupling_to(self, other: "ToroidalAtom") -> torch.Tensor: | |
| """Pairwise coupling K_ij = kappa_ij * cos(phi_i - phi_j).""" | |
| dphi = self.phase_distance(other) | |
| return self.kappa * other.kappa * torch.cos(dphi) | |
| def clone(self) -> "ToroidalAtom": | |
| return ToroidalAtom( | |
| r=self.r.clone(), | |
| phi=self.phi.clone(), | |
| omega=self.omega.clone(), | |
| E=self.E.clone(), | |
| kappa=self.kappa.clone(), | |
| M=self.M.clone(), | |
| tau=self.tau.clone(), | |
| rho=self.rho.clone(), | |
| payload=self.payload, | |
| ) | |
| def to(self, device): | |
| return ToroidalAtom( | |
| r=self.r.to(device), | |
| phi=self.phi.to(device), | |
| omega=self.omega.to(device), | |
| E=self.E.to(device), | |
| kappa=self.kappa.to(device), | |
| M=self.M.to(device), | |
| tau=self.tau.to(device), | |
| rho=self.rho.to(device), | |
| payload=self.payload, | |
| ) | |
| class ToroidalAtomCollection(nn.Module): | |
| """ | |
| A batch/collection of toroidal atoms. | |
| Stores atoms as stacked tensors for vectorized operations. | |
| """ | |
| def __init__(self, atoms: list[ToroidalAtom] | None = None) -> None: | |
| super().__init__() | |
| self.atoms: list[ToroidalAtom] = atoms or [] | |
| self._buffers: dict[str, torch.Tensor] = {} | |
| self._dirty = True | |
| if atoms: | |
| self._rebuild_buffers() | |
| def _rebuild_buffers(self) -> None: | |
| if not self.atoms: | |
| self._buffers = {} | |
| self._dirty = False | |
| return | |
| self._buffers = { | |
| "r": torch.stack([a.r for a in self.atoms]), | |
| "phi": torch.stack([a.phi for a in self.atoms]), | |
| "omega": torch.stack([a.omega for a in self.atoms]), | |
| "E": torch.stack([a.E for a in self.atoms]), | |
| "kappa": torch.stack([a.kappa for a in self.atoms]), | |
| "M": torch.stack([a.M for a in self.atoms]), | |
| "tau": torch.stack([a.tau for a in self.atoms]), | |
| "rho": torch.stack([a.rho for a in self.atoms]), | |
| } | |
| self._dirty = False | |
| def _ensure_buffers(self) -> None: | |
| if getattr(self, "_dirty", True): | |
| self._rebuild_buffers() | |
| def r(self) -> torch.Tensor: | |
| self._ensure_buffers() | |
| return self._buffers.get("r", torch.zeros(0)) | |
| def phi(self) -> torch.Tensor: | |
| self._ensure_buffers() | |
| return self._buffers.get("phi", torch.zeros(0)) | |
| def omega(self) -> torch.Tensor: | |
| self._ensure_buffers() | |
| return self._buffers.get("omega", torch.zeros(0)) | |
| def E(self) -> torch.Tensor: | |
| self._ensure_buffers() | |
| return self._buffers.get("E", torch.zeros(0)) | |
| def kappa(self) -> torch.Tensor: | |
| self._ensure_buffers() | |
| return self._buffers.get("kappa", torch.zeros(0)) | |
| def M(self) -> torch.Tensor: | |
| self._ensure_buffers() | |
| return self._buffers.get("M", torch.zeros(0)) | |
| def tau(self) -> torch.Tensor: | |
| self._ensure_buffers() | |
| return self._buffers.get("tau", torch.zeros(0)) | |
| def rho(self) -> torch.Tensor: | |
| self._ensure_buffers() | |
| return self._buffers.get("rho", torch.zeros(0)) | |
| def __len__(self) -> int: | |
| return len(self.atoms) | |
| def add(self, atom: ToroidalAtom) -> None: | |
| self.atoms.append(atom) | |
| self._dirty = True | |
| def extend(self, atoms: list[ToroidalAtom]) -> None: | |
| if atoms: | |
| self.atoms.extend(atoms) | |
| self._dirty = True | |
| def remove(self, indices: list[int]) -> None: | |
| # Remove in reverse order to preserve indices | |
| for i in sorted(indices, reverse=True): | |
| self.atoms.pop(i) | |
| self._dirty = True | |
| def top_by_energy(self, k: int) -> list[ToroidalAtom]: | |
| energies = self.E.sum(dim=-1) if self.E.numel() else torch.zeros(len(self.atoms)) | |
| top_k = energies.topk(min(k, len(self.atoms)), dim=0).indices | |
| return [self.atoms[i.item()] for i in top_k] | |
| def to(self, device): | |
| new_atoms = [a.to(device) for a in self.atoms] | |
| collection = ToroidalAtomCollection(new_atoms) | |
| return collection | |
| def state_dict(self): | |
| return { | |
| "atoms": [ | |
| { | |
| "r": a.r.cpu(), "phi": a.phi.cpu(), "omega": a.omega.cpu(), | |
| "E": a.E.cpu(), "kappa": a.kappa.cpu(), "M": a.M.cpu(), | |
| "tau": a.tau.cpu(), "rho": a.rho.cpu(), | |
| "payload": bytes(getattr(a, "payload", b"") or b""), | |
| } | |
| for a in self.atoms | |
| ] | |
| } | |
| def load_state_dict(self, state: dict) -> None: | |
| self.atoms = [] | |
| for sd in state["atoms"]: | |
| self.atoms.append(ToroidalAtom( | |
| r=sd["r"], phi=sd["phi"], omega=sd["omega"], | |
| E=sd["E"], kappa=sd["kappa"], M=sd["M"], | |
| tau=sd["tau"], rho=sd["rho"], | |
| payload=bytes(sd.get("payload", b"") or b""), | |
| )) | |
| self._rebuild_buffers() | |