{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Stochastic constrained optimization dynamics.\nSets up simple 2-d problems on Linf balls to visualize dynamics of various\nstochastic constrained optimization algorithms.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\nimport numpy as np\nimport torch\n\nfrom chop.constraints import LinfBall\nfrom chop.stochastic import PGD, PGDMadry, FrankWolfe\n\ntorch.random.manual_seed(0)\n\nOPTIMIZER_CLASSES = [PGD, PGDMadry, FrankWolfe]\n\n\ndef setup_problem(make_nonconvex=False):\n    radius = 1.\n    x_star = torch.tensor([radius, radius/2])\n    x_0 = torch.zeros_like(x_star)\n\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(radius)\n\n    return x_0, x_star, loss_func, constraint\n\n\ndef optimize(x_0, loss_func, constraint, optimizer_class, iterations=10):\n    x = x_0.detach().clone()\n    x.requires_grad = True\n    # Use Madry's heuristic for step size\n    lr = {\n        FrankWolfe: 2.5 / iterations,\n        PGD: 2.5 * constraint.alpha / iterations * 2.,\n        PGDMadry: 2.5 / iterations\n    }\n    optimizer = optimizer_class([x], constraint, lr=lr[optimizer_class])\n    iterates = [x.data.numpy().copy()]\n    losses = []\n\n    for _ in range(iterations):\n        optimizer.zero_grad()\n        loss = loss_func(x)\n        loss.backward()\n        optimizer.step()\n        losses.append(loss.item())\n        iterates.append(x.data.numpy().copy())\n\n    loss = loss_func(x)\n    losses.append(loss.item())\n    return losses, iterates\n\n\nif __name__ == \"__main__\":\n\n    x_0, x_star, loss_func, constraint = setup_problem(make_nonconvex=False)\n    iterations = 10\n    losses_all = {}\n    iterates_all = {}\n    for opt_class in OPTIMIZER_CLASSES:\n        losses_, iterates_ = optimize(x_0,\n                                      loss_func,\n                                      constraint,\n                                      opt_class,\n                                      iterations)\n        losses_all[opt_class.name] = losses_\n        iterates_all[opt_class.name] = iterates_\n    # print(losses)\n    fig, ax = plt.subplots()\n    for opt_class in OPTIMIZER_CLASSES:\n        ax.plot(np.arange(iterations + 1), losses_all[opt_class.name],\n                label=opt_class.name)\n    fig.legend()\n    plt.show()\n\n    fig, axes = plt.subplots(nrows=2, ncols=2, sharex=True, sharey=True)\n    for ax, opt_class in zip(axes.reshape(-1), OPTIMIZER_CLASSES):\n        ax.plot(*zip(*iterates_all[opt_class.name]), '-o', label=opt_class.name, alpha=.6)\n        ax.set_xlim(-1, 1)\n        ax.set_ylim(-1, 1)\n        ax.legend(loc='lower left')\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
}