Fall back to BayBE fit criterion for unspecified cases

Author: AdrianSosic

baybe/surrogates/gaussian_process/components/fit_criterion.py

@@ -5,10 +5,16 @@
 from enum import Enum
 from typing import TYPE_CHECKING
 
+from attrs import define
+from typing_extensions import override
+
 if TYPE_CHECKING:
     from gpytorch.likelihoods import Likelihood as GPyTorchLikelihood
     from gpytorch.mlls import MarginalLogLikelihood
     from gpytorch.models import GP as GPyTorchModel
+    from torch import Tensor
+
+    from baybe.searchspace.core import SearchSpace
 
 
 class FitCriterion(Enum):
@@ -44,3 +50,25 @@ def to_gpytorch(
 
 PlainFitCriterionFactory = PlainGPComponentFactory[FitCriterion]
 """A trivial factory that returns a fixed fit criterion."""
+
+
+@define
+class _MLLForNonTLFitCriterionFactory(FitCriterionFactoryProtocol):
+    """A fit criterion factory switching between MLL and BayBE depending on the context.
+
+    In transfer learning contexts, delegates to
+    :class:`baybe.surrogates.gaussian_process.presets.baybe.BayBEFitCriterionFactory`.
+    """
+
+    @override
+    def __call__(
+        self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
+    ) -> FitCriterion:
+        if searchspace.task_idx is None:
+            return FitCriterion.MARGINAL_LOG_LIKELIHOOD
+
+        from baybe.surrogates.gaussian_process.presets.baybe import (
+            BayBEFitCriterionFactory,
+        )
+
+        return BayBEFitCriterionFactory()(searchspace, train_x, train_y)

baybe/surrogates/gaussian_process/presets/chen.py

@@ -19,8 +19,7 @@
 )
 from baybe.priors.basic import GammaPrior
 from baybe.surrogates.gaussian_process.components.fit_criterion import (
-    FitCriterion,
-    PlainFitCriterionFactory,
+    _MLLForNonTLFitCriterionFactory,
 )
 from baybe.surrogates.gaussian_process.components.kernel import (
     _PureKernelFactory,
@@ -72,7 +71,7 @@ def _make(
         )
 
 
-CHENFitCriterionFactory = PlainFitCriterionFactory(FitCriterion.MARGINAL_LOG_LIKELIHOOD)
+CHENFitCriterionFactory = _MLLForNonTLFitCriterionFactory()
 """A factory providing fitting criteria for the CHEN preset."""
 
 # Collect leftover original slotted classes processed by `attrs.define`

baybe/surrogates/gaussian_process/presets/edbo.py

@@ -22,8 +22,7 @@
 from baybe.priors.basic import GammaPrior
 from baybe.searchspace.discrete import SubspaceDiscrete
 from baybe.surrogates.gaussian_process.components.fit_criterion import (
-    FitCriterion,
-    PlainFitCriterionFactory,
+    _MLLForNonTLFitCriterionFactory,
 )
 from baybe.surrogates.gaussian_process.components.kernel import (
     _PureKernelFactory,
@@ -179,10 +178,10 @@ def __call__(
         return likelihood
 
 
-# Collect leftover original slotted classes processed by `attrs.define`
-EDBOFitCriterionFactory = PlainFitCriterionFactory(FitCriterion.MARGINAL_LOG_LIKELIHOOD)
+EDBOFitCriterionFactory = _MLLForNonTLFitCriterionFactory()
 """A factory providing fitting criteria for the EDBO preset."""
 
+# Collect leftover original slotted classes processed by `attrs.define`
 gc.collect()
 
 # Preset defaults

baybe/surrogates/gaussian_process/presets/edbo_smoothed.py

@@ -19,8 +19,7 @@
 )
 from baybe.priors.basic import GammaPrior
 from baybe.surrogates.gaussian_process.components.fit_criterion import (
-    FitCriterion,
-    PlainFitCriterionFactory,
+    _MLLForNonTLFitCriterionFactory,
 )
 from baybe.surrogates.gaussian_process.components.kernel import (
     _PureKernelFactory,
@@ -130,9 +129,7 @@ def __call__(
         return likelihood
 
 
-SmoothedEDBOFitCriterionFactory = PlainFitCriterionFactory(
-    FitCriterion.MARGINAL_LOG_LIKELIHOOD
-)
+SmoothedEDBOFitCriterionFactory = _MLLForNonTLFitCriterionFactory()
 """A factory providing fitting criteria for the smoothed EDBO preset."""
 
 # Collect leftover original slotted classes processed by `attrs.define`