Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Remove all hardcoded Float64 from the Advection module #4053

Merged
merged 8 commits into from
Jan 23, 2025
Merged

Remove all hardcoded Float64 from the Advection module #4053

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
521fcc7
remove all float64
simone-silvestri Jan 20, 2025
bb77bff
small fix
simone-silvestri Jan 20, 2025
5e7eaeb
disambiguate
simone-silvestri Jan 21, 2025
b94872e
add a sum to one
simone-silvestri Jan 21, 2025
ba58d8c
compute coefficients in high precision
simone-silvestri Jan 21, 2025
ca13e92
Fix doctests
simone-silvestri Jan 21, 2025
75f485c
implement suggestions
simone-silvestri Jan 22, 2025
ab8132d
Merge branch 'ss/fix-float-32' of github.com:CliMA/Oceananigans.jl in…
simone-silvestri Jan 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/Advection/Advection.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,9 @@ abstract type AbstractUpwindBiasedAdvectionScheme{B, FT} <: AbstractAdvectionSch
# Note that it is not possible to compile schemes for `advection_buffer = 41` or higher.
const advection_buffers = [1, 2, 3, 4, 5, 6]

# To add support for new floating types add them here
const supported_float_types = [Float32, Float64]
simone-silvestri marked this conversation as resolved.
Show resolved Hide resolved

@inline Base.eltype(::AbstractAdvectionScheme{<:Any, FT}) where FT = FT

@inline required_halo_size_x(::AbstractAdvectionScheme{B}) where B = B
Expand Down
14 changes: 7 additions & 7 deletions src/Advection/centered_reconstruction.jl
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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
Loading