(0.96.28) Small improvements to NaNChecker

Project.toml

@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.96.27"
+version = "0.96.28"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/Models/Models.jl

@@ -176,7 +176,7 @@ function default_nan_checker(model::OceananigansModels)
 
     first_name = first(keys(model_fields))
     field_to_check_nans = NamedTuple{tuple(first_name)}(model_fields)
-    nan_checker = NaNChecker(field_to_check_nans)
+    nan_checker = NaNChecker(; fields=field_to_check_nans)
     return nan_checker
 end
 

src/Models/nan_checker.jl

@@ -1,11 +1,12 @@
+using CUDA: @allowscalar
 using Oceananigans.Utils: prettykeys
 
-mutable struct NaNChecker{F}
+mutable struct NaNChecker{F, H}
     fields :: F
+    hasnan :: H
     erroring :: Bool
 end
 
-NaNChecker(fields) = NaNChecker(fields, false) # default
 default_nan_checker(model) = nothing
 
 function Base.summary(nc::NaNChecker)
@@ -23,19 +24,30 @@ Base.show(io, nc::NaNChecker) = print(io, summary(nc))
     NaNChecker(; fields, erroring=false)
 
 Return a `NaNChecker`, which sets `sim.running=false` if a `NaN` is detected
-in any member of `fields` when `NaNChecker(sim)` is called. `fields` should be
-a container with key-value pairs like a dictionary or `NamedTuple`.
+in any member of `fields` when `(::NaNChecker)(sim)` is called. `fields` should be
+a container with key-value pairs like a dictionary or `NamedTuple`. 
+`(::NaNChecker)(sim)` also returns a boolean indicating whether NaNs were found.
 
 If `erroring=true`, the `NaNChecker` will throw an error on NaN detection.
 """
-NaNChecker(; fields, erroring=false) = NaNChecker(fields, erroring)
+function NaNChecker(; fields, erroring=false)
+    first_field = first(fields)
+    hasnan = Field{Nothing, Nothing, Nothing}(first_field.grid, Bool)
+    return NaNChecker(fields, hasnan, erroring)
+end
 
-hasnan(field) = any(isnan, parent(field))
-hasnan(model::AbstractModel) = hasnan(first(fields(model)))
+function hasnan(field, checker)
+    any!(isnan, checker.hasnan, field)
+    return @allowscalar first(checker.hasnan)
+end
+
+hasnan(model::AbstractModel, checker) = hasnan(first(fields(model)), checker)
 
 function (nc::NaNChecker)(simulation)
+    found_nan = false
     for (name, field) in pairs(nc.fields)
-        if hasnan(field)
+        found_nan = hasnan(field, nc)
+        if found_nan
             simulation.running = false
             clock = simulation.model.clock
             t = time(simulation)
@@ -48,7 +60,7 @@ function (nc::NaNChecker)(simulation)
             end
         end
     end
-    return nothing
+    return found_nan
 end
 
 """