docs/source/tutorials/tfim_vqe.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# VQE on 1D TFIM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "The main aim of this tutorial is not about the physics perspective of VQE, instead, we demonstrate\n",
    "the main ingredients of TensorCircuit by this simple VQE toy model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Background\n",
    "\n",
    "Basically, we train a parameterized quantum circuit with repetitions of $e^{i\\theta} ZZ$ and $e^{i\\theta X}$ layers as $U(\\rm{\\theta})$. And the objective to be minimized is this task is $\\mathcal{L}(\\rm{\\theta})=\\langle 0^n\\vert U(\\theta)^\\dagger H U(\\theta)\\vert 0^n\\rangle$. The Hamiltonian is from TFIM as $H = \\sum_{i} Z_iZ_{i+1} -\\sum_i X_i$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import jax\n",
    "from jax.config import config\n",
    "\n",
    "config.update(\"jax_enable_x64\", True)\n",
    "from jax import numpy as jnp\n",
    "from jax.experimental import optimizers\n",
    "import tensorcircuit as tc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To enable automatic differentiation support, we should set the TensorCircuit backend beyond the default one \"NumPy\".\n",
    "And we can also set the high precision complex128 for the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "tc.set_backend(\"tensorflow\")\n",
    "tc.set_dtype(\"complex128\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "complex dtype of simulation: complex128 \n",
      "real dtype of simulation: float64 \n",
      "backend package of simulation: tensorflow\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    \"complex dtype of simulation:\",\n",
    "    tc.dtypestr,\n",
    "    \"\\nreal dtype of simulation:\",\n",
    "    tc.rdtypestr,\n",
    "    \"\\nbackend package of simulation:\",\n",
    "    tc.backend.name,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 1.  0.  0.  0.]\n",
      " [ 0. -1.  0. -0.]\n",
      " [ 0.  0. -1. -0.]\n",
      " [ 0. -0. -0.  1.]]\n"
     ]
    }
   ],
   "source": [
    "# zz gate matrix to be utilized\n",
    "zz = np.kron(tc.gates._z_matrix, tc.gates._z_matrix)\n",
    "print(zz)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Higher-level API"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first design the Hamiltonian energy expectation function with the input as quantum circuit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tfi_energy(c: tc.Circuit, j: float = 1.0, h: float = -1.0):\n",
    "    e = 0.0\n",
    "    n = c._nqubits\n",
    "    for i in range(n):\n",
    "        e += h * c.expectation((tc.gates.x(), [i]))  # <X_i>\n",
    "    for i in range(n - 1):  # OBC\n",
    "        e += j * c.expectation(\n",
    "            (tc.gates.z(), [i]), (tc.gates.z(), [(i + 1) % n])\n",
    "        )  # <Z_iZ_{i+1}>\n",
    "    return tc.backend.real(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we make the quantum function with $\\rm{\\theta}$ as input and energy expectation $\\mathcal{L}$ as output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vqe_tfim(param, n, nlayers):\n",
    "    c = tc.Circuit(n)\n",
    "    paramc = tc.backend.cast(\n",
    "        param, tc.dtypestr\n",
    "    )  # We assume the input param with dtype float64\n",
    "    for i in range(n):\n",
    "        c.H(i)\n",
    "    for j in range(nlayers):\n",
    "        for i in range(n - 1):\n",
    "            c.exp1(i, i + 1, unitary=zz, theta=paramc[2 * j, i])\n",
    "        for i in range(n):\n",
    "            c.rx(i, theta=paramc[2 * j + 1, i])\n",
    "    e = tfi_energy(c)\n",
    "    return e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To train the parameterized circuit, we should utilize the gradient information $\\frac{\\partial \\mathcal{L}}{\\partial \\rm{\\theta}}$ with gradient descent.\n",
    "We also use ``jit`` to wrap the value and grad function for a substantial speedup. Note how (1, 2) args of ``vqe_tfim`` are labeled as static since they are just integers for qubit number and layer number instead of tensors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "vqe_tfim_vag = tc.backend.jit(\n",
    "    tc.backend.value_and_grad(vqe_tfim), static_argnums=(1, 2)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(-5.2044235531710905, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.168460907768175, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.229007202330065, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.2368387790165105, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.246179597523659, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.262673966580785, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.286129364991173, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.291252895716095, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.2930457160020765, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.293225326335964, shape=(), dtype=float64)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=float64, numpy=-7.293297606006469>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def train_step_tf(n, nlayers, maxiter=10000):\n",
    "    param = tf.Variable(\n",
    "        initial_value=tf.random.normal(\n",
    "            shape=[nlayers * 2, n], stddev=0.1, dtype=getattr(tf, tc.rdtypestr)\n",
    "        )\n",
    "    )\n",
    "    opt = tf.keras.optimizers.Adam(1e-2)\n",
    "    for i in range(maxiter):\n",
    "        e, grad = vqe_tfim_vag(param, n, nlayers)\n",
    "        opt.apply_gradients([(grad, param)])\n",
    "        if i % 200 == 0:\n",
    "            print(e)\n",
    "    return e\n",
    "\n",
    "\n",
    "train_step_tf(6, 3, 2000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Batched VQE Example\n",
    "\n",
    "We can even run a batched version of VQE optimization, namely, we simultaneously optimize parameterized circuits for different random initializations, so that we can try our best to avoid local minimums and locate the best of the converged energies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "vqe_tfim_vvag = tc.backend.jit(\n",
    "    tc.backend.vectorized_value_and_grad(vqe_tfim), static_argnums=(1, 2)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[-4.56780182 -5.32411397 -5.34948039 -5.49728838 -5.51974631 -4.89464895\n",
      " -5.23113926 -5.70097167 -5.4384308  -5.27898261 -4.73926061 -5.43748391\n",
      " -5.02246224 -4.46749643 -5.34320604 -5.29828815], shape=(16,), dtype=float64)\n",
      "tf.Tensor(\n",
      "[-7.15906597 -7.20867528 -7.16615816 -7.16164269 -7.15427498 -7.17176534\n",
      " -7.15677645 -7.19769858 -7.1876547  -7.17160745 -7.14313137 -7.16458417\n",
      " -7.12556993 -7.1043696  -7.17233218 -7.17955502], shape=(16,), dtype=float64)\n",
      "tf.Tensor(\n",
      "[-7.22332735 -7.28775096 -7.22854626 -7.28800389 -7.22006811 -7.2773814\n",
      " -7.22241623 -7.23446324 -7.23115651 -7.23081143 -7.25399986 -7.26564648\n",
      " -7.16463543 -7.27854832 -7.23574558 -7.28935649], shape=(16,), dtype=float64)\n",
      "tf.Tensor(\n",
      "[-7.23956454 -7.29093555 -7.23464822 -7.2914774  -7.22326999 -7.29014637\n",
      " -7.24891067 -7.2505597  -7.23879431 -7.23826618 -7.28737831 -7.29193732\n",
      " -7.22649018 -7.29136679 -7.25276205 -7.29214669], shape=(16,), dtype=float64)\n",
      "tf.Tensor(\n",
      "[-7.24561853 -7.29413883 -7.23950499 -7.29230127 -7.22749993 -7.29051998\n",
      " -7.28702174 -7.289441   -7.25016979 -7.26370483 -7.29320874 -7.29451577\n",
      " -7.22882824 -7.29213765 -7.27040912 -7.29358236], shape=(16,), dtype=float64)\n",
      "tf.Tensor(\n",
      "[-7.24997971 -7.294748   -7.25420008 -7.29271584 -7.24577837 -7.29082466\n",
      " -7.29171805 -7.29016935 -7.28645108 -7.29170429 -7.29499124 -7.29520514\n",
      " -7.23115011 -7.29305292 -7.28793637 -7.2949226 ], shape=(16,), dtype=float64)\n",
      "tf.Tensor(\n",
      "[-7.25300306 -7.29512508 -7.28240557 -7.29287622 -7.28264095 -7.29125472\n",
      " -7.29399162 -7.29066326 -7.29233232 -7.29290676 -7.29521188 -7.29530935\n",
      " -7.23475933 -7.29429836 -7.29053038 -7.29559969], shape=(16,), dtype=float64)\n",
      "tf.Tensor(\n",
      "[-7.25706762 -7.29527205 -7.29168082 -7.29292216 -7.29221412 -7.29183888\n",
      " -7.29474989 -7.29119684 -7.29306204 -7.29300514 -7.29525079 -7.29538907\n",
      " -7.24226472 -7.29544118 -7.29086393 -7.29576418], shape=(16,), dtype=float64)\n",
      "tf.Tensor(\n",
      "[-7.26218683 -7.29529443 -7.29438024 -7.29294969 -7.29426403 -7.29243576\n",
      " -7.29491073 -7.29171828 -7.29384631 -7.29301505 -7.29527304 -7.29545727\n",
      " -7.25772904 -7.29581457 -7.2910381  -7.29582192], shape=(16,), dtype=float64)\n",
      "tf.Tensor(\n",
      "[-7.26654138 -7.29529938 -7.29515899 -7.29297612 -7.29499748 -7.29281659\n",
      " -7.29500531 -7.29227399 -7.29455691 -7.29302087 -7.29529063 -7.29551808\n",
      " -7.2892432  -7.29593328 -7.29120759 -7.29585771], shape=(16,), dtype=float64)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(16,), dtype=float64, numpy=\n",
       "array([-7.29011428, -7.29530356, -7.29549915, -7.29300424, -7.29529224,\n",
       "       -7.29296047, -7.29508027, -7.29270021, -7.29499648, -7.29302725,\n",
       "       -7.29530574, -7.2955733 , -7.29593608, -7.29604964, -7.29138482,\n",
       "       -7.29589095])>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def batched_train_step_tf(batch, n, nlayers, maxiter=10000):\n",
    "    param = tf.Variable(\n",
    "        initial_value=tf.random.normal(\n",
    "            shape=[batch, nlayers * 2, n], stddev=0.1, dtype=getattr(tf, tc.rdtypestr)\n",
    "        )\n",
    "    )\n",
    "    opt = tf.keras.optimizers.Adam(1e-2)\n",
    "    for i in range(maxiter):\n",
    "        e, grad = vqe_tfim_vvag(param, n, nlayers)\n",
    "        opt.apply_gradients([(grad, param)])\n",
    "        if i % 200 == 0:\n",
    "            print(e)\n",
    "    return e\n",
    "\n",
    "\n",
    "batched_train_step_tf(16, 6, 3, 2000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Different Backends\n",
    "\n",
    "We can change the backends at runtime without even changing one line of the code!\n",
    "\n",
    "However, in normal user cases, we strongly recommend the users stick to one backend in one jupyter or python script.\n",
    "One can enjoy the facility provided by other backends by changing the ``set_backend`` line and running the same script again. This approach is much safer than using multiple backends in the same file unless you know the lower-level details of TensorCircuit enough."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "tc.set_backend(\"jax\")  # change to jax backend"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "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": [
      "[-5.67575948 -5.44768444 -5.7821556  -5.36699503 -5.00485098 -5.59416181\n",
      " -5.13421084 -5.70462279 -5.73699416 -5.25819658 -4.70729299 -5.82823766\n",
      " -5.69154358 -5.51112311 -5.46091316 -5.31649863]\n",
      "[-7.16831387 -7.17873365 -7.21905991 -7.17714641 -7.21910053 -7.17729778\n",
      " -7.23594046 -7.1978075  -7.2311691  -7.18566164 -7.15141273 -7.1760751\n",
      " -7.20727055 -7.22174427 -7.15227955 -7.15343225]\n",
      "[-7.24047827 -7.23486717 -7.26382185 -7.25267406 -7.23938877 -7.24135079\n",
      " -7.28655961 -7.24413064 -7.28070556 -7.24825735 -7.23400189 -7.25234153\n",
      " -7.25756263 -7.2505181  -7.22647645 -7.2589444 ]\n",
      "[-7.28642159 -7.23707926 -7.28988032 -7.28627451 -7.28716418 -7.25068739\n",
      " -7.29122589 -7.2510777  -7.2906953  -7.25976327 -7.23891735 -7.29227009\n",
      " -7.28973637 -7.26238069 -7.245065   -7.29155041]\n",
      "[-7.29198674 -7.24196434 -7.29188725 -7.29243688 -7.2926968  -7.26254168\n",
      " -7.29233808 -7.26729904 -7.29277165 -7.28066403 -7.24315235 -7.29344766\n",
      " -7.2920645  -7.26717433 -7.26959622 -7.29307748]\n",
      "[-7.29320541 -7.27162341 -7.29245991 -7.2934821  -7.29360574 -7.27103573\n",
      " -7.29311302 -7.29213505 -7.29356392 -7.29162927 -7.24981922 -7.29384423\n",
      " -7.2930642  -7.27323089 -7.29156252 -7.293684  ]\n",
      "[-7.29384132 -7.29130123 -7.29333191 -7.29442748 -7.29396942 -7.29011734\n",
      " -7.2936549  -7.2929216  -7.29430937 -7.29288465 -7.27758063 -7.29446428\n",
      " -7.2939575  -7.27958774 -7.29260011 -7.29408982]\n",
      "[-7.29433271 -7.29273603 -7.29467579 -7.29507754 -7.29438517 -7.29298445\n",
      " -7.2940647  -7.29297896 -7.29494014 -7.29348052 -7.29142477 -7.29524322\n",
      " -7.29455267 -7.28904449 -7.29319333 -7.29437335]\n",
      "[-7.29480095 -7.29294578 -7.29507043 -7.29530148 -7.29479745 -7.29390014\n",
      " -7.29439439 -7.29304737 -7.29520739 -7.29397645 -7.29276852 -7.2955135\n",
      " -7.29492853 -7.29145529 -7.29390729 -7.29462868]\n",
      "[-7.29530498 -7.29299558 -7.29520643 -7.29537379 -7.29505127 -7.2945667\n",
      " -7.29466477 -7.29312996 -7.29529573 -7.29447499 -7.29327567 -7.29560796\n",
      " -7.29520605 -7.29179062 -7.29484071 -7.29489978]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DeviceArray([-7.29575181, -7.29302602, -7.29529976, -7.29541094,\n",
       "             -7.29517778, -7.29488194, -7.29487651, -7.29323608,\n",
       "             -7.29532772, -7.29494698, -7.29369784, -7.29567791,\n",
       "             -7.29534388, -7.29187906, -7.29536221, -7.29516005],            dtype=float64)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vqe_tfim_vvag = tc.backend.jit(\n",
    "    tc.backend.vectorized_value_and_grad(vqe_tfim), static_argnums=(1, 2)\n",
    ")\n",
    "\n",
    "\n",
    "def batched_train_step_jax(batch, n, nlayers, maxiter=10000):\n",
    "    key = jax.random.PRNGKey(42)\n",
    "    param = jax.random.normal(key, shape=[batch, nlayers * 2, n]) * 0.1\n",
    "    opt_init, opt_update, get_params = optimizers.adam(step_size=1e-2)\n",
    "    opt_state = opt_init(param)\n",
    "\n",
    "    def update(i, opt_state):\n",
    "        param = get_params(opt_state)\n",
    "        (value, gradient) = vqe_tfim_vvag(param, n, nlayers)\n",
    "        return value, opt_update(i, gradient, opt_state)\n",
    "\n",
    "    for i in range(maxiter):\n",
    "        value, opt_state = update(i, opt_state)\n",
    "        param = get_params(opt_state)\n",
    "        if i % 200 == 0:\n",
    "            print(value)\n",
    "    return value\n",
    "\n",
    "\n",
    "batched_train_step_jax(16, 6, 3, 2000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lower-level API\n",
    "\n",
    "The higher-level API under the namespace of ``TensorCircuit`` provides a unified framework to do linear algebra and automatic differentiation which is backend agnostic.\n",
    "\n",
    "One may also use the related APIs (ops, AD-related, jit-related) directly provided by TensorFlow or Jax, as long as one is ok to stick with one fixed backend. See the tensorflow backend example below.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "tc.set_backend(\"tensorflow\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tfi_energy(c: tc.Circuit, j: float = 1.0, h: float = -1.0):\n",
    "    e = 0.0\n",
    "    n = c._nqubits\n",
    "    for i in range(n):\n",
    "        e += h * c.expectation((tc.gates.x(), [i]))  # <X_i>\n",
    "    for i in range(n - 1):  # OBC\n",
    "        e += j * c.expectation(\n",
    "            (tc.gates.z(), [i]), (tc.gates.z(), [(i + 1) % n])\n",
    "        )  # <Z_iZ_{i+1}>\n",
    "    return tf.math.real(e)\n",
    "\n",
    "\n",
    "def vqe_tfim(param, n, nlayers):\n",
    "    c = tc.Circuit(n)\n",
    "    paramc = tf.cast(param, tf.complex128)\n",
    "    for i in range(n):\n",
    "        c.H(i)\n",
    "    for j in range(nlayers):\n",
    "        for i in range(n - 1):\n",
    "            c.exp1(i, i + 1, unitary=zz, theta=paramc[2 * j, i])\n",
    "        for i in range(n):\n",
    "            c.rx(i, theta=paramc[2 * j + 1, i])\n",
    "    e = tfi_energy(c)\n",
    "    return e\n",
    "\n",
    "\n",
    "@tf.function\n",
    "def vqe_tfim_vag(param, n, nlayers):\n",
    "    with tf.GradientTape() as tape:\n",
    "        tape.watch(param)\n",
    "        v = vqe_tfim(param, n, nlayers)\n",
    "    grad = tape.gradient(v, param)\n",
    "    return v, grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(-5.5454151788179376, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.167693061786028, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.254761404891117, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.290050014550046, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.29133881232428, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.2918048286324915, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.292590929769901, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.294195015132205, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.295013538531699, shape=(), dtype=float64)\n",
      "tf.Tensor(-7.2951174084838835, shape=(), dtype=float64)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=float64, numpy=-7.295170441938537>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_step_tf(6, 3, 2000)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "97eeaec2e58aad59d290a91b8916e9d88dd2ea1991ab7757ddd07fd169009203"
  },
  "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": 2
}