remove allocations in combine_surfaces

src/FieldExchanger.jl

@@ -145,35 +145,31 @@ 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
 - `csf`: [NamedTuple] containing coupler fields.
 - `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(:surface_temperature))
+    combine_surfaces!(csf.emissivity, csf, model_sims, Val(:emissivity))
     combine_surfaces!(
         csf.surface_direct_albedo,
+        csf,
         model_sims,
         Val(:surface_direct_albedo),
-        csf.scalar_temp1,
     )
     combine_surfaces!(
         csf.surface_diffuse_albedo,
+        csf,
         model_sims,
         Val(:surface_diffuse_albedo),
-        csf.scalar_temp1,
     )
+    combine_surfaces_temperature!(csf.T_sfc, csf, model_sims)
     return nothing
 end
 
@@ -309,55 +305,71 @@ function step_model_sims!(cs::Interfacer.CoupledSimulation)
 end
 
 """
-    combine_surfaces!(combined_field::CC.Fields.Field, sims, field_name::Val, temp1)
+    combine_surfaces!(combined_field, csf, sims, 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`, 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`: [CC.Fields.Field] output object containing weighted values.
+- `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] 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] 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!(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
         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)
+            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})
+
+"""
+    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
 
-    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 +378,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 +390,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

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.random, csf, sims, var_name)
         @test combined_field == fill(FT(sum(fractions .* fields)), test_space)
     end