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)