{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "281f452f",
   "metadata": {},
   "source": [
    "# VQE on 1D TFIM with Different Hamiltonian Representation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5a444b8",
   "metadata": {},
   "source": [
    "## Overview"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73b70145",
   "metadata": {},
   "source": [
    "For the ground state preparation of a Hamiltonian $H$ in VQE, we need to calculate the expectation value of Hamiltonian $H$, i.e., $\\langle 0^N \\vert U^{\\dagger}(\\theta) H U(\\theta) \\vert 0^N \\rangle$ and update the parameters $\\theta$ in $U(\\theta)$ based on gradient descent. In this tutorial, we will show four ways supported in ``TensorCircuit`` to calculate $\\langle H \\rangle$: \n",
    "\n",
    "1, $\\langle H \\rangle = \\sum_{i} \\langle h_{i} \\rangle$, where $h_{i}$ are the Pauli-string operators;\n",
    "\n",
    "2, $\\langle H \\rangle$ where $H$ is a sparse matrix;\n",
    "\n",
    "3, $\\langle H \\rangle$ where $H$ is a dense matrix;\n",
    "\n",
    "4, expectation value of the Matrix Product Operator (MPO).\n",
    "\n",
    "We consider transverse field ising model (TFIM) here, which reads\n",
    "$$\n",
    "H = \\sum_{i} \\sigma_{i}^{x} \\sigma_{i+1}^{x} - \\sum_{i} \\sigma_{i}^{z},\n",
    "$$\n",
    "where $\\sigma_{i}^{x,z}$ are Pauli matrices of the $i$-th qubit."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cf5e5e5-cd06-47cb-831a-5ef79216df84",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d47fc431",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import tensorcircuit as tc\n",
    "import tensornetwork as tn\n",
    "from tensorcircuit.templates.measurements import operator_expectation\n",
    "from tensorcircuit.quantum import quantum_constructor\n",
    "\n",
    "tc.set_backend(\"tensorflow\")\n",
    "tc.set_dtype(\"complex128\")\n",
    "dtype = np.complex128\n",
    "\n",
    "xx = tc.gates._xx_matrix  # xx gate matrix to be utilized"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28f41a7e-6961-43da-bcd3-d9ed4440cbbd",
   "metadata": {},
   "source": [
    "## Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "95049a99",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 4  # The number of qubits\n",
    "nlayers = 2  # The number of circuit layers\n",
    "ntrials = 2  # The number of random circuit instances"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a8b70bb-a312-4934-813a-035547ca8090",
   "metadata": {},
   "source": [
    "## Parameterized Quantum Circuits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "087d85b4-0db9-428c-b3d9-0b2d7076db67",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tfi_circuit(param, n, nlayers):\n",
    "    c = tc.Circuit(n)\n",
    "    for j in range(nlayers):\n",
    "        for i in range(n - 1):\n",
    "            c.exp1(i, i + 1, unitary=xx, theta=param[2 * j, i])\n",
    "        for i in range(n):\n",
    "            c.rz(i, theta=param[2 * j + 1, i])\n",
    "    return c"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99343a0a",
   "metadata": {},
   "source": [
    "##  Pauli-string Operators"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5506cf20-c011-4b27-af67-10eac4f8e3bf",
   "metadata": {},
   "source": [
    "### Energy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8e6f96a5",
   "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.z(), [i]))  # <Z_i>\n",
    "    for i in range(n - 1):  # OBC\n",
    "        e += j * c.expectation(\n",
    "            (tc.gates.x(), [i]), (tc.gates.x(), [(i + 1) % n])\n",
    "        )  # <X_iX_{i+1}>\n",
    "    return tc.backend.real(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c4d4895d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def vqe_tfim_paulistring(param, n, nlayers):\n",
    "    c = tfi_circuit(param, n, nlayers)\n",
    "    e = tfi_energy(c)\n",
    "    return e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bb0bfb56-2cb0-4a5b-978f-267292f96ea4",
   "metadata": {},
   "outputs": [],
   "source": [
    "vqe_tfim_paulistring_vvag = tc.backend.jit(\n",
    "    tc.backend.vectorized_value_and_grad(vqe_tfim_paulistring)\n",
    ")  # use vvag to get losses and gradients of different random circuit instances"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80cc3f92-4d76-470d-b1c0-65c2736c6383",
   "metadata": {},
   "source": [
    "### Main Optimization Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7ae7e2f0-4a42-405c-b4bd-63f5d1ca8ef9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-03-16 14:09:12.304188: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor([-4.00557571 -3.97372412], shape=(2,), dtype=float64)\n",
      "tf.Tensor([-4.68208061 -4.684804  ], shape=(2,), dtype=float64)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2,), dtype=float64, numpy=array([-4.75683202, -4.73689914])>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def batched_train_step_paulistring_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",
    "    )  # initial parameters\n",
    "    opt = tf.keras.optimizers.Adam(1e-2)\n",
    "    for i in range(maxiter):\n",
    "        e, grad = vqe_tfim_paulistring_vvag(\n",
    "            param.value(), n, nlayers\n",
    "        )  # energy and gradients\n",
    "        opt.apply_gradients([(grad, param)])\n",
    "        if i % 200 == 0:\n",
    "            print(e)\n",
    "    return e\n",
    "\n",
    "\n",
    "batched_train_step_paulistring_tf(ntrials, n, nlayers, 400)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f39f9640-e637-4b78-9519-bde9fed3e2fb",
   "metadata": {},
   "source": [
    "## Sparse Matrix, Dense Matrix, and MPO"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f20dfad-a208-4ef0-af26-623f5530162b",
   "metadata": {},
   "source": [
    "### Hamiltonian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ef40b932-7c02-4ee9-a9e7-866438017150",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tfi_hamiltonian():\n",
    "    h = []\n",
    "    w = []\n",
    "\n",
    "    ### Z\n",
    "    for i in range(n):\n",
    "        h.append([])\n",
    "        w.append(-1.0)  # weight\n",
    "        for j in range(n):\n",
    "            if j == i:\n",
    "                h[i].append(3)\n",
    "            else:\n",
    "                h[i].append(0)\n",
    "\n",
    "    ### XX\n",
    "    for i in range(n - 1):\n",
    "        h.append([])\n",
    "        w.append(1.0)  # weight\n",
    "        for j in range(n):\n",
    "            if j == (i + 1) % n or j == i:\n",
    "                h[i + n].append(1)\n",
    "            else:\n",
    "                h[i + n].append(0)\n",
    "\n",
    "    hamiltonian_sparse = tc.quantum.PauliStringSum2COO(\n",
    "        tf.constant(h, dtype=tf.complex128), tf.constant(w, dtype=tf.complex128)\n",
    "    )  # sparse matrix\n",
    "\n",
    "    hamiltonian_dense = tc.quantum.PauliStringSum2Dense(\n",
    "        tf.constant(h, dtype=tf.complex128), tf.constant(w, dtype=tf.complex128)\n",
    "    )  # dense matrix\n",
    "    return hamiltonian_sparse, hamiltonian_dense"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac737a0b-30bd-472a-a322-21702e3018cc",
   "metadata": {},
   "source": [
    "### Generate QuOperator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "27e58ec1-1424-4b40-bd67-c73b3f20a410",
   "metadata": {},
   "outputs": [],
   "source": [
    "def quoperator_mpo(tfi_mpo):\n",
    "    tfi_mpo = tfi_mpo.tensors\n",
    "\n",
    "    mpo = []\n",
    "    for i in range(n):\n",
    "        mpo.append(tn.Node(tfi_mpo[i]))\n",
    "\n",
    "    for i in range(n - 1):\n",
    "        tn.connect(mpo[i][1], mpo[i + 1][0])\n",
    "\n",
    "    tfi_mpo = quantum_constructor(\n",
    "        [mpo[i][-1] for i in range(n)],  # out_edges\n",
    "        [mpo[i][-2] for i in range(n)],  # in_edges\n",
    "        [],\n",
    "        [mpo[0][0], mpo[-1][1]],  # ignore_edges\n",
    "    )\n",
    "    return tfi_mpo"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75ab67dd-b332-4958-8508-f6cb145debf4",
   "metadata": {},
   "source": [
    "### Energy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d91d0b0a-6ceb-4df1-8b9d-8adce794ba49",
   "metadata": {},
   "outputs": [],
   "source": [
    "def vqe_tfim(param, n, nlayers, hamiltonian):\n",
    "    c = tfi_circuit(param, n, nlayers)\n",
    "    e = operator_expectation(\n",
    "        c, hamiltonian\n",
    "    )  # in operator_expectation, the \"hamiltonian\" can be sparse matrix, dense matrix or mpo\n",
    "    return e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "cf349d64-c3e1-4c92-ad8e-614078bbb6f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "vqe_tfim_vvag = tc.backend.jit(\n",
    "    tc.backend.vectorized_value_and_grad(vqe_tfim)\n",
    ")  # use vvag to get losses and gradients of different random circuit instances"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a23af8cd-a734-448b-b248-56ead572affb",
   "metadata": {},
   "source": [
    "### Main Optimization Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "17a9a6dc-97df-4f43-be6f-a866b771b655",
   "metadata": {},
   "outputs": [],
   "source": [
    "def batched_train_step_tf(batch, n, nlayers, hamiltonian, 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",
    "    )  # initial parameters\n",
    "\n",
    "    opt = tf.keras.optimizers.Adam(1e-2)\n",
    "    for i in range(maxiter):\n",
    "        e, grad = vqe_tfim_vvag(\n",
    "            param.value(), n, nlayers, hamiltonian\n",
    "        )  # energy and gradients\n",
    "        opt.apply_gradients([(grad, param)])\n",
    "        if i % 200 == 0:\n",
    "            print(e)\n",
    "    return e"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "949180b3-e40d-42f8-a224-721b0bc67d0b",
   "metadata": {},
   "source": [
    "### Sparse Matrix, Dense Matrix, and MPO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "cfe4d26e-60c9-46e7-894f-77c8ce4c2d18",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-03-16 14:09:30.874680: I tensorflow/compiler/xla/service/service.cc:171] XLA service 0x7fd94503abc0 initialized for platform Host (this does not guarantee that XLA will be used). Devices:\n",
      "2022-03-16 14:09:30.874726: I tensorflow/compiler/xla/service/service.cc:179]   StreamExecutor device (0): Host, Default Version\n",
      "2022-03-16 14:09:31.014341: I tensorflow/compiler/jit/xla_compilation_cache.cc:351] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.\n"
     ]
    }
   ],
   "source": [
    "(\n",
    "    hamiltonian_sparse,\n",
    "    hamiltonian_dense,\n",
    ") = tfi_hamiltonian()  # hamiltonian: sparse matrix, dense matrix\n",
    "\n",
    "Jx = np.array([1.0 for _ in range(n - 1)])  # strength of xx interaction (OBC)\n",
    "Bz = np.array([1.0 for _ in range(n)])  # strength of transverse field\n",
    "hamiltonian_mpo = tn.matrixproductstates.mpo.FiniteTFI(\n",
    "    Jx, Bz, dtype=dtype\n",
    ")  # matrix product operator\n",
    "hamiltonian_mpo = quoperator_mpo(hamiltonian_mpo)  # generate QuOperator from mpo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "9ac024d4-8ae6-4e35-ad2d-9505aa8353e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Using a while_loop for converting SparseTensorDenseMatMul\n",
      "tf.Tensor([-4.04418884 -3.22012342], shape=(2,), dtype=float64)\n",
      "tf.Tensor([-4.67668625 -4.66761143], shape=(2,), dtype=float64)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2,), dtype=float64, numpy=array([-4.74512239, -4.69965641])>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batched_train_step_tf(ntrials, n, nlayers, hamiltonian_sparse, 400)  # sparse matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d5127255-3cbf-4fa5-8624-28151dd5e158",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor([-3.72705324 -3.99225849], shape=(2,), dtype=float64)\n",
      "tf.Tensor([-4.70773521 -4.7330719 ], shape=(2,), dtype=float64)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2,), dtype=float64, numpy=array([-4.74236986, -4.7559722 ])>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batched_train_step_tf(ntrials, n, nlayers, hamiltonian_dense, 400)  # dense matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "cc67826a-c350-4813-99e9-b0b09e52d7b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor([-3.9129593  -3.44283879], shape=(2,), dtype=float64)\n",
      "tf.Tensor([-4.68271695 -4.67584305], shape=(2,), dtype=float64)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2,), dtype=float64, numpy=array([-4.75283209, -4.75535872])>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batched_train_step_tf(ntrials, n, nlayers, hamiltonian_mpo, 400)  # mpo"
   ]
  }
 ],
 "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
}