{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Geometric 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(S₀::Float64, σ::Float64)\n",
    "    \n",
    "    # analytical solution\n",
    "    F_ana(t1, t2) = S₀^2 * exp(mu * (t1 + t2)) * (exp(σ^2 * min(t1, t2)) - 1);\n",
    "    \n",
    "    # right-hand side for the \"vertical\" evolution\n",
    "    function fv!(out, _, _, _, t1, t2)\n",
    "      out[1] = 0\n",
    "      out[2] = mu * F[t1, t2]\n",
    "    end\n",
    "    \n",
    "    # right-hand side for the \"diagonal\" evolution\n",
    "    function fd!(out, _, _, _, t1, t2)\n",
    "      out[1] = mu * S[t1, t2]\n",
    "      out[2] = 2mu * F[t1, t2] + σ^2 * (S[t1, t2]^2 + F[t1, t2])\n",
    "    end\n",
    "  \n",
    "    S = GreenFunction(S₀ * ones(1,1), Symmetrical)\n",
    "    F = GreenFunction(0.0 * ones(1,1), Symmetrical)\n",
    "    \n",
    "    sol = kbsolve!(fv!, fd!, [S, F], (0.0, T), atol=1e-9, rtol=1e-7, dtini=1e-10)\n",
    "         \n",
    "    return [sol.t, S.data, 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 = 1.0\n",
    "\n",
    "# drift\n",
    "mu = 1.\n",
    "\n",
    "# time-scale\n",
    "t_scale = (mu == 0. ? 1. : abs(mu))\n",
    "\n",
    "# different initial values\n",
    "S₀ = [1., 2., 5.]\n",
    "\n",
    "# different diffusion strengths\n",
    "sigma = [1., .5, .1]\n",
    "\n",
    "s = mapreduce(solution, hcat, S₀, sigma);"
   ]
  },
  {
   "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": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "figure(figsize=(7, 3))\n",
    "\n",
    "ax = subplot(121)\n",
    "for k in eachindex(sigma)\n",
    "    if k == 1\n",
    "        plot(s[1,k], diag(s[3,k]), label=L\"F(t, t)\", lw=1.5, c=colors[k], ls=lss[k])\n",
    "        plot(s[1,k], diag(s[2,k]), label=L\"\\langle X(t) \\rangle\", lw=3.5, c=colors[k], ls=lss[k])\n",
    "    else\n",
    "        plot(s[1,k], diag(s[3,k]), lw=1.5, c=colors[k], ls=lss[k])\n",
    "        plot(s[1,k], diag(s[2,k]), lw=3.5, c=colors[k], ls=lss[k])\n",
    "    end\n",
    "end\n",
    "ax.set_xlim(0, t_scale * T)\n",
    "ax.set_ylim(0, 15)\n",
    "ax.set_xlabel(L\"\\mu t\")\n",
    "ax.legend(loc=\"best\", handlelength=1.5, frameon=false, borderpad=0, labelspacing=0.25, fontsize=\"small\")\n",
    "ax.get_legend().legendHandles[1].set_color(\"k\")\n",
    "ax.get_legend().legendHandles[2].set_color(\"k\")\n",
    "\n",
    "ax = subplot(122)\n",
    "semilogy([], [], label=L\"\\sigma / \\mu\", c=\"w\")\n",
    "for k in eachindex(sigma)\n",
    "    semilogy(s[1,k], abs.((diag(s[3,k])  .- diag(s[4,k]))), lw=1.5, c=colors[k], label=L\"%$(string(sigma[k]/(mu)))\", ls=lss[k])\n",
    "end\n",
    "\n",
    "ax.set_xlim(0, t_scale * T)\n",
    "ax.set_ylim(1e-9, 1e-5)\n",
    "ax.set_xlabel(L\"\\mu t\")\n",
    "ax.set_ylabel(L\"\\left|F(t, t) - \\mathcal{F}(t, t)\\right|\", labelpad=10)\n",
    "\n",
    "# ax.yaxis.set_ticks_position(\"right\")\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",
    "\n",
    "tight_layout(pad=0.25, w_pad=1, h_pad=0)\n",
    "\n",
    "# savefig(\"geometric_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 = \n",
    "ax = subplot(111)\n",
    "heatmap = ax.pcolormesh(t_scale * X, t_scale * Y, s[3,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(\"\\$\\\\lambda t\\$\")\n",
    "ax.set_ylabel(\"\\$\\\\lambda t'\\$\")\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",
    "tight_layout(pad=0.75, w_pad=0.25, h_pad=0)\n",
    "# savefig(\"geometric_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
}