{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b01c5570",
   "metadata": {},
   "source": [
    "# The UCCSD Wavefunction ansatz\n",
    "\n",
    "In quantum chemistry, the Unitary Coupled Cluster with Single and Double excitations (UCCSD) wavefunction ansatz [1](https://arxiv.org/pdf/1701.02691), [2](https://arxiv.org/pdf/2109.15176) is a powerful method used to approximate the ground state of a quantum system. This approach is particularly valuable in the context of quantum computing, where it is employed to solve the electronic structure problem. In this tutorial, we will learn how to implement the UCCSD in CUDA-Q to calclate the ground state energy for chemical systems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad507956",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Requires pyscf to be installed\n",
    "%pip install pyscf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "68a9f31e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cudaq\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\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c610617",
   "metadata": {},
   "source": [
    "We begin by calculating the electronic hamiltonian and converting it to qubit hamiltonian. To learn more, see this [tutorial](https://nvidia.github.io/cuda-quantum/latest/applications/python/generate_fermionic_ham.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "732bcd12",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "overwrite output file: Zn 0-pyscf.log\n",
      "[pyscf] Total number of orbitals =  57\n",
      "[pyscf] Total number of electrons =  38\n",
      "[pyscf] HF energy =  -1852.5618486460523\n",
      "[pyscf] R-CASCI energy using molecular orbitals=  -1852.589855127006\n",
      "[pyscf] R-CCSD energy of the active space using molecular orbitals=  -1852.5897242853828\n"
     ]
    }
   ],
   "source": [
    "from qchem.classical_pyscf import get_mol_hamiltonian\n",
    "from qchem.hamiltonian import jordan_wigner_fermion\n",
    "\n",
    "#geometry = 'H 0.0 0.0 0.0; H 0.0 0.0 0.7474'\n",
    "#geometry = 'Li 0.3925 0.0 0.0; H -1.1774 0.0 0.0'\n",
    "geometry =  'Zn 0.0 0.0 00; O 0.0 0.0 1.7047'\n",
    "\n",
    "# When not using active space.\n",
    "#molecular_data = get_mol_hamiltonian(xyz=geometry, spin=0, charge=0, basis='sto3g', ccsd=True, verbose=True)\n",
    "\n",
    "# When using active space.\n",
    "molecular_data = get_mol_hamiltonian(xyz=geometry, spin=0, charge=0, basis='ccpvdz', nele_cas=4, norb_cas=4, \n",
    "                                     ccsd=True, casci=True, verbose=True)\n",
    "\n",
    "obi = molecular_data[0]\n",
    "tbi = molecular_data[1]\n",
    "const = molecular_data[2]\n",
    "electron_count = molecular_data[3]\n",
    "norbitals = molecular_data[4]\n",
    "qubits_num = 2 * norbitals\n",
    "\n",
    "# Here, we are excluding the const from the jordan wigner transformation. \n",
    "# When calculating the expectation value of the Hamiltonian, \n",
    "# it is recommend to remove the identity operator from the \n",
    "# Hamiltonian and add its coefficient as a constant shift to the energy.\n",
    "# In exact arithmetic, this should not make a difference because the expectation value \n",
    "# for the identity is just the norm, which is 1.\n",
    "# However, in floating-point arithmetic, this avoids the round-off error \n",
    "# associated with that term.  This error can be significant because the \n",
    "# coefficient in this term corresponds to core energy, which is usually the largest \n",
    "# coefficient in the Hamiltonian by an order of magnitude.\n",
    "hamiltonian = jordan_wigner_fermion(obi, tbi, 0.0, tolerance = 1e-15)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecb2f0c1",
   "metadata": {},
   "source": [
    "### What is UCCSD?\n",
    "\n",
    "The UCCSD ansatz is an extension of the traditional coupled cluster method, which is widely used in classical computational chemistry. The key idea behind UCCSD is to construct the wavefunction as an exponential of an anti-Hermitian operator that includes single and double excitations. Mathematically, this can be expressed as:\n",
    "\n",
    "$$\n",
    "|\\Psi_{\\mathrm{UCCSD}}\\rangle = e^{T - T^\\dagger} |\\Phi\\rangle\n",
    "$$\n",
    "\n",
    "where $T$ is the cluster operator defined as:\n",
    "\n",
    "$$\n",
    "T = T_1 + T_2\n",
    "$$\n",
    "\n",
    "Here, $T_1$ and $T_2$ represent the single and double excitation operators, respectively.\n",
    "\n",
    "The UCCSD ansatz maintains the unitarity of the wavefunction, which is crucial for quantum computations. This unitarity ensures that the norm of the wavefunction remains constant, preserving the physical properties of the system.\n",
    "\n",
    "### Implementation in Quantum Computing\n",
    "\n",
    "In the context of quantum computing, the UCCSD ansatz is implemented using quantum circuits. The exponential operator $e^{T-T^\\dagger}$ is decomposed into a sequence of quantum gates that can be executed on a quantum computer. This decomposition typically involves the Trotter-Suzuki approximation, which approximates the exponential of a sum of operators by a product of exponentials of the individual operators.\n",
    "\n",
    "Below, we show how to get the UCCSD operators and how to prepare the quantum circuit using UCCSD ansatz. User can try to employ only single excitations, double excitations, or both singles and doubles excitations.\n",
    "\n",
    "Note: the `cudaq.kernels.uccsd(qubits, thetas, electron_count, qubits_num)` available in cudaq.kernels uses cnot and rotation gates to implement the exponential of pauli operators in the UCCSD, see this [tutorial](https://nvidia.github.io/cuda-quantum/latest/applications/python/vqe_advanced.html). The version we use in this tutorial employs the `exp_pauli()` gate available in CUDA-Q. Using `exp_pauli()` gate instead of decomposing it when running on classical simulator reduces the number of gates dramatically leading to improvements in perfomance. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2d355130",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 8 singles, 18 doubles, 26 total\n",
      "word_single: ['YZZZXIII', 'XZZZYIII', 'YZZZZZXI', 'XZZZZZYI', 'IIYZXIII', 'IIXZYIII', 'IIYZZZXI', 'IIXZZZYI', 'IYZZZXII', 'IXZZZYII', 'IYZZZZZX', 'IXZZZZZY', 'IIIYZXII', 'IIIXZYII', 'IIIYZZZX', 'IIIXZZZY']\n",
      "word_double: ['XXIIXYII', 'XXIIYXII', 'XYIIYYII', 'YXIIYYII', 'XYIIXXII', 'YXIIXXII', 'YYIIXYII', 'YYIIYXII', 'XXIIIXYI', 'XXIIIYXI', 'XYIIIYYI', 'YXIIIYYI', 'XYIIIXXI', 'YXIIIXXI', 'YYIIIXYI', 'YYIIIYXI', 'XXIIXZZY', 'XXIIYZZX', 'XYIIYZZY', 'YXIIYZZY', 'XYIIXZZX', 'YXIIXZZX', 'YYIIXZZY', 'YYIIYZZX', 'XXIIIIXY', 'XXIIIIYX', 'XYIIIIYY', 'YXIIIIYY', 'XYIIIIXX', 'YXIIIIXX', 'YYIIIIXY', 'YYIIIIYX', 'XZZXXYII', 'XZZXYXII', 'XZZYYYII', 'YZZXYYII', 'XZZYXXII', 'YZZXXXII', 'YZZYXYII', 'YZZYYXII', 'XZZXIXYI', 'XZZXIYXI', 'XZZYIYYI', 'YZZXIYYI', 'XZZYIXXI', 'YZZXIXXI', 'YZZYIXYI', 'YZZYIYXI', 'XZZXXZZY', 'XZZXYZZX', 'XZZYYZZY', 'YZZXYZZY', 'XZZYXZZX', 'YZZXXZZX', 'YZZYXZZY', 'YZZYYZZX', 'XZZXIIXY', 'XZZXIIYX', 'XZZYIIYY', 'YZZXIIYY', 'XZZYIIXX', 'YZZXIIXX', 'YZZYIIXY', 'YZZYIIYX', 'IXXIXYII', 'IXXIYXII', 'IXYIYYII', 'IYXIYYII', 'IXYIXXII', 'IYXIXXII', 'IYYIXYII', 'IYYIYXII', 'IXXIIXYI', 'IXXIIYXI', 'IXYIIYYI', 'IYXIIYYI', 'IXYIIXXI', 'IYXIIXXI', 'IYYIIXYI', 'IYYIIYXI', 'IXXIXZZY', 'IXXIYZZX', 'IXYIYZZY', 'IYXIYZZY', 'IXYIXZZX', 'IYXIXZZX', 'IYYIXZZY', 'IYYIYZZX', 'IXXIIIXY', 'IXXIIIYX', 'IXYIIIYY', 'IYXIIIYY', 'IXYIIIXX', 'IYXIIIXX', 'IYYIIIXY', 'IYYIIIYX', 'IIXXXYII', 'IIXXYXII', 'IIXYYYII', 'IIYXYYII', 'IIXYXXII', 'IIYXXXII', 'IIYYXYII', 'IIYYYXII', 'IIXXIXYI', 'IIXXIYXI', 'IIXYIYYI', 'IIYXIYYI', 'IIXYIXXI', 'IIYXIXXI', 'IIYYIXYI', 'IIYYIYXI', 'IIXXXZZY', 'IIXXYZZX', 'IIXYYZZY', 'IIYXYZZY', 'IIXYXZZX', 'IIYXXZZX', 'IIYYXZZY', 'IIYYYZZX', 'IIXXIIXY', 'IIXXIIYX', 'IIXYIIYY', 'IIYXIIYY', 'IIXYIIXX', 'IIYXIIXX', 'IIYYIIXY', 'IIYYIIYX', 'XZXIXZYI', 'XZXIYZXI', 'XZYIYZYI', 'YZXIYZYI', 'XZYIXZXI', 'YZXIXZXI', 'YZYIXZYI', 'YZYIYZXI', 'IXZXIXZY', 'IXZXIYZX', 'IXZYIYZY', 'IYZXIYZY', 'IXZYIXZX', 'IYZXIXZX', 'IYZYIXZY', 'IYZYIYZX']\n",
      "coef_single: [0.5, -0.5, 0.5, -0.5, 0.5, -0.5, 0.5, -0.5, 0.5, -0.5, 0.5, -0.5, 0.5, -0.5, 0.5, -0.5]\n",
      "coef_double: [0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125, 0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125, 0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125, 0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125, 0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125, 0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125, 0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125, 0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125, 0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125, 0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125, 0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125, 0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125, 0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125, 0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125, 0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125, 0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125, 0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125, 0.125, 0.125, 0.125, 0.125, -0.125, -0.125, -0.125, -0.125]\n"
     ]
    }
   ],
   "source": [
    "\n",
    "from qchem.uccsd import get_uccsd_op, uccsd_circuit, uccsd_parameter_size\n",
    "from qchem.uccsd import uccsd_circuit_double, uccsd_circuit_single\n",
    "\n",
    "spin_mult, only_singles, only_doubles = 0, False, False\n",
    "\n",
    "# Get the number of parameters\n",
    "singles, doubles, total = uccsd_parameter_size(electron_count, qubits_num, spin_mult)\n",
    "print(f\"Number of parameters: {singles} singles, {doubles} doubles, {total} total\")\n",
    "\n",
    "\n",
    "# Get the UCCSD pool\n",
    "if not only_singles and not only_doubles:\n",
    "    word_single, word_double, coef_single, coef_double = get_uccsd_op(electron_count, qubits_num, \n",
    "                                                                  spin_mult = 0, only_singles = False, only_doubles = False)\n",
    "    print(f\"word_single: {word_single}\")\n",
    "    print(f\"word_double: {word_double}\")\n",
    "    print(f\"coef_single: {coef_single}\")\n",
    "    print(f\"coef_double: {coef_double}\")\n",
    "\n",
    "elif only_singles and not only_doubles:\n",
    "    word_single, coef_single = get_uccsd_op(electron_count, qubits_num, spin_mult = 0, only_singles = True, only_doubles = False)\n",
    "    \n",
    "    print(f\"word_single: {word_single}\")\n",
    "    print(f\"coef_single: {coef_single}\")\n",
    "\n",
    "elif only_doubles and not only_singles:\n",
    "    word_double, coef_double = get_uccsd_op(electron_count, qubits_num, spin_mult = 0, only_singles = False, only_doubles = True)\n",
    "    \n",
    "    print(f\"word_double: {word_double}\")\n",
    "    print(f\"coef_double: {coef_double}\")\n",
    "else:\n",
    "    raise ValueError(\"Invalid option for only_singles and only_doubles\")\n",
    "\n",
    "\n",
    "# Get the UCCSD circuit (singles and doubles excitation are included)\n",
    "@cudaq.kernel\n",
    "def uccsd_kernel(qubits_num: int, electron_count: int, theta: list[float], \n",
    "                 word_single: list[cudaq.pauli_word], word_double: list[cudaq.pauli_word], coef_single: list[float], coef_double: list[float]):\n",
    "    \"\"\"\n",
    "    UCCSD kernel\n",
    "    \"\"\"\n",
    "    # Prepare the statefrom qchem.uccsd import get_uccsd_op, uccsd_circuit, uccsd_parameter_size\n",
    "\n",
    "    qubits = cudaq.qvector(qubits_num)\n",
    "\n",
    "    # Initialize the qubits\n",
    "    for i in range(electron_count):\n",
    "        x(qubits[i])\n",
    "    \n",
    "    # Apply the UCCSD circuit\n",
    "    uccsd_circuit(qubits, theta, word_single, coef_single, word_double, coef_double)\n",
    "\n",
    "\n",
    "# Get the UCCSD circuit (only doubles excitations are included)\n",
    "@cudaq.kernel\n",
    "def uccsd_double_kernel(qubits_num: int, electron_count: int, theta: list[float], \n",
    "                word_double: list[cudaq.pauli_word], coef_double: list[float]):\n",
    "    \"\"\"\n",
    "    UCCSD kernel\n",
    "    \"\"\"\n",
    "    # Prepare the state\n",
    "    qubits = cudaq.qvector(qubits_num)\n",
    "\n",
    "    # Initialize the qubits\n",
    "    for i in range(electron_count):\n",
    "        x(qubits[i])\n",
    "    \n",
    "    # Apply the UCCSD circuit\n",
    "    uccsd_circuit_double(qubits, theta, word_double, coef_double)\n",
    "\n",
    "# Get the UCCSD circuit (only singles excitations are included)\n",
    "@cudaq.kernel\n",
    "def uccsd_single_kernel(qubits_num: int, electron_count: int, theta: list[float], \n",
    "                word_single: list[cudaq.pauli_word], coef_single: list[float]):\n",
    "    \"\"\"\n",
    "    UCCSD kernel\n",
    "    \"\"\"\n",
    "    # Prepare the state\n",
    "    qubits = cudaq.qvector(qubits_num)\n",
    "\n",
    "    # Initialize the qubits\n",
    "    for i in range(electron_count):\n",
    "        x(qubits[i])\n",
    "    \n",
    "    # Apply the UCCSD circuit\n",
    "    uccsd_circuit_single(qubits, theta, word_single, coef_single)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fee19b5",
   "metadata": {},
   "source": [
    "### Run VQE\n",
    "\n",
    "To learn about VQE in CUDA-Q, check this [tutorial](https://nvidia.github.io/cuda-quantum/latest/applications/python/vqe_advanced.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c7436df",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total energy: -1852.5895309905 Hartree\n",
      "optimized parameters: [ 0.00000000e+00  1.91918143e-03 -1.18739821e-01 -4.56567309e-18\n",
      "  0.00000000e+00  1.93627077e-03 -1.18755472e-01  1.54888818e-21\n",
      "  9.53533311e-03  0.00000000e+00  0.00000000e+00  1.40615591e-03\n",
      "  0.00000000e+00 -1.02368050e-02  2.78804812e-03  0.00000000e+00\n",
      "  0.00000000e+00  2.78862284e-03 -1.01582650e-02  0.00000000e+00\n",
      "  3.21798016e-01 -9.94788673e-15  9.94560544e-15  1.54838228e-02\n",
      " -4.22692720e-03 -4.25296495e-03]\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from scipy.optimize import minimize\n",
    "\n",
    "# Initial guess for the parameters \n",
    "if not only_singles and not only_doubles: \n",
    "    theta = [0.0] * total\n",
    "elif only_singles and not only_doubles:\n",
    "    theta = [0.0] * singles\n",
    "elif only_doubles and not only_singles:\n",
    "    theta = [0.0] * doubles\n",
    "else:\n",
    "    raise ValueError(\"Invalid option for only_singles and only_doubles\")\n",
    "\n",
    "#print(cudaq.draw(kernel, qubits_num, electron_count, theta, word_single, word_double, coef_single, coef_double))\n",
    "\n",
    "# The ansatz is a variational circuit, so we need to optimize the parameters\n",
    "\n",
    "# Parameter shift to compute the gradient\n",
    "def parameter_shift(theta):\n",
    "            \n",
    "    parameter_count = len(theta)\n",
    "    epsilon = np.pi / 4\n",
    "    # The gradient is calculated using parameter shift.\n",
    "    grad = np.zeros(parameter_count)\n",
    "    theta2 = theta.copy()\n",
    "\n",
    "    for i in range(parameter_count):\n",
    "        theta2[i] = theta[i] + epsilon\n",
    "        exp_val_plus = cost(theta2)\n",
    "        theta2[i] = theta[i] - epsilon\n",
    "        exp_val_minus = cost(theta2)\n",
    "        grad[i] = (exp_val_plus - exp_val_minus) / (2 * epsilon)\n",
    "        theta2[i] = theta[i]\n",
    "    return grad\n",
    "\n",
    "def cost(theta):\n",
    "            \n",
    "    theta=theta.tolist()\n",
    "    \n",
    "    if not only_singles and not only_doubles:\n",
    "        energy = cudaq.observe(uccsd_kernel, hamiltonian, qubits_num, electron_count, \n",
    "                            theta, word_single, word_double, coef_single, coef_double).expectation()\n",
    "    \n",
    "    elif only_singles and not only_doubles:\n",
    "        energy = cudaq.observe(uccsd_single_kernel, hamiltonian, qubits_num, electron_count, \n",
    "                            theta, word_single, coef_single).expectation()\n",
    "    elif only_doubles and not only_singles:\n",
    "        energy = cudaq.observe(uccsd_double_kernel, hamiltonian, qubits_num, electron_count, \n",
    "                            theta, word_double, coef_double).expectation()\n",
    "    else:\n",
    "        raise ValueError(\"Invalid option for only_singles and only_doubles\")\n",
    "    \n",
    "    return energy\n",
    "\n",
    "#result_vqe=minimize(cost, theta, method='L-BFGS-B', jac='2-point', tol=1e-7)\n",
    "result_vqe=minimize(cost, theta, method='L-BFGS-B', jac=parameter_shift, tol=1e-7)\n",
    "\n",
    "total_energy = result_vqe.fun + const\n",
    "print(f\"Total energy: {total_energy:.10f} Hartree\")\n",
    "# Print the optimized parameters: first n are singles, then doubles.\n",
    "print(f\"optimized parameters: {result_vqe.x}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e20b6410",
   "metadata": {},
   "source": [
    "### Challenges and consideration\n",
    "\n",
    "While UCCSD is a powerful method, it is not without its challenges. The primary difficulty lies in the efficient implementation of the quantum circuits, as the number of gates required can grow significantly with the size of the system. Additionally, the accuracy of the Trotter-Suzuki approximation can affect the overall precision of the results."
   ]
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
