Merge pull request #498 from SciML/s/sparse_jacobian

Fast Jacobian Sparsity

Author: ChrisRackauckas

src/differentials.jl

@@ -48,19 +48,36 @@ function expand_derivatives(O::Operation,simplify=true)
         isa(o, Operation)   || return O
         isa(o.op, Variable) && return O
 
-        x = sum(1:length(o.args)) do i
+        l = length(o.args)
+        exprs = Expression[]
+        c = 0
+
+        for i in 1:l
             t2 = expand_derivatives(D(o.args[i]),false)
 
-            if _iszero(t2)
-                return t2
+            x = if _iszero(t2)
+                t2
             elseif _isone(t2)
-                return derivative(o, i)
+                derivative(o, i)
             else
-                return derivative(o, i) * t2
+                derivative(o, i) * t2
+            end
+
+            if _iszero(x)
+                continue
+            elseif x isa Expression
+                push!(exprs, x)
+            else
+                c += x
             end
         end
 
-        return simplify ? ModelingToolkit.simplify(x) : x
+        if isempty(exprs)
+            return c
+        else
+            x = Operation(+, !iszero(c) ? vcat(c, exprs) : exprs)
+            return simplify ? ModelingToolkit.simplify(x) : x
+        end
     end
 
     return simplify ? ModelingToolkit.simplify(O) : O

src/direct.jl

@@ -22,6 +22,55 @@ function jacobian(ops::AbstractVector{<:Expression}, vars::AbstractVector{<:Expr
     [expand_derivatives(Differential(v)(O),simplify) for O in ops, v in vars]
 end
 
+"""
+```julia
+sparsejacobian(ops::AbstractVector{<:Expression}, vars::AbstractVector{<:Expression}; simplify = true)
+```
+
+A helper function for computing the sparse Jacobian of an array of expressions with respect to
+an array of variable expressions.
+"""
+function sparsejacobian(ops::AbstractVector{<:Expression}, vars::AbstractVector{<:Expression}; simplify = true)
+    I = Int[]
+    J = Int[]
+    du = Expression[]
+
+    sp = jacobian_sparsity(ops, vars)
+    I,J,_ = findnz(sp)
+
+    exprs = Expression[]
+
+    for (i,j) in zip(I, J)
+        push!(exprs, expand_derivatives(Differential(vars[j])(ops[i])))
+    end
+    sparse(I,J, exprs, length(ops), length(vars))
+end
+
+using SymbolicUtils: @rule, Rewriters
+
+function jacobian_sparsity(du, u)
+    dict = Dict(zip(to_symbolic.(u), 1:length(u)))
+
+    i = Ref(1)
+    I = Int[]
+    J = Int[]
+
+    # This rewriter notes down which u's appear in a
+    # given du (whose index is stored in the `i` Ref)
+    r = [@rule ~x::(x->haskey(dict, x)) => begin
+        push!(I, i[])
+        push!(J, dict[~x])
+        nothing
+    end] |> Rewriters.Chain |> Rewriters.Postwalk
+
+    for ii = 1:length(du)
+        i[] = ii
+        r(to_symbolic(du[ii]))
+    end
+
+    sparse(I, J, true, length(du), length(u))
+end
+
 """
 ```julia
 hessian(O::Expression, vars::AbstractVector{<:Expression}; simplify = true)