Improve emulator re-initialisation

autoemulate/calibration/history_matching.py

@@ -12,10 +12,11 @@
 from autoemulate.core.device import TorchDeviceMixin
 from autoemulate.core.logging_config import get_configured_logger
 from autoemulate.core.plotting import display_figure
+from autoemulate.core.reinitialize import fit_from_reinitialized
 from autoemulate.core.results import Result
 from autoemulate.core.types import DeviceLike, DistributionLike, TensorLike
 from autoemulate.data.utils import set_random_seed
-from autoemulate.emulators import TransformedEmulator, get_emulator_class
+from autoemulate.emulators import Emulator
 from autoemulate.simulations.base import Simulator
 
 logger = logging.getLogger("autoemulate")
@@ -284,13 +285,15 @@ class HistoryMatchingWorkflow(HistoryMatching):
     def __init__(
         self,
         simulator: Simulator,
-        result: Result,
         observations: dict[str, tuple[float, float]] | dict[str, float],
+        emulator: Emulator | None = None,
+        result: Result | None = None,
         threshold: float = 3.0,
         model_discrepancy: float = 0.0,
         rank: int = 1,
         train_x: TensorLike | None = None,
         train_y: TensorLike | None = None,
+        transformed_emulator_params: dict | None = None,
         calibration_params: list[str] | None = None,
         device: DeviceLike | None = None,
         random_seed: int | None = None,
@@ -303,12 +306,17 @@ def __init__(
         ----------
         simulator: Simulator
             A simulator.
-        result: Result
-            A Result object containing the pre-trained emulator and its hyperparameters.
         observations: dict[str, tuple[float, float] | dict[str, float]
             For each output variable, specifies observed [value, noise] (with noise
             specified as variances). In case of no uncertainty in observations, provides
             just the observed value.
+        emulator: Emulator | None
+            An Emulator object containing the pre-trained emulator. If not provided, a
+            Result object must be provided instead. Defaults to None.
+        result: Result | None
+            A Result object containing the pre-trained emulator and its hyperparameters.
+            If not provided, an Emulator object must be provided instead.
+            Defaults to None.
         threshold: float
             Implausibility threshold (query points with implausibility scores that
             exceed this value are ruled out). Defaults to 3, which is considered
@@ -324,6 +332,11 @@ def __init__(
             Optional tensor of input data the emulator was trained on.
         train_y: TensorLike | None
             Optional tensor of output data the emulator was trained on.
+        transformed_emulator_params: dict | None
+            Optional dictionary of parameters for TransformedEmulator. These are
+            already contained in a Result object, so only needed if a
+            TransformedEmulator instance with non-default params is provided.
+            Defaults to None.
         calibration_params: list[str] | None
             Optional subset of parameters to calibrate. These have to correspond to the
             parameters that the emulator was trained on. If None, calibrate all
@@ -342,8 +355,15 @@ def __init__(
             set_random_seed(seed=random_seed)
         self.logger, self.progress_bar = get_configured_logger(log_level)
 
-        self.result = result
-        self.emulator = result.model
+        if result is not None:
+            self.emulator = result.model
+        elif emulator is not None:
+            self.emulator = emulator
+        else:
+            msg = "Either `emulator` or `result` must be provided."
+            raise ValueError(msg)
+
+        self.transformed_emulator_params = transformed_emulator_params or {}
         self.emulator.device = self.device
 
         # New data is simulated in `run()` and appended here
@@ -623,19 +643,16 @@ def refit_emulator(self, x: TensorLike, y: TensorLike) -> None:
         y: TensorLike
             Tensor of output data to refit the emulator on.
         """
-        # Create a fresh model with the same configuration
-        self.emulator = TransformedEmulator(
-            x.float(),
-            y.float(),
-            model=get_emulator_class(self.result.model_name),
-            x_transforms=self.result.x_transforms,
-            y_transforms=self.result.y_transforms,
+        x = x.float().to(self.device)
+        y = y.float().to(self.device)
+        self.emulator = fit_from_reinitialized(
+            x,
+            y,
+            emulator=self.emulator,
+            transformed_emulator_params=self.transformed_emulator_params,
             device=self.device,
-            **self.result.params,
         )
 
-        self.emulator.fit(x, y)
-
     def run(
         self,
         n_simulations: int = 100,

autoemulate/core/compare.py

@@ -19,6 +19,7 @@
     display_figure,
     plot_xy,
 )
+from autoemulate.core.reinitialize import fit_from_reinitialized
 from autoemulate.core.results import Result, Results
 from autoemulate.core.save import ModelSerialiser
 from autoemulate.core.tuner import Tuner
@@ -566,34 +567,21 @@ def fit_from_reinitialized(
         # Get the result to use
         result = self.best_result() if result_id is None else self.get_result(result_id)
 
-        # Set the random seed for initialization
-        if random_seed is not None:
-            set_random_seed(seed=random_seed)
-
         # Convert and move the new data to device
         x_tensor, y_tensor = self._convert_to_tensors(x, y)
         x_tensor, y_tensor = self._move_tensors_to_device(x_tensor, y_tensor)
 
-        # Get the model class from the model name
-        model_class = get_emulator_class(result.model_name)
-
-        # Create a fresh model with the same configuration
-        fresh_model = TransformedEmulator(
+        # NOTE: function passes data to the Emulator model which handles conversion to
+        # tensors and device handling
+        return fit_from_reinitialized(
             x_tensor,
             y_tensor,
-            model=model_class,
-            x_transforms=result.x_transforms,
-            y_transforms=result.y_transforms,
+            emulator=result.model,
+            transformed_emulator_params=transformed_emulator_params,
             device=self.device,
-            **result.params,
-            **transformed_emulator_params,
+            random_seed=random_seed,
         )
 
-        # Fit the fresh model on the new data
-        fresh_model.fit(x_tensor, y_tensor)
-
-        return fresh_model
-
     def plot(  # noqa: PLR0912, PLR0915
         self,
         model_obj: int | Emulator | Result,

autoemulate/core/reinitialize.py

@@ -0,0 +1,99 @@
+import inspect
+
+from autoemulate.core.types import DeviceLike, TensorLike
+from autoemulate.data.utils import set_random_seed
+from autoemulate.emulators import Emulator, TransformedEmulator, get_emulator_class
+
+
+def fit_from_reinitialized(
+    x: TensorLike,
+    y: TensorLike,
+    emulator: Emulator,
+    transformed_emulator_params: dict | None = None,
+    device: DeviceLike | None = None,
+    random_seed: int | None = None,
+):
+    """
+    Fit a fresh model with reinitialized parameters using the best configuration.
+
+    This method creates a new model instance with the same configuration as the
+    best (or specified) model from the comparison, but with freshly initialized
+    parameters fitted on the provided data.
+
+    Parameters
+    ----------
+    x: TensorLike
+        Input features for training the fresh model.
+    y: TensorLike
+        Target values for training the fresh model.
+    emulator: Emulator
+        An Emulator object containing the pre-trained emulator.
+    transformed_emulator_params: None | TransformedEmulatorParams
+        Parameters for the transformed emulator. When None, the same parameters as
+        used when identifying the best model are used. Defaults to None.
+    device: str | None
+        Device to use for model fitting (e.g., 'cpu' or 'cuda'). If None, the default
+        device is used. Defaults to None.
+    random_seed: int | None
+        Random seed for parameter initialization. Defaults to None.
+
+    Returns
+    -------
+    TransformedEmulator
+        A new model instance with the same configuration but fresh parameters
+        fitted on the provided data.
+
+    Notes
+    -----
+    Unlike TransformedEmulator.refit() which retrains an existing model,
+    this method creates a completely new model instance with reinitialized
+    parameters. This ensures that when fitting on new data that the same
+    initialization conditions are applied. This can have an affect for example
+    given kernel initialization in Gaussian Processes or weight initialization in
+    neural networks.
+    """
+    if random_seed is not None:
+        set_random_seed(seed=random_seed)
+
+    # Extract emulator and its parameters from Emulator instance
+    if isinstance(emulator, TransformedEmulator):
+        model = emulator.model
+        emulator_name = emulator.untransformed_model_name
+        x_transforms = emulator.x_transforms
+        y_transforms = emulator.y_transforms
+    else:
+        model = emulator
+        emulator_name = emulator.model_name()
+        x_transforms = None
+        y_transforms = None
+
+    # Extract parameters from the provided emulator instance
+    model_cls = get_emulator_class(emulator_name)
+    init_sig = inspect.signature(model_cls.__init__)
+    emulator_params = {}
+    for param_name in init_sig.parameters:
+        if param_name in ["self", "x", "y", "device"]:
+            continue
+        # NOTE: some emulators have standardize_x/y params option
+        # this is different to TransformedEmulator x/y transforms
+        if param_name == "standardize_x":
+            emulator_params["standardize_x"] = bool(model.x_transform)
+        if param_name == "standardize_y":
+            emulator_params["standardize_y"] = bool(model.y_transform)
+        if hasattr(model, param_name):
+            emulator_params[param_name] = getattr(model, param_name)
+
+    transformed_emulator_params = transformed_emulator_params or {}
+
+    new_emulator = TransformedEmulator(
+        x.float(),
+        y.float(),
+        model=model_cls,
+        x_transforms=x_transforms,
+        y_transforms=y_transforms,
+        device=device,
+        **emulator_params,
+    )
+
+    new_emulator.fit(x.float(), y.float())
+    return new_emulator