{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gibbs State Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<em> Copyright (c) 2021 Institute for Quantum Computing, Baidu Inc. All Rights Reserved. </em>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial will show how to train a quantum neural network (QNN) through Paddle Quantum to prepare a quantum Gibbs state."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Background"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The frontiers of quantum computing include quantum machine learning and quantum optimization. In these two areas, the preparation of specific quantum states is a fundamental problem. In particular, the preparation of Gibbs state is a necessary step to realize many quantum algorithms and is widely used in:\n",
    "- Learning of restricted Boltzmann machines in quantum machine learning [1]\n",
    "- Solving optimization problems such as convex optimization and positive semidefinite programming [2]\n",
    "- Combinatorial optimization problem [3]\n",
    "\n",
    "The Gibbs state is defined as follows: Given the Hamiltonian $H$ of an $n$-qubit system (generally this is a Hermitian matrix of $2^n\\times2^n$), the Gibbs state at temperature $T$ is\n",
    "$$\n",
    "\\rho_G = \\frac{{{e^{-\\beta H}}}}{{\\text{tr}({e^{-\\beta H}})}},\n",
    "\\tag{1}\n",
    "$$\n",
    "where ${e^{-\\beta H}}$ is the matrix exponential of matrix $-\\beta H$. $\\beta = \\frac{1}{{kT}}$ is the inverse temperature parameter of the system, where $T $ Is the temperature parameter and $k$ is Boltzmann's constant (in this tutorial, we take $k = 1$)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Paddle Quantum Implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let us import the necessary libraries and packages through the following code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:56:40.078886Z",
     "start_time": "2021-04-30T08:56:36.868403Z"
    }
   },
   "outputs": [],
   "source": [
    "import scipy\n",
    "import paddle\n",
    "from numpy import trace as np_trace\n",
    "import paddle_quantum as pq\n",
    "from paddle_quantum.ansatz import Circuit\n",
    "from paddle_quantum.state import zero_state\n",
    "from paddle_quantum.qinfo import state_fidelity, partial_trace, pauli_str_to_matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a hands-on example, here we consider a 3-qubit Hamiltonian and its Gibbs state:\n",
    "\n",
    "$$\n",
    "H = -Z \\otimes Z \\otimes I - I \\otimes Z \\otimes Z - Z \\otimes I \\otimes Z, \\quad I=\\left [\n",
    "\\begin{matrix}\n",
    "1 & 0  \\\\\n",
    "0 & 1  \\\\\n",
    "\\end{matrix} \n",
    "\\right ], \\quad \n",
    "Z=\\left [\n",
    "\\begin{matrix}\n",
    "1 & 0  \\\\\n",
    "0 & -1  \\\\\n",
    "\\end{matrix} \n",
    "\\right ].\n",
    "\\tag{2}\n",
    "$$\n",
    "\n",
    "In this example, we set the inverse temperature parameter to $\\beta = 1.5$. Besides, to test the final results, we have generated the ideal Gibbs state $\\rho_G$ in advance according to the definition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:56:40.099967Z",
     "start_time": "2021-04-30T08:56:40.082138Z"
    }
   },
   "outputs": [],
   "source": [
    "N = 4                               # The width of the QNN\n",
    "N_SYS_B = 3                         # The number of qubits of subsystem B used to generate the Gibbs state\n",
    "SEED = 16                           # Fixed random seed\n",
    "beta = 1.5                          # Set the inverse temperature parameter beta\n",
    "pq.set_backend('density_matrix')    # Set density matrix backend\n",
    "pq.set_dtype('complex128')          # Set calculation precision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:56:40.201757Z",
     "start_time": "2021-04-30T08:56:40.106315Z"
    }
   },
   "outputs": [],
   "source": [
    "# Generate a specific Hamiltonian represented by a Pauli string\n",
    "H = [[-1.0,'z0,z1'], [-1.0,'z1,z2'], [-1.0,'z0,z2']]\n",
    "\n",
    "# Generate matrix information of Hamiltonian\n",
    "hamiltonian = pauli_str_to_matrix(H, N_SYS_B).numpy()\n",
    "\n",
    "# Generate the ideal target Gibbs state rho\n",
    "rho_G = scipy.linalg.expm(-1 * beta * hamiltonian) / np_trace(scipy.linalg.expm(-1 * beta * hamiltonian))\n",
    "\n",
    "# Set to the data type supported by Paddle Quantum\n",
    "hamiltonian = hamiltonian.astype('complex128')\n",
    "rho_G = paddle.to_tensor(rho_G, dtype='complex128')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Building a quantum neural network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- In this example, we will prepare the Gibbs state by training the QNN (also can be understood as a parameterized quantum circuit). Here, we provide a simple 4-qubit quantum circuit as follows:\n",
    "\n",
    "  ![Ugibbs.jpg](https://release-data.cdn.bcebos.com/PIC%2FUgibbs.jpg)\n",
    "\n",
    "- We need to preset some circuit parameters. For example, the circuit has 4 qubits. The first qubit is the ancillary system, and the 2-4th qubits are the subsystems used to generate the Gibbs state.\n",
    "- Initialize the variable ${\\bf{\\theta }}$ that represents the vector of parameters in our QNN.\n",
    "         \n",
    "\n",
    "Next, we use Paddle Quantum's `Circuit` class and the built-in `real_entangled_layer` circuit template to build a QNN based on the circuit design in the above figure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:56:40.213503Z",
     "start_time": "2021-04-30T08:56:40.205311Z"
    }
   },
   "outputs": [],
   "source": [
    "def U_theta(num_qubits: int, depth: int) -> Circuit:\n",
    "    \"\"\"\n",
    "    Quantum Neural Network\n",
    "    \"\"\"\n",
    "    # Initialize the QNN according to the number of qubits\n",
    "    cir = Circuit(num_qubits)\n",
    "    \n",
    "    # Built-in {R_y + CNOT} circuit template\n",
    "    cir.real_entangled_layer(depth=depth)\n",
    "    \n",
    "    # The QNN acts on a given initial state\n",
    "    cir.ry()\n",
    "    \n",
    "    return cir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuring the training model: loss function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Now that we have the data and QNN architecture, we also need to define appropriate training parameters, models, and loss functions to achieve our goals.\n",
    "\n",
    "- Specifically, we refer to the method in the paper [4]. The core idea is to use the Gibbs state to achieve the minimum free energy.\n",
    "\n",
    "- By applying the QNN $U(\\theta)$ on the initial state, we can get the output state $\\left| {\\psi \\left( {\\bf{\\theta }} \\right)} \\right\\rangle $. Its state in the 2-4th qubits is recorded as $\\rho_B(\\theta)$.\n",
    "\n",
    "- Set the loss function in the training model. In Gibbs state learning, we use the truncation of the von Neumann entropy function to estimate the free energy, and the corresponding loss function, as in reference [4], can be set as $loss = {L_1} + {L_2} + {L_3} $, where\n",
    "\n",
    "$$\n",
    "{L_1}= \\text{tr}(H\\rho_B), \\quad {L_2} = 2{\\beta^{-1}}{\\text{tr}}(\\rho_B^2), \\quad L_3 = - {\\beta ^{ - 1}}\\big(\\text{tr}(\\rho_B^3) + 3\\big)/2.\n",
    "\\tag{3}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:56:40.238960Z",
     "start_time": "2021-04-30T08:56:40.228348Z"
    }
   },
   "outputs": [],
   "source": [
    " # define loss function\n",
    "def loss_func(cir: Circuit, Hamiltonian: paddle.Tensor, N_SYS_B: int) -> paddle.Tensor:\n",
    "    # Apply QNN\n",
    "    rho_AB = cir(zero_state(N))\n",
    "    \n",
    "    # Calculate partial trace to obtain the quantum state rho_B of subsystem B\n",
    "    rho_B = partial_trace(rho_AB.data, 2 ** (N - N_SYS_B), 2 ** (N_SYS_B), 1)\n",
    "    \n",
    "    # Calculate the three parts of the loss function\n",
    "    rho_B_squre = rho_B @ rho_B\n",
    "    loss1 = paddle.real(paddle.trace(rho_B @ Hamiltonian))\n",
    "    loss2 = paddle.real(paddle.trace(rho_B_squre)) * 2 / beta\n",
    "    loss3 = -(paddle.real(paddle.trace(rho_B_squre @ rho_B)) + 3) / (2 * beta)\n",
    "    \n",
    "    # 最终的损失函数\n",
    "    loss = loss1 + loss2 + loss3  \n",
    "    \n",
    "    return loss, rho_B"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure training model: model parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before training the QNN, we also need to set some training hyperparameters, mainly the learning rate (LR), the number of iterations (ITR), and the depth (D) of the QNN. Here we set the learning rate to 0.5 and the number of iterations to 50. Readers may wish to adjust by themselves to explore the influence of hyperparameter adjustment on the training effect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:56:40.966455Z",
     "start_time": "2021-04-30T08:56:40.959537Z"
    }
   },
   "outputs": [],
   "source": [
    "ITR = 50 # Set the total number of iterations of training\n",
    "LR = 0.5 # Set the learning rate\n",
    "D = 1 # Set the depth of the QNN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- After all the training model parameters are set, we convert the data into Tensor in PaddlePaddle and then train the QNN.\n",
    "- During training, we use [Adam Optimizer](https://www.paddlepaddle.org.cn/documentation/docs/en/api/paddle/optimizer/adam/Adam_en.html). Other optimizers are also provided in PaddlePaddle.\n",
    "- We output the results of the training process in turn.\n",
    "- In particular, we sequentially output the fidelity of the quantum state $\\rho_B(\\theta)$ and Gibbs state $\\rho_G$ we learned. The higher the fidelity, the closer the QNN output state is to Gibbs state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:56:50.509486Z",
     "start_time": "2021-04-30T08:56:47.508586Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iter: 10 loss: -3.1085 fid: 0.9241\n",
      "iter: 20 loss: -3.3375 fid: 0.9799\n",
      "iter: 30 loss: -3.3692 fid: 0.9897\n",
      "iter: 40 loss: -3.3990 fid: 0.9929\n",
      "iter: 50 loss: -3.4133 fid: 0.9959\n",
      "\n",
      "The trained circuit:  \n",
      "--Ry(6.290)----*--------------x----Ry(0.747)--\n",
      "               |              |               \n",
      "--Ry(4.745)----x----*---------|----Ry(6.249)--\n",
      "                    |         |               \n",
      "--Ry(-0.01)---------x----*----|----Ry(-0.05)--\n",
      "                         |    |               \n",
      "--Ry(0.017)--------------x----*----Ry(6.310)--\n",
      "                                              \n"
     ]
    }
   ],
   "source": [
    "paddle.seed(SEED)\n",
    "    \n",
    "# Convert Numpy array to Tensor supported in PaddlePaddle\n",
    "H = paddle.to_tensor(hamiltonian)\n",
    "\n",
    "# Determine the parameter dimension of the network\n",
    "circuit = U_theta(N, D)\n",
    "\n",
    "# Generally speaking, we use Adam optimizer to get relatively good convergence\n",
    "# Of course, it can be changed to SGD or RMS prop.\n",
    "opt = paddle.optimizer.Adam(learning_rate=LR, parameters=circuit.parameters())\n",
    "\n",
    "# Optimization loops\n",
    "for itr in range(1, ITR + 1):\n",
    "        \n",
    "    # Run forward propagation to calculate the loss function and return the generated quantum state rho_B\n",
    "    loss, rho_B = loss_func(circuit, H, N_SYS_B)\n",
    "        \n",
    "    # Run back propagation to minimize the loss function\n",
    "    loss.backward()\n",
    "    opt.minimize(loss)\n",
    "    opt.clear_grad()\n",
    "\n",
    "    # Convert to Numpy array to calculate the fidelity of the quantum state F(rho_B, rho_G)\n",
    "    fid = state_fidelity(rho_B, rho_G)\n",
    "        \n",
    "    # Print training results\n",
    "    if itr % 10 == 0:\n",
    "        print('iter:', itr, 'loss:', '%.4f' % loss.numpy(), 'fid:', '%.4f' % fid.numpy())\n",
    "    if itr == ITR:\n",
    "        print(\"\\nThe trained circuit:  \")\n",
    "        print(circuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "According to the results obtained from the above training, after about 50 iterations, we can achieve a high-precision Gibbs state with a fidelity higher than 99.5% and complete the preparation of the Gibbs state efficiently and accurately. We can output the QNN's learned parameters and its output state through the print function.\n",
    "\n",
    "_______\n",
    "\n",
    "## References\n",
    "\n",
    "[1] Kieferová, M. & Wiebe, N. Tomography and generative training with quantum Boltzmann machines. [Phys. Rev. A 96, 062327 (2017).](https://journals.aps.org/pra/abstract/10.1103/PhysRevA.96.062327)\n",
    "\n",
    "[2] Brandao, F. G. S. L. & Svore, K. M. Quantum Speed-Ups for Solving Semidefinite Programs. [in 2017 IEEE 58th Annual Symposium on Foundations of Computer Science (FOCS) 415–426 (IEEE, 2017). ](https://ieeexplore.ieee.org/abstract/document/8104077)\n",
    "\n",
    "[3] Somma, R. D., Boixo, S., Barnum, H. & Knill, E. Quantum Simulations of Classical Annealing Processes. [Phys. Rev. Lett. 101, 130504 (2008).](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.101.130504)\n",
    "\n",
    "[4] Wang, Y., Li, G. & Wang, X. Variational quantum Gibbs state preparation with a truncated Taylor series. [Phys. Rev. A 16, 054035 (2021).](https://journals.aps.org/prapplied/abstract/10.1103/PhysRevApplied.16.054035)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.13 ('py3.7_pq2.2.1')",
   "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.13"
  },
  "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
  },
  "vscode": {
   "interpreter": {
    "hash": "4e4e2eb86ad73936e915e7c7629a18a8ca06348106cf3e66676b9578cb1a47dd"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}