Modified particles to conserve tracers if sinking below zero, and fixed light

src/Light/2band.jl

@@ -5,9 +5,9 @@
 
     z = grid.zᵃᵃᶜ[grid.Nz]
 
-    ∫chlᵉʳ = (P[i, j, grid.Nz]*params.Rd_chl/params.r_pig)^params.e_r*z
-    ∫chlᵉᵇ = (P[i, j, grid.Nz]*params.Rd_chl/params.r_pig)^params.e_b*z
-    PAR[i, j, grid.Nz] =  sp*(exp(params.k_r0 * z + params.Χ_rp * ∫chlᵉʳ).+ exp(params.k_b0 * z + params.Χ_bp * ∫chlᵉᵇ))/2
+    ∫chlᵉʳ = -(P[i, j, grid.Nz]*params.Rd_chl/params.r_pig)^params.e_r*z
+    ∫chlᵉᵇ = -(P[i, j, grid.Nz]*params.Rd_chl/params.r_pig)^params.e_b*z
+    PAR[i, j, grid.Nz] =  sp*(exp(params.k_r0 * z - params.Χ_rp * ∫chlᵉʳ).+ exp(params.k_b0 * z - params.Χ_bp * ∫chlᵉᵇ))/2
     for k=grid.Nz-1:-1:1
         z = grid.zᵃᵃᶜ[k]
         dz = grid.zᵃᵃᶜ[k+1] - z 
@@ -16,7 +16,7 @@
         ∫chlᵉʳ += (mean_pig^params.e_r)*dz
         ∫chlᵉᵇ += (mean_pig^params.e_b)*dz
 
-        PAR[i, j, k] =  sp*(exp(params.k_r0 * z + params.Χ_rp * ∫chlᵉʳ) + exp(params.k_b0 * z + params.Χ_bp * ∫chlᵉᵇ))/2
+        PAR[i, j, k] =  sp*(exp(params.k_r0 * z - params.Χ_rp * ∫chlᵉʳ) + exp(params.k_b0 * z - params.Χ_bp * ∫chlᵉᵇ))/2
     end
 end 
 

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, Float64}}) = 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_light.jl

@@ -2,8 +2,8 @@ using OceanBioME, Test, Oceananigans
 
 grid = RectilinearGrid(size=(1,1,2), extent=(1,1,2))
 
-params=LOBSTER.default
-
+params=LOBSTER.defaults
+#=
 @testset "Light attenuation" begin
     PAR = Oceananigans.Fields.Field{Center, Center, Center}(grid)
 
@@ -22,5 +22,23 @@ params=LOBSTER.default
 
     results_PAR=convert(Array, model.auxiliary_fields.PAR)[1, 1, :]
 
-    @test all(results_PAR.==expected_PAR)
-end
+    @test all(results_PAR .≈ expected_PAR)
+end=#
+PAR = Oceananigans.Fields.Field{Center, Center, Center}(grid)
+
+    model = NonhydrostaticModel(; grid, timestepper=:RungeKutta3, tracers=(:P, ), auxiliary_fields = (PAR=PAR, ))
+    Pᵢ(x,y,z)=2.5+z
+    set!(model, u=0, v=0, w=0, P=Pᵢ)
+    sim = Simulation(model, Δt=0.1, stop_time=1)
+    surface_PAR(t) = 1
+    sim.callbacks[:update_par] = Callback(Light.update_2λ!, IterationInterval(1), merge(params, (surface_PAR=surface_PAR,)));
+    run!(sim)
+
+    expected_PAR = [
+        (exp(-params.k_r0*0.5-0.5*params.Χ_rp*(2*params.Rd_chl/params.r_pig)^params.e_r)*exp(-params.k_r0*1-1*params.Χ_rp*(1.5*params.Rd_chl/params.r_pig)^params.e_r)+exp(-params.k_b0*0.5-0.5*params.Χ_bp*(2*params.Rd_chl/params.r_pig)^params.e_b)*exp(-params.k_b0*1-1*params.Χ_bp*(1.5*params.Rd_chl/params.r_pig)^params.e_b))/2,
+        (exp(-params.k_r0*0.5-0.5*params.Χ_rp*(2*params.Rd_chl/params.r_pig)^params.e_r)+exp(-params.k_b0*0.5-0.5*params.Χ_bp*(2*params.Rd_chl/params.r_pig)^params.e_b))/2
+    ]
+
+    results_PAR=convert(Array, model.auxiliary_fields.PAR)[1, 1, :]
+
+    @test all(results_PAR .≈ expected_PAR)

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,53 @@ 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
+
+    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=(property=:A, constant=0.5)), ), #other property dependant scale factor
+        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] ≈ -0.25
+
     @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 +139,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 +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
         )
 
@@ -137,7 +185,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 +207,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
         )