Lotka_Volterra_Bayes.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d40ee5ff-bd66-48ae-84bb-0609271ec4bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "using DifferentialEquations, Plots, StatsPlots\n",
    "using CSV, DataFrames\n",
    "using Turing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3498cd5d-74ef-45d2-874a-c48fec4598a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "function lotka_volterra!(du, u, p, t)\n",
    "\n",
    "  # Unpack the values so that they have clearer meaning\n",
    "  x, y = u\n",
    "  bx, mx, by, my = p\n",
    "\n",
    "  # Define the ODE\n",
    "  du[1] = (bx - y * mx) * x\n",
    "  du[2] = (x*by - my) * y\n",
    "\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b90f5ef4-dbf5-4587-afeb-5266cf0c150f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model parameters\n",
    "p = [1.1, 0.5, 0.1, 0.2]\n",
    "\n",
    "# Initial conditions\n",
    "u0 = [1, 1]\n",
    "\n",
    "# Timespan of the solution\n",
    "tspan = (0.0, 40.0)\n",
    "\n",
    "prob = ODEProblem(lotka_volterra!, u0, tspan, p)\n",
    "\n",
    "sol = solve(prob)\n",
    "\n",
    "plot(sol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba87d54e-9986-467a-8087-b2f013417b5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load data and determine parameters\n",
    "\n",
    "data = CSV.read(\"./data/lv_pop_data.csv\", DataFrame)\n",
    "pop_data = Array(data)'\n",
    "\n",
    "time_plot=0:2:30;\n",
    "plot(time_plot, pop_data[1, :], label=false);\n",
    "plot!(time_plot, pop_data[2, :], label=false);\n",
    "scatter!(time_plot, pop_data[1, :], label=\"Prey\");\n",
    "scatter!(time_plot, pop_data[2, :], label=\"Pred\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbcfd66d-f6db-42bd-9592-80e6c8492da2",
   "metadata": {},
   "outputs": [],
   "source": [
    "@model function fitlv(data)\n",
    "\n",
    "    σ ~ InverseGamma(2, 3)\n",
    "\n",
    "    bx ~ truncated(Normal(1, 0.5), 0, 2)\n",
    "    mx ~ truncated(Normal(1, 0.5), 0, 2)\n",
    "    by ~ truncated(Normal(1, 0.5), 0, 2)\n",
    "    my ~ truncated(Normal(1, 0.5), 0, 2)\n",
    "\n",
    "    param = [bx, mx, by, my]\n",
    "    prob = ODEProblem(lotka_volterra!, u0, (0.0, 30), param)\n",
    "    predicted = solve(prob, Tsit5(), saveat=2)\n",
    "\n",
    "    for i = 1:length(predicted)\n",
    "        data[:, i] ~ MvNormal(predicted[i], σ)\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "379a82da-dc3e-4765-a5aa-e7135bcaab80",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = fitlv(pop_data)\n",
    "\n",
    "posterior = sample(model, NUTS(0.6), 10000) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddb64f32-ac52-4e71-9ebb-68b354830ba0",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(posterior)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4081d347-8d38-4498-b3a9-8481079611e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "birth_prey = sample(posterior[:bx], 100)\n",
    "mort_prey = sample(posterior[:mx], 100)\n",
    "birth_pred = sample(posterior[:by], 100)\n",
    "mort_pred = sample(posterior[:my], 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68199eb2-f36a-47d6-948f-5298792ae28f",
   "metadata": {},
   "outputs": [],
   "source": [
    "solutions = []\n",
    "\n",
    "for i in 1:length(birth_prey)\n",
    "\n",
    "    p = [birth_prey[i], mort_prey[i], birth_pred[i], mort_pred[i]];\n",
    "    problem = ODEProblem(lotka_volterra!, u0, (0.0, 30.0), p);\n",
    "    push!(solutions, solve(problem, saveat = 0.1));\n",
    "\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d676ae1-b4b7-40c1-aba3-9a1761b14769",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_mean = [mean(birth_prey), mean(mort_prey), mean(birth_pred), mean(mort_pred)];\n",
    "\n",
    "problem_mean = ODEProblem(lotka_volterra!, u0, (0.0,30.0), p_mean);\n",
    "push!(solutions, solve(problem_mean, saveat = 0.1));\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "009d2343-8b7f-4050-a857-e70e3cb05f71",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(solutions[1], alpha=0.2, color=\"blue\");\n",
    "\n",
    "for i in 2:(length(solutions) - 1)\n",
    "    plot!(solutions[i], alpha=0.2, legend=false, color=\"blue\");\n",
    "end\n",
    "\n",
    "plot!(solutions[end], lw = 2, color=\"red\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfb16ee9-b3b2-4e4a-89b0-f78d724e4f21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Comparing inference with the data\n",
    "scatter!(time_plot, pop_data[1, :], color = \"blue\");\n",
    "scatter!(time_plot, pop_data[2, :], color = \"orange\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56e17ba2-5dde-49c9-b1e0-4e1e2c522f85",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.10.2",
   "language": "julia",
   "name": "julia-1.10"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.10.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}