Add specific component dt's to cli_options and parse (#1059)

* Add specific component dt's to cli_options and parse

If both dt and component dts are passed in, the component dts are chosen.
When both are passed, or if the coupler dt is not divisible by a component dt,
a relavant warning or error message is passed.

* Add experiment for component dts to pipeline

* Improve comments and change config_dict logic for dts

Unused dt's are now removed from the config.
IF dt_atmos is used, it is temporarily duplicated
as dt during get_atmos_config_dict

* Fix how get_atmos_config_dict accesses configs

get_atmos_config_dict accessed config_dict from the global
scope instead of coupler_dict, which is passed as an argument

* Add dt through merge rather than modifying coupler_dict

.buildkite/pipeline.yml

@@ -327,6 +327,13 @@ steps:
           slurm_ntasks: 1
           slurm_mem: 20GB
 
+      - label: "AMIP - Component dts test"
+        command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/target_amip_component_dts.yml --job_id target_amip_component_dts"
+        artifact_paths: "experiments/ClimaEarth/output/amip/target_amip_component_dts/artifacts/*"
+        agents:
+          slurm_ntasks: 1
+          slurm_mem: 20GB
+
       - label: "MPI AMIP"
         command: "srun julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/coarse_mpi_n4.yml --job_id coarse_mpi_n4"
         artifact_paths: "experiments/ClimaEarth/output/amip/coarse_mpi_n4/artifacts/*"

config/ci_configs/target_amip_component_dts.yml

@@ -0,0 +1,22 @@
+apply_limiter: false
+dt_atmos: "150secs"
+dt_land: "50secs"
+dt_ocean: "30secs"
+dt_seaice: "75secs"
+dt_cpl: 150
+dt_rad: "1hours"
+dt_save_to_sol: "1days"
+dz_bottom: 30
+dz_top: 3000
+energy_check: false
+h_elem: 16
+mode_name: "amip"
+moist: "equil"
+mono_surface: false
+precip_model: "0M"
+rad: "gray"
+rayleigh_sponge: true
+t_end: "600secs"
+vert_diff: "true"
+z_elem: 50
+z_stretch: false

experiments/ClimaEarth/cli_options.jl

@@ -7,6 +7,18 @@ function argparse_settings()
         help = " Coupling time step in seconds"
         arg_type = Int
         default = 400
+        "--dt_atmos"
+        help = " Atmos simulation time step"
+        arg_type = String
+        "--dt_land"
+        help = " Land simulation time step"
+        arg_type = String
+        "--dt_ocean"
+        help = " Ocean simulation time step"
+        arg_type = String
+        "--dt_seaice"
+        help = " Sea Ice simulation time step"
+        arg_type = String
         "--anim"
         help = "Boolean flag indicating whether to make animations"
         arg_type = Bool

experiments/ClimaEarth/components/atmosphere/climaatmos.jl

@@ -350,8 +350,11 @@ function get_atmos_config_dict(coupler_dict::Dict, job_id::String)
 
     # merge configs
     # (if there are common keys, the last dictionary in the `merge` arguments takes precedence)
-    atmos_config = merge(atmos_config, Dict("output_dir" => atmos_output_dir))
-    coupler_config = merge(atmos_config, config_dict)
+    # The Atmos `get_simulation` function expects the atmos config to contains its timestep size
+    # in the `dt` field. If there is a `dt_atmos` field in coupler_dict, we add it to the atmos config as `dt`
+    dt_atmos = haskey(coupler_dict, "dt_atmos") ? coupler_dict["dt_atmos"] : coupler_dict["dt"]
+    atmos_config = merge(atmos_config, Dict("output_dir" => atmos_output_dir, "dt" => dt_atmos))
+    coupler_config = merge(atmos_config, coupler_dict)
 
     # set restart file to the initial file saved in this location if it is not nothing
     # TODO this is hardcoded and should be fixed once we have a better restart system

experiments/ClimaEarth/run_amip.jl

@@ -155,9 +155,6 @@ if mode_name == "amip" && use_coupler_diagnostics
     )
 end
 
-## get component model dictionaries (if applicable)
-atmos_config_dict, config_dict = get_atmos_config_dict(config_dict, job_id)
-atmos_config_object = CA.AtmosConfig(atmos_config_dict)
 
 ## read in some parsed command line arguments, required by this script
 energy_check = config_dict["energy_check"]
@@ -166,8 +163,35 @@ land_sim_name = "bucket"
 t_end = Float64(time_to_seconds(config_dict["t_end"]))
 t_start = 0.0
 tspan = (t_start, t_end)
-Δt_component = Float64(time_to_seconds(config_dict["dt"]))
 Δt_cpl = Float64(config_dict["dt_cpl"])
+component_dt_names = [:dt_atmos, :dt_land, :dt_ocean, :dt_seaice]
+# check if all component dt's are specified
+if all(key -> !isnothing(config_dict[String(key)]), component_dt_names)
+    # when all component dt's are specified, ignore the dt field
+    if haskey(config_dict, "dt")
+        @warn "Removing dt in favor of individual component dt's"
+        delete!(config_dict, "dt")
+    end
+    for key in component_dt_names
+        component_dt = Float64(time_to_seconds(config_dict[String(key)]))
+        @assert Δt_cpl % component_dt == 0.0 "Coupler dt must be divisible by all component dt's\n dt_cpl = $Δt_cpl\n $key = $component_dt"
+        eval(:($key = $component_dt))
+    end
+else
+    # when not all component dt's are specified, use the dt field
+    @assert haskey(config_dict, "dt") "dt or (dt_atmos, dt_land, dt_ocean, and dt_seaice) must be specified"
+    for key in component_dt_names
+        if !isnothing(config_dict[String(key)])
+            @warn "Removing $key from config in favor of dt because not all component dt's are specified"
+        end
+        delete!(config_dict, String(key))
+        eval(:($key = Float64(time_to_seconds(config_dict["dt"]))))
+    end
+end
+## get component model dictionaries (if applicable)
+atmos_config_dict, config_dict = get_atmos_config_dict(config_dict, job_id)
+atmos_config_object = CA.AtmosConfig(atmos_config_dict)
+
 saveat = Float64(time_to_seconds(config_dict["dt_save_to_sol"]))
 date0 = date = Dates.DateTime(config_dict["start_date"], Dates.dateformat"yyyymmdd")
 mono_surface = config_dict["mono_surface"]
@@ -307,7 +331,7 @@ if mode_name == "amip"
         config_dict["land_albedo_type"],
         config_dict["land_temperature_anomaly"],
         dir_paths;
-        dt = Δt_component,
+        dt = dt_land,
         space = boundary_space,
         saveat = saveat,
         area_fraction = land_area_fraction,
@@ -361,7 +385,7 @@ if mode_name == "amip"
     ice_sim = ice_init(
         FT;
         tspan = tspan,
-        dt = Δt_component,
+        dt = dt_seaice,
         space = boundary_space,
         saveat = saveat,
         area_fraction = ice_fraction,
@@ -405,7 +429,7 @@ elseif mode_name in ("slabplanet", "slabplanet_aqua", "slabplanet_terra")
         config_dict["land_albedo_type"],
         config_dict["land_temperature_anomaly"],
         dir_paths;
-        dt = Δt_component,
+        dt = dt_land,
         space = boundary_space,
         saveat = saveat,
         area_fraction = land_area_fraction,
@@ -420,7 +444,7 @@ elseif mode_name in ("slabplanet", "slabplanet_aqua", "slabplanet_terra")
     ocean_sim = ocean_init(
         FT;
         tspan = tspan,
-        dt = Δt_component,
+        dt = dt_ocean,
         space = boundary_space,
         saveat = saveat,
         area_fraction = (FT(1) .- land_area_fraction), ## NB: this ocean fraction includes areas covered by sea ice (unlike the one contained in the cs)
@@ -455,7 +479,7 @@ elseif mode_name == "slabplanet_eisenman"
         config_dict["land_albedo_type"],
         config_dict["land_temperature_anomaly"],
         dir_paths;
-        dt = Δt_component,
+        dt = dt_land,
         space = boundary_space,
         saveat = saveat,
         area_fraction = land_area_fraction,
@@ -470,7 +494,7 @@ elseif mode_name == "slabplanet_eisenman"
     ocean_sim = ocean_init(
         FT;
         tspan = tspan,
-        dt = Δt_component,
+        dt = dt_ocean,
         space = boundary_space,
         saveat = saveat,
         area_fraction = zeros(boundary_space), # zero, since ML is calculated below
@@ -483,7 +507,7 @@ elseif mode_name == "slabplanet_eisenman"
         tspan,
         space = boundary_space,
         area_fraction = (FT(1) .- land_area_fraction),
-        dt = Δt_component,
+        dt = dt_seaice,
         saveat = saveat,
         thermo_params = thermo_params,
     )