Implement adjacency_matrix and expression in _FunctionStorage

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,

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
 

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()