From 06ba365dda0189f4e936915533614eb80626d2e7 Mon Sep 17 00:00:00 2001
From: "Bryan M. Li" <bryanlimy@gmail.com>
Date: Wed, 13 Mar 2024 12:59:30 +0000
Subject: [PATCH] lower max learning rate value to avoid gradient explode. set
 verbose and error_score to both Searcher

---
 autoemulate/emulators/neural_networks/mlp.py |  6 +++---
 autoemulate/emulators/neural_networks/rbf.py |  6 +++---
 autoemulate/hyperparam_searching.py          | 14 +++++++++++---
 3 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/autoemulate/emulators/neural_networks/mlp.py b/autoemulate/emulators/neural_networks/mlp.py
index fbc89b9c..b685f10e 100644
--- a/autoemulate/emulators/neural_networks/mlp.py
+++ b/autoemulate/emulators/neural_networks/mlp.py
@@ -50,18 +50,18 @@ def get_grid_params(search_type: str = "random"):
                 nn.Sigmoid,
                 nn.GELU,
             ],
-            "optimizer": [torch.optim.AdamW, torch.optim.LBFGS],
+            "optimizer": [torch.optim.SGD, torch.optim.AdamW, torch.optim.LBFGS],
             "optimizer__weight_decay": (1 / 10 ** np.arange(1, 9)).tolist(),
         }
         match search_type:
             case "random":
                 param_space |= {
-                    "lr": loguniform(1e-06, 1e-2),
+                    "lr": loguniform(1e-6, 1e-4),
                 }
             case "bayes":
                 param_space |= {
                     "optimizer": Categorical(param_space["optimizer"]),
-                    "lr": Real(1e-06, 1e-2, prior="log-uniform"),
+                    "lr": Real(1e-6, 1e-4, prior="log-uniform"),
                 }
             case _:
                 raise ValueError(f"Invalid search type: {search_type}")
diff --git a/autoemulate/emulators/neural_networks/rbf.py b/autoemulate/emulators/neural_networks/rbf.py
index 8738e2c1..285034dc 100644
--- a/autoemulate/emulators/neural_networks/rbf.py
+++ b/autoemulate/emulators/neural_networks/rbf.py
@@ -324,18 +324,18 @@ def get_grid_params(search_type: str = "random"):
                 rbf_inverse_quadratic,
                 rbf_inverse_multiquadric,
             ],
-            "optimizer": [torch.optim.AdamW, torch.optim.SGD],
+            "optimizer": [torch.optim.SGD, torch.optim.AdamW, torch.optim.LBFGS],
             "optimizer__weight_decay": (1 / 10 ** np.arange(1, 9)).tolist(),
         }
         match search_type:
             case "random":
                 param_space |= {
-                    "lr": loguniform(1e-06, 1e-2),
+                    "lr": loguniform(1e-06, 1e-3),
                 }
             case "bayes":
                 param_space |= {
                     "optimizer": Categorical(param_space["optimizer"]),
-                    "lr": Real(1e-06, 1e-2, prior="log-uniform"),
+                    "lr": Real(1e-06, 1e-3, prior="log-uniform"),
                 }
             case _:
                 raise ValueError(f"Invalid search type: {search_type}")
diff --git a/autoemulate/hyperparam_searching.py b/autoemulate/hyperparam_searching.py
index 783c4e52..89464555 100644
--- a/autoemulate/hyperparam_searching.py
+++ b/autoemulate/hyperparam_searching.py
@@ -23,6 +23,8 @@ def _optimize_params(
     param_space=None,
     n_jobs=None,
     logger=None,
+    error_score=np.nan,
+    verbose=0,
 ):
     """Performs hyperparameter search for the provided model.
 
@@ -49,7 +51,11 @@ def _optimize_params(
     n_jobs : int
         Number of jobs to run in parallel.
     logger : logging.Logger
-        Logger instance.
+        Logger instance
+    error_score: 'raise' or numeric
+        Value to assign to the score if an error occurs in estimator fitting.
+    verbose: int
+        Verbosity level for the searcher
 
     Returns
     -------
@@ -68,7 +74,8 @@ def _optimize_params(
             cv=cv,
             n_jobs=n_jobs,
             refit=True,
-            verbose=0,
+            error_score=error_score,
+            verbose=verbose,
         )
     # Bayes search
     elif search_type == "bayes":
@@ -79,7 +86,8 @@ def _optimize_params(
             cv=cv,
             n_jobs=n_jobs,
             refit=True,
-            verbose=0,
+            error_score=error_score,
+            verbose=verbose,
         )
     elif search_type == "grid":
         raise NotImplementedError("Grid search not available yet.")