Improving Oceananigans' split barotropic solver

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")

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

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.ζ³

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
 
 #####

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} =

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

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])