CADRE implementation for PyHealth

[natalie-erj]

Contributors:

Natalie Erjavec (Email: erjavec2@illinois.edu / NetID:erjavec2)
Austin Offenberger (Email: austino3@illinois.edu / NetID:austino3)
Prasanna Kumar Muralimanohar(Email: pkm7@illinois.edu / NetID:pkm7)

Contribution Type:

Full path (Dataset + Task + Model + Paper Replication + Extension)

Overview

A reproduction and extension of CADRE (Contextual Attention-based Drug REsponse) within the PyHealth framework, originally published as:

Tao, Y., Ren, S., Ding, M.Q., Schwartz, R., & Lu, X. (2020). Predicting Drug Sensitivity of Cancer Cell Lines via Collaborative Filtering with Contextual Attention. Proceedings of Machine Learning Research, 126, 456–477. PMLR (MLHC 2020).

Original code: github.com/yifengtao/CADRE

Description

This PR adds support for cancer drug sensitivity prediction to PyHealth, including two datasets, two tasks, and two model architectures. The implementation is based on the CADRE paper (Tao et al., MLHC 2020) with an additional dot-product attention extension.

New Datasets

• GDSCDataset (pyhealth/datasets/gdsc.py)

Wraps the Genomics of Drug Sensitivity in Cancer (GDSC) dataset for drug sensitivity prediction. Each patient is a cancer cell line; the single visit represents its genomic profile.

• CCLEDataset (pyhealth/datasets/ccle.py)

Wraps the Cancer Cell Line Encyclopedia (CCLE) dataset using the same interface as
GDSCDataset.

New Tasks

• DrugSensitivityPredictionGDSC
• DrugSensitivityPredictionCCLE

New Models

CADRE (pyhealth/models/cadre.py)

Re-implementation of Tao et al. (2020). Predicts drug sensitivity via:

1. ExpEncoder — frozen Gene2Vec embeddings + multi-head contextual attention. Each
   drug's target pathway embedding conditions the attention scores, producing a
   drug-specific cell-line representation
2. DrugDecoder — collaborative filtering dot product between cell-line and learned
   drug embeddings
3. Masked BCE loss — computed only on tested (cell line, drug) pairs

The freeze_gene_emb parameter (default True) supports the CADRE pretrain ablation
from the paper by optionally unfreezing the Gene2Vec embeddings during training.

CADREDotAttn (pyhealth/models/cadre_dot_attn.py)

Extension replacing CADRE's additive contextual attention with transformer-style
scaled dot-product attention. Drug embeddings serve as queries against
gene key vectors, so gradients flow through both the prediction path (decoder dot
product) and the attention alignment path (encoder)

Tests

20 unit tests using synthetic in-memory CSV data:

• tests/core/test_gdsc_dataset.py
• tests/core/test_ccle_dataset.py
• tests/core/test_drug_sensitivity_gdsc_task.py
• tests/core/test_cadre_model.py

References

• Tao, Y. et al. (2020). Predicting Drug Sensitivity of Cancer Cell Lines via
  Collaborative Filtering with Contextual Attention. MLHC 2020, PMLR 126:456–477.
• Yang, W. et al. (2013). Genomics of Drug Sensitivity in Cancer (GDSC): a
  resource for therapeutic biomarker discovery in cancer cells. Nucleic Acids
  Research, 41(D1), D955–D961.
• Barretina, J. et al. (2012). The Cancer Cell Line Encyclopedia enables
  predictive modelling of anticancer drug sensitivity. Nature, 483(7391), 603–607.