Add TimesNet exposure encoder

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: