Pyomo遍历模型变量时如何获取对应集合信息？

问题背景 #

我正在处理一个包含不同参数和变量的大型Pyomo优化问题，这些参数和变量关联不同的集合，示例变量定义代码如下：

mdl.A = Var(mdl.t, mdl.X, domain=NonNegativeReals)
mdl.B = Var(mdl.X, domain=NonNegativeReals)
mdl.C = Var(mdl.t, mdl.X, mdl.Y, domain=NonNegativeReals)

求解模型后，通过以下代码获取并遍历所有变量：

model_vars = mdl.component_map(ctype=pyo.Var)
for k in model_vars.keys():
    v = model_vars[k]

尝试用v.index_set().name获取关联集合时，部分场景能得到正确集合名，但部分场景输出类似“variablename_index”的自动生成名称，需要一种能准确获取变量对应集合的方法。

解决方案 #

方法1：直接访问变量的_index属性（推荐） #

Pyomo变量的_index属性会直接返回定义时传入的原始集合对象（或多集合组成的元组），可以直接提取集合名称：

import pyomo.environ as pyo

for var_name, var in model_vars.items():
    index_sets = var._index
    # 区分单集合与多集合场景
    if isinstance(index_sets, pyo.Set):
        set_names = [index_sets.name]
    else:
        set_names = [s.name for s in index_sets]
    print(f"变量 {var_name} 关联集合：{set_names}")

该方法能直接拿到定义变量时传入的集合，完全避免自动生成索引名的问题。

方法2：通过索引集合的_sets属性提取 #

如果变量的索引是由多个集合交叉生成的复合集合，可以通过index_set()._sets获取原始集合：

for var_name, var in model_vars.items():
    idx_set = var.index_set()
    # 复合集合会包含_sets属性，存储原始集合列表
    if hasattr(idx_set, '_sets'):
        set_names = [s.name for s in idx_set._sets]
    else:
        set_names = [idx_set.name]
    print(f"变量 {var_name} 关联集合：{set_names}")

方法3：定义变量时手动添加集合注释 #

对于需要长期维护的大型模型，可以在定义变量时通过doc参数记录关联集合，后续直接读取：

# 定义变量时添加注释
mdl.A = Var(mdl.t, mdl.X, domain=NonNegativeReals, doc="关联集合：t, X")
mdl.B = Var(mdl.X, domain=NonNegativeReals, doc="关联集合：X")
mdl.C = Var(mdl.t, mdl.X, mdl.Y, domain=NonNegativeReals, doc="关联集合：t, X, Y")

# 遍历读取注释
for var_name, var in model_vars.items():
    print(f"变量 {var_name} 关联集合：{var.doc}")

该方法可读性强，适合团队协作或复杂模型的维护。

内容的提问来源于stack exchange，提问作者Mathias Heiker