Merge #483

483: Generalize callbacks r=LenkaNovak a=LenkaNovak



Co-authored-by: Lenka  Novakova <lenka@dhcp-72-159.caltech.edu>

.buildkite/longruns/pipeline.yml

@@ -14,6 +14,9 @@ env:
   CONFIG_PATH: "config/longrun_configs"
   PERF_CONFIG_PATH: "config/perf_configs"
 
+  LD_LIBRARY_PATH: "/central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/cpu/artifacts/e9feeabfa989549ac1dac61e4df5f2849149ad6d/lib/:${LD_LIBRARY_PATH}"
+
+
 timeout_in_minutes: 1440
 
 steps:

.buildkite/pipeline.yml

@@ -16,6 +16,8 @@ env:
   PERF_CONFIG_PATH: "config/perf_configs"
   MPI_CONFIG_PATH: "config/mpi_configs"
 
+  LD_LIBRARY_PATH: "/central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/cpu/artifacts/e9feeabfa989549ac1dac61e4df5f2849149ad6d/lib/:${LD_LIBRARY_PATH}"
+
 timeout_in_minutes: 1440
 
 steps:
@@ -240,21 +242,21 @@ steps:
       #   command: "julia --color=yes --project=experiments/AMIP/modular/ experiments/AMIP/modular/coupler_driver_modular.jl --coupled true --surface_setup PrescribedSurface --moist equil --vert_diff true --rad allskywithclear --rayleigh_sponge false --energy_check true --mode_name slabplanet --t_end 10days --dt_save_to_sol 3600secs --dt_cpl 21600 --dt 200secs --dt_rad 6hours --idealized_insolation true --mono_surface true --h_elem 4 --precip_model 0M --run_name target_params_in_slab1"
       #   artifact_paths: "experiments/AMIP/modular/output/slabplanet/target_params_in_slab1_artifacts/total_energy*.png"
 
-      - label: "AMIP - modular Float64 + monthly checkpoint"
+      - label: "AMIP - modular Float64 + hourly checkpoint"
         key: "modular_amip"
-        command: "julia --color=yes --project=experiments/AMIP/modular/ experiments/AMIP/modular/coupler_driver_modular.jl --config_file $CONFIG_PATH/coarse_single_modular_ft64_monthly_checkpoints.yml"
-        artifact_paths: "experiments/AMIP/modular/output/amip/coarse_single_modular_ft64_monthly_checkpoints_artifacts/*"
+        command: "julia --color=yes --project=experiments/AMIP/modular/ experiments/AMIP/modular/coupler_driver_modular.jl --config_file $CONFIG_PATH/coarse_single_modular_ft64_hourly_checkpoints.yml"
+        artifact_paths: "experiments/AMIP/modular/output/amip/coarse_single_modular_ft64_hourly_checkpoints_artifacts/*"
         env:
           FLAME_PLOT: ""
           BUILD_HISTORY_HANDLE: ""
         agents:
           slurm_ntasks: 1
           slurm_mem: 20GB
 
-      - label: "AMIP - modular Float64 + monthly checkpoint + co2"
-        key: "coarse_single_modular_ft64_monthly_checkpoints_co2"
-        command: "julia --color=yes --project=experiments/AMIP/modular/ experiments/AMIP/modular/coupler_driver_modular.jl --config_file $CONFIG_PATH/coarse_single_modular_ft64_monthly_checkpoints_co2.yml"
-        artifact_paths: "experiments/AMIP/modular/output/amip/coarse_single_modular_ft64_monthly_checkpoints_co2_artifacts/*"
+      - label: "AMIP - modular Float64 + hourly checkpoint + co2"
+        key: "coarse_single_modular_ft64_hourly_checkpoints_co2"
+        command: "julia --color=yes --project=experiments/AMIP/modular/ experiments/AMIP/modular/coupler_driver_modular.jl --config_file $CONFIG_PATH/coarse_single_modular_ft64_hourly_checkpoints_co2.yml"
+        artifact_paths: "experiments/AMIP/modular/output/amip/coarse_single_modular_ft64_hourly_checkpoints_co2_artifacts/*"
         env:
           FLAME_PLOT: ""
           BUILD_HISTORY_HANDLE: ""

Project.toml

@@ -38,7 +38,7 @@ Insolation = "0.6"
 JLD2 = "0.4"
 NCDatasets = "0.11, 0.12"
 OrdinaryDiffEq = "5, 6"
-Plots = "1"
+Plots = "1.39.0"
 SciMLBase = "1"
 StaticArrays = "1"
 SurfaceFluxes = "0.7"

config/longrun_configs/dyamond_target.yml

@@ -11,4 +11,4 @@ dt_save_to_disk: "0.5days"
 turb_flux_partition: "CombinedStateFluxes"
 atmos_config_file: "config/longrun_configs/longrun_aquaplanet_dyamond.yml"
 atmos_toml_file: "toml/longrun_aquaplanet_dyamond.toml"
-monthly_checkpoint: false
+hourly_checkpoint: false

config/model_configs/coarse_single_modular_ft64_monthly_checkpoints.yml → config/model_configs/coarse_single_modular_ft64_hourly_checkpoints.yml

@@ -1,5 +1,5 @@
-run_name: "coarse_single_modular_ft64_monthly_checkpoints"
-monthly_checkpoint: true
+run_name: "coarse_single_modular_ft64_hourly_checkpoints"
+hourly_checkpoint: true
 moist: "equil"
 vert_diff: "true"
 rad: "gray"
@@ -15,4 +15,4 @@ h_elem: 6
 dt_save_restart: "10days"
 precip_model: "0M"
 apply_limiter: false
-job_id: "coarse_single_modular_ft64_monthly_checkpoints"
+job_id: "coarse_single_modular_ft64_hourly_checkpoints"

config/model_configs/coarse_single_modular_ft64_monthly_checkpoints_co2.yml → config/model_configs/coarse_single_modular_ft64_hourly_checkpoints_co2.yml

@@ -1,5 +1,5 @@
-run_name: "coarse_single_modular_ft64_monthly_checkpoints_co2"
-monthly_checkpoint: true
+run_name: "coarse_single_modular_ft64_hourly_checkpoints_co2"
+hourly_checkpoint: true
 moist: "equil"
 vert_diff: "true"
 rad: "allskywithclear"
@@ -15,4 +15,4 @@ h_elem: 6
 dt_save_restart: "10days"
 precip_model: "0M"
 apply_limiter: false
-job_id: "coarse_single_modular_ft64_monthly_checkpoints_co2"
+job_id: "coarse_single_modular_ft64_hourly_checkpoints_co2"

config/model_configs/interactive_debug.yml

@@ -14,4 +14,4 @@ dt_save_restart: "5days"
 precip_model: "0M"
 run_name: "interactive_debug_run"
 job_id: "interactive_debug_run"
-monthly_checkpoint: true
+hourly_checkpoint: true

docs/src/timemanager.md

@@ -13,4 +13,8 @@ ClimaCoupler.TimeManager.datetime_to_strdate
 ClimaCoupler.TimeManager.trigger_callback
 ClimaCoupler.TimeManager.Monthly
 ClimaCoupler.TimeManager.EveryTimestep
+ClimaCoupler.TimeManager.trigger_callback!
+ClimaCoupler.TimeManager.HourlyCallback
+ClimaCoupler.TimeManager.MonthlyCallback
+ClimaCoupler.TimeManager.update_firstdayofmonth!
 ```

experiments/AMIP/modular/cli_options.jl

@@ -30,8 +30,8 @@ function argparse_settings()
         help = "Method to partition turbulent fluxes. [`PartitionedStateFluxes`, `CombinedStateFluxes`]"
         arg_type = String
         default = "CombinedStateFluxes"
-        "--monthly_checkpoint" # TODO generalize to any frequency
-        help = "Boolean flag indicating whether to checkpoint monthly"
+        "--hourly_checkpoint"
+        help = "Boolean flag indicating whether to checkpoint at intervals of 1 hour or multiple hours"
         arg_type = Bool
         default = false
         "--restart_dir"

experiments/AMIP/modular/components/atmosphere/climaatmos_init.jl

@@ -8,8 +8,8 @@ import ClimaCoupler.FluxCalculator:
     extrapolate_ρ_to_sfc,
     get_surface_params
 using ClimaCore: Fields.level, Geometry
-import ClimaCoupler.FieldExchanger: get_thermo_params
-import ClimaCoupler.Interfacer: get_field, update_field!, name, get_model_state_vector
+import ClimaCoupler.Interfacer: get_field, update_field!, name
+import ClimaCoupler.Checkpointer: get_model_state_vector
 using StaticArrays
 
 # the clima atmos `integrator` is now defined

experiments/AMIP/modular/coupler_driver_modular.jl

@@ -76,7 +76,16 @@ import ClimaCoupler.ConservationChecker:
 import ClimaCoupler.Utilities: swap_space!
 import ClimaCoupler.BCReader:
     bcfile_info_init, float_type_bcf, update_midmonth_data!, next_date_in_file, interpolate_midmonth_to_daily
-import ClimaCoupler.TimeManager: current_date, datetime_to_strdate, trigger_callback, Monthly, EveryTimestep
+import ClimaCoupler.TimeManager:
+    current_date,
+    datetime_to_strdate,
+    trigger_callback,
+    Monthly,
+    EveryTimestep,
+    HourlyCallback,
+    MonthlyCallback,
+    update_firstdayofmonth!,
+    trigger_callback!
 import ClimaCoupler.Diagnostics: get_var, init_diagnostics, accumulate_diagnostics!, save_diagnostics, TimeMean
 import ClimaCoupler.PostProcessor: postprocess
 
@@ -90,8 +99,7 @@ import ClimaCoupler.Interfacer:
     LandModelSimulation,
     OceanModelSimulation,
     get_field,
-    update_field!,
-    update_sim!
+    update_field!
 import ClimaCoupler.FluxCalculator:
     PartitionedStateFluxes,
     CombinedStateFluxes,
@@ -156,7 +164,7 @@ tspan = (Int(0), t_end)
 saveat = time_to_seconds(config_dict["dt_save_to_sol"])
 date0 = date = DateTime(config_dict["start_date"], dateformat"yyyymmdd")
 mono_surface = config_dict["mono_surface"]
-monthly_checkpoint = config_dict["monthly_checkpoint"]
+hourly_checkpoint = config_dict["hourly_checkpoint"]
 restart_dir = config_dict["restart_dir"]
 restart_t = Int(config_dict["restart_t"])
 
@@ -439,6 +447,13 @@ if energy_check
     conservation_checks = (; energy = EnergyConservationCheck(model_sims), water = WaterConservationCheck(model_sims))
 end
 
+dir_paths = (; output = COUPLER_OUTPUT_DIR, artifacts = COUPLER_ARTIFACTS_DIR)
+checkpoint_cb =
+    HourlyCallback(dt = FT(48), func = checkpoint_sims, ref_date = [dates.date[1]], active = hourly_checkpoint) # bi-daily
+update_firstdayofmonth!_cb =
+    MonthlyCallback(dt = FT(1), func = update_firstdayofmonth!, ref_date = [dates.date1[1]], active = true) # for BCReader
+callbacks = (; checkpoint = checkpoint_cb, update_firstdayofmonth! = update_firstdayofmonth!_cb)
+
 ## coupler simulation
 cs = CoupledSimulation{FT}(
     comms_ctx,
@@ -454,6 +469,8 @@ cs = CoupledSimulation{FT}(
     model_sims,
     mode_specifics,
     diagnostics,
+    callbacks,
+    dir_paths,
 );
 
 #=
@@ -572,13 +589,13 @@ function solve_coupler!(cs)
 
         ## run component models sequentially for one coupling timestep (Δt_cpl)
         ClimaComms.barrier(comms_ctx)
+        update_surface_fractions!(cs)
         update_model_sims!(cs.model_sims, cs.fields, turbulent_fluxes)
 
         ## step sims
         step_model_sims!(cs.model_sims, t)
 
         ## exchange combined fields and (if specified) calculate fluxes using combined states
-        update_surface_fractions!(cs)
         import_combined_surface_fields!(cs.fields, cs.model_sims, cs.boundary_space, turbulent_fluxes) # i.e. T_sfc, albedo, z0, beta
         if turbulent_fluxes isa CombinedStateFluxes
             combined_turbulent_fluxes!(cs.model_sims, cs.fields, turbulent_fluxes) # this updates the surface thermo state, sfc_ts, in ClimaAtmos (but also unnecessarily calculates fluxes)
@@ -594,19 +611,11 @@ function solve_coupler!(cs)
 
         import_atmos_fields!(cs.fields, cs.model_sims, cs.boundary_space, turbulent_fluxes) # radiative and/or turbulent
 
-        ## monthly callbacks
-        if trigger_callback(cs, Monthly())
-            ## step to the next calendar month
-            cs.dates.date1[1] += Dates.Month(1)
-            ## checkpoint model state
-            if monthly_checkpoint
-                for sim in cs.model_sims
-                    if get_model_state_vector(sim) !== nothing
-                        checkpoint_model_state(sim, comms_ctx, Int(t), output_dir = COUPLER_ARTIFACTS_DIR)
-                    end
-                end
-            end
-        end
+        ## callback to update the fist day of month if needed (for BCReader)
+        trigger_callback!(cs, cs.callbacks.update_firstdayofmonth!)
+
+        ## callback to checkpoint model state
+        trigger_callback!(cs, cs.callbacks.checkpoint)
 
     end
     @show walltime

experiments/AMIP/modular/user_io/user_logging.jl

@@ -1,3 +1,5 @@
+using ClimaCoupler: Checkpointer
+
 """
     Base.show(io::IO, dict::Dict)
 
@@ -8,3 +10,19 @@ function Base.show(io::IO, dict::Dict)
         println(io, " $k => $(dict[k])")
     end
 end
+
+# user callbacks
+"""
+    checkpoint_sims(cs::CoupledSimulation, _)
+
+This is a callback function that checkpoints all simulations defined in the current coupled simulation.
+"""
+function checkpoint_sims(cs::CoupledSimulation, _)
+    for sim in cs.model_sims
+        if get_model_state_vector(sim) !== nothing
+            t = Dates.datetime2epochms(cs.dates.date[1])
+            t0 = Dates.datetime2epochms(cs.dates.date0[1])
+            Checkpointer.checkpoint_model_state(sim, cs.comms_ctx, Int((t - t0) / 1e3), output_dir = cs.dirs.artifacts)
+        end
+    end
+end

src/Checkpointer.jl

@@ -8,7 +8,6 @@ module Checkpointer
 using ClimaCore: Fields, InputOutput
 using ClimaCoupler: Interfacer
 using Dates
-using ClimaCoupler.TimeManager: AbstractFrequency, Monthly, EveryTimestep, trigger_callback
 using ClimaComms
 export get_model_state_vector, checkpoint_model_state, restart_model_state!
 

src/Interfacer.jl

@@ -46,6 +46,9 @@ struct CoupledSimulation{
     NTSM <: NamedTuple,
     NTMS <: NamedTuple,
     NTM <: NamedTuple,
+    TD <: Tuple,
+    NTC <: NamedTuple,
+    NTP <: NamedTuple,
 }
     comms_ctx::X
     dates::D
@@ -59,10 +62,12 @@ struct CoupledSimulation{
     surface_fractions::NTSM
     model_sims::NTMS
     mode::NTM
-    diagnostics::Tuple
+    diagnostics::TD
+    callbacks::NTC
+    dirs::NTP
 end
 
-CoupledSimulation{FT}(args...) where {FT} = CoupledSimulation{FT, typeof.(args[1:12])...}(args...)
+CoupledSimulation{FT}(args...) where {FT} = CoupledSimulation{FT, typeof.(args[1:end])...}(args...)
 
 """
     float_type(::CoupledSimulation)