Using alternative array container types in JuMP

Specific Domains Optimization (Mathematical)
1

I am trying to force the container type used by a JuMP model. My example uses KeyedArray but I’ve seen a similar error when specifying other types from outside JuMP (e.g. AxisArray).

using JuMP, AxisKeys, GLPK
vals = ["A", "B", "C", "D"]
model = Model(GLPK.Optimizer)
@variable(model, x[v in vals], Bin, container=KeyedArray)

The error is:

ERROR: UndefVarError: KeyedArray not defined
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/JuMP/zn6NT/src/macros.jl:142 [inlined]
 [2] top-level scope
   @ REPL[26]:1

I am wondering why KeyedArray is not recognised since AxisKeys has been loaded. I am using Julia 1.7 and JuMP 0.23.1.

3 Likes
2

This exposed a bug in our container interface. Let me make a fix and I’ll get back to you with the syntax.

2 Likes
3

Okay, this needs https://github.com/jump-dev/JuMP.jl/pull/2916, which will be in the next JuMP release (but likely not today or in the next few days).

Once it lands, you’ll need to implement the Containers.container function to let JuMP know how to handle the new type. This can be a but tricky as there are a few different indices that can be passed. But a good first attempt is:

julia> using JuMP, AxisKeys

julia> function Containers.container(f::Function, indices, ::Type{AxisKeys.KeyedArray})
           set = collect(indices)
           data = [f(s...) for s in set]
           return AxisKeys.KeyedArray(data, set)
       end

julia> S = ["A", "B", "C", "D"]
4-element Vector{String}:
 "A"
 "B"
 "C"
 "D"

julia> model = Model();

julia> @variable(model, x[S], Bin, container = KeyedArray)
1-dimensional KeyedArray(...) with keys:
↓   4-element Vector{Tuple{String}}
And data, 4-element Vector{VariableRef}:
  ("A",)   x[A]
  ("B",)   x[B]
  ("C",)   x[C]
  ("D",)   x[D]

Note that the keys are x(("A",)), but you could fix that in the Containers.container function.

In general though, this syntax isn’t intended for widespread usage or extension (which is why it had a bug and isn’t well documented.

A much better work-around is to go:

julia> model = Model();

julia> @variable(model, x[1:4], Bin)
4-element Vector{VariableRef}:
 x[1]
 x[2]
 x[3]
 x[4]

julia> X = KeyedArray(x, S)
1-dimensional KeyedArray(...) with keys:
↓   4-element Vector{String}
And data, 4-element Vector{VariableRef}:
 ("A")  x[1]
 ("B")  x[2]
 ("C")  x[3]
 ("D")  x[4]

You should think of the built-in JuMP types as a good starting point. But there’s nothing stopping you from creating your own types outside JuMP without trying to interact with the complicated macro machinery.

6 Likes
4

I see, thank you for the explanation!

5

Thanks for the explanation @odow, this is super helpful :slight_smile:

Our use case involves a large model that is built via several functions which each add a part, something like

function build_model(data)
    m = Model(...)
    add_variables!(m, data)
    add_constraints!(m, data)
    add_objective!(m, data)
    return m
end

The downside of your suggested solution is that X (i.e., the KeyedArray wrapper for x) is not part of the model’s object dictionary, and therefore we can’t easily fetch X from m when analysing the results. One solution would be to create an extra variable/constraint for each one of them such as

X = model[:X] = KeyedArray(model[:x], indices=idx)

Which works fine, but requires us to have twice as many variables/constraints in the model and to keep track of which are the KeyedArray ones and which are not. Do you have any better suggestion for this case?

6

This has been released in JuMP 0.23.2: Extensions · JuMP
I originally added this as a request from @oxinabox so I don’t think the “most users shouldn’t do this” applies to you.

If your keyed arrays are just vectors, then you can use the Containers.container function above. Otherwise, you’ll need to write something a bit more complicated.

Which works fine, but requires us to have twice as many variables/constraints in the model and to keep track of which are the KeyedArray ones and which are not.

This is what I had in mind. Note that this isn’t creating twice as many decision variables. It’s just a new wrapper. You could also call unregister(model, :x) to get rid of the :x name, or you could just over-write it if you didn’t want to refer to it later.

2 Likes