Key not found error

[baggepinnen]

The following example generates a ERROR: KeyError: key 228 not found. The example forms the dynamics of a simple mechanical system and integrates it a few steps using a hand-written RK4 integrator. The failing call occurs when calling sparse_hessian

import FastDifferentiation as fad
using StaticArrays

function cartpole(x, u, p=0, t=0)
    mc, mp, l, g = 1.0, 0.2, 0.5, 9.81

    q  = x[SA[1, 2]]
    qd = x[SA[3, 4]]

    s = sin(q[2])
    c = cos(q[2])

    H = [mc+mp mp*l*c; mp*l*c mp*l^2]
    C = [0 -mp*qd[2]*l*s; 0 0]
    G = [0, mp * g * l * s]
    B = [1, 0]
    # qdd = (-H) \ (C * qd + G - B * u[1]) # Backsolve errors #20 
    den = (H[1, 1]*H[2, 2] - H[1, 2]*H[2, 1])
    xdd = [H[2,2] -H[1,2]; -H[2,1] H[1,1]]
    qdd = xdd * (C * qd + G - B * u[1])
    return [qd; qdd]
end

function rk4(f::F, Ts0; supersample::Integer = 1) where {F}
    supersample ≥ 1 || throw(ArgumentError("supersample must be positive."))
    # Runge-Kutta 4 method
    Ts = Ts0 / supersample # to preserve type stability in case Ts0 is an integer
    let Ts = Ts
        function (x, u, p=0, t=0)
            T = typeof(x)
            f1 = f(x, u, p, t)
            f2 = f(x + Ts / 2 * f1, u, p, t + Ts / 2)
            f3 = f(x + Ts / 2 * f2, u, p, t + Ts / 2)
            f4 = f(x + Ts * f3, u, p, t + Ts)
            add = Ts / 6 * (f1 + 2 * f2 + 2 * f3 + f4)
            # This gymnastics with changing the name to y is to ensure type stability when x + add is not the same type as x. The compiler is smart enough to figure out the type of y
            y = x + add
            for i in 2:supersample
                f1 = f(y, u, p, t)
                f2 = f(y + Ts / 2 * f1, u, p, t + Ts / 2)
                f3 = f(y + Ts / 2 * f2, u, p, t + Ts / 2)
                f4 = f(y + Ts * f3, u, p, t + Ts)
                add = Ts / 6 * (f1 + 2 * f2 + 2 * f3 + f4)
                y += add
            end
            return y
        end
    end
end

function rollout(f, x0::AbstractVector, u)
    T = promote_type(eltype(x0), eltype(u))
    x = zeros(T, length(x0), size(u, 2))
    x[:, 1] .= x0
    rollout!(f, x, u)
end
function rollout!(f, x, u)
    for i = 1:size(x, 2)-1
        x[:, i+1] = f(x[:, i], u[:, i]) 
    end
    x, u
end

Ts = 0.01
N = 2 # Scale this number up to benchmark larger problems
u = reshape(fad.make_variables(:u, 2*N), 2, N)
x0 = fad.make_variables(:x0, 4)

discrete_cartpole = rk4(cartpole, Ts)
x, _ = rollout(discrete_cartpole, x0, u) # Never finishes with Symbolics

vars = [x0; vec(u)]
c = sum(abs2, x) + sum(abs2, u)

hs = fad.sparse_hessian(c, vars) # Errors
# h = fad.make_function(hs, vars)
# @btime h(randn(6))

julia> hs = fad.sparse_hessian(c, vars)
ERROR: KeyError: key 228 not found
Stacktrace:
  [1] getindex
    @ ./dict.jl:484 [inlined]
  [2] _evaluate_branching_subgraph(subgraph::FastDifferentiation.FactorableSubgraph{Int64, FastDifferentiation.DominatorSubgraph}, sum::FastDifferentiation.Node{Float64, 0}, current_vertex::Int64, sub_edges::Set{FastDifferentiation.PathEdge{Int64}}, counts::Dict{Int64, Int64}, vertex_sums::Dict{Int64, FastDifferentiation.Node})
    @ FastDifferentiation ~/.julia/packages/FastDifferentiation/vh3e2/src/Factoring.jl:468
  [3] _evaluate_branching_subgraph(subgraph::FastDifferentiation.FactorableSubgraph{Int64, FastDifferentiation.DominatorSubgraph}, sum::FastDifferentiation.Node{Int64, 0}, current_vertex::Int64, sub_edges::Set{FastDifferentiation.PathEdge{Int64}}, counts::Dict{Int64, Int64}, vertex_sums::Dict{Int64, FastDifferentiation.Node})
    @ FastDifferentiation ~/.julia/packages/FastDifferentiation/vh3e2/src/Factoring.jl:474
  [4] evaluate_branching_subgraph(subgraph::FastDifferentiation.FactorableSubgraph{Int64, FastDifferentiation.DominatorSubgraph})
    @ FastDifferentiation ~/.julia/packages/FastDifferentiation/vh3e2/src/Factoring.jl:454
  [5] factor_subgraph!(subgraph::FastDifferentiation.FactorableSubgraph{Int64, FastDifferentiation.DominatorSubgraph})
    @ FastDifferentiation ~/.julia/packages/FastDifferentiation/vh3e2/src/Factoring.jl:394
  [6] factor!(a::FastDifferentiation.DerivativeGraph{Int64})
    @ FastDifferentiation ~/.julia/packages/FastDifferentiation/vh3e2/src/Factoring.jl:497
  [7] _sparse_symbolic_jacobian!(graph::FastDifferentiation.DerivativeGraph{Int64}, partial_variables::Vector{FastDifferentiation.Node})
    @ FastDifferentiation ~/.julia/packages/FastDifferentiation/vh3e2/src/Jacobian.jl:86
  [8] sparse_jacobian
    @ ~/.julia/packages/FastDifferentiation/vh3e2/src/Jacobian.jl:110 [inlined]
  [9] sparse_hessian(expression::FastDifferentiation.Node{typeof(+), 2}, variable_order::Vector{FastDifferentiation.Node})
    @ FastDifferentiation ~/.julia/packages/FastDifferentiation/vh3e2/src/Jacobian.jl:258
 [10] top-level scope
    @ REPL[17]:1

[brianguenter]

This code doesn't run on my machine. SA is not defined. Can you post the rest of the code that defines SA?

[baggepinnen]

Sorry for that, I edited the original post, it was missing using StaticArrays

[brianguenter]

This appears to be an error in tagging reachable roots and variables in the graph factorization code. May take a while to track down.

[moble]

Out of curiosity, I ran the example above on the latest v0.3.14. It still errors, though the error is different:

AssertionError: Should only be one path from root 1 to variable 2. Instead have 3 children from node 122 on the path

I also narrowed it down to a simpler MWE by printing out the c expression that's causing the error, removing as much of it as possible (~98%) and setting as many variables to zero as possible while still getting an error. Also, I split the remaining expression into two ~equal parts, each of which differentiates just fine, but the sum of which results in an error.

import FastDifferentiation as fad

fad.@variables x02 x04 u1

vars = [x02, x04, u1]

c1 = -(cos((x02 + (x04 + ((-(cos(x02)) * (((-(x04) * sin(x02)) * x04) - u1)) + sin(x02))))))
c2 = ((-((x04 + (-(cos((x02 + x04))) * (((-((x04 + (-(cos(x02)) * (((-(x04) * sin(x02)) * x04) - u1)))) * sin((x02 + x04))) * x04) - u1)))) * sin(x02)) * x04)

fad.jacobian([c1], vars)  # This is fine
fad.jacobian([c2], vars)  # This is fine
fad.jacobian([c1+c2], vars)  # This errors

The error and stacktrace is

ERROR: AssertionError: Should only be one path from root 1 to variable 3. Instead have 2 children from node 31 on the path
Stacktrace:
 [1] follow_path(a::FastDifferentiation.DerivativeGraph{Int64}, root_index::Int64, var_index::Int64)
   @ FastDifferentiation ~/.julia/packages/FastDifferentiation/dkSCb/src/Factoring.jl:525
 [2] evaluate_path
   @ ~/.julia/packages/FastDifferentiation/dkSCb/src/Factoring.jl:555 [inlined]
 [3] _symbolic_jacobian!(graph::FastDifferentiation.DerivativeGraph{…}, partial_variables::Vector{…})
   @ FastDifferentiation ~/.julia/packages/FastDifferentiation/dkSCb/src/Jacobian.jl:47
 [4] _symbolic_jacobian
   @ ~/.julia/packages/FastDifferentiation/dkSCb/src/Jacobian.jl:61 [inlined]
 [5] jacobian(terms::Vector{FastDifferentiation.Node}, partial_variables::Vector{FastDifferentiation.Node})
   @ FastDifferentiation ~/.julia/packages/FastDifferentiation/dkSCb/src/Jacobian.jl:90
 [6] top-level scope
   @ ~/FastDifferentiation_experiments/issue23simplified.jl:12
Some type information was truncated. Use `show(err)` to see complete types.

Looks likely related to #65. Maybe the snippet above could be a good test case for #66.

[moble]

Another clue: I noticed that there are a few terms like

((-(x04) * sin(x02)) * x04)

The error goes away if I replace that by either

((-(x04^2) * sin(x02)))

or

-(((x04) * sin(x02)) * x04)

in one place or another, but not both.

[brianguenter]

Hopefully this bug will be fixed by #65. Thanks for making the small MWE. These make it much easier to debug and verify the fix.