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``/
+   ``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,7 +32,10 @@ def prepare_drug_task_data():
     print("info")
     mimicvi.info()
 
-    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,10 +14,13 @@
 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,
     DrugRecommendationMIMIC4,
+    DrugRecommendationOMOP as DrugRecommendationOMOP,
     drug_recommendation_mimic3_fn,
     drug_recommendation_mimic4_fn,
     drug_recommendation_omop_fn,

pyhealth/tasks/drug_recommendation.py

@@ -1,4 +1,6 @@
 from typing import Any, Dict, Iterable, List, Optional
+from typing import ClassVar  # keep off line 1: merging would re-flag pre-existing UP035/I001
+from collections import defaultdict
 
 import polars as pl
 
@@ -644,6 +646,170 @@ 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 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 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",
+        ...     ],
+        ... )
+        >>> 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]]:
+        """Processes a patient into drug recommendation samples.
+
+        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 exposing ``get_events``.
+
+        Returns:
+            List of samples with patient_id, visit_id, cumulative conditions
+            and procedures, leak-free drugs history, and the target drugs.
+        """
+        visits = patient.get_events(event_type="visit_occurrence")
+        if len(visits) < 2:
+            return []
+
+        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 not (conditions and procedures and drugs):
+                continue
+            samples.append(
+                {
+                    "visit_id": visit_id,
+                    "patient_id": patient.patient_id,
+                    "conditions": conditions,
+                    "procedures": procedures,
+                    "drugs": drugs,
+                }
+            )
+
+        # Exclude patients with less than 2 valid visits
+        if len(samples) < 2:
+            return []
+
+        # 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
+
+
 def drug_recommendation_omop_fn(patient: Patient):
     """Processes a single patient for the drug recommendation task.
 
@@ -709,6 +875,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