fixed scalarIndexing performance issue

lib/ClimaLandSimulations/src/Fluxnet/fluxnet_utilities/make_drivers.jl

@@ -87,8 +87,10 @@ function make_drivers(site_ID, setup, config, params, context)
         end
     end
     if length(missing_drivers) != 0
-        error("Driver data missing for columns: $([missing_drivers[i] * " " for
-            i in 1:length(missing_drivers)]...)")
+        error(
+            "Driver data missing for columns: $([missing_drivers[i] * " " for
+          i in 1:length(missing_drivers)]...)",
+        )
     end
 
     thermo_params = LP.thermodynamic_parameters(earth_param_set)

src/shared_utilities/utils.jl

@@ -593,22 +593,23 @@ function count_nans_state(
     mask = nothing,
     verbose = false,
 )
-    # Note: this code uses `parent`; this pattern should not be replicated
-    num_nans = 0
-    ClimaComms.allowscalar(ClimaComms.device()) do
-        num_nans =
-            isnothing(mask) ? Int(sum(isnan.(parent(state)))) :
-            Int(sum(isnan.(parent(state)) .* parent(mask)))
+    if isnothing(mask)
+        num_nans = count(isnan, parent(state))
+    else
+        num_nans = mapreduce(
+            (s, m) -> m != 0 && isnan(s),
+            Base.add_sum,
+            parent(state),
+            parent(mask),
+        )
     end
-    if isapprox(num_nans, 0)
+    if num_nans == 0
         verbose && @info "No NaNs found"
     else
         @warn "$num_nans NaNs found"
     end
     return nothing
 end
-
-
 """
     NaNCheckCallback(nancheck_frequency::Union{AbstractFloat, Dates.Period},
                         start_date, dt)