{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fermi-Hubbard Model I"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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, BlockArrays\n",
    "\n",
    "using JLD\n",
    "\n",
    "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",
    "$$"
   ]
  },
  {
   "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": [
    "# 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)\n",
    "\n",
    "# Initial conditions\n",
    "N_u = zeros(L)\n",
    "N_d = zeros(L)\n",
    "\n",
    "# From the paper\n",
    "N_u[1:4] = [0.7, 0.0, 0.7, 0.0]\n",
    "N_d[1:4] = [0.0, 0.25, 0.0, 0.25]\n",
    "\n",
    "N_u[5:8] = [0.0, 0.4, 0.0, 0.4]\n",
    "N_d[5:8] = [0.65, 0.0, 0.65, 0.0]\n",
    "\n",
    "# From the docs\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 FermiHubbardData2B{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",
    "end\n",
    "\n",
    "data = FermiHubbardData2B(GL_u=GL_u, GG_u=GG_u, GL_d=GL_d, GG_d=GG_d);"
   ]
  },
  {
   "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",
    "# Relatively small interaction parameter\n",
    "const U₀ = 0.25\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",
    "# Right-hand side for the \"diagonal\" evolution\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": "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": {},
   "outputs": [],
   "source": [
    "tmax = 32;\n",
    "atol = 1e-8\n",
    "rtol = 1e-6;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "@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...) -> second_Born!(model, data, x...),\n",
    "    atol = atol,\n",
    "    rtol = rtol,\n",
    "    stop = x -> (println(\"t: $(x[end])\"); flush(stdout); false)\n",
    ");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save(\"FH_3D_sol_U_\"*string(U₀)*\"_tmax_\"*string(tmax)*\"_atol_\"*string(atol)*\"_rtol_\"*string(rtol)*\".jld\", \"solution\", sol)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example plots"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loaded_sol = load(\"FH_3D_sol_U_\"*string(U₀)*\"_tmax_\"*string(tmax)*\"_atol_\"*string(atol)*\"_rtol_\"*string(rtol)*\".jld\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import PyCall\n",
    "const is = PyCall.pyimport(\"mpl_toolkits.axes_grid1.inset_locator\")\n",
    "const inset_axes = is.inset_axes;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "let\n",
    "    num_points = loaded_sol[\"solution\"].t |> length\n",
    "    \n",
    "    xpad = 8\n",
    "    ypad = 5\n",
    "\n",
    "    fig, (ax1, ax2) = subplots(1, 2, figsize = (8, 3))\n",
    "    idx_1 = 1\n",
    "    idx_2 = 8\n",
    "\n",
    "    ax1.plot(loaded_sol[\"solution\"].t, [imag(data.GL_u.data[idx_1, idx_1, k, k] .+ data.GL_d.data[idx_1, idx_1, k, k]) for k = 1:num_points], \n",
    "        label = \"\\$ i=1\\$\", lw=1.5, ls = \"--\", c = \"#438E6A\")\n",
    "\n",
    "    ax1.plot(loaded_sol[\"solution\"].t, [imag(data.GL_u.data[idx_2, idx_2, k, k] .+ data.GL_d.data[idx_2, idx_2, k, k]) for k = 1:num_points], \n",
    "        label = \"\\$ i=8\\$\", lw=1.5, ls = \"-\", c = \"#2D5FAA\")\n",
    "\n",
    "    ax1.set_xlim(0, tmax)\n",
    "    ax1.set_xticks([0, 8, 16, 24, 32])\n",
    "    ax1.set_ylim(-0, 1)\n",
    "    ax1.set_xlabel(\"\\$J t\\$\")\n",
    "    ax1.set_ylabel(\"Charge\", labelpad = 8)\n",
    "    ax1.xaxis.set_tick_params(pad = xpad)\n",
    "    ax1.yaxis.set_tick_params(pad = ypad)\n",
    "    ax1.legend(loc = \"best\", handlelength = 1.9, frameon = false, borderpad = 0, labelspacing = 0.25)\n",
    "\n",
    "    # ax2 = subplot(122)\n",
    "\n",
    "    ax2.plot(loaded_sol[\"solution\"].t, [imag(data.GL_u.data[idx_1, idx_1, k, k] .- data.GL_d.data[idx_1, idx_1, k, k]) for k = 1:num_points], \n",
    "        label = \"\\$ i=2\\$\", lw=1.5, ls = \"--\", c = \"#438E6A\")\n",
    "\n",
    "    ax2.plot(loaded_sol[\"solution\"].t, [imag(data.GL_u.data[idx_2, idx_2, k, k] .- data.GL_d.data[idx_2, idx_2, k, k]) for k = 1:num_points], \n",
    "        label = \"\\$ i=2\\$\",  lw=1.5, ls = \"-\", c = \"#2D5FAA\")\n",
    "\n",
    "    ax2.set_xlim(0, tmax)\n",
    "    ax2.set_ylim(-1, 1)\n",
    "    ax2.set_xticks([0, 8, 16, 24, 32])\n",
    "    ax2.set_xlabel(\"\\$J t\\$\")\n",
    "    ax2.set_ylabel(\"Spin\", labelpad = 16)\n",
    "    ax2.xaxis.set_tick_params(pad = xpad)\n",
    "    ax2.yaxis.set_tick_params(pad = ypad)\n",
    "    ax2.set_axisbelow(false)\n",
    "    ax2.yaxis.set_label_position(\"right\")\n",
    "\n",
    "    axins1 = inset_axes(ax1, width=1.1 * 1.0, height=0.8*3/4, loc=4, \n",
    "        bbox_to_anchor=(0.96, .0, .0, .0),\n",
    "        bbox_transform=ax1.transAxes)\n",
    "\n",
    "    axins1.plot(loaded_sol[\"solution\"].t, [(tr(data.GL_u.data[:, :, k, k]) |> imag) .+ (tr(data.GL_d.data[:, :, k, k]) |> imag) for k = 1:num_points] \n",
    "        .- sum(N_u .+ N_d)  .|> abs, \n",
    "        label = \"\\$ c \\$\", ls = \"-\", c = \"k\")\n",
    "    axins1.set_xlim(0, tmax)\n",
    "    axins1.set_xticks([0, 8, 16, 24, 32])\n",
    "    axins1.set_xticklabels([])\n",
    "    axins1.set_yticks([k for k in 0:0.5:1] .* 5e-15)\n",
    "    axins1.set_ylim([0.0, 1] .* 5e-15)\n",
    "    axins1.set_ylabel(L\"Q(t) - Q_0\", fontdict = Dict(:fontsize=>10))\n",
    "    axins1.tick_params(axis=\"y\", labelsize=10)\n",
    "    axins1.yaxis.get_offset_text().set_fontsize(10)\n",
    "    axins1.yaxis.set_label_position(\"right\")\n",
    "\n",
    "    axins2 = inset_axes(ax2, width=1.1 * 1.0, height=0.8*3/4, loc=4, \n",
    "        bbox_to_anchor=(0.96, .0, .0, .0),\n",
    "        bbox_transform=ax2.transAxes)\n",
    "\n",
    "    axins2.plot(loaded_sol[\"solution\"].t, [(tr(data.GL_u.data[:, :, k, k]) |> imag) .- (tr(data.GL_d.data[:, :, k, k]) |> imag) for k = 1:num_points] \n",
    "        .- sum(N_u .- N_d) .|> abs, \n",
    "        label = \"\\$ c \\$\", ls = \"-\", c = \"k\")\n",
    "    axins2.set_xlim(0, tmax)\n",
    "    axins2.set_xticks([0, 8, 16, 24, 32])\n",
    "#     axins2.set_yticks([k for k in 0:0.5:1] .* 3e-15)\n",
    "    axins2.set_xticklabels([])\n",
    "    axins2.set_ylim([0.0, 1] .* 12e-15)\n",
    "    axins2.set_ylabel(L\"S(t) - S_0\", fontdict = Dict(:fontsize=>10))\n",
    "    axins2.tick_params(axis=\"y\", labelsize=10)\n",
    "    axins2.yaxis.get_offset_text().set_fontsize(10)\n",
    "    axins2.yaxis.set_label_position(\"right\")\n",
    "\n",
    "    tight_layout(pad = 0.1, w_pad = 0.5, h_pad = 0)\n",
    "#     savefig(\"fermi_hubbard_T.pdf\")\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "Y, X = meshgrid(loaded_sol[\"solution\"].t, loaded_sol[\"solution\"].t);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "let\n",
    "    # quantum number to look at\n",
    "    idx = 1\n",
    "    shift = 0\n",
    "\n",
    "    ρτ, (τs, ts) = wigner_transform_itp((data.GG_u.data[idx, idx, :, :] - data.GL_u.data[idx, idx, :, :]), \n",
    "        loaded_sol[\"solution\"].t[1+shift:end-shift], fourier=false);\n",
    "    ρω, (ωs, ts) = wigner_transform_itp((data.GG_u.data[idx, idx, :, :] - data.GL_u.data[idx, idx, :, :]), \n",
    "        loaded_sol[\"solution\"].t[1+shift:end-shift], fourier=true);\n",
    "    cmap = \"gist_heat\";\n",
    "    \n",
    "    xpad = 8\n",
    "    ypad = 5\n",
    "    \n",
    "    figure(figsize=(7, 3))\n",
    "    \n",
    "    t_scale = 1\n",
    "    vmin = -1.0\n",
    "    vmax = 1.0\n",
    "\n",
    "    center = floor(length(ts) / 2) |> Int\n",
    "\n",
    "    ax = subplot(121)\n",
    "    plot(τs, -ρτ[:, center] |> imag, ls=\"-\", c=\"C0\", lw=1.5)\n",
    "    ax.set_xlabel(\"\\$J \\\\tau\\$\")\n",
    "    ax.set_xlim(-tmax / 2, t_scale * tmax / 2)\n",
    "    ax.set_ylim(-1, 1)\n",
    "    ax.set_xticks([-tmax/2, -tmax/4, 0, tmax/4, tmax/2])\n",
    "    ax.xaxis.set_tick_params(pad=xpad)\n",
    "    ax.yaxis.set_tick_params(pad=ypad)\n",
    "    ax.set_ylabel(\"\\$  A_{11, \\\\uparrow}(T_{\\\\mathrm{max}}/2, \\\\tau)_W \\$\")\n",
    "\n",
    "    ax = subplot(122)\n",
    "    heatmap = ax.pcolormesh(X, Y, imag(data.GL_u.data[1, 1, :, :]) .- imag(data.GG_u.data[1, 1, :, :]), cmap=cmap, rasterized=true, vmin=vmin, vmax=vmax)\n",
    "    heatmap.set_edgecolor(\"face\")\n",
    "    ax.set_aspect(\"equal\")\n",
    "    cbar = colorbar(mappable=heatmap)\n",
    "    cbar.formatter.set_powerlimits((0, 0))\n",
    "    ax.set_xlabel(\"\\$J t\\$\")\n",
    "    ax.set_ylabel(\"\\$J t'\\$\")\n",
    "    ax.set_xlim(0, t_scale * tmax)\n",
    "    ax.set_ylim(0, t_scale * tmax)\n",
    "    ax.set_xticks(t_scale .* [0, tmax/2, tmax])\n",
    "    ax.set_yticks(t_scale .* [0, tmax/2, tmax])\n",
    "\n",
    "    tight_layout(pad=0.75, w_pad=0.5, h_pad=0)\n",
    "\n",
    "    # savefig(\"fermi_hubbard_example_two_times.pdf\")\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quench"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "const U₀ = 10.\n",
    "model = FermiHubbardModel(U = t -> -U₀ * [(-1)^k * (1 + exp(-20(t - 2k)))^(-1) for k in 1:tmax-1] |> sum);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tmax = 16;\n",
    "atol = 1e-5\n",
    "rtol = 1e-3\n",
    "\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...) -> second_Born!(model, data, x...),\n",
    "    atol = atol,\n",
    "    rtol = rtol,\n",
    "    stop = x -> (println(\"t: $(x[end])\"); flush(stdout); false)\n",
    ");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save(\"quenched_FH_3D_sol_U_\"*string(U₀)*\"_tmax_\"*string(tmax)*\"_atol_\"*string(atol)*\"_rtol_\"*string(rtol)*\".jld\", \"solution\", sol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loaded_sol = load(\"quenched_FH_3D_sol_U_\"*string(U₀)*\"_tmax_\"*string(tmax)*\"_atol_\"*string(atol)*\"_rtol_\"*string(rtol)*\".jld\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "let\n",
    "    tt = loaded_sol[\"solution\"].t\n",
    "    fig = figure(figsize=(8, 4))\n",
    "    ax = subplot(221)\n",
    "    ax.plot(tt[2:end], map(t -> t[2] - t[1], zip(tt[1:end-1], tt[2:end])), \"-s\", ms=3, markeredgecolor=\"#22577c\")\n",
    "    ax.set_xticks([0, 4, 8, 12, 16])\n",
    "    xlim(0, tmax)\n",
    "    ylim(0, 0.2)\n",
    "    ax.set_xticklabels([])\n",
    "    ylabel(\"\\$J h\\$\")\n",
    "\n",
    "    ax = subplot(223)\n",
    "    ax.plot(tt, model.U.(tt), \"-k\")\n",
    "    ax.set_xticks([0, 4, 8, 12, 16])\n",
    "    xlim(0, tmax)\n",
    "    xlabel(\"\\$Jt\\$\")\n",
    "    ylabel(\"\\$U(t)\\$\")\n",
    "\n",
    "    # quantum number to look at\n",
    "    idx = 1\n",
    "\n",
    "    ρτ, (τs, ts) = wigner_transform_itp((data.GG_u.data[idx, idx, :, :] - data.GL_u.data[idx, idx, :, :]), \n",
    "        loaded_sol[\"solution\"].t[1:end], fourier=false);\n",
    "    ρω, (ωs, ts) = wigner_transform_itp((data.GG_u.data[idx, idx, :, :] - data.GL_u.data[idx, idx, :, :]), \n",
    "        loaded_sol[\"solution\"].t[1:end], fourier=true);\n",
    "\n",
    "    t_scale = 1\n",
    "    ω_scale = 1;\n",
    "\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",
    "    Y, X = meshgrid(loaded_sol[\"solution\"].t, loaded_sol[\"solution\"].t);\n",
    "\n",
    "    cmap = \"gist_heat\";\n",
    "\n",
    "    # fig = figure(figsize=(7, 3))\n",
    "    t_scale = 1\n",
    "    vmin = -1.0\n",
    "    vmax = 1.0\n",
    "\n",
    "    ax = subplot(122)\n",
    "    heatmap = ax.pcolormesh(X, Y, imag(data.GL_u.data[1, 1, :, :]) .- imag(data.GG_u.data[1, 1, :, :]), cmap=cmap, rasterized=true, vmin=vmin, vmax=vmax)\n",
    "    heatmap.set_edgecolor(\"face\")\n",
    "    ax.set_aspect(\"equal\")\n",
    "    cbar = colorbar(mappable=heatmap)\n",
    "    cbar.formatter.set_powerlimits((0, 0))\n",
    "    ax.set_xlabel(\"\\$J t\\$\")\n",
    "    ax.set_ylabel(\"\\$J t'\\$\")\n",
    "    ax.set_xlim(0, t_scale * tmax)\n",
    "    ax.set_ylim(0, t_scale * tmax)\n",
    "    ax.set_xlim(0, t_scale * 8)\n",
    "    ax.set_ylim(0, t_scale * 8)\n",
    "\n",
    "    ax.set_xticks(t_scale .* [0, 2, 4, 6, 8])\n",
    "    ax.set_yticks(t_scale .* [0, 2, 4, 6, 8])\n",
    "\n",
    "    tight_layout(pad=0.75, w_pad=0.5, h_pad=0)\n",
    "#     savefig(\"quenched_fermi_hubbard_example_two_times.pdf\")\n",
    "    fig\n",
    "end;"
   ]
  },
  {
   "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
}