Coverage for backend/ahuora-builder/src/ahuora_builder/methods/adapter.py: 91%
159 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-07-22 05:22 +0000
1import traceback
2from typing import TYPE_CHECKING, Any
4from pyomo.environ import Block, value as pyo_value, Component, Var, ScalarVar, Reference, Expression
5from pyomo.core.base.expression import ScalarExpression
6from pyomo.core.base.constraint import ScalarConstraint
7from pyomo.core.base.var import VarData
8from pyomo.core.base.units_container import units
9from idaes.core import FlowsheetBlock
10from ahuora_builder.properties_manager import PropertiesManager
11from ahuora_builder_types.unit_model_schema import SolvedPropertyValueSchema
12from .units_handler import attach_unit, get_attached_unit, get_attached_unit_str, ValueWithUnits
13from ahuora_builder_types import PropertiesSchema, PropertySchema
14from typing import Dict
15from pyomo.core.base.indexed_component import UnindexedComponent_set, IndexedComponent
16from pyomo.core.base.indexed_component_slice import (
17 IndexedComponent_slice,
18)
19import numpy as np
20from ahuora_builder_types.id_types import PropertyValueId
22if TYPE_CHECKING:
23 from ahuora_builder.methods.BlockContext import BlockContext
25def get_component(blk: Block, key: str):
26 """
27 Get a component from a block, given a key. Doesn't handle indexes.
28 """
29 try:
30 # allow key to be split by "." to access nested properties eg. "hot_side.deltaP"
31 key_split = key.split(".")
32 b = blk
33 for k in key_split:
34 b = getattr(b, k)
35 return b
36 except AttributeError:
37 raise ValueError(
38 f"Property {key} not found in block `{blk}`. "
39 f"Available properties are {[x for x in blk.component_map().keys()]}"
40 )
42def add_to_property_map(vars: IndexedComponent, id: PropertyValueId, fs):
43 try:
44 name = next(vars.values()).name # maybe we should deprectate the name?
45 fs.properties_map.add(id, vars, name)
46 except StopIteration:
47 pass # The only reason this would happen is if there are no values in the indexed component, e.g in milk solids, there is no gas phase. In these case, we just skip adding to the property map.
48 #return c
50def add_corresponding_constraint(fs: FlowsheetBlock,c, id: PropertyValueId):
51 fs.properties_map.add_constraint(id, c)
54def soft_cast_float(value: Any) -> float:
55 """
56 Softly cast a value to float, returning None if the value is None or cannot be cast.
57 This is needed because not everything is indexed by time at all.
58 """
59 try:
60 return float(value)
61 except (ValueError, TypeError):
62 return None
64def items_by_time(s: Dict[str, Any]) -> list[tuple[str, Any]]:
65 """
66 Converts a dictionary of items indexed by time (as strings) to a list of tuples, and fixes the ordering.
67 This is because if "11" is ordered before "2" in a dictionary of strings, but we want the results ordered by time.
68 """
69 return sorted(s.items(), key=lambda x: soft_cast_float(x[0])) # sort by time index
73def get_index_set_shape(component: IndexedComponent) -> tuple[int,...]:
74 index_set = component.index_set()
75 if index_set.dimen == 0: 75 ↛ 76line 75 didn't jump to line 76 because the condition on line 75 was never true
76 return () # I don't think we ever actually have this case (we early return if there are no items in the index set with Warning: No variables found)
77 elif index_set.dimen == 1:
78 return (len(index_set),)
79 elif index_set.dimen > 1: 79 ↛ exitline 79 didn't return from function 'get_index_set_shape' because the condition on line 79 was always true
80 shape = [len(s) for s in index_set.subsets()]
81 return tuple(shape)
84def serialize_properties_map(fs: FlowsheetBlock) -> list[SolvedPropertyValueSchema]:
85 properties : list[SolvedPropertyValueSchema] = []
86 # TODO: collate properties by timestep.
87 properties_map: PropertiesManager = fs.properties_map
88 for (uid, s) in properties_map.items():
89 if uid == -1:
90 # skip
91 continue
92 s.component
94 shape = get_index_set_shape(s.component)
96 values = [pyo_value(c) for c in s.component.values()]
98 if shape == (): 98 ↛ 100line 98 didn't jump to line 100 because the condition on line 98 was never true
99 # not indexed at all, just return the value (not as an array)
100 items = values[0]
101 else:
102 # Convert to an ndarray
103 items = np.array(values).tolist()
105 property_dict = SolvedPropertyValueSchema(
106 id=uid,
107 name=s.name, # for debugging
108 value=items,
109 unit=get_attached_unit_str(s.component),
110 )
111 if s.unknown_units:
112 property_dict.unknown_units = s.unknown_units
114 properties.append(property_dict)
116 return properties
118def slice_is_indexed(blk: IndexedComponent | IndexedComponent_slice) -> bool:
119 """
120 Check if a block, variable, or expression is not indexed.
121 """
122 if isinstance(blk, IndexedComponent_slice):
123 # IndexedComponent_slice doesn't have is_indexed() method.
124 # Instead, it runs is_indexed() on each of the underlying components.
125 # We can assume that they are all the same, and just check the first one.
126 return blk.is_indexed()[0]
127 return blk.is_indexed()
129def slice_index_dimen(blk: IndexedComponent | IndexedComponent_slice) -> int:
130 """
131 Get the dimension of the index set of a block, variable, or expression.
132 """
133 if isinstance(blk, IndexedComponent_slice):
134 # IndexedComponent_slice doesn't have index_set() method.
135 # Instead, it runs index_set() on each of the underlying components.
136 # We can assume that they are all the same, and just check the first one.
137 return blk.index_set()[0].dimen
138 return blk.index_set().dimen
141def get_sliced_version(block: Block | Var | Expression | IndexedComponent_slice) -> Block | Var | Expression | IndexedComponent_slice:
142 """
143 Get a sliced version of a block, variable, or expression.
144 A very loose way of thinking of a sliced version is it "references all the indexes at the same time".
146 in a scalar block, you can do scalar_block.property.
147 in an indexed block, you can't do indexed_block.property, because you need to define the index you're looking at.
148 However, if you want to define the indexes later, such as in a reference, you can do indexed_block[:].property.
149 I.e you're creating a slice to all the indexes, and then accessing the property on that slice.
151 You can't really use a sliced version directly, but you can use it to create a Reference with reference(sliced_version).
152 It'll put all the indexes back together.
154 The main advantage is that you can use get_sliced_version on a indexed subattribute of a slice, and it will collate the indexes together.
155 """
156 if not slice_is_indexed(block):
157 return block
159 block_dimen = slice_index_dimen(block)
160 # we want to get a slice to all items in a block, like block[:,:]
161 # To do this programmatically:
162 # the ":" is represented by slice(None)
163 # so for a 2D block, we want (slice(None), slice(None))
164 # we can use a tuple comprehension to create this for any dimen
165 block_slice = block[tuple(slice(None) for _ in range(block_dimen))]
166 return block_slice
169def collate_indexes(block: Block, property_key: str) -> IndexedComponent:
170 """
171 Returns a reference to the property with all the indexes together.
172 This is because some properties have a different way of doing indexes:
173 e.g
174 - properties_out[t].temperature -> Reference[t]
175 - block.heat_duty[t] -> Reference[t]
176 - properties_out[t].mole_frac_comp[compound] -> Reference[t, compound]
177 - block.area -> Reference[None]
178 """
179 block_slice = get_sliced_version(block)
180 # Now we have the sliced version of the block, we can access the property using the property key.
181 # As we support nested properties, e.g property_key could be "hot_side.deltaP",
182 # we can use the get_component function rather than just getattr.
183 block_property = get_component(block_slice, property_key)
184 if block_property is None: 184 ↛ 185line 184 didn't jump to line 185 because the condition on line 184 was never true
185 raise ValueError(f"Property {property_key} not found in block {block}.")
186 # Now we have the property, we need to check if it has any indexes, and get a slice to all those indexes.
187 property_slice = get_sliced_version(block_property)
188 #property_slice = block_property
189 # Calling is_indexed() on a property_slice doesn't really make sense, it returns an array of bools. slices are always indexed though so
190 # that array gets cast to True, so its' kinda okay.
191 if property_slice.is_indexed(): # and not isinstance(property_slice, IndexedComponent_slice):
192 # if the property is indexed, we need to return a reference to the slice.
193 # This will collate all the indexes together.
194 # Note that we need to special case DerivativeVar, as Reference doesn't support DerivativeVar directly. ( pyomo dae will try differentiate it again.)
195 return Reference(property_slice, ctype=IndexedComponent) ## if isinstance(block_property,DerivativeVar) else NOTSET
196 else: # this is not indexed at all lol, so we can just return the property itself. Creating a reference will add a index[None] which is unnecessary.
197 return property_slice
201def fix_var(blk: Block, var : ScalarVar | ScalarExpression, value : ValueWithUnits) -> ScalarVar | ScalarConstraint:#
202 # returns: the expression that a constraint was added for, or the fixed var
203 if hasattr(blk, "constrain_component"):
204 constraint = blk.constrain_component(var, value)
205 return constraint
206 elif isinstance(var, ScalarExpression): 206 ↛ 208line 206 didn't jump to line 208 because the condition on line 206 was never true
207 # This is only used in tests right now I think. Our only expressions are currently in the property packages, and they use the constrain_component method.
208 constraint = ScalarConstraint(expr=var == value)
209 blk.add_component(f"{var.name}_constraint", constraint)
210 return constraint
211 else:
212 #print(var)
213 var.fix(value)
214 return var
216def fix_slice(var_slice: IndexedComponent | IndexedComponent_slice, values: list[ValueWithUnits]) -> list[ScalarVar | ScalarConstraint]:
217 # Fix a slice of a variable to the given values.
218 # the var_slice should be a Reference or other indexed component, or a scalar variable
219 # values should be a list of values to fix the variable to.
220 results: list[ScalarVar | ScalarConstraint] = []
221 for var, value in zip(var_slice.values(), values):
222 blk = var.parent_block()
223 constraint = fix_var(blk, var, value)
224 results.append(constraint)
225 return results
228def deactivate_components(components: list[ScalarVar | ScalarConstraint ]):
229 # Deactivate a "PropertyValue" (which may have multiple subcomponents if it's indexed)
230 for c in components:
231 deactivate_component(c)
233def deactivate_component(c: ScalarVar | ScalarConstraint):
234 # deactivate "guess" variables: fixed for initialisation
235 # and unfixed for a control constraint
236 if isinstance(c, ScalarConstraint): 236 ↛ 237line 236 didn't jump to line 237 because the condition on line 236 was never true
237 c.deactivate()
238 else:
239 c.unfix()
241def deactivate_fixed_guesses(guess_vars: list[list[ScalarVar | ScalarConstraint]]):
243 for c in guess_vars:
244 deactivate_components(c)
247def load_initial_guess(component: Component, value: float):
248 """Load an initial value into scalar or indexed variable data."""
249 if isinstance(component, VarData):
250 component.set_value(value)
253def load_initial_guesses(components: IndexedComponent, values: list[float]):
254 for c, v in zip(components.values(), values):
255 load_initial_guess(c, v)
258def fix_block(
259 block: Block,
260 properties_schema: PropertiesSchema,
261 fs: FlowsheetBlock,
262 block_ctx: "BlockContext",
263) -> None:
264 """
265 Fix the properties of a block based on the properties schema.
267 Args:
268 - block: The block to fix the properties of.
269 - properties_schema: The schema of the properties to fix.
270 - fs: Used to store the properties in the properties map.
271 """
272 property_key: str
273 property_info: PropertySchema
274 for property_key, property_info in properties_schema.items():
275 # Key is e.g "enth_mol"
277 # TODO: Handle transformers
278 # indexed_data = extract_indexes(
279 # property_info.data,
280 # property_info.unit,
281 # transformers,
282 # )
283 property_reference = collate_indexes(block, property_key)
285 for property_value in property_info.data:
286 discrete_indexes = property_value.discrete_indexes or []
288 pv_id = property_value.id
289 num_discrete_indexes = len(discrete_indexes)
290 num_property_indexes = property_reference.index_set().dimen if property_reference.is_indexed() else 0
291 num_continuous_indexes = num_property_indexes - num_discrete_indexes # This is the dimension of the property_value.value ndarray.
293 # We have the convention that all the continuous indexes come first in property_reference, and then the discrete indexes.
294 # This is what idaes normally does.
297 # We need to get a slice to the current set of discrete indexes.
298 if len(discrete_indexes) == 0:
299 property_slice = property_reference # no indexes to worry about.
300 else:
301 property_slice = property_reference[tuple(
302 list(slice(None) for _ in range(num_continuous_indexes)) + discrete_indexes
303 )]
304 # Now property_slice is only indexed by the continuous indexes.
306 # get a reference to the variable/expression. This will also add an index [None] if it is not indexed at all, i.e no continuous indexes.
307 variable_references = Reference(property_slice, ctype=IndexedComponent) #Ctype=IndexedComponent avoids problems with DerivativeVar
309 add_to_property_map(variable_references, pv_id, fs)
311 # Because both pyomo and numpy flatten arrays with the last index changing fastest, we can just flatten both the index set and the values, and then iterate through them together.
312 variable_indexes = list(variable_references.index_set())
314 if (len(variable_indexes) == 0):
315 print(f"Warning: No variables found for {property_key} with indexes {discrete_indexes}. This may be expected in the milk property package, which doesn't have a gas phase.")
316 continue
318 variable_transformed = variable_references
321 if property_value.value is not None:
322 variable_values = np.array([property_value.value]).flatten() # we put the value in an array to handle the case where it is a scalar.
323 variable_values_with_units = [attach_unit(v, property_info.unit) for v in variable_values]
324 variable_values_converted = [
325 units.convert(v, get_attached_unit(var)) for v, var in zip(variable_values_with_units, variable_references.values())
326 ]
328 #value = units.convert(property_value.value, get_attached_unit(var))
329 if property_value.constraint is not None:
330 # add the constraint to the list of constraints to be added
331 # at the flowsheet level. The var value should also
332 # be a guess to maintain the degrees of freedom.
333 expr = property_value.constraint
334 fs.constraint_exprs.append((variable_transformed, expr, pv_id))
335 if property_value.controlled:
336 # this is a set point. To maintain degrees of freedom,
337 # use the manipulated variable as a guess during initialisation.
338 # this is a weird case because we might have guesses for both the
339 # manipulated variable and this variable, and we want to end up
340 # using the formula/constraint. If we were to try use this guess
341 # by eliminating the manipulated variable guess, we would run into
342 # degrees of freedom issues at the flowsheet level if elimination
343 # failed (since it then adds another flowsheet level constraint).
344 load_initial_guesses(variable_references, variable_values_converted)
345 else:
346 # not a set point, use this guess during initialization.
347 component = fix_slice(variable_references, variable_values_converted)
348 fs.guess_vars.append(component)
349 elif property_value.controlled is not None:
350 # this value is a controlling variable, so don't fix it now
351 # it should be fixed after initialisation
352 block_ctx.add_controlled_var(variable_references, variable_values_converted, pv_id, property_value.controlled)
353 elif property_value.guess:
354 block_ctx.add_guess_var( variable_references, variable_values_converted, pv_id)
355 else:
356 components = fix_slice(variable_references, variable_values_converted)
357 # add_corresponding_constraint used to take the constraint returned by constrain_component if you are constraining an expression. TODO: Do we need to add this back in?
358 # from c = fix_var()
359 add_corresponding_constraint(fs, components, pv_id)