add NaN check callbacks to longruns

Author: juliasloan25

experiments/long_runs/bucket.jl

@@ -144,7 +144,11 @@ function setup_prob(t0, tf, Δt; outdir = outdir, nelements = (101, 7))
     diag_cb = ClimaDiagnostics.DiagnosticsCallback(diagnostic_handler)
 
     driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
-    return prob, SciMLBase.CallbackSet(driver_cb, diag_cb)
+
+    nancheck_freq = Dates.Month(1)
+    nancheck_cb = ClimaLand.NaNCheckCallback(nancheck_freq, start_date, t0, Δt)
+
+    return prob, SciMLBase.CallbackSet(driver_cb, diag_cb, nancheck_cb)
 end
 
 function setup_and_solve_problem(; greet = false)

experiments/long_runs/land.jl

@@ -374,7 +374,10 @@ function setup_prob(t0, tf, Δt; outdir = outdir, nelements = (101, 15))
     diag_cb = ClimaDiagnostics.DiagnosticsCallback(diagnostic_handler)
 
     driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
-    return prob, SciMLBase.CallbackSet(driver_cb, diag_cb)
+
+    nancheck_freq = Dates.Month(1)
+    nancheck_cb = ClimaLand.NaNCheckCallback(nancheck_freq, start_date, t0, Δt)
+    return prob, SciMLBase.CallbackSet(driver_cb, diag_cb, nancheck_cb)
 end
 
 function setup_and_solve_problem(; greet = false)

experiments/long_runs/land_region.jl

@@ -389,7 +389,10 @@ function setup_prob(t0, tf, Δt; outdir = outdir, nelements = (10, 10, 15))
     diag_cb = ClimaDiagnostics.DiagnosticsCallback(diagnostic_handler)
 
     driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
-    return prob, SciMLBase.CallbackSet(driver_cb, diag_cb)
+
+    nancheck_freq = Dates.Month(1)
+    nancheck_cb = ClimaLand.NaNCheckCallback(nancheck_freq, start_date, t0, Δt)
+    return prob, SciMLBase.CallbackSet(driver_cb, diag_cb, nancheck_cb)
 end
 
 function setup_and_solve_problem(; greet = false)

experiments/long_runs/snowy_land.jl

@@ -396,7 +396,11 @@ function setup_prob(t0, tf, Δt; outdir = outdir, nelements = (101, 15))
     end
     report_cb = SciMLBase.DiscreteCallback(every1000steps, report)
 
-    return prob, SciMLBase.CallbackSet(driver_cb, diag_cb, report_cb)
+    nancheck_freq = Dates.Month(1)
+    nancheck_cb = ClimaLand.NaNCheckCallback(nancheck_freq, start_date, t0, Δt)
+
+    return prob,
+    SciMLBase.CallbackSet(driver_cb, diag_cb, report_cb, nancheck_cb)
 end
 
 function setup_and_solve_problem(; greet = false)

experiments/long_runs/soil.jl

@@ -215,7 +215,10 @@ function setup_prob(t0, tf, Δt; outdir = outdir, nelements = (101, 15))
     diag_cb = ClimaDiagnostics.DiagnosticsCallback(diagnostic_handler)
 
     driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
-    return prob, SciMLBase.CallbackSet(driver_cb, diag_cb)
+
+    nancheck_freq = Dates.Month(1)
+    nancheck_cb = ClimaLand.NaNCheckCallback(nancheck_freq, start_date, t0, Δt)
+    return prob, SciMLBase.CallbackSet(driver_cb, diag_cb, nancheck_cb)
 end
 
 function setup_and_solve_problem(; greet = false)

src/shared_utilities/utils.jl

@@ -320,7 +320,7 @@ function CheckpointCallback(
 
     if !isnothing(dt)
         dt_period = Dates.Millisecond(1000dt)
-        if !isdivisible(checkpoint_frequency_period / dt_period)
+        if !isdivisible(checkpoint_frequency_period, dt_period)
             @warn "Checkpoint frequency ($(checkpoint_frequency_period)) is not an integer multiple of dt $(dt_period)"
         end
     end
@@ -529,3 +529,66 @@ function count_nans_state(state::ClimaCore.Fields.Field, mask = nothing)
     end
     return nothing
 end
+
+"""
+    NaNCheckCallback(nancheck_frequency::Union{AbstractFloat, Dates.Period},
+                        output_dir, start_date, t_start; model, dt)
+
+Constructs a DiscreteCallback which counts the number of NaNs in the state
+and produces a warning if any are found.
+
+# Arguments
+- `nancheck_frequency`: The frequency at which the state is checked for NaNs.
+  Can be specified as a float (in seconds) or a `Dates.Period`.
+- `start_date`: The start date of the simulation.
+- `t_start`: The starting time of the simulation (in seconds).
+- `dt`: The timestep of the model (optional), used to check for consistency.
+
+The callback uses `ClimaDiagnostics.EveryCalendarDtSchedule` to determine when
+to save checkpoints based on the `nancheck_frequency`. The schedule is
+initialized with the `start_date` and `t_start` to ensure that the first
+checkpoint is saved at the correct time.
+
+The `save_checkpoint` function is called with the current state vector `u`, the
+current time `t`, and the `output_dir` to save the checkpoint to disk.
+"""
+function NaNCheckCallback(
+    nancheck_frequency::Union{AbstractFloat, Dates.Period},
+    start_date,
+    t_start,
+    dt,
+)
+    # TODO: Move to a more general callback system. For the time being, we use
+    # the ClimaDiagnostics one because it is flexible and it supports calendar
+    # dates.
+
+    if nancheck_frequency isa AbstractFloat
+        # Assume it is in seconds, but go through Millisecond to support
+        # fractional seconds
+        nancheck_frequency_period = Dates.Millisecond(1000nancheck_frequency)
+    else
+        nancheck_frequency_period = nancheck_frequency
+    end
+
+    schedule = EveryCalendarDtSchedule(
+        nancheck_frequency_period;
+        start_date,
+        date_last = start_date + Dates.Millisecond(1000t_start),
+    )
+
+    if !isnothing(dt)
+        dt_period = Dates.Millisecond(1000dt)
+        if !isdivisible(nancheck_frequency_period, dt_period)
+            @warn "Callback frequency ($(nancheck_frequency_period)) is not an integer multiple of dt $(dt_period)"
+        end
+    end
+
+    cond = let schedule = schedule
+        (u, t, integrator) -> schedule(integrator)
+    end
+    affect! = let output_dir = output_dir, model = model
+        (integrator) -> count_nans_state(integrator.u)
+    end
+
+    SciMLBase.DiscreteCallback(cond, affect!)
+end