{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MathematicalProgram debugging tips\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": [
    "After constructing and solving an optimization problem through Drake's `MathematicalProgram` interface, you might not get the desired results. For example, you might expect the problem to have a solution, while `MathematicalProgram` reports that the problem is not solved succesfully. In this tutorial we provide some tips to debug `MathematicalProgram` when it doesn't behave desirably.\n",
    "\n",
    "First you should understand whether the optimization problem is convex or not. For a convex problem (like LP, QP, SDP), when the problem is feasible, then theoretically the solver should always find a solution; on the other hand if the problem is non-convex and solved through gradient-based solvers (like SNOPT/IPOPT), then the solver might fail to terminate at a feasible solution even if one exists. When the gradient-based solver (like SNOPT/IPOPT) reports the problem being infeasible, it only means that the solver gets stuck at an infeasible value, and it doesn't know how to reduce the infeasibility in the local neighbourhood. A solution could exist far away from where the solver gets stuck, but the solver is trapped and can't jump to the distant solution. One possible solution is to choose a different initial guess of the optimization program.\n",
    "\n",
    "Here is an example to show the importance of initial guess in nonlinear optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from matplotlib import cm\n",
    "import pydrake.solvers.mathematicalprogram as mp\n",
    "from pydrake.solvers.ipopt import IpoptSolver\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def constraint(x):\n",
    "    return [np.cos(x[0]) + 2 * np.cos(x[0] - x[1])]\n",
    "\n",
    "# Find a solution satisfying\n",
    "# 1 <= cos(x[0]) + 2 * cos(x[0] - x[1]) <= 2\n",
    "# This problem has infinitely many solutions, for example, x = [0, pi/2]\n",
    "\n",
    "# To visualize the constraint, I draw the landscape of cos(x[0]) + 2 * cos(x[0] - x[1])), and\n",
    "# also highlight the point x = [0, 0]. You could see that x = [0, 0] is at a\n",
    "# peak of the landscape.\n",
    "def draw_constraint_landscape():\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111, projection='3d')\n",
    "    x_mesh, y_mesh = np.meshgrid(np.linspace(-np.pi, np.pi, 31), np.linspace(-np.pi, np.pi, 31))\n",
    "    constraint_val = np.cos(x_mesh) + 2 * np.cos(x_mesh - y_mesh)\n",
    "    surf = ax.plot_surface(x_mesh, y_mesh, constraint_val, cmap=cm.coolwarm, alpha=0.8)\n",
    "    ax.plot([0], [0], [3], marker='.', color='g', markersize=20)\n",
    "    ax.set_xlabel(\"x[0]\")\n",
    "    ax.set_ylabel(\"x[1]\")\n",
    "    ax.set_zlabel(\"cos(x[0]) + 2 * cos(x[0]-x[1])\")\n",
    "    fig.show()\n",
    "    \n",
    "\n",
    "draw_constraint_landscape()\n",
    "\n",
    "prog = mp.MathematicalProgram()\n",
    "x = prog.NewContinuousVariables(2, \"x\")\n",
    "\n",
    "prog.AddConstraint(constraint, [1], [2], x)\n",
    "\n",
    "solver = IpoptSolver()\n",
    "# With the initial guess being (0, 0), the solver cannot find a solution.\n",
    "prog.SetInitialGuess(x, [0, 0])\n",
    "result = solver.Solve(prog)\n",
    "print(f\"Starting from x=[0, 0], the solver result is {result.get_solution_result()}\")\n",
    "print(f\"The solver gets stuck at x={result.GetSolution(x)}\")\n",
    "\n",
    "# With a different initial guess, the solver can find the solution\n",
    "prog.SetInitialGuess(x, [0.1, 0.5])\n",
    "result = solver.Solve(prog)\n",
    "print(f\"Starting from x=[0.1, 0.5], the solver result is {result.get_solution_result()}\")\n",
    "print(f\"The found solution is x={result.GetSolution(x)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the example above, you could see that with a bad initial guess $x = [0, 0]$, the solver gets stuck and reports the problem being infeasible (the reason for getting stuck is that the gradient of the constraint function $cos(x[0]) + 2cos(x[0]-x[1])$ is zero at the initial guess $x=[0, 0]$, hence the gradient-based solver doesn't know how to move the decision variables. This phenomenon also appears when solving an inverse-kinematics problem with an initial pose at singularity, or solving a unit-length constraint $x^Tx=1$ with initial guess $x=0$); but by changing the initial guess, the solver can find a solution.\n",
    "\n",
    "Note that even if the solver (like SNOPT) has found a feasible solution during the optimization process, it could jump to an infeasible value in the next iteration. SNOPT doesn't guarantee to stay within the feasible region during the optimization process.\n",
    "\n",
    "Sometimes the problem is infeasible because the constraint is imposed incorrectly. To understand why the optimization fails, we provide some debugging tips. You could use these tips to diagonose the problematic constraint/initial guesses.\n",
    "\n",
    "## Name your costs/constraints/variables\n",
    "It often helps to print out some constraints/costs for diagonosis. In order to get a meaningful print out message, you could name the costs/constraints/variables.\n",
    "### 1. Name the variables\n",
    "When you create the variables through `NewContinuousVariables` (or `NewBinaryVariables`), you can pass in a string as the variable name. Here is an example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prog = mp.MathematicalProgram()\n",
    "x = prog.NewContinuousVariables(2, \"point\")\n",
    "print(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Name the constraint\n",
    "You could use `set_description()` function to name a constraint. Here is an example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prog = mp.MathematicalProgram()\n",
    "x = prog.NewContinuousVariables(2, \"point\")\n",
    "constraint = prog.AddConstraint(lambda z: [np.sum(z**2)], [1.], [1.], x)\n",
    "constraint.evaluator().set_description(\"unit-length constraint\")\n",
    "print(constraint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Name the cost\n",
    "Similarly you could use `set_description()` function to name a cost. Here is an example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prog = mp.MathematicalProgram()\n",
    "x = prog.NewContinuousVariables(2, \"point\")\n",
    "# Add the cost on the distance to (1, 2)\n",
    "cost1 = prog.AddCost(lambda z: np.sqrt((z[0]-1)**2 + (z[1]-2)**2), x)\n",
    "cost1.evaluator().set_description(\"distance to (1, 2)\")\n",
    "# Add the cost on the distance to (3, -1)\n",
    "cost2 = prog.AddCost(lambda z: np.sqrt((z[0]-3)**2 + (z[1] + 1)**2), x)\n",
    "cost2.evaluator().set_description(\"distance to (3, -1)\")\n",
    "print(f\"cost1: {cost1}\")\n",
    "print(f\"cost2: {cost2}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we will see in the next section, we can print out the infeasible constraints. By naming the variables/constraints/costs, the print out message becomes more meaningful.\n",
    "\n",
    "## Call GetInfeasibleConstraints()\n",
    "When `MathematicalProgram` is solved through a gradient-based solver (like SNOPT/IPOPT) and reports that the problem being infeasible, the solver returns the decision variable values where it gets stuck. You could call [MathematicalProgramResult::GetInfeasibleConstraints()](https://drake.mit.edu/pydrake/pydrake.solvers.mathematicalprogram.html?highlight=getinfeasibleconstraints#pydrake.solvers.mathematicalprogram.MathematicalProgramResult.GetInfeasibleConstraints) or [MathematicalProgramResult::GetInfeasibleConstraintNames()](https://drake.mit.edu/pydrake/pydrake.solvers.mathematicalprogram.html?highlight=getinfeasibleconstraints#pydrake.solvers.mathematicalprogram.MathematicalProgramResult.GetInfeasibleConstraintNames) to retrieve the constraint with large violations at that variable value. You could then diagonose the retrieved infeasible constraint and improve the constraint/initial guess accordingly.\n",
    "\n",
    "Here is an example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pydrake.autodiffutils import (\n",
    "    initializeAutoDiff,\n",
    "    autoDiffToGradientMatrix,\n",
    ")\n",
    "prog = mp.MathematicalProgram()\n",
    "x = prog.NewContinuousVariables(2, \"x\")\n",
    "\n",
    "# Add the constraint dist(x, 0) >= 1\n",
    "constraint1 = prog.AddConstraint(lambda z: [z.dot(z)], [1], [np.inf], x)\n",
    "constraint1.evaluator().set_description(\"outside unit circle\")\n",
    "\n",
    "# Add the constraint x[0]**2 + 4 * x[1]**2 <= 4\n",
    "constraint2 = prog.AddConstraint(lambda z: [z[0]**2 + 4 * z[1]**2], [0], [4], x)\n",
    "constraint2.evaluator().set_description(\"inside ellipsoid 1\")\n",
    "\n",
    "solver = IpoptSolver()\n",
    "prog.SetInitialGuess(x, [0, 0])\n",
    "result = solver.Solve(prog)\n",
    "print(\"Start from initial guess x = [0, 0]\")\n",
    "print(f\"optimization status: {result.get_solution_result()}\")\n",
    "infeasible_constraints = result.GetInfeasibleConstraints(prog)\n",
    "for c in infeasible_constraints:\n",
    "    print(f\"infeasible constraint: {c}\")\n",
    "x_stuck = result.GetSolution(x)\n",
    "print(f\"x_stuck={x_stuck.T}\")\n",
    "# Now evaluate the gradient of the constraint at x_stuck (where the solver gets stuck)\n",
    "print(f\"Gradient of the infeasible constraint at x_stuck: {autoDiffToGradientMatrix(infeasible_constraints[0].evaluator().Eval(initializeAutoDiff(x_stuck)))}\")\n",
    "\n",
    "# For a different initial state, the constraint that was infeasible now has non-zero gradient\n",
    "x_new = np.array([0.1, 0.2])\n",
    "print(f\"\\nStart from initial guess x_new = {x_new}\")\n",
    "print(f\"Gradient of the infeasible constraint at x_new: {autoDiffToGradientMatrix(infeasible_constraints[0].evaluator().Eval(initializeAutoDiff(x_new)))}\")\n",
    "prog.SetInitialGuess(x, x_new)\n",
    "# With this new initial guess, the solver will be able to find the solution\n",
    "result = solver.Solve(prog)\n",
    "print(f\"optimization status: {result.get_solution_result()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set solver verbosity level\n",
    "Many solvers can print out the progress in each iterations, and allow the users to set the verbosity levels. Here is a list of method to set the print out for each solver\n",
    "\n",
    "1. For Mosek, call [MosekSolver::set_stream_logging(True)](https://drake.mit.edu/pydrake/pydrake.solvers.mosek.html?highlight=moseksolver#pydrake.solvers.mosek.MosekSolver.set_stream_logging).\n",
    "2. For SNOPT, call [SetSolverOption(SnoptSolver().solver_id(), \"print file\", FILE_NAME)](https://drake.mit.edu/pydrake/pydrake.solvers.mathematicalprogram.html?highlight=setsolveroption#pydrake.solvers.mathematicalprogram.MathematicalProgram.SetSolverOption) to print the result to a file `FILE_NAME`.\n",
    "3. For IPOPT, call [SetSolverOption(IpoptSolver().solver_id(), options, options_val)](https://drake.mit.edu/pydrake/pydrake.solvers.mathematicalprogram.html?highlight=setsolveroption#pydrake.solvers.mathematicalprogram.MathematicalProgram.SetSolverOption) to set the options relating to the solver output. You could refer to the [IPOPT documentation](https://coin-or.github.io/Ipopt/OPTIONS.html) on the supported options. For example, to adjust print level, call `prog.SetSolverOptions(IpoptSolver.id(), \"print_level\", 5)` to set the print_level to 5. \n",
    "4. For Gurobi, call `prog.SetSolverOption(GurobiSolver().solver_id(), \"OutputFlag\", True)`.\n",
    "5. For SCS, call `prog.SetSolverOption(ScsSolver().solver_id(), \"verbose\", 1)`.\n",
    "6. For OSQP, call `prog.SetSolverOption(OsqpSolver().solver_id(), \"verbose\", 1)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Demonstrate setting solver verbosity level.\n",
    "prog = mp.MathematicalProgram()\n",
    "x = prog.NewContinuousVariables(2)\n",
    "# Add the constraint 1 <= squared_norm(x) <= 2\n",
    "prog.AddConstraint(lambda z: [z.dot(z)], [1], [2], x)\n",
    "# Add the constraint 2 <= x[0]**2 + 4 * x[0]*x[1] + 4*x[1]**2 + 2 * x[0] <= 5\n",
    "prog.AddConstraint(lambda z: [z[0] ** 2 + 4 * z[0] * z[1] + 4 * z[1]**2 + 2 * z[0]], [2], [5], x)\n",
    "\n",
    "ipopt_solver = IpoptSolver()\n",
    "# Set print level to 5. \n",
    "prog.SetSolverOption(ipopt_solver.solver_id(), \"print_level\", 5)\n",
    "prog.SetInitialGuess(x, [0, 0])\n",
    "# If you are running this code in jupyter notebook, the solver prints out the\n",
    "# progress in the console from which this notebook is launched.\n",
    "result = ipopt_solver.Solve(prog)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add Callback\n",
    "Some solvers support adding a callback function, which is executed in each iteration of the optimization process. You could use this callback to visualize the progress."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the solver progress in each iteration through a callback\n",
    "# Find the closest point on a curve to a desired point.\n",
    "from pydrake.solvers.mathematicalprogram import Solve\n",
    "\n",
    "fig = plt.figure()\n",
    "curve_x = np.linspace(1, 10, 100)\n",
    "ax = plt.gca()\n",
    "ax.plot(curve_x, 9./curve_x)\n",
    "ax.plot(-curve_x, -9./curve_x)\n",
    "ax.plot(0, 0, 'o')\n",
    "x_init = [4., 5.]\n",
    "point_x, = ax.plot(x_init[0], x_init[1], 'x')\n",
    "ax.axis('equal')\n",
    "\n",
    "def visualization_callback(x):\n",
    "    global iter_count\n",
    "    point_x.set_xdata(x[0])\n",
    "    point_x.set_ydata(x[1])\n",
    "    ax.set_title(f\"iteration {iter_count}\")\n",
    "    fig.canvas.draw()\n",
    "    fig.canvas.flush_events()\n",
    "    # Also update the iter_count variable in the callback.\n",
    "    # This shows we can do more than just visualization in\n",
    "    # callback.\n",
    "    iter_count += 1\n",
    "    plt.pause(0.1)\n",
    "    \n",
    "iter_count = 0\n",
    "prog = mp.MathematicalProgram()\n",
    "x = prog.NewContinuousVariables(2)\n",
    "prog.AddConstraint(x[0] * x[1] == 9)\n",
    "prog.AddCost(x[0]**2 + x[1]**2)\n",
    "prog.AddVisualizationCallback(visualization_callback, x)\n",
    "result = Solve(prog, x_init)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use EvalBinding\n",
    "For each individual constraint/cost, you could call [EvalBinding(binding, x)](https://drake.mit.edu/pydrake/pydrake.solvers.mathematicalprogram.html?highlight=evalbinding#pydrake.solvers.mathematicalprogram.MathematicalProgram.EvalBinding) to evaluate the constraint/cost `binding` with the program decision variables set to `x`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## Demonstrate EvalBinding function\n",
    "prog = mp.MathematicalProgram()\n",
    "p1 = prog.NewContinuousVariables(2, \"p1\")\n",
    "p2 = prog.NewContinuousVariables(2, \"p2\")\n",
    "\n",
    "# Add the constraint that p1 is in an ellipsoid (p1(0)-1)**2 + 4*(p1(1)-2)**2 <= 1\n",
    "constraint1 = prog.AddConstraint(lambda z: [(z[0]-1)**2 + 4 * (z[1]-2)**2], [0], [1], p1)\n",
    "# Add the constraint that p2 is in an ellipsoid (p2(0) + 2)**2 + 0.25*(p2(1)+ 1)**2) <= 1\n",
    "constraint2 = prog.AddConstraint(lambda z: [(z[0]+2)**2 + 0.25*(z[1]+1)**2], [0], [1], p2)\n",
    "# Add a cost to minimize the distance between p1 and p2\n",
    "cost = prog.AddCost((p1-p2).dot(p1-p2))\n",
    "\n",
    "# Evaluate the constraint and cost at a guess p1=[0, 1], p2 = [-1, -4]\n",
    "p1_val = [0, 1]\n",
    "p2_val = [-1, -4]\n",
    "prog.SetInitialGuess(p1, p1_val)\n",
    "prog.SetInitialGuess(p2, p2_val)\n",
    "print(f\"constraint 1 evaluated at p1={p1_val}, p2={p2_val} is {prog.EvalBinding(constraint1, prog.initial_guess())}\")\n",
    "print(f\"constraint 2 evaluated at p1={p1_val}, p2={p2_val} is {prog.EvalBinding(constraint2, prog.initial_guess())}\")\n",
    "print(f\"cost evaluated at p1={p1_val}, p2={p2_val} is {prog.EvalBinding(cost, prog.initial_guess())}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Removing/relaxing constraints\n",
    "When the solver reports the problem being infeasible, you could remove or relax the infeasible constraint(s), and solve the problem again. Removing the constraint is trivial, you just need to comment out the line that added the constraint in the first place. To relax the constraint bound, you can use the function [UpdateLowerBound()](https://drake.mit.edu/pydrake/pydrake.solvers.mathematicalprogram.html?highlight=updateupperbound#pydrake.solvers.mathematicalprogram.PyFunctionConstraint.UpdateLowerBound) or [UpdateUpperBound()](https://drake.mit.edu/pydrake/pydrake.solvers.mathematicalprogram.html?highlight=updateupperbound#pydrake.solvers.mathematicalprogram.PyFunctionConstraint.UpdateUpperBound). Here is a quick example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Relaxing the constraint\n",
    "prog = mp.MathematicalProgram()\n",
    "x = prog.NewContinuousVariables(2)\n",
    "\n",
    "# Add the constraint x^T * x <= 1\n",
    "constraint1 = prog.AddConstraint(lambda z: [z.dot(z)], [0], [1], x)\n",
    "constraint1.evaluator().set_description(\"inside unit circle\")\n",
    "\n",
    "# Add the constraint norm(x-[3, 0]) <= 1\n",
    "constraint2 = prog.AddConstraint(lambda z: [np.sum((z - np.array([3, 0]))**2)], [0], [1], x)\n",
    "constraint2.evaluator().set_description(\"distance to [3, 0] less than 1\")\n",
    "\n",
    "prog.SetInitialGuess(x, [1, 0])\n",
    "solver = IpoptSolver()\n",
    "result = solver.Solve(prog)\n",
    "print(f\"For the original problem, the solver status is {result.get_solution_result()}\")\n",
    "print(f\"x is stuck at {result.GetSolution(x)}\")\n",
    "# Now get the infeasible constraint\n",
    "infeasible_constraints = result.GetInfeasibleConstraints(prog)\n",
    "for c in infeasible_constraints:\n",
    "    print(f\"infeasible constraint: {c}\")\n",
    "\n",
    "# Now update the upper bound of the first infeasible constraint\n",
    "infeasible_constraints[0].evaluator().UpdateUpperBound([4])\n",
    "# I also update the description of the constraint. Without updating the description, the\n",
    "# problem still solves fine, but it would be confusing if you print out this constraint.\n",
    "infeasible_constraints[0].evaluator().set_description(\"inside a circle with radius=2\")\n",
    "result = solver.Solve(prog)\n",
    "print(f\"For the relaxed problem, the solver status is {result.get_solution_result()}\")\n",
    "print(f\"Solution is x = {result.GetSolution(x)}\")"
   ]
  }
 ],
 "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
}