{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "64448bd9-7802-43aa-b8cd-8265dc474c24",
   "metadata": {},
   "source": [
    "# Superconducting Qubits\n",
    "\n",
    "These examples all cover dynamical simulation of superconducting 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": [
    "## Cavity QED\n",
    "This example demonstrates a simulation of cavity quantum electrodynamics (interaction between light confined in a reflective cavity and atoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bdf1a7fc-95a6-4814-84c0-f9c83a8d5ed8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'With decay')"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 118080 characters]",
      "text/plain": [
       "<Figure size 1800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import cudaq\n",
    "from cudaq import boson, 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",
    "# System dimensions: atom (2-level system) and cavity (10-level system)\n",
    "dimensions = {0: 2, 1: 10}\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",
    "# This is the so-called Jaynes-Cummings model:\n",
    "# https://en.wikipedia.org/wiki/Jaynes%E2%80%93Cummings_model\n",
    "hamiltonian = 2 * np.pi * boson.number(1) + 2 * np.pi * boson.number(\n",
    "    0) + 2 * np.pi * 0.25 * (sm * a_dag + sm_dag * a)\n",
    "\n",
    "# Initial state of the system\n",
    "# Atom 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 5 photons initially\n",
    "cavity_state = cp.zeros((10, 10), dtype=cp.complex128)\n",
    "cavity_state[5][5] = 1.0\n",
    "rho0 = cudaq.State.from_data(cp.kron(cavity_state, qubit_state))\n",
    "\n",
    "steps = np.linspace(0, 10, 201)\n",
    "schedule = Schedule(steps, [\"time\"])\n",
    "\n",
    "# First, evolve the system without any collapse operators (ideal).\n",
    "evolution_result = cudaq.evolve(hamiltonian,\n",
    "                                dimensions,\n",
    "                                schedule,\n",
    "                                rho0,\n",
    "                                observables=[boson.number(1),\n",
    "                                             boson.number(0)],\n",
    "                                collapse_operators=[],\n",
    "                                store_intermediate_results=cudaq.IntermediateResultSave.EXPECTATION_VALUE,\n",
    "                                integrator=ScipyZvodeIntegrator())\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=[np.sqrt(0.1) * 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",
    "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((\"Cavity Photon Number\", \"Atom Excitation Probability\"))\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((\"Cavity Photon Number\", \"Atom Excitation Probability\"))\n",
    "plt.title(\"With decay\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "371030f1-8931-4765-95b8-3f7558f946fa",
   "metadata": {},
   "source": [
    "## Cross Resonance\n",
    "\n",
    "This example simulates cross-resonance interactions between superconducting qubits. Cross-resonance interaction is key to implementing two-qubit conditional gates, e.g., the CNOT gate. The example comes from a [simple all-microwave entangling gate for fixed-frequency superconducting qubits](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.107.080502)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "59af9caf-beca-4169-88d3-409668199148",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7fc0f29528f0>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 225224 characters]",
      "text/plain": [
       "<Figure size 1800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import cudaq\n",
    "from cudaq import spin, 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",
    "# Detuning between two qubits\n",
    "delta = 100 * 2 * np.pi\n",
    "# Static coupling between qubits\n",
    "J = 7 * 2 * np.pi\n",
    "# spurious electromagnetic `crosstalk` due to stray electromagnetic coupling in the device circuit and package\n",
    "# see (Physical Review Letters 107, 080502)\n",
    "m_12 = 0.2\n",
    "# Drive strength\n",
    "Omega = 20 * 2 * np.pi\n",
    "\n",
    "# Qubit Hamiltonian (in the rotating frame w.r.t. the target qubit)\n",
    "hamiltonian = delta / 2 * spin.z(0) + J * (\n",
    "    spin.minus(1) * spin.plus(0) +\n",
    "    spin.plus(1) * spin.minus(0)) + Omega * spin.x(0) + m_12 * Omega * spin.x(1)\n",
    "\n",
    "# Dimensions of sub-system\n",
    "dimensions = {0: 2, 1: 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",
    "# Two initial states: |00> and |10>.\n",
    "# We show the 'conditional' evolution when controlled qubit is in |1> state.\n",
    "psi_00 = cudaq.State.from_data(\n",
    "    cp.array([1.0, 0.0, 0.0, 0.0], dtype=cp.complex128))\n",
    "psi_10 = cudaq.State.from_data(\n",
    "    cp.array([0.0, 1.0, 0.0, 0.0], dtype=cp.complex128))\n",
    "\n",
    "# Schedule of time steps.\n",
    "steps = np.linspace(0.0, 1.0, 1001)\n",
    "schedule = Schedule(steps, [\"time\"])\n",
    "\n",
    "# Run the simulation.\n",
    "# Control bit = 0\n",
    "evolution_result_00 = cudaq.evolve(hamiltonian,\n",
    "                                   dimensions,\n",
    "                                   schedule,\n",
    "                                   psi_00,\n",
    "                                   observables=[\n",
    "                                       spin.x(0),\n",
    "                                       spin.y(0),\n",
    "                                       spin.z(0),\n",
    "                                       spin.x(1),\n",
    "                                       spin.y(1),\n",
    "                                       spin.z(1)\n",
    "                                   ],\n",
    "                                   collapse_operators=[],\n",
    "                                   store_intermediate_results=cudaq.IntermediateResultSave.EXPECTATION_VALUE,\n",
    "                                   integrator=ScipyZvodeIntegrator())\n",
    "\n",
    "# Control bit = 1\n",
    "evolution_result_10 = cudaq.evolve(hamiltonian,\n",
    "                                   dimensions,\n",
    "                                   schedule,\n",
    "                                   psi_10,\n",
    "                                   observables=[\n",
    "                                       spin.x(0),\n",
    "                                       spin.y(0),\n",
    "                                       spin.z(0),\n",
    "                                       spin.x(1),\n",
    "                                       spin.y(1),\n",
    "                                       spin.z(1)\n",
    "                                   ],\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",
    "results_00 = [\n",
    "    get_result(0, evolution_result_00),\n",
    "    get_result(1, evolution_result_00),\n",
    "    get_result(2, evolution_result_00),\n",
    "    get_result(3, evolution_result_00),\n",
    "    get_result(4, evolution_result_00),\n",
    "    get_result(5, evolution_result_00)\n",
    "]\n",
    "results_10 = [\n",
    "    get_result(0, evolution_result_10),\n",
    "    get_result(1, evolution_result_10),\n",
    "    get_result(2, evolution_result_10),\n",
    "    get_result(3, evolution_result_10),\n",
    "    get_result(4, evolution_result_10),\n",
    "    get_result(5, evolution_result_10)\n",
    "]\n",
    "\n",
    "# The changes in recession frequencies of the target qubit when control qubit is in |1> state allow us to implement two-qubit conditional gates.\n",
    "fig = plt.figure(figsize=(18, 6))\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(steps, results_00[5])\n",
    "plt.plot(steps, results_10[5])\n",
    "plt.ylabel(r\"$\\langle Z_2\\rangle$\")\n",
    "plt.xlabel(\"Time\")\n",
    "plt.legend((r\"$|\\psi_0\\rangle=|00\\rangle$\", r\"$|\\psi_0\\rangle=|10\\rangle$\"))\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(steps, results_00[4])\n",
    "plt.plot(steps, results_10[4])\n",
    "plt.ylabel(r\"$\\langle Y_2\\rangle$\")\n",
    "plt.xlabel(\"Time\")\n",
    "plt.legend((r\"$|\\psi_0\\rangle=|00\\rangle$\", r\"$|\\psi_0\\rangle=|10\\rangle$\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc31f7a6-0fdc-4f95-a3ed-0cd2fd74fbd9",
   "metadata": {},
   "source": [
    "## Transmon Resonator\n",
    "\n",
    "This example demonstrates a simulation of a superconducting transmon qubit coupled to a resonator (i.e., cavity). Based on refs from \"[Charge-insensitive qubit design derived from the Cooper pair box](https://journals.aps.org/pra/abstract/10.1103/PhysRevA.76.042319)\" and [https://github.com/jrjohansson/qutip-lectures/blob/master/Lecture-10-cQED-dispersive-regime.ipynb](https://github.com/jrjohansson/qutip-lectures/blob/master/Lecture-10-cQED-dispersive-regime.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e328ebc4-209e-439a-9314-e336535a1d63",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'Resonator Quadrature')"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 69224 characters]",
      "text/plain": [
       "<Figure size 1800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import cudaq\n",
    "from cudaq import operators, spin, operators, 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",
    "# Number of cavity photons\n",
    "N = 20\n",
    "\n",
    "# System dimensions: transmon + cavity\n",
    "dimensions = {0: 2, 1: N}\n",
    "\n",
    "# See III.B of PRA 76, 042319\n",
    "# System parameters\n",
    "# Unit: GHz\n",
    "omega_01 = 3.0 * 2 * np.pi  # transmon qubit frequency\n",
    "omega_r = 2.0 * 2 * np.pi  # resonator frequency\n",
    "# Dispersive shift\n",
    "chi_01 = 0.025 * 2 * np.pi\n",
    "chi_12 = 0.0\n",
    "\n",
    "omega_01_prime = omega_01 + chi_01\n",
    "omega_r_prime = omega_r - chi_12 / 2.0\n",
    "chi = chi_01 - chi_12 / 2.0\n",
    "\n",
    "# System Hamiltonian\n",
    "hamiltonian = 0.5 * omega_01_prime * spin.z(0) + (\n",
    "    omega_r_prime + chi * spin.z(0)) * operators.number(1)\n",
    "\n",
    "# Initial state of the system\n",
    "# Transmon in a superposition state\n",
    "transmon_state = cp.array([1. / np.sqrt(2.), 1. / np.sqrt(2.)],\n",
    "                          dtype=cp.complex128)\n",
    "\n",
    "\n",
    "# Helper to create a coherent state in Fock basis truncated at `num_levels`.\n",
    "# Note: There are a couple of ways of generating a coherent state,\n",
    "# e.g., see https://qutip.readthedocs.io/en/v5.0.3/apidoc/functions.html#qutip.core.states.coherent\n",
    "# or https://en.wikipedia.org/wiki/Coherent_state\n",
    "# Here, in this example, we use a the formula: `|alpha> = D(alpha)|0>`,\n",
    "# i.e., apply the displacement operator on a zero (or vacuum) state to compute the corresponding coherent state.\n",
    "def coherent_state(num_levels, amplitude):\n",
    "    displace_mat = operators.displace(0).to_matrix({0: num_levels},\n",
    "                                                   displacement=amplitude)\n",
    "    # `D(alpha)|0>` is the first column of `D(alpha)` matrix\n",
    "    return cp.array(np.transpose(displace_mat)[0])\n",
    "\n",
    "\n",
    "# Cavity in a coherent state\n",
    "cavity_state = coherent_state(N, 2.0)\n",
    "psi0 = cudaq.State.from_data(cp.kron(transmon_state, cavity_state))\n",
    "\n",
    "steps = np.linspace(0, 250, 1000)\n",
    "schedule = Schedule(steps, [\"time\"])\n",
    "\n",
    "# Evolve the system\n",
    "evolution_result = cudaq.evolve(hamiltonian,\n",
    "                                dimensions,\n",
    "                                schedule,\n",
    "                                psi0,\n",
    "                                observables=[\n",
    "                                    operators.number(1),\n",
    "                                    operators.number(0),\n",
    "                                    operators.position(1),\n",
    "                                    operators.position(0)\n",
    "                                ],\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",
    "count_results = [\n",
    "    get_result(0, evolution_result),\n",
    "    get_result(1, evolution_result)\n",
    "]\n",
    "\n",
    "quadrature_results = [\n",
    "    get_result(2, evolution_result),\n",
    "    get_result(3, evolution_result)\n",
    "]\n",
    "\n",
    "fig = plt.figure(figsize=(18, 6))\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(steps, count_results[0])\n",
    "plt.plot(steps, count_results[1])\n",
    "plt.ylabel(\"n\")\n",
    "plt.xlabel(\"Time [ns]\")\n",
    "plt.legend((\"Cavity Photon Number\", \"Transmon Excitation Probability\"))\n",
    "plt.title(\"Excitation Numbers\")\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(steps, quadrature_results[0])\n",
    "plt.ylabel(\"x\")\n",
    "plt.xlabel(\"Time [ns]\")\n",
    "plt.legend((\"cavity\"))\n",
    "plt.title(\"Resonator Quadrature\")"
   ]
  }
 ],
 "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
}
