Coverage for backend/ahuora-builder/src/ahuora_builder/custom/PySMOModel.py: 87%
104 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-05-13 02:47 +0000
1# Methods and helper functions to train and use a surrogate valve.
2from copy import deepcopy
3from pyomo.core.base.expression import ScalarExpression
4from idaes.core import MaterialBalanceType, ControlVolume0DBlock, declare_process_block_class, EnergyBalanceType, MomentumBalanceType, MaterialBalanceType, useDefault, UnitModelBlockData
5from pyomo.common.config import ConfigBlock, ConfigValue, In
6from idaes.core.surrogate.surrogate_block import SurrogateBlock
7from idaes.core.util.config import is_physical_parameter_block
8from idaes.core.util.exceptions import InitializationError
9from idaes.core.util.model_statistics import degrees_of_freedom
10from idaes.core.solvers import get_solver
11import idaes.logger as idaeslog
12from ahuora_builder.methods.expression_parsing import get_property_from_id
13from ahuora_builder.state_args import extract_state_args
14from idaes.core.surrogate.pysmo_surrogate import PysmoSurrogate
15from io import StringIO
16import json
17from pyomo.environ import Var, Constraint, check_optimal_termination
18from pyomo.environ import (
19 Constraint,
20 Set,
21 Var,
22 Suffix,
23 units as pyunits,
24)
26def make_control_volume(unit, name, config):
27 if config.dynamic is not False: 27 ↛ 28line 27 didn't jump to line 28 because the condition on line 27 was never true
28 raise ValueError('SurrogateValve does not support dynamics')
29 if config.has_holdup is not False: 29 ↛ 30line 29 didn't jump to line 30 because the condition on line 29 was never true
30 raise ValueError('SurrogateValve does not support holdup')
32 control_volume = ControlVolume0DBlock(
33 dynamic=config.dynamic,
34 has_holdup=config.has_holdup,
35 property_package=config.property_package,
36 property_package_args=config.property_package_args)
38 # Add the control volume block to the unit
39 setattr(unit, name, control_volume)
41 control_volume.add_state_blocks(has_phase_equilibrium=config.has_phase_equilibrium)
42 # control_volume.add_material_balances(balance_type=config.material_balance_type,
43 # has_phase_equilibrium=config.has_phase_equilibrium)
44 # control_volume.add_total_enthalpy_balances(has_heat_of_reaction=False,
45 # has_heat_transfer=False,
46 # has_work_transfer=False)
47@declare_process_block_class("PySMOModel")
48class PySMOModelData(UnitModelBlockData):
49 CONFIG = UnitModelBlockData.CONFIG()
50 # Declare all the standard config arguments for the control_volume
51 CONFIG.declare("material_balance_type", ConfigValue(default=MaterialBalanceType.componentPhase, domain=In(MaterialBalanceType)))
52 CONFIG.declare("energy_balance_type", ConfigValue(default=EnergyBalanceType.enthalpyTotal, domain=In([EnergyBalanceType.enthalpyTotal])))
53 CONFIG.declare("momentum_balance_type", ConfigValue(default=MomentumBalanceType.none, domain=In([MomentumBalanceType.none])))
54 CONFIG.declare("has_phase_equilibrium", ConfigValue(default=False, domain=In([False])))
55 CONFIG.declare("has_pressure_change", ConfigValue(default=False, domain=In([False])))
56 CONFIG.declare("property_package", ConfigValue(default=useDefault, domain=is_physical_parameter_block))
57 CONFIG.declare("property_package_args", ConfigBlock(implicit=True))
58 # no other args need to be declared, we are just hardcoding the valve model.
59 CONFIG.declare("model", ConfigValue())
60 CONFIG.declare("ids", ConfigValue())
61 CONFIG.declare("unitopNames", ConfigValue())
62 CONFIG.declare(
63 "num_inlets",
64 ConfigValue(
65 default=False,
66 domain=int,
67 description="Number of inlets to add",
68 doc="Number of inlets to add",
69 ),
70 )
71 CONFIG.declare(
72 "num_outlets",
73 ConfigValue(
74 default=False,
75 domain=int,
76 description="Number of outlets to add",
77 doc="Number of outlets to add",
78 ),
79 )
81 def build(self):
82 super(PySMOModelData, self).build()
83 self.CONFIG.dynamic = False
84 self.CONFIG.has_holdup = False
85 # This function handles adding the control volume block to the unit,
86 # and addiing the necessary material and energy balances.
87 make_control_volume(self, "control_volume", self.CONFIG)
89 # self.add_inlet_port()
90 # self.add_outlet_port()
92 # Defining parameters of state block class
93 tmp_dict = dict(**self.config.property_package_args)
94 tmp_dict["parameters"] = self.config.property_package
95 tmp_dict["defined_state"] = True # inlet block is an inlet
97 # Add state blocks for inlet, outlet, and waste
98 # These include the state variables and any other properties on demand
99 num_inlets = self.config.num_inlets
100 self.inlet_list = [ "inlet_" + str(i+1) for i in range(num_inlets) ]
101 self.inlet_set = Set(initialize=self.inlet_list)
102 self.inlet_blocks = []
104 for name in self.inlet_list:
105 # add properties_inlet_1, properties_inlet2 etc
106 state_block = self.config.property_package.state_block_class(
107 self.flowsheet().config.time, doc="inlet ml", **tmp_dict
108 )
109 self.inlet_blocks.append(state_block)
110 # Dynamic equivalent to self.properties_inlet_1 = stateblock
111 setattr(self,"properties_" + name, state_block)
112 # also add the port
113 self.add_port(name=name,block=state_block)
116 # Add outlet state blocks
118 num_outlets = self.config.num_outlets
119 self.outlet_list = [ "outlet_" + str(i+1) for i in range(num_outlets) ]
120 self.outlet_set = Set(initialize=self.outlet_list)
121 self.outlet_blocks = []
123 for name in self.outlet_list:
124 tmp_dict["defined_state"] = False
125 state_block = self.config.property_package.state_block_class(
126 self.flowsheet().config.time, doc="outlet ml", **tmp_dict
127 )
128 self.outlet_blocks.append(state_block)
129 setattr(self,"properties_" + name, state_block)
130 self.add_port(name=name,block=state_block)
134 # Add variables for custom properties
135 names = self.config.unitopNames
136 for name in names:
137 self.add_component(
138 name, Var(self.flowsheet().time, initialize=10)
139 )
141 def post_load(self):
142 # Build the model
143 model_data = self.config.model
144 json_str = json.dumps(model_data)
145 f = StringIO(json_str)
147 model = PysmoSurrogate.load(f)
148 ids = self.config.ids
149 fs = self.flowsheet()
150 # TODO: Make surrogate models work with dynamics. This involves making a surrogate model for each time step.
151 # Not sure if idaes has a framework for doing this or not.
152 # For now, we are just assuming steady state and time_index=0.
153 inputs = [get_property_from_id(fs, i,0) for i in ids["input"]]
154 outputs = [get_property_from_id(fs, i,0) for i in ids["output"]]
156 self.check_is_expression(inputs)
157 self.check_is_expression(outputs)
159 self.surrogate = SurrogateBlock(concrete=True)
160 self.surrogate.build_model(model,input_vars=inputs, output_vars=outputs)
162 def initialize(self, *args, **kwargs):
163 return None
166 def check_is_expression(self, vars):
167 for index, var in enumerate(vars):
168 if isinstance(var, ScalarExpression):
169 name = f"{var.name}_{index}"
170 new_var = Var(self.flowsheet().time, initialize=1)
171 self.add_component(name, new_var)
172 def constraint_rule(model, t):
173 if var.is_indexed(): 173 ↛ 174line 173 didn't jump to line 174 because the condition on line 173 was never true
174 return new_var[t] == var[t]
175 else:
176 return new_var[t] == var
178 self.add_component(f"{name}_constraint", Constraint(self.flowsheet().time, rule=constraint_rule))
179 vars[index] = new_var
181 def _get_stream_table_contents(self, time_point=0):
183 io_dict = {}
184 for inlet_name in self.inlet_list:
185 io_dict[inlet_name] = getattr(self, inlet_name) # get a reference to the port
187 out_dict = {}
188 for outlet_name in self.outlet_list:
189 out_dict[outlet_name] = getattr(self, outlet_name) # get a reference to the port