Use get_heterogeneous_feature_mapping in LFI argparse dispatcher

Summary: Refactors the learned imputation argument dispatcher to delegate feature index computation to the dataset's built-in mapping utility. This eliminates duplicated feature-ordering logic and ensures consistency with the canonical ordering convention.

Differential Revision: D102197138

Author: Carl Hvarfner

ax/generators/torch/botorch_modular/input_constructors/input_transforms.py

@@ -357,40 +357,12 @@ def _input_transform_argparse_learned_feature_imputation(
         )
     input_transform_options = input_transform_options or {}
 
-    # Order datasets: target first, then remaining (same as HeterogeneousMTGP).
-    child_datasets = dataset.datasets.copy()
-    target_dataset = child_datasets.pop(dataset.target_outcome_name)
-    all_datasets = [target_dataset] + list(child_datasets.values())
-
-    # The feature_names[:task_feature_index] slice only works when the task
-    # column is the last column (index == -1). Guard against other positions
-    # the same way ImputedMultiTaskGP.construct_inputs does.
+    # Delegate feature ordering and index computation to MultiTaskDataset.
+    all_datasets, feature_indices_list, d = dataset.get_heterogeneous_feature_mapping()
+    feature_indices = dict(enumerate(feature_indices_list))
     task_feature_index = (
-        dataset.task_feature_index if (dataset.task_feature_index is not None) else -1
+        dataset.task_feature_index if dataset.task_feature_index is not None else -1
     )
-    if task_feature_index != -1:
-        raise NotImplementedError(
-            "LearnedFeatureImputation argparse only supports "
-            "task_feature_index == -1. Got "
-            f"task_feature_index={task_feature_index}."
-        )
-
-    # Use target's feature order as canonical (NO alphabetical sort).
-    # Source-only features are appended at the end.
-    all_features: list[str] = list(target_dataset.feature_names[:task_feature_index])
-    for ds in all_datasets[1:]:
-        for fn in ds.feature_names[:task_feature_index]:
-            if fn not in all_features:
-                all_features.append(fn)
-    d = len(all_features)
-
-    # Map each task's features to indices in the global feature space.
-    feature_indices = {
-        task_idx: [
-            all_features.index(fn) for fn in ds.feature_names[:task_feature_index]
-        ]
-        for task_idx, ds in enumerate(all_datasets)
-    }
 
     dtype = torch_dtype or torch.float64
     # Constrain imputation values to [0, 1] since the preceding Normalize

ax/generators/torch/tests/test_input_transform_argparse.py

@@ -475,9 +475,7 @@ def test_argparse_learned_feature_imputation(self) -> None:
                 target_outcome_name="y0",
                 task_feature_index=0,
             )
-            with self.assertRaisesRegex(
-                NotImplementedError, "task_feature_index == -1"
-            ):
+            with self.assertRaisesRegex(NotImplementedError, "task_feature_index.*-1"):
                 input_transform_argparse(
                     LearnedFeatureImputation,
                     dataset=bad_ds,