Upload code/model.py

code/model.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import nn
+
+
+def fourier_loglam(loglam: torch.Tensor, n_freq: int = 32) -> torch.Tensor:
+    # Center roughly around optical wavelengths to improve conditioning.
+    x = loglam - math.log(6000.0)
+    freqs = torch.logspace(0.0, math.log10(128.0), n_freq, device=loglam.device, dtype=loglam.dtype)
+    angles = x.unsqueeze(-1) * freqs
+    return torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1)
+
+
+class NativeSpecZMAE(nn.Module):
+    def __init__(
+        self,
+        feature_dim: int = 6,
+        target_aux_dim: int = 4,
+        d_model: int = 256,
+        nhead: int = 8,
+        encoder_layers: int = 4,
+        decoder_layers: int = 3,
+        dim_feedforward: int | None = None,
+        dropout: float = 0.1,
+        fourier_freqs: int = 32,
+        z_bins: int = 64,
+        y_min: float = 0.0,
+        y_max: float = math.log1p(6.0),
+    ):
+        super().__init__()
+        self.fourier_freqs = fourier_freqs
+        self.y_min = y_min
+        self.y_max = y_max
+        self.z_bins = z_bins
+        pos_dim = fourier_freqs * 2
+        ff = dim_feedforward or d_model * 4
+
+        self.input_proj = nn.Sequential(
+            nn.Linear(feature_dim + pos_dim, d_model),
+            nn.LayerNorm(d_model),
+            nn.GELU(),
+            nn.Linear(d_model, d_model),
+        )
+        self.target_proj = nn.Sequential(
+            nn.Linear(target_aux_dim + pos_dim, d_model),
+            nn.LayerNorm(d_model),
+            nn.GELU(),
+            nn.Linear(d_model, d_model),
+        )
+        self.cls = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
+        self.z_query = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
+
+        enc_layer = nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=ff,
+            dropout=dropout,
+            batch_first=True,
+            norm_first=True,
+            activation="gelu",
+        )
+        self.encoder = nn.TransformerEncoder(enc_layer, num_layers=encoder_layers)
+        dec_layer = nn.TransformerDecoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=ff,
+            dropout=dropout,
+            batch_first=True,
+            norm_first=True,
+            activation="gelu",
+        )
+        self.decoder = nn.TransformerDecoder(dec_layer, num_layers=decoder_layers)
+        self.rec_head = nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, d_model // 2), nn.GELU(), nn.Linear(d_model // 2, 1))
+        self.z_head = nn.Sequential(nn.LayerNorm(d_model * 2), nn.Linear(d_model * 2, d_model), nn.GELU(), nn.Linear(d_model, 2))
+        self.z_bin_head = nn.Sequential(nn.LayerNorm(d_model * 2), nn.Linear(d_model * 2, d_model), nn.GELU(), nn.Linear(d_model, z_bins))
+
+    def forward(
+        self,
+        enc_features: torch.Tensor,
+        enc_loglam: torch.Tensor,
+        enc_padding: torch.Tensor,
+        target_loglam: torch.Tensor,
+        target_aux: torch.Tensor,
+    ) -> dict[str, torch.Tensor]:
+        bsz = enc_features.shape[0]
+        enc_pos = fourier_loglam(enc_loglam, self.fourier_freqs)
+        tokens = self.input_proj(torch.cat([enc_features, enc_pos], dim=-1))
+        cls = self.cls.expand(bsz, -1, -1)
+        z_query = self.z_query.expand(bsz, -1, -1)
+        memory_in = torch.cat([cls, z_query, tokens], dim=1)
+        special_padding = torch.zeros((bsz, 2), dtype=torch.bool, device=enc_padding.device)
+        memory_padding = torch.cat([special_padding, enc_padding], dim=1)
+        memory = self.encoder(memory_in, src_key_padding_mask=memory_padding)
+
+        cls_out = memory[:, 0]
+        z_out = memory[:, 1]
+        z_feat = torch.cat([z_out, cls_out], dim=-1)
+        z_params = self.z_head(z_feat)
+        z_bins = self.z_bin_head(z_feat)
+
+        tgt_pos = fourier_loglam(target_loglam, self.fourier_freqs)
+        tgt = self.target_proj(torch.cat([target_aux, tgt_pos], dim=-1))
+        decoded = self.decoder(tgt=tgt, memory=memory, memory_key_padding_mask=memory_padding)
+        rec = self.rec_head(decoded).squeeze(-1)
+        return {
+            "rec": rec,
+            "y_mu": z_params[:, 0],
+            "y_logvar": torch.clamp(z_params[:, 1], -8.0, 4.0),
+            "z_bin_logits": z_bins,
+        }
+
+    def y_to_bin(self, y: torch.Tensor) -> torch.Tensor:
+        scaled = (y - self.y_min) / max(self.y_max - self.y_min, 1e-6)
+        return torch.clamp((scaled * self.z_bins).long(), 0, self.z_bins - 1)