{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "64448bd9-7802-43aa-b8cd-8265dc474c24",
   "metadata": {},
   "source": [
    "# Control\n",
    "\n",
    "These examples all cover dynamical simulation of systems related to device control. These examples use small systems, but note that the GPU will only provide an advantage for systems with total dimension of O(1000). "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccb79b71-dbbc-46e5-b3e7-1c7cd46845dc",
   "metadata": {},
   "source": [
    "## Gate Calibration\n",
    "\n",
    "This example demonstrates how to perform pulse optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7a9901f3-9572-4b92-95a2-d965c7e4d379",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gaussian amplitude = 0.0, derivative amplitude = 0.0, Overlap: 0.7071067811865475\n",
      "Gaussian amplitude = 157.07963267948966, derivative amplitude = 0.0, Overlap: 0.7678402941325417\n",
      "Gaussian amplitude = 157.07963267948966, derivative amplitude = 157.07963267948966, Overlap: 0.6797941803182259\n",
      "Gaussian amplitude = 0.0, derivative amplitude = -157.07963267948966, Overlap: 0.9296240579072232\n",
      "Gaussian amplitude = -78.53981633974483, derivative amplitude = -314.1592653589793, Overlap: 0.7162211356808167\n",
      "Gaussian amplitude = 157.07963267948966, derivative amplitude = -157.07963267948966, Overlap: 0.44596190760545296\n",
      "Gaussian amplitude = 39.269908169872416, derivative amplitude = -39.269908169872416, Overlap: 0.9398996786747706\n",
      "Gaussian amplitude = -117.80972450961724, derivative amplitude = -196.34954084936206, Overlap: 0.5878521963356885\n",
      "Gaussian amplitude = 88.35729338221293, derivative amplitude = -49.087385212340514, Overlap: 0.6575502829271433\n",
      "Gaussian amplitude = 98.17477042468103, derivative amplitude = -19.634954084936208, Overlap: 0.33247544702132076\n",
      "Gaussian amplitude = 19.634954084936208, derivative amplitude = -98.17477042468103, Overlap: 0.8409827814368478\n",
      "Gaussian amplitude = -39.269908169872416, derivative amplitude = -117.80972450961724, Overlap: 0.9445817360291182\n",
      "Gaussian amplitude = -107.99224746714913, derivative amplitude = -166.89710972195778, Overlap: 0.7586919940233878\n",
      "Gaussian amplitude = -19.634954084936208, derivative amplitude = -58.90486225480862, Overlap: 0.7406581261969767\n",
      "Gaussian amplitude = 9.817477042468104, derivative amplitude = -88.35729338221293, Overlap: 0.8651921848480539\n",
      "Gaussian amplitude = -9.817477042468104, derivative amplitude = -68.72233929727672, Overlap: 0.8181626678528717\n",
      "Gaussian amplitude = 4.908738521234052, derivative amplitude = -83.44855486097889, Overlap: 0.8663447451357891\n",
      "Gaussian amplitude = -4.908738521234052, derivative amplitude = -73.63107781851078, Overlap: 0.8433090770060998\n",
      "Gaussian amplitude = 2.454369260617026, derivative amplitude = -80.99418560036186, Overlap: 0.8638676425264983\n",
      "Gaussian amplitude = 0.0, derivative amplitude = -78.53981633974483, Overlap: 0.8592468550642993\n",
      "Gaussian amplitude = -17.18058482431918, derivative amplitude = -100.62913968529807, Overlap: 0.9218675201964192\n",
      "Gaussian amplitude = -56.45049299419159, derivative amplitude = -139.8990478551705, Overlap: 0.9265972065467258\n",
      "Gaussian amplitude = -78.53981633974483, derivative amplitude = -157.07963267948966, Overlap: 0.835400572144582\n",
      "Gaussian amplitude = -32.520392703175595, derivative amplitude = -114.74176293384598, Overlap: 0.9480750666520233\n",
      "Gaussian amplitude = -15.339807878856416, derivative amplitude = -92.65243958829275, Overlap: 0.9007114410060454\n",
      "Gaussian amplitude = -46.1728217153578, derivative amplitude = -128.08739578845106, Overlap: 0.9490749424029333\n",
      "Gaussian amplitude = -39.42330624866097, derivative amplitude = -125.0194342126798, Overlap: 0.9590912872207543\n",
      "Gaussian amplitude = -39.50000528805525, derivative amplitude = -128.6242890642111, Overlap: 0.9644209828123245\n",
      "Gaussian amplitude = -53.15243430023746, derivative amplitude = -141.96992191881617, Overlap: 0.9371109174510482\n",
      "Gaussian amplitude = -37.67840310244106, derivative amplitude = -121.54880268008851, Overlap: 0.9556739509139307\n",
      "Gaussian amplitude = -31.005586675138513, derivative amplitude = -122.08569595584855, Overlap: 0.9654586085651646\n",
      "Gaussian amplitude = -23.421969155028872, derivative amplitude = -119.0848460395473, Overlap: 0.963623453294681\n",
      "Gaussian amplitude = -32.82718886075271, derivative amplitude = -129.1611823399711, Overlap: 0.9745284542544127\n",
      "Gaussian amplitude = -30.401581739908533, derivative amplitude = -132.9673721699124, Overlap: 0.9801467465441793\n",
      "Gaussian amplitude = -21.9071631269918, derivative amplitude = -126.42877906154987, Overlap: 0.9754292788007772\n",
      "Gaussian amplitude = -21.303158191761824, derivative amplitude = -137.31045527561375, Overlap: 0.9831092973124017\n",
      "Gaussian amplitude = -16.451943950073478, derivative amplitude = -144.92283493549633, Overlap: 0.9767359198754701\n",
      "Gaussian amplitude = -29.79757680467855, derivative amplitude = -143.84904838397628, Overlap: 0.9800104018090544\n",
      "Gaussian amplitude = -27.82497338525686, derivative amplitude = -139.49398105336968, Overlap: 0.983130133355889\n",
      "Gaussian amplitude = -18.726549837110152, derivative amplitude = -143.83706415907102, Overlap: 0.9795024686649386\n",
      "Gaussian amplitude = -27.48282376420894, derivative amplitude = -135.68479516720205, Overlap: 0.9830113579207141\n",
      "Gaussian amplitude = -21.645307812809744, derivative amplitude = -141.11964116178137, Overlap: 0.9826376028076026\n",
      "Gaussian amplitude = -26.02344477635914, derivative amplitude = -137.0435066658469, Overlap: 0.9836739278237355\n",
      "Gaussian amplitude = -32.545259969854186, derivative amplitude = -139.2270324436028, Overlap: 0.9801447148031301\n",
      "Gaussian amplitude = -24.113683636284915, derivative amplitude = -137.78959956761102, Overlap: 0.9838140571803666\n",
      "Gaussian amplitude = -22.31215502738719, derivative amplitude = -135.33912518008825, Overlap: 0.9830417565995295\n",
      "Gaussian amplitude = -26.446768795789446, derivative amplitude = -138.45526708504931, Overlap: 0.9836520843944795\n",
      "Gaussian amplitude = -23.690359616854607, derivative amplitude = -136.3778391484086, Overlap: 0.9836030636444814\n",
      "Gaussian amplitude = -25.75766650105573, derivative amplitude = -137.93591010088915, Overlap: 0.9837756871844154\n",
      "Gaussian amplitude = -23.847905360981503, derivative amplitude = -138.68200300265326, Overlap: 0.9837642033752482\n",
      "Gaussian amplitude = -24.39179021482591, derivative amplitude = -138.27237891845166, Overlap: 0.9838287992237971\n",
      "Gaussian amplitude = -22.747807350055094, derivative amplitude = -138.12606838517354, Overlap: 0.983605007962648\n",
      "Gaussian amplitude = -25.00520171330557, derivative amplitude = -137.98344967196024, Overlap: 0.9838349659123633\n",
      "Gaussian amplitude = -25.283308291846573, derivative amplitude = -138.46622902280092, Overlap: 0.9838084449662036\n",
      "Gaussian amplitude = -24.40608980017533, derivative amplitude = -137.9587569314085, Overlap: 0.9838333961345107\n",
      "Gaussian amplitude = -25.01950129865499, derivative amplitude = -137.6698276849171, Overlap: 0.9838262991248778\n",
      "Gaussian amplitude = -24.54871798578318, derivative amplitude = -138.12174111006803, Overlap: 0.9838371837908917\n",
      "Gaussian amplitude = -25.14782989891342, derivative amplitude = -138.14643385061981, Overlap: 0.9838281816540192\n",
      "Gaussian amplitude = -24.59152482485985, derivative amplitude = -138.0056761612113, Overlap: 0.9838386325114293\n",
      "Gaussian amplitude = -24.135041097337464, derivative amplitude = -138.1439675993191, Overlap: 0.9838178216495735\n",
      "Gaussian amplitude = -24.78766155931354, derivative amplitude = -138.02357915379997, Overlap: 0.9838394947681391\n",
      "Gaussian amplitude = -24.83046839839021, derivative amplitude = -137.90751420494328, Overlap: 0.9838382084236699\n",
      "Gaussian amplitude = -24.760030795238453, derivative amplitude = -137.96107093122444, Overlap: 0.983839415811919\n",
      "Gaussian amplitude = -24.95616752969215, derivative amplitude = -137.97897392381313, Overlap: 0.9838364330493514\n",
      "Gaussian amplitude = -24.682685501067926, derivative amplitude = -137.99900060186175, Overlap: 0.9838395809932626\n",
      "Gaussian amplitude = -24.710316265143017, derivative amplitude = -138.06150882443728, Overlap: 0.9838396318122055\n",
      "Gaussian amplitude = -24.6854590000953, derivative amplitude = -138.1117277710437, Overlap: 0.9838391065292872\n",
      "Gaussian amplitude = -24.6053402068974, derivative amplitude = -138.03693027249906, Overlap: 0.9838388699265836\n",
      "Gaussian amplitude = -24.742081221209506, derivative amplitude = -138.02691693347472, Overlap: 0.983839704126317\n",
      "Gaussian amplitude = -24.769711985284598, derivative amplitude = -138.08942515605025, Overlap: 0.9838393352545376\n",
      "Gaussian amplitude = -24.704442122122096, derivative amplitude = -138.02160674040888, Overlap: 0.9838396943815279\n",
      "Gaussian amplitude = -24.73620707818858, derivative amplitude = -137.98701484944633, Overlap: 0.9838396407934783\n",
      "Gaussian amplitude = -24.729734374927194, derivative amplitude = -138.00563834319405, Overlap: 0.9838397063405587\n",
      "optimal overlap = 0.9838397063405587\n",
      "optimal parameters = [-0.24729734374927193, -1.3800563834319406]\n"
     ]
    }
   ],
   "source": [
    "import cudaq\n",
    "from cudaq import boson, Schedule, ScalarOperator, ScipyZvodeIntegrator\n",
    "import numpy as np\n",
    "import cupy as cp\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Set the target to our dynamics simulator\n",
    "cudaq.set_target(\"dynamics\")\n",
    "\n",
    "# Sample device parameters\n",
    "# Assuming a simple transmon device Hamiltonian in rotating frame.\n",
    "detuning = 0.0  # Detuning of the drive; assuming resonant drive\n",
    "anharmonicity = -340.0  # Anharmonicity\n",
    "sigma = 0.01  # sigma of the Gaussian pulse\n",
    "cutoff = 4.0 * sigma  # total length of drive pulse\n",
    "\n",
    "# Dimensions of sub-system\n",
    "# We model `transmon` as a 3-level system to account for leakage.\n",
    "dimensions = {0: 3}\n",
    "\n",
    "# Initial state of the system (ground state).\n",
    "psi0 = cudaq.State.from_data(cp.array([1.0, 0.0, 0.0], dtype=cp.complex128))\n",
    "\n",
    "\n",
    "def gaussian(t):\n",
    "    \"\"\"\n",
    "    Gaussian shape with cutoff. Starts at t = 0, amplitude normalized to one\n",
    "    \"\"\"\n",
    "    val = (np.exp(-((t-cutoff/2)/sigma)**2/2)-np.exp(-(cutoff/sigma)**2/8)) \\\n",
    "           / (1-np.exp(-(cutoff/sigma)**2/8))\n",
    "    return val\n",
    "\n",
    "\n",
    "def dgaussian(t):\n",
    "    \"\"\"\n",
    "    Derivative of Gaussian. Starts at t = 0, amplitude normalized to one\n",
    "    \"\"\"\n",
    "    return -(t - cutoff / 2) / sigma * np.exp(-(\n",
    "        (t - cutoff / 2) / sigma)**2 / 2 + 0.5)\n",
    "\n",
    "\n",
    "# Schedule of time steps.\n",
    "steps = np.linspace(0.0, cutoff, 201)\n",
    "schedule = Schedule(steps, [\"t\"])\n",
    "\n",
    "# We optimize for a X(pi/2) rotation\n",
    "target_state = np.array([1.0 / np.sqrt(2), -1j / np.sqrt(2), 0.0],\n",
    "                        dtype=cp.complex128)\n",
    "\n",
    "\n",
    "# Optimize the amplitude of the drive pulse (DRAG - Derivative Removal by Adiabatic Gate)\n",
    "def cost_function(amps):\n",
    "    amplitude = 100 * amps[0]\n",
    "    drag_amp = 100 * amps[1]\n",
    "    # Qubit Hamiltonian\n",
    "    hamiltonian = detuning * boson.number(0) + (\n",
    "        anharmonicity / 2) * boson.create(0) * boson.create(\n",
    "            0) * boson.annihilate(0) * boson.annihilate(0)\n",
    "    # Drive term\n",
    "    hamiltonian += amplitude * ScalarOperator(gaussian) * (boson.create(0) +\n",
    "                                                           boson.annihilate(0))\n",
    "\n",
    "    # Drag term (leakage reduction)\n",
    "    hamiltonian += 1j * drag_amp * ScalarOperator(dgaussian) * (\n",
    "        boson.annihilate(0) - boson.create(0))\n",
    "\n",
    "    # We optimize for a X(pi/2) rotation\n",
    "    evolution_result = cudaq.evolve(hamiltonian,\n",
    "                                    dimensions,\n",
    "                                    schedule,\n",
    "                                    psi0,\n",
    "                                    observables=[],\n",
    "                                    collapse_operators=[],\n",
    "                                    store_intermediate_results=cudaq.IntermediateResultSave.NONE,\n",
    "                                    integrator=ScipyZvodeIntegrator())\n",
    "    final_state = evolution_result.final_state()\n",
    "\n",
    "    overlap = np.abs(final_state.overlap(target_state))\n",
    "    print(\n",
    "        f\"Gaussian amplitude = {amplitude}, derivative amplitude = {drag_amp}, Overlap: {overlap}\"\n",
    "    )\n",
    "    return 1.0 - overlap\n",
    "\n",
    "\n",
    "# Specify the optimizer\n",
    "optimizer = cudaq.optimizers.NelderMead()\n",
    "optimal_error, optimal_parameters = optimizer.optimize(dimensions=2,\n",
    "                                                       function=cost_function)\n",
    "\n",
    "print(\"optimal overlap =\", 1.0 - optimal_error)\n",
    "print(\"optimal parameters =\", optimal_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbd91da3-1f18-43db-b986-ce5c5078c0f0",
   "metadata": {},
   "source": [
    "## Pulse\n",
    "\n",
    "This example demonstrates how a transmon qubit is driven by a pulse (modulated signal with a Gaussian envelop) in the lab frame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2002b8f9-7271-4613-ae8d-ecaa3cb79b2e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'Qubit State')"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 105080 characters]",
      "text/plain": [
       "<Figure size 600x1600 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import cudaq\n",
    "from cudaq import spin, boson, ScalarOperator, Schedule, ScipyZvodeIntegrator\n",
    "import numpy as np\n",
    "import cupy as cp\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Set the target to our dynamics simulator\n",
    "cudaq.set_target(\"dynamics\")\n",
    "\n",
    "# Device parameters\n",
    "# Strength of the Rabi-rate in GHz.\n",
    "rabi_rate = 0.1\n",
    "\n",
    "# Frequency of the qubit transition in GHz.\n",
    "omega = 5.0 * 2 * np.pi\n",
    "\n",
    "# Define Gaussian envelope function to approximately implement a `rx(pi/2)` gate.\n",
    "amplitude = 1. / 2.0  # Pi/2 rotation\n",
    "sigma = 1.0 / rabi_rate / amplitude\n",
    "pulse_duration = 6 * sigma\n",
    "\n",
    "\n",
    "def gaussian(t, duration, amplitude, sigma):\n",
    "    # Gaussian envelope function\n",
    "    return amplitude * np.exp(-0.5 * (t - duration / 2)**2 / (sigma)**2)\n",
    "\n",
    "\n",
    "def signal(t):\n",
    "    # Modulated signal\n",
    "    return np.cos(omega * t) * gaussian(t, pulse_duration, amplitude, sigma)\n",
    "\n",
    "\n",
    "# Qubit Hamiltonian\n",
    "hamiltonian = omega * spin.z(0) / 2\n",
    "# Add modulated driving term to the Hamiltonian\n",
    "hamiltonian += np.pi * rabi_rate * ScalarOperator(signal) * spin.x(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}\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.\n",
    "# Since this is a lab-frame simulation, the time step must be small to accurately capture the modulated signal.\n",
    "dt = 1 / omega / 20\n",
    "n_steps = int(np.ceil(pulse_duration / dt)) + 1\n",
    "steps = np.linspace(0, pulse_duration, n_steps)\n",
    "schedule = Schedule(steps, [\"t\"])\n",
    "\n",
    "# Run the simulation.\n",
    "# First, we run the simulation without any collapse operators (no decoherence).\n",
    "evolution_result = cudaq.evolve(hamiltonian,\n",
    "                                dimensions,\n",
    "                                schedule,\n",
    "                                psi0,\n",
    "                                observables=[boson.number(0)],\n",
    "                                collapse_operators=[],\n",
    "                                store_intermediate_results=cudaq.IntermediateResultSave.EXPECTATION_VALUE,\n",
    "                                integrator=ScipyZvodeIntegrator())\n",
    "\n",
    "pop1 = [\n",
    "    exp_vals[0].expectation()\n",
    "    for exp_vals in evolution_result.expectation_values()\n",
    "]\n",
    "pop0 = [1.0 - x for x in pop1]\n",
    "fig = plt.figure(figsize=(6, 16))\n",
    "envelop = [gaussian(t, pulse_duration, amplitude, sigma) for t in steps]\n",
    "\n",
    "plt.subplot(3, 1, 1)\n",
    "plt.plot(steps, envelop)\n",
    "plt.ylabel(\"Amplitude\")\n",
    "plt.xlabel(\"Time\")\n",
    "plt.title(\"Envelope\")\n",
    "\n",
    "modulated_signal = [\n",
    "    np.cos(omega * t) * gaussian(t, pulse_duration, amplitude, sigma)\n",
    "    for t in steps\n",
    "]\n",
    "plt.subplot(3, 1, 2)\n",
    "plt.plot(steps, modulated_signal)\n",
    "plt.ylabel(\"Amplitude\")\n",
    "plt.xlabel(\"Time\")\n",
    "plt.title(\"Signal\")\n",
    "\n",
    "plt.subplot(3, 1, 3)\n",
    "plt.plot(steps, pop0)\n",
    "plt.plot(steps, pop1)\n",
    "plt.ylabel(\"Population\")\n",
    "plt.xlabel(\"Time\")\n",
    "plt.legend((\"Population in |0>\", \"Population in |1>\"))\n",
    "plt.title(\"Qubit State\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "179336e7-2336-4d4d-9fb8-90d51b404303",
   "metadata": {},
   "source": [
    "## Qubit Control\n",
    "\n",
    "This example simulates time evolution of a qubit (`transmon`) being driven close to resonance in the presence of noise (decoherence) and exhibiting Rabi oscillations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1e73a3b9-4a3b-4caf-a6a6-05719b635373",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'With decoherence')"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 206832 characters]",
      "text/plain": [
       "<Figure size 1800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import cudaq\n",
    "from cudaq import spin, ScalarOperator, Schedule, ScipyZvodeIntegrator\n",
    "import numpy as np\n",
    "import cupy as cp\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Set the target to our dynamics simulator\n",
    "cudaq.set_target(\"dynamics\")\n",
    "\n",
    "# Qubit Hamiltonian reference: https://qiskit-community.github.io/qiskit-dynamics/tutorials/Rabi_oscillations.html\n",
    "# Device parameters\n",
    "# Qubit resonant frequency\n",
    "omega_z = 10.0 * 2 * np.pi\n",
    "# Transverse term\n",
    "omega_x = 2 * np.pi\n",
    "# Harmonic driving frequency\n",
    "# Note: we chose a frequency slightly different from the resonant frequency to demonstrate the off-resonance effect.\n",
    "omega_drive = 0.99 * omega_z\n",
    "\n",
    "# Qubit Hamiltonian\n",
    "hamiltonian = 0.5 * omega_z * spin.z(0)\n",
    "# Add modulated driving term to the Hamiltonian\n",
    "hamiltonian += omega_x * ScalarOperator(\n",
    "    lambda t: np.cos(omega_drive * t)) * spin.x(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}\n",
    "\n",
    "# Initial state of the system (ground state).\n",
    "rho0 = cudaq.State.from_data(\n",
    "    cp.array([[1.0, 0.0], [0.0, 0.0]], dtype=cp.complex128))\n",
    "\n",
    "# Schedule of time steps.\n",
    "t_final = np.pi / omega_x\n",
    "dt = 2.0 * np.pi / omega_drive / 100\n",
    "n_steps = int(np.ceil(t_final / dt)) + 1\n",
    "steps = np.linspace(0, t_final, n_steps)\n",
    "schedule = Schedule(steps, [\"t\"])\n",
    "\n",
    "# Run the simulation.\n",
    "# First, we run the simulation without any collapse operators (no decoherence).\n",
    "evolution_result = cudaq.evolve(hamiltonian,\n",
    "                                dimensions,\n",
    "                                schedule,\n",
    "                                rho0,\n",
    "                                observables=[spin.x(0),\n",
    "                                             spin.y(0),\n",
    "                                             spin.z(0)],\n",
    "                                collapse_operators=[],\n",
    "                                store_intermediate_results=cudaq.IntermediateResultSave.EXPECTATION_VALUE,\n",
    "                                integrator=ScipyZvodeIntegrator())\n",
    "\n",
    "# Now, run the simulation with qubit decoherence\n",
    "gamma_sm = 4.0\n",
    "gamma_sz = 1.0\n",
    "evolution_result_decay = cudaq.evolve(\n",
    "    hamiltonian,\n",
    "    dimensions,\n",
    "    schedule,\n",
    "    rho0,\n",
    "    observables=[spin.x(0), spin.y(0), spin.z(0)],\n",
    "    collapse_operators=[\n",
    "        np.sqrt(gamma_sm) * spin.plus(0),\n",
    "        np.sqrt(gamma_sz) * spin.z(0)\n",
    "    ],\n",
    "    store_intermediate_results=cudaq.IntermediateResultSave.EXPECTATION_VALUE,\n",
    "    integrator=ScipyZvodeIntegrator())\n",
    "\n",
    "get_result = lambda idx, res: [\n",
    "    exp_vals[idx].expectation() for exp_vals in res.expectation_values()\n",
    "]\n",
    "ideal_results = [\n",
    "    get_result(0, evolution_result),\n",
    "    get_result(1, evolution_result),\n",
    "    get_result(2, evolution_result)\n",
    "]\n",
    "decoherence_results = [\n",
    "    get_result(0, evolution_result_decay),\n",
    "    get_result(1, evolution_result_decay),\n",
    "    get_result(2, evolution_result_decay)\n",
    "]\n",
    "\n",
    "fig = plt.figure(figsize=(18, 6))\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(steps, ideal_results[0])\n",
    "plt.plot(steps, ideal_results[1])\n",
    "plt.plot(steps, ideal_results[2])\n",
    "plt.ylabel(\"Expectation value\")\n",
    "plt.xlabel(\"Time\")\n",
    "plt.legend((\"Sigma-X\", \"Sigma-Y\", \"Sigma-Z\"))\n",
    "plt.title(\"No decoherence\")\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(steps, decoherence_results[0])\n",
    "plt.plot(steps, decoherence_results[1])\n",
    "plt.plot(steps, decoherence_results[2])\n",
    "plt.ylabel(\"Expectation value\")\n",
    "plt.xlabel(\"Time\")\n",
    "plt.legend((\"Sigma-X\", \"Sigma-Y\", \"Sigma-Z\"))\n",
    "plt.title(\"With decoherence\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f1acadd-a724-4b33-bf20-27e8740cfa03",
   "metadata": {},
   "source": [
    "## Qubit Dynamics "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ecfbd887-ded4-4905-80de-02a460550b74",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'With decay')"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 136872 characters]",
      "text/plain": [
       "<Figure size 1800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import cudaq\n",
    "from cudaq import spin, Schedule, RungeKuttaIntegrator\n",
    "import numpy as np\n",
    "import cupy as cp\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Set the target to our dynamics simulator\n",
    "cudaq.set_target(\"dynamics\")\n",
    "\n",
    "# Qubit Hamiltonian\n",
    "hamiltonian = 2 * np.pi * 0.1 * spin.x(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}\n",
    "\n",
    "# Initial state of the system (ground state).\n",
    "rho0 = cudaq.State.from_data(\n",
    "    cp.array([[1.0, 0.0], [0.0, 0.0]], dtype=cp.complex128))\n",
    "\n",
    "# Schedule of time steps.\n",
    "steps = np.linspace(0, 10, 101)\n",
    "schedule = Schedule(steps, [\"time\"])\n",
    "\n",
    "# Run the simulation.\n",
    "# First, we run the simulation without any collapse operators (ideal).\n",
    "evolution_result = cudaq.evolve(hamiltonian,\n",
    "                                dimensions,\n",
    "                                schedule,\n",
    "                                rho0,\n",
    "                                observables=[spin.y(0), spin.z(0)],\n",
    "                                collapse_operators=[],\n",
    "                                store_intermediate_results=cudaq.IntermediateResultSave.EXPECTATION_VALUE,\n",
    "                                integrator=RungeKuttaIntegrator())\n",
    "\n",
    "# Now, run the simulation with qubit decaying due to the presence of a collapse operator.\n",
    "evolution_result_decay = cudaq.evolve(\n",
    "    hamiltonian,\n",
    "    dimensions,\n",
    "    schedule,\n",
    "    rho0,\n",
    "    observables=[spin.y(0), spin.z(0)],\n",
    "    collapse_operators=[np.sqrt(0.05) * spin.x(0)],\n",
    "    store_intermediate_results=cudaq.IntermediateResultSave.EXPECTATION_VALUE,\n",
    "    integrator=RungeKuttaIntegrator())\n",
    "\n",
    "get_result = lambda idx, res: [\n",
    "    exp_vals[idx].expectation() for exp_vals in res.expectation_values()\n",
    "]\n",
    "ideal_results = [\n",
    "    get_result(0, evolution_result),\n",
    "    get_result(1, evolution_result)\n",
    "]\n",
    "decay_results = [\n",
    "    get_result(0, evolution_result_decay),\n",
    "    get_result(1, evolution_result_decay)\n",
    "]\n",
    "\n",
    "fig = plt.figure(figsize=(18, 6))\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(steps, ideal_results[0])\n",
    "plt.plot(steps, ideal_results[1])\n",
    "plt.ylabel(\"Expectation value\")\n",
    "plt.xlabel(\"Time\")\n",
    "plt.legend((\"Sigma-Y\", \"Sigma-Z\"))\n",
    "plt.title(\"No decay\")\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(steps, decay_results[0])\n",
    "plt.plot(steps, decay_results[1])\n",
    "plt.ylabel(\"Expectation value\")\n",
    "plt.xlabel(\"Time\")\n",
    "plt.legend((\"Sigma-Y\", \"Sigma-Z\"))\n",
    "plt.title(\"With decay\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2f3112f-8589-41d7-88b5-a16c9200e31f",
   "metadata": {},
   "source": [
    "## Landau-Zenner\n",
    "\n",
    "\n",
    "This example simulates the so-called Landau–Zener transition: given a time-dependent Hamiltonian such that the energy separationof the two states is a linear function of time, an analytical formula exists to calculate the probability of finding the system in the excited state after the transition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f82ea3eb-db07-48c5-9398-6eae4d1a069a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7fa9f4718af0>"
      ]
     },
     "execution_count": 5,
     "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": [
    "import cudaq\n",
    "from cudaq import spin, boson, ScalarOperator, Schedule, ScipyZvodeIntegrator\n",
    "import numpy as np\n",
    "import cupy as cp\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# References:\n",
    "# - https://en.wikipedia.org/wiki/Landau%E2%80%93Zener_formula\n",
    "# - `The Landau-Zener formula made simple`, `Eric P Glasbrenner and Wolfgang P Schleich 2023 J. Phys. B: At. Mol. Opt. Phys. 56 104001`\n",
    "# - QuTiP notebook: https://github.com/qutip/qutip-notebooks/blob/master/examples/landau-zener.ipynb\n",
    "\n",
    "# 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 = boson.annihilate(0)\n",
    "sm_dag = boson.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}\n",
    "\n",
    "# Landau–Zener Hamiltonian:\n",
    "# `[[-alpha*t, g], [g, alpha*t]] = g * pauli_x - alpha * t * pauli_z`\n",
    "g = 2 * np.pi\n",
    "# Analytical equation:\n",
    "# `P(0) = exp(-pi * g ^ 2/ alpha)`\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)\n",
    "\n",
    "# Hamiltonian\n",
    "hamiltonian = g * sx - alpha * ScalarOperator(lambda t: t) * sz\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\"])\n",
    "\n",
    "# Run the simulation.\n",
    "evolution_result = cudaq.evolve(hamiltonian,\n",
    "                                dimensions,\n",
    "                                schedule,\n",
    "                                psi0,\n",
    "                                observables=[boson.number(0)],\n",
    "                                collapse_operators=[],\n",
    "                                store_intermediate_results=cudaq.IntermediateResultSave.EXPECTATION_VALUE,\n",
    "                                integrator=ScipyZvodeIntegrator())\n",
    "\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)\n"
   ]
  }
 ],
 "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
}
