diff --git a/src/diff.jl b/src/diff.jl index 32b37b876..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,25 +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(op!,output::Array{T},input::Array{T}) 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 `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) -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) @@ -578,13 +615,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 @@ -617,17 +647,69 @@ 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 +""" +$(TYPEDSIGNATURES) + +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, args...; kwargs...) + vars = ArrayInterfaceCore.restructure(input, map(variable, eachindex(input))) + expr = f(vars, args...; kwargs...) + hessian_sparsity(expr, 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]