{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to CUDA-Q Dynamics (Jaynes-Cummings Model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Why dynamics simulations vs. circuit simulations?\n",
    "CUDA-Q Dynamics is a powerful library enabling GPU-accelerated simulations of the time evolution of open and closed quantum systems. In contrast to quantum circuit simulation, which simulates the algorithmic output of the application of discrete logical quantum gates (eg X, Y, Z, CNOT, H, etc), dynamics simulations represent how quantum systems evolve in *time*, revealing important insights about the non-idealities of a quantum system, including how they interact with the environment and other quantum systems.\n",
    "\n",
    "To make a classical analogy, the logic of a classical computer can be modeled using binary logic gates applied to transistor states, analogous to simulating the application of quantum logic gates to qubits. Dynamics simulations are an abstraction layer below this. Classically, to design better and higher-performing transistors, it's useful to fully model the device physics of the transistor itself, including fluctuations in voltage, capacitance, and current. Analogously, quantum dyanmics simulations can fully capture the physics of components beyond just the qubits, including drive pulses, resonators, couplers, and noise."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![alt text 2](images/Benchmarks.png)\n",
    "\n",
    "Figure 1. Classical computing can be a helpful analogy for understanding the difference between algorithm and dynamics simulations. Image credits: Fredrik Brange (top right) and qutip-qip (bottom right)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Functionality\n",
    "\n",
    "CUDA-Q dynamics provides an easy API, `evolve`, to solve the Schrodinger and Lindblad Master equation, as well as number of commonly used [operators](https://nvidia.github.io/cuda-quantum/latest/using/dynamics.html#operator) to construct Hamiltonians. The basic order of operations is similar to other quantum dynamics frameworks. \n",
    "1. Define:\n",
    "    - The dimensions of the system\n",
    "    - The Hamiltonian\n",
    "    - The initial state of the system\n",
    "    - Any dissipation terms\n",
    "    - The timesteps desired\n",
    "2. Evolve!\n",
    "\n",
    "Conveniently, this is the same programming model and API used to submit CUDA-Q jobs to analog quantum computers like QuEra and Pasqal. For simulation workloads, we also provide several commonly used [numerical integrators](https://nvidia.github.io/cuda-quantum/latest/using/dynamics.html#numerical-integrators).\n",
    "\n",
    "In addition to this higher level functionality, we also provide a lower level library called [cuDensityMat](https://docs.nvidia.com/cuda/cuquantum/latest/cudensitymat/index.html) which provide the building blocks for accelerating custom solvers and other dynamics frameworks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Performance\n",
    "CUDA-Q Dynamics shines for *large* system sizes, as is the case for most GPU workloads. We will see maximal speed-ups over CPU implementations with Hilbert spaces larger than $O(1000)$ levels. For systems of this size and larger, it is common to see speed-ups of >1000x. A benchmark modeling 100 timesteps of a system consisting of a 32-level transmon, 256-level resonator, and 4-level Purcell filter yielded a 1,140x speed-up in execution from single CPU (dual-socket Intel Xeon 8480CL) to single GPU (H100), or 1.5 days to 2 minutes.\n",
    "\n",
    "CUDA-Q Dynamics also provides multi-GPU and multi-node capabilities to both increase the performance and scale of simulations. Additional GPUs can be pooled for additional memory to simulate larger systems (weak-scaling). Alternatively, with a fixed system size, additional GPUs can be pooled for increase performance (strong-scaling).\n",
    "\n",
    "Full details of the benchmarks can be found in the blog [here](https://developer.nvidia.com/blog/accelerating-googles-qpu-development-with-new-quantum-dynamics-capabilities/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 1 - Simulating the Jaynes-Cummings Hamiltonian\n",
    "The [Jaynes-Cummings](https://en.wikipedia.org/wiki/Jaynes%E2%80%93Cummings_model) (JC) Hamiltonian is a important theoretical model describing a two-level system interacting with a resonator. It describes how an atom or qubit fundamentally interacts with light. It is a powerful building block for modeling and designing many superconducting qubit systems and can be extended to describe many of the interactions in a superconducting qubit system. In practice, it describes the physics by which qubits can be controlled by external control pulses and the readout process used to determine their state.\n",
    "\n",
    "![alt text 1](images/JC_model.png)\n",
    "\n",
    "Figure 2. A two-level quantum system such as a qubit or an atom interacting with a light field confined by a resonator. The JC model describes this interaction. Image credit: Wikipedia.\n",
    "\n",
    "\n",
    "The Hamiltonian can be simplified into the following three principal terms, by taking the rotating wave approximation, and dropping constant terms.\n",
    "\n",
    "$\\hat{H}_{JC} = \\hat{H}_{cavity} +\\hat{H}_{qubit} +\\hat{H}_{interaction} $\n",
    "\n",
    "with \n",
    "\n",
    "$\\hat{H}_{cavity} = \\hbar \\omega_c \\hat a^\\dagger_c \\hat a_c$ \n",
    "\n",
    "$\\hat{H}_{qubit} = \\hbar \\omega_q \\hat \\sigma_+ \\hat \\sigma_-$ \n",
    "\n",
    "$\\hat{H}_{interaction} = \\hbar g_c (\\hat \\sigma_+ \\hat a_c + \\hat \\sigma_- \\hat a^\\dagger_c)$ \n",
    "\n",
    "In this section, we will simulate JC Hamiltonian. You will learn to:\n",
    "- Define the dimensions of the system\n",
    "- Construct the Hamiltonian with CUDA-Q `operators`\n",
    "- Define the initial state of the system\n",
    "- Execute the simulation with `cudaq.evolve()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary packages\n",
    "import cudaq\n",
    "from cudaq import boson, Schedule, ScipyZvodeIntegrator, spin, RungeKuttaIntegrator\n",
    "import numpy as np\n",
    "import cupy as cp\n",
    "import os\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start by setting the backend target to `dynamics`. Then, we define the degrees of freedom, specifying the number of levels associated with each one. We can then construct the Hamiltonian by applying the correct operators to the relevant degrees of freedom."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the target to our dynamics simulator\n",
    "cudaq.set_target(\"dynamics\")\n",
    "\n",
    "# System dimensions: atom (2-level system) and cavity (2-level system)\n",
    "dimensions = {0: 2, 1: 2}\n",
    "\n",
    "# Alias for commonly used operators\n",
    "# Cavity operators\n",
    "a = boson.annihilate(1)\n",
    "a_dag = boson.create(1)\n",
    "\n",
    "# Atom operators\n",
    "sm = boson.annihilate(0)\n",
    "sm_dag = boson.create(0)\n",
    "\n",
    "# Defining the Hamiltonian for the system: self-energy terms and cavity-atom interaction term.\n",
    "H_cav = 2 * np.pi * boson.number(1)\n",
    "H_qubit = 2 * np.pi * boson.number(0) \n",
    "H_int = 2 * np.pi * 0.25 * (sm * a_dag + sm_dag * a)\n",
    "hamiltonian = H_cav + H_qubit + H_int\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize the states of the system in cupy arrays and define the time schedule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initial state of the system\n",
    "# Qubit in ground state\n",
    "qubit_state = cp.array([[1.0, 0.0], [0.0, 0.0]], dtype=cp.complex128)\n",
    "\n",
    "# Cavity in a state which has 1 photon initially\n",
    "cavity_state = cp.zeros((2, 2), dtype=cp.complex128)\n",
    "cavity_state[1][1] = 1.0\n",
    "\n",
    "# Construct the Density Matrix\n",
    "rho0 = cudaq.State.from_data(cp.kron(cavity_state, qubit_state))\n",
    "\n",
    "# Define our timesteps\n",
    "steps = np.linspace(0, 10, 201)\n",
    "schedule = Schedule(steps, [\"time\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evolve! Various observables can be collected at each time step. For this, we will track the population of both the qubit and the cavity in the absence of noise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evolve the system\n",
    "evolution_result = cudaq.evolve(\n",
    "    hamiltonian,\n",
    "    dimensions,\n",
    "    schedule,\n",
    "    rho0,\n",
    "    observables=[boson.number(1), boson.number(0)],\n",
    "    collapse_operators=[],\n",
    "    store_intermediate_results=cudaq.IntermediateResultSave.EXPECTATION_VALUE,\n",
    "    integrator=ScipyZvodeIntegrator())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'No decay')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 77788 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Collect results\n",
    "get_result = lambda idx, res: [\n",
    "    exp_vals[idx].expectation() for exp_vals in res.expectation_values()\n",
    "]\n",
    "\n",
    "ideal_results = [\n",
    "    get_result(0, evolution_result),\n",
    "    get_result(1, evolution_result)\n",
    "]\n",
    "\n",
    "# Plot the results\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((\"Cavity Photon Number\", \"Qubit Excitation Probability\"))\n",
    "plt.title(\"No decay\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see the qubit and cavity coherently exchanging population with a Rabi period of 2. Intuitively, a photon is being absorbed by the qubit and then re-emitted to the cavity. \n",
    "\n",
    "Recall that we set the coupling rate in the interation Hamiltonian to $g=\\pi/2$. From performing some analysis on the JC Hamiltonian, one can derive that this oscillation rate is $\\Omega_n = 2 g \\sqrt{n +1}$, where the effective subspace is spanned by $\\{\\ket{e,n},\\ket{g,n+1}\\}$. Is this consistent with what we observe?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 1 - Simulating a many-photon Jaynes-Cummings Hamiltonian\n",
    "\n",
    "Let us now adjust the code to simulate a larger system. Adjust the code below to now simulate 20 photons in the cavity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# System dimensions: atom (2-level system) and cavity (20 photon system)\n",
    "dimensions = {0: 2, 1: 21}\n",
    "\n",
    "# Initial state of the system\n",
    "# Qubit in ground state\n",
    "qubit_state = cp.array([[1.0, 0.0], [0.0, 0.0]], dtype=cp.complex128)\n",
    "\n",
    "\n",
    "# Cavity in a state which has 20 photons initially\n",
    "cavity_state = cp.zeros((21, 21), dtype=cp.complex128)\n",
    "cavity_state[20][20] = 1.0\n",
    "\n",
    "# Construct the Density Matrix\n",
    "rho0 = cudaq.State.from_data(cp.kron(cavity_state, qubit_state))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evolve the system\n",
    "evolution_result = cudaq.evolve(\n",
    "    hamiltonian,\n",
    "    dimensions,\n",
    "    schedule,\n",
    "    rho0,\n",
    "    observables=[boson.number(1), boson.number(0)],\n",
    "    collapse_operators=[],\n",
    "    store_intermediate_results=cudaq.IntermediateResultSave.EXPECTATION_VALUE,\n",
    "    integrator=ScipyZvodeIntegrator())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, '20 Photons')"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 53384 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Collect results\n",
    "get_result = lambda idx, res: [\n",
    "    exp_vals[idx].expectation() for exp_vals in res.expectation_values()\n",
    "]\n",
    "\n",
    "ideal_results = [\n",
    "    get_result(0, evolution_result),\n",
    "    get_result(1, evolution_result)\n",
    "]\n",
    "\n",
    "# Plot the results\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((\"Cavity Photon Number\", \"Qubit Excitation Probability\"))\n",
    "plt.title(\"20 Photons\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What do you observe about the new Rabi frequency? Is this consistent with what you expect?\n",
    "\n",
    "The Rabi frequency is proportional to $2 g \\sqrt{n+1}$, so $\\sqrt{20}=4.5$ times faster than the 1 photon case "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 2 - Simulating open quantum systems with the `collapse_operators`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the previous section we simulated an idealized quantum system, isolated from the environment, with no noise. This is equivalent to solving the Schrodinger equation. Let us now add the effects of the environment, by adding noise to the system via the `collapse_operators`. We will now solve the Lindbladian:\n",
    "\n",
    "$\\dot \\rho = -\\frac{i}{\\hbar} [H,\\rho]+ \\sum_i \\gamma_i (L_i \\rho L^\\dagger_i-\\frac{1}{2}\\{L^\\dagger_i L_i,\\rho\\})$\n",
    "\n",
    "where we can assign the dissipation terms $L_i$ and rates $\\gamma_i$.\n",
    "\n",
    "In this section we will learn to use the `collapse_operators` to model noise in the system.\n",
    "\n",
    "We can add a list of dissipation terms to the `collapse_operators` argument in `cudaq.evolve`. Let's add a photon loss mechanism with rate 0.1 with the `a` annihilation operator and observe what happens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'With decay')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 43808 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "gamma=np.sqrt(0.1)\n",
    "\n",
    "# Then, evolve the system with a collapse operator modeling cavity decay (leaking photons)\n",
    "evolution_result_decay = cudaq.evolve(\n",
    "    hamiltonian,\n",
    "    dimensions,\n",
    "    schedule,\n",
    "    rho0,\n",
    "    observables=[boson.number(1), boson.number(0)],\n",
    "    collapse_operators=[gamma * a],\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",
    "decay_results = [\n",
    "    get_result(0, evolution_result_decay),\n",
    "    get_result(1, evolution_result_decay)\n",
    "]\n",
    "\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((\"Cavity Photon Number\", \"Qubit Excitation Probability\"))\n",
    "plt.title(\"With decay\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 2 - Adding additional jump operators $L_i$\n",
    "\n",
    "2a. Now instead of considering a photon loss mechanism, consider a photon excitation mechanism with rate 0.1. Initialize the cavity with one photon."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# System dimensions: atom (2-level system) and cavity (20 photon system)\n",
    "dimensions = {0: 2, 1: 21}\n",
    "\n",
    "# Initial state of the system\n",
    "# Qubit in ground state\n",
    "qubit_state = cp.array([[1.0, 0.0], [0.0, 0.0]], dtype=cp.complex128)\n",
    "\n",
    "# Cavity in a state which has 1 photon initially\n",
    "cavity_state = cp.zeros((21, 21), dtype=cp.complex128)\n",
    "cavity_state[1][1] = 1.0\n",
    "\n",
    "# Construct the Density Matrix\n",
    "rho0 = cudaq.State.from_data(cp.kron(cavity_state, qubit_state))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'With excitation')"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 52576 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "gamma=np.sqrt(0.1)\n",
    "collapse_operators=[gamma * a_dag]\n",
    "\n",
    "# Then, evolve the system with a collapse operator modeling cavity decay (leaking photons)\n",
    "evolution_result_excitation = cudaq.evolve(\n",
    "    hamiltonian,\n",
    "    dimensions,\n",
    "    schedule,\n",
    "    rho0,\n",
    "    observables=[boson.number(1), boson.number(0)],\n",
    "    collapse_operators=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",
    "decay_results = [\n",
    "    get_result(0, evolution_result_excitation),\n",
    "    get_result(1, evolution_result_excitation)\n",
    "]\n",
    "\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((\"Cavity Photon Number\", \"Qubit Excitation Probability\"))\n",
    "plt.title(\"With excitation\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2b. Now add a dephasing term to qubit with rate `dephasing=0.4`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'With excitation and qubit dephasing')"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 47560 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "gamma=np.sqrt(0.1)\n",
    "dephasing=np.sqrt(0.4)\n",
    "collapse_operators=[gamma * a_dag, dephasing * spin.z(0)]\n",
    "\n",
    "# Then, evolve the system with a collapse operator modeling cavity decay (leaking photons)\n",
    "evolution_result_excitation = cudaq.evolve(\n",
    "    hamiltonian,\n",
    "    dimensions,\n",
    "    schedule,\n",
    "    rho0,\n",
    "    observables=[boson.number(1), boson.number(0)],\n",
    "    collapse_operators=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",
    "decay_results = [\n",
    "    get_result(0, evolution_result_excitation),\n",
    "    get_result(1, evolution_result_excitation)\n",
    "]\n",
    "\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((\"Cavity Photon Number\", \"Qubit Excitation Probability\"))\n",
    "plt.title(\"With excitation and qubit dephasing\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 3 - Many qubits coupled to the resonator\n",
    "Let us now construct a model of several qubits coupled to a resonator. This is known as the Tavis-Cummings model. We define the dimensions of the system, construct the Hamiltonian assuming a distribution of qubit-resonator couplings near $\\pi/2$, qubit energies near $2\\pi$, and photon energies of $2\\pi$. Initialize all qubits in the ground state and with the resonator with 10 photons. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# System dimensions: atom (2-level system) and cavity (2-level system)\n",
    "no_qubits=6\n",
    "dimensions = {k: 11 if k == 0 else 2 for k in range(no_qubits+1)}\n",
    "\n",
    "# Define our timesteps\n",
    "steps = np.linspace(0, 10, 101)\n",
    "schedule = Schedule(steps, [\"time\"])\n",
    "\n",
    "# Cavity operators\n",
    "a = boson.annihilate(0)\n",
    "a_dag = boson.create(0)\n",
    "\n",
    "# Defining the Hamiltonian for the system: self-energy terms and cavity-atom interaction term.\n",
    "H_cav = 2 * np.pi * boson.number(0)\n",
    "\n",
    "# Qubit and interaction Hamiltonians\n",
    "def H_qubit_i(qubit, energy):\n",
    "    return energy * boson.number(qubit) \n",
    "def H_int_i(qubit, coupling):\n",
    "    return coupling * (boson.annihilate(qubit) * a_dag + boson.create(qubit) * a)\n",
    "\n",
    "# Qubit and coupling energies\n",
    "coupling = [np.pi/4 * random.gauss(1,0.1) for i in range(no_qubits)]\n",
    "qubit_energy = [2 * np.pi * random.gauss(1,0.1) for i in range(no_qubits)]\n",
    "\n",
    "hamiltonian = H_cav\n",
    "for qubit in range(1,no_qubits+1):\n",
    "    hamiltonian+=H_qubit_i(qubit, qubit_energy[qubit-1])+H_int_i(qubit, coupling[qubit-1])\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initial state of the system\n",
    "\n",
    "# Cavity in a state which has 10 photons initially\n",
    "cavity_state = cp.zeros((11, 11), dtype=cp.complex128)\n",
    "cavity_state[10][10] = 1.0\n",
    "\n",
    "# Qubits in ground state\n",
    "qubit_g = cp.array([[1.0, 0.0], [0.0, 0.0]], dtype=cp.complex128)\n",
    "\n",
    "# Qubit in excited state\n",
    "qubit_e = cp.array([[0.0, 0.0], [0.0, 1.0]], dtype=cp.complex128)\n",
    "\n",
    "composite_state = cavity_state\n",
    "for i in range(1,no_qubits+1):\n",
    "    composite_state = cp.kron(qubit_g, composite_state)\n",
    "\n",
    "# Construct the Density Matrix\n",
    "rho0 = cudaq.State.from_data(composite_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional: Add dissipation terms for each term\n",
    "T1=[abs(random.gauss(100,100)) for i in range(no_qubits+1)]\n",
    "T1[0]=50\n",
    "\n",
    "gamma=[1/np.sqrt(T1[i]) for i in range(no_qubits+1)]\n",
    "collapse_operators=[gamma[i]*boson.annihilate(i) for i in range(no_qubits+1)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the collapse operators\n",
    "collapse_operators=[]\n",
    "\n",
    "# Evolve the system\n",
    "evolution_result = cudaq.evolve(\n",
    "    hamiltonian,\n",
    "    dimensions,\n",
    "    schedule,\n",
    "    rho0,\n",
    "    observables=[boson.number(i) for i in range(no_qubits+1)],\n",
    "    collapse_operators=collapse_operators,\n",
    "    store_intermediate_results=cudaq.IntermediateResultSave.EXPECTATION_VALUE,\n",
    "    integrator=ScipyZvodeIntegrator())\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'Many qubits coupled to a resonator')"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 82028 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "get_result = lambda idx, res: [\n",
    "    exp_vals[idx].expectation() for exp_vals in res.expectation_values()\n",
    "]\n",
    "\n",
    "tc_results = [\n",
    "    get_result(i, evolution_result) for i in range(no_qubits+1)\n",
    "]\n",
    "plt.plot()\n",
    "\n",
    "for i in range(no_qubits):\n",
    "    plt.plot(steps, tc_results[i])\n",
    "plt.ylabel(\"Expectation value\")\n",
    "plt.xlabel(\"Time\")\n",
    "plt.legend((\"Cavity Photon Number\", \"Qubit Excitation Probability\"))\n",
    "plt.title(\"Many qubits coupled to a resonator\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now see the effect of multiple qubits exchanging excitations at different rates with the photons in the resonator"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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": 4
}
