From 9b4265ce01cafbba94c71b1582c98b8eb0f5707c Mon Sep 17 00:00:00 2001 From: John Siirola Date: Fri, 29 May 2026 00:40:18 -0600 Subject: [PATCH 1/5] Support a reusable 'seen' cache within the IdentifyVariableVisitor --- pyomo/core/expr/visitor.py | 45 ++++++++++++++----- pyomo/core/tests/unit/test_visitor.py | 19 ++++++++ pyomo/util/vars_from_expressions.py | 65 ++++++++++++++++++++++----- 3 files changed, 109 insertions(+), 20 deletions(-) diff --git a/pyomo/core/expr/visitor.py b/pyomo/core/expr/visitor.py index 210c0ec32e7..ca20cdf6e0a 100644 --- a/pyomo/core/expr/visitor.py +++ b/pyomo/core/expr/visitor.py @@ -1395,32 +1395,53 @@ def identify_components(expr, component_types): class IdentifyVariableVisitor(StreamBasedExpressionVisitor): - def __init__(self, include_fixed=False, named_expression_cache=None): + def __init__(self, include_fixed=False, named_expression_cache=None, seen=None): """Visitor that collects all unique variables participating in an expression - Args: - include_fixed (bool): Whether to include fixed variables - named_expression_cache (optional, dict): Dict mapping ids of named - expressions to a tuple of the list of all variables and the - set of all variable ids contained in the named expression. + Parameters + ---------- + include_fixed : bool + If True, fixed variables will be reported + + named_expression_cache : dict + Cache of named expressions that have been visited by this + walker. The value includes the variables within the named + expression as well as information for detecting when the + named expression has changed (for cache invalidation). + + seen : dict[int, VarData] + Dict mapping the :func:`id()` of variables to + :class:`VarData` for all variables "seen" by this walker. + If provided, this dictionary is preserved between calls to + :meth:`walk_expression` (so repeated variables are not + returned more than once) """ super().__init__() self._include_fixed = include_fixed + # cache of visited named expressions. This dict maps + # {eid: (seen, exprs)}. + # - eid is the id() of the named expression + # - seen is the processed result for the named expression + # (including any nested named expressions) + # - exprs is used for automatically invalidating the cache (see below). self._cache = named_expression_cache # Stack of named expressions. This holds the tuple # (eid, _seen, _exprs) # where eid is the id() of the subexpression we are currently # processing, and _seen and _exprs are from the parent context. self._expr_stack = [] - # The following attributes will be added by initializeWalker: - # self._seen: dict(eid: obj) + # cache of "seen" variables: dict(eid: VarData) + self._seen = seen + # The following attribute will be added by initializeWalker: # self._exprs: list of (e, e.expr) for any (nested) named expressions def initializeWalker(self, expr): assert not self._expr_stack - self._seen = {} + if self._seen is None: + self._seen = {} + self._expr_stack.append(None) self._exprs = None if not self.beforeChild(None, expr, 0)[0]: return False, self.finalizeResult(None) @@ -1452,8 +1473,12 @@ def exitNode(self, node, data): self._merge_obj_lists(_seen, _exprs) def finalizeResult(self, result): + seen = self._seen + if self._expr_stack: + assert self._expr_stack.pop() is None + self._seen = None assert not self._expr_stack - return self._seen.values() + return seen.values() def _merge_obj_lists(self, _seen, _exprs): self._seen.update(_seen) diff --git a/pyomo/core/tests/unit/test_visitor.py b/pyomo/core/tests/unit/test_visitor.py index 6268afd45ad..ba825e1c2ee 100644 --- a/pyomo/core/tests/unit/test_visitor.py +++ b/pyomo/core/tests/unit/test_visitor.py @@ -61,6 +61,7 @@ SimpleExpressionVisitor, StreamBasedExpressionVisitor, ExpressionReplacementVisitor, + IdentifyVariableVisitor, evaluate_expression, expression_to_string, replace_expressions, @@ -252,6 +253,24 @@ def test_identify_vars_linear_expression(self): expr = quicksum([m.x, m.x], linear=True) self.assertEqual(list(identify_variables(expr, include_fixed=False)), [m.x]) + def test_identify_vars_seen_cache(self): + m = ConcreteModel() + m.x = Var() + m.y = Var() + m.z = Var() + + e1 = m.x + m.y + e2 = m.y + m.z + + v = IdentifyVariableVisitor() + self.assertEqual(list(v.walk_expression(e1)), [m.x, m.y]) + self.assertEqual(list(v.walk_expression(e2)), [m.y, m.z]) + + seen = {} + v = IdentifyVariableVisitor(seen=seen) + self.assertEqual(list(v.walk_expression(e2)), [m.y, m.z]) + self.assertEqual(list(v.walk_expression(e1)), [m.y, m.z, m.x]) + class TestIdentifyParams(unittest.TestCase): def test_identify_params_numeric(self): diff --git a/pyomo/util/vars_from_expressions.py b/pyomo/util/vars_from_expressions.py index 6ad2f80b56b..27907ffc51f 100644 --- a/pyomo/util/vars_from_expressions.py +++ b/pyomo/util/vars_from_expressions.py @@ -8,11 +8,11 @@ # ____________________________________________________________________________________ """ -This module contains a function to generate a list of the Vars which appear +This module contains functions to generate a list of the Vars appearing in expressions in the active tree. Note this is not the same as -component_data_objects(Var) because it does not look for Var objects which are +``component_data_objects(Var)`` because it does not look for Var objects which are not used in any expressions and it does not care if the Vars it finds are -actually in the subtree or not. +actually in the Block subtree or not. """ from pyomo.core import Block @@ -28,7 +28,7 @@ def get_vars_from_components( descend_into=Block, descent_order=None, ): - """Returns a generator of all the Var objects which are used in Constraint + """Returns a generator of all the Var objects which are used in expressions on the block. By default, this recurses into sub-blocks. Args: @@ -40,9 +40,10 @@ def get_vars_from_components( sort: sort method for iterating through Constraint objects descend_into: Ctypes to descend into when finding Constraints descent_order: Traversal strategy for finding the objects of type ctype + """ - visitor = IdentifyVariableVisitor(include_fixed, {}) - seen = set() + seen = {} + visitor = IdentifyVariableVisitor(include_fixed, {}, seen=seen) for constraint in block.component_data_objects( ctype, active=active, @@ -50,7 +51,51 @@ def get_vars_from_components( descend_into=descend_into, descent_order=descent_order, ): - for var in visitor.walk_expression(constraint.expr): - if id(var) not in seen: - seen.add(id(var)) - yield var + visitor.walk_expression(constraint.expr) + return seen.values() + + +def get_vars( + block, + include_fixed=False, + active=True, + sort=False, + descend_into=Block, + descent_order=None, +): + """Return all vars referenced through expressions in the specified block. + + This is a simple wrapper around :func:`get_vars_from_components()` + that gathers all variables referenced by :class:`Constraint` and + :class:`Objective` objects within the specified block. Note that as + it is designed to return the "variables used in the current model," + it uses different defaults for `active` and `include_fixed`. + + Parameters + ---------- + include_fixed : bool + If True, both fixed and free variables will be returned + + active : bool | None + If True, only variables accessible through the active component + tree will be erturned + + sort: SortOrder | None + sort method for iterating through Constraint objects + + descend_into : None | type | tuple[type] + Ctypes to descend into when finding Constraints + + descent_order : None | TraversalStrategy + Traversal strategy for walking the block hierarchy + + """ + return get_vars_from_components( + block, + ctype=(Constraint, Objective), + include_fixed=include_fixed, + active=active, + sort=sort, + descend_into=descend_into, + descent_order=descent_order, + ) From 3dd824591c83f61cd95382f5396a625fb2e8e944 Mon Sep 17 00:00:00 2001 From: John Siirola Date: Fri, 29 May 2026 00:48:26 -0600 Subject: [PATCH 2/5] NFC: fix typo --- pyomo/util/vars_from_expressions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyomo/util/vars_from_expressions.py b/pyomo/util/vars_from_expressions.py index 27907ffc51f..6941532fc67 100644 --- a/pyomo/util/vars_from_expressions.py +++ b/pyomo/util/vars_from_expressions.py @@ -78,7 +78,7 @@ def get_vars( active : bool | None If True, only variables accessible through the active component - tree will be erturned + tree will be returned sort: SortOrder | None sort method for iterating through Constraint objects From 037981ab5dc972e9caf212446337a14a05bc25bd Mon Sep 17 00:00:00 2001 From: John Siirola Date: Fri, 29 May 2026 09:01:31 -0600 Subject: [PATCH 3/5] Fix import, add tests --- .../util/tests/test_vars_from_expressions.py | 59 +++++++++++++++++++ pyomo/util/vars_from_expressions.py | 2 +- 2 files changed, 60 insertions(+), 1 deletion(-) create mode 100644 pyomo/util/tests/test_vars_from_expressions.py diff --git a/pyomo/util/tests/test_vars_from_expressions.py b/pyomo/util/tests/test_vars_from_expressions.py new file mode 100644 index 00000000000..bd60fb70f66 --- /dev/null +++ b/pyomo/util/tests/test_vars_from_expressions.py @@ -0,0 +1,59 @@ +# ____________________________________________________________________________________ +# +# Pyomo: Python Optimization Modeling Objects +# Copyright (c) 2008-2026 National Technology and Engineering Solutions of Sandia, LLC +# Under the terms of Contract DE-NA0003525 with National Technology and Engineering +# Solutions of Sandia, LLC, the U.S. Government retains certain rights in this +# software. This software is distributed under the 3-clause BSD License. +# ____________________________________________________________________________________ + +import pyomo.environ as pyo +from pyomo.common import unittest +from pyomo.util.vars_from_expressions import get_vars, get_vars_from_components + + +class TestVarsFromExpressions(unittest.TestCase): + def test_get_vars(self): + m = pyo.ConcreteModel() + m.x = pyo.Var(list(range(5))) + m.c1 = pyo.Constraint(expr=m.x[0] + m.x[1] == 0) + m.c2 = pyo.Constraint(expr=m.x[1] + m.x[2] == 0) + m.obj = pyo.Objective(expr=m.x[3] + m.x[4]) + + self.assertEqual(list(get_vars(m)), [m.x[0], m.x[1], m.x[2], m.x[3], m.x[4]]) + # verify the default values for active and include_fixed + m.x[0].fix(0) + m.c2.deactivate() + self.assertEqual(list(get_vars(m)), [m.x[1], m.x[3], m.x[4]]) + + def test_get_vars_from_components(self): + m = pyo.ConcreteModel() + m.x = pyo.Var(list(range(5))) + m.c1 = pyo.Constraint(expr=m.x[0] + m.x[1] == 0) + m.c2 = pyo.Constraint(expr=m.x[1] + m.x[2] == 0) + m.obj = pyo.Objective(expr=m.x[3] + m.x[4]) + + self.assertEqual( + list(get_vars_from_components(m, pyo.Constraint)), [m.x[0], m.x[1], m.x[2]] + ) + self.assertEqual( + list(get_vars_from_components(m, pyo.Objective)), [m.x[3], m.x[4]] + ) + self.assertEqual( + list(get_vars_from_components(m, (pyo.Constraint, pyo.Objective))), + [m.x[0], m.x[1], m.x[2], m.x[3], m.x[4]], + ) + + # verify the default values for active and include_fixed + m.x[0].fix(0) + m.c2.deactivate() + self.assertEqual( + list(get_vars_from_components(m, pyo.Constraint)), [m.x[0], m.x[1], m.x[2]] + ) + self.assertEqual( + list(get_vars_from_components(m, pyo.Objective)), [m.x[3], m.x[4]] + ) + self.assertEqual( + list(get_vars_from_components(m, (pyo.Constraint, pyo.Objective))), + [m.x[0], m.x[1], m.x[2], m.x[3], m.x[4]], + ) diff --git a/pyomo/util/vars_from_expressions.py b/pyomo/util/vars_from_expressions.py index 6941532fc67..5f0d2a6fdd3 100644 --- a/pyomo/util/vars_from_expressions.py +++ b/pyomo/util/vars_from_expressions.py @@ -15,7 +15,7 @@ actually in the Block subtree or not. """ -from pyomo.core import Block +from pyomo.core import Block, Constraint, Objective from pyomo.core.expr.visitor import IdentifyVariableVisitor From 01b9e7bb517fb70f71ce3214fc34331e0831f917 Mon Sep 17 00:00:00 2001 From: John Siirola Date: Fri, 29 May 2026 12:21:45 -0600 Subject: [PATCH 4/5] Rename 'seen' -> 'var_cache' for clarity/consistency --- pyomo/core/expr/visitor.py | 16 +++++++++------- pyomo/core/tests/unit/test_visitor.py | 4 ++-- pyomo/util/vars_from_expressions.py | 10 +++++----- 3 files changed, 16 insertions(+), 14 deletions(-) diff --git a/pyomo/core/expr/visitor.py b/pyomo/core/expr/visitor.py index ca20cdf6e0a..64aaa44737c 100644 --- a/pyomo/core/expr/visitor.py +++ b/pyomo/core/expr/visitor.py @@ -1395,7 +1395,9 @@ def identify_components(expr, component_types): class IdentifyVariableVisitor(StreamBasedExpressionVisitor): - def __init__(self, include_fixed=False, named_expression_cache=None, seen=None): + def __init__( + self, include_fixed=False, named_expression_cache=None, var_cache=None + ): """Visitor that collects all unique variables participating in an expression @@ -1410,12 +1412,12 @@ def __init__(self, include_fixed=False, named_expression_cache=None, seen=None): expression as well as information for detecting when the named expression has changed (for cache invalidation). - seen : dict[int, VarData] + var_cache : dict[int, VarData] Dict mapping the :func:`id()` of variables to - :class:`VarData` for all variables "seen" by this walker. - If provided, this dictionary is preserved between calls to - :meth:`walk_expression` (so repeated variables are not - returned more than once) + :class:`VarData` for all variables that have been "seen" by + this walker. If provided, this dictionary is preserved + between calls to :meth:`walk_expression` (so repeated + variables are not returned more than once). """ super().__init__() @@ -1433,7 +1435,7 @@ def __init__(self, include_fixed=False, named_expression_cache=None, seen=None): # processing, and _seen and _exprs are from the parent context. self._expr_stack = [] # cache of "seen" variables: dict(eid: VarData) - self._seen = seen + self._seen = var_cache # The following attribute will be added by initializeWalker: # self._exprs: list of (e, e.expr) for any (nested) named expressions diff --git a/pyomo/core/tests/unit/test_visitor.py b/pyomo/core/tests/unit/test_visitor.py index ba825e1c2ee..50d90077f88 100644 --- a/pyomo/core/tests/unit/test_visitor.py +++ b/pyomo/core/tests/unit/test_visitor.py @@ -253,7 +253,7 @@ def test_identify_vars_linear_expression(self): expr = quicksum([m.x, m.x], linear=True) self.assertEqual(list(identify_variables(expr, include_fixed=False)), [m.x]) - def test_identify_vars_seen_cache(self): + def test_identify_vars_var_cache(self): m = ConcreteModel() m.x = Var() m.y = Var() @@ -267,7 +267,7 @@ def test_identify_vars_seen_cache(self): self.assertEqual(list(v.walk_expression(e2)), [m.y, m.z]) seen = {} - v = IdentifyVariableVisitor(seen=seen) + v = IdentifyVariableVisitor(var_cache=seen) self.assertEqual(list(v.walk_expression(e2)), [m.y, m.z]) self.assertEqual(list(v.walk_expression(e1)), [m.y, m.z, m.x]) diff --git a/pyomo/util/vars_from_expressions.py b/pyomo/util/vars_from_expressions.py index 5f0d2a6fdd3..e93fa1dcdb3 100644 --- a/pyomo/util/vars_from_expressions.py +++ b/pyomo/util/vars_from_expressions.py @@ -42,17 +42,17 @@ def get_vars_from_components( descent_order: Traversal strategy for finding the objects of type ctype """ - seen = {} - visitor = IdentifyVariableVisitor(include_fixed, {}, seen=seen) - for constraint in block.component_data_objects( + var_cache = {} + visitor = IdentifyVariableVisitor(include_fixed, {}, var_cache=var_cache) + for component in block.component_data_objects( ctype, active=active, sort=sort, descend_into=descend_into, descent_order=descent_order, ): - visitor.walk_expression(constraint.expr) - return seen.values() + visitor.walk_expression(component.expr) + return var_cache.values() def get_vars( From 33064ab8055cce6ed43f1a9833333a6595524fec Mon Sep 17 00:00:00 2001 From: John Siirola Date: Fri, 29 May 2026 12:22:36 -0600 Subject: [PATCH 5/5] Only return newly-found variables from IdentifyVariableVisitor.walk_expression() --- pyomo/core/expr/visitor.py | 22 ++++++++++++++++++---- pyomo/core/tests/unit/test_visitor.py | 4 +++- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/pyomo/core/expr/visitor.py b/pyomo/core/expr/visitor.py index 64aaa44737c..6bd3e4a0ee0 100644 --- a/pyomo/core/expr/visitor.py +++ b/pyomo/core/expr/visitor.py @@ -9,6 +9,7 @@ import inspect +import itertools import logging import sys from copy import deepcopy @@ -1401,6 +1402,12 @@ def __init__( """Visitor that collects all unique variables participating in an expression + :meth:`walk_expression` returns a generator of the unique + variables found in the expression. If `var_cache` was + specified, then only the *new* variables found in `expr` are + returned (the full list of all variables is maintained in the + `var_cache` dict). + Parameters ---------- include_fixed : bool @@ -1444,6 +1451,8 @@ def initializeWalker(self, expr): if self._seen is None: self._seen = {} self._expr_stack.append(None) + else: + self._expr_stack.append(len(self._seen)) self._exprs = None if not self.beforeChild(None, expr, 0)[0]: return False, self.finalizeResult(None) @@ -1476,11 +1485,16 @@ def exitNode(self, node, data): def finalizeResult(self, result): seen = self._seen - if self._expr_stack: - assert self._expr_stack.pop() is None - self._seen = None + initial_num_seen = self._expr_stack.pop() assert not self._expr_stack - return seen.values() + if initial_num_seen is None: + self._seen = None + return seen.values() + else: + # Only return the *new* variables found on this walk. This + # relies on dict iteration being in insertion order (which, + # since python 3.7, it is) + return itertools.islice(seen.values(), initial_num_seen, len(seen)) def _merge_obj_lists(self, _seen, _exprs): self._seen.update(_seen) diff --git a/pyomo/core/tests/unit/test_visitor.py b/pyomo/core/tests/unit/test_visitor.py index 50d90077f88..b2a18fa8ceb 100644 --- a/pyomo/core/tests/unit/test_visitor.py +++ b/pyomo/core/tests/unit/test_visitor.py @@ -269,7 +269,9 @@ def test_identify_vars_var_cache(self): seen = {} v = IdentifyVariableVisitor(var_cache=seen) self.assertEqual(list(v.walk_expression(e2)), [m.y, m.z]) - self.assertEqual(list(v.walk_expression(e1)), [m.y, m.z, m.x]) + self.assertEqual(list(seen.values()), [m.y, m.z]) + self.assertEqual(list(v.walk_expression(e1)), [m.x]) + self.assertEqual(list(seen.values()), [m.y, m.z, m.x]) class TestIdentifyParams(unittest.TestCase):