diff --git a/src/BuoyancyFormulations/buoyancy_force.jl b/src/BuoyancyFormulations/buoyancy_force.jl index 903872fac2..9d494470e6 100644 --- a/src/BuoyancyFormulations/buoyancy_force.jl +++ b/src/BuoyancyFormulations/buoyancy_force.jl @@ -1,21 +1,25 @@ -using Oceananigans.Grids: NegativeZDirection, validate_unit_vector +using Oceananigans.Utils +using Oceananigans.Fields +using Oceananigans.Grids: NegativeZDirection, validate_unit_vector, architecture +using Oceananigans.BoundaryConditions + +using KernelAbstractions: @kernel, @index using Adapt -struct BuoyancyForce{M, G} +struct BuoyancyForce{M, G, B} formulation :: M gravity_unit_vector :: G + gradients :: B end -Adapt.adapt_structure(to, bf::BuoyancyForce) = - BuoyancyForce(adapt(to, bf.formulation), - adapt(to, bf.gravity_unit_vector)) - """ - BuoyancyForce(formulation; gravity_unit_vector=NegativeZDirection()) + BuoyancyForce(grid, formulation::AbstractBuoyancyFormulation; gravity_unit_vector=NegativeZDirection(), materialize_gradients=true) -Construct a `buoyancy` given a buoyancy `model`. Optional keyword argument `gravity_unit_vector` +Construct a `buoyancy` given a buoyancy `formulation`. Optional keyword argument `gravity_unit_vector` can be used to specify the direction of gravity (default `NegativeZDirection()`). The buoyancy acceleration acts in the direction opposite to gravity. +If `materialize_gradients` is true (default), the buoyancy gradients will be precomputed and stored in fields for +performance reasons. For `materialize_gradients=true`, the `grid` argument must be provided. Example ======= @@ -44,11 +48,31 @@ NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0) └── coriolis: Nothing ``` """ -function BuoyancyForce(formulation; gravity_unit_vector=NegativeZDirection()) +function BuoyancyForce(grid, formulation::AbstractBuoyancyFormulation; gravity_unit_vector=NegativeZDirection(), materialize_gradients=true) gravity_unit_vector = validate_unit_vector(gravity_unit_vector) - return BuoyancyForce(formulation, gravity_unit_vector) + + if materialize_gradients + ∂x_b = XFaceField(grid) + ∂y_b = YFaceField(grid) + ∂z_b = ZFaceField(grid) + + gradients = (; ∂x_b, ∂y_b, ∂z_b) + else + gradients = nothing + end + + return BuoyancyForce(formulation, gravity_unit_vector, gradients) end +# Fallback for when no grid is available, we overwrite `materialize_gradients` to false +BuoyancyForce(formulation::AbstractBuoyancyFormulation; materialize_gradients=false, kwargs...) = + BuoyancyForce(nothing, formulation; materialize_gradients=false, kwargs...) + +Adapt.adapt_structure(to, bf::BuoyancyForce) = + BuoyancyForce(Adapt.adapt(to, bf.formulation), + Adapt.adapt(to, bf.gravity_unit_vector), + Adapt.adapt(to, bf.gradients)) + @inline ĝ_x(bf) = @inbounds - bf.gravity_unit_vector[1] @inline ĝ_y(bf) = @inbounds - bf.gravity_unit_vector[2] @inline ĝ_z(bf) = @inbounds - bf.gravity_unit_vector[3] @@ -65,14 +89,18 @@ end @inline get_temperature_and_salinity(bf::BuoyancyForce, C) = get_temperature_and_salinity(bf.formulation, C) -@inline ∂x_b(i, j, k, grid, b::BuoyancyForce, C) = ∂x_b(i, j, k, grid, b.formulation, C) -@inline ∂y_b(i, j, k, grid, b::BuoyancyForce, C) = ∂y_b(i, j, k, grid, b.formulation, C) -@inline ∂z_b(i, j, k, grid, b::BuoyancyForce, C) = ∂z_b(i, j, k, grid, b.formulation, C) +@inline ∂x_b(i, j, k, grid, b::BuoyancyForce{<:Any, <:Any, Nothing}, C) = ∂x_b(i, j, k, grid, b.formulation, C) +@inline ∂y_b(i, j, k, grid, b::BuoyancyForce{<:Any, <:Any, Nothing}, C) = ∂y_b(i, j, k, grid, b.formulation, C) +@inline ∂z_b(i, j, k, grid, b::BuoyancyForce{<:Any, <:Any, Nothing}, C) = ∂z_b(i, j, k, grid, b.formulation, C) @inline top_buoyancy_flux(i, j, grid, b::BuoyancyForce, args...) = top_buoyancy_flux(i, j, grid, b.formulation, args...) -regularize_buoyancy(bf) = bf -regularize_buoyancy(formulation::AbstractBuoyancyFormulation) = BuoyancyForce(formulation) +materialize_buoyancy(bf, grid; kw...) = bf +materialize_buoyancy(formulation::AbstractBuoyancyFormulation, grid; kw...) = BuoyancyForce(grid, formulation; kw...) + +# Fallback +compute_buoyancy_gradients!(::BuoyancyForce{<:Any, <:Any, <:Nothing}, grid, tracers; kw...) = nothing +compute_buoyancy_gradients!(::Nothing, grid, tracers; kw...) = nothing Base.summary(bf::BuoyancyForce) = string(summary(bf.formulation), " with ĝ = ", @@ -89,3 +117,29 @@ function Base.show(io::IO, bf::BuoyancyForce) "├── formulation: ", prettysummary(bf.formulation), '\n', "└── gravity_unit_vector: ", summarize_vector(bf.gravity_unit_vector)) end + +##### +##### Some performance optimizations for models that compute gradients over and over... +##### + +@inline ∂x_b(i, j, k, grid, b::BuoyancyForce, C) = @inbounds b.gradients.∂x_b[i, j, k] +@inline ∂y_b(i, j, k, grid, b::BuoyancyForce, C) = @inbounds b.gradients.∂y_b[i, j, k] +@inline ∂z_b(i, j, k, grid, b::BuoyancyForce, C) = @inbounds b.gradients.∂z_b[i, j, k] + +function compute_buoyancy_gradients!(buoyancy, grid, tracers; parameters=:xyz) + gradients = buoyancy.gradients + formulation = buoyancy.formulation + launch!(architecture(grid), grid, parameters, _compute_buoyancy_gradients!, gradients, grid, formulation, tracers) + fill_halo_regions!(gradients, only_local_halos=true) + + return nothing +end + +@kernel function _compute_buoyancy_gradients!(g, grid, b, C) + i, j, k = @index(Global, NTuple) + @inbounds begin + g.∂x_b[i, j, k] = ∂x_b(i, j, k, grid, b, C) + g.∂y_b[i, j, k] = ∂y_b(i, j, k, grid, b, C) + g.∂z_b[i, j, k] = ∂z_b(i, j, k, grid, b, C) + end +end diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl index fb39c66cf8..48a174fb63 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl @@ -3,7 +3,7 @@ using OrderedCollections: OrderedDict using Oceananigans.DistributedComputations using Oceananigans.Architectures: AbstractArchitecture using Oceananigans.Advection: AbstractAdvectionScheme, Centered, VectorInvariant, adapt_advection_order -using Oceananigans.BuoyancyFormulations: validate_buoyancy, regularize_buoyancy, SeawaterBuoyancy +using Oceananigans.BuoyancyFormulations: validate_buoyancy, materialize_buoyancy, SeawaterBuoyancy using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions using Oceananigans.Biogeochemistry: validate_biogeochemistry, AbstractBiogeochemistry, biogeochemical_auxiliary_fields using Oceananigans.Fields: Field, CenterField, tracernames, VelocityFields, TracerFields @@ -163,7 +163,7 @@ function HydrostaticFreeSurfaceModel(; grid, advection = NamedTuple(name => adapt_advection_order(scheme, grid) for (name, scheme) in pairs(advection)) validate_buoyancy(buoyancy, tracernames(tracers)) - buoyancy = regularize_buoyancy(buoyancy) + buoyancy = materialize_buoyancy(buoyancy, grid) # Collect boundary conditions for all model prognostic fields and, if specified, some model # auxiliary fields. Boundary conditions are "regularized" based on the _name_ of the field: diff --git a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl index 16587f9d77..840d511f22 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl @@ -4,6 +4,7 @@ using Oceananigans.BoundaryConditions using Oceananigans: UpdateStateCallsite using Oceananigans.Biogeochemistry: update_biogeochemical_state! using Oceananigans.BoundaryConditions: update_boundary_conditions! +using Oceananigans.BuoyancyFormulations: compute_buoyancy_gradients! using Oceananigans.TurbulenceClosures: compute_diffusivities! using Oceananigans.ImmersedBoundaries: mask_immersed_field!, mask_immersed_field_xy!, inactive_node using Oceananigans.Models: update_model_field_time_series! @@ -83,6 +84,9 @@ function compute_auxiliaries!(model::HydrostaticFreeSurfaceModel; w_parameters = P = model.pressure.pHY′ arch = architecture(grid) + # Maybe compute buoyancy gradients + compute_buoyancy_gradients!(buoyancy, grid, tracers; parameters = κ_parameters) + # Update the vertical velocity to comply with the barotropic correction step update_grid_vertical_velocity!(model, grid, model.vertical_coordinate) diff --git a/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl b/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl index 561e689c36..1c8da0d263 100644 --- a/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl +++ b/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl @@ -3,7 +3,7 @@ using OrderedCollections: OrderedDict using Oceananigans.Architectures: AbstractArchitecture using Oceananigans.DistributedComputations: Distributed using Oceananigans.Advection: Centered, adapt_advection_order -using Oceananigans.BuoyancyFormulations: validate_buoyancy, regularize_buoyancy, SeawaterBuoyancy +using Oceananigans.BuoyancyFormulations: validate_buoyancy, materialize_buoyancy, SeawaterBuoyancy using Oceananigans.BoundaryConditions: MixedBoundaryCondition using Oceananigans.Biogeochemistry: validate_biogeochemistry, AbstractBiogeochemistry, biogeochemical_auxiliary_fields using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions @@ -188,7 +188,7 @@ function NonhydrostaticModel(; grid, all_auxiliary_fields = merge(auxiliary_fields, biogeochemical_auxiliary_fields(biogeochemistry)) tracers, auxiliary_fields = validate_biogeochemistry(tracers, all_auxiliary_fields, biogeochemistry, grid, clock) validate_buoyancy(buoyancy, tracernames(tracers)) - buoyancy = regularize_buoyancy(buoyancy) + buoyancy = materialize_buoyancy(buoyancy, grid) # Adjust advection scheme to be valid on a particular grid size. i.e. if the grid size # is smaller than the advection order, reduce the order of the advection in that particular diff --git a/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl b/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl index ac16bcad60..a663f30cdd 100644 --- a/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl +++ b/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl @@ -3,6 +3,7 @@ using Oceananigans.Architectures using Oceananigans.BoundaryConditions using Oceananigans.Biogeochemistry: update_biogeochemical_state! using Oceananigans.BoundaryConditions: update_boundary_conditions! +using Oceananigans.BuoyancyFormulations: compute_buoyancy_gradients! using Oceananigans.TurbulenceClosures: compute_diffusivities! using Oceananigans.Fields: compute! using Oceananigans.ImmersedBoundaries: mask_immersed_field! @@ -55,15 +56,23 @@ function update_state!(model::NonhydrostaticModel, callbacks=[]; compute_tendenc return nothing end -function compute_auxiliaries!(model::NonhydrostaticModel; p_parameters = tuple(p_kernel_parameters(model.grid)), - κ_parameters = tuple(:xyz)) +function compute_auxiliaries!(model::NonhydrostaticModel; p_parameters = p_kernel_parameters(model.grid), + κ_parameters = :xyz) + grid = model.grid closure = model.closure diffusivity = model.diffusivity_fields + tracers = model.tracers + buoyancy = model.buoyancy + + # Maybe compute buoyancy gradients + compute_buoyancy_gradients!(buoyancy, grid, tracers; parameters = κ_parameters) + + # Compute diffusivities + compute_diffusivities!(diffusivity, closure, model; parameters = κ_parameters) + + # Update hydrostatic pressure + update_hydrostatic_pressure!(model; parameters = p_parameters) - for (ppar, κpar) in zip(p_parameters, κ_parameters) - compute_diffusivities!(diffusivity, closure, model; parameters = κpar) - update_hydrostatic_pressure!(model; parameters = ppar) - end return nothing end diff --git a/src/MultiRegion/multi_region_models.jl b/src/MultiRegion/multi_region_models.jl index a7e8fe826f..fa83c4a3b6 100644 --- a/src/MultiRegion/multi_region_models.jl +++ b/src/MultiRegion/multi_region_models.jl @@ -1,5 +1,6 @@ using Oceananigans.Models: AbstractModel using Oceananigans.Advection: WENO, VectorInvariant +using Oceananigans.BuoyancyFormulations: BuoyancyForce, NegativeZDirection, AbstractBuoyancyFormulation using Oceananigans.Models.HydrostaticFreeSurfaceModels: AbstractFreeSurface using Oceananigans.TimeSteppers: AbstractTimeStepper, QuasiAdamsBashforth2TimeStepper using Oceananigans.Models: PrescribedVelocityFields @@ -39,6 +40,7 @@ Types = (:HydrostaticFreeSurfaceModel, :PrescribedVelocityFields, :ConjugateGradientSolver, :CrossAndSelfUpwinding, + :BuoyancyForce, :OnlySelfUpwinding, :GridFittedBoundary, :GridFittedBottom, @@ -53,6 +55,10 @@ for T in Types end end +# TODO: For the moment, buoyancy gradients cannot be precomputed in MultiRegionModels +BuoyancyForce(grid::MultiRegionGrids, formulation::AbstractBuoyancyFormulation; gravity_unit_vector=NegativeZDirection(), materialize_gradients=true) = + BuoyancyForce(grid, formulation; gravity_unit_vector=gravity_unit_vector, materialize_gradients=false) + @inline isregional(pv::PrescribedVelocityFields) = isregional(pv.u) | isregional(pv.v) | isregional(pv.w) @inline regions(pv::PrescribedVelocityFields) = regions(pv[findfirst(isregional, (pv.u, pv.v, pv.w))]) diff --git a/test/regression_tests/rayleigh_benard_regression_test.jl b/test/regression_tests/rayleigh_benard_regression_test.jl index 363c115f9c..b37270cdeb 100644 --- a/test/regression_tests/rayleigh_benard_regression_test.jl +++ b/test/regression_tests/rayleigh_benard_regression_test.jl @@ -1,5 +1,6 @@ using Oceananigans.Grids: xnode, znode using Oceananigans.TimeSteppers: update_state! +using Oceananigans.BuoyancyFormulations: BuoyancyForce using Oceananigans.DistributedComputations: cpu_architecture, partition, reconstruct_global_grid function run_rayleigh_benard_regression_test(arch, grid_type) @@ -48,13 +49,13 @@ function run_rayleigh_benard_regression_test(arch, grid_type) bottom = BoundaryCondition(Value(), Δb)) model = NonhydrostaticModel(; grid, - timestepper = :QuasiAdamsBashforth2, - closure = ScalarDiffusivity(ν=ν, κ=κ), - tracers = (:b, :c), - buoyancy = BuoyancyTracer(), - boundary_conditions = (; b=bbcs), - hydrostatic_pressure_anomaly = CenterField(grid), - forcing = (; c=cforcing)) + timestepper = :QuasiAdamsBashforth2, + closure = ScalarDiffusivity(ν=ν, κ=κ), + tracers = (:b, :c), + buoyancy = BuoyancyForce(BuoyancyTracer()), + boundary_conditions = (; b=bbcs), + hydrostatic_pressure_anomaly = CenterField(grid), + forcing = (; c=cforcing)) # Lz/Nz will work for both the :regular and :vertically_unstretched grids. Δt = 0.01 * min(model.grid.Δxᶜᵃᵃ, model.grid.Δyᵃᶜᵃ, Lz/Nz)^2 / ν