{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Classroom problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#using Pkg\n",
    "#Pkg.add(\"StatsBase\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## General code for solving Itô integrals with the Euler-Maruyama method (no boundary conditions (BCs) included)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "de_EM_solver = function(f, g, x₀, p, T, Δt)\n",
    "\n",
    "    \"\"\"\n",
    "    f = f(x, p, t) and g = g(x, p, t), where x is the process variable, p stands for a vector of parameters and t is time.    \n",
    "    Outputs:\n",
    "        tt - a partitioned time interval [0, T]\n",
    "        xx - a solution to the stochastic integral dX_t = f(X_t, p, t)dt + g(X_t, p, t)dW\n",
    "    \"\"\"\n",
    "    \n",
    "    x = copy(x₀)\n",
    "    t = 0.0\n",
    "\n",
    "    Nx = length(x) \n",
    "    tt = 0.0:Δt:T\n",
    "    xx = zeros(length(tt), Nx)\n",
    "\n",
    "    xx[1, :] .= x\n",
    "    \n",
    "    Nw = length(p)-1\n",
    "    \n",
    "    for n in 1:length(tt)-1\n",
    "        t⁻ = tt[n]\n",
    "        x⁻ = xx[n, :]\n",
    "        xx[n+1, :] = x⁻ + f(x⁻, p, t⁻).*Δt + sqrt(Δt)*g(x⁻, p, t⁻)*randn(Nw, 1)\n",
    "    end\n",
    "    \n",
    "    return tt, xx\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. System with multiple Favourable States (Homework from week 3 revisited)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To simplify the problem, let us work with the interval $t\\in[0,5]$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\frac{dx}{dt} = -\\frac{κ₁}{6V^2}*x^3+\\frac{\\kappa_2}{2V}*x^2-κ₃*x+κ₄V,\\quad x_0 = 0$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SM = [-1 1 -1 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Δt = 0.01\n",
    "κ₁ = 0.0015\n",
    "κ₂ = 0.36\n",
    "κ₃ = 37.5\n",
    "κ₄ = 2200\n",
    "V = 0.5\n",
    "\n",
    "p = [κ₁,κ₂,κ₃,κ₄, V]\n",
    "x₀ = zeros(1)\n",
    "T = 5.0\n",
    "\n",
    "\n",
    "fode = function(x, p, t) \n",
    "    κ₁,κ₂,κ₃,κ₄,V = p\n",
    "    return ####                                         # drift term\n",
    "end\n",
    "gode = function(x, p, t)\n",
    "    Nx = length(x)\n",
    "    Nw = length(p)-1\n",
    "    return (zeros(Nx, Nw))\n",
    "end\n",
    "tt, xx_ode = de_EM_solver(fode, gode, x₀, p, T, Δt);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Plots\n",
    "plot(tt, xx_ode, colour=:black, ylabel=\"X(t)\", xlabel=\"t\", label = \"ODE\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using LinearAlgebra\n",
    "\n",
    "fsde = function(x, p, t) \n",
    "    κ₁,κ₂,κ₃,κ₄,V = p\n",
    "    return ####                                         # drift term (same as for ODE)\n",
    "end\n",
    "gsde = function(x, p, t)\n",
    "    κ₁,κ₂,κ₃,κ₄,V = p\n",
    "    Nw = length(p)-1\n",
    "    return ####                                         # diffusion matrix (1x4)\n",
    "end\n",
    "tt, xx_sde = de_EM_solver(fode, gsde, x₀, p, T, Δt);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(tt, xx_sde, label=\"SDE\", xlabel=\"t\", ylabel=\"X(t)\")\n",
    "plot!(tt, xx_ode, colour=:black, linewidth=5, label=\"ODE\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, change the volume $V$ to $20V$. What is the effect of chaning the volume $V$?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = ####\n",
    "p = [κ₁,κ₂,κ₃,κ₄, V]\n",
    "\n",
    "tt, xx_ode = de_EM_solver(fode, gode, x₀, p, T, Δt)\n",
    "tt, xx_sde = de_EM_solver(fsde, gsde, x₀, p, T, Δt);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(tt, xx_sde, label=\"SDE\", xlabel=\"t\", ylabel=\"X(t)\")\n",
    "plot!(tt, xx_ode, colour=:black, linewidth=5, label=\"ODE\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Approximation of the solution to the chemical master equation with multiple realisations of the solution to the SDE."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let the vessel have its original volume $V=0.5$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = ####\n",
    "p = [κ₁,κ₂,κ₃,κ₄, V]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## General code for solving Itô integrals with the Euler-Maruyama method (BCs included)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "de_EM_solver_bc = function(f, g, x₀, p, T, Δt)\n",
    "\n",
    "    \"\"\"\n",
    "    f = f(x, p, t) and g = g(x, p, t), where x is the process variable, p stands for a vector of parameters and t is time.    \n",
    "    Outputs:\n",
    "        tt - a partitioned time interval [0, T]\n",
    "        xx - a solution to the stochastic integral dX_t = f(X_t, p, t)dt + g(X_t, p, t)dW\n",
    "    \"\"\"\n",
    "    \n",
    "    tt = 0.0:Δt:T\n",
    "    Nt = length(tt)\n",
    "\n",
    "    x = copy(x₀)\n",
    "    Nx = length(x) \n",
    "    xx = zeros(length(tt), Nx)\n",
    "\n",
    "    xx[1, :] .= x\n",
    "    \n",
    "    Nw = length(p)-1\n",
    "    \n",
    "    for n in 1:length(tt)-1\n",
    "        t⁻ = tt[n]\n",
    "        x⁻ = xx[n, :]\n",
    "        \n",
    "        ####\n",
    "        #### check whether the particle has escaped from the positive quadrant\n",
    "        ####\n",
    "        \n",
    "    end\n",
    "    \n",
    "    return tt, xx\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Nr = 5000\n",
    "\n",
    "tt = 0.0:Δt:T\n",
    "Nt = length(tt)\n",
    "\n",
    "Nx = length(x₀)\n",
    "\n",
    "using SharedArrays\n",
    "count_sde = SharedArray{Float64}(Nr, Nt, Nx);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for nr in 1:Nr\n",
    "    tt, xx_sde = de_EM_solver_bc(fode, gsde, x₀, p, T, Δt);\n",
    "    count_sde[nr, :] = xx_sde\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "count_sde[1, :]'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Place the particle positions in histogram bins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Nb = ####                                           # number of bins\n",
    "upperbound = ###                                   # upper limit for the copy number of A\n",
    "xx = range(.0, stop = upperbound, length = Nb+1)   # bin limits\n",
    "Δx = xx[2]-xx[1]\n",
    "xx_centre = xx[1:end-1] .+ Δx/2                    # bin centroids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "particles_counter = function(x::AbstractVector, xedges::AbstractRange)\n",
    "    # Calculate bin index from x value\n",
    "    N = size(x, 1)\n",
    "    nxbins = length(xedges)-1\n",
    "    xmin, xmax = extrema(xedges)\n",
    "    δ𝑖δx = nxbins / (xmax - xmin)\n",
    "    \n",
    "    particle_density = zeros(nxbins)\n",
    "   \n",
    "    # Calculate the means for each bin, ignoring NaNs\n",
    "    @inbounds for n ∈ eachindex(x)\n",
    "        𝑖 = (x[n] - xmin) * δ𝑖δx\n",
    "        if (0 <= 𝑖 < nxbins)  \n",
    "           i = ceil(Int, 𝑖)\n",
    "           particle_density[i] += ####\n",
    "        end\n",
    "    end\n",
    "  \n",
    "    Δx = xedges[2]-xedges[1]\n",
    "    return ####\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### For the full evolution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pp_sde = zeros(Nt, Nb) # number or time intervals and number of centroids\n",
    "\n",
    "for t_id in 1:Nt\n",
    "    pp_sde[t_id, :] = particles_counter(count_sde[:, t_id], xx)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(xx_centre, pp_sde[2, :]./maximum(pp_sde[2, :]))\n",
    "plot!(xx_centre, pp_sde[20, :]./maximum(pp_sde[20, :]))\n",
    "plot!(xx_centre, pp_sde[501, :]./maximum(pp_sde[501, :]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Plots\n",
    "anim = @animate for k in 1:length(tt)\n",
    "    plot(xx_centre, pp_sde[k, :], linetype=:steppost,\n",
    "            title=\"t: \".*\"$(round(tt[k], digits = 1))\".*\" minutes\",\n",
    "            label = \"Empirical distribution with SDEs\",\n",
    "            colour = :black,\n",
    "            xlabel = \"nA\",\n",
    "            ylabel = \"p(nA)\") # 2:end since end = 1, periodic condition\n",
    "end\n",
    "gif(anim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fokker-Planck for the same scenario\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#import Pkg; Pkg.add([\"DiffEqOperators\", \"DifferentialEquations\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "LoadError",
     "evalue": "UndefVarError: xx_centre not defined",
     "output_type": "error",
     "traceback": [
      "UndefVarError: xx_centre not defined",
      "",
      "Stacktrace:",
      " [1] top-level scope",
      "   @ In[3]:7",
      " [2] eval",
      "   @ ./boot.jl:360 [inlined]",
      " [3] include_string(mapexpr::typeof(REPL.softscope), mod::Module, code::String, filename::String)",
      "   @ Base ./loading.jl:1116"
     ]
    }
   ],
   "source": [
    "x₀ = zeros(1)\n",
    "\n",
    "using Distributions\n",
    "init_dist = Uniform(x₀[]-1.0, x₀[]+1.0)\n",
    "#init_dist = Dirac(x₀[])\n",
    "prior_pdf = x -> pdf(init_dist, x)\n",
    "pp₀ = prior_pdf.(xx_centre);\n",
    "pp₀ ./= sum(pp₀*Δx)\n",
    "\n",
    "\n",
    "ffunc, gfunc = fsde, gsde\n",
    "\n",
    "M = size(Dfunc(zeros(length(x₀)), p, 0.), 2)\n",
    "ff = Array{Float64}(undef, Nb)\n",
    "gg = Array{Float64}(undef, Nb, M, M)\n",
    "\n",
    "for n in 1:Nb   \n",
    "    ff[n] = ffunc(xx_centre[n], p, 0.) \n",
    "    auxtensor = gfunc(xx_centre[n], p, 0.)    \n",
    "    auxtensor = diagm(auxtensor[1,:])\n",
    "    gg[n, :, :] = auxtensor*auxtensor'\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "LoadError",
     "evalue": "UndefVarError: Δx not defined",
     "output_type": "error",
     "traceback": [
      "UndefVarError: Δx not defined",
      "",
      "Stacktrace:",
      " [1] top-level scope",
      "   @ In[4]:8",
      " [2] eval",
      "   @ ./boot.jl:360 [inlined]",
      " [3] include_string(mapexpr::typeof(REPL.softscope), mod::Module, code::String, filename::String)",
      "   @ Base ./loading.jl:1116"
     ]
    }
   ],
   "source": [
    "using DiffEqOperators\n",
    "\n",
    "# Second order approximation to the second derivative\n",
    "order = 2\n",
    "\n",
    "deriv = 1\n",
    "#deriv1 = UpwindDifference(deriv, order, Δx, Nx+1)\n",
    "deriv1 = CenteredDifference(deriv, order, Δx, Nb)\n",
    "\n",
    "deriv = 2\n",
    "deriv2 = CenteredDifference(deriv, order, Δx, Nb)\n",
    "\n",
    "#bc = DirichletBC(0.,0.)\n",
    "#nc = NeumannBC(0.,0.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "fpsystem! (generic function with 1 method)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "FP_assemble = function(pp, μ, D)\n",
    "    pp ./= sum(pp*Δx)\n",
    "    \n",
    "    # Drift terms    \n",
    "    μμ = broadcast(*, μ,pp)\n",
    "    #μμ = bc*μμ\n",
    "    μμ = [0; μμ; 0]\n",
    "    \n",
    "    μμ = deriv1*μμ\n",
    "\n",
    "    # Diffusion terms\n",
    "    DD = zeros(size(μμ))\n",
    "    for m in 1:M\n",
    "        DDm = broadcast(*,0.5*D[:, m, m],pp)\n",
    "        # DDm = bc*DDm\n",
    "        DDm = [0; DDm;0]\n",
    "        DDm = deriv2*DDm\n",
    "        DD += DDm\n",
    "    end\n",
    "    \n",
    "    return -μμ+DD\n",
    "end\n",
    "\n",
    "## Construct the RHS of the PDE\n",
    "function fpsystem!(du,u,p,t)\n",
    "    du .= FP_assemble(u, ff, gg)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using DifferentialEquations\n",
    "\n",
    "sol = solve(ODEProblem(fpsystem!, pp₀ , (0, T)), Tsit5(), saveat = tt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Plots\n",
    "anim = @animate for t_idx in 1:length(tt)\n",
    "    plot(xx_centre, pp_sde[t_idx, :], linetype=:steppre,\n",
    "            title=\"t: \".*\"$(round(tt[t_idx], digits = 1))\".*\" minutes\",\n",
    "            label = \"Empirical distribution with SDEs\",\n",
    "            colour = :black,\n",
    "            xlabel = \"nA\",\n",
    "            ylabel = \"p(nA)\") # 2:end since end = 1, periodic condition\n",
    "    plot!(xx_centre, sol(tt[t_idx]), label = \"Solution to the FP equation\", colour = :darkorange)\n",
    "end\n",
    "gif(anim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Homework"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Distributions\n",
    "\n",
    "function velocityjump_2d(x₀, θ₀, dτ, κ, maxnumberofjumps)\n",
    "\n",
    "  x = copy(x₀)\n",
    "  θ = copy(θ₀)\n",
    "\n",
    "  θ_dist = VonMises(####, ####)\n",
    "\n",
    "  t = 0.0\n",
    "  tt = [t]\n",
    "  \n",
    "  xx = copy(x)     \n",
    "  \n",
    "  k=1\n",
    "  while (k <= maxnumberofjumps)\n",
    "      \n",
    "      τ = ####                                   \n",
    "      t += τ\n",
    "    \n",
    "      θ += rand(θ_dist) \n",
    "      θ_dist = VonMises(####, ####)\n",
    "        \n",
    "      x = x .+ ####\n",
    "    \n",
    "      append!(tt, t)\n",
    "      xx = cat(xx, x, dims=2)\n",
    "      k += 1\n",
    "  end\n",
    "  \n",
    "  return tt, xx'\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "λ = 1\n",
    "s = 1\n",
    "\n",
    "# 2D Brownian motion nearly uncorrelated\n",
    "dτ, κ = 0, 0.1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Case 1\n",
    "Initial position and direction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x₀ = zeros(2)\n",
    "θ₀ = 0;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tt, xx = velocityjump_2d(####)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Plots\n",
    "xbound = maximum(abs.(xx[:, 1]))\n",
    "ybound = maximum(abs.(xx[:, 2]))\n",
    "plot(xx[:, 2], xx[:, 1],  xlims=(-ybound, ybound), ylims = (-xbound, xbound), label=\"Randowm walk\", xlabel = \"y\", ylabel = \"x\", \n",
    "    title = \"dθₖ ~ vonMises(μ, κ),             μ=-dτ*sin(θₖ₋₁), dτ: \".*\"$(dτ)\".*\" and κ: \".*\"$(κ)\", size=(800,500))\n",
    "scatter!([xx[1, 1]], [xx[1, 2]], markersize=7, label=\"Initial position\", legend=:bottomright)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Case 2\n",
    "Initial position and direction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# correlated, but not preferred direction\n",
    "#dτ, κ = 0, 2\n",
    "\n",
    "# correlated, with preferred direction\n",
    "dτ, κ = 0.3, 4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tt, xx = velocityjump_2d(####)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x₀ = zeros(2)\n",
    "θ₀ = 0;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Plots\n",
    "xbound = maximum(abs.(xx[:, 1]))\n",
    "ybound = maximum(abs.(xx[:, 2]))\n",
    "plot(xx[:, 2], xx[:, 1], label=\"Randowm walk\", xlabel = \"y\", ylabel = \"x\", \n",
    "    title = \"dθₖ ~ vonMises(μ, κ),             μ=-dτ*sin(θₖ₋₁), dτ: \".*\"$(dτ)\".*\" and κ: \".*\"$(κ)\", size=(800,500))\n",
    "scatter!([xx[1, 1]], [xx[1, 2]], markersize=7, label=\"Initial position\", legend=:bottomright)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.6.2",
   "language": "julia",
   "name": "julia-1.6"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}