keyword_axis_vars.md

(based on work by @tkoolen)

Update:

This macro has been improved and moved into a new package: https://github.com/rdeits/AxisArrayVariables.jl

Original post:

We can define a macro to build up AxisArrays of JuMP variables:

macro axis_variables(m, var, axis_args...)
    axes = []
    for arg in axis_args
        if isa(arg, Expr) && (arg.head == :kw || arg.head == :(=))  # arguments parse as :kw in 0.5 and := in 0.6
            name = arg.args[1]
            domain = arg.args[2]
            push!(axes, Expr(:call, Expr(:curly, :Axis, Expr(:quote, name)), domain))
        else
            push!(axes, arg)
        end
    end
    ranges = [:(1:length($a)) for a in axes]
    quote
        vars = let
            local $(esc(var))
            @variable $m $var[$(ranges...)]
        end
        $(esc(var)) = $(Expr(:call, :AxisArray, :vars, axes...))
    end
end

which you can run like this:

julia> @axis_variables(m, x, time=1:5, side=[:left, :right])
2-dimensional AxisArray{JuMP.Variable,2,...} with axes:
    :time, 1:5
    :side, Symbol[:left,:right]
And data, a 5×2 Array{JuMP.Variable,2}:
 x[1,1]  x[1,2]
 x[2,1]  x[2,2]
 x[3,1]  x[3,2]
 x[4,1]  x[4,2]
 x[5,1]  x[5,2]

A more complete example:

m = Model()
@axis_variables(m, x, time=1:5, side=[:left, :right])
@objective(m, Min, sum(x.^2))

solve(m)

# Extract the values from x and store them in v. 
# v still has all the axis information:
v = similar(x, Float64)
v .= getvalue.(x)

Now we can get the values using the named axes:

v[:, :left]

You can also add Axis axes directly, and mix and match with keywords:

m = Model()
time = Axis{:time}(1:5)
@axis_variables(m, x, time, side=[:left, :right])

or inline:

m = Model()
@axis_variables(m, x, Axis{:time}(1:5), side=[:left, :right])

Cool! What's the main advantage of this approach, as opposed to

@variable(m, x[time=1:5, side=[:left,:right])

?

@joehuchette, looks like it has better slicing support.

Yeah, for example you can do

m = Model()
time = Axis{:time}(1:5)
@axis_variables(m, y, coord=[:x, :y], time, side=[:left, :right])
y[time(3)]

output:

2-dimensional AxisArray{JuMP.Variable,2,...} with axes:
    :coord, Symbol[:x,:y]
    :side, Symbol[:left,:right]
And data, a 2×2 Array{JuMP.Variable,2}:
 y[1,3,1]  y[1,3,2]
 y[2,3,1]  y[2,3,2]

or

y[:x]

output:

2-dimensional AxisArray{JuMP.Variable,2,...} with axes:
    :time, 1:5
    :side, Symbol[:left,:right]
And data, a 5×2 Array{JuMP.Variable,2}:
 y[1,1,1]  y[1,1,2]
 y[1,2,1]  y[1,2,2]
 y[1,3,1]  y[1,3,2]
 y[1,4,1]  y[1,4,2]
 y[1,5,1]  y[1,5,2]