sdt.py

from data_prep import *
import warnings
import concurrent.futures
import pickle as pickle
import numpy as np
import math
import itertools
from scipy.optimize import linprog
import random
import seaborn as sns
from scipy import stats
from tqdm import tqdm
from sklearn.model_selection import KFold, GridSearchCV
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, precision_recall_fscore_support, make_scorer
from scipy.stats.mstats_basic import rankdata
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.neighbors import KNeighborsClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import SVC
from sklearn.neural_network import MLPClassifier
import xgboost as xgb


SEED = 0
warnings.filterwarnings('ignore') 
sns.set(color_codes=True)

_, allAffect, _, allPredLabel, allTrueLabel, _, _ = load_data()
firstAffect, _, _, _, firstPredLabel, firstTrueLabel = load_data(mode=1)
secondAffect, _, _, _, secondPredLabel, secondTrueLabel = load_data(mode=2)
thirdAffect, _, _, _, thirdPredLabel, thirdTrueLabel = load_data(mode=3)

def sdt_element(true, pred, hp='1', show=False):
    
    num_agent = sum([1 for l in true if l == 0])
    num_human = sum([1 for l in true if l == 2])
    num_agent_agent = sum([1 for i, _ in enumerate(true) if true[i] == 0 and pred[i] == 0])
    num_human_agent = sum([1 for i, _ in enumerate(true) if true[i] == 2 and pred[i] == 0])
    num_agent_not_sure = sum([1 for i, _ in enumerate(true) if true[i] == 0 and pred[i] == 1])
    num_human_not_sure = sum([1 for i, _ in enumerate(true) if true[i] == 2 and pred[i] == 1])
    num_human_human = sum([1 for i, _ in enumerate(true) if true[i] == 2 and pred[i] == 2])
    num_agent_human = sum([1 for i, _ in enumerate(true) if true[i] == 0 and pred[i] == 2])
    
    if hp == '1':
        c2_hit_rate = num_human_human / num_human 
        c2_false_alarm_rate = num_agent_human / num_agent
        c1_hit_rate = c2_hit_rate + num_human_not_sure / num_human
        c1_false_alarm_rate = c2_false_alarm_rate + num_agent_not_sure / num_agent

    if hp == '2':
        c2_hit_rate = num_agent_agent / num_agent 
        c2_false_alarm_rate = num_human_agent / num_human
        c1_hit_rate = c2_hit_rate + num_agent_not_sure / num_agent
        c1_false_alarm_rate = c2_false_alarm_rate + num_human_not_sure / num_human
    
    zs_c2_hit_rate = -stats.norm.ppf(c2_hit_rate)
    zs_c2_false_alarm_rate = -stats.norm.ppf(c2_false_alarm_rate)
    zs_c1_hit_rate = -stats.norm.ppf(c1_hit_rate)
    zs_c1_false_alarm_rate = -stats.norm.ppf(c1_false_alarm_rate)
    
    if hp == '1':
        c2_human_criterion = zs_c2_hit_rate
        c2_agent_criterion = zs_c2_false_alarm_rate
        c1_human_criterion = zs_c1_hit_rate
        c1_agent_criterion = zs_c1_false_alarm_rate

    if hp == '2':
        c2_agent_criterion = zs_c2_hit_rate
        c2_human_criterion = zs_c2_false_alarm_rate
        c1_agent_criterion = zs_c1_hit_rate
        c1_human_criterion = zs_c1_false_alarm_rate
    
    c2_d_prime = zs_c2_false_alarm_rate - zs_c2_hit_rate
    c1_d_prime = zs_c1_false_alarm_rate - zs_c1_hit_rate
   
    d_prime = (c2_d_prime + c1_d_prime) / 2

    if show:
        return (c1_agent_criterion, c2_agent_criterion, c1_human_criterion, c2_human_criterion, c2_d_prime, c1_d_prime, d_prime)
    else:
        return (c1_agent_criterion, c2_agent_criterion, c1_human_criterion, c2_human_criterion, d_prime)


def sdt_model(av_human_m, av_agent_m, av_human_sd, av_agent_sd, criterion, X_test_av, X_test_true, hp):
    
    if math.isnan(av_human_m) or math.isnan(av_agent_m) or math.isnan(av_human_sd) or math.isnan(av_agent_sd):

        return 1, (1, X_test_true, 0)

    d_prime = criterion[-1]
    
    if hp == '1':
        order = [0, 1, 2]
    else:
        order = [2, 1, 0]
        d_prime *= -1
        
    if X_test_true == 2:

        if av_human_sd != 0:
            z_av = (X_test_av - av_human_m) / av_human_sd
        else:
            z_av = 0

        if z_av <= criterion[2]:
            return order[0], (0, z_av, z_av+d_prime)
        if z_av > criterion[2] and z_av < criterion[3]:
            return order[1], (1, z_av, z_av+d_prime) 
        if z_av >= criterion[3]:
            return order[2], (2, z_av, z_av+d_prime)
        
    else:
        if av_agent_sd != 0:
            z_av = (X_test_av - av_agent_m) / av_agent_sd
        else:
            z_av = 0
        
        if z_av <= criterion[0]:
            return order[0], (0, z_av, z_av) 
        if z_av > criterion[0] and z_av < criterion[1]:
            return order[1], (1, z_av, z_av) 
        if z_av >= criterion[1]:
            return order[2], (2, z_av, z_av)


def evaluation(y_true, result, show=True):
    
    y_pred = [r[1] for r in result]
    y_m = [r[2][0] for r in result]
    
    accuracy = accuracy_score(y_true, y_pred)
    p_macro, r_macro, f1_macro, _ = precision_recall_fscore_support(y_true, y_pred, average='macro')
    r1, p1 = stats.spearmanr(y_true, y_pred, alternative='greater')
    r2, p2 = stats.spearmanr(y_true, y_m)
    
    if show:

        print('Nested LOOCV')
        print('    accuracy:              %.4f' % (accuracy))
        print('    precision-macro:       %.4f' % (p_macro))
        print('    recall-macro:          %.4f' % (r_macro))
        print('    f1-marcro:             %.4f' % (f1_macro))
        print('    spearman-label:        %.4f %.4f' % (r1, p1))
        print('    av magnitude')
        print('        spearman_all:      %.4f %.4f' % (r2, p2))
    else:
        return accuracy, p_macro, r_macro, f1_macro, r1, p1, r2, p2


def inner_cv(av, av_element, criterion_all, X_train_all_av, X_train_all_true, y_train_all, cv_inner, model, HP):
    
    X_train_all_tuning_av = [x[av] for x in X_train_all_av]
    
    av_tuning_human_m = av_element[0][av]
    av_tuning_agent_m = av_element[1][av]
    av_tuning_human_sd = av_element[2][av]
    av_tuning_agent_sd = av_element[3][av]
    
    criterion = criterion_all[av]
    hp = HP[av]
    
    result = [model(av_tuning_human_m, av_tuning_agent_m, av_tuning_human_sd, av_tuning_agent_sd, criterion, X_train_all_tuning_av[val_ix[0]], X_train_all_true[val_ix[0]], hp)[0] for _, val_ix in cv_inner.split(X_train_all_tuning_av, y_train_all)]    
    r, p = stats.spearmanr(y_train_all, result, alternative='greater')
    s = r

    return av, s


def optimized_result(parameter, av_element, criterion_all, model, X_test_av, X_test_true, HP):
    
    av = parameter
    
    av_tuned_human_m = av_element[0][av]
    av_tuned_agent_m = av_element[1][av]
    av_tuned_human_sd = av_element[2][av]
    av_tuned_agent_sd = av_element[3][av]
    
    criterion = criterion_all[av]
    hp = HP[av]
   
    X_test_tuned_av = X_test_av[av]
    
    return model(av_tuned_human_m, av_tuned_agent_m, av_tuned_human_sd, av_tuned_agent_sd, criterion, X_test_tuned_av, X_test_true, hp)


def outer_cv(av_mode, cv_outer, cv_inner, X_av, X_true, y, av_element, criterion_all, model, train_all_ix, test_ix, HP):
    
    X_train_all_av, X_test_av = [X_av[i] for i in train_all_ix], X_av[test_ix[0]]
    X_train_all_true, X_test_true = [X_true[i] for i in train_all_ix], X_true[test_ix[0]]
    y_train_all = [y[i] for i in train_all_ix]
    score = [inner_cv(av, av_element, criterion_all, X_train_all_av, X_train_all_true, y_train_all, cv_inner, model, HP) for av in range(len(av_mode))]
    parameter = sorted(score, key=lambda x: x[1], reverse=True)[0][0]
    yHat, yM = optimized_result(parameter, av_element, criterion_all, model, X_test_av, X_test_true, HP)

    return (parameter, HP[parameter]), yHat, yM, score

         
def nested_cv(X, y, h=None, show=True, mode=None, model=None, av_mode=None):
    
    if mode == 'naive_random':

        accuracy, p_macro, r_macro, f1_macro, spearmanr = [], [], [], [], []
        for _ in tqdm(range(1000)):
            y1, y2 = [], []
            cv_outer = KFold(n_splits=len(X))
            for train_ix, test_ix in cv_outer.split(X, y):
                y_train, y_test = [y[i] for i in train_ix], [y[i] for i in test_ix]
                y_hat = np.random.randint(0, 3, size=len(y_test))
                y1.append(y_hat); y2.append(y_test)
            accuracy.append(accuracy_score(y2, y1))
            p_macro.append(precision_score(y2, y1, average='macro'))
            r_macro.append(recall_score(y2, y1, average='macro'))
            f1_macro.append(f1_score(y2, y1, average='macro'))
            spearmanr.append(stats.spearmanr(y2, y1, alternative='greater'))
        print(mode)
        print('accuracy:             %.4f±%.4f' % (np.mean(accuracy), np.std(accuracy)))
        print('precision-macro:      %.4f±%.4f' % (np.mean(p_macro), np.std(p_macro)))
        print('recall-macro:         %.4f±%.4f' % (np.mean(r_macro), np.std(r_macro)))
        print('f1-macro:             %.4f±%.4f' % (np.mean(f1_macro), np.std(f1_macro)))
        print('spearman:             %.4f±%.4f %.4f±%.4f' % (np.mean([s[0] for s in spearmanr]), np.std([s[0] for s in spearmanr]), np.mean([s[1] for s in spearmanr]), np.std([s[1] for s in spearmanr])))
        
    if mode == 'naive_probability':

        accuracy, p_macro, r_macro, f1_macro, spearmanr = [], [], [], [], []
        for _ in tqdm(range(1000)):
            y1, y2 = [], []
            cv_outer = KFold(n_splits=len(X))
            for train_ix, test_ix in cv_outer.split(X, y):
                y_train, y_test = [y[i] for i in train_ix], [y[i] for i in test_ix]
                y_hat = random.sample(y, k=len(y_test))
                y1.append(y_hat); y2.append(y_test)
            accuracy.append(accuracy_score(y2, y1))
            p_macro.append(precision_score(y2, y1, average='macro'))
            r_macro.append(recall_score(y2, y1, average='macro'))
            f1_macro.append(f1_score(y2, y1, average='macro'))
            spearmanr.append(stats.spearmanr(y2, y1, alternative='greater'))
        print(mode)
        print('accuracy:             %.4f±%.4f' % (np.mean(accuracy), np.std(accuracy)))
        print('precision-macro:      %.4f±%.4f' % (np.mean(p_macro), np.std(p_macro)))
        print('recall-macro:         %.4f±%.4f' % (np.mean(r_macro), np.std(r_macro)))
        print('f1-macro:             %.4f±%.4f' % (np.mean(f1_macro), np.std(f1_macro)))
        print('spearman:             %.4f±%.4f %.4f±%.4f' % (np.mean([s[0] for s in spearmanr]), np.std([s[0] for s in spearmanr]), np.mean([s[1] for s in spearmanr]), np.std([s[1] for s in spearmanr])))
    
    if mode == 'god':

        y1, y2 = [], []
        cv_outer = KFold(n_splits=len(X))
        for train_ix, test_ix in tqdm(cv_outer.split(X, y)):
            X_test, y_test = [X[i] for i in test_ix], [y[i] for i in test_ix]
            y_hat = X_test
            y1.append(y_hat); y2.append(y_test)
        accuracy = accuracy_score(y2, y1)
        p_macro = precision_score(y2, y1, average='macro')
        r_macro = recall_score(y2, y1, average='macro')
        f1_macro = f1_score(y2, y1, average='macro')
        r, p = stats.spearmanr(y2, y1, alternative='greater')
        print(mode)
        print('accuracy:             %.4f' % (accuracy))
        print('precision-macro:      %.4f' % (p_macro))
        print('recall-macro:         %.4f' % (r_macro))
        print('f1-macro:             %.4f' % (f1_macro))
        print('spearman:             %.4f %.4f' % (r, p))
        
    if mode in ['original', 'plm']:

        X_av, X_true = X
        
        X_av_human = [x for i, x in enumerate(X_av) if X_true[i] == 2]
        X_av_agent = [x for i, x in enumerate(X_av) if X_true[i] == 0]

        av_human_m = np.mean(X_av_human, axis=0)
        av_agent_m = np.mean(X_av_agent, axis=0)
        av_human_sd = np.std(X_av_human, axis=0)
        av_agent_sd = np.std(X_av_agent, axis=0)
        
        av_element = (av_human_m, av_agent_m, av_human_sd, av_agent_sd)

        if h == '1':
            HP = ['1' for i in range(len(av_human_m))]
        elif h == '2':
            HP = ['2' for i in range(len(av_human_m))]
        
        criterion_all = [sdt_element(X_true, y, hp=hp) for hp in HP]
        
        cv_outer = KFold(n_splits=len(y))
        cv_inner = KFold(n_splits=len(y)-1)
        
        #result = [outer_cv(av_mode, cv_outer, cv_inner, X_av, X_true, y, av_element, criterion_all, model, train_all_ix, test_ix, HP) for train_all_ix, test_ix in cv_outer.split(X_av, y)]
        
        with concurrent.futures.ProcessPoolExecutor() as executor:
            parameter = [(av_mode, cv_outer, cv_inner, X_av, X_true, y, av_element, criterion_all, model, train_all_ix, test_ix, HP) for train_all_ix, test_ix in cv_outer.split(X_av, y)]
            if show:
                result = tqdm(executor.map(outer_cv_job, parameter))
            else:
                result = executor.map(outer_cv_job, parameter)
        
        result = [r for r in result]
        score = evaluation(y, result, show=show)
        
        return result, score

    if mode == None:

        def spearmanr(y_true, y_pred):

            x = np.column_stack((y_true, y_pred))
            x_ranked = np.apply_along_axis(rankdata, 0, x)
    
            return np.corrcoef(x_ranked, rowvar=0)[0][1]

        X_score, X_true = X
    
        try:
            X = [sum(X_score[i], [])+[X_true[i]] for i in range(len(X_score))]
        except:
            X = [X_score[i]+[X_true[i]] for i in range(len(X_score))]        

        standardScaler = StandardScaler()
        standardScaler.fit(X)
        X = standardScaler.transform(X)

        cv_outer = KFold(n_splits=10)
        cv_inner = KFold(n_splits=5)

        y_true, y_pred = [], []

        if model == 'mlr':
    
            model = LogisticRegression(random_state=2022)
            space = [
                {
                    'penalty': ['none'], 
                    'class_weight': [None, 'balanced'], 
                    'max_iter': [10, 50, 100, 500, 1000, 5000, 10000], 
                    'warm_start': [True, False]
                    },
                {
                    'penalty': ['l2'], 
                    'class_weight': [None, 'balanced'], 
                    'C': [0.0001, 0.001, 0.01, 0.1, 1.0], 
                    'max_iter': [10, 50, 100, 500, 1000, 5000, 10000], 
                    'warm_start': [True, False]},
            ]

        if model == 'knn':

            model = KNeighborsClassifier()
            space = {
                'n_neighbors': [i for i in range(1, 21)],
                'weights': ['uniform', 'distance']
            }

        if model == 'rf':

            model = RandomForestClassifier(random_state=2022)
            space = {
                'n_estimators': [10, 50, 100, 500, 1000],
                'max_depth': [1, 2, 3, 4, 5, 6, 7, None],
                'warm_start': [True, False],
                'class_weight': [None, 'balanced', 'balanced_subsample'],
                'max_samples': [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
            }

        if model == 'svc':

            model = SVC(cache_size=1000, random_state=2022)
            space = [
                {
                    'C': [0.0001, 0.001, 0.01, 0.1, 1.0],
                    'kernel': ['linear', 'poly', 'rbf', 'sigmoid'],
                    'gamma': ['scale', 'auto', 1e-3, 1e-4],
                    'shrinking': [True, False],
                    'class_weight': [None, 'balanced'],
                    'max_iter': [-1, 10, 50, 100, 500, 1000, 5000, 10000]
                },
                {
                    'C': [0.0001, 0.001, 0.01, 0.1, 1.0],
                    'kernel': ['linear'],
                    'shrinking': [True, False],
                    'class_weight': [None, 'balanced'],
                    'max_iter': [-1, 10, 50, 100, 500, 1000, 5000, 10000]
                },
            ]

        if model == 'mlp':

            model = MLPClassifier(random_state=2022)
            space = [
                {
                    'hidden_layer_sizes': [(5, 5), (10,), (10, 10), (20,), (20, 20), (50,), (50, 50), (100,), (100, 100)],
                    'activation': ['identity', 'logistic', 'tanh', 'relu'],
                    'solver': ['lbfgs'],
                    'alpha': [0.0001, 0.001, 0.01, 0.1, 1.0],
                    'max_iter': [20, 50, 100, 200, 300, 400, 500],
                    'warm_start': [True, False],
                },
                {
                    'hidden_layer_sizes': [(5, 5), (10,), (10, 10), (20,), (20, 20), (50,), (50, 50), (100,), (100, 100)],
                    'activation': ['identity', 'logistic', 'tanh', 'relu'],
                    'solver': ['sgd'],
                    'alpha': [0.0001, 0.001, 0.01, 0.1, 1.0],
                    'learning_rate': ['constant', 'invscaling', 'adaptive'],
                    'learning_rate_init': [0.001, 0.002],
                    'max_iter': [20, 50, 100, 200, 300, 400, 500],
                    'shuffle': [True, False],
                    'warm_start': [True, False],
                },
                {
                    'hidden_layer_sizes': [(5, 5), (10,), (10, 10), (20,), (20, 20), (50,), (50, 50), (100,), (100, 100)],
                    'activation': ['identity', 'logistic', 'tanh', 'relu'],
                    'solver': ['adam'],
                    'alpha': [0.0001, 0.001, 0.01, 0.1, 1.0],
                    'learning_rate_init': [0.001, 0.002],
                    'max_iter': [20, 50, 100, 200, 300, 400, 500],
                    'shuffle': [True, False],
                    'warm_start': [True, False],
                }
            ]

        if model == 'xgboost':

            model = xgb.XGBClassifier(random_state=2022)
            space = {
                'n_estimators': [10, 50, 100, 500, 1000],
                'max_depth': [2, 3, 4, 5, 6, 7],
                'min_child_weight': [1, 2, 3],
                'learning_rate': [0.01, 0.05, 0.1, 0.2, 0.3],
                'booster': ['gbtree'],
                'tree_method': ['exact', 'approx', 'hist', 'prune', 'refresh', 'sync'],
                'gamma': [0, 0.1, 0.2],
                'subsample': [0.7, 1],
                'colsample_bytree': [0.75, 1],
                'reg_lambda': [0.0001, 0.001, 0.01, 0.1, 1.0],
            }
    
        for train_ix, test_ix in cv_outer.split(X, y):
        
            X_train, y_train = [X[i] for i in train_ix], [y[i] for i in train_ix]
            X_test, y_test = [X[i] for i in test_ix], [y[i] for i in test_ix]
            search = GridSearchCV(model, space, scoring=make_scorer(spearmanr), n_jobs=-1, cv=cv_inner, refit=True)
            result = search.fit(X_train, y_train)
            best_model = result.best_estimator_
            yHat = best_model.predict(X_test)
            y_true.extend(y_test); y_pred.extend(yHat)
    
        r, p = stats.spearmanr(y_true, y_pred, alternative='greater')
        return r, p


def baseline_ml_item(data, model):

    affect, trueLabel, predLabel = data
    r_aa, p_aa = nested_cv([affect, trueLabel], predLabel, model=model)
    r_aa_pre, p_aa_pre = nested_cv([[i[0] for i in affect], trueLabel], predLabel, model=model)
    r_aa_post, p_aa_post = nested_cv([[i[1] for i in affect], trueLabel], predLabel, model=model)
    r_pa, p_pa = nested_cv([[[[i[0], i[1], i[2], i[-1]] for i in j] for j in affect], trueLabel], predLabel, model=model)
    r_pa_pre, p_pa_pre = nested_cv([[i[0] for i in [[[i[0], i[1], i[2], i[-1]] for i in j] for j in affect]], trueLabel], predLabel, model=model)
    r_pa_post, p_pa_post = nested_cv([[i[1] for i in [[[i[0], i[1], i[2], i[-1]] for i in j] for j in affect]], trueLabel], predLabel, model=model)
    r_na, p_na = nested_cv([[[[i[3], i[4]] for i in j] for j in affect], trueLabel], predLabel, model=model)
    r_na_pre, p_na_pre = nested_cv([[i[0] for i in [[[i[3], i[4]] for i in j] for j in affect]], trueLabel], predLabel, model=model)
    r_na_post, p_na_post = nested_cv([[i[1] for i in [[[i[3], i[4]] for i in j] for j in affect]], trueLabel], predLabel, model=model)

    with open('/'.join(os.getcwd().split('/')[:-1])+'./table2/ml_baselines/'+model+'.txt', 'a+') as f:

        print('         AA:      %.4f %.4f' % (r_aa, p_aa), file=f)
        print('         AA_pre:  %.4f %.4f' % (r_aa_pre, p_aa_pre), file=f)
        print('         AA_post: %.4f %.4f' % (r_aa_post, p_aa_post), file=f)
        print('         PA:      %.4f %.4f' % (r_pa, p_pa), file=f)
        print('         PA_pre:  %.4f %.4f' % (r_pa_pre, p_pa_pre), file=f)
        print('         PA_post: %.4f %.4f' % (r_pa_post, p_pa_post), file=f)
        print('         NA:      %.4f %.4f' % (r_na, p_na), file=f)
        print('         NA_pre:  %.4f %.4f' % (r_na_pre, p_na_pre), file=f)
        print('         NA_post: %.4f %.4f' % (r_na_post, p_na_post), file=f)
        f.close()


def baseline_ml(model):

    with open('/'.join(os.getcwd().split('/')[:-1])+'/table2/ml_baselines/'+model+'.txt', 'a+') as f:

        print(model.upper(), file=f)
        print('     First stage:', file=f)
        f.close()

    baseline_ml_item([firstAffect, firstTrueLabel, firstPredLabel], model)

    with open('/'.join(os.getcwd().split('/')[:-1])+'/table2/ml_baselines/'+model+'.txt', 'a+') as f:

        print('--------------------------------', file=f)
        print('     Second stage:', file=f)
        f.close()

    baseline_ml_item([secondAffect, secondTrueLabel, secondPredLabel], model)

    with open('/'.join(os.getcwd().split('/')[:-1])+'/table2/ml_baselines/'+model+'.txt', 'a+') as f:

        print('--------------------------------', file=f)
        print('     Third stage:', file=f)
        f.close()

    baseline_ml_item([thirdAffect, thirdTrueLabel, thirdPredLabel], model)

    with open('/'.join(os.getcwd().split('/')[:-1])+'/table2/ml_baselines/'+model+'.txt', 'a+') as f:

        print('--------------------------------', file=f)
        print('     All stages:', file=f)
        f.close()

    baseline_ml_item([allAffect, allTrueLabel, allPredLabel], model)

    with open('/'.join(os.getcwd().split('/')[:-1])+'/table2/ml_baselines/'+model+'.txt', 'a+') as f:

        print(' ', file=f)
        f.close()