use fit_from_initialized in AL, change types to DistributionLike from…

… GaussianLike to match TranformedEmulator predict type

autoemulate/learners/base.py

@@ -4,15 +4,15 @@
 
 import torch
 from anytree import Node, RenderTree
-from torch.distributions import MultivariateNormal
 from torcheval.metrics import MeanSquaredError, R2Score
 
 from autoemulate.core.logging_config import get_configured_logger
+from autoemulate.core.reinitialize import fit_from_reinitialized
 from autoemulate.data.utils import ValidationMixin
 from autoemulate.emulators.base import Emulator
 from autoemulate.simulations.base import Simulator
 
-from ..core.types import GaussianLike, TensorLike
+from ..core.types import DistributionLike, TensorLike
 
 
 @dataclass(kw_only=True)
@@ -48,7 +48,12 @@ def __post_init__(self):
         log_level = getattr(self, "log_level", "progress_bar")
         self.logger, self.progress_bar = get_configured_logger(log_level)
         self.logger.info("Initializing Learner with training data.")
-        self.emulator.fit(self.x_train, self.y_train)
+        self.emulator = fit_from_reinitialized(
+            self.x_train,
+            self.y_train,
+            emulator=self.emulator,
+            device=self.emulator.device,
+        )
         self.logger.info("Emulator fitted with initial training data.")
         self.in_dim = self.x_train.shape[1]
         self.out_dim = self.y_train.shape[1]
@@ -140,11 +145,9 @@ def fit(self, *args):
         x, output, extra = self.query(*args)
         if isinstance(output, TensorLike):
             y_pred = output
-        elif isinstance(output, GaussianLike):
+        elif isinstance(output, DistributionLike):
             assert output.variance.ndim == 2
-            y_pred, _ = output.mean, output.variance
-        elif isinstance(output, GaussianLike):
-            y_pred, _ = output.loc, None
+            y_pred = output.mean
         else:
             msg = (
                 f"Output must be either `Tensor` or `MultivariateNormal` but got "
@@ -159,7 +162,12 @@ def fit(self, *args):
             assert isinstance(y_true, TensorLike)
             self.x_train = torch.cat([self.x_train, x])
             self.y_train = torch.cat([self.y_train, y_true])
-            self.emulator.fit(self.x_train, self.y_train)
+            self.emulator = fit_from_reinitialized(
+                self.x_train,
+                self.y_train,
+                emulator=self.emulator,
+                device=self.emulator.device,
+            )
             self.mse.update(y_pred, y_true)
             self.r2.update(y_pred, y_true)
             self.n_queries += 1
@@ -181,7 +189,7 @@ def fit(self, *args):
 
         # If Gaussian output
         # TODO: check generality for other GPs (e.g. with full covariance)
-        if isinstance(output, MultivariateNormal):
+        if isinstance(output, DistributionLike):
             assert isinstance(output.variance, TensorLike)
             assert output.variance.ndim == 2
             assert output.variance.shape[1] == self.out_dim
@@ -257,7 +265,11 @@ def summary(self):
     @abstractmethod
     def query(
         self, x: TensorLike | None = None
-    ) -> tuple[TensorLike | None, TensorLike | GaussianLike, dict[str, float]]:
+    ) -> tuple[
+        TensorLike | None,
+        TensorLike | DistributionLike,
+        dict[str, float],
+    ]:
         """
         Abstract method to query new samples.
 
@@ -268,7 +280,7 @@ def query(
 
         Returns
         -------
-        tuple[TensorLike or None, TensorLike, TensorLike, Dict[str, list[Any]]]
+        tuple[TensorLike | None, TensorLike | DistributionLike, dict[str, float]]
             A tuple containing:
             - The queried samples (or None if no query is made),
             - The predicted outputs,

autoemulate/learners/stream.py

@@ -7,7 +7,7 @@
 
 from autoemulate.data.utils import set_random_seed
 
-from ..core.types import GaussianLike, TensorLike
+from ..core.types import DistributionLike, TensorLike
 from .base import Active
 
 
@@ -26,7 +26,7 @@ class Stream(Active):
     @abstractmethod
     def query(
         self, x: TensorLike | None = None
-    ) -> tuple[TensorLike | None, TensorLike | GaussianLike, dict[str, float]]:
+    ) -> tuple[TensorLike | None, TensorLike | DistributionLike, dict[str, float]]:
         """
         Abstract method to query new samples from a stream.
 
@@ -37,7 +37,7 @@ def query(
 
         Returns
         -------
-        Tuple[torch.Tensor or None, torch.Tensor, torch.Tensor, Dict[str, List[Any]]]
+        tuple[TensorLike | None, TensorLike | DistributionLike, dict[str, float]]
             A tuple containing:
             - The queried samples (or None if no query is made),
             - The predicted outputs,
@@ -64,9 +64,7 @@ def fit_samples(self, x: torch.Tensor):
             )
         ):
             self.fit(xi.reshape(1, -1))
-            pb.set_postfix(
-                ordered_dict={key: val[-1] for key, val in self.metrics.items()}
-            )
+            pb.set_postfix(ordered_dict={k: v[-1] for k, v in self.metrics.items()})
 
     def fit_batches(self, x: torch.Tensor, batch_size: int):
         """
@@ -114,7 +112,11 @@ def query(
         self,
         x: TensorLike | None = None,
         random_seed: int | None = None,
-    ) -> tuple[torch.Tensor | None, TensorLike | GaussianLike, dict[str, float]]:
+    ) -> tuple[
+        torch.Tensor | None,
+        TensorLike | DistributionLike,
+        dict[str, float],
+    ]:
         """
         Query new samples randomly based on a fixed probability.
 
@@ -125,7 +127,7 @@ def query(
 
         Returns
         -------
-        Tuple[torch.Tensor or None, torch.Tensor, torch.Tensor, Dict[str, List[Any]]]
+        tuple[TensorLike | None, TensorLike | DistributionLike, dict[str, float]]
             A tuple containing:
             - The queried samples (or None if the random condition is not met),
             - The predicted outputs,
@@ -135,8 +137,7 @@ def query(
         assert isinstance(x, TensorLike)
         # TODO: move handling to check method in base class
         output = self.emulator.predict(x)
-        assert isinstance(output, TensorLike | GaussianLike)
-        # assert isinstance(output, TensorLike | DistributionLike)
+        assert isinstance(output, TensorLike | DistributionLike)
         if random_seed is not None:
             set_random_seed(seed=random_seed)
         x = x if np.random.rand() < self.p_query else None
@@ -199,7 +200,7 @@ def query(
         self, x: TensorLike | None = None
     ) -> tuple[
         torch.Tensor | None,
-        torch.Tensor | GaussianLike,
+        torch.Tensor | DistributionLike,
         dict[str, float],
     ]:
         """
@@ -212,7 +213,7 @@ def query(
 
         Returns
         -------
-        Tuple[torch.Tensor or None, torch.Tensor, torch.Tensor, Dict[str, List[Any]]]
+        tuple[TensorLike | None, TensorLike | DistributionLike, dict[str, float]]
             A tuple containing:
             - The queried samples (or None if the score does not exceed the threshold),
             - The predicted outputs,
@@ -222,7 +223,7 @@ def query(
         # TODO: move handling to check method in base class
         assert isinstance(x, torch.Tensor)
         output = self.emulator.predict(x)
-        assert isinstance(output, GaussianLike)
+        assert isinstance(output, DistributionLike)
         assert isinstance(output.variance, torch.Tensor)
         score = self.score(x, output.mean, output.variance)
         x = x if score > self.threshold else None
@@ -457,7 +458,11 @@ def __post_init__(self):
 
     def query(
         self, x: TensorLike | None = None
-    ) -> tuple[TensorLike | None, TensorLike | GaussianLike, dict[str, float]]:
+    ) -> tuple[
+        TensorLike | None,
+        TensorLike | DistributionLike,
+        dict[str, float],
+    ]:
         """
         Query new samples based on the adaptive threshold.
 
@@ -471,7 +476,7 @@ def query(
 
         Returns
         -------
-        Tuple[torch.Tensor or None, torch.Tensor, torch.Tensor, Dict[str, List[Any]]]
+        tuple[TensorLike | None, TensorLike | DistributionLike, dict[str, float]]
             A tuple containing:
             - The queried samples,
             - The predicted outputs,