[Containers] support ArrayInterface.jl traits

[ParadaCarleton]

As far as I'm aware, JuMP's containers/labeled arrays aren't used outside of JuMP itself. This can cause a lot of problems and incompatibilities when trying to work with data from packages other than JuMP.

DimensionalData.jl is more widely-used, e.g. throughout the Julia Geostats ecosystem and in ArviZ.jl. Would it make sense to switch JuMP's containers out with DimensionalData.jl DimArrays for compatibility?

[jd-foster]

I guess the short answer is that the principle in JuMP / MathOptInterface is to try to keep external dependencies to a minimum. In this regard, there seems to be small gains from "compatibility" in the abstract considering that the required functionality is fully implemented within the package with an AbstractArray interface.

[blegat]

What we need is that the array directly generalizes Array so that the user does not have to know about whether the container Array or DenseAxisArray or SparseAxisArray.
This is what prevents us to use a Dict for SpaceAxisArray, because a Dict does not just behave like an Array with custom indices with missing entries. What we want there is really like a SparseArray with custom indexing.
The issue with AxisArrays is that if you pass integer indices, it assumes that you are directly indexing the underlying array:

julia> using AxisArrays

julia> A = AxisArray(1:2; i = 2:3)
1-dimensional AxisArray{Int64,1,...} with axes:
    :i, 2:3
And data, a 2-element UnitRange{Int64}:
 1
 2

julia> A[2]
2

We would need A[2] to be 1 here.
There is the same issue with DimensionalData:

julia> A = ones(X(2:3), Y(3:4))
2×2 DimArray{Float64,2} with dimensions: 
  X Sampled{Int64} 2:3 ForwardOrdered Regular Points,
  Y Sampled{Int64} 3:4 ForwardOrdered Regular Points
    3    4
 2  1.0  1.0
 3  1.0  1.0

julia> A[2, 3]
ERROR: BoundsError: attempt to access 2×2 Matrix{Float64} at index [2, 3]
Stacktrace:
 [1] getindex
   @ ./array.jl:925 [inlined]
 [2] getindex(::DimArray{Float64, 2, Tuple{X{DimensionalData.Dimensions.LookupArrays.Sampled{Int64, UnitRange{Int64}, DimensionalData.Dimensions.LookupArrays.ForwardOrdered, DimensionalData.Dimensions.LookupArrays.Regular{Int64}, DimensionalData.Dimensions.LookupArrays.Points, DimensionalData.Dimensions.LookupArrays.NoMetadata}}, Y{DimensionalData.Dimensions.LookupArrays.Sampled{Int64, UnitRange{Int64}, DimensionalData.Dimensions.LookupArrays.ForwardOrdered, DimensionalData.Dimensions.LookupArrays.Regular{Int64}, DimensionalData.Dimensions.LookupArrays.Points, DimensionalData.Dimensions.LookupArrays.NoMetadata}}}, Tuple{}, Matrix{Float64}, DimensionalData.NoName, DimensionalData.Dimensions.LookupArrays.NoMetadata}, ::Int64, ::Int64)
   @ DimensionalData ~/.julia/dev/DimensionalData/src/array/indexing.jl:6
 [3] top-level scope
   @ REPL[10]:1

julia> A[At(2), At(3)]
1.0

Here, it seems you need to use At to get the behavior we would want by default.

If there was a package that would define arrays that we can use without having to change the syntax to our users, I would be in favor of using it but there is none at the moment as far as I know which is why we have to define our own containers.

[odow]

The largest reason against doing this (for the default behavior) is that it would be breaking.

However, adding support for @variable(model, x[2:3, 3:4], container = DimArray) is do-able in an extension package in the upcoming Julia v1.9.

[mlubin]

Just for some historical context, we did evaluate AxisArrays as a replacement for JuMP's custom containers. The indexing issue that @blegat described was a blocker: JuliaArrays/AxisArrays.jl#117.

[metab0t]

The extension method is an attractive approach to support custom container type.
@odow What are the functions we should specialize in extension to support multiple container types？

[odow]

See https://jump.dev/JuMP.jl/stable/developers/extensions/#Creating-new-container-types

[ParadaCarleton]

Just for some historical context, we did evaluate AxisArrays as a replacement for JuMP's custom containers. The indexing issue that @blegat described was a blocker: JuliaArrays/AxisArrays.jl#117.

I think this is a reasonable concern. To deal with this, could we just wrap a DimArray around an OffsetArray?

[odow]

To deal with this, could we just wrap a DimArray around an OffsetArray?

I'm not sure I understand how this would work. JuMP's DenseAxisArray can use arbitrary indices of any Julia type. For example:

julia> model = Model();

julia> a, b = [1, 2], [3, 4]
([1, 2], [3, 4])

julia> @variable(model, x[1:2, [:a, :b], ["c", "d"], [a, b]])
4-dimensional DenseAxisArray{VariableRef,4,...} with index sets:
    Dimension 1, Base.OneTo(2)
    Dimension 2, [:a, :b]
    Dimension 3, ["c", "d"]
    Dimension 4, [[1, 2], [3, 4]]
And data, a 2×2×2×2 Array{VariableRef, 4}:
[:, :, "c", [1, 2]] =
 x[1,a,c,[1, 2]]  x[1,b,c,[1, 2]]
 x[2,a,c,[1, 2]]  x[2,b,c,[1, 2]]

[:, :, "d", [1, 2]] =
 x[1,a,d,[1, 2]]  x[1,b,d,[1, 2]]
 x[2,a,d,[1, 2]]  x[2,b,d,[1, 2]]

[:, :, "c", [3, 4]] =
 x[1,a,c,[3, 4]]  x[1,b,c,[3, 4]]
 x[2,a,c,[3, 4]]  x[2,b,c,[3, 4]]

[:, :, "d", [3, 4]] =
 x[1,a,d,[3, 4]]  x[1,b,d,[3, 4]]
 x[2,a,d,[3, 4]]  x[2,b,d,[3, 4]]

julia> x[2, :a, "d", b]
x[2,a,d,[3, 4]]

Since there's a public interface to defining new container types (#3214 (comment)), I think the first step is to demonstrate what it would look like in a separate package. I haven't looked at the details, but I assume the necessary code is something like:

function Containers.container(f::Function, indices, ::Type{DimensionalData.DimArray}, names)
    data = [f(i...) for i in indices]
    return DimensionalData.DimArray(data, names)
end

[ParadaCarleton]

I'm not sure I understand how this would work. JuMP's DenseAxisArray can use arbitrary indices of any Julia type. For example:

Right, so what I mean is that the container macro would create a DimArray wrapped around an OffsetArray. If the user creates an index with unusual integer indices, e.g. -12:100 or 0:2:100, the indices of the underlying array will be offset to match the requested labels. Then, indexing with x[-12] will work as usual, since the integer indices still refer to the right thing.

[odow]

The axis could also be 3:-1:1, or [1, 5, 7, 2], so I don't think this works in general?

[ParadaCarleton]

The axis could also be 3:-1:1, or [1, 5, 7, 2], so I don't think this works in general?

Hmm, I know OffsetArrays handles the first case, but I don't know if it handles the second. If not, it's possible we could modify DimensionalData.jl to fix this.

[odow]

If not, it's possible we could modify DimensionalData.jl to fix this.

Potentially. But it's another argument to leave JuMP's defaults alone and work on the JuMP<->DimensionalData interface in a separate opt-in package.

[rafaqz]

1.9 should make this fairly easy to do.

But I wasn't aware you had an in-house axis array implementation here. Its a bit tragic how much of this work we have duplicated in the last few years, although I understand the desire to keep things in-house for JuMP. FWIW AbstractDimArray (in DimensionalData.jl) is extensible and could have custom JuMP specific logic added without having to implement everything here.

There are also some efforts at ArrayInterface.jl to add a compatability layer for these array labelling packages. It may be that we can add a weak dependency for that instead and make things generic so e.g. AxisKeys.jl will work as well.

@odow DimensionalData.jl lookups can contain any julia objects, 3:-1:1, or [1, 5, 7, 2] are no problem at all. The vector will be classified as Irregular Unordered and will use findfirst rather than searchsortedfirst when indexing.