dwi_parse_z.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#%% Imports
"""
Created on Tue Sep 24 09:58:05 2024

@author: andmar
"""

import os
import sys
import copy
import pickle
import numpy as np
import pandas as pd
import statsmodels.api as sm
from scipy import stats

from matplotlib import pyplot as plt
import seaborn as sns
import statsmodels 

root_dir = '/Users/andmar/data/sairut/braincharts'
#w_dir = os.path.join(root_dir,'models','lifespan_FA_24K_19sites'); pkl = False
w_dir = '/Users/andmar/data/ramcir/FA_models_race_clean_nodhcp_6'; pkl = True

zthresh = 2

#%% Load  data

with open(os.path.join(root_dir,'docs','phenotypes_fa.txt')) as f:
    idp_ids_fa = f.read().splitlines()

symptoms = ['m1negtsx', 'm2postsx','m3distsx','panss_total']
symptom = symptoms[-1]

if pkl:
    with open(os.path.join(w_dir,'df_te.pkl'),'rb') as f:
        #df_te = pickle.load(f)
        df_te = pd.read_pickle(f)
else:
    df_te = pd.read_csv(os.path.join(w_dir,'df_te.csv'),index_col=0)
df_te = df_te.loc[:, ~df_te.columns.str.startswith('Zscore')]

# load panss scores
panss = pd.read_csv('/Users/andmar/data/datasets/HCP_EP/panss01.txt',index_col=0,skiprows=[1],delimiter='\t')
panss['id'] = 'sub-' + panss['src_subject_id'].astype(str) + '_01_MR'
panss_dwi = panss.loc[panss['id'].isin(df_te['id'])]

# select only HCPEP sites
hcpep_te = ((df_te['site'] == 'HCPEP_3010') | (df_te['site'] == 'HCPEP_3011') | (df_te['site'] == 'HCPEP_3012') | (df_te['site'] == 'HCPEP_3013'))
df_te_ep = df_te.loc[hcpep_te]

# load z-statistics
col_names = []
for i, idp in enumerate(idp_ids_fa):
    z = np.loadtxt(os.path.join(w_dir,idp,'Z_estimate.txt'))
    df_te_ep['Z_'+idp] = z[hcpep_te]
    col_names.append('Z_'+idp)
colz_fa = col_names

#%% Run t-test (DWI)

# first do a t-test agaoinst diagnosis
pat = df_te_ep['diagnosis'] == 10.
t = stats.ttest_ind(df_te_ep[colz_fa].loc[pat], df_te_ep[colz_fa].loc[~pat])
t2 = stats.ttest_ind(df_te_ep[idp_ids_fa].loc[pat], df_te_ep[idp_ids_fa].loc[~pat])
h, pfdr = statsmodels.stats.multitest.fdrcorrection(t.pvalue)

print('\n','nominally significant effects:')
print('\n'.join([idp_ids_fa[i] for i in np.where(t.pvalue < 0.05)[0]]))
print('\n','FDR corrected effects (diagnosis):')
print('\n'.join([idp_ids_fa[i] for i in np.where(h)[0]]))

#plt.bar(range(len(t.statistic)),t.statistic)
plt.plot(t.statistic)
plt.plot(t2.statistic)
plt.legend(('z','raw'))
plt.xlabel('ROI')
plt.ylabel('t')
plt.title('diagnsosis effect (FA)')
plt.show()

# join dataframes in include symptoms and DWI metrics
df = pd.concat([df_te_ep.set_index('id'),panss_dwi.set_index('id')], axis=1, join='inner')

#df[col_names + ['panss_total']].corr()['pans_total']
p = np.zeros(len(idp_ids_fa))
r = np.zeros(len(idp_ids_fa))
for i, idp in enumerate(col_names):
    r[i], p[i] = stats.spearmanr(df[idp], df[symptom])

r_raw = np.zeros(len(idp_ids_fa))
p_raw = np.zeros(len(idp_ids_fa))
for i, idp in enumerate(idp_ids_fa):
    r_raw[i], p_raw[i] = stats.spearmanr(df[idp], df[symptom])
    
plt.plot(r)
plt.plot(r_raw)
plt.legend(('z','raw'))
plt.title(symptom + ' (FA)')
plt.ylabel('correlation')
plt.xlabel('ROI')
plt.show()

h, pfdr = statsmodels.stats.multitest.fdrcorrection(p)

print('\n','nominally significant effects (symptoms):')
print('\n'.join([idp_ids_fa[i] for i in np.where(p < 0.05)[0]]))
print('\n','FDR corrected effects (symptoms):')
print('\n'.join([idp_ids_fa[i] for i in np.where(h)[0]]))

#%% Cortical thickness - load data from raw z -scores 

w_dir_ct = os.path.join(root_dir,'models','lifespan_29K_82sites_pat')
df_te_ct = pd.read_csv('/Users/andmar/data/sairut/data/lifespan_big_patients_te.csv', index_col=0)

hcpep_te_ct = ((df_te_ct['site'] == 'HCP_EP_IU') | (df_te_ct['site'] == 'HCP_EP_BWH') | (df_te_ct['site'] == 'HCP_EP_MGH') | (df_te_ct['site'] == 'HCP_EP_McL'))
df_te_ct_ep = df_te_ct.loc[hcpep_te_ct]

with open(os.path.join(root_dir,'docs','phenotypes_ct_lh.txt')) as f:
    idp_ids_lh = f.read().splitlines()
with open(os.path.join(root_dir,'docs','phenotypes_ct_rh.txt')) as f:
    idp_ids_rh = f.read().splitlines()
with open(os.path.join(root_dir,'docs','phenotypes_sc.txt')) as f:
    idp_ids_sc = f.read().splitlines()

idp_ids = idp_ids_lh + idp_ids_rh + idp_ids_sc

# remove bad IDPs
idp_ids = [i for i in idp_ids if not 'vessel' in i]
idp_ids = [i for i in idp_ids if not 'choroid-plexus' in i]
idp_ids = [i for i in idp_ids if not 'TotalGrayVol' in i]
idp_ids = [i for i in idp_ids if not 'SupraTentorialVolNotVent' in i]
idp_ids = [i for i in idp_ids if not 'EstimatedTotalIntraCranialVol' in i]

#%% load CT z-scores
# col_names = []
# for i, idp in enumerate(idp_ids):
#     z = np.loadtxt(os.path.join(w_dir_ct,idp,'Z_estimate.txt'))

#     #df_te_ct_ep['Z_'+idp] = z[hcpep_te_ct]
#     col_names.append('Z_'+idp)
    
#     df_te_idp = pd.read_csv(os.path.join(w_dir_ct, idp,'df_te.csv'),index_col=0)

#     hcpep_te_idp = ((df_te_idp['site'] == 'HCP_EP_IU') | (df_te_idp['site'] == 'HCP_EP_BWH') | (df_te_idp['site'] == 'HCP_EP_MGH') | (df_te_idp['site'] == 'HCP_EP_McL'))
   
#     zh = z[hcpep_te_idp]
#     dfh = df_te_idp.loc[hcpep_te_idp]
    
#     dfh['Z_'+idp] = zh
        
#     df_te_ct_ep['Z_'+idp] =  dfh['Z_'+idp]

# colz_ct = col_names
# df_te_ct_ep.to_csv('df_te_ct_ep.csv')

#%% load CT z scores - from data frame

df_te_ct_ep = pd.read_csv('/Users/andmar/data/sairut/helper_scripts/df_te_ct_ep.csv',index_col=0)
colz_ct = df_te_ct_ep.columns.to_list()[211:]
col_names = colz_ct

#%% Cortical thickness
panss_ct = panss.loc[panss['id'].isin(df_te_ct.index)]
dfc = pd.concat([df_te_ct_ep,panss_dwi.set_index('id')], axis=1, join='inner')

#df[col_names + ['panss_total']].corr()['pans_total']
p = np.zeros(len(idp_ids))
r = np.zeros(len(idp_ids))
for i, idp in enumerate(col_names):
    r[i], p[i] = stats.spearmanr(dfc[idp], dfc[symptom])

r_raw = np.zeros(len(idp_ids))
p_raw = np.zeros(len(idp_ids))
for i, idp in enumerate(idp_ids):
    r_raw[i], p_raw[i] = stats.spearmanr(dfc[idp], dfc[symptom])
    
plt.plot(r)
plt.plot(r_raw)
plt.legend(('z','raw'))
plt.title(symptom + ' (CT + SC)')
plt.ylabel('correlation')
plt.xlabel('ROI')
plt.show()

h, pfdr = statsmodels.stats.multitest.fdrcorrection(p)

print('\n','nominally significant effects (symptoms):')
print('\n'.join([idp_ids[i] for i in np.where(p < 0.05)[0]]))

print('\n','FDR corrected effects (symptoms):')
print('\n'.join([idp_ids[i] for i in np.where(h)[0]]))

#df_all = df[idp_ids_fa].join(dfc[idp_ids + symptoms + ['sex','age'] ])
df_all = df[colz_fa].join(dfc[colz_ct + symptoms + ['sex','age'] ])

#%% Run msCCA

rootdir = '/Users/andmar/data/sscca'

sys.path.append(os.path.join(rootdir, 'saccade'))
from scca import MSCCA
from utils import deflate

def run_mscca(X, l1, niter, sign, rank, n_views, tr_frac):
    # initialise the weights (to store for visualisation)
    W = []
    for v in range(n_views):
        W.append( np.zeros((X[v].shape[1], rank)) )

    tr = np.random.uniform(size=X[0].shape[0]) < 0.7
    te = ~tr
       
    # standardize
    Xtr = []
    Xte = []
    for v in range(n_views):
        m = np.mean(X[v][tr,:], axis = 0)
        s = np.std(X[v][tr,:], axis = 0)
        Xtr.append( (X[v][tr,:] - m) / s )
        Xte.append( (X[v][te,:] - m) / s )
    
    Cm = MSCCA(n_components=rank, n_views=n_views)
    Cm.fit(Xtr, l1=l1, sign=sign, verbose=False)
    
    scores_te = Cm.transform(Xte)
    if np.isnan(scores_te).any():
        Cm.fit(Xtr, l1=l1, sign=sign, verbose=False)   

    # compute the canonical correlations
    R12 = np.zeros(rank)
    R13 = np.zeros(rank)
    R23 = np.zeros(rank) 
    R = np.zeros(rank)   
    for r in range(rank):
          R12[r] = np.corrcoef(scores_te[0][:,r].ravel(), scores_te[1][:,r].ravel())[0][1] 
          R13[r] = np.corrcoef(scores_te[0][:,r].ravel(), scores_te[2][:,r].ravel())[0][1] 
          # this is not included in the objective function
          R23[r] = np.corrcoef(scores_te[1][:,r].ravel(), scores_te[2][:,r].ravel())[0][1] 
          
          R[r] =  (R12[r] +  R13[r])/2
          
    # save the weights
    #for v in range(n_views):
    #    W[v][:,:] = Cm.W[v]

    return R, Cm, R12, R13, R23

# msCCA paramaters
l1 = [0.9, 0.2, 0.2] # sparsity parameters
niter = 1000;        # number of iterations 
sign = [1., 0, 0]     # sign contraints
rank = 3             # number of components
n_views = 3          # number of data modalities. view 0 = symptoms
n_splits = 1000

# remove some bad data
df_all = df_all.fillna(df_all.median(numeric_only=True))

#configure data matrix
X = [ df_all[symptoms].to_numpy(),
      df_all[colz_fa].to_numpy(), 
      df_all[colz_ct].to_numpy() ]

# initialise the weights (to store for visualisation)
W = []
for v in range(n_views):
    W.append( np.zeros((X[v].shape[1], rank, n_splits)) )

# run msCCA
R = np.zeros((n_splits,rank))
R_all = np.zeros((n_splits,n_views,rank))
for i in range(n_splits):
    #print('split', i, 'fitting scca...')

    r, Cm,  R12, R13, R23 = run_mscca(X, l1, niter, sign, rank, n_views, 0.7)
    R[i,:] = r
    for r in range(rank):
        R_all[i,0,r] = R12[r]
        R_all[i,1,r] = R13[r]
        R_all[i,2,r] = R23[r]
         
    # save the weights
    for v in range(n_views):
        W[v][:,:,i] = Cm.W[v]
    
    print('r_test =',R[i,:])

sns.displot(R,kind='kde', fill=True)
plt.xlabel('r (test)')
plt.title("Canonical correlation (overall)")
plt.show()

sns.displot(R_all[:,0,:],kind='kde',fill=True)
plt.xlabel('r (test)')
plt.title("Canonical correlation (view 0 and view 1)")
plt.show()

sns.displot(R_all[:,1,:],kind='kde',fill=True)
plt.xlabel('r (test)')
plt.title("Canonical correlation (view 0 and view 2)")
plt.show()

sns.displot(R_all[:,2,:],kind='kde',fill=True)
plt.xlabel('r (test)')
plt.title("Canonical correlation (view 1 and view 2)")
plt.show()

print('overall', f'r_test = {np.mean(R,axis=0)}, std ={np.std(R,axis=0)} ') 


#%% plot Weights using stability selection
select_thresh = 0.

for i in range(n_views):
    plt.imshow(W[i][:,0,:].T,aspect='auto')
    plt.colorbar()
    plt.title('weights view ' + str(i))
    plt.xlabel('variable')
    plt.ylabel('iteration')
    plt.show()
    
    if i > 0:
        p_selection = np.sum(np.abs(W[i][:,0,:]) > 0.0001,axis = 1) / n_splits
        plt.plot(p_selection)
        plt.title('Selection probability')
        plt.show()
        
        print('\n','Stable features view',i,':')
        if i == 1:
            print('\n'.join([idp_ids_fa[i] for i in np.where(p_selection > select_thresh)[0]]))
        elif i == 2:
            print('\n'.join([idp_ids[i] for i in np.where(p_selection > select_thresh)[0]]))

sns.violinplot(pd.DataFrame(W[0][:,0,:].T, columns=symptoms))
plt.title('symptom weights (view 0)')
plt.show()

for v in range(1,3):
    plt.errorbar(range(W[v].shape[0]),np.mean(W[v][:,0,:],axis=1), np.std(W[v][:,0,:],axis=1),linestyle="", marker='o')
    plt.title(f'weights (view {v})')
    plt.show()


#%% Permutation test

n_perm = 1000
n_splits = 1

#R = np.zeros((n_splits,rank))
#for i in range(n_splits):
#    r, Cm, R12, R13, R23 = run_mscca(X, l1, niter, sign, rank, n_views, 0.7)
#    R[i,:] = r
r_true = np.mean(R,axis=0)

Xp = copy.deepcopy(X)

R_perm = np.zeros((n_perm,rank))
for p in range(n_perm):
    print('permutation',p)
    Xp[0] = np.random.permutation(X[0])
    Rp = np.zeros((n_splits,rank))
    for i in range(n_splits):

        r, Cm,  R12, R13, R23 = run_mscca(Xp, l1, niter, sign, rank, n_views, 0.7)
        Rp[i,:] = r
    R_perm[p,:] = np.mean(Rp,axis=0)

print('p-values:', np.sum(R_perm > r_true, axis=0) / n_perm )
# %%