docs/source/whitepaper/4-gradient-optimization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1dd80d43",
   "metadata": {},
   "source": [
    "# Gradient and Variational Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66167c3f",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "TensorCircuit is designed to make optimization of parameterized quantum gates easy, fast, and convenient. In this note, we review how to obtain circuit gradients and run variational optimization."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97754f01",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6af689b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.optimize as optimize\n",
    "import tensorflow as tf\n",
    "import tensorcircuit as tc\n",
    "\n",
    "K = tc.set_backend(\"tensorflow\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba8d0487",
   "metadata": {},
   "source": [
    "## PQC\n",
    "\n",
    " Consider a variational circuit acting on $n$ qubits, and consisting of $k$ layers, where each layer comprises parameterized $e^{i\\theta X\\otimes X}$ gates between neighboring qubits followed by a sequence of single qubit parameterized $Z$ and $X$ rotations. We now show how to implement such circuits in TensorCircuit, and how to use one of the machine learning backends to compute cost functions and gradients easily and efficiently.\n",
    "\n",
    "The circuit for general $n,k$ and set of parameters can be defined as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "21cef70c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def qcircuit(n, k, params):\n",
    "    c = tc.Circuit(n)\n",
    "    for j in range(k):\n",
    "        for i in range(n - 1):\n",
    "            c.exp1(\n",
    "                i, i + 1, theta=params[j * (3 * n - 1) + i], unitary=tc.gates._xx_matrix\n",
    "            )\n",
    "        for i in range(n):\n",
    "            c.rz(i, theta=params[j * (3 * n - 1) + n - 1 + i])\n",
    "            c.rx(i, theta=params[j * (3 * n - 1) + 2 * n - 1 + i])\n",
    "    return c"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae87b1a7",
   "metadata": {},
   "source": [
    "As an example, we take $n=3, k=2$, set TensorFlow as our backend, and define an energy cost function to minimize\n",
    "$$E = \\langle X_0 X_1\\rangle_\\theta + \\langle X_1 X_2\\rangle_\\theta.$$ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c75e7d80",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 3\n",
    "k = 2\n",
    "\n",
    "\n",
    "def energy(params):\n",
    "    c = qcircuit(n, k, params)\n",
    "    e = c.expectation_ps(x=[0, 1]) + c.expectation_ps(x=[1, 2])\n",
    "    return K.real(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08baf09f",
   "metadata": {},
   "source": [
    "## Grad and JIT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2390d0c5",
   "metadata": {},
   "source": [
    "Using the ML backend support for automatic differentiation, we can now quickly compute both the energy and the gradient of the energy with respect to the parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5aedf3cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "energy_val_grad = K.value_and_grad(energy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27da172d",
   "metadata": {},
   "source": [
    "This creates a function that given a set of parameters as input, returns both the energy and the gradient of the energy. If only the gradient is desired, then this can be computed by ``K.grad(energy)``. While we could run the above code directly on a set of parameters, if multiple evaluations of the energy will be performed, significant time savings can be had by using a just-in-time compiled version of the function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0bd4f9da",
   "metadata": {},
   "outputs": [],
   "source": [
    "energy_val_grad_jit = K.jit(energy_val_grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bdff95f",
   "metadata": {},
   "source": [
    "With ``K.jit``, the initial evaluation of the energy and gradient may take longer, but subsequent evaluations will be noticeably faster than non-jitted code. We recommend always using ``jit`` as long as the function is \"tensor-in, tensor-out\", and we have worked hard to make all aspects of the circuit simulator compatible with JIT."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50dbd1a5",
   "metadata": {},
   "source": [
    "## Optimization via ML Backend\n",
    "\n",
    "With the energy function and gradients available, optimization of the parameters is straightforward.  Below is an example of how to do this via stochastic gradient descent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1283344e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i=0, energy=0.11897378414869308\n",
      "i=10, energy=-0.3692811131477356\n",
      "i=20, energy=-0.7194114923477173\n",
      "i=30, energy=-0.904697597026825\n",
      "i=40, energy=-1.013866662979126\n",
      "i=50, energy=-1.1042678356170654\n",
      "i=60, energy=-1.1998062133789062\n",
      "i=70, energy=-1.308410406112671\n",
      "i=80, energy=-1.4276418685913086\n",
      "i=90, energy=-1.5474387407302856\n"
     ]
    }
   ],
   "source": [
    "learning_rate = 2e-2\n",
    "opt = K.optimizer(tf.keras.optimizers.SGD(learning_rate))\n",
    "\n",
    "\n",
    "def grad_descent(params, i):\n",
    "    val, grad = energy_val_grad_jit(params)\n",
    "    params = opt.update(grad, params)\n",
    "    if i % 10 == 0:\n",
    "        print(f\"i={i}, energy={val}\")\n",
    "    return params\n",
    "\n",
    "\n",
    "params = K.implicit_randn(k * (3 * n - 1))\n",
    "for i in range(100):\n",
    "    params = grad_descent(params, i)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f97011f8",
   "metadata": {},
   "source": [
    "## Optimization via Scipy Interface\n",
    "\n",
    "An alternative to using the machine learning backends for the optimization is to use SciPy.\n",
    "This can be done via the ``scipy_interface`` API call and allows for gradient-based (e.g. BFGS) and non-gradient-based (e.g. COBYLA) optimizers to be used, which are not available via the ML backends."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b4cc8845",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/shixin/Cloud/newwork/quantum-information/codebases/tensorcircuit/tensorcircuit/interfaces.py:237: ComplexWarning: Casting complex values to real discards the imaginary part\n",
      "  scipy_gs = scipy_gs.astype(np.float64)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "      fun: -2.000000476837158\n",
       " hess_inv: <16x16 LbfgsInvHessProduct with dtype=float64>\n",
       "      jac: array([ 2.43186951e-04, -1.50322914e-04,  8.94665718e-05,  1.18807920e-05,\n",
       "        2.95639038e-05,  1.19209290e-07, -5.96046448e-08, -2.98023224e-08,\n",
       "        0.00000000e+00, -1.19209290e-07,  3.90738450e-07,  9.34305717e-07,\n",
       "       -8.22039729e-05,  1.19209290e-07,  0.00000000e+00,  0.00000000e+00])\n",
       "  message: 'CONVERGENCE: REL_REDUCTION_OF_F_<=_FACTR*EPSMCH'\n",
       "     nfev: 60\n",
       "      nit: 19\n",
       "     njev: 60\n",
       "   status: 0\n",
       "  success: True\n",
       "        x: array([ 2.35625520e+00,  7.85409154e-01,  1.57088576e+00,  2.10625989e-05,\n",
       "       -1.57088425e+00, -1.70256902e+00, -5.33743572e-01,  3.11436816e-01,\n",
       "        1.26543793e+00,  1.91663337e+00, -1.15901008e-07, -1.76623396e-05,\n",
       "       -1.59972887e-04, -8.97072367e-01,  1.79929630e+00, -9.67278961e-01])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_scipy = tc.interfaces.scipy_interface(energy, shape=[k * (3 * n - 1)], jit=True)\n",
    "params = K.implicit_randn(k * (3 * n - 1))\n",
    "r = optimize.minimize(f_scipy, params, method=\"L-BFGS-B\", jac=True)\n",
    "r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6041d1e7",
   "metadata": {},
   "source": [
    "The first line above specifies the shape of the parameters to be supplied to the function to be minimized, which here is the energy function.  The ``jit=True`` argument automatically takes care of jitting the energy function.  Gradient-free optimization can similarly be performed efficiently by supplying the ``gradient=False`` argument to  ``scipy_interface``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2f09d59e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "     fun: -1.9999911785125732\n",
       "   maxcv: 0.0\n",
       " message: 'Optimization terminated successfully.'\n",
       "    nfev: 386\n",
       "  status: 1\n",
       " success: True\n",
       "       x: array([ 7.87597857e-01, -5.14158452e-01, -1.56560250e+00, -3.15230777e-04,\n",
       "        9.91532990e-01,  5.95588091e-01,  1.38523058e+00, -3.59642968e-04,\n",
       "       -3.23365306e-01, -4.16465772e-01, -7.32259085e-03,  6.53997758e-05,\n",
       "        7.71203778e-01,  2.46256921e+00,  8.78602039e-01, -3.51989842e-01])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_scipy = tc.interfaces.scipy_interface(\n",
    "    energy, shape=[k * (3 * n - 1)], jit=True, gradient=False\n",
    ")\n",
    "params = K.implicit_randn(k * (3 * n - 1))\n",
    "r = optimize.minimize(f_scipy, params, method=\"COBYLA\")\n",
    "r"
   ]
  }
 ],
 "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
}