diff --git a/src/Advection/Advection.jl b/src/Advection/Advection.jl index ca45d6a892c..060586bcf77 100644 --- a/src/Advection/Advection.jl +++ b/src/Advection/Advection.jl @@ -27,7 +27,7 @@ using Oceananigans.Utils: StaticStridesOffsetArray using Oceananigans: Oceananigans, fully_supported_float_types using Oceananigans.Architectures: Architectures, architecture, on_architecture, CPU -using Oceananigans.Grids: Grids, AbstractGrid, Center, Face, Flat, with_halo, horizontally_extended_interior_kernel_parameters, +using Oceananigans.Grids: Grids, AbstractGrid, Center, Face, Flat, horizontally_extended_interior_kernel_parameters, required_halo_size_x, required_halo_size_y, required_halo_size_z, XFlatGrid, YFlatGrid, ZFlatGrid using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid @@ -67,7 +67,7 @@ include("centered_reconstruction.jl") include("upwind_biased_reconstruction.jl") include("weno_reconstruction.jl") include("weno_interpolants.jl") -include("stretched_weno_smoothness.jl") +include("stretched_weno.jl") include("multi_dimensional_reconstruction.jl") include("vector_invariant_upwinding.jl") include("vector_invariant_advection.jl") diff --git a/src/Advection/adapt_advection_order.jl b/src/Advection/adapt_advection_order.jl index 182f53eeb4b..c164d15ef8e 100644 --- a/src/Advection/adapt_advection_order.jl +++ b/src/Advection/adapt_advection_order.jl @@ -91,6 +91,6 @@ function adapt_advection_order(advection::WENO{B}, N::Int, grid::AbstractGrid) w if N >= B return advection else - return WENO(eltype(grid); order=2N-1) + return WENO(eltype(grid); order=2N-1, stretched=something(advection.coefficients, false)) end end diff --git a/src/Advection/adaptive_implicit_vertical_advection.jl b/src/Advection/adaptive_implicit_vertical_advection.jl index 26230cefd6e..93857a76537 100644 --- a/src/Advection/adaptive_implicit_vertical_advection.jl +++ b/src/Advection/adaptive_implicit_vertical_advection.jl @@ -1,7 +1,7 @@ using Oceananigans.Operators: Δz⁻¹ᶜᶜᶜ, Δz⁻¹ᶜᶜᶠ, Δz⁻¹ᶠᶜᶠ, Δz⁻¹ᶜᶠᶠ, Az_qᶜᶜᶠ, Azᶜᶜᶠ, ℑxᶠᵃᵃ, ℑyᵃᶠᵃ using Oceananigans.Grids: Center, Face using Oceananigans.BoundaryConditions: BoundaryConditions, _unwrap_for_gpu -using Oceananigans.TimeSteppers: SplitRungeKuttaTimeStepper, RungeKutta3TimeStepper +using Oceananigans.TimeSteppers: SplitRungeKuttaTimeStepper, SSPRungeKuttaTimeStepper, RungeKutta3TimeStepper using Oceananigans.Utils: foreach_name const AVID = AdaptiveVerticallyImplicitDiscretization @@ -164,6 +164,9 @@ end return Δt / timestepper.β[nstage] end +# Every SSP stage is a forward-Euler step over the full Δt, so the next substep has the same Δτ. +@inline adaptive_advection_timestep(timestepper::SSPRungeKuttaTimeStepper, clock) = clock.last_stage_Δt + @inline sum_rk3_coefficients(ts, ::Val{1}) = ts.γ¹ @inline sum_rk3_coefficients(ts, ::Val{2}) = ts.γ² + ts.ζ² @inline sum_rk3_coefficients(ts, ::Val{3}) = ts.γ³ + ts.ζ³ diff --git a/src/Advection/bounds_preserving_tracer_advection_operators.jl b/src/Advection/bounds_preserving_tracer_advection_operators.jl index c0ee5c38d1c..6fde3ebb687 100644 --- a/src/Advection/bounds_preserving_tracer_advection_operators.jl +++ b/src/Advection/bounds_preserving_tracer_advection_operators.jl @@ -132,7 +132,7 @@ end @inline without_bounds_preservation(scheme) = scheme @inline function without_bounds_preservation(scheme::WENO{N, FT, WCT}) where {N, FT, WCT} - return WENO{N, FT, WCT}(nothing, scheme.buffer_scheme, scheme.advecting_velocity_scheme, scheme.time_discretization) + return WENO{N, FT, WCT}(nothing, scheme.buffer_scheme, scheme.advecting_velocity_scheme, scheme.time_discretization, scheme.coefficients) end ##### diff --git a/src/Advection/materialize_advection.jl b/src/Advection/materialize_advection.jl index bccc15fb8cf..cfcee53b843 100644 --- a/src/Advection/materialize_advection.jl +++ b/src/Advection/materialize_advection.jl @@ -47,14 +47,16 @@ materialize_advection(weno::WENO{N, FT, WCT}, grid) where {N, FT, WCT} = WENO{N, FT, WCT}(materialize_bounds(weno.bounds, grid), materialize_advection(without_bounds_preservation(weno.buffer_scheme), grid), materialize_advection(weno.advecting_velocity_scheme, grid), - weno.time_discretization) + weno.time_discretization, + materialize_coefficients(weno.coefficients, weno, grid)) function materialize_advection(weno::WENO{N, FT, Nothing}, grid) where {N, FT} WTC = default_weno_weight_computation(architecture(grid)) return WENO{N, FT, WTC}(materialize_bounds(weno.bounds, grid), materialize_advection(without_bounds_preservation(weno.buffer_scheme), grid), materialize_advection(weno.advecting_velocity_scheme, grid), - weno.time_discretization) + weno.time_discretization, + materialize_coefficients(weno.coefficients, weno, grid)) end materialize_advection(scheme::UpwindBiased{N, FT}, grid) where {N, FT} = diff --git a/src/Advection/reconstruction_coefficients.jl b/src/Advection/reconstruction_coefficients.jl index 9ffeb42c979..394d2c5f234 100644 --- a/src/Advection/reconstruction_coefficients.jl +++ b/src/Advection/reconstruction_coefficients.jl @@ -1,5 +1,3 @@ -using Oceananigans.Grids: ξnodes, ηnodes, rnodes - # Generic reconstruction methods valid for all reconstruction schemes # Unroll the functions to pass the coordinates in case of a stretched grid """ same as [`symmetric_interpolate_xᶠᵃᵃ`](@ref) but on `Center`s instead of `Face`s """ @@ -163,123 +161,3 @@ julia> calc_reconstruction_stencil(Float32, 3, :left, :x) end return Expr(:call, :+, stencil_full...) end - -##### -##### Shenanigans for stretched directions -##### - -@inline function reconstruction_stencil(buffer, shift, dir, func::Bool = false) - N = buffer * 2 - order = shift == :symmetric ? N : N - 1 - if shift != :symmetric - N = N .- 1 - end - rng = 1:N - if shift == :right - rng = rng .+ 1 - end - stencil_full = Vector(undef, N) - for (idx, n) in enumerate(rng) - c = n - buffer - 1 - if func - stencil_full[idx] = dir == :x ? - :(ψ(i + $c, j, k, grid, args...)) : - dir == :y ? - :(ψ(i, j + $c, k, grid, args...)) : - :(ψ(i, j, k + $c, grid, args...)) - else - stencil_full[idx] = dir == :x ? - :(ψ[i + $c, j, k]) : - dir == :y ? - :(ψ[i, j + $c, k]) : - :(ψ[i, j, k + $c]) - end - end - return :($(reverse(stencil_full)...),) -end - -@inline function compute_reconstruction_coefficients(grid, FT, scheme; order) - - method = scheme == :Centered ? 1 : scheme == :Upwind ? 2 : 3 - - if grid isa Nothing - coeff_xᶠᵃᵃ = nothing - coeff_xᶜᵃᵃ = nothing - coeff_yᵃᶠᵃ = nothing - coeff_yᵃᶜᵃ = nothing - coeff_zᵃᵃᶠ = nothing - coeff_zᵃᵃᶜ = nothing - else - arch = architecture(grid) - Hx, Hy, Hz = halo_size(grid) - new_grid = with_halo((Hx+1, Hy+1, Hz+1), grid) - - ξᶠᵃᵃ = ξnodes(new_grid, Face(), with_halos=true) - ηᵃᶠᵃ = ηnodes(new_grid, Face(), with_halos=true) - rᵃᵃᶠ = rnodes(new_grid, Face(), with_halos=true) - - ξᶜᵃᵃ = ξnodes(new_grid, Center(), with_halos=true) - ηᵃᶜᵃ = ηnodes(new_grid, Center(), with_halos=true) - rᵃᵃᶜ = rnodes(new_grid, Center(), with_halos=true) - - coeff_xᶠᵃᵃ = reconstruction_coefficients(FT, ξᶠᵃᵃ, arch, new_grid.Nx, Val(method); order) - coeff_xᶜᵃᵃ = reconstruction_coefficients(FT, ξᶜᵃᵃ, arch, new_grid.Nx, Val(method); order) - coeff_yᵃᶠᵃ = reconstruction_coefficients(FT, ηᵃᶠᵃ, arch, new_grid.Ny, Val(method); order) - coeff_yᵃᶜᵃ = reconstruction_coefficients(FT, ηᵃᶜᵃ, arch, new_grid.Ny, Val(method); order) - coeff_zᵃᵃᶠ = reconstruction_coefficients(FT, rᵃᵃᶠ, arch, new_grid.Nz, Val(method); order) - coeff_zᵃᵃᶜ = reconstruction_coefficients(FT, rᵃᵃᶜ, arch, new_grid.Nz, Val(method); order) - end - - return (coeff_xᶠᵃᵃ, coeff_xᶜᵃᵃ, coeff_yᵃᶠᵃ, coeff_yᵃᶜᵃ, coeff_zᵃᵃᶠ, coeff_zᵃᵃᶜ) -end - -# Fallbacks for uniform or Flat directions -for val in [1, 2, 3] - @eval begin - @inline reconstruction_coefficients(FT, coord::StaticStridesOffsetArray{<:Any, <:Any, <:Any, <:Any, <:AbstractRange}, arch, N, ::Val{$val}; order) = nothing - @inline reconstruction_coefficients(FT, coord::AbstractRange, arch, N, ::Val{$val}; order) = nothing - @inline reconstruction_coefficients(FT, coord::Nothing, arch, N, ::Val{$val}; order) = nothing - @inline reconstruction_coefficients(FT, coord::Number, arch, N, ::Val{$val}; order) = nothing - end -end - -# Stretched reconstruction coefficients for `Centered` schemes -@inline function reconstruction_coefficients(FT, coord, arch, N, ::Val{1}; order) - cpu_coord = on_architecture(CPU(), coord) - r = ((order + 1) ÷ 2) - 1 - s = create_reconstruction_coefficients(FT, r, cpu_coord, arch, N; order) - return s -end - -# Stretched reconstruction coefficients for `UpwindBiased` schemes -@inline function reconstruction_coefficients(FT, coord, arch, N, ::Val{2}; order) - cpu_coord = on_architecture(CPU(), coord) - rleft = ((order + 1) ÷ 2) - 2 - rright = ((order + 1) ÷ 2) - 1 - s = [] - for r in [rleft, rright] - push!(s, create_reconstruction_coefficients(FT, r, cpu_coord, arch, N; order)) - end - return tuple(s...) -end - -# Stretched reconstruction coefficients for `WENO` schemes -@inline function reconstruction_coefficients(FT, coord, arch, N, ::Val{3}; order) - cpu_coord = on_architecture(CPU(), coord) - s = [] - for r in -1:order-1 - push!(s, create_reconstruction_coefficients(FT, r, cpu_coord, arch, N; order)) - end - return tuple(s...) -end - -# general reconstruction coefficients for order `order` and stencil `r` where r -@inline function create_reconstruction_coefficients(FT, r, cpu_coord, arch, N; order) - stencil = NTuple{order, FT}[] - @inbounds begin - for i = 0:N+1 - push!(stencil, stencil_coefficients(FT, i, r, cpu_coord, cpu_coord; order)) - end - end - return StaticStridesOffsetArray(on_architecture(arch, stencil), -1) -end diff --git a/src/Advection/stretched_weno.jl b/src/Advection/stretched_weno.jl new file mode 100644 index 00000000000..d2cd4160fb9 --- /dev/null +++ b/src/Advection/stretched_weno.jl @@ -0,0 +1,84 @@ +using Oceananigans.Grids: ξnodes, ηnodes, rnodes +using Oceananigans.ImmersedBoundaries: underlying_grid + +##### +##### WENO on stretched directions: reconstruction coefficients, smoothness indicators and optimal weights computed +##### exactly from the cell widths and tabulated along each direction +##### + +# The weights of the averages over the cells between `edges` that give the value at x★ of the polynomial with those averages +function cell_average_weights(edges, x★) + n = length(edges) - 1 + averages = [(edges[j+1]^(q+1) - edges[j]^(q+1)) / ((q + 1) * (edges[j+1] - edges[j])) for q in 0:n-1, j in 1:n] + return averages \ [x★^q for q in 0:n-1] +end + +# The coefficients of the 2N - 1 cells between the 2N `edges`, reconstructing at the right edge of the anchor cell N +function biased_coefficients(FT, N, edges) + R = Rational{BigInt} + edges = R.(edges) + ξ = (edges .- (edges[N] + edges[N+1]) / 2) ./ (edges[N+1] - edges[N]) + candidate(s) = ξ[N-s:2N-s] + weights = [cell_average_weights(candidate(s), R(1, 2)) for s in 0:N-1] + + # β is invariant to a constant shift of the stencil, so it is a quadratic form in the first differences + T = [R(i < m) for m in 1:N, i in 1:N-1] + normalization = reference_smoothness_normalization[N - 1] + β = Tuple(FT.(ldlt_smoothness_coefficients(T' * smoothness_matrix(candidate(s)) * T * normalization)) for s in 0:N-1) + p = Tuple(FT.(Tuple(difference_coefficients(weights[s+1], N, s))) for s in 0:N-1) + + # The optimal weights combine the candidates into the reconstruction from the whole stencil + A = zeros(R, 2N - 1, N) + for s in 0:N-1, m in 1:N + A[N - s + m - 1, s + 1] = weights[s + 1][m] + end + C★ = (A' * A) \ (A' * cell_average_weights(ξ, R(1, 2))) + + return (; p, β, C★ = Tuple(FT.(C★))) +end + +# The coefficients at every point of `nodes`, the faces or the centers of a direction with its halos, from the cells +# between consecutive nodes. Stencils that leave the halos reuse the nearest complete one. +function coefficient_table(FT, N, nodes, arch) + entries = map(firstindex(nodes):lastindex(nodes)) do p + p = clamp(p, firstindex(nodes) + N, lastindex(nodes) - N) + left = biased_coefficients(FT, N, [nodes[q] for q in p-N:p+N-1]) + right = biased_coefficients(FT, N, [-nodes[q] for q in p+N:-1:p-N+1]) # mirrored + (; left, right) + end + return StaticStridesOffsetArray(on_architecture(arch, entries), firstindex(nodes) - 1) +end + +materialize_coefficients(coefficients, scheme, grid) = coefficients + +function materialize_coefficients(directions::Tuple, ::WENO{N, FT}, grid) where {N, FT} + arch = architecture(grid) + cpu_grid = on_architecture(CPU(), underlying_grid(grid)) + + # Flat and uniform directions, and the curvilinear horizontal ones (without ξnodes and ηnodes), keep the constants + function direction_coefficients(direction, nodes) + faces = direction ∈ directions && applicable(nodes, cpu_grid, Face()) ? nodes(cpu_grid, Face(); with_halos = true) : nothing + isnothing(faces) && return nothing + Δ = diff(collect(faces)) + all(≈(first(Δ)), Δ) && return nothing + centers = nodes(cpu_grid, Center(); with_halos = true) + return (faces = coefficient_table(FT, N, faces, arch), centers = coefficient_table(FT, N, centers, arch)) + end + + return NamedTuple{(:x, :y, :z)}(map(direction_coefficients, (:x, :y, :z), (ξnodes, ηnodes, rnodes))) +end + +##### +##### Reconstruction with tabulated coefficients +##### + +@inline weno_coefficients(c::NamedTuple, ::Val{d}, location, n, bias) where d = weno_coefficients(getfield(c, d), location, n, bias) + +@inline function weno_coefficients(c::NamedTuple, ::Val{location}, n, bias) where location + entry = @inbounds getfield(c, location)[n] + return ifelse(bias == LeftBias, entry.left, entry.right) +end + +@inline coeff_p(scheme, ::Val{s}, C) where s = @inbounds C.p[s+1] +@inline C★(scheme, ::Val{s}, C) where s = @inbounds C.C★[s+1] +@inline smoothness_indicator(δ, scheme, ::Val{s}, C) where s = smoothness_operation(scheme, δ, @inbounds C.β[s+1]) diff --git a/src/Advection/stretched_weno_smoothness.jl b/src/Advection/stretched_weno_smoothness.jl deleted file mode 100644 index a624e18a384..00000000000 --- a/src/Advection/stretched_weno_smoothness.jl +++ /dev/null @@ -1,111 +0,0 @@ -using KernelAbstractions.Extras.LoopInfo: @unroll - -##### -##### Stretched smoothness indicators gathered from precomputed values. -##### The stretched values for β coefficients are calculated from -##### Shu, NASA/CR-97-206253, ICASE Report No. 97-65 -##### by hardcoding that p(x) is a 2nd order polynomial -##### - -@inline function biased_left_β(ψ, scheme, r, dir, i, location) - @inbounds begin - stencil = retrieve_left_smooth(scheme, r, dir, i, location) - wᵢᵢ = stencil[1] - wᵢⱼ = stencil[2] - result = 0 - @unroll for j = 1:3 - result += ψ[j] * ( wᵢᵢ[j] * ψ[j] + wᵢⱼ[j] * dagger(ψ)[j] ) - end - end - return result -end - -@inline function biased_right_β(ψ, scheme, r, dir, i, location) - @inbounds begin - stencil = retrieve_right_smooth(scheme, r, dir, i, location) - wᵢᵢ = stencil[1] - wᵢⱼ = stencil[2] - result = 0 - @unroll for j = 1:3 - result += ψ[j] * ( wᵢᵢ[j] * ψ[j] + wᵢⱼ[j] * dagger(ψ)[j] ) - end - end - return result -end - -@inline left_biased_β₀(FT, ψ, T, scheme, args...) = biased_left_β(ψ, scheme, 0, args...) -@inline left_biased_β₁(FT, ψ, T, scheme, args...) = biased_left_β(ψ, scheme, 1, args...) -@inline left_biased_β₂(FT, ψ, T, scheme, args...) = biased_left_β(ψ, scheme, 2, args...) - -@inline right_biased_β₀(FT, ψ, T, scheme, args...) = biased_right_β(ψ, scheme, 2, args...) -@inline right_biased_β₁(FT, ψ, T, scheme, args...) = biased_right_β(ψ, scheme, 1, args...) -@inline right_biased_β₂(FT, ψ, T, scheme, args...) = biased_right_β(ψ, scheme, 0, args...) - -@inline retrieve_left_smooth(scheme, r, ::Val{1}, i, ::Type{Face}) = @inbounds scheme.smooth_xᶠᵃᵃ[r+1][i] -@inline retrieve_left_smooth(scheme, r, ::Val{1}, i, ::Type{Center}) = @inbounds scheme.smooth_xᶜᵃᵃ[r+1][i] -@inline retrieve_left_smooth(scheme, r, ::Val{2}, i, ::Type{Face}) = @inbounds scheme.smooth_yᵃᶠᵃ[r+1][i] -@inline retrieve_left_smooth(scheme, r, ::Val{2}, i, ::Type{Center}) = @inbounds scheme.smooth_yᵃᶜᵃ[r+1][i] -@inline retrieve_left_smooth(scheme, r, ::Val{3}, i, ::Type{Face}) = @inbounds scheme.smooth_zᵃᵃᶠ[r+1][i] -@inline retrieve_left_smooth(scheme, r, ::Val{3}, i, ::Type{Center}) = @inbounds scheme.smooth_zᵃᵃᶜ[r+1][i] - -@inline retrieve_right_smooth(scheme, r, ::Val{1}, i, ::Type{Face}) = @inbounds scheme.smooth_xᶠᵃᵃ[r+4][i] -@inline retrieve_right_smooth(scheme, r, ::Val{1}, i, ::Type{Center}) = @inbounds scheme.smooth_xᶜᵃᵃ[r+4][i] -@inline retrieve_right_smooth(scheme, r, ::Val{2}, i, ::Type{Face}) = @inbounds scheme.smooth_yᵃᶠᵃ[r+4][i] -@inline retrieve_right_smooth(scheme, r, ::Val{2}, i, ::Type{Center}) = @inbounds scheme.smooth_yᵃᶜᵃ[r+4][i] -@inline retrieve_right_smooth(scheme, r, ::Val{3}, i, ::Type{Face}) = @inbounds scheme.smooth_zᵃᵃᶠ[r+4][i] -@inline retrieve_right_smooth(scheme, r, ::Val{3}, i, ::Type{Center}) = @inbounds scheme.smooth_zᵃᵃᶜ[r+4][i] - -@inline calc_smoothness_coefficients(FT, ::Val{false}, args...; kwargs...) = nothing -@inline calc_smoothness_coefficients(FT, ::Val{true}, coord::StaticStridesOffsetArray{<:Any, <:Any, <:Any, <:Any, <:AbstractRange}, arch, N; order) = nothing -@inline calc_smoothness_coefficients(FT, ::Val{true}, coord::AbstractRange, arch, N; order) = nothing - -function calc_smoothness_coefficients(FT, beta, coord, arch, N; order) - - cpu_coord = on_architecture(CPU(), coord) - - order == 3 || throw(ArgumentError("The stretched smoothness coefficients are only implemented for order == 3")) - - s1 = create_smoothness_coefficients(FT, 0, -, cpu_coord, arch, N; order) - s2 = create_smoothness_coefficients(FT, 1, -, cpu_coord, arch, N; order) - s3 = create_smoothness_coefficients(FT, 2, -, cpu_coord, arch, N; order) - s4 = create_smoothness_coefficients(FT, 0, +, cpu_coord, arch, N; order) - s5 = create_smoothness_coefficients(FT, 1, +, cpu_coord, arch, N; order) - s6 = create_smoothness_coefficients(FT, 2, +, cpu_coord, arch, N; order) - - return (s1, s2, s3, s4, s5, s6) -end - -function create_smoothness_coefficients(FT, r, op, cpu_coord, arch, N; order) - - # derivation written on overleaf - stencil = NTuple{2, NTuple{order, FT}}[] - @inbounds begin - for i = 0:N+1 - - bias1 = Int(op == +) - bias2 = bias1 - 1 - - Δcᵢ = cpu_coord[i + bias1] - cpu_coord[i + bias2] - - Bᵢ = stencil_coefficients(i, r, cpu_coord, cpu_coord; order, op, shift = bias1, der = Primitive()) - bᵢ = stencil_coefficients(i, r, cpu_coord, cpu_coord; order, op, shift = bias1) - bₓᵢ = stencil_coefficients(i, r, cpu_coord, cpu_coord; order, op, shift = bias1, der = FirstDerivative()) - Aᵢ = stencil_coefficients(i, r, cpu_coord, cpu_coord; order, op, shift = bias2, der = Primitive()) - aᵢ = stencil_coefficients(i, r, cpu_coord, cpu_coord; order, op, shift = bias2) - aₓᵢ = stencil_coefficients(i, r, cpu_coord, cpu_coord; order, op, shift = bias2, der = FirstDerivative()) - pₓₓ = stencil_coefficients(i, r, cpu_coord, cpu_coord; order, op, shift = bias1, der = SecondDerivative()) - - Pᵢ = (Bᵢ .- Aᵢ) - - wᵢᵢ = Δcᵢ .* (bᵢ .* bₓᵢ .- aᵢ .* aₓᵢ .- pₓₓ .* Pᵢ) .+ Δcᵢ^4 .* (pₓₓ .* pₓₓ) - wᵢⱼ = Δcᵢ .* (star(bᵢ, bₓᵢ) .- star(aᵢ, aₓᵢ) .- star(pₓₓ, Pᵢ)) .+ Δcᵢ^4 .* star(pₓₓ, pₓₓ) - - push!(stencil, (wᵢᵢ, wᵢⱼ)) - end - end - - return StaticStridesOffsetArray(on_architecture(arch, stencil), -1) -end - -@inline dagger(ψ) = (ψ[2], ψ[3], ψ[1]) -@inline star(ψ₁, ψ₂) = (ψ₁ .* dagger(ψ₂) .+ dagger(ψ₁) .* ψ₂) diff --git a/src/Advection/upwind_biased_reconstruction.jl b/src/Advection/upwind_biased_reconstruction.jl index 26bb44a2c8b..5c54ee77733 100644 --- a/src/Advection/upwind_biased_reconstruction.jl +++ b/src/Advection/upwind_biased_reconstruction.jl @@ -23,11 +23,6 @@ function UpwindBiased(FT::DataType = Oceananigans.defaults.FloatType; advecting_velocity_scheme = Centered(FT; order = 2) if N > 1 - # coefficients = Tuple(nothing for i in 1:6) - # Stretched coefficient seem to be more unstable that constant spacing ones for - # linear (non-WENO) upwind reconstruction. We keep constant coefficients for the moment - # Some tests are needed to verify why this is the case (and if it is expected) - # coefficients = compute_reconstruction_coefficients(grid, FT, :Upwind; order) if buffer_scheme isa DecreasingOrderAdvectionScheme if order ≤ minimum_buffer_upwind_order # At minimum order, switch to Centered scheme diff --git a/src/Advection/weno_interpolants.jl b/src/Advection/weno_interpolants.jl index 8a11a17cad6..d7bc57e653e 100644 --- a/src/Advection/weno_interpolants.jl +++ b/src/Advection/weno_interpolants.jl @@ -103,11 +103,12 @@ for FT in fully_supported_float_types end end +@inline C★(scheme, stencil::Val, ::Nothing) = C★(scheme, stencil) + # Reconstruction coefficients on the first differences `δ[i] = S[i+1] - S[i]`, relative to the anchor -# `S[buffer]`. The change of basis is exact because `∑ⱼ cⱼ = 1`. -function difference_reconstruction_coefficients(FT, buffer, stencil) - c = stencil_coefficients(BigFloat, 50, stencil, collect(1:100), collect(1:100); order=buffer) - d = zeros(BigFloat, buffer - 1) +# `S[buffer]`, from the coefficients `c` of the cell values. The change of basis is exact because `∑ⱼ cⱼ = 1`. +function difference_coefficients(c, buffer, stencil) + d = zeros(eltype(c), buffer - 1) for k in 1:buffer - 1 i = buffer - stencil + k - 1 # the difference δ[i] that d[k] multiplies @@ -123,9 +124,12 @@ function difference_reconstruction_coefficients(FT, buffer, stencil) end end - return FT.(tuple(d...)) + return d end +difference_reconstruction_coefficients(FT, buffer, stencil) = + FT.(tuple(difference_coefficients(collect(stencil_coefficients(BigFloat, 50, stencil, collect(1:100), collect(1:100); order=buffer)), buffer, stencil)...)) + # ENO reconstruction procedure per stencil for buffer in advection_buffers[2:end] # WENO{<:Any, 1} does not exist for stencil in collect(0:1:buffer-1) @@ -156,12 +160,14 @@ for buffer in advection_buffers[2:end] # WENO{<:Any, 1} does not exist where ``dᵢ`` is computed from the function `coeff_p` """ - @inline biased_p(scheme::WENO{$buffer}, ::Val{$stencil}, δ) = - @inbounds @muladd $(Expr(:call, :+, (:(coeff_p(scheme, Val($stencil))[$i] * δ[$i]) for i in 1:buffer-1)...)) + @inline biased_p(scheme::WENO{$buffer}, ::Val{$stencil}, δ, coefficients = nothing) = + @inbounds @muladd $(Expr(:call, :+, (:(coeff_p(scheme, Val($stencil), coefficients)[$i] * δ[$i]) for i in 1:buffer-1)...)) end end end +@inline coeff_p(scheme, stencil::Val, ::Nothing) = coeff_p(scheme, stencil) + ##### ##### Smoothness indicators ##### @@ -189,15 +195,18 @@ References * Jiang, G.-S., & Shu, C.-W. (1996). "Efficient implementation of weighted ENO schemes." Journal of Computational Physics, 126(1), 202-228. """ -function exact_smoothness_matrix(buffer, stencil) - R = Rational{BigInt} - k = buffer +exact_smoothness_matrix(buffer, stencil) = + smoothness_matrix([Rational{BigInt}(m - (stencil + 1)) - 1//2 for m in 1:buffer+1]) * reference_smoothness_normalization[buffer - 1] - # Cell averages of the monomials xⁿ over the cells of the stencil, centered at m - (stencil + 1) +# The Jiang & Shu quadratic form of the cells between consecutive `edges`, in units of the reconstructed cell [-1/2, 1/2] +function smoothness_matrix(edges) + R = eltype(edges) + k = length(edges) - 1 + + # Cell averages of the monomials xⁿ over the cells A = zeros(R, k, k) for m in 1:k, n in 0:k-1 - c = R(m - (stencil + 1)) - A[m, n+1] = ((c + 1//2)^(n+1) - (c - 1//2)^(n+1)) / (n + 1) + A[m, n+1] = (edges[m+1]^(n+1) - edges[m]^(n+1)) / ((n + 1) * (edges[m+1] - edges[m])) end B = inv(A) # polynomial coefficients aₙ = ∑ⱼ B[n+1, j] ψⱼ @@ -216,7 +225,7 @@ function exact_smoothness_matrix(buffer, stencil) Q += Dₗ' * M * Dₗ end - return Q * reference_smoothness_normalization[buffer - 1] + return Q end # LDLᵀ factorization of the quadratic form β = δᵀ Qδ δ, packed as, for each l, the strictly-lower column of L @@ -349,6 +358,8 @@ for buffer in advection_buffers[2:end] # WENO{<:Any, 1} does not exist end end +@inline smoothness_indicator(δ, scheme, stencil::Val, ::Nothing) = smoothness_indicator(δ, scheme, stencil) + # Shenanigans for WENO weights calculation for vector invariant formulation -> [β[i] = 0.5 * (βᵤ[i] + βᵥ[i]) for i in 1:buffer] @inline function metaprogrammed_beta_sum(buffer) elem = Vector(undef, buffer) @@ -366,7 +377,7 @@ stencil_differences(buffer, stencil) = Expr(:tuple, (:(δ[$i]) for i in (buffer @inline function metaprogrammed_beta_loop(buffer) elem = Vector(undef, buffer) for stencil = 1:buffer - elem[stencil] = :(smoothness_indicator($(stencil_differences(buffer, stencil-1)), scheme, Val($(stencil-1)))) + elem[stencil] = :(smoothness_indicator($(stencil_differences(buffer, stencil-1)), scheme, Val($(stencil-1)), coefficients)) end return :($(elem...),) @@ -376,7 +387,7 @@ end @inline function metaprogrammed_zweno_alpha_loop(buffer) elem = Vector(undef, buffer) for stencil = 1:buffer - elem[stencil] = :(C★(scheme, Val($(stencil-1))) * zweno_ratio_term(newton_div(WCT, τ, β[$stencil] + ϵ★))) + elem[stencil] = :(C★(scheme, Val($(stencil-1)), coefficients) * zweno_ratio_term(newton_div(WCT, τ, β[$stencil] + ϵ★))) end return quote @@ -398,8 +409,8 @@ end for buffer in advection_buffers[2:end] @eval begin @inline beta_sum(scheme::WENO{$buffer, FT}, β₁, β₂) where FT = @inbounds $(metaprogrammed_beta_sum(buffer)) - @inline beta_loop(scheme::WENO{$buffer, FT}, δ) where FT = @inbounds $(metaprogrammed_beta_loop(buffer)) - @inline zweno_alpha_loop(scheme::WENO{$buffer, FT, WCT}, β, τ) where {FT, WCT} = @inbounds $(metaprogrammed_zweno_alpha_loop(buffer)) + @inline beta_loop(scheme::WENO{$buffer, FT}, δ, coefficients = nothing) where FT = @inbounds $(metaprogrammed_beta_loop(buffer)) + @inline zweno_alpha_loop(scheme::WENO{$buffer, FT, WCT}, β, τ, coefficients = nothing) where {FT, WCT} = @inbounds $(metaprogrammed_zweno_alpha_loop(buffer)) end end @@ -444,25 +455,25 @@ where The ``α`` values are normalized before returning """ -@inline function biased_weno_weights(δ, grid, scheme::WENO{N, FT}, args...) where {N, FT} - β = beta_loop(scheme, δ) +@inline function biased_weno_weights(δ, grid, scheme::WENO{N, FT}, coefficients = nothing, args...) where {N, FT} + β = beta_loop(scheme, δ, coefficients) τ = global_smoothness_indicator(Val(N), β) - α = zweno_alpha_loop(scheme, β, τ) + α = zweno_alpha_loop(scheme, β, τ, coefficients) Σα⁻¹ = weight_normalization(scheme, α) return α .* Σα⁻¹ end -@inline function biased_weno_weights(ijk, grid, scheme::WENO{N, FT}, bias, dir, ::VelocityStencil, u, v) where {N, FT} +@inline function biased_weno_weights(ijk, grid, scheme::WENO{N, FT}, coefficients, bias, dir, ::VelocityStencil, u, v) where {N, FT} i, j, k = ijk uₛ = tangential_stencil_u(i, j, k, grid, scheme, bias, dir, u) vₛ = tangential_stencil_v(i, j, k, grid, scheme, bias, dir, v) - βᵤ = beta_loop(scheme, weno_differences(scheme, uₛ)) - βᵥ = beta_loop(scheme, weno_differences(scheme, vₛ)) + βᵤ = beta_loop(scheme, weno_differences(scheme, uₛ), coefficients) + βᵥ = beta_loop(scheme, weno_differences(scheme, vₛ), coefficients) β = beta_sum(scheme, βᵤ, βᵥ) τ = global_smoothness_indicator(Val(N), β) - α = zweno_alpha_loop(scheme, β, τ) + α = zweno_alpha_loop(scheme, β, τ, coefficients) Σα⁻¹ = weight_normalization(scheme, α) return α .* Σα⁻¹ end @@ -556,7 +567,7 @@ end @inline function metaprogrammed_weno_reconstruction(buffer) elem = Vector(undef, buffer) for stencil = 1:buffer - elem[stencil] = :(ω[$stencil] * biased_p(scheme, Val($(stencil-1)), $(stencil_differences(buffer, stencil-1)))) + elem[stencil] = :(ω[$stencil] * biased_p(scheme, Val($(stencil-1)), $(stencil_differences(buffer, stencil-1)), coefficients)) end return Expr(:call, :+, :ψ₀, elem...) @@ -584,58 +595,67 @@ Here, [`biased_p`](@ref) is the function that computes the linear reconstruction # Calculation of WENO reconstructed value v⋆ = ∑ᵣ(wᵣv̂ᵣ) for buffer in advection_buffers[2:end] - @eval @inline weno_reconstruction(scheme::WENO{$buffer}, ψ₀, δ, ω) = @inbounds @muladd $(metaprogrammed_weno_reconstruction(buffer)) + @eval @inline weno_reconstruction(scheme::WENO{$buffer}, ψ₀, δ, ω, coefficients = nothing) = @inbounds @muladd $(metaprogrammed_weno_reconstruction(buffer)) end -# Interpolation functions -for (interp, dir, val) in zip([:xᶠᵃᵃ, :yᵃᶠᵃ, :zᵃᵃᶠ], [:x, :y, :z], [1, 2, 3]) +@inline weno_coefficients(::Nothing, args...) = nothing + +# Interpolation functions: the Center reconstruction at i is the Face reconstruction at i + 1, with the coefficients of the Center at i +for (dir, val) in ((:x, 1), (:y, 2), (:z, 3)), (location, interp, shifted) in ((:faces, (:xᶠᵃᵃ, :yᵃᶠᵃ, :zᵃᵃᶠ)[val], false), + (:centers, (:xᶜᵃᵃ, :yᵃᶜᵃ, :zᵃᵃᶜ)[val], true)) interpolate_func = Symbol(:biased_interpolate_, interp) stencil = Symbol(:weno_stencil_, dir) + index = (:i, :j, :k)[val] + i′, j′, k′ = ntuple(d -> shifted && d == val ? :($((:i, :j, :k)[d]) + 1) : (:i, :j, :k)[d], 3) @eval begin @inline function $interpolate_func(i, j, k, grid, scheme::WENO{N, FT}, bias, ψ, args...) where {N, FT} - S = $stencil(i, j, k, grid, scheme, bias, ψ, args...) + C = weno_coefficients(scheme.coefficients, Val($(QuoteNode(dir))), Val($(QuoteNode(location))), $index, bias) + S = $stencil($i′, $j′, $k′, grid, scheme, bias, ψ, args...) ψ₀ = weno_anchor(scheme, S) δ = weno_differences(scheme, S) - ω = biased_weno_weights(δ, grid, scheme, bias, args...) - return weno_reconstruction(scheme, ψ₀, δ, ω) + ω = biased_weno_weights(δ, grid, scheme, C, bias, args...) + return weno_reconstruction(scheme, ψ₀, δ, ω, C) end @inline function $interpolate_func(i, j, k, grid, scheme::WENO{N, FT}, bias, ψ, VI::AbstractSmoothnessStencil, args...) where {N, FT} - S = $stencil(i, j, k, grid, scheme, bias, ψ, args...) + C = weno_coefficients(scheme.coefficients, Val($(QuoteNode(dir))), Val($(QuoteNode(location))), $index, bias) + S = $stencil($i′, $j′, $k′, grid, scheme, bias, ψ, args...) ψ₀ = weno_anchor(scheme, S) δ = weno_differences(scheme, S) - ω = biased_weno_weights(δ, grid, scheme, bias, VI, args...) - return weno_reconstruction(scheme, ψ₀, δ, ω) + ω = biased_weno_weights(δ, grid, scheme, C, bias, VI, args...) + return weno_reconstruction(scheme, ψ₀, δ, ω, C) end @inline function $interpolate_func(i, j, k, grid, scheme::WENO{N, FT}, bias, ψ, VI::VelocityStencil, u, v, args...) where {N, FT} - S = $stencil(i, j, k, grid, scheme, bias, ψ, u, v, args...) + C = weno_coefficients(scheme.coefficients, Val($(QuoteNode(dir))), Val($(QuoteNode(location))), $index, bias) + S = $stencil($i′, $j′, $k′, grid, scheme, bias, ψ, u, v, args...) ψ₀ = weno_anchor(scheme, S) δ = weno_differences(scheme, S) - ω = biased_weno_weights((i, j, k), grid, scheme, bias, Val($val), VI, u, v) - return weno_reconstruction(scheme, ψ₀, δ, ω) + ω = biased_weno_weights(($i′, $j′, $k′), grid, scheme, C, bias, Val($val), VI, u, v) + return weno_reconstruction(scheme, ψ₀, δ, ω, C) end @inline function $interpolate_func(i, j, k, grid, scheme::WENO{N, FT}, bias, ψ, VI::FunctionStencil, args...) where {N, FT} - S = $stencil(i, j, k, grid, scheme, bias, ψ, args...) + C = weno_coefficients(scheme.coefficients, Val($(QuoteNode(dir))), Val($(QuoteNode(location))), $index, bias) + S = $stencil($i′, $j′, $k′, grid, scheme, bias, ψ, args...) ψ₀ = weno_anchor(scheme, S) δ = weno_differences(scheme, S) - Sₛ = $stencil(i, j, k, grid, scheme, bias, VI.func, args...) - ω = biased_weno_weights(weno_differences(scheme, Sₛ), grid, scheme, bias, VI, args...) - return weno_reconstruction(scheme, ψ₀, δ, ω) + Sₛ = $stencil($i′, $j′, $k′, grid, scheme, bias, VI.func, args...) + ω = biased_weno_weights(weno_differences(scheme, Sₛ), grid, scheme, C, bias, VI, args...) + return weno_reconstruction(scheme, ψ₀, δ, ω, C) end end end diff --git a/src/Advection/weno_reconstruction.jl b/src/Advection/weno_reconstruction.jl index da9c82793c4..193fe2996d4 100644 --- a/src/Advection/weno_reconstruction.jl +++ b/src/Advection/weno_reconstruction.jl @@ -2,16 +2,18 @@ ##### Weighted Essentially Non-Oscillatory (WENO) advection scheme ##### -struct WENO{N, FT, WCT, TD, PP, CA, SI} <: AbstractUpwindBiasedAdvectionScheme{N, FT, TD} +struct WENO{N, FT, WCT, TD, PP, CA, SI, SC} <: AbstractUpwindBiasedAdvectionScheme{N, FT, TD} bounds :: PP buffer_scheme :: CA advecting_velocity_scheme :: SI time_discretization :: TD + coefficients :: SC function WENO{N, FT, WCT}(bounds::PP, buffer_scheme::CA, advecting_velocity_scheme :: SI, - time_discretization :: TD) where {N, FT, WCT, PP, CA, SI, TD} + time_discretization :: TD, + coefficients :: SC = nothing) where {N, FT, WCT, PP, CA, SI, TD, SC} - return new{N, FT, WCT, TD, PP, CA, SI}(bounds, buffer_scheme, advecting_velocity_scheme, time_discretization) + return new{N, FT, WCT, TD, PP, CA, SI, SC}(bounds, buffer_scheme, advecting_velocity_scheme, time_discretization, coefficients) end end @@ -23,7 +25,8 @@ end time_discretization = ExplicitTimeDiscretization(), bounds = nothing, maximum_courant_number = 5//18, - minimum_buffer_upwind_order = 3) + minimum_buffer_upwind_order = 3, + stretched = false) Construct a weighted essentially non-oscillatory advection scheme of order `order` with precision `FT`. @@ -52,6 +55,10 @@ Keyword arguments scheme order reaches this value, subsequent buffers use `Centered(order=2)` instead of continuing to decrease the upwind order. Default: 3. +- `stretched`: The directions, among `:x`, `:y` and `:z`, along which the reconstruction coefficients, the smoothness + indicators and the optimal weights are computed from the cell widths of the grid, built when the grid is + known (see `materialize_advection`): a tuple such as `(:z,)` or `(:x, :z)`, `true` for all three or + `false` for none. Uniform and unselected directions keep the constant coefficients. Default: `false`. Examples ======== @@ -124,7 +131,8 @@ function WENO(FT::DataType=Oceananigans.defaults.FloatType; time_discretization = ExplicitTimeDiscretization(), bounds = nothing, maximum_courant_number = 5//18, - minimum_buffer_upwind_order = 3) + minimum_buffer_upwind_order = 3, + stretched = false) mod(order, 2) == 0 && throw(ArgumentError("WENO reconstruction scheme is defined only for odd orders")) @@ -142,7 +150,7 @@ function WENO(FT::DataType=Oceananigans.defaults.FloatType; # At minimum order, switch to Centered scheme buffer_scheme = Centered(FT; order=2) else - buffer_scheme = WENO(FT; order=order-2, bounds, maximum_courant_number, minimum_buffer_upwind_order, weight_computation) + buffer_scheme = WENO(FT; order=order-2, bounds, maximum_courant_number, minimum_buffer_upwind_order, weight_computation, stretched) end end @@ -150,14 +158,20 @@ function WENO(FT::DataType=Oceananigans.defaults.FloatType; preserved_bounds = isnothing(bounds) ? nothing : BoundsPreservation(convert(FT, bounds[1]), convert(FT, bounds[2]), convert(FT, maximum_courant_number), nothing) - return WENO{N, FT, weight_computation}(preserved_bounds, buffer_scheme, advecting_velocity_scheme, time_discretization) + directions = stretched === true ? (:x, :y, :z) : stretched === false ? () : stretched + all(∈((:x, :y, :z)), directions) || throw(ArgumentError("stretched directions must be among :x, :y and :z, found $directions")) + coefficients = isempty(directions) ? nothing : directions + + return WENO{N, FT, weight_computation}(preserved_bounds, buffer_scheme, advecting_velocity_scheme, time_discretization, coefficients) end end weno_order(::WENO{N}) where N = 2N-1 Base.eltype(::WENO{N, FT}) where {N, FT} = FT -Base.summary(a::WENO{N, FT, WCT, TD, Nothing}) where {N, FT, WCT, TD} = string("WENO{$N, $FT, $WCT}(order=", 2N-1, ")") -Base.summary(a::WENO{N, FT, WCT, TD, PP}) where {N, FT, WCT, TD, PP} = string("WENO{$N, $FT, $WCT}(order=", 2N-1, ", bounds=", string(a.bounds), ")") +Base.summary(a::WENO{N, FT, WCT, TD, Nothing}) where {N, FT, WCT, TD} = + string("WENO{$N, $FT, $WCT}(order=", 2N-1, isnothing(a.coefficients) ? "" : ", stretched", ")") +Base.summary(a::WENO{N, FT, WCT, TD, PP}) where {N, FT, WCT, TD, PP} = + string("WENO{$N, $FT, $WCT}(order=", 2N-1, ", bounds=", string(a.bounds), isnothing(a.coefficients) ? "" : ", stretched", ")") function Base.show(io::IO, a::WENO) print(io, summary(a), '\n') @@ -177,13 +191,15 @@ Adapt.adapt_structure(to, scheme::WENO{N, FT, WCT}) where {N, FT, WCT} = WENO{N, FT, WCT}(Adapt.adapt(to, scheme.bounds), Adapt.adapt(to, scheme.buffer_scheme), Adapt.adapt(to, scheme.advecting_velocity_scheme), - Adapt.adapt(to, scheme.time_discretization)) + Adapt.adapt(to, scheme.time_discretization), + Adapt.adapt(to, scheme.coefficients)) Architectures.on_architecture(to, scheme::WENO{N, FT, WCT}) where {N, FT, WCT} = WENO{N, FT, WCT}(on_architecture(to, scheme.bounds), on_architecture(to, scheme.buffer_scheme), on_architecture(to, scheme.advecting_velocity_scheme), - on_architecture(to, scheme.time_discretization)) + on_architecture(to, scheme.time_discretization), + on_architecture(to, scheme.coefficients)) # Select the default WENO weight computation # Specific backends may override diff --git a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl index 33cd88dcbab..376a10d2483 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl @@ -21,8 +21,8 @@ using Oceananigans.Grids: Grids, AbstractGrid, OrthogonalSphericalShellGrid, Per using Oceananigans.Models: Models, default_free_surface_boundary_conditions, materialize_free_surface using Oceananigans.Operators: Δzᶜᶠᶜ, Δzᶠᶜᶜ using Oceananigans.Simulations: Simulations, timestepper -using Oceananigans.TimeSteppers: TimeSteppers, QuasiAdamsBashforth2TimeStepper, SplitRungeKuttaName, - SplitRungeKuttaTimeStepper, step_lagrangian_particles!, +using Oceananigans.TimeSteppers: TimeSteppers, MultiStageTimeStepper, QuasiAdamsBashforth2TimeStepper, SplitRungeKuttaName, + SSPRungeKuttaTimeStepper, SplitRungeKuttaTimeStepper, step_lagrangian_particles!, update_lagrangian_particle_state! using Oceananigans.TurbulenceClosures: TurbulenceClosures using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: TKEBasedVerticalDiffusivities @@ -200,6 +200,7 @@ include("update_hydrostatic_free_surface_model_state.jl") include("deferred_barotropic_acceleration.jl") include("hydrostatic_free_surface_ab2_step.jl") include("hydrostatic_free_surface_rk_step.jl") +include("hydrostatic_free_surface_ssp_step.jl") include("cache_hydrostatic_free_surface_tendencies.jl") include("prescribed_hydrostatic_velocity_fields.jl") include("single_column_model_mode.jl") diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl index b7d2f75c45e..69c18304b87 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl @@ -1,9 +1,14 @@ module SplitExplicitFreeSurfaces -export SplitExplicitFreeSurface, ForwardBackwardScheme +export SplitExplicitFreeSurface, ForwardBackwardScheme, RungeKutta3Scheme export FixedSubstepNumber, FixedTimeStepSize +export ConstantAveragingKernel, CosineAveragingKernel, LowDissipationAveragingKernel +export SymmetricTrigAveragingKernel, WideTrigAveragingKernel +export OptimizedSymmetricAveragingKernel, OptimizedAsymmetricAveragingKernel +export FrozenSlowForcing, ProgressiveSlowForcing using DocStringExtensions: TYPEDSIGNATURES +using LinearAlgebra: LinearAlgebra using KernelAbstractions: @index, @kernel, @localmem, @synchronize, StaticSize using Oceananigans: Oceananigans @@ -27,6 +32,8 @@ using Oceananigans.Units: Time using Oceananigans.Utils: Utils, @apply_regionally, KernelParameters, add_time_interval, configure_kernel, launch! include("split_explicit_timesteppers.jl") +include("slow_forcing_reconstruction.jl") +include("split_explicit_averaging_kernels.jl") include("split_explicit_free_surface.jl") include("distributed_split_explicit_free_surface.jl") include("initialize_split_explicit_substepping.jl") diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl index 80c1f965895..87eca949caf 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl @@ -17,9 +17,29 @@ @inbounds Gⱽ[i, j, 1] = δV / Δt end -# Note that for AB2, `transport_velocities` holds the value of the prognostic velocities at `tⁿ` +# Note that for AB2 and for an SSP stage, `transport_velocities` holds the value of the prognostic +# velocities at `tⁿ` and at the stage input respectively. @inline baseline_velocities(model, ::SplitRungeKuttaTimeStepper) = model.timestepper.Ψ⁻ @inline baseline_velocities(model, ::QuasiAdamsBashforth2TimeStepper) = model.transport_velocities +@inline baseline_velocities(model, ::SSPRungeKuttaTimeStepper) = model.transport_velocities + +weight_slow_forcing!(GUⁿ, GVⁿ, grid, timestepper, stage) = nothing + +# G★ accumulates βₘ Gᵐ over the stages, and the last stage drives its solve with the full sum. +function weight_slow_forcing!(GUⁿ, GVⁿ, grid, ts::SSPRungeKuttaTimeStepper, stage) + FT = eltype(grid) + β = convert(FT, ssp_quadrature_weights(ts.coefficients)[stage]) + keep = convert(FT, stage > 1) + Sᵁ, Sⱽ = stage == ts.Nstages ? (GUⁿ, GVⁿ) : (ts.G★.U, ts.G★.V) + launch!(architecture(grid), grid, :xy, _accumulate_slow_forcing!, Sᵁ, Sⱽ, ts.G★.U, ts.G★.V, GUⁿ, GVⁿ, β, keep) + return nothing +end + +@kernel function _accumulate_slow_forcing!(Sᵁ, Sⱽ, Ĝᵁ, Ĝⱽ, Gᵁ, Gⱽ, β, keep) + i, j = @index(Global, NTuple) + @inbounds Sᵁ[i, j, 1] = keep * Ĝᵁ[i, j, 1] + β * Gᵁ[i, j, 1] + @inbounds Sⱽ[i, j, 1] = keep * Ĝⱽ[i, j, 1] + β * Gⱽ[i, j, 1] +end function HydrostaticFreeSurfaceModels.compute_free_surface_tendency!(grid, model, free_surface::SplitExplicitFreeSurface, Δt) GUⁿ = model.timestepper.Gⁿ.U @@ -29,9 +49,15 @@ function HydrostaticFreeSurfaceModels.compute_free_surface_tendency!(grid, model baseline = baseline_velocities(model, model.timestepper) u⁻, v⁻ = baseline.u, baseline.v - @apply_regionally launch!(architecture(grid), grid, :xy, _compute_realized_barotropic_tendency!, GUⁿ, GVⁿ, grid, u, v, u⁻, v⁻, Δt) + @apply_regionally begin + launch!(architecture(grid), grid, :xy, _compute_realized_barotropic_tendency!, GUⁿ, GVⁿ, grid, u, v, u⁻, v⁻, Δt) + weight_slow_forcing!(GUⁿ, GVⁿ, grid, model.timestepper, model.clock.stage) + end fill_halo_regions!((GUⁿ, GVⁿ); async=true) + # Stash the first two stage values for the quadratic reconstruction formed on the final stage. + @apply_regionally cache_stage_slow_forcing!(free_surface.slow_forcing, GUⁿ, GVⁿ, model.clock.stage) + return nothing end diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl index f0ebaa8bc00..e23ea9d9ddd 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl @@ -1,5 +1,5 @@ using Oceananigans.Grids: peripheral_node -using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper, SplitRungeKuttaTimeStepper +using Oceananigans.TimeSteppers: MultiStageTimeStepper, QuasiAdamsBashforth2TimeStepper, SplitRungeKuttaTimeStepper using Oceananigans.Utils: Utils # This file contains two different methods performed at different stages of the simulation. @@ -46,7 +46,7 @@ function initialize_free_surface_state!(free_surface, baroclinic_timestepper, ti end # At the last stage we reset the velocities and perform the complete substepping from n to n+1 -function initialize_free_surface_state!(free_surface, baroclinic_ts::SplitRungeKuttaTimeStepper, barotropic_ts) +function initialize_free_surface_state!(free_surface, baroclinic_ts::MultiStageTimeStepper, barotropic_ts) η = free_surface.displacement U, V = free_surface.barotropic_velocities diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/slow_forcing_reconstruction.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/slow_forcing_reconstruction.jl new file mode 100644 index 00000000000..9c19b72ca8f --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/slow_forcing_reconstruction.jl @@ -0,0 +1,203 @@ +using Adapt: Adapt +using StaticArrays: SMatrix +using Oceananigans.Grids: Face, Center +using Oceananigans.TimeSteppers: SplitRungeKuttaTimeStepper, SSPRungeKuttaTimeStepper, + MultiStageTimeStepper, ssp_quadrature_weights, ssp_stage_fractions + +##### +##### Time reconstruction of the slow forcing across the barotropic sub-cycle. +##### +##### Reconstructing the depth-integrated slow forcing in time does not lift the order-two cap of the +##### split-explicit scheme, but it removes the barotropic-baroclinic resonance near μ₀ = k c₀ Δt ≈ 5.7 at +##### which the frozen forcing loses about 40% of its damping. +##### + +""" + struct FrozenSlowForcing + +Hold the slow forcing fixed across the barotropic sub-cycle. This is the classical split-explicit +treatment and the default. +""" +struct FrozenSlowForcing end + +""" + ProgressiveSlowForcing(timestepper = SplitRungeKuttaTimeStepper(stages = 3)) + +Reconstruct the slow forcing *progressively*: each stage carries the best polynomial the already-computed +stage forcings support -- the first stage frozen, the second linear, the third and any later stage +quadratic -- evaluated at the midpoint of each barotropic substep. + +The `timestepper` must match the one that steps the model. It fixes the nodes at which the stage forcings are +sampled and how many of them are stored: the stage in progress reads its forcing live, so `Nstages - 1` fields +per horizontal velocity component. + +The intermediate stages do not change the order but they do change the stability, because the perturbation +the sub-cycle injects into the first baroclinic mode collects the forcing error of every stage. +[`ModifiedRungeKutta4TimeStepper`](@ref) requires `ProgressiveSlowForcing(ModifiedRungeKutta4TimeStepper())`. +""" +struct ProgressiveSlowForcing{U, V, W, F, T} + Gᵁ :: U # the stage forcings preceding the current stage, one field each + Gⱽ :: V + weights :: W # dimensionless reconstruction weights, one entry per stage, all of one type + frozen :: F # whether each stage's polynomial is a constant, and so takes the frozen path + nodes :: T # the composition's sample times +end + +function ProgressiveSlowForcing(timestepper = SplitRungeKuttaTimeStepper(stages = 3)) + M = timestepper.Nstages + M > 1 || throw(ArgumentError("A slow-forcing reconstruction needs at least two stages, got $M.")) + + τ = stage_sample_times(timestepper) + N = M - 1 + + degrees = Tuple(reconstruction_degree(m) for m in 1:M) + + # All stages carry the same weight type so that indexing with the running stage stays type-stable; a + # stage whose polynomial is a constant stores zeros and is flagged to take the frozen path instead. + weights = Tuple(degrees[m] == 1 ? zero_reconstruction_weights(Val(N), Float64) : + dimensionless_reconstruction_weights(τ[1:m], Val(degrees[m]), Val(N), Float64) + for m in 1:M) + + frozen = Tuple(D == 1 for D in degrees) + + return ProgressiveSlowForcing(Tuple(nothing for n in 1:N), Tuple(nothing for n in 1:N), weights, frozen, τ) +end + +@inline cached_stages(sf::ProgressiveSlowForcing) = length(sf.Gᵁ) + +materialize_slow_forcing(::FrozenSlowForcing, grid, args...) = FrozenSlowForcing() + +function materialize_slow_forcing(sf::ProgressiveSlowForcing, grid, u_bcs, v_bcs) + N = cached_stages(sf) + Gᵁ = ntuple(n -> Field{Face, Center, Nothing}(grid, boundary_conditions = u_bcs), N) + Gⱽ = ntuple(n -> Field{Center, Face, Nothing}(grid, boundary_conditions = v_bcs), N) + return ProgressiveSlowForcing(Gᵁ, Gⱽ, sf.weights, sf.frozen, sf.nodes) +end + +Adapt.adapt_structure(to, sf::ProgressiveSlowForcing) = + ProgressiveSlowForcing(Adapt.adapt(to, sf.Gᵁ), + Adapt.adapt(to, sf.Gⱽ), + Adapt.adapt(to, sf.weights), + Adapt.adapt(to, sf.frozen), + Adapt.adapt(to, sf.nodes)) + +# The stage values are rebuilt within every time step. +Oceananigans.prognostic_state(::FrozenSlowForcing) = nothing +Oceananigans.prognostic_state(::ProgressiveSlowForcing) = nothing + +##### +##### The nodes: where in the step each stage forcing was sampled +##### + +""" + stage_sample_times(timestepper) + +Times, as fractions of `Δt` from the start of the step, that the state entering each Runge-Kutta stage +approximates. The low-storage form advances from `Ψⁿ` by `γₘ = Δt / βₘ`, so the state entering stage `m` +approximates `γₘ₋₁` and `β = (3, 2, 1)` gives `(0, 1/3, 1/2)`; the Shu-Osher form sits at `(0, 1, 1/2)`. +""" +@inline stage_sample_times(ts::SplitRungeKuttaTimeStepper{<:NTuple{M, Any}}) where M = (0.0, ntuple(m -> 1 / ts.β[m], Val(M - 1))...) +@inline stage_sample_times(::SSPRungeKuttaTimeStepper) = (0.0, 1.0, 1/2) + +##### +##### The rule: how many coefficients each stage's polynomial carries +##### + +@inline reconstruction_degree(stage) = min(stage, 3) + +##### +##### The weights, built once per stage on the host +##### + +# Square node set: the polynomial passes through every sample. Rectangular: least-squares fit. +@inline invert_vandermonde(V::SMatrix{N, N}) where N = inv(V) +@inline invert_vandermonde(V::SMatrix{M, D}) where {M, D} = inv(V' * V) * V' + +""" + dimensionless_reconstruction_weights(τ, Val(D), Val(N), FT) + +Weights `w` of the polynomial `F(σ) = F₀ + F₁ σ + F₂ σ²` through the slow forcing sampled at the `M` +dimensionless nodes `τ`, with `σ = s/Δt`, laid out in the slots the substep kernel reads: + + Fₖ₋₁ = w[k][1] G¹ + … + w[k][N] Gᴺ + w[k][N+1] Gᴹ + +The first `N` slots are the cached stage forcings and the last is the live forcing of the current stage, which +is always the final sample. Slots beyond the samples the stage has carry a zero weight, so the tuple length is +fixed by the composition rather than by the stage. + +The nodes are `τⱼ Δt`, so the fit factors into a shape fixed by the composition and a scaling `Δt^{-(k-1)}` +on the `k`-th coefficient; the shape is built once and rescaled at each stage. +""" +function dimensionless_reconstruction_weights(τ::NTuple{M, Any}, ::Val{D}, ::Val{N}, ::Type{FT}) where {M, D, N, FT} + V = SMatrix{M, D, FT}(ntuple(n -> FT(τ[mod1(n, M)])^(cld(n, M) - 1), Val(M * D))) + X = invert_vandermonde(V) # D × M + + # Sample M is the live forcing rather than a cache, so it lands in the last slot whatever the stage. + @inline weight(k, j) = j == N + 1 ? X[k, M] : (j < M ? X[k, j] : zero(FT)) + + return ntuple(Val(3)) do k + k > D ? ntuple(j -> zero(FT), Val(N + 1)) : ntuple(j -> FT(weight(k, j)), Val(N + 1)) + end +end + +""" + scale_reconstruction_weights(w, Δt) + +Restore the units of the stored dimensionless weights: the `k`-th coefficient of `F(s) = F₀ + F₁ s + F₂ s²` +carries `1/time^(k-1)`, so `Δt` enters only through that power. +""" +@inline scale_reconstruction_weights(w::NTuple{3, Any}, Δt) = ntuple(k -> map(x -> x / Δt^(k - 1), w[k]), Val(3)) + +##### +##### Evaluation inside the substep kernel +##### + +# `s` is the midpoint of the current barotropic substep, measured from tⁿ: sampling at either endpoint +# injects an O(Δτ) error and costs two orders. +@inline slow_forcing(i, j, Gᴴ, ::Nothing, w, s) = @inbounds Gᴴ[i, j, 1] + +@inline function slow_forcing(i, j, Gᴴ, Gᶜ::NTuple{N}, w, s) where N + @inbounds begin + Gᴹ = Gᴴ[i, j, 1] + F₀ = reconstruction_sum(w[1], Gᶜ, Gᴹ, i, j, Val(N)) + F₁ = reconstruction_sum(w[2], Gᶜ, Gᴹ, i, j, Val(N)) + F₂ = reconstruction_sum(w[3], Gᶜ, Gᴹ, i, j, Val(N)) + end + + return F₀ + F₁ * s + F₂ * s^2 +end + +@inline reconstruction_sum(w, Gᶜ, Gᴹ, i, j, ::Val{N}) where N = + @inbounds w[N+1] * Gᴹ + sum(ntuple(n -> w[n] * Gᶜ[n][i, j, 1], Val(N))) + +##### +##### Capturing the stage values, and selecting when the reconstruction is active +##### + +@inline cache_stage_slow_forcing!(::FrozenSlowForcing, args...) = nothing + +@inline function cache_stage_slow_forcing!(sf::ProgressiveSlowForcing, GUⁿ, GVⁿ, stage) + if stage ≤ cached_stages(sf) + parent(sf.Gᵁ[stage]) .= parent(GUⁿ) + parent(sf.Gⱽ[stage]) .= parent(GVⁿ) + end + return nothing +end + +@inline zero_reconstruction_weights(::Val{N}, ::Type{FT}) where {N, FT} = + ntuple(k -> ntuple(j -> zero(FT), Val(N + 1)), Val(3)) + +@inline reconstruction_timestep(::SSPRungeKuttaTimeStepper, stage, Δτ) = Δτ +@inline reconstruction_timestep(ts::SplitRungeKuttaTimeStepper, stage, Δτ) = @inbounds Δτ * ts.β[stage] + +# `nothing` caches specialize the substep kernel on the frozen case, keeping the plain `Gᴴ[i, j, 1]` load. +@inline stage_reconstruction(::FrozenSlowForcing, timestepper, stage, Δτ::FT) where FT = (nothing, nothing, zero_reconstruction_weights(Val(1), FT)) + +# Anything that is not a multi-stage scheme has no stage values to reconstruct from. +@inline stage_reconstruction(sf::ProgressiveSlowForcing, timestepper, stage, Δτ::FT) where FT = (nothing, nothing, zero_reconstruction_weights(Val(length(sf.Gᵁ)), FT)) + +@inline function stage_reconstruction(sf::ProgressiveSlowForcing, ts::MultiStageTimeStepper, stage, Δτ::FT) where FT + @inbounds sf.frozen[stage] && return (nothing, nothing, zero_reconstruction_weights(Val(length(sf.Gᵁ)), FT)) + @inbounds w̃ = sf.weights[stage] + return (sf.Gᵁ, sf.Gⱽ, scale_reconstruction_weights(w̃, reconstruction_timestep(ts, stage, Δτ))) +end diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_averaging_kernels.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_averaging_kernels.jl new file mode 100644 index 00000000000..ea6a8e89e60 --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_averaging_kernels.jl @@ -0,0 +1,197 @@ +#### +#### Averaging kernels for the SplitExplicitFreeSurface +#### +# +# An averaging kernel is a function of τ, the barotropic substep normalized by the baroclinic step. +# The moments of the resulting weights wₘ set the order of the filtered barotropic mode: +# +# μ₀ = ∑ wₘ = 1 and μ₁ = ∑ wₘ τₘ = 1 (conservation and consistency, always enforced) +# μ₂ = ∑ wₘ (τₘ - 1)² (diffusion) +# μ₃ = ∑ wₘ (τₘ - 1)³ (dispersion) +# +# μ₂ = 0 gives second order and μ₂ = μ₃ = 0 third order, realized only if the barotropic substep is at +# least as accurate (`RungeKutta3Scheme`). Kernels with μ₂ = 0 are less stable at strong stratification. +# +# - `nothing`: no averaging, only for simple tests +# - `ConstantAveragingKernel()`: μ₂ ≫ 0 +# - `CosineAveragingKernel()`: μ₂ > 0, μ₃ = 0 +# - `averaging_shape_function`: μ₂ > 0, Shchepetkin and McWilliams (2005), the default +# - `LowDissipationAveragingKernel()`: μ₂ = 0, the same shape with `r` retuned +# - `SymmetricTrigAveragingKernel()`: μ₂ = μ₃ = 0, narrowest window +# - `WideTrigAveragingKernel()`: μ₂ = μ₃ = 0, wider window, needs substeps % 8 == 0 +# - `OptimizedSymmetricAveragingKernel()`: μ₂ = μ₃ = 0, minimax shape, stratification robust +# - `OptimizedAsymmetricAveragingKernel()`: μ₂ = μ₃ = 0, minimax shape, lowest total variation + +struct ConstantAveragingKernel <: Function end +struct CosineAveragingKernel <: Function end +struct LowDissipationAveragingKernel <: Function end +struct SymmetricTrigAveragingKernel <: Function end +struct WideTrigAveragingKernel <: Function end +struct OptimizedSymmetricAveragingKernel <: Function end +struct OptimizedAsymmetricAveragingKernel <: Function end + +##### +##### From a shape function to the substep weights +##### + +# Sample `shape` on the substep grid, keep the leading weights that `retain` accepts, normalize them, and +# rescale the substep so that μ₁ = 1. +function trimmed_weights(FT, substeps, shape, retain = !=(0)) + τᶠ = range(FT(0), FT(2), length = substeps+1) + raw = map(shape, τᶠ[2:end]) + M★ = something(findlast(retain, raw), firstindex(raw)) + + w = collect(raw[1:M★]) + w ./= sum(w) + + Δτ = τᶠ[2] - τᶠ[1] + return FT(Δτ / (sum(w .* (1:M★)) * Δτ)), w +end + +@inline kernel_moment(Δτ, w, p) = sum(w[m] * (m * Δτ - 1)^p for m in eachindex(w)) + +# The transport weights are the reversed cumulative sum of the averaging weights. +function weights_from_shape(FT, substeps, shape, retain = !=(0)) + Δτ, w = trimmed_weights(FT, substeps, shape, retain) + t = [Δτ * sum(@view w[m:end]) for m in eachindex(w)] + return Δτ, FT.(w), FT.(t) +end + +# Bisect the leading amplitude of `shape(a)` so that μ₂ vanishes; the higher amplitudes are fixed. +function vanishing_second_moment(FT, substeps, shape, lo, hi, retain = !=(0)) + μ₂(a) = kernel_moment(trimmed_weights(FT, substeps, shape(a), retain)..., 2) + + flo = μ₂(lo) + for _ in 1:80 + mid = (lo + hi) / 2 + fmid = μ₂(mid) + if flo * fmid ≤ 0 + hi = mid + else + lo, flo = mid, fmid + end + end + + return (lo + hi) / 2 +end + +# A user-supplied kernel is trimmed where it turns negative, rather than at the end of its support. +weights_from_substeps(FT, substeps, averaging_kernel) = weights_from_shape(FT, substeps, averaging_kernel, >(0)) + +# Without an averaging kernel we take the endpoint. +@inline function weights_from_substeps(FT, substeps, ::Nothing) + fractional_step = one(FT) / substeps + averaging_weights = FT[m == substeps ? one(FT) : zero(FT) for m in 1:substeps] + transport_weights = FT[one(FT) / substeps for _ in 1:substeps] + return fractional_step, averaging_weights, transport_weights +end + +##### +##### Shchepetkin and McWilliams (2005), and its μ₂ = 0 retuning +##### + +@inline (::CosineAveragingKernel)(τ::FT) where FT = τ ≥ 0.5 && τ ≤ 1.5 ? convert(FT, 1 + cos(2π * (τ - 1))) : zero(FT) +@inline (::ConstantAveragingKernel)(τ::FT) where FT = convert(FT, 1) + +# (p = 2, q = 4) minimize dispersion error from Shchepetkin and McWilliams (2005): https://doi.org/10.1016/j.ocemod.2004.08.002 +@inline function averaging_shape_function(τ::FT; p = 2, q = 4, r = FT(0.18927)) where FT + τ₀ = (p + 2) * (p + q + 2) / (p + 1) / (p + q + 1) + return (τ / τ₀)^p * (1 - (τ / τ₀)^q) - r * (τ / τ₀) +end + +function weights_from_substeps(FT, substeps, ::LowDissipationAveragingKernel) + shape(r) = τ -> averaging_shape_function(τ; r = convert(FT, r)) + μ₂(r) = kernel_moment(trimmed_weights(FT, substeps, shape(r), >(0))..., 2) + + rₗ, rₕ = FT(0.18927), FT(0.9) + r = μ₂(rₗ) * μ₂(rₕ) > 0 ? FT(0.285) : vanishing_second_moment(FT, substeps, shape, rₗ, rₕ, >(0)) + + return weights_from_shape(FT, substeps, shape(r), >(0)) +end + +##### +##### Trigonometric μ₂ = μ₃ = 0 windows +##### + +# Both owe μ₃ = 0 to being symmetric about their own centre, and that symmetry survives the discretisation +# only when the window edges land on the substep grid τₖ = 2k/substeps. The symmetric window is centred on +# τ = 1, a grid point for any even `substeps`; the wide one is centred on τ = 9/8 and is not. +required_substep_multiple(averaging_kernel) = 1 +required_substep_multiple(::WideTrigAveragingKernel) = 8 + +@inline symmetric_trig_shape(τ::FT, a1, a2, a3) where FT = + ifelse((τ ≥ FT(1//2)) & (τ ≤ FT(3//2)), 1 + a1 * cospi(2*(τ-1)) + a2 * cospi(4*(τ-1)) + a3 * cospi(6*(τ-1)), zero(FT)) + +function weights_from_substeps(FT, substeps, ::SymmetricTrigAveragingKernel) + a2, a3 = FT(-27//20), FT(1//2) + shape(a1) = τ -> symmetric_trig_shape(τ, FT(a1), a2, a3) + a1 = vanishing_second_moment(FT, substeps, shape, FT(1//5), FT(16//5)) + return weights_from_shape(FT, substeps, shape(a1)) +end + +const wide_trig_window_end = 7//4 + +@inline function wide_trig_shape(τ::FT, a) where FT + c = FT((1//2 + wide_trig_window_end) / 2) + L = FT(wide_trig_window_end - 1//2) + s = one(FT) + for k in eachindex(a) + s += a[k] * cospi(2k * (τ - c) / L) + end + return ifelse((τ ≥ FT(1//2)) & (τ ≤ FT(wide_trig_window_end)), s, zero(FT)) +end + +function weights_from_substeps(FT, substeps, ::WideTrigAveragingKernel) + higher = (FT(-1), FT(1//2), FT(-1), FT(1//2)) + shape(a1) = τ -> wide_trig_shape(τ, (FT(a1), higher...)) + a1 = vanishing_second_moment(FT, substeps, shape, FT(-3), FT(4)) + return weights_from_shape(FT, substeps, shape(a1)) +end + +##### +##### Minimax-optimized μ₂ = μ₃ = 0 kernels +##### + +# shape(s) = 1 + ∑ₖ aₖ cos(πks/R) on |s| ≤ R, with s = τ - 1. a₅ is fixed by shape(±R) = 0 so the shape +# tapers to zero rather than being truncated, and a₁ is bisected so that μ₂ = 0. +const optimized_symmetric_radius = 0.752661 +const optimized_symmetric_higher_amplitudes = (-1.289381, -0.674748, 0.605401) + +@inline function optimized_symmetric_amplitudes(a1, higher) + a = (a1, higher...) + return (a..., 1 + sum(a[k] * (-1)^k for k in 1:4)) +end + +@inline function optimized_symmetric_shape(τ::FT, R, a) where FT + s = τ - 1 + v = one(FT) + for k in eachindex(a) + v += a[k] * cospi(k * s / R) + end + return ifelse(abs(s) ≤ FT(R), v, zero(FT)) +end + +function weights_from_substeps(FT, substeps, ::OptimizedSymmetricAveragingKernel) + R = FT(optimized_symmetric_radius) + shape(a1) = τ -> optimized_symmetric_shape(τ, R, map(FT, optimized_symmetric_amplitudes(a1, optimized_symmetric_higher_amplitudes))) + a1 = vanishing_second_moment(FT, substeps, shape, FT(-4), FT(4)) + return weights_from_shape(FT, substeps, shape(a1)) +end + +# v(τ) = (1 - u²)^p ∑_{k=0}^{4} bₖ uᵏ with u = (τ - c)/R. The moments are linear in b, so μ₀ = 1 and +# μ₁ = μ₂ = μ₃ = 0 are four equations on the five bₖ, solved on whatever substep grid is supplied: the +# moments are then exact for every `substeps`, with no grid-alignment condition. +const optimized_asymmetric_radius = 0.864972 +const optimized_asymmetric_center = 0.909165 +const optimized_asymmetric_power = 0.595115 + +function weights_from_substeps(FT, substeps, ::OptimizedAsymmetricAveragingKernel) + R, c, p = FT(optimized_asymmetric_radius), FT(optimized_asymmetric_center), FT(optimized_asymmetric_power) + basis(τ, k) = (u = (τ - c) / R; abs(u) ≤ 1 ? (1 - u^2)^p * u^k : zero(FT)) + + τᶠ = collect(range(FT(0), FT(2), length = substeps+1)[2:end]) + A = FT[sum(basis.(τᶠ, k) .* (τᶠ .- 1).^(j-1)) for j in 1:4, k in 0:4] + b = A \ FT[1, 0, 0, 0] + + return weights_from_shape(FT, substeps, τ -> sum(b[k+1] * basis(τ, k) for k in 0:4)) +end diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl index dd0fdb3aae7..17c0469dc2a 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl @@ -39,7 +39,7 @@ function substep_halo_filling(extend_halos::Bool, TX, TY, bcs) return open_boundaries ? LocalHaloFilling() : ExtendedHalos() end -struct SplitExplicitFreeSurface{E, H, U, M, FT, K, S, T, B} <: AbstractFreeSurface{H, FT} +struct SplitExplicitFreeSurface{E, H, U, M, FT, K, S, T, W, B} <: AbstractFreeSurface{H, FT} displacement :: H barotropic_velocities :: U # A namedtuple with U, V filtered_state :: M # A namedtuple with η, U, V averaged throughout the substepping @@ -47,10 +47,11 @@ struct SplitExplicitFreeSurface{E, H, U, M, FT, K, S, T, B} <: AbstractFreeSurfa kernel_parameters :: K substepping :: S # Either `FixedSubstepNumber` or `FixedTimeStepSize` timestepper :: T # Contains all auxiliary field and settings necessary to the particular timestepping + slow_forcing :: W # How the slow forcing is represented across the barotropic sub-cycle boundary_transport :: B # Target transports of the sides with a targeted `GravityWaveRadiation` condition (or `nothing`) - function SplitExplicitFreeSurface{E}(η::H, u::U, m::M, g::FT, k::K, s::S, t::T, b::B) where {E, H, U, M, FT, K, S, T, B} - return new{E, H, U, M, FT, K, S, T, B}(η, u, m, g, k, s, t, b) + function SplitExplicitFreeSurface{E}(η::H, u::U, m::M, g::FT, k::K, s::S, t::T, w::W, b::B) where {E, H, U, M, FT, K, S, T, W, B} + return new{E, H, U, M, FT, K, S, T, W, B}(η, u, m, g, k, s, t, w, b) end end @@ -156,7 +157,8 @@ Keyword Arguments to the two consecutive baroclinic timesteps between which the barotropic timestepping occurs and ``M_*`` corresponds to the last barotropic time step for which the `averaging_kernel > 0`. By default, the averaging kernel described by - [Shchepetkin and McWilliams (2005)](@cite Shchepetkin2005) is used. + [Shchepetkin and McWilliams (2005)](@cite Shchepetkin2005) is used; see + `split_explicit_averaging_kernels.jl` for the higher-order alternatives. - `timestepper`: Time stepping scheme used for the barotropic advancement. Only one supported: * `ForwardBackwardScheme()` (default): `η = f(U)` then `U = f(η)`, @@ -173,7 +175,8 @@ function SplitExplicitFreeSurface(grid = nothing; fixed_Δt = nothing, extend_halos = true, averaging_kernel = averaging_shape_function, - timestepper = ForwardBackwardScheme()) + timestepper = ForwardBackwardScheme(), + slow_forcing = FrozenSlowForcing()) if !isnothing(grid) FT = eltype(grid) @@ -198,19 +201,22 @@ function SplitExplicitFreeSurface(grid = nothing; nothing, substepping, timestepper, + slow_forcing, nothing) end # A free surface where halos are explicitly filled at each substep const FillHaloSplitExplicit = SplitExplicitFreeSurface{<:Union{LocalHaloFilling, CompleteHaloFilling}} -# Substeps between two halo exchanges: a halo of H cells in a Connected direction holds H - 2 substeps +# Substeps between two halo exchanges: a halo of H cells in a Connected direction holds (H - 2) ÷ stages substeps @inline halo_exchange_interval(::SplitExplicitFreeSurface{CompleteHaloFilling}, grid, Nsubsteps) = 1 function halo_exchange_interval(free_surface, grid, Nsubsteps) TX, TY, _ = topology(grid) Hx, Hy, _ = halo_size(grid) - return min(TX() isa ConnectedTopology ? Hx - 2 : Nsubsteps, TY() isa ConnectedTopology ? Hy - 2 : Nsubsteps) + stages = stages_per_substep(free_surface.timestepper) + return min(TX() isa ConnectedTopology ? (Hx - 2) ÷ stages : Nsubsteps, + TY() isa ConnectedTopology ? (Hy - 2) ÷ stages : Nsubsteps) end # Extended halos start the substeps filled, while `CompleteHaloFilling` exchanges before every substep @@ -234,6 +240,11 @@ end function split_explicit_substepping(cfl, ::Nothing, fixed_Δt, grid, averaging_kernel, gravitational_acceleration) substepping = split_explicit_substepping(cfl, nothing, nothing, grid, averaging_kernel, gravitational_acceleration) substeps = ceil(Int, 2 * fixed_Δt / substepping.Δt_barotropic) + + # Rounding up only shortens the substep, so the requested `cfl` is still met. + multiple = required_substep_multiple(averaging_kernel) + substeps = multiple * cld(substeps, multiple) + substepping = split_explicit_substepping(nothing, substeps, nothing, grid, averaging_kernel, gravitational_acceleration) return substepping end @@ -274,10 +285,12 @@ function HydrostaticFreeSurfaceModels.materialize_free_surface(free_surface::Spl strategy = substep_halo_filling(extend_halos, TX, TY, bcs) + stages = stages_per_substep(free_surface.timestepper) + maybe_extended_grid = if strategy isa CompleteHaloFilling grid else - maybe_extend_halos(TX, TY, grid, substepping) + maybe_extend_halos(TX, TY, grid, substepping; stages) end # The barotropic fields live on `maybe_extended_grid`, whose halos size the exterior states of their open boundaries @@ -306,7 +319,9 @@ function HydrostaticFreeSurfaceModels.materialize_free_surface(free_surface::Spl materialize_kernel_parameters(TX, TY, maybe_extended_grid, substepping, strategy) end - timestepper = materialize_timestepper(free_surface.timestepper, maybe_extended_grid, free_surface, velocities, bcs.U, bcs.V) + timestepper = materialize_timestepper(free_surface.timestepper, maybe_extended_grid, free_surface, velocities, bcs) + + slow_forcing = materialize_slow_forcing(free_surface.slow_forcing, maybe_extended_grid, bcs.U, bcs.V) return SplitExplicitFreeSurface{typeof(strategy)}(η, barotropic_velocities, @@ -315,6 +330,7 @@ function HydrostaticFreeSurfaceModels.materialize_free_surface(free_surface::Spl kernel_parameters, substepping, timestepper, + slow_forcing, boundary_transport) end @@ -349,15 +365,6 @@ end @inline companion_at_gravity_wave_side(::Nothing, side, companion) = nothing @inline companion_at_gravity_wave_side(bcs, side, companion) = gravity_wave_boundary_condition(getproperty(bcs, side)) ? companion : nothing -# (p = 2, q = 4, r = 0.18927) minimize dispersion error from Shchepetkin and McWilliams (2005): https://doi.org/10.1016/j.ocemod.2004.08.002 -@inline function averaging_shape_function(τ::FT; p = 2, q = 4, r = FT(0.18927)) where FT - τ₀ = (p + 2) * (p + q + 2) / (p + 1) / (p + q + 1) - return (τ / τ₀)^p * (1 - (τ / τ₀)^q) - r * (τ / τ₀) -end - -@inline cosine_averaging_kernel(τ::FT) where FT = τ ≥ 0.5 && τ ≤ 1.5 ? convert(FT, 1 + cos(2π * (τ - 1))) : zero(FT) -@inline constant_averaging_kernel(τ::FT) where FT = convert(FT, 1) - """ An internal type for the `SplitExplicitFreeSurface` that allows substepping with a fixed `Δt_barotropic` based on a CFL condition """ struct FixedTimeStepSize{B, F} @@ -373,6 +380,8 @@ struct FixedSubstepNumber{B, F} transport_weights :: F end +Adapt.adapt_structure(to, s::FixedSubstepNumber) = FixedSubstepNumber(s.fractional_step_size, nothing, nothing) + function FixedTimeStepSize(grid; cfl = 0.7, averaging_kernel = averaging_shape_function, @@ -391,26 +400,6 @@ function FixedTimeStepSize(grid; return FixedTimeStepSize(Δt_barotropic, averaging_kernel) end -@inline function weights_from_substeps(FT, substeps, averaging_kernel) - τᶠ = range(FT(0), FT(2), length = substeps+1) - Δτ = τᶠ[2] - τᶠ[1] - - averaging_weights = map(averaging_kernel, τᶠ[2:end]) - # Find the latest allowable weight - M★ = something(findlast(>(0), averaging_weights), firstindex(averaging_weights)) - - trimmed_weights = averaging_weights[1:M★] - trimmed_weights ./= sum(trimmed_weights) - - # Rescale the substep size so the trimmed weights' first moment lands exactly on the baroclinic step - barycenter = sum(trimmed_weights .* (1:M★)) * Δτ - Δτ = Δτ / barycenter - - transport_weights = [sum(trimmed_weights[i:M★]) for i in 1:M★] .* Δτ - - return FT(Δτ), map(FT, tuple(trimmed_weights...)), map(FT, tuple(transport_weights...)) -end - Base.summary(s::FixedTimeStepSize) = string("FixedTimeStepSize($(prettytime(s.Δt_barotropic)))") Base.summary(s::FixedSubstepNumber) = string("FixedSubstepNumber($(length(s.averaging_weights)))") @@ -423,27 +412,29 @@ Base.show(io::IO, sefs::SplitExplicitFreeSurface) = print(io, "$(summary(sefs))\ ##### # Extending halos is not allowed with variable time-stepping -maybe_extend_halos(TX, TY, grid, ::FixedTimeStepSize) = grid +maybe_extend_halos(TX, TY, grid, ::FixedTimeStepSize; stages = 1) = grid -function maybe_extend_halos(TX, TY, grid, substepping::FixedSubstepNumber) +# A substep erodes one cell of halo per stage. +function maybe_extend_halos(TX, TY, grid, substepping::FixedSubstepNumber; stages = 1) old_halos = halo_size(grid) Nsubsteps = length(substepping.averaging_weights) + step_halo = stages * Nsubsteps + 2 # Every rank shares the same halo, so it is capped at the smallest local size Nx = all_reduce(min, size(grid, 1), architecture(grid)) Ny = all_reduce(min, size(grid, 2), architecture(grid)) - Hx = TX() isa ConnectedTopology ? min(max(Nsubsteps+2, old_halos[1]), Nx) : old_halos[1] - Hy = TY() isa ConnectedTopology ? min(max(Nsubsteps+2, old_halos[2]), Ny) : old_halos[2] + Hx = TX() isa ConnectedTopology ? min(max(step_halo, old_halos[1]), Nx) : old_halos[1] + Hy = TY() isa ConnectedTopology ? min(max(step_halo, old_halos[2]), Ny) : old_halos[2] new_halos = (Hx, Hy, old_halos[3]) - x_capped = TX() isa ConnectedTopology && Hx < Nsubsteps + 2 - y_capped = TY() isa ConnectedTopology && Hy < Nsubsteps + 2 + x_capped = TX() isa ConnectedTopology && Hx < step_halo + y_capped = TY() isa ConnectedTopology && Hy < step_halo if x_capped || y_capped - exchange_interval = min(TX() isa ConnectedTopology ? Hx - 2 : Nsubsteps, - TY() isa ConnectedTopology ? Hy - 2 : Nsubsteps) + exchange_interval = min(TX() isa ConnectedTopology ? (Hx - 2) ÷ stages : Nsubsteps, + TY() isa ConnectedTopology ? (Hy - 2) ÷ stages : Nsubsteps) Nexchanges = cld(Nsubsteps, exchange_interval) - 1 @root @info "SplitExplicitFreeSurface: the extended halos $((Hx, Hy)) are capped by the smallest local grid size $((Nx, Ny)), " * "so the $Nsubsteps barotropic substeps fill the halos $Nexchanges more time(s) per time step." maxlog=1 @@ -480,6 +471,7 @@ Adapt.adapt_structure(to, free_surface::SplitExplicitFreeSurface{extend_halos}) nothing, Adapt.adapt(to, free_surface.substepping), Adapt.adapt(to, free_surface.timestepper), + Adapt.adapt(to, free_surface.slow_forcing), Adapt.adapt(to, free_surface.boundary_transport)) for Type in (SplitExplicitFreeSurface, @@ -501,7 +493,8 @@ end function Oceananigans.prognostic_state(fs::SplitExplicitFreeSurface) return (displacement = prognostic_state(fs.displacement), barotropic_velocities = prognostic_state(fs.barotropic_velocities), - timestepper = prognostic_state(fs.timestepper)) + timestepper = prognostic_state(fs.timestepper), + slow_forcing = prognostic_state(fs.slow_forcing)) end function Oceananigans.restore_prognostic_state!(restored::SplitExplicitFreeSurface, from) diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl index 79e243af2fc..d8f7eebd680 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl @@ -21,27 +21,89 @@ struct ForwardBackwardScheme end materialize_timestepper(::ForwardBackwardScheme, grid, args...) = ForwardBackwardScheme() ##### -##### Timestepper extrapolations and utils +##### Multi-stage barotropic substep integrators ##### -function materialize_timestepper(name::Symbol, args...) - fullname = Symbol(name, :Scheme) - TS = getglobal(@__MODULE__, fullname) - return materialize_timestepper(TS, args...) +""" + struct RungeKutta3Scheme + +Low-storage three-stage Runge-Kutta (1/3, 1/2, 1) substep integrator for the split-explicit barotropic mode. +Within each substep the fast oscillator is advanced from the substep-start state (ηⁿ, Uⁿ, Vⁿ): +```math +U^{(m)} = U^n + γ_m (- g H ∂ η^{(m-1)} + G), \\qquad η^{(m)} = η^n - γ_m ∇ ⋅ U^{(m-1)}, +``` +with ``γ = (Δτ/3, Δτ/2, Δτ)`` and both stage right-hand sides evaluated at the previous stage. +""" +struct RungeKutta3Scheme{H, U, V} + η¹ :: H + U¹ :: U + V¹ :: V + η² :: H + U² :: U + V² :: V +end + +RungeKutta3Scheme() = RungeKutta3Scheme(nothing, nothing, nothing, nothing, nothing, nothing) + +# Cells of halo eroded per barotropic substep when halos are not filled in between. +@inline stages_per_substep(::ForwardBackwardScheme) = 1 +@inline stages_per_substep(::RungeKutta3Scheme) = 3 + +function materialize_timestepper(::RungeKutta3Scheme, grid, free_surface, velocities, bcs) + η_bcs = get(bcs, :η, nothing) + η¹ = free_surface_displacement_field(velocities, free_surface, grid; boundary_conditions = η_bcs) + η² = free_surface_displacement_field(velocities, free_surface, grid; boundary_conditions = η_bcs) + U¹ = Field{Face, Center, Nothing}(grid, boundary_conditions = bcs.U) + V¹ = Field{Center, Face, Nothing}(grid, boundary_conditions = bcs.V) + U² = Field{Face, Center, Nothing}(grid, boundary_conditions = bcs.U) + V² = Field{Center, Face, Nothing}(grid, boundary_conditions = bcs.V) + return RungeKutta3Scheme(η¹, U¹, V¹, η², U², V²) end +@inline stage_parameters(::RungeKutta3Scheme, Δτ::FT) where FT = (Δτ / 3, Δτ / 2, Δτ) + +##### +##### Timestepper extrapolations and utils +##### + +# A substep integrator may be named rather than instantiated, as `timestepper = :ForwardBackward`. +@inline barotropic_timestepper(name::Symbol) = getglobal(@__MODULE__, Symbol(name, :Scheme))() + +materialize_timestepper(name::Symbol, args...) = materialize_timestepper(barotropic_timestepper(name), args...) + +@inline stages_per_substep(name::Symbol) = stages_per_substep(barotropic_timestepper(name)) + initialize_free_surface_timestepper!(::ForwardBackwardScheme, args...) = nothing +initialize_free_surface_timestepper!(::RungeKutta3Scheme, args...) = nothing -# The functions `η★` `U★` and `V★` represent the value of free surface, barotropic zonal and meridional velocity at time step m+1/2 +# `η★` enters the velocity update and `U★`/`V★` the free-surface update; a substep integrator carrying +# history extrapolates instead of returning the current field. @inline U★(i, j, k, grid, ::ForwardBackwardScheme, Uᵐ) = @inbounds Uᵐ[i, j, k] +@inline V★(i, j, k, grid, ::ForwardBackwardScheme, Vᵐ) = @inbounds Vᵐ[i, j, k] @inline η★(i, j, k, grid, ::ForwardBackwardScheme, ηᵐ⁺¹) = @inbounds ηᵐ⁺¹[i, j, k] @inline cache_previous_free_surface!(::ForwardBackwardScheme, i, j, k, η) = nothing @inline cache_previous_velocities!(::ForwardBackwardScheme, i, j, k, U, V) = nothing +##### +##### Adapt +##### + +Adapt.adapt_structure(to, ts::RungeKutta3Scheme) = + RungeKutta3Scheme(Adapt.adapt(to, ts.η¹), + Adapt.adapt(to, ts.U¹), + Adapt.adapt(to, ts.V¹), + Adapt.adapt(to, ts.η²), + Adapt.adapt(to, ts.U²), + Adapt.adapt(to, ts.V²)) + ##### ##### Checkpointing ##### Oceananigans.prognostic_state(::ForwardBackwardScheme) = nothing Oceananigans.restore_prognostic_state!(restored::ForwardBackwardScheme, ::Nothing) = restored + +# The Runge-Kutta scratch is rebuilt at every substep. +Oceananigans.prognostic_state(::RungeKutta3Scheme) = nothing +Oceananigans.restore_prognostic_state!(restored::RungeKutta3Scheme, ::Nothing) = restored diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl index 05dbfc70a03..316cab11be5 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl @@ -1,5 +1,3 @@ -using KernelAbstractions.Extras.LoopInfo: @unroll - using Oceananigans: fields using Oceananigans.DistributedComputations: maybe_distributed_fill_halo_regions! @@ -40,8 +38,9 @@ end # ∂t(U) = - gH∇η + f # # The free surface field η and its average η̄ are located on `Face`s at the surface (grid.Nz +1). All other intermediate -# variables (U, V, Ū, V̄) are barotropic fields (`ReducedField`) for which a k index is not defined. -@kernel function _split_explicit_barotropic_velocity!(transport_weight, grid, filled_halos, Δτ, η, U, V, Gᵁ, Gⱽ, g, Ũ, Ṽ, timestepper) +# variables (U, V, Ū, V̄) are barotropic fields (`ReducedField`) for which a k index is not defined. +@kernel function _split_explicit_barotropic_velocity!(averaging_weight, sᵐ, grid, filled_halos, Δτ, η, U, V, + Gᵁ, Gⱽ, Gᵁᶜ, Gⱽᶜ, w, g, U̅, V̅, timestepper) i, j = @index(Global, NTuple) k_top = grid.Nz+1 @@ -52,20 +51,26 @@ end ∂xᵣ = x_derivative_operator(filled_halos) ∂yᵣ = y_derivative_operator(filled_halos) - # ∂τ(U) = - ∇η + G - # Note: use ∂xᵣT and ∂yᵣT (derivatives at constant r) for the free surface, - # since η lives on the surface and doesn't have vertical structure + # With the frozen forcing this dispatches back to a plain `Gᵁ[i, j, 1]` load. + Gᵁˢ = slow_forcing(i, j, Gᵁ, Gᵁᶜ, w, sᵐ) + Gⱽˢ = slow_forcing(i, j, Gⱽ, Gⱽᶜ, w, sᵐ) + + # ∂τ(U) = - ∇η + G, using ∂xᵣT/∂yᵣT (derivatives at constant r) since η lives on the surface and has + # no vertical structure. @inbounds begin - U[i, j, 1] += Δτ * (- g * Hᶠᶜ * ∂xᵣ(i, j, k_top, grid, η★, timestepper, η) + Gᵁ[i, j, 1]) - V[i, j, 1] += Δτ * (- g * Hᶜᶠ * ∂yᵣ(i, j, k_top, grid, η★, timestepper, η) + Gⱽ[i, j, 1]) + Uᵐ⁺¹ = U[i, j, 1] + Δτ * (- g * Hᶠᶜ * ∂xᵣ(i, j, k_top, grid, η★, timestepper, η) + Gᵁˢ) + Vᵐ⁺¹ = V[i, j, 1] + Δτ * (- g * Hᶜᶠ * ∂yᵣ(i, j, k_top, grid, η★, timestepper, η) + Gⱽˢ) - # Averaging the transport - Ũ[i, j, 1] += transport_weight * U[i, j, 1] - Ṽ[i, j, 1] += transport_weight * V[i, j, 1] + U[i, j, 1] = Uᵐ⁺¹ + V[i, j, 1] = Vᵐ⁺¹ + + # Time-averaging the barotropic velocity + U̅[i, j, 1] += averaging_weight * Uᵐ⁺¹ + V̅[i, j, 1] += averaging_weight * Vᵐ⁺¹ end end -@kernel function _split_explicit_free_surface!(averaging_weight, grid, filled_halos, Δτ, η, U, V, F, clock, η̅, U̅, V̅, timestepper) +@kernel function _split_explicit_free_surface!(averaging_weight, transport_weight, grid, filled_halos, Δτ, η, U, V, F, clock, η̅, Ũ, Ṽ, timestepper) i, j = @index(Global, NTuple) k_top = grid.Nz+1 @@ -75,15 +80,125 @@ end δy = y_difference_operator(filled_halos) δh_U = (δx(i, j, grid.Nz, grid, Δy_qᶠᶜᶠ, U★, timestepper, U) + - δy(i, j, grid.Nz, grid, Δx_qᶜᶠᶠ, U★, timestepper, V)) * Az⁻¹ᶜᶜᶠ(i, j, k_top, grid) + δy(i, j, grid.Nz, grid, Δx_qᶜᶠᶠ, V★, timestepper, V)) * Az⁻¹ᶜᶜᶠ(i, j, k_top, grid) @inbounds begin - η[i, j, k_top] += Δτ * (F(i, j, k_top, grid, clock, (; η, U, V)) - δh_U) + ηᵐ⁺¹ = η[i, j, k_top] + Δτ * (F(i, j, k_top, grid, clock, (; η, U, V)) - δh_U) + η[i, j, k_top] = ηᵐ⁺¹ - # Time-averaging - η̅[i, j, k_top] += averaging_weight * η[i, j, k_top] - U̅[i, j, 1] += averaging_weight * U[i, j, 1] - V̅[i, j, 1] += averaging_weight * V[i, j, 1] + # Time-averaging η and the transport U★/V★ that advanced it, which is what constancy needs. + η̅[i, j, k_top] += averaging_weight * ηᵐ⁺¹ + Ũ[i, j, 1] += transport_weight * U★(i, j, 1, grid, timestepper, U) + Ṽ[i, j, 1] += transport_weight * V★(i, j, 1, grid, timestepper, V) + end +end + +##### +##### Multi-stage substep kernel (RungeKutta3Scheme) +##### +##### Each barotropic substep runs three RK stages from the substep-start state (ηⁿ, Uⁿ, Vⁿ). Free surface and +##### velocity commute within a stage, since both depend only on the previous-stage state, so one kernel does +##### both; the previous stage needs its own buffer to keep the δx stencil from racing on the live fields. +##### The averages (η̅, U̅, V̅, Ũ, Ṽ) accumulate only on the final stage, which is what preserves constancy. +##### +##### One kernel per position in the stage sequence, so that each reads and writes exactly what it needs: the +##### opening stage reads the substep-start state as its previous stage, the closing stage writes back into it +##### and carries the averages, and the stages in between carry neither. + +# The state triples travel as tuples: `isregional` recurses element by element over the argument list and +# stops inferring on a long one. + +# ∂τ(U) = - gH ∇η + G and ∂τ(η) = - ∇⋅U, both evaluated at `previous_state` and applied from `substep_state`. +@inline function barotropic_stage_tendencies(i, j, grid, filled_halos, ηᵖ, Uᵖ, Vᵖ, slow_forcings, g, F, clock, sᵐ) + k_top = grid.Nz + 1 + + Gᵁ, Gⱽ, Gᵁᶜ, Gⱽᶜ, w = slow_forcings + + Hᶠᶜ = x_column_depth(i, j, k_top, grid, filled_halos, ηᵖ) + Hᶜᶠ = y_column_depth(i, j, k_top, grid, filled_halos, ηᵖ) + ∂xᵣ = x_derivative_operator(filled_halos) + ∂yᵣ = y_derivative_operator(filled_halos) + δx = x_difference_operator(filled_halos) + δy = y_difference_operator(filled_halos) + + # With the frozen forcing this dispatches back to a plain `Gᵁ[i, j, 1]` load. + Gᵁˢ = slow_forcing(i, j, Gᵁ, Gᵁᶜ, w, sᵐ) + Gⱽˢ = slow_forcing(i, j, Gⱽ, Gⱽᶜ, w, sᵐ) + + δh_U = (δx(i, j, grid.Nz, grid, Δy_qᶠᶜᶠ, Uᵖ) + δy(i, j, grid.Nz, grid, Δx_qᶜᶠᶠ, Vᵖ)) * Az⁻¹ᶜᶜᶠ(i, j, k_top, grid) + + Gη = F(i, j, k_top, grid, clock, (; η = ηᵖ, U = Uᵖ, V = Vᵖ)) - δh_U + GU = - g * Hᶠᶜ * ∂xᵣ(i, j, k_top, grid, ηᵖ) + Gᵁˢ + GV = - g * Hᶜᶠ * ∂yᵣ(i, j, k_top, grid, ηᵖ) + Gⱽˢ + + return Gη, GU, GV +end + +# Opening stage: the previous stage *is* the substep-start state, so one state triple travels instead of two. +@kernel function _first_barotropic_stage!(sᵐ, Δτ, grid, filled_halos, substep_state, stage_state, slow_forcings, g, F, clock) + + i, j = @index(Global, NTuple) + k_top = grid.Nz + 1 + + ηⁿ, Uⁿ, Vⁿ = substep_state + η, U, V = stage_state + + Gη, GU, GV = barotropic_stage_tendencies(i, j, grid, filled_halos, ηⁿ, Uⁿ, Vⁿ, slow_forcings, g, F, clock, sᵐ) + + @inbounds begin + η[i, j, k_top] = ηⁿ[i, j, k_top] + Δτ * Gη + U[i, j, 1] = Uⁿ[i, j, 1] + Δτ * GU + V[i, j, 1] = Vⁿ[i, j, 1] + Δτ * GV + end +end + +# Intermediate stages: no average accumulates here, so the filtered state does not travel. +@kernel function _barotropic_stage!(sᵐ, Δτ, grid, filled_halos, substep_state, previous_state, stage_state, slow_forcings, g, F, clock) + + i, j = @index(Global, NTuple) + k_top = grid.Nz + 1 + + ηⁿ, Uⁿ, Vⁿ = substep_state + ηᵖ, Uᵖ, Vᵖ = previous_state + η, U, V = stage_state + + Gη, GU, GV = barotropic_stage_tendencies(i, j, grid, filled_halos, ηᵖ, Uᵖ, Vᵖ, slow_forcings, g, F, clock, sᵐ) + + @inbounds begin + η[i, j, k_top] = ηⁿ[i, j, k_top] + Δτ * Gη + U[i, j, 1] = Uⁿ[i, j, 1] + Δτ * GU + V[i, j, 1] = Vⁿ[i, j, 1] + Δτ * GV + end +end + +# Closing stage: it writes back into the substep-start state, and it alone feeds the time averages. +@kernel function _final_barotropic_stage!(averaging_weight, transport_weight, sᵐ, Δτ, grid, filled_halos, substep_state, previous_state, slow_forcings, g, F, clock, averages) + + i, j = @index(Global, NTuple) + k_top = grid.Nz + 1 + + η, U, V = substep_state + ηᵖ, Uᵖ, Vᵖ = previous_state + η̅, U̅, V̅, Ũ, Ṽ = averages + + Gη, GU, GV = barotropic_stage_tendencies(i, j, grid, filled_halos, ηᵖ, Uᵖ, Vᵖ, slow_forcings, g, F, clock, sᵐ) + + @inbounds begin + ηⁿ⁺¹ = η[i, j, k_top] + Δτ * Gη + Uⁿ⁺¹ = U[i, j, 1] + Δτ * GU + Vⁿ⁺¹ = V[i, j, 1] + Δτ * GV + + η[i, j, k_top] = ηⁿ⁺¹ + U[i, j, 1] = Uⁿ⁺¹ + V[i, j, 1] = Vⁿ⁺¹ + + η̅[i, j, k_top] += averaging_weight * ηⁿ⁺¹ + U̅[i, j, 1] += averaging_weight * Uⁿ⁺¹ + V̅[i, j, 1] += averaging_weight * Vⁿ⁺¹ + + # The transport in the tracer-continuity average is the flux that advanced η, i.e. the previous stage. + Ũ[i, j, 1] += transport_weight * Uᵖ[i, j, 1] + Ṽ[i, j, 1] += transport_weight * Vᵖ[i, j, 1] end end @@ -101,113 +216,254 @@ const MINIMUM_SUBSTEPS = 5 @inline calculate_adaptive_settings(substepping::FNS, substeps) = substepping.fractional_step_size, substepping.averaging_weights, substepping.transport_weights @inline calculate_adaptive_settings(substepping::FTS, substeps) = weights_from_substeps(eltype(substepping.Δt_barotropic), substeps, substepping.averaging_kernel) -iterate_split_explicit!(free_surface, grid, GUⁿ, GVⁿ, Δτᴮ, F, clock, weights, transport_weights, ::Val{Nsubsteps}) where Nsubsteps = - @apply_regionally iterate_split_explicit_in_halo!(free_surface, grid, GUⁿ, GVⁿ, Δτᴮ, F, clock, weights, transport_weights, Val(Nsubsteps)) +##### +##### Halo filling in between barotropic kernels +##### +##### `ExtendedHalos` substeps into the halo rather than filling it, so its argument groups are empty. -function iterate_split_explicit!(free_surface::FillHaloSplitExplicit, grid, GUⁿ, GVⁿ, Δτᴮ, F, clock, weights, transport_weights, ::Val{Nsubsteps}) where Nsubsteps - arch = architecture(grid) +@inline fill_barotropic_halos!(arch, substep_clock, ::Tuple{}, halo_exchange) = nothing - η = free_surface.displacement - grid = free_surface.displacement.grid - arch = architecture(grid) - state = free_surface.filtered_state - timestepper = free_surface.timestepper - g = free_surface.gravitational_acceleration - parameters = free_surface.kernel_parameters +# The halo arguments are built once per barotropic step and close with the reference fields; the substep clock +# is spliced in right before them. +@inline function fill_barotropic_halos!(arch, substep_clock, halo_args::Tuple, halo_exchange) + field_args = first(halo_args) + maybe_distributed_fill_halo_regions!(arch, Base.front(field_args)..., substep_clock, last(field_args); only_local_halos = !halo_exchange) + return fill_barotropic_halos!(arch, substep_clock, Base.tail(halo_args), halo_exchange) +end + +@inline exchange_extended_halos!(free_surface, halo_exchange) = nothing + +function exchange_extended_halos!(free_surface::SplitExplicitFreeSurface{ExtendedHalos}, halo_exchange) + U, V = free_surface.barotropic_velocities + halo_exchange && fill_halo_regions!((U, V, free_surface.displacement)) + return nothing +end + +# A substep clock with a smaller Δτ that advances through the substeps, for time-dependent boundary conditions. +@inline barotropic_substep_clock(clock, start_time, substep, Δτᴮ) = + (; time = add_time_interval(start_time, Δτᴮ, substep - 1), iteration = clock.iteration, stage = 0, last_stage_Δt = Δτᴮ) - # Unpack state quantities, parameters and forcing terms. +@inline barotropic_halo_arguments(::SplitExplicitFreeSurface{ExtendedHalos}, arch, grid, reference_fields, fields_to_fill::Tuple) = () +@inline barotropic_halo_arguments(::FillHaloSplitExplicit, arch, grid, reference_fields, ::Tuple{}) = () + +@inline function barotropic_halo_arguments(free_surface::FillHaloSplitExplicit, arch, grid, reference_fields, fields_to_fill::Tuple) + @apply_regionally halo_args = build_halo_fill_args(arch, first(fields_to_fill), grid, reference_fields) + remaining = barotropic_halo_arguments(free_surface, arch, grid, reference_fields, Base.tail(fields_to_fill)) + return (halo_args, remaining...) +end + +# `filled_halos` selects the topology-aware operators when the halos go stale in between kernels. +@inline substep_filled_halos(::SplitExplicitFreeSurface{ExtendedHalos}) = Val(false) +@inline substep_filled_halos(::FillHaloSplitExplicit) = Val(true) + +##### +##### Barotropic substeppers +##### +##### A substepper collects, once per baroclinic step, everything that does not change from substep to substep. +##### We perform ~50 substeps of ~100 very small kernels, whose latency is dominated by argument conversion. + +# One method per leading-scalar count: a `weights...` vararg splats through a dynamic call and boxes the +# whole argument list on every launch. +@inline apply_barotropic_kernel!(kernel, args, w₁) = kernel(w₁, args...) +@inline apply_barotropic_kernel!(kernel, args, w₁, w₂) = kernel(w₁, w₂, args...) +@inline apply_barotropic_kernel!(kernel, args, w₁, w₂, w₃) = kernel(w₁, w₂, w₃, args...) + +function barotropic_substepper(timestepper::ForwardBackwardScheme, free_surface, arch, grid, parameters, substep_arguments) + + η = free_surface.displacement U, V = free_surface.barotropic_velocities + state = free_surface.filtered_state η̅, U̅, V̅ = state.η̅, state.U̅, state.V̅ - Ũ, Ṽ = state.Ũ, state.Ṽ + Ũ, Ṽ = state.Ũ, state.Ṽ + + (; Δτᴮ, clock, F, g, GUⁿ, GVⁿ, Gᵁᶜ, Gⱽᶜ, w) = substep_arguments + + filled_halos = substep_filled_halos(free_surface) @apply_regionally workspec = Utils.possibly_load_active_cells_map(grid, parameters, false) @apply_regionally velocity_kernel!, _ = configure_kernel(arch, grid, workspec, _split_explicit_barotropic_velocity!) @apply_regionally free_surface_kernel!, _ = configure_kernel(arch, grid, workspec, _split_explicit_free_surface!) - U_args = (grid, Val(true), Δτᴮ, η, U, V, GUⁿ, GVⁿ, g, Ũ, Ṽ, timestepper) - η_args = (grid, Val(true), Δτᴮ, η, U, V, F, clock, η̅, U̅, V̅, timestepper) + U_args = (grid, filled_halos, Δτᴮ, η, U, V, GUⁿ, GVⁿ, Gᵁᶜ, Gⱽᶜ, w, g, U̅, V̅, timestepper) + η_args = (grid, filled_halos, Δτᴮ, η, U, V, F, clock, η̅, Ũ, Ṽ, timestepper) + + @apply_regionally converted_U_args = convert_to_device(arch, U_args) + @apply_regionally converted_η_args = convert_to_device(arch, η_args) barotropic_model_fields = (; U, V, η) - # Builds also a separate "sub-stepping" clock to account for time dependent forcing and boundary conditions - @apply_regionally U_halo_args = build_halo_fill_args(arch, U, grid, barotropic_model_fields) - @apply_regionally V_halo_args = build_halo_fill_args(arch, V, grid, barotropic_model_fields) - @apply_regionally η_halo_args = build_halo_fill_args(arch, η, grid, barotropic_model_fields) + velocity_halos = barotropic_halo_arguments(free_surface, arch, grid, barotropic_model_fields, (U, V)) + free_surface_halos = barotropic_halo_arguments(free_surface, arch, grid, barotropic_model_fields, (η, )) - exchange_interval = halo_exchange_interval(free_surface, grid, Nsubsteps) + face_pins = configure_face_pins(arch, grid, U, V, free_surface.boundary_transport) - GC.@preserve U_args η_args begin - # We need to perform ~50 time-steps which means launching ~100 very small kernels: we are limited by latency of - # argument conversion to GPU-compatible values. To alleviate this penalty we convert first and then we substep! - @apply_regionally converted_U_args = convert_to_device(arch, U_args) - @apply_regionally converted_η_args = convert_to_device(arch, η_args) - face_pins = configure_face_pins(arch, grid, U, V, free_surface.boundary_transport) + return (; free_surface_kernel!, velocity_kernel!, η_args = converted_η_args, U_args = converted_U_args, + velocity_halos, free_surface_halos, face_pins, barotropic_model_fields, Δτᴮ) +end - @unroll for substep in 1:Nsubsteps - @inbounds averaging_weight = weights[substep] - @inbounds transport_weight = transport_weights[substep] +function barotropic_substepper(timestepper::RungeKutta3Scheme, free_surface, arch, grid, parameters, substep_arguments) + + η = free_surface.displacement + U, V = free_surface.barotropic_velocities + state = free_surface.filtered_state + η̅, U̅, V̅ = state.η̅, state.U̅, state.V̅ + Ũ, Ṽ = state.Ũ, state.Ṽ - substep_clock = (; time = add_time_interval(clock.time, Δτᴮ, substep - 1), iteration = clock.iteration, stage = 0, last_stage_Δt = Δτᴮ) - only_local_halos = !exchange_halos(free_surface, substep, exchange_interval) + (; Δτᴮ, clock, F, g, GUⁿ, GVⁿ, Gᵁᶜ, Gⱽᶜ, w) = substep_arguments - maybe_distributed_fill_halo_regions!(arch, η_halo_args[1:end-1]..., substep_clock, η_halo_args[end]; only_local_halos) - @apply_regionally cache_boundary_conditions!((U, V, η), substep_clock, barotropic_model_fields) - @apply_regionally apply_barotropic_kernel!(velocity_kernel!, transport_weight, converted_U_args) - @apply_regionally step_boundary_conditions!((U, V), substep_clock, barotropic_model_fields, Δτᴮ) + filled_halos = substep_filled_halos(free_surface) - maybe_distributed_fill_halo_regions!(arch, U_halo_args[1:end-1]..., substep_clock, U_halo_args[end]; only_local_halos) - maybe_distributed_fill_halo_regions!(arch, V_halo_args[1:end-1]..., substep_clock, V_halo_args[end]; only_local_halos) - pin_barotropic_faces!(face_pins) - @apply_regionally apply_barotropic_kernel!(free_surface_kernel!, averaging_weight, converted_η_args) - @apply_regionally step_boundary_conditions!((η,), substep_clock, barotropic_model_fields, Δτᴮ) + @apply_regionally workspec = Utils.possibly_load_active_cells_map(grid, parameters, false) + @apply_regionally first_stage_kernel!, _ = configure_kernel(arch, grid, workspec, _first_barotropic_stage!) + @apply_regionally stage_kernel!, _ = configure_kernel(arch, grid, workspec, _barotropic_stage!) + @apply_regionally final_stage_kernel!, _ = configure_kernel(arch, grid, workspec, _final_barotropic_stage!) + + stages_Δτ = stage_parameters(timestepper, Δτᴮ) + Nstages = length(stages_Δτ) + + substep_state = (η, U, V) + first_buffer = (timestepper.η¹, timestepper.U¹, timestepper.V¹) + second_buffer = (timestepper.η², timestepper.U², timestepper.V²) + + # The buffers rotate in the three stages, the last one writing back into the substep-start state. + previous_state = (substep_state, first_buffer, second_buffer) + stage_state = (first_buffer, second_buffer, substep_state) + + slow_forcings = (GUⁿ, GVⁿ, Gᵁᶜ, Gⱽᶜ, w) + averages = (η̅, U̅, V̅, Ũ, Ṽ) + + # `ntuple` over a `Val` splices `stage` in as a literal, so each branch is resolved at compile time. + stage_args = ntuple(Val(Nstages)) do stage + args = if stage == 1 + (stages_Δτ[stage], grid, filled_halos, substep_state, stage_state[stage], slow_forcings, g, F, clock) + elseif stage == Nstages + (stages_Δτ[stage], grid, filled_halos, substep_state, previous_state[stage], slow_forcings, g, F, clock, averages) + else + (stages_Δτ[stage], grid, filled_halos, substep_state, previous_state[stage], stage_state[stage], slow_forcings, g, F, clock) end + + @apply_regionally converted_args = convert_to_device(arch, args) + converted_args + end + + stage_halos = ntuple(Val(Nstages)) do stage + ηᵖ, Uᵖ, Vᵖ = previous_state[stage] + barotropic_halo_arguments(free_surface, arch, grid, (; U = Uᵖ, V = Vᵖ, η = ηᵖ), (Uᵖ, Vᵖ, ηᵖ)) end + # Each stage advances η with the transport of its previous stage, so that is the transport pinned to the target. + stage_pins = ntuple(Val(Nstages)) do stage + ηᵖ, Uᵖ, Vᵖ = previous_state[stage] + configure_face_pins(arch, grid, Uᵖ, Vᵖ, free_surface.boundary_transport) + end + + # Each stage holds its own exterior states; all of them step from the one cached at the substep start. + stage_boundaries = ntuple(Val(Nstages)) do stage + ηˢ, Uˢ, Vˢ = stage_state[stage] + (; stage_fields = (; U = Uˢ, V = Vˢ, η = ηˢ), Δτ = stages_Δτ[stage]) + end + + barotropic_model_fields = (; U, V, η) + + return (; first_stage_kernel!, stage_kernel!, final_stage_kernel!, stage_args, stage_halos, stage_pins, + stage_boundaries, barotropic_model_fields) +end + +function barotropic_substep!(::ForwardBackwardScheme, substepper, arch, substep_clock, averaging_weight, transport_weight, sᵐ, halo_exchange) + + (; barotropic_model_fields, Δτᴮ) = substepper + (; U, V, η) = barotropic_model_fields + + fill_barotropic_halos!(arch, substep_clock, substepper.free_surface_halos, halo_exchange) + @apply_regionally cache_boundary_conditions!((U, V, η), substep_clock, barotropic_model_fields) + @apply_regionally apply_barotropic_kernel!(substepper.velocity_kernel!, substepper.U_args, averaging_weight, sᵐ) + @apply_regionally step_boundary_conditions!((U, V), substep_clock, barotropic_model_fields, Δτᴮ) + + fill_barotropic_halos!(arch, substep_clock, substepper.velocity_halos, halo_exchange) + pin_barotropic_faces!(substepper.face_pins) + @apply_regionally apply_barotropic_kernel!(substepper.free_surface_kernel!, substepper.η_args, averaging_weight, transport_weight) + @apply_regionally step_boundary_conditions!((η,), substep_clock, barotropic_model_fields, Δτᴮ) + + return nothing +end + +@inline function step_stage_boundaries!(stage_boundary, substep_clock) + (; stage_fields, Δτ) = stage_boundary + @apply_regionally step_boundary_conditions!(stage_fields, substep_clock, stage_fields, Δτ) return nothing end -@inline apply_barotropic_kernel!(kernel, weight, args) = kernel(weight, args...) +# The stages recurse over their argument tuples rather than indexing them with a running counter: the tuples +# are heterogeneous, one entry per position in the sequence, and recursion keeps every index a literal. +@inline function barotropic_stages!(substepper, arch, substep_clock, args::Tuple{Any}, halos::Tuple{Any}, + pins::Tuple{Any}, boundaries::Tuple{Any}, averaging_weight, transport_weight, sᵐ, halo_exchange) + + fill_barotropic_halos!(arch, substep_clock, first(halos), halo_exchange) + pin_barotropic_faces!(first(pins)) + @apply_regionally apply_barotropic_kernel!(substepper.final_stage_kernel!, first(args), averaging_weight, transport_weight, sᵐ) + step_stage_boundaries!(first(boundaries), substep_clock) + return nothing +end + +@inline function barotropic_stages!(substepper, arch, substep_clock, args::Tuple, halos::Tuple, + pins::Tuple, boundaries::Tuple, averaging_weight, transport_weight, sᵐ, halo_exchange) + + fill_barotropic_halos!(arch, substep_clock, first(halos), halo_exchange) + pin_barotropic_faces!(first(pins)) + @apply_regionally apply_barotropic_kernel!(substepper.stage_kernel!, first(args), sᵐ) + step_stage_boundaries!(first(boundaries), substep_clock) + + return barotropic_stages!(substepper, arch, substep_clock, Base.tail(args), Base.tail(halos), + Base.tail(pins), Base.tail(boundaries), averaging_weight, transport_weight, sᵐ, halo_exchange) +end + +@noinline function barotropic_substep!(::RungeKutta3Scheme, substepper, arch, substep_clock, averaging_weight, transport_weight, sᵐ, halo_exchange) + + stage_args = substepper.stage_args + stage_halos = substepper.stage_halos + stage_pins = substepper.stage_pins + stage_boundaries = substepper.stage_boundaries + barotropic_model_fields = substepper.barotropic_model_fields + + fill_barotropic_halos!(arch, substep_clock, first(stage_halos), halo_exchange) + pin_barotropic_faces!(first(stage_pins)) + @apply_regionally cache_boundary_conditions!(barotropic_model_fields, substep_clock, barotropic_model_fields) + @apply_regionally apply_barotropic_kernel!(substepper.first_stage_kernel!, first(stage_args), sᵐ) + step_stage_boundaries!(first(stage_boundaries), substep_clock) + + return barotropic_stages!(substepper, arch, substep_clock, Base.tail(stage_args), Base.tail(stage_halos), + Base.tail(stage_pins), Base.tail(stage_boundaries), averaging_weight, transport_weight, sᵐ, halo_exchange) +end -function iterate_split_explicit_in_halo!(free_surface, grid, GUⁿ, GVⁿ, Δτᴮ, F::Fη, clock, weights, transport_weights, ::Val{Nsubsteps}) where {Fη, Nsubsteps} - arch = architecture(grid) +# `Δt` is the step of the current baroclinic stage and `Δτᴮ` the barotropic substep. +function iterate_split_explicit!(free_surface, grid, GUⁿ, GVⁿ, Δτᴮ, Δt, F, clock, weights, transport_weights, + baroclinic_timestepper, Nsubsteps::Int, start_time = clock.time) - η = free_surface.displacement + # Note: `free_surface.displacement.grid` carries the extended halos, where `grid` may not. grid = free_surface.displacement.grid - state = free_surface.filtered_state + arch = architecture(grid) timestepper = free_surface.timestepper - g = free_surface.gravitational_acceleration parameters = free_surface.kernel_parameters + g = free_surface.gravitational_acceleration - # Unpack state quantities, parameters and forcing terms. - U, V = free_surface.barotropic_velocities - η̅, U̅, V̅ = state.η̅, state.U̅, state.V̅ - Ũ, Ṽ = state.Ũ, state.Ṽ + Gᵁᶜ, Gⱽᶜ, w = stage_reconstruction(free_surface.slow_forcing, baroclinic_timestepper, clock.stage, Δt) - workspec = Utils.possibly_load_active_cells_map(grid, parameters, false) - barotropic_velocity_kernel!, _ = configure_kernel(arch, grid, workspec, _split_explicit_barotropic_velocity!) - free_surface_kernel!, _ = configure_kernel(arch, grid, workspec, _split_explicit_free_surface!) + substep_arguments = (; Δτᴮ, clock, F, g, GUⁿ, GVⁿ, Gᵁᶜ, Gⱽᶜ, w) + substepper = barotropic_substepper(timestepper, free_surface, arch, grid, parameters, substep_arguments) exchange_interval = halo_exchange_interval(free_surface, grid, Nsubsteps) - U_args = (grid, Val(false), Δτᴮ, η, U, V, GUⁿ, GVⁿ, g, Ũ, Ṽ, timestepper) - η_args = (grid, Val(false), Δτᴮ, η, U, V, F, clock, η̅, U̅, V̅, timestepper) - - GC.@preserve U_args η_args begin - # We need to perform ~50 time-steps which means launching ~100 very small kernels: we are limited by latency of - # argument conversion to GPU-compatible values. To alleviate this penalty we convert first and then we substep! - converted_U_args = convert_to_device(arch, U_args) - converted_η_args = convert_to_device(arch, η_args) - - @unroll for substep in 1:Nsubsteps + # The substepper holds device values that do not root the fields they point into. + GC.@preserve free_surface GUⁿ GVⁿ Gᵁᶜ Gⱽᶜ F begin + for substep in 1:Nsubsteps @inbounds averaging_weight = weights[substep] @inbounds transport_weight = transport_weights[substep] + sᵐ = (substep - oneunit(substep) / 2) * Δτᴮ # MIDPOINT of the substep, measured from tⁿ + substep_clock = barotropic_substep_clock(clock, start_time, substep, Δτᴮ) + halo_exchange = exchange_halos(free_surface, substep, exchange_interval) - if exchange_halos(free_surface, substep, exchange_interval) - fill_halo_regions!((U, V, η)) - end - - barotropic_velocity_kernel!(transport_weight, converted_U_args...) - free_surface_kernel!(averaging_weight, converted_η_args...) + exchange_extended_halos!(free_surface, halo_exchange) + barotropic_substep!(timestepper, substepper, arch, substep_clock, averaging_weight, transport_weight, sᵐ, halo_exchange) end end @@ -231,11 +487,16 @@ end # Open boundaries read model fields while filling the barotropic halos; `ExtendedHalos` has none, so it # fills without threading them, which avoids a per-step allocation on distributed grids. -@inline fill_barotropic_state_halos!(field, ::SplitExplicitFreeSurface{ExtendedHalos}, model) = - fill_halo_regions!(field; async=true) +@inline fill_barotropic_state_halos!(field, ::SplitExplicitFreeSurface{ExtendedHalos}, model) = fill_halo_regions!(field; async=true) +@inline fill_barotropic_state_halos!(field, ::FillHaloSplitExplicit, model) = fill_halo_regions!(field, model.clock, fields(model); async=true) + +# Every SSP solve restarts from ψⁿ at tⁿ, while the clock has reached the time of the stage +@inline barotropic_start_time(clock, timestepper, Δt) = clock.time -@inline fill_barotropic_state_halos!(field, ::FillHaloSplitExplicit, model) = - fill_halo_regions!(field, model.clock, fields(model); async=true) +@inline function barotropic_start_time(clock, ts::SSPRungeKuttaTimeStepper, Δt) + τ = (0, ssp_stage_fractions(ts.coefficients)...) + return add_time_interval(clock.time, τ[clock.stage] * Δt, -1) +end function HydrostaticFreeSurfaceModels.step_free_surface!(free_surface::SplitExplicitFreeSurface, model, baroclinic_timestepper, Δt) # Note: free_surface.displacement.grid != model.grid for DistributedSplitExplicitFreeSurface since @@ -275,14 +536,17 @@ function HydrostaticFreeSurfaceModels.step_free_surface!(free_surface::SplitExpl @apply_regionally initialize_free_surface_state!(free_surface, baroclinic_timestepper, barotropic_timestepper) # Substep halo fills read time series up to the last substep time, which can lie past tⁿ⁺¹ - t = model.clock.time + t = barotropic_start_time(model.clock, baroclinic_timestepper, Δt) update_model_field_time_series!(model, Time(t), Time(add_time_interval(t, Δτᴮ, Nsubsteps - 1))) # Solve for the free surface at tⁿ⁺¹. - iterate_split_explicit!(free_surface, free_surface_grid, GUⁿ, GVⁿ, Δτᴮ, F, model.clock, weights, transport_weights, Val(Nsubsteps)) + # Substep times start from `t`; the η forcing `F` reads `model.clock`, the time of the stage. + iterate_split_explicit!(free_surface, free_surface_grid, GUⁿ, GVⁿ, Δτᴮ, Δt, F, model.clock, weights, + transport_weights, baroclinic_timestepper, Nsubsteps, t) # Update eta and velocities for the next timestep. The halos are updated in the `update_state!` function. @apply_regionally launch!(architecture(free_surface_grid), free_surface_grid, :xy, _update_split_explicit_state!, η, U, V, free_surface_grid, filtered_state) + @apply_regionally HydrostaticFreeSurfaceModels.blend_free_surface!(free_surface, baroclinic_timestepper, model.clock.stage) # Fill all the barotropic state. fill_barotropic_state_halos!((filtered_state.Ũ, filtered_state.Ṽ), free_surface, model) diff --git a/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl index 80293f82ebd..ddc34f7fb9a 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl @@ -73,12 +73,22 @@ function step_free_surface!(free_surface::ExplicitFreeSurface, model, timesteppe return nothing end -function step_free_surface!(free_surface::ExplicitFreeSurface, model, timestepper::SplitRungeKuttaTimeStepper, Δt) +function step_free_surface!(free_surface::ExplicitFreeSurface, model, timestepper::MultiStageTimeStepper, Δt) @apply_regionally explicit_rk3_step_free_surface!(free_surface, model, Δt) fill_halo_regions!(free_surface.displacement; async=true) return nothing end +function step_free_surface!(free_surface::ExplicitFreeSurface, model, timestepper::SSPRungeKuttaTimeStepper, Δt) + FT = eltype(model.grid) + a, b = timestepper.coefficients[model.clock.stage] + @apply_regionally launch!(model.architecture, model.grid, :xy, _ssp_step_explicit_free_surface!, + free_surface.displacement, convert(FT, Δt), timestepper.Gⁿ.η, timestepper.Ψ⁻.η, + convert(FT, a), convert(FT, b), size(model.grid, 3)) + fill_halo_regions!(free_surface.displacement; async=true) + return nothing +end + explicit_rk3_step_free_surface!(free_surface, model, Δt) = launch!(model.architecture, model.grid, :xy, _explicit_rk3_step_free_surface!, free_surface.displacement, Δt, @@ -98,6 +108,12 @@ explicit_ab2_step_free_surface!(free_surface, model, Δt) = @inbounds η[i, j, Nz+1] = η⁻[i, j, Nz+1] + Δt * Gⁿ[i, j, Nz+1] end +# ηᵐ = a ηⁿ + b (ηᵐ⁻¹ + Δt Gᵐ) +@kernel function _ssp_step_explicit_free_surface!(η, Δt, Gⁿ, ηⁿ, a, b, Nz) + i, j = @index(Global, NTuple) + @inbounds η[i, j, Nz+1] = a * ηⁿ[i, j, Nz+1] + b * (η[i, j, Nz+1] + Δt * Gⁿ[i, j, Nz+1]) +end + @kernel function _explicit_ab2_step_free_surface!(η, Δt, χ, Gηⁿ, Gη⁻, Nz) i, j = @index(Global, NTuple) FT0 = typeof(χ) diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl index 3344b07c52f..7f87767afc1 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl @@ -73,7 +73,7 @@ mutable struct HydrostaticFreeSurfaceModel{TS, E, A<:AbstractArchitecture, S, boundary_transport :: BM # Transport fields for targeted open boundary conditions (or `nothing`) end -supported_timesteppers = (:QuasiAdamsBashforth2, :SplitRungeKutta2, :SplitRungeKutta3, :SplitRungeKutta4, :SplitRungeKutta5) +supported_timesteppers = (:QuasiAdamsBashforth2, :SplitRungeKutta2, :SplitRungeKutta3, :SplitRungeKutta4, :SplitRungeKutta5, :SSPRungeKutta3, :ModifiedRungeKutta4) default_free_surface(grid::XYRegularStaticRG; gravitational_acceleration=defaults.gravitational_acceleration) = ImplicitFreeSurface(; gravitational_acceleration) @@ -331,6 +331,7 @@ function build_hydrostatic_free_surface_model(grid, ::Val{tracer_names}, timeste G⁻ = previous_hydrostatic_tendency_fields(timestepper, velocities, free_surface, grid, tracer_names, boundary_conditions) cached_state = previous_hydrostatic_state_fields(timestepper, velocities, free_surface, tracers) timestepper = TimeStepper(timestepper, grid, prognostic_fields; implicit_solver, Gⁿ, G⁻, cached_state...) + validate_timestepper_free_surface(timestepper, free_surface) materialize_clock!(clock, timestepper) # Materialize forcing for model tracer and velocity fields. diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ssp_step.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ssp_step.jl new file mode 100644 index 00000000000..8ae2bb74d39 --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ssp_step.jl @@ -0,0 +1,235 @@ +using Oceananigans.TimeSteppers: _ssp_euler_substep_field!, _ssp_blend_field!, _ssp_blend_tracer_field!, + SSPBarotropicForcing, implicit_step!, ssp_stage_fractions, ssp_substep! + +# `ImplicitFreeSurface` solves on a predictor velocity, so the Shu-Osher blend would have to be deferred +# until after a pressure correction that has already overwritten Ψᵐ⁻¹. +validate_timestepper_free_surface(timestepper, free_surface) = nothing + +validate_timestepper_free_surface(::SSPRungeKuttaTimeStepper, ::ImplicitFreeSurface) = + throw(ArgumentError("SSPRungeKuttaTimeStepper does not support ImplicitFreeSurface")) + +function validate_timestepper_free_surface(timestepper::MultiStageTimeStepper, free_surface::SplitExplicitFreeSurface) + slow_forcing = free_surface.slow_forcing + slow_forcing isa SplitExplicitFreeSurfaces.ProgressiveSlowForcing || return nothing + + τ = SplitExplicitFreeSurfaces.stage_sample_times(timestepper) + nodes = slow_forcing.nodes + length(nodes) == length(τ) && all(isapprox.(nodes, τ; rtol=1e-12)) && return nothing + + throw(ArgumentError("$(nameof(typeof(slow_forcing))) samples the slow forcing at $nodes, but the timestepper samples at $τ")) +end + +""" +$(TYPEDSIGNATURES) + +Perform one strong-stability-preserving Runge-Kutta stage for `HydrostaticFreeSurfaceModel`. + +The baroclinic velocities, tracers and grid are advanced by a full `Δt` and blended with the cached state +by the Shu-Osher pair `(a, b)`. Each stage is an implicit-explicit forward-Euler step: the explicit tendency +and the vertical solve are applied to `Ψᵐ⁻¹`, and the blend follows. + +The barotropic pair is blended by *increment* rather than by state: every sub-cycle restarts from +`(ηⁿ, Uⁿ, Vⁿ)` and spans the full `Δt`, and the blend applies the increment it produced to the previous +stage, `Ψᵐ = a Ψⁿ + b (Ψᵐ⁻¹ + Ψ⋆ - Ψⁿ)`. Because the Shu-Osher weights sum to one, the free barotropic +propagator is applied exactly once per step, so the over-integration that blending the barotropic *state* +would incur -- the Shu-Osher weights are derived for forward-Euler stages, not for near-exact advances -- +does not arise, while the barotropic stage values are the ones the Shu-Osher recursion requires. The last +stage is driven by the stage-weighted slow forcing, so its sub-cycle is simultaneously the predictor that +supplies that stage's transport and the corrector that sets the barotropic pair at `n+1`. +""" +TimeSteppers.ssp_substep!(model::HydrostaticFreeSurfaceModel, Δt, a, b, callbacks) = + ssp_substep!(model, model.free_surface, model.grid, Δt, a, b, callbacks) + +@inline function TimeSteppers.ssp_substep!(model, free_surface, grid, Δt, a, b, callbacks) + timestepper = model.timestepper + + # As in split Runge-Kutta, exterior states step from their value at tⁿ to the time the stage reaches + stage_Δt = convert(eltype(grid), ssp_stage_fractions(timestepper.coefficients)[model.clock.stage] * Δt) + + @apply_regionally begin + update_transport_velocities!(model.transport_velocities, model.velocities, free_surface) + compute_momentum_flux_bcs!(model) + ssp_euler_substep_velocities!(model.velocities, model, Δt) + end + + compute_free_surface_tendency!(grid, model, free_surface, Δt) + + foreach((ψᵐ⁻¹, ψ) -> ψᵐ⁻¹ .= ψ, previous_stage_fields(timestepper.G★), free_surface_fields(free_surface)) + step_free_surface!(free_surface, model, timestepper, Δt) + + @apply_regionally begin + ssp_blend_velocities!(model.velocities, model, a, b, stage_Δt) + mask_immersed_horizontal_velocities!(model.velocities) + compute_transport_velocities!(model, free_surface) + end + + u, v, _ = model.velocities + fill_halo_regions!((u, v), model.clock, fields(model); async=true) + + @apply_regionally compute_tracer_tendencies!(model) + execute_tracer_tendency_callbacks!(model, callbacks) + + @apply_regionally begin + rk_substep_grid!(grid, model, model.vertical_coordinate, Δt) + + # Reconciled at every stage: deferring it to the final stage breaks tracer constancy. + correct_barotropic_mode!(model, Δt) + ssp_substep_tracers!(model.tracers, model, Δt, a, b) + step_boundary_conditions!(model.tracers, model.clock, fields(model), stage_Δt) + end + + return nothing +end + +##### +##### Velocities +##### + +function ssp_euler_substep_velocities!(velocities, model, Δt) + ssp_euler_substep_velocity!(velocities, model, Δt, Val(:u)) + ssp_euler_substep_velocity!(velocities, model, Δt, Val(:v)) + return nothing +end + +@inline function ssp_euler_substep_velocity!(velocities, model, Δt, ::Val{name}) where name + grid = model.grid + FT = eltype(grid) + + Gⁿ = model.timestepper.Gⁿ[name] + velocity_field = velocities[name] + + launch!(architecture(grid), grid, :xyz, _ssp_euler_substep_field!, velocity_field, convert(FT, Δt), Gⁿ; exclude_periphery=true) + + implicit_step!(velocity_field, + model.timestepper.implicit_solver, + model.closure, + model.closure_fields, + nothing, + model.clock, + fields(model), + Δt, + model.advection.momentum, + model.velocities) + return nothing +end + +function ssp_blend_velocities!(velocities, model, a, b, stage_Δt) + grid = model.grid + FT = eltype(grid) + Ψ⁻ = model.timestepper.Ψ⁻ + a, b = convert(FT, a), convert(FT, b) + + launch!(architecture(grid), grid, :xyz, _ssp_blend_field!, velocities.u, Ψ⁻.u, a, b; exclude_periphery=true) + launch!(architecture(grid), grid, :xyz, _ssp_blend_field!, velocities.v, Ψ⁻.v, a, b; exclude_periphery=true) + step_boundary_conditions!(velocities, model.clock, fields(model), stage_Δt) + + return nothing +end + +##### +##### Tracers +##### + +ssp_substep_tracers!(::EmptyNamedTuple, model, Δt, a, b) = nothing + +function ssp_substep_tracers!(tracers, model, Δt, a, b) + ssp_substep_tracers!(model, Δt, a, b, Val(1), Val(propertynames(tracers))) + return nothing +end + +@inline ssp_substep_tracers!(model, Δt, a, b, ::Val, ::Val{()}) = nothing + +@inline function ssp_substep_tracers!(model, Δt, a, b, ::Val{tracer_index}, ::Val{names}) where {tracer_index, names} + ssp_substep_tracer!(model, Δt, a, b, Val(tracer_index), Val(first(names))) + ssp_substep_tracers!(model, Δt, a, b, Val(tracer_index + 1), Val(Base.tail(names))) + return nothing +end + +@inline function ssp_substep_tracer!(model, Δt, a, b, ::Val{tracer_index}, ::Val{tracer_name}) where {tracer_index, tracer_name} + closure = model.closure + (hasclosure(closure, FlavorOfCATKE) && tracer_name == :e) && return nothing + + grid = model.grid + FT = eltype(grid) + + Gⁿ = model.timestepper.Gⁿ[tracer_name] + Ψ⁻ = model.timestepper.Ψ⁻[tracer_name] + c = model.tracers[tracer_name] + + launch!(architecture(grid), grid, :xyz, _ssp_euler_substep_tracer_field!, c, grid, convert(FT, Δt), Gⁿ) + + # The adaptive implicit advection must see the same total velocity as the explicit flux, drift included + @inbounds c_advection = model.advection[tracer_name] + c_velocities = tracer_advecting_velocities(model.transport_velocities, model.biogeochemistry, closure, + model.closure_fields, model.forcing[tracer_name], Val(tracer_name)) + + implicit_step!(c, + model.timestepper.implicit_solver, + closure, + model.closure_fields, + Val(tracer_index), + model.clock, + fields(model), + Δt, + c_advection, + c_velocities) + + launch!(architecture(grid), grid, :xyz, _ssp_blend_tracer_field!, c, grid, Ψ⁻, convert(FT, a), convert(FT, b)) + + return nothing +end + +# The Euler stage advances the thickness-weighted tracer, (σĉ) = (σc)ᵐ⁻¹ + Δt Gᵐ, so the previous-stage +# term carries σᵐ⁻¹ rather than σᵐ. +@kernel function _ssp_euler_substep_tracer_field!(c, grid, Δt, Gⁿ) + i, j, k = @index(Global, NTuple) + σ⁻¹ᶜᶜⁿ = σ⁻¹ⁿ(i, j, k, grid, Center(), Center(), Center()) + σᶜᶜ⁻ = σ⁻(i, j, k, grid, Center(), Center(), Center()) + @inbounds c[i, j, k] = (σᶜᶜ⁻ * c[i, j, k] + Δt * Gⁿ[i, j, k]) * σ⁻¹ᶜᶜⁿ +end + +##### +##### Blending the barotropic state into the Shu-Osher combination +##### +##### Every sub-cycle restarts from (ηⁿ, Uⁿ, Vⁿ) and spans the full Δt, so its output is Ψ⋆ = Ψⁿ + ΔΨ and +##### the blend applies ΔΨ to the previous stage. Unrolled, Ψⁿ⁺¹ = Ψⁿ + Σₘ βₘ ΔΨᵐ with Σₘ βₘ = 1: the free +##### barotropic propagator acts once per step, and only the responses to the stage forcings are combined. +##### Blending the barotropic *state* instead would compose three near-exact advances and over-integrate. + +@inline previous_stage_fields(::Nothing) = NamedTuple() +@inline previous_stage_fields(G★::SSPBarotropicForcing) = (η = G★.η, U = G★.Uᵐ, V = G★.Vᵐ) + +blend_free_surface!(free_surface, timestepper, stage) = nothing + +""" +$(TYPEDSIGNATURES) + +Blend the barotropic state into the Shu-Osher combination, `Ψ ← a Ψⁿ + b (Ψᵐ⁻¹ + Ψ⋆ - Ψⁿ)`, where `Ψ⋆` is +what the sub-cycle just produced from `Ψⁿ`. + +The free surface carries the identity: the thickness follows `η`, so its increment +`b (η⋆ - ηⁿ) = -b Δt ∇·Ũ` must carry the same factor `b` as the tracer update, which is what keeps `σ` +consistent with the transport that advected the tracers. The barotropic velocity carries the order: with +the same blend, `Uᵐ` is the Shu-Osher stage value the quadrature `β` expects, so a slow forcing that +depends on the velocity -- Coriolis, drag, momentum advection of the barotropic flow -- is integrated at +the order of the composition rather than at first order. +""" +function blend_free_surface!(free_surface, timestepper::SSPRungeKuttaTimeStepper, stage) + grid = free_surface.displacement.grid + FT = eltype(grid) + a, b = timestepper.coefficients[stage] + ψ = free_surface_fields(free_surface) + launch!(architecture(grid), grid, :xy, _blend_free_surface!, ψ, timestepper.Ψ⁻[keys(ψ)], + previous_stage_fields(timestepper.G★), convert(FT, a), convert(FT, b), grid.Nz + 1) + return nothing +end + +# ψ holds the output of the solve that restarted from ψⁿ. +@kernel function _blend_free_surface!(ψ, ψⁿ, ψᵐ⁻¹, a, b, k_top) + i, j = @index(Global, NTuple) + @inbounds begin + ψ.η[i, j, k_top] = a * ψⁿ.η[i, j, k_top] + b * (ψᵐ⁻¹.η[i, j, k_top] + ψ.η[i, j, k_top] - ψⁿ.η[i, j, k_top]) + ψ.U[i, j, 1] = a * ψⁿ.U[i, j, 1] + b * (ψᵐ⁻¹.U[i, j, 1] + ψ.U[i, j, 1] - ψⁿ.U[i, j, 1]) + ψ.V[i, j, 1] = a * ψⁿ.V[i, j, 1] + b * (ψᵐ⁻¹.V[i, j, 1] + ψ.V[i, j, 1] - ψⁿ.V[i, j, 1]) + end +end diff --git a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl index d8612252cc4..035ad83a50c 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl @@ -152,7 +152,7 @@ function step_free_surface!(free_surface::ImplicitFreeSurface, model, timesteppe return nothing end -function step_free_surface!(free_surface::ImplicitFreeSurface, model, timestepper::SplitRungeKuttaTimeStepper, Δt) +function step_free_surface!(free_surface::ImplicitFreeSurface, model, timestepper::MultiStageTimeStepper, Δt) set!(free_surface.displacement, timestepper.Ψ⁻.η) step_free_surface!(free_surface, model, nothing, Δt) return nothing diff --git a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl index c1c46d80a4b..d5450eb72a9 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl @@ -123,6 +123,8 @@ free_surface_names(::SplitExplicitFreeSurface, ::PrescribedVelocityFields, grid) ab2_step_velocities!(::PrescribedVelocityFields, args...) = nothing rk_substep_velocities!(::PrescribedVelocityFields, args...) = nothing +ssp_euler_substep_velocities!(::PrescribedVelocityFields, args...) = nothing +ssp_blend_velocities!(::PrescribedVelocityFields, args...) = nothing step_free_surface!(::Nothing, model, timestepper, Δt) = nothing compute_w_from_continuity!(::PrescribedVelocityFields, args...; kwargs...) = nothing mask_immersed_horizontal_velocities!(::PrescribedVelocityFields) = nothing diff --git a/src/Models/VarianceDissipationComputations/VarianceDissipationComputations.jl b/src/Models/VarianceDissipationComputations/VarianceDissipationComputations.jl index f897fd1e40d..17fe024c621 100644 --- a/src/Models/VarianceDissipationComputations/VarianceDissipationComputations.jl +++ b/src/Models/VarianceDissipationComputations/VarianceDissipationComputations.jl @@ -4,17 +4,19 @@ export VarianceDissipation, flatten_dissipation_fields using DocStringExtensions: TYPEDSIGNATURES using KernelAbstractions: @index, @kernel -using OffsetArrays: OffsetArrays using Oceananigans: Oceananigans, UpdateStateCallsite using Oceananigans.Advection: _advective_tracer_flux_x, _advective_tracer_flux_y, _advective_tracer_flux_z -using Oceananigans.Fields: CenterField, VelocityFields, XFaceField, YFaceField, ZFaceField +using Oceananigans.Fields: CenterField, Field, VelocityFields, XFaceField, YFaceField, ZFaceField using Oceananigans.Grids: AbstractGrid, Center, Face, architecture using Oceananigans.Operators: Axᶠᶜᶜ, Ayᶜᶠᶜ, Azᶜᶜᶠ, δxᶠᶜᶜ, δyᶜᶠᶜ, δzᶜᶜᶠ, σ⁻, σ⁻¹ⁿ, σⁿ using Oceananigans.Simulations: Simulations, Callback, validate_schedule -using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper, +using Oceananigans.TimeSteppers: MultiStageTimeStepper, + SSPRungeKuttaTimeStepper, + ssp_quadrature_weights, + QuasiAdamsBashforth2TimeStepper, RungeKutta3TimeStepper, SplitRungeKuttaTimeStepper using Oceananigans.TurbulenceClosures: _diffusive_flux_x, @@ -82,8 +84,17 @@ function VarianceDissipation(tracer_name, grid; Fⁿ = c_grid_vector(grid) Fⁿ⁻¹ = c_grid_vector(grid) cⁿ⁻¹ = CenterField(grid) + cᵐ⁻¹ = CenterField(grid) - previous_state = (; cⁿ⁻¹, Uⁿ⁻¹, Uⁿ) + # stretching of the horizontal face areas at the end of the previous stage + σᵐ⁻¹ = (x = Field{Face, Center, Nothing}(grid), + y = Field{Center, Face, Nothing}(grid)) + + for σ in σᵐ⁻¹ + fill!(parent(σ), 1) + end + + previous_state = (; cⁿ⁻¹, cᵐ⁻¹, Uⁿ⁻¹, Uⁿ, σᵐ⁻¹) advective_fluxes = (; Fⁿ, Fⁿ⁻¹) diffusive_fluxes = (; Vⁿ, Vⁿ⁻¹) @@ -107,11 +118,7 @@ function (ϵ::VarianceDissipation)(model) throw(ArgumentError("Model must have a tracer called $(ϵ.tracer_name).")) end - # First we compute the dissipation from previously computed fluxes - compute_dissipation!(ϵ, model, ϵ.tracer_name) - - # Then we update the fluxes to be used in the next time step - cache_fluxes!(ϵ, model, ϵ.tracer_name) + update_variance_budget!(ϵ, model, transport_timing(model)) return nothing end diff --git a/src/Models/VarianceDissipationComputations/advective_dissipation.jl b/src/Models/VarianceDissipationComputations/advective_dissipation.jl index de47e528f41..96aa58d0137 100644 --- a/src/Models/VarianceDissipationComputations/advective_dissipation.jl +++ b/src/Models/VarianceDissipationComputations/advective_dissipation.jl @@ -1,4 +1,4 @@ -@kernel function _assemble_ab2_advective_dissipation!(P, grid, χ, Fⁿ, Fⁿ⁻¹, Uⁿ, Uⁿ⁻¹, cⁿ⁺¹, cⁿ) +@kernel function _assemble_advective_dissipation!(P, grid, C₁, C₂, Fⁿ, Fⁿ⁻¹, Uⁿ, Uⁿ⁻¹, cⁿ⁺¹, cⁿ) i, j, k = @index(Global, NTuple) δˣc★ = δxᶠᶜᶜ(i, j, k, grid, c★, cⁿ⁺¹, cⁿ) @@ -10,82 +10,43 @@ δᶻc★ = δzᶜᶜᶠ(i, j, k, grid, c★, cⁿ⁺¹, cⁿ) δᶻc² = δzᶜᶜᶠ(i, j, k, grid, c², cⁿ⁺¹, cⁿ) - C₁ = convert(eltype(grid), 1.5 + χ) - C₂ = convert(eltype(grid), 0.5 + χ) - @inbounds begin - u₁ = C₁ * Uⁿ.u[i, j, k] * σ⁻¹ⁿ(i, j, k, grid, f, c, c) - v₁ = C₁ * Uⁿ.v[i, j, k] * σ⁻¹ⁿ(i, j, k, grid, c, f, c) - w₁ = C₁ * Uⁿ.w[i, j, k] * σ⁻¹ⁿ(i, j, k, grid, c, c, f) - - u₂ = C₂ * Uⁿ⁻¹.u[i, j, k] / σ⁻(i, j, k, grid, f, c, c) - v₂ = C₂ * Uⁿ⁻¹.v[i, j, k] / σ⁻(i, j, k, grid, c, f, c) - w₂ = C₂ * Uⁿ⁻¹.w[i, j, k] / σ⁻(i, j, k, grid, c, c, f) + fx = C₁ * Fⁿ.x[i, j, k] - C₂ * Fⁿ⁻¹.x[i, j, k] + fy = C₁ * Fⁿ.y[i, j, k] - C₂ * Fⁿ⁻¹.y[i, j, k] + fz = C₁ * Fⁿ.z[i, j, k] - C₂ * Fⁿ⁻¹.z[i, j, k] - fx₁ = C₁ * Fⁿ.x[i, j, k] * σ⁻¹ⁿ(i, j, k, grid, f, c, c) - fy₁ = C₁ * Fⁿ.y[i, j, k] * σ⁻¹ⁿ(i, j, k, grid, c, f, c) - fz₁ = C₁ * Fⁿ.z[i, j, k] * σ⁻¹ⁿ(i, j, k, grid, c, c, f) + u = C₁ * Uⁿ.u[i, j, k] - C₂ * Uⁿ⁻¹.u[i, j, k] + v = C₁ * Uⁿ.v[i, j, k] - C₂ * Uⁿ⁻¹.v[i, j, k] + w = C₁ * Uⁿ.w[i, j, k] - C₂ * Uⁿ⁻¹.w[i, j, k] - fx₂ = C₂ * Fⁿ⁻¹.x[i, j, k] / σ⁻(i, j, k, grid, f, c, c) - fy₂ = C₂ * Fⁿ⁻¹.y[i, j, k] / σ⁻(i, j, k, grid, c, f, c) - fz₂ = C₂ * Fⁿ⁻¹.z[i, j, k] / σ⁻(i, j, k, grid, c, c, f) - - P.x[i, j, k] = 2 * δˣc★ * (fx₁ - fx₂) - δˣc² * (u₁ - u₂) - P.y[i, j, k] = 2 * δʸc★ * (fy₁ - fy₂) - δʸc² * (v₁ - v₂) - P.z[i, j, k] = 2 * δᶻc★ * (fz₁ - fz₂) - δᶻc² * (w₁ - w₂) + P.x[i, j, k] = 2 * δˣc★ * fx - δˣc² * u + P.y[i, j, k] = 2 * δʸc★ * fy - δʸc² * v + P.z[i, j, k] = 2 * δᶻc★ * fz - δᶻc² * w end end -@kernel function _assemble_rk3_advective_dissipation!(P, grid, Fⁿ, Uⁿ, cⁿ⁺¹, cⁿ) - i, j, k = @index(Global, NTuple) - - δˣc★ = δxᶠᶜᶜ(i, j, k, grid, c★, cⁿ⁺¹, cⁿ) - δˣc² = δxᶠᶜᶜ(i, j, k, grid, c², cⁿ⁺¹, cⁿ) - - δʸc★ = δyᶜᶠᶜ(i, j, k, grid, c★, cⁿ⁺¹, cⁿ) - δʸc² = δyᶜᶠᶜ(i, j, k, grid, c², cⁿ⁺¹, cⁿ) - - δᶻc★ = δzᶜᶜᶠ(i, j, k, grid, c★, cⁿ⁺¹, cⁿ) - δᶻc² = δzᶜᶜᶠ(i, j, k, grid, c², cⁿ⁺¹, cⁿ) +##### +##### Fluxes and transports, area-weighted, on the grid the stage computed its tendencies on. +##### The vertical face area does not stretch; the horizontal ones are taken back to the stretching `σᵐ⁻¹`. +##### - @inbounds begin - u₁ = Uⁿ.u[i, j, k] * σ⁻¹ⁿ(i, j, k, grid, f, c, c) - v₁ = Uⁿ.v[i, j, k] * σ⁻¹ⁿ(i, j, k, grid, c, f, c) - w₁ = Uⁿ.w[i, j, k] * σ⁻¹ⁿ(i, j, k, grid, c, c, f) - - fx₁ = Fⁿ.x[i, j, k] * σ⁻¹ⁿ(i, j, k, grid, f, c, c) - fy₁ = Fⁿ.y[i, j, k] * σ⁻¹ⁿ(i, j, k, grid, c, f, c) - fz₁ = Fⁿ.z[i, j, k] * σ⁻¹ⁿ(i, j, k, grid, c, c, f) - - P.x[i, j, k] = 2 * δˣc★ * fx₁ - δˣc² * u₁ - P.y[i, j, k] = 2 * δʸc★ * fy₁ - δʸc² * v₁ - P.z[i, j, k] = 2 * δᶻc★ * fz₁ - δᶻc² * w₁ - end -end - -@kernel function _cache_advective_fluxes!(Fⁿ, Fⁿ⁻¹, grid::AbstractGrid, advection, U, tracer) +@kernel function _accumulate_advective_fluxes!(F, U★, grid, advection, U, tracer, σᵐ⁻¹, weight, keep) i, j, k = @index(Global, NTuple) - @inbounds begin - # Save previous advective fluxes - Fⁿ⁻¹.x[i, j, k] = Fⁿ.x[i, j, k] - Fⁿ⁻¹.y[i, j, k] = Fⁿ.y[i, j, k] - Fⁿ⁻¹.z[i, j, k] = Fⁿ.z[i, j, k] - - # Calculate new advective fluxes - Fⁿ.x[i, j, k] = _advective_tracer_flux_x(i, j, k, grid, advection, U.u, tracer) * σⁿ(i, j, k, grid, f, c, c) - Fⁿ.y[i, j, k] = _advective_tracer_flux_y(i, j, k, grid, advection, U.v, tracer) * σⁿ(i, j, k, grid, c, f, c) - Fⁿ.z[i, j, k] = _advective_tracer_flux_z(i, j, k, grid, advection, U.w, tracer) * σⁿ(i, j, k, grid, c, c, f) - end -end + rˣ = @inbounds σᵐ⁻¹.x[i, j, 1] / σⁿ(i, j, k, grid, f, c, c) + rʸ = @inbounds σᵐ⁻¹.y[i, j, 1] / σⁿ(i, j, k, grid, c, f, c) -@kernel function _cache_advective_fluxes!(Fⁿ, grid::AbstractGrid, advection, U, tracer) - i, j, k = @index(Global, NTuple) + fx = _advective_tracer_flux_x(i, j, k, grid, advection, U.u, tracer) * rˣ + fy = _advective_tracer_flux_y(i, j, k, grid, advection, U.v, tracer) * rʸ + fz = _advective_tracer_flux_z(i, j, k, grid, advection, U.w, tracer) @inbounds begin - # Calculate new advective fluxes - Fⁿ.x[i, j, k] = _advective_tracer_flux_x(i, j, k, grid, advection, U.u, tracer) * σⁿ(i, j, k, grid, f, c, c) - Fⁿ.y[i, j, k] = _advective_tracer_flux_y(i, j, k, grid, advection, U.v, tracer) * σⁿ(i, j, k, grid, c, f, c) - Fⁿ.z[i, j, k] = _advective_tracer_flux_z(i, j, k, grid, advection, U.w, tracer) * σⁿ(i, j, k, grid, c, c, f) + F.x[i, j, k] = keep * F.x[i, j, k] + weight * fx + F.y[i, j, k] = keep * F.y[i, j, k] + weight * fy + F.z[i, j, k] = keep * F.z[i, j, k] + weight * fz + + U★.u[i, j, k] = keep * U★.u[i, j, k] + weight * U.u[i, j, k] * Axᶠᶜᶜ(i, j, k, grid) * rˣ + U★.v[i, j, k] = keep * U★.v[i, j, k] + weight * U.v[i, j, k] * Ayᶜᶠᶜ(i, j, k, grid) * rʸ + U★.w[i, j, k] = keep * U★.w[i, j, k] + weight * U.w[i, j, k] * Azᶜᶜᶠ(i, j, k, grid) end end diff --git a/src/Models/VarianceDissipationComputations/compute_dissipation.jl b/src/Models/VarianceDissipationComputations/compute_dissipation.jl index 84588ac9c1d..75a17f97116 100644 --- a/src/Models/VarianceDissipationComputations/compute_dissipation.jl +++ b/src/Models/VarianceDissipationComputations/compute_dissipation.jl @@ -4,21 +4,19 @@ """ $(TYPEDSIGNATURES) -Compute the numerical dissipation for tracer `tracer_name`, from the previously calculated advective and diffusive fluxes, -the formulation is: +Compute the numerical dissipation for tracer `tracer_name` from the cached advective and diffusive fluxes: A = 2 * δc★ * F - U δc² # For advective dissipation D = 2 * δc★ * F # For diffusive dissipation -Where ``F'' is the flux associated with the particular process,``U'' is the advecting velocity, -while ``c★'' and ``c²'' are functions defined above. -Note that ``F'' and ``U'' need to be numerically accurate for the budgets to close, -i.e. for the AB2 scheme: +where `F` is the flux associated with the particular process, `U` the transport that advects the tracer, and `c★` +and `c²` the functions defined above. `F` and `U` are those that advanced the tracer over the step: for the AB2 +scheme - F = 1.5 Fⁿ - 0.5 Fⁿ⁻¹ - U = 1.5 Uⁿ - 0.5 Uⁿ⁻¹ + F = (1.5 + χ) Fⁿ - (0.5 + χ) Fⁿ⁻¹ + U = (1.5 + χ) Uⁿ - (0.5 + χ) Uⁿ⁻¹ -For a RK3 method (not implemented at the moment), the whole substepping procedure needs to be accounted for. +and for the Runge-Kutta schemes the flux and transport of the final stage, or their quadrature over the stages. """ function compute_dissipation!(dissipation, model, tracer_name) @@ -58,12 +56,19 @@ function compute_dissipation!(dissipation, model, tracer_name) return nothing end -assemble_advective_dissipation!(P, grid, ts::QuasiAdamsBashforth2TimeStepper, substep, Fⁿ, Fⁿ⁻¹, Uⁿ, Uⁿ⁻¹, cⁿ⁺¹, cⁿ) = - launch!(architecture(grid), grid, :xyz, _assemble_ab2_advective_dissipation!, P, grid, ts.χ, Fⁿ, Fⁿ⁻¹, Uⁿ, Uⁿ⁻¹, cⁿ⁺¹, cⁿ) +function assemble_advective_dissipation!(P, grid, ts::QuasiAdamsBashforth2TimeStepper, substep, Fⁿ, Fⁿ⁻¹, Uⁿ, Uⁿ⁻¹, cⁿ⁺¹, cⁿ) + FT = eltype(grid) + C₁ = convert(FT, 1.5 + ts.χ) + C₂ = convert(FT, 0.5 + ts.χ) + launch!(architecture(grid), grid, :xyz, _assemble_advective_dissipation!, P, grid, C₁, C₂, Fⁿ, Fⁿ⁻¹, Uⁿ, Uⁿ⁻¹, cⁿ⁺¹, cⁿ) + return nothing +end -function assemble_advective_dissipation!(P, grid, ts::SplitRungeKuttaTimeStepper, substep, Fⁿ, Fⁿ⁻¹, Uⁿ, Uⁿ⁻¹, cⁿ⁺¹, cⁿ) +function assemble_advective_dissipation!(P, grid, ts::MultiStageTimeStepper, substep, Fⁿ, Fⁿ⁻¹, Uⁿ, Uⁿ⁻¹, cⁿ⁺¹, cⁿ) if substep == ts.Nstages - launch!(architecture(grid), grid, :xyz, _assemble_rk3_advective_dissipation!, P, grid, Fⁿ, Uⁿ, cⁿ⁺¹, cⁿ) + FT = eltype(grid) + launch!(architecture(grid), grid, :xyz, _assemble_advective_dissipation!, P, grid, one(FT), zero(FT), + Fⁿ, Fⁿ⁻¹, Uⁿ, Uⁿ⁻¹, cⁿ⁺¹, cⁿ) end return nothing end @@ -77,3 +82,11 @@ function assemble_diffusive_dissipation!(K, grid, ts::SplitRungeKuttaTimeStepper end return nothing end + +# As for the advective production: the β-weighted sum of the stage fluxes, at the endpoints of the step. +function assemble_diffusive_dissipation!(K, grid, ts::SSPRungeKuttaTimeStepper, substep, Vⁿ, V̄, cⁿ⁺¹, cⁿ) + if substep == ts.Nstages + launch!(architecture(grid), grid, :xyz, _assemble_rk3_diffusive_dissipation!, K, grid, V̄, cⁿ⁺¹, cⁿ) + end + return nothing +end diff --git a/src/Models/VarianceDissipationComputations/diffusive_dissipation.jl b/src/Models/VarianceDissipationComputations/diffusive_dissipation.jl index 51944c150f3..85bdee3db87 100644 --- a/src/Models/VarianceDissipationComputations/diffusive_dissipation.jl +++ b/src/Models/VarianceDissipationComputations/diffusive_dissipation.jl @@ -69,6 +69,23 @@ end compute_diffusive_fluxes!(Vⁿ, i, j, k, grid, clo, K, b, c, c_id, clk, fields) end +# `Vᵐ` is scratch for this stage's flux and `V̄` the running `Σₘ βₘ Vᵐ`. +@kernel function _accumulate_ssp_diffusive_fluxes!(V̄, Vᵐ, grid::AbstractGrid, clo, K, b, c, c_id, clk, fields, β, keep) + i, j, k = @index(Global, NTuple) + + Vᵐ.x[i, j, k] = zero(grid) + Vᵐ.y[i, j, k] = zero(grid) + Vᵐ.z[i, j, k] = zero(grid) + + compute_diffusive_fluxes!(Vᵐ, i, j, k, grid, clo, K, b, c, c_id, clk, fields) + + @inbounds begin + V̄.x[i, j, k] = keep * V̄.x[i, j, k] + β * Vᵐ.x[i, j, k] + V̄.y[i, j, k] = keep * V̄.y[i, j, k] + β * Vᵐ.y[i, j, k] + V̄.z[i, j, k] = keep * V̄.z[i, j, k] + β * Vᵐ.z[i, j, k] + end +end + # Deal with tuples of closures and closure_fields @inline compute_diffusive_fluxes!(Vⁿ, i, j, k, grid, clo::Tuple{<:Any}, K, args...) = compute_diffusive_fluxes!(Vⁿ, i, j, k, grid, clo[1], K[1], args...) diff --git a/src/Models/VarianceDissipationComputations/update_fluxes.jl b/src/Models/VarianceDissipationComputations/update_fluxes.jl index 6dc5619d16b..880663ee9f8 100644 --- a/src/Models/VarianceDissipationComputations/update_fluxes.jl +++ b/src/Models/VarianceDissipationComputations/update_fluxes.jl @@ -2,89 +2,148 @@ using Oceananigans: fields using Oceananigans.Grids: Flat, topology using Oceananigans.Models.HydrostaticFreeSurfaceModels: HydrostaticFreeSurfaceModel -# Store advective and diffusive fluxes for dissipation computation -function cache_fluxes!(dissipation, model, tracer_name) - grid = model.grid - sz = size(model.tracers[1].data) - of = OffsetArrays.offsets(model.tracers[1].data) - - params = KernelParameters(sz, of) - - Uⁿ = dissipation.previous_state.Uⁿ - Uⁿ⁻¹ = dissipation.previous_state.Uⁿ⁻¹ - U = model isa HydrostaticFreeSurfaceModel ? model.transport_velocities : model.velocities - timestepper = model.timestepper - stage = model.clock.stage - - update_transport!(Uⁿ, Uⁿ⁻¹, grid, params, timestepper, stage, U) - tracer_id = findfirst(x -> x == tracer_name, keys(model.tracers)) - cache_fluxes!(dissipation, model, tracer_name, Val(tracer_id)) +""" + StageTransport() + +The tracers of a stage are advected by a transport that the stage itself builds, after its barotropic step, from +the velocities the previous stage ended with. The flux of a stage is therefore known only at its end, from the +tracer and the grid the previous stage ended with. +""" +struct StageTransport end + +""" + NextStageTransport() + +The tracers of a stage are advected by the velocities the previous stage ended with. The flux of the next stage is +known at the end of the current one. +""" +struct NextStageTransport end + +transport_timing(model) = NextStageTransport() +transport_timing(::HydrostaticFreeSurfaceModel) = StageTransport() + +advecting_velocities(model) = model.velocities +advecting_velocities(model::HydrostaticFreeSurfaceModel) = model.transport_velocities + +function update_variance_budget!(ϵ, model, timing::StageTransport) + cache_advective_fluxes!(ϵ, model, timing) + compute_dissipation!(ϵ, model, ϵ.tracer_name) + cache_diffusive_fluxes!(ϵ, model) + store_stage_state!(ϵ, model) + return nothing +end +function update_variance_budget!(ϵ, model, timing::NextStageTransport) + compute_dissipation!(ϵ, model, ϵ.tracer_name) + store_stage_state!(ϵ, model) + cache_advective_fluxes!(ϵ, model, timing) + cache_diffusive_fluxes!(ϵ, model) return nothing end -function flux_parameters(grid) +function face_ranges(grid) Nx, Ny, Nz = size(grid) TX, TY, TZ = topology(grid) Fx = ifelse(TX == Flat, 1:1, 1:Nx+1) Fy = ifelse(TY == Flat, 1:1, 1:Ny+1) Fz = ifelse(TZ == Flat, 1:1, 1:Nz+1) - return KernelParameters(Fx, Fy, Fz) + return Fx, Fy, Fz end -function cache_fluxes!(dissipation, model, tracer_name::Symbol, tracer_id) +flux_parameters(grid) = KernelParameters(face_ranges(grid)...) - # Grab tracer properties - c = model.tracers[tracer_name] - cⁿ⁻¹ = dissipation.previous_state.cⁿ⁻¹ +function horizontal_flux_parameters(grid) + Fx, Fy, _ = face_ranges(grid) + return KernelParameters(Fx, Fy) +end - grid = model.grid - U = model.velocities - params = flux_parameters(grid) - stage = model.clock.stage - timestepper = model.timestepper +""" + stage_weights(timestepper, stage, timing, FT) + +The weight with which the flux cached at the end of `stage` enters the flux of the step, and whether the sum +already under way is kept, or `nothing` if that flux does not enter the step. +""" +stage_weights(::QuasiAdamsBashforth2TimeStepper, stage, timing, FT) = (one(FT), zero(FT)) - #### - #### Update the advective fluxes and compute gradient squared - #### +stage_weights(ts::SplitRungeKuttaTimeStepper, stage, ::StageTransport, FT) = + stage == ts.Nstages ? (one(FT), zero(FT)) : nothing - Fⁿ = dissipation.advective_fluxes.Fⁿ - Fⁿ⁻¹ = dissipation.advective_fluxes.Fⁿ⁻¹ - advection = getadvection(model.advection, tracer_name) +stage_weights(ts::SplitRungeKuttaTimeStepper, stage, ::NextStageTransport, FT) = + stage == ts.Nstages - 1 ? (one(FT), zero(FT)) : nothing - cache_advective_fluxes!(Fⁿ, Fⁿ⁻¹, grid, params, timestepper, stage, advection, U, c) +function stage_weights(ts::SSPRungeKuttaTimeStepper, stage, ::StageTransport, FT) + β = ssp_quadrature_weights(ts.coefficients) + return convert(FT, β[stage]), convert(FT, stage == 1 ? 0 : 1) +end - #### - #### Update the diffusive fluxes - #### +stage_weights(ts::SSPRungeKuttaTimeStepper, stage, ::NextStageTransport, FT) = ssp_accumulation_weights(ts, stage, FT) - Vⁿ = dissipation.diffusive_fluxes.Vⁿ - Vⁿ⁻¹ = dissipation.diffusive_fluxes.Vⁿ⁻¹ +""" + ssp_accumulation_weights(ts, stage, FT) + +The weight this stage's flux carries in the running sum `F̄ = Σₘ βₘ Fᵐ`, and whether the sum already +under way is kept or restarted. + +The callback fires at the end of stage `m`, where the state is `yᵐ`, the argument of the tendency of +stage `m+1`, hence weight `β_{m+1}`. At the final stage the assembly has already consumed the sum, so +that flux restarts it with `β₁` and seeds the next step. +""" +@inline function ssp_accumulation_weights(ts::SSPRungeKuttaTimeStepper, stage, FT) + β = ssp_quadrature_weights(ts.coefficients) + seeds = stage == ts.Nstages + w = seeds ? β[1] : β[stage+1] + return convert(FT, w), convert(FT, seeds ? 0 : 1) +end - D = model.closure_fields - B = model.buoyancy - clk = model.clock - clo = model.closure - model_fields = fields(model) +# Adams-Bashforth combines the flux of the step with the one of the previous step +store_previous_advective_fluxes!(dissipation, timestepper) = nothing - cache_diffusive_fluxes(Vⁿ, Vⁿ⁻¹, grid, params, timestepper, stage, clo, D, B, c, tracer_id, clk, model_fields) +function store_previous_advective_fluxes!(dissipation, ::QuasiAdamsBashforth2TimeStepper) + Fⁿ, Fⁿ⁻¹ = dissipation.advective_fluxes + Uⁿ, Uⁿ⁻¹ = dissipation.previous_state.Uⁿ, dissipation.previous_state.Uⁿ⁻¹ - if timestepper isa QuasiAdamsBashforth2TimeStepper - parent(cⁿ⁻¹) .= parent(c) - elseif (timestepper isa SplitRungeKuttaTimeStepper) && (stage == timestepper.Nstages) - parent(cⁿ⁻¹) .= parent(c) + for (previous, current) in zip((Fⁿ⁻¹..., Uⁿ⁻¹...), (Fⁿ..., Uⁿ...)) + parent(previous) .= parent(current) end return nothing end -cache_advective_fluxes!(Fⁿ, Fⁿ⁻¹, grid, params, ::QuasiAdamsBashforth2TimeStepper, stage, advection, U, c) = - launch!(architecture(grid), grid, params, _cache_advective_fluxes!, Fⁿ, Fⁿ⁻¹, grid, advection, U, c) +function cache_advective_fluxes!(dissipation, model, timing) + grid = model.grid + timestepper = model.timestepper + weights = stage_weights(timestepper, model.clock.stage, timing, eltype(grid)) + isnothing(weights) && return nothing -function cache_advective_fluxes!(Fⁿ, Fⁿ⁻¹, grid, params, ts::SplitRungeKuttaTimeStepper, stage, advection, U, c) - if stage == ts.Nstages-1 - launch!(architecture(grid), grid, params, _cache_advective_fluxes!, Fⁿ, grid, advection, U, c) - end + advection = getadvection(model.advection, dissipation.tracer_name) + advection isa Nothing && return nothing + + weight, keep = weights + state = dissipation.previous_state + + store_previous_advective_fluxes!(dissipation, timestepper) + + launch!(architecture(grid), grid, flux_parameters(grid), _accumulate_advective_fluxes!, + dissipation.advective_fluxes.Fⁿ, state.Uⁿ, grid, advection, advecting_velocities(model), + state.cᵐ⁻¹, state.σᵐ⁻¹, weight, keep) + + return nothing +end + +function cache_diffusive_fluxes!(dissipation, model) + grid = model.grid + params = flux_parameters(grid) + tracer_name = dissipation.tracer_name + tracer_id = findfirst(x -> x == tracer_name, keys(model.tracers)) + c = model.tracers[tracer_name] + + Vⁿ = dissipation.diffusive_fluxes.Vⁿ + Vⁿ⁻¹ = dissipation.diffusive_fluxes.Vⁿ⁻¹ + + cache_diffusive_fluxes(Vⁿ, Vⁿ⁻¹, grid, params, model.timestepper, model.clock.stage, model.closure, + model.closure_fields, model.buoyancy, c, Val(tracer_id), model.clock, fields(model)) + + return nothing end cache_diffusive_fluxes(Vⁿ, Vⁿ⁻¹, grid, params, ::QuasiAdamsBashforth2TimeStepper, stage, clo, D, B, c, tracer_id, clk, model_fields) = @@ -96,32 +155,34 @@ function cache_diffusive_fluxes(Vⁿ, Vⁿ⁻¹, grid, params, ts::SplitRungeKut end end -update_transport!(Uⁿ, Uⁿ⁻¹, grid, params, ::QuasiAdamsBashforth2TimeStepper, stage, U) = - launch!(architecture(grid), grid, params, _update_transport!, Uⁿ, Uⁿ⁻¹, grid, U) - -function update_transport!(Uⁿ, Uⁿ⁻¹, grid, params, ts::SplitRungeKuttaTimeStepper, stage, U) - if stage == ts.Nstages-1 - launch!(architecture(grid), grid, params, _update_transport!, Uⁿ, grid, U) - end +function cache_diffusive_fluxes(Vⁿ, V̄, grid, params, ts::SSPRungeKuttaTimeStepper, stage, clo, D, B, c, tracer_id, clk, model_fields) + β, keep = ssp_accumulation_weights(ts, stage, eltype(grid)) + launch!(architecture(grid), grid, params, _accumulate_ssp_diffusive_fluxes!, + V̄, Vⁿ, grid, clo, D, B, c, tracer_id, clk, model_fields, β, keep) end -@kernel function _update_transport!(Uⁿ, Uⁿ⁻¹, grid, U) - i, j, k = @index(Global, NTuple) +end_of_step(::QuasiAdamsBashforth2TimeStepper, stage) = true +end_of_step(ts::MultiStageTimeStepper, stage) = stage == ts.Nstages - @inbounds begin - Uⁿ⁻¹.u[i, j, k] = Uⁿ.u[i, j, k] - Uⁿ⁻¹.v[i, j, k] = Uⁿ.v[i, j, k] - Uⁿ⁻¹.w[i, j, k] = Uⁿ.w[i, j, k] - Uⁿ.u[i, j, k] = U.u[i, j, k] * Axᶠᶜᶜ(i, j, k, grid) - Uⁿ.v[i, j, k] = U.v[i, j, k] * Ayᶜᶠᶜ(i, j, k, grid) - Uⁿ.w[i, j, k] = U.w[i, j, k] * Azᶜᶜᶠ(i, j, k, grid) +function store_stage_state!(dissipation, model) + grid = model.grid + state = dissipation.previous_state + c = model.tracers[dissipation.tracer_name] + + parent(state.cᵐ⁻¹) .= parent(c) + launch!(architecture(grid), grid, horizontal_flux_parameters(grid), _store_horizontal_stretching!, state.σᵐ⁻¹, grid) + + if end_of_step(model.timestepper, model.clock.stage) + parent(state.cⁿ⁻¹) .= parent(c) end -end -@kernel function _update_transport!(Uⁿ, grid, U) - i, j, k = @index(Global, NTuple) + return nothing +end - @inbounds Uⁿ.u[i, j, k] = U.u[i, j, k] * Axᶠᶜᶜ(i, j, k, grid) - @inbounds Uⁿ.v[i, j, k] = U.v[i, j, k] * Ayᶜᶠᶜ(i, j, k, grid) - @inbounds Uⁿ.w[i, j, k] = U.w[i, j, k] * Azᶜᶜᶠ(i, j, k, grid) +@kernel function _store_horizontal_stretching!(σ, grid) + i, j = @index(Global, NTuple) + @inbounds begin + σ.x[i, j, 1] = σⁿ(i, j, 1, grid, f, c, c) + σ.y[i, j, 1] = σⁿ(i, j, 1, grid, c, f, c) + end end diff --git a/src/MultiRegion/cubed_sphere_grid.jl b/src/MultiRegion/cubed_sphere_grid.jl index 3070eb14d71..eb1746c9a75 100644 --- a/src/MultiRegion/cubed_sphere_grid.jl +++ b/src/MultiRegion/cubed_sphere_grid.jl @@ -466,12 +466,12 @@ radius(mrg::ConformalCubedSphereGridOfSomeKind) = first(mrg).radius Grids.grid_name(mrg::ConformalCubedSphereGridOfSomeKind) = "ConformalCubedSphereGrid" -function SplitExplicitFreeSurfaces.maybe_extend_halos(TX, TY, grid::ConformalCubedSphereGridOfSomeKind, substepping::FixedSubstepNumber) +function SplitExplicitFreeSurfaces.maybe_extend_halos(TX, TY, grid::ConformalCubedSphereGridOfSomeKind, substepping::FixedSubstepNumber; stages = 1) old_halos = halo_size(grid) - Nsubsteps = length(substepping.averaging_weights) + step_halo = stages * length(substepping.averaging_weights) + 2 - Hx = TX() isa ConnectedTopology ? max(Nsubsteps+2, old_halos[1]) : old_halos[1] - Hy = TY() isa ConnectedTopology ? max(Nsubsteps+2, old_halos[2]) : old_halos[2] + Hx = TX() isa ConnectedTopology ? max(step_halo, old_halos[1]) : old_halos[1] + Hy = TY() isa ConnectedTopology ? max(step_halo, old_halos[2]) : old_halos[2] new_halos = (Hx, Hy, old_halos[3]) diff --git a/src/MultiRegion/multi_region_grid.jl b/src/MultiRegion/multi_region_grid.jl index cd164450177..b92cc6b7ddb 100644 --- a/src/MultiRegion/multi_region_grid.jl +++ b/src/MultiRegion/multi_region_grid.jl @@ -311,11 +311,11 @@ const MRG = MultiRegionGrid @inline multiregion_split_explicit_halos(old_halos, step_halo, ::XPartition) = (max(step_halo, old_halos[1]), old_halos[2], old_halos[3]) @inline multiregion_split_explicit_halos(old_halos, step_halo, ::YPartition) = (old_halos[1], max(step_halo, old_halos[2]), old_halos[3]) -function SplitExplicitFreeSurfaces.maybe_extend_halos(TX, TY, grid::MultiRegionGrids, substepping::FixedSubstepNumber) +function SplitExplicitFreeSurfaces.maybe_extend_halos(TX, TY, grid::MultiRegionGrids, substepping::FixedSubstepNumber; stages = 1) old_halos = halo_size(grid) - Nsubsteps = length(substepping.averaging_weights) + step_halo = stages * length(substepping.averaging_weights) + 2 - new_halos = multiregion_split_explicit_halos(old_halos, Nsubsteps+2, grid.partition) + new_halos = multiregion_split_explicit_halos(old_halos, step_halo, grid.partition) if new_halos == old_halos return grid diff --git a/src/MultiRegion/multi_region_models.jl b/src/MultiRegion/multi_region_models.jl index 4ed607104b6..9ee8c0f0308 100644 --- a/src/MultiRegion/multi_region_models.jl +++ b/src/MultiRegion/multi_region_models.jl @@ -7,10 +7,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels: HydrostaticFreeSurfaceMo using Oceananigans.Fields: XFaceField, YFaceField using Oceananigans.Operators: rotation_angle using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces: SplitExplicitFreeSurfaces, - FillHaloSplitExplicit, - apply_barotropic_kernel!, - _split_explicit_barotropic_velocity!, - _split_explicit_free_surface! + FillHaloSplitExplicit using Oceananigans.TurbulenceClosures: TurbulenceClosures, VerticallyImplicitTimeDiscretization, implicit_diffusion_solver, add_closure_specific_boundary_conditions, closure_required_tracers, step_closure_prognostics! @@ -18,8 +15,6 @@ using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: FlavorOfCAT using Oceananigans.Advection: OnlySelfUpwinding, CrossAndSelfUpwinding using Oceananigans.ImmersedBoundaries: GridFittedVerticalBoundaries, PartialCellVerticalBoundaries, GridFittedBoundary using Oceananigans.Solvers: ConjugateGradientSolver -using Oceananigans.Utils: configure_kernel -using KernelAbstractions.Extras.LoopInfo: @unroll const MultiRegionModel = HydrostaticFreeSurfaceModel{<:Any, <:Any, <:AbstractArchitecture, <:Any, <:MultiRegionGrids} const CubedSphereModel = HydrostaticFreeSurfaceModel{<:Any, <:Any, <:AbstractArchitecture, <:Any, <:ConformalCubedSphereGridOfSomeKind} @@ -80,6 +75,7 @@ end getregion(fs.kernel_parameters, r), getregion(fs.substepping, r), getregion(fs.timestepper, r), + getregion(fs.slow_forcing, r), getregion(fs.boundary_transport, r)) @inline Utils.getregion(fs::SplitExplicitFreeSurface{E}, r) where {E} = @@ -90,6 +86,7 @@ end _getregion(fs.kernel_parameters, r), _getregion(fs.substepping, r), _getregion(fs.timestepper, r), + _getregion(fs.slow_forcing, r), _getregion(fs.boundary_transport, r)) @inline Utils.isregional(fs::SplitExplicitFreeSurface) = isregional(fs.displacement) @@ -177,11 +174,11 @@ Advection.WENO(mrg::MultiRegionGrid, args...; kwargs...) = construct_regionally( # The bounds-preserving limiter of a WENO scheme on a multi-region grid is a multi-region field @inline Utils.getregion(weno::WENO{N, FT, WCT}, r) where {N, FT, WCT} = WENO{N, FT, WCT}(_getregion(weno.bounds, r), _getregion(weno.buffer_scheme, r), - _getregion(weno.advecting_velocity_scheme, r), weno.time_discretization) + _getregion(weno.advecting_velocity_scheme, r), weno.time_discretization, weno.coefficients) @inline Utils._getregion(weno::WENO{N, FT, WCT}, r) where {N, FT, WCT} = WENO{N, FT, WCT}(getregion(weno.bounds, r), getregion(weno.buffer_scheme, r), - getregion(weno.advecting_velocity_scheme, r), weno.time_discretization) + getregion(weno.advecting_velocity_scheme, r), weno.time_discretization, weno.coefficients) function Advection.cell_advection_timescale(grid::MultiRegionGrids, velocities) Δt = construct_regionally(cell_advection_timescale, grid, velocities) @@ -227,44 +224,16 @@ end HydrostaticFreeSurfaceModels.default_free_surface(grid::YinYangGridOfSomeKind; gravitational_acceleration=Oceananigans.defaults.gravitational_acceleration) = SplitExplicitFreeSurface(; substeps=50, extend_halos=false, gravitational_acceleration) -SplitExplicitFreeSurfaces.maybe_extend_halos(TX, TY, grid::YinYangGridOfSomeKind, ::FixedSubstepNumber) = +SplitExplicitFreeSurfaces.maybe_extend_halos(TX, TY, grid::YinYangGridOfSomeKind, ::FixedSubstepNumber; stages = 1) = throw(ArgumentError("The SplitExplicitFreeSurface on a YinYangGrid needs extend_halos = false.")) -## Split explicit extension for complete halo filling -function SplitExplicitFreeSurfaces.iterate_split_explicit!(free_surface::FillHaloSplitExplicit, grid::CubedSphereOrYinYangGrid, GUⁿ, GVⁿ, Δτᴮ, F, clock, weights, transport_weights, ::Val{Nsubsteps}) where Nsubsteps - arch = architecture(grid) +# On panelled grids the halos are filled through the region connectivity, on whole fields rather than on +# per-region argument tuples. +SplitExplicitFreeSurfaces.barotropic_halo_arguments(::FillHaloSplitExplicit, arch, grid::CubedSphereOrYinYangGrid, + reference_fields, ::Tuple{}) = () - η = free_surface.displacement - grid = free_surface.displacement.grid - state = free_surface.filtered_state - timestepper = free_surface.timestepper - g = free_surface.gravitational_acceleration - parameters = free_surface.kernel_parameters - - # Unpack state quantities, parameters and forcing terms. - U, V = free_surface.barotropic_velocities - η̅, U̅, V̅ = state.η̅, state.U̅, state.V̅ - Ũ, Ṽ = state.Ũ, state.Ṽ - - @apply_regionally velocity_kernel!, _ = configure_kernel(arch, grid, parameters, _split_explicit_barotropic_velocity!) - @apply_regionally free_surface_kernel!, _ = configure_kernel(arch, grid, parameters, _split_explicit_free_surface!) - - U_args = (grid, Val(true), Δτᴮ, η, U, V, GUⁿ, GVⁿ, g, Ũ, Ṽ, timestepper) - η_args = (grid, Val(true), Δτᴮ, η, U, V, F, clock, η̅, U̅, V̅, timestepper) - - @unroll for substep in 1:Nsubsteps - @inbounds averaging_weight = weights[substep] - @inbounds transport_weight = transport_weights[substep] - - fill_halo_regions!(η) - @apply_regionally apply_barotropic_kernel!(velocity_kernel!, transport_weight, U_args) - - fill_halo_regions!((U, V)) - @apply_regionally apply_barotropic_kernel!(free_surface_kernel!, averaging_weight, η_args) - end - - return nothing -end +SplitExplicitFreeSurfaces.barotropic_halo_arguments(::FillHaloSplitExplicit, arch, grid::CubedSphereOrYinYangGrid, + reference_fields, fields_to_fill::Tuple) = ((fields_to_fill, reference_fields),) # The barotropic face integral behind `target_transport` is region-local, so targets are refused here. function HydrostaticFreeSurfaceModels.validate_free_surface_boundary_conditions(::SplitExplicitFreeSurfaces.SplitExplicitFreeSurface, boundary_conditions, ::MultiRegionGrids) diff --git a/src/TimeSteppers/TimeSteppers.jl b/src/TimeSteppers/TimeSteppers.jl index 8d0f8c3b449..b33e5a34c24 100644 --- a/src/TimeSteppers/TimeSteppers.jl +++ b/src/TimeSteppers/TimeSteppers.jl @@ -4,6 +4,12 @@ export QuasiAdamsBashforth2TimeStepper, RungeKutta3TimeStepper, SplitRungeKuttaTimeStepper, + ModifiedRungeKutta4TimeStepper, + SSPRungeKuttaTimeStepper, + ssp_substep!, + SSPRK3_COEFFICIENTS, + MultiStageTimeStepper, + ssp_quadrature_weights, time_step!, Clock, convert_time, @@ -21,6 +27,8 @@ using Oceananigans: Oceananigans, AbstractModel, fields, initialize!, prognostic using Oceananigans.Architectures: architecture using Oceananigans.BoundaryConditions: BoundaryConditions, cache_boundary_conditions!, initialize_boundary_conditions! using Oceananigans.Fields: Fields +using Oceananigans.Grids: Center +using Oceananigans.Operators: σ⁻¹ⁿ using Oceananigans.Units: Units using Oceananigans.Utils: @apply_regionally, AbstractTimeDiscretization, @@ -82,6 +90,7 @@ include("clock.jl") include("quasi_adams_bashforth_2.jl") include("runge_kutta_3.jl") include("split_runge_kutta.jl") +include("ssp_runge_kutta.jl") """ $(TYPEDSIGNATURES) @@ -117,6 +126,16 @@ for stages in 2:5 SplitRungeKuttaTimeStepper(args...; coefficients=$coefficients, kwargs...) end +TimeStepper(::Val{:SSPRungeKutta3}, args...; kwargs...) = + SSPRungeKuttaTimeStepper(args...; coefficients=SSPRK3_COEFFICIENTS, kwargs...) + +# The four-stage composition damped for split-explicit use; see `ModifiedRungeKutta4TimeStepper` for the choice of `a`. +TimeStepper(::Val{:ModifiedRungeKutta4}, args...; kwargs...) = + SplitRungeKuttaTimeStepper(args...; coefficients=ModifiedRungeKutta4TimeStepper().β, kwargs...) + +TimeStepper(ts::SSPRungeKuttaTimeStepper, grid, prognostic_fields; kw...) = + SSPRungeKuttaTimeStepper(grid, prognostic_fields; coefficients=ts.coefficients, kw...) + TimeStepper(ts::SplitRungeKuttaTimeStepper, grid, prognostic_fields; kw...) = SplitRungeKuttaTimeStepper(grid, prognostic_fields; coefficients=ts.β, kw...) diff --git a/src/TimeSteppers/split_runge_kutta.jl b/src/TimeSteppers/split_runge_kutta.jl index e80715bb80b..5eab95b3910 100644 --- a/src/TimeSteppers/split_runge_kutta.jl +++ b/src/TimeSteppers/split_runge_kutta.jl @@ -116,6 +116,42 @@ function SplitRungeKuttaTimeStepper(; coefficients = nothing, stages = 3) return SplitRungeKuttaTimeStepper{typeof(coefficients), Nothing, Nothing, Nothing}(length(coefficients), coefficients, nothing, nothing, nothing) end +""" + ModifiedRungeKutta4TimeStepper(; a = 0.22) + +Return the four-stage `SplitRungeKuttaTimeStepper` of Silvestri et al. (2026), with stage fractions +`γ = a Δt, Δt/3, Δt/2, Δt`, that is `β = (1/a, 3, 2, 1)`. + +The composition is third order for every `a`, which sets the free quartic coefficient `c₄ = a/6` of the +stability function `R(z) = 1 + z + z²/2 + z³/6 + c₄ z⁴`. `a = 1/4` recovers the classical fourth-order +polynomial; smaller `a` trades that order, which the split-explicit coupling caps at two anyway, for the +quartic dissipation `|R(iθ)| ≈ 1 - (1/24 - c₄) θ⁴` that absorbs the perturbation the barotropic sub-cycle +injects into the first baroclinic mode. The band closes when `1/24 - c₄ ≥ σ₀`, i.e. `a ≤ 1/4 - 6 σ₀`, with +`σ₀` the injection measured on the configuration at hand. + +Must be paired with `ProgressiveSlowForcing(ModifiedRungeKutta4TimeStepper())`. + +Examples +======== + +```jldoctest +julia> using Oceananigans.TimeSteppers + +julia> ModifiedRungeKutta4TimeStepper() +SplitRungeKuttaTimeStepper +├── stages: 4 +├── β: (4.545454545454546, 3, 2, 1) +└── implicit_solver: nothing +``` + +References +========== + +* Silvestri, S., Campin, J.-M., Wagner, G. L., Constantinou, N. C., Lee, X. K., and Ferrari, R. (2026). + A low-storage Runge-Kutta framework for nonlinear free-surface ocean models. J. Adv. Model. Earth Sy. +""" +ModifiedRungeKutta4TimeStepper(; a = 0.22) = SplitRungeKuttaTimeStepper(coefficients = (1/a, 3, 2, 1)) + """ $(TYPEDSIGNATURES) diff --git a/src/TimeSteppers/ssp_runge_kutta.jl b/src/TimeSteppers/ssp_runge_kutta.jl new file mode 100644 index 00000000000..806e1ee97f9 --- /dev/null +++ b/src/TimeSteppers/ssp_runge_kutta.jl @@ -0,0 +1,232 @@ +""" + SSPRungeKuttaTimeStepper{C, TG, PF, TI} <: AbstractTimeStepper + +Hold parameters and tendency fields for a strong-stability-preserving Runge-Kutta time-stepping scheme +in the split-explicit arrangement of [Lan et al. (2022)](@cite Lan2022). + +Fields +====== +- `Nstages`: total number of stages +- `coefficients`: Shu-Osher pairs `(a, b)`, one per stage; the stage advance is + `Ψᵐ = a Ψⁿ + b (Ψᵐ⁻¹ + Δt Gᵐ)` +- `Gⁿ`: Tendency fields at the current stage. +- `Ψ⁻`: Prognostic fields cached at the beginning of the time step (before the stages). +- `implicit_solver`: Solver for implicit time stepping of diffusion (or `nothing`). + +Unlike `SplitRungeKuttaTimeStepper`, every stage advances by the *full* `Δt` and the stages are combined +by the Shu-Osher convex weights rather than starting afresh from `Ψⁿ` with a reduced increment. +""" +struct SSPRungeKuttaTimeStepper{C, TG, PF, TI, GS} <: AbstractTimeStepper + Nstages :: Int + coefficients :: C + Gⁿ :: TG + Ψ⁻ :: PF + implicit_solver :: TI + G★ :: GS # the previous-stage (η, U, V) and the (U, V) corrector forcing when split-explicit, or `nothing` +end + +""" + ssp_quadrature_weights(coefficients) + +Return the effective quadrature weights `βₘ` on the stage tendencies, i.e. the coefficients such that +`Ψⁿ⁺¹ = Ψⁿ + Δt Σₘ βₘ Gᵐ`. For the Shu-Osher recursion `Ψᵐ = a Ψⁿ + b (Ψᵐ⁻¹ + Δt Gᵐ)` the weight of stage +`m` is the product of the `b` coefficients of stage `m` and every stage after it, `(1/6, 1/6, 2/3)` for +the three-stage scheme. +""" +@inline function ssp_quadrature_weights(coefficients) + N = length(coefficients) + return ntuple(m -> prod(coefficients[j][2] for j in m:N), N) +end + +""" + ssp_stage_fractions(coefficients) + +Return the fractions `τᵐ` of `Δt` reached by the state after each stage of the Shu-Osher recursion, +`τᵐ = bᵐ (τᵐ⁻¹ + 1)` with `τ⁰ = 0`: `(1, 1/2, 1)` for the three-stage scheme. +""" +ssp_stage_fractions(coefficients) = accumulate((τ, (a, b)) -> b * (τ + 1), coefficients; init = 0) + +""" +Auxiliary storage the split-explicit arrangement needs beyond `Gⁿ` and `Ψ⁻`: scratch copies of the previous +stage's barotropic state, since every solve restarts from `(ηⁿ, Uⁿ, Vⁿ)` whereas the blend applies its increment +to `(ηᵐ⁻¹, Uᵐ⁻¹, Vᵐ⁻¹)`, and the `(U, V)` accumulator for the stage-weighted slow forcing that drives the +corrector sub-cycle. Returns `nothing` for any other free surface. +""" +@inline function ssp_auxiliary_state(Gⁿ::NamedTuple, Ψ⁻) + (:U in keys(Gⁿ)) && (:V in keys(Gⁿ)) || return nothing + return (U = similar(Gⁿ.U), V = similar(Gⁿ.V), Uᵐ = similar(Ψ⁻.U), Vᵐ = similar(Ψ⁻.V), η = similar(Ψ⁻.η)) +end + +@inline ssp_auxiliary_state(Gⁿ, Ψ⁻) = nothing + +const SSPBarotropicForcing = NamedTuple{(:U, :V, :Uᵐ, :Vᵐ, :η)} + +# Shu-Osher pairs (a, b) for the classical three-stage strong-stability-preserving scheme. +const SSPRK3_COEFFICIENTS = ((0//1, 1//1), (3//4, 1//4), (1//3, 2//3)) + +""" + SSPRungeKuttaTimeStepper(grid, prognostic_fields, args...; + implicit_solver = nothing, + coefficients = SSPRK3_COEFFICIENTS, + Gⁿ = map(similar, prognostic_fields), + Ψ⁻ = map(similar, prognostic_fields)) + +Return a strong-stability-preserving `SSPRungeKuttaTimeStepper` on `grid`. + +The three-stage scheme is the classical Shu-Osher form, +```math +\\begin{aligned} +Ψ^{(1)} &= Ψ^n + Δt\\,G(Ψ^n) \\\\ +Ψ^{(2)} &= \\tfrac34 Ψ^n + \\tfrac14 \\left[Ψ^{(1)} + Δt\\,G(Ψ^{(1)})\\right] \\\\ +Ψ^{n+1} &= \\tfrac13 Ψ^n + \\tfrac23 \\left[Ψ^{(2)} + Δt\\,G(Ψ^{(2)})\\right] +\\end{aligned} +``` +which for a linear problem has the same amplification polynomial as the low-storage form of +`SplitRungeKuttaTimeStepper`, and therefore the same un-split stability limit. Paired with the barotropic +sub-cycling arrangement of [Lan et al. (2022)](@cite Lan2022) it has no resonance near a barotropic +Courant number of `5.7`, where the low-storage composition loses about a third of its damping. + +!!! note "The barotropic mode is blended by increment, not by state" + The Shu-Osher weights assume every stage is a forward-Euler step, whereas a sub-cycled barotropic + solve is a near-exact advance over the full `Δt`, so composing three such advances would + over-integrate the barotropic mode. The blend therefore acts on the *increment* each sub-cycle + produces, `Ψᵐ = a Ψⁿ + b (Ψᵐ⁻¹ + Ψ⋆ - Ψⁿ)`: every sub-cycle restarts from `Ψⁿ` and the weights sum + to one, so the free barotropic propagator is applied exactly once per step while the responses to + the stage forcings are combined as the Shu-Osher recursion requires. +""" +function SSPRungeKuttaTimeStepper(grid, prognostic_fields, args...; + implicit_solver::TI = nothing, + coefficients = SSPRK3_COEFFICIENTS, + Gⁿ::TG = map(similar, prognostic_fields), + Ψ⁻::PF = map(similar, prognostic_fields), + kwargs...) where {TI, TG, PF} + + !isnothing(implicit_solver) && + @warn("Implicit-explicit time-stepping with an SSP Runge-Kutta scheme is not tested.") + + G★ = ssp_auxiliary_state(Gⁿ, Ψ⁻) + C, GS = typeof(coefficients), typeof(G★) + return SSPRungeKuttaTimeStepper{C, TG, PF, TI, GS}(length(coefficients), coefficients, + Gⁿ, Ψ⁻, implicit_solver, G★) +end + +""" + SSPRungeKuttaTimeStepper(; coefficients = SSPRK3_COEFFICIENTS) + +Construct an `SSPRungeKuttaTimeStepper` from its Shu-Osher `coefficients`, defaulting to the classical +three-stage scheme. + +The fields are left empty: the result carries the `coefficients` to a model constructor, which calls +`TimeStepper` on it once the grid and prognostic fields are known. +""" +function SSPRungeKuttaTimeStepper(; coefficients = SSPRK3_COEFFICIENTS) + C = typeof(coefficients) + return SSPRungeKuttaTimeStepper{C, Nothing, Nothing, Nothing, Nothing}(length(coefficients), coefficients, + nothing, nothing, nothing, nothing) +end + +Base.summary(ts::SSPRungeKuttaTimeStepper) = "SSPRungeKuttaTimeStepper($(ts.Nstages) stages)" + +function Base.show(io::IO, ts::SSPRungeKuttaTimeStepper) + print(io, summary(ts), "\n") + print(io, "├── stages: ", ts.Nstages, "\n") + print(io, "├── coefficients: ", ts.coefficients, "\n") + print(io, "└── implicit_solver: ", summary(ts.implicit_solver)) +end + +""" +$(TYPEDSIGNATURES) + +Step forward `model` one time step `Δt` with a strong-stability-preserving Runge-Kutta scheme. + +Every stage advances by the full `Δt` and is blended with the cached state `Ψⁿ` by the Shu-Osher +coefficients. +""" +function time_step!(model::AbstractModel{<:SSPRungeKuttaTimeStepper}, Δt; callbacks=[]) + + kernel_Δt = kernel_time_step(architecture(model.grid), model.grid, Δt) + + maybe_prepare_first_time_step!(model, Δt, callbacks) + + cache_current_fields!(model) + cache_boundary_conditions!(model) + + stage_times = map(τ -> next_time(model.clock, τ * Δt), ssp_stage_fractions(model.timestepper.coefficients)) + + for (stage, (a, b)) in enumerate(model.timestepper.coefficients) + + model.clock.stage = stage + model.clock.last_stage_Δt = Δt + + ssp_substep!(model, kernel_Δt, a, b, callbacks) + step_closure_prognostics!(model, kernel_Δt) + + model.clock.time = stage_times[stage] + + if stage == model.timestepper.Nstages + model.clock.last_Δt = Δt + end + + update_state!(model, callbacks) + end + + step_lagrangian_particles!(model, kernel_Δt) + + model.clock.iteration += 1 + + # Once per time step: the particle state is not reset to tⁿ between stages, so per-stage updates would compound + update_lagrangian_particle_state!(model) + + return nothing +end + +""" +$(TYPEDSIGNATURES) + +Perform a single strong-stability-preserving Runge-Kutta stage, advancing the model state by `Δt` and +blending it with the cached state by the Shu-Osher pair `(a, b)`. + +Implemented by each model type, as `rk_substep!` is. +""" +ssp_substep!(model::AbstractModel, Δt, a, b, callbacks) = + error("ssp_substep! not implemented for $(typeof(model))") + +function maybe_prepare_first_time_step!(model::AbstractModel{<:SSPRungeKuttaTimeStepper}, Δt, callbacks) + if model.clock.iteration == 0 + model.clock.last_Δt = Δt + model.clock.last_stage_Δt = Δt + reconcile_state!(model) + initialize_boundary_conditions!(model) + update_state!(model, callbacks) + end + return nothing +end + +##### +##### The stage update: the Euler step Ψ̂ᵐ = Ψᵐ⁻¹ + Δt Gᵐ, and the Shu-Osher blend Ψᵐ = a Ψⁿ + b Ψ̂ᵐ. +##### + +@kernel function _ssp_euler_substep_field!(field, Δt, Gⁿ) + i, j, k = @index(Global, NTuple) + @inbounds field[i, j, k] = field[i, j, k] + Δt * Gⁿ[i, j, k] +end + +@kernel function _ssp_blend_field!(field, Ψ⁻, a, b) + i, j, k = @index(Global, NTuple) + @inbounds field[i, j, k] = a * Ψ⁻[i, j, k] + b * field[i, j, k] +end + +# (σc)ᵐ = a (σc)ⁿ + b (σĉ), with (σc)ⁿ cached in σc⁻ +@kernel function _ssp_blend_tracer_field!(c, grid, σc⁻, a, b) + i, j, k = @index(Global, NTuple) + σ⁻¹ᶜᶜⁿ = σ⁻¹ⁿ(i, j, k, grid, Center(), Center(), Center()) + @inbounds c[i, j, k] = a * σc⁻[i, j, k] * σ⁻¹ᶜᶜⁿ + b * c[i, j, k] +end + +""" + const MultiStageTimeStepper + +Timesteppers that advance the state through Runge-Kutta stages, as distinct from the two-level +Adams-Bashforth scheme. +""" +const MultiStageTimeStepper = Union{SplitRungeKuttaTimeStepper, SSPRungeKuttaTimeStepper} diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/time_step_catke_equation.jl b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/time_step_catke_equation.jl index 7b34e334064..55b4885b221 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/time_step_catke_equation.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/time_step_catke_equation.jl @@ -1,7 +1,8 @@ using Oceananigans: fields using Oceananigans.Grids: bottommost_active_node using Oceananigans.Operators: σ⁻, σ⁻¹ⁿ, σⁿ -using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper, SplitRungeKuttaTimeStepper, implicit_step! +using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper, SplitRungeKuttaTimeStepper, SSPRungeKuttaTimeStepper, + _ssp_blend_tracer_field!, implicit_step! get_time_step(closure::CATKEVerticalDiffusivity) = closure.tke_time_step @@ -145,6 +146,90 @@ function time_step_catke_equation!(model, ::SplitRungeKuttaTimeStepper, Δt) return nothing end +##### +##### Strong-stability-preserving stages +##### +##### `e` carries the Shu-Osher blend the other tracers get, `(σe)ᵐ = a (σe)ⁿ + b [(σe)ᵐ⁻¹ + Δt Gᵉ]`. +##### The stage advance is over the full `Δt` from `eᵐ⁻¹`, so the inner substeps are all Euler increments +##### and the blend with `(σe)ⁿ` is applied once, after them: splitting `b` across the substeps would +##### weight only the first increment. +##### + +function time_step_catke_equation!(model, timestepper::SSPRungeKuttaTimeStepper, Δt) + + # TODO: properly handle closure tuples + if model.closure isa Tuple + closure = first(model.closure) + closure_fields = first(model.closure_fields) + else + closure = model.closure + closure_fields = model.closure_fields + end + + e = model.tracers.e + arch = model.architecture + grid = model.grid + Gⁿ = timestepper.Gⁿ.e + σe⁻ = timestepper.Ψ⁻.e + + κe = closure_fields.κe + Le = closure_fields.Le + previous_velocities = closure_fields.previous_velocities + tracer_index = findfirst(k -> k == :e, keys(model.tracers)) + implicit_solver = timestepper.implicit_solver + + FT = eltype(grid) + a, b = timestepper.coefficients[model.clock.stage] + + Δτ = get_time_step(closure) + + if isnothing(Δτ) + Δτ = Δt + M = 1 + else + M = ceil(Int, Δt / Δτ) # number of substeps + Δτ = Δt / M + end + + tracers = buoyancy_tracers(model) + buoyancy = buoyancy_force(model) + + for m = 1:M + # Compute the linear implicit component of the RHS (closure_fields, L)... + launch!(arch, grid, :xyz, + compute_TKE_diffusivity!, + κe, grid, closure, + model.velocities, tracers, buoyancy, closure_fields) + + if m == 1 + # The first substep carries the thickness rescaling σᵐ⁻¹/σᵐ that the others do not need. + launch!(arch, grid, :xyz, + _ssp_substep_turbulent_kinetic_energy!, + Le, grid, closure, + model.velocities, previous_velocities, + tracers, buoyancy, closure_fields, + Δτ, Gⁿ) + else + launch!(arch, grid, :xyz, + _rk_euler_substep_turbulent_kinetic_energy!, + Le, grid, closure, + model.velocities, previous_velocities, + tracers, buoyancy, closure_fields, + Δτ, Gⁿ) + end + + implicit_step!(e, implicit_solver, closure, + closure_fields, Val(tracer_index), + model.clock, + fields(model), + Δτ) + end + + launch!(arch, grid, :xyz, _ssp_blend_tracer_field!, e, grid, σe⁻, convert(FT, a), convert(FT, b)) + + return nothing +end + const c = Center() @kernel function compute_TKE_diffusivity!(κe, grid, closure, @@ -309,6 +394,32 @@ end end end +@kernel function _ssp_substep_turbulent_kinetic_energy!(Le, grid, closure, + next_velocities, previous_velocities, + tracers, buoyancy, closure_fields, + Δτ, slow_Gⁿe) + + i, j, k = @index(Global, NTuple) + + e = tracers.e + + fast_Gⁿe = fast_tke_tendency(i, j, k, grid, Le, closure, + next_velocities, previous_velocities, + tracers, buoyancy, closure_fields) + + σᶜᶜⁿ = σⁿ(i, j, k, grid, Center(), Center(), Center()) + σᶜᶜ⁻ = σ⁻(i, j, k, grid, Center(), Center(), Center()) + active = !inactive_cell(i, j, k, grid) + + FT = eltype(grid) + Δτ = convert(FT, Δτ) + + @inbounds begin + total_Gⁿ = slow_Gⁿe[i, j, k] + fast_Gⁿe * σᶜᶜⁿ + e[i, j, k] = (σᶜᶜ⁻ * e[i, j, k] + Δτ * total_Gⁿ * active) / σᶜᶜⁿ + end +end + @inline function TurbulenceClosures.implicit_linear_coefficient(i, j, k, grid, closure::FlavorOfCATKE{<:VITD}, K, ::Val{id}, args...) where id L = K._tupled_implicit_linear_coefficients[id] return @inbounds L[i, j, k] diff --git a/test/hydrostatic_free_surface/open_boundary_conditions.jl b/test/hydrostatic_free_surface/open_boundary_conditions.jl index b494de810bf..0b5bf2f669a 100644 --- a/test/hydrostatic_free_surface/open_boundary_conditions.jl +++ b/test/hydrostatic_free_surface/open_boundary_conditions.jl @@ -4,6 +4,7 @@ using Oceananigans using Oceananigans.BoundaryConditions: GravityWaveRadiation, NormalRadiation, GravityWaveRadiationBoundaryCondition, SurfaceWaveRadiationBoundaryCondition, fill_halo_regions! using Oceananigans.BoundaryConditions: ObliqueRadiation, oblique_radiation_update, normal_radiation_update using Oceananigans.BoundaryConditions: TracerReservoir, reservoir_update +using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces: RungeKutta3Scheme using Oceananigans.Units using Oceananigans.MultiRegion: MultiRegionGrid, XPartition using Statistics: mean @@ -772,6 +773,101 @@ function test_gravity_wave_target_transport() return pinned && flat && drift && inert && wet_only && rejected end +const boundary_read_times = Float64[] + +function record_boundary_read_time(j, k, grid, clock, model_fields) + push!(boundary_read_times, clock.time) + return zero(grid) +end + +# The latest time one time step reads barotropic boundary data at, in units of Δt from tⁿ +function latest_boundary_read(timestepper; Δt = 0.5) + grid = RectilinearGrid(size = (20, 1, 1), x = (0, 1000), y = (0, 100), z = (-100, 0), + topology = ((Closed, Open), Periodic, Bounded)) + + u_bcs = FieldBoundaryConditions(east = NormalFlowBoundaryCondition(0; scheme = NormalRadiation(outflow_timescale = 100.0))) + U_bcs = FieldBoundaryConditions(grid, (Face(), Center(), nothing); + east = GravityWaveRadiationBoundaryCondition((0.0, record_boundary_read_time))) + + model = HydrostaticFreeSurfaceModel(grid; timestepper, + free_surface = SplitExplicitFreeSurface(grid; substeps = 10), + boundary_conditions = (u = u_bcs, U = U_bcs), + buoyancy = nothing, tracers = ()) + + time_step!(model, Δt) + empty!(boundary_read_times) + tⁿ = model.clock.time + time_step!(model, Δt) + + return (maximum(boundary_read_times) - tⁿ) / Δt +end + +const exterior_target_reads = Ref(0) + +function count_exterior_target_read(j, k, grid, clock, model_fields) + exterior_target_reads[] += 1 + return zero(grid) +end + +# Stepping is the only reader of a radiation target, and zero timescales send the exterior state to the target at once: +# returns the exterior tracer after one step, starting from c = 1, and the reads of the normal velocity target per step +function stepped_exterior_states(timestepper) + grid = RectilinearGrid(size = (8, 1, 1), x = (0, 1000), y = (0, 100), z = (-100, 0), topology = (Open, Periodic, Bounded)) + + relaxed = NormalRadiation(inflow_timescale = 0, outflow_timescale = 0) + open_transport = GravityWaveRadiationBoundaryCondition((0.0, 0.0)) + + boundary_conditions = (u = FieldBoundaryConditions(west = NormalFlowBoundaryCondition(count_exterior_target_read; discrete_form = true, scheme = relaxed), + east = NormalFlowBoundaryCondition(0; scheme = NormalRadiation(outflow_timescale = 100.0))), + U = FieldBoundaryConditions(grid, (Face(), Center(), nothing); west = open_transport, east = open_transport), + c = FieldBoundaryConditions(west = ValueBoundaryCondition(0; scheme = relaxed))) + + model = HydrostaticFreeSurfaceModel(grid; timestepper, boundary_conditions, buoyancy = nothing, tracers = :c, + free_surface = SplitExplicitFreeSurface(grid; substeps = 10)) + set!(model, c = 1) + time_step!(model, 0.5) + exterior_target_reads[] = 0 + time_step!(model, 0.5) + + return model.tracers.c[0, 1, 1], exterior_target_reads[] +end + +# The energy a multi-stage model leaves in the domain after radiating a surface pulse +function radiated_pulse_energy(timestepper; free_surface_keywords = (;)) + Nx, Ny, Nz = 60, 1, 1 + Lx, Ly, H = 1000.0, 100.0, 100.0 + + grid = RectilinearGrid(size = (Nx, Ny, Nz), + x = (0, Lx), y = (0, Ly), z = (-H, 0), + topology = (Open, Periodic, Bounded)) + + u_bcs = FieldBoundaryConditions(east = NormalFlowBoundaryCondition(0; scheme = NormalRadiation(outflow_timescale = 100.0)), + west = NormalFlowBoundaryCondition(0; scheme = NormalRadiation(outflow_timescale = 100.0))) + U_bcs = FieldBoundaryConditions(grid, (Face(), Center(), nothing); + east = GravityWaveRadiationBoundaryCondition((0.0, 0.0)), + west = GravityWaveRadiationBoundaryCondition((0.0, 0.0))) + η_bcs = FieldBoundaryConditions(grid, (Center(), Center(), Face()); + east = SurfaceWaveRadiationBoundaryCondition(), + west = SurfaceWaveRadiationBoundaryCondition()) + + model = HydrostaticFreeSurfaceModel(grid; timestepper, + free_surface = SplitExplicitFreeSurface(grid; substeps = 10, free_surface_keywords...), + boundary_conditions = (u = u_bcs, U = U_bcs, η = η_bcs), + buoyancy = nothing, tracers = ()) + + σ = Lx / 10 + set!(model, η = (x, y, z) -> 0.01 * exp(-(x - Lx/2)^2 / (2σ^2))) + + η = model.free_surface.displacement + E₀ = sum(interior(η) .^ 2) + + for _ in 1:100 + time_step!(model, 0.5) + end + + return sum(interior(η) .^ 2) / E₀ +end + @testset "Open Boundary Conditions for HydrostaticFreeSurfaceModel" begin @testset "GravityWaveRadiation with target_transport" begin @test test_gravity_wave_target_transport() @@ -809,6 +905,29 @@ end @test test_implicit_gravity_wave_radiation() end + @testset "SSPRungeKutta3 barotropic solves read boundary data from tⁿ at every stage" begin + # A solve starting at the time of stage 2, tⁿ + Δt, would read past tⁿ + 2Δt + @test latest_boundary_read(:SSPRungeKutta3) < 2 + end + + @testset "Multi-stage time steppers step the exterior states of open boundaries [$timestepper]" for + timestepper in (:SplitRungeKutta3, :SSPRungeKutta3) + + exterior_tracer, velocity_target_reads = stepped_exterior_states(timestepper) + @test exterior_tracer == 0 + @test velocity_target_reads == 3 + end + + @testset "SSPRungeKutta3 radiates a surface pulse" begin + @test radiated_pulse_energy(:SSPRungeKutta3) < 1e-3 + end + + @testset "RungeKutta3Scheme steps the exterior states at every stage [$timestepper]" for + timestepper in (:SplitRungeKutta3, :SSPRungeKutta3) + + @test radiated_pulse_energy(timestepper; free_surface_keywords = (; timestepper = RungeKutta3Scheme())) < 1e-3 + end + @testset "GravityWaveRadiation–SurfaceWaveRadiation default pairing" begin @test test_gravity_wave_pairing() end diff --git a/test/hydrostatic_free_surface/split_explicit_solver.jl b/test/hydrostatic_free_surface/split_explicit_solver.jl index b912865dd40..e4d0e1e1792 100644 --- a/test/hydrostatic_free_surface/split_explicit_solver.jl +++ b/test/hydrostatic_free_surface/split_explicit_solver.jl @@ -6,9 +6,17 @@ using Oceananigans.Models: materialize_free_surface using Oceananigans.Models.HydrostaticFreeSurfaceModels using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces: calculate_substeps, calculate_adaptive_settings, - constant_averaging_kernel, + ConstantAveragingKernel, SplitExplicitFreeSurface, - iterate_split_explicit! + iterate_split_explicit!, + weights_from_substeps, + LowDissipationAveragingKernel, + SymmetricTrigAveragingKernel, + WideTrigAveragingKernel, + OptimizedSymmetricAveragingKernel, + OptimizedAsymmetricAveragingKernel, + ForwardBackwardScheme, + RungeKutta3Scheme @inline noforcing(args...) = 0 @@ -35,7 +43,7 @@ clock = Clock{Float64}(time=0) velocities = VelocityFields(grid) - sefs = SplitExplicitFreeSurface(substeps = 200, averaging_kernel = constant_averaging_kernel) + sefs = SplitExplicitFreeSurface(substeps = 200, averaging_kernel = ConstantAveragingKernel()) sefs = materialize_free_surface(sefs, velocities, grid, barotropic_boundary_conditions(grid)) sefs.displacement .= 0 @@ -56,7 +64,7 @@ clock = Clock{Float64}(time=0) Nsubsteps = calculate_substeps(sefs.substepping, 1) fractional_Δt, weights, transport_weights = calculate_adaptive_settings(sefs.substepping, Nsubsteps) # barotropic time step in fraction of baroclinic step and averaging weights - iterate_split_explicit!(sefs, grid, GU, GV, Δτ, noforcing, clock, weights, transport_weights, Val(1)) + iterate_split_explicit!(sefs, grid, GU, GV, Δτ, Δτ, noforcing, clock, weights, transport_weights, nothing, 1) U_computed = Array(U.data.parent)[2:Nx+1, 2:Ny+1] U_exact = (reshape(-cos.(grid.xᶠᵃᵃ), (length(grid.xᶜᵃᵃ), 1)).+reshape(0 * grid.yᵃᶜᵃ, (1, length(grid.yᵃᶜᵃ))))[2:Nx+1, 2:Ny+1] @@ -75,7 +83,7 @@ clock = Clock{Float64}(time=0) Nt = floor(Int, T / Δτ) Δτ_end = T - Nt * Δτ - sefs = SplitExplicitFreeSurface(substeps = Nt, averaging_kernel = constant_averaging_kernel) + sefs = SplitExplicitFreeSurface(substeps = Nt, averaging_kernel = ConstantAveragingKernel()) sefs = materialize_free_surface(sefs, velocities, grid, barotropic_boundary_conditions(grid)) # set!(η, f(x, y)) @@ -93,9 +101,9 @@ clock = Clock{Float64}(time=0) weights = sefs.substepping.averaging_weights for _ in 1:Nt - iterate_split_explicit!(sefs, grid, GU, GV, Δτ, noforcing, clock, weights, weights, Val(1)) + iterate_split_explicit!(sefs, grid, GU, GV, Δτ, Δτ, noforcing, clock, weights, weights, nothing, 1) end - iterate_split_explicit!(sefs, grid, GU, GV, Δτ, noforcing, clock, weights, weights, Val(1)) + iterate_split_explicit!(sefs, grid, GU, GV, Δτ, Δτ, noforcing, clock, weights, weights, nothing, 1) U_computed = Array(deepcopy(interior(U))) η_computed = Array(deepcopy(interior(η))) @@ -108,7 +116,7 @@ clock = Clock{Float64}(time=0) @test maximum(abs.(η_computed - η_exact)) < max(100eps(FT), 1e-6) end - sefs = SplitExplicitFreeSurface(substeps = 200, averaging_kernel = constant_averaging_kernel) + sefs = SplitExplicitFreeSurface(substeps = 200, averaging_kernel = ConstantAveragingKernel()) sefs = materialize_free_surface(sefs, velocities, grid, barotropic_boundary_conditions(grid)) sefs.displacement .= 0 @@ -140,7 +148,7 @@ clock = Clock{Float64}(time=0) fractional_Δt, weights, transport_weights = calculate_adaptive_settings(sefs.substepping, Nsubsteps) # barotropic time step in fraction of baroclinic step and averaging weights for step in 1:Nsubsteps - iterate_split_explicit!(sefs, grid, GU, GV, Δτ, noforcing, clock, weights, transport_weights, Val(1)) + iterate_split_explicit!(sefs, grid, GU, GV, Δτ, Δτ, noforcing, clock, weights, transport_weights, nothing, 1) end U_computed = Array(deepcopy(interior(U))) @@ -174,7 +182,7 @@ clock = Clock{Float64}(time=0) Nt = floor(Int, T / Δτ) Δτ_end = T - Nt * Δτ - sefs = SplitExplicitFreeSurface(grid; substeps = Nt + 1, averaging_kernel = constant_averaging_kernel) + sefs = SplitExplicitFreeSurface(grid; substeps = Nt + 1, averaging_kernel = ConstantAveragingKernel()) sefs = materialize_free_surface(sefs, velocities, grid, barotropic_boundary_conditions(grid)) state = sefs.filtered_state @@ -201,9 +209,9 @@ clock = Clock{Float64}(time=0) weights = sefs.substepping.averaging_weights for i in 1:Nt - iterate_split_explicit!(sefs, grid, GU, GV, Δτ, noforcing, clock, weights, weights, Val(1)) + iterate_split_explicit!(sefs, grid, GU, GV, Δτ, Δτ, noforcing, clock, weights, weights, nothing, 1) end - iterate_split_explicit!(sefs, grid, GU, GV, Δτ, noforcing, clock, weights, weights, Val(1)) + iterate_split_explicit!(sefs, grid, GU, GV, Δτ, Δτ, noforcing, clock, weights, weights, nothing, 1) η_mean_after = mean(Array(interior(η))) @@ -250,7 +258,7 @@ clock = Clock{Float64}(time=0) end # end of testset loop @testset "extend_halos vs fill_halos consistency" begin - for arch in archs + for arch in archs, timestepper in (ForwardBackwardScheme(), RungeKutta3Scheme()) topology = (Periodic, Periodic, Bounded) Nx, Ny, Nz = 32, 32, 1 Lx = Ly = 2π @@ -265,13 +273,13 @@ end # end of testset loop Nsubsteps = 30 # Create two free surfaces: one with extended halos, one that fills halos each substep - sefs_extend = SplitExplicitFreeSurface(grid; substeps = Nsubsteps, - averaging_kernel = constant_averaging_kernel, + sefs_extend = SplitExplicitFreeSurface(grid; substeps = Nsubsteps, timestepper, + averaging_kernel = ConstantAveragingKernel(), extend_halos = true) sefs_extend = materialize_free_surface(sefs_extend, velocities, grid, barotropic_boundary_conditions(grid)) - sefs_fill = SplitExplicitFreeSurface(grid; substeps = Nsubsteps, - averaging_kernel = constant_averaging_kernel, + sefs_fill = SplitExplicitFreeSurface(grid; substeps = Nsubsteps, timestepper, + averaging_kernel = ConstantAveragingKernel(), extend_halos = false) sefs_fill = materialize_free_surface(sefs_fill, velocities, grid, barotropic_boundary_conditions(grid)) @@ -296,11 +304,11 @@ end # end of testset loop Δτ = 1.0 fractional_Δt, weights, transport_weights = calculate_adaptive_settings(sefs_extend.substepping, Nsubsteps) - iterate_split_explicit!(sefs_extend, sefs_extend.displacement.grid, GU, GV, Δτ, noforcing, clock, weights, transport_weights, Val(Nsubsteps)) + iterate_split_explicit!(sefs_extend, sefs_extend.displacement.grid, GU, GV, Δτ, Δτ, noforcing, clock, weights, transport_weights, nothing, Nsubsteps) fractional_Δt, weights, transport_weights = calculate_adaptive_settings(sefs_fill.substepping, Nsubsteps) - iterate_split_explicit!(sefs_fill, grid, GU, GV, Δτ, noforcing, clock, weights, transport_weights, Val(Nsubsteps)) + iterate_split_explicit!(sefs_fill, grid, GU, GV, Δτ, Δτ, noforcing, clock, weights, transport_weights, nothing, Nsubsteps) # Compare: both should give the same interior result η_extend = Array(interior(sefs_extend.displacement)) @@ -330,7 +338,9 @@ end Lz = 1 / Oceananigans.defaults.gravitational_acceleration bottom_height(λ, φ) = max(cosine_of_pole_distance(λ, φ, 70), cosine_of_pole_distance(λ, φ, 250)) > cosd(3) ? 0 : -Lz - for arch in archs, (fold_topology, pivot) in ((CenterFolded, UPivot), (CenterFolded, TPivot), (FaceFolded, FPivot)) + for arch in archs, timestepper in (ForwardBackwardScheme(), RungeKutta3Scheme()), + (fold_topology, pivot) in ((CenterFolded, UPivot), (CenterFolded, TPivot), (FaceFolded, FPivot)) + underlying_grid = TripolarGrid(arch; size = (16, 10, 1), z = (-Lz, 0), north_poles_latitude = 55, first_pole_longitude = 70, fold_topology, pivot) grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height)) @@ -346,12 +356,13 @@ end Δτ = 1e4 extended, filled = map((true, false)) do extend_halos - free_surface = SplitExplicitFreeSurface(grid; substeps = Nsubsteps, extend_halos, averaging_kernel = constant_averaging_kernel) + free_surface = SplitExplicitFreeSurface(grid; substeps = Nsubsteps, timestepper, extend_halos, + averaging_kernel = ConstantAveragingKernel()) free_surface = materialize_free_surface(free_surface, velocities, grid, boundary_conditions) set!(free_surface.displacement, (λ, φ, z) -> cosd(φ) * sind(λ - 20) + cosd(φ)^3 * cosd(3λ)) fractional_Δt, weights, transport_weights = calculate_adaptive_settings(free_surface.substepping, Nsubsteps) - iterate_split_explicit!(free_surface, grid, GU, GV, Δτ, noforcing, clock, weights, transport_weights, Val(Nsubsteps)) + iterate_split_explicit!(free_surface, grid, GU, GV, Δτ, Δτ, noforcing, clock, weights, transport_weights, nothing, Nsubsteps) return free_surface end @@ -362,3 +373,235 @@ end @test Array(interior(extended.filtered_state.η̅)) ≈ Array(interior(filled.filtered_state.η̅)) end end + +@testset "Averaging kernel moments" begin + kernel_moment(Δτ, w, p) = sum(w[m] * (m * Δτ - 1)^p for m in eachindex(w)) + for FT in float_types + # multiples of 8 land the wide-kernel window edges on the substep grid → exact μ₃ = 0 + for substeps in (48, 64) + tol = sqrt(eps(FT)) + for (kernel, third_order) in ((LowDissipationAveragingKernel(), false), + (SymmetricTrigAveragingKernel(), true), + (WideTrigAveragingKernel(), true), + (OptimizedSymmetricAveragingKernel(), true), + (OptimizedAsymmetricAveragingKernel(), true)) + Δτ, w, transport_weights = weights_from_substeps(FT, substeps, kernel) + @test sum(w) ≈ 1 atol=tol # μ₀ + @test kernel_moment(Δτ, w, 1) ≈ 0 atol=tol # μ₁ = 1 (barycenter on the baroclinic step) + @test kernel_moment(Δτ, w, 2) ≈ 0 atol=tol # μ₂ = 0 + third_order && @test kernel_moment(Δτ, w, 3) ≈ 0 atol=tol # μ₃ = 0 + @test sum(transport_weights) ≈ 1 atol=tol # reversed-cumsum transport ⇒ tracer constancy + end + + # widening the window deepens μ₄, i.e. more low-frequency dissipation + μ₄_symmetric = kernel_moment(weights_from_substeps(FT, substeps, SymmetricTrigAveragingKernel())[1:2]..., 4) + μ₄_wide = kernel_moment(weights_from_substeps(FT, substeps, WideTrigAveragingKernel())[1:2]..., 4) + @test μ₄_wide < μ₄_symmetric < 0 + end + + # The optimized asymmetric kernel imposes μ₀ = 1 and μ₁ = μ₂ = μ₃ = 0 directly on the substep grid + # rather than inheriting them from a continuous symmetry, so unlike the trigonometric kernels its moments + # are exact for EVERY substep count, not only where the window edges land on the grid. + for substeps in (30, 44, 50, 60) + tol = sqrt(eps(FT)) + Δτ, w, transport_weights = weights_from_substeps(FT, substeps, OptimizedAsymmetricAveragingKernel()) + @test sum(w) ≈ 1 atol=tol + @test kernel_moment(Δτ, w, 1) ≈ 0 atol=tol + @test kernel_moment(Δτ, w, 2) ≈ 0 atol=tol + @test kernel_moment(Δτ, w, 3) ≈ 0 atol=tol + @test sum(transport_weights) ≈ 1 atol=tol + end + end +end + +@testset "Slow-forcing reconstruction across the barotropic sub-cycle" begin + using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces: + dimensionless_reconstruction_weights, scale_reconstruction_weights, stage_reconstruction, + reconstruction_degree, stage_sample_times, cached_stages, + FrozenSlowForcing, ProgressiveSlowForcing, RungeKutta3Scheme + using Oceananigans.TimeSteppers: SplitRungeKuttaTimeStepper, SSPRungeKuttaTimeStepper, + ModifiedRungeKutta4TimeStepper, SSPRK3_COEFFICIENTS + + # Evaluate F(s) from the weights and the samples, in the layout the substep kernel uses: the leading + # slots are the cached stage forcings and the last is the live forcing of the stage in progress. + poly(w, samples, s) = sum(sum(w[k][j] * samples[j] for j in eachindex(samples)) * s^(k-1) for k in 1:3) + + # Lay `M` stage forcings out the way the substep kernel sees them: samples 1 … M-1 in the cache slots, + # and sample M -- the forcing of the stage in progress, which is never cached -- in the last slot. + slots(F, N) = ntuple(j -> j == N + 1 ? F[end] : (j < length(F) ? F[j] : 0.0), N + 1) + + # The shape is stored on the reconstruction and only rescaled by Δt at each stage; the tests exercise + # exactly that composition. + weights(τ, D, N, Δt) = scale_reconstruction_weights(dimensionless_reconstruction_weights(τ, D, N, Float64), Δt) + + @testset "Interpolation through the stage values" begin + Δt = 600.0 + F¹, F², F³ = 1.5, 2.25, 2.5 + + # the low-storage RK3 nodes: the quadratic passes through the three samples, at s = 0, Δt/3, Δt/2 + w = weights((0.0, 1/3, 1/2), Val(3), Val(2), Δt) + Fs = slots((F¹, F², F³), 2) + @test poly(w, Fs, 0.0) ≈ F¹ rtol=1e-12 + @test poly(w, Fs, Δt/3) ≈ F² rtol=1e-12 + @test poly(w, Fs, Δt/2) ≈ F³ rtol=1e-12 + + # a forcing that is exactly quadratic in time is reproduced away from the nodes too + q(s) = 3.0 - 0.5s + 2e-4 * s^2 + qs = slots((q(0), q(Δt/3), q(Δt/2)), 2) + for s in (0.0, 100.0, Δt/2, Δt, 2Δt) + @test poly(w, qs, s) ≈ q(s) rtol=1e-10 + end + + # the Shu-Osher nodes (0, Δt, Δt/2) must interpolate at *their* points + wS = weights((0.0, 1.0, 1/2), Val(3), Val(2), Δt) + @test poly(wS, Fs, 0.0) ≈ F¹ rtol=1e-12 + @test poly(wS, Fs, Δt) ≈ F² rtol=1e-12 + @test poly(wS, Fs, Δt/2) ≈ F³ rtol=1e-12 + + # a constant forcing must reduce to the frozen value at every s + for s in (0.0, 137.0, Δt, 2Δt) + @test poly(w, slots((7.0, 7.0, 7.0), 2), s) ≈ 7.0 rtol=1e-12 + end + end + + @testset "Progressive reconstruction on the four-stage nodes" begin + Δt = 600.0 + a = 0.22 + τ = (0.0, a, 1/3, 1/2) + q(s) = 3.0 - 0.5s + 2e-4 * s^2 + + # stage 2 is linear through two samples: exact for a linear forcing, and it must interpolate + ℓ(s) = 2.0 + 0.01s + w2 = weights(τ[1:2], Val(2), Val(3), Δt) + ℓ2 = slots((ℓ(0), ℓ(a*Δt)), 3) + @test poly(w2, ℓ2, 0.0) ≈ ℓ(0) rtol=1e-12 + @test poly(w2, ℓ2, a*Δt) ≈ ℓ(a*Δt) rtol=1e-12 + @test poly(w2, ℓ2, Δt) ≈ ℓ(Δt) rtol=1e-10 + + # the cache slots the stage does not use carry a zero weight, so their (stale) contents cannot leak in + @test all(w2[k][2] == 0 && w2[k][3] == 0 for k in 1:3) + + # stage 3 is the square quadratic through three samples + w3 = weights(τ[1:3], Val(3), Val(3), Δt) + q3 = slots((q(0), q(a*Δt), q(Δt/3)), 3) + for s in (0.0, a*Δt, Δt/3, Δt, 2Δt) + @test poly(w3, q3, s) ≈ q(s) rtol=1e-10 + end + @test all(w3[k][3] == 0 for k in 1:3) + + # stage 4 has four samples for three coefficients: a least-squares fit, which is still *exact* + # when the data are exactly quadratic + w4 = weights(τ, Val(3), Val(3), Δt) + q4 = slots((q(0), q(a*Δt), q(Δt/3), q(Δt/2)), 3) + for s in (0.0, a*Δt, Δt/3, Δt/2, Δt, 2Δt) + @test poly(w4, q4, s) ≈ q(s) rtol=1e-8 + end + + # and it reduces to the frozen value on a constant forcing + for s in (0.0, 137.0, Δt, 2Δt) + @test poly(w4, slots((7.0, 7.0, 7.0, 7.0), 3), s) ≈ 7.0 rtol=1e-10 + end + end + + @testset "Nodes follow the composition" begin + # RK3 keeps the nodes it always had + @test all(stage_sample_times(SplitRungeKuttaTimeStepper(stages=3)) .≈ (0.0, 1/3, 1/2)) + @test all(stage_sample_times(SSPRungeKuttaTimeStepper(coefficients=SSPRK3_COEFFICIENTS)) .≈ (0.0, 1.0, 1/2)) + + # MRK4's leading node *is* its damping fraction, read off the composition rather than stored + for a in (0.22, 1/4, 0.18) + ts = ModifiedRungeKutta4TimeStepper(; a) + @test ts.Nstages == 4 + @test all(stage_sample_times(ts) .≈ (0.0, a, 1/3, 1/2)) + end + + # the classical four-stage polynomial is the a = 1/4 member + @test all(ModifiedRungeKutta4TimeStepper(a = 1/4).β .≈ (4, 3, 2, 1)) + end + + @testset "Degree rules" begin + @test cached_stages(ProgressiveSlowForcing(ModifiedRungeKutta4TimeStepper())) == 3 + @test cached_stages(ProgressiveSlowForcing()) == 2 + @test all(ProgressiveSlowForcing().nodes .≈ (0.0, 1/3, 1/2)) + + # constant, linear, quadratic, quadratic + @test [reconstruction_degree(m) for m in 1:4] == [1, 2, 3, 3] + + @test_throws ArgumentError ProgressiveSlowForcing(SplitRungeKuttaTimeStepper(coefficients=(1,))) + + # the stored nodes are the composition's own + @test all(ProgressiveSlowForcing(ModifiedRungeKutta4TimeStepper(a=0.22)).nodes .≈ (0.0, 0.22, 1/3, 1/2)) + end + + @testset "Stage weights are scaled by the baroclinic step" begin + # Each stage passes its own step, Δt / βₘ for split Runge-Kutta and Δt for SSP, while the nodes are fractions of Δt + Δt = 600.0 + ℓ(s) = 2.0 + 0.01s + for ts in (ModifiedRungeKutta4TimeStepper(), SplitRungeKuttaTimeStepper(stages=3), + SSPRungeKuttaTimeStepper(coefficients=SSPRK3_COEFFICIENTS)) + sf = ProgressiveSlowForcing(ts) + τ = stage_sample_times(ts) + for m in 2:ts.Nstages + Δτ = ts isa SSPRungeKuttaTimeStepper ? Δt : Δt / ts.β[m] + _, _, w = stage_reconstruction(sf, ts, m, Δτ) + ℓs = slots(Tuple(ℓ(τⱼ * Δt) for τⱼ in τ[1:m]), cached_stages(sf)) + for s in (0.0, Δτ / 2, Δτ) + @test poly(w, ℓs, s) ≈ ℓ(s) rtol=1e-10 + end + end + end + end + + @testset "The model rejects what its time stepper cannot run" begin + grid = RectilinearGrid(size=(4, 4, 2), extent=(1000, 1000, 100), topology=(Periodic, Periodic, Bounded)) + model_with(timestepper, free_surface) = HydrostaticFreeSurfaceModel(grid; timestepper, free_surface, tracers=(), buoyancy=nothing) + + built_for_four_stages = SplitExplicitFreeSurface(grid; substeps=4, slow_forcing=ProgressiveSlowForcing(ModifiedRungeKutta4TimeStepper())) + @test_throws ArgumentError model_with(:SplitRungeKutta3, built_for_four_stages) + @test_throws ArgumentError model_with(:SSPRungeKutta3, ImplicitFreeSurface()) + + model = model_with(:SplitRungeKutta3, SplitExplicitFreeSurface(grid; substeps=4, slow_forcing=ProgressiveSlowForcing())) + time_step!(model, 10) + @test model.clock.iteration == 1 + end + + @testset "Model integration and tracer constancy" begin + underlying = RectilinearGrid(size=(16, 8, 8), x=(0, 2e5), y=(0, 1e5), + topology=(Periodic, Periodic, Bounded), + z=MutableVerticalDiscretization((-2000, 0)), halo=(4, 4, 4)) + seamount(x, y) = -2000 + 1200 * exp(-((x - 1e5)^2 + (y - 5e4)^2) / (2.5e4)^2) + grid = ImmersedBoundaryGrid(underlying, GridFittedBottom(seamount)) + + function run(slow_forcing; extend_halos=true) + fs = SplitExplicitFreeSurface(grid; substeps=24, timestepper=RungeKutta3Scheme(), slow_forcing, extend_halos) + model = HydrostaticFreeSurfaceModel(grid; free_surface=fs, buoyancy=BuoyancyTracer(), + tracers=(:b, :c), timestepper=:SplitRungeKutta3, + vertical_coordinate=ZStarCoordinate()) + bᵢ(x, y, z) = 1e-5 * z + 1e-3 * exp(-((x - 6e4)^2 + (y - 5e4)^2) / (2e4)^2) + set!(model, b=bᵢ, c=1.0, u=0.1) + for _ in 1:20 + time_step!(model, 120.0) + end + return model + end + + frozen = run(FrozenSlowForcing()) + progressive = run(ProgressiveSlowForcing()) + filled = run(FrozenSlowForcing(); extend_halos=false) + + # The reconstruction touches the momentum forcing only, never the transport that advects tracers, + # so constancy on the z★ grid must survive untouched -- for either halo strategy, since `filled` + # substeps with the plain operators over freshly filled halos rather than into the halo itself. + for model in (frozen, progressive, filled) + c = interior(model.tracers.c) + wet = c .!= 0 + @test maximum(abs, c[wet] .- 1) < 1e-12 + end + + # It must change the answer, but only at the level of a second-order scheme difference. + ηf = interior(frozen.free_surface.displacement) + ηp = interior(progressive.free_surface.displacement) + @test ηf != ηp + @test maximum(abs, ηp .- ηf) < 0.05 * maximum(abs, ηf) + end +end diff --git a/test/lagrangian_particles/tracking.jl b/test/lagrangian_particles/tracking.jl index 8bb2ecefd94..7f7a73c6476 100644 --- a/test/lagrangian_particles/tracking.jl +++ b/test/lagrangian_particles/tracking.jl @@ -79,7 +79,7 @@ end @testset "Tracked fields at final positions [HydrostaticFreeSurfaceModel, $arch, $FT, $timestepper]" for arch in archs, FT in float_types, - timestepper in (:QuasiAdamsBashforth2, :SplitRungeKutta2, :SplitRungeKutta3, :SplitRungeKutta4, :SplitRungeKutta5) + timestepper in (:QuasiAdamsBashforth2, :SplitRungeKutta2, :SplitRungeKutta3, :SplitRungeKutta4, :SplitRungeKutta5, :SSPRungeKutta3) grid = RectilinearGrid(arch, FT; size=(4, 4, 4), x=(0, 1000), y=(0, 1000), z=(-1000, 0)) concentration = CenterField(grid) @@ -88,7 +88,8 @@ end values = (x = FT[300], y = FT[300], z = FT[-500], concentration = FT[-1]) properties = StructArray(on_architecture(arch, values)) particles = LagrangianParticles(properties; tracked_fields=(; concentration)) - model = HydrostaticFreeSurfaceModel(grid; particles, timestepper, buoyancy=nothing, tracers=()) + free_surface_keywords = timestepper == :SSPRungeKutta3 ? (; free_surface = SplitExplicitFreeSurface(grid; substeps=4)) : (;) + model = HydrostaticFreeSurfaceModel(grid; particles, timestepper, buoyancy=nothing, tracers=(), free_surface_keywords...) set!(model, u = 10) time_step!(model, FT(1)) diff --git a/test/simulation/implicit_diffusion.jl b/test/simulation/implicit_diffusion.jl index 041b6fea557..acc052ff5d0 100644 --- a/test/simulation/implicit_diffusion.jl +++ b/test/simulation/implicit_diffusion.jl @@ -136,7 +136,7 @@ end @testset "Implicit Diffusion Diagnostic" begin schedules = [IterationInterval(1), IterationInterval(10), IterationInterval(100)] - timesteppers = (:QuasiAdamsBashforth2, :SplitRungeKutta2, :SplitRungeKutta3, :SplitRungeKutta5) + timesteppers = (:QuasiAdamsBashforth2, :SplitRungeKutta2, :SplitRungeKutta3, :SplitRungeKutta5, :SSPRungeKutta3) for arch in archs, schedule in schedules, timestepper in timesteppers @testset "Implicit Diffusion on $schedule schedule and $timestepper, in $dim-direction [$(typeof(arch))]" for dim in (:x, :y, :z) test_implicit_diffusion_diagnostic(arch, dim, timestepper, schedule) diff --git a/test/time_stepping/time_stepping.jl b/test/time_stepping/time_stepping.jl index 788f11f8364..d6642169273 100644 --- a/test/time_stepping/time_stepping.jl +++ b/test/time_stepping/time_stepping.jl @@ -3,7 +3,7 @@ include(joinpath(@__DIR__, "..", "setup", "dependencies_for_runtests.jl")) using Adapt: Adapt using Dates: Millisecond using TimesDates: TimeDate -using Oceananigans: UpdateStateCallsite +using Oceananigans: TracerTendencyCallsite, UpdateStateCallsite using Oceananigans.Grids: topological_tuple_length using Oceananigans.TimeSteppers: Clock, AbstractLagrangianParticles using Oceananigans.Utils: kernel_time_step @@ -30,6 +30,35 @@ end (recorder::ΔtTypeRecorder)(particles, model, Δt) = (recorder.Δt_type = typeof(Δt); nothing) +η_cosine_forcing(i, j, k, grid, clock, model_fields) = cos(2π * clock.time / 20) + +# z-star reads the same forcing for the grid motion, so a constant tracer stays constant only if the solve and the grid agree +function η_forcing_constancy_drift(arch, timestepper) + grid = RectilinearGrid(arch; size=(4, 4, 4), x=(0, 1000), y=(0, 1000), z=MutableVerticalDiscretization((-100, 0)), + topology=(Periodic, Periodic, Bounded)) + model = HydrostaticFreeSurfaceModel(grid; timestepper, buoyancy=nothing, tracers=:c, vertical_coordinate=ZStarCoordinate(), + free_surface=SplitExplicitFreeSurface(grid; substeps=20), + forcing=(; η=Forcing(η_cosine_forcing; discrete_form=true))) + set!(model, c=1) + for _ in 1:40 + time_step!(model, 0.5) + end + return maximum(abs, Array(interior(model.tracers.c)) .- 1) +end + +# η(T) = (20 / 2π) sin(2π T / 20) for a uniform forcing: no pressure gradient, so only the time quadrature matters +function η_forcing_error(arch, timestepper, Δt; stop_time = 5) + grid = RectilinearGrid(arch; size=(4, 4, 1), x=(0, 1000), y=(0, 1000), z=(-100, 0), topology=(Periodic, Periodic, Bounded)) + model = HydrostaticFreeSurfaceModel(grid; timestepper, buoyancy=nothing, tracers=(), + free_surface=SplitExplicitFreeSurface(grid; substeps=20), + forcing=(; η=Forcing(η_cosine_forcing; discrete_form=true))) + for _ in 1:round(Int, stop_time / Δt) + time_step!(model, Δt) + end + η = first(Array(interior(model.free_surface.displacement))) + return abs(η - 20 / 2π * sin(2π * stop_time / 20)) +end + # Particles that advance their own state by `clock.last_stage_Δt` in `update_lagrangian_particle_state!` mutable struct ElapsedTimeParticles <: AbstractLagrangianParticles elapsed_time :: Float64 @@ -510,6 +539,39 @@ timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3) end end + @testset "SSPRungeKutta3 integrates the η forcing with the stage quadrature [$arch]" for arch in archs + errors = [η_forcing_error(arch, :SSPRungeKutta3, Δt) for Δt in (0.2, 0.1, 0.05)] + @test all(log2(errors[i] / errors[i+1]) > 3 for i in 1:length(errors)-1) + @test η_forcing_constancy_drift(arch, :SSPRungeKutta3) < 1e-10 + end + + @testset "Multi-stage time steppers execute TracerTendencyCallsite callbacks at every stage [$arch, $timestepper]" for + arch in archs, timestepper in (:SplitRungeKutta3, :SSPRungeKutta3) + + grid = RectilinearGrid(arch; size=(4, 4, 4), extent=(1, 1, 1), topology=(Periodic, Periodic, Bounded)) + model = HydrostaticFreeSurfaceModel(grid; timestepper, tracers=:c, buoyancy=nothing, + free_surface=SplitExplicitFreeSurface(grid; substeps=4)) + simulation = Simulation(model; Δt=1e-3, stop_iteration=2, verbose=false) + executions = Ref(0) + add_callback!(simulation, sim -> (executions[] += 1; nothing), IterationInterval(1); callsite=TracerTendencyCallsite()) + run!(simulation) + @test executions[] == 6 + end + + @testset "CATKE steps with multi-stage time steppers [$arch, $timestepper]" for + arch in archs, timestepper in (:SplitRungeKutta3, :SSPRungeKutta3) + + grid = RectilinearGrid(arch; size=(4, 4, 8), extent=(1000, 1000, 100), topology=(Periodic, Periodic, Bounded)) + model = with_logger(NullLogger()) do + HydrostaticFreeSurfaceModel(grid; timestepper, tracers=:b, buoyancy=BuoyancyTracer(), + closure=CATKEVerticalDiffusivity(), + free_surface=SplitExplicitFreeSurface(grid; substeps=4)) + end + set!(model, b=(x, y, z) -> 1e-5 * z, e=1e-5) + time_step!(model, 10) + @test all(isfinite, Array(interior(model.tracers.e))) + end + @testset "SplitRungeKutta stages evaluate time-dependent forcing at the stage time" begin split_runge_kutta_timesteppers = (:SplitRungeKutta2, :SplitRungeKutta3, :SplitRungeKutta4, :SplitRungeKutta5) @@ -521,7 +583,7 @@ timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3) # If every stage sees tⁿ instead, c(t) = t (t - Δt) / 2. forcing = Forcing((x, y, z, t) -> t) - for timestepper in split_runge_kutta_timesteppers + for timestepper in (split_runge_kutta_timesteppers..., :SSPRungeKutta3) model = HydrostaticFreeSurfaceModel(grid; timestepper, tracers=:c, forcing=(; c=forcing), velocities=PrescribedVelocityFields()) for n in 1:3 @@ -537,7 +599,7 @@ timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3) # Float32 round-off (~1e-7) requires coarse steps to resolve the 4th- and 5th-order errors. Δts = FT == Float32 ? (1.0, 0.5, 0.25) : (0.2, 0.1, 0.05) - for (timestepper, order) in zip(split_runge_kutta_timesteppers, (2, 3, 4, 5)) + for (timestepper, order) in zip((split_runge_kutta_timesteppers..., :SSPRungeKutta3), (2, 3, 4, 5, 3)) errors = [time_dependent_forcing_error(arch, FT, timestepper, Δt) for Δt in Δts] observed_orders = [log2(errors[i] / errors[i+1]) for i in 1:length(Δts)-1] @test all(observed_orders .> order - 0.3) diff --git a/test/unit/adaptive_implicit_vertical_advection.jl b/test/unit/adaptive_implicit_vertical_advection.jl index bdc691e16dd..744e416f5f2 100644 --- a/test/unit/adaptive_implicit_vertical_advection.jl +++ b/test/unit/adaptive_implicit_vertical_advection.jl @@ -13,7 +13,8 @@ using Oceananigans.Advection: AdaptiveImplicitVerticalAdvection, using Oceananigans.Grids: Center, Face, znode using Oceananigans.Operators: volume, ℑzᵃᵃᶠ using Oceananigans.TimeSteppers: AdaptiveVerticallyImplicitDiscretization, ExplicitTimeDiscretization, - RungeKutta3TimeStepper, time_discretization, implicit_step!, reset! + RungeKutta3TimeStepper, time_discretization, implicit_step!, reset!, + SSPRungeKuttaTimeStepper, SplitRungeKuttaTimeStepper, QuasiAdamsBashforth2TimeStepper using Oceananigans.TurbulenceClosures: implicit_diffusion_solver, VerticallyImplicitTimeDiscretization @testset "AdaptiveVerticallyImplicitDiscretization construction" begin @@ -387,6 +388,31 @@ end @test all(isfinite, interior(scheme.bounds.limiter)) end +@testset "adaptive_advection_timestep across timesteppers" begin + grid = RectilinearGrid(CPU(), size=(1, 1, 4), extent=(1, 1, 1)) + Δt = 60.0 + + # Each SSP stage is a forward-Euler step over the full Δt, so the explicit scaling and the implicit solve + # must see the same Δt at every stage. + ssp = SSPRungeKuttaTimeStepper() + for stage in 1:3 + clock = Clock(time=0.0, last_Δt=Δt, last_stage_Δt=Δt, stage=stage) + @test adaptive_advection_timestep(ssp, clock) == Δt + end + + # The low-storage stages are Euler steps of Δt / βᵐ, and the scheme must be primed with the *next* one. + split = SplitRungeKuttaTimeStepper(stages=3) + for stage in 1:3 + clock = Clock(time=0.0, last_Δt=Δt, last_stage_Δt=Δt/split.β[stage], stage=stage) + nstage = stage < split.Nstages ? stage + 1 : 1 + @test adaptive_advection_timestep(split, clock) ≈ Δt / split.β[nstage] + end + + qab2 = QuasiAdamsBashforth2TimeStepper(grid, (; c = CenterField(grid))) + clock = Clock(time=0.0, last_Δt=Δt, last_stage_Δt=Δt, stage=1) + @test adaptive_advection_timestep(qab2, clock) == Δt +end + @testset "AIVA Δt is the Δt of the upcoming RK3 stage" begin for arch in archs, FT in float_types grid = RectilinearGrid(arch, FT, size=(1, 1, 1), extent=(1, 1, 1)) diff --git a/test/unit/stretched_weno.jl b/test/unit/stretched_weno.jl new file mode 100644 index 00000000000..3a762de6ceb --- /dev/null +++ b/test/unit/stretched_weno.jl @@ -0,0 +1,108 @@ +include(joinpath(@__DIR__, "..", "setup", "dependencies_for_runtests.jl")) + +using Oceananigans.Advection: biased_coefficients, weno_reconstruction, coeff_p, smoothness_coefficients, C★, + adapt_advection_order, materialize_advection, biased_interpolate_xᶠᵃᵃ, + biased_interpolate_xᶜᵃᵃ, LeftBias, RightBias +using Oceananigans.Utils: getregion, _getregion + +function stretched_reconstruction_error(order, stretched, Nx, bias = LeftBias) + stretching(ξ) = ξ + 7 * sin(π * ξ) / 20π + primitive(x) = -cos(2π * x) / 2π + grid = RectilinearGrid(CPU(); size = Nx, x = i -> stretching((i - 1) / Nx), halo = 6, topology = (Bounded, Flat, Flat)) + scheme = materialize_advection(WENO(Float64; order, stretched), grid) + xᶠ = xnodes(grid, Face()) + c = CenterField(grid) + set!(c, reshape([(primitive(xᶠ[i+1]) - primitive(xᶠ[i])) / (xᶠ[i+1] - xᶠ[i]) for i in 1:Nx], Nx, 1, 1)) + return maximum(abs(biased_interpolate_xᶠᵃᵃ(i, 1, 1, grid, scheme, bias, c) - sin(2π * xᶠ[i])) for i in order+1:Nx-order+1) +end + +function stretched_center_reconstruction_error(order, stretched, Nx, bias = LeftBias) + stretching(ξ) = ξ + 7 * sin(π * ξ) / 20π + primitive(x) = -cos(2π * x) / 2π + grid = RectilinearGrid(CPU(); size = Nx, x = i -> stretching((i - 1) / Nx), halo = 6, topology = (Bounded, Flat, Flat)) + scheme = materialize_advection(WENO(Float64; order, stretched), grid) + xᶜ = xnodes(grid, Center()) + u = XFaceField(grid) + for i in 2:Nx + u[i, 1, 1] = (primitive(xᶜ[i]) - primitive(xᶜ[i-1])) / (xᶜ[i] - xᶜ[i-1]) + end + return maximum(abs(biased_interpolate_xᶜᵃᵃ(i, 1, 1, grid, scheme, bias, u) - sin(2π * xᶜ[i])) for i in order+1:Nx-order+1) +end + +@testset "Stretched WENO" begin + @testset "Uniform cells give the constant coefficients [WENO$(2N - 1)]" for N in 2:5 + coefficients = biased_coefficients(Float64, N, 0:2N-1) + scheme = WENO(Float64; order = 2N - 1) + + for s in 0:N-1 + @test all(isapprox.(coefficients.p[s+1], coeff_p(scheme, Val(s)), atol = 1e-14)) + @test all(isapprox.(coefficients.β[s+1], smoothness_coefficients(Val(Float64), Val(N), Val(s)), atol = 1e-14)) + @test coefficients.C★[s+1] ≈ C★(scheme, Val(s)) atol = 1e-14 + end + end + + @testset "Polynomials on geometric cells are reconstructed exactly [WENO$(2N - 1)]" for N in (3, 4) + edges = cumsum([0.0; [1.07^m for m in 1:2N-1]]) + coefficients = biased_coefficients(Float64, N, edges) + scheme = WENO(Float64; order = 2N - 1) + + for degree in 0:2N-2 + average(v) = (edges[v+1]^(degree+1) - edges[v]^(degree+1)) / ((degree + 1) * (edges[v+1] - edges[v])) + S = Tuple(average(v) for v in 1:2N-1) + δ = Tuple(S[m+1] - S[m] for m in 1:2N-2) + @test weno_reconstruction(scheme, S[N], δ, coefficients.C★, coefficients) ≈ edges[N+1]^degree rtol = 1e-12 + end + end + + @testset "Order on a stretched grid [$bias, $(nameof(error_of))]" for bias in (LeftBias, RightBias), + error_of in (stretched_reconstruction_error, stretched_center_reconstruction_error) + uniform = [error_of(5, false, Nx, bias) for Nx in (64, 128)] + stretched = [error_of(5, true, Nx, bias) for Nx in (64, 128)] + @test log2(uniform[1] / uniform[2]) < 2.2 + @test log2(stretched[1] / stretched[2]) > 4.5 + @test stretched[2] < uniform[2] / 1000 + end + + @testset "Selected directions" begin + stretching(i) = (ξ = (i - 1) / 8; ξ + 7 * sin(π * ξ) / 20π) + grid = RectilinearGrid(CPU(); size = (8, 8, 8), x = stretching, y = stretching, z = stretching, halo = (3, 3, 3)) + tabulated(scheme) = map(!isnothing, values(materialize_advection(scheme, grid).coefficients)) + + @test tabulated(WENO(order = 5, stretched = true)) == (true, true, true) + @test tabulated(WENO(order = 5, stretched = (:z,))) == (false, false, true) + @test tabulated(WENO(order = 5, stretched = (:x, :z))) == (true, false, true) + @test isnothing(materialize_advection(WENO(order = 5), grid).coefficients) + @test_throws ArgumentError WENO(order = 5, stretched = (:w,)) + end + + @testset "Rebuilt schemes keep their coefficients" begin + stretching(i) = (ξ = (i - 1) / 8; ξ + 7 * sin(π * ξ) / 20π) + grid = RectilinearGrid(CPU(); size = (8, 8, 8), x = stretching, y = stretching, z = stretching, halo = (3, 3, 3)) + scheme = materialize_advection(WENO(order = 5, stretched = (:z,)), grid) + + @test getregion(scheme, 1).coefficients === scheme.coefficients + @test _getregion(scheme, 1).coefficients === scheme.coefficients + @test adapt_advection_order(WENO(order = 5, stretched = (:z,)), 2, grid).coefficients == (:z,) + @test isnothing(adapt_advection_order(WENO(order = 5), 2, grid).coefficients) + end + + @testset "Time stepping with every smoothness stencil" begin + stretching(i) = (ξ = (i - 1) / 12; ξ + 7 * sin(π * ξ) / 20π) + grid = RectilinearGrid(CPU(); size = (12, 12, 12), x = stretching, y = stretching, z = i -> stretching(i) - 1, halo = (6, 6, 6), + topology = (Bounded, Bounded, Bounded)) + + for momentum_advection in (WENO(order = 5, stretched = true), WENOVectorInvariant(stretched = true)) + model = HydrostaticFreeSurfaceModel(grid; momentum_advection, tracer_advection = WENO(order = 7, stretched = true), tracers = :c, + buoyancy = nothing, free_surface = ExplicitFreeSurface(gravitational_acceleration = 1e-2)) + set!(model, u = (x, y, z) -> 1e-2 * sin(2π * y), v = (x, y, z) -> 1e-2 * sin(2π * x), c = (x, y, z) -> exp(-20 * (x - 1/2)^2)) + time_step!(model, 1e-1) + @test all(isfinite, interior(model.velocities.u)) + @test maximum(abs, interior(model.tracers.c)) < 1.01 + end + end + + @testset "Summary" begin + @test summary(WENO(order = 5, stretched = true)) == "WENO{3, Float64, Nothing}(order=5, stretched)" + @test summary(WENO(order = 5)) == "WENO{3, Float64, Nothing}(order=5)" + end +end