[Nonlinear] improve parsing performance of Scalar(Affine,Quadratic)Function (#2487)

Author: odow

src/Nonlinear/parse.jl

@@ -265,22 +265,98 @@ end
 function parse_expression(
     data::Model,
     expr::Expression,
-    x::MOI.ScalarAffineFunction,
+    f::MOI.ScalarAffineFunction,
     parent_index::Int,
 )
-    f = convert(MOI.ScalarNonlinearFunction, x)
-    parse_expression(data, expr, f, parent_index)
+    if isempty(f.terms)
+        parse_expression(data, expr, f.constant, parent_index)
+        return
+    elseif iszero(f.constant) && length(f.terms) == 1
+        # Expression of for `a * x`
+        parse_expression(data, expr, only(f.terms), parent_index)
+        return
+    end
+    id_plus = data.operators.multivariate_operator_to_id[:+]
+    push!(expr.nodes, Node(NODE_CALL_MULTIVARIATE, id_plus, parent_index))
+    new_parent = length(expr.nodes)
+    for term in f.terms
+        parse_expression(data, expr, term, new_parent)
+    end
+    if !iszero(f.constant)
+        parse_expression(data, expr, f.constant, new_parent)
+    end
     return
 end
 
 function parse_expression(
     data::Model,
     expr::Expression,
-    x::MOI.ScalarQuadraticFunction,
+    f::MOI.ScalarQuadraticFunction,
     parent_index::Int,
 )
-    f = convert(MOI.ScalarNonlinearFunction, x)
-    parse_expression(data, expr, f, parent_index)
+    if isempty(f.quadratic_terms) && isempty(f.affine_terms)
+        parse_expression(data, expr, f.constant, parent_index)
+        return
+    elseif iszero(f.constant)
+        if length(f.quadratic_terms) == 1 && isempty(f.affine_terms)
+            parse_expression(data, expr, only(f.quadratic_terms), parent_index)
+            return
+        elseif isempty(f.quadratic_terms) && length(f.affine_terms) == 1
+            parse_expression(data, expr, only(f.affine_terms), parent_index)
+            return
+        end
+    end
+    id_plus = data.operators.multivariate_operator_to_id[:+]
+    push!(expr.nodes, Node(NODE_CALL_MULTIVARIATE, id_plus, parent_index))
+    new_parent = length(expr.nodes)
+    for term in f.quadratic_terms
+        parse_expression(data, expr, term, new_parent)
+    end
+    for term in f.affine_terms
+        parse_expression(data, expr, term, new_parent)
+    end
+    if !iszero(f.constant)
+        parse_expression(data, expr, f.constant, new_parent)
+    end
+    return
+end
+
+function parse_expression(
+    data::Model,
+    expr::Expression,
+    x::MOI.ScalarAffineTerm,
+    parent_index::Int,
+)
+    if isone(x.coefficient)
+        parse_expression(data, expr, x.variable, parent_index)
+    else
+        id_mul = data.operators.multivariate_operator_to_id[:*]
+        push!(expr.nodes, Node(NODE_CALL_MULTIVARIATE, id_mul, parent_index))
+        mul_parent = length(expr.nodes)
+        parse_expression(data, expr, x.coefficient, mul_parent)
+        parse_expression(data, expr, x.variable, mul_parent)
+    end
+    return
+end
+
+function parse_expression(
+    data::Model,
+    expr::Expression,
+    x::MOI.ScalarQuadraticTerm,
+    parent_index::Int,
+)
+    id_mul = data.operators.multivariate_operator_to_id[:*]
+    push!(expr.nodes, Node(NODE_CALL_MULTIVARIATE, id_mul, parent_index))
+    mul_parent = length(expr.nodes)
+    coef = x.coefficient
+    if x.variable_1 == x.variable_2
+        coef /= 2
+    end
+    if !isone(coef)
+        parse_expression(data, expr, coef, mul_parent)
+    end
+    parse_expression(data, expr, x.variable_1, mul_parent)
+    parse_expression(data, expr, x.variable_2, mul_parent)
     return
 end
 

test/Nonlinear/Nonlinear.jl

@@ -1006,24 +1006,49 @@ end
 function test_scalar_nonlinear_function_parse_scalaraffinefunction()
     model = MOI.Utilities.Model{Float64}()
     x = MOI.add_variable(model)
-    f = 1.0 * x + 2.0
-    nlp_model = MOI.Nonlinear.Model()
-    e1 = MOI.Nonlinear.add_expression(nlp_model, f)
-    e2 = MOI.Nonlinear.add_expression(nlp_model, :(1.0 * $x + 2.0))
-    @test nlp_model[e1] == nlp_model[e2]
+    y = MOI.add_variable(model)
+    terms = MOI.ScalarAffineTerm{Float64}[]
+    for (f, expr) in (
+        MOI.ScalarAffineFunction(terms, 0.0) => :(0.0),
+        MOI.ScalarAffineFunction(terms, 1.0) => :(1.0),
+        (1.0 * x + 2.0) => :($x + 2.0),
+        (2.0 * x + 2.0) => :(2.0 * $x + 2.0),
+        (2.0 * x + -3.0 * y) => :(2.0 * $x + -3.0 * $y),
+        (2.0 * x) => :(2.0 * $x),
+        (1.0 * x) => :($x),
+    )
+        nlp_model = MOI.Nonlinear.Model()
+        f1 = MOI.Nonlinear.add_expression(nlp_model, f)
+        f2 = MOI.Nonlinear.add_expression(nlp_model, expr)
+        @test nlp_model[f1] == nlp_model[f2]
+    end
     return
 end
 
 function test_scalar_nonlinear_function_parse_scalarquadraticfunction()
     model = MOI.Utilities.Model{Float64}()
     x = MOI.add_variable(model)
     y = MOI.add_variable(model)
-    f = 1.5 * x * x + 2.5 * x * y + 3.5 * x + 2.0
-    nlp_model = MOI.Nonlinear.Model()
-    e1 = MOI.Nonlinear.add_expression(nlp_model, f)
-    f_expr = :(1.5 * $x * $x + 2.5 * $x * $y + 3.5 * $x + 2.0)
-    e2 = MOI.Nonlinear.add_expression(nlp_model, f_expr)
-    @test nlp_model[e1] == nlp_model[e2]
+    terms = MOI.ScalarAffineTerm{Float64}[]
+    qterms = MOI.ScalarQuadraticTerm{Float64}[]
+    aterm = MOI.ScalarAffineTerm(2.0, x)
+    for (f, expr) in (
+        MOI.ScalarQuadraticFunction(qterms, terms, 0.0) => :(0.0),
+        MOI.ScalarQuadraticFunction(qterms, terms, 1.0) => :(1.0),
+        MOI.ScalarQuadraticFunction(qterms, [aterm], 0.0) => :(2.0 * $x),
+        (1.0 * x * x + 1.0 * x + 1.0) => :($x * $x + $x + 1),
+        (1.0 * x * x + 1.0 * x) => :($x * $x + $x),
+        (1.0 * x * x + 2.0 * x) => :($x * $x + 2.0 * $x),
+        (2.0 * x * x + 2.0 * x) => :(2.0 * $x * $x + 2.0 * $x),
+        (1.0 * x * x) => :($x * $x),
+        (1.5 * x * x + 2.5 * x * y + 3.5 * x + 2.0) =>
+            :(1.5 * $x * $x + 2.5 * $x * $y + 3.5 * $x + 2.0),
+    )
+        nlp_model = MOI.Nonlinear.Model()
+        f1 = MOI.Nonlinear.add_expression(nlp_model, f)
+        f2 = MOI.Nonlinear.add_expression(nlp_model, expr)
+        @test nlp_model[f1] == nlp_model[f2]
+    end
     return
 end