File size: 3,276 Bytes
b9301f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | """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)
@torch.no_grad()
def add_atom_contribution(self, r, phi, omega, E, kappa) -> None:
self.alpha.add_(self._projection(r, phi, omega, E, kappa))
@torch.no_grad()
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
|