Update GP config, remove assignment

Author: sgreenbury

autoemulate/experimental/emulators/gaussian_process/exact.py

@@ -8,7 +8,7 @@
 from gpytorch.kernels import MultitaskKernel, ScaleKernel
 from gpytorch.likelihoods import MultitaskGaussianLikelihood
 from gpytorch.means import MultitaskMean
-from torch import nn, optim
+from torch import optim
 from torch.optim.lr_scheduler import LRScheduler
 
 from autoemulate.experimental.callbacks.early_stopping import (
@@ -63,7 +63,6 @@ def __init__(  # noqa: PLR0913 allow too many arguments since all currently requ
         mean_module_fn: MeanModuleFn = constant_mean,
         covar_module_fn: CovarModuleFn = rbf,
         epochs: int = 50,
-        activation: type[nn.Module] = nn.ReLU,
         lr: float = 1e-1,
         early_stopping: EarlyStopping | None = None,
         device: DeviceLike | None = None,
@@ -87,8 +86,6 @@ def __init__(  # noqa: PLR0913 allow too many arguments since all currently requ
             Function to create the covariance module.
         epochs : int, default=50
             Number of training epochs.
-        activation : type[nn.Module], default=nn.ReLU
-            Activation function to use in the model.
         lr : float, default=2e-1
             Learning rate for the optimizer.
         device : DeviceLike | None, default=None
@@ -130,7 +127,6 @@ def __init__(  # noqa: PLR0913 allow too many arguments since all currently requ
         self.covar_module = covar_module
         self.epochs = epochs
         self.lr = lr
-        self.activation = activation
         self.optimizer = self.optimizer_cls(self.parameters(), lr=self.lr)  # type: ignore[call-arg] since all optimizers include lr
         # Extract scheduler-specific kwargs if present
         scheduler_kwargs = kwargs.pop("scheduler_kwargs", {})
@@ -206,8 +202,27 @@ def _predict(self, x: TensorLike, with_grad: bool) -> GaussianProcessLike:
             x = x.to(self.device)
             return self(x)
 
+    @classmethod
+    def scheduler_config(cls) -> dict:
+        """
+        Returns a random configuration for the learning rate scheduler.
+        This should be added to the `get_tune_config()` method of subclasses
+        to allow tuning of the scheduler parameters.
+        """
+        all_params = [
+            {"scheduler_cls": None, "scheduler_kwargs": None},
+            {
+                "scheduler_cls": [LRScheduler],
+                "scheduler_kwargs": [
+                    {"policy": "ReduceLROnPlateau", "patience": 5, "factor": 0.5}
+                ],
+            },
+        ]
+        return np.random.choice(all_params)
+
     @staticmethod
     def get_tune_config():
+        scheduler_params = GaussianProcessExact.scheduler_config()
         return {
             "mean_module_fn": [
                 constant_mean,
@@ -226,12 +241,10 @@ def get_tune_config():
                 rbf_times_linear,
             ],
             "epochs": [50, 100, 200],
-            "activation": [
-                nn.ReLU,
-                nn.GELU,
-            ],
-            "lr": list(np.logspace(-3, -1)),
+            "lr": list(np.logspace(-3, 0, 100)),
             "likelihood_cls": [MultitaskGaussianLikelihood],
+            "scheduler_cls": scheduler_params["scheduler_cls"],
+            "scheduler_kwargs": scheduler_params["scheduler_kwargs"],
         }
 
 
@@ -255,7 +268,6 @@ def __init__(  # noqa: PLR0913 allow too many arguments since all currently requ
         mean_module_fn: MeanModuleFn = constant_mean,
         covar_module_fn: CovarModuleFn = rbf,
         epochs: int = 50,
-        activation: type[nn.Module] = nn.ReLU,
         lr: float = 2e-1,
         early_stopping: EarlyStopping | None = None,
         seed: int | None = None,
@@ -332,7 +344,6 @@ def __init__(  # noqa: PLR0913 allow too many arguments since all currently requ
         self.covar_module = covar_module
         self.epochs = epochs
         self.lr = lr
-        self.activation = activation
         self.optimizer = self.optimizer_cls(self.parameters(), lr=self.lr)  # type: ignore[call-arg] since all optimizers include lr
         # Extract scheduler-specific kwargs if present
         scheduler_kwargs = kwargs.pop("scheduler_kwargs", {})