jiarui_li/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add evaluation scripts.

Browse Source
This commit is contained in:
Jiarui Li
2025-10-18 13:26:56 +08:00
parent 082c719975
commit dbc3000192
4 changed files with 3519 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1271
delphi_labels_chapters_colours_icd.csv Normal file
View File

File diff suppressed because it is too large Load Diff

501
evaluate_auc.py Normal file
View File

@@ -0,0 +1,501 @@
import scipy.stats
import scipy
import warnings
import torch
from model import DelphiConfig, Delphi
from tqdm import tqdm
import pandas as pd
import numpy as np
import argparse
from utils import get_batch, get_p2i
from pathlib import Path
def auc(x1, x2):
n1 = len(x1)
n2 = len(x2)
R1 = np.concatenate([x1, x2]).argsort().argsort()[:n1].sum() + n1
U1 = n1 * n2 + 0.5 * n1 * (n1 + 1) - R1
if n1 == 0 or n2 == 0:
return np.nan
return U1 / n1 / n2
def get_common_diseases(delphi_labels, filter_min_total=100):
chapters_of_interest = [
"I. Infectious Diseases",
"II. Neoplasms",
"III. Blood & Immune Disorders",
"IV. Metabolic Diseases",
"V. Mental Disorders",
"VI. Nervous System Diseases",
"VII. Eye Diseases",
"VIII. Ear Diseases",
"IX. Circulatory Diseases",
"X. Respiratory Diseases",
"XI. Digestive Diseases",
"XII. Skin Diseases",
"XIII. Musculoskeletal Diseases",
"XIV. Genitourinary Diseases",
"XV. Pregnancy & Childbirth",
"XVI. Perinatal Conditions",
"XVII. Congenital Abnormalities",
"Death",
]
labels_df = delphi_labels[
delphi_labels["ICD-10 Chapter (short)"].isin(chapters_of_interest) * (delphi_labels["count"] > filter_min_total)
]
return labels_df["index"].tolist()
def optimized_bootstrapped_auc_gpu(case, control, n_bootstrap=1000):
"""
Computes bootstrapped AUC estimates using PyTorch on CUDA.
Parameters:
case: 1D tensor of scores for positive cases
control: 1D tensor of scores for controls
n_bootstrap: Number of bootstrap replicates
Returns:
Tensor of shape (n_bootstrap,) containing AUC for each bootstrap replicate
"""
if not torch.cuda.is_available():
raise RuntimeError("CUDA is not available. This function requires a GPU.")
# Convert inputs to CUDA tensors
if not torch.is_tensor(case):
case = torch.tensor(case, device="cuda", dtype=torch.float32)
else:
case = case.to("cuda", dtype=torch.float32)
if not torch.is_tensor(control):
control = torch.tensor(control, device="cuda", dtype=torch.float32)
else:
control = control.to("cuda", dtype=torch.float32)
n_case = case.size(0)
n_control = control.size(0)
total = n_case + n_control
# Generate bootstrap samples
boot_idx_case = torch.randint(0, n_case, (n_bootstrap, n_case), device="cuda")
boot_idx_control = torch.randint(0, n_control, (n_bootstrap, n_control), device="cuda")
boot_case = case[boot_idx_case]
boot_control = control[boot_idx_control]
combined = torch.cat([boot_case, boot_control], dim=1)
# Mask to identify case entries
mask = torch.zeros((n_bootstrap, total), dtype=torch.bool, device="cuda")
mask[:, :n_case] = True
# Compute ranks and AUC
ranks = combined.argsort(dim=1).argsort(dim=1)
case_ranks_sum = torch.sum(ranks.float() * mask.float(), dim=1)
min_case_rank_sum = n_case * (n_case - 1) / 2.0
U = case_ranks_sum - min_case_rank_sum
aucs = U / (n_case * n_control)
return aucs.cpu().tolist()
# AUC comparison adapted from
# https://github.com/Netflix/vmaf/
def compute_midrank(x):
"""Computes midranks.
Args:
x - a 1D numpy array
Returns:
array of midranks
"""
J = np.argsort(x)
Z = x[J]
N = len(x)
T = np.zeros(N, dtype=np.float32)
i = 0
while i < N:
j = i
while j < N and Z[j] == Z[i]:
j += 1
T[i:j] = 0.5 * (i + j - 1)
i = j
T2 = np.empty(N, dtype=np.float32)
# Note(kazeevn) +1 is due to Python using 0-based indexing
# instead of 1-based in the AUC formula in the paper
T2[J] = T + 1
return T2
def fastDeLong(predictions_sorted_transposed, label_1_count):
"""
The fast version of DeLong's method for computing the covariance of
unadjusted AUC.
Args:
predictions_sorted_transposed: a 2D numpy.array[n_classifiers, n_examples]
sorted such as the examples with label "1" are first
Returns:
(AUC value, DeLong covariance)
Reference:
@article{sun2014fast,
title={Fast Implementation of DeLong's Algorithm for
Comparing the Areas Under Correlated Receiver Operating Characteristic Curves},
author={Xu Sun and Weichao Xu},
journal={IEEE Signal Processing Letters},
volume={21},
number={11},
pages={1389--1393},
year={2014},
publisher={IEEE}
}
"""
# Short variables are named as they are in the paper
m = label_1_count
n = predictions_sorted_transposed.shape[1] - m
positive_examples = predictions_sorted_transposed[:, :m]
negative_examples = predictions_sorted_transposed[:, m:]
k = predictions_sorted_transposed.shape[0]
tx = np.empty([k, m], dtype=np.float32)
ty = np.empty([k, n], dtype=np.float32)
tz = np.empty([k, m + n], dtype=np.float32)
for r in range(k):
tx[r, :] = compute_midrank(positive_examples[r, :])
ty[r, :] = compute_midrank(negative_examples[r, :])
tz[r, :] = compute_midrank(predictions_sorted_transposed[r, :])
aucs = tz[:, :m].sum(axis=1) / m / n - float(m + 1.0) / 2.0 / n
v01 = (tz[:, :m] - tx[:, :]) / n
v10 = 1.0 - (tz[:, m:] - ty[:, :]) / m
sx = np.cov(v01)
sy = np.cov(v10)
delongcov = sx / m + sy / n
return aucs, delongcov
def compute_ground_truth_statistics(ground_truth):
assert np.array_equal(np.unique(ground_truth), [0, 1])
order = (-ground_truth).argsort()
label_1_count = int(ground_truth.sum())
return order, label_1_count
def get_auc_delong_var(healthy_scores, diseased_scores):
"""
Computes ROC AUC value and variance using DeLong's method
Args:
healthy_scores: Values for class 0 (healthy/controls)
diseased_scores: Values for class 1 (diseased/cases)
Returns:
AUC value and variance
"""
# Create ground truth labels (1 for diseased, 0 for healthy)
ground_truth = np.array([1] * len(diseased_scores) + [0] * len(healthy_scores))
predictions = np.concatenate([diseased_scores, healthy_scores])
# Compute statistics needed for DeLong method
order, label_1_count = compute_ground_truth_statistics(ground_truth)
predictions_sorted_transposed = predictions[np.newaxis, order]
# Calculate AUC and covariance
aucs, delongcov = fastDeLong(predictions_sorted_transposed, label_1_count)
assert len(aucs) == 1, "There is a bug in the code, please forward this to the developers"
return aucs[0], delongcov
def get_calibration_auc(j, k, d, p, offset=365.25, age_groups=range(45, 80, 5), precomputed_idx=None, n_bootstrap=1, use_delong=False):
age_step = age_groups[1] - age_groups[0]
# Indexes of cases with disease k
wk = np.where(d[2] == k)
if len(wk[0]) < 2:
return None
# For controls, we need to exclude cases with disease k
wc = np.where((d[2] != k) * (~(d[2] == k).any(-1))[..., None])
wall = (np.concatenate([wk[0], wc[0]]), np.concatenate([wk[1], wc[1]])) # All cases and controls
# We need to take into account the offset t and use the tokens for prediction that are at least t before the event
if precomputed_idx is None:
pred_idx = (d[1][wall[0]] <= d[3][wall].reshape(-1, 1) - offset).sum(1) - 1
else:
pred_idx = precomputed_idx[wall] # It's actually much faster to precompute this
z = d[1][(wall[0], pred_idx)] # Times of the tokens for prediction
z = z[pred_idx != -1]
zk = d[3][wall] # Target times
zk = zk[pred_idx != -1]
# x = np.exp(p[..., j][(wall[0], pred_idx)]) * 365.25
# x = 1 - np.exp(-x * age_step) # the function is monotinic, so we don't need to do this for the AUC
x = p[..., j][(wall[0], pred_idx)]
x = x[pred_idx != -1]
wk = (wk[0][pred_idx[: len(wk[0])] != -1], wk[1][pred_idx[: len(wk[0])] != -1])
p_idx = wall[0][pred_idx != -1]
out = []
for i, aa in enumerate(age_groups):
a = np.logical_and(z / 365.25 >= aa, z / 365.25 < aa + age_step)
# Optionally, add extra filtering on the time difference, for example:
# a *= (zk - z < 365.25)
selected_groups = p_idx[a]
perm = np.random.permutation(len(selected_groups))
_, indices = np.unique(selected_groups[perm], return_index=True)
indices = perm[indices]
selected = np.zeros(np.sum(a), dtype=bool)
selected[indices] = True
a[a] = selected
control = x[len(wk[0]) :][a[len(wk[0]) :]]
case = x[: len(wk[0])][a[: len(wk[0])]]
if len(control) == 0 or len(case) == 0:
continue
if use_delong:
auc_value_delong, auc_variance_delong = get_auc_delong_var(control, case)
auc_delong_dict = {"auc_delong": auc_value_delong, "auc_variance_delong": auc_variance_delong}
else:
auc_delong_dict = {}
if n_bootstrap > 1:
aucs_bootstrapped = optimized_bootstrapped_auc_gpu(case, control, n_bootstrap)
for bootstrap_idx in range(n_bootstrap):
y = auc_value_delong if n_bootstrap == 1 else aucs_bootstrapped[bootstrap_idx]
out_item = {
"token": k,
"auc": y,
"age": aa,
"n_healthy": len(control),
"n_diseased": len(case),
}
out.append(out_item | auc_delong_dict)
if n_bootstrap > 1:
out_item["bootstrap_idx"] = bootstrap_idx
return out
# New internal function that performs the AUC evaluation pipeline.
def evaluate_auc_pipeline(
model,
d100k,
output_path,
delphi_labels,
diseases_of_interest=None,
filter_min_total=100,
disease_chunk_size=200,
age_groups=np.arange(40, 80, 5),
offset=0.1,
batch_size=128,
device="cpu",
seed=1337,
n_bootstrap=1,
meta_info={},
):
"""
Runs the AUC evaluation pipeline.
Args:
model (torch.nn.Module): The loaded model set to eval().
d100k (tuple): Data batch from get_batch.
delphi_labels (pd.DataFrame): DataFrame with label info (token names, etc. "delphi_labels_chapters_colours_icd.csv").
output_path (str | None): Directory where CSV files will be written. If None, files will not be saved.
diseases_of_interest (np.ndarray or list, optional): If provided, these disease indices are used.
filter_min_total (int): Minimum total token count to include a token.
disease_chunk_size (int): Maximum chunk size for processing diseases.
age_groups (np.ndarray): Age groups to use in calibration.
offset (float): Offset used in get_calibration_auc.
batch_size (int): Batch size for model forwarding.
device (str): Device identifier.
seed (int): Random seed for reproducibility.
n_bootstrap (int): Number of bootstrap samples. (1 for no bootstrap)
Returns:
tuple: (df_auc_unpooled, df_auc, df_both) DataFrames.
"""
assert n_bootstrap > 0, "n_bootstrap must be greater than 0"
# Set random seeds
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
# Get common diseases
if diseases_of_interest is None:
diseases_of_interest = get_common_diseases(delphi_labels, filter_min_total)
# Split diseases into chunks for processing
num_chunks = (len(diseases_of_interest) + disease_chunk_size - 1) // disease_chunk_size
diseases_chunks = np.array_split(diseases_of_interest, num_chunks)
# Precompute prediction indices for calibration
pred_idx_precompute = (d100k[1][:, :, np.newaxis] < d100k[3][:, np.newaxis, :] - offset).sum(1) - 1
all_aucs = []
tqdm_options = {"desc": "Processing disease chunks", "total": len(diseases_chunks)}
for disease_chunk_idx, diseases_chunk in tqdm(enumerate(diseases_chunks), **tqdm_options):
p100k = []
model.to(device)
with torch.no_grad():
# Process the evaluation data in batches
for dd in tqdm(
zip(*[torch.split(x, batch_size) for x in d100k]),
desc=f"Model inference, chunk {disease_chunk_idx}",
total=d100k[0].shape[0] // batch_size + 1,
):
dd = [x.to(device) for x in dd]
outputs = model(*dd)[0].cpu().detach().numpy()
# Keep only the columns corresponding to the current disease chunk
p100k.append(outputs[:, :, diseases_chunk].astype("float16")) # enough to store logits, but not rates
p100k = np.vstack(p100k)
# Loop over each disease (token) in the current chunk, sexes separately
for sex, sex_idx in [("female", 2), ("male", 3)]:
sex_mask = ((d100k[0] == sex_idx).sum(1) > 0).cpu().detach().numpy()
p_sex = p100k[sex_mask]
d100k_sex = [d_[sex_mask].cpu().detach().numpy() for d_ in d100k]
precomputed_idx_subset = pred_idx_precompute[sex_mask].cpu().detach().numpy()
for j, k in tqdm(
list(enumerate(diseases_chunk)), desc=f"Processing diseases in chunk {disease_chunk_idx}, {sex}"
):
# Get calibration AUC for the current disease token.
out = get_calibration_auc(
j,
k,
d100k_sex,
p_sex,
age_groups=age_groups,
offset=offset,
precomputed_idx=precomputed_idx_subset,
n_bootstrap=n_bootstrap,
use_delong=True,
)
if out is None:
# print(f"No data for disease {k} and sex {sex}")
continue
for out_item in out:
out_item["sex"] = sex
all_aucs.append(out_item)
df_auc_unpooled = pd.DataFrame(all_aucs)
for key, value in meta_info.items():
df_auc_unpooled[key] = value
delphi_labels_subset = delphi_labels[['index', 'ICD-10 Chapter (short)', 'name', 'color', 'count']]
df_auc_unpooled_merged = df_auc_unpooled.merge(delphi_labels_subset, left_on="token", right_on="index", how="inner")
def aggregate_age_brackets_delong(group):
# For normal distributions, when averaging n of them:
# The variance of the sum is the sum of variances
# The variance of the average is the sum of variances divided by n^2
n = len(group)
mean = group['auc_delong'].mean()
# Since we're taking the average, divide combined variance by n^2
var = group['auc_variance_delong'].sum() / (n**2)
return pd.Series({
'auc': mean,
'auc_variance_delong': var,
'n_samples': n,
'n_diseased': group['n_diseased'].sum(),
'n_healthy': group['n_healthy'].sum(),
})
print('Using DeLong method to calculate AUC confidence intervals..')
df_auc = df_auc_unpooled.groupby(["token"]).apply(aggregate_age_brackets_delong).reset_index()
df_auc_merged = df_auc.merge(delphi_labels, left_on="token", right_on="index", how="inner")
if output_path is not None:
Path(output_path).mkdir(exist_ok=True, parents=True)
df_auc_merged.to_parquet(f"{output_path}/df_both.parquet", index=False)
df_auc_unpooled_merged.to_parquet(f"{output_path}/df_auc_unpooled.parquet", index=False)
return df_auc_unpooled_merged, df_auc_merged
def main():
parser = argparse.ArgumentParser(description="Evaluate AUC")
parser.add_argument("--input_path", type=str, help="Path to the dataset")
parser.add_argument("--output_path", type=str, help="Path to the output")
parser.add_argument("--model_ckpt_path", type=str, help="Path to the model weights")
parser.add_argument("--no_event_token_rate", type=int, help="No event token rate")
parser.add_argument(
"--health_token_replacement_prob", default=0.0, type=float, help="Health token replacement probability"
)
parser.add_argument("--dataset_subset_size", type=int, default=-1, help="Dataset subset size for evaluation")
parser.add_argument("--n_bootstrap", type=int, default=1, help="Number of bootstrap samples")
# Optional filtering/chunking parameters:
parser.add_argument("--filter_min_total", type=int, default=100, help="Minimum total count to filter tokens")
parser.add_argument("--disease_chunk_size", type=int, default=200, help="Chunk size for processing diseases")
args = parser.parse_args()
input_path = args.input_path
output_path = args.output_path
no_event_token_rate = args.no_event_token_rate
health_token_replacement_prob = args.health_token_replacement_prob
dataset_subset_size = args.dataset_subset_size
# Create output folder if it doesn't exist.
Path(output_path).mkdir(exist_ok=True, parents=True)
device = "cuda"
seed = 1337
# Load model checkpoint and initialize model.
ckpt_path = args.model_ckpt_path
checkpoint = torch.load(ckpt_path, map_location=device)
conf = DelphiConfig(**checkpoint["model_args"])
model = Delphi(conf)
state_dict = checkpoint["model"]
model.load_state_dict(state_dict)
model.eval()
model = model.to(device)
# Load training and validation data.
val = np.fromfile(f"{input_path}/val.bin", dtype=np.uint32).reshape(-1, 3).astype(np.int64)
val_p2i = get_p2i(val)
if dataset_subset_size == -1:
dataset_subset_size = len(val_p2i)
# Get a subset batch for evaluation.
d100k = get_batch(
range(dataset_subset_size),
val,
val_p2i,
select="left",
block_size=80,
device=device,
padding="random",
no_event_token_rate=no_event_token_rate,
health_token_replacement_prob=health_token_replacement_prob,
)
# Load labels (external) to be passed in.
delphi_labels = pd.read_csv("delphi_labels_chapters_colours_icd.csv")
# Call the internal evaluation function.
df_auc_unpooled, df_auc_merged = evaluate_auc_pipeline(
model,
d100k,
output_path,
delphi_labels,
diseases_of_interest=None,
filter_min_total=args.filter_min_total,
disease_chunk_size=args.disease_chunk_size,
device=device,
seed=seed,
n_bootstrap=args.n_bootstrap,
)
if __name__ == "__main__":
main()

1118
evaluate_delphi.py Normal file
View File

File diff suppressed because it is too large Load Diff

629
shap_analysis.py Normal file
View File

@@ -0,0 +1,629 @@
#!/usr/bin/env python
# coding: utf-8
# # Looking inside Delphi with SHAP values
#
# Welcome to the Delphi SHAP notebook!
#
# Delphi is a generative autoregressive model that not only predicts the future disease rates, but also sample entire disease trajectories one step at a time.
#
# In this notebook, we will use SHAP (SHapley Additive exPlanations) framework to analyse which interaction between diseases that Delphi learned from data and how these interaction influence its predicitons.
#
# Let's start by looking at what SHAP values mean:
#
# ## SHAP Values and Delphi
#
# SHAP (SHapley Additive exPlanations) values help us understand how a machine learning model makes its predictions by showing the contribution of each input feature.
#
#
# ### Example: Patient Trajectory
#
# Consider a simplified patient trajectory:
# `Male, Migraine, Common cold, Brain cancer`
#
# For this trajectory, Delphi would predict a very high mortality risk (aka high rate for the Death token being next). Say, 95% chance of death within a year. Why? Technically, we don't know, since neural networks are black boxes.
#
# Let's try masking several tokens and predicting the next token again.
#
# `Male, [Masked: Migraine], Common cold, Brain cancer` -> 95% chance of death within a year, no change
#
# `Male, Migraine, Common cold, [Masked: Brain cancer]` -> 5% chance of death within a year, risk drops significantly
#
# Without speaking about causality, we can assume that there is *some* connection between brain cancer and death risk. SHAP framework allows using such masking to systematically assess the contribution of each token to the prediction. We can perform this analysis for all trajectories in the dataset and evaluate how, on average, a given disease influences the risk of any other disease.
#
# In case of Delphi, masking means replacing a disease token with "no event" token for all input tokens, except for the sex token that is inverted.
# In[2]:
import os
import pickle
import torch
from model import DelphiConfig, Delphi
from tqdm import tqdm
import pandas as pd
import numpy as np
import textwrap
import warnings
import matplotlib.pyplot as plt
get_ipython().run_line_magic('config', "InlineBackend.figure_format='retina'")
plt.rcParams['figure.facecolor'] = 'white'
plt.rcParams.update({'axes.grid': True,
'grid.linestyle': ':',
'axes.spines.bottom': False,
'axes.spines.left': False,
'axes.spines.right': False,
'axes.spines.top': False})
plt.rcParams['figure.dpi']= 72
plt.rcParams['pdf.fonttype'] = 42
delphi_labels = pd.read_csv('delphi_labels_chapters_colours_icd.csv')
# ## Load model
# In[3]:
out_dir = 'Delphi-2M'
device = 'cuda' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1', etc.
dtype ='float32' #'bfloat16' # 'float32' or 'bfloat16' or 'float16'
seed = 1337
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
device_type = 'cuda' if 'cuda' in device else 'cpu'
dtype = {'float32': torch.float32, 'float64': torch.float64, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
# In[4]:
ckpt_path = os.path.join(out_dir, 'ckpt.pt')
checkpoint = torch.load(ckpt_path, map_location=device)
conf = DelphiConfig(**checkpoint['model_args'])
model = Delphi(conf)
state_dict = checkpoint['model']
model.load_state_dict(state_dict)
model.eval()
model = model.to(device)
# ## Load data
# In[5]:
from utils import get_batch, get_p2i
train = np.fromfile('data/ukb_simulated_data/train.bin', dtype=np.uint32).reshape(-1,3)
val = np.fromfile('data/ukb_simulated_data/val.bin', dtype=np.uint32).reshape(-1,3)
train_p2i = get_p2i(train)
val_p2i = get_p2i(val)
# In[6]:
# define a random example health trajectory
person = [('Male', 0),
('B01 Varicella [chickenpox]',2),
('L20 Atopic dermatitis',3),
('Healthy', 5),
('Healthy', 10),
('Healthy', 15),
('Healthy', 20),
('G43 Migraine', 20),
('E73 Lactose intolerance', 21),
('B27 Infectious mononucleosis', 22),
('Healthy', 25),
('J11 Influenza, virus not identified', 28),
('Healthy', 30),
('Healthy', 35),
('C25 Malignant neoplasm of pancreas', 38),
('Healthy', 40),
('Smoking low', 41),
('BMI mid', 41),
('Alcohol high', 41),
('Healthy', 42),
]
person = [(a, b * 365.25) for a,b in person]
# ### Individual SHAP values
# In[7]:
# define helper functions
id_to_token = delphi_labels['name'].to_dict()
token_to_id = {v:k for k, v in id_to_token.items()}
def tokens_to_ids(tokens):
return [token_to_id[t] for t in tokens]
def ids_to_tokens(ids):
return [id_to_token[int(id_)] for id_ in ids]
def split_person(p):
tokens = [i[0] for i in p]
ages = [i[1] for i in p]
return tokens, ages
def get_person(idx):
x, y, _, time = get_batch([idx], val, val_p2i,
select='left', block_size=48,
device=device, padding='random')
x, y = x[y > -1], y[y > -1]
person = []
for token_id, date in zip(x, y):
person.append((id_to_token[token_id.item()], date.item()))
return person, y, time[0][-1]
# In[8]:
from utils import shap_custom_tokenizer, shap_model_creator
import shap
# person_to_process = get_person(137)[0]
person_to_process = person
diseases_of_interest = [1269, 46, 95, 1168, 374, 173, 214, 305, 505, 584]
person_tokens, person_ages = split_person(person_to_process)
person_tokens_ids = tokens_to_ids(person_tokens)
masker = shap.maskers.Text(shap_custom_tokenizer, output_type='str', mask_token='10000', collapse_mask_token=False)
model_shap = shap_model_creator(model, diseases_of_interest, person_tokens_ids, person_ages, device)
explainer = shap.Explainer(model_shap, masker, output_names=delphi_labels['name'].values[diseases_of_interest])
shap_values = explainer([' '.join(list(map(lambda x: str(token_to_id[x]), person_tokens)))])
shap_values.data = np.array([list(map(lambda x: f"{delphi_labels['name'].values[token_to_id[x[0]]]}({x[1]/365:.1f} years) ", person_to_process))])
# In[9]:
out = shap.plots.text(shap_values, display=True) # sometimes this interactive plot can't be rendered well (eg in VS Code, feel free to skip it)
# In[10]:
# SHAP values can be interpreted as how much each input token changes predicted logit corresponding to a particular disease.
# As Delphi logits are log-disease rates, we can convent SHAP values to the disease-specific fold risk changes.
# Shown below is a waterfall plot, showing SHAP values for the most "influential" diseases within
# a single trajectory.
from plotting import waterfall
with plt.style.context('default'):
plt.rcParams['pdf.fonttype'] = 42
plt.rcParams['figure.dpi'] = 150
plt.rcParams['font.size'] = 4
waterfall(shap_values[0, ..., 0], max_display=7, show=False, ages=person_ages)
plt.gca().set_title('Impact of diseases on mortality', fontweight=1, size=18)
plt.show()
# ## Pre-computed many cases
#
# The small synthetic dataset is not enough to properly run following part; if you have access to the full dataset, run `shap-agg-eval.py` to evaluate SHAP values for the entire dataset.
# In[11]:
import pickle
with open('shap_agg.pickle', 'rb') as f:
shap_pkl = pickle.load(f)
all_tokens = shap_pkl['tokens']
all_values = shap_pkl['values']
# In[12]:
import pandas as pd
import seaborn as sns
import numpy as np
import matplotlib.pyplot as plt
df_shap = pd.DataFrame(all_values)
df_shap['token'] = all_tokens.astype('int')
# In[13]:
token_count_dict = df_shap['token'].value_counts().sort_index().to_dict()
N_min = 5 # we will only consider diseases that have at least 5 calculated SHAP values, otherwise they are too noisy
columns_more_N = [c for c in df_shap.columns if c == 1269 or (c in token_count_dict and token_count_dict[c] >= N_min)]
df_shap_agg = df_shap[df_shap['token'].apply(lambda x: token_count_dict[x] > N_min)].groupby('token').mean()
# Since we now have calculated SHAP values for the entire dataset, we can use them to analyse "connections" between diseases.
#
# For every "predicted"-"predictor" pair, we average all SHAP values for the given pair.
#
# We can further analyse them in sevaral directions:
# - For a given disease in the past medical history, which disease rates are most increased by it?
# - For a given potential future disease, having which disease in the past medical history would most increase its rate?
#
# Let's see which diseases increase disease risk the most and also which diseases are most influenced by being a heavy smoker.
# In[14]:
import matplotlib.pyplot as plt
import numpy as np
def plot_shap_distribution(df_melted, y_axis_labels, group_by_col_name, title, x_lim_tuple, highlight_last_dot=True):
"""
Generates a plot showing median and quartile ranges of SHAP values for different tokens.
"""
plt.figure(figsize=(3, 6), facecolor='w')
ax = plt.gca()
for i, label_for_y_tick in enumerate(y_axis_labels):
data_for_label = df_melted[df_melted[group_by_col_name] == label_for_y_tick]['value']
if not data_for_label.empty:
median = np.median(data_for_label)
quartiles = np.percentile(data_for_label, [25, 75])
dot_color = 'red' if highlight_last_dot and i == len(y_axis_labels) - 1 else 'black'
ax.plot(median, i, 'o', color=dot_color, zorder=3)
ax.hlines(i, quartiles[0], quartiles[1], color='gray', linestyles='solid', linewidth=1)
plt.title(title)
plt.yticks(range(len(y_axis_labels)), y_axis_labels)
plt.xticks(rotation=25, ha='right')
plt.xscale('log')
plt.xlim(*x_lim_tuple)
plt.xlabel('Risk increase, folds', size=11, labelpad=10)
plt.show()
target_token = 1269
n_first = 20
plot = True
selected_context_tokens1 = df_shap_agg[target_token].nlargest(n_first).index[::-1]
df_plot_source = df_shap[df_shap['token'].isin(selected_context_tokens1)]
df_plot_melted = df_plot_source[[target_token, 'token']].reset_index(drop=True).melt(id_vars=['token'], value_vars=[target_token])
df_plot_melted['context_token_label'] = df_plot_melted['token'].map(id_to_token)
df_plot_melted['value'] = np.exp(df_plot_melted['value'])
y_axis_labels1 = [id_to_token[token] for token in selected_context_tokens1]
title1 = 'Mortality factors'
xlim1 = (1, 1000)
plot_shap_distribution(
df_melted=df_plot_melted,
y_axis_labels=y_axis_labels1,
group_by_col_name='context_token_label',
title=title1,
x_lim_tuple=xlim1
)
# In[15]:
target_token = 9
n_first = 20
shap_values_for_context = df_shap_agg.loc[target_token]
selected_feature_tokens2 = shap_values_for_context.sort_values(ascending=False).index[:n_first][::-1]
df_plot_source = df_shap[df_shap['token'] == target_token]
df_plot_melted = df_plot_source[[*selected_feature_tokens2, 'token']].reset_index(drop=True).melt(
id_vars=['token'],
value_vars=selected_feature_tokens2,
var_name='feature_token_id',
value_name='raw_shap_value'
)
df_plot_melted['feature_label'] = df_plot_melted['feature_token_id'].map(id_to_token)
df_plot_melted['value'] = np.exp(df_plot_melted['raw_shap_value'])
y_axis_labels = [id_to_token[token] for token in selected_feature_tokens2]
title = 'Consequences of\nsmoking heavily'
xlim = (1, 11)
plot_shap_distribution(
df_melted=df_plot_melted,
y_axis_labels=y_axis_labels,
group_by_col_name='feature_label',
title=title,
x_lim_tuple=xlim
)
# ### Time-resolved SHAP analysis
#
# Before, we aggregated the calculated SHAP values in a fairly simple way: just averaged them within all "predictor-predicted" pairs. This is an oversimplification, since the context in which these two diseases occur is also important.
#
# For instance, the amount of time passed since the "predictor" disease occured is important, since some acute conditions may have vastly different effects compared to their chronic forms.
#
# Now, we will aggregate SHAP values within the pairs, additionally separating them by the time between the "predictor" and "predicted" diseases.
# In[16]:
d = get_batch(range(len(np.unique(shap_pkl['people']))), val, val_p2i,
select='left', block_size=48,
device='cpu', padding='regular')
# In[17]:
has_gender = torch.isin(d[0], torch.tensor([2, 3])).any(dim=1).numpy()
is_male = torch.isin(d[0], torch.tensor([3])).any(dim=1).numpy()
is_female = torch.isin(d[0], torch.tensor([2])).any(dim=1).numpy()
# In[18]:
def get_person(idx):
x, y, _, time = get_batch([idx], val, val_p2i,
select='left', block_size=64,
device=device, padding='random',
cut_batch=True)
x, y = x[y > -1], y[y > -1]
person = []
for token_id, date in zip(x, y):
person.append((id_to_token[token_id.item()], date.item()))
return person, y, time[0][-1]
# In[19]:
# the shap result pickle does not contain time, so we need to add it
persons_lengths = []
ages = []
reg_times = []
for p in tqdm(np.unique(shap_pkl['people'])):
pers = get_person(p)
reg_time_idx = np.where(np.isin(tokens_to_ids(np.array(pers[0])[:, 0]), np.arange(4, 13)))[0]
if len(reg_time_idx) > 0:
reg_time = pers[0][reg_time_idx[0]][1]
else:
reg_time = -1
reg_times += [reg_time] * len(pers[0])
persons_lengths += [p] * len(pers[0])
ages += [pers[-1].item()] * len(pers[0])
assert len(ages) == len(df_shap)
# In[20]:
all_tokens = shap_pkl['tokens']
all_values = shap_pkl['values']
all_times = shap_pkl['times']
df_shap = pd.DataFrame(all_values)
df_shap['token'] = all_tokens
df_shap['time'] = all_times
df_shap['person'] = shap_pkl['people']
df_shap['age'] = np.array(ages) / 365.25
df_shap['reg_time_years'] = np.array(reg_times) / 365.25
df_shap['Time, years'] = df_shap['time'] / 365.25
df_shap['age_at_token'] = df_shap['age'] - df_shap['time'] / 365.25
df_shap = df_shap[df_shap['reg_time_years'] > 0]
token_count_dict = df_shap['token'].value_counts().sort_index().to_dict()
# In[22]:
import numpy as np
def bins_avg(x, y, grid_size=3):
'''Filter out regions wiht few data points'''
x, y = np.array(x), np.array(y)
bin_edges = np.arange(np.min(x), np.max(x), grid_size)
bin_indices = np.digitize(x, bin_edges)
bin_avgs = np.array([y[bin_indices == i].mean() for i in range(1, len(bin_edges)+1)])
return bin_edges, bin_avgs
# In[23]:
tokens_of_interest = [46, 95, 1168, 1188, 173, 214, 305, 505, 584]
n_groups = len(tokens_of_interest) // 5 + 1
palette_faint = [sns.color_palette("Paired")[0], sns.color_palette("Paired")[2], sns.color_palette("Paired")[4]]
palette_bright = [sns.color_palette("Paired")[1], sns.color_palette("Paired")[3], sns.color_palette("Paired")[5]]
for num_g, token_group in enumerate(np.array_split(tokens_of_interest, n_groups)):
fig, axs = plt.subplots(1, 5, figsize=(12, 2), sharey=True)
for num, (ax, token_id) in enumerate(zip(axs.flatten(), token_group)):
df_trait = df_shap[df_shap['token'] == token_id].copy()
df_trait[1269] = np.exp(df_trait[1269].values)
df_trait['Time, years'] = df_trait['time'] / 365.25
df_trait = df_trait.head(2000)
if len(df_trait) < 2:
continue
sns.scatterplot(data=df_trait, x='Time, years', y=1269, ax=ax, color=palette_faint[0], alpha=0.7, rasterized=True)
x, y = df_trait['Time, years'], df_trait[1269]
n = 3
with warnings.catch_warnings():
warnings.simplefilter('ignore')
x, y = bins_avg(x, y, grid_size=n)
ax.plot(x, y, color=palette_bright[0], linewidth=1.5)
ax.set_ylim(0.5, 500)
ax.set_xlim(0.1, 10)
ax.set_ylabel('Impact on mortality');
ax.set_title(textwrap.fill(id_to_token[token_id], width=15) + f' {token_id}', size=9)
# ax.set_xscale('log')
ax.set_yscale('log')
plt.show()
# As shown is the graph above, some diseases (pancteatic cancer, miocardial infarction, septiceamia) have a much higher impact on mortality if they occur in the recent past, while others (diabetis, depression) don't have a clear time-dependence.
# ## Interaction heatmap
#
# To analyse the interactions between diseases more systematically, we can plot a heatmap of the SHAP values for all "predictor-predicted" pairs, sorted by ICD-10 chapter.
#
# Let's plot two separate heatmaps, one for the cases where the "predictor" disease occured in the past 5 years (with the "predicted disease being the reference) and one for the cases where it occured more than 10 years ago.
# In[24]:
N_min = 5
token_count_dict_below_5y = df_shap[df_shap['Time, years'] < 5]['token'].value_counts().sort_index().to_dict()
token_count_dict_over_10y = df_shap[df_shap['Time, years'] > 10]['token'].value_counts().sort_index().to_dict()
for d in [token_count_dict_below_5y, token_count_dict_over_10y]:
for i in range(1300):
if i not in d:
d[i] = 0
columns_more_N = [c for c in df_shap.columns if c == 1269 or(c in token_count_dict_below_5y and token_count_dict_below_5y[c] >= N_min and
c in token_count_dict_over_10y and token_count_dict_over_10y[c] >= N_min)]
df_shap_agg_below_5y = df_shap[df_shap['token'].apply(lambda x: x in columns_more_N) & (df_shap['Time, years'] < 5)].groupby('token').mean()[columns_more_N]
df_shap_agg_over_10y = df_shap[df_shap['token'].apply(lambda x: x in columns_more_N) & (df_shap['Time, years'] > 10)].groupby('token').mean()[columns_more_N]
# In[25]:
from matplotlib.colors import LogNorm
to_exclude_predicted = ['Technical', 'Smoking, Alcohol and BMI', 'Sex', 'XVI. Perinatal Conditions']
to_exclude_predictor = ['Technical', 'Smoking, Alcohol and BMI', 'Sex', 'XVI. Perinatal Conditions', 'Death']
chapter_order = ['I. Infectious Diseases', 'II. Neoplasms',
'III. Blood & Immune Disorders', 'IV. Metabolic Diseases',
'V. Mental Disorders', 'VI. Nervous System Diseases',
'VII. Eye Diseases', 'VIII. Ear Diseases', 'IX. Circulatory Diseases', 'X. Respiratory Diseases',
'XI. Digestive Diseases', 'XII. Skin Diseases',
'XIII. Musculoskeletal Diseases', 'XIV. Genitourinary Diseases',
'XV. Pregnancy & Childbirth', 'XVI. Perinatal Conditions',
'XVII. Congenital Abnormalities', 'Death']
def get_tick_coords(arr):
return np.where(arr[1:] != arr[:-1])[0]
def plot_full_shap_heatmap(cur_df, title):
new_death_rows = 10
for c in range(1269+1, 1269+new_death_rows+1):
cur_df[c] = cur_df[1269]
delphi_labels = pd.read_csv("delphi_labels_chapters_colours_icd.csv", index_col=0)
death_df = delphi_labels[delphi_labels['ICD-10 Chapter (short)']=="Death"].sample(new_death_rows, replace=True)
death_df['name'] = death_df['name'].apply(lambda x: x + str(np.random.randint(0, 100000)))
death_df.index = pd.Index(range(1269+1, 1269+new_death_rows+1))
delphi_labels = pd.concat([delphi_labels, death_df])
to_exclude_predicted_idx = delphi_labels[~delphi_labels['ICD-10 Chapter (short)'].isin(to_exclude_predicted)].index
to_exclude_predictor_idx = delphi_labels[~delphi_labels['ICD-10 Chapter (short)'].isin(to_exclude_predictor)].index
to_exclude_predicted_idx = to_exclude_predicted_idx[to_exclude_predicted_idx.isin(cur_df.columns)]
to_exclude_predicted_idx = sorted(to_exclude_predicted_idx, key=lambda x: (chapter_order.index(delphi_labels.loc[x, 'ICD-10 Chapter (short)']), x))
to_exclude_predicted_idx = pd.Index(to_exclude_predicted_idx)
to_exclude_predictor_idx = to_exclude_predictor_idx[to_exclude_predictor_idx.isin(cur_df.index.values)]
to_exclude_predictor_idx = sorted(to_exclude_predictor_idx, key=lambda x: (chapter_order.index(delphi_labels.loc[x, 'ICD-10 Chapter (short)']), x))
to_exclude_predictor_idx = pd.Index(to_exclude_predictor_idx)
cur_df = cur_df.loc[to_exclude_predictor_idx, to_exclude_predicted_idx]
row_colors = delphi_labels.iloc[cur_df.index.values]['color'].to_numpy()
col_colors = delphi_labels.iloc[cur_df.columns]['color'].to_numpy()
y_tick_coords = get_tick_coords(delphi_labels.iloc[cur_df.index.values]['color'].to_numpy())
x_tick_coords = get_tick_coords(delphi_labels.iloc[cur_df.columns]['color'].to_numpy())
g = sns.clustermap(np.exp(cur_df.values), row_cluster=False, col_cluster=False,
row_colors=row_colors, col_colors=col_colors,
# norm=LogNorm(vmin=5e-2, vmax=2e1),
norm=LogNorm(vmin=1e-1, vmax=1e1),
cmap='RdBu_r',
figsize=(8.5, 8.5),
rasterized=True,
)
g.ax_heatmap.set_xticks(x_tick_coords)
g.ax_heatmap.set_yticks(y_tick_coords)
g.ax_heatmap.tick_params(length=0, width=0.5, labelsize=8, grid_alpha=0.6, grid_linewidth=0.35, grid_color='gray')
g.ax_cbar.tick_params(length=0.5, width=0.6, labelsize=8, grid_alpha=0.45, grid_linewidth=0.45)
for ch, color in delphi_labels[['ICD-10 Chapter (short)', 'color']].drop_duplicates('color').values:
col_loc = np.where(col_colors == color)[0].mean() if (col_colors == color).any() else np.nan
g.ax_heatmap.text(col_loc - 10, -60, ch, va='bottom', rotation=90, ha='left', size=8)
row_loc = np.where(row_colors == color)[0].mean() if (row_colors == color).any() else np.nan
g.ax_heatmap.text(-70, row_loc, ch, va='center', ha='right', size=9)
from matplotlib.patches import Patch
# Create legend for chapter colors
chapter_color_map = delphi_labels[['ICD-10 Chapter (short)', 'color']].drop_duplicates('color')
chapter_color_map = chapter_color_map[~chapter_color_map['ICD-10 Chapter (short)'].isin(to_exclude_predicted)]
handles = [Patch(facecolor=color) for color in chapter_color_map['color']]
plt.legend(handles, chapter_color_map['ICD-10 Chapter (short)'], title='ICD-10 Chapters',
bbox_transform=plt.gcf().transFigure, loc='center left', bbox_to_anchor=(.96, 0.5))
plt.suptitle(title, y=1.1, size=10, x=0.5)
plt.show()
# In[26]:
plot_full_shap_heatmap(df_shap_agg_below_5y, 'Influence of tokens from below 5 years,\nrisk increase, folds')
# In[27]:
plot_full_shap_heatmap(df_shap_agg_over_10y, 'Influence of tokens from above 10 years,\nrisk increase, folds')
# Interestingly, the resulting heatmap has a block diagonal structure, meaning that within a chapter the interactions between diseases tend to be stronger than between chapters.
#
# The second ("above 10 years") heatmap is also more pale, meaning that most of the disease-disease interactions get weaker over time.