Transforms (#355)

* less max_parallelism

* fixed test parametrization

* removed max_parallel-test

* added valid/invalid serializations to tests

* register invalid

* separate transform data model

* ignore vscode's devcontainer

* separate api file

* added init and TF constructor

* rename to `DropDataTransform`

* Apply suggestions from code review

Co-authored-by: Johannes P. Dürholt <[email protected]>

* Apply suggestions from code review

Co-authored-by: Johannes P. Dürholt <[email protected]>

* docstring, register valid specs, rename to untransform_candidates

* move test to separate file and extend

---------

Co-authored-by: Johannes P. Dürholt <[email protected]>

.gitignore

@@ -116,6 +116,7 @@ venv.bak/
 
 # vscode project settings
 .vscode/
+.devcontainer
 
 # Rope project settings
 .ropeproject
@@ -133,3 +134,4 @@ dmypy.json
 
 # generated version file
 bofire/version.py
+

bofire/data_models/strategies/api.py

@@ -22,6 +22,7 @@
 from bofire.data_models.strategies.space_filling import SpaceFillingStrategy
 from bofire.data_models.strategies.stepwise.conditions import (  # noqa: F401
     AlwaysTrueCondition,
+    AnyCondition,
     CombiCondition,
     NumberOfExperimentsCondition,
 )
@@ -30,6 +31,10 @@
     StepwiseStrategy,
 )
 from bofire.data_models.strategies.strategy import Strategy
+from bofire.data_models.transforms.api import (  # noqa: F401
+    AnyTransform,
+    DropDataTransform,
+)
 
 AbstractStrategy = Union[
     Strategy,
@@ -67,7 +72,4 @@
 ]
 
 
-AnyCondition = Union[NumberOfExperimentsCondition, CombiCondition, AlwaysTrueCondition]
-
-
 AnyLocalSearchConfig = LSRBO

bofire/data_models/strategies/stepwise/conditions.py

@@ -1,9 +1,18 @@
-from typing import List, Literal, Union
+from abc import abstractmethod
+from typing import List, Literal, Optional, Union
 
+import pandas as pd
 from pydantic import Field, field_validator
 from typing_extensions import Annotated
 
 from bofire.data_models.base import BaseModel
+from bofire.data_models.domain.api import Domain
+
+
+class EvaluateableCondition:
+    @abstractmethod
+    def evaluate(self, domain: Domain, experiments: Optional[pd.DataFrame]) -> bool:
+        pass
 
 
 class Condition(BaseModel):
@@ -14,16 +23,28 @@ class SingleCondition(BaseModel):
     type: str
 
 
-class NumberOfExperimentsCondition(SingleCondition):
+class NumberOfExperimentsCondition(SingleCondition, EvaluateableCondition):
     type: Literal["NumberOfExperimentsCondition"] = "NumberOfExperimentsCondition"
     n_experiments: Annotated[int, Field(ge=1)]
 
+    def evaluate(self, domain: Domain, experiments: Optional[pd.DataFrame]) -> bool:
+        if experiments is None:
+            n_experiments = 0
+        else:
+            n_experiments = len(
+                domain.outputs.preprocess_experiments_all_valid_outputs(experiments)
+            )
+        return n_experiments <= self.n_experiments
+
 
-class AlwaysTrueCondition(SingleCondition):
+class AlwaysTrueCondition(SingleCondition, EvaluateableCondition):
     type: Literal["AlwaysTrueCondition"] = "AlwaysTrueCondition"
 
+    def evaluate(self, domain: Domain, experiments: Optional[pd.DataFrame]) -> bool:
+        return True
 
-class CombiCondition(Condition):
+
+class CombiCondition(Condition, EvaluateableCondition):
     type: Literal["CombiCondition"] = "CombiCondition"
     conditions: Annotated[
         List[
@@ -41,3 +62,15 @@ def validate_n_required_conditions(cls, v, info):
                 "Number of required conditions larger than number of conditions."
             )
         return v
+
+    def evaluate(self, domain: Domain, experiments: Optional[pd.DataFrame]) -> bool:
+        n_matched_conditions = 0
+        for c in self.conditions:
+            if c.evaluate(domain, experiments):
+                n_matched_conditions += 1
+        if n_matched_conditions >= self.n_required_conditions:
+            return True
+        return False
+
+
+AnyCondition = Union[NumberOfExperimentsCondition, CombiCondition, AlwaysTrueCondition]

bofire/data_models/strategies/stepwise/stepwise.py

@@ -1,11 +1,12 @@
-from typing import List, Literal, Type, Union
+from typing import List, Literal, Optional, Type, Union
 
 from pydantic import Field, field_validator
 from typing_extensions import Annotated
 
 from bofire.data_models.base import BaseModel
 from bofire.data_models.constraints.api import Constraint
 from bofire.data_models.features.api import Feature
+from bofire.data_models.strategies.api import AnyCondition
 from bofire.data_models.strategies.doe import DoEStrategy
 from bofire.data_models.strategies.factorial import FactorialStrategy
 from bofire.data_models.strategies.predictives.mobo import MoboStrategy
@@ -21,12 +22,9 @@
 from bofire.data_models.strategies.random import RandomStrategy
 from bofire.data_models.strategies.shortest_path import ShortestPathStrategy
 from bofire.data_models.strategies.space_filling import SpaceFillingStrategy
-from bofire.data_models.strategies.stepwise.conditions import (
-    AlwaysTrueCondition,
-    CombiCondition,
-    NumberOfExperimentsCondition,
-)
+from bofire.data_models.strategies.stepwise.conditions import AlwaysTrueCondition
 from bofire.data_models.strategies.strategy import Strategy
+from bofire.data_models.transforms.api import AnyTransform
 
 AnyStrategy = Union[
     SoboStrategy,
@@ -44,14 +42,12 @@
     ShortestPathStrategy,
 ]
 
-AnyCondition = Union[NumberOfExperimentsCondition, CombiCondition, AlwaysTrueCondition]
-
 
 class Step(BaseModel):
     type: Literal["Step"] = "Step"
     strategy_data: AnyStrategy
     condition: AnyCondition
-    max_parallelism: Annotated[int, Field(ge=-1)]  # -1 means no restriction at all
+    transform: Optional[AnyTransform] = None
 
 
 class StepwiseStrategy(Strategy):

bofire/data_models/transforms/api.py

@@ -0,0 +1,3 @@
+from bofire.data_models.transforms.drop_data import DropDataTransform
+
+AnyTransform = DropDataTransform

bofire/data_models/transforms/drop_data.py

@@ -0,0 +1,9 @@
+from typing import List, Literal, Optional
+
+from pydantic import BaseModel
+
+
+class DropDataTransform(BaseModel):
+    type: Literal["DropDataTransform"] = "DropDataTransform"
+    to_be_removed_experiments: Optional[List[int]] = None
+    to_be_removed_candidates: Optional[List[int]] = None

bofire/strategies/stepwise/stepwise.py

@@ -1,10 +1,11 @@
-from typing import Dict, Tuple, Type
+from typing import Dict, Literal, Optional, Tuple, Type, TypeVar, Union
 
 import pandas as pd
+from pydantic import PositiveInt
 
 import bofire.data_models.strategies.api as data_models
-import bofire.strategies.stepwise.conditions as conditions
-from bofire.data_models.strategies.api import Step
+import bofire.transforms.api as transforms
+from bofire.data_models.domain.api import Domain
 from bofire.data_models.strategies.api import StepwiseStrategy as data_model
 from bofire.strategies.doe_strategy import DoEStrategy
 from bofire.strategies.factorial import FactorialStrategy
@@ -22,6 +23,7 @@
 from bofire.strategies.shortest_path import ShortestPathStrategy
 from bofire.strategies.space_filling import SpaceFillingStrategy
 from bofire.strategies.strategy import Strategy
+from bofire.transforms.transform import Transform
 
 # we have to duplicate the map functionality due to prevent circular imports
 STRATEGY_MAP: Dict[Type[data_models.Strategy], Type[Strategy]] = {
@@ -41,40 +43,67 @@
 }
 
 
-def map(data_model: data_models.Strategy) -> Strategy:
+def _map(data_model: data_models.Strategy) -> Strategy:
     cls = STRATEGY_MAP[data_model.__class__]
     return cls.from_spec(data_model=data_model)
 
 
+T = TypeVar("T", pd.DataFrame, Domain)
+
+
+TfData = Union[Literal["experiments"], Literal["candidates"], Literal["domain"]]
+
+
+def _apply_tf(
+    data: Optional[T],
+    transform: Optional[Transform],
+    tf: TfData,
+) -> Optional[T]:
+    if data is not None and transform is not None:
+        return getattr(transform, f"transform_{tf}")(data)
+
+
 class StepwiseStrategy(Strategy):
     def __init__(self, data_model: data_model, **kwargs):
         super().__init__(data_model, **kwargs)
-        self.steps = data_model.steps
+        self.strategies = [_map(s.strategy_data) for s in data_model.steps]
+        self.conditions = [s.condition for s in data_model.steps]
+        self.transforms = [
+            s.transform and transforms.map(s.transform) for s in data_model.steps
+        ]
 
     def has_sufficient_experiments(self) -> bool:
         return True
 
-    def _get_step(self) -> Tuple[int, Step]:  # type: ignore
-        for i, step in enumerate(self.steps):
-            condition = conditions.map(step.condition)
+    def _get_step(self) -> Tuple[Strategy, Optional[Transform]]:
+        """Returns the strategy at the current step and the corresponding transform if given."""
+        for i, condition in enumerate(self.conditions):
             if condition.evaluate(self.domain, experiments=self.experiments):
-                return i, step
+                return self.strategies[i], self.transforms[i]
         raise ValueError("No condition could be satisfied.")
 
-    def _ask(self, candidate_count: int) -> pd.DataFrame:
-        # we have to decide here w
-        istep, step = self._get_step()
-        if (step.max_parallelism > 0) and (candidate_count > step.max_parallelism):
-            raise ValueError(
-                f"Maximum number of candidates for step {istep} is {step.max_parallelism}."
-            )
-        # map it
-        strategy = map(step.strategy_data)
+    def _ask(self, candidate_count: Optional[PositiveInt]) -> pd.DataFrame:
+        strategy, transform = self._get_step()
+
+        candidate_count = candidate_count or 1
+
+        # handle a possible transform
+        tf_domain = _apply_tf(self.domain, transform, "domain")
+        transformed_domain = tf_domain or self.domain
+        strategy.domain = transformed_domain
+        tf_exp = _apply_tf(self.experiments, transform, "experiments")
+        transformed_experiments = self.experiments if tf_exp is None else tf_exp
+        tf_cand = _apply_tf(self.candidates, transform, "candidates")
+        transformed_candidates = self.candidates if tf_cand is None else tf_cand
         # tell the experiments
-        if self.num_experiments > 0:
-            strategy.tell(experiments=self.experiments)
+        if transformed_experiments is not None and len(transformed_experiments) > 0:
+            strategy.tell(experiments=transformed_experiments, replace=True)
         # tell pending
-        if self.num_candidates > 0:
-            strategy.set_candidates(self.candidates)
+        if transformed_candidates is not None and len(transformed_candidates) > 0:
+            strategy.set_candidates(transformed_candidates)
         # ask and return
-        return strategy.ask(candidate_count=candidate_count)
+        candidates = strategy.ask(candidate_count=candidate_count)
+        if transform is not None:
+            return transform.untransform_candidates(candidates)
+        else:
+            return candidates

bofire/transforms/api.py

@@ -0,0 +1,13 @@
+from typing import Dict, Type
+
+import bofire.data_models.transforms.api as data_models
+from bofire.transforms.drop_data import DropDataTransform
+from bofire.transforms.transform import Transform
+
+TRANSFORM_MAP: Dict[Type[data_models.AnyTransform], Type[Transform]] = {
+    data_models.DropDataTransform: DropDataTransform
+}
+
+
+def map(data_model: data_models.AnyTransform) -> Transform:
+    return TRANSFORM_MAP[data_model.__class__](data_model)

bofire/transforms/drop_data.py

@@ -0,0 +1,16 @@
+import pandas as pd
+
+from bofire.data_models.transforms.api import DropDataTransform as DataModel
+from bofire.transforms.transform import Transform
+
+
+class DropDataTransform(Transform):
+    def __init__(self, data_model: DataModel):
+        self.to_be_removed_experiments = data_model.to_be_removed_experiments or []
+        self.to_be_removed_candidates = data_model.to_be_removed_candidates or []
+
+    def transform_experiments(self, experiments: pd.DataFrame) -> pd.DataFrame:
+        return experiments.drop(self.to_be_removed_experiments)
+
+    def transform_candidates(self, candidates: pd.DataFrame) -> pd.DataFrame:
+        return candidates.drop(self.to_be_removed_candidates)

bofire/transforms/transform.py

@@ -0,0 +1,20 @@
+import pandas as pd
+
+from bofire.data_models.domain.api import Domain
+
+
+class Transform:
+    def __init__(self, *_args, **_kwargs) -> None:
+        pass
+
+    def transform_experiments(self, experiments: pd.DataFrame) -> pd.DataFrame:
+        return experiments
+
+    def transform_candidates(self, candidates: pd.DataFrame) -> pd.DataFrame:
+        return candidates
+
+    def transform_domain(self, domain: Domain) -> Domain:
+        return domain
+
+    def untransform_candidates(self, experiments: pd.DataFrame) -> pd.DataFrame:
+        return experiments