Include background tracer and velocities in closure tendency terms

[glwagner]

cc @hdrake @liuchihl

[glwagner]

Probably before merging this we also want to use the total velocities in the stress computation

[glwagner]

But I wonder also if we should add a feature to BackgroundFields that allows this to be toggled on and off

[glwagner]

Ok @hdrake @liuchihl I have added the option to include or not the background field when computing closure fluxes.

They are noit included by default. If you want to include them you need to build the BackgroundFields explicitly by writing something like

background_fields = BackgroundFields(; background_closure_fluxes=true, b=B)

where B is the background buoyancy field as before. Then pass this to the model constructor instead of passing a NamedTuple.

Let me know if this seems like a good interface and also if it works.

[hdrake]

Looks like a good interface to me.

But is it on purpose that there is only support for background fields in the NonhydrostaticModel and not for the HydrostaticFreeSurfaceModel?

@liuchihl will test it in our configurations.

[glwagner]

Looks like a good interface to me.

But is it on purpose that there is only support for background fields in the NonhydrostaticModel and not for the HydrostaticFreeSurfaceModel?

@liuchihl will test it in our configurations.

Well yes, it's substantial effort to support background fields. So we implemented it in the nonhydrostatic model first. Nobody has requested having background fields for the hydrostatic model. It's not impossible but might require some thinking if it's going to work with the more complicated turbulence closures (like CATKE or k-epsilon) that sometimes get used for hydrostatic applications.

Since the nonhydrostatic model is fast (at least on one GPU) the hydrostatic model is mostly important for simulations on the sphere (although this statement needs to be evaluated more carefully for complex domains when we have a proper nonhydrostatic solver).

[hdrake]

@glwagner, can you help us understand how the diffusive flux calculations here interact with gradient boundary conditions?

We've run some simulations and seem to be getting the expected behavior in the interior, where buoyancy tendencies are due to the convergence of both the perturbation and background diffusive fluxes.

However, we are not seeing the expected behavior from the flux convergence due to boundary conditions (on the domain, not even considering immersed boundary conditions yet). My expectation was that the GradientBoundaryCondition would be applied to whichever tracer fields are passed along to the tendency functions:

             - ∇_dot_qᶜ(i, j, k, grid, closure, diffusivities, val_tracer_index, closure_c, clock, model_fields, buoyancy)
             - immersed_∇_dot_qᶜ(i, j, k, grid, closure_c, c_immersed_bc, closure, diffusivities, val_tracer_index, clock, model_fields)

which in our case should be the sum of the perturbation and background tracer fields. Instead, it seems that our solutions are behaving as though the GradientBoundaryCondition is only being applied to the perturbation fluxes at the boundaries. Indeed, when we remove the background field from the gradient values passed to GradientBoundaryCondition, we get the behavior we are looking for.

[glwagner]

@glwagner, can you help us understand how the diffusive flux calculations here interact with gradient boundary conditions?

We've run some simulations and seem to be getting the expected behavior in the interior, where buoyancy tendencies are due to the convergence of both the perturbation and background diffusive fluxes.

However, we are not seeing the expected behavior from the flux convergence due to boundary conditions (on the domain, not even considering immersed boundary conditions yet). My expectation was that the GradientBoundaryCondition would be applied to whichever tracer fields are passed along to the tendency functions:

             - ∇_dot_qᶜ(i, j, k, grid, closure, diffusivities, val_tracer_index, closure_c, clock, model_fields, buoyancy)
             - immersed_∇_dot_qᶜ(i, j, k, grid, closure_c, c_immersed_bc, closure, diffusivities, val_tracer_index, clock, model_fields)

which in our case should be the sum of the perturbation and background tracer fields. Instead, it seems that our solutions are behaving as though the GradientBoundaryCondition is only being applied to the perturbation fluxes at the boundaries. Indeed, when we remove the background field from the gradient values passed to GradientBoundaryCondition, we get the behavior we are looking for.

GradientBoundaryCondition isn't applied to the fields at all. The tracer gradient across the boundary is the quantity we need in order to compute fluxes across boundaries. Therefore when the immersed boundary flux is computed, we call

Oceananigans.jl/src/ImmersedBoundaries/immersed_boundary_condition.jl

Lines 100 to 101 in 738d172

	@inline right_gradient(i, j, k, ibg, κ, Δ, bc::GBC, c, clock, fields) = getbc(bc, i, j, k, ibg, clock, fields)
	@inline left_gradient(i, j, k, ibg, κ, Δ, bc::GBC, c, clock, fields) = getbc(bc, i, j, k, ibg, clock, fields)

Contrast this with the routine for ValueBoundaryCondition which is a bit more involved:

Oceananigans.jl/src/ImmersedBoundaries/immersed_boundary_condition.jl

Lines 103 to 107 in 738d172

	@inline function right_gradient(i, j, k, ibg, κ, Δ, bc::VBC, c, clock, fields)
	cᵇ = getbc(bc, i, j, k, ibg, clock, fields)
	cⁱʲᵏ = @inbounds c[i, j, k]
	return 2 * (cᵇ - cⁱʲᵏ) / Δ
	end

These gradients are then used to compute the flux, for example

Oceananigans.jl/src/ImmersedBoundaries/immersed_boundary_condition.jl

Lines 117 to 122 in 738d172

	@inline function _west_ib_flux(i, j, k, ibg, bc::VBCorGBC, (LX, LY, LZ), c, closure::ASD, K, id, clock, fields)
	Δ = Δx(index_left(i, LX), j, k, ibg, LX, LY, LZ)
	κ = h_diffusivity(i, j, k, ibg, flip(LX), LY, LZ, closure, K, id, clock)
	∇c = left_gradient(i, j, k, ibg, κ, Δ, bc, c, clock, fields)
	return - κ * ∇c
	end

Instead, it seems that our solutions are behaving as though the GradientBoundaryCondition is only being applied to the perturbation fluxes at the boundaries.

I don't know exactly what it means for a gradient to be applied to the field. Can you please clarify?

[glwagner]

Additionally can you help me understand why you are using GradientBoundaryCondition at all? I had thought that it's usually preferred to use FluxBoundaryCondition since this is often the preferred parameterization for large scale motions and basically we only use ValueBoundaryCondition for DNS (and GradientBoundaryCondition almost never). Curious to understand the class of problems that benefit from GradientBoundaryCondition.

[liuchihl]

@glwagner Thanks for implementing the total tracer diffusive flux at a high level. After running several tests, I found it to work exceptionally well! I conducted a series of tests: 1) comparing 1D vs 3D, 2) with and without the Coriolis force, and 3) with and without the immersed boundary. Everything looks great! Here are some simple examples on a rotated coordinate:

• 1D test with a small f:

nonconstantdiffusivity250days-theta.0.002_Nx4_Ny4_smallf_zlargerf.mp4

• 3D simulation with immersed grids:

nonconstantdiffusivity8days-theta.0.2_Nx4_Ny4_immersed_3Dfields_withcrossflux.mp4

The only caveat mentioned by @hdrake is that GradientBoundaryCondition is only being applied to the perturbation fluxes at the boundaries, i.e., GradientBoundaryCondition(-N^2*cos(θ)) is needed to make the total buoyancy gradient to be 0.

[hdrake]

Additionally can you help me understand why you are using GradientBoundaryCondition at all? [- @glwagner]

Flux boundary conditions are both conceptually and practically easier to implement when the boundary fluxes are zero or constant. They can be trickier when they depend on interior flow variables. In our case, for example, the boundary condition on the perturbation variable $b'$ (the buoyancy tracer in Oceananigans) is that the total diffusive flux should vanish, which means the perturbation diffusive flux needs to be minus the background diffusive flux.

Imposing a flux boundary condition requires knowing the diffusivity $\kappa$ right at the boundary. It is obvious how to implement this if the diffusivity is a constant, because the background diffusive flux is also a known constant, but less obvious how to do it when using a subgrid turbulence closure that yields a diffusivity that varies in space and time. @tomchor pointed out to use that we can sidestep this complexity if we just divide both sides of the boundary condition by $\kappa$, because then the boundary condition simply becomes that the buoyancy gradient across the boundary should just be equal to minus the background buoyancy gradient—a known constant in our problem.

[hdrake]

I don't know exactly what it means for a gradient to be applied to the field. Can you please clarify?

I just meant where in the code the gradient boundary conditions get imposed, which you've shown us is in the calculation of the gradients that feed into the downgradient diffusive fluxes that are used in the diffusive flux divergence contribution to the tracer tendencies. Thanks!

[glwagner]

I don't know exactly what it means for a gradient to be applied to the field. Can you please clarify?

I just meant where in the code the gradient boundary conditions get imposed, which you've shown us is in the calculation of the gradients that feed into the downgradient diffusive fluxes that are used in the diffusive flux divergence contribution to the tracer tendencies. Thanks!

Okay great. I would only add, I think it's clearer to think of the gradient as being used to diagnose the cross-boundary flux (rather than imposed).

I guess the point here is that there is actually an apparent flux of tracer into the perturbation field because of the presence of the background. So we are imagining that the background is being maintained by some large scale circulation which is ultimately the source of this apparent flux. FluxBoundaryCondition can be used if you know the diffusivity a priori but otherwise it does look like it will be simpler to use GradientBoundaryCondition to add this apparent flux contribution to the evolution of the perturbation.

[glwagner]

@glwagner Thanks for implementing the total tracer diffusive flux at a high level. After running several tests, I found it to work exceptionally well! I conducted a series of tests: 1) comparing 1D vs 3D, 2) with and without the Coriolis force, and 3) with and without the immersed boundary. Everything looks great! Here are some simple examples on a rotated coordinate:

• 1D test with a small f:

nonconstantdiffusivity250days-theta.0.002_Nx4_Ny4_smallf_zlargerf.mp4

• 3D simulation with immersed grids:

nonconstantdiffusivity8days-theta.0.2_Nx4_Ny4_immersed_3Dfields_withcrossflux.mp4
The only caveat mentioned by @hdrake is that GradientBoundaryCondition is only being applied to the perturbation fluxes at the boundaries, i.e., GradientBoundaryCondition(-N^2*cos(θ)) is needed to make the total buoyancy gradient to be 0.

Thanks @liuchihl !

[glwagner]

Hi @hdrake @liuchihl please review this PR and let me know if things look good.

I've tried to consistently add the background fields to the fields that go into computing both tracer fluxes and subgrid stresses. It makes things a little more complicated, but the code is more general now. I think with this change, things like a background shear will be used in computing a subgrid stress divergence. But the shear is still not used when computing a nonlinear diffusivity (this could be added though...)

[hdrake]

Thanks @glwagner! We haven't been using background velocities in our setups, so won't have as much to say on that yet. But we'll continue testing the background tracer fluxes.

@samlewin, are you using background shear in addition to background tracer fields in your configurations, or is the shear just in your initial conditions? This PR might be relevant.

[liuchihl]

What code are you running?

I was testing my internal tide example, where the code is written as a function, and the run script is used to adjust the parameters.
(I could work on an MWE if necessary.)

[liuchihl]

@liuchihl Have you tried running on the CPU to see if you get a more useful error message?

Thanks, I'll give that a try. I could be wrong, but my guess is that the issue might be related to the use of GPU.

[liuchihl]

@ali-ramadhan
Thanks for the suggestion. I tried running it on the CPU and got a different error message, which is shown below. Additionally, the error occurs when defining background fields in the model as such

background_fields = Oceananigans.BackgroundFields(; background_closure_fluxes=true, b=B̄_field)

But it runs fine like the usual way

background_fields = (; b=B̄_field)

[ Info: Initializing simulation...
ERROR: LoadError: type BackgroundFields has no field u
Stacktrace:
  [1] getproperty
    @ ./Base.jl:37 [inlined]
  [2] assemble_closure_velocities
    @ ~/.julia/packages/Oceananigans/xmqSH/src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl:35 [inlined]
  [3] u_velocity_tendency
    @ ~/.julia/packages/Oceananigans/xmqSH/src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl:92 [inlined]
  [4] cpu_compute_Gu!
    @ ~/.julia/packages/KernelAbstractions/491pi/src/macros.jl:291 [inlined]
  [5] __thread_run(tid::Int64, len::Int64, rem::Int64, obj::KernelAbstractions.Kernel{…}, ndrange::Nothing, iterspace::KernelAbstractions.NDIteration.NDRange{…}, args::Tuple{…}, dynamic::KernelAbstractions.NDIteration.DynamicCheck)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/491pi/src/cpu.jl:144
  [6] __run(obj::KernelAbstractions.Kernel{…}, ndrange::Nothing, iterspace::KernelAbstractions.NDIteration.NDRange{…}, args::Tuple{…}, dynamic::KernelAbstractions.NDIteration.DynamicCheck, static_threads::Bool)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/491pi/src/cpu.jl:111
  [7] (::KernelAbstractions.Kernel{…})(::Field{…}, ::Vararg{…}; ndrange::Nothing, workgroupsize::Nothing)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/491pi/src/cpu.jl:46
  [8] (::KernelAbstractions.Kernel{…})(::Field{…}, ::Vararg{…})
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/491pi/src/cpu.jl:39
  [9] launch!(::CPU, ::ImmersedBoundaryGrid{…}, ::Symbol, ::typeof(Oceananigans.Models.NonhydrostaticModels.compute_Gu!), ::Field{…}, ::Vararg{…}; include_right_boundaries::Bool, reduced_dimensions::Tuple{}, location::Nothing, active_cells_map::Nothing, kwargs::@Kwargs{})
    @ Oceananigans.Utils ~/.julia/packages/Oceananigans/xmqSH/src/Utils/kernel_launching.jl:168
 [10] launch!
    @ ~/.julia/packages/Oceananigans/xmqSH/src/Utils/kernel_launching.jl:154 [inlined]
 [11] #compute_interior_tendency_contributions!#26
    @ ~/.julia/packages/Oceananigans/xmqSH/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl:102 [inlined]
 [12] compute_interior_tendency_contributions!
    @ ~/.julia/packages/Oceananigans/xmqSH/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl:54 [inlined]
 [13] compute_tendencies!(model::NonhydrostaticModel{…}, callbacks::Vector{…})
    @ Oceananigans.Models.NonhydrostaticModels ~/.julia/packages/Oceananigans/xmqSH/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl:32
 [14] #apply_regionally!#56
    @ ~/.julia/packages/Oceananigans/xmqSH/src/Utils/multi_region_transformation.jl:121 [inlined]
 [15] apply_regionally!
    @ ~/.julia/packages/Oceananigans/xmqSH/src/Utils/multi_region_transformation.jl:118 [inlined]
 [16] macro expansion
    @ ~/.julia/packages/Oceananigans/xmqSH/src/Utils/multi_region_transformation.jl:206 [inlined]
 [17] update_state!(model::NonhydrostaticModel{…}, callbacks::Vector{…}; compute_tendencies::Bool)
    @ Oceananigans.Models.NonhydrostaticModels ~/.julia/packages/Oceananigans/xmqSH/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl:52
 [18] update_state! (repeats 2 times)
    @ ~/.julia/packages/Oceananigans/xmqSH/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl:20 [inlined]
 [19] initialize!(sim::Simulation{…})
    @ Oceananigans.Simulations ~/.julia/packages/Oceananigans/xmqSH/src/Simulations/run.jl:191
 [20] time_step!(sim::Simulation{…})
    @ Oceananigans.Simulations ~/.julia/packages/Oceananigans/xmqSH/src/Simulations/run.jl:112
 [21] run!(sim::Simulation{…}; pickup::Bool)
    @ Oceananigans.Simulations ~/.julia/packages/Oceananigans/xmqSH/src/Simulations/run.jl:97
 [22] top-level scope
    @ /scratch/bcpi/cliu28/internal-tide-mixing/run_internal_tide_testing.jl:34
 [23] include(fname::String)
    @ Base.MainInclude ./client.jl:489
 [24] top-level scope
    @ REPL[2]:1
in expression starting at /scratch/bcpi/cliu28/internal-tide-mixing/run_internal_tide_testing.jl:34

[glwagner]

@liuchihl I think the best way forward is to write a simple test that illustrates the error. Then I can help fix the error to make the test pass. Once that is done, we may be ready to merge this PR. What do you think?

PS it is always best to work with minimal examples, and to paste code directly into a discussion stream (rather than providing links). This will help us keep up an efficient workflow.

[liuchihl]

@liuchihl I think the best way forward is to write a simple test that illustrates the error. Then I can help fix the error to make the test pass. Once that is done, we may be ready to merge this PR. What do you think?

Sure, I agree with that! I will work on that and let you know how it goes.

PS it is always best to work with minimal examples, and to paste code directly into a discussion stream (rather than providing links). This will help us keep up an efficient workflow.

For sure, sorry about that, I haven't been able to create an MWE for this specific issue because I don't understand the problem yet.

[glwagner]

This does it:

using Oceananigans
grid = RectilinearGrid(size=1, x=(0, 1), topology=(Periodic, Flat, Flat))
B(args...) = 0
background_fields = Oceananigans.BackgroundFields(; background_closure_fluxes=true, b=B)
model = NonhydrostaticModel(; grid, background_fields)

You can use this in a test, something like

B(args...) = 0
function time_step_background_fields_with_closure_fluxes(arch)
    grid = RectilinearGrid(arch, size=1, x=(0, 1), topology=(Periodic, Flat, Flat))
    background_fields = Oceananigans.BackgroundFields(; background_closure_fluxes=true, b=B)
    model = NonhydrostaticModel(; grid, background_fields)
    time_step!(model, 1)
    return true
end

@test time_step_background_fields_with_closure_fluxes(arch)

If it throws an error, the test fails.

This is a basic unit test which is always a good idea to write first. The next step is to write a test that confirms the functionality works correctly. This is an example of a hierarchy of tests proceeding from simple to complex.

[glwagner]

Since the code is just a few lines, it's easy to figure out what's going on in the REPL. The top of the error message says:

[2] assemble_closure_velocities
    @ ~/.julia/packages/Oceananigans/xmqSH/src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl:35 [inlined]

So let's go look at this line. It is:

Oceananigans.jl/src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl

Line 35 in f89445d

u = sum_fields(velocities.u, background_fields.u)

Ok. Does background_fields have the property u? At the REPL, using the MWE, I find:

julia> background_fields.u
ERROR: type BackgroundFields has no field u
Stacktrace:
 [1] getproperty(x::Oceananigans.Models.NonhydrostaticModels.BackgroundFields{true, @NamedTuple{…}, @NamedTuple{…}}, f::Symbol)
   @ Base ./Base.jl:37
 [2] top-level scope
   @ REPL[7]:1
Some type information was truncated. Use `show(err)` to see complete types.

Ok, we found the error. But where is u? writing julia> background_fields. and pressing TAB yields

julia> background_fields.
tracers
velocities

Ok, so maybe we want background_fields.velocities.u. Looking into the source code for background fields:

Oceananigans.jl/src/Models/NonhydrostaticModels/background_fields.jl

Lines 32 to 38 in f89445d

	struct BackgroundFields{Q, U, C}
	velocities :: U
	tracers :: C
	function BackgroundFields{Q}(velocities::U, tracers::C) where {Q, U, C}
	return new{Q, U, C}(velocities, tracers)
	end
	end

confirms what we found, that BackgroundFields.u does not exist. Just to confirm, we go back to the REPL and try:

julia> background_fields.velocities.u
ZeroField{Int64}

so yes, using velocities.u should fix this particular bug.

This is an easy fix. Just to go through the motions and get a hang for how this works, let's add the test first, and then make the fix.

[glwagner]

I fixed the conflicts so this PR should be ready for work

[liuchihl]

Thanks for explaining these in detail!

[hdrake]

To summarize where we're at: we need at least two unit tests:

• a functionality test, like as proposed and implemented by @glwagner [here]:(Include background tracer and velocities in closure tendency terms #3646 (comment))

B(args...) = 0
function time_step_background_fields_with_closure_fluxes(arch)
    grid = RectilinearGrid(arch, size=1, x=(0, 1), topology=(Periodic, Flat, Flat))
    background_fields = Oceananigans.BackgroundFields(; background_closure_fluxes=true, b=B)
    model = NonhydrostaticModel(; grid, background_fields)
    time_step!(model, 1)
    return true
end

@test time_step_background_fields_with_closure_fluxes(arch)

• a correctness test, like the 1D diffusion example I proposed here but have not yet implemented:

The simplest test would be a 1D simulation with a background linear tracer profile with boundary conditions of Flux = − κ ∂ z c at the top and bottom. The correct solution would be that nothing happens when the simulation is time-stepped.

@liuchihl, can you create a PR into Greg's branch that adds both of these tests to the test suite?

[liuchihl]

Sure, I will work on the unit tests.

[glwagner]

And note that we only really need a unit test to merge this, same for many things.

Correctness is a high bar and it's ok if we can't come up wtih someone right away. Sometimes we don't really have a way to do correctness and our best option is something like a regression test (eg we verified it worked at one point, so we just make sure that it keeps returning that same result).

One way to go partway towards functional test but not all the way to "correctness" is to simply test that a simple set up returns a different result when closure fluxes are included vs not (for example).

[glwagner]

Is there still interest in this PR?

[liuchihl]

@glwagner Thank you for following up. Yes, we’re still interested in this great PR. Apologies for the delay—I’m planning to work on the unit test later this month.

[glwagner]

No worries at all. Since you're interested I will try to keep it up to date.

[liuchihl]

@glwagner I've added a unit test for the background flux divergence in this PR. The test validates the model result by comparing cases with and without background flux divergence enabled. I'd appreciate your thoughts and review on this unit test!