Multipop dispersal

.gitignore

@@ -2,3 +2,5 @@
 *.jl.*.cov
 *.jl.mem
 docs/build
+*.json
+Manifest.toml

Project.toml

@@ -11,12 +11,14 @@ DynamicGrids = "a5dba43e-3abc-5203-bfc5-584ca68d3f5b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Stencils = "264155e8-78a8-466a-aa59-c9b28c34d21a"
 
 [compat]
 Distributions = "0.21, 0.22, 0.23, 0.24, 0.25"
-DocStringExtensions = "0.8"
+DocStringExtensions = "0.8, 0.9"
 DynamicGrids = "0.21"
 Reexport = "0.2, 1.0"
+Stencils = "0.3"
 julia = "1"
 
 [extras]

src/Dispersal.jl

@@ -14,7 +14,9 @@ using DynamicGrids.ConstructionBase
 using DynamicGrids.DimensionalData
 using DynamicGrids.ModelParameters
 using DynamicGrids.Setfield
+using DynamicGrids.Stencils
 using DynamicGrids.StaticArrays
+using Stencils
 
 using DynamicGrids.ModelParameters: params
 

src/allee.jl

@@ -23,3 +23,10 @@ AlleeExtinction{R,W}(; minfounders=Param(5.0, bounds=(1.0, 200.0))) where {R,W}
 
     return N >= f ? N : zero(N)
 end
+@inline function applyrule(data, rule::AlleeExtinction, Ns::AbstractArray, I)
+    fs = get(data, rule.minfounders, I...)
+
+    return broadcast(Ns, fs) do N, f
+        N >= f ? N : zero(N)
+    end
+end

src/growth.jl

@@ -53,7 +53,7 @@ modifyrule(rule::ExponentialGrowth, data) = precalc_nsteps(rule, data)
     N > zero(N) || return zero(N)
     rt = get(data, rule.rate, I...) * rule.nsteps
 
-    return @fastmath N * exp(rt)
+    return N * exp(rt)
 end
 
 """
@@ -106,11 +106,25 @@ modifyrule(rule::LogisticGrowth, data) = precalc_nsteps(rule, data)
     rt = get(data, rule.rate, I...) * rule.nsteps
     k = get(data, rule.carrycap, I...)
 
-    if rt > zero(rt)
-        return @fastmath (N * k) / (N + (k - N) * exp(-rt))
+    new_N = if rt > zero(rt)
+        (N * k) / (N + (k - N) * exp(-rt))
     else
-        return @fastmath N * exp(rt)
+        N * exp(rt)
     end
+    return min(max(zero(new_N), new_N), k)
+end
+@inline function applyrule(data, rule::LogisticGrowth, Ns::AbstractArray, I)
+    rts = get(data, rule.rate, I...) .* rule.nsteps
+    ks = get(data, rule.carrycap, I...)
+
+    new_Ns = map(Ns, rts, ks) do N, rt, k
+        if rt > zero(rt)
+            (N * k) / (N + (k - N) * exp(-rt))
+        else
+            N * exp(rt)
+        end
+    end
+    return min.(max.(zero(eltype(new_Ns)), new_Ns), ks)
 end
 
 """

src/human.jl

@@ -306,7 +306,7 @@ Base.@kwdef struct HeirarchicalGroups{S} <: TransportMode
 end
 
 @inline function disperse!(
-    data::DG.WritableGridData, mode::HeirarchicalGroups, rule::HumanDispersal,
+    data::DG.GridData, mode::HeirarchicalGroups, rule::HumanDispersal,
     shortlist, dispersalprob, N, I
 )
     dispersed = zero(N)
@@ -320,7 +320,7 @@ end
     return dispersed
 end
 @inline function disperse!(
-    data::DG.WritableGridData, mode::BatchGroups, rule::HumanDispersal,
+    data::DG.GridData, mode::BatchGroups, rule::HumanDispersal,
     shortlist, dispersalprob, N, I
 )
     # Find the expected number of dispersers given N and dispersal prob
@@ -345,7 +345,7 @@ end
 unless it falls outside the grid or is masked, in which case we
 say the event just never happened.
 =#
-function disperse2dest!(data::DG.WritableGridData, rule, shortlist, maybedispersing)
+function disperse2dest!(data::DG.GridData, rule, shortlist, maybedispersing)
     dest_id = min(length(shortlist), searchsortedfirst(shortlist, rand()))
     # Randomise cell destination within upsampled destination cells
     upsample = upsample_index(shortlist[dest_id].index, rule.scale)

src/kernel/distancemethods.jl

@@ -33,7 +33,8 @@ Joseph D. Chipperfield et al 2011"
 """
 struct CentroidToCentroid <: DistanceMethod end
 
-dispersalprob(f, ::CentroidToCentroid, x, y, cellsize) = f(sqrt(x^2 + y^2) * cellsize)
+dispersalprob(f, ::CentroidToCentroid, x, y, cellsize::T) where T =
+    f(T(sqrt(x^2 + y^2)) * cellsize)
 
 """
     AreaToCentroid <: DistanceMethod
@@ -52,15 +53,15 @@ Base.@kwdef struct AreaToCentroid{SS<:Number} <: DistanceMethod
     subsampling::SS = Param(10.0; bounds=(2.0, 40.0))
 end
 
-@inline function dispersalprob(f, dm::AreaToCentroid, x, y, cellsize)
-    prob = zero(cellsize)
-    centerfirst = 1 / subsampling(dm) / 2 - 0.5
+@inline function dispersalprob(f, dm::AreaToCentroid, x, y, cellsize::T) where T
+    prob = zero(T)
+    centerfirst = 1 / subsampling(dm) / 2 - oneunit(T) / 2
     centerlast = centerfirst * -1
     range = LinRange(centerfirst, centerlast, round(Int, subsampling(dm)))
     for j in range, i in range
         prob += sqrt((x + i)^2 + (y + j)^2) * cellsize |> f
     end
-    prob / subsampling(dm)^2
+    return T(prob / subsampling(dm)^2)
 end
 
 """
@@ -79,16 +80,16 @@ Base.@kwdef struct AreaToArea{SS<:Number} <: DistanceMethod
     subsampling::SS = Param(10.0; bounds=(2.0, 40.0))
 end
 
-@inline function dispersalprob(f, dm::AreaToArea, x, y, cellsize)
-    prob = zero(cellsize)
+@inline function dispersalprob(f, dm::AreaToArea, x, y, cellsize::T) where T
+    prob = zero(T)
     # Get the center point of the first cell (for both dimensions)
-    centerfirst = 1 / subsampling(dm) / 2 - 0.5
+    centerfirst = 1 / subsampling(dm) / 2 - oneunit(T) / 2
     centerlast = centerfirst * -1
     range = LinRange(centerfirst, centerlast, round(Int, subsampling(dm)))
     for i in range, j in range
         for a in range, b in range
-            prob += sqrt((x + i + a)^2 + (y + j + b)^2) * cellsize |> f
+            prob += T(sqrt((x + i + a)^2 + (y + j + b)^2) * cellsize) |> f
         end
     end
-    prob / subsampling(dm)^4
+    return T(prob / subsampling(dm)^4)
 end

src/kernel/formulations.jl

@@ -30,7 +30,7 @@ where λ is a shape parameter.
 Base.@kwdef struct ExponentialKernel{P} <: KernelFormulation
     λ::P = Param(1.0, bounds=(0.0, 2.0))
 end
-(f::ExponentialKernel)(d) = exp(-d / f.λ)
+(f::ExponentialKernel)(d) = exp.(-d ./ f.λ)
 
 """
     GeometricKernel <: KernelFormulation
@@ -50,7 +50,7 @@ where α is a shape parameter.
 Base.@kwdef struct GeometricKernel{P} <: KernelFormulation
     α::P = Param(1.0, bounds=(-1000.0, 1000.0))
 end
-(f::GeometricKernel)(d) = (1 + d)^f.α * ((f.α + 1)*(f.α + 2)) / (2 * π)
+(f::GeometricKernel)(d) = (1 + d) .^ f.α .* ((f.α .+ 1) .* (f.α .+ 2)) ./ (2 * π)
 
 """
     GaussianKernel <: KernelFormulation
@@ -68,12 +68,12 @@ where α is a positive parameter.
 Base.@kwdef struct GaussianKernel{P} <: KernelFormulation
     α::P = Param(1.0, bounds=(0.0, 1000.0))
 end
-(f::GaussianKernel)(d) = 1 / (π * f.α^2) * exp(- d^2 / f.α^2)
+(f::GaussianKernel)(d) = 1 ./ (π * f.α .^ 2) .* exp.(-d ^ 2 ./ f.α .^ 2)
 
 """
     WeibullKernel <: KernelFormulation
 
-    WeibullKernel(α,β)
+    WeibullKernel(α, β)
 
 Probability density function of distance ``d``.
 
@@ -87,4 +87,4 @@ Base.@kwdef struct WeibullKernel{A,B} <: KernelFormulation
     α::A = Param(1.0, bounds=(0.0, 1000.0))
     β::B = Param(2.0, bounds=(0.0, 1000.0))
 end
-(f::WeibullKernel)(d) = f.β / (2 * π * f.α^2) * d^(f.β - 2) * exp(- d^f.β / f.α^f.β)
+(f::WeibullKernel)(d) = f.β ./ (2 * π .* f.α .^ 2) .* d .^ (f.β .- 2) * exp.(-d .^ f.β ./ f.α .^ f.β)

src/kernel/inwards.jl

@@ -1,9 +1,9 @@
 """
-    InwardsPopulationDispersal <: NeighborhoodRule
+    InwardsDispersal <: NeighborhoodRule
 
-    InwardsPopulationDispersal(; kw...)
-    InwardsPopulationDispersal{R}(; kw...)
-    InwardsPopulationDispersal{R,W}(; kw...)
+    InwardsDispersal(; kw...)
+    InwardsDispersal{R}(; kw...)
+    InwardsDispersal{R,W}(; kw...)
 
 Implements deterministic dispersal from populations in neighboring cells to the current 
 cell.
@@ -28,8 +28,8 @@ grid.
 
 Pass grid `Symbol`s to `R` or both `R` and `W` type parameters to use to specific grids.
 """
-struct InwardsDispersal{R,W,N<:AbstractKernelNeighborhood} <: NeighborhoodRule{R,W}
-    neighborhood::N
+struct InwardsDispersal{R,W,S<:Stencils.AbstractKernelStencil} <: NeighborhoodRule{R,W}
+    stencil::S
 end
 function InwardsDispersal{R,W}(; kw...) where {R,W}
     InwardsDispersal{R,W}(DispersalKernel(; kw...))

src/kernel/kernel.jl

@@ -1,95 +1,102 @@
 """
-    DispersalKernel <: AbstractKernelNeighborhood
+    DispersalKernel <: AbstractKernelStencil
 
     DispersalKernel(; kw...)
 
-Dispersal kernel for taking the dot product of the neighborhood and a matching
-kernel of weights. May hold any `Neighborhood` object: the kernel
+Dispersal kernel for taking the dot product of the stencil and a matching
+kernel of weights. May hold any `Stencil` object: the kernel
 will be built to match the shape, using the `folumation`, `cellsize`
 and `distancemethod`.
 
 # Keyword Arguments
 
-- `neighborhood`: Any DynamicGrids.jl `Neighborhood`, or an
+- `stencil`: Any DynamicGrids.jl `Stencil`, or an
     already constructed [`DispersalKernel`](@ref). Using this keyword means `radius` is
     ignored, and for a `DispersalKernel`, all other keywords are ignored.
-- `neighborhood`: `Neighborhood` object specifying the range from the origin of the
-    discretised dispersal kernal. Defaults to `Window(radius)`.
 - `formulation`: kernel formulation object holding the exact form of the kernal.
     Default [`ExponentialKernel`](@ref).
 - `cellsize`: the cell size of the discretised kernal (i.e. simulation grid size).
     Default is 1.0.
 - `distancemethod`: [`DistanceMethod`](@ref) object for calculating distance between cells.
     The default is [`CentroidToCentroid`](@ref).
 """
-struct DispersalKernel{R,N,L,H<:Neighborhood{R,N,L},K,F,C,D<:DistanceMethod} <: AbstractKernelNeighborhood{R,N,L,H}
-    neighborhood::H
+struct DispersalKernel{R,N,L,T,H<:Stencil{R,N,L,T},K,F,C,D<:DistanceMethod} <: Stencils.AbstractKernelStencil{R,N,L,T,H}
+    stencil::H
     kernel::K
     formulation::F
     cellsize::C
     distancemethod::D
     function DispersalKernel(
-        hood::H, kernel, formulation::F, cellsize::C, distancemethod::D
-    ) where {H<:Neighborhood{R,N,L},F,C,D<:DistanceMethod} where {R,N,L}
-        if hood isa DispersalKernel
-            hood
+        stencil::H, kernel, formulation::F, cellsize::C, distancemethod::D
+    ) where {H<:Stencil{R,N,L,T},F,C,D<:DistanceMethod} where {R,N,L,T}
+        if stencil isa DispersalKernel
+            stencil
         else
             # Build the kernel matrix
-            newkernel = scale(buildkernel(hood, formulation, distancemethod, cellsize))
-            new{R,N,L,H,typeof(newkernel),F,C,D}(
-                hood, newkernel, formulation, cellsize, distancemethod
+            newkernel = buildkernel(stencil, formulation, distancemethod, cellsize) |> scale
+            new{R,N,L,T,H,typeof(newkernel),F,C,D}(
+                stencil, newkernel, formulation, cellsize, distancemethod
             )
         end
     end
-    function DispersalKernel{R,N,L,H,K,F,C,D}(
-        hood::H, kernel::K, formulation::F, cellsize::C, distancemethod::D
-    ) where {R,N,L,H,K,F,C,D}
-        new{R,N,L,H,K,F,C,D}(hood, kernel, formulation, cellsize, distancemethod)
+    function DispersalKernel{R,N,L,T,H,K,F,C,D}(
+        stencil::H, kernel::K, formulation::F, cellsize::C, distancemethod::D
+    ) where {R,N,L,T,H,K,F,C,D}
+        new{R,N,L,T,H,K,F,C,D}(stencil, kernel, formulation, cellsize, distancemethod)
     end
 end
 function DispersalKernel(;
     radius=1,
-    neighborhood=Window{radius}(),
+    stencil=Window{radius}(),
     formulation=ExponentialKernel(),
     cellsize=1.0,
     distancemethod=CentroidToCentroid(),
 )
-    DispersalKernel(neighborhood, nothing, formulation, cellsize, distancemethod)
+    DispersalKernel(stencil, nothing, formulation, cellsize, distancemethod)
 end
 DispersalKernel{R}(; radius=R, kw...) where R = DispersalKernel(; radius=radius, kw...)
 
 ConstructionBase.constructorof(::Type{<:DispersalKernel}) = DispersalKernel
 
-function DG.setwindow(n::DispersalKernel{R,N,L,<:Any,K,F,C,D}, buffer) where {R,N,L,K,F,C,D}
-    newhood = DG.setwindow(neighborhood(n), buffer)
-    DispersalKernel{R,N,L,typeof(newhood),K,F,C,D}(
-        newhood, kernel(n), formulation(n), cellsize(n), distancemethod(n)
+function DG.Stencils.rebuild(n::DispersalKernel{R,N,L,<:Any,<:Any,K,F,C,D}, buffer) where {R,N,L,K,F,C,D}
+    newstencil = Stencils.rebuild(stencil(n), buffer)
+    DispersalKernel{R,N,L,eltype(newstencil),typeof(newstencil),K,F,C,D}(
+        newstencil, kernel(n), formulation(n), cellsize(n), distancemethod(n)
     )
 end
+DG.Stencils.neighbors(hood::DispersalKernel) = neighbors(hood.stencil)
 
-cellsize(hood::DispersalKernel) = hood.cellsize
-distancemethod(hood::DispersalKernel) = hood.distancemethod
-formulation(hood::DispersalKernel) = hood.formulation
+cellsize(stencil::DispersalKernel) = stencil.cellsize
+distancemethod(stencil::DispersalKernel) = stencil.distancemethod
+formulation(stencil::DispersalKernel) = stencil.formulation
 
-function buildkernel(window::Window{R}, formulation, distancemethod, cellsize) where R
-    # The radius doesn't include the center cell, so add it
-    S = 2R + 1
-    kernel = zeros(typeof(cellsize), S, S)
-    r1 = R + one(R)
-    # Paper: l. 97
-    for x = 0:R, y = 0:R
-        # Calculate the distance effect from the center cell to this cell
-        prob = dispersalprob(formulation, distancemethod, x, y, cellsize)
-        # Update the kernel value based on the formulation and distance
-        kernel[ x + r1,  y + r1] = prob
-        kernel[ x + r1, -y + r1] = prob
-        kernel[-x + r1,  y + r1] = prob
-        kernel[-x + r1, -y + r1] = prob
+# function buildkernel(window::Window{R}, formulation, distancemethod, cellsize) where R
+#     # The radius doesn't include the center cell, so add it
+#     S = 2R + 1
+#     kernel = zeros(typeof(cellsize), S, S)
+#     r1 = R + one(R)
+#     # Paper: l. 97
+#     for x = 0:R, y = 0:R
+#         # Calculate the distance effect from the center cell to this cell
+#         prob = dispersalprob(formulation, distancemethod, x, y, cellsize)
+#         # Update the kernel value based on the formulation and distance
+#         kernel[ x + r1,  y + r1] = prob
+#         kernel[ x + r1, -y + r1] = prob
+#         kernel[-x + r1,  y + r1] = prob
+#         kernel[-x + r1, -y + r1] = prob
+#     end
+#     SMatrix{S,S}(kernel)
+# end
+function buildkernel(stencil::Stencil{<:Any,<:Any,L}, f, dm, cellsize) where L
+    probs = map(offsets(stencil)) do (x, y)
+        dispersalprob(f, dm, x, y, cellsize)
     end
-    SMatrix{S,S}(kernel)
-end
-function buildkernel(window::Neighborhood{<:Any,<:Any,L}, f, dm, cellsize) where L
-    SVector{L}(Tuple(dispersalprob(f, dm, x, y, cellsize) for (x, y) in offsets(window)))
+    return SVector{L}(probs)
 end
 
-scale(x) = x ./ sum(x)
+function scale(xs)
+    s = sum(xs)
+    map(xs) do x 
+        x ./ s
+    end
+end