How can I index a variable in Pyomo with two lists or Sets, but have the second list be different for each element of the first list?

Question

I am learning Pyomo and I want to try to build a variable that is indexed by two lists or set. This is easy if for each element of the list N the elements of the list M are part of it.

from pyomo.environ import*
model = ConcreteModel
N=[1,2,3] 
U=['A','B','C','D']
model.X = Var(N,U,within=NonNegativeReals)
model.X.pprint()

Retrieved from

X : Size=12, Index=X_index
Key      : Lower : Value : Upper : Fixed : Stale : Domain
(1, 'A') :     0 :  None :  None : False :  True : NonNegativeReals
(1, 'B') :     0 :  None :  None : False :  True : NonNegativeReals
(1, 'C') :     0 :  None :  None : False :  True : NonNegativeReals
(1, 'D') :     0 :  None :  None : False :  True : NonNegativeReals
(2, 'A') :     0 :  None :  None : False :  True : NonNegativeReals
(2, 'B') :     0 :  None :  None : False :  True : NonNegativeReals
(2, 'C') :     0 :  None :  None : False :  True : NonNegativeReals
(2, 'D') :     0 :  None :  None : False :  True : NonNegativeReals
(3, 'A') :     0 :  None :  None : False :  True : NonNegativeReals
(3, 'B') :     0 :  None :  None : False :  True : NonNegativeReals
(3, 'C') :     0 :  None :  None : False :  True : NonNegativeReals
(3, 'D') :     0 :  None :  None : False :  True : NonNegativeReals

This is easy and I know how to do it. But now I want that for each element N[i] the indexing set U is different as well as of different sizes. I have tried to use a list of lists but I don't get what I want.

from pyomo.environ import*
model = ConcreteModel
N=[1,2,3] 
U=U=[['A','B','C','D'],['B','C','E','G'],['A','C','D']]
model.X = Var(N,U,within=NonNegativeReals)
model.X.pprint()

The result is:

The value=['A', 'C', 'D'] has dimension 3 and is not valid for Set OrderedScalarSet which has dimen=4

And I want to get

X : Size=12, Index=X_index
Key      : Lower : Value : Upper : Fixed : Stale : Domain
(1, 'A') :     0 :  None :  None : False :  True : NonNegativeReals
(1, 'B') :     0 :  None :  None : False :  True : NonNegativeReals
(1, 'C') :     0 :  None :  None : False :  True : NonNegativeReals
(1, 'D') :     0 :  None :  None : False :  True : NonNegativeReals
(2, 'B') :     0 :  None :  None : False :  True : NonNegativeReals
(2, 'C') :     0 :  None :  None : False :  True : NonNegativeReals
(2, 'E') :     0 :  None :  None : False :  True : NonNegativeReals
(2, 'G') :     0 :  None :  None : False :  True : NonNegativeReals
(3, 'A') :     0 :  None :  None : False :  True : NonNegativeReals
(3, 'C') :     0 :  None :  None : False :  True : NonNegativeReals
(3, 'D') :     0 :  None :  None : False :  True : NonNegativeReals

How can I get it?

Thank you very much and best regards

Answer 1

Create a custom index set that includes the desired combinations of indices from the two lists:

from pyomo.environ import *

model = ConcreteModel()

N = [1, 2, 3]
U = [['A', 'B', 'C', 'D'], ['B', 'C', 'E', 'G'], ['A', 'C', 'D']]

# Create a custom index set
index_set = [(n, u) for i, n in enumerate(N) for u in U[i]]

# Create the variable with the custom index set
model.X = Var(index_set, within=NonNegativeReals)
model.X.pprint()