{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Full-gradient constrained optimization dynamics.\nSets up simple 2-d problems on Linf balls to visualize dynamics of various\nconstrained optimization algorithms.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from functools import partial\nimport matplotlib.pyplot as plt\nimport torch\n\nfrom chop.constraints import LinfBall\nfrom chop.optim import minimize_frank_wolfe, minimize_pgd, minimize_pgd_madry, minimize_three_split\nfrom chop import utils \n\ntorch.random.manual_seed(0)\n\n\ndef setup_problem(make_nonconvex=False):\n    alpha = 1.\n    x_star = torch.tensor([alpha, alpha/2]).unsqueeze(0)\n    x_0 = torch.zeros_like(x_star)\n\n    @utils.closure\n    def loss_func(x):\n        val = .5 * ((x - x_star) ** 2).sum()\n        if make_nonconvex:\n            val += .1 * torch.sin(50 * torch.norm(x, p=1) + .1)\n        return val\n\n    constraint = LinfBall(alpha)\n\n    return x_0, x_star, loss_func, constraint\n\n\ndef log(kwargs, iterates, losses):\n    x= kwargs['x'].squeeze().data\n    iterates.append(x)\n    val = kwargs['closure'](x, return_jac=False).data\n    losses.append(val)\n\n\n\nif __name__ == \"__main__\":\n\n    x_0, x_star, loss_func, constraint = setup_problem(make_nonconvex=True)\n    iterations = 10\n\n    iterates_pgd = [x_0.squeeze().data]\n    iterates_pgd_madry = [x_0.squeeze().data]\n    iterates_splitting = [x_0.squeeze().data]\n    iterates_fw = [x_0.squeeze().data]\n\n    losses_pgd = [loss_func(x_0, return_jac=False).data]\n    losses_pgd_madry = [loss_func(x_0, return_jac=False).data]\n    losses_splitting = [loss_func(x_0, return_jac=False).data]\n    losses_fw = [loss_func(x_0, return_jac=False).data]\n\n    log_pgd = partial(log, iterates=iterates_pgd, losses=losses_pgd)\n    log_pgd_madry = partial(log, iterates=iterates_pgd_madry, losses=losses_pgd_madry)\n    log_splitting = partial(log, iterates=iterates_splitting, losses=losses_splitting)\n    log_fw = partial(log, iterates=iterates_fw, losses=losses_fw)\n\n    sol_pgd = minimize_pgd(loss_func, x_0, constraint.prox,\n                           max_iter=iterations,\n                           callback=log_pgd)\n\n    sol_pgd_madry = minimize_pgd_madry(loss_func, x_0, constraint.prox,\n                                       constraint.lmo,\n                                       step=2. / iterations,\n                                       max_iter=iterations,\n                                       callback=log_pgd_madry)\n\n    sol_splitting = minimize_three_split(loss_func, x_0, prox1=constraint.prox, \n                                         max_iter=iterations, callback=log_splitting)\n\n    sol_fw = minimize_frank_wolfe(loss_func, x_0, constraint.lmo, callback=log_fw,\n                                  max_iter=iterations)\n\n    fig, ax = plt.subplots()\n    ax.plot(losses_pgd, label=\"PGD\")\n    ax.plot(losses_pgd_madry, label=\"PGD Madry\")\n    ax.plot(losses_splitting, label=\"Operator Splitting\")\n    ax.plot(losses_fw, label=\"Frank-Wolfe\")\n    fig.legend()\n    plt.show()\n\n    fig, ax = plt.subplots(ncols=2, nrows=2, sharex=True, sharey=True)\n    ax = ax.flatten()\n    ax[0].plot(*zip(*iterates_pgd), '-o', label=\"PGD\", alpha=.6)\n    ax[0].set_xlim(-1, 1)\n    ax[0].set_ylim(-1, 1)\n    ax[0].legend()\n\n    ax[1].plot(*zip(*iterates_pgd_madry), '-o', label=\"PGD Madry\", alpha=.6)\n    ax[1].set_xlim(-1, 1)\n    ax[1].set_ylim(-1, 1)\n    ax[1].legend()\n\n    ax[2].plot(*zip(*iterates_splitting), '-o', label=\"Operator Splitting\", alpha=.6)\n    ax[2].set_xlim(-1, 1)\n    ax[2].set_ylim(-1, 1)\n    ax[2].legend()\n\n    ax[3].plot(*zip(*iterates_fw), '-o', label=\"Frank-Wolfe\", alpha=.6)\n    ax[3].set_xlim(-1, 1)\n    ax[3].set_ylim(-1, 1)\n    ax[3].legend()\n\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.8.5"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}