Universal constraint sampler (#328)

Author: Osburg

bofire/data_models/strategies/api.py

@@ -21,6 +21,9 @@
 from bofire.data_models.strategies.samplers.polytope import PolytopeSampler
 from bofire.data_models.strategies.samplers.rejection import RejectionSampler
 from bofire.data_models.strategies.samplers.sampler import SamplerStrategy
+from bofire.data_models.strategies.samplers.universal_constraint import (
+    UniversalConstraintSampler,
+)
 from bofire.data_models.strategies.stepwise.conditions import (  # noqa: F401
     AlwaysTrueCondition,
     CombiCondition,
@@ -50,6 +53,7 @@
     QparegoStrategy,
     PolytopeSampler,
     RejectionSampler,
+    UniversalConstraintSampler,
     RandomStrategy,
     DoEStrategy,
     StepwiseStrategy,
@@ -68,7 +72,7 @@
     MoboStrategy,
 ]
 
-AnySampler = Union[PolytopeSampler, RejectionSampler]
+AnySampler = Union[PolytopeSampler, RejectionSampler, UniversalConstraintSampler]
 
 
 AnyCondition = Union[NumberOfExperimentsCondition, CombiCondition, AlwaysTrueCondition]

bofire/data_models/strategies/samplers/universal_constraint.py

@@ -0,0 +1,49 @@
+from typing import Annotated, Literal, Type
+
+from pydantic import Field
+
+from bofire.data_models.constraints.api import (
+    LinearEqualityConstraint,
+    LinearInequalityConstraint,
+    NChooseKConstraint,
+    NonlinearEqualityConstraint,
+    NonlinearInequalityConstraint,
+)
+from bofire.data_models.features.api import (
+    ContinuousInput,
+    ContinuousOutput,
+    Feature,
+)
+from bofire.data_models.strategies.strategy import Strategy
+
+
+class UniversalConstraintSampler(Strategy):
+    """Sampler that generates samples by optimization in IPOPT.
+
+    Attributes:
+        domain (Domain): Domain defining the constrained input space
+        sampling_fraction (float, optional): Fraction of sampled points to total points generated in
+            the sampling process. Defaults to 0.3.
+        ipopt_options (dict, optional): Dictionary containing options for the IPOPT solver. Defaults to {"maxiter":200, "disp"=0}.
+    """
+
+    type: Literal["UniversalConstraintSampler"] = "UniversalConstraintSampler"
+    sampling_fraction: Annotated[float, Field(gt=0, lt=1)] = 0.3
+    ipopt_options: dict = {"maxiter": 200, "disp": 0}
+
+    @classmethod
+    def is_constraint_implemented(cls, my_type: Type[Feature]) -> bool:
+        return my_type in [
+            LinearEqualityConstraint,
+            LinearInequalityConstraint,
+            NonlinearInequalityConstraint,
+            NonlinearEqualityConstraint,
+            NChooseKConstraint,
+        ]
+
+    @classmethod
+    def is_feature_implemented(cls, my_type: Type[Feature]) -> bool:
+        return my_type in [
+            ContinuousInput,
+            ContinuousOutput,
+        ]

bofire/strategies/api.py

@@ -15,5 +15,8 @@
 from bofire.strategies.samplers.polytope import PolytopeSampler  # noqa: F401
 from bofire.strategies.samplers.rejection import RejectionSampler  # noqa: F401
 from bofire.strategies.samplers.sampler import SamplerStrategy  # noqa: F401
+from bofire.strategies.samplers.universal_constraint import (  # noqa: F401
+    UniversalConstraintSampler,
+)
 from bofire.strategies.stepwise.stepwise import StepwiseStrategy  # noqa: F401
 from bofire.strategies.strategy import Strategy  # noqa: F401

bofire/strategies/mapper.py

@@ -19,6 +19,9 @@
 from bofire.strategies.samplers.polytope import PolytopeSampler  # noqa: F401
 from bofire.strategies.samplers.rejection import RejectionSampler  # noqa: F401
 from bofire.strategies.samplers.sampler import SamplerStrategy  # noqa: F401
+from bofire.strategies.samplers.universal_constraint import (  # noqa: F401
+    UniversalConstraintSampler,
+)
 from bofire.strategies.stepwise.stepwise import StepwiseStrategy
 from bofire.strategies.strategy import Strategy  # noqa: F401
 
@@ -33,6 +36,7 @@
     data_models.QparegoStrategy: QparegoStrategy,
     data_models.PolytopeSampler: PolytopeSampler,
     data_models.RejectionSampler: RejectionSampler,
+    data_models.UniversalConstraintSampler: UniversalConstraintSampler,
     data_models.DoEStrategy: DoEStrategy,
     data_models.StepwiseStrategy: StepwiseStrategy,
     data_models.FactorialStrategy: FactorialStrategy,

bofire/strategies/samplers/universal_constraint.py

@@ -0,0 +1,55 @@
+import pandas as pd
+
+from bofire.data_models.strategies.api import UniversalConstraintSampler as DataModel
+from bofire.strategies.doe.design import find_local_max_ipopt
+from bofire.strategies.enum import OptimalityCriterionEnum
+from bofire.strategies.strategy import Strategy
+
+
+class UniversalConstraintSampler(Strategy):
+    """Sampler that generates samples by optimization in IPOPT.
+
+    Attributes:
+        domain (Domain): Domain defining the constrained input space
+        sampling_fraction (float, optional): Fraction of sampled points to total points generated in
+            the sampling process. Defaults to 0.3.
+        ipopt_options (dict, optional): Dictionary containing options for the IPOPT solver. Defaults to {"maxiter":200, "disp"=0}.
+    """
+
+    def __init__(
+        self,
+        data_model: DataModel,
+        **kwargs,
+    ):
+        super().__init__(data_model=data_model, **kwargs)
+        assert data_model.sampling_fraction > 0 and data_model.sampling_fraction <= 1
+        self.sampling_fraction = data_model.sampling_fraction
+        self.ipopt_options = data_model.ipopt_options
+
+    def _ask(self, candidate_count: int) -> pd.DataFrame:
+        samples = find_local_max_ipopt(
+            domain=self.domain,
+            model_type="linear",  # dummy model
+            n_experiments=self.num_candidates
+            + int(candidate_count / self.sampling_fraction),
+            ipopt_options=self.ipopt_options,
+            objective=OptimalityCriterionEnum.SPACE_FILLING,
+            fixed_experiments=self.candidates,
+        )
+
+        samples = samples.iloc[
+            self.num_candidates :,
+        ]
+        samples = samples.sample(
+            n=candidate_count,
+            replace=False,
+            ignore_index=True,
+            random_state=self._get_seed(),
+        )
+
+        self.domain.validate_experiments(samples)
+
+        return samples
+
+    def has_sufficient_experiments(self) -> bool:
+        return True

tests/bofire/data_models/specs/strategies.py

@@ -149,6 +149,15 @@
         "seed": 42,
     },
 )
+specs.add_valid(
+    strategies.UniversalConstraintSampler,
+    lambda: {
+        "domain": domain.valid().obj().dict(),
+        "sampling_fraction": 0.3,
+        "ipopt_options": {"maxiter": 200, "disp": 0},
+        "seed": 42,
+    },
+)
 
 tempdomain = domain.valid().obj().dict()
 

tests/bofire/strategies/test_samplers.py

@@ -1,4 +1,5 @@
 import pytest
+from pandas import concat
 
 import bofire.data_models.strategies.api as data_models
 import bofire.strategies.api as strategies
@@ -7,10 +8,62 @@
     LinearEqualityConstraint,
     LinearInequalityConstraint,
     NChooseKConstraint,
+    NonlinearEqualityConstraint,
+    NonlinearInequalityConstraint,
 )
 from bofire.data_models.domain.api import Constraints, Domain, Inputs
 from bofire.data_models.features.api import CategoricalInput, ContinuousInput
 
+inputs = [ContinuousInput(key=f"if{i}", bounds=(0, 1)) for i in range(1, 4)]
+c1 = LinearInequalityConstraint(
+    features=["if1", "if2", "if3"], coefficients=[1, 1, 1], rhs=1
+)
+c2 = LinearEqualityConstraint(
+    features=["if1", "if2", "if3"], coefficients=[1, 1, 1], rhs=1
+)
+c3 = NonlinearEqualityConstraint(
+    expression="if1**2 + if2**2 - if3", features=["if1", "if2", "if3"]
+)
+c4 = NonlinearInequalityConstraint(
+    expression="if1**2 + if2**2 - if3", features=["if1", "if2", "if3"]
+)
+c5 = NChooseKConstraint(
+    features=["if1", "if2", "if3"], min_count=0, max_count=1, none_also_valid=True
+)
+
+
+domains = [
+    Domain.from_lists(inputs=inputs, constraints=[c1]),
+    Domain.from_lists(inputs=inputs, constraints=[c2]),
+    Domain.from_lists(inputs=inputs, constraints=[c3]),
+    Domain.from_lists(inputs=inputs, constraints=[c4]),
+    Domain.from_lists(inputs=inputs, constraints=[c5]),
+]
+
+
+@pytest.mark.parametrize(
+    "domain, num_samples",
+    [(domain, candidate_count) for domain in domains for candidate_count in [1, 16]],
+)
+def test_UniversalConstraintSampler(domain, num_samples):
+    data_model = data_models.UniversalConstraintSampler(domain=domain)
+    sampler = strategies.UniversalConstraintSampler(data_model=data_model)
+    samples = sampler.ask(num_samples)
+    assert len(samples) == num_samples
+
+
+def test_UniversalConstraintSampler_pending_candidates():
+    data_model = data_models.UniversalConstraintSampler(domain=domains[0])
+    sampler = strategies.UniversalConstraintSampler(data_model=data_model)
+    pending_candidates = sampler.ask(2, add_pending=True)
+    samples = sampler.ask(1)
+    assert len(samples) == 1
+    all_samples = concat(
+        [samples, pending_candidates], axis=0, ignore_index=True
+    ).drop_duplicates()
+    assert len(all_samples) == 3
+
+
 inputs = Inputs(
     features=[ContinuousInput(key=f"if{i}", bounds=(0, 1)) for i in range(1, 4)]
 )