{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c9f44aeb-4585-468b-bf8c-57f0c23fd070",
   "metadata": {},
   "source": [
    "# Example 2: PennyLane VQE calculation on water\n",
    "\n",
    "Unlike the previous case, this example showcases the use of `qc2` in combination with `PennyLane` to perform a oo-VQE run on water molecule. All circuit evaluations are conducted using the `default.qubit` state simulator device, which provides exact expectation values."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06892d44-96c8-4ebd-8f17-2bcc8cf81935",
   "metadata": {},
   "source": [
    "### Import required packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "07462ed4-add5-47b2-8a3a-e3b5dcfa153a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess\n",
    "\n",
    "# ASE molecule object\n",
    "from ase.build import molecule\n",
    "\n",
    "# PennyLane-related packages\n",
    "import pennylane as qml\n",
    "from pennylane import numpy as np\n",
    "\n",
    "# ignore package import warnings\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "# qc2 packages\n",
    "from qc2.data import qc2Data\n",
    "from qc2.ase import Psi4\n",
    "from qc2.algorithms.utils import ActiveSpace\n",
    "from qc2.algorithms.pennylane import oo_VQE"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e461edc-f8a7-48c8-a855-6f86d0af742d",
   "metadata": {},
   "source": [
    "### Instantiate `qc2Data` class and run qc2-ASE calculator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "830abd9b-0315-4e8f-b1ee-1cc435a6616e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Python driver attempted to set threads to 1.\n",
      "  Psi4 was compiled without OpenMP, setting threads to 1.\n",
      "  Python driver attempted to set threads to 1.\n",
      "  Psi4 was compiled without OpenMP, setting threads to 1.\n",
      "* Reference energy (Hartree): -74.96449224627995\n",
      "* Saving qchem data in h2o.hdf5\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# set Atoms object\n",
    "mol = molecule('H2O')\n",
    "\n",
    "# instantiate qc2Data class\n",
    "qc2data = qc2Data(\n",
    "    molecule=mol,\n",
    "    filename='h2o.hdf5',\n",
    "    schema='qcschema'\n",
    ")\n",
    "\n",
    "# specify the qchem calculator\n",
    "qc2data.molecule.calc = Psi4(method='hf', basis='sto-3g')\n",
    "\n",
    "# run calculation and save qchem data in the fcidump file\n",
    "qc2data.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b8e4a40-3e7c-4b53-8490-a669ea40ec69",
   "metadata": {},
   "source": [
    "This will generate a `h2o.hdf5` file containing all relevant qchem data according to the `QCSchema`. Like before, let's check its data structure via the shell command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2dc82e2c-ad46-42ee-af96-567584fbb21c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HDF5 \"h2o.hdf5\" {\n",
      "FILE_CONTENTS {\n",
      " group      /\n",
      " group      /keywords\n",
      " group      /model\n",
      " group      /molecule\n",
      " group      /properties\n",
      " group      /provenance\n",
      " group      /wavefunction\n",
      " dataset    /wavefunction/localized_orbitals_a\n",
      " dataset    /wavefunction/localized_orbitals_b\n",
      " dataset    /wavefunction/scf_eigenvalues_a\n",
      " dataset    /wavefunction/scf_eigenvalues_b\n",
      " dataset    /wavefunction/scf_eri\n",
      " dataset    /wavefunction/scf_eri_mo_aa\n",
      " dataset    /wavefunction/scf_eri_mo_ab\n",
      " dataset    /wavefunction/scf_eri_mo_ba\n",
      " dataset    /wavefunction/scf_eri_mo_bb\n",
      " dataset    /wavefunction/scf_fock_a\n",
      " dataset    /wavefunction/scf_fock_mo_a\n",
      " dataset    /wavefunction/scf_fock_mo_b\n",
      " dataset    /wavefunction/scf_orbitals_a\n",
      " dataset    /wavefunction/scf_orbitals_b\n",
      " }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "!h5dump -n h2o.hdf5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a546a43-5307-4c56-bdcc-0140b61c32d1",
   "metadata": {},
   "source": [
    "You could also use `!h5dump h2o.hdf5` to inspect the numerical data contained within each `group` and `dataset`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f687841-e96a-41c1-b4c8-25e0184e8fdb",
   "metadata": {},
   "source": [
    "### Instantiate `qc2.algorithms.oo_VQE` class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6a416526-7bd5-448d-b6ec-be67c349dfea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up activate space\n",
    "active_space = ActiveSpace(\n",
    "    num_active_electrons=(2, 2),  # => (n_alpha, n_beta)\n",
    "    num_active_spatial_orbitals=4 # => active orbitals\n",
    ")\n",
    "\n",
    "# instantiate oo_VQE class\n",
    "qc2data.algorithm = oo_VQE(\n",
    "    active_space=active_space,\n",
    "    mapper=\"jw\",                                           # => use Jordan-Wigner mapper\n",
    "    optimizer=qml.GradientDescentOptimizer(stepsize=0.5),  # => GradientDescentOptimizer optimizer from PennyLane\n",
    "    device=\"default.qubit\",\n",
    "    max_iterations=100\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06792fc5-396b-41c2-8b1c-a0bbb2ded6b0",
   "metadata": {},
   "source": [
    "### Run oo-VQE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "619e3c98-85e6-4a19-abb0-3ebbdaf7c49a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ">>> Optimizing circuit and orbital parameters...\n",
      "iter = 000, energy = -74.964404821247 Ha\n",
      "iter = 001, energy = -74.974654732364 Ha\n",
      "iter = 002, energy = -74.977136100996 Ha\n",
      "iter = 003, energy = -74.977654273872 Ha\n",
      "iter = 004, energy = -74.977797944648 Ha\n",
      "iter = 005, energy = -74.977844599529 Ha\n",
      "iter = 006, energy = -74.977860462172 Ha\n",
      "iter = 007, energy = -74.977866022280 Ha\n",
      "iter = 008, energy = -74.977867940065 Ha\n",
      "iter = 009, energy = -74.977868636996 Ha\n",
      "iter = 010, energy = -74.977868887832 Ha\n",
      "iter = 011, energy = -74.977868978674 Ha\n",
      "optimization finished.\n",
      "\n",
      "=== PENNYLANE oo-VQE RESULTS ===\n",
      "* Total ground state energy (Hartree): -74.977868978674\n"
     ]
    }
   ],
   "source": [
    "# add extra options to `device` and `QNode` if needed...\n",
    "results = qc2data.algorithm.run(\n",
    "    device_kwargs={\"shots\": None},\n",
    "    qnode_kwargs={\"diff_method\": \"best\"}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "dca08041-486a-44cc-911d-c2f233fb1005",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "* Optimized circuit parameters:\n",
      "[ 2.07765036e-02 -1.11833627e-16 -2.07765036e-02 -1.57494321e-17\n",
      " -4.82754622e-17 -8.98925217e-17 -6.71880313e-17 -3.46826285e-17\n",
      "  1.56019786e-01 -5.54938723e-17 -2.59957419e-17  4.35653437e-02\n",
      "  1.45114627e-18 -9.95617852e-18  1.18705568e-17  1.02981524e-17\n",
      "  1.21106275e-17  6.38136431e-18  9.59524869e-18  9.84554818e-18\n",
      "  3.12784437e-19  3.55066164e-18  5.68534374e-02  3.11924957e-18\n",
      " -1.73185321e-19  3.13803385e-02] \n",
      "\n",
      "* Optimized orbital parameters:\n",
      "[0.0026712593299834616, -0.6190212888834982, 6.275542869417149e-10, -1.113516977506888e-16, 3.8213383832866345e-15, 5.586006655411282e-14, 7.555308746399609e-15, -2.7210673457436863e-05, -0.0019899144798891403, 9.447639902049476e-12, 0.0019118957526852315, 2.3253980286437208e-15, 2.0963033505657805e-15, 1.8613420600403597e-12, -0.002273727905803901, 2.9862067670493623e-12, -6.067518761511912e-16, 5.340727372484386e-10] \n",
      "\n",
      "* oo-VQE energy (Hartree):\n",
      "-74.97786897867395 \n",
      "\n"
     ]
    }
   ],
   "source": [
    "# print optimized circuit parameters\n",
    "print(f'* Optimized circuit parameters:')\n",
    "print(results.optimal_circuit_params, \"\\n\")\n",
    "\n",
    "# orbital parameters\n",
    "print(f'* Optimized orbital parameters:')\n",
    "print(results.optimal_orbital_params, \"\\n\")\n",
    "\n",
    "print(f'* oo-VQE energy (Hartree):')\n",
    "print(results.optimal_energy, \"\\n\")\n",
    "\n",
    "# vectors containing all intermediate results\n",
    "# results.optimizer_evals, results.energy, results.circuit_parameters, results.orbital_parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c008445-2fde-4e78-960c-8465a8846835",
   "metadata": {},
   "source": [
    "### Compare oo-VQE result with classical qchem calculations\n",
    "\n",
    "Once again, let's compare our oo-VQE energy with the one obtained from Psi4 calculations alone:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "af2a9819-ad51-4c38-aaed-ec77da1fa8e0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "* Final oo-VQE energy:        -74.977869 hartree\n",
      "* Final CASSCF/sto-3g energy: -74.977870 hartree\n",
      "* Final FCI/sto-3g energy:    -75.015429 hartree\n"
     ]
    }
   ],
   "source": [
    "import psi4\n",
    "\n",
    "# Set Psi4 to run in serial mode\n",
    "psi4.set_num_threads(1)\n",
    "\n",
    "# Define the molecule in XYZ format\n",
    "h2o = psi4.geometry(\"\"\"\n",
    "O 0.0000  0.000000000  0.1192622641\n",
    "H 0.0000  0.763237638 -0.4770469398\n",
    "H 0.0000 -0.763237638 -0.4770469398\n",
    "symmetry c1\n",
    "noreorient\n",
    "nocom\n",
    "\"\"\")\n",
    "\n",
    "# Setup CAS space (Active space definition)\n",
    "n_active_elec = 4    # Number of active electrons\n",
    "n_active_orb = 4     # Number of active orbitals\n",
    "\n",
    "# Set computation options\n",
    "psi4.set_options({\n",
    "    'basis': 'sto-3g',\n",
    "    'reference': 'rhf'\n",
    "})\n",
    "\n",
    "# run reference Hartree-Fock\n",
    "e_scf, wfn_scf = psi4.energy('scf', return_wfn=True)\n",
    "\n",
    "# Perform FCI calculation\n",
    "e_fci, wfn_fci = psi4.energy('fci/{}'.format(n_active_orb, n_active_elec), ref_wfn=wfn_scf, return_wfn=True)\n",
    "\n",
    "# Set options for CASSCF\n",
    "psi4.set_options({\n",
    "    'icore': 1,\n",
    "    'restricted_docc': 3\n",
    "})\n",
    "\n",
    "# Perform CASSCF calculation\n",
    "e_casscf, wfn_casscf = psi4.energy('casscf/{}'.format(n_active_orb, n_active_elec), ref_wfn=wfn_scf, return_wfn=True)\n",
    "\n",
    "# Print results\n",
    "print('')\n",
    "print(\"* Final oo-VQE energy:        {:.6f} hartree\".format(results.optimal_energy))\n",
    "print(\"* Final CASSCF/sto-3g energy: {:.6f} hartree\".format(e_casscf))\n",
    "print(\"* Final FCI/sto-3g energy:    {:.6f} hartree\".format(e_fci))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "399c3078-b91c-4ed3-b155-b8fb2c2144ac",
   "metadata": {},
   "source": [
    "As seen, our final oo-VQE energy is exactly the one expected when performing a CASSCF/sto-3g single-point calculation with the chosen active space."
   ]
  }
 ],
 "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}