Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Relax type annotations in Jacobian output parsing #217

Merged
merged 5 commits into from
Nov 25, 2024
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Rename functions
adrhill committed Nov 25, 2024
commit 435aa8f823052ea4afc90b369fde12f45d4448cf
2 changes: 1 addition & 1 deletion src/SparseConnectivityTracer.jl
Original file line number Diff line number Diff line change
@@ -31,7 +31,7 @@ include("overloads/arrays.jl")
include("overloads/ambiguities.jl")

include("trace_functions.jl")
include("parse_outputs_to_mat.jl")
include("parse_outputs_to_matrix.jl")
include("adtypes_interface.jl")

export TracerSparsityDetector
22 changes: 12 additions & 10 deletions src/parse_outputs_to_mat.jl → src/parse_outputs_to_matrix.jl
Original file line number Diff line number Diff line change
@@ -4,7 +4,7 @@
# Jacobian #
#==========#

function jacobian_pattern_to_mat(
function jacobian_tracers_to_matrix(
xt::AbstractArray{T}, yt::AbstractArray{<:Real}
) where {T<:GradientTracer}
n, m = length(xt), length(yt)
@@ -23,17 +23,17 @@ function jacobian_pattern_to_mat(
return sparse(I, J, V, m, n)
end

function jacobian_pattern_to_mat(
function jacobian_tracers_to_matrix(
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(tracer.(xt), _tracer_or_number.(yt))
end

#=========#
# Hessian #
#=========#

function hessian_pattern_to_mat(xt::AbstractArray{T}, yt::T) where {T<:HessianTracer}
function hessian_tracers_to_matrix(xt::AbstractArray{T}, yt::T) where {T<:HessianTracer}
n = length(xt)
I = Int[] # row indices
J = Int[] # column indices
@@ -54,18 +54,20 @@ function hessian_pattern_to_mat(xt::AbstractArray{T}, yt::T) where {T<:HessianTr
return h
end

function hessian_pattern_to_mat(
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_pattern_to_mat(tracer.(xt), tracer(yt))
return hessian_tracers_to_matrix(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))
function hessian_tracers_to_matrix(
xt::AbstractArray{T}, yt::Number
) where {T<:HessianTracer}
return hessian_tracers_to_matrix(xt, myempty(T))
end

function hessian_pattern_to_mat(
function hessian_tracers_to_matrix(
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(tracer.(xt), myempty(T))
end
16 changes: 8 additions & 8 deletions src/trace_functions.jl
Original file line number Diff line number Diff line change
@@ -2,7 +2,7 @@
1) creating tracers from inputs
2) evaluating the function with the created tracers
The resulting output is parsed in `src/parse_outputs_to_mat.jl`.
The resulting output is parsed in `src/parse_outputs_to_matrix.jl`.
=#

#==================#
@@ -72,15 +72,15 @@ 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)`.
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`.
@@ -89,7 +89,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`.
@@ -98,7 +98,7 @@ 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))
return jacobian_tracers_to_matrix(to_array(xt), to_array(yt))
end

#=========#
@@ -108,7 +108,7 @@ end
# 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`.
@@ -117,5 +117,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
return hessian_tracers_to_matrix(to_array(xt), yt)
end