Merge pull request #900 from alan-turing-institute/879-registry

Add emulator registry (#879)

Author: sgreenbury

autoemulate/core/compare.py

@@ -37,12 +37,7 @@
     TransformedEmulatorParams,
 )
 from autoemulate.data.utils import ConversionMixin, set_random_seed
-from autoemulate.emulators import (
-    ALL_EMULATORS,
-    DEFAULT_EMULATORS,
-    PYTORCH_EMULATORS,
-    get_emulator_class,
-)
+from autoemulate.emulators import _default_registry, get_emulator_class
 from autoemulate.emulators.base import Emulator, ProbabilisticEmulator
 from autoemulate.emulators.transformed.base import TransformedEmulator
 from autoemulate.transforms.base import AutoEmulateTransform
@@ -202,22 +197,26 @@ def __init__(
     @staticmethod
     def all_emulators() -> list[type[Emulator]]:
         """Return a list of all available emulators."""
-        return ALL_EMULATORS
+        return _default_registry.all_emulators
 
     @staticmethod
     def default_emulators() -> list[type[Emulator]]:
         """Return a list of default emulators used by AutoEmulate."""
-        return DEFAULT_EMULATORS
+        return _default_registry.default_emulators
 
     @staticmethod
     def pytorch_emulators() -> list[type[Emulator]]:
         """Return a list of all available PyTorch emulators."""
-        return PYTORCH_EMULATORS
+        return _default_registry.pytorch_emulators
 
     @staticmethod
     def probablistic_emulators() -> list[type[Emulator]]:
         """Return a list of all available probabilistic emulators."""
-        return [emulator for emulator in ALL_EMULATORS if emulator.supports_uq]
+        return [
+            emulator
+            for emulator in _default_registry.all_emulators
+            if emulator.supports_uq
+        ]
 
     @staticmethod
     def list_emulators(default_only: bool = True) -> pd.DataFrame:
@@ -243,7 +242,11 @@ def list_emulators(default_only: bool = True) -> pd.DataFrame:
                 - 'Uncertainty_Quantification',
                 - 'Automatic_Differentiation`
         """
-        emulator_set = DEFAULT_EMULATORS if default_only else ALL_EMULATORS
+        emulator_set = (
+            _default_registry.default_emulators
+            if default_only
+            else _default_registry.all_emulators
+        )
         return pd.DataFrame(
             {
                 "Emulator": [emulator.model_name() for emulator in emulator_set],

autoemulate/emulators/__init__.py

@@ -1,4 +1,4 @@
-from .base import Emulator, GaussianProcessEmulator
+from .base import Emulator
 from .ensemble import EnsembleMLP, EnsembleMLPDropout
 from .gaussian_process.exact import (
     GaussianProcessCorrelatedMatern32,
@@ -11,74 +11,22 @@
 from .polynomials import PolynomialRegression
 from .radial_basis_functions import RadialBasisFunctions
 from .random_forest import RandomForest
+from .registry import Registry, _default_registry, get_emulator_class, register
 from .svm import SupportVectorMachine
 from .transformed.base import TransformedEmulator
 
-DEFAULT_EMULATORS: list[type[Emulator]] = [
-    GaussianProcessMatern32,
-    GaussianProcessRBF,
-    RadialBasisFunctions,
-    PolynomialRegression,
-    MLP,
-    EnsembleMLP,
-]
-
-# listing non pytorch emulators as we do not expect this list to grow
-NON_PYTORCH_EMULATORS: list[type[Emulator]] = [
-    LightGBM,
-    SupportVectorMachine,
-    RandomForest,
-]
-
-ALL_EMULATORS: list[type[Emulator]] = [
-    *DEFAULT_EMULATORS,
-    *NON_PYTORCH_EMULATORS,
-    GaussianProcessCorrelatedMatern32,
-    GaussianProcessCorrelatedRBF,
-    EnsembleMLPDropout,
-]
-
-PYTORCH_EMULATORS: list[type[Emulator]] = [
-    emulator for emulator in ALL_EMULATORS if emulator not in NON_PYTORCH_EMULATORS
-]
-GAUSSIAN_PROCESS_EMULATORS: list[type[Emulator]] = [
-    emulator
-    for emulator in ALL_EMULATORS
-    if issubclass(emulator, GaussianProcessEmulator)
-]
-
-EMULATOR_REGISTRY = {em_cls.model_name().lower(): em_cls for em_cls in ALL_EMULATORS}
-EMULATOR_REGISTRY_SHORT_NAME = {em_cls.short_name(): em_cls for em_cls in ALL_EMULATORS}
-
-
-def get_emulator_class(name: str) -> type[Emulator]:
-    """
-    Get the emulator class by name.
-
-    Parameters
-    ----------
-    name: str
-        The name of the emulator class.
-
-    Returns
-    -------
-    type[Emulator] | None
-        The emulator class if found, None otherwise.
-    """
-    emulator_cls = EMULATOR_REGISTRY.get(
-        name.lower()
-    ) or EMULATOR_REGISTRY_SHORT_NAME.get(name.lower())
-
-    if emulator_cls is None:
-        raise ValueError(
-            f"Unknown emulator name: {name}.Available: {list(EMULATOR_REGISTRY.keys())}"
-        )
-
-    return emulator_cls
-
+# Module-level constants for backward compatibility and simplified public access
+DEFAULT_EMULATORS = _default_registry._default_emulators
+NON_PYTORCH_EMULATORS = _default_registry._non_pytorch_emulators
+ALL_EMULATORS = _default_registry._all_emulators
+PYTORCH_EMULATORS = _default_registry._pytorch_emulators
+GAUSSIAN_PROCESS_EMULATORS = _default_registry._gaussian_process_emulators
+EMULATOR_REGISTRY = _default_registry._emulator_registry
+EMULATOR_REGISTRY_SHORT_NAME = _default_registry._emulator_registry_short_name
 
 __all__ = [
     "MLP",
+    "Emulator",
     "EnsembleMLP",
     "EnsembleMLPDropout",
     "GaussianProcessCorrelatedMatern32",
@@ -89,6 +37,9 @@ def get_emulator_class(name: str) -> type[Emulator]:
     "PolynomialRegression",
     "RadialBasisFunctions",
     "RandomForest",
+    "Registry",
     "SupportVectorMachine",
     "TransformedEmulator",
+    "get_emulator_class",
+    "register",
 ]

autoemulate/emulators/gaussian_process/exact.py

@@ -476,18 +476,13 @@ def forward(self, x):
         return GaussianProcessLike(mean_x, covar_x)
 
 
-# GP registry to raise exception if duplicate created
-GP_REGISTRY = {
-    "GaussianProcess": GaussianProcess,
-    "GaussianProcessCorrelated": GaussianProcessCorrelated,
-}
-
-
 def create_gp_subclass(
     name: str,
     gp_base_class: type[GaussianProcess],
     covar_module_fn: CovarModuleFn,
     mean_module_fn: MeanModuleFn = constant_mean,
+    auto_register: bool = True,
+    overwrite: bool = True,
     **fixed_kwargs,
 ) -> type[GaussianProcess]:
     """
@@ -496,6 +491,9 @@ def create_gp_subclass(
     This function creates a subclass of GaussianProcess where certain parameters
     are fixed to specific values, reducing the parameter space for tuning.
 
+    The created subclass is automatically registered with the main emulator Registry
+    (unless auto_register=False), making it discoverable by AutoEmulate.
+
     Parameters
     ----------
     name : str
@@ -506,6 +504,13 @@ def create_gp_subclass(
         Covariance module function to use in the subclass.
     mean_module_fn : MeanModuleFn
         Mean module function to use in the subclass. Defaults to `constant_mean`.
+    auto_register : bool
+        Whether to automatically register the created subclass with the main emulator
+        Registry. Defaults to True.
+    overwrite : bool
+        Whether to allow overwriting an existing class with the same name in the
+        main Registry. Useful for interactive development in notebooks. Defaults to
+        True.
     **fixed_kwargs
         Keyword arguments to fix in the subclass. These parameters will be
         set to the provided values and excluded from hyperparameter tuning.
@@ -516,17 +521,21 @@ def create_gp_subclass(
         A new subclass of GaussianProcess with the specified parameters fixed.
         The returned class can be pickled and used like any other GP emulator.
 
+    Raises
+    ------
+    ValueError
+        If `name` matches `model_name()` or `short_name()` of an already registered
+        emulator in the main Registry and `overwrite=False`.
+
     Notes
     -----
-    Fixed parameters are automatically excluded from `get_tune_params()` to
-    prevent them from being included in hyperparameter optimization.
+    - Fixed parameters are automatically excluded from `get_tune_params()` to prevent
+    them from being included in hyperparameter optimization.
+    - Pickling: The created subclass is registered in the caller's module namespace,
+    ensuring it can be pickled and unpickled correctly even when created in downstream
+    code that uses autoemulate as a dependency.
+    - If auto_register=True (default), the class is also added to the main Registry.
     """
-    if name in GP_REGISTRY:
-        raise ValueError(
-            f"A GP class named '{name}' already exists. "
-            f"Use a unique name or delete the existing class from GP_REGISTRY."
-        )
-
     standardize_x = fixed_kwargs.get("standardize_x", False)
     standardize_y = fixed_kwargs.get("standardize_y", True)
     fixed_mean_params = fixed_kwargs.get("fixed_mean_params", False)
@@ -614,30 +623,50 @@ def get_tune_params():
     model training and are not fixed.
     """
 
-    # Set the provided name for the class
+    # Determine the caller's module for proper pickling support.
+    # When called from autoemulate itself, use __name__.
+    # When called from user code, use the caller's module
+    caller_frame = sys._getframe(1)
+    caller_module_name = caller_frame.f_globals.get("__name__", __name__)
+
+    # Set the class name and module
     GaussianProcessSubclass.__name__ = name
     GaussianProcessSubclass.__qualname__ = name
-    GaussianProcessSubclass.__module__ = __name__
+    GaussianProcessSubclass.__module__ = caller_module_name
+
+    # Register class in the caller's module globals for pickling
+    # This ensures the class can be pickled/unpickled correctly
+    caller_frame.f_globals[name] = GaussianProcessSubclass
+
+    # Also register in the caller's module if it's a real module (not __main__)
+    if caller_module_name in sys.modules and caller_module_name != "__main__":
+        setattr(sys.modules[caller_module_name], name, GaussianProcessSubclass)
+
+    # Automatically register with the main emulator Registry if requested
+    if auto_register:
+        # Lazy import to avoid circular dependency with __init__.py
+        from autoemulate.emulators import register  # noqa: PLC0415
 
-    # Register class in the module's globals so can be pickled
-    setattr(sys.modules[__name__], name, GaussianProcessSubclass)
-    # Register subclass
-    GP_REGISTRY[name] = GaussianProcessSubclass
+        register(GaussianProcessSubclass, overwrite=overwrite)
 
     return GaussianProcessSubclass
 
 
+# Built-in GP subclasses - auto_register=False as already registered in Registry init:
+# autoemulate/emulators/__init__.py
 GaussianProcessRBF = create_gp_subclass(
     "GaussianProcessRBF",
     GaussianProcess,
     covar_module_fn=rbf_kernel,
     mean_module_fn=constant_mean,
+    auto_register=False,
 )
 GaussianProcessMatern32 = create_gp_subclass(
     "GaussianProcessMatern32",
     GaussianProcess,
     covar_module_fn=matern_3_2_kernel,
     mean_module_fn=constant_mean,
+    auto_register=False,
 )
 
 # correlated GP kernels
@@ -646,10 +675,12 @@ def get_tune_params():
     GaussianProcessCorrelated,
     covar_module_fn=rbf_kernel,
     mean_module_fn=constant_mean,
+    auto_register=False,
 )
 GaussianProcessCorrelatedMatern32 = create_gp_subclass(
     "GaussianProcessCorrelatedMatern32",
     GaussianProcessCorrelated,
     covar_module_fn=matern_3_2_kernel,
     mean_module_fn=constant_mean,
+    auto_register=False,
 )