Skip to content

Commit

Permalink
Merge branch 'main' into ss/remove-allowscalar-from-tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@simone-silvestri
simone-silvestri authored Jan 23, 2025
2 parents 6e36c9e + 373647f commit 6f277a3
Show file tree
Hide file tree
Showing 11 changed files with 215 additions and 212 deletions.
2 changes: 1 addition & 1 deletion src/Advection/Advection.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,10 +28,10 @@ using DocStringExtensions
using Base: @propagate_inbounds
using Adapt
using OffsetArrays

using Oceananigans.Grids
using Oceananigans.Grids: with_halo
using Oceananigans.Architectures: architecture, CPU
using Oceananigans: supported_float_types

using Oceananigans.Operators
using Oceananigans.Operators: flux_div_xyᶜᶜᶜ, Γᶠᶠᶜ, ∂t_σ
Expand Down
14 changes: 7 additions & 7 deletions src/Advection/centered_reconstruction.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -98,16 +98,16 @@ const ACAS = AbstractCenteredAdvectionScheme
@inline biased_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme::ACAS, bias, c, args...) = symmetric_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme, c, args...)

# uniform centered reconstruction
for buffer in advection_buffers
for buffer in advection_buffers, FT in supported_float_types
@eval begin
@inline inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme::Centered{$buffer, FT, <:Nothing}, ψ, idx, loc, args...) where FT = @inbounds $(calc_reconstruction_stencil(buffer, :symmetric, :x, false))
@inline inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme::Centered{$buffer, FT, <:Nothing}, ψ::Function, idx, loc, args...) where FT = @inbounds $(calc_reconstruction_stencil(buffer, :symmetric, :x, true))
@inline inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme::Centered{$buffer, $FT, <:Nothing}, ψ, idx, loc, args...) = @inbounds $(calc_reconstruction_stencil(FT, buffer, :symmetric, :x, false))
@inline inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme::Centered{$buffer, $FT, <:Nothing}, ψ::Function, idx, loc, args...) = @inbounds $(calc_reconstruction_stencil(FT, buffer, :symmetric, :x, true))

@inline inner_symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme::Centered{$buffer, FT, XT, <:Nothing}, ψ, idx, loc, args...) where {FT, XT} = @inbounds $(calc_reconstruction_stencil(buffer, :symmetric, :y, false))
@inline inner_symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme::Centered{$buffer, FT, XT, <:Nothing}, ψ::Function, idx, loc, args...) where {FT, XT} = @inbounds $(calc_reconstruction_stencil(buffer, :symmetric, :y, true))
@inline inner_symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme::Centered{$buffer, $FT, XT, <:Nothing}, ψ, idx, loc, args...) where XT = @inbounds $(calc_reconstruction_stencil(FT, buffer, :symmetric, :y, false))
@inline inner_symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme::Centered{$buffer, $FT, XT, <:Nothing}, ψ::Function, idx, loc, args...) where XT = @inbounds $(calc_reconstruction_stencil(FT, buffer, :symmetric, :y, true))

@inline inner_symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme::Centered{$buffer, FT, XT, YT, <:Nothing}, ψ, idx, loc, args...) where {FT, XT, YT} = @inbounds $(calc_reconstruction_stencil(buffer, :symmetric, :z, false))
@inline inner_symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme::Centered{$buffer, FT, XT, YT, <:Nothing}, ψ::Function, idx, loc, args...) where {FT, XT, YT} = @inbounds $(calc_reconstruction_stencil(buffer, :symmetric, :z, true))
@inline inner_symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme::Centered{$buffer, $FT, XT, YT, <:Nothing}, ψ, idx, loc, args...) where {XT, YT} = @inbounds $(calc_reconstruction_stencil(FT, buffer, :symmetric, :z, false))
@inline inner_symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme::Centered{$buffer, $FT, XT, YT, <:Nothing}, ψ::Function, idx, loc, args...) where {XT, YT} = @inbounds $(calc_reconstruction_stencil(FT, buffer, :symmetric, :z, true))
end
end

Expand Down
96 changes: 42 additions & 54 deletions src/Advection/reconstruction_coefficients.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -97,8 +97,8 @@ Positional Arguments
On a uniform `grid`, the coefficients are independent of the `xr` and `xi` values.
"""
@inline function stencil_coefficients(i, r, xr, xi; shift = 0, op = Base.:(-), order = 3, der = nothing)
coeffs = zeros(order)
@inline function stencil_coefficients(FT, i, r, xr, xi; shift = 0, op = Base.:(-), order = 3, der = nothing)
coeffs = zeros(BigFloat, order)
@inbounds begin
for j in 0:order-1
for m in j+1:order
Expand All @@ -109,51 +109,38 @@ On a uniform `grid`, the coefficients are independent of the `xr` and `xi` value
end
end

return tuple(coeffs...)
coeffs = FT.(coeffs)[1:end-1]

return tuple(coeffs..., 1-sum(coeffs)) # Coefficients should sum to 1!
end

"""
Coefficients for uniform centered and upwind schemes
uniform_reconstruction_coefficients(FT, Val(bias), buffer)
Returns coefficients for finite volume reconstruction used in linear advection schemes (`Centered` and `UpwindBiased`).
`FT` is the floating type (e.g. `Float32`, `Float64`), `bias` is either `:symmetric`, `:left`, or `:right`,
and `buffer` is the buffer size which determines the order of the reconstruction.
symmetric coefficients are for centered reconstruction (dispersive, even order),
left and right are for upwind biased (diffusive, odd order)
examples:
julia> using Oceananigans.Advection: coeff2_symmetric, coeff3_left, coeff3_right, coeff4_symmetric, coeff5_left
```julia
julia> using Oceananigans.Advection: uniform_reconstruction_coefficients
julia> coeff2_symmetric
julia> uniform_reconstruction_coefficients(Float64, Val(:symmetric), 1)
(0.5, 0.5)
julia> coeff3_left, coeff3_right
((0.33333333333333337, 0.8333333333333334, -0.16666666666666666), (-0.16666666666666669, 0.8333333333333333, 0.3333333333333333))
julia> coeff4_symmetric
(-0.08333333333333333, 0.5833333333333333, 0.5833333333333333, -0.08333333333333333)
julia> uniform_reconstruction_coefficients(Float32, Val(:left), 3)
(-0.05f0, 0.45f0, 0.78333336f0, -0.21666667f0, 0.033333335f0)
julia> coeff5_left
(-0.049999999999999926, 0.45000000000000007, 0.7833333333333333, -0.21666666666666667, 0.03333333333333333)
julia> uniform_reconstruction_coefficients(Float16, Val(:right), 4)
(Float16(-0.00714), Float16(0.0595), Float16(-0.2405), Float16(0.76), Float16(0.51), Float16(-0.09045), Float16(0.00952))
```
"""
const coeff1_left = 1.0
const coeff1_right = 1.0

# buffer in [1:6] allows up to Centered(order = 12) and UpwindBiased(order = 11)
for buffer in advection_buffers
order_bias = 2buffer - 1
order_symm = 2buffer

coeff_symm = Symbol(:coeff, order_symm, :_symmetric)
coeff_left = Symbol(:coeff, order_bias, :_left)
coeff_right = Symbol(:coeff, order_bias, :_right)
@eval begin
const $coeff_symm = stencil_coefficients(50, $(buffer - 1), collect(1:100), collect(1:100); order = $order_symm)
if $order_bias > 1
const $coeff_left = stencil_coefficients(50, $(buffer - 2), collect(1:100), collect(1:100); order = $order_bias)
const $coeff_right = stencil_coefficients(50, $(buffer - 1), collect(1:100), collect(1:100); order = $order_bias)
end
end
end
uniform_reconstruction_coefficients(FT, ::Val{:symmetric}, buffer) = stencil_coefficients(FT, 50, buffer-1, collect(1:100), collect(1:100); order = 2buffer)
uniform_reconstruction_coefficients(FT, ::Val{:left}, buffer) = buffer==1 ? (one(FT),) : stencil_coefficients(FT, 50, buffer-2, collect(1:100), collect(1:100); order = 2buffer-1)
uniform_reconstruction_coefficients(FT, ::Val{:right}, buffer) = buffer==1 ? (one(FT),) : stencil_coefficients(FT, 50, buffer-1, collect(1:100), collect(1:100); order = 2buffer-1)

"""
calc_reconstruction_stencil(buffer, shift, dir, func::Bool = false)
calc_reconstruction_stencil(FT, buffer, shift, dir, func::Bool = false)
Stencils for reconstruction calculations (note that WENO has its own reconstruction stencils)
Expand All @@ -167,23 +154,23 @@ Examples
```jldoctest
julia> using Oceananigans.Advection: calc_reconstruction_stencil
julia> calc_reconstruction_stencil(1, :right, :x)
:(+(convert(FT, coeff1_right[1]) * ψ[i + 0, j, k]))
julia> calc_reconstruction_stencil(Float32, 1, :right, :x)
:(+(1.0f0 * ψ[i + 0, j, k]))
julia> calc_reconstruction_stencil(1, :left, :x)
:(+(convert(FT, coeff1_left[1]) * ψ[i + -1, j, k]))
julia> calc_reconstruction_stencil(Float64, 1, :left, :x)
:(+(1.0 * ψ[i + -1, j, k]))
julia> calc_reconstruction_stencil(1, :symmetric, :x)
:(convert(FT, coeff2_symmetric[2]) * ψ[i + -1, j, k] + convert(FT, coeff2_symmetric[1]) * ψ[i + 0, j, k])
julia> calc_reconstruction_stencil(Float64, 1, :symmetric, :y)
:(0.5 * ψ[i, j + -1, k] + 0.5 * ψ[i, j + 0, k])
julia> calc_reconstruction_stencil(2, :symmetric, :x)
:(convert(FT, coeff4_symmetric[4]) * ψ[i + -2, j, k] + convert(FT, coeff4_symmetric[3]) * ψ[i + -1, j, k] + convert(FT, coeff4_symmetric[2]) * ψ[i + 0, j, k] + convert(FT, coeff4_symmetric[1]) * ψ[i + 1, j, k])
julia> calc_reconstruction_stencil(Float32, 2, :symmetric, :x)
:(-0.083333254f0 * ψ[i + -2, j, k] + 0.5833333f0 * ψ[i + -1, j, k] + 0.5833333f0 * ψ[i + 0, j, k] + -0.083333336f0 * ψ[i + 1, j, k])
julia> calc_reconstruction_stencil(3, :left, :x)
:(convert(FT, coeff5_left[5]) * ψ[i + -3, j, k] + convert(FT, coeff5_left[4]) * ψ[i + -2, j, k] + convert(FT, coeff5_left[3]) * ψ[i + -1, j, k] + convert(FT, coeff5_left[2]) * ψ[i + 0, j, k] + convert(FT, coeff5_left[1]) * ψ[i + 1, j, k])
julia> calc_reconstruction_stencil(Float32, 3, :left, :x)
:(0.0333333f0 * ψ[i + -3, j, k] + -0.21666667f0 * ψ[i + -2, j, k] + 0.78333336f0 * ψ[i + -1, j, k] + 0.45f0 * ψ[i + 0, j, k] + -0.05f0 * ψ[i + 1, j, k])
```
"""
@inline function calc_reconstruction_stencil(buffer, shift, dir, func::Bool = false)
@inline function calc_reconstruction_stencil(FT, buffer, shift, dir, func::Bool = false)
N = buffer * 2
order = shift == :symmetric ? N : N - 1
if shift != :symmetric
Expand All @@ -194,21 +181,22 @@ julia> calc_reconstruction_stencil(3, :left, :x)
rng = rng .+ 1
end
stencil_full = Vector(undef, N)
coeff = Symbol(:coeff, order, :_, shift)
coeff = uniform_reconstruction_coefficients(FT, Val(shift), buffer)
for (idx, n) in enumerate(rng)
c = n - buffer - 1
C = coeff[order - idx + 1]
if func
stencil_full[idx] = dir == :x ?
:(convert(FT, $coeff[$(order - idx + 1)]) * ψ(i + $c, j, k, grid, args...)) :
:($C * ψ(i + $c, j, k, grid, args...)) :
dir == :y ?
:(convert(FT, $coeff[$(order - idx + 1)]) * ψ(i, j + $c, k, grid, args...)) :
:(convert(FT, $coeff[$(order - idx + 1)]) * ψ(i, j, k + $c, grid, args...))
:($C * ψ(i, j + $c, k, grid, args...)) :
:($C * ψ(i, j, k + $c, grid, args...))
else
stencil_full[idx] = dir == :x ?
:(convert(FT, $coeff[$(order - idx + 1)]) * ψ[i + $c, j, k]) :
:($C * ψ[i + $c, j, k]) :
dir == :y ?
:(convert(FT, $coeff[$(order - idx + 1)]) * ψ[i, j + $c, k]) :
:(convert(FT, $coeff[$(order - idx + 1)]) * ψ[i, j, k + $c])
:($C * ψ[i, j + $c, k]) :
:($C * ψ[i, j, k + $c])
end
end
return Expr(:call, :+, stencil_full...)
Expand Down Expand Up @@ -329,7 +317,7 @@ end
stencil = NTuple{order, FT}[]
@inbounds begin
for i = 0:N+1
push!(stencil, stencil_coefficients(i, r, cpu_coord, cpu_coord; order))
push!(stencil, stencil_coefficients(FT, i, r, cpu_coord, cpu_coord; order))
end
end
return OffsetArray(on_architecture(arch, stencil), -1)
Expand Down
42 changes: 21 additions & 21 deletions src/Advection/upwind_biased_reconstruction.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -111,31 +111,31 @@ const UY{N, FT} = UpwindBiased{N, FT, <:Any, <:Nothing} where {N, FT}
const UZ{N, FT} = UpwindBiased{N, FT, <:Any, <:Any, <:Nothing} where {N, FT}

# Uniform upwind biased reconstruction
for buffer in advection_buffers
for buffer in advection_buffers, FT in supported_float_types
@eval begin
@inline inner_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, ::UX{$buffer, FT}, bias, ψ, idx, loc, args...) where FT =
@inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(buffer, :left, :x, false)),
$(calc_reconstruction_stencil(buffer, :right, :x, false)))
@inline inner_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, ::UX{$buffer, $FT}, bias, ψ, idx, loc, args...) =
@inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(FT, buffer, :left, :x, false)),
$(calc_reconstruction_stencil(FT, buffer, :right, :x, false)))

@inline inner_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, ::UX{$buffer, FT}, bias, ψ::Function, idx, loc, args...) where FT =
@inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(buffer, :left, :x, true)),
$(calc_reconstruction_stencil(buffer, :right, :x, true)))
@inline inner_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, ::UX{$buffer, $FT}, bias, ψ::Function, idx, loc, args...) =
@inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(FT, buffer, :left, :x, true)),
$(calc_reconstruction_stencil(FT, buffer, :right, :x, true)))

@inline inner_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, ::UY{$buffer, FT}, bias, ψ, idx, loc, args...) where FT =
@inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(buffer, :left, :y, false)),
$(calc_reconstruction_stencil(buffer, :right, :y, false)))
@inline inner_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, ::UY{$buffer, $FT}, bias, ψ, idx, loc, args...) =
@inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(FT, buffer, :left, :y, false)),
$(calc_reconstruction_stencil(FT, buffer, :right, :y, false)))

@inline inner_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, ::UY{$buffer, FT}, bias, ψ::Function, idx, loc, args...) where FT =
@inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(buffer, :left, :y, true)),
$(calc_reconstruction_stencil(buffer, :right, :y, true)))
@inline inner_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, ::UY{$buffer, $FT}, bias, ψ::Function, idx, loc, args...) =
@inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(FT, buffer, :left, :y, true)),
$(calc_reconstruction_stencil(FT, buffer, :right, :y, true)))

@inline inner_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, ::UZ{$buffer, FT}, bias, ψ, idx, loc, args...) where FT =
@inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(buffer, :left, :z, false)),
$(calc_reconstruction_stencil(buffer, :right, :z, false)))
@inline inner_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, ::UZ{$buffer, $FT}, bias, ψ, idx, loc, args...) =
@inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(FT, buffer, :left, :z, false)),
$(calc_reconstruction_stencil(FT, buffer, :right, :z, false)))

@inline inner_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, ::UZ{$buffer, FT}, bias, ψ::Function, idx, loc, args...) where FT =
@inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(buffer, :left, :z, true)),
$(calc_reconstruction_stencil(buffer, :right, :z, true)))
@inline inner_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, ::UZ{$buffer, $FT}, bias, ψ::Function, idx, loc, args...) =
@inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(FT, buffer, :left, :z, true)),
$(calc_reconstruction_stencil(FT, buffer, :right, :z, true)))
end
end

Expand All @@ -147,11 +147,11 @@ for (dir, ξ, val) in zip((:xᶠᵃᵃ, :yᵃᶠᵃ, :zᵃᵃᶠ), (:x, :y, :z),
@eval begin
@inline $stencil(i, j, k, grid, scheme::UpwindBiased{$buffer, FT}, bias, ψ, idx, loc, args...) where FT =
@inbounds ifelse(bias isa LeftBias, sum($(reconstruction_stencil(buffer, :left, ξ, false)) .* retrieve_coeff(scheme, Val(1), Val($val), idx, loc)),
sum($(reconstruction_stencil(buffer, :right, ξ, false)) .* retrieve_coeff(scheme, Val(2), Val($val), idx, loc)))
sum($(reconstruction_stencil(buffer, :right, ξ, false)) .* retrieve_coeff(scheme, Val(2), Val($val), idx, loc)))

@inline $stencil(i, j, k, grid, scheme::UpwindBiased{$buffer, FT}, bias, ψ::Function, idx, loc, args...) where FT =
@inbounds ifelse(bias isa LeftBias, sum($(reconstruction_stencil(buffer, :left, ξ, true)) .* retrieve_coeff(scheme, Val(1), Val($val), idx, loc)),
sum($(reconstruction_stencil(buffer, :right, ξ, true)) .* retrieve_coeff(scheme, Val(2), Val($val), idx, loc)))
sum($(reconstruction_stencil(buffer, :right, ξ, true)) .* retrieve_coeff(scheme, Val(2), Val($val), idx, loc)))
end
end
end
Expand Down
Loading

0 comments on commit 6f277a3

Please sign in to comment.