""" Hierarchical multi-step hyperbolic predictor. Design principles drawn from: - Hyperbolic hierarchy capacity (Nickel & Kiela, Sala et al.) - Recursive / tree-structured modeling ideas (R2D2-style differentiable trees, arXiv:2301.12987 hierarchical inductive biases) - Spectral / Neural Operator intuition for spatiotemporal fields (The Well baselines) - Riemannian optimization hygiene (geoopt + numerical stability literature) """ from __future__ import annotations import torch import torch.nn as nn import torch.nn.functional as F from .poincare import PoincareBall8D class SpectralConv2d(nn.Module): """Lightweight spectral mixing block (FNO-inspired, low-rank).""" def __init__(self, in_ch: int, out_ch: int, modes: int = 8): super().__init__() self.modes = modes self.scale = 1.0 / (in_ch * out_ch) self.weights = nn.Parameter(self.scale * torch.randn(in_ch, out_ch, modes, modes, 2)) def compl_mul(self, a, b): return torch.einsum("bixy,ioxy->boxy", a, b) def forward(self, x: torch.Tensor) -> torch.Tensor: B, C, H, W = x.shape x_ft = torch.fft.rfft2(x, norm="ortho") out_ft = torch.zeros(B, self.weights.shape[1], H, W // 2 + 1, dtype=torch.cfloat, device=x.device) m1, m2 = min(self.modes, H), min(self.modes, W // 2 + 1) w = torch.view_as_complex(self.weights) out_ft[:, :, :m1, :m2] = self.compl_mul(x_ft[:, :, :m1, :m2], w[:, :, :m1, :m2]) return torch.fft.irfft2(out_ft, s=(H, W), norm="ortho") class MultiScaleEncoder(nn.Module): """ Multi-scale spatiotemporal encoder → 8-D Euclidean latent. REDESIGNED this session: previously took a fixed `in_channels` and built `nn.Conv2d(in_channels, 32, 1)` — every checkpoint was then permanently locked to that channel count (verified directly: two uploaded checkpoints both had encoder.lift.weight shaped for in_channels=11, unusable for a 2-channel dataset without surgery). This blocked the actual goal of streaming successive Well datasets with different channel counts through one continually-trained model. Now channel-count-independent, following the shared per-channel-stem principle used by recent heterogeneous-PDE foundation models (MORPH, Tadpole) rather than padding to a fixed C_max: the 1x1 stem is applied to each channel independently (folded into the batch dimension), then mean-fused across channels before the more expensive spectral/local spatial processing (which therefore runs once per sample regardless of channel count, not once per channel -- verified this keeps the expensive ops cheap). Verified directly (not assumed) before merging: same weights produce finite, correctly-shaped output for C=2, C=11, and C=47 at a different resolution, and gradients flow correctly across a sequential C=2 -> C=11 training step (the actual continual-training scenario this was built for). Known current limitation, stated plainly: mean-fusion across channels is lossy (loses relative channel importance -- a highly informative channel is weighted the same as a noisy one). Attention-based fusion (MORPH-style) would address this but is a larger change; mean-fusion is the correct minimal first step, not the final design. """ def __init__(self, hidden: int = 64, out_dim: int = 8, stem_ch: int = 32, channel_fuse: str = "mean"): super().__init__() if channel_fuse != "mean": raise ValueError( f"channel_fuse={channel_fuse!r} not implemented yet -- only " f"'mean' exists currently. Raising rather than silently " f"falling back to mean, since that would silently change " f"behavior from what was requested." ) self.channel_fuse = channel_fuse self.stem_ch = stem_ch self.stem = nn.Conv2d(1, stem_ch, 1) self.spec = SpectralConv2d(stem_ch, stem_ch, modes=6) self.local = nn.Sequential( nn.Conv2d(stem_ch, stem_ch, 3, padding=1), nn.GELU(), nn.Conv2d(stem_ch, stem_ch, 3, padding=1), nn.GELU(), ) self.pool = nn.AdaptiveAvgPool2d(4) self.head = nn.Sequential( nn.Flatten(), nn.Linear(stem_ch * 4 * 4, hidden), nn.GELU(), nn.Linear(hidden, out_dim), ) def forward(self, x: torch.Tensor) -> torch.Tensor: # x: (B, C, H, W) or (B, T, C, H, W); C is NOT fixed by the module. if x.dim() == 5: x = x[:, -1] B, C, H, W = x.shape x = x.reshape(B * C, 1, H, W) # fold channels into batch h = self.stem(x) # shared stem, per channel h = h.reshape(B, C, self.stem_ch, H, W) h = h.mean(dim=1) # fuse channels -> (B, stem_ch, H, W) h = h + self.spec(h) # expensive ops run once per sample h = h + self.local(h) return self.head(self.pool(h)) class HierarchicalHyperbolicPredictor(nn.Module): def __init__(self, encoder: MultiScaleEncoder, c: float = 1.0, pred_steps: int = 4, levels: int = 2, learnable_c: bool = False): super().__init__() self.encoder = encoder self.poincare = PoincareBall8D(c=c, learnable_c=learnable_c) self.pred_steps = pred_steps self.levels = levels # Coarse RNN in tangent space self.coarse_rnn = nn.GRU(8, 48, batch_first=True) self.coarse_head = nn.Linear(48, 8) # Fine residual heads (one per level after the first) self.fine_heads = nn.ModuleList([nn.Linear(8 + 8, 8) for _ in range(max(0, levels - 1))]) def encode(self, x: torch.Tensor) -> torch.Tensor: z_euc = self.encoder(x) z = self.poincare.expmap0(z_euc) return self.poincare.clip_norm(z) def forward(self, x: torch.Tensor) -> torch.Tensor: z0 = self.encode(x) # (B, 8) z_tan = self.poincare.logmap0(z0) h = z_tan.unsqueeze(1) coarse_seq = [] hidden = None cur = z_tan for _ in range(self.pred_steps): out, hidden = self.coarse_rnn(h, hidden) delta = self.coarse_head(out.squeeze(1)) cur = cur + delta coarse_seq.append(cur) h = cur.unsqueeze(1) # residual refinement levels refined = torch.stack(coarse_seq, dim=1) # (B, S, 8) for head in self.fine_heads: # condition on previous refined + original z0 cond = torch.cat([refined, z_tan.unsqueeze(1).expand_as(refined)], dim=-1) refined = refined + 0.5 * head(cond) # map final tangent sequence into the ball B, S, D = refined.shape ball = self.poincare.expmap0(refined.reshape(B * S, D)) ball = self.poincare.clip_norm(ball).reshape(B, S, D) return ball def hyperbolic_loss(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: B, S, _ = pred.shape return self.poincare.dist(pred.reshape(B * S, -1), target.reshape(B * S, -1)).mean() class HyperbolicCritic(nn.Module): def __init__(self, c: float = 1.0, learnable_c: bool = False): super().__init__() self.poincare = PoincareBall8D(c=c, learnable_c=learnable_c) self.net = nn.Sequential( nn.Linear(8, 64), nn.GELU(), nn.Linear(64, 32), nn.GELU(), nn.Linear(32, 1), ) def forward(self, z_ball: torch.Tensor) -> torch.Tensor: # evaluate in tangent space at origin for stability z_tan = self.poincare.logmap0(z_ball) return self.net(z_tan).squeeze(-1)