Trouble Saving Terms in a Tracer Budget

[kenflat2]

Hi all,

I'm hoping to compute a temporal average of a tracer budget for my simulation. I would like to compute the value of each of the terms on the right of the following budget

$\langle\partial_t c\rangle = - \langle\nabla\cdot(uc)\rangle + \langle\nabla\cdot(\kappa\nabla c)\rangle + \langle F_c\rangle$

for every grid cell in the simulation.

To do this, I wrote the following function which takes the current simulation as input and adds an output writer that saves the fields of interest. I am using a HydrostaticFreeSurfaceModel and lifted code from its internals in an effort to recompute the fields exactly.

function add_buoyancy_budget_diagnostics!(simulation, checkpoint_time, budget_duration, directory)

    model = simulation.model

    arch = model.architecture
    grid = model.grid

    tracer_index = 1
    tracer_name = :T

    @inbounds c_tendency    = model.timestepper.Gⁿ[tracer_name]
    @inbounds c_advection   = model.advection[tracer_name]
    @inbounds c_forcing     = model.forcing[tracer_name]
    @inbounds c_immersed_bc = immersed_boundary_condition(model.tracers[tracer_name])

    # Recreate arguments exactly as they are in Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendency_contributions!
    val_tracer_index = Val(tracer_index)
    val_tracer_name = Val(tracer_name)
    advection = c_advection
    closure = model.closure
    c_immersed_bc = c_immersed_bc
    buoyancy = model.buoyancy
    biogeochemistry = model.biogeochemistry
    velocities = model.velocities
    free_surface = model.free_surface
    tracers = model.tracers
    diffusivities = model.diffusivity_fields
    auxiliary_fields = model.auxiliary_fields
    clock = model.clock
    forcing = c_forcing

    # Code here is lifted verbatim from hydrostatic_free_surface_tracer_tendency in 
    # the file Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl
    @inbounds c = tracers[tracer_index]
    model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers),
                         auxiliary_fields,
                         biogeochemical_auxiliary_fields(biogeochemistry))

    biogeochemical_velocities = biogeochemical_drift_velocity(biogeochemistry, val_tracer_name)
    closure_velocities = closure_turbulent_velocity(closure, diffusivities, val_tracer_name)

    total_velocities = sum_of_velocities(velocities, biogeochemical_velocities, closure_velocities)
    total_velocities = with_advective_forcing(forcing, total_velocities)

    # The tracer tendency is the sum of the following features...

    #  - div_Uc(i, j, k, grid, advection, total_velocities, c)
    #  - ∇_dot_qᶜ(i, j, k, grid, closure, diffusivities, val_tracer_index, c, clock, model_fields, buoyancy)
    #  - immersed_∇_dot_qᶜ(i, j, k, grid, c, c_immersed_bc, closure, diffusivities, val_tracer_index, clock, model_fields)
    #  + biogeochemical_transition(i, j, k, grid, biogeochemistry, val_tracer_name, clock, model_fields)
    #  + forcing(i, j, k, grid, clock, model_fields))

    # ... which we recreate in kernel functions
    ∇uc = KernelFunctionOperation{Center, Center, Center}(div_Uc, grid, advection, total_velocities, c)
    ∇qc = KernelFunctionOperation{Center, Center, Center}(∇_dot_qᶜ, grid, closure, diffusivities, val_tracer_index, c, clock, model_fields, buoyancy)
    immersed_∇qc = KernelFunctionOperation{Center, Center, Center}(immersed_∇_dot_qᶜ, grid, c, c_immersed_bc, closure, diffusivities, val_tracer_index, clock, model_fields)
    Fc  = KernelFunctionOperation{Center, Center, Center}(forcing, grid, clock, model_fields)

    writer_fluxes = JLD2Writer(model, (; ∇uc, ∇qc, immersed_∇qc, Fc), 
                        schedule = AveragedTimeInterval(budget_duration),
                        filename = joinpath(directory, "buoyancy_budget_averaged.jld2"))

    simulation.output_writers[:buoyancy_budget_averaged] = writer_fluxes
end

However, the code above only accurately saves the div_Uc term. The other terms ∇_dot_qᶜ, immersed_∇_dot_qᶜ and forcing are all zero. My model has both an explicit vertical diffusivity and a tracer forcing at the surface, so I expect ∇_dot_qᶜ and Fc to be nonzero. I'm not sure what could be different about the div_Uc term compared to the others, so I'm lost at the moment.

Any help would be greatly appreciated, thank you!

[glwagner]

Maybe this is a bug for averaging multiple terms. What happens if you save each term with a different output writer?

[glwagner]

Can you use @which to figure out which method is being called?

[kenflat2]

The version of the diffusive flux divergence being called is

∇_dot_qᶜ(i, j, k, grid, closures::Tuple{Any, Any}, Ks, args...) @ Oceananigans.TurbulenceClosures ~/.julia/packages/Oceananigans/Xawvn/src/TurbulenceClosures/closure_tuples.jl:41

which I believe leads to a call to diffusive_flux_z down the line with the following type signature

diffusive_flux_z(i, j, k, grid, ::Oceananigans.TurbulenceClosures.AbstractTimeDiscretization, args...) @ Oceananigans.TurbulenceClosures ~/.julia/packages/Oceananigans/Xawvn/src/TurbulenceClosures/implicit_explicit_time_discretization.jl:41

that returns zero incorrectly. I found that the diffusive_flux_z type with type signature

diffusive_flux_z(i, j, k, grid, closures::Tuple{Any, Any}, Ks, args...) @ Oceananigans.TurbulenceClosures ~/.julia/packages/Oceananigans/Xawvn/src/TurbulenceClosures/closure_tuples.jl:41

returns nonzero values.

[glwagner]

Ah. If your closure is using an implicit time discretization, then indeed it does not make an explicit contribution to the tendency. You can compute the explicit contribution by reforming the closure with ExplicitTimeDiscretization().

[glwagner]

The immersed contribution is only non-zero if you are using ValueBoundaryCondition or GradientBoundaryCondition on an immersed boundary in conjunction with ScalarDiffusivity

[kenflat2]

Would it be reasonable to just change the closures in the model to ExplicitTimeDiscretization(), or would that meaningfully change the model physics?

[kenflat2]

I was able to save the budget terms by reforming the closures as suggested, thank you for that tip!

At present the budget doesn't close. Our model is a reentrant channel and the budget is averaged horizontally so the vertical eddy flux of buoyancy should balance the vertical diffusive flux, but the former is roughly double the latter. The diffusive flux comes from an imposed high vertical diffusivity in the top 150m.

We also use a convective adjustment parameterization, could this cause discrepancies between the diffusive fluxes used by the model and what is stored in the output writer?