From 1b72a4fc9087bb00f0d8dfb80c32e12b8560d45b Mon Sep 17 00:00:00 2001 From: lehendo Date: Tue, 25 Aug 2026 14:46:40 -0500 Subject: [PATCH] Fix drug_recommendation_omop_fn self-leakage and add DrugRecommendationOMOP - drug_recommendation_omop_fn never excluded the current visit's own drugs from drugs_all, unlike every sibling drug-recommendation function/class, so the last history entry was identical to the prediction target. Fixed to match the established pattern. - Add DrugRecommendationOMOP, a current-API, leak-free class-based replacement, since drug_recommendation_omop_fn (and its mimic3/mimic4 siblings) cannot actually run through set_task() under the current Patient/Visit API -- OMOP drug recommendation had no working path at all before this. Verified end-to-end against real OMOP demo data. - Document the legacy function family's set_task() incompatibility and flag the same issue in the one live example that hits it. --- .../pyhealth.tasks.drug_recommendation.rst | 15 ++ .../drug_recommendation_mimic4_gamenet.py | 4 + pyhealth/tasks/__init__.py | 1 + pyhealth/tasks/drug_recommendation.py | 143 ++++++++++++++++++ .../test_drug_recommendation_omop_leakage.py | 143 ++++++++++++++++++ 5 files changed, 306 insertions(+) create mode 100644 tests/core/test_drug_recommendation_omop_leakage.py diff --git a/docs/api/tasks/pyhealth.tasks.drug_recommendation.rst b/docs/api/tasks/pyhealth.tasks.drug_recommendation.rst index fe4c9f8ff..eaf23ded5 100644 --- a/docs/api/tasks/pyhealth.tasks.drug_recommendation.rst +++ b/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 \ No newline at end of file diff --git a/examples/drug_recommendation/drug_recommendation_mimic4_gamenet.py b/examples/drug_recommendation/drug_recommendation_mimic4_gamenet.py index bd5b33cb0..04eed866e 100644 --- a/examples/drug_recommendation/drug_recommendation_mimic4_gamenet.py +++ b/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]) diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index df8411db0..f28b8ca76 100644 --- a/pyhealth/tasks/__init__.py +++ b/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, diff --git a/pyhealth/tasks/drug_recommendation.py b/pyhealth/tasks/drug_recommendation.py index ec113e1dd..3e4fea3bb 100644 --- a/pyhealth/tasks/drug_recommendation.py +++ b/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 diff --git a/tests/core/test_drug_recommendation_omop_leakage.py b/tests/core/test_drug_recommendation_omop_leakage.py new file mode 100644 index 000000000..34ac6c0ce --- /dev/null +++ b/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()