import os
import sys
import glob
import numpy as np
import json
import deepdish as dd
import pandas as pd
import copy
from collections import defaultdict
from sklearn.model_selection import cross_val_score
from sklearn.metrics import mean_squared_error
import connectivity.io as cio
from connectivity import data as cdata
import connectivity.constants as const
import connectivity.model as model
import connectivity.evaluation as ev
import connectivity.nib_utils as nio
import warnings
warnings.filterwarnings("ignore")
np.seterr(divide="ignore", invalid="ignore")
"""Main module for training and evaluating connectivity models.
@authors: Maedbh King, Jörn Diedrichsen
Typical usage example:
config = get_default_train_config()
models = train_models(config=config)
config = get_default_eval_config()
models = eval_models(config)
"""
def get_default_train_config():
"""Defaults for training model(s).
Returns:
A dict mapping keys to training parameters.
"""
# defaults training config:
config = {
"name": "L2_WB162_A1", # Model name - determines the directory
"model": "L2regression", # Model class name (must be in model.py)
"param": {"alpha": 1}, # Parameter to model constructor
"sessions": [1, 2], # Sessions used for training data
"glm": "glm7", # GLM used for training data
"train_exp": "sc1", # Experiment used for training data
"averaging": "sess", # Averaging scheme for X and Y (see data.py)
"weighting": True, # "none" "full" "diag"
# "incl_inst": True,
"X_data": "tessels0162",
"Y_data": "cerebellum_suit",
"validate_model": True,
"cv_fold": None, #TO IMPLEMENT: "sess", "run" (None is "tasks")
"subjects": [
"s01",
"s03",
"s04",
"s06",
"s08",
"s09",
"s10",
"s12",
"s14",
"s15",
"s17",
"s18",
"s19",
"s20",
"s21",
"s22",
"s24",
"s25",
"s26",
"s27",
"s28",
"s29",
"s30",
"s31",
],
"mode": "crossed",
"save_weights": False, #Training mode
}
return config
def get_default_eval_config():
"""Defaults for evaluating model(s).
Returns:
A dict mapping keys to evaluation parameters.
"""
config = {
"name": "L2_tessels1002_A1",
"sessions": [1, 2],
"glm": "glm7",
"train_exp": "sc1",
"eval_exp": "sc2",
"averaging": "sess",
"weighting": True, # 0: none, 1: by regr., 2: by full matrix???
"incl_inst": True,
"X_data": "tessels1002",
"Y_data": "cerebellum_suit",
"subjects": [
"s01",
"s03",
"s04",
"s06",
"s08",
"s09",
"s10",
"s12",
"s14",
"s15",
"s17",
"s18",
"s19",
"s20",
"s21",
"s22",
"s24",
"s25",
"s26",
"s27",
"s28",
"s29",
"s30",
"s31",
],
"mode": "crossed",
"splitby": None,
"save_maps": False,
"threshold": 0.1
}
return config
def train_metrics(model, X, Y):
"""computes training metrics (rmse and R) on X and Y
Args:
model (class instance): must be fitted model
X (nd-array):
Y (nd-array):
Returns:
rmse_train (scalar), R_train (scalar)
"""
Y_pred = model.predict(X)
# get train rmse and R
rmse_train = mean_squared_error(Y, Y_pred, squared=False)
R_train, _ = ev.calculate_R(Y, Y_pred)
return rmse_train, R_train
def validate_metrics(model, X, Y, X_info, cv_fold):
"""computes CV training metrics (rmse and R) on X and Y
Args:
model (class instance): must be fitted model
X (nd-array):
Y (nd-array):
cv_fold (int): number of CV folds
Returns:
rmse_cv (scalar), R_cv (scalar)
"""
# get cv rmse and R
rmse_cv_all = np.sqrt(cross_val_score(model, X, Y, scoring="neg_mean_squared_error", cv=cv_fold) * -1)
# TO DO: implement train/validate splits for "sess", "run"
r_cv_all = cross_val_score(model, X, Y, scoring=ev.calculate_R_cv, cv=cv_fold)
return np.nanmean(rmse_cv_all), np.nanmean(r_cv_all)
def eval_models(config):
"""Evaluates a specific model class on X and Y data from a specific experiment for subjects listed in config.
Args:
config (dict): Evaluation configuration, returned from get_default_eval_config()
Returns:
models (pd dataframe): evaluation of different models on the data
"""
eval_all = defaultdict(list)
eval_voxels = defaultdict(list)
for idx, subj in enumerate(config["subjects"]):
print(f"Evaluating model on {subj}")
# get data
Y, Y_info, X, X_info = _get_data(config=config, exp=config["eval_exp"], subj=subj)
# Get the model from file
fname = _get_model_name(train_name=config["name"], exp=config["train_exp"], subj_id=subj)
fitted_model = dd.io.load(fname)
# Get model predictions
Y_pred = fitted_model.predict(X)
if config["mode"] == "crossed":
Y_pred = np.r_[Y_pred[Y_info.sess == 2, :], Y_pred[Y_info.sess == 1, :]]
# get rmse
rmse = mean_squared_error(Y, Y_pred, squared=False)
data = {"rmse_eval": rmse, "subj_id": subj}
# Copy over all scalars or strings to eval_all dataframe:
for key, value in config.items():
if type(value) is not list:
data.update({key: value})
# add evaluation (summary)
evals = _get_eval(Y=Y, Y_pred=Y_pred, Y_info=Y_info, X_info=X_info)
data.update(evals)
# # add sparsity metric (voxels)
# sparsity_results = _get_sparsity(config, fitted_model)
# data.update(sparsity_results)
# add evaluation (voxels)
if config["save_maps"]:
for k, v in data.items():
if "vox" in k:
eval_voxels[k].append(v)
# don't save voxel data to summary
data = {k: v for k, v in data.items() if "vox" not in k}
# append data for each subj
for k, v in data.items():
eval_all[k].append(v)
# Return list of models
return pd.DataFrame.from_dict(eval_all), eval_voxels
def _get_eval(Y, Y_pred, Y_info, X_info):
"""Compute evaluation, returning summary and voxel data.
Args:
Y (np array):
Y_pred (np array):
Y_info (pd dataframe):
X_info (pd dataframe):
Returns:
dict containing evaluations (R, R2, noise).
"""
# initialise dictionary
data = {}
# Add the evaluation
data["R_eval"], data["R_vox"] = ev.calculate_R(Y=Y, Y_pred=Y_pred)
# R between predicted and observed
data["R2"], data["R2_vox"] = ev.calculate_R2(Y=Y, Y_pred=Y_pred)
# R2 between predicted and observed
(
data["noise_Y_R"],
data["noise_Y_R_vox"],
data["noise_Y_R2"],
data["noise_Y_R2_vox"],
) = ev.calculate_reliability(Y=Y, dataframe=Y_info)
# Noise ceiling for cerebellum (squared)
(
data["noise_X_R"],
data["noise_X_R_vox"],
data["noise_X_R2"],
data["noise_X_R2_vox"],
) = ev.calculate_reliability(Y=Y_pred, dataframe=X_info)
# calculate noise ceiling
data["noiseceiling_Y_R_vox"] = np.sqrt(data["noise_Y_R_vox"])
data["noiseceiling_XY_R_vox"] = np.sqrt(data["noise_Y_R_vox"] * np.sqrt(data["noise_X_R_vox"]))
# # Noise ceiling for cortex (squared)
# pass
return data
def _get_data(config, exp, subj):
"""get X and Y data for exp and subj
Args:
config (dict): must contain keys for glm, Y_data, X_data, averaging, weighting
exp (str): 'sc1' or 'sc2'
subj (str): default subjs are set in constants.py
Returns:
Y (nd array), Y_info (pd dataframe), X (nd array), X_info (pd dataframe)
"""
# Get the data
Ydata = cdata.Dataset(
experiment=exp,
glm=config["glm"],
subj_id=subj,
roi=config["Y_data"],
)
# load mat
Ydata.load_mat()
Y, Y_info = Ydata.get_data(averaging=config["averaging"], weighting=config["weighting"])
Xdata = cdata.Dataset(
experiment=exp,
glm=config["glm"],
subj_id=subj,
roi=config["X_data"],
)
# load mat
Xdata.load_mat()
X, X_info = Xdata.get_data(averaging=config["averaging"], weighting=config["weighting"])
return Y, Y_info, X, X_info
def _get_model_name(train_name, exp, subj_id):
"""returns path/name for connectivity training model outputs.
Args:
train_name (str): Name of trained model
exp (str): Experiment name
subj_id (str): Subject id
Returns:
fpath (str): Full path and name to connectivity output for model training.
"""
dirs = const.Dirs(exp_name=exp)
fname = f"{train_name}_{subj_id}.h5"
fpath = os.path.join(dirs.conn_train_dir, train_name, fname)
return fpath
# functions for running and checking winner-n-take-all models
def select_winners(config, winner_model = None, save = True):
dirs = const.Dirs(exp_name=config["train_exp"], glm=config["glm"])
# Store the training configuration in model directory
if save:
fpath = os.path.join(dirs.conn_train_dir, config["name"])
cio.make_dirs(fpath)
# not saving the config
# cio.save_dict_as_JSON(os.path.join(fpath, "train_config.json"), config)
# Loop over subjects and train
for subj in config["subjects"]:
print(f"selecting winner features for {subj}")
# get data
Y, Y_info, X, X_info = _get_data(config=config, exp=config["train_exp"], subj=subj)
# Generate new model and put in the list
new_model = getattr(model, config["model"])(**config["param"])
# cross the sessions
if config["mode"] == "crossed":
Y = np.r_[Y[Y_info.sess == 2, :], Y[Y_info.sess == 1, :]]
# Fit model, get train and validate metrics
if winner_model is None:
support_ = None
else:
support_ = winner_model.support_
new_model.set_support_(X, Y, support_ = support_)
# new_model = new_model.select_winners(X, Y, winner_model = winner_model)
# Save the class of selected regressors
if save:
fname = _get_model_name(train_name=config["name"], exp=config["train_exp"], subj_id=subj)
dd.io.save(fname, new_model, compression=None)
# return the model to be used for the next model
return new_model
def train_wnta(config, save = True):
"""Trains a winner-n-take-all model class on X and Y data from a specific experiment for subjects listed in config.
Args:
config (dict): Training configuration, returned from get_default_train_config()
save (bool): Optional; Save fitted models automatically to disk.
Returns:
models (list): list of trained models for subjects listed in config.
train_all (pd dataframe): dataframe containing
"""
dirs = const.Dirs(exp_name=config["train_exp"], glm=config["glm"])
models = []
train_all = defaultdict(list)
# Store the training configuration in model directory
if save:
fpath = os.path.join(dirs.conn_train_dir, config["name"])
cio.make_dirs(fpath)
cio.save_dict_as_JSON(os.path.join(fpath, "train_config.json"), config)
# path to folder with n-winner feature models
n = config['param']['n_features_to_select']
winner_name = f"winners_{config['X_data']}_N{n:.0f}"
# Loop over subjects and train
for subj in config["subjects"]:
print(f"Training model on {subj}")
# get data
Y, Y_info, X, X_info = _get_data(config=config, exp=config["train_exp"], subj=subj)
# Generate new model and put in the list
new_model = getattr(model, config["model"])(**config["param"])
# cross the sessions
if config["mode"] == "crossed":
Y = np.r_[Y[Y_info.sess == 2, :], Y[Y_info.sess == 1, :]]
# load in the Winner-n model for the n parameter
winner_path = _get_model_name(train_name=winner_name, exp=config["train_exp"], subj_id=subj)
winner_model = dd.io.load(winner_path)
# update the winner_model attribute of the wnta model
new_model.winner_model = winner_model
models.append(new_model)
# fit the model
models[-1].fit(X, Y)
models[-1].rmse_train, models[-1].R_train = train_metrics(models[-1], X, Y)
# collect train metrics (rmse and R)
data = {
"subj_id": subj,
"rmse_train": models[-1].rmse_train,
"R_train": models[-1].R_train
}
# run cross validation and collect metrics (rmse and R)
if config['validate_model']:
models[-1].rmse_cv, models[-1].R_cv = validate_metrics(models[-1], X, Y, X_info, config["cv_fold"])
data.update({"rmse_cv": models[-1].rmse_cv,
"R_cv": models[-1].R_cv
})
# Copy over all scalars or strings from config to eval dict:
for key, value in config.items():
if not isinstance(value, (list, dict)):
data.update({key: value})
for k, v in data.items():
train_all[k].append(v)
# Save the fitted model to disk if required
if save:
fname = _get_model_name(train_name=config["name"], exp=config["train_exp"], subj_id=subj)
dd.io.save(fname, models[-1], compression=None)
return models, pd.DataFrame.from_dict(train_all)
def load_model(config):
dirs = const.Dirs(exp_name=config["train_exp"], glm=config["glm"])
models = []
train_all = defaultdict(list)
# path to the models
# fpath = os.path.join(dirs.conn_train_dir, config["name"])
for subj in config["subjects"]:
# Get the model from file
fname = _get_model_name(train_name=config["name"], exp=config["train_exp"], subj_id=subj)
fitted_model = dd.io.load(fname)
models.append(fitted_model)
return models