Merge pull request #4 from mckenziephagen/save_bic

Save bic

Author: mckenziephagen

.gitignore

@@ -0,0 +1,3 @@
+*ipynb* 
+*cache*
+*egg*  

src/model_fc/models.py

@@ -1,28 +1,27 @@
 import numpy as np
-from mpi4py import MPI
 from nilearn.connectome import ConnectivityMeasure
 from pyuoi.linear_model import UoI_Lasso
-from sklearn.linear_model import ElasticNetCV, LassoCV, LassoLarsIC
+from sklearn.linear_model import ElasticNetCV, LassoCV, LassoLarsIC, RidgeCV
 from sklearn.metrics import r2_score
 
 
 def run_model(train_ts, test_ts, n_rois, model, **kwargs):
     """Calculate a model based functional connectivity matrix.
 
-
     train_ts: training timeseries
     test_ts: testing timeseries
     n_rois: number of rois in parcellation
     model: model object
 
-
     """
     assert train_ts.shape[1] == n_rois == test_ts.shape[1]
-    fc_mat = np.zeros((n_rois, n_rois))
+    fc_mat = np.empty((n_rois, n_rois))
 
-    inner_rsq_dict = {"train": [], "test": []}
+    results_dict = {}
 
     for target_idx in range(train_ts.shape[1]):
+        results_dict[f"node_{target_idx}"] = {}
+        print(f"*****ECHO {target_idx}***********")
         y_train = np.array(train_ts[:, target_idx])
         X_train = np.delete(train_ts, target_idx, axis=1)
 
@@ -31,13 +30,15 @@ def run_model(train_ts, test_ts, n_rois, model, **kwargs):
 
         model.fit(X=X_train, y=y_train)
 
-        fc_mat[target_idx, :] = np.insert(model.coef_, target_idx, 0)
+        fc_mat[target_idx, :] = np.insert(model.coef_, target_idx, 1)
         test_rsq, train_rsq = eval_metrics(X_train, y_train, X_test, y_test, model)
 
-        inner_rsq_dict["test"].append(test_rsq)
-        inner_rsq_dict["train"].append(train_rsq)
+        results_dict[f"node_{target_idx}"]["train_r2"] = train_rsq
+        results_dict[f"node_{target_idx}"]["test_r2"] = test_rsq
+
+    results_dict["fc_matrix"] = fc_mat
 
-    return (fc_mat, inner_rsq_dict, model)
+    return results_dict
 
 
 def eval_metrics(X_train, y_train, X_test, y_test, model):
@@ -53,25 +54,19 @@ def init_model(
     model_str, max_iter, random_state, stability_selection=16, selection_frac=0.7
 ):
     """Initialize model object for FC calculations."""
-    if model_str == "uoi-lasso":
+    if model_str == "uoiLasso":
         uoi_lasso = UoI_Lasso(estimation_score="BIC")
-        comm = MPI.COMM_WORLD
-
         uoi_lasso.selection_frac = selection_frac
         uoi_lasso.stability_selection = stability_selection
         uoi_lasso.copy_X = True
         uoi_lasso.estimation_target = None
         uoi_lasso.logger = None
-        uoi_lasso.warm_start = False
-        uoi_lasso.comm = comm
-        uoi_lasso.random_state = 1
+        uoi_lasso.warm_start = True
         uoi_lasso.n_lambdas = 100
         uoi_lasso.max_iter = max_iter
-        uoi_lasso.random_state = random_state
-
         model = uoi_lasso
 
-    elif model_str == "lasso-cv":
+    elif model_str == "lassoCV":
         lasso = LassoCV(
             fit_intercept=True,
             cv=5,
@@ -81,8 +76,12 @@ def init_model(
         )
 
         model = lasso
+    elif model_str == "ridgeCV":
+        ridge = RidgeCV(fit_intercept=True)
+
+        model = ridge
 
-    elif model_str == "lasso-bic":
+    elif model_str == "lassoBIC":
         lasso = LassoLarsIC(criterion="bic", fit_intercept=True, max_iter=max_iter)
 
         model = lasso