docs/source/tutorials/vqe_h2o.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4009fbd4",
   "metadata": {},
   "source": [
    "# Variational Quantum Eigensolver (VQE) on Molecules"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0b04a30",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Overview\n",
    "\n",
    "VQE is a variational algorithm for calculating the ground state of some given hamiltonian H which we call it $\\psi_g$ that satisfied $H \\left|\\psi_g\\right> =E_g\\left|\\psi_g\\right>$. For an arbitrary normalized wavefunction $\\psi_f$, the expectation value $\\left<\\psi_f|H|\\psi_f \\right>$ is always not lower than the ground state energy unless $\\psi_f = \\psi_g$ to some phase factor (here we assume there is no degeneracy in ground state). Based on that fact, if we use a parameterized wavefunction $\\psi_\\theta$, e.g. given by a parameterized quantum circuit (PQC) with parameters $\\theta$, we can give an approximation for the ground state enery and wavefunction by minimizing the expectation value of $H$. In practical quantum hardware, this algorithm can be realized in a quantum-neural hybrid paradigm with the gradient calculated using finite difference or paremeter shift in quantum hardware and the optimization using gradient descent method in classical computer. While in a numerical simulation, we can just calculate the gradients using automatic differentiation. \n",
    "\n",
    "Calculating the ground state energy for molecules is often important for quantum chemistry tasks since it can be used to find out the atom structure of the molecules. In the simulation of molecules, we do not consider the motion of nuclei which means we fix the nuclear coordinates of its constituent atoms. We only consider the electrons in the molecules since the nuclei are way heavier than the electrons and thus the energy carried by phonons is negligible or can be reconsidered using Born-Oppenheimer approximation. Strictly speaking, the eletrons lives in continuous space, thus the Hilbert space is of infinite dimensions. To conduct a practical calculation, we only preserve some important single-particle basis, e.g. the low energy atomic orbitals. In the second quantization formalism, we can represent these atomic orbitals as $c_i^\\dagger|0>$. By considering the interactions of nuclei and electrons as background and also the electron-electron interaction, a molecules hamiltonian can in generally be represented as $H = \\sum_{i, j} h_{i,j} c_i^\\dagger c_j + \\sum_{i, j, k, l} \\alpha_{i, j, k, l} c_i^\\dagger c_j^\\dagger c_k c_l$. Notice that the spin index is also absorbed into the orbital index. There are many softwares that can give these parameters in H such as `pyscf` which we will use later in this tutorial. Now we have a fermionic description for moleculars. By using a mapping from fermionic operators to spin operators such as Jordan-Wigner transformation or Bravyi-Kitaev transformation, we can map the fermionic hamiltonian to a spin hamiltonian which is more compatible with quantum computer. For a spin hamiltonian, we can easily use a PQC to construct a trail wavefunction and conduct the VQE algorithm. In the following part of this tutorial, we will demonstrate a complete example of how to use TensorCircuit to simulate VQE algorithm on Molecules."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12886697",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "We should first ``pip install openfermion openfermionpyscf`` to generate fermionic and qubit Hamiltonian of H2O molecule based on quantum chemistry computation provided by openfermion and pyscf."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9e28c969",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from openfermion.chem import MolecularData\n",
    "from openfermion.transforms import (\n",
    "    get_fermion_operator,\n",
    "    jordan_wigner,\n",
    "    binary_code_transform,\n",
    "    checksum_code,\n",
    "    reorder,\n",
    ")\n",
    "from openfermion.chem import geometry_from_pubchem\n",
    "from openfermion.utils import up_then_down\n",
    "from openfermion.linalg import LinearQubitOperator\n",
    "from openfermionpyscf import run_pyscf\n",
    "import tensorflow as tf\n",
    "\n",
    "import tensorcircuit as tc\n",
    "\n",
    "K = tc.set_backend(\"tensorflow\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8464f4f5",
   "metadata": {},
   "source": [
    "## Generate Hamiltonian"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6677497",
   "metadata": {},
   "source": [
    "* Get molecule energy info and molecule orbitals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "40574813",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-75.0155301894916 -75.01540899923558 -74.96444758276998\n"
     ]
    }
   ],
   "source": [
    "multiplicity = 1\n",
    "basis = \"sto-3g\"\n",
    "# 14 spin orbitals for H2O\n",
    "geometry = geometry_from_pubchem(\"h2o\")\n",
    "description = \"h2o\"\n",
    "molecule = MolecularData(geometry, basis, multiplicity, description=description)\n",
    "molecule = run_pyscf(molecule, run_mp2=True, run_cisd=True, run_ccsd=True, run_fci=True)\n",
    "print(molecule.fci_energy, molecule.ccsd_energy, molecule.hf_energy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc75e810",
   "metadata": {},
   "source": [
    "* Get Fermionic Hamiltonian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2c526993",
   "metadata": {},
   "outputs": [],
   "source": [
    "mh = molecule.get_molecular_hamiltonian()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "df6f82e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "fh = get_fermion_operator(mh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d77ca301",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-32.68991541360029\n"
     ]
    }
   ],
   "source": [
    "print(fh.terms[((0, 1), (0, 0))])  # coefficient of C0^\\dagger C_0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3603a455",
   "metadata": {},
   "source": [
    "* Transform into qubit Hamiltonian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0d91ce08",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "14"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# The normal transformation such as JW or BK requires 14 qubits for H2O's 14 orbitals\n",
    "\n",
    "a = jordan_wigner(fh)\n",
    "LinearQubitOperator(a).n_qubits"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9c6c84e",
   "metadata": {},
   "source": [
    "We can use binary code to save two further qubits, as the number of spin up and spin down filling is both 5 (5/odd electrons in 7 orbitals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "12ee73f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "12"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b = binary_code_transform(reorder(fh, up_then_down), 2 * checksum_code(7, 1))\n",
    "# 7 is 7 spin polarized orbitals, and 1 is for odd occupation\n",
    "LinearQubitOperator(b).n_qubits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3dd4b878",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12.412562749393349\n"
     ]
    }
   ],
   "source": [
    "print(b.terms[((0, \"Z\"),)])  # coefficient of Z_0 Pauli-string"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ecd89c0",
   "metadata": {},
   "source": [
    "* Transform the qubit Hamiltonian in openfermion to the format in TensorCircuit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "66fdc39d",
   "metadata": {},
   "outputs": [],
   "source": [
    "lsb, wb = tc.templates.chems.get_ps(b, 12)\n",
    "lsa, wa = tc.templates.chems.get_ps(a, 14)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8263f1a",
   "metadata": {},
   "source": [
    "* Inspect Hamiltonian in matrix form"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "88b116e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "ma = tc.quantum.PauliStringSum2COO_numpy(lsa, wa)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "63b92ae4",
   "metadata": {},
   "outputs": [],
   "source": [
    "mb = tc.quantum.PauliStringSum2COO_numpy(lsb, wb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4f76dfbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "mad, mbd = ma.todense(), mb.todense()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b43d5c1",
   "metadata": {},
   "source": [
    "The corresponding Hartree Fock product state in these two types of Hamiltonian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a907144e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0b11111111110000'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bin(np.argmin(np.diag(mad)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "a56f6191",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0b111110111110'"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bin(np.argmin(np.diag(mbd)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "994d944b",
   "metadata": {},
   "source": [
    "## VQE Setup\n",
    "\n",
    "We can in principle evaluate each Pauli string of the Hamiltonian as an expectation measurement, but it costs lots of simulation time, instead we fuse them as a Hamiltonian matrix as shown above to run the VQE."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b0f4d11",
   "metadata": {},
   "source": [
    "* Using dense matrix expectation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "6f1e065e",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 12\n",
    "depth = 4\n",
    "mbd_tf = tc.array_to_tensor(mbd)\n",
    "\n",
    "\n",
    "def vqe(param):\n",
    "    c = tc.Circuit(n)\n",
    "    for i in [0, 1, 2, 3, 4, 6, 7, 8, 9, 10]:\n",
    "        c.X(i)\n",
    "    for j in range(depth):\n",
    "        for i in range(n - 1):\n",
    "            c.exp1(i, i + 1, unitary=tc.gates._xx_matrix, theta=param[j, i, 0])\n",
    "        for i in range(n):\n",
    "            c.rx(i, theta=param[j, i, 1])\n",
    "        for i in range(n):\n",
    "            c.ry(i, theta=param[j, i, 2])\n",
    "        for i in range(n):\n",
    "            c.rx(i, theta=param[j, i, 3])\n",
    "    return tc.templates.measurements.operator_expectation(c, mbd_tf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "58fee777",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(-74.76671, shape=(), dtype=float32)\n",
      "tf.Tensor(-74.95493, shape=(), dtype=float32)\n",
      "tf.Tensor(-74.95319, shape=(), dtype=float32)\n",
      "tf.Tensor(-74.954315, shape=(), dtype=float32)\n",
      "tf.Tensor(-74.956116, shape=(), dtype=float32)\n",
      "tf.Tensor(-74.95809, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "vags = tc.backend.jit(tc.backend.value_and_grad(vqe))\n",
    "lr = tf.keras.optimizers.schedules.ExponentialDecay(\n",
    "    decay_rate=0.5, decay_steps=300, initial_learning_rate=0.5e-2\n",
    ")\n",
    "opt = tc.backend.optimizer(tf.keras.optimizers.Adam(lr))\n",
    "\n",
    "param = tc.backend.implicit_randn(shape=[depth, n, 4], stddev=0.02, dtype=\"float32\")\n",
    "for i in range(600):\n",
    "    e, g = vags(param)\n",
    "    param = opt.update(g, param)\n",
    "    if i % 100 == 0:\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94918b36",
   "metadata": {},
   "source": [
    "* Using sparse matrix expectation\n",
    "\n",
    "We can also use the sparse Hamiltonian matrix for circuit expectation evaluation, the only difference is to replace ``mbd_tf`` with ``mb_tf``"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "721f4494",
   "metadata": {},
   "outputs": [],
   "source": [
    "mb_tf = tc.backend.coo_sparse_matrix(\n",
    "    np.transpose(np.stack([mb.row, mb.col])), mb.data, shape=(2**n, 2**n)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8847fe3",
   "metadata": {},
   "source": [
    "A micro-benchmark between sparse matrix evaluation and dense matrix evaluation for expectation in terms of time, sparse always wins in terms of space, of course."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "8a60c6c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dense_expt(param):\n",
    "    c = tc.Circuit(n)\n",
    "    for i in range(n):\n",
    "        c.H(i)\n",
    "        c.rx(i, theta=param[i])\n",
    "    return tc.templates.measurements.operator_expectation(c, mbd_tf)\n",
    "\n",
    "\n",
    "def sparse_expt(param):\n",
    "    c = tc.Circuit(n)\n",
    "    for i in range(n):\n",
    "        c.H(i)\n",
    "        c.rx(i, theta=param[i])\n",
    "    return tc.templates.measurements.operator_expectation(c, mb_tf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "2764bbca",
   "metadata": {},
   "outputs": [],
   "source": [
    "dense_vag = tc.backend.jit(tc.backend.value_and_grad(dense_expt))\n",
    "sparse_vag = tc.backend.jit(tc.backend.value_and_grad(sparse_expt))\n",
    "\n",
    "v0, g0 = dense_vag(tc.backend.ones([n]))\n",
    "v1, g1 = sparse_vag(tc.backend.ones([n]))\n",
    "\n",
    "# consistency check\n",
    "\n",
    "np.testing.assert_allclose(v0, v1, atol=1e-5)\n",
    "np.testing.assert_allclose(g0, g1, atol=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "bc0265ee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "30.7 ms ± 1.45 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit dense_vag(tc.backend.ones([n]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "8913d16d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.6 ms ± 63 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit sparse_vag(tc.backend.ones([n]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3eafaa67",
   "metadata": {},
   "source": [
    "Therefore, sparse matrix evaluation also saves time apart from space, which is always recommended."
   ]
  }
 ],
 "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}