Materialize buoyancy gradients for use in parameterizations

src/BuoyancyFormulations/buoyancy_force.jl

@@ -1,15 +1,17 @@
-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())
 
@@ -44,11 +46,31 @@ NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
 └── coriolis: Nothing
 ```
 """
-function BuoyancyForce(formulation; gravity_unit_vector=NegativeZDirection())
+function BuoyancyForce(grid; formulation, 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)
+
+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 +87,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 +115,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/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.Biogeochemistry: validate_biogeochemistry, AbstractBiogeochemistry, biogeochemical_auxiliary_fields
 using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions
 using Oceananigans.Fields: Field, tracernames, VelocityFields, TracerFields, CenterField
@@ -171,7 +171,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