{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 概览\n",
    "- 在这个案例中，我们将展示如何通过Paddle Quantum训练量子神经网络来将量子态进行对角化。\n",
    "\n",
    "- 首先，让我们通过下面几行代码引入必要的library和package。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import numpy\n",
    "from numpy import diag\n",
    "import scipy\n",
    "from paddle import fluid\n",
    "from paddle_quantum.circuit import UAnsatz\n",
    "from paddle.complex import matmul, trace, transpose"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "$\\newcommand{\\ket}[1]{|{#1}\\rangle}$\n",
    "$\\newcommand{\\bra}[1]{\\langle{#1}|}$\n",
    "\n",
    "## 背景\n",
    "量子态对角化算法（VQSD）[1-3] 目标是输出一个量子态的特征谱，即其所有特征值。求解量子态的特征值在量子计算中有着诸多应用，比如可以用于计算保真度和冯诺依曼熵，也可以用于主成分分析。\n",
    "- 量子态通常是一个混合态，表示如下 $\\rho_{\\text{mixed}} = \\sum_i P_i \\ket{\\psi_i}\\bra{\\psi_i}$\n",
    "- 作为一个简单的例子，我们考虑一个2量子位的量子态，它的特征谱为 $(0.5, 0.3, 0.1, 0.1)$, 我们先通过随机作用一个酉矩阵来生成具有这样特征谱的量子态。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "scipy.random.seed(1)\n",
    "V = scipy.stats.unitary_group.rvs(4) # 随机生成一个酉矩阵\n",
    "D = diag([0.5, 0.3, 0.1, 0.1]) # 输入目标态rho的谱\n",
    "V_H = V.conj().T \n",
    "rho = V @ D @ V_H # 生成 rho\n",
    "print(rho) # 打印量子态 rho\n",
    "rho = rho.astype('complex64')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 搭建量子神经网络\n",
    "- 在这个案例中，我们将通过训练量子神经网络QNN（也可以理解为参数化量子电路）来学习量子态的特征谱。这里，我们提供一个预设的2量子位量子电路。\n",
    "\n",
    "- 我们预设一些该参数化电路的参数，比如宽度为2量子位。\n",
    "\n",
    "- 初始化其中的变量参数，${\\bf{\\theta }}$代表我们量子神经网络中的参数组成的向量。\n",
    "         "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%% \n"
    }
   },
   "outputs": [],
   "source": [
    "N = 2 # 量子神经网络的宽度\n",
    "SEED = 1 # 种子\n",
    "THETA_SIZE = 14 # 网络中的参数\n",
    "\n",
    "def U_theta(theta, N):\n",
    "    \"\"\"\n",
    "    U_theta\n",
    "    \"\"\"\n",
    "\n",
    "    cir = UAnsatz(N)\n",
    "    cir.rz(theta[0], 1)\n",
    "    cir.ry(theta[1], 1)\n",
    "    cir.rz(theta[2], 1)\n",
    "\n",
    "    cir.rz(theta[3], 2)\n",
    "    cir.ry(theta[4], 2)\n",
    "    cir.rz(theta[5], 2)\n",
    "\n",
    "    cir.cnot([2, 1])\n",
    "\n",
    "    cir.rz(theta[6], 1)\n",
    "    cir.ry(theta[7], 2)\n",
    "\n",
    "    cir.cnot([1, 2])\n",
    "\n",
    "    cir.rz(theta[8], 1)\n",
    "    cir.ry(theta[9], 1)\n",
    "    cir.rz(theta[10], 1)\n",
    "\n",
    "    cir.rz(theta[11], 2)\n",
    "    cir.ry(theta[12], 2)\n",
    "    cir.rz(theta[13], 2)\n",
    "\n",
    "    return cir.state\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 配置训练模型 - 损失函数\n",
    "- 现在我们已经有了数据和量子神经网络的架构，我们将进一步定义训练参数、模型和损失函数。\n",
    "- 通过作用量子神经网络$U(\\theta)$在量子态$\\rho$后得到的量子态记为$\\tilde\\rho$，我们设定损失函数为$\\tilde\\rho$与用来标记的量子态$\\sigma=0.1 \\ket{00}\\bra{00} + 0.2 \\ket{01}\\bra{01} + 0.3 \\ket{10}\\bra{10} + 0.4 \\ket{11}\\bra{11}$的内积。\n",
    "- 具体的，设定损失函数为 $\\mathcal{L}(\\boldsymbol{\\theta})  = Tr(\\tilde\\rho\\sigma) .$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "sigma = diag([0.1, 0.2, 0.3, 0.4]) # 输入用来标记的量子态sigma\n",
    "sigma = sigma.astype('complex64')\n",
    "\n",
    "class Net(fluid.dygraph.Layer):\n",
    "    \"\"\"\n",
    "    Construct the model net\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, shape, rho, sigma, param_attr=fluid.initializer.Uniform(low=0.0, high=2 * numpy.pi, seed=SEED),\n",
    "                 dtype='float32'):\n",
    "        super(Net, self).__init__()\n",
    "\n",
    "        self.rho = fluid.dygraph.to_variable(rho)\n",
    "        self.sigma = fluid.dygraph.to_variable(sigma)\n",
    "\n",
    "        self.theta = self.create_parameter(shape=shape, attr=param_attr, dtype=dtype, is_bias=False)\n",
    "\n",
    "    def forward(self, N):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "        Returns:\n",
    "            The loss.\n",
    "        \"\"\"\n",
    "\n",
    "        out_state = U_theta(self.theta, N)\n",
    "\n",
    "        # rho_tilde 是将U_theta作用在rho后得到的量子态 \n",
    "        rho_tilde = matmul(\n",
    "            matmul(transpose(\n",
    "                fluid.framework.ComplexVariable(out_state.real, -out_state.imag),\n",
    "                perm=[1, 0]\n",
    "            ), self.rho), out_state\n",
    "        )\n",
    "\n",
    "        # record the new loss\n",
    "        loss = trace(matmul(self.sigma, rho_tilde))\n",
    "\n",
    "        return loss.real, rho_tilde"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 配置训练模型 - 模型参数\n",
    "在进行量子神经网络的训练之前，我们还需要进行一些训练（超）参数的设置，例如学习速率与迭代次数。\n",
    "- 设定学习速率（learning rate）为0.1;\n",
    "- 设定迭代次数为50次。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ITR = 50 #训练的总的迭代次数\n",
    "\n",
    "LR = 0.1 #学习速率"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 进行训练\n",
    "\n",
    "- 当训练模型的各项参数都设置完成后，我们将数据转化为Paddle动态图中的变量，进而进行量子神经网络的训练。\n",
    "- 过程中我们用的是Adam Optimizer，也可以调用Paddle中提供的其他优化器。\n",
    "- 我们将训练过程中的结果依次输出。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%% \n"
    }
   },
   "outputs": [],
   "source": [
    "with fluid.dygraph.guard():\n",
    "    # net\n",
    "    net = Net(shape=[THETA_SIZE], rho=rho, sigma=sigma)\n",
    "\n",
    "    # optimizer\n",
    "    opt = fluid.optimizer.AdagradOptimizer(learning_rate=LR, parameter_list=net.parameters())\n",
    "    # gradient descent loop\n",
    "    for itr in range(ITR):\n",
    "        loss, rho_tilde = net(N)\n",
    "\n",
    "        rho_tilde_np = rho_tilde.numpy()\n",
    "        loss.backward()\n",
    "        opt.minimize(loss)\n",
    "        net.clear_gradients()\n",
    "\n",
    "        print('iter:', itr, 'loss:', '%.4f' % loss.numpy()[0])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 总结\n",
    "根据上面训练得到的结果，通过大概50次迭代，我们就比较好的完成了对角化。\n",
    "我们可以通过打印$\n",
    "\\tilde{\\rho} = U(\\boldsymbol{\\theta})\\rho U^\\dagger(\\boldsymbol{\\theta})\n",
    "$\n",
    "的来验证谱分解的效果。特别的，我们可以验证它的对角线与我们目标谱是非常接近的。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "print(rho_tilde_np)\n",
    "\n",
    "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": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 参考文献\n",
    "\n",
    "[1] R. Larose, A. Tikku, É. O. Neel-judy, L. Cincio, and P. J. Coles, “Variational quantum state diagonalization,” npj Quantum Inf., no. November 2018, 2019.\n",
    "\n",
    "[2] K. M. Nakanishi, K. Mitarai, and K. Fujii, “Subspace-search variational quantum eigensolver for excited states,” Phys. Rev. Res., vol. 1, no. 3, p. 033062, Oct. 2019.\n",
    "\n",
    "[3] M. Cerezo, K. Sharma, A. Arrasmith, and P. J. Coles, “Variational Quantum State Eigensolver,” arXiv:2004.01372, no. 1, pp. 1–14, Apr. 2020.\n"
   ]
  }
 ],
 "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": 1
}