How to use ForwardDiff.jl with inv and matrix exponential?

General Usage
, , ,
1

I have a question listed here: Why does ForwardDiff.jl give me all-zero Jacobian matrix?, and I found it might be due to part of the function which contains inv and exp (matrix exponential…) from ExponentialUtilities. I am not sure whether ForwardDiff.jl could fail with inv of matrix or exp?

one example of my function could see below:

# ------------------------------------------------------------------
# solve_odes_water: matrix-exponential solution of dW/dt = A × W + b.
#
# Analytical solution over one timestep ΔT:
#   W(t+ΔT) = e^(A·ΔT) · (W₀ + A⁻¹·b)  −  A⁻¹·b
#
# Derivation: substituting u = W + A⁻¹b → du/dt = A·u → u(t) = e^(At)·u₀
#   → W(t+ΔT) = e^(A·ΔT)·(W₀ + A⁻¹·b) − A⁻¹·b
#
# Falls back to W(t+ΔT) = W₀ when A is singular (isSingular_water = true),
# which avoids NaN propagation at very low carbon / dry spinup.
#
# Mathematically identical to solve_odes_3_water in original.
# ------------------------------------------------------------------
function solve_odes_water(new_vegW, deriv_mat, deriv_const, init_cond, deltaT, z_zero, helpers)
    if isSingular_water(deriv_mat, z_zero)
        for ix ∈ eachindex(new_vegW)
            @rep_elem init_cond[ix] ⇒ (new_vegW, ix, :vegW)
        end
    else
        inv_deriv = inv(deriv_mat)
        new_vegW  = real(-inv_deriv * deriv_const +
                         exp(deltaT * deriv_mat) * (init_cond + inv_deriv * deriv_const))
    end
    return new_vegW
end

Btw, in my code/models, I am using static arrays and static matrix to reduce allocations, and at some places I forced to use Float32. In the final cost function, I set it as dual type. I guess this is not the problem? because I run same codes with different parameter values and input time series data (parameter values coming from optimization over different sites using different input data) the model over 40 independent sites (40 independent runs), only 5-8 sites return zero-valued Jacobians…others are fine…

2

It’s easy to try — if you set A = rand(3,3), you can easily see that ForwardDiff.jacobian(inv, A) succeeds, and indeed matches the analytical result (-A^T \otimes A):

julia> ForwardDiff.jacobian(inv, A) ≈ -kron(inv(A)', inv(A))
true

but with the matrix exponential it throws a MethodError:

julia> ForwardDiff.jacobian(exp, A)
ERROR: MethodError: no method matching exp!(::Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(exp), Float64}, Float64, 9}})

and indeed there is an open issue for this: support for `Base.expm` (need advice implementing) · Issue #174 · JuliaDiff/ForwardDiff.jl · GitHub

(ChainRules.jl has rules for the matrix exponential, however, so if you use something like Zygote.jl or Enzyme.jl that supports ChainRules it should be easier to get it working. Or use a different implementation of exp like in ExponentialUtilities.jl)