From c96101a6dec43af3071e086c3c484cf799ee406b Mon Sep 17 00:00:00 2001
From: Julia Sloan <jsloan@caltech.edu>
Date: Wed, 24 Apr 2024 17:00:45 -0700
Subject: [PATCH] use atmos config in coupler

---
 config/benchmark_configs/amip_diagedmf.yml    |  3 +-
 .../benchmark_configs/gpu_amip_diagedmf.yml   |  3 +-
 experiments/AMIP/cli_options.jl               |  4 +++
 .../AMIP/components/atmosphere/climaatmos.jl  | 28 +++++++++++++++----
 4 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/config/benchmark_configs/amip_diagedmf.yml b/config/benchmark_configs/amip_diagedmf.yml
index fc890ffe80..2b78ef6ad0 100644
--- a/config/benchmark_configs/amip_diagedmf.yml
+++ b/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"
diff --git a/config/benchmark_configs/gpu_amip_diagedmf.yml b/config/benchmark_configs/gpu_amip_diagedmf.yml
index b798e82967..84f8038863 100644
--- a/config/benchmark_configs/gpu_amip_diagedmf.yml
+++ b/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"
diff --git a/experiments/AMIP/cli_options.jl b/experiments/AMIP/cli_options.jl
index 7406c5f1fd..af58607704 100644
--- a/experiments/AMIP/cli_options.jl
+++ b/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`]"
diff --git a/experiments/AMIP/components/atmosphere/climaatmos.jl b/experiments/AMIP/components/atmosphere/climaatmos.jl
index 23b702ef96..a9ad2acec1 100644
--- a/experiments/AMIP/components/atmosphere/climaatmos.jl
+++ b/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