CliMA · simone-silvestri · Jul 25, 2022 · Jun 8, 2022 · Jun 13, 2022 · Jun 13, 2022
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.76.7"
+version = "0.77.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/README.md b/README.md
@@ -95,7 +95,7 @@ Let's run a two-dimensional, horizontally-periodic simulation of turbulence usin
 ```julia
 using Oceananigans
 grid = RectilinearGrid(CPU(), size=(128, 128), x=(0, 2π), y=(0, 2π), topology=(Periodic, Periodic, Flat))
-model = NonhydrostaticModel(; grid, advection=WENO5())
+model = NonhydrostaticModel(; grid, advection=WENO())
 ϵ(x, y, z) = 2rand() - 1
 set!(model, u=ϵ, v=ϵ)
 simulation = Simulation(model; Δt=0.01, stop_time=4)

diff --git a/benchmark/benchmark_advection_schemes.jl b/benchmark/benchmark_advection_schemes.jl
@@ -24,7 +24,7 @@ end
 # Benchmark parameters
 
 Architectures = has_cuda() ? [CPU, GPU] : [CPU]
-Schemes = (CenteredSecondOrder, CenteredFourthOrder, UpwindBiasedThirdOrder, UpwindBiasedFifthOrder, WENO5)
+Schemes = (CenteredSecondOrder, CenteredFourthOrder, UpwindBiasedThirdOrder, UpwindBiasedFifthOrder, WENO)
 
 # Run and summarize benchmarks
 

diff --git a/benchmark/benchmark_multi_GPU.jl b/benchmark/benchmark_multi_GPU.jl
@@ -70,7 +70,7 @@ function run_solid_body_rotation(; architecture = CPU(),
                                         free_surface = free_surface,
                                         coriolis = coriolis,
                                         tracers = (:T, :b),
-                                        tracer_advection = WENO5(),
+                                        tracer_advection = WENO(),
                                         buoyancy = BuoyancyTracer(),                                        
                                         closure = closure)
 

diff --git a/docs/src/appendix/benchmarks.md b/docs/src/appendix/benchmarks.md
@@ -12,7 +12,7 @@ format the benchmark results.
 ## Shallow Water Model
 
 This benchmark tests the performance of the shallow water model run in a doubly periodic domain (`topology = (Periodic, Periodic, Flat)`)
-on a CPU versus a GPU.  We find that with the `WENO5` advection scheme we get a maximum speedup of more than 400 times on a `16384^2` grid.
+on a CPU versus a GPU.  We find that with the `WENO` advection scheme we get a maximum speedup of more than 400 times on a `16384^2` grid.
 
 ```
 Oceananigans v0.58.1
@@ -86,7 +86,7 @@ The time graph below shows that execution times on GPU are negligebly small up u
 
 ## Nonhydrostatic Model
 
-Similar to to shallow water model, the nonhydrostatic model benchmark tests for its performance on both a CPU and a GPU. It was also benchmarked with the `WENO5` advection scheme. The nonhydrostatic model is 3-dimensional unlike the 2-dimensional shallow water model. Total number of grid points is Ns cubed.
+Similar to to shallow water model, the nonhydrostatic model benchmark tests for its performance on both a CPU and a GPU. It was also benchmarked with the `WENO` advection scheme. The nonhydrostatic model is 3-dimensional unlike the 2-dimensional shallow water model. Total number of grid points is Ns cubed.
 
 ```
 Oceananigans v0.58.8

diff --git a/docs/src/model_setup/forcing_functions.md b/docs/src/model_setup/forcing_functions.md
@@ -303,7 +303,7 @@ w_sediment = 2/9 * Δb / ν_molecular * r_sediment^2 # m s⁻¹
 sinking = AdvectiveForcing(UpwindBiasedFifthOrder(), w=w_sediment)
 
 # output
-AdvectiveForcing with the UpwindBiasedFifthOrder scheme:
+AdvectiveForcing with the UpwindBiased scheme:
 ├── u: ZeroField{Int64}
 ├── v: ZeroField{Int64}
 └── w: ConstantField(-0.00352102)
@@ -328,10 +328,10 @@ slip_bcs = FieldBoundaryConditions(grid, (Center, Center, Face),
                                    top=no_penetration, bottom=no_penetration)
 
 w_slip = ZFaceField(grid, boundary_conditions=slip_bcs)
-sinking = AdvectiveForcing(WENO5(; grid), w=w_slip)
+sinking = AdvectiveForcing(WENO(; grid), w=w_slip)
 
 # output
-AdvectiveForcing with the WENO5 scheme:
+AdvectiveForcing with the WENO scheme:
 ├── u: ZeroField{Int64}
 ├── v: ZeroField{Int64}
 └── w: 32×32×33 Field{Center, Center, Face} on RectilinearGrid on CPU

diff --git a/docs/src/quick_start.md b/docs/src/quick_start.md
@@ -6,7 +6,7 @@ This code:
 using Oceananigans
 
 grid = RectilinearGrid(size=(128, 128), x=(0, 2π), y=(0, 2π), topology=(Periodic, Periodic, Flat))
-model = NonhydrostaticModel(; grid, advection=WENO5())
+model = NonhydrostaticModel(; grid, advection=WENO())
 
 ϵ(x, y, z) = 2rand() - 1
 set!(model, u=ϵ, v=ϵ)
@@ -50,7 +50,7 @@ Fine, we'll re-run this code on the GPU then:
 using Oceananigans
 
 grid = RectilinearGrid(GPU(), size=(128, 128), x=(0, 2π), y=(0, 2π), topology=(Periodic, Periodic, Flat))
-model = NonhydrostaticModel(; grid, advection=WENO5())
+model = NonhydrostaticModel(; grid, advection=WENO())
 
 ϵ(x, y, z) = 2rand() - 1
 set!(model, u=ϵ, v=ϵ)

diff --git a/docs/src/simulation_tips.md b/docs/src/simulation_tips.md
@@ -185,7 +185,7 @@ For large simulations on the GPU, careful management of memory allocation may be
 
 - Try to use higher-order advection schemes. In general when you use a higher-order scheme you need
   fewer grid points to achieve the same accuracy that you would with a lower-order one. Oceananigans
-  provides two high-order advection schemes: 5th-order WENO method (WENO5) and 3rd-order upwind.
+  provides two high-order advection schemes: 5th-order WENO method (WENO) and 3rd-order upwind.
 
 - Manually define scratch space to be reused in diagnostics. By default, every time a user-defined
   diagnostic is calculated the compiler reserves a new chunk of memory for that calculation, usually

diff --git a/examples/baroclinic_adjustment.jl b/examples/baroclinic_adjustment.jl
@@ -63,8 +63,8 @@ model = HydrostaticFreeSurfaceModel(; grid,
                                     buoyancy = BuoyancyTracer(),
                                     tracers = :b,
                                     closure = (vertical_diffusive_closure, horizontal_diffusive_closure),
-                                    momentum_advection = WENO5(),
-                                    tracer_advection = WENO5(),
+                                    momentum_advection = WENO(),
+                                    tracer_advection = WENO(),
                                     free_surface = ImplicitFreeSurface())
 
 # We want to initialize our model with a baroclinically unstable front plus some small-amplitude

diff --git a/examples/horizontal_convection.jl b/examples/horizontal_convection.jl
@@ -91,7 +91,7 @@ nothing # hide
 # Runge-Kutta time-stepping scheme, and a `BuoyancyTracer`.
 
 model = NonhydrostaticModel(; grid,
-                            advection = WENO5(),
+                            advection = WENO(),
                             timestepper = :RungeKutta3,
                             tracers = :b,
                             buoyancy = BuoyancyTracer(),

diff --git a/examples/langmuir_turbulence.jl b/examples/langmuir_turbulence.jl
@@ -127,7 +127,7 @@ coriolis = FPlane(f=1e-4) # s⁻¹
 # we use `UniformStokesDrift`, which expects Stokes drift functions of ``z, t`` only.
 
 model = NonhydrostaticModel(; grid, coriolis,
-                            advection = WENO5(),
+                            advection = WENO(),
                             timestepper = :RungeKutta3,
                             tracers = :b,
                             buoyancy = BuoyancyTracer(),

diff --git a/examples/shallow_water_Bickley_jet.jl b/examples/shallow_water_Bickley_jet.jl
@@ -40,7 +40,7 @@ grid = RectilinearGrid(size = (Nx, Ny),
 
 # ## Building a `ShallowWaterModel`
 #
-# We build a `ShallowWaterModel` with the `WENO5` advection scheme,
+# We build a `ShallowWaterModel` with the `WENO` advection scheme,
 # 3rd-order Runge-Kutta time-stepping, non-dimensional Coriolis and
 # gravitational acceleration
 
@@ -49,7 +49,7 @@ coriolis = FPlane(f=1)
 
 model = ShallowWaterModel(; grid, coriolis, gravitational_acceleration,
                           timestepper = :RungeKutta3,
-                          momentum_advection = WENO5())
+                          momentum_advection = WENO())
 
 # Use `architecture = GPU()` to run this problem on a GPU.
 

diff --git a/src/Advection/Advection.jl b/src/Advection/Advection.jl
@@ -16,41 +16,49 @@ export
     advective_tracer_flux_y,
     advective_tracer_flux_z,
 
-    CenteredSecondOrder,
-    UpwindBiasedFirstOrder,
-    UpwindBiasedThirdOrder,
-    UpwindBiasedFifthOrder,
-    CenteredFourthOrder,
-    WENO5,
+    AdvectionScheme,
+    Centered, CenteredSecondOrder, CenteredFourthOrder,
+    UpwindBiased, UpwindBiasedFirstOrder, UpwindBiasedThirdOrder, UpwindBiasedFifthOrder,
+    WENO, WENOThirdOrder, WENOFifthOrder,
     VectorInvariant,
     EnergyConservingScheme,
     EnstrophyConservingScheme
 
 using DocStringExtensions
 
+using Base: @propagate_inbounds
+using Adapt 
+using OffsetArrays
+using KernelAbstractions.Extras.LoopInfo: @unroll
+
 using Oceananigans.Grids
+using Oceananigans.Grids: with_halo, return_metrics
+using Oceananigans.Architectures: arch_array, architecture, CPU
+
 using Oceananigans.Operators
 
+import Base: show, summary
 import Oceananigans.Grids: required_halo_size
 
-abstract type AbstractAdvectionScheme{Buffer} end
-abstract type AbstractCenteredAdvectionScheme{Buffer} <: AbstractAdvectionScheme{Buffer} end
-abstract type AbstractUpwindBiasedAdvectionScheme{Buffer} <: AbstractAdvectionScheme{Buffer} end
+abstract type AbstractAdvectionScheme{Buffer, FT} end
+abstract type AbstractCenteredAdvectionScheme{Buffer, FT} <: AbstractAdvectionScheme{Buffer, FT} end
+abstract type AbstractUpwindBiasedAdvectionScheme{Buffer, FT} <: AbstractAdvectionScheme{Buffer, FT} end
 
-required_halo_size(scheme::AbstractAdvectionScheme{Buffer}) where Buffer = Buffer + 1
+@inline boundary_buffer(::AbstractAdvectionScheme{N}) where N = N
+@inline required_halo_size(scheme::AbstractAdvectionScheme{Buffer}) where Buffer = Buffer 
 
 include("centered_advective_fluxes.jl")
 include("upwind_biased_advective_fluxes.jl")
 include("flat_advective_fluxes.jl")
 
-include("upwind_biased_first_order.jl")
-include("centered_second_order.jl")
-include("upwind_biased_third_order.jl")
-include("centered_fourth_order.jl")
-include("upwind_biased_fifth_order.jl")
-include("weno_fifth_order.jl")
-include("vector_invariant_advection.jl")
+include("reconstruction_coefficients.jl")
+include("centered_reconstruction.jl")
+include("upwind_biased_reconstruction.jl")
+include("weno_reconstruction.jl")
+include("weno_interpolants.jl")
+include("stretched_weno_smoothness.jl")
 
+include("vector_invariant_advection.jl")
 include("topologically_conditional_interpolation.jl")
 
 include("momentum_advection_operators.jl")

diff --git a/src/Advection/centered_advective_fluxes.jl b/src/Advection/centered_advective_fluxes.jl
@@ -4,30 +4,30 @@
 ##### See topologically_conditional_interpolation.jl for an explanation of the underscore-prepended functions _symmetric_interpolate_*.
 #####
 
-const Centered = AbstractCenteredAdvectionScheme
+const CenteredScheme = AbstractCenteredAdvectionScheme
 
 #####
 ##### Advective momentum flux operators
 #####
 ##### Note the convention "advective_momentum_flux_Ua" corresponds to the advection _of_ a _by_ U.
 #####
 
-@inline advective_momentum_flux_Uu(i, j, k, grid, scheme::Centered, U, u) = Axᶜᶜᶜ(i, j, k, grid) * _symmetric_interpolate_xᶜᵃᵃ(i, j, k, grid, scheme, U) * _symmetric_interpolate_xᶜᵃᵃ(i, j, k, grid, scheme, u)
-@inline advective_momentum_flux_Vu(i, j, k, grid, scheme::Centered, V, u) = Ayᶠᶠᶜ(i, j, k, grid) * _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, V) * _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, u)
-@inline advective_momentum_flux_Wu(i, j, k, grid, scheme::Centered, W, u) = Azᶠᶜᶠ(i, j, k, grid) * _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, W) * _symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme, u)
+@inline advective_momentum_flux_Uu(i, j, k, grid, scheme::CenteredScheme, U, u) = @inbounds Axᶜᶜᶜ(i, j, k, grid) * _symmetric_interpolate_xᶜᵃᵃ(i, j, k, grid, scheme, U) * _symmetric_interpolate_xᶜᵃᵃ(i, j, k, grid, scheme, u)
+@inline advective_momentum_flux_Vu(i, j, k, grid, scheme::CenteredScheme, V, u) = @inbounds Ayᶠᶠᶜ(i, j, k, grid) * _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, V) * _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, u)
+@inline advective_momentum_flux_Wu(i, j, k, grid, scheme::CenteredScheme, W, u) = @inbounds Azᶠᶜᶠ(i, j, k, grid) * _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, W) * _symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme, u)
 
-@inline advective_momentum_flux_Uv(i, j, k, grid, scheme::Centered, U, v) = Axᶠᶠᶜ(i, j, k, grid) * _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, U) * _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, v)
-@inline advective_momentum_flux_Vv(i, j, k, grid, scheme::Centered, V, v) = Ayᶜᶜᶜ(i, j, k, grid) * _symmetric_interpolate_yᵃᶜᵃ(i, j, k, grid, scheme, V) * _symmetric_interpolate_yᵃᶜᵃ(i, j, k, grid, scheme, v)
-@inline advective_momentum_flux_Wv(i, j, k, grid, scheme::Centered, W, v) = Azᶜᶠᶠ(i, j, k, grid) * _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, W) * _symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme, v)
+@inline advective_momentum_flux_Uv(i, j, k, grid, scheme::CenteredScheme, U, v) = @inbounds Axᶠᶠᶜ(i, j, k, grid) * _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, U) * _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, v)
+@inline advective_momentum_flux_Vv(i, j, k, grid, scheme::CenteredScheme, V, v) = @inbounds Ayᶜᶜᶜ(i, j, k, grid) * _symmetric_interpolate_yᵃᶜᵃ(i, j, k, grid, scheme, V) * _symmetric_interpolate_yᵃᶜᵃ(i, j, k, grid, scheme, v)
+@inline advective_momentum_flux_Wv(i, j, k, grid, scheme::CenteredScheme, W, v) = @inbounds Azᶜᶠᶠ(i, j, k, grid) * _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, W) * _symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme, v)
 
-@inline advective_momentum_flux_Uw(i, j, k, grid, scheme::Centered, U, w) = Axᶠᶜᶠ(i, j, k, grid) * _symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme, U) * _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, w)
-@inline advective_momentum_flux_Vw(i, j, k, grid, scheme::Centered, V, w) = Ayᶜᶠᶠ(i, j, k, grid) * _symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme, V) * _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, w)
-@inline advective_momentum_flux_Ww(i, j, k, grid, scheme::Centered, W, w) = Azᶜᶜᶜ(i, j, k, grid) * _symmetric_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme, W) * _symmetric_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme, w)
+@inline advective_momentum_flux_Uw(i, j, k, grid, scheme::CenteredScheme, U, w) = @inbounds Axᶠᶜᶠ(i, j, k, grid) * _symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme, U) * _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, w)
+@inline advective_momentum_flux_Vw(i, j, k, grid, scheme::CenteredScheme, V, w) = @inbounds Ayᶜᶠᶠ(i, j, k, grid) * _symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme, V) * _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, w)
+@inline advective_momentum_flux_Ww(i, j, k, grid, scheme::CenteredScheme, W, w) = @inbounds Azᶜᶜᶜ(i, j, k, grid) * _symmetric_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme, W) * _symmetric_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme, w)
 
 #####
 ##### Advective tracer flux operators
 #####
 
-@inline advective_tracer_flux_x(i, j, k, grid, scheme::Centered, U, c) = @inbounds Ax_qᶠᶜᶜ(i, j, k, grid, U) * _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, c)
-@inline advective_tracer_flux_y(i, j, k, grid, scheme::Centered, V, c) = @inbounds Ay_qᶜᶠᶜ(i, j, k, grid, V) * _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, c)
-@inline advective_tracer_flux_z(i, j, k, grid, scheme::Centered, W, c) = @inbounds Az_qᶜᶜᶠ(i, j, k, grid, W) * _symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme, c)
+@inline advective_tracer_flux_x(i, j, k, grid, scheme::CenteredScheme, U, c) = @inbounds Ax_qᶠᶜᶜ(i, j, k, grid, U) * _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, c)
+@inline advective_tracer_flux_y(i, j, k, grid, scheme::CenteredScheme, V, c) = @inbounds Ay_qᶜᶠᶜ(i, j, k, grid, V) * _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, c)
+@inline advective_tracer_flux_z(i, j, k, grid, scheme::CenteredScheme, W, c) = @inbounds Az_qᶜᶜᶠ(i, j, k, grid, W) * _symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme, c)
diff --git a/src/Advection/centered_fourth_order.jl b/src/Advection/centered_fourth_order.jl