goldilocks-pysustain.py

#!/usr/bin/env python3
# Demonstration of Goldilocks DPM framework for pySuStaIn
# Neil Oxtoby, UCL, 2023

run_locally = False
rundate = '20231211.2'

if rundate=='20231211':
    dataset_name  = "SynthADNI"
    config        = "vanilla"
    nom           = ""
elif rundate=='20231211.2':
    dataset_name  = "SynthADNI"
    config        = "goldilocks"
    nom           = "_rerun"

from pathlib import Path
path_to_goldilocks_dpm = Path.cwd().parent / "goldilocks-dpm"
path_to_demos = path_to_goldilocks_dpm / "goldilocks_dpm" / "demos"
output_folder = path_to_demos / "cluster"

import sys,os
# sys.path.append(path_to_goldilocks_dpm)
os.chdir(path_to_goldilocks_dpm)

import pandas as pd, numpy as np, seaborn as sn
from matplotlib import pyplot as plt

from goldilocks_dpm import goldilocks_ZscoreSustain

from pySuStaIn.ZscoreSustain  import ZscoreSustain

# Load synthetic data
csv_main = path_to_demos / "ADNIMERGE2023_synthetic.csv"
assert csv_main.exists(), f"ERROR: {csv_main.name} file not found."
df = pd.read_csv(csv_main,low_memory=False)

# Save CV indices for cluster
import pickle
from sklearn.model_selection import StratifiedKFold
N_folds = 10
pickle_filename_cv = output_folder / csv_main.name.split('/')[-1].replace('.csv','-cvindices.pickle')

biomarkers         = ["ABETA","TAU","ADAS13","MMSE","Ventricles_ICV"]
direction_abnormal = [     -1,    1,       1,    -1,               1]

# Patients and controls
dx_col = "DX"
id_col = "ID"
df['y'] = df[dx_col].values
controls  = (df[dx_col]==0).values
cases     = (df[dx_col]==1).values
prodromal = (df[dx_col]==2).values

# Raw data
X_synth = df[biomarkers].values
y_synth = df[dx_col].values
id_synth = df[id_col].values
print(f"Your data contains {X_synth.shape[0]} samples from {X_synth.shape[1]} biomarkers: \n{biomarkers}")

# Create Goldilocks DPM object, version 1: without a pySuStaIn object
gdpm = goldilocks_ZscoreSustain(
    classes = y_synth,
    dpmData = X_synth,
    output_folder = output_folder,
    robust_zscores = False,
    case_label = 1,
    ctrl_label = 0, 
    direction_abnormal = direction_abnormal,
    biomarker_labels = biomarkers
)

# Scroll down to "ALTERNATIVE" for how to create Goldilocks DPM object using an existing pySuStaIn object

gdpm.run_goldilocks(plot=True, plot_format="png", verbose=False)

print(gdpm.Z_vals)
print(gdpm.Z_max)

# Add z-scores to dataframe
sustain_features = [f"{b}_z" for b in biomarkers]
Z_synth = gdpm.Z.copy()
df[sustain_features] = Z_synth

data = Z_synth
y    = y_synth
id_  = id_synth

sustain_output_folder = '_'.join( [dataset_name,'zscore',rundate] )
sequence_without_controls = True
sustainType             = 'zscore'
y_for_sequencing = list(df[dx_col].values!=0)
if sequence_without_controls:
    df_train = df.loc[y_for_sequencing].copy()
    data_sequencing = data[y_for_sequencing,:]
    y_sequencing    = y[y_for_sequencing,]
    id_sequencing   = id_[(y_for_sequencing)]
    sustain_output_folder = output_folder / f'{sustain_output_folder}_sequence_without_controls'
else:
    df_train = df.copy()
    data_sequencing = data
    y_sequencing    = y
    id_sequencing   = id_
    sustain_output_folder = output_folder / sustain_output_folder
os.system(f'mkdir -p {sustain_output_folder}')


# number of starting points
N_startpoints           = 25
# maximum number of subtypes
N_S_max                 = 4
N_iterations_MCMC       = int(1e4)  #int(1e6)
# cross-validation
validate                = True
N_folds                 = 10

M_seq, N_seq            = data_sequencing.shape     # number of individuals, events

## pySuStaIn v1: defaults
z_events                = [1,2,3]
Z_vals_default          = np.array([z_events]*N_seq) # Z-scores for each biomarker
Z_max_default           = np.array([5]*N_seq)        # maximum z-score

## Goldilocks-informed pySuStaIn
Z_vals_goldilocks       = gdpm.Z_vals                # Z-scores for each biomarker
Z_max_goldilocks        = np.array(gdpm.Z_max)       # maximum z-score

# Avoid Z_max < Z_vals
Z_max_default    = np.max([np.reshape(np.max(Z_vals_default,axis=1),(-1,1)),np.reshape(Z_max_default,(5,1))],axis=0).flatten()
Z_max_goldilocks = np.max([np.reshape(np.max(Z_vals_goldilocks,axis=1),(-1,1)),np.reshape(Z_max_goldilocks,(5,1))],axis=0).flatten()


sustain_default = ZscoreSustain(
    data_sequencing, 
    Z_vals_default, Z_max_default, 
    biomarkers, 
    N_startpoints, N_S_max, N_iterations_MCMC, 
    output_folder, 'SynthADNI_sustain_default',
    use_parallel_startpoints=False, seed=42
)

sustain_goldilocks = ZscoreSustain(
    data_sequencing,
    Z_vals_goldilocks, Z_max_goldilocks,
    biomarkers,
    N_startpoints, N_S_max, N_iterations_MCMC,
    output_folder, 'SynthADNI_sustain_goldilocks',
    use_parallel_startpoints=False, seed=42
)

## ALTERNATIVE: create Goldilocks DPM object directly from a pySuStaIn object
## TODO: make this work.
# gdpm_alt = goldilocks_ZscoreSustain(
#     sustain_object = sustain_default,
#     #dpmData = None,
#     classes = y_sequencing,
#     output_folder = output_folder,
#     robust_zscores = False,
#     case_label = 1,
#     ctrl_label = 0,
#     direction_abnormal = direction_abnormal,
#     biomarker_labels = biomarkers
# )
# gdpm_alt.run_goldilocks(plot=True, plot_format="png", verbose=False)



if ~run_locally:
    print("Preparing data for SuStaIn to be run elsewhere, e.g., on a compute cluster")
    # Training data
    if sum(df_train['y'].isnull())==0:
        df_train['y'] = df_train['y'].astype(int).values
    else:
        print(f'Missing data in y (cases/controls labels): df_train["y"].isnull().sum(): {df_train["y"].isnull().sum()} (of {df_train["y"].shape[0]}):')
        print(df_train.loc[df_train["y"].isnull()].groupby('Diagnosis')[id_col].count())
    # Write out to CSV for the cluster
    colz = ['y',id_col,dx_col] + sustain_features
    colz = colz + [c for c in df_train.columns.tolist() if c not in colz]
    df[colz].to_csv(str(csv_main).replace('.csv','_wrangled-all.csv'),index=False)
    df_train[colz].to_csv(str(csv_main).replace('.csv','_wrangled-train.csv'),index=False)
else:
    # Don't worry, if run_locally==True, SuStaIn is run further down if it hasn't been run already
    print("")



if pickle_filename_cv.exists():
    pickle_file                 = open(pickle_filename_cv, 'rb')
    variables_to_pickle         = pickle.load(pickle_file)
    test_idxs = variables_to_pickle['test_idxs']
    pickle_file.close()
    for k in range(len(test_idxs)):
        id_train = id_sequencing[test_idxs[k]] #[id_sequencing[j] for j in t]
        rowz = df["ID"].isin(id_train).values
        df.loc[rowz,'cv_idx'] = k
else:
    df['cv_idx'] = np.nan
    #* k-fold cross validation
    test_idxs  = []
    train_idxs  = []
    cv         = StratifiedKFold(n_splits=N_folds, shuffle=True)
    cv_it      = cv.split(data_sequencing, y_sequencing)
    
    k = 0
    for train, test in cv_it:
        test_idxs.append(test)
        train_idxs.append(train)
        k += 1
    test_idxs = np.array(test_idxs,dtype='object')
    train_idxs = np.array(train_idxs,dtype='object')
    
    for k in range(len(test_idxs)):
        id_train = id_sequencing[test_idxs[k]] #[id_sequencing[j] for j in t]
        rowz = df["ID"].isin(id_train).values
        df.loc[rowz,'cv_idx'] = k
    
    pickle_filepath             = Path(pickle_filename_cv)
    pickle_file                 = open(pickle_filename_cv, 'wb')
    variables_to_pickle         = {
        'test_idxs': test_idxs
    }
    po = pickle.dump(variables_to_pickle,pickle_file)
    pickle_file.close()





if run_locally:
    ## Default SuStaIn
    # get the start time
    st_default = time.process_time()
    sustain_default.run_sustain_algorithm(plot=False)
    # get the end time
    et_default = time.process_time()
    # get execution time
    res_default = et_default - st_default
    print('Default SuStaIn CPU Execution time:', res_default/60, 'minutes')


    ## Goldilocks SuStaIn
    # get the start time
    st_goldilocks = time.process_time()
    sustain_goldilocks.run_sustain_algorithm(plot=False) # plot=True gives an error: pySuStaIn plotting only handles identical z-event sets per biomarker
    # get the end time
    et_goldilocks = time.process_time()
    # get execution time
    res_goldilocks = et_goldilocks - st_goldilocks
    print('Goldilocks-informed SuStaIn CPU Execution time:', res_goldilocks/60, 'minutes')





if run_locally:
    ## TODO: plot staging distributions => should get better (flatter) spread for Goldilocks
    data_inference = gdpm.Z
    y_inference = gdpm.classes
    s = 1

    pickle_filename_s           = output_folder + '/pickle_files/' + sustain_default.dataset_name + '_subtype' + str(s) + '.pickle'
    pickle_filepath             = Path(pickle_filename_s)
    pickle_file                 = open(pickle_filename_s, 'rb')
    loaded_variables            = pickle.load(pickle_file)
    ml_subtype                  = loaded_variables["ml_subtype"]
    prob_ml_subtype             = loaded_variables["prob_ml_subtype"]
    ml_stage                    = loaded_variables["ml_stage"]
    prob_ml_stage               = loaded_variables["prob_ml_stage"]
    prob_subtype                = loaded_variables["prob_subtype"]
    prob_stage                  = loaded_variables["prob_stage"]
    prob_subtype_stage          = loaded_variables["prob_subtype_stage"]
    samples_sequence            = loaded_variables["samples_sequence"]
    samples_f                   = loaded_variables["samples_f"]
    pickle_file.close()

    N_samples            = data_inference.shape[0]
    ml_subtype,          \
    prob_ml_subtype,     \
    ml_stage,            \
    prob_ml_stage,       \
    prob_subtype,        \
    prob_stage,          \
    prob_subtype_stage   = sustain_default.subtype_and_stage_individuals_newData(
        gdpm.Z,
        samples_sequence,
        samples_f,
        N_samples
    )

    # Test out new plotting code from ChatGPT
    all_zscores = np.unique(np.sort(list(sustain_goldilocks.stage_zscore) + list(sustain_default.stage_zscore)))
    # import matplotlib.colors as mcolors
    # print(list(mcolors.XKCD_COLORS.keys())[0:len(all_zscores)])

    from matplotlib import cm
    viridis = cm.get_cmap('viridis', len(all_zscores))
    colour_mat_all = viridis(all_zscores/all_zscores.max())
    # Select subset of colours
    zscores_defaults   = np.unique(sustain_default.stage_zscore)
    zscores_goldilocks = np.unique(sustain_goldilocks.stage_zscore)
    rowz_default    = [True if z in zscores_defaults   else False for z in all_zscores]
    rowz_goldilocks = [True if z in zscores_goldilocks else False for z in all_zscores]
    colour_mat_default    = colour_mat_all[rowz_default,   :]
    colour_mat_goldilocks = colour_mat_all[rowz_goldilocks,:]

    plot_order = np.arange(len(biomarkers))

    subtype_labels = ['0','1']

    fig,ax,cbar = plot_SuStaIn_model_arbitrarycolours(
        samples_sequence, 
        samples_f,
        biomarkers,
        sustain_goldilocks.stage_zscore,
        sustain_goldilocks.stage_biomarker_index,
        colour_mat_all,
        plot_order,
        subtype_labels,
        all_zscores
    )
    xt = all_zscores/np.max(all_zscores)
    xtl = [str(k) for k in all_zscores]
    cbar.set_ticks(xt)
    cbar.set_ticklabels(xtl)


    #* Plot subtypes and stages by diagnostic group
    stages_bins = np.arange(-0.5,1+data_inference.shape[1]*np.max(Z_max_default))
    # dx_list = np.sort(list(dx_column_mapper.keys()))
    # dx_list_label = [dx_column_mapper[dx][-1] for dx in dx_list]
    dx_list = np.unique(df[dx_col].values)
    dx_list_label = [str(dx) for dx in dx_list]
    fig,ax = plt.subplots(2,s+1,figsize=(16,8),sharey=False)
    axflat = ax.flatten()
    for k in range(s):
        st_k = (ml_subtype==k).flatten()
        st_nonzero = (ml_stage>0).flatten()
        ax[0,k].hist(y_inference[st_k & st_nonzero],bins=np.arange(-0.5,max(dx_list)+0.5))
        ax[0,k].set_xticks(dx_list)
        ax[0,k].set_xticklabels(dx_list_label,fontsize=18)
        ax[0,k].set_title('Subtype %i' % (k+1),fontsize=20)
        stages_tmp = []
        for dx_k in dx_list:
            if sum((y_inference==dx_k))==0:
                continue
            else:
                # FIXME: handle empty cases where no individuals with dx_k are subtyped/staged
                ax[1,k].hist( ml_stage[st_k & st_nonzero & (y_inference==dx_k)], label=dx_list_label,bins=stages_bins)
                ax[1,k].set_xlabel('Stage',fontsize=20)
    ax.flatten()[0].set_ylabel('Count',fontsize=20)
    ax.flatten()[0].set_xlabel('DX',fontsize=20)
    # ax[1,0].set_ylabel('Count',fontsize=20)
    # ax[1,0].set_ylim([0,50])
    ax.flatten()[0].legend(fontsize=20)
    fig.tight_layout()
    fig.show()