From 71339c37a5d96024fb8b529583e004082605847d Mon Sep 17 00:00:00 2001 From: shudson Date: Tue, 28 Oct 2025 13:21:04 -0500 Subject: [PATCH 1/3] Add botorch_mfkg use-case. Works in batch. May need to update for async. --- .../persistent_botorch_mfkg_branin.py | 187 ++++++++++++++++++ libensemble/sim_funcs/augmented_branin.py | 38 ++++ .../run_botorch_mfkg_branin.py | 79 ++++++++ 3 files changed, 304 insertions(+) create mode 100644 libensemble/gen_funcs/persistent_botorch_mfkg_branin.py create mode 100644 libensemble/sim_funcs/augmented_branin.py create mode 100644 libensemble/tests/regression_tests/run_botorch_mfkg_branin.py 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..5391fea12 --- /dev/null +++ b/libensemble/gen_funcs/persistent_botorch_mfkg_branin.py @@ -0,0 +1,187 @@ +""" +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 bounds +dim = 3 +bounds = torch.tensor([[0.0] * dim, [1.0] * dim], **tkwargs) + +# 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): + """ + Wrapper to convert tensor input to numpy and send to libE for evaluation. + + Args: + X: tensor of shape (n, 3) where columns are [x0, x1, fidelity] + ps: PersistentSupport object for communication + gen_specs: Generator specifications + + Returns: + tensor of shape (n,) with objective values + """ + # 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] + + tag, Work, calc_in = ps.send_recv(H_o) + + # Convert results back to tensor + if calc_in is None or len(calc_in) == 0: + return None, tag + + train_obj = torch.tensor(calc_in["f"], **tkwargs).unsqueeze(-1) + + return train_obj, tag + + +# Function to generate training data +def generate_initial_data(n, ps, gen_specs): # Jeff: Initial sample size is twice this value of n + train_x = 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) + train_obj, tag = problem(train_x_full, ps, gen_specs) + return train_x_full, 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): + + 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): + # 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 + cost = cost_model(candidates).sum() + new_x = candidates.detach() + new_obj, tag = problem(new_x, 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): + mll, model = initialize_model(train_x, train_obj) + fit_gpytorch_mll(mll) + mfkg_acqf = get_mfkg(model) + new_x, new_obj, _, tag = optimize_mfkg_and_get_observation(mfkg_acqf, q, ps, gen_specs) + + 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" evaluation happens, and needs to be communicated back to 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 = gen_specs["user"]["ub"] + lb = gen_specs["user"]["lb"] + 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) + + # 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) + + 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..5937db730 --- /dev/null +++ b/libensemble/sim_funcs/augmented_branin.py @@ -0,0 +1,38 @@ +""" +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 math +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 * math.pi**2) - 0.1 * (1 - fidelity)) * (15 * x0 - 5) ** 2 + 5 / math.pi * (15 * x0 - 5) - 6 + t2 = 10 * (1 - 1 / (8 * math.pi)) * np.cos(15 * x0 - 5) + result = t1**2 + t2 + 10 + + return -result # negate for maximization diff --git a/libensemble/tests/regression_tests/run_botorch_mfkg_branin.py b/libensemble/tests/regression_tests/run_botorch_mfkg_branin.py new file mode 100644 index 000000000..b65eb13ca --- /dev/null +++ b/libensemble/tests/regression_tests/run_botorch_mfkg_branin.py @@ -0,0 +1,79 @@ +""" +Example of multi-fidelity optimization using a persistent BoTorch MFKG gen_func. + +This test uses the gen_on_manager option (persistent generator runs on +a thread). Therefore nworkers is the number of simulation workers. + +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 + +When running with the above commands, the number of concurrent evaluations of +the objective function will be 3. + +""" + +# 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 add_unique_random_streams, parse_args, save_libE_output + +# 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_specs["gen_on_manager"] = True + + sim_specs = { + "sim_f": augmented_branin, + "in": ["x", "fidelity"], + "out": [("f", float)], + } + + gen_specs = { + "gen_f": persistent_botorch_mfkg, + # "in": ["sim_id", "x", "f", "fidelity"], + "persis_in": ["sim_id", "x", "f", "fidelity"], + "out": [ + ("x", float, (2,)), + ("fidelity", float), + ], + "user": { + "lb": np.array([0.0, 0.0]), + "ub": np.array([1.0, 1.0]), + "n_init_samples": 4, + "q": 2, + }, + } + + alloc_specs = { + "alloc_f": only_persistent_gens, + "user": {"async_return": False}, + } + + # libE logger + logger.set_level("INFO") + + # Exit criteria + exit_criteria = {"sim_max": 12} # Exit after running sim_max simulations + + # Create a different random number stream for each worker and the manager + persis_info = add_unique_random_streams({}, nworkers + 1) + + # Run LibEnsemble, and store results in history array H + H, persis_info, flag = libE(sim_specs, gen_specs, exit_criteria, persis_info, alloc_specs, libE_specs) + + # Save results to numpy file + if is_manager: + save_libE_output(H, persis_info, __file__, nworkers) + From 50ee0dfab773ba97f8dd0aa8c2d955e48e458f39 Mon Sep 17 00:00:00 2001 From: Jeffrey Larson Date: Wed, 29 Oct 2025 10:33:32 -0500 Subject: [PATCH 2/3] Noting Ax must be 0.5.0 --- libensemble/gen_funcs/persistent_ax_multitask.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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) From a041af921291c33d8cab9042d5f08ea1bca97866 Mon Sep 17 00:00:00 2001 From: Jeffrey Larson Date: Fri, 5 Dec 2025 08:48:59 -0600 Subject: [PATCH 3/3] adding note --- libensemble/tests/regression_tests/run_botorch_mfkg_branin.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libensemble/tests/regression_tests/run_botorch_mfkg_branin.py b/libensemble/tests/regression_tests/run_botorch_mfkg_branin.py index b65eb13ca..be599ae70 100644 --- a/libensemble/tests/regression_tests/run_botorch_mfkg_branin.py +++ b/libensemble/tests/regression_tests/run_botorch_mfkg_branin.py @@ -51,7 +51,7 @@ "user": { "lb": np.array([0.0, 0.0]), "ub": np.array([1.0, 1.0]), - "n_init_samples": 4, + "n_init_samples": 4, # Each of these points will have a high-fidelity and low-fidelity evaluation "q": 2, }, }