Add DrugRecommendationOMOP task

docs/api/tasks/pyhealth.tasks.drug_recommendation.rst

@@ -19,9 +19,24 @@ Task Classes
     :undoc-members:
     :show-inheritance:
 
+.. autoclass:: pyhealth.tasks.drug_recommendation.DrugRecommendationOMOP
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
 Task Functions (Legacy)
 ------------------------
 
+.. note::
+
+   These functions predate the current dataset API: they expect an
+   indexable, ``len()``-able ``Patient`` with ``Visit.get_code_list(table)``,
+   neither of which the current ``pyhealth.data.Patient``/``Visit`` classes
+   provide (``Visit`` is now a deprecated no-op stub). As a result they
+   cannot currently be run through ``BaseDataset.set_task()``. Prefer the
+   task classes above (``DrugRecommendationMIMIC3``/``MIMIC4``/``EICU``),
+   which use the current API and are actively maintained.
+
 .. autofunction:: pyhealth.tasks.drug_recommendation.drug_recommendation_mimic3_fn
 .. autofunction:: pyhealth.tasks.drug_recommendation.drug_recommendation_mimic4_fn
 .. autofunction:: pyhealth.tasks.drug_recommendation.drug_recommendation_omop_fn

examples/drug_recommendation/drug_recommendation_mimic4_gamenet.py

@@ -32,6 +32,10 @@ def prepare_drug_task_data():
     print("info")
     mimicvi.info()
 
+    # NOTE: drug_recommendation_mimic4_fn is a legacy, pre-2.0 task function
+    # (expects an indexable Patient with Visit.get_code_list()) and is not
+    # compatible with the current BaseDataset.set_task(), which requires a
+    # BaseTask instance. Use DrugRecommendationMIMIC4() instead.
     mimic4_sample = mimicvi.set_task(drug_recommendation_mimic4_fn)
     print(mimic4_sample[0])
 

pyhealth/tasks/__init__.py

@@ -18,6 +18,7 @@
     DrugRecommendationEICU,
     DrugRecommendationMIMIC3,
     DrugRecommendationMIMIC4,
+    DrugRecommendationOMOP as DrugRecommendationOMOP,
     drug_recommendation_mimic3_fn,
     drug_recommendation_mimic4_fn,
     drug_recommendation_omop_fn,

pyhealth/tasks/drug_recommendation.py

@@ -1,4 +1,5 @@
 from typing import Any, Dict, Iterable, List, Optional
+from typing import ClassVar
 
 import polars as pl
 
@@ -644,6 +645,143 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]:
         return samples
 
 
+class DrugRecommendationOMOP(BaseTask):
+    """Task for drug recommendation using an OMOP CDM dataset.
+
+    Drug recommendation aims at recommending a set of drugs given the patient health
+    history (e.g., conditions and procedures). This task creates samples with
+    cumulative history, where each visit includes all previous visit information.
+
+    Features key-value pairs:
+    - using condition_occurrence table as condition codes
+    - using procedure_occurrence table as procedure codes
+    - using drug_exposure table as drug codes
+
+    Attributes:
+        task_name (str): The name of the task.
+        input_schema (Dict[str, str]): The schema for input data:
+            - conditions: Nested list of condition concept ids (history + current)
+            - procedures: Nested list of procedure concept ids (history + current)
+            - drugs_hist: Nested list of drug concept ids from history (current
+              visit excluded)
+        output_schema (Dict[str, str]): The schema for output data:
+            - drugs: List of drug concept ids to predict for current visit
+
+    Examples:
+        >>> from pyhealth.datasets import OMOPDataset
+        >>> from pyhealth.tasks import DrugRecommendationOMOP
+        >>> dataset = OMOPDataset(
+        ...     root="/path/to/omop",
+        ...     tables=["condition_occurrence", "procedure_occurrence", "drug_exposure"],
+        ... )
+        >>> task = DrugRecommendationOMOP()
+        >>> sample_dataset = dataset.set_task(task)
+    """
+
+    task_name: str = "DrugRecommendationOMOP"
+    input_schema: ClassVar[dict[str, str]] = {
+        "conditions": "nested_sequence",
+        "procedures": "nested_sequence",
+        "drugs_hist": "nested_sequence",
+    }
+    output_schema: ClassVar[dict[str, str]] = {"drugs": "multilabel"}
+
+    def __call__(self, patient: Any) -> list[dict[str, Any]]:
+        """Process a patient to create drug recommendation samples.
+
+        Creates one sample per visit (after first visit) with cumulative history.
+        Each sample includes all previous visits' conditions, procedures, and drugs.
+
+        Args:
+            patient: Patient object with get_events method
+
+        Returns:
+            List of samples, each with patient_id, visit_id, conditions history,
+            procedures history, drugs history, and target drugs
+        """
+        samples = []
+
+        # Get all visit occurrences
+        visit_occurrences = patient.get_events(event_type="visit_occurrence")
+        if len(visit_occurrences) < 2:
+            # Need at least 2 visits for history-based prediction
+            return []
+
+        # Process each visit
+        for visit in visit_occurrences:
+            condition_events = patient.get_events(
+                event_type="condition_occurrence",
+                filters=[("visit_occurrence_id", "==", visit.visit_occurrence_id)],
+            )
+            conditions = [
+                str(event.condition_concept_id)
+                for event in condition_events
+                if getattr(event, "condition_concept_id", None) is not None
+            ]
+
+            procedure_events = patient.get_events(
+                event_type="procedure_occurrence",
+                filters=[("visit_occurrence_id", "==", visit.visit_occurrence_id)],
+            )
+            procedures = [
+                str(event.procedure_concept_id)
+                for event in procedure_events
+                if getattr(event, "procedure_concept_id", None) is not None
+            ]
+
+            drug_events = patient.get_events(
+                event_type="drug_exposure",
+                filters=[("visit_occurrence_id", "==", visit.visit_occurrence_id)],
+            )
+            drugs = [
+                str(event.drug_concept_id)
+                for event in drug_events
+                if getattr(event, "drug_concept_id", None) is not None
+            ]
+
+            # Exclude visits without condition, procedure, or drug code
+            if len(conditions) * len(procedures) * len(drugs) == 0:
+                continue
+
+            samples.append(
+                {
+                    "visit_id": visit.visit_occurrence_id,
+                    "patient_id": patient.patient_id,
+                    "conditions": conditions,
+                    "procedures": procedures,
+                    "drugs": drugs,
+                    "drugs_hist": drugs,
+                }
+            )
+
+        # Exclude patients with less than 2 valid visits
+        if len(samples) < 2:
+            return []
+
+        # Add cumulative history for first sample
+        samples[0]["conditions"] = [samples[0]["conditions"]]
+        samples[0]["procedures"] = [samples[0]["procedures"]]
+        samples[0]["drugs_hist"] = [samples[0]["drugs_hist"]]
+
+        # Add cumulative history for subsequent samples
+        for i in range(1, len(samples)):
+            samples[i]["conditions"] = samples[i - 1]["conditions"] + [
+                samples[i]["conditions"]
+            ]
+            samples[i]["procedures"] = samples[i - 1]["procedures"] + [
+                samples[i]["procedures"]
+            ]
+            samples[i]["drugs_hist"] = samples[i - 1]["drugs_hist"] + [
+                samples[i]["drugs_hist"]
+            ]
+
+        # Remove target drug from history (set current visit drugs_hist to empty)
+        for i in range(len(samples)):
+            samples[i]["drugs_hist"][i] = []
+
+        return samples
+
+
 def drug_recommendation_omop_fn(patient: Patient):
     """Processes a single patient for the drug recommendation task.
 
@@ -709,6 +847,11 @@ def drug_recommendation_omop_fn(patient: Patient):
             samples[i]["drugs_all"]
         ]
 
+    # remove the target drug from the history (mirrors drugs_hist handling
+    # in the other drug_recommendation_*_fn / DrugRecommendation* tasks)
+    for i in range(len(samples)):
+        samples[i]["drugs_all"][i] = []
+
     return samples
 
 

tests/core/test_drug_recommendation_omop_leakage.py

@@ -0,0 +1,143 @@
+import unittest
+from pathlib import Path
+
+from pyhealth.datasets import OMOPDataset
+from pyhealth.tasks import DrugRecommendationOMOP
+from pyhealth.tasks.drug_recommendation import drug_recommendation_omop_fn
+
+
+class _MockVisit:
+    """Minimal stand-in for the legacy pyhealth.data.Visit interface that
+    drug_recommendation_omop_fn expects (visit_id, get_code_list(table)).
+    """
+
+    def __init__(self, visit_id, codes):
+        self.visit_id = visit_id
+        self._codes = codes
+
+    def get_code_list(self, table):
+        return self._codes[table]
+
+
+class _MockPatient:
+    """Minimal stand-in for the legacy indexable/len-able Patient interface
+    that drug_recommendation_omop_fn expects (patient_id, len(), [i]).
+    """
+
+    def __init__(self, patient_id, visits):
+        self.patient_id = patient_id
+        self._visits = visits
+
+    def __len__(self):
+        return len(self._visits)
+
+    def __getitem__(self, i):
+        return self._visits[i]
+
+
+class TestDrugRecommendationOMOPLeakage(unittest.TestCase):
+    """Regression test for the drugs_all self-leakage bug.
+
+    drug_recommendation_omop_fn built a "drugs_all" history feature by
+    accumulating each visit's own drugs without ever excluding the current
+    visit -- unlike its class-based siblings (DrugRecommendationMIMIC3/4/
+    EICU) and the drug_recommendation_mimic3_fn/mimic4_fn functions, all of
+    which zero out the current visit's slot in the history sequence. That
+    meant the last entry of "drugs_all" was identical to the "drugs" target
+    for every sample, so a model could trivially copy it instead of
+    predicting from history.
+    """
+
+    def setUp(self):
+        self.visits = [
+            _MockVisit(
+                "v1",
+                {
+                    "condition_occurrence": ["C1"],
+                    "procedure_occurrence": ["P1"],
+                    "drug_exposure": ["D1"],
+                },
+            ),
+            _MockVisit(
+                "v2",
+                {
+                    "condition_occurrence": ["C2"],
+                    "procedure_occurrence": ["P2"],
+                    "drug_exposure": ["D2"],
+                },
+            ),
+            _MockVisit(
+                "v3",
+                {
+                    "condition_occurrence": ["C3"],
+                    "procedure_occurrence": ["P3"],
+                    "drug_exposure": ["D3"],
+                },
+            ),
+        ]
+        self.patient = _MockPatient("pt1", self.visits)
+
+    def test_drugs_all_excludes_current_visit_drugs(self):
+        samples = drug_recommendation_omop_fn(self.patient)
+        self.assertEqual(len(samples), 3)
+
+        for i, sample in enumerate(samples):
+            with self.subTest(visit=sample["visit_id"]):
+                self.assertEqual(
+                    sample["drugs_all"][i],
+                    [],
+                    "current visit's own drugs must not leak into its own "
+                    "history slot",
+                )
+
+    def test_drugs_all_preserves_prior_visit_history(self):
+        samples = drug_recommendation_omop_fn(self.patient)
+
+        # visit 2's history should still contain visit 1's drugs
+        self.assertEqual(samples[1]["drugs_all"][0], ["D1"])
+        # visit 3's history should still contain visits 1 and 2's drugs
+        self.assertEqual(samples[2]["drugs_all"][0], ["D1"])
+        self.assertEqual(samples[2]["drugs_all"][1], ["D2"])
+
+    def test_drugs_target_unaffected(self):
+        samples = drug_recommendation_omop_fn(self.patient)
+        self.assertEqual(samples[0]["drugs"], ["D1"])
+        self.assertEqual(samples[1]["drugs"], ["D2"])
+        self.assertEqual(samples[2]["drugs"], ["D3"])
+
+
+class TestDrugRecommendationOMOP(unittest.TestCase):
+    """DrugRecommendationOMOP is the current-API, leak-free replacement for
+    the legacy drug_recommendation_omop_fn (which cannot even run under the
+    current dataset API -- see the docs note on this task family). Verified
+    against real demo OMOP data.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        root = str(Path(__file__).parents[2] / "test-resources" / "omop")
+        tables = ["condition_occurrence", "procedure_occurrence", "drug_exposure"]
+        cls.dataset = OMOPDataset(root=root, tables=tables)
+
+    def test_drugs_hist_excludes_current_visit_and_preserves_history(self):
+        # person_id "1" has 4 chronological visits (ids "1".."4"), each with
+        # exactly one condition/procedure/drug code (all coded "1").
+        patient = self.dataset.get_patient("1")
+        samples = DrugRecommendationOMOP()(patient)
+        self.assertEqual(len(samples), 4)
+
+        for i, sample in enumerate(samples):
+            with self.subTest(visit=sample["visit_id"]):
+                self.assertEqual(sample["drugs"], ["1"])
+                self.assertEqual(
+                    sample["drugs_hist"][i],
+                    [],
+                    "current visit's own drugs must not leak into its own "
+                    "history slot",
+                )
+                for j in range(i):
+                    self.assertEqual(sample["drugs_hist"][j], ["1"])
+
+
+if __name__ == "__main__":
+    unittest.main()