axbo_extensions 1.4.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- axbo_extensions-1.4.0/LICENSE +21 -0
- axbo_extensions-1.4.0/PKG-INFO +32 -0
- axbo_extensions-1.4.0/README.md +10 -0
- axbo_extensions-1.4.0/axbo_extensions/__init__.py +12 -0
- axbo_extensions-1.4.0/axbo_extensions/dimreduction.py +121 -0
- axbo_extensions-1.4.0/axbo_extensions/multifidelity.py +267 -0
- axbo_extensions-1.4.0/axbo_extensions/regionscreening.py +116 -0
- axbo_extensions-1.4.0/axbo_extensions/robustness.py +558 -0
- axbo_extensions-1.4.0/axbo_extensions/saasbo.py +33 -0
- axbo_extensions-1.4.0/axbo_extensions/trialdeps.py +182 -0
- axbo_extensions-1.4.0/axbo_extensions.egg-info/PKG-INFO +32 -0
- axbo_extensions-1.4.0/axbo_extensions.egg-info/SOURCES.txt +22 -0
- axbo_extensions-1.4.0/axbo_extensions.egg-info/dependency_links.txt +1 -0
- axbo_extensions-1.4.0/axbo_extensions.egg-info/requires.txt +11 -0
- axbo_extensions-1.4.0/axbo_extensions.egg-info/top_level.txt +1 -0
- axbo_extensions-1.4.0/pyproject.toml +45 -0
- axbo_extensions-1.4.0/setup.cfg +4 -0
- axbo_extensions-1.4.0/tests/test_dimreduction.py +23 -0
- axbo_extensions-1.4.0/tests/test_multifidelity.py +282 -0
- axbo_extensions-1.4.0/tests/test_placeholder.py +7 -0
- axbo_extensions-1.4.0/tests/test_regionscreening.py +48 -0
- axbo_extensions-1.4.0/tests/test_robustness.py +418 -0
- axbo_extensions-1.4.0/tests/test_saasbo.py +25 -0
- axbo_extensions-1.4.0/tests/test_trialdeps.py +97 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: axbo_extensions
|
|
3
|
+
Version: 1.4.0
|
|
4
|
+
Summary: Ax/BoTorch GLUE for multi-fidelity + robust optimization.
|
|
5
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
6
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
7
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
8
|
+
Requires-Python: >=3.11
|
|
9
|
+
Description-Content-Type: text/markdown
|
|
10
|
+
License-File: LICENSE
|
|
11
|
+
Requires-Dist: ax-platform==1.3.1
|
|
12
|
+
Requires-Dist: botorch==0.18.1
|
|
13
|
+
Requires-Dist: pydantic>=2.0
|
|
14
|
+
Requires-Dist: numpy
|
|
15
|
+
Requires-Dist: sympy
|
|
16
|
+
Provides-Extra: dev
|
|
17
|
+
Requires-Dist: pytest; extra == "dev"
|
|
18
|
+
Requires-Dist: mypy; extra == "dev"
|
|
19
|
+
Requires-Dist: black; extra == "dev"
|
|
20
|
+
Requires-Dist: ruff; extra == "dev"
|
|
21
|
+
Dynamic: license-file
|
|
22
|
+
|
|
23
|
+
# Bayesian (BOTorch) extensions for the Ax platform
|
|
24
|
+
|
|
25
|
+
> This is largely an experimental library basically glueing **multi-fidelity** and **robust** optimization
|
|
26
|
+
> into the ax-platform internals. Refactored out of
|
|
27
|
+
> [foamBO](https://github.com/FoamScience/OpenFOAM-Multi-Objective-Optimization) to decouple it from its orchestration capabilities.
|
|
28
|
+
|
|
29
|
+
Extracts foamBO's self-contained (historic) `robustness.py` (qMultiFidHVKG cost loop,
|
|
30
|
+
`SubstituteContextFeatures`, `RobustAcquisition` CVaR/MARS, `augment_generator_specs` /
|
|
31
|
+
`cycle_context`) into a separate package so more tools can have these capabilities
|
|
32
|
+
without getting foamBO's orchestration if they don't want it.
|
|
@@ -0,0 +1,10 @@
|
|
|
1
|
+
# Bayesian (BOTorch) extensions for the Ax platform
|
|
2
|
+
|
|
3
|
+
> This is largely an experimental library basically glueing **multi-fidelity** and **robust** optimization
|
|
4
|
+
> into the ax-platform internals. Refactored out of
|
|
5
|
+
> [foamBO](https://github.com/FoamScience/OpenFOAM-Multi-Objective-Optimization) to decouple it from its orchestration capabilities.
|
|
6
|
+
|
|
7
|
+
Extracts foamBO's self-contained (historic) `robustness.py` (qMultiFidHVKG cost loop,
|
|
8
|
+
`SubstituteContextFeatures`, `RobustAcquisition` CVaR/MARS, `augment_generator_specs` /
|
|
9
|
+
`cycle_context`) into a separate package so more tools can have these capabilities
|
|
10
|
+
without getting foamBO's orchestration if they don't want it.
|
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
"""axbo_extensions — shared Ax/BoTorch GLUE for multi-fidelity + robust optimization.
|
|
2
|
+
|
|
3
|
+
Home for foamBO's ported MultiFid/robust glue (qMultiFidHVKG cost loop, SubstituteContextFeatures,
|
|
4
|
+
RobustAcquisition CVaR/MARS, augment_generator_specs/cycle_context).
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from importlib.metadata import PackageNotFoundError, version
|
|
8
|
+
|
|
9
|
+
try:
|
|
10
|
+
__version__ = version("axbo_extensions")
|
|
11
|
+
except PackageNotFoundError: # running from source without an install
|
|
12
|
+
__version__ = "0.0.0"
|
|
@@ -0,0 +1,121 @@
|
|
|
1
|
+
"""Sobol-sensitivity dimensionality reduction (generic port of foamBO
|
|
2
|
+
``_maybe_reduce_dimensions``).
|
|
3
|
+
|
|
4
|
+
After the GP is fit, fix low-influence parameters at a constant to shrink the search space.
|
|
5
|
+
The latch lives in a caller-owned ``state`` dict so it survives the detached ask/run/tell
|
|
6
|
+
boundary (foamBO used an in-memory ``nonlocal``). Must run AFTER ``get_next_trials`` — Sobol
|
|
7
|
+
needs a fitted model.
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
from __future__ import annotations
|
|
11
|
+
|
|
12
|
+
import logging
|
|
13
|
+
|
|
14
|
+
log = logging.getLogger("axbo_extensions.dimreduction")
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
def select_params_to_fix(
|
|
18
|
+
importance: dict[str, float], min_importance: float, max_fix_fraction: float
|
|
19
|
+
) -> list[str]:
|
|
20
|
+
"""Low-influence params to fix: importance < threshold, ascending, capped so >=1 stays
|
|
21
|
+
active. Pure (no model) — the testable core of the reduction decision."""
|
|
22
|
+
names = list(importance)
|
|
23
|
+
cap = min(int(len(names) * max_fix_fraction), len(names) - 1)
|
|
24
|
+
ordered = sorted(names, key=lambda n: importance[n])
|
|
25
|
+
return [n for n in ordered if importance[n] < min_importance][:cap]
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def _cv_fit_ok(gs) -> bool:
|
|
29
|
+
"""False if model fit is too poor to trust Sobol (mean CV error > 50% of error range)."""
|
|
30
|
+
try:
|
|
31
|
+
import numpy as np
|
|
32
|
+
from ax.adapter.cross_validation import cross_validate
|
|
33
|
+
|
|
34
|
+
errs = []
|
|
35
|
+
for r in cross_validate(adapter=gs.adapter):
|
|
36
|
+
for m, obs in r.observed.data.means_dict.items():
|
|
37
|
+
errs.append(abs(obs - r.predicted.means_dict.get(m, obs)))
|
|
38
|
+
if errs:
|
|
39
|
+
rng = max(errs)
|
|
40
|
+
if rng > 0 and float(np.mean(errs)) / rng > 0.5:
|
|
41
|
+
return False
|
|
42
|
+
except Exception: # noqa: BLE001 - CV unavailable -> proceed with Sobol
|
|
43
|
+
pass
|
|
44
|
+
return True
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
def _sobol_importance(gs, exp) -> dict[str, float]:
|
|
48
|
+
"""Total-order Sobol indices of the fitted GP mean, per range parameter (Ax-private)."""
|
|
49
|
+
import torch
|
|
50
|
+
from ax.utils.sensitivity.sobol_measures import SobolSensitivityGPMean
|
|
51
|
+
|
|
52
|
+
model = gs.adapter.generator.surrogate.model
|
|
53
|
+
names = [n for n, p in exp.search_space.parameters.items() if hasattr(p, "lower")]
|
|
54
|
+
bounds = torch.zeros(2, len(names), dtype=torch.float64)
|
|
55
|
+
bounds[1] = 1.0 # Ax models operate in normalized [0,1] space
|
|
56
|
+
total = SobolSensitivityGPMean(
|
|
57
|
+
model=model, bounds=bounds, num_mc_samples=1000
|
|
58
|
+
).total_order_indices()
|
|
59
|
+
return {n: float(total[i].abs()) for i, n in enumerate(names)}
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def maybe_reduce_dimensions(client, cfg, state: dict) -> None:
|
|
63
|
+
"""Fix low-influence parameters once enough trials are in. Mutates the experiment search
|
|
64
|
+
space in place (persists when the Client is saved); idempotent via ``state['done']``.
|
|
65
|
+
|
|
66
|
+
cfg: a DimensionalityReductionConfig-like; state: ``{"done": bool, "attempts": int}``
|
|
67
|
+
(the caller persists it). Wrapped in a 3-strike try/except over the Ax-private paths (§6).
|
|
68
|
+
"""
|
|
69
|
+
if state.get("done") or not getattr(cfg, "enabled", False):
|
|
70
|
+
return
|
|
71
|
+
from ax.core.base_trial import TrialStatus
|
|
72
|
+
|
|
73
|
+
exp = client._experiment
|
|
74
|
+
n_done = sum(1 for t in exp.trials.values() if t.status == TrialStatus.COMPLETED)
|
|
75
|
+
if n_done < cfg.after_trials:
|
|
76
|
+
return
|
|
77
|
+
gs = client._generation_strategy
|
|
78
|
+
if gs is None or getattr(gs, "adapter", None) is None:
|
|
79
|
+
return
|
|
80
|
+
|
|
81
|
+
try:
|
|
82
|
+
if not _cv_fit_ok(gs):
|
|
83
|
+
log.info("dim-reduction deferred: model fit too poor")
|
|
84
|
+
return
|
|
85
|
+
importance = _sobol_importance(gs, exp)
|
|
86
|
+
except Exception as e: # noqa: BLE001
|
|
87
|
+
state["attempts"] = state.get("attempts", 0) + 1
|
|
88
|
+
if state["attempts"] >= 3:
|
|
89
|
+
state["done"] = True
|
|
90
|
+
log.warning("dim-reduction disabled after 3 failed attempts: %s", e)
|
|
91
|
+
return
|
|
92
|
+
|
|
93
|
+
to_fix = select_params_to_fix(importance, cfg.min_importance, cfg.max_fix_fraction)
|
|
94
|
+
if not to_fix:
|
|
95
|
+
state["done"] = True
|
|
96
|
+
return
|
|
97
|
+
|
|
98
|
+
best: dict = {}
|
|
99
|
+
if cfg.fix_at == "best":
|
|
100
|
+
try:
|
|
101
|
+
best = client.get_best_parameterization(use_model_predictions=False)[0]
|
|
102
|
+
except Exception: # noqa: BLE001 - fall back to center
|
|
103
|
+
best = {}
|
|
104
|
+
|
|
105
|
+
from ax.core.parameter import FixedParameter
|
|
106
|
+
|
|
107
|
+
for n in to_fix:
|
|
108
|
+
p = exp.search_space.parameters[n]
|
|
109
|
+
if cfg.fix_at == "best" and best.get(n) is not None:
|
|
110
|
+
val = best[n]
|
|
111
|
+
elif hasattr(p, "lower"):
|
|
112
|
+
val = (p.lower + p.upper) / 2
|
|
113
|
+
elif hasattr(p, "values"):
|
|
114
|
+
val = p.values[0]
|
|
115
|
+
else:
|
|
116
|
+
continue
|
|
117
|
+
exp.search_space.update_parameter(
|
|
118
|
+
FixedParameter(name=n, parameter_type=p.parameter_type, value=val)
|
|
119
|
+
)
|
|
120
|
+
log.info("dim-reduction: fixed %s=%s (importance=%.4f)", n, val, importance[n])
|
|
121
|
+
state["done"] = True
|
|
@@ -0,0 +1,267 @@
|
|
|
1
|
+
"""Multi-fidelity GLUE for Ax + BoTorch (ported from foamBO ``robustness.py`` /
|
|
2
|
+
``optimize.py``).
|
|
3
|
+
|
|
4
|
+
- ``MultiFidHVKGAcquisition``: Ax ``Acquisition`` subclass that strips ``X_pending``
|
|
5
|
+
(qMultiFidHVKG's input constructor rejects it).
|
|
6
|
+
- ``update_cost_state``: recompute ``cost_intercept`` / ``fidelity_weights`` for the
|
|
7
|
+
``AffineFidelityCostModel`` qMultiFidHVKG builds, from the observed is_cost metric. The
|
|
8
|
+
±0.5 positivity-floor + normalization math is load-bearing — copied verbatim.
|
|
9
|
+
- ``apply_fidelity_parameters``: stamp ``_is_fidelity`` / ``_target_value`` onto Ax
|
|
10
|
+
search-space parameters after ``configure_experiment``.
|
|
11
|
+
- ``augment_specs_for_multifidelity``: inject ``SingleTaskMultiFidelityGP`` surrogate
|
|
12
|
+
+ qMultiFidHVKG acqf + ``MultiFidHVKGAcquisition`` into a Client's generation-strategy nodes.
|
|
13
|
+
|
|
14
|
+
Pinned to ax-platform 1.3.x / botorch 0.18.x.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
from __future__ import annotations
|
|
18
|
+
|
|
19
|
+
import logging
|
|
20
|
+
from collections import defaultdict
|
|
21
|
+
from typing import Any, Optional
|
|
22
|
+
|
|
23
|
+
from ax.generators.torch.botorch_modular.acquisition import Acquisition
|
|
24
|
+
|
|
25
|
+
log = logging.getLogger("axbo_extensions.multifidelity")
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
class MultiFidHVKGAcquisition(Acquisition):
|
|
29
|
+
"""Acquisition subclass that strips kwargs incompatible with qMultiFidHVKG.
|
|
30
|
+
|
|
31
|
+
``construct_inputs_qMultiFidHVKG`` doesn't accept ``X_pending`` (unlike most other
|
|
32
|
+
BoTorch input constructors). This subclass clears it before the parent passes
|
|
33
|
+
acqf_options to the input constructor.
|
|
34
|
+
"""
|
|
35
|
+
|
|
36
|
+
def _construct_botorch_acquisition(self, botorch_acqf_class, botorch_acqf_options, model):
|
|
37
|
+
_saved = self.X_pending
|
|
38
|
+
self.X_pending = None
|
|
39
|
+
try:
|
|
40
|
+
return super()._construct_botorch_acquisition(
|
|
41
|
+
botorch_acqf_class=botorch_acqf_class,
|
|
42
|
+
botorch_acqf_options=botorch_acqf_options,
|
|
43
|
+
model=model,
|
|
44
|
+
)
|
|
45
|
+
finally:
|
|
46
|
+
self.X_pending = _saved
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
def _fidelity_feature_index(client, fid_param_name: str) -> int:
|
|
50
|
+
"""Feature index of the fidelity parameter in the search space."""
|
|
51
|
+
param_names = list(client._experiment.search_space.parameters.keys())
|
|
52
|
+
return param_names.index(fid_param_name)
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def update_cost_state(client, multifid_cost: dict, log=None) -> None:
|
|
56
|
+
"""Recompute cost-model params from observed is_cost metric data.
|
|
57
|
+
|
|
58
|
+
Reads observed values of the cost metric, computes per-fidelity means, then derives
|
|
59
|
+
``cost_intercept`` and ``fidelity_weights`` for the ``AffineFidelityCostModel`` that
|
|
60
|
+
qMultiFidHVKG's input constructor builds internally. Writes directly into the acqf_opts
|
|
61
|
+
dict (``multifid_cost['acqf_opts_ref']``) so the next acquisition construction picks up
|
|
62
|
+
updated values.
|
|
63
|
+
|
|
64
|
+
AffineFidelityCostModel: cost(x) = cost_intercept + Σ(w_i × s_i). With two fidelity
|
|
65
|
+
levels (s=0 cheap, s=1 expensive): cost_intercept = mean_cost_at_s0,
|
|
66
|
+
fidelity_weight = mean_cost_at_s1 - mean_cost_at_s0.
|
|
67
|
+
"""
|
|
68
|
+
from ax.core.base_trial import TrialStatus as _TS
|
|
69
|
+
|
|
70
|
+
cost_metric_name = multifid_cost["metric"]
|
|
71
|
+
fidelity_param_name = multifid_cost["fidelity_param"]
|
|
72
|
+
# Collect observed costs per fidelity level (keyed by raw fidelity value)
|
|
73
|
+
sums: dict = defaultdict(float)
|
|
74
|
+
counts: dict = defaultdict(int)
|
|
75
|
+
df = client._experiment.lookup_data().df
|
|
76
|
+
for trial in client._experiment.trials.values():
|
|
77
|
+
if trial.status not in (_TS.COMPLETED, _TS.EARLY_STOPPED):
|
|
78
|
+
continue
|
|
79
|
+
fid_val = trial.arm.parameters.get(fidelity_param_name)
|
|
80
|
+
if fid_val is None:
|
|
81
|
+
continue
|
|
82
|
+
sub = df[(df.trial_index == trial.index) & (df.metric_name == cost_metric_name)]
|
|
83
|
+
if sub.empty:
|
|
84
|
+
continue
|
|
85
|
+
sums[fid_val] += float(sub["mean"].iloc[-1])
|
|
86
|
+
counts[fid_val] += 1
|
|
87
|
+
per_f = {fv: sums[fv] / counts[fv] for fv in sums if counts[fv] > 0}
|
|
88
|
+
if not per_f or per_f == multifid_cost["state"]["per_fidelity"]:
|
|
89
|
+
return # no change
|
|
90
|
+
multifid_cost["state"]["per_fidelity"] = per_f
|
|
91
|
+
|
|
92
|
+
# Build numeric mapping for AffineFidelityCostModel.
|
|
93
|
+
# For numeric fidelity: use values directly.
|
|
94
|
+
# For categorical (str) fidelity: map to ordinal indices sorted by cost.
|
|
95
|
+
raw_keys = sorted(per_f.keys(), key=lambda k: per_f[k])
|
|
96
|
+
is_numeric = all(isinstance(k, (int, float)) for k in raw_keys)
|
|
97
|
+
if is_numeric:
|
|
98
|
+
fid_vals = sorted(float(k) for k in raw_keys)
|
|
99
|
+
costs = [per_f[k] for k in sorted(per_f.keys())]
|
|
100
|
+
else:
|
|
101
|
+
# Map categorical levels to 0..N-1 ordered by ascending cost
|
|
102
|
+
fid_vals = list(range(len(raw_keys)))
|
|
103
|
+
costs = [per_f[k] for k in raw_keys]
|
|
104
|
+
|
|
105
|
+
cost_intercept = min(costs) # cheapest fidelity floor
|
|
106
|
+
# Weights: slope per fidelity unit. For discrete {0,1}: w = cost_high - cost_low.
|
|
107
|
+
# For continuous: linear fit.
|
|
108
|
+
fid_range = fid_vals[-1] - fid_vals[0] if len(fid_vals) > 1 else 0
|
|
109
|
+
if fid_range > 0:
|
|
110
|
+
fid_weight = (max(costs) - min(costs)) / fid_range
|
|
111
|
+
else:
|
|
112
|
+
fid_weight = 1.0
|
|
113
|
+
# Positivity floor: BoTorch's InverseCostWeightedUtility requires cost(X) > 0 over
|
|
114
|
+
# the entire acqf domain. Ax extends INT-fidelity model-space bounds by ±0.5, so
|
|
115
|
+
# AffineFidelityCostModel sees fid ∈ [-0.5, 1.5]; cost(X) = intercept + weight × X
|
|
116
|
+
# must stay > 0 across that range → intercept ≥ 0.5 × weight + ε. Normalize to keep
|
|
117
|
+
# costs O(1) (qMultiFidHVKG tolerates large absolute values poorly).
|
|
118
|
+
max_cost = max(costs) if max(costs) > 0 else 1.0
|
|
119
|
+
cost_intercept_norm = cost_intercept / max_cost
|
|
120
|
+
fid_weight_norm = fid_weight / max_cost
|
|
121
|
+
cost_intercept_safe = max(cost_intercept_norm, 0.5 * fid_weight_norm + 1e-3)
|
|
122
|
+
opts = multifid_cost["acqf_opts_ref"]
|
|
123
|
+
opts["cost_intercept"] = max(cost_intercept_safe, 1e-6)
|
|
124
|
+
if log:
|
|
125
|
+
log.debug(
|
|
126
|
+
"MultiFid cost normalized: raw intercept=%.3f weight=%.3f max_cost=%.3f → "
|
|
127
|
+
"intercept=%.3f weight=%.3f",
|
|
128
|
+
cost_intercept,
|
|
129
|
+
fid_weight,
|
|
130
|
+
max_cost,
|
|
131
|
+
cost_intercept_safe,
|
|
132
|
+
fid_weight_norm,
|
|
133
|
+
)
|
|
134
|
+
fid_feature_idx = _fidelity_feature_index(client, multifid_cost["fidelity_param"])
|
|
135
|
+
opts["fidelity_weights"] = {fid_feature_idx: fid_weight_norm}
|
|
136
|
+
if log:
|
|
137
|
+
log.info(
|
|
138
|
+
"Multi-fidelity cost updated: intercept=%.2f, weight=%.2f "
|
|
139
|
+
"(from %d fidelity levels: %s)",
|
|
140
|
+
cost_intercept,
|
|
141
|
+
fid_weight,
|
|
142
|
+
len(per_f),
|
|
143
|
+
list(per_f.keys()),
|
|
144
|
+
)
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
def apply_fidelity_parameters(client, fidelity_params: dict, log=None) -> None:
|
|
148
|
+
"""Stamp ``_is_fidelity`` / ``_target_value`` onto Ax search-space parameters.
|
|
149
|
+
|
|
150
|
+
``fidelity_params`` maps parameter name -> fidelity_target. Call after
|
|
151
|
+
``configure_experiment`` (Ax's parameter configs don't carry these attrs).
|
|
152
|
+
"""
|
|
153
|
+
from ax.core.parameter import FixedParameter
|
|
154
|
+
|
|
155
|
+
for pname, target_val in fidelity_params.items():
|
|
156
|
+
p = client._experiment.search_space.parameters.get(pname)
|
|
157
|
+
if p is None or isinstance(p, FixedParameter):
|
|
158
|
+
continue
|
|
159
|
+
p._is_fidelity = True
|
|
160
|
+
p._target_value = target_val
|
|
161
|
+
if log:
|
|
162
|
+
log.info("Fidelity parameter: %s (fidelity_target=%s)", pname, target_val)
|
|
163
|
+
|
|
164
|
+
|
|
165
|
+
def augment_specs_for_multifidelity(
|
|
166
|
+
client, cost_metric_name: Optional[str] = None, log=None
|
|
167
|
+
) -> Optional[dict]:
|
|
168
|
+
"""Inject MultiFid surrogate + qMultiFidHVKG acqf + ``MultiFidHVKGAcquisition`` into the Client's
|
|
169
|
+
generation-strategy nodes.
|
|
170
|
+
|
|
171
|
+
Requires fidelity parameters already stamped (``apply_fidelity_parameters``) and a
|
|
172
|
+
generation strategy already set. Returns the ``multifid_cost`` state dict (to feed
|
|
173
|
+
``update_cost_state`` each batch) when ``cost_metric_name`` is given, else ``None``.
|
|
174
|
+
"""
|
|
175
|
+
from ax.adapter.registry import Generators
|
|
176
|
+
from ax.generators.torch.botorch_modular.surrogate import SurrogateSpec
|
|
177
|
+
from ax.generators.torch.botorch_modular.utils import ModelConfig
|
|
178
|
+
from botorch.acquisition.multi_objective.hypervolume_knowledge_gradient import (
|
|
179
|
+
qMultiFidelityHypervolumeKnowledgeGradient,
|
|
180
|
+
)
|
|
181
|
+
from botorch.models.gp_regression_fidelity import SingleTaskMultiFidelityGP
|
|
182
|
+
|
|
183
|
+
fidelity_params = [
|
|
184
|
+
p
|
|
185
|
+
for p in client._experiment.search_space.parameters.values()
|
|
186
|
+
if getattr(p, "is_fidelity", False)
|
|
187
|
+
]
|
|
188
|
+
if not fidelity_params:
|
|
189
|
+
return None
|
|
190
|
+
|
|
191
|
+
gs = client._generation_strategy
|
|
192
|
+
multifid_cost: Optional[dict] = None
|
|
193
|
+
for node in gs._nodes:
|
|
194
|
+
for spec in node.generator_specs:
|
|
195
|
+
if spec.generator_enum != Generators.BOTORCH_MODULAR:
|
|
196
|
+
continue
|
|
197
|
+
gk = spec.generator_kwargs
|
|
198
|
+
if "botorch_acqf_class" not in gk:
|
|
199
|
+
gk["botorch_acqf_class"] = qMultiFidelityHypervolumeKnowledgeGradient
|
|
200
|
+
if log:
|
|
201
|
+
log.info("Multi-fidelity: auto-selected qMultiFidHVKG")
|
|
202
|
+
# qMultiFidHVKG's fantasy-based acqf optimizer diverges on raw multi-scale bounds
|
|
203
|
+
# (out-of-bounds candidates / hangs). MBM deliberately DROPS UnitX to keep discrete
|
|
204
|
+
# params discrete, leaving the acqf in raw parameter space — fine for qLogNEHVI/MARS,
|
|
205
|
+
# fatal for KG. Force UnitX-inclusive transforms so the acqf optimizes in [0,1].
|
|
206
|
+
# Idempotent; range-int params get continuous-relaxed + rounded (acceptable for KG).
|
|
207
|
+
from ax.adapter.registry import Cont_X_trans, Y_trans
|
|
208
|
+
|
|
209
|
+
gk["transforms"] = Cont_X_trans + Y_trans
|
|
210
|
+
# Inject surrogate only if not already set (preserves an explicit one and is
|
|
211
|
+
# idempotent on a reloaded client whose strategy already carries it).
|
|
212
|
+
if gk.get("surrogate_spec") is None:
|
|
213
|
+
# allow_batched_models=False forces a ModelListGP of per-output STMFGPs
|
|
214
|
+
# (qMultiFidHVKG requirement — load-bearing).
|
|
215
|
+
gk["surrogate_spec"] = SurrogateSpec(
|
|
216
|
+
model_configs=[ModelConfig(botorch_model_class=SingleTaskMultiFidelityGP)],
|
|
217
|
+
allow_batched_models=False,
|
|
218
|
+
)
|
|
219
|
+
gk["acquisition_class"] = MultiFidHVKGAcquisition
|
|
220
|
+
if log:
|
|
221
|
+
log.info(
|
|
222
|
+
"Multi-fidelity: SingleTaskMultiFidelityGP surrogate (fidelity param: %s)",
|
|
223
|
+
fidelity_params[0].name,
|
|
224
|
+
)
|
|
225
|
+
# Build cost state regardless of fresh/reloaded surrogate so the batch loop's
|
|
226
|
+
# update_cost_state has a live acqf_opts_ref after a JSON reload.
|
|
227
|
+
if cost_metric_name is not None:
|
|
228
|
+
multifid_cost = {
|
|
229
|
+
"state": {"per_fidelity": {}},
|
|
230
|
+
"metric": cost_metric_name,
|
|
231
|
+
"fidelity_param": fidelity_params[0].name,
|
|
232
|
+
"acqf_opts_ref": gk.setdefault("botorch_acqf_options", {}),
|
|
233
|
+
}
|
|
234
|
+
multifid_cost["acqf_opts_ref"].setdefault("cost_intercept", 1.0)
|
|
235
|
+
return multifid_cost
|
|
236
|
+
|
|
237
|
+
|
|
238
|
+
# ---------------------------------------------------------------------------
|
|
239
|
+
# Ax JSON serialization registry — run at import so saved MultiFid experiments serialize.
|
|
240
|
+
# MultiFidHVKGAcquisition is an Ax Acquisition wrapper; qMultiFidHVKG/MOMF are BoTorch
|
|
241
|
+
# AcquisitionFunctions not in Ax's default registry. Without this, save_to_json_file
|
|
242
|
+
# raises "Class qMultiFidelityHypervolumeKnowledgeGradient not in registry".
|
|
243
|
+
# ---------------------------------------------------------------------------
|
|
244
|
+
try:
|
|
245
|
+
from ax.storage.botorch_modular_registry import (
|
|
246
|
+
CLASS_TO_REGISTRY,
|
|
247
|
+
CLASS_TO_REVERSE_REGISTRY,
|
|
248
|
+
register_acquisition,
|
|
249
|
+
)
|
|
250
|
+
from botorch.acquisition import AcquisitionFunction as _AcqFn
|
|
251
|
+
|
|
252
|
+
register_acquisition(MultiFidHVKGAcquisition)
|
|
253
|
+
for _mod, _name in [
|
|
254
|
+
(
|
|
255
|
+
"botorch.acquisition.multi_objective.hypervolume_knowledge_gradient",
|
|
256
|
+
"qMultiFidelityHypervolumeKnowledgeGradient",
|
|
257
|
+
),
|
|
258
|
+
("botorch.acquisition.multi_objective.multi_fidelity", "MOMF"),
|
|
259
|
+
]:
|
|
260
|
+
try:
|
|
261
|
+
_cls = getattr(__import__(_mod, fromlist=[_name]), _name)
|
|
262
|
+
CLASS_TO_REGISTRY[_AcqFn][_cls] = _name
|
|
263
|
+
CLASS_TO_REVERSE_REGISTRY[_AcqFn][_name] = _cls
|
|
264
|
+
except (ImportError, AttributeError): # pragma: no cover
|
|
265
|
+
pass
|
|
266
|
+
except ImportError: # pragma: no cover - Ax storage registry path drift
|
|
267
|
+
log.warning("Ax storage registry unavailable; MultiFid JSON (de)serialization skipped")
|
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
"""Region screening — inactivity decision core (port of foamBO _maybe_screen_regions).
|
|
2
|
+
|
|
3
|
+
Only the **pure decision** is ported: given a region's per-trial scalars, is its spread
|
|
4
|
+
below the inactivity threshold, and how long has the streak run. That math is durable and
|
|
5
|
+
reusable.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
import logging
|
|
11
|
+
from typing import Iterable, Optional
|
|
12
|
+
|
|
13
|
+
log = logging.getLogger("axbo_extensions.regionscreening")
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def region_inactive(
|
|
17
|
+
scalars: Iterable[float],
|
|
18
|
+
min_delta: Optional[float] = None,
|
|
19
|
+
min_delta_frac: Optional[float] = None,
|
|
20
|
+
) -> dict:
|
|
21
|
+
"""Is the region's objective-scalar spread below its inactivity threshold?
|
|
22
|
+
|
|
23
|
+
``min_delta`` = absolute threshold; else ``min_delta_frac`` is relative to the largest
|
|
24
|
+
magnitude observed. Returns ``{samples, spread, threshold, inactive}``.
|
|
25
|
+
"""
|
|
26
|
+
vals = list(scalars)
|
|
27
|
+
spread = (max(vals) - min(vals)) if vals else 0.0
|
|
28
|
+
ref = max((abs(v) for v in vals), default=0.0)
|
|
29
|
+
threshold = min_delta if min_delta is not None else (min_delta_frac or 0.0) * ref
|
|
30
|
+
return {
|
|
31
|
+
"samples": len(vals),
|
|
32
|
+
"spread": spread,
|
|
33
|
+
"threshold": threshold,
|
|
34
|
+
"inactive": bool(vals) and spread < threshold,
|
|
35
|
+
}
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def next_streak(streak: int, inactive: bool) -> int:
|
|
39
|
+
"""Consecutive-inactive streak: increment when inactive, reset to 0 otherwise."""
|
|
40
|
+
return streak + 1 if inactive else 0
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
def shrunk_bounds(lo: float, hi: float, center: float, shrink_factor: float) -> tuple[float, float]:
|
|
44
|
+
"""New ``[lo, hi]`` after a shrink: a ``shrink_factor * range`` window around ``center``,
|
|
45
|
+
clamped to the original bounds. Pure — the shrink geometry."""
|
|
46
|
+
half = (shrink_factor * (hi - lo)) / 2.0
|
|
47
|
+
return max(lo, center - half), min(hi, center + half)
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
def _group_params(parameter_groups: dict, group: str) -> list[str]:
|
|
51
|
+
return [n for n, gs in (parameter_groups or {}).items() if group in (gs or [])]
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def maybe_screen_regions(
|
|
55
|
+
client, cfg, scalars_by_region: dict, parameter_groups: dict, state: dict
|
|
56
|
+
) -> None:
|
|
57
|
+
"""Advise/shrink parameter groups whose objective-scalar spread stays below threshold for
|
|
58
|
+
``confirm_passes`` consecutive passes (plan §2.A2). ``scalars_by_region``: ``{group:
|
|
59
|
+
{trial_idx: scalar}}`` (caller reads these from the executor's results — the generic
|
|
60
|
+
result-field provider). ``state`` persists streaks / shrunk groups / original bounds /
|
|
61
|
+
advisories across the detached ask/run/tell boundary. Shrink mutates the search space
|
|
62
|
+
in place (Ax-private update_parameter, gated)."""
|
|
63
|
+
if not getattr(cfg, "enabled", False) or not cfg.regions:
|
|
64
|
+
return
|
|
65
|
+
streaks = state.setdefault("streaks", {})
|
|
66
|
+
shrunk = state.setdefault("shrunk_groups", [])
|
|
67
|
+
advisories = state.setdefault("advisories", [])
|
|
68
|
+
|
|
69
|
+
ready = []
|
|
70
|
+
for region in cfg.regions:
|
|
71
|
+
vals = list(scalars_by_region.get(region.group, {}).values())
|
|
72
|
+
if len(vals) < cfg.after_trials:
|
|
73
|
+
continue
|
|
74
|
+
inactive = region_inactive(vals, region.min_delta, region.min_delta_frac)["inactive"]
|
|
75
|
+
streaks[region.group] = next_streak(streaks.get(region.group, 0), inactive)
|
|
76
|
+
if streaks[region.group] >= cfg.confirm_passes and region.group not in shrunk:
|
|
77
|
+
ready.append(region.group)
|
|
78
|
+
if not ready:
|
|
79
|
+
return
|
|
80
|
+
|
|
81
|
+
if cfg.mode == "advise":
|
|
82
|
+
for g in ready:
|
|
83
|
+
advisories.append({"group": g, "params": _group_params(parameter_groups, g)})
|
|
84
|
+
shrunk.append(g) # recorded so we advise once
|
|
85
|
+
log.info("region screening (advise): group %s inactive", g)
|
|
86
|
+
return
|
|
87
|
+
|
|
88
|
+
# shrink mode — narrow each group param's range around the best point, once per group.
|
|
89
|
+
try:
|
|
90
|
+
from ax.core.parameter import RangeParameter
|
|
91
|
+
|
|
92
|
+
exp = client._experiment
|
|
93
|
+
orig_bounds = state.setdefault("original_bounds", {})
|
|
94
|
+
try:
|
|
95
|
+
best = client.get_best_parameterization(use_model_predictions=False)[0]
|
|
96
|
+
except Exception: # noqa: BLE001
|
|
97
|
+
best = {}
|
|
98
|
+
n_total = len(exp.search_space.parameters)
|
|
99
|
+
budget = max(min(cfg.max_shrinkable or n_total, n_total - 1) - len(orig_bounds), 0)
|
|
100
|
+
for g in ready:
|
|
101
|
+
for pname in _group_params(parameter_groups, g):
|
|
102
|
+
if budget <= 0:
|
|
103
|
+
break
|
|
104
|
+
p = exp.search_space.parameters.get(pname)
|
|
105
|
+
if not isinstance(p, RangeParameter):
|
|
106
|
+
continue
|
|
107
|
+
lo, hi = float(p.lower), float(p.upper)
|
|
108
|
+
orig_bounds.setdefault(pname, (lo, hi))
|
|
109
|
+
center = best.get(pname, (lo + hi) / 2.0)
|
|
110
|
+
new_lo, new_hi = shrunk_bounds(lo, hi, center, cfg.shrink_factor)
|
|
111
|
+
p.update_range(lower=new_lo, upper=new_hi)
|
|
112
|
+
budget -= 1
|
|
113
|
+
log.info("region screening: shrank %s -> [%.4g, %.4g]", pname, new_lo, new_hi)
|
|
114
|
+
shrunk.append(g)
|
|
115
|
+
except Exception as e: # noqa: BLE001 - Ax-private mutation path, stay non-fatal
|
|
116
|
+
log.warning("region screening shrink failed: %s", e)
|