add NaN check for state

Adds a function that quantitatively checks how many NaNs are present in the state, and displays this information to the user. This can be used to inspect sol.u[end] after a simulation has run, to see if any NaNs were produced at the end of the simulation.
If no NaNs are found, this information is logged in an info statement when run in verbose mode. If NaNs are found, this information is logged in a warn statement.

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

@@ -3,7 +3,7 @@ import SciMLBase
 import ClimaDiagnostics.Schedules: EveryCalendarDtSchedule
 import Dates
 
-export FTfromY
+export FTfromY, count_nans_state
 
 """
      heaviside(x::FT)::FT where {FT}
@@ -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
@@ -492,3 +492,109 @@ function isdivisible(
     # have any common divisor)
     return isinteger(Dates.Day(1) / dt_small)
 end
+
+"""
+    count_nans_state(state, mask::ClimaCore.Fields.Field = nothing, verbose = false)
+
+Count the number of NaNs in the state, e.g. the FieldVector given by
+`sol.u[end]` after calling `solve`. This function is useful for
+debugging simulations to determine quantitatively if a simulation is stable.
+
+If this function is called on a FieldVector, it will recursively call itself
+on each Field in the FieldVector. If it is called on a Field, it will count
+the number of NaNs in the Field and produce a warning if any are found.
+
+If a ClimaCore Field is provided as `mask`, the function will only count NaNs
+in the state variables where the mask is 1. This is intended to be used with
+the land/sea mask, to avoid counting NaNs over the ocean. Note this assumes
+the mask is 1 over land and 0 over ocean.
+
+The `verbose` argument toggles whether the function produces output when no
+NaNs are found.
+"""
+function count_nans_state(
+    state::ClimaCore.Fields.FieldVector;
+    mask = nothing,
+    verbose = false,
+)
+    for pn in propertynames(state)
+        state_new = getproperty(state, pn)
+        @info "Checking NaNs in $pn"
+        count_nans_state(state_new; mask, verbose)
+    end
+    return nothing
+end
+
+function count_nans_state(
+    state::ClimaCore.Fields.Field;
+    mask = nothing,
+    verbose = false,
+)
+    num_nans =
+        isnothing(mask) ? count(==(1), isnan.(parent(state))) :
+        count(==(1), isnan.(parent(state)) .* parent(mask))
+    if num_nans > 0
+        @warn "$num_nans NaNs found"
+    else
+        verbose && @info "No NaNs found"
+    end
+    return nothing
+end
+
+"""
+    NaNCheckCallback(nancheck_frequency::Union{AbstractFloat, Dates.Period},
+                        start_date, t_start, 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 check for NaNs based on the `nancheck_frequency`. The schedule is
+initialized with the `start_date` and `t_start` to ensure that it is first
+called at the correct time.
+"""
+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! = (integrator) -> count_nans_state(integrator.u)
+
+    SciMLBase.DiscreteCallback(cond, affect!)
+end

test/shared_utilities/utilities.jl

@@ -2,7 +2,6 @@ using Test
 import ClimaComms
 ClimaComms.@import_required_backends
 using ClimaCore: Spaces, Geometry, Fields
-import ClimaComms
 using ClimaLand
 using ClimaLand: Domains, condition, SavingAffect, saving_initialize
 
@@ -256,3 +255,108 @@ end
         @test Y.subfields.subfield2 == Y_copy.subfields.subfield2
     end
 end
+
+@testset "count_nans_state, FT = $FT" begin
+    # Test on a 3D spherical domain
+    domain = ClimaLand.Domains.SphericalShell(;
+        radius = FT(2),
+        depth = FT(1.0),
+        nelements = (10, 5),
+        npolynomial = 3,
+    )
+
+    # Construct some fields
+    space = domain.space.subsurface
+    var1 = Fields.zeros(space)
+    var2 = Fields.zeros(space)
+    var3 = Fields.zeros(space)
+    fieldvec = Fields.FieldVector(var2 = var2, var3 = var3)
+
+    # Construct a FieldVector containing the fields and a nested FieldVector
+    Y = Fields.FieldVector(var1 = var1, fieldvec = fieldvec)
+
+    # Count and log the number of NaNs in the state (test verbose and non-verbose cases)
+    @test_logs (:info, "Checking NaNs in var1") (:info, "No NaNs found") (
+        :info,
+        "Checking NaNs in fieldvec",
+    ) (:info, "Checking NaNs in var2") (:info, "No NaNs found") (
+        :info,
+        "Checking NaNs in var3",
+    ) (:info, "No NaNs found") ClimaLand.count_nans_state(Y, verbose = true)
+
+    @test_logs (:info, "Checking NaNs in var1") (
+        :info,
+        "Checking NaNs in fieldvec",
+    ) (:info, "Checking NaNs in var2") (:info, "Checking NaNs in var3") ClimaLand.count_nans_state(
+        Y,
+    )
+
+    # Add some NaNs to the fields without scalar indexing
+    ArrayType = ClimaComms.array_type(ClimaComms.device())
+    var1_copy = Array(parent(var1))
+    var1_copy[1] = NaN
+    parent(var1) .= ArrayType(var1_copy)
+    var2_copy = Array(parent(var1))
+    var2_copy[1] = NaN
+    var2_copy[2] = NaN
+    parent(var2) .= ArrayType(var2_copy)
+
+    # Count and log the number of NaNs in the state (test verbose and non-verbose cases)
+    @test_logs (:info, "Checking NaNs in var1") (:warn, "1 NaNs found") (
+        :info,
+        "Checking NaNs in fieldvec",
+    ) (:info, "Checking NaNs in var2") (:warn, "2 NaNs found") (
+        :info,
+        "Checking NaNs in var3",
+    ) (:info, "No NaNs found") ClimaLand.count_nans_state(Y, verbose = true)
+
+    @test_logs (:info, "Checking NaNs in var1") (:warn, "1 NaNs found") (
+        :info,
+        "Checking NaNs in fieldvec",
+    ) (:info, "Checking NaNs in var2") (:warn, "2 NaNs found") (
+        :info,
+        "Checking NaNs in var3",
+    ) ClimaLand.count_nans_state(Y)
+
+    # Test with a mask
+    mask_zeros = Fields.zeros(space)
+    @test_logs (:info, "Checking NaNs in var1") (:info, "No NaNs found") (
+        :info,
+        "Checking NaNs in fieldvec",
+    ) (:info, "Checking NaNs in var2") (:info, "No NaNs found") (
+        :info,
+        "Checking NaNs in var3",
+    ) (:info, "No NaNs found") ClimaLand.count_nans_state(
+        Y,
+        mask = mask_zeros,
+        verbose = true,
+    )
+    @test_logs (:info, "Checking NaNs in var1") (
+        :info,
+        "Checking NaNs in fieldvec",
+    ) (:info, "Checking NaNs in var2") (:info, "Checking NaNs in var3") ClimaLand.count_nans_state(
+        Y,
+        mask = mask_zeros,
+    )
+
+    mask_ones = Fields.ones(space)
+    @test_logs (:info, "Checking NaNs in var1") (:warn, "1 NaNs found") (
+        :info,
+        "Checking NaNs in fieldvec",
+    ) (:info, "Checking NaNs in var2") (:warn, "2 NaNs found") (
+        :info,
+        "Checking NaNs in var3",
+    ) (:info, "No NaNs found") ClimaLand.count_nans_state(
+        Y,
+        mask = mask_ones,
+        verbose = true,
+    )
+
+    @test_logs (:info, "Checking NaNs in var1") (:warn, "1 NaNs found") (
+        :info,
+        "Checking NaNs in fieldvec",
+    ) (:info, "Checking NaNs in var2") (:warn, "2 NaNs found") (
+        :info,
+        "Checking NaNs in var3",
+    ) ClimaLand.count_nans_state(Y, mask = mask_ones)
+end