use atmos config in coupler

config/benchmark_configs/amip_diagedmf.yml

@@ -1,6 +1,7 @@
 FLOAT_TYPE: "Float32"
 anim: false
-atmos_config_file: "config/gpu_configs/gpu_aquaplanet_diagedmf.yml"
+atmos_config_file: "config/benchmark_configs/climaatmos_diagedmf.yml"
+atmos_config_repo: "ClimaCoupler"
 dt_cpl: 120
 dt_save_state_to_disk: "Inf"
 dt_save_to_sol: "Inf"

config/benchmark_configs/gpu_amip_diagedmf.yml

@@ -1,6 +1,7 @@
 FLOAT_TYPE: "Float32"
 anim: false
-atmos_config_file: "config/gpu_configs/gpu_aquaplanet_diagedmf.yml"
+atmos_config_file: "config/benchmark_configs/gpu_climaatmos_diagedmf.yml"
+atmos_config_repo: "ClimaCoupler"
 dt_cpl: 120
 dt_save_state_to_disk: "Inf"
 dt_save_to_sol: "Inf"

experiments/AMIP/cli_options.jl

@@ -93,6 +93,10 @@ function argparse_settings()
         help = "Type of albedo model. [`ConstantAlbedo` (default), `RegressionFunctionAlbedo`, `CouplerAlbedo`]"
         arg_type = String
         default = "CouplerAlbedo"
+        "--atmos_config_repo"
+        help = "The repository containing the ClimaAtmos configuration file to use [`ClimaAtmos` (default), `ClimaCoupler`]"
+        arg_type = String
+        default = "ClimaAtmos"
         # ClimaLand specific
         "--land_albedo_type"
         help = "Access land surface albedo information from data file. [`function`, `map_static`, `map_temporal`]"

experiments/AMIP/components/atmosphere/climaatmos.jl

@@ -299,20 +299,36 @@ FluxCalculator.get_surface_params(sim::ClimaAtmosSimulation) = CAP.surface_fluxe
 
 Returns the specified atmospheric configuration (`atmos_config`) overwitten by arguments
 in the coupler dictionary (`config_dict`). The returned dictionary will then be passed to CA.AtmosConfig().
+The `atmos_config_repo` flag allows us to
+use a configuration specified within the ClimaCoupler repo, which is useful for direct
+coupled/atmos-only comparisons.
 """
 function get_atmos_config_dict(coupler_dict)
     atmos_config_file = coupler_dict["atmos_config_file"]
+    atmos_config_repo = coupler_dict["atmos_config_repo"]
     # override default or specified configs with coupler arguments, and set the correct atmos config_file
-    if isnothing(atmos_config_file)
-        @info "Using Atmos default configuration"
-        atmos_config = merge(CA.default_config_dict(), coupler_dict, Dict("config_file" => atmos_config_file))
-    else
-        @info "Using Atmos configuration from $atmos_config_file"
+    if atmos_config_repo == "ClimaCoupler"
+        @assert !isnothing(atmos_config_file) "Must specify `atmos_config_file` within ClimaCoupler."
+        @info "Using Atmos configuration from ClimaCoupler in $atmos_config_file"
         atmos_config = merge(
-            CA.override_default_config(joinpath(pkgdir(CA), atmos_config_file)),
+            CA.override_default_config(joinpath(pkgdir(ClimaCoupler), atmos_config_file)),
             coupler_dict,
             Dict("config_file" => atmos_config_file),
         )
+    elseif atmos_config_repo == "ClimaAtmos"
+        if isnothing(atmos_config_file)
+            @info "Using Atmos default configuration"
+            atmos_config = merge(CA.default_config_dict(), coupler_dict, Dict("config_file" => atmos_config_file))
+        else
+            @info "Using Atmos configuration from $atmos_config_file"
+            atmos_config = merge(
+                CA.override_default_config(joinpath(pkgdir(CA), atmos_config_file)),
+                coupler_dict,
+                Dict("config_file" => atmos_config_file),
+            )
+        end
+    else
+        error("Invalid `atmos_config_repo`; please use \"ClimaCoupler\" or \"ClimaAtmos\"")
     end
 
     # use coupler toml if atmos is not defined

experiments/AMIP/coupler_driver.jl

@@ -689,8 +689,7 @@ function solve_coupler!(cs)
 
     ClimaComms.iamroot(comms_ctx) && @info("Starting coupling loop")
     ## step in time
-    ## Use ClimaComms.@elapsed to time the simulation on both CPU and GPU
-    walltime = ClimaComms.@elapsed comms_ctx.device for t in ((tspan[begin] + Δt_cpl):Δt_cpl:tspan[end])
+    for t in ((tspan[begin] + Δt_cpl):Δt_cpl:tspan[end])
 
         cs.dates.date[1] = TimeManager.current_date(cs, t)
 
@@ -777,18 +776,30 @@ function solve_coupler!(cs)
         TimeManager.trigger_callback!(cs, cs.callbacks.checkpoint)
 
     end
-    ClimaComms.iamroot(comms_ctx) && @show(walltime)
 
-    return walltime
+    return nothing
 end
 
 ## exit if running performance anaysis #hide
 if haskey(ENV, "CI_PERF_SKIP_COUPLED_RUN") #hide
     throw(:exit_profile_init) #hide
 end #hide
 
-## run the coupled simulation
-walltime = solve_coupler!(cs);
+## run the coupled simulation for one timestep to precompile everything before timing
+cs.tspan[2] = Δt_cpl * 2
+solve_coupler!(cs)
+
+## run the coupled simulation for the full timespan and time it
+cs.tspan[1] = Δt_cpl * 2
+cs.tspan[2] = tspan[2]
+
+## Use ClimaComms.@elapsed to time the simulation on both CPU and GPU
+walltime = ClimaComms.@elapsed comms_ctx.device begin
+    s = CA.@timed_str begin
+        solve_coupler!(cs)
+    end
+end
+ClimaComms.iamroot(comms_ctx) && @show(walltime)
 
 ## Use ClimaAtmos calculation to show the simulated years per day of the simulation (SYPD)
 es = CA.EfficiencyStats(tspan, walltime)