parameter_shift_rule.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2ef3e146",
   "metadata": {},
   "source": [
    "## Parameter-Shift Rule for Rotation gates\n",
    "\n",
    "In this tutorial we use the Parameter Shift Rule (PSR) [1, 2] for evaluating the gradients of a variational quantum circuit with respect a variational parameter.\n",
    "\n",
    "#### The parameter shift rule in a nutshell\n",
    "\n",
    "Let's consider a parametrized circuit $\\mathcal{U}(\\vec{\\theta})$, in which we build up an unitary gate of the form:\n",
    "\n",
    "$$ \\mathcal{G} = \\exp \\bigl[-i\\mu G \\bigr] $$\n",
    "\n",
    "which has at most two eigenvalues $\\pm r$. Let's consider an observable $B$ and, finally, let $|q_f \\rangle$ be the state we obtain by applying $\\mathcal{U}$ to $|0\\rangle$.\n",
    "We are interested in evaluating the gradients of the following expression:\n",
    "\n",
    "$$ f(\\mu) \\equiv \\langle q_f | B | q_f \\rangle, $$\n",
    "\n",
    "where we specify that $f$ depends directly on the parameter $\\mu$. We are interested in this result because the expectation value of $B$ is typically involved in computing predictions in quantum machine learning problems. The PSR allows us to calculate the derivative of $f(\\mu)$ with respect to a evaluating $f$ twice more:\n",
    "\n",
    "$$ \\partial_{\\mu} = \\frac{1}{r} \\bigl( f(\\mu^+) - f(\\mu^-) \\bigr), $$\n",
    "\n",
    "where $\\mu^{\\pm} = \\mu \\pm s$ and $s = \\pi / 4r$. Finally, if we pick $G$ from the rotations generators we can use $s=\\pi/2$ and $r=1/2$.\n",
    "\n",
    "#### Loading required features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ebaa326",
   "metadata": {},
   "outputs": [],
   "source": [
    "import qibo\n",
    "import numpy as np\n",
    "from qibo import hamiltonians, gates\n",
    "from qibo.models import Circuit\n",
    "from qibo.hamiltonians.abstract import AbstractHamiltonian\n",
    "from qibo.config import raise_error"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "942c8771",
   "metadata": {},
   "source": [
    "Now we have can write a `parameter_shift` function, in which we take into account an hamiltonian (which is our $B$ observable), the index which identify the target variational parameter, the initial state of the circuit and the wigenvalues of the target observable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57e1fbbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def parameter_shift(\n",
    "    circuit, hamiltonian, parameter_index, generator_eigenval, initial_state=None\n",
    "):\n",
    "    \n",
    "    # inheriting hamiltonian's backend\n",
    "    backend = hamiltonian.backend\n",
    "        \n",
    "    # defining the shift according to the psr\n",
    "    s = np.pi / (4 * generator_eigenval)\n",
    "\n",
    "    # saving original parameters and making a copy\n",
    "    original = np.asarray(circuit.get_parameters()).copy()\n",
    "    shifted = original.copy()\n",
    "\n",
    "    # forward shift and evaluation\n",
    "    shifted[parameter_index] += s\n",
    "    circuit.set_parameters(shifted)\n",
    "\n",
    "    forward = hamiltonian.expectation(backend.execute_circuit(circuit=circuit, initial_state=initial_state).state())\n",
    "\n",
    "    # backward shift and evaluation\n",
    "    shifted[parameter_index] -= 2 * s\n",
    "    circuit.set_parameters(shifted)\n",
    "\n",
    "    backward = hamiltonian.expectation(backend.execute_circuit(circuit=circuit, initial_state=initial_state).state())\n",
    "\n",
    "    # restoring the original circuit\n",
    "    circuit.set_parameters(original)\n",
    "\n",
    "    return generator_eigenval * (forward - backward)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ed630a5",
   "metadata": {},
   "source": [
    "Now we have a `parameter_shift` function and we can use it for calculating the gradients of the expected value of $H$ on the final state with respect to $\\mu$. In order to check the results, we compare them with the same variables evaluated using the `GradientTape()` module of `tensorflow`.\n",
    "\n",
    "For doing this, we need to load `tensorflow` and to activate the appropriate `qibo`'s backend."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ec5c059",
   "metadata": {},
   "outputs": [],
   "source": [
    "# in order to see the difference with tf gradients\n",
    "import tensorflow as tf\n",
    "qibo.set_backend('tensorflow')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1a2af5d",
   "metadata": {},
   "source": [
    "Now we can define the hamiltonian (in this case we use a Pauli Z as observable) and a parametrized circuit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d1f1c7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# defining an observable\n",
    "def hamiltonian(nqubits = 1):\n",
    "    m0 = (1/nqubits)*hamiltonians.Z(nqubits).matrix\n",
    "    ham = hamiltonians.Hamiltonian(nqubits, m0)\n",
    "    return ham\n",
    "\n",
    "# defining a dummy circuit\n",
    "def circuit(nqubits = 1):\n",
    "    c = Circuit(nqubits = 1)\n",
    "    c.add(gates.RY(q = 0, theta = 0))\n",
    "    c.add(gates.RX(q = 0, theta = 0))\n",
    "    c.add(gates.M(0))\n",
    "    return c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bc3ac09",
   "metadata": {},
   "outputs": [],
   "source": [
    "# drawing the circuit\n",
    "c = circuit(1)\n",
    "c.draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2af2b1fc",
   "metadata": {},
   "source": [
    "This is the moment to write a function which returns the `tensorflow` values of the gradients. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19051140",
   "metadata": {},
   "outputs": [],
   "source": [
    "# using GradientTape to benchmark\n",
    "def gradient_tape(params):\n",
    "    params = tf.Variable(params)\n",
    "    \n",
    "    with tf.GradientTape() as tape:\n",
    "        c = circuit(nqubits = 1)\n",
    "        c.set_parameters(params)\n",
    "        h = hamiltonian()\n",
    "        expected_value = h.expectation(c.execute().state()) \n",
    "    \n",
    "    grads = tape.gradient(expected_value, [params])\n",
    "    return grads"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5493f89",
   "metadata": {},
   "source": [
    "In order to check the difference, we randomly generate some parameters and we impose them as variational parameters of the circuit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1956f455",
   "metadata": {},
   "outputs": [],
   "source": [
    "# some parameters\n",
    "test_params = np.random.randn(2)\n",
    "c.set_parameters(test_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a29d4218",
   "metadata": {},
   "source": [
    "Here we are!\n",
    "\n",
    "Now we can calculate the gradients using the two methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad97dc3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "test_hamiltonian = hamiltonian()\n",
    "\n",
    "# running the psr with respect to the two parameters\n",
    "grad_0 = parameter_shift(circuit = c, hamiltonian = test_hamiltonian, parameter_index = 0, generator_eigenval = 0.5)\n",
    "grad_1 = parameter_shift(circuit = c, hamiltonian = test_hamiltonian, parameter_index = 1, generator_eigenval = 0.5)\n",
    "\n",
    "tf_grads = gradient_tape(test_params)\n",
    "\n",
    "print('Test gradient with respect params[0] with PSR: ', grad_0.numpy())\n",
    "print('Test gradient with respect params[0] with tf:  ', tf_grads[0][0].numpy())\n",
    "print('Test gradient with respect params[0] with PSR: ', grad_1.numpy())\n",
    "print('Test gradient with respect params[0] with tf:  ', tf_grads[0][1].numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08887a64",
   "metadata": {},
   "source": [
    "As you can see, the values are identical!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "492d5726",
   "metadata": {},
   "source": [
    "### References\n",
    "\n",
    "[1] Kosuke Mitarai, Makoto Negoro, Masahiro Kitagawa, Keisuke Fujii, *Quantum Circuit Learning*, (2018), [arXiv:1803.00745v3](https://arxiv.org/abs/1803.00745)\n",
    "\n",
    "[2] Maria Schuld, Ville Bergholm, Christian Gogolin, Josh Izaac, Nathan Killoran, *Evaluating analytic gradients on quantum hardware*, (2018), [arXiv:1811.11184v1](https://arxiv.org/abs/1811.11184)"
   ]
  }
 ],
 "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}