Merge pull request #3710 from CliMA/gb/non_zero_t_start

Add option to restart from a non zero t and bump to 0.29.1

Author: Sbozzolo

NEWS.md

@@ -4,6 +4,15 @@ ClimaAtmos.jl Release Notes
 main
 -------
 
+v0.29.1
+-------
+
+### Add support for non-zero `t_start`
+
+Passing a non zero `t_start` is useful in conditions where one wants to have a
+specific `start_date`, but start the simulation from a different point. This is
+used by `ClimaCoupler` to restart simulations.
+
 v0.29.0
 -------
 

Project.toml

@@ -1,7 +1,7 @@
 name = "ClimaAtmos"
 uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 authors = ["Climate Modeling Alliance"]
-version = "0.29.0"
+version = "0.29.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

config/default_configs/default_config.yml

@@ -123,6 +123,9 @@ FLOAT_TYPE:
 dt:
   help: "Simulation time step. Examples: [`10secs`, `1hours`]"
   value: "600secs"
+t_start:
+  help: "Simulation start time. This is typically useful in the context of manually restarting a simulation (e.g., but reinitializing the initial state). It is overwritten when the simulation is restarted with the checkpointing system. Examples: [`0secs`, `40secs`]"
+  value: "0secs"
 t_end:
   help: "Simulation end time. Examples: [`1200days`, `40secs`]"
   value: "10days"

docs/src/restarts.md

@@ -47,6 +47,12 @@ is started.
 If is also possible to manually specify a restart file. In this case, this will
 override any file automatically detected.
 
+`ClimaAtmos` also comes with a configuration,` t_start`, to change the initial
+time of the simulation without changing the start date. This option can be
+useful manually restarting a simulation (e.g., by overwriting the initial
+conditions). When a simulation is restarted from a checkpoint, this option
+becomes redundant.
+
 ### Accumulated Diagnostics
 
 At the moment, `ClimaAtmos` does not support working with accumulated

src/callbacks/callbacks.jl

@@ -290,7 +290,7 @@ NVTX.@annotate function rrtmgp_model_callback!(integrator)
 
     set_surface_albedo!(Y, p, t, p.atmos.surface_albedo)
 
-    RRTMGPI.update_fluxes!(rrtmgp_model, UInt32(t / integrator.p.dt))
+    RRTMGPI.update_fluxes!(rrtmgp_model, UInt32(floor(t / integrator.p.dt)))
     Fields.field2array(ᶠradiation_flux) .= rrtmgp_model.face_flux
     return nothing
 end

src/callbacks/get_callbacks.jl

@@ -1,14 +1,6 @@
-function get_diagnostics(
-    parsed_args,
-    atmos_model,
-    Y,
-    p,
-    sim_info,
-    t_start,
-    output_dir,
-)
-
-    (; dt, start_date) = sim_info
+function get_diagnostics(parsed_args, atmos_model, Y, p, sim_info, output_dir)
+
+    (; dt, t_start, start_date) = sim_info
 
     FT = Spaces.undertype(axes(Y.c))
     context = ClimaComms.context(axes(Y.c))
@@ -249,16 +241,16 @@ function checkpoint_frequency_from_parsed_args(dt_save_state_to_disk::String)
 end
 
 
-function get_callbacks(config, sim_info, atmos, params, Y, p, t_start)
+function get_callbacks(config, sim_info, atmos, params, Y, p)
     (; parsed_args, comms_ctx) = config
     FT = eltype(params)
-    (; dt, output_dir, start_date) = sim_info
+    (; dt, output_dir, start_date, t_start, t_end) = sim_info
 
     callbacks = ()
     if parsed_args["log_progress"]
         @info "Progress logging enabled"
         walltime_info = WallTimeInfo()
-        tot_steps = ceil(Int, (sim_info.t_end - t_start) / dt)
+        tot_steps = ceil(Int, (t_end - t_start) / dt)
         five_percent_steps = ceil(Int, 0.05 * tot_steps)
         cond = let schedule = CappedGeometricSeriesSchedule(five_percent_steps)
             (u, t, integrator) -> schedule(integrator)
@@ -340,7 +332,7 @@ function get_callbacks(config, sim_info, atmos, params, Y, p, t_start)
             dt isa ITime ?
             ITime(time_to_seconds(parsed_args["dt_cloud_fraction"])) :
             FT(time_to_seconds(parsed_args["dt_cloud_fraction"]))
-        dt_cf, _, _, _ = promote(dt_cf, t_start, dt, sim_info.t_end)
+        dt_cf, _, _, _ = promote(dt_cf, t_start, dt, t_end)
         callbacks =
             (callbacks..., call_every_dt(cloud_fraction_model_callback!, dt_cf))
     end
@@ -349,7 +341,7 @@ function get_callbacks(config, sim_info, atmos, params, Y, p, t_start)
         dt_rad =
             dt isa ITime ? ITime(time_to_seconds(parsed_args["dt_rad"])) :
             FT(time_to_seconds(parsed_args["dt_rad"]))
-        dt_rad, _, _, _ = promote(dt_rad, t_start, dt, sim_info.t_end)
+        dt_rad, _, _, _ = promote(dt_rad, t_start, dt, t_end)
         # We use Millisecond to support fractional seconds, eg. 0.1
         dt_rad_ms = Dates.Millisecond(1_000 * float(dt_rad))
         if parsed_args["dt_save_state_to_disk"] != "Inf" &&

src/solver/type_getters.jl

@@ -604,12 +604,19 @@ function get_sim_info(config::AtmosConfig)
     isnothing(restart_file) ||
         @info "Restarting simulation from file $restart_file"
     epoch = DateTime(parsed_args["start_date"], dateformat"yyyymmdd")
+    t_start_int = time_to_seconds(parsed_args["t_start"])
+    if !isnothing(restart_file) && t_start_int != 0
+        @warn "Non zero `t_start` passed with a restarting simulation. The provided `t_start` will be ignored."
+    end
     if parsed_args["use_itime"]
         dt = ITime(time_to_seconds(parsed_args["dt"]))
+        t_start = ITime(time_to_seconds(parsed_args["t_start"]), epoch = epoch)
         t_end = ITime(time_to_seconds(parsed_args["t_end"]), epoch = epoch)
-        (dt, t_end) = promote(dt, t_end)
+        # ITime(0) is added for backward compatibility (since t_start used to always be 0)
+        (dt, t_start, t_end, _) = promote(dt, t_start, t_end, ITime(0))
     else
         dt = FT(time_to_seconds(parsed_args["dt"]))
+        t_start = FT(time_to_seconds(parsed_args["t_start"]))
         t_end = FT(time_to_seconds(parsed_args["t_end"]))
     end
     sim = (;
@@ -619,12 +626,14 @@ function get_sim_info(config::AtmosConfig)
         job_id,
         dt = dt,
         start_date = epoch,
+        t_start = t_start,
         t_end = t_end,
     )
-    n_steps = floor(Int, sim.t_end / sim.dt)
+    n_steps = floor(Int, (sim.t_end - sim.t_start) / sim.dt)
     @info(
         "Time info:",
         dt = parsed_args["dt"],
+        t_start = parsed_args["t_start"],
         t_end = parsed_args["t_end"],
         floor_n_steps = n_steps,
     )
@@ -734,6 +743,8 @@ function get_simulation(config::AtmosConfig)
                 hash(atmos),
             )
             spaces = get_spaces_restart(Y)
+            # Fix the t_start in sim_info with the one from the restart
+            sim_info = merge(sim_info, (; t_start))
         end
         @info "Allocating Y: $s"
     else
@@ -749,14 +760,6 @@ function get_simulation(config::AtmosConfig)
                 spaces.center_space,
                 spaces.face_space,
             )
-            if sim_info.dt isa ITime
-                t_start = ITime(
-                    Spaces.undertype(axes(Y.c))(0);
-                    epoch = sim_info.start_date,
-                )
-            else
-                t_start = Spaces.undertype(axes(Y.c))(0)
-            end
         end
         @info "Allocating Y: $s"
 
@@ -800,7 +803,7 @@ function get_simulation(config::AtmosConfig)
     @info "ode_configuration: $s"
 
     s = @timed_str begin
-        callback = get_callbacks(config, sim_info, atmos, params, Y, p, t_start)
+        callback = get_callbacks(config, sim_info, atmos, params, Y, p)
     end
     @info "get_callbacks: $s"
 
@@ -814,7 +817,6 @@ function get_simulation(config::AtmosConfig)
                     Y,
                     p,
                     sim_info,
-                    t_start,
                     output_dir,
                 )
         end
@@ -853,7 +855,7 @@ function get_simulation(config::AtmosConfig)
     @info "n_steps_per_cycle_per_cb (non diagnostics): $steps_cycle_non_diag"
     @info "n_steps_per_cycle (non diagnostics): $steps_cycle"
 
-    tspan = (t_start, sim_info.t_end)
+    tspan = (sim_info.t_start, sim_info.t_end)
     s = @timed_str begin
         integrator_args, integrator_kwargs = args_integrator(
             config.parsed_args,

test/coupler_compatibility.jl

@@ -213,6 +213,8 @@ end
     # Verify that using PrescribedSurface does not break the initialization of
     # RRTMGP or diagnostic EDMF. We currently need a moisture model in order to
     # use diagnostic EDMF.
+    #
+    # Also verify we can start with a different t_start than 0
     config = CA.AtmosConfig(
         Dict(
             "surface_setup" => "PrescribedSurface",
@@ -225,6 +227,7 @@ end
             # remove the following line and check that the test runs in less than a few
             # minutes on GitHub
             "output_default_diagnostics" => false,
+            "t_start" => "1secs",
         );
         job_id = "coupler_compatibility3",
     )