ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | """Hierarchical aggregation and coherent MERGE driven by toroidal phase/energy.""" | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class AggregationEngine(nn.Module): | |
| def __init__( | |
| self, | |
| d_model: int = 256, | |
| phase_coherence_threshold: float = 0.45, | |
| energy_threshold: float = 0.3, | |
| max_depth: int = 5, | |
| merge_energy_floor: float = 0.08, | |
| ) -> None: | |
| super().__init__() | |
| self.d_model = d_model | |
| self.phase_coherence_threshold = phase_coherence_threshold | |
| self.energy_threshold = energy_threshold | |
| self.max_depth = max_depth | |
| self.merge_energy_floor = merge_energy_floor | |
| def compute_coherence(self, phi: torch.Tensor, E: torch.Tensor) -> torch.Tensor: | |
| dphi = phi[:, None, :] - phi[None, :, :] | |
| phase = torch.cos(dphi).mean(dim=-1) | |
| energy = torch.sqrt((E[:, None, :] * E[None, :, :]).clamp_min(0)) | |
| return phase * energy.mean(dim=-1) | |
| def find_clusters(self, coherence: torch.Tensor, min_size: int = 2) -> list[list[int]]: | |
| n = coherence.shape[0] | |
| if n == 0: | |
| return [] | |
| mask = coherence > self.phase_coherence_threshold | |
| mask.fill_diagonal_(False) | |
| visited = torch.zeros(n, dtype=torch.bool, device=coherence.device) | |
| clusters: list[list[int]] = [] | |
| for i in range(n): | |
| if visited[i]: | |
| continue | |
| queue = [i] | |
| visited[i] = True | |
| cluster: list[int] = [] | |
| while queue: | |
| node = queue.pop() | |
| cluster.append(node) | |
| neighbours = torch.nonzero(mask[node] & ~visited, as_tuple=False).flatten().tolist() | |
| for j in neighbours: | |
| visited[j] = True | |
| queue.append(j) | |
| if len(cluster) >= min_size: | |
| clusters.append(cluster) | |
| return clusters | |
| def _circular_mean(phi: torch.Tensor, weights: torch.Tensor | None = None) -> torch.Tensor: | |
| if weights is None: | |
| return torch.atan2(torch.sin(phi).mean(dim=0), torch.cos(phi).mean(dim=0)) | |
| w = weights.reshape(-1, *([1] * (phi.ndim - 1))) | |
| return torch.atan2((torch.sin(phi) * w).sum(dim=0), (torch.cos(phi) * w).sum(dim=0)) | |
| def _energy_weights(self, E: torch.Tensor, idx: torch.Tensor) -> torch.Tensor: | |
| weights = E[idx].mean(dim=-1).clamp_min(1e-6) | |
| return weights / weights.sum() | |
| def _combine_cluster( | |
| self, cluster, r, phi, omega, E, kappa, M, tau, rho | |
| ) -> dict: | |
| """Energy-weighted honest combine of r/φ/ω/E/κ/M/τ/ρ into one heavier atom.""" | |
| idx = torch.tensor(cluster, device=r.device, dtype=torch.long) | |
| weights = self._energy_weights(E, idx) | |
| r_m = F.normalize((r[idx] * weights[:, None]).sum(dim=0), dim=-1) | |
| phi_m = self._circular_mean(phi[idx], weights) | |
| omega_m = (omega[idx] * weights[:, None]).sum(dim=0) | |
| e_sum = (E[idx] * weights[:, None]).sum(dim=0) | |
| # Structural mass boost while staying in [0, 1]. | |
| e_boost = (1.0 - (1.0 - e_sum).clamp(0, 1) ** len(cluster)).clamp(0, 1) | |
| kappa_m = (kappa[idx] * weights[:, None]).sum(dim=0) | |
| M_m = F.normalize((M[idx] * weights[:, None]).sum(dim=0), dim=-1) | |
| tau_m = (tau[idx] * weights[:, None]).sum(dim=0) | |
| rho_m = rho[idx].max(dim=0).values + 1 | |
| return { | |
| "r": r_m, | |
| "phi": phi_m, | |
| "omega": omega_m, | |
| "E": e_boost, | |
| "kappa": kappa_m, | |
| "M": M_m, | |
| "tau": tau_m, | |
| "rho": rho_m, | |
| "constituents": list(cluster), | |
| } | |
| def aggregate(self, r, phi, omega, E, kappa, M, tau, rho) -> tuple[list[dict], list[list[int]]]: | |
| tensors = [x.unsqueeze(0) if x.dim() == 1 else x for x in (r, phi, omega, E, kappa, M, tau, rho)] | |
| if tensors[0].dim() == 3 and tensors[0].shape[-2] == 1: | |
| tensors = [x.squeeze(-2) for x in tensors] | |
| r, phi, omega, E, kappa, M, tau, rho = tensors | |
| if r.shape[0] < 2: | |
| return [], [] | |
| coherence = self.compute_coherence(phi, E) | |
| clusters = self.find_clusters(coherence) | |
| aggregates = [ | |
| self._combine_cluster(c, r, phi, omega, E, kappa, M, tau, rho) for c in clusters | |
| ] | |
| return aggregates, clusters | |
| def merge_coherent( | |
| self, | |
| r, phi, omega, E, kappa, M, tau, rho, | |
| *, | |
| max_merges: int = 8, | |
| ) -> tuple[list[dict], list[int], int, dict]: | |
| """MERGE phase-coherent, energetic atoms into heavier structure. | |
| Returns (merged_atom_dicts, remove_indices, merge_count, diag). | |
| """ | |
| tensors = [x.unsqueeze(0) if x.dim() == 1 else x for x in (r, phi, omega, E, kappa, M, tau, rho)] | |
| if tensors[0].dim() == 3 and tensors[0].shape[-2] == 1: | |
| tensors = [x.squeeze(-2) for x in tensors] | |
| r, phi, omega, E, kappa, M, tau, rho = tensors | |
| n = int(r.shape[0]) | |
| if n < 2: | |
| return [], [], 0, {"max_phase_coherence": 0.0, "n_pairs_above_energy_floor": 0} | |
| coherence = self.compute_coherence(phi, E) | |
| available = torch.ones(n, dtype=torch.bool, device=r.device) | |
| mean_E = E.mean(dim=-1) | |
| merged: list[dict] = [] | |
| remove: list[int] = [] | |
| merge_count = 0 | |
| tri = torch.triu(coherence, diagonal=1) | |
| # Read-only merge diagnostics: max coherence among candidate pairs | |
| # BEFORE the phase_coherence_threshold filter (default 0.45 so mph~0.5 fires). | |
| max_phase_coherence = float(tri.max().item()) if tri.numel() else 0.0 | |
| geo = torch.sqrt((mean_E[:, None] * mean_E[None, :]).clamp_min(0.0)) | |
| both_low = (mean_E[:, None] < self.merge_energy_floor) & (mean_E[None, :] < self.merge_energy_floor) | |
| geo_ok = geo >= (self.merge_energy_floor * 0.5) | |
| energy_ok = (~both_low) & geo_ok | |
| energy_ok = torch.triu(energy_ok, diagonal=1) | |
| n_pairs_above_energy_floor = int(energy_ok.sum().item()) | |
| flat = tri.reshape(-1) | |
| k = min(max(n * 2, 8), flat.numel()) | |
| top_vals, top_idx = torch.topk(flat, k=k) | |
| for score, flat_i in zip(top_vals.tolist(), top_idx.tolist()): | |
| if merge_count >= max_merges: | |
| break | |
| if score < self.phase_coherence_threshold: | |
| break | |
| i = flat_i // n | |
| j = flat_i % n | |
| if i == j or not bool(available[i] and available[j]): | |
| continue | |
| if float(mean_E[i]) < self.merge_energy_floor and float(mean_E[j]) < self.merge_energy_floor: | |
| continue | |
| if float(torch.sqrt(mean_E[i] * mean_E[j])) < self.merge_energy_floor * 0.5: | |
| continue | |
| cluster = [i, j] | |
| merged.append(self._combine_cluster(cluster, r, phi, omega, E, kappa, M, tau, rho)) | |
| remove.extend(cluster) | |
| available[i] = False | |
| available[j] = False | |
| merge_count += 1 | |
| diag = { | |
| "max_phase_coherence": max_phase_coherence, | |
| "n_pairs_above_energy_floor": n_pairs_above_energy_floor, | |
| } | |
| return merged, sorted(set(remove)), merge_count, diag | |
| def hierarchical_aggregate(self, r, phi, omega, E, kappa, M, tau, rho, depth: int = 3): | |
| depth = min(depth, self.max_depth) | |
| hierarchy = [] | |
| current = [x.clone() for x in (r, phi, omega, E, kappa, M, tau, rho)] | |
| for _ in range(depth): | |
| aggregates, _ = self.aggregate(*current) | |
| if not aggregates: | |
| break | |
| hierarchy.append(aggregates) | |
| if len(aggregates) < 2: | |
| break | |
| current = [ | |
| torch.stack([a[key] for a in aggregates]) | |
| for key in ("r", "phi", "omega", "E", "kappa", "M", "tau", "rho") | |
| ] | |
| return hierarchy | |