{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set Partitioning on Dirac\n",
    "#### Device: Dirac-1\n",
    "The set partitioning problem is an optimization problem which selects sets $S_i$ from a collection $S$ such that each member $x\\in X=\\bigcup_i S_i$ is included in exactly one $S_i$ of the selected sets (see [Operations Research Journal Vol. 17, No. 5](https://doi.org/10.1287/opre.17.5.848)). It has applications in logistics, design, manufacturing and scheduling, among others.\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Formulation\n",
    "The set partitioning problem is formulated by creating constraints which specify that only one out of all the sets $S_i$ of which a particular member $s$ is selected. This looks like\n",
    "$$\n",
    "\\sum_{i,x\\in S_i} s_i = 1\\quad \\forall x\\in X\n",
    "$$\n",
    "where $s_i\\in\\{0,1\\}$ indicates set $S_i$ is selected. In addition to a constraint for each member, there is an objective function which measures the cost, weight or benefit of selecting certain sets from $S$. The objective function could be any form, but we can solve linear or quadratic objective functions with Dirac. A quadratic objective function has coefficients $c_{ij},c_{i}$ for quadratic and linear terms, respectively. With the variables $s_i$, we have\n",
    "$$\n",
    "\\sum_j\\sum_i c_{ij}s_is_j + \\sum_i c_is_i.\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementation\n",
    "With constraint-based solvers, it would be sufficient to implement the objective and constraints directly in the modeling or matrix format required by the solver, but an addition step is required for unconstrained solving. Penalties can be created from the constraints $As=b$ using\n",
    "$$\n",
    "P(s)=s^T(A^TA)s-(2b^TA)^Ts+b^Tb.\n",
    "$$\n",
    "\n",
    "When all constraints are met, $P(s)=0$ and when any constraint is violated, $P(s)>0$. There is a difficulty in combining the objective function and penalties in that a scalar value of the objective function for a constraint violating solution could be less than 0 or at least less than the penalty, which results in a value of the total function less than 0. This will result in an optimizer finding an infeasible solution, unless a multiplier is applied to $P(s)$. Sufficiently large multipliers will guarantee that no infeasible solution will take on a value of the total function which is less than any value for a feasible solution.\n",
    "\n",
    "Here we have a value which is known to be sufficient, $\\alpha$. We apply it to $P(s)$ to get the function\n",
    "$$\n",
    "Q(s)=\\sum_j\\sum_i c_{ij}s_is_j + \\sum_i c_is_i + \\alpha P(s)\n",
    "$$\n",
    "\n",
    "## Demonstration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from qci_client import QciClient"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating the Data\n",
    "$S$ is a collection of 9 different sets. The members of the sets are the letters A through F. $W$ are the weights of each subset $S_i$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(21)\n",
    "S = [{\"A\", \"B\", \"C\"}, {\"D\", \"E\", \"F\"}, {\"A\", \"F\"}, {\"B\", \"E\"}, {\"C\", \"D\"}, {\"A\"},\n",
    "     {\"B\"}, {\"C\", \"D\", \"E\"}, {\"B\", \"C\"}]\n",
    "W = [100 * np.random.random() for S_i in S]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$X$ is the union of all $S_i$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'A', 'B', 'C', 'D', 'E', 'F'}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X = set()\n",
    "for S_i in S:\n",
    "    X = X.union(S_i)\n",
    "X"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create the constraints by indicating if a member is in a subset with a 1 in the position for the variable $s_i$ for every $x$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1, 0, 1, 0, 0, 1, 0, 0, 0],\n",
       "       [0, 1, 1, 0, 0, 0, 0, 0, 0],\n",
       "       [1, 0, 0, 1, 0, 0, 1, 0, 1],\n",
       "       [0, 1, 0, 0, 1, 0, 0, 1, 0],\n",
       "       [1, 0, 0, 0, 1, 0, 0, 1, 1],\n",
       "       [0, 1, 0, 1, 0, 0, 0, 1, 0]])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = []\n",
    "for x in X:\n",
    "    row = [1 if x in S_i else 0 for S_i in S]\n",
    "    A.append(row)\n",
    "A = np.array(A)\n",
    "A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1., 1., 1., 1., 1., 1.])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b = np.ones((A.shape[0],))\n",
    "b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[3, 0, 1, 1, 1, 1, 1, 1, 2],\n",
       "        [0, 3, 1, 1, 1, 0, 0, 2, 0],\n",
       "        [1, 1, 2, 0, 0, 1, 0, 0, 0],\n",
       "        [1, 1, 0, 2, 0, 0, 1, 1, 1],\n",
       "        [1, 1, 0, 0, 2, 0, 0, 2, 1],\n",
       "        [1, 0, 1, 0, 0, 1, 0, 0, 0],\n",
       "        [1, 0, 0, 1, 0, 0, 1, 0, 1],\n",
       "        [1, 2, 0, 1, 2, 0, 0, 3, 1],\n",
       "        [2, 0, 0, 1, 1, 0, 1, 1, 2]]),\n",
       " array([[-6.],\n",
       "        [-6.],\n",
       "        [-4.],\n",
       "        [-4.],\n",
       "        [-4.],\n",
       "        [-2.],\n",
       "        [-2.],\n",
       "        [-6.],\n",
       "        [-4.]]),\n",
       " 6.0)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "J = A.T@A\n",
    "h = -2 * b.T@A\n",
    "n = h.shape[0]\n",
    "h = h.reshape((n, 1))\n",
    "offset = b.T@b\n",
    "J, h, offset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solving the problem without an objective function should reveal if there are multiple solutions to the exact cover problem.\n",
    "\n",
    "First, create a connection to the REST API using `QciClient`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "token = \"token_here\"\n",
    "api_url = \"https://api.qci-prod.com\"\n",
    "client = QciClient(api_token=token, url=api_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next line creates a QUBO from the quadratic and linear portions of the penalty function by adding all the linear terms in the diagonal of the quadratic operator. This file is uploaded to the API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-3.,  0.,  1.,  1.,  1.,  1.,  1.,  1.,  2.],\n",
       "       [ 0., -3.,  1.,  1.,  1.,  0.,  0.,  2.,  0.],\n",
       "       [ 1.,  1., -2.,  0.,  0.,  1.,  0.,  0.,  0.],\n",
       "       [ 1.,  1.,  0., -2.,  0.,  0.,  1.,  1.,  1.],\n",
       "       [ 1.,  1.,  0.,  0., -2.,  0.,  0.,  2.,  1.],\n",
       "       [ 1.,  0.,  1.,  0.,  0., -1.,  0.,  0.,  0.],\n",
       "       [ 1.,  0.,  0.,  1.,  0.,  0., -1.,  0.,  1.],\n",
       "       [ 1.,  2.,  0.,  1.,  2.,  0.,  0., -3.,  1.],\n",
       "       [ 2.,  0.,  0.,  1.,  1.,  0.,  1.,  1., -2.]])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "P = J + np.diag(h.T[0])\n",
    "P"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "qubo_file = {\"file_name\": \"penalty-only-sp-qubo\", \"file_config\": {\"qubo\": {\"data\":P}}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_id = client.upload_file(file=qubo_file)[\"file_id\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the file ID returned by the upload method, build the job body requesting the job to run on eqc1 (aka Dirac-1), returning 5 samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "body = client.build_job_body(job_type=\"sample-qubo\", qubo_file_id=file_id, job_params={\n",
    "    \"device_type\": \"dirac-1\", \"num_samples\": 5\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The job type `sample-qubo` converts the QUBO into an Ising Hamiltonian before sending to Dirac-1 for sampling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-04-29 16:29:09 - Dirac allocation balance = 0 s (unmetered)\n",
      "2024-04-29 16:29:09 - Job submitted: job_id='662fbcc5e15a79bd9d02c4a7'\n",
      "2024-04-29 16:29:10 - QUEUED\n",
      "2024-04-29 16:29:12 - RUNNING\n",
      "2024-04-29 16:30:48 - COMPLETED\n",
      "2024-04-29 16:30:51 - Dirac allocation balance = 0 s (unmetered)\n"
     ]
    }
   ],
   "source": [
    "response = client.process_job(job_body=body)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The iteration over the samples tests if the sample selected all the members of $X$ exactly once. 100% set coverage indicates that all members were included and a partition was found if no member is included in multiple sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'job_info': {'job_id': '662fbcc5e15a79bd9d02c4a7',\n",
       "  'job_submission': {'problem_config': {'quadratic_unconstrained_binary_optimization': {'qubo_file_id': '662fbc9198263204a365428a'}},\n",
       "   'device_config': {'dirac-1': {'num_samples': 5}}},\n",
       "  'job_status': {'submitted_at_rfc3339nano': '2024-04-29T15:29:09.875Z',\n",
       "   'queued_at_rfc3339nano': '2024-04-29T15:29:09.88Z',\n",
       "   'running_at_rfc3339nano': '2024-04-29T15:29:10.369Z',\n",
       "   'completed_at_rfc3339nano': '2024-04-29T15:30:47.916Z'},\n",
       "  'job_result': {'file_id': '662fbd2798263204a3654295', 'device_usage_s': 9}},\n",
       " 'status': 'COMPLETED',\n",
       " 'results': {'counts': [3, 2],\n",
       "  'energies': [-6, -6],\n",
       "  'solutions': [[0, 0, 1, 0, 0, 0, 1, 1, 0], [0, 1, 0, 0, 0, 1, 0, 0, 1]]}}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_result(results, objective=None):\n",
    "    samples = results[\"solutions\"]\n",
    "    energies = results[\"energies\"]\n",
    "    counts = results[\"counts\"]\n",
    "    for j, sample in enumerate(samples):\n",
    "        print(\"QUBO value:\", energies[j])\n",
    "        print(\"Times sample found\", counts[j])\n",
    "        sample = np.array(sample)\n",
    "        if objective is not None:\n",
    "            print(\"Objective Function value:\", sample.T@objective@sample)\n",
    "        print(\"Selected Sets\")\n",
    "        testX = set()\n",
    "        members = []\n",
    "        for i in range(len(S)):\n",
    "            if sample[i] == 1:\n",
    "                print(f\"S_{i}\", end=\" \")\n",
    "                testX = testX.union(S[i])\n",
    "                members.extend(S[i])\n",
    "        print()\n",
    "        print(f\"Set coverage {100*len(testX)/len(X)}%\")\n",
    "        print(f\"Partition found: {len(testX)==len(members)}\")\n",
    "\n",
    "def get_results(response):\n",
    "    if \"results\" in response and response[\"results\"] is not None:\n",
    "        results = response[\"results\"]\n",
    "    else:\n",
    "        if \"job_info\" in response and \"job_result\" in response[\"job_info\"]:\n",
    "            details = response[\"job_info\"][\"job_result\"]\n",
    "        else:\n",
    "            details = None\n",
    "        raise RuntimeError(f\"Execution failed. See details: {details}\")\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "QUBO value: -6\n",
      "Times sample found 3\n",
      "Selected Sets\n",
      "S_2 S_6 S_7 \n",
      "Set coverage 100.0%\n",
      "Partition found: True\n",
      "QUBO value: -6\n",
      "Times sample found 2\n",
      "Selected Sets\n",
      "S_1 S_5 S_8 \n",
      "Set coverage 100.0%\n",
      "Partition found: True\n"
     ]
    }
   ],
   "source": [
    "results = get_results(response)\n",
    "check_result(results)\n",
    "# save the response\n",
    "penalty_only_response = response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The objective function specified is a linear function. It is built into the diagonal of a matrix, as the linear portion of the penalty function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 4.87248808,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ],\n",
       "       [ 0.        , 28.91096598,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ],\n",
       "       [ 0.        ,  0.        , 72.09663468,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ],\n",
       "       [ 0.        ,  0.        ,  0.        ,  2.16162499,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ],\n",
       "       [ 0.        ,  0.        ,  0.        ,  0.        , 20.59227653,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ],\n",
       "       [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         5.07732567,  0.        ,  0.        ,  0.        ],\n",
       "       [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        , 30.2271894 ,  0.        ,  0.        ],\n",
       "       [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        , 66.39102946,  0.        ],\n",
       "       [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        , 30.81143932]])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "objective = np.diag([W[i] for i in range(len(S))])\n",
    "objective"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "72.09663468312299"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alpha = np.max(objective)\n",
    "alpha"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-211.41741597,    0.        ,   72.09663468,   72.09663468,\n",
       "          72.09663468,   72.09663468,   72.09663468,   72.09663468,\n",
       "         144.19326937],\n",
       "       [   0.        , -187.37893807,   72.09663468,   72.09663468,\n",
       "          72.09663468,    0.        ,    0.        ,  144.19326937,\n",
       "           0.        ],\n",
       "       [  72.09663468,   72.09663468,  -72.09663468,    0.        ,\n",
       "           0.        ,   72.09663468,    0.        ,    0.        ,\n",
       "           0.        ],\n",
       "       [  72.09663468,   72.09663468,    0.        , -142.03164437,\n",
       "           0.        ,    0.        ,   72.09663468,   72.09663468,\n",
       "          72.09663468],\n",
       "       [  72.09663468,   72.09663468,    0.        ,    0.        ,\n",
       "        -123.60099284,    0.        ,    0.        ,  144.19326937,\n",
       "          72.09663468],\n",
       "       [  72.09663468,    0.        ,   72.09663468,    0.        ,\n",
       "           0.        ,  -67.01930901,    0.        ,    0.        ,\n",
       "           0.        ],\n",
       "       [  72.09663468,    0.        ,    0.        ,   72.09663468,\n",
       "           0.        ,    0.        ,  -41.86944529,    0.        ,\n",
       "          72.09663468],\n",
       "       [  72.09663468,  144.19326937,    0.        ,   72.09663468,\n",
       "         144.19326937,    0.        ,    0.        , -149.89887459,\n",
       "          72.09663468],\n",
       "       [ 144.19326937,    0.        ,    0.        ,   72.09663468,\n",
       "          72.09663468,    0.        ,   72.09663468,   72.09663468,\n",
       "        -113.38183004]])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Q = objective + alpha * P\n",
    "Q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "qubo_file = {\"file_name\": \"full-sp-qubo\", \"file_config\": {\"qubo\": {\"data\": Q}}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_id = client.upload_file(file=qubo_file)[\"file_id\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "body = client.build_job_body(job_type=\"sample-qubo\", qubo_file_id=file_id, job_params={\n",
    "    \"device_type\": \"dirac-1\", \"num_samples\": 5\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-04-29 16:36:42 - Dirac allocation balance = 0 s (unmetered)\n",
      "2024-04-29 16:36:43 - Job submitted: job_id='662fbe8be15a79bd9d02c4a9'\n",
      "2024-04-29 16:36:44 - RUNNING\n",
      "2024-04-29 16:38:21 - COMPLETED\n",
      "2024-04-29 16:38:24 - Dirac allocation balance = 0 s (unmetered)\n"
     ]
    }
   ],
   "source": [
    "response = client.process_job(job_body=body)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "QUBO value: -398.7963365771029\n",
      "Times sample found 5\n",
      "Objective Function value: 33.78345405989442\n",
      "Selected Sets\n",
      "S_0 S_1 \n",
      "Set coverage 100.0%\n",
      "Partition found: True\n"
     ]
    }
   ],
   "source": [
    "result = get_results(response)\n",
    "check_result(result, objective)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the samples from the saved response with the objective function included to manually validate the minimization. Note: The QUBO value does not change because it represents the penalty-only QUBO and not including the objective function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "QUBO value: -6\n",
      "Times sample found 3\n",
      "Objective Function value: 168.71485354205467\n",
      "Selected Sets\n",
      "S_2 S_6 S_7 \n",
      "Set coverage 100.0%\n",
      "Partition found: True\n",
      "QUBO value: -6\n",
      "Times sample found 2\n",
      "Objective Function value: 64.79973097220724\n",
      "Selected Sets\n",
      "S_1 S_5 S_8 \n",
      "Set coverage 100.0%\n",
      "Partition found: True\n"
     ]
    }
   ],
   "source": [
    "check_result(get_results(penalty_only_response), objective)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "testing",
   "language": "python",
   "name": "testing"
  },
  "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}