add verbose option

Author: juliasloan25

src/shared_utilities/utils.jl

@@ -494,7 +494,7 @@ function isdivisible(
 end
 
 """
-    count_nans_state(state, mask::ClimaCore.Fields.Field = nothing)
+    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
@@ -508,24 +508,35 @@ 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)
+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)
+        count_nans_state(state_new; mask, verbose)
     end
     return nothing
 end
 
-function count_nans_state(state::ClimaCore.Fields.Field, mask = nothing)
+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
-        @info "No NaNs found"
+        verbose && @info "No NaNs found"
     end
     return nothing
 end

test/shared_utilities/utilities.jl

@@ -275,14 +275,21 @@ end
     # 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
+    # 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)
+    ) (: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())
@@ -294,14 +301,22 @@ end
     var2_copy[2] = NaN
     parent(var2) .= ArrayType(var2_copy)
 
-    # Count and log the number of NaNs in the state
+    # 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",
-    ) (:info, "No NaNs found") ClimaLand.count_nans_state(Y)
+    ) ClimaLand.count_nans_state(Y)
 
     # Test with a mask
     mask_zeros = Fields.zeros(space)
@@ -311,7 +326,18 @@ end
     ) (:info, "Checking NaNs in var2") (:info, "No NaNs found") (
         :info,
         "Checking NaNs in var3",
-    ) (:info, "No NaNs found") ClimaLand.count_nans_state(Y, mask_zeros)
+    ) (: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") (
@@ -320,5 +346,17 @@ end
     ) (:info, "Checking NaNs in var2") (:warn, "2 NaNs found") (
         :info,
         "Checking NaNs in var3",
-    ) (:info, "No NaNs found") ClimaLand.count_nans_state(Y, mask_ones)
+    ) (: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