Merge pull request #24 from TobyBoyne/feature/constraintclass

Constraint Class

docs/notebooks/constraint_classes.ipynb

@@ -0,0 +1,220 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Constraint Classes\n",
+    "\n",
+    "To make applying constraints to your model easier, some constraints have been \n",
+    "provided as a part of ENTMOOT."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from entmoot.problem_config import ProblemConfig\n",
+    "from entmoot.models.enting import Enting\n",
+    "from entmoot.optimizers.pyomo_opt import PyomoOptimizer"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### NChooseKConstraint\n",
+    "\n",
+    "This constraint is often used in the design of experiments. This applies a bound on the \n",
+    "number of non-zero variables."
+   ]
+  },
+  {
+   "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"
+     ]
+    }
+   ],
+   "source": [
+    "from entmoot.benchmarks import build_reals_only_problem, eval_reals_only_testfunc\n",
+    "\n",
+    "# standard setting up of problem\n",
+    "problem_config = ProblemConfig(rnd_seed=73)\n",
+    "build_reals_only_problem(problem_config)\n",
+    "rnd_sample = problem_config.get_rnd_sample_list(num_samples=50)\n",
+    "testfunc_evals = eval_reals_only_testfunc(rnd_sample)\n",
+    "\n",
+    "params = {\"unc_params\": {\"dist_metric\": \"l1\", \"acq_sense\": \"penalty\"}}\n",
+    "enting = Enting(problem_config, params=params)\n",
+    "# fit tree ensemble\n",
+    "enting.fit(rnd_sample, testfunc_evals)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Set parameter Username\n",
+      "Academic license - for non-commercial use only - expires 2024-09-06\n",
+      "Read LP format model from file C:\\Users\\tobyb\\AppData\\Local\\Temp\\tmpzrofd3mo.pyomo.lp\n",
+      "Reading time = 0.02 seconds\n",
+      "x1: 2774 rows, 1913 columns, 9130 nonzeros\n",
+      "Gurobi Optimizer version 10.0.2 build v10.0.2rc0 (win64)\n",
+      "\n",
+      "CPU model: 11th Gen Intel(R) Core(TM) i7-1165G7 @ 2.80GHz, instruction set [SSE2|AVX|AVX2|AVX512]\n",
+      "Thread count: 4 physical cores, 8 logical processors, using up to 8 threads\n",
+      "\n",
+      "Optimize a model with 2774 rows, 1913 columns and 9130 nonzeros\n",
+      "Model fingerprint: 0x31e842ff\n",
+      "Variable types: 1292 continuous, 621 integer (621 binary)\n",
+      "Coefficient statistics:\n",
+      "  Matrix range     [1e-06, 1e+06]\n",
+      "  Objective range  [1e+00, 2e+00]\n",
+      "  Bounds range     [1e+00, 5e+00]\n",
+      "  RHS range        [1e-04, 5e+00]\n",
+      "Presolve removed 273 rows and 260 columns\n",
+      "Presolve time: 0.05s\n",
+      "Presolved: 2501 rows, 1653 columns, 8080 nonzeros\n",
+      "Variable types: 1282 continuous, 371 integer (371 binary)\n",
+      "Found heuristic solution: objective 24.8382603\n",
+      "Found heuristic solution: objective 19.1248452\n",
+      "\n",
+      "Root relaxation: objective 2.576224e+00, 577 iterations, 0.01 seconds (0.01 work units)\n",
+      "\n",
+      "    Nodes    |    Current Node    |     Objective Bounds      |     Work\n",
+      " Expl Unexpl |  Obj  Depth IntInf | Incumbent    BestBd   Gap | It/Node Time\n",
+      "\n",
+      "     0     0    2.57622    0   18   19.12485    2.57622  86.5%     -    0s\n",
+      "H    0     0                       3.6397521    2.57622  29.2%     -    0s\n",
+      "H    0     0                       3.3869028    3.24087  4.31%     -    0s\n",
+      "     0     0    3.38690    0    6    3.38690    3.38690  0.00%     -    0s\n",
+      "\n",
+      "Cutting planes:\n",
+      "  Gomory: 4\n",
+      "  Cover: 83\n",
+      "  Implied bound: 402\n",
+      "  Clique: 185\n",
+      "  MIR: 4\n",
+      "  Flow cover: 13\n",
+      "  Network: 8\n",
+      "  RLT: 30\n",
+      "  Relax-and-lift: 98\n",
+      "  PSD: 4\n",
+      "\n",
+      "Explored 1 nodes (877 simplex iterations) in 0.14 seconds (0.13 work units)\n",
+      "Thread count was 8 (of 8 available processors)\n",
+      "\n",
+      "Solution count 4: 3.3869 3.63975 19.1248 24.8383 \n",
+      "\n",
+      "Optimal solution found (tolerance 1.00e-04)\n",
+      "Best objective 3.386902819809e+00, best bound 3.386902819809e+00, gap 0.0000%\n"
+     ]
+    }
+   ],
+   "source": [
+    "from entmoot.constraints import NChooseKConstraint\n",
+    "model_pyo = problem_config.get_pyomo_model_core()\n",
+    "\n",
+    "# define the constraint\n",
+    "# then immediately apply it to the model\n",
+    "model_pyo.nchoosek = NChooseKConstraint(\n",
+    "    feature_keys=[\"x1\", \"x2\", \"x3\", \"x4\", \"x5\"], \n",
+    "    min_count=1,\n",
+    "    max_count=3,\n",
+    "    none_also_valid=True\n",
+    ").as_pyomo_constraint(model_pyo, problem_config.feat_list)\n",
+    "\n",
+    "\n",
+    "# optimise the model\n",
+    "params_pyomo = {\"solver_name\": \"gurobi\"}\n",
+    "opt_pyo = PyomoOptimizer(problem_config, params=params_pyomo)\n",
+    "res_pyo = opt_pyo.solve(enting, model_core=model_pyo)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[0.0, 0.0, 4.088297585401641, 4.888952150927435, 4.944564863420855]\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(res_pyo.opt_point)\n",
+    "assert 1 <= sum(x > 1e-6 for x in res_pyo.opt_point) <= 3"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Defining your own constraint\n",
+    "\n",
+    "We have provided some constraints already as a part of ENTMOOT. If these do not \n",
+    "fit your needs, then you can define your own!\n",
+    "\n",
+    "The easiest approach is to subclass ExpressionConstraint, and define some custom expression\n",
+    "that is a function of the variables. From that, you should be able to use the constraint \n",
+    "as shown above. This needs to return a pyomo.Expression object. If you need to do \n",
+    "a more involved procedure that modifies the model, you can use a FunctionalConstraint \n",
+    "instead (see NChooseKConstraint)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from entmoot.constraints import ExpressionConstraint\n",
+    "\n",
+    "class SumLessThanTen(ExpressionConstraint):\n",
+    "    \"\"\"A constraint that enforces all features to be equal.\"\"\"\n",
+    "    def _get_expr(self, features):\n",
+    "        return sum(features) <= 10"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "enttest",
+   "language": "python",
+   "name": "enttest"
+  },
+  "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
+}

entmoot/benchmarks.py

@@ -149,3 +149,24 @@ def compute_objectives(xi: Iterable, no_cat=False):
             f"You provided the illegal value {n_obj} for the number of objectives. "
             f"Allowed values are 1 and 2"
         )
+
+def build_reals_only_problem(problem_config: ProblemConfig):
+    """A problem containing only real values, as used to demonstrate the NChooseK 
+    constraint.
+    
+    The minimum is (1.0, 2.0, 3.0, ...)"""
+
+    problem_config.add_feature("real", (0.0, 5.0), name="x1")
+    problem_config.add_feature("real", (0.0, 5.0), name="x2")
+    problem_config.add_feature("real", (0.0, 5.0), name="x3")
+    problem_config.add_feature("real", (0.0, 5.0), name="x4")
+    problem_config.add_feature("real", (0.0, 5.0), name="x5")
+    problem_config.add_min_objective()
+
+def eval_reals_only_testfunc(X: ArrayLike):
+    """The function (x1 - 1)**2 + (x2 - 2)**2 + ..."""
+    x = np.array(X)
+    xbar = np.ones_like(x)
+    xbar *= (np.arange(x.shape[1]) + 1)[None, :]
+    y = np.sum((x - xbar)**2, axis=1)
+    return y.reshape(-1, 1)

entmoot/constraints.py

@@ -0,0 +1,137 @@
+from typing import TYPE_CHECKING, Callable
+from abc import ABC, abstractmethod
+
+import pyomo.environ as pyo
+
+if TYPE_CHECKING:
+    from problem_config import FeatureType
+
+ConstraintFunctionType = Callable[[pyo.ConcreteModel, int], pyo.Expression]
+
+
+class Constraint(ABC):
+    """A constraint to be applied to a model.
+
+    Implements a user-friendly way to construct constraints to an optimisation problem.
+
+    Attributes:
+        feature_keys: A list of the string names of the features to be constrained"""
+
+    def __init__(self, feature_keys: list[str]):
+        self.feature_keys = feature_keys
+
+    def _get_feature_vars(
+        self, model: pyo.ConcreteModel, feat_list: list["FeatureType"]
+    ) -> list[pyo.Var]:
+        """Return a list of all the pyo.Vars, in the order of the constraint definition"""
+        all_keys = [feat.name for feat in feat_list]
+        feat_idxs = [all_keys.index(key) for key in self.feature_keys]
+        features = [model._all_feat[i] for i in feat_idxs]
+        return features
+
+    @abstractmethod
+    def as_pyomo_constraint(
+        self, model: pyo.ConcreteModel, feat_list: list["FeatureType"]
+    ) -> pyo.Constraint:
+        """Convert to a pyomo.Constraint object.
+
+        This requires the model (to access the variables), and the feat_list (to access the feature names)
+        """
+        pass
+
+
+class ExpressionConstraint(Constraint):
+    """Constraints defined by pyomo.Expressions.
+
+    For constraints that can be simply defined by an expression of variables.
+    """
+
+    def as_pyomo_constraint(
+        self, model: pyo.ConcreteModel, feat_list: list["FeatureType"]
+    ) -> pyo.Constraint:
+        features = self._get_feature_vars(model, feat_list)
+        return pyo.Constraint(expr=self._get_expr(features))
+
+    @abstractmethod
+    def _get_expr(self, features) -> pyo.Expression:
+        pass
+
+
+class FunctionalConstraint(Constraint):
+    """A constraint that uses a functional approach.
+
+    For constraints that require creating intermediate variables and access to the model.
+    """
+
+    def as_pyomo_constraint(
+        self, model: pyo.ConcreteModel, feat_list: list["FeatureType"]
+    ) -> pyo.Constraint:
+        features = self._get_feature_vars(model, feat_list)
+        return pyo.Constraint(rule=self._get_function(model, features))
+
+    @abstractmethod
+    def _get_function(self, features) -> ConstraintFunctionType:
+        pass
+
+
+class LinearConstraint(ExpressionConstraint):
+    """Constraint that is a function of X @ C, where X is the feature list, and C
+    is the list of coefficients."""
+
+    def __init__(self, feature_keys: list[str], coefficients: list[float], rhs: float):
+        self.coefficients = coefficients
+        self.rhs = rhs
+        super().__init__(feature_keys)
+
+    def _get_lhs(self, features: pyo.ConcreteModel) -> pyo.Expression:
+        """Get the left-hand side of the linear constraint"""
+        return sum(f * c for f, c in zip(features, self.coefficients))
+
+
+class LinearEqualityConstraint(LinearConstraint):
+    def _get_expr(self, features):
+        return self._get_lhs(features) == self.rhs
+
+
+class LinearInequalityConstraint(LinearConstraint):
+    def _get_expr(self, features):
+        return self._get_lhs(features) <= self.rhs
+
+
+class NChooseKConstraint(FunctionalConstraint):
+    """Constrain the number of active features to be bounded by min_count and max_count."""
+
+    tol: float = 1e-6
+    M: float = 1e6
+
+    def __init__(
+        self,
+        feature_keys: list[str],
+        min_count: int,
+        max_count: int,
+        none_also_valid: bool = False,
+    ):
+        self.min_count = min_count
+        self.max_count = max_count
+        self.none_also_valid = none_also_valid
+        super().__init__(feature_keys)
+
+    def _get_function(self, model, features):
+        # constrain the features using the binary variable y
+        # where y indicates whether the feature is selected
+        # y * tol <= x <= y * M
+        # tol is sufficiently small, M is sufficiently large
+        model.feat_selected = pyo.Var(
+            range(len(features)), domain=pyo.Binary, initialize=0
+        )
+        model.ub_selected = pyo.ConstraintList()
+        model.lb_selected = pyo.ConstraintList()
+
+        for i in range(len(features)):
+            model.ub_selected.add(expr=model.feat_selected[i] * self.M >= features[i])
+            model.lb_selected.add(expr=model.feat_selected[i] * self.tol <= features[i])
+
+        def inner(model, i):
+            return sum(model.feat_selected.values()) <= self.max_count
+
+        return inner