QAOA_En.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## English | [简体中文](./QAOA.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preparation\n",
    "\n",
    "This document provides an interactive experience to show how the quantum approximate optimization algorithm (QAOA) [1] works in the Paddle Quantum. \n",
    "\n",
    "To get started, let us import some necessary libraries and functions："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "from paddle import fluid\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import networkx as nx\n",
    "\n",
    "from numpy import matmul as np_matmul\n",
    "from paddle.complex import matmul as pp_matmul\n",
    "from paddle.complex import transpose\n",
    "from paddle_quantum.circuit import UAnsatz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Background\n",
    "\n",
    "QAOA is one of quantum algorithms which can be implemented on near-term quantum processors, also called as noisy intermediate-scale quantum (NISQ) processors, and may have wide applications in solving hard computational problems.  For example, it could be applied to tackle a large family of optimization problems, named as the quadratic unconstrained binary optimization (QUBO) which is ubiquitous in the computer science and operation research. Basically, this class can be modeled with the form of\n",
    "\n",
    "$$F=\\max_{z_i\\in\\{-1,1\\}} \\sum q_{ij}(1-z_iz_j)=-\\min_{z_i\\in\\{-1,1\\}} \\sum q_{ij}z_iz_j+ \\sum q_{ij} $$\n",
    "\n",
    "\n",
    "where $z_i$s are  binary parameters and coefficients $q_{ij}$ refer to the weight associated to $x_i x_j$. Indeed, it is usually extremely difficult for classical computers to give the exact optimal solution, while QAOA provides an alternative approach which may have a speedup advantage over classical ones to solve these hard problems.\n",
    "\n",
    "QAOA works as follows: The above optimization problem is first mapped to another problem of finding the ground energy or/and the corresponding ground state for a complex many-body Hamiltonian, e.g., the well-known Ising model or spin-glass model in many-body physics. Essentially, it is transformed to find the smallest eigenvalue and the corresponding eigenvector(s) for a real diagonal matrix $H$. Then, QAOA designates a specific routine with adjustable parameters to approximately find the best solution. Moreover, to accomplish the task, these parameters could be updated via some rules set by fast classical algorithms, such as gradient-free or gradient-based methods. Thus, it is also a quantum-classical hybrid algorithm just as the variational quantum eigensolver (VQE).\n",
    "\n",
    "Here, the Max-Cut problem in graph theory is used to explain how QAOA works."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# Example\n",
    "\n",
    "##  1. Max-Cut problem\n",
    "\n",
    "Given a graph $G$ composed of $N$ nodes and $M$ edges, the problem is to find a cut protocol which divides the node set into two complementary subsets $S$ and $S^\\prime$ such that the number of edges between these sets is as large as possible. For example, consider the ring case with four nodes as shown in the figure.\n",
    "\n",
    " ![ring4.png](https://release-data.cdn.bcebos.com/PIC%2FMaxCut.png) \n",
    " \n",
    "\n",
    "Thus, given a cut protocol, if the node $i$ belongs to the set $S$, then it is assigned to $z_i =1$, while $z_j= -1$ for $j \\in S^\\prime$. Then, for any edge connecting nodes $i$ and $j$, if both nodes are in the same set $S$ or $S^\\prime$, then there is $z_iz_j=1$; otherwise, $z_iz_j=-1$. Hence, the cut problem can be formulated as the optimization problem \n",
    "\n",
    "$$ F=\\min_{z_i\\in\\{-1, 1\\}} z_1 z_2+z_2z_3+z_3z_4+z_4z_1.$$\n",
    "\n",
    "Here, the weight $q_{ij}s$ are set to $1$ for all edges. Indeed, any feasible solution to the above problem can be described by a bitstring $ \\boldsymbol{z}=z_1z_2z_3z_4 \\in \\{-1, 1\\}^4$. Moreover, we need to search over all possible bitstrings of $2^N$ to find its optimal solution, which becomes computionally hard for classical algorithms.\n",
    "\n",
    "Two methods are provided to pre-process this optimization problem, i.e., to input the given graph with/without weights: \n",
    "\n",
    "- Method 1 generates the graph via its full description of nodes and edges,\n",
    "- Method 2 specifies the graph via its adjacency matrix.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_graph(N, GRAPHMETHOD):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        N: number of nodes (vertices) in the graph\n",
    "        METHOD: choose which method to generate a graph\n",
    "    Return:\n",
    "        the specific graph and its adjacency matrix\n",
    "    \"\"\"\n",
    "    # Method 1 generates a graph by self-definition\n",
    "    if GRAPHMETHOD == 1:\n",
    "        print(\"Method 1 generates the graph from self-definition using EDGE description\")\n",
    "        graph = nx.Graph()\n",
    "        graph_nodelist=range(N)\n",
    "        graph.add_edges_from([(0, 1), (1, 2), (2, 3), (3, 0)])\n",
    "        graph_adjacency = nx.to_numpy_matrix(graph, nodelist=graph_nodelist)\n",
    "    # Method 2 generates a graph by using its adjacency matrix directly\n",
    "    elif GRAPHMETHOD == 2:\n",
    "        print(\"Method 2 generates the graph from networks using adjacency matrix\")\n",
    "        graph_adjacency = np.array([[0, 1, 0, 1], [1, 0, 1, 0], [0, 1, 0, 1], [1, 0, 1, 0]])\n",
    "        graph = nx.Graph(graph_adjacency)\n",
    "    else:\n",
    "        print(\"Method doesn't exist \")\n",
    "\n",
    "    return graph, graph_adjacency"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, Method 1 is used to process and then visualize the given graph. Note that the node label starts from $0$ to $ N-1$ in both methods for an $N$-node graph. \n",
    "\n",
    "Here, we need to specify:\n",
    "\n",
    "- number of nodes: $N=4$\n",
    "- which method to preprocess the graph: GRAPHMETHOD = 1 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# number of qubits or number of nodes in the graph\n",
    "N=4  \n",
    "classical_graph, classical_graph_adjacency= generate_graph(N, GRAPHMETHOD=1)\n",
    "print(classical_graph_adjacency)\n",
    "\n",
    "pos = nx.circular_layout(classical_graph)\n",
    "nx.draw(classical_graph, pos, width=4, with_labels=True, font_weight='bold')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Encoding\n",
    "\n",
    "This step encodes the classical optimization problem to its quantum version. Using the transformation $z=1\\rightarrow |0\\rangle = \\begin{bmatrix}1 \\\\ 0\\end{bmatrix}$ and $z=-1\\rightarrow |1\\rangle = \\begin{bmatrix}0 \\\\ 1\\end{bmatrix}$, the binary parameter $z$ is encoded as the eigenvalues of the Pauli-Z operator acting on the single qubit, i.e., $z\\rightarrow Z=\\begin{bmatrix} 1 & 0\\\\ 0 & -1\\end{bmatrix}$. It yields that the objective function in the classical optimization problem is transformed to the Hamiltonian\n",
    "\n",
    "$$H_c= Z_1Z_2+Z_2Z_3+Z_3Z_4+Z_4Z_1.$$\n",
    "\n",
    "Here, for simplicity $Z_iZ_{j}$ stands for the tensor product $Z_i\\otimes Z_j$ which represents that Pauli-Z operator acts on each qubit $i, j$ and the identity operation is imposed on the rest. And the Max-Cut problem is mapped to the following quantum optimization problem\n",
    "\n",
    "$$ F=\\min_{|\\psi\\rangle}\\, \\langle \\psi| H_c |\\psi\\rangle$$\n",
    "\n",
    "where the state vector $|\\psi\\rangle$ describes a $4$-dimensional complex vector which is normalized to $1$, and $\\langle \\psi|$ is its conjugate transpose form. It is equivalent to find the smallest eigenvalue $F$ and the corresponding eigenstate(s) for the matrix $H_c$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def H_generator(N, adjacency_matrix):\n",
    "    \"\"\"\n",
    "    This function maps the given graph via its adjacency matrix to the corresponding Hamiltiona H_c.\n",
    "    \n",
    "    Args:\n",
    "        N: number of qubits, or number of nodes in the graph, or number of parameters in the classical problem\n",
    "        adjacency_matrix:  the adjacency matrix generated from the graph encoding the classical problem\n",
    "    Return:\n",
    "        H_graph: the problem-based Hamiltonian H generated from the graph_adjacency matrix for the given graph\n",
    "        H_graph_diag: the real part of the problem-based Hamiltonian H_graph\n",
    "    \"\"\"\n",
    "\n",
    "    sigma_Z = np.array([[1, 0], [0, -1]])\n",
    "    H = np.zeros([2 ** N, 2 ** N])\n",
    "    # Generate the Hamiltonian H_c from the graph via its adjacency matrix\n",
    "    for row in range(N):\n",
    "        for col in range(N):\n",
    "            if abs(adjacency_matrix[N - row - 1, N - col - 1]) and row < col:\n",
    "                identity_1 = np.diag(np.ones([2 ** row]))\n",
    "                identity_2 = np.diag(np.ones([2 ** (col - row - 1)]))\n",
    "                identity_3 = np.diag(np.ones([2 ** (N - col - 1)]))\n",
    "                H += adjacency_matrix[N - row - 1, N - col - 1] * np.kron(\n",
    "                    np.kron(np.kron(np.kron(identity_1, sigma_Z), identity_2), sigma_Z),\n",
    "                    identity_3,\n",
    "                )\n",
    "\n",
    "    H_graph = H.astype(\"complex64\")\n",
    "    H_graph_diag = np.diag(H_graph).real\n",
    "\n",
    "    return H_graph, H_graph_diag"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The explicit form of the matrix $H_c$, including its maximal and minimal eigenvalues, can be imported, which later could be used to benchmark the performance of QAOA. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "_, H_problem_diag = H_generator(N, classical_graph_adjacency)\n",
    "\n",
    "H_graph_max = np.max(H_problem_diag)\n",
    "H_graph_min = np.min(H_problem_diag)\n",
    "\n",
    "print(H_problem_diag)\n",
    "print('H_max:', H_graph_max, '  H_min:', H_graph_min)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Building\n",
    "\n",
    "This part is to build up the parameterized quantum circuit of QAOA to perform the computation process. Particularly, the QAOA circuit is constructed by alternatively placing two parameterized modules\n",
    "\n",
    "$$U_x(\\beta_P)U_c(\\gamma_P)\\dots U_x(\\beta_1)U_c(\\gamma_1)$$\n",
    "\n",
    "where $P$ is the number of layers to place these two modules. Particularly, one is governed by the encoding matrix $H_c$ via the unitary transformation\n",
    "\n",
    "$$U_c(\\gamma)=e^{-i \\gamma H_c }$$\n",
    "\n",
    "where $i$ is the imaginary unit, and $\\gamma\\in [0, \\pi]$ is to be optimized. The other one is \n",
    "\n",
    "$$U_x(\\beta)=e^{-i \\beta H_x },$$\n",
    "\n",
    "where $\\beta\\in [0, \\pi]$ and the driving Hamiltonian or matrix $H_x$ adimits an explicit form of\n",
    "\n",
    "$$H_x =X_1+X_2+X_3+X_4 $$\n",
    "\n",
    "where the operator $X=\\begin{bmatrix} 0 & 1\\\\ 1& 0\\end{bmatrix}$ defines the Pauli-X operation acting on the qubit.\n",
    "\n",
    "\n",
    "\n",
    "Further, each module in the QAOA circuit can be decomposed into a series of operations acting on single qubits and two qubits. In particular, the first has the decomposition of $U_c(\\gamma)=\\prod_{(i, j)}e^{-i \\gamma Z_i\\otimes Z_j }$ while there is $U_x(\\beta)=\\prod_{i}e^{-i \\beta X_i}$ for the second. This is illustrated in the following figure.\n",
    "\n",
    " ![Circ](https://release-data.cdn.bcebos.com/PIC%2FQAOACir.png) \n",
    "\n",
    "Then, based on\n",
    "\n",
    "- initial state of QAOA circuits \n",
    "- adjacency matrix describing the graph\n",
    "- number of qubits\n",
    "- number of layers\n",
    "\n",
    "we are able to construct the standard QAOA circuit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "def circuit_QAOA(theta, input_state, adjacency_matrix, N, P):\n",
    "    \"\"\"\n",
    "    This function constructs the parameterized QAOA circuit which is composed of P layers of two blocks:\n",
    "    one block is U_theta[layer][0] based on the problem Hamiltonian H which encodes the classical problem,\n",
    "    and the other is U_theta[layer][1] constructed from the driving Hamiltonian describing the rotation around Pauli X\n",
    "    acting on each qubit. It finally outputs the final state of the QAOA circuit.\n",
    "    \n",
    "    Args:\n",
    "         theta: parameters to be optimized in the QAOA circuit\n",
    "         input_state: initial state of the QAOA circuit which usually is the uniform superposition of 2^N bit-strings \n",
    "         in the computataional basis\n",
    "         adjacency_matrix:  the adjacency matrix of the graph encoding the classical problem\n",
    "         N: number of qubits, or equivalently, the number of parameters in the original classical problem\n",
    "         P: number of layers of two blocks in the QAOA circuit\n",
    "    Returns:\n",
    "        the final state of the QAOA circuit: cir.state\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    cir = UAnsatz(N, input_state=input_state)\n",
    "    \n",
    "    # This loop defines the QAOA circuit with P layers of two blocks\n",
    "    for layer in range(P):\n",
    "        # The second and third loops construct the first block U_theta[layer][0] which involves two-qubit operation\n",
    "        #  e^{-i\\beta Z_iZ_j} acting on a pair of qubits or nodes i and j in the circuit in each layer.\n",
    "        for row in range(N):\n",
    "            for col in range(N):\n",
    "                if abs(adjacency_matrix[row, col]) and row < col:\n",
    "                    cir.cnot([row + 1, col + 1])\n",
    "                    cir.rz(\n",
    "                        theta=theta[layer][0]\n",
    "                        * adjacency_matrix[row, col],\n",
    "                        which_qubit=col + 1,\n",
    "                    )\n",
    "                    cir.cnot([row + 1, col + 1])\n",
    "        # This loop constructs the second block U_theta only involving the single-qubit operation e^{-i\\beta X}.\n",
    "        for i in range(1, N + 1):\n",
    "            cir.rx(theta=theta[layer][1], which_qubit=i)\n",
    "\n",
    "    return cir.state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Indeed, the QAOA circuit could be extended to other structures by replacing the modules in the above standard circuit to improve QAOA performance. Here, we provide one candidate extension in which the Pauli-X rotation $R_x(\\beta) $ on each qubit in the driving matrix $H_x$ is replaced by an arbitrary rotation described by $R_z(\\beta_1)R_x(\\beta_2)R_z(\\beta_3)$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "def circuit_extend_QAOA(theta, input_state, adjacency_matrix, N, P):\n",
    "    \"\"\"\n",
    "    This is an extended version of the QAOA circuit, and the main difference is U_theta[layer]([1]-[3]) constructed\n",
    "    from the driving Hamiltonian describing the rotation around an arbitrary direction on each qubit.\n",
    "\n",
    "    Args:\n",
    "        theta: parameters to be optimized in the QAOA circuit\n",
    "        input_state: input state of the QAOA circuit which usually is the uniform superposition of 2^N bit-strings\n",
    "                     in the computational basis\n",
    "        adjacency_matrix:  the adjacency matrix of the problem graph encoding the original problem\n",
    "        N: number of qubits, or equivalently, the number of parameters in the original classical problem\n",
    "        P: number of layers of two blocks in the QAOA circuit\n",
    "    Returns:\n",
    "        final state of the QAOA circuit: cir.state\n",
    "\n",
    "    Note: If this U_extend_theta function is used to construct QAOA circuit, then we need to change the parameter layer\n",
    "           in the Net function defined below from the Net(shape=[D, 2]) for U_theta function to Net(shape=[D, 4])\n",
    "           because the number of parameters doubles in each layer in this QAOA circuit.\n",
    "    \"\"\"\n",
    "    cir = UAnsatz(N, input_state=input_state)\n",
    "\n",
    "    for layer in range(P):\n",
    "        for row in range(N):\n",
    "            for col in range(N):\n",
    "                if abs(adjacency_matrix[row, col]) and row < col:\n",
    "                    cir.cnot([row + 1, col + 1])\n",
    "                    cir.rz(\n",
    "                        theta=theta[layer][0]\n",
    "                        * adjacency_matrix[row, col],\n",
    "                        which_qubit=col + 1,\n",
    "                    )\n",
    "                    cir.cnot([row + 1, col + 1])\n",
    "\n",
    "        for i in range(1, N + 1):\n",
    "            cir.rz(theta=theta[layer][1], which_qubit=i)\n",
    "            cir.rx(theta=theta[layer][2], which_qubit=i)\n",
    "            cir.rz(theta=theta[layer][3], which_qubit=i)\n",
    "\n",
    "    return cir.state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, the QAOA circuit outputs\n",
    "\n",
    "$$|\\psi(\\boldsymbol{\\beta},\\boldsymbol{\\gamma}, P)\\rangle=U_x(\\beta_P)U_c(\\gamma_P)\\dots U_x(\\beta_1)U_c(\\gamma_1)|+\\rangle_1\\dots|+\\rangle_N$$\n",
    "\n",
    "where each qubit is initialized as the superposition state $|+\\rangle=\\frac{1}{\\sqrt{2}}\\left(|0\\rangle+|1\\rangle\\right)$. And we are able to obtain the loss function for the QAOA circuit\n",
    "\n",
    "$$F_P=\\min_{\\boldsymbol{\\beta},\\boldsymbol{\\gamma}} \\langle \\psi(\\boldsymbol{\\beta},\\boldsymbol{\\gamma}, P)| H_c|\\psi(\\boldsymbol{\\beta},\\boldsymbol{\\gamma}, P)\\rangle.$$\n",
    "\n",
    "Additionally, we may tend to fast classical algorithms to update the parameter vectors $\\boldsymbol{\\beta},\\boldsymbol{\\gamma}$ to achieve the optimal value for the above quantum optimization problem. \n",
    "\n",
    "In Paddle Quantum, this process is accomplished in the Net function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "class Net(fluid.dygraph.Layer):\n",
    "    \"\"\"\n",
    "    It constructs the net for QAOA which combines the  QAOA circuit with the classical optimizer which sets rules\n",
    "    to update parameters described by theta introduced in the QAOA circuit.\n",
    "\n",
    "    \"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        shape,\n",
    "        param_attr=fluid.initializer.Uniform(low=0.0, high=np.pi, seed=1024),\n",
    "        dtype=\"float32\",\n",
    "    ):\n",
    "        super(Net, self).__init__()\n",
    "\n",
    "        self.theta = self.create_parameter(\n",
    "            shape=shape, attr=param_attr, dtype=dtype, is_bias=False\n",
    "        )\n",
    "\n",
    "    def forward(self, input_state, adjacency_matrix, out_state_store, N, P, METHOD):\n",
    "        \"\"\"\n",
    "        This function constructs the loss function for the QAOA circuit.\n",
    "\n",
    "        Args:\n",
    "            self: the free parameters to be optimized in the QAOA circuit and defined in the above function\n",
    "            input_state: initial state of the QAOA circuit which usually is the uniform superposition of 2^N bit-strings\n",
    "                         in the computational basis $|0\\rangle, |1\\rangle$\n",
    "            adjacency_matrix: the adjacency matrix generated from the graph encoding the classical problem\n",
    "            out_state_store: the output state of the QAOA circuit\n",
    "            N: number of qubits\n",
    "            P: number of layers\n",
    "            METHOD: which version of QAOA is chosen to solve the problem, i.e., standard version labeled by 1 or\n",
    "            extended version by 2.\n",
    "        Returns:\n",
    "            The loss function for the parameterized QAOA circuit.\n",
    "        \"\"\"\n",
    "        \n",
    "        # Generate the problem_based quantum Hamiltonian H_problem based on the classical problem in paddle\n",
    "        H, _ = H_generator(N, adjacency_matrix)\n",
    "        H_problem = fluid.dygraph.to_variable(H)\n",
    "\n",
    "        # The standard QAOA circuit: the function circuit_QAOA is used to construct the circuit, indexed by METHOD 1.\n",
    "        if METHOD == 1:\n",
    "            out_state = circuit_QAOA(self.theta, input_state, adjacency_matrix, N, P)\n",
    "        # The extended QAOA circuit: the function circuit_extend_QAOA is used to construct the net, indexed by METHOD 2.\n",
    "        elif METHOD == 2:\n",
    "            out_state = circuit_extend_QAOA(self.theta, input_state, adjacency_matrix, N, P)\n",
    "        else:\n",
    "            raise ValueError(\"Wrong method called!\")\n",
    "\n",
    "        out_state_store.append(out_state.numpy())\n",
    "        loss = pp_matmul(\n",
    "            pp_matmul(out_state, H_problem),\n",
    "            transpose(\n",
    "                fluid.framework.ComplexVariable(out_state.real, -out_state.imag),\n",
    "                perm=[1, 0],\n",
    "            ),\n",
    "        )\n",
    "\n",
    "        return loss.real\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Training\n",
    "\n",
    "In this part, the QAOA circuit is trained to find the \"optimal\" solution to the optimization problem.\n",
    "\n",
    "First, let us specify some parameters:\n",
    "\n",
    "- number of qubits: N\n",
    "- number of layes: P\n",
    "- iteration steps: ITR\n",
    "- learning rate: LR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "N = 4 # number of qubits, or number of nodes in the graph\n",
    "P = 4 # number of layers \n",
    "ITR = 120  # number of iteration steps\n",
    "LR = 0.1  # learning rate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, with the following inputs\n",
    "\n",
    "- initial state: each qubit is initialized as $\\frac{1}{\\sqrt{2}}\\left(|0\\rangle+|1\\rangle\\right)$\n",
    "- Standard QAOA circuit (METHOD = 1) or Extended QAOA (METHOD = 2) \n",
    "- Classical optimizer: Adam optimizer\n",
    "\n",
    "we are able to train the whole net："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "def Paddle_QAOA(classical_graph_adjacency, N, P, METHOD, ITR, LR):\n",
    "    \"\"\"\n",
    "    This is the core function to run QAOA.\n",
    "\n",
    "     Args:\n",
    "         classical_graph_adjacency: adjacency matrix to describe the graph which encodes the classical problem\n",
    "         N: number of qubits (default value N=4)\n",
    "         P: number of layers of blocks in the QAOA circuit (default value P=4)\n",
    "         METHOD: which version of the QAOA circuit is used: 1, standard circuit (default); 2, extended circuit\n",
    "         ITR: number of iteration steps for QAOA (default value ITR=120)\n",
    "         LR: learning rate for the gradient-based optimization method (default value LR=0.1)\n",
    "     Returns:\n",
    "         optimized parameters theta and the bitstrings sampled from the output state with maximal probability\n",
    "    \"\"\"\n",
    "    out_state_store = []\n",
    "    with fluid.dygraph.guard():\n",
    "        # Preparing the initial state\n",
    "        _initial_state = np.ones([1, 2 ** N]).astype(\"complex64\") / np.sqrt(2 ** N)\n",
    "        initial_state = fluid.dygraph.to_variable(_initial_state)\n",
    "\n",
    "        # Construct the net or QAOA circuits based on the standard modules\n",
    "        if METHOD == 1:\n",
    "            net = Net(shape=[P, 2])\n",
    "        # Construct the net or QAOA circuits based on the extended modules\n",
    "        elif METHOD == 2:\n",
    "            net = Net(shape=[P, 4])\n",
    "        else:\n",
    "            raise ValueError(\"Wrong method called!\")\n",
    "\n",
    "        # Classical optimizer\n",
    "        opt = fluid.optimizer.AdamOptimizer(learning_rate=LR, parameter_list=net.parameters())\n",
    "\n",
    "        # Gradient descent loop\n",
    "        summary_iter, summary_loss = [], []\n",
    "        for itr in range(1, ITR + 1):\n",
    "            loss = net(\n",
    "                initial_state, classical_graph_adjacency, out_state_store, N, P, METHOD\n",
    "            )\n",
    "            loss.backward()\n",
    "            opt.minimize(loss)\n",
    "            net.clear_gradients()\n",
    "\n",
    "            print(\"iter:\", itr, \"  loss:\", \"%.4f\" % loss.numpy())\n",
    "            summary_loss.append(loss[0][0].numpy())\n",
    "            summary_iter.append(itr)\n",
    "\n",
    "        theta_opt = net.parameters()[0].numpy()\n",
    "        print(theta_opt)\n",
    "\n",
    "        os.makedirs(\"output\", exist_ok=True)\n",
    "        np.savez(\"./output/summary_data\", iter=summary_iter, energy=summary_loss)\n",
    "\n",
    "    # Output the measurement probability distribution which is sampled from the output state of optimized QAOA circuit.\n",
    "    prob_measure = np.zeros([1, 2 ** N]).astype(\"complex\")\n",
    "\n",
    "    rho_out = out_state_store[-1]\n",
    "    rho_out = np_matmul(np.conjugate(rho_out).T, rho_out).astype(\"complex\")\n",
    "\n",
    "    for index in range(0, 2 ** N):\n",
    "        comput_basis = np.zeros([1, 2 ** N])\n",
    "        comput_basis[0][index] = 1\n",
    "        prob_measure[0][index] = np.real(np_matmul(np_matmul(comput_basis, rho_out), comput_basis.T))\n",
    "\n",
    "    return prob_measure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After the completion of training, the QAOA outputs the results, including the optimal parameters $\\boldsymbol{\\beta}^*$ and $\\boldsymbol{\\gamma}^*$. By contrast, its performance can be evaluated with the true value of the optimization problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "classical_graph, classical_graph_adjacency = generate_graph(N, 1)\n",
    "\n",
    "prob_measure_dis = Paddle_QAOA(classical_graph_adjacency, N =4, P=4, METHOD=1, ITR=120, LR=0.1)\n",
    "\n",
    "# Load the data of QAOA\n",
    "x1 = np.load('./output/summary_data.npz')\n",
    "\n",
    "H_min = np.ones([len(x1['iter'])]) * H_graph_min\n",
    "\n",
    "# Plot it\n",
    "\n",
    "loss_QAOA, = plt.plot(x1['iter'], x1['energy'], \\\n",
    "                                        alpha=0.7, marker='', linestyle=\"--\", linewidth=2, color='m')\n",
    "benchmark, = plt.plot(x1['iter'], H_min, alpha=0.7, marker='', linestyle=\":\", linewidth=2, color='b')\n",
    "plt.xlabel('Number of iteration')\n",
    "plt.ylabel('Performance of the loss function for QAOA')\n",
    "\n",
    "plt.legend(handles=[\n",
    "    loss_QAOA,\n",
    "    benchmark\n",
    "],\n",
    "    labels=[\n",
    "            r'Loss function $\\left\\langle {\\psi \\left( {\\bf{\\theta }} \\right)} '\n",
    "            r'\\right|H\\left| {\\psi \\left( {\\bf{\\theta }} \\right)} \\right\\rangle $',\n",
    "            'The benchmark result',\n",
    "    ], loc='best')\n",
    "\n",
    "# Show the picture\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## 5. Decoding \n",
    "\n",
    "However, the output of optimized QAOA circuits \n",
    "\n",
    "$$|\\psi(\\boldsymbol{\\beta}^*,\\boldsymbol{\\gamma}^*, P)\\rangle=\\sum_{i=1}^{2^4}\\lambda_i |\\boldsymbol{x}_i\\rangle$$\n",
    "\n",
    "does not give us the answer to the Max-Cut problem directly. Instead, each bitstring $\\boldsymbol{x}_i=x_1x_2x_3 x_4\\in \\{0, 1\\}^4$ in the state $|\\psi(\\boldsymbol{\\beta}^*,\\boldsymbol{\\gamma}^*, P)$ represents a possible classical solution. Thus, we need to decode the ouptut of QAOA circuits. \n",
    "\n",
    "The task of decoding quantum answer can be accomplished via measurement. Given the output state, the measurement statistics for each bitstring obeys the probability distribution\n",
    "\n",
    "$$ p(\\boldsymbol{x})=|\\langle \\boldsymbol{x}|\\psi(\\boldsymbol{\\beta}^*,\\boldsymbol{\\gamma}^*,P)\\rangle|^2.$$\n",
    "              \n",
    "And this distribution is plotted using the following function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "prob_measure = prob_measure_dis.flatten()\n",
    "\n",
    "pos = nx.circular_layout(classical_graph)\n",
    "# when N is large, it is not suggested to plot this figure\n",
    "name_list = [np.binary_repr(index, width=N) for index in range(0, 2 ** N)]\n",
    "plt.bar(\n",
    "        range(len(np.real(prob_measure))),\n",
    "        np.real(prob_measure),\n",
    "        width=0.7,\n",
    "        tick_label=name_list,\n",
    "    )\n",
    "plt.xticks(rotation=90)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, using the relation $|x \\rangle\\rightarrow z=2x-1$, we are able to obtain a classical answer from the quantum state. Specifically, assume that $z_i=-1$ for $i \\in S$ and $z_j=-1$ for $j \\in S^\\prime$. Thus, one bistring sampled from the output state of QAOA corresponds to one feasible cut to the given graph. And it is highly possible that the higher probability the bitstring is, the more likely it gives rise to the max cut protocol.\n",
    "\n",
    "The bistring with the largest probability is picked up, and then mapped back to solution to the Max-Cut problem ：\n",
    "\n",
    "- the node set $S$ is in blue color\n",
    "- the node set $S^\\prime$ is in red color\n",
    "- the dashed lines represent the cut edges "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "# Find the position of max value in the measure_prob_distribution\n",
    "max_prob_pos_list = np.where(prob_measure == max(prob_measure))\n",
    "# Store the max value from ndarray to list\n",
    "max_prob_list = max_prob_pos_list[0].tolist()\n",
    "# Change it to the  binary format\n",
    "solution_list = [np.binary_repr(index, width=N) for index in max_prob_list]\n",
    "print(\"The output bitstring:\", solution_list)\n",
    "\n",
    "# Draw the graph representing the first bitstring in the solution_list to the MaxCut-like problem\n",
    "head_bitstring = solution_list[0]\n",
    "\n",
    "node_cut = [\"blue\" if head_bitstring[node] == \"1\" else \"red\" for node in classical_graph]\n",
    "\n",
    "edge_cut = [\n",
    "    \"solid\" if head_bitstring[node_row] == head_bitstring[node_col] else \"dashed\"\n",
    "    for node_row, node_col in classical_graph.edges()\n",
    "    ]\n",
    "nx.draw(\n",
    "        classical_graph,\n",
    "        pos,\n",
    "        node_color=node_cut,\n",
    "        style=edge_cut,\n",
    "        width=4,\n",
    "        with_labels=True,\n",
    "        font_weight=\"bold\",\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "source": [
    "# References\n",
    "\n",
    "[1] E. Farhi, J. Goldstone, and S. Gutman. 2014. A quantum approximate optimization algorithm. arXiv:1411.4028 "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.10"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}