From 3c5048d74a30df95a50e4311da4392b08f19d333 Mon Sep 17 00:00:00 2001 From: ArnoStrouwen Date: Sun, 9 Oct 2022 15:16:50 +0200 Subject: [PATCH 1/3] sparsity helper functions --- src/diff.jl | 42 +++++++++++++++++++++++++++++++++--------- test/diff.jl | 10 ++++++++-- 2 files changed, 41 insertions(+), 11 deletions(-) diff --git a/src/diff.jl b/src/diff.jl index 32b37b876..ac0224ebc 100644 --- a/src/diff.jl +++ b/src/diff.jl @@ -530,17 +530,23 @@ end """ ```julia -jacobian_sparsity(op!,output::Array{T},input::Array{T}) where T<:Number +jacobian_sparsity(f!,output::AbstractVector{T},input::AbstractVector{T}, args...;kwargs...) where T<:Number ``` -Return the sparsity pattern of the Jacobian of the mutating function `op!(output,input,args...)`. -""" -function jacobian_sparsity(op!,output::Array{T},input::Array{T}, args...) where T<:Number - eqs=similar(output,Num) - fill!(eqs,false) - vars=ArrayInterfaceCore.restructure(input,[variable(i) for i in eachindex(input)]) - op!(eqs,vars, args...) - jacobian_sparsity(eqs,vars) +Return the sparsity pattern of the Jacobian of the mutating function `f!(output,input,args...;kwargs...)`. +""" +function jacobian_sparsity( + f!::Function, + output::Array{T}, + input::Array{T}, + args...; + kwargs..., +) where {T<:Number} + eqs = similar(output, Num) + fill!(eqs, false) + vars = ArrayInterfaceCore.restructure(input, [variable(i) for i in eachindex(input)]) + f!(eqs, vars, args...; kwargs...) + jacobian_sparsity(eqs, vars) end @@ -629,6 +635,24 @@ let end end +""" +```julia +hessian_sparsity(f,input::AbstractVector{T}, args...;kwargs...) where T<:Number +``` + +Return the sparsity pattern of the Hessian of the function `f(input,args...;kwargs...)`. +""" +function hessian_sparsity( + f::Function, + input::AbstractVector{T}, + args...; + kwargs..., +) where {T<:Number} + vars = ArrayInterfaceCore.restructure(input, [variable(i) for i in eachindex(input)]) + eq = f(vars, args...; kwargs...) + hessian_sparsity(eq, vars) +end + """ $(SIGNATURES) diff --git a/test/diff.jl b/test/diff.jl index 5cd2717ae..5ae5313f3 100644 --- a/test/diff.jl +++ b/test/diff.jl @@ -189,7 +189,7 @@ function f!(res,u) (x,y,z)=u res.=[x^2, y^3, x^4, sin(y), x+y, x+z^2, z+x, x+y^2+sin(z)] end -function f1!(res,u,a,b,c) +function f1!(res,u,a,b;c) (x,y,z)=u res.=[a*x^2, y^3, b*x^4, sin(y), c*x+y, x+z^2, a*z+x, x+y^2+sin(z)] end @@ -198,7 +198,7 @@ input=rand(3) output=rand(8) findnz(Symbolics.jacobian_sparsity(f!, output, input))[[1,2]] == findnz(reference_jac)[[1,2]] -findnz(Symbolics.jacobian_sparsity(f1!, output, input,1,2,3))[[1,2]] == findnz(reference_jac)[[1,2]] +findnz(Symbolics.jacobian_sparsity(f1!, output, input,1,2,c=3))[[1,2]] == findnz(reference_jac)[[1,2]] input = rand(2,2) function f2!(res,u,a,b,c) @@ -228,15 +228,21 @@ using Symbolics rosenbrock(X) = sum(1:length(X)-1) do i 100 * (X[i+1] - X[i]^2)^2 + (1 - X[i])^2 end +rosenbrock2(X,a;b) = sum(1:length(X)-1) do i + a * (X[i+1] - X[i]^2)^2 + (b - X[i])^2 +end @variables a,b X = [a,b] +input = rand(2) spoly(x) = simplify(x, expand=true) rr = rosenbrock(X) reference_hes = Symbolics.hessian(rr, X) @test findnz(sparse(reference_hes))[1:2] == findnz(Symbolics.hessian_sparsity(rr, X))[1:2] +@test findnz(sparse(reference_hes))[1:2] == findnz(Symbolics.hessian_sparsity(rosenbrock, input))[1:2] +@test findnz(sparse(reference_hes))[1:2] == findnz(Symbolics.hessian_sparsity(rosenbrock2, input,100,b=1))[1:2] sp_hess = Symbolics.sparsehessian(rr, X) @test findnz(sparse(reference_hes))[1:2] == findnz(sp_hess)[1:2] From b53fd4f17f28a4bc9902e8ef0686e6c9d4b6560e Mon Sep 17 00:00:00 2001 From: "Bowen S. Zhu" Date: Sun, 9 Oct 2022 13:03:47 -0400 Subject: [PATCH 2/3] Update `hessian_sparsity` docstring --- src/diff.jl | 79 +++++++++++++++++++++++++++++++++++------------------ 1 file changed, 53 insertions(+), 26 deletions(-) diff --git a/src/diff.jl b/src/diff.jl index ac0224ebc..7e9f25310 100644 --- a/src/diff.jl +++ b/src/diff.jl @@ -584,13 +584,6 @@ end isidx(x) = x isa TermCombination -""" - hessian_sparsity(op, vars::AbstractVector) - -Return the sparsity pattern of the Hessian of an expression with respect to -an array of variable expressions. -""" -function hessian_sparsity end basic_simterm(t, g, args; kws...) = Term{Any}(g, args) let @@ -623,34 +616,68 @@ let global hessian_sparsity - function hessian_sparsity(f, u) - @assert !(f isa AbstractArray) - f = value(f) - u = map(value, u) + @doc """ + $(TYPEDSIGNATURES) + + Return the sparsity pattern of the Hessian of an expression with respect to + an array of variable expressions. + + # Arguments + - `expr`: a symbolic expression. + - `vars`: a vector of symbolic variables. + + # Examples + ```jldoctest + julia> using Symbolics + + julia> vars = @variables x₁ x₂; + + julia> expr = 3x₁^2 + 4x₁ * x₂; + + julia> Symbolics.hessian_sparsity(expr, vars) + 2×2 SparseArrays.SparseMatrixCSC{Bool, Int64} with 3 stored entries: + 1 1 + 1 ⋅ + ``` + """ + function hessian_sparsity(expr, vars::AbstractVector) + @assert !(expr isa AbstractArray) + expr = value(expr) + u = map(value, vars) idx(i) = TermCombination(Set([Dict(i=>1)])) dict = Dict(u .=> idx.(1:length(u))) - f = Rewriters.Prewalk(x->haskey(dict, x) ? dict[x] : x; similarterm=basic_simterm)(f) + f = Rewriters.Prewalk(x->haskey(dict, x) ? dict[x] : x; similarterm=basic_simterm)(expr) lp = linearity_propagator(f) _sparse(lp, length(u)) end end - """ -```julia -hessian_sparsity(f,input::AbstractVector{T}, args...;kwargs...) where T<:Number -``` +$(TYPEDSIGNATURES) -Return the sparsity pattern of the Hessian of the function `f(input,args...;kwargs...)`. +Return the sparsity pattern of the Hessian of the given function `f`. + +# Arguments +- `f`: an out-of-place function `f(input, args...; kwargs...)`. +- `input`: a vector of input values whose [eltype](https://docs.julialang.org/en/v1/base/collections/#Base.eltype) can be either symbolic or [primitive](https://docs.julialang.org/en/v1/manual/types/#Primitive-Types). + +# Examples +```jldoctest +julia> using Symbolics + +julia> f(x) = 4x[1] * x[2] - 5x[2]^2; + +julia> input = Vector{Float64}(undef, 2); + +julia> Symbolics.hessian_sparsity(f, input) +2×2 SparseArrays.SparseMatrixCSC{Bool, Int64} with 3 stored entries: + ⋅ 1 + 1 1 +``` """ -function hessian_sparsity( - f::Function, - input::AbstractVector{T}, - args...; - kwargs..., -) where {T<:Number} - vars = ArrayInterfaceCore.restructure(input, [variable(i) for i in eachindex(input)]) - eq = f(vars, args...; kwargs...) - hessian_sparsity(eq, vars) +function hessian_sparsity(f::Function, input::AbstractVector, args...; kwargs...) + vars = ArrayInterfaceCore.restructure(input, map(variable, eachindex(input))) + expr = f(vars, args...; kwargs...) + hessian_sparsity(expr, vars) end """ From d3d4d0a3e6664983b47c18127060c81584c39866 Mon Sep 17 00:00:00 2001 From: "Bowen S. Zhu" Date: Sun, 9 Oct 2022 13:08:17 -0400 Subject: [PATCH 3/3] Update `jacobian_sparsity` docstring --- src/diff.jl | 81 ++++++++++++++++++++++++++++++++++++----------------- 1 file changed, 56 insertions(+), 25 deletions(-) diff --git a/src/diff.jl b/src/diff.jl index 7e9f25310..e5cf7bbb2 100644 --- a/src/diff.jl +++ b/src/diff.jl @@ -489,16 +489,33 @@ function sparsejacobian_vals(ops::AbstractVector, vars::AbstractVector, I::Abstr end """ -```julia -jacobian_sparsity(ops::AbstractVector, vars::AbstractVector) -``` +$(TYPEDSIGNATURES) Return the sparsity pattern of the Jacobian of an array of expressions with respect to an array of variable expressions. + +# Arguments +- `exprs`: an array of symbolic expressions. +- `vars`: an array of symbolic variables. + +# Examples +```jldoctest +julia> using Symbolics + +julia> vars = @variables x₁ x₂; + +julia> exprs = [2x₁, 3x₂, 4x₁ * x₂]; + +julia> Symbolics.jacobian_sparsity(exprs, vars) +3×2 SparseArrays.SparseMatrixCSC{Bool, Int64} with 4 stored entries: + 1 ⋅ + ⋅ 1 + 1 1 +``` """ -function jacobian_sparsity(du, u) - du = map(value, du) - u = map(value, u) +function jacobian_sparsity(exprs::AbstractArray, vars::AbstractArray) + du = map(value, exprs) + u = map(value, vars) dict = Dict(zip(u, 1:length(u))) i = Ref(1) @@ -526,31 +543,45 @@ function jacobian_sparsity(du, u) sparse(I, J, true, length(du), length(u)) end +""" +$(TYPEDSIGNATURES) +Return the sparsity pattern of the Jacobian of the mutating function `f!`. -""" -```julia -jacobian_sparsity(f!,output::AbstractVector{T},input::AbstractVector{T}, args...;kwargs...) where T<:Number -``` +# Arguments +- `f!`: an in-place function `f!(output, input, args...; kwargs...)`. +- `output`: output array. +- `input`: input array. -Return the sparsity pattern of the Jacobian of the mutating function `f!(output,input,args...;kwargs...)`. -""" -function jacobian_sparsity( - f!::Function, - output::Array{T}, - input::Array{T}, - args...; - kwargs..., -) where {T<:Number} - eqs = similar(output, Num) - fill!(eqs, false) - vars = ArrayInterfaceCore.restructure(input, [variable(i) for i in eachindex(input)]) - f!(eqs, vars, args...; kwargs...) - jacobian_sparsity(eqs, vars) -end +The [eltype](https://docs.julialang.org/en/v1/base/collections/#Base.eltype) +of `output` and `input` can be either symbolic or +[primitive](https://docs.julialang.org/en/v1/manual/types/#Primitive-Types). + +# Examples +```jldoctest +julia> using Symbolics +julia> f!(y, x) = y .= [x[2], 2x[1], 3x[1] * x[2]]; +julia> output = Vector{Float64}(undef, 3); +julia> input = Vector{Float64}(undef, 2); + +julia> Symbolics.jacobian_sparsity(f!, output, input) +3×2 SparseArrays.SparseMatrixCSC{Bool, Int64} with 4 stored entries: + ⋅ 1 + 1 ⋅ + 1 1 +``` +""" +function jacobian_sparsity(f!::Function, output::AbstractArray, input::AbstractArray, + args...; kwargs...) + exprs = similar(output, Num) + fill!(exprs, false) + vars = ArrayInterfaceCore.restructure(input, map(variable, eachindex(input))) + f!(exprs, vars, args...; kwargs...) + jacobian_sparsity(exprs, vars) +end """ exprs_occur_in(exprs::Vector, expr)