Implement ENTMOOT in Bofire

bofire/data_models/strategies/api.py

@@ -1,6 +1,7 @@
 from typing import Union
 
 from bofire.data_models.strategies.doe import DoEStrategy
+from bofire.data_models.strategies.enting import EntingStrategy
 from bofire.data_models.strategies.factorial import FactorialStrategy
 from bofire.data_models.strategies.predictives.botorch import BotorchStrategy
 from bofire.data_models.strategies.predictives.multiobjective import (

bofire/data_models/strategies/enting.py

@@ -0,0 +1,41 @@
+from typing import Any, Dict, Literal, Type
+
+from bofire.data_models.constraints.api import (
+    Constraint,
+    LinearEqualityConstraint,
+    LinearInequalityConstraint,
+)
+from bofire.data_models.features.api import (
+    CategoricalDescriptorInput,
+    CategoricalInput,
+    ContinuousInput,
+    ContinuousOutput,
+    DiscreteInput,
+    Feature,
+)
+from bofire.data_models.objectives.api import MinimizeObjective, Objective
+from bofire.data_models.strategies.strategy import Strategy
+
+
+class EntingStrategy(Strategy):
+    type: Literal["EntingStrategy"] = "EntingStrategy"
+    enting_params: Dict[str, Any] = {}
+    solver_params: Dict[str, Any] = {}
+
+    @classmethod
+    def is_constraint_implemented(cls, my_type: Type[Constraint]) -> bool:
+        return my_type in [LinearEqualityConstraint, LinearInequalityConstraint]
+
+    @classmethod
+    def is_feature_implemented(cls, my_type: Type[Feature]) -> bool:
+        return my_type in [
+            CategoricalInput,
+            DiscreteInput,
+            CategoricalDescriptorInput,
+            ContinuousInput,
+            ContinuousOutput,
+        ]
+
+    @classmethod
+    def is_objective_implemented(cls, my_type: Type[Objective]) -> bool:
+        return my_type in [MinimizeObjective]

bofire/strategies/entmoot/enting.py

@@ -0,0 +1,67 @@
+import pandas as pd
+import pyomo.environ as pyo
+from entmoot.models.enting import Enting
+from entmoot.optimizers.pyomo_opt import PyomoOptimizer
+from entmoot.problem_config import ProblemConfig
+from pydantic import PositiveInt
+
+import bofire.data_models.strategies.api as data_models
+from bofire.strategies.entmoot.problem_config import domain_to_problem_config
+from bofire.strategies.strategy import Strategy
+
+
+class EntingStrategy(Strategy):
+    """Strategy for selecting new candidates using ENTMOOT"""
+
+    def __init__(
+        self,
+        data_model: data_models.EntingStrategy,
+        **kwargs,
+    ):
+        super().__init__(data_model=data_model, **kwargs)
+        self._init_problem_config()
+        self._enting = Enting(self._problem_config, data_model.enting_params)
+        self._solver_params = data_model.solver_params
+
+    def _init_problem_config(self) -> None:
+        cfg = domain_to_problem_config(self.domain)
+        self._problem_config: ProblemConfig = cfg[0]
+        self._model_pyo: pyo.ConcreteModel = cfg[1]
+
+    def _ask(self, candidate_count: PositiveInt) -> pd.DataFrame:
+        if candidate_count > 1:
+            raise NotImplementedError("Can currently only handle one at a time")
+        opt_pyo = PyomoOptimizer(self._problem_config, params=self._solver_params)
+        res = opt_pyo.solve(tree_model=self._enting, model_core=self._model_pyo)
+        candidate = res.opt_point
+        objective_value = res.opt_val
+        unc_unscaled = res.unc_unscaled
+
+        keys = [feat.name for feat in self._problem_config.feat_list]
+        candidates = pd.DataFrame(
+            data=[candidate + [objective_value, unc_unscaled]],
+            index=range(candidate_count),
+            columns=keys + ["y_pred", "y_sd"],
+        )
+
+        return candidates
+
+    def _tell(self):
+        input_keys = self.domain.inputs.get_keys()
+        output_keys = self.domain.outputs.get_keys()
+
+        X = self._experiments[input_keys].values
+        y = self._experiments[output_keys].values
+        self._enting.fit(X, y)
+
+    def has_sufficient_experiments(self) -> bool:
+        if self.experiments is None:
+            return False
+        return (
+            len(
+                self.domain.outputs.preprocess_experiments_all_valid_outputs(
+                    experiments=self.experiments
+                )
+            )
+            > 1
+        )

bofire/strategies/entmoot/entmoot_benchmarks.py

@@ -0,0 +1,42 @@
+from bofire.data_models.domain.api import Domain, Inputs, Outputs
+from bofire.data_models.features.api import (
+    CategoricalInput,
+    ContinuousOutput,
+    DiscreteInput,
+)
+from bofire.data_models.objectives.api import MaximizeObjective
+
+
+def build_multi_obj_categorical_problem(n_obj: int = 2, no_cat=False) -> Domain:
+    """
+    Builds a small test example which is used in Entmoot tests.
+    """
+
+    cat_feat = (
+        []
+        if no_cat
+        else [CategoricalInput(key="x0", categories=("blue", "orange", "gray"))]
+    )
+    input_features = Inputs(
+        features=cat_feat
+        + [
+            DiscreteInput(key="x1", values=[5, 6]),
+            DiscreteInput(key="x2", values=[0, 1]),  # binary
+            ContinuousOutput(key="x3", values=[5.0, 6.0]),
+            ContinuousOutput(key="x4", values=[4.6, 6.0]),
+            ContinuousOutput(key="x5", values=[5.0, 6.0]),
+        ]
+    )
+
+    output_features = Outputs(
+        features=[
+            ContinuousOutput(
+                key=f"y{i}", objective=MaximizeObjective(w=1.0, bounds=[0.0, 1.0])
+            )
+            for i in range(n_obj)
+        ]
+    )
+
+    domain = Domain(inputs=input_features, outputs=output_features)
+
+    return domain

bofire/strategies/entmoot/entmoot_example.ipynb

@@ -0,0 +1,217 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Demonstration of the Entmoot API\n",
+    "\n",
+    "This notebook includes comparisons of the Entmoot strategy to other Bofire \n",
+    "strategies."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Single objective Bayesian Optimization strategy\n",
+    "\n",
+    "This section includes a comparison to the Sobo strategy given in the \"Getting \n",
+    "Started\" docs. The API is identical, with additional options given to the \n",
+    "data model. Note that the EntingStrategy only supports one candidate, as \n",
+    "each generated candidate is optimal (so generating multiple would generate \n",
+    "duplicates)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "c:\\users\\tobyb\\phd\\bofire\\bofire\\utils\\cheminformatics.py:19: UserWarning: mordred not installed. Mordred molecular descriptors cannot be used.\n",
+      "  warnings.warn(\n",
+      "c:\\users\\tobyb\\phd\\bofire\\bofire\\surrogates\\xgb.py:12: UserWarning: xgboost not installed, BoFire's `XGBoostSurrogate` cannot be used.\n",
+      "  warnings.warn(\"xgboost not installed, BoFire's `XGBoostSurrogate` cannot be used.\")\n"
+     ]
+    }
+   ],
+   "source": [
+    "from bofire.benchmarks.single import Himmelblau\n",
+    "import bofire.strategies.mapper as strategy_mapper\n",
+    "\n",
+    "benchmark = Himmelblau()\n",
+    "\n",
+    "samples = benchmark.domain.inputs.sample(10)\n",
+    "experiments = benchmark.f(samples, return_complete=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>x_1</th>\n",
+       "      <th>x_2</th>\n",
+       "      <th>y_pred</th>\n",
+       "      <th>y_sd</th>\n",
+       "      <th>y_des</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>-0.178689</td>\n",
+       "      <td>2.418258</td>\n",
+       "      <td>39.226887</td>\n",
+       "      <td>237.849118</td>\n",
+       "      <td>-39.226887</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "        x_1       x_2     y_pred        y_sd      y_des\n",
+       "0 -0.178689  2.418258  39.226887  237.849118 -39.226887"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from bofire.data_models.strategies.api import SoboStrategy\n",
+    "from bofire.data_models.acquisition_functions.api import qNEI\n",
+    "\n",
+    "sobo_strategy_data_model = SoboStrategy(domain=benchmark.domain, acquisition_function=qNEI())\n",
+    "sobo_strategy = strategy_mapper.map(sobo_strategy_data_model)\n",
+    "\n",
+    "sobo_strategy.tell(experiments=experiments)\n",
+    "sobo_strategy.ask(candidate_count=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "C:\\Users\\tobyb\\phd\\entmoot\\entmoot\\models\\mean_models\\tree_ensemble.py:23: UserWarning: No 'train_params' for tree ensemble training specified. Switch training to default params!\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>x_1</th>\n",
+       "      <th>x_2</th>\n",
+       "      <th>y_pred</th>\n",
+       "      <th>y_sd</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>0.101473</td>\n",
+       "      <td>-6.0</td>\n",
+       "      <td>26.755872</td>\n",
+       "      <td>0.232471</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "        x_1  x_2     y_pred      y_sd\n",
+       "0  0.101473 -6.0  26.755872  0.232471"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from bofire.data_models.strategies.api import EntingStrategy\n",
+    "\n",
+    "enting_params = {\"unc_params\": {\"dist_metric\": \"l1\", \"acq_sense\": \"exploration\"}}\n",
+    "solver_params = {\"solver_name\": \"gurobi\"}\n",
+    "\n",
+    "enting_strategy_data_model = EntingStrategy(domain=benchmark.domain, enting_params=enting_params, solver_params=solver_params)\n",
+    "enting_strategy = strategy_mapper.map(enting_strategy_data_model)\n",
+    "\n",
+    "enting_strategy.tell(experiments=experiments)\n",
+    "enting_strategy.ask(candidate_count=1)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.9"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}