100km atmos scaling run

switch off coupler diagnostics rm cpu bind no atmos diags no atmos diags 300d clean clean 250 d fix changed atmos albedo flag wip runs try disable nonideal insol albedo initial = 0.38 revert reshape add debug plots add debug plots try with gpu debug plots fix debug plots try 70d 1 proc 120s 30GB 1 ntask 30GB disable diags and debug carry over from ln/100km-test (1.082 SYPD w/o diags) format rm debug for gpu debug for CPUSingleThreaded try sponge setup 50d 1 proc clean singleton clean clean 120d test git push no sponge edmf params sponge + edmf params + 200d 205d running revs deps pip rm old file
CliMA · May 4, 2024 · 9a6993c · 9a6993c
1 parent 6a5c89c
commit 9a6993c
Show file tree

Hide file tree

Showing 15 changed files with 88 additions and 43 deletions.
diff --git a/.buildkite/longruns/pipeline.yml b/.buildkite/longruns/pipeline.yml
@@ -294,7 +294,7 @@ steps:
           slurm_gpus_per_task: 1
           slurm_cpus_per_task: 4
           slurm_ntasks: 4
-          slurm_mem: 20GB
+          slurm_mem: 30GB
           modules: common
         soft_fail: true
 
@@ -339,12 +339,11 @@ steps:
 
       - label: ":envelope: Slack report: target AMIP"
         command:
-          - slack-upload -c "#coupler-report" -f experiments/AMIP/output/amip/gpu_amip_target_topo_diagedmf_artifacts/amip_ncep.png -m png -n amip_fine -x "90d Target AMIP v NCEP Last Month Mean"
+          - slack-upload -c "#coupler-report" -f experiments/AMIP/output/amip/gpu_amip_target_topo_diagedmf_artifacts/amip_ncep.png -m png -n amip_fine -x "300d Target AMIP v NCEP Last Month Mean"
           - |
             find experiments/AMIP/output/amip/gpu_amip_target_topo_diagedmf_artifacts/ -type f -name 'bias*.png' -print0 | while IFS= read -r -d '' file; do
               slack-upload -c "#coupler-report" -f "$$file" -m png -n "$$(basename "$$file" .png)" -x "$$(basename "$$file" .png)"
             done
         agents:
           queue: clima
           modules: common
-
diff --git a/config/longrun_configs/amip_target_topo.yml b/config/longrun_configs/amip_target_topo.yml
@@ -1,4 +1,5 @@
 FLOAT_TYPE: "Float32"
+albedo_model: "CouplerAlbedo"
 anim: true
 apply_limiter: false
 atmos_config_file: "config/longrun_configs/longrun_aquaplanet_rhoe_equil_55km_nz63_clearsky_0M_earth.yml"

diff --git a/config/longrun_configs/amip_target_topo_diagedmf.yml b/config/longrun_configs/amip_target_topo_diagedmf.yml
@@ -1,26 +1,30 @@
 FLOAT_TYPE: "Float32"
-anim: true
+albedo_model: "CouplerAlbedo"
+anim: false
 apply_limiter: false
-atmos_config_file: "config/longrun_configs/longrun_aquaplanet_rhoe_equil_55km_nz63_allsky_diagedmf_0M.yml"
-coupler_toml_file: "toml/amip_target_topo_diagedmf.toml"
-dt: "100secs"
+atmos_config_file: "config/gpu_configs/gpu_aquaplanet_diagedmf.yml"
+dt: "120secs"
 dt_cloud_fraction: "1hours"
-dt_cpl: 100
+dt_cpl: 120
 dt_rad: "1hours"
 dt_save_state_to_disk: "20days"
 dt_save_to_sol: "10days"
 energy_check: false
-hourly_checkpoint: true
+hourly_checkpoint: false
 idealized_insolation: false
 job_id: "amip_target_topo_diagedmf"
 land_albedo_type: "map_temporal"
 mode_name: "amip"
 mono_surface: false
 netcdf_output_at_levels: true
+output_default_diagnostics: true
+rayleigh_sponge: true
 run_name: "amip_target_topo_diagedmf"
 start_date: "19790301"
 surface_setup: "PrescribedSurface"
-t_end: "40days"
+t_end: "30days"
 topo_smoothing: true
 topography: "Earth"
+coupler_toml_file: "toml/amip_target_topo_diagedmf.toml"
 turb_flux_partition: "CombinedStateFluxes"
+use_reference_state: false
diff --git a/config/longrun_configs/gpu_amip_target_topo.yml b/config/longrun_configs/gpu_amip_target_topo.yml
@@ -1,5 +1,6 @@
 FLOAT_TYPE: "Float32"
 anim: false
+albedo_model: "CouplerAlbedo"
 apply_limiter: false
 atmos_config_file: "config/longrun_configs/longrun_aquaplanet_rhoe_equil_55km_nz63_clearsky_0M_earth.yml"
 dt: "100secs"

diff --git a/config/longrun_configs/gpu_amip_target_topo_diagedmf.yml b/config/longrun_configs/gpu_amip_target_topo_diagedmf.yml
@@ -1,26 +1,30 @@
 FLOAT_TYPE: "Float32"
+albedo_model: "CouplerAlbedo"
 anim: false
 apply_limiter: false
-atmos_config_file: "config/longrun_configs/longrun_aquaplanet_rhoe_equil_55km_nz63_allsky_diagedmf_0M.yml"
-dt: "100secs"
+atmos_config_file: "config/gpu_configs/gpu_aquaplanet_diagedmf.yml"
+dt: "120secs"
 dt_cloud_fraction: "1hours"
-dt_cpl: 100
+dt_cpl: 120
 dt_rad: "1hours"
 dt_save_state_to_disk: "20days"
 dt_save_to_sol: "10days"
 energy_check: false
-hourly_checkpoint: true
+hourly_checkpoint: false
 idealized_insolation: false
 job_id: "gpu_amip_target_topo_diagedmf"
 land_albedo_type: "map_temporal"
 mode_name: "amip"
 mono_surface: false
 netcdf_output_at_levels: true
+output_default_diagnostics: true
+rayleigh_sponge: true
 run_name: "gpu_amip_target_topo_diagedmf"
 start_date: "19790301"
 surface_setup: "PrescribedSurface"
-t_end: "90days"
+t_end: "300days"
 topo_smoothing: true
 topography: "Earth"
+coupler_toml_file: "toml/amip_target_topo_diagedmf.toml"
 turb_flux_partition: "CombinedStateFluxes"
 use_reference_state: false
diff --git a/config/model_configs/interactive_debug.yml b/config/model_configs/interactive_debug.yml
@@ -7,16 +7,16 @@ dt_save_to_sol: "10days"
 energy_check: true
 evolving_ocean: true
 h_elem: 4
-job_id: "interactive_debug_run"
+job_id: "interactive_debug_run154"
 land_albedo_type: "function"
 mode_name: "slabplanet"
 moist: "equil"
 mono_surface: true
 precip_model: "0M"
 rad: "gray"
-run_name: "interactive_debug_run"
+run_name: "interactive_debug_run154"
 start_date: "19790301"
 surface_setup: "PrescribedSurface"
-t_end: "20days"
+t_end: "2days"
 turb_flux_partition: "CombinedStateFluxes"
 vert_diff: "true"
diff --git a/experiments/AMIP/cli_options.jl b/experiments/AMIP/cli_options.jl
@@ -85,6 +85,10 @@ function argparse_settings()
         default = "PrescribedSurface"
         "--atmos_config_file"
         help = "A yaml file used to set the atmospheric model configuration. If nothing is specified, the default configuration is used."
+        "--albedo_model"
+        help = "Type of albedo model. [`ConstantAlbedo` (default), `RegressionFunctionAlbedo`, `CouplerAlbedo`]"
+        arg_type = String
+        default = "CouplerAlbedo"
         # 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
@@ -23,7 +23,6 @@ Interfacer.name(::ClimaAtmosSimulation) = "ClimaAtmosSimulation"
 
 function atmos_init(::Type{FT}, atmos_config_dict::Dict) where {FT}
     # By passing `parsed_args` to `AtmosConfig`, `parsed_args` overwrites the default atmos config
-    atmos_config_dict["surface_albedo"] = "CouplerAlbedo"
     atmos_config = CA.AtmosConfig(atmos_config_dict)
     simulation = CA.get_simulation(atmos_config)
     (; integrator) = simulation
@@ -64,6 +63,23 @@ function atmos_init(::Type{FT}, atmos_config_dict::Dict) where {FT}
     return sim
 end
 
+# Extension of CA.set_surface_albedo! to set the surface albedo to 0.38 at the beginning of the simulation,
+# so the initial callback initialization doesn't lead to NaNs in the radiation model.
+# Subsequently, the surface albedo will be updated by the coupler, via water_albedo_from_atmosphere!.
+function CA.set_surface_albedo!(Y, p, t, ::CA.CouplerAlbedo)
+    if t == 0
+        FT = eltype(Y)
+        # set initial insolation initial conditions
+        !p.radiation.idealized_insolation && CA.set_insolation_variables!(Y, p, t)
+        # set surface albedo to 0.38
+        @warn "Setting surface albedo to 0.38 at the beginning of the simulation"
+        p.radiation.radiation_model.direct_sw_surface_albedo .= FT(0.38)
+        p.radiation.radiation_model.diffuse_sw_surface_albedo .= FT(0.38)
+    else
+        nothing
+    end
+end
+
 """
     Checkpointer.get_model_prog_state(sim::ClimaAtmosSimulation)
 
@@ -208,6 +224,13 @@ Interfacer.step!(sim::ClimaAtmosSimulation, t) = Interfacer.step!(sim.integrator
 Interfacer.reinit!(sim::ClimaAtmosSimulation) = Interfacer.reinit!(sim.integrator)
 
 function FieldExchanger.update_sim!(atmos_sim::ClimaAtmosSimulation, csf, turbulent_fluxes)
+
+    u = atmos_sim.integrator.u
+    p = atmos_sim.integrator.p
+    t = atmos_sim.integrator.t
+
+    !p.radiation.idealized_insolation && CA.set_insolation_variables!(u, p, t)
+
     Interfacer.update_field!(atmos_sim, Val(:surface_direct_albedo), csf.surface_direct_albedo)
     Interfacer.update_field!(atmos_sim, Val(:surface_diffuse_albedo), csf.surface_diffuse_albedo)
     Interfacer.update_field!(atmos_sim, Val(:surface_temperature), csf)
@@ -404,14 +427,11 @@ function FluxCalculator.water_albedo_from_atmosphere!(
     λ = FT(0) # spectral wavelength (not used for now)
 
     # update for current zenith angle
-    CA.set_insolation_variables!(Y, p, t)
-
     bottom_coords = CC.Fields.coordinate_field(CC.Spaces.level(Y.c, 1))
     μ = CA.RRTMGPI.array2field(radiation_model.cos_zenith, axes(bottom_coords))
     FT = eltype(atmos_sim.integrator.u)
     α_model = CA.RegressionFunctionAlbedo{FT}()
 
-
     # set the direct and diffuse surface albedos
     direct_albedo .= CA.surface_albedo_direct(α_model).(λ, μ, LinearAlgebra.norm.(CC.Fields.level(Y.c.uₕ, 1)))
     diffuse_albedo .= CA.surface_albedo_diffuse(α_model).(λ, μ, LinearAlgebra.norm.(CC.Fields.level(Y.c.uₕ, 1)))

diff --git a/experiments/AMIP/coupler_driver.jl b/experiments/AMIP/coupler_driver.jl
@@ -77,6 +77,7 @@ include("components/ocean/eisenman_seaice.jl")
 ## helpers for user-specified IO
 include("user_io/user_diagnostics.jl")
 include("user_io/user_logging.jl")
+include("user_io/debug_plots.jl")
 
 #=
 ### Configuration Dictionaries
@@ -92,7 +93,6 @@ parsed_args = parse_commandline(argparse_settings())
 
 ## modify parsed args for fast testing from REPL #hide
 if isinteractive()
-    include("user_io/debug_plots.jl")
     parsed_args["config_file"] =
         isnothing(parsed_args["config_file"]) ? joinpath(pkg_dir, "config/model_configs/interactive_debug.yml") :
         parsed_args["config_file"]
@@ -805,6 +805,8 @@ if ClimaComms.iamroot(comms_ctx)
 
     ## plotting AMIP results
     if cs.mode.name == "amip"
+        ## plot data that correspond to the model's last save_hdf5 call (i.e., last month)
+
         @info "AMIP plots"
 
         ## ClimaESM
@@ -857,9 +859,9 @@ if ClimaComms.iamroot(comms_ctx)
             COUPLER_OUTPUT_DIR,
             output_dir = COUPLER_ARTIFACTS_DIR,
             month_date = cs.dates.date[1],
-        ) ## plot data that correspond to the model's last save_hdf5 call (i.e., last month)
+        )
 
-        ## combined plots
+        ## combine AMIP and NCEP plots
         plot_combined = Plots.plot(fig_amip, fig_ncep, layout = (2, 1), size = (1400, 1800))
         Plots.png(joinpath(COUPLER_ARTIFACTS_DIR, "amip_ncep.png"))
 
@@ -893,12 +895,14 @@ if ClimaComms.iamroot(comms_ctx)
         make_plots(Val(:general_ci_plots), [joinpath(COUPLER_OUTPUT_DIR, "clima_atmos")], COUPLER_ARTIFACTS_DIR)
     end
 
+    ## plot all model states and coupler fields (useful for debugging) TODO: make MPI & GPU compatible
+    comms_ctx.device == ClimaComms.CPUSingleThreaded() &&
+        comms_ctx isa ClimaComms.SingletonCommsContext &&
+        debug(cs, joinpath(COUPLER_ARTIFACTS_DIR, "end_"))
+
     if isinteractive()
         ## clean up for interactive runs, retain all output otherwise
         rm(COUPLER_OUTPUT_DIR; recursive = true, force = true)
-
-        ## plot all model states and coupler fields (useful for debugging)
-        debug(cs)
     end
 
 end
diff --git a/experiments/AMIP/user_io/debug_plots.jl b/experiments/AMIP/user_io/debug_plots.jl
@@ -34,13 +34,17 @@ function debug(cs_fields::NamedTuple, dir, cs_fields_ref = nothing)
         :F_radiative,
         :F_turb_energy,
         :F_turb_moisture,
+        :F_turb_ρτxz,
+        :F_turb_ρτyz,
         :P_liq,
+        :P_snow,
         :T_S,
         :ρ_sfc,
         :q_sfc,
         :beta,
         :z0b_S,
         :z0m_S,
+        :radiative_energy_flux_toa,
     )
     all_plots = []
     for field_name in field_names

diff --git a/experiments/ClimaCore/Project.toml b/experiments/ClimaCore/Project.toml
@@ -2,6 +2,7 @@
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
 ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c"
 ClimaTimeSteppers = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
+CommonDataModel = "1fbeeb36-5f17-413c-809b-666fb144f157"
 DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
 FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
 IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
@@ -19,6 +20,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 [compat]
 ClimaParams = "0.10"
 ClimaTimeSteppers = "0.7"
+CommonDataModel = "=0.3.5"
 DifferentialEquations = "7.10"
 FileIO = "1.16"
 IntervalSets = "0.6"

diff --git a/src/ConservationChecker.jl b/src/ConservationChecker.jl
@@ -277,7 +277,9 @@ function plot_global_conservation(
 
     # check that the relative error is small (TODO: reduce this to sqrt(eps(FT)))
     if !softfail
-        @assert rse[end] < 2e-3
+        @show typeof(cc)
+        @show rse[end]
+        @assert rse[end] < 3e-3
     end
 end
 

diff --git a/src/Regridder.jl b/src/Regridder.jl
@@ -25,7 +25,8 @@ export write_to_hdf5,
     combine_surfaces_from_sol!,
     binary_mask,
     nans_to_zero,
-    cgll2latlonz
+    cgll2latlonz,
+    truncate_dataset
 
 
 #= Converts NaNs to zeros of the same type. =#
@@ -673,7 +674,7 @@ function truncate_dataset(
 end
 
 """
-    find_idx_bounding_dates(dates, date_start, date_end) 
+    find_idx_bounding_dates(dates, date_start, date_end)
 
 Returns the index range from dates that contains date_start to date_end
 """
@@ -687,10 +688,10 @@ function find_idx_bounding_dates(dates, date_start, date_end)
     elseif date_start > last(dates)
         start_id = length(dates)
         # if the simulation start date falls within the range of the dataset
-        # find the closest date to the start date and truncate there 
+        # find the closest date to the start date and truncate there
     else
         (~, start_id) = findmin(x -> abs(x - date_start), dates)
-        # if the closest date is after the start date, add one more date before 
+        # if the closest date is after the start date, add one more date before
         if dates[start_id] > date_start
             start_id = start_id - 1
         end
@@ -701,11 +702,11 @@ function find_idx_bounding_dates(dates, date_start, date_end)
     if date_end < dates[1]
         end_id = 1
         # if the simulation end date is after the last date in the dataset
-        # leave the end of the dataset as is 
+        # leave the end of the dataset as is
     elseif date_end > last(dates)
         end_id = length(dates)
         # if the simulation end date falls within the range of the dataset
-        # find the closest date to the end date and truncate there 
+        # find the closest date to the end date and truncate there
     else
         (~, end_id) = findmin(x -> abs(x - date_end), dates)
         # if the closest date is before the end date, add one more date after

diff --git a/test/Project.toml b/test/Project.toml
@@ -11,6 +11,7 @@ ClimaCoupler = "4ade58fe-a8da-486c-bd89-46df092ec0c7"
 ClimaLand = "08f4d4ce-cf43-44bb-ad95-9d2d5f413532"
 ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c"
 ClimaTimeSteppers = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
+CommonDataModel = "1fbeeb36-5f17-413c-809b-666fb144f157"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 GeoMakie = "db073c08-6b98-4ee5-b6a4-5efafb3259c6"
@@ -40,6 +41,7 @@ ArtifactWrappers = "0.2"
 ClimaAtmos = "0.22"
 ClimaLand = "0.11"
 ClimaParams = "0.10"
+CommonDataModel = "=0.3.5"
 Dates = "1"
 Downloads = "1"
 IntervalSets = "0.5, 0.6, 0.7"

diff --git a/toml/amip_target_topo_diagedmf.toml b/toml/amip_target_topo_diagedmf.toml
@@ -1,15 +1,12 @@
+[alpha_rayleigh_w]
+value = 3.0
+
 [zd_rayleigh]
 value = 35000.0
 
 [alpha_rayleigh_uh]
 value = 0.0
 
-[alpha_rayleigh_w]
-value = 3.0
-
-[precipitation_timescale]
-value = 600
-
 [entr_inv_tau]
 value = 0.002
 
@@ -23,10 +20,10 @@ value = 0
 value = 0.12
 
 [detr_vertdiv_coeff]
-value = 0
+value = 0.6
 
 [min_area_limiter_scale]
 value = 0
 
 [max_area_limiter_scale]
-value = 0
+value = 0