Skip to content

Commit

Permalink
Add more advection scheme tests; fix order adapting with tracer-speci…
Browse files Browse the repository at this point in the history 
…fic schemes; fix WENOVectorInvariant (#3864)

* Add a test for FluxFormAdvection + different schemes for different tracers

* Reconfigure hydrostatic free surface model constructor

* Bump Project version

* better FluxFormAdvection docstring

* Add missing comma

* Fix syntax error in test file

* More bugfixes

* Fix bug in WENOVectorInvariant constructor

* Add WENOVectorInvariant to tests

* Anoter test bug

* Fix issue with automagic bgc tracers

---------

Co-authored-by: Navid C. Constantinou <[email protected]>
  • Loading branch information
@glwagner @navidcy
glwagner and navidcy authored Oct 24, 2024
1 parent 12fdd64 commit 30e3e05
Show file tree
Hide file tree
Showing 5 changed files with 100 additions and 78 deletions.
2 changes: 1 addition & 1 deletion Project.toml
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "Oceananigans"
uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
authors = ["Climate Modeling Alliance and contributors"]
version = "0.92.1"
version = "0.92.2"

[deps]
Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
Expand Down
6 changes: 3 additions & 3 deletions src/Advection/flux_form_advection.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,10 +10,10 @@ struct FluxFormAdvection{N, FT, A, B, C} <: AbstractAdvectionScheme{N, FT}
end

"""
function FluxFormAdvection(x, y, z)
FluxFormAdvection(x_advection, y_advection, z_advection)
Builds a `FluxFormAdvection` type with reconstructions schemes `x`, `y`, and `z` to be applied in
the x, y, and z direction, respectively.
Return a `FluxFormAdvection` type with reconstructions schemes `x_advection`, `y_advection`,
and `z_advection` to be applied in the ``x``, ``y``, and ``z`` directions, respectively.
"""
function FluxFormAdvection(x_advection, y_advection, z_advection)
Hx = required_halo_size_x(x_advection)
Expand Down
11 changes: 8 additions & 3 deletions src/Advection/vector_invariant_advection.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -178,7 +178,6 @@ nothing_to_default(user_value; default) = isnothing(user_value) ? default : user
WENOVectorInvariant(; upwinding = nothing,
multi_dimensional_stencil = false,
weno_kw...)
"""
function WENOVectorInvariant(FT::DataType = Float64;
upwinding = nothing,
Expand Down Expand Up @@ -211,8 +210,14 @@ function WENOVectorInvariant(FT::DataType = Float64;
default_upwinding = OnlySelfUpwinding(cross_scheme = divergence_scheme)
upwinding = nothing_to_default(upwinding; default = default_upwinding)

schemes = (vorticity_scheme, vertical_scheme, kinetic_energy_gradient_scheme, divergence_scheme)
N = maximum(required_halo_size(s) for s in schemes)
N = max(required_halo_size_x(vorticity_scheme),
required_halo_size_y(vorticity_scheme),
required_halo_size_x(divergence_scheme),
required_halo_size_y(divergence_scheme),
required_halo_size_x(kinetic_energy_gradient_scheme),
required_halo_size_y(kinetic_energy_gradient_scheme),
required_halo_size_z(vertical_scheme))

FT = eltype(vorticity_scheme) # assumption

return VectorInvariant{N, FT, multi_dimensional_stencil}(vorticity_scheme,
Expand Down
103 changes: 51 additions & 52 deletions src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,22 +56,21 @@ default_free_surface(grid; gravitational_acceleration=g_Earth) =
"""
HydrostaticFreeSurfaceModel(; grid,
clock = Clock{eltype(grid)}(time = 0),
momentum_advection = VectorInvariant(),
tracer_advection = CenteredSecondOrder(),
buoyancy = SeawaterBuoyancy(eltype(grid)),
coriolis = nothing,
free_surface = default_free_surface(grid, gravitational_acceleration=g_Earth),
forcing::NamedTuple = NamedTuple(),
closure = nothing,
boundary_conditions::NamedTuple = NamedTuple(),
tracers = (:T, :S),
particles::ParticlesOrNothing = nothing,
biogeochemistry::AbstractBGCOrNothing = nothing,
velocities = nothing,
pressure = nothing,
diffusivity_fields = nothing,
auxiliary_fields = NamedTuple(),
)
momentum_advection = VectorInvariant(),
tracer_advection = CenteredSecondOrder(),
buoyancy = SeawaterBuoyancy(eltype(grid)),
coriolis = nothing,
free_surface = default_free_surface(grid, gravitational_acceleration=g_Earth),
forcing::NamedTuple = NamedTuple(),
closure = nothing,
boundary_conditions::NamedTuple = NamedTuple(),
tracers = (:T, :S),
particles::ParticlesOrNothing = nothing,
biogeochemistry::AbstractBGCOrNothing = nothing,
velocities = nothing,
pressure = nothing,
diffusivity_fields = nothing,
auxiliary_fields = NamedTuple())
Construct a hydrostatic model with a free surface on `grid`.
Expand Down Expand Up @@ -103,38 +102,42 @@ Keyword arguments
- `auxiliary_fields`: `NamedTuple` of auxiliary fields. Default: `nothing`.
"""
function HydrostaticFreeSurfaceModel(; grid,
clock = Clock{eltype(grid)}(time = 0),
momentum_advection = VectorInvariant(),
tracer_advection = CenteredSecondOrder(),
buoyancy = nothing,
coriolis = nothing,
free_surface = default_free_surface(grid, gravitational_acceleration=g_Earth),
tracers = nothing,
forcing::NamedTuple = NamedTuple(),
closure = nothing,
boundary_conditions::NamedTuple = NamedTuple(),
particles::ParticlesOrNothing = nothing,
biogeochemistry::AbstractBGCOrNothing = nothing,
clock = Clock{eltype(grid)}(time = 0),
momentum_advection = VectorInvariant(),
tracer_advection = CenteredSecondOrder(),
buoyancy = nothing,
coriolis = nothing,
free_surface = default_free_surface(grid, gravitational_acceleration=g_Earth),
tracers = nothing,
forcing::NamedTuple = NamedTuple(),
closure = nothing,
boundary_conditions::NamedTuple = NamedTuple(),
particles::ParticlesOrNothing = nothing,
biogeochemistry::AbstractBGCOrNothing = nothing,
velocities = nothing,
pressure = nothing,
diffusivity_fields = nothing,
auxiliary_fields = NamedTuple()
)
pressure = nothing,
diffusivity_fields = nothing,
auxiliary_fields = NamedTuple())

# Check halos and throw an error if the grid's halo is too small
@apply_regionally validate_model_halo(grid, momentum_advection, tracer_advection, closure)

arch = architecture(grid)

@apply_regionally momentum_advection = validate_momentum_advection(momentum_advection, grid)

# Validate biogeochemistry (add biogeochemical tracers automagically)
tracers = tupleit(tracers) # supports tracers=:c keyword argument (for example)
biogeochemical_fields = merge(auxiliary_fields, biogeochemical_auxiliary_fields(biogeochemistry))
tracers, auxiliary_fields = validate_biogeochemistry(tracers, biogeochemical_fields, biogeochemistry, grid, clock)

# Reduce the advection order in directions that do not have enough grid points
momentum_advection = adapt_advection_order(momentum_advection, grid)
tracer_advection = adapt_advection_order(tracer_advection, grid)
@apply_regionally momentum_advection = validate_momentum_advection(momentum_advection, grid)
default_tracer_advection, tracer_advection = validate_tracer_advection(tracer_advection, grid)
default_generator(name, tracer_advection) = default_tracer_advection

# Generate tracer advection scheme for each tracer
tracer_advection_tuple = with_tracers(tracernames(tracers), tracer_advection, default_generator, with_velocities=false)
momentum_advection_tuple = (; momentum = momentum_advection)
advection = merge(momentum_advection_tuple, tracer_advection_tuple)
advection = NamedTuple(name => adapt_advection_order(scheme, grid) for (name, scheme) in pairs(advection))

tracers, auxiliary_fields = validate_biogeochemistry(tracers, merge(auxiliary_fields, biogeochemical_auxiliary_fields(biogeochemistry)), biogeochemistry, grid, clock)
validate_buoyancy(buoyancy, tracernames(tracers))
buoyancy = regularize_buoyancy(buoyancy)

Expand All @@ -152,7 +155,8 @@ function HydrostaticFreeSurfaceModel(; grid,

# Next, we form a list of default boundary conditions:
prognostic_field_names = (:u, :v, :w, tracernames(tracers)..., :η, keys(auxiliary_fields)...)
default_boundary_conditions = NamedTuple{prognostic_field_names}(Tuple(FieldBoundaryConditions() for name in prognostic_field_names))
default_boundary_conditions = NamedTuple{prognostic_field_names}(Tuple(FieldBoundaryConditions()
for name in prognostic_field_names))

# Then we merge specified, embedded, and default boundary conditions. Specified boundary conditions
# have precedence, followed by embedded, followed by default.
Expand All @@ -161,7 +165,11 @@ function HydrostaticFreeSurfaceModel(; grid,

# Finally, we ensure that closure-specific boundary conditions, such as
# those required by CATKEVerticalDiffusivity, are enforced:
boundary_conditions = add_closure_specific_boundary_conditions(closure, boundary_conditions, grid, tracernames(tracers), buoyancy)
boundary_conditions = add_closure_specific_boundary_conditions(closure,
boundary_conditions,
grid,
tracernames(tracers),
buoyancy)

# Ensure `closure` describes all tracers
closure = with_tracers(tracernames(tracers), closure)
Expand All @@ -177,6 +185,7 @@ function HydrostaticFreeSurfaceModel(; grid,

@apply_regionally validate_velocity_boundary_conditions(grid, velocities)

arch = architecture(grid)
free_surface = validate_free_surface(arch, free_surface)
free_surface = materialize_free_surface(free_surface, velocities, grid)

Expand All @@ -190,17 +199,7 @@ function HydrostaticFreeSurfaceModel(; grid,
# Regularize forcing for model tracer and velocity fields.
model_fields = merge(hydrostatic_prognostic_fields(velocities, free_surface, tracers), auxiliary_fields)
forcing = model_forcing(model_fields; forcing...)

default_tracer_advection, tracer_advection = validate_tracer_advection(tracer_advection, grid)

# Advection schemes
tracer_advection_tuple = with_tracers(tracernames(tracers),
tracer_advection,
(name, tracer_advection) -> default_tracer_advection,
with_velocities=false)

advection = merge((momentum=momentum_advection,), tracer_advection_tuple)


model = HydrostaticFreeSurfaceModel(arch, grid, clock, advection, buoyancy, coriolis,
free_surface, forcing, closure, particles, biogeochemistry, velocities, tracers,
pressure, diffusivity_fields, timestepper, auxiliary_fields)
Expand Down
56 changes: 37 additions & 19 deletions test/test_hydrostatic_free_surface_models.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,23 +3,24 @@ include("dependencies_for_runtests.jl")
using Oceananigans.Models.HydrostaticFreeSurfaceModels: VectorInvariant, PrescribedVelocityFields
using Oceananigans.Models.HydrostaticFreeSurfaceModels: ExplicitFreeSurface, ImplicitFreeSurface
using Oceananigans.Models.HydrostaticFreeSurfaceModels: SingleColumnGrid
using Oceananigans.Advection: EnergyConserving, EnstrophyConserving
using Oceananigans.Advection: EnergyConserving, EnstrophyConserving, FluxFormAdvection
using Oceananigans.TurbulenceClosures
using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity

function time_step_hydrostatic_model_works(grid;
coriolis = nothing,
free_surface = ExplicitFreeSurface(),
momentum_advection = nothing,
tracers = [:b],
tracer_advection = nothing,
closure = nothing,
velocities = nothing)

tracers = [:b]
buoyancy = BuoyancyTracer()
closure isa CATKEVerticalDiffusivity && push!(tracers, :e)

model = HydrostaticFreeSurfaceModel(; grid, coriolis, tracers, velocities, buoyancy,
momentum_advection, free_surface, closure)
momentum_advection, tracer_advection, free_surface, closure)

simulation = Simulation(model, Δt=1.0, stop_iteration=1)

Expand Down Expand Up @@ -159,6 +160,7 @@ topos_3d = ((Periodic, Periodic, Bounded),
end

for arch in archs

for topo in topos_3d
grid = RectilinearGrid(arch, size=(1, 1, 1), extent=(1, 1, 1), topology=topo)

Expand All @@ -170,14 +172,18 @@ topos_3d = ((Periodic, Periodic, Bounded),

z_face_generator(; Nz=1, p=1, H=1) = k -> -H + (k / (Nz+1))^p # returns a generating function

rectilinear_grid = RectilinearGrid(arch, size=(3, 3, 1), extent=(1, 1, 1), halo=(3, 3, 3))
vertically_stretched_grid = RectilinearGrid(arch, size=(3, 3, 1), x=(0, 1), y=(0, 1), z=z_face_generator(), halo=(3, 3, 3))
H = 7
halo = (7, 7, 7)
rectilinear_grid = RectilinearGrid(arch; size=(H, H, 1), extent=(1, 1, 1), halo)
vertically_stretched_grid = RectilinearGrid(arch; size=(H, H, 1), x=(0, 1), y=(0, 1), z=z_face_generator(), halo=(H, H, H))

lat_lon_sector_grid = LatitudeLongitudeGrid(arch, size=(3, 3, 3), longitude=(0, 60), latitude=(15, 75), z=(-1, 0), precompute_metrics=true)
lat_lon_strip_grid = LatitudeLongitudeGrid(arch, size=(3, 3, 3), longitude=(-180, 180), latitude=(15, 75), z=(-1, 0), precompute_metrics=true)
precompute_metrics = true
lat_lon_sector_grid = LatitudeLongitudeGrid(arch; size=(H, H, H), longitude=(0, 60), latitude=(15, 75), z=(-1, 0), precompute_metrics, halo)
lat_lon_strip_grid = LatitudeLongitudeGrid(arch; size=(H, H, H), longitude=(-180, 180), latitude=(15, 75), z=(-1, 0), precompute_metrics, halo)

lat_lon_sector_grid_stretched = LatitudeLongitudeGrid(arch, size=(3, 3, 3), longitude=(0, 60), latitude=(15, 75), z=z_face_generator(), precompute_metrics=true)
lat_lon_strip_grid_stretched = LatitudeLongitudeGrid(arch, size=(3, 3, 3), longitude=(-180, 180), latitude=(15, 75), z=z_face_generator(), precompute_metrics=true)
z = z_face_generator()
lat_lon_sector_grid_stretched = LatitudeLongitudeGrid(arch; size=(H, H, H), longitude=(0, 60), latitude=(15, 75), z, precompute_metrics, halo)
lat_lon_strip_grid_stretched = LatitudeLongitudeGrid(arch; size=(H, H, H), longitude=(-180, 180), latitude=(15, 75), z, precompute_metrics, halo)

grids = (rectilinear_grid, vertically_stretched_grid,
lat_lon_sector_grid, lat_lon_strip_grid,
Expand Down Expand Up @@ -210,22 +216,34 @@ topos_3d = ((Periodic, Periodic, Bounded),
HydrostaticSphericalCoriolis(scheme=EnstrophyConserving()))

@testset "Time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(coriolis))]" begin
@test time_step_hydrostatic_model_works(lat_lon_sector_grid, coriolis=coriolis)
@test time_step_hydrostatic_model_works(lat_lon_strip_grid, coriolis=coriolis)
@test time_step_hydrostatic_model_works(lat_lon_sector_grid; coriolis)
@test time_step_hydrostatic_model_works(lat_lon_strip_grid; coriolis)
end
end

for momentum_advection in (VectorInvariant(), WENOVectorInvariant(), CenteredSecondOrder(), WENO())
@testset "Time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(momentum_advection))]" begin
@info " Testing time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(momentum_advection))]..."
@test time_step_hydrostatic_model_works(rectilinear_grid; momentum_advection)
end
end

for momentum_advection in (VectorInvariant(), CenteredSecondOrder(), WENO())
for momentum_advection in (VectorInvariant(), WENOVectorInvariant())
@testset "Time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(momentum_advection))]" begin
@info " Testing time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(momentum_advection))]..."
@test time_step_hydrostatic_model_works(rectilinear_grid, momentum_advection=momentum_advection)
@test time_step_hydrostatic_model_works(lat_lon_sector_grid; momentum_advection)
end
end

momentum_advection = VectorInvariant()
@testset "Time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(momentum_advection))]" begin
@info " Testing time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(momentum_advection))]..."
@test time_step_hydrostatic_model_works(lat_lon_sector_grid, momentum_advection=momentum_advection)
for tracer_advection in [WENO(),
FluxFormAdvection(WENO(), WENO(), Centered()),
(b=WENO(), c=nothing)]

T = typeof(tracer_advection)
@testset "Time-stepping HydrostaticFreeSurfaceModels with tracer advection [$arch, $T]" begin
@info " Testing time-stepping HydrostaticFreeSurfaceModels with tracer advection [$arch, $T]..."
@test time_step_hydrostatic_model_works(rectilinear_grid; tracer_advection, tracers=[:b, :c])
end
end

for closure in (ScalarDiffusivity(),
Expand All @@ -238,8 +256,8 @@ topos_3d = ((Periodic, Periodic, Bounded),
@testset "Time-stepping Curvilinear HydrostaticFreeSurfaceModels [$arch, $(typeof(closure).name.wrapper)]" begin
@info " Testing time-stepping Curvilinear HydrostaticFreeSurfaceModels [$arch, $(typeof(closure).name.wrapper)]..."
@test_skip time_step_hydrostatic_model_works(arch, vertically_stretched_grid, closure=closure)
@test time_step_hydrostatic_model_works(lat_lon_sector_grid, closure=closure)
@test time_step_hydrostatic_model_works(lat_lon_strip_grid, closure=closure)
@test time_step_hydrostatic_model_works(lat_lon_sector_grid; closure)
@test time_step_hydrostatic_model_works(lat_lon_strip_grid; closure)
end
end

Expand Down

2 comments on commit 30e3e05

@glwagner
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator register

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/118009

Tip: Release Notes

Did you know you can add release notes too? Just add markdown formatted text underneath the comment after the text
"Release notes:" and it will be added to the registry PR, and if TagBot is installed it will also be added to the
release that TagBot creates. i.e.

@JuliaRegistrator register

Release notes:

## Breaking changes

- blah

To add them here just re-invoke and the PR will be updated.

Tagging

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v0.92.2 -m "<description of version>" 30e3e05fe661cb5f7a0db37ab6c3f42ac8c5f80e
git push origin v0.92.2

Please sign in to comment.