{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "a6f294e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.nn.init as init\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import networkx as nx\n",
    "\n",
    "from pyqtorch.core.circuit import QuantumCircuit\n",
    "from pyqtorch.ansatz import AlternateLayerAnsatz\n",
    "from pyqtorch.embedding import SingleLayerEncoding\n",
    "from pyqtorch.core.operation import Z, RX"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb4585dc",
   "metadata": {},
   "source": [
    "## Solve MIS for QAOA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "3c4e17c3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 40484 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "np.random.seed(0)\n",
    "n_nodes = 5\n",
    "\n",
    "graph = nx.gnp_random_graph(n_nodes, .25, seed=42)\n",
    "nx.draw(graph, with_labels=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "a1b223ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyqtorch.matrices import generate_ising_from_graph, sum_N\n",
    "ising_matrix = generate_ising_from_graph(graph, type_ising='N')\n",
    "ising_cost = 1.2*ising_matrix - sum_N(n_nodes)\n",
    "\n",
    "ising_matrix = ising_matrix.reshape([2] * n_nodes + [1])\n",
    "ising_cost = ising_cost.reshape([2] * n_nodes + [1])\n",
    "\n",
    "\n",
    "class MIS(QuantumCircuit):\n",
    "    def __init__(self, n_qubits, n_layers):\n",
    "        super().__init__(n_qubits)\n",
    "        self.beta = nn.Parameter(torch.empty(n_layers,))\n",
    "        self.gamma = nn.Parameter(torch.empty(n_layers,))\n",
    "        self.reset_parameters()\n",
    "        \n",
    "    def reset_parameters(self):\n",
    "        init.uniform_(self.beta, -2 * np.pi, 2 * np.pi)\n",
    "        init.uniform_(self.gamma, -2 * np.pi, 2 * np.pi)\n",
    "        \n",
    "    def forward(self, return_cost=False):\n",
    "        state = self.uniform_state()\n",
    "        for b, g in zip(self.beta, self.gamma):\n",
    "            state = state * torch.exp(-1j * g * ising_matrix)\n",
    "            for i in range(self.n_qubits):\n",
    "                state = RX(b, state, [i], self.n_qubits)\n",
    "        if return_cost:\n",
    "            return torch.real(torch.sum(torch.abs(state)**2 * ising_cost))\n",
    "        else:\n",
    "            state = state.reshape((2**self.n_qubits,))\n",
    "            return torch.abs(state)**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "c84f30d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0 | Loss -0.5000000000000002\n",
      "Epoch 50 | Loss -0.5000000000000002\n"
     ]
    }
   ],
   "source": [
    "model = MIS(n_nodes, 20)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=.02)\n",
    "epochs = 100\n",
    "\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    optimizer.zero_grad()\n",
    "    loss = model(True)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    if epoch%50 == 0:\n",
    "        print(f\"Epoch {epoch} | Loss {loss}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "4f7a6f49",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0011111001'"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prob = model()\n",
    "mis = torch.argmax(prob)\n",
    "format(mis, '010b')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "7bf6d8b0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(249)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "bfd64d73",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1'"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"{0:b}\".format(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "15ce4608",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(6.+0.j, dtype=torch.complex128)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sum_N(n_nodes)[mis]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cb69112",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 ('pynn')",
   "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.10.4"
  },
  "vscode": {
   "interpreter": {
    "hash": "6f77e4d31a6b2b6cebe1094e9c25b1f29bee03e4b42142e04de4cbf5445e2748"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
