{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d817ca15",
   "metadata": {},
   "source": [
    "## Exercise 13: The Variational Quantum Eigensolver (VQE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bfb70ac",
   "metadata": {},
   "source": [
    "In this exercise we are going to see how variational methods can be used in a hybrid quantum-classical scheme in order to approximate quantum states.\n",
    "\n",
    "In particular the aim of the Variational Quantum Eigensolver (VQE) is to optimize a set of parameterized quantum gates to approximate the ground state of a quantum system, given that the expectation value of the energy and its derivatives with respect to the parameters are evaluated on a quantum hardware in a scalable way.\n",
    "\n",
    "Preparing ground state is an exponentially complex task even on a quantum computer (QMA), but the VQE is expected to give some advantage for specific tasks in physics, chemistry and/or material sciences."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcaedb85",
   "metadata": {},
   "source": [
    "### Classical simulation of $H_2$ \n",
    "\n",
    "Here we will recall the code for the classical simulation of the dissociation process of $H_2$ diatomic molecule."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1bbd29a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install pyscf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9a522498",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pyscf import gto,scf,ao2mo,mp,cc,fci,tools\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "np.set_printoptions(precision=2,suppress=1e-8,linewidth=120)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d868346",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Define a set of distances between the two atoms\n",
    "\n",
    "distances = np.arange(0.3, 4, .05)\n",
    "\n",
    "## and a basis in which the calculations will be made\n",
    "\n",
    "basis = 'sto-6g' #'6-31g' 'cc-pvdz' 'aug-cc-pvdz' ,...\n",
    "\n",
    "## Define a dictionary containing the energies derived with different methods\n",
    "energies = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6eb5cbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "energies[\"HF_\"+basis]   = []\n",
    "energies[\"FCI_\"+basis]  = []\n",
    "energies[\"CCSD_\"+basis] = []\n",
    "\n",
    "for (i,r) in enumerate(distances):\n",
    "    # Build the 3D geometry of the molecule\n",
    "    geometry = \"H .0 .0 .0; H .0 .0 \"+str(r)\n",
    "    # And pass it to the gaussian orbitals generator (this is an alternative way to create the molecule object)\n",
    "    mol = gto.M(atom=geometry,charge=0,spin=0,basis=basis,symmetry=True,verbose=0)\n",
    "    \n",
    "    mf  = scf.RHF(mol) # Initialise a Restricted HF calculation using the molecule constructed\n",
    "    Ehf = mf.kernel()  #<- calling the kernel we compute the energy using the orbitals obtained\n",
    "    \n",
    "    fci_h2 = fci.FCI(mf) # FCI calculation\n",
    "    Efci = fci_h2.kernel()[0]\n",
    "    \n",
    "    \n",
    "    ccsd_h2 = cc.CCSD(mf) # CCSD calculation\n",
    "    e_ccsd  = ccsd_h2.kernel()[0]\n",
    "    e_ccsd += Ehf # <- this lines is mandatory because CCSD computes the energy difference with HF\n",
    "    \n",
    "    \n",
    "    # Save energies\n",
    "    energies[\"HF_\"+basis].append(Ehf)\n",
    "    energies[\"FCI_\"+basis].append(Efci)\n",
    "    energies[\"CCSD_\"+basis].append(e_ccsd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bf81f26",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Now plot the result\n",
    "\n",
    "plt.plot(distances,energies[\"HF_\"+basis],label=\"HF - \"+basis)\n",
    "plt.plot(distances,energies[\"FCI_\"+basis],label=\"FCI - \"+basis)\n",
    "plt.plot(distances,energies[\"CCSD_\"+basis],label=\"CCSD - \"+basis,linestyle=\"\",marker=\".\")\n",
    "\n",
    "plt.xlabel(r\"$r$ [$\\AA$]\")\n",
    "plt.ylabel(r\"$E$ [Hartree]\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77281a76",
   "metadata": {},
   "source": [
    "As can be seen, the CCSD calculation is as accurate as FCI in this small system."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "983098bb",
   "metadata": {},
   "source": [
    "### Quantum simulation of $H_2$\n",
    "\n",
    "Now we will proceed to the simulation on the quantum computer.\n",
    "\n",
    "As showed [here](https://www.nature.com/articles/nature23879) the Variational quantum eigensolver can be used to calculate dissociation processes of small molecules on quantum processor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e1881b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qiskit import QuantumCircuit\n",
    "from qiskit_aer import Aer\n",
    "from qiskit.circuit import ParameterVector\n",
    "from qiskit.quantum_info import SparsePauliOp # To create operators\n",
    "from qiskit.primitives import Estimator       # to estimate expectation values"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68358331",
   "metadata": {},
   "source": [
    "#### Preparing the Hamiltonian\n",
    "\n",
    "First, since we want to study a fermionic system on a set of qubits, we must find a mapping between the two.\n",
    "\n",
    "The oldest mapping proposed is the Jordan-Wigner, which was illustrated also in the lectures, but there are many others.\n",
    "\n",
    "Two noticeable alternatives are the [parity mapping](https://iopscience.iop.org/article/10.1088/1367-2630/aac54f) and the [Bravyi-Kitaev mapping](https://pubs.acs.org/doi/10.1021/acs.jctc.8b00450).\n",
    "\n",
    "In particular, we are going to focus on the parity mapping for the $H_2$ system. \n",
    "\n",
    "Instead of storing the occupation of a fermionic state $f_{j}$, as Jordan-Wigner does, this isomorphism uses the qubits to store the quantity \n",
    "\\begin{equation}\n",
    "\\label{eq:parity}\n",
    "p_{j}=\\sum_{i=0}^{j-1}f_{i} \\mod 2 \\quad ,\n",
    "\\end{equation}\n",
    "called parity of the set of occupation numbers $f_{j-1}\\dots f_{0}$. \n",
    "\n",
    "This is convenient for system like $H_2$ which conserve the total number of electrons with fixed spin orientation, meaning that we can get rid of two qubits out of the box.\n",
    "\n",
    "Manually performing the mapping can be a very tedious and prone-to-error task, for this reason we will use the very practical Qiskit PySCF driver, which recalls the creation of a PySCF molecule object that we have seen in previous lectures but gives a qubit operator ready to be measured on a quantum computer.\n",
    "\n",
    "This driver can be found in the Qiskit optional package `nature`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb5858c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install 'qiskit-nature'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb7b76e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qiskit_nature.second_q.drivers import PySCFDriver\n",
    "from qiskit_nature.second_q.formats.molecule_info import MoleculeInfo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e49bc488",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialise the molecule using the PySCF driver\n",
    "\n",
    "molecule = MoleculeInfo([\"H\", \"H\"], [(0.0, 0.0, 0.0), (0.0, 0.0, 0.735)], charge=0, multiplicity=1)\n",
    "driver = PySCFDriver.from_molecule(molecule, basis=\"sto6g\")\n",
    "problem = driver.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d9a0db5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now the problem has also the nuclear repulsion shift of the Born-Oppenheimer approximation\n",
    "problem.nuclear_repulsion_energy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8c7fed6",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Calling .second_q_ops will generate a bunch of operators such as dipole, magnetization, ...\n",
    "## We only need the Hamiltonian\n",
    "hamiltonian = problem.second_q_ops()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73847e36",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(hamiltonian)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46688c08",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Now we will map the problem to a qubit operator\n",
    "from qiskit_nature.second_q.mappers import JordanWignerMapper, BravyiKitaevMapper, ParityMapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a961a14",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Create the qubit operator \n",
    "#mapper = JordanWignerMapper()\n",
    "#mapper = BravyiKitaevMapper()\n",
    "mapper = ParityMapper()\n",
    "qubit_op = mapper.map(hamiltonian)\n",
    "print(qubit_op)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6576ac3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Create the qubit operator (With 2 qubit reduction)\n",
    "mapper = ParityMapper(num_particles=problem.num_particles) # Parity mapper requires the spin up and down electrons\n",
    "qubit_op = mapper.map(hamiltonian)\n",
    "print(qubit_op)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bdaed5a",
   "metadata": {},
   "source": [
    "Now we have seen how to create a qubit operator for a single configuration, we will need to repeat this procedure for every $r$ in order to study the dissociation process of $H_2$.\n",
    "\n",
    "#### Initial state\n",
    "\n",
    "We need a starting point for our calculations, that in this case will be the Hartree-Fock (HF) determinant.\n",
    "This can be easily prepared on the quantum circuit, since we are already in the basis of the HF orbitals.\n",
    "This means that we will need to put an $X$ gate if the corresponding spin-orbital is occupied (Jordan-Wigner) or it has a odd occupation number (Parity). \n",
    "\n",
    "Qiskit provides a useful `HartreeFock` function to create this initial state, that then we will pass to the VQE function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bc7ddd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qiskit_nature.second_q.circuit.library.initial_states import HartreeFock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9d634d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "init_state = HartreeFock(num_spatial_orbitals=problem.num_spatial_orbitals, num_particles=problem.num_particles, qubit_mapper=mapper)\n",
    "init_state.draw(\"mpl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b498c95e",
   "metadata": {},
   "source": [
    "#### Variational wave function\n",
    "\n",
    "Now that we have the Hamiltonian operator, we will need an ansatz for our wavefunction. Since we need a variational ansatz, some gates will contains parameters that are going to be iteratively optimized (as an example, rotation on the $x,y,z$ axis).\n",
    "\n",
    "There is no analytical way to choose an ansatz for the system: there are empirical rules based on similarity with what we are studying.\n",
    "Some ans&auml;tze come from classical computational chemistry, such as the highly accurate [q-UCCSD](https://arxiv.org/pdf/1506.00443.pdf),  but mostly we have to consider some circuits that can be run on current devices, so they have to contain few two qubits gates and be relatively shallow: these ans&auml;tze are called hardware-efficient.\n",
    "\n",
    "\n",
    "What we are going to consider is one of the so-called \"hardware-efficient ans&auml;tze\". \n",
    "The system does not contain many qubits, so our trial ansatz will be very simple : a layer of rotations around the $y$-axis followed by CNOTs and again a layer of rotations. This simple structure can be easily extended both in depth (adding more CNOTs and rotation) and in width, to study bigger system, therefore is widely used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab10c0bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def variational_ansatz(n_qubits,params):\n",
    "    \n",
    "    qc = QuantumCircuit(n_qubits)\n",
    "    \n",
    "    for i in range(n_qubits):\n",
    "        qc.ry(params[i],i)\n",
    "    qc.barrier()\n",
    "    for i in range(n_qubits-1):\n",
    "        qc.cx(i,i+1)\n",
    "    qc.barrier()\n",
    "    for i in range(n_qubits):\n",
    "        qc.ry(params[n_qubits+i],i)\n",
    "    return qc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80741ad5",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Let's plot an example\n",
    "\n",
    "params = ParameterVector('θ',10)\n",
    "variational_ansatz(5, params).draw(\"mpl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "332d5c92",
   "metadata": {},
   "source": [
    "#### Evaluating energy and the gradient\n",
    "\n",
    "Now we focus on the core part of the VQE algorithm: the evaluation of the energy and of its gradient on the quantum hardware.\n",
    "\n",
    "Qiskit has built-in methods for this, but we are going to implement ours from scratch."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dca6b018",
   "metadata": {},
   "source": [
    "First, we want a circuit to measure the expecation value of the Hamiltonian on our trial state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3442af4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def energy(quantum_circuit, hamiltonian, parameters, estimator):\n",
    "    ## This function evaluates the energy during a VQE calculation\n",
    "    return estimator.run(quantum_circuit, hamiltonian, parameters).result().values[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99c303c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Example energy estimation\n",
    "param_vec = ParameterVector(\"θ\",4)\n",
    "q_circuit = variational_ansatz(2,param_vec)\n",
    "estimator = Estimator()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5a3dafb",
   "metadata": {},
   "outputs": [],
   "source": [
    "energy(q_circuit, qubit_op, np.random.rand(4), estimator)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efa2c5a9",
   "metadata": {},
   "source": [
    "Now things get interesting! In order to optimize our parameters $\\theta$, a quantum computer has to give to classical optimizer important information such as first- or second-order derivatives! But how can we measure the derivative wrt to a specific parameter on a quantum circuit?\n",
    "\n",
    "$$\n",
    "\\frac{\\partial }{\\partial \\theta_i} L(\\theta) = \\frac{\\partial }{\\partial \\theta_i} \\langle \\psi(\\theta) | H | \\psi(\\theta) \\rangle\n",
    "$$\n",
    "\n",
    "many different methods have been proposed recently, in this case we are going to see a method  that is called _parameter shift_:\n",
    "\n",
    "- assume that every parametrized gate is of the form $$ U_j(\\theta_j) = e^{-i \\theta_j G_j}= \\cos(\\theta_j)\\mathbf{I} -i \\sin(\\theta_j)G_j$$ where $G_j$ is an operator such that $G_{j}^{2} = \\mathbf{I}$\n",
    "\n",
    "- then the derivative can be expressed as $$ \\frac{\\partial }{\\partial \\theta_i} L(\\theta) = \\frac{L(\\theta+ e_i s)-L(\\theta-e_is)}{2\\sin(s)}  $$ where $s \\in \\mathbf{R}$ and $e_i$ indicates the versor in the $i$-th direction. In our case we are going to consider $s= \\frac{\\pi}{2}$.\n",
    "\n",
    "This means that to calculate the gradient with $N_p$ parameters, we have to measure $H$ on $2N_p$ different circuits, but it's possible!\n",
    "\n",
    "Note that the assumption we made is quite general: $G_j$ could be every tensor product of Pauli operators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d40bafd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#useful function to shift the parameters\n",
    "def ei(i,n):\n",
    "    '''unit vector in direction i of size n'''\n",
    "    vi = np.zeros(n)\n",
    "    vi[i] = 1.0\n",
    "    return vi[:]\n",
    "\n",
    "\n",
    "def gradient(quantum_circuit, hamiltonian, parameters, estimator):\n",
    "\n",
    "    # TODO implement the gradient with the parameter-shift rule\n",
    "    # you need to pass the estimator object to the energy(...) function defined above.\n",
    "\n",
    "    # g = ...\n",
    "    # return g\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "867b7184",
   "metadata": {},
   "source": [
    "#### The VQE algorithm\n",
    "\n",
    "Now we declare a function to repeat iteratively the procedure of measuring energy, its gradient and then optimizing the parameters using a standard gradient descent technique, namely $$ \\theta_{new}= \\theta_{old} - \\eta \\nabla_{\\theta} L(\\theta)$$ with $\\eta \\in \\mathbf{R}$ as the learning rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af9861a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def VQE(operator=None, ansatz=None, init_params=None, init_state=None, estimator=None, max_iter=100, lr=0.1):\n",
    "    \"\"\"\n",
    "    Args: \n",
    "        operator: the hamiltonian as a qubit operator\n",
    "        ansatz: the variational ansatz\n",
    "        init_params: the initital parameters\n",
    "        init_state: circuit to prepare the initial state (e.g. the Hartree-Fock state)\n",
    "        estimator: qiskit Estimator object, to be passed to energy(...) and gradient(...)\n",
    "        max_iter: number of gradient descent iterations\n",
    "        lr: learning rate η\n",
    "    Returns:\n",
    "        energies: list of estimated energies at each step of the vqe \n",
    "        gradients:  list of estimated gradients at each step of the vqe \n",
    "    \"\"\"\n",
    "    # compose initial state preparation and ansatz\n",
    "    quantum_circuit = init_state.compose(ansatz(operator.num_qubits, ParameterVector('θ', len(init_params))))\n",
    "    \n",
    "    # TODO implement the vqe using the energy(...) and gradient(...) functions defined above\n",
    "    \n",
    "    # energies = ...\n",
    "    # gradients = ...\n",
    "    # return energies, gradients"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0eb21f92",
   "metadata": {},
   "source": [
    "#### Perform the optimization on a single configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "820a4d86",
   "metadata": {},
   "outputs": [],
   "source": [
    "molecule    = MoleculeInfo([\"H\", \"H\"], [(0.0, 0.0, 0.0), (0.0, 0.0, 0.735)], charge=0, multiplicity=1)\n",
    "driver      = PySCFDriver.from_molecule(molecule, basis=\"sto6g\")\n",
    "problem     = driver.run()\n",
    "hamiltonian = problem.second_q_ops()[0]\n",
    "    \n",
    "nucl_shift  = problem.nuclear_repulsion_energy\n",
    "\n",
    "\n",
    "mapper      = ParityMapper(num_particles=problem.num_particles) # Parity mapper requires the spin up and down electrons\n",
    "qubit_op    = mapper.map(hamiltonian)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9a280f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now create the circuit\n",
    "n_qubits    = qubit_op.num_qubits\n",
    "init_params = np.random.rand(2*n_qubits)\n",
    "init_state  = HartreeFock(num_spatial_orbitals=problem.num_spatial_orbitals, num_particles=problem.num_particles, qubit_mapper=mapper)\n",
    "\n",
    "\n",
    "n_reps = 50\n",
    "lr = 0.5\n",
    "estimator = Estimator()\n",
    "\n",
    "energies_vqe, gradients_vqe = VQE(qubit_op,variational_ansatz,init_params,init_state,estimator,n_reps,lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a51420e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "steps = list(range(n_reps))\n",
    "\n",
    "plt.figure(figsize=(4.8,3.4),dpi=100)\n",
    "plt.errorbar(steps,np.array(energies_vqe)+nucl_shift,marker='o',linestyle='dashed',label=\"VQE\")\n",
    "plt.hlines(-1.145, xmin= -10, xmax= 1000,label='Theoretical',linestyle ='dashed',color='black')\n",
    "plt.xlabel('Step')\n",
    "plt.xlim(xmin=-1,xmax=51)\n",
    "plt.ylabel('Energy')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88c12825",
   "metadata": {},
   "source": [
    "#### Quantum simulation\n",
    "\n",
    "Now we are going to perform the VQE algorithm for every configuration of the $H_2$ molecule."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57d4d448",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare the backend\n",
    "estimator = Estimator()\n",
    "vqe_distances = list(np.arange(0.3, 1.5, .1)) + list(np.arange(1.5, 4.0, .2))\n",
    "\n",
    "energies[\"VQE_\"+basis] = []\n",
    "energies[\"VQE_error\"+basis] = []\n",
    "\n",
    "# Repeat for every r\n",
    "for (i,r) in enumerate(vqe_distances):\n",
    "    print(\"\\n=======================================\")\n",
    "    print(f\" DISTANCE: {r:0.3}\")\n",
    "    print(\"=======================================\\n\")\n",
    "    \n",
    "    molecule    = MoleculeInfo([\"H\", \"H\"], [(0.0, 0.0, 0.0), (0.0, 0.0, r)], charge=0, multiplicity=1)\n",
    "    driver      = PySCFDriver.from_molecule(molecule, basis=\"sto6g\")\n",
    "    problem     = driver.run()\n",
    "    nucl_shift  = problem.nuclear_repulsion_energy\n",
    "    \n",
    "    \n",
    "    # Build the (reduced) qubit operator\n",
    "    hamiltonian = problem.second_q_ops()[0]\n",
    "    mapper     = ParityMapper(num_particles=problem.num_particles)\n",
    "    qubit_op   = mapper.map(hamiltonian)\n",
    "    \n",
    "    # Now create the circuit\n",
    "    n_qubits   = qubit_op.num_qubits\n",
    "    params     = np.random.rand(2*n_qubits)\n",
    "    init_state = HartreeFock(num_spatial_orbitals=problem.num_spatial_orbitals, num_particles=problem.num_particles, qubit_mapper=mapper)\n",
    "    \n",
    "    #Run the algorithm\n",
    "    n_reps  = 150\n",
    "    lr = 0.5\n",
    "    energies_vqe, gradients_vqe = VQE(qubit_op,variational_ansatz,params,init_state,estimator,n_reps,lr)\n",
    "    \n",
    "    print(\"Final energy: \",energies_vqe[-1]+nucl_shift)\n",
    "    \n",
    "    energies[\"VQE_\"+basis].append(energies_vqe[-1]+nucl_shift)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa826ead",
   "metadata": {},
   "source": [
    "#### Plot the final results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49766423",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(distances,energies[\"HF_\"+basis],linestyle=\"dashed\",label=\"HF - \"+basis)\n",
    "plt.plot(distances,energies[\"FCI_\"+basis],label=\"FCI - \"+basis)\n",
    "plt.errorbar(vqe_distances,energies[\"VQE_\"+basis],label=\"VQE - \"+basis,linestyle=\"\",marker=\"o\",markersize=5,mew=0.5,mec=\"black\")\n",
    "\n",
    "plt.xlabel(r\"$r$ [$\\AA$]\")\n",
    "plt.ylabel(r\"$E$ [Hartree]\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
