Support ScalarNonlinearFunction

Project.toml

@@ -12,6 +12,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 [compat]
 Graphs = "1"
 JuMP = "1"
+MathOptInterface = "1"
 julia = "1"
 
 [extras]

src/identify_variables.jl

@@ -28,12 +28,6 @@ import MathOptInterface as MOI
 import MathProgIncidence: get_equality_constraints
 
 
-# TODO: This file implements functions that filter duplicates from the
-# vectors of identified variables. It may be useful at somepoint to
-# identify variables, preserving duplicates. This can be implemented
-# when/if there is a need.
-
-
 """
     identify_unique_variables(constraints::Vector)::Vector{JuMP.VariableRef}
 
@@ -207,50 +201,149 @@ function identify_unique_variables(
     return _filter_duplicates(refs)
 end
 
-
 """
     identify_unique_variables(fcn)::Vector{JuMP.VariableIndex}
 
 Return the variables that appear in the provided MathOptInterface function.
 
+No variable will appear more than once.
+
 # Implementation
-Only `ScalarQuadraticFunction` and `ScalarAffineFunction` are supported.
-This can be changed there is demand for other functions. For each type of
-supported function, the `_get_variable_terms` function should be defined.
-Then, for the type of each term, an additional `identify_unique_variables`
-function should be implemented.
+Only `ScalarNonlinearFunction`, `ScalarQuadraticFunction`, and
+`ScalarAffineFunction` are supported. This can be changed there is demand for
+other functions.
+These methods are implemented by first identifying all variables that
+participate in the function, then filtering out duplicate variables.
 
 """
 function identify_unique_variables(
-    fcn::Union{MOI.ScalarQuadraticFunction, MOI.ScalarAffineFunction},
+    fcn::Union{
+        MOI.ScalarAffineFunction,
+        MOI.ScalarQuadraticFunction,
+        MOI.ScalarNonlinearFunction,
+    },
 )::Vector{MOI.VariableIndex}
-    variables = Vector{MOI.VariableIndex}()
-    for terms in _get_variable_terms(fcn)
-        for term in terms
-            for var in identify_unique_variables(term)
-                push!(variables, var)
-            end
-        end
-    end
+    variables = _identify_variables(fcn)
     return _filter_duplicates(variables)
 end
 
+# This method is used to handle function-in-set constraints.
+function identify_unique_variables(
+    var::MOI.VariableIndex
+)::Vector{MOI.VariableIndex}
+    return [var]
+end
+
+# NOTE: We will get a MethodError if this is called with a non-vector,
+# non-Affine/Quadratic/Nonlinear function. Should probably implement
+# some default method to catch that.
 function identify_unique_variables(
-    fcn::T
-)::Vector{MOI.VariableIndex} where {T<:MOI.AbstractVectorFunction}
+    fcn::MOI.AbstractVectorFunction
+)::Vector{MOI.VariableIndex}
     throw(TypeError(
         fcn,
-        Union{MOI.ScalarQuadraticFunction, MOI.ScalarAffineFunction},
+        Union{
+            MOI.ScalarAffineFunction,
+            MOI.ScalarQuadraticFunction,
+            MOI.ScalarNonlinearFunction,
+        },
         typeof(fcn),
     ))
 end
 
-function identify_unique_variables(
-    var::MOI.VariableIndex
+"""
+    _identify_variables(fcn)::Vector{JuMP.VariableIndex}
+
+Return all variables that appear in the provided MathOptInterface function.
+
+Duplicates may be present in the resulting vector of variables.
+If there is a use case, these methods can be made public.
+
+# Implementation
+Implemented for `ScalarAffineFunction`, `ScalarQuadraticFunction`, and
+`ScalarNonlinearFunction`. Relies on an underlying _collect_variables!
+method that (potentially recursively) builds up a vector of variables
+in-place.
+
+"""
+function _identify_variables(
+    fcn::Union{
+        MOI.ScalarAffineFunction,
+        MOI.ScalarQuadraticFunction,
+        MOI.ScalarNonlinearFunction,
+    },
 )::Vector{MOI.VariableIndex}
-    return [var]
+    variables = Vector{MOI.VariableIndex}()
+    _collect_variables!(variables, fcn)
+    return variables
 end
 
+"""
+    _collect_variables!(variables, fcn)::Vector{JuMP.VariableIndex}
+
+Add variables from `fcn` to the `variables` vector
+
+# Implementation
+Implemented for `ScalarAffineFunction`, `ScalarQuadraticFunction`, and
+`ScalarNonlinearFunction`. For affine and quadratic functions, we iterate
+over terms with the `_get_variable_terms` function. For nonlinear functions,
+we recurse until pushing a root node onto the variable stack (or hitting
+an affine/quadratic function). Methods may need to be added if more node types
+are added to `ScalarNonlinearFunction` in the future.
+
+"""
+function _collect_variables!(
+    variables::Vector{MOI.VariableIndex},
+    fcn::MOI.ScalarNonlinearFunction,
+)::Vector{MOI.VariableIndex}
+    for arg in fcn.args
+        _collect_variables!(variables, arg)
+    end
+    return variables
+end
+
+function _collect_variables!(
+    variables::Vector{MOI.VariableIndex},
+    fcn::Union{MOI.ScalarQuadraticFunction, MOI.ScalarAffineFunction},
+)::Vector{MOI.VariableIndex}
+    for terms in _get_variable_terms(fcn)
+        for term in terms
+            _collect_variables!(variables, term)
+        end
+    end
+    return variables
+end
+
+function _collect_variables!(
+    variables::Vector{MOI.VariableIndex},
+    var::MOI.VariableIndex,
+)::Vector{MOI.VariableIndex}
+    push!(variables, var)
+    return variables
+end
+
+function _collect_variables!(
+    variables::Vector{MOI.VariableIndex},
+    var::Float64,
+)::Vector{MOI.VariableIndex}
+    return variables
+end
+
+function _collect_variables!(
+    variables::Vector{MOI.VariableIndex},
+    term::MOI.ScalarAffineTerm,
+)::Vector{MOI.VariableIndex}
+    push!(variables, term.variable)
+    return variables
+end
+
+function _collect_variables!(
+    variables::Vector{MOI.VariableIndex},
+    term::MOI.ScalarQuadraticTerm,
+)::Vector{MOI.VariableIndex}
+    push!(variables, term.variable_1, term.variable_2)
+    return variables
+end
 
 """
     _get_variable_terms(fcn)
@@ -301,7 +394,6 @@ function identify_unique_variables(
     end
 end
 
-
 """
     _filter_duplicates
 
@@ -324,7 +416,6 @@ function _filter_duplicates(
     return filtered
 end
 
-
 function _filter_duplicates(
     variables::Vector{JuMP.VariableRef},
 )::Vector{JuMP.VariableRef}

test/identify_variables.jl

@@ -26,8 +26,8 @@ using MathProgIncidence: identify_unique_variables
 # Local import of JuMP models for testing
 include("models.jl") # make_degenerate_flow_model, make_simple_model
 
-function test_linear()
-    m = make_degenerate_flow_model()
+function test_linear(model_function=make_degenerate_flow_model)
+    m = model_function()
     variables = identify_unique_variables(m[:sum_comp_eqn])
     # What is happening when we hash a VariableRef? I.e. how does
     # the model get hashed?
@@ -36,16 +36,16 @@ function test_linear()
     @test(pred_var_set == Set(variables))
 end
 
-function test_quadratic()
-    m = make_degenerate_flow_model()
+function test_quadratic(model_function=make_degenerate_flow_model)
+    m = model_function()
     variables = identify_unique_variables(m[:comp_dens_eqn][1])
     pred_var_set = Set([m[:rho], m[:x][1]])
     @test(length(variables) == length(pred_var_set))
     @test(pred_var_set == Set(variables))
 end
 
-function test_nonlinear()
-    m = make_degenerate_flow_model()
+function test_nonlinear(model_function=make_degenerate_flow_model)
+    m = model_function()
     variables = identify_unique_variables(m[:bulk_dens_eqn])
     pred_var_set = Set([m[:rho], m[:x][1], m[:x][2], m[:x][3]])
     @test(length(variables) == length(pred_var_set))
@@ -61,8 +61,8 @@ function test_nonlinear_with_potential_duplicates()
     @test(Set(variables) == Set([m[:var][1], m[:var][2]]))
 end
 
-function test_several_equalities()
-    m = make_degenerate_flow_model()
+function test_several_equalities(model_function=make_degenerate_flow_model)
+    m = model_function()
     constraints = [m[:comp_flow_eqn][1], m[:sum_comp_eqn], m[:bulk_dens_eqn]]
     variables = identify_unique_variables(constraints)
     pred_var_set = Set([
@@ -72,8 +72,8 @@ function test_several_equalities()
     @test(pred_var_set == Set(variables))
 end
 
-function test_several_constraints_with_ineq()
-    m = make_degenerate_flow_model()
+function test_several_constraints_with_ineq(model_function=make_degenerate_flow_model)
+    m = model_function()
     @JuMP.constraint(m, ineq1, m[:flow_comp][2] >= 1.0)
     @JuMP.constraint(m, ineq2, m[:flow_comp][1]^2 + m[:flow_comp][3]^2 <= 1.0)
     constraints = [
@@ -98,8 +98,8 @@ function test_several_constraints_with_ineq()
     @test(pred_var_set == Set(variables))
 end
 
-function test_model()
-    m = make_degenerate_flow_model()
+function test_model(model_function=make_degenerate_flow_model)
+    m = model_function()
     @JuMP.variable(m, dummy)
     @JuMP.NLconstraint(m, dummy^3.0 <= 5)
     # Note that include_inequalities=false by default.
@@ -118,8 +118,8 @@ function test_model()
     @test(pred_var_set == Set(variables))
 end
 
-function test_model_with_ineq()
-    m = make_degenerate_flow_model()
+function test_model_with_ineq(model_function=make_degenerate_flow_model)
+    m = model_function()
     @JuMP.variable(m, dummy)
     @JuMP.NLconstraint(m, dummy^3.0 <= 5)
     variables = identify_unique_variables(m, include_inequality=true)
@@ -149,8 +149,8 @@ function test_function_with_variable_squared()
     @test(Set(variables) == Set([m[:dummy1].index, m[:dummy2].index]))
 end
 
-function test_model_bad_constr()
-    m = make_degenerate_flow_model()
+function test_model_bad_constr(model_function=make_degenerate_flow_model)
+    m = model_function()
     @JuMP.variable(m, dummy[1:2])
     @JuMP.constraint(m, vectorcon, dummy in MOI.Nonnegatives(2))
     @test_throws(
@@ -159,8 +159,8 @@ function test_model_bad_constr()
     )
 end
 
-function test_model_bad_constr_no_ineq()
-    m = make_degenerate_flow_model()
+function test_model_bad_constr_no_ineq(model_function=make_degenerate_flow_model)
+    m = model_function()
     @JuMP.variable(m, dummy[1:2])
     @JuMP.constraint(m, vectorcon, dummy in MOI.Nonnegatives(2))
     # Note that we don't throw an error because we don't attempt to
@@ -208,8 +208,8 @@ function test_two_constraints_same_type()
     @test Set(vars) == pred_var_set
 end
 
-function test_variables_in_inequalities()
-    m = make_simple_model()
+function test_variables_in_inequalities(model_function=make_simple_model)
+    m = model_function()
     @JuMP.variable(m, y >= 1)
     x = m[:x]
     @JuMP.objective(m, Min, x[1] + 2*x[2] + 3*x[3] + y^2)
@@ -224,18 +224,39 @@ end
 
 @testset "get-equality" begin
     test_linear()
+    test_linear(make_degenerate_flow_model_with_ScalarNonlinearFunction)
+
     test_quadratic()
+    test_quadratic(make_degenerate_flow_model_with_ScalarNonlinearFunction)
+
     test_nonlinear()
+    test_nonlinear(make_degenerate_flow_model_with_ScalarNonlinearFunction)
+
     test_nonlinear_with_potential_duplicates()
+
     test_several_equalities()
+    test_several_equalities(make_degenerate_flow_model_with_ScalarNonlinearFunction)
+
     test_several_constraints_with_ineq()
+    test_several_constraints_with_ineq(make_degenerate_flow_model_with_ScalarNonlinearFunction)
+
     test_model()
+    test_model(make_degenerate_flow_model_with_ScalarNonlinearFunction)
+
     test_model_with_ineq()
+    test_model_with_ineq(make_degenerate_flow_model_with_ScalarNonlinearFunction)
+
     test_model_bad_constr()
+    test_model_bad_constr(make_degenerate_flow_model_with_ScalarNonlinearFunction)
+
     test_model_bad_constr_no_ineq()
+    test_model_bad_constr_no_ineq(make_degenerate_flow_model_with_ScalarNonlinearFunction)
+
     test_function_with_variable_squared()
     test_fixing_constraint()
     test_inequality_with_bounds()
     test_two_constraints_same_type()
+
     test_variables_in_inequalities()
+    test_variables_in_inequalities(make_simple_model_with_ScalarNonlinearFunction)
 end