{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using the Augmented Lagrangian function.\n",
    "\n",
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/cooper-org/cooper/blob/main/docs/source/notebooks/plot_augmented_lagrangian.ipynb)\n",
    "\n",
    "\n",
    "\n",
    "This tutorial demonstrates how to use the {py:class}`~cooper.formulations.AugmentedLagrangian` formulation to solve constrained optimization problems in Cooper. We illustrate its usage and advantages over the {py:class}`~cooper.formulations.QuadraticPenalty` formulation with a simple 2D example from {cite:p}`nocedal2006NumericalOptimization`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "%pip install cooper-optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "\n",
    "import cooper\n",
    "\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem Formulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We consider the following constrained optimization problem (Problem 17.3 from {cite:t}`nocedal2006NumericalOptimization`) with a single equality constraint:\n",
    "\n",
    "$$\n",
    "\\min_{\\boldsymbol{x} \\in \\mathbb{R}^2} f(\\boldsymbol{x}) = x_1 + x_2 \\quad \\text{s.t.} \\quad x_1^2 + x_2^2 = 2.\n",
    "$$\n",
    "\n",
    "The unique solution is $\\boldsymbol{x}^* = (-1, -1)$.\n",
    "\n",
    "The code below implements this constrained minimization problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Problem2D(cooper.ConstrainedMinimizationProblem):\n",
    "    def __init__(self, formulation_type):\n",
    "        super().__init__()\n",
    "\n",
    "        if not formulation_type.expects_multiplier:\n",
    "            self.multiplier = None\n",
    "        else:\n",
    "            self.multiplier = cooper.multipliers.DenseMultiplier(num_constraints=1, device=DEVICE)\n",
    "\n",
    "        if not formulation_type.expects_penalty_coefficient:\n",
    "            self.penalty = None\n",
    "        else:\n",
    "            self.penalty = cooper.penalty_coefficients.DensePenaltyCoefficient(\n",
    "                init=torch.tensor(0.0, device=DEVICE),\n",
    "            )\n",
    "\n",
    "        self.constraint = cooper.Constraint(\n",
    "            constraint_type=cooper.ConstraintType.EQUALITY,\n",
    "            formulation_type=formulation_type,\n",
    "            multiplier=self.multiplier,\n",
    "            penalty_coefficient=self.penalty,\n",
    "        )\n",
    "\n",
    "    def compute_cmp_state(self, x):\n",
    "        objective = torch.sum(x)  # x1 + x2\n",
    "        constraint = torch.sum(x**2) - 2  # x1^2 + x2^2 - 2 = 0\n",
    "\n",
    "        constraint_state = cooper.ConstraintState(violation=constraint)\n",
    "\n",
    "        return cooper.CMPState(loss=objective, observed_constraints={self.constraint: constraint_state})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Augmented Lagrangian function associated with this problem is:\n",
    "\n",
    "$$\n",
    "\\mathcal{L}_{c}(\\boldsymbol{x}, \\mu) = x_1 + x_2 + \\mu (x_1^2 + x_2^2 - 2) + \\frac{c}{2} (x_1^2 + x_2^2 - 2)^2,\n",
    "$$\n",
    "\n",
    "where $\\mu$ is the Lagrange multiplier associated with the equality constraint, and $c$ is the penalty parameter.\n",
    "\n",
    "We will also consider the Quadratic Penalty function associated with this problem:\n",
    "\n",
    "$$\n",
    "\\mathcal{P}_{c}(\\boldsymbol{x}) = x_1 + x_2 + \\frac{c}{2} (x_1^2 + x_2^2 - 2)^2.\n",
    "$$\n",
    "\n",
    "Both of these formulations are instantiated in the following code block:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solving the Problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_STAR = torch.tensor([-1.0, -1.0], device=DEVICE)\n",
    "\n",
    "\n",
    "def train(problem, x, primal_lr, momentum, dual_lr, penalty_increment, n_steps):\n",
    "    has_dual_variables = problem.multiplier is not None\n",
    "\n",
    "    primal_optimizer = torch.optim.SGD([x], lr=primal_lr, momentum=momentum)\n",
    "\n",
    "    if has_dual_variables:\n",
    "        dual_optimizer = torch.optim.SGD(problem.dual_parameters(), lr=dual_lr, maximize=True)\n",
    "\n",
    "        constrained_optimizer = cooper.optim.SimultaneousOptimizer(\n",
    "            cmp=problem,\n",
    "            primal_optimizers=primal_optimizer,\n",
    "            dual_optimizers=dual_optimizer,\n",
    "        )\n",
    "    else:\n",
    "        # Formulations without dual variables, such as the Quadratic Penalty\n",
    "        # formulation, do not require a dual optimizer\n",
    "        constrained_optimizer = cooper.optim.UnconstrainedOptimizer(\n",
    "            cmp=problem,\n",
    "            primal_optimizers=primal_optimizer,\n",
    "        )\n",
    "\n",
    "    # Increase the penalty coefficient by `increment` if the constraint is violate by more\n",
    "    # than `violation_tolerance`\n",
    "    penalty_scheduler = cooper.penalty_coefficients.AdditivePenaltyCoefficientUpdater(\n",
    "        increment=penalty_increment,\n",
    "        violation_tolerance=1e-3,\n",
    "    )\n",
    "\n",
    "    dist_2_x_star, multipliers, penalty_coefficients = [], [], []\n",
    "\n",
    "    for _ in range(n_steps):\n",
    "        roll_out = constrained_optimizer.roll(compute_cmp_state_kwargs={\"x\": x})\n",
    "\n",
    "        # Update the penalty coefficient\n",
    "        constraint_state = roll_out.cmp_state.observed_constraints[problem.constraint]\n",
    "        penalty_scheduler.update_penalty_coefficient_(problem.constraint, constraint_state)\n",
    "\n",
    "        multiplier_value = problem.multiplier.weight.item() if has_dual_variables else None\n",
    "        penalty_coefficient_value = problem.constraint.penalty_coefficient().item()\n",
    "\n",
    "        dist_2_x_star.append(torch.norm(x - X_STAR).item())\n",
    "        multipliers.append(multiplier_value)\n",
    "        penalty_coefficients.append(penalty_coefficient_value)\n",
    "\n",
    "    return x, dist_2_x_star, multipliers, penalty_coefficients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "primal_lr = 1e-2\n",
    "momentum = 0.9\n",
    "penalty_increment = 1e-2\n",
    "n_steps = 5_000\n",
    "\n",
    "# Only required for Augmented Lagrangian formulation\n",
    "dual_lr = 1e-2\n",
    "\n",
    "x_qp = torch.tensor([1.0, 1.0], device=DEVICE, requires_grad=True)\n",
    "x_al = torch.tensor([1.0, 1.0], device=DEVICE, requires_grad=True)\n",
    "\n",
    "QP_Problem = Problem2D(cooper.formulations.QuadraticPenalty)\n",
    "AL_Problem = Problem2D(cooper.formulations.AugmentedLagrangian)\n",
    "\n",
    "x_qp_final, dist_qp, qp_multipliers, qp_penalty_coefficients = train(\n",
    "    QP_Problem, x_qp, primal_lr, momentum, dual_lr, penalty_increment, n_steps\n",
    ")\n",
    "x_al_final, dist_al, al_multipliers, al_penalty_coefficients = train(\n",
    "    AL_Problem, x_al, primal_lr, momentum, dual_lr, penalty_increment, n_steps\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The plot below illustrates that the Quadratic Penalty function requires a significantly larger penalty coefficient $c$ to enforce feasibility compared to the Augmented Lagrangian function. This large penalty leads to numerical instability, preventing the algorithm's convergence.\n",
    "\n",
    "In contrast, the Augmented Lagrangian function achieves convergence with a finite penalty coefficient, avoiding these stability issues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 3))\n",
    "\n",
    "# Plot distance to optimal solution\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(dist_qp, label=\"Quadratic Penalty\", linewidth=3)\n",
    "plt.plot(dist_al, label=\"Augmented Lagrangian\", linewidth=3, linestyle=\"dashed\")\n",
    "plt.yscale(\"log\")\n",
    "plt.xlabel(\"Iteration\", fontsize=12)\n",
    "plt.ylabel(\"Distance to Optimal Solution\", fontsize=12)\n",
    "plt.title(\"Convergence Comparison\", fontsize=14)\n",
    "plt.legend()\n",
    "plt.grid(True, linestyle=\"--\", alpha=0.6)\n",
    "\n",
    "# Plot penalty coefficients\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(qp_penalty_coefficients, label=\"Quadratic Penalty\", linewidth=3)\n",
    "plt.plot(al_penalty_coefficients, label=\"Augmented Lagrangian\", linewidth=3, linestyle=\"dashed\")\n",
    "plt.yscale(\"log\")\n",
    "plt.xlabel(\"Iteration\", fontsize=12)\n",
    "plt.ylabel(\"Penalty Coefficient\", fontsize=12)\n",
    "plt.title(\"Penalty Coefficient Evolution\", fontsize=14)\n",
    "plt.legend()\n",
    "plt.grid(True, linestyle=\"--\", alpha=0.6)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md:myst"
  },
  "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}