Fix broken example scripts and fairness metrics import

docs/api/metrics/pyhealth.metrics.fairness.rst

@@ -1,6 +1,10 @@
 pyhealth.metrics.fairness
 ===================================
 
+``fairness_metrics_fn`` is also importable directly from the top-level
+``pyhealth.metrics`` package (``from pyhealth.metrics import
+fairness_metrics_fn``), not just from this submodule.
+
 .. currentmodule:: pyhealth.metrics.fairness
 
 .. autofunction:: fairness_metrics_fn

examples/benchmark_perf/loc/minimal_los.py

@@ -1,5 +1,6 @@
 from datetime import datetime; from pyhealth.datasets import MIMIC4Dataset; from pyhealth.tasks.base_task import BaseTask
-def categorize_los(d): return 0 if d<1 else (d if d<=7 else (8 if d<=14 else 9))class LengthOfStayPredictionMIMIC4(BaseTask):
+def categorize_los(d): return 0 if d<1 else (d if d<=7 else (8 if d<=14 else 9))
+class LengthOfStayPredictionMIMIC4(BaseTask):
     task_name="LengthOfStayPredictionMIMIC4"
     input_schema={"conditions":"sequence","procedures":"sequence","drugs":"sequence"}
     output_schema={"los":"multiclass"}

examples/readmission/readmission_mimic3_fairness.py

@@ -6,47 +6,47 @@
 from pyhealth.trainer import Trainer
 from pyhealth.metrics.fairness_utils.utils import sensitive_attributes_from_patient_ids
 
-# STEP 1: load data
-base_dataset = MIMIC3Dataset(
+# This script must be run under an `if __name__ == "__main__":` guard: the
+# dataloaders below use multiprocessing workers, and without the guard each
+# worker re-imports this module as if it were the main script, re-running
+# everything at the top level (including reconstructing MIMIC3Dataset) --
+# causing a runaway process-spawning hang instead of ever completing.
+if __name__ == "__main__":
+    # STEP 1: load data
+    base_dataset = MIMIC3Dataset(
         root="https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/",
         tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
-)
-base_dataset.stats()
-
-# STEP 2: set task
-sample_dataset = base_dataset.set_task(ReadmissionPredictionMIMIC3(exclude_minors=False)) # Must include minors to get any readmission samples on the synthetic dataset
-
-train_dataset, val_dataset, test_dataset = split_by_patient(sample_dataset, [0.8, 0.1, 0.1])
-train_dataloader = get_dataloader(train_dataset, batch_size=32, shuffle=True)
-val_dataloader = get_dataloader(val_dataset, batch_size=32, shuffle=False)
-test_dataloader = get_dataloader(test_dataset, batch_size=32, shuffle=False)
-
-# STEP 3: define model
-model = Transformer(
-    dataset=sample_dataset,
-    # look up what are available for "feature_keys" and "label_keys" in dataset.samples[0]
-    feature_keys=["conditions", "procedures"],
-    label_key="label",
-    mode="binary",
-)
-
-# STEP 4: define trainer
-trainer = Trainer(model=model)
-trainer.train(
-    train_dataloader=train_dataloader,
-    val_dataloader=val_dataloader,
-    epochs=3,
-    monitor="pr_auc",
-)
-
-# STEP 5: inference, return patient_ids
-y_true, y_prob, loss, patient_ids = trainer.inference(test_dataloader, return_patient_ids=True)
-
-# STEP 6: get sensitive attribute array from patient_ids
-sensitive_attribute_array = sensitive_attributes_from_patient_ids(base_dataset, patient_ids,
-                                                                  'gender', 'F')
-
-# STEP 7: use pyhealth.metrics to evaluate fairness
-fairness_metrics = fairness_metrics_fn(y_true, y_prob, sensitive_attribute_array,
-                                       favorable_outcome=0)
-print(fairness_metrics)
+    )
+    base_dataset.stats()
+
+    # STEP 2: set task
+    sample_dataset = base_dataset.set_task(ReadmissionPredictionMIMIC3(exclude_minors=False)) # Must include minors to get any readmission samples on the synthetic dataset
+
+    train_dataset, val_dataset, test_dataset = split_by_patient(sample_dataset, [0.8, 0.1, 0.1])
+    train_dataloader = get_dataloader(train_dataset, batch_size=32, shuffle=True)
+    val_dataloader = get_dataloader(val_dataset, batch_size=32, shuffle=False)
+    test_dataloader = get_dataloader(test_dataset, batch_size=32, shuffle=False)
+
+    # STEP 3: define model
+    model = Transformer(dataset=sample_dataset)
+
+    # STEP 4: define trainer
+    trainer = Trainer(model=model)
+    trainer.train(
+        train_dataloader=train_dataloader,
+        val_dataloader=val_dataloader,
+        epochs=3,
+        monitor="pr_auc",
+    )
+
+    # STEP 5: inference, return patient_ids
+    y_true, y_prob, loss, patient_ids = trainer.inference(test_dataloader, return_patient_ids=True)
+
+    # STEP 6: get sensitive attribute array from patient_ids
+    sensitive_attribute_array = sensitive_attributes_from_patient_ids(base_dataset, patient_ids,
+                                                                      'gender', 'F')
+
+    # STEP 7: use pyhealth.metrics to evaluate fairness
+    fairness_metrics = fairness_metrics_fn(y_true, y_prob, sensitive_attribute_array,
+                                           favorable_outcome=0)
+    print(fairness_metrics)

pyhealth/metrics/__init__.py

@@ -15,10 +15,9 @@
     SufficiencyMetric,
     evaluate_attribution,
 )
+from .fairness import fairness_metrics_fn
 from .multiclass import multiclass_metrics_fn
 from .multilabel import multilabel_metrics_fn
-
-# from .fairness import fairness_metrics_fn
 from .ranking import ranking_metrics_fn
 from .regression import regression_metrics_fn
 
@@ -36,6 +35,7 @@
     "RemovalBasedMetric",
     "Evaluator",
     "evaluate_attribution",
+    "fairness_metrics_fn",
     "multiclass_metrics_fn",
     "multilabel_metrics_fn",
     "ranking_metrics_fn",

pyhealth/metrics/fairness_utils/utils.py

@@ -2,28 +2,57 @@
 from typing import List
 import numpy as np
 
-from pyhealth.datasets import BaseEHRDataset
+from pyhealth.datasets import BaseDataset
 
-def sensitive_attributes_from_patient_ids(dataset: BaseEHRDataset,
+def sensitive_attributes_from_patient_ids(dataset: BaseDataset,
                                           patient_ids: List[str],
                                           sensitive_attribute: str,
                                           protected_group: str) -> np.ndarray:
     """
     Returns the desired sensitive attribute array from patient_ids.
 
     Args:
-        dataset: Dataset object.
+        dataset: Dataset object (must implement ``get_patient(patient_id)``
+            returning a :class:`~pyhealth.data.data.Patient` with a
+            ``"patients"``-typed demographic event).
         patient_ids: List of patient IDs.
         sensitive_attribute: Sensitive attribute to extract.
         protected_group: Value of the protected group.
-    
+
     Returns:
         Sensitive attribute array of shape (n_samples,).
+
+    Examples:
+        >>> import polars as pl
+        >>> from datetime import datetime
+        >>> from pyhealth.data import Patient
+        >>> event_df = pl.DataFrame({
+        ...     "patient_id": ["patient-0", "patient-1"],
+        ...     "event_type": ["patients", "patients"],
+        ...     "timestamp": [datetime(2020, 1, 1), datetime(2020, 1, 1)],
+        ...     "patients/gender": ["F", "M"],
+        ... })
+        >>> class ToyDataset:
+        ...     def get_patient(self, patient_id):
+        ...         return Patient(
+        ...             patient_id=patient_id,
+        ...             data_source=event_df.filter(pl.col("patient_id") == patient_id),
+        ...         )
+        >>> sensitive_attributes_from_patient_ids(
+        ...     ToyDataset(), ["patient-0", "patient-1"], "gender", "F"
+        ... )
+        array([1., 0.])
     """
 
     sensitive_attribute_array = np.zeros(len(patient_ids))
     for idx, patient_id in enumerate(patient_ids):
-        sensitive_attribute_value = getattr(dataset.patients[patient_id], sensitive_attribute)
+        patient = dataset.get_patient(patient_id)
+        demographic_events = patient.get_events(event_type="patients")
+        sensitive_attribute_value = (
+            demographic_events[0].attr_dict.get(sensitive_attribute)
+            if demographic_events
+            else None
+        )
         if sensitive_attribute_value == protected_group:
             sensitive_attribute_array[idx] = 1
     return sensitive_attribute_array