From 8d08aa3811d86428f1641fb72d64ef3cc4ead106 Mon Sep 17 00:00:00 2001 From: Sebastian Ament Date: Wed, 5 Feb 2025 16:26:20 -0800 Subject: [PATCH 1/2] Adding `fit_only_completed_map_metrics` option to `Adaptor` (#3312) Summary: This commit introduces the option to fit to data from uncompleted map metrics with the base `Adaptor`. Differential Revision: D69213286 --- ax/modelbridge/base.py | 34 +++++++++++--- ax/modelbridge/tests/test_base_modelbridge.py | 44 +++++++++++++++++++ ax/modelbridge/tests/test_registry.py | 1 + ax/modelbridge/torch.py | 2 + 4 files changed, 74 insertions(+), 7 deletions(-) diff --git a/ax/modelbridge/base.py b/ax/modelbridge/base.py index 7c6179b73ad..e2274621d64 100644 --- a/ax/modelbridge/base.py +++ b/ax/modelbridge/base.py @@ -23,11 +23,13 @@ from ax.core.data import Data from ax.core.experiment import Experiment from ax.core.generator_run import extract_arm_predictions, GeneratorRun +from ax.core.map_data import MapData from ax.core.observation import ( Observation, ObservationData, ObservationFeatures, observations_from_data, + observations_from_map_data, recombine_observations, separate_observations, ) @@ -112,6 +114,7 @@ def __init__( fit_abandoned: bool = False, fit_tracking_metrics: bool = True, fit_on_init: bool = True, + fit_only_completed_map_metrics: bool = True, ) -> None: """ Applies transforms and fits model. @@ -156,6 +159,9 @@ def __init__( To fit the model afterwards, use `_process_and_transform_data` to get the transformed inputs and call `_fit_if_implemented` with the transformed inputs. + fit_only_completed_map_metrics: Whether to fit a model to map metrics only + when the trial is completed. This is useful for applications like + modeling partially completed learning curves in AutoML. """ t_fit_start = time.monotonic() transforms = transforms or [] @@ -184,6 +190,7 @@ def __init__( self._fit_out_of_design = fit_out_of_design self._fit_abandoned = fit_abandoned self._fit_tracking_metrics = fit_tracking_metrics + self._fit_only_completed_map_metrics = fit_only_completed_map_metrics self.outcomes: list[str] = [] self._experiment_has_immutable_search_space_and_opt_config: bool = ( experiment is not None and experiment.immutable_search_space_and_opt_config @@ -292,12 +299,21 @@ def _prepare_observations( ) -> list[Observation]: if experiment is None or data is None: return [] - return observations_from_data( - experiment=experiment, - data=data, - statuses_to_include=self.statuses_to_fit, - statuses_to_include_map_metric=self.statuses_to_fit_map_metric, - ) + if not self._fit_only_completed_map_metrics and isinstance(data, MapData): + return observations_from_map_data( + experiment=experiment, + map_data=data, + map_keys_as_parameters=True, + statuses_to_include=self.statuses_to_fit, + statuses_to_include_map_metric=self.statuses_to_fit_map_metric, + ) + else: + return observations_from_data( + experiment=experiment, + data=data, + statuses_to_include=self.statuses_to_fit, + statuses_to_include_map_metric=self.statuses_to_fit_map_metric, + ) def _transform_data( self, @@ -557,7 +573,11 @@ def statuses_to_fit(self) -> set[TrialStatus]: @property def statuses_to_fit_map_metric(self) -> set[TrialStatus]: """Statuses to fit the model on.""" - return {TrialStatus.COMPLETED} + return ( + {TrialStatus.COMPLETED} + if self._fit_only_completed_map_metrics + else self.statuses_to_fit + ) @training_in_design.setter def training_in_design(self, training_in_design: list[bool]) -> None: diff --git a/ax/modelbridge/tests/test_base_modelbridge.py b/ax/modelbridge/tests/test_base_modelbridge.py index 6cc13a12c4a..133ffc5f335 100644 --- a/ax/modelbridge/tests/test_base_modelbridge.py +++ b/ax/modelbridge/tests/test_base_modelbridge.py @@ -15,6 +15,7 @@ from ax.core.arm import Arm from ax.core.data import Data from ax.core.experiment import Experiment +from ax.core.map_data import MapData from ax.core.metric import Metric from ax.core.objective import Objective, ScalarizedObjective from ax.core.observation import ObservationData, ObservationFeatures @@ -1040,3 +1041,46 @@ def test_SetModelSpace(self) -> None: ) self.assertEqual(sum(m.training_in_design), 7) self.assertEqual(m.model_space.parameters["x2"].upper, 20) + + @mock.patch( + "ax.modelbridge.base.observations_from_map_data", + autospec=True, + return_value=([get_observation1()]), + ) + @mock.patch( + "ax.modelbridge.base.observations_from_data", + autospec=True, + return_value=([get_observation1(), get_observation2()]), + ) + def test_fit_only_completed_map_metrics( + self, mock_observations_from_data: Mock, mock_observations_from_map_data: Mock + ) -> None: + # NOTE: If empty data object is not passed, observations are not + # extracted, even with mock. + # _prepare_observations is called in the constructor and itself calls + # observations_from_map_data. + Adapter( + search_space=get_search_space_for_value(), + model=0, + experiment=get_experiment_for_value(), + data=MapData(), + status_quo_name="1_1", + fit_only_completed_map_metrics=False, + ) + self.assertTrue(mock_observations_from_map_data.called) + self.assertFalse(mock_observations_from_data.called) + + # calling without map data calls regular observations_from_data even + # if fit_only_completed_map_metrics is False + mock_observations_from_data.reset_mock() + mock_observations_from_map_data.reset_mock() + Adapter( + search_space=get_search_space_for_value(), + model=0, + experiment=get_experiment_for_value(), + data=Data(), + status_quo_name="1_1", + fit_only_completed_map_metrics=False, + ) + self.assertFalse(mock_observations_from_map_data.called) + self.assertTrue(mock_observations_from_data.called) diff --git a/ax/modelbridge/tests/test_registry.py b/ax/modelbridge/tests/test_registry.py index 11a8e3fc6f6..543c64b3aa2 100644 --- a/ax/modelbridge/tests/test_registry.py +++ b/ax/modelbridge/tests/test_registry.py @@ -171,6 +171,7 @@ def test_enum_sobol_legacy_GPEI(self) -> None: "fit_tracking_metrics": True, "fit_on_init": True, "default_model_gen_options": None, + "fit_only_completed_map_metrics": True, }, ) prior_kwargs = {"lengthscale_prior": GammaPrior(6.0, 6.0)} diff --git a/ax/modelbridge/torch.py b/ax/modelbridge/torch.py index 318c2532c0b..48df6b04efe 100644 --- a/ax/modelbridge/torch.py +++ b/ax/modelbridge/torch.py @@ -119,6 +119,7 @@ def __init__( fit_tracking_metrics: bool = True, fit_on_init: bool = True, default_model_gen_options: TConfig | None = None, + fit_only_completed_map_metrics: bool = True, ) -> None: # This warning is being added while we are on 0.4.3, so it will be # released in 0.4.4 or 0.5.0. The `torch_dtype` argument can be removed @@ -161,6 +162,7 @@ def __init__( fit_abandoned=fit_abandoned, fit_tracking_metrics=fit_tracking_metrics, fit_on_init=fit_on_init, + fit_only_completed_map_metrics=fit_only_completed_map_metrics, ) def feature_importances(self, metric_name: str) -> dict[str, float]: From bf1f4704096462c17f950fb7a80e256da2256461 Mon Sep 17 00:00:00 2001 From: Sebastian Ament Date: Wed, 5 Feb 2025 16:26:20 -0800 Subject: [PATCH 2/2] Separating out `get_map_values` helper from `MapUnitX` transform (#3313) Summary: This commit separates out a `get_map_values` helper function from the `MapUnitX` transform. Differential Revision: D69213291 --- ax/modelbridge/transforms/map_unit_x.py | 36 +++++++++++++++++++------ 1 file changed, 28 insertions(+), 8 deletions(-) diff --git a/ax/modelbridge/transforms/map_unit_x.py b/ax/modelbridge/transforms/map_unit_x.py index 732502ed464..6fd389a4535 100644 --- a/ax/modelbridge/transforms/map_unit_x.py +++ b/ax/modelbridge/transforms/map_unit_x.py @@ -41,14 +41,7 @@ def __init__( ) -> None: assert observations is not None, "MapUnitX requires observations" assert search_space is not None, "MapUnitX requires search space" - # Loop through observation features and identify parameters that - # are not part of the search space. Store all observed values to - # infer bounds - map_values = defaultdict(list) - for obs in observations: - for p in obs.features.parameters: - if p not in search_space.parameters: - map_values[p].append(obs.features.parameters[p]) + map_values = get_map_values(search_space, observations) # pyre-fixme[24]: Generic type `list` expects 1 type parameter, use # `typing.List` to avoid runtime subscripting errors. @@ -81,3 +74,30 @@ def untransform_observation_features( scale_fac = (u - l) / self.target_range obsf.parameters[p_name] = scale_fac * (param - self.target_lb) + l return observation_features + + +def get_map_values( + search_space: SearchSpace, + observations: list[Observation], +) -> dict[str, list[float]]: + """Computes a dictionary mapping the name of a map parameter to its associated + progression values, in the same order as they occur in the observations. + + Args: + search_space: The search space. + observations: A list of observations associated with the search space. + + Returns: + The dictionary mapping the name of a map metric to the associated values, + in the same order they occur in `observations`. + """ + # Loop through observation features and identify parameters that + # are not part of the search space. Store all observed values to + # infer bounds + map_values = defaultdict(list) + for obs in observations: + # if we had access to the original data object, could loop over data.map_keys + for p in obs.features.parameters: + if p not in search_space.parameters: + map_values[p].append(obs.features.parameters[p]) + return map_values