{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "26cef5af",
   "metadata": {},
   "source": [
    "# PGD Planted Benchmarking"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6937af4",
   "metadata": {},
   "source": [
    "### import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "94968b5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34f2ad22",
   "metadata": {},
   "source": [
    "### Define multistart projected gradient descent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "fdae4b10",
   "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": "f3923bba",
   "metadata": {},
   "source": [
    "### Define function to load instances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "fe11fa83",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def loadQ(path, n):\n",
    "    tri = np.load(path, mmap_mode=\"r\")\n",
    "    Q = np.empty((n,n), dtype=tri.dtype)\n",
    "    iu = np.triu_indices(n)\n",
    "    Q[iu] = tri\n",
    "    Q.T[iu] = tri\n",
    "    return Q"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43622cd2",
   "metadata": {},
   "source": [
    "### Solve with multistart PGD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3e64448b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loaded file sucessfully.\n"
     ]
    }
   ],
   "source": [
    "# Instance details\n",
    "optimal_energy = 12000\n",
    "sum_constraint = 100\n",
    "num_var = 2000\n",
    "k = 44\n",
    "ub =10\n",
    "seed =100\n",
    "\n",
    "name = f\"STQP_n_{num_var}_k_{k}_R_{sum_constraint}_seed_{seed}_ub_{ub}\"\n",
    "instance_path = os.path.join(f\"Instances/{name}.npy\")\n",
    "c = np.zeros(num_var)\n",
    "try:\n",
    "    Q= loadQ(instance_path,num_var)\n",
    "    print(\"loaded file sucessfully.\")\n",
    "except FileNotFoundError:\n",
    "    print(f\"File {instance_path} does not exist.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "41ea7c26",
   "metadata": {},
   "outputs": [],
   "source": [
    "# solver parameters\n",
    "restarts = 10\n",
    "learning_rate = 0.01\n",
    "max_iter = 2 * 10**4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a6ca2285",
   "metadata": {},
   "outputs": [],
   "source": [
    "# solve using multistart PGD\n",
    "pgd_start = time.time()\n",
    "solutions, energies, times, total_iters = pgd_multistart(\n",
    "    Q=Q,\n",
    "    c=c,\n",
    "    sum_constraint=sum_constraint,\n",
    "    restarts=restarts,\n",
    "    lr=learning_rate,\n",
    "    max_iter=max_iter,\n",
    ")\n",
    "pgd_time = time.time() - pgd_start\n",
    "best_energy = min(energies)\n",
    "best_solution = solutions[int(np.argmin(energies))]\n",
    "indices = np.where(best_solution>1e-6)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "496c68d6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "restarts:10, total PGD iterations:16695\n",
      "best energy over restarts:12815.673891\n",
      "support size:14\n",
      "time taken bt PGD 152.516615152359\n"
     ]
    }
   ],
   "source": [
    "print(f\"restarts:{restarts}, total PGD iterations:{total_iters}\")\n",
    "print(f\"best energy over restarts:{best_energy:.6f}\")\n",
    "print(f\"support size:{len(indices)}\")\n",
    "print(f\"time taken by PGD {pgd_time}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "fdb5c5d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute relative gap\n",
    "tol = 1e-7\n",
    "abs_gap = best_energy-optimal_energy\n",
    "relative_gap = (round(abs_gap,6)*100)/round(optimal_energy,6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "1dcac44d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gradient descent solution is not optimal\n",
      "Absolute Gap:815.6738909034775\n",
      "Relative Gap(%):6.797282425\n"
     ]
    }
   ],
   "source": [
    "# print result\n",
    "if abs_gap>tol:\n",
    "    print(f\"Gradient descent solution is not optimal\")\n",
    "    print(f\"Absolute Gap:{abs_gap}\")\n",
    "    print(f\"Relative Gap(%):{abs(relative_gap)}\")\n",
    "else:\n",
    "    print(f\"Optimal solution found:{best_energy}.\")\n"
   ]
  }
 ],
 "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
}