Add cont_kernel_factory argument to SingleTaskMultifidelityGP (meta-pytorch#3167)

Summary:
## Motivation

See meta-pytorch#3158

### Have you read the [Contributing Guidelines on pull requests](https://github.com/meta-pytorch/botorch/blob/main/CONTRIBUTING.md#pull-requests)?

Yes

Pull Request resolved: meta-pytorch#3167

Test Plan:
Written tests in `test_gp_regression_fidelity` (specifically `test_cont_kernel_factory`) to verify behaviour

## Related PRs

N/A

Reviewed By: saitcakmak

Differential Revision: D92831698

Pulled By: hvarfner

fbshipit-source-id: e7284931f9fe235de69f7db206354b80632db61b

Author: TobyBoyne

botorch/models/gp_regression_fidelity.py

@@ -44,6 +44,7 @@
 from gpytorch.kernels.kernel import ProductKernel
 from gpytorch.kernels.scale_kernel import ScaleKernel
 from gpytorch.likelihoods.likelihood import Likelihood
+from gpytorch.module import Module
 from gpytorch.priors.torch_priors import GammaPrior
 from torch import Tensor
 
@@ -72,6 +73,7 @@ def __init__(
         data_fidelities: Sequence[int] | None = None,
         linear_truncated: bool = True,
         nu: float = 2.5,
+        covar_module: Module | None = None,
         likelihood: Likelihood | None = None,
         outcome_transform: OutcomeTransform | _DefaultType | None = DEFAULT,
         input_transform: InputTransform | None = None,
@@ -93,6 +95,8 @@ def __init__(
                 of the default kernel.
             nu: The smoothness parameter for the Matern kernel: either 1/2, 3/2, or
                 5/2. Only used when ``linear_truncated=True``.
+            covar_module: The module for computing the covariance matrix between
+                the non-fidelity features. Defaults to ``RBFKernel``.
             likelihood: A likelihood. If omitted, use a standard GaussianLikelihood
                 with inferred noise level.
             outcome_transform: An outcome transform that is applied to the
@@ -130,6 +134,7 @@ def __init__(
             data_fidelities=data_fidelities,
             linear_truncated=linear_truncated,
             nu=nu,
+            data_covar_module=covar_module,
         )
         super().__init__(
             train_X=train_X,
@@ -143,9 +148,10 @@ def __init__(
         # Used for subsetting along the output dimension. See Model.subset_output.
         self._subset_batch_dict = {
             "mean_module.raw_constant": -1,
-            "covar_module.raw_outputscale": -1,
             **subset_batch_dict,
         }
+        if linear_truncated:
+            self._subset_batch_dict["covar_module.raw_outputscale"] = -1
         if train_Yvar is None:
             self._subset_batch_dict["likelihood.noise_covar.raw_noise"] = -2
         self.to(train_X)
@@ -174,7 +180,8 @@ def _setup_multifidelity_covar_module(
     data_fidelities: Sequence[int] | None,
     linear_truncated: bool,
     nu: float,
-) -> tuple[ScaleKernel, dict]:
+    data_covar_module: Module | None = None,
+) -> tuple[ScaleKernel | ProductKernel, dict[str, int]]:
     """Helper function to get the covariance module and associated subset_batch_dict
     for the multifidelity setting.
 
@@ -190,7 +197,8 @@ def _setup_multifidelity_covar_module(
             of the default kernel.
         nu: The smoothness parameter for the Matern kernel: either 1/2, 3/2, or
             5/2. Only used when ``linear_truncated=True``.
-
+        data_covar_module: The module for computing the covariance matrix between
+                the non-fidelity features. Defaults to ``RBFKernel``.
     Returns:
         The covariance module and subset_batch_dict.
     """
@@ -205,6 +213,12 @@ def _setup_multifidelity_covar_module(
 
     kernels = []
 
+    if linear_truncated and data_covar_module is not None:
+        raise ValueError(
+            "Non-fidelity covariance module cannot be specified when using a linear "
+            "truncated kernel."
+        )
+
     if linear_truncated:
         leading_dims = [iteration_fidelity] if iteration_fidelity is not None else []
         trailing_dims = (
@@ -225,13 +239,18 @@ def _setup_multifidelity_covar_module(
         if iteration_fidelity is not None:
             non_active_dims.add(iteration_fidelity)
         active_dimsX = sorted(set(range(dim)) - non_active_dims)
-        kernels.append(
-            get_covar_module_with_dim_scaled_prior(
-                ard_num_dims=len(active_dimsX),
-                batch_shape=aug_batch_shape,
-                active_dims=active_dimsX,
+
+        if data_covar_module is None:
+            kernels.append(
+                get_covar_module_with_dim_scaled_prior(
+                    ard_num_dims=len(active_dimsX),
+                    batch_shape=aug_batch_shape,
+                    active_dims=active_dimsX,
+                )
             )
-        )
+        else:
+            kernels.append(data_covar_module)
+
         if iteration_fidelity is not None:
             kernels.append(
                 ExponentialDecayKernel(
@@ -255,11 +274,15 @@ def _setup_multifidelity_covar_module(
 
     kernel = ProductKernel(*kernels)
 
-    covar_module = ScaleKernel(
-        kernel, batch_shape=aug_batch_shape, outputscale_prior=GammaPrior(2.0, 0.15)
-    )
+    if linear_truncated:
+        covar_module = ScaleKernel(
+            kernel, batch_shape=aug_batch_shape, outputscale_prior=GammaPrior(2.0, 0.15)
+        )
+        key_prefix = "covar_module.base_kernel.kernels"
+    else:
+        covar_module = kernel
+        key_prefix = "covar_module.kernels"
 
-    key_prefix = "covar_module.base_kernel.kernels"
     if linear_truncated:
         subset_batch_dict = {}
         for i in range(len(kernels)):
@@ -271,9 +294,15 @@ def _setup_multifidelity_covar_module(
                 }
             )
     else:
-        subset_batch_dict = {
-            f"{key_prefix}.0.raw_lengthscale": -3,
-        }
+        subset_batch_dict = {}
+
+        # Only set the subset_batch_dict if using the default kernel. See SingleTaskGP.
+        if data_covar_module is None:
+            subset_batch_dict.update(
+                {
+                    f"{key_prefix}.0.raw_lengthscale": -3,
+                }
+            )
 
         if iteration_fidelity is not None:
             subset_batch_dict.update(

test/models/test_gp_regression_fidelity.py

@@ -17,6 +17,9 @@
 from botorch.sampling import SobolQMCNormalSampler
 from botorch.utils.datasets import SupervisedDataset
 from botorch.utils.testing import BotorchTestCase, get_random_data
+from gpytorch.kernels.kernel import ProductKernel
+from gpytorch.kernels.linear_kernel import LinearKernel
+from gpytorch.kernels.rbf_kernel import RBFKernel
 from gpytorch.kernels.scale_kernel import ScaleKernel
 from gpytorch.likelihoods import FixedNoiseGaussianLikelihood
 from gpytorch.means import ConstantMean
@@ -147,7 +150,9 @@ def test_gp(self) -> None:
 
                 # test init
                 self.assertIsInstance(model.mean_module, ConstantMean)
-                self.assertIsInstance(model.covar_module, ScaleKernel)
+                self.assertIsInstance(
+                    model.covar_module, ScaleKernel if lin_trunc else ProductKernel
+                )
                 if use_octf:
                     self.assertIsInstance(model.outcome_transform, Standardize)
                 if use_intf:
@@ -427,6 +432,68 @@ def test_construct_inputs(self):
                 self.assertTrue(kwargs["train_X"].equal(data_dict["train_X"]))
                 self.assertTrue(kwargs["train_Y"].equal(data_dict["train_Y"]))
 
+    def test_custom_data_covar_module(self):
+        iteration_fidelity, data_fidelities = self.FIDELITY_TEST_PAIRS[0]
+        d = 1
+        m = 2
+        batch_shape = torch.Size([2, 3])
+        aug_batch_shape = batch_shape + torch.Size([m])
+        tkwargs = {"device": self.device, "dtype": torch.double}
+
+        _, model_kwargs = self._get_model_and_data(
+            iteration_fidelity=iteration_fidelity,
+            data_fidelities=data_fidelities,
+            batch_shape=batch_shape,
+            m=m,
+            lin_truncated=False,
+            **tkwargs,
+        )
+
+        active_dimsX = [0]
+        linear_kernel = LinearKernel(
+            ard_num_dims=len(active_dimsX),
+            batch_shape=aug_batch_shape,
+            active_dims=active_dimsX,
+        )
+
+        expected_kernels = (RBFKernel, LinearKernel)
+        expected_subset_batch_dict = {
+            "mean_module.raw_constant": -1,
+            "covar_module.kernels.1.raw_power": -2,
+            "covar_module.kernels.1.raw_offset": -2,
+        }
+        if model_kwargs.get("train_Yvar", None) is None:
+            expected_subset_batch_dict["likelihood.noise_covar.raw_noise"] = -2
+
+        for covar_module, expected_kernel in zip(
+            (None, linear_kernel), expected_kernels
+        ):
+            model = SingleTaskMultiFidelityGP(**model_kwargs, covar_module=covar_module)
+            non_fid_kernel = model.covar_module.kernels[0]
+            self.assertIsInstance(non_fid_kernel, expected_kernel)
+            self.assertEqual(non_fid_kernel.ard_num_dims, d)
+            self.assertEqual(non_fid_kernel.batch_shape, aug_batch_shape)
+            self.assertEqual(
+                non_fid_kernel.active_dims,
+                torch.as_tensor(list(range(d)), device=self.device),
+            )
+
+            extra_subset_batch_dict = (
+                {
+                    "covar_module.kernels.0.raw_lengthscale": -3,
+                }
+                if covar_module is None
+                else {}
+            )
+            self.assertEqual(
+                model._subset_batch_dict,
+                {**expected_subset_batch_dict, **extra_subset_batch_dict},
+            )
+
+        model_kwargs.update({"linear_truncated": True})
+        with self.assertRaises(ValueError):
+            model = SingleTaskMultiFidelityGP(**model_kwargs, covar_module=covar_module)
+
 
 class TestFixedNoiseSingleTaskMultiFidelityGP(TestSingleTaskMultiFidelityGP):
     def _get_model_and_data(