Fix comprehensiveness/sufficiency crash on code-sequence inputs

docs/api/metrics/pyhealth.metrics.interpretability.rst

@@ -29,6 +29,14 @@ For binary classifiers, a sample filter can mark class-0 predictions as
 from the class-0 perspective. Each percentage is evaluated independently, so a
 percentage's score does not depend on the other requested percentages.
 
+Integer inputs, such as the code ids of ``sequence`` and ``nested_sequence``
+fields, are always ablated by setting the removed codes to the padding id 0,
+whatever the ``ablation_strategy``; averaging or adding noise to code ids has
+no meaning. A sample whose codes would all be removed keeps its first code, so
+every sequence has at least one non-padding position. Float inputs use the
+chosen strategy (``"zero"``, ``"mean"`` or ``"noise"``). See
+``examples/interpretability/comprehensiveness_code_sequences.py``.
+
 Base Class
 ^^^^^^^^^^
 

examples/interpretability/comprehensiveness_code_sequences.py

@@ -0,0 +1,63 @@
+"""Comprehensiveness and sufficiency for a model on code sequences.
+
+Removal-based metrics ablate the code ids of ``sequence`` inputs by setting
+them to padding. This example trains an MLP on a small synthetic dataset and
+scores Integrated Gradients attributions; it runs in seconds on a CPU and needs
+no downloads.
+
+Usage:
+    python examples/interpretability/comprehensiveness_code_sequences.py
+"""
+
+import random
+
+import torch
+
+from pyhealth.datasets import create_sample_dataset, get_dataloader
+from pyhealth.interpret.methods import IntegratedGradients
+from pyhealth.metrics.interpretability import (
+    evaluate_attribution,
+    threshold_sample_filter,
+)
+from pyhealth.models import MLP
+from pyhealth.trainer import Trainer
+
+
+def main():
+    random.seed(0)
+    torch.manual_seed(0)
+    # The label is 1 when code "risk" is present, so a faithful explanation
+    # should point at it.
+    samples = []
+    for i in range(400):
+        codes = random.sample([f"c{k}" for k in range(30)], 6)
+        label = i % 2
+        if label:
+            codes[random.randrange(6)] = "risk"
+        samples.append({"patient_id": f"p{i}", "codes": codes, "label": label})
+    dataset = create_sample_dataset(
+        samples=samples,
+        input_schema={"codes": "sequence"},
+        output_schema={"label": "binary"},
+        dataset_name="synthetic_codes",
+    )
+    loader = get_dataloader(dataset, batch_size=64, shuffle=True)
+
+    model = MLP(dataset=dataset)
+    Trainer(model=model, enable_logging=False).train(
+        train_dataloader=loader, epochs=5
+    )
+    model.eval()
+
+    scores = evaluate_attribution(
+        model,
+        get_dataloader(dataset, batch_size=64),
+        IntegratedGradients(model, steps=16),
+        percentages=[10, 20, 50],
+        sample_filter=threshold_sample_filter(0.5),
+    )
+    print(scores)
+
+
+if __name__ == "__main__":
+    main()

pyhealth/metrics/interpretability/base.py

@@ -266,37 +266,9 @@ def _apply_ablation(
 
                     mask = masks[key]
 
-                    # Check if values are integers (discrete features)
-                    is_discrete = x_values.dtype in [
-                        torch.long,
-                        torch.int,
-                        torch.int32,
-                        torch.int64,
-                    ]
-
                     # Apply ablation to values part
-                    if is_discrete:
-                        # For discrete features (codes), multiply by (1-mask)
-                        # Where mask=1 (ablate): set to 0 (padding index)
-                        # Where mask=0 (keep): preserve original value
-                        ablated_values = x_values * (1 - mask).long()
-
-                        # Safety: prevent complete ablation of sequences
-                        # Complete ablation (all zeros) causes issues in
-                        # StageNet:
-                        # - All embeddings become zero (padding_idx=0)
-                        # - Mask becomes all zeros
-                        # - get_last_visit() tries to index with -1
-                        # Solution: keep at least one non-zero element
-                        for b in range(ablated_values.shape[0]):
-                            if ablated_values[b].sum() == 0:
-                                non_zero_mask = x_values[b] != 0
-                                if non_zero_mask.any():
-                                    # Keep first non-zero element
-                                    first_idx = non_zero_mask.nonzero()[0]
-                                    ablated_values[b][tuple(first_idx)] = x_values[b][
-                                        tuple(first_idx)
-                                    ]
+                    if _is_discrete(x_values):
+                        ablated_values = _ablate_codes(x_values, mask)
                     else:
                         # For continuous features, apply standard ablation
                         if self.ablation_strategy == "zero":
@@ -332,8 +304,13 @@ def _apply_ablation(
 
             mask = masks[key]
 
+            # Code ids can only be removed (set to padding), whatever the
+            # strategy; averaging or adding noise to ids is meaningless.
+            if _is_discrete(x):
+                ablated_inputs[key] = _ablate_codes(x, mask)
+
             # Apply ablation strategy
-            if self.ablation_strategy == "zero":
+            elif self.ablation_strategy == "zero":
                 # Set ablated features to 0
                 ablated_inputs[key] = x * (1 - mask)
 
@@ -527,3 +504,39 @@ def compute(
             # Average across percentages
             metric_scores = metric_scores / len(self.percentages) # type: ignore
             return metric_scores, val_mask
+
+
+def _is_discrete(x: torch.Tensor) -> bool:
+    """Whether ``x`` holds discrete values such as code ids."""
+    return not (x.is_floating_point() or x.is_complex() or x.dtype == torch.bool)
+
+
+def _ablate_codes(x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    """Sets the masked code ids to padding (0), keeping the dtype.
+
+    A sample whose codes would all be removed keeps its first non-padding
+    code, because models such as StageNet and RNN need at least one
+    non-padding position per sequence.
+
+    Args:
+        x: Integer code ids, shape ``(batch, ...)``.
+        mask: Same shape as ``x``; 1 marks the codes to remove.
+
+    Returns:
+        A tensor like ``x`` with the masked codes set to 0.
+
+    Examples:
+        >>> import torch
+        >>> _ablate_codes(torch.tensor([[3, 4, 5]]), torch.tensor([[1.0, 0.0, 1.0]]))
+        tensor([[0, 4, 0]])
+        >>> _ablate_codes(torch.tensor([[3, 4, 0]]), torch.ones(1, 3))
+        tensor([[3, 0, 0]])
+    """
+    ablated = torch.where(mask.bool(), torch.zeros_like(x), x)
+    for b in range(ablated.shape[0]):
+        if not ablated[b].any():
+            nonzero = (x[b] != 0).nonzero()
+            if len(nonzero):
+                first = tuple(nonzero[0])
+                ablated[b][first] = x[b][first]
+    return ablated

tests/core/test_interp_ablation_codes.py

@@ -0,0 +1,89 @@
+"""Removal-based metrics on code-sequence inputs (integer code ids)."""
+
+import math
+import unittest
+
+import torch
+
+from pyhealth.datasets import create_sample_dataset, get_dataloader
+from pyhealth.interpret.methods import IntegratedGradients
+from pyhealth.metrics.interpretability import (
+    ComprehensivenessMetric,
+    evaluate_attribution,
+    threshold_sample_filter,
+)
+from pyhealth.metrics.interpretability.base import _ablate_codes
+from pyhealth.models import MLP, Transformer
+
+
+def _dataset():
+    samples = [
+        {
+            "patient_id": f"p{i}",
+            "codes": [f"c{(i + j) % 7}" for j in range(3 + i % 3)],
+            "label": i % 2,
+        }
+        for i in range(20)
+    ]
+    return create_sample_dataset(
+        samples=samples,
+        input_schema={"codes": "sequence"},
+        output_schema={"label": "binary"},
+        dataset_name="ablation_codes",
+    )
+
+
+class TestAblateCodes(unittest.TestCase):
+    def test_masked_codes_become_padding_and_keep_dtype(self):
+        x = torch.tensor([[3, 4, 5, 0]])
+        out = _ablate_codes(x, torch.tensor([[1.0, 0.0, 1.0, 0.0]]))
+        self.assertEqual(out.dtype, torch.long)
+        self.assertEqual(out.tolist(), [[0, 4, 0, 0]])
+
+    def test_fully_ablated_sample_keeps_first_code(self):
+        out = _ablate_codes(torch.tensor([[0, 6, 7]]), torch.ones(1, 3))
+        self.assertEqual(out.tolist(), [[0, 6, 0]])
+
+    def test_float_inputs_still_use_the_strategy(self):
+        model = MLP(dataset=_dataset())
+        metric = ComprehensivenessMetric(
+            model, ablation_strategy="zero", sample_filter=threshold_sample_filter()
+        )
+        x = torch.tensor([[1.5, 2.5]])
+        out = metric._apply_ablation({"x": x}, {"x": torch.tensor([[1.0, 0.0]])})
+        self.assertTrue(torch.equal(out["x"], torch.tensor([[0.0, 2.5]])))
+
+
+class TestRemovalMetricsOnCodeSequences(unittest.TestCase):
+    """Comprehensiveness/sufficiency used to crash: code ids became floats."""
+
+    @classmethod
+    def setUpClass(cls):
+        torch.manual_seed(0)
+        cls.dataset = _dataset()
+
+    def _check(self, model_class):
+        model = model_class(dataset=self.dataset)
+        model.eval()
+        loader = get_dataloader(self.dataset, batch_size=8)
+        for strategy in ["zero", "mean", "noise"]:
+            with self.subTest(model=model_class.__name__, strategy=strategy):
+                scores = evaluate_attribution(
+                    model,
+                    loader,
+                    IntegratedGradients(model, steps=5),
+                    ablation_strategy=strategy,
+                    sample_filter=threshold_sample_filter(0.0),
+                )
+                for name, value in scores.items():
+                    self.assertTrue(math.isfinite(value), (name, value))
+
+    def test_mlp(self):
+        self._check(MLP)
+
+    def test_transformer(self):
+        self._check(Transformer)
+
+
+if __name__ == "__main__":
+    unittest.main()