diff --git a/ax/modelbridge/generation_strategy.py b/ax/modelbridge/generation_strategy.py index db60abb9b11..22e4c717539 100644 --- a/ax/modelbridge/generation_strategy.py +++ b/ax/modelbridge/generation_strategy.py @@ -8,8 +8,8 @@ from __future__ import annotations +import warnings from collections.abc import Callable - from copy import deepcopy from functools import wraps from logging import Logger @@ -33,7 +33,7 @@ from ax.modelbridge.generation_node_input_constructors import InputConstructorPurpose from ax.modelbridge.model_spec import FactoryFunctionModelSpec from ax.modelbridge.transition_criterion import TrialBasedCriterion -from ax.utils.common.logger import _round_floats_for_logging, get_logger +from ax.utils.common.logger import get_logger from ax.utils.common.typeutils import assert_is_instance_list from pyre_extensions import none_throws @@ -277,54 +277,18 @@ def trials_as_df(self) -> pd.DataFrame | None: """Puts information on individual trials into a data frame for easy viewing. - For example for a GenerationStrategy composed of GenerationSteps: - Gen. Step | Models | Trial Index | Trial Status | Arm Parameterizations - [0] | [Sobol] | 0 | RUNNING | {"0_0":{"x":9.17...}} + THIS METHOD IS DEPRECATED AND WILL BE REMOVED IN A FUTURE RELEASE. + Please use `Experiment.to_df()` instead. """ - logger.info( - "Note that parameter values in dataframe are rounded to 2 decimal " - "points; the values in the dataframe are thus not the exact ones " - "suggested by Ax in trials." + warnings.warn( + "`GenerationStrategy.trials_as_df` is deprecated and will be removed in " + "a future release. Please use `Experiment.to_df()` instead.", + DeprecationWarning, + stacklevel=2, ) - if self._experiment is None or len(self.experiment.trials) == 0: + if self._experiment is None: return None - - step_or_node_col = ( - "Generation Nodes" if self.is_node_based else "Generation Step" - ) - records = [ - { - step_or_node_col: [ - ( - gr._generation_node_name - if gr.generator_run_type != "MANUAL" - else "MANUAL" - ) - for gr in trial.generator_runs - ], - "Generation Model(s)": [ - (gr._model_key if gr.generator_run_type != "MANUAL" else "MANUAL") - for gr in trial.generator_runs - ], - "Trial Index": trial_idx, - "Trial Status": trial.status.name, - "Arm Parameterizations": { - arm.name: _round_floats_for_logging(arm.parameters) - for arm in trial.arms - }, - } - for trial_idx, trial in self.experiment.trials.items() - ] - - return pd.DataFrame.from_records(records).reindex( - columns=[ - step_or_node_col, - "Generation Model(s)", - "Trial Index", - "Trial Status", - "Arm Parameterizations", - ] - ) + return self.experiment.to_df() @property def optimization_complete(self) -> bool: diff --git a/ax/modelbridge/tests/test_generation_strategy.py b/ax/modelbridge/tests/test_generation_strategy.py index 22d37e6e3db..d8efac00486 100644 --- a/ax/modelbridge/tests/test_generation_strategy.py +++ b/ax/modelbridge/tests/test_generation_strategy.py @@ -743,24 +743,14 @@ def test_trials_as_df(self) -> None: GenerationStep(model=Models.SOBOL, num_trials=3), ] ) - # No trials yet, so the DF will be None. - self.assertIsNone(sobol_generation_strategy.trials_as_df) - # Now the trial should appear in the DF. - trial = exp.new_trial(sobol_generation_strategy.gen(experiment=exp)) - trials_df = none_throws(sobol_generation_strategy.trials_as_df) - self.assertFalse(trials_df.empty) - self.assertEqual(trials_df.head()["Trial Status"][0], "CANDIDATE") - # Changes in trial status should be reflected in the DF. - trial._status = TrialStatus.RUNNING - trials_df = none_throws(sobol_generation_strategy.trials_as_df) - self.assertEqual(trials_df.head()["Trial Status"][0], "RUNNING") - # Check that rows are present for step 0 and 1 after moving to step 1 - for _i in range(3): - # attach necessary trials to fill up the Generation Strategy - trial = exp.new_trial(sobol_generation_strategy.gen(experiment=exp)) - trials_df = none_throws(sobol_generation_strategy.trials_as_df) - self.assertEqual(trials_df.head()["Generation Step"][0], ["GenerationStep_0"]) - self.assertEqual(trials_df.head()["Generation Step"][2], ["GenerationStep_1"]) + # No experiment attached to the GS, should be None. + with self.assertWarnsRegex(DeprecationWarning, "trials_as_df"): + self.assertIsNone(sobol_generation_strategy.trials_as_df) + # Experiment attached with a trial, should match Experiment.to_df(). + exp.new_trial(sobol_generation_strategy.gen(experiment=exp)) + with self.assertWarnsRegex(DeprecationWarning, "trials_as_df"): + trials_df = none_throws(sobol_generation_strategy.trials_as_df) + self.assertTrue(trials_df.equals(exp.to_df())) def test_max_parallelism_reached(self) -> None: exp = get_branin_experiment() @@ -1847,41 +1837,6 @@ def test_node_gs_with_auto_transitions_three_phase(self) -> None: trial = exp.new_batch_trial(generator_runs=gs_2._gen_with_multiple_nodes(exp)) self.assertEqual(trial.generator_runs[0]._generation_node_name, "sobol_4") - def test_trials_as_df_node_gs(self) -> None: - exp = get_branin_experiment() - gs = self.complex_multinode_per_trial_gs - arms_per_node = { - "sobol": 3, - "mbm": 1, - "sobol_2": 2, - "sobol_3": 3, - "sobol_4": 4, - } - gs.experiment = exp - self.assertIsNone(gs.trials_as_df) - # Now the trial should appear in the DF. - trial = exp.new_batch_trial( - generator_runs=gs._gen_with_multiple_nodes(exp, arms_per_node=arms_per_node) - ) - trials_df = none_throws(gs.trials_as_df) - self.assertFalse(trials_df.empty) - self.assertEqual(trials_df.head()["Trial Status"][0], "CANDIDATE") - self.assertEqual(trials_df.head()["Generation Model(s)"][0], ["Sobol"]) - # Changes in trial status should be reflected in the DF. - trial.run() - self.assertEqual( - none_throws(gs.trials_as_df).head()["Trial Status"][0], "RUNNING" - ) - # Add a new trial which will be generated from multiple nodes, and check that - # is properly reflected in the DF. - trial = exp.new_batch_trial( - generator_runs=gs._gen_with_multiple_nodes(exp, arms_per_node=arms_per_node) - ) - self.assertEqual( - none_throws(gs.trials_as_df).head()["Generation Nodes"][1], - ["mbm", "sobol_2", "sobol_3"], - ) - def test_gs_with_fixed_features_constructor(self) -> None: exp = get_branin_experiment() sobol_criterion = [