Merge pull request #41 from JuliaDiffEq/refactor

Refactor and update README to include sparsity detection

Author: ChrisRackauckas

README.md

@@ -29,30 +29,35 @@ end
 ```
 
 For this function, we know that the sparsity pattern of the Jacobian is a
-`Tridiagonal` matrix. We represent our sparsity by that matrix:
+`Tridiagonal` matrix. However, if we didn't know the sparsity pattern for
+the Jacobian, we could use the `sparsity!` function to automatically
+detect the sparsity pattern. We declare that it outputs a length 30 vector
+and takes in a length 30 vector, and it spits out a `Sparsity` object
+which we can turn into a `SparseMatrixCSC`:
 
 ```julia
-sparsity_pattern = Tridiagonal(ones(29),ones(30),ones(29))
+sparsity_pattern = sparsity!(f,output,input)
+jac = Float64.(sparse(sparsity_pattern))
 ```
 
 Now we call `matrix_colors` to get the color vector for that matrix:
 
 ```julia
-colors = matrix_colors(sparsity_pattern)
+colors = matrix_colors(jac)
 ```
 
 Since `maximum(colors)` is 3, this means that finite differencing can now
 compute the Jacobian in just 4 `f`-evaluations:
 
 ```julia
-J = DiffEqDiffTools.finite_difference_jacobian(f, rand(30), color=colors)
+DiffEqDiffTools.finite_difference_jacobian!(jac, f, rand(30), color=colors)
 @show fcalls # 4
 ```
 
 In addition, a faster forward-mode autodiff call can be utilized as well:
 
 ```julia
-forwarddiff_color_jacobian!(sparsity_pattern, f, x, color = colors)
+forwarddiff_color_jacobian!(jac, f, x, color = colors)
 ```
 
 If one only need to compute products, one can use the operators. For example,
@@ -83,6 +88,36 @@ gmres!(res,J,v)
 
 ## Documentation
 
+### Automated Sparsity Detection
+
+Automated sparsity detection is provided by the `sparsity!` function whose
+syntax is:
+
+```julia
+`sparsity!(f, Y, X, args...; sparsity=Sparsity(length(X), length(Y)), verbose=true)`
+```
+
+The arguments are:
+
+- `f`: the function
+- `Y`: the output array
+- `X`: the input array
+- `args`: trailing arguments to `f`. They are considered subject to change, unless wrapped as `Fixed(arg)`
+- `S`: (optional) the sparsity pattern
+- `verbose`: (optional) whether to describe the paths taken by the sparsity detection.
+
+The function `f` is assumed to take arguments of the form `f(dx,x,args...)`.
+`sparsity!` returns a `Sparsity` object which describes where the non-zeros
+of the Jacobian occur. `sparse(::Sparsity)` transforms the pattern into
+a sparse matrix.
+
+This function utilizes non-standard interpretation, which we denote
+combinatoric concolic analysis, to directly realize the sparsity pattern from the program's AST. It requires that the function `f` is a Julia function. It does not
+work numerically, meaning that it is not prone to floating point error or
+cancelation. It allows for branching and will automatically check all of the
+branches. However, a while loop of indeterminate length which is dependent
+on the input argument is not allowed.
+
 ### Matrix Coloring
 
 Matrix coloring allows you to reduce the number of times finite differencing

src/SparseDiffTools.jl

@@ -1,10 +1,16 @@
 module SparseDiffTools
 
 using SparseArrays, LinearAlgebra, BandedMatrices, BlockBandedMatrices,
-      LightGraphs, VertexSafeGraphs, DiffEqDiffTools, ForwardDiff, Zygote
+      LightGraphs, VertexSafeGraphs, DiffEqDiffTools, ForwardDiff, Zygote,
+      SparseArrays
 using BlockBandedMatrices:blocksize,nblocks
 using ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD
 
+using Cassette
+import Cassette: tag, untag, Tagged, metadata, hasmetadata, istagged, canrecurse
+import Cassette: tagged_new_tuple, ContextTagged, BindingMeta, DisableHooks, nametype
+import Core: SSAValue
+
 export  contract_color,
         greedy_d1,
         matrix2graph,
@@ -20,7 +26,8 @@ export  contract_color,
         auto_hesvecgrad,auto_hesvecgrad!,
         numback_hesvec,numback_hesvec!,
         autoback_hesvec,autoback_hesvec!,
-        JacVec,HesVec,HesVecGrad
+        JacVec,HesVec,HesVecGrad,
+        Sparsity, sparsity!
 
 
 include("coloring/high_level.jl")
@@ -30,5 +37,8 @@ include("coloring/matrix2graph.jl")
 include("differentiation/compute_jacobian_ad.jl")
 include("differentiation/jaches_products.jl")
 include("program_sparsity/program_sparsity.jl")
+include("program_sparsity/sparsity_tracker.jl")
+include("program_sparsity/path.jl")
+include("program_sparsity/take_all_branches.jl")
 
 end # module

src/program_sparsity/program_sparsity.jl

@@ -1,13 +1,9 @@
-include("sparsity_tracker.jl")
-include("path.jl")
-include("take_all_branches.jl")
-
 struct Fixed
     value
 end
 
 """
-`sparsity!(f, Y, X, args...; sparsity=Sparsity(length(X), length(Y)))`
+`sparsity!(f, Y, X, args...; sparsity=Sparsity(length(X), length(Y)), verbose=true)`
 
 Execute the program that figures out the sparsity pattern of
 the jacobian of the function `f`.
@@ -18,10 +14,12 @@ the jacobian of the function `f`.
 - `X`: the input array
 - `args`: trailing arguments to `f`. They are considered subject to change, unless wrapped as `Fixed(arg)`
 - `S`: (optional) the sparsity pattern
+- `verbose`: (optional) whether to describe the paths taken by the sparsity detection.
 
 Returns a `Sparsity`
 """
-function sparsity!(f!, Y, X, args...; sparsity=Sparsity(length(Y), length(X)))
+function sparsity!(f!, Y, X, args...; sparsity=Sparsity(length(Y), length(X)),
+                                      verbose = true)
     path = Path()
     ctx = SparsityContext(metadata=(sparsity, path), pass=BranchesPass)
     ctx = Cassette.enabletagging(ctx, f!)
@@ -36,7 +34,7 @@ function sparsity!(f!, Y, X, args...; sparsity=Sparsity(length(Y), length(X)))
                          map(arg -> arg isa Fixed ?
                              arg.value : tag(arg, ctx, ProvinanceSet(())), args)...)
 
-        println("Explored path: ", path)
+        verbose && println("Explored path: ", path)
         alldone(path) && break
         reset!(path)
     end

src/program_sparsity/sparsity_tracker.jl

@@ -1,10 +1,3 @@
-using Cassette
-import Cassette: tag, untag, Tagged, metadata, hasmetadata, istagged, canrecurse
-import Core: SSAValue
-using SparseArrays
-
-export Sparsity, sparsity!
-
 """
 The sparsity pattern.