Fix step size calculation if some results are non-deterministic (inf derivative) or nan

Author: wsmoses

src/methods.jl

@@ -371,11 +371,20 @@ function estimate_step(
     return _limit_step(m, x, step, acc)
 end
 
+function finite_or_zero(fs::AbstractArray{<:Number})
+   ifelse.(isfinite.(fs), fs, zero(fs))
+end
+
+function finite_or_zero(fs::AbstractArray{<:AbstractArray})
+   finite_or_zero.(fs)
+end
+
 function _estimate_magnitudes(
     m::FiniteDifferenceMethod{P,Q}, f::TF, x::T,
 ) where {P,Q,TF,T<:AbstractFloat}
     step = first(estimate_step(m, f, x))
     fs = _eval_function(m, f, x, step)
+    fs = finite_or_zero(fs)
     # Estimate magnitude of `∇f` in a neighbourhood of `x`.
     ∇fs = SVector{3}(
         _compute_estimate(m, fs, x, step, m.coefs_neighbourhood[1]),

test/grad.jl

@@ -217,3 +217,35 @@ using FiniteDifferences: grad, jacobian, _jvp, jvp, j′vp, _j′vp, to_vec
         @test [real(ȳ), imag(ȳ)] ≈ Jy'z̄_vec
     end
 end
+
+using LinearAlgebra
+
+function partial_nan_returning(x)
+    y = Matrix{Float64}(undef, 5, 5)
+    y .= NaN
+    y = Hermitian(y)
+    y .= x
+    return parent(y)
+end
+
+randvar = 1
+function partial_nondet_returning(x)
+    global randvar
+    y = Matrix{Float64}(undef, 5, 5)
+    y .= randvar
+    randvar += 1
+    y = Hermitian(y)
+    y .= x
+    return parent(y)
+end
+
+@testset "jvp: Estimate step correctly for when some terms are nan/infinite" begin
+    fdm = FiniteDifferences.central_fdm(5, 1)
+    res = jvp(fdm, partial_nan_returning, 3.1, 2.7)
+    @show res
+    @test Hermitian(res) .≈ 2.7
+    
+    res = jvp(fdm, partial_nondet_returning, 3.1, 2.7)
+    @show res
+    @test Hermitian(res) .≈ 2.7
+end