{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to CUDA-Q Dynamics (Time Dependent Hamiltonians)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Landau-Zener model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this lab we will learn how to implement time dependent Hamiltonian terms and custom operators.\n",
    "This is important functionality for modeling the application of control pulses to qubits or other interactions that maybe changing in time. \n",
    "\n",
    "In the first section we will use the example of a Landau-Zener transition to demonstrate two ways of implementing time dependent terms.\n",
    "\n",
    "The [Landau-Zener model](https://en.wikipedia.org/wiki/Landau%E2%80%93Zener_formula) describes the transition probability of a quantum system with two energy levels that cross each other as a function of an external parameter. For example, one can consider two spins weakly coupled to each other ($g$) while an external magnetic field is swept. If the two spins have an energy spectrum that cross each other as a function of the magnetic field, an avoided crossing is created due to the interaction of the spins with each other.\n",
    "\n",
    "If the magnetic field is swept slowly (adiabatically), no transitions across the splitting occurs. However, if it is swept quickly (diabatically), the transition hops over the splitting, effectively ignoring the splitting. If swept at an intermediate rate, the transition probability is given by a well defined probability. \n",
    "\n",
    "![alt text 2](images/LZ_transition.png)\n",
    "\n",
    "Figure 1. Sketch of two crossing energy levels with an avoided crossing along a parameter z. Image credit: Wikipedia"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 1 - Implementing time dependent terms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this lab we will learn to implement time dependent Hamiltonian terms with:\n",
    "- a time dependent scalar coefficient via `ScalarOperator`\n",
    "- a `tensor callback` function for custom operators\n",
    "\n",
    "We consider the following simple Hamiltonian representing the lowest two levels of the system. Note we consider first a small system as an instructional example to demonstrate the above features. A system this small won't experience a good acceleration run on GPU, but the features here can easily be scaled towards much larger systems which will experience a significant GPU speed-up.\n",
    "\n",
    "$$ \n",
    "\\hat H = \\begin{bmatrix}\n",
    "-\\alpha t & g \\\\\n",
    "g & \\alpha t\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "Here the coupling between levels is denoted by $g$ and the time dependent parameter is denoted by $\\alpha$. The above Hamiltonian can be of course decomposed into the following Pauli operators:\n",
    "\n",
    "$\\hat H = g\\hat\\sigma_x - \\alpha t \\hat\\sigma_z$\n",
    "\n",
    "We can implement the Pauli decomposition easily by constructing our Hamiltonian in CUDA-Q Dynamics with our typical spin operators and by leveraging a `ScalarOperator` with a time dependence and multiplying that by the $\\sigma_z$ term. The transition probability is known to be $P = e^{-\\pi g^2/\\alpha}$ but we will observe this by performing the following simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cudaq\n",
    "from cudaq import spin, operators, ScalarOperator, Schedule, ScipyZvodeIntegrator\n",
    "import numpy as np\n",
    "import cupy as cp\n",
    "import os\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the target to our dynamics simulator\n",
    "cudaq.set_target(\"dynamics\")\n",
    "\n",
    "# Define some shorthand operators\n",
    "sx = spin.x(0)\n",
    "sz = spin.z(0)\n",
    "sm = operators.annihilate(0)\n",
    "sm_dag = operators.create(0)\n",
    "\n",
    "# Dimensions of sub-system. We only have a single degree of freedom of dimension 2 (two-level system).\n",
    "dimensions = {0: 2}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's set a target transition probability of 0.75 to inform our value of $\\alpha$. We will see that computing the simulation will yield a transition result predicted by the theory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the Hamiltonian with spin operators and a ScalarOperator\n",
    "# Set the coupling to 2pi\n",
    "g = 2 * np.pi\n",
    "# The target ground state probability that we want to achieve\n",
    "target_p0 = 0.75\n",
    "# Compute `alpha` parameter:\n",
    "alpha = (-np.pi * g**2) / np.log(target_p0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define our Hamiltonian below and use the `ScalarOperator` to multiply the $\\sigma_z$ term by $t$ in a lambda function. This functionality also enables the easy implementation of time dependent drive terms and envelop functions by simply replacing the function in the `ScalarOperator` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hamiltonian\n",
    "hamiltonian = g * sx - alpha * ScalarOperator(lambda t: t) * sz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initial state of the system (ground state)\n",
    "psi0 = cudaq.State.from_data(cp.array([1.0, 0.0], dtype=cp.complex128))\n",
    "\n",
    "# Schedule of time steps (simulating a long time range)\n",
    "steps = np.linspace(-2.0, 2.0, 5000)\n",
    "schedule = Schedule(steps, [\"t\"])\n",
    "\n",
    "# Run the simulation.\n",
    "evolution_result = cudaq.evolve(hamiltonian,\n",
    "                                dimensions,\n",
    "                                schedule,\n",
    "                                psi0,\n",
    "                                observables=[operators.number(0)],\n",
    "                                collapse_operators=[],\n",
    "                                store_intermediate_results=cudaq.IntermediateResultSave.EXPECTATION_VALUE,\n",
    "                                integrator=ScipyZvodeIntegrator())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7b6ad27635c0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 62196 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Collect the results and plot\n",
    "prob1 = [\n",
    "    exp_vals[0].expectation()\n",
    "    for exp_vals in evolution_result.expectation_values()\n",
    "]\n",
    "\n",
    "prob0 = [1 - val for val in prob1]\n",
    "fig, ax = plt.subplots(figsize=(12, 8))\n",
    "ax.plot(steps, prob1, 'b', steps, prob0, 'r')\n",
    "ax.plot(steps, (1.0 - target_p0) * np.ones(np.shape(steps)), 'k')\n",
    "ax.plot(steps, target_p0 * np.ones(np.shape(steps)), 'm')\n",
    "ax.set_xlabel(\"Time\")\n",
    "ax.set_ylabel(\"Occupation probability\")\n",
    "ax.set_title(\"Landau-Zener transition\")\n",
    "ax.legend((\"Excited state\", \"Ground state\", \"LZ formula (Excited state)\",\n",
    "           \"LZ formula (Ground state)\"),\n",
    "          loc=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see the transition probability of the ground state converges to the predicted probability of 0.75."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 2 - Implementing custom operators"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the previous section we decomposed the Hamiltonian into spin operators to apply the time dependence to the $\\sigma_z$ terms. Alternatively, we can directly represent the Hamiltonian\n",
    "$$ \n",
    "\\hat H = \\begin{bmatrix}\n",
    "-\\alpha t & g \\\\\n",
    "g & \\alpha t\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "with a `tensor callback` and by defining an `ElementaryOperator`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define the above Hamiltonian in a matrix in the function `callback_tensor(t)` which accepts time as an argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cudaq import operators\n",
    "\n",
    "def callback_tensor(t):\n",
    "    return np.array([[-alpha * t, g], [g, alpha * t]], dtype=np.complex128)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now define an operator with this callback tensor, labeled `lz_op`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's define the control term as a callback tensor that acts on 2-level systems\n",
    "operators.define(\"lz_op\", [2], callback_tensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can define our hamiltonian now with this single `lz_op` `ElementaryOperator`, applied to the 0th degree of freedom. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hamiltonian\n",
    "hamiltonian = operators.instantiate(\"lz_op\", [0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can run the rest of the simulation as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dimensions of sub-system. We only have a single degree of freedom of dimension 2 (two-level system).\n",
    "dimensions = {0: 2}\n",
    "\n",
    "# Initial state of the system (ground state)\n",
    "psi0 = cudaq.State.from_data(cp.array([1.0, 0.0], dtype=cp.complex128))\n",
    "\n",
    "# Schedule of time steps (simulating a long time range)\n",
    "steps = np.linspace(-2.0, 2.0, 5000)\n",
    "schedule = Schedule(steps, [\"t\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the simulation.\n",
    "evolution_result = cudaq.evolve(hamiltonian,\n",
    "                                dimensions,\n",
    "                                schedule,\n",
    "                                psi0,\n",
    "                                observables=[operators.number(0)],\n",
    "                                collapse_operators=[],\n",
    "                                store_intermediate_results=cudaq.IntermediateResultSave.EXPECTATION_VALUE,\n",
    "                                integrator=ScipyZvodeIntegrator())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7b6ad1fce000>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 62196 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "prob1 = [\n",
    "    exp_vals[0].expectation()\n",
    "    for exp_vals in evolution_result.expectation_values()\n",
    "]\n",
    "\n",
    "prob0 = [1 - val for val in prob1]\n",
    "fig, ax = plt.subplots(figsize=(12, 8))\n",
    "ax.plot(steps, prob1, 'b', steps, prob0, 'r')\n",
    "ax.plot(steps, (1.0 - target_p0) * np.ones(np.shape(steps)), 'k')\n",
    "ax.plot(steps, target_p0 * np.ones(np.shape(steps)), 'm')\n",
    "ax.set_xlabel(\"Time\")\n",
    "ax.set_ylabel(\"Occupation probability\")\n",
    "ax.set_title(\"Landau-Zener transition\")\n",
    "ax.legend((\"Excited state\", \"Ground state\", \"LZ formula (Excited state)\",\n",
    "           \"LZ formula (Ground state)\"),\n",
    "          loc=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 3 - Heisenberg Model with a time-varying magnetic field"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section  we will now look at a much larger quantum system, considering a 1-D chain of 10 spins according to the Heisenberg Hamiltonian, adding a time-dependent transverse magnetic field after a delay. We can use this system to model quantum quench dynamics.\n",
    "\n",
    "$\\hat H = J_{ij}\\sum S_i \\cdot S_j + B_x (t) \\sum \\sigma^x_i $\n",
    "\n",
    " For this example, we consider the case where the $\\sigma^z_i \\sigma^z_{i+1}$ coupling is dominant, and prepare the system in an alternating spin up and down state, the approximate ground state.\n",
    "\n",
    "  ![alt text 1](images/Spin_chain.png)\n",
    "\n",
    "image credit: Wikipedia\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the target to our dynamics simulator\n",
    "cudaq.set_target(\"dynamics\")\n",
    "\n",
    "# In this example, we solve the Quantum Heisenberg model (https://en.wikipedia.org/wiki/Quantum_Heisenberg_model),\n",
    "# which exhibits the so-called quantum quench effect.\n",
    "\n",
    "# Number of spins\n",
    "N = 10\n",
    "\n",
    "dimensions = {}\n",
    "for i in range(N):\n",
    "    dimensions[i] = 2\n",
    "\n",
    "# Initial state: alternating spin up and down\n",
    "spin_state = ''\n",
    "for i in range(N):\n",
    "    spin_state += str(int(i % 2))\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 1 - Define a time-varying magnetic field\n",
    "\n",
    "Let us construct the Hamiltonian with a step function B-field that turns on with magnitude 2.0 after a time delay=2.5. Modify the `B_field` function below to reflect this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a time dependent B field\n",
    "delay=2.5\n",
    "def B_field(time):\n",
    "    if abs(time)<delay:\n",
    "        B=0.0\n",
    "    else:\n",
    "        B=2.0\n",
    "        # B=2.0*np.sin(2*np.pi*time)\n",
    "    return B"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now add the magnetic field term to the Hamiltonian. Assume the B-field lies in the x direction. The spin-spin Hamiltonian terms are given already."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Heisenberg model spin coupling strength\n",
    "Jx = -0.2\n",
    "Jy = -0.2\n",
    "Jz = -1.0\n",
    "\n",
    "# Construct the Hamiltonian\n",
    "H = operators.zero()\n",
    "\n",
    "# Spin-spin terms\n",
    "for i in range(N - 1):\n",
    "    H += Jx * spin.x(i) * spin.x(i + 1)\n",
    "    H += Jy * spin.y(i) * spin.y(i + 1)\n",
    "    H += Jz * spin.z(i) * spin.z(i + 1)\n",
    "\n",
    "# Magnetic Field terms\n",
    "for i in range(N):\n",
    "    H += ScalarOperator(B_field) * spin.x(i) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now execute the simulation, measuring the $\\sigma^z_i$ expectation. Other common observables for this system are spin correlators $\\langle \\sigma^z_i \\sigma^z_j \\rangle$ and the magnetization $M = \\frac{1}{N} \\sum \\langle\\sigma^z_i\\rangle$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the time schedule\n",
    "steps = np.linspace(0.0, 5, 1000)\n",
    "schedule = Schedule(steps, [\"time\"])\n",
    "\n",
    "# Prepare the initial state vector\n",
    "psi0_ = cp.zeros(2**N, dtype=cp.complex128)\n",
    "psi0_[int(spin_state, 2)] = 1.0\n",
    "psi0 = cudaq.State.from_data(psi0_)\n",
    "\n",
    "# Run the simulation\n",
    "evolution_result = cudaq.evolve(H,\n",
    "                                dimensions,\n",
    "                                schedule,\n",
    "                                psi0,\n",
    "                                observables=[spin.z(i) for i in range(N)],\n",
    "                                collapse_operators=[],\n",
    "                                store_intermediate_results=cudaq.IntermediateResultSave.EXPECTATION_VALUE,\n",
    "                                integrator=ScipyZvodeIntegrator())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0, 'Time')"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 138120 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the results\n",
    "get_result = lambda idx, res: [\n",
    "    exp_vals[idx].expectation() for exp_vals in res.expectation_values()\n",
    "]\n",
    "\n",
    "exp_val = [get_result(i, evolution_result) for i in range(N)]\n",
    "\n",
    "# Plot the results\n",
    "fig = plt.figure(figsize=(12, 6))\n",
    "for i in range(N):  \n",
    "    plt.plot(steps, exp_val[i],label=f'spin {i}')\n",
    "plt.legend(fontsize=12)\n",
    "plt.ylabel(\"$\\\\sigma^z $\")\n",
    "plt.xlabel(\"Time\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the prior to the addition of the magnetic field, the system stays relatively stationary. With the transverse field, the system quickly mixes. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 2 (Optional) \n",
    "\n",
    "a. Adjust the time dependent magnetic field to now be a sine function with amplitude 2.0, and frequency $2\\pi$.\n",
    "\n",
    "b. Consider now a disordered spin chain with varying couplings\n",
    "\n",
    "c. Observe the behavior of the following observables: spin correlators $\\langle \\sigma^z_i \\sigma^z_j \\rangle$ and the magnetization $M = \\frac{1}{N} \\sum \\langle\\sigma^z_i\\rangle$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a time dependent B field\n",
    "delay=2.5\n",
    "def B_field(time):\n",
    "    if time<delay:\n",
    "        B=0\n",
    "    else:\n",
    "        B=2.0\n",
    "        # B=2.0*np.sin(2*np.pi*time)\n",
    "    return B\n",
    "\n",
    "observables=[spin.z(0)*spin.z(1)]"
   ]
  }
 ],
 "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": 4
}
