Add BONSAI-style greedy pruning for NChooseK constraints

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:

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

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]

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]
 

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

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())

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.

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(

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):