add changes to hyak

Author: mckenziephagen

src/model_fc/models.py

@@ -1,5 +1,4 @@
 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, RidgeCV
@@ -14,7 +13,6 @@ def run_model(train_ts, test_ts, n_rois, model, **kwargs):
     n_rois: number of rois in parcellation
     model: model object
 
-
     """
     assert train_ts.shape[1] == n_rois == test_ts.shape[1]
     fc_mat = np.empty((n_rois, n_rois))
@@ -23,7 +21,7 @@ def run_model(train_ts, test_ts, n_rois, model, **kwargs):
 
     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)
 
@@ -35,7 +33,6 @@ def run_model(train_ts, test_ts, n_rois, model, **kwargs):
         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)
 
-        results_dict[f"node_{target_idx}"]["model"] = model
         results_dict[f"node_{target_idx}"]["train_r2"] = train_rsq
         results_dict[f"node_{target_idx}"]["test_r2"] = test_rsq
 
@@ -57,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,
@@ -85,12 +76,12 @@ def init_model(
         )
 
         model = lasso
-    elif model_str == "ridge-cv":
-        ridge = RidgeCV(fit_intercept=True, max_iter=max_iter)
+    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