Merge pull request #2139 from CliMA/ck/fix_missing_debug_call

Add missing `post_op_callback` calls

Author: charleskawczynski

NEWS.md

@@ -4,6 +4,14 @@ ClimaCore.jl Release Notes
 main
 -------
 
+### ![][badge-✨feature/enhancement] A new DebugOnly module can help find where
+     NaNs/Inf come from. PRs [2115](https://github.com/CliMA/ClimaCore.jl/pull/2115) and
+     [2139](https://github.com/CliMA/ClimaCore.jl/pull/2139)
+
+A new `ClimaCore.DebugOnly` module was added, which can help users find where
+NaNs or Infs come from in a simulation, interactively. Documentation, with a
+simple representative example can be found [here](https://clima.github.io/ClimaCore.jl/dev/debugging/#Infiltrating).
+
 ### ![][badge-✨feature/enhancement] Various improvements to `Remapper` [2060](https://github.com/CliMA/ClimaCore.jl/pull/2060)
 
 The `ClimaCore.Remapping` module received two improvements. First, `Remapper` is

docs/src/debugging.md

@@ -78,15 +78,21 @@ import ClimaCore
 import Infiltrator # must be in your default environment
 ClimaCore.DebugOnly.call_post_op_callback() = true
 function ClimaCore.DebugOnly.post_op_callback(result, args...; kwargs...)
-    if any(isnan, parent(result))
-        println("NaNs found!")
+    has_nans = any(isnan, parent(result))
+    has_inf = any(isinf, parent(result))
+    if has_nans || has_inf
+        has_nans && println("NaNs found!")
+        has_inf && println("Infs found!")
         # Let's define the stack trace so that we know where this came from
         st = stacktrace()
 
         # Let's use Infiltrator.jl to exfiltrate to drop into the REPL.
         # Now, `Infiltrator.safehouse` will be a NamedTuple
-        # containing `result`, `args` and `kwargs`.
+        # containing `result`, `args`, `kwargs` and `st`.
         Infiltrator.@exfiltrate
+        # sometimes code execution doesn't stop, I'm not sure why. Let's
+        # make sure we exfiltrate immediately with the data we want.
+        error("Exfiltrating.")
     end
 end
 
@@ -97,7 +103,7 @@ data = ClimaCore.DataLayouts.VIJFH{FT}(Array{FT}, zeros; Nv=5, Nij=2, Nh=2)
 (;result, args, kwargs, st) = Infiltrator.safehouse;
 
 # You can print the stack trace, to see where the NaNs were found:
-ClimaCore.DebugOnly.print_depth_limited_stack_trace(st;maxtypedepth=1)
+ClimaCore.DebugOnly.print_depth_limited_stack_trace(st; maxtypedepth=1)
 
 # Once there, you can see that the call lead you to `copyto!`,
 # Inspecting `args` shows that the `Broadcasted` object used to populate the
@@ -122,3 +128,10 @@ Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0}}(identity, (NaN,)
     `post_op_callback` is called in many places, so this is a
     performance-critical code path and expensive operations performed in
     `post_op_callback` may significantly slow down your code.
+
+!!! warn
+
+    It is _highly_ recommended to use `post_op_callback` _without_ `@testset`,
+    as Test.jl may continue running through code execution, until all of the
+    tests in a given `@testset` are complete, and the result will be that you
+    will get the _last_ observed instance of `NaN` or `Inf`.

src/CommonSpaces/CommonSpaces.jl

@@ -91,7 +91,7 @@ space = ExtrudedCubedSphereSpace(;
     z_max = 1,
     radius = 10,
     h_elem = 10,
-    n_quad_points = 4
+    n_quad_points = 4,
     staggering = CellCenter()
 )
 ```

src/Fields/fieldvector.jl

@@ -325,10 +325,14 @@ end
     parent(getfield(_values(fv), symb))
 @inline transform_broadcasted(x, symb, axes) = x
 
-@inline Base.copyto!(
+@inline function Base.copyto!(
     dest::FieldVector,
     bc::Union{FieldVector, Base.Broadcast.Broadcasted{FieldVectorStyle}},
-) = copyto_per_field!(dest, bc)
+)
+    copyto_per_field!(dest, bc)
+    call_post_op_callback() && post_op_callback(dest, dest, bc)
+    return dest
+end
 
 @inline function copyto_per_field!(
     dest::FieldVector,
@@ -351,18 +355,29 @@ end
     return dest
 end
 
-@inline Base.copyto!(
+@inline function Base.copyto!(
     dest::FieldVector,
     bc::Base.Broadcast.Broadcasted{<:Base.Broadcast.Style{Tuple}},
-) = copyto_per_field_scalar!(dest, bc)
+)
+    copyto_per_field_scalar!(dest, bc)
+    call_post_op_callback() && post_op_callback(dest, dest, bc)
+    return dest
+end
 
-@inline Base.copyto!(
+@inline function Base.copyto!(
     dest::FieldVector,
     bc::Base.Broadcast.Broadcasted{<:Base.Broadcast.AbstractArrayStyle{0}},
-) = copyto_per_field_scalar!(dest, bc)
+)
+    copyto_per_field_scalar!(dest, bc)
+    call_post_op_callback() && post_op_callback(dest, dest, bc)
+    return dest
+end
 
-@inline Base.copyto!(dest::FieldVector, bc::Real) =
+@inline function Base.copyto!(dest::FieldVector, bc::Real)
     copyto_per_field_scalar!(dest, bc)
+    call_post_op_callback() && post_op_callback(dest, dest, bc)
+    return dest
+end
 
 @inline function copyto_per_field_scalar!(dest::FieldVector, bc)
     map(propertynames(dest)) do symb