Can you define a JuMP variable using the keys of an existing DenseAxisArray?

Specific Domains Optimization (Mathematical)
1

Say I have a variable defined over two axes:

@variable(m, x[1:3,[:a, :b]])

2-dimensional DenseAxisArray{VariableRef,2,...} with index sets:
    Dimension 1, Base.OneTo(3)
    Dimension 2, [:a, :b]
And data, a 3×2 Matrix{VariableRef}:
 x[1,a]  x[1,b]
 x[2,a]  x[2,b]
 x[3,a]  x[3,b]

Is there a way to dynamically define a second variable using the keys of the first?
For example, one might expect @variable(m, y[axes(x)...]) to work, but this unfortunately results in an error: “cannot use splatting operator ... in the definition of an index set.”

Using keys(x) as the index of y results in a variable that cannot be indexed in the same way as x. E.g. y[1,:a] will not work:

@variable(m, y[keys(x)])

1-dimensional DenseAxisArray{VariableRef,1,...} with index sets:
    Dimension 1, JuMP.Containers.DenseAxisArrayKey{Tuple{Int64, Symbol}}[JuMP.Containers.DenseAxisArrayKey{Tuple{Int64, Symbol}}((1, :a)), JuMP.Containers.DenseAxisArrayKey{Tuple{Int64, Symbol}}((2, :a)), JuMP.Containers.DenseAxisArrayKey{Tuple{Int64, Symbol}}((3, :a)), JuMP.Containers.DenseAxisArrayKey{Tuple{Int64, Symbol}}((1, :b)), JuMP.Containers.DenseAxisArrayKey{Tuple{Int64, Symbol}}((2, :b)), JuMP.Containers.DenseAxisArrayKey{Tuple{Int64, Symbol}}((3, :b))]
And data, a 6-element Vector{VariableRef}:
 y[JuMP.Containers.DenseAxisArrayKey{Tuple{Int64, Symbol}}((1, :a))]
 y[JuMP.Containers.DenseAxisArrayKey{Tuple{Int64, Symbol}}((2, :a))]
 y[JuMP.Containers.DenseAxisArrayKey{Tuple{Int64, Symbol}}((3, :a))]
 y[JuMP.Containers.DenseAxisArrayKey{Tuple{Int64, Symbol}}((1, :b))]
 y[JuMP.Containers.DenseAxisArrayKey{Tuple{Int64, Symbol}}((2, :b))]
 y[JuMP.Containers.DenseAxisArrayKey{Tuple{Int64, Symbol}}((3, :b))]
2

Hmm. You can’t do it via the macros because they use the syntax as typed, not the value of the variable in axes(x). What’s the use-case for this?

It’s not pretty, but:

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

julia> @variable(m, x[1:3, [:a, :b]])
2-dimensional DenseAxisArray{VariableRef,2,...} with index sets:
    Dimension 1, Base.OneTo(3)
    Dimension 2, [:a, :b]
And data, a 3×2 Matrix{VariableRef}:
 x[1,a]  x[1,b]
 x[2,a]  x[2,b]
 x[3,a]  x[3,b]

julia> N = length.(axes(x))
(3, 2)

julia> Y = reshape([VariableRef(m) for _ in 1:prod(N)], N)
3×2 Matrix{VariableRef}:
 _[7]  _[10]
 _[8]  _[11]
 _[9]  _[12]

julia> y = Containers.DenseAxisArray(Y, axes(x)...)
2-dimensional DenseAxisArray{VariableRef,2,...} with index sets:
    Dimension 1, Base.OneTo(3)
    Dimension 2, [:a, :b]
And data, a 3×2 Matrix{VariableRef}:
 _[7]  _[10]
 _[8]  _[11]
 _[9]  _[12]
3

Thanks for the example! I might be able to put something together using snippets from there.

I wanted to generate a residual variable for each of my non-linear equations, so that the equations can effectively be turned off by unfixing the residual. I have a mostly working macro written similarly to the existing JuMP macros, but I find it quite difficult to debug some of the macro hygiene edge cases and am trying to refactor as a macro that simply calls a function.

BUT, a much simpler use case is actually right there in the manual:

julia> N = 10
10

julia> S = [(1, 1, 1), (N, N, N)]
2-element Vector{Tuple{Int64, Int64, Int64}}:
 (1, 1, 1)
 (10, 10, 10)

julia> @time @variable(model, x1[i=1:N, j=1:N, k=1:N; (i, j, k) in S])
  0.203861 seconds (392.22 k allocations: 23.977 MiB, 99.10% compilation time)
JuMP.Containers.SparseAxisArray{VariableRef, 3, Tuple{Int64, Int64, Int64}} with 2 entries:
  [1, 1, 1   ]  =  x1[1,1,1]
  [10, 10, 10]  =  x1[10,10,10]

julia> @time @variable(model, x2[S])
  0.045407 seconds (65.24 k allocations: 3.771 MiB, 99.15% compilation time)
1-dimensional DenseAxisArray{VariableRef,1,...} with index sets:
    Dimension 1, [(1, 1, 1), (10, 10, 10)]
And data, a 2-element Vector{VariableRef}:
 x2[(1, 1, 1)]
 x2[(10, 10, 10)]

Note that the second example while more performant results in a different notation.
E.g. x2 cannot be accessed as x2[1, 1, 1].