{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Variational Quantum State Diagonalization \n",
    "\n",
    "<em> Copyright (c) 2021 Institute for Quantum Computing, Baidu Inc. All Rights Reserved. </em>\n",
    "\n",
    "## Overview\n",
    "\n",
    "- In this tutorial, we will train a quantum neural network (QNN) through Paddle Quantum to complete the diagonalization of quantum states.\n",
    "\n",
    "- First, import the following packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T09:16:42.117977Z",
     "start_time": "2021-04-30T09:16:38.578847Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy\n",
    "from numpy import diag\n",
    "from numpy import pi as PI\n",
    "import scipy\n",
    "import scipy.stats\n",
    "import paddle\n",
    "from paddle import matmul, trace\n",
    "from paddle_quantum.circuit import UAnsatz\n",
    "from paddle_quantum.utils import dagger"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Background\n",
    "The Variational Quantum State Diagonalization [1-3] aims to output the eigen-spectrum (eigenvalues) of a quantum state. Solving the eigenvalues ​​of quantum states has many applications in quantum computation, such as calculating fidelity and Von Neumann entropy.\n",
    "\n",
    "- Quantum state is usually a mixed state which can be expressed as follows: \n",
    "\n",
    "$$\n",
    "\\rho_{\\text{mixed}} = \\sum_i P_i |\\psi_i\\rangle\\langle\\psi_i|. \\tag{1}\n",
    "$$\n",
    "\n",
    "- As an example, we consider a mixed 2-qubit quantum state with eigen-spectrum $[0.5, 0.3, 0.1, 0.1]$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T09:16:42.139676Z",
     "start_time": "2021-04-30T09:16:42.121581Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.2569+0.j     -0.012 +0.0435j -0.0492-0.0055j -0.0548+0.0682j]\n",
      " [-0.012 -0.0435j  0.2959-0.j      0.1061-0.0713j -0.0392-0.0971j]\n",
      " [-0.0492+0.0055j  0.1061+0.0713j  0.2145-0.j      0.0294-0.1132j]\n",
      " [-0.0548-0.0682j -0.0392+0.0971j  0.0294+0.1132j  0.2327+0.j    ]]\n"
     ]
    }
   ],
   "source": [
    "# Fixed random seed\n",
    "scipy.random.seed(13) \n",
    "V = scipy.stats.unitary_group.rvs(4)  # Randomly generate a unitary matrix\n",
    "D = diag([0.5, 0.3, 0.1, 0.1])        # Input the spectrum of the target state rho\n",
    "V_H = V.conj().T                      # Conjugate transpose operation\n",
    "rho = V @ D @ V_H                     # Generate rho by inverse spectral decomposition\n",
    "print(numpy.around(rho, 4))           # Print quantum state rho"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building a quantum neural network\n",
    "\n",
    "- In this case, we will learn the eigen-spectrum of quantum state $\\rho$ defined above by training a QNN (also known as the parameterized quantum circuit). Here, we provide a predefined 2-qubit quantum circuit.\n",
    "\n",
    "- One can randomly initialize the QNN parameters ${\\bf{\\vec{\\theta }}}$ containing 15 parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T09:16:42.152252Z",
     "start_time": "2021-04-30T09:16:42.144220Z"
    }
   },
   "outputs": [],
   "source": [
    "N = 2           # The width of the quantum neural network\n",
    "SEED = 14       # Fixed random seed\n",
    "THETA_SIZE = 15 # The number of parameters in the quantum neural network\n",
    "\n",
    "def U_theta(theta, N):\n",
    "    \"\"\"\n",
    "    Quantum Neural Network\n",
    "    \"\"\"\n",
    "    # Initialize the quantum neural network according to the number of qubits/network width\n",
    "    cir = UAnsatz(N)\n",
    "    # Call the built-in quantum neural network template\n",
    "    cir.universal_2_qubit_gate(theta, [0, 1])\n",
    "    # Return the circuit of the quantum neural network\n",
    "    return cir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training model and loss function\n",
    "\n",
    "- After setting up the quantum state and the QNN architecture, we will further define the parameters to be trained, loss function, and optimization methods. \n",
    "- The quantum state obtained by acting the quantum neural network $U(\\theta)$ on $\\rho$  is denoted by $\\tilde\\rho$, and we set the loss function to be the inner product of the quantum state $\\sigma$ and $\\tilde\\rho$ where\n",
    "\n",
    "$$\n",
    "\\sigma=0.1 |00\\rangle\\langle 00| + 0.2 |01\\rangle \\langle 01| + 0.3 |10\\rangle \\langle10| + 0.4 |11 \\rangle\\langle 11|, \\tag{2}\n",
    "$$\n",
    "\n",
    "- In specific, the loss function is defined as the state overlap \n",
    "\n",
    "$$\n",
    "\\mathcal{L}(\\boldsymbol{\\theta}) = \\text{Tr}(\\tilde\\rho\\sigma). \\tag{3}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T09:16:42.172611Z",
     "start_time": "2021-04-30T09:16:42.155186Z"
    }
   },
   "outputs": [],
   "source": [
    "# Enter the quantum state sigma \n",
    "sigma = diag([0.1, 0.2, 0.3, 0.4]).astype('complex128')\n",
    "\n",
    "class Net(paddle.nn.Layer):\n",
    "    \"\"\"\n",
    "    Construct the model net\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, shape, rho, sigma, dtype='float64'):\n",
    "        super(Net, self).__init__()\n",
    "        \n",
    "        \n",
    "        \n",
    "        # Convert Numpy array to Tensor supported in Paddle \n",
    "        self.rho = paddle.to_tensor(rho)\n",
    "        self.sigma = paddle.to_tensor(sigma)\n",
    "        \n",
    "        # Initialize the theta parameter list and fill the initial value with the uniform distribution of [0, 2*pi]\n",
    "        self.theta = self.create_parameter(shape=shape, \n",
    "                                           default_initializer=paddle.nn.initializer.Uniform(low=0.0, high=2*PI),\n",
    "                                           dtype=dtype, is_bias=False)\n",
    "\n",
    "    # Define loss function and forward propagation mechanism\n",
    "    def forward(self, N):\n",
    "        \n",
    "        # Apply quantum neural network\n",
    "        cir = U_theta(self.theta, N)\n",
    "        U = cir.U\n",
    "\n",
    "        # rho_tilde is the quantum state U*rho*U^dagger \n",
    "        rho_tilde = matmul(matmul(U, self.rho), dagger(U))\n",
    "\n",
    "        # Calculate the loss function\n",
    "        loss = trace(matmul(self.sigma, rho_tilde))\n",
    "\n",
    "        return paddle.real(loss), rho_tilde, cir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hyper-parameters\n",
    "\n",
    "Before training the quantum neural network, we also need to set up several hyper-parameters, mainly the learning rate LR and the number of iterations ITR. Here we set the learning rate to be LR = 0.1 and the number of iterations to ITR = 50. One can adjust these hyper-parameters accordingly and check how they influence the training performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T09:16:42.468056Z",
     "start_time": "2021-04-30T09:16:42.457529Z"
    }
   },
   "outputs": [],
   "source": [
    "ITR = 50 # Set the total number of iterations of training\n",
    "LR = 0.1 # Set the learning rate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training process\n",
    "\n",
    "- After setting all the parameters of SSVQE model, we need to convert all the data into Tensor in the PaddlePaddle, and then train the quantum neural network.\n",
    "- We used Adam Optimizer in training, and one can also call other optimizers provided in Paddle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T09:16:50.340636Z",
     "start_time": "2021-04-30T09:16:48.778551Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iter: 0 loss: 0.2494\n",
      "iter: 10 loss: 0.1958\n",
      "iter: 20 loss: 0.1843\n",
      "iter: 30 loss: 0.1816\n",
      "iter: 40 loss: 0.1805\n",
      "\n",
      "The trained circuit:\n",
      "--U----X----Rz(1.489)----*-----------------X----U--\n",
      "       |                 |                 |       \n",
      "--U----*----Ry(1.367)----X----Ry(2.749)----*----U--\n",
      "                                                   \n"
     ]
    }
   ],
   "source": [
    "paddle.seed(SEED)\n",
    "\n",
    "# Determine the parameter dimension of the network\n",
    "net = Net(shape=[THETA_SIZE], rho=rho, sigma=sigma)\n",
    "\n",
    "# We use Adam optimizer for better performance\n",
    "# One can change it to SGD or RMSprop.\n",
    "opt = paddle.optimizer.Adam(learning_rate=LR, parameters=net.parameters())\n",
    "\n",
    "# Optimization loop\n",
    "for itr in range(ITR):\n",
    "\n",
    "    # Forward propagation calculates the loss function and returns the estimated energy spectrum\n",
    "    loss, rho_tilde, cir = net(N)\n",
    "    rho_tilde_np = rho_tilde.numpy()\n",
    "\n",
    "    # Back propagation minimizes the loss function\n",
    "    loss.backward()\n",
    "    opt.minimize(loss)\n",
    "    opt.clear_grad()\n",
    "\n",
    "    # Print training results\n",
    "    if itr% 10 == 0:\n",
    "        print('iter:', itr,'loss:','%.4f'% loss.numpy()[0])\n",
    "    if itr == ITR - 1:\n",
    "        print(\"\\nThe trained circuit:\")\n",
    "        print(cir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Benchmarking\n",
    "\n",
    "After 50 iterations, we have completed the diagonalization procedure. The next step is to verify the results of spectral decomposition by printing out $\\tilde{\\rho} = U(\\boldsymbol{\\theta})\\rho U^\\dagger(\\boldsymbol{\\theta})$. One can see the results are very close to what we expect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T09:16:52.099611Z",
     "start_time": "2021-04-30T09:16:52.087176Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The estimated spectrum is: [0.49938069 0.29916354 0.10103808 0.10041768]\n",
      "The target spectrum is: [0.5 0.3 0.1 0.1]\n"
     ]
    }
   ],
   "source": [
    "print(\"The estimated spectrum is:\", numpy.real(numpy.diag(rho_tilde_np)))\n",
    "print(\"The target spectrum is:\", numpy.diag(D))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_______\n",
    "\n",
    "\n",
    "## References\n",
    "\n",
    "[1] Larose, R., Tikku, A., Neel-judy, É. O., Cincio, L. & Coles, P. J. Variational quantum state diagonalization. [npj Quantum Inf. (2019) doi:10.1038/s41534-019-0167-6.](https://www.nature.com/articles/s41534-019-0167-6)\n",
    "\n",
    "[2] Nakanishi, K. M., Mitarai, K. & Fujii, K. Subspace-search variational quantum eigensolver for excited states. [Phys. Rev. Res. 1, 033062 (2019).](https://journals.aps.org/prresearch/pdf/10.1103/PhysRevResearch.1.033062)\n",
    "\n",
    "[3] Cerezo, M., Sharma, K., Arrasmith, A. & Coles, P. J. Variational Quantum State Eigensolver. [arXiv:2004.01372 (2020).](https://arxiv.org/pdf/2004.01372.pdf)"
   ]
  }
 ],
 "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.8.3"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}