add state NaN checker

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.
If NaNs are found, this information is logged in a `warn` statement.

Author: juliasloan25

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}
@@ -492,3 +492,34 @@ function isdivisible(
     # have any common divisor)
     return isinteger(Dates.Day(1) / dt_small)
 end
+
+"""
+    count_nans_state(sol)
+
+Count the number of NaNs in the state variables. 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.
+
+Input: `state` - e.g. the FieldVector given by `sol.u[end]` after calling `solve`
+"""
+function count_nans_state(state::ClimaCore.Fields.FieldVector)
+    for pn in propertynames(state)
+        state_new = getproperty(state, pn)
+        @info "Checking NaNs in $pn"
+        count_nans_state(state_new)
+    end
+    return nothing
+end
+
+function count_nans_state(state::ClimaCore.Fields.Field)
+    num_nans = count(isnan.(Array(parent(state))))
+    if num_nans > 0
+        @warn "$num_nans NaNs found"
+    else
+        @info "No NaNs found"
+    end
+    return nothing
+end

test/shared_utilities/utilities.jl

@@ -256,3 +256,47 @@ 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_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)
+
+    # Add some NaNs to the fields
+    parent(var1)[1] = NaN
+    parent(var2)[1] = NaN
+    parent(var2)[2] = NaN
+
+    # Count and log the number of NaNs in the state
+    @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)
+
+end