Reformulate hydrostatic model timestepping

.buildkite/pipeline.yml

@@ -84,7 +84,7 @@ steps:
           - "🦢 shallow_water"
           - "🐳 simulation"
           - "🍂 lagrangian_particles"
-          - "🧅 multi_region"
+          # - "🧅 multi_region"
           - "🦧 scripts"
           - "👺 enzyme"
           - "🍱 xesmf"

ext/OceananigansNCDatasetsExt.jl

@@ -7,9 +7,11 @@ using Dates: AbstractTime, UTC, now
 using Printf: @sprintf
 using OrderedCollections: OrderedDict
 
+using Oceananigans: fields
 using Oceananigans: initialize!, prettytime, pretty_filesize, AbstractModel
 using Oceananigans.Architectures: CPU, GPU, on_architecture
 using Oceananigans.AbstractOperations: KernelFunctionOperation, AbstractOperation
+using Oceananigans.DistributedComputations: synchronize_communication!
 using Oceananigans.BuoyancyFormulations: BuoyancyForce, BuoyancyTracer, SeawaterBuoyancy, LinearEquationOfState
 using Oceananigans.Fields
 using Oceananigans.Fields: Reduction, reduced_dimensions, reduced_location, location
@@ -1352,6 +1354,12 @@ Write output to netcdf file `output_writer.filepath` at specified intervals. Inc
 every time an output is written to the file.
 """
 function write_output!(ow::NetCDFWriter, model::AbstractModel)
+
+    # Synchronize model state if needed
+    for field in fields(model)
+        synchronize_communication!(field)
+    end
+
     # Start a new file if the file_splitting(model) is true
     ow.file_splitting(model) && start_next_file(model, ow)
     update_file_splitting_schedule!(ow.file_splitting, ow.filepath)

ext/OceananigansReactantExt/TimeSteppers.jl

@@ -12,8 +12,6 @@ using Oceananigans.Utils: @apply_regionally, apply_regionally!
 using Oceananigans.TimeSteppers:
     update_state!,
     tick!,
-    compute_pressure_correction!,
-    correct_velocities_and_cache_previous_tendencies!,
     step_lagrangian_particles!,
     QuasiAdamsBashforth2TimeStepper,
     compute_flux_bc_tendencies!
@@ -87,8 +85,7 @@ function time_step!(model::ReactantModel{<:QuasiAdamsBashforth2TimeStepper{FT}},
     ab2_timestepper.χ = χ
 
     # Full step for tracers, fractional step for velocities.
-    compute_flux_bc_tendencies!(model)
-    ab2_step!(model, Δt)
+    ab2_step!(model, Δt, callbacks)
 
     tick!(model.clock, Δt)
 
@@ -105,10 +102,7 @@ function time_step!(model::ReactantModel{<:QuasiAdamsBashforth2TimeStepper{FT}},
         model.clock.last_stage_Δt = Δt
     end
 
-    compute_pressure_correction!(model, Δt)
-    correct_velocities_and_cache_previous_tendencies!(model, Δt)
-
-    update_state!(model, callbacks; compute_tendencies=true)
+    update_state!(model, callbacks)
     step_lagrangian_particles!(model, Δt)
 
     # Return χ to initial value

src/Advection/Advection.jl

@@ -44,6 +44,8 @@ import Base: show, summary, Callable
 import Oceananigans.Grids: required_halo_size_x, required_halo_size_y, required_halo_size_z
 import Oceananigans.Architectures: on_architecture
 
+using Oceananigans.Grids: XFlatGrid, YFlatGrid, ZFlatGrid
+
 abstract type AbstractAdvectionScheme{B, FT} end
 abstract type AbstractCenteredAdvectionScheme{B, FT} <: AbstractAdvectionScheme{B, FT} end
 abstract type AbstractUpwindBiasedAdvectionScheme{B, FT} <: AbstractAdvectionScheme{B, FT} end
@@ -52,11 +54,11 @@ abstract type AbstractUpwindBiasedAdvectionScheme{B, FT} <: AbstractAdvectionSch
 # are constructed via metaprogramming. (The `advection_buffer` is the width of
 # the halo region required for an advection scheme on a non-immersed-boundary grid.)
 # An upper limit of `advection_buffer = 6` means we can build advection schemes up to
-# `Centered(order=12`) and `UpwindBiased(order=11)`. The list can be extended in order to
+# `Centered(order=10`) and `UpwindBiased(order=9)`. The list can be extended in order to
 # compile schemes with higher orders; for example `advection_buffers = [1, 2, 3, 4, 5, 6, 8]`
 # will compile schemes for `advection_buffer=8` and thus `Centered(order=16)` and `UpwindBiased(order=15)`.
 # Note that it is not possible to compile schemes for `advection_buffer = 41` or higher.
-const advection_buffers = [1, 2, 3, 4, 5, 6]
+const advection_buffers = [1, 2, 3, 4, 5]
 
 @inline Base.eltype(::AbstractAdvectionScheme{<:Any, FT}) where FT = FT
 

src/Advection/adapt_advection_order.jl

@@ -72,25 +72,26 @@ adapt_advection_order(advection, N, grid) = advection
 ##### Directional adapt advection order
 #####
 
-function adapt_advection_order(advection::Centered{B}, N::Int, grid::AbstractGrid) where B
+function adapt_advection_order(advection::Centered{B, FT}, N::Int, grid::AbstractGrid) where {B, FT}
     if N >= B
         return advection
     else
-        return Centered(; order=2N)
+        return Centered(FT; order=2N)
     end
 end
 
-function adapt_advection_order(advection::UpwindBiased{B}, N::Int, grid::AbstractGrid) where B
+function adapt_advection_order(advection::UpwindBiased{B, FT}, N::Int, grid::AbstractGrid) where {B, FT}
     if N >= B
         return advection
     else
-        return UpwindBiased(; order=2N-1)
+        return UpwindBiased(FT; order=2N-1)
     end
 end
-function adapt_advection_order(advection::WENO{B}, N::Int, grid::AbstractGrid) where B
+
+function adapt_advection_order(advection::WENO{B, FT}, N::Int, grid::AbstractGrid) where {B, FT}
     if N >= B
         return advection
     else
-        return WENO(order=2N-1)
+        return WENO(FT; order=2N-1)
     end
-end
+end

src/Advection/centered_reconstruction.jl

@@ -2,54 +2,62 @@
 ##### Centered advection scheme
 #####
 
-struct Centered{N, FT, CA} <: AbstractCenteredAdvectionScheme{N, FT}
-    buffer_scheme :: CA
-
-    Centered{N, FT}(buffer_scheme::CA) where {N, FT, CA} = new{N, FT, CA}(buffer_scheme)
-end
+struct Centered{N, FT} <: AbstractCenteredAdvectionScheme{N, FT} end
 
 function Centered(FT::DataType=Oceananigans.defaults.FloatType; order = 2)
+
     mod(order, 2) != 0 && throw(ArgumentError("Centered reconstruction scheme is defined only for even orders"))
 
-    N  = Int(order ÷ 2)
-    if N > 1
-        buffer_scheme = Centered(FT; order=order-2)
-    else
-        buffer_scheme = nothing
-    end
-
-    return Centered{N, FT}(buffer_scheme)
+    N = Int(order ÷ 2)
+    return Centered{N, FT}()
 end
 
 Base.summary(a::Centered{N}) where N = string("Centered(order=", 2N, ")")
-
-Base.show(io::IO, a::Centered{N, FT}) where {N, FT} =
-    print(io, summary(a), " \n",
-              "└── buffer_scheme: ", summary(a.buffer_scheme))
-
-
-Adapt.adapt_structure(to, scheme::Centered{N, FT}) where {N, FT} = Centered{N, FT}(Adapt.adapt(to, scheme.buffer_scheme))
-on_architecture(to, scheme::Centered{N, FT}) where {N, FT} = Centered{N, FT}(on_architecture(to, scheme.buffer_scheme))
+Base.show(io::IO, a::Centered{N, FT}) where {N, FT} = summary(a)
 
 const ACAS = AbstractCenteredAdvectionScheme
 
 # left and right biased for Centered reconstruction are just symmetric!
-@inline biased_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme::ACAS, bias, args...) = symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, args...)
-@inline biased_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme::ACAS, bias, args...) = symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, args...)
-@inline biased_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme::ACAS, bias, args...) = symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme, args...)
-@inline biased_interpolate_xᶜᵃᵃ(i, j, k, grid, scheme::ACAS, bias, args...) = symmetric_interpolate_xᶜᵃᵃ(i, j, k, grid, scheme, args...)
-@inline biased_interpolate_yᵃᶜᵃ(i, j, k, grid, scheme::ACAS, bias, args...) = symmetric_interpolate_yᵃᶜᵃ(i, j, k, grid, scheme, args...)
-@inline biased_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme::ACAS, bias, args...) = symmetric_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme, args...)
-
-# uniform centered reconstruction
-for buffer in advection_buffers, FT in fully_supported_float_types
-    @eval begin
-        @inline symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, ::Centered{$buffer, $FT}, ψ, args...) = @inbounds @muladd $(calc_reconstruction_stencil(FT, buffer, :symmetric, :x, false))
-        @inline symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, ::Centered{$buffer, $FT}, ψ, args...) = @inbounds @muladd $(calc_reconstruction_stencil(FT, buffer, :symmetric, :y, false))
-        @inline symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, ::Centered{$buffer, $FT}, ψ, args...) = @inbounds @muladd $(calc_reconstruction_stencil(FT, buffer, :symmetric, :z, false))
-
-        @inline symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, ::Centered{$buffer, $FT}, ψ::Callable, args...) = @inbounds @muladd $(calc_reconstruction_stencil(FT, buffer, :symmetric, :x,  true))
-        @inline symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, ::Centered{$buffer, $FT}, ψ::Callable, args...) = @inbounds @muladd $(calc_reconstruction_stencil(FT, buffer, :symmetric, :y,  true))
-        @inline symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, ::Centered{$buffer, $FT}, ψ::Callable, args...) = @inbounds @muladd $(calc_reconstruction_stencil(FT, buffer, :symmetric, :z,  true))
+@inline _biased_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme::ACAS, bias, c, args...) = _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, c, args...)
+@inline _biased_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme::ACAS, bias, c, args...) = _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, c, args...)
+@inline _biased_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme::ACAS, bias, c, args...) = _symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme, c, args...)
+@inline _biased_interpolate_xᶜᵃᵃ(i, j, k, grid, scheme::ACAS, bias, c, args...) = _symmetric_interpolate_xᶜᵃᵃ(i, j, k, grid, scheme, c, args...)
+@inline _biased_interpolate_yᵃᶜᵃ(i, j, k, grid, scheme::ACAS, bias, c, args...) = _symmetric_interpolate_yᵃᶜᵃ(i, j, k, grid, scheme, c, args...)
+@inline _biased_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme::ACAS, bias, c, args...) = _symmetric_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme, c, args...)
+
+for (side, dir) in zip((:ᶠᵃᵃ, :ᵃᶠᵃ, :ᵃᵃᶠ), (:x, :y, :z))
+    for (F, bool) in zip((:Any, :(Base.Callable)), (false, true))
+        for FT in fully_supported_float_types
+            interp = Symbol(:symmetric_interpolate_, dir, side)
+            @eval begin
+                @inline $interp(i, j, k, grid, ::Centered{1, $FT}, red_order::Int, ψ::$F, args...) = @muladd $(stencil_reconstruction(FT, 1, :symmetric, dir, bool))
+
+                @inline function $interp(i, j, k, grid, ::Centered{2, $FT}, red_order::Int, ψ::$F, args...)          
+                    @muladd ifelse(red_order==1, $(stencil_reconstruction(FT, 1, :symmetric, dir, bool)),
+                                                 $(stencil_reconstruction(FT, 2, :symmetric, dir, bool)))
+                end
+
+                @inline function $interp(i, j, k, grid, ::Centered{3, $FT}, red_order::Int, ψ::$F, args...)          
+                    @muladd ifelse(red_order==1, $(stencil_reconstruction(FT, 1, :symmetric, dir, bool)),
+                            ifelse(red_order==2, $(stencil_reconstruction(FT, 2, :symmetric, dir, bool)),
+                                                 $(stencil_reconstruction(FT, 3, :symmetric, dir, bool))))
+                end
+
+                @inline function $interp(i, j, k, grid, ::Centered{4, $FT}, red_order::Int, ψ::$F, args...)          
+                    @muladd ifelse(red_order==1, $(stencil_reconstruction(FT, 1, :symmetric, dir, bool)),
+                            ifelse(red_order==2, $(stencil_reconstruction(FT, 2, :symmetric, dir, bool)),
+                            ifelse(red_order==3, $(stencil_reconstruction(FT, 3, :symmetric, dir, bool)),
+                                                 $(stencil_reconstruction(FT, 4, :symmetric, dir, bool)))))
+                end
+
+                @inline function $interp(i, j, k, grid, ::Centered{5, $FT}, red_order::Int, ψ::$F, args...)          
+                    @muladd ifelse(red_order==1, $(stencil_reconstruction(FT, 1, :symmetric, dir, bool)),
+                            ifelse(red_order==2, $(stencil_reconstruction(FT, 2, :symmetric, dir, bool)),
+                            ifelse(red_order==3, $(stencil_reconstruction(FT, 3, :symmetric, dir, bool)),
+                            ifelse(red_order==4, $(stencil_reconstruction(FT, 4, :symmetric, dir, bool)),
+                                                 $(stencil_reconstruction(FT, 5, :symmetric, dir, bool))))))
+                end
+            end
+        end
     end
 end

src/Advection/flat_advective_fluxes.jl

@@ -34,19 +34,15 @@ FlatGrids = [:XFG, :YFG, :ZFG, :XFG, :YFG, :ZFG]
 for (dir, GridType) in zip((:xᶠᵃᵃ, :yᵃᶠᵃ, :zᵃᵃᶠ, :xᶜᵃᵃ, :yᵃᶜᵃ, :zᵃᵃᶜ), FlatGrids)
     alt_symm_interp   = Symbol(:_symmetric_interpolate_, dir)
     alt_biased_interp = Symbol(:_biased_interpolate_, dir)
-    @eval begin
-        @inline $alt_symm_interp(i, j, k, grid::$GridType, ::HOADV, ψ, args...) = @inbounds ψ[i, j, k]
-        @inline $alt_symm_interp(i, j, k, grid::$GridType, ::LOADV, ψ, args...) = @inbounds ψ[i, j, k]
-        @inline $alt_symm_interp(i, j, k, grid::$GridType, ::HOADV, ψ::Callable, args...) = @inbounds ψ(i, j, k, grid, args...)
-        @inline $alt_symm_interp(i, j, k, grid::$GridType, ::LOADV, ψ::Callable, args...) = @inbounds ψ(i, j, k, grid, args...)
-        @inline $alt_symm_interp(i, j, k, grid::$GridType, ::HOADV, ψ::Callable, ::AS, args...) = @inbounds ψ(i, j, k, grid, args...)
-        @inline $alt_symm_interp(i, j, k, grid::$GridType, ::LOADV, ψ::Callable, ::AS, args...) = @inbounds ψ(i, j, k, grid, args...)
-
-        @inline $alt_biased_interp(i, j, k, grid::$GridType, ::HOADV, ψ, args...) = @inbounds ψ[i, j, k]
-        @inline $alt_biased_interp(i, j, k, grid::$GridType, ::LOADV, ψ, args...) = @inbounds ψ[i, j, k]
-        @inline $alt_biased_interp(i, j, k, grid::$GridType, ::HOADV, bias, ψ::Callable, args...) = ψ(i, j, k, grid, args...)
-        @inline $alt_biased_interp(i, j, k, grid::$GridType, ::LOADV, bias, ψ::Callable, args...) = ψ(i, j, k, grid, args...)
-        @inline $alt_biased_interp(i, j, k, grid::$GridType, ::HOADV, bias, ψ::Callable, ::AS, args...) = ψ(i, j, k, grid, args...)
-        @inline $alt_biased_interp(i, j, k, grid::$GridType, ::LOADV, bias, ψ::Callable, ::AS, args...) = ψ(i, j, k, grid, args...)
+    for Scheme in [:Centered, :UpwindBiased, :WENO]
+        @eval begin
+            @inline $alt_symm_interp(i, j, k, grid::$GridType, ::$Scheme, ψ, args...) = @inbounds ψ[i, j, k]
+            @inline $alt_symm_interp(i, j, k, grid::$GridType, ::$Scheme, ψ::Callable, args...) = @inbounds ψ(i, j, k, grid, args...)
+            @inline $alt_symm_interp(i, j, k, grid::$GridType, ::$Scheme, ψ::Callable, ::AS, args...) = @inbounds ψ(i, j, k, grid, args...)
+
+            @inline $alt_biased_interp(i, j, k, grid::$GridType, ::$Scheme, ψ, args...) = @inbounds ψ[i, j, k]
+            @inline $alt_biased_interp(i, j, k, grid::$GridType, ::$Scheme, bias, ψ::Callable, args...) = ψ(i, j, k, grid, args...)
+            @inline $alt_biased_interp(i, j, k, grid::$GridType, ::$Scheme, bias, ψ::Callable, ::AS, args...) = ψ(i, j, k, grid, args...)
+        end
     end
-end
+end