# ____________________________________________________________________________________
#
# 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 logging
logger = logging.getLogger('pyomo.core')
from pyomo.common.collections import ComponentMap, ComponentSet
from pyomo.common.log import is_debug_set
import pyomo.core.expr as EXPR
from pyomo.core.base import (
Transformation,
TransformationFactory,
Connector,
Constraint,
ConstraintList,
Var,
SortComponents,
)
from pyomo.core.base.connector import ConnectorData, ScalarConnector
[docs]
@TransformationFactory.register(
'core.expand_connectors',
doc="Expand all connectors in the model to simple constraints",
)
class ExpandConnectors(Transformation):
def _apply_to(self, instance, **kwds):
if is_debug_set(logger): # pragma:nocover
logger.debug("Calling ConnectorExpander")
connectorsFound = False
for c in instance.component_data_objects(Connector):
connectorsFound = True
break
if not connectorsFound:
return
if is_debug_set(logger): # pragma:nocover
logger.debug(" Connectors found!")
#
# At this point, there are connectors in the model, so we must
# look for constraints that involve connectors and expand them.
#
# List of the connectors in the order in which we found them
# (this should be deterministic, provided that the user's model
# is deterministic)
connector_list = []
# list of constraints with connectors: tuple(constraint, connector_set)
# (this should be deterministic, provided that the user's model
# is deterministic)
constraint_list = []
# ID of the next connector group (set of matched connectors)
groupID = 0
# connector_groups stars out as a dict of {id(set): (groupID, set)}
# If you sort by the groupID, then this will be deterministic.
connector_groups = dict()
# map of connector to the set of connectors that must match it
matched_connectors = ComponentMap()
# The set of connectors found in the current constraint
found = ComponentSet()
connector_types = set([ScalarConnector, ConnectorData])
for constraint in instance.component_data_objects(
Constraint, sort=SortComponents.deterministic
):
ref = None
for c in EXPR.identify_components(constraint.body, connector_types):
found.add(c)
if c in matched_connectors:
if ref is None:
# The first connector in this constraint has
# already been seen. We will use that Set as
# the reference
ref = matched_connectors[c]
elif ref is not matched_connectors[c]:
# We already have a reference group; merge this
# new group into it.
#
# Optimization: this merge is linear in the size
# of the src set. If the reference set is
# smaller, save time by switching to a new
# reference set.
src = matched_connectors[c]
if len(ref) < len(src):
ref, src = src, ref
ref.update(src)
for _ in src:
matched_connectors[_] = ref
del connector_groups[id(src)]
# else: pass
# The new group *is* the reference group;
# there is nothing to do.
else:
# The connector has not been seen before.
connector_list.append(c)
if ref is None:
# This is the first connector in the constraint:
# start a new reference set.
ref = ComponentSet()
connector_groups[id(ref)] = (groupID, ref)
groupID += 1
# This connector hasn't been seen. Record it.
ref.add(c)
matched_connectors[c] = ref
if ref is not None:
constraint_list.append((constraint, found))
found = ComponentSet()
# Validate all connector sets and expand the empty ones
known_conn_sets = {}
for groupID, conn_set in sorted(connector_groups.values()):
known_conn_sets[id(conn_set)] = self._validate_and_expand_connector_set(
conn_set
)
# Expand each constraint
for constraint, conn_set in constraint_list:
cList = ConstraintList()
constraint.parent_block().add_component(
'%s.expanded' % (constraint.getname(fully_qualified=False),), cList
)
connId = next(iter(conn_set))
ref = known_conn_sets[id(matched_connectors[connId])]
for k, v in sorted(ref.items()):
if v[1] >= 0:
_iter = v[0]
else:
_iter = (v[0],)
for idx in _iter:
substitution = {}
for c in conn_set:
if v[1] >= 0:
new_v = c.vars[k][idx]
elif k in c.aggregators:
new_v = c.vars[k].add()
else:
new_v = c.vars[k]
substitution[id(c)] = new_v
cList.add(
(
constraint.lower,
EXPR.clone_expression(constraint.body, substitution),
constraint.upper,
)
)
constraint.deactivate()
# Now, go back and implement VarList aggregators
for conn in connector_list:
block = conn.parent_block()
for var, aggregator in conn.aggregators.items():
c = Constraint(expr=aggregator(block, conn.vars[var]))
block.add_component(
'%s.%s.aggregate' % (conn.getname(fully_qualified=True), var), c
)
def _validate_and_expand_connector_set(self, connectors):
ref = {}
# First, go through the connectors and get the superset of all fields
for c in connectors:
for k, v in c.vars.items():
if k in ref:
# We have already seen this var
continue
if v is None:
# This is an implicit var
continue
# OK: New var, so add it to the reference list
_len = (
# -3 if v is None else
-2
if k in c.aggregators
else (
-1
if not hasattr(v, 'is_indexed') or not v.is_indexed()
else len(v)
)
)
ref[k] = (v, _len, c)
if not ref:
logger.warning(
"Cannot identify a reference connector: no connectors "
"in the connector set have assigned variables:\n\t(%s)"
% (', '.join(sorted(c.name for c in connectors)),)
)
return ref
# Now make sure that connectors match
empty_or_partial = []
for c in connectors:
c_is_partial = False
if not c.vars:
# This is an empty connector and should be defined with
# "auto" vars
empty_or_partial.append(c)
continue
for k, v in ref.items():
if k not in c.vars:
raise ValueError(
"Connector mismatch: Connector '%s' missing variable "
"'%s' (appearing in reference connector '%s')"
% (c.name, k, v[2].name)
)
_v = c.vars[k]
if _v is None:
if not c_is_partial:
empty_or_partial.append(c)
c_is_partial = True
continue
_len = (
-3
if _v is None
else (
-2
if k in c.aggregators
else (
-1
if not hasattr(_v, 'is_indexed') or not _v.is_indexed()
else len(_v)
)
)
)
if (_len >= 0) ^ (v[1] >= 0):
raise ValueError(
"Connector mismatch: Connector variable '%s' mixing "
"indexed and non-indexed targets on connectors '%s' "
"and '%s'" % (k, v[2].name, c.name)
)
if _len >= 0 and _len != v[1]:
raise ValueError(
"Connector mismatch: Connector variable '%s' index "
"mismatch (%s elements in reference connector '%s', "
"but %s elements in connector '%s')"
% (k, v[1], v[2].name, _len, c.name)
)
if v[1] >= 0 and len(v[0].index_set() ^ _v.index_set()):
raise ValueError(
"Connector mismatch: Connector variable '%s' has "
"mismatched indices on connectors '%s' and '%s'"
% (k, v[2].name, c.name)
)
# as we are adding things to the model, sort by key so that
# the order things are added is deterministic
sorted_refs = sorted(ref.items())
if len(empty_or_partial) > 1:
# This is expensive (names aren't cheap), but does result in
# a deterministic ordering
empty_or_partial.sort(key=lambda x: x.getname(fully_qualified=True))
# Fill in any empty connectors
for c in empty_or_partial:
block = c.parent_block()
for k, v in sorted_refs:
if k in c.vars and c.vars[k] is not None:
continue
if v[1] >= 0:
idx = (v[0].index_set(),)
else:
idx = ()
var_args = {}
try:
var_args['domain'] = v[0].domain
except AttributeError:
pass
try:
var_args['bounds'] = v[0].bounds
except AttributeError:
pass
new_var = Var(*idx, **var_args)
block.add_component(
'%s.auto.%s' % (c.getname(fully_qualified=True), k), new_var
)
if idx:
for i in idx[0]:
new_var[i].domain = v[0][i].domain
new_var[i].setlb(v[0][i].lb)
new_var[i].setub(v[0][i].ub)
c.vars[k] = new_var
return ref