From 71339c37a5d96024fb8b529583e004082605847d Mon Sep 17 00:00:00 2001
From: shudson <shudson@anl.gov>
Date: Tue, 28 Oct 2025 13:21:04 -0500
Subject: [PATCH 1/7] 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 <jmlarson@anl.gov>
Date: Wed, 29 Oct 2025 10:33:32 -0500
Subject: [PATCH 2/7] 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 <jmlarson@anl.gov>
Date: Fri, 5 Dec 2025 08:48:59 -0600
Subject: [PATCH 3/7] 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,
         },
     }

From 2ff60bf73a960399a68338503102f44c2ce8c4a6 Mon Sep 17 00:00:00 2001
From: Jeffrey Larson <jmlarson@anl.gov>
Date: Tue, 7 Jul 2026 15:20:04 -0500
Subject: [PATCH 4/7] black

---
 .../persistent_botorch_mfkg_branin.py          | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/libensemble/gen_funcs/persistent_botorch_mfkg_branin.py b/libensemble/gen_funcs/persistent_botorch_mfkg_branin.py
index 5391fea12..5e52d6a00 100644
--- a/libensemble/gen_funcs/persistent_botorch_mfkg_branin.py
+++ b/libensemble/gen_funcs/persistent_botorch_mfkg_branin.py
@@ -55,12 +55,12 @@ def project(X):
 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
     """
@@ -69,15 +69,15 @@ def problem(X, ps, gen_specs):
     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
 
 
@@ -154,12 +154,12 @@ def do_iteration(train_x, train_obj, q, ps, gen_specs):
     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
+    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
 

From 8ff42c3aa4075528e12a19f0d6196ca024a44f24 Mon Sep 17 00:00:00 2001
From: Jeffrey Larson <jmlarson@anl.gov>
Date: Tue, 7 Jul 2026 15:30:34 -0500
Subject: [PATCH 5/7] black

---
 .../persistent_botorch_mfkg_branin.py         | 28 +++++++++----------
 libensemble/sim_funcs/augmented_branin.py     |  5 ++--
 .../run_botorch_mfkg_branin.py                |  3 +-
 3 files changed, 16 insertions(+), 20 deletions(-)

diff --git a/libensemble/gen_funcs/persistent_botorch_mfkg_branin.py b/libensemble/gen_funcs/persistent_botorch_mfkg_branin.py
index 5e52d6a00..5b592439d 100644
--- a/libensemble/gen_funcs/persistent_botorch_mfkg_branin.py
+++ b/libensemble/gen_funcs/persistent_botorch_mfkg_branin.py
@@ -34,10 +34,6 @@
     "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}
 
@@ -82,8 +78,8 @@ def problem(X, ps, gen_specs):
 
 
 # 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)
+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)
@@ -101,7 +97,7 @@ def initialize_model(train_x, train_obj):
 
 
 # Multifidelity Knowledge Gradient acquisition function
-def get_mfkg(model):
+def get_mfkg(model, bounds):
 
     curr_val_acqf = FixedFeatureAcquisitionFunction(
         acq_function=PosteriorMean(model),
@@ -129,7 +125,7 @@ def get_mfkg(model):
 
 
 # Optimization step
-def optimize_mfkg_and_get_observation(mfkg_acqf, q, ps, gen_specs):
+def optimize_mfkg_and_get_observation(mfkg_acqf, q, ps, gen_specs, bounds):
     # Generate new candidates
     candidates, _ = optimize_acqf_mixed(
         acq_function=mfkg_acqf,
@@ -149,11 +145,11 @@ def optimize_mfkg_and_get_observation(mfkg_acqf, q, ps, gen_specs):
 
 
 # Function to perform a single iteration
-def do_iteration(train_x, train_obj, q, ps, gen_specs):
+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)
-    new_x, new_obj, _, tag = optimize_mfkg_and_get_observation(mfkg_acqf, q, ps, gen_specs)
+    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
@@ -171,17 +167,19 @@ def persistent_botorch_mfkg(H, persis_info, gen_specs, libE_info):
     ps = PersistentSupport(libE_info, EVAL_GEN_TAG)
 
     # Extract user parameters
-    ub = gen_specs["user"]["ub"]
-    lb = gen_specs["user"]["lb"]
+    ub = np.asarray(gen_specs["user"]["ub"], dtype=float)
+    lb = np.asarray(gen_specs["user"]["lb"], dtype=float)
+    bounds = torch.tensor([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)
+    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)
+        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
index 5937db730..8aa1300cd 100644
--- a/libensemble/sim_funcs/augmented_branin.py
+++ b/libensemble/sim_funcs/augmented_branin.py
@@ -6,7 +6,6 @@
 
 __all__ = ["augmented_branin", "augmented_branin_func"]
 
-import math
 import numpy as np
 
 
@@ -31,8 +30,8 @@ def augmented_branin_func(x, fidelity):
     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)
+    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/run_botorch_mfkg_branin.py b/libensemble/tests/regression_tests/run_botorch_mfkg_branin.py
index be599ae70..506bbae9e 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,  # Each of these points will have a high-fidelity and low-fidelity evaluation 
+            "n_init_samples": 4,  # Each of these points will have a high-fidelity and low-fidelity evaluation
             "q": 2,
         },
     }
@@ -76,4 +76,3 @@
     # Save results to numpy file
     if is_manager:
         save_libE_output(H, persis_info, __file__, nworkers)
-

From ea1d255f1483f41715f9ce363e869e037490e841 Mon Sep 17 00:00:00 2001
From: Jeffrey Larson <jmlarson@anl.gov>
Date: Wed, 8 Jul 2026 07:37:31 -0500
Subject: [PATCH 6/7] Removing old options/imports

---
 .../tests/regression_tests/run_botorch_mfkg_branin.py        | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/libensemble/tests/regression_tests/run_botorch_mfkg_branin.py b/libensemble/tests/regression_tests/run_botorch_mfkg_branin.py
index 506bbae9e..2e220b968 100644
--- a/libensemble/tests/regression_tests/run_botorch_mfkg_branin.py
+++ b/libensemble/tests/regression_tests/run_botorch_mfkg_branin.py
@@ -27,12 +27,11 @@
 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
+from libensemble.tools import 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,
@@ -68,7 +67,7 @@
     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)
+    persis_info = {}
 
     # 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)

From dae7ff5e1e9883fdafb45c31ebc5ef56e13e30ad Mon Sep 17 00:00:00 2001
From: Jeffrey Larson <jmlarson@anl.gov>
Date: Thu, 16 Jul 2026 11:16:50 -0500
Subject: [PATCH 7/7] Hoping to get the async case tested

---
 .../persistent_botorch_mfkg_branin.py         |  56 +++++--
 .../run_botorch_mfkg_branin.py                |  77 ----------
 .../regression_tests/test_mfkg_branin.py      | 141 ++++++++++++++++++
 3 files changed, 182 insertions(+), 92 deletions(-)
 delete mode 100644 libensemble/tests/regression_tests/run_botorch_mfkg_branin.py
 create mode 100644 libensemble/tests/regression_tests/test_mfkg_branin.py

diff --git a/libensemble/gen_funcs/persistent_botorch_mfkg_branin.py b/libensemble/gen_funcs/persistent_botorch_mfkg_branin.py
index 5b592439d..2a4081c06 100644
--- a/libensemble/gen_funcs/persistent_botorch_mfkg_branin.py
+++ b/libensemble/gen_funcs/persistent_botorch_mfkg_branin.py
@@ -50,7 +50,15 @@ def project(X):
 # 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.
+    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]
@@ -58,23 +66,39 @@ def problem(X, ps, gen_specs):
         gen_specs: Generator specifications
 
     Returns:
-        tensor of shape (n,) with objective values
+        (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
+    # 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)
 
-    tag, Work, calc_in = ps.send_recv(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"]))
 
-    # Convert results back to tensor
-    if calc_in is None or len(calc_in) == 0:
-        return None, tag
+    if len(objs) == 0:
+        return None, None, tag
 
-    train_obj = torch.tensor(calc_in["f"], **tkwargs).unsqueeze(-1)
+    new_x = torch.tensor(np.array(xs), **tkwargs)
+    new_obj = torch.tensor(np.array(objs), **tkwargs).unsqueeze(-1)
 
-    return train_obj, tag
+    return new_x, new_obj, tag
 
 
 # Function to generate training data
@@ -85,8 +109,10 @@ def generate_initial_data(n, ps, gen_specs, bounds):  # Jeff: Initial sample siz
     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
+    # 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
@@ -137,10 +163,10 @@ def optimize_mfkg_and_get_observation(mfkg_acqf, q, ps, gen_specs, bounds):
         options={"batch_limit": 10, "maxiter": 10},  # Jeff: I decreased this to make libE development faster
     )
 
-    # Observe new values
+    # 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 = candidates.detach()
-    new_obj, tag = problem(new_x, ps, gen_specs)
+    new_x, new_obj, tag = problem(candidates.detach(), ps, gen_specs)
     return new_x, new_obj, cost, tag
 
 
@@ -169,7 +195,7 @@ def persistent_botorch_mfkg(H, persis_info, gen_specs, libE_info):
     # 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.append(lb, 0.0), np.append(ub, 1.0)], **tkwargs)
+    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"]
diff --git a/libensemble/tests/regression_tests/run_botorch_mfkg_branin.py b/libensemble/tests/regression_tests/run_botorch_mfkg_branin.py
deleted file mode 100644
index 2e220b968..000000000
--- a/libensemble/tests/regression_tests/run_botorch_mfkg_branin.py
+++ /dev/null
@@ -1,77 +0,0 @@
-"""
-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 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()
-
-    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,  # Each of these points will have a high-fidelity and low-fidelity evaluation
-            "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 = {}
-
-    # 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)
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)