some sparsity utilities

Author: mohdibntarek

Project.toml

@@ -1,7 +1,7 @@
 name = "NonconvexCore"
 uuid = "035190e5-69f1-488f-aaab-becca2889735"
 authors = ["Mohamed Tarek <[email protected]> and contributors"]
-version = "1.0.6"
+version = "1.0.7"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/NonconvexCore.jl

@@ -46,6 +46,7 @@ include("utilities/params.jl")
 include("functions/functions.jl")
 include("functions/value_jacobian.jl")
 include("functions/counting_function.jl")
+include("functions/sparse.jl")
 include("common.jl")
 include("utilities/callbacks.jl")
 include("utilities/convergence.jl")

src/functions/sparse.jl

@@ -0,0 +1,39 @@
+# workaround Zygote.jacobian and ForwardDiff.jacobian not supporting sparse jacobians
+
+function sparse_jacobian(f, x)
+    val, pb = Zygote.pullback(f, x)
+    M, N = length(val), length(x)
+    T = eltype(val)
+    return copy(mapreduce(hcat, 1:M, init = spzeros(T, N, 0)) do i
+        pb(I(M)[:, i])[1]
+    end')
+end
+function sparse_hessian(f, x)
+    return sparse_fd_jacobian(x -> Zygote.gradient(f, x)[1], x)
+end
+
+function sparse_fd_jacobian(f, x)
+    pf = pushforward_function(f, x)
+    M, N = length(f(x)), length(x)
+    init = pf(I(M)[:, 1])[1]
+    M = length(init)
+    return mapreduce(hcat, 2:N; init) do i
+        pf(I(M)[:, i])[1]
+    end
+end
+
+# from AbstractDifferentiation
+function pushforward_function(f, xs...)
+    return function pushforward(vs)
+        if length(xs) == 1
+            v = vs isa Tuple ? only(vs) : vs
+            return (ForwardDiff.derivative(h -> f(step_toward(xs[1], v, h)), 0),)
+        else
+            return ForwardDiff.derivative(h -> f(step_toward.(xs, vs, h)...), 0)
+        end
+    end
+end
+
+@inline step_toward(x::Number, v::Number, h) = x + h * v
+# support arrays and tuples
+@noinline step_toward(x, v, h) = x .+ h .* v

src/models/moi.jl

@@ -222,6 +222,15 @@ nvalues(::Nothing) = 0
 nvalues(J::Matrix) = length(J)
 nvalues(H::LowerTriangular{<:Real, <:Matrix}) = (size(H, 1) + 1) * size(H, 1) ÷ 2
 nvalues(H::SparseMatrixCSC) = length(H.nzval)
+function nvalues(HL::LowerTriangular{<:Real, <:SparseMatrixCSC})
+    nvalues = 0
+    H = HL.data
+    for col in 1:length(H.colptr)-1
+        indices = [i for i in H.colptr[col]:H.colptr[col+1]-1 if H.rowval[i] >= col]
+        nvalues += length(indices)
+    end
+    return nvalues
+end
 
 # Implement these for sparse matrices
 function fill_indices!(rows, cols, J0::Matrix; offset = 0, row_offset = 0)
@@ -252,6 +261,17 @@ function fill_indices!(rows, cols, HL::SparseMatrixCSC; offset = 0, row_offset =
     end
     return rows, cols
 end
+function fill_indices!(rows, cols, HL::LowerTriangular{<:Real, <:SparseMatrixCSC}; offset = 0, row_offset = 0)
+    H = HL.data
+    for col in 1:length(H.colptr)-1
+        indices = [i for i in H.colptr[col]:H.colptr[col+1]-1 if H.rowval[i] >= col]
+        nvars = length(indices)
+        cols[offset + 1 : offset + nvars] .= col
+        rows[offset + 1 : offset + nvars] = row_offset .+ H.rowval[indices]
+        offset += nvars
+    end
+    return rows, cols
+end
 
 function add_values!(values, J::Matrix; offset = 0)
     nvars = length(J)
@@ -271,6 +291,16 @@ function add_values!(values, HL::SparseMatrixCSC; factor = 1, offset = 0)
     values[offset+1:offset+nvars] .= HL.nzval .* factor
     return values
 end
+function add_values!(values, HL::LowerTriangular{<:Real, <:SparseMatrixCSC}; factor = 1, offset = 0)
+    H = HL.data
+    for col in 1:length(H.colptr)-1
+        indices = [i for i in H.colptr[col]:H.colptr[col+1]-1 if H.rowval[i] >= col]
+        nvars = length(indices)
+        values[offset + 1 : offset + nvars] .= H.nzval[indices] .* factor
+        offset += nvars
+    end
+    return values
+end
 
 _dot(f, x, y) = dot(f(x), y)
 _dot(::Nothing, ::Any, ::Any) = 0.0