def test_get_model_from_generator_run(self):
"""Tests that it is possible to restore a model from a generator run it
produced, if `Models` registry was used.
"""
exp = get_branin_experiment()
initial_sobol = Models.SOBOL(experiment=exp, seed=239)
gr = initial_sobol.gen(n=1)
# Restore the model as it was before generation.
sobol = get_model_from_generator_run(generator_run=gr,
experiment=exp,
data=exp.fetch_data())
self.assertEqual(sobol.model.init_position, 0)
self.assertEqual(sobol.model.seed, 239)
# Restore the model as it was after generation (to resume generation).
sobol_after_gen = get_model_from_generator_run(generator_run=gr,
experiment=exp,
data=exp.fetch_data(),
after_gen=True)
self.assertEqual(sobol_after_gen.model.init_position, 1)
self.assertEqual(sobol_after_gen.model.seed, 239)
self.assertEqual(
initial_sobol.gen(n=1).arms,
sobol_after_gen.gen(n=1).arms)
exp.new_trial(generator_run=gr)
# Check restoration of GPEI, to ensure proper restoration of callable kwargs
gpei = Models.GPEI(experiment=exp, data=get_branin_data())
# Punch GPEI model + bridge kwargs into the Sobol generator run, to avoid
# a slow call to `gpei.gen`.
gr._model_key = "GPEI"
gr._model_kwargs = gpei._model_kwargs
gr._bridge_kwargs = gpei._bridge_kwargs
gpei_restored = get_model_from_generator_run(gr,
experiment=exp,
data=get_branin_data())
for key in gpei.__dict__:
self.assertIn(key, gpei_restored.__dict__)
original, restored = gpei.__dict__[key], gpei_restored.__dict__[
key]
# Fit times are set in instantiation so not same and model compared below
if key in ["fit_time", "fit_time_since_gen", "model"]:
continue # Fit times are set in instantiation so won't be same
if isinstance(original, OrderedDict) and isinstance(
restored, OrderedDict):
original, restored = list(original.keys()), list(
restored.keys())
if isinstance(original, Model) and isinstance(restored, Model):
continue # Model equality is tough to compare.
self.assertEqual(original, restored)
for key in gpei.model.__dict__:
self.assertIn(key, gpei_restored.model.__dict__)
original, restored = (
gpei.model.__dict__[key],
gpei_restored.model.__dict__[key],
)
# Botorch model equality is tough to compare and training data
# is unnecessary to compare, because data passed to model was the same
if key in ["model", "warm_start_refitting", "Xs", "Ys"]:
continue
self.assertEqual(original, restored)
def test_use_update(self, mock_lookup_data, mock_update):
exp = get_branin_experiment()
sobol_gs_with_update = GenerationStrategy(steps=[
GenerationStep(model=Models.SOBOL, num_trials=-1, use_update=True)
])
sobol_gs_with_update._experiment = exp
self.assertEqual(
sobol_gs_with_update._find_trials_completed_since_last_gen(),
set(),
)
with self.assertRaises(NotImplementedError):
# `BraninMetric` is available while running by default, which should
# raise an error when use with `use_update=True` on a generation step, as we
# have not yet properly addressed that edge case (for lack of use case).
sobol_gs_with_update.gen(experiment=exp)
core_stubs_module = get_branin_experiment.__module__
with patch(
f"{core_stubs_module}.BraninMetric.is_available_while_running",
return_value=False,
):
# Try without passing data (GS looks up data on experiment).
trial = exp.new_trial(generator_run=sobol_gs_with_update.gen(
experiment=exp))
mock_update.assert_not_called()
trial._status = TrialStatus.COMPLETED
for i in range(3):
gr = sobol_gs_with_update.gen(experiment=exp)
self.assertEqual(
mock_lookup_data.call_args[1].get("trial_indices"), {i})
trial = exp.new_trial(generator_run=gr)
trial._status = TrialStatus.COMPLETED
# `_seen_trial_indices_by_status` is set during `gen`, to the experiment's
# `trial_indices_by_Status` at the time of candidate generation.
self.assertNotEqual(
sobol_gs_with_update._seen_trial_indices_by_status,
exp.trial_indices_by_status,
)
# Try with passing data.
sobol_gs_with_update.gen(
experiment=exp, data=get_branin_data(trial_indices=range(4)))
# Now `_seen_trial_indices_by_status` should be set to experiment's,
self.assertEqual(
sobol_gs_with_update._seen_trial_indices_by_status,
exp.trial_indices_by_status,
)
# Only the data for the last completed trial should be considered new and passed
# to `update`.
self.assertEqual(
set(mock_update.call_args[1].get(
"new_data").df["trial_index"].values), {3})
# Try with passing same data as before; no update should be performed.
with patch.object(sobol_gs_with_update,
"_update_current_model") as mock_update:
sobol_gs_with_update.gen(
experiment=exp, data=get_branin_data(trial_indices=range(4)))
mock_update.assert_not_called()
def testDecodeGenerationStrategy(self):
generation_strategy = get_generation_strategy()
experiment = get_branin_experiment()
gs_json = object_to_json(generation_strategy)
new_generation_strategy = generation_strategy_from_json(gs_json)
self.assertEqual(generation_strategy, new_generation_strategy)
self.assertGreater(len(new_generation_strategy._steps), 0)
self.assertIsInstance(new_generation_strategy._steps[0].model, Models)
# Model has not yet been initialized on this GS since it hasn't generated
# anything yet.
self.assertIsNone(new_generation_strategy.model)
# Check that we can encode and decode the generation strategy after
# it has generated some generator runs.
generation_strategy = new_generation_strategy
gr = generation_strategy.gen(experiment)
gs_json = object_to_json(generation_strategy)
new_generation_strategy = generation_strategy_from_json(gs_json)
self.assertEqual(generation_strategy, new_generation_strategy)
self.assertIsInstance(new_generation_strategy._steps[0].model, Models)
# Since this GS has now generated one generator run, model should have
# been initialized and restored when decoding from JSON.
self.assertIsInstance(new_generation_strategy.model, ModelBridge)
# Check that we can encode and decode the generation strategy after
# it has generated some trials and been updated with some data.
generation_strategy = new_generation_strategy
experiment.new_trial(gr) # Add previously generated GR as trial.
# Make generation strategy aware of the trial's data via `gen`.
generation_strategy.gen(experiment, data=get_branin_data())
gs_json = object_to_json(generation_strategy)
new_generation_strategy = generation_strategy_from_json(gs_json)
self.assertEqual(generation_strategy, new_generation_strategy)
self.assertIsInstance(new_generation_strategy._steps[0].model, Models)
self.assertIsInstance(new_generation_strategy.model, ModelBridge)
def test_transform_callback_int(self, *_):
exp = get_branin_experiment(with_batch=True)
data = get_branin_data(trial_indices=exp.trials)
parameters = [
RangeParameter(name="x1",
parameter_type=ParameterType.INT,
lower=1,
upper=10),
RangeParameter(name="x2",
parameter_type=ParameterType.INT,
lower=5,
upper=15),
]
gpei = TorchModelBridge(
experiment=exp,
data=data,
search_space=SearchSpace(parameters=parameters),
model=BotorchModel(),
transforms=[IntToFloat],
torch_dtype=torch.double,
fit_out_of_design=True,
)
transformed = gpei._transform_callback([5.4, 7.6])
self.assertTrue(np.allclose(transformed, [5, 8]))
np_mb = ArrayModelBridge(
experiment=exp,
data=exp.fetch_data(),
search_space=SearchSpace(parameters=parameters),
model=NumpyModel(),
transforms=[IntToFloat],
)
transformed = np_mb._transform_callback(np.array([5.4, 7.6]))
self.assertTrue(np.allclose(transformed, [5, 8]))
def test_use_update(self, mock_fetch_trials_data, mock_update):
exp = get_branin_experiment()
sobol_gs_with_update = GenerationStrategy(
steps=[GenerationStep(model=Models.SOBOL, num_trials=-1, use_update=True)]
)
# Try without passing data (generation strategy fetches data from experiment).
trial = exp.new_trial(generator_run=sobol_gs_with_update.gen(experiment=exp))
mock_update.assert_not_called()
trial._status = TrialStatus.COMPLETED
for i in range(3):
trial = exp.new_trial(
generator_run=sobol_gs_with_update.gen(experiment=exp)
)
self.assertEqual(
mock_fetch_trials_data.call_args[1].get("trial_indices"), {i}
)
trial._status = TrialStatus.COMPLETED
# Try with passing data.
sobol_gs_with_update.gen(
experiment=exp, data=get_branin_data(trial_indices=range(4))
)
# Only the data for the last completed trial should be considered new and passed
# to `update`.
self.assertEqual(
set(mock_update.call_args[1].get("new_data").df["trial_index"].values), {3}
)
def testUpdateGenerationStrategyIncrementally(self):
experiment = get_branin_experiment()
generation_strategy = choose_generation_strategy(
experiment.search_space)
save_experiment(experiment=experiment)
save_generation_strategy(generation_strategy=generation_strategy)
# add generator runs, save, reload
generator_runs = []
for i in range(7):
data = get_branin_data() if i > 0 else None
gr = generation_strategy.gen(experiment, data=data)
generator_runs.append(gr)
trial = experiment.new_trial(generator_run=gr).mark_running(
no_runner_required=True)
trial.mark_completed()
save_experiment(experiment=experiment)
update_generation_strategy(generation_strategy=generation_strategy,
generator_runs=generator_runs)
loaded_generation_strategy = load_generation_strategy_by_experiment_name(
experiment_name=experiment.name)
self.assertEqual(generation_strategy._curr.index,
loaded_generation_strategy._curr.index, 1)
self.assertEqual(len(loaded_generation_strategy._generator_runs), 7)
def _run_GS_for_N_rounds(self, gs: GenerationStrategy, exp: Experiment,
num_rounds: int) -> List[int]:
could_gen = []
for _ in range(num_rounds):
(
num_trials_to_gen,
opt_complete,
) = gs.current_generator_run_limit()
self.assertFalse(opt_complete)
could_gen.append(num_trials_to_gen)
trials = []
for _ in range(num_trials_to_gen):
gr = gs.gen(
experiment=exp,
pending_observations=get_pending(experiment=exp),
)
trials.append(
exp.new_trial(gr).mark_running(no_runner_required=True))
for trial in trials:
exp.attach_data(get_branin_data(trial_indices=[trial.index]))
trial.mark_completed()
return could_gen
def test_current_generator_run_limit(self):
NUM_INIT_TRIALS = 5
SECOND_STEP_PARALLELISM = 3
NUM_ROUNDS = 4
exp = get_branin_experiment()
sobol_gs_with_parallelism_limits = GenerationStrategy(steps=[
GenerationStep(
model=Models.SOBOL,
num_trials=NUM_INIT_TRIALS,
min_trials_observed=3,
),
GenerationStep(
model=Models.SOBOL,
num_trials=-1,
max_parallelism=SECOND_STEP_PARALLELISM,
),
])
sobol_gs_with_parallelism_limits._experiment = exp
could_gen = []
for _ in range(NUM_ROUNDS):
(
num_trials_to_gen,
opt_complete,
) = sobol_gs_with_parallelism_limits.current_generator_run_limit()
self.assertFalse(opt_complete)
could_gen.append(num_trials_to_gen)
trials = []
for _ in range(num_trials_to_gen):
gr = sobol_gs_with_parallelism_limits.gen(
experiment=exp,
pending_observations=get_pending(experiment=exp),
)
trials.append(
exp.new_trial(gr).mark_running(no_runner_required=True))
for trial in trials:
exp.attach_data(get_branin_data(trial_indices=[trial.index]))
trial.mark_completed()
# We expect trials from first generation step + trials from remaining rounds in
# batches limited by parallelism setting in the second step.
self.assertEqual(
len(exp.trials),
NUM_INIT_TRIALS + (NUM_ROUNDS - 1) * SECOND_STEP_PARALLELISM,
)
self.assertTrue(all(t.status.is_completed
for t in exp.trials.values()))
self.assertEqual(could_gen, [NUM_INIT_TRIALS] +
[SECOND_STEP_PARALLELISM] * (NUM_ROUNDS - 1))
def testUpdateGenerationStrategy(self):
generation_strategy = get_generation_strategy()
save_generation_strategy(generation_strategy=generation_strategy)
experiment = get_branin_experiment()
generation_strategy = get_generation_strategy()
save_experiment(experiment)
# add generator run, save, reload
experiment.new_trial(generator_run=generation_strategy.gen(experiment))
save_generation_strategy(generation_strategy=generation_strategy)
loaded_generation_strategy = load_generation_strategy_by_experiment_name(
experiment_name=experiment.name
)
# `_seen_trial_indices_by_status` attribute of a GS is not saved in DB,
# so it will be None in the restored version of the GS.
# Hackily removing it from the original GS to check equality.
generation_strategy._seen_trial_indices_by_status = None
self.assertEqual(generation_strategy, loaded_generation_strategy)
# add another generator run, save, reload
experiment.new_trial(
generator_run=generation_strategy.gen(experiment, data=get_branin_data())
)
save_generation_strategy(generation_strategy=generation_strategy)
save_experiment(experiment)
loaded_generation_strategy = load_generation_strategy_by_experiment_name(
experiment_name=experiment.name
)
# `_seen_trial_indices_by_status` attribute of a GS is not saved in DB,
# so it will be None in the restored version of the GS.
# Hackily removing it from the original GS to check equality.
generation_strategy._seen_trial_indices_by_status = None
self.assertEqual(generation_strategy, loaded_generation_strategy)
# make sure that we can update the experiment too
experiment.description = "foobar"
save_experiment(experiment)
loaded_generation_strategy = load_generation_strategy_by_experiment_name(
experiment_name=experiment.name
)
self.assertEqual(generation_strategy, loaded_generation_strategy)
self.assertEqual(
generation_strategy._experiment.description, experiment.description
)
self.assertEqual(
generation_strategy._experiment.description,
loaded_generation_strategy._experiment.description,
)
def testUpdateGenerationStrategy(self):
generation_strategy = get_generation_strategy()
save_generation_strategy(generation_strategy=generation_strategy)
# Add data, save, reload
generation_strategy._data = Data(
df=pd.DataFrame.from_records([{
"metric_name": "foo",
"mean": 1,
"arm_name": "bar"
}]))
save_generation_strategy(generation_strategy=generation_strategy)
loaded_generation_strategy = load_generation_strategy_by_id(
gs_id=generation_strategy._db_id)
self.assertEqual(generation_strategy, loaded_generation_strategy)
experiment = get_branin_experiment()
generation_strategy = get_generation_strategy()
save_experiment(experiment)
# add generator run, save, reload
experiment.new_trial(generator_run=generation_strategy.gen(experiment))
save_generation_strategy(generation_strategy=generation_strategy)
loaded_generation_strategy = load_generation_strategy_by_experiment_name(
experiment_name=experiment.name)
self.assertEqual(generation_strategy, loaded_generation_strategy)
# add another generator run, save, reload
experiment.new_trial(generator_run=generation_strategy.gen(
experiment, new_data=get_branin_data()))
save_generation_strategy(generation_strategy=generation_strategy)
save_experiment(experiment)
loaded_generation_strategy = load_generation_strategy_by_experiment_name(
experiment_name=experiment.name)
self.assertEqual(generation_strategy, loaded_generation_strategy)
# make sure that we can update the experiment too
experiment.description = "foobar"
save_experiment(experiment)
loaded_generation_strategy = load_generation_strategy_by_experiment_name(
experiment_name=experiment.name)
self.assertEqual(generation_strategy, loaded_generation_strategy)
self.assertEqual(generation_strategy._experiment.description,
experiment.description)
self.assertEqual(
generation_strategy._experiment.description,
loaded_generation_strategy._experiment.description,
)
def testEncodeDecodeGenerationStrategy(self):
# Cannot load generation strategy before it has been saved
with self.assertRaises(ValueError):
load_generation_strategy_by_id(gs_id=0)
# Check that we can encode and decode the generation strategy *before*
# it has generated some trials and been updated with some data.
generation_strategy = get_generation_strategy()
# Check that we can save a generation strategy without an experiment
# attached.
save_generation_strategy(generation_strategy=generation_strategy)
# Also try restoring this generation strategy by its ID in the DB.
new_generation_strategy = load_generation_strategy_by_id(
gs_id=generation_strategy._db_id
)
self.assertEqual(generation_strategy, new_generation_strategy)
self.assertIsNone(generation_strategy._experiment)
# Cannot load generation strategy before it has been saved
experiment = get_branin_experiment()
save_experiment(experiment)
with self.assertRaises(ValueError):
load_generation_strategy_by_experiment_name(experiment_name=experiment.name)
# Check that we can encode and decode the generation strategy *after*
# it has generated some trials and been updated with some data.
generation_strategy = new_generation_strategy
experiment.new_trial(generation_strategy.gen(experiment=experiment))
experiment.new_trial(
generation_strategy.gen(experiment, data=get_branin_data())
)
save_generation_strategy(generation_strategy=generation_strategy)
save_experiment(experiment)
# Try restoring the generation strategy using the experiment its
# attached to.
new_generation_strategy = load_generation_strategy_by_experiment_name(
experiment_name=experiment.name
)
# `_seen_trial_indices_by_status` attribute of a GS is not saved in DB,
# so it will be None in the restored version of the GS.
# Hackily removing it from the original GS to check equality.
generation_strategy._seen_trial_indices_by_status = None
self.assertEqual(generation_strategy, new_generation_strategy)
self.assertIsInstance(new_generation_strategy._steps[0].model, Models)
self.assertIsInstance(new_generation_strategy.model, ModelBridge)
self.assertEqual(len(new_generation_strategy._generator_runs), 2)
self.assertEqual(new_generation_strategy._experiment._name, experiment._name)
def testEncodeDecodeGenerationStrategy(self):
# Cannot load generation strategy before it has been saved
with self.assertRaises(ValueError):
load_generation_strategy_by_id(gs_id=0)
# Check that we can encode and decode the generation strategy *before*
# it has generated some trials and been updated with some data.
generation_strategy = get_generation_strategy()
# Check that we can save a generation strategy without an experiment
# attached.
save_generation_strategy(generation_strategy=generation_strategy)
# Also try restoring this generation strategy by its ID in the DB.
new_generation_strategy = load_generation_strategy_by_id(
gs_id=generation_strategy._db_id)
self.assertEqual(generation_strategy, new_generation_strategy)
self.assertIsNone(generation_strategy._experiment)
self.assertEqual(len(generation_strategy._generated), 0)
self.assertEqual(len(generation_strategy._observed), 0)
# Cannot load generation strategy before it has been saved
experiment = get_branin_experiment()
save_experiment(experiment)
with self.assertRaises(ValueError):
load_generation_strategy_by_experiment_name(
experiment_name=experiment.name)
# Check that we can encode and decode the generation strategy *after*
# it has generated some trials and been updated with some data.
generation_strategy = new_generation_strategy
experiment.new_trial(generator_run=generation_strategy.gen(experiment))
experiment.new_trial(
generation_strategy.gen(experiment, new_data=get_branin_data()))
self.assertGreater(len(generation_strategy._generated), 0)
self.assertGreater(len(generation_strategy._observed), 0)
save_generation_strategy(generation_strategy=generation_strategy)
save_experiment(experiment)
# Try restoring the generation strategy using the experiment its
# attached to.
new_generation_strategy = load_generation_strategy_by_experiment_name(
experiment_name=experiment.name)
self.assertEqual(generation_strategy, new_generation_strategy)
self.assertIsInstance(new_generation_strategy._steps[0].model, Models)
self.assertIsInstance(new_generation_strategy.model, ModelBridge)
self.assertEqual(len(new_generation_strategy._generator_runs), 2)
self.assertEqual(new_generation_strategy._experiment._name,
experiment._name)
def testUpdateGenerationStrategy(self):
generation_strategy = get_generation_strategy()
save_generation_strategy(generation_strategy=generation_strategy)
experiment = get_branin_experiment()
generation_strategy = get_generation_strategy()
save_experiment(experiment)
# add generator run, save, reload
experiment.new_trial(generator_run=generation_strategy.gen(experiment))
save_generation_strategy(generation_strategy=generation_strategy)
loaded_generation_strategy = load_generation_strategy_by_experiment_name(
experiment_name=experiment.name)
self.assertEqual(generation_strategy, loaded_generation_strategy)
# add another generator run, save, reload
experiment.new_trial(generator_run=generation_strategy.gen(
experiment, data=get_branin_data()))
save_generation_strategy(generation_strategy=generation_strategy)
save_experiment(experiment)
loaded_generation_strategy = load_generation_strategy_by_experiment_name(
experiment_name=experiment.name)
# During restoration of generation strategy's model from its last generator
# run, we set `_seen_trial_indices_by_status` to that of the experiment,
# from which we are grabbing the data to restore the model with. When the
# experiment was updated more recently than the last `gen` from generation
# strategy, the generation strategy prior to save might not have 'seen'
# some recently added trials, so we update the mappings to match and check
# that the generation strategies are equal otherwise.
generation_strategy._seen_trial_indices_by_status[
TrialStatus.CANDIDATE].add(1)
self.assertEqual(generation_strategy, loaded_generation_strategy)
# make sure that we can update the experiment too
experiment.description = "foobar"
save_experiment(experiment)
loaded_generation_strategy = load_generation_strategy_by_experiment_name(
experiment_name=experiment.name)
self.assertEqual(generation_strategy, loaded_generation_strategy)
self.assertEqual(generation_strategy._experiment.description,
experiment.description)
self.assertEqual(
generation_strategy._experiment.description,
loaded_generation_strategy._experiment.description,
)
def testDecodeGenerationStrategy(self):
generation_strategy = get_generation_strategy()
experiment = get_branin_experiment()
gs_json = object_to_json(generation_strategy)
new_generation_strategy = generation_strategy_from_json(gs_json)
self.assertEqual(generation_strategy, new_generation_strategy)
self.assertGreater(len(new_generation_strategy._steps), 0)
self.assertIsInstance(new_generation_strategy._steps[0].model, Models)
# Model has not yet been initialized on this GS since it hasn't generated
# anything yet.
self.assertIsNone(new_generation_strategy.model)
# Check that we can encode and decode the generation strategy after
# it has generated some trials.
generation_strategy = new_generation_strategy
experiment.new_trial(generator_run=generation_strategy.gen(experiment))
gs_json = object_to_json(generation_strategy)
new_generation_strategy = generation_strategy_from_json(gs_json)
# `_seen_trial_indices_by_status` attribute of a GS is not saved in DB,
# so it will be None in the restored version of the GS.
# Hackily removing it from the original GS to check equality.
generation_strategy._seen_trial_indices_by_status = None
self.assertEqual(generation_strategy, new_generation_strategy)
self.assertIsInstance(new_generation_strategy._steps[0].model, Models)
# Since this GS has now generated one generator run, model should have
# been initialized and restored when decoding from JSON.
self.assertIsInstance(new_generation_strategy.model, ModelBridge)
# Check that we can encode and decode the generation strategy after
# it has generated some trials and been updated with some data.
generation_strategy = new_generation_strategy
experiment.new_trial(
generation_strategy.gen(experiment, data=get_branin_data()))
gs_json = object_to_json(generation_strategy)
new_generation_strategy = generation_strategy_from_json(gs_json)
# `_seen_trial_indices_by_status` attribute of a GS is not saved in DB,
# so it will be None in the restored version of the GS.
# Hackily removing it from the original GS to check equality.
generation_strategy._seen_trial_indices_by_status = None
self.assertEqual(generation_strategy, new_generation_strategy)
self.assertIsInstance(new_generation_strategy._steps[0].model, Models)
self.assertIsInstance(new_generation_strategy.model, ModelBridge)
def testUpdateGenerationStrategy(self):
generation_strategy = get_generation_strategy()
save_generation_strategy(generation_strategy=generation_strategy)
experiment = get_branin_experiment()
generation_strategy = get_generation_strategy()
save_experiment(experiment)
# add generator run, save, reload
experiment.new_trial(generator_run=generation_strategy.gen(experiment))
save_generation_strategy(generation_strategy=generation_strategy)
loaded_generation_strategy = load_generation_strategy_by_experiment_name(
experiment_name=experiment.name)
self.assertEqual(generation_strategy, loaded_generation_strategy)
# add another generator run, save, reload
experiment.new_trial(generator_run=generation_strategy.gen(
experiment, data=get_branin_data()))
save_generation_strategy(generation_strategy=generation_strategy)
save_experiment(experiment)
loaded_generation_strategy = load_generation_strategy_by_experiment_name(
experiment_name=experiment.name)
self.assertEqual(generation_strategy, loaded_generation_strategy)
# make sure that we can update the experiment too
experiment.description = "foobar"
save_experiment(experiment)
loaded_generation_strategy = load_generation_strategy_by_experiment_name(
experiment_name=experiment.name)
self.assertEqual(generation_strategy, loaded_generation_strategy)
self.assertEqual(generation_strategy._experiment.description,
experiment.description)
self.assertEqual(
generation_strategy._experiment.description,
loaded_generation_strategy._experiment.description,
)
def test_candidate_metadata_propagation(self, mock_model_fit,
mock_model_update, mock_model_gen):
exp = get_branin_experiment(with_status_quo=True, with_batch=True)
# Check that the metadata is correctly re-added to observation
# features during `fit`.
preexisting_batch_gr = exp.trials[0]._generator_run_structs[
0].generator_run
preexisting_batch_gr._candidate_metadata_by_arm_signature = {
preexisting_batch_gr.arms[0].signature: {
"preexisting_batch_cand_metadata": "some_value"
}
}
modelbridge = ArrayModelBridge(
search_space=exp.search_space,
experiment=exp,
model=NumpyModel(),
data=get_branin_data(),
)
self.assertTrue(
np.array_equal(
mock_model_fit.call_args[1].get("Xs"),
np.array([[list(exp.trials[0].arms[0].parameters.values())]]),
))
self.assertEqual(
mock_model_fit.call_args[1].get("candidate_metadata"),
[[{
"preexisting_batch_cand_metadata": "some_value"
}]],
)
# Check that `gen` correctly propagates the metadata to the GR.
gr = modelbridge.gen(n=1)
self.assertEqual(
gr.candidate_metadata_by_arm_signature,
{
gr.arms[0].signature: {
"some_key": "some_value_0"
},
gr.arms[1].signature: {
"some_key": "some_value_1"
},
},
)
# Check that the metadata is correctly re-added to observation
# features during `update`.
batch = exp.new_batch_trial(generator_run=gr)
modelbridge.update(
experiment=exp,
new_data=get_branin_data(trial_indices=[batch.index]))
self.assertTrue(
np.array_equal(
mock_model_update.call_args[1].get("Xs"),
np.array(
[[list(exp.trials[0].arms[0].parameters.values()), [1,
2]]]),
))
self.assertEqual(
mock_model_update.call_args[1].get("candidate_metadata"),
[[
{
"preexisting_batch_cand_metadata": "some_value"
},
# new data contained data just for arm '1_0', not for '1_1',
# so we don't expect to see '{"some_key": "some_value_1"}'
# in candidate metadata.
{
"some_key": "some_value_0"
},
]],
)
# Check that `None` candidate metadata is handled correctly.
mock_model_gen.return_value = (
np.array([[2, 4], [3, 5]]),
np.array([1, 2]),
None,
{},
)
gr = modelbridge.gen(n=1)
self.assertIsNone(gr.candidate_metadata_by_arm_signature)
# Check that the metadata is correctly re-added to observation
# features during `update`.
batch = exp.new_batch_trial(generator_run=gr)
modelbridge.update(
experiment=exp,
new_data=get_branin_data(trial_indices=[batch.index]))
self.assertTrue(
np.array_equal(
mock_model_update.call_args[1].get("Xs"),
np.array([[
list(exp.trials[0].arms[0].parameters.values()), [1, 2],
[2, 4]
]]),
))
self.assertEqual(
mock_model_update.call_args[1].get("candidate_metadata"),
[[
{
"preexisting_batch_cand_metadata": "some_value"
},
{
"some_key": "some_value_0"
},
{},
]],
)
# Check that no candidate metadata is handled correctly.
exp = get_branin_experiment(with_status_quo=True)
modelbridge = ArrayModelBridge(search_space=exp.search_space,
experiment=exp,
model=NumpyModel())
# Hack in outcome names to bypass validation (since we instantiated model
# without data).
modelbridge.outcomes = modelbridge._metric_names = next(
iter(exp.metrics))
gr = modelbridge.gen(n=1)
self.assertIsNone(
mock_model_fit.call_args[1].get("candidate_metadata"))
self.assertIsNone(gr.candidate_metadata_by_arm_signature)
batch = exp.new_batch_trial(generator_run=gr)
modelbridge.update(
experiment=exp,
new_data=get_branin_data(trial_indices=[batch.index]))
def testDecodeGenerationStrategy(self):
generation_strategy = get_generation_strategy()
experiment = get_branin_experiment()
gs_json = object_to_json(
generation_strategy,
encoder_registry=DEPRECATED_ENCODER_REGISTRY,
class_encoder_registry=DEPRECATED_CLASS_ENCODER_REGISTRY,
)
new_generation_strategy = generation_strategy_from_json(
gs_json,
decoder_registry=DEPRECATED_DECODER_REGISTRY,
class_decoder_registry=DEPRECATED_CLASS_DECODER_REGISTRY,
)
self.assertEqual(generation_strategy, new_generation_strategy)
self.assertGreater(len(new_generation_strategy._steps), 0)
self.assertIsInstance(new_generation_strategy._steps[0].model, Models)
# Model has not yet been initialized on this GS since it hasn't generated
# anything yet.
self.assertIsNone(new_generation_strategy.model)
# Check that we can encode and decode the generation strategy after
# it has generated some generator runs. Since we now need to `gen`,
# we remove the fake callable kwarg we added, since model does not
# expect it.
generation_strategy = get_generation_strategy(with_callable_model_kwarg=False)
gr = generation_strategy.gen(experiment)
gs_json = object_to_json(
generation_strategy,
encoder_registry=DEPRECATED_ENCODER_REGISTRY,
class_encoder_registry=DEPRECATED_CLASS_ENCODER_REGISTRY,
)
new_generation_strategy = generation_strategy_from_json(
gs_json,
decoder_registry=DEPRECATED_DECODER_REGISTRY,
class_decoder_registry=DEPRECATED_CLASS_DECODER_REGISTRY,
)
self.assertEqual(generation_strategy, new_generation_strategy)
self.assertIsInstance(new_generation_strategy._steps[0].model, Models)
# Since this GS has now generated one generator run, model should have
# been initialized and restored when decoding from JSON.
self.assertIsInstance(new_generation_strategy.model, ModelBridge)
# Check that we can encode and decode the generation strategy after
# it has generated some trials and been updated with some data.
generation_strategy = new_generation_strategy
experiment.new_trial(gr) # Add previously generated GR as trial.
# Make generation strategy aware of the trial's data via `gen`.
generation_strategy.gen(experiment, data=get_branin_data())
gs_json = object_to_json(
generation_strategy,
encoder_registry=DEPRECATED_ENCODER_REGISTRY,
class_encoder_registry=DEPRECATED_CLASS_ENCODER_REGISTRY,
)
new_generation_strategy = generation_strategy_from_json(
gs_json,
decoder_registry=DEPRECATED_DECODER_REGISTRY,
class_decoder_registry=DEPRECATED_CLASS_DECODER_REGISTRY,
)
self.assertEqual(generation_strategy, new_generation_strategy)
self.assertIsInstance(new_generation_strategy._steps[0].model, Models)
self.assertIsInstance(new_generation_strategy.model, ModelBridge)
