def testEmpiricalBayesThompsonSamplerGen(self):
generator = EmpiricalBayesThompsonSampler(min_weight=0.0)
generator.fit(
Xs=self.Xs,
Ys=self.Ys,
Yvars=self.Yvars,
parameter_values=self.parameter_values,
outcome_names=self.outcome_names,
)
with patch(
"ax.models.discrete.eb_thompson.EmpiricalBayesThompsonSampler",
return_value=(
[[4, 4], [3, 3], [2, 2], [1, 1]],
[
2.67,
0,
0.25,
0.07,
],
),
):
arms, weights, _ = generator.gen(
n=5,
parameter_values=self.parameter_values,
objective_weights=np.array([1, 0]),
)
self.assertEqual(arms, [[4, 4], [3, 3], [2, 2], [1, 1]])
for weight, expected_weight in zip(
weights, [4 * i for i in [0.66, 0.25, 0.07, 0.02]]):
self.assertAlmostEqual(weight, expected_weight, 1)
def testEmpiricalBayesThompsonSamplerValidation(self):
generator = EmpiricalBayesThompsonSampler(min_weight=0.01)
with self.assertRaises(ValueError):
generator.fit(
Xs=[[[1, 1], [2, 2], [3, 3], [4, 4]], [[1, 1], [2, 2], [4,
4]]],
Ys=self.Ys,
Yvars=self.Yvars,
parameter_values=self.parameter_values,
outcome_names=self.outcome_names,
)
def testEmpiricalBayesThompsonSamplerWarning(self):
generator = EmpiricalBayesThompsonSampler(min_weight=0.0)
generator.fit(
Xs=[x[:-1] for x in self.Xs],
Ys=[y[:-1] for y in self.Ys],
Yvars=[y[:-1] for y in self.Yvars],
parameter_values=self.parameter_values,
outcome_names=self.outcome_names,
)
arms, weights = generator.gen(
n=5,
parameter_values=self.parameter_values,
objective_weights=np.array([1, 0]),
)
self.assertEqual(arms, [[3, 3], [2, 2], [1, 1]])
for weight, expected_weight in zip(weights, [0.74, 0.21, 0.05]):
self.assertAlmostEqual(weight, expected_weight, 1)
def testEmpiricalBayesThompsonSamplerGen(self):
generator = EmpiricalBayesThompsonSampler(min_weight=0.0)
generator.fit(
Xs=self.Xs,
Ys=self.Ys,
Yvars=self.Yvars,
parameter_values=self.parameter_values,
outcome_names=self.outcome_names,
)
arms, weights = generator.gen(
n=5,
parameter_values=self.parameter_values,
objective_weights=np.array([1, 0]),
)
self.assertEqual(arms, [[4, 4], [3, 3], [2, 2], [1, 1]])
for weight, expected_weight in zip(weights, [0.66, 0.25, 0.07, 0.02]):
self.assertAlmostEqual(weight, expected_weight, 1)
def testEmpiricalBayesThompsonSamplerFit(self):
generator = EmpiricalBayesThompsonSampler(min_weight=0.0)
generator.fit(
Xs=self.Xs,
Ys=self.Ys,
Yvars=self.Yvars,
parameter_values=self.parameter_values,
outcome_names=self.outcome_names,
)
self.assertEqual(generator.X, self.Xs[0])
self.assertTrue(
np.allclose(
np.array(generator.Ys),
np.array([[1.3, 2.1, 2.9, 3.7], [0.25, 0.25, 0.25, 0.25]]),
))
self.assertTrue(
np.allclose(
np.array(generator.Yvars),
np.array([[1.03, 0.87, 0.87, 1.03],
[0.375, 0.375, 0.375, 1.375]]),
))
def get_empirical_bayes_thompson(
experiment: Experiment,
data: Data,
search_space: Optional[SearchSpace] = None,
num_samples: int = 10000,
min_weight: Optional[float] = None,
uniform_weights: bool = False,
) -> DiscreteModelBridge:
"""Instantiates an empirical Bayes / Thompson sampling model."""
if data.df.empty: # pragma: no cover
raise ValueError(
"Empirical Bayes Thompson sampler requires non-empty data.")
model = EmpiricalBayesThompsonSampler(num_samples=num_samples,
min_weight=min_weight,
uniform_weights=uniform_weights)
return DiscreteModelBridge(
experiment=experiment,
search_space=search_space
if search_space is not None else experiment.search_space,
data=data,
model=model,
transforms=TS_trans,
)
def testEmpiricalBayesThompsonSamplerPredict(self):
generator = EmpiricalBayesThompsonSampler(min_weight=0.0)
generator.fit(
Xs=self.Xs,
Ys=self.Ys,
Yvars=self.Yvars,
parameter_values=self.parameter_values,
outcome_names=self.outcome_names,
)
f, cov = generator.predict([[1, 1], [3, 3]])
self.assertTrue(np.allclose(f, np.array([[1.3, 0.25], [2.9, 0.25]])))
# first element of cov is the cov matrix for the first prediction
# the element at 0,0 is for the first outcome
# the element at 1,1 is for the second outcome
self.assertTrue(
np.allclose(
cov,
np.array([[[1.03, 0.0], [0.0, 0.375]],
[[0.87, 0.0], [0.0, 0.375]]]),
))
with self.assertRaises(ValueError):
generator.predict([[1, 2]])
