MCTS based ACQF Optimization for Handling of NChooseKs.

bofire/benchmarks/benchmark.py

@@ -243,22 +243,49 @@ def get_optima(self) -> pd.DataFrame:
 
 
 class SpuriousFeaturesWrapper(Benchmark):
-    """Wrapper that adds spurious features to a benchmark, that are ignored on evaluation."""
+    """Wrapper that adds spurious features to a benchmark, that are ignored on evaluation.
 
-    def __init__(self, benchmark: Benchmark, n_spurious_features: int = 1, **kwargs):
+    Args:
+        benchmark: The benchmark to wrap.
+        n_spurious_features: Number of spurious features to add.
+        max_count: If provided, adds an NChooseKConstraint on all input features
+            (original + spurious) limiting the number of non-zero features.
+    """
+
+    def __init__(
+        self,
+        benchmark: Benchmark,
+        n_spurious_features: int = 1,
+        max_count: Optional[int] = None,
+        **kwargs,
+    ):
         super().__init__(**kwargs)
         assert n_spurious_features >= 1, "n_spurious_features must be >= 1."
         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 = list(self._benchmark.domain.constraints.constraints)
+        if max_count is not None:
+            constraints.append(
+                NChooseKConstraint(
+                    features=inputs.get_keys(),
+                    max_count=max_count,
+                    min_count=0,
+                    none_also_valid=True,
+                )
+            )
+
         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(constraints=constraints),
         )
 
     def _f(self, candidates: pd.DataFrame, **kwargs) -> pd.DataFrame:

bofire/data_models/strategies/random.py

@@ -25,6 +25,7 @@ class RandomStrategy(Strategy):
     n_thinning: Annotated[int, Field(ge=1)] = 32
     num_base_samples: Optional[Annotated[int, Field(gt=0)]] = None
     max_iters: Annotated[int, Field(gt=0)] = 1000
+    max_combinations: Annotated[int, Field(gt=0)] = 64
     sampler_kwargs: Optional[dict] = None
 
     def is_constraint_implemented(self, my_type: Type[Constraint]) -> bool:

bofire/strategies/predictives/acqf_optimization.py

@@ -47,6 +47,7 @@
 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.optimize_mcts import optimize_acqf_mcts
 from bofire.strategies.random import RandomStrategy
 from bofire.strategies.shortest_path import ShortestPathStrategy
 from bofire.utils.torch_tools import (
@@ -63,6 +64,7 @@ class OptimizerEnum(str, Enum):
     OPTIMIZE_ACQF = "OPTIMIZE_ACQF"
     OPTIMIZE_ACQF_MIXED = "OPTIMIZE_ACQF_MIXED"
     OPTIMIZE_ACQF_MIXED_ALTERNATING = "OPTIMIZE_ACQF_MIXED_ALTERNATING"
+    OPTIMIZE_ACQF_MCTS = "OPTIMIZE_ACQF_MCTS"
 
 
 # Threshold for switching between optimizers optimize_acqf_mixed
@@ -357,6 +359,36 @@ class _OptimizeAcqfMixedAlternatingInput(_OptimizeAcqfInputBase):
     equality_constraints: list[tuple[Tensor, Tensor, float]] | None
 
 
+class _OptimizeAcqfMctsInput(_OptimizeAcqfInputBase):
+    acq_function: Callable
+    bounds: Tensor
+    nchooseks: list[tuple[list[int], int, int]] | None
+    cat_dims: Dict[int, List[float]] | None
+    nig_alpha0: float
+    ts_prior_var: float
+    adaptive_prior_var: bool
+    cache_hit_mode: str
+    variance_decay: float
+    rollout_mode: str
+    adaptive_n0: bool
+    p_stop_rollout: float
+    num_iterations: int
+    pw_k0: float
+    pw_alpha: float
+    max_rollout_retries: int
+    use_cache: bool
+    n_sobol_samples: int
+    top_k_refine: int
+    screening_num_iterations: int | None
+    q: int
+    raw_samples: int
+    num_restarts: int
+    fixed_features: dict[int, float] | None
+    inequality_constraints: list[tuple[Tensor, Tensor, float]] | None
+    equality_constraints: list[tuple[Tensor, Tensor, float]] | None
+    seed: int | None
+
+
 class BotorchOptimizer(AcquisitionOptimizer):
     def __init__(self, data_model: BotorchOptimizerDataModel):
         self.n_restarts = data_model.n_restarts
@@ -454,6 +486,7 @@ def _optimize_acqf_continuous(
             OptimizerEnum.OPTIMIZE_ACQF: optimize_acqf,
             OptimizerEnum.OPTIMIZE_ACQF_MIXED: optimize_acqf_mixed,
             OptimizerEnum.OPTIMIZE_ACQF_MIXED_ALTERNATING: optimize_acqf_mixed_alternating,
+            OptimizerEnum.OPTIMIZE_ACQF_MCTS: optimize_acqf_mcts,
         }
         candidates, acqf_vals = optimizer_mapping[optimizer](
             **optimizer_input.model_dump()
@@ -482,6 +515,12 @@ def _get_optimizer_options(self, domain: Domain) -> Dict[str, int]:
         }
 
     def _determine_optimizer(self, domain: Domain, n_acqfs) -> OptimizerEnum:
+        # Check if we have NChooseK constraints - if so, use MCTS optimizer
+        if len(domain.constraints.get([NChooseKConstraint])) > 0 or any(
+            isinstance(feat, ContinuousInput) and feat.allow_zero
+            for feat in domain.inputs.get(ContinuousInput)
+        ):
+            return OptimizerEnum.OPTIMIZE_ACQF_MCTS
         if n_acqfs > 1:
             return OptimizerEnum.OPTIMIZE_ACQF_LIST
         n_categorical_combinations = (
@@ -508,6 +547,7 @@ def _get_arguments_for_optimizer(
         | _OptimizeAcqfMixedInput
         | _OptimizeAcqfListInput
         | _OptimizeAcqfMixedAlternatingInput
+        | _OptimizeAcqfMctsInput
     ):
         input_preprocessing_specs = self._input_preprocessing_specs(domain)
         features2idx = self._features2idx(domain)
@@ -616,6 +656,68 @@ def _get_arguments_for_optimizer(
                     if feat.key not in fixed_keys
                 },
             )
+        elif optimizer == OptimizerEnum.OPTIMIZE_ACQF_MCTS:
+            # Convert NChooseKConstraint to tuples for MCTS
+            nchoosek_constraints = domain.constraints.get([NChooseKConstraint])
+            nchooseks_list = []
+            nchoosek_feature_keys: set[str] = set()
+            for constraint in nchoosek_constraints:
+                # Get feature indices for the constraint
+                feature_indices = [
+                    features2idx[feat_key][0] for feat_key in constraint.features
+                ]
+                # Create tuple (features, min_count, max_count)
+                nchooseks_list.append(
+                    (feature_indices, constraint.min_count, constraint.max_count)
+                )
+                nchoosek_feature_keys.update(constraint.features)
+            # Continuous features with allow_zero=True are treated as NChooseK
+            # constraints where min_count=0 and max_count=1, but only if they
+            # are not already part of an explicit NChooseK constraint.
+            for feat in domain.inputs.get(ContinuousInput):
+                assert isinstance(feat, ContinuousInput)
+                if feat.allow_zero and feat.key not in nchoosek_feature_keys:
+                    feature_index = features2idx[feat.key][0]
+                    nchooseks_list.append(([feature_index], 0, 1))
+
+            # Get categorical dimensions (same as mixed_alternating)
+            fixed_keys = domain.inputs.get_fixed().get_keys()
+
+            return _OptimizeAcqfMctsInput(
+                acq_function=acqfs[0],
+                bounds=bounds,
+                nchooseks=nchooseks_list if nchooseks_list else None,
+                cat_dims={
+                    features2idx[feat.key][0]: feat.to_ordinal_encoding(  # type: ignore
+                        pd.Series(feat.get_allowed_categories())  # type: ignore
+                    ).tolist()
+                    for feat in domain.inputs.get(CategoricalInput)
+                    if feat.key not in fixed_keys
+                },
+                nig_alpha0=1.0,
+                ts_prior_var=1.0,
+                adaptive_prior_var=True,
+                cache_hit_mode="variance_inflation",
+                variance_decay=0.95,
+                rollout_mode="ts_group_action",
+                adaptive_n0=False,
+                p_stop_rollout=0.35,
+                num_iterations=300,
+                pw_k0=2.0,
+                pw_alpha=0.6,
+                max_rollout_retries=3,
+                use_cache=True,
+                n_sobol_samples=64,
+                top_k_refine=8,
+                screening_num_iterations=None,
+                q=candidate_count,
+                raw_samples=self.n_raw_samples,
+                num_restarts=self.n_restarts,
+                fixed_features=self.get_fixed_features(domain=domain),
+                inequality_constraints=inequality_constraints,
+                equality_constraints=equality_constraints,
+                seed=None,
+            )
         else:
             raise ValueError(f"Unknown optimizer: {optimizer}")