{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "64448bd9-7802-43aa-b8cd-8265dc474c24",
   "metadata": {},
   "source": [
    "# Trapped Ion Qubits\n",
    "\n",
    "These examples all cover dynamical simulation of trapped ion 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": "e4fd221c-e31e-4082-b2ca-3b9418d050ca",
   "metadata": {},
   "source": [
    "## GHZ state\n",
    "This example shows how to prepare a GHZ (`Greenberger-Horne-Zeilinger`) state using trapped ions and CUDA-Q. We'll use the effective spin model from the famous `Sørensen-Mølmer` [paper.](https://www.quantumoptics.at/images/publications/dissertation/ThesisMvM_Web.pdf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bdf1a7fc-95a6-4814-84c0-f9c83a8d5ed8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running time evolution...\n",
      "\n",
      "Results at GHZ time (χt = 0.393):\n",
      "P(|gggggggg⟩) = 0.500\n",
      "P(|eeeeeeee⟩) = 0.500\n",
      "Total in extremes: 0.999\n",
      "GHZ state quality: 1.000 (1.0 = perfect)\n",
      "Perfect GHZ would have P(gg) = P(ee) = 0.5\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 89692 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\"\"\"\n",
    "What we're doing:\n",
    "- Start with N ions all in ground state `|gg...g⟩`\n",
    "- Apply an effective Hamiltonian H = 4χJ_x²  \n",
    "- Watch the system evolve into a GHZ superposition: `(|gg...g⟩ + |ee...e⟩)/√2`\n",
    "\n",
    "This tutorial uses the simplified effective model (`Eq. 2` from the paper).\n",
    "\"\"\"\n",
    "\n",
    "import cudaq\n",
    "from cudaq import spin, Schedule, RungeKuttaIntegrator\n",
    "import numpy as np\n",
    "import cupy as cp\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "cudaq.set_target(\"dynamics\")\n",
    "\n",
    "# Physical parameters, these come from the `Sørensen-Mølmer` paper\n",
    "nu = 1.0  # Trap frequency (our reference)\n",
    "delta = 0.9 * nu  # Laser detuning\n",
    "Omega = 0.1 * nu  # Laser strength\n",
    "eta = 0.1  # How strongly lasers couple to motion\n",
    "\n",
    "# The effective coupling strength (this is the key parameter!)\n",
    "chi = (eta**2 * Omega**2 * nu) / (2 * (nu**2 - delta**2))\n",
    "\n",
    "N = 8  # Number of ions (start small!)\n",
    "evolution_time = 1.0 / chi  # How long to evolve\n",
    "num_steps = 100  # Time resolution\n",
    "\n",
    "dimensions = {i: 2 for i in range(N)}  # Each ion is a 2-level system\n",
    "\n",
    "# collective spin operator J_x\n",
    "J_x = spin.empty()\n",
    "for i in range(N):\n",
    "    J_x += spin.x(i)\n",
    "J_x /= 2  # Normalize\n",
    "\n",
    "hamiltonian = 4 * chi * J_x * J_x\n",
    "\n",
    "# Set up initial state and time evolution\n",
    "# Start with all ions in ground state `|gg...g⟩`\n",
    "initial_state_vector = cp.zeros(2**N, dtype=cp.complex128)\n",
    "initial_state_vector[0] = 1.0  # |00...0⟩ = `|gg...g⟩`\n",
    "initial_state = cudaq.State.from_data(initial_state_vector)\n",
    "\n",
    "# Set up time points for evolution\n",
    "times = np.linspace(0, evolution_time, num_steps)\n",
    "chi_times = chi * times  # Dimensionless time χt\n",
    "schedule = Schedule(times, [\"t\"])\n",
    "\n",
    "# Create observables to track the populations\n",
    "# We want to measure the probability of being in `|gg...g⟩` and `|ee...e⟩`\n",
    "\n",
    "# Projector onto `|gg...g⟩ = |00...0⟩`\n",
    "P_ground = spin.empty()\n",
    "for i in range(N):\n",
    "    if i == 0:\n",
    "        P_ground = (spin.identity(i) - spin.z(i)) / 2  # `|0⟩⟨0|`\n",
    "    else:\n",
    "        P_ground = P_ground * (spin.identity(i) - spin.z(i)) / 2\n",
    "\n",
    "# Projector onto `|ee...e⟩ = |11...1⟩`\n",
    "P_excited = spin.empty()\n",
    "for i in range(N):\n",
    "    if i == 0:\n",
    "        P_excited = (spin.identity(i) + spin.z(i)) / 2  # `|1⟩⟨1|`\n",
    "    else:\n",
    "        P_excited = P_excited * (spin.identity(i) + spin.z(i)) / 2\n",
    "\n",
    "observables = [P_ground, P_excited]\n",
    "\n",
    "# Run the simulation!\n",
    "print(\"Running time evolution...\")\n",
    "evolution_result = cudaq.evolve(\n",
    "    hamiltonian,\n",
    "    dimensions,\n",
    "    schedule,\n",
    "    initial_state,\n",
    "    observables=observables,\n",
    "    collapse_operators=[],  # No decoherence for this tutorial\n",
    "    store_intermediate_results=cudaq.IntermediateResultSave.EXPECTATION_VALUE,  # Save expectation values\n",
    "    integrator=RungeKuttaIntegrator())\n",
    "\n",
    "# Extract the results\n",
    "exp_vals = evolution_result.expectation_values()\n",
    "pop_ground = [exp_vals[i][0].expectation() for i in range(len(times))]\n",
    "pop_excited = [exp_vals[i][1].expectation() for i in range(len(times))]\n",
    "\n",
    "# The GHZ state appears at a special time: χt = π/8\n",
    "ghz_chi_t = np.pi / 8\n",
    "ghz_time_idx = np.argmin(np.abs(chi_times - ghz_chi_t))\n",
    "ghz_pop_ground = pop_ground[ghz_time_idx]\n",
    "ghz_pop_excited = pop_excited[ghz_time_idx]\n",
    "\n",
    "print(f\"\\nResults at GHZ time (χt = {ghz_chi_t:.3f}):\")\n",
    "print(f\"P(|{'g'*N}⟩) = {ghz_pop_ground:.3f}\")\n",
    "print(f\"P(|{'e'*N}⟩) = {ghz_pop_excited:.3f}\")\n",
    "print(f\"Total in extremes: {ghz_pop_ground + ghz_pop_excited:.3f}\")\n",
    "\n",
    "# Check GHZ state quality\n",
    "# For perfect GHZ state: `P(gg) = P(ee) = 0.5`\n",
    "ghz_quality = 1 - 2 * abs(\n",
    "    ghz_pop_ground - 0.5)  # Distance from ideal probability\n",
    "print(f\"GHZ state quality: {ghz_quality:.3f} (1.0 = perfect)\")\n",
    "print(f\"Perfect GHZ would have P(gg) = P(ee) = 0.5\")\n",
    "\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.plot(chi_times,\n",
    "         pop_ground,\n",
    "         'b-',\n",
    "         linewidth=2,\n",
    "         label=f\"|{'g'*N}⟩ (all ground)\")\n",
    "plt.plot(chi_times,\n",
    "         pop_excited,\n",
    "         'r-',\n",
    "         linewidth=2,\n",
    "         label=f\"|{'e'*N}⟩ (all excited)\")\n",
    "plt.axvline(ghz_chi_t,\n",
    "            color='gray',\n",
    "            linestyle='--',\n",
    "            alpha=0.7,\n",
    "            label=\"GHZ time\")\n",
    "plt.axhline(0.5, color='black', linestyle=':', alpha=0.5, label=\"Perfect GHZ\")\n",
    "\n",
    "plt.xlabel(\"Dimensionless time (χt)\")\n",
    "plt.ylabel(\"Population\")\n",
    "plt.title(f\"Evolution to GHZ State ({N} trapped ions)\")\n",
    "plt.legend()\n",
    "plt.grid(True, alpha=0.3)\n",
    "plt.xlim(0, 1)\n",
    "plt.ylim(0, 1)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "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
}
