Adds split timestepping for physics and BGC

src/Biogeochemistry.jl

@@ -59,10 +59,12 @@ is a subtype of `AbstractBioeochemistry`:
 abstract type AbstractBiogeochemistry end
 
 # Returns the forcing for discrete form models
-@inline biogeochemical_transition(i, j, k, grid, bgc, val_tracer_name, clock, fields) =
+@inline biogeochemical_transition(i, j, k, grid, bgc, val_tracer_name, clock, fields, val_compute_bgc) =
     bgc(i, j, k, grid, val_tracer_name, clock, fields)
 
-@inline biogeochemical_transition(i, j, k, grid, ::Nothing, val_tracer_name, clock, fields) = zero(grid)
+@inline biogeochemical_transition(i, j, k, grid, ::Nothing, val_tracer_name, clock, fields, val_compute_bgc) = zero(grid)
+
+@inline biogeochemical_transition(i, j, k, grid, bgc, val_tracer_name, clock, fields, ::Val{false}) = zero(grid)
 
 # Required for when a model is defined but not for all tracers
 @inline (bgc::AbstractBiogeochemistry)(i, j, k, grid, val_tracer_name, clock, fields) = zero(grid)
@@ -106,7 +108,7 @@ abstract type AbstractContinuousFormBiogeochemistry <: AbstractBiogeochemistry e
 
 """Return the biogeochemical forcing for `val_tracer_name` for continuous form when model is called."""
 @inline function biogeochemical_transition(i, j, k, grid, bgc::AbstractContinuousFormBiogeochemistry,
-                                           val_tracer_name, clock, fields)
+                                           val_tracer_name, clock, fields, val_compute_bgc)
 
     names_to_extract = tuple(required_biogeochemical_tracers(bgc)...,
                              required_biogeochemical_auxiliary_fields(bgc)...)
@@ -120,6 +122,8 @@ abstract type AbstractContinuousFormBiogeochemistry <: AbstractBiogeochemistry e
     return bgc(val_tracer_name, x, y, z, clock.time, fields_ijk...)
 end
 
+@inline biogeochemical_transition(i, j, k, grid, bgc::AbstractContinuousFormBiogeochemistry, val_tracer_name, clock, fields, ::Val{false}) = zero(grid)
+
 @inline (bgc::AbstractContinuousFormBiogeochemistry)(val_tracer_name, x, y, z, t, fields...) = zero(t)
 
 tracernames(tracers) = keys(tracers)
@@ -171,4 +175,70 @@ const AbstractBGCOrNothing = Union{Nothing, AbstractBiogeochemistry}
 required_biogeochemical_tracers(::AbstractBGCOrNothing) = ()
 required_biogeochemical_auxiliary_fields(::AbstractBGCOrNothing) = ()
 
+#####
+##### BGC only stepping
+#####
+
+using Oceananigans: interior_tendency_kernel_parameters
+using Oceananigans.Utils: launch!
+
+using KernelAbstractions: @kernel, @index
+
+import Oceananigans: compute_biogeochemical_tendencies!
+
+function compute_biogeochemical_tendencies!(model, tendencies; active_cells_map = nothing)
+
+    arch                 = model.architecture
+    grid                 = model.grid
+    biogeochemistry      = model.biogeochemistry
+    velocities           = model.velocities
+    tracers              = model.tracers
+    auxiliary_fields     = model.auxiliary_fields
+    clock                = model.clock
+
+    kernel_parameters = tuple(interior_tendency_kernel_parameters(grid))
+
+    tracer_kernel_args   = (biogeochemistry, velocities, tracers, auxiliary_fields)
+
+    for tracer_index in 1:length(tracers)
+        @inbounds c_tendency = tendencies[tracer_index + 3]
+        @inbounds tracer_name = keys(tracers)[tracer_index]
+
+        args = tuple(Val(tracer_name), tracer_kernel_args..., clock)
+
+        for parameters in kernel_parameters
+            launch!(arch, grid, parameters, compute_bgc_G!, 
+                    c_tendency, grid, active_cells_map, args;
+                    active_cells_map)
+        end
+    end
+
+    return nothing
+end
+
+""" Calculate the right-hand-side of the tracer advection-diffusion equation. """
+@kernel function compute_bgc_G!(Gc, grid, ::Nothing, args)
+    i, j, k = @index(Global, NTuple)
+    @inbounds Gc[i, j, k] = biogeochemical_tendency(i, j, k, grid, args...)
+end
+
+@kernel function compute_bgc_G!(Gc, grid, interior_map, args) 
+    idx = @index(Global, Linear)
+    i, j, k = active_linear_index_to_tuple(idx, interior_map)
+    @inbounds Gc[i, j, k] = biogeochemical_tendency(i, j, k, grid, args...)
+end
+
+@inline function biogeochemical_tendency(i, j, k, grid,
+                                         val_tracer_name,
+                                         biogeochemistry,
+                                         velocities,
+                                         tracers,
+                                         auxiliary_fields,
+                                         clock)
+
+    model_fields = merge(velocities, tracers, auxiliary_fields)
+
+    return biogeochemical_transition(i, j, k, grid, biogeochemistry, val_tracer_name, clock, model_fields, Val(true))
+end
+
 end # module

src/Fields/sum_of_fields.jl

@@ -0,0 +1,73 @@
+using Base: @propagate_inbounds
+using Oceananigans.Fields: AbstractField, location
+
+import Adapt: adapt_structure
+import Base: getindex, show, summary
+import Oceananigans.BoundaryConditions: fill_halo_regions!
+
+"""
+    SumOfFields{N, F}
+
+`SumOfFields` objects hold `N` field/fields and return their sum when indexed.
+"""
+struct SumOfFields{N, LX, LY, LZ, G, T, F} <: AbstractField{LX, LY, LZ, G, T, 3}
+    fields :: F
+      grid :: G
+
+    function SumOfFields{N}(fields...) where N
+        grid = first(fields).grid
+        loc = location(first(fields))
+
+        all(f.grid == grid for f in fields) || 
+            throw(ArgumentError("All `fields` in `SumOfFields` must be on the same grid"))
+
+        all(location(f) == loc for f in fields) || 
+            throw(ArgumentError("All `fields` in `SumOfFields` must be on the same location"))
+
+        T = eltype(first(fields).data)
+        F = typeof(fields)
+        G = typeof(grid)
+
+        return new{N, loc..., G, T, F}(fields, grid)
+    end
+end
+
+adapt_structure(to, sum::SumOfFields{N}) where N = SumOfFields{N}((adapt_structure(to, field) for field in sum.fields)...)
+
+grid_name(field::SumOfFields) = grid_name(field.grid)
+
+function Base.summary(field::SumOfFields{N, LX, LY, LZ}) where {N, LX, LY, LZ}
+    prefix = string(size_summary(size(field)), " SumOfFields{$N, $LX, $LY, $LZ}")
+
+    suffix = string(" on ", grid_name(field), " on ", summary(architecture(field)))
+
+    return string(prefix, suffix)
+end
+
+Base.show(io::IO, sof::SumOfFields) =
+    print(io,
+          summary(sof), "\n",
+          "└── grid: ", summary(sof.grid))
+
+@propagate_inbounds function getindex(s::SumOfFields{N, LX, LY, LZ, G, T, F}, i...) where {N, LX, LY, LZ, G, T, F}
+    first = getindex(SumOfFields{3, F, LX, LY, LZ, G, T}((s.fields[1], s.fields[2], s.fields[3]), s.grid), i...)
+    last = getindex(SumOfFields{N - 3, F, LX, LY, LZ, G, T}(s.fields[4:N], s.grid), i...)
+    return first + last
+end
+
+@propagate_inbounds getindex(s::SumOfFields{1}, i...) = getindex(s.fields[1], i...)
+@propagate_inbounds getindex(s::SumOfFields{2}, i...) = getindex(s.fields[1], i...) + getindex(s.fields[2], i...)
+
+@propagate_inbounds getindex(s::SumOfFields{3}, i...) = 
+    getindex(s.fields[1], i...) + getindex(s.fields[2], i...) + getindex(s.fields[3], i...)
+
+@propagate_inbounds getindex(s::SumOfFields{4}, i...) = 
+    getindex(s.fields[1], i...) + getindex(s.fields[2], i...) + getindex(s.fields[3], i...) + getindex(s.fields[4], i...)
+
+function fill_halo_regions!(sof::SumOfFields)
+    for f in sof.fields
+        fill_halo_regions!(f)
+    end
+
+    return nothing
+end

src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl

@@ -10,6 +10,8 @@ using Oceananigans.Fields: immersed_boundary_condition
 using Oceananigans.Biogeochemistry: update_tendencies!
 using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: FlavorOfCATKE, FlavorOfTD
 
+using Oceananigans.TimeSteppers: compute_bgc_with_physics
+
 using Oceananigans.ImmersedBoundaries: retrieve_interior_active_cells_map, ActiveCellsIBG, 
                                        active_linear_index_to_tuple
 
@@ -88,7 +90,8 @@ function compute_hydrostatic_free_surface_tendency_contributions!(model, kernel_
                      model.diffusivity_fields,
                      model.auxiliary_fields,
                      c_forcing,
-                     model.clock)
+                     model.clock,
+                     Val(compute_bgc_with_physics(model.timestepper)))
 
         for parameters in kernel_parameters
             launch!(arch, grid, parameters,

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl

@@ -117,7 +117,8 @@ where `c = C[tracer_index]`.
                                                           diffusivities,
                                                           auxiliary_fields,
                                                           forcing,
-                                                          clock) where tracer_index
+                                                          clock,
+                                                          val_compute_bgc) where tracer_index
 
     @inbounds c = tracers[tracer_index]
     model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields)
@@ -133,7 +134,7 @@ where `c = C[tracer_index]`.
     return ( - div_Uc(i, j, k, grid, advection, total_velocities, c)
              - ∇_dot_qᶜ(i, j, k, grid, closure, diffusivities, val_tracer_index, c, clock, model_fields, buoyancy)
              - immersed_∇_dot_qᶜ(i, j, k, grid, c, c_immersed_bc, closure, diffusivities, val_tracer_index, clock, model_fields)
-             + biogeochemical_transition(i, j, k, grid, biogeochemistry, val_tracer_name, clock, model_fields)
+             + biogeochemical_transition(i, j, k, grid, biogeochemistry, val_tracer_name, clock, model_fields, val_compute_bgc)
              + forcing(i, j, k, grid, clock, model_fields))
 end
 

src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl

@@ -6,6 +6,8 @@ using Oceananigans.Models: complete_communication_and_compute_boundary!, interio
 using Oceananigans.ImmersedBoundaries: retrieve_interior_active_cells_map, ActiveCellsIBG, 
                                        active_linear_index_to_tuple
 
+using Oceananigans.TimeSteppers: compute_bgc_with_physics, timestepper_tendencies
+
 import Oceananigans.TimeSteppers: compute_tendencies!
 
 """
@@ -34,7 +36,7 @@ function compute_tendencies!(model::NonhydrostaticModel, callbacks)
 
     # Calculate contributions to momentum and tracer tendencies from user-prescribed fluxes across the
     # boundaries of the domain
-    compute_boundary_tendency_contributions!(model.timestepper.Gⁿ,
+    compute_boundary_tendency_contributions!(timestepper_tendencies(model.timestepper),
                                              model.architecture,
                                              model.velocities,
                                              model.tracers,
@@ -53,7 +55,7 @@ end
 """ Store previous value of the source term and compute current source term. """
 function compute_interior_tendency_contributions!(model, kernel_parameters; active_cells_map = nothing)
 
-    tendencies           = model.timestepper.Gⁿ
+    tendencies           = timestepper_tendencies(model.timestepper)
     arch                 = model.architecture
     grid                 = model.grid
     advection            = model.advection
@@ -126,7 +128,8 @@ function compute_interior_tendency_contributions!(model, kernel_parameters; acti
                      start_tracer_kernel_args..., 
                      c_immersed_bc,
                      end_tracer_kernel_args...,
-                     forcing, clock)
+                     forcing, clock,
+                     Val(compute_bgc_with_physics(model.timestepper)))
 
         for parameters in kernel_parameters
             launch!(arch, grid, parameters, compute_Gc!, 

src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl

@@ -248,7 +248,8 @@ velocity components, tracer fields, and precalculated diffusivities where applic
                                  auxiliary_fields,
                                  diffusivities,
                                  forcing,
-                                 clock) where tracer_index
+                                 clock,
+                                 val_compute_bgc) where tracer_index
 
     @inbounds c = tracers[tracer_index]
     @inbounds background_fields_c = background_fields.tracers[tracer_index]
@@ -266,6 +267,6 @@ velocity components, tracer fields, and precalculated diffusivities where applic
              - div_Uc(i, j, k, grid, advection, velocities, background_fields_c)
              - ∇_dot_qᶜ(i, j, k, grid, closure, diffusivities, val_tracer_index, c, clock, model_fields, buoyancy)
              - immersed_∇_dot_qᶜ(i, j, k, grid, c, c_immersed_bc, closure, diffusivities, val_tracer_index, clock, model_fields)
-             + biogeochemical_transition(i, j, k, grid, biogeochemistry, val_tracer_name, clock, model_fields)
+             + biogeochemical_transition(i, j, k, grid, biogeochemistry, val_tracer_name, clock, model_fields, val_compute_bgc)
              + forcing(i, j, k, grid, clock, model_fields))
 end

src/Models/ShallowWaterModels/rk3_substep_shallow_water_model.jl

@@ -14,14 +14,14 @@ function rk3_substep!(model::ShallowWaterModel, Δt, γⁿ, ζⁿ)
 
     substep_solution_kernel!(model.solution,
                              Δt, γⁿ, ζⁿ,
-                             model.timestepper.Gⁿ,
-                             model.timestepper.G⁻)
+                             timestepper_tendencies(timestepper),
+                             timestepper_previous_tendencies(timestepper))
 
 
     for i in 1:length(model.tracers)
         @inbounds c = model.tracers[i]
-        @inbounds Gcⁿ = model.timestepper.Gⁿ[i+3]
-        @inbounds Gc⁻ = model.timestepper.G⁻[i+3]
+        @inbounds Gcⁿ = timestepper_tendencies(timestepper)[i+3]
+        @inbounds Gc⁻ = timestepper_previous_tendencies(timestepper)[i+3]
 
         substep_tracer_kernel!(c, Δt, γⁿ, ζⁿ, Gcⁿ, Gc⁻)
     end

src/Models/interleave_communication_and_computation.jl

@@ -6,6 +6,8 @@ using Oceananigans.DistributedComputations
 using Oceananigans.DistributedComputations: DistributedGrid
 using Oceananigans.DistributedComputations: synchronize_communication!, SynchronizedDistributed
 
+import Oceananigans: interior_tendency_kernel_parameters
+
 function complete_communication_and_compute_boundary!(model, ::DistributedGrid, arch)
 
     # Iterate over the fields to clear _ALL_ possible architectures

src/Oceananigans.jl

@@ -190,6 +190,8 @@ function fields end
 function prognostic_fields end
 function tracer_tendency_kernel_function end
 function boundary_conditions end
+function compute_biogeochemical_tendencies! end
+function interior_tendency_kernel_parameters end
 
 #####
 ##### Include all the submodules

src/TimeSteppers/TimeSteppers.jl

@@ -55,9 +55,13 @@ step_lagrangian_particles!(model, Δt) = nothing
 reset!(timestepper) = nothing
 implicit_step!(field, ::Nothing, args...; kwargs...) = nothing
 
+compute_bgc_with_physics(timestepper) = true
+timestepper_tendencies(timestepper) = timestepper.Gⁿ
+timestepper_previous_tendencies(timestepper) = timestepper.G⁻
+
 include("clock.jl")
-include("store_tendencies.jl")
 include("quasi_adams_bashforth_2.jl")
 include("runge_kutta_3.jl")
-
+include("strange_splitting.jl")
+include("store_tendencies.jl")
 end # module

src/TimeSteppers/quasi_adams_bashforth_2.jl

@@ -73,16 +73,15 @@ The steps of the Quasi-Adams-Bashforth second-order (AB2) algorithm are:
 6. Update the model state.
 7. Compute tendencies for the next time step
 """
-function time_step!(model::AbstractModel{<:QuasiAdamsBashforth2TimeStepper}, Δt;
-                    callbacks=[], euler=false)
+time_step!(model::AbstractModel{<:QuasiAdamsBashforth2TimeStepper}, Δt; callbacks=[], euler=false) = 
+    time_step!(model.timestepper, model, Δt; callbacks, euler)
 
+function time_step!(ab2_timestepper::QuasiAdamsBashforth2TimeStepper, model, Δt; callbacks=[], euler=false)
     Δt == 0 && @warn "Δt == 0 may cause model blowup!"
 
     # Be paranoid and update state at iteration 0
     model.clock.iteration == 0 && update_state!(model, callbacks)
 
-    ab2_timestepper = model.timestepper
-
     # Change the default χ if necessary, which occurs if:
     #   * We detect that the time-step size has changed.
     #   * We detect that this is the "first" time-step, which means we
@@ -151,8 +150,8 @@ function ab2_step!(model, Δt)
     for (i, field) in enumerate(model_fields)
 
         step_field_kernel!(field, Δt, χ,
-                           model.timestepper.Gⁿ[i],
-                           model.timestepper.G⁻[i])
+                           timestepper_tendencies(timestepper)[i],
+                           timestepper_previous_tendencies(timestepper)[i])
 
         # TODO: function tracer_index(model, field_index) = field_index - 3, etc...
         tracer_index = Val(i - 3) # assumption