From ec7facf623b19719bf3d5d06e45b65a77e6372fd Mon Sep 17 00:00:00 2001
From: Rafael Schouten <rafaelschouten@gmail.com>
Date: Fri, 4 Oct 2024 15:16:02 +0200
Subject: [PATCH] generalise dispersal to multiple populations

---
 src/kernel/formulations.jl | 10 +++++-----
 src/kernel/inwards.jl      |  8 ++++----
 src/kernel/kernel.jl       | 14 +++++++++++---
 src/kernel/outwards.jl     | 18 +++++++++---------
 4 files changed, 29 insertions(+), 21 deletions(-)

diff --git a/src/kernel/formulations.jl b/src/kernel/formulations.jl
index c5c0efb..4e0db47 100644
--- a/src/kernel/formulations.jl
+++ b/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.β)
diff --git a/src/kernel/inwards.jl b/src/kernel/inwards.jl
index 7e5647a..64d2f11 100644
--- a/src/kernel/inwards.jl
+++ b/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.
diff --git a/src/kernel/kernel.jl b/src/kernel/kernel.jl
index 0a07ab7..29a808f 100644
--- a/src/kernel/kernel.jl
+++ b/src/kernel/kernel.jl
@@ -33,7 +33,7 @@ struct DispersalKernel{R,N,L,T,H<:Stencil{R,N,L,T},K,F,C,D<:DistanceMethod} <: S
             stencil
         else
             # Build the kernel matrix
-            newkernel = scale(buildkernel(stencil, formulation, distancemethod, cellsize))
+            newkernel = buildkernel(stencil, formulation, distancemethod, cellsize) |> scale
             new{R,N,L,T,H,typeof(newkernel),F,C,D}(
                 stencil, newkernel, formulation, cellsize, distancemethod
             )
@@ -88,7 +88,15 @@ formulation(stencil::DispersalKernel) = stencil.formulation
 #     SMatrix{S,S}(kernel)
 # end
 function buildkernel(stencil::Stencil{<:Any,<:Any,L}, f, dm, cellsize) where L
-    SVector{L}(Tuple(dispersalprob(f, dm, x, y, cellsize) for (x, y) in offsets(stencil)))
+    probs = map(offsets(stencil)) do (x, y)
+        dispersalprob(f, dm, x, y, cellsize)
+    end
+    return SVector{L}(probs)
 end
 
-scale(x) = x ./ sum(x)
+function scale(xs)
+    s = sum(xs)
+    map(xs) do x 
+        x ./ s
+    end
+end
diff --git a/src/kernel/outwards.jl b/src/kernel/outwards.jl
index 97b5baa..d9b68c0 100644
--- a/src/kernel/outwards.jl
+++ b/src/kernel/outwards.jl
@@ -1,9 +1,9 @@
 """
-    OutwardsPopulationDispersal <: SetNeighborhoodRule
+    OutwardsDispersal <: SetNeighborhoodRule
 
-    OutwardsPopulationDispersal(; kw...)
-    OutwardsPopulationDispersal{R}(; kw...)
-    OutwardsPopulationDispersal{R,W}(; kw...)
+    OutwardsDispersal(; kw...)
+    OutwardsDispersal{R}(; kw...)
+    OutwardsDispersal{R,W}(; kw...)
 
 Implements deterministic dispersal from the current cell to populations in neighboring
 cells.
@@ -29,10 +29,10 @@ is occupied.
     Default is 1.0.
 - `distancemethod`: [`DistanceMethod`](@ref) object for calculating distance between cells.
     The default is [`CentroidToCentroid`](@ref).
-- `maskbehavior`: The default setting is `IgnoreMaskEdges()`. Use `CheckMaskEdges()` to indicate that the grid is 
-    masked, enabling the rule to perform boundary checking at mask edges. Not using 
-    `CheckMaskEdges()` on a masked grid may result in the loss of individuals at the edges, but it comes
-    at a performance cost.
+- `maskbehavior`: The default is `IgnoreMaskEdges()`. Use `CheckMaskEdges()` 
+    to enabling the rule to perform boundary checking at mask edges. 
+    Using `IgnoreMaskEdges()` on a masked grid may result in the loss of individuals at the edges, 
+    while using `CheckMaskEdges()` should not. However, this comes at a performance cost.
 
 Pass grid name `Symbol`s to `R` and `W` type parameters to use specific grids.
 """
@@ -69,7 +69,7 @@ end
         target = I .+ offset
         (target_mod, inbounds) = DynamicGrids.inbounds(data, target)
         if inbounds && (isnothing(mask_data) || mask_data[target_mod...])
-            @inbounds propagules = N * k  
+            @inbounds propagules = N .* k
             @inbounds add!(data[W], propagules, target_mod...)  
             sum += propagules
         end