diff --git a/ax/benchmark/benchmark.py b/ax/benchmark/benchmark.py index 3abb9bbeae6..d479f9e007f 100644 --- a/ax/benchmark/benchmark.py +++ b/ax/benchmark/benchmark.py @@ -51,8 +51,10 @@ from ax.core.utils import get_model_times from ax.generation_strategy.generation_strategy import GenerationStrategy from ax.service.orchestrator import Orchestrator -from ax.service.utils.best_point import get_trace -from ax.service.utils.best_point_mixin import BestPointMixin +from ax.service.utils.best_point import ( + get_best_parameters_from_model_predictions_with_trial_index, + get_trace, +) from ax.service.utils.orchestrator_options import OrchestratorOptions, TrialType from ax.utils.common.logger import DEFAULT_LOG_LEVEL, get_logger from ax.utils.common.random import with_rng_seed @@ -326,9 +328,8 @@ def get_best_parameters( trial_indices: Iterable[int] | None = None, ) -> TParameterization | None: """ - Get the most promising point. - - Only SOO is supported. It will return None if no best point can be found. + Get the most promising point. Returns None if no point is predicted to + satisfy all outcome constraints. Args: experiment: The experiment to get the data from. This should contain @@ -338,9 +339,9 @@ def get_best_parameters( best point. trial_indices: Use data from only these trials. If None, use all data. """ - result = BestPointMixin._get_best_trial( + result = get_best_parameters_from_model_predictions_with_trial_index( experiment=experiment, - generation_strategy=generation_strategy, + adapter=generation_strategy.adapter, trial_indices=trial_indices, ) if result is None: @@ -507,7 +508,7 @@ def run_optimization_with_orchestrator( orchestrator = Orchestrator( experiment=experiment, - generation_strategy=method.generation_strategy.clone_reset(), + generation_strategy=method.generation_strategy, options=orchestrator_options, ) @@ -562,6 +563,8 @@ def benchmark_replication( Return: ``BenchmarkResult`` object. """ + # Reset the generation strategy to ensure that it is in an unused state. + method.generation_strategy = method.generation_strategy.clone_reset() experiment = run_optimization_with_orchestrator( problem=problem, method=method, diff --git a/ax/benchmark/tests/test_benchmark.py b/ax/benchmark/tests/test_benchmark.py index bfdcc036709..0e45c2568e2 100644 --- a/ax/benchmark/tests/test_benchmark.py +++ b/ax/benchmark/tests/test_benchmark.py @@ -17,7 +17,6 @@ import numpy as np import torch -from ax.adapter.factory import get_sobol from ax.adapter.registry import Generators from ax.benchmark.benchmark import ( _get_oracle_value_of_params, @@ -38,7 +37,6 @@ from ax.benchmark.benchmark_problem import ( BenchmarkProblem, create_problem_from_botorch, - get_continuous_search_space, get_moo_opt_config, get_soo_opt_config, ) @@ -275,8 +273,6 @@ def _test_replication_async(self, map_data: bool) -> None: map_data: If True, the test function produces time-series data with just one step, so behavior is the same as when map_data=False. """ - method = get_async_benchmark_method() - complete_out_of_order_runtimes = { 0: 2, 1: 1, @@ -345,6 +341,8 @@ def _test_replication_async(self, map_data: bool) -> None: } for case_name, step_runtime_fn in step_runtime_fns.items(): + method = get_async_benchmark_method() + with self.subTest(case_name, step_runtime_fn=step_runtime_fn): problem = get_async_benchmark_problem( map_data=map_data, @@ -417,12 +415,9 @@ def _test_replication_async(self, map_data: bool) -> None: }, f"Failure for trial {trial_index} with {case_name}", ) - self.assertFalse(np.isnan(result.inference_trace).any()) - self.assertEqual( - result.inference_trace.tolist(), - expected_traces[case_name], - msg=case_name, - ) + # Evaluating inference trace here is not relevant since the GS + # is not model-based. + self.assertTrue(np.isnan(result.inference_trace).all()) self.assertEqual( result.oracle_trace.tolist(), expected_traces[case_name], @@ -491,6 +486,7 @@ def test_run_optimization_with_orchestrator(self) -> None: none_throws(runner.simulated_backend_runner).simulator._verbose_logging ) + method = get_async_benchmark_method() with self.subTest("Logs not produced by default"), self.assertNoLogs( level=logging.INFO, logger=logger ), self.assertNoLogs(logger=logger): @@ -618,9 +614,9 @@ def test_early_stopping(self) -> None: self.assertEqual(max_run, {0: 4, 1: 2, 2: 2, 3: 2}) def test_replication_variable_runtime(self) -> None: - method = get_async_benchmark_method(max_pending_trials=1) for map_data in [False, True]: with self.subTest(map_data=map_data): + method = get_async_benchmark_method(max_pending_trials=1) problem = get_async_benchmark_problem( map_data=map_data, step_runtime_fn=lambda params: params["x0"] + 1, @@ -793,7 +789,11 @@ def test_replication_mbm(self) -> None: acquisition_cls=qLogNoisyExpectedImprovement, distribute_replications=False, ), - get_augmented_branin_problem(fidelity_or_task="fidelity"), + get_single_objective_benchmark_problem( + observe_noise_sd=False, + num_trials=6, + report_inference_value_as_trace=True, + ), "MBM::SingleTaskGP_qLogNEI", ), ]: @@ -1196,6 +1196,7 @@ def test_get_opt_trace_by_cumulative_epochs(self) -> None: ): get_opt_trace_by_steps(experiment=experiment) + method = get_sobol_benchmark_method(distribute_replications=False) with self.subTest("Constrained"): problem = get_benchmark_problem("constrained_gramacy_observed_noise") experiment = self.run_optimization_with_orchestrator( @@ -1237,72 +1238,20 @@ def test_get_benchmark_result_with_cumulative_steps(self) -> None: self.assertLessEqual(transformed.score_trace.min(), result.score_trace.min()) def test_get_best_parameters(self) -> None: - """ - Whether this produces the correct values is tested more thoroughly in - other tests such as `test_replication_with_inference_value` and - `test_get_inference_trace_from_params`. Setting up an experiment with - data and trials without just running a benchmark is a pain, so in those - tests, we just run a benchmark. - """ - gs = get_sobol_generation_strategy() - - search_space = get_continuous_search_space(bounds=[(0, 1)]) - moo_config = get_moo_opt_config(outcome_names=["a", "b"], ref_point=[0, 0]) - experiment = Experiment( - name="test", - is_test=True, - search_space=search_space, - optimization_config=moo_config, + experiment = get_experiment_with_observations(observations=[[1.0, 2.0]]) + generation_strategy = get_sobol_generation_strategy() + mock_function = ( + "ax.benchmark.benchmark." + "get_best_parameters_from_model_predictions_with_trial_index" ) - with self.subTest("MOO not supported"), self.assertRaisesRegex( - NotImplementedError, "Please use `get_pareto_optimal_parameters`" - ): - get_best_parameters(experiment=experiment, generation_strategy=gs) - - soo_config = get_soo_opt_config(outcome_names=["a"]) - with self.subTest("Empty experiment"): - result = get_best_parameters( - experiment=experiment.clone_with(optimization_config=soo_config), - generation_strategy=gs, - ) + with patch(mock_function, return_value=None): + result = get_best_parameters(experiment, generation_strategy) self.assertIsNone(result) - with self.subTest("All constraints violated"): - experiment = get_experiment_with_observations( - observations=[[1, -1], [2, -1]], - constrained=True, - ) - best_point = get_best_parameters( - experiment=experiment, generation_strategy=gs - ) - self.assertIsNone(best_point) - - with self.subTest("No completed trials"): - experiment = get_experiment_with_observations(observations=[]) - sobol_generator = get_sobol(search_space=experiment.search_space) - for _ in range(3): - trial = experiment.new_trial(generator_run=sobol_generator.gen(n=1)) - trial.run() - best_point = get_best_parameters( - experiment=experiment, generation_strategy=gs - ) - self.assertIsNone(best_point) - - experiment = get_experiment_with_observations( - observations=[[1], [2]], constrained=False - ) - with self.subTest("Working case"): - best_point = get_best_parameters( - experiment=experiment, generation_strategy=gs - ) - self.assertEqual(best_point, experiment.trials[1].arms[0].parameters) - - with self.subTest("Trial indices"): - best_point = get_best_parameters( - experiment=experiment, generation_strategy=gs, trial_indices=[0] - ) - self.assertEqual(best_point, experiment.trials[0].arms[0].parameters) + with patch(mock_function, return_value=(0, {"x": 1.0}, None)): + result = get_best_parameters(experiment, generation_strategy) + self.assertEqual(result, {"x": 1.0}) def test_get_benchmark_result_from_experiment_and_gs(self) -> None: problem = get_single_objective_benchmark_problem() diff --git a/ax/service/tests/test_best_point_utils.py b/ax/service/tests/test_best_point_utils.py index cc7306706a4..ff332532a06 100644 --- a/ax/service/tests/test_best_point_utils.py +++ b/ax/service/tests/test_best_point_utils.py @@ -10,16 +10,16 @@ from itertools import product from typing import Any from unittest import mock -from unittest.mock import patch, PropertyMock +from unittest.mock import Mock, patch, PropertyMock import numpy as np import pandas as pd import torch -from ax.adapter.cross_validation import AssessModelFitResult from ax.adapter.registry import Generators from ax.adapter.torch import TorchAdapter from ax.core.arm import Arm +from ax.core.data import Data from ax.core.experiment import Experiment from ax.core.generator_run import GeneratorRun from ax.core.metric import Metric @@ -54,6 +54,8 @@ get_branin_search_space, get_experiment_with_map_data, get_experiment_with_observations, + get_multi_objective_optimization_config, + get_sobol, ) from ax.utils.testing.mock import mock_botorch_optimize from pyre_extensions import assert_is_instance, none_throws @@ -302,99 +304,6 @@ def test_get_trace_by_arm_pull_from_data(self) -> None: self.assertEqual(set(result.columns), {"trial_index", "arm_name", "value"}) self.assertEqual(result["value"].tolist(), [0.0, 0.0, 2.0]) - @mock_botorch_optimize - def test_best_from_model_prediction(self) -> None: - exp = get_branin_experiment() - gs = choose_generation_strategy_legacy( - search_space=exp.search_space, - num_initialization_trials=3, - suggested_model_override=Generators.BOTORCH_MODULAR, - ) - - for _ in range(3): - generator_run = gs.gen_single_trial(experiment=exp, n=1) - trial = exp.new_trial(generator_run=generator_run) - trial.run().mark_completed() - exp.attach_data(exp.fetch_data()) - - generator_run = gs.gen_single_trial(experiment=exp, n=1) - trial = exp.new_trial(generator_run=generator_run) - trial.run().mark_completed() - - with patch.object( - TorchAdapter, - "model_best_point", - return_value=( - ( - exp.trials[0].arms[0], - ( - {"branin": 34.76260622783635}, - {"branin": {"branin": 0.00028306433439807734}}, - ), - ) - ), - ) as mock_model_best_point, self.assertLogs( - logger=best_point_logger, level="WARN" - ) as lg: - # Test bad model fit causes function to resort back to raw data - with patch( - "ax.service.utils.best_point.assess_model_fit", - return_value=AssessModelFitResult( - good_fit_metrics_to_fisher_score={}, - bad_fit_metrics_to_fisher_score={ - "branin": 0, - }, - ), - ): - self.assertIsNotNone( - get_best_parameters_from_model_predictions_with_trial_index( - experiment=exp, adapter=gs.adapter - ) - ) - self.assertTrue( - any("Model fit is poor" in warning for warning in lg.output), - msg=lg.output, - ) - mock_model_best_point.assert_not_called() - - # Test model best point is used when fit is good - with patch( - "ax.service.utils.best_point.assess_model_fit", - return_value=AssessModelFitResult( - good_fit_metrics_to_fisher_score={ - "branin": 0, - }, - bad_fit_metrics_to_fisher_score={}, - ), - ): - self.assertIsNotNone( - get_best_parameters_from_model_predictions_with_trial_index( - experiment=exp, adapter=gs.adapter - ) - ) - mock_model_best_point.assert_called() - - # Assert the non-mocked method works correctly as well - res = get_best_parameters_from_model_predictions_with_trial_index( - experiment=exp, adapter=gs.adapter - ) - trial_index, best_params, predict_arm = none_throws(res) - self.assertIsNotNone(best_params) - self.assertIsNotNone(trial_index) - self.assertIsNotNone(predict_arm) - # It works even when there are no predictions already stored on the - # GeneratorRun - for trial in exp.trials.values(): - trial.generator_run._best_arm_predictions = None - res = get_best_parameters_from_model_predictions_with_trial_index( - experiment=exp, adapter=gs.adapter - ) - trial_index, best_params_no_gr, predict_arm_no_gr = none_throws(res) - self.assertEqual(best_params, best_params_no_gr) - self.assertEqual(predict_arm, predict_arm_no_gr) - self.assertIsNotNone(trial_index) - self.assertIsNotNone(predict_arm) - def test_best_raw_objective_point(self) -> None: with self.subTest("Only early-stopped trials"): exp = get_experiment_with_map_data() @@ -892,19 +801,8 @@ def test_get_best_point_with_model_prediction( ax_client.get_next_trial() ax_client.complete_trial(i, raw_data={"y": i}) - # Mock with no bad fir metrics ensures that the model is used - # to extract the best point. - with patch( - f"{best_point_module}.assess_model_fit", - return_value=AssessModelFitResult( - good_fit_metrics_to_fisher_score={"y": 1}, - bad_fit_metrics_to_fisher_score={}, - ), - ) as mock_model_fit: - best_index, best_params, predictions = none_throws( - ax_client.get_best_trial() - ) - mock_model_fit.assert_called_once() + # Get the best trial using the model + best_index, best_params, predictions = none_throws(ax_client.get_best_trial()) self.assertEqual(best_index, idx) self.assertEqual(best_params, params) # We should get both mean & covariance predictions. @@ -919,6 +817,229 @@ def test_get_best_point_with_model_prediction( self.assertEqual(best_params, params) self.assertEqual(predictions, ({"y": mock.ANY}, {"y": {"y": mock.ANY}})) + @mock_botorch_optimize + def test_get_best_parameters_from_model_predictions_with_trial_index( + self, + ) -> None: + # Setup experiment + exp = get_branin_experiment() + gs = choose_generation_strategy_legacy( + search_space=exp.search_space, + num_initialization_trials=3, + suggested_model_override=Generators.BOTORCH_MODULAR, + ) + + # Add some trials with data + for _ in range(4): + generator_run = gs.gen_single_trial(experiment=exp, n=1) + trial = exp.new_trial(generator_run=generator_run) + trial.run().mark_completed() + exp.attach_data(exp.fetch_data()) + + # Test 1: No adapter (None) - should fall back to generator run + with self.subTest("No adapter - fallback to generator run"): + result = get_best_parameters_from_model_predictions_with_trial_index( + experiment=exp, adapter=None + ) + self.assertIsNotNone(result) + trial_index, params, _ = none_throws(result) + arm_params = result[1] + self.assertIn(arm_params, [v.parameters for v in exp.arms_by_name.values()]) + + # Test 2: Non-TorchAdapter - should fall back to generator run + # Then, the recommendation should be in-sample + with self.subTest("Non-TorchAdapter - fallback to generator run"): + non_torch_adapter = Mock() # Not a TorchAdapter + result = none_throws( + get_best_parameters_from_model_predictions_with_trial_index( + experiment=exp, adapter=non_torch_adapter + ) + ) + arm_params = result[1] + self.assertIn(arm_params, [v.parameters for v in exp.arms_by_name.values()]) + + # Test 3: TorchAdapter - should use adapter + with self.subTest("TorchAdapter - should use adapter"): + with patch.object( + TorchAdapter, + "model_best_point", + return_value=( + exp.trials[0].arms[0], + ({"branin": 1.0}, {"branin": {"branin": 0.1}}), + ), + ) as mock_model_best_point: + result = get_best_parameters_from_model_predictions_with_trial_index( + experiment=exp, adapter=gs.adapter + ) + mock_model_best_point.assert_called_once() + self.assertIsNotNone(result) + + with self.subTest( + "TorchAdapter with model_best_point returning None" + ), patch.object( + TorchAdapter, "model_best_point", return_value=None + ) as mock_model_best_point: + result = get_best_parameters_from_model_predictions_with_trial_index( + experiment=exp, adapter=gs.adapter + ) + mock_model_best_point.assert_called_once() + # Should still return a result from generator run fallback + self.assertIsNotNone(result) + + # Test 5: No generator run available - should return None + with self.subTest("No generator run available"): + # Create experiment with no generator runs + empty_exp = get_branin_experiment() + empty_exp.new_trial().run().mark_completed() + + result = get_best_parameters_from_model_predictions_with_trial_index( + experiment=empty_exp, adapter=None + ) + self.assertIsNone(result) + + # Test 6: Trial indices subset - should work with subset of data + with self.subTest("Trial indices subset"): + # Create experiment with specific data to test trial_indices functionality + test_exp = get_branin_experiment() + + # Add trials with specific objective values to control which is "best" + # Trial 0: objective = 5.0 (worse) + # Trial 1: objective = 3.0 (medium) + # Trial 2: objective = 1.0 (best overall) + values = [5.0, 3.0, 1.0, 2.0] + infer_best_idcs = [0, 0, 0, 2] + for i, (obj_value, infer_idx) in enumerate(zip(values, infer_best_idcs)): + trial = ( + test_exp.new_trial( + generator_run=GeneratorRun( + arms=[Arm(parameters={"x1": float(i), "x2": float(i)})] + ) + ) + .run() + .mark_completed() + ) + if i > 0: + trial.generator_run._best_arm_predictions = ( + test_exp.trials[infer_idx].arms[0], + None, # Don't need predictions for this test + ) + # Create specific data with controlled objective values + data = pd.DataFrame.from_records( + [ + { + "trial_index": trial.index, + "metric_name": "branin", + "arm_name": trial.arms[0].name, + "mean": obj_value, + "sem": 0.1, + } + ] + ) + + test_exp.attach_data(Data(df=data)) + + # First test: without trial_indices restriction - should select trial 2 + result_all = get_best_parameters_from_model_predictions_with_trial_index( + experiment=test_exp, adapter=None + ) + self.assertIsNotNone(result_all) + trial_index_all, _, _ = none_throws(result_all) + + # Second test: with trial_indices=[0, 1] - should select trial 0 or 1 + result_subset = get_best_parameters_from_model_predictions_with_trial_index( + experiment=test_exp, adapter=None, trial_indices=[0, 1] + ) + self.assertIsNotNone(result_subset) + trial_index_subset, _, _ = none_throws(result_subset) + + # Verify the trial_indices parameter affects the result + self.assertGreaterEqual( + trial_index_all, 2, "Best trial should be ≥2 when all trials considered" + ) + self.assertLessEqual( + trial_index_subset, + 1, + "Best trial should be ≤1 when restricted to [0,1]", + ) + + # Test 7: No optimization config - should raise ValueError + with self.subTest("No optimization config"): + exp_no_opt_config = get_branin_experiment(has_optimization_config=False) + with self.assertRaisesRegex( + ValueError, + "Cannot identify the best point without an optimization config", + ): + get_best_parameters_from_model_predictions_with_trial_index( + experiment=exp_no_opt_config, adapter=None, optimization_config=None + ) + + # Test 8: Multi-objective optimization config - should log warning + with self.subTest("Multi-objective optimization config"): + moo_config = get_multi_objective_optimization_config( + custom_metric=True, + outcome_constraint=False, + relative=False, + ) + with self.assertLogs(logger=best_point_logger, level="WARN") as lg: + result = get_best_parameters_from_model_predictions_with_trial_index( + experiment=exp, adapter=None, optimization_config=moo_config + ) + + # Should log deprecation warning for MOO + self.assertTrue( + any( + "get_best_parameters_from_model_predictions is deprecated for " + "multi-objective optimization configs" in log + for log in lg.output + ), + msg=lg.output, + ) + + # Test 9: BatchTrial with generator_run_structs - should handle BatchTrial case + with self.subTest("BatchTrial with generator_run_structs"): + # Create experiment and manually add BatchTrials with best_arm_predictions + batch_exp = get_branin_experiment() + gs = choose_generation_strategy_legacy( + batch_exp.search_space, + num_initialization_trials=2, + ) + for _ in range(4): + generator_run = gs.gen_single_trial(experiment=batch_exp, n=2) + trial = batch_exp.new_batch_trial(generator_run=generator_run) + trial.run().mark_completed() + batch_exp.attach_data(exp.fetch_data()) + + result = get_best_parameters_from_model_predictions_with_trial_index( + experiment=batch_exp, adapter=None + ) + + # Should return a result from BatchTrial's generator_run_structs + self.assertIsNotNone(result) + trial_index, params, _predictions = none_throws(result) + self.assertIsInstance(trial_index, int) + self.assertIsInstance(params, dict) + + with self.subTest("All constraints violated"): + experiment = get_experiment_with_observations( + observations=[[1, -1], [2, -1]], + constrained=True, + ) + best_point = get_best_parameters_from_model_predictions_with_trial_index( + experiment=experiment, adapter=gs.adapter + ) + self.assertIsNone(best_point) + + with self.subTest("No completed trials"): + experiment = get_experiment_with_observations(observations=[]) + sobol_generator = get_sobol(search_space=experiment.search_space) + for _ in range(3): + trial = experiment.new_trial(generator_run=sobol_generator.gen(n=1)) + trial.run() + best_point = get_best_parameters_from_model_predictions_with_trial_index( + experiment=experiment, adapter=gs.adapter + ) + self.assertIsNone(best_point) + def _repeat_elements(list_to_replicate: list[Any], n_repeats: int) -> pd.Series: return pd.Series([item for item in list_to_replicate for _ in range(n_repeats)]) diff --git a/ax/service/utils/best_point.py b/ax/service/utils/best_point.py index 1f127d8b800..3cec4d00dfa 100644 --- a/ax/service/utils/best_point.py +++ b/ax/service/utils/best_point.py @@ -8,6 +8,7 @@ from collections import OrderedDict from collections.abc import Iterable, Mapping +from copy import copy from logging import Logger import numpy as np @@ -19,11 +20,6 @@ predicted_pareto_frontier as predicted_pareto, ) from ax.adapter.base import Adapter -from ax.adapter.cross_validation import ( - assess_model_fit, - compute_diagnostics, - cross_validate, -) from ax.adapter.registry import Generators from ax.adapter.torch import TorchAdapter from ax.adapter.transforms.derelativize import Derelativize @@ -45,7 +41,6 @@ from ax.plot.pareto_utils import get_tensor_converter_adapter from ax.utils.common.logger import get_logger from botorch.utils.multi_objective.box_decompositions import DominatedPartitioning -from numpy import nan from numpy.typing import NDArray from pyre_extensions import assert_is_instance, none_throws @@ -241,8 +236,7 @@ def get_best_parameters_from_model_predictions_with_trial_index( and its ``predict`` method (if implemented). If ``adapter`` is not a ``TorchAdapter``, the best point is extracted from the (first) generator run of the latest trial. If the latest trial doesn't have a generator run, returns - None. If the model fit assessment returns bad fit for any of the metrics, this - will fall back to returning the best point based on raw observations. + None. TModelPredictArm is of the form: ({metric_name: mean}, {metric_name_1: {metric_name_2: cov_1_2}}) @@ -271,10 +265,51 @@ def get_best_parameters_from_model_predictions_with_trial_index( "multi-objective optimization configs. This method will return an " "arbitrary point on the pareto frontier." ) - gr = None data = experiment.lookup_data(trial_indices=trial_indices) - # Extract the latest GR from the experiment. - for _, trial in sorted(experiment.trials.items(), key=lambda x: x[0], reverse=True): + + # 1. If the adapter is a TorchAdapter, use it to get best points + if isinstance(adapter, TorchAdapter): + use_subset_of_data = trial_indices is not None and trial_indices != list( + experiment.trials.keys() + ) + if use_subset_of_data: + # get the adapter for the subset of data that we have access to for + # best point prediction + available_trial_observations, search_space = ( + adapter._process_and_transform_data(experiment=experiment, data=data) + ) + best_point_adapter = copy(adapter) + best_point_adapter._fit_if_implemented( + search_space=search_space, + experiment_data=available_trial_observations, + time_so_far=0.0, + ) + else: + best_point_adapter = adapter + + res = best_point_adapter.model_best_point() + # Check if that fails + if res is not None: + best_arm, best_arm_predictions = res + # Map the arm to the trial index of the first trial that contains it. + for trial_index, trial in experiment.trials.items(): + if best_arm in trial.arms: + return ( + trial_index, + none_throws(best_arm).parameters, + best_arm_predictions, + ) + + # 2. Otherwise, _extract_best_arm_from_gr (fallback) + # Extract the latest generator run from the experiment + gr = None + if trial_indices is not None: + filtered_trials = { + k: v for k, v in experiment.trials.items() if k in trial_indices + } + else: + filtered_trials = experiment.trials + for _, trial in sorted(filtered_trials.items(), key=lambda x: x[0], reverse=True): if isinstance(trial, Trial): gr = trial.generator_run elif isinstance(trial, BatchTrial): @@ -284,51 +319,9 @@ def get_best_parameters_from_model_predictions_with_trial_index( if gr is not None: break - if not isinstance(adapter, TorchAdapter): - if gr is None: - return None - return _extract_best_arm_from_gr(gr=gr, trials=experiment.trials) - - # Check to see if the adapter is worth using. - cv_results = cross_validate(model=adapter) - diagnostics = compute_diagnostics(result=cv_results) - assess_model_fit_results = assess_model_fit(diagnostics=diagnostics) - objective_name = optimization_config.objective.metric.name - # If model fit is bad use raw results - if objective_name in assess_model_fit_results.bad_fit_metrics_to_fisher_score: - logger.warning("Model fit is poor; falling back on raw data for best point.") - - if not _is_all_noiseless(df=data.df, metric_name=objective_name): - logger.warning( - "Model fit is poor and data on objective metric " - + f"{objective_name} is noisy; interpret best points " - + "results carefully." - ) - - return get_best_by_raw_objective_with_trial_index( - experiment=experiment, - optimization_config=optimization_config, - trial_indices=trial_indices, - ) - - res = adapter.model_best_point() - if res is None: - if gr is None: - return None - return _extract_best_arm_from_gr(gr=gr, trials=experiment.trials) - - best_arm, best_arm_predictions = res - - # Map the arm to the trial index of the first trial that contains it. - for trial_index, trial in experiment.trials.items(): - if best_arm in trial.arms: - return ( - trial_index, - none_throws(best_arm).parameters, - best_arm_predictions, - ) - - return None + if gr is None: + return None + return _extract_best_arm_from_gr(gr=gr, trials=experiment.trials) def get_best_by_raw_objective_with_trial_index( @@ -607,17 +600,6 @@ def tag_feasible_arms( ) -def _is_all_noiseless(df: pd.DataFrame, metric_name: str) -> bool: - """Noiseless is defined as SEM = 0 or SEM = NaN on a given metric (usually - the objective). - """ - - name_mask = df["metric_name"] == metric_name - df_metric_arms_sems = df[name_mask]["sem"] - - return ((df_metric_arms_sems == 0) | df_metric_arms_sems == nan).all() - - def get_values_of_outcomes_single_or_scalarized_objective( df_wide: pd.DataFrame, objective: Objective ) -> NDArray: