Coding_Challenges/qchem_500_MindTheGap_template/qchem_500_MindTheGap.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "name": "qchem_500_MindTheGap.ipynb",
   "private_outputs": true,
   "provenance": [],
   "authorship_tag": "ABX9TyPMfjgkGnQo+j+OaGQXjU6w",
   "include_colab_link": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "view-in-github",
    "colab_type": "text"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/JavierPerez21/QHack2022/blob/master/Coding_Challenges/qchem_500_MindTheGap_template/qchem_500_MindTheGap.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WNKQPhtGlc98"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "!pip install pennylane"
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "import pennylane as qml\n",
    "from pennylane import numpy as np\n",
    "from pennylane import hf"
   ],
   "metadata": {
    "id": "8ybJLQCGqeC2"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "The objective of this challenge is to find the ground state energy as well as the energy of the first excited state of Hydrogen. We can represent the $H_2$ molecule. We can represent the $H_2$ molecule in dirac notation as a state:\n",
    "$$|H_{2, init}\\rangle = \\begin{pmatrix} 1 \\\\ 1 \\\\ 0 \\\\ 0 \\end{pmatrix}$$\n",
    "where each 1 represents an electron.\n",
    "\n",
    "Since we are dealing with the representation of a molecule, the number of particles must be preserved, so we must use particle preserving gates, namely [Givens rotations](https://pennylane.ai/qml/demos/tutorial_givens_rotations.html).\n",
    "\n",
    "Given the Hamiltonian of $H_2$, $\\hat{H}$, the energy of the ground state can be found using DoubleExcitationGates as:\n",
    "$$E(\\theta) = \\langle \\Psi_0(\\theta) | \\hat{H} | \\Psi_0(\\theta) \\rangle = \\langle H_{2, init}| G_2(\\theta)^\\dagger \\; \\hat{H} \\; G_2(\\theta) |H_{2, init}\\rangle$$  \n",
    "\n",
    "One just needs to implement [VQE](https://pennylane.ai/qml/demos/tutorial_vqe.html) to find the correct angle $\\theta$ that characterizes the ground state."
   ],
   "metadata": {
    "id": "X39ilDpOmQI3"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "def ground_state_VQE(H):\n",
    "    \"\"\"Perform VQE to find the ground state of the H2 Hamiltonian.\n",
    "    Args:\n",
    "        - H (qml.Hamiltonian): The Hydrogen (H2) Hamiltonian\n",
    "    Returns:\n",
    "        - (float): The ground state energy\n",
    "        - (np.ndarray): The ground state calculated through your optimization routine\n",
    "    \"\"\"\n",
    "\n",
    "    # QHACK #\n",
    "    \n",
    "    # Initialize thhe fevice\n",
    "    num_qubits = len(H.wires)\n",
    "    qubits = H.wires\n",
    "    num_param_sets = (2 ** num_qubits) - 1\n",
    "    dev = qml.device(\"default.qubit\", wires=num_qubits)\n",
    "\n",
    "    # Initialize the energy, theta and hi state\n",
    "    energy = 0\n",
    "    hi = np.array([1, 1, 0, 0])\n",
    "    theta = np.random.uniform(low=-np.pi / 2, high=np.pi / 2, size=(1))\n",
    "\n",
    "    # Create circuit for VQE\n",
    "    def circuit(theta):\n",
    "        qml.DoubleExcitation(theta[0], wires=qubits)\n",
    "\n",
    "    # Define cost function and a function to return the ground_state\n",
    "    @qml.qnode(dev)\n",
    "    def cost_fn(theta):\n",
    "        qml.BasisState(hi, wires=qubits)\n",
    "        circuit(theta)\n",
    "        return qml.expval(H)\n",
    "    @qml.qnode(dev)\n",
    "    def get_ground_state(theta):\n",
    "        qml.BasisState(hi, wires=qubits)\n",
    "        circuit(theta)\n",
    "        return qml.state()\n",
    "\n",
    "    # Define the optimizer, parameters of the VQE and lists to keep track of parameter and results at every iteration\n",
    "    opt = qml.optimize.AdamOptimizer(0.1)\n",
    "    max_iterations = 200\n",
    "    conv_tol = 1e-06\n",
    "    energy = [cost_fn(theta)]\n",
    "    theta_col = [theta]\n",
    "    states = [get_ground_state(theta)]\n",
    "\n",
    "    # Perform VQE to find the angle theta\n",
    "    for n in range(max_iterations):\n",
    "        theta, prev_energy = opt.step_and_cost(cost_fn, theta)\n",
    "\n",
    "        energy.append(cost_fn(theta))\n",
    "        theta_col.append(theta)\n",
    "        states.append(get_ground_state(theta))\n",
    "\n",
    "        conv = np.abs(energy[-1] - prev_energy)\n",
    "\n",
    "        if n % 2 == 0:\n",
    "            print(f\"Step = {n},  Energy = {energy[-1]:.8f} HH\")\n",
    "\n",
    "        if conv <= conv_tol:\n",
    "            break\n",
    "    print(\"\\n\" f\"Final value of the ground-state energy = {energy[-1]:.8f} HH\")\n",
    "    print(\"\\n\" f\"Optimal value of the circuit parameters = {theta_col[-1]}\")\n",
    "    print(\"\\n\" f\"Ground-state state vector = {states[-1]}\")\n",
    "\n",
    "    # Return the ground state energy and the ground state\n",
    "    return energy[-1], states[-1]\n",
    "    # QHACK #\n"
   ],
   "metadata": {
    "id": "8U-nudIenFb-"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "As illustrated in the problem pdf, given the ground state $|\\lambda_0\\rangle = |\\Psi_0(\\theta^*)\\rangle$ and the Ground state Hamiltonian $\\hat{H}$ we can calculate the excited state Hamiltonian as:\n",
    "$$\\hat{H}_1=\\hat{H} + \\beta |\\lambda_0\\rangle \\langle \\lambda_0 | = (E_0 + \\beta)  |\\lambda_0\\rangle \\langle \\lambda_0 |  + \\sum_i E_ i |\\lambda_i\\rangle \\langle \\lambda_i | $$\n",
    "\n",
    "This can be done using the decomposing the matrix $(E_0 + \\beta)  |\\lambda_0\\rangle \\langle \\lambda_0 |$ into observables and coefficients\n",
    "\n",
    "Then one can calculate the Ground state energy as done before and perfor VQE to find the new correct angle $\\theta$. The only difference is that, in this case, since we are working with the first excited states, one of the electrons of the molecules will be in its second energy level so we should use Single Excitation Givens gates.\n",
    "\n",
    "$$E_1(\\theta) = \\langle \\Psi_1(\\theta) | \\hat{H}_1 | \\Psi_1(\\theta) \\rangle = \\langle H_{2, init}| G_1(\\theta)^\\dagger \\; \\hat{H}_1 \\; G_1(\\theta) |H_{2, init}\\rangle$$  "
   ],
   "metadata": {
    "id": "iLkGSCt-sbwk"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "def create_H1(ground_state, beta, H):\n",
    "    \"\"\"Create the H1 matrix, then use `qml.Hermitian(matrix)` to return an observable-form of H1.\n",
    "    Args:\n",
    "        - ground_state (np.ndarray): from the ground state VQE calculation\n",
    "        - beta (float): the prefactor for the ground state projector term\n",
    "        - H (qml.Hamiltonian): the result of hf.generate_hamiltonian(mol)()\n",
    "    Returns:\n",
    "        - (qml.Observable): The result of qml.Hermitian(H1_matrix)\n",
    "    \"\"\"\n",
    "\n",
    "    # QHACK #\n",
    "    # Define function to decompose hamiltonian into observables and coefficients\n",
    "    from functools import reduce\n",
    "    from itertools import product\n",
    "    from operator import matmul\n",
    "    def decompose_hamiltonian(H):\n",
    "        N = int(np.log2(len(H)))\n",
    "        paulis = [qml.Identity, qml.PauliX, qml.PauliY, qml.PauliZ]\n",
    "\n",
    "        obs = []\n",
    "        coeffs = []\n",
    "\n",
    "        for term in product(paulis, repeat=N):\n",
    "            matrices = [i._matrix() for i in term]\n",
    "            coeff = np.trace(reduce(np.kron, matrices) @ H) / (2 ** N)\n",
    "\n",
    "            if not np.allclose(coeff, 0):\n",
    "                coeffs.append(coeff)\n",
    "\n",
    "                if not all(t is qml.Identity for t in term):\n",
    "                    obs.append(reduce(matmul, [t(i) for i, t in enumerate(term) if t is not qml.Identity]))\n",
    "                else:\n",
    "                    obs.append(reduce(matmul, [t(i) for i, t in enumerate(term)]))\n",
    "\n",
    "        return coeffs, obs\n",
    "\n",
    "    # Get new coefficients and new observables to add to H\n",
    "    gs = ground_state.reshape(-1, 1)\n",
    "    new_coefficients, new_observables = decompose_hamiltonian(beta * (gs @ gs.T))\n",
    "\n",
    "    # Add new observables and new coefficients\n",
    "    new_coefficients = np.tensor(new_coefficients)\n",
    "    new_coefficients = np.concatenate((new_coefficients, H.terms[0]))\n",
    "    new_observables += H.terms[1]\n",
    "    \n",
    "    # Create new H1 hamiltonian out of new lists of coefficients and observables\n",
    "    H1 = qml.Hamiltonian(new_coefficients, new_observables)\n",
    "    return H1\n",
    "    # QHACK #"
   ],
   "metadata": {
    "id": "GJyabL_1vFuZ"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "def excited_state_VQE(H1):\n",
    "    \"\"\"Perform VQE using the \"excited state\" Hamiltonian.\n",
    "    Args:\n",
    "        - H1 (qml.Observable): result of create_H1\n",
    "    Returns:\n",
    "        - (float): The excited state energy\n",
    "    \"\"\"\n",
    "\n",
    "    # QHACK #\n",
    "    # Initialize thhe fevice\n",
    "    num_qubits = len(H1.wires)\n",
    "    qubits = H1.wires\n",
    "    num_param_sets = (2 ** num_qubits) - 1\n",
    "    dev = qml.device(\"default.qubit\", wires=num_qubits)\n",
    "\n",
    "    # Initialize the energy, theta and hi state\n",
    "    energy = 0\n",
    "    hi = np.array([1, 1, 0, 0])\n",
    "    theta = np.random.uniform(low=-np.pi / 2, high=np.pi / 2, size=(1))\n",
    "\n",
    "    # Create circuit for VQE\n",
    "    def circuit(theta):\n",
    "        qml.SingleExcitation(theta[0], wires=[1, 2])\n",
    "\n",
    "    # Define cost function and a function to return the ground_state\n",
    "    @qml.qnode(dev)\n",
    "    def cost_fn(theta):\n",
    "        qml.BasisState(hi, wires=qubits)\n",
    "        circuit(theta)\n",
    "        return qml.expval(H1)\n",
    "    @qml.qnode(dev)\n",
    "    def get_ground_state(theta):\n",
    "        qml.BasisState(hi, wires=qubits)\n",
    "        circuit(theta)\n",
    "        return qml.state()\n",
    "\n",
    "    # Define the optimizer, parameters of the VQE and lists to keep track of parameter and results at every iteration\n",
    "    opt = qml.optimize.MomentumOptimizer()\n",
    "    max_iterations = 300\n",
    "    conv_tol = 1e-05\n",
    "    energy = [cost_fn(theta)]\n",
    "    theta_col = [theta]\n",
    "    states = [get_ground_state(theta)]\n",
    "\n",
    "    # Perform VQE to find the angle theta\n",
    "    for n in range(max_iterations):\n",
    "        theta, prev_energy = opt.step_and_cost(cost_fn, theta)\n",
    "\n",
    "        energy.append(cost_fn(theta))\n",
    "        theta_col.append(theta)\n",
    "        states.append(get_ground_state(theta))\n",
    "\n",
    "        conv = np.abs(energy[-1] - prev_energy)\n",
    "\n",
    "        if n % 2 == 0:\n",
    "            print(f\"Step = {n},  Energy = {energy[-1]:.8f} HH\")\n",
    "\n",
    "        if conv <= conv_tol:\n",
    "            break\n",
    "    print(\"\\n\" f\"Final value of the ground-state energy = {energy[-1]:.8f} HH\")\n",
    "    print(\"\\n\" f\"Optimal value of the circuit parameters = {theta_col[-1]}\")\n",
    "    print(\"\\n\" f\"Ground-state state vector = {states[-1]}\")\n",
    "\n",
    "    # Return the ground state energy and the ground state\n",
    "    return energy[-1]\n",
    "    # QHACK #"
   ],
   "metadata": {
    "id": "OtZhF_mOwNAR"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Here we check the inputs of 1.in"
   ],
   "metadata": {
    "id": "FKgKRfKUw6Xf"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "coord = 0.6614\n",
    "symbols = [\"H\", \"H\"]\n",
    "geometry = np.array([[0.0, 0.0, -coord], [0.0, 0.0, coord]], requires_grad=False)\n",
    "mol = hf.Molecule(symbols, geometry)\n",
    "\n",
    "H = hf.generate_hamiltonian(mol)()\n",
    "E0, ground_state = ground_state_VQE(H)\n",
    "\n",
    "beta = 15.0 + E0\n",
    "H1 = create_H1(ground_state, beta, H)\n",
    "E1 = excited_state_VQE(H1)\n",
    "\n",
    "print(f\"Ground state energy E_0 = {E0.item()}\")\n",
    "print(f\"First excited state energy E_0 = {E1.item()}\")"
   ],
   "metadata": {
    "id": "5fVGfR36w0Nn"
   },
   "execution_count": null,
   "outputs": []
  }
 ]
}