Enums for array indexing

I am building up a stochastic transition matrix, and I want to use constants for the indices of various states to avoid errors, eg

P[state_A, state_B] = some_probability
P[state_A, sink] = death_probability

I can do it with const state_A = 1 etc.

However, using @enum States state_A=1 ... would be so much nicer, but this does not work directly. I can define

Base.getindex(A::AbstractArray, inds::Enum...) =
    getindex(A, [convert(Int, e) for e in inds]...)
Base.setindex!(A::AbstractArray, X, inds::Enum...) =
    setindex!(A, X, [convert(Int, e) for e in inds]...)

but it is a bit heavy-handed, and does not support mixtures of enums and ints, let alone :.

Is there an elegant trick that would make it work, or should I stick to consts, maybe with a nice macro like

macro const_indexes(names...)
        $([:(const $(esc(name)) = $(index)) for (index, name) in enumerate(names)]...)

I think you should stick to constants. The purpose of Enum is to use types and rely on dispatch. Here, you most definitely won’t know what all of the dispatches will be at compile-time (since your current state is random!), so no matter what solution you have, I think Enums would be slow because of dynamic dispatching.

What do you mean with dynamic dispatching? All the state enums would be of the same type, namely States.

The purpose of Enum is to use types and rely on dispatch

I disagree, the purpose of an enum is to be an element in a finite discrete space. Seems to me exactly what OP is intending to use it for.


Yes, I seem to have this wrong. I incorrectly remembered using it with value types for related dispatches, but you can definitely use them for a finite space. The example right below here is pretty much spot on for this:


So take my solution thing away: I was wrong.

1 Like

No, you still get the solution thingy, because you gave me the right practical advice: at the moment, use constants. I agree with @kristoffer.carlsson that an enum-like solution would be the “right” thing for this case, with some extensions, eg if I could get the number of levels I could use that for matrix size etc. But if they are not simple to implement, the thing I can do now is use constants.

If you would be satisfied with the first index to be always an enum, but the rest being a mixture of enums and integers, you might use a union of enum and integer starting at the second index. Thus your line would read something like

Base.getindex(A::AbstractArray, idx1::Enum, inds::Union{Enum,Int}...) = getindex(A, convert(Int, idx1), [convert(Int, e) for e in inds]...)

Unfortunately, using the union for all indices leads to error in the further evaluation, which I’d suspect to be the result of ambiguous method definition of the pure integer version.

julia> @enum States state_1=1 state_2=2

julia> a = [1 2; 3 4]
2×2 Array{Int64,2}:
 1  2
 3  4

julia> a[state_2,1]

julia> a[state_2,state_1]

julia> a[0,state_1]
------ ErrorException ------------------ Stacktrace (most recent call last)

 [1] — getindex(::Array{Int64,2}, ::Int64, ::States) at abstractarray.jl:752

indexing Array{Int64,2} with types Tuple{Int64,States} is not supported

This is kinda-sorta possible, but it’s going to be fragile. You want to define something like getindex(::AbstractArray, ::Union{Enum, OfficallySupportedIndexTypes}...), convert the enums, and then re-dispatch to getindex(::AbstractArray, ::OfficallySupportedIndexTypes...). But the trouble is that the methods defined in base are defined as getindex(::AbstractArray, ::Any...), so suddenly your new definition is more specific and will hit stack overflow errors. I think the only way around this is to hook into the internal machinery in an undocumented manner.

The other issue is that splatting a comprehension is going to be terrible for performance, but that’s an easy fix. The better way here is with a recursive function:

enums_to_index() = ()
enums_to_index(e::Enum, rest...) = (convert(Int, e), enums_to_index(rest...)...)
enums_to_index(x, rest...) = (x, enums_to_index(rest...)...)

Now you can use enums_to_index with any combination of argument types — be they arrays or integers or whatever. This construct is very efficient.

Just as an update: this is now extremely easy and uniform on version 0.6.

@enum States state_1=1 state_2=2
Base.to_index(s::States) = Int(s)
A = [1 2; 3 4]

julia> A[state_2, 1]

julia> A[2, state_1] = 42

julia> @view A[:, state_1]
2-element SubArray{Int64,1,Array{Int64,2},Tuple{Base.Slice{Base.OneTo{Int64}},Int64},true}: