diff --git a/src/Nonlinear/ReverseAD/types.jl b/src/Nonlinear/ReverseAD/types.jl index fc599accb0..6f16ac7a5d 100644 --- a/src/Nonlinear/ReverseAD/types.jl +++ b/src/Nonlinear/ReverseAD/types.jl @@ -42,6 +42,12 @@ struct _SubexpressionStorage end end +function MOI.Nonlinear.expression(expr::_SubexpressionStorage) + return MOI.Nonlinear.Expression(expr.nodes, expr.const_values) +end + +MOI.Nonlinear.adjacency_matrix(expr::_SubexpressionStorage) = expr.adj + struct _FunctionStorage nodes::Vector{Nonlinear.Node} adj::SparseArrays.SparseMatrixCSC{Bool,Int} @@ -136,6 +142,12 @@ struct _FunctionStorage end end +function MOI.Nonlinear.expression(expr::_FunctionStorage) + return MOI.Nonlinear.Expression(expr.nodes, expr.const_values) +end + +MOI.Nonlinear.adjacency_matrix(expr::_FunctionStorage) = expr.adj + """ NLPEvaluator( model::Nonlinear.Model, diff --git a/src/Nonlinear/types.jl b/src/Nonlinear/types.jl index 36ac83a245..1cd2f3e5cf 100644 --- a/src/Nonlinear/types.jl +++ b/src/Nonlinear/types.jl @@ -76,9 +76,10 @@ tree. struct Expression nodes::Vector{Node} values::Vector{Float64} - Expression() = new(Node[], Float64[]) end +Expression() = Expression(Node[], Float64[]) + function Base.:(==)(x::Expression, y::Expression) return x.nodes == y.nodes && x.values == y.values end @@ -107,6 +108,8 @@ struct Constraint } end +expression(c::Constraint) = c.expression + """ ParameterIndex diff --git a/test/Nonlinear/ReverseAD.jl b/test/Nonlinear/ReverseAD.jl index 608f28d4a0..d46d0836f2 100644 --- a/test/Nonlinear/ReverseAD.jl +++ b/test/Nonlinear/ReverseAD.jl @@ -1421,6 +1421,27 @@ function test_hessian_reinterpret_unsafe() return end +function test_expression_and_adjacency_matrix() + model = Nonlinear.Model() + x = MOI.VariableIndex(1) + expr = Nonlinear.add_expression(model, :($x^2 + 1)) + Nonlinear.set_objective(model, :($expr^2)) + con = Nonlinear.add_constraint(model, :(2 * $expr), MOI.LessThan(1.0)) + @test Nonlinear.expression(model[con]) isa Nonlinear.Expression + evaluator = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x]) + MOI.initialize(evaluator, [:Grad]) + d = evaluator.backend + @test Nonlinear.expression(d.objective) isa Nonlinear.Expression + A = SparseArrays.sparse([2, 3], [1, 1], Bool[1, 1], 3, 3) + @test Nonlinear.adjacency_matrix(d.objective) == A + @test Nonlinear.expression(only(d.constraints)) isa Nonlinear.Expression + @test Nonlinear.adjacency_matrix(only(d.constraints)) == A + A = SparseArrays.sparse([2, 5, 3, 4], [1, 1, 2, 2], Bool[1, 1, 1, 1], 5, 5) + @test Nonlinear.expression(only(d.subexpressions)) isa Nonlinear.Expression + @test Nonlinear.adjacency_matrix(only(d.subexpressions)) == A + return +end + end # module TestReverseAD.runtests()