I_ACVI_HBG.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_GkYPL_Y3C1l"
      },
      "source": [
        "# High dimensional bilinear game on the simplex\n",
        "\n",
        "Results on the following constrained bilinear game, where $\\mathbf{x}_1, \\mathbf{x}_2 \\in \\mathbb{R}^{500}$, and $\\eta \\in (0,1)$:\n",
        "\n",
        "$ \n",
        "min_{\\mathbf{x}_1} max_{\\mathbf{x}_2} \\quad \\eta \\cdot \\mathbf{x}_1^\\intercal + (1-\\eta) \\cdot \\mathbf{x}_1^\\intercal \\mathbf{x}_2 - \\eta \\cdot \\mathbf{x}_2^\\intercal \\mathbf{x}_2\n",
        "$\n",
        "\n",
        "$\n",
        "\\hspace{2em} \\text{s.t.} \\qquad \\mathbf{x} \\geq \\mathbf{0}\n",
        "$\n",
        "\n",
        "$\n",
        "\\hspace{5em} \\mathbf{e}^\\intercal \\mathbf{x}_1 = 1\n",
        "$\n",
        "\n",
        "$\n",
        "\\hspace{5em} \\mathbf{e}^\\intercal \\mathbf{x}_2 = 1\n",
        "$"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Imports"
      ],
      "metadata": {
        "id": "QT9jrXv4YDGy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import time\n",
        "import scipy.optimize as opt\n",
        "from scipy import linalg\n",
        "import matplotlib.pyplot as plt\n",
        "from cvxopt import matrix,solvers\n",
        "solvers.options['show_progress'] = False"
      ],
      "metadata": {
        "id": "PXe_GVJ9YEnp"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E52xV0sp-8kn"
      },
      "source": [
        "# Experiments on conditioning"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## I-ACVI"
      ],
      "metadata": {
        "id": "6mWScvLbWVPN"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KXSb6ourOxv8"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import time \n",
        "\n",
        "eps = .02  # target relative error\n",
        "dim = 500  # dim(x1) == dim(x2) == dim\n",
        "\n",
        "for conditioning in range(1,11):\n",
        "\n",
        "  beta, mu, delta, K, lr, T = .5, 1e-5, .5, 50, .003, 200  \n",
        "  if conditioning in [1,2,3]:\n",
        "    l = 20\n",
        "  elif conditioning in [4,5,6]:\n",
        "    l = 50\n",
        "  else:\n",
        "    l = 100\n",
        "\n",
        "  # Building HBG-v2 matrices\n",
        "  eigs = np.linspace(1, conditioning, dim)\n",
        "  A1 = np.concatenate((np.zeros((dim, dim)), np.diag(eigs)), axis=1)\n",
        "  A2 = np.concatenate((-np.diag(eigs), np.zeros((dim, dim))), axis=1)\n",
        "  A  = np.concatenate((A1, A2), axis=0)\n",
        "  x_opt = 1/eigs\n",
        "  x_opt /= x_opt.sum()\n",
        "  x_opt = np.concatenate((x_opt, x_opt), axis=0).reshape(-1, 1)\n",
        "\n",
        "  # Build projection matrix Pc\n",
        "  temp1 = np.concatenate((np.ones((1,dim)), np.zeros((1,dim))), axis=1)\n",
        "  temp2 = np.concatenate((np.zeros((1,dim)), np.ones((1,dim))), axis=1)\n",
        "  C = np.concatenate((temp1,temp2), axis=0)\n",
        "  d = np.ones((2,1))\n",
        "  temp = np.linalg.inv(np.dot(C, C.T))\n",
        "  temp = np.dot(C.T, temp)\n",
        "  dc = np.dot(temp, d)\n",
        "  Pc = np.identity(2*dim) - np.dot(temp,C)\n",
        "\n",
        "  np.random.seed(0)\n",
        "\n",
        "  # Initialize players\n",
        "  init = np.random.rand(2*dim, 1)\n",
        "  init[:dim] = init[:dim] / np.sum(init[:dim]) # ensuring it is part of the simplex\n",
        "  init[dim:] = init[dim:] / np.sum(init[dim:])  \n",
        "  z_x = np.copy(init)\n",
        "  z_y = np.copy(init)\n",
        "  z_lmd = np.zeros(init.shape)\n",
        "\n",
        "  finished, cnt, t0 = False, 0, time.time()\n",
        "    \n",
        "  for _ in range(T):\n",
        "    mu *= delta\n",
        "    \n",
        "    for _ in range(K):\n",
        "      cnt += 1\n",
        "      # Solve approximately the X problem (line 8 of algorithm)\n",
        "      for _ in range(l):\n",
        "        g = z_x + 1/beta * np.dot(Pc, np.dot(A,z_x)) - np.dot(Pc, z_y) + 1/beta * np.dot(Pc, z_lmd) - dc\n",
        "        z_x -= lr * g\n",
        "\n",
        "      if np.linalg.norm(z_x-x_opt)/np.linalg.norm(x_opt) <= eps:\n",
        "        finished = True\n",
        "        print(f\"Reached a relative error of {eps} after {cnt} iterations in {time.time()-t0:.2f} sec.\")\n",
        "        break\n",
        "\n",
        "      # Solve approximately the Y problem (line 9 of algorithm)\n",
        "      for _ in range(l): \n",
        "        assert all(z_y > 0) # ensuring the log terms are positive\n",
        "        g = - mu * 1/z_y + beta*(z_y - z_x - z_lmd/beta)\n",
        "        z_y -= lr * g\n",
        "        \n",
        "      # Update lambdas (line 10 of algorithm)\n",
        "      z_lmd += beta*(z_x-z_y)\n",
        "\n",
        "    if finished:\n",
        "      break"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## EG"
      ],
      "metadata": {
        "id": "Bb9VgcoKWiuv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import time \n",
        "from cvxopt import matrix,solvers\n",
        "solvers.options['show_progress'] = False\n",
        "\n",
        "eps = .02  # target relative error\n",
        "dim = 500  # dim(x1) == dim(x2) == dim\n",
        "max_time = 700\n",
        "\n",
        "lr = 0.3\n",
        "\n",
        "for conditioning in range(1,11):\n",
        "\n",
        "  lr *= 0.9   \n",
        "  T = 10000    \n",
        "\n",
        "  # Building HBG matrices\n",
        "  eta = 0.0\n",
        "  eigs = np.linspace(1, conditioning, dim)\n",
        "  A1 = np.concatenate((np.zeros((dim, dim)), np.diag(eigs)), axis=1)\n",
        "  A2 = np.concatenate((-np.diag(eigs), np.zeros((dim, dim))), axis=1)\n",
        "  A  = np.concatenate((A1, A2), axis=0)\n",
        "  x_opt = 1/eigs\n",
        "  x_opt /= x_opt.sum()\n",
        "  x_opt = np.concatenate((x_opt, x_opt), axis=0).reshape(-1, 1)\n",
        "\n",
        "  # Build projection matrix Pc\n",
        "  temp1 = np.concatenate((np.ones((1,dim)), np.zeros((1,dim))), axis=1)\n",
        "  temp2 = np.concatenate((np.zeros((1,dim)), np.ones((1,dim))), axis=1)\n",
        "  C = np.concatenate((temp1,temp2), axis=0)\n",
        "  d = np.ones((2,1))\n",
        "  temp = np.linalg.inv(np.dot(C, C.T))\n",
        "  temp = np.dot(C.T, temp)\n",
        "  dc = np.dot(temp, d)\n",
        "  Pc = np.identity(2*dim) - np.dot(temp,C)\n",
        "\n",
        "  G = matrix(-np.identity(2*dim), tc='d')\n",
        "  h = matrix(np.zeros(((2*dim,1))), tc='d')\n",
        "  AA = matrix(C, tc='d')\n",
        "  b = matrix(d, tc='d')\n",
        "\n",
        "  def proj(z):\n",
        "    P = matrix(np.identity(2*dim), tc='d')\n",
        "    q = matrix(-z, tc='d')\n",
        "    res = solvers.qp(P, q, G, h, AA, b)\n",
        "    return res['x']\n",
        "\n",
        "  np.random.seed(0)\n",
        "\n",
        "  # Initialize players\n",
        "  init = np.random.rand(2*dim, 1)\n",
        "  init[:dim] = init[:dim] / np.sum(init[:dim]) # ensuring it is part of the simplex\n",
        "  init[dim:] = init[dim:] / np.sum(init[dim:])  \n",
        "\n",
        "  finished, cnt, t0 = False, 0, time.time()\n",
        "\n",
        "  # # EG ------------\n",
        "  z = np.copy(init)\n",
        "  for _ in range(T):\n",
        "    cnt += 1\n",
        "    zz = proj(z-lr*np.dot(A,z))\n",
        "    z = proj(z-lr*np.dot(A,zz))\n",
        "    if np.linalg.norm(z-x_opt)/np.linalg.norm(x_opt) <= eps or time.time()-t0 > max_time:\n",
        "      print(f\"Reached a relative error of {eps} after {cnt} iterations in {time.time()-t0:.2f} sec.\")\n",
        "      break\n"
      ],
      "metadata": {
        "id": "6F0fomP08X3H"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Experiment on iterations to reach a fixed relative error"
      ],
      "metadata": {
        "id": "jipkmJVoXJ-C"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2e4z_p8f-1SK"
      },
      "outputs": [],
      "source": [
        "dim = 500  # dim(x1) == dim(x2) == dim\n",
        "x_opt = np.ones((2*dim,1))/dim\n",
        "_seed = 0\n",
        "max_iters = 300\n",
        "\n",
        "temp1 = np.concatenate((np.ones((1,dim)), np.zeros((1,dim))),axis=1)\n",
        "temp2 = np.concatenate((np.zeros((1,dim)), np.ones((1,dim))),axis=1)\n",
        "C = np.concatenate((temp1,temp2), axis=0)\n",
        "d = np.ones((2,1))\n",
        "temp = np.linalg.inv(np.dot(C, C.T))\n",
        "temp = np.dot(C.T, temp)\n",
        "dc = np.dot(temp, d)\n",
        "Pc = np.identity(2*dim) - np.dot(temp,C)\n",
        "\n",
        "G = matrix(-np.identity(2*dim), tc='d')\n",
        "h = matrix(np.zeros(((2*dim,1))), tc='d')\n",
        "AA = matrix(C, tc='d')\n",
        "b = matrix(d, tc='d')\n",
        "\n",
        "def f_y(y,c,mu): \n",
        "  \"\"\" Update y in Alg. ACVI. \"\"\"\n",
        "  return -mu * np.sum(np.log(y)) + beta/2 * np.sum((y-c)**2)\n",
        "\n",
        "def proj(z):\n",
        "  P = matrix(np.identity(2*dim), tc='d')\n",
        "  q = matrix(-z, tc='d')\n",
        "  res = solvers.qp(P, q, G, h, AA, b)\n",
        "  return res['x']\n",
        "\n",
        "np.random.seed(_seed)\n",
        "init = np.random.rand(2*dim, 1)\n",
        "init[:dim] = init[:dim] / np.sum(init[:dim])\n",
        "init[dim:] = init[dim:] / np.sum(init[dim:])\n",
        "\n",
        "eps = .02\n",
        "gamma = .3  # step-size for EG, OGDA, LA\n",
        "num_acvi, num_i_acvi, num_gda, num_eg, num_la, num_ogda = [], [], [], [], [], []  \n",
        "rot = np.linspace(.01, .99, 30) \n",
        "\n",
        "for a in rot:\n",
        "  A1 = np.concatenate((a*np.identity(dim), (1-a)*np.identity(dim)), axis=1)\n",
        "  A2 = np.concatenate((-(1-a)*np.identity(dim), a*np.identity(dim)), axis=1)\n",
        "  A = np.concatenate((A1, A2), axis=0)\n",
        "\n",
        "  # Inexact ACVI ------------\n",
        "  beta, mu, delta, K, gamma_, T = .5, 1e-6, .8, 10, 0.05, 100\n",
        "  y, lmd = init, np.zeros(((2*dim,1)))\n",
        "  finished, cnt = False, 0\n",
        "  z_x = np.copy(init)\n",
        "  z_y = np.copy(init)\n",
        "  z_lmd = np.zeros(init.shape)\n",
        "  for t in range(T):\n",
        "    mu *= delta\n",
        "    flag = 1\n",
        "    for k in range(K):\n",
        "      cnt += 1\n",
        "      # Solve approximately the X problem\n",
        "      for _ in range(10):\n",
        "        g = z_x + 1/beta * np.dot(Pc, np.dot(A,z_x)) - np.dot(Pc, z_y) + 1/beta * np.dot(Pc, z_lmd) - dc\n",
        "        z_x -= gamma_ * g\n",
        "\n",
        "      if np.linalg.norm(z_x-x_opt)/np.linalg.norm(x_opt) <= eps or cnt == max_iters:\n",
        "        finished = True\n",
        "        num_i_acvi.append(cnt)\n",
        "        break\n",
        "\n",
        "      # Solve approximately the Y problem\n",
        "      for _ in range(10):\n",
        "        assert all(z_y > 0)\n",
        "        g = - mu * 1/z_y + beta*(z_y - z_x - z_lmd/beta)\n",
        "        z_y -= gamma_ * g\n",
        "        \n",
        "      # Update lambdas\n",
        "      z_lmd += beta*(z_x-z_y)\n",
        "\n",
        "    if flag == 0 or finished:\n",
        "      break\n",
        "\n",
        "  # ACVI ------------\n",
        "  beta, mu, delta, T, K = .5, 1e-6, .5, 10, 1\n",
        "  y, lmd = init, np.zeros(((2*dim,1)))\n",
        "  finished, cnt = False, 0\n",
        "  for t in range(T):\n",
        "    mu *= delta\n",
        "    flag = 1\n",
        "    for k in range(K if t+1 < T else 290):\n",
        "      cnt += 1\n",
        "      x = linalg.solve(np.identity(2*dim) + np.dot(Pc,A)/beta,\n",
        "                       np.dot(Pc, y-lmd/beta) + dc)\n",
        "      if np.linalg.norm(x-x_opt)/np.linalg.norm(x_opt) <= eps or \\\n",
        "         cnt == max_iters:\n",
        "        finished = True\n",
        "        num_acvi.append(cnt)\n",
        "        break\n",
        "      res = opt.minimize(lambda z: f_y(z,x+lmd/beta,mu), y, method='SLSQP')\n",
        "      if not res.success:\n",
        "        print('Failed to update y. t={}. k={}'.format(t, k))\n",
        "        flag = 0\n",
        "        break\n",
        "      y = res.x.reshape(-1,1)\n",
        "      lmd += beta*(x-y)\n",
        "    if flag == 0 or finished:\n",
        "      break\n",
        "\n",
        "  # GDA ------------\n",
        "  z = np.copy(init)\n",
        "  for t in range(max_iters):\n",
        "    z = proj(z - gamma * np.dot(A,z))\n",
        "    if np.linalg.norm(z-x_opt)/np.linalg.norm(x_opt) <= eps or t+1 >= max_iters:\n",
        "      num_gda.append(t+1)\n",
        "      break\n",
        "\n",
        "  # EG ------------\n",
        "  z = np.copy(init)\n",
        "  for t in range(max_iters):\n",
        "    zz = proj(z-gamma*np.dot(A,z))\n",
        "    z = proj(z-gamma*np.dot(A,zz))\n",
        "    if np.linalg.norm(z-x_opt)/np.linalg.norm(x_opt) <= eps or \\\n",
        "       (t+1)*2 >= max_iters:\n",
        "      num_eg.append(2*(t+1))\n",
        "      break\n",
        "\n",
        "  # OGDA ------------\n",
        "  z0 = np.copy(init)\n",
        "  z1 = np.copy(init)\n",
        "  for t in range(max_iters):\n",
        "    z = proj(z1 - 2*gamma*np.dot(A,z1) + gamma*np.dot(A,z0))\n",
        "    z0 = np.copy(z1)\n",
        "    z1 = np.copy(z)\n",
        "    if np.linalg.norm(z-x_opt) / np.linalg.norm(x_opt) <= eps or \\\n",
        "       t+1 >= max_iters:\n",
        "      num_ogda.append(t+1)\n",
        "      break\n",
        "\n",
        "  # LA ------------\n",
        "  z = np.copy(init)\n",
        "  la_k, la_alpha = 5, 0.5\n",
        "  for t in range(max_iters):\n",
        "    z_old = np.copy(z)\n",
        "    for _ in range(la_k):\n",
        "      #z = z - gamma * np.dot(A,z)\n",
        "      z = proj( z - gamma * np.dot(A, z))\n",
        "    z = proj(la_alpha * z_old + (1-la_alpha) * z)\n",
        "    if np.linalg.norm(z - x_opt) / np.linalg.norm(x_opt) <= eps or \\\n",
        "       (t+1) * la_k >= max_iters:\n",
        "      num_la.append(min(la_k*(t+1), max_iters))\n",
        "      break"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Experiment on CPU-time for varying relative error"
      ],
      "metadata": {
        "id": "8TRvSiOEZDKE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "dim = 500  # dim(x1) == dim(x2) == dim\n",
        "x_opt = np.ones((2*dim,1))/dim\n",
        "_seed = 0\n",
        "max_iters = 300\n",
        "\n",
        "temp1 = np.concatenate((np.ones((1,dim)), np.zeros((1,dim))),axis=1)\n",
        "temp2 = np.concatenate((np.zeros((1,dim)), np.ones((1,dim))),axis=1)\n",
        "C = np.concatenate((temp1,temp2), axis=0)\n",
        "d = np.ones((2,1))\n",
        "temp = np.linalg.inv(np.dot(C, C.T))\n",
        "temp = np.dot(C.T, temp)\n",
        "dc = np.dot(temp, d)\n",
        "Pc = np.identity(2*dim) - np.dot(temp,C)\n",
        "\n",
        "G = matrix(-np.identity(2*dim), tc='d')\n",
        "h = matrix(np.zeros(((2*dim,1))), tc='d')\n",
        "AA = matrix(C, tc='d')\n",
        "b = matrix(d, tc='d')\n",
        "\n",
        "def f_y(y,c,mu): \n",
        "  \"\"\" Update y in Alg. ACVI. \"\"\"\n",
        "  return -mu * np.sum(np.log(y)) + beta/2 * np.sum((y-c)**2)\n",
        "\n",
        "def proj(z):\n",
        "  P = matrix(np.identity(2*dim), tc='d')\n",
        "  q = matrix(-z, tc='d')\n",
        "  res = solvers.qp(P, q, G, h, AA, b)\n",
        "  return res['x']\n",
        "\n",
        "np.random.seed(_seed)\n",
        "init = np.random.rand(2*dim, 1)\n",
        "init[:dim] = init[:dim] / np.sum(init[:dim])\n",
        "init[dim:] = init[dim:] / np.sum(init[dim:])\n",
        "\n",
        "aa = .05\n",
        "max_cpu_time = 1500  # in sec\n",
        "max_iters = 3000\n",
        "la_k, la_alpha = 5, .5  # hyperparameters for LA-GDA \n",
        "gamma = .1  # Step-size for GDA, EG, OGDA, LA-GDA\n",
        "A1 = np.concatenate((aa*np.identity(dim),(1-aa)*np.identity(dim)),axis=1)\n",
        "A2 = np.concatenate((-(1-aa)*np.identity(dim),aa*np.identity(dim)),axis=1)\n",
        "A = np.concatenate((A1,A2), axis=0)\n",
        "t_acvi, t_i_acvi, t_gda, t_eg, t_ogda, t_lak = [], [], [], [], [], []\n",
        "eps = [.5, .1, .05, .01, .001, .0005]\n",
        "\n",
        "for de in eps:\n",
        "\n",
        "  # Inexact ACVI ------------\n",
        "  beta, mu, delta, K, gamma_, T = .5, 1e-6, .8, 10, 0.05, 100\n",
        "  y, lmd = init, np.zeros(((2*dim,1)))\n",
        "  finished, cnt = False, 0\n",
        "  z_x = np.copy(init)\n",
        "  z_y = np.copy(init)\n",
        "  z_lmd = np.zeros(init.shape)\n",
        "  t0 = time.time()\n",
        "  for t in range(T):\n",
        "    mu *= delta\n",
        "    for k in range(K):\n",
        "      cnt += 1\n",
        "      # Solve approximately the X problem\n",
        "      for _ in range(10):\n",
        "        g = z_x + 1/beta * np.dot(Pc, np.dot(A,z_x)) - np.dot(Pc, z_y) + 1/beta * np.dot(Pc, z_lmd) - dc\n",
        "        z_x -= gamma_ * g\n",
        "      if np.linalg.norm(z_x-x_opt)/np.linalg.norm(x_opt) <= de or time.time()-t0 >= max_cpu_time:\n",
        "        finished = True\n",
        "        t_i_acvi.append(time.time()-t0)\n",
        "        break\n",
        "      # Solve approximately the Y problem\n",
        "      for _ in range(10):\n",
        "        assert all(z_y >= 0)\n",
        "        g = - mu * 1/z_y + beta*(z_y - z_x - z_lmd/beta)\n",
        "        z_y -= gamma_ * g\n",
        "      # Update lambdas\n",
        "      z_lmd += beta*(z_x-z_y)\n",
        "    if finished:\n",
        "      break\n",
        "\n",
        "  # ACVI ------------\n",
        "  beta, mu, delta, T, K = .5, 1e-6, .5, 10, 1\n",
        "  y, lmd = init,np.zeros(((2*dim,1)))\n",
        "  t0 = time.time()\n",
        "  for t in range(T):\n",
        "    mu *= delta\n",
        "    flag = 1\n",
        "    for k in range(K if t+1 < T else 290):\n",
        "      x = linalg.solve(np.identity(2*dim)+np.dot(Pc,A)/beta,np.dot(Pc,y-lmd/beta)+dc)\n",
        "      if np.linalg.norm(x-x_opt)/np.linalg.norm(x_opt) <= de or time.time()-t0 >= max_cpu_time:\n",
        "        flag = 0\n",
        "        t_acvi.append(time.time()-t0)\n",
        "        break\n",
        "      res = opt.minimize(lambda z: f_y(z,x+lmd/beta,mu), y, method='SLSQP', jac=None)\n",
        "      if not res.success:\n",
        "        print('Failed to update y. t={}. k={}'.format(t,k))\n",
        "        flag = 0\n",
        "        break\n",
        "      y = res.x.reshape(-1,1)\n",
        "      lmd += beta*(x-y)\n",
        "    if flag == 0:\n",
        "      break\n",
        "\n",
        "  # GDA ------------\n",
        "  z = np.copy(init)\n",
        "  t0 = time.time()\n",
        "  for t in range(max_iters):\n",
        "    z = proj(z-gamma*np.dot(A,z))\n",
        "    if np.linalg.norm(z-x_opt) / np.linalg.norm(x_opt) <= de or \\\n",
        "        time.time()-t0 >= max_cpu_time:\n",
        "      t_gda.append(time.time()-t0)\n",
        "      break\n",
        "\n",
        "  # EG ------------\n",
        "  z = np.copy(init)\n",
        "  t0 = time.time()\n",
        "  for t in range(max_iters):\n",
        "    zz = proj(z-gamma*np.dot(A,z))\n",
        "    z = proj(z-gamma*np.dot(A,zz))\n",
        "    if np.linalg.norm(z-x_opt)/np.linalg.norm(x_opt) <= de or \\\n",
        "        time.time()-t0 >= max_cpu_time:\n",
        "      t_eg.append(time.time()-t0)\n",
        "      break\n",
        "  # OGDA ------------\n",
        "  z0 = np.copy(init)\n",
        "  z1 = np.copy(init)\n",
        "  t0 = time.time()\n",
        "  for t in range(max_iters):\n",
        "    z = proj(z1 - 2*gamma*np.dot(A,z1) + gamma*np.dot(A,z0))\n",
        "    z0 = np.copy(z1)\n",
        "    z1 = np.copy(z)\n",
        "    if np.linalg.norm(z-x_opt)/np.linalg.norm(x_opt) <= de or \\\n",
        "        time.time()-t0 >= max_cpu_time:\n",
        "      t_ogda.append(time.time()-t0)\n",
        "      break\n",
        "\n",
        "  # LAK-GDA ------------\n",
        "  la_k, la_alpha = 5, .5 \n",
        "  z = np.copy(init)\n",
        "  t0 = time.time()\n",
        "  for t in range(max_iters):\n",
        "    z_old = z\n",
        "    for i in range(la_k):\n",
        "      z = proj(z-gamma*np.dot(A, z))\n",
        "    z = proj(la_alpha*z_old + (1-la_alpha)*z)\n",
        "    if np.linalg.norm(z-x_opt) / np.linalg.norm(x_opt) <= de or \\\n",
        "        time.time()-t0 >= max_cpu_time:\n",
        "      t_lak.append(time.time()-t0)\n",
        "      break"
      ],
      "metadata": {
        "id": "ks85tikVZigC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Experiment on warm-up"
      ],
      "metadata": {
        "id": "faJRdLwZyvng"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from numpy.core.arrayprint import format_float_scientific\n",
        "import time\n",
        "\n",
        "eps = 1e-5\n",
        "max_iters = 3000\n",
        "a = 0.05\n",
        "\n",
        "A1 = np.concatenate((a*np.identity(dim), (1-a)*np.identity(dim)), axis=1)\n",
        "A2 = np.concatenate((-(1-a)*np.identity(dim), a*np.identity(dim)), axis=1)\n",
        "A = np.concatenate((A1, A2), axis=0)\n",
        "\n",
        "warmup = [5,10,20,30,40,50,60,70,80,90,100,110,120,130]\n",
        "\n",
        "num_grads = []\n",
        "all_times = []\n",
        "\n",
        "for K_ in [1,5,10,20,30,40,50,60,70]:\n",
        "\n",
        "  num_grad = []\n",
        "  t_i_acvi = []\n",
        "\n",
        "  for K_0 in warmup:\n",
        "\n",
        "    print(f\"K={K_}, K0={K_0}\")\n",
        "\n",
        "    # Inexact ACVI ------------\n",
        "    beta, mu, delta, gamma_, T_ = .5, 1e-6, .8, 0.05, 5000\n",
        "    n_grads = 0\n",
        "    finished, cnt = False, 0\n",
        "    z_x = np.copy(init)\n",
        "    z_y = np.copy(init)\n",
        "    z_lmd = np.zeros(init.shape)\n",
        "    t0 = time.time()\n",
        "    for t in range(0, T_):\n",
        "      mu *= delta\n",
        "      #print(mu)\n",
        "      for k in range(0, K_0 if t == 0 else K_): # if t+1 < T_ else last_K):\n",
        "        cnt += 1\n",
        "        # Solve approximately the X problem\n",
        "        for _ in range(10): \n",
        "          n_grads += 1\n",
        "          g = z_x + 1/beta * np.dot(Pc, np.dot(A,z_x)) - np.dot(Pc, z_y) + 1/beta * np.dot(Pc, z_lmd) - dc\n",
        "          z_x -= gamma_ * g\n",
        "\n",
        "        if np.linalg.norm(z_x-x_opt)/np.linalg.norm(x_opt) <= eps or cnt == max_iters:\n",
        "          finished = True\n",
        "          num_grad.append(n_grads)\n",
        "          t_i_acvi.append(time.time()-t0)\n",
        "          print(f\"num grad to reach distance of 1e-5: {n_grads}, time: {t_i_acvi[-1]}s\")\n",
        "          break\n",
        "\n",
        "        # Solve approximately the Y problem\n",
        "        for _ in range(10):\n",
        "          n_grads += 1\n",
        "          assert all(z_y > 0)\n",
        "          g = - mu * 1/z_y + beta*(z_y - z_x - z_lmd/beta)\n",
        "          z_y -= gamma_ * g\n",
        "          \n",
        "        # Update lambdas\n",
        "        z_lmd += beta*(z_x-z_y)\n",
        "\n",
        "      if finished:\n",
        "        break\n",
        "\n",
        "  num_grads.append(num_grad)\n",
        "  all_times.append(t_i_acvi)\n",
        "\n",
        "print(num_grads)"
      ],
      "metadata": {
        "id": "cVFxHUzAZil_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "uEPV68cIyzSu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "JuQjeSlFyzU7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "gEJkiB2WyzWj"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}