model_utils.py

#  ___________________________________________________________________________
#
#  Pyomo: Python Optimization Modeling Objects
#  Copyright (c) 2008-2024
#  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

import pyomo.environ as pyo
from pyomo.core.expr import replace_expressions, identify_mutable_parameters
from pyomo.core.base.var import IndexedVar
from pyomo.core.base.param import IndexedParam
from pyomo.common.collections import ComponentMap

from pyomo.environ import ComponentUID

logger = logging.getLogger(__name__)


def convert_params_to_vars(model, param_CUIDs=None, fix_vars=False):
    """
    Convert select Params to Vars

    Parameters
    ----------
    model : Pyomo concrete model
        Original model
    param_CUIDs : list of strings
        List of parameter CUIDs to convert, if None then all Params are converted
    fix_vars : bool
        Fix the new variables, default is False

    Returns
    -------
    model : Pyomo concrete model
        Model with select Params converted to Vars
    """

    model = model.clone()

    if param_CUIDs is None:
        param_CUIDs = [
            ComponentUID(param) for param in model.component_data_objects(pyo.Param)
        ]

    # keep a list of the parameter CUIDs in the case of indexing
    indexed_param_CUIDs = []

    # Convert Params to Vars, unfix Vars, and create a substitution map
    substitution_map = {}
    comp_map = ComponentMap()
    for param_CUID in param_CUIDs:

        # Leverage the parser in ComponentUID to locate the component.
        theta_object = param_CUID.find_component_on(model)

        # Param
        if theta_object.is_parameter_type():

            # Delete Param, add Var
            vals = theta_object.extract_values()
            model.del_component(theta_object)
            model.add_component(theta_object.name, pyo.Var(initialize=vals[None]))

            # Update substitution map
            theta_var_cuid = ComponentUID(theta_object.name)
            theta_var_object = theta_var_cuid.find_component_on(model)
            substitution_map[id(theta_object)] = theta_var_object
            comp_map[theta_object] = theta_var_object

        # Indexed Param -- Delete Param, add Var
        elif isinstance(theta_object, IndexedParam):

            # save Param values
            vals = theta_object.extract_values()

            # get indexed Params
            param_theta_objects = [theta_obj for _, theta_obj in theta_object.items()]

            # get indexed Param names
            indexed_param_CUIDs.extend(
                ComponentUID(theta_obj) for theta_obj in theta_object.values()
            )

            # delete Param
            model.del_component(theta_object)

            # add Var w/ previous Param values
            index_name = theta_object.index_set().name
            index_cuid = ComponentUID(index_name)
            index_object = index_cuid.find_component_on(model)
            model.add_component(
                theta_object.name, pyo.Var(index_object, initialize=vals)
            )

            # Update substitution map (map each indexed param to indexed var)
            theta_var_cuid = ComponentUID(theta_object.name)
            theta_var_object = theta_var_cuid.find_component_on(model)
            comp_map[theta_object] = theta_var_object
            var_theta_objects = list(theta_var_object.values())
            for param_theta_obj, var_theta_obj in zip(
                param_theta_objects, var_theta_objects
            ):
                substitution_map[id(param_theta_obj)] = var_theta_obj
                comp_map[param_theta_obj] = var_theta_obj

        # Var or Indexed Var
        elif isinstance(theta_object, IndexedVar) or theta_object.is_variable_type():
            theta_var_object = theta_object

        else:
            logger.warning("%s is not a Param or Var on the model", (param_name))
            return model

        if fix_vars:
            theta_var_object.fix()
        else:
            theta_var_object.unfix()

    # If no substitutions are needed, return the model
    if len(substitution_map) == 0:
        return model

    # Update the list of param_CUIDs if the parameters were indexed
    if len(indexed_param_CUIDs) > 0:
        param_CUIDs = indexed_param_CUIDs

    # convert to a set for look up efficiency
    param_CUIDs_set = set(param_CUIDs)

    # Convert Params to Vars in Expressions
    for expr in model.component_data_objects(pyo.Expression):
        if expr.active and any(
            ComponentUID(v) in param_CUIDs_set
            for v in identify_mutable_parameters(expr)
        ):
            new_expr = replace_expressions(expr=expr, substitution_map=substitution_map)
            model.del_component(expr)
            model.add_component(expr.name, pyo.Expression(rule=new_expr))

    # Convert Params to Vars in Constraint expressions
    num_constraints = len(list(model.component_objects(pyo.Constraint, active=True)))
    if num_constraints > 0:
        model.constraints = pyo.ConstraintList()
        for c in model.component_data_objects(pyo.Constraint):
            if c.active and any(
                ComponentUID(v) in param_CUIDs_set
                for v in identify_mutable_parameters(c.expr)
            ):
                if c.equality:
                    model.constraints.add(
                        replace_expressions(
                            expr=c.lower, substitution_map=substitution_map
                        )
                        == replace_expressions(
                            expr=c.body, substitution_map=substitution_map
                        )
                    )
                elif c.lower is not None:
                    model.constraints.add(
                        replace_expressions(
                            expr=c.lower, substitution_map=substitution_map
                        )
                        <= replace_expressions(
                            expr=c.body, substitution_map=substitution_map
                        )
                    )
                elif c.upper is not None:
                    model.constraints.add(
                        replace_expressions(
                            expr=c.upper, substitution_map=substitution_map
                        )
                        >= replace_expressions(
                            expr=c.body, substitution_map=substitution_map
                        )
                    )
                else:
                    raise ValueError(
                        "Unable to parse constraint to convert params to vars."
                    )
                c.deactivate()

    # Convert Params to Vars in Objective expressions
    for obj in model.component_data_objects(pyo.Objective):
        if obj.active and any(
            ComponentUID(v) in param_CUIDs_set for v in identify_mutable_parameters(obj)
        ):
            expr = replace_expressions(expr=obj.expr, substitution_map=substitution_map)
            model.del_component(obj)
            model.add_component(obj.name, pyo.Objective(rule=expr, sense=obj.sense))

    # Convert Params to Vars in Suffixes
    for s in model.component_objects(pyo.Suffix):
        current_keys = list(s.keys())
        for c in current_keys:
            if c in comp_map:
                s[comp_map[c]] = s.pop(c)

        assert len(current_keys) == len(s.keys())

    # print('--- Updated Model ---')
    # model.pprint()
    # solver = pyo.SolverFactory('ipopt')
    # solver.solve(model)

    return model