fix GPU compat

Author: juliasloan25

src/shared_utilities/utils.jl

@@ -494,28 +494,34 @@ function isdivisible(
 end
 
 """
-    count_nans_state(sol)
+    count_nans_state(state, mask::ClimaCore.Fields.Field = nothing)
 
-Count the number of NaNs in the state variables. This function is useful for
+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.
 
-Input: `state` - e.g. the FieldVector given by `sol.u[end]` after calling `solve`
+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.
 """
-function count_nans_state(state::ClimaCore.Fields.FieldVector)
+function count_nans_state(state::ClimaCore.Fields.FieldVector, mask = nothing)
     for pn in propertynames(state)
         state_new = getproperty(state, pn)
         @info "Checking NaNs in $pn"
-        count_nans_state(state_new)
+        count_nans_state(state_new, mask)
     end
     return nothing
 end
 
-function count_nans_state(state::ClimaCore.Fields.Field)
-    num_nans = count(isnan.(Array(parent(state))))
+function count_nans_state(state::ClimaCore.Fields.Field, mask = nothing)
+    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

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
 
@@ -285,10 +284,15 @@ end
         "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
+    # 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_logs (:info, "Checking NaNs in var1") (:warn, "1 NaNs found") (
@@ -299,4 +303,22 @@ end
         "Checking NaNs in var3",
     ) (:info, "No NaNs found") 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_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_ones)
 end