Update contextual models for use in MBM (pytorch#2206)

Summary:
X-link: facebook/Ax#2206

Pull Request resolved: pytorch#2206

Updates for contextual models to be used in MBM

Reviewed By: saitcakmak

Differential Revision: D52452235

fbshipit-source-id: 6fa479e1498789d9330cc7b07ac0500612bfd505

Author: bletham

botorch/models/contextual.py

@@ -121,7 +121,7 @@ def __init__(
     def construct_inputs(
         cls,
         training_data: SupervisedDataset,
-        decomposition: Dict[str, List[int]],
+        decomposition: Dict[str, List[str]],
         train_embedding: bool = True,
         cat_feature_dict: Optional[Dict] = None,
         embs_feature_dict: Optional[Dict] = None,
@@ -133,7 +133,7 @@ def construct_inputs(
 
         Args:
             training_data: A `SupervisedDataset` containing the training data.
-            decomposition: Dictionary of context names and their indexes of the
+            decomposition: Dictionary of context names and the names of the
                 corresponding active context parameters.
             train_embedding: Whether to train the embedding layer or not.
             cat_feature_dict: Keys are context names and values are list of categorical
@@ -148,9 +148,13 @@ def construct_inputs(
             context_weight_dict: Known population weights of each context.
         """
         base_inputs = super().construct_inputs(training_data=training_data, **kwargs)
+        index_decomp = {
+            c: [training_data.feature_names.index(i) for i in v]
+            for c, v in decomposition.items()
+        }
         return {
             **base_inputs,
-            "decomposition": decomposition,
+            "decomposition": index_decomp,
             "train_embedding": train_embedding,
             "cat_feature_dict": cat_feature_dict,
             "embs_feature_dict": embs_feature_dict,

test/models/test_contextual.py

@@ -135,10 +135,10 @@ def test_LCEAGP_construct_inputs(self):
         for dtype in (torch.float, torch.double):
             tkwargs = {"device": self.device, "dtype": dtype}
             datasets, (train_X, train_Y, train_Yvar) = _gen_datasets(**tkwargs)
-            decomposition = {"1": [0, 1], "2": [2, 3]}
+            decomposition = {"1": ["x0", "x1"], "2": ["x2", "x3"]}
+            decomposition_index = {"1": [0, 1], "2": [2, 3]}
 
-            model = LCEAGP(train_X, train_Y, train_Yvar, decomposition)
-            data_dict = model.construct_inputs(
+            data_dict = LCEAGP.construct_inputs(
                 training_data=datasets,
                 decomposition=decomposition,
                 train_embedding=False,
@@ -147,5 +147,5 @@ def test_LCEAGP_construct_inputs(self):
             self.assertTrue(train_X.equal(data_dict["train_X"]))
             self.assertTrue(train_Y.equal(data_dict["train_Y"]))
             self.assertTrue(train_Yvar.equal(data_dict["train_Yvar"]))
-            self.assertDictEqual(data_dict["decomposition"], decomposition)
+            self.assertDictEqual(data_dict["decomposition"], decomposition_index)
             self.assertFalse(data_dict["train_embedding"])