Materialize buoyancy gradients for use in parameterizations (#4837)

* cmopute

* fix tests

* overwrite

* Remove @const from kernel function parameters

* Update src/Models/NonhydrostaticModels/nonhydrostatic_model.jl

Co-authored-by: Gregory L. Wagner <[email protected]>

* Update src/BuoyancyFormulations/buoyancy_force.jl

Co-authored-by: Gregory L. Wagner <[email protected]>

* Update src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl

Co-authored-by: Gregory L. Wagner <[email protected]>

* regularize -> materialize

* regularize -> materialize again

* more precompute -> materialize

* bugfixes

* one fix

* bugfix

* another bugfix

* fix multiregion

* do not do this

* make sure the correct parameters are passed

* fix

* Update multi_region_models.jl

* fix all bugs

---------

Co-authored-by: Gregory L. Wagner <[email protected]>

Author: simone-silvestri

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

src/BuoyancyFormulations/nonlinear_equation_of_state.jl

@@ -1,5 +1,5 @@
 using Oceananigans.Fields: AbstractField
-using Oceananigans.Grids: znode
+using Oceananigans.Grids: znode, ZFlatGrid
 using Oceananigans.Operators: Δzᶜᶜᶠ, Δzᶜᶜᶜ
 
 const c = Center()
@@ -16,6 +16,11 @@ const f = Face()
     ifelse(k < 1,           znode(i, j,           1, grid, c, c, f) + (1 - k) * Δzᶜᶜᶜ(i, j, 1, grid),
     ifelse(k > grid.Nz + 1, znode(i, j, grid.Nz + 1, grid, c, c, f) + (k - grid.Nz - 1) * Δzᶜᶜᶜ(i, j, grid.Nz, grid),
                             znode(i, j,           k, grid, c, c, f)))
+
+# Fallbacks for ZFlatGrids (assumed to be at z = 0 for buoyancy purposes)
+@inline Zᶜᶜᶜ(i, j, k, grid::ZFlatGrid) = zero(grid)
+@inline Zᶜᶜᶠ(i, j, k, grid::ZFlatGrid) = zero(grid)
+
 # Dispatch shenanigans
 @inline θ_and_sᴬ(i, j, k, θ::AbstractArray, sᴬ::AbstractArray) = @inbounds θ[i, j, k], sᴬ[i, j, k]
 @inline θ_and_sᴬ(i, j, k, θ::Number,        sᴬ::AbstractArray) = @inbounds θ, sᴬ[i, j, k]

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:

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)
 

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

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

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, validate_unit_vector
 using Oceananigans.Models.HydrostaticFreeSurfaceModels: AbstractFreeSurface
 using Oceananigans.TimeSteppers: AbstractTimeStepper, QuasiAdamsBashforth2TimeStepper
 using Oceananigans.Models: PrescribedVelocityFields
@@ -53,6 +54,15 @@ for T in Types
     end
 end
 
+# TODO: For the moment, buoyancy gradients cannot be precomputed in MultiRegionModels
+function BuoyancyForce(grid::MultiRegionGrids, formulation::AbstractBuoyancyFormulation; 
+                       gravity_unit_vector=NegativeZDirection(), 
+                       materialize_gradients=false) 
+
+    gravity_unit_vector = validate_unit_vector(gravity_unit_vector)
+    return BuoyancyForce(formulation, gravity_unit_vector, nothing)
+end
+
 @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))])
 

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 / ν