Support multiple data fidelity dimensions in MultiFidelityGP models

botorch/models/gp_regression_fidelity.py

@@ -66,7 +66,7 @@ def __init__(
         train_X: Tensor,
         train_Y: Tensor,
         iteration_fidelity: Optional[int] = None,
-        data_fidelity: Optional[int] = None,
+        data_fidelity: Optional[int | List[int]] = None,
         linear_truncated: bool = True,
         nu: float = 2.5,
         likelihood: Optional[Likelihood] = None,
@@ -81,8 +81,9 @@ def __init__(
             train_Y: A `batch_shape x n x m` tensor of training observations.
             iteration_fidelity: The column index for the training iteration fidelity
                 parameter (optional).
-            data_fidelity: The column index for the downsampling fidelity parameter
-                (optional).
+            data_fidelity: The column index for the downsampling fidelity parameter.
+                If multiple indices are provided, a kernel will be constructed for
+                each index (optional).
             linear_truncated: If True, use a `LinearTruncatedFidelityKernel` instead
                 of the default kernel.
             nu: The smoothness parameter for the Matern kernel: either 1/2, 3/2, or
@@ -117,7 +118,7 @@ def __init__(
             dim=transformed_X.size(-1),
             aug_batch_shape=self._aug_batch_shape,
             iteration_fidelity=iteration_fidelity,
-            data_fidelity=data_fidelity,
+            data_fidelities=data_fidelity,
             linear_truncated=linear_truncated,
             nu=nu,
         )
@@ -150,11 +151,8 @@ def construct_inputs(
             training_data: Dictionary of `SupervisedDataset`.
             fidelity_features: Index of fidelity parameter as input columns.
         """
-        if len(fidelity_features) != 1:
-            raise UnsupportedError("Multiple fidelity features not supported.")
-
         inputs = super().construct_inputs(training_data=training_data, **kwargs)
-        inputs["data_fidelity"] = fidelity_features[0]
+        inputs["data_fidelity"] = fidelity_features
         return inputs
 
 
@@ -185,7 +183,7 @@ def __init__(
         train_Y: Tensor,
         train_Yvar: Tensor,
         iteration_fidelity: Optional[int] = None,
-        data_fidelity: Optional[int] = None,
+        data_fidelity: Optional[int | List[int]] = None,
         linear_truncated: bool = True,
         nu: float = 2.5,
         outcome_transform: Optional[OutcomeTransform] = None,
@@ -200,8 +198,9 @@ def __init__(
             train_Yvar: A `batch_shape x n x m` tensor of observed measurement noise.
             iteration_fidelity: The column index for the training iteration fidelity
                 parameter (optional).
-            data_fidelity: The column index for the downsampling fidelity parameter
-                (optional).
+            data_fidelity: The column index for the downsampling fidelity parameter.
+                If multiple indices are provided, a kernel will be constructed for
+                each index (optional).
             linear_truncated: If True, use a `LinearTruncatedFidelityKernel` instead
                 of the default kernel.
             nu: The smoothness parameter for the Matern kernel: either 1/2, 3/2, or
@@ -213,6 +212,13 @@ def __init__(
             input_transform: An input transform that is applied in the model's
                 forward pass.
         """
+        self._init_args = {
+            "iteration_fidelity": iteration_fidelity,
+            "data_fidelity": data_fidelity,
+            "linear_truncated": linear_truncated,
+            "nu": nu,
+            "outcome_transform": outcome_transform,
+        }
         if iteration_fidelity is None and data_fidelity is None:
             raise UnsupportedError(
                 "FixedNoiseMultiFidelityGP requires at least one fidelity parameter."
@@ -226,7 +232,7 @@ def __init__(
             dim=transformed_X.size(-1),
             aug_batch_shape=self._aug_batch_shape,
             iteration_fidelity=iteration_fidelity,
-            data_fidelity=data_fidelity,
+            data_fidelities=data_fidelity,
             linear_truncated=linear_truncated,
             nu=nu,
         )
@@ -259,19 +265,16 @@ def construct_inputs(
             training_data: Dictionary of `SupervisedDataset`.
             fidelity_features: Column indices of fidelity features.
         """
-        if len(fidelity_features) != 1:
-            raise UnsupportedError("Multiple fidelity features not supported.")
-
         inputs = super().construct_inputs(training_data=training_data, **kwargs)
-        inputs["data_fidelity"] = fidelity_features[0]
+        inputs["data_fidelity"] = fidelity_features
         return inputs
 
 
 def _setup_multifidelity_covar_module(
     dim: int,
     aug_batch_shape: torch.Size,
     iteration_fidelity: Optional[int],
-    data_fidelity: Optional[int],
+    data_fidelities: Optional[int | List[int]],
     linear_truncated: bool,
     nu: float,
 ) -> Tuple[ScaleKernel, Dict]:
@@ -284,8 +287,8 @@ def _setup_multifidelity_covar_module(
             `BatchedMultiOutputGPyTorchModel`.
         iteration_fidelity: The column index for the training iteration fidelity
             parameter (optional).
-        data_fidelity: The column index for the downsampling fidelity parameter
-            (optional).
+        data_fidelities: The column indices for the downsampling fidelity parameters
+            (optional). A single index can be provided as an int.
         linear_truncated: If True, use a `LinearTruncatedFidelityKernel` instead
             of the default kernel.
         nu: The smoothness parameter for the Matern kernel: either 1/2, 3/2, or
@@ -295,78 +298,114 @@ def _setup_multifidelity_covar_module(
         The covariance module and subset_batch_dict.
     """
 
-    if iteration_fidelity is not None and iteration_fidelity < 0:
-        iteration_fidelity = dim + iteration_fidelity
-    if data_fidelity is not None and data_fidelity < 0:
-        data_fidelity = dim + data_fidelity
+    if iteration_fidelity is not None:
+        if iteration_fidelity < 0:
+            iteration_fidelity = dim + iteration_fidelity
+    if data_fidelities is not None:
+        if isinstance(data_fidelities, int):
+            data_fidelities = [data_fidelities]
+        for i in range(len(data_fidelities)):
+            if data_fidelities[i] < 0:
+                data_fidelities[i] = dim + data_fidelities[i]
+
+    kernels = []
 
     if linear_truncated:
-        fidelity_dims = [
-            i for i in (iteration_fidelity, data_fidelity) if i is not None
-        ]
-        kernel = LinearTruncatedFidelityKernel(
-            fidelity_dims=fidelity_dims,
-            dimension=dim,
-            nu=nu,
-            batch_shape=aug_batch_shape,
-            power_prior=GammaPrior(3.0, 3.0),
-        )
+        if data_fidelities is not None:
+            for data_fidelity in data_fidelities:
+                if iteration_fidelity is not None:
+                    fidelity_dims = [iteration_fidelity, data_fidelity]
+                else:
+                    fidelity_dims = [data_fidelity]
+                kernels.append(
+                    LinearTruncatedFidelityKernel(
+                        fidelity_dims=fidelity_dims,
+                        dimension=dim,
+                        nu=nu,
+                        batch_shape=aug_batch_shape,
+                        power_prior=GammaPrior(3.0, 3.0),
+                    )
+                )
+        else:
+            kernels.append(
+                LinearTruncatedFidelityKernel(
+                    fidelity_dims=[iteration_fidelity],
+                    dimension=dim,
+                    nu=nu,
+                    batch_shape=aug_batch_shape,
+                    power_prior=GammaPrior(3.0, 3.0),
+                )
+            )
     else:
-        active_dimsX = [
-            i for i in range(dim) if i not in {iteration_fidelity, data_fidelity}
-        ]
-        kernel = RBFKernel(
-            ard_num_dims=len(active_dimsX),
-            batch_shape=aug_batch_shape,
-            lengthscale_prior=GammaPrior(3.0, 6.0),
-            active_dims=active_dimsX,
-        )
-        additional_kernels = []
-        if iteration_fidelity is not None:
-            exp_kernel = ExponentialDecayKernel(
+        if iteration_fidelity is not None and data_fidelities is not None:
+            active_dimsX = [
+                i
+                for i in range(dim)
+                if (i != iteration_fidelity and i not in data_fidelities)
+            ]
+        elif iteration_fidelity is not None:
+            active_dimsX = [i for i in range(dim) if i != iteration_fidelity]
+        elif data_fidelities is not None:
+            active_dimsX = [i for i in range(dim) if i not in data_fidelities]
+        kernels.append(
+            RBFKernel(
+                ard_num_dims=len(active_dimsX),
                 batch_shape=aug_batch_shape,
                 lengthscale_prior=GammaPrior(3.0, 6.0),
-                offset_prior=GammaPrior(3.0, 6.0),
-                power_prior=GammaPrior(3.0, 6.0),
-                active_dims=[iteration_fidelity],
+                active_dims=active_dimsX,
             )
-            additional_kernels.append(exp_kernel)
-        if data_fidelity is not None:
-            ds_kernel = DownsamplingKernel(
-                batch_shape=aug_batch_shape,
-                offset_prior=GammaPrior(3.0, 6.0),
-                power_prior=GammaPrior(3.0, 6.0),
-                active_dims=[data_fidelity],
+        )
+        if iteration_fidelity is not None:
+            kernels.append(
+                ExponentialDecayKernel(
+                    batch_shape=aug_batch_shape,
+                    lengthscale_prior=GammaPrior(3.0, 6.0),
+                    offset_prior=GammaPrior(3.0, 6.0),
+                    power_prior=GammaPrior(3.0, 6.0),
+                    active_dims=[iteration_fidelity],
+                )
             )
-            additional_kernels.append(ds_kernel)
-        kernel = ProductKernel(kernel, *additional_kernels)
+        if data_fidelities is not None:
+            for data_fidelity in data_fidelities:
+                kernels.append(
+                    DownsamplingKernel(
+                        batch_shape=aug_batch_shape,
+                        offset_prior=GammaPrior(3.0, 6.0),
+                        power_prior=GammaPrior(3.0, 6.0),
+                        active_dims=[data_fidelity],
+                    )
+                )
+
+    kernel = ProductKernel(*kernels)
 
     covar_module = ScaleKernel(
         kernel, batch_shape=aug_batch_shape, outputscale_prior=GammaPrior(2.0, 0.15)
     )
 
+    key_prefix = "covar_module.base_kernel.kernels"
     if linear_truncated:
-        subset_batch_dict = {
-            "covar_module.base_kernel.raw_power": -2,
-            "covar_module.base_kernel.covar_module_unbiased.raw_lengthscale": -3,
-            "covar_module.base_kernel.covar_module_biased.raw_lengthscale": -3,
-        }
+        subset_batch_dict = {}
+        for i in range(len(kernels)):
+            subset_batch_dict[f"{key_prefix}.{i}.raw_power"] = -2
+            subset_batch_dict[
+                f"{key_prefix}.{i}.covar_module_unbiased.raw_lengthscale"
+            ] = -3
+            subset_batch_dict[
+                f"{key_prefix}.{i}.covar_module_biased.raw_lengthscale"
+            ] = -3
     else:
         subset_batch_dict = {
-            "covar_module.base_kernel.kernels.0.raw_lengthscale": -3,
-            "covar_module.base_kernel.kernels.1.raw_power": -2,
-            "covar_module.base_kernel.kernels.1.raw_offset": -2,
+            f"{key_prefix}.0.raw_lengthscale": -3,
         }
+
         if iteration_fidelity is not None:
-            subset_batch_dict = {
-                "covar_module.base_kernel.kernels.1.raw_lengthscale": -3,
-                **subset_batch_dict,
-            }
-            if data_fidelity is not None:
-                subset_batch_dict = {
-                    "covar_module.base_kernel.kernels.2.raw_power": -2,
-                    "covar_module.base_kernel.kernels.2.raw_offset": -2,
-                    **subset_batch_dict,
-                }
+            subset_batch_dict[f"{key_prefix}.1.raw_power"] = -2
+            subset_batch_dict[f"{key_prefix}.1.raw_offset"] = -2
+            subset_batch_dict[f"{key_prefix}.1.raw_lengthscale"] = -3
+        if data_fidelities is not None:
+            start_idx = 2 if iteration_fidelity is not None else 1
+            for i in range(start_idx, len(data_fidelities) + start_idx):
+                subset_batch_dict[f"{key_prefix}.{i}.raw_power"] = -2
+                subset_batch_dict[f"{key_prefix}.{i}.raw_offset"] = -2
 
     return covar_module, subset_batch_dict