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