{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c610ec7d",
   "metadata": {},
   "source": [
    "# Gradient Evaluation Efficiency Comparison"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edb4fc7f",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "In this tutorial, we compare the efficiency of gradient and gradient-like object (such as quantum Fisher information) evaluation via automatical differentiation framework provided by TensorCircuit and the traditional parameter shift framework provided by Qiskit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3f56164",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "We import necessary packages and modules from Qiskit and 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 Gradient Framework Benchmark\n",
    "\n",
    "Since Qiskit is **TOO** slow in terms of gradient evaluation, we use small systems to do the benchmark in Jupyter to save time,\n",
    "for larger size and deep circuits, the efficiency difference will become more evident.\n",
    "\n",
    "The three gradient like tasks are Gradient, Quantum Fisher Information(QFI), and Hessian evaluation."
   ]
  },
  {
   "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",
    "    # Define the expectation value corresponding to the energy\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",
    "    # Define the expectation value corresponding to the energy\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)  # gradient"
   ]
  },
  {
   "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 Automatic Differentiation Benchmark\n",
    "\n",
    "We benchmark on the same problems defined in the Qiskit part above, and we can see the speed boost!\n",
    "In fact, for a moderate 10-qubit 4-blocks system, QFI evaluation is accelerated more than $10^6$ times!\n",
    "(Note how staging time for jit can be amortized and only running time relevant. In the Qiskit case, there is no jit and thus the running time is the same as the staging time.)"
   ]
  },
  {
   "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": [
    "The results obtained from the two methods are consistent by the following checks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4432beba",
   "metadata": {},
   "source": [
    "* Gradient"
   ]
  },
  {
   "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": [
    "* Quantum Fisher Information(QFI)"
   ]
  },
  {
   "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
}