{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Subspace-Search Variational Quantum Eigensolver\n",
    "As explained in section 5-1, the variational quantum eigensolver (VQE) is an algorithm for searching the ground state of a quantum system. The ground state is an important state that determines the properties of various substances, but there are also phenomena such as photoresponse that are difficult to describe without the excited states.\n",
    "\n",
    "Therefore, algorithms for generating a wave function of an excited state on a quantum computer have been actively studied. Among them, here a new method proposed in 2018, Subspace-Search VQE (SSVQE) [1] is introduced.\n",
    "\n",
    "This notebook is translated from https://dojo.qulacs.org/ja/latest/notebooks/6.3_subspace_search_VQE.html\n",
    "\n",
    "## Algorithnm\n",
    "The procedures for SSVQE is as following:\n",
    "\n",
    "1. Prepare $k$ initial states $\\{\\left|\\varphi_i\\right>\\}_{i=0}^{k-1}$ orthogonal to each other.\n",
    "2. For each of them, operate an appropriate quantum circuit $U(\\theta)$ to generate a trial state. (Note: $\\left|\\psi_i(\\theta)\\right>=U(\\theta)\\left|\\varphi_i\\right>$)\n",
    "3. Adjust $\\theta$ to minimize cost function $L(\\theta)=\\sum_i w_i \\left<\\psi_i(\\theta)\\left|H\\right|\\psi_i(\\theta)\\right>$. Here when $w_i$ is positive and $i>j$, make sure $w_i<w_j$. \n",
    "\n",
    "When $U(\\theta)$ has sufficient expressiveness, $\\left|\\psi_i(\\theta)\\right>$ becomes the $i$th excited state at converged $\\theta$. This is because when the $i$th excited state of the Hamiltonian is $\\left|E_i\\right>$, the global minimum of the cost function is $\\left|\\psi_i\\right>=\\left|E_i\\right>$. (Refer to the original paper [1] for details)\n",
    "\n",
    "The following shows an example of simulating SSVQE using Qulacs.\n",
    "\n",
    "## Implementation of SSVQE\n",
    "Implement SSVQE to search ground state and first excited state of hydrogen molecule. When hydrogen molecule is treated with sto-3g minimal basis set, 4 qubit Hamiltonian can be obtained. Therefore, two initial states, $\\left|0000\\right>$ and $\\left|0001\\right>$, which are flipped by one qubit, are used as orthogonal initial states required for SSVQE.\n",
    "\n",
    "* Note for those who are familiar with quantum chemistry: The first excited state here is the first excited state of the Hamiltonian of a hydrogen molecule without considering the number of electrons, and in chemistry terms, the ground state of a hydrogen molecular ion.\n",
    "\n",
    "### Create Hamiltonian\n",
    "Follow the same procedure as explained before to calculate Hamiltonian using PySCF + OpenFermion. However, the excited state is considered in SSVQE, so `scipy.sparse.linalg.eigsh` is used to obtain the exact solution of excited state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Please execute if various libraries are not installed\n",
    "## When use Google Colaboratory, please ignore 'You must restart the runtime in order to use newly installed versions'.\n",
    "## Crash when restarting runtime.\n",
    "!pip install qulacs pyscf openfermion openfermionpyscf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import qulacs\n",
    "from openfermion.transforms import get_fermion_operator, jordan_wigner\n",
    "from openfermion.transforms import get_sparse_operator\n",
    "from openfermion.hamiltonians import MolecularData\n",
    "from openfermionpyscf import run_pyscf\n",
    "from scipy.optimize import minimize\n",
    "from pyscf import fci\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "basis = \"sto-3g\"\n",
    "multiplicity = 1\n",
    "charge = 0\n",
    "distance  = 0.977\n",
    "geometry = [[\"H\", [0,0,0]],[\"H\", [0,0,distance]]]\n",
    "description  = \"tmp\"\n",
    "molecule = MolecularData(geometry, basis, multiplicity, charge, description)\n",
    "molecule = run_pyscf(molecule,run_scf=1,run_fci=1)\n",
    "n_qubit = molecule.n_qubits\n",
    "n_electron = molecule.n_electrons\n",
    "fermionic_hamiltonian = get_fermion_operator(molecule.get_molecular_hamiltonian())\n",
    "jw_hamiltonian = jordan_wigner(fermionic_hamiltonian)\n",
    "hamiltonian_matrix = get_sparse_operator(jw_hamiltonian)\n",
    "from scipy.sparse.linalg import eigsh\n",
    "eigval, eigvec = eigsh(hamiltonian_matrix, k=2, which=\"SA\")\n",
    "from qulacs import Observable\n",
    "from qulacs.observable import create_observable_from_openfermion_text\n",
    "qulacs_hamiltonian = create_observable_from_openfermion_text(str(jw_hamiltonian))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct ansatz\n",
    "Create quantum circuit on Qulacs. Here, the quantum circuit is modeled after the experiments with superconducting qubits (A. Kandala et. al. , “Hardware-efficient variational quantum eigensolver for small molecules and quantum magnets“, Nature **549**, 242–246)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qulacs import QuantumState, QuantumCircuit\n",
    "from qulacs.gate import CZ, RY, RZ, merge\n",
    "\n",
    "depth = n_qubit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def he_ansatz_circuit(n_qubit, depth, theta_list):\n",
    "    \"\"\"he_ansatz_circuit\n",
    "    Returns hardware efficient ansatz circuit.\n",
    "\n",
    "    Args:\n",
    "        n_qubit (:class:`int`):\n",
    "            the number of qubit used (equivalent to the number of fermionic modes)\n",
    "        depth (:class:`int`):\n",
    "            depth of the circuit.\n",
    "        theta_list (:class:`numpy.ndarray`):\n",
    "            rotation angles.\n",
    "    Returns:\n",
    "        :class:`qulacs.QuantumCircuit`\n",
    "    \"\"\"\n",
    "    circuit = QuantumCircuit(n_qubit)\n",
    "    circuit.add_gate(RY(0, theta_list[-2]))\n",
    "    circuit.add_gate(RZ(0, theta_list[-1]))\n",
    "    for d in range(depth):\n",
    "        for i in range(n_qubit):\n",
    "            circuit.add_gate(merge(RY(i, theta_list[2*i+2*n_qubit*d]), RZ(i, theta_list[2*i+1+2*n_qubit*d])))\n",
    "        for i in range(n_qubit//2):\n",
    "            circuit.add_gate(CZ(2*i, 2*i+1))\n",
    "        for i in range(n_qubit//2-1):\n",
    "            circuit.add_gate(CZ(2*i+1, 2*i+2))\n",
    "    for i in range(n_qubit):\n",
    "        circuit.add_gate(merge(RY(i, theta_list[2*i+2*n_qubit*depth]), RZ(i, theta_list[2*i+1+2*n_qubit*depth])))\n",
    "\n",
    "    return circuit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define cost function of SSVQE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.21631044919594045"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_exp(state, theta_list):\n",
    "    circuit = he_ansatz_circuit(n_qubit, depth, theta_list) #Create quantum circuit\n",
    "    circuit.update_quantum_state(state) #Operate on quantum state\n",
    "    return qulacs_hamiltonian.get_expectation_value(state)\n",
    "\n",
    "def cost(theta_list):\n",
    "    state0 = QuantumState(n_qubit) #Prepare |00000>\n",
    "    state1 = QuantumState(n_qubit); state1.set_computational_basis(1) #Prepare |00001>\n",
    "    return get_exp(state0, theta_list)+0.5*get_exp(state1, theta_list)\n",
    "\n",
    "init_theta_list = np.random.random(2*n_qubit*(depth+1)+2)*1e-1\n",
    "cost(init_theta_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run SSVQE\n",
    "Now everthing is prepared, run SSVQE. For optimization, the BFGS method implemented in scipy is applied, and initial parameters are randomly selected. This process should end in tens of seconds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_history0 = []\n",
    "exp_history1 = []\n",
    "def callback(theta_list):\n",
    "    state0 = QuantumState(n_qubit) #Prepare |0000>\n",
    "    state1 = QuantumState(n_qubit); state1.set_computational_basis(1) #Prepare |0001>\n",
    "    exp_history0.append(get_exp(state0, theta_list))\n",
    "    exp_history1.append(get_exp(state1, theta_list))\n",
    "init_theta_list = np.random.random(2*n_qubit*(depth+1)+2)*1e-1\n",
    "method = \"BFGS\"\n",
    "options = {\"disp\": True, \"maxiter\": 50, \"gtol\": 1e-6}\n",
    "opt = minimize(cost, init_theta_list,\n",
    "               method=method,\n",
    "               callback=callback)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The results can be plotted to see that they converge to the correct solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 44448 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.rcParams[\"font.size\"] = 18\n",
    "plt.plot(exp_history0, label=r\"input $|0000\\rangle$\")\n",
    "plt.plot(exp_history1, label=r\"input $|0001\\rangle$\")\n",
    "plt.plot(range(len(exp_history0)), [molecule.fci_energy]*len(exp_history0), linestyle=\"dashed\", color=\"black\", label=\"Exact ground state energy\")\n",
    "plt.plot(range(len(exp_history1)), [eigval[1]]*len(exp_history1), linestyle=\"-.\", color=\"black\", label=\"Exact 1st excited state energy\")\n",
    "plt.xlabel(\"Iteration\")\n",
    "plt.ylabel(\"Energy expectation value\")\n",
    "plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0, fontsize=18)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reference\n",
    "[1] K. M Nakanishi, K. Mitarai, and K. Fujii, “Subspace-search variational quantum eigensolver for excited states”, Phys. Rev. Research 1, 033062 (2019), https://arxiv.org/abs/1810.09434\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}