From 02804bf44be58655364579a69b1e400bd8d40774 Mon Sep 17 00:00:00 2001 From: kmdeck Date: Tue, 4 Nov 2025 13:10:15 -0800 Subject: [PATCH 1/3] reduce allocations in integrated land coupling step --- .../components/land/climaland_integrated.jl | 37 +++++++++++-------- 1 file changed, 21 insertions(+), 16 deletions(-) diff --git a/experiments/ClimaEarth/components/land/climaland_integrated.jl b/experiments/ClimaEarth/components/land/climaland_integrated.jl index df61d45565..d10c54af2a 100644 --- a/experiments/ClimaEarth/components/land/climaland_integrated.jl +++ b/experiments/ClimaEarth/components/land/climaland_integrated.jl @@ -581,15 +581,17 @@ function FluxCalculator.compute_surface_fluxes!( # Combine turbulent energy fluxes from each component of the land model # Use temporary variables to avoid allocating + @. p.scratch1 = canopy_dest.lhf + soil_dest.lhf * (1 .- p.snow.snow_cover_fraction) + + p.snow.snow_cover_fraction * snow_dest.lhf Interfacer.remap!( csf.scalar_temp1, - canopy_dest.lhf .+ soil_dest.lhf .* (1 .- p.snow.snow_cover_fraction) .+ - p.snow.snow_cover_fraction .* snow_dest.lhf, + p.scratch1, ) + @. p.scratch1 = canopy_dest.shf + soil_dest.shf * (1 - p.snow.snow_cover_fraction) + + p.snow.snow_cover_fraction * snow_dest.shf Interfacer.remap!( csf.scalar_temp2, - canopy_dest.shf .+ soil_dest.shf .* (1 .- p.snow.snow_cover_fraction) .+ - p.snow.snow_cover_fraction .* snow_dest.shf, + p.scratch1 ) # Zero out the fluxes where the area fraction is zero @@ -605,14 +607,14 @@ function FluxCalculator.compute_surface_fluxes!( # Combine turbulent moisture fluxes from each component of the land model # Note that we multiply by ρ_liq to convert from m s-1 to kg m-2 s-1 ρ_liq = (LP.ρ_cloud_liq(sim.model.soil.parameters.earth_param_set)) + @. p.scratch1 = (canopy_dest.transpiration + + (soil_dest.vapor_flux_liq + soil_dest.vapor_flux_ice) * + (1 - p.snow.snow_cover_fraction) + + p.snow.snow_cover_fraction * snow_dest.vapor_flux + ) * ρ_liq Interfacer.remap!( csf.scalar_temp1, - ( - canopy_dest.transpiration .+ - (soil_dest.vapor_flux_liq .+ soil_dest.vapor_flux_ice) .* - (1 .- p.snow.snow_cover_fraction) .+ - p.snow.snow_cover_fraction .* snow_dest.vapor_flux - ) .* ρ_liq, + p.scratch1 ) @. csf.scalar_temp1 = ifelse(area_fraction == 0, zero(csf.scalar_temp1), csf.scalar_temp1) @@ -621,19 +623,21 @@ function FluxCalculator.compute_surface_fluxes!( # Combine turbulent momentum fluxes from each component of the land model # Note that we exclude the canopy component here for now, since we can have nonzero momentum fluxes # where there is zero LAI. This should be fixed in ClimaLand. + @. p.scratch1 = soil_dest.ρτxz * (1 - p.snow.snow_cover_fraction) + + p.snow.snow_cover_fraction * snow_dest.ρτxz Interfacer.remap!( csf.scalar_temp1, - soil_dest.ρτxz .* (1 .- p.snow.snow_cover_fraction) .+ - p.snow.snow_cover_fraction .* snow_dest.ρτxz, + p.scratch1, ) @. csf.scalar_temp1 = ifelse(area_fraction == 0, zero(csf.scalar_temp1), csf.scalar_temp1) @. csf.F_turb_ρτxz += csf.scalar_temp1 * area_fraction + @. p.scratch1 = soil_dest.ρτyz * (1 - p.snow.snow_cover_fraction) + + p.snow.snow_cover_fraction * snow_dest.ρτyz Interfacer.remap!( csf.scalar_temp1, - soil_dest.ρτyz .* (1 .- p.snow.snow_cover_fraction) .+ - p.snow.snow_cover_fraction .* snow_dest.ρτyz, + p.scratch1 ) @. csf.scalar_temp1 = ifelse(area_fraction == 0, zero(csf.scalar_temp1), csf.scalar_temp1) @@ -642,10 +646,11 @@ function FluxCalculator.compute_surface_fluxes!( # Combine the buoyancy flux from each component of the land model # Note that we exclude the canopy component here for now, since ClimaLand doesn't # include its extra resistance term in the buoyancy flux calculation. + @. p.scratch1 = soil_dest.buoy_flux * (1 - p.snow.snow_cover_fraction) + + p.snow.snow_cover_fraction * snow_dest.buoy_flux Interfacer.remap!( csf.scalar_temp1, - soil_dest.buoy_flux .* (1 .- p.snow.snow_cover_fraction) .+ - p.snow.snow_cover_fraction .* snow_dest.buoy_flux, + p.scratch1 ) @. csf.scalar_temp1 = ifelse(area_fraction == 0, zero(csf.scalar_temp1), csf.scalar_temp1) From 94ac42107afbb9ce499d1a8d163d71abb06f7a06 Mon Sep 17 00:00:00 2001 From: kmdeck Date: Tue, 4 Nov 2025 13:46:33 -0800 Subject: [PATCH 2/3] update yml --- config/amip_configs/amip.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/config/amip_configs/amip.yml b/config/amip_configs/amip.yml index 2f104c7e1a..6bb49b2ffc 100644 --- a/config/amip_configs/amip.yml +++ b/config/amip_configs/amip.yml @@ -1,12 +1,12 @@ FLOAT_TYPE: "Float32" albedo_model: "CouplerAlbedo" atmos_config_file: "config/atmos_configs/climaatmos_diagedmf.yml" -bucket_albedo_type: "map_temporal" checkpoint_dt: "366days" coupler_toml: ["toml/amip.toml"] dt: "120secs" dt_cpl: "120secs" energy_check: false +land_model: "integrated" mode_name: "amip" netcdf_output_at_levels: true output_default_diagnostics: true From 35212ba1a346bc1a9e1259fb532cff2ad64bd7a8 Mon Sep 17 00:00:00 2001 From: kmdeck Date: Tue, 4 Nov 2025 15:06:13 -0800 Subject: [PATCH 3/3] remove remap --- .../components/land/climaland_integrated.jl | 67 +++++++++++-------- 1 file changed, 38 insertions(+), 29 deletions(-) diff --git a/experiments/ClimaEarth/components/land/climaland_integrated.jl b/experiments/ClimaEarth/components/land/climaland_integrated.jl index d10c54af2a..5db44880f9 100644 --- a/experiments/ClimaEarth/components/land/climaland_integrated.jl +++ b/experiments/ClimaEarth/components/land/climaland_integrated.jl @@ -554,11 +554,11 @@ function FluxCalculator.compute_surface_fluxes!( @. coupled_atmos.u = StaticArrays.SVector(p.scratch1, p.scratch2) # Use scratch space for remapped atmospheric fields to avoid allocations - Interfacer.remap!(p.scratch1, csf.ρ_atmos) - Interfacer.remap!(p.scratch2, csf.T_atmos) - Interfacer.remap!(p.scratch3, csf.q_atmos) + #Interfacer.remap!(p.scratch1, csf.ρ_atmos) + #Interfacer.remap!(p.scratch2, csf.T_atmos) + #Interfacer.remap!(p.scratch3, csf.q_atmos) @. coupled_atmos.thermal_state = - TD.PhaseEquil_ρTq(thermo_params, p.scratch1, p.scratch2, p.scratch3) + TD.PhaseEquil_ρTq(thermo_params, csf.ρ_atmos, csf.T_atmos, csf.q_atmos) # set the same atmosphere state for all sub-components @assert sim.model.soil.boundary_conditions.top.atmos === @@ -583,17 +583,18 @@ function FluxCalculator.compute_surface_fluxes!( # Use temporary variables to avoid allocating @. p.scratch1 = canopy_dest.lhf + soil_dest.lhf * (1 .- p.snow.snow_cover_fraction) + p.snow.snow_cover_fraction * snow_dest.lhf - Interfacer.remap!( - csf.scalar_temp1, - p.scratch1, - ) + #Interfacer.remap!( + # csf.scalar_temp1, + # p.scratch1, + #) + csf.scalar_temp1 .= p.scratch1 @. p.scratch1 = canopy_dest.shf + soil_dest.shf * (1 - p.snow.snow_cover_fraction) + p.snow.snow_cover_fraction * snow_dest.shf - Interfacer.remap!( - csf.scalar_temp2, - p.scratch1 - ) - + #Interfacer.remap!( + # csf.scalar_temp2, + # p.scratch1 + #) + csf.scalar_temp2 .= p.scratch1 # Zero out the fluxes where the area fraction is zero @. csf.scalar_temp1 = ifelse(area_fraction == 0, zero(csf.scalar_temp1), csf.scalar_temp1) @@ -612,10 +613,12 @@ function FluxCalculator.compute_surface_fluxes!( (1 - p.snow.snow_cover_fraction) + p.snow.snow_cover_fraction * snow_dest.vapor_flux ) * ρ_liq - Interfacer.remap!( - csf.scalar_temp1, - p.scratch1 - ) + #Interfacer.remap!( + # csf.scalar_temp1, + # p.scratch1 + #) + csf.scalar_temp1 .= p.scratch1 + @. csf.scalar_temp1 = ifelse(area_fraction == 0, zero(csf.scalar_temp1), csf.scalar_temp1) @. csf.F_turb_moisture += csf.scalar_temp1 * area_fraction @@ -625,20 +628,24 @@ function FluxCalculator.compute_surface_fluxes!( # where there is zero LAI. This should be fixed in ClimaLand. @. p.scratch1 = soil_dest.ρτxz * (1 - p.snow.snow_cover_fraction) + p.snow.snow_cover_fraction * snow_dest.ρτxz - Interfacer.remap!( - csf.scalar_temp1, - p.scratch1, - ) + #Interfacer.remap!( + # csf.scalar_temp1, + # p.scratch1, + #) + csf.scalar_temp1 .= p.scratch1 + @. csf.scalar_temp1 = ifelse(area_fraction == 0, zero(csf.scalar_temp1), csf.scalar_temp1) @. csf.F_turb_ρτxz += csf.scalar_temp1 * area_fraction @. p.scratch1 = soil_dest.ρτyz * (1 - p.snow.snow_cover_fraction) + p.snow.snow_cover_fraction * snow_dest.ρτyz - Interfacer.remap!( - csf.scalar_temp1, - p.scratch1 - ) + #Interfacer.remap!( + # csf.scalar_temp1, + # p.scratch1 + #) + csf.scalar_temp1 .= p.scratch1 + @. csf.scalar_temp1 = ifelse(area_fraction == 0, zero(csf.scalar_temp1), csf.scalar_temp1) @. csf.F_turb_ρτyz += csf.scalar_temp1 * area_fraction @@ -648,10 +655,12 @@ function FluxCalculator.compute_surface_fluxes!( # include its extra resistance term in the buoyancy flux calculation. @. p.scratch1 = soil_dest.buoy_flux * (1 - p.snow.snow_cover_fraction) + p.snow.snow_cover_fraction * snow_dest.buoy_flux - Interfacer.remap!( - csf.scalar_temp1, - p.scratch1 - ) + #Interfacer.remap!( + # csf.scalar_temp1, + # p.scratch1 + #) + csf.scalar_temp1 .= p.scratch1 + @. csf.scalar_temp1 = ifelse(area_fraction == 0, zero(csf.scalar_temp1), csf.scalar_temp1) @. csf.buoyancy_flux += csf.scalar_temp1 * area_fraction