Feat/mlp depth loss weight bootstrap

corebehrt/functional/causal/cohort_sampler.py

@@ -14,34 +14,28 @@ def sample_cohort(
     sample_fraction: float = None,
     sample_size: int = None,
     seed: int = 42,
+    method: str = "resample",
 ) -> torch.Tensor:
     """
-    Randomly sample patient IDs by fraction or absolute size without replacement.
+    Randomly sample patient IDs by fraction or absolute size.
 
     Args:
         full_pids: Tensor of all patient IDs in the full cohort
         sample_fraction: Fraction of patients to sample (0 < fraction <= 1).
                         Ignored if sample_size is provided.
         sample_size: Absolute number of patients to sample. Takes precedence over sample_fraction.
         seed: Random seed for reproducibility
+        method: Sampling method - "resample" (without replacement) or "bootstrap" (with replacement)
 
     Returns:
         Tensor of sampled patient IDs
 
     Raises:
         ValueError: If neither or both parameters are provided, or values are invalid
-
-    Examples:
-        >>> full_pids = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
-        >>> # Sample by fraction
-        >>> sampled = sample_cohort(full_pids, sample_fraction=0.5, seed=42)
-        >>> len(sampled) == 5
-        True
-        >>> # Sample by absolute size
-        >>> sampled = sample_cohort(full_pids, sample_size=3, seed=42)
-        >>> len(sampled) == 3
-        True
     """
+    if method not in ("resample", "bootstrap"):
+        raise ValueError(f"method must be 'resample' or 'bootstrap', got '{method}'")
+
     # Validate inputs
     if sample_size is None and sample_fraction is None:
         raise ValueError("Either 'sample_size' or 'sample_fraction' must be provided")
@@ -54,12 +48,13 @@ def sample_cohort(
 
     # Calculate number of patients to sample
     n_total = len(full_pids)
+    replace = method == "bootstrap"
 
     if sample_size is not None:
         # Use absolute size
         if sample_size <= 0:
             raise ValueError(f"sample_size must be positive, got {sample_size}")
-        if sample_size > n_total:
+        if not replace and sample_size > n_total:
             raise ValueError(
                 f"sample_size {sample_size} exceeds available patients {n_total}"
             )
@@ -76,8 +71,8 @@ def sample_cohort(
                 f"sample_fraction {sample_fraction} results in 0 samples from {n_total} patients"
             )
 
-    # Sample indices without replacement
-    sampled_indices = rng.choice(n_total, size=n_sample, replace=False)
+    # Sample indices
+    sampled_indices = rng.choice(n_total, size=n_sample, replace=replace)
 
     # Return sampled PIDs
     return full_pids[sampled_indices]

corebehrt/main_causal/simulate_with_sampling.py

@@ -57,17 +57,22 @@ def main_simulate(config_path):
     # Sample subset of PIDs (either by size or fraction)
     sample_size = sampling_cfg.get("size", None)
     sample_fraction = sampling_cfg.get("fraction", None)
+    method = sampling_cfg.get("method", "resample")
     seed = cfg.get("seed", 42)
 
     if sample_size is not None:
-        logger.info(f"Sampling {sample_size} patients with seed {seed}")
-        sampled_pids = sample_cohort(full_pids, sample_size=sample_size, seed=seed)
+        logger.info(
+            f"Sampling {sample_size} patients with seed {seed} (method={method})"
+        )
+        sampled_pids = sample_cohort(
+            full_pids, sample_size=sample_size, seed=seed, method=method
+        )
     elif sample_fraction is not None:
         logger.info(
-            f"Sampling {sample_fraction * 100:.1f}% of patients with seed {seed}"
+            f"Sampling {sample_fraction * 100:.1f}% of patients with seed {seed} (method={method})"
         )
         sampled_pids = sample_cohort(
-            full_pids, sample_fraction=sample_fraction, seed=seed
+            full_pids, sample_fraction=sample_fraction, seed=seed, method=method
         )
     else:
         raise ValueError(

corebehrt/modules/model/causal/heads.py

@@ -73,17 +73,22 @@ class MLPHead(nn.Module):
     """
 
     def __init__(
-        self, input_size: int, hidden_size_ratio: int = 2, dropout_prob: float = 0.1
+        self,
+        input_size: int,
+        hidden_size_ratio: int = 2,
+        dropout_prob: float = 0.1,
+        num_hidden_layers: int = 1,
     ):
         super().__init__()
-        intermediate_size = input_size // hidden_size_ratio
-        self.classifier = nn.Sequential(
-            nn.LayerNorm(input_size),
-            nn.Linear(input_size, intermediate_size),
-            nn.GELU(),
-            # nn.Dropout(dropout_prob),
-            nn.Linear(intermediate_size, 1, bias=True),
-        )
+        layers = [nn.LayerNorm(input_size)]
+        current_size = input_size
+        for _ in range(num_hidden_layers):
+            next_size = current_size // hidden_size_ratio
+            layers.append(nn.Linear(current_size, next_size))
+            layers.append(nn.GELU())
+            current_size = next_size
+        layers.append(nn.Linear(current_size, 1, bias=True))
+        self.classifier = nn.Sequential(*layers)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """

corebehrt/modules/model/causal/model.py

@@ -48,6 +48,7 @@ def __init__(self, config):
         self.shared_representation = self.head_config.get("shared_representation", True)
         self.bidirectional = self.head_config.get("bidirectional", True)
         self.bottleneck_dim = self.head_config.get("bottleneck_dim", 128)
+        self.exposure_loss_weight = self.head_config.get("exposure_loss_weight", 1.0)
         self.l1_lambda = self.head_config.get("l1_lambda", 0.0)
         if self.l1_lambda > 0:
             logger.info(f"Applying L1 regularization with lambda={self.l1_lambda}")
@@ -103,12 +104,23 @@ def _setup_mlp_heads(self, config):
         else:
             head_input_size = config.hidden_size
 
-        self.exposure_head = MLPHead(input_size=head_input_size)
+        mlp_num_hidden_layers = self.head_config.get("mlp_num_hidden_layers", 1)
+        mlp_hidden_size_ratio = self.head_config.get("mlp_hidden_size_ratio", 2)
+
+        self.exposure_head = MLPHead(
+            input_size=head_input_size,
+            num_hidden_layers=mlp_num_hidden_layers,
+            hidden_size_ratio=mlp_hidden_size_ratio,
+        )
 
         # Create separate heads for each outcome
         self.outcome_heads = nn.ModuleDict()
         for outcome_name in self.outcome_names:
-            self.outcome_heads[outcome_name] = MLPHead(input_size=head_input_size + 1)
+            self.outcome_heads[outcome_name] = MLPHead(
+                input_size=head_input_size + 1,
+                num_hidden_layers=mlp_num_hidden_layers,
+                hidden_size_ratio=mlp_hidden_size_ratio,
+            )
 
     def _get_loss_fn(self, loss_name: str, loss_params: dict):
         """Returns the loss function instance based on the name."""
@@ -256,9 +268,11 @@ def _compute_losses(self, outputs, batch):
                 batch[EXPOSURE_TARGET].view(-1),
             )
             outputs.exposure_loss = exposure_loss
-            total_loss += exposure_loss
+            total_loss += self.exposure_loss_weight * exposure_loss
 
         # Only compute outcome losses for available labels
+        outcome_loss_sum = 0
+        n_outcomes = 0
         for outcome_name in self.outcome_names:
             if outcome_name not in batch:
                 continue
@@ -269,7 +283,11 @@ def _compute_losses(self, outputs, batch):
             )
 
             outputs.outcome_losses[outcome_name] = outcome_loss
-            total_loss += outcome_loss
+            outcome_loss_sum += outcome_loss
+            n_outcomes += 1
+
+        if n_outcomes > 0:
+            total_loss += outcome_loss_sum / n_outcomes
 
         # Add L1 regularization on the bottleneck representation
         if self.shared_representation and self.l1_lambda > 0:

corebehrt/modules/trainer/causal/trainer.py

@@ -140,15 +140,18 @@ def _train_step(self, batch: dict):
             # Collect individual task losses for PCGrad
             task_losses = []
 
-            # Add exposure loss if available
+            # Add exposure loss if available (weighted)
             if hasattr(outputs, "exposure_loss") and outputs.exposure_loss is not None:
-                task_losses.append(outputs.exposure_loss)
+                task_losses.append(
+                    self.model.exposure_loss_weight * outputs.exposure_loss
+                )
 
-            # Add outcome losses if available
+            # Add outcome losses if available (normalized by count)
             if hasattr(outputs, "outcome_losses") and outputs.outcome_losses:
+                n_outcomes = len(outputs.outcome_losses)
                 for outcome_loss in outputs.outcome_losses.values():
                     if outcome_loss is not None:
-                        task_losses.append(outcome_loss)
+                        task_losses.append(outcome_loss / n_outcomes)
 
             # If we don't have individual losses, fall back to total loss
             if not task_losses and hasattr(outputs, "loss"):

tests/test_functional/test_causal/test_cohort_sampler.py

@@ -0,0 +1,70 @@
+"""Unit tests for cohort sampling utilities."""
+
+import unittest
+
+import torch
+
+from corebehrt.functional.causal.cohort_sampler import sample_cohort
+
+
+class TestSampleCohort(unittest.TestCase):
+    def setUp(self):
+        self.full_pids = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
+
+    def test_default_method_is_resample(self):
+        """Default method should be resample (without replacement)."""
+        sampled = sample_cohort(self.full_pids, sample_size=5, seed=42)
+        self.assertEqual(len(sampled), 5)
+        # No duplicates
+        self.assertEqual(len(set(sampled.tolist())), 5)
+
+    def test_resample_no_duplicates(self):
+        """Resample method should produce no duplicate PIDs."""
+        sampled = sample_cohort(
+            self.full_pids, sample_size=10, seed=42, method="resample"
+        )
+        self.assertEqual(len(set(sampled.tolist())), 10)
+
+    def test_resample_rejects_oversize(self):
+        """Resample method should reject sample_size > n_total."""
+        with self.assertRaises(ValueError):
+            sample_cohort(self.full_pids, sample_size=20, seed=42, method="resample")
+
+    def test_bootstrap_allows_duplicates(self):
+        """Bootstrap method should sample with replacement (duplicates possible)."""
+        # Use a large sample to make duplicates very likely
+        sampled = sample_cohort(
+            self.full_pids, sample_size=100, seed=42, method="bootstrap"
+        )
+        self.assertEqual(len(sampled), 100)
+        # With 100 draws from 10 items, duplicates are virtually guaranteed
+        self.assertLess(len(set(sampled.tolist())), 100)
+
+    def test_bootstrap_accepts_oversize(self):
+        """Bootstrap method should accept sample_size > n_total."""
+        sampled = sample_cohort(
+            self.full_pids, sample_size=20, seed=42, method="bootstrap"
+        )
+        self.assertEqual(len(sampled), 20)
+
+    def test_invalid_method_raises(self):
+        """Invalid method should raise ValueError."""
+        with self.assertRaises(ValueError):
+            sample_cohort(self.full_pids, sample_size=5, seed=42, method="invalid")
+
+    def test_bootstrap_with_fraction(self):
+        """Bootstrap with fraction should work."""
+        sampled = sample_cohort(
+            self.full_pids, sample_fraction=0.5, seed=42, method="bootstrap"
+        )
+        self.assertEqual(len(sampled), 5)
+
+    def test_resample_reproducibility(self):
+        """Same seed should produce same results."""
+        s1 = sample_cohort(self.full_pids, sample_size=5, seed=123)
+        s2 = sample_cohort(self.full_pids, sample_size=5, seed=123)
+        self.assertTrue(torch.equal(s1, s2))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_modules/test_model/test_causal/test_head.py

@@ -218,6 +218,49 @@ def test_output_range(self):
         # Check that outputs are not extremely large (arbitrary threshold)
         self.assertTrue(torch.abs(output).max() < 100)
 
+    def test_default_matches_original_architecture(self):
+        """Test that default (num_hidden_layers=1) produces 4 modules: LayerNorm, Linear, GELU, Linear."""
+        head = MLPHead(self.input_size)
+        modules = list(head.classifier)
+        self.assertEqual(len(modules), 4)
+        self.assertIsInstance(modules[0], torch.nn.LayerNorm)
+        self.assertIsInstance(modules[1], torch.nn.Linear)
+        self.assertIsInstance(modules[2], torch.nn.GELU)
+        self.assertIsInstance(modules[3], torch.nn.Linear)
+
+    def test_multi_layer_forward_pass(self):
+        """Test forward pass with different depths."""
+        for depth in [1, 2, 3]:
+            with self.subTest(depth=depth):
+                head = MLPHead(self.input_size, num_hidden_layers=depth).to(self.device)
+                x = torch.randn(self.batch_size, self.input_size, device=self.device)
+                output = head(x)
+                self.assertEqual(output.shape, (self.batch_size, 1))
+
+    def test_multi_layer_module_count(self):
+        """Test that deeper heads have more modules."""
+        for depth in [1, 2, 3]:
+            with self.subTest(depth=depth):
+                head = MLPHead(self.input_size, num_hidden_layers=depth)
+                modules = list(head.classifier)
+                # LayerNorm + (Linear + GELU) * depth + Linear
+                expected = 1 + 2 * depth + 1
+                self.assertEqual(len(modules), expected)
+
+    def test_gradient_flow_deep_head(self):
+        """Test that gradients flow through a deep MLPHead."""
+        head = MLPHead(self.input_size, num_hidden_layers=3).to(self.device)
+        x = torch.randn(
+            self.batch_size, self.input_size, device=self.device, requires_grad=True
+        )
+
+        output = head(x)
+        loss = output.sum()
+        loss.backward()
+
+        self.assertIsNotNone(x.grad)
+        self.assertFalse(torch.allclose(x.grad, torch.zeros_like(x.grad)))
+
 
 if __name__ == "__main__":
     unittest.main()