diff --git a/Project.toml b/Project.toml index f415eed..9ecc2ee 100644 --- a/Project.toml +++ b/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" diff --git a/docs/src/jvp.md b/docs/src/jvp.md index 301f876..07cf18e 100644 --- a/docs/src/jvp.md +++ b/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: diff --git a/src/jvp.jl b/src/jvp.jl index 3db5eb9..d194bea 100644 --- a/src/jvp.jl +++ b/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) diff --git a/test/jvp_accuracy_tests.jl b/test/jvp_accuracy_tests.jl new file mode 100644 index 0000000..bf42c93 --- /dev/null +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index 1b3f5b9..4e9619d 100644 --- a/test/runtests.jl +++ b/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