{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4889fbd3",
   "metadata": {},
   "source": [
    "# Dirac-3S QPLIB Benchmarking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "563202ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "from eqc_direct.client import EqcClient\n",
    "from eqc_direct.utils import *\n",
    "import numpy as np\n",
    "import time\n",
    "import os\n",
    "\n",
    "ip_address=\"172.18.15.110\"\n",
    "client = EqcClient(ip_address=ip_address) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c589d169",
   "metadata": {},
   "source": [
    "### Set up dirac-3S solver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5f13ec3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def direct_dirac3(indices, coefficients, relaxation_schedule, sum_constraint, num_samples, ip_address):\n",
    "    eqc_client = EqcClient(ip_address=ip_address) #S9 \"172.18.41.45\", #S3 \"172.18.41.173\"\n",
    "    lock_id, start_ts, end_ts=eqc_client.wait_for_lock()\n",
    "    try:\n",
    "        result_dict = eqc_client.solve_sum_constrained(\n",
    "            lock_id=lock_id,\n",
    "            poly_indices = indices,\n",
    "            poly_coefficients = coefficients,\n",
    "            relaxation_schedule = relaxation_schedule,\n",
    "            sum_constraint = sum_constraint,\n",
    "            num_samples = num_samples,\n",
    "        )\n",
    "        total_time = np.array(result_dict[\"preprocessing_time\"])+np.array(result_dict[\"postprocessing_time\"])+np.array(result_dict[\"runtime\"])\n",
    "        print(f\"Total execution time(s):{total_time}\")\n",
    "    finally:\n",
    "        # release lock when finished using the device\n",
    "        lock_release_out = eqc_client.release_lock(lock_id=lock_id)\n",
    "    return result_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cec6c91",
   "metadata": {},
   "source": [
    "### load instance function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b6398f8f",
   "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": "9e75f398",
   "metadata": {},
   "source": [
    "### solve using dirac-3S"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b0e2be89",
   "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",
    "\n",
    "poly_indices, poly_coefficients = convert_hamiltonian_to_poly_format(\n",
    "    linear_terms=c,\n",
    "    quadratic_terms=M.copy(),   # the util mutates its input\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3ee7e923",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Submitting job to Dirac-3S\n",
      "Total execution time(s):[2.60927448 3.92017277 3.66367798 5.65877393 3.50433762 3.31685099\n",
      " 4.9969642  2.65017603 2.94848334 3.96039353 4.76025672 5.69540859\n",
      " 3.56813195 4.78678866 3.77307181 3.79501493 3.80302066 3.06004788\n",
      " 3.61366897 3.79009187 3.36575746 2.88107067 3.71976098 4.5675675\n",
      " 3.07030962 4.16581185 4.04716632 4.73760869 3.84494819 2.83097038\n",
      " 3.98634813 4.67104459 3.73876965 3.85616669 2.68340677 4.18464988\n",
      " 3.98370418 4.75931722 3.90973101 4.86702825 4.46201827 2.94292931\n",
      " 4.27756047 3.44825469 3.7916764  2.74540909 3.15322383 4.18542743\n",
      " 2.83082731 3.97312964 2.9065745  4.4457614  3.55981524 3.19108973\n",
      " 3.93139399 3.77664767 3.94058345 2.63672721 3.83822486 3.5642059\n",
      " 3.71930424 3.01746587 2.71147231 5.44961055 3.03422832 3.50443695\n",
      " 3.14953573 4.63366989 3.4780005  4.37385515 4.02699331 4.06825119\n",
      " 3.804987   2.96381695 3.53761173 3.83295924 3.04909765 3.63119259\n",
      " 4.00034818 3.93047291 4.8892151  2.73325073 2.97676449 4.75357767\n",
      " 3.2108481  3.2900839  3.74374173 3.93122042 3.67027874 5.30804687\n",
      " 4.6108509  4.86014954 4.89530536 3.84987223 3.32743471 3.97914636\n",
      " 3.97654084 4.17954905 3.96836448 5.34683647]\n",
      "Dirac-3 Run complete. Response received.\n",
      "Dirac-3 all samples energies:[0.09, 0.04, 0.06, 0.06, 0.09, 0.04, 0.04, 0.11, 0.09, 0.08, 0.07, 0.04, 0.13, 0.05, 0.11, 0.07, 0.07, 0.06, 0.0, 0.08, 0.05, 0.05, 0.06, 0.07, 0.06, 0.0, 0.0, 0.05, 0.11, 0.15, 0.07, 0.07, 0.11, 0.11, 0.14, 0.11, 0.04, 0.07, 0.05, 0.05, 0.0, 0.1, 0.0, 0.04, 0.05, 0.04, 0.09, 0.0, 0.1, 0.0, 0.07, 0.07, 0.07, 0.04, 0.11, 0.0, 0.05, 0.1, 0.04, 0.0, 0.05, 0.06, 0.07, 0.05, 0.09, 0.0, 0.05, 0.07, 0.0, 0.05, 0.06, 0.04, 0.05, 0.06, 0.02, 0.04, 0.05, 0.0, 0.01, 0.0, 0.04, 0.06, 0.06, 0.0, 0.07, 0.04, 0.04, 0.09, 0.04, 0.0, 0.0, 0.0, 0.04, 0.0, 0.06, 0.05, 0.0, 0.07, 0.0, 0.05]\n",
      "Dirac-3 all samples runtimes:[2.38, 3.69, 3.43, 5.43, 3.27, 3.09, 4.77, 2.42, 2.72, 3.73, 4.53, 5.47, 3.34, 4.56, 3.54, 3.57, 3.57, 2.83, 3.38, 3.56, 3.14, 2.65, 3.49, 4.34, 2.84, 3.94, 3.82, 4.51, 3.61, 2.6, 3.76, 4.44, 3.51, 3.63, 2.45, 3.95, 3.75, 4.53, 3.68, 4.64, 4.23, 2.71, 4.05, 3.22, 3.56, 2.52, 2.92, 3.96, 2.6, 3.74, 2.68, 2.63, 3.33, 2.96, 3.7, 3.55, 3.71, 2.41, 3.61, 3.33, 3.49, 2.79, 2.48, 3.64, 2.8, 3.27, 2.92, 4.4, 3.25, 4.14, 3.8, 3.84, 3.58, 2.73, 3.31, 3.6, 2.82, 3.4, 3.77, 3.7, 4.66, 2.5, 2.75, 4.52, 2.98, 3.06, 3.51, 3.7, 3.44, 5.08, 4.38, 4.63, 4.67, 3.62, 3.1, 3.75, 3.75, 3.95, 3.74, 5.12]\n"
     ]
    }
   ],
   "source": [
    "def objective(x):\n",
    "    \"\"\"QPLIB objective, same convention as the poly submitted to Dirac-3.\"\"\"\n",
    "    x = np.asarray(x, dtype=float)\n",
    "    return float(x @ M @ x + c @ x) + inst[\"obj_const\"]\n",
    "\n",
    "sum_constraint = float(inst[\"rhs\"][0])   # sum_i x_i = 1\n",
    "relaxation_schedule = 1\n",
    "num_samples = 100\n",
    "\n",
    "\n",
    "poly_indices, poly_coefficients = convert_hamiltonian_to_poly_format(\n",
    "                    linear_terms=c,\n",
    "                    quadratic_terms=M.copy(),   # the util mutates its input\n",
    "                )\n",
    "\n",
    "dirac_start = time.time()\n",
    "print(\"Submitting job to Dirac-3S\")\n",
    "response = direct_dirac3(indices=poly_indices,\n",
    "                        coefficients=poly_coefficients,\n",
    "                        relaxation_schedule=relaxation_schedule,\n",
    "                        sum_constraint=sum_constraint,\n",
    "                        ip_address=ip_address,\n",
    "                        num_samples=num_samples)\n",
    "print(f\"Dirac-3 Run complete. Response received.\")\n",
    "energies = response['energy']\n",
    "solutions = response['solution']\n",
    "runtimes = response['runtime']\n",
    "preprocessing_time = response['preprocessing_time']\n",
    "postprocessing_time = response['postprocessing_time']\n",
    "\n",
    "print(f\"Dirac-3 all samples energies:{[round(e, 2) for e in energies]}\")\n",
    "print(f\"Dirac-3 all samples runtimes:{[round(t, 2) for t in runtimes]}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "26dd42e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Eenergy value obtained by Dirac-3:0.0010485\n",
      "Solution sum:1.0000001 (target 1.0), support:2\n",
      "Preprocessing Time:0.22956259548664093\n",
      "Time taken for this run:3.384004592895508\n"
     ]
    }
   ],
   "source": [
    "best = min(zip(energies, solutions, runtimes), key=lambda x: x[0])\n",
    "dirac_energy, best_solution, best_runtime = best\n",
    "\n",
    "\n",
    "print(f\"Eenergy value obtained by Dirac-3:{dirac_energy}\")\n",
    "print(f\"Solution sum:{sum(best_solution)} (target {sum_constraint}), support:{np.count_nonzero(best_solution)}\")\n",
    "print(f\"Preprocessing Time:{preprocessing_time}\")\n",
    "print(f\"Time taken for this run:{best_runtime}\")"
   ]
  }
 ],
 "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
}