diff --git a/src/sparse_jacobian.jl b/src/sparse_jacobian.jl new file mode 100644 index 0000000000000000000000000000000000000000..4429a266478f3463724aee0a672272b5cb08b305 --- /dev/null +++ b/src/sparse_jacobian.jl @@ -0,0 +1,114 @@ +module SparseJacobian + +export sparse_jacobian + +using SparseArrays + +using ForwardDiff: derivative +using LightGraphs: SimpleGraph, add_edge!, greedy_color + +""" + TReal{T} <: Real + +Tracked real variable of type `T`, storing dependent variables as a vector of +indices. +""" +struct TReal{T} <: Real + x::T + origin::Vector{Int} +end + +TReal(x) = TReal(x, Int[]) +TReal(x::TReal) = x +TReal{T}(x) where T = TReal(convert(T, x), Int[]) + +Base.promote_rule(::Type{TReal{S}}, ::Type{T}) where {T, S} = TReal{promote_type(T, S)} + +for op in [:+, :-, :exp, :log, :sin, :cos, :sqrt] + @eval Base.$op(a::TReal) = TReal(Base.$op(a.x), a.origin) +end +for op in [:+, :-, :*, :/, :^, :max, :min] + @eval Base.$op(a::TReal, b::TReal) = TReal(Base.$op(a.x, b.x), union(a.origin, b.origin)) +end + +# FIXME: this construction breaks for non-trivial array types +track(a::AbstractVector) = [TReal(a[i], Int[i]) for i in eachindex(a)] + +function sparsity(f, x) + tfx = f(track(x)) + + I = Int[] + J = Int[] + V = Bool[] + for i in eachindex(tfx) + for j in tfx[i].origin + push!(I, i) + push!(J, j) + push!(V, true) + end + end + + return sparse(I, J, V, length(tfx), length(x)) +end + +""" + sparsity_graph(a::SparseMatrixCSC) + +Create a graph representation of the sparsity pattern in `a`. +Nodes represent columns of `a`, while two are connected by an edge if and only +if both corresponding columns have a common non-zero row. +""" +function sparsity_graph(a::SparseMatrixCSC) + g = SimpleGraph(size(a, 2)) + + rowcols = [Int[] for _ in axes(a, 1)] + rows = rowvals(a) + for j = 1:size(a, 2) + for k in nzrange(a, j) + i = rows[k] + for l in rowcols[i] + add_edge!(g, j, l) + end + push!(rowcols[i], j) + end + end + + return g +end + +""" + sparse_jacobian(f, x; + sp = sparsity(f, x), gcol = greedy_color(sparsity_graph(sp))) + +Determine the jacobian of `f` at `x` optionally by prescribing the sparsity +pattern `sp` and/or the graph coloring `gcol`. +""" +function sparse_jacobian(f, x; sp = sparsity(f, x), gcol = greedy_color(sparsity_graph(sp))) + rows = rowvals(sp) + + I = Int[] + J = Int[] + V = eltype(x)[] + + y = similar(x) + fdir = ε -> f(x .+ ε .* y) + for col in 1:gcol.num_colors + y .= gcol.colors .== col + fd = derivative(fdir, 0) + + for j = 1:size(sp, 2) + gcol.colors[j] == col || continue + for k in nzrange(sp, j) + i = rows[k] + push!(I, i) + push!(J, j) + push!(V, fd[i]) + end + end + end + + return sparse(I, J, V, size(sp)...) +end + + +end