Merge pull request #4162 from CliMA/zs/restart

add an option for reproducible restart

Author: szy21

NEWS.md

@@ -4,6 +4,10 @@ ClimaAtmos.jl Release Notes
 main
 -------
 
+PR [#4162](https://github.com/CliMA/ClimaAtmos.jl/pull/4162) adds an option for
+reproducible restart. It is set to false by default. This shouldn't affect restart
+in the coupler as the coupler save the cache for restarting.
+
 PR [#4021](https://github.com/CliMA/ClimaAtmos.jl/pull/4021) uses ClimaCore
 convenience constructors to create spaces without an AtmosConfig.
 

config/default_configs/default_config.yml

@@ -335,11 +335,14 @@ deep_atmosphere:
   help: "If true, use deep atmosphere equations and metric terms, otherwise assume columns are cylindrical (shallow atmosphere) [`true` (default), `false`]"
   value: true
 restart_file:
-  help: "Path to HDF5 file to use as simulation starting point"
+  help: "Path to HDF5 file to use as simulation starting point. Note that the simulation can only be restarted in a reproducible way when `reproducible_restart` is true."
   value: ~
 detect_restart_file:
   help: "When true, try finding a restart file and use it to restart the simulation. Only works with ActiveLink."
   value: false
+reproducible_restart:
+  help: "If true, the simulation is reproducible when restarting from a restart file. Disable this option when running production runs."
+  value: false
 prescribed_aerosols:
   help: "Which aerosols to add. List of keys from the data file (e.g., CB1, CB2)."
   value: []

docs/src/restarts.md

@@ -26,8 +26,12 @@ particularly useful for
 
 !!! note
 
-    The simulation can only be restarted in a reproducible way when using `GridScaleCloud`
-    for `cloud_model`.
+    By default, the simulation cannot be restarted in a reproducible way. To
+    enable reproducible restarts, you need to set `reproducible_restart` to `true`. 
+    When `reproducible_restart` is true, `ClimaAtmos` recalculates the grid_scale 
+    cloud fraction and uses it in the buoyancy gradient calculation to ensure deterministic 
+    behavior across restarts. We recommend disabling this option for production runs.
+
 
 ### How Restarts Work
 

src/cache/precomputed_quantities.jl

@@ -533,9 +533,9 @@ NVTX.@annotate function set_explicit_precomputed_quantities!(Y, p, t)
     # depends on the mixing length, which depends on the buoyancy gradient.
     # We break this circular dependency by using cloud fraction from the previous time step in the
     # buoyancy gradient calculation. This breaks reproducible restart in general,
-    # but we support reproducible restart with grid-scale cloud by recalculating the cloud fraction here.
-    if cloud_model isa GridScaleCloud
-        set_cloud_fraction!(Y, p, moisture_model, cloud_model)
+    # but we support reproducible restart by recalculating the cloud fraction with GridScaleCloud here.
+    if p.atmos.numerics.reproducible_restart isa ReproducibleRestart
+        set_cloud_fraction!(Y, p, moisture_model, GridScaleCloud())
     end
 
     if turbconv_model isa PrognosticEDMFX

src/solver/type_getters.jl

@@ -188,9 +188,12 @@ function get_scale_blending_method(parsed_args)
 end
 
 function get_numerics(parsed_args, FT)
-    test_dycore =
+    test_dycore_consistency =
         parsed_args["test_dycore_consistency"] ? TestDycoreConsistency() :
         nothing
+    reproducible_restart =
+        parsed_args["reproducible_restart"] ? ReproducibleRestart() :
+        nothing
 
     energy_q_tot_upwinding = Val(Symbol(parsed_args["energy_q_tot_upwinding"]))
     tracer_upwinding = Val(Symbol(parsed_args["tracer_upwinding"]))
@@ -227,7 +230,8 @@ function get_numerics(parsed_args, FT)
         edmfx_sgsflux_upwinding,
         edmfx_tracer_upwinding,
         limiter,
-        test_dycore_consistency = test_dycore,
+        test_dycore_consistency,
+        reproducible_restart,
         diff_mode,
         hyperdiff,
     )

src/solver/types.jl

@@ -466,6 +466,7 @@ function get_ρu₃qₜ_surface(flow::ShipwayHill2012VelocityProfile, thermo_par
 end
 
 struct TestDycoreConsistency end
+struct ReproducibleRestart end
 
 abstract type AbstractTimesteppingMode end
 struct Explicit <: AbstractTimesteppingMode end
@@ -478,7 +479,7 @@ struct SmoothMinimumBlending <: AbstractScaleBlendingMethod end
 struct HardMinimumBlending <: AbstractScaleBlendingMethod end
 Base.broadcastable(x::AbstractScaleBlendingMethod) = tuple(x)
 
-Base.@kwdef struct AtmosNumerics{EN_UP, TR_UP, ED_UP, SG_UP, ED_TR_UP, TDC, LIM, DM, HD}
+Base.@kwdef struct AtmosNumerics{EN_UP, TR_UP, ED_UP, SG_UP, ED_TR_UP, TDC, RR, LIM, DM, HD}
 
     """Enable specific upwinding schemes for specific equations"""
     energy_q_tot_upwinding::EN_UP
@@ -489,6 +490,8 @@ Base.@kwdef struct AtmosNumerics{EN_UP, TR_UP, ED_UP, SG_UP, ED_TR_UP, TDC, LIM,
 
     """Add NaNs to certain equations to track down problems"""
     test_dycore_consistency::TDC
+    """Whether the simulation is reproducible when restarting from a restart file"""
+    reproducible_restart::RR
 
     limiter::LIM
 
@@ -941,6 +944,7 @@ const _DEFAULT_ATMOS_MODEL_KWARGS = (
     edmfx_sgsflux_upwinding = Val(:none),
     edmfx_tracer_upwinding = Val(:first_order),
     test_dycore_consistency = nothing,
+    reproducible_restart = nothing,
     limiter = nothing,
     diff_mode = Explicit(),
     hyperdiff = nothing,

test/restart.jl

@@ -389,6 +389,7 @@ if MANYTESTS
                         job_id = "$(configuration)_$(moisture)_$(precip)_$(topography)_$(radiation)_$(turbconv_mode)"
                         test_dict = Dict(
                             "test_dycore_consistency" => true, # We will add NaNs to the cache, just to make sure
+                            "reproducible_restart" => true,
                             "check_nan_every" => 3,
                             "log_progress" => false,
                             "moist" => moisture,
@@ -445,6 +446,7 @@ else
             "h_elem" => 4,
             "z_elem" => 15,
             "test_dycore_consistency" => true, # We will add NaNs to the cache, just to make sure
+            "reproducible_restart" => true,
             "check_nan_every" => 3,
             "log_progress" => false,
             "dt" => "1secs",
@@ -455,7 +457,6 @@ else
             "output_dir" => joinpath(amip_output_loc, amip_job_id),
             "dt_cloud_fraction" => "1secs",
             "rad" => "allskywithclear",
-            "cloud_model" => "grid_scale",
             "toml" => [
                 joinpath(@__DIR__, "../toml/longrun_aquaplanet_diagedmf.toml"),
             ],