{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fermi-Hubbard Model II"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Pkg; Pkg.activate()\n",
    "\n",
    "using KadanoffBaym, FFTW, Interpolations\n",
    "using LinearAlgebra, BlockArrays\n",
    "\n",
    "using JLD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using PyPlot\n",
    "# PyPlot.plt.style.use(\"./paper.mplstyle\")\n",
    "using LaTeXStrings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hamiltonian\n",
    "\n",
    "$$\n",
    "\\begin{align}\\begin{split}\n",
    "    \\hat{H} &= - J \\sum_{\\langle{i,\\,j}\\rangle}\\sum_\\sigma \\hat{c}^{\\dagger}_{i,\\sigma} \\hat{c}^{\\phantom{\\dagger}}_{i+1,\\sigma} + U\\sum_{i=1}^L  \\hat{c}^{\\dagger}_{i,\\uparrow} \\hat{c}^{\\phantom{\\dagger}}_{i,\\uparrow}   \\hat{c}^{\\dagger}_{i,\\downarrow} \\hat{c}^{\\phantom{\\dagger}}_{i,\\downarrow}, \n",
    "\\end{split}\\end{align}\n",
    "$$\n",
    "\n",
    "### Green functions\n",
    "\n",
    "$$\n",
    "    G^>_{\\uparrow,ij}(t, t') = -i \\left\\langle \\hat{c}^{\\phantom{\\dagger}}_{i,\\uparrow}(t) \\hat{c}^{{\\dagger}}_{i,\\uparrow}(t') \\right\\rangle\\\\\n",
    "    G^>_{\\downarrow,ij}(t, t') = -i \\left\\langle \\hat{c}^{\\phantom{\\dagger}}_{i,\\downarrow}(t) \\hat{c}^{{\\dagger}}_{i,\\downarrow}(t') \\right\\rangle\\\\\n",
    "$$\n",
    "\n",
    "### Self-energies\n",
    "\n",
    "Hartree-Fock:\n",
    "$$\n",
    "    \\Sigma^{\\mathrm{HF}}_{\\uparrow,\\,ij}(t, t') = {\\mathrm{i}}\\delta_{ij}\\delta(t - t') G^<_{\\downarrow,ii}(t, t)\\\\\n",
    "    \\Sigma^{\\mathrm{HF}}_{\\downarrow,\\,ij}(t, t') = {\\mathrm{i}}\\delta_{ij}\\delta(t - t') G^<_{\\uparrow,ii}(t, t)\n",
    "$$\n",
    "\n",
    "\n",
    "Second-order Born approximation:\n",
    "$$\n",
    "    \\Sigma_{ij, \\uparrow}  (t, t') = U^2 G_{ij, \\uparrow}(t, t') G_{ij, \\downarrow}(t, t') G_{ji, \\downarrow}(t', t),\\\\\n",
    "    \\Sigma_{ij, \\downarrow}(t, t') = U^2 G_{ij, \\downarrow}(t, t') G_{ij, \\uparrow}(t, t') G_{ji, \\uparrow}(t', t)\n",
    "$$\n",
    "\n",
    "\n",
    "$T$-matrix approximation:\n",
    "$$\n",
    "    \\Sigma_{ij, \\uparrow}  (t, t') = i U^2 T_{ij}(t, t') G_{ji, \\downarrow}(t', t),\\\\\n",
    "    \\Sigma_{ij, \\downarrow}(t, t') = i U^2 T_{ij}(t, t') G_{ji, \\uparrow}(t', t)\n",
    "$$\n",
    "\n",
    "$$\n",
    "   \\boldsymbol{T}(t, t') =  \\boldsymbol{\\Phi}(t, t') - U \\int_{\\mathcal{C}}\\mathrm{d}s\\; \\boldsymbol{\\Phi}(t, s) \\boldsymbol{T}(s, t')\n",
    "$$\n",
    "\n",
    "$$\n",
    "    \\Phi_{ij}(t, t') = -i G_{ij, \\uparrow}(t, t') G_{ij, \\downarrow}(t, t')\n",
    "$$\n",
    "\n",
    "$$\n",
    "   \\boldsymbol{T}^\\lessgtr(t, t') =  \\boldsymbol{\\Phi}^\\lessgtr(t, t') - U \\int_{0}^{t}\\mathrm{d}s\\; [\\boldsymbol{\\Phi}^>(t, s) - \\boldsymbol{\\Phi}^<(t, s)] \\boldsymbol{T}^\\lessgtr(s, t') - U \\int_{0}^{t'}\\mathrm{d}s\\; \\boldsymbol{\\Phi}^\\lessgtr(t, s)[\\boldsymbol{T}^<(s, t') - \\boldsymbol{T}^>(s, t')]\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solving"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function integrate1(hs::Vector, t1, t2, A::GreenFunction, B::GreenFunction, C::GreenFunction; tmax=t1)\n",
    "    retval = zero(A[t1,t1])\n",
    "\n",
    "    @inbounds for k in 1:tmax\n",
    "        @views LinearAlgebra.mul!(retval, A[t1, k] - B[t1, k], C[k, t2], hs[k], 1.0)\n",
    "    end\n",
    "    return retval\n",
    "end\n",
    "\n",
    "function integrate2(hs::Vector, t1, t2, A::GreenFunction, B::GreenFunction, C::GreenFunction; tmax=t2)\n",
    "    retval = zero(A[t1,t1])\n",
    "\n",
    "    @inbounds for k in 1:tmax\n",
    "        @views LinearAlgebra.mul!(retval, A[t1, k], B[k, t2] - C[k, t2], hs[k], 1.0)\n",
    "    end\n",
    "    return retval\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function fixed_point(F::Function, x0::AbstractArray; \n",
    "        mixing::Float64=0.5, \n",
    "        abstol::Float64=1e-12, \n",
    "        maxiter::Int=1000, \n",
    "        verbose::Bool=true, \n",
    "        norm=x -> LinearAlgebra.norm(x, Inf)\n",
    "    )\n",
    "    \n",
    "    x_old = copy(x0)\n",
    "\n",
    "    step = 0\n",
    "    while step < maxiter\n",
    "        x = F(x_old)\n",
    "        res = norm(x - x_old)\n",
    "        if verbose\n",
    "            @info \"step: $step // res: $res\"\n",
    "        end\n",
    "        if res < abstol\n",
    "            break\n",
    "        end\n",
    "        @. x_old = mixing * x + (1.0 - mixing) * x_old\n",
    "        step += 1\n",
    "    end\n",
    "\n",
    "    if step == maxiter\n",
    "        @warn \"No convergence reached.\"\n",
    "    end\n",
    "\n",
    "    return x_old\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lattice size\n",
    "L = 8\n",
    "\n",
    "# Allocate the initial Green functions (time arguments at the end)\n",
    "GL_u = GreenFunction(zeros(ComplexF64, L, L, 1, 1), SkewHermitian)\n",
    "GG_u = GreenFunction(zeros(ComplexF64, L, L, 1, 1), SkewHermitian)\n",
    "GL_d = GreenFunction(zeros(ComplexF64, L, L, 1, 1), SkewHermitian)\n",
    "GG_d = GreenFunction(zeros(ComplexF64, L, L, 1, 1), SkewHermitian);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initial conditions\n",
    "N_u = zeros(L)\n",
    "N_d = zeros(L)\n",
    "\n",
    "N_u[1:4] = 0.1 .* [1, 1, 1, 1]\n",
    "N_d[1:4] = 0.1 .* [1, 1, 1, 1]\n",
    "\n",
    "N_u[5:8] = 0.0 .* [1, 1, 1, 1]\n",
    "N_d[5:8] = 0.0 .* [1, 1, 1, 1]\n",
    "\n",
    "GL_u[1, 1] = 1.0im * diagm(N_u)\n",
    "GG_u[1, 1] = -1.0im * (I - diagm(N_u))\n",
    "GL_d[1, 1] = 1.0im * diagm(N_d)\n",
    "GG_d[1, 1] = -1.0im * (I - diagm(N_d));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Base.@kwdef struct FermiHubbardDataTM{T}\n",
    "    GL_u::T\n",
    "    GG_u::T\n",
    "    GL_d::T\n",
    "    GG_d::T\n",
    "\n",
    "    ΣL_u::T = zero(GL_u)\n",
    "    ΣG_u::T = zero(GG_u)\n",
    "    ΣL_d::T = zero(GL_d)\n",
    "    ΣG_d::T = zero(GG_d)\n",
    "\n",
    "    TL::T = zero(GL_u)\n",
    "    TG::T = zero(GG_u)\n",
    "\n",
    "    ΦL::T = zero(GL_u)\n",
    "    ΦG::T = zero(GG_u)\n",
    "end\n",
    "\n",
    "data = FermiHubbardDataTM(GL_u=GL_u, GG_u=GG_u, GL_d=GL_d, GG_d=GG_d)\n",
    "\n",
    "# Initialize T-matrix\n",
    "data.TL[1, 1] = -1.0im .* GL_u[1, 1] .* GL_d[1, 1]\n",
    "data.TG[1, 1] = -1.0im .* GG_u[1, 1] .* GG_d[1, 1]\n",
    "data.ΦL[1, 1] = data.TL[1, 1]\n",
    "data.ΦG[1, 1] = data.TG[1, 1];"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Base.@kwdef struct FermiHubbardModel{T}\n",
    "    # interaction strength\n",
    "    U::T\n",
    "\n",
    "    # 8-site 3D cubic lattice\n",
    "    h = begin\n",
    "        h = BlockArray{ComplexF64}(undef_blocks, [4, 4], [4, 4])\n",
    "        diag_block = [0 -1 0 -1; -1 0 -1 0; 0 -1 0 -1; -1 0 -1 0]\n",
    "        setblock!(h, diag_block, 1, 1)\n",
    "        setblock!(h, diag_block, 2, 2)\n",
    "        setblock!(h, Diagonal(-1 .* ones(4)), 1, 2)\n",
    "        setblock!(h, Diagonal(-1 .* ones(4)), 2, 1)\n",
    "\n",
    "        h |> Array\n",
    "    end\n",
    "\n",
    "    H_u = h\n",
    "    H_d = h\n",
    "end\n",
    "\n",
    "# Interaction parameter\n",
    "const U₀ = 2.0\n",
    "model = FermiHubbardModel(U = t -> U₀);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Right-hand side for the \"vertical\" evolution\n",
    "function fv!(model, data, out, times, h1, h2, t, t′)\n",
    "    # Unpack data and model\n",
    "    (; GL_u, GG_u, GL_u, GG_d, ΣL_u, ΣG_u, ΣL_d, ΣG_d) = data\n",
    "    (; H_u, H_d, U) = model\n",
    "\n",
    "    # Real-time collision integrals\n",
    "    ∫dt1(A, B, C) = integrate1(h1, t, t′, A, B, C)\n",
    "    ∫dt2(A, B, C) = integrate2(h2, t, t′, A, B, C)\n",
    "    \n",
    "    # The interaction varies as a function of the forward time (t+t')/2\n",
    "    U_t = U((times[t] + times[t′])/2)\n",
    "    \n",
    "    # Hartree-Fock self-energies\n",
    "    ΣHF_u(t, t′) = im * U_t * Diagonal(GL_d[t, t])\n",
    "    ΣHF_d(t, t′) = im * U_t * Diagonal(GL_u[t, t])\n",
    "    \n",
    "    # Equations of motion\n",
    "    out[1] = -1.0im * ((H_u + ΣHF_u(t, t′)) * GL_u[t, t′] + \n",
    "            ∫dt1(ΣG_u, ΣL_u, GL_u) + ∫dt2(ΣL_u, GL_u, GG_u)\n",
    "        )\n",
    "\n",
    "    out[2] = -1.0im * ((H_u + ΣHF_u(t, t′)) * GG_u[t, t′] + \n",
    "            ∫dt1(ΣG_u, ΣL_u, GG_u) + ∫dt2(ΣG_u, GL_u, GG_u)\n",
    "        )\n",
    "\n",
    "    out[3] = -1.0im * ((H_d + ΣHF_d(t, t′)) * GL_d[t, t′] + \n",
    "            ∫dt1(ΣG_d, ΣL_d, GL_d) + ∫dt2(ΣL_d, GL_d, GG_d)\n",
    "        )\n",
    "\n",
    "    out[4] = -1.0im * ((H_d + ΣHF_d(t, t′)) * GG_d[t, t′] +\n",
    "            ∫dt1(ΣG_d, ΣL_d, GG_d) + ∫dt2(ΣG_d, GL_d, GG_d)\n",
    "        )  \n",
    "    \n",
    "    return out\n",
    "end\n",
    "\n",
    "function fd!(model, data, out, times, h1, h2, t, t′)\n",
    "    fv!(model, data, out, times, h1, h2, t, t)\n",
    "    out .-= adjoint.(out)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## T-matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Callback function for the self-energies\n",
    "function T_matrix!(model, data, times, h1, h2, t, t′)\n",
    "    # Unpack data and model\n",
    "    (; GL_u, GG_u, GL_d, GG_d, TL, TG, ΣL_u, ΣG_u, ΣL_d, ΣG_d, ΦL, ΦG) = data\n",
    "    (; U) = model\n",
    "        \n",
    "    # Real-time collision integral\n",
    "    ∫dt1(A, B, C) = integrate1(h1, t, t′, A, B, C)\n",
    "    ∫dt2(A, B, C) = integrate2(h2, t, t′, A, B, C)\n",
    "    \n",
    "    # Resize self-energies etc. when Green functions are resized\n",
    "    if (n = size(GL_u, 3)) > size(ΣL_u, 3)\n",
    "        resize!(ΣL_u, n)\n",
    "        resize!(ΣG_u, n)\n",
    "        resize!(ΣL_d, n)\n",
    "        resize!(ΣG_d, n)\n",
    "        \n",
    "        resize!(TL, n)\n",
    "        resize!(TG, n)\n",
    "        resize!(ΦL, n)\n",
    "        resize!(ΦG, n)\n",
    "    end\n",
    "    \n",
    "    # The interaction varies as a function of the forward time (t+t')/2\n",
    "    U_t = U((times[t] + times[t′])/2)\n",
    "    \n",
    "    # Set all Φs at the very first t′ since they are all known by then\n",
    "    if t′ == 1\n",
    "        for t′ in 1:t\n",
    "            ΦL[t, t′] = -1.0im .* GL_u[t, t′] .* GL_d[t, t′]\n",
    "            ΦG[t, t′] = -1.0im .* GG_u[t, t′] .* GG_d[t, t′]\n",
    "        end\n",
    "    end   \n",
    "    \n",
    "    # Solve VIEs implicitly (uncached version)\n",
    "    TL[t, t′], TG[t, t′] = fixed_point([ΦL[t, t′], ΦG[t, t′]]; mixing=0.5, verbose=false) do x\n",
    "        TL[t, t′], TG[t, t′] = x[1], x[2]\n",
    "        \n",
    "        [\n",
    "            ΦL[t, t′] - U_t * (∫dt1(ΦG, ΦL, TL) + ∫dt2(ΦL, TL, TG)),\n",
    "            ΦG[t, t′] - U_t * (∫dt1(ΦG, ΦL, TG) + ∫dt2(ΦG, TL, TG))\n",
    "        ]\n",
    "    end\n",
    "    \n",
    "#     # Solve VIEs implicitly (cached version)\n",
    "#     TL[t, t′], TG[t, t′] = let\n",
    "#         ∫dt_(x...) = (x[1] < 1 || x[2] < 1) ? zero(x[3][1,1]) : ∫dt(x...)\n",
    "        \n",
    "#         ΦL_ = ΦL[t, t′]\n",
    "#         ΦG_ = ΦG[t, t′]\n",
    "        \n",
    "#         I1 = integrate1(h1, t, t′, ΦG, ΦL, TL; tmax=t-1)\n",
    "#         I2 = integrate2(h2, t, t′, ΦL, TL, TG; tmax=t′-1)\n",
    "#         I3 = integrate1(h1, t, t′, ΦG, ΦL, TG; tmax=t-1)\n",
    "#         I4 = integrate2(h2, t, t′, ΦG, TL, TG; tmax=t′-1)\n",
    "        \n",
    "#         L_ = ΦL[t, t′] - U_t * (I1 + I2)\n",
    "#         G_ = ΦG[t, t′] - U_t * (I3 + I4)\n",
    "        \n",
    "#         fixed_point([L_, G_]; mixing=0.5, verbose=false) do x\n",
    "#             TL[t, t′], TG[t, t′] = x[1], x[2]\n",
    "\n",
    "#             [\n",
    "#                 L_ - U_t * (h1[t] * (ΦG[t,t] - ΦL[t,t]) * TL[t,t′] + h2[t′] * ΦL[t,t′] * (TL[t′,t′] - TG[t′,t′])),\n",
    "#                 G_ - U_t * (h1[t] * (ΦG[t,t] - ΦL[t,t]) * TG[t,t′] + h2[t′] * ΦG[t,t′] * (TL[t′,t′] - TG[t′,t′]))\n",
    "#             ]\n",
    "#         end\n",
    "#     end           \n",
    "    \n",
    "    # Define the self-energies\n",
    "    ΣL_u[t, t′] = 1.0im .* U_t^2 .* TL[t, t′] .* transpose(GG_d[t′, t])\n",
    "    ΣL_d[t, t′] = 1.0im .* U_t^2 .* TL[t, t′] .* transpose(GG_u[t′, t])\n",
    "    \n",
    "    ΣG_u[t, t′] = 1.0im .* U_t^2 .* TG[t, t′] .* transpose(GL_d[t′, t])\n",
    "    ΣG_d[t, t′] = 1.0im .* U_t^2 .* TG[t, t′] .* transpose(GL_u[t′, t])\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmax = 8 # 32;\n",
    "tols = [(1e-8, 1e-10), (1e-6, 1e-8), (1e-4, 1e-6)];"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for (rtol, atol) in tols\n",
    "    @time sol = kbsolve!(\n",
    "        (x...) -> fv!(model, data, x...),\n",
    "        (x...) -> fd!(model, data, x...),\n",
    "        [data.GL_u, data.GG_u, data.GL_d, data.GG_d],\n",
    "        (0.0, tmax);\n",
    "        callback = (x...) -> T_matrix!(model, data, x...),\n",
    "        atol = atol,\n",
    "        rtol = rtol,\n",
    "        dtini=1e-10,\n",
    "        stop = x -> (println(\"t: $(x[end])\"); flush(stdout); false)\n",
    "    );\n",
    "\n",
    "    save(\"FH_3D_T_matrix_sol_U_$(U₀)_tmax_$(tmax)_atol_$(atol)_rtol_$(rtol).jld\",\n",
    "         \"solution\", sol,\n",
    "         \"GFs\", data.GL_u.data .+ data.GL_d.data)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Second Born approx."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Callback function for the self-energies\n",
    "function second_Born!(model, data, times, _, _, t, t′)\n",
    "    # Unpack data and model\n",
    "    (; GL_u, GG_u, GL_d, GG_d, ΣL_u, ΣG_u, ΣL_d, ΣG_d) = data\n",
    "    (; U) = model\n",
    "        \n",
    "    # Resize self-energies when Green functions are resized    \n",
    "    if (n = size(GL_u, 3)) > size(ΣL_u, 3)\n",
    "        resize!(ΣL_u, n)\n",
    "        resize!(ΣG_u, n)\n",
    "        resize!(ΣL_d, n)\n",
    "        resize!(ΣG_d, n)        \n",
    "    end\n",
    "    \n",
    "    # The interaction varies as a function of the forward time (t+t')/2\n",
    "    U_t = U((times[t] + times[t′])/2)\n",
    "    \n",
    "    # Define the self-energies\n",
    "    ΣL_u[t, t′] = U_t^2 .* GL_u[t, t′] .* GL_d[t, t′] .* transpose(GG_d[t′, t])\n",
    "    ΣL_d[t, t′] = U_t^2 .* GL_u[t, t′] .* GL_d[t, t′] .* transpose(GG_u[t′, t])\n",
    "    \n",
    "    ΣG_u[t, t′] = U_t^2 .* GG_u[t, t′] .* GG_d[t, t′] .* transpose(GL_d[t′, t])\n",
    "    ΣG_d[t, t′] = U_t^2 .* GG_u[t, t′] .* GG_d[t, t′] .* transpose(GL_u[t′, t])\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Callback function for the self-energies\n",
    "function second_Born!(model, data, times, _, _, t, t′)\n",
    "    # Unpack data and model\n",
    "    (; GL_u, GG_u, GL_d, GG_d, ΣL_u, ΣG_u, ΣL_d, ΣG_d) = data\n",
    "    (; U) = model\n",
    "        \n",
    "    # Resize self-energies when Green functions are resized    \n",
    "    if (n = size(GL_u, 3)) > size(ΣL_u, 3)\n",
    "        resize!(ΣL_u, n)\n",
    "        resize!(ΣG_u, n)\n",
    "        resize!(ΣL_d, n)\n",
    "        resize!(ΣG_d, n)        \n",
    "    end\n",
    "    \n",
    "    # The interaction varies as a function of the forward time (t+t')/2\n",
    "    U_t = U((times[t] + times[t′])/2)\n",
    "    \n",
    "    # Define the self-energies\n",
    "    ΣL_u[t, t′] = U_t^2 .* GL_u[t, t′] .* GL_d[t, t′] .* transpose(GG_d[t′, t])\n",
    "    ΣL_d[t, t′] = U_t^2 .* GL_u[t, t′] .* GL_d[t, t′] .* transpose(GG_u[t′, t])\n",
    "    \n",
    "    ΣG_u[t, t′] = U_t^2 .* GG_u[t, t′] .* GG_d[t, t′] .* transpose(GL_d[t′, t])\n",
    "    ΣG_d[t, t′] = U_t^2 .* GG_u[t, t′] .* GG_d[t, t′] .* transpose(GL_u[t′, t])\n",
    "end\n",
    "\n",
    "# Solve second Born for comparison\n",
    "@time sol_2B = kbsolve!(\n",
    "    (x...) -> fv!(model, data, x...),\n",
    "    (x...) -> fd!(model, data, x...),\n",
    "    [data.GL_u, data.GG_u, data.GL_d, data.GG_d],\n",
    "    (0.0, tmax);\n",
    "    callback = (x...) -> second_Born!(model, data, x...),\n",
    "    atol = 1e-8,\n",
    "    rtol = 1e-6,\n",
    "    dtini=1e-10,\n",
    "    stop = x -> (println(\"t: $(x[end])\"); flush(stdout); false)\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "GFs = []\n",
    "ts = []\n",
    "\n",
    "for (i, (rtol, atol)) in enumerate(tols)\n",
    "    loaded_data = load(\"FH_3D_T_matrix_sol_U_$(U₀)_tmax_$(tmax)_atol_$(atol)_rtol_$(rtol).jld\")\n",
    "    push!(GFs, loaded_data[\"GFs\"])\n",
    "    push!(ts, loaded_data[\"solution\"].t)\n",
    "end;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "xpad = 8\n",
    "ypad = 5\n",
    "\n",
    "figure(figsize = (8, 3))\n",
    "\n",
    "ax = subplot(121)\n",
    "idx_1 = 1\n",
    "ms = 4\n",
    "\n",
    "colors = [\"C2\", \"C0\", \"C3\"]\n",
    "styles = [\"s\", \"D\", \"o\"]\n",
    "ls = [\"-\", \"-.\", \"--\"]\n",
    "\n",
    "plot([], [], label=\"\\$\\\\texttt{rtol}\\$\", c=\"w\")\n",
    "for (i, (rtol, atol)) in enumerate(tols)\n",
    "    plot(ts[i], [imag(GFs[i][idx_1, idx_1, k, k]) for k in eachindex(ts[i])], \n",
    "        label = L\"\\texttt{%$(rtol)}\", lw=1.5, styles[i], ms=ms, c = colors[i])\n",
    "end\n",
    "\n",
    "plot(sol_2B.t, [imag(data.GL_u.data[idx_1, idx_1, k, k] .+ data.GL_d.data[idx_1, idx_1, k, k]) for k in eachindex(sol_2B.t)], \n",
    "    lw=3, ls = \"--\", c = \"k\", alpha=0.75)\n",
    "\n",
    "xlim(0, tmax)\n",
    "ax.set_xticks(0:tmax / 4:tmax)\n",
    "ylim(0, 0.3)\n",
    "xlabel(\"\\$J t\\$\")\n",
    "ylabel(\"Charge on site 1\", labelpad = 8)\n",
    "ax.xaxis.set_tick_params(pad = xpad)\n",
    "ax.yaxis.set_tick_params(pad = ypad)\n",
    "ticklabel_format(axis = \"y\", style = \"sci\", scilimits = (-0, 0))\n",
    "ax.legend(loc = \"best\", handlelength = 1, frameon = false, borderpad = 0, labelspacing = 0.25, ncol=2)\n",
    "\n",
    "ax = subplot(122)\n",
    "for (i, (rtol, atol)) in enumerate(reverse(tols))\n",
    "    semilogy(ts[i], [tr(GFs[i][:, :, k, k]) |> imag for k in eachindex(ts[i])] .- sum(N_u .+ N_d) .|> abs, \n",
    "        ls = ls[i], c = colors[i])\n",
    "end\n",
    "\n",
    "xlim(0, tmax)\n",
    "ylim(1e-10, 2e-4)\n",
    "ax.set_xticks(0:tmax / 4:tmax)\n",
    "ylabel(\"\\$Q(t) - Q_0\\$\", labelpad = 16)\n",
    "xlabel(\"\\$J t\\$\")\n",
    "ax.xaxis.set_tick_params(pad = xpad)\n",
    "ax.yaxis.set_tick_params(pad = ypad)\n",
    "# ticklabel_format(axis = \"y\", style = \"sci\", scilimits = (-0, 0))\n",
    "ax.yaxis.set_label_position(\"right\")\n",
    "\n",
    "tight_layout(pad = 0.1, w_pad = 0.5, h_pad = 0)\n",
    "# savefig(\"fermi_hubbard_T_matrix_convergence.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": 2
}
