Derived keys for indexing into JuMP containers

I am seeking to better understand JuMP’s indexing functionality. As I understand it, when a set of decision variables or constraints in a container is created, JuMP automatically uses the elements of the set as the indexing key for these variables or constraints.

However, I am wondering if we could instruct JuMP to use a derived key instead. To illustrate, consider the following set up:


model = Model()

my_set = Set([MyStruct(1), MyStruct(2)])

@variable(model, x[my_set]) # here, elements of `my_set` serve as indices

What I would like to explore is the possibility of using an element’s internal id (or similar property) as an index:


JuMP.derived_key(o::MyStruct) = o.internal_id

@variable(model, y[my_set]) # here, we'd like indices to be based on internal id's

Is this possible with JuMP?

@slwu89 is also interested.

1 Like

Unfortunately this is not possible and there is no workaround.

If you want to use a different index you must explicitly create and use that index.

2 Likes

Thank you for your swift response!

1 Like

It would nevertheless be possible to do something like

derived_key(x) = x

# Collect iterators with unknown size so they can be used as axes.
_collect(::Base.SizeUnknown, x) = collect(derived_key.(x))

_collect(::Any, x) = derived_key.(x)

around https://github.com/jump-dev/JuMP.jl/blob/12ffa9110700f05fe7504312a903eea13db31324/src/Containers/vectorized_product_iterator.jl#L46

This is evidently a quick “fix” that almost certainly fails to address some edge cases.

For a short demo, it would allow me to do

model = Model()

struct MyStruct
    x
end

JuMP.Containers.derived_key(s::MyStruct) = s.x ^ 2

my_set = [MyStruct(i) for i=1:3]

@variable(model, x[my_set])

resulting in

1-dimensional DenseAxisArray{VariableRef,1,...} with index sets:
    Dimension 1, [1, 4, 9]
And data, a 3-element Vector{VariableRef}:
 x[1]
 x[4]
 x[9]

The question is, would you consider this beneficial to incorporate into JuMP? If so, I would refine it further and submit a PR.

1 Like

I don’t think we want to support this in JuMP.

The main reason is that it introduces another layer of complexity between the code as written and what gets executed.

In your example, it is much more explicit what you are doing if you do something like this:

struct MyStruct
    x::Int
end
my_set = Dict(i^2 => MyStruct(i) for i in 1:3)
model = Model()
@variable(model, x[keys(my_set)])
2 Likes

Sure, thanks for clarification - and for the awesome work on JuMP in general!

1 Like