Symbolic differentiation of sparse function gives incorrect jacobian

Hi all,

I am trying to differentiate a vector function with respect to a subset of the variables symbolically, but I am running into problems. Below is the problematic code:

using Symbolics, SparseArrays, LinearAlgebra, ForwardDiff

# setup data
θ = rand(3)
x = rand(11)
λ = rand(24)
S = sprand(Float64, 6, 7, 0.1)
Cp = sprand(Float64, 1, 11, 0.4)
_h = rand(24)
_d = spzeros(10)
λ = zeros(24)
θ = [30.0, 70.0, 90.0]
x = [1.9873817034700316, 3.9747634069400632, 1.9873817034700316, 1.9873817034700316, 1.9873817034700316, 3.9747634069400632, 1.9873817034700316, 0.028391167192429023, 0.022082018927444796, 0.13249211356466878, 0.13249211356466878]
λ[end] = -1.9873817034700316
S = [
    1.0  0.0  0.0  0.0  0.0  0.0  -1.0
    0.0  1.0  0.0  0.0  0.0  0.0  -2.0
    0.0  0.0  1.0  0.0  0.0  0.0  -1.0
    0.0  0.0  0.0  1.0  0.0  0.0  -1.0
    0.0  0.0  0.0  0.0  1.0  0.0  -1.0
    0.0  0.0  0.0  0.0  0.0  1.0  -2.0
]
se_vars = ((2, -1.0), (3, -1.0),(3, -3.0), (1, -2.0), (1, -1.0),)

# create functions
Se(θ) = sparse(
    [1, 2, 3, 4, 3],
    [3, 4, 4, 5, 5],
    [k / θ[idx] for (idx, k) in se_vars],
    4,
    7,
)

E(θ) = [
    S spzeros(6, 4)
    Se(θ) I(4)
]

d = _ -> _d

M = _ -> [
    -I(11)
    I(11)
    -Cp
    Cp
]

h = _ -> _h

# make symbolic variables
Symbolics.@variables begin
    sx[1:11]
    sλ[1:24]
    sθ[1:3]
end
 
# define function to differentiate
sparse_F(x, λ, θ) = [
    E(θ) * x - d(θ)
    λ .* (M(θ) * x - h(θ))
]

sj = Symbolics.jacobian(sparse_F(sx, sλ, sθ), collect(sθ)) # jacobian wrt θ only (collecting comes from the linked to other issue on discourse)
# this yields a matrix of all zeros

# check the answer with ForwardDiff
dense_F(x, λ, θ) = [
    Array(E(θ)) * x - Array(d(θ))
    λ .* (Array(M(θ)) * x - Array(h(θ)))
]

j = ForwardDiff.jacobian(θ -> dense_F(x, λ, θ), θ) # this is different :( 

Weirdly, if I don’t write the function out nicely (as above), but group all the variables together, symbolics correctly calculates the Jacobian. From here I can just extract the part I care about. However, for the application I am working on this is a huge matrix., and evaluating the Jacobian this way causes Julia to crash… Here is the code that works, but calculates all the derivatives and crashes for big systems:

xidxs = 1:11
λidxs = last(xidxs) .+ (1:24)
θidxs = last(λidxs) .+ (1:3)

sparse_F2(z) = [
    E(z[θidxs]) * z[xidxs] - d(z[θidxs])
    spdiagm(z[λidxs]) * (M(z[θidxs]) * z[xidxs] - h(z[θidxs]))
]
sz = [sx; sλ; sθ]
sparse_F2(sz)

sj = Symbolics.jacobian(sparse_F2(sz), sz)[:, end-3:end] # need to separate it out 

Does anybody have any ideas on how to make the first code snippet work? I think this may also be related to this problem.

Open an issue so @shashi sees it.

1 Like

Will do, thanks! The issue is here for future reference.

This was fixed (very quickly) in this issue. Thanks to all for the help!

2 Likes