diff --git a/bofire/benchmarks/benchmark.py b/bofire/benchmarks/benchmark.py index ff1f24835..147d010e6 100644 --- a/bofire/benchmarks/benchmark.py +++ b/bofire/benchmarks/benchmark.py @@ -131,6 +131,7 @@ def __init__( n_filler_features: int = 1, n_features_per_original_feature: int = 1, max_count: Optional[int] = None, + min_count: int = 0, **kwargs, ): super().__init__(**kwargs) @@ -191,13 +192,9 @@ def __init__( if max_count is not None: constraints.append( NChooseKConstraint( - features=[ - key - for key in inputs.get_keys() - if not key.startswith("x_filler_") - ], + features=inputs.get_keys(), max_count=max_count, - min_count=0, + min_count=min_count, none_also_valid=True, ) ) @@ -245,20 +242,40 @@ def get_optima(self) -> pd.DataFrame: class SpuriousFeaturesWrapper(Benchmark): """Wrapper that adds spurious features to a benchmark, that are ignored on evaluation.""" - def __init__(self, benchmark: Benchmark, n_spurious_features: int = 1, **kwargs): + def __init__( + self, + benchmark: Benchmark, + n_spurious_features: int = 1, + max_count: Optional[int] = None, + min_count: int = 0, + **kwargs, + ): super().__init__(**kwargs) assert n_spurious_features >= 1, "n_spurious_features must be >= 1." + assert len(benchmark.domain.constraints) == 0, "Constraints not supported yet." self._benchmark = benchmark + inputs = Inputs( + features=benchmark.domain.inputs.features # ty: ignore[unsupported-operator] + + [ + ContinuousInput(key=f"x_spurious_{i}", bounds=(0, 1)) + for i in range(n_spurious_features) + ] + ) + constraints = Constraints( + constraints=[ + NChooseKConstraint( + features=inputs.get_keys(), + max_count=max_count, + min_count=min_count, + none_also_valid=False, + ) + ] + ) + self._domain = Domain( - inputs=Inputs( - features=benchmark.domain.inputs.features # ty: ignore[unsupported-operator] - + [ - ContinuousInput(key=f"x_spurious_{i}", bounds=(0, 1)) - for i in range(n_spurious_features) - ] - ), + inputs=inputs, outputs=self._benchmark.domain.outputs, - constraints=self._benchmark.domain.constraints, + constraints=constraints, ) def _f(self, candidates: pd.DataFrame, **kwargs) -> pd.DataFrame: diff --git a/bofire/data_models/domain/features.py b/bofire/data_models/domain/features.py index 0b0d2c269..92eea206f 100644 --- a/bofire/data_models/domain/features.py +++ b/bofire/data_models/domain/features.py @@ -823,6 +823,7 @@ def get_bounds( specs: InputTransformSpecs, experiments: Optional[pd.DataFrame] = None, reference_experiment: Optional[pd.Series] = None, + relax_allow_zero: bool = False, ) -> Tuple[List[float], List[float]]: """Returns the boundaries of the optimization problem based on the transformations defined in the `specs` dictionary. @@ -837,6 +838,9 @@ def get_bounds( then the local bounds based on a local search region are provided as reference to the reference experiment. Currently only supported for continuous inputs. For more details, it is referred to https://www.merl.com/publications/docs/TR2023-057.pdf. Defaults to None. + relax_allow_zero (bool, optional): If True, semi-continuous continuous inputs + (`allow_zero=True` with positive lower bound) report a relaxed lower bound of 0. + Other input types ignore this flag. Defaults to False. Raises: ValueError: If a feature type is not known. @@ -866,6 +870,7 @@ def get_bounds( if reference_experiment is not None else None ), + relax_allow_zero=relax_allow_zero, ) lower += lo upper += up diff --git a/bofire/data_models/features/categorical.py b/bofire/data_models/features/categorical.py index c5f6cbc31..ea4c95273 100644 --- a/bofire/data_models/features/categorical.py +++ b/bofire/data_models/features/categorical.py @@ -383,7 +383,9 @@ def get_bounds( transform_type: TTransform, values: Optional[pd.Series] = None, reference_value: Optional[str] = None, + relax_allow_zero: bool = False, ) -> Tuple[List[float], List[float]]: + # relax_allow_zero is only meaningful for ContinuousInput; ignored here. assert isinstance(transform_type, CategoricalEncodingEnum) if transform_type == CategoricalEncodingEnum.ORDINAL: return [0], [len(self.categories) - 1] diff --git a/bofire/data_models/features/continuous.py b/bofire/data_models/features/continuous.py index b621d6e68..20693ea68 100644 --- a/bofire/data_models/features/continuous.py +++ b/bofire/data_models/features/continuous.py @@ -221,14 +221,27 @@ def get_bounds( transform_type: Optional[TTransform] = None, values: Optional[pd.Series] = None, reference_value: Optional[float] = None, + relax_allow_zero: bool = False, ) -> Tuple[List[float], List[float]]: assert transform_type is None if reference_value is not None and values is not None: raise ValueError("Only one can be used, `local_value` or `values`.") + # Effective lower bound: 0 for semi-continuous features when the + # caller asks for the convex-relaxation view. The `is_fixed` case + # short-circuits below (a fixed feature reports its single value). + effective_lower = self.lower_bound + if ( + relax_allow_zero + and self.allow_zero + and self.lower_bound > 0 + and not self.is_fixed() + ): + effective_lower = 0.0 + if values is None: if reference_value is None or self.is_fixed(): - return [self.lower_bound], [self.upper_bound] + return [effective_lower], [self.upper_bound] local_relative_bounds = self.local_relative_bounds or ( math.inf, @@ -238,7 +251,7 @@ def get_bounds( return [ max( reference_value - local_relative_bounds[0], - self.lower_bound, + effective_lower, ), ], [ min( @@ -247,7 +260,7 @@ def get_bounds( ), ] - lower = min(self.lower_bound, values.min()) + lower = min(effective_lower, values.min()) upper = max(self.upper_bound, values.max()) return [lower], [upper] diff --git a/bofire/data_models/features/descriptor.py b/bofire/data_models/features/descriptor.py index 51aee7b72..8b4f1317a 100644 --- a/bofire/data_models/features/descriptor.py +++ b/bofire/data_models/features/descriptor.py @@ -164,7 +164,9 @@ def get_bounds( transform_type: TTransform, values: Optional[pd.Series] = None, reference_value: Optional[str] = None, + relax_allow_zero: bool = False, ) -> Tuple[List[float], List[float]]: + # relax_allow_zero is only meaningful for ContinuousInput; ignored here. if transform_type != CategoricalEncodingEnum.DESCRIPTOR: return super().get_bounds(transform_type, values) # in case that values is None, we return the optimization bounds diff --git a/bofire/data_models/features/discrete.py b/bofire/data_models/features/discrete.py index 016522e8b..9edf292cf 100644 --- a/bofire/data_models/features/discrete.py +++ b/bofire/data_models/features/discrete.py @@ -171,8 +171,10 @@ def get_bounds( transform_type: Optional[TTransform] = None, values: Optional[pd.Series] = None, reference_value: Optional[float] = None, + relax_allow_zero: bool = False, ) -> Tuple[List[float], List[float]]: assert transform_type is None + # relax_allow_zero is only meaningful for ContinuousInput; ignored here. if values is None: return [self.lower_bound], [self.upper_bound] lower = min(self.lower_bound, values.min()) diff --git a/bofire/data_models/features/feature.py b/bofire/data_models/features/feature.py index 6c0321b32..fb681626f 100644 --- a/bofire/data_models/features/feature.py +++ b/bofire/data_models/features/feature.py @@ -137,6 +137,7 @@ def get_bounds( transform_type: Optional[TTransform] = None, values: Optional[pd.Series] = None, reference_value: Optional[Union[float, str]] = None, + relax_allow_zero: bool = False, ) -> Tuple[List[float], List[float]]: """Returns the bounds of an input feature depending on the requested transform type. @@ -147,6 +148,10 @@ def get_bounds( reference_value (Optional[float], optional): If a reference value is provided, then the local bounds based on a local search region are provided. Currently only supported for continuous inputs. For more details, it is referred to https://www.merl.com/publications/docs/TR2023-057.pdf. + relax_allow_zero (bool, optional): If True, semi-continuous continuous inputs (`allow_zero=True` + with a positive lower bound) report a relaxed lower bound of 0, exposing the convex + relaxation `[0, ub]` to downstream optimisers. Other input types ignore this flag. + Defaults to False. Returns: Tuple[List[float], List[float]]: List of lower bound values, list of upper bound values. diff --git a/bofire/data_models/features/molecular.py b/bofire/data_models/features/molecular.py index 2d9b6901b..2354afbdd 100644 --- a/bofire/data_models/features/molecular.py +++ b/bofire/data_models/features/molecular.py @@ -104,7 +104,9 @@ def get_bounds( transform_type: Union[CategoricalEncodingEnum, AnyMolFeatures], values: Optional[pd.Series] = None, reference_value: Optional[str] = None, + relax_allow_zero: bool = False, ) -> Tuple[List[float], List[float]]: + # relax_allow_zero is only meaningful for ContinuousInput; ignored here. if isinstance(transform_type, CategoricalEncodingEnum): # we are just using the standard categorical transformations return super().get_bounds( diff --git a/bofire/data_models/strategies/predictives/acqf_optimization.py b/bofire/data_models/strategies/predictives/acqf_optimization.py index 4ec2617fe..381d5d976 100644 --- a/bofire/data_models/strategies/predictives/acqf_optimization.py +++ b/bofire/data_models/strategies/predictives/acqf_optimization.py @@ -113,6 +113,26 @@ class BotorchOptimizer(AcquisitionOptimizer): # local search region params local_search_config: Optional[AnyLocalSearchConfig] = None + # NChooseK / semi-continuous pruning hyperparameters. See + # `bofire.strategies.predictives._nchoosek_pruning.prune_nchoosek` for + # full semantics. + per_step_local_reopt: bool = Field( + default=False, + description=( + "When pruning is applicable, refine each per-feature variant " + "via optimize_acqf in addition to the QP projection. More " + "accurate but multiplies cost by the local solver." + ), + ) + final_local_reopt: bool = Field( + default=True, + description=( + "When pruning is applicable, run a single optimize_acqf " + "clean-up after the greedy loop with zeroed and active " + "semi-continuous features pinned." + ), + ) + @field_validator("batch_limit") @classmethod def validate_batch_limit(cls, batch_limit: int, info): diff --git a/bofire/strategies/predictives/_nchoosek_pruning.py b/bofire/strategies/predictives/_nchoosek_pruning.py new file mode 100644 index 000000000..16966836f --- /dev/null +++ b/bofire/strategies/predictives/_nchoosek_pruning.py @@ -0,0 +1,1021 @@ +"""Greedy NChooseK pruning, based on the BONSAI algorithm +(https://arxiv.org/abs/2602.07144). + +This module is a self-contained, side-effect-free implementation of +the greedy pruning post-processing step that turns the output of an +unconstrained acquisition function maximisation into a candidate that +satisfies all NChooseK constraints and per-feature semi-continuity +constraints (``allow_zero=True`` with ``lb > 0``). + +The full conceptual write-up lives in ``docs/pruning.md``. The module +operates on tensors throughout — the caller is responsible for +``domain.inputs.transform`` / ``inverse_transform``, and for building +the ``features2idx``, ``bounds`` and linear-constraint tensors that +the algorithm consumes. + +The module assumes that the columns of ``X``/``bounds`` are aligned +with the indices reported by ``features2idx`` and used inside +``inequality_constraints`` / ``equality_constraints``. For +pure-continuous domains (Phase 1's scope) this holds by construction; +mixing NChooseK with one-hot-encoded categorical features is out of +scope. +""" + +from typing import Any, Dict, List, Optional, Sequence, Set, Tuple + +import torch +from botorch.acquisition.acquisition import AcquisitionFunction +from botorch.optim.optimize import optimize_acqf +from botorch.optim.parameter_constraints import project_to_feasible_space_via_slsqp +from torch import Tensor + +from bofire.data_models.constraints.api import ( + InterpointConstraint, + LinearEqualityConstraint, + LinearInequalityConstraint, + NChooseKConstraint, + NonlinearConstraint, + ProductConstraint, +) +from bofire.data_models.domain.api import Domain +from bofire.data_models.features.api import ContinuousInput + + +class PruningInfeasibleError(RuntimeError): + """Raised when greedy pruning cannot satisfy every NChooseK and + semi-continuity constraint — typically because the per-constraint + ``min_count`` guard empties the action set before all + ``max_count`` constraints are met. + """ + + +# --------------------------------------------------------------------------- +# Domain-level applicability gates +# --------------------------------------------------------------------------- + + +def has_semicontinuous_features(domain: Domain) -> bool: + """True iff any continuous input has ``allow_zero=True`` and a positive + lower bound — i.e., its feasible region per coordinate is the + disconnected union ``{0} ∪ [lb, ub]``. + """ + for feat in domain.inputs.get(ContinuousInput): + assert isinstance(feat, ContinuousInput) + if feat.allow_zero and feat.bounds[0] > 0: + return True + return False + + +def has_nchoosek_linear_overlap(domain: Domain) -> bool: + """True iff any NChooseK feature also appears in a linear (equality or + inequality) constraint. Used to determine whether QP projection is + needed during pruning. + """ + nchoosek_features: Set[str] = set() + for c in domain.constraints.get(NChooseKConstraint): + assert isinstance(c, NChooseKConstraint) + nchoosek_features.update(c.features) + + linear_features: Set[str] = set() + for c in domain.constraints.get( + includes=[LinearEqualityConstraint, LinearInequalityConstraint] + ): + linear_features.update(c.features) + + return bool(nchoosek_features.intersection(linear_features)) + + +def _features_in_blocking_constraints(domain: Domain) -> Set[str]: + blocking: Set[str] = set() + for c in domain.constraints.get( + includes=[ProductConstraint, NonlinearConstraint, InterpointConstraint] + ): + blocking.update(c.features) + return blocking + + +def is_nchoosek_pruning_applicable(domain: Domain) -> bool: + """True iff greedy pruning can be safely applied for the domain's + NChooseK constraints (BONSAI algorithm). + + Pruning is applicable when at least one NChooseK constraint exists and + no NChooseK feature appears in any nonlinear (Product, Nonlinear) or + interpoint constraint. Overlap with linear equality/inequality + constraints is allowed and handled via QP projection inside the + pruning loop. + """ + nchoosek_constraints = domain.constraints.get(NChooseKConstraint) + if len(nchoosek_constraints) == 0: + return False + + blocking = _features_in_blocking_constraints(domain) + for c in nchoosek_constraints: + assert isinstance(c, NChooseKConstraint) + if blocking.intersection(c.features): + return False + return True + + +def semicontinuous_specs_from_domain( + domain: Domain, + features2idx: Dict[str, Tuple[int, ...]], +) -> Dict[int, Tuple[float, float]]: + """Build the ``semicontinuous_specs`` mapping for ``prune_nchoosek``. + + Returns a dict ``{tensor_column_index: (lb, ub)}`` for every continuous + input with ``allow_zero=True`` and a positive lower bound. + """ + specs: Dict[int, Tuple[float, float]] = {} + for feat in domain.inputs.get(ContinuousInput): + assert isinstance(feat, ContinuousInput) + if feat.allow_zero and feat.bounds[0] > 0: + for col in features2idx[feat.key]: + specs[col] = (float(feat.bounds[0]), float(feat.bounds[1])) + return specs + + +def is_pruning_applicable(domain: Domain) -> bool: + """Unified gate: pruning runs if either NChooseK pruning is + applicable, or the domain has standalone semi-continuous features — + and no semi-continuous feature appears in a blocking nonlinear / + interpoint constraint. + """ + if is_nchoosek_pruning_applicable(domain): + return True + if not has_semicontinuous_features(domain): + return False + + blocking = _features_in_blocking_constraints(domain) + for feat in domain.inputs.get(ContinuousInput): + assert isinstance(feat, ContinuousInput) + if feat.allow_zero and feat.bounds[0] > 0 and feat.key in blocking: + return False + return True + + +# --------------------------------------------------------------------------- +# State helpers +# --------------------------------------------------------------------------- + + +def _classify_features_for_row( + x: Tensor, + semicontinuous_specs: Dict[int, Tuple[float, float]], + tol: float, +) -> Tuple[Set[int], Set[int], Set[int]]: + """Partition tensor column indices into (zero, fractional, active). + + A column is *fractional* when the feature is semi-continuous + (``j ∈ semicontinuous_specs``) and the current value lies in the + open interval ``(0, lb_j)``. *Zero* and *active* are defined for + every column. + """ + d = x.shape[0] + zero: Set[int] = set() + fractional: Set[int] = set() + active: Set[int] = set() + for j in range(d): + v = float(x[j].abs().item()) + if v <= tol: + zero.add(j) + continue + if j in semicontinuous_specs: + lb_j, _ = semicontinuous_specs[j] + if v < lb_j - tol: + fractional.add(j) + continue + active.add(j) + return zero, fractional, active + + +def _is_nchoosek_fulfilled( + x: Tensor, + nchoosek_constraints: Sequence[NChooseKConstraint], + features2idx: Dict[str, Tuple[int, ...]], + tol: float = 1e-6, +) -> bool: + """Tensor-native NChooseK fulfilment check for a single candidate. + + Ported from ``BotorchStrategy._nchoosek_fulfilled_tensor``. Counts + non-zero columns per constraint (using ``|x[idx]| > tol``) and + honours ``none_also_valid`` when the count is zero. + """ + for c in nchoosek_constraints: + indices: List[int] = [] + for feat_key in c.features: + indices.extend(features2idx[feat_key]) + count = int((x[indices].abs() > tol).sum().item()) + if count > c.max_count: + return False + if count < c.min_count and not (c.none_also_valid and count == 0): + return False + return True + + +def _active_counts( + x: Tensor, + nchoosek_constraints: Sequence[NChooseKConstraint], + features2idx: Dict[str, Tuple[int, ...]], + tol: float, +) -> Dict[int, int]: + """Per-constraint active counts, keyed by constraint position.""" + counts: Dict[int, int] = {} + for c_idx, c in enumerate(nchoosek_constraints): + indices: List[int] = [] + for feat_key in c.features: + indices.extend(features2idx[feat_key]) + counts[c_idx] = int((x[indices].abs() > tol).sum().item()) + return counts + + +def _violated_constraints( + active_counts: Dict[int, int], + nchoosek_constraints: Sequence[NChooseKConstraint], +) -> Set[int]: + """Constraints whose count exceeds ``max_count`` (positions only).""" + violated: Set[int] = set() + for c_idx, c in enumerate(nchoosek_constraints): + if active_counts[c_idx] > c.max_count: + violated.add(c_idx) + return violated + + +def _zero_action_blocked_by_min_count( + j_idx: int, + active_counts: Dict[int, int], + nchoosek_constraints: Sequence[NChooseKConstraint], + features2idx: Dict[str, Tuple[int, ...]], +) -> bool: + """True iff zeroing ``j_idx`` would push some ``a_c`` below + ``min_count_c`` for any constraint ``c`` containing it. + + ``none_also_valid`` exempts the case where the post-commit count + is exactly zero. + """ + for c_idx, c in enumerate(nchoosek_constraints): + constraint_indices: Set[int] = set() + for feat_key in c.features: + constraint_indices.update(features2idx[feat_key]) + if j_idx not in constraint_indices: + continue + post = active_counts[c_idx] - 1 + if post < c.min_count and not (c.none_also_valid and post == 0): + return True + return False + + +def _features_eligible_for_zero( + fractional: Set[int], + active: Set[int], + violated: Set[int], + nchoosek_constraints: Sequence[NChooseKConstraint], + features2idx: Dict[str, Tuple[int, ...]], +) -> Set[int]: + """Indices admissible as targets of a ``zero`` action. + + Every fractional feature is admissible (regardless of whether it + sits in a violated NChooseK), since the algorithm must resolve + semi-continuity. Active features are admissible only if they + participate in at least one currently violated constraint — + zeroing a feature outside every violated constraint costs AF + without contributing to feasibility. + """ + eligible: Set[int] = set(fractional) + if not violated: + return eligible + + violated_indices: Set[int] = set() + for c_idx in violated: + c = nchoosek_constraints[c_idx] + for feat_key in c.features: + violated_indices.update(features2idx[feat_key]) + eligible.update(active.intersection(violated_indices)) + return eligible + + +# --------------------------------------------------------------------------- +# Variant builders +# --------------------------------------------------------------------------- +# +# Both variant builders accept the loop's current `active_set` and the +# `semicontinuous_specs` mapping. Before the QP projection or +# `optimize_acqf` call, they tighten bounds on every column in +# `(active_set ∩ semicontinuous_specs.keys()) − {j_idx}` to its +# `[lb_i, ub_i]` band: +# +# for i in (active_set & semicontinuous_specs.keys()) - {j_idx}: +# tightened[0, i] = lb_i +# tightened[1, i] = ub_i +# +# This prevents previously-committed active semi-continuous features +# from drifting back into their gap `(0, lb_i)` when the optimiser +# redistributes mass to satisfy shared linear constraints (mixture +# `Σ x = 1`, etc.). Without this, the per-step reopt for action +# `active(j_{k+1})` could push `x_{j_k}` (committed active in iter k) +# below its own `lb`, leaving the state tracker and tensor +# inconsistent until `_final_local_reopt` cleaned it up. +# +# The `− {j_idx}` exclusion preserves the ability to deactivate a +# previously-active feature. If `j_idx ∈ active_set` and is itself +# semi-continuous, we want it pinned to 0 via `fixed_features`, not +# held in `[lb_j, ub_j]` — those two would conflict. The exclusion +# means "the feature being acted on is exempt from the tightening, +# because the action itself defines its target value/bounds." +# +# Consequences: +# - per-step reopt is internally consistent across iterations, +# - `_final_local_reopt` is now genuinely a polish (not a +# feasibility rescue), so `final_local_reopt`'s role is +# well-defined regardless of whether semi-continuous features +# are in play. +# --------------------------------------------------------------------------- +# +# TODO: introduce an "activate-zero" action for non-semi-continuous +# features when `min_count > 0` is violated. +# +# Today the action set only contains: +# - zero(j) for any j currently active and in some violated NChooseK, +# - active(j) for any j ∈ fractional (snap a fractional semi-continuous +# feature into [lb_j, ub_j]). +# +# Both move features *out* of the active set or resolve fractional +# states. Neither moves a currently-zero feature *into* the active set. +# So if the AF maximiser produces a candidate with `a_c < min_count_c` +# for some constraint `c` (with `none_also_valid=False` or `count > 0`) +# and there are no fractional features, the eligibility set +# `_features_eligible_for_zero(...)` collapses to ∅ — fractional is +# empty, no constraint is violated by max_count, no zero action helps. +# The loop bails with `PruningInfeasibleError`. +# +# This regime is reachable any time: +# - `min_count_c > 0` and `none_also_valid_c = False`, AND +# - the AF (e.g. qLogEI on a Map-SAAS posterior) concentrates mass on +# fewer than `min_count_c` features. +# +# It is the dominant failure mode on real formulation domains (mixture +# `Σ x = 1` + NChooseK with `min_count = 6` over 12 features). With +# only `min_count = 0` allowed, BoFire users can't model "you must use +# at least N components in this formulation" — which is the most common +# real cardinality constraint. +# +# Principled fix: add a third action category, mirror of the active +# variant but for non-semi-continuous features. +# +# def _build_activate_variant( +# x_i, j_idx, ub_j, bounds, ineq, eq, acqf, per_step_local_reopt, +# active_set, pinned_zero_indices, fixed_features, tol, +# ) -> Tuple[Tensor, bool]: +# """Snap currently-zero feature j into a positive value. +# Same machinery as `_build_active_variant` but with the per- +# feature lower bound set to `tol` (any positive value) instead +# of the semi-continuous `lb_j`. The QP enforces sum=1 etc. +# The optimiser picks where in `[tol, ub_j]` to place x_j. +# """ +# tightened = bounds.clone() +# tightened[0, j_idx] = tol +# tightened[1, j_idx] = ub_j +# ... # rest mirrors _build_active_variant +# +# Eligibility: +# - activate(j) is admissible only when some `a_c < min_count_c` +# (with the usual `none_also_valid` carve-out) AND `j` is currently +# in `zero_set` AND `j ∈ features(c)` for some min-count-violated +# constraint `c`. +# - Each iteration's action set then becomes +# {zero(j) : eligible_for_zero} ∪ +# {active(j) : j ∈ fractional ∩ semicontinuous_specs} ∪ +# {activate(j) : eligible_for_activate}. +# - Termination unchanged: stop when no fractional, all `a_c` in +# `[min_count_c, max_count_c]`. +# +# Selection rule: same — argmin AF reduction across all action kinds. +# An activate variant typically *increases* AF (we're adding a new +# active feature to a sparse candidate that the surrogate likes), so +# its reduction is small or negative; the greedy will prefer it +# whenever min_count is the binding violation, which is the right +# behaviour. +# +# State updates: +# - activate(j): zero_set.discard(j); active_set.add(j); a_c +# increments for every c ∋ j. +# +# Edge cases: +# - Mutual infeasibility (e.g. one constraint demands ≥ 2 actives, +# another forbids both). The eligibility set still empties out; +# `PruningInfeasibleError` still fires. The error class stays; +# only the *common* min-count-violation case ceases to raise. +# - Interaction with `per_step_local_reopt`: the activate variant +# calls `optimize_acqf` with `bounds[0, j_idx] = tol` (or the +# partial-drift fix's tightened bounds for other actives). Same +# plumbing as the active variant. +# - Cost: the action set size grows by up to |zero_set| in +# min-count-violated iterations. Manageable. +# +# Plumbing requirement: `_features_eligible_for_activate(...)` helper, +# and the loop's per-iteration state needs to track per-constraint +# min-count violations explicitly (currently `_violated_constraints` +# only tracks max-count violations). +# +# Order of operations: do the partial-drift fix above first (cleaner +# foundation, smaller scope), then this. Both fixes touch the variant +# builders' signatures, so doing them together would conflate two +# design discussions in the same diff; sequential is cleaner. +# --------------------------------------------------------------------------- + + +_OPTIMIZE_ACQF_DEFAULTS: Dict[str, Any] = { + "q": 1, + "num_restarts": 1, + "raw_samples": None, +} + + +def _local_optacqf( + initial: Tensor, + bounds: Tensor, + fixed_features: Optional[Dict[int, float]], + inequality_constraints: List[Tuple[Tensor, Tensor, float]], + equality_constraints: List[Tuple[Tensor, Tensor, float]], + acqf: AcquisitionFunction, +) -> Optional[Tensor]: + """Single-restart ``optimize_acqf`` warm-started from ``initial``. + + Returns the refined ``(d,)`` tensor on success or ``None`` on + optimizer failure (caller falls back to ``initial``). + """ + try: + local_candidate, _ = optimize_acqf( + acq_function=acqf, + bounds=bounds, + batch_initial_conditions=initial.unsqueeze(0).unsqueeze(0), + fixed_features=fixed_features, + inequality_constraints=inequality_constraints + if inequality_constraints + else None, + equality_constraints=equality_constraints if equality_constraints else None, + **_OPTIMIZE_ACQF_DEFAULTS, + ) + except Exception: + return None + return local_candidate.squeeze(0).squeeze(0) + + +def _build_zero_variant( + x_i: Tensor, + j_idx: int, + bounds: Tensor, + inequality_constraints: List[Tuple[Tensor, Tensor, float]], + equality_constraints: List[Tuple[Tensor, Tensor, float]], + acqf: AcquisitionFunction, + per_step_local_reopt: bool, + pinned_zero_indices: Optional[Set[int]] = None, + fixed_features: Optional[Dict[int, float]] = None, + active_set: Optional[Set[int]] = None, + semicontinuous_specs: Optional[Dict[int, Tuple[float, float]]] = None, +) -> Tuple[Tensor, bool]: + """Construct the candidate variant where ``x_j`` is forced to zero. + + ``pinned_zero_indices`` lists tensor columns that have been + committed to zero in earlier greedy iterations. They are passed + as additional ``fixed_features`` to the QP projection and (when + enabled) the local re-optimisation, preventing the linear- + constraint redistribution from resurrecting previously-zeroed + features. + + ``fixed_features`` is the caller-supplied mapping of features + that must remain at fixed values throughout pruning (e.g., + inputs whose bounds collapse to a single value). They are merged + into the projection's and local re-optimiser's ``fixed_features`` + so SLSQP and ``optimize_acqf`` never move them. + + ``active_set`` and ``semicontinuous_specs`` together let the + builder tighten bounds for previously-committed active + semi-continuous features (every ``i ∈ active_set ∩ + semicontinuous_specs.keys()`` other than ``j_idx``), so the QP + projection / ``optimize_acqf`` cannot push them back into the + semi-continuity gap ``(0, lb_i)`` while satisfying the linear + constraints. The ``− {j_idx}`` exclusion preserves the ability + to deactivate a previously-active feature: if ``j_idx ∈ + active_set`` and is itself semi-continuous, we want it pinned to + 0 (via ``fixed_features``), not held in ``[lb_j, ub_j]``. + + Returns ``(variant, valid)``: + + - ``valid=True`` if the variant is a feasible point (linear- + constraint compatible); + - ``valid=False`` if the QP projection failed (mutually + infeasible bounds + linear constraints). The variant is still + returned with ``x[j_idx] = 0`` so the caller has a well-formed + tensor; the caller must replace its AF value with ``-inf`` so + it is never selected. + """ + has_linear = bool(inequality_constraints) or bool(equality_constraints) + fixed: Dict[int, float] = dict(fixed_features or {}) + for j in pinned_zero_indices or set(): + fixed[j] = 0.0 + fixed[j_idx] = 0.0 + + tightened = bounds.clone() + if active_set and semicontinuous_specs: + for i in (set(active_set) & semicontinuous_specs.keys()) - {j_idx}: + lb_i, ub_i = semicontinuous_specs[i] + tightened[0, i] = lb_i + tightened[1, i] = ub_i + + if not has_linear: + variant = x_i.clone() + for j, v in fixed.items(): + variant[j] = v + if per_step_local_reopt: + refined = _local_optacqf( + variant, + tightened, + fixed, + inequality_constraints, + equality_constraints, + acqf, + ) + if refined is not None: + variant = refined + return variant, True + + try: + projected = project_to_feasible_space_via_slsqp( + X=x_i.unsqueeze(0), + bounds=tightened, + inequality_constraints=inequality_constraints or None, + equality_constraints=equality_constraints or None, + fixed_features=fixed, + ).squeeze(0) + except Exception: + fallback = x_i.clone() + for j, v in fixed.items(): + fallback[j] = v + return fallback, False + + # Safety net: SLSQP short-circuits when both lists are None, + # and we trust the fixed_features convention rather than the + # solver's idle path. + projected = projected.clone() + for j, v in fixed.items(): + projected[j] = v + + if per_step_local_reopt: + refined = _local_optacqf( + projected, + tightened, + fixed, + inequality_constraints, + equality_constraints, + acqf, + ) + if refined is not None: + return refined, True + + return projected, True + + +def _build_active_variant( + x_i: Tensor, + j_idx: int, + lb_j: float, + ub_j: float, + bounds: Tensor, + inequality_constraints: List[Tuple[Tensor, Tensor, float]], + equality_constraints: List[Tuple[Tensor, Tensor, float]], + acqf: AcquisitionFunction, + per_step_local_reopt: bool, + pinned_zero_indices: Optional[Set[int]] = None, + fixed_features: Optional[Dict[int, float]] = None, + active_set: Optional[Set[int]] = None, + semicontinuous_specs: Optional[Dict[int, Tuple[float, float]]] = None, +) -> Tuple[Tensor, bool]: + """Construct the variant where ``x_j`` is snapped into ``[lb_j, ub_j]``. + + The local bounds are tightened on column ``j_idx`` so the + optimiser is constrained to the active branch of the + semi-continuous feature. Other columns retain the global bounds. + + ``pinned_zero_indices`` keeps previously-zeroed coordinates pinned + via ``fixed_features`` so the projection cannot resurrect them. + ``fixed_features`` is the caller-supplied mapping of features + that must remain at their fixed values throughout pruning. + + ``active_set`` and ``semicontinuous_specs`` together let the + builder also tighten bounds for previously-committed active + semi-continuous features (every ``i ∈ active_set ∩ + semicontinuous_specs.keys()`` other than ``j_idx``), preventing + them from drifting back into ``(0, lb_i)`` during the QP + projection / ``optimize_acqf`` call. The ``− {j_idx}`` exclusion + is harmless here (``j_idx`` is fractional pre-commit, so not in + ``active_set``) but kept for symmetry with the zero-variant + builder. + """ + tightened = bounds.clone() + tightened[0, j_idx] = lb_j + tightened[1, j_idx] = ub_j + if active_set and semicontinuous_specs: + for i in (set(active_set) & semicontinuous_specs.keys()) - {j_idx}: + lb_i, ub_i = semicontinuous_specs[i] + tightened[0, i] = lb_i + tightened[1, i] = ub_i + + starting = x_i.clone() + if float(starting[j_idx].item()) < lb_j: + starting[j_idx] = lb_j + + fixed: Dict[int, float] = dict(fixed_features or {}) + for j in pinned_zero_indices or set(): + fixed[j] = 0.0 + for j, v in fixed.items(): + starting[j] = v + fixed_arg: Optional[Dict[int, float]] = fixed if fixed else None + + has_linear = bool(inequality_constraints) or bool(equality_constraints) + + if not has_linear: + variant = starting + if per_step_local_reopt: + refined = _local_optacqf( + variant, + tightened, + fixed_arg, + inequality_constraints, + equality_constraints, + acqf, + ) + if refined is not None: + variant = refined + return variant, True + + try: + projected = project_to_feasible_space_via_slsqp( + X=starting.unsqueeze(0), + bounds=tightened, + inequality_constraints=inequality_constraints or None, + equality_constraints=equality_constraints or None, + fixed_features=fixed_arg, + ).squeeze(0) + except Exception: + return starting, False + + projected = projected.clone() + for j, v in fixed.items(): + projected[j] = v + + if per_step_local_reopt: + refined = _local_optacqf( + projected, + tightened, + fixed_arg, + inequality_constraints, + equality_constraints, + acqf, + ) + if refined is not None: + return refined, True + + return projected, True + + +def _final_local_reopt( + x: Tensor, + zero_set: Set[int], + active_set: Set[int], + semicontinuous_specs: Dict[int, Tuple[float, float]], + bounds: Tensor, + inequality_constraints: List[Tuple[Tensor, Tensor, float]], + equality_constraints: List[Tuple[Tensor, Tensor, float]], + acqf: AcquisitionFunction, + fixed_features: Optional[Dict[int, float]] = None, +) -> Tensor: + """Run a single local ``optimize_acqf`` with the loop's decisions + frozen. + + Zeroed features are pinned via ``fixed_features``; semi-continuous + features in the active set get bounds tightened to + ``[lb_j, ub_j]``. Caller-supplied ``fixed_features`` are merged in + so they remain pinned during the clean-up. All other coordinates + keep their global bounds. Falls back to ``x`` on optimiser failure. + """ + tightened = bounds.clone() + for j in active_set: + if j in semicontinuous_specs: + lb_j, ub_j = semicontinuous_specs[j] + tightened[0, j] = lb_j + tightened[1, j] = ub_j + + fixed: Dict[int, float] = dict(fixed_features or {}) + for j in zero_set: + fixed[j] = 0.0 + + refined = _local_optacqf( + x, + tightened, + fixed if fixed else None, + inequality_constraints, + equality_constraints, + acqf, + ) + return refined if refined is not None else x + + +# --------------------------------------------------------------------------- +# AF evaluation +# --------------------------------------------------------------------------- + + +def _evaluate_variants_with_prefix( + X_prefix: Tensor, + variants: Tensor, + acqf: AcquisitionFunction, +) -> Tensor: + """Batched AF evaluation. ``X_prefix`` is ``(i, d)`` (possibly + empty); ``variants`` is ``(R, d)``; returns ``(R,)``. + """ + R = variants.shape[0] + if X_prefix.shape[0] > 0: + prefix = X_prefix.unsqueeze(0).expand(R, -1, -1) + eval_X = torch.cat([prefix, variants.unsqueeze(1)], dim=1) + else: + eval_X = variants.unsqueeze(1) + with torch.no_grad(): + return acqf(eval_X) + + +def _af_reduction( + dense_af: Tensor, + variant_af: Tensor, +) -> Tensor: + """Absolute BONSAI AF reduction: ``g_j = α(x_dense) − α(x_pruned_j)``. + + Smaller is better. ``argmin`` picks the variant whose pruning costs the + least acquisition value — the BONSAI greedy rule from + https://arxiv.org/abs/2602.07144. + + Equivalent to ``argmax variant_af`` since ``dense_af`` is the same for + every variant in a single iteration; we keep the explicit subtraction + to mirror the paper's notation. We deliberately do not normalise by + the dense incremental because (a) the paper only normalises in the + termination criterion (the ρ threshold), not in selection, and (b) we + terminate by NChooseK satisfaction rather than ρ, so the + normalisation has no semantic role here. The relative form is also + ill-defined when ``dense_af ≤ base_af`` (qLogEI on a data-starved + candidate, etc.); the absolute form has no such pathology. + """ + return dense_af - variant_af + + +# --------------------------------------------------------------------------- +# Main loop +# --------------------------------------------------------------------------- + + +def prune_nchoosek( + X: Tensor, + acqf: AcquisitionFunction, + nchoosek_constraints: Sequence[NChooseKConstraint], + features2idx: Dict[str, Tuple[int, ...]], + bounds: Tensor, + inequality_constraints: List[Tuple[Tensor, Tensor, float]], + equality_constraints: List[Tuple[Tensor, Tensor, float]], + semicontinuous_specs: Dict[int, Tuple[float, float]], + *, + fixed_features: Optional[Dict[int, float]] = None, + per_step_local_reopt: bool = False, + final_local_reopt: bool = True, + tol: float = 1e-6, +) -> Tensor: + """Greedy BONSAI pruning of a q-batch of acquisition-optimised + candidates. + + For every candidate row ``X[i]``, iteratively commit the action + (zero a feature, or snap a fractional feature into its active + band) whose acquisition reduction is smallest, until every + NChooseK constraint is satisfied and no feature remains in the + semi-continuous gap. Later candidates condition on already-pruned + earlier ones via the prefixed AF evaluation. + + Args: + X: ``(q, d)`` candidate tensor. + acqf: BoTorch acquisition function. Called as ``acqf(X)`` with + ``(b, q', d)`` shaped input. + nchoosek_constraints: NChooseK constraints (filtered list, + i.e. no other constraint types). + features2idx: Mapping from feature key to a tuple of tensor + column indices (one entry per column, length 1 for plain + ``ContinuousInput``). + bounds: ``(2, d)`` tensor of lower and upper bounds. + inequality_constraints: BoTorch-style ``(indices, coefficients, + rhs)`` triples for ``A x >= b``. May be empty. + equality_constraints: same format, for ``A x = b``. May be empty. + semicontinuous_specs: ``{j_idx: (lb_j, ub_j)}`` for every + tensor column whose feature is ``allow_zero=True`` with + ``lb_j > 0``. May be empty when no semi-continuous + features are in the domain. + fixed_features: ``{j_idx: value}`` for caller-supplied features + that must remain at the given value throughout pruning + (e.g., inputs whose bounds collapse to a single value). + These features are never proposed as zero or active + actions, and they are forwarded to both the QP projection + and ``optimize_acqf`` so they cannot drift. May be empty + or ``None``. + per_step_local_reopt: When ``True``, every variant + (zero or active) is locally re-optimised via + ``optimize_acqf`` after the QP projection. This is more + accurate but doubles or triples the per-iteration cost. + When ``False``, the projection result is used directly. + final_local_reopt: When ``True``, after the greedy loop + completes a single ``optimize_acqf`` clean-up pass is run + with zeroed features fixed and active semi-continuous + features bounded to their active band. + tol: Tolerance used for "is this feature zero?" classification + and for the fulfilment check. + + Returns: + ``(q, d)`` tensor of pruned candidates. The tensor is a + clone — the input ``X`` is not mutated. + + Raises: + PruningInfeasibleError: if the greedy loop cannot satisfy all + constraints (typically because the ``min_count`` guard + empties the action set before the ``max_count`` constraints + are met). + """ + # Defensive guard: NChooseK features must each map to a single tensor + # column. The NChooseKConstraint validator already restricts to + # ContinuousInput, but the optimizer's input_preprocessing_specs could + # in principle multi-encode an ill-typed feature; fail loudly if so. + for c in nchoosek_constraints: + for feat_key in c.features: + cols = features2idx.get(feat_key, ()) + if len(cols) != 1: + raise NotImplementedError( + f"Pruning requires every NChooseK feature to map to a " + f"single tensor column. Feature {feat_key!r} maps to " + f"{cols} via features2idx — this typically means it " + f"was one-hot or otherwise multi-column encoded, which " + f"is out of scope for the BONSAI pruning module." + ) + + X = X.clone() + q = X.shape[0] + if q == 0: + return X + + fixed_keys: Set[int] = set((fixed_features or {}).keys()) + + for i in range(q): + zero_set, frac_set, active_set = _classify_features_for_row( + X[i], + semicontinuous_specs, + tol, + ) + # Caller-fixed features are never proposed as actions. + frac_set -= fixed_keys + active_set -= fixed_keys + # Caller-fixed features whose value is exactly zero get + # tracked as part of zero_set for the action eligibility + # filter, but are not emitted by the loop. + for j, v in (fixed_features or {}).items(): + if abs(v) <= tol: + zero_set.add(j) + + while True: + if not frac_set and _is_nchoosek_fulfilled( + X[i], + nchoosek_constraints, + features2idx, + tol, + ): + break + + counts = _active_counts( + X[i], + nchoosek_constraints, + features2idx, + tol, + ) + violated = _violated_constraints(counts, nchoosek_constraints) + eligible = _features_eligible_for_zero( + frac_set, + active_set, + violated, + nchoosek_constraints, + features2idx, + ) + + # Caller-fixed features are never moved by the loop. + eligible = eligible - fixed_keys + + actions: List[Tuple[int, str, Tensor, bool]] = [] + for j in sorted(eligible): + if _zero_action_blocked_by_min_count( + j, + counts, + nchoosek_constraints, + features2idx, + ): + continue + variant, valid = _build_zero_variant( + X[i], + j, + bounds, + inequality_constraints, + equality_constraints, + acqf, + per_step_local_reopt, + pinned_zero_indices=zero_set, + fixed_features=fixed_features, + active_set=active_set, + semicontinuous_specs=semicontinuous_specs, + ) + actions.append((j, "zero", variant, valid)) + + for j in sorted(frac_set): + if j not in semicontinuous_specs: + continue + lb_j, ub_j = semicontinuous_specs[j] + variant, valid = _build_active_variant( + X[i], + j, + lb_j, + ub_j, + bounds, + inequality_constraints, + equality_constraints, + acqf, + per_step_local_reopt, + pinned_zero_indices=zero_set, + fixed_features=fixed_features, + active_set=active_set, + semicontinuous_specs=semicontinuous_specs, + ) + actions.append((j, "active", variant, valid)) + + if not actions: + raise PruningInfeasibleError( + "Greedy pruning could not satisfy NChooseK + " + "semi-continuity constraints: action set empty " + "while constraints still violated." + ) + + variants = torch.stack([a[2] for a in actions]) + af_values = _evaluate_variants_with_prefix(X[:i], variants, acqf) + valid_flags = torch.tensor( + [a[3] for a in actions], + device=af_values.device, + ) + af_values = torch.where( + valid_flags, + af_values, + torch.full_like(af_values, float("-inf")), + ) + + dense_af = acqf(X[: i + 1].unsqueeze(0)).detach() + af_red = _af_reduction(dense_af, af_values) + + # Stable tie-break: smallest af_reduction first; on ties, + # prefer "zero" over "active", then smaller j_idx. + best_k = 0 + best_key = ( + float(af_red[0].item()), + 0 if actions[0][1] == "zero" else 1, + actions[0][0], + ) + for k in range(1, len(actions)): + key = ( + float(af_red[k].item()), + 0 if actions[k][1] == "zero" else 1, + actions[k][0], + ) + if key < best_key: + best_key = key + best_k = k + + j_pick, kind_pick, variant_pick, _ = actions[best_k] + X[i] = variant_pick + + # Update per-candidate state. + if kind_pick == "zero": + frac_set.discard(j_pick) + active_set.discard(j_pick) + zero_set.add(j_pick) + else: + frac_set.discard(j_pick) + active_set.add(j_pick) + + if final_local_reopt: + X[i] = _final_local_reopt( + X[i], + zero_set - fixed_keys, + active_set, + semicontinuous_specs, + bounds, + inequality_constraints, + equality_constraints, + acqf, + fixed_features=fixed_features, + ) + + return X diff --git a/bofire/strategies/predictives/acqf_optimization.py b/bofire/strategies/predictives/acqf_optimization.py index 0a5938397..5faf8c495 100644 --- a/bofire/strategies/predictives/acqf_optimization.py +++ b/bofire/strategies/predictives/acqf_optimization.py @@ -47,6 +47,12 @@ from bofire.data_models.strategies.shortest_path import has_local_search_region from bofire.data_models.types import InputTransformSpecs from bofire.strategies import utils +from bofire.strategies.predictives._nchoosek_pruning import ( + is_nchoosek_pruning_applicable, + is_pruning_applicable, + prune_nchoosek, + semicontinuous_specs_from_domain, +) from bofire.strategies.random import RandomStrategy from bofire.strategies.shortest_path import ShortestPathStrategy from bofire.utils.torch_tools import ( @@ -54,6 +60,7 @@ get_interpoint_constraints, get_linear_constraints, get_nonlinear_constraints, + get_torch_bounds_from_domain, tkwargs, ) @@ -367,6 +374,9 @@ def __init__(self, data_model: BotorchOptimizerDataModel): self.local_search_config = data_model.local_search_config + self.per_step_local_reopt = data_model.per_step_local_reopt + self.final_local_reopt = data_model.final_local_reopt + super().__init__(data_model) def _setup(self): @@ -379,8 +389,21 @@ def _optimize( domain: Domain, experiments: Optional[pd.DataFrame] = None, ) -> pd.DataFrame: + pruning_applicable = is_pruning_applicable(domain) + if pruning_applicable and self.local_search_config is not None: + raise NotImplementedError( + "Local search region BO combined with NChooseK / " + "semi-continuous pruning is not yet supported. Disable " + "either `local_search_config` or the constraints that " + "trigger pruning." + ) + input_preprocessing_specs = self._input_preprocessing_specs(domain) - bounds = utils.get_torch_bounds_from_domain(domain, input_preprocessing_specs) + bounds = get_torch_bounds_from_domain( + domain, + input_preprocessing_specs, + relax_allow_zero=pruning_applicable, + ) # setup local bounds assert experiments is not None @@ -397,6 +420,17 @@ def _optimize( acqfs=acqfs, bounds=bounds, ) + # print(candidates) + + if pruning_applicable: + candidates = self._prune_if_applicable( + candidates=candidates, + acqfs=acqfs, + domain=domain, + bounds=bounds, + ) + + # print(candidates) candidates = self._candidates_tensor_to_dataframe(candidates, domain) @@ -434,6 +468,49 @@ def _optimize( return candidates + def _prune_if_applicable( + self, + candidates: Tensor, + acqfs: List[AcquisitionFunction], + domain: Domain, + bounds: Tensor, + ) -> Tensor: + """Apply BONSAI greedy pruning to the AF-winning candidate tensor. + + Caller must have already verified ``is_pruning_applicable(domain)``. + Reuses the bounds (relaxed for semi-continuous features) the AF + maximiser used, and the same linear-constraint and + fixed-feature accessors. + + Forwards ``acqfs[0]`` to the greedy AF evaluation. Multi-AF + per-candidate pruning is not yet implemented. + # TODO: per-AF pruning for multi-objective. + """ + features2idx = self._features2idx(domain) + inequality_constraints = get_linear_constraints( + domain, constraint=LinearInequalityConstraint + ) + equality_constraints = get_linear_constraints( + domain, constraint=LinearEqualityConstraint + ) + fixed_features = self.get_fixed_features(domain) + semicontinuous_specs = semicontinuous_specs_from_domain(domain, features2idx) + nchoosek_constraints = list(domain.constraints.get(NChooseKConstraint)) + + return prune_nchoosek( + X=candidates, + acqf=acqfs[0], + nchoosek_constraints=nchoosek_constraints, + features2idx=features2idx, + bounds=bounds, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + semicontinuous_specs=semicontinuous_specs, + fixed_features=fixed_features, + per_step_local_reopt=self.per_step_local_reopt, + final_local_reopt=self.final_local_reopt, + ) + def _optimize_acqf_continuous( self, domain: Domain, @@ -469,13 +546,14 @@ def _get_optimizer_options(self, domain: Domain) -> Dict[str, int]: """ assert self.batch_limit is not None + pruning_applicable = is_nchoosek_pruning_applicable(domain) + constraint_types = [ProductConstraint] + if not pruning_applicable: + constraint_types.append(NChooseKConstraint) return { "batch_limit": ( self.batch_limit - if len( - domain.constraints.get([NChooseKConstraint, ProductConstraint]), - ) - == 0 + if len(domain.constraints.get(constraint_types)) == 0 else 1 ), "maxiter": self.maxiter, @@ -487,11 +565,30 @@ def _determine_optimizer(self, domain: Domain, n_acqfs) -> OptimizerEnum: n_categorical_combinations = ( domain.inputs.get_number_of_categorical_combinations() ) + # When pruning is applicable, semi-continuous features + # (`allow_zero=True` with `lb > 0`) are handled by the post-AF + # pruning step rather than by enumerating their on/off states + # at AF-optimisation time. Divide them out of the combination + # count so a pure-continuous semi-continuous domain routes to + # `optimize_acqf` rather than `optimize_acqf_mixed`. + if is_pruning_applicable(domain): + n_semi = sum( + 1 + for f in domain.inputs.get(ContinuousInput) + if isinstance(f, ContinuousInput) + and f.allow_zero + and f.bounds[0] > 0 + and not f.is_fixed() + ) + if n_semi > 0: + n_categorical_combinations //= 2**n_semi if n_categorical_combinations == 1: return OptimizerEnum.OPTIMIZE_ACQF + exclude_nchoosek = is_nchoosek_pruning_applicable(domain) if ( n_categorical_combinations <= ALTERNATING_OPTIMIZER_THRESHOLD - or len(get_nonlinear_constraints(domain)) > 0 + or len(get_nonlinear_constraints(domain, exclude_nchoosek=exclude_nchoosek)) + > 0 ): return OptimizerEnum.OPTIMIZE_ACQF_MIXED return OptimizerEnum.OPTIMIZE_ACQF_MIXED_ALTERNATING @@ -517,7 +614,15 @@ def _get_arguments_for_optimizer( equality_constraints = get_linear_constraints( domain, constraint=LinearEqualityConstraint ) - if len(nonlinear_constraints := get_nonlinear_constraints(domain)) == 0: + exclude_nchoosek = is_nchoosek_pruning_applicable(domain) + if ( + len( + nonlinear_constraints := get_nonlinear_constraints( + domain, exclude_nchoosek=exclude_nchoosek + ) + ) + == 0 + ): ic_generator = None ic_gen_kwargs = {} else: diff --git a/bofire/strategies/random.py b/bofire/strategies/random.py index 2f22bbb3a..d50aea1ee 100644 --- a/bofire/strategies/random.py +++ b/bofire/strategies/random.py @@ -3,11 +3,13 @@ import warnings from collections import Counter from copy import deepcopy -from typing import Dict, List, Optional, Tuple, cast +from typing import Dict, Iterable, List, Optional, Tuple, cast import numpy as np import pandas as pd +import scipy.optimize import torch +from botorch.exceptions.errors import InfeasibilityError from botorch.optim.initializers import sample_q_batches_from_polytope from botorch.optim.parameter_constraints import _generate_unfixed_lin_constraints from pydantic.types import PositiveInt @@ -115,6 +117,106 @@ def _ask(self, candidate_count: PositiveInt) -> pd.DataFrame: n_iters += 1 return pd.concat(valid_samples, ignore_index=True).iloc[:candidate_count] + @staticmethod + def _subset_lp_feasible( + domain: Domain, + active_keys: Iterable[str], + ) -> bool: + """Check whether pinning the non-``active_keys`` continuous features + to zero is compatible with the domain's linear equality and + inequality constraints. + + Solves a small LP with the trivial objective ``min 0`` (only the + constraints matter) via ``scipy.optimize.linprog`` and returns + ``True`` iff scipy reports ``status == 0`` (a feasible point + exists). NChooseK constraints are *not* part of the LP — the + active subset has already been chosen by the caller; the LP only + checks that the resulting reduced polytope is non-empty. + + TODO: re-evaluate against ``botorch.utils.sampling.find_interior_point``. + + ``find_interior_point`` is the LP that ``HitAndRunPolytopeSampler`` + invokes internally and would be the more semantically aligned + check — it searches for a *strictly interior* point with positive + inequality slack and so rejects degenerate polytopes (single + points, lower-dim faces) that hit-and-run can't sample, which + ``linprog(c=0, ...)`` would otherwise let through. Two reasons we + currently roll our own LP instead: + + 1. ``find_interior_point`` has an undocumented precondition: it + adds a slack variable ``s`` (extending ``c`` and ``A_ub`` to + ``d+1`` columns), but does not extend ``A_eq``. Every internal + call site (``PolytopeSampler.find_interior_point`` at + ``botorch/utils/sampling.py:559``) pre-pads ``A_eq`` with a + zero column for ``s``; the free function does not. Passing an + unpadded ``A_eq`` (the natural shape) triggers a scipy size + mismatch ``ValueError``. Using the helper from outside botorch + therefore requires replicating the padding workaround. + 2. ``linprog(c=0, A_eq=..., A_ub=..., bounds=...)`` handles + per-feature bounds, equalities, and inequalities in one call + without any slack-variable plumbing, so the code is shorter. + + The right long-term move is probably to swap to + ``find_interior_point`` (or reuse the ``PolytopeSampler.find_interior_point`` + method directly) once we either upstream the padding fix or + accept the local workaround. Worth revisiting if a domain + produces a feasible-but-degenerate polytope that this LP accepts + but ``HitAndRunPolytopeSampler`` then can't sample from. + """ + active_set = set(active_keys) + keys = domain.inputs.get_keys(ContinuousInput) + n = len(keys) + if n == 0: + return True + key_to_idx = {k: i for i, k in enumerate(keys)} + + bounds: List[Tuple[float, float]] = [] + for k in keys: + if k in active_set: + feat = domain.inputs.get_by_key(k) + assert isinstance(feat, ContinuousInput) + bounds.append((float(feat.lower_bound), float(feat.upper_bound))) + else: + bounds.append((0.0, 0.0)) + + a_eq_rows: List[np.ndarray] = [] + b_eq_vals: List[float] = [] + for c in domain.constraints.get(LinearEqualityConstraint): + assert isinstance(c, LinearEqualityConstraint) + row = np.zeros(n) + for f, coef in zip(c.features, c.coefficients): + if f in key_to_idx: + row[key_to_idx[f]] = coef + a_eq_rows.append(row) + b_eq_vals.append(float(c.rhs)) + + a_ub_rows: List[np.ndarray] = [] + b_ub_vals: List[float] = [] + for c in domain.constraints.get(LinearInequalityConstraint): + assert isinstance(c, LinearInequalityConstraint) + row = np.zeros(n) + for f, coef in zip(c.features, c.coefficients): + if f in key_to_idx: + row[key_to_idx[f]] = coef + a_ub_rows.append(row) + b_ub_vals.append(float(c.rhs)) + + a_eq = np.array(a_eq_rows) if a_eq_rows else None + b_eq = np.array(b_eq_vals) if b_eq_vals else None + a_ub = np.array(a_ub_rows) if a_ub_rows else None + b_ub = np.array(b_ub_vals) if b_ub_vals else None + + result = scipy.optimize.linprog( + c=np.zeros(n), + A_ub=a_ub, + b_ub=b_ub, + A_eq=a_eq, + b_eq=b_eq, + bounds=bounds, + method="highs", + ) + return bool(result.status == 0) + @staticmethod def _get_zeroable_keys(domain: Domain) -> Tuple[set[str], set[str]]: """Collect feature keys that can take the value zero. @@ -185,16 +287,32 @@ def _sample_with_nchooseks( feat = domain.inputs.get_by_key(key=key) assert isinstance(feat, ContinuousInput) feat.bounds = [0.0, 0.0] - samples.append( - self._sample_from_polytope( - domain=domain, - fallback_sampling_method=self.fallback_sampling_method, - n_burnin=self.n_burnin, - n_thinning=self.n_thinning, - seed=self._get_seed(), - n=count * sampling_multiplier, - sampler_kwargs=self.sampler_kwargs, - ), + try: + samples.append( + self._sample_from_polytope( + domain=domain, + fallback_sampling_method=self.fallback_sampling_method, + n_burnin=self.n_burnin, + n_thinning=self.n_thinning, + seed=self._get_seed(), + n=count * sampling_multiplier, + sampler_kwargs=self.sampler_kwargs, + ), + ) + except InfeasibilityError: + # Safety net: the principled LP rejection in + # `sample_valid_nchoosek_features` should already have + # filtered this combination, but numerical edge cases + # (linprog vs. botorch's interior-point solver + # disagreeing) could still slip through. Drop the + # combination and continue. + continue + if not samples: + raise ValueError( + "All sampled NChooseK subsets produced infeasible " + "polytopes. Check the interaction between NChooseK and " + "linear constraints (e.g. a mixture sum=1 forbids the " + "empty subset)." ) samples = pd.concat(samples, axis=0, ignore_index=True) return samples.sample( @@ -262,10 +380,17 @@ def sample_valid_nchoosek_features( weights = [math.comb(len(con.features), k) for k in ks] groups.append((con.features, ks, weights)) con_feature_sets.append(set(con.features)) - _, allow_zero_keys = RandomStrategy._get_zeroable_keys(domain) + zeroable_keys, allow_zero_keys = RandomStrategy._get_zeroable_keys(domain) + zeroable_keys = zeroable_keys | allow_zero_keys for key in allow_zero_keys: groups.append(([key], [0, 1], [1, 1])) + all_continuous_keys = set(domain.inputs.get_keys(ContinuousInput)) + has_linear_constraints = ( + len(domain.constraints.get(LinearEqualityConstraint)) > 0 + or len(domain.constraints.get(LinearInequalityConstraint)) > 0 + ) + results: List[Tuple[str, ...]] = [] for _ in range(n): for _ in range(max_iters): @@ -273,13 +398,23 @@ def sample_valid_nchoosek_features( for features, ks, weights in groups: k = rng.choices(ks, weights=weights, k=1)[0] active.update(rng.sample(features, k)) - if all( + if not all( (con.none_also_valid and len(active & fset) == 0) or con.min_count <= len(active & fset) <= con.max_count for con, fset in zip(nchoosek_cons, con_feature_sets) ): - results.append(tuple(sorted(active))) - break + continue + # Principled rejection: the chosen subset, after pinning all + # other zeroable features to zero, must be compatible with + # the linear equality and inequality constraints. We solve a + # small LP with a zero objective to test feasibility. + if has_linear_constraints: + fixed_to_zero = zeroable_keys - active + variable_keys = all_continuous_keys - fixed_to_zero + if not RandomStrategy._subset_lp_feasible(domain, variable_keys): + continue + results.append(tuple(sorted(active))) + break else: raise ValueError( f"Failed to sample a valid NChooseK combination after " diff --git a/bofire/strategies/utils.py b/bofire/strategies/utils.py index b5def7423..0d9bb063f 100644 --- a/bofire/strategies/utils.py +++ b/bofire/strategies/utils.py @@ -531,16 +531,6 @@ def _do(self, problem, X, **kwargs): return X_corrected -def get_torch_bounds_from_domain( - domain: Domain, input_preprocessing_specs: InputTransformSpecs -) -> torch.Tensor: - """Get the bounds for the optimization problem in the format required by BoTorch.""" - lower, upper = domain.inputs.get_bounds( - specs=input_preprocessing_specs, - ) - return torch.tensor([lower, upper], **tkwargs) - - def get_ga_problem_and_algorithm( data_model: GeneticAlgorithmDataModel, domain: Domain, diff --git a/bofire/utils/torch_tools.py b/bofire/utils/torch_tools.py index 0612a40a1..d049df44a 100644 --- a/bofire/utils/torch_tools.py +++ b/bofire/utils/torch_tools.py @@ -55,6 +55,31 @@ } +def get_torch_bounds_from_domain( + domain: Domain, + input_preprocessing_specs: InputTransformSpecs, + relax_allow_zero: bool = False, +) -> Tensor: + """Get the bounds for the optimization problem in the format required by BoTorch. + + Args: + domain: Optimization problem definition. + input_preprocessing_specs: Per-feature transform specifications. + relax_allow_zero: When True, semi-continuous continuous inputs + (`allow_zero=True` with positive lower bound) report a relaxed + lower bound of 0, exposing the convex relaxation `[0, ub]` to + downstream optimisers. Defaults to False. + + Returns: + A `(2, d)` tensor of lower and upper bounds. + """ + lower, upper = domain.inputs.get_bounds( + specs=input_preprocessing_specs, + relax_allow_zero=relax_allow_zero, + ) + return torch.tensor([lower, upper], **tkwargs) + + def get_linear_constraints( domain: Domain, constraint: Union[Type[LinearEqualityConstraint], Type[LinearInequalityConstraint]], @@ -258,12 +283,15 @@ def product_constraint(indices: Tensor, exponents: Tensor, rhs: float, sign: int def get_nonlinear_constraints( domain: Domain, includes: Optional[List[Type[Constraint]]] = None, + exclude_nchoosek: bool = False, ) -> List[Tuple[Callable[[Tensor], float], bool]]: """Returns a list of callable functions that represent the nonlinear constraints for the given domain that can be processed by botorch. Args: domain (Domain): The domain for which to generate the nonlinear constraints. + includes: List of constraint types to include. Defaults to NChooseK and ProductInequality. + exclude_nchoosek: If True, NChooseK constraints are excluded even if in includes. Returns: List[Callable[[Tensor], float]]: A list of callable functions that take a tensor @@ -276,7 +304,7 @@ def get_nonlinear_constraints( ), "Only NChooseK and ProductInequality constraints are supported." callables = [] - if NChooseKConstraint in includes: + if NChooseKConstraint in includes and not exclude_nchoosek: callables += get_nchoosek_constraints(domain) if ProductInequalityConstraint in includes: callables += get_product_constraints(domain) diff --git a/docs/pruning.md b/docs/pruning.md new file mode 100644 index 000000000..82935eca0 --- /dev/null +++ b/docs/pruning.md @@ -0,0 +1,369 @@ +# Greedy Pruning for NChooseK Constraints + +This document describes the greedy pruning post-processing step used to enforce +NChooseK (cardinality) constraints during acquisition function optimization. The +algorithm is a direct adaptation of the BONSAI procedure +(Daulton et al., *Bayesian Optimization with Natural Simplicity and +Interpretability*, arXiv:2602.07144) to BoFire's setting of mixed continuous / +semi-continuous features and (optionally) overlapping linear constraints. + +## Motivation + +NChooseK constraints are non-convex: they require that, out of a set of +candidate features, at most `max_count` (and at least `min_count`) take a +strictly positive value. Encoding this directly in the acquisition function +optimizer turns the problem into a mixed-integer program, which is expensive, +brittle, and forces the optimizer onto disconnected feasible regions. + +The BONSAI insight is that we don't need to encode the cardinality constraint +inside the optimizer at all. Instead: + +1. Optimize the acquisition function on a *convex relaxation* of the feasible + set, ignoring the cardinality constraint. +2. Post-process the resulting candidate by greedily pushing features off the + active set until the cardinality constraint is satisfied, choosing at each + step the feature whose removal costs the least acquisition value. + +This is the same pattern the original BONSAI paper applies to "deviations from a +default configuration", with the default value here taken to be zero. + +## Setting + +Each feature `x_j` has + +- a lower bound `lb_j ≥ 0` and upper bound `ub_j > lb_j`, +- optionally an `allow_zero` flag, indicating that the value `0` is also + feasible even when `lb_j > 0` (i.e. the feasible set per feature is the + disconnected union `{0} ∪ [lb_j, ub_j]`). + +A feature is called **semi-continuous** when `allow_zero` is true and +`lb_j > 0`. Note that a single semi-continuous feature is structurally a +one-feature NChooseK constraint with `min_count = 0`, `max_count = 1`, +`none_also_valid = True`. The algorithm treats both cases uniformly. + +The feasible set may also include linear equality and inequality constraints +that *overlap* the NChooseK feature set (e.g. a mixture constraint +`Σ x_j = 1` over the same features that participate in the NChooseK). +Nonlinear or interpoint constraints overlapping the NChooseK feature set are +**not supported** by this algorithm — they are blocked at the domain level and +fall back to the standard nonlinear-constrained acquisition optimization. + +## Algorithm + +### Step 0 — Acquisition function optimization on the convex relaxation + +The acquisition function is optimized over the box `[0, ub_j]` for every +feature, with the linear constraints kept active. The semi-continuity gap +`(0, lb_j)` and the NChooseK cardinality constraint are *both* relaxed away at +this stage. The optimizer therefore sees a convex (or at least well-behaved) +feasible set and can use gradients freely. + +The result is a dense candidate `x*` in which every feature falls into one of +three states: + +- **zero:** `x*_j = 0`. No further decision needed. +- **fractional:** `x*_j ∈ (0, lb_j)`. Violates semi-continuity. Must be either + zeroed or snapped into `[lb_j, ub_j]`. +- **cleanly active:** `x*_j ∈ [lb_j, ub_j]`. No semi-continuity issue, but + may still need to be zeroed to satisfy NChooseK. + +Let `a` denote the active count `|{j : x*_j > 0}|`. + +### Step 1 — Greedy pruning loop + +The loop maintains a current candidate that starts at `x*` and, at each +iteration, commits one *action* that moves a single feature into a terminal +state (`0` or `[lb_j, ub_j]`). Two kinds of actions are considered. + +**Zero action `zero(j)`** — pin `x_j = 0` and re-establish feasibility of the +remaining features against the linear constraints. If a feature `j` is currently +active (fractional or cleanly active), this is always available. Committing +`zero(j)` decrements the active count by one. + +**Active action `active(j)`** — only available for currently fractional +features. Set the bounds of `j` to `[lb_j, ub_j]` and re-establish feasibility. +Committing `active(j)` resolves the semi-continuity violation for `j` without +changing the active count. + +For each candidate action, a per-action *variant* candidate is constructed (see +hyperparameters below). The acquisition function is evaluated at each variant. +The action with the **highest acquisition value** (equivalently, the smallest +acquisition reduction relative to the dense AF value) is committed. + +### Step 2 — Termination and the min_count guard + +The loop terminates when *both* of the following hold: + +- no feature is fractional, and +- `min_count ≤ a ≤ max_count`. + +The upper bound `a ≤ max_count` is enforced by exhaustion: once all fractional +features have been resolved, the action set contains only zero variants, and +each commit decrements `a` by one until the bound is met. + +The lower bound `a ≥ min_count` is **not** automatically respected by the +greedy rule, because the rule is budget-blind: it always picks the action with +smallest acquisition loss, even if that action would drop `a` below `min_count`. +This is what the **min_count guard** prevents. + +The guard is a filter applied to the action set at every iteration: a +`zero(j)` action is removed from the set whenever committing it would bring +`a − 1 < min_count`. (When `none_also_valid = True` the guard also tolerates +`a = 0` as a special case.) Active actions are never filtered. If the filtered +action set is empty before the constraint is met — for example, because every +remaining variant is QP-infeasible — the algorithm raises rather than returning +an infeasible candidate. + +### Properties + +- Each iteration commits exactly one action, and each fractional feature can + appear in at most one active action. The loop therefore terminates in at most + `n` iterations, where `n` is the number of features in the NChooseK group. +- Linear feasibility is enforced at every step by construction: each variant is + built with `x_j` either fixed at zero or constrained to `[lb_j, ub_j]`, and + projected onto the linear constraint set. Variants whose projection is + infeasible are dropped from the action set for that iteration. +- The acquisition value is monotonically traded for feasibility in the BONSAI + sense: at every step, the smallest available acquisition loss is incurred. +- The procedure is deterministic given a fixed acquisition function and a + fixed dense candidate `x*`. + +## Multiple NChooseK constraints + +The algorithm extends to domains with several NChooseK constraints — possibly +on disjoint feature sets, possibly overlapping — without changing its +selection rule. Three book-keeping changes are needed. + +### Per-constraint active counts + +Instead of a single active count `a`, the loop tracks an active count `a_c` +for every NChooseK constraint `c`. A commit on a feature `j` updates `a_c` +for every constraint `c` whose feature set contains `j`. Disjoint +constraints update independently; overlapping constraints update in lock-step +on shared features. + +### Termination is a conjunction over all constraints + +The loop terminates when *no* feature is fractional and *every* NChooseK +constraint is satisfied: + +``` +fractional == ∅ AND ∀ c : min_count_c ≤ a_c ≤ max_count_c +``` + +If even one constraint is still violated, the loop continues. + +### The min_count guard becomes per-constraint + +A `zero(j)` action is filtered out of the action set whenever it would push +`a_c` below `min_count_c` for any constraint `c` that contains `j` (with the +usual `none_also_valid` exception that allows `a_c = 0`). Active actions are +never filtered. As before, an empty filtered action set before all +`max_count_c` are satisfied indicates a mutually infeasible configuration: +the algorithm raises rather than returning an infeasible candidate. + +### A useful efficiency restriction + +When NChooseK constraints overlap with cleanly-active features, it is +worthwhile to further restrict the zero-action set to features that +participate in at least one *currently violated* constraint +(`a_c > max_count_c` for some `c ∋ j`). Zeroing a feature outside every +violated constraint costs acquisition value without contributing to +feasibility, and the greedy rule would never pick such an action anyway — +making the restriction a free pruning of the action set rather than a change +in behaviour. + +### Termination bound + +Each commit either resolves a fractional feature (one active or zero +commit per fractional feature) or zeros a cleanly-active feature whose +removal closes a `max_count_c` gap. The total number of iterations is +therefore bounded by + +``` +|fractional| + Σ_c max(0, a_c − max_count_c) +``` + +evaluated at the dense relaxation candidate. + +### Disjoint vs. overlapping feature sets + +Disjoint NChooseK constraints are essentially independent: the loop iterates +over the union of their feature sets, but no commit on one constraint's +features ever affects another constraint's `a_c`. The interesting case is +overlapping constraints, where a single zero commit can simultaneously +reduce the violation of several constraints (the favourable interaction) +or simultaneously block several min_count guards (the failure mode that +forces the algorithm to raise). + +## Hyperparameters + +The algorithm exposes two knobs that trade computational cost against the +fidelity of the per-action acquisition value estimate. + +### `per_step_local_reopt: bool` + +Controls how each action's variant candidate is constructed. + +- **`False` (project only).** When linear constraints overlap the NChooseK + feature set, the dense candidate is QP-projected onto the linear feasible set + with the per-action bound (`x_j = 0` for `zero`, `x_j ∈ [lb_j, ub_j]` for + `active`) applied. The acquisition value is evaluated directly at this + projected point. When no linear overlap exists, the projection collapses to + simple coordinate replacement (zeroing `x_j` or clamping it into `[lb_j, + ub_j]`). +- **`True` (project then locally reoptimize).** After QP projection, the + acquisition function is locally re-optimized starting from the projection, + with the per-action bound retained. The acquisition value of each action then + reflects the *best achievable* AF value under that bound, rather than the + acquisition at an arbitrary feasible warm-start. + +Project-only is cheaper by a factor equal to the per-action local solver cost +and is usually good enough when the acquisition surface is smooth on the scale +of the linear-projection displacement. Local reoptimization is more accurate +and is recommended when the acquisition is sharp near the boundary of `lb_j` +(e.g. peaks immediately above `lb_j`). + +### `final_local_reopt: bool` + +Controls whether the final pruned candidate is locally reoptimized as a +clean-up step. + +- **`False`.** The candidate returned by the greedy loop is returned as-is. + Its components are guaranteed to lie in `{0} ∪ [lb_j, ub_j]` and to satisfy + all linear and NChooseK constraints by construction. +- **`True`.** A single local acquisition optimization is run with the active + set frozen — features that the loop committed to zero stay at zero, and + features that the loop committed to active remain in `[lb_j, ub_j]`. This + cleans up small drifts introduced by the per-step QP projections and lets the + optimizer settle into the local AF maximum within the now-feasible region. + +The cost of `final_local_reopt = True` is one extra local solve, independent of +the size of the NChooseK group, and it is usually worth enabling unless the +acquisition optimizer call is expensive. + +The two knobs are independent. A useful default combination is +`per_step_local_reopt = False, final_local_reopt = True`: cheap per-action +estimates drive the greedy decisions, and a single end-of-pipeline solve +absorbs the accumulated rounding from the QP projections. + +## Worked example + +Consider a domain with four features `x1, x2, x3, x4`, each with bounds +`(0, 1)`, `allow_zero = True` and `lb = 0.2`, subject to + +- a mixture equality `x1 + x2 + x3 + x4 = 1`, +- an NChooseK constraint on `{x1, x2, x3, x4}` with `min_count = 1`, + `max_count = 2`, `none_also_valid = False`. + +### Step 0 + +Acquisition function maximization on the convex relaxation +`x_j ∈ [0, 1], Σ x_j = 1` returns + +``` +x* = (0.60, 0.30, 0.05, 0.05) +``` + +State of features: + +| feature | value | state | +|---------|-------|--------------------| +| x1 | 0.60 | cleanly active | +| x2 | 0.30 | cleanly active | +| x3 | 0.05 | fractional | +| x4 | 0.05 | fractional | + +Active count `a = 4`, exceeding `max_count = 2`. + +### Iteration 1 + +The exit condition fails: `fractional = {x3, x4}` is non-empty and `a = 4` +exceeds the budget. + +The action set has zero variants for every active feature and active variants +for the fractional ones. After QP projection (and, optionally, local +reoptimization) onto the mixture constraint, illustrative acquisition values +are: + +| action | candidate after projection | AF value | +|-------------|----------------------------------|----------| +| zero(x1) | (0.00, 0.60, 0.20, 0.20) | 0.71 | +| zero(x2) | (0.70, 0.00, 0.20, 0.10) | 0.78 | +| zero(x3) | (0.62, 0.32, 0.00, 0.06) | 0.84 | +| zero(x4) | (0.62, 0.32, 0.06, 0.00) | 0.84 | +| active(x3) | (0.55, 0.25, 0.20, 0.00) | 0.82 | +| active(x4) | (0.55, 0.25, 0.00, 0.20) | 0.82 | + +The min_count guard would block any action that brings `a` below 1; here every +zero action leaves `a ≥ 3` so none are filtered. + +The greedy rule picks the action with the highest acquisition value. Both +`zero(x3)` and `zero(x4)` are tied at `0.84`. Break the tie deterministically +(say, in feature order) and commit `zero(x3)`. + +``` +candidate ← (0.62, 0.32, 0.00, 0.06) a = 3, fractional = {x4} +``` + +Notice that `zero(x3)` resolved both the semi-continuity violation on `x3` +*and* freed one NChooseK slot in a single move — the BONSAI rule arbitrates +between the two purposes naturally. + +### Iteration 2 + +The exit condition still fails: `fractional = {x4}` and `a = 3 > 2`. + +Re-evaluate the action set against the new candidate. `x3` is now permanently +zero and is removed from consideration. + +| action | candidate after projection | AF value | +|-------------|----------------------------------|----------| +| zero(x1) | (0.00, 0.70, 0.00, 0.30) | 0.69 | +| zero(x2) | (0.80, 0.00, 0.00, 0.20) | 0.75 | +| zero(x4) | (0.66, 0.34, 0.00, 0.00) | 0.83 | +| active(x4) | (0.55, 0.25, 0.00, 0.20) | 0.82 | + +The greedy rule picks `zero(x4)` at `0.83` and commits. + +``` +candidate ← (0.66, 0.34, 0.00, 0.00) a = 2, fractional = ∅ +``` + +### Iteration 3 + +The exit condition holds: no feature is fractional and `a = 2` lies in +`[min_count, max_count] = [1, 2]`. + +If `final_local_reopt` is enabled, a single local acquisition optimization is +run with `x3 = x4 = 0` fixed and `x1, x2 ∈ [0.2, 1]`. This may slightly refine +the values of `x1` and `x2` while preserving every constraint. Otherwise the +candidate is returned as is: + +``` +final candidate = (0.66, 0.34, 0.00, 0.00) +``` + +The candidate satisfies the mixture equality (`Σ = 1.0`), the per-feature +semi-continuity (`x ∈ {0} ∪ [0.2, 1]` for every coordinate), and the NChooseK +constraint (`a = 2 ≤ max_count = 2` and `a = 2 ≥ min_count = 1`). + +## Comparison to the published BONSAI algorithm + +The algorithm above differs from the published BONSAI procedure in two +respects: + +1. **Default value.** BONSAI prunes deviations from an arbitrary user-supplied + default configuration; here, the default is fixed at zero, which is the + natural reference value for cardinality constraints. +2. **Termination.** BONSAI terminates when the relative acquisition loss + exceeds a user-specified threshold `ρ`. The variant described here + terminates when the NChooseK and semi-continuity constraints are exactly + satisfied, with no AF-loss budget. This makes the procedure feasibility-first + rather than acquisition-loss-first, which is the appropriate trade-off when + the cardinality constraint is a hard requirement of the experimental design. + The min_count guard is the corresponding feasibility-first analogue of + BONSAI's "do not prune below the floor" rule. + +The greedy selection rule itself — at each step, prefer the action with +smallest acquisition reduction — is unchanged. diff --git a/tests/bofire/data_models/domain/test_domain.py b/tests/bofire/data_models/domain/test_domain.py index d5d6d4d4a..f4b97e138 100644 --- a/tests/bofire/data_models/domain/test_domain.py +++ b/tests/bofire/data_models/domain/test_domain.py @@ -549,3 +549,9 @@ def test_domain_to_description_with_constraints(): "\n## Constraints (candidates MUST satisfy all of these)\n" "- 1.0*x1 + 1.0*x2 <= 1.5 — Budget constraint" ) + + +# Note: pruning-applicability and gate tests live in +# tests/bofire/strategies/test_pruning_gate.py since the gate functions +# moved from Domain methods to free functions in +# bofire.strategies.predictives._nchoosek_pruning. diff --git a/tests/bofire/strategies/test_nchoosek_pruning.py b/tests/bofire/strategies/test_nchoosek_pruning.py new file mode 100644 index 000000000..5058139cf --- /dev/null +++ b/tests/bofire/strategies/test_nchoosek_pruning.py @@ -0,0 +1,1314 @@ +"""Unit tests for the standalone NChooseK pruning module. + +Phase 1 verification of `bofire/strategies/predictives/_nchoosek_pruning.py`. +""" + +from typing import Any, Dict, List, Tuple, cast + +import pytest +import torch +from botorch.acquisition.acquisition import AcquisitionFunction +from botorch.utils.testing import MockAcquisitionFunction +from torch import Tensor + +from bofire.data_models.constraints.api import ( + LinearEqualityConstraint, + LinearInequalityConstraint, + NChooseKConstraint, +) +from bofire.data_models.domain.api import Domain +from bofire.data_models.features.api import ContinuousInput, ContinuousOutput +from bofire.data_models.objectives.api import MaximizeObjective +from bofire.strategies.predictives import _nchoosek_pruning as ncp +from bofire.strategies.predictives._nchoosek_pruning import ( + PruningInfeasibleError, + _active_counts, + _build_active_variant, + _build_zero_variant, + _classify_features_for_row, + _features_eligible_for_zero, + _is_nchoosek_fulfilled, + _violated_constraints, + _zero_action_blocked_by_min_count, + prune_nchoosek, +) +from bofire.utils.torch_tools import ( + get_linear_constraints, + get_torch_bounds_from_domain, + tkwargs, +) + + +# --------------------------------------------------------------------------- +# Test acquisition functions +# --------------------------------------------------------------------------- + + +class WeightedSumAcqf: + """Deterministic AF over the q-batch. + + For input ``X`` of shape ``(b, q, d)``, returns ``(b,)`` where each + value is the maximum over the q-batch of ``(X * weights).sum(-1)``. + Useful for tests that need to assert *which* feature was zeroed. + """ + + def __init__(self, weights: Tensor): + self.weights = weights + self.model = None + self.X_pending = None + + def __call__(self, X: Tensor) -> Tensor: + return (X * self.weights).sum(-1).max(dim=-1).values + + def set_X_pending(self, X_pending=None): + self.X_pending = X_pending + + +class ConstantAcqf: + """AF returning a constant for every input — exercises the case + where every variant has the same absolute AF and the tie-break runs. + """ + + def __init__(self, value: float = 1.0): + self.value = value + self.model = None + self.X_pending = None + + def __call__(self, X: Tensor) -> Tensor: + b = X.shape[0] + return torch.full((b,), self.value, **tkwargs) + + def set_X_pending(self, X_pending=None): + self.X_pending = X_pending + + +# --------------------------------------------------------------------------- +# Domain → algorithm-input helper +# --------------------------------------------------------------------------- + + +def _inputs_from_domain(domain: Domain) -> Dict[str, Any]: + """Build the kwargs for ``prune_nchoosek`` from a Domain. + + Mirrors what ``BotorchOptimizer`` will do in Phase 2. + """ + specs: Dict[str, Any] = {} # pure continuous => no encoding + features2idx, _ = domain.inputs._get_transform_info(specs) + + bounds = get_torch_bounds_from_domain(domain, specs) + + inequality_constraints = get_linear_constraints( + domain, constraint=LinearInequalityConstraint + ) + equality_constraints = get_linear_constraints( + domain, constraint=LinearEqualityConstraint + ) + + nchoosek_constraints = list(domain.constraints.get(NChooseKConstraint)) + + semicontinuous_specs: Dict[int, Tuple[float, float]] = {} + for feat in domain.inputs.get(ContinuousInput): + assert isinstance(feat, ContinuousInput) + if feat.allow_zero and feat.bounds[0] > 0: + for col in features2idx[feat.key]: + semicontinuous_specs[col] = ( + float(feat.bounds[0]), + float(feat.bounds[1]), + ) + # Relax the lower bound for AF-optimisation / pruning: + # the convex relaxation feasible region is [0, ub] for + # semi-continuous features. The active variant inside + # prune_nchoosek tightens this back to [lb, ub] per + # feature when needed. + bounds[0, col] = 0.0 + + return { + "nchoosek_constraints": nchoosek_constraints, + "features2idx": features2idx, + "bounds": bounds, + "inequality_constraints": inequality_constraints, + "equality_constraints": equality_constraints, + "semicontinuous_specs": semicontinuous_specs, + } + + +def _row_to_tensor(row: List[float]) -> Tensor: + return torch.tensor(row, **tkwargs) + + +def _stack_to_tensor(rows: List[List[float]]) -> Tensor: + return torch.tensor(rows, **tkwargs) + + +def _make_simple_domain( + *, + n_features: int = 4, + allow_zero: bool = False, + lb: float = 0.0, + ub: float = 1.0, + constraints: List = None, +) -> Domain: + if constraints is None: + constraints = [] + inputs = [ + ContinuousInput( + key=f"x{i + 1}", + bounds=(lb, ub), + allow_zero=allow_zero, + ) + for i in range(n_features) + ] + outputs = [ContinuousOutput(key="y", objective=MaximizeObjective(w=1.0))] + return Domain.from_lists(inputs=inputs, outputs=outputs, constraints=constraints) + + +# --------------------------------------------------------------------------- +# Pure helper tests +# --------------------------------------------------------------------------- + + +class TestPureHelpers: + def _basic_constraint(self, **kw) -> NChooseKConstraint: + defaults = { + "features": ["x1", "x2", "x3"], + "min_count": 1, + "max_count": 2, + "none_also_valid": False, + } + defaults.update(kw) + return NChooseKConstraint(**defaults) # type: ignore + + def _f2i_continuous(self, n: int) -> Dict[str, Tuple[int, ...]]: + return {f"x{i + 1}": (i,) for i in range(n)} + + # ----- _classify_features_for_row ----- + + def test_classify_no_semicontinuous_partitions_zero_and_active(self): + x = _row_to_tensor([0.0, 0.5, 1e-9, 0.7]) + zero, frac, active = _classify_features_for_row(x, {}, tol=1e-6) + assert zero == {0, 2} + assert frac == set() + assert active == {1, 3} + + def test_classify_semicontinuous_value_in_gap_is_fractional(self): + x = _row_to_tensor([0.0, 0.1, 0.5, 0.0]) + semi = {1: (0.2, 1.0), 2: (0.2, 1.0)} + zero, frac, active = _classify_features_for_row(x, semi, tol=1e-6) + assert zero == {0, 3} + assert frac == {1} + assert active == {2} + + def test_classify_value_at_lb_is_active_not_fractional(self): + x = _row_to_tensor([0.2, 0.0, 0.5, 0.0]) + semi = {0: (0.2, 1.0)} + zero, frac, active = _classify_features_for_row(x, semi, tol=1e-6) + assert frac == set() + assert active == {0, 2} + assert zero == {1, 3} + + # ----- _is_nchoosek_fulfilled ----- + + def test_fulfilled_max_count_violated(self): + c = self._basic_constraint(max_count=2) + x = _row_to_tensor([0.5, 0.5, 0.5]) + assert _is_nchoosek_fulfilled(x, [c], self._f2i_continuous(3)) is False + + def test_fulfilled_min_count_violated(self): + c = self._basic_constraint(min_count=2, max_count=3) + x = _row_to_tensor([0.5, 0.0, 0.0]) + assert _is_nchoosek_fulfilled(x, [c], self._f2i_continuous(3)) is False + + def test_fulfilled_none_also_valid_zero_count(self): + c = self._basic_constraint(min_count=1, max_count=2, none_also_valid=True) + x = _row_to_tensor([0.0, 0.0, 0.0]) + assert _is_nchoosek_fulfilled(x, [c], self._f2i_continuous(3)) is True + + def test_fulfilled_none_also_valid_does_not_allow_low_nonzero(self): + c = self._basic_constraint(min_count=2, max_count=3, none_also_valid=True) + x = _row_to_tensor([0.5, 0.0, 0.0]) # one active, below min_count + assert _is_nchoosek_fulfilled(x, [c], self._f2i_continuous(3)) is False + + def test_fulfilled_tol_classifies_small_values_as_zero(self): + c = self._basic_constraint(min_count=1, max_count=2) + x = _row_to_tensor([1e-9, 1e-9, 0.5]) + assert _is_nchoosek_fulfilled(x, [c], self._f2i_continuous(3)) is True + + def test_fulfilled_within_bounds(self): + c = self._basic_constraint(min_count=1, max_count=2) + x = _row_to_tensor([0.5, 0.5, 0.0]) + assert _is_nchoosek_fulfilled(x, [c], self._f2i_continuous(3)) is True + + # ----- _active_counts ----- + + def test_active_counts_per_constraint(self): + c1 = self._basic_constraint(features=["x1", "x2"], max_count=1) + c2 = self._basic_constraint(features=["x2", "x3", "x4"], max_count=2) + x = _row_to_tensor([0.5, 0.5, 0.5, 0.0]) + f2i = self._f2i_continuous(4) + counts = _active_counts(x, [c1, c2], f2i, tol=1e-6) + assert counts == {0: 2, 1: 2} + + # ----- _violated_constraints ----- + + def test_violated_constraints_only_max_violations(self): + c1 = self._basic_constraint(features=["x1", "x2"], max_count=1, min_count=0) + c2 = self._basic_constraint(features=["x3", "x4"], max_count=2, min_count=2) + counts = {0: 2, 1: 1} # c1 violates max, c2 violates min (but excluded) + violated = _violated_constraints(counts, [c1, c2]) + assert violated == {0} + + # ----- _zero_action_blocked_by_min_count ----- + + def test_min_count_guard_blocks_zero(self): + c = self._basic_constraint(min_count=2, max_count=3, none_also_valid=False) + f2i = self._f2i_continuous(3) + counts = {0: 2} # at the floor; zeroing any feature drops below + assert _zero_action_blocked_by_min_count(0, counts, [c], f2i) is True + + def test_min_count_guard_allows_zero_above_floor(self): + c = self._basic_constraint(min_count=2, max_count=3, none_also_valid=False) + f2i = self._f2i_continuous(3) + counts = {0: 3} + assert _zero_action_blocked_by_min_count(0, counts, [c], f2i) is False + + def test_min_count_guard_none_also_valid_allows_drop_to_zero(self): + c = self._basic_constraint(min_count=2, max_count=3, none_also_valid=True) + f2i = self._f2i_continuous(3) + counts = {0: 1} # dropping to zero is allowed by none_also_valid + assert _zero_action_blocked_by_min_count(0, counts, [c], f2i) is False + + def test_min_count_guard_none_also_valid_blocks_drop_to_intermediate(self): + c = self._basic_constraint(min_count=2, max_count=3, none_also_valid=True) + f2i = self._f2i_continuous(3) + counts = {0: 2} # dropping to 1 is blocked, 1 < min_count and not zero + assert _zero_action_blocked_by_min_count(0, counts, [c], f2i) is True + + def test_min_count_guard_ignores_disjoint_constraints(self): + # j_idx=0 (x1) is not in c2 — c2 should not affect the verdict + c1 = self._basic_constraint(features=["x1", "x2"], min_count=0, max_count=2) + c2 = self._basic_constraint(features=["x3", "x4"], min_count=2, max_count=2) + f2i = self._f2i_continuous(4) + counts = {0: 2, 1: 2} + assert _zero_action_blocked_by_min_count(0, counts, [c1, c2], f2i) is False + + # ----- _features_eligible_for_zero ----- + + def test_eligible_includes_all_fractional(self): + c = self._basic_constraint() + f2i = self._f2i_continuous(3) + eligible = _features_eligible_for_zero( + fractional={0}, + active={1}, + violated=set(), + nchoosek_constraints=[c], + features2idx=f2i, + ) + assert eligible == {0} + + def test_eligible_includes_active_in_violated_constraint(self): + c = self._basic_constraint(features=["x1", "x2"], max_count=1) + f2i = self._f2i_continuous(3) + eligible = _features_eligible_for_zero( + fractional=set(), + active={0, 1, 2}, + violated={0}, + nchoosek_constraints=[c], + features2idx=f2i, + ) + # x3 is not in c, so excluded; x1 and x2 are + assert eligible == {0, 1} + + def test_eligible_excludes_active_outside_any_violated(self): + c = self._basic_constraint(features=["x1", "x2"], max_count=2) + f2i = self._f2i_continuous(3) + eligible = _features_eligible_for_zero( + fractional=set(), + active={0, 1, 2}, + violated=set(), + nchoosek_constraints=[c], + features2idx=f2i, + ) + assert eligible == set() + + +# --------------------------------------------------------------------------- +# Variant construction tests +# --------------------------------------------------------------------------- + + +class TestVariantConstruction: + def _bounds(self, n: int = 4) -> Tensor: + return torch.tensor([[0.0] * n, [1.0] * n], **tkwargs) + + def _acqf(self) -> Any: + return cast(AcquisitionFunction, MockAcquisitionFunction()) + + def test_zero_variant_no_constraints_zeroes_in_place(self): + x = _row_to_tensor([0.5, 0.5, 0.5, 0.5]) + variant, valid = _build_zero_variant( + x, + j_idx=2, + bounds=self._bounds(), + inequality_constraints=[], + equality_constraints=[], + acqf=self._acqf(), + per_step_local_reopt=False, + ) + assert valid is True + assert variant[2].item() == pytest.approx(0.0) + # other dims unchanged + for j in (0, 1, 3): + assert variant[j].item() == pytest.approx(0.5) + + def test_zero_variant_with_mixture_eq_projects_feasibly(self): + domain = _make_simple_domain( + n_features=4, + constraints=[ + LinearEqualityConstraint( + features=["x1", "x2", "x3", "x4"], + coefficients=[1.0, 1.0, 1.0, 1.0], + rhs=1.0, + ), + ], + ) + inp = _inputs_from_domain(domain) + x = _row_to_tensor([0.4, 0.3, 0.2, 0.1]) + variant, valid = _build_zero_variant( + x, + j_idx=2, + bounds=inp["bounds"], + inequality_constraints=inp["inequality_constraints"], + equality_constraints=inp["equality_constraints"], + acqf=self._acqf(), + per_step_local_reopt=False, + ) + assert valid is True + assert variant[2].item() == pytest.approx(0.0, abs=1e-6) + assert variant.sum().item() == pytest.approx(1.0, abs=1e-4) + assert (variant >= 0).all() + assert (variant <= 1).all() + + def test_zero_variant_with_inequality_only_respects_constraint(self): + # x1 + x2 <= 0.6 → rewritten by get_linear_constraints as -x1 - x2 >= -0.6 + domain = _make_simple_domain( + n_features=4, + constraints=[ + LinearInequalityConstraint( + features=["x1", "x2"], + coefficients=[1.0, 1.0], + rhs=0.6, + ), + ], + ) + inp = _inputs_from_domain(domain) + x = _row_to_tensor([0.4, 0.4, 0.5, 0.5]) # already infeasible: x1+x2=0.8 > 0.6 + variant, valid = _build_zero_variant( + x, + j_idx=0, + bounds=inp["bounds"], + inequality_constraints=inp["inequality_constraints"], + equality_constraints=inp["equality_constraints"], + acqf=self._acqf(), + per_step_local_reopt=False, + ) + assert valid is True + assert variant[0].item() == pytest.approx(0.0, abs=1e-6) + assert variant[0].item() + variant[1].item() <= 0.6 + 1e-4 + + def test_zero_variant_qp_infeasible_returns_invalid_flag(self): + # x1 + x2 = 1.0 with x2 bounds [0, 0.4]. Zeroing x1 forces + # x2 = 1.0, which violates the upper bound — infeasible. + inputs = [ + ContinuousInput(key="x1", bounds=(0.0, 1.0)), + ContinuousInput(key="x2", bounds=(0.0, 0.4)), + ] + outputs = [ContinuousOutput(key="y", objective=MaximizeObjective(w=1.0))] + domain = Domain.from_lists( + inputs=inputs, + outputs=outputs, + constraints=[ + LinearEqualityConstraint( + features=["x1", "x2"], + coefficients=[1.0, 1.0], + rhs=1.0, + ), + ], + ) + inp = _inputs_from_domain(domain) + x = _row_to_tensor([0.6, 0.4]) + variant, valid = _build_zero_variant( + x, + j_idx=0, + bounds=inp["bounds"], + inequality_constraints=inp["inequality_constraints"], + equality_constraints=inp["equality_constraints"], + acqf=self._acqf(), + per_step_local_reopt=False, + ) + assert valid is False + assert variant[0].item() == pytest.approx(0.0) + + def test_active_variant_with_mixture_eq_snaps_into_band(self): + domain = _make_simple_domain( + n_features=4, + allow_zero=True, + lb=0.2, + ub=1.0, + constraints=[ + LinearEqualityConstraint( + features=["x1", "x2", "x3", "x4"], + coefficients=[1.0, 1.0, 1.0, 1.0], + rhs=1.0, + ), + ], + ) + inp = _inputs_from_domain(domain) + # x3 starts fractional at 0.05; snap to >= 0.2 + x = _row_to_tensor([0.5, 0.4, 0.05, 0.05]) + variant, valid = _build_active_variant( + x, + j_idx=2, + lb_j=0.2, + ub_j=1.0, + bounds=inp["bounds"], + inequality_constraints=inp["inequality_constraints"], + equality_constraints=inp["equality_constraints"], + acqf=self._acqf(), + per_step_local_reopt=False, + ) + assert valid is True + assert variant[2].item() >= 0.2 - 1e-6 + assert variant.sum().item() == pytest.approx(1.0, abs=1e-4) + + def test_active_variant_no_constraints_keeps_starting_or_lb(self): + # No linear constraints — projection short-circuits, we just clamp + # the fractional column to lb. + x = _row_to_tensor([0.5, 0.05, 0.5, 0.5]) + variant, valid = _build_active_variant( + x, + j_idx=1, + lb_j=0.2, + ub_j=1.0, + bounds=torch.tensor([[0.0] * 4, [1.0] * 4], **tkwargs), + inequality_constraints=[], + equality_constraints=[], + acqf=self._acqf(), + per_step_local_reopt=False, + ) + assert valid is True + assert variant[1].item() == pytest.approx(0.2) + + +# --------------------------------------------------------------------------- +# End-to-end pruning tests +# --------------------------------------------------------------------------- + + +class TestPruneNchoosekEndToEnd: + # ----- Single NChooseK without semi-continuity (matches today's path) ----- + + def test_q1_max_count_violation_zeros_smallest_weight(self): + domain = _make_simple_domain( + n_features=4, + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=1, + max_count=2, + none_also_valid=False, + ), + ], + ) + inp = _inputs_from_domain(domain) + # Weights: x2 and x3 have lowest, but only one needs to be zeroed + weights = torch.tensor([10.0, 1.0, 2.0, 8.0], **tkwargs) + acqf = WeightedSumAcqf(weights) + X = _stack_to_tensor([[0.5, 0.5, 0.5, 0.5]]) + out = prune_nchoosek( + X=X, + acqf=cast(AcquisitionFunction, acqf), + **inp, + final_local_reopt=False, + ) + # exactly two non-zero features + nz = (out[0].abs() > 1e-6).sum().item() + assert nz == 2 + # the zeroed features should be the lowest-weight ones (x2 then x3) + assert out[0, 1].item() == pytest.approx(0.0) + assert out[0, 2].item() == pytest.approx(0.0) + + def test_q1_already_feasible_is_noop(self): + domain = _make_simple_domain( + n_features=4, + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=1, + max_count=2, + none_also_valid=False, + ), + ], + ) + inp = _inputs_from_domain(domain) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + X = _stack_to_tensor([[0.5, 0.5, 0.0, 0.0]]) + out = prune_nchoosek( + X=X, + acqf=acqf, + **inp, + final_local_reopt=False, + ) + assert torch.allclose(out, X) + + def test_q2_processes_each_candidate(self): + # Both candidates violate; both get pruned to <= 2 active. + domain = _make_simple_domain( + n_features=4, + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=1, + max_count=2, + none_also_valid=False, + ), + ], + ) + inp = _inputs_from_domain(domain) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + X = _stack_to_tensor( + [ + [0.5, 0.5, 0.5, 0.5], + [0.4, 0.4, 0.4, 0.4], + ] + ) + out = prune_nchoosek( + X=X, + acqf=acqf, + **inp, + final_local_reopt=False, + ) + for i in (0, 1): + nz = (out[i].abs() > 1e-6).sum().item() + assert nz <= 2 + + # ----- Single NChooseK with linear overlap (today's QP path) ----- + + def test_q1_with_mixture_constraint_qp_path(self): + domain = _make_simple_domain( + n_features=4, + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=1, + max_count=2, + none_also_valid=False, + ), + LinearEqualityConstraint( + features=["x1", "x2", "x3", "x4"], + coefficients=[1.0, 1.0, 1.0, 1.0], + rhs=1.0, + ), + ], + ) + inp = _inputs_from_domain(domain) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + X = _stack_to_tensor([[0.4, 0.3, 0.2, 0.1]]) + out = prune_nchoosek( + X=X, + acqf=acqf, + **inp, + final_local_reopt=False, + ) + nz = (out[0].abs() > 1e-6).sum().item() + assert nz <= 2 + assert out[0].sum().item() == pytest.approx(1.0, abs=1e-3) + + # ----- Semi-continuous features ----- + + def test_standalone_semicontinuous_no_nchoosek(self): + # No NChooseK — only semi-continuity to enforce. The fractional + # column must be resolved to either 0 or in [lb, ub]. + domain = _make_simple_domain( + n_features=3, + allow_zero=True, + lb=0.2, + ub=1.0, + ) + inp = _inputs_from_domain(domain) + # x2 starts at 0.1, in the gap (0, 0.2) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + X = _stack_to_tensor([[0.5, 0.1, 0.5]]) + out = prune_nchoosek( + X=X, + acqf=acqf, + **inp, + final_local_reopt=False, + ) + v = float(out[0, 1].abs().item()) + assert v <= 1e-6 or v >= 0.2 - 1e-6 + + def test_semicontinuous_in_nchoosek_zero_variant_wins(self): + # AF heavily favours x1 — fractional x4 should be zeroed (cheap). + domain = _make_simple_domain( + n_features=4, + allow_zero=True, + lb=0.2, + ub=1.0, + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=1, + max_count=2, + none_also_valid=False, + ), + LinearEqualityConstraint( + features=["x1", "x2", "x3", "x4"], + coefficients=[1.0, 1.0, 1.0, 1.0], + rhs=1.0, + ), + ], + ) + inp = _inputs_from_domain(domain) + weights = torch.tensor([10.0, 5.0, 0.5, 0.1], **tkwargs) + acqf = WeightedSumAcqf(weights) + X = _stack_to_tensor([[0.5, 0.4, 0.05, 0.05]]) + out = prune_nchoosek( + X=X, + acqf=cast(AcquisitionFunction, acqf), + **inp, + final_local_reopt=False, + ) + # x4 (lowest weight, fractional) should end at zero + assert out[0, 3].item() == pytest.approx(0.0, abs=1e-3) + # NChooseK satisfied + nz = (out[0].abs() > 1e-6).sum().item() + assert nz <= 2 + # mixture preserved + assert out[0].sum().item() == pytest.approx(1.0, abs=1e-3) + + # ----- Multiple NChooseK ----- + + def test_disjoint_nchoosek_constraints(self): + domain = _make_simple_domain( + n_features=4, + constraints=[ + NChooseKConstraint( + features=["x1", "x2"], + min_count=0, + max_count=1, + none_also_valid=True, + ), + NChooseKConstraint( + features=["x3", "x4"], + min_count=0, + max_count=1, + none_also_valid=True, + ), + ], + ) + inp = _inputs_from_domain(domain) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + X = _stack_to_tensor([[0.5, 0.5, 0.5, 0.5]]) + out = prune_nchoosek( + X=X, + acqf=acqf, + **inp, + final_local_reopt=False, + ) + active_x12 = (out[0, 0:2].abs() > 1e-6).sum().item() + active_x34 = (out[0, 2:4].abs() > 1e-6).sum().item() + assert active_x12 <= 1 + assert active_x34 <= 1 + + def test_overlapping_nchoosek_one_zero_resolves_both(self): + # Both constraints share x2. Zeroing x2 fixes both at once. + # We force this by giving x1, x3 high weights and x2 low weight. + domain = _make_simple_domain( + n_features=3, + constraints=[ + NChooseKConstraint( + features=["x1", "x2"], + min_count=0, + max_count=1, + none_also_valid=True, + ), + NChooseKConstraint( + features=["x2", "x3"], + min_count=0, + max_count=1, + none_also_valid=True, + ), + ], + ) + inp = _inputs_from_domain(domain) + weights = torch.tensor([10.0, 0.1, 10.0], **tkwargs) + acqf = WeightedSumAcqf(weights) + X = _stack_to_tensor([[0.5, 0.5, 0.5]]) + out = prune_nchoosek( + X=X, + acqf=cast(AcquisitionFunction, acqf), + **inp, + final_local_reopt=False, + ) + # x2 should be zeroed — the cheapest, fixes both + assert out[0, 1].item() == pytest.approx(0.0) + # both constraints satisfied + assert (out[0, 0:2].abs() > 1e-6).sum().item() <= 1 + assert (out[0, 1:3].abs() > 1e-6).sum().item() <= 1 + + def test_conflicting_nchoosek_raises_pruning_infeasible_error(self): + # c1 demands min 2 active over {x1, x2}; c2 forbids both. + domain = _make_simple_domain( + n_features=2, + constraints=[ + NChooseKConstraint( + features=["x1", "x2"], + min_count=2, + max_count=2, + none_also_valid=False, + ), + NChooseKConstraint( + features=["x1", "x2"], + min_count=0, + max_count=1, + none_also_valid=True, + ), + ], + ) + inp = _inputs_from_domain(domain) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + X = _stack_to_tensor([[0.5, 0.5]]) + with pytest.raises(PruningInfeasibleError): + prune_nchoosek( + X=X, + acqf=acqf, + **inp, + final_local_reopt=False, + ) + + # ----- min_count guard ----- + + def test_min_count_guard_blocks_zero_below_floor(self): + domain = _make_simple_domain( + n_features=4, + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=2, + max_count=4, + none_also_valid=False, + ), + ], + ) + inp = _inputs_from_domain(domain) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + X = _stack_to_tensor([[0.5, 0.5, 0.0, 0.0]]) # at floor + out = prune_nchoosek( + X=X, + acqf=acqf, + **inp, + final_local_reopt=False, + ) + # nothing should change — no max violation, no fractional + assert torch.allclose(out, X) + + def test_min_count_guard_none_also_valid_allows_drop_to_zero(self): + domain = _make_simple_domain( + n_features=4, + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=2, + max_count=2, + none_also_valid=True, + ), + ], + ) + inp = _inputs_from_domain(domain) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + # Three actives — must zero one. Guard: post=2, OK. + X = _stack_to_tensor([[0.5, 0.5, 0.5, 0.0]]) + out = prune_nchoosek( + X=X, + acqf=acqf, + **inp, + final_local_reopt=False, + ) + nz = (out[0].abs() > 1e-6).sum().item() + assert nz == 2 + + # ----- Hyperparameters ----- + + def test_final_local_reopt_does_not_violate_constraints(self): + domain = _make_simple_domain( + n_features=4, + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=1, + max_count=2, + none_also_valid=False, + ), + LinearEqualityConstraint( + features=["x1", "x2", "x3", "x4"], + coefficients=[1.0, 1.0, 1.0, 1.0], + rhs=1.0, + ), + ], + ) + inp = _inputs_from_domain(domain) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + X = _stack_to_tensor([[0.4, 0.3, 0.2, 0.1]]) + out = prune_nchoosek( + X=X, + acqf=acqf, + **inp, + final_local_reopt=True, + ) + nz = (out[0].abs() > 1e-6).sum().item() + assert nz <= 2 + assert out[0].sum().item() == pytest.approx(1.0, abs=1e-3) + + def test_per_step_local_reopt_runs_without_error(self): + # Smoke-test: per_step_local_reopt=True doesn't break anything. + domain = _make_simple_domain( + n_features=4, + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=1, + max_count=2, + none_also_valid=False, + ), + LinearEqualityConstraint( + features=["x1", "x2", "x3", "x4"], + coefficients=[1.0, 1.0, 1.0, 1.0], + rhs=1.0, + ), + ], + ) + inp = _inputs_from_domain(domain) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + X = _stack_to_tensor([[0.4, 0.3, 0.2, 0.1]]) + out = prune_nchoosek( + X=X, + acqf=acqf, + **inp, + per_step_local_reopt=True, + final_local_reopt=False, + ) + nz = (out[0].abs() > 1e-6).sum().item() + assert nz <= 2 + assert out[0].sum().item() == pytest.approx(1.0, abs=1e-3) + + # ----- Failure paths and edge cases ----- + + def test_q0_returns_unchanged(self): + domain = _make_simple_domain( + n_features=2, + constraints=[ + NChooseKConstraint( + features=["x1", "x2"], + min_count=0, + max_count=1, + none_also_valid=True, + ), + ], + ) + inp = _inputs_from_domain(domain) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + X = torch.empty((0, 2), **tkwargs) + out = prune_nchoosek(X=X, acqf=acqf, **inp, final_local_reopt=False) + assert out.shape == (0, 2) + + def test_constant_af_deterministic_selection(self): + # ConstantAcqf returns the same value for every input → every + # variant has identical absolute AF, every reduction is zero, + # action selection falls back to the deterministic tie-break. + domain = _make_simple_domain( + n_features=4, + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=1, + max_count=2, + none_also_valid=False, + ), + ], + ) + inp = _inputs_from_domain(domain) + acqf = ConstantAcqf(value=0.5) + X = _stack_to_tensor([[0.5, 0.5, 0.5, 0.5]]) + out = prune_nchoosek( + X=X, + acqf=cast(AcquisitionFunction, acqf), + **inp, + final_local_reopt=False, + ) + nz = (out[0].abs() > 1e-6).sum().item() + assert nz == 2 + + def test_qp_failure_does_not_crash(self, monkeypatch): + # Force every QP call to raise; the pruning loop should still + # return a tensor or raise PruningInfeasibleError, never crash. + def fake_project(*args, **kwargs): + raise RuntimeError("QP failed") + + monkeypatch.setattr( + ncp, + "project_to_feasible_space_via_slsqp", + fake_project, + ) + domain = _make_simple_domain( + n_features=4, + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=1, + max_count=2, + none_also_valid=False, + ), + LinearEqualityConstraint( + features=["x1", "x2", "x3", "x4"], + coefficients=[1.0, 1.0, 1.0, 1.0], + rhs=1.0, + ), + ], + ) + inp = _inputs_from_domain(domain) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + X = _stack_to_tensor([[0.4, 0.3, 0.2, 0.1]]) + # Either succeeds with -inf fallback or raises infeasibility — + # but never an unhandled RuntimeError. + try: + prune_nchoosek( + X=X, + acqf=acqf, + **inp, + final_local_reopt=False, + ) + except PruningInfeasibleError: + pass + + def test_argmin_tie_break_returns_first_occurrence(self): + # ConstantAcqf → all af_reductions are zero. Tie-break is + # smallest j_idx with "zero" preferred over "active". With + # 4 active features and max_count=2, we expect the first two + # iterations to zero x1 and x2 (lowest j_idx). + domain = _make_simple_domain( + n_features=4, + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=1, + max_count=2, + none_also_valid=False, + ), + ], + ) + inp = _inputs_from_domain(domain) + acqf = ConstantAcqf(value=1.0) + X = _stack_to_tensor([[0.5, 0.5, 0.5, 0.5]]) + out = prune_nchoosek( + X=X, + acqf=cast(AcquisitionFunction, acqf), + **inp, + final_local_reopt=False, + ) + # x1 and x2 should be zeroed (lowest indices, tie-break) + assert out[0, 0].item() == pytest.approx(0.0) + assert out[0, 1].item() == pytest.approx(0.0) + # remaining are active + assert out[0, 2].item() > 0 + assert out[0, 3].item() > 0 + + # ----- Caller-supplied fixed_features ----- + + def test_fixed_feature_outside_nchoosek_does_not_move(self): + # 5-feature domain: x5 is fixed at 0.7 by the caller. NChooseK + # acts only on x1..x4. Mixture x1+...+x4 = 1 (does not include + # the fixed feature). The pruning should not touch x5. + inputs = [ContinuousInput(key=f"x{i + 1}", bounds=(0.0, 1.0)) for i in range(4)] + inputs.append(ContinuousInput(key="x5", bounds=(0.0, 1.0))) + outputs = [ContinuousOutput(key="y", objective=MaximizeObjective(w=1.0))] + domain = Domain.from_lists( + inputs=inputs, + outputs=outputs, + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=1, + max_count=2, + none_also_valid=False, + ), + LinearEqualityConstraint( + features=["x1", "x2", "x3", "x4"], + coefficients=[1.0, 1.0, 1.0, 1.0], + rhs=1.0, + ), + ], + ) + inp = _inputs_from_domain(domain) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + X = _stack_to_tensor([[0.4, 0.3, 0.2, 0.1, 0.7]]) + out = prune_nchoosek( + X=X, + acqf=acqf, + fixed_features={4: 0.7}, + **inp, + final_local_reopt=True, + ) + assert out[0, 4].item() == pytest.approx(0.7, abs=1e-6) + nz = (out[0, :4].abs() > 1e-6).sum().item() + assert nz <= 2 + assert out[0, :4].sum().item() == pytest.approx(1.0, abs=1e-3) + + def test_fixed_zero_feature_does_not_count_or_move(self): + # x5 is fixed to zero by the caller — it should never be a + # zero-action target (it is already zero by external decree) + # and the final value must remain exactly 0. + inputs = [ContinuousInput(key=f"x{i + 1}", bounds=(0.0, 1.0)) for i in range(5)] + outputs = [ContinuousOutput(key="y", objective=MaximizeObjective(w=1.0))] + domain = Domain.from_lists( + inputs=inputs, + outputs=outputs, + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=1, + max_count=2, + none_also_valid=False, + ), + ], + ) + inp = _inputs_from_domain(domain) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + X = _stack_to_tensor([[0.5, 0.5, 0.5, 0.5, 0.0]]) + out = prune_nchoosek( + X=X, + acqf=acqf, + fixed_features={4: 0.0}, + **inp, + final_local_reopt=True, + ) + assert out[0, 4].item() == pytest.approx(0.0) + nz = (out[0, :4].abs() > 1e-6).sum().item() + assert nz <= 2 + + def test_fixed_feature_inside_nchoosek_excluded_from_actions(self): + # x1 is fixed at 0.5 even though it appears in the NChooseK + # constraint. The pruning must exclude it from the action set + # (cannot zero a fixed feature) and instead zero among + # {x2, x3, x4}. + domain = _make_simple_domain( + n_features=4, + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=1, + max_count=2, + none_also_valid=False, + ), + ], + ) + inp = _inputs_from_domain(domain) + # Use weights that would otherwise make x1 the prime zero + # candidate (smallest weight). The fixed-features filter must + # override. + weights = torch.tensor([0.1, 5.0, 5.0, 10.0], **tkwargs) + acqf = WeightedSumAcqf(weights) + X = _stack_to_tensor([[0.5, 0.5, 0.5, 0.5]]) + out = prune_nchoosek( + X=X, + acqf=cast(AcquisitionFunction, acqf), + fixed_features={0: 0.5}, + **inp, + final_local_reopt=False, + ) + assert out[0, 0].item() == pytest.approx(0.5) + nz = (out[0].abs() > 1e-6).sum().item() + assert nz <= 2 + + +# --------------------------------------------------------------------------- +# Partial-drift fix tests +# --------------------------------------------------------------------------- + + +class TestPartialDriftFix: + """Verify that variant builders tighten bounds for previously-committed + active semi-continuous features so they cannot drift into the + semi-continuity gap during a later iteration's optimize_acqf or QP + projection. Covers both `_build_zero_variant` and + `_build_active_variant`, plus the deactivation case (the `−{j_idx}` + exclusion) and an end-to-end `prune_nchoosek` run with + `final_local_reopt=False`. + """ + + def _setup_mixture_4(self): + """4 semi-continuous features (lb=0.2, ub=1) in a mixture + Σ x = 1. Returns (bounds, ineq, eq, semi_specs) ready for use + with the variant builders. + """ + d = 4 + bounds = torch.tensor( + [[0.0] * d, [1.0] * d], + **tkwargs, + ) + # Σ x = 1, expressed in get_linear_constraints' format + # (indices, -coefficients, -rhs): + eq = [ + ( + torch.tensor([0, 1, 2, 3]), + torch.tensor([-1.0, -1.0, -1.0, -1.0], **tkwargs), + -1.0, + ) + ] + ineq: list = [] + semi = {i: (0.2, 1.0) for i in range(d)} + return bounds, ineq, eq, semi + + def test_zero_variant_keeps_active_semi_in_band(self): + """Building a zero variant for a feature while two semi-continuous + features are committed-active: those committed actives must stay + in `[0.2, 1]` after the projection, not drift into `(0, 0.2)`. + """ + bounds, ineq, eq, semi = self._setup_mixture_4() + # x_1 and x_2 are already in their active bands [0.2, 1]; x_0 + # is the deactivation target; x_3 is zero. The candidate is + # mixture-feasible (sums to 1). + x_i = torch.tensor([0.2, 0.4, 0.4, 0.0], **tkwargs) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + variant, valid = _build_zero_variant( + x_i, + j_idx=0, + bounds=bounds, + inequality_constraints=ineq, + equality_constraints=eq, + acqf=acqf, + per_step_local_reopt=False, + active_set={1, 2}, + semicontinuous_specs=semi, + ) + assert valid is True + # x_0 was the target → 0 + assert variant[0].item() == pytest.approx(0.0, abs=1e-6) + # x_1 and x_2 were committed-active → must stay in [0.2, 1] + assert variant[1].item() >= 0.2 - 1e-6 + assert variant[1].item() <= 1.0 + 1e-6 + assert variant[2].item() >= 0.2 - 1e-6 + assert variant[2].item() <= 1.0 + 1e-6 + # mixture preserved + assert variant.sum().item() == pytest.approx(1.0, abs=1e-3) + + def test_active_variant_keeps_other_active_semi_in_band(self): + """Building an active variant for a fractional feature while + another semi-continuous feature is already committed-active: + the committed-active one must keep its band. + """ + bounds, ineq, eq, semi = self._setup_mixture_4() + # x_1 already active at 0.4; x_2 is fractional (0.1) and we'll + # commit it active. x_0 free, x_3 zero. + x_i = torch.tensor([0.5, 0.4, 0.1, 0.0], **tkwargs) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + variant, valid = _build_active_variant( + x_i, + j_idx=2, + lb_j=0.2, + ub_j=1.0, + bounds=bounds, + inequality_constraints=ineq, + equality_constraints=eq, + acqf=acqf, + per_step_local_reopt=False, + active_set={1}, + semicontinuous_specs=semi, + ) + assert valid is True + # x_2 (target) snapped into [0.2, 1] + assert variant[2].item() >= 0.2 - 1e-6 + # x_1 (committed-active) stays in [0.2, 1] + assert variant[1].item() >= 0.2 - 1e-6 + assert variant[1].item() <= 1.0 + 1e-6 + # mixture preserved + assert variant.sum().item() == pytest.approx(1.0, abs=1e-3) + + def test_zero_variant_can_deactivate_active_semi(self): + """Building a zero variant for a feature that is itself + currently in `active_set` and is semi-continuous: the `− {j_idx}` + exclusion must let it leave its band and pin to 0, while *other* + active semi features stay locked. + """ + bounds, ineq, eq, semi = self._setup_mixture_4() + # Three actives all at 0.33; x_3 zero. We deactivate x_2. + x_i = torch.tensor( + [0.33, 0.34, 0.33, 0.0], + **tkwargs, + ) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + variant, valid = _build_zero_variant( + x_i, + j_idx=2, + bounds=bounds, + inequality_constraints=ineq, + equality_constraints=eq, + acqf=acqf, + per_step_local_reopt=False, + active_set={0, 1, 2}, + semicontinuous_specs=semi, + ) + assert valid is True + # x_2 (the deactivation target) pinned to 0 — exclusion did its + # job; we did NOT keep x_2 ≥ 0.2. + assert variant[2].item() == pytest.approx(0.0, abs=1e-6) + # x_0 and x_1 (still active) stay in [0.2, 1] + assert variant[0].item() >= 0.2 - 1e-6 + assert variant[1].item() >= 0.2 - 1e-6 + # mixture preserved + assert variant.sum().item() == pytest.approx(1.0, abs=1e-3) + + def test_pruning_no_drift_with_flr_false(self): + """End-to-end: a 4-feature semi-continuous mixture domain. + With `final_local_reopt=False`, every coordinate of the final + candidate must lie in `{0} ∪ [lb, ub]`. Without the partial- + drift fix the per-step reopts could leave coordinates in + `(0, lb)` and `flr=False` wouldn't catch them. + """ + domain = _make_simple_domain( + n_features=4, + allow_zero=True, + lb=0.2, + ub=1.0, + constraints=[ + LinearEqualityConstraint( + features=["x1", "x2", "x3", "x4"], + coefficients=[1.0, 1.0, 1.0, 1.0], + rhs=1.0, + ), + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=0, + max_count=4, + none_also_valid=True, + ), + ], + ) + inp = _inputs_from_domain(domain) + acqf = cast(AcquisitionFunction, MockAcquisitionFunction()) + # Initial candidate: all four equal at 0.25 (every coordinate + # is fractional, in (0, 0.2)? Actually 0.25 > 0.2 so they're + # cleanly active. Use a fractional starting point instead so + # the active variant is exercised.) + X = _stack_to_tensor([[0.1, 0.4, 0.4, 0.1]]) + out = prune_nchoosek( + X=X, + acqf=acqf, + **inp, + final_local_reopt=False, + ) + for j in range(4): + v = float(out[0, j].abs().item()) + assert v <= 1e-6 or v >= 0.2 - 1e-6, ( + f"x_{j + 1}={v} fell into the (0, 0.2) gap; " + "partial-drift fix didn't keep it in band." + ) diff --git a/tests/bofire/strategies/test_pruning_gate.py b/tests/bofire/strategies/test_pruning_gate.py new file mode 100644 index 000000000..6584ef671 --- /dev/null +++ b/tests/bofire/strategies/test_pruning_gate.py @@ -0,0 +1,367 @@ +"""Tests for the domain-level pruning applicability gates. + +These functions live in `bofire.strategies.predictives._nchoosek_pruning` +as free functions (moved out of `Domain` in Phase 2 of the BONSAI +pruning refactor). +""" + +from bofire.data_models.constraints.api import ( + LinearEqualityConstraint, + LinearInequalityConstraint, + NChooseKConstraint, + ProductInequalityConstraint, +) +from bofire.data_models.domain.api import Domain +from bofire.data_models.features.api import ContinuousInput +from bofire.strategies.predictives._nchoosek_pruning import ( + has_nchoosek_linear_overlap, + has_semicontinuous_features, + is_nchoosek_pruning_applicable, + is_pruning_applicable, +) + + +class TestIsNchoosekPruningApplicable: + def test_no_nchoosek(self): + domain = Domain.from_lists( + inputs=[ + ContinuousInput(key="x1", bounds=(0, 1)), + ContinuousInput(key="x2", bounds=(0, 1)), + ], + ) + assert is_nchoosek_pruning_applicable(domain) is False + + def test_nchoosek_only(self): + domain = Domain.from_lists( + inputs=[ + ContinuousInput(key="x1", bounds=(0, 1)), + ContinuousInput(key="x2", bounds=(0, 1)), + ContinuousInput(key="x3", bounds=(0, 1)), + ], + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3"], + min_count=0, + max_count=2, + none_also_valid=True, + ), + ], + ) + assert is_nchoosek_pruning_applicable(domain) is True + + def test_nchoosek_with_overlapping_linear_inequality(self): + domain = Domain.from_lists( + inputs=[ + ContinuousInput(key="x1", bounds=(0, 1)), + ContinuousInput(key="x2", bounds=(0, 1)), + ContinuousInput(key="x3", bounds=(0, 1)), + ], + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3"], + min_count=0, + max_count=2, + none_also_valid=True, + ), + LinearInequalityConstraint( + features=["x1", "x2"], + coefficients=[1.0, 1.0], + rhs=1.0, + ), + ], + ) + assert is_nchoosek_pruning_applicable(domain) is True + assert has_nchoosek_linear_overlap(domain) is True + + def test_nchoosek_with_overlapping_linear_equality(self): + domain = Domain.from_lists( + inputs=[ + ContinuousInput(key="x1", bounds=(0, 1)), + ContinuousInput(key="x2", bounds=(0, 1)), + ContinuousInput(key="x3", bounds=(0, 1)), + ], + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3"], + min_count=0, + max_count=2, + none_also_valid=True, + ), + LinearEqualityConstraint( + features=["x1", "x2"], + coefficients=[1.0, 1.0], + rhs=1.0, + ), + ], + ) + assert is_nchoosek_pruning_applicable(domain) is True + assert has_nchoosek_linear_overlap(domain) is True + + def test_nchoosek_with_partial_linear_overlap(self): + domain = Domain.from_lists( + inputs=[ + ContinuousInput(key="x1", bounds=(0, 1)), + ContinuousInput(key="x2", bounds=(0, 1)), + ContinuousInput(key="x3", bounds=(0, 1)), + ContinuousInput(key="x4", bounds=(0, 1)), + ], + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3"], + min_count=0, + max_count=2, + none_also_valid=True, + ), + LinearInequalityConstraint( + features=["x3", "x4"], + coefficients=[1.0, 1.0], + rhs=1.0, + ), + ], + ) + assert is_nchoosek_pruning_applicable(domain) is True + assert has_nchoosek_linear_overlap(domain) is True + + def test_nchoosek_with_truly_disjoint_linear(self): + domain = Domain.from_lists( + inputs=[ + ContinuousInput(key="x1", bounds=(0, 1)), + ContinuousInput(key="x2", bounds=(0, 1)), + ContinuousInput(key="x3", bounds=(0, 1)), + ContinuousInput(key="x4", bounds=(0, 1)), + ], + constraints=[ + NChooseKConstraint( + features=["x1", "x2"], + min_count=0, + max_count=1, + none_also_valid=True, + ), + LinearInequalityConstraint( + features=["x3", "x4"], + coefficients=[1.0, 1.0], + rhs=1.0, + ), + ], + ) + assert is_nchoosek_pruning_applicable(domain) is True + assert has_nchoosek_linear_overlap(domain) is False + + def test_nchoosek_blocked_by_product_constraint(self): + domain = Domain.from_lists( + inputs=[ + ContinuousInput(key="x1", bounds=(0, 1)), + ContinuousInput(key="x2", bounds=(0, 1)), + ContinuousInput(key="x3", bounds=(0, 1)), + ], + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3"], + min_count=0, + max_count=2, + none_also_valid=True, + ), + ProductInequalityConstraint( + features=["x1", "x2"], + exponents=[1.0, 1.0], + rhs=0.5, + sign=1, + ), + ], + ) + assert is_nchoosek_pruning_applicable(domain) is False + + def test_multiple_nchoosek_one_overlapping_linear(self): + domain = Domain.from_lists( + inputs=[ContinuousInput(key=f"x{i}", bounds=(0, 1)) for i in range(1, 5)], + constraints=[ + NChooseKConstraint( + features=["x1", "x2"], + min_count=0, + max_count=1, + none_also_valid=True, + ), + NChooseKConstraint( + features=["x3", "x4"], + min_count=0, + max_count=1, + none_also_valid=True, + ), + LinearInequalityConstraint( + features=["x3", "x4"], + coefficients=[1.0, 1.0], + rhs=1.0, + ), + ], + ) + assert is_nchoosek_pruning_applicable(domain) is True + + def test_multiple_nchoosek_blocked_by_product(self): + domain = Domain.from_lists( + inputs=[ContinuousInput(key=f"x{i}", bounds=(0, 1)) for i in range(1, 5)], + constraints=[ + NChooseKConstraint( + features=["x1", "x2"], + min_count=0, + max_count=1, + none_also_valid=True, + ), + NChooseKConstraint( + features=["x3", "x4"], + min_count=0, + max_count=1, + none_also_valid=True, + ), + ProductInequalityConstraint( + features=["x3", "x4"], + exponents=[1.0, 1.0], + rhs=0.5, + sign=1, + ), + ], + ) + assert is_nchoosek_pruning_applicable(domain) is False + + def test_multiple_nchoosek_no_blockers(self): + domain = Domain.from_lists( + inputs=[ContinuousInput(key=f"x{i}", bounds=(0, 1)) for i in range(1, 5)], + constraints=[ + NChooseKConstraint( + features=["x1", "x2"], + min_count=0, + max_count=1, + none_also_valid=True, + ), + NChooseKConstraint( + features=["x3", "x4"], + min_count=0, + max_count=1, + none_also_valid=True, + ), + ], + ) + assert is_nchoosek_pruning_applicable(domain) is True + + +class TestHasSemicontinuousFeatures: + def test_no_continuous_inputs_with_allow_zero(self): + domain = Domain.from_lists( + inputs=[ + ContinuousInput(key="x1", bounds=(0.0, 1.0)), + ContinuousInput(key="x2", bounds=(0.0, 1.0), allow_zero=False), + ], + ) + assert has_semicontinuous_features(domain) is False + + def test_allow_zero_with_positive_lb(self): + domain = Domain.from_lists( + inputs=[ + ContinuousInput(key="x1", bounds=(0.0, 1.0)), + ContinuousInput(key="x2", bounds=(0.2, 1.0), allow_zero=True), + ], + ) + assert has_semicontinuous_features(domain) is True + + +class TestIsPruningApplicable: + def test_neither_trigger(self): + domain = Domain.from_lists( + inputs=[ + ContinuousInput(key="x1", bounds=(0, 1)), + ContinuousInput(key="x2", bounds=(0, 1)), + ], + ) + assert is_pruning_applicable(domain) is False + + def test_nchoosek_trigger(self): + domain = Domain.from_lists( + inputs=[ + ContinuousInput(key="x1", bounds=(0, 1)), + ContinuousInput(key="x2", bounds=(0, 1)), + ], + constraints=[ + NChooseKConstraint( + features=["x1", "x2"], + min_count=0, + max_count=1, + none_also_valid=True, + ), + ], + ) + assert is_pruning_applicable(domain) is True + + def test_semicontinuous_trigger(self): + domain = Domain.from_lists( + inputs=[ + ContinuousInput(key="x1", bounds=(0.2, 1.0), allow_zero=True), + ContinuousInput(key="x2", bounds=(0.0, 1.0)), + ], + ) + assert is_pruning_applicable(domain) is True + + def test_both_triggers(self): + domain = Domain.from_lists( + inputs=[ + ContinuousInput(key="x1", bounds=(0.2, 1.0), allow_zero=True), + ContinuousInput(key="x2", bounds=(0.0, 1.0)), + ContinuousInput(key="x3", bounds=(0.0, 1.0)), + ], + constraints=[ + NChooseKConstraint( + features=["x1", "x2", "x3"], + min_count=0, + max_count=2, + none_also_valid=True, + ), + ], + ) + assert is_pruning_applicable(domain) is True + + def test_semicontinuous_blocked_by_product(self): + domain = Domain.from_lists( + inputs=[ + ContinuousInput(key="x1", bounds=(0.2, 1.0), allow_zero=True), + ContinuousInput(key="x2", bounds=(0.0, 1.0)), + ], + constraints=[ + ProductInequalityConstraint( + features=["x1", "x2"], + exponents=[1.0, 1.0], + rhs=0.5, + sign=1, + ), + ], + ) + # x1 (semicontinuous) is in a product constraint → not applicable + assert is_pruning_applicable(domain) is False + + def test_nchoosek_blocked_overrides_semicontinuous(self): + # Even though semi-continuous is present, the overall pruning is + # blocked when NChooseK is itself blocked (defensive: caller should + # not encounter this in practice). + domain = Domain.from_lists( + inputs=[ + ContinuousInput(key="x1", bounds=(0.2, 1.0), allow_zero=True), + ContinuousInput(key="x2", bounds=(0.0, 1.0)), + ContinuousInput(key="x3", bounds=(0.0, 1.0)), + ], + constraints=[ + NChooseKConstraint( + features=["x2", "x3"], + min_count=0, + max_count=1, + none_also_valid=True, + ), + ProductInequalityConstraint( + features=["x2", "x3"], + exponents=[1.0, 1.0], + rhs=0.5, + sign=1, + ), + ], + ) + # NChooseK pruning blocked. But semi-continuous x1 is not in any + # blocking constraint, so the semi-continuous path still allows + # pruning. + assert is_nchoosek_pruning_applicable(domain) is False + assert is_pruning_applicable(domain) is True diff --git a/tests/bofire/strategies/test_sobo.py b/tests/bofire/strategies/test_sobo.py index be074d64e..43f106188 100644 --- a/tests/bofire/strategies/test_sobo.py +++ b/tests/bofire/strategies/test_sobo.py @@ -48,6 +48,10 @@ from bofire.data_models.strategies.predictives.acqf_optimization import LSRBO from bofire.data_models.unions import to_list from bofire.strategies.api import CustomSoboStrategy, RandomStrategy, SoboStrategy +from bofire.strategies.predictives._nchoosek_pruning import ( + has_nchoosek_linear_overlap, + is_nchoosek_pruning_applicable, +) def test_SOBO_not_fitted(): @@ -547,9 +551,10 @@ def test_sobo_get_optimizer_options(): acquisition_optimizer=BotorchOptimizer(maxiter=500, batch_limit=4), ) strategy = SoboStrategy(data_model=strategy_data) + # NChooseK-only domain: pruning is applicable, so batch_limit is not forced to 1 assert strategy.acqf_optimizer._get_optimizer_options(strategy.domain) == { "maxiter": 500, - "batch_limit": 1, + "batch_limit": 4, } @@ -564,3 +569,149 @@ def test_sobo_interpoint(): strategy = SoboStrategy(data_model=strategy_data) strategy.tell(experiments) strategy.ask(2) + + +# test_nchoosek_fulfilled_tensor moved to +# tests/bofire/strategies/test_nchoosek_pruning.py:TestPureHelpers +# (the helper now lives in _nchoosek_pruning.py rather than as a static +# method on BotorchStrategy). + + +def test_nchoosek_pruning_sobo(): + """End-to-end test: SOBO with NChooseK constraint produces valid candidates.""" + domain = Domain( + inputs=[ # type: ignore + ContinuousInput(key="x1", bounds=(0, 1)), + ContinuousInput(key="x2", bounds=(0, 1)), + ContinuousInput(key="x3", bounds=(0, 1)), + ContinuousInput(key="x4", bounds=(0, 1)), + ], + outputs=[ContinuousOutput(key="y", objective=MaximizeObjective(w=1.0))], # type: ignore + constraints=[ # type: ignore + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=1, + max_count=2, + none_also_valid=False, + ), + ], + ) + assert is_nchoosek_pruning_applicable(domain) is True + + # Generate some training data + random_strategy = RandomStrategy( + data_model=RandomStrategyDataModel(domain=domain), + ) + experiments = random_strategy.ask(10) + experiments["y"] = np.random.default_rng(42).standard_normal(10) + experiments["valid_y"] = 1 + + strategy_data = data_models.SoboStrategy(domain=domain) + strategy = SoboStrategy(data_model=strategy_data) + strategy.tell(experiments) + + candidates = strategy.ask(1, raise_validation_error=False) + + # Check that NChooseK constraint is fulfilled + nchoosek = domain.constraints.get(NChooseKConstraint)[0] + assert isinstance(nchoosek, NChooseKConstraint) + assert nchoosek.is_fulfilled(candidates[domain.inputs.get_keys()]).all() + + +def test_nchoosek_pruning_with_mixture_constraint(): + """End-to-end test: SOBO with NChooseK + linear equality (mixture) constraint. + + This tests the QP + local optimize_acqf pruning path. + """ + from bofire.data_models.constraints.api import LinearEqualityConstraint + + domain = Domain( + inputs=[ # type: ignore + ContinuousInput(key="x1", bounds=(0, 1)), + ContinuousInput(key="x2", bounds=(0, 1)), + ContinuousInput(key="x3", bounds=(0, 1)), + ContinuousInput(key="x4", bounds=(0, 1)), + ], + outputs=[ContinuousOutput(key="y", objective=MaximizeObjective(w=1.0))], # type: ignore + constraints=[ # type: ignore + NChooseKConstraint( + features=["x1", "x2", "x3", "x4"], + min_count=1, + max_count=2, + none_also_valid=False, + ), + LinearEqualityConstraint( + features=["x1", "x2", "x3", "x4"], + coefficients=[1.0, 1.0, 1.0, 1.0], + rhs=1.0, + ), + ], + ) + assert is_nchoosek_pruning_applicable(domain) is True + assert has_nchoosek_linear_overlap(domain) is True + + # Generate feasible training data (sum to 1, at most 2 active) + rng = np.random.default_rng(42) + experiments_list = [] + for _ in range(10): + x = np.zeros(4) + active = rng.choice(4, size=2, replace=False) + vals = rng.dirichlet(np.ones(2)) + x[active] = vals + experiments_list.append(x) + + experiments = pd.DataFrame(experiments_list, columns=["x1", "x2", "x3", "x4"]) + experiments["y"] = rng.standard_normal(10) + experiments["valid_y"] = 1 + + strategy_data = data_models.SoboStrategy(domain=domain) + strategy = SoboStrategy(data_model=strategy_data) + strategy.tell(experiments) + + candidates = strategy.ask(1, raise_validation_error=False) + + # Check NChooseK constraint is fulfilled + nchoosek = domain.constraints.get(NChooseKConstraint)[0] + assert isinstance(nchoosek, NChooseKConstraint) + assert nchoosek.is_fulfilled(candidates[domain.inputs.get_keys()]).all() + + # Check mixture constraint is approximately fulfilled (sum ≈ 1) + input_keys = domain.inputs.get_keys() + assert np.isclose(candidates[input_keys].sum(axis=1).values[0], 1.0, atol=1e-3) + + +def test_sobo_semicontinuous_through_pruning(): + """SOBO end-to-end on a domain with semi-continuous features + (allow_zero=True with positive lower bound). Every returned + coordinate must lie in `{0} ∪ [lb, ub]`. + """ + from bofire.data_models.features.api import ContinuousInput, ContinuousOutput + from bofire.data_models.objectives.api import MaximizeObjective + + lb = 0.2 + ub = 1.0 + domain = Domain.from_lists( + inputs=[ + ContinuousInput(key=f"x{i + 1}", bounds=(lb, ub), allow_zero=True) + for i in range(3) + ], + outputs=[ContinuousOutput(key="y", objective=MaximizeObjective(w=1.0))], + ) + + random_strategy = RandomStrategy( + data_model=RandomStrategyDataModel(domain=domain), + ) + experiments = random_strategy.ask(8) + experiments["y"] = np.random.default_rng(0).standard_normal(8) + experiments["valid_y"] = 1 + + strategy_data = data_models.SoboStrategy(domain=domain) + strategy = SoboStrategy(data_model=strategy_data) + strategy.tell(experiments) + candidates = strategy.ask(1, raise_validation_error=False) + + for key in domain.inputs.get_keys(): + v = float(candidates[key].iloc[0]) + assert v == pytest.approx(0.0, abs=1e-3) or ( + lb - 1e-3 <= v <= ub + 1e-3 + ), f"{key}={v} fell into the semi-continuity gap (0, {lb})"