JuMP Developers Meetup

Description: This notebook gives a quick overview of the process for extension packages to work with JuMP's macros, and was done for the JuMP developers meetup event in June 2017. We begin with a few code snippets to illustrate the general concept, and then have a quick peek under the hood of what happens in a JuMP macro call.

Author: Yeesian Ng

License:

This work is licensed under a Creative Commons Attribution-ShareAlike 4.0 International License.

Overview

  • Why Macros?
  • Quick Example
    • src/macros.jl
      • @variable (implementation: 250lines)
      • @constraint (implementation: 150lines)
    • some utilities in JuMP
      • for parsing and modifying expressions
      • "high-level-interface" to MathProgBase
  • What we talk about when we talk about JuMP

Why Macros?

cleaner syntax

You might be currently writing code like

p_fr = PowerModels.getvariable(pm.model, :p)[f_idx]
t_fr = PowerModels.getvariable(pm.model, :t)[f_bus]
t_to = PowerModels.getvariable(pm.model, :t)[t_bus]
[email protected](pm.model, p_fr == -b*(t_fr - t_to))

when you prefer to be writing code like

[email protected](pm, p[f_idx] == -b*(t[f_bus] - t[t_bus]))

some correlation between developers who understand JuMP's Macros

Quick Example

In [1]:
import JuMPChance, Distributions

m = JuMPChance.ChanceModel()
JuMPChance.@indepnormal(m, x, mean=0, var=1)

JuMP.@variable(m, z)
JuMP.@objective(m, Min, z)
JuMP.@constraint(m, z*x >= -1, with_probability=0.95)

JuMP.solve(m, method=:Reformulate, silent=true)
isapprox(
    JuMP.getvalue(z),
    -1/Distributions.quantile(Distributions.Normal(0,1),0.95)
)
Out[1]:
true

Working with JuMP's Macro Utilities

A quick tour through JuMPChance's @indepnormal macro.

(1) Expressions and Evaluation

we begin with the following macro call:

In [2]:
macro indepnormal(m, x, mean, var)
    m, x, mean, var
end

m, x, mean, var = @indepnormal(m, x, mean=1, var=1)
Out[2]:
(:m,:x,:(mean=1),:(var=1))

that we'd like to rewrite into code like the following:

In [3]:
quote
    $x = IndepNormal($m,$mean,$var,$(string(x)))
end
Out[3]:
quote  # In[3], line 2:
    x = IndepNormal(m,mean=1,var=1,"x")
end
In [4]:
m = esc(m)
mean = esc(mean.args[2])
var = esc(var.args[2])

quote
    $(esc(x)) = IndepNormal($m,$mean,$var,$(string(x)))
end
Out[4]:
quote  # In[4], line 6:
    $(Expr(:escape, :x)) = IndepNormal($(Expr(:escape, :m)),$(Expr(:escape, 1)),$(Expr(:escape, 1)),"x")
end

(2) JuMP.buildrefsets

Let's try a more involved example:

In [5]:
m, x, mean, var = @indepnormal(m, ω[i=1:4,j=1:3], mean=1, var=1)
@show m
@show x
@show mean
@show var;
m = :m
x = :(ω[i = 1:4,j = 1:3])
mean = :(mean=1)
var = :(var=1)
In [8]:
dump(x)
Expr
  head: Symbol ref
  args: Array{Any}((3,))
    1: Symbol ω
    2: Expr
      head: Symbol =
      args: Array{Any}((2,))
        1: Symbol i
        2: Expr
          head: Symbol :
          args: Array{Any}((2,))
            1: Int64 1
            2: Int64 4
          typ: Any
      typ: Any
    3: Expr
      head: Symbol =
      args: Array{Any}((2,))
        1: Symbol j
        2: Expr
          head: Symbol :
          args: Array{Any}((2,))
            1: Int64 1
            2: Int64 3
          typ: Any
      typ: Any
  typ: Any

Would our previous macro work here?

In [6]:
quote
    $x = IndepNormal($m,$mean,$var,$(string(x)))
end
Out[6]:
quote  # In[6], line 2:
    ω[i = 1:4,j = 1:3] = IndepNormal(m,mean=1,var=1,"ω[i = 1:4,j = 1:3]")
end
In [7]:
###############################################################################
# buildrefsets
# Unexported. Takes as input an object representing a name, associated index
# sets, and conditions on those sets, for example
# buildrefsets(:(x[i=1:3,[:red,:blue]],k=S; i+k <= 6))
# Used internally in macros to build JuMPContainers and constraints. Returns
#       refcall:  Expr to reference a particular element, e.g. :(x[i,j,k])
#       idxvars:  Index names used in referencing, e.g.g {:i,:j,:k}
#       idxsets:  Index sets for indexing, e.g. {1:3, [:red,:blue], S}
#       idxpairs: Vector of IndexPair
#       condition: Expr containing any condition present for indexing
# Note in particular that it does not actually evaluate the condition, and so
# it returns just the cartesian product of possible indices.

refcall, idxvars, idxsets, idxpairs, condition = JuMP.buildrefsets(x)
@show refcall
@show idxvars
@show idxsets
@show idxpairs
@show condition;
refcall = :(ω[$(Expr(:escape, :i)),$(Expr(:escape, :j))])
idxvars = Any[:i,:j]
idxsets = Any[:($(Expr(:escape, :(1:4)))),:($(Expr(:escape, :(1:3))))]
idxpairs = JuMP.IndexPair[JuMP.IndexPair(:i,:(1:4)),JuMP.IndexPair(:j,:(1:3))]
condition = :(())

In order to construct the string correspond to the variable name "ω[$i,$j]":

In [9]:
varname = JuMP.getname(x)
varstr = :(string($(string(varname)),"["))
for idxvar in idxvars
    push!(varstr.args,:(string($(esc(idxvar)))))
    push!(varstr.args,",")
end
deleteat!(varstr.args,length(varstr.args))
push!(varstr.args,"]")
varstr
Out[9]:
:(string("ω","[",string($(Expr(:escape, :i))),",",string($(Expr(:escape, :j))),"]"))

(3) JuMP.getloopedcode

In [10]:
###############################################################################
# getloopedcode
# Unexported. Takes a bit of code and corresponding looping information and
# returns that code nested in corresponding loops, along with preceding code
# to construct an appropriate container. Input is:
#       c: symbolic representation of name and appropriate indexing sets, if
#          any. E.g. :(myvar) or :(x[i=1:3,[:red,:blue]])
#       code: inner loop code kernel to be nested in the loops
#       condition: a boolean expression to be evaluated before each kernel.
#                  If none, pass :().
#       idxvars: As defined for buildrefsets
#       idxsets: As defined for buildrefsets
#       idxpairs: As defined for buildrefsets
#       sym: A symbol or expression containing the element type of the
#            resulting container, e.g. :AffExpr or :Variable

variable = gensym()
code = :( $(refcall) = IndepNormal($m, $mean, $var, $varstr ) )
JuMP.getloopedcode(variable, code, condition, idxvars, idxsets, idxpairs, :IndepNormal)
Out[10]:
quote  # /Users/yeesian/.julia/v0.5/JuMP/src/macros.jl, line 198:
    ##280 = Array{IndepNormal}((length($(Expr(:escape, :(1:4)))),length($(Expr(:escape, :(1:3)))))...) # /Users/yeesian/.julia/v0.5/JuMP/src/macros.jl, line 199:
    begin  # /Users/yeesian/.julia/v0.5/JuMP/src/macros.jl, line 181:
        let  # /Users/yeesian/.julia/v0.5/JuMP/src/macros.jl, line 182:
            begin 
                local $(Expr(:escape, :i))
            end # /Users/yeesian/.julia/v0.5/JuMP/src/macros.jl, line 183:
            for $(Expr(:escape, :i)) = $(Expr(:escape, :(1:4))) # /Users/yeesian/.julia/v0.5/JuMP/src/macros.jl, line 184:
                begin  # /Users/yeesian/.julia/v0.5/JuMP/src/macros.jl, line 181:
                    let  # /Users/yeesian/.julia/v0.5/JuMP/src/macros.jl, line 182:
                        begin 
                            local $(Expr(:escape, :j))
                        end # /Users/yeesian/.julia/v0.5/JuMP/src/macros.jl, line 183:
                        for $(Expr(:escape, :j)) = $(Expr(:escape, :(1:3))) # /Users/yeesian/.julia/v0.5/JuMP/src/macros.jl, line 184:
                            ω[$(Expr(:escape, :i)),$(Expr(:escape, :j))] = IndepNormal(m,mean=1,var=1,string("ω","[",string($(Expr(:escape, :i))),",",string($(Expr(:escape, :j))),"]"))
                        end
                    end
                end
            end
        end
    end # /Users/yeesian/.julia/v0.5/JuMP/src/macros.jl, line 200:
    nothing
end

Working with JuMP's Macros

In [11]:
using Base.Meta # for `quot`
In [12]:
macro variable(args...)
    args
end
Out[12]:
@variable (macro with 1 method)
In [13]:
args = @variable(m, x >= 0)
Out[13]:
(:m,:(x >= 0))
In [14]:
m = esc(args[1])
var = args[2].args[2]
lb = 0
ub = Inf
t = quot(:Default)
quotvarname = quot(JuMP.getname(var))
escvarname  = esc(JuMP.getname(var))
value = NaN

variablecall = :( constructvariable!($m, $JuMP._error, $lb, $ub, $t, string($quotvarname), $value) )
code = :($variable = $variablecall)
code = quote
    $code
    registervar($m, $quotvarname, $variable)
    $escvarname = $variable
end
JuMP.assert_validmodel(m, code)
Out[14]:
quote  # /Users/yeesian/.julia/v0.5/JuMP/src/macros.jl, line 243:
    validmodel($(Expr(:escape, :m)),:m) # /Users/yeesian/.julia/v0.5/JuMP/src/macros.jl, line 244:
    begin  # In[14], line 13:
        ##280 = constructvariable!($(Expr(:escape, :m)),JuMP._error,0,Inf,:Default,string(:x),NaN) # In[14], line 14:
        registervar($(Expr(:escape, :m)),:x,##280) # In[14], line 15:
        $(Expr(:escape, :x)) = ##280
    end
end

and now for a more involved example

In [15]:
args = @variable(m, x[1:N,1:N], Symmetric, Poly(X))
Out[15]:
(:m,:(x[1:N,1:N]),:Symmetric,:(Poly(X)))
In [16]:
var = args[2]
variable = gensym()
quotvarname = quot(JuMP.getname(var))
escvarname  = esc(JuMP.getname(var))

refcall, idxvars, idxsets, idxpairs, condition = JuMP.buildrefsets(var, variable)
@show refcall
@show idxvars
@show idxsets
@show idxpairs
@show condition;
refcall = :(##281[$(Expr(:escape, Symbol("##282"))),$(Expr(:escape, Symbol("##283")))])
idxvars = Any[Symbol("##282"),Symbol("##283")]
idxsets = Any[:($(Expr(:escape, :(1:N)))),:($(Expr(:escape, :(1:N))))]
idxpairs = JuMP.IndexPair[JuMP.IndexPair(nothing,:(1:N)),JuMP.IndexPair(nothing,:(1:N))]
condition = :(())

Homework: convince yourself that @variable(m, x[1:N,1:N], Symmetric, Poly(X)) is equivalent to:

In [17]:
x = Matrix{...}(N, N)
for i in 1:N
    for j in 1:N
        x[i,j] = x[j,i] = constructvariable!(m, Poly(X),
            msg -> error("In @variable(m, x[1:N,1:N], Symmetric, Poly(X)): ", msg), -Inf, Inf, :Cont, "", NaN
        )
    end
end
syntax: invalid identifier name "..."
In [18]:
macro constraint(args...)
    args
end
Out[18]:
@constraint (macro with 1 method)
In [19]:
args = @constraint(m, a*x <= 5)
Out[19]:
(:m,:(a * x <= 5))
In [20]:
@show m = args[1]
@show x = args[2]
@show extra = args[3:end];
m = args[1] = :m
x = args[2] = :(a * x <= 5)
extra = args[3:end] = ()
In [21]:
m = esc(m)
# Two formats:
    # - @constraint(m, a*x <= 5)
    # - @constraint(m, myref[a=1:5], a*x <= 5)
# Canonicalize the arguments
c = length(extra) == 1 ? x        : gensym()
x = length(extra) == 1 ? extra[1] : x
@show c
@show x

variable = gensym()
quotvarname = quot(JuMP.getname(c))
escvarname  = esc(JuMP.getname(c));
c = Symbol("##284")
x = :(a * x <= 5)
In [22]:
refcall, idxvars, idxsets, idxpairs, condition = JuMP.buildrefsets(c, variable)
Out[22]:
(Symbol("##285"),Any[],Any[],JuMP.IndexPair[],:(()))
In [23]:
(sense,vectorized) = JuMP._canonicalize_sense(x.args[1])
Out[23]:
(:(<=),false)
In [24]:
lhs = :($(x.args[2]) - $(x.args[3]))
Out[24]:
:(a * x - 5)
In [25]:
addconstr = (vectorized ? :addVectorizedConstraint : :addconstraint)
Out[25]:
:addconstraint
In [26]:
newaff, parsecode = JuMP.parseExprToplevel(lhs, :q)
constraintcall = :($addconstr($m, constructconstraint!($newaff,$(quot(sense)))))
code = quote
    q = zero(AffExpr)
    $parsecode
    $(refcall) = $constraintcall
end
Out[26]:
quote  # In[26], line 4:
    q = zero(AffExpr) # In[26], line 5:
    begin 
        begin 
            ##287 = addtoexpr_reorder(q,$(Expr(:escape, :a)),$(Expr(:escape, :x)))
        end
        ##286 = addtoexpr_reorder(##287,-1.0,$(Expr(:escape, 5)))
    end # In[26], line 6:
    ##285 = addconstraint($(Expr(:escape, :m)),constructconstraint!(##286,:(<=)))
end
In [27]:
JuMP.assert_validmodel(m, quote
    $(JuMP.getloopedcode(variable, code, condition, idxvars, idxsets, idxpairs, :ConstraintRef))
    $(quote
            registercon($m, $quotvarname, $variable)
            $escvarname = $variable
    end)
end)
Out[27]:
quote  # /Users/yeesian/.julia/v0.5/JuMP/src/macros.jl, line 243:
    validmodel($(Expr(:escape, :m)),:m) # /Users/yeesian/.julia/v0.5/JuMP/src/macros.jl, line 244:
    begin  # In[27], line 2:
        begin  # In[26], line 4:
            q = zero(AffExpr) # In[26], line 5:
            begin 
                begin 
                    ##287 = addtoexpr_reorder(q,$(Expr(:escape, :a)),$(Expr(:escape, :x)))
                end
                ##286 = addtoexpr_reorder(##287,-1.0,$(Expr(:escape, 5)))
            end # In[26], line 6:
            ##285 = addconstraint($(Expr(:escape, :m)),constructconstraint!(##286,:(<=)))
        end # In[27], line 3:
        begin  # In[27], line 4:
            registercon($(Expr(:escape, :m)),Symbol("##284"),##285) # In[27], line 5:
            $(Expr(:escape, Symbol("##284"))) = ##285
        end
    end
end

Note: JuMP also accepts constraint syntax of the form @constraint(m, foo in bar).

This will be rewritten to a call to constructconstraint!(foo, bar). To extend JuMP to accept set-based constraints of this form, it is necessary to add the corresponding methods to constructconstraint!.

Putting everything together

In [ ]:
function solve(m::Model; suppress_warnings=false,
                ignore_solve_hook=(m.solvehook===nothing),
                relaxation=false,
                kwargs...)
    # If the user or an extension has provided a solve hook, call
    # that instead of solving the model ourselves
    if !ignore_solve_hook
        return m.solvehook(m; suppress_warnings=suppress_warnings, kwargs...)::Symbol
    end
    # [...]
end

see JuMPChance.Model and JuMPChance.solvehook() for a working example.

JuMP.build()

visit the link ^ if you wish to see how the conversion from a JuMP.Model into a MathProgBase model is done.

JuMP.Model

See JuMPChance.ChanceModel() for an example of how the .ext attribute in JuMP.Model() might be used later on (e.g. [1], [2], [3], [4]).

Recap

  • Why Macros?
  • Quick Example
    • src/macros.jl
      • @variable (implementation: 250lines)
      • @constraint (implementation: 150lines)
    • some utilities in JuMP
      • for parsing and modifying expressions
      • "high-level-interface" to MathProgBase
  • What we talk about when we talk about JuMP