From ccf9959b31acd92175092961c510b3625b311668 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Wed, 18 Sep 2024 20:25:01 -0700
Subject: [PATCH] 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     |  2 +-
 ...ngle_column_radiative_equilibrium_gray.yml |  2 +-
 src/callbacks/get_callbacks.jl                | 44 ++++++++++++++-----
 src/solver/type_getters.jl                    |  2 +-
 4 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml
index c5f73c362e..56abd972b9 100644
--- a/config/default_configs/default_config.yml
+++ b/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)]"
diff --git a/config/model_configs/single_column_radiative_equilibrium_gray.yml b/config/model_configs/single_column_radiative_equilibrium_gray.yml
index 5e25e0fb8b..02b096b8fa 100644
--- a/config/model_configs/single_column_radiative_equilibrium_gray.yml
+++ b/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
diff --git a/src/callbacks/get_callbacks.jl b/src/callbacks/get_callbacks.jl
index f9b6408188..45ac611530 100644
--- a/src/callbacks/get_callbacks.jl
+++ b/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,38 @@ 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)
-        callbacks = (
-            callbacks...,
-            call_every_dt(
-                (integrator) ->
-                    save_state_to_disk_func(integrator, output_dir),
-                dt_save_state_to_disk;
-                skip_first = sim_info.restart,
-            ),
+    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),
+        )
+        cond = let schedule = schedule
+            (u, t, integrator) -> schedule(integrator)
+        end
+        affect! = let output_dir = output_dir
+            (integrator) -> save_state_to_disk_func(integrator, output_dir)
+        end
+        callbacks = (callbacks..., SciMLBase.DiscreteCallback(cond, affect!))
+    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)
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 2a59e7311c..20f437a911 100644
--- a/src/solver/type_getters.jl
+++ b/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"