google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Tests for scores."""
17
import functools
18

19
from absl.testing import absltest
20
import jax
21
import jax.numpy as jnp
22
import jax.scipy as jscipy
23
from spaceopt import scores
24

25

26
def mocked_fantasize_y_values(key,
27
                              x_locations_for_p_min,
28
                              params,
29
                              x_obs,
30
                              y_obs,
31
                              y_fantasized_num_for_p_min=50):
32
  del params
33
  del x_obs
34
  del y_obs
35
  mu = jnp.zeros((x_locations_for_p_min.shape[0], 1))
36
  cov = jnp.eye(x_locations_for_p_min.shape[0])
37

38
  y_rand = jax.random.normal(key, (cov.shape[0], y_fantasized_num_for_p_min))
39
  chol = jscipy.linalg.cholesky(cov + jnp.eye(cov.shape[0]) * 1e-4, lower=True)
40
  fantasized_y = jnp.dot(chol, y_rand) + mu
41

42
  return fantasized_y
43

44

45
def draw_x_location_fn(key, search_space, budget):
46
  return jax.random.uniform(
47
      key,
48
      shape=(budget, 1),
49
      minval=search_space[:, 0],
50
      maxval=search_space[:, 1])
51

52

53
partial_fantasize_y_values = functools.partial(
54
    mocked_fantasize_y_values, params=None, x_obs=None, y_obs=None)
55

56

57
def draw_y_values_fn(key, x_locations, y_drawn_num):
58
  return partial_fantasize_y_values(
59
      key=key,
60
      x_locations_for_p_min=x_locations,
61
      y_fantasized_num_for_p_min=y_drawn_num)
62

63

64
class ScoresTest(absltest.TestCase):
65

66
  def setUp(self):
67

68
    super(ScoresTest, self).setUp()
69

70
    x_locations = jnp.linspace(-4., 4., 7)[:, None]
71
    self.x_obs = jnp.linspace(-1., 1., 2)[:, None]
72
    self.y_obs = jnp.array([[7.], [11.]])
73

74
    self.x_batch = jnp.linspace(1., 3., 4)[:, None]
75
    self.y_batch = jnp.array([[1., 5., 9.],
76
                              [2., 6., 10.],
77
                              [3., 7., 11.],
78
                              [4., 8., 12.]])
79
    key = jax.random.PRNGKey(0)
80
    self.utility_measure = scores.UtilityMeasure(
81
        incumbent=6.,
82
        x_locations_for_p_min=x_locations,
83
        params=None,
84
        fantasize_y_values_for_p_min=mocked_fantasize_y_values,
85
        y_fantasized_num_for_p_min=50,
86
        initial_entropy_key=key)
87

88
    self.key = jax.random.PRNGKey(0)
89
    self.budget = 10
90
    self.search_space = jnp.array([[-4., 4.]])
91

92
  def test_is_improvement_shape(self):
93
    """Test that the is_improvement output has the right shape."""
94
    is_improvement = self.utility_measure.is_improvement(self.y_batch)
95
    self.assertEqual(is_improvement.shape, (self.y_batch.shape[1],))
96

97
  def test_is_improvement_values(self):
98
    """Test that the is_improvement output has the right value."""
99
    is_improvement = self.utility_measure.is_improvement(self.y_batch)
100
    self.assertTrue((is_improvement == jnp.array([True, True, False])).all())
101

102
  def test_improvement_shape(self):
103
    """Test that the improvement output has the right shape."""
104
    improvement = self.utility_measure.improvement(self.y_batch)
105
    self.assertEqual(improvement.shape, (self.y_batch.shape[1],))
106

107
  def test_improvement_values(self):
108
    """Test that the improvement output has the right value."""
109
    improvement = self.utility_measure.improvement(self.y_batch)
110
    self.assertTrue((improvement == jnp.array([5., 1., 0.])).all())
111

112
  def test_information_gain_shape(self):
113
    """Test that the information_gain output has the right shape."""
114
    key = jax.random.PRNGKey(1)
115
    information_gain = self.utility_measure.information_gain(
116
        key=key,
117
        x_obs=self.x_obs,
118
        y_obs=self.y_obs,
119
        x_batch=self.x_batch,
120
        y_batch=self.y_batch)
121
    self.assertEqual(information_gain.shape, (self.y_batch.shape[1],))
122

123
  def test_score_values(self):
124
    """Test that the improvement-based score cannot be negative."""
125
    # pylint: disable=unused-argument
126
    def utility_is_imp_fn(key, x_batch, y_batch):
127
      return self.utility_measure.is_improvement(y_batch)
128

129
    def utility_imp_fn(key, x_batch, y_batch):
130
      return self.utility_measure.improvement(y_batch)
131
    # pylint: enable=unused-argument
132
    statistics_fns = [jnp.mean, jnp.median]
133
    key = jax.random.PRNGKey(1)
134

135
    mean_utility = scores.mean_utility(
136
        key,
137
        self.search_space,
138
        self.budget,
139
        utility_is_imp_fn,
140
        draw_y_values_fn,
141
        x_drawn_num=100,
142
        y_drawn_num=100)
143
    scores_dict = scores.scores(mean_utility, statistics_fns)
144

145
    for j in range(len(statistics_fns)):
146
      self.assertGreaterEqual(scores_dict[statistics_fns[j].__name__], 0.)
147

148
    mean_utility = scores.mean_utility(
149
        key,
150
        self.search_space,
151
        self.budget,
152
        utility_imp_fn,
153
        draw_y_values_fn,
154
        x_drawn_num=100,
155
        y_drawn_num=100)
156
    scores_dict = scores.scores(mean_utility, statistics_fns)
157

158
    for j in range(len(statistics_fns)):
159
      self.assertGreaterEqual(scores_dict[statistics_fns[j].__name__], 0.)
160

161

162
if __name__ == '__main__':
163
  absltest.main()
164