diff --git a/Project.toml b/Project.toml index af37712..4703089 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "SparseConnectivityTracer" uuid = "9f842d2f-2579-4b1d-911e-f412cf18a3f5" authors = ["Adrian Hill "] -version = "0.6.8" +version = "0.6.9-DEV" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" diff --git a/src/SparseConnectivityTracer.jl b/src/SparseConnectivityTracer.jl index bfa4701..5eddc83 100644 --- a/src/SparseConnectivityTracer.jl +++ b/src/SparseConnectivityTracer.jl @@ -31,6 +31,7 @@ include("overloads/arrays.jl") include("overloads/ambiguities.jl") include("trace_functions.jl") +include("parse_outputs_to_matrix.jl") include("adtypes_interface.jl") export TracerSparsityDetector diff --git a/src/parse_outputs_to_matrix.jl b/src/parse_outputs_to_matrix.jl new file mode 100644 index 0000000..69687dd --- /dev/null +++ b/src/parse_outputs_to_matrix.jl @@ -0,0 +1,76 @@ +#= Parse tracers from function evaluation in `src/trace_functions.jl` into an output matrix. =# + +_tracer_or_number(x::Real) = x +_tracer_or_number(d::Dual) = tracer(d) + +#==========# +# Jacobian # +#==========# + +function jacobian_tracers_to_matrix( + xt::AbstractArray{T}, yt::AbstractArray +) where {T<:GradientTracer} + n, m = length(xt), length(yt) + I = Int[] # row indices + J = Int[] # column indices + V = Bool[] # values + for (i, y) in enumerate(yt) + if y isa T && !isemptytracer(y) + for j in gradient(y) + push!(I, i) + push!(J, j) + push!(V, true) + end + end + end + return sparse(I, J, V, m, n) +end + +function jacobian_tracers_to_matrix( + xt::AbstractArray{D}, yt::AbstractArray +) where {P,T<:GradientTracer,D<:Dual{P,T}} + return jacobian_tracers_to_matrix(tracer.(xt), _tracer_or_number.(yt)) +end + +#=========# +# Hessian # +#=========# + +function hessian_tracers_to_matrix(xt::AbstractArray{T}, yt::T) where {T<:HessianTracer} + n = length(xt) + I = Int[] # row indices + J = Int[] # column indices + V = Bool[] # values + + if !isemptytracer(yt) + for (i, j) in tuple_set(hessian(yt)) + push!(I, i) + push!(J, j) + push!(V, true) + # TODO: return `Symmetric` instead on next breaking release + push!(I, j) + push!(J, i) + push!(V, true) + end + end + h = sparse(I, J, V, n, n) + return h +end + +function hessian_tracers_to_matrix( + xt::AbstractArray{D1}, yt::D2 +) where {P1,P2,T<:HessianTracer,D1<:Dual{P1,T},D2<:Dual{P2,T}} + return hessian_tracers_to_matrix(tracer.(xt), tracer(yt)) +end + +function hessian_tracers_to_matrix( + xt::AbstractArray{T}, yt::Number +) where {T<:HessianTracer} + return hessian_tracers_to_matrix(xt, myempty(T)) +end + +function hessian_tracers_to_matrix( + xt::AbstractArray{D1}, yt::Number +) where {P1,T<:HessianTracer,D1<:Dual{P1,T}} + return hessian_tracers_to_matrix(tracer.(xt), myempty(T)) +end diff --git a/src/trace_functions.jl b/src/trace_functions.jl index 7d2b52d..d678119 100644 --- a/src/trace_functions.jl +++ b/src/trace_functions.jl @@ -1,7 +1,8 @@ #= This file handles the actual tracing of functions: 1) creating tracers from inputs 2) evaluating the function with the created tracers -3) parsing the resulting tracers into an output matrix + +The resulting output is parsed in `src/parse_outputs_to_matrix.jl`. =# #==================# @@ -58,20 +59,16 @@ end to_array(x::Real) = [x] to_array(x::AbstractArray) = x -# Utilities -_tracer_or_number(x::Real) = x -_tracer_or_number(d::Dual) = tracer(d) - -#================# -# GradientTracer # -#================# +#==========# +# Jacobian # +#==========# # Compute the sparsity pattern of the Jacobian of `y = f(x)`. function _jacobian_sparsity( f, x, ::Type{T}=DEFAULT_GRADIENT_TRACER ) where {T<:GradientTracer} xt, yt = trace_function(T, f, x) - return jacobian_pattern_to_mat(to_array(xt), to_array(yt)) + return jacobian_tracers_to_matrix(to_array(xt), to_array(yt)) end # Compute the sparsity pattern of the Jacobian of `f!(y, x)`. @@ -79,7 +76,7 @@ function _jacobian_sparsity( f!, y, x, ::Type{T}=DEFAULT_GRADIENT_TRACER ) where {T<:GradientTracer} xt, yt = trace_function(T, f!, y, x) - return jacobian_pattern_to_mat(to_array(xt), to_array(yt)) + return jacobian_tracers_to_matrix(to_array(xt), to_array(yt)) end # Compute the local sparsity pattern of the Jacobian of `y = f(x)` at `x`. @@ -88,7 +85,7 @@ function _local_jacobian_sparsity( ) where {T<:GradientTracer} D = Dual{eltype(x),T} xt, yt = trace_function(D, f, x) - return jacobian_pattern_to_mat(to_array(xt), to_array(yt)) + return jacobian_tracers_to_matrix(to_array(xt), to_array(yt)) end # Compute the local sparsity pattern of the Jacobian of `f!(y, x)` at `x`. @@ -97,42 +94,17 @@ function _local_jacobian_sparsity( ) where {T<:GradientTracer} D = Dual{eltype(x),T} xt, yt = trace_function(D, f!, y, x) - return jacobian_pattern_to_mat(to_array(xt), to_array(yt)) -end - -function jacobian_pattern_to_mat( - xt::AbstractArray{T}, yt::AbstractArray{<:Real} -) where {T<:GradientTracer} - n, m = length(xt), length(yt) - I = Int[] # row indices - J = Int[] # column indices - V = Bool[] # values - for (i, y) in enumerate(yt) - if y isa T && !isemptytracer(y) - for j in gradient(y) - push!(I, i) - push!(J, j) - push!(V, true) - end - end - end - return sparse(I, J, V, m, n) -end - -function jacobian_pattern_to_mat( - xt::AbstractArray{D}, yt::AbstractArray{<:Real} -) where {P,T<:GradientTracer,D<:Dual{P,T}} - return jacobian_pattern_to_mat(tracer.(xt), _tracer_or_number.(yt)) + return jacobian_tracers_to_matrix(to_array(xt), to_array(yt)) end -#===============# -# HessianTracer # -#===============# +#=========# +# Hessian # +#=========# # Compute the sparsity pattern of the Hessian of a scalar function `y = f(x)`. function _hessian_sparsity(f, x, ::Type{T}=DEFAULT_HESSIAN_TRACER) where {T<:HessianTracer} xt, yt = trace_function(T, f, x) - return hessian_pattern_to_mat(to_array(xt), yt) + return hessian_tracers_to_matrix(to_array(xt), yt) end # Compute the local sparsity pattern of the Hessian of a scalar function `y = f(x)` at `x`. @@ -141,42 +113,5 @@ function _local_hessian_sparsity( ) where {T<:HessianTracer} D = Dual{eltype(x),T} xt, yt = trace_function(D, f, x) - return hessian_pattern_to_mat(to_array(xt), yt) -end - -function hessian_pattern_to_mat(xt::AbstractArray{T}, yt::T) where {T<:HessianTracer} - n = length(xt) - I = Int[] # row indices - J = Int[] # column indices - V = Bool[] # values - - if !isemptytracer(yt) - for (i, j) in tuple_set(hessian(yt)) - push!(I, i) - push!(J, j) - push!(V, true) - # TODO: return `Symmetric` instead on next breaking release - push!(I, j) - push!(J, i) - push!(V, true) - end - end - h = sparse(I, J, V, n, n) - return h -end - -function hessian_pattern_to_mat( - xt::AbstractArray{D1}, yt::D2 -) where {P1,P2,T<:HessianTracer,D1<:Dual{P1,T},D2<:Dual{P2,T}} - return hessian_pattern_to_mat(tracer.(xt), tracer(yt)) -end - -function hessian_pattern_to_mat(xt::AbstractArray{T}, yt::Number) where {T<:HessianTracer} - return hessian_pattern_to_mat(xt, myempty(T)) -end - -function hessian_pattern_to_mat( - xt::AbstractArray{D1}, yt::Number -) where {P1,T<:HessianTracer,D1<:Dual{P1,T}} - return hessian_pattern_to_mat(tracer.(xt), myempty(T)) + return hessian_tracers_to_matrix(to_array(xt), yt) end diff --git a/test/test_gradient.jl b/test/test_gradient.jl index 2c515e1..87dc127 100644 --- a/test/test_gradient.jl +++ b/test/test_gradient.jl @@ -123,6 +123,36 @@ J(f, x) = jacobian_sparsity(f, x, detector) x -> x[1] > x[2] ? x[3] : x[4], [1.0, 2.0, 3.0, 4.0] ) == [0 0 1 1;] end + + @testset "Output of type Vector{Any}" begin + function f(x::AbstractVector) + n = length(x) + ret = [] # return type will be Vector{Any} + for i in 1:(n - 1) + append!( + ret, + abs2(x[i + 1]) - abs2(x[i]) + abs2(x[n - i]) - abs2(x[n - i + 1]), + ) + end + return ret + end + x = [ + 0.263914 + 0.605532 + 1.281598 + 1.413663 + 0.178133 + -1.705427 + ] + @test J(f, x) == [ + 1 1 0 0 1 1 + 0 1 1 1 1 0 + 0 0 1 1 0 0 + 0 1 1 1 1 0 + 1 1 0 0 1 1 + ] + end + yield() end end @@ -248,6 +278,35 @@ end @test J(x -> log(det(x)), [1.0 -1.0; 2.0 2.0]) == [1 1 1 1] @test J(x -> dot(x[1:2], x[4:5]), [0, 1, 0, 1, 0]) == [1 0 0 0 1] end + @testset "Output of type Vector{Any}" begin + function f(x::AbstractVector) + n = length(x) + ret = [] # return type will be Vector{Any} + for i in 1:(n - 1) + append!( + ret, + abs2(x[i + 1]) - abs2(x[i]) + abs2(x[n - i]) - abs2(x[n - i + 1]), + ) + end + return ret + end + x = [ + 0.263914 + 0.605532 + 1.281598 + 1.413663 + 0.178133 + -1.705427 + ] + @test J(f, x) == [ + 1 1 0 0 1 1 + 0 1 1 1 1 0 + 0 0 1 1 0 0 + 0 1 1 1 1 0 + 1 1 0 0 1 1 + ] + end + yield() end end