{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MathematicalProgram debugging tips\n",
    "For instructions on how to run these tutorial notebooks, please see the [index](./index.ipynb).\n"
   ]
  },
  {
   "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 successfully. 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": [
    "from IPython.display import Markdown, display\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from matplotlib import cm\n",
    "from pydrake.solvers import MathematicalProgram\n",
    "from pydrake.solvers 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 = 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",
    "## Print a summary of the MathematicalProgram\n",
    "Especially for small problems, it can be extremely helpful to display the MathematicalProgram as a string.  This will give you a list of decision variables, costs, and constraints that have been added to the program.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A sample (quadratic) program\n",
    "prog = MathematicalProgram()\n",
    "x = prog.NewContinuousVariables(3, \"x\")\n",
    "prog.AddQuadraticCost(x[0] * x[0] + 2 * x[0] + 3)\n",
    "prog.Add2NormSquaredCost(A = [[1, 3], [2, 4]], b=[1, 4], vars=[x[1], x[2]])\n",
    "prog.AddLinearEqualityConstraint(x[0] + 2*x[1] == 5)\n",
    "prog.AddLinearConstraint(x[0] + 4 *x[1] <= 10)\n",
    "prog.AddBoundingBoxConstraint(-1, 10, x)\n",
    "\n",
    "# Now print a summary:\n",
    "print(prog)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also display the program in LaTeX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(Markdown(prog.ToLatex()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 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 = 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 = 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 = 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.html#pydrake.solvers.MathematicalProgramResult.GetInfeasibleConstraints) or [MathematicalProgramResult::GetInfeasibleConstraintNames()](https://drake.mit.edu/pydrake/pydrake.solvers.html#pydrake.solvers.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",
    "    ExtractGradient,\n",
    ")\n",
    "prog = 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: {ExtractGradient(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: {ExtractGradient(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": [
    "## Enable solver verbosity\n",
    "Many solvers can print out the progress in each iteration. To enable printouts, enable either [kPrintToConsole](https://drake.mit.edu/pydrake/pydrake.solvers.html#pydrake.solvers.CommonSolverOption) or [kPrintFileName](https://drake.mit.edu/pydrake/pydrake.solvers.html#pydrake.solvers.CommonSolverOption) in the [SolverOptions](https://drake.mit.edu/pydrake/pydrake.solvers.html#pydrake.solvers.SolverOptions).\n",
    "\n",
    "Note that when running code in a Jupyter notebook, solvers print to the console from which the notebook is launched, rather than the notebook window. For online notebooks (e.g., Deepnote), this means that the console output will not be visible. In that case, printing to a file is better."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.solvers import CommonSolverOption, SolverOptions\n",
    "\n",
    "# Create a simple program consisting of\n",
    "#  constraint: 1 <= squared_norm(x) <= 2.\n",
    "#  constraint: 2 <= x[0]**2 + 4 * x[0]*x[1] + 4*x[1]**2 + 2 * x[0] <= 5.\n",
    "prog = MathematicalProgram()\n",
    "x = prog.NewContinuousVariables(2)\n",
    "prog.AddConstraint(lambda z: [z.dot(z)], [1], [2], x)\n",
    "prog.AddConstraint(lambda z: [z[0] ** 2 + 4 * z[0] * z[1] + 4 * z[1]**2 + 2 * z[0]], [2], [5], x)\n",
    "prog.SetInitialGuess(x, [0, 0])\n",
    "\n",
    "# Solve with printing.\n",
    "ipopt_solver = IpoptSolver()\n",
    "filename = \"/tmp/debug.txt\"\n",
    "solver_options = SolverOptions()\n",
    "solver_options.SetOption(CommonSolverOption.kPrintFileName, filename)\n",
    "result = ipopt_solver.Solve(prog, solver_options=solver_options)\n",
    "with open(filename) as f:\n",
    "    print(f.read())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some solvers offer more fine-tuning for the progress output.  For example, the [IPOPT options](https://coin-or.github.io/Ipopt/OPTIONS.html) offer an IPOPT-specific `print_level` for console output or `file_print_level` for file output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Enable a slightly more verbose level of printing this time.\n",
    "# The options can be attached to the program, instead of passed to Solve().\n",
    "solver_options.SetOption(IpoptSolver.id(), \"print_level\", 5)\n",
    "prog.SetSolverOptions(solver_options)\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 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",
    "ax.plot(x_init[0], x_init[1], 'x')\n",
    "ax.axis('equal')\n",
    "\n",
    "def visualization_callback(x):\n",
    "    ax.plot(x[0], x[1], 'x', color='gray')\n",
    "    \n",
    "prog = 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.html#pydrake.solvers.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 = 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.html#pydrake.solvers.PyFunctionConstraint.UpdateLowerBound) or [UpdateUpperBound()](https://drake.mit.edu/pydrake/pydrake.solvers.html#pydrake.solvers.PyFunctionConstraint.UpdateUpperBound). Here is a quick example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Relaxing the constraint\n",
    "prog = 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
}