Modified particles to conserve tracers if pulling below zero

Now instead of allowing tracers to go below zero the model requires
fallback properties of the particles to be defined where the deficite
is taken from. For example, in the SLatissima model this is N (gN/dm^2)
so we also require a "scale factor" which in this case is the A property
and the unit conversion from mmolN to gN.

Also updated tests.

src/Models/Macroalgae/SLatissima.jl

@@ -111,11 +111,11 @@ end
 source_fields = ((tracer=:NO₃, property=:NO₃, scalefactor=1.0), 
                         (tracer=:NH₄, property=:NH₄, scalefactor=1.0),
                         (tracer=:PAR, property=:PAR, scalefactor=1.0))
-sink_fields = ((tracer=:NO₃, property=:j_NO₃, scalefactor=-1.0), 
-                    (tracer=:NH₄, property=:j_NH₄, scalefactor=-1.0), 
-                    (tracer=:DIC, property=:p, scalefactor=-1.0),
-                    (tracer=:DOM, property=:e, scalefactor=1.0/6.56),#Rd_dom from LOBSTER
-                    (tracer=:DD, property=:ν, scalefactor=1.0))
+sink_fields = ((tracer=:NO₃, property=:j_NO₃, scalefactor=-1.0, fallback=:N, fallback_scalefactor=(property=:A, constant=14*0.001/defaults.K_A)), 
+                    (tracer=:NH₄, property=:j_NH₄, scalefactor=-1.0, fallback=:N, fallback_scalefactor=(property=:A, constant=14*0.001/defaults.K_A)), 
+                    (tracer=:DIC, property=:p, scalefactor=-1.0, fallback=:C, fallback_scalefactor=(property=:A, constant=14*0.001)),
+                    (tracer=:DOM, property=:e, scalefactor=1.0/6.56, fallback=:A, fallback_scalefactor=0),#Rd_dom from LOBSTER, placeholder fallbacks because if these are taking away something has gone very wrong
+                    (tracer=:DD, property=:ν, scalefactor=1.0, fallback=:A, fallback_scalefactor=0))
 
 function defineparticles(initials, n)
     x̄₀ = []

src/Particles/Particles.jl

@@ -1,38 +1,59 @@
 module Particles
 using KernelAbstractions, Oceananigans, StructArrays
-using Oceananigans.Architectures: device
+using Oceananigans.Architectures: device, arch_array
+using Oceananigans.Operators: Vᶜᶜᶜ
 
-@kernel function source_sink!(model, loc, source_properties, scalefactors, tracers, Δt)
-    #this implimentation over doing it for single tracer at a time (to better match Oceananigans kernal function format) as identifying the particle mask takes time
-    #don't like all the looping in this fucntion
+function getmask(i, j, k, fi, fj, fk, grid)
+    nodesi=arch_array(grid.architecture, zeros(Int, 2, 2, 2))
+    nodesj=arch_array(grid.architecture, zeros(Int, 2, 2, 2))
+    nodesk=arch_array(grid.architecture, zeros(Int, 2, 2, 2))
+
+    d=arch_array(grid.architecture, zeros(2,2,2))
+
+    for (a, di) in enumerate([fi, 1-fi]), (b, dj) in enumerate([fj, 1-fj]), (c, dk) in enumerate([fk, 1-fk])
+        d[a,b,c] = sqrt(di^2+dj^2+dk^2)
+        nodesi[a, b, c] = min(max(i+a, 1), grid.Nx)
+        nodesj[a, b, c] = min(max(j+b, 1), grid.Ny)
+        nodesk[a, b, c] = min(max(k+c, 1), grid.Nz)
+    end
+
+    normfactor=1/sum(1 ./ d)
+
+    d ./= normfactor
+    return nodesi, nodesj, nodesk, d
+end
+
+function apply_sinks!(model, particles, p, i, j, k, d, Δt)
+    for sink in particles.parameters.sink_fields
+        value =particles.parameters.density * sink.scalefactor * Δt * getproperty(model.particles.properties, sink.property)[p] / (d * Vᶜᶜᶜ(i, j, k, model.grid))
+        if abs(value) > 0
+            res = value + getproperty(model.tracers, sink.tracer)[i, j, k]
+            if res < 0
+                @warn "Particle tried to take tracer below zero, conserving tracer (may cause adverse effects on particle dynamics)"
+                getproperty(model.tracers, sink.tracer)[i, j, k] = 0
+                getproperty(particles.properties, sink.fallback)[p] += fallback(particles, p, res, sink.fallback_scalefactor)*Vᶜᶜᶜ(i, j, k, model.grid)
+            else
+                getproperty(model.tracers, sink.tracer)[i, j, k] += value
+            end
+        end
+    end
+end
+
+fallback(particles, p, diff, fallback_scalefactor::Number) = diff*fallback_scalefactor
+
+fallback(particles, p, diff, fallback_scalefactor::NamedTuple{(:property, :constant), Tuple{Symbol, Number}}) = diff*fallback_scalefactor.constant/(getproperty(particles.properties, fallback_scalefactor.property)[p])
+
+@kernel function source_sink!(model, loc, Δt)
     p = @index(Global)
     (fi, i::Int), (fj, j::Int), (fk, k::Int) = modf.(Oceananigans.Fields.fractional_indices(model.particles.properties.x[p], model.particles.properties.y[p], model.particles.properties.z[p], loc, model.grid))
-   #possibly not the most efficient implimentation but benchmarked 10x faster than a weird vector solution        
-    if (fi, fj, fk) != (0, 0, 0)
-        d=zeros(2,2,2)
-        vols=zeros(2,2,2)
-
-        nodesi=zeros(Int, 2, 2, 2)
-        nodesj=zeros(Int, 2, 2, 2)
-        nodesk=zeros(Int, 2, 2, 2)
-        for (a, di) in enumerate([fi, 1-fi]), (b, dj) in enumerate([fj, 1-fj]), (c, dk) in enumerate([fk, 1-fk])
-            d[a,b,c] = sqrt(di^2+dj^2+dk^2)
-            nodesi[a, b, c] = min(max(i+a, 1), model.grid.Nx)
-            nodesj[a, b, c] = min(max(j+b, 1), model.grid.Ny)
-            nodesk[a, b, c] = min(max(k+c, 1), model.grid.Nz)
-        end
 
-        normfactor=1/sum(1 ./ d)
-        #have todo second loop to have the normfactor (I think)
+    if (fi, fj, fk) != (0, 0, 0)
+        nodesi, nodesj, nodesk, d = getmask(i, j, k, fi, fj, fk, model.grid)
         for (ind, i) in enumerate(nodesi)
-            for (l, tracer) in enumerate(tracers)
-                getproperty(model.tracers, tracer)[i, nodesj[ind], nodesk[ind]] += scalefactors[l] * Δt * normfactor * getproperty(model.particles.properties, source_properties[l])[p] / (d[ind] * Oceananigans.Operators.Vᶜᶜᶜ(i, nodesj[ind], nodesk[ind], model.grid))
-            end
+            apply_sinks!(model, model.particles, p, i, nodesj[ind], nodesk[ind], d[ind], Δt)
         end
     else
-        for (l, tracer) in enumerate(tracers)
-            getproperty(model.tracers, tracer)[i+1, j+1, k+1] += scalefactors[l] * Δt * getproperty(model.particles.properties, source_properties[l])[p] / (Oceananigans.Operators.Vᶜᶜᶜ(i, j, k, model.grid))
-        end
+        apply_sinks!(model, model.particles, p, i+1, j+1, k+1, 1, Δt)
     end
 end
 
@@ -89,12 +110,9 @@ function dynamics!(particles, model, Δt)
     if length(particles.parameters.sink_fields)>0#probably not the most julian way todo this and would also be fixed if the next line was more generic
         LX, LY, LZ = location(getproperty(model.tracers, particles.parameters.sink_fields[1].tracer))#this will not be right if they are different between the differnet fields
 
-        source_properties = getproperty.(particles.parameters.sink_fields, :property)
-        source_tracers = getproperty.(particles.parameters.sink_fields, :tracer)
-        source_scalefactors = getproperty.(particles.parameters.sink_fields, :scalefactor) .* model.particles.parameters.density
-
         sourcesink_kernal! = source_sink!(device(model.architecture), workgroup, worksize)
-        sourcesink_event = sourcesink_kernal!(model, (LX(),LY(),LZ()), source_properties, source_scalefactors, source_tracers, Δt)
+        sourcesink_event = sourcesink_kernal!(model, (LX(),LY(),LZ()), Δt)
+
         wait(sourcesink_event)
     end
     particles.parameters.custom_dynamics(particles, model, Δt)
@@ -107,7 +125,7 @@ function setup(particles::StructArray, equation::Function, equation_arguments::N
     equation_parameters::NamedTuple, integrals::NTuple{N, Symbol} where N, 
     tracked::NTuple{N, Symbol} where N, 
     tracked_fields::NTuple{N, NamedTuple{(:tracer, :property, :scalefactor), Tuple{Symbol, Symbol, Float64}}} where N, 
-    sink_fields::NTuple{N, NamedTuple{(:tracer, :property, :scalefactor), Tuple{Symbol, Symbol, Float64}}} where N, 
+    sink_fields, #::NTuple{N, NamedTuple{(:tracer, :property, :scalefactor, :fallback, :fallback_scalefactor), Tuple{Symbol, Symbol, Float64, Symbol, NamedTuple}}} where N - got too complicated
     density::AbstractFloat, custom_dynamics=no_dynamics)
 
     for property in (equation_arguments..., integrals..., tracked...)

test/runtests.jl

@@ -1 +1,2 @@
-include("test_particle_updating.jl")
+include("test_particle_updating.jl")
+include("test_light.jl")

test/test_particle_updating.jl

@@ -60,6 +60,7 @@ end
     function particleupdate(x, y, z, t, A, B, params, Δt)
         return (A=0.1, B=t)
     end
+    particlestruct=StructArray{CustomParticle}(([0.25], [0.25], [-0.25], [1.0], [0.0], [0.0]))
     particles = Particles.setup(
         particlestruct, 
         particleupdate, 
@@ -68,7 +69,7 @@ end
         (), 
         (:A, :B), 
         ((tracer=:C, property=:C, scalefactor=1.0), ),
-        ((tracer=:C, property=:A, scalefactor=-1.0), ),
+        ((tracer=:C, property=:A, scalefactor=-1.0, fallback=:A, fallback_scalefactor=0), ),#placeholder fallback scale factor as testing further down), ),
         1.0
     )
 
@@ -79,6 +80,32 @@ end
 
     @test model.tracers.C[1, 1, 1] ≈ 0.9 atol = 0.05
 
+    #pull tracer below zero
+    function particleupdate(x, y, z, t, A, params, Δt)
+        return (A=2.0, )
+    end
+
+    particlestruct=StructArray{CustomParticle}(([0.25], [0.25], [-0.25], [1.0], [0.0], [0.0]))
+    particles = Particles.setup(
+        particlestruct, 
+        particleupdate, 
+        (:A, ), 
+        NamedTuple(), 
+        (), 
+        (:A, ), 
+        ((tracer=:C, property=:C, scalefactor=1.0), ),
+        ((tracer=:C, property=:A, scalefactor=-1.0, fallback=:B, fallback_scalefactor=1.0), ),
+        1.0
+    )
+
+    model = NonhydrostaticModel(; grid, timestepper=:RungeKutta3, particles=particles, tracers=(:C, ))
+    set!(model, u=0, v=0, w=0, C=1)
+    sim = Simulation(model, Δt=1, stop_time=1)
+    run!(sim)
+
+    @test model.tracers.C[1, 1, 1] ≈ 0.0
+    @test model.particles.properties.B[1] ≈ -1.0
+
     @testset "Larger grid for point assignment" begin
         grid = RectilinearGrid(size=(2,1,1), extent=(2,1,1), topology=(Periodic, Periodic, Periodic))
         particlestruct=StructArray{CustomParticle}(([1], [0.25], [-0.25], [1.0], [0.0], [0.0]))
@@ -91,7 +118,7 @@ end
             (), 
             (:A, :B), 
             ((tracer=:C, property=:C, scalefactor=1.0), ),
-            ((tracer=:C, property=:A, scalefactor=-1.0), ),
+            ((tracer=:C, property=:A, scalefactor=-1.0, fallback=:A, fallback_scalefactor=0), ),#placeholder fallback scale factor as testing further down), ),
             1.0
         )
 
@@ -116,7 +143,7 @@ end
             (), 
             (:A, :B), 
             ((tracer=:C, property=:C, scalefactor=1.0), ),
-            ((tracer=:C, property=:A, scalefactor=-1.0), ),
+            ((tracer=:C, property=:A, scalefactor=-1.0, fallback=:A, fallback_scalefactor=0), ),#placeholder fallback scale factor as testing further down), ),
             1.0
         )
 
@@ -137,7 +164,7 @@ end
             (), 
             (:A, :B), 
             ((tracer=:C, property=:C, scalefactor=1.0), ),
-            ((tracer=:C, property=:A, scalefactor=-1.0), ),
+            ((tracer=:C, property=:A, scalefactor=-1.0, fallback=:A, fallback_scalefactor=0), ),#placeholder fallback scale factor as testing further down), ),
             1.0
         )
 
@@ -159,7 +186,7 @@ end
             (), 
             (:A, :B), 
             ((tracer=:C, property=:C, scalefactor=1.0), ),
-            ((tracer=:C, property=:A, scalefactor=-1.0), ),
+            ((tracer=:C, property=:A, scalefactor=-1.0, fallback=:A, fallback_scalefactor=0), ),#placeholder fallback scale factor as testing further down), ),
             10.0
         )