Refactoring of DoEStrategy data model class

bofire/data_models/strategies/api.py

@@ -1,7 +1,17 @@
 from typing import Union
 
 from bofire.data_models.strategies.actual_strategy_type import ActualStrategy
-from bofire.data_models.strategies.doe import DoEStrategy
+from bofire.data_models.strategies.doe import (
+    AnyDoEOptimalityCriterion,
+    AnyOptimalityCriterion,
+    AOptimalityCriterion,
+    DoEStrategy,
+    DOptimalityCriterion,
+    EOptimalityCriterion,
+    GOptimalityCriterion,
+    KOptimalityCriterion,
+    SpaceFillingCriterion,
+)
 from bofire.data_models.strategies.factorial import FactorialStrategy
 from bofire.data_models.strategies.fractional_factorial import (
     FractionalFactorialStrategy,

bofire/data_models/strategies/doe.py

@@ -1,24 +1,98 @@
-from typing import Literal, Optional, Type, Union
+from typing import Dict, Literal, Optional, Type, Union
 
+from formulaic import Formula
+from formulaic.errors import FormulaSyntaxError
+from pydantic import Field, field_validator
+
+from bofire.data_models.base import BaseModel
 from bofire.data_models.constraints.api import Constraint
 from bofire.data_models.features.api import Feature, MolecularInput
 from bofire.data_models.objectives.api import Objective
 from bofire.data_models.strategies.strategy import Strategy
 from bofire.data_models.types import Bounds
-from bofire.strategies.enum import OptimalityCriterionEnum
 
 
-class DoEStrategy(Strategy):
-    type: Literal["DoEStrategy"] = "DoEStrategy"
+PREDEFINED_MODEL_TYPES = Literal[
+    "linear",
+    "linear-and-quadratic",
+    "linear-and-interactions",
+    "fully-quadratic",
+]
+
+
+class OptimalityCriterion(BaseModel):
+    type: str
+    delta: float = 1e-6
+    transform_range: Optional[Bounds] = None
+
+
+class SpaceFillingCriterion(OptimalityCriterion):
+    type: Literal["SpaceFillingCriterion"] = "SpaceFillingCriterion"  # type: ignore
+
+
+class DoEOptimalityCriterion(OptimalityCriterion):
+    type: str
     formula: Union[
-        Literal[
-            "linear",
-            "linear-and-quadratic",
-            "linear-and-interactions",
-            "fully-quadratic",
-        ],
+        PREDEFINED_MODEL_TYPES,
         str,
     ]
+    """
+    model_type (str, Formula): keyword or formulaic Formula describing the model. Known keywords
+    are "linear", "linear-and-interactions", "linear-and-quadratic", "fully-quadratic".
+    """
+
+    @field_validator("formula")
+    @classmethod
+    def validate_formula(cls, formula: str) -> str:
+        if formula not in PREDEFINED_MODEL_TYPES.__args__:  # type: ignore
+            # check that it is a valid formula
+            try:
+                Formula(formula)
+            except FormulaSyntaxError:
+                raise ValueError(f"Invalid formula: {formula}")
+        return formula
+
+
+class DOptimalityCriterion(DoEOptimalityCriterion):
+    type: Literal["DOptimalityCriterion"] = "DOptimalityCriterion"  # type: ignore
+
+
+class EOptimalityCriterion(DoEOptimalityCriterion):
+    type: Literal["EOptimalityCriterion"] = "EOptimalityCriterion"  # type: ignore
+
+
+class AOptimalityCriterion(DoEOptimalityCriterion):
+    type: Literal["AOptimalityCriterion"] = "AOptimalityCriterion"  # type: ignore
+
+
+class GOptimalityCriterion(DoEOptimalityCriterion):
+    type: Literal["GOptimalityCriterion"] = "GOptimalityCriterion"  # type: ignore
+
+
+class KOptimalityCriterion(DoEOptimalityCriterion):
+    type: Literal["KOptimalityCriterion"] = "KOptimalityCriterion"  # type: ignore
+
+
+AnyDoEOptimalityCriterion = Union[
+    KOptimalityCriterion,
+    GOptimalityCriterion,
+    AOptimalityCriterion,
+    EOptimalityCriterion,
+    DOptimalityCriterion,
+]
+
+AnyOptimalityCriterion = Union[
+    AnyDoEOptimalityCriterion,
+    SpaceFillingCriterion,
+]
+
+
+class DoEStrategy(Strategy):
+    type: Literal["DoEStrategy"] = "DoEStrategy"  # type: ignore
+
+    criterion: AnyOptimalityCriterion = Field(
+        default_factory=lambda: DOptimalityCriterion(formula="fully-quadratic")
+    )
     optimization_strategy: Literal[
         "default",
         "exhaustive",
@@ -28,11 +102,8 @@ class DoEStrategy(Strategy):
         "iterative",
     ] = "default"
 
-    verbose: bool = False
-
-    objective: OptimalityCriterionEnum = OptimalityCriterionEnum.D_OPTIMALITY
-
-    transform_range: Optional[Bounds] = None
+    verbose: bool = False  # get rid of this at a later stage
+    ipopt_options: Optional[Dict] = None
 
     @classmethod
     def is_constraint_implemented(cls, my_type: Type[Constraint]) -> bool:

bofire/data_models/strategies/space_filling.py

@@ -31,7 +31,7 @@ class SpaceFillingStrategy(Strategy):
 
     """
 
-    type: Literal["SpaceFillingStrategy"] = "SpaceFillingStrategy"
+    type: Literal["SpaceFillingStrategy"] = "SpaceFillingStrategy"  # type: ignore
     sampling_fraction: Annotated[float, Field(gt=0, lt=1)] = 0.3
     ipopt_options: dict = {"maxiter": 200, "disp": 0}