modify sb constrain methods to work with refactored `fix_state_bloc…

…k` method in the platform

property_packages/helmholtz/helmholtz_extended.py

@@ -78,22 +78,33 @@ def build(self, *args):
         add_extra_expressions(self)
         # Add a block for constraints, so we can disable or enable them in bulk
         self.constraints = Block()
-
+
+
     def constrain(self, name: str, value: float) -> Constraint | Var | None:
-        # Value must be a float. TODO: Handle unit conversion.
-        var = getattr(self,name)
-        if type(var) == ScalarExpression:
-            c = Constraint(expr=var == value)
-            self.constraints.add_component(name, c)
+        """constrain a component by name to a value"""
+        # TODO: handle unit conversion
+        var = getattr(self, name)
+        return self.constrain_component(var, value)
+
+
+    def constrain_component(self, component: Var | Expression, value: float) -> Constraint | Var | None:
+        """
+        Constrain a component to a value
+        """
+        if type(component) == ScalarExpression:
+            c = Constraint(expr=component == value)
+            self.constraints.add_component(component.local_name, c)
             return c
-        elif type(var) in (ScalarVar, _GeneralVarData, VarData):
-            var.fix(value)
-            return var
-        elif type(var) in ( _GeneralExpressionData, ExpressionData) :
+        elif type(component) in (ScalarVar, _GeneralVarData, VarData, IndexedVar):
+            component.fix(value)
+            return component
+        elif type(component) in (_GeneralExpressionData, ExpressionData):
             # allowed, but we don't need to fix it (eg. mole_frac_comp in helmholtz)
             return None
         else:
-            raise Exception(f"Variable {self} {name} is not a Var or Expression: {type(var)}")
+            raise Exception(
+                f"Component {component} is not a Var or Expression: {type(component)}"
+            )
 
 @declare_process_block_class("HelmholtzExtendedParameterBlock")
 class HelmholtzExtendedParameterBlockData(HelmholtzParameterBlockData):

property_packages/modular/modular_extended.py

@@ -1,6 +1,6 @@
 from pyomo.environ import Block, Constraint
 from pyomo.core.base.expression import ScalarExpression
-from pyomo.core.base.var import ScalarVar, _GeneralVarData, VarData
+from pyomo.core.base.var import ScalarVar, _GeneralVarData, VarData, IndexedVar
 from idaes.core import declare_process_block_class
 from idaes.core.util.exceptions import ConfigurationError
 from idaes.models.properties.modular_properties.base.generic_property import (
@@ -133,6 +133,7 @@ def fix_state_vars(blk, state_args=None):
             "entr_mol",
             "entr_mass",
             "smooth_temperature",
+            "total_energy_flow",
         ]
         for n, v in b.define_state_vars().items():
             fix_var = True
@@ -634,20 +635,28 @@ def build(self, *args):
         # Add a block for constraints, so we can disable or enable them in bulk
         self.constraints = Block()
 
+
     def constrain(self, name: str, value: float) -> Constraint | Var | None:
-        # Value must be a float. TODO: Handle unit conversion.
+        """constrain a component by name to a value"""
         var = getattr(self, name)
-        if type(var) == ScalarExpression:
-            c = Constraint(expr=var == value)
+        return self.constrain_component(var, value)
+
+
+    def constrain_component(self, component: Var | Expression, value: float) -> Constraint | Var | None:
+        """
+        Constrain a component to a value
+        """
+        if type(component) == ScalarExpression:
+            c = Constraint(expr=component == value)
             c.defining_state_var = True
-            self.constraints.add_component(name, c)
+            self.constraints.add_component(component.local_name, c)
             return c
-        elif type(var) in (ScalarVar, _GeneralVarData, VarData):
-            var.fix(value)
-            return var
+        elif type(component) in (ScalarVar, _GeneralVarData, VarData, IndexedVar):
+            component.fix(value)
+            return component
         else:
             raise Exception(
-                f"Variable {self} {name} is not a Var or Expression: {type(var)}"
+                f"Component {component} is not a Var or Expression: {type(component)}"
             )