| """nika-5: grid β latent-mesh world model of the Sun. Pure torch. |
| |
| Observations live on a dense grid; state lives on a coarse mesh. One encoder |
| shared across all AIA channels conditioned on log-T, one for the field vector. |
| The readout is a function of (cell latent, log-T): the slider is the readout. |
| """ |
| from dataclasses import dataclass |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from lib import AIA_CHANNELS, AIA_SCALE, B_SCALE, LOGT, fourier, mlp, norm_logt, scatter_mean |
|
|
|
|
| |
|
|
| @dataclass |
| class SolarGraph: |
| grid_pos: torch.Tensor |
| mu: torch.Tensor |
| g2m: torch.Tensor |
| g2m_feat: torch.Tensor |
| m2m: torch.Tensor |
| m2m_feat: torch.Tensor |
| n_mesh: int |
|
|
| def to(self, device): |
| moved = {f: getattr(self, f).to(device) |
| for f in ("grid_pos", "mu", "g2m", "g2m_feat", "m2m", "m2m_feat")} |
| return SolarGraph(**moved, n_mesh=self.n_mesh) |
|
|
|
|
| def _lattice(side): |
| ax = (torch.arange(side) + 0.5) / side |
| y, x = torch.meshgrid(ax, ax, indexing="ij") |
| return torch.stack([x.flatten(), y.flatten()], -1) |
|
|
|
|
| def _mesh_edges(side, offsets): |
| idx = torch.arange(side * side).view(side, side) |
| src, dst = [], [] |
| for dy, dx in offsets: |
| s = idx[max(0, -dy): side - max(0, dy), max(0, -dx): side - max(0, dx)] |
| d = idx[max(0, dy): side + min(0, dy), max(0, dx): side + min(0, dx)] |
| src.append(s.flatten()) |
| dst.append(d.flatten()) |
| e = torch.stack([torch.cat(src), torch.cat(dst)], -1) |
| return torch.cat([e, e.flip(-1)]) |
|
|
|
|
| def build_graph(res=256, mesh_side=48): |
| grid = _lattice(res) |
| r = ((grid - 0.5).square().sum(-1)).sqrt() / 0.5 |
| mu = (1 - r.square()).clamp(min=0).sqrt() |
|
|
| mesh = _lattice(mesh_side) |
| near = [(0, 1), (1, 0), (1, 1), (1, -1)] |
| far = [(0, 2), (2, 0), (0, 4), (4, 0)] |
| m2m = _mesh_edges(mesh_side, near + far) |
| scale = 1.0 / mesh_side |
| rel = (mesh[m2m[:, 0]] - mesh[m2m[:, 1]]) / scale |
| m2m_feat = torch.cat([rel, rel.norm(dim=-1, keepdim=True)], -1) |
|
|
| cell = grid * mesh_side - 0.5 |
| base = cell.floor().long() |
| corners = torch.tensor([[0, 0], [0, 1], [1, 0], [1, 1]]) |
| m_yx = (base[:, None] + corners[None]).clamp(0, mesh_side - 1) |
| m_idx = (m_yx[..., 1] * mesh_side + m_yx[..., 0]).flatten() |
| g_idx = torch.arange(len(grid)).repeat_interleave(4) |
| g2m = torch.stack([g_idx, m_idx], -1) |
| rel = (grid[g_idx] - mesh[m_idx]) / scale |
| g2m_feat = torch.cat([rel, rel.norm(dim=-1, keepdim=True)], -1) |
|
|
| return SolarGraph(grid, mu, g2m, g2m_feat, m2m, m2m_feat, n_mesh=mesh_side**2) |
|
|
|
|
| |
|
|
| @dataclass |
| class Observation: |
| aia: torch.Tensor |
| field: torch.Tensor |
| present: torch.Tensor |
|
|
| @staticmethod |
| def pack(frames, res=256): |
| pool = lambda x: torch.as_tensor(x).view(res, 512 // res, res, 512 // res).mean((1, 3)) |
| aia = torch.stack([ |
| torch.stack([pool(f[c]) if c in f else torch.zeros(res, res) for c in AIA_CHANNELS]) |
| for f in frames]).flatten(2) |
| field = torch.stack([ |
| torch.stack([pool(f[c]) for c in ("Bx", "By", "Bz")]) for f in frames]).flatten(2) |
| present = torch.tensor([[c in f for c in AIA_CHANNELS] for f in frames]) |
| return Observation(torch.asinh(aia / AIA_SCALE[:, None]), torch.tanh(field / B_SCALE), present) |
|
|
| def masked(self, hide): |
| vis = self.present.clone() |
| vis[:, hide] = False |
| return Observation(self.aia * vis[:, :, None], self.field, vis) |
|
|
| def to(self, device): |
| return Observation(self.aia.to(device), self.field.to(device), self.present.to(device)) |
|
|
|
|
| |
|
|
| class ChannelEncoder(nn.Module): |
| """One encoder for every AIA channel: (intensity, log-T) -> d. Sum over present.""" |
|
|
| def __init__(self, d): |
| super().__init__() |
| self.net = mlp(2, d, d) |
|
|
| def forward(self, aia, present): |
| logt = norm_logt(LOGT.to(aia.device))[:, None].expand_as(aia) |
| return (self.net(torch.stack([aia, logt], -1)) * present[:, :, None, None]).sum(1) |
|
|
|
|
| class FieldEncoder(nn.Module): |
| """Magnetic vector -> d. The causal modality gets its own pathway.""" |
|
|
| def __init__(self, d): |
| super().__init__() |
| self.net = mlp(3, d, d) |
|
|
| def forward(self, field): |
| return self.net(field.transpose(1, 2)) |
|
|
|
|
| class GeometryEmbed(nn.Module): |
| """ΞΌ + Fourier position -> d. Static, computed once per graph.""" |
|
|
| def __init__(self, d): |
| super().__init__() |
| self.net = nn.Linear(17, d) |
|
|
| def forward(self, graph): |
| return self.net(torch.cat([graph.mu[:, None], fourier(graph.grid_pos)], -1)) |
|
|
|
|
| class GridToMesh(nn.Module): |
| """Lift cell tokens onto the latent mesh along g2m edges.""" |
|
|
| def __init__(self, d): |
| super().__init__() |
| self.edge = mlp(d + 3, d, d) |
|
|
| def forward(self, cells, graph): |
| g, m = graph.g2m[:, 0], graph.g2m[:, 1] |
| feat = graph.g2m_feat.expand(cells.shape[0], -1, -1) |
| return scatter_mean(self.edge(torch.cat([cells[:, g], feat], -1)), m, graph.n_mesh) |
|
|
|
|
| class DynamicsRound(nn.Module): |
| """One tick: field-gated messages update a velocity; state integrates it.""" |
|
|
| def __init__(self, d, damp=0.9, dt=0.2): |
| super().__init__() |
| self.msg = mlp(2 * d + 9, d, d) |
| self.gate = mlp(2 * d + 9, d // 2, 1) |
| self.upd = mlp(2 * d, d, d) |
| self.damp, self.dt = damp, dt |
|
|
| def forward(self, h, v, edge_field, graph): |
| src, dst = graph.m2m[:, 0], graph.m2m[:, 1] |
| feat = graph.m2m_feat.expand(h.shape[0], -1, -1) |
| pair = torch.cat([h[:, dst], h[:, src], feat, edge_field], -1) |
| gate = torch.sigmoid(self.gate(pair)) |
| agg = scatter_mean(gate * self.msg(pair), dst, graph.n_mesh) |
| v = self.damp * v + self.upd(torch.cat([h, agg], -1)) |
| return h + self.dt * v, v, gate.squeeze(-1) |
|
|
|
|
| class MeshDynamics(nn.Module): |
| """Second-order rollout over the mesh; exposes mean gates for physics probes.""" |
|
|
| def __init__(self, d, rounds): |
| super().__init__() |
| self.rounds = nn.ModuleList(DynamicsRound(d) for _ in range(rounds)) |
|
|
| def forward(self, h, field_mesh, graph): |
| src, dst = graph.m2m[:, 0], graph.m2m[:, 1] |
| edge_field = torch.cat([field_mesh[:, src], field_mesh[:, dst]], -1) |
| v = torch.zeros_like(h) |
| gates = [] |
| for tick in self.rounds: |
| h, v, g = tick(h, v, edge_field, graph) |
| gates.append(g) |
| return h, torch.stack(gates).mean(0) |
|
|
|
|
| class MeshToGrid(nn.Module): |
| """Project mesh state back to cell latents along the same edges.""" |
|
|
| def __init__(self, d): |
| super().__init__() |
| self.edge = mlp(d + 3, d, d) |
|
|
| def forward(self, h, n_cells, graph): |
| g, m = graph.g2m[:, 0], graph.g2m[:, 1] |
| feat = graph.g2m_feat.expand(h.shape[0], -1, -1) |
| return scatter_mean(self.edge(torch.cat([h[:, m], feat], -1)), g, n_cells) |
|
|
|
|
| class TemperatureReadout(nn.Module): |
| """Intensity as a function of (cell latent, log-T). Query any temperature.""" |
|
|
| def __init__(self, d_in, d): |
| super().__init__() |
| self.net = mlp(d_in + 1, d, 1, depth=3) |
|
|
| def forward(self, cells, logt_q): |
| q = norm_logt(logt_q.to(cells.device))[None, :, None, None] |
| z = cells[:, None].expand(-1, len(logt_q), -1, -1) |
| return self.net(torch.cat([z, q.expand(*z.shape[:3], 1)], -1)).squeeze(-1) |
|
|
|
|
| |
|
|
| @dataclass |
| class Config: |
| dim: int = 128 |
| rounds: int = 6 |
| res: int = 256 |
| mesh_side: int = 48 |
|
|
|
|
| class Nika5(nn.Module): |
| def __init__(self, cfg: Config): |
| super().__init__() |
| self.cfg = cfg |
| d = cfg.dim |
| self.channels = ChannelEncoder(d) |
| self.field = FieldEncoder(d) |
| self.geometry = GeometryEmbed(d) |
| self.norm = nn.LayerNorm(d) |
| self.lift = GridToMesh(d) |
| self.dynamics = MeshDynamics(d, cfg.rounds) |
| self.project = MeshToGrid(d) |
| |
| |
| |
| self.readout = TemperatureReadout(2 * d + 9, d) |
|
|
| def forward(self, obs: Observation, graph: SolarGraph, logt_q): |
| cells = self.norm(self.channels(obs.aia, obs.present) |
| + self.field(obs.field) |
| + self.geometry(graph)) |
| state = self.lift(cells, graph) |
| field_mesh = scatter_mean(obs.field.transpose(1, 2)[:, graph.g2m[:, 0]], |
| graph.g2m[:, 1], graph.n_mesh) |
| state, self.gates = self.dynamics(state, field_mesh, graph) |
| context = self.project(state, cells.shape[1], graph) |
| raw = (obs.aia * obs.present[:, :, None]).transpose(1, 2) |
| return self.readout(torch.cat([context, cells, raw], -1), logt_q) |
|
|
