diff --git a/Project.toml b/Project.toml index 4d7abdc..104a77d 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "NonconvexCore" uuid = "035190e5-69f1-488f-aaab-becca2889735" authors = ["Mohamed Tarek and contributors"] -version = "1.0.6" +version = "1.0.7" [deps] ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" diff --git a/src/NonconvexCore.jl b/src/NonconvexCore.jl index bcceb64..6ee3e0d 100644 --- a/src/NonconvexCore.jl +++ b/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") diff --git a/src/functions/sparse.jl b/src/functions/sparse.jl new file mode 100644 index 0000000..17a07a5 --- /dev/null +++ b/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 diff --git a/src/models/moi.jl b/src/models/moi.jl index 4be0a50..e27c8fd 100644 --- a/src/models/moi.jl +++ b/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