{
 "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": null,
   "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": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from entmoot.constraints import ExpressionConstraint\n",
    "\n",
    "class SumLessThanTen(ExpressionConstraint):\n",
    "    \"\"\"A constraint that enforces selected features to sum to less than ten.\"\"\"\n",
    "    def _get_expr(self, features):\n",
    "        return sum(features) <= 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constraint Lists\n",
    "\n",
    "For a problem definition, it may be easier to define a set of constraints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from entmoot.constraints import LinearInequalityConstraint, ConstraintList\n",
    "import pyomo.environ as pyo\n",
    "model_pyo = problem_config.get_pyomo_model_core()\n",
    "\n",
    "# define the constraint\n",
    "# then immediately apply it to the model\n",
    "constraints = [\n",
    "    NChooseKConstraint(\n",
    "        feature_keys=[\"x1\", \"x2\", \"x3\", \"x4\", \"x5\"], \n",
    "        min_count=1,\n",
    "        max_count=4,\n",
    "        none_also_valid=True\n",
    "    ),\n",
    "    LinearInequalityConstraint(\n",
    "        feature_keys=[\"x3\", \"x4\", \"x5\"],\n",
    "        coefficients=[1, 1, 1],\n",
    "        rhs=12.0\n",
    "    )\n",
    "]\n",
    "\n",
    "model_pyo.problem_constraints = pyo.ConstraintList()\n",
    "ConstraintList(constraints).apply_pyomo_constraints(\n",
    "    model_pyo, problem_config.feat_list, model_pyo.problem_constraints\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(res_pyo.opt_point)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}