{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Brownian Motion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# wrapper function for the KB solver\n",
    "function solution(N₀::Float64, D::Float64; args...)\n",
    "    \n",
    "    # analytical solution\n",
    "    function F_ana(t1, t2)\n",
    "        return (F[1, 1] - D/(2 * theta)) * exp(-theta * (t1 + t2)) + D/(2 * theta) * exp(-theta * abs(t1 - t2))\n",
    "    end\n",
    "    \n",
    "    # right-hand side for the \"vertical\" evolution\n",
    "    function fv!(out, _, _, _, t1, t2)\n",
    "        out[1] = -theta * F[t1, t2]\n",
    "    end\n",
    "    \n",
    "    # right-hand side for the \"diagonal\" evolution\n",
    "    function fd!(out, _, _, _, t1, t2)\n",
    "        out[1] = -theta * 2F[t1, t2] + D\n",
    "    end\n",
    "    \n",
    "    F = GreenFunction(N₀ * ones(1,1), Symmetrical)\n",
    "    \n",
    "    sol = kbsolve!(fv!, fd!, [F], (0.0, T); args...)\n",
    "\n",
    "    return [sol.t, F.data, [F_ana(t1, t2) for t1 in sol.t, t2 in sol.t]]\n",
    "end;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# final time\n",
    "T = 4.0\n",
    "\n",
    "# drift\n",
    "theta = 1.\n",
    "\n",
    "# time-scale\n",
    "t_scale = (iszero(theta) ? one(theta) : abs(theta))\n",
    "\n",
    "# different initial values\n",
    "N₀ = [1., 3., 5.]\n",
    "\n",
    "# different diffusion strengths\n",
    "Ds = [8., 4., 1.]\n",
    "\n",
    "s = mapreduce((x...) -> solution(x...; atol=1e-9, rtol=1e-7, dtini=1e-10), hcat, N₀, Ds);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cmap = \"gist_heat\";\n",
    "colors = [\"C0\", \"C3\", \"C2\"];\n",
    "lss = [\"-\", \"--\", \"-.\"];"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = figure(figsize=(7, 3))\n",
    "\n",
    "ax = subplot(121)\n",
    "for k in eachindex(N₀)\n",
    "    plot(s[1,k], diag(s[2,k]), label=\"\\$xxx\\$\", lw=1.5, c=colors[k], ls=lss[k])\n",
    "end\n",
    "ax.set_xlim(0, theta * T)\n",
    "ax.set_ylim(0, 1.1 * Ds[1]/(2theta))\n",
    "ax.set_yticks([0, 1, 2, 3, 4, 5])\n",
    "ax.set_xlabel(L\"\\theta t\")\n",
    "ax.set_ylabel(L\"F(t, t)\")#, pad=15)\n",
    "\n",
    "ax = subplot(122)\n",
    "plot([], [], label=L\"D/\\theta\", c=\"w\")\n",
    "for k in eachindex(N₀)\n",
    "    semilogy(s[1,k], abs.((diag(s[2,k])  .- diag(s[3,k])) ./ diag(s[3,k])), lw=1.5, c=colors[k], label=L\"%$(string(Ds[k]/(theta)))\", ls=lss[k])\n",
    "end\n",
    "\n",
    "ax.set_xlim(0, theta * T)\n",
    "ax.set_ylim(1e-9, 1e-6)\n",
    "ax.set_xlabel(L\"\\theta t\")\n",
    "ax.set_ylabel(L\"\\left|\\left(F(t,t) - \\mathcal{F}(t,t)\\right) / \\mathcal{F}(t,t)\\right|\", labelpad=10)\n",
    "ax.yaxis.set_label_position(\"right\")\n",
    "\n",
    "ax.legend(loc=\"best\", frameon=false, labelspacing=0.0, borderpad=0, handlelength=1.5, fontsize=\"small\")\n",
    "tight_layout(pad=0.25, w_pad=1, h_pad=0)\n",
    "\n",
    "# savefig(\"brownian_motion_example_1.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2D plot"
   ]
  },
  {
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y, X = meshgrid(s[1,1], s[1,1]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "figure(figsize=(4, 3))\n",
    "\n",
    "vmin = 0\n",
    "vmax = Ds[1]/(2theta)\n",
    "ax = subplot(111)\n",
    "heatmap = ax.pcolormesh(t_scale * X, t_scale * Y, s[2,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(L\"\\theta t\")\n",
    "ax.set_ylabel(L\"\\theta t^\\prime\")\n",
    "ax.set_xlim(0, t_scale * T)\n",
    "ax.set_ylim(0, t_scale * T)\n",
    "ax.set_xticks(t_scale .* [0, T/2, T])\n",
    "ax.set_yticks(t_scale .* [0, T/2, T])\n",
    "\n",
    "\n",
    "tight_layout(pad=0.75, w_pad=0.25, h_pad=0)\n",
    "# savefig(\"brownian_motion_example_2.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.7.0",
   "language": "julia",
   "name": "julia-1.7"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}