diff --git a/vizier/_src/algorithms/designers/gp_bandit.py b/vizier/_src/algorithms/designers/gp_bandit.py index 9f97e49bc..a59cdead0 100644 --- a/vizier/_src/algorithms/designers/gp_bandit.py +++ b/vizier/_src/algorithms/designers/gp_bandit.py @@ -50,6 +50,13 @@ from vizier.pyvizier.converters import padding from vizier.utils import profiler + +# The maximum number of feasible values to use for the trust region. If a +# discrete or integer parameter has more than this many feasible values, it is +# considered continuous and always included in the trust region computation. +_MAX_NUM_FEASIBLE_VALUES_FOR_TRUST_REGION = 1000 + + default_acquisition_optimizer_factory = vb.VectorizedOptimizerFactory( strategy_factory=es.VectorizedEagleStrategyFactory( eagle_config=es.EagleStrategyConfig() @@ -246,7 +253,9 @@ def acq_fn_factory(data: types.ModelData) -> acq_lib.AcquisitionFunction: scoring_fn = self._scoring_function_factory( empty_data, predictive, - self._converter.continuous_feasible_values, + self._converter.continuous_feasible_values( + max_num_feasible_values=_MAX_NUM_FEASIBLE_VALUES_FOR_TRUST_REGION + ), self._use_trust_region, ) if ( @@ -532,7 +541,9 @@ def suggest(self, count: int = 1) -> Sequence[vz.TrialSuggestion]: scoring_fn = self._scoring_function_factory( data, gp, - self._converter.continuous_feasible_values, + self._converter.continuous_feasible_values( + max_num_feasible_values=_MAX_NUM_FEASIBLE_VALUES_FOR_TRUST_REGION + ), self._use_trust_region, ) logging.info('Optimizing acquisition: %s', scoring_fn) diff --git a/vizier/_src/algorithms/designers/gp_ucb_pe.py b/vizier/_src/algorithms/designers/gp_ucb_pe.py index 965955078..c39c971d5 100644 --- a/vizier/_src/algorithms/designers/gp_ucb_pe.py +++ b/vizier/_src/algorithms/designers/gp_ucb_pe.py @@ -54,6 +54,12 @@ tfd = tfp.distributions +# The maximum number of feasible values to use for the trust region. If a +# discrete or integer parameter has more than this many feasible values, it is +# considered continuous and always included in the trust region computation. +_MAX_NUM_FEASIBLE_VALUES_FOR_TRUST_REGION = 1000 + + class MultimetricPromisingRegionPenaltyType(enum.Enum): """The type of penalty to apply to the points outside the promising region. @@ -1390,7 +1396,9 @@ def suggest( ) tr = acquisitions.TrustRegion( trusted=tr_features, - continuous_feasible_values=self._converter.continuous_feasible_values, + continuous_feasible_values=self._converter.continuous_feasible_values( + max_num_feasible_values=_MAX_NUM_FEASIBLE_VALUES_FOR_TRUST_REGION + ), ) acquisition_problem = copy.deepcopy(self._problem) diff --git a/vizier/_src/algorithms/designers/gp_ucb_pe_test.py b/vizier/_src/algorithms/designers/gp_ucb_pe_test.py index b552520c9..a66047ded 100644 --- a/vizier/_src/algorithms/designers/gp_ucb_pe_test.py +++ b/vizier/_src/algorithms/designers/gp_ucb_pe_test.py @@ -615,6 +615,10 @@ def test_discrete_parameters_are_explored( root.add_discrete_param('discrete_double_1', [-3.0, 7.0]) root.add_discrete_param('discrete_double_single_feasible_value', [0.37]) root.add_float_param('double', min_value=-5.0, max_value=5.0) + root.add_int_param('integer_with_many_feasible_values', -1, 6 * 10**7) + root.add_discrete_param( + 'discrete_double_many_feasible_values', np.linspace(-10.0, 10.0, 1000) + ) problem = vz.ProblemStatement(space) problem.metric_information.append( vz.MetricInformation( diff --git a/vizier/pyvizier/converters/jnp_converters.py b/vizier/pyvizier/converters/jnp_converters.py index b3c395064..716665c2e 100644 --- a/vizier/pyvizier/converters/jnp_converters.py +++ b/vizier/pyvizier/converters/jnp_converters.py @@ -246,29 +246,52 @@ def output_specs(self): # TODO: Add back pytype def metric_specs(self) -> Sequence[vz.MetricInformation]: return self._impl.metric_specs - @property - def continuous_feasible_values(self) -> list[jax.Array]: + def continuous_feasible_values( + self, max_num_feasible_values: int | None = None + ) -> list[jax.Array]: """Returns a list of feasible values for each continuified parameter. The list is ordered the same as the parameters in the search space. An empty array means that the parameter is continuous and all values within its range are feasible. The returned feasible values are in the scaled space. - Categorical parameters are ignored. + Categorical parameters are ignored. If `max_num_feasible_values` is + specified, discrete or integer parameters with more than + `max_num_feasible_values` feasible values are considered continuous and + empty arrays are returned for them. + + Args: + max_num_feasible_values: If specified, discrete or integer parameters with + more than this many feasible values are considered continuous and empty + arrays are returned for them. + + Returns: + A list of feasible values for each continuified parameter. The list is + ordered the same as the parameters in the search space. """ continuous_feasible_values = [] for param in self._problem.search_space.parameters: if param.type in [vz.ParameterType.DISCRETE, vz.ParameterType.INTEGER]: - converted_feasible_values = self.to_features([ - vz.TrialSuggestion(parameters={param.name: value}) - for value in param.feasible_values - ]).continuous.unpad() - continuous_feasible_values.append( - # `self.to_features` returns NaNs for missing parameters (we call it - # with one specific parameter instead of all parameters in the - # search space), so we remove them here. - converted_feasible_values[~jnp.isnan(converted_feasible_values)] - ) + if param.type == vz.ParameterType.INTEGER: + num_feasible_values = param.bounds[1] - param.bounds[0] + 1 + else: + num_feasible_values = len(param.feasible_values) + if ( + max_num_feasible_values is not None + and num_feasible_values > max_num_feasible_values + ): + continuous_feasible_values.append(jnp.asarray([])) + else: + converted_feasible_values = self.to_features([ + vz.TrialSuggestion(parameters={param.name: value}) + for value in param.feasible_values + ]).continuous.unpad() + continuous_feasible_values.append( + # `self.to_features` returns NaNs for missing parameters (we call + # it with one specific parameter instead of all parameters in the + # search space), so we remove them here. + converted_feasible_values[~jnp.isnan(converted_feasible_values)] + ) elif param.type == vz.ParameterType.DOUBLE: continuous_feasible_values.append(jnp.asarray([])) elif param.type == vz.ParameterType.CUSTOM: diff --git a/vizier/pyvizier/converters/jnp_converters_test.py b/vizier/pyvizier/converters/jnp_converters_test.py index da08dc83c..5baca4d68 100644 --- a/vizier/pyvizier/converters/jnp_converters_test.py +++ b/vizier/pyvizier/converters/jnp_converters_test.py @@ -110,7 +110,7 @@ def setUp(self): self.maxDiff = None def test_continuous_feasible_values(self): - """Tests the `continuous_feasible_values` property. + """Tests the `continuous_feasible_values` method. The returned feasible values are expected to be in the scaled space, and categorical parameters are ignored. @@ -125,6 +125,7 @@ def test_continuous_feasible_values(self): ) root.add_int_param('integer', -2, 2) root.add_float_param('float', -3, 4) + root.add_int_param('integer_with_many_feasible_values', -1, 6 * 10**7) converter = jnpc.TrialToModelInputConverter.from_problem( vz.ProblemStatement( search_space=space, @@ -135,19 +136,22 @@ def test_continuous_feasible_values(self): ], ) ) - self.assertLen(converter.continuous_feasible_values, 5) + continuous_feasible_values = converter.continuous_feasible_values( + max_num_feasible_values=1000 + ) + self.assertLen(continuous_feasible_values, 6) self.assertSequenceAlmostEqual( - converter.continuous_feasible_values[0], + continuous_feasible_values[0], # The original feasible values are linearly scaled to [0, 1]. (np.asarray([-0.4, -0.3, 0.4, 1.1, 1.15]) + 0.4) / 1.55, ) self.assertSequenceAlmostEqual( # The original feasible values [5, 9] are linearly scaled to [0, 1]. - converter.continuous_feasible_values[1], + continuous_feasible_values[1], np.asarray([0.0, 1.0]), ) self.assertSequenceAlmostEqual( - converter.continuous_feasible_values[2], + continuous_feasible_values[2], # The original feasible values [1e-5, 1e-2, 1e-1] are log-then-linearly # scaled to [0, 1]. ( @@ -156,13 +160,15 @@ def test_continuous_feasible_values(self): ), ) self.assertSequenceAlmostEqual( - converter.continuous_feasible_values[3], + continuous_feasible_values[3], # The integer parameter with range [-2, 2] are internally scaled to # evenly spaced values in [0, 1]. np.asarray([0.0, 0.25, 0.5, 0.75, 1.0]), ) # Feasible values for float parameters are empty. - self.assertEmpty(converter.continuous_feasible_values[4]) + self.assertEmpty(continuous_feasible_values[4]) + # Feasible values for int parameters with many feasible values are empty. + self.assertEmpty(continuous_feasible_values[5]) def test_back_to_back_conversion(self): converter = jnpc.TrialToContinuousAndCategoricalConverter.from_study_config(