PredictKeyTemps.py

import pickle

from sklearn.svm import SVR
from sklearn import linear_model
from sklearn.preprocessing import StandardScaler
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.ensemble import AdaBoostRegressor
from sklearn.ensemble import RandomForestRegressor
from sklearn.neighbors import KNeighborsRegressor
from sklearn.model_selection import RandomizedSearchCV, KFold
from sklearn.feature_selection import RFECV
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error, mean_absolute_percentage_error
import json
import numpy as np
import re
from cubist import Cubist
import matplotlib.pyplot as plt
import seaborn as sb
from xgboost import XGBRegressor
from sklearn.neural_network import MLPRegressor
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)
import pandas as pd


def read_x_data(x_data_file):
    print("Reading X data")
    x_df = pd.read_csv(x_data_file, index_col=0)
    x_matrix = x_df.to_numpy()
    x_ids = x_df.index.values
    x_vars = x_df.columns.values
    return x_matrix, x_ids, x_vars


def read_y_data(y_data_file):
    print("Reading Y data")
    y_df = pd.read_csv(y_data_file, index_col=1)
    datasets = y_df['dataset']
    y_df.drop(['dataset'], axis=1, inplace=True)
    y_matrix = y_df.to_numpy()
    y_ids = y_df.index.values
    y_vars = y_df.columns.values
    return y_matrix, y_ids, y_vars, datasets


def prepare_y_data(y_data, y_ids, x_ids, datasets):
    print("Preparing Y data")

    # find IDs that match with UniProt ID regex pattern
    y_ids = [find_uniprot_id(x) for x in y_ids]

    # re-arrange y-data according to ordering of IDs in x-data
    match_idx = [y_ids.index(x) for x in x_ids if x in y_ids]
    y_data = y_data[match_idx, :]
    y_ids = [y_ids[i] for i in match_idx]
    datasets = [datasets[i] for i in match_idx]
    return y_data, y_ids, datasets


def prepare_x_data(x_data, x_ids, y_ids):
    print("Preparing X data")
    match_idx = [np.where(x_ids == x)[0][0] for x in y_ids if x in x_ids]
    x_data = x_data[match_idx, :]
    x_ids = x_ids[match_idx]
    return x_data, x_ids


def prepare_data(x_data_file, y_data_file, t_r2=0.6):
    x_matrix, x_ids, x_vars = read_x_data(x_data_file)
    y_matrix, y_ids, y_vars, datasets = read_y_data(y_data_file)

    y_matrix, y_ids, datasets = prepare_y_data(y_matrix, y_ids, x_ids, datasets)
    x_matrix, x_ids = prepare_x_data(x_matrix, x_ids, y_ids)

    # remove rows with nan values and with insufficient R2 (< 0.6)
    print("Cleaning datasets")
    nan_row_idx = [np.isnan(x_matrix[row, :]).any() or np.isnan(y_matrix[row, :]).any()
                   for row in range(0, x_matrix.shape[0])]

    if t_r2 is not None and t_r2 < 1:
        r2_col_idx = [x == "r2_adj" for x in y_vars]
        suff_r2_idx = [x >= t_r2 for x in y_matrix[:, r2_col_idx]]
        keep_row_idx = [i for i in range(0, len(nan_row_idx)) if suff_r2_idx[i] and not nan_row_idx[i]]
    else:
        keep_row_idx = [i for i in range(0, len(nan_row_idx)) if not nan_row_idx[i]]

    y_matrix = y_matrix[keep_row_idx, :]
    y_df = pd.DataFrame(y_matrix, columns=y_vars)

    x_matrix = x_matrix[keep_row_idx, :]
    x_df = pd.DataFrame(x_matrix, columns=x_vars)

    datasets = [datasets[i] for i in keep_row_idx]

    return x_df, y_df, datasets


def write_clean_data(x_data, y_data, datasets, x_filename, y_filename, group_filename):
    print("Writing clean data to file")
    x_data.to_csv(x_filename, index=False)
    y_data.to_csv(y_filename, index=False)
    with open(group_filename, "w", newline="\n") as f:
        for line in datasets:
            f.write(line + "\n")


def read_clean_data(x_filename, y_filename, group_filename):
    print("Reading data from file")
    x_data = pd.read_csv(x_filename)
    y_data = pd.read_csv(y_filename)
    groups = []
    with open(group_filename, "r") as f:
        for line in f:
            groups.append(line.strip("\n"))
    return x_data, y_data, groups


def remove_y_outliers_per_group(x_data, y_data, groups):
    print("Removing outliers for each group (outside median +/- 3*sd)")
    uniq_groups = set(groups)
    for g in uniq_groups:
        group_idx = np.asarray([i for i in range(0, len(groups)) if groups[i] == g])
        med, sd = np.median(y_data.iloc[group_idx]), np.std(y_data.iloc[group_idx])
        sd_threshold = 3 * sd
        t_lower, t_upper = med - sd_threshold, med + sd_threshold
        remove_idx = np.where(np.logical_or((t_lower > y_data.iloc[group_idx]), (y_data.iloc[group_idx] > t_upper)))[0]
        y_data.drop(group_idx[remove_idx], inplace=True)
        y_data.reset_index(inplace=True, drop=True)
        x_data.drop(group_idx[remove_idx], axis=0, inplace=True)
        x_data.reset_index(inplace=True, drop=True)

        for i in sorted(remove_idx, reverse=True):
            del groups[group_idx[i]]

    return x_data, y_data, groups


def split_data_train_test(x_data, y_data, pct_test=40):
    X_train, X_test, y_train, y_test, idx_train, idx_test = train_test_split(x_data, y_data, range(0, len(y_data)),
                                                                             test_size=pct_test/100, random_state=42)
    return X_train, X_test, y_train, y_test, idx_test


def get_scaler(X_train, method="standard"):
    if method == "standard":
        scaler = StandardScaler().fit(X_train)
    elif method == "minmax":
        scaler = MinMaxScaler().fit(X_train)
    else:
        scaler = None
    return scaler


def feature_selection_rfecv(X, y, regr, n_jobs=1, step=1, cv=5, rfecv_result_file="feature_selection/rfecv_results.csv",
                            feature_support_file="feature_selection/feature_ranking_support.csv"):
    print("Running RFECV")
    # print("JOBS=%d, STEP=%d, CV=%d" % (n_jobs, step, cv))
    rfecv = RFECV(
        estimator=regr,
        step=step,
        cv=cv,
        n_jobs=n_jobs,
        scoring=r2_scorer,
        min_features_to_select=1,
        verbose=2
    )
    rfecv = rfecv.fit(X, y)

    res_df = pd.DataFrame(rfecv.cv_results_)
    res_df.to_csv(rfecv_result_file)

    feat_dict = {"ranking": rfecv.ranking_,
                 "support": rfecv.support_}
    feat_df = pd.DataFrame(feat_dict)
    feat_df.to_csv(feature_support_file)

    print(str(rfecv.n_features_) + " features selected.")

    return rfecv


def plot_rfecv_scores(rfecv_scores_file, n, step):
    rfecv_scores = pd.read_csv(rfecv_scores_file)
    av = rfecv_scores['mean_test_score'][1:]
    sd = rfecv_scores['std_test_score'][1:]
    n_features = list(range(n, 0, -step))
    plt.plot(n_features, av)
    plt.fill_between(n_features, av-sd, av+sd, alpha=0.3, facecolor="blue")
    plt.xlabel("Number of features")
    plt.ylabel("$R^2 (5x CV)$")
    plt.savefig("feature_selection/rfecv_scores.png")


def write_rfecv_features(rfecv_feature_support_file, feature_name_file):
    fs = pd.read_csv(rfecv_feature_support_file)
    fn = pd.read_csv(feature_name_file, header=None, names=["Feature"])
    fs = pd.concat([fn, fs], axis=1)
    fs.sort_values(by="ranking", inplace=True)
    fs.to_csv("feature_selection/feature_ranking_rfecv_sorted.csv", index=False, lineterminator="\n")


def write_selected_features(rfecv_feature_support_file, feature_input_file, feature_output_file):
    fs = pd.read_csv(rfecv_feature_support_file)
    feature_df = pd.read_csv(feature_input_file)
    feature_df = feature_df.iloc[:, np.append(0, np.where(fs['support'])[0]+1)]
    feature_df.to_csv(feature_output_file, lineterminator="\n", index=False)

def r2_scorer(estimator, X_test, y_test):
    y_pred = estimator.predict(X_test)
    return r2_score(y_test, y_pred)


def score_prediction(y_pred, y_test):
    # RMSE
    print("%.2f" % np.sqrt(mean_squared_error(y_test, y_pred)), end="\t")
    # MAE
    print("%.2f" % np.sqrt(mean_absolute_error(y_test, y_pred)), end="\t")
    # MAPE
    print("%.2f" % np.sqrt(mean_absolute_percentage_error(y_test, y_pred)), end="\t")
    # R2
    print("%.2f" % r2_score(y_test, y_pred), end="\t")
    # rhoP
    print("%.2f" % np.corrcoef(y_test, y_pred)[0][1])


def scatter_test_pred(y_pred, y_test, groups=None, fname="topt_scatter_plots/scatterplot.png"):
    df = pd.DataFrame({"y_test": y_test, "y_pred": y_pred, "groups": groups})
    ax = sb.scatterplot(data=df, x="y_test", y="y_pred", hue="groups")
    ax.set_ylim((0, 100))
    ax.set_xlim((0, 100))
    ax.set_xlabel("Test")
    ax.set_ylabel("Prediction")

    sb.move_legend(ax, "upper left", bbox_to_anchor=(1, 1))

    f = plt.gcf()
    f.savefig(fname, bbox_inches='tight')
    plt.close(f)


def perform_regression(regr, x_data, y_data, groups=None, plot=False, fout="topt_scatter_plots/scatterplot.png"):
    X_train, X_test, y_train, y_test, idx_test = split_data_train_test(x_data, y_data)

    scaler = get_scaler(X_train)
    X_train_transformed = pd.DataFrame(scaler.transform(X_train), columns=X_train.columns)
    regr.fit(X_train_transformed, y_train)
    X_test_transformed = pd.DataFrame(scaler.transform(X_test), columns=X_test.columns)
    y_pred = regr.predict(X_test_transformed)
    score_prediction(y_pred, y_test)

    if plot:
        groups = [groups[i] for i in idx_test]
        scatter_test_pred(y_pred, y_test, groups, fout)

    return regr


def svr_rbf(x_data, y_data):
    print("SVR (RBF kernel)", end="\t")
    regr = SVR(kernel="rbf")
    regr = perform_regression(regr, x_data, y_data)
    return regr


def svr_linear(x_data, y_data):
    print("SVR (linear kernel)", end="\t")
    regr = SVR(kernel="linear")
    regr = perform_regression(regr, x_data, y_data)
    return regr


def svr_poly(x_data, y_data):
    print("SVR (polynomial kernel)", end="\t")
    regr = SVR(kernel="poly")
    regr = perform_regression(regr, x_data, y_data)
    return regr


def ridge(x_data, y_data, alpha=.5):
    print("Ridge (alpha=" + str(alpha) + ")", end="\t")
    regr = linear_model.Ridge(alpha=alpha)
    regr = perform_regression(regr, x_data, y_data)
    return regr


def lasso(x_data, y_data, alpha=.01):
    print("Lasso (alpha=" + str(alpha) + ")", end="\t")
    regr = linear_model.Lasso(alpha=alpha)
    regr = perform_regression(regr, x_data, y_data)
    return regr


def gbdt(x_data, y_data, groups=None):
    print("GBDT", end="\t")
    regr = GradientBoostingRegressor()
    regr = perform_regression(regr, x_data, y_data, plot=True, fout="topt_scatter_plots/scatter_gbdt.png", groups=groups)
    return regr


def adaboost(x_data, y_data):
    print("AdaBoost", end="\t")
    regr = AdaBoostRegressor(random_state=42)
    regr = perform_regression(regr, x_data, y_data)
    return regr


def knn(x_data, y_data, groups=None, k=10):
    print("KNN", end="\t")
    regr = KNeighborsRegressor(n_neighbors=k)
    regr = perform_regression(regr, x_data, y_data, plot=True, fout="topt_scatter_plots/scatter_knn.png", groups=groups)
    return regr


def cubist_reg(x_data, y_data, groups=None):
    print("Cubist", end="\t")
    regr = Cubist()
    regr = perform_regression(regr, x_data, y_data, plot=True, fout="topt_scatter_plots/scatter_cubist.png", groups=groups)
    return regr


def bayesian_ridge(x_data, y_data, groups=None):
    print("Bayesian ridge", end="\t")
    regr = linear_model.BayesianRidge()
    regr = perform_regression(regr, x_data, y_data, plot=True, fout="topt_scatter_plots/scatter_bayesian_ridge.png", groups=groups)
    return regr


def xgboost_reg(x_data, y_data, groups=None):
    print("XGBoost", end="\t")
    regr = XGBRegressor()
    regr = perform_regression(regr, x_data, y_data, plot=True, fout="topt_scatter_plots/scatter_xgboost.png", groups=groups)
    return regr


def mlpreg(x_data, y_data, groups=None):
    print("MLP", end="\t")
    regr = MLPRegressor()
    regr = perform_regression(regr, x_data, y_data, plot=True, fout="topt_scatter_plots/scatter_mlp.png", groups=groups)
    return regr


def randomforest(x_data, y_data, groups=None, optimal=True, hyperparam=False, fselect=False, n_jobs=1,
                 rfecv_result_file="feature_selection/rfecv_results.csv",
                 feature_support_file="feature_selection/feature_ranking_support.csv",
                 x_scaler_file="model/x_scaler.pkl"):
    X_train, X_test, y_train, y_test, idx_test = split_data_train_test(x_data, y_data)
    groups = [groups[i] for i in idx_test]
    scaler = get_scaler(X_train)
    pickle.dump(scaler, open(x_scaler_file, 'wb'))
    X_train_transformed = pd.DataFrame(scaler.transform(X_train), columns=X_train.columns)
    X_test_transformed = pd.DataFrame(scaler.transform(X_test), columns=X_test.columns)

    if optimal:
        hyperparam = False
        fselect = False

        # read optimal parameters from file
        opt_params = json.load(open("hyperparameter_tuning/rf_optimal_hyperparams.json", "r"))
        regr = RandomForestRegressor(**opt_params, random_state=42, n_jobs=n_jobs)
    else:
        regr = RandomForestRegressor(random_state=42, n_jobs=n_jobs)

    if fselect:
        hyperparam = False
        scaler = get_scaler(x_data)
        x_data_transformed = scaler.transform(x_data)
        kfsplit = KFold(n_splits=5, shuffle=True, random_state=42)
        feature_selection_rfecv(x_data_transformed, y_data, regr, n_jobs=n_jobs, step=10, cv=kfsplit,
                                rfecv_result_file=rfecv_result_file,
                                feature_support_file=feature_support_file)

    if hyperparam:
        print("\nHyperparameter fitting")
        # Number of trees in random forest
        n_estimators = [int(x) for x in np.linspace(start=200, stop=2000, num=10)]
        # Number of features to consider at every split
        max_features = ['auto', 'sqrt']
        # Maximum number of levels in tree
        max_depth = [int(x) for x in np.linspace(10, 150, num=5)]
        max_depth.append(None)
        # Minimum number of samples required to split a node
        min_samples_split = [2, 4, 6, 8]
        # Minimum number of samples required at each leaf node
        min_samples_leaf = [1, 2, 4]
        # Method of selecting samples for training each tree
        bootstrap = [True, False]
        # Create the random grid
        random_grid = {'n_estimators': n_estimators,
                       'max_features': max_features,
                       'max_depth': max_depth,
                       'min_samples_split': min_samples_split,
                       'min_samples_leaf': min_samples_leaf,
                       'bootstrap': bootstrap}

        kfsplit = KFold(n_splits=5, shuffle=True, random_state=42)
        rf_random = RandomizedSearchCV(estimator=regr, param_distributions=random_grid, n_iter=100, cv=kfsplit, verbose=2,
                                       random_state=42, n_jobs=n_jobs)
        rf_random.fit(X_train_transformed, y_train)
        print(rf_random.best_params_)
        print("\nRandom Forest (optimized)", end="\t")
        y_pred = rf_random.best_estimator_.predict(X_test_transformed)
        score_prediction(y_pred, y_test)
        scatter_test_pred(y_pred, y_test, groups, 'topt_scatter_plots/scatter_rf_optimized.png')

        # write optimal parameters to file
        json.dump(rf_random.best_params_, open("hyperparameter_tuning/rf_optimal_hyperparams.json", "w"))

    print("Random Forest", end="\t")
    regr.fit(X_train_transformed, y_train)

    y_pred = regr.predict(X_test_transformed)
    score_prediction(y_pred, y_test)
    scatter_test_pred(y_pred, y_test, groups, 'topt_scatter_plots/scatter_rf.png')

    return regr


def find_uniprot_id(prot_id):
    match = re.search(r"[OPQ][0-9][A-Z0-9]{3}[0-9]|[A-NR-Z][0-9]([A-Z][A-Z0-9]{2}[0-9]){1,2}", str(prot_id))
    if match is not None:
        return match[0]
    else:
        return match


def assert_nan(a):
    for i in range(0, len(a)):
        if not np.isnan(a[i]):
            return True
    return False