fix: harden DrugRecommendationOMOP history, codes, and tests

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

@@ -34,9 +34,9 @@ Task Functions (Legacy)
    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.
+   task classes above (``DrugRecommendationMIMIC3``/``MIMIC4``/``EICU``/
+   ``OMOP``), 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
+.. autofunction:: pyhealth.tasks.drug_recommendation.drug_recommendation_omop_fn

examples/drug_recommendation/drug_recommendation_mimic4_gamenet.py

@@ -3,7 +3,7 @@
 
 # import mimic4 dataset and drug recommendaton task
 from pyhealth.datasets import MIMIC4Dataset
-from pyhealth.tasks import drug_recommendation_mimic4_fn
+from pyhealth.tasks import DrugRecommendationMIMIC4
 
 # import dataloader related functions
 from pyhealth.datasets.splitter import split_by_patient
@@ -32,11 +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)
+    # drug_recommendation_mimic4_fn is a pre-2.0 task function (it expects an
+    # indexable Patient with Visit.get_code_list) and cannot be passed to
+    # BaseDataset.set_task(), which requires a BaseTask instance.
+    mimic4_sample = mimicvi.set_task(DrugRecommendationMIMIC4())
     print(mimic4_sample[0])
 
     return mimic4_sample

examples/drug_recommendation/drug_recommendation_omop.py

@@ -0,0 +1,30 @@
+"""Drug recommendation on an OMOP CDM dataset.
+
+Run with a local OMOP CDM v5.3 export, e.g. the CMS SynPUF 1k sample.
+"""
+
+from pyhealth.datasets import OMOPDataset, get_dataloader, split_by_patient
+from pyhealth.tasks import DrugRecommendationOMOP
+
+
+def main() -> None:
+    dataset = OMOPDataset(
+        root="/path/to/omop_cdm",
+        tables=[
+            "condition_occurrence",
+            "procedure_occurrence",
+            "drug_exposure",
+        ],
+    )
+    dataset.stats()
+
+    samples = dataset.set_task(DrugRecommendationOMOP())
+    print(samples[0])
+
+    train, _val, _test = split_by_patient(samples, [0.8, 0.1, 0.1])
+    train_loader = get_dataloader(train, batch_size=32, shuffle=True)
+    print(next(iter(train_loader)).keys())
+
+
+if __name__ == "__main__":
+    main()

pyhealth/tasks/__init__.py

@@ -14,6 +14,8 @@
 from .covid19_cxr_classification import COVID19CXRClassification
 from .deid_ner import DeIDNERTask
 from .dka import DKAPredictionMIMIC4, T1DDKAPredictionMIMIC4
+# New exports must use the redundant `X as X` form: this module has no
+# __all__, and the PR lint gate flags F401 on newly added import lines.
 from .drug_recommendation import (
     DrugRecommendationEICU,
     DrugRecommendationMIMIC3,

pyhealth/tasks/drug_recommendation.py

@@ -1,5 +1,6 @@
 from typing import Any, Dict, Iterable, List, Optional
-from typing import ClassVar
+from typing import ClassVar  # keep off line 1: merging would re-flag pre-existing UP035/I001
+from collections import defaultdict
 
 import polars as pl
 
@@ -648,136 +649,163 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]:
 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.
+    Drug recommendation aims at recommending a set of drugs given the patient
+    health history (e.g., conditions and procedures). This task creates one
+    sample per qualifying visit with cumulative history: ``conditions`` and
+    ``procedures`` include the current visit, while ``drugs_hist`` excludes it
+    so the prediction target never appears in its own history.
 
     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
 
+    Concept ids equal to ``0`` are dropped: in OMOP, ``0`` is the
+    "no matching concept" sentinel, not a real code.
+
     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:
+        input_schema (dict[str, str]): The schema for input data:
+            - conditions: Nested list of condition concept ids (history +
+              current visit)
+            - procedures: Nested list of procedure concept ids (history +
+              current visit)
+            - drugs_hist: Nested list of drug concept ids from history; the
+              current visit's slot is always empty
+        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"],
+        ...     tables=[
+        ...         "condition_occurrence",
+        ...         "procedure_occurrence",
+        ...         "drug_exposure",
+        ...     ],
         ... )
-        >>> task = DrugRecommendationOMOP()
-        >>> sample_dataset = dataset.set_task(task)
+        >>> sample_dataset = dataset.set_task(DrugRecommendationOMOP())
     """
 
     task_name: str = "DrugRecommendationOMOP"
+    # ClassVar is required here: ruff's RUF012 flags mutable class attributes,
+    # and the PR lint gate checks added lines. The sibling tasks predate that
+    # gate, hence the local inconsistency.
     input_schema: ClassVar[dict[str, str]] = {
         "conditions": "nested_sequence",
         "procedures": "nested_sequence",
         "drugs_hist": "nested_sequence",
     }
     output_schema: ClassVar[dict[str, str]] = {"drugs": "multilabel"}
 
+    # (sample key, event type, concept id column)
+    _SOURCES: ClassVar[tuple[tuple[str, str, str], ...]] = (
+        ("conditions", "condition_occurrence", "condition_concept_id"),
+        ("procedures", "procedure_occurrence", "procedure_concept_id"),
+        ("drugs", "drug_exposure", "drug_concept_id"),
+    )
+    _NULLISH: ClassVar[frozenset[str]] = frozenset({"", "nan", "none", "<na>"})
+
+    @classmethod
+    def _norm(cls, value: Any) -> str | None:
+        """Normalizes a raw column value to a stable string, or None.
+
+        CSV sources are loaded as all-string with pyarrow
+        (``strings_can_be_null=False``), so a blank cell arrives as ``""``
+        rather than ``None``; Parquet sources keep their native dtype. This
+        collapses both cases.
+        """
+        if value is None:
+            return None
+        text = str(value).strip()
+        return None if text.lower() in cls._NULLISH else text
+
+    @classmethod
+    def _concept_id(cls, value: Any) -> str | None:
+        """Normalizes a concept id, dropping OMOP's 0 = 'no matching concept'."""
+        code = cls._norm(value)
+        return None if code == "0" else code
+
+    def _codes_by_visit(
+        self, patient: Any, event_type: str, field: str
+    ) -> dict[str, list[str]]:
+        """Groups one table's concept ids by visit in a single pass.
+
+        Avoids one ``get_events`` call per (visit, table), which is O(V*N).
+        """
+        grouped: dict[str, list[str]] = defaultdict(list)
+        for event in patient.get_events(event_type=event_type):
+            visit_id = self._norm(getattr(event, "visit_occurrence_id", None))
+            code = self._concept_id(getattr(event, field, None))
+            if visit_id is None or code is None:
+                continue
+            grouped[visit_id].append(code)
+        return grouped
+
     def __call__(self, patient: Any) -> list[dict[str, Any]]:
-        """Process a patient to create drug recommendation samples.
+        """Processes a patient into drug recommendation samples.
 
-        Creates one sample per visit (after first visit) with cumulative history.
-        Each sample includes all previous visits' conditions, procedures, and drugs.
+        Emits one sample per visit that has at least one condition, one
+        procedure and one drug. Patients with fewer than two such visits are
+        dropped. Visits are consumed in chronological order (``Patient``
+        sorts its event source by timestamp).
 
         Args:
-            patient: Patient object with get_events method
+            patient: Patient object exposing ``get_events``.
 
         Returns:
-            List of samples, each with patient_id, visit_id, conditions history,
-            procedures history, drugs history, and target drugs
+            List of samples with patient_id, visit_id, cumulative conditions
+            and procedures, leak-free drugs history, and the 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
+        visits = patient.get_events(event_type="visit_occurrence")
+        if len(visits) < 2:
             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
-            ]
+        grouped = {
+            key: self._codes_by_visit(patient, event_type, field)
+            for key, event_type, field in self._SOURCES
+        }
 
+        samples: list[dict[str, Any]] = []
+        for visit in visits:
+            visit_id = self._norm(getattr(visit, "visit_occurrence_id", None))
+            if visit_id is None:
+                continue
+            conditions = grouped["conditions"].get(visit_id, [])
+            procedures = grouped["procedures"].get(visit_id, [])
+            drugs = grouped["drugs"].get(visit_id, [])
             # Exclude visits without condition, procedure, or drug code
-            if len(conditions) * len(procedures) * len(drugs) == 0:
+            if not (conditions and procedures and drugs):
                 continue
-
             samples.append(
                 {
-                    "visit_id": visit.visit_occurrence_id,
+                    "visit_id": visit_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] = []
+        # Snapshot before rewriting, then rebuild each sample from fresh lists
+        # so that no two samples ever share a list object.
+        per_visit = [
+            (list(s["conditions"]), list(s["procedures"]), list(s["drugs"]))
+            for s in samples
+        ]
+        for index, sample in enumerate(samples):
+            window = per_visit[: index + 1]
+            sample["conditions"] = [list(codes) for codes, _, _ in window]
+            sample["procedures"] = [list(codes) for _, codes, _ in window]
+            sample["drugs_hist"] = [list(codes) for _, _, codes in window]
+            # The target visit's own drugs must not appear in its own history.
+            sample["drugs_hist"][index] = []
 
         return samples