Add TimesNet exposure encoder

backbones.py

@@ -326,3 +326,328 @@ class AgeSinusoidalEncoding(nn.Module):
         output[:, :, 1::2] = torch.sin(args)
         output = self.linear(output)
         return output
+
+
+class GRN2d(nn.Module):
+    """Global Response Normalization from ConvNeXt V2 for NCHW tensors.
+
+    Reference:
+        Woo et al., "ConvNeXt V2: Co-designing and Scaling ConvNets with
+        Masked Autoencoders", CVPR 2023. https://arxiv.org/abs/2301.00808
+    """
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, dim, 1, 1))
+        self.beta = nn.Parameter(torch.zeros(1, dim, 1, 1))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gx = torch.norm(x, p=2, dim=(2, 3), keepdim=True)
+        nx = gx / (gx.mean(dim=1, keepdim=True) + self.eps)
+        return x + self.gamma * (x * nx) + self.beta
+
+
+class LayerNorm2d(nn.Module):
+    """Channel-wise LayerNorm for NCHW tensors."""
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.bias = nn.Parameter(torch.zeros(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.permute(0, 2, 3, 1)
+        x = F.layer_norm(x, (self.weight.numel(),), self.weight, self.bias, self.eps)
+        return x.permute(0, 3, 1, 2)
+
+
+class ConvNeXtV2Block2d(nn.Module):
+    """Lightweight ConvNeXt V2-style 2D block for TimesNet period images.
+
+    This is intentionally a block, not the full ConvNeXt V2 image backbone:
+    TimesNet's 2D tensors are reshaped time-series period maps rather than
+    natural images, so aggressive visual downsampling would destroy axis
+    semantics.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        kernel_size: int = 7,
+        mlp_ratio: float = 4.0,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        padding = kernel_size // 2
+        hidden_dim = int(dim * mlp_ratio)
+        self.dwconv = nn.Conv2d(
+            dim,
+            dim,
+            kernel_size=kernel_size,
+            padding=padding,
+            groups=dim,
+        )
+        self.norm = LayerNorm2d(dim)
+        self.pwconv1 = nn.Conv2d(dim, hidden_dim, kernel_size=1)
+        self.act = nn.GELU()
+        self.grn = GRN2d(hidden_dim)
+        self.pwconv2 = nn.Conv2d(hidden_dim, dim, kernel_size=1)
+        self.drop = nn.Dropout(dropout)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.normal_(self.dwconv.weight, mean=0.0, std=0.02)
+        nn.init.zeros_(self.dwconv.bias)
+        nn.init.normal_(self.pwconv1.weight, mean=0.0, std=0.02)
+        nn.init.zeros_(self.pwconv1.bias)
+        nn.init.normal_(self.pwconv2.weight, mean=0.0, std=0.02)
+        nn.init.zeros_(self.pwconv2.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = self.dwconv(x)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        x = self.drop(x)
+        return residual + x
+
+
+class TimesNetBlock(nn.Module):
+    """TimesNet block with ConvNeXt V2-style 2D extraction.
+
+    The block follows TimesNet's idea: discover dominant periods with FFT,
+    reshape a 1D sequence into period-wise 2D maps, run a 2D convolutional
+    extractor, then fuse the top-k period branches.
+
+    Reference:
+        Wu et al., "TimesNet: Temporal 2D-Variation Modeling for General Time
+        Series Analysis", ICLR 2023. https://arxiv.org/abs/2210.02186
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        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__()
+        if top_k <= 0:
+            raise ValueError(f"top_k must be > 0, got {top_k}")
+        self.top_k = top_k
+        self.norm = nn.LayerNorm(d_model)
+        self.extractor = nn.Sequential(*[
+            ConvNeXtV2Block2d(
+                dim=d_model,
+                kernel_size=conv_kernel_size,
+                mlp_ratio=mlp_ratio,
+                dropout=dropout,
+            )
+            for _ in range(n_convnext_blocks)
+        ])
+
+    def _select_periods(self, x: torch.Tensor) -> tuple[list[int], torch.Tensor]:
+        B, T, C = x.shape
+        spectrum = torch.fft.rfft(x.float(), dim=1)
+        amplitude = spectrum.abs().mean(dim=(0, 2))
+        if amplitude.numel() <= 1:
+            return [max(T, 1)], x.new_ones(1)
+
+        amplitude = amplitude.clone()
+        amplitude[0] = 0.0
+        k = min(self.top_k, amplitude.numel() - 1)
+        weights, indices = torch.topk(amplitude, k=k)
+        periods = [max(1, T // int(idx.item())) for idx in indices]
+        return periods, weights.to(dtype=x.dtype, device=x.device)
+
+    def _period_branch(self, x: torch.Tensor, period: int) -> torch.Tensor:
+        B, T, C = x.shape
+        if T % period != 0:
+            padded_len = ((T // period) + 1) * period
+            pad = x.new_zeros(B, padded_len - T, C)
+            x_pad = torch.cat([x, pad], dim=1)
+        else:
+            padded_len = T
+            x_pad = x
+
+        n_periods = padded_len // period
+        x_2d = x_pad.reshape(B, n_periods, period, C).permute(0, 3, 1, 2)
+        y = self.extractor(x_2d)
+        y = y.permute(0, 2, 3, 1).reshape(B, padded_len, C)
+        return y[:, :T, :]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x_norm = self.norm(x)
+        periods, weights = self._select_periods(x_norm)
+        branches = torch.stack(
+            [self._period_branch(x_norm, period) for period in periods],
+            dim=-1,
+        )
+        weights = torch.softmax(weights, dim=0).view(1, 1, 1, -1)
+        return residual + (branches * weights).sum(dim=-1)
+
+
+class TimesNetEncoder(nn.Module):
+    """Encode a multivariate time series into one fixed-size embedding."""
+
+    def __init__(
+        self,
+        input_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,
+        append_observed_mask: bool = True,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.append_observed_mask = append_observed_mask
+        in_dim = input_dim * 2 if append_observed_mask else input_dim
+        self.input_proj = nn.Linear(in_dim, 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.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.normal_(self.input_proj.weight, mean=0.0, std=0.02)
+        nn.init.zeros_(self.input_proj.bias)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        observed_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if x.dim() != 3:
+            raise ValueError(f"x must have shape (B, T, C), got {tuple(x.shape)}")
+        if x.size(-1) != self.input_dim:
+            raise ValueError(
+                f"last dim must be input_dim={self.input_dim}, got {x.size(-1)}"
+            )
+
+        finite_mask = torch.isfinite(x)
+        x = torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)
+        if observed_mask is None:
+            observed_mask = finite_mask
+        elif observed_mask.shape == x.shape[:2]:
+            observed_mask = observed_mask.unsqueeze(-1).expand_as(x)
+        elif observed_mask.shape != x.shape:
+            raise ValueError(
+                "observed_mask must have shape (B, T) or (B, T, C), got "
+                f"{tuple(observed_mask.shape)}"
+            )
+
+        observed_mask = observed_mask.to(device=x.device, dtype=x.dtype)
+        x = x * observed_mask
+        if self.append_observed_mask:
+            x = torch.cat([x, observed_mask], dim=-1)
+
+        h = self.input_proj(x)
+        for block in self.blocks:
+            h = block(h)
+        h = self.final_ln(h)
+
+        valid_time = observed_mask.amax(dim=-1)
+        pooled = (h * valid_time.unsqueeze(-1)).sum(dim=1)
+        denom = valid_time.sum(dim=1, keepdim=True).clamp_min(1.0)
+        return pooled / denom
+
+
+class TimesNetExposureEncoder(nn.Module):
+    """Encode pre-onset environmental exposure into an event-level embedding.
+
+    Expected inputs:
+        daily:   (B, 1826, 4) for tmean, tmax, tmin, rhmean
+        monthly: (B, 241, 2) for tmean, rhmean
+
+    Output:
+        (B, n_embd), suitable for adding to a disease event embedding in the
+        route-2 single-stream event-enhancement setup.
+    """
+
+    def __init__(
+        self,
+        n_embd: int,
+        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,
+        use_gate: bool = True,
+    ):
+        super().__init__()
+        d_model = n_embd if d_model is None else d_model
+        self.daily_encoder = TimesNetEncoder(
+            input_dim=daily_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,
+            append_observed_mask=True,
+        )
+        self.monthly_encoder = TimesNetEncoder(
+            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,
+            append_observed_mask=True,
+        )
+        self.out_proj = nn.Sequential(
+            nn.LayerNorm(2 * d_model),
+            nn.Linear(2 * d_model, n_embd),
+            nn.GELU(),
+            nn.Linear(n_embd, n_embd),
+        )
+        self.gate = nn.Parameter(torch.tensor(-2.0)) if use_gate else None
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        for module in self.out_proj:
+            if isinstance(module, nn.Linear):
+                nn.init.normal_(module.weight, mean=0.0, std=0.02)
+                nn.init.zeros_(module.bias)
+
+    def forward(
+        self,
+        daily: torch.Tensor,
+        monthly: torch.Tensor,
+        daily_mask: torch.Tensor | None = None,
+        monthly_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        h_daily = self.daily_encoder(daily, observed_mask=daily_mask)
+        h_monthly = self.monthly_encoder(monthly, observed_mask=monthly_mask)
+        h = self.out_proj(torch.cat([h_daily, h_monthly], dim=-1))
+        if self.gate is not None:
+            h = torch.sigmoid(self.gate) * h
+        return h

models.py

@@ -8,6 +8,7 @@ from backbones import (
     AgeSinusoidalEncoding,
     GPTBlock,
     GaussianRBFTimeBasis,
+    TimesNetExposureEncoder,
     TimeRoPE,
     TokenAutoDiscretization,
 )
@@ -160,6 +161,16 @@ class DeepHealth(nn.Module):
         dist_mode: str = "exponential",  # "exponential", "weibull" or "mixed"
         extra_pool_reduce: str = "mean",
         dropout: float = 0.0,
+        use_exposure_encoder: bool = False,
+        exposure_daily_input_dim: int = 4,
+        exposure_monthly_input_dim: int = 2,
+        exposure_d_model: int | None = None,
+        exposure_n_layers: int = 2,
+        exposure_top_k: int = 3,
+        exposure_n_convnext_blocks: int = 2,
+        exposure_conv_kernel_size: int = 7,
+        exposure_mlp_ratio: float = 4.0,
+        exposure_use_gate: bool = True,
     ):
         super().__init__()
         if target_mode not in ["next_token", "all_future"]:
@@ -189,8 +200,26 @@ class DeepHealth(nn.Module):
         self.time_mode = time_mode
         self.dist_mode = dist_mode
         self.extra_pool_reduce = extra_pool_reduce
+        self.use_exposure_encoder = use_exposure_encoder
         self.n_embd = n_embd
         self.vocab_size = vocab_size
+        self.exposure_encoder = (
+            TimesNetExposureEncoder(
+                n_embd=n_embd,
+                daily_input_dim=exposure_daily_input_dim,
+                monthly_input_dim=exposure_monthly_input_dim,
+                d_model=exposure_d_model,
+                n_layers=exposure_n_layers,
+                top_k=exposure_top_k,
+                n_convnext_blocks=exposure_n_convnext_blocks,
+                conv_kernel_size=exposure_conv_kernel_size,
+                mlp_ratio=exposure_mlp_ratio,
+                dropout=dropout,
+                use_gate=exposure_use_gate,
+            )
+            if use_exposure_encoder
+            else None
+        )
         nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
         nn.init.zeros_(self.token_embedding.weight[0])
         nn.init.normal_(self.gender_embedding.weight, mean=0.0, std=0.02)
@@ -317,6 +346,94 @@ class DeepHealth(nn.Module):
         pooled_mask = arange_groups.unsqueeze(0) < group_count.unsqueeze(1)
         return pooled_h, pooled_time, pooled_mask
 
+    def _encode_event_exposure(
+        self,
+        exposure_daily: torch.Tensor | None,
+        exposure_monthly: torch.Tensor | None,
+        exposure_daily_mask: torch.Tensor | None,
+        exposure_monthly_mask: torch.Tensor | None,
+        event_shape: tuple[int, int],
+    ) -> torch.Tensor | None:
+        if self.exposure_encoder is None:
+            if exposure_daily is not None or exposure_monthly is not None:
+                raise ValueError(
+                    "Exposure tensors were provided but use_exposure_encoder=False"
+                )
+            return None
+        if exposure_daily is None or exposure_monthly is None:
+            raise ValueError(
+                "exposure_daily and exposure_monthly are required when "
+                "use_exposure_encoder=True"
+            )
+
+        batch_size, event_len = event_shape
+        if exposure_daily.shape[:2] != event_shape:
+            raise ValueError(
+                "exposure_daily must have shape (B, L, T_daily, C_daily), got "
+                f"{tuple(exposure_daily.shape)} for event shape {event_shape}"
+            )
+        if exposure_monthly.shape[:2] != event_shape:
+            raise ValueError(
+                "exposure_monthly must have shape (B, L, T_monthly, C_monthly), got "
+                f"{tuple(exposure_monthly.shape)} for event shape {event_shape}"
+            )
+        if exposure_daily.dim() != 4 or exposure_monthly.dim() != 4:
+            raise ValueError(
+                "exposure_daily and exposure_monthly must both be 4D tensors"
+            )
+
+        def flatten_mask(mask: torch.Tensor | None, ref: torch.Tensor) -> torch.Tensor | None:
+            if mask is None:
+                return None
+            if mask.shape[:2] != event_shape:
+                raise ValueError(
+                    "exposure mask must start with shape (B, L), got "
+                    f"{tuple(mask.shape)}"
+                )
+            if mask.dim() not in {3, 4}:
+                raise ValueError(
+                    "exposure mask must have shape (B, L, T) or (B, L, T, C), "
+                    f"got {tuple(mask.shape)}"
+                )
+            if mask.shape[2] != ref.shape[2]:
+                raise ValueError(
+                    "exposure mask time dimension does not match exposure tensor"
+                )
+            if mask.dim() == 4 and mask.shape[3] != ref.shape[3]:
+                raise ValueError(
+                    "exposure mask channel dimension does not match exposure tensor"
+                )
+            return mask.reshape(batch_size * event_len, *mask.shape[2:])
+
+        daily = exposure_daily.reshape(
+            batch_size * event_len,
+            exposure_daily.size(2),
+            exposure_daily.size(3),
+        )
+        monthly = exposure_monthly.reshape(
+            batch_size * event_len,
+            exposure_monthly.size(2),
+            exposure_monthly.size(3),
+        )
+        daily_mask = flatten_mask(exposure_daily_mask, exposure_daily)
+        monthly_mask = flatten_mask(exposure_monthly_mask, exposure_monthly)
+
+        param = next(self.exposure_encoder.parameters())
+        daily = daily.to(device=param.device, dtype=param.dtype)
+        monthly = monthly.to(device=param.device, dtype=param.dtype)
+        if daily_mask is not None:
+            daily_mask = daily_mask.to(device=param.device)
+        if monthly_mask is not None:
+            monthly_mask = monthly_mask.to(device=param.device)
+
+        exposure_emb = self.exposure_encoder(
+            daily=daily,
+            monthly=monthly,
+            daily_mask=daily_mask,
+            monthly_mask=monthly_mask,
+        )
+        return exposure_emb.reshape(batch_size, event_len, self.n_embd)
+
     def _forward_shared(
         self,
         event_seq: torch.LongTensor,
@@ -329,6 +446,10 @@ class DeepHealth(nn.Module):
         other_value: torch.Tensor | None = None,
         other_value_kind: torch.LongTensor | None = None,
         other_time: torch.FloatTensor | None = None,
+        exposure_daily: torch.Tensor | None = None,
+        exposure_monthly: torch.Tensor | None = None,
+        exposure_daily_mask: torch.Tensor | None = None,
+        exposure_monthly_mask: torch.Tensor | None = None,
         return_output: bool = False,
         **unused_kwargs,
     ) -> torch.Tensor | DeepHealthOutput:
@@ -357,6 +478,16 @@ class DeepHealth(nn.Module):
 
         event_len = event_seq.size(1)
         h_disease = self.token_embedding(event_seq)
+        h_exposure = self._encode_event_exposure(
+            exposure_daily=exposure_daily,
+            exposure_monthly=exposure_monthly,
+            exposure_daily_mask=exposure_daily_mask,
+            exposure_monthly_mask=exposure_monthly_mask,
+            event_shape=(event_seq.size(0), event_len),
+        )
+        if h_exposure is not None:
+            h_exposure = h_exposure.to(device=event_seq.device, dtype=h_disease.dtype)
+            h_disease = h_disease + h_exposure
         t_disease = time_seq
 
         if other_time.shape != other_type.shape: