Source code for pyomo.core.plugins.transform.nonnegative_transform

# ____________________________________________________________________________________
#
# 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.core.expr as EXPR

from pyomo.core import (
    nonpyomo_leaf_types,
    TransformationFactory,
    IntegerSet,
    Integers,
    PositiveIntegers,
    NonPositiveIntegers,
    NegativeIntegers,
    NonNegativeIntegers,
    Reals,
    PositiveReals,
    NonNegativeReals,
    NegativeReals,
    NonPositiveReals,
    PercentFraction,
    RealSet,
    Var,
    Set,
    value,
    Binary,
    Constraint,
    Objective,
)
from pyomo.core.base.misc import create_name
from pyomo.core.plugins.transform.util import partial
from pyomo.core.plugins.transform.hierarchy import IsomorphicTransformation
from pyomo.core.plugins.transform.util import collectAbstractComponents


import logging

logger = logging.getLogger('pyomo.core')




class VarmapVisitor(EXPR.ExpressionReplacementVisitor):


    def __init__(self, varmap):
        super(VarmapVisitor, self).__init__()
        self.varmap = varmap


    def visiting_potential_leaf(self, node):
        if node.__class__ in nonpyomo_leaf_types:
            return True, node
        #
        # Clone leaf nodes in the expression tree
        #
        if node.is_variable_type():
            if node.local_name in self.varmap:
                return True, self.varmap[node.local_name]
            else:
                return True, node

        if isinstance(node, EXPR.LinearExpression):
            with EXPR.nonlinear_expression() as expr:
                for c, v in zip(node.linear_coefs, node.linear_vars):
                    if hasattr(v, 'local_name'):
                        expr += c * self.varmap.get(v.local_name)
                    else:
                        expr += c * v
            return True, expr

        return False, None



def _walk_expr(expr, varMap):
    """
    Walks an expression tree, making the replacements defined in varMap
    """
    visitor = VarmapVisitor(varMap)
    return visitor.dfs_postorder_stack(expr)




@TransformationFactory.register(
    "core.nonnegative_vars",
    doc="Create an equivalent model in which all variables lie in the nonnegative orthant.",
)
class NonNegativeTransformation(IsomorphicTransformation):
    """
    Creates a new, equivalent model by forcing all variables to lie in
    the nonnegative orthant by introducing auxiliary variables.
    """



    def __init__(self, **kwds):
        kwds["name"] = kwds.pop("name", "vars")
        super(NonNegativeTransformation, self).__init__(**kwds)

        self.realSets = (
            Reals,
            PositiveReals,
            NonNegativeReals,
            NegativeReals,
            NonPositiveReals,
            PercentFraction,
            RealSet,
        )
        # Intentionally leave out Binary, Boolean, BinarySet, and BooleanSet;
        # we check for those explicitly
        self.discreteSets = (
            IntegerSet,
            Integers,
            PositiveIntegers,
            NonPositiveIntegers,
            NegativeIntegers,
            NonNegativeIntegers,
        )


    def _create_using(self, model, **kwds):
        """
        Force all variables to lie in the nonnegative orthant.

        Required arguments:
            model       The model to transform.

        Optional keyword arguments:
          pos_suffix    The suffix applied to the 'positive' component of
                            converted variables. Default is '_plus'.
          neg_suffix    The suffix applied to the 'positive' component of
                            converted variables. Default is '_minus'.
        """
        #
        # Optional naming schemes
        #
        pos_suffix = kwds.pop("pos_suffix", "_plus")
        neg_suffix = kwds.pop("neg_suffix", "_minus")
        #
        # We first perform an abstract problem transformation. Then, if model
        # data is available, we instantiate the new model. If not, we construct
        # a mapping that can later be used to populate the new model.
        #
        nonneg = model.clone()
        components = collectAbstractComponents(nonneg)

        # Map from variable base names to a {index, rule} map
        constraint_rules = {}

        # Map from variable base names to a rule defining the domains for that
        # variable
        domain_rules = {}

        # Map from variable base names to its set of indices
        var_indices = {}

        # Map from fully qualified variable names to replacement expressions.
        # For now, it is actually a map from a variable name to a closure that
        # must later be evaluated with a model containing the replacement
        # variables.
        var_map = {}

        #
        # Get the constraints that enforce the bounds and domains of each
        # variable
        #
        for var_name in components["Var"]:
            var = nonneg.__getattribute__(var_name)

            # Individual bounds and domains
            orig_bounds = {}
            orig_domain = {}

            # New indices
            indices = set()

            # Map from constraint names to a constraint rule.
            constraints = {}

            # Map from variable indices to a domain
            domains = {}

            for ndx in var:
                # Fully qualified variable name
                vname = create_name(str(var_name), ndx)

                # We convert each index to a string to avoid difficult issues
                # regarding appending a suffix to tuples.
                #
                # If the index is None, this casts the index to a string,
                # which doesn't match up with how Pyomo treats None indices
                # internally. Replace with "" to be consistent.
                if ndx is None:
                    v_ndx = ""
                else:
                    v_ndx = str(ndx)

                # Get the variable bounds
                lb = value(var[ndx].lb)
                ub = value(var[ndx].ub)
                orig_bounds[ndx] = (lb, ub)

                # Get the variable domain
                if var[ndx].domain is not None:
                    orig_domain[ndx] = var[ndx].domain
                else:
                    orig_domain[ndx] = var.domain

                # Determine the replacement expression. Either a new single
                # variable with the same attributes, or a sum of two new
                # variables.
                #
                # If both the bounds and domain allow for negative values,
                # replace the variable with the sum of nonnegative ones.

                bounds_neg = (
                    orig_bounds[ndx] == (None, None)
                    or orig_bounds[ndx][0] is None
                    or orig_bounds[ndx][0] < 0
                )
                domain_neg = (
                    orig_domain[ndx] is None
                    or orig_domain[ndx].bounds()[0] is None
                    or orig_domain[ndx].bounds()[0] < 0
                )
                if bounds_neg and domain_neg:
                    # Make two new variables.
                    posVarSuffix = "%s%s" % (v_ndx, pos_suffix)
                    negVarSuffix = "%s%s" % (v_ndx, neg_suffix)

                    new_indices = (posVarSuffix, negVarSuffix)

                    # Replace the original variable with a sum expression
                    expr_dict = {posVarSuffix: 1, negVarSuffix: -1}
                else:
                    # Add the new index. Lie if is 'None', since Pyomo treats
                    # 'None' specially as a key.
                    #
                    # More lies: don't let a blank index exist. Replace it with
                    # '_'. I don't actually have a justification for this other
                    # than that allowing "" as a key will eventually almost
                    # certainly lead to a strange bug.
                    if v_ndx is None:
                        t_ndx = "None"
                    elif v_ndx == "":
                        t_ndx = "_"
                    else:
                        t_ndx = v_ndx

                    new_indices = (t_ndx,)

                    # Replace the original variable with a sum expression
                    expr_dict = {t_ndx: 1}

                # Add the new indices
                for x in new_indices:
                    indices.add(x)

                # Replace the original variable with an expression
                var_map[vname] = partial(self.sumRule, var_name, expr_dict)

                # Enforce bounds as constraints
                if orig_bounds[ndx] != (None, None):
                    cname = "%s_%s" % (vname, "bounds")
                    tmp = orig_bounds[ndx]
                    constraints[cname] = partial(
                        self.boundsConstraintRule, tmp[0], tmp[1], var_name, expr_dict
                    )

                # Enforce the bounds of the domain as constraints
                if orig_domain[ndx] != None:
                    cname = "%s_%s" % (vname, "domain_bounds")
                    tmp = orig_domain[ndx].bounds()
                    constraints[cname] = partial(
                        self.boundsConstraintRule, tmp[0], tmp[1], var_name, expr_dict
                    )

                # Domain will either be NonNegativeReals, NonNegativeIntegers,
                # or Binary. We consider Binary because some solvers may
                # optimize over binary variables.
                if var[ndx].is_continuous():
                    for x in new_indices:
                        domains[x] = NonNegativeReals
                elif var[ndx].is_binary():
                    for x in new_indices:
                        domains[x] = Binary
                elif var[ndx].is_integer():
                    for x in new_indices:
                        domains[x] = NonNegativeIntegers
                else:
                    logger.warning(
                        "Warning: domain '%s' not recognized, "
                        "defaulting to 'NonNegativeReals'" % (var.domain,)
                    )
                    for x in new_indices:
                        domains[x] = NonNegativeReals

            constraint_rules[var_name] = constraints
            domain_rules[var_name] = partial(self.exprMapRule, domains)
            var_indices[var_name] = indices

        # Remove all existing variables.
        toRemove = []
        for attr_name, attr in nonneg.__dict__.items():
            if isinstance(attr, Var):
                toRemove.append(attr_name)
        for attr_name in toRemove:
            nonneg.__delattr__(attr_name)

        # Add the sets defining the variables, then the variables
        for k, v in var_indices.items():
            sname = "%s_indices" % k
            nonneg.__setattr__(sname, Set(initialize=v))
            nonneg.__setattr__(
                k,
                Var(
                    nonneg.__getattribute__(sname),
                    domain=domain_rules[k],
                    bounds=(0, None),
                ),
            )

        # Construct the model to get the variables and their indices
        # recognized in the model
        ##nonneg = nonneg.create()

        # Safe to evaluate the modifiedVars mapping
        for var in var_map:
            var_map[var] = var_map[var](nonneg)

        # Map from constraint base names to maps from indices to expressions
        constraintExprs = {}

        #
        # Convert all modified variables in all constraints in the original
        # problem
        #
        for conName in components["Constraint"]:
            con = nonneg.__getattribute__(conName)

            # Map from constraint indices to a corrected expression
            exprMap = {}

            for ndx, cdata in con._data.items():
                lower = _walk_expr(cdata.lower, var_map)
                body = _walk_expr(cdata.body, var_map)
                upper = _walk_expr(cdata.upper, var_map)

                # Lie if ndx is None. Pyomo treats 'None' indices specially.
                if ndx is None:
                    ndx = "None"

                # Cast indices to strings, otherwise tuples ruin everything
                exprMap[str(ndx)] = (lower, body, upper)

            # Add to list of expression maps
            constraintExprs[conName] = exprMap

        # Map from constraint base names to maps from indices to expressions
        objectiveExprs = {}

        #
        # Convert all modified variables in all objectives in the original
        # problem
        #
        for objName in components["Objective"]:
            obj = nonneg.__getattribute__(objName)

            # Map from objective indices to a corrected expression
            exprMap = {}

            for ndx, odata in obj._data.items():
                exprMap[ndx] = _walk_expr(odata.expr, var_map)

            # Add to list of expression maps
            objectiveExprs[objName] = exprMap

        # Make the modified original constraints
        for conName, ruleMap in constraintExprs.items():
            # Make the set of indices
            sname = conName + "_indices"
            _set = Set(initialize=ruleMap.keys())
            nonneg.__setattr__(sname, _set)
            _set.construct()

            # Define the constraint
            _con = Constraint(
                nonneg.__getattribute__(sname), rule=partial(self.exprMapRule, ruleMap)
            )
            nonneg.__setattr__(conName, _con)
            _con.construct()

        # Make the bounds constraints
        for varName, ruleMap in constraint_rules.items():
            conName = varName + "_constraints"
            # Make the set of indices
            sname = conName + "_indices"
            _set = Set(initialize=ruleMap.keys())
            nonneg.__setattr__(sname, _set)
            _set.construct()

            # Define the constraint
            _con = Constraint(
                nonneg.__getattribute__(sname),
                rule=partial(self.delayedExprMapRule, ruleMap),
            )
            nonneg.__setattr__(conName, _con)
            _con.construct()

        # Make the objectives
        for objName, ruleMap in objectiveExprs.items():
            # Make the set of indices
            sname = objName + "_indices"
            _set = Set(initialize=ruleMap.keys())
            nonneg.__setattr__(sname, _set)
            _set.construct()

            # Define the constraint
            _obj = Objective(
                nonneg.__getattribute__(sname), rule=partial(self.exprMapRule, ruleMap)
            )
            nonneg.__setattr__(objName, _obj)
            _obj.construct()

        return nonneg



    @staticmethod
    def boundsConstraintRule(lb, ub, attr, vars, model):
        """
        Produces 'lb < x^+ - x^- < ub' style constraints. Designed to
        be made a closer through functools.partial, across lb, ub, attr,
        and vars. vars is a {varname: coefficient} dictionary. attr is the
        base variable name; that is, X[1] would be referenced by

          model.__getattribute__('X')[1]

        and so attr='X', and 1 is a key of vars.

        """
        if lb is None and ub is None:
            return Constraint.Skip
        return (
            lb,
            sum(c * model.__getattribute__(attr)[v] for (v, c) in vars.items()),
            ub,
        )


    @staticmethod
    def noConstraint(*args):
        return None



    @staticmethod
    def sumRule(attr, vars, model):
        """
        Returns a sum expression.
        """
        return sum(c * model.__getattribute__(attr)[v] for (v, c) in vars.items())




    @staticmethod
    def exprMapRule(ruleMap, model, ndx=None):
        """Rule intended to return expressions from a lookup table"""
        return ruleMap[ndx]




    @staticmethod
    def delayedExprMapRule(ruleMap, model, ndx=None):
        """
        Rule intended to return expressions from a lookup table. Each entry
        in the lookup table is a functor that needs to be evaluated before
        returning.
        """
        return ruleMap[ndx](model)