Merge pull request #171 from alan-turing-institute/main

Pulling in latest edits to docs (and Bryan's latest update)

CONTRIBUTING.md

@@ -30,7 +30,9 @@ The easiest way to get involved with the active development of AutoEmulate is to
 
 ## How to Contribute
 
-We welcome contributions of all kinds, be it code, documentation, or community engagement. We encourage you to read through the following sections to learn more about how you can contribute to
+We welcome contributions of all kinds, be it code, documentation, or community engagement. We encourage you to read through the following sections to learn more about how you can contribute to the package.
+
+We are always interested in adding more simulations or simulation input/output datasets from any field (see https://github.com/alan-turing-institute/autoemulate/issues/4).
 
 ## How to Submit Changes
 

autoemulate/emulators/neural_net_torch.py

@@ -8,6 +8,7 @@
 import torch
 from scipy.sparse import issparse
 from sklearn.exceptions import DataConversionWarning
+from sklearn.exceptions import NotFittedError
 from skorch import NeuralNetRegressor
 from skorch.callbacks import Callback
 
@@ -49,13 +50,13 @@ def __init__(
         self,
         module: str = "mlp",
         criterion=torch.nn.MSELoss,
-        optimizer=torch.optim.Adam,
-        lr: float = 0.01,
+        optimizer=torch.optim.AdamW,
+        lr: float = 1e-3,
         batch_size: int = 128,
         max_epochs: int = 1,
         module__input_size: int = 2,
         module__output_size: int = 1,
-        optimizer__weight_decay: float = 0.0001,
+        optimizer__weight_decay: float = 0.0,
         iterator_train__shuffle: bool = True,
         callbacks: List[Callback] = [InputShapeSetter()],
         train_split: bool = False,  # to run cross_validate without splitting the data
@@ -65,17 +66,8 @@ def __init__(
         if "random_state" in kwargs:
             setattr(self, "random_state", kwargs.pop("random_state"))
             set_random_seed(self.random_state)
-        # get all arguments for module initialization
-        module_args = {
-            "input_size": module__input_size,
-            "output_size": module__output_size,
-        }
-        for k, v in kwargs.items():
-            if k.startswith("module__"):
-                module_args[k.replace("module__", "")] = v
-
         super().__init__(
-            module=get_module(module, module_args),
+            module=get_module(module),
             criterion=criterion,
             optimizer=optimizer,
             lr=lr,
@@ -90,8 +82,7 @@ def __init__(
             verbose=verbose,
             **kwargs,
         )
-        self._initialize_module()
-        self._initialize_optimizer()
+        self.initialize()
 
     def set_params(self, **params):
         if "random_state" in params:
@@ -101,8 +92,7 @@ def set_params(self, **params):
             else:
                 setattr(self, "random_state", random_state)
             set_random_seed(self.random_state)
-            self._initialize_module()
-            self._initialize_optimizer()
+            self.initialize()
         return super().set_params(**params)
 
     def initialize_module(self, reason=None):
@@ -127,9 +117,8 @@ def _more_tags(self):
                 "check_no_attributes_set_in_init": "skorch initialize attributes in __init__.",
                 "check_regressors_no_decision_function": "skorch NeuralNetRegressor class implements the predict_proba.",
                 "check_parameters_default_constructible": "skorch NeuralNet class callbacks parameter expects a list of callables.",
+                "check_methods_subset_invariance": "the assert_allclose check is done in float64 while Torch models operate in float32. The max absolute difference is 1.1920929e-07.",
                 "check_dont_overwrite_parameters": "the change of public attribute module__input_size is needed to support dynamic input size.",
-                "check_estimators_overwrite_params": "module parameters changes upon fitting the estimator hence produce non-identical result.",
-                "check_estimators_unfitted": "NeuralNetTorch does not support prediction without initializing the module.",
             },
         }
 
@@ -183,6 +172,8 @@ def fit(self, X, y, **fit_params):
 
     @torch.inference_mode()
     def predict_proba(self, X):
+        if not hasattr(self, "n_features_in_"):
+            raise NotFittedError
         dtype = X.dtype if hasattr(X, "dtype") else None
         X, _ = self.check_data(X)
         y_pred = super().predict_proba(X)

autoemulate/emulators/neural_networks/get_module.py

@@ -2,16 +2,16 @@
 from autoemulate.emulators.neural_networks.mlp import MLPModule
 
 
-def get_module(module: str | TorchModule, module_args) -> TorchModule:
+def get_module(module: str | TorchModule) -> TorchModule:
     """
-    Return the module instance for NeuralNetRegressor. If `module` is a string,
-    then initialize a TorchModule with the same registered name. If `module` is
+    Return the module class for NeuralNetRegressor. If `module` is
     already a TorchModule, then return it as is.
     """
-    if not isinstance(module, TorchModule):
-        match module:
-            case "mlp":
-                module = MLPModule(**module_args)
-            case _:
-                raise NotImplementedError(f"Module {module} not implemented.")
+    if not isinstance(module, str):
+        return module
+    match module:
+        case "mlp":
+            module = MLPModule
+        case _:
+            raise NotImplementedError(f"Module {module} not implemented.")
     return module

autoemulate/emulators/neural_networks/mlp.py

@@ -1,5 +1,6 @@
 from typing import Tuple
 
+import numpy as np
 import torch
 from scipy.stats import loguniform
 from skopt.space import Integer
@@ -17,7 +18,9 @@ def __init__(
         input_size: int = None,
         output_size: int = None,
         random_state: int = None,
-        hidden_sizes: Tuple[int] = (100,),
+        hidden_layers: int = 1,
+        hidden_size: int = 100,
+        hidden_activation: Tuple[callable] = nn.ReLU,
     ):
         super(MLPModule, self).__init__(
             module_name="mlp",
@@ -26,32 +29,34 @@ def __init__(
             random_state=random_state,
         )
         modules = []
-        for hidden_size in hidden_sizes:
+        assert hidden_layers >= 1
+        for _ in range(hidden_layers):
             modules.append(nn.Linear(in_features=input_size, out_features=hidden_size))
-            modules.append(nn.ReLU())
+            modules.append(hidden_activation())
             input_size = hidden_size
         modules.append(nn.Linear(in_features=input_size, out_features=output_size))
         self.model = nn.Sequential(*modules)
 
     def get_grid_params(self, search_type: str = "random"):
+        param_space = {
+            "max_epochs": np.arange(10, 110, 10).tolist(),
+            "batch_size": np.arange(2, 128, 2).tolist(),
+            "module__hidden_layers": np.arange(1, 4).tolist(),
+            "module__hidden_size": np.arange(50, 250, 50).tolist(),
+            "module__hidden_activation": [
+                nn.ReLU,
+                nn.Tanh,
+                nn.Sigmoid,
+                nn.GELU,
+            ],
+            "optimizer": [torch.optim.AdamW, torch.optim.SGD],
+            "optimizer__weight_decay": (1 / 10 ** np.arange(1, 9)).tolist(),
+        }
         match search_type:
             case "random":
-                param_space = {
-                    "lr": loguniform(1e-4, 1e-2),
-                    "max_epochs": [10, 20, 30],
-                    "module__hidden_sizes": [
-                        (50,),
-                        (100,),
-                        (100, 50),
-                        (100, 100),
-                        (200, 100),
-                    ],
-                }
+                param_space |= {"lr": loguniform(1e-06, 1e-2)}
             case "bayes":
-                param_space = {
-                    "lr": Real(1e-4, 1e-2, prior="log-uniform"),
-                    "max_epochs": Integer(10, 30),
-                }
+                param_space |= {"lr": Real(1e-06, 1e-2, prior="log-uniform")}
             case _:
                 raise ValueError(f"Invalid search type: {search_type}")
 

docs/_toc.yml

@@ -11,11 +11,17 @@ chapters:
   # - file: getting-started/data_requirements
   - file: getting-started/installation
   - file: getting-started/quickstart
+  - file: getting-started/choosing-design
 
 - file: tutorials/index
   sections:
   - file: tutorials/01_start
 
+- file: community/index
+  sections:
+  - file: community/contributing
+  - file: community/code-of-conduct
+
 - file: reference/index
   sections:
   - file: reference/compare

docs/community/code-of-conduct.md

@@ -0,0 +1,2 @@
+```{include} ../../CODE_OF_CONDUCT.md
+```

docs/community/contributing.md

@@ -0,0 +1,2 @@
+```{include} ../../CONTRIBUTING.md
+```

docs/community/index.md

@@ -0,0 +1,5 @@
+# Community
+
+The AutoEmulate community is a group of users and contributors who are interested in the development of AutoEmulate. The community is open to anyone who is interested in AutoEmulate. The community is a place to ask questions, discuss ideas, and share your work.
+
+We are currently building out our community infrastructure, but meanwhile we are very much welcoming contributions to the package. To get started, please read the [Contributing Guide](./contributing.md).

docs/getting-started/choosing-design.md

@@ -0,0 +1,5 @@
+# Set up your sampling/experimental design
+
+<!--
+To add here (see #6): the user currently has the responsibility for sampling / experimental design, and autoemulate only does the model fitting
+-->