diff --git a/ax/core/search_space.py b/ax/core/search_space.py index 1268f9bdeef..fcff0269f3d 100644 --- a/ax/core/search_space.py +++ b/ax/core/search_space.py @@ -761,7 +761,6 @@ def _gen_dummy_values_to_complete_flat_parameterization( dummy_values_to_inject[param_name] = choice(param.values) elif isinstance(param, RangeParameter): val = uniform(param.lower, param.upper) - print(val) if param.parameter_type is ParameterType.INT: val += 0.5 dummy_values_to_inject[param_name] = param.cast(val) diff --git a/ax/core/tests/test_utils.py b/ax/core/tests/test_utils.py index 084e7a47bbb..4904df54906 100644 --- a/ax/core/tests/test_utils.py +++ b/ax/core/tests/test_utils.py @@ -4,16 +4,20 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. +from unittest.mock import patch + import numpy as np import pandas as pd +from ax.core.arm import Arm +from ax.core.base_trial import TrialStatus + from ax.core.data import Data +from ax.core.generator_run import GeneratorRun from ax.core.metric import Metric -from ax.core.objective import MultiObjective, Objective -from ax.core.optimization_config import ( - MultiObjectiveOptimizationConfig, - OptimizationConfig, -) -from ax.core.outcome_constraint import ObjectiveThreshold, OutcomeConstraint +from ax.core.objective import Objective +from ax.core.observation import ObservationFeatures +from ax.core.optimization_config import OptimizationConfig +from ax.core.outcome_constraint import OutcomeConstraint from ax.core.types import ComparisonOp from ax.core.utils import ( best_feasible_objective, @@ -21,16 +25,64 @@ get_missing_metrics_by_name, get_model_times, get_model_trace_of_times, + get_pending_observation_features, + get_pending_observation_features_based_on_trial_status as get_pending_status, MissingMetrics, ) -from ax.modelbridge.modelbridge_utils import feasible_hypervolume + +from ax.utils.common.constants import Keys from ax.utils.common.testutils import TestCase -from ax.utils.testing.core_stubs import get_robust_branin_experiment +from ax.utils.testing.core_stubs import ( + get_experiment, + get_hierarchical_search_space_experiment, + get_robust_branin_experiment, +) from pyre_extensions import none_throws class UtilsTest(TestCase): def setUp(self) -> None: + self.experiment = get_experiment() + self.arm = Arm({"x": 5, "y": "foo", "z": True, "w": 5}) + self.trial = self.experiment.new_trial(GeneratorRun([self.arm])) + self.experiment_2 = get_experiment() + self.batch_trial = self.experiment_2.new_batch_trial(GeneratorRun([self.arm])) + self.batch_trial.set_status_quo_with_weight(self.experiment_2.status_quo, 1) + self.obs_feat = ObservationFeatures.from_arm( + arm=self.trial.arm, trial_index=np.int64(self.trial.index) + ) + self.hss_arm = Arm({"model": "XGBoost", "num_boost_rounds": 12}) + self.hss_exp = get_hierarchical_search_space_experiment() + self.hss_gr = GeneratorRun( + arms=[self.hss_arm], + candidate_metadata_by_arm_signature={ + self.hss_arm.signature: { + Keys.FULL_PARAMETERIZATION: { + "model_name": "XGBoost", + "num_boost_rounds": 12, + "learning_rate": 0.01, + "l2_reg_weight": 0.0001, + } + } + }, + ) + self.hss_trial = self.hss_exp.new_trial(self.hss_gr) + self.hss_cand_metadata = self.hss_trial._get_candidate_metadata( + arm_name=self.hss_arm.name + ) + self.hss_full_parameterization = self.hss_cand_metadata.get( + Keys.FULL_PARAMETERIZATION + ).copy() + self.hss_obs_feat = ObservationFeatures.from_arm( + arm=self.hss_arm, + trial_index=np.int64(self.hss_trial.index), + metadata=self.hss_cand_metadata, + ) + self.hss_obs_feat_all_params = ObservationFeatures.from_arm( + arm=Arm(self.hss_full_parameterization), + trial_index=np.int64(self.hss_trial.index), + metadata={Keys.FULL_PARAMETERIZATION: self.hss_full_parameterization}, + ) self.df = pd.DataFrame( [ { @@ -134,62 +186,6 @@ def test_best_feasible_objective(self) -> None: ) self.assertEqual(list(bfo), [1.0, 1.0, 2.0]) - def test_feasible_hypervolume(self) -> None: - ma = Metric(name="a", lower_is_better=False) - mb = Metric(name="b", lower_is_better=True) - mc = Metric(name="c", lower_is_better=False) - optimization_config = MultiObjectiveOptimizationConfig( - objective=MultiObjective(metrics=[ma, mb]), - outcome_constraints=[ - OutcomeConstraint( - mc, - op=ComparisonOp.GEQ, - bound=0, - relative=False, - ) - ], - objective_thresholds=[ - ObjectiveThreshold( - ma, - bound=1.0, - ), - ObjectiveThreshold( - mb, - bound=1.0, - ), - ], - ) - feas_hv = feasible_hypervolume( - optimization_config, - values={ - "a": np.array( - [ - 1.0, - 3.0, - 2.0, - 2.0, - ] - ), - "b": np.array( - [ - 0.0, - 1.0, - 0.0, - 0.0, - ] - ), - "c": np.array( - [ - 0.0, - -0.0, - 1.0, - -2.0, - ] - ), - }, - ) - self.assertEqual(list(feas_hv), [0.0, 0.0, 1.0, 1.0]) - def test_get_model_times(self) -> None: exp = get_robust_branin_experiment(num_sobol_trials=2) fit_times, gen_times = get_model_trace_of_times(exp) @@ -200,3 +196,348 @@ def test_get_model_times(self) -> None: self.assertTrue(all(elt >= 0 for elt in gen_times_not_none)) self.assertEqual(sum(fit_times_not_none), total_fit_time) self.assertEqual(sum(gen_times_not_none), total_gen_time) + + def test_get_pending_observation_features(self) -> None: + # Pending observations should be none if there aren't any. + self.assertIsNone(get_pending_observation_features(self.experiment)) + self.trial.mark_running(no_runner_required=True) + # Now that the trial is deployed, it should become a pending trial on the + # experiment and appear as pending for all metrics. + self.assertEqual( + get_pending_observation_features(self.experiment), + {"tracking": [self.obs_feat], "m2": [self.obs_feat], "m1": [self.obs_feat]}, + ) + # With `fetch_data` on trial returning data for metric "m2", that metric + # should no longer have pending observation features. + with patch.object( + self.trial, + "lookup_data", + return_value=Data.from_evaluations( + {self.trial.arm.name: {"m2": (1, 0)}}, trial_index=self.trial.index + ), + ): + self.assertEqual( + get_pending_observation_features(self.experiment), + {"tracking": [self.obs_feat], "m2": [], "m1": [self.obs_feat]}, + ) + # When a trial is marked failed, it should no longer appear in pending. + self.trial.mark_failed() + self.assertIsNone(get_pending_observation_features(self.experiment)) + # When a trial is abandoned, it should appear in pending features whether + # or not there is data for it. + self.trial._status = TrialStatus.ABANDONED # Cannot re-mark a failed trial. + self.assertEqual( + get_pending_observation_features(self.experiment), + {"tracking": [self.obs_feat], "m2": [self.obs_feat], "m1": [self.obs_feat]}, + ) + # When an arm is abandoned, it should appear in pending features whether + # or not there is data for it. + self.batch_trial.mark_arm_abandoned(arm_name="0_0") + # Checking with data for all metrics. + with patch.object( + self.batch_trial, + "fetch_data", + return_value=Metric._wrap_trial_data_multi( + data=Data.from_evaluations( + { + self.batch_trial.arms[0].name: { + "m1": (1, 0), + "m2": (1, 0), + "tracking": (1, 0), + } + }, + trial_index=self.trial.index, + ), + ), + ): + self.assertEqual( + get_pending_observation_features(self.experiment), + { + "tracking": [self.obs_feat], + "m2": [self.obs_feat], + "m1": [self.obs_feat], + }, + ) + # Checking with data for all metrics. + with patch.object( + self.trial, + "fetch_data", + return_value=Metric._wrap_trial_data_multi( + data=Data.from_evaluations( + { + self.trial.arm.name: { + "m1": (1, 0), + "m2": (1, 0), + "tracking": (1, 0), + } + }, + trial_index=self.trial.index, + ), + ), + ): + self.assertEqual( + get_pending_observation_features(self.experiment), + { + "tracking": [self.obs_feat], + "m2": [self.obs_feat], + "m1": [self.obs_feat], + }, + ) + + def test_get_pending_observation_features_hss(self) -> None: + # Pending observations should be none if there aren't any. + self.assertIsNone(get_pending_observation_features(self.hss_exp)) + self.hss_trial.mark_running(no_runner_required=True) + # Now that the trial is deployed, it should become a pending trial on the + # experiment and appear as pending for all metrics. + pending = get_pending_observation_features(self.hss_exp) + self.assertEqual( + pending, + { + "m1": [self.hss_obs_feat], + "m2": [self.hss_obs_feat], + }, + ) + + # Check that transforming observation features works correctly (it should inject + # full parameterization into resulting obs.feats.) + for p in none_throws(pending).values(): + for pf in p: + self.assertEqual( + none_throws(pf.metadata), + none_throws(self.hss_gr.candidate_metadata_by_arm_signature)[ + self.hss_arm.signature + ], + ) + + # With `fetch_data` on trial returning data for metric "m2", that metric + # should no longer have pending observation features. + with patch.object( + self.hss_trial, + "lookup_data", + return_value=Data.from_evaluations( + {self.hss_trial.arm.name: {"m2": (1, 0)}}, + trial_index=self.hss_trial.index, + ), + ): + self.assertEqual( + get_pending_observation_features(self.hss_exp), + {"m2": [], "m1": [self.hss_obs_feat]}, + ) + # When a trial is marked failed, it should no longer appear in pending. + self.hss_trial.mark_failed() + self.assertIsNone(get_pending_observation_features(self.hss_exp)) + + # When an arm is abandoned, it should appear in pending features whether + # or not there is data for it. + hss_exp = get_hierarchical_search_space_experiment() + hss_batch_trial = hss_exp.new_batch_trial(generator_run=self.hss_gr) + hss_batch_trial.mark_arm_abandoned(hss_batch_trial.arms[0].name) + # Checking with data for all metrics. + with patch.object( + hss_batch_trial, + "fetch_data", + return_value=Metric._wrap_trial_data_multi( + data=Data.from_evaluations( + { + hss_batch_trial.arms[0].name: { + "m1": (1, 0), + "m2": (1, 0), + } + }, + trial_index=hss_batch_trial.index, + ), + ), + ): + pending = get_pending_observation_features(hss_exp) + self.assertEqual( + pending, + { + "m1": [self.hss_obs_feat], + "m2": [self.hss_obs_feat], + }, + ) + # Check that candidate metadata is property propagated for abandoned arm. + for p in none_throws(pending).values(): + for pf in p: + self.assertEqual( + none_throws(pf.metadata), + none_throws(self.hss_gr.candidate_metadata_by_arm_signature)[ + self.hss_arm.signature + ], + ) + + # Checking with data for all metrics. + with patch.object( + hss_batch_trial, + "fetch_data", + return_value=Metric._wrap_trial_data_multi( + data=Data.from_evaluations( + { + hss_batch_trial.arms[0].name: { + "m1": (1, 0), + "m2": (1, 0), + } + }, + trial_index=hss_batch_trial.index, + ), + ), + ): + self.assertEqual( + get_pending_observation_features(hss_exp), + { + "m2": [self.hss_obs_feat], + "m1": [self.hss_obs_feat], + }, + ) + + def test_get_pending_observation_features_batch_trial(self) -> None: + # Check the same functionality for batched trials. + self.assertIsNone(get_pending_observation_features(self.experiment_2)) + self.batch_trial.mark_running(no_runner_required=True) + # Status quo of this experiment is out-of-design, so it shouldn't be + # among the pending points. + self.assertEqual( + get_pending_observation_features(self.experiment_2), + { + "tracking": [self.obs_feat], + "m2": [self.obs_feat], + "m1": [self.obs_feat], + }, + ) + + # Status quo of this experiment is out-of-design, so it shouldn't be + # among the pending points. + sq_obs_feat = ObservationFeatures.from_arm( + self.batch_trial.arms_by_name.get("status_quo"), + trial_index=self.batch_trial.index, + ) + self.assertEqual( + get_pending_observation_features( + self.experiment_2, + include_out_of_design_points=True, + ), + { + "tracking": [self.obs_feat, sq_obs_feat], + "m2": [self.obs_feat, sq_obs_feat], + "m1": [self.obs_feat, sq_obs_feat], + }, + ) + self.batch_trial.mark_completed() + + # Set SQ to in-design; then we can expect it to appear among the pending + # points without specifying `include_out_of_design_points=True`. + exp = get_experiment(with_status_quo=False) + in_design_status_quo = Arm( + name="in_design_status_quo", + parameters={"w": 5.45, "x": 5, "y": "bar", "z": True}, + ) + exp.status_quo = in_design_status_quo + batch = exp.new_batch_trial().add_arm(self.arm) + batch.set_status_quo_with_weight(exp.status_quo, 1) + self.assertEqual(batch.status_quo, in_design_status_quo) + self.assertTrue( + exp.search_space.check_membership( + in_design_status_quo.parameters, raise_error=True + ) + ) + batch.mark_running(no_runner_required=True) + sq_obs_feat = ObservationFeatures.from_arm( + in_design_status_quo, + trial_index=batch.index, + ) + self.assertEqual( + get_pending_observation_features(exp), + { + "tracking": [self.obs_feat, sq_obs_feat], + "m2": [self.obs_feat, sq_obs_feat], + "m1": [self.obs_feat, sq_obs_feat], + }, + ) + + def test_get_pending_observation_features_based_on_trial_status(self) -> None: + # Pending observations should be none if there aren't any as trial is + # candidate. + self.assertTrue(self.trial.status.is_candidate) + self.assertIsNone(get_pending_status(self.experiment)) + self.trial.mark_staged() + # Now that the trial is staged, it should become a pending trial on the + # experiment and appear as pending for all metrics. + self.assertEqual( + get_pending_status(self.experiment), + {"tracking": [self.obs_feat], "m2": [self.obs_feat], "m1": [self.obs_feat]}, + ) + # Same should be true for running trial. + # NOTE: Can't mark a staged trial running unless it uses a runner that + # specifically requires staging; hacking around that here since the marking + # logic does not matter for this test. + self.trial._status = TrialStatus.RUNNING + # Now that the trial is staged, it should become a pending trial on the + # experiment and appear as pending for all metrics. + self.assertEqual( + get_pending_status(self.experiment), + {"tracking": [self.obs_feat], "m2": [self.obs_feat], "m1": [self.obs_feat]}, + ) + # When a trial is marked failed, it should no longer appear in pending. + self.trial.mark_failed() + self.assertIsNone(get_pending_status(self.experiment)) + # And if the trial is abandoned, it should always appear in pending features. + self.trial._status = TrialStatus.ABANDONED # Cannot re-mark a failed trial. + self.assertEqual( + get_pending_status(self.experiment), + {"tracking": [self.obs_feat], "m2": [self.obs_feat], "m1": [self.obs_feat]}, + ) + + def test_get_pending_observation_features_based_on_trial_status_hss(self) -> None: + self.assertTrue(self.hss_trial.status.is_candidate) + self.assertIsNone(get_pending_status(self.hss_exp)) + self.hss_trial.mark_staged() + # Now that the trial is staged, it should become a pending trial on the + # experiment and appear as pending for all metrics. + pending = get_pending_status(self.hss_exp) + self.assertEqual( + pending, + { + "m1": [self.hss_obs_feat], + "m2": [self.hss_obs_feat], + }, + ) + + # Same should be true for running trial. + # NOTE: Can't mark a staged trial running unless it uses a runner that + # specifically requires staging; hacking around that here since the marking + # logic does not matter for this test. + self.hss_trial._status = TrialStatus.RUNNING + # Now that the trial is staged, it should become a pending trial on the + # experiment and appear as pending for all metrics. + pending = get_pending_status(self.hss_exp) + self.assertEqual( + pending, + { + "m1": [self.hss_obs_feat], + "m2": [self.hss_obs_feat], + }, + ) + # When a trial is marked failed, it should no longer appear in pending. + self.hss_trial.mark_failed() + self.assertIsNone(get_pending_status(self.hss_exp)) + # And if the trial is abandoned, it should always appear in pending features. + self.hss_trial._status = TrialStatus.ABANDONED # Cannot re-mark a failed trial. + self.assertEqual( + pending, + { + "m1": [self.hss_obs_feat], + "m2": [self.hss_obs_feat], + }, + ) + + # Check that transforming observation features works correctly (it should inject + # full parameterization into resulting obs.feats.) + for p in none_throws(pending).values(): + for pf in p: + self.assertEqual( + none_throws(pf.metadata), + none_throws(self.hss_gr.candidate_metadata_by_arm_signature)[ + self.hss_arm.signature + ], + ) diff --git a/ax/core/utils.py b/ax/core/utils.py index 7a110a55333..7f3bfe6e128 100644 --- a/ax/core/utils.py +++ b/ax/core/utils.py @@ -8,6 +8,7 @@ from typing import Dict, Iterable, List, NamedTuple, Optional, Set, Tuple import numpy as np +from ax.core.arm import Arm from ax.core.base_trial import BaseTrial, TrialStatus from ax.core.batch_trial import BatchTrial from ax.core.data import Data @@ -198,7 +199,9 @@ def get_model_times(experiment: Experiment) -> Tuple[float, float]: def get_pending_observation_features( - experiment: Experiment, include_failed_as_pending: bool = False + experiment: Experiment, + *, + include_out_of_design_points: bool = False, ) -> Optional[Dict[str, List[ObservationFeatures]]]: """Computes a list of pending observation features (corresponding to arms that have been generated and deployed in the course of the experiment, but have not @@ -210,14 +213,23 @@ def get_pending_observation_features( Args: experiment: Experiment, pending features on which we seek to compute. - include_failed_as_pending: Whether to include failed trials as pending - (for example, to avoid the model suggesting them again). + include_out_of_design_points: By default, this function will not include + "out of design" points (those that are not in the search space) among + the pending points. This is because pending points are generally used to + help the model avoid re-suggesting the same points again. For points + outside of the search space, this will not happen, so they typically do + not need to be included. However, if the user wants to include them, + they can be included by setting this flag to ``True``. Returns: An optional mapping from metric names to a list of observation features, pending for that metric (i.e. do not have evaluation data for that metric). If there are no pending features for any of the metrics, return is None. """ + + def _is_in_design(arm: Arm) -> bool: + return experiment.search_space.check_membership(parameterization=arm.parameters) + pending_features = {} # Note that this assumes that if a metric appears in fetched data, the trial is # not pending for the metric. Where only the most recent data matters, this will @@ -227,50 +239,37 @@ def get_pending_observation_features( for metric_name in experiment.metrics: if metric_name not in pending_features: pending_features[metric_name] = [] - include_since_failed = include_failed_as_pending and trial.status.is_failed - if isinstance(trial, BatchTrial): - if trial.status.is_abandoned or ( - (trial.status.is_deployed or include_since_failed) - and metric_name not in dat.df.metric_name.values - and trial.arms is not None - ): - for arm in trial.arms: - not_none(pending_features.get(metric_name)).append( - ObservationFeatures.from_arm( - arm=arm, - trial_index=np.int64(trial_index), - metadata=trial._get_candidate_metadata( - arm_name=arm.name - ), - ) - ) - abandoned_arms = trial.abandoned_arms - for abandoned_arm in abandoned_arms: - not_none(pending_features.get(metric_name)).append( + + if trial.status.is_abandoned or ( + trial.status.is_deployed + and metric_name not in dat.df.metric_name.values + and trial.arms is not None + ): + for arm in trial.arms: + # Do not add out-of-design points unless requested. + if not include_out_of_design_points and not _is_in_design(arm=arm): + continue + pending_features[metric_name].append( ObservationFeatures.from_arm( - arm=abandoned_arm, + arm=arm, trial_index=np.int64(trial_index), - metadata=trial._get_candidate_metadata( - arm_name=abandoned_arm.name - ), + metadata=trial._get_candidate_metadata(arm_name=arm.name), ) ) - if isinstance(trial, Trial): - if trial.status.is_abandoned or ( - (trial.status.is_deployed or include_since_failed) - and metric_name not in dat.df.metric_name.values - and trial.arm is not None - ): + # Also add abandoned arms as pending for all metrics. + if isinstance(trial, BatchTrial): + for arm in trial.abandoned_arms: + if not include_out_of_design_points and not _is_in_design(arm=arm): + continue not_none(pending_features.get(metric_name)).append( ObservationFeatures.from_arm( - arm=not_none(trial.arm), + arm=arm, trial_index=np.int64(trial_index), - metadata=trial._get_candidate_metadata( - arm_name=not_none(trial.arm).name - ), + metadata=trial._get_candidate_metadata(arm_name=arm.name), ) ) + return pending_features if any(x for x in pending_features.values()) else None diff --git a/ax/modelbridge/tests/test_convert_metric_names.py b/ax/modelbridge/tests/test_convert_metric_names_transform.py similarity index 100% rename from ax/modelbridge/tests/test_convert_metric_names.py rename to ax/modelbridge/tests/test_convert_metric_names_transform.py diff --git a/ax/modelbridge/tests/test_inverse_gaussian_cdf_y.py b/ax/modelbridge/tests/test_inverse_gaussian_cdf_y_transform.py similarity index 100% rename from ax/modelbridge/tests/test_inverse_gaussian_cdf_y.py rename to ax/modelbridge/tests/test_inverse_gaussian_cdf_y_transform.py diff --git a/ax/modelbridge/tests/test_power_transform_y.py b/ax/modelbridge/tests/test_power_y_transform.py similarity index 100% rename from ax/modelbridge/tests/test_power_transform_y.py rename to ax/modelbridge/tests/test_power_y_transform.py diff --git a/ax/modelbridge/tests/test_robust.py b/ax/modelbridge/tests/test_robust_modelbridge.py similarity index 100% rename from ax/modelbridge/tests/test_robust.py rename to ax/modelbridge/tests/test_robust_modelbridge.py diff --git a/ax/modelbridge/tests/test_rounding.py b/ax/modelbridge/tests/test_rounding_transform.py similarity index 100% rename from ax/modelbridge/tests/test_rounding.py rename to ax/modelbridge/tests/test_rounding_transform.py diff --git a/ax/modelbridge/tests/test_stratified_standardize_y.py b/ax/modelbridge/tests/test_stratified_standardize_y_transform.py similarity index 100% rename from ax/modelbridge/tests/test_stratified_standardize_y.py rename to ax/modelbridge/tests/test_stratified_standardize_y_transform.py diff --git a/ax/modelbridge/tests/test_torch_modelbridge_moo.py b/ax/modelbridge/tests/test_torch_moo_modelbridge.py similarity index 100% rename from ax/modelbridge/tests/test_torch_modelbridge_moo.py rename to ax/modelbridge/tests/test_torch_moo_modelbridge.py diff --git a/ax/modelbridge/tests/test_utils.py b/ax/modelbridge/tests/test_utils.py index d0a46f0c4d5..cf25143b3a9 100644 --- a/ax/modelbridge/tests/test_utils.py +++ b/ax/modelbridge/tests/test_utils.py @@ -9,25 +9,23 @@ import numpy as np from ax.core.arm import Arm -from ax.core.base_trial import TrialStatus from ax.core.data import Data from ax.core.generator_run import GeneratorRun from ax.core.metric import Metric from ax.core.objective import MultiObjective, Objective -from ax.core.observation import Observation, ObservationData, ObservationFeatures +from ax.core.observation import ObservationData, ObservationFeatures +from ax.core.optimization_config import MultiObjectiveOptimizationConfig from ax.core.outcome_constraint import ( ObjectiveThreshold, OutcomeConstraint, ScalarizedOutcomeConstraint, ) from ax.core.types import ComparisonOp -from ax.core.utils import ( - get_pending_observation_features, - get_pending_observation_features_based_on_trial_status as get_pending_status, -) +from ax.core.utils import get_pending_observation_features from ax.modelbridge.modelbridge_utils import ( extract_objective_thresholds, extract_outcome_constraints, + feasible_hypervolume, observation_data_to_array, pending_observations_as_array_list, ) @@ -39,6 +37,8 @@ get_experiment, get_hierarchical_search_space_experiment, ) +from pyre_extensions import none_throws + TEST_PARAMETERIZATON_LIST = ["5", "foo", "True", "5"] @@ -82,422 +82,6 @@ def setUp(self) -> None: metadata={Keys.FULL_PARAMETERIZATION: self.hss_full_parameterization}, ) - def test_get_pending_observation_features(self) -> None: - # Pending observations should be none if there aren't any. - self.assertIsNone(get_pending_observation_features(self.experiment)) - self.trial.mark_running(no_runner_required=True) - # Now that the trial is deployed, it should become a pending trial on the - # experiment and appear as pending for all metrics. - self.assertEqual( - get_pending_observation_features(self.experiment), - {"tracking": [self.obs_feat], "m2": [self.obs_feat], "m1": [self.obs_feat]}, - ) - # With `fetch_data` on trial returning data for metric "m2", that metric - # should no longer have pending observation features. - with patch.object( - self.trial, - "lookup_data", - return_value=Data.from_evaluations( - {self.trial.arm.name: {"m2": (1, 0)}}, trial_index=self.trial.index - ), - ): - self.assertEqual( - get_pending_observation_features(self.experiment), - {"tracking": [self.obs_feat], "m2": [], "m1": [self.obs_feat]}, - ) - # When a trial is marked failed, it should no longer appear in pending... - self.trial.mark_failed() - self.assertIsNone(get_pending_observation_features(self.experiment)) - # ... unless specified to include failed trials in pending observations. - self.assertEqual( - get_pending_observation_features( - self.experiment, include_failed_as_pending=True - ), - {"tracking": [self.obs_feat], "m2": [self.obs_feat], "m1": [self.obs_feat]}, - ) - # When a trial is abandoned, it should appear in pending features whether - # or not there is data for it. - self.trial._status = TrialStatus.ABANDONED # Cannot re-mark a failed trial. - self.assertEqual( - get_pending_observation_features( - self.experiment, include_failed_as_pending=True - ), - {"tracking": [self.obs_feat], "m2": [self.obs_feat], "m1": [self.obs_feat]}, - ) - # When an arm is abandoned, it should appear in pending features whether - # or not there is data for it. - self.batch_trial.mark_arm_abandoned(arm_name="0_0") - # Checking with data for all metrics. - with patch.object( - self.batch_trial, - "fetch_data", - return_value=Metric._wrap_trial_data_multi( - data=Data.from_evaluations( - { - self.batch_trial.arms[0].name: { - "m1": (1, 0), - "m2": (1, 0), - "tracking": (1, 0), - } - }, - trial_index=self.trial.index, - ), - ), - ): - self.assertEqual( - get_pending_observation_features( - self.experiment, include_failed_as_pending=True - ), - { - "tracking": [self.obs_feat], - "m2": [self.obs_feat], - "m1": [self.obs_feat], - }, - ) - # Checking with data for all metrics. - with patch.object( - self.trial, - "fetch_data", - return_value=Metric._wrap_trial_data_multi( - data=Data.from_evaluations( - { - self.trial.arm.name: { - "m1": (1, 0), - "m2": (1, 0), - "tracking": (1, 0), - } - }, - trial_index=self.trial.index, - ), - ), - ): - self.assertEqual( - get_pending_observation_features(self.experiment), - { - "tracking": [self.obs_feat], - "m2": [self.obs_feat], - "m1": [self.obs_feat], - }, - ) - - def test_get_pending_observation_features_hss(self) -> None: - # Pending observations should be none if there aren't any. - self.assertIsNone(get_pending_observation_features(self.hss_exp)) - self.hss_trial.mark_running(no_runner_required=True) - # Now that the trial is deployed, it should become a pending trial on the - # experiment and appear as pending for all metrics. - pending = get_pending_observation_features(self.hss_exp) - self.assertEqual( - pending, - { - "m1": [self.hss_obs_feat], - "m2": [self.hss_obs_feat], - }, - ) - - # Check that transforming observation features works correctly since this - # is applying `Cast` transform, it should inject full parameterization into - # resulting obs.feats.). Therefore, transforming the extracted pending features - # and observation features made from full parameterization should be the same. - obsd = ObservationData( - metric_names=["m1"], means=np.array([1.0]), covariance=np.array([[1.0]]) - ) - self.assertEqual( - self.hss_sobol._transform_data( - observations=[ - Observation(data=obsd, features=pending["m1"][0]) # pyre-ignore - ], - search_space=self.hss_exp.search_space, - transforms=self.hss_sobol._raw_transforms, - transform_configs=None, - ), - self.hss_sobol._transform_data( - observations=[ - Observation( - data=obsd, features=self.hss_obs_feat_all_params.clone() - ) - ], - search_space=self.hss_exp.search_space, - transforms=self.hss_sobol._raw_transforms, - transform_configs=None, - ), - ) - # With `fetch_data` on trial returning data for metric "m2", that metric - # should no longer have pending observation features. - with patch.object( - self.hss_trial, - "lookup_data", - return_value=Data.from_evaluations( - {self.hss_trial.arm.name: {"m2": (1, 0)}}, - trial_index=self.hss_trial.index, - ), - ): - self.assertEqual( - get_pending_observation_features(self.hss_exp), - {"m2": [], "m1": [self.hss_obs_feat]}, - ) - # When a trial is marked failed, it should no longer appear in pending... - self.hss_trial.mark_failed() - self.assertIsNone(get_pending_observation_features(self.hss_exp)) - # ... unless specified to include failed trials in pending observations. - self.assertEqual( - get_pending_observation_features( - self.hss_exp, include_failed_as_pending=True - ), - { - "m1": [self.hss_obs_feat], - "m2": [self.hss_obs_feat], - }, - ) - - # When an arm is abandoned, it should appear in pending features whether - # or not there is data for it. - hss_exp = get_hierarchical_search_space_experiment() - hss_batch_trial = hss_exp.new_batch_trial(generator_run=self.hss_gr) - hss_batch_trial.mark_arm_abandoned(hss_batch_trial.arms[0].name) - # Checking with data for all metrics. - with patch.object( - hss_batch_trial, - "fetch_data", - return_value=Metric._wrap_trial_data_multi( - data=Data.from_evaluations( - { - hss_batch_trial.arms[0].name: { - "m1": (1, 0), - "m2": (1, 0), - } - }, - trial_index=hss_batch_trial.index, - ), - ), - ): - pending = get_pending_observation_features( - hss_exp, include_failed_as_pending=True - ) - self.assertEqual( - pending, - { - "m1": [self.hss_obs_feat], - "m2": [self.hss_obs_feat], - }, - ) - # Check that candidate metadata is property propagated for abandoned arm. - self.assertEqual( - self.hss_sobol._transform_data( - observations=[Observation(data=obsd, features=pending["m1"][0])], - search_space=hss_exp.search_space, - transforms=self.hss_sobol._raw_transforms, - transform_configs=None, - ), - self.hss_sobol._transform_data( - observations=[ - Observation( - data=obsd, features=self.hss_obs_feat_all_params.clone() - ) - ], - search_space=hss_exp.search_space, - transforms=self.hss_sobol._raw_transforms, - transform_configs=None, - ), - ) - # Checking with data for all metrics. - with patch.object( - hss_batch_trial, - "fetch_data", - return_value=Metric._wrap_trial_data_multi( - data=Data.from_evaluations( - { - hss_batch_trial.arms[0].name: { - "m1": (1, 0), - "m2": (1, 0), - } - }, - trial_index=hss_batch_trial.index, - ), - ), - ): - self.assertEqual( - get_pending_observation_features(hss_exp), - { - "m2": [self.hss_obs_feat], - "m1": [self.hss_obs_feat], - }, - ) - - def test_get_pending_observation_features_batch_trial(self) -> None: - # Check the same functionality for batched trials. - self.assertIsNone(get_pending_observation_features(self.experiment_2)) - self.batch_trial.mark_running(no_runner_required=True) - sq_obs_feat = ObservationFeatures.from_arm( - self.batch_trial.arms_by_name.get("status_quo"), - trial_index=self.batch_trial.index, - ) - self.assertEqual( - get_pending_observation_features(self.experiment_2), - { - "tracking": [self.obs_feat, sq_obs_feat], - "m2": [self.obs_feat, sq_obs_feat], - "m1": [self.obs_feat, sq_obs_feat], - }, - ) - - def test_get_pending_observation_features_based_on_trial_status(self) -> None: - # Pending observations should be none if there aren't any as trial is - # candidate. - self.assertTrue(self.trial.status.is_candidate) - self.assertIsNone(get_pending_status(self.experiment)) - self.trial.mark_staged() - # Now that the trial is staged, it should become a pending trial on the - # experiment and appear as pending for all metrics. - self.assertEqual( - get_pending_status(self.experiment), - {"tracking": [self.obs_feat], "m2": [self.obs_feat], "m1": [self.obs_feat]}, - ) - # Same should be true for running trial. - # NOTE: Can't mark a staged trial running unless it uses a runner that - # specifically requires staging; hacking around that here since the marking - # logic does not matter for this test. - self.trial._status = TrialStatus.RUNNING - # Now that the trial is staged, it should become a pending trial on the - # experiment and appear as pending for all metrics. - self.assertEqual( - get_pending_status(self.experiment), - {"tracking": [self.obs_feat], "m2": [self.obs_feat], "m1": [self.obs_feat]}, - ) - # When a trial is marked failed, it should no longer appear in pending. - self.trial.mark_failed() - self.assertIsNone(get_pending_status(self.experiment)) - # And if the trial is abandoned, it should always appear in pending features. - self.trial._status = TrialStatus.ABANDONED # Cannot re-mark a failed trial. - self.assertEqual( - get_pending_status(self.experiment), - {"tracking": [self.obs_feat], "m2": [self.obs_feat], "m1": [self.obs_feat]}, - ) - - def test_get_pending_observation_features_based_on_trial_status_hss(self) -> None: - self.assertTrue(self.hss_trial.status.is_candidate) - self.assertIsNone(get_pending_status(self.hss_exp)) - self.hss_trial.mark_staged() - # Now that the trial is staged, it should become a pending trial on the - # experiment and appear as pending for all metrics. - pending = get_pending_status(self.hss_exp) - self.assertEqual( - pending, - { - "m1": [self.hss_obs_feat], - "m2": [self.hss_obs_feat], - }, - ) - - # Same should be true for running trial. - # NOTE: Can't mark a staged trial running unless it uses a runner that - # specifically requires staging; hacking around that here since the marking - # logic does not matter for this test. - self.hss_trial._status = TrialStatus.RUNNING - # Now that the trial is staged, it should become a pending trial on the - # experiment and appear as pending for all metrics. - pending = get_pending_status(self.hss_exp) - self.assertEqual( - pending, - { - "m1": [self.hss_obs_feat], - "m2": [self.hss_obs_feat], - }, - ) - # When a trial is marked failed, it should no longer appear in pending. - self.hss_trial.mark_failed() - self.assertIsNone(get_pending_status(self.hss_exp)) - # And if the trial is abandoned, it should always appear in pending features. - self.hss_trial._status = TrialStatus.ABANDONED # Cannot re-mark a failed trial. - self.assertEqual( - pending, - { - "m1": [self.hss_obs_feat], - "m2": [self.hss_obs_feat], - }, - ) - - # Check that transforming observation features works correctly since this - # is applying `Cast` transform, it should inject full parameterization into - # resulting obs.feats.). Therefore, transforming the extracted pending features - # and observation features made from full parameterization should be the same. - obsd = ObservationData( - metric_names=["m1"], means=np.array([1.0]), covariance=np.array([[1.0]]) - ) - self.assertEqual( - self.hss_sobol._transform_data( - observations=[ - Observation(data=obsd, features=pending["m1"][0]) # pyre-ignore - ], - search_space=self.hss_exp.search_space, - transforms=self.hss_sobol._raw_transforms, - transform_configs=None, - ), - self.hss_sobol._transform_data( - observations=[ - Observation( - data=obsd, features=self.hss_obs_feat_all_params.clone() - ) - ], - search_space=self.hss_exp.search_space, - transforms=self.hss_sobol._raw_transforms, - transform_configs=None, - ), - ) - - def test_pending_observations_as_array_list(self) -> None: - # Mark a trial dispatched so that there are pending observations. - self.trial.mark_running(no_runner_required=True) - # If outcome names are respected, unlisted metrics should be filtered out. - self.assertEqual( - [ - x.tolist() - # pyre-fixme[16]: Optional type has no attribute `__iter__`. - for x in pending_observations_as_array_list( - # pyre-fixme[6]: For 1st param expected `Dict[str, - # List[ObservationFeatures]]` but got `Optional[Dict[str, - # List[ObservationFeatures]]]`. - pending_observations=get_pending_observation_features( - self.experiment - ), - outcome_names=["m2", "m1"], - param_names=["x", "y", "z", "w"], - ) - ], - [[TEST_PARAMETERIZATON_LIST], [TEST_PARAMETERIZATON_LIST]], - ) - self.experiment.attach_data( - Data.from_evaluations( - {self.trial.arm.name: {"m2": (1, 0)}}, trial_index=self.trial.index - ) - ) - # With `fetch_data` on trial returning data for metric "m2", that metric - # should no longer have pending observation features. - with patch.object( - self.trial, - "fetch_data", - return_value=Data.from_evaluations( - {self.trial.arm.name: {"m2": (1, 0)}}, trial_index=self.trial.index - ), - ): - pending = get_pending_observation_features(self.experiment) - # There should be no pending observations for metric m2 now, since the - # only trial there is, has been updated with data for it. - self.assertEqual( - [ - x.tolist() - for x in pending_observations_as_array_list( - # pyre-fixme[6]: For 1st param expected `Dict[str, - # List[ObservationFeatures]]` but got `Optional[Dict[str, - # List[ObservationFeatures]]]`. - pending_observations=pending, - outcome_names=["m2", "m1"], - param_names=["x", "y", "z", "w"], - ) - ], - [[], [TEST_PARAMETERIZATON_LIST]], - ) - def test_extract_outcome_constraints(self) -> None: outcomes = ["m1", "m2", "m3"] # pass no outcome constraints @@ -639,3 +223,109 @@ def test_observation_data_to_array(self) -> None: np.array_equal(Ycov, np.array([[[5, 6, 4], [8, 9, 7], [2, 3, 1]]])) ) mock_warn.assert_called_once() + + def test_feasible_hypervolume(self) -> None: + ma = Metric(name="a", lower_is_better=False) + mb = Metric(name="b", lower_is_better=True) + mc = Metric(name="c", lower_is_better=False) + optimization_config = MultiObjectiveOptimizationConfig( + objective=MultiObjective(metrics=[ma, mb]), + outcome_constraints=[ + OutcomeConstraint( + mc, + op=ComparisonOp.GEQ, + bound=0, + relative=False, + ) + ], + objective_thresholds=[ + ObjectiveThreshold( + ma, + bound=1.0, + ), + ObjectiveThreshold( + mb, + bound=1.0, + ), + ], + ) + feas_hv = feasible_hypervolume( + optimization_config, + values={ + "a": np.array( + [ + 1.0, + 3.0, + 2.0, + 2.0, + ] + ), + "b": np.array( + [ + 0.0, + 1.0, + 0.0, + 0.0, + ] + ), + "c": np.array( + [ + 0.0, + -0.0, + 1.0, + -2.0, + ] + ), + }, + ) + self.assertEqual(list(feas_hv), [0.0, 0.0, 1.0, 1.0]) + + def test_pending_observations_as_array_list(self) -> None: + # Mark a trial dispatched so that there are pending observations. + self.trial.mark_running(no_runner_required=True) + # If outcome names are respected, unlisted metrics should be filtered out. + self.assertEqual( + [ + x.tolist() + for x in none_throws( + pending_observations_as_array_list( + pending_observations=none_throws( + get_pending_observation_features(self.experiment) + ), + outcome_names=["m2", "m1"], + param_names=["x", "y", "z", "w"], + ) + ) + ], + [[TEST_PARAMETERIZATON_LIST], [TEST_PARAMETERIZATON_LIST]], + ) + self.experiment.attach_data( + Data.from_evaluations( + {self.trial.arm.name: {"m2": (1, 0)}}, trial_index=self.trial.index + ) + ) + # With `fetch_data` on trial returning data for metric "m2", that metric + # should no longer have pending observation features. + with patch.object( + self.trial, + "fetch_data", + return_value=Data.from_evaluations( + {self.trial.arm.name: {"m2": (1, 0)}}, trial_index=self.trial.index + ), + ): + pending = none_throws(get_pending_observation_features(self.experiment)) + # There should be no pending observations for metric m2 now, since the + # only trial there is, has been updated with data for it. + self.assertEqual( + [ + x.tolist() + for x in none_throws( + pending_observations_as_array_list( + pending_observations=pending, + outcome_names=["m2", "m1"], + param_names=["x", "y", "z", "w"], + ) + ) + ], + [[], [TEST_PARAMETERIZATON_LIST]], + ) diff --git a/ax/modelbridge/tests/test_winsorize_transform_legacy.py b/ax/modelbridge/tests/test_winsorize_legacy_transform.py similarity index 100% rename from ax/modelbridge/tests/test_winsorize_transform_legacy.py rename to ax/modelbridge/tests/test_winsorize_legacy_transform.py