Fixing boundary checks on masked grids

.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, 0.9"
 DynamicGrids = "0.21"
 Reexport = "0.2, 1.0"
+Stencils = "0.3"
 julia = "1"
 
 [extras]

src/Dispersal.jl

@@ -16,6 +16,7 @@ using DynamicGrids.ModelParameters
 using DynamicGrids.Setfield
 using DynamicGrids.Stencils
 using DynamicGrids.StaticArrays
+using Stencils
 
 using DynamicGrids.ModelParameters: params
 

src/kernel/outwards.jl

@@ -29,23 +29,50 @@ is occupied.
     Default is 1.0.
 - `distancemethod`: [`DistanceMethod`](@ref) object for calculating distance between cells.
     The default is [`CentroidToCentroid`](@ref).
+- `mask_flag`: The default setting is `NoMask()`. Use `Mask()` to indicate that the grid is 
+    masked, enabling the rule to perform boundary checking at mask edges. Not using 
+    `Mask()` on a masked grid may result in the loss of individuals at the edges, but it comes
+    at a performance cost.
 
 Pass grid name `Symbol`s to `R` and `W` type parameters to use specific grids.
 """
-struct OutwardsDispersal{R,W,S<:Stencils.AbstractKernelStencil} <: SetNeighborhoodRule{R,W}
+
+struct Mask end
+struct NoMask end
+
+struct OutwardsDispersal{R,W,S<:Stencils.AbstractKernelStencil, M} <: SetNeighborhoodRule{R,W}
     stencil::S
+    mask_flag::M
+end
+
+# Constructors for OutwardsDispersal
+function OutwardsDispersal{R,W}(stencil::S; mask_flag::Union{Mask, NoMask}=NoMask()) where {R,W,S<:Stencils.AbstractKernelStencil}
+    OutwardsDispersal{R,W,S,typeof(mask_flag)}(stencil, mask_flag)
 end
-function OutwardsDispersal{R,W}(; kw...) where {R,W}
-    OutwardsDispersal{R,W}(DispersalKernel(; kw...))
+
+function OutwardsDispersal{R,W}(; mask_flag::Union{Mask, NoMask}=NoMask(), kw...) where {R,W}
+    stencil = DispersalKernel(; kw...)
+    OutwardsDispersal{R,W,typeof(stencil),typeof(mask_flag)}(stencil, mask_flag)
 end
 
 @inline function applyrule!(data, rule::OutwardsDispersal{R,W}, N, I) where {R,W}
     N == zero(N) && return nothing
+
+    # Check if the current cell is masked, skip if it is
+    mask_data = if rule.mask_flag === NoMask() nothing else DynamicGrids.mask(data) end
+    if !isnothing(mask_data) && !mask_data[I...]
+        return nothing
+    end
+
     sum = zero(N)
     for (offset, k) in zip(offsets(rule), kernel(rule))
-        @inbounds propagules = N * k
-        @inbounds add!(data[W], propagules, I .+ offset...)
-        sum += propagules
+        target = I .+ offset
+        (target_mod, inbounds) = DynamicGrids.inbounds(data, target)
+        if inbounds && (isnothing(mask_data) || mask_data[target_mod...])
+            @inbounds propagules = N * k  
+            @inbounds add!(data[W], propagules, target_mod...)  
+            sum += propagules
+        end
     end
     @inbounds sub!(data[W], sum, I...)
     return nothing

test/integration.jl

@@ -169,3 +169,47 @@ end
     end
 
 end
+
+@testset "OutwardsDispersal Test" begin
+    # Define a mask
+    mask_data = [i == 1 || i == 10 || j == 1 || j == 10 ? false : true for i in 1:10, j in 1:10]
+
+    # Create OutwardsDispersal with a mask
+    outdisp_with_mask = OutwardsDispersal(
+        formulation=ExponentialKernel(λ=0.0125),
+        distancemethod=AreaToArea(30),
+        mask_flag=Mask()
+    )
+
+    # Create OutwardsDispersal without a mask, NoMask is default
+    outdisp_without_mask = OutwardsDispersal(
+        formulation=ExponentialKernel(λ=0.0125),
+        distancemethod=AreaToArea(30)
+    )
+
+    # Create a grid with empty borders matching the mask
+    init = map(x -> x ? 100.0 : 0.0, mask_data)
+
+    # Create ruleset and outputs
+    rule_with_mask = Ruleset(outdisp_with_mask; boundary=Reflect())
+    rule_without_mask = Ruleset(outdisp_without_mask; boundary=Reflect())
+
+    # Run the simulation with a mask
+    output_with_mask = ArrayOutput(init; tspan=1:1000, mask=mask_data)
+    a = sim!(output_with_mask, rule_with_mask)
+
+    @test sum(a[1]) ≈ sum(a[1000]) # Floating error should be smaller than 1.0
+
+    # Run the simulation without a mask to check default works fine
+    output_without_mask = ArrayOutput(init; tspan=1:1000)
+    b = sim!(output_without_mask, rule_without_mask)
+
+    @test sum(b[1]) ≈ sum(b[1000]) # Floating error should be smaller than 1.0
+
+    # Run the simulation with a mask but outdisp_without_mask
+    output_without_mask = ArrayOutput(init; tspan=1:1000, mask=mask_data)
+    b = sim!(output_without_mask, rule_without_mask)
+
+    @test sum(b[1]) > sum(b[1000]) #= Floating error should be larger than 1.0
+    because this does not identify the mask properly =#
+end