Strongly typed indices

[mlubin]

This is more of a Julia ecosystem issue than a JuMP issue, but it would be nice to have better syntax to do the following:

item_count = 3
factory_count = 3
struct ItemIndex
    value::Int
end

struct FactoryIndex
    value::Int
end

item_set = [ItemIndex(i) for i in 1:item_count]
factory_set = [FactoryIndex(i) for i in 1:factory_count]

m = Model()
@variable(m, item_is_produced_at_factory[item_set, factory_set], Bin)
item_is_produced_at_factory[ItemIndex(1), FactoryIndex(1)] # Ok
item_is_produced_at_factory[FactoryIndex(1), ItemIndex(1)] # Error

[odow]

Was this a request for this functionality? Or better syntax? Because this works as advertised.

[mlubin]

It was a request for better syntax. The right way to go about it would be to have a separate Julia package that provides strongly typed integers (e.g., https://embeddedartistry.com/blog/2018/04/23/namedtype-the-easy-way-to-use-strong-types-in-c/) and then provide JuMP integration with that package.

[odow]

So something like this?

julia> module TypedIndices

       struct TypedIndex{Name,T}
           data::T
       end
       function Base.show(io::IO, x::TypedIndex{Name,T}) where {Name,T}
           print(io, "$(Name)($(repr(x.data)))")
       end

       function TypedIndex(name::Symbol, iter::AbstractVector)
           return map(i -> TypedIndex{name,eltype(iter)}(i), iter)
       end

       function (x::Vector{TypedIndex{Name,T}})(key::T) where {Name,T}
           typed_key = TypedIndex{Name,T}(key)
           if !(typed_key in x)
               throw(KeyError(key))
           end
           return typed_key
       end

       macro index(expr)
           @assert expr.head == :ref
           @assert length(expr.args) == 2
           key, iter = expr.args[1], esc(expr.args[2])
           q_key = QuoteNode(key)
           return :($(esc(key)) = TypedIndex($q_key, $iter))
       end

       end
Main.TypedIndices

julia> using JuMP

julia> model = Model()
A JuMP Model
Feasibility problem with:
Variables: 0
Model mode: AUTOMATIC
CachingOptimizer state: NO_OPTIMIZER
Solver name: No optimizer attached.

julia> TypedIndices.@index Items[1:4]
4-element Array{Main.TypedIndices.TypedIndex{:Items,Int64},1}:
 Items(1)
 Items(2)
 Items(3)
 Items(4)

julia> TypedIndices.@index Factories[["a", "b", :c]]
3-element Array{Main.TypedIndices.TypedIndex{:Factories,Any},1}:
 Factories("a")
 Factories("b")
 Factories(:c)

julia> @variable(model, x[Items, Factories], Bin)
2-dimensional DenseAxisArray{VariableRef,2,...} with index sets:
    Dimension 1, Main.TypedIndices.TypedIndex{:Items,Int64}[Items(1), Items(2), Items(3), Items(4)]
    Dimension 2, Main.TypedIndices.TypedIndex{:Factories,Any}[Factories("a"), Factories("b"), Factories(:c)]
And data, a 4×3 Array{VariableRef,2}:
 x[Items(1),Factories("a")]  x[Items(1),Factories("b")]  x[Items(1),Factories(:c)]
 x[Items(2),Factories("a")]  x[Items(2),Factories("b")]  x[Items(2),Factories(:c)]
 x[Items(3),Factories("a")]  x[Items(3),Factories("b")]  x[Items(3),Factories(:c)]
 x[Items(4),Factories("a")]  x[Items(4),Factories("b")]  x[Items(4),Factories(:c)]

julia> x[Items(1), Factories("a")]
x[Items(1),Factories("a")]

julia> x[Factories("a"), Items(1)]
ERROR: KeyError: key Factories("a") not found
Stacktrace:
 [1] getindex at ./dict.jl:467 [inlined]
 [2] lookup_index at /Users/oscar/.julia/packages/JuMP/qhoVb/src/Containers/DenseAxisArray.jl:140 [inlined]
 [3] _to_index_tuple at /Users/oscar/.julia/packages/JuMP/qhoVb/src/Containers/DenseAxisArray.jl:149 [inlined]
 [4] to_index at /Users/oscar/.julia/packages/JuMP/qhoVb/src/Containers/DenseAxisArray.jl:167 [inlined]
 [5] getindex(::JuMP.Containers.DenseAxisArray{VariableRef,2,Tuple{Array{Main.TypedIndices.TypedIndex{:Items,Int64},1},Array{Main.TypedIndices.TypedIndex{:Factories,Any},1}},Tuple{Dict{Main.TypedIndices.TypedIndex{:Items,Int64},Int64},Dict{Main.TypedIndices.TypedIndex{:Factories,Any},Int64}}}, ::Main.TypedIndices.TypedIndex{:Factories,Any}, ::Main.TypedIndices.TypedIndex{:Items,Int64}) at /Users/oscar/.julia/packages/JuMP/qhoVb/src/Containers/DenseAxisArray.jl:182
 [6] top-level scope at REPL[8]:1

julia> @constraint(model, [f in Factories], sum(x[:, f]) <= 1)
1-dimensional DenseAxisArray{ConstraintRef{Model,MathOptInterface.ConstraintIndex{MathOptInterface.ScalarAffineFunction{Float64},MathOptInterface.LessThan{Float64}},ScalarShape},1,...} with index sets:
    Dimension 1, Main.TypedIndices.TypedIndex{:Factories,Any}[Factories("a"), Factories("b"), Factories(:c)]
And data, a 3-element Array{ConstraintRef{Model,MathOptInterface.ConstraintIndex{MathOptInterface.ScalarAffineFunction{Float64},MathOptInterface.LessThan{Float64}},ScalarShape},1}:
 x[Items(1),Factories("a")] + x[Items(2),Factories("a")] + x[Items(3),Factories("a")] + x[Items(4),Factories("a")] ≤ 1.0
 x[Items(1),Factories("b")] + x[Items(2),Factories("b")] + x[Items(3),Factories("b")] + x[Items(4),Factories("b")] ≤ 1.0
 x[Items(1),Factories(:c)] + x[Items(2),Factories(:c)] + x[Items(3),Factories(:c)] + x[Items(4),Factories(:c)] ≤ 1.0

[mlubin]

That's most of the way there, but I'd also expect typed integers to be as fast as integers. In the example above, Items(1) is treated like a key in a dictionary within DenseAxisArrays.

[odow]

Ah. The issue is really the lookup in DenseAxisArray, which is slow even if the axis is integer:

JuMP.jl/src/Containers/DenseAxisArray.jl

Lines 18 to 29 in 3ccbf39

	function build_lookup(ax)
	d = Dict{eltype(ax),Int}()
	cnt = 1
	for el in ax
	if haskey(d, el)
	error("Repeated index $el. Index sets must have unique elements.")
	end
	d[el] = cnt
	cnt += 1
	end
	return d
	end

It's probably worth checking if the axis is equivalent to a Base.OneTo(N) to give a fast lookup that avoids the dictionary.

Then something like

@variable(model, x[1:2, 2:3])

would be fast for the first index, and slow(er) for the second index. And you could have a fast fallback for TypedIndex{<:Any,<:Integer}.

[odow]

One option from #3214 is to support names in containers, and then support kwarg indexing (but in the right order).

something like

m = Model()
@variable(m, x[item = 1:3, factory = 1:3], Bin)
x[item = 1, factory = 2]  # ok
x[factory = 2, item = 1]  # error?

[odow]

Miles says go for it

[blegat]

I was a bit worried about the argument raised yesterday that it won't work with Array. I think most users don't even notice the change of Containers at the moment