Scale the JVP finite difference step by norm(x)/norm(v)

Project.toml

@@ -1,6 +1,6 @@
 name = "FiniteDiff"
 uuid = "6a86dc24-6348-571c-b903-95158fe2bd41"
-version = "2.32.1"
+version = "2.33.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

docs/src/jvp.md

@@ -11,6 +11,21 @@ The JVP computes `J(x) * v` where `J(x)` is the Jacobian of function `f` at poin
 
 where `h` is the step size and `v` is the direction vector.
 
+## Step Size
+
+Because `f` is evaluated at `x + h*v`, the step `h` has units of `[x]/[v]`: it is not a
+step in `x` and it is not dimensionless. The step used is
+
+```julia
+h = max(relstep * norm(x), absstep) * dir / norm(v)
+```
+
+so that the perturbation `h*v` is a `relstep` relative change of `x`, floored at
+`absstep`, and the computed JVP is invariant (up to that floor) to rescaling of `v`.
+This matters for matrix-free Krylov solvers, where `v` is a unit-norm basis vector and
+`norm(x)` may be large: a step that ignores `norm(x)` is stuck at the `absstep` floor
+and the JVP loses accuracy proportionally to `norm(x)`.
+
 ## Performance Benefits
 
 JVP functions are particularly efficient when you only need directional derivatives:

src/jvp.jl

@@ -79,6 +79,31 @@ function JVPCache(
     JVPCache{typeof(x), typeof(fx), fdtype}(x, fx)
 end
 
+"""
+    jvp_epsilon(fdtype, x, v, relstep::Real, absstep::Real, dir::Real)
+
+Compute the finite difference step size used by `finite_difference_jvp` and
+`finite_difference_jvp!`.
+
+The JVP evaluates `f(x + ϵ*v)`, so `ϵ` must carry units of `[x]/[v]` in order for
+`ϵ*v` to be a perturbation of `x`. The step is therefore the ordinary scalar step
+applied to the magnitude of `x`, rescaled by the magnitude of `v`:
+
+    ϵ = max(relstep*norm(x), absstep) * dir / norm(v)
+
+so that `norm(ϵ*v) == max(relstep*norm(x), absstep)`, i.e. the perturbation is a
+`relstep` relative change of `x` (floored at `absstep`), independent of how `v` is
+scaled. `norm(v) == 0` and non-finite `norm(x)` fall back to the unscaled step
+rather than producing `Inf`/`NaN` steps.
+"""
+@inline function jvp_epsilon(fdtype, x, v, relstep::Real, absstep::Real, dir::Real)
+    nx = norm(_vec(x))
+    nv = norm(_vec(v))
+    isfinite(nx) || (nx = zero(nx))
+    epsilon = compute_epsilon(fdtype, nx, relstep, absstep, dir)
+    (iszero(nv) || !isfinite(nv)) ? epsilon : epsilon / nv
+end
+
 """
     FiniteDiff.finite_difference_jvp(
         f,
@@ -106,6 +131,10 @@ computing the full Jacobian when only `J*v` is needed.
 - `relstep`: Relative step size (default: method-dependent optimal value)
 - `absstep=relstep`: Absolute step size fallback
 
+The step is `ϵ = max(relstep*norm(x), absstep)*dir/norm(v)`, i.e. `relstep` and
+`absstep` set the size of the perturbation `ϵ*v` relative to `x`, and the result is
+unchanged (up to the `absstep` floor) if `v` is rescaled.
+
 # Returns
 - Vector `J(x) * v` representing the Jacobian-vector product
 
@@ -169,8 +198,7 @@ function finite_difference_jvp(
         ArgumentError("finite_difference_jvp doesn't support :complex-mode finite diff")
     end
 
-    tmp = sqrt(abs(dot(_vec(x), _vec(v))))
-    epsilon = compute_epsilon(fdtype, tmp, relstep, absstep, dir)
+    epsilon = jvp_epsilon(fdtype, x, v, relstep, absstep, dir)
     if fdtype == Val(:forward)
         fx = f_in isa Nothing ? f(x) : f_in
         x1 = @. x + epsilon * v
@@ -250,8 +278,7 @@ function finite_difference_jvp!(
     end
 
     (; x1, fx1) = cache
-    tmp = sqrt(abs(dot(_vec(x), _vec(v))))
-    epsilon = compute_epsilon(fdtype, tmp, relstep, absstep, dir)
+    epsilon = jvp_epsilon(fdtype, x, v, relstep, absstep, dir)
     if fdtype == Val(:forward)
         if f_in isa Nothing
             f(fx1, x)

test/jvp_accuracy_tests.jl

@@ -0,0 +1,70 @@
+using FiniteDiff, LinearAlgebra, Test
+
+# f_i(x) = x_i * x_{i+1} (cyclic): smooth, varies on the scale of x, and has an
+# exactly known Jacobian, so the reference J*v carries no discretization error.
+const N_JVP = 200
+nxt(i) = i == N_JVP ? 1 : i + 1
+fcyc(x) = [x[i] * x[nxt(i)] for i in eachindex(x)]
+fcyc!(y, x) = (for i in eachindex(x)
+        y[i] = x[i] * x[nxt(i)]
+    end; y)
+Jv_exact(x, v) = [x[nxt(i)] * v[i] + x[i] * v[nxt(i)] for i in eachindex(x)]
+
+relerr(a, b) = norm(a - b) / norm(b)
+
+x_big = [100.0 * (1 + i / N_JVP) for i in 1:N_JVP]
+v_unit = [sin(float(i)) for i in 1:N_JVP]
+v_unit ./= norm(v_unit)
+ref_big = Jv_exact(x_big, v_unit)
+
+@testset "JVP step size and accuracy" begin
+    @testset "JVP step size scales with norm(x)" begin
+        @test relerr(FiniteDiff.finite_difference_jvp(fcyc, copy(x_big), v_unit), ref_big) < 1e-7
+        jvp = similar(x_big)
+        FiniteDiff.finite_difference_jvp!(jvp, fcyc!, copy(x_big), v_unit)
+        @test relerr(jvp, ref_big) < 1e-7
+
+        @test relerr(FiniteDiff.finite_difference_jvp(fcyc, copy(x_big), v_unit, Val{:central}),
+            ref_big) < 1e-10
+        FiniteDiff.finite_difference_jvp!(jvp, fcyc!, copy(x_big), v_unit, Val{:central})
+        @test relerr(jvp, ref_big) < 1e-10
+    end
+
+    @testset "JVP is invariant to the scaling of v" begin
+        base = FiniteDiff.finite_difference_jvp(fcyc, copy(x_big), v_unit)
+        for s in (1.0e-4, 1.0e2, 1.0e4)
+            scaled = FiniteDiff.finite_difference_jvp(fcyc, copy(x_big), s .* v_unit) ./ s
+            @test relerr(scaled, base) < 1e-8
+            @test relerr(scaled, ref_big) < 1e-7
+        end
+    end
+
+    @testset "JVP degenerate directions and states" begin
+        zerojvp = FiniteDiff.finite_difference_jvp(fcyc, copy(x_big), zero(v_unit))
+        @test all(isfinite, zerojvp)
+        @test iszero(zerojvp)
+
+        # norm(x) == 0 must fall back to the absstep floor rather than a zero step
+        zerox = zeros(N_JVP)
+        atzero = FiniteDiff.finite_difference_jvp(fcyc, zerox, v_unit)
+        @test all(isfinite, atzero)
+        @test norm(atzero) < 1e-6
+
+        # a non-finite entry of x must not contaminate the step, and hence every output
+        nanx = copy(x_big)
+        nanx[1] = NaN
+        nanjvp = FiniteDiff.finite_difference_jvp(fcyc, nanx, v_unit)
+        @test count(isnan, nanjvp) == 2
+    end
+
+    @testset "JVP complex-valued state" begin
+        xc = [(1 + im) * 100.0 * (1 + i / N_JVP) for i in 1:N_JVP]
+        vc = v_unit .+ im .* reverse(v_unit)
+        vc ./= norm(vc)
+        refc = Jv_exact(xc, vc)
+        @test relerr(FiniteDiff.finite_difference_jvp(fcyc, copy(xc), vc), refc) < 1e-7
+        jvpc = similar(xc)
+        FiniteDiff.finite_difference_jvp!(jvpc, fcyc!, copy(xc), vc)
+        @test relerr(jvpc, refc) < 1e-7
+    end
+end

test/runtests.jl

@@ -17,6 +17,7 @@ if GROUP == "All" || GROUP == "Core"
   @time @safetestset "Public API Tests" begin include("public_api_tests.jl") end
   @time @safetestset "FiniteDiff Standard Tests" begin include("finitedifftests.jl") end
   @time @safetestset "Color Differentiation Tests" begin include("coloring_tests.jl") end
+  @time @safetestset "JVP Accuracy Tests" begin include("jvp_accuracy_tests.jl") end
   @time @safetestset "Out of Place Tests" begin include("out_of_place_tests.jl") end
   @time @safetestset "Cache Reuse Safety Tests" begin include("cache_reuse_tests.jl") end
 end