{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "855f0802",
   "metadata": {},
   "source": [
    "# PGD QPLIB Benchmarking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2e1bbd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bfb8bcd",
   "metadata": {},
   "source": [
    "### Define multistart projected gradient descent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "56605a5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def project_simplex(v, sum_constraint):\n",
    "    \"\"\"Exact Euclidean projection onto {x >= 0, sum(x) == R}.\n",
    "    \"\"\"\n",
    "    R = float(sum_constraint)\n",
    "    v = np.asarray(v, dtype=float)\n",
    "    u = np.sort(v)[::-1]\n",
    "    cumulative = np.cumsum(u) - R\n",
    "    ind = np.arange(1, v.size + 1)\n",
    "    rho = ind[u - cumulative / ind > 0][-1]\n",
    "    return np.maximum(v - cumulative[rho - 1] / rho, 0.0)\n",
    "\n",
    "\n",
    "def pgd(Q, c, sum_constraint, lr=0.01, max_iter=10**4, tol=1e-9, x0=None):\n",
    "    \"\"\"Projected gradient descent on {x >= 0, sum(x) == R}.\n",
    "    \"\"\"\n",
    "    Q = np.asarray(Q, dtype=float)\n",
    "    c = np.asarray(c, dtype=float).ravel()\n",
    "    n = Q.shape[0]\n",
    "    R = float(sum_constraint)\n",
    "\n",
    "    x = np.full(n, R / n) if x0 is None else np.asarray(x0, dtype=float)\n",
    "    x = project_simplex(x, R)\n",
    "    QT = Q + Q.T                                   # gradient of x'Qx is (Q+Q')x\n",
    "\n",
    "    it = 0\n",
    "    for it in range(max_iter):\n",
    "        grad = QT @ x + c\n",
    "        x_next = project_simplex(x - lr * grad, R)\n",
    "        delta = np.linalg.norm(x_next - x)\n",
    "        x = x_next                                 # keep the newer iterate\n",
    "        if delta < tol:\n",
    "            break\n",
    "\n",
    "    return x, float(x @ Q @ x + c @ x), it + 1\n",
    "\n",
    "\n",
    "def pgd_multistart(Q, c, sum_constraint, restarts=32, seed=0, **kwargs):\n",
    "    \"\"\"PGD from the simplex centre plus random starts.\n",
    "    \"\"\"\n",
    "    rng = np.random.default_rng(seed)\n",
    "    n = Q.shape[0]\n",
    "    solutions, energies, times, total_iters = [], [], [], 0\n",
    "    for r in range(restarts):\n",
    "        start = time.time()\n",
    "        x0 = None\n",
    "        if r > 0:\n",
    "            x0 = rng.random(n)\n",
    "            x0 *= sum_constraint / x0.sum()\n",
    "        x, energy, iters = pgd(Q, c, sum_constraint, x0=x0, **kwargs)\n",
    "        pgd_t = time.time()-start\n",
    "        solutions.append(x)\n",
    "        energies.append(energy)\n",
    "        times.append(pgd_t)\n",
    "        total_iters += iters\n",
    "    return solutions, energies, times, total_iters\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce05228d",
   "metadata": {},
   "source": [
    "### load instance function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c2c4a8be",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_qplib(path):\n",
    "    with open(path) as fh:\n",
    "        toks = [t for t in (line.split(\"#\")[0].strip() for line in fh) if t]\n",
    "\n",
    "    pos = 0\n",
    "\n",
    "    def nxt():\n",
    "        nonlocal pos\n",
    "        val = toks[pos]\n",
    "        pos += 1\n",
    "        return val\n",
    "\n",
    "    def defaulted(length):\n",
    "        \"\"\"Read a `default value` / `number of non-defaults` / entries block.\"\"\"\n",
    "        arr = np.full(length, float(nxt()))\n",
    "        for _ in range(int(nxt())):\n",
    "            i, v = nxt().split()\n",
    "            arr[int(i) - 1] = float(v)\n",
    "        return arr\n",
    "\n",
    "    name = nxt()\n",
    "    probtype = nxt()\n",
    "    sense = nxt().lower()\n",
    "    n = int(nxt())\n",
    "    m = int(nxt())\n",
    "\n",
    "    # quadratic terms of the objective (each unordered pair listed once)\n",
    "    Q = np.zeros((n, n), dtype=np.float32)\n",
    "    for _ in range(int(nxt())):\n",
    "        i, j, v = nxt().split()\n",
    "        i, j, v = int(i) - 1, int(j) - 1, float(v)\n",
    "        if i == j:\n",
    "            Q[i, i] = v\n",
    "        else:\n",
    "            Q[i, j] = Q[j, i] = 0.5 * v   # split the pair coefficient over both triangles\n",
    "\n",
    "    b = defaulted(n)            # linear terms of the objective\n",
    "    obj_const = float(nxt())\n",
    "\n",
    "    # linear terms of the constraints\n",
    "    A = np.zeros((m, n))\n",
    "    for _ in range(int(nxt())):\n",
    "        k, i, v = nxt().split()\n",
    "        A[int(k) - 1, int(i) - 1] = float(v)\n",
    "\n",
    "    inf = float(nxt())\n",
    "    lhs, rhs = defaulted(m), defaulted(m)\n",
    "    lb, ub = defaulted(n), defaulted(n)\n",
    "\n",
    "    return dict(name=name, type=probtype, sense=sense, n=n, m=m,\n",
    "                Q=Q, b=b, obj_const=obj_const,\n",
    "                A=A, lhs=lhs, rhs=rhs, lb=lb, ub=ub, inf=inf)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a826a87",
   "metadata": {},
   "source": [
    "### solve using multi-start PGD\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3a38e512",
   "metadata": {},
   "outputs": [],
   "source": [
    "inst_dir = \"Instances/\"\n",
    "inst_name = \"QPLIB_2761.qplib\"\n",
    "inst = read_qplib(os.path.join(inst_dir, inst_name))\n",
    "n, Q, b = inst[\"n\"], inst[\"Q\"], inst[\"b\"]\n",
    "M = 0.5 * Q\n",
    "c = b.copy()\n",
    "sum_constraint = float(inst[\"rhs\"][0])   # sum_i x_i = 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "aee9a4c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "restarts =100\n",
    "pgd_solutions, pgd_energies, times,total_iters = pgd_multistart(Q=M,\n",
    "                                                  c=c,\n",
    "                                                  sum_constraint=sum_constraint,\n",
    "                                                  restarts=restarts,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "52ed5cba",
   "metadata": {},
   "outputs": [],
   "source": [
    "pgd_energies_arr= np.array(pgd_energies)\n",
    "lo = min(pgd_energies_arr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "13d5404a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Multi-start PGD all energies: [0.10097685334659422, 0.0010485291160228403, 0.07039351549786885, 0.06693547209452383, 0.0010485291160227453, 0.01768112059763443, 0.09006403320120959, 0.0010485291160227796, 0.0010485291160227885, 0.06693263741421218, 0.0010485291160227468, 0.0716207293298459, 0.06606244717371348, 0.028507570296855817, 0.038198279039776746, 0.017681120597634478, 0.0010485291160228102, 0.0010485291160228138, 0.0010485291160228386, 0.09859646900306727, 0.06606244717371344, 0.06726658906968412, 0.07284124272182768, 0.07404253782634901, 0.001048529116022831, 0.06672539316893514, 0.0010485291160227575, 0.003990561779633381, 0.03833167156964212, 0.020248759338627555, 0.05296478049682289, 0.06606244717371344, 0.0010485291160228234, 0.010117643197959296, 0.09641535524734914, 0.0039905617796334265, 0.05732327918611205, 0.12691332537961625, 0.07020143907506993, 0.001048529116022746, 0.09641535524734914, 0.07222367403229096, 0.003990561779633378, 0.06606244717371343, 0.04578102721482212, 0.0010485291160227629, 0.11999728202282746, 0.09006403320120959, 0.0010485291160227557, 0.001048529116022779, 0.120796649804321, 0.0010485291160227798, 0.0010485291160227787, 0.048849061347131705, 0.001048529116022755, 0.0010485291160228336, 0.001048529116022795, 0.06531440522549319, 0.003990561779633435, 0.05728837829848721, 0.0010485291160227683, 0.0010485291160227826, 0.0010485291160228294, 0.04624217095385548, 0.017681120597634464, 0.06672539316893508, 0.00399056177963337, 0.0736028973621164, 0.1114976268045212, 0.0010485291160227995, 0.0010485291160228138, 0.10097685334659422, 0.001048529116022808, 0.06325572427054121, 0.0601078862513565, 0.01768112059763439, 0.001048529116022766, 0.003990561779633373, 0.11363332831651837, 0.0010485291160228023, 0.017681120597634395, 0.0010485291160227757, 0.0010485291160228308, 0.04489618562362864, 0.060010792088498695, 0.10456220462874036, 0.001048529116022779, 0.08748024698693273, 0.0010485291160227603, 0.09246536949307427, 0.06325572427054117, 0.0010485291160228062, 0.09819787932686895, 0.0010485291160227848, 0.014390150851714398, 0.10097685334659424, 0.07222367403229091, 0.12079664980432094, 0.09489267716053129, 0.01768112059763447]\n",
      "PGD best energy:0.0010485291160227453\n"
     ]
    }
   ],
   "source": [
    "print(f\"Multi-start PGD all energies: {pgd_energies}\")\n",
    "print(f\"PGD best energy:{lo}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ba9ab58",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}