Support ScalarNonlinearFunction

[Robbybp]

This PR still needs:

• Tests
• Updates to examples/readme Since the README is always public, and the documentation is for the current main branch, I'll defer these until a release, as I don't want the docs to diverge from the latest release.
• Update to JuMP version in compat
• [Nonlinear] add new nonlinear interface jump-dev/JuMP.jl#3106 to get merged
• Possibly documentation that identify_unique_variables supports Float64 (to handle arbitrary nodes in ScalarNonlinearFunction) Marking as complete as identify_unique_variables no longer has to support Float64
• Benchmarks on pglib instances (e.g. Performance issue in Dulmage-Mendelsohn #8) with WIP: update to new JuMP nonlinear syntax PowerModels.jl#858 to make sure performance is not heavily affected

[Robbybp]

It might also be nice to include a test with ScalarNonlinearFunction and a user-defined function, just to make sure/demonstrate that it works.

[Robbybp]

Benchmarks on pglib. This is the Rosetta-OPF implementation, with NLconstraint replaced with constraint.

Instance	Graph construction time	Dulmage-Mendelsohn time	N. variables
pglib_opf_case3_lmbd.m	4.45e-04	2.49e-04	19
pglib_opf_case5_pjm.m	7.31e-04	3.64e-04	35
pglib_opf_case14_ieee.m	1.85e-03	6.22e-04	109
pglib_opf_case24_ieee_rts.m	3.42e-03	8.95e-04	201
pglib_opf_case30_as.m	3.40e-03	9.21e-04	225
pglib_opf_case30_ieee.m	3.32e-03	9.10e-04	225
pglib_opf_case39_epri.m	3.83e-03	1.02e-03	263
pglib_opf_case57_ieee.m	6.83e-03	1.68e-03	435
pglib_opf_case60_c.m	7.00e-03	1.78e-03	473
pglib_opf_case73_ieee_rts.m	1.07e-02	2.12e-03	627
pglib_opf_case118_ieee.m	1.58e-02	1.69e-02	981
pglib_opf_case89_pegase.m	1.64e-02	3.59e-03	1019
pglib_opf_case200_activ.m	2.01e-02	4.85e-03	1381
pglib_opf_case179_goc.m	2.09e-02	2.35e-02	1411
pglib_opf_case162_ieee_dtc.m	2.34e-02	5.42e-03	1461
pglib_opf_case300_ieee.m	3.32e-02	8.53e-03	2245
pglib_opf_case240_pserc.m	3.60e-02	8.98e-03	2273
pglib_opf_case500_goc.m	6.01e-02	1.48e-02	3913
pglib_opf_case588_sdet.m	5.57e-02	1.64e-02	3921
pglib_opf_case793_goc.m	7.53e-02	2.37e-02	5239
pglib_opf_case1354_pegase.m	1.61e-01	9.14e-02	10673
pglib_opf_case1888_rte.m	2.00e-01	5.83e-02	13901
pglib_opf_case1951_rte.m	2.12e-01	5.92e-02	14287
pglib_opf_case2383wp_k.m	3.22e-01	6.77e-02	16351
pglib_opf_case2312_goc.m	2.46e-01	8.00e-02	16677
pglib_opf_case2000_goc.m	3.88e-01	9.59e-02	18533
pglib_opf_case2736sp_k.m	2.69e-01	8.68e-02	18549
pglib_opf_case2737sop_k.m	2.67e-01	8.53e-02	18551
pglib_opf_case2746wp_k.m	2.72e-01	8.59e-02	18609
pglib_opf_case2746wop_k.m	2.71e-01	9.27e-02	18721
pglib_opf_case3012wp_k.m	2.86e-01	9.20e-02	20313
pglib_opf_case2848_rte.m	3.02e-01	9.74e-02	20801
pglib_opf_case2868_rte.m	4.29e-01	9.41e-02	20969
pglib_opf_case3120sp_k.m	2.98e-01	1.06e-01	21013
pglib_opf_case2853_sdet.m	4.40e-01	9.04e-02	21391
pglib_opf_case3022_goc.m	3.34e-01	1.04e-01	22585
pglib_opf_case3375wp_k.m	3.47e-01	1.02e-01	23393
pglib_opf_case2869_pegase.m	4.88e-01	1.07e-01	24067
pglib_opf_case2742_goc.m	3.83e-01	1.16e-01	24177
pglib_opf_case4661_sdet.m	4.85e-01	1.69e-01	33311
pglib_opf_case3970_goc.m	6.99e-01	2.03e-01	34505
pglib_opf_case4020_goc.m	5.62e-01	2.21e-01	35993
pglib_opf_case4917_goc.m	7.24e-01	2.13e-01	36739
pglib_opf_case4601_goc.m	7.29e-01	2.44e-01	37999
pglib_opf_case4837_goc.m	7.81e-01	2.73e-01	40735
pglib_opf_case4619_goc.m	8.41e-01	2.75e-01	41839
pglib_opf_case6468_rte.m	8.96e-01	2.94e-01	48937
pglib_opf_case6470_rte.m	8.93e-01	2.63e-01	48961
pglib_opf_case6495_rte.m	7.25e-01	4.48e-01	49067
pglib_opf_case6515_rte.m	9.32e-01	3.06e-01	49179
pglib_opf_case10000_goc.m	1.86e+00	1.50e+00	72773
pglib_opf_case8387_pegase.m	1.40e+00	4.39e-01	75019
pglib_opf_case9241_pegase.m	1.58e+00	5.80e-01	82679
pglib_opf_case9591_goc.m	1.52e+00	8.85e-01	82843
pglib_opf_case10480_goc.m	1.86e+00	7.88e-01	95197
pglib_opf_case13659_pegase.m	2.00e+00	7.91e-01	109187
pglib_opf_case19402_goc.m	3.41e+00	1.97e+00	177621
pglib_opf_case24464_goc.m	3.96e+00	2.24e+00	200193
pglib_opf_case30000_goc.m	3.59e+00	1.83e+00	201573

Interestingly, the graph construction time here is slower than what I report here. Maybe I'm not parsing ScalarNonlinearFunction as efficiently as I could be? (I use a pretty naive recursive function to identify variables, and filter duplicate variables eagerly.)

[Robbybp]

JuMP has released the nonlinear expression system, but this needs some profiling (as well as updated examples) before it can be merged. In particular, we see a 4x slowdown in incidence graph construction on the 20-30k bus pglib instances when using ScalarNonlinearFunction. This is probably due to the method used to "flatten" the expression graph, but I'm not 100% sure what needs to be done to fix this.

[Robbybp]

A large portion of the overhead appears to be in identify_unique_variables(fcn::ScalarNonlinearFunction):

  1457       583 @MathProgIncidence/src/identify_variables.jl               228 identify_unique_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
  1625        32 @MathProgIncidence/src/identify_variables.jl               155 identify_unique_variables(constraint::JuMP.ConstraintRef{JuMP.Model, MathOptInterface.ConstraintIndex{MathOptInterface.ScalarQuadrat...
  1706         2 @MathProgIncidence/src/identify_variables.jl               132 identify_unique_variables(constraint::JuMP.ConstraintRef{JuMP.Model, MathOptInterface.ConstraintIndex{MathOptInterface.ScalarQuadrat...
  1846         0 @MathProgIncidence/src/incidence_graph.jl                  138 #get_bipartite_incidence_graph#14
  2000         0 @MathProgIncidence/src/incidence_graph.jl                  122 get_bipartite_incidence_graph
  2039         0 @MathProgIncidence/src/interface.jl                        134 MathProgIncidence.IncidenceGraphInterface(m::JuMP.Model)
  2039         0 @MathProgIncidence/src/interface.jl                        134 #IncidenceGraphInterface#61
  2041         0 @Base/client.jl                                            522 _start()
  2041         0 @Base/client.jl                                            307 exec_options(opts::Base.JLOptions)
  2041         0 @Base/Base.jl                                              457 include(mod::Module, _path::String)
  2041         0 @Base/loading.jl                                          1963 _include(mapexpr::Function, mod::Module, _path::String)
  2041         0 @Base/loading.jl                                          1903 include_string(mapexpr::typeof(identity), mod::Module, code::String, filename::String)
  2041         2 @Base/boot.jl                                              370 eval

It's a little difficult to tell exactly what is taking the most time within this function:

 12╎    ╎    ╎   825  @MathProgIncidence/src/identify_variables.jl:155; identify_unique_variables(constraint::JuMP.ConstraintRef{JuMP.Model, MathOptInterface.ConstraintIndex{Mat...
 87╎    ╎    ╎    726  @MathProgIncidence/src/identify_variables.jl:228; identify_unique_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
 73╎    ╎    ╎     592  @MathProgIncidence/src/identify_variables.jl:228; identify_unique_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
100╎    ╎    ╎    ╎ 438  @MathProgIncidence/src/identify_variables.jl:228; identify_unique_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
 47╎    ╎    ╎    ╎  192  @MathProgIncidence/src/identify_variables.jl:228; identify_unique_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
   ╎    ╎    ╎    ╎   74   @MathProgIncidence/src/identify_variables.jl:246; identify_unique_variables(fcn::MathOptInterface.ScalarAffineFunction{Float64})
   ╎    ╎    ╎    ╎  68   @MathProgIncidence/src/identify_variables.jl:246; identify_unique_variables(fcn::MathOptInterface.ScalarQuadraticFunction{Float64})
  1╎    ╎    ╎    58   @MathProgIncidence/src/identify_variables.jl:232; identify_unique_variables(fcn::MathOptInterface.ScalarNonlinearFunction)

 25╎    ╎    ╎   824  @MathProgIncidence/src/identify_variables.jl:155; identify_unique_variables(constraint::JuMP.ConstraintRef{JuMP.Model, MathOptInterface.ConstraintIndex{Mat...
 66╎    ╎    ╎    745  @MathProgIncidence/src/identify_variables.jl:228; identify_unique_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
 85╎    ╎    ╎     606  @MathProgIncidence/src/identify_variables.jl:228; identify_unique_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
 86╎    ╎    ╎    ╎ 405  @MathProgIncidence/src/identify_variables.jl:228; identify_unique_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
 40╎    ╎    ╎    ╎  194  @MathProgIncidence/src/identify_variables.jl:228; identify_unique_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
   ╎    ╎    ╎    ╎   75   @MathProgIncidence/src/identify_variables.jl:246; identify_unique_variables(fcn::MathOptInterface.ScalarQuadraticFunction{Float64})
   ╎    ╎    ╎    ╎    59   @MathProgIncidence/src/identify_variables.jl:332; _filter_duplicates(indices::Vector{MathOptInterface.VariableIndex})
   ╎    ╎    ╎    ╎     59   @Base/set.jl:45; Set
  2╎    ╎    ╎    ╎    ╎ 59   @Base/dict.jl:70; Dict{MathOptInterface.VariableIndex, Nothing}()
   ╎    ╎    ╎    ╎  50   @MathProgIncidence/src/identify_variables.jl:246; identify_unique_variables(fcn::MathOptInterface.ScalarQuadraticFunction{Float64})
   ╎    ╎    ╎    ╎ 72   @MathProgIncidence/src/identify_variables.jl:232; identify_unique_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
   ╎    ╎    ╎    ╎  54   @MathProgIncidence/src/identify_variables.jl:332; _filter_duplicates(indices::Vector{MathOptInterface.VariableIndex})
   ╎    ╎    ╎    ╎   54   @Base/set.jl:45; Set
  2╎    ╎    ╎    ╎    53   @Base/dict.jl:70; Dict{MathOptInterface.VariableIndex, Nothing}()

_filter_duplicates appears to be at least some of it, but a non-trivial portion seems to be the overhead of dynamically building the array of variables in identify_unique_variables itself.

[Robbybp]

We can get rid of the _filter_duplicates overhead by not eagerly filtering variables as soon as we have them, and just filtering duplicates once at the end:

    82         0 @Base/array.jl                                            1061 push!
    83         0 @Base/array.jl                                             126 vect
    83         0 @Base/array.jl                                             674 _array_for
    89         0 @Base/array.jl                                             671 _array_for
    89         0 @Base/abstractarray.jl                                     883 similar
    89         0 @Base/abstractarray.jl                                     884 similar
    92         0 @Base/array.jl                                            1061 push!(a::Vector{MathOptInterface.VariableIndex}, item::MathOptInterface.VariableIndex)
   108         0 @MathProgIncidence/src/get_equality.jl                     183 get_equality_constraints
   108         0 @JuMP/src/constraints.jl                                  1422 all_constraints
   108        16 @JuMP/src/constraints.jl                                  1429 all_constraints(model::JuMP.Model; include_variable_in_set_constraints::Bool)
   118         0 @Base/boot.jl                                              486 Array
   121        63 @MathProgIncidence/src/identify_variables.jl               239 _identify_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
   154       154 @MathProgIncidence/src/identify_variables.jl               240 _identify_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
   174         0 @MathProgIncidence/src/incidence_graph.jl                  128 #get_bipartite_incidence_graph#14
   174         0 @MathProgIncidence/src/get_equality.jl                     313 _get_constraints
   174         0 @MathProgIncidence/src/get_equality.jl                     324 _get_constraints(model::JuMP.Model; include_inequality::Bool, include_active_inequalities::Bool, tolerance::Float64)
   175       175 @Base/array.jl                                            1014 _growend!
   217       217 @Base/boot.jl                                              477 Array
   730        11 @MathProgIncidence/src/identify_variables.jl                79 identify_unique_variables(constraints::Vector{JuMP.ConstraintRef})
   731         9 @MathProgIncidence/src/incidence_graph.jl                  171 get_bipartite_incidence_graph(constraints::Vector{JuMP.ConstraintRef}, variables::Vector{JuMP.VariableRef})
   746         0 @MathProgIncidence/src/incidence_graph.jl                  144 get_bipartite_incidence_graph
   804         0 @MathProgIncidence/src/incidence_graph.jl                  147 get_bipartite_incidence_graph
  1004       666 @MathProgIncidence/src/identify_variables.jl               238 _identify_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
  1138         2 @MathProgIncidence/src/identify_variables.jl               229 identify_unique_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
  1329        69 @MathProgIncidence/src/identify_variables.jl               158 identify_unique_variables(constraint::JuMP.ConstraintRef{JuMP.Model, MathOptInterface.ConstraintIndex{MathOptInterface.ScalarQuadrat...
  1437         2 @MathProgIncidence/src/identify_variables.jl               135 identify_unique_variables(constraint::JuMP.ConstraintRef{JuMP.Model, MathOptInterface.ConstraintIndex{MathOptInterface.ScalarQuadrat...
  1550         0 @MathProgIncidence/src/incidence_graph.jl                  138 #get_bipartite_incidence_graph#14
  1724         0 @MathProgIncidence/src/incidence_graph.jl                  122 get_bipartite_incidence_graph
  1750         0 @MathProgIncidence/src/interface.jl                        134 MathProgIncidence.IncidenceGraphInterface(m::JuMP.Model)
  1750         0 @MathProgIncidence/src/interface.jl                        134 #IncidenceGraphInterface#61
  1772         0 @Base/client.jl                                            522 _start()
  1772         0 @Base/client.jl                                            307 exec_options(opts::Base.JLOptions)
  1772         0 @Base/Base.jl                                              457 include(mod::Module, _path::String)
  1772         0 @Base/loading.jl                                          1963 _include(mapexpr::Function, mod::Module, _path::String)
  1772         0 @Base/loading.jl                                          1903 include_string(mapexpr::typeof(identity), mod::Module, code::String, filename::String)
  1772        22 @Base/boot.jl                                              370 eval

The comparison between the two approaches is now:

Case	Nonlinear implementation	Time to construct incidence graph
pglib_opf_case30000_goc.m	nonlinear-expr	2.40e+00
pglib_opf_case30000_goc.m	@NL*	8.63e-01

Breakdown of waht is happening inside of _identify_variables(fcn::ScalarNonlinearFunction):

   ╎    ╎    1550 @MathProgIncidence/src/incidence_graph.jl:138; #get_bipartite_incidence_graph#14
   ╎    ╎     746  @MathProgIncidence/src/incidence_graph.jl:144; get_bipartite_incidence_graph
 11╎    ╎    ╎ 730  @MathProgIncidence/src/identify_variables.jl:79; identify_unique_variables(constraints::Vector{JuMP.ConstraintRef})
  1╎    ╎    ╎  713  @MathProgIncidence/src/identify_variables.jl:135; identify_unique_variables(constraint::JuMP.ConstraintRef{JuMP.Model, MathOptInterface.ConstraintIndex{Mat...
 39╎    ╎    ╎   675  @MathProgIncidence/src/identify_variables.jl:158; identify_unique_variables(constraint::JuMP.ConstraintRef{JuMP.Model, MathOptInterface.ConstraintIndex{Mat...
   ╎    ╎    ╎    574  @MathProgIncidence/src/identify_variables.jl:229; identify_unique_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
 94╎    ╎    ╎     480  @MathProgIncidence/src/identify_variables.jl:238; _identify_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
 87╎    ╎    ╎    ╎ 335  @MathProgIncidence/src/identify_variables.jl:238; _identify_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
 70╎    ╎    ╎    ╎  194  @MathProgIncidence/src/identify_variables.jl:238; _identify_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
 37╎    ╎    ╎    ╎   68   @MathProgIncidence/src/identify_variables.jl:238; _identify_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
   ╎    ╎    ╎    ╎   28   @MathProgIncidence/src/identify_variables.jl:303; _identify_variables(var::Float64)
   ╎    ╎    ╎    ╎  28   @MathProgIncidence/src/identify_variables.jl:251; _identify_variables(fcn::MathOptInterface.ScalarQuadraticFunction{Float64})
   ╎    ╎    ╎    ╎   28   @MathProgIncidence/src/identify_variables.jl:361; _identify_variables
   ╎    ╎    ╎    ╎    28   @Base/array.jl:126; vect
   ╎    ╎    ╎    ╎     28   @Base/array.jl:674; _array_for
   ╎    ╎    ╎    ╎    ╎ 28   @Base/array.jl:671; _array_for
   ╎    ╎    ╎    ╎    ╎  28   @Base/abstractarray.jl:883; similar
   ╎    ╎    ╎    ╎    ╎   28   @Base/abstractarray.jl:884; similar
   ╎    ╎    ╎    ╎    ╎    28   @Base/boot.jl:486; Array
 28╎    ╎    ╎    ╎    ╎     28   @Base/boot.jl:477; Array

We probably need to be more efficient about recursively building up the array of variables (i.e. build it in-place instead of allocating new arrays all the time).

[Robbybp]

Once we build up the variable array in-place, the comparison starts to look a lot better:

Case	Nonlinear implementation	Time to construct incidence graph
pglib_opf_case30000_goc.m	nonlinear-expr	1.32e+00
pglib_opf_case30000_goc.m	@NL*	8.21e-01

We can still see a good amount of overhead in the new _collect_variables! function:

    95        95 @Base/boot.jl                                              477 Array
   146         0 @MathProgIncidence/src/incidence_graph.jl                  128 #get_bipartite_incidence_graph#14
   146         0 @MathProgIncidence/src/get_equality.jl                     313 _get_constraints
   146         0 @MathProgIncidence/src/get_equality.jl                     324 _get_constraints(model::JuMP.Model; include_inequality::Bool, include_active_inequalities::Bool, tolerance::Float64)
   244        11 @MathProgIncidence/src/incidence_graph.jl                  171 get_bipartite_incidence_graph(constraints::Vector{JuMP.ConstraintRef}, variables::Vector{JuMP.VariableRef})
   268       239 @MathProgIncidence/src/identify_variables.jl               246 _collect_variables!(variables::Vector{MathOptInterface.VariableIndex}, fcn::MathOptInterface.ScalarNonlinearFunction)
   271         0 @MathProgIncidence/src/identify_variables.jl               237 _identify_variables
   277         2 @MathProgIncidence/src/identify_variables.jl                79 identify_unique_variables(constraints::Vector{JuMP.ConstraintRef})
   278         0 @MathProgIncidence/src/identify_variables.jl               229 identify_unique_variables(fcn::MathOptInterface.ScalarNonlinearFunction)
   294         0 @MathProgIncidence/src/incidence_graph.jl                  144 get_bipartite_incidence_graph
   336         0 @MathProgIncidence/src/incidence_graph.jl                  147 get_bipartite_incidence_graph
   407        14 @MathProgIncidence/src/identify_variables.jl               158 identify_unique_variables(constraint::JuMP.ConstraintRef{JuMP.Model, MathOptInterface.ConstraintIndex{MathOptInterface.ScalarQuadrat...
   508         1 @MathProgIncidence/src/identify_variables.jl               135 identify_unique_variables(constraint::JuMP.ConstraintRef{JuMP.Model, MathOptInterface.ConstraintIndex{MathOptInterface.ScalarQuadrat...
   630         0 @MathProgIncidence/src/incidence_graph.jl                  138 #get_bipartite_incidence_graph#14
   776         0 @MathProgIncidence/src/incidence_graph.jl                  122 get_bipartite_incidence_graph
   798         0 @MathProgIncidence/src/interface.jl                        134 MathProgIncidence.IncidenceGraphInterface(m::JuMP.Model)
   798         0 @MathProgIncidence/src/interface.jl                        134 #IncidenceGraphInterface#61
   825         0 @Base/client.jl                                            522 _start()
   825         0 @Base/client.jl                                            307 exec_options(opts::Base.JLOptions)
   825         0 @Base/Base.jl                                              457 include(mod::Module, _path::String)
   825         0 @Base/loading.jl                                          1963 _include(mapexpr::Function, mod::Module, _path::String)
   825         0 @Base/loading.jl                                          1903 include_string(mapexpr::typeof(identity), mod::Module, code::String, filename::String)
   825        27 @Base/boot.jl                                              370 eval
Total snapshots: 1195. Utilization: 100% across all threads and tasks. Use the `groupby` kwarg to break down by thread and/or task.

We now have a ~50% performance hit in incidence graph construction when using ScalarNonlinearFunction. I think this makes sense, because we now have to construct the flattened data structure ourselves, instead of it already being constructed in model.nlp_model.constraints[idx].expression. Putting together the full workflow, we have:

Case	Nonlinear implementation	Model build time	Incidence graph time	Dulmage-Mendelsohn time	Total time
pglib_opf_case30000_goc.m	nonlinear-expr	9.09e+00	1.48e+00	2.28e+00	1.29e+01
pglib_opf_case30000_goc.m	@NL*	6.40e+00	9.31e-01	2.39e+00	9.72e+00

With ScalarNonlinearFunction, we are about 30% slower overall, and the majority of that is in model construction. This looks fine.