1"""Patch tokenisation with gated feature projection, resolution and temporal-gap embeddings."""2from __future__ import annotations34import math56import torch7import torch.nn as nn8import torch.nn.functional as F91011class GatedProjection(nn.Module):12 """x -> W2(GELU(W1 x) * sigmoid(Wg x)) : a gated linear unit over the flattened patch features."""1314 def __init__(self, d_in: int, d_model: int, dropout: float = 0.0):15 super().__init__()16 self.lin = nn.Linear(d_in, 2 * d_model)17 self.out = nn.Linear(d_model, d_model)18 self.drop = nn.Dropout(dropout)1920 def forward(self, x: torch.Tensor) -> torch.Tensor:21 a, g = self.lin(x).chunk(2, dim=-1)22 return self.out(self.drop(F.gelu(a) * torch.sigmoid(g)))232425def sinusoidal(positions: torch.Tensor, d: int) -> torch.Tensor:26 """positions: [..., ] float in arbitrary units -> [..., d] sinusoidal features."""27 half = d // 228 freq = torch.exp(-math.log(10000.0) * torch.arange(half, device=positions.device, dtype=torch.float32) / half)29 ang = positions.unsqueeze(-1).float() * freq30 return torch.cat([torch.sin(ang), torch.cos(ang)], dim=-1)313233class PatchEmbed(nn.Module):34 """Groups `patch` consecutive bars (each with f features) into one token: [B, L, f] -> [B, L/patch, d_model].35 Adds: learned resolution embedding, sinusoidal position within the window (in patch units), temporal-gap36 embedding (max log gap inside the patch), and a 'padded / missing fraction' scalar."""3738 def __init__(self, f_bar: int, patch: int, d_model: int, n_resolutions: int = 3, dropout: float = 0.1):39 super().__init__()40 self.patch, self.f_bar, self.d = patch, f_bar, d_model41 self.proj = GatedProjection(patch * f_bar + 2, d_model, dropout)42 self.res_emb = nn.Embedding(n_resolutions, d_model)43 self.gap_proj = nn.Linear(d_model, d_model)44 self.pos_proj = nn.Linear(d_model, d_model)45 self.norm = nn.LayerNorm(d_model)4647 def forward(self, x: torch.Tensor, mask: torch.Tensor, res_id: int, gap_feat_index: int = 10,48 missing_feat_index: int = 8) -> tuple[torch.Tensor, torch.Tensor]:49 B, L, f = x.shape50 P = L // self.patch51 xp = x[:, :P * self.patch].reshape(B, P, self.patch * f)52 mp = mask[:, :P * self.patch].reshape(B, P, self.patch).float()53 frac_present = mp.mean(-1, keepdim=True)54 gap = x[:, :P * self.patch, gap_feat_index].reshape(B, P, self.patch).max(-1, keepdim=True).values55 tok = self.proj(torch.cat([xp, frac_present, gap], dim=-1))56 pos = torch.arange(P, device=x.device, dtype=torch.float32).flip(0) # distance to origin in patches57 tok = tok + self.pos_proj(sinusoidal(pos, self.d)).unsqueeze(0)58 tok = tok + self.gap_proj(sinusoidal(gap.squeeze(-1) * 8.0, self.d))59 tok = tok + self.res_emb.weight[res_id]60 token_valid = (frac_present.squeeze(-1) > 0) # a patch with no present bar is still a token (masked by flag)61 return self.norm(tok), token_valid62