Skip to content

Add reset_fluxes! function to FluxCalculator #1908

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
akshaysridhar wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add reset_fluxes! function to FluxCalculator #1908

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions ext/ClimaCouplerCMIPExt/clima_seaice.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -429,6 +429,21 @@ function FluxCalculator.update_turbulent_fluxes!(sim::ClimaSeaIceSimulation, fie
return nothing
end

function FluxCalculator.reset_fluxes!(sim::ClimaSeaIceSimulation)
si_flux_heat = sim.ice.model.external_heat_fluxes.top
if si_flux_heat isa OC.Field
OC.interior(si_flux_heat, :, :, 1) .= 0
end

if !isnothing(sim.ice.model.dynamics)
si_flux_u = sim.ice.model.dynamics.external_momentum_stresses.top.u
si_flux_v = sim.ice.model.dynamics.external_momentum_stresses.top.v
si_flux_u isa OC.Field && (OC.interior(si_flux_u, :, :, 1) .= 0)
si_flux_v isa OC.Field && (OC.interior(si_flux_v, :, :, 1) .= 0)
end
return nothing
end

function Interfacer.update_field!(sim::ClimaSeaIceSimulation, ::Val{:area_fraction}, field)
sim.area_fraction .= field
return nothing
Expand Down
17 changes: 11 additions & 6 deletions ext/ClimaCouplerCMIPExt/oceananigans.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -576,6 +576,14 @@ function Interfacer.update_field!(sim::OceananigansSimulation, ::Val{:area_fract
return nothing
end

function FluxCalculator.reset_fluxes!(sim::OceananigansSimulation)

@juliasloan25 juliasloan25 May 7, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could you add a brief docstring please?

oc_flux_T = surface_flux(sim.ocean.model.tracers.T)
OC.interior(oc_flux_T, :, :, 1) .= 0
oc_flux_S = surface_flux(sim.ocean.model.tracers.S)
OC.interior(oc_flux_S, :, :, 1) .= 0
return nothing
end

"""
FieldExchanger.update_sim!(sim::OceananigansSimulation, csf)

Expand All @@ -585,8 +593,9 @@ by the coupler.
Update the portion of the surface_fluxes for T and S that is due to radiation and
precipitation. The rest will be updated in `update_turbulent_fluxes!`.

This function sets the surface fluxes directly, overwriting any previous values.
Additional contributions will be made in `update_turbulent_fluxes!` and `ocean_seaice_fluxes!`.
Flux tendencies are expected to be reset at the start of the coupler timestep via
[`FluxCalculator.reset_fluxes!`](@ref). Additional contributions are made in `update_turbulent_fluxes!`
and `ocean_seaice_fluxes!`.

A note on sign conventions:
ClimaAtmos and Oceananigans both use the convention that a positive flux is an upward flux.
Expand All @@ -600,12 +609,8 @@ function FieldExchanger.update_sim!(sim::OceananigansSimulation, csf)
grid = sim.ocean.model.grid
Nz = grid.Nz
ice_concentration = OC.interior(sim.ice_concentration, :, :, 1)

# Reset fluxes to 0 at the start of the step
oc_flux_T = surface_flux(sim.ocean.model.tracers.T)
OC.interior(oc_flux_T, :, :, 1) .= 0
oc_flux_S = surface_flux(sim.ocean.model.tracers.S)
OC.interior(oc_flux_S, :, :, 1) .= 0

# for masking out the poles
polar_mask = sim.remapping.polar_mask
Expand Down
20 changes: 20 additions & 0 deletions ext/ClimaCouplerClimaAtmosExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -539,6 +539,26 @@ function FluxCalculator.update_turbulent_fluxes!(sim::ClimaAtmosSimulation, fiel
return nothing
end

function FluxCalculator.reset_fluxes!(sim::ClimaAtmosSimulation)
integrator = sim.integrator
sfc = integrator.p.precomputed.sfc_conditions

fill!(sfc.ρ_flux_h_tot, zero(eltype(sfc.ρ_flux_h_tot)))
fill!(sfc.ρ_flux_uₕ, zero(eltype(sfc.ρ_flux_uₕ)))

if hasmoisture(integrator)
fill!(sfc.ρ_flux_q_tot, zero(eltype(sfc.ρ_flux_q_tot)))
fill!(integrator.p.precomputed.surface_rain_flux, zero(eltype(integrator.p.precomputed.surface_rain_flux)))
fill!(integrator.p.precomputed.surface_snow_flux, zero(eltype(integrator.p.precomputed.surface_snow_flux)))
end

fill!(sfc.obukhov_length, zero(eltype(sfc.obukhov_length)))
fill!(sfc.ustar, zero(eltype(sfc.ustar)))
fill!(sfc.buoyancy_flux, zero(eltype(sfc.buoyancy_flux)))

return nothing
end

"""
Extend Interfacer.add_coupler_fields! to add the fields required for ClimaAtmosSimulation.

Expand Down
4 changes: 4 additions & 0 deletions src/FieldExchanger.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -405,6 +405,10 @@ function exchange!(cs::Interfacer.CoupledSimulation)
import_atmos_fields!(cs.fields, cs.model_sims)
import_combined_surface_fields!(cs.fields, cs.model_sims)

# NOTE: Component-model flux accumulators are expected to have been zeroed by
# `FluxCalculator.reset_fluxes!(cs)` in the top-level coupler step. We do not
# re-zero here so the ordering remains explicit in `SimCoordinator.step!`.

# Update the component model simulations with the coupler fields
update_model_sims!(cs.model_sims, cs.fields)
return nothing
Expand Down
27 changes: 27 additions & 0 deletions src/FluxCalculator.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,37 @@ import ..Interfacer, ..Utilities

export turbulent_fluxes!,
get_surface_params,
reset_fluxes!,
update_turbulent_fluxes!,
compute_surface_fluxes!,
ocean_seaice_fluxes!

"""
reset_fluxes!(cs::Interfacer.CoupledSimulation)
reset_fluxes!(sim::Interfacer.AbstractComponentSimulation)

Reset any internally accumulated surface flux tendencies on each component model in `cs`
(or on a single `sim`) at the start of a coupler timestep, before flux contributions are
added in `FieldExchanger.exchange!`/`FieldExchanger.update_sim!`,
[`FluxCalculator.update_turbulent_fluxes!`](@ref), and/or
[`FluxCalculator.ocean_seaice_fluxes!`](@ref).

The `CoupledSimulation` method is intended to be called explicitly from the top-level
coupler loop (`SimCoordinator.step!`) between `update_surface_fractions!` and `exchange!`,
so the reset ordering is visible at the call site rather than hidden inside the exchange.

By default, the per-`sim` method is a no-op. Component models that accumulate fluxes across
multiple coupler updates should extend that method.
"""
function reset_fluxes!(cs::Interfacer.CoupledSimulation)
for sim in cs.model_sims
reset_fluxes!(sim)
end
return nothing
end

reset_fluxes!(sim::Interfacer.AbstractComponentSimulation) = nothing

function turbulent_fluxes!(cs::Interfacer.CoupledSimulation)
return turbulent_fluxes!(cs.fields, cs.model_sims, cs.thermo_params)
end
Expand Down
5 changes: 5 additions & 0 deletions src/SimCoordinator.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,11 @@ function step!(cs::Interfacer.CoupledSimulation)
# Update the surface fractions for surface models
FieldExchanger.update_surface_fractions!(cs)

# Zero any internally-accumulated surface flux tendencies on each component model
# before this coupler step starts adding new contributions in `exchange!`,
# `turbulent_fluxes!`, and `ocean_seaice_fluxes!`.
FluxCalculator.reset_fluxes!(cs)

# Exchange all non-turbulent flux fields between models, including radiative and precipitation fluxes
FieldExchanger.exchange!(cs)

Expand Down
Loading