{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ad5f297b",
   "metadata": {},
   "source": [
    "# Bosonic Dimer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9dee761",
   "metadata": {},
   "outputs": [],
   "source": [
    "using Pkg; Pkg.activate()\n",
    "\n",
    "using KadanoffBaym\n",
    "\n",
    "using FFTW, Interpolations\n",
    "function wigner_transform_itp(x::AbstractMatrix, ts::Vector; fourier = true, ts_lin = range(first(ts), last(ts); length = length(ts)))\n",
    "    itp = interpolate((ts, ts), x, Gridded(Linear()))\n",
    "    return wigner_transform([itp(t1, t2) for t1 in ts_lin, t2 in ts_lin]; ts = ts_lin, fourier = fourier)\n",
    "end\n",
    "\n",
    "using LinearAlgebra\n",
    "\n",
    "using PyPlot\n",
    "# PyPlot.plt.style.use(\"./paper.mplstyle\")\n",
    "using LaTeXStrings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ee0896c",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac99780f",
   "metadata": {},
   "source": [
    "A nice example to illustrate how one can use `KadanoffBaym.jl` to study *open systems* is a Bose dimer. It consists of two bosonic modes $\\omega_{1,2}$ (you can imagine two single-mode cavities at different frequencies), which are coupled with strength $J$. Additionally, each mode is coupled to its own reservoir at inverse temperature $\\beta_{1,2}$, respectively. Such a system is described by the master equation\n",
    "\n",
    "\\begin{align*}\n",
    "\t\\partial_{t} \\hat{\\rho}=-i\\left[\\hat{H} \\hat{\\rho}-\\hat{\\rho} \\hat{H}^{\n",
    "\t\t\\dagger}\\right]\n",
    "\t+ \\lambda\\sum_{i=1}^L \\left[(N_i + 1)\\hat{a}^{\\phantom{\\dagger}}_i \\hat{\\rho}  \\hat{a}^{\\dagger}_i + N_i \\hat{a}^{\\dagger}_i\\hat{\\rho}   \\hat{a}^{\\phantom{\\dagger}}_i \\right]\n",
    "\\end{align*}\n",
    "\n",
    "for $[\\hat{a}^{\\phantom{\\dagger}}_i, \\hat{a}^{\\dagger}_i]=1$, $i=1, ..., L$ and $L=2$. The $N_i=1/(e^{\\beta_i \\omega_i}-1)$ denote the thermal occupations of the reservoirs. The in this case *non-Hermitian* Hamiltonian is given by\n",
    "\n",
    "\\begin{align*}\n",
    "    \\hat{H}=\\sum_{i=1}^L (\\omega_{i}-i \\lambda (N_i + 1/2)) \\hat{a}^{\\dagger}_i \\hat{a}^{\\phantom{\\dagger}}_i.\n",
    "\\end{align*}\n",
    "\n",
    "The bosonic *lesser* and *greater* Green functions are \n",
    "\n",
    "\\begin{align*}\n",
    "\t\\left[\\boldsymbol{G}^<(t, t')\\right]_{ij} &= G^<_{ij}(t, t') = -i\\left\\langle{\\hat{a}_j^{{\\dagger}}(t')\\hat{a}_i^{\\phantom{\\dagger}}(t)}\\right\\rangle, \\\\\n",
    "\t\\left[\\boldsymbol{G}^>(t, t')\\right]_{ij} &= G^>_{ij}(t, t') = -i\\left\\langle{\\hat{a}_i^{\\phantom{\\dagger}}(t)\\hat{a}_j^{{\\dagger}}(t')}\\right\\rangle.\n",
    "\\end{align*}\n",
    "\n",
    "For convenience, we also introduce the (anti-) time-ordered Green functions\n",
    "\n",
    "\\begin{align*}\n",
    "\tG^{T}_{ij}(t, t') &= \\Theta(t - t') G^>_{ij}(t, t') +  \\Theta(t' - t) G^<_{ij}(t, t'), \\\\\n",
    "\tG^{\\tilde{T}}_{ij}(t, t') &= \\Theta(t - t') G^<_{ij}(t, t') +  \\Theta(t' - t) G^>_{ij}(t, t').\n",
    "\\end{align*}\n",
    "\n",
    "With the help of these, we can express the \"vertical\" equations of motion compactly:\n",
    "\n",
    "\\begin{align*}\n",
    "\t\\partial_t \\boldsymbol{G}^<(t, t') &= -i \\boldsymbol{H} \\boldsymbol{G}^<(t, t') + \\lambda \\operatorname{diag} (N_1, ..., N_L) \\boldsymbol{G}^{\\tilde{T}}(t, t') , \\\\\n",
    "\t\\partial_t \\boldsymbol{G}^>(t, t') &= -i \\boldsymbol{H}^\\dagger \\boldsymbol{G}^>(t, t') - \\lambda \\operatorname{diag} (N_1 + 1, ..., N_L + 1) \\boldsymbol{G}^{{T}}(t, t') ,\n",
    "\\end{align*}\n",
    "\n",
    "where $\\boldsymbol{H} = \\operatorname{diag}(\\omega_{1}- i\\lambda (N_1 + 1/2), ..., \\omega_{L}- i\\lambda(N_L + 1/2))$. We also need the equations in the \"diaognal\" time direction, which in the present case become\n",
    "\n",
    "\\begin{align*}\n",
    "\t\\partial_T {G}_{{ij}}^<(T, 0)_W &= \n",
    "\t-i \\left[\\boldsymbol{H} \\boldsymbol{G}^<(T, 0)_W - \\boldsymbol{G}^<(T, 0)_W \\boldsymbol{H}^\\dagger\\right]_{ij} \\\\\n",
    "\t&+ \\frac{i\\lambda}{2} (N_i + N_j) ({G}_{{ij}}^<(T, 0)_W + {G}_{{ij}}^>(T, 0)_W), \\\\\n",
    "\t\\partial_T {G}_{{ij}}^>(T, 0)_W &= -i \\left[\\boldsymbol{H}^\\dagger \\boldsymbol{G}^>(T, 0)_W - \\boldsymbol{G}^>(T, 0)_W \\boldsymbol{H}\\right]_{ij} \n",
    "\t\\\\\n",
    "\t&- \\frac{i\\lambda}{2} (N_i + N_j + 2) ({G}_{{ij}}^<(T, 0)_W + {G}_{{ij}}^>(T, 0)_W),\n",
    "\\end{align*}\n",
    "\n",
    "where the subscript $W$ again indicates *Wigner coordinates*."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c74db597",
   "metadata": {},
   "source": [
    "## Defining the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b16fea3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# quantum numbers\n",
    "dim = 2\n",
    "\n",
    "# Allocate the initial Green functions (time arguments at the end)\n",
    "GL = GreenFunction(zeros(ComplexF64, dim, dim, 1, 1), SkewHermitian)\n",
    "GG = GreenFunction(zeros(ComplexF64, dim, dim, 1, 1), SkewHermitian)\n",
    "\n",
    "# initial condition\n",
    "GL[1, 1] = -im * diagm([0.0, 2])\n",
    "GG[1, 1] = -im * I(2) + GL[1,1];"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1ccd8bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Non-Hermitian Hamiltonian and jump operator\n",
    "ω₁ = 2.5\n",
    "ω₂ = 0.0\n",
    "J = pi / 4\n",
    "\n",
    "λ = 1\n",
    "\n",
    "N₁ = 1.\n",
    "N₂ = 0.1\n",
    "\n",
    "H = [ω₁ - 0.5im * λ * ((N₁ + 1) + N₁) J; J ω₂ - 0.5im * λ * ((N₂ + 1) + N₂)]\n",
    "\n",
    "# right-hand side for the \"vertical\" evolution\n",
    "function fv!(out, _, _, _, t, t′)\n",
    "    out[1] = -1.0im * (H * GL[t, t′] + λ * [[1.0im * N₁, 0] [0, 1.0im * N₂]] * GL[t, t′])\n",
    "    out[2] = -1.0im * (adjoint(H) * GG[t, t′] - 1.0im * λ * [[(N₁ + 1), 0] [0, (N₂ + 1)]] * GG[t, t′])\n",
    "end\n",
    "\n",
    "# right-hand side for the \"diagonal\" evolution\n",
    "function fd!(out, _, _, _, t, t′)\n",
    "    out[1] = (-1.0im * (H * GL[t, t] - GL[t, t] * adjoint(H)\n",
    "             + 1.0im * λ * [[N₁ * (GL[1, 1, t, t] + GG[1, 1, t, t]), (N₁ + N₂) * (GL[2, 1, t, t] + GG[2, 1, t, t]) / 2] [(N₁ + N₂) * (GL[1, 2, t, t] + GG[1, 2, t, t]) / 2, N₂ * (GL[2, 2, t, t] + GG[2, 2, t, t])]])\n",
    "             )\n",
    "    out[2] = (-1.0im * (adjoint(H) * GG[t, t] - GG[t, t] * H\n",
    "             - 1.0im * λ * [[(N₁ + 1) * (GL[1, 1, t, t] + GG[1, 1, t, t]), (N₁ + N₂ + 2) * (GG[2, 1, t, t] + GL[2, 1, t, t]) / 2] [(N₁ + N₂ + 2) * (GG[1, 2, t, t] + GL[1, 2, t, t]) / 2, (N₂ + 1) * (GL[2, 2, t, t] + GG[2, 2, t, t])]])\n",
    "             )\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc1e11fb",
   "metadata": {},
   "source": [
    "## Solving an example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32403551",
   "metadata": {},
   "outputs": [],
   "source": [
    "# call the solver\n",
    "sol = kbsolve!(fv!, fd!, [GL, GG], (0.0, 32.0); atol=1e-6, rtol=1e-4);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79c63f9f",
   "metadata": {},
   "source": [
    "## Wigner coordinates and Fourier transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf651c70",
   "metadata": {},
   "outputs": [],
   "source": [
    "ρ_11_wigner, (taus, ts) = wigner_transform_itp((GG.data - GL.data)[1, 1, :, :], sol.t; fourier=false);\n",
    "\n",
    "ρ_22_wigner, (taus, ts) = wigner_transform_itp((GG.data - GL.data)[2, 2, :, :], sol.t; fourier=false);\n",
    "\n",
    "ρ_11_FFT, (ωs, ts) = wigner_transform_itp((GG.data - GL.data)[1, 1, :, :], sol.t; fourier=true);\n",
    "\n",
    "ρ_22_FFT, (ωs, ts) = wigner_transform_itp((GG.data - GL.data)[2, 2, :, :], sol.t; fourier=true);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80efc51a",
   "metadata": {},
   "source": [
    "## Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "809f849c",
   "metadata": {},
   "outputs": [],
   "source": [
    "let\n",
    "    function meshgrid(xin, yin)\n",
    "      nx=length(xin)\n",
    "      ny=length(yin)\n",
    "      xout=zeros(ny, nx)\n",
    "      yout=zeros(ny, nx)\n",
    "      for jx=1:nx\n",
    "          for ix=1:ny\n",
    "              xout[ix, jx]=xin[jx]\n",
    "              yout[ix, jx]=yin[ix]\n",
    "          end\n",
    "      end\n",
    "      return (x=xout, y=yout)\n",
    "    end\n",
    "\n",
    "    steps = 1\n",
    "    cmap = \"gist_heat\";\n",
    "\n",
    "    Y, X = meshgrid(sol.t[1:steps:end], sol.t[1:steps:end]);\n",
    "\n",
    "    xpad = 8\n",
    "    ypad = 5\n",
    "\n",
    "    fig = figure(figsize=(7, 3))\n",
    "\n",
    "    ax = subplot(121)\n",
    "\n",
    "    plot(sol.t, [-imag(GL.data[1, 1, k, k]) for k in 1:length(sol.t)], ls=\"--\", c=\"C3\", label=L\"i=1\", lw=1.5)\n",
    "    plot(sol.t, [-imag(GL.data[2, 2, k, k]) for k in 1:length(sol.t)], ls=\"-\", c=\"C0\", label=L\"i=2\", lw=1.5)\n",
    "\n",
    "    T = sol.t[end]\n",
    "\n",
    "    ax.set_xlim(0, T/2)\n",
    "    ax.set_xticks([0, T/4, T/2])\n",
    "    ax.set_ylim(0, 2.0)\n",
    "    ax.set_xlabel(L\"\\lambda T\")\n",
    "    ax.set_ylabel(L\"-\\mathrm{Im}\\; G^<_{ii}(t, t)\")\n",
    "    ax.xaxis.set_tick_params(pad=xpad)\n",
    "    ax.yaxis.set_tick_params(pad=ypad)\n",
    "    ax.legend(loc=\"best\", handlelength=1.9, frameon=false, borderpad=0, labelspacing=0.25)\n",
    "\n",
    "    ax = subplot(122)\n",
    "    X, Y = meshgrid(ts[1:steps:end], taus[1:steps:end]);\n",
    "    vmin = 1.0\n",
    "    vmax = -0.5\n",
    "\n",
    "    heatmap = ax.pcolormesh(X, Y, -ρ_11_wigner[1:steps:end, 1:steps:end] |> imag, cmap=cmap, vmin=vmin, vmax=vmax, rasterized=true)\n",
    "    heatmap.set_edgecolor(\"face\")\n",
    "\n",
    "    ax.set_xlabel(L\"\\lambda T\")\n",
    "    ax.set_ylabel(L\"\\lambda \\tau\")\n",
    "    ax.set_xlim(0, T)\n",
    "    ax.set_xticks([0, T/2, T])\n",
    "    ax.set_ylim(-T/2, T/2)\n",
    "    ax.set_yticks([-T/2, 0, T/2])\n",
    "    colorbar(mappable=heatmap)\n",
    "    ax.set_aspect(\"equal\")\n",
    "\n",
    "    tight_layout()\n",
    "    # savefig(\"boson_example_1.pdf\")\n",
    "    fig\n",
    "end;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c876e1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "let\n",
    "    xpad = 8\n",
    "    ypad = 5\n",
    "\n",
    "    T = sol.t[end]\n",
    "\n",
    "    fig = figure(figsize=(7, 3))\n",
    "\n",
    "    ax = subplot(121)\n",
    "    plot(taus, -ρ_11_wigner[:, Int(floor(length(taus)/2))] |> imag, ls=\"--\", c=\"C3\", label=L\"i=1\", lw=1.5) # fixed T\n",
    "    plot(taus, -ρ_22_wigner[:, Int(floor(length(taus)/2))] |> imag, ls=\"-\", c=\"C0\", label=L\"i=2\", lw=1.5)\n",
    "    ax.set_xlabel(L\"\\lambda \\tau\")\n",
    "    ax.set_xlim(-T/2, T/2)\n",
    "    ax.set_ylim(-0.5, 1.0)\n",
    "    ax.set_xticks([-T/2, -T/4, 0, T/4, T/2])\n",
    "    ax.xaxis.set_tick_params(pad=xpad)\n",
    "    ax.yaxis.set_tick_params(pad=ypad)\n",
    "#     ax.set_ylabel(L\"-\\textrm{Im}\\, A_{ii}(T, \\tau)_W\")\n",
    "    ax.legend(loc=\"best\", handlelength=1.4, frameon=false, borderpad=0, labelspacing=0.25)\n",
    "\n",
    "    ax = subplot(122)\n",
    "    plot(ωs, -ρ_11_FFT[:, Int(floor(length(taus)/2))] |> imag, ls=\"--\", c=\"C3\", lw=1.5)\n",
    "    plot(ωs, -ρ_22_FFT[:, Int(floor(length(taus)/2))] |> imag, ls=\"-\", c=\"C0\", lw=1.5)\n",
    "    ax.set_xlabel(L\"\\omega/\\lambda\")\n",
    "    ax.set_xlim(10 .* (-1, 1))\n",
    "    ax.set_xticks([-10, -5, 0, 5, 10])\n",
    "    ax.xaxis.set_tick_params(pad=xpad)\n",
    "    ax.yaxis.set_tick_params(pad=ypad)\n",
    "#     ax.set_ylabel(L\"-\\textrm{Im}\\, A_{ii}(T, \\omega)_{\\tilde{W}}\", labelpad=16)\n",
    "    ax.yaxis.set_label_position(\"right\")\n",
    "    ax.yaxis.set_ticks_position(\"both\")\n",
    "\n",
    "    tight_layout(pad=0.1, w_pad=0.5, h_pad=0)\n",
    "    # savefig(\"boson_example_2.pdf\")\n",
    "    fig\n",
    "end;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfd6711a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.9.4",
   "language": "julia",
   "name": "julia-1.9"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.9.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}