|
1 | | -struct JacobianData{T1 <: SparseMatrixCSC,T2 <: ForwardColorJacCache} |
2 | | - f!::Function |
3 | | - jac::T1 |
4 | | - jac_cache::T2 |
5 | | - length_jac::Int |
6 | | - sparsity::SparseArrays.SparseMatrixCSC{Bool,Int64} |
7 | | - |
8 | | - function JacobianData(f!, output, input) |
9 | | - # TODO Proper way to do it (currently not working) |
10 | | - # Issue: https://github.com/SciML/SparsityDetection.jl/issues/41 |
11 | | - # sparsity = jacobian_sparsity(f!, output, input) |
12 | | - |
13 | | - # Workaround |
| 1 | +import SparseArrays |
| 2 | +import SparseDiffTools |
| 3 | +# import Symbolics |
| 4 | + |
| 5 | +""" |
| 6 | +`JacobianData` is a utility strcture to hold Jacobian information associated with a given function. |
| 7 | +""" |
| 8 | +struct JacobianData{T,F<:Function,C<:SparseDiffTools.ForwardColorJacCache} |
| 9 | + f!::F |
| 10 | + jac::SparseArrays.SparseMatrixCSC{T,Int} |
| 11 | + jac_cache::C |
| 12 | + jac_length::Int |
| 13 | + jac_sparsity::SparseArrays.SparseMatrixCSC{Bool,Int} |
| 14 | + |
| 15 | + @doc """ |
| 16 | + JacobianData(f!, output, input) |
| 17 | +
|
| 18 | + Create a new `JacobianData`. |
| 19 | +
|
| 20 | + # Arguments |
| 21 | + - `f!`: the method associated with this `JacobianData` |
| 22 | + - `output`: input vector to be passed to `f!` |
| 23 | + - `input`: output vector expected from `f!` |
| 24 | + """ |
| 25 | + function JacobianData(f!::F, output::Vector{T}, input::Vector{T}) where {T,F<:Function} |
| 26 | + # # Issue: https://github.com/SciML/SparsityDetection.jl/issues/41 |
| 27 | + # # Calculate the sparsity pattern of the Jacobian |
| 28 | + # jac_sparsity = Symbolics.jacobian_sparsity(f!, output, input) |
| 29 | + |
| 30 | + # [HACKY!] Calculate the sparsity pattern of the Jacobian |
14 | 31 | num_samples = 50 |
15 | 32 | jac_samples = [ForwardDiff.jacobian(f!, rand!(output), rand!(input)) for _ in 1:num_samples] |
16 | | - sparsity = sparse(sum(jac_samples) .≠ 0) |
| 33 | + jac_sparsity = SparseArrays.sparse(sum(jac_samples) .≠ 0) |
17 | 34 |
|
18 | | - jac = convert.(Float64, sparse(sparsity)) |
| 35 | + # Placeholder for the actual Jacobian |
| 36 | + jac = T.(jac_sparsity) |
19 | 37 |
|
20 | | - jac_cache = ForwardColorJacCache(f!, input, dx=output, |
21 | | - colorvec=matrix_colors(jac), |
22 | | - sparsity=sparsity) |
| 38 | + # Color the sparse matrix using graphical techniques (colorvec-assisted differentiation is significantly faster) |
| 39 | + colorvec = SparseDiffTools.matrix_colors(jac) |
23 | 40 |
|
24 | | - length_jac = nnz(jac) |
| 41 | + # Construct the color cache in advance, in order to not compute it each time the Jacobian needs to be evaluated |
| 42 | + jac_cache = SparseDiffTools.ForwardColorJacCache(f!, input, dx=output, colorvec=colorvec, sparsity=jac_sparsity) |
25 | 43 |
|
26 | | - T1 = typeof(jac) |
27 | | - T2 = typeof(jac_cache) |
| 44 | + # The length of the Jacobian, i.e., the number of non-zero elements |
| 45 | + jac_length = SparseArrays.nnz(jac) |
28 | 46 |
|
29 | | - new{T1,T2}(f!, jac, jac_cache, length_jac, sparsity) |
| 47 | + C = typeof(jac_cache) |
| 48 | + |
| 49 | + new{T,F,C}(f!, jac, jac_cache, jac_length, jac_sparsity) |
30 | 50 | end |
31 | 51 | end |
32 | 52 |
|
33 | | -function (jd::JacobianData)(x) |
34 | | - forwarddiff_color_jacobian!(jd.jac, jd.f!, x, jd.jac_cache) |
35 | | -end |
| 53 | +""" |
| 54 | +Evaluate this `JacobianData`'s function `f!` for point `x` using forward-mode automatic differentiation. |
| 55 | +
|
| 56 | +This method's syntax is "special". For more info on Function-like-objects, read the |
| 57 | +[docs](https://docs.julialang.org/en/v1/manual/methods/#Function-like-objects). |
| 58 | +""" |
| 59 | +(jd::JacobianData)(x) = SparseDiffTools.forwarddiff_color_jacobian!(jd.jac, jd.f!, x, jd.jac_cache) |
0 commit comments