{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quantum Chemistry with Paddle Quantum's qchem\n",
    "*Copyright (c) 2021 Institute for Quantum Computing, Baidu Inc. All Rights Reserved.*\n",
    "\n",
    "qchem, which builds on top of Paddle Quantum, is designed to be a toolkit for quantum chemistry research in quantum computing era. It provides high level APIs for researchers who are interested in leveraging their current quantum chemsitry calculation with quantum computing power. And it also allows experts to write customized code. qchem is currently under active development, feel free to join us by opening issues or pull requests."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate ground state energy of a molecule\n",
    "qchem provides `run_chem` function to calculate the ground state energy and ground state wave function. For example, let's try to calculate the ground state energy of hydrogen molecule. First, we need to import qchem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from paddle_quantum import qchem"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we need to provide chemical information required by `run_chem` function, including geometry, charge, basis function etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define the geometry of hydrogen molecule, length unit use angstrom.\n",
    "h2_geometry = [(\"H\", [0.0, 0.0, 0.0]), (\"H\", [0.0, 0.0, 0.74])]\n",
    "charge = 0\n",
    "basis_set = \"sto-3g\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's choose an ansatz of many-electron wave function for our ground state energy calculation. Currently, you can choose to use \"hardware efficient\" [<sup>1</sup>](#refer-1) or \"hartree fock\" [<sup>2</sup>](#refer-2) ansatz, more functionalities will be released in the future. Let's choose \"hardware efficient\" ansatz."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# call run_chem function with \"hardware efficient\" ansatz.\n",
    "h2_gs_en, h2_wf_model = qchem.run_chem(\n",
    "    h2_geometry,\n",
    "    \"hardware efficient\",\n",
    "    basis_set, \n",
    "    charge\n",
    ")\n",
    "\n",
    "# additional information for optimizer can be passed using `optimizer_option` keyword argument.\n",
    "h2_gs_en, h2_wf_model = qchem.run_chem(\n",
    "    h2_geometry,\n",
    "    \"hardware efficient\",\n",
    "    basis_set, \n",
    "    charge,\n",
    "    optimizer_option={\"learning_rate\": 0.6, \"weight_decay\": 0.9}\n",
    ")\n",
    "\n",
    "# additional information for ansatz can be passed using `ansatz_option` keyword argument, e.g.\n",
    "# \"hardware efficient\" ansatz has a parameter \"cir_depth\", which can be used to specify the depth\n",
    "# of quantum circuit.\n",
    "h2_gs_en, h2_wf_model = qchem.run_chem(\n",
    "    h2_geometry,\n",
    "    \"hardware efficient\",\n",
    "    basis_set, \n",
    "    charge,\n",
    "    ansatz_option={\"cir_depth\": 5}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also specify `max_iters` and `a_tol` keyword arguments to control the maximum iteration cycles and convergence criteria of the optimization process. \n",
    "\n",
    "To try another ansatz, you just need to replace the ansatz argument to, e.g. \"hartree fock\", and run similar command.\n",
    "\n",
    "To see the quantum circuit for hydrogen molecule's hardware efficient ansatz, you can run `print(h2_wf_model.circuit)`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Design your own ansatz\n",
    "For those who want to try an ansatz that isn't currently included in qchem, we provide method to write your own ansatz. Writing your own ansatz is similar to defining an neural network in paddlepaddle, except that your ansatz should inherit from `Qmodel`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from paddle_quantum.circuit import UAnsatz\n",
    "from paddle_quantum import qchem\n",
    "from paddle_quantum.qchem.layers import CrossResonanceLayer, EulerRotationLayer\n",
    "\n",
    "\n",
    "# Your own model should inherit from `Qmodel`.\n",
    "## NOTE: THIS MODEL IS ONLY DEFINED FOR DEMONSTRATION PURPOSE! \n",
    "class MyAnsatz(qchem.QModel):\n",
    "    def __init__(self, n_qubits):\n",
    "        super().__init__(n_qubits)\n",
    "\n",
    "        self.entangle = CrossResonanceLayer(self._n_qubit)\n",
    "        self.rot = EulerRotationLayer(self._n_qubit)\n",
    "\n",
    "    def forward(self, state):\n",
    "        self._circuit = UAnsatz(self.n_qubit)\n",
    "\n",
    "        out = self.entangle(state)\n",
    "        self._circuit += self.entangle.circuit\n",
    "\n",
    "        out = self.rot(out)\n",
    "        self._circuit += self.rot.circuit\n",
    "\n",
    "        out = self.entangle(out)\n",
    "        self._circuit += self.entangle.circuit\n",
    "\n",
    "        return out\n",
    "\n",
    "# instantiate your model\n",
    "my_cir = MyAnsatz(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can then follow the optimization procedure [here](https://qml.baidu.com/tutorials/quantum-simulation/variational-quantum-eigensolver.html) and use any paddlepaddle optimizer to train the ansatz you have built."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use paddlepaddle's optimizer\n",
    "import numpy as np \n",
    "import paddle\n",
    "\n",
    "optimizer = paddle.optimizer.Adam(parameters=my_cir.parameters(), learning_rate=0.08)\n",
    "\n",
    "# define the loss function\n",
    "## NOTE: THIS LOSS FUNCTION IS ONLY DEFINED FOR DEMONSTRATION PURPOSE!\n",
    "def loss_fn(state: paddle.Tensor) -> paddle.Tensor:\n",
    "    return paddle.norm(state.real())\n",
    "\n",
    "# start learning\n",
    "s0 = np.zeros((2**5,), dtype=np.complex128)\n",
    "s0[0] = 1.0+0.0j\n",
    "s0 = paddle.to_tensor(s0)\n",
    "\n",
    "for i in range(10):\n",
    "    loss = loss_fn(my_cir(s0))\n",
    "    print(f\"At {i:>d}th step: loss={loss.item():>.5f}.\")\n",
    "\n",
    "    optimizer.clear_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## References\n",
    "\n",
    "[1] Kandala, Abhinav, et al. \"Hardware-efficient variational quantum eigensolver for small molecules and quantum magnets.\" [Nature 549.7671 (2017): 242-246.](https://www.nature.com/articles/nature23879)\n",
    "\n",
    "[2] Arute, Frank, et al. \"Hartree-Fock on a superconducting qubit quantum computer.\" [Science 369.6507 (2020): 1084-1089.](https://www.science.org/doi/10.1126/science.abb9811)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "7b343e878babc25549085ff27e754b596ec1b81bbbd70d50b28da4f6023d5bd9"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}