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.
@@ -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)