Different advection schemes in different directions

[simone-silvestri]

This was already possible for tracers with the TracerAdvection scheme. This PR changes the TracerAdvection scheme into a FluxFormAdvection scheme that can be used also for momentum

closes #3699
closes #3622

[glwagner]

Does this close #3622?

[glwagner]

I preferred it as it was before without metaprogramming

[glwagner]

Using this branch I get

julia> using Oceananigans
Precompiling Oceananigans
  1 dependency successfully precompiled in 11 seconds. 139 already precompiled.

julia> grid = RectilinearGrid(size = (16, 1, 16),
                              halo = (3, 1, 3),
                              x = (0, 1),
                              y = (0, 1),
                              z = (0, 1),
                              topology = (Periodic, Periodic, Bounded))
16×1×16 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×1×3 halo
├── Periodic x ∈ [0.0, 1.0) regularly spaced with Δx=0.0625
├── Periodic y ∈ [0.0, 1.0) regularly spaced with Δy=1.0
└── Bounded  z ∈ [0.0, 1.0] regularly spaced with Δz=0.0625

julia> model = NonhydrostaticModel(; grid, advection = WENO(order=5))
┌ Warning: Inflating model grid halo size to (3, 3, 3) and recreating grid. Note that an ImmersedBoundaryGrid requires an extra halo point in all non-flat directions compared to a non-immersed boundary grid.
└ @ Oceananigans.Models.NonhydrostaticModels ~/Projects/Oceananigans.jl/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl:245
ERROR: ArgumentError: halo must be ≤ size for coordinate y
Stacktrace:
 [1] validate_halo(TX::Type, TY::Type, TZ::Type, size::Tuple{Int64, Int64, Int64}, halo::Tuple{Int64, Int64, Int64})
   @ Oceananigans.Grids ~/Projects/Oceananigans.jl/src/Grids/input_validation.jl:87
 [2] validate_rectilinear_grid_args(topology::Tuple{…}, size::Tuple{…}, halo::Tuple{…}, FT::Type, extent::Nothing, x::Tuple{…}, y::Tuple{…}, z::Tuple{…})
   @ Oceananigans.Grids ~/Projects/Oceananigans.jl/src/Grids/rectilinear_grid.jl:292
 [3] RectilinearGrid(architecture::CPU, FT::DataType; size::Tuple{…}, x::Tuple{…}, y::Tuple{…}, z::Tuple{…}, halo::Tuple{…}, extent::Nothing, topology::Tuple{…})
   @ Oceananigans.Grids ~/Projects/Oceananigans.jl/src/Grids/rectilinear_grid.jl:269
 [4] with_halo(new_halo::Tuple{…}, old_grid::RectilinearGrid{…})
   @ Oceananigans.Grids ~/Projects/Oceananigans.jl/src/Grids/rectilinear_grid.jl:389
 [5] inflate_grid_halo_size(::RectilinearGrid{…}, ::WENO{…}, ::Vararg{…})
   @ Oceananigans.Models.NonhydrostaticModels ~/Projects/Oceananigans.jl/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl:250
 [6] NonhydrostaticModel(; grid::RectilinearGrid{…}, clock::Clock{…}, advection::WENO{…}, buoyancy::Nothing, coriolis::Nothing, stokes_drift::Nothing, forcing::@NamedTuple{}, closure::Nothing, boundary_conditions::@NamedTuple{}, tracers::Tuple{}, timestepper::Symbol, background_fields::@NamedTuple{}, particles::Nothing, biogeochemistry::Nothing, velocities::Nothing, hydrostatic_pressure_anomaly::Oceananigans.Models.NonhydrostaticModels.DefaultHydrostaticPressureAnomaly, nonhydrostatic_pressure::Field{…}, diffusivity_fields::Nothing, pressure_solver::Nothing, immersed_boundary::Nothing, auxiliary_fields::@NamedTuple{})
   @ Oceananigans.Models.NonhydrostaticModels ~/Projects/Oceananigans.jl/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl:181
 [7] top-level scope
   @ REPL[3]:1
Some type information was truncated. Use `show(err)` to see complete types.

[simone-silvestri]

you cannot use WENO(; order =5) in the y-direction if you have only one grid point.
In this case the script would have to change to

julia> using Oceananigans
Precompiling Oceananigans
  1 dependency successfully precompiled in 11 seconds. 139 already precompiled.

julia> grid = RectilinearGrid(size = (16, 1, 16),
                              halo = (3, 1, 3),
                              x = (0, 1),
                              y = (0, 1),
                              z = (0, 1),
                              topology = (Periodic, Periodic, Bounded))
16×1×16 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×1×3 halo
├── Periodic x ∈ [0.0, 1.0) regularly spaced with Δx=0.0625
├── Periodic y ∈ [0.0, 1.0) regularly spaced with Δy=1.0
└── Bounded  z ∈ [0.0, 1.0] regularly spaced with Δz=0.0625

julia> advection = FluxFormAdvection(WENO(order=5), nothing, WENO(order=5))

julia> model = NonhydrostaticModel(; grid, advection)

I think spitting out the error should be the correct behavior because we want to make sure that people know what scheme is begin used in the different directions, and correct the advection scheme accordingly. We can probably change the error message to be a bit more descriptive

[glwagner]

i think we should support that syntax, its unambiguous in my opinion what is supposed to happen and falls within our philosophy of "friendly" (ie friendly defaults, making life easy)

[glwagner]

This also fails:

julia> grid = RectilinearGrid(size = (16, 2, 16),
                              halo = (3, 2, 3),
                              x = (0, 1),
                              y = (0, 1),
                              z = (0, 1),
                              topology = (Periodic, Periodic, Bounded))
16×2×16 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×2×3 halo
├── Periodic x ∈ [0.0, 1.0) regularly spaced with Δx=0.0625
├── Periodic y ∈ [0.0, 1.0) regularly spaced with Δy=0.5
└── Bounded  z ∈ [0.0, 1.0] regularly spaced with Δz=0.0625

julia> model = NonhydrostaticModel(; grid, advection = WENO())
┌ Warning: Inflating model grid halo size to (3, 3, 3) and recreating grid. Note that an ImmersedBoundaryGrid requires an extra halo point in all non-flat directions compared to a non-immersed boundary grid.
└ @ Oceananigans.Models.NonhydrostaticModels ~/Projects/Oceananigans.jl/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl:245
ERROR: ArgumentError: halo must be ≤ size for coordinate y
Stacktrace:
 [1] validate_halo(TX::Type, TY::Type, TZ::Type, size::Tuple{Int64, Int64, Int64}, halo::Tuple{Int64, Int64, Int64})
   @ Oceananigans.Grids ~/Projects/Oceananigans.jl/src/Grids/input_validation.jl:87
 [2] validate_rectilinear_grid_args(topology::Tuple{…}, size::Tuple{…}, halo::Tuple{…}, FT::Type, extent::Nothing, x::Tuple{…}, y::Tuple{…}, z::Tuple{…})
   @ Oceananigans.Grids ~/Projects/Oceananigans.jl/src/Grids/rectilinear_grid.jl:292
 [3] RectilinearGrid(architecture::CPU, FT::DataType; size::Tuple{…}, x::Tuple{…}, y::Tuple{…}, z::Tuple{…}, halo::Tuple{…}, extent::Nothing, topology::Tuple{…})
   @ Oceananigans.Grids ~/Projects/Oceananigans.jl/src/Grids/rectilinear_grid.jl:269
 [4] with_halo(new_halo::Tuple{…}, old_grid::RectilinearGrid{…})
   @ Oceananigans.Grids ~/Projects/Oceananigans.jl/src/Grids/rectilinear_grid.jl:389
 [5] inflate_grid_halo_size(::RectilinearGrid{…}, ::WENO{…}, ::Vararg{…})
   @ Oceananigans.Models.NonhydrostaticModels ~/Projects/Oceananigans.jl/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl:250
 [6] NonhydrostaticModel(; grid::RectilinearGrid{…}, clock::Clock{…}, advection::WENO{…}, buoyancy::Nothing, coriolis::Nothing, stokes_drift::Nothing, forcing::@NamedTuple{}, closure::Nothing, boundary_conditions::@NamedTuple{}, tracers::Tuple{}, timestepper::Symbol, background_fields::@NamedTuple{}, particles::Nothing, biogeochemistry::Nothing, velocities::Nothing, hydrostatic_pressure_anomaly::Oceananigans.Models.NonhydrostaticModels.DefaultHydrostaticPressureAnomaly, nonhydrostatic_pressure::Field{…}, diffusivity_fields::Nothing, pressure_solver::Nothing, immersed_boundary::Nothing, auxiliary_fields::@NamedTuple{})
   @ Oceananigans.Models.NonhydrostaticModels ~/Projects/Oceananigans.jl/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl:181
 [7] top-level scope
   @ REPL[6]:1
Some type information was truncated. Use `show(err)` to see complete types.

ie when we don't specify an order at all for WENO. Why should this fail? It looks perfectly fine, reading the code.

[simone-silvestri]

@glwagner after the last commit the code should have the expected behavior of reducing the advection order to comply with grid-size limitations. I am not doing the same check for the vector invariant scheme because it might get a bit complicated.

[simone-silvestri]

This now spits out

julia> grid = RectilinearGrid(size = (16, 2, 16),
                                     halo = (3, 1, 3),
                                     x = (0, 1),
                                     y = (0, 1),
                                     z = (0, 1),
                                     topology = (Periodic, Periodic, Bounded))
16×2×16 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×1×3 halo
├── Periodic x ∈ [0.0, 1.0) regularly spaced with Δx=0.0625
├── Periodic y ∈ [0.0, 1.0) regularly spaced with Δy=0.5
└── Bounded  z ∈ [0.0, 1.0] regularly spaced with Δz=0.0625

julia> model = NonhydrostaticModel(; grid, advection = WENO());
[ Info: User-defined advection scheme WENO reconstruction order 5 reduced to WENO reconstruction order 3 in the y-direction to comply with grid-size limitations.
┌ Warning: Inflating model grid halo size to (3, 2, 3) and recreating grid. Note that an ImmersedBoundaryGrid requires an extra halo point in all non-flat directions compared to a non-immersed boundary grid.
└ @ Oceananigans.Models.NonhydrostaticModels ~/development/Oceananigans.jl/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl:250

julia> grid = RectilinearGrid(size = (16, 1, 16),
                                     halo = (3, 1, 3),
                                     x = (0, 1),
                                     y = (0, 1),
                                     z = (0, 1),
                                     topology = (Periodic, Periodic, Bounded))
16×1×16 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×1×3 halo
├── Periodic x ∈ [0.0, 1.0) regularly spaced with Δx=0.0625
├── Periodic y ∈ [0.0, 1.0) regularly spaced with Δy=1.0
└── Bounded  z ∈ [0.0, 1.0] regularly spaced with Δz=0.0625

julia> model = NonhydrostaticModel(; grid, advection = WENO());
[ Info: User-defined advection scheme WENO reconstruction order 5 reduced to Nothing in the y-direction to comply with grid-size limitations.

[simone-silvestri]

This should be ready to go

[glwagner]

Please restore this code to how it was!

[simone-silvestri]

There was a problem with this test. The constructor of the ScalarDiffusivity is not type-stable (like most constructors). I have documented this problem with comments in this PR. This test should produce the same exact result just avoiding the need to reconstruct the ScalarDiffusivity. If we want a different solution, that's okay. The problem is that it will be difficult to enforce type stability on all our constructors.

[glwagner]

I think it would be better to mark this as @test_broken then. We need the constructor to be type stable to do this work!

[glwagner]

Please understand what I am saying: the point of this test is not to get a numerical result. It's supposed to verify that we can use this method for doing parameter estimation with this specific setup. So if the code breaks then the test fails.

[simone-silvestri]

ok I can do that

[glwagner]

There was a problem with this test. The constructor of the ScalarDiffusivity is not type-stable (like most constructors)

Many constructors are type stable, and we have to work to make them all type stable to support AD.

This test passed previously, so not sure what changed in this PR?

[glwagner]

Please restore the enzyme tests!

[glwagner]

This should be ready to go

Please just restore the enzyme tests because we would like to maintain a parameter-estimation style test. For a closure like CATKE, we have to use the pattern where we set model.closure = new_closure. So despite that using an array for the diffusivity is functionally equivalent to what was written, it is definitely not equivalent in terms of the code that is exercised for taking a gradient.