diff --git a/libensemble/gen_funcs/persistent_ax_multitask.py b/libensemble/gen_funcs/persistent_ax_multitask.py index 0a2e07f20..30ce3634e 100644 --- a/libensemble/gen_funcs/persistent_ax_multitask.py +++ b/libensemble/gen_funcs/persistent_ax_multitask.py @@ -8,7 +8,7 @@ This `gen_f` is meant to be used with the `alloc_f` function `only_persistent_gens` -Requires: Ax>=0.5.0 +Requires: Ax==0.5.0 Ax notes: Each arm = a set of simulation inputs (a sim_id) diff --git a/libensemble/gen_funcs/persistent_botorch_mfkg_branin.py b/libensemble/gen_funcs/persistent_botorch_mfkg_branin.py new file mode 100644 index 000000000..2a4081c06 --- /dev/null +++ b/libensemble/gen_funcs/persistent_botorch_mfkg_branin.py @@ -0,0 +1,211 @@ +""" +This file defines the persistent generator function for multi-fidelity Bayesian +optimization using BoTorch's Multi-Fidelity Knowledge Gradient (MFKG) acquisition function. + +This gen_f is meant to be used with the alloc_f function `only_persistent_gens`. +""" + +import numpy as np +import torch +from botorch import fit_gpytorch_mll +from botorch.acquisition import PosteriorMean +from botorch.acquisition.cost_aware import InverseCostWeightedUtility +from botorch.acquisition.fixed_feature import FixedFeatureAcquisitionFunction +from botorch.acquisition.knowledge_gradient import qMultiFidelityKnowledgeGradient +from botorch.acquisition.utils import project_to_target_fidelity +from botorch.models.cost import AffineFidelityCostModel +from botorch.models.gp_regression_fidelity import SingleTaskMultiFidelityGP +from botorch.models.transforms.outcome import Standardize +from botorch.optim.optimize import optimize_acqf, optimize_acqf_mixed +from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood + +from libensemble.message_numbers import EVAL_GEN_TAG, FINISHED_PERSISTENT_GEN_TAG, PERSIS_STOP, STOP_TAG +from libensemble.tools.persistent_support import PersistentSupport + +__all__ = ["persistent_botorch_mfkg"] + +# Set seeds for reproducibility +np.random.seed(42) +torch.manual_seed(42) + +# Torch settings +tkwargs = { + "dtype": torch.double, + "device": torch.device("cuda" if torch.cuda.is_available() else "cpu"), +} + +# Specify target fidelity +target_fidelities = {2: 1.0} + +# Specify cost model +cost_model = AffineFidelityCostModel(fidelity_weights={2: 0.9}, fixed_cost=0.1) +cost_aware_utility = InverseCostWeightedUtility(cost_model=cost_model) + + +# Custom function to project posterior to target fidelity (defer to default) +def project(X): + return project_to_target_fidelity(X=X, target_fidelities=target_fidelities) + + +# Wrapper function for compatibility with existing code +def problem(X, ps, gen_specs): + """ + Send a batch of points to libE for evaluation and collect the results. + + Results are gathered with a ``recv`` loop rather than a single ``send_recv`` + so this works in both batch return mode (``async_return`` off, where the + whole batch arrives in one message) and asynchronous return mode + (``async_return`` on with ``active_recv_gen``, where results stream back one + or a few at a time, possibly out of order). The evaluated points are + reconstructed from the returned rows so the (x, fidelity) -> f pairing stays + correct regardless of the order results arrive in. + + Args: + X: tensor of shape (n, 3) where columns are [x0, x1, fidelity] + ps: PersistentSupport object for communication + gen_specs: Generator specifications + + Returns: + (new_x, new_obj, tag): tensor of the evaluated points (shape (m, 3)), + tensor of their objective values (shape (m, 1)), and the last message + tag. new_x/new_obj are None if no results were received (e.g. on stop). + """ + # Send points to be evaluated. + X_np = X.cpu().numpy() + H_o = np.zeros(len(X), dtype=gen_specs["out"]) + H_o["x"] = X_np[:, :2] + H_o["fidelity"] = X_np[:, 2] + ps.send(H_o) + + # Collect results until the whole batch is back (or we are told to stop). + n_expected = len(X) + xs = [] + objs = [] + tag = None + while len(objs) < n_expected: + tag, Work, calc_in = ps.recv() + if tag in [STOP_TAG, PERSIS_STOP]: + break + if calc_in is None or len(calc_in) == 0: + continue + for row in calc_in: + xs.append(np.append(row["x"], row["fidelity"])) + objs.append(float(row["f"])) + + if len(objs) == 0: + return None, None, tag + + new_x = torch.tensor(np.array(xs), **tkwargs) + new_obj = torch.tensor(np.array(objs), **tkwargs).unsqueeze(-1) + + return new_x, new_obj, tag + + +# Function to generate training data +def generate_initial_data(n, ps, gen_specs, bounds): # Jeff: Initial sample size is twice this value of n + train_x = bounds[0, :-1] + (bounds[1, :-1] - bounds[0, :-1]) * torch.rand(n, 2, **tkwargs) + train_lf = torch.zeros(n, 1) + train_hf = torch.ones(n, 1) + train_x_full_lf = torch.cat((train_x, train_lf), dim=1) + train_x_full_hf = torch.cat((train_x, train_hf), dim=1) + train_x_full = torch.cat((train_x_full_lf, train_x_full_hf), dim=0) + # problem returns the evaluated points (reconstructed from the results) so + # that x and objective stay paired even when results arrive out of order. + new_x, train_obj, tag = problem(train_x_full, ps, gen_specs) + return new_x, train_obj, tag + + +# Function to initialize a botorch model +def initialize_model(train_x, train_obj): + model = SingleTaskMultiFidelityGP(train_x, train_obj, outcome_transform=Standardize(m=1), data_fidelities=[2]) + mll = ExactMarginalLogLikelihood(model.likelihood, model) + return mll, model + + +# Multifidelity Knowledge Gradient acquisition function +def get_mfkg(model, bounds): + + curr_val_acqf = FixedFeatureAcquisitionFunction( + acq_function=PosteriorMean(model), + d=3, + columns=[2], + values=[1], + ) + + _, current_value = optimize_acqf( + acq_function=curr_val_acqf, + bounds=bounds[:, :-1], + q=1, # Jeff: Don't adjust this for some reason + num_restarts=1, # Jeff: I decreased this to make libE development faster + raw_samples=10, # Jeff: I decreased this to make libE development faster + options={"batch_limit": 10, "maxiter": 10}, # Jeff: I decreased this to make libE development faster + ) + + return qMultiFidelityKnowledgeGradient( + model=model, + num_fantasies=128, + current_value=current_value, + cost_aware_utility=cost_aware_utility, + project=project, + ) + + +# Optimization step +def optimize_mfkg_and_get_observation(mfkg_acqf, q, ps, gen_specs, bounds): + # Generate new candidates + candidates, _ = optimize_acqf_mixed( + acq_function=mfkg_acqf, + bounds=bounds, + fixed_features_list=[{2: 0.0}, {2: 1.0}], + q=q, # Jeff: This is the number of new samples to make + num_restarts=1, # Jeff: I decreased this to make libE development faster + raw_samples=10, # Jeff: I decreased this to make libE development faster + options={"batch_limit": 10, "maxiter": 10}, # Jeff: I decreased this to make libE development faster + ) + + # Observe new values. problem returns the actually-evaluated points paired + # with their objectives (order-independent), which we use for training. + cost = cost_model(candidates).sum() + new_x, new_obj, tag = problem(candidates.detach(), ps, gen_specs) + return new_x, new_obj, cost, tag + + +# Function to perform a single iteration +def do_iteration(train_x, train_obj, q, ps, gen_specs, bounds): + mll, model = initialize_model(train_x, train_obj) + fit_gpytorch_mll(mll) + mfkg_acqf = get_mfkg(model, bounds) + new_x, new_obj, _, tag = optimize_mfkg_and_get_observation(mfkg_acqf, q, ps, gen_specs, bounds) + + if new_obj is None: + return model, train_x, train_obj, tag + + train_x = torch.cat([train_x, new_x]) + train_obj = torch.cat([train_obj, new_obj]) # Jeff: This is where the "sim" eval happens, to be send the manager + + return model, train_x, train_obj, tag + + +def persistent_botorch_mfkg(H, persis_info, gen_specs, libE_info): + """ + Persistent generator function for multi-fidelity Bayesian optimization using BoTorch's MFKG. + """ + ps = PersistentSupport(libE_info, EVAL_GEN_TAG) + + # Extract user parameters + ub = np.asarray(gen_specs["user"]["ub"], dtype=float) + lb = np.asarray(gen_specs["user"]["lb"], dtype=float) + bounds = torch.tensor(np.array([np.append(lb, 0.0), np.append(ub, 1.0)]), **tkwargs) + + n_init_samples = gen_specs["user"]["n_init_samples"] + q = gen_specs["user"]["q"] + + # ## Perform Multifidelity Bayesian Optimization + # Generate initial data + train_x, train_obj, tag = generate_initial_data(n_init_samples, ps, gen_specs, bounds) + + # Step + while tag not in [STOP_TAG, PERSIS_STOP]: + model, train_x, train_obj, tag = do_iteration(train_x, train_obj, q, ps, gen_specs, bounds) + + return None, persis_info, FINISHED_PERSISTENT_GEN_TAG diff --git a/libensemble/sim_funcs/augmented_branin.py b/libensemble/sim_funcs/augmented_branin.py new file mode 100644 index 000000000..8aa1300cd --- /dev/null +++ b/libensemble/sim_funcs/augmented_branin.py @@ -0,0 +1,37 @@ +""" +This module evaluates the augmented Branin function for multi-fidelity optimization. + +Augmented Branin is a modified version of the Branin function with a fidelity parameter. +""" + +__all__ = ["augmented_branin", "augmented_branin_func"] + +import numpy as np + + +def augmented_branin(H, persis_info, sim_specs, libE_info): + """ + Evaluates the augmented Branin function for a collection of points given in ``H["x"]`` + with fidelity values in ``H["fidelity"]``. + """ + batch = len(H["x"]) + H_o = np.zeros(batch, dtype=sim_specs["out"]) + + for i in range(batch): + x = H["x"][i] + fidelity = H["fidelity"][i] + H_o["f"][i] = augmented_branin_func(x.reshape(1, -1), fidelity)[0] + + return H_o, persis_info + + +def augmented_branin_func(x, fidelity): + """Augmented Branin function for multi-fidelity optimization.""" + x0 = x[:, 0] + x1 = x[:, 1] + + t1 = 15 * x1 - (5.1 / (4 * np.pi**2) - 0.1 * (1 - fidelity)) * (15 * x0 - 5) ** 2 + 5 / np.pi * (15 * x0 - 5) - 6 + t2 = 10 * (1 - 1 / (8 * np.pi)) * np.cos(15 * x0 - 5) + result = t1**2 + t2 + 10 + + return -result # negate for maximization diff --git a/libensemble/tests/regression_tests/test_mfkg_branin.py b/libensemble/tests/regression_tests/test_mfkg_branin.py new file mode 100644 index 000000000..d28dc0bc1 --- /dev/null +++ b/libensemble/tests/regression_tests/test_mfkg_branin.py @@ -0,0 +1,141 @@ +""" +Example of multi-fidelity optimization using a persistent BoTorch MFKG gen_func. + +One worker runs the persistent MFKG generator; the remaining workers evaluate +the (augmented Branin) objective. The generator collects a full batch of ``q`` +points before refitting its GP model, so a batch size of ``q`` keeps up to +``q`` simulation workers busy at once. + +This script runs the ensemble twice to exercise both ways of returning results +to the generator: + +* ``async_return=False`` (batch): the whole batch of ``q`` evaluations is + returned to the generator in one message once all have completed. +* ``async_return=True`` with ``active_recv_gen=True``: each evaluation is + handed back as soon as it finishes (possibly out of order); the generator + collects the batch as results stream in. + +Execute via one of the following commands: + mpiexec -np 5 python run_botorch_mfkg_branin.py + python run_botorch_mfkg_branin.py --nworkers 4 + python run_botorch_mfkg_branin.py --nworkers 4 --comms tcp + +With ``--nworkers 4`` one worker is the generator and three concurrently +evaluate the objective. + +""" + +# Do not change these lines - they are parsed by run-tests.sh +# TESTSUITE_COMMS: local mpi +# TESTSUITE_NPROCS: 4 +# TESTSUITE_EXTRA: true +# TESTSUITE_OS_SKIP: OSX + +import numpy as np + +from libensemble import logger +from libensemble.alloc_funcs.start_only_persistent import only_persistent_gens +from libensemble.gen_funcs.persistent_botorch_mfkg_branin import persistent_botorch_mfkg +from libensemble.libE import libE +from libensemble.sim_funcs.augmented_branin import augmented_branin +from libensemble.tools import parse_args, save_libE_output + +LB = np.array([0.0, 0.0]) +UB = np.array([1.0, 1.0]) +N_INIT_SAMPLES = 4 # Each initial point gets a high- and a low-fidelity evaluation +Q = 4 # Batch size per MFKG iteration (keeps up to q sim workers busy) +SIM_MAX = 16 # Exit after running this many simulations + + +def run_mfkg(async_return, nworkers, is_manager, libE_specs): + """Run the MFKG ensemble once and check the results. + + ``async_return`` selects batch return (False) or asynchronous return + (True, with the generator in active-receive mode). + """ + sim_specs = { + "sim_f": augmented_branin, + "in": ["x", "fidelity"], + "out": [("f", float)], + } + + gen_specs = { + "gen_f": persistent_botorch_mfkg, + "persis_in": ["sim_id", "x", "f", "fidelity"], + "out": [ + ("x", float, (2,)), + ("fidelity", float), + ], + "user": { + "lb": LB, + "ub": UB, + "n_init_samples": N_INIT_SAMPLES, + "q": Q, + }, + } + + alloc_specs = { + "alloc_f": only_persistent_gens, + "user": { + # When async, return each evaluation to the generator as soon as it + # completes and let the generator receive while still active. When + # not async, the whole batch is returned in one message. + "async_return": async_return, + "active_recv_gen": async_return, + }, + } + + exit_criteria = {"sim_max": SIM_MAX} + + # Fresh RNG streams for each run. + persis_info = {} + + H, persis_info, flag = libE(sim_specs, gen_specs, exit_criteria, persis_info, alloc_specs, libE_specs) + + if is_manager: + save_libE_output(H, persis_info, __file__, nworkers) + + mode = "async" if async_return else "batch" + + # -- Sanity checks that the run behaved as expected -------------------- + ended = H["sim_ended"] + n_ended = int(np.sum(ended)) + + # libE finished cleanly. + assert flag == 0, f"[{mode}] libE returned a nonzero exit flag: {flag}" + + # We completed at least sim_max evaluations, and the MFKG loop produced + # points beyond the (2 * n_init_samples) initial sample. + assert n_ended >= SIM_MAX, f"[{mode}] Expected >= {SIM_MAX} completed sims, got {n_ended}" + assert n_ended > 2 * N_INIT_SAMPLES, f"[{mode}] MFKG did not generate points beyond the initial sample" + + # Every completed evaluation returned a finite objective value. + assert np.all(np.isfinite(H["f"][ended])), f"[{mode}] Found non-finite objective value(s)" + + # The generator only ever requests the two discrete fidelities (low=0, high=1). + fids = np.round(H["fidelity"][ended], 6) + assert np.all(np.isin(fids, [0.0, 1.0])), f"[{mode}] Unexpected fidelity values: {np.unique(fids)}" + + # Generated points stayed within [lb, ub]. + xs = H["x"][ended] + assert np.all(xs >= LB - 1e-9) and np.all(xs <= UB + 1e-9), f"[{mode}] Generated point(s) outside the bounds" + + # augmented_branin is negated for maximization; the best attainable value is + # ~ -0.397887 (the negated Branin global minimum). We must never exceed it. + best_f = float(np.max(H["f"][ended])) + assert best_f <= -0.397887 + 1e-3, f"[{mode}] Objective exceeded the known optimum: {best_f}" + assert best_f < 0.0, f"[{mode}] Best objective is unexpectedly non-negative: {best_f}" + + print(f"[{mode}] assertions passed: completed sims: {n_ended}, best f: {best_f:.4f}") + + +# Main block is necessary only when using local comms with spawn start method (default on macOS and Windows). +if __name__ == "__main__": + nworkers, is_manager, libE_specs, _ = parse_args() + + # libE logger + logger.set_level("INFO") + + # Exercise both batch and asynchronous return of results to the generator. + for async_return in [False, True]: + run_mfkg(async_return, nworkers, is_manager, libE_specs)