MIMIC-III AMA prediction task + examples + tests

[Madhav-Kanda]

Contributors

• Names: Aryan Taneja, Khushi Garg, Madhav Kanda
• NetID / email: aryant2, kagarg2, madhav3

Type of contribution

• Primary: task — new AMAPredictionMIMIC3 task for against-medical-advice (AMA) discharge on MIMIC-III.
• Secondary: examples + tests + docs (API documentation for the task).

Paper (reproducibility)

Boag et al., Racial Disparities and Mistrust in End-of-Life Care, 3rd Machine Learning for Healthcare Conference, PMLR 85:587–602, 2018.

• PMLR: https://proceedings.mlr.press/v85/boag18a.html
• arXiv (optional): https://arxiv.org/abs/1808.03827

The AMA discharge setup in this PR follows MIMIC-III discharge_location and the demographic / administrative feature ablations discussed in that work.

High-level description

This PR adds AMAPredictionMIMIC3, a MIMIC-III task that predicts AMA discharge from administrative and demographic features only (gender + insurance, age, length of stay, optional normalized race, optional substance-use flag from admission diagnosis text). No ICD / prescription sequences are required for the task samples.

It supports three ablation feature sets aligned with the paper’s AMA analysis:

1. BASELINE: demographics, age, los
2. BASELINE+RACE: adds race
3. BASELINE+RACE+SUBSTANCE: adds has_substance_use

Examples

• examples/mimic3_ama_prediction_logistic_regression.py — logistic regression ablation with AUROC, subgroup metrics (race, age band, insurance), and fairness-style reporting (demographic parity, equal opportunity). Includes generate_synthetic_mimic3 for local CSVs: exhaustive cross-product coverage or random stochastic data; CLI flags --data-source (auto / synthetic / real), --synthetic-mode, --root, --patients, --seed, etc.
• examples/mimic3_ama_prediction_rnn.py — same data path and CLI pattern using an RNN; loads the synthetic generator from the logistic example via importlib.

Tests

All AMA-related tests live in tests/core/test_mimic3_ama_prediction.py: helpers / mocks, ablation sanity checks, synthetic MIMIC3Dataset pipeline tests, short training smoke tests, and an exhaustive synthetic row-count check.

File guide (what to review)

Area	Path
Task implementation	pyhealth/tasks/ama_prediction.py
Public export	pyhealth/tasks/__init__.py
API docs	docs/api/tasks.rst, docs/api/tasks/pyhealth.tasks.ama_prediction.rst
Examples	examples/mimic3_ama_prediction_logistic_regression.py, examples/mimic3_ama_prediction_rnn.py
Tests	tests/core/test_mimic3_ama_prediction.py

Contribution guide alignment (CONTRIBUTING.md)

• Unit / integration tests included.
• Documentation updated for the new task in the API docs tree.
• Style: PEP 8 (88 character line length) and Google-style docstrings as described in CONTRIBUTING.md.

How to verify

# Tests
python -m unittest PyHealth.tests.core.test_mimic3_ama_prediction -v

# Example — synthetic (default exhaustive grid when no --root)
python examples/mimic3_ama_prediction_logistic_regression.py --splits 1 --epochs 1

# Example — real MIMIC-III
python examples/mimic3_ama_prediction_logistic_regression.py \
  --data-source real --root /path/to/mimic-iii/1.4