shap_plot_protein.py

import torch
import shap
import os
import matplotlib.pyplot as plt
import numpy as np
from ukb_clomics.eval.survival import CoxPL
from ukb_clomics.eval.cox_model import load_single_disease_data
from ukb_clomics.eval.cox_model import Survivaldata
import pandas as pd
import pytorch_lightning as pl

pl.seed_everything(42)


class RiskModel(torch.nn.Module):
    def __init__(self, cox_model):
        super().__init__()
        self.cox_model = cox_model
        self.cox_model.eval()

    def forward(self, x):
        log_hazard = self.cox_model.model(x).squeeze(-1)
        return log_hazard.unsqueeze(-1)


def load_indices(path_list):
    indices = set()
    for path in path_list:
        df_indices = pd.read_csv(path, usecols=[0])
        indices.update(df_indices.iloc[:, 0].astype(int).tolist())
    return indices


def load_features(cfg):
    if isinstance(cfg.train_path, list):
        X_train = pd.concat([pd.read_csv(path, index_col=0) for path in cfg.train_path], axis=0)
    else:
        raise ValueError('train_path must be a list of paths')

    if isinstance(cfg.test_path, list):
        X_test = pd.concat([pd.read_csv(path, index_col=0) for path in cfg.test_path], axis=0)
    else:
        raise ValueError('test_path must be a list of paths')

    if cfg.use_indices is not None:
        indices_to_use = load_indices(cfg.use_indices)
        X_train_used_indices = X_train.index.intersection(indices_to_use)
        X_test_used_indices = X_test.index.intersection(indices_to_use)
        used_indices = X_train_used_indices.union(X_test_used_indices)
    else:
        if cfg.exclude_indices:
            indices_to_exclude = load_indices(cfg.exclude_indices)
            X_train_used_indices = X_train.index.difference(indices_to_exclude)
            X_test_used_indices = X_test.index.difference(indices_to_exclude)
            used_indices = X_train_used_indices.union(X_test_used_indices)
        else:
            used_indices = X_train.index.union(X_test.index)
            X_train_used_indices = X_train.index
            X_test_used_indices = X_test.index

    X_train = X_train.loc[X_train_used_indices, :]
    X_test = X_test.loc[X_test_used_indices, :]

    return X_train, X_test, used_indices


def load_disease_related(cfg, used_indices):
    df_disease = pd.read_csv(cfg.disease_path, index_col=0)
    df_cov = pd.read_csv(cfg.cov_path, index_col=0)

    df_disease = df_disease.loc[used_indices, :]
    df_cov = df_cov.loc[used_indices, :]
    df_cov.columns = ['age', 'sex', 'bmi', 'assessment_date', 'death_date', 'center']
    df_cov = df_cov[['age', 'sex', 'bmi', 'assessment_date', 'death_date']]

    df_disease_code_map = pd.read_csv(cfg.disease_code_map_path)

    return df_disease, df_cov, df_disease_code_map


def add_covariates(X_train, X_test, df_cov, used_indices):
    df_cov_used = df_cov.loc[used_indices, :][['age', 'sex', 'bmi']]
    X_train_cov = df_cov_used.loc[X_train.index.intersection(used_indices), :]
    X_test_cov = df_cov_used.loc[X_test.index.intersection(used_indices), :]
    X_train = pd.concat([X_train, X_train_cov], axis=1)
    X_test = pd.concat([X_test, X_test_cov], axis=1)
    return X_train, X_test


def prepare_explain_data(X, df_disease_code_map, df_disease, df_cov, cfg):
    columns = df_disease_code_map[df_disease_code_map['code'].isin(cfg.disease_code)]['index'].values
    column = [f'participant.{c}' for c in columns][0]
    df_crt_disease = load_single_disease_data(
        column, df_disease, df_disease_code_map, df_cov, years=cfg.years
    )
    test_df = pd.concat((X, df_crt_disease.iloc[:, -2:]), axis=1).dropna()
    test_df = test_df.sample(500)

    test_df_surv = Survivaldata(
        test_df.iloc[:, :-2].to_numpy().astype(np.float32),
        test_df.iloc[:, -1].to_numpy().astype(np.float32),
        test_df.iloc[:, -2].to_numpy().astype(int),
    )

    explain_tensor = test_df_surv.features.to('cuda')
    background_indices = np.random.choice(explain_tensor.shape[0], 200, replace=False)
    background_tensor = explain_tensor[background_indices].float()

    exp_omics = explain_tensor[:, :-3]
    exp_cov = explain_tensor[:, -3:]

    return background_tensor, exp_omics, exp_cov


def get_protein_feature_names(cfg):
    """Read column names from the proteomics CSV, strip the 'olink_instance_0.' prefix, and uppercase."""
    header = pd.read_csv(cfg.train_path[0], index_col=0, nrows=0).columns.tolist()
    return [col.replace('olink_instance_0.', '').upper() for col in header]


def explain_protein_cox_model(background_tensor, exp_omics, exp_cov, df_disease_code_map, cfg):
    os.makedirs(cfg.save_dir, exist_ok=True)

    explain_data_combined = np.concatenate((exp_omics.cpu().numpy(), exp_cov.cpu().numpy()), axis=1)
    explain_tensor = torch.from_numpy(explain_data_combined).float().to('cuda')

    cox_model = CoxPL.load_from_checkpoint(cfg.cox_ckpt)
    risk_model = RiskModel(cox_model)

    explainer = shap.GradientExplainer(risk_model, background_tensor)
    n_features = explain_data_combined.shape[1]
    shap_values = explainer.shap_values(explain_tensor).reshape(-1, n_features)

    shap_omics = shap_values[:, :-3]
    shap_cov = shap_values[:, -3:]

    protein_feature_names = get_protein_feature_names(cfg)
    feature_names = protein_feature_names + ['Age', 'Sex', 'BMI']

    shap_values_combined = np.concatenate((shap_omics, shap_cov), axis=1)

    shap.summary_plot(
        shap_values_combined,
        features=explain_data_combined,
        feature_names=feature_names,
        max_display=15,
        show=False,
    )

    disease = df_disease_code_map[df_disease_code_map['code'] == cfg.disease_code[0]]['disease'].values[0]
    disease_name = disease.capitalize()
    plt.title(f'{disease_name} Proteomics')
    plt.savefig(f'{cfg.save_dir}/{disease}_proteomics.png', dpi=300, bbox_inches='tight')
    plt.close()


def main(cfg):
    X_train, X_test, used_indices = load_features(cfg)
    df_disease, df_cov, df_disease_code_map = load_disease_related(cfg, used_indices)
    X_train, X_test = add_covariates(X_train, X_test, df_cov, used_indices)

    background_tensor, exp_omics, exp_cov = prepare_explain_data(
        X_test, df_disease_code_map, df_disease, df_cov, cfg
    )
    explain_protein_cox_model(background_tensor, exp_omics, exp_cov, df_disease_code_map, cfg)


if __name__ == '__main__':

    class Config:
        cox_ckpt = 'ukb-clomics/logs/selected_diseases/paired/risk_score_with_cov/proteomics_baseline/cox_model/N03/epoch=19-step=5440.ckpt'

        train_path = ['data/processed_data/proteomics.train.csv',
                      'data/processed_data/proteomics.val.csv']
        test_path = ['data/processed_data/proteomics.test.csv']

        use_indices = None
        exclude_indices = None

        disease_path = 'data/raw_data/first_outcome.raw.csv'
        disease_code_map_path = 'data/ukb_disease_code_map.csv'
        cov_path = 'data/raw_data/covariants.raw.csv'

        save_dir = 'logs/shap/protein_plots'

        years = 10
        disease_code = ['N03']

    cfg = Config()
    main(cfg)