def _dim_two_gp(mean_shift: tuple[float, float] = (0.0, 0.0)) -> GaussianProcess: matern52 = tfp.math.psd_kernels.MaternFiveHalves( amplitude=tf.cast(2.3, tf.float64), length_scale=tf.cast(0.5, tf.float64)) return GaussianProcess( [ lambda x: mean_shift[0] + branin(x), lambda x: mean_shift[1] + quadratic(x) ], [matern52, rbf()], )
def _mo_test_model( num_obj: int, *kernel_amplitudes: float | TensorType | None) -> GaussianProcess: means = [ quadratic, lambda x: tf.reduce_sum(x, axis=-1, keepdims=True), quadratic ] kernels = [ tfp.math.psd_kernels.ExponentiatedQuadratic(k_amp) for k_amp in kernel_amplitudes ] return GaussianProcess(means[:num_obj], kernels[:num_obj])
def test_expected_constrained_improvement_is_less_for_constrained_points() -> None: class _Constraint(AcquisitionFunctionBuilder): def prepare_acquisition_function( self, datasets: Mapping[str, Dataset], models: Mapping[str, ProbabilisticModel] ) -> AcquisitionFunction: return lambda x: tf.cast(x >= 0, x.dtype) def two_global_minima(x: tf.Tensor) -> tf.Tensor: return x ** 4 / 4 - x ** 2 / 2 initial_query_points = tf.constant([[-2.0], [0.0], [1.2]]) data = {"foo": Dataset(initial_query_points, two_global_minima(initial_query_points))} models_ = {"foo": GaussianProcess([two_global_minima], [rbf()])} eci = ExpectedConstrainedImprovement("foo", _Constraint()).prepare_acquisition_function( data, models_ ) npt.assert_array_less(eci(tf.constant([-1.0])), eci(tf.constant([1.0])))