{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Linear Program (LP) Tutorial\n",
    "For instructions on how to run these tutorial notebooks, please see the [README](https://github.com/RobotLocomotion/drake/blob/master/tutorials/README.md)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Important Note\n",
    "Please refer to [mathematical program tutorial](./mathematical_program.ipynb) for constructing and solving a general optimization program in Drake."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Linear Program\n",
    "A linear program (LP) is a special type of optimization problem. The cost and constraints in an LP is a linear (affine) function of decision variables. The mathematical formulation of a general LP is\n",
    "\\begin{align}\n",
    "\\min_x \\;c^Tx + d\\\\\n",
    "\\text{subject to } Ax\\leq b\n",
    "\\end{align}\n",
    "\n",
    "A linear program can be solved by many open source or commercial solvers. Drake supports some solvers including SCS, Gurobi, Mosek, etc. Please see our [Doxygen page]( https://drake.mit.edu/doxygen_cxx/group__solvers.html) for a complete list of supported solvers. Note that some commercial solvers (such as Gurobi and Mosek) are not included in the pre-compiled Drake binaries, and therefore not on Binder/Colab. \n",
    "    \n",
    "Drake's API supports multiple functions to add linear cost and constraints. We briefly go through some of the functions in this tutorial. For a complete list of functions, please check our [Doxygen](https://drake.mit.edu/doxygen_cxx/classdrake_1_1solvers_1_1_mathematical_program.html).\n",
    "\n",
    "### Add linear cost\n",
    "The easiest way to add linear cost is to call `AddLinearCost` function. We first demonstrate how to construct an optimization program with 2 decision variables, then we will call `AddLinearCost` to add the cost."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.solvers.mathematicalprogram import MathematicalProgram, Solve\n",
    "import numpy as np\n",
    "\n",
    "# Create an empty MathematicalProgram named prog (with no decision variables,\n",
    "# constraints or costs)\n",
    "prog = MathematicalProgram()\n",
    "# Add two decision variables x[0], x[1].\n",
    "x = prog.NewContinuousVariables(2, \"x\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can call `AddLinearCost(expression)` to add a new linear cost. `expression` is a symbolic linear expression of the decision variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a symbolic linear expression as the cost.\n",
    "cost1 = prog.AddLinearCost(x[0] + 3 * x[1] + 2)\n",
    "# Print the newly added cost\n",
    "print(cost1)\n",
    "# The newly added cost is stored in prog.linear_costs().\n",
    "print(prog.linear_costs()[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we call `AddLinearCost` again, the total cost stored in `prog` is the summation of all the costs. You can see that `prog.linear_costs()` will have two entries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cost2 = prog.AddLinearCost(2 * x[1] + 3)\n",
    "print(f\"number of linear cost objects: {len(prog.linear_costs())}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you know the coefficient of the linear cost as a vector, you could also add the cost by calling `AddLinearCost(e, f, x)` which will add a linear cost $e^Tx + f$ to the optimization program"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We add a linear cost 3 * x[0] + 4 * x[1] + 5 to prog by specifying the coefficients\n",
    "# [3., 4] and the constant 5 in AddLinearCost\n",
    "cost3 = prog.AddLinearCost([3., 4.], 5., x)\n",
    "print(cost3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, the user can call `AddCost` to add a linear expression to the linear cost. Drake will analyze the structure of the expression, if Drake determines the expression is linear, then the added cost is linear."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"number of linear cost objects before calling AddCost: {len(prog.linear_costs())}\")\n",
    "# Call AddCost to add a linear expression as linear cost. After calling this function,\n",
    "# len(prog.linear_costs()) will increase by 1.\n",
    "cost4 = prog.AddCost(x[0] + 3 * x[1] + 5)\n",
    "print(f\"number of linear cost objects after calling AddCost: {len(prog.linear_costs())}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add linear constraints\n",
    "We have three types of linear constraints\n",
    "  * Bounding box constraint. A lower/upper bound on the decision variable: $ lower \\le x \\le upper $.\n",
    "  * Linear equality constraint: $Ax = b$.\n",
    "  * Linear inequality constraint: $lower <= Ax <= upper$.\n",
    "  \n",
    "#### AddLinearConstraint and AddConstraint function\n",
    "The easiest way to add linear constraints is to call `AddConstraint` or `AddLinearConstraint` function, which can handle all three types of linear constraint. Compared to the generic `AddConstraint` function, `AddLinearConstraint` does more sanity will refuse to add the constraint if it is not linear."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prog = MathematicalProgram()\n",
    "x = prog.NewContinuousVariables(2, \"x\")\n",
    "y = prog.NewContinuousVariables(3, \"y\")\n",
    "\n",
    "# Call AddConstraint to add a bounding box constraint x[0] >= 1\n",
    "bounding_box1 = prog.AddConstraint(x[0] >= 1)\n",
    "print(f\"number of bounding box constraint objects: {len(prog.bounding_box_constraints())}\")\n",
    "\n",
    "# Call AddLinearConstraint to add a bounding box constraint x[1] <= 2\n",
    "bounding_box2 = prog.AddLinearConstraint(x[1] <= 2)\n",
    "print(f\"number of bounding box constraint objects: {len(prog.bounding_box_constraints())}\")\n",
    "\n",
    "# Call AddConstraint to add a linear equality constraint x[0] + y[1] == 3\n",
    "linear_eq1 = prog.AddConstraint(x[0] + y[1] == 3.)\n",
    "print(f\"number of linear equality constraint objects: {len(prog.linear_equality_constraints())}\")\n",
    "\n",
    "# Call AddLinearConstraint to add a linear equality constraint x[1] + 2 * y[2] == 1\n",
    "linear_eq2 = prog.AddLinearConstraint(x[1] + 2 * y[2] == 1)\n",
    "print(f\"number of linear equality constraint objects: {len(prog.linear_equality_constraints())}\")\n",
    "\n",
    "# Call AddConstraint to add a linear inequality constraint x[0] + 3*x[1] + 2*y[2] <= 4\n",
    "linear_ineq1 = prog.AddConstraint(x[0] + 3*x[1] + 2*y[2] <= 4)\n",
    "print(f\"number of linear inequality constraint objects: {len(prog.linear_constraints())}\")\n",
    "\n",
    "# Call AddLinearConstraint to add a linear inequality constraint x[1] + 4 * y[1] >= 2\n",
    "linear_ineq2 = prog.AddLinearConstraint(x[1] + 4 * y[1] >= 2)\n",
    "print(f\"number of linear inequality constraint objects: {len(prog.linear_constraints())}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`AddLinearConstraint` will check if the constraint is actually linear, and throw an exception if the constraint is not linear."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a nonlinear constraint square(x[0]) == 2 by calling AddLinearConstraint. This should\n",
    "# throw an exception\n",
    "try:\n",
    "    prog.AddLinearConstraint(x[0] ** 2 == 2)\n",
    "except RuntimeError as err:\n",
    "    print(err.args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the users know the coefficients of the constraint as a matrix, they could also call `AddLinearConstraint(A, lower, upper, x)` to add a constraint $lower \\le Ax \\le upper$. This version of the method does not construct any symbolic representations, and will be more efficient especially when `A` is very large."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a linear constraint 2x[0] + 3x[1] <= 2, 1 <= 4x[1] + 5y[2] <= 3.\n",
    "# This is equivalent to lower <= A * [x;y[2]] <= upper with\n",
    "# lower = [-inf, 1], upper = [2, 3], A = [[2, 3, 0], [0, 4, 5]].\n",
    "linear_constraint = prog.AddLinearConstraint(\n",
    "    A=[[2., 3., 0], [0., 4., 5.]],\n",
    "    lb=[-np.inf, 1],\n",
    "    ub=[2., 3.],\n",
    "    vars=np.hstack((x, y[2])))\n",
    "print(linear_constraint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### AddBoundingBoxConstraint\n",
    "If your constraint is a bounding box constraint (i.e. $lower \\le x \\le upper$), apart from calling `AddConstraint` or `AddLinearConstraint`, you could also call `AddBoundingBoxConstraint(lower, upper, x)`, which  will be slightly faster than `AddConstraint` and `AddLinearConstraint`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a bounding box constraint -1 <= x[0] <= 2, 3 <= x[1] <= 5\n",
    "bounding_box3 = prog.AddBoundingBoxConstraint([-1, 3], [2, 5], x)\n",
    "print(bounding_box3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the variables share the same lower or upper bound, you could use a scalar `lower` or `upper` value in `AddBoundingBoxConstraint`. For example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a bounding box constraint 3 <= y[i] <= 5 for all i.\n",
    "bounding_box4 = prog.AddBoundingBoxConstraint(3, 5, y)\n",
    "print(bounding_box4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### AddLinearEqualityConstraint\n",
    "If your constraint is a linear equality constraint (i.e. $ Ax = b$), apart from calling `AddConstraint` or `AddLinearConstraint`, you could also call `AddLinearEqualityConstraint` to be more specific (and slightly faster than `AddConstraint` and `AddLinearConstraint`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a linear equality constraint 4 * x[0] + 5 * x[1] == 1\n",
    "linear_eq3 = prog.AddLinearEqualityConstraint(np.array([[4, 5]]), np.array([1]), x)\n",
    "print(linear_eq3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solving Linear Program.\n",
    "\n",
    "Once all the constraints and costs are added to the program, we can call `Solve` function to solve the program and call `GetSolution` to obtain the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve an optimization program\n",
    "# min -3x[0] - x[1] - 5x[2] -x[3] + 2\n",
    "# s.t 3x[0] + x[1] + 2x[2] = 30\n",
    "#     2x[0] + x[1] + 3x[2] + x[3] >= 15\n",
    "#     2x[1] + 3x[3] <= 25\n",
    "#     -100 <= x[0] + 2x[2] <= 40\n",
    "#   x[0], x[1], x[2], x[3] >= 0, x[1] <= 10\n",
    "prog = MathematicalProgram()\n",
    "# Declare x as decision variables.\n",
    "x = prog.NewContinuousVariables(4)\n",
    "# Add linear costs. To show that calling AddLinearCosts results in the sum of each individual\n",
    "# cost, we add two costs -3x[0] - x[1] and -5x[2]-x[3]+2\n",
    "prog.AddLinearCost(-3*x[0] -x[1])\n",
    "prog.AddLinearCost(-5*x[2] - x[3] + 2)\n",
    "# Add linear equality constraint 3x[0] + x[1] + 2x[2] == 30\n",
    "prog.AddLinearConstraint(3*x[0] + x[1] + 2*x[2] == 30)\n",
    "# Add Linear inequality constraints\n",
    "prog.AddLinearConstraint(2*x[0] + x[1] + 3*x[2] + x[3] >= 15)\n",
    "prog.AddLinearConstraint(2*x[1] + 3*x[3] <= 25)\n",
    "# Add linear inequality constraint -100 <= x[0] + 2x[2] <= 40\n",
    "prog.AddLinearConstraint(A=[[1., 2.]], lb=[-100], ub=[40], vars=[x[0], x[2]])\n",
    "prog.AddBoundingBoxConstraint(0, np.inf, x)\n",
    "prog.AddLinearConstraint(x[1] <= 10)\n",
    "\n",
    "# Now solve the program.\n",
    "result = Solve(prog)\n",
    "print(f\"Is solved succesfully: {result.is_success()}\")\n",
    "print(f\"x optimal value: {result.GetSolution(x)}\")\n",
    "print(f\"optimal cost: {result.get_optimal_cost()}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}