Add support for saving checkpoints by calendar dates

This commit adds support for providing a calendar frequency to the function that
saves to the state to disk. This is required to ensure that diagnostics and
states are synced and output at the same time. To enable this, I use the same
infrastructure used by diagnostic module. Everything was already there, I just
needed to add the YAML parsing.

config/default_configs/default_config.yml

@@ -155,7 +155,7 @@ test_dycore_consistency:
   help: "Test dycore consistency [`false` (default), `true`]"
   value: false
 dt_save_state_to_disk:
-  help: "Time between saving the state to disk. Examples: [`10secs`, `1hours`, `Inf` (do not save)]"
+  help: "Time between saving the state to disk. Examples: [`10secs`, `1hours`, `1months`, `Inf` (do not save)]"
   value: "Inf"
 dt_save_to_sol:
   help: "Time between saving solution. Examples: [`10days`, `1hours`, `Inf` (do not save)]"

config/model_configs/single_column_radiative_equilibrium_gray.yml

@@ -1,4 +1,4 @@
-dt_save_state_to_disk: "100days"
+dt_save_state_to_disk: "3months"
 initial_condition: "IsothermalProfile"
 hyperdiff: false
 # It seems radiative equilibrium needs a larger dz near the top

src/callbacks/get_callbacks.jl

@@ -172,7 +172,7 @@ function get_diagnostics(parsed_args, atmos_model, Y, p, dt)
     return diagnostics, writers
 end
 
-function get_callbacks(config, sim_info, atmos, params, Y, p)
+function get_callbacks(config, sim_info, atmos, params, Y, p, t_start)
     (; parsed_args, comms_ctx) = config
     FT = eltype(params)
     (; dt, output_dir) = sim_info
@@ -209,18 +209,41 @@ function get_callbacks(config, sim_info, atmos, params, Y, p)
         ),
     )
 
-    dt_save_state_to_disk =
-        time_to_seconds(parsed_args["dt_save_state_to_disk"])
-    if !(dt_save_state_to_disk == Inf)
+    if occursin("months", parsed_args["dt_save_state_to_disk"])
+        months = match(r"^(\d+)months$", parsed_args["dt_save_state_to_disk"])
+        isnothing(months) && error(
+            "$(period_str) has to be of the form <NUM>months, e.g. 2months for 2 months",
+        )
+        period_dates = Dates.Month(parse(Int, first(months)))
+        schedule = CAD.EveryCalendarDtSchedule(
+            period_dates;
+            reference_date = p.start_date,
+            date_last = p.start_date + Dates.Second(t_start),
+        )
         callbacks = (
             callbacks...,
-            call_every_dt(
+            SciMLBase.DiscreteCallback(
+                # Condition
+                (u, t, integrator) -> schedule(integrator),
+                # Action
                 (integrator) ->
                     save_state_to_disk_func(integrator, output_dir),
-                dt_save_state_to_disk;
-                skip_first = sim_info.restart,
             ),
         )
+    else
+        dt_save_state_to_disk =
+            time_to_seconds(parsed_args["dt_save_state_to_disk"])
+        if !(dt_save_state_to_disk == Inf)
+            callbacks = (
+                callbacks...,
+                call_every_dt(
+                    (integrator) ->
+                        save_state_to_disk_func(integrator, output_dir),
+                    dt_save_state_to_disk;
+                    skip_first = sim_info.restart,
+                ),
+            )
+        end
     end
 
     if is_distributed(comms_ctx)

src/solver/type_getters.jl

@@ -669,7 +669,7 @@ function get_simulation(config::AtmosConfig)
     @info "ode_configuration: $s"
 
     s = @timed_str begin
-        callback = get_callbacks(config, sim_info, atmos, params, Y, p)
+        callback = get_callbacks(config, sim_info, atmos, params, Y, p, t_start)
     end
     @info "get_callbacks: $s"