{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "95df216f-42b0-4bc4-8cb2-600fd8067064",
   "metadata": {},
   "source": [
    "# Spin Qubits\n",
    "\n",
    "These examples all cover dynamical simulation of spin qubits. 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": "fdd6afdb-d328-49df-84dc-ca7a018d2fce",
   "metadata": {},
   "source": [
    "## Silicon Spin Qubit\n",
    "\n",
    "This example demonstrates simulation of an electrically-driven silicon spin qubit taken from [\"Nonlinear Response and Crosstalk of Electrically Driven Silicon Spin Qubits\"](https://journals.aps.org/prapplied/pdf/10.1103/PhysRevApplied.19.044078)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e34ad263-766d-4470-9ec6-af294c90afdb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f61b2d4f9d0>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 44228 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import cudaq\n",
    "from cudaq import spin, boson, Schedule, ScalarOperator, ScipyZvodeIntegrator\n",
    "import numpy as np\n",
    "import cupy as cp\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "\n",
    "# Set the target to our dynamics simulator\n",
    "cudaq.set_target(\"dynamics\")\n",
    "\n",
    "dimensions = {0: 2}\n",
    "resonance_frequency = 2 * np.pi * 10  # 10 Ghz\n",
    "\n",
    "# Run the simulation:\n",
    "evolution_results = []\n",
    "get_result = lambda idx, res: [\n",
    "    exp_vals[idx].expectation() for exp_vals in res.expectation_values()\n",
    "]\n",
    "# Sweep the amplitude\n",
    "amplitudes = np.linspace(0.0, 0.5, 20)\n",
    "# Construct a list of Hamiltonian operator for each amplitude so that we can batch them all together\n",
    "batched_hamiltonian = []\n",
    "for amplitude in amplitudes:\n",
    "    # Electric dipole spin resonance (`EDSR`) Hamiltonian\n",
    "    H = 0.5 * resonance_frequency * spin.z(0) + amplitude * ScalarOperator(\n",
    "        lambda t: 0.5 * np.sin(resonance_frequency * t)) * spin.x(0)\n",
    "    # Append the Hamiltonian to the batched list\n",
    "    # This allows us to compute the dynamics for all amplitudes in a single simulation run\n",
    "    batched_hamiltonian.append(H)\n",
    "\n",
    "# Initial state is the ground state of the spin qubit\n",
    "# We run all simulations for the same initial state, but with different Hamiltonian operators.\n",
    "psi0 = cudaq.State.from_data(cp.array([1.0, 0.0], dtype=cp.complex128))\n",
    "\n",
    "# Simulation schedule\n",
    "t_final = 100\n",
    "dt = 0.005\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",
    "results = cudaq.evolve(\n",
    "    batched_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",
    "get_result = lambda idx, res: [\n",
    "    exp_vals[idx].expectation() for exp_vals in res.expectation_values()\n",
    "]\n",
    "\n",
    "evolution_results = []\n",
    "for result in results:\n",
    "    evolution_results.append(get_result(0, result))\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "im = ax.contourf(steps, amplitudes, evolution_results)\n",
    "ax.set_xlabel(\"Time (ns)\")\n",
    "ax.set_ylabel(f\"Amplitude (a.u.)\")\n",
    "fig.suptitle(f\"Excited state probability\")\n",
    "fig.colorbar(im)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3791ecd9-beb0-4407-a387-852d2e8372b7",
   "metadata": {},
   "source": [
    "## Heisenberg Model\n",
    "\n",
    "This example solves the [Quantum Heisenberg model](https://en.wikipedia.org/wiki/Quantum_Heisenberg_model),which exhibits the so-called quantum quench effect. See [\"Quantum quenches in the anisotropic spin-1/2 Heisenberg chain: different approaches to many-body dynamics far from equilibrium\"](https://iopscience.iop.org/article/10.1088/1367-2630/12/5/055017) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dfe45c6f-eb8c-4504-ba65-deb1a590a44d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0, 'Time')"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 105388 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import cudaq\n",
    "from cudaq import spin, Schedule, ScipyZvodeIntegrator\n",
    "\n",
    "import numpy as np\n",
    "import cupy as cp\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "\n",
    "# Set the target to our dynamics simulator\n",
    "cudaq.set_target(\"dynamics\")\n",
    "\n",
    "# Number of spins\n",
    "N = 9\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",
    "# Observable is the staggered magnetization operator\n",
    "staggered_magnetization_op = spin.empty()\n",
    "for i in range(N):\n",
    "    if i % 2 == 0:\n",
    "        staggered_magnetization_op += spin.z(i)\n",
    "    else:\n",
    "        staggered_magnetization_op -= spin.z(i)\n",
    "\n",
    "staggered_magnetization_op /= N\n",
    "\n",
    "observe_results = []\n",
    "batched_hamiltonian = []\n",
    "anisotropy_parameters = [0.0, 0.25, 4.0]\n",
    "for g in anisotropy_parameters:\n",
    "    # Heisenberg model spin coupling strength\n",
    "    Jx = 1.0\n",
    "    Jy = 1.0\n",
    "    Jz = g\n",
    "\n",
    "    # Construct the Hamiltonian\n",
    "    H = spin.empty()\n",
    "\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",
    "    # Append the Hamiltonian to the batched list\n",
    "    batched_hamiltonian.append(H)\n",
    "\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 in batched mode\n",
    "evolution_results = cudaq.evolve(\n",
    "    batched_hamiltonian,\n",
    "    dimensions,\n",
    "    schedule,\n",
    "    psi0,  # Same initial state for all Hamiltonian operators\n",
    "    observables=[staggered_magnetization_op],\n",
    "    collapse_operators=[],\n",
    "    store_intermediate_results=cudaq.IntermediateResultSave.EXPECTATION_VALUE,\n",
    "    integrator=ScipyZvodeIntegrator())\n",
    "\n",
    "for g, evolution_result in zip(anisotropy_parameters, evolution_results):\n",
    "    exp_val = [\n",
    "        exp_vals[0].expectation()\n",
    "        for exp_vals in evolution_result.expectation_values()\n",
    "    ]\n",
    "\n",
    "    observe_results.append((g, exp_val))\n",
    "\n",
    "# Plot the results\n",
    "fig = plt.figure(figsize=(12, 6))\n",
    "for g, exp_val in observe_results:\n",
    "    plt.plot(steps, exp_val, label=f'$ g = {g}$')\n",
    "plt.legend(fontsize=16)\n",
    "plt.ylabel(\"Staggered Magnetization\")\n",
    "plt.xlabel(\"Time\")"
   ]
  }
 ],
 "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
}
