Scale RandomStrategy NChooseK sampling and add allow_zero support

bofire/data_models/constraints/nchoosek.py

@@ -41,6 +41,10 @@ def validate_inputs(self, inputs: Inputs):
             assert isinstance(
                 feature_, ContinuousInput
             ), f"Feature {f} is not a ContinuousInput."
+            if feature_.bounds[0] < 0:
+                raise ValueError(
+                    f"Feature {f} must have a lower bound of >=0, but has {feature_.bounds[0]}",
+                )
 
     @model_validator(mode="after")
     def validate_counts(self):

bofire/data_models/domain/domain.py

@@ -1,5 +1,7 @@
 import collections.abc
 import itertools
+import math
+import random
 import warnings
 from collections.abc import Sequence
 from typing import Any, Dict, Literal, Optional, Tuple, Union
@@ -236,6 +238,78 @@ def get_nchoosek_combinations(self, exhaustive: bool = False):
 
         return used_features_list_final, unused_features_list
 
+    def sample_valid_nchoosek_features(
+        self,
+        seed: Optional[int] = None,
+        n: int = 1,
+        max_iters: int = 1000,
+    ) -> list[Tuple[str, ...]]:
+        """Sample sets of active feature keys uniformly from all valid subsets.
+
+        Includes (a) one group per ``NChooseKConstraint`` (respecting
+        ``min_count``, ``max_count``, and ``none_also_valid``) and (b) one
+        singleton group per ``ContinuousInput`` with ``allow_zero=True`` that
+        is not already part of any ``NChooseKConstraint``.
+
+        Within each group the subset size ``k`` is drawn with probability
+        proportional to ``C(n, k)`` and ``k`` features are then chosen
+        uniformly, so the per-group distribution is uniform over all valid
+        subsets. When ``NChooseKConstraint``s share features, the per-group
+        union may violate one of the constraints; in that case rejection
+        sampling is used (up to ``max_iters`` attempts per drawn combination).
+
+        Args:
+            seed: Random seed used to initialise the internal sampler.
+                Defaults to ``None`` (non-deterministic).
+            n: Number of combinations to draw. Defaults to 1.
+            max_iters: Maximum number of rejection-sampling attempts per
+                drawn combination. Defaults to 1000.
+
+        Returns:
+            A list of ``n`` sorted tuples of active feature keys.
+
+        Raises:
+            ValueError: If a valid combination is not found within
+                ``max_iters`` attempts.
+        """
+        rng = random.Random(seed)
+        groups: list[Tuple[list[str], list[int], list[int]]] = []
+        nchoosek_keys: set[str] = set()
+        nchoosek_cons = list(self.constraints.get(NChooseKConstraint))
+        for con in nchoosek_cons:
+            assert isinstance(con, NChooseKConstraint)
+            ks = list(range(con.min_count, con.max_count + 1))
+            if con.none_also_valid and 0 not in ks:
+                ks.insert(0, 0)
+            weights = [math.comb(len(con.features), k) for k in ks]
+            groups.append((list(con.features), ks, weights))
+            nchoosek_keys.update(con.features)
+        for feat in self.inputs.get(ContinuousInput):
+            assert isinstance(feat, ContinuousInput)
+            if feat.allow_zero and feat.key not in nchoosek_keys:
+                groups.append(([feat.key], [0, 1], [1, 1]))
+
+        results: list[Tuple[str, ...]] = []
+        for _ in range(n):
+            for _ in range(max_iters):
+                active: set[str] = set()
+                for features, ks, weights in groups:
+                    k = rng.choices(ks, weights=weights, k=1)[0]
+                    active.update(rng.sample(features, k))
+                if all(
+                    (con.none_also_valid and len(active & set(con.features)) == 0)
+                    or con.min_count <= len(active & set(con.features)) <= con.max_count
+                    for con in nchoosek_cons
+                ):
+                    results.append(tuple(sorted(active)))
+                    break
+            else:
+                raise ValueError(
+                    f"Failed to sample a valid NChooseK combination after "
+                    f"{max_iters} attempts.",
+                )
+        return results
+
     def coerce_invalids(self, experiments: pd.DataFrame) -> pd.DataFrame:
         """Coerces all invalid output measurements to np.nan
 

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/random.py

@@ -59,6 +59,7 @@ def __init__(
         self.fallback_sampling_method = data_model.fallback_sampling_method
         self.n_burnin = data_model.n_burnin
         self.n_thinning = data_model.n_thinning
+        self.max_combinations = data_model.max_combinations
         self.sampler_kwargs = data_model.sampler_kwargs
 
     def has_sufficient_experiments(self) -> bool:
@@ -124,42 +125,55 @@ def _sample_with_nchooseks(
             pd.DataFrame: A DataFrame containing the sampled data.
 
         """
-        if len(self.domain.constraints.get(NChooseKConstraint)) > 0:
-            _, unused = self.domain.get_nchoosek_combinations()
-
-            if candidate_count <= len(unused):
-                sampled_combinations = [
-                    unused[i]
-                    for i in np.random.default_rng(self._get_seed()).choice(
-                        len(unused),
-                        size=candidate_count,
-                        replace=False,
-                    )
-                ]
-                num_samples_per_it = 1
-            else:
-                sampled_combinations = unused
-                num_samples_per_it = math.ceil(candidate_count / len(unused))
+        nchoosek_feature_keys: set[str] = set()
+        for constraint in self.domain.constraints.get(NChooseKConstraint):
+            assert isinstance(constraint, NChooseKConstraint)
+            nchoosek_feature_keys.update(constraint.features)
+        allow_zero_feature_keys = {
+            feat.key
+            for feat in self.domain.inputs.get(ContinuousInput)
+            if isinstance(feat, ContinuousInput) and feat.allow_zero
+        }
+        zeroable_keys = nchoosek_feature_keys | allow_zero_feature_keys
+
+        if zeroable_keys:
+            # Draw a uniform sample of valid active-feature subsets (one per
+            # NChooseK constraint plus one per allow_zero singleton). We draw
+            # at most `max_combinations` distinct subsets; their multiplicities
+            # in `drawn` determine how many polytope samples each subset gets.
+            n_combos = min(self.max_combinations, candidate_count)
+            drawn = self.domain.sample_valid_nchoosek_features(
+                seed=self._get_seed(),
+                n=n_combos,
+            )
+            combinations: Dict[tuple, int] = {}
+            for combo in drawn:
+                combinations[combo] = combinations.get(combo, 0) + 1
+
+            # Each sampled subset gets `count * sampling_multiplier` polytope
+            # samples, so the total before final resampling is at least
+            # `candidate_count`.
+            sampling_multiplier = math.ceil(candidate_count / n_combos)
 
             samples = []
-            for u in sampled_combinations:
-                # create new domain without the nchoosekconstraints
+            for combo, count in combinations.items():
+                # Clone the domain and turn the NChooseK problem into a plain
+                # polytope problem: drop the NChooseK constraints, then pin
+                # every zeroable feature that wasn't selected to bounds [0, 0].
                 domain = deepcopy(self.domain)
                 domain.constraints = domain.constraints.get(excludes=NChooseKConstraint)
-                # fix the unused features
-                for key in u:
+                for key in zeroable_keys - set(combo):
                     feat = domain.inputs.get_by_key(key=key)
                     assert isinstance(feat, ContinuousInput)
                     feat.bounds = [0.0, 0.0]
-                # setup then sampler for this situation
                 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=num_samples_per_it,
+                        n=count * sampling_multiplier,
                         sampler_kwargs=self.sampler_kwargs,
                     ),
                 )
@@ -277,7 +291,7 @@ def _sample_from_polytope(
             samples = pd.DataFrame(
                 data=np.nan,
                 index=range(n),
-                columns=domain.inputs.get_keys(),
+                columns=domain.inputs.get_keys(ContinuousInput),
             )
         else:
             bounds = torch.tensor([lower, upper]).to(**tkwargs)
@@ -334,10 +348,13 @@ def _sample_from_polytope(
             )
 
         # setup the categoricals and discrete ones as uniform sampled vals
+        # we have to make sure here that no fixed ones occur here
         samples = pd.concat(
             [
                 samples,
-                domain.inputs.get([CategoricalInput, DiscreteInput]).sample(
+                domain.inputs.get([CategoricalInput, DiscreteInput])
+                .get_free()
+                .sample(
                     n,
                     method=fallback_sampling_method,
                     seed=seed,
@@ -350,6 +367,9 @@ def _sample_from_polytope(
         # setup the fixed continuous ones
         for key, value in fixed_features.items():
             samples[key] = value
+        # setup the fixed discrete/categorical ones
+        for feat in domain.inputs.get([CategoricalInput, DiscreteInput]).get_fixed():
+            samples[feat.key] = feat.fixed_value()[0]
 
         return samples[domain.inputs.get_keys()]
 
@@ -362,6 +382,7 @@ def make(
         n_thinning: int | None = None,
         num_base_samples: int | None = None,
         max_iters: int | None = None,
+        max_combinations: int | None = None,
         seed: int | None = None,
         sampler_kwargs: Optional[dict] = None,
     ) -> Self:
@@ -373,6 +394,7 @@ def make(
             n_thinning: The thinning factor for the polytope sampler.
             num_base_samples: The number of base samples for rejection sampling.
             max_iters: The maximum number of iterations for rejection sampling.
+            max_combinations: The maximum number of distinct NChooseK feature combinations to draw per ask.
             seed: The seed value for random number generation.
             sampler_kwargs: Additional arguments for the sampler. Defaults to None.
         Returns:

tests/bofire/data_models/domain/test_domain.py

@@ -507,3 +507,75 @@ def test_is_fulfilled():
         domain.is_fulfilled(experiments),
         pd.Series([True, False, False], index=experiments.index),
     )
+
+
+def test_sample_valid_nchoosek_features_uniform_over_subsets():
+    """With one NChooseK on n=5 features and k in [1, 3], there are
+    C(5,1)+C(5,2)+C(5,3) = 25 valid subsets. With uniform sampling each
+    should appear with frequency ~1/25.
+    """
+    inputs = [ContinuousInput(key=f"x{i}", bounds=(0, 1)) for i in range(5)]
+    constraint = NChooseKConstraint(
+        features=[f"x{i}" for i in range(5)],
+        min_count=1,
+        max_count=3,
+        none_also_valid=False,
+    )
+    domain = Domain(
+        inputs=Inputs(features=inputs),
+        constraints=Constraints(constraints=[constraint]),
+    )
+    n_samples = 25_000
+    samples = domain.sample_valid_nchoosek_features(seed=0, n=n_samples)
+    counts: dict = {}
+    for s in samples:
+        counts[s] = counts.get(s, 0) + 1
+    assert len(counts) == 25, f"Expected 25 unique subsets, got {len(counts)}"
+    expected = n_samples / 25
+    for subset, count in counts.items():
+        rel = abs(count - expected) / expected
+        assert (
+            rel < 0.20
+        ), f"Subset {subset} count {count} too far from expected {expected:.0f}"
+
+
+def test_sample_valid_nchoosek_features_none_also_valid():
+    """When none_also_valid=True, the empty subset is in the support."""
+    inputs = [ContinuousInput(key=f"x{i}", bounds=(0, 1)) for i in range(3)]
+    constraint = NChooseKConstraint(
+        features=["x0", "x1", "x2"],
+        min_count=2,
+        max_count=3,
+        none_also_valid=True,
+    )
+    domain = Domain(
+        inputs=Inputs(features=inputs),
+        constraints=Constraints(constraints=[constraint]),
+    )
+    samples = domain.sample_valid_nchoosek_features(seed=1, n=2000)
+    unique = set(samples)
+    # Valid subsets: () + C(3,2) + C(3,3) = 1 + 3 + 1 = 5
+    assert len(unique) == 5
+    assert () in unique
+
+
+def test_sample_valid_nchoosek_features_allow_zero_singletons():
+    """Without any NChooseK, allow_zero=True features form singleton groups."""
+    inputs = [
+        ContinuousInput(key="a", bounds=(0.1, 1), allow_zero=True),
+        ContinuousInput(key="b", bounds=(0.1, 1), allow_zero=True),
+        ContinuousInput(key="c", bounds=(0.1, 1)),
+    ]
+    domain = Domain(inputs=Inputs(features=inputs))
+    samples = domain.sample_valid_nchoosek_features(seed=2, n=2000)
+    unique = set(samples)
+    # Each of {a, b} can be on or off independently -> 4 subsets
+    assert unique == {(), ("a",), ("b",), ("a", "b")}
+
+
+def test_sample_valid_nchoosek_features_empty_returns_empty_tuple():
+    """Domain without NChooseK and without allow_zero features yields ()."""
+    inputs = [ContinuousInput(key="x", bounds=(0, 1))]
+    domain = Domain(inputs=Inputs(features=inputs))
+    samples = domain.sample_valid_nchoosek_features(seed=3, n=4)
+    assert samples == [(), (), (), ()]

tests/bofire/strategies/test_random.py

@@ -213,8 +213,8 @@ def test_nchoosek():
     If7 = ContinuousInput(bounds=(1, 1), key="If7")
 
     c2 = LinearInequalityConstraint.from_greater_equal(
-        features=["if1", "if2"],
-        coefficients=[1.0, 1.0],
+        features=["if1", "if2", "if3"],
+        coefficients=[1.0, 1.0, 1.0],
         rhs=0.2,
     )
 
@@ -225,8 +225,8 @@ def test_nchoosek():
         none_also_valid=False,
     )
     c7 = LinearEqualityConstraint(
-        features=["if1", "if2"],
-        coefficients=[1.0, 1.0],
+        features=["if1", "if2", "if3"],
+        coefficients=[1.0, 1.0, 1.0],
         rhs=1.0,
     )
     domain = Domain.from_lists(
@@ -239,6 +239,47 @@ def test_nchoosek():
     assert len(samples) == 50
 
 
+def test_allow_zero_without_nchoosek():
+    """Test random sampling with allow_zero features but no NChooseK constraint."""
+    if1 = ContinuousInput(bounds=(0.1, 1), key="if1", allow_zero=True)
+    if2 = ContinuousInput(bounds=(0.1, 1), key="if2", allow_zero=True)
+    if3 = ContinuousInput(bounds=(0.1, 1), key="if3")
+    domain = Domain.from_lists(inputs=[if1, if2, if3])
+    data_model = data_models.RandomStrategy(domain=domain)
+    sampler = strategies.RandomStrategy(data_model=data_model)
+    samples = sampler.ask(50)
+    assert len(samples) == 50
+    # if3 should never be zero (not allow_zero)
+    assert (samples["if3"] != 0.0).all()
+    # if1 and if2 should have some zeros (allow_zero)
+    assert (samples["if1"] == 0.0).any() or (samples["if2"] == 0.0).any()
+
+
+def test_allow_zero_with_nchoosek():
+    """Test that allow_zero features already in NChooseK don't get duplicate groups."""
+    if1 = ContinuousInput(bounds=(0, 1), key="if1")
+    if2 = ContinuousInput(bounds=(0, 1), key="if2")
+    if3 = ContinuousInput(bounds=(0, 1), key="if3")
+    if4 = ContinuousInput(bounds=(0.1, 1), key="if4", allow_zero=True)
+    c = NChooseKConstraint(
+        features=["if1", "if2", "if3"],
+        min_count=1,
+        max_count=2,
+        none_also_valid=False,
+    )
+    domain = Domain.from_lists(inputs=[if1, if2, if3, if4], constraints=[c])
+    data_model = data_models.RandomStrategy(domain=domain)
+    sampler = strategies.RandomStrategy(data_model=data_model)
+    samples = sampler.ask(50)
+    assert len(samples) == 50
+    # At most 2 features should be non-zero per sample (from NChooseK)
+    nonzero_counts = (samples[["if1", "if2", "if3"]] != 0.0).sum(axis=1)
+    assert (nonzero_counts >= 1).all()
+    assert (nonzero_counts <= 2).all()
+    # if4 (allow_zero, not in NChooseK) should have some zeros
+    assert (samples["if4"] == 0.0).any()
+
+
 def test_sample_from_polytope():
     if1 = ContinuousInput(
         bounds=(0, 1),