From bb6e4955439e30b92f5082fe742295179fc69eeb Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Thu, 4 Dec 2025 16:53:04 -0800 Subject: [PATCH 1/3] remove allocations in combine_surfaces --- src/FieldExchanger.jl | 77 +++++++++++++++-------------------- test/field_exchanger_tests.jl | 10 ++++- 2 files changed, 41 insertions(+), 46 deletions(-) diff --git a/src/FieldExchanger.jl b/src/FieldExchanger.jl index d3ec982352..09a55ab19a 100644 --- a/src/FieldExchanger.jl +++ b/src/FieldExchanger.jl @@ -145,14 +145,10 @@ import_atmos_fields!(csf, ::Interfacer.ComponentModelSimulation, atmos_sim) = no import_combined_surface_fields!(csf, model_sims) Updates the coupler with the surface properties. The `Interfacer.get_field` -functions for (`:surface_temperature`, `:surface_direct_albedo`, +functions for (`:emissivity`, `:surface_temperature`, `:surface_direct_albedo`, `:surface_diffuse_albedo`) need to be specified for each surface model. -Note: The calculation of surface humidity done here uses atmospheric properties stored in -the coupled fields. For these values to be correct, this function should be called -after `import_atmos_fields!` in a timestep. - -Note 2: Not all surface fields are imported here. Some quantities are retrieved +Note: Not all surface fields are imported here. Some quantities are retrieved from each surface model when surface fluxes are computed, in `compute_surface_fluxes!`. # Arguments @@ -160,20 +156,10 @@ from each surface model when surface fluxes are computed, in `compute_surface_fl - `model_sims`: [NamedTuple] containing `ComponentModelSimulation`s. """ function import_combined_surface_fields!(csf, model_sims) - combine_surfaces!(csf.emissivity, model_sims, Val(:emissivity), csf.scalar_temp1) + combine_surfaces!(csf, model_sims, Val(:emissivity)) combine_surfaces!(csf, model_sims, Val(:surface_temperature)) - combine_surfaces!( - csf.surface_direct_albedo, - model_sims, - Val(:surface_direct_albedo), - csf.scalar_temp1, - ) - combine_surfaces!( - csf.surface_diffuse_albedo, - model_sims, - Val(:surface_diffuse_albedo), - csf.scalar_temp1, - ) + combine_surfaces!(csf, model_sims, Val(:surface_direct_albedo)) + combine_surfaces!(csf, model_sims, Val(:surface_diffuse_albedo)) return nothing end @@ -309,55 +295,55 @@ function step_model_sims!(cs::Interfacer.CoupledSimulation) end """ - combine_surfaces!(combined_field::CC.Fields.Field, sims, field_name::Val, temp1) + combine_surfaces!(csf, sims, field_name_val::Val{field_name}) where {field_name} -Sums the fields, specified by `field_name`, weighted by the respective area fractions of all -surface simulations. THe result is saved in `combined_field`. +Sums the surface fields specified by `field_name_val`, weighted by the respective area fractions +of all surface simulations. The result is saved in the coupler field specified by `field_name_val`. For surface temperature, upward longwave radiation is computed from the temperatures of each surface, weighted by their area fractions, and then the combined temperature is computed from the combined upward longwave radiation. # Arguments -- `combined_field`: [CC.Fields.Field] output object containing weighted values. +- `csf`: [NamedTuple] containing coupler fields. Note: For the surface temperature, all coupler fields are passed in a NamedTuple. -- `sims`: [NamedTuple] containing simulations . -- `field_name`: [Val] containing the name Symbol of the field t be extracted by the `Interfacer.get_field` functions. -- `scalar_temp`: [CC.Fields.Field] temporary scalar-valued field for intermediate calculations. - Omitted for surface temperature method. +- `sims`: [NamedTuple] containing simulations. +- `field_name_val`: [Val] containing the name Symbol of the field to be extracted by the `Interfacer.get_field` functions. # Example - `combine_surfaces!(temp_field, cs.model_sims, Val(:emissivity))` """ -function combine_surfaces!(combined_field, sims, field_name, scalar_temp) - boundary_space = axes(combined_field) +function combine_surfaces!(csf, sims, field_name_val::Val{field_name}) where {field_name} + # Extract the coupler field we are updating + combined_field = getproperty(csf, field_name) combined_field .= 0 + for sim in sims if sim isa Interfacer.SurfaceModelSimulation - # Store the area fraction of this simulation in `scalar_temp` - Interfacer.get_field!(scalar_temp, sim, Val(:area_fraction)) + # Store the area fraction of this simulation in `scalar_temp` and rename for clarity + Interfacer.get_field!(csf.scalar_temp1, sim, Val(:area_fraction)) + area_fraction = csf.scalar_temp1 + + # Remap the surface field onto a coupler temporary field to avoid allocation + Interfacer.get_field!(csf.scalar_temp2, sim, field_name_val) + surface_field = csf.scalar_temp2 + # Zero out the contribution from this surface if the area fraction is zero. # Note that multiplying by `area_fraction` is not sufficient in the case of NaNs combined_field .+= - scalar_temp .* - ifelse.( - scalar_temp .≈ 0, - zero(combined_field), - Interfacer.get_field(sim, field_name, boundary_space), - ) + area_fraction .* + ifelse.(area_fraction .≈ 0, zero(combined_field), surface_field) end end return nothing end -function combine_surfaces!(csf, sims, field_name::Val{:surface_temperature}) +function combine_surfaces!(csf, sims, field_name_val::Val{:surface_temperature}) # extract the coupler fields we need to get the surface temperature T_sfc = csf.T_sfc emissivity_sfc = csf.emissivity - boundary_space = axes(T_sfc) - FT = CC.Spaces.undertype(boundary_space) - - T_sfc .= FT(0) + FT = eltype(T_sfc) + T_sfc .= zero(FT) for sim in sims if sim isa Interfacer.SurfaceModelSimulation # Store the area fraction and emissivity of this simulation in temp fields @@ -366,6 +352,10 @@ function combine_surfaces!(csf, sims, field_name::Val{:surface_temperature}) Interfacer.get_field!(csf.scalar_temp2, sim, Val(:emissivity)) emissivity_sim = csf.scalar_temp2 + # Remap the surface field onto a coupler temporary field to avoid allocation + Interfacer.get_field!(csf.scalar_temp3, sim, Val(:surface_temperature)) + T_sfc_sim = csf.scalar_temp3 + # Zero out the contribution from this surface if the area fraction is zero. # Note that multiplying by `area_fraction` is not sufficient in the case of NaNs # Compute upward longwave radiation from surface temperature for this simulation @@ -374,8 +364,7 @@ function combine_surfaces!(csf, sims, field_name::Val{:surface_temperature}) ifelse.( area_fraction .≈ 0, zero(T_sfc), - emissivity_sim .* - Interfacer.get_field(sim, field_name, boundary_space) .^ FT(4), + emissivity_sim .* T_sfc_sim .^ FT(4), ) end end diff --git a/test/field_exchanger_tests.jl b/test/field_exchanger_tests.jl index a2b1b33f70..9ca5995de7 100644 --- a/test/field_exchanger_tests.jl +++ b/test/field_exchanger_tests.jl @@ -269,9 +269,15 @@ for FT in (Float32, Float64) Interfacer.get_field(sims.c, var_name), Interfacer.get_field(sims.d, var_name), ) - temp_field = CC.Fields.zeros(test_space) - FieldExchanger.combine_surfaces!(combined_field, sims, var_name, temp_field) + # Create a coupler fields NamedTuple with the field we want to combine + csf = (; + random = combined_field, + scalar_temp1 = CC.Fields.zeros(test_space), + scalar_temp2 = CC.Fields.zeros(test_space), + ) + + FieldExchanger.combine_surfaces!(csf, sims, var_name) @test combined_field == fill(FT(sum(fractions .* fields)), test_space) end From 3de55c3654e57680c342e273b5ad1aa23711ec43 Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Fri, 5 Dec 2025 10:12:47 -0800 Subject: [PATCH 2/3] try to improve performance --- src/FieldExchanger.jl | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/FieldExchanger.jl b/src/FieldExchanger.jl index 09a55ab19a..91d3d0f7e8 100644 --- a/src/FieldExchanger.jl +++ b/src/FieldExchanger.jl @@ -295,10 +295,10 @@ function step_model_sims!(cs::Interfacer.CoupledSimulation) end """ - combine_surfaces!(csf, sims, field_name_val::Val{field_name}) where {field_name} + combine_surfaces!(csf, sims, ::Val{field_name}) where {field_name} -Sums the surface fields specified by `field_name_val`, weighted by the respective area fractions -of all surface simulations. The result is saved in the coupler field specified by `field_name_val`. +Sums the surface fields specified by `field_name`, weighted by the respective area fractions +of all surface simulations. The result is saved in the coupler field specified by `field_name`. For surface temperature, upward longwave radiation is computed from the temperatures of each surface, weighted by their area fractions, and then the combined temperature @@ -313,7 +313,7 @@ is computed from the combined upward longwave radiation. # Example - `combine_surfaces!(temp_field, cs.model_sims, Val(:emissivity))` """ -function combine_surfaces!(csf, sims, field_name_val::Val{field_name}) where {field_name} +function combine_surfaces!(csf, sims, ::Val{field_name}) where {field_name} # Extract the coupler field we are updating combined_field = getproperty(csf, field_name) combined_field .= 0 @@ -325,7 +325,7 @@ function combine_surfaces!(csf, sims, field_name_val::Val{field_name}) where {fi area_fraction = csf.scalar_temp1 # Remap the surface field onto a coupler temporary field to avoid allocation - Interfacer.get_field!(csf.scalar_temp2, sim, field_name_val) + Interfacer.get_field!(csf.scalar_temp2, sim, Val(field_name)) surface_field = csf.scalar_temp2 # Zero out the contribution from this surface if the area fraction is zero. @@ -337,7 +337,7 @@ function combine_surfaces!(csf, sims, field_name_val::Val{field_name}) where {fi end return nothing end -function combine_surfaces!(csf, sims, field_name_val::Val{:surface_temperature}) +function combine_surfaces!(csf, sims, ::Val{:surface_temperature}) # extract the coupler fields we need to get the surface temperature T_sfc = csf.T_sfc emissivity_sfc = csf.emissivity From 4a630a723e997ca458e57ffb5e5adddfe8457c3d Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Fri, 5 Dec 2025 10:29:16 -0800 Subject: [PATCH 3/3] remove runtime dispatch --- src/FieldExchanger.jl | 56 +++++++++++++++++++++++++---------- test/field_exchanger_tests.jl | 2 +- 2 files changed, 42 insertions(+), 16 deletions(-) diff --git a/src/FieldExchanger.jl b/src/FieldExchanger.jl index 91d3d0f7e8..1fe0c8d2d8 100644 --- a/src/FieldExchanger.jl +++ b/src/FieldExchanger.jl @@ -156,10 +156,20 @@ from each surface model when surface fluxes are computed, in `compute_surface_fl - `model_sims`: [NamedTuple] containing `ComponentModelSimulation`s. """ function import_combined_surface_fields!(csf, model_sims) - combine_surfaces!(csf, model_sims, Val(:emissivity)) - combine_surfaces!(csf, model_sims, Val(:surface_temperature)) - combine_surfaces!(csf, model_sims, Val(:surface_direct_albedo)) - combine_surfaces!(csf, model_sims, Val(:surface_diffuse_albedo)) + combine_surfaces!(csf.emissivity, csf, model_sims, Val(:emissivity)) + combine_surfaces!( + csf.surface_direct_albedo, + csf, + model_sims, + Val(:surface_direct_albedo), + ) + combine_surfaces!( + csf.surface_diffuse_albedo, + csf, + model_sims, + Val(:surface_diffuse_albedo), + ) + combine_surfaces_temperature!(csf.T_sfc, csf, model_sims) return nothing end @@ -295,27 +305,26 @@ function step_model_sims!(cs::Interfacer.CoupledSimulation) end """ - combine_surfaces!(csf, sims, ::Val{field_name}) where {field_name} + combine_surfaces!(combined_field, csf, sims, field_name) Sums the surface fields specified by `field_name`, weighted by the respective area fractions of all surface simulations. The result is saved in the coupler field specified by `field_name`. -For surface temperature, upward longwave radiation is computed from the temperatures -of each surface, weighted by their area fractions, and then the combined temperature -is computed from the combined upward longwave radiation. +Note that even though `combined_field` is contained in `csf`, it is passed as a separate +argument to avoid runtime dispatch on the field name when accessing the `csf` NamedTuple. # Arguments +- `combined_field`: [Field] coupler field save the combined surface fields to. - `csf`: [NamedTuple] containing coupler fields. Note: For the surface temperature, all coupler fields are passed in a NamedTuple. - `sims`: [NamedTuple] containing simulations. -- `field_name_val`: [Val] containing the name Symbol of the field to be extracted by the `Interfacer.get_field` functions. +- `field_name`: [Val] containing the name Symbol of the field to be extracted by the `Interfacer.get_field` functions. # Example - `combine_surfaces!(temp_field, cs.model_sims, Val(:emissivity))` """ -function combine_surfaces!(csf, sims, ::Val{field_name}) where {field_name} - # Extract the coupler field we are updating - combined_field = getproperty(csf, field_name) +function combine_surfaces!(combined_field, csf, sims, field_name) + # Set the combined field to zero before accumulating across all surface models combined_field .= 0 for sim in sims @@ -325,7 +334,7 @@ function combine_surfaces!(csf, sims, ::Val{field_name}) where {field_name} area_fraction = csf.scalar_temp1 # Remap the surface field onto a coupler temporary field to avoid allocation - Interfacer.get_field!(csf.scalar_temp2, sim, Val(field_name)) + Interfacer.get_field!(csf.scalar_temp2, sim, field_name) surface_field = csf.scalar_temp2 # Zero out the contribution from this surface if the area fraction is zero. @@ -337,9 +346,26 @@ function combine_surfaces!(csf, sims, ::Val{field_name}) where {field_name} end return nothing end -function combine_surfaces!(csf, sims, ::Val{:surface_temperature}) + +""" + combine_surfaces_temperature!(combined_field, csf, sims) + +Computes the combined surface temperature from the combined upward longwave radiation. + +Upward longwave radiation is computed from the temperatures of each surface, weighted by +their area fractions, and then the combined temperature is computed from the combined +upward longwave radiation. + +We have a separate function for surface temperature to avoid runtime dispatch on the field name. + +# Arguments +- `combined_field`: [Field] coupler field save the combined surface temperature to. +- `csf`: [NamedTuple] containing coupler fields. +- `sims`: [NamedTuple] containing simulations. +""" +function combine_surfaces_temperature!(combined_field, csf, sims) # extract the coupler fields we need to get the surface temperature - T_sfc = csf.T_sfc + T_sfc = combined_field emissivity_sfc = csf.emissivity FT = eltype(T_sfc) diff --git a/test/field_exchanger_tests.jl b/test/field_exchanger_tests.jl index 9ca5995de7..b3125b0e32 100644 --- a/test/field_exchanger_tests.jl +++ b/test/field_exchanger_tests.jl @@ -277,7 +277,7 @@ for FT in (Float32, Float64) scalar_temp2 = CC.Fields.zeros(test_space), ) - FieldExchanger.combine_surfaces!(csf, sims, var_name) + FieldExchanger.combine_surfaces!(csf.random, csf, sims, var_name) @test combined_field == fill(FT(sum(fractions .* fields)), test_space) end