{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# QuTiP example: Floquet formalism for dynamics and steadystates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "J.R. Johansson and P.D. Nation\n",
    "\n",
    "For more information about QuTiP see [http://qutip.org](http://qutip.org)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qutip import *\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Floquet modes and quasi energies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Find the floquet modes and quasi energies for a driven system and plot the floquet states/quasienergies for one period of the driving.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hamiltonian_t(t, args):\n",
    "    \"\"\" evaluate the hamiltonian at time t. \"\"\"\n",
    "    H0 = args['H0']\n",
    "    H1 = args['H1']\n",
    "    w  = args['w']\n",
    "\n",
    "    return H0 + sin(w * t) * H1\n",
    "\n",
    "def H1_coeff_t(t, args):\n",
    "    return sin(args['w'] * t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def qubit_integrate(delta, eps0_vec, A, omega, gamma1, gamma2, psi0, T, option):\n",
    "\n",
    "    # Hamiltonian\n",
    "    sx = sigmax()\n",
    "    sz = sigmaz()\n",
    "    sm = destroy(2)\n",
    "\n",
    "    # collapse operators\n",
    "    c_op_list = []\n",
    "\n",
    "    n_th = 0.0 # zero temperature\n",
    "\n",
    "    # relaxation\n",
    "    rate = gamma1 * (1 + n_th)\n",
    "    if rate > 0.0:\n",
    "        c_op_list.append(sqrt(rate) * sm)\n",
    "\n",
    "    # excitation\n",
    "    rate = gamma1 * n_th\n",
    "    if rate > 0.0:\n",
    "        c_op_list.append(sqrt(rate) * sm.dag())\n",
    "\n",
    "    # dephasing \n",
    "    rate = gamma2\n",
    "    if rate > 0.0:\n",
    "        c_op_list.append(sqrt(rate) * sz)\n",
    "\n",
    "\n",
    "    quasi_energies = np.zeros((len(eps0_vec), 2))\n",
    "    f_gnd_prob     = np.zeros((len(eps0_vec), 2))\n",
    "    wf_gnd_prob    = np.zeros((len(eps0_vec), 2))\n",
    "\n",
    "    for idx, eps0 in enumerate(eps0_vec):\n",
    "\n",
    "        H0 = - delta/2.0 * sx - eps0/2.0 * sz\n",
    "        H1 = A/2.0 * sz\n",
    "       \n",
    "        # H = H0 + H1 * sin(w * t) in the 'list-string' format\n",
    "        H = [H0, [H1, 'sin(w * t)']]\n",
    "        Hargs = {'w': omega}\n",
    "            \n",
    "        # find the floquet modes\n",
    "        f_modes,f_energies = floquet_modes(H, T, Hargs)\n",
    "\n",
    "        quasi_energies[idx,:] = f_energies\n",
    "\n",
    "        f_gnd_prob[idx, 0] = expect(sm.dag() * sm, f_modes[0])\n",
    "        f_gnd_prob[idx, 1] = expect(sm.dag() * sm, f_modes[1])\n",
    "\n",
    "        f_states = floquet_states_t(f_modes, f_energies, 0, H, T, Hargs)\n",
    "\n",
    "        wf_gnd_prob[idx, 0] = expect(sm.dag() * sm, f_states[0])\n",
    "        wf_gnd_prob[idx, 1] = expect(sm.dag() * sm, f_states[1])\n",
    "        \n",
    "    return quasi_energies, f_gnd_prob, wf_gnd_prob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta = 0.2 * 2 * np.pi  # qubit sigma_x coefficient\n",
    "eps0  = 0.5 * 2 * np.pi  # qubit sigma_z coefficient\n",
    "gamma1 = 0.0        # relaxation rate\n",
    "gamma2 = 0.0         # dephasing  rate\n",
    "A      = 2.0 * 2 * np.pi \n",
    "psi0   = basis(2,0)    # initial state\n",
    "omega  = 1.0 * 2 * np.pi # driving frequency\n",
    "T      = (2*np.pi)/omega  # driving period\n",
    "\n",
    "param  = np.linspace(-5.0, 5.0, 200) * 2 * np.pi \n",
    "\n",
    "eps0 = param"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = time.time()\n",
    "q_energies, f_gnd_prob, wf_gnd_prob = qubit_integrate(delta, eps0, A, omega, gamma1, gamma2, psi0, T, \"dynamics\")\n",
    "print('dynamics: time elapsed = ' + str(time.time() - start_time))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(param, np.real(q_energies[:,0]) / delta, 'b',\n",
    "        param, np.real(q_energies[:,1]) / delta, 'r')\n",
    "ax.set_xlabel('A or e')\n",
    "ax.set_ylabel('Quasienergy')\n",
    "ax.set_title('Floquet quasienergies');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(param, np.real(f_gnd_prob[:,0]), 'b', \n",
    "        param, np.real(f_gnd_prob[:,1]), 'r')\n",
    "ax.set_xlabel('A or e')\n",
    "ax.set_ylabel('Occ. prob.')\n",
    "ax.set_title('Floquet modes excitation probability');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(param, np.real(wf_gnd_prob[:,0]), 'b',\n",
    "        param, np.real(wf_gnd_prob[:,1]), 'r')\n",
    "ax.set_xlabel('A or e')\n",
    "ax.set_ylabel('Occ. prob.')\n",
    "ax.set_title('Floquet states excitation probability');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Floquet modes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hamiltonian_t(t, args):\n",
    "    \"\"\" evaluate the hamiltonian at time t. \"\"\"\n",
    "    H0 = args[0]\n",
    "    H1 = args[1]\n",
    "    w  = args[2]\n",
    "\n",
    "    return H0 + cos(w * t) * H1\n",
    "\n",
    "def H1coeff_t(t, args):\n",
    "    w  = args['w']\n",
    "    return sin(w * t)           "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def qubit_integrate(delta, eps0, A, omega, psi0, tlist):\n",
    "\n",
    "    # Hamiltonian\n",
    "    sx = sigmax()\n",
    "    sz = sigmaz()\n",
    "    sm = destroy(2)\n",
    "\n",
    "    H0 = - delta/2.0 * sx - eps0/2.0 * sz\n",
    "    H1 =   A/2.0 * sz\n",
    "        \n",
    "    #H_args = (H0, H1, omega)\n",
    "    H_args = {'w': omega}\n",
    "    H = [H0, [H1, 'sin(w * t)']]\n",
    "    #H = [H0, [H1, H1coeff_t]]\n",
    "    \n",
    "    # find the propagator for one driving period\n",
    "    T = 2*np.pi / omega\n",
    "       \n",
    "    f_modes_0,f_energies = floquet_modes(H, T, H_args)\n",
    "\n",
    "    p_ex_0 = np.zeros(shape(tlist))\n",
    "    p_ex_1 = np.zeros(shape(tlist))\n",
    "\n",
    "    p_00 = np.zeros(shape(tlist), dtype=complex)\n",
    "    p_01 = np.zeros(shape(tlist), dtype=complex)    \n",
    "    p_10 = np.zeros(shape(tlist), dtype=complex)\n",
    "    p_11 = np.zeros(shape(tlist), dtype=complex)    \n",
    "    \n",
    "    e_0 = np.zeros(shape(tlist))\n",
    "    e_1 = np.zeros(shape(tlist))\n",
    "         \n",
    "    f_modes_table_t = floquet_modes_table(f_modes_0, f_energies, tlist, H, T, H_args) \n",
    "\n",
    "    for idx, t in enumerate(tlist):\n",
    "        #f_modes_t = floquet_modes_t(f_modes_0, f_energies, t, H, T, H_args) \n",
    "        f_modes_t = floquet_modes_t_lookup(f_modes_table_t, t, T) \n",
    "\n",
    "        p_ex_0[idx] = expect(sm.dag() * sm, f_modes_t[0])\n",
    "        p_ex_1[idx] = expect(sm.dag() * sm, f_modes_t[1])\n",
    "\n",
    "        p_00[idx] = f_modes_t[0].full()[0][0]\n",
    "        p_01[idx] = f_modes_t[0].full()[1][0]\n",
    "        p_10[idx] = f_modes_t[1].full()[0][0]\n",
    "        p_11[idx] = f_modes_t[1].full()[1][0]\n",
    "\n",
    "        #evals = hamiltonian_t(t, H_args).eigenenergies()\n",
    "        evals = qobj_list_evaluate(H, t, H_args).eigenenergies()\n",
    "        e_0[idx] = min(np.real(evals))\n",
    "        e_1[idx] = max(np.real(evals))\n",
    "        \n",
    "    return p_ex_0, p_ex_1, e_0, e_1, f_energies, p_00, p_01, p_10, p_11"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta = 0.2 * 2 * np.pi  # qubit sigma_x coefficient\n",
    "eps0  = 1.0 * 2 * np.pi  # qubit sigma_z coefficient\n",
    "A     = 2.5 * 2 * np.pi  # sweep rate\n",
    "psi0   = basis(2,0)   # initial state\n",
    "omega  = 1.0 * 2 * np.pi # driving frequency\n",
    "T      = (2*np.pi)/omega # driving period\n",
    "\n",
    "tlist = np.linspace(0.0, T, 101)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = time.time()\n",
    "p_ex_0, p_ex_1, e_0, e_1, f_e, p_00, p_01, p_10, p_11 = qubit_integrate(delta, eps0, A, omega, psi0, tlist)\n",
    "print('dynamics: time elapsed = ' + str(time.time() - start_time))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=[8,10])\n",
    "\n",
    "ax1.plot(tlist, np.real(p_ex_0), 'b', tlist, np.real(p_ex_1), 'r')\n",
    "ax1.set_xlabel('Time ($T$)')\n",
    "ax1.set_ylabel('Excitation probabilities')\n",
    "ax1.set_title('Floquet modes')\n",
    "ax1.legend((\"Floquet mode 1\", \"Floquet mode 2\"))\n",
    "\n",
    "ax2.plot(tlist, np.real(e_0), 'c', tlist, np.real(e_1), 'm')\n",
    "ax2.plot(tlist, np.ones(shape(tlist)) * f_e[0], 'b', tlist, np.ones(shape(tlist)) * f_e[1], 'r')\n",
    "ax2.set_xlabel('Time ($T$)')\n",
    "ax2.set_ylabel('Energy [GHz]')\n",
    "ax2.set_title('Eigen- and quasi-energies')\n",
    "ax2.legend((\"Ground state\", \"Excited state\", \"Quasienergy 1\", \"Quasienergy 2\"));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(3, 1, figsize=[8,12])\n",
    "\n",
    "axes[0].plot(tlist, np.real(p_00), 'b', tlist, np.real(p_01), 'r')\n",
    "axes[0].plot(tlist, np.real(p_10), 'c', tlist, np.real(p_11), 'm')\n",
    "axes[0].set_xlabel('Time ($T$)')\n",
    "axes[0].set_ylabel('real')\n",
    "axes[0].set_title('Floquet modes')\n",
    "axes[0].legend((\"FM1 - gnd\", \"FM1 - exc\", \"FM2 - gnd\", \"FM2 - exc\"))\n",
    "\n",
    "axes[1].plot(tlist, np.imag(p_00), 'b', tlist, np.imag(p_01), 'r')\n",
    "axes[1].plot(tlist, np.imag(p_10), 'c', tlist, np.imag(p_11), 'm')\n",
    "axes[1].set_xlabel('Time ($T$)')\n",
    "axes[1].set_ylabel('imag')\n",
    "axes[1].legend((\"FM1 - gnd\", \"FM1 - exc\", \"FM2 - gnd\", \"FM2 - exc\"))\n",
    "\n",
    "axes[2].plot(tlist, abs(p_00), 'b', tlist, abs(p_01), 'r.')\n",
    "axes[2].plot(tlist, abs(p_10), 'c', tlist, abs(p_11), 'm.')\n",
    "axes[2].set_xlabel('Time ($T$)')\n",
    "axes[2].set_ylabel('abs')\n",
    "axes[2].legend((\"FM1 - gnd\", \"FM1 - exc\", \"FM2 - gnd\", \"FM2 - exc\"));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Floquet evolution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Find the floquet modes and quasi energies for a driven system and evolve the wavefunction \"stroboscopically\", i.e., only by evaluating at time mupliples of the driving period t = n * T for integer n.\n",
    "\n",
    "The system is a strongly driven two-level atom.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hamiltonian_t(t, args):\n",
    "    \"\"\" evaluate the hamiltonian at time t. \"\"\"\n",
    "    H0 = args[0]\n",
    "    H1 = args[1]\n",
    "    w  = args[2]\n",
    "    return H0 + np.sin(w * t) * H1\n",
    "    \n",
    "def H1coeff_t(t, args):\n",
    "    w  = args['w']\n",
    "    return np.sin(w * t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def qubit_integrate(delta, eps0, A, omega, psi0, tlist, T, option):\n",
    "\n",
    "    # Hamiltonian\n",
    "    sx = sigmax()\n",
    "    sz = sigmaz()\n",
    "    sm = destroy(2)\n",
    "\n",
    "    H0 = - delta/2.0 * sx - eps0/2.0 * sz\n",
    "    H1 =   A/2.0 * sz        \n",
    "    #H_args = (H0, H1, omega)\n",
    "    H_args = {'w': omega}\n",
    "    H = [H0, [H1, 'sin(w * t)']]\n",
    "    #H = [H0, [H1, H1coeff_t]]  \n",
    "    \n",
    "\n",
    "    if option == \"floquet\":\n",
    "\n",
    "        # find the floquet modes for the time-dependent hamiltonian        \n",
    "        f_modes_0,f_energies = floquet_modes(H, T, H_args)\n",
    "\n",
    "        # decompose the inital state in the floquet modes (=floquet states at t=0)\n",
    "        f_coeff = floquet_state_decomposition(f_modes_0, f_energies, psi0)\n",
    "\n",
    "                    \n",
    "        # only evaluate the wavefunction at multiples of the driving period\n",
    "        # i.e. a \"stroboscopic\" evolution\n",
    "        N = max(tlist)/T\n",
    "        p_ex = np.zeros(N)                    \n",
    "        tlist = []\n",
    "        \n",
    "        # calculate the wavefunctions at times t=nT, for integer n, by using \n",
    "        # the floquet modes and quasienergies\n",
    "        for n in np.arange(N):            \n",
    "            psi_t = floquet_wavefunction_t(f_modes_0, f_energies, f_coeff, n*T, H, T, H_args)            \n",
    "            p_ex[n] = expect(sm.dag() * sm, psi_t)\n",
    "            tlist.append(n*T)   \n",
    "    \n",
    "    else:\n",
    "    \n",
    "        # for reference: evolve and calculate expectation using the full ode solver\n",
    "        #H_args = (H0, H1, omega)\n",
    "        #expt_list = mesolve(hamiltonian_t, psi0, tlist, [], [sm.dag() * sm], H_args)\n",
    "        output = mesolve(H, psi0, tlist, [], [sm.dag() * sm], H_args)\n",
    "        p_ex = output.expect[0]\n",
    "        \n",
    "    return tlist, p_ex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta = 0.2 * 2 * np.pi  # qubit sigma_x coefficient\n",
    "eps0  = 0.1 * 2 * np.pi  # qubit sigma_z coefficient\n",
    "A     = 1.0 * 2 * np.pi  # driving amplitude\n",
    "psi0   = basis(2,0)      # initial state\n",
    "omega  = 1.0 * 2 * np.pi # driving frequency\n",
    "\n",
    "T      = (2*np.pi)/omega # driving period\n",
    "tlist  = np.linspace(0.0, 25 * T, 500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = time.time()\n",
    "tlist1, p_ex = qubit_integrate(delta, eps0, A, omega, psi0, tlist, T, \"dynamics\")\n",
    "print('dynamics: time elapsed = ' + str(time.time() - start_time))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = time.time()\n",
    "tlist2, f_ex = qubit_integrate(delta, eps0, A, omega, psi0, tlist, T, \"floquet\")\n",
    "print('floquet: time elapsed = ' + str(time.time() - start_time))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=[12, 6])\n",
    "ax.plot(tlist1, np.real(p_ex),   'b')\n",
    "ax.plot(tlist1, np.real(1-p_ex), 'r')\n",
    "ax.plot(tlist2, np.real(f_ex),   'bo', linewidth=2.0)\n",
    "ax.plot(tlist2, np.real(1-f_ex), 'ro', linewidth=2.0)\n",
    "\n",
    "ax.set_xlabel('Time')\n",
    "ax.set_ylabel('Probability')\n",
    "ax.set_title('Stroboscopic time-evolution with Floquet states')\n",
    "ax.legend((\"ode $P_1$\", \"ode $P_0$\", \"Floquet $P_1$\", \"Floquet $P_0$\"));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Floquet-Markov master equation: comparison with other solvers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gamma1 = 0.0015          # relaxation rate\n",
    "gamma2 = 0.0           # dephasing  rate\n",
    "\n",
    "def J_cb(omega):\n",
    "    \"\"\" Noise spectral density \"\"\"\n",
    "    #print \"evaluate J_cb for omega =\", omega\n",
    "    return 0.5 * gamma1 * omega/(2*np.pi)\n",
    "    \n",
    "def H1_coeff_t(t, args):\n",
    "    return np.sin(args['w'] * t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def qubit_integrate(delta, eps0, A, w, gamma1, gamma2, psi0, tlist):\n",
    "\n",
    "    # Hamiltonian\n",
    "    sx = sigmax()\n",
    "    sz = sigmaz()\n",
    "    sm = destroy(2)\n",
    "\n",
    "    H0 = - delta/2.0 * sx - eps0/2.0 * sz\n",
    "    H1 = - A * sx\n",
    "\n",
    "    args = {'w': w}\n",
    "    H = [H0, [H1, 'sin(w * t)']]\n",
    "\n",
    "\n",
    "    # --------------------------------------------------------------------------\n",
    "    # 1) time-dependent hamiltonian\n",
    "    # \n",
    "    #    \n",
    "    c_op_list = [np.sqrt(gamma1) * sx, np.sqrt(gamma2) * sz]\n",
    "\n",
    "    start_time = time.time()\n",
    "    output = mesolve(H, psi0, tlist, c_op_list, [sm.dag() * sm], args=args)      \n",
    "    expt_list1 = output.expect\n",
    "    print('Method 1: time elapsed = ' + str(time.time() - start_time))\n",
    "\n",
    "    # --------------------------------------------------------------------------\n",
    "    # 2) Constant hamiltonian\n",
    "    #\n",
    "    H_rwa = - delta/2.0 * sx - A * sx / 2\n",
    "    c_op_list = [np.sqrt(gamma1) * sx, np.sqrt(gamma2) * sz]\n",
    "    \n",
    "    start_time = time.time()\n",
    "    output = mesolve(H_rwa, psi0, tlist, c_op_list, [sm.dag() * sm])  \n",
    "    expt_list2 = output.expect\n",
    "    print('Method 2: time elapsed = ' + str(time.time() - start_time))\n",
    "\n",
    "\n",
    "    # --------------------------------------------------------------------------\n",
    "    # 3) Floquet unitary evolution\n",
    "    #\n",
    "    qutip.solver.config.reset()\n",
    "    \n",
    "    start_time = time.time()\n",
    "       \n",
    "    T = 2*np.pi / w       \n",
    "    f_modes_0,f_energies = floquet_modes(H, T, args)    \n",
    "    \n",
    "    # decompose the initial state vector in terms of the floquet modes (basis\n",
    "    # transformation). used to calculate psi_t below.\n",
    "    f_coeff = floquet_state_decomposition(f_modes_0, f_energies, psi0)\n",
    "    \n",
    "    # --------------------------------------------------------------------------\n",
    "    # 4) Floquet markov master equation dynamics\n",
    "    #       \n",
    "    kmax = 1\n",
    "    temp = 25e-3\n",
    "    w_th = temp * (1.38e-23 / 6.626e-34) * 2 * np.pi * 1e-9   \n",
    "    \n",
    "    f_modes_table_t = floquet_modes_table(f_modes_0, f_energies, np.linspace(0, T, 500+1), H, T, args) \n",
    "    \n",
    "    # calculate the rate-matrices for the floquet-markov master equation\n",
    "    Delta, X, Gamma, Amat = floquet_master_equation_rates(f_modes_0, f_energies, sx, \n",
    "                                                          H, T, args, J_cb, w_th, kmax, f_modes_table_t)\n",
    "   \n",
    "    # the floquet-markov master equation tensor\n",
    "    R = floquet_master_equation_tensor(Amat, f_energies)\n",
    "    \n",
    "    rho_list = floquet_markov_mesolve(R, f_modes_0, psi0, tlist, []).states\n",
    "\n",
    "    expt_list3 = np.zeros(shape(expt_list2), dtype=complex)\n",
    "    expt_list4 = np.zeros(shape(expt_list2), dtype=complex)\n",
    "    for idx, t in enumerate(tlist):            \n",
    "    \n",
    "        # unitary floquet evolution\n",
    "        psi_t = floquet_wavefunction_t(f_modes_0, f_energies, f_coeff, t, H, T, args)            \n",
    "        expt_list3[0][idx] = expect(sm.dag() * sm, psi_t) \n",
    "        \n",
    "        # the rho_list returned by the floquet master equation is defined in the\n",
    "        # floquet basis, so to transform it back to the computational basis\n",
    "        # before we calculate expectation values.\n",
    "        #f_modes_t = floquet_modes_t(f_modes_0, f_energies, t, H, T, args)\n",
    "        f_modes_t = floquet_modes_t_lookup(f_modes_table_t, t, T) \n",
    "        expt_list4[0][idx] = expect((sm.dag() * sm), rho_list[idx].transform(f_modes_t, True)) \n",
    "    \n",
    "    print('Method 3+4: time elapsed = ' + str(time.time() - start_time))\n",
    "\n",
    "    # calculate the steadystate density matrix according to the floquet-markov\n",
    "    # master equation\n",
    "    rho_ss_fb = floquet_master_equation_steadystate(H0, Amat) # in floquet basis\n",
    "    rho_ss_cb = rho_ss_fb.transform(f_modes_0, True)         # in computational basis\n",
    "    expt_list5 = np.ones(np.shape(expt_list2), dtype=complex) * expect(sm.dag() * sm, rho_ss_cb)\n",
    "    \n",
    "    return expt_list1[0], expt_list2[0], expt_list3[0], expt_list4[0], expt_list5[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta = 0.0 * 2 * np.pi   # qubit sigma_x coefficient\n",
    "eps0  = 1.0 * 2 * np.pi   # qubit sigma_z coefficient\n",
    "A     = 0.05 * 2 * np.pi  # driving amplitude (reduce to make the RWA more accurate)\n",
    "w     = 1.0 * 2 * np.pi   # driving frequency\n",
    "psi0 = (0.3*basis(2,0) + 0.7*basis(2,1)).unit()      # initial state\n",
    "\n",
    "tlist = np.linspace(0, 30.0, 500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p_ex1, p_ex2, p_ex3, p_ex4, p_ex5 = qubit_integrate(delta, eps0, A, w, gamma1, gamma2, psi0, tlist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(tlist, np.real(p_ex1), 'b', tlist, np.real(p_ex2), 'g-') # lindblad\n",
    "ax.plot(tlist, np.real(p_ex3), 'r', tlist, np.real(p_ex4), 'm-', tlist, np.real(p_ex5), 'c-') # floquet markov\n",
    "ax.set_xlabel('Time')\n",
    "ax.set_ylabel('Occupation probability')\n",
    "ax.set_title('Comparison between time-dependent master equations')\n",
    "ax.legend((\"TD Hamiltonian, Std ME\", \"RWA Hamiltonian, Std ME\", \n",
    "        \"Unitary Floquet evol.\", \"Floquet-Markov ME\", \"F-M ME steady state\"));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Floquet-Markov master equation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def J_cb(omega):\n",
    "    \"\"\" Noise spectral density \"\"\"\n",
    "    return omega\n",
    "    \n",
    "def hamiltonian_t(t, args):\n",
    "    \"\"\" evaluate the hamiltonian at time t. \"\"\"\n",
    "    H0 = args[0]\n",
    "    H1 = args[1]\n",
    "    w  = args[2]\n",
    "\n",
    "    return H0 + np.cos(w * t) * H1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def qubit_integrate(delta, eps0, A, omega, psi0, tlist):\n",
    "\n",
    "    # Hamiltonian\n",
    "    sx = sigmax()\n",
    "    sz = sigmaz()\n",
    "    sm = destroy(2)\n",
    "\n",
    "    H0 = - delta/2.0 * sx - eps0/2.0 * sz\n",
    "    H1 =   A/2.0 * sz\n",
    "        \n",
    "    #H_args = (H0, H1, omega)\n",
    "    H_args = {'w': omega}\n",
    "    H = [H0, [H1, 'sin(w * t)']]\n",
    "    \n",
    "    # find the propagator for one driving period\n",
    "    T = 2*np.pi / omega\n",
    "       \n",
    "    f_modes_0,f_energies = floquet_modes(H, T, H_args)\n",
    "\n",
    "    c_op = sigmax()\n",
    "\n",
    "    kmax = 1\n",
    "\n",
    "    temp = 25e-3\n",
    "    w_th = temp * (1.38e-23 / 6.626e-34) * 2 * np.pi * 1e-9\n",
    "    \n",
    "    Delta, X, Gamma, A = floquet_master_equation_rates(f_modes_0, f_energies, c_op, H, T, H_args, J_cb, w_th, kmax)\n",
    "\n",
    "    k_idx = 0\n",
    "    for k in range(-kmax,kmax+1, 1):\n",
    "        print(\"X[\",k,\"] =\\n\", X[:,:,k_idx])\n",
    "        k_idx += 1\n",
    "\n",
    "    k_idx = 0\n",
    "    for k in range(-kmax,kmax+1, 1):\n",
    "        print(\"Delta[\",k,\"] =\\n\", Delta[:,:,k_idx])\n",
    "        k_idx += 1\n",
    "\n",
    "    k_idx = 0\n",
    "    for k in range(-kmax,kmax+1, 1):\n",
    "        print(\"Gamma[\",k,\"] =\\n\", Gamma[:,:,k_idx])\n",
    "        k_idx += 1\n",
    "        \n",
    "    print(\"A =\\n\", A)\n",
    "\n",
    "    rho_ss = floquet_master_equation_steadystate(H0, A)\n",
    "    \n",
    "    \n",
    "    R = floquet_master_equation_tensor(A, f_energies)\n",
    "    \n",
    "    print(\"Floquet-Markov master equation tensor\")\n",
    "    \n",
    "    print(\"R =\\n\", R)\n",
    "    \n",
    "    print(\"Floquet-Markov master equation steady state =\\n\", rho_ss)\n",
    "\n",
    "    p_ex_0 = np.zeros(shape(tlist))\n",
    "    p_ex_1 = np.zeros(shape(tlist))\n",
    "    \n",
    "    e_0 = np.zeros(shape(tlist))\n",
    "    e_1 = np.zeros(shape(tlist))\n",
    "        \n",
    "    f_modes_table_t = floquet_modes_table(f_modes_0, f_energies, tlist, H, T, H_args) \n",
    "    for idx, t in enumerate(tlist):\n",
    "        f_modes_t = floquet_modes_t_lookup(f_modes_table_t, t, T) \n",
    "\n",
    "        p_ex_0[idx] = expect(sm.dag() * sm, f_modes_t[0])\n",
    "        p_ex_1[idx] = expect(sm.dag() * sm, f_modes_t[1])\n",
    "\n",
    "        #evals = hamiltonian_t(t, H_args).eigenenergies()\n",
    "        evals = qobj_list_evaluate(H, t, H_args).eigenenergies()\n",
    "        e_0[idx] = min(np.real(evals))\n",
    "        e_1[idx] = max(np.real(evals))\n",
    "        \n",
    "    return p_ex_0, p_ex_1, e_0, e_1, f_energies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta = 0.2 * 2 * np.pi  # qubit sigma_x coefficient\n",
    "eps0  = 1.0 * 2 * np.pi  # qubit sigma_z coefficient\n",
    "A     = 2.5 * 2 * np.pi  # sweep rate\n",
    "psi0   = basis(2,0)      # initial state\n",
    "omega  = 1.0 * 2 * np.pi # driving frequency\n",
    "T      = (2*np.pi)/omega # driving period\n",
    "\n",
    "tlist = np.linspace(0.0, 1 * T, 101)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = time.time()\n",
    "p_ex_0, p_ex_1, e_0, e_1, f_e = qubit_integrate(delta, eps0, A, omega, psi0, tlist)\n",
    "print('dynamics: time elapsed = ' + str(time.time() - start_time))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, 1, figsize=[8,10])\n",
    "\n",
    "axes[0].plot(tlist, np.real(p_ex_0), 'b', tlist, np.real(p_ex_1), 'r')\n",
    "axes[0].set_xlabel('Time ($T$)')\n",
    "axes[0].set_ylabel('Excitation probabilities')\n",
    "axes[0].set_title('Floquet modes')\n",
    "axes[0].legend((\"Floquet mode 1\", \"Floquet mode 2\"))\n",
    "\n",
    "axes[1].plot(tlist, np.real(e_0), 'c', tlist, np.real(e_1), 'm')\n",
    "axes[1].plot(tlist, np.ones_like(tlist) * f_e[0], 'b', tlist, np.ones_like(tlist) * f_e[1], 'r')\n",
    "axes[1].set_xlabel('Time ($T$)')\n",
    "axes[1].set_ylabel('Energy [GHz]')\n",
    "axes[1].set_title('Eigen- and quasi-energies')\n",
    "axes[1].legend((\"Ground state\", \"Excited state\", \"Quasienergy 1\", \"Quasienergy 2\"));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Steadystate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def J_cb(omega):\n",
    "    \"\"\" Noise spectral density \"\"\"\n",
    "    return omega\n",
    "    \n",
    "def hamiltonian_t(t, args):\n",
    "    \"\"\" evaluate the hamiltonian at time t. \"\"\"\n",
    "    H0 = args['H0']\n",
    "    H1 = args['H1']\n",
    "    w  = args['w']\n",
    "\n",
    "    return H0 + np.sin(w * t) * H1\n",
    "\n",
    "def H1_coeff_t(t, args):\n",
    "    return np.sin(args['w'] * t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def qubit_integrate(delta, eps0_vec, A, omega, gamma1, gamma2, psi0, T, option):\n",
    "\n",
    "    # Hamiltonian\n",
    "    sx = sigmax()\n",
    "    sz = sigmaz()\n",
    "    sm = destroy(2)\n",
    "\n",
    "    quasi_energies = np.zeros((len(eps0_vec), 2))\n",
    "    f_gnd_prob     = np.zeros((len(eps0_vec), 2))\n",
    "    wf_gnd_prob    = np.zeros((len(eps0_vec), 2))\n",
    "    ss_prob1       = np.zeros(len(eps0_vec))\n",
    "    ss_prob2       = np.zeros(len(eps0_vec))\n",
    "\n",
    "    Hargs = {'w': omega}\n",
    "\n",
    "    for idx, eps0 in enumerate(eps0_vec):\n",
    "\n",
    "        H0 = - delta/2.0 * sx - eps0/2.0 * sz\n",
    "        H1 = A/2.0 * sz\n",
    "        H = [H0, [H1, 'sin(w * t)']]\n",
    "            \n",
    "        f_modes,f_energies = floquet_modes(H, T, Hargs)\n",
    "\n",
    "        quasi_energies[idx,:] = f_energies\n",
    "\n",
    "        f_gnd_prob[idx, 0] = expect(sm.dag() * sm, f_modes[0])\n",
    "        f_gnd_prob[idx, 1] = expect(sm.dag() * sm, f_modes[1])\n",
    "\n",
    "        f_states = floquet_states_t(f_modes, f_energies, 0, H, T, Hargs)\n",
    "\n",
    "        wf_gnd_prob[idx, 0] = expect(sm.dag() * sm, f_states[0])\n",
    "        wf_gnd_prob[idx, 1] = expect(sm.dag() * sm, f_states[1])\n",
    "\n",
    "        c_op = sigmax()\n",
    "        kmax = 5\n",
    "        temp = 0e-3\n",
    "        w_th = temp * (1.38e-23 / 6.626e-34) * 2 * np.pi * 1e-9    \n",
    "        Delta, X, Gamma, Amat = floquet_master_equation_rates(f_modes, f_energies, c_op, H, T, Hargs, J_cb, w_th, kmax)\n",
    "\n",
    "        rho_ss_fb = floquet_master_equation_steadystate(H0, Amat) # floquet basis\n",
    "        rho_ss_cb = rho_ss_fb.transform(f_modes, True) #False           # computational basis\n",
    "        \n",
    "        ss_prob1[idx] = expect(sm.dag() * sm, rho_ss_fb)\n",
    "        ss_prob2[idx] = expect(sm.dag() * sm, rho_ss_cb)\n",
    "\n",
    "    return quasi_energies, f_gnd_prob, wf_gnd_prob, ss_prob1, ss_prob2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta = 0.1 * 2 * np.pi  # qubit sigma_x coefficient\n",
    "eps0  = 1.0 * 2 * np.pi  # qubit sigma_z coefficient\n",
    "gamma1 = 0.0             # relaxation rate\n",
    "gamma2 = 0.0             # dephasing  rate\n",
    "A      = 2.0 * 2 * np.pi \n",
    "psi0   = basis(2,0)    # initial state\n",
    "omega  = np.sqrt(delta**2 + eps0**2) # driving frequency\n",
    "T      = (2*np.pi)/omega  # driving period\n",
    "\n",
    "param  = np.linspace(-2.0, 2.0, 100) * 2 * np.pi \n",
    "\n",
    "eps0 = param"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = time.time()\n",
    "q_energies, f_gnd_prob, wf_gnd_prob, ss_prob1, ss_prob2 = qubit_integrate(delta, eps0, A, omega, \n",
    "                                                                          gamma1, gamma2, psi0, T, \"dynamics\")\n",
    "print('dynamics: time elapsed = ' + str(time.time() - start_time))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(param, np.real(q_energies[:,0]) / delta, 'b',\n",
    "        param, np.real(q_energies[:,1]) / delta, 'r')\n",
    "ax.set_xlabel('A or e')\n",
    "ax.set_ylabel('Quasienergy')\n",
    "ax.set_title('Floquet quasienergies')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(param, np.real(f_gnd_prob[:,0]), 'b',\n",
    "        param, np.real(f_gnd_prob[:,1]), 'r')\n",
    "ax.set_xlabel('A or e')\n",
    "ax.set_ylabel('Occ. prob.')\n",
    "ax.set_title('Floquet modes excitation probability')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(param, np.real(wf_gnd_prob[:,0]), 'b',\n",
    "        param, np.real(wf_gnd_prob[:,1]), 'r')\n",
    "ax.set_xlabel('A or e')\n",
    "ax.set_ylabel('Occ. prob.')\n",
    "ax.set_title('Floquet states excitation probability')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(param, np.real(ss_prob1), 'r')\n",
    "ax.plot(param, np.real(ss_prob2), 'b')\n",
    "ax.set_xlabel('A or e')\n",
    "ax.set_ylabel('Occ. prob. in steady state')\n",
    "ax.set_title('Steady state excitation probability')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Versions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qutip.ipynbtools import version_table\n",
    "\n",
    "version_table()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}