docs/source/tutorials/gradient_benchmark_cn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c610ec7d",
   "metadata": {},
   "source": [
    "# 梯度计算效率比较"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edb4fc7f",
   "metadata": {},
   "source": [
    "## 概述\n",
    "\n",
    "在本教程中，我们比较了通过 TensorCircuit 提供的自动微分框架和 Qiskit 提供的传统参数平移框架对梯度和类梯度对象（例如量子费雪信息）进行估测的效率。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3f56164",
   "metadata": {},
   "source": [
    "## 设置\n",
    "\n",
    "我们从 Qiskit 和 TensorCircuit 导入必要的包和模块。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1ab97b03",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "from functools import reduce\n",
    "from operator import xor\n",
    "\n",
    "from qiskit.opflow import I, X, Z, StateFn, CircuitStateFn, SummedOp\n",
    "from qiskit.circuit import QuantumCircuit, ParameterVector\n",
    "from scipy.optimize import minimize\n",
    "from qiskit.opflow.gradients import Gradient, QFI, Hessian\n",
    "\n",
    "import tensorcircuit as tc\n",
    "from tensorcircuit import experimental"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9870a872",
   "metadata": {},
   "source": [
    "## Qiskit 梯度框架基准\n",
    "\n",
    "由于 Qiskit 在梯度估测方面**太**慢了，我们使用小型系统在 Jupyter 中进行基准测试以节省时间，\n",
    "对于更大尺寸和更深的电路，效率差异将变得更加明显。\n",
    "\n",
    "三个类梯度的任务分别是梯度、量子费学信息和 Hessian 评估。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "01180b47",
   "metadata": {},
   "outputs": [],
   "source": [
    "def benchmark(f, *args, trials=10):\n",
    "    time0 = time.time()\n",
    "    r = f(*args)\n",
    "    time1 = time.time()\n",
    "    for _ in range(trials):\n",
    "        r = f(*args)\n",
    "    time2 = time.time()\n",
    "    if trials > 0:\n",
    "        time21 = (time2 - time1) / trials\n",
    "    else:\n",
    "        time21 = 0\n",
    "    ts = (time1 - time0, time21)\n",
    "    print(\"staging time: %.6f s\" % ts[0])\n",
    "    if trials > 0:\n",
    "        print(\"running time: %.6f s\" % ts[1])\n",
    "    return r, ts\n",
    "\n",
    "\n",
    "def grad_qiskit(n, l, trials=2):\n",
    "    hamiltonian = reduce(xor, [X for _ in range(n)])\n",
    "    wavefunction = QuantumCircuit(n)\n",
    "    params = ParameterVector(\"theta\", length=3 * n * l)\n",
    "    for j in range(l):\n",
    "        for i in range(n - 1):\n",
    "            wavefunction.cnot(i, i + 1)\n",
    "        for i in range(n):\n",
    "            wavefunction.rx(params[3 * n * j + i], i)\n",
    "        for i in range(n):\n",
    "            wavefunction.rz(params[3 * n * j + i + n], i)\n",
    "        for i in range(n):\n",
    "            wavefunction.rx(params[3 * n * j + i + 2 * n], i)\n",
    "\n",
    "    # 定义能量对应的期望值\n",
    "    op = ~StateFn(hamiltonian) @ StateFn(wavefunction)\n",
    "    grad = Gradient().convert(operator=op, params=params)\n",
    "\n",
    "    def get_grad_qiskit(values):\n",
    "        value_dict = {params: values}\n",
    "        grad_result = grad.assign_parameters(value_dict).eval()\n",
    "        return grad_result\n",
    "\n",
    "    return benchmark(get_grad_qiskit, np.ones([3 * n * l]), trials=trials)\n",
    "\n",
    "\n",
    "def qfi_qiskit(n, l, trials=0):\n",
    "    hamiltonian = reduce(xor, [X for _ in range(n)])\n",
    "    wavefunction = QuantumCircuit(n)\n",
    "    params = ParameterVector(\"theta\", length=3 * n * l)\n",
    "    for j in range(l):\n",
    "        for i in range(n - 1):\n",
    "            wavefunction.cnot(i, i + 1)\n",
    "        for i in range(n):\n",
    "            wavefunction.rx(params[3 * n * j + i], i)\n",
    "        for i in range(n):\n",
    "            wavefunction.rz(params[3 * n * j + i + n], i)\n",
    "        for i in range(n):\n",
    "            wavefunction.rx(params[3 * n * j + i + 2 * n], i)\n",
    "\n",
    "    nat_grad = QFI().convert(operator=StateFn(wavefunction), params=params)\n",
    "\n",
    "    def get_qfi_qiskit(values):\n",
    "        value_dict = {params: values}\n",
    "        grad_result = nat_grad.assign_parameters(value_dict).eval()\n",
    "        return grad_result\n",
    "\n",
    "    return benchmark(get_qfi_qiskit, np.ones([3 * n * l]), trials=trials)\n",
    "\n",
    "\n",
    "def hessian_qiskit(n, l, trials=0):\n",
    "    hamiltonian = reduce(xor, [X for _ in range(n)])\n",
    "    wavefunction = QuantumCircuit(n)\n",
    "    params = ParameterVector(\"theta\", length=3 * n * l)\n",
    "    for j in range(l):\n",
    "        for i in range(n - 1):\n",
    "            wavefunction.cnot(i, i + 1)\n",
    "        for i in range(n):\n",
    "            wavefunction.rx(params[3 * n * j + i], i)\n",
    "        for i in range(n):\n",
    "            wavefunction.rz(params[3 * n * j + i + n], i)\n",
    "        for i in range(n):\n",
    "            wavefunction.rx(params[3 * n * j + i + 2 * n], i)\n",
    "\n",
    "    # 定义能量对应的期望值\n",
    "    op = ~StateFn(hamiltonian) @ StateFn(wavefunction)\n",
    "    grad = Hessian().convert(operator=op, params=params)\n",
    "\n",
    "    def get_hs_qiskit(values):\n",
    "        value_dict = {params: values}\n",
    "        grad_result = grad.assign_parameters(value_dict).eval()\n",
    "        return grad_result\n",
    "\n",
    "    return benchmark(get_hs_qiskit, np.ones([3 * n * l]), trials=trials)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0acaf881",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "staging time: 1.665786 s\n",
      "running time: 1.474930 s\n"
     ]
    }
   ],
   "source": [
    "g0, _ = grad_qiskit(6, 3)  # 梯度"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6d44b352",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "staging time: 47.131374 s\n"
     ]
    }
   ],
   "source": [
    "qfi0, _ = qfi_qiskit(6, 3)  # QFI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3ae25c46",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "staging time: 80.495983 s\n"
     ]
    }
   ],
   "source": [
    "hs0, _ = hessian_qiskit(6, 3)  # Hessian"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0069d34e",
   "metadata": {},
   "source": [
    "## TensorCircuit 自动微分基准测试\n",
    "\n",
    "我们对上面 Qiskit 部分中定义的相同问题进行基准测试，我们可以看到速度提升！\n",
    "事实上，对于一个中等的 10 量子比特 4 块系统，QFI 评估加速超过 $10^6$ 倍！\n",
    "（注意即时编译的第一次运行的编译时间是可以被平摊的，只需要比较运行时间。在 Qiskit 的情况下，没有即时编译，因此运行时间与暂存时间相同。）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "11cdf0c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def grad_tc(n, l, trials=10):\n",
    "    def f(params):\n",
    "        c = tc.Circuit(n)\n",
    "        for j in range(l):\n",
    "            for i in range(n - 1):\n",
    "                c.cnot(i, i + 1)\n",
    "            for i in range(n):\n",
    "                c.rx(i, theta=params[3 * n * j + i])\n",
    "            for i in range(n):\n",
    "                c.rz(i, theta=params[3 * n * j + i + n])\n",
    "            for i in range(n):\n",
    "                c.rx(i, theta=params[3 * n * j + i + 2 * n])\n",
    "        return tc.backend.real(c.expectation(*[[tc.gates.x(), [i]] for i in range(n)]))\n",
    "\n",
    "    get_grad_tc = tc.backend.jit(tc.backend.grad(f))\n",
    "    return benchmark(get_grad_tc, tc.backend.ones([3 * n * l], dtype=\"float32\"))\n",
    "\n",
    "\n",
    "def qfi_tc(n, l, trials=10):\n",
    "    def s(params):\n",
    "        c = tc.Circuit(n)\n",
    "        for j in range(l):\n",
    "            for i in range(n - 1):\n",
    "                c.cnot(i, i + 1)\n",
    "            for i in range(n):\n",
    "                c.rx(i, theta=params[3 * n * j + i])\n",
    "            for i in range(n):\n",
    "                c.rz(i, theta=params[3 * n * j + i + n])\n",
    "            for i in range(n):\n",
    "                c.rx(i, theta=params[3 * n * j + i + 2 * n])\n",
    "        return c.state()\n",
    "\n",
    "    get_qfi_tc = tc.backend.jit(experimental.qng(s, mode=\"fwd\"))\n",
    "    return benchmark(get_qfi_tc, tc.backend.ones([3 * n * l], dtype=\"float32\"))\n",
    "\n",
    "\n",
    "def hessian_tc(n, l, trials=10):\n",
    "    def f(params):\n",
    "        c = tc.Circuit(n)\n",
    "        for j in range(l):\n",
    "            for i in range(n - 1):\n",
    "                c.cnot(i, i + 1)\n",
    "            for i in range(n):\n",
    "                c.rx(i, theta=params[3 * n * j + i])\n",
    "            for i in range(n):\n",
    "                c.rz(i, theta=params[3 * n * j + i + n])\n",
    "            for i in range(n):\n",
    "                c.rx(i, theta=params[3 * n * j + i + 2 * n])\n",
    "        return tc.backend.real(c.expectation(*[[tc.gates.x(), [i]] for i in range(n)]))\n",
    "\n",
    "    get_hs_tc = tc.backend.jit(tc.backend.hessian(f))\n",
    "    return benchmark(get_hs_tc, tc.backend.ones([3 * n * l], dtype=\"float32\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c3673976",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------\n",
      "tensorflow backend\n",
      "gradient\n",
      "staging time: 15.889095 s\n",
      "running time: 0.001126 s\n",
      "quantum Fisher information\n",
      "WARNING:tensorflow:The dtype of the watched primal must be floating (e.g. tf.float32), got tf.complex64\n",
      "staging time: 53.973453 s\n",
      "running time: 0.002332 s\n",
      "Hessian\n",
      "staging time: 96.066412 s\n",
      "running time: 0.004685 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------\n",
      "jax backend\n",
      "gradient\n",
      "staging time: 4.696845 s\n",
      "running time: 0.000105 s\n",
      "quantum Fisher information\n",
      "staging time: 4.618631 s\n",
      "running time: 0.000386 s\n",
      "Hessian\n",
      "staging time: 23.591966 s\n",
      "running time: 0.001681 s\n"
     ]
    }
   ],
   "source": [
    "for k in [\"tensorflow\", \"jax\"]:\n",
    "    with tc.runtime_backend(k):\n",
    "        print(\"---------------\")\n",
    "        print(\"%s backend\" % k)\n",
    "        print(\"gradient\")\n",
    "        g1, _ = grad_tc(6, 3)\n",
    "        print(\"quantum Fisher information\")\n",
    "        qfi1, _ = qfi_tc(6, 3)\n",
    "        print(\"Hessian\")\n",
    "        hs1, _ = hessian_tc(6, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88ccec20",
   "metadata": {},
   "source": [
    "通过以下检查，两种方法得到的结果是一致的。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4432beba",
   "metadata": {},
   "source": [
    "* 梯度"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8dd03639",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(g0, g1, atol=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37843c96",
   "metadata": {},
   "source": [
    "* 量子费雪信息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "fcf33c7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(qfi0, 4.0 * qfi1, atol=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c775431b",
   "metadata": {},
   "source": [
    "* Hessian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "67b4ce48",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(hs0, hs1, atol=1e-4)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}