{
 "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.\n"
   ]
  },
  {
   "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-01-09T10:26:18.115601Z",
     "start_time": "2021-01-09T10:26:14.563581Z"
    }
   },
   "outputs": [],
   "source": [
    "import scipy\n",
    "from numpy import array, concatenate, zeros\n",
    "from numpy import pi as PI\n",
    "from numpy import trace as np_trace\n",
    "from paddle import fluid\n",
    "from paddle.complex import matmul, trace\n",
    "from paddle_quantum.circuit import UAnsatz\n",
    "from paddle_quantum.state import density_op\n",
    "from paddle_quantum.utils 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-01-09T10:26:18.156435Z",
     "start_time": "2021-01-09T10:26:18.118621Z"
    }
   },
   "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 = 14    # Fixed random seed\n",
    "beta = 1.5   # Set the inverse temperature parameter beta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-01-09T10:26:18.185219Z",
     "start_time": "2021-01-09T10:26:18.161336Z"
    }
   },
   "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)\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 = rho_G.astype(\"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 `UAnsatz` class and the built-in `real_entangled_layer(theta, D)` 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-01-09T10:26:18.202627Z",
     "start_time": "2021-01-09T10:26:18.188216Z"
    }
   },
   "outputs": [],
   "source": [
    "def U_theta(initial_state, theta, N, D):\n",
    "    \"\"\"\n",
    "    Quantum Neural Network\n",
    "    \"\"\"\n",
    "    \n",
    "    # Initialize the QNN according to the number of qubits/network width\n",
    "    cir = UAnsatz(N)\n",
    "    \n",
    "    # Built-in {R_y + CNOT} circuit template\n",
    "    cir.real_entangled_layer(theta[:D], D)\n",
    "    \n",
    "    # Add the last layer of R_y rotation gates\n",
    "    for i in range(N):\n",
    "        cir.ry(theta=theta[D][i][0], which_qubit=i)\n",
    "        \n",
    "    # The QNN acts on a given initial state\n",
    "    final_state = cir.run_density_matrix(initial_state)\n",
    "\n",
    "    return final_state"
   ]
  },
  {
   "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-01-09T10:26:18.589450Z",
     "start_time": "2021-01-09T10:26:18.574464Z"
    }
   },
   "outputs": [],
   "source": [
    "class Net(fluid.dygraph.Layer):\n",
    "    def __init__(self, shape, param_attr=fluid.initializer.Uniform(\n",
    "        low=0.0, high=2*PI, seed=SEED), dtype='float64'):\n",
    "        super(Net, self).__init__()\n",
    "        \n",
    "        # Initialize parameters theta with the uniform distribution of [0, 2*pi]\n",
    "        self.theta = self.create_parameter(shape=shape,\n",
    "                     attr=param_attr, dtype=dtype, is_bias=False)\n",
    "        \n",
    "        # Initialize the density matrix rho = |0..0><0..0|\n",
    "        self.initial_state=fluid.dygraph.to_variable(density_op(N))\n",
    "\n",
    "    # Define loss function and forward propagation mechanism\n",
    "    def forward(self, H, N, N_SYS_B, D):\n",
    "\n",
    "        # Apply QNN\n",
    "        rho_AB = U_theta(self.initial_state, self.theta, N, D)\n",
    "\n",
    "        # Calculate partial trace to obtain the quantum state rho_B of subsystem B\n",
    "        rho_B = partial_trace(rho_AB, 2 ** (N-N_SYS_B), 2 ** (N_SYS_B), 1)\n",
    "        \n",
    "        # Calculate the three parts of the loss function\n",
    "        rho_B_squre = matmul(rho_B, rho_B)\n",
    "        loss1 = (trace(matmul(rho_B, H))).real\n",
    "        loss2 = (trace(rho_B_squre)).real * 2 / beta\n",
    "        loss3 =-((trace(matmul(rho_B_squre, rho_B))).real + 3) / (2 * beta)\n",
    "        \n",
    "        # Final loss function\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-01-09T10:26:20.193223Z",
     "start_time": "2021-01-09T10:26:20.180849Z"
    }
   },
   "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 variables in the PaddlePaddle dynamic graph and then train the QNN.\n",
    "- During training, we use [Adam Optimizer](https://www.paddlepaddle.org.cn/documentation/docs/en/api/optimizer/AdamOptimizer.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": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-01-09T10:26:24.585128Z",
     "start_time": "2021-01-09T10:26:21.249352Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iter: 10 loss: -3.1189 fid: 0.9504\n",
      "iter: 20 loss: -3.3502 fid: 0.9846\n",
      "iter: 30 loss: -3.3630 fid: 0.9873\n",
      "iter: 40 loss: -3.4087 fid: 0.9948\n",
      "iter: 50 loss: -3.4110 fid: 0.9953\n"
     ]
    }
   ],
   "source": [
    "# Initialize PaddlePaddle dynamic graph mechanism\n",
    "with fluid.dygraph.guard():\n",
    "    \n",
    "    # Convert Numpy array to variable supported in PaddlePaddle dynamic graph mode\n",
    "    H = fluid.dygraph.to_variable(hamiltonian)\n",
    "\n",
    "    # Determine the parameter dimension of the network\n",
    "    net = Net(shape=[D + 1, N, 1])\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 = fluid.optimizer.AdamOptimizer(learning_rate=LR,\n",
    "                          parameter_list=net.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 = net(H, N, N_SYS_B, D)\n",
    "        \n",
    "        # Under the dynamic graph mechanism, run back propagation to minimize the loss function\n",
    "        loss.backward()\n",
    "        opt.minimize(loss)\n",
    "        net.clear_gradients()\n",
    "\n",
    "        # Convert to Numpy array to calculate the fidelity of the quantum state F(rho_B, rho_G)\n",
    "        rho_B = rho_B.numpy()\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(),\n",
    "                                 'fid:','%.4f'% fid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[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. [arXiv:2005.08797 (2020).](https://arxiv.org/pdf/2005.08797.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.6.10"
  },
  "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": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}