{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Simulating an open two-level quantum system\n",
    "\n",
    "This tutorial provides a step-by-step guide to simulating the dynamics of a simple two-level quantum system subject to dissipation using the `qruise-toolset`. We will begin by simulating the system on its own, and then explore its behaviour when coupled to a leaky single-mode cavity.\n",
    "\n",
    "## 1. Simulating a two-level system\n",
    "The two-level system (TLS) we are considering is described by the following Hamiltonian:\n",
    "\n",
    "$$\n",
    "H = 2 \\pi \\beta \\sigma_x , \\tag{1}\n",
    "$$\n",
    "\n",
    "where $\\beta$ is the tunnelling rate, describing the coupling strength that drives transitions between the TLS's ground and excited states. $\\sigma_x$ is the Pauli-X operator, which we'll import from the `QuTip` library.\n",
    "\n",
    "While closed quantum systems are governed by the Schrödinger equation, open systems require the Lindblad master equation, which account for both unitary dynamics and dissipative processes. This is given by\n",
    "\n",
    "$$\n",
    "\\frac{\\partial \\rho(t)}{\\partial t} = -i \\left[ H(t), \\rho(t) \\right] + \\sum_i \\gamma_i \\left( L_i \\rho(t) L_i^\\dagger - \\frac{1}{2} \\left\\{ L_i^\\dagger L_i, \\rho(t) \\right\\} \\right) , \\tag{2}\n",
    "$$\n",
    "\n",
    "where $\\rho(t)$ is the density matrix of the system at time $t$, $\\gamma_i$ are decay rates (associated with decoherence or dissipation), and $L_i$ are the corresponding collapse operators. In our system, we'll use a single collapse operator to describe dissipation, with $\\gamma_i = \\gamma = 0.05 \\text{ s}^{-1}$ and $L_i = L =\\sigma_x$.\n",
    "\n",
    "---\n",
    "\n",
    "**Tip:** Make sure that your environment is correctly set to handle `float64` precision by setting `JAX_ENABLE_X64=True` or add the following codeblock to your scripts' preamble: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "\n",
    "jax.config.update(\"jax_enable_x64\", True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "We can start by defining the time parameters and decay rates for the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax.numpy as jnp\n",
    "\n",
    "t0 = 0.0  # start time of simulation\n",
    "tfinal = 10.0  # end time of simulation\n",
    "N = 200  # number of time points\n",
    "ts = jnp.linspace(t0, tfinal, N)  # time array\n",
    "\n",
    "beta = 0.1  # tunnelling rate\n",
    "gamma = 0.05  # dissipation rate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "We then create our Hamiltonian and collapse operator (including the dissipation rate)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import qruise.toolset as qr\n",
    "import qutip as qt\n",
    "\n",
    "# define Hamiltonian and collapse operator\n",
    "H = qr.Hamiltonian(2 * jnp.pi * beta * qt.sigmax())\n",
    "c_op = jnp.sqrt(gamma) * qt.sigmax()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "To solve the Lindblad equation, we need to specify the equations and parameters that govern the system. In the `qruise-toolset`, we combine these using a `Problem`, which is instantiated with:\n",
    "\n",
    "- the Hamiltonian (`H`)\n",
    "- the initial qubit state (`pho0`, here the ground state)\n",
    "- the time range of the simulation (`t0` to `tfinal`).\n",
    "\n",
    "We can then use Qruise's `ODESolver` with the collapse operator (`c_op`) to solve the Lindblad master equation and calculate the wavefunction of the system at each timestamp. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffrax import Tsit5\n",
    "\n",
    "# define initial qubit state (in matrix form)\n",
    "rho0 = qt.basis(2, 0)\n",
    "rho0 = rho0 @ rho0.dag()  # density matrix\n",
    "\n",
    "# define Problem\n",
    "prob = qr.Problem(H, rho0, {}, (t0, tfinal))\n",
    "\n",
    "# use ODEsolver to solve Lindblad master equation\n",
    "solver = qr.ODESolver(\n",
    "    Tsit5(), dt0=None, saveat={\"ts\": ts}, pid_args={\"atol\": 1e-8, \"rtol\": 1e-6}\n",
    ")\n",
    "solver.set_system(prob.lindblad([c_op]))\n",
    "_, rhos = solver.evolve(*prob.problem())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "We can now calculate and plot the expectation values of $\\sigma_z$ and $\\sigma_y$ ($\\langle \\sigma_z \\rangle$ and $\\langle \\sigma_y \\rangle$, respectively). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate expectation values of sigmaz and sigma y\n",
    "sigmaz_exp = qr.op_expectation(qt.sigmaz(), rhos)\n",
    "sigmay_exp = qr.op_expectation(qt.sigmay(), rhos)\n",
    "\n",
    "# plot expectation values against time\n",
    "canvas = qr.PlotUtil(\n",
    "    x_axis_label=\"Time [a.u.]\", y_axis_label=\"Operator expectation value\", notebook=True\n",
    ")\n",
    "canvas.plot(ts, sigmaz_exp, labels=[\"<sigma_z>\"])\n",
    "canvas.plot(ts, sigmay_exp, labels=[\"<sigma_y>\"])\n",
    "canvas.show_canvas()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "Nice! You can see that $\\langle \\sigma_z \\rangle$ and  $\\langle \\sigma_y \\rangle$ exhibit periodic oscillations driven by tunnelling (governed by $\\beta$) between the $|0\\rangle$ and $|1\\rangle$ states, with their amplitudes decaying over time due to dissipation (governed by $\\gamma$)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## 2. Simulating a two-level system coupled to a cavity\n",
    "\n",
    "Now let's take this one step further by considering our TLS as a two-level atom that we'll couple to a leaky single-mode cavity via a dipole-type interaction. This supports a coherent exchange of energy between the TLS and the cavity.\n",
    "\n",
    "The Hamiltonian for this system is:\n",
    "\n",
    "$$\n",
    "H = \\underbrace{2 \\pi \\sigma^+ \\sigma^- \\vphantom{\\big|}}_{\\text{TLS energy}}\n",
    "+ \\underbrace{2 \\pi a^\\dagger a \\vphantom{\\big|}}_{\\text{cavity energy}} \n",
    "+ \\underbrace{2 \\pi g (a \\sigma^++ a^\\dagger \\sigma^-) \\vphantom{\\big|}}_{\\text{TLS-cavity interaction}} , \\tag{3}\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "- $\\sigma^+$ and $\\sigma^-$ are the raising and lowering operators for the TLS\n",
    "- $a^\\dagger$ and $a$ are the creation and annihilation operators for the single-mode cavity\n",
    "- $g$ is the TLS-cavity coupling strength.\n",
    "\n",
    "The Lindblad master equation (Equation 2) stays the same, except we use the new Hamiltonian and a new collapse operator. Here, we have $L_i=L=a$ (instead of $L=\\sigma_x$), and $\\gamma_i = \\kappa$, which represents the leakage rate from the cavity. \n",
    "\n",
    "We'll use the same time parameters as in the previous simulation, but we need to define the system parameters $g$ and $\\kappa$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "g = 0.25  # TLS-cavity coupling strength\n",
    "kappa = 0.1  # cavity leakage rate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "Then we can define our operators and Hamiltonian."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define TLS and cavity operators\n",
    "\n",
    "atomic = qt.tensor(qt.destroy(2), qt.qeye(10))  # TLS lowering operator\n",
    "cavity = qt.tensor(qt.qeye(2), qt.destroy(10))  # cavity annihilation operator\n",
    "\n",
    "# define Hamiltonian and collapse operator\n",
    "H = qr.Hamiltonian(\n",
    "    2 * jnp.pi * atomic.dag() * atomic\n",
    "    + 2 * jnp.pi * cavity.dag() * cavity\n",
    "    + 2 * jnp.pi * g * (cavity * atomic.dag() + cavity.dag() * atomic)\n",
    ")\n",
    "c_op = jnp.sqrt(kappa) * cavity"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "**Note:** The Hilbert dimension of the TLS is 2 ($|0\\rangle$ and $|1\\rangle$) and we'll set it to 10 for the cavity (up to 10 photons).\n",
    "\n",
    "---\n",
    "\n",
    "Now we just need to define our initial system state (`rho0`), which we'll take as the ground state ($|0\\rangle$) for the TLS and as a 5-photon Fock state ($|5\\rangle$) for the cavity. We can then define the `Problem` and solve the Lindblad master equation using Qruise's `ODESolver` as we did in the previous section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define inital system state (in matrix form)\n",
    "rho0 = qt.tensor(qt.fock(2, 0), qt.fock(10, 5))\n",
    "rho0 = rho0 @ rho0.dag()\n",
    "\n",
    "# define Problem\n",
    "prob = qr.Problem(H, rho0, {}, (t0, tfinal))\n",
    "\n",
    "# use ODEsolver to solve Lindblad master equation\n",
    "solver = qr.ODESolver(\n",
    "    Tsit5(), dt0=None, saveat={\"ts\": ts}, pid_args={\"atol\": 1e-8, \"rtol\": 1e-6}\n",
    ")\n",
    "solver.set_system(prob.lindblad([c_op]))\n",
    "_, rhos = solver.evolve(*prob.problem())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "We can now calculate and plot the expectation values of $\\sigma^+ \\sigma^-$ and $a^\\dagger a$ ($\\langle \\sigma^+ \\sigma^- \\rangle$ and $\\langle a^\\dagger a \\rangle$, respectively). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# expectation values of sigma^+ sigma^- and a^\\dagger a\n",
    "atomic_dag_atomic_exp = qr.op_expectation(atomic.dag() * atomic, rhos)\n",
    "cavity_dag_cavity_exp = qr.op_expectation(cavity.dag() * cavity, rhos)\n",
    "\n",
    "# plotting the expectation values\n",
    "canvas.init_canvas(\n",
    "    x_axis_label=\"Time [a.u.]\", y_axis_label=\"Operator expectation value\"\n",
    ")\n",
    "canvas.plot(ts, cavity_dag_cavity_exp, labels=[\"cavity photon number\"])\n",
    "canvas.plot(ts, atomic_dag_atomic_exp, labels=[\"atomic excitation probability\"])\n",
    "canvas.show_canvas()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "Great! You can see we still see periodic oscillation, driven by the coupling between the TLS and the cavity ($g$), with decaying amplitudes due to dissipation ($\\gamma$). \n",
    "\n",
    "Congratulations! You now know how to simulate an open two-level quantum system."
   ]
  }
 ],
 "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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}