From 2938d30ad0ff45b3bc7f96b0b85273e9bd48d992 Mon Sep 17 00:00:00 2001 From: Chansoo Lee Date: Tue, 28 Nov 2023 14:09:42 -0800 Subject: [PATCH] Add support for conditionals in `TrialToArrayConverter`. PiperOrigin-RevId: 586099252 --- .../_src/pyvizier/shared/parameter_config.py | 64 ++++++++++++++----- .../pyvizier/shared/parameter_config_test.py | 20 ++++++ vizier/pyvizier/converters/core.py | 44 +++++-------- vizier/pyvizier/converters/core_test.py | 16 +++++ 4 files changed, 100 insertions(+), 44 deletions(-) diff --git a/vizier/_src/pyvizier/shared/parameter_config.py b/vizier/_src/pyvizier/shared/parameter_config.py index c3ff14c10..aa272e2ac 100644 --- a/vizier/_src/pyvizier/shared/parameter_config.py +++ b/vizier/_src/pyvizier/shared/parameter_config.py @@ -17,7 +17,7 @@ """ParameterConfig wraps ParameterConfig and ParameterSpec protos.""" import collections -from typing import Collection, Set as AbstractSet, Sized +from typing import Iterable, Set as AbstractSet, Sized import copy import enum import json @@ -410,7 +410,7 @@ def child_parameter_configs(self) -> List['ParameterConfig']: def subspaces( self, - ) -> Collection[Tuple[ParameterValueTypes, 'SearchSpace']]: + ) -> Iterable[Tuple[ParameterValueTypes, 'SearchSpace']]: return self._children.items() # TODO: TO BE DEPRECATED. @@ -724,24 +724,26 @@ def subspace(self, value: ParameterValueTypes) -> 'SearchSpace': return self._children[value] -ParameterConfigOrConfigs = Union[ParameterConfig, Collection[ParameterConfig]] - - @attr.define(init=False) -class ParameterConfigSelector(Sized): +class ParameterConfigSelector(Iterable[ParameterConfig], Sized): """Holds a reference to ParameterConfigs.""" # Selected configs. - _selected: tuple[ParameterConfig] = attr.field(init=True) + _selected: tuple[ParameterConfig] = attr.field(init=True, converter=tuple) + + def __iter__(self) -> Iterator[ParameterConfig]: + return iter(self._selected) def __len__(self) -> int: return len(self._selected) - def __init__(self, selected: ParameterConfigOrConfigs): - if isinstance(selected, Collection): - self.__attrs_init__(tuple(selected)) - else: + def __init__( + self, selected: Union[ParameterConfig, Iterable[ParameterConfig]], / + ): + if isinstance(selected, ParameterConfig): self.__attrs_init__(tuple([selected])) + else: + self.__attrs_init__(tuple(selected)) def select_values( self, values: MonotypeParameterSequence @@ -762,13 +764,29 @@ def select_values( spaces.append(config.subspace(value)) return SearchSpaceSelector(spaces) + def merge(self) -> 'ParameterConfigSelector': + """Merge by taking the union of the parameter configs with the same name. + + Returns: + The returned ParameterConfigSelector does not contain parameters with + duplicate names. Their feasible set (either as a range or discrete set) is + the union of all feasible sets under the same parameter name. + """ + merged_configs = {} + for parameter_config in self: + name = parameter_config.name # Alias + existing_config = merged_configs.setdefault(name, parameter_config) + merged_configs[name] = ParameterConfig.merge( + existing_config, parameter_config + ) + return ParameterConfigSelector(merged_configs.values()) + class InvalidParameterError(ValueError): """Error thrown when parameter values are invalid.""" ################### Main Classes ################### -SearchSpaceOrSpaces = Union['SearchSpace', Collection['SearchSpace']] @attr.define(init=False) @@ -783,11 +801,13 @@ class SearchSpaceSelector: def __len__(self) -> int: return len(self._selected) - def __init__(self, selected: SearchSpaceOrSpaces): - if isinstance(selected, Collection): - self.__attrs_init__(tuple(selected)) - else: + def __init__( + self, selected: Union['SearchSpace', Iterable['SearchSpace']], / + ): + if isinstance(selected, SearchSpace): self.__attrs_init__(tuple([selected])) + else: + self.__attrs_init__(tuple(selected)) def add_float_param( self, @@ -1239,7 +1259,7 @@ def parse_multi_dimensional_parameter_name( # TODO: Add def extend(space: SearchSpace) def _add_parameters( - self, parameters: List[ParameterConfig] + self, parameters: Iterable[ParameterConfig] ) -> ParameterConfigSelector: """Adds deepcopy of the ParameterConfigs. @@ -1249,6 +1269,7 @@ def _add_parameters( Returns: A list of SearchSpaceSelectors, one for each parameters added. """ + parameters = list(parameters) logging.info( 'Adding child parameters %s to %s subspaces ', set(p.name for p in parameters), @@ -1262,6 +1283,15 @@ def _add_parameters( return ParameterConfigSelector(added) + def select_all(self) -> ParameterConfigSelector: + """Select all parameters at all levels.""" + all_parameter_configs = [] + for space in self._selected: + for top_level_config in space.parameters: + all_parameter_configs.extend(list(top_level_config.traverse())) + + return ParameterConfigSelector(all_parameter_configs) + @attr.define(frozen=False, init=True, slots=True, kw_only=True) class SearchSpace: diff --git a/vizier/_src/pyvizier/shared/parameter_config_test.py b/vizier/_src/pyvizier/shared/parameter_config_test.py index 1d6833d90..0733f3dde 100644 --- a/vizier/_src/pyvizier/shared/parameter_config_test.py +++ b/vizier/_src/pyvizier/shared/parameter_config_test.py @@ -18,8 +18,10 @@ from typing import Any +from absl import logging from vizier._src.pyvizier.shared import parameter_config as pc from vizier._src.pyvizier.shared import trial +from vizier.testing import test_studies from absl.testing import absltest from absl.testing import parameterized @@ -557,6 +559,24 @@ def testValidateCategoricalInput(self): root.add_categorical_param('categorical', ['3.2', '2', 5]) +class FlattenAndMergeTest(absltest.TestCase): + + def testFlattenAndMerge(self): + space = test_studies.conditional_automl_space() + parameters = space.root.select_all().merge() + logging.info('Merged: %s', parameters) + self.assertCountEqual( + [p.name for p in parameters], + [ + 'model_type', + 'learning_rate', + 'optimizer_type', + 'use_special_logic', + 'special_logic_parameter', + ], + ) + + class SearchSpaceContainsTest(absltest.TestCase): def _space(self): diff --git a/vizier/pyvizier/converters/core.py b/vizier/pyvizier/converters/core.py index 715986e91..0701cce11 100644 --- a/vizier/pyvizier/converters/core.py +++ b/vizier/pyvizier/converters/core.py @@ -20,7 +20,6 @@ import copy import dataclasses import enum -import itertools from typing import Any, Callable, Collection, Dict, Iterator, List, Mapping, Optional, Sequence, Tuple, Type, Union from absl import logging @@ -1147,29 +1146,17 @@ def from_study_configs( Returns: `DefaultTrialConverter`. """ - # Cache ParameterConfigs. - # Traverse through all parameter configs and merge the same-named ones. - parameter_configs: Dict[str, pyvizier.ParameterConfig] = dict() - for study_config in study_configs: - all_parameter_configs = itertools.chain.from_iterable( - [ - top_level_config.traverse() - for top_level_config in study_config.search_space.parameters - ] - ) - for parameter_config in all_parameter_configs: - name = parameter_config.name # Alias - existing_config = parameter_configs.get(name, None) - if existing_config is None: - parameter_configs[name] = parameter_config - else: - parameter_configs[name] = pyvizier.ParameterConfig.merge( - existing_config, parameter_config - ) + # Merge parameter configs by name. + merged_configs = list( + pyvizier.SearchSpaceSelector([sc.search_space for sc in study_configs]) + .select_all() + .merge() + ) - parameter_converters = [] - for pc in parameter_configs.values(): - parameter_converters.append(DefaultModelInputConverter(pc)) + merged_configs = {pc.name: pc for pc in merged_configs} + parameter_converters = [ + DefaultModelInputConverter(pc) for pc in merged_configs.values() + ] # Append study id feature if configured to do so. if use_study_id_feature: @@ -1188,17 +1175,17 @@ def from_study_configs( 'had study id configured.' ) use_study_id_feature = False - elif STUDY_ID_FIELD in parameter_configs: + elif STUDY_ID_FIELD in merged_configs: raise ValueError( 'Dataset name conflicts with a ParameterConfig ' - 'that already exists: {}'.format(parameter_configs[STUDY_ID_FIELD]) + 'that already exists: {}'.format(merged_configs[STUDY_ID_FIELD]) ) # Create new parameter config. parameter_config = pyvizier.ParameterConfig.factory( STUDY_ID_FIELD, feasible_values=list(study_ids) ) - parameter_configs[STUDY_ID_FIELD] = parameter_config + merged_configs[STUDY_ID_FIELD] = parameter_config logging.info('Created a new ParameterConfig %s', parameter_config) # Create converter. @@ -1309,7 +1296,10 @@ def create_output_converter(metric): sc = study_config # alias, to keep pylint quiet in the next line. converter = DefaultTrialConverter( - [create_input_converter(p) for p in sc.search_space.parameters], + [ + create_input_converter(p) + for p in sc.search_space.root.select_all().merge() + ], [create_output_converter(m) for m in sc.metric_information], ) return cls(converter) diff --git a/vizier/pyvizier/converters/core_test.py b/vizier/pyvizier/converters/core_test.py index 833efb17b..2a6347e5a 100644 --- a/vizier/pyvizier/converters/core_test.py +++ b/vizier/pyvizier/converters/core_test.py @@ -31,6 +31,22 @@ Trial = pyvizier.Trial +class TrialToArrayConverterConditionalSpaceTest(parameterized.TestCase): + + def test_automl_study(self): + space = test_studies.conditional_automl_space() + converter = core.TrialToArrayConverter.from_study_config( + pyvizier.ProblemStatement(search_space=space) + ) + features = converter.to_features([pyvizier.Trial()]) + np.testing.assert_equal( + features, + np.array( + [[0.0, 0.0, 1.0, np.nan, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, np.nan]] + ), + ) + + class TrialToArrayConverterTest(parameterized.TestCase): """Test TrialToArrayConverter class."""