def test_expected_improvement(variance_scale: float,
num_samples_per_point: int, best: tf.Tensor,
rtol: float, atol: float) -> None:
variance_scale = tf.constant(variance_scale, tf.float64)
best = tf.cast(best, dtype=tf.float64)
x_range = tf.linspace(0.0, 1.0, 11)
x_range = tf.cast(x_range, dtype=tf.float64)
xs = tf.reshape(
tf.stack(tf.meshgrid(x_range, x_range, indexing="ij"), axis=-1),
(-1, 2))
kernel = tfp.math.psd_kernels.MaternFiveHalves(variance_scale,
length_scale=0.25)
model = GaussianProcess([branin], [kernel])
mean, variance = model.predict(xs)
samples = tfp.distributions.Normal(
mean, tf.sqrt(variance)).sample(num_samples_per_point)
samples_improvement = tf.where(samples < best, best - samples, 0)
ei_approx = tf.reduce_mean(samples_improvement, axis=0)
ei = expected_improvement(model, best, xs)
npt.assert_allclose(ei, ei_approx, rtol=rtol, atol=atol)
def test_expected_improvement(variance_scale: float,
num_samples_per_point: int, best: tf.Tensor,
rtol: float, atol: float) -> None:
variance_scale = tf.constant(variance_scale, tf.float64)
best = tf.cast(best, dtype=tf.float64)
x_range = tf.linspace(0.0, 1.0, 11)
x_range = tf.cast(x_range, dtype=tf.float64)
xs = tf.reshape(
tf.stack(tf.meshgrid(x_range, x_range, indexing="ij"), axis=-1),
(-1, 2))
class _Model(GaussianMarginal):
kernel = tfp.math.psd_kernels.MaternFiveHalves(variance_scale,
length_scale=0.25).apply
def predict(self,
query_points: TensorType) -> Tuple[TensorType, TensorType]:
return branin(query_points), self.kernel(query_points,
query_points)[:, None]
mean, variance = _Model().predict(xs)
samples = tfp.distributions.Normal(
mean, tf.sqrt(variance)).sample(num_samples_per_point)
samples_improvement = tf.where(samples < best, best - samples, 0)
ei_approx = tf.reduce_mean(samples_improvement, axis=0)
ei = expected_improvement(_Model(), best, xs)
npt.assert_allclose(ei, ei_approx, rtol=rtol, atol=atol)
def test_expected_improvement_builder_builds_expected_improvement(
query_at: tf.Tensor) -> None:
dataset = Dataset(tf.constant([[-2.], [-1.], [0.], [1.], [2.]]),
tf.zeros([5, 1]))
model = QuadraticWithUnitVariance()
builder = ExpectedImprovement()
acq_fn = builder.prepare_acquisition_function(dataset, model)
expected = expected_improvement(model, tf.constant([0.]), query_at)
npt.assert_array_almost_equal(acq_fn(query_at), expected)
def test_expected_improvement_builder_builds_expected_improvement_using_best_from_model() -> None:
dataset = Dataset(
tf.constant([[-2.0], [-1.0], [0.0], [1.0], [2.0]]),
tf.constant([[4.1], [0.9], [0.1], [1.1], [3.9]]),
)
model = QuadraticMeanAndRBFKernel()
acq_fn = ExpectedImprovement().prepare_acquisition_function(dataset, model)
xs = tf.linspace([-10.0], [10.0], 100)
expected = expected_improvement(model, tf.constant([0.0]), xs)
npt.assert_allclose(acq_fn(xs), expected)
def test_expected_improvement() -> None:
def _ei(x: tf.Tensor) -> tf.Tensor:
n = tfp.distributions.Normal(0, 1)
return -x * n.cdf(-x) + n.prob(-x)
query_at = tf.constant([[-2.0], [-1.5], [-1.0], [-0.5], [0.0], [0.5],
[1.0], [1.5], [2.0]])
actual = expected_improvement(QuadraticWithUnitVariance(),
tf.constant([0.]), query_at)
npt.assert_array_almost_equal(actual, _ei(query_at**2))
def test_reducers_on_ei(reducer):
m = 6
zero = tf.convert_to_tensor([0.0], dtype=tf.float64)
model = QuadraticWithUnitVariance()
acqs = [ExpectedImprovement().using("foo") for _ in range(m)]
acq = reducer.type_class(*acqs)
acq_fn = acq.prepare_acquisition_function({"foo": reducer.dataset}, {"foo": model})
individual_ei = [expected_improvement(model, zero, reducer.query_point) for _ in range(m)]
expected = reducer.raw_reduce_op(individual_ei)
desired = acq_fn(reducer.query_point)
np.testing.assert_array_almost_equal(desired, expected)
def test_expected_improvement_raises_for_invalid_batch_size(at: TensorType) -> None:
ei = expected_improvement(QuadraticMeanAndRBFKernel(), tf.constant([1.0]))
with pytest.raises(TF_DEBUGGING_ERROR_TYPES):
ei(at)