KKT Transform

pyomo/core/plugins/transform/__init__.py

@@ -21,4 +21,5 @@
     scaling,
     logical_to_linear,
     lp_dual,
+    kkt,
 )

pyomo/core/plugins/transform/kkt.py

@@ -0,0 +1,374 @@
+# ____________________________________________________________________________________
+#
+# 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.
+# ____________________________________________________________________________________
+
+
+from pyomo.common.autoslots import AutoSlots
+from pyomo.common.collections import ComponentMap, ComponentSet
+from pyomo.common.config import ConfigDict, ConfigValue
+from pyomo.core import (
+    Block,
+    Constraint,
+    ConstraintList,
+    Expression,
+    NonNegativeReals,
+    Objective,
+    RangeSet,
+    Reals,
+    Set,
+    TransformationFactory,
+    Var,
+    maximize,
+    minimize,
+)
+from pyomo.core.base.constraint import ConstraintData
+from pyomo.core.base.var import VarData
+from pyomo.core.expr.calculus.diff_with_pyomo import reverse_sd
+from pyomo.mpec import ComplementarityList, complements
+from pyomo.util.vars_from_expressions import get_vars_from_components
+from pyomo.util.config_domains import ComponentDataSet
+
+
+class _KKTReformulationData(AutoSlots.Mixin):
+    __slots__ = (
+        "equality_multiplier_from_con",
+        "equality_con_from_multiplier",
+        "inequality_multiplier_from_con",
+        "inequality_con_from_multiplier",
+        "var_bound_multiplier_index_to_con",
+        "equality_multiplier_index_to_con",
+        "inequality_multiplier_index_to_con",
+        "equality_con_to_expr",
+        "inequality_con_to_expr",
+        "ranged_constraints",
+    )
+
+    def __init__(self):
+        self.equality_multiplier_from_con = ComponentMap()
+        self.equality_con_from_multiplier = ComponentMap()
+        self.inequality_multiplier_from_con = ComponentMap()
+        self.inequality_con_from_multiplier = ComponentMap()
+
+        self.var_bound_multiplier_index_to_con = {}
+        self.equality_multiplier_index_to_con = {}
+        self.inequality_multiplier_index_to_con = {}
+        self.equality_con_to_expr = ComponentMap()
+        self.inequality_con_to_expr = ComponentMap()
+
+        self.ranged_constraints = ComponentSet()
+
+
+Block.register_private_data_initializer(_KKTReformulationData)
+
+
+@TransformationFactory.register(
+    'core.kkt', 'Generate KKT reformulation of the given model'
+)
+class NonLinearProgrammingKKT:
+    CONFIG = ConfigDict("core.kkt")
+    CONFIG.declare(
+        'kkt_block_name',
+        ConfigValue(
+            default='kkt',
+            doc="""
+           Name of the block on which the kkt variables and constraints will be stored. 
+           """,
+        ),
+    )
+    CONFIG.declare(
+        'parametrize_wrt',
+        ConfigValue(
+            default=[],
+            domain=ComponentDataSet(Var),
+            description='Vars to treat as data for the purposes of generating KKT reformulation',
+            doc="""
+            Optional list of Vars to be treated as data while generating the KKT reformulation.
+            """,
+        ),
+    )
+
+    def apply_to(self, model, **kwds):
+        """
+        Reformulate model with KKT conditions.
+        """
+        config = self.CONFIG(kwds.pop('options', {}))
+        config.set_value(kwds)
+
+        if hasattr(model, config.kkt_block_name):
+            raise ValueError(
+                f"""model already has an attribute with the 
+                specified kkt_block_name: '{config.kkt_block_name}'"""
+            )
+
+        # we should check that all vars the user fixed are included
+        # in parametrize_wrt
+        params = config.parametrize_wrt
+        vars_in_cons = ComponentSet(
+            get_vars_from_components(model, Constraint, active=True, descend_into=True)
+        )
+        vars_in_obj = ComponentSet(
+            get_vars_from_components(model, Objective, active=True, descend_into=True)
+        )
+        vars_in_model = vars_in_cons | vars_in_obj
+        fixed_vars_in_model = ComponentSet(v for v in vars_in_model if v.is_fixed())
+        missing = fixed_vars_in_model - params
+        if missing:
+            raise ValueError("All fixed variables must be included in parametrize_wrt.")
+
+        # we should also check that all vars the user passes in parametrize_wrt
+        # exist on an active constraint or objective within the model
+        unknown = params - vars_in_model
+        if unknown:
+            raise ValueError(
+                "A variable passed in parametrize_wrt does not exist on an "
+                "active constraint or objective within the model."
+            )
+
+        kkt_block = Block(concrete=True)
+        kkt_block.parametrize_wrt = params
+        self._reformulate(model, kkt_block)
+        model.add_component(config.kkt_block_name, kkt_block)
+        return model
+
+    def _reformulate(self, model, kkt_block):
+        active_objs = list(
+            model.component_data_objects(Objective, active=True, descend_into=True)
+        )
+        if len(active_objs) != 1:
+            raise ValueError(
+                f"model must have only one active objective; found {len(active_objs)}"
+            )
+
+        self._construct_lagrangean(model, kkt_block)
+        self._enforce_stationarity_conditions(kkt_block)
+        self._enforce_complementarity_conditions(model, kkt_block)
+
+        active_objs[0].deactivate()
+        kkt_block.dummy_obj = Objective(expr=1)
+
+        return model
+
+    def _construct_lagrangean(self, model, kkt_block):
+        # we need to loop through the model and store the
+        # equality and inequality constraints
+        info = model.private_data()
+        equality_cons = []
+        inequality_cons = []
+        for con in model.component_data_objects(
+            Constraint, descend_into=True, active=True
+        ):
+            lower, body, upper = con.to_bounded_expression()
+            if con.equality:
+                equality_cons.append(con)
+                info.equality_con_to_expr[con] = upper - body
+            else:
+                if lower is not None:
+                    inequality_cons.append((con, "lb"))
+                    info.inequality_con_to_expr.setdefault(con, {})["lb"] = lower - body
+                if upper is not None:
+                    inequality_cons.append((con, "ub"))
+                    info.inequality_con_to_expr.setdefault(con, {})["ub"] = body - upper
+
+                    # lower is not None and upper is not None -> ranged constraint
+                    if lower is not None:
+                        # we want to keep track of the ranged constraints because the mapping between
+                        # multipliers and ranged constraints will be a tuple (to indicate bound as well)
+                        # instead of simply the model object
+                        info.ranged_constraints.add(con)
+
+        kkt_block.equality_cons_list = list(equality_cons)
+        kkt_block.gamma_set = RangeSet(0, len(equality_cons) - 1)
+        kkt_block.gamma = Var(kkt_block.gamma_set, domain=Reals)
+        info.equality_multiplier_index_to_con = dict(
+            enumerate(kkt_block.equality_cons_list)
+        )
+        kkt_block.inequality_cons_list = list(inequality_cons)
+        kkt_block.alpha_con_set = RangeSet(0, len(inequality_cons) - 1)
+        kkt_block.alpha_con = Var(kkt_block.alpha_con_set, domain=NonNegativeReals)
+        info.inequality_multiplier_index_to_con = dict(
+            enumerate(kkt_block.inequality_cons_list)
+        )
+
+        # we also need to consider inequality constraints
+        # formed by the user specifying variable bounds
+        var_bound_sides = []
+        vars_in_cons = ComponentSet(
+            get_vars_from_components(model, Constraint, active=True, descend_into=True)
+        )
+        vars_in_obj = ComponentSet(
+            get_vars_from_components(model, Objective, active=True, descend_into=True)
+        )
+        kkt_block.var_set = vars_in_cons | vars_in_obj
+        kkt_block.var_set = kkt_block.var_set - kkt_block.parametrize_wrt
+        for var in kkt_block.var_set:
+            if var.has_lb():
+                var_bound_sides.append((var, "lb"))
+            if var.has_ub():
+                var_bound_sides.append((var, "ub"))
+        kkt_block.var_bound_set = RangeSet(0, len(var_bound_sides) - 1)
+        kkt_block.alpha_var_bound = Var(
+            kkt_block.var_bound_set, domain=NonNegativeReals
+        )
+
+        info.var_bound_multiplier_index_to_con = dict(enumerate(var_bound_sides))
+
+        # indexing the inequality constraint expressions will help
+        # with constructing the lagrangean later
+        def _var_bound_expr_rule(kkt, i):
+            var, side = info.var_bound_multiplier_index_to_con[i]
+            return (var.lb - var) if side == "lb" else (var - var.ub)
+
+        kkt_block.var_bound_expr = Expression(
+            kkt_block.var_bound_set, rule=_var_bound_expr_rule
+        )
+
+        # we will construct the lagrangean by first adding the objective,
+        # and then looping through and adding the product of each constraint and
+        # the corresponding multiplier
+        obj = list(
+            model.component_data_objects(Objective, active=True, descend_into=True)
+        )
+        # Note: maximize is -1 and minimize is +1
+        lagrangean = obj[0].sense * obj[0].expr
+
+        for index, con in enumerate(kkt_block.equality_cons_list):
+            lagrangean += info.equality_con_to_expr[con] * kkt_block.gamma[index]
+            info.equality_con_from_multiplier[kkt_block.gamma[index]] = con
+            info.equality_multiplier_from_con[con] = kkt_block.gamma[index]
+
+        for index, (con, bound) in enumerate(kkt_block.inequality_cons_list):
+            lagrangean += (
+                info.inequality_con_to_expr[con][bound] * kkt_block.alpha_con[index]
+            )
+            # mappings for ranged constraints will consider bounds as well
+            if con in info.ranged_constraints:
+                info.inequality_con_from_multiplier[kkt_block.alpha_con[index]] = (
+                    con,
+                    bound,
+                )
+                info.inequality_multiplier_from_con[(con, bound)] = kkt_block.alpha_con[
+                    index
+                ]
+            else:
+                info.inequality_con_from_multiplier[kkt_block.alpha_con[index]] = con
+                info.inequality_multiplier_from_con[con] = kkt_block.alpha_con[index]
+
+        for i in kkt_block.var_bound_set:
+            lagrangean += kkt_block.var_bound_expr[i] * kkt_block.alpha_var_bound[i]
+            # mappings for ranged constraints built from variable bounds
+            var, bound = info.var_bound_multiplier_index_to_con[i]
+            info.inequality_con_from_multiplier[kkt_block.alpha_var_bound[i]] = (
+                var,
+                bound,
+            )
+            info.inequality_multiplier_from_con[(var, bound)] = (
+                kkt_block.alpha_var_bound[i]
+            )
+
+        kkt_block.lagrangean = Expression(expr=lagrangean)
+
+    def _enforce_stationarity_conditions(self, kkt_block):
+        deriv_lagrangean = reverse_sd(kkt_block.lagrangean.expr)
+        kkt_block.stationarity_conditions = ConstraintList()
+        for var in kkt_block.var_set:
+            kkt_block.stationarity_conditions.add(deriv_lagrangean[var] == 0)
+
+    def _enforce_complementarity_conditions(self, model, kkt_block):
+        info = model.private_data()
+        kkt_block.complements = ComplementarityList()
+        for index, (con, bound) in enumerate(kkt_block.inequality_cons_list):
+            kkt_block.complements.add(
+                complements(
+                    kkt_block.alpha_con[index] >= 0,
+                    info.inequality_con_to_expr[con][bound] <= 0,
+                )
+            )
+        # we also need to consider the inequality constraints
+        # formed from the user specifying the variable bounds
+        for i in kkt_block.var_bound_set:
+            kkt_block.complements.add(
+                complements(
+                    kkt_block.alpha_var_bound[i] >= 0, kkt_block.var_bound_expr[i] <= 0
+                )
+            )
+
+    def get_constraint_from_multiplier(self, model, multiplier_var):
+        """
+        Return the constraint or variable bound corresponding to a KKT multiplier variable.
+
+        Parameters
+        ----------
+        model: ConcreteModel
+            The model on which the kkt transformation was applied
+        multiplier_var: Var
+            A KKT multiplier created by the transformation.
+
+        Returns
+        -------
+        Constraint or Tuple
+            - Constraint object for simple constraints
+            - (Constraint, bound) tuple for ranged constraints
+            - (Var, bound) tuple for variable bounds
+        """
+
+        info = model.private_data()
+        if multiplier_var in info.equality_con_from_multiplier:
+            return info.equality_con_from_multiplier[multiplier_var]
+        if multiplier_var in info.inequality_con_from_multiplier:
+            # if this multiplier var maps to a ranged constraint, we will return a tuple
+            # so that we can indicate which bound the multiplier var maps to
+            return info.inequality_con_from_multiplier[multiplier_var]
+        raise ValueError(
+            f"The KKT multiplier: {multiplier_var.name}, does not exist on {model.name}."
+        )
+
+    def get_multiplier_from_constraint(self, model, component):
+        """
+        Return the multiplier for the constraint.
+
+        Parameters
+        ----------
+        model: ConcreteModel
+            The model on which the kkt transformation was applied to
+        component: Constraint or Variable
+
+        Returns
+        -------
+        VarData | tuple[VarData | None, VarData | None]
+            The KKT multiplier(s) corresponding to the component.
+            For ranged constraints/variables, returns (lb_mult, ub_mult),
+            where an entry is None if that bound doesn't exist.
+        """
+        info = model.private_data()
+        if isinstance(component, ConstraintData):
+            con = component
+            if con in info.equality_multiplier_from_con:
+                return info.equality_multiplier_from_con[con]
+            elif con in info.inequality_multiplier_from_con:
+                return info.inequality_multiplier_from_con[con]
+            elif con in info.ranged_constraints:
+                lb_mult = info.inequality_multiplier_from_con.get((con, 'lb'))
+                ub_mult = info.inequality_multiplier_from_con.get((con, 'ub'))
+                return (lb_mult, ub_mult)
+            else:
+                raise ValueError(
+                    f"Constraint '{con.name}' does not exist on {model.name}."
+                )
+        elif isinstance(component, VarData):
+            var = component
+            lb_mult = info.inequality_multiplier_from_con.get((var, 'lb'))
+            ub_mult = info.inequality_multiplier_from_con.get((var, 'ub'))
+            if lb_mult is None and ub_mult is None:
+                raise ValueError(
+                    f"No multipliers exist for variable '{var.name}' on {model.name}."
+                )
+            return (lb_mult, ub_mult)
+        else:
+            raise ValueError(
+                f"Component '{component.name}' does not exist on {model.name}."
+            )