makes box models work with non-trivial coordinates

OceanBioME · Sep 10, 2024 · 725cc5b · 725cc5b · jagoosw · Sep 10, 2024
1 parent 7f14d19
commit 725cc5b
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 11 deletions.
diff --git a/src/BoxModel/timesteppers.jl b/src/BoxModel/timesteppers.jl
@@ -1,5 +1,6 @@
 using Oceananigans.Architectures: device
 using Oceananigans.Biogeochemistry: update_tendencies!, biogeochemical_auxiliary_fields
+using Oceananigans.Grids: nodes, Center
 using Oceananigans.TimeSteppers: rk3_substep_field!, store_field_tendencies!, RungeKutta3TimeStepper, QuasiAdamsBashforth2TimeStepper
 using Oceananigans.Utils: work_layout, launch!
 
@@ -39,7 +40,7 @@ function compute_tendencies!(model::BoxModel, callbacks)
     for tracer in required_biogeochemical_tracers(model.biogeochemistry)
         forcing = @inbounds model.forcing[tracer]
 
-        @inbounds Gⁿ[tracer][1, 1, 1] = tracer_tendency(Val(tracer), model.biogeochemistry, forcing, model.clock.time, model.field_values)
+        @inbounds Gⁿ[tracer][1, 1, 1] = tracer_tendency(Val(tracer), model.biogeochemistry, forcing, model.clock.time, model.field_values, model.grid)
     end
 
     for callback in callbacks
@@ -51,10 +52,14 @@ function compute_tendencies!(model::BoxModel, callbacks)
     return nothing
 end
 
-@inline tracer_tendency(val_name, biogeochemistry, forcing, time, model_fields) =
-    biogeochemistry(val_name, 0, 0, 0, time, model_fields...) + forcing(time, model_fields...)
+@inline boxmodel_xyz(nodes, grid) = map(n->boxmodel_coordinate(n, grid), nodes)
+@inline boxmodel_coordinate(::Nothing, grid) = zero(grid)
+@inline boxmodel_coordinate(nodes, grid) = @inbounds nodes[1]
 
-@inline tracer_tendency(::Val{:T}, biogeochemistry, forcing, time, model_fields) = 0
+@inline tracer_tendency(val_name, biogeochemistry, forcing, time, model_fields, grid) =
+    biogeochemistry(val_name, boxmodel_xyz(nodes(grid, Center(), Center(), Center()), grid)..., time, model_fields...) + forcing(time, model_fields...)
+
+@inline tracer_tendency(::Val{:T}, biogeochemistry, forcing, time, model_fields, grid) = 0
 
 function rk3_substep!(model::BoxModel, Δt, γⁿ, ζⁿ)
     model_fields = prognostic_fields(model)

diff --git a/test/test_boxmodel.jl b/test/test_boxmodel.jl
@@ -1,14 +1,18 @@
 #include("dependencies_for_runtests.jl")
 
+using OceanBioME.BoxModels: boxmodel_xyz
+
 using Oceananigans: AbstractGrid, AbstractModel, AbstractOutputWriter
 using Oceananigans.Fields: FunctionField
+using Oceananigans.Grids: nodes
 
 PAR(t) = (60 * (1 - cos((t + 15days) * 2π / 365days)) * (1 / (1 + 0.2 * exp(-((mod(t, 365days) - 200days) / 50days)^2))) + 2) * exp(-2)
 
 defaults = (NO₃ = 10.0, NH₄ = 0.1, P = 0.1, Z = 0.01)
 
-function simple_box_model()
-    grid = BoxModelGrid()
+function simple_box_model(; x = nothing, y = nothing, z = nothing)
+    grid = BoxModelGrid(; x, y, z)
+
     clock = Clock(time = zero(grid))
 
     light_attenuation_model = PrescribedPhotosyntheticallyActiveRadiation(FunctionField{Center, Center, Center}(PAR, grid; clock))
@@ -21,11 +25,14 @@ function simple_box_model()
 end
 
 @testset "Construct `BoxModel`" begin
-    model = simple_box_model()
+    for z in (nothing, -100)
+        model = simple_box_model(; z)
 
-    @test grid isa AbstractGrid
-    @test model isa AbstractModel
-    @test set!(model; defaults...) isa Nothing
+        @test model.grid isa AbstractGrid
+        @test model isa AbstractModel
+        @test set!(model; defaults...) isa Nothing
+        @test boxmodel_xyz(nodes(model.grid, Center(), Center(), Center()), model.grid) == (0, 0, ifelse(isnothing(z), 0, z))
+    end
 end
 
 @testset "Timestep `BoxModel`" begin
@@ -38,8 +45,9 @@ end
     end
 
     # values are being updated
-    for (name, field) in enumerate(model.fields)
+    for (name, field) in pairs(model.fields)
         default = name in keys(defaults) ? defaults[name] : 0
+
         @test field[1, 1, 1] != default
     end
 end