Quantum_Autoencoder_Tutorial_CN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 量子变分自编码器 （Quantum Autoencoder）\n",
    "\n",
    "<em> Copyright (c) Institute for Quantum Computing, Baidu Inc. All Rights Reserved. </em>\n",
    "\n",
    "## 概览\n",
    "\n",
    "在这个案例中，我们将展示如何训练量子自编码器来压缩和重建给定的量子态（混合态）[1]。量子自编码器的形式与经典自编码器非常相似，同样由编码器 $E$（encoder）和解码器 $D$（decoder）组成。对于输入的 $N$ 量子比特系统的量子态 $\\rho_{in}$（这里我们采用量子力学的密度算符表示来描述混合态），先用编码器 $E = U(\\theta)$ 将信息编码到其中的部分量子比特上。这部分量子比特记为**系统 $A$**。对剩余的量子比特 （这部分记为**系统 $B$**）进行测量并丢弃后，我们就得到了压缩后的量子态 $\\rho_{encode}$！压缩后的量子态维度和量子系统 $A$ 的维度大小保持一致。假设我们需要 $N_A$ 个量子比特来描述系统 $A$ ，那么编码后量子态 $\\rho_{encode}$ 的维度就是 $2^{N_A}\\times 2^{N_A}$。 这里比较特殊的是, 对应我们这一步测量并丢弃的操作的数学操作是 partial trace。读者在直观上可以理解为张量积 $\\otimes$ 的逆运算。\n",
    "\n",
    "我们不妨来看一个具体的例子来理解：\n",
    "给定一个由 $N_A$ 个量子比特描述的量子态 $\\rho_A$ 和另外一个由 $N_B$ 个量子比特描述的量子态 $\\rho_B$, 那么由 $A、B$ 两个子系统构成的的整体量子系统 $N = N_A+ N_B$ 的量子态就可以描述为: $\\rho_{AB} = \\rho_A \\otimes \\rho_B$。现在我们让整个量子系统在酉矩阵 $U$ 的作用下演化一段时间后，得到了一个新的量子态 $\\tilde{\\rho_{AB}} = U\\rho_{AB}U^\\dagger$。 那么如果这时候我们只想得到量子子系统 $A$ 所处于的新的量子态 $\\tilde{\\rho_A}$， 我们该怎么做呢？很简单，只需要测量量子子系统 $B$ 然后丢弃测量结果。运算上这一步由 partial trace 来完成 $\\tilde{\\rho_A} = \\text{Tr}_B (\\tilde{\\rho_{AB}})$。在量桨上，我们可以用内置的 `partial_trace(rho_AB, 2**N_A, 2**N_B, 2)` 指令来完成这一运算。**注意：** 其中最后一个输入为2，表示我们想丢弃量子系统 $B$ 的量子态。\n",
    "\n",
    "<img src=\"figures/encoder_pipeline.png\" width=\"950\" >\n",
    "\n",
    "在讨论完编码的过程后，我们来看看如何解码。为了解码量子态 $\\rho_{encode}$，我们需要引入与系统 $B$ 维度相同的系统 $C$ 并且初始态取为全0态。再用解码器 $D = U^\\dagger(\\theta)$ 作用在整个量子系统 $A+C$ 上对系统 $A$ 中压缩的信息进行解码。我们希望最后得到的量子态 $\\rho_{out}$ 与 $\\rho_{in}$ 尽可能相似并用 Uhlmann-Josza 保真度 $F$ （Fidelity）来衡量他们之间的相似度。\n",
    "\n",
    "$$ F(\\rho_{in}, \\rho_{out}) = \\left(\\operatorname{tr} \\sqrt{\\sqrt{\\rho_{in}} \\rho_{out} \\sqrt{\\rho_{in}}} \\right)^{2}$$\n",
    "\n",
    "最后，通过优化编码器里的参数，我们就可以尽可能地提高 $\\rho_{in}$ 与 $\\rho_{out}$ 的保真度。这里我们先引入必要的 package。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from numpy import diag\n",
    "import scipy\n",
    "\n",
    "import paddle\n",
    "from paddle import fluid\n",
    "from paddle_quantum.circuit import UAnsatz\n",
    "from paddle.complex import matmul, trace, kron\n",
    "from paddle_quantum.utils import dagger, state_fidelity, partial_trace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 生成初始态\n",
    "\n",
    "下面我们通过一个简单的例子来展示量子自编码器的工作流程和原理。\n",
    "\n",
    "我们考虑 $N = 3$ 个量子比特上的量子态 $\\rho_{in}$。我们先通过编码器将其信息编码到下方的两个量子比特（系统 $A$），之后对第一个量子比特（系统 $B$）测量并丢弃，之后引入一个处于0态的量子比特（新的参考系统 $C$）来替代丢弃的量子比特 $B$ 的维度。最后通过解码器，将 $A$ 中压缩的信息复原成 $\\rho_{out}$。在这里，我们假设初态是一个混合态并且 $\\rho_{in}$ 的特征谱为 $\\lambda_i \\in \\{0.4， 0.2， 0.2， 0.1， 0.1, \\,0, \\,0, \\,0 \\}$，然后通过作用一个随机的酉变换来生成初态 $\\rho_{in}$。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "N_A = 2        # 系统 A 的量子比特数\n",
    "N_B = 1        # 系统 B 的量子比特数\n",
    "N = N_A + N_B  # 总的量子比特数\n",
    "\n",
    "scipy.random.seed(1)                            # 固定随机种子\n",
    "V = scipy.stats.unitary_group.rvs(2**N)         # 随机生成一个酉矩阵\n",
    "D = diag([0.4, 0.2, 0.2, 0.1, 0.1, 0, 0, 0])    # 输入目标态rho的谱\n",
    "V_H = V.conj().T                                # 进行厄尔米特转置\n",
    "rho_in = (V @ D @ V_H).astype('complex128')      # 生成 rho_in\n",
    "\n",
    "# 将 C 量子系统初始化为\n",
    "rho_C = np.diag([1,0]).astype('complex128')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 搭建量子神经网络\n",
    "\n",
    "在这里，我们用量子神经网络来作为编码器和解码器。假设系统 $A$ 有 $N_A$ 个量子比特，系统 $B$ 和 $C$ 分别有$N_B$ 个量子比特，量子神经网络的深度为 D。编码器 $E$ 作用在系统 $A$ 和 $B$ 共同构成的总系统上，解码器 $D$ 作用在$A$ 和 $C$ 共同构成的总系统上。在我们的问题里，$N_{A} = 2$，$N_{B} = 1$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置电路参数\n",
    "cir_depth = 6                        # 电路深度\n",
    "block_len = 2                        # 每个模组的长度\n",
    "theta_size = N*block_len*cir_depth   # 网络参数 theta 的大小\n",
    "\n",
    "\n",
    "# 搭建编码器 Encoder E\n",
    "def Encoder(theta):\n",
    "\n",
    "    # 用 UAnsatz 初始化网络\n",
    "    cir = UAnsatz(N)\n",
    "    \n",
    "    # 搭建层级结构：\n",
    "    for layer_num in range(cir_depth):\n",
    "        \n",
    "        for which_qubit in range(N):\n",
    "            cir.ry(theta[block_len*layer_num*N + which_qubit], which_qubit)\n",
    "            cir.rz(theta[(block_len*layer_num + 1)*N + which_qubit], which_qubit)\n",
    "\n",
    "        for which_qubit in range(N-1):\n",
    "            cir.cnot([which_qubit, which_qubit + 1])\n",
    "        cir.cnot([N-1, 0])\n",
    "\n",
    "    return cir.U"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 配置训练模型-损失函数\n",
    "\n",
    "在这里，我们定义的损失函数为 \n",
    "\n",
    "$$\n",
    "Loss = 1 - \\langle 0...0|\\rho_{trash}|0...0\\rangle\n",
    "$$\n",
    "\n",
    "其中 $\\rho_{trash}$ 是经过编码后丢弃的 $B$ 系统的量子态。接着我们通过飞桨的自动微分框架来训练量子神经网络，最小化损失函数。如果损失函数最后达到 0，则输入态和输出态完全相同。这就意味着我们完美地实现了压缩和解压缩，这时初态和末态的保真度为  $F(\\rho_{in}, \\rho_{out}) = 1$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iter: 10 loss: 0.1795 fid: 0.8076\n",
      "iter: 20 loss: 0.1374 fid: 0.8573\n",
      "iter: 30 loss: 0.1149 fid: 0.8812\n",
      "iter: 40 loss: 0.1073 fid: 0.8886\n",
      "iter: 50 loss: 0.1044 fid: 0.8910\n",
      "iter: 60 loss: 0.1029 fid: 0.8921\n",
      "iter: 70 loss: 0.1019 fid: 0.8927\n",
      "iter: 80 loss: 0.1013 fid: 0.8931\n",
      "iter: 90 loss: 0.1009 fid: 0.8933\n",
      "iter: 100 loss: 0.1006 fid: 0.8935\n"
     ]
    }
   ],
   "source": [
    "# 超参数设置\n",
    "N_A = 2        # 系统 A 的量子比特数\n",
    "N_B = 1        # 系统 B 的量子比特数\n",
    "N = N_A + N_B  # 总的量子比特数\n",
    "LR = 0.2       # 设置学习速率\n",
    "ITR = 100      # 设置迭代次数\n",
    "SEED = 14      # 固定初始化参数用的随机数种子\n",
    "\n",
    "class NET(fluid.dygraph.Layer):\n",
    "    \"\"\"\n",
    "    Construct the model net\n",
    "    \"\"\"\n",
    "    def __init__(self, shape, param_attr=fluid.initializer.Uniform(\n",
    "        low=0.0, high=2 * np.pi, seed = SEED), dtype='float64'):\n",
    "        super(NET, self).__init__()\n",
    "        \n",
    "        # 我们需要将 Numpy array 转换成 Paddle 动态图模式中支持的 variable\n",
    "        self.rho_in = fluid.dygraph.to_variable(rho_in)\n",
    "        self.rho_C = fluid.dygraph.to_variable(rho_C)\n",
    "        self.theta = self.create_parameter(shape=shape, \n",
    "                     attr=param_attr, dtype=dtype, is_bias=False)\n",
    "    \n",
    "    # 定义损失函数和前向传播机制\n",
    "    def forward(self):\n",
    "        # 生成初始的编码器 E 和解码器 D\\n\",\n",
    "        E = Encoder(self.theta)\n",
    "        E_dagger = dagger(E)\n",
    "        D = E_dagger\n",
    "        D_dagger = E\n",
    "\n",
    "        # 编码量子态 rho_in\n",
    "        rho_BA = matmul(matmul(E, self.rho_in), E_dagger)\n",
    "        \n",
    "        # 取 partial_trace() 获得 rho_encode 与 rho_trash\n",
    "        rho_encode = partial_trace(rho_BA, 2 ** N_B, 2 ** N_A, 1)\n",
    "        rho_trash = partial_trace(rho_BA, 2 ** N_B, 2 ** N_A, 2)\n",
    "\n",
    "        # 解码得到量子态 rho_out\n",
    "        rho_CA = kron(self.rho_C, rho_encode)\n",
    "        rho_out = matmul(matmul(D, rho_CA), D_dagger)\n",
    "        \n",
    "        # 通过 rho_trash 计算损失函数\n",
    "        \n",
    "        zero_Hamiltonian = fluid.dygraph.to_variable(np.diag([1,0]).astype('complex128'))\n",
    "        loss = 1 - (trace(matmul(zero_Hamiltonian, rho_trash))).real\n",
    "\n",
    "        return loss, self.rho_in, rho_out\n",
    "\n",
    "\n",
    "# 初始化paddle动态图机制  \n",
    "with fluid.dygraph.guard():\n",
    "\n",
    "    # 生成网络\n",
    "    net = NET([theta_size])\n",
    "\n",
    "    # 一般来说，我们利用Adam优化器来获得相对好的收敛，当然你可以改成SGD或者是RMS prop.\n",
    "    opt = fluid.optimizer.AdagradOptimizer(learning_rate=LR, \n",
    "                          parameter_list=net.parameters())\n",
    "\n",
    "    # 优化循环\n",
    "    for itr in range(1, ITR + 1):\n",
    "        \n",
    "        #  前向传播计算损失函数\n",
    "        loss, rho_in, rho_out = net()\n",
    "        \n",
    "        # 在动态图机制下，反向传播极小化损失函数\n",
    "        loss.backward()\n",
    "        opt.minimize(loss)\n",
    "        net.clear_gradients()\n",
    "        \n",
    "        # 计算并打印保真度\n",
    "        fid = state_fidelity(rho_in.numpy(), rho_out.numpy())\n",
    "\n",
    "        if itr % 10 == 0:\n",
    "            print('iter:', itr, 'loss:', '%.4f' % loss, 'fid:', '%.4f' % np.square(fid))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 总结\n",
    "\n",
    "如果系统 $A$ 的维度为 $d_A$，容易证明量子自编码器能实现的压缩重建最大保真度为 $\\rho_{in}$ 最大的 $d_A$ 个本征值之和。在我们的示例里 $d_A = 4$，最大保真度为 \n",
    "\n",
    "$$ F_{limit} = 0.4 + 0.2 + 0.2 + 0.1 = 0.9$$\n",
    "\n",
    "通过 100 次迭代，我们训练出的量子自编码器保真度达到 0.89 以上，和最优情况非常接近。\n",
    "\n",
    "\n",
    "## 参考文献\n",
    "\n",
    "[1] [Romero, J., Olson, J. P. & Aspuru-Guzik, A. Quantum autoencoders for efficient compression of quantum data. Quantum Sci. Technol. 2, 045001 (2017).](https://iopscience.iop.org/article/10.1088/2058-9565/aa8072)"
   ]
  }
 ],
 "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}