{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quantum Enhanced Auxiliary Field Quantum Monte Carlo\n",
    "\n",
    "This work was done in collaboration with the Next Generation Computing team at [BASF](https://www.basf.com/global/en.html).\n",
    "\n",
    "In this tutorial we implement a quantum-classical hybrid workflow for computing the ground state energies of a strongly interacting molecular system. The algorithm consists of two parts:\n",
    "\n",
    "\n",
    "1. A variational quantum eigensolver that uses the quantum-number-preserving ansatz proposed by [Anselmetti et al. (2021)](https://doi.org/10.1088/1367-2630/ac2cb3) to generate a quantum trial wave function $|\\Psi_T\\rangle$ using CUDA Quantum.\n",
    "\n",
    "2. An Auxiliary-Field Quantum Monte Carlo simulation that realizes a classical imaginary time evolution and collects the ground state energy estimates.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "# Package installs\n",
    "!pip install pyscf>=2.7 openfermion==1.6.1 ipie==0.7.1 numba>=0.61 -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Relevant imports\n",
    "\n",
    "import cudaq\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from pyscf import gto, scf, ao2mo, mcscf\n",
    "\n",
    "from afqmc_src.vqe_cudaq_qnp import VQE, get_cudaq_hamiltonian\n",
    "from afqmc_src.utils_ipie import get_coeff_wf, gen_ipie_input_from_pyscf_chk\n",
    "\n",
    "from ipie.hamiltonians.generic import Generic as HamGeneric\n",
    "from ipie.qmc.afqmc import AFQMC\n",
    "from ipie.systems.generic import Generic\n",
    "from ipie.trial_wavefunction.particle_hole import ParticleHole\n",
    "from ipie.analysis.extraction import extract_observable\n",
    "\n",
    "from ipie.config import config\n",
    "\n",
    "if cudaq.num_available_gpus() > 0 and cudaq.has_target(\"nvidia\"):\n",
    "    cudaq.set_target(\"nvidia\")\n",
    "else:\n",
    "    print(\"CUDA or GPU support is unavailable. Running with CPU simulator. Performance may be significantly reduced.\")\n",
    "    cudaq.set_target(\"qpp-cpu\")\n",
    "\n",
    "# Ipie has recently added GPU support however this remains a bit tricky to use as it requires manual installation of several packages.\n",
    "# Once this is streamlined, we can set the GPU option to True in the tutorial.\n",
    "config.update_option(\"use_gpu\", False)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start by defining the structure of the molecule, the basis set, and its spin. We build the molecule object with PySCF and run a preliminary Hartree-Fock computation. Here we choose as an example a [chelating agent](https://doi.org/10.1021/acs.jctc.3c01375) representing a relevant class of substances industrially produced at large scales. Their use ranges, among the others, from water softeners in cleaning applications, modulators of redox behaviour in oxidative bleaching, scale suppressants, soil remediation and ligands for catalysts. In particular we focus here in a Fe(III)-NTA complex whose structure is given in the file imported below.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the molecular structure and the basis set for the Fenta molecule.\n",
    "\n",
    "atom = \"afqmc_src/geo_fenta.xyz\"\n",
    "basis = \"cc-pVTZ\"\n",
    "spin = 1\n",
    "num_active_orbitals = 5\n",
    "num_active_electrons = 5\n",
    "\n",
    "# You can swap to O3 which is a smaller system and takes less computational resources and time to run.\n",
    "\n",
    "atom = \"afqmc_src/geo_o3.xyz\"\n",
    "basis = \"cc-pVTZ\"\n",
    "spin = 0\n",
    "num_active_orbitals = 9\n",
    "num_active_electrons = 12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-224.34048064812268"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# PYSCF helps to build the molecule and run Hartree-Fock.\n",
    "\n",
    "# Define the molecule.\n",
    "molecule = gto.M(atom=atom, spin=spin, basis=basis, verbose=0)\n",
    "\n",
    "# Restriced open shell HF.\n",
    "hartee_fock = scf.ROHF(molecule)\n",
    "hartee_fock.chkfile = \"afqmc_src/output.chk\"\n",
    "\n",
    "# Run Hartree-Fock.\n",
    "hartee_fock.kernel()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hamiltonian preparation for VQE\n",
    "\n",
    "Since this molecule contains of around 600 orbitals (which would correspond to 1200 qubits) and 143 total electrons, it is impossible to perform a full VQE with full statevector simulation. Therefore, we need to identify an active space with fewer orbitals and electrons that contribute to the strongly interacting part of the whole molecule. We then run a post Hartree-Fock computation with the PySCF's built-in CASCI method in order to obtain the one-body ($t_{pq}$) and two-body \n",
    "($V_{prqs}$) integrals that define the molecular Hamiltonian in the active space:\n",
    "\n",
    "$$ H= \\sum_{pq}t_{pq}\\hat{a}_{p}^\\dagger \\hat {a}_{q}+\\sum_{pqrs}  V_{prqs}\\hat a_{p}^\\dagger \\hat a_{q}^\\dagger \\hat a_{s}\\hat a_{r} \\tag{1}$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openfermion.transforms import jordan_wigner\n",
    "from openfermion import generate_hamiltonian\n",
    "\n",
    "# Run a CASCI simulation for computing the Hamiltonian in the active space.\n",
    "casci = mcscf.CASCI(hartee_fock, num_active_orbitals, num_active_electrons)\n",
    "casci.fix_spin_(ss=(molecule.spin / 2 * (molecule.spin / 2 + 1)))\n",
    "\n",
    "# Executes the kernel to compute the hamiltonian in the active space.\n",
    "casci.kernel()\n",
    "\n",
    "# Compute the one-body (h1) and two-body integrals (tbi) as shown in equation 1.\n",
    "h1, energy_core = casci.get_h1eff()\n",
    "\n",
    "h2 = casci.get_h2eff()\n",
    "h2_no_symmetry = ao2mo.restore('1', h2, num_active_orbitals)\n",
    "\n",
    "# V_pqrs terms in H.\n",
    "tbi = np.asarray(h2_no_symmetry.transpose(0, 2, 3, 1), order='C')\n",
    "\n",
    "# Compute the hamiltonian and convert it to a CUDA-Q operator.\n",
    "mol_ham = generate_hamiltonian(h1, tbi, energy_core.item())\n",
    "jw_hamiltonian = jordan_wigner(mol_ham)\n",
    "hamiltonian, constant_term = get_cudaq_hamiltonian(jw_hamiltonian)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run VQE with CUDA-Q\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now execute the VQE algorithm using the quantum number preserving ansatz. At the end of the VQE, we store the final statevector that will be used in the classical AFQMC computation as an initial guess.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Using cudaq optimizer\n",
      "# Num Params: 16\n",
      "# Qubits: 18\n",
      "# N_layers: 1\n",
      "# Energy after the VQE: -224.39453699035025\n"
     ]
    }
   ],
   "source": [
    "# Define some options for the VQE.\n",
    "options = {\n",
    "    'n_vqe_layers': 1,\n",
    "    'maxiter': 100,\n",
    "    'energy_core': constant_term,\n",
    "    'return_final_state_vec': True\n",
    "}\n",
    "\n",
    "n_qubits = 2 * num_active_orbitals\n",
    "\n",
    "vqe = VQE(n_qubits=n_qubits,\n",
    "          num_active_electrons=num_active_electrons,\n",
    "          spin=spin,\n",
    "          options=options)\n",
    "\n",
    "results = vqe.execute(hamiltonian)\n",
    "\n",
    "# Best energy from VQE.\n",
    "optimized_energy = results['energy_optimized']\n",
    "\n",
    "# Final state vector.\n",
    "final_state_vector = results[\"state_vec\"]\n",
    "\n",
    "# Energies during the optimization loop.\n",
    "vqe_energies = results[\"callback_energies\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Auxiliary Field Quantum Monte Carlo (AFQMC)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "AFQMC is a numerical method for computing relevant properties of strongly interacting molecules. AFQMC is a type of Quantum Monte Carlo method that combines the use of random walks with an auxiliary field to simulate the imaginary-time evolution of a quantum system and drive it to the lowest energy state. This method can provide accurate results for ground-state properties of a wide range of physical systems, including atoms, molecules, and solids. Here we summarize the main features of AFQMC while a detailed introduction to can be found [here](https://www.cond-mat.de/events/correl13/manuscripts/zhang.pdf).\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We consider the electronic Hamiltonian in the second quantization\n",
    "\\begin{equation}\n",
    "H = {H}_1 + {H}_2 \n",
    "=\\sum_{pq} h_{pq} {a}_{p}^{\\dagger} {a}_{q} + \\frac{1}{2} \\sum_{pqrs} v_{pqrs}{a}_{p}^{\\dagger} {a}_r {a}^{\\dagger}_{q} {a}_s \\tag{2}\n",
    "\\end{equation}\n",
    "where ${a}_{p}^{\\dagger}$ and ${a}_{q}$ are fermionic creation and annihilation operators of orbitals $p$ and $q$, respectively. The terms $h_{pq} $ and \n",
    "$v_{pqrs}$ are the matrix elements of the one-body, $H_1$, and two-body, $H_2$, interactions of $H$, respectively. Here, we omit the spin indices for simplicity.\n",
    "\n",
    "AFQMC realizes an imaginary time propagation of an initial state (chosen as a Slater determinant) $\\ket{\\Psi_{I}}$ towards the ground state $\\ket{\\Psi_0}$ of a given hamiltonian, $H$, with\n",
    "\\begin{equation}\n",
    "\\ket{\\Psi_0} \\sim\\lim_{n \\to \\infty} \\left[ e^{-\\Delta\\tau H  }  \\right]^{n} \\ket{\\Psi_{I}}\n",
    "\\tag{3}\n",
    "\\end{equation} \n",
    "where $\\Delta\\tau$ is the imaginary time step.\n",
    "\n",
    "AFQMC relies on decomposing the two-body interactions $H_2$ in terms of sum of squares of one-body operators ${v}_\\gamma$ such that the Hamiltonian ${H}$ becomes\n",
    "\\begin{equation}\n",
    "H = v_0 - \\frac{1}{2}\\sum_{\\gamma=1}^{N_\\gamma} {v}_\\gamma^2\n",
    "\\tag{4}\n",
    "\\end{equation}\n",
    "with ${v}_0 = {H}_1 $ and $\n",
    "{v}_\\gamma = i \\sum_{pq} L^{\\gamma}_{pq} {a}_{p}^{\\dagger}{a}_{q}.\n",
    "$\n",
    "The $N_\\gamma$ matrices $L^{\\gamma}_{pq}$ are called Cholesky vectors as they are  obtained via a Cholesky decomposition of the two-body matrix elements \n",
    "$v_{pqrs}$ via $v_{pqrs} = \\sum_{\\gamma=1}^{N_\\gamma} L^{\\gamma}_{pr} L^{\\gamma}_{qs}$.\n",
    "\n",
    "The imaginary time propagation evolves an ensemble of walkers $\\{\\phi^{(n)}\\}$ (that are Slater determinants) and allows one to access observables of the system. For example, the local energy\n",
    "\\begin{equation}\n",
    "\\mathcal{E}_{\\text{loc}}(\\phi^{(n)}) = \\frac{\\bra{\\Psi_\\mathrm{T}}H\\ket{\\phi^{(n)}}}{\\braket{\\Psi_\\mathrm{T}| \\phi^{(n)}}}\n",
    "\\tag{5}\n",
    "\\end{equation}\n",
    "defined as the mixed expectation value of the Hamiltonian with the trial wave function $\\ket{\\Psi_\\mathrm{T}}$.\n",
    "\n",
    "\n",
    "The trial wavefunction can be in general a single or a multi-Slater determinant coming from VQE for example. This might help in achieving more accurate ground state energy estimates.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "The implementation of AFQMC we use here is from [ipie](https://github.com/JoonhoLee-Group/ipie) that supports both CPUs and GPUs and requires the following steps:\n",
    "\n",
    "\n",
    "1. Preparation of the molecular Hamiltonian by performing the Cholesky decomposition\n",
    "\n",
    "2. Preparation of the trial state from the VQE wavefunction\n",
    "\n",
    "3. Executing AFQMC\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preparation of the molecular Hamiltonian\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Number of electrons in simulation: (12, 12)\n"
     ]
    }
   ],
   "source": [
    "# AFQMC.\n",
    "\n",
    "# Generate the input Hamiltonian for ipie from the checkpoint file from pyscf.\n",
    "ipie_hamiltonian = gen_ipie_input_from_pyscf_chk(hartee_fock.chkfile,\n",
    "                                                 mcscf=True,\n",
    "                                                 chol_cut=1e-5)\n",
    "\n",
    "h1e, cholesky_vectors, e0 = ipie_hamiltonian\n",
    "\n",
    "num_basis = cholesky_vectors.shape[1]\n",
    "num_chol = cholesky_vectors.shape[0]\n",
    "\n",
    "system = Generic(nelec=molecule.nelec)\n",
    "\n",
    "afqmc_hamiltonian = HamGeneric(\n",
    "    np.array([h1e, h1e]),\n",
    "    cholesky_vectors.transpose((1, 2, 0)).reshape(\n",
    "        (num_basis * num_basis, num_chol)), e0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preparation of the trial wave function\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build the trial wavefunction from the state vector computed via VQE.\n",
    "wavefunction = get_coeff_wf(final_state_vector,\n",
    "                            n_active_elec=num_active_electrons,\n",
    "                            spin=spin)\n",
    "\n",
    "trial = ParticleHole(wavefunction,\n",
    "                     molecule.nelec,\n",
    "                     num_basis,\n",
    "                     num_dets_for_props=len(wavefunction[0]),\n",
    "                     verbose=False)\n",
    "\n",
    "trial.compute_trial_energy = True\n",
    "trial.build()\n",
    "trial.half_rotate(afqmc_hamiltonian)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup of the AFQMC parameters\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we can choose the input options like the timestep $\\Delta\\tau$, the total number of walkers `num_walkers` and the total number of AFQMC iterations `num_blocks`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# random seed is 96264512\n",
      "            Block                   Weight            WeightFactor            HybridEnergy                  ENumer                  EDenom                  ETotal                  E1Body                  E2Body\n",
      "                0   1.0000000000000000e+02  1.0000000000000000e+02  0.0000000000000000e+00 -2.2439508600587662e+04  1.0000000000000000e+02 -2.2439508600587664e+02 -3.7639365207102333e+02  1.5199856606514666e+02\n",
      "                1   4.2257056483540680e+02  1.4118069021027802e+03 -1.1708693041911519e+02 -2.2474996636792086e+04  1.0000000000000000e+02 -2.2474996636792088e+02 -3.7647319974320789e+02  1.5172323337528707e+02\n",
      "                2   1.0031477304650311e+02  3.8296193043358181e+02 -1.1739718565992376e+02 -2.2490598976544803e+04  1.0000000000000000e+02 -2.2490598976544803e+02 -3.7651529453780955e+02  1.5160930477236158e+02\n",
      "                3   9.9899069244275523e+01  1.0007683865659165e+02 -1.1733421170204294e+02 -2.2497093535653130e+04  9.9999999999999986e+01 -2.2497093535653136e+02 -3.7660556251095613e+02  1.5163462715442475e+02\n",
      "                4   1.0009053912662398e+02  1.0004689248382901e+02 -1.1745467967836868e+02 -2.2496851740853570e+04  1.0000000000000001e+02 -2.2496851740853566e+02 -3.7677033675853681e+02  1.5180181935000121e+02\n",
      "                5   9.9997252340375638e+01  1.0010501097562071e+02 -1.1749279331833796e+02 -2.2503958046591404e+04  1.0000000000000000e+02 -2.2503958046591404e+02 -3.7664566538980034e+02  1.5160608492388630e+02\n",
      "                6   1.0011587922408398e+02  1.0018079748584451e+02 -1.1766351406966876e+02 -2.2514170340022600e+04  1.0000000000000001e+02 -2.2514170340022596e+02 -3.7662567666476224e+02  1.5148397326453633e+02\n",
      "                7   9.9931815987329159e+01  9.9916747721404647e+01 -1.1761065926981941e+02 -2.2517005364646804e+04  1.0000000000000000e+02 -2.2517005364646803e+02 -3.7661792399494180e+02  1.5144787034847369e+02\n",
      "                8   9.9903303729090240e+01  9.9912226290917729e+01 -1.1757442061417670e+02 -2.2518386757775326e+04  9.9999999999999986e+01 -2.2518386757775332e+02 -3.7676538247523837e+02  1.5158151489748514e+02\n",
      "                9   1.0011958041774147e+02  1.0012012496437015e+02 -1.1776042000332971e+02 -2.2513320955662537e+04  9.9999999999999986e+01 -2.2513320955662542e+02 -3.7678749021253810e+02  1.5165428065591271e+02\n",
      "               10   9.9637479319398665e+01  9.9239880452225833e+01 -1.1745892832071745e+02 -2.2518486392115545e+04  1.0000000000000000e+02 -2.2518486392115545e+02 -3.7689296595539395e+02  1.5170810203423866e+02\n"
     ]
    }
   ],
   "source": [
    "# Setup the AFQMC parameters.\n",
    "afqmc_msd = AFQMC.build(molecule.nelec,\n",
    "                        afqmc_hamiltonian,\n",
    "                        trial,\n",
    "                        num_walkers=100,\n",
    "                        num_steps_per_block=25,\n",
    "                        num_blocks=10,\n",
    "                        timestep=0.005,\n",
    "                        stabilize_freq=5,\n",
    "                        seed=96264512,\n",
    "                        pop_control_freq=5,\n",
    "                        verbose=False)\n",
    "\n",
    "# Run the AFQMC.\n",
    "afqmc_msd.run(estimator_filename='afqmc_src/estimates.0.h5')\n",
    "afqmc_msd.finalise(verbose=False)\n",
    "\n",
    "# Extract the energies.\n",
    "qmc_data = extract_observable(afqmc_msd.estimators.filename, \"energy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x79af5871de20>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 42836 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the energies.\n",
    "\n",
    "vqe_y = vqe_energies\n",
    "vqe_x = list(range(len(vqe_y)))\n",
    "plt.plot(vqe_x, vqe_y, label=\"VQE\")\n",
    "\n",
    "afqmc_y = list(qmc_data[\"ETotal\"])\n",
    "afqmc_x = [i + vqe_x[-1] for i in list(range(len(afqmc_y)))]\n",
    "plt.plot(afqmc_x, afqmc_y, label=\"AFQMC\")\n",
    "\n",
    "plt.xlabel(\"Optimization steps\")\n",
    "plt.ylabel(\"Energy [Ha]\")\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CUDA-Q Version proto-0.8.0 (https://github.com/NVIDIA/cuda-quantum 03d503c8eab430b4e2cd6882778467e1c3a5b27b)\n"
     ]
    }
   ],
   "source": [
    "print(cudaq.__version__)"
   ]
  }
 ],
 "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
