{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fa5b6ab9",
   "metadata": {},
   "source": [
    "# Fermionic Dimer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d45bf03",
   "metadata": {},
   "outputs": [],
   "source": [
    "using Pkg; Pkg.activate()\n",
    "using KadanoffBaym\n",
    "using LinearAlgebra\n",
    "\n",
    "using PyPlot\n",
    "PyPlot.plt.style.use(\"./paper.mplstyle\")\n",
    "using LaTeXStrings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b485778d",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f066217b",
   "metadata": {},
   "source": [
    "### Tight-binding Hamiltonian:\n",
    "\n",
    "\\begin{align}\\begin{split}\n",
    "\t\\hat{H} &= \\sum_{i=1}^L \\varepsilon_i \\hat{c}_i^{\\dagger} \\hat{c}_i^{\\phantom{\\dagger}} + J \\sum_{\\langle i, j\\rangle}\\left(\\hat{c}_i^{\\dagger} \\hat{c}_j^{\\phantom{\\dagger}} + \\hat{c}_j^{\\dagger} \\hat{c}_i^{\\phantom{\\dagger}}\\right)\n",
    "\\end{split}\\end{align}\n",
    "\n",
    "### Green functions\n",
    "\n",
    "\\begin{align}\\begin{split}\n",
    "\t\\left[\\boldsymbol{G}^>(t, t')\\right]_{ij} &= G^>_{ij}(t, t') = -i\\left\\langle{\\hat{c}_i^{\\phantom{\\dagger}}(t)\\hat{c}_j^{{\\dagger}}(t')}\\right\\rangle \\\\\n",
    "\t\\left[\\boldsymbol{G}^<(t, t')\\right]_{ij} &= G^<_{ij}(t, t') = \\phantom{-} i\\left\\langle{\\hat{c}_j^{{\\dagger}}(t')\\hat{c}_i^{\\phantom{\\dagger}}(t)}\\right\\rangle\n",
    "\\end{split}\\end{align}\n",
    "\n",
    "### Equations of motion\n",
    "\n",
    "\\begin{align}\\begin{split}\n",
    "\ti\\partial_t \\boldsymbol{G}^{\\lessgtr}(t, t') &= \\boldsymbol{H} \\boldsymbol{G}^{\\lessgtr}(t, t') \\\\\n",
    "\ti\\partial_T \\boldsymbol{G}^{\\lessgtr}(T, 0)_W &= [\\boldsymbol{H},\\boldsymbol{G}^{\\lessgtr}(T, 0)_W]\n",
    "\\end{split}\\end{align}\n",
    "\n",
    "\\begin{align}\\begin{split}\n",
    "\t\\boldsymbol{H} &= \n",
    "\t\\begin{pmatrix}\n",
    "\t\t\\varepsilon_1 & J      &        &   \\\\\n",
    "\t\tJ             & \\ddots & \\ddots &   \\\\\n",
    "\t\t              & \\ddots & \\ddots & J \\\\\n",
    "\t\t              &        & J & \\varepsilon_L \n",
    "\t\\end{pmatrix}\n",
    "\\end{split}\\end{align}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63686633",
   "metadata": {},
   "source": [
    "## Defining the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a754f1d4",
   "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",
    "N_0 = 1.0\n",
    "GL[1, 1] = +im * diagm([N_0, 0.0])\n",
    "GG[1, 1] = -im * I(2) + GL[1, 1];"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84bd2462",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hamiltonian matrix\n",
    "ε₁ = 1.0/20\n",
    "ε₂ = -1.0/20\n",
    "H = ComplexF64[ε₁ 1.0; 1.0 ε₂];\n",
    "\n",
    "# right-hand side for the \"vertical\" evolution\n",
    "function fv!(out, _, _, _, t1, t2)\n",
    "    out[1] = -1.0im * H * GL[t1, t2]\n",
    "    out[2] = -1.0im * H * GG[t1, t2]\n",
    "end\n",
    "\n",
    "# right-hand side for the \"diagonal\" evolution\n",
    "function fd!(out, h1, h2, times, t1, t2)\n",
    "  fv!(out, times, h1, h2, t1, t2)\n",
    "  out[1] .-= adjoint(out[1])\n",
    "  out[2] .-= adjoint(out[2])\n",
    "end\n",
    "\n",
    "# Analytic result\n",
    "ana(i1, i2, t1, t2) = (exp(-1.0im * H * t1) * GL[:, :, 1, 1] * exp(1.0im * H * t2))[i1, i2];"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08b8abd5",
   "metadata": {},
   "source": [
    "## Solving an example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b51e38ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# call the solver\n",
    "sol = kbsolve!(fv!, fd!, [GL, GG], (0.0, 100.0); atol=1e-9, rtol=1e-7);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3169157c",
   "metadata": {},
   "outputs": [],
   "source": [
    "let\n",
    "    # two quantum numbers to plot\n",
    "    idx1 = 1\n",
    "    idx2 = 1;\n",
    "\n",
    "    xpad = 8\n",
    "    ypad = 5\n",
    "\n",
    "    fig = figure(figsize=(7, 3))\n",
    "\n",
    "    ax = subplot(121)\n",
    "    plot(sol.t, [imag(GL.data[idx1, idx2, k, k]) for k in 1:length(sol.t)], marker=\"\", ms=3.0, ls=\"-\", c=\"C0\")\n",
    "    ax.set_xlim(0, 20.)\n",
    "    ax.set_ylim(0, N_0)\n",
    "    # ax.set_xticks(J .* [0, 0.5, 1])\n",
    "    ax.set_yticks([0, 0.5, 1])\n",
    "    ax.set_xlabel(L\"J t\")\n",
    "    ax.set_ylabel(L\"\\mathrm{Im}G^<_{11}(t, t)\")\n",
    "    ax.xaxis.set_tick_params(pad=xpad)\n",
    "    ax.yaxis.set_tick_params(pad=ypad)\n",
    "    ax.set_axisbelow(false)\n",
    "    ticklabel_format(axis=\"y\", style=\"sci\", scilimits=(-0, 0))\n",
    "\n",
    "    ax = subplot(122)\n",
    "    plot(sol.t, [imag(GL.data[idx1, idx2, k, k] - ana(idx1, idx2, sol.t[k], sol.t[k])) for k in eachindex(sol.t)], marker=\"\", ms=3.0, ls=\"-\", c=\"r\")\n",
    "    # ax.set_xlim(0, J * 5)\n",
    "    ax.set_ylim((-5, 5) .* 1e-4)\n",
    "    # ax.set_xticks(J .* [0, 1, 2, 3, 4, 5])\n",
    "    ax.set_xlabel(L\"J t\")\n",
    "    ax.set_ylabel(L\"\\mathrm{Im}\\left[G^<_{11}(t, t) - \\mathcal{G}^<_{11}(t, t)\\right]\", labelpad=16)\n",
    "    ax.xaxis.set_tick_params(pad=xpad)\n",
    "    ax.yaxis.set_tick_params(pad=ypad)\n",
    "    ax.yaxis.set_label_position(\"right\")\n",
    "    ax.set_axisbelow(false)\n",
    "    ticklabel_format(axis=\"y\", style=\"sci\", scilimits=(-0, 0))\n",
    "\n",
    "    tight_layout(pad=0.1, w_pad=0.75, h_pad=0)\n",
    "    # savefig(\"fermion_example_1.pdf\")\n",
    "    fig\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b668451",
   "metadata": {},
   "source": [
    "## Error analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e67de28f",
   "metadata": {},
   "outputs": [],
   "source": [
    "using LsqFit\n",
    "\n",
    "idx1 = 1\n",
    "idx2 = 1\n",
    "\n",
    "epsilons = [10^(-k) for k in range(3, 10; length=30)]\n",
    "\n",
    "err_data = []\n",
    "p_norm = 1\n",
    "\n",
    "for (k, eps) in enumerate(epsilons)\n",
    "    print(\"$k, \")\n",
    "\n",
    "    sol = kbsolve!(fv!, fd!, [GL, GG], (0.0, 100.0); dtini=1e-10, atol=1e-2eps, rtol=eps, kmax=9);\n",
    "\n",
    "    s = [ana(i, j, t1, t2) for i in 1:2, j in 1:2, t1 in sol.t, t2 in sol.t]\n",
    "    push!(err_data, (length(sol.t), KadanoffBaym.norm(GL.data - s), eps))\n",
    "end\n",
    "\n",
    "xdata = log10.([x[1] for x in err_data])\n",
    "ydata = log10.([x[2] for x in err_data]);\n",
    "\n",
    "fit_func = (n, p) -> -p[1] .* n .+ p[2];\n",
    "fit_result = curve_fit(fit_func, xdata, ydata, [2.0, 1]);\n",
    "coef(fit_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "143562e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "let\n",
    "    fig = figure(figsize=(7, 3))\n",
    "\n",
    "    ax = subplot(121)\n",
    "    plot(xdata, map(x -> fit_func(x, coef(fit_result)), xdata), \"--k\", lw=2,\n",
    "         label=L\"\\mathcal{O}(h^{%$(floor(coef(fit_result)[1], sigdigits=4))})\")\n",
    "    plot(xdata, ydata, \"o\", lw=0, ms=6,\n",
    "         markerfacecolor=\"C0\", markeredgewidth=0.25, markeredgecolor=\"#2D5FAA\")\n",
    "    ax.set_xlabel(L\"\\log(n)\")\n",
    "    ax.set_ylabel(L\"\\log(\\epsilon)\")\n",
    "    legend(loc=\"best\", handlelength=1.8, frameon=false, borderpad=0, labelspacing=0)\n",
    "\n",
    "    ax = subplot(122)\n",
    "    plot(log10.([x[3] for x in err_data]), log10.([x[1] for x in err_data]), \"o\", ms=5,\n",
    "         markerfacecolor=\"C0\", markeredgewidth=0.25, markeredgecolor=\"#2D5FAA\")\n",
    "    ax.set_xlim(-2, -10.3)\n",
    "    ax.set_xticks([-2, -4, -6, -8, -10])\n",
    "    ax.yaxis.set_label_position(\"right\")\n",
    "    ax.set_xlabel(L\"\\log(\\texttt{rtol})\")\n",
    "    ax.set_ylabel(L\"\\log(n)\", labelpad=16)\n",
    "\n",
    "    tight_layout(pad=0.1, w_pad=0.75, h_pad=0)\n",
    "    # savefig(\"fermion_example_error_scaling.pdf\")\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07cdb63d",
   "metadata": {},
   "source": [
    "## Comparison with fixed-step scheme"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cc5485d",
   "metadata": {},
   "outputs": [],
   "source": [
    "begin\n",
    "    # short final time\n",
    "    T = 0.4;\n",
    "\n",
    "    sol_adaptive = kbsolve!(fv!, fd!, [GL, GG], (0.0, T); \n",
    "        dtini=1e-6, rtol=1e-5, atol=1e-12, γ=99/100);\n",
    "    dts_adaptive = diff(sol_adaptive.t);\n",
    "    @show size(dts_adaptive)\n",
    "    @show error = let\n",
    "      s = [ana(i, j, t1, t2) for i in 1:2, j in 1:2, t1 in sol_adaptive.t, t2 in sol_adaptive.t]\n",
    "      KadanoffBaym.norm(GL.data - s)\n",
    "    end\n",
    "\n",
    "\n",
    "    sol_fixed = kbsolve!(fv!, fd!, [GL, GG], (0.0, T); \n",
    "        dtini=1e-6, rtol=1e-5, dtmax=1e-2, atol=1e-12, γ=9999/10000)\n",
    "    dts_fixed = diff(sol_fixed.t);\n",
    "    @show size(dts_fixed)\n",
    "    @show error = let\n",
    "      s = [ana(i, j, t1, t2) for i in 1:2, j in 1:2, t1 in sol_fixed.t, t2 in sol_fixed.t]\n",
    "      KadanoffBaym.norm(GL.data - s)\n",
    "    end;\n",
    "\n",
    "    sol_fixed2 = kbsolve!(fv!, fd!, [GL, GG], (0.0, T); \n",
    "        dtini=1e-2, rtol=1e0, dtmax=1e-2, atol=1e-12, γ=1)\n",
    "    dts_fixed2 = diff(sol_fixed2.t);\n",
    "    @show size(dts_fixed2)\n",
    "    @show error = let\n",
    "      s = [ana(i, j, t1, t2) for i in 1:2, j in 1:2, t1 in sol_fixed2.t, t2 in sol_fixed2.t]\n",
    "      KadanoffBaym.norm(GL.data - s)\n",
    "    end;\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63be8fc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "let\n",
    "    fig = figure(figsize=(3.5, 3))\n",
    "    ax = fig.add_subplot(111)\n",
    "\n",
    "    ax.tick_params(axis=\"y\", which=\"minor\")\n",
    "    ax.set_xlim(0, T)\n",
    "    ax.set_ylim(8e-5, 2e-1)\n",
    "\n",
    "    n = 7\n",
    "\n",
    "    # we discard the points that are outside of ylim (since we are drawing on top of the canvas)\n",
    "    ax.semilogy(sol_adaptive.t[n+1:end], dts_adaptive[n:end], \"s-C0\", lw=1.5, ms=5, label=\"adaptive\", markeredgecolor=\"#22577c\", clip_on=false, zorder=100)\n",
    "    ax.semilogy(sol_fixed.t[n+1:end], dts_fixed[n:end], \"o-C3\", lw=1, ms=4, label=\"semi-fixed\", alpha=0.9, clip_on=false, zorder=100)\n",
    "    ax.semilogy(sol_fixed2.t[2:end], dts_fixed2, \"o-k\", lw=1, ms=2, label=\"fixed\", clip_on=false, zorder=100)\n",
    "\n",
    "    xlabel(L\"J t\")\n",
    "    ylabel(L\"J h\")\n",
    "    legend(loc=\"lower right\", frameon=false, labelspacing=0.2, borderpad=0.0, handlelength=1.5, fontsize=15)\n",
    "\n",
    "    locmin = matplotlib.ticker.LogLocator(base=10.0,subs=(0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9),numticks=10)\n",
    "    ax.yaxis.set_minor_locator(locmin)\n",
    "    ax.yaxis.set_minor_formatter(matplotlib.ticker.NullFormatter())\n",
    "\n",
    "    tight_layout(pad=0.3)\n",
    "    # savefig(\"fermion_example_adaptive_dt.pdf\")\n",
    "    fig\n",
    "end;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e98536b8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.7.2",
   "language": "julia",
   "name": "julia-1.7"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}