Coverage for backend/idaes_service/solver/custom/PySMOModel.py: 86%

1# Methods and helper functions to train and use a surrogate valve. 

2from pyomo.core.base.expression import ScalarExpression 

3from idaes.core import MaterialBalanceType, ControlVolume0DBlock, declare_process_block_class, EnergyBalanceType, MomentumBalanceType, MaterialBalanceType, useDefault, UnitModelBlockData 

4from pyomo.common.config import ConfigBlock, ConfigValue, In 

5from idaes.core.surrogate.surrogate_block import SurrogateBlock 

6from idaes.core.util.config import is_physical_parameter_block 

7from idaes_service.solver.methods.expression_parsing import get_property_from_id 

8from idaes.core.surrogate.pysmo_surrogate import PysmoSurrogate 

9from io import StringIO 

10import json 

11from pyomo.environ import Var, Constraint 

12from pyomo.environ import ( 

13 Constraint, 

14 Set, 

15 Var, 

16 Suffix, 

17 units as pyunits, 

18) 

19 

20def make_control_volume(unit, name, config): 

21 if config.dynamic is not False: 21 ↛ 22line 21 didn't jump to line 22 because the condition on line 21 was never true

22 raise ValueError('SurrogateValve does not support dynamics') 

23 if config.has_holdup is not False: 23 ↛ 24line 23 didn't jump to line 24 because the condition on line 23 was never true

24 raise ValueError('SurrogateValve does not support holdup') 

25 

26 control_volume = ControlVolume0DBlock( 

27 dynamic=config.dynamic, 

28 has_holdup=config.has_holdup, 

29 property_package=config.property_package, 

30 property_package_args=config.property_package_args) 

31 

32 # Add the control volume block to the unit 

33 setattr(unit, name, control_volume) 

34 

35 control_volume.add_state_blocks(has_phase_equilibrium=config.has_phase_equilibrium) 

36 # control_volume.add_material_balances(balance_type=config.material_balance_type, 

37 # has_phase_equilibrium=config.has_phase_equilibrium) 

38 # control_volume.add_total_enthalpy_balances(has_heat_of_reaction=False,  

39 # has_heat_transfer=False,  

40 # has_work_transfer=False) 

41@declare_process_block_class("PySMOModel") 

42class PySMOModelData(UnitModelBlockData): 

43 CONFIG = UnitModelBlockData.CONFIG() 

44 # Declare all the standard config arguments for the control_volume 

45 CONFIG.declare("material_balance_type", ConfigValue(default=MaterialBalanceType.componentPhase, domain=In(MaterialBalanceType))) 

46 CONFIG.declare("energy_balance_type", ConfigValue(default=EnergyBalanceType.enthalpyTotal, domain=In([EnergyBalanceType.enthalpyTotal]))) 

47 CONFIG.declare("momentum_balance_type", ConfigValue(default=MomentumBalanceType.none, domain=In([MomentumBalanceType.none]))) 

48 CONFIG.declare("has_phase_equilibrium", ConfigValue(default=False, domain=In([False]))) 

49 CONFIG.declare("has_pressure_change", ConfigValue(default=False, domain=In([False]))) 

50 CONFIG.declare("property_package", ConfigValue(default=useDefault, domain=is_physical_parameter_block)) 

51 CONFIG.declare("property_package_args", ConfigBlock(implicit=True)) 

52 # no other args need to be declared, we are just hardcoding the valve model. 

53 CONFIG.declare("model", ConfigValue()) 

54 CONFIG.declare("ids", ConfigValue()) 

55 CONFIG.declare("unitopNames", ConfigValue()) 

56 CONFIG.declare( 

57 "num_inlets", 

58 ConfigValue( 

59 default=False, 

60 domain=int, 

61 description="Number of inlets to add", 

62 doc="Number of inlets to add", 

63 ), 

64 ) 

65 CONFIG.declare( 

66 "num_outlets", 

67 ConfigValue( 

68 default=False, 

69 domain=int, 

70 description="Number of outlets to add", 

71 doc="Number of outlets to add", 

72 ), 

73 ) 

74 

75 def build(self): 

76 super(PySMOModelData, self).build() 

77 self.CONFIG.dynamic = False 

78 self.CONFIG.has_holdup = False 

79 # This function handles adding the control volume block to the unit, 

80 # and addiing the necessary material and energy balances. 

81 make_control_volume(self, "control_volume", self.CONFIG) 

82 

83 # self.add_inlet_port() 

84 # self.add_outlet_port() 

85 

86 # Defining parameters of state block class 

87 tmp_dict = dict(**self.config.property_package_args) 

88 tmp_dict["parameters"] = self.config.property_package 

89 tmp_dict["defined_state"] = True # inlet block is an inlet 

90 

91 # Add state blocks for inlet, outlet, and waste 

92 # These include the state variables and any other properties on demand 

93 num_inlets = self.config.num_inlets 

94 self.inlet_list = [ "inlet_" + str(i+1) for i in range(num_inlets) ] 

95 self.inlet_set = Set(initialize=self.inlet_list) 

96 self.inlet_blocks = [] 

97 

98 for name in self.inlet_list: 

99 # add properties_inlet_1, properties_inlet2 etc 

100 state_block = self.config.property_package.state_block_class( 

101 self.flowsheet().config.time, doc="inlet ml", **tmp_dict 

102 ) 

103 self.inlet_blocks.append(state_block) 

104 # Dynamic equivalent to self.properties_inlet_1 = stateblock 

105 setattr(self,"properties_" + name, state_block) 

106 # also add the port 

107 self.add_port(name=name,block=state_block) 

108 

109 

110 # Add outlet state blocks 

111 

112 num_outlets = self.config.num_outlets 

113 self.outlet_list = [ "outlet_" + str(i+1) for i in range(num_outlets) ] 

114 self.outlet_set = Set(initialize=self.outlet_list) 

115 self.outlet_blocks = [] 

116 

117 for name in self.outlet_list: 

118 tmp_dict["defined_state"] = False 

119 state_block = self.config.property_package.state_block_class( 

120 self.flowsheet().config.time, doc="outlet ml", **tmp_dict 

121 ) 

122 self.outlet_blocks.append(state_block) 

123 setattr(self,"properties_" + name, state_block) 

124 self.add_port(name=name,block=state_block) 

125 

126 

127 

128 # Add variables for custom properties  

129 names = self.config.unitopNames 

130 for name in names: 

131 self.add_component( 

132 name, Var(self.flowsheet().time, initialize=10) 

133 ) 

134 

135 def initialize(self, *args, **kwargs): 

136 model_data = self.config.model 

137 json_str = json.dumps(model_data) 

138 f = StringIO(json_str) 

139 

140 model = PysmoSurrogate.load(f) 

141 ids = self.config.ids 

142 fs = self.flowsheet() 

143 # TODO: Make surrogate models work with dynamics. This involves making a surrogate model for each time step. 

144 # Not sure if idaes has a framework for doing this or not. 

145 # For now, we are just assuming steady state and time_index=0. 

146 inputs = [get_property_from_id(fs, i,0) for i in ids["input"]] 

147 outputs = [get_property_from_id(fs, i,0) for i in ids["output"]] 

148 

149 self.check_is_expression(inputs) 

150 self.check_is_expression(outputs) 

151 

152 self.surrogate = SurrogateBlock(concrete=True) 

153 self.surrogate.build_model(model,input_vars=inputs, output_vars=outputs) 

154 

155 def check_is_expression(self, vars): 

156 for index, var in enumerate(vars): 

157 if isinstance(var, ScalarExpression): 

158 name = f"{var.name}_{index}" 

159 new_var = Var(self.flowsheet().time, initialize=1) 

160 self.add_component(name, new_var) 

161 def constraint_rule(model, t): 

162 if var.is_indexed(): 162 ↛ 163line 162 didn't jump to line 163 because the condition on line 162 was never true

163 return new_var[t] == var[t] 

164 else: 

165 return new_var[t] == var 

166 

167 self.add_component(f"{name}_constraint", Constraint(self.flowsheet().time, rule=constraint_rule)) 

168 vars[index] = new_var 

169 

170 def _get_stream_table_contents(self, time_point=0): 

171 

172 io_dict = {} 

173 for inlet_name in self.inlet_list: 

174 io_dict[inlet_name] = getattr(self, inlet_name) # get a reference to the port 

175 

176 out_dict = {} 

177 for outlet_name in self.outlet_list: 

178 out_dict[outlet_name] = getattr(self, outlet_name) # get a reference to the port