Indexing Nonlinear Expressions with nested indices in JuMP

I’m trying to define a JuMP @NLexpression with embedded indexing and I need to use numerical indices instead of collections. I would have thought that the below would work but I’m getting an error saying that F was already defined… Might anyone know how to circumvent this problem?

arr_1 = [A,B,C,...] ## A,B,C are structs that have arr_2 as fields
for a = 1:length(arr_1)
	arr_2_a = arr_1[a].arr_2 ## arr_1[a].arr_2 = [x_a, y_a, z_a, ...], each of which is another struct
	
	for b = 1:length(arr_2_a)
		@NLexpression(
			model,
			F[a,b] = some_function_of_arr_1[a]_and_arr_2_a[b]
			)
	end
end

@NLexpressions have the same collection syntax as normal @expressions.

Here I the documentation:
https://www.juliaopt.org/JuMP.jl/stable/nlp/#Syntax-notes-1
https://www.juliaopt.org/JuMP.jl/stable/nlp/#JuMP.@NLexpression

# Instead of:
for a = 1:4
    @NLexpression(model, F[a], x^a)
end
# use
@NLexpression(model, F[a = 1:4], x^a)

Right but my issue here is that the indexing is embedded – I need to index by a=1:4 and then by b=1:f(a) where f(a) is a number in a dict

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> @variable(model, x)
x

julia> f = Dict(1 => 2, 2 => 4)
Dict{Int64,Int64} with 2 entries:
  2 => 4
  1 => 2

julia> @NLexpression(model, F[a=1:2, b=1:f[a]], b * x^a)
JuMP.Containers.SparseAxisArray{NonlinearExpression,2,Tuple{Int64,Int64}} with 6 entries:
  [1, 2]  =  "Reference to nonlinear expression #2"
  [2, 3]  =  "Reference to nonlinear expression #5"
  [2, 4]  =  "Reference to nonlinear expression #6"
  [2, 2]  =  "Reference to nonlinear expression #4"
  [1, 1]  =  "Reference to nonlinear expression #1"
  [2, 1]  =  "Reference to nonlinear expression #3"
2 Likes