diff --git a/.claude/skills/generate-scripts/references/results_metadata.md b/.claude/skills/generate-scripts/references/results_metadata.md index 97a71c14b..2a35a095e 100644 --- a/.claude/skills/generate-scripts/references/results_metadata.md +++ b/.claude/skills/generate-scripts/references/results_metadata.md @@ -39,4 +39,3 @@ If the minimum objective value is exactly 0.0, check whether those rows have `sim_ended == True`. Unevaluated rows often have fields initialized to zero. This is common for the last few rows when the simulation budget is exhausted — they were allocated by the generator but never evaluated. - diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 69c11918b..de0438650 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -37,4 +37,4 @@ repos: rev: v1.19.1 hooks: - id: mypy - exclude: ^docs/conf\.py$|libensemble/utils/(launcher|loc_stack|runners|pydantic|output_directory)\.py$|libensemble/tests/(regression_tests|functionality_tests|unit_tests|scaling_tests)/.* + pass_filenames: false diff --git a/libensemble/gen_classes/aposmm.py b/libensemble/gen_classes/aposmm.py index 0ecad8968..00a29cb9c 100644 --- a/libensemble/gen_classes/aposmm.py +++ b/libensemble/gen_classes/aposmm.py @@ -1,17 +1,18 @@ import copy import warnings from math import gamma, pi, sqrt -from typing import Any, Dict, List, Optional +from typing import Any import numpy as np from gest_api.vocs import VOCS from numpy import typing as npt -from libensemble.generators import PersistentGenInterfacer -from libensemble.message_numbers import EVAL_GEN_TAG, PERSIS_STOP +from libensemble.generators import LibensembleGenerator +from libensemble.message_numbers import FINISHED_PERSISTENT_GEN_TAG +from libensemble.utils.misc import np_to_list_dicts, unmap_numpy_array -class APOSMM(PersistentGenInterfacer): +class APOSMM(LibensembleGenerator): """ APOSMM coordinates multiple local optimization runs, dramatically reducing time for discovering multiple minima on parallel systems. @@ -175,27 +176,44 @@ def __init__( max_active_runs: int, initial_sample_size: int, History: npt.NDArray = [], - sample_points: Optional[npt.NDArray] = None, + sample_points: npt.NDArray = None, localopt_method: str = "scipy_Nelder-Mead", - rk_const: Optional[float] = None, + rk_const: float | None = None, xtol_abs: float = 1e-6, ftol_abs: float = 1e-6, opt_return_codes: list[int] = [0], mu: float = 1e-8, nu: float = 1e-8, - dist_to_bound_multiple: float = 0.05, + dist_to_bound_multiple: float = 0.5, random_seed: int = 1, **kwargs, ) -> None: - from libensemble.gen_funcs.persistent_aposmm import aposmm + from libensemble.gen_funcs.aposmm_localopt_support import LocalOptInterfacer + from libensemble.gen_funcs.persistent_aposmm import ( + add_k_sample_points_to_local_H, + add_to_local_H, + decide_where_to_start_localopt, + initialize_APOSMM, + initialize_children, + initialize_dists_and_inds, + update_history_dist, + update_history_optimal, + ) + + # Store references to the functions we'll call later + self._add_k_sample_points = add_k_sample_points_to_local_H + self._add_to_local_H = add_to_local_H + self._decide_where_to_start = decide_where_to_start_localopt + self._initialize_dists_and_inds = initialize_dists_and_inds + self._update_history_dist = update_history_dist + self._update_history_optimal = update_history_optimal + self._LocalOptInterfacer = LocalOptInterfacer self.vocs = vocs - gen_specs: Dict[str, Any] = {} - gen_specs["user"] = {} - libE_info: Dict[str, Any] = {} - gen_specs["gen_f"] = aposmm + gen_specs: dict[str, Any] = {"user": {}} + persis_info: dict[str, Any] = {} n = len(list(vocs.variables.keys())) if not rk_const: @@ -221,7 +239,10 @@ def __init__( if val is not None: gen_specs["user"][k] = val - super().__init__(vocs, History, {}, gen_specs, libE_info, **kwargs) + super().__init__(vocs, History, persis_info, gen_specs, {}, **kwargs) + + # APOSMM manages sim_id internally — don't remap to _id + self.variables_mapping.pop("sim_id", None) # Set bounds using the correct x mapping x_mapping = self.variables_mapping["x"] @@ -267,118 +288,317 @@ def __init__( if "components" in kwargs or "components" in gen_specs.get("user", {}): gen_specs["persis_in"].append("fvec") - # SH - Need to know if this is gen_on_manager or not. - self.persis_info["nworkers"] = gen_specs["user"].get("max_active_runs") - self.all_local_minima: List[npt.NDArray] = [] - self._suggest_idx = 0 - self._last_suggest: Optional[npt.NDArray] = None - self._ingest_buf: Optional[npt.NDArray] = None - self._n_buffd_results = 0 + # Initialize APOSMM internal state directly (no subprocess) + user_specs = gen_specs["user"] + libE_info: dict[str, Any] = {"comm": []} # no comm needed in direct mode + self._n, self._n_s, self._rk_const, self._ld, self._mu, self._nu, _, self.local_H = initialize_APOSMM( + History, user_specs, libE_info + ) + ( + self._local_opters, + self._sim_id_to_child_inds, + self._run_order, + self._run_pts, + self._total_runs, + self._ended_runs, + self._fields_to_pass, + ) = initialize_children(user_specs) + + self._user_specs = user_specs + self._max_active_runs = max_active_runs + self._rng = np.random.default_rng(random_seed) + + # Build reverse mapping: VOCS field name -> (internal_name, index) + self._reverse_mapping = {} + for internal_name, vocs_names in self.variables_mapping.items(): + for i, vocs_name in enumerate(vocs_names): + self._reverse_mapping[vocs_name] = (internal_name, i, len(vocs_names)) + + self.all_local_minima: list[npt.NDArray] = [] self._told_initial_sample = False - self._first_called_method: Optional[str] = None - self._last_call: Optional[str] = None - self._last_num_points = 0 + self._first_called_method: str | None = None + self._pending_results = None + self._first_pass = True + self._n_r = 0 # number of results received in last ingest + self._initial_sample_generated = False + self._initial_suggest_idx = 0 # tracks how many initial sample points have been handed out + self.gen_result: tuple[npt.NDArray | list | None, dict | None, int | None] | None = None + + def _map_to_internal(self, results): + """Map VOCS-named structured array to internal APOSMM field names (x, x_on_cube, f, sim_id).""" + if results is None or len(results) == 0: + return results + # If already has internal names, return as-is + if "x" in results.dtype.names and "f" in results.dtype.names: + return results + + n_rows = len(results) + # Build dtype for internal array + internal_fields = [] + added = set() + for vocs_name in results.dtype.names: + if vocs_name in self._reverse_mapping: + internal_name, _, size = self._reverse_mapping[vocs_name] + if internal_name not in added: + if size > 1: + internal_fields.append((internal_name, float, size)) + else: + internal_fields.append((internal_name, float)) + added.add(internal_name) + elif vocs_name == "_id": + if "sim_id" not in added: + internal_fields.append(("sim_id", int)) + added.add("sim_id") + elif vocs_name == "sim_id": + if "sim_id" not in added: + internal_fields.append(("sim_id", int)) + added.add("sim_id") + + out = np.zeros(n_rows, dtype=internal_fields) + has_sim_id = "sim_id" in results.dtype.names + for vocs_name in results.dtype.names: + if vocs_name in self._reverse_mapping: + internal_name, idx, size = self._reverse_mapping[vocs_name] + if size > 1: + out[internal_name][:, idx] = results[vocs_name] + else: + out[internal_name] = results[vocs_name] + elif vocs_name == "sim_id": + out["sim_id"] = results["sim_id"] + elif vocs_name == "_id" and not has_sim_id: + out["sim_id"] = results["_id"] + + return out def _slot_in_data(self, results): - """Slot in libE_calc_in and trial data into corresponding array fields. *Initial sample only!!*""" - for name in results.dtype.names: - if name == "_id": - self._ingest_buf["sim_id"][self._n_buffd_results : self._n_buffd_results + len(results)] = results[ - "_id" - ] - else: - self._ingest_buf[name][self._n_buffd_results : self._n_buffd_results + len(results)] = results[name] - - def _enough_initial_sample(self): - return ( - self._n_buffd_results >= int(self.gen_specs["user"]["initial_sample_size"]) - ) or self._told_initial_sample - - def _ready_to_suggest_genf(self): - """ - We're presumably ready to be suggested IF: - - When we're working on the initial sample: - - We have no _last_suggest cached - - all points given out have returned AND we've been suggested *at least* as many points as we cached - - When we're done with the initial sample: - - we've been suggested *at least* as many points as we cached - - we've just ingested some results - """ - if not self._told_initial_sample and self._last_suggest is not None: - cond = all([i in self._ingest_buf["sim_id"] for i in self._last_suggest["sim_id"]]) - else: - cond = True - return self._last_suggest is None or (cond and (self._suggest_idx >= len(self._last_suggest))) - - def suggest_numpy(self, num_points: int = 0) -> npt.NDArray: + """Slot ingested results into local_H during initial sample phase.""" + n_s_before = self._n_s + n_new = len(results) + old_len = len(self.local_H) + needed = n_s_before + n_new + if needed > old_len: + self.local_H.resize(needed, refcheck=False) + self._initialize_dists_and_inds(self.local_H, needed - old_len) + + for i, row in enumerate(results): + idx = n_s_before + i + self.local_H["sim_id"][idx] = idx + for name in results.dtype.names: + if name == "sim_id": + continue + if name in self.local_H.dtype.names: + self.local_H[name][idx] = row[name] + self.local_H["sim_ended"][idx] = True + self._n_s += n_new + self._update_history_dist(self.local_H, self._n) + + def suggest_numpy(self, num_points: int | None = 0) -> npt.NDArray: """Request the next set of points to evaluate, as a NumPy array.""" + out_fields = [i[0] for i in self.gen_specs["out"]] + num_points = num_points or 0 if self._first_called_method is None: self._first_called_method = "suggest" - self.gen_specs["user"]["generate_sample_points"] = True - if self._ready_to_suggest_genf(): - self._suggest_idx = 0 - if self._last_call == "suggest" and num_points == 0 and self._last_num_points == 0: - self.finalize() + # Initial sample phase: generate random points once, return in batches + if not self._told_initial_sample: + if not self._initial_sample_generated: + total = self._user_specs["initial_sample_size"] + self._add_k_sample_points( + total, + self._user_specs, + self.persis_info, + self._n, + [], + self.local_H, + self._sim_id_to_child_inds, + self._rng, + ) + self._initial_sample_generated = True + self._initial_suggest_idx = 0 + + if self._initial_suggest_idx >= self._user_specs["initial_sample_size"]: raise RuntimeError("Cannot suggest points since APOSMM is currently expecting to receive a sample") - self._last_suggest = super().suggest_numpy(num_points) - assert self._last_suggest is not None - if self._last_suggest["local_min"].any(): # filter out local minima rows - min_idxs = self._last_suggest["local_min"] - self.all_local_minima.append(self._last_suggest[min_idxs]) - self._last_suggest = self._last_suggest[~min_idxs] + k = num_points if num_points > 0 else (self._user_specs["initial_sample_size"] - self._initial_suggest_idx) + start = self._initial_suggest_idx + end = min(start + k, self._user_specs["initial_sample_size"]) + result = self.local_H[start:end][out_fields].copy() + self._initial_suggest_idx = end + return unmap_numpy_array(result, self.variables_mapping) + + # Main optimization phase + new_opt_inds: list[int] = [] + new_inds: list[int] = [] + + # Process any pending ingested results through local optimizers + if self._pending_results is not None: + from libensemble.gen_funcs.aposmm_localopt_support import ConvergedMsg + + calc_in = self._pending_results + self._pending_results = None + + # Update local_H with received results + for row in calc_in: + sim_id = int(row["sim_id"]) + self.local_H[sim_id]["sim_ended"] = True + for name in calc_in.dtype.names: + if name in self.local_H.dtype.names: + self.local_H[name][sim_id] = row[name] + self._n_s = int(np.sum(~self.local_H["local_pt"][: len(self.local_H)])) + self._update_history_dist(self.local_H, self._n) + + for row in calc_in: + sim_id = int(row["sim_id"]) + if self._sim_id_to_child_inds.get(sim_id): + for child_idx in self._sim_id_to_child_inds[sim_id]: + if child_idx not in self._local_opters: + continue + x_new = self._local_opters[child_idx].iterate(row[self._fields_to_pass]) + if isinstance(x_new, ConvergedMsg): + x_opt = x_new.x + opt_flag = x_new.opt_flag + opt_ind = self._update_history_optimal( + x_opt, + opt_flag, + self.local_H, + self._run_order[child_idx], + ) + new_opt_inds.append(opt_ind) + self._local_opters.pop(child_idx) + self._ended_runs.append(child_idx) + else: + self._add_to_local_H(self.local_H, x_new, self._user_specs, local_flag=1, on_cube=True) + new_inds.append(len(self.local_H) - 1) + self._run_order[child_idx].append(self.local_H[-1]["sim_id"]) + self._run_pts[child_idx].append(x_new) + sid = self.local_H[-1]["sim_id"] + if sid in self._sim_id_to_child_inds: + self._sim_id_to_child_inds[sid] += (child_idx,) + else: + self._sim_id_to_child_inds[sid] = (child_idx,) + + # Decide where to start new local optimization runs + starting_inds = self._decide_where_to_start( + self.local_H, + self._n, + self._n_s, + self._rk_const, + self._ld, + self._mu, + self._nu, + ) + + for ind in starting_inds: + if len([p for p in self._local_opters.values() if p.is_running]) < self._max_active_runs: + self.local_H["started_run"][ind] = 1 + local_opter = self._LocalOptInterfacer( + self._user_specs, + self.local_H[ind]["x_on_cube"], + self.local_H[ind]["f"] if "f" in self._fields_to_pass else self.local_H[ind]["fvec"], + self.local_H[ind]["grad"] if "grad" in self._fields_to_pass else None, + ) + self._local_opters[self._total_runs] = local_opter + x_new = local_opter.iterate(self.local_H[ind][self._fields_to_pass]) + self._add_to_local_H(self.local_H, x_new, self._user_specs, local_flag=1, on_cube=True) + new_inds.append(len(self.local_H) - 1) + self._run_order[self._total_runs] = [ind, self.local_H[-1]["sim_id"]] + self._run_pts[self._total_runs] = [self.local_H["x_on_cube"], x_new] + sid = self.local_H[-1]["sim_id"] + if sid in self._sim_id_to_child_inds: + self._sim_id_to_child_inds[sid] += (self._total_runs,) + else: + self._sim_id_to_child_inds[sid] = (self._total_runs,) + self._total_runs += 1 + + # Fill remaining slots with sample points + if self._first_pass: + num_samples = self._max_active_runs - 1 - len(new_inds) + self._first_pass = False + else: + num_samples = self._n_r - len(new_inds) + + if num_samples > 0: + self._add_k_sample_points( + num_samples, + self._user_specs, + self.persis_info, + self._n, + [], + self.local_H, + self._sim_id_to_child_inds, + self._rng, + ) + new_inds = new_inds + list(range(len(self.local_H) - num_samples, len(self.local_H))) - if num_points > 0: # we've been suggested for a selection of the last suggest - assert self._last_suggest is not None - results = np.copy(self._last_suggest[self._suggest_idx : self._suggest_idx + num_points]) - self._suggest_idx += num_points + all_inds = new_inds + new_opt_inds + if len(all_inds) == 0: + return np.zeros(0, dtype=[(name, self.local_H.dtype[name]) for name in out_fields]) - else: - results = np.copy(self._last_suggest) - self._last_suggest = None + result = self.local_H[all_inds][out_fields].copy() - self._last_call = "suggest" - self._last_num_points = num_points - return results + # Track local minima for suggest_updates() + if result["local_min"].any(): + min_idxs = result["local_min"] + self.all_local_minima.append(result[min_idxs].copy()) - def ingest_numpy(self, results: npt.NDArray, tag: int = EVAL_GEN_TAG) -> None: + return unmap_numpy_array(result, self.variables_mapping) - if self._first_called_method is None: - self._first_called_method = "ingest" - self.gen_specs["user"]["generate_sample_points"] = False + def ingest_numpy(self, results: npt.NDArray, tag: int = 0) -> None: + """Send the results of evaluations to the generator.""" - if (results is None and tag == PERSIS_STOP) or self._told_initial_sample: - super().ingest_numpy(results, tag) - self._last_call = "ingest" + if results is None: return - # Initial sample buffering here: - - if self._n_buffd_results == 0: - # Create a dtype that includes sim_id but excludes _id - descr = [d for d in results.dtype.descr if d[0] != "_id"] - if "sim_id" not in [d[0] for d in descr]: - descr.append(("sim_id", int)) - self._ingest_buf = np.zeros(self.gen_specs["user"]["initial_sample_size"], dtype=descr) + if self._first_called_method is None: + self._first_called_method = "ingest" - if not self._enough_initial_sample(): - self._slot_in_data(np.copy(results)) - self._n_buffd_results += len(results) + results = self._map_to_internal(results) - if self._enough_initial_sample(): - assert self._ingest_buf is not None - if "sim_id" in results.dtype.names and not self._told_initial_sample: - self._ingest_buf["sim_id"] = range(len(self._ingest_buf)) - super().ingest_numpy(self._ingest_buf, tag) - self._told_initial_sample = True - self._n_buffd_results = 0 + if not self._told_initial_sample: + # Initial sample phase: slot data into local_H + self._slot_in_data(results) + if self._n_s >= self._user_specs["initial_sample_size"]: + self._told_initial_sample = True + return - self._last_call = "ingest" + # Main phase: buffer results for processing in next suggest call + self._n_r = len(results) + self._pending_results = results.copy() - def suggest_updates(self) -> List[npt.NDArray]: + def suggest_updates(self) -> list[npt.NDArray]: """Request a list of NumPy arrays containing entries that have been identified as minima.""" minima = copy.deepcopy(self.all_local_minima) self.all_local_minima = [] return minima + + def setup(self) -> None: + """Reject legacy setup calls; direct APOSMM initializes in its constructor.""" + raise RuntimeError("Direct APOSMM does not support setup().") + + def finalize(self) -> None: + """Stop all local optimizer processes.""" + if self._first_called_method is None: + raise RuntimeError("Generator has not been started.") + for _, p in self._local_opters.items(): + p.destroy() + self._local_opters.clear() + self.persis_info["run_order"] = self._run_order + self.gen_result = (self.local_H, self.persis_info, FINISHED_PERSISTENT_GEN_TAG) + + def export( + self, vocs_field_names: bool = False, as_dicts: bool = False + ) -> tuple[npt.NDArray | list | None, dict | None, int | None]: + """Return the APOSMM history, persistent information, and exit tag.""" + if self.gen_result is None: + return (None, None, None) + + local_H, persis_info, tag = self.gen_result + if vocs_field_names and local_H is not None and self.variables_mapping: + local_H = unmap_numpy_array(local_H, self.variables_mapping) + if as_dicts and local_H is not None: + if vocs_field_names and self.variables_mapping: + local_H = np_to_list_dicts(local_H, self.variables_mapping) + else: + local_H = np_to_list_dicts(local_H) + return (local_H, persis_info, tag) diff --git a/libensemble/gen_classes/sampling.py b/libensemble/gen_classes/sampling.py index 0b0662448..ff097105a 100644 --- a/libensemble/gen_classes/sampling.py +++ b/libensemble/gen_classes/sampling.py @@ -156,9 +156,7 @@ class UniformSampleWithVariableResources(LibensembleGenerator): path was tested with the default alloc. """ - def __init__( - self, vocs: VOCS, max_resource_sets: int, random_seed: int = 1, *args, **kwargs - ): + def __init__(self, vocs: VOCS, max_resource_sets: int, random_seed: int = 1, *args, **kwargs): super().__init__(vocs, *args, **kwargs) self.rng = np.random.default_rng(random_seed) self.max_rsets = max_resource_sets diff --git a/libensemble/tests/regression_tests/test_aposmm_nlopt.py b/libensemble/tests/regression_tests/test_aposmm_nlopt.py index 40ebcc497..65f1583da 100644 --- a/libensemble/tests/regression_tests/test_aposmm_nlopt.py +++ b/libensemble/tests/regression_tests/test_aposmm_nlopt.py @@ -92,12 +92,13 @@ def six_hump_camel_func(x): H, _, _ = workflow.run() if workflow.is_manager: - print("[Manager]:", H[np.where(H["local_min"])]["x"]) + x_min = np.column_stack([H[H["local_min"]]["x0"], H[H["local_min"]]["x1"]]) + print("[Manager]:", x_min) print("[Manager]: Time taken =", time() - start_time, flush=True) tol = 1e-5 for m in minima: # The minima are known on this test problem. # We use their values to test APOSMM has identified all minima - print(np.min(np.sum((H[H["local_min"]]["x"] - m) ** 2, 1)), flush=True) - assert np.min(np.sum((H[H["local_min"]]["x"] - m) ** 2, 1)) < tol + print(np.min(np.sum((x_min - m) ** 2, 1)), flush=True) + assert np.min(np.sum((x_min - m) ** 2, 1)) < tol diff --git a/libensemble/tests/regression_tests/test_aposmm_scipy.py b/libensemble/tests/regression_tests/test_aposmm_scipy.py index 7a24588d2..10fc2d383 100644 --- a/libensemble/tests/regression_tests/test_aposmm_scipy.py +++ b/libensemble/tests/regression_tests/test_aposmm_scipy.py @@ -89,7 +89,8 @@ def six_hump_camel_func(x): H, _, _ = workflow.run() if workflow.is_manager: - print("[Manager]:", H[np.where(H["local_min"])]["x"]) + x_min = np.column_stack([H[H["local_min"]]["x0"], H[H["local_min"]]["x1"]]) + print("[Manager]:", x_min) print("[Manager]: Time taken =", time() - start_time, flush=True) tol = 1e-3 @@ -97,7 +98,7 @@ def six_hump_camel_func(x): for m in minima: # The minima are known on this test problem. # We use their values to test APOSMM has identified all minima - print(np.min(np.sum((H[H["local_min"]]["x"] - m) ** 2, 1)), flush=True) - if np.min(np.sum((H[H["local_min"]]["x"] - m) ** 2, 1)) < tol: + print(np.min(np.sum((x_min - m) ** 2, 1)), flush=True) + if np.min(np.sum((x_min - m) ** 2, 1)) < tol: min_found += 1 assert min_found >= 2, f"Found {min_found} minima" diff --git a/libensemble/tests/unit_tests/test_persistent_aposmm.py b/libensemble/tests/unit_tests/test_persistent_aposmm.py index 8e777af32..6c31540aa 100644 --- a/libensemble/tests/unit_tests/test_persistent_aposmm.py +++ b/libensemble/tests/unit_tests/test_persistent_aposmm.py @@ -470,7 +470,7 @@ def test_asktell_ingest_first(): @pytest.mark.extra def test_asktell_consecutive_during_sample(): - """Test consecutive suggest and ingest during sample""" + """Test consecutive ingest during sampling""" from gest_api.vocs import VOCS @@ -503,12 +503,11 @@ def test_asktell_consecutive_during_sample(): dist_to_bound_multiple=0.01, ) - # Test consecutive suggest first = my_APOSMM.suggest(1) first[0]["energy"] = six_hump_camel_func(np.array([first[0]["core"], first[0]["edge"]])) my_APOSMM.ingest(first) - second = my_APOSMM.suggest(1) - second += my_APOSMM.suggest(4) + second = my_APOSMM.suggest(5) + for point in second: point["energy"] = six_hump_camel_func(np.array([point["core"], point["edge"]])) # Test consecutive ingest @@ -522,8 +521,8 @@ def test_asktell_consecutive_during_sample(): while total_evals < eval_max: - sample, detected_minima = my_APOSMM.suggest(3), my_APOSMM.suggest_updates() - sample += my_APOSMM.suggest(3) + sample = my_APOSMM.suggest(6) + detected_minima = my_APOSMM.suggest_updates() if len(detected_minima): for m in detected_minima: potential_minima.append(m) diff --git a/libensemble/tools/test_support.py b/libensemble/tools/test_support.py index bad9c2a78..7e0727125 100644 --- a/libensemble/tools/test_support.py +++ b/libensemble/tools/test_support.py @@ -244,9 +244,9 @@ def check_gpu_setting(task, assert_setting=True, print_setting=False, resources= if assert_setting: if isinstance(expected, dict): for key, value in expected.items(): - assert key in gpu_setting, ( - f"Worker {task.workerID}: Expected env key '{key}' not found in GPU setting: {gpu_setting}" - ) + assert ( + key in gpu_setting + ), f"Worker {task.workerID}: Expected env key '{key}' not found in GPU setting: {gpu_setting}" assert gpu_setting[key] == value, ( f"Worker {task.workerID}: GPU setting key '{key}' has value '{gpu_setting[key]}', " f"expected '{value}'" diff --git a/libensemble/utils/runners.py b/libensemble/utils/runners.py index 45f99435d..0f3ae16b8 100644 --- a/libensemble/utils/runners.py +++ b/libensemble/utils/runners.py @@ -200,20 +200,22 @@ def _result(self, calc_in: npt.NDArray, persis_info: dict, libE_info: dict) -> ( class LibensembleGenRunner(StandardGenRunner): - def _get_initial_suggest(self, libE_info) -> npt.NDArray: - """Get initial batch from a LibensembleGenerator. + def _get_mapping_for_outputs(self) -> dict: + """Return mappings whose internal fields are declared generator outputs.""" + output_names = {field[0] for field in self.specs.get("out", [])} + return { + name: fields for name, fields in getattr(self.gen, "variables_mapping", {}).items() if name in output_names + } - LibensembleGenerator.suggest_numpy emits VOCS-field-named structured arrays - (e.g. x0/x1, energy). The manager-side history expects mapped fields (x, f) - unless the user explicitly requested otherwise. - """ + def _get_initial_suggest(self, libE_info) -> npt.NDArray: + """Get initial batch from a LibensembleGenerator.""" initial_batch = self.specs.get("initial_batch_size") or self.specs.get("batch_size") or libE_info["batch_size"] H_out = self.gen.suggest_numpy(initial_batch) - return map_numpy_array(H_out, mapping=getattr(self.gen, "variables_mapping", {})) + return map_numpy_array(H_out, mapping=self._get_mapping_for_outputs()) def _get_points_updates(self, batch_size: int) -> (npt.NDArray, list): numpy_out = self.gen.suggest_numpy(batch_size) - numpy_out = map_numpy_array(numpy_out, mapping=getattr(self.gen, "variables_mapping", {})) + numpy_out = map_numpy_array(numpy_out, mapping=self._get_mapping_for_outputs()) if callable(getattr(self.gen, "suggest_updates", None)): updates = self.gen.suggest_updates() else: @@ -221,10 +223,10 @@ def _get_points_updates(self, batch_size: int) -> (npt.NDArray, list): return numpy_out, updates def _convert_ingest(self, x: npt.NDArray) -> list: - self.gen.ingest_numpy(unmap_numpy_array(x, mapping=getattr(self.gen, "variables_mapping", {}))) + self.gen.ingest_numpy(unmap_numpy_array(x, mapping=self._get_mapping_for_outputs())) def _convert_initial_ingest(self, x: npt.NDArray) -> list: - self.gen.ingest_numpy(unmap_numpy_array(x, mapping=getattr(self.gen, "variables_mapping", {}))) + self.gen.ingest_numpy(unmap_numpy_array(x, mapping=self._get_mapping_for_outputs())) class LibensembleGenThreadRunner(StandardGenRunner):