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": {},
   "source": [
    "## Overview\n",
    "\n",
    "Some quantum chemistry backgrounds.\n",
    "\n",
    "[WIP]"
   ]
  },
  {
   "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 obitals)"
   ]
  },
  {
   "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 Paulistring"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ecd89c0",
   "metadata": {},
   "source": [
    "* Transform the qubit Hamiltoian 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 corrsponding Hartree Fock product state in these two types of Hamiltonians"
   ]
  },
  {
   "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 princinple evaluate each Pauli string of the Hamiltonian as an expectation measurment, but it costs lots of simulation time, instead we fuse them together 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 win 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 recommeded."
   ]
  }
 ],
 "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
}