JuMP.jl icon indicating copy to clipboard operation
JuMP.jl copied to clipboard

Nonlinear subexpressions

Open odow opened this issue 10 months ago • 4 comments

Background

  • JuMP has long supported nonlinear subexpressions created with @NLexpression.
  • The AD system in MOI.Nonlinear has support for representing expressions, and dealing with them efficiently
  • When we created ScalarNonlinearFunction, we did not create an analogous ScalarNonlinearExpression object.
  • In most cases, nonlinear expressions make little difference, or are a minor win.
  • An upside and a downside is that it's up to users to decide which expressions should be common. I have seen, anecdotally, that users on the forum either use no @NLexpression, or they made every possible expression a @NLexpression
  • Best performance is to judiciously use common subexpressions where it would make a difference.
  • The best real-world example we have is @mitchphillipson's MCP model: https://github.com/jump-dev/JuMP.jl/issues/3729#issuecomment-2060586587

The question is how to achieve this in JuMP.

I opened an issue in MOI: https://github.com/jump-dev/MathOptInterface.jl/issues/2488

Short of rewriting much of the MOI.Nonlinear module to use a single DAG of expressions (instead of the current tree), we could pass the expressions through to MOI, and then attempt to detect common subexpressions. This would rely on a heuristic of when it was beneficial to do so.

A simpler approach would be to add a "nonlinear expression" set to MOI (and JuMP), just like we've done for Parameter.

A crude API would be:

@variable(model, y in CommonSubexpression())
@constraint(model, [y; f(x)] in CommonSubexpression())

with the fallback to a bridge

@variable(model, y)
@constraint(model, y == f(x))

and maybe one for MCP:

@variable(model, y)
@constraint(model, f(x) - y ⟂ y)

We could come up with nicer syntax, for example:

@expression(model, y, f(x), NonlinearSubexpression())

odow avatar Apr 25 '24 19:04 odow

@common_expression(model, y, f(x))

@variable(model, y)
@constraint(model, y :== f(x))

odow avatar May 23 '24 19:05 odow

What is stopping @expression(model, y, f(x)) from adopting the common-subexpression behavior you're describing? Would this require breaking changes to MOI?

Robbybp avatar May 24 '24 18:05 Robbybp

What is stopping @expression(model, y, f(x)) from adopting the common-subexpression behavior you're describing?

Yes, it's a breaking change. We discussed this on the monthly developer call. @mlubin is in favor of looking into a :subexpression node. I need to just try out some options.

odow avatar May 25 '24 00:05 odow

Some related discussion from https://github.com/jump-dev/MathOptInterface.jl/issues/2488:

Some possible options:

  1. Create a single unified expression graph (as opposed to the current expression tree). Change quite a lot of how ReverseAD processes functions etc.
  2. Don't do anything. Add a new decision variable with y = f(x) for expressions. This makes things very explicit.
  3. Add something new type/index to JuMP/MOI. But we're breaking the abstraction quite a bit.
  4. Add a univariate :subexpression operator to MOI. So: @expression(model, expr, subexpression(sin(x) + 1)). Solvers would see this during parsing, and could choose to ignore, or store it in a common expression.

I don't have a strong opinion yet. But our current approach requires us to smack the modeler on the hand and tell them they're holding it wrong and that they should do (2): https://github.com/jump-dev/JuMP.jl/issues/3729#issuecomment-2060586587

odow avatar May 25 '24 00:05 odow