Repository: jvalegre/robert (language: python)

Issue:
Merging v1.0



Discussion / hints:

Fix patch:
diff --git a/.circleci/config.yml b/.circleci/config.yml
index 95c677e..764ae71 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -36,12 +36,16 @@ jobs:
           # app-dir: ~/project/package-directory/  # If your requirements.txt isn't in the root directory.
           # pip-dependency-file: test-requirements.txt  # if you have a different name for your requirements file, maybe one that combines your runtime and test requirements.
       - run: |
-          # install aqme and dependencies of ROBERT
-          conda install -y -c conda-forge aqme
-          conda install -y -c conda-forge robert
-          # update ROBERT with pip
-          python -m pip install --upgrade pip
+          # install AQME, ROBERT and their dependencies
+          python -m pip install --upgrade pip          
+          conda install -y -c conda-forge rdkit
+          conda install -y -c conda-forge openbabel
+          conda install -y -c conda-forge xtb
+          conda install -y -c conda-forge crest
+          pip install aqme
           pip install .
+          # uninstall robert to avoid codecov issues (only running from the robert folder)
+          pip uninstall -y robert
           # install and run pytest
           pip install pytest
           pip install pytest-cov
diff --git a/README.md b/README.md
index a7521a4..88e1790 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,6 @@
 ![](Logos/Robert_logo.jpg)
 #
-## <p align="center"> Robert (Refiner and Optimizer of a Bunch of Existing Regression Tools)</p>
+## <p align="center"> ROBERT (Refiner and Optimizer of a Bunch of Existing Regression Tools)</p>
 
 
 [![CircleCI](https://img.shields.io/circleci/build/github/jvalegre/robert?label=Circle%20CI&logo=circleci)](https://app.circleci.com/pipelines/github/jvalegre/robert)
diff --git a/meta.yaml b/meta.yaml
index ca47d17..7cf677f 100644
--- a/meta.yaml
+++ b/meta.yaml
@@ -19,7 +19,7 @@ requirements:
     - pip
     - python >=3.6
   run:
-    - matplotlib-base
+    - matplotlib-base>=3.7.1
     - pandas
     - numpy<=1.23.1
     - progress
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..e69de29
diff --git a/robert/aqme.py b/robert/aqme.py
index ed1dfde..9fbabf6 100644
--- a/robert/aqme.py
+++ b/robert/aqme.py
@@ -14,7 +14,7 @@
      y : str, default=''
          Name of the column containing the response variable in the input CSV file (i.e. 'solubility'). 
      qdescp_keywords : str, default=''    
-         Add extra keywords to the AQME-QDESCP run (i.e. qdescp_keywords='--qdescp_atom P')
+         Add extra keywords to the AQME-QDESCP run (i.e. qdescp_keywords="--qdescp_atoms ['P']")
      csearch_keywords : str, default:''
          Add extra keywords to the AQME-CSEARCH run (i.e. csearch_keywords='--sample 10')
 """
@@ -24,6 +24,7 @@
 #####################################################.
 
 import os
+import glob
 import subprocess
 import time
 import shutil
@@ -68,12 +69,12 @@ def __init__(self, **kwargs):
 
         # if no qdesc_atom is set, only keep molecular properties and discard atomic properties
         aqme_db = f'AQME-ROBERT_{self.args.csv_name}'
-        if 'qdescp_atom' not in self.args.qdescp_keywords:
+        if 'qdescp_atoms' not in self.args.qdescp_keywords:
             _ = filter_atom_prop(aqme_db)
 
         # ensure that the AQME database was successfully created
         if not os.path.exists(aqme_db):
-            self.log.write(f"\nx  The initial AQME descriptor protocol did not create any CSV output!")
+            self.args.log.write(f"\nx  The initial AQME descriptor protocol did not create any CSV output!")
             sys.exit()
 
         # move AQME output files (remove previous runs as well)
@@ -102,13 +103,13 @@ def init_aqme(self):
         try:
             from aqme.qprep import qprep
         except ModuleNotFoundError:
-            self.log.write("x  AQME is not installed (required for the --aqme option)! You can install the program with 'conda install -c conda-forge aqme'")
+            self.args.log.write("x  AQME is not installed (required for the --aqme option)! You can install the program with 'conda install -c conda-forge aqme'")
             sys.exit()       
 
 
 def filter_atom_prop(aqme_db):
     '''
-    Function that filters off atomic properties if no atom was selected in the --qdescp_atom option
+    Function that filters off atomic properties if no atom was selected in the --qdescp_atoms option
     '''
     
     aqme_df = pd.read_csv(aqme_db)
@@ -127,10 +128,11 @@ def move_aqme():
     Move raw data from AQME-CSEARCH and -QDESCP runs into the AQME folder
     '''
     
-    for file in ['CSEARCH','QDESCP','CSEARCH_data.dat','CSEARCH-Data-solubility.csv','QDESCP_boltz_descriptors.csv','QDESCP_data.dat']:
-        if os.path.exists(f'AQME/{file}'):
-            if len(file.split('.')) == 1:
-                shutil.rmtree(f'AQME/{file}')
-            else:
-                os.remove(f'AQME/{file}')
-        shutil.move(file, f'AQME/{file}')
\ No newline at end of file
+    for file in glob.glob(f'*'):
+        if 'CSEARCH' in file or 'QDESCP' in file:
+            if os.path.exists(f'AQME/{file}'):
+                if len(file.split('.')) == 1:
+                    shutil.rmtree(f'AQME/{file}')
+                else:
+                    os.remove(f'AQME/{file}')
+            shutil.move(file, f'AQME/{file}')
\ No newline at end of file
diff --git a/robert/argument_parser.py b/robert/argument_parser.py
index 8abe6ed..334de0a 100644
--- a/robert/argument_parser.py
+++ b/robert/argument_parser.py
@@ -20,15 +20,16 @@
     "ignore" : [],
     "categorical" : "onehot",
     "corr_filter" : True,
-    "thres_y" : 0.02,
-    "thres_x" : 0.85,
+    "desc_thres" : 25,
+    "thres_y" : 0.001,
+    "thres_x" : 0.9,
     "train" : [60,70,80,90],
     "split" : "KN",
     "model" : ['RF','GB','NN','MVL'],
     "custom_params" : None,
     "type" : "reg",
     "seed" : 0,
-    "epochs" : 50,
+    "epochs" : 80,
     "error_type" : "rmse",
     "pfi_epochs" : 30,
     "pfi_threshold" : 0.04,
@@ -43,7 +44,7 @@
     "pfi_show" : 10,
     "names" : '',
     "qdescp_keywords" : '',
-    "csearch_keywords": '',
+    "csearch_keywords": '--sample 100',
     "report_modules" : ['AQME','CURATE','GENERATE','VERIFY','PREDICT']
 }
 
diff --git a/robert/curate.py b/robert/curate.py
index 23b0609..4428208 100644
--- a/robert/curate.py
+++ b/robert/curate.py
@@ -27,12 +27,15 @@
      corr_filter : bool, default=True
          Activate the correlation filters of descriptors. Two filters will be performed based on the correlation
          of the descriptors with other descriptors (x filter) and the y values (y filter).
-     thres_x : float, default=0.85
+     desc_thres : float, default=25
+         Threshold for the descriptor-to-datapoints ratio to deactivate the correlation filter. By default,
+         the correlation filter is disabled if there are 25 times more datapoints than descriptors.
+     thres_x : float, default=0.9
          Thresolhold to discard descriptors based on high R**2 correlation with other descriptors (i.e. 
-         if thres_x=0.85, variables that show R**2 > 0.85 will be discarded).
-     thres_y : float, default=0.02
+         if thres_x=0.9, variables that show R**2 > 0.9 will be discarded).
+     thres_y : float, default=0.001
          Thresolhold to discard descriptors with poor correlation with the y values based on R**2 (i.e.
-         if thres_y=0.02, variables that show R**2 < 0.02 will be discarded).
+         if thres_y=0.001, variables that show R**2 < 0.01 will be discarded).
 """
 #####################################################.
 #         This file stores the CURATE class         #
@@ -40,6 +43,7 @@
 #####################################################.
 
 import time
+import os
 import pandas as pd
 import numpy as np
 from scipy import stats
@@ -169,42 +173,46 @@ def correlation_filter(self, csv_df):
         """
 
         txt_corr = f'\no  Correlation filter activated with these thresholds: thres_x = {self.args.thres_x}, thres_y = {self.args.thres_y}'
-        txt_corr += f'\n   Excluded descriptors:'
-
+        # don't remove descriptors if there are too few descriptors (after 25 times more datapoints)
+        n_descps = len(csv_df.columns)-len(self.args.ignore)-1 # all columns - ignored - y
         descriptors_drop = []
-        for i,column in enumerate(csv_df.columns):
-            if column not in descriptors_drop and column not in self.args.ignore and column != self.args.y:
-                # finds the descriptors with low correlation to the response values
-                try:
-                    _, _, r_value_y, _, _ = stats.linregress(csv_df[column],csv_df[self.args.y])
-                    rsquared_y = r_value_y**2
-                    if rsquared_y < self.args.thres_y:
+        if not n_descps*self.args.desc_thres < len(csv_df[self.args.y]):
+            txt_corr += f'\n   Excluded descriptors:'
+            for i,column in enumerate(csv_df.columns):
+                if column not in descriptors_drop and column not in self.args.ignore and column != self.args.y:
+                    # finds the descriptors with low correlation to the response values
+                    try:
+                        _, _, r_value_y, _, _ = stats.linregress(csv_df[column],csv_df[self.args.y])
+                        rsquared_y = r_value_y**2
+                        if rsquared_y < self.args.thres_y:
+                            descriptors_drop.append(column)
+                            txt_corr += f'\n   - {column}: R**2 = {round(rsquared_y,2)} with the {self.args.y} values'
+                    except ValueError: # this avoids X descriptors where the majority of the values are the same
                         descriptors_drop.append(column)
-                        txt_corr += f'\n   - {column}: R**2 = {round(rsquared_y,2)} with the {self.args.y} values'
-                except ValueError: # this avoids X descriptors where the majority of the values are the same
-                    descriptors_drop.append(column)
-                    txt_corr += f'\n   - {column}: error in R**2 with the {self.args.y} values (are all the values the same?)'
-
-                # finds correlated descriptors
-                if column != csv_df.columns[-1] and column not in descriptors_drop:
-                    for j,column2 in enumerate(csv_df.columns):
-                        if j > i and column2 not in self.args.ignore and column2 not in descriptors_drop and column2 != self.args.y:
-                            _, _, r_value_x, _, _ = stats.linregress(csv_df[column],csv_df[column2])
-                            rsquared_x = r_value_x**2
-                            if rsquared_x > self.args.thres_x:
-                                # discard the column with less correlation with the y values
-                                _, _, r_value_y2, _, _ = stats.linregress(csv_df[column2],csv_df[self.args.y])
-                                rsquared_y2 = r_value_y2**2
-                                if rsquared_y >= rsquared_y2:
-                                    descriptors_drop.append(column2)
-                                    txt_corr += f'\n   - {column2}: R**2 = {round(rsquared_x,2)} with {column}'
-                                else:
-                                    descriptors_drop.append(column)
-                                    txt_corr += f'\n   - {column}: R**2 = {round(rsquared_x,2)} with {column2}'
-        
-        if len(descriptors_drop) == 0:
-            txt_corr += f'\n   -  No descriptors were removed'
-    
+                        txt_corr += f'\n   - {column}: error in R**2 with the {self.args.y} values (are all the values the same?)'
+
+                    # finds correlated descriptors
+                    if column != csv_df.columns[-1] and column not in descriptors_drop:
+                        for j,column2 in enumerate(csv_df.columns):
+                            if j > i and column2 not in self.args.ignore and column not in descriptors_drop and column2 not in descriptors_drop and column2 != self.args.y:
+                                _, _, r_value_x, _, _ = stats.linregress(csv_df[column],csv_df[column2])
+                                rsquared_x = r_value_x**2
+                                if rsquared_x > self.args.thres_x:
+                                    # discard the column with less correlation with the y values
+                                    _, _, r_value_y2, _, _ = stats.linregress(csv_df[column2],csv_df[self.args.y])
+                                    rsquared_y2 = r_value_y2**2
+                                    if rsquared_y >= rsquared_y2:
+                                        descriptors_drop.append(column2)
+                                        txt_corr += f'\n   - {column2}: R**2 = {round(rsquared_x,2)} with {column}'
+                                    else:
+                                        descriptors_drop.append(column)
+                                        txt_corr += f'\n   - {column}: R**2 = {round(rsquared_x,2)} with {column2}'
+            
+            if len(descriptors_drop) == 0:
+                txt_corr += f'\n   -  No descriptors were removed'
+        else:
+            txt_corr += f'\n   x The number of descriptors ({n_descps}) is {self.args.desc_thres} times lower than the number of datapoints ({len(csv_df[self.args.y])}), the correlation filter was disabled! The ratio of this filter can be adjusted with the desc_thres option'
+
         self.args.log.write(txt_corr)
 
         # drop descriptors that did not pass the filters
@@ -222,7 +230,7 @@ def save_curate(self,csv_df):
         '''
         
         # saves curated database
-        csv_basename = f'{self.args.csv_name}'.split('.')[0]
+        csv_basename = os.path.basename(f'{self.args.csv_name}').split('.')[0]
         csv_curate_name = f'{csv_basename}_CURATE.csv'
         csv_curate_name = self.args.destination.joinpath(csv_curate_name)
         _ = csv_df.to_csv(f'{csv_curate_name}', index = None, header=True)
@@ -250,12 +258,30 @@ def pearson_map(self,csv_df_pearson):
         mask = np.zeros_like(corr_matrix, dtype=np.bool)
         mask[np.triu_indices_from(mask)]= True
 
-        _, ax = plt.subplots(figsize=(7.45,6))
+        # these size ranges avoid matplot errors
+        if len(csv_df_pearson.columns) > 100:
+            _, ax = plt.subplots(figsize=(120,72))
+            size_title = 56
+            size_font = 18
+        elif len(csv_df_pearson.columns) > 60:
+            _, ax = plt.subplots(figsize=(45,36))
+            size_title = 42
+            size_font = 12
+        elif len(csv_df_pearson.columns) > 30:
+            _, ax = plt.subplots(figsize=(15,12))
+            size_title = 28
+            size_font = 6
+        else:
+            _, ax = plt.subplots(figsize=(7.45,6))
+            size_title = 14
+            size_font = 14-2*((len(csv_df_pearson.columns)/5))
 
         sb.set(font_scale=1.2, style='ticks')
 
         # determines size of the letters inside the boxes (approx.)
-        size_font = 14-2*((len(csv_df_pearson.columns)/5))
+        annot = True
+        if len(csv_df_pearson.columns) > 30:
+            annot = False
 
         _ = sb.heatmap(corr_matrix,
                             mask = mask,
@@ -267,7 +293,7 @@ def pearson_map(self,csv_df_pearson):
                                         'ticks' : [-1, -.5, 0, 0.5, 1]},
                             vmin = -1,
                             vmax = 1,
-                            annot = True,
+                            annot = annot,
                             annot_kws = {'size': size_font})
         plt.tick_params(labelsize=size_font)
         #add the column names as labels
@@ -275,7 +301,7 @@ def pearson_map(self,csv_df_pearson):
         ax.set_xticklabels(corr_matrix.columns)
 
         title_fig = 'Pearson\'s r heatmap'
-        plt.title(title_fig, y=1.04, fontsize = 14, fontweight="bold")
+        plt.title(title_fig, y=1.04, fontsize = size_title, fontweight="bold")
         sb.set_style({'xtick.bottom': True}, {'ytick.left': True})
 
         heatmap_name = 'Pearson_heatmap.png'
diff --git a/robert/generate.py b/robert/generate.py
index fd8b1da..75bd114 100644
--- a/robert/generate.py
+++ b/robert/generate.py
@@ -27,14 +27,15 @@
          Mode for splitting data. Options: 
             1. 'KN' (k-neighbours clustering-based splitting)
             2. 'RND' (random splitting).  
-     model : list, default=['RF','GB','NN','ADAB']
-         ML models used in the ML scan. Options: 
+     model : list, default=['RF','GB','NN','MVL']
+         ML models available: 
             1. 'RF' (Random forest)
             2. 'MVL' (Multivariate lineal models)
             3. 'GB' (Gradient boosting)
-            4. 'AdaB' (AdaBoost regressor)
-            5. 'NN' (MLP regressor neural network)
-            6. 'VR' (Voting regressor combining RF, GB and NN)
+            4. 'NN' (MLP neural network)
+            5. 'GP' (Gaussian Process)
+            6. 'AdaB' (AdaBoost)
+            7. 'VR' (Voting regressor combining RF, GB and NN)
      custom_params : str, default=None
          Define new parameters for the ML models used in the hyperoptimization workflow. The path
          to the folder containing all the yaml files should be specified (i.e. custom_params='YAML_FOLDER')
@@ -44,7 +45,7 @@
             2. 'clas' (Classifier)
      seed : int, default=8,
          Random seed used in the ML predictor models, data splitting and other protocols.
-     epochs : int, default=500,
+     epochs : int, default=80,
          Number of epochs for the hyperopt optimization.
      error_type : str, default: rmse (regression), acc (classification)
          Target value used during the hyperopt optimization. Options:
diff --git a/robert/generate_utils.py b/robert/generate_utils.py
index cdc7a99..529b206 100644
--- a/robert/generate_utils.py
+++ b/robert/generate_utils.py
@@ -31,14 +31,14 @@ def hyperopt_workflow(self, csv_df, ML_model, size, Xy_data_hp):
 
     # edit this function to modify the hyperopt parameter optimization (i.e. the 
     # lists represent values to include in the grid search)
-    space4rf = hyperopt_params(self, ML_model, Xy_data_hp['X_train_scaled'])
+    space4rf = hyperopt_params(self, ML_model)
 
     # Run hyperopt
     trials = Trials()
 
     # This best high initial number is a dummy value used as the start of the optimization
     # (otherwise it does not run more than once since best becomes a dictionary)
-    best = 100000
+    best = float('inf')
 
     hyperopt_data = {'best': best, 'model': ML_model.upper(),
                 'type': self.args.type.lower(),
@@ -64,6 +64,14 @@ def hyperopt_workflow(self, csv_df, ML_model, size, Xy_data_hp):
     else:
         max_evals = self.args.epochs
     try:
+        # this part allows users to modify exploitation/exploration, replaces the fmin() function below
+        # from hyperopt import partial, mix, anneal, rand
+        # best = fmin(f, space4rf, algo=partial(mix.suggest,
+        #         p_suggest=[
+        #         (.1, rand.suggest),
+        #         (.1, anneal.suggest),
+        #         (.8, tpe.suggest),]),
+        #         max_evals=max_evals, trials=trials, rstate=np.random.default_rng(hyperopt_data['seed']))
         best = fmin(f, space4rf, algo=tpe.suggest, max_evals=max_evals, trials=trials, rstate=np.random.default_rng(hyperopt_data['seed']))
         if os.path.exists('hyperopt.json'):
             os.remove('hyperopt.json')
@@ -82,43 +90,79 @@ def hyperopt_workflow(self, csv_df, ML_model, size, Xy_data_hp):
 
 
 # generates initial parameters for the hyperopt optimization
-def hyperopt_params(self, model_type, X_hyperopt):
+def hyperopt_params(self, model_type):
     # load the parameters of the models from their corresponding yaml files
-    params = load_from_yaml(self,model_type,X_hyperopt)
+    params = load_params(self,model_type)
 
     if model_type.upper() == 'RF':
         space4rf_hyperopt = {'max_depth': hp.choice('max_depth', params['max_depth']),
                 'max_features': hp.choice('max_features', params['max_features']),
-                'n_estimators': hp.choice('n_estimators', params['n_estimators'])}  
+                'n_estimators': hp.choice('n_estimators', params['n_estimators']),
+                'min_samples_split': hp.choice('min_samples_split', params['min_samples_split']),
+                'min_samples_leaf': hp.choice('min_samples_leaf', params['min_samples_leaf']),
+                'min_weight_fraction_leaf': hp.choice('min_weight_fraction_leaf', params['min_weight_fraction_leaf']),
+                'oob_score': hp.choice('oob_score', params['oob_score']),
+                'ccp_alpha': hp.choice('ccp_alpha', params['ccp_alpha']),
+                'max_samples': hp.choice('max_samples', params['max_samples']),
+                }  
 
     elif model_type.upper() == 'GB':
         space4rf_hyperopt = {'max_depth': hp.choice('max_depth', params['max_depth']),
                 'max_features': hp.choice('max_features', params['max_features']),
                 'n_estimators': hp.choice('n_estimators', params['n_estimators']),
                 'learning_rate': hp.choice('learning_rate', params['learning_rate']),
-                'validation_fraction': hp.choice('validation_fraction', params['validation_fraction'])}  
-
-    elif model_type.upper() == 'ADAB':
-        space4rf_hyperopt = {'n_estimators': hp.choice('n_estimators', params['n_estimators']),
-            'learning_rate': hp.choice('learning_rate', params['learning_rate'])}  
+                'validation_fraction': hp.choice('validation_fraction', params['validation_fraction']),
+                'subsample': hp.choice('subsample', params['subsample']),
+                'min_samples_split': hp.choice('min_samples_split', params['min_samples_split']),
+                'min_samples_leaf': hp.choice('min_samples_leaf', params['min_samples_leaf']),
+                'min_weight_fraction_leaf': hp.choice('min_weight_fraction_leaf', params['min_weight_fraction_leaf']),
+                'ccp_alpha': hp.choice('ccp_alpha', params['ccp_alpha'])}
 
     elif model_type.upper() == 'NN':
         space4rf_hyperopt = {'batch_size': hp.choice('batch_size', params['batch_size']),
                 'hidden_layer_sizes': hp.choice('hidden_layer_sizes', params['hidden_layer_sizes']),
                 'learning_rate_init': hp.choice('learning_rate_init', params['learning_rate_init']),
                 'max_iter': hp.choice('max_iter', params['max_iter']),
-                'validation_fraction': hp.choice('validation_fraction', params['validation_fraction'])}
+                'validation_fraction': hp.choice('validation_fraction', params['validation_fraction']),
+                'alpha': hp.choice('alpha', params['alpha']),
+                'shuffle': hp.choice('shuffle', params['shuffle']),
+                'tol': hp.choice('tol', params['tol']),
+                'early_stopping': hp.choice('early_stopping', params['early_stopping']),
+                'beta_1': hp.choice('beta_1', params['beta_1']),
+                'beta_2': hp.choice('beta_2', params['beta_2']),
+                'epsilon': hp.choice('epsilon', params['epsilon'])}
+
+    elif model_type.upper() == 'ADAB':
+        space4rf_hyperopt = {'n_estimators': hp.choice('n_estimators', params['n_estimators']),
+            'learning_rate': hp.choice('learning_rate', params['learning_rate'])}  
+
+    elif model_type.upper() == 'GP':
+        space4rf_hyperopt = {'n_restarts_optimizer': hp.choice('n_restarts_optimizer', params['n_restarts_optimizer'])}  
 
     elif model_type.upper() == 'VR':
         space4rf_hyperopt = {'max_depth': hp.choice('max_depth', params['max_depth']),
                 'max_features': hp.choice('max_features', params['max_features']),
                 'n_estimators': hp.choice('n_estimators', params['n_estimators']),
+                'min_samples_split': hp.choice('min_samples_split', params['min_samples_split']),
+                'min_samples_leaf': hp.choice('min_samples_leaf', params['min_samples_leaf']),
+                'min_weight_fraction_leaf': hp.choice('min_weight_fraction_leaf', params['min_weight_fraction_leaf']),
+                'oob_score': hp.choice('oob_score', params['oob_score']),
+                'ccp_alpha': hp.choice('ccp_alpha', params['ccp_alpha']),
+                'subsample': hp.choice('subsample', params['subsample']),
+                'max_samples': hp.choice('max_samples', params['max_samples']),
                 'learning_rate': hp.choice('learning_rate', params['learning_rate']),
                 'validation_fraction': hp.choice('validation_fraction', params['validation_fraction']),
                 'batch_size': hp.choice('batch_size', params['batch_size']),
                 'hidden_layer_sizes': hp.choice('hidden_layer_sizes', params['hidden_layer_sizes']),
                 'learning_rate_init': hp.choice('learning_rate_init', params['learning_rate_init']),
-                'max_iter': hp.choice('max_iter', params['max_iter'])}  
+                'max_iter': hp.choice('max_iter', params['max_iter']),
+                'alpha': hp.choice('alpha', params['alpha']),
+                'shuffle': hp.choice('shuffle', params['shuffle']),
+                'tol': hp.choice('tol', params['tol']),
+                'early_stopping': hp.choice('early_stopping', params['early_stopping']),
+                'beta_1': hp.choice('beta_1', params['beta_1']),
+                'beta_2': hp.choice('beta_2', params['beta_2']),
+                'epsilon': hp.choice('epsilon', params['epsilon'])}  
 
     elif model_type.upper() == 'MVL':
         space4rf_hyperopt = {'max_features': hp.choice('max_features', params['max_features'])}
@@ -153,12 +197,14 @@ def f(params):
             layer_arrays = ele
         params['hidden_layer_sizes'] = (layer_arrays)
 
-    opt_target = load_n_predict(params, hyperopt_data, hyperopt=True)
-
-    # since the hyperoptimizer aims to minimize the target values, the code needs to use negative
-    # values for R2, accuracy, F1 score and MCC (these values are inverted again before storing them)
-    if params['error_type'].lower() in ['r2', 'mcc', 'f1', 'acc']:
-        opt_target = -opt_target
+    try:
+        opt_target = load_n_predict(params, hyperopt_data, hyperopt=True)
+        # since the hyperoptimizer aims to minimize the target values, the code needs to use negative
+        # values for R2, accuracy, F1 score and MCC (these values are inverted again before storing them)
+        if params['error_type'].lower() in ['r2', 'mcc', 'f1', 'acc']:
+            opt_target = -opt_target
+    except RuntimeError:
+        opt_target = float('inf')
 
     if opt_target < best:
         # The "best" optimizing value is updated in an external JSON file after each hyperopt cycle
@@ -188,9 +234,18 @@ def f(params):
             csv_hyperopt['n_estimators'] = params['n_estimators']
             csv_hyperopt['max_depth'] = params['max_depth']
             csv_hyperopt['max_features'] = params['max_features']
+            csv_hyperopt['min_samples_split'] = params['min_samples_split']
+            csv_hyperopt['min_samples_leaf'] = params['min_samples_leaf']
+            csv_hyperopt['min_weight_fraction_leaf'] = params['min_weight_fraction_leaf']
+            csv_hyperopt['ccp_alpha'] = params['ccp_alpha']
+
+            if hyperopt_data['model'].upper() in ['RF','VR']:
+                csv_hyperopt['oob_score'] = params['oob_score']
+                csv_hyperopt['max_samples'] = params['max_samples']
 
             if hyperopt_data['model'].upper() in ['GB','VR']:
                 csv_hyperopt['learning_rate'] = params['learning_rate']
+                csv_hyperopt['subsample'] = params['subsample']
 
             if hyperopt_data['model'].upper() == 'GB':
                 csv_hyperopt['validation_fraction'] = params['validation_fraction']
@@ -201,11 +256,21 @@ def f(params):
                 csv_hyperopt['learning_rate_init'] = params['learning_rate_init']
                 csv_hyperopt['max_iter'] = params['max_iter']
                 csv_hyperopt['validation_fraction'] = params['validation_fraction']
+                csv_hyperopt['alpha'] = params['alpha']
+                csv_hyperopt['shuffle'] = params['shuffle']
+                csv_hyperopt['tol'] = params['tol']
+                csv_hyperopt['early_stopping'] = params['early_stopping']
+                csv_hyperopt['beta_1'] = params['beta_1']
+                csv_hyperopt['beta_2'] = params['beta_2']
+                csv_hyperopt['epsilon'] = params['epsilon']
     
         elif hyperopt_data['model'].upper() == 'ADAB':
                 csv_hyperopt['n_estimators'] = params['n_estimators']
                 csv_hyperopt['learning_rate'] = params['learning_rate']
-            
+
+        elif hyperopt_data['model'].upper() == 'GP':
+                csv_hyperopt['n_restarts_optimizer'] = params['n_restarts_optimizer']
+
         csv_hyperopt[hyperopt_data['error_type']] = opt_target
 
         # save into a csv file
@@ -219,7 +284,7 @@ def f(params):
     return {'loss': best, 'status': STATUS_OK}
 
 
-def load_from_yaml(self,model_type,X_hyperopt):
+def load_params(self,model_type):
     """
     Loads the parameters for the calculation from a yaml if specified. Otherwise
     does nothing.
@@ -240,16 +305,11 @@ def load_from_yaml(self,model_type,X_hyperopt):
             self.args.log.write(f'\nx  Error while reading {varfile}. Edit the yaml file and try again (i.e. use ":" instead of "=" to specify variables)')
             sys.exit()
 
-    # number of descriptors to scan, from 1 to the max amount of descriptors using the interval
-    # specified with interval_descriptors
-    if model_type in ['RF','GB','VR','MVL']:
-        max_features = [1]
-        n_descriptors = len(X_hyperopt.columns)
-        interval_descriptors = int((n_descriptors+1)/params["interval_descriptors"])
-        for num in range(max_features[0],n_descriptors,interval_descriptors):
-            max_features.append(num)
-        max_features.append(n_descriptors)
-        params['max_features'] = max_features
+    for param in params:
+        if params[param] == 'True':
+            params[param] = True
+        elif params[param] == 'False':
+            params[param] = False
 
     return params
 
@@ -327,7 +387,7 @@ def k_neigh(self,X_scaled,csv_y,size):
     try:
         kmeans.fit(X_scaled_array)
     except ValueError:
-        self.args.log.write("\nx  The K-means clustering process failed! This is probably due to having NaN or strings as descriptors. To avoid this issue, curate the data first with the CURATE module!")
+        self.args.log.write("\nx  The K-means clustering process failed! This might be due to having NaN or strings as descriptors (curate the data first with CURATE) or having too few datapoints!")
         sys.exit()
     _ = kmeans.predict(X_scaled_array)
     centers = kmeans.cluster_centers_
diff --git a/robert/model_params/GB_params.yaml b/robert/model_params/GB_params.yaml
index 5861688..ae98ba5 100644
--- a/robert/model_params/GB_params.yaml
+++ b/robert/model_params/GB_params.yaml
@@ -1,16 +1,29 @@
-# range to divide descriptors, including min and max (i.e., if you used 30 descriptors with interval_descriptors : 3, the amount of descriptors screened would be 1, 11, 21 and 30)
-interval_descriptors : 3
+# features used in the model
+max_features : [1.0,0.75,0.5,0.25]
 
 # learning rate shrinks the contribution of each tree by learning_rate
-learning_rate : [0.01,0.05,0.1,0.2,0.5]
+learning_rate : [0.01,0.05,0.1,0.5]
 
 # proportion of training data to set aside as validation set for early stopping
 validation_fraction : [0.1,0.2,0.3]
 
 # maximum depth of the individual regression estimators
-max_depth : [20,40,80,160,320]
+max_depth : [5,20,40,60,100]
 
 # number of boosting stages to perform
-n_estimators : [20,40,80,160,320]
+n_estimators : [5,20,40,60,100]
 
+# fraction of samples for the base learners
+subsample : [1.0]
 
+# min number of samples in each split
+min_samples_split : [2]
+
+# min number of samples for leaf nodes
+min_samples_leaf : [1]
+
+# min weight to be leaf node
+min_weight_fraction_leaf: [0]
+
+# complexity parameter for pruning
+ccp_alpha : [0]
\ No newline at end of file
diff --git a/robert/model_params/GP_params.yaml b/robert/model_params/GP_params.yaml
new file mode 100644
index 0000000..47ab3e9
--- /dev/null
+++ b/robert/model_params/GP_params.yaml
@@ -0,0 +1,2 @@
+# number of restarts of the optimizer
+n_restarts_optimizer : [0, 1, 2, 5, 10, 20, 50, 100]
\ No newline at end of file
diff --git a/robert/model_params/MVL_params.yaml b/robert/model_params/MVL_params.yaml
index 3a32388..d970874 100644
--- a/robert/model_params/MVL_params.yaml
+++ b/robert/model_params/MVL_params.yaml
@@ -1,4 +1,2 @@
-# range to divide descriptors, including min and max (i.e., if you used 30 descriptors with interval_descriptors : 3, the amount of descriptors screened would be 1, 11, 21 and 30)
-interval_descriptors : 3
-
-
+# features used in the model
+max_features : [1.0,0.9,0.8,0.7,0.6,0.5,0.4,0.3,0.2,0.1]
diff --git a/robert/model_params/NN_params.yaml b/robert/model_params/NN_params.yaml
index 01097c9..a262b7c 100644
--- a/robert/model_params/NN_params.yaml
+++ b/robert/model_params/NN_params.yaml
@@ -1,14 +1,35 @@
 # size of minibatches for stochastic optimizers
-batch_size : [20,40,80,160,320]
+batch_size : [4,8,24,64,128]
 
 # the ith element represents the number of neurons in the ith hidden layer
-hidden_layer_sizes : ["32","64","128","64,32","64,32,16"]
+hidden_layer_sizes : ["16","64","128","16,16","64,64"]
 
 # initial learning rate used
-learning_rate_init : [0.0001,0.0005,0.001,0.005,0.01]
+learning_rate_init : [0.001,0.01]
 
 # maximum number of iterations
-max_iter : [200,500]
+max_iter : [50,200]
 
 # proportion of training data to set aside as validation set for early stopping
-validation_fraction : [0.1,0.2]
+validation_fraction : [0.1,0.2,0.3]
+
+# L2 regularization term
+alpha : [0.0001]
+
+# shuffle samples
+shuffle : [True,False]
+
+# tolerance for opt
+tol : [0.0001]
+
+# enable early stop
+early_stopping : [False]
+
+# decay rate of first momentum moment
+beta_1 : [0.999,0.9,0.8]
+
+# decay rate of second momentum moment
+beta_2 : [0.999]
+
+# numerical stability
+epsilon : [0.00000001]
\ No newline at end of file
diff --git a/robert/model_params/RF_params.yaml b/robert/model_params/RF_params.yaml
index 964df5b..a5a4f42 100644
--- a/robert/model_params/RF_params.yaml
+++ b/robert/model_params/RF_params.yaml
@@ -1,10 +1,26 @@
-# range to divide descriptors, including min and max (i.e., if you used 30 descriptors with interval_descriptors : 3, the amount of descriptors screened would be 1, 11, 21 and 30)
-interval_descriptors : 3
+# features used in the model
+max_features : [1.0,0.75,0.5,0.25]
 
 # the maximum depth of the trees
-max_depth : [20,40,80,160,320]
+max_depth : [5,20,40,60,100]
 
 # the number of trees in the forest
-n_estimators : [1,3,20,40,80]
+n_estimators : [5,20,40,60,100]
 
+# min number of samples in each split
+min_samples_split : [2]
 
+# min number of samples for leaf nodes
+min_samples_leaf : [1]
+
+# min weight to be leaf node
+min_weight_fraction_leaf: [0]
+
+# use oob samples
+oob_score: [True,False]
+
+# complexity parameter for pruning
+ccp_alpha : [0]
+
+# number of samples to draw for bootstrap
+max_samples : [0.25,0.5,0.75,1]
\ No newline at end of file
diff --git a/robert/model_params/VR_params.yaml b/robert/model_params/VR_params.yaml
index 8e4d830..0b7b301 100644
--- a/robert/model_params/VR_params.yaml
+++ b/robert/model_params/VR_params.yaml
@@ -1,14 +1,28 @@
 # options for GB and RF models
-max_depth : [20,40,80,160,320]
-n_estimators : [5,20,40,160,320]
-interval_descriptors : 3
+max_depth : [5,20,40,60,100]
+n_estimators : [5,20,40,60,100]
+max_features : [1.0,0.75,0.5,0.25]
+min_samples_split : [2]
+min_samples_leaf : [1]
+min_weight_fraction_leaf: [0]
+oob_score: [True,False]
+ccp_alpha : [0]
+max_samples : [0.25,0.5,0.75,1]
+subsample : [1.0,0.75,0.5,0.25]
 
 # options for GB and NN
-learning_rate : [0.01,0.05,0.1,0.2,0.5]
-validation_fraction : [0.1,0.2]
+learning_rate : [0.01,0.05,0.1,0.5]
+validation_fraction : [0.1,0.2,0.3]
 
 # options for NN
-batch_size : [20,40,160,320]
-hidden_layer_sizes : ["32","64","64,32","64,32,16"]
-learning_rate_init : [0.0001,0.0005,0.001,0.01]
-max_iter : [200,500]
\ No newline at end of file
+batch_size : [4,8,24,64,128]
+hidden_layer_sizes : ["16","64","128","16,16","64,64"]
+learning_rate_init : [0.001,0.01]
+max_iter : [50,200]
+alpha : [0.0001]
+shuffle : [True,False]
+tol : [0.0001]
+early_stopping : [False]
+beta_1 : [0.999,0.9,0.8]
+beta_2 : [0.999]
+epsilon : [0.00000001]
\ No newline at end of file
diff --git a/robert/robert.py b/robert/robert.py
index 9c4c856..2728986 100644
--- a/robert/robert.py
+++ b/robert/robert.py
@@ -79,6 +79,7 @@ def main():
             ignore=args.ignore,
             categorical=args.categorical,
             corr_filter=args.corr_filter,
+            desc_thres=args.desc_thres,
             thres_x=args.thres_x,
             thres_y=args.thres_y,
         )
diff --git a/robert/utils.py b/robert/utils.py
index 3e252d5..a27a472 100644
--- a/robert/utils.py
+++ b/robert/utils.py
@@ -29,12 +29,13 @@
     GradientBoostingClassifier,
     AdaBoostClassifier,
     VotingClassifier)
+from sklearn.gaussian_process import GaussianProcessRegressor, GaussianProcessClassifier
 from sklearn.neural_network import MLPRegressor, MLPClassifier
 from sklearn.linear_model import LinearRegression
 import warnings # this avoids warnings from sklearn
 warnings.filterwarnings("ignore")
 
-robert_version = "0.0.1"
+robert_version = "1.1.0"
 time_run = time.strftime("%Y/%m/%d %H:%M:%S", time.localtime())
 robert_ref = f"ROBERT v {robert_version}, Dalmau, D.; Alegre-Requena, J. V., 2023. https://github.com/jvalegre/robert"
 
@@ -92,19 +93,6 @@ def write(self, message):
         self.log.write(f"{message}\n")
         print(f"{message}\n")
 
-    def fatal(self, message):
-        """
-        Writes the message to the file. Closes the file and raises an error exit
-
-        Parameters
-        ----------
-        message : str
-           text to be written in the log file.
-        """
-        self.write(message)
-        self.finalize()
-        raise SystemExit(1)
-
     def finalize(self):
         """
         Closes the file
@@ -180,7 +168,7 @@ def format_lists(value):
             value = ast.literal_eval(value)
         except (SyntaxError, ValueError):
             # this line fixes issues when using "[X]" or ["X"] instead of "['X']" when using lists
-            value = value.replace('[',']').replace(',',']').split(']')
+            value = value.replace('[',']').replace(',',']').replace("'",']').split(']')
             while('' in value):
                 value.remove('')
     return value
@@ -235,6 +223,7 @@ def load_variables(kwargs, robert_module):
             
             self.thres_x = float(self.thres_x)
             self.thres_y = float(self.thres_y)
+            self.desc_thres = float(self.desc_thres)
 
         elif robert_module.upper() == 'GENERATE':
             self.log.write(f"\no  Starting generation of ML models with the GENERATE module")
@@ -366,7 +355,7 @@ def sanity_checks(self, type_checks, module, columns_csv):
                 curate_valid = False
 
             for model_type in self.model:
-                if model_type.upper() not in ['RF','MVL','GB','ADAB','NN','VR'] or len(self.model) == 0:
+                if model_type.upper() not in ['RF','MVL','GB','GP','ADAB','NN','VR'] or len(self.model) == 0:
                     self.log.write(f"\nx  The model option used is not valid! Options: 'RF', 'MVL', 'GB', 'ADAB', 'NN', 'VR'")
                     curate_valid = False
                 if model_type.upper() == 'MVL' and self.type.lower() == 'clas':
@@ -405,8 +394,8 @@ def sanity_checks(self, type_checks, module, columns_csv):
                 curate_valid = False
         
         if module.lower() == 'verify':
-            if float(self.thres_test) > 1 or float(self.thres_test) < 0:
-                self.log.write(f"\nx  The thres_test should be between 0 and 1!")
+            if float(self.thres_test) < 0:
+                self.log.write(f"\nx  The thres_test should be higher than 0!")
                 curate_valid = False
 
         if module.lower() in ['verify','predict']:
@@ -543,55 +532,69 @@ def load_model(params):
 
 
 def load_model_reg(params):
-    if params['model'].upper() == 'RF':     
+    if params['model'].upper() in ['RF','VR']:
         loaded_model = RandomForestRegressor(max_depth=params['max_depth'],
                                 max_features=params['max_features'],
                                 n_estimators=params['n_estimators'],
+                                min_samples_split=params['min_samples_split'],
+                                min_samples_leaf=params['min_samples_leaf'],
+                                min_weight_fraction_leaf=params['min_weight_fraction_leaf'],
+                                oob_score=params['oob_score'],
+                                ccp_alpha=params['ccp_alpha'],
+                                max_samples=params['max_samples'],
                                 random_state=params['seed'],
                                 n_jobs=-1)
 
-    elif params['model'].upper() == 'GB':    
+        if params['model'].upper() == 'VR':
+            r1 = loaded_model
+
+    if params['model'].upper() in ['GB','VR']:    
         loaded_model = GradientBoostingRegressor(max_depth=params['max_depth'], 
                                 max_features=params['max_features'],
                                 n_estimators=params['n_estimators'],
                                 learning_rate=params['learning_rate'],
                                 validation_fraction=params['validation_fraction'],
+                                subsample=params['subsample'],
+                                min_samples_split=params['min_samples_split'],
+                                min_samples_leaf=params['min_samples_leaf'],
+                                min_weight_fraction_leaf=params['min_weight_fraction_leaf'],
+                                ccp_alpha=params['ccp_alpha'],
                                 random_state=params['seed'])
 
-    elif params['model'].upper() == 'ADAB':
-        loaded_model = AdaBoostRegressor(n_estimators=params['n_estimators'],
-                                learning_rate=params['learning_rate'],
-                                random_state=params['seed'])
+        if params['model'].upper() == 'VR':
+            r2 = loaded_model
 
-    elif params['model'].upper() == 'NN':
+    if params['model'].upper() in ['NN','VR']:
         loaded_model = MLPRegressor(batch_size=params['batch_size'],
                                 hidden_layer_sizes=params['hidden_layer_sizes'],
                                 learning_rate_init=params['learning_rate_init'],
                                 max_iter=params['max_iter'],
                                 validation_fraction=params['validation_fraction'],
-                                random_state=params['seed'])                    
+                                alpha=params['alpha'],
+                                shuffle=params['shuffle'],
+                                tol=params['tol'],
+                                early_stopping=params['early_stopping'],
+                                beta_1=params['beta_1'],
+                                beta_2=params['beta_2'],
+                                epsilon=params['epsilon'],
+                                random_state=params['seed'])  
+
+        if params['model'].upper() == 'VR':
+            r3 = loaded_model      
             
-    elif params['model'].upper() == 'VR':
-        r1 = GradientBoostingRegressor(max_depth=params['max_depth'], 
-                                max_features=params['max_features'],
-                                n_estimators=params['n_estimators'],
+    if params['model'].upper() == 'ADAB':
+        loaded_model = AdaBoostRegressor(n_estimators=params['n_estimators'],
                                 learning_rate=params['learning_rate'],
-                                validation_fraction=params['validation_fraction'],
                                 random_state=params['seed'])
-        r2 = RandomForestRegressor(max_depth=params['max_depth'],
-                                max_features=params['max_features'],
-                                n_estimators=params['n_estimators'],
-                                random_state=params['seed'],
-                                n_jobs=-1)
-        r3 = MLPRegressor(batch_size=params['batch_size'],
-                                hidden_layer_sizes=params['hidden_layer_sizes'],
-                                learning_rate_init=params['learning_rate_init'],
-                                max_iter=params['max_iter'],
-                                validation_fraction=params['validation_fraction'],
+
+    if params['model'].upper() == 'GP':
+        loaded_model = GaussianProcessRegressor(n_restarts_optimizer=params['n_restarts_optimizer'],
                                 random_state=params['seed'])
-        loaded_model = VotingRegressor([('gb', r1), ('rf', r2), ('nn', r3)])
 
-    elif params['model'].upper() == 'MVL':
+    if params['model'].upper() == 'VR':
+        loaded_model = VotingRegressor([('rf', r1), ('gb', r2), ('nn', r3)])
+
+    if params['model'].upper() == 'MVL':
         loaded_model = LinearRegression(n_jobs=-1)
 
     return loaded_model
@@ -599,54 +602,68 @@ def load_model_reg(params):
 
 def load_model_clas(params):
 
-    if params['model'].upper() == 'RF':     
+    if params['model'].upper() in ['RF','VR']:     
         loaded_model = RandomForestClassifier(max_depth=params['max_depth'],
                                 max_features=params['max_features'],
                                 n_estimators=params['n_estimators'],
+                                min_samples_split=params['min_samples_split'],
+                                min_samples_leaf=params['min_samples_leaf'],
+                                min_weight_fraction_leaf=params['min_weight_fraction_leaf'],
+                                oob_score=params['oob_score'],
+                                ccp_alpha=params['ccp_alpha'],
+                                max_samples=params['max_samples'],
                                 random_state=params['seed'],
                                 n_jobs=-1)
 
-    elif params['model'].upper() == 'GB':    
+        if params['model'].upper() == 'VR':
+            r1 = loaded_model
+
+    if params['model'].upper() in ['GB','VR']:    
         loaded_model = GradientBoostingClassifier(max_depth=params['max_depth'], 
                                 max_features=params['max_features'],
                                 n_estimators=params['n_estimators'],
                                 learning_rate=params['learning_rate'],
                                 validation_fraction=params['validation_fraction'],
+                                subsample=params['subsample'],
+                                min_samples_split=params['min_samples_split'],
+                                min_samples_leaf=params['min_samples_leaf'],
+                                min_weight_fraction_leaf=params['min_weight_fraction_leaf'],
+                                ccp_alpha=params['ccp_alpha'],
                                 random_state=params['seed'])
 
-    elif params['model'].upper() == 'ADAB':
+        if params['model'].upper() == 'VR':
+            r2 = loaded_model
+
+    if params['model'].upper() == 'ADAB':
             loaded_model = AdaBoostClassifier(n_estimators=params['n_estimators'],
                                 learning_rate=params['learning_rate'],
                                 random_state=params['seed'])
 
-    elif params['model'].upper() == 'NN':
+    if params['model'].upper() in ['NN','VR']:
         loaded_model = MLPClassifier(batch_size=params['batch_size'],
                                 hidden_layer_sizes=params['hidden_layer_sizes'],
                                 learning_rate_init=params['learning_rate_init'],
                                 max_iter=params['max_iter'],
                                 validation_fraction=params['validation_fraction'],
+                                alpha=params['alpha'],
+                                shuffle=params['shuffle'],
+                                tol=params['tol'],
+                                early_stopping=params['early_stopping'],
+                                beta_1=params['beta_1'],
+                                beta_2=params['beta_2'],
+                                epsilon=params['epsilon'],
                                 random_state=params['seed'])
 
-    elif params['model'].upper() == 'VR':
-        r1 = GradientBoostingClassifier(max_depth=params['max_depth'], 
-                                max_features=params['max_features'],
-                                n_estimators=params['n_estimators'],
-                                learning_rate=params['learning_rate'],
-                                validation_fraction=params['validation_fraction'],
-                                random_state=params['seed'])
-        r2 = RandomForestClassifier(max_depth=params['max_depth'],
-                                max_features=params['max_features'],
-                                n_estimators=params['n_estimators'],
+        if params['model'].upper() == 'VR':
+            r3 = loaded_model
+
+    if params['model'].upper() == 'GP':
+        loaded_model = GaussianProcessClassifier(n_restarts_optimizer=params['n_restarts_optimizer'],
                                 random_state=params['seed'],
                                 n_jobs=-1)
-        r3 = MLPClassifier(batch_size=params['batch_size'],
-                                hidden_layer_sizes=params['hidden_layer_sizes'],
-                                learning_rate_init=params['learning_rate_init'],
-                                max_iter=params['max_iter'],
-                                validation_fraction=params['validation_fraction'],
-                                random_state=params['seed'])
 
-        loaded_model = VotingClassifier([('gb', r1), ('rf', r2), ('nn', r3)])
+    if params['model'].upper() == 'VR':
+        loaded_model = VotingClassifier([('rf', r1), ('gb', r2), ('nn', r3)])
 
     return loaded_model
 
@@ -658,7 +675,7 @@ def load_n_predict(params, data, hyperopt=False):
     loaded_model = load_model(params)
 
     # Fit the model with the training set
-    loaded_model.fit(np.array(data['X_train_scaled']).tolist(), np.array(data['y_train']).tolist())
+    loaded_model.fit(data['X_train_scaled'], data['y_train'])
     # store the predicted values for training
     data['y_pred_train'] = loaded_model.predict(data['X_train_scaled']).tolist()
 
diff --git a/setup.py b/setup.py
index ed88f00..67eaa31 100644
--- a/setup.py
+++ b/setup.py
@@ -1,5 +1,5 @@
 from setuptools import setup, find_packages
-version="1.0.0"
+version="1.1.0"
 setup(
     name="robert",
     packages=find_packages(exclude=["tests"]),
@@ -36,7 +36,7 @@
         "pandas",
         "progress",
         "numpy<=1.23.1",
-        "matplotlib",
+        "matplotlib>=3.7.1",
         "seaborn",
         "scipy",
         "scikit-learn",


Test patch:
diff --git a/tests/Robert_example.csv b/tests/Robert_example.csv
index 4362921..d4a3436 100644
--- a/tests/Robert_example.csv
+++ b/tests/Robert_example.csv
@@ -19,7 +19,7 @@ Name,Target_values,x1,x2,x3,x4,x5,x6,x7,x8,x9,x10,x11,xtest,ynoise
 18,2.312908191,-0.46,110.959137,58.39613696,Csub-Csub,78.64178467,43.36078467,2,0,0,0,2,8,8
 19,1.16234335,-30.93,110.8476944,27.8146944,Csub-Csub,76.10786438,10.35686438,1,0,0,0,1,9,9
 20,1.989080521,-32.41,111.2339783,26.72097827,Csub-H,82.79759216,5.392952454,2,0,0,0,2,4,4
-21,2.801329141,-5.9,110.3707199,52.36771991,Csub-Csub,78.09907532,37.37807532,2,0,0,0,2,4,405
+21,2.801329141,-5.9,110.3707199,52.36771991,Csub-Csub,78.09907532,37.37807532,2,0,0,0,2,4,500
 22,2.391560251,-5.9,110.9408264,52.93782642,Csub-Csub,78.66916656,37.94816656,2,0,0,0,2,3,3
 23,2.011592734,-5.9,110.4220352,52.41903522,Csub-Csub,78.39561462,37.67461462,2,0,0,0,2,2,2
 24,0.998633019,-27.6,110.4682465,30.76524646,Csub-O,66.3848114,3.963811401,1,1,0,1,2,1,1
diff --git a/tests/Robert_example_NaNs.csv b/tests/Robert_example_NaNs.csv
index bb59f47..85ef571 100644
--- a/tests/Robert_example_NaNs.csv
+++ b/tests/Robert_example_NaNs.csv
@@ -19,7 +19,7 @@ Name,Target_values,x1,x2,x3,x4,x5,x6,x7,x8,x9,x10,x11,xtest,ynoise
 18,2.312908191,-0.46,110.959137,58.39613696,Csub-Csub,78.64178467,43.36078467,2,0,0,0,2,8,8
 19,1.16234335,-30.93,110.8476944,27.8146944,Csub-Csub,76.10786438,10.35686438,1,0,0,0,1,9,9
 20,1.989080521,-32.41,111.2339783,26.72097827,Csub-H,82.79759216,5.392952454,2,0,0,0,2,4,4
-21,2.801329141,-5.9,110.3707199,52.36771991,Csub-Csub,78.09907532,37.37807532,2,0,0,0,2,4,405
+21,2.801329141,-5.9,110.3707199,52.36771991,Csub-Csub,78.09907532,37.37807532,2,0,0,0,2,4,500
 22,2.391560251,-5.9,110.9408264,52.93782642,Csub-Csub,78.66916656,37.94816656,2,0,0,0,2,3,3
 23,2.011592734,-5.9,110.4220352,52.41903522,Csub-Csub,78.39561462,37.67461462,2,0,0,0,2,2,2
 24,0.998633019,-27.6,110.4682465,30.76524646,Csub-O,66.3848114,3.963811401,1,1,0,1,2,1,1
diff --git a/tests/params.yaml b/tests/params.yaml
new file mode 100644
index 0000000..769bf63
--- /dev/null
+++ b/tests/params.yaml
@@ -0,0 +1,6 @@
+curate : True
+y : Target_values
+ignore : ['Name']
+discard : ['xtest']
+thres_x : 0.999
+thres_y : 0.000001
\ No newline at end of file
diff --git a/tests/solubility.csv b/tests/solubility.csv
index eac0b65..494deb5 100644
--- a/tests/solubility.csv
+++ b/tests/solubility.csv
@@ -1,26 +1,25 @@
 code_name,solub,smiles
-mol_5,-1.33,c1ccsc1
-mol_6,-1.5,c2ccc1scnc1c2 
-mol_11,1.07,O=C1CCCN1
-mol_13,-2.68,CCCC=C
-mol_16,-1.41,CC(C)Cl
-mol_17,-0.47,CCC(C)CO
-mol_18,-1,N#Cc1ccccc1
-mol_23,-4.11,CCOP(=S)(OCC)SCSCC
-mol_24,-2.35,CCOc1ccc(NC(=O)C)cc1
-mol_26,-1.81,CCCCCCCO
-mol_30,-3.17,CCCCCCCC(=O)OC
-mol_31,-3.75,CCc1ccc(CC)cc1
-mol_34,-1.64,ClC(=C)Cl
-mol_36,-0.85,CCCCC=O
-mol_40,-4.81,CCCCCCCI
-mol_42,-0.15,OCCCC=C
-mol_45,-0.63,CC(C)OC=O
-mol_46,-1.55,CCCCCC(C)O
-mol_51,-3.66,CCCCCCC#C
-mol_56,-0.62,Cc1ccccc1O
-mol_58,-4.14,Cc1ccc(C)c2ccccc12
-mol_60,-0.19,CCCC(=O)C
-mol_62,-0.82,CCCOC(=O)CC
-mol_64,-0.8,Nc1ccc(O)cc1
-mol_1127,-3.18,CCC(C)C
+mol_5,-1.33,C
+mol_16,-1.41,CF
+mol_23,-4.11,CP
+mol_34,-1.64,CS
+mol_41,-2.93,CO
+mol_42,-2,CN
+mol_43,-3,PCP
+mol_44,-4,CCl
+mol_45,-1.5,OCN
+mol_46,-2.5,FCP
+mol_60,-4,OCO
+mol_61,-2,NCS
+mol_62,-3,SCS
+mol_63,-2,PCO
+mol_64,-3,ClCO
+mol_66,-4,FCO
+mol_67,-3,FCF
+mol_60,-1,FCN
+mol_61,-2,NCP
+mol_62,-3,SCN
+mol_63,-2,PCS
+mol_64,-3,ClCN
+mol_66,-4,FCCl
+mol_67,-3,OCS
diff --git a/tests/test_1curate.py b/tests/test_1curate.py
index 8263cf7..3dc4206 100644
--- a/tests/test_1curate.py
+++ b/tests/test_1curate.py
@@ -12,8 +12,8 @@
 import pandas as pd
 
 # saves the working directory
-path_main = os.getcwd() + "/tests"
-path_curate = path_main + "/CURATE"
+path_tests = os.getcwd() + "/tests"
+path_curate = os.getcwd() + "/CURATE"
 
 # CURATE tests
 @pytest.mark.parametrize(
@@ -31,6 +31,9 @@
         (
             "filter_thres"
         ),  # test to check the thresholds of the correlation filters
+        (
+            "filter_thres_yaml"
+        ),  # test to check the thresholds of the correlation filters with a yaml file
         (
             "standard"
         ),  # standard test
@@ -38,12 +41,9 @@
 )
 def test_CURATE(test_job):
 
-    # start in the tests folder
-    os.chdir(path_main)
-
     # leave the folders as they were initially to run a different batch of tests
-    if os.path.exists(f"{path_main}/CURATE"):
-        shutil.rmtree(f"{path_main}/CURATE")
+    if os.path.exists(f"{path_curate}"):
+        shutil.rmtree(f"{path_curate}")
         # remove DAT and CSV files generated by CURATE
         dat_files = glob.glob("*.dat")
         for dat_file in dat_files:
@@ -56,7 +56,7 @@ def test_CURATE(test_job):
         "-m",
         "robert",
         "--curate",
-        "--csv_name", "Robert_example.csv",
+        "--csv_name", f"{path_tests}/Robert_example.csv",
         '--y', 'Target_values',
         "--ignore", "['Name']",
         "--discard", "['xtest']"
@@ -66,7 +66,7 @@ def test_CURATE(test_job):
         subprocess.run(cmd_robert)
 
         # check that the DAT file is created inside the CURATE folder
-        assert not os.path.exists(f"{path_main}/CURATE_data.dat")
+        assert not os.path.exists("CURATE_data.dat")
         outfile = open(f"{path_curate}/CURATE_data.dat", "r")
         outlines = outfile.readlines()
         outfile.close()
@@ -122,7 +122,7 @@ def test_CURATE(test_job):
             "-m",
             "robert",
             "--curate",
-            "--csv_name", "Robert_example_NaNs.csv",
+            "--csv_name", f"{path_tests}/Robert_example_NaNs.csv",
             '--y', 'Target_values',
             "--ignore", "['Name']",
             "--discard", "['xtest']"
@@ -141,7 +141,7 @@ def test_CURATE(test_job):
             "-m",
             "robert",
             "--curate",
-            "--csv_name", "Robert_example.csv",
+            "--csv_name", f"{path_tests}/Robert_example.csv",
             '--y', 'Target_values',
             "--ignore", "['Name']",
             "--discard", "['xtest']",
@@ -155,19 +155,28 @@ def test_CURATE(test_job):
         for var in discard_vars:
             assert var in db_final.columns
 
-    elif test_job == 'filter_thres':
-        cmd_robert = [
-            "python",
-            "-m",
-            "robert",
-            "--curate",
-            "--csv_name", "Robert_example.csv",
-            '--y', 'Target_values',
-            "--ignore", "['Name']",
-            "--discard", "['xtest']",
-            "--thres_x", "0.999",
-            "--thres_y", "0.001"
-        ]
+    elif test_job in ['filter_thres','filter_thres_yaml']:
+        if test_job == 'filter_thres':
+            cmd_robert = [
+                "python",
+                "-m",
+                "robert",
+                "--curate",
+                "--csv_name", f"{path_tests}/Robert_example.csv",
+                '--y', 'Target_values',
+                "--ignore", "['Name']",
+                "--discard", "['xtest']",
+                "--thres_x", "0.999",
+                "--thres_y", "0.000001"
+            ]
+        elif test_job == 'filter_thres_yaml':
+            cmd_robert = [
+                "python",
+                "-m",
+                "robert",
+                "--varfile", f"{path_tests}/params.yaml",
+                "--csv_name", f"{path_tests}/Robert_example.csv",
+            ]
         subprocess.run(cmd_robert)
         
         # check that descriptors aren't removed
diff --git a/tests/test_2generate.py b/tests/test_2generate.py
index 8c20d45..fbad701 100644
--- a/tests/test_2generate.py
+++ b/tests/test_2generate.py
@@ -13,8 +13,8 @@
 from pathlib import Path
 
 # saves the working directory
-path_main = os.getcwd() + "/tests"
-path_generate = path_main + "/GENERATE"
+path_main = os.getcwd()
+path_generate = os.getcwd() + "/GENERATE"
 
 # GENERATE tests
 @pytest.mark.parametrize(
@@ -45,12 +45,9 @@
 )
 def test_GENERATE(test_job):
 
-    # start in the tests folder
-    os.chdir(path_main)
-
     # leave the folders as they were initially to run a different batch of tests
-    if os.path.exists(f"{path_main}/GENERATE"):
-        shutil.rmtree(f"{path_main}/GENERATE")
+    if os.path.exists(f"{path_generate}"):
+        shutil.rmtree(f"{path_generate}")
         # remove DAT and CSV files generated by GENERATE
         dat_files = glob.glob("*.dat")
         for dat_file in dat_files:
@@ -61,7 +58,7 @@ def test_GENERATE(test_job):
 
     # runs the program with the different tests
     if test_job == 'reduced_clas':
-        csv_name = 'Robert_example_clas.csv'
+        csv_name = 'tests/Robert_example_clas.csv'
     else:
         csv_name = 'CURATE/Robert_example_CURATE.csv'
     cmd_robert = [
@@ -79,8 +76,11 @@ def test_GENERATE(test_job):
         if test_job != 'reduced_others':
             model_list = ['RF']
         else:
-            model_list = ['Adab','VR']
-        train_list = [60]
+            model_list = ['Adab','VR','GP']
+        if test_job == 'reduced_clas':
+            train_list = [80]
+        else:
+            train_list = [60]
 
         if test_job == "reduced_noPFI":
             cmd_robert = cmd_robert + ["--pfi_filter", "False"]
@@ -146,7 +146,7 @@ def test_GENERATE(test_job):
                 elif test_job =='reduced_PFImax':
                     desc_list = ['x6', 'x7']
                 elif test_job == 'reduced_random':
-                    desc_list = ['x6', 'x7', 'x10', 'Csub-Csub', 'Csub-H']
+                    desc_list = ['x6', 'x7', 'x10']
                     assert db_best['Set'][0] == 'Validation'
                     assert db_best['Set'][1] == 'Validation'
                     assert db_best['Set'][2] == 'Training'
@@ -163,5 +163,5 @@ def test_GENERATE(test_job):
         assert not os.path.exists(f'{path_generate}/Raw_data/Heatmap ML models with PFI filter.png')
 
     if test_job == 'reduced_clas':
-        filepath = Path(f"{path_main}/GENERATE")
+        filepath = Path(f"{path_generate}")
         filepath.rename(f"{path_main}/GENERATE_clas")
diff --git a/tests/test_3verify.py b/tests/test_3verify.py
index 1ebe9a4..a12a3a2 100644
--- a/tests/test_3verify.py
+++ b/tests/test_3verify.py
@@ -12,7 +12,7 @@
 from pathlib import Path
 
 # saves the working directory
-path_main = os.getcwd() + "/tests"
+path_main = os.getcwd()
 path_verify = path_main + "/VERIFY"
 
 # VERIFY tests
@@ -35,12 +35,9 @@
 )
 def test_VERIFY(test_job):
 
-    # start in the tests folder
-    os.chdir(path_main)
-
     # leave the folders as they were initially to run a different batch of tests
-    if os.path.exists(f"{path_main}/VERIFY"):
-        shutil.rmtree(f"{path_main}/VERIFY")
+    if os.path.exists(f"{path_verify}"):
+        shutil.rmtree(f"{path_verify}")
         # remove DAT and CSV files generated by VERIFY
         dat_files = glob.glob("*.dat")
         for dat_file in dat_files:
@@ -73,11 +70,10 @@ def test_VERIFY(test_job):
     ]
 
     if test_job == "thres_test":
-        cmd_robert = cmd_robert + ["--thres_test", "1"]
+        cmd_robert = cmd_robert + ["--thres_test", "10"]
     elif test_job == "kfold":
         cmd_robert = cmd_robert + ["--kfold", "10"]
 
-
     subprocess.run(cmd_robert)
 
     # check that the DAT file is created
@@ -90,8 +86,8 @@ def test_VERIFY(test_job):
         if 'Results of the VERIFY tests:' in line:
             if test_job == "thres_test":
                 assert "- 5-fold CV: NOT DETERMINED" in outlines[i+2]
-                assert "o y_mean: PASSED" in outlines[i+4]
-                assert "o y_shuffle: PASSED" in outlines[i+5]
+                assert "x y_mean: FAILED" in outlines[i+4]
+                assert "x y_shuffle: FAILED" in outlines[i+5]
                 assert "x onehot: FAILED" in outlines[i+6]
             elif test_job == "kfold":
                 assert "- 10-fold CV: NOT DETERMINED" in outlines[i+2]
diff --git a/tests/test_4predict.py b/tests/test_4predict.py
index 434518f..6c762cc 100644
--- a/tests/test_4predict.py
+++ b/tests/test_4predict.py
@@ -12,7 +12,7 @@
 from pathlib import Path
 
 # saves the working directory
-path_main = os.getcwd() + "/tests"
+path_main = os.getcwd()
 path_predict = path_main + "/PREDICT"
 
 # PREDICT tests
@@ -38,12 +38,9 @@
 )
 def test_PREDICT(test_job):
 
-    # start in the tests folder
-    os.chdir(path_main)
-
     # leave the folders as they were initially to run a different batch of tests
-    if os.path.exists(f"{path_main}/PREDICT"):
-        shutil.rmtree(f"{path_main}/PREDICT")
+    if os.path.exists(f"{path_predict}"):
+        shutil.rmtree(f"{path_predict}")
         # remove DAT and CSV files generated by PREDICT
         dat_files = glob.glob("*.dat")
         for dat_file in dat_files:
@@ -80,7 +77,7 @@ def test_PREDICT(test_job):
     if test_job == "names":
         cmd_robert = cmd_robert + ["--names", "Name"]
     if test_job == "csv_test":
-        cmd_robert = cmd_robert + ["--csv_test", "Robert_example_test.csv"]
+        cmd_robert = cmd_robert + ["--csv_test", "tests/Robert_example_test.csv"]
 
     subprocess.run(cmd_robert)
 
@@ -93,24 +90,22 @@ def test_PREDICT(test_job):
     categor_test = False
     for i,line in enumerate(outlines):
         if 'Results saved in' in line and 'No_PFI.dat' in line:
-            if test_job not in ["csv_test","clas"]:
+            if test_job not in ["csv_test"]:
                 assert 'Points Train:Validation = 29:8' in outlines[i+1]
-            elif test_job == "csv_test":
-                assert 'Points Train:Validation:Test = 29:8:9' in outlines[i+1]
             else:
-                assert 'Points Train:Validation = 22:15' in outlines[i+1]
+                assert 'Points Train:Validation:Test = 29:8:9' in outlines[i+1]
         elif 'x  There are missing descriptors' in line:
             categor_test = True
         if 'Outlier values saved in' in line and 'No_PFI' in line:
             if test_job == "names":
-                assert '-  21 (' in outlines[i+2]
+                assert '-  19 (' in outlines[i+2]
                 break
             elif test_job != "clas":
                 assert 'x  No names option' in outlines[i+1]             
                 if test_job == "t_value":
                     assert 'Train: 0 outliers' in outlines[i+2]
                 else:
-                    assert 'Train: 1 outliers' in outlines[i+2]
+                    assert 'Train: 2 outliers' in outlines[i+2]
                 assert 'Validation: 0 outliers' in outlines[i+3]
                 if test_job == "csv_test":
                     assert 'Test: 0 outliers' in outlines[i+4]
diff --git a/tests/test_5aqme_n_full.py b/tests/test_5aqme_n_full.py
index 9900839..365e26d 100644
--- a/tests/test_5aqme_n_full.py
+++ b/tests/test_5aqme_n_full.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 
 ######################################################.
-# 	         Testing PREDICT with pytest 	         #
+# 	          Testing AQME with pytest   	         #
 ######################################################.
 
 import os
@@ -12,10 +12,10 @@
 import pandas as pd
 
 # saves the working directory
-path_main = os.getcwd() + "/tests"
+path_main = os.getcwd()
 path_aqme = path_main + "/AQME"
 
-# PREDICT tests
+# AQME and full workflow tests
 @pytest.mark.parametrize(
     "test_job",
     [
@@ -27,12 +27,9 @@
         ),  # test for a full workflow starting from AQME
     ],
 )
-def test_PREDICT(test_job):
+def test_AQME(test_job):
 
-    # start in the tests folder
-    os.chdir(path_main)
-
-    # reset the folder
+    # reset the folders
     folders = ['CURATE','GENERATE','GENERATE_reg','GENERATE_clas','PREDICT','VERIFY','AQME']
     for folder in folders:
         if os.path.exists(f"{path_main}/{folder}"):
@@ -41,11 +38,13 @@ def test_PREDICT(test_job):
     # runs the program with the different tests
     if test_job == 'full_workflow':
         y_var = 'Target_values'
-        csv_var = "Robert_example.csv"
+        csv_var = "tests/Robert_example.csv"
         ignore_var = "['Name']"
 
     elif test_job == 'aqme':
         y_var = 'solub'
+        # for AQME-ROBERT workflows, the CSV file must be in the working dir
+        shutil.copy(f"{path_main}/tests/solubility.csv", f"{path_main}/solubility.csv")
         csv_var = "solubility.csv"
         ignore_var = "['smiles','code_name']"
 
@@ -63,7 +62,9 @@ def test_PREDICT(test_job):
         cmd_robert = cmd_robert + ["--discard", "['xtest']"]
 
     elif test_job == 'aqme':
-        cmd_robert = cmd_robert + ["--aqme","--csearch_keywords", "--sample 2", "--qdescp_keywords", "--qdescp_atom P"]
+        cmd_robert = cmd_robert + ["--aqme","--csearch_keywords", "--sample 2", 
+                    "--qdescp_keywords", "--qdescp_atoms ['C']", "--model", "['RF']",
+                    "--train", "[60]"]
 
     subprocess.run(cmd_robert)
 
@@ -77,7 +78,10 @@ def test_PREDICT(test_job):
     folders_gen = ['No_PFI','PFI']
     for folder in folders_gen:
         csv_amount = glob.glob(f'{path_main}/GENERATE/Raw_data/{folder}/*.csv')
-        assert len(csv_amount) == 32
+        if test_job == 'aqme':
+            assert len(csv_amount) == 2
+        else:
+            assert len(csv_amount) == 32
         best_amount = glob.glob(f'{path_main}/GENERATE/Best_model/{folder}/*.csv')
         assert len(best_amount) == 2
 
@@ -93,7 +97,7 @@ def test_PREDICT(test_job):
     if test_job == 'aqme':
         assert os.path.exists(f'{path_main}/AQME-ROBERT_solubility.csv')
         db_aqme = pd.read_csv(f'{path_main}/AQME-ROBERT_solubility.csv')
-        descps = ['code_name','solub','smiles','FUKUI+','DBSTEP_Vbur','TPSA']
+        descps = ['code_name','solub','smiles','C_FUKUI+','C_DBSTEP_Vbur','MolLogP']
         for descp in descps:
             assert descp in db_aqme.columns
 
@@ -112,4 +116,5 @@ def test_PREDICT(test_job):
         if os.path.exists(f"{path_main}/{folder}"):
             shutil.rmtree(f"{path_main}/{folder}")
     if os.path.exists(f'{path_main}/AQME-ROBERT_solubility.csv'):
-        os.remove(f'{path_main}/AQME-ROBERT_solubility.csv')
\ No newline at end of file
+        os.remove(f'{path_main}/AQME-ROBERT_solubility.csv')
+