{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Nonlinear Program (NLP) Tutorial\n",
    "For instructions on how to run these tutorial notebooks, please see the [index](./index.ipynb).\n",
    "\n",
    "## Important Note\n",
    "Please refer to the [MathematicalProgram Tutorial](./mathematical_program.ipynb) for constructing and solving a general optimization program in Drake.\n",
    "\n",
    "## Nonlinear Program\n",
    "A Nonlinear Programming (NLP) problem is a special type of optimization problem. The cost and/or constraints in an NLP are nonlinear functions of decision variables. The mathematical formulation of a general NLP is\n",
    "\n",
    "$\\begin{aligned} \\min_x&\\; f(x)\\\\ \\text{subject to }& g_i(x)\\leq 0 \\end{aligned}$\n",
    "\n",
    "where $f(x)$ is the cost function, and $g_i(x)$ is the i'th constraint.\n",
    "\n",
    "An NLP is typically solved through gradient-based optimization (like gradient descent, SQP, interior point methods, etc). These methods rely on the gradient of the cost/constraints $\\partial f/\\partial x, \\partial g_i/\\partial x$. pydrake can compute the gradient of many functions through automatic differentiation, so very often the user doesn't need to manually provide the gradient.\n",
    "\n",
    "## Setting the objective\n",
    "The user can call `AddCost` function to add a nonlinear cost into the program. Note that the user can call `AddCost` repeatedly, and the program will evaluate the *summation* of each individual cost as the total cost."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding a cost through a python function\n",
    "We can define a cost through a python function, and then add this python function to the objective through `AddCost` function. When adding a cost, we should provide the variable associated with that cost, the syntax is `AddCost(cost_evaluator, vars=associated_variables)`, which means the cost is evaluated on the `associated_variables`.In the code example below, We first demonstrate how to construct an optimization program with 3 decision variables, then we show how to add a cost through a python function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.solvers 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 three decision variables x[0], x[1], x[2]\n",
    "x = prog.NewContinuousVariables(3, \"x\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cost_fun(z):\n",
    "    cos_z = np.cos(z[0] + z[1])\n",
    "    sin_z = np.sin(z[0] + z[1])\n",
    "    return cos_z**2 + cos_z + sin_z\n",
    "# Add the cost evaluated with x[0] and x[1].\n",
    "cost1 = prog.AddCost(cost_fun, vars=[x[0], x[1]])\n",
    "print(cost1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that by changing the argument `vars` in `AddCost` function, we can add the cost to a different set of variables. In the code example below, we use the same python function `cost_fun`, but impose this cost on the variable `x[0], x[2]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cost2 = prog.AddCost(cost_fun, vars=[x[0], x[2]])\n",
    "print(cost2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding cost through a lambda function\n",
    "A more compact approach to add a cost is through a lambda function. For example, the code below adds a cost $x[1]^2 + x[0]$ to the optimization program."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a cost x[1]**2 + x[0] using a lambda function.\n",
    "cost3 = prog.AddCost(lambda z: z[0]**2 + z[1], vars = [x[1], x[0]])\n",
    "print(cost3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we change the associated variables, then it represents a different cost. For example, we can use the same lambda function, but add the cost $x[1]^2 + x[2]$ to the program by changing the argument to `vars`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cost4 = prog.AddCost(lambda z: z[0]**2 + z[1], vars = x[1:])\n",
    "print(cost4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding quadratic cost\n",
    "In NLP, adding a quadratic cost $0.5x^TQx+ b'x+c$ is very common. pydrake provides multiple functions to add quadratic cost, including\n",
    "- `AddQuadraticCost`\n",
    "- `AddQuadraticErrorCost`\n",
    "- `Add2NormSquaredCost`\n",
    "\n",
    "### AddQuadraticCost\n",
    "We can add a simple quadratic expression as a cost."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cost4 = prog.AddQuadraticCost(x[0]**2 + 3 * x[1]**2 + 2*x[0]*x[1] + 2*x[1] * x[0] + 1)\n",
    "print(cost4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the user knows the matrix form of `Q` and `b`, then it is faster to pass in these matrices to `AddQuadraticCost`, instead of using the symbolic quadratic expression as above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a cost x[0]**2 + 2*x[1]**2 + x[0]*x[1] + 3*x[1] + 1.\n",
    "cost5 = prog.AddQuadraticCost(\n",
    "    Q=np.array([[2., 1], [1., 4.]]),\n",
    "    b=np.array([0., 3.]),\n",
    "    c=1.,\n",
    "    vars=x[:2])\n",
    "print(cost5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### AddQuadraticErrorCost\n",
    "This function adds a cost of the form $(x - x_{des})^TQ(x-x_{des})$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cost6 = prog.AddQuadraticErrorCost(\n",
    "    Q=np.array([[1, 0.5], [0.5, 1]]),\n",
    "    x_desired=np.array([1., 2.]),\n",
    "    vars=x[1:])\n",
    "print(cost6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add2NormSquaredCost\n",
    "This function adds a quadratic cost of the form $(Ax-b)^T(Ax-b)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add the L2 norm cost on (A*x[:2] - b).dot(A*x[:2]-b)\n",
    "cost7 = prog.Add2NormSquaredCost(\n",
    "    A=np.array([[1., 2.], [2., 3], [3., 4]]),\n",
    "    b=np.array([2, 3, 1.]),\n",
    "    vars=x[:2])\n",
    "print(cost7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding constraints\n",
    "\n",
    "Drake supports adding constraints in the following form\n",
    "$$\n",
    "\\begin{aligned}\n",
    "lower \\leq g(x) \\leq upper\n",
    "\\end{aligned}\n",
    "$$\n",
    "where $g(x)$ returns a numpy vector.\n",
    "\n",
    "The user can call `AddConstraint(g, lower, upper, vars=x)` to add the constraint. Here `g` must be a python function (or a lambda function)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Define a python function to add the constraint x[0]**2 + 2x[1]<=1, -0.5<=sin(x[1])<=0.5\n",
    "def constraint_evaluator1(z):\n",
    "    return np.array([z[0]**2+2*z[1], np.sin(z[1])])\n",
    "\n",
    "constraint1 = prog.AddConstraint(\n",
    "    constraint_evaluator1,\n",
    "    lb=np.array([-np.inf, -0.5]),\n",
    "    ub=np.array([1., 0.5]),\n",
    "    vars=x[:2])\n",
    "print(constraint1)\n",
    "\n",
    "# Add another constraint using lambda function.\n",
    "constraint2 = prog.AddConstraint(\n",
    "    lambda z: np.array([z[0]*z[1]]),\n",
    "    lb=[0.],\n",
    "    ub=[1.],\n",
    "    vars=[x[2]])\n",
    "print(constraint2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solving the nonlinear program\n",
    "\n",
    "Once all the constraints and costs are added to the program, we can call the `Solve` function to solve the program and call `GetSolution` to obtain the results. Solving an NLP requires an initial guess on all the decision variables. If the user doesn't specify an initial guess, we will use a zero vector by default."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Solve a simple nonlinear \n",
    "# min               -x0\n",
    "# subject to x1 - exp(x0) >= 0\n",
    "#            x2 - exp(x1) >= 0\n",
    "#            0 <= x0 <= 100\n",
    "#            0 <= x1 <= 100\n",
    "#            0 <= x2 <= 10\n",
    "prog = MathematicalProgram()\n",
    "x = prog.NewContinuousVariables(3)\n",
    "# The cost is a linear function, so we call AddLinearCost\n",
    "prog.AddLinearCost(-x[0])\n",
    "# Now add the constraint x1-exp(x0)>=0 and x2-exp(x1)>=0\n",
    "prog.AddConstraint(\n",
    "    lambda z: np.array([z[1]-np.exp(z[0]), z[2]-np.exp(z[1])]),\n",
    "    lb=[0, 0],\n",
    "    ub=[np.inf, np.inf],\n",
    "    vars=x)\n",
    "# Add the bounding box constraint 0<=x0<=100, 0<=x1<=100, 0<=x2<=10\n",
    "prog.AddBoundingBoxConstraint(0, 100, x[:2])\n",
    "prog.AddBoundingBoxConstraint(0, 10, x[2])\n",
    "\n",
    "# Now solve the program with initial guess x=[1, 2, 3]\n",
    "result = Solve(prog, np.array([1.,2.,3.]))\n",
    "print(f\"Is optimization successful? {result.is_success()}\")\n",
    "print(f\"Solution to x: {result.GetSolution(x)}\")\n",
    "print(f\"optimal cost: {result.get_optimal_cost()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting the initial guess\n",
    "Some NLPs might have many decision variables. In order to set the initial guess for these decision variables, we provide a function `SetInitialGuess` to set the initial guess of a subset of decision variables. For example, in the problem below, we want to find the two closest points $p_1$ and $p_2$, where $p_1$ is on the unit circle, and $p_2$ is on the curve $y=x^2$, we can set the initial guess for these two variables separately by calling `SetInitialGuess`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "prog = MathematicalProgram()\n",
    "p1 = prog.NewContinuousVariables(2, \"p1\")\n",
    "p2 = prog.NewContinuousVariables(2, \"p2\")\n",
    "\n",
    "# Add the constraint that p1 is on the unit circle centered at (0, 2)\n",
    "prog.AddConstraint(\n",
    "    lambda z: [z[0]**2 + (z[1]-2)**2],\n",
    "    lb=np.array([1.]),\n",
    "    ub=np.array([1.]),\n",
    "    vars=p1)\n",
    "\n",
    "# Add the constraint that p2 is on the curve y=x*x\n",
    "prog.AddConstraint(\n",
    "    lambda z: [z[1] - z[0]**2],\n",
    "    lb=[0.],\n",
    "    ub=[0.],\n",
    "    vars=p2)\n",
    "\n",
    "# Add the cost on the distance between p1 and p2\n",
    "prog.AddQuadraticCost((p1-p2).dot(p1-p2))\n",
    "\n",
    "# Set the value of p1 in initial guess to be [0, 1]\n",
    "prog.SetInitialGuess(p1, [0., 1.])\n",
    "# Set the value of p2 in initial guess to be [1, 1]\n",
    "prog.SetInitialGuess(p2, [1., 1.])\n",
    "\n",
    "# Now solve the program\n",
    "result = Solve(prog)\n",
    "print(f\"Is optimization successful? {result.is_success()}\")\n",
    "p1_sol = result.GetSolution(p1)\n",
    "p2_sol = result.GetSolution(p2)\n",
    "print(f\"solution to p1 {p1_sol}\")\n",
    "print(f\"solution to p2 {p2_sol}\")\n",
    "print(f\"optimal cost {result.get_optimal_cost()}\")\n",
    "\n",
    "# Plot the solution.\n",
    "plt.figure()\n",
    "plt.plot(np.cos(np.linspace(0, 2*np.pi, 100)), 2+np.sin(np.linspace(0, 2*np.pi, 100)))\n",
    "plt.plot(np.linspace(-2, 2, 100), np.power(np.linspace(-2, 2, 100), 2))\n",
    "plt.plot(p1_sol[0], p1_sol[1], '*')\n",
    "plt.plot(p2_sol[0], p2_sol[1], '*')\n",
    "plt.axis('equal')\n",
    "plt.show()"
   ]
  }
 ],
 "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
}