fix: multi-fidelity + robust optimization composition

numeric fidelity, qMFHVKG wiring, constant-outcome crashes/bugs

Author: Elwardi

src/foambo/common.py

@@ -323,8 +323,12 @@ def prepare_case(
             param_name = entry['parameter']
             template_path = entry['file_path'].lstrip('/')
             if param_name in parameters:
+                val = str(parameters[param_name])
+                value_map = entry.get('value_map')
+                if value_map:
+                    val = value_map.get(val, val)
                 shutil.copy2(
-                    os.path.join(case.path, template_path + "." + parameters[param_name]),
+                    os.path.join(case.path, template_path + "." + val),
                     os.path.join(case.path, template_path)
                 )
     # Process parameters with foamlib

src/foambo/constraints.py

@@ -148,6 +148,10 @@ def _patched_optimize_acqf(*args, **kwargs):
     # --- Inject nonlinear constraints ---
     if _active_nl_constraints and "nonlinear_inequality_constraints" not in kwargs:
         kwargs["nonlinear_inequality_constraints"] = _active_nl_constraints
+        # ic_generator is required when nonlinear constraints are present
+        if "ic_generator" not in kwargs:
+            from botorch.optim.initializers import gen_batch_initial_conditions
+            kwargs["ic_generator"] = gen_batch_initial_conditions
         log.debug("Injected %d nonlinear constraint(s) into optimize_acqf",
                   len(_active_nl_constraints))
 

src/foambo/default_config.py

@@ -911,22 +911,23 @@ def __init__(self, **kwargs):
             nodes are only needed to override the default (e.g. pick MOMF).
 
             **Runner dispatch (recommended):** use ``file_substitution`` with
-            a string ``ChoiceParameter`` fidelity. Place ``Allrun.meanline``
-            + ``Allrun.CFD`` in the template case:
+            ``value_map`` to map numeric fidelity to file suffixes. Place
+            ``Allrun.meanline`` + ``Allrun.CFD`` in the template case:
             ```yaml
             experiment:
               parameters:
                 - name: fidelity
-                  values: ["meanline", "CFD"]
-                  parameter_type: str
+                  bounds: [0, 1]
+                  parameter_type: int
                   is_fidelity: true
-                  target_value: "CFD"
+                  target_value: 1
 
             optimization:
               case_runner:
                 file_substitution:
                   - parameter: fidelity
                     file_path: /Allrun
+                    value_map: {0: "meanline", 1: "CFD"}
             ```
 
             Other resolvable names: ``qMultiFidelityKnowledgeGradient``

src/foambo/docs_concepts.py

@@ -655,15 +655,17 @@
 
 **Setup in foamBO:**
 
-1. Add a fidelity parameter with ``is_fidelity: true`` and ``target_value``:
+1. Add a **numeric** fidelity parameter with ``is_fidelity: true`` and ``target_value``.
+   Ax requires fidelity to be a ``RangeParameter`` (numeric), not a ``ChoiceParameter``.
+   Use ``int`` for discrete levels or ``float`` for continuous:
 ```yaml
 experiment:
   parameters:
     - name: fidelity
-      parameter_type: str
-      values: ["coarse", "fine"]
+      parameter_type: int
+      bounds: [0, 1]        # 0=cheap, 1=expensive
       is_fidelity: true
-      target_value: "fine"
+      target_value: 1
 ```
 
 2. Mark one metric as the cost signal with ``is_cost: true``:
@@ -719,29 +721,22 @@
 - foamBO does NOT auto-wire cost for MOMF; pass ``cost_call`` manually
   via ``botorch_acqf_options`` if needed.
 
-**Runner dispatch** via ``file_substitution`` (recommended):
-Use a string ``ChoiceParameter`` for fidelity and foamBO's built-in file
-substitution to swap the runner script (or a portion of it) per fidelity level. Place
-``Allrun.coarse`` and ``Allrun.fine`` in the template case:
+**Runner dispatch** via ``file_substitution`` with ``value_map`` (recommended):
+Since fidelity must be numeric, use ``value_map`` to map numeric values to
+human-readable file suffixes. Place ``Allrun.meanline`` and ``Allrun.CFD``
+in the template case:
 ```yaml
 optimization:
   case_runner:
     file_substitution:
       - parameter: fidelity
         file_path: /Allrun
+        value_map: {0: "meanline", 1: "CFD"}
 ```
-When ``fidelity=coarse``, the runner copies ``Allrun.coarse`` →
-``Allrun`` before execution. When ``fidelity=fine``, copies
-``Allrun.fine`` → ``Allrun``. No if/else branching in scripts needed.
-
-**Alternative** (continuous fidelity via env var):
-```bash
-if [ "$FIDELITY" = "0" ] || [ "$FIDELITY" = "0.0" ]; then
-    ./Allrun.coarse
-else
-    ./Allrun.fine
-fi
-```
+When ``fidelity=0``, copies ``Allrun.meanline`` → ``Allrun``.
+When ``fidelity=1``, copies ``Allrun.CFD`` → ``Allrun``.
+Without ``value_map``, the raw numeric value is used as the suffix
+(e.g. ``Allrun.0``, ``Allrun.1``).
 
 **Cost model evolution:** before any trials complete, uniform cost is
 assumed (cost ratio = 1). As ``is_cost`` data arrives, foamBO recomputes

src/foambo/optimize.py

@@ -85,6 +85,11 @@ def _update_cost_state(client, mf_cost, log=None):
     """
     from ax.core.base_trial import TrialStatus as _TS
     from collections import defaultdict
+
+    def _get_fidelity_feature_index(cl, fid_param_name):
+        """Return the feature index of the fidelity parameter in the search space."""
+        param_names = list(cl._experiment.search_space.parameters.keys())
+        return param_names.index(fid_param_name)
     cost_metric_name = mf_cost["metric"]
     fidelity_param_name = mf_cost["fidelity_param"]
     # Collect observed costs per fidelity level (keyed by raw fidelity value)
@@ -128,12 +133,15 @@ def _update_cost_state(client, mf_cost, log=None):
         fid_weight = (max(costs) - min(costs)) / fid_range
     else:
         fid_weight = 1.0
-    # Write into acqf_opts — Ax passes these to the input constructor
+    # Write into acqf_opts — Ax passes these to the input constructor.
+    # fidelity_weights is keyed by feature index (same as target_fidelities),
+    # NOT by fidelity values.  Ax extracts the fidelity feature index from
+    # search_space_digest.fidelity_features.
     opts = mf_cost["acqf_opts_ref"]
     opts["cost_intercept"] = max(cost_intercept, 1e-6)
-    opts["fidelity_weights"] = {
-        int(fv): fid_weight for fv in fid_vals
-    }
+    # Determine the fidelity parameter's feature index
+    fid_feature_idx = _get_fidelity_feature_index(client, mf_cost["fidelity_param"])
+    opts["fidelity_weights"] = {fid_feature_idx: fid_weight}
     if log:
         log.info("Multi-fidelity cost updated: intercept=%.2f, weight=%.2f "
                  "(from %d fidelity levels: %s)", cost_intercept, fid_weight,
@@ -273,12 +281,15 @@ def _patched_object_to_json(obj, **kwargs):
     _fidelity_cfg = getattr(exp_cfg, '_fidelity_params', None) or \
                     getattr(type(exp_cfg), '_fidelity_params', None)
     if _fidelity_cfg:
+        from ax.core.parameter import FixedParameter
         for pname, target_val in _fidelity_cfg.items():
             p = client._experiment.search_space.parameters.get(pname)
-            if p is not None:
+            if p is not None and not isinstance(p, FixedParameter):
                 p._is_fidelity = True
                 p._target_value = target_val
                 log.info("Fidelity parameter: %s (target_value=%s)", pname, target_val)
+            elif isinstance(p, FixedParameter):
+                log.info("Fidelity parameter %s is fixed (specialized) — skipping MF wiring", pname)
 
     ## 2.0 Trial generation
     # Suppress Ax's verbose GenerationStrategy repr log
@@ -321,9 +332,16 @@ def _patched_object_to_json(obj, **kwargs):
                 try:
                     from botorch.models.gp_regression_fidelity import SingleTaskMultiFidelityGP
                     from ax.generators.torch.botorch_modular.surrogate import SurrogateSpec
+                    from ax.generators.torch.botorch_modular.utils import ModelConfig
+                    from foambo.robustness import MFHVKGAcquisition
                     gk["surrogate_spec"] = SurrogateSpec(
-                        botorch_model_class=SingleTaskMultiFidelityGP,
+                        model_configs=[ModelConfig(
+                            botorch_model_class=SingleTaskMultiFidelityGP,
+                        )],
+                        allow_batched_models=False,
                     )
+                    # qMFHVKG input constructor doesn't accept X_pending
+                    spec.generator_kwargs["acquisition_class"] = MFHVKGAcquisition
                     log.info("Multi-fidelity: auto-selected SingleTaskMultiFidelityGP surrogate "
                              "(fidelity param: %s)", _fidelity_params[0].name)
                     # Wire cost model for MF acquisition.

src/foambo/orchestrate.py

@@ -917,11 +917,26 @@ class VariableSubstOptions(FoamBOBaseModel):
 
 
 class FileSubstOptions(FoamBOBaseModel):
-    """Replaces a case file with a variant based on a choice parameter value."""
-    parameter: str = Field(description="Name of the choice parameter whose value selects the file variant",
+    """Replaces a case file with a variant based on a parameter value."""
+    parameter: str = Field(description="Name of the parameter whose value selects the file variant",
                             examples=["y"])
     file_path: str = Field(description="Relative path to the file (e.g. `/system/fvSolution`). Variants must exist as `<file_path>.<value>`",
                             examples=["/constant/y"])
+    value_map: Dict[str, str] | None = Field(default=None,
+        description=(
+            "Optional mapping from parameter values to file suffixes. "
+            "Useful for numeric parameters (e.g. fidelity) where the file "
+            "variant has a human-readable name.\n\n"
+            "Example: ``{0: meanline, 1: CFD}`` maps fidelity=0 to "
+            "``Allrun.meanline`` and fidelity=1 to ``Allrun.CFD``."
+        ))
+
+    @field_validator("value_map", mode="before")
+    @classmethod
+    def coerce_value_map_keys(cls, v):
+        if isinstance(v, dict | DictConfig):
+            return {str(k): str(val) for k, val in v.items()}
+        return v
 
 
 class FoamJobRunnerOptions(FoamBOBaseModel):
@@ -1043,8 +1058,14 @@ def to_objective_metrics_dict(self):
         }
 
     def to_tracking_metrics_dict(self):
+        # Register non-objective, non-cost metrics as tracking so they appear
+        # in experiment.signature_to_metric (required by early stopping).
+        # is_cost metrics are excluded — they inflate objective_weights
+        # (size mismatch with GP model outputs) and their data is fetched
+        # by the runner anyway (just not modeled by the GP).
         return {
-            "metrics": [m.to_metric() for m in self.metrics if not m.name in self.objective]
+            "metrics": [m.to_metric() for m in self.metrics
+                        if m.name not in self.objective and not m.is_cost]
         }
 
     def to_runner_dict(self):

src/foambo/preflight.py

@@ -263,6 +263,31 @@ def _check_config_coherence(cfg: DictConfig, r: PreflightResult) -> PreflightRes
             r.failed("MF cost metric",
                      f"multiple is_cost metrics: {cost_metrics}; exactly one required")
 
+    # MF + robust + string choice params: qMFKG requires continuous optimizer,
+    # string categoricals can't be relaxed to continuous.
+    has_robust = cfg.get("robust_optimization") is not None
+    if has_fidelity:
+        choice_params = [p["name"] for p in exp.get("parameters", [])
+                         if "values" in p and p.get("parameter_type") == "str"]
+        if choice_params:
+            if has_robust:
+                r.failed("MF + robust + string params",
+                         f"string choice parameters {choice_params} are incompatible with "
+                         "multi-fidelity + robust optimization (qMFKG requires continuous optimizer). "
+                         "Use numeric encoding with file_substitution.value_map instead")
+            else:
+                r.warned("MF + string params",
+                         f"string choice parameters {choice_params} may cause issues with "
+                         "multi-fidelity acquisition (KG-family requires continuous optimizer). "
+                         "Consider numeric encoding with file_substitution.value_map")
+        # Fidelity param itself must be numeric
+        fid_params = [p for p in exp.get("parameters", []) if p.get("is_fidelity")]
+        for fp in fid_params:
+            if "values" in fp:
+                r.failed(f"Fidelity parameter '{fp['name']}'",
+                         "must be numeric (bounds + parameter_type int/float), not a choice parameter. "
+                         "Ax requires numeric fidelity for the GP kernel")
+
     param_names = {p["name"] for p in exp.get("parameters", [])}
     for constraint in exp.get("parameter_constraints", []):
         # Validate with sympy if available, else fall back to token check