use ClimaUtilities CallbackManager

CliMA · Sep 10, 2024 · 6fbb583 · 6fbb583
1 parent e078085
commit 6fbb583
Show file tree

Hide file tree

Showing 17 changed files with 84 additions and 471 deletions.
diff --git a/docs/make.jl b/docs/make.jl
@@ -65,7 +65,6 @@ interface_pages = [
     "interfacer.md",
     "postprocessor.md",
     "regridder.md",
-    "timemanager.md",
     "utilities.md",
 ]
 performance_pages = ["performance.md"]

diff --git a/docs/src/timemanager.md b/docs/src/timemanager.md
diff --git a/experiments/ClimaEarth/run_amip.jl b/experiments/ClimaEarth/run_amip.jl
@@ -48,19 +48,12 @@ import ClimaCore as CC
 # ## Coupler specific imports
 import ClimaCoupler
 import ClimaCoupler:
-    ConservationChecker,
-    Checkpointer,
-    Diagnostics,
-    FieldExchanger,
-    FluxCalculator,
-    Interfacer,
-    Regridder,
-    TimeManager,
-    Utilities
+    ConservationChecker, Checkpointer, Diagnostics, FieldExchanger, FluxCalculator, Interfacer, Regridder, Utilities
 
 import ClimaUtilities.SpaceVaryingInputs: SpaceVaryingInput
 import ClimaUtilities.TimeVaryingInputs: TimeVaryingInput, evaluate!
 import ClimaUtilities.ClimaArtifacts: @clima_artifact
+import ClimaUtilities: CallbackManager
 import Interpolations
 
 pkg_dir = pkgdir(ClimaCoupler)
@@ -472,7 +465,7 @@ if use_coupler_diagnostics
     monthly_3d_diags = Diagnostics.init_diagnostics(
         (:T, :u, :q_tot, :q_liq_ice),
         atmos_sim.domain.center_space;
-        save = TimeManager.Monthly(),
+        save = CallbackManager.Monthly(),
         operations = (; accumulate = Diagnostics.TimeMean([Int(0)])),
         output_dir = dir_paths.output,
         name_tag = "monthly_mean_3d_",
@@ -481,7 +474,7 @@ if use_coupler_diagnostics
     monthly_2d_diags = Diagnostics.init_diagnostics(
         (:precipitation_rate, :toa_fluxes, :T_sfc, :turbulent_energy_fluxes),
         boundary_space;
-        save = TimeManager.Monthly(),
+        save = CallbackManager.Monthly(),
         operations = (; accumulate = Diagnostics.TimeMean([Int(0)])),
         output_dir = dir_paths.output,
         name_tag = "monthly_mean_2d_",
@@ -530,20 +523,20 @@ The currently implemented callbacks are:
   NB: Eventually, we will call all of radiation from the coupler, in addition to the albedo calculation.
 =#
 
-checkpoint_cb = TimeManager.HourlyCallback(
+checkpoint_cb = CallbackManager.HourlyCallback(
     dt = hourly_checkpoint_dt,
     func = checkpoint_sims,
     ref_date = [dates.date[1]],
     active = hourly_checkpoint,
 ) # 20 days
-update_firstdayofmonth!_cb = TimeManager.MonthlyCallback(
+update_firstdayofmonth!_cb = CallbackManager.MonthlyCallback(
     dt = FT(1),
-    func = TimeManager.update_firstdayofmonth!,
+    func = CallbackManager.update_firstdayofmonth!,
     ref_date = [dates.date1[1]],
     active = true,
 )
 dt_water_albedo = parse(FT, filter(x -> !occursin(x, "hours"), dt_rad))
-albedo_cb = TimeManager.HourlyCallback(
+albedo_cb = CallbackManager.HourlyCallback(
     dt = dt_water_albedo,
     func = FluxCalculator.water_albedo_from_atmosphere!,
     ref_date = [dates.date[1]],
@@ -683,7 +676,7 @@ function solve_coupler!(cs)
     ## step in time
     for t in ((tspan[begin] + Δt_cpl):Δt_cpl:tspan[end])
 
-        cs.dates.date[1] = TimeManager.current_date(cs, t)
+        cs.dates.date[1] = Interfacer.current_date(cs, t)
 
         ## print date on the first of month
         if cs.dates.date[1] >= cs.dates.date1[1]
@@ -717,7 +710,7 @@ function solve_coupler!(cs)
 
         ## update water albedo from wind at dt_water_albedo
         ## (this will be extended to a radiation callback from the coupler)
-        TimeManager.trigger_callback!(cs, cs.callbacks.water_albedo)
+        CallbackManager.trigger_callback!(cs, cs.callbacks.water_albedo)
 
 
         ## update the surface fractions for surface models,
@@ -752,10 +745,10 @@ function solve_coupler!(cs)
         FieldExchanger.import_atmos_fields!(cs.fields, cs.model_sims, cs.boundary_space, cs.turbulent_fluxes) # radiative and/or turbulent
 
         ## callback to update the fist day of month if needed
-        TimeManager.trigger_callback!(cs, cs.callbacks.update_firstdayofmonth!)
+        CallbackManager.trigger_callback!(cs, cs.callbacks.update_firstdayofmonth!)
 
         ## callback to checkpoint model state
-        TimeManager.trigger_callback!(cs, cs.callbacks.checkpoint)
+        CallbackManager.trigger_callback!(cs, cs.callbacks.checkpoint)
     end
     return nothing
 end

diff --git a/experiments/ClimaEarth/run_cloudless_aquaplanet.jl b/experiments/ClimaEarth/run_cloudless_aquaplanet.jl
@@ -32,7 +32,7 @@ import ClimaCoupler:
     FluxCalculator,
     Interfacer,
     Regridder,
-    TimeManager,
+    CallbackManager,
     Utilities
 
 pkg_dir = pkgdir(ClimaCoupler)
@@ -217,20 +217,20 @@ dates = (; date = [date], date0 = [date0], date1 = [Dates.firstdayofmonth(date0)
 ## Initialize Callbacks
 =#
 
-checkpoint_cb = TimeManager.HourlyCallback(
+checkpoint_cb = CallbackManager.HourlyCallback(
     dt = FT(480),
     func = checkpoint_sims,
     ref_date = [dates.date[1]],
     active = hourly_checkpoint,
 ) # 20 days
-update_firstdayofmonth!_cb = TimeManager.MonthlyCallback(
+update_firstdayofmonth!_cb = CallbackManager.MonthlyCallback(
     dt = FT(1),
-    func = TimeManager.update_firstdayofmonth!,
+    func = CallbackManager.update_firstdayofmonth!,
     ref_date = [dates.date1[1]],
     active = true,
 )
 dt_water_albedo = parse(FT, filter(x -> !occursin(x, "hours"), dt_rad))
-albedo_cb = TimeManager.HourlyCallback(
+albedo_cb = CallbackManager.HourlyCallback(
     dt = dt_water_albedo,
     func = FluxCalculator.water_albedo_from_atmosphere!,
     ref_date = [dates.date[1]],
@@ -325,7 +325,7 @@ function solve_coupler!(cs)
     ## step in time
     for t in ((tspan[begin] + Δt_cpl):Δt_cpl:tspan[end])
 
-        cs.dates.date[1] = TimeManager.current_date(cs, t)
+        cs.dates.date[1] = Interfacer.current_date(cs, t)
 
         ## print date on the first of month
         if cs.dates.date[1] >= cs.dates.date1[1]
@@ -335,7 +335,7 @@ function solve_coupler!(cs)
         ClimaComms.barrier(comms_ctx)
 
         ## update water albedo from wind at dt_water_albedo (this will be extended to a radiation callback from the coupler)
-        TimeManager.trigger_callback!(cs, cs.callbacks.water_albedo)
+        CallbackManager.trigger_callback!(cs, cs.callbacks.water_albedo)
 
         ## run component models sequentially for one coupling timestep (Δt_cpl)
         FieldExchanger.update_model_sims!(cs.model_sims, cs.fields, cs.turbulent_fluxes)
@@ -350,10 +350,10 @@ function solve_coupler!(cs)
         FieldExchanger.import_atmos_fields!(cs.fields, cs.model_sims, cs.boundary_space, cs.turbulent_fluxes) # radiative and/or turbulent
 
         ## callback to update the fist day of month if needed
-        TimeManager.trigger_callback!(cs, cs.callbacks.update_firstdayofmonth!)
+        CallbackManager.trigger_callback!(cs, cs.callbacks.update_firstdayofmonth!)
 
         ## callback to checkpoint model state
-        TimeManager.trigger_callback!(cs, cs.callbacks.checkpoint)
+        CallbackManager.trigger_callback!(cs, cs.callbacks.checkpoint)
 
     end
 

diff --git a/experiments/ClimaEarth/run_cloudy_aquaplanet.jl b/experiments/ClimaEarth/run_cloudy_aquaplanet.jl
@@ -32,7 +32,7 @@ import ClimaCoupler:
     FluxCalculator,
     Interfacer,
     Regridder,
-    TimeManager,
+    CallbackManager,
     Utilities
 
 pkg_dir = pkgdir(ClimaCoupler)
@@ -238,15 +238,15 @@ dates = (; date = [date], date0 = [date0], date1 = [Dates.firstdayofmonth(date0)
 ## Initialize Callbacks
 =#
 
-checkpoint_cb = TimeManager.HourlyCallback(
+checkpoint_cb = CallbackManager.HourlyCallback(
     dt = FT(480),
     func = checkpoint_sims,
     ref_date = [dates.date[1]],
     active = hourly_checkpoint,
 ) # 20 days
-update_firstdayofmonth!_cb = TimeManager.MonthlyCallback(
+update_firstdayofmonth!_cb = CallbackManager.MonthlyCallback(
     dt = FT(1),
-    func = TimeManager.update_firstdayofmonth!,
+    func = CallbackManager.update_firstdayofmonth!,
     ref_date = [dates.date1[1]],
     active = true,
 )
@@ -337,7 +337,7 @@ function solve_coupler!(cs)
     ## step in time
     for t in ((tspan[begin] + Δt_cpl):Δt_cpl:tspan[end])
 
-        cs.dates.date[1] = TimeManager.current_date(cs, t)
+        cs.dates.date[1] = Interfacer.current_date(cs, t)
 
         ## print date on the first of month
         if cs.dates.date[1] >= cs.dates.date1[1]
@@ -359,10 +359,10 @@ function solve_coupler!(cs)
         FieldExchanger.import_atmos_fields!(cs.fields, cs.model_sims, cs.boundary_space, cs.turbulent_fluxes) # radiative and/or turbulent
 
         ## callback to update the fist day of month if needed
-        TimeManager.trigger_callback!(cs, cs.callbacks.update_firstdayofmonth!)
+        CallbackManager.trigger_callback!(cs, cs.callbacks.update_firstdayofmonth!)
 
         ## callback to checkpoint model state
-        TimeManager.trigger_callback!(cs, cs.callbacks.checkpoint)
+        CallbackManager.trigger_callback!(cs, cs.callbacks.checkpoint)
 
     end
 

diff --git a/experiments/ClimaEarth/run_cloudy_slabplanet.jl b/experiments/ClimaEarth/run_cloudy_slabplanet.jl
@@ -32,7 +32,7 @@ import ClimaCoupler:
     FluxCalculator,
     Interfacer,
     Regridder,
-    TimeManager,
+    CallbackManager,
     Utilities
 
 pkg_dir = pkgdir(ClimaCoupler)
@@ -284,20 +284,20 @@ model_sims = (atmos_sim = atmos_sim, ocean_sim = ocean_sim);
 ## Initialize Callbacks
 =#
 
-checkpoint_cb = TimeManager.HourlyCallback(
+checkpoint_cb = CallbackManager.HourlyCallback(
     dt = FT(480),
     func = checkpoint_sims,
     ref_date = [dates.date[1]],
     active = hourly_checkpoint,
 ) # 20 days
-update_firstdayofmonth!_cb = TimeManager.MonthlyCallback(
+update_firstdayofmonth!_cb = CallbackManager.MonthlyCallback(
     dt = FT(1),
-    func = TimeManager.update_firstdayofmonth!,
+    func = CallbackManager.update_firstdayofmonth!,
     ref_date = [dates.date1[1]],
     active = true,
 )
 dt_water_albedo = parse(FT, filter(x -> !occursin(x, "hours"), dt_rad))
-albedo_cb = TimeManager.HourlyCallback(
+albedo_cb = CallbackManager.HourlyCallback(
     dt = dt_water_albedo,
     func = FluxCalculator.water_albedo_from_atmosphere!,
     ref_date = [dates.date[1]],
@@ -389,7 +389,7 @@ function solve_coupler!(cs)
     ## step in time
     for t in ((tspan[begin] + Δt_cpl):Δt_cpl:tspan[end])
 
-        cs.dates.date[1] = TimeManager.current_date(cs, t)
+        cs.dates.date[1] = Interfacer.current_date(cs, t)
 
         ## print date on the first of month
         if cs.dates.date[1] >= cs.dates.date1[1]
@@ -399,7 +399,7 @@ function solve_coupler!(cs)
         ClimaComms.barrier(comms_ctx)
 
         ## update water albedo from wind at dt_water_albedo (this will be extended to a radiation callback from the coupler)
-        TimeManager.trigger_callback!(cs, cs.callbacks.water_albedo)
+        CallbackManager.trigger_callback!(cs, cs.callbacks.water_albedo)
 
         ## run component models sequentially for one coupling timestep (Δt_cpl)
         FieldExchanger.update_model_sims!(cs.model_sims, cs.fields, cs.turbulent_fluxes)
@@ -414,10 +414,10 @@ function solve_coupler!(cs)
         FieldExchanger.import_atmos_fields!(cs.fields, cs.model_sims, cs.boundary_space, cs.turbulent_fluxes) # radiative and/or turbulent
 
         ## callback to update the fist day of month if needed
-        TimeManager.trigger_callback!(cs, cs.callbacks.update_firstdayofmonth!)
+        CallbackManager.trigger_callback!(cs, cs.callbacks.update_firstdayofmonth!)
 
         ## callback to checkpoint model state
-        TimeManager.trigger_callback!(cs, cs.callbacks.checkpoint)
+        CallbackManager.trigger_callback!(cs, cs.callbacks.checkpoint)
 
     end
 

diff --git a/experiments/ClimaEarth/run_dry_held_suarez.jl b/experiments/ClimaEarth/run_dry_held_suarez.jl
@@ -26,7 +26,7 @@ import ClimaCore
 
 ## Coupler specific imports
 import ClimaCoupler
-import ClimaCoupler: Checkpointer, FieldExchanger, Interfacer, TimeManager, Utilities
+import ClimaCoupler: Checkpointer, FieldExchanger, Interfacer, CallbackManager, Utilities
 
 pkg_dir = pkgdir(ClimaCoupler)
 
@@ -179,15 +179,15 @@ dates = (; date = [date], date0 = [date0], date1 = [Dates.firstdayofmonth(date0)
 #=
 ## Initialize Callbacks
 =#
-checkpoint_cb = TimeManager.HourlyCallback(
+checkpoint_cb = CallbackManager.HourlyCallback(
     dt = FT(480),
     func = checkpoint_sims,
     ref_date = [dates.date[1]],
     active = hourly_checkpoint,
 ) # 20 days TODO: not GPU friendly
-update_firstdayofmonth!_cb = TimeManager.MonthlyCallback(
+update_firstdayofmonth!_cb = CallbackManager.MonthlyCallback(
     dt = FT(1),
-    func = TimeManager.update_firstdayofmonth!,
+    func = CallbackManager.update_firstdayofmonth!,
     ref_date = [dates.date1[1]],
     active = true,
 )
@@ -237,7 +237,7 @@ function solve_coupler!(cs)
     ## step in time
     walltime = @elapsed for t in ((tspan[begin] + Δt_cpl):Δt_cpl:tspan[end])
 
-        cs.dates.date[1] = TimeManager.current_date(cs, t)
+        cs.dates.date[1] = Interfacer.current_date(cs, t)
 
         ## print date on the first of month
         if cs.dates.date[1] >= cs.dates.date1[1]
@@ -250,10 +250,10 @@ function solve_coupler!(cs)
         FieldExchanger.import_atmos_fields!(cs.fields, cs.model_sims, cs.boundary_space, cs.turbulent_fluxes) # radiative and/or turbulent
 
         ## callback to update the fist day of month if needed
-        TimeManager.trigger_callback!(cs, cs.callbacks.update_firstdayofmonth!)
+        CallbackManager.trigger_callback!(cs, cs.callbacks.update_firstdayofmonth!)
 
         ## callback to checkpoint model state
-        TimeManager.trigger_callback!(cs, cs.callbacks.checkpoint)
+        CallbackManager.trigger_callback!(cs, cs.callbacks.checkpoint)
 
     end
     ClimaComms.iamroot(comms_ctx) ? @show(walltime) : nothing