Coverage for backend/ahuora-compounds/ahuora_property_packages/base/flexible_state_block.py: 21%

1 

2from pyomo.environ import Constraint, check_optimal_termination 

3from idaes.core.util.initialization import solve_indexed_blocks 

4from idaes.core.util.model_statistics import degrees_of_freedom, number_unfixed_variables 

5from idaes.core.util.exceptions import InitializationError 

6import idaes.logger as idaeslog 

7from pyomo.core.base.expression import ScalarExpression, IndexedExpression, Expression, ExpressionData 

8from pyomo.environ import Var, ScalarVar 

9from pyomo.core.base.var import IndexedVar 

10import pyomo.environ as pyo 

11from idaes.core import ( 

12 StateBlock, 

13) 

14from pyomo.contrib.incidence_analysis import IncidenceGraphInterface 

15from pyomo.environ import Constraint 

16from pyomo.util.subsystems import create_subsystem_block 

17from pyomo.environ import SolverFactory 

18 

19class ExpressionConversionError(Exception): 

20 pass 

21 

22 

23def solve_square_subsets(block): 

24 """ 

25 Whatever part of the block is square, solve it. 

26 This gets everything that we can calculate exactly to the true values. 

27 """ 

28 igraph = IncidenceGraphInterface(block) 

29 var_dm, con_dm = igraph.dulmage_mendelsohn() 

30 # get the square portion of the block 

31 square_vars = var_dm.square 

32 square_cons = con_dm.square 

33 

34 solver = SolverFactory("ipopt") 

35 

36 sub = create_subsystem_block( 

37 constraints=square_cons, 

38 variables=square_vars, 

39 ) 

40 solver.solve(sub) 

41 

42class FlexibleStateBlockData(): 

43 

44 def build(blk, *args): 

45 blk.add_extra_expressions() 

46 # We initialise vars_to_deactivate here with __setattr__ instead of doing it 

47 # in the constructor of _SeawaterStateBlockConstraints as blk.vars_to_deactivate = [], 

48 # because on state blocks, missing attributes are not treated like normal Python objects. 

49 # Why the below works is because it bypasses the custom __setattr__ logic (which makes it a metadata) 

50 # on the block and writes directly to the object. 

51 # Also, initializing here guarantees every state block has its own list before constrain_component() runs. 

52 object.__setattr__(blk, "vars_to_deactivate", []) 

53 

54 def constrain(blk, name: str, value: float) -> Constraint | Var | None: 

55 """constrain a component by name to a value""" 

56 var = getattr(blk, name) 

57 return blk.constrain_component(var, value) 

58 

59 def constrain_component(blk, component: Var | Expression, value: float) -> Var | None: 

60 """ 

61 Constrain a component to a value 

62 """ 

63 try: 

64 variable = _convert_expression_to_var(component) 

65 except ExpressionConversionError: 

66 variable = component # already a Var, just fix it directly 

67 

68 variable.fix(value) 

69 

70 if isinstance(variable, IndexedVar): 

71 for i in variable.index_set(): 

72 # direct dictionary access avoids intercepted attribute resolution 

73 blk.__dict__["vars_to_deactivate"].append(variable[i]) 

74 else: 

75 blk.__dict__["vars_to_deactivate"].append(variable) 

76 

77 return variable 

78 

79 def add_extra_expressions(blk): 

80 """ 

81 IDAES state blocks don't support all the properties 

82 we need, so we add some extra expressions here. 

83 

84 This method can be overridden in a subclass to add 

85 additional expressions specific to the property package. 

86 """ 

87 if not hasattr(blk, "enth_mass"): 

88 blk.add_component("enth_mass", Expression(expr=(blk.flow_mol * blk.enth_mol) / blk.flow_mass)) 

89 if not hasattr(blk, "entr_mass"): 

90 blk.add_component("entr_mass", Expression(expr=(blk.flow_mol * blk.entr_mol) / blk.flow_mass)) 

91 if not hasattr(blk, "entr_mol"): 

92 blk.add_component("entr_mol", Expression(expr=(blk.flow_mol * blk.entr_mass) / blk.flow_mass)) 

93 if not hasattr(blk, "total_energy_flow"): 

94 blk.add_component("total_energy_flow", Expression(expr=blk.flow_mass * blk.enth_mass)) 

95 

96 

97def _convert_expression_to_var(expr: ScalarExpression | IndexedExpression): 

98 if isinstance(expr, ScalarExpression) or isinstance(expr, ExpressionData): 

99 var = Var(units=pyo.units.get_units(expr)) 

100 constraint = Constraint(expr= var == expr) 

101 elif isinstance(expr, IndexedExpression): 

102 var = Var(expr.index_set(), units=pyo.units.get_units(expr.units)) 

103 def rule(b, i): 

104 return var[i] == expr[i] 

105 constraint = Constraint(expr= rule) 

106 else: 

107 raise ExpressionConversionError(f"Expression {expr} is not a ScalarExpression or IndexedExpression: {type(expr)}") 

108 block = expr.parent_block() 

109 block.add_component(f"{expr.local_name}_var", var) 

110 block.add_component(f"{expr.local_name}_constraint", constraint) 

111 return var 

112 

113 

114def _solve_block(self, solve_log, init_log, opt, step_name): 

115 skip_solve = True # skip solve if only state variables are present 

116 for k in self.keys(): 

117 if number_unfixed_variables(self[k]) != 0: 

118 skip_solve = False 

119 

120 if not skip_solve: 

121 # Initialize properties 

122 with idaeslog.solver_log(solve_log, idaeslog.DEBUG) as slc: 

123 results = solve_indexed_blocks(opt, [self], tee=slc.tee) 

124 init_log.info_high( 

125 f"Property initialization {step_name}: {idaeslog.condition(results)}" 

126 ) 

127 

128 if (not skip_solve) and (not check_optimal_termination(results)): 

129 raise InitializationError( 

130 f"{self.name} {step_name} failed to initialize successfully. Please " 

131 f"check the output logs for more information." 

132 ) 

133 

134 

135class FlexibleStateBlock(StateBlock): 

136 """ 

137 This adds some additional methods for handling extra variables that have been fixed with constrain_component() 

138 during initialization. These wrapper methods can be called before and after the normal initialization routine 

139 to handle deactivating and reactivating constraints added with constrain_component(). 

140 """ 

141 

142 def _deactivate_additional_constraints(self): 

143 # Temporarily deactivate platform constraints added with 

144 # StateBlockConstraints.constrain()) so they don't interfere with 

145 # initialization. 

146 deactivated_vars: list[tuple[ScalarVar,float]] = [] 

147 

148 for k in self.keys(): 

149 blk = self[k] 

150 

151 for var in blk.__dict__["vars_to_deactivate"]: 

152 if var.is_fixed(): 

153 var.unfix() 

154 # Store the original value so we can reactivate and fix back to the original value later. 

155 deactivated_vars.append((var, var.value)) 

156 

157 self.deactivated_vars = deactivated_vars 

158 

159 def _reactivate_additional_constraints(self): 

160 for var, value in getattr(self, "deactivated_vars", []): 

161 var.fix(value) 

162 

163 def initialize(self, 

164 *args, 

165 state_vars_fixed=False, 

166 hold_state=False, 

167 outlvl=idaeslog.NOTSET, 

168 solver=None, 

169 optarg=None, 

170 **kwargs): 

171 """ 

172 Initialize with platform-added property constraints relaxed first, then 

173 solve again after reactivating them. 

174 

175 Native property package initialize/release methods own their state flag 

176 format. This keeps the flexible layer generic across modular, Helmholtz, 

177 and other property packages. 

178 """ 

179 original_dof = degrees_of_freedom(self) 

180 self.pre_initialize( 

181 *args, 

182 state_vars_fixed=state_vars_fixed, 

183 hold_state=hold_state, 

184 outlvl=outlvl, 

185 solver=solver, 

186 optarg=optarg, 

187 **kwargs, 

188 ) 

189 flags = super().initialize( 

190 *args, 

191 state_vars_fixed=state_vars_fixed, 

192 hold_state=True, 

193 outlvl=outlvl, 

194 solver=solver, 

195 optarg=optarg, 

196 **kwargs, 

197 ) 

198 

199 # if original dof is zero, we don't need to hold state as conditions are already fixed. 

200 if not state_vars_fixed and (not hold_state or original_dof == 0): 

201 self.release_state( 

202 flags, 

203 outlvl=outlvl, 

204 solver=solver, 

205 optarg=optarg, 

206 ) 

207 

208 return flags 

209 

210 def release_state(self, flags, outlvl=idaeslog.NOTSET,solver=None,optarg=None, **kwargs): 

211 super().release_state(flags, outlvl=outlvl, **kwargs) 

212 self.post_release_state() # Reactivate any additional constraints that were deactivated during initialization 

213 

214 logger = idaeslog.getInitLogger(self.name, outlvl, tag="properties") 

215 for blk in self.values(): 

216 try: 

217 # after reactivating the constraints, some of the initial guesses may have not 

218 # been quite right. E.g if flow_mass was fixed to a value that is different to flow_mol from initialisation. 

219 # however, other parts may not be fully specified, so we can't solve the whole block. 

220 # Instead, we solve the square subsets of the block to ensure all constraints that we can already calculate 

221 # are satisfied. 

222 solve_square_subsets(blk) 

223 except ValueError as e: 

224 # log the error but continue with initialization. 

225 logger.warning(f"Failed to solve square subsets for block {blk.name}: {e}") 

226 

227 

228 def pre_initialize(self, *args, **kwargs): 

229 """ 

230 Deactivate any additional constraints added with StateBlockConstraints.constrain() during initialization. 

231 Run the normal initialisation routine. 

232 Reactivate the constraints and solve again to ensure they are satisfied. 

233 If hold_state is True, restore the state to what it was 

234 """ 

235 self._deactivate_additional_constraints() 

236 

237 def post_release_state(self): 

238 # Reactivate platform constraints that were deferred during initialize 

239 self._reactivate_additional_constraints()