{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6f01675f-cd1a-4712-9ca4-3cc1dbbd9056",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Bring your own constraints!\n",
    "The following example shows how your own constraints can be enforced on the input\n",
    "variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c8da843-670c-4f07-bd66-471ec19d3601",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from entmoot import Enting, ProblemConfig, GurobiOptimizer, PyomoOptimizer\n",
    "from entmoot.benchmarks import (\n",
    "    build_multi_obj_categorical_problem,\n",
    "    eval_multi_obj_cat_testfunc,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20d1e71a-2eff-49d1-8d45-412ebdc5412a",
   "metadata": {},
   "source": [
    "This part is pretty standard."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "311c88fa-f757-44f3-8ae5-555f715fc1b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define problem\n",
    "problem_config = ProblemConfig(rnd_seed=73)\n",
    "# number of objectives\n",
    "number_objectives = 2\n",
    "build_multi_obj_categorical_problem(problem_config, n_obj=number_objectives)\n",
    "\n",
    "# sample data\n",
    "rnd_sample = problem_config.get_rnd_sample_list(num_samples=20)\n",
    "testfunc_evals = eval_multi_obj_cat_testfunc(rnd_sample, n_obj=number_objectives)\n",
    "\n",
    "params = {\"unc_params\": {\"dist_metric\": \"l1\", \"acq_sense\": \"exploration\"}}\n",
    "enting = Enting(problem_config, params=params)\n",
    "# fit tree ensemble\n",
    "enting.fit(rnd_sample, testfunc_evals)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "818791f2-9784-4b33-acf7-fba8a4747a93",
   "metadata": {},
   "source": [
    "How to add constraints depends on wether you are using Gurobi or Pyomo"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "943e15bb-c253-4d2f-8937-4708cd090fec",
   "metadata": {},
   "source": [
    "### Gurobi Version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0dc70de-f14f-42a4-9202-0e4777d33bec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get optimization model\n",
    "model_gur = problem_config.get_gurobi_model_core()\n",
    "# extract decision variables\n",
    "x = model_gur._all_feat[3]\n",
    "y = model_gur._all_feat[4]\n",
    "z = model_gur._all_feat[5]\n",
    "# add constraint that all variables should coincide\n",
    "model_gur.addConstr(x == y)\n",
    "model_gur.addConstr(y == z)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae8f39bd-224d-4bd7-9b6b-5bbc6a01d355",
   "metadata": {},
   "source": [
    "It is important to update the Gurobi model. Otherwise the constraints will not be added to the model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61a8e860-92ee-4535-ad73-9b2acf132ad8",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_gur.update()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0d9a901-c7b4-46d9-a932-f1b2292476cb",
   "metadata": {},
   "source": [
    "Now you can run the optimization and verify that the variable values indeed equal each other."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15bc6c5d-011f-4c57-8af0-4da9704af41e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build GurobiOptimizer object and solve optimization problem\n",
    "params_gurobi = {\"MIPGap\": 0}\n",
    "opt_gur = GurobiOptimizer(problem_config, params=params_gurobi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e987e159-d855-4476-b7f8-3a041ab45d5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_gur = opt_gur.solve(enting, model_core=model_gur)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ae2f026-dded-4376-866b-6ca34adc5bf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build GurobiOptimizer object and solve optimization problem\n",
    "params_gurobi = {\"MIPGap\": 0}\n",
    "opt_gur = GurobiOptimizer(problem_config, params=params_gurobi)\n",
    "\n",
    "res_gur = opt_gur.solve(enting, model_core=model_gur)\n",
    "x_opt, y_opt, z_opt = res_gur.opt_point[3:]\n",
    "\n",
    "assert round(x_opt, 5) == round(y_opt, 5) and round(y_opt, 5) == round(z_opt, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24b3d335-d601-41c1-ad46-d80949c0cfcc",
   "metadata": {},
   "source": [
    "### Pyomo Version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1644fd93-d413-4aaa-9d91-2467f2c59900",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pyomo version\n",
    "import pyomo.environ as pyo\n",
    "\n",
    "model_pyo = problem_config.get_pyomo_model_core()\n",
    "# extract decision variables\n",
    "x = model_pyo._all_feat[3]\n",
    "y = model_pyo._all_feat[4]\n",
    "z = model_pyo._all_feat[5]\n",
    "# add constraint that all variables should coincide\n",
    "model_pyo.xy_equal_constr = pyo.Constraint(expr=x == y)\n",
    "model_pyo.yz_equal_constr = pyo.Constraint(expr=y == z)\n",
    "\n",
    "# Build GurobiOptimizer object and solve optimization problem\n",
    "params_pyomo = {\"solver_name\": \"gurobi\"}\n",
    "opt_pyo = PyomoOptimizer(problem_config, params=params_pyomo)\n",
    "\n",
    "res_pyo = opt_pyo.solve(enting, model_core=model_pyo)\n",
    "x_opt, y_opt, z_opt = res_pyo.opt_point[3:]\n",
    "\n",
    "assert round(x_opt, 5) == round(y_opt, 5) and round(y_opt, 5) == round(z_opt, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9f1cad4-8d2b-4bc7-bc86-de06d4a141c3",
   "metadata": {},
   "source": [
    "Note that no model update is required in the Pyomo version in contrast to the Gurobi variant"
   ]
  }
 ],
 "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}