Changes sediment timestepping to correctly use the models timestepper + bug fix

Project.toml

@@ -1,6 +1,6 @@
 name = "OceanBioME"
 uuid = "a49af516-9db8-4be4-be45-1dad61c5a376"
-version = "0.18.7"
+version = "0.18.8"
 authors = ["Jago Strong-Wright <js2430@damtp.cam.ac.uk> and contributors"]
 
 [deps]

src/Models/AdvectedPopulations/NutrientsPlanktonDetritus/Plankton/Plankton.jl

@@ -37,9 +37,10 @@ using ..NutrientsPlanktonDetritusModels.DetritusModels:
 import Adapt: adapt_structure
 import Base: show, summary
 
-import Oceananigans.Biogeochemistry: 
+import Oceananigans.Biogeochemistry:
     required_biogeochemical_tracers,
-    required_biogeochemical_auxiliary_fields
+    required_biogeochemical_auxiliary_fields,
+    biogeochemical_drift_velocity
 
 import OceanBioME: chlorophyll
 

src/Sediments/Sediments.jl

@@ -45,7 +45,9 @@ include("show.jl")
 
 function BiogeochemicalSediment(grid, biogeochemistry;
                                 clock = Clock(time = zero(grid)),
-                                timestepper = :QuasiAdamsBashforth2)
+                                timestepper = nothing)
+
+    deprecate_sediment_timestepper(timestepper)
 
     bottom_indices = calculate_bottom_indices(grid)
 
@@ -55,7 +57,7 @@ function BiogeochemicalSediment(grid, biogeochemistry;
     prognostic_fields = NamedTuple{sediment_field_names}(map(n -> Field{Center, Center, Nothing}(grid), 1:length(sediment_field_names)))
     tracked_fields = NamedTuple{tracked_field_names}(map(n -> Field{Center, Center, Nothing}(grid), 1:length(tracked_field_names)))
 
-    timestepper = TimeStepper(timestepper, grid, prognostic_fields)
+    timestepper = TimeStepper(:QuasiAdamsBashforth2, grid, prognostic_fields)
 
     return BiogeochemicalSediment(architecture(grid), biogeochemistry, timestepper, clock, grid, prognostic_fields, tracked_fields, bottom_indices)
 end

src/Sediments/timesteppers.jl

@@ -1,22 +1,68 @@
 using Oceananigans.Architectures: architecture
-using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper, RungeKutta3TimeStepper
+using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper, RungeKutta3TimeStepper,
+                                 time_step!, tick!
 using Oceananigans.Utils: launch!
 
 import Oceananigans.TimeSteppers: ab2_step!, rk3_substep!,
                                   cache_previous_tendencies!, compute_flux_bc_tendencies!
 
-const VALID_TIMESTEPPERS = Union{<:QuasiAdamsBashforth2TimeStepper, <:RungeKutta3TimeStepper}
+const SUPPORTED_TIMESTEPPERS = Union{QuasiAdamsBashforth2TimeStepper, RungeKutta3TimeStepper}
 
-validate_sediment_timestepper(timestepper) = throw(ArgumentError("$(typeof(timestepper)) is not configured for sediment models"))
-validate_sediment_timestepper(::VALID_TIMESTEPPERS) = nothing
+function step_sediment!(sediment, model, parent_timestepper::SUPPORTED_TIMESTEPPERS)
+    substep_sediment!(sediment, model, parent_timestepper)
+
+    cache_previous_tendencies!(sediment)
+
+    tick!(sediment.clock, model.clock.last_stage_Δt)
+
+    return nothing
+end
+
+function step_sediment!(sediment, model, parent_timestepper)
+    @warn "$(nameof(typeof(parent_timestepper))) is not supported for sediment models, " *
+          "so the sediment is being stepped independently of $(nameof(typeof(model))), " *
+          "which is neither mass conserving nor higher than first order accurate" maxlog = 1
+
+    time_step!(sediment, model.clock.last_stage_Δt)
+
+    return nothing
+end
+
+function substep_sediment!(sediment, model, parent_timestepper::QuasiAdamsBashforth2TimeStepper)
+    ab2_step!(sediment, model.clock.last_Δt, parent_timestepper.χ, nothing)
+
+    return nothing
+end
+
+function substep_sediment!(sediment, model, parent_timestepper::RungeKutta3TimeStepper)
+    γ¹ = parent_timestepper.γ¹
+    γ² = parent_timestepper.γ²
+    γ³ = parent_timestepper.γ³
+    ζ² = parent_timestepper.ζ²
+    ζ³ = parent_timestepper.ζ³
+
+    stage_Δt = model.clock.last_stage_Δt
+
+    if model.clock.stage == 2
+        rk3_substep!(sediment, stage_Δt / γ¹, γ¹, nothing, nothing)
+    elseif model.clock.stage == 3
+        rk3_substep!(sediment, stage_Δt / (γ² + ζ²), γ², ζ², nothing)
+    else
+        rk3_substep!(sediment, model.clock.last_Δt, γ³, ζ³, nothing)
+    end
+
+    return nothing
+end
 
 # AB2 methods
 
-function ab2_step!(model::BiogeochemicalSediment, Δt, callbacks)
+ab2_step!(model::BiogeochemicalSediment, Δt, callbacks) =
+    ab2_step!(model, Δt, model.timestepper.χ, callbacks)
+
+function ab2_step!(model::BiogeochemicalSediment, Δt, χ, callbacks)
     grid = model.grid
     arch = architecture(grid)
     model_fields = prognostic_fields(model)
-    χ = model.timestepper.χ
 
     for (i, field) in enumerate(model_fields)
         kernel_args = (field, Δt, χ, model.timestepper.Gⁿ[i], model.timestepper.G⁻[i])
@@ -94,3 +140,10 @@ function cache_previous_tendencies!(model::BiogeochemicalSediment)
 end
 
 compute_flux_bc_tendencies!(model::BiogeochemicalSediment) = nothing
+
+deprecate_sediment_timestepper(::Nothing) = nothing
+
+deprecate_sediment_timestepper(timestepper) =
+    @warn "The `timestepper` keyword argument to sediment models is deprecated and ignored: " *
+          "the sediment is stepped with the time stepping coefficients of the model it is " *
+          "coupled to, so `timestepper = :$(timestepper)` has no effect" maxlog = 1

src/Sediments/update_state.jl

@@ -4,14 +4,14 @@ using Oceananigans.TimeSteppers: time_step!
 import Oceananigans.TimeSteppers: update_state!
 
 function update_biogeochemical_state!(model, sediment_model::BiogeochemicalSediment)
-    Δt = model.clock.last_stage_Δt
-
     update_tracked_fields!(sediment_model, model)
 
     if !((model.clock.iteration == 0) & (model.clock.stage == 1))
-        time_step!(sediment_model, Δt;)
+        step_sediment!(sediment_model, model, model.timestepper)
     end
 
+    update_state!(sediment_model)
+
     return nothing
 end
 

test/test_sediments.jl

@@ -4,31 +4,35 @@ using OceanBioME.Models.SedimentModels: InstantRemineralisation, SimpleMultiG
 using OceanBioME.Sediments: BiogeochemicalSediment
 
 display_name(::NutrientsPlanktonDetritus) = "NutrientsPlanktonDetritus"
-display_name(::BiogeochemicalSediment{<:SimpleMultiG}) = "Multi-G"
-display_name(::BiogeochemicalSediment{<:OceanBioME.Models.SedimentModels.InstantRemineralisation}) = "Instant remineralisation"
+display_name(::BiogeochemicalSediment{<:Any, <:SimpleMultiG}) = "Multi-G"
+display_name(::BiogeochemicalSediment{<:Any, <:InstantRemineralisation}) = "Instant remineralisation"
 display_name(::RectilinearGrid) = "Rectilinear grid"
 display_name(::LatitudeLongitudeGrid) = "Latitude longitude grid"
 display_name(::ImmersedBoundaryGrid) = "Immersed boundary grid"
 
-function display_name(architecture, grid, sediment_model, biogeochemistry, timestepper)
+function display_name(architecture, grid, sediment_model, biogeochemistry, model_name)
     arch_name = typeof(architecture)
     sediment_name = display_name(sediment_model)
     bgc_name = display_name(biogeochemistry.underlying_biogeochemistry)
     grid_name = display_name(grid)
 
-    @info "Testing sediment on $arch_name with $timestepper and $sediment_name on $bgc_name with $grid_name"
+    @info "Testing sediment on $arch_name with $model_name and $sediment_name on $bgc_name with $grid_name"
 
-    return "$architecture, $timestepper, $sediment_name, $bgc_name, $grid_name"
+    return "$architecture, $model_name, $sediment_name, $bgc_name, $grid_name"
 end
 
-set_sinkers!(::NutrientsPlanktonDetritus{<:Any, <:Any, <:Detritus}, model) = set!(model, D = 1)
-set_sinkers!(::NutrientsPlanktonDetritus{<:Any, <:Any, <:DissolvedParticulate{1, 2}}, model) = set!(model, sPOM = 1, bPOM = 1)
-#=set_sinkers!(::NutrientsPlanktonDetritus{<:Any, <:Any, <:DissolvedPar}, model) =
-    set!(model, sPON = 1, bPON = 1, sPOC = 6.56, bPOC = 6.56)=#
+set_sinkers!(::NutrientsPlanktonDetritus{<:Any, <:Any, <:Any, <:Detritus}, model) = set!(model, D = 1)
+set_sinkers!(::NutrientsPlanktonDetritus{<:Any, <:Any, <:Any, <:DissolvedParticulate{1, 2}}, model) = set!(model, sPOM = 1, bPOM = 1)
+set_sinkers!(::NutrientsPlanktonDetritus{<:Any, <:Any, <:Any, <:CarbonNitrogenDissolvedParticulate}, model) =
+    set!(model, sPON = 1, bPON = 1, sPOC = 6.56, bPOC = 6.56)
 
 sum_of_volume_integrals(biogeochemistry, tracers) = sum(map(f -> Field(Integral(f)), values(tracers)))
-#=sum_of_volume_integrals(::NutrientsPlanktonDetritus{<:Any, <:Any, <:VariableRedfieldDetritus}, tracers) =
-    sum([Field(Integral(f)) for (n, f) in pairs(tracers) if n in (:NO₃, :NH₄, :P, :Z, :sPON, :bPON, :DON)])=#
+sum_of_volume_integrals(::NutrientsPlanktonDetritus{<:Any, <:Any, <:Any, <:CarbonNitrogenDissolvedParticulate}, tracers) =
+    sum([Field(Integral(f)) for (n, f) in pairs(tracers) if n in (:NO₃, :NH₄, :P, :Z, :sPON, :bPON, :DON)])
+# `O₂` carries no nitrogen and isn't conserved against the sediment inventory, so it has to be
+# excluded from what this test calls "total nitrogen"
+sum_of_volume_integrals(::NutrientsPlanktonDetritus{<:Any, <:Any, <:Any, <:DissolvedParticulate, <:Any, <:Oxygen}, tracers) =
+    sum([Field(Integral(f)) for (n, f) in pairs(tracers) if n != :O₂])
 
 sum_of_area_integrals(sediment, fields) = sum(map(f -> Field(Integral(f, dims = (1, 2))), values(fields)))
 sum_of_area_integrals(::SimpleMultiG{Nothing}, fields) =
@@ -71,10 +75,9 @@ function test_sediment(grid, biogeochemistry, model_name, advection = WENO(order
 
     final_total_nitrogen = CUDA.@allowscalar total_nitrogen[1, 1, 1]
 
-    # simple multi-G is only good to this precision, IR is fine to default
-    @test isapprox(initial_total_nitrogen, final_total_nitrogen,rtol = 0.2e-6)
+    @test isapprox(initial_total_nitrogen, final_total_nitrogen, rtol = 1e-8)#0.2e-6)
 
-    @test all(interior(sediment_nitrogen) .!= 0)
+    @test CUDA.@allowscalar all(interior(sediment_nitrogen) .!= 0)
 
     return model
 end
@@ -101,42 +104,40 @@ immersed_latlon_grid = ImmersedBoundaryGrid(
 )
 
 grids = (rectilinear_grid, latlon_grid, immersed_latlon_grid)
-sediment_timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3)
-models = (NonhydrostaticModel, HydrostaticFreeSurfaceModel) # I don't think we need to test on both models anymore
-#=
+models = (NonhydrostaticModel, HydrostaticFreeSurfaceModel) # exercises both `substep_sediment!` methods (RK3 and AB2)
+
 @testset "Sediment integration" begin
-    for grid in grids, timestepper in sediment_timesteppers
-        npzd_ir = NutrientPhytoplanktonZooplanktonDetritus(;
-            grid,
-            sediment_model = InstantRemineralisationSediment(grid; timestepper)
+    for grid in grids
+        npzd_ir = NPZD(
+            grid;
+            sediment = InstantRemineralisationSediment(grid)
         )
 
-        lobster_ir = LOBSTER(;
-            grid,
-            sediment_model = InstantRemineralisationSediment(
+        lobster_ir = LOBSTER(
+            grid;
+            sediment = InstantRemineralisationSediment(
                 grid;
                 sinking_tracers = (:sPOM, :bPOM),
-                remineralisation_reciever = :NH₄,
-                timestepper
+                remineralisation_reciever = :NH₄
             )
         )
 
-        simple_lobster_multi_g = LOBSTER(;
-            grid,
+        simple_lobster_multi_g = LOBSTER(
+            grid;
             sediment = SimpleMultiGSediment(grid),
-            oxygen = true
+            oxygen = Oxygen()
         )
 
-        full_lobster_multi_g = LOBSTER(;
-            grid,
+        full_lobster_multi_g = LOBSTER(
+            grid;
+            detritus = CarbonNitrogenDissolvedParticulate(grid; open_bottom = true),
             sediment = SimpleMultiGSediment(
                 grid;
                 sinking_nitrogen = (:sPON, :bPON),
                 sinking_carbon = (:sPOC, :bPOC)
             ),
-            oxygen = true,
-            carbonates = true,
-            variable_redfield = true
+            oxygen = Oxygen(),
+            inorganic_carbon = CarbonateSystem()
         )
 
         bgcs = [npzd_ir, lobster_ir, simple_lobster_multi_g, full_lobster_multi_g]
@@ -152,12 +153,15 @@ models = (NonhydrostaticModel, HydrostaticFreeSurfaceModel) # I don't think we n
                 continue
             end
 
-            test_name = display_name(architecture, grid, biogeochemistry.sediment, biogeochemistry, timestepper)
+            test_name = display_name(architecture, grid, biogeochemistry.sediment, biogeochemistry, model)
 
             @testset "$(test_name)" begin
-                test_sediment(grid, biogeochemistry, model)
+                # `InstantRemineralisation` defines its tendency method for the chosen
+                # `remineralisation_reciever` via `eval` at construction time, so a top-level
+                # loop that both builds the biogeochemistry and calls `test_sediment` in the
+                # same compiled thunk needs `invokelatest` to see it
+                Base.invokelatest(test_sediment, grid, biogeochemistry, model)
             end
         end
     end
 end
-=#