From 3621198d02195b723195b043e86738cd5c3b8e40 Mon Sep 17 00:00:00 2001 From: Elena Kashtelyan Date: Wed, 2 Sep 2020 10:53:06 -0700 Subject: [PATCH] =?UTF-8?q?`Model.construct=5Finputs`=20expects=20a=20sing?= =?UTF-8?q?le=20`TrainingData`=20now=20=E2=80=93=E2=80=93=20BoTorch=20chan?= =?UTF-8?q?ges=20(#529)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: Pull Request resolved: https://github.com/pytorch/botorch/pull/529 Change the structure of `TrainingData` and `construct_inputs` in BoTorch. **Previously:** ``` class TrainingData: Xs: List[Tensor] Ys: List[Tensor] Yvars: Optional[List[Tensor]] = None training_data: TrainingData ``` **Now:** ``` class TrainingData: X: Tensor Y: Tensor Yvar: Optional[Tensor] = None training_data: Dict[str, TrainingData] ``` Reviewed By: Balandat Differential Revision: D23380460 fbshipit-source-id: 1e1dbf1e2c19f06b9b0edb1b2042dc0226e52926 --- .../acquisition/multi_objective/objective.py | 3 +- botorch/models/gp_regression.py | 47 ++++---- botorch/models/gp_regression_fidelity.py | 70 +++++------- botorch/models/model.py | 2 +- botorch/models/multitask.py | 51 ++++++++- botorch/utils/containers.py | 10 +- test/models/test_gp_regression.py | 106 ++---------------- test/models/test_gp_regression_fidelity.py | 68 ++--------- test/models/test_multitask.py | 46 +++++++- test/utils/test_containers.py | 22 ++-- 10 files changed, 184 insertions(+), 241 deletions(-) diff --git a/botorch/acquisition/multi_objective/objective.py b/botorch/acquisition/multi_objective/objective.py index 262d6f9011..31a50caa48 100644 --- a/botorch/acquisition/multi_objective/objective.py +++ b/botorch/acquisition/multi_objective/objective.py @@ -31,7 +31,8 @@ def forward(self, samples: Tensor, **kwargs) -> Tensor: Returns: A `sample_shape x batch_shape x q x m'`-dim Tensor of objective values with - `m'` the output dimension. This assumes maximization in each output dimension). + `m'` the output dimension. This assumes maximization in each output + dimension). This method is usually not called directly, but via the objectives diff --git a/botorch/models/gp_regression.py b/botorch/models/gp_regression.py index 8a44709692..5757b31172 100644 --- a/botorch/models/gp_regression.py +++ b/botorch/models/gp_regression.py @@ -141,15 +141,14 @@ def forward(self, x: Tensor) -> MultivariateNormal: def construct_inputs( cls, training_data: TrainingData, **kwargs: Any ) -> Dict[str, Any]: - r"""Standardize kwargs of the model constructor.""" - Xs = training_data.Xs - Ys = training_data.Ys - if len(Xs) == len(Ys) == 1: - return {"train_X": Xs[0], "train_Y": Ys[0]} - if all(torch.equal(Xs[0], X) for X in Xs[1:]): - # Use batched multioutput, single task GP. - return {"train_X": Xs[0], "train_Y": torch.cat(Ys, dim=-1)} - raise ValueError("Unexpected training data format.") + r"""Construct kwargs for the `Model` from `TrainingData`. + + Args: + training_data: `TrainingData` container with data for single outcome + or for multiple outcomes for batched multi-output case. + **kwargs: None expected for this class. + """ + return {"train_X": training_data.X, "train_Y": training_data.Y} class FixedNoiseGP(BatchedMultiOutputGPyTorchModel, ExactGP): @@ -296,22 +295,20 @@ def subset_output(self, idcs: List[int]) -> BatchedMultiOutputGPyTorchModel: def construct_inputs( cls, training_data: TrainingData, **kwargs: Any ) -> Dict[str, Any]: - r"""Standardize kwargs of the model constructor.""" - if training_data.Yvars is None: - raise ValueError(f"Yvars required for {cls.__name__}.") - Xs = training_data.Xs - Ys = training_data.Ys - Yvars = training_data.Yvars - if len(Xs) == len(Ys) == 1: - return {"train_X": Xs[0], "train_Y": Ys[0], "train_Yvar": Yvars[0]} - if all(torch.equal(Xs[0], X) for X in Xs[1:]): - # Use batched multioutput, single task GP. - return { - "train_X": Xs[0], - "train_Y": torch.cat(Ys, dim=-1), - "train_Yvar": torch.cat(Yvars, dim=-1), - } - raise ValueError("Unexpected training data format.") + r"""Construct kwargs for the `Model` from `TrainingData`. + + Args: + training_data: `TrainingData` container with data for single outcome + or for multiple outcomes for batched multi-output case. + **kwargs: None expected for this class. + """ + if training_data.Yvar is None: + raise ValueError(f"Yvar required for {cls.__name__}.") + return { + "train_X": training_data.X, + "train_Y": training_data.Y, + "train_Yvar": training_data.Yvar, + } class HeteroskedasticSingleTaskGP(SingleTaskGP): diff --git a/botorch/models/gp_regression_fidelity.py b/botorch/models/gp_regression_fidelity.py index 29dbc680f9..7b81ed196e 100644 --- a/botorch/models/gp_regression_fidelity.py +++ b/botorch/models/gp_regression_fidelity.py @@ -117,26 +117,23 @@ def __init__( @classmethod def construct_inputs(cls, training_data: TrainingData, **kwargs) -> Dict[str, Any]: - r"""Standardize kwargs of the model constructor.""" + r"""Construct kwargs for the `Model` from `TrainingData`. + + Args: + training_data: `TrainingData` container with data for single outcome + or for multiple outcomes for batched multi-output case. + **kwargs: Options, expected for this class: + - fidelity_features: List of columns of X that are fidelity parameters. + """ fidelity_features = kwargs.get("fidelity_features") if fidelity_features is None: raise ValueError(f"Fidelity features required for {cls.__name__}.") - Xs = training_data.Xs - Ys = training_data.Ys - if len(Xs) == len(Ys) == 1: - return { - "train_X": Xs[0], - "train_Y": Ys[0], - "data_fidelity": fidelity_features[0], - } - if all(torch.equal(Xs[0], X) for X in Xs[1:]): - # Use batched multioutput, single task GP. - return { - "train_X": Xs[0], - "train_Y": torch.cat(Ys, dim=-1), - "data_fidelity": fidelity_features[0], - } - raise ValueError("Unexpected training data format.") + + return { + "train_X": training_data.X, + "train_Y": training_data.Y, + "data_fidelity": fidelity_features[0], + } class FixedNoiseMultiFidelityGP(FixedNoiseGP): @@ -220,31 +217,26 @@ def __init__( @classmethod def construct_inputs(cls, training_data: TrainingData, **kwargs) -> Dict[str, Any]: - r"""Standardize kwargs of the model constructor.""" - if training_data.Yvars is None: - raise ValueError(f"Yvars required for {cls.__name__}.") - Xs = training_data.Xs - Ys = training_data.Ys - Yvars = training_data.Yvars + r"""Construct kwargs for the `Model` from `TrainingData`. + + Args: + training_data: `TrainingData` container with data for single outcome + or for multiple outcomes for batched multi-output case. + **kwargs: Options, expected for this class: + - fidelity_features: List of columns of X that are fidelity parameters. + """ fidelity_features = kwargs.get("fidelity_features") if fidelity_features is None: raise ValueError(f"Fidelity features required for {cls.__name__}.") - if len(Xs) == len(Ys) == 1: - return { - "train_X": Xs[0], - "train_Y": Ys[0], - "train_Yvar": Yvars[0], - "data_fidelity": fidelity_features[0], - } - if all(torch.equal(Xs[0], X) for X in Xs[1:]): - # Use batched multioutput, single task GP. - return { - "train_X": Xs[0], - "train_Y": torch.cat(Ys, dim=-1), - "train_Yvar": torch.cat(Yvars, dim=-1), - "data_fidelity": fidelity_features[0], - } - raise ValueError("Unexpected training data format.") + if training_data.Yvar is None: + raise ValueError(f"Yvar required for {cls.__name__}.") + + return { + "train_X": training_data.X, + "train_Y": training_data.Y, + "train_Yvar": training_data.Yvar, + "data_fidelity": fidelity_features[0], + } def _setup_multifidelity_covar_module( diff --git a/botorch/models/model.py b/botorch/models/model.py index e7779bbacf..cd22fca1b4 100644 --- a/botorch/models/model.py +++ b/botorch/models/model.py @@ -129,7 +129,7 @@ def fantasize( def construct_inputs( cls, training_data: TrainingData, **kwargs: Any ) -> Dict[str, Any]: - r"""Standardize kwargs of the model constructor.""" + r"""Construct kwargs for the `Model` from `TrainingData`.""" raise NotImplementedError( f"`construct_inputs` not implemented for {cls.__name__}." ) diff --git a/botorch/models/multitask.py b/botorch/models/multitask.py index c73a92a1a6..edb6241167 100644 --- a/botorch/models/multitask.py +++ b/botorch/models/multitask.py @@ -10,10 +10,11 @@ from __future__ import annotations -from typing import List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple import torch from botorch.models.gpytorch import MultiTaskGPyTorchModel +from botorch.utils.containers import TrainingData from gpytorch.distributions.multivariate_normal import MultivariateNormal from gpytorch.kernels.index_kernel import IndexKernel from gpytorch.kernels.matern_kernel import MaternKernel @@ -166,6 +167,29 @@ def get_all_tasks( all_tasks = train_X[:, task_feature].unique().to(dtype=torch.long).tolist() return all_tasks, task_feature, d + @classmethod + def construct_inputs(cls, training_data: TrainingData, **kwargs) -> Dict[str, Any]: + r"""Construct kwargs for the `Model` from `TrainingData`. + + Args: + training_data: `TrainingData` container with data for single outcome + or for multiple outcomes for batched multi-output case. + **kwargs: Additional options for the model that pertain to the + training data: + - `task_features` – indices of the input columns containing the task + features (expected list of length 1). + """ + + task_features = kwargs.get("task_features") + if task_features is None: + raise ValueError(f"task features required for {cls.__name__}.") + + return { + "train_X": training_data.X, + "train_Y": training_data.Y, + "task_feature": task_features[0], + } + class FixedNoiseMultiTaskGP(MultiTaskGP): r"""Multi-Task GP model using an ICM kernel, with known observation noise. @@ -226,3 +250,28 @@ def __init__( ) self.likelihood = FixedNoiseGaussianLikelihood(noise=train_Yvar.squeeze(-1)) self.to(train_X) + + @classmethod + def construct_inputs(cls, training_data: TrainingData, **kwargs) -> Dict[str, Any]: + r"""Construct kwargs for the `Model` from `TrainingData`. + + Args: + training_data: `TrainingData` container with data for single outcome + or for multiple outcomes for batched multi-output case. + **kwargs: Additional options for the model that pertain to the + training data: + - `task_features` – indices of task features in X. + """ + + task_features = kwargs.get("task_features") + if task_features is None: + raise ValueError(f"task features required for {cls.__name__}.") + if training_data.Yvar is None: + raise ValueError(f"Yvar required for {cls.__name__}.") + + return { + "train_X": training_data.X, + "train_Y": training_data.Y, + "train_Yvar": training_data.Yvar, + "task_feature": task_features[0], + } diff --git a/botorch/utils/containers.py b/botorch/utils/containers.py index 28c7ce20ed..c464ee3f99 100644 --- a/botorch/utils/containers.py +++ b/botorch/utils/containers.py @@ -8,14 +8,14 @@ Containers to standardize inputs into models and acquisition functions. """ -from typing import List, NamedTuple, Optional +from typing import NamedTuple, Optional from torch import Tensor class TrainingData(NamedTuple): - r"""Standardized struct of model training data.""" + r"""Standardized struct of model training data for a single outcome.""" - Xs: List[Tensor] - Ys: List[Tensor] - Yvars: Optional[List[Tensor]] = None + X: Tensor + Y: Tensor + Yvar: Optional[Tensor] = None diff --git a/test/models/test_gp_regression.py b/test/models/test_gp_regression.py index c724580842..d76ceae0e3 100644 --- a/test/models/test_gp_regression.py +++ b/test/models/test_gp_regression.py @@ -280,47 +280,12 @@ def test_construct_inputs(self): model, model_kwargs = self._get_model_and_data( batch_shape=batch_shape, m=2, **tkwargs ) - # len(Xs) == len(Ys) == 1 training_data = TrainingData( - Xs=[model_kwargs["train_X"][0]], Ys=[model_kwargs["train_Y"][0]] - ) - data_dict = model.construct_inputs(training_data) - self.assertTrue( - torch.equal(data_dict["train_X"], model_kwargs["train_X"][0]) - ) - self.assertTrue( - torch.equal(data_dict["train_Y"], model_kwargs["train_Y"][0]) - ) - # all X's are equal - training_data = TrainingData( - Xs=[model_kwargs["train_X"], model_kwargs["train_X"]], - Ys=[model_kwargs["train_Y"], model_kwargs["train_Y"]], + X=model_kwargs["train_X"], Y=model_kwargs["train_Y"] ) data_dict = model.construct_inputs(training_data) self.assertTrue(torch.equal(data_dict["train_X"], model_kwargs["train_X"])) - self.assertTrue( - torch.equal( - data_dict["train_Y"], - torch.cat( - [model_kwargs["train_Y"], model_kwargs["train_Y"]], dim=-1 - ), - ) - ) - # unexpected data format - training_data = TrainingData( - Xs=[model_kwargs["train_X"], torch.add(model_kwargs["train_X"], 1)], - Ys=[model_kwargs["train_Y"], model_kwargs["train_Y"]], - ) - with self.assertRaises(ValueError): - model.construct_inputs(training_data) - # make sure Yvar is not added to dict - training_data = TrainingData( - Xs=[model_kwargs["train_X"]], - Ys=[model_kwargs["train_Y"]], - Yvars=[torch.full_like(model_kwargs["train_Y"], 0.01)], - ) - data_dict = model.construct_inputs(training_data) - self.assertTrue("train_Yvar" not in data_dict) + self.assertTrue(torch.equal(data_dict["train_Y"], model_kwargs["train_Y"])) class TestFixedNoiseGP(TestSingleTaskGP): @@ -363,73 +328,20 @@ def test_construct_inputs(self): batch_shape=batch_shape, m=2, **tkwargs ) training_data = TrainingData( - Xs=[model_kwargs["train_X"][0]], - Ys=[model_kwargs["train_Y"][0]], - Yvars=[model_kwargs["train_Yvar"][0]], + X=model_kwargs["train_X"], + Y=model_kwargs["train_Y"], + Yvar=model_kwargs["train_Yvar"], ) data_dict = model.construct_inputs(training_data) self.assertTrue("train_Yvar" in data_dict) - self.assertTrue( - torch.equal(data_dict["train_X"], model_kwargs["train_X"][0]) - ) - self.assertTrue( - torch.equal(data_dict["train_Y"], model_kwargs["train_Y"][0]) - ) - self.assertTrue( - torch.equal(data_dict["train_Yvar"], model_kwargs["train_Yvar"][0]) - ) - # if Yvars is missing, then raise error - training_data = TrainingData( - Xs=model_kwargs["train_X"], Ys=model_kwargs["train_Y"] - ) - with self.assertRaises(ValueError): - model.construct_inputs(training_data) - - # len(Xs) == len(Ys) == 1 - training_data = TrainingData( - Xs=[model_kwargs["train_X"][0]], - Ys=[model_kwargs["train_Y"][0]], - Yvars=[model_kwargs["train_Yvar"][0]], - ) - data_dict = model.construct_inputs(training_data) - self.assertTrue( - torch.equal(data_dict["train_X"], model_kwargs["train_X"][0]) - ) - self.assertTrue( - torch.equal(data_dict["train_Y"], model_kwargs["train_Y"][0]) - ) - self.assertTrue( - torch.equal(data_dict["train_Yvar"], model_kwargs["train_Yvar"][0]) - ) - # all X's are equal - training_data = TrainingData( - Xs=[model_kwargs["train_X"], model_kwargs["train_X"]], - Ys=[model_kwargs["train_Y"], model_kwargs["train_Y"]], - Yvars=[model_kwargs["train_Yvar"], model_kwargs["train_Yvar"]], - ) - data_dict = model.construct_inputs(training_data) self.assertTrue(torch.equal(data_dict["train_X"], model_kwargs["train_X"])) + self.assertTrue(torch.equal(data_dict["train_Y"], model_kwargs["train_Y"])) self.assertTrue( - torch.equal( - data_dict["train_Y"], - torch.cat( - [model_kwargs["train_Y"], model_kwargs["train_Y"]], dim=-1 - ), - ) - ) - self.assertTrue( - torch.equal( - data_dict["train_Yvar"], - torch.cat( - [model_kwargs["train_Yvar"], model_kwargs["train_Yvar"]], dim=-1 - ), - ) + torch.equal(data_dict["train_Yvar"], model_kwargs["train_Yvar"]) ) - # unexpected data format + # if Yvars is missing, then raise error training_data = TrainingData( - Xs=[model_kwargs["train_X"], torch.add(model_kwargs["train_X"], 1)], - Ys=[model_kwargs["train_Y"], model_kwargs["train_Y"]], - Yvars=[model_kwargs["train_Yvar"]], + X=model_kwargs["train_X"], Y=model_kwargs["train_Y"] ) with self.assertRaises(ValueError): model.construct_inputs(training_data) diff --git a/test/models/test_gp_regression_fidelity.py b/test/models/test_gp_regression_fidelity.py index b31a02c7d1..d4bbe917a8 100644 --- a/test/models/test_gp_regression_fidelity.py +++ b/test/models/test_gp_regression_fidelity.py @@ -348,9 +348,9 @@ def test_construct_inputs(self): ) # len(Xs) == len(Ys) == 1 training_data = TrainingData( - Xs=[model_kwargs["train_X"][0]], - Ys=[model_kwargs["train_Y"][0]], - Yvars=[torch.full_like(model_kwargs["train_Y"], 0.01)[0]], + X=model_kwargs["train_X"], + Y=model_kwargs["train_Y"], + Yvar=torch.full_like(model_kwargs["train_Y"], 0.01), ) # missing fidelity features with self.assertRaises(ValueError): @@ -360,36 +360,12 @@ def test_construct_inputs(self): self.assertTrue("data_fidelity" in data_dict) self.assertEqual(data_dict["data_fidelity"], 1) data_dict = model.construct_inputs(training_data, fidelity_features=[1]) - self.assertTrue( - torch.equal(data_dict["train_X"], model_kwargs["train_X"][0]) - ) - self.assertTrue( - torch.equal(data_dict["train_Y"], model_kwargs["train_Y"][0]) - ) - # all X's are equal - training_data = TrainingData( - Xs=[model_kwargs["train_X"], model_kwargs["train_X"]], - Ys=[model_kwargs["train_Y"], model_kwargs["train_Y"]], - ) - data_dict = model.construct_inputs(training_data, fidelity_features=[1]) self.assertTrue( torch.equal(data_dict["train_X"], model_kwargs["train_X"]) ) self.assertTrue( - torch.equal( - data_dict["train_Y"], - torch.cat( - [model_kwargs["train_Y"], model_kwargs["train_Y"]], dim=-1 - ), - ) - ) - # unexpected data format - training_data = TrainingData( - Xs=[model_kwargs["train_X"], torch.add(model_kwargs["train_X"], 1)], - Ys=[model_kwargs["train_Y"], model_kwargs["train_Y"]], + torch.equal(data_dict["train_Y"], model_kwargs["train_Y"]) ) - with self.assertRaises(ValueError): - model.construct_inputs(training_data, fidelity_features=[1]) class TestFixedNoiseMultiFidelityGP(TestSingleTaskMultiFidelityGP): @@ -474,16 +450,16 @@ def test_construct_inputs(self): **tkwargs, ) training_data = TrainingData( - Xs=[model_kwargs["train_X"][0]], Ys=[model_kwargs["train_Y"][0]] + X=model_kwargs["train_X"], Y=model_kwargs["train_Y"] ) # missing Yvars with self.assertRaises(ValueError): model.construct_inputs(training_data, fidelity_features=[1]) # len(Xs) == len(Ys) == 1 training_data = TrainingData( - Xs=[model_kwargs["train_X"][0]], - Ys=[model_kwargs["train_Y"][0]], - Yvars=[torch.full_like(model_kwargs["train_Y"], 0.01)[0]], + X=model_kwargs["train_X"], + Y=model_kwargs["train_Y"], + Yvar=torch.full_like(model_kwargs["train_Y"], 0.01), ) # missing fidelity features with self.assertRaises(ValueError): @@ -493,35 +469,9 @@ def test_construct_inputs(self): self.assertTrue("data_fidelity" in data_dict) self.assertEqual(data_dict["data_fidelity"], 1) data_dict = model.construct_inputs(training_data, fidelity_features=[1]) - self.assertTrue( - torch.equal(data_dict["train_X"], model_kwargs["train_X"][0]) - ) - self.assertTrue( - torch.equal(data_dict["train_Y"], model_kwargs["train_Y"][0]) - ) - # all X's are equal - training_data = TrainingData( - Xs=[model_kwargs["train_X"], model_kwargs["train_X"]], - Ys=[model_kwargs["train_Y"], model_kwargs["train_Y"]], - Yvars=[model_kwargs["train_Yvar"], model_kwargs["train_Yvar"]], - ) - data_dict = model.construct_inputs(training_data, fidelity_features=[1]) self.assertTrue( torch.equal(data_dict["train_X"], model_kwargs["train_X"]) ) self.assertTrue( - torch.equal( - data_dict["train_Y"], - torch.cat( - [model_kwargs["train_Y"], model_kwargs["train_Y"]], dim=-1 - ), - ) - ) - # unexpected data format - training_data = TrainingData( - Xs=[model_kwargs["train_X"], torch.add(model_kwargs["train_X"], 1)], - Ys=[model_kwargs["train_Y"], model_kwargs["train_Y"]], - Yvars=[model_kwargs["train_Yvar"], model_kwargs["train_Yvar"]], + torch.equal(data_dict["train_Y"], model_kwargs["train_Y"]) ) - with self.assertRaises(ValueError): - model.construct_inputs(training_data, fidelity_features=[1]) diff --git a/test/models/test_multitask.py b/test/models/test_multitask.py index e72e37f572..3422f45e88 100644 --- a/test/models/test_multitask.py +++ b/test/models/test_multitask.py @@ -12,6 +12,7 @@ from botorch.fit import fit_gpytorch_model from botorch.models.multitask import FixedNoiseMultiTaskGP, MultiTaskGP from botorch.posteriors import GPyTorchPosterior +from botorch.utils.containers import TrainingData from botorch.utils.testing import BotorchTestCase from gpytorch.distributions import MultitaskMultivariateNormal, MultivariateNormal from gpytorch.kernels import IndexKernel, MaternKernel, ScaleKernel @@ -37,9 +38,13 @@ def _get_random_mt_data(**tkwargs): def _get_model(**tkwargs): + return _get_model_and_training_data(**tkwargs)[0] + + +def _get_model_and_training_data(**tkwargs): train_X, train_Y = _get_random_mt_data(**tkwargs) model = MultiTaskGP(train_X, train_Y, task_feature=1) - return model.to(**tkwargs) + return model.to(**tkwargs), train_X, train_Y def _get_model_single_output(**tkwargs): @@ -49,10 +54,14 @@ def _get_model_single_output(**tkwargs): def _get_fixed_noise_model(**tkwargs): + return _get_fixed_noise_model_and_training_data(**tkwargs)[0] + + +def _get_fixed_noise_model_and_training_data(**tkwargs): train_X, train_Y = _get_random_mt_data(**tkwargs) train_Yvar = torch.full_like(train_Y, 0.05) model = FixedNoiseMultiTaskGP(train_X, train_Y, train_Yvar, task_feature=1) - return model.to(**tkwargs) + return model.to(**tkwargs), train_X, train_Y, train_Yvar def _get_fixed_noise_model_single_output(**tkwargs): @@ -327,3 +336,36 @@ def test_FixedNoiseMultiTaskGP_fixed_prior(self): self.assertIsInstance( model.task_covar_module.IndexKernelPrior, LKJCovariancePrior ) + + def test_MultiTaskGP_construct_inputs(self): + for dtype in (torch.float, torch.double): + tkwargs = {"device": self.device, "dtype": dtype} + model, train_X, train_Y = _get_model_and_training_data(**tkwargs) + training_data = TrainingData(X=train_X, Y=train_Y) + # if task_features is missing, then raise error + with self.assertRaises(ValueError): + model.construct_inputs(training_data) + data_dict = model.construct_inputs(training_data, task_features=[0]) + self.assertTrue(torch.equal(data_dict["train_X"], train_X)) + self.assertTrue(torch.equal(data_dict["train_Y"], train_Y)) + + def test_FixedNoiseMultiTaskGP_construct_inputs(self): + for dtype in (torch.float, torch.double): + tkwargs = {"device": self.device, "dtype": dtype} + ( + model, + train_X, + train_Y, + train_Yvar, + ) = _get_fixed_noise_model_and_training_data(**tkwargs) + td_no_Yvar = TrainingData(X=train_X, Y=train_Y) + # if task_features is missing, then raise error + with self.assertRaisesRegex(ValueError, "task features required"): + model.construct_inputs(td_no_Yvar) + # if task_features is missing, then raise error + with self.assertRaisesRegex(ValueError, "Yvar required"): + model.construct_inputs(td_no_Yvar, task_features=[0]) + training_data = TrainingData(X=train_X, Y=train_Y, Yvar=train_Yvar) + data_dict = model.construct_inputs(training_data, task_features=[0]) + self.assertTrue(torch.equal(data_dict["train_X"], train_X)) + self.assertTrue(torch.equal(data_dict["train_Y"], train_Y)) diff --git a/test/utils/test_containers.py b/test/utils/test_containers.py index ba527a2468..e92d34005c 100644 --- a/test/utils/test_containers.py +++ b/test/utils/test_containers.py @@ -11,16 +11,16 @@ class TestConstructContainers(BotorchTestCase): def test_TrainingData(self): - Xs = torch.tensor([[-1.0, 0.0, 0.0], [0.0, 1.0, 1.0]]) - Ys = torch.tensor([[-1.0, 0.0, 0.0], [0.0, 1.0, 1.0]]) - Yvars = torch.tensor([[-1.0, 0.0, 0.0], [0.0, 1.0, 1.0]]) + X = torch.tensor([[-1.0, 0.0, 0.0], [0.0, 1.0, 1.0]]) + Y = torch.tensor([[-1.0, 0.0, 0.0], [0.0, 1.0, 1.0]]) + Yvar = torch.tensor([[-1.0, 0.0, 0.0], [0.0, 1.0, 1.0]]) - training_data = TrainingData(Xs, Ys) - self.assertTrue(torch.equal(training_data.Xs, Xs)) - self.assertTrue(torch.equal(training_data.Ys, Ys)) - self.assertEqual(training_data.Yvars, None) + training_data = TrainingData(X, Y) + self.assertTrue(torch.equal(training_data.X, X)) + self.assertTrue(torch.equal(training_data.Y, Y)) + self.assertEqual(training_data.Yvar, None) - training_data = TrainingData(Xs, Ys, Yvars) - self.assertTrue(torch.equal(training_data.Xs, Xs)) - self.assertTrue(torch.equal(training_data.Ys, Ys)) - self.assertTrue(torch.equal(training_data.Yvars, Yvars)) + training_data = TrainingData(X, Y, Yvar) + self.assertTrue(torch.equal(training_data.X, X)) + self.assertTrue(torch.equal(training_data.Y, Y)) + self.assertTrue(torch.equal(training_data.Yvar, Yvar))