Colored Jacobian uses one step per color, inaccurate for variables of different magnitudes #241

Goal

We compute Jacobians of the dynamics of nonlinear models by finite differences, for use in nonlinear model-predictive control. The Jacobians are sparse, so we would like to evaluate them with a column coloring. The models are badly scaled in the sense that the variables differ in magnitude by several orders, for example a pressure of order 1e5 Pa next to a variable of order one or smaller.

Observation

With colorvec and sparsity, finite_difference_jacobian! perturbs all columns of a color by one common step, computed in src/jacobians.jl as

@. x2 = x1 * (_color == color_i)
tmp = norm(x2)
epsilon = compute_epsilon(Val(:forward), sqrt(tmp), relstep, absstep, dir)

(lines 567-569 for forward differences, 608-610 for central differences, and the corresponding lines 344 and 385 of the out-of-place method). The step is therefore determined by the square root of the norm of the variables of the color. When a color contains variables of very different magnitudes, this step differs by orders of magnitude from the step that the dense Jacobian takes for an individual column (compute_epsilon(fdtype, x[i], relstep, absstep)), and the colored Jacobian differs from the dense one in the leading digits.

Example

The Jacobian is diagonal, so both columns are given one color.

using FiniteDiff, SparseArrays, LinearAlgebra

f!(y, x) = (y[1] = x[1]^3; y[2] = x[2]^3; y)   # diagonal Jacobian, both columns in one color
Jexact(x) = Diagonal(3 .* x .^ 2)

for x in ([1e5, 1e-2], [1e10, 1.0]), fdtype in (Val(:forward), Val(:central))
    Jd = zeros(2, 2)
    FiniteDiff.finite_difference_jacobian!(Jd, f!, x, fdtype)
    Js = sparse([1, 2], [1, 2], [1.0, 1.0])
    FiniteDiff.finite_difference_jacobian!(Js, f!, x, fdtype; colorvec = [1, 1], sparsity = Js)
    println("x = $x, $fdtype")
    println("  dense   relative error: ", diag(Jd) ./ diag(Jexact(x)) .- 1)
    println("  colored relative error: ", diag(Js) ./ diag(Jexact(x)) .- 1)
end

Output:

x = [100000.0, 0.01], Val{:forward}()
  dense   relative error: [1.495313073363036e-8, 1.4900886828872473e-6]
  colored relative error: sparsevec([1, 2], [3.362815435004052e-7, 0.0004712901064343189], 2)
x = [100000.0, 0.01], Val{:central}()
  dense   relative error: [1.7613910330283034e-11, 1.2222844958387213e-7]
  colored relative error: sparsevec([1, 2], [1.985967612583295e-9, 0.012222842875003748], 2)
x = [1.0e10, 1.0], Val{:forward}()
  dense   relative error: [1.5047466161988154e-8, 1.4901161193847656e-8]
  colored relative error: sparsevec([1, 2], [-0.000432427792591783, 0.0014908562680582271], 2)
x = [1.0e10, 1.0], Val{:central}()
  dense   relative error: [-1.7607026947530358e-12, 8.275380380950992e-12]
  colored relative error: sparsevec([1, 2], [-5.824787592256797e-7, 0.12222842875003481], 2)

The dense Jacobian is accurate to the expected order in every case; the colored Jacobian has a relative error of about 1e-2 to 1e-1 in the entries of the small variable when central differences are used.

FiniteDiff.finite_difference_jvp!, which DifferentiationInterface uses for AutoSparse(AutoFiniteDiff()), behaves similarly, since jvp_epsilon (src/jvp.jl#L99-L105) takes the step relative to norm(x) of the whole point. I will report the DifferentiationInterface side separately.

What we tried

• relstep and absstep: these set one scalar step for all columns of a color, so no choice is accurate for both a variable of order 1e5 and a variable of order 1e-2.
• Forward and central differences: both are affected; central differences more so.
• As a workaround we wrote our own colored finite differences, in which every column of a color is perturbed by the step that the dense Jacobian takes for that column, and each entry is divided by the step of its column. The result equals the dense Jacobian up to rounding, at the cost of one function evaluation per color.

Question

Is there an intended way to obtain colored finite-difference Jacobians with a per-column step (the step of the dense Jacobian) from FiniteDiff, or is the common step per color a deliberate design choice? If the latter, is there a recommended approach for badly scaled problems other than rescaling the model variables?

Versions

Julia 1.12.7, FiniteDiff v2.33.0.

Host: demeter2 · Session: c59a36f8-c4c7-4a84-9858-91adc076743c · Working directory: ~/.julia/dev/MPCComponents

🤖 Generated with Claude Code