{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "introduction",
   "metadata": {},
   "source": [
    "# Spin-Hamiltonian Simulation Using CUDA-Q\n",
    "\n",
    "This tutorial demonstrates how to perform one-dimensional spin-Hamiltonian simulations. We will focus on simulating the time evolution of spin chains with the Heisenberg and Transverse Field Ising Model (TFIM) Hamiltonians using Trotter-Suzuki decomposition.\n",
    "\n",
    "## Introduction\n",
    "\n",
    "Simulating quantum systems is a key problem to solve in quantum computing. The Heisenberg and TFIM Hamiltonians are models that describe interactions in spin chains, which are essential in understanding, as an example, magnetic properties of materials and quantum phase transitions.\n",
    "\n",
    "### Heisenberg Hamiltonian\n",
    "\n",
    "The Heisenberg Hamiltonian for a one-dimensional spin chain with nearest-neighbor interactions is given by:\n",
    "\n",
    "$$\n",
    "H = \\sum_{i} \\left( J_x S_i^x S_{i+1}^x + J_y S_i^y S_{i+1}^y + J_z S_i^z S_{i+1}^z \\right)\n",
    "$$\n",
    "\n",
    "- $S_i^\\alpha$ are the spin operators (Pauli matrices) acting on spin $i$ in the $ \\alpha$-direction ( $\\alpha = x, y, z$).\n",
    "- $J_x, J_y, J_z$ are the coupling coefficients between spins.\n",
    "\n",
    "---\n",
    "\n",
    "### Transverse Field Ising Model (TFIM)\n",
    "\n",
    "The TFIM Hamiltonian is a simplified model focusing on the competition between an external transverse magnetic field and nearest-neighbor spin interactions:\n",
    "\n",
    "$$\n",
    "H = -h \\sum_i S_i^x - \\sum_i S_i^z S_{i+1}^z\n",
    "$$\n",
    "\n",
    "- $h$ is the strength of the transverse magnetic field.\n",
    "- The first term represents the interaction of spins with the external field in the $x$-direction.\n",
    "- The second term represents the interaction between neighboring spins in the $z$-direction.\n",
    "\n",
    "---\n",
    "\n",
    "### Time Evolution and Trotter Decomposition\n",
    "\n",
    "To simulate the time evolution $e^{-iHt}$ of these Hamiltonians, we use the Trotter-Suzuki decomposition, which approximates the exponential of a sum of non-commuting operators by a product of exponentials of individual terms:\n",
    "\n",
    "$$\n",
    "e^{-iHt} \\approx \\left( e^{-iH_1 \\Delta t} e^{-iH_2 \\Delta t} \\cdots e^{-iH_n \\Delta t} \\right)^K\n",
    "$$\n",
    "\n",
    "- $\\Delta t = t / K$ is the time step.\n",
    "- $K$ is the number of Trotter steps.\n",
    "- $H_i$ are the individual terms in the Hamiltonian.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "import-libraries",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import Required Libraries\n",
    "import cudaq\n",
    "import numpy as np\n",
    "import time\n",
    "import sys\n",
    "from cudaq import spin\n",
    "import matplotlib.pyplot as plt\n",
    "from typing import List"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "function-definitions",
   "metadata": {},
   "source": [
    "## Key steps\n",
    "\n",
    "### 1. Prepare initial state\n",
    "\n",
    "Prepare the initial quantum state of the system. For this simulation, we create the state $|0101\\ldots\\rangle$, where every odd qubit is in the $|1\\rangle$ state, and every even qubit is in the $|0\\rangle$ state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "get-initial-state",
   "metadata": {},
   "outputs": [],
   "source": [
    "@cudaq.kernel\n",
    "def get_initial_state(n_spins: int):\n",
    "    \"\"\"Create initial state |1010...>\"\"\"\n",
    "    qubits = cudaq.qvector(n_spins)\n",
    "    for i in range(0, n_spins, 2):\n",
    "        x(qubits[i])\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "trotter-step",
   "metadata": {},
   "source": [
    "### 2. Hamiltonian Trotterization\n",
    "\n",
    "This function performs a single Trotter step for simulating a spin chain. If `_use_XXYYZZ_gate` is True, it constructs specific two-qubit gates using decomposition for exponentiation. It is a custom and efficient exponentiation combination of XX+YY+ZZ gates as a single operation. Otherwise, it uses a standrd CUDA-Q exp_pauli operation for the exponentiation of Pauli gates in Hamiltonian.\n",
    "Input parameters:\n",
    "- **state:** The current quantum state.  \n",
    "- **dt:** Time step for Trotter evolution.  \n",
    "- **Jx, Jy, Jz:** Coupling constants for the spin-spin interactions along the x, y, and z axes.  \n",
    "- **h_x, h_y, h_z:** Local magnetic field strengths (unused in the code provided).  \n",
    "- **`_use_XXYYZZ_gate`:** Flag to determine if custom XX+YY+ZZ gates should be used.  \n",
    "- **coefficients:** Coefficients of Pauli terms for exponential evolution.  \n",
    "- **words:** Corresponding Pauli operators (e.g., X, Y, Z) for evolution.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "trotter-step-code",
   "metadata": {},
   "outputs": [],
   "source": [
    "@cudaq.kernel\n",
    "def trotter_step(state: cudaq.State, dt: float, Jx: float, Jy: float, Jz: float,\n",
    "                 h_x: list[float], h_y: list[float], h_z: list[float], _use_XXYYZZ_gate: bool,\n",
    "                 coefficients: List[complex], words: List[cudaq.pauli_word]):\n",
    "    \"\"\"Perform single Trotter step\"\"\"\n",
    "    qubits = cudaq.qvector(state)\n",
    "    n_spins = len(qubits)\n",
    "   \n",
    "    # Apply two-qubit interaction terms\n",
    "    if _use_XXYYZZ_gate:\n",
    "        for j in range(2):\n",
    "            for i in range(j % 2, n_spins - 1, 2):\n",
    "                rx(-np.pi/2,qubits[i])\n",
    "                rx(np.pi/2,qubits[i+1])\n",
    "                x.ctrl(qubits[i], qubits[i+1])\n",
    "                h(qubits[i])\n",
    "                s(qubits[i])\n",
    "                rz(-2*Jy*dt,qubits[i+1])\n",
    "                x.ctrl(qubits[i], qubits[i+1])\n",
    "                h(qubits[i])\n",
    "                rx(2*Jx*dt,qubits[i])\n",
    "                rz(-2*Jz*dt,qubits[i+1])\n",
    "                x.ctrl(qubits[i], qubits[i+1])\n",
    "    else:\n",
    "        for i in range(len(coefficients)):\n",
    "            exp_pauli(coefficients[i].real * dt, qubits, words[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "compute-overlap-probability",
   "metadata": {},
   "source": [
    "### 3. `Compute overlap`\n",
    "\n",
    "Calculate the probability of the evolved state overlapping with the initial state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "compute-overlap-code",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_overlap_probability(initial_state: cudaq.State, evolved_state: cudaq.State):\n",
    "    \"\"\"Compute probability of the overlap with the initial state\"\"\"\n",
    "    overlap = initial_state.overlap(evolved_state)\n",
    "    return np.abs(overlap)**2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "create-heisenberg-hamiltonian",
   "metadata": {},
   "source": [
    "### 4. Construct Heisenberg Hamiltonian\n",
    "\n",
    "Construct the Heisenberg Hamiltonian as a `cudaq.SpinOperator` object, considering nearest-neighbor interactions along X, Y, and Z directions. It supports arbitrary coupling constants Jx, Jy, and Jz. See formula above.\n",
    "Input:\n",
    "- **`n_spins`**: Number of spins (qubits) in the system.  \n",
    "- **`Jx`, `Jy`, `Jz`**: Coupling constants for interactions in the X, Y, and Z directions.  \n",
    "- **`h_x`, `h_y`, `h_z`**: Magnetic field components (currently unused).  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "heisenberg-code",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_hamiltonian_heisenberg(n_spins: int, Jx: float, Jy: float, Jz: float, h_x: list[float], h_y: list[float], h_z: list[float]):\n",
    "    \"\"\"Create the Hamiltonian operator\"\"\"\n",
    "    ham = 0\n",
    "\n",
    "    # Add two-qubit interaction terms for Heisenberg Hamiltonian\n",
    "    for i in range(0, n_spins - 1):\n",
    "        ham += Jx * spin.x(i) * spin.x(i + 1)\n",
    "        ham += Jy * spin.y(i) * spin.y(i + 1)\n",
    "        ham += Jz * spin.z(i) * spin.z(i + 1)\n",
    "   \n",
    "    return ham"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "create-tfim-hamiltonian",
   "metadata": {},
   "source": [
    "### 5. Construct TFIM Hamiltonian\n",
    "\n",
    "Construct the TFIM Hamiltonian as a `cudaq.SpinOperator` object. See formula above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "tfim-code",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_hamiltonian_tfim(n_spins: int, h_field: float):\n",
    "    \"\"\"Create the Hamiltonian operator\"\"\"\n",
    "    ham = 0\n",
    "   \n",
    "    # Add single-qubit terms\n",
    "    for i in range(0, n_spins):\n",
    "        ham += -1 * h_field * spin.x(i)\n",
    "\n",
    "    # Add two-qubit interaction terms for Ising Hamiltonian\n",
    "    for i in range(0, n_spins-1):\n",
    "        ham += -1 * spin.z(i) * spin.z(i + 1)\n",
    "   \n",
    "    return ham"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "extract-coefficients-and-words",
   "metadata": {},
   "source": [
    "### 6. Extract coefficients and Pauli words\n",
    "\n",
    "Extract the coefficients and Pauli words from the provided Hamiltonian for use in the Trotter step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "extract-code",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extractCoefficients(hamiltonian: cudaq.SpinOperator):\n",
    "    result = []\n",
    "    for term in hamiltonian:\n",
    "        result.append(term.evaluate_coefficient())\n",
    "    return result\n",
    "\n",
    "def extractWords(hamiltonian: cudaq.SpinOperator):\n",
    "    # Our kernel uses these words to apply exp_pauli to the entire state.\n",
    "    # we hence ensure that each pauli word covers the entire space.\n",
    "    n_spins = hamiltonian.qubit_count\n",
    "    result = []\n",
    "    for term in hamiltonian:\n",
    "        result.append(term.get_pauli_word(n_spins))\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "main-simulation",
   "metadata": {},
   "source": [
    "## Main code\n",
    "\n",
    "Import required libraries, set up the simulation parameters, and perform the time evolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "main-code",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initialize state\n",
      "Circuit execution time: 2.244 seconds\n"
     ]
    }
   ],
   "source": [
    "# Import Required Libraries\n",
    "import cudaq\n",
    "import numpy as np\n",
    "import time\n",
    "import sys\n",
    "from cudaq import spin\n",
    "import matplotlib.pyplot as plt\n",
    "from typing import List\n",
    "\n",
    "# Parameters\n",
    "n_spins = 4  # Number of spins in the chain\n",
    "ham_type = \"heisenberg\"  # Choose between \"heisenberg\" and \"tfim\"\n",
    "Jx, Jy, Jz = 1.0, 1.0, 1.0  # Coupling coefficients for Heisenberg Hamiltonian\n",
    "h_field = 1.0  # Transverse field strength for TFIM\n",
    "K = 100  # Number of Trotter steps\n",
    "t = np.pi  # Total evolution time\n",
    "dt = t / K  # Time step size\n",
    "\n",
    "# Optimized XXYYZZ exponentiation. Works only for Heisenberg Hamiltonian\n",
    "_use_XXYYZZ_gate = False\n",
    "if _use_XXYYZZ_gate == True and ham_type == \"tfim\":\n",
    "    print (\"XXYYZZ exponentiation works only for Heisenberg\")\n",
    "    sys.exit(0)\n",
    "\n",
    "# Create Hamiltonian\n",
    "if ham_type == \"heisenberg\":\n",
    "    # Initialize field for Heisenberg Hamiltonian\n",
    "    h_x = np.ones(n_spins)\n",
    "    h_y = np.ones(n_spins)\n",
    "    h_z = np.ones(n_spins)\n",
    "    hamiltonian = create_hamiltonian_heisenberg(n_spins, Jx, Jy, Jz, h_x, h_y,h_z)\n",
    "elif ham_type == \"tfim\":\n",
    "    hamiltonian = create_hamiltonian_tfim(n_spins, h_field)\n",
    "else:\n",
    "    raise ValueError(\"Invalid Hamiltonian type. Choose 'heisenberg' or 'tfim'.\")\n",
    "\n",
    "# Extract coefficients and words\n",
    "coefficients = extractCoefficients(hamiltonian)\n",
    "words = extractWords(hamiltonian)\n",
    "\n",
    "# Initialize and save the initial state\n",
    "print (\"Initialize state\")\n",
    "initial_state = cudaq.get_state(get_initial_state, n_spins)\n",
    "state = initial_state\n",
    "\n",
    "# Store probabilities over time\n",
    "probabilities = []\n",
    "probabilities.append(1.0)\n",
    "\n",
    "# Time evolution\n",
    "start_time = time.time()\n",
    "for k in range(1, K):\n",
    "\n",
    "    # Apply single Trotter step\n",
    "    state = cudaq.get_state(trotter_step, state, dt, Jx, Jy, Jz, h_x, h_y, h_z, _use_XXYYZZ_gate, coefficients, words)\n",
    "\n",
    "    # Calculate probability between initial and current states\n",
    "    probability = compute_overlap_probability(initial_state, state)\n",
    "    probabilities.append(probability)\n",
    "\n",
    "total_time = time.time() - start_time\n",
    "print(f\"Circuit execution time: {total_time:.3f} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "results-and-visualization",
   "metadata": {},
   "source": [
    "## Visualization of probablity over time\n",
    "\n",
    "Plot the probability of the system remaining in the initial state as a function of time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "plot-code",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 43704 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot probability over time\n",
    "plt.plot(range(K), probabilities, marker='o')\n",
    "plt.xlabel('Step (k)')\n",
    "plt.ylabel('Probability of initial state')\n",
    "plt.title('Probability vs. Step in Trotter Evolution')\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da2c333a",
   "metadata": {},
   "source": [
    "## Expectation value over time: \n",
    "\n",
    "Set up the simulation parameters, perform the time evolution, and visualize results for expectation value over time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9f09a43e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize list to store expectation values\n",
    "exp_values = []\n",
    "\n",
    "# Time evolution\n",
    "start_time = time.time()\n",
    "for k in range(1, K):\n",
    "    # Apply single Trotter step\n",
    "    state = cudaq.get_state(trotter_step, state, dt, Jx, Jy, Jz, h_x, h_y, h_z, _use_XXYYZZ_gate, coefficients, words)\n",
    "\n",
    "    # Calculate expectation value\n",
    "    exp_val = cudaq.observe(trotter_step, hamiltonian, state, dt, Jx, Jy, Jz, h_x, h_y, h_z, _use_XXYYZZ_gate, coefficients, words).expectation()\n",
    "    exp_values.append(exp_val.real)\n",
    "    #print(f\"Step {k}, Energy: {exp_val.real:.6f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3621ecdf",
   "metadata": {},
   "source": [
    "## Visualization of expectation over time\n",
    "\n",
    "Plot the expectation value over time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f0941a30",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 54580 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot expectation value over time\n",
    "x = np.arange(1, K)\n",
    "plt.plot(x, exp_values, marker='o')\n",
    "plt.xlabel('Step (k)')\n",
    "plt.ylabel('Expectation Value (Energy)')\n",
    "plt.title('Expectation Value vs. Step in Trotter Evolution')\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f5c059c",
   "metadata": {},
   "source": [
    "## Additional information\n",
    "\n",
    "In the code above, we relied on CUDA-Q expectation of Pauli gates. If we needed to do it manually here is how it is done."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "99b2f46e",
   "metadata": {},
   "outputs": [],
   "source": [
    "@cudaq.kernel\n",
    "def xx_gate(qubits: cudaq.qview, tau: float) -> None:\n",
    "    \"\"\"XX gate implementation\"\"\"\n",
    "    h(qubits[0])\n",
    "    h(qubits[1])\n",
    "    x.ctrl(qubits[0], qubits[1])\n",
    "    rz(np.pi * tau, qubits[1])\n",
    "    x.ctrl(qubits[0], qubits[1])\n",
    "    h(qubits[0])\n",
    "    h(qubits[1])\n",
    "\n",
    "@cudaq.kernel\n",
    "def yy_gate(qubits: cudaq.qview, tau: float) -> None:\n",
    "    \"\"\"YY gate implementation\"\"\"\n",
    "    # S gates (equivalent to rz(pi/2))\n",
    "    rz(np.pi/2, qubits[0])\n",
    "    rz(np.pi/2, qubits[1])\n",
    "    \n",
    "    h(qubits[0])\n",
    "    h(qubits[1])\n",
    "    x.ctrl(qubits[0], qubits[1])\n",
    "    rz(np.pi * tau, qubits[1])\n",
    "    x.ctrl(qubits[0], qubits[1])\n",
    "    h(qubits[0])\n",
    "    h(qubits[1])\n",
    "    \n",
    "    # S inverse gates (equivalent to rz(-pi/2))\n",
    "    rz(-np.pi/2, qubits[0])\n",
    "    rz(-np.pi/2, qubits[1])\n",
    "\n",
    "@cudaq.kernel\n",
    "def zz_gate(qubits: cudaq.qview, tau: float) -> None:\n",
    "    \"\"\"ZZ gate implementation\"\"\"\n",
    "    x.ctrl(qubits[0], qubits[1])\n",
    "    rz(np.pi * tau, qubits[1])\n",
    "    x.ctrl(qubits[0], qubits[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdea4560",
   "metadata": {},
   "source": [
    "## Relevant references\n",
    "\n",
    "- https://nvidia.github.io/cuda-quantum/latest/examples/python/operators.html\n",
    "- https://arxiv.org/abs/quant-ph/9503016"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "22ecee93",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CUDA-Q Version proto-0.8.0 (https://github.com/NVIDIA/cuda-quantum d5e0513d54809a835d1c2c108c0692be10d7d1bb)\n"
     ]
    }
   ],
   "source": [
    "print(cudaq.__version__)"
   ]
  }
 ],
 "metadata": {
  "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
