shap_plot.py

import argparse
import torch
import shap
import os
import matplotlib.pyplot as plt
import numpy as np
from ukb_clomics.model.align import CLIP
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)

# def set_pub_style():
#     """
#     Sets Matplotlib params to meet standard biomedical journal requirements:
#     - Arial/Helvetica font
#     - 7pt size for main text
#     - Editable text (Type 42)
#     - Clean 'classic' aesthetics
#     """
#     plt.rcParams['font.family'] = 'sans-serif'
#     plt.rcParams['font.sans-serif'] = ['Arial', 'Helvetica', 'DejaVu Sans']
#     plt.rcParams['font.size'] = 7
#     plt.rcParams['axes.labelsize'] = 7
#     plt.rcParams['axes.titlesize'] = 8
#     plt.rcParams['xtick.labelsize'] = 6
#     plt.rcParams['ytick.labelsize'] = 6
#     plt.rcParams['legend.fontsize'] = 6
#     plt.rcParams['pdf.fonttype'] = 42 
#     plt.rcParams['ps.fonttype'] = 42
#     plt.rcParams['axes.spines.top'] = False
#     plt.rcParams['axes.spines.right'] = False
#     plt.rcParams['lines.linewidth'] = 1



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

    # ACCEPT TWO SEPARATE INPUTS NOW
    def forward(self, x):
        # 3. Predict log-hazard
        log_hazard = self.cox_model.model(x).squeeze(-1) 
        
        return log_hazard.unsqueeze(-1)
    
## helper functions
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:
        print('Loading the specified indices for evaluation ...')
        indices_to_use = load_indices(cfg.use_indices)
        # new_used_indices = used_indices.intersection(indices_to_use)
        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_indices(path_list):
    """
    Load the indices from a list of paths
    
    Args:
        path_list (List[str]): list of paths to the index files
    Returns:
    set: set of indices
    """
    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_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)
    # make it slightly faster by only keeping used indices
    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_indices = df_cov.loc[used_indices,:][['age','sex', 'bmi',]]
    # train intersection
    X_train_used_indices = X_train.index.intersection(used_indices)
    X_test_used_indices = X_test.index.intersection(used_indices)
    
    X_train_cov = df_cov_used_indices.loc[X_train_used_indices,:]
    X_test_cov = df_cov_used_indices.loc[X_test_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.{column}' for column 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() ## data to explain
    
    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 explain_dl_model(background_tensor,exp_omics,exp_cov,df_disease_code_map,cfg):
    os.makedirs(cfg.save_dir, exist_ok=True)
    clip_encoder = CLIP.load_from_checkpoint(cfg.clip_ckpt)
    cox_model = CoxPL.load_from_checkpoint(cfg.clip_cox_ckpt)
    clip_encoder.eval()

    bg_omics = background_tensor[:, :-3]
    bg_cov = background_tensor[:, -3:]

    # Pre-compute CLIP latent representations (no gradients needed)
    with torch.no_grad():
        bg_latent = clip_encoder.model(bg_omics)
        exp_latent = clip_encoder.model(exp_omics)

    # Concatenate latent reps with covariates
    bg_latent_cov = torch.cat((bg_latent, bg_cov), dim=1)
    exp_latent_cov = torch.cat((exp_latent, exp_cov), dim=1)

    # Run SHAP only on the Cox model using latent + covariates as input
    full_model = BaselineRiskModel(cox_model)
    explainer = shap.GradientExplainer(full_model, bg_latent_cov)
    shap_values = explainer.shap_values(exp_latent_cov)

    latent_dim = exp_latent.shape[1]
    shap_values = np.array(shap_values).reshape(-1, latent_dim + 3)

    shap_latent = shap_values[:, :-3]

    if cfg.proteins_path is not None:
        protein_names = pd.read_csv(cfg.proteins_path, header=None)[0].str.upper().tolist()
        latent_feature_names = protein_names[:latent_dim]
    else:
        latent_feature_names = [f'latent_{i}' for i in range(latent_dim)]

    # Determine top features by mean |SHAP| (same order as SHAP summary_plot)
    max_display = 10
    mean_abs_shap = np.abs(shap_latent).mean(axis=0)
    sorted_indices = np.argsort(-mean_abs_shap)[:max_display]
    top_proteins = [latent_feature_names[i] for i in sorted_indices]

    # Exclude covariates (Age, Sex, BMI) from the SHAP plot
    shap.summary_plot(
        shap_latent,
        features=exp_latent.cpu().numpy(),
        feature_names=latent_feature_names,
        max_display=max_display,
        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} CLIP')
    plt.savefig(f'{cfg.save_dir}/{disease}_clip.png', dpi=300, bbox_inches='tight')
    plt.close()

    return top_proteins, disease
        


def main(cfg):

    ## load data once
    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)
    
    print(f'X_train shape: {X_train.shape}')

    for code in cfg.disease_codes:
        print(f'\n=== Processing disease code: {code} ===')
        cfg.disease_code = [code]
        cfg.clip_cox_ckpt = find_checkpoint(cfg.clip_cox_base, code)

        background_tensor, exp_omics, exp_cov = prepare_explain_data(X_test, df_disease_code_map, df_disease, df_cov, cfg)
        top_proteins, disease = explain_dl_model(background_tensor, exp_omics, exp_cov, df_disease_code_map, cfg)
    



def find_checkpoint(base_dir, disease_code):
    ckpt_dir = os.path.join(base_dir, disease_code)
    ckpts = [f for f in os.listdir(ckpt_dir) if f.endswith('.ckpt')]
    if len(ckpts) != 1:
        raise ValueError(f"Expected exactly 1 checkpoint in {ckpt_dir}, found {len(ckpts)}: {ckpts}")
    return os.path.join(ckpt_dir, ckpts[0])


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--clip_ckpt',
                        default='logs/align/clip/clip/version_0/checkpoints/epoch=68-step=18837.ckpt',
                        help='Path to CLIP checkpoint')
    parser.add_argument('--clip_cox_base',
                        default='logs/shap/clip_concat/cox_model',
                        help='Base folder for CLIP-Cox checkpoints')
    parser.add_argument('--train_path', nargs='+',
                        default=['data/processed_data/metabolomics.train.csv',
                                 'data/processed_data/metabolomics.val.csv'],
                        help='Paths to training CSV files')
    parser.add_argument('--test_path', nargs='+',
                        default=['data/processed_data/metabolomics.test.csv'],
                        help='Paths to test CSV files')
    parser.add_argument('--exclude_indices', nargs='+',
                        default=None,
                        help='Paths to index files to exclude')
    parser.add_argument('--use_indices', nargs='+', default=['data/processed_data/proteomics.train.csv',
                                 'data/processed_data/proteomics.val.csv',
                                 'data/processed_data/proteomics.test.csv'],
                        help='Paths to index files to use (overrides exclude_indices)')
    parser.add_argument('--disease_path', default='data/raw_data/first_outcome.raw.csv')
    parser.add_argument('--disease_code_map_path', default='data/ukb_disease_code_map.csv')
    parser.add_argument('--cov_path', default='data/raw_data/covariants.raw.csv')
    parser.add_argument('--proteins_path', default='data/proteins.txt')
    parser.add_argument('--clip_protein_corr_path', default='clip_protein_correlations.csv')
    parser.add_argument('--save_dir', default='logs/shap/plots')
    parser.add_argument('--years', type=int, default=10)
    parser.add_argument('--disease_codes', nargs='+', required=True,
                        help='One or more disease codes (e.g. K71 N03 I50)')
    args = parser.parse_args()

    main(args)