move checkpoint_sims to Checkpointer; delete user_logging.jl

CliMA · Nov 7, 2024 · 73853e1 · 73853e1
1 parent 4d614cf
commit 73853e1
Show file tree

Hide file tree

Showing 10 changed files with 19 additions and 45 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -46,3 +46,5 @@ This PR does a few cleanup tasks in the experiments/ClimaEarth/ directory:
 - Delete the `viz_explorer.jl`, which was previously used to create animations of simulation fields, but is not currently being used. Note that these animations will no longer appear in buildkite output.
 - Move functions in the `io_helpers.jl` file to the Utilities module and delete this file.
 This should not change any behavior.
+- Move `checkpoint_sims` function to the Checkpointer module and delete the `user_logging.jl` file.
+This should not change any behavior.
diff --git a/docs/src/checkpointer.md b/docs/src/checkpointer.md
@@ -8,4 +8,5 @@ This module contains general functions for logging the model states and restarti
     ClimaCoupler.Checkpointer.get_model_prog_state
     ClimaCoupler.Checkpointer.restart_model_state!
     ClimaCoupler.Checkpointer.checkpoint_model_state
+    ClimaCoupler.Checkpointer.checkpoint_sims
 ```
diff --git a/experiments/ClimaEarth/run_amip.jl b/experiments/ClimaEarth/run_amip.jl
@@ -76,7 +76,6 @@ include("components/ocean/prescr_seaice.jl")
 include("components/ocean/eisenman_seaice.jl")
 
 ## helpers for user-specified IO
-include("user_io/user_logging.jl")
 include("user_io/debug_plots.jl")
 
 #=

diff --git a/experiments/ClimaEarth/run_cloudless_aquaplanet.jl b/experiments/ClimaEarth/run_cloudless_aquaplanet.jl
@@ -37,9 +37,6 @@ pkg_dir = pkgdir(ClimaCoupler)
 include("components/atmosphere/climaatmos.jl")
 include("components/ocean/slab_ocean.jl")
 
-## helpers for user-specified IO
-include("user_io/user_logging.jl")
-
 #=
 ### Setup simulation parameters
 =#

diff --git a/experiments/ClimaEarth/run_cloudy_aquaplanet.jl b/experiments/ClimaEarth/run_cloudy_aquaplanet.jl
@@ -37,9 +37,6 @@ pkg_dir = pkgdir(ClimaCoupler)
 include("components/atmosphere/climaatmos.jl")
 include("components/ocean/slab_ocean.jl")
 
-## helpers for user-specified IO
-include("user_io/user_logging.jl")
-
 #=
 ### Setup simulation parameters
 =#

diff --git a/experiments/ClimaEarth/run_cloudy_slabplanet.jl b/experiments/ClimaEarth/run_cloudy_slabplanet.jl
@@ -38,9 +38,6 @@ include("components/atmosphere/climaatmos.jl")
 include("components/ocean/slab_ocean.jl")
 include("components/land/climaland_bucket.jl")
 
-## helpers for user-specified IO
-include("user_io/user_logging.jl")
-
 #=
 ### Setup simulation parameters
 =#

diff --git a/experiments/ClimaEarth/run_dry_held_suarez.jl b/experiments/ClimaEarth/run_dry_held_suarez.jl
@@ -37,9 +37,6 @@ pkg_dir = pkgdir(ClimaCoupler)
 ## helpers for component models
 include("components/atmosphere/climaatmos.jl")
 
-## helpers for user-specified IO
-include("user_io/user_logging.jl")
-
 #=
 ### Setup simulation parameters
 Here we follow ClimaCore's dry Held-Suarez `held_suarez_rhoe` example.

diff --git a/experiments/ClimaEarth/run_moist_held_suarez.jl b/experiments/ClimaEarth/run_moist_held_suarez.jl
@@ -39,9 +39,6 @@ pkg_dir = pkgdir(ClimaCoupler)
 ## helpers for component models
 include("components/atmosphere/climaatmos.jl")
 
-## helpers for user-specified IO
-include("user_io/user_logging.jl")
-
 #=
 ### Setup simulation parameters
 Here we follow Thatcher and Jablonowski (2016).

diff --git a/experiments/ClimaEarth/user_io/user_logging.jl b/experiments/ClimaEarth/user_io/user_logging.jl
diff --git a/src/Checkpointer.jl b/src/Checkpointer.jl
@@ -9,7 +9,7 @@ import ClimaComms
 import ClimaCore as CC
 import ..Interfacer
 
-export get_model_prog_state, checkpoint_model_state, restart_model_state!
+export get_model_prog_state, checkpoint_model_state, restart_model_state!, checkpoint_sims
 
 """
     get_model_prog_state(sim::Interfacer.ComponentModelSimulation)
@@ -71,4 +71,19 @@ function restart_model_state!(
     Y .= Y_new
 end
 
+"""
+    checkpoint_sims(cs::CoupledSimulation, _)
+
+This is a callback function that checkpoints all simulations defined in the current coupled simulation.
+"""
+function checkpoint_sims(cs::Interfacer.CoupledSimulation, _)
+    for sim in cs.model_sims
+        if Checkpointer.get_model_prog_state(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
+
 end # module