Add exposure autoencoder pretraining

backbones.py

@@ -474,3 +474,117 @@ class TimesNetExposureEncoder(nn.Module):
         if self.gate is not None:
             h = torch.sigmoid(self.gate) * h
         return h
+
+
+class TimesNetSequenceDecoder(nn.Module):
+    """Decode a fixed-size latent vector into a multivariate time series."""
+
+    def __init__(
+        self,
+        output_dim: int,
+        latent_dim: int,
+        d_model: int,
+        n_layers: int = 2,
+        top_k: int = 3,
+        n_convnext_blocks: int = 2,
+        conv_kernel_size: int = 7,
+        mlp_ratio: float = 4.0,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        self.latent_proj = nn.Linear(latent_dim, d_model)
+        self.position_proj = nn.Linear(3, d_model)
+        self.blocks = nn.ModuleList([
+            TimesNetBlock(
+                d_model=d_model,
+                top_k=top_k,
+                n_convnext_blocks=n_convnext_blocks,
+                conv_kernel_size=conv_kernel_size,
+                mlp_ratio=mlp_ratio,
+                dropout=dropout,
+            )
+            for _ in range(n_layers)
+        ])
+        self.final_ln = nn.LayerNorm(d_model)
+        self.output_proj = nn.Linear(d_model, output_dim)
+
+    def forward(self, latent: torch.Tensor, length: int) -> torch.Tensor:
+        if latent.dim() != 2:
+            raise ValueError(
+                f"latent must have shape (B, D), got {tuple(latent.shape)}"
+            )
+        position = torch.linspace(
+            0.0, 1.0, length, device=latent.device, dtype=latent.dtype
+        )
+        position = torch.stack(
+            [position, torch.sin(2 * torch.pi * position),
+             torch.cos(2 * torch.pi * position)],
+            dim=-1,
+        )
+        h = self.latent_proj(latent).unsqueeze(1)
+        h = h + self.position_proj(position).unsqueeze(0)
+        for block in self.blocks:
+            h = block(h)
+        return self.output_proj(self.final_ln(h))
+
+
+class TimesNetExposureAutoencoder(nn.Module):
+    """Dual-resolution exposure autoencoder with a reusable event encoder."""
+
+    def __init__(
+        self,
+        n_embd: int = 120,
+        daily_input_dim: int = 4,
+        monthly_input_dim: int = 2,
+        d_model: int | None = None,
+        n_layers: int = 2,
+        top_k: int = 3,
+        n_convnext_blocks: int = 2,
+        conv_kernel_size: int = 7,
+        mlp_ratio: float = 4.0,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        d_model = n_embd if d_model is None else d_model
+        encoder_kwargs = dict(
+            n_embd=n_embd,
+            daily_input_dim=daily_input_dim,
+            monthly_input_dim=monthly_input_dim,
+            d_model=d_model,
+            n_layers=n_layers,
+            top_k=top_k,
+            n_convnext_blocks=n_convnext_blocks,
+            conv_kernel_size=conv_kernel_size,
+            mlp_ratio=mlp_ratio,
+            dropout=dropout,
+            use_gate=True,
+        )
+        decoder_kwargs = dict(
+            latent_dim=n_embd,
+            d_model=d_model,
+            n_layers=n_layers,
+            top_k=top_k,
+            n_convnext_blocks=n_convnext_blocks,
+            conv_kernel_size=conv_kernel_size,
+            mlp_ratio=mlp_ratio,
+            dropout=dropout,
+        )
+        self.encoder = TimesNetExposureEncoder(**encoder_kwargs)
+        self.daily_decoder = TimesNetSequenceDecoder(
+            output_dim=daily_input_dim, **decoder_kwargs
+        )
+        self.monthly_decoder = TimesNetSequenceDecoder(
+            output_dim=monthly_input_dim, **decoder_kwargs
+        )
+
+    def forward(
+        self,
+        daily: torch.Tensor,
+        monthly: torch.Tensor,
+        daily_mask: torch.Tensor | None = None,
+        monthly_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        latent = self.encoder(daily, monthly, daily_mask, monthly_mask)
+        daily_reconstruction = self.daily_decoder(latent, daily.size(1))
+        monthly_reconstruction = self.monthly_decoder(latent, monthly.size(1))
+        return daily_reconstruction, monthly_reconstruction, latent