Adding a custom NaN checker that report also where the NaNs were found

[navidcy]

Often when you run big simulation and get NaN it's useful if you get a sense of where the NaN first started. I have in mind eg of a global ocean simulation with realistic bathymetry. In this case whether the NaN originate from the surface or the interior, or whether it came from shallow ocean regions or not it might give an indication for what the issue is. Even if many NaNs are found, still reporting where they came from might be useful. E.g.

error("time = $t, iteration = $iter: NaN found in field $name. Aborting simulation.")

@info "NaN were found in in field $name at indices:
for nan_indices in nan_indiceses
    i, j, k = nan_indices
    @info "i, j, k = $i, $j, $k
end

What's a good way to do that without adding too much overhead to the nan_checker?
I guess one should import and extend

Oceananigans.jl/src/Models/nan_checker.jl

Lines 22 to 52 in c3a5654

	"""
	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`.
	If `erroring=true`, the `NaNChecker` will throw an error on NaN detection.
	"""
	NaNChecker(; fields, erroring=false) = NaNChecker(fields, erroring)
	hasnan(field) = any(isnan, parent(field))
	hasnan(model::AbstractModel) = hasnan(first(fields(model)))
	function (nc::NaNChecker)(simulation)
	for (name, field) in pairs(nc.fields)
	if hasnan(field)
	simulation.running = false
	clock = simulation.model.clock
	t = time(simulation)
	iter = iteration(simulation)
	if nc.erroring
	error("time = $t, iteration = $iter: NaN found in field $name. Aborting simulation.")
	else
	@info "time = $t, iteration = $iter: NaN found in field $name. Stopping simulation."
	end
	end
	end
	return nothing
	end

right?

cc @taimoorsohail

[glwagner]

@navidcy check out this issue which proposes another feature for debuggint that's similar: #4181. The idea is to add something to arch. Either a wrapper or property in the architecture which, if activated, adds a check for nans after executing a kernel within launch!. The tricky part is figuring out where to look for NaNs. You could start by checking the args. Or, using an interface. But the nice thing is that it could report the kernel that produced the NaNs and I suppose also the index.

[navidcy]

oh I didn't notice that! great!

[glwagner]

One can also use the NaNChecker to do that. I think you want to use the erroring version then.

[glwagner]

Something like this might work

function find_nans(sim)
    nans_in_u = isnan.(interior(sim.model.velocities.u))
    @show findall(nans_in_u)
end