Figures_S3-5.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Packages and Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Distributions\n",
    "using Plots\n",
    "using Random\n",
    "using Statistics\n",
    "using JLD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "simpler_jump (generic function with 2 methods)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function simpler_jump(x, y, a = 1)\n",
    "    theta = rand(Uniform(0, 2 * π))\n",
    "    u = rand(Uniform(0, 1))\n",
    "\n",
    "    r = a/(3 * u)^(1/3)\n",
    "\n",
    "    xfl = r * sin(theta)\n",
    "    yfl = r * cos(theta)\n",
    "\n",
    "    (floor(Int, x + xfl + 1/2),\n",
    "    floor(Int, y + yfl + 1/2), r)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "gridoffsetandsize (generic function with 1 method)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function gridoffsetandsize(infectors, removed, susceptibles)\n",
    "    xoffset = yoffset = 0\n",
    "\n",
    "    data = vcat(infectors, removed, susceptibles)\n",
    "\n",
    "    xcoords = map(x -> x[1], data)\n",
    "    ycoords = map(x -> x[2], data)\n",
    "\n",
    "    (xmin, xmax) = (minimum(xcoords), maximum(xcoords))\n",
    "    (ymin, ymax) = (minimum(ycoords), maximum(ycoords))\n",
    "\n",
    "    if xmin < 1\n",
    "        xoffset = abs(xmin)\n",
    "    end\n",
    "\n",
    "    if ymin < 1\n",
    "        yoffset = abs(ymin)\n",
    "    end\n",
    "\n",
    "    gridsize = max(xmax + xoffset, ymax + yoffset)\n",
    "\n",
    "    # add 1 because you need to include the point 0\n",
    "    (xoffset + 1, yoffset + 1, gridsize + 1)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "visualizegrid (generic function with 1 method)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function creategrid(infectors, removed, visited, xoffset, yoffset, gridsize)\n",
    "    grid = zeros(gridsize, gridsize)\n",
    "\n",
    "    infector_coords = map(x -> x[1:2], infectors)\n",
    "    removed_coords = map(x -> x[1:2], removed)\n",
    "    visited_coords = map(x -> x[1:2], visited)\n",
    "    \n",
    "    for (x, y) in visited_coords\n",
    "        grid[x + xoffset, y + yoffset] = 1\n",
    "    end\n",
    "\n",
    "    for (x, y) in removed_coords\n",
    "        grid[x + xoffset, y + yoffset] = 2\n",
    "    end\n",
    "\n",
    "    for (x, y) in infector_coords\n",
    "        grid[x + xoffset, y + yoffset] = 3\n",
    "    end\n",
    "    \n",
    "    if !(1 in grid)\n",
    "        for (x, y) in removed_coords\n",
    "            grid[x + xoffset, y + yoffset] = 1\n",
    "        end\n",
    "\n",
    "        for (x, y) in infector_coords\n",
    "            grid[x + xoffset, y + yoffset] = 2\n",
    "        end\n",
    "    end\n",
    "\n",
    "    grid\n",
    "end\n",
    "\n",
    "function visualizegrid(grid, time)\n",
    "    if 1 in grid && 2 in grid && 3 in grid\n",
    "        colorgradient = cgrad([:black, :blue, :white, :red])\n",
    "    elseif 1 in grid && 2 in grid\n",
    "        colorgradient = cgrad([:black, :white, :red])\n",
    "    else\n",
    "        colorgradient = cgrad([:black, :white])\n",
    "    end\n",
    "\n",
    "    heatmap(1:size(grid, 1),\n",
    "            1:size(grid, 2),\n",
    "            grid,\n",
    "            legend = :none,\n",
    "            title=\"Time step $(time).\",\n",
    "            c = colorgradient,\n",
    "            size=(650, 600), colorbar = false)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "outbreakwithjumps_varp (generic function with 2 methods)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function outbreakwithjumps_varp(maxtime :: Int,\n",
    "                           maxtau :: Int,\n",
    "                           nextstepfcn :: Function,\n",
    "                           probinfected :: Function,\n",
    "                           p_arr :: Array{Float64, 1},\n",
    "                           graphcapturetimes :: Array{Int, 1},\n",
    "                           maxgridsize :: Int,\n",
    "                           addinfectors = true) :: Tuple{Array{Int, 1},\n",
    "                                                   Array{Array{Int}, 1},\n",
    "                                                   Array{Array{Int}, 1},\n",
    "                                                   Array{Any, 1},\n",
    "                                                   Array{Int64, 1},\n",
    "                                                   Array{Int64, 1}, \n",
    "                                                   Array{Int64, 1}}\n",
    "    # Tracking the different agents\n",
    "    x_start = 0\n",
    "    y_start = 0\n",
    "\n",
    "    newinfections = Array{Int64, 1}()\n",
    "    infectors = [[x_start, y_start, 0, 0]]\n",
    "    numinfected = [length(infectors)]\n",
    "    removed = Set([[x_start, y_start]])\n",
    "    visited = Set([[x_start, y_start]])\n",
    "    numvisited = zeros(Int, maxtime)\n",
    "    numsusceptiblesvisited = zeros(Int, maxtime)\n",
    "    r_inf = Array{Tuple{Int64, Int64}, 1}()\n",
    "    graphcomponents = Array{Any, 1}()\n",
    "    \n",
    "\n",
    "    for time in 1:maxtime\n",
    "        ninfectors = length(infectors)\n",
    "        ninfected = 0\n",
    "        newinfections_now = 0\n",
    "        \n",
    "        if time != 1\n",
    "            numvisited[time] = numvisited[time - 1]\n",
    "            numsusceptiblesvisited[time] = numsusceptiblesvisited[time - 1]\n",
    "        end\n",
    "\n",
    "        for i in 1:ninfectors\n",
    "            infector = infectors[i]\n",
    "            (infector[1], infector[2]) = simpler_jump(infector[1], infector[2])\n",
    "            infector[3] += 1\n",
    "\n",
    "            inremoved = infector[1:2] in removed\n",
    "            alreadyvisited = infector[1:2] in visited\n",
    "            \n",
    "            if !alreadyvisited\n",
    "                push!(visited, infector[1:2])\n",
    "                numvisited[time] += 1\n",
    "            end\n",
    "            \n",
    "            isinfecting = rand() <= probinfected(infector[1], p_arr[1], p_arr[2])\n",
    "            \n",
    "            if !alreadyvisited && !isinfecting\n",
    "                numsusceptiblesvisited[time] += 1\n",
    "            end\n",
    "            \n",
    "            if !inremoved && isinfecting && alreadyvisited\n",
    "                numsusceptiblesvisited[time] -= 1\n",
    "            end\n",
    "            \n",
    "            if !inremoved && isinfecting\n",
    "                push!(removed, infector[1:2])\n",
    "            end\n",
    "\n",
    "            if !inremoved && isinfecting\n",
    "                newinfections_now += 1\n",
    "                if !addinfectors\n",
    "                    ninfected += 1\n",
    "                else\n",
    "                    push!(infectors, [infector[1:2]..., 0, 0])\n",
    "            end\n",
    "\n",
    "                infector[4] += 1\n",
    "            end\n",
    "        end\n",
    "        \n",
    "        push!(newinfections, newinfections_now)\n",
    "\n",
    "        infectedtoremove = filter(x -> x[3] >= maxtau, infectors)\n",
    "        for infector in infectedtoremove\n",
    "            push!(r_inf, (infector[4], time))\n",
    "        end\n",
    "\n",
    "        filter!(x -> x[3] < maxtau, infectors)\n",
    "\n",
    "        if addinfectors\n",
    "            push!(numinfected, length(infectors))\n",
    "        else\n",
    "            push!(numinfected, ninfected)\n",
    "        end\n",
    "        \n",
    "        #infector_coords = map(x -> x[1:2], infectors)\n",
    "        \n",
    "        #visited_now = filter(x -> (x !in removed), visited)\n",
    "        #visited_now = filter(x -> (x !in infector_coords), visited_now)\n",
    "        \n",
    "        if time in graphcapturetimes\n",
    "            push!(graphcomponents, [copy(time), deepcopy(infectors), collect(removed), collect(visited)])\n",
    "        end\n",
    "    end\n",
    "    \n",
    "    \n",
    "    removed = collect(removed)\n",
    "#=\n",
    "    comparison of the old code (first two lines) and the new one which calculated the gridsize at t = 180 \n",
    "    (second two lines). the important result is that there were occasionally different results which led to \n",
    "    the bounds error.\n",
    "(xoffset, yoffset, gridsize) = reduce(map(x -> gridoffsetandsize(x[2:end]...), graphcomponents)) do x, y\n",
    "        x[3] > y[3] ? x : y\n",
    "    end\n",
    "\n",
    "(_, _, gridsizeat180) = gridoffsetandsize(graphcomponents[5][2], graphcomponents[5][3])\n",
    "\n",
    "(gridsizeat180, gridsize)=#\n",
    "\n",
    "    \n",
    "    graphs = Array{Any, 1}()\n",
    "    for i in 1:length(graphcomponents)\n",
    "        (xoffset, yoffset, gridsize) = gridoffsetandsize(graphcomponents[i][2], graphcomponents[i][3], graphcomponents[i][4])\n",
    "        xoffset = round(Int, maxgridsize/2)\n",
    "        yoffset = round(Int, maxgridsize/2)\n",
    "        gridsize = maxgridsize\n",
    "        push!(graphs, creategrid(graphcomponents[i][2], graphcomponents[i][3], graphcomponents[i][4], xoffset, yoffset, gridsize))\n",
    "    end\n",
    "\n",
    "    (numinfected, infectors, removed, graphs, newinfections, numsusceptiblesvisited, numvisited)\n",
    " \n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "probinfected (generic function with 1 method)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function probinfected(x, p1, p2)\n",
    "    if x <= 0 \n",
    "        return p1\n",
    "    else\n",
    "        return p2\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "createpaddedgrid (generic function with 1 method)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function createpaddedgrid(infectors, removed, visited)\n",
    "    xoffset = yoffset = 0\n",
    "\n",
    "    data = vcat(infectors, removed, visited)\n",
    "\n",
    "    xcoords = map(x -> x[1], data)\n",
    "    ycoords = map(x -> x[2], data)\n",
    "\n",
    "    (xmin, xmax) = (minimum(xcoords), maximum(xcoords))\n",
    "    (ymin, ymax) = (minimum(ycoords), maximum(ycoords))\n",
    "\n",
    "    if xmin < 1\n",
    "        xoffset = abs(xmin)\n",
    "    end\n",
    "\n",
    "    if ymin < 1\n",
    "        yoffset = abs(ymin)\n",
    "    end\n",
    "\n",
    "    gridsize = max(xmax + xoffset, ymax + yoffset)\n",
    "\n",
    "    # add 1 because you need to include the point 0\n",
    "    xoffset = xoffset + 1\n",
    "    yoffset = yoffset + 1\n",
    "    gridsize = gridsize + 1\n",
    "\n",
    "    xoffset = xoffset * 2\n",
    "    yoffset = yoffset * 2\n",
    "    gridsize = gridsize * 2 # for padding\n",
    "\n",
    "    grid = zeros(gridsize, gridsize)\n",
    "\n",
    "    infector_coords = map(x -> x[1:2], infectors)\n",
    "    removed_coords = map(x -> x[1:2], removed)\n",
    "    visited_coords = map(x -> x[1:2], visited)\n",
    "\n",
    "    for (x, y) in visited_coords\n",
    "        grid[x + xoffset, y + yoffset] = 1\n",
    "    end\n",
    "\n",
    "    for (x, y) in removed_coords\n",
    "        grid[x + xoffset, y + yoffset] = 2\n",
    "    end\n",
    "\n",
    "    for (x, y) in infector_coords\n",
    "        grid[x + xoffset, y + yoffset] = 3\n",
    "    end\n",
    "\n",
    "    if !(1 in grid)\n",
    "        for (x, y) in removed_coords\n",
    "            grid[x + xoffset, y + yoffset] = 1\n",
    "        end\n",
    "\n",
    "        for (x, y) in infector_coords\n",
    "            grid[x + xoffset, y + yoffset] = 2\n",
    "        end\n",
    "    end\n",
    "\n",
    "    grid\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "visualizegrid_inf (generic function with 1 method)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function visualizegrid(grid, time)\n",
    "    if 1 in grid && 2 in grid && 3 in grid\n",
    "        colorgradient = cgrad([:black, :blue, :white, :red])\n",
    "    elseif 1 in grid && 2 in grid\n",
    "        colorgradient = cgrad([:black, :white, :red])\n",
    "    else\n",
    "        colorgradient = cgrad([:black, :white])\n",
    "    end\n",
    "\n",
    "    heatmap(1:size(grid, 1),\n",
    "            1:size(grid, 2),\n",
    "            grid,\n",
    "            legend = :none,\n",
    "            title=\"Time step $(time)\",\n",
    "            ticks=nothing,\n",
    "            c = colorgradient,\n",
    "            size=(650, 600), colorbar = false)\n",
    "end\n",
    "\n",
    "function visualizegrid_rem(grid, time)\n",
    "    if 1 in grid && 2 in grid && 3 in grid\n",
    "        colorgradient = cgrad([:black, :black, :white, :black])\n",
    "    elseif 1 in grid && 2 in grid\n",
    "        colorgradient = cgrad([:black, :white, :black])\n",
    "    else\n",
    "        colorgradient = cgrad([:black, :white])\n",
    "    end\n",
    "\n",
    "    heatmap(1:size(grid, 1),\n",
    "            1:size(grid, 2),\n",
    "            grid,\n",
    "            legend = :none,\n",
    "            title=\"Time step $(time)\",\n",
    "            ticks=nothing,\n",
    "            c = colorgradient,\n",
    "            size=(650, 600), colorbar = false)\n",
    "end\n",
    "\n",
    "function visualizegrid_inf(grid, time)\n",
    "    if 1 in grid && 2 in grid && 3 in grid\n",
    "        colorgradient = cgrad([:black, :black, :black, :white])\n",
    "    elseif 1 in grid && 2 in grid\n",
    "        colorgradient = cgrad([:black, :black, :white])\n",
    "    end\n",
    "\n",
    "    heatmap(1:size(grid, 1),\n",
    "            1:size(grid, 2),\n",
    "            grid,\n",
    "            legend = :none,\n",
    "            title=\"Time step $(time)\",\n",
    "            ticks=nothing,\n",
    "            c = colorgradient,\n",
    "            size=(650, 600), colorbar = false)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The code used to generate the panels in Figures S3 - S5.\n",
    "# The simulations are stochastic, so your plots may look substantially different from the figures.\n",
    "# Alternatively, you can directly load the .jld data to plot the same figures.\n",
    "\n",
    "maxtime = 1500\n",
    "times = map(x -> x, 300:300:maxtime)\n",
    "# p = 0.48\n",
    "p_arr = [0.06, 0.06] # p in the left region (x <= 0), p in the right region\n",
    "τ = 50\n",
    "maxgridsize = 1500\n",
    "\n",
    "# (infections, infectors, removed, graphs) = outbreakwithjumps(time, τ, simpler_jump, times, p)\n",
    "(infections, infectors, removed, graphs, newinfections) = outbreakwithjumps_varp(maxtime, τ, simpler_jump, probinfected, p_arr, times, maxgridsize)\n",
    "\n",
    "filename = \"p0pt06\"\n",
    "\n",
    "save(string(filename, \".jld\"), \"graphs\", graphs, \"infectors\", infectors, \"removed\", removed, \n",
    "        \"infections\", infections, \"newinfections\", newinfections)\n",
    "\n",
    "# All ---------------------------------------------------------------------\n",
    "\n",
    "plots = Array{Any, 1}()\n",
    "for (idx, time) in enumerate(times)\n",
    "    push!(plots, visualizegrid(graphs[idx], time))\n",
    "end\n",
    "\n",
    "#plot([plots1, plots2]..., size=(1400, 250), layout=(2,length(times)))\n",
    "p1 = plot(plots..., size=(1400, 250), layout=(1,length(times)))\n",
    "\n",
    "plot(p1)\n",
    "\n",
    "savefig(string(filename, \".pdf\"))\n",
    "\n",
    "# Removed Only -------------------------------------------------------------\n",
    "\n",
    "plots = Array{Any, 1}()\n",
    "for (idx, time) in enumerate(times)\n",
    "    push!(plots, visualizegrid_rem(graphs[idx], time))\n",
    "end\n",
    "\n",
    "#plot([plots1, plots2]..., size=(1400, 250), layout=(2,length(times)))\n",
    "p1 = plot(plots..., size=(1400, 250), layout=(1,length(times)))\n",
    "\n",
    "plot(p1)\n",
    "\n",
    "savefig(string(filename, \"_rem.pdf\"))\n",
    "\n",
    "# Infectors Only ------------------------------------------------------------\n",
    "\n",
    "plots = Array{Any, 1}()\n",
    "for (idx, time) in enumerate(times)\n",
    "    push!(plots, visualizegrid_inf(graphs[idx], time))\n",
    "end\n",
    "\n",
    "#plot([plots1, plots2]..., size=(1400, 250), layout=(2,length(times)))\n",
    "p1 = plot(plots..., size=(1400, 250), layout=(1,length(times)))\n",
    "\n",
    "plot(p1)\n",
    "\n",
    "savefig(string(filename, \"_inf.pdf\"))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.6.0",
   "language": "julia",
   "name": "julia-1.6"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.6.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}