Coverage for backend/django/idaes_factory/idaes_factory.py: 95%
227 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-06-23 21:51 +0000
1import traceback
3from core.auxiliary.enums.unitOpData import SimulationObjectClass
4from core.auxiliary.models import MLModel
5from flowsheetInternals.unitops.config.objects.machine_learning_block_config import IDAdapter, MLPropetiesAdapter, unitopNamesAdapter, get_ml_properties, get_id_mappings, get_unitop_names
6from idaes_factory.adapters.property_package_adapter import PropertyPackageAdapter
7from idaes_factory.adapters.generic_adapters import NumInletsAdapter, NumOutletsAdapter
8from opentelemetry import trace
9from CoreRoot import settings
10from ahuora_builder_types.unit_model_schema import SolvedPropertyValueSchema, UnitModelSchema, ValueArgSchema
11import dotenv
12from typing import Any
13from django.db import transaction
14from ahuora_builder_types.scenario_schema import UnfixedVariableSchema, OptimizationSchema
15from core.auxiliary.models.Scenario import Scenario, OptimizationDegreesOfFreedom
16from core.auxiliary.property_state import (
17 validate_objective_property,
18 validate_optimization_dof_property_value,
19)
20from core.exceptions import DetailedException
21from flowsheetInternals.graphicData.models.groupingModel import Grouping
22from flowsheetInternals.unitops.models.SimulationObject import SimulationObject
23from ahuora_builder_types import FlowsheetSchema
24from core.auxiliary.models.PropertyInfo import PropertyInfo
25from core.auxiliary.models.PropertyValue import PropertyValue
26from core.auxiliary.models.Solution import Solution
27from core.auxiliary.enums.unitsLibrary import units_library
28from .adapters import arc_adapter
29from .adapters.convert_expression import convert_expression
30from .idaes_factory_context import IdaesFactoryContext, LiveSolveParams
31from .queryset_lookup import get_value_object
32from .unit_conversion import convert_value
33from idaes_factory.unit_conversion.unit_conversion import can_convert
34from core.auxiliary.models.Scenario import Scenario, SolverOptionEnum
35from idaes_factory.build_hooks import IdaesBuildHookContext, run_before_context_load_hooks
37dotenv.load_dotenv()
39# Todo: replace these with literal types from the Compounds/PP library
40Compound = str
41PropertyPackage = str
44class IdaesFactoryBuildException(DetailedException):
45 pass
48tracer = trace.get_tracer(settings.OPEN_TELEMETRY_TRACER_NAME)
51class IdaesFactory:
52 """
53 The IdaesFactory class is the core class for building
54 a flowsheet (JSON schema) that can be sent to the IDAES
55 solver, and for storing the results back in the database.
56 """
58 def __init__(
59 self,
60 group_id: int,
61 scenario: Scenario | None = None,
62 require_variables_fixed: bool = True,
63 solve_index: int | None = None,
65 ) -> None:
66 """Prepare a factory capable of serialising the requested flowsheet.
68 Args:
69 group_id: Identifier of the flowsheet to serialise.
70 scenario: Optional scenario providing solve configuration settings.
71 require_variables_fixed: Whether adapters should enforce fixed variables.
72 solve_index: Optional multi-steady-state index to bind to the context.
73 """
75 self.solve_index = solve_index
76 self.scenario = scenario
77 self._context_load_hooks_ran = False
79 if scenario is not None:
80 is_dynamic = scenario.enable_dynamics
81 step_size = scenario.simulation_length / \
82 float(scenario.num_time_steps)
83 enable_rating = scenario.enable_rating
85 if scenario.enable_optimization:
86 # If we are doing optimization, we solve from the root
87 group_id = scenario.flowsheet.rootGrouping.id
88 else:
89 is_dynamic = False
90 step_size = 1 # Just need a placeholder value.
91 enable_rating = False
93 time_steps = list([int(i) * step_size for i in range(0,
94 scenario.num_time_steps)]) if is_dynamic else [0]
96 self.flowsheet = FlowsheetSchema(
97 group_id=group_id,
98 dynamic=is_dynamic,
99 time_set=time_steps,
100 property_packages=[],
101 unit_models=[],
102 arcs=[],
103 expressions=[],
104 optimizations=[],
105 is_rating_mode=enable_rating,
106 disable_initialization=getattr(
107 scenario, "disable_initialization", False),
108 skip_initialization_for_units_with_initial_values=getattr(
109 scenario,
110 "skip_initialization_for_units_with_initial_values",
111 False,
112 ),
113 solver_option=getattr(scenario, "solver_option", "ipopt"),
114 )
116 self._run_before_context_load_hooks_once(group_id, solve_index)
118 # factory context
119 self.context = IdaesFactoryContext(
120 group_id,
121 require_variables_fixed=require_variables_fixed,
122 solve_index=solve_index,
123 time_steps=time_steps,
124 time_step_size=step_size,
125 scenario=scenario,
126 )
128 def _run_before_context_load_hooks_once(self, group_id: int, solve_index: int | None) -> None:
129 """Run pre-load hooks exactly once for this factory instance."""
131 if self._context_load_hooks_ran: 131 ↛ 132line 131 didn't jump to line 132 because the condition on line 131 was never true
132 return
133 if self.scenario is not None:
134 active_flowsheet = self.scenario.flowsheet
135 else:
136 active_flowsheet = (
137 Grouping.objects.select_related("flowsheet")
138 .get(id=group_id)
139 .flowsheet
140 )
141 run_before_context_load_hooks(
142 IdaesBuildHookContext(
143 flowsheet=active_flowsheet,
144 scenario=self.scenario,
145 group_id=group_id,
146 solve_index=solve_index,
147 )
148 )
149 self._context_load_hooks_ran = True
151 # Updates the context to use a different solve index.
152 # build() should be called after this to update the extracted flowsheet data.
153 def use_with_solve_index(self, solve_index: int) -> None:
154 """Rebind the factory to a different multi steady-state solve index.
156 Args:
157 solve_index: Index of the solve configuration within the scenario.
158 """
159 self.solve_index = solve_index
160 self.context.update_solve_index(self.solve_index)
162 @tracer.start_as_current_span("build_flowsheet")
163 def build(self):
164 """Populate the flowsheet schema with units, arcs, expressions, and metadata.
166 Raises:
167 IdaesFactoryBuildException: If any adapter fails during serialisation.
168 """
169 try:
170 self.setup_unit_models()
171 self.create_arcs()
172 self.add_property_packages()
173 self.add_expressions()
174 self.add_optimizations()
175 self.check_dependencies()
176 except Exception as e:
177 raise IdaesFactoryBuildException(e, "idaes_factory_build") from e
179 def clear_flowsheet(self) -> None:
180 """Reset the in-memory flowsheet while preserving configuration metadata."""
181 self.flowsheet = FlowsheetSchema(
182 group_id=self.flowsheet.group_id,
183 dynamic=self.flowsheet.dynamic,
184 time_set=self.flowsheet.time_set,
185 property_packages=[],
186 unit_models=[],
187 arcs=[],
188 expressions=[],
189 optimizations=[],
190 is_rating_mode=self.flowsheet.is_rating_mode,
191 disable_initialization=self.flowsheet.disable_initialization,
192 skip_initialization_for_units_with_initial_values=(
193 self.flowsheet.skip_initialization_for_units_with_initial_values
194 ),
195 solver_option=self.flowsheet.solver_option
196 )
198 def add_property_packages(self) -> None:
199 """Attach any property packages collected during context loading."""
200 self.flowsheet.property_packages = self.context.property_packages
202 def setup_unit_models(self):
203 """Serialise all unit operations."""
204 # add all unit models
205 exclude = {"stream", "recycle", "specificationBlock",
206 "energy_stream", "ac_stream", "humid_air_stream", "transformer_stream"}
207 for unit_model in self.context.exclude_object_type(exclude):
208 self.add_unit_model(unit_model)
210 def add_ml_model_properties(self, unit_model: SimulationObject, ml_model: MLModel) -> UnitModelSchema:
211 """Serialise an attached ML model without creating independent ports."""
212 return UnitModelSchema(
213 id=ml_model.pk*-1, # use negative ids for attached ML models to avoid conflicts with SimulationObject ids
214 type=SimulationObjectClass.MachineLearningBlock,
215 name=unit_model.componentName + str(ml_model.pk),
216 args={
217 "property_package": PropertyPackageAdapter().serialise(self.context, unit_model),
218 "model": ValueArgSchema(value=ml_model.surrogate_model),
219 "ids": get_id_mappings(ml_model),
220 "unitopNames": get_unitop_names(ml_model),
221 "num_inlets": NumInletsAdapter().serialise(self.context, unit_model),
222 "num_outlets": NumOutletsAdapter().serialise(self.context, unit_model),
223 },
224 properties=get_ml_properties(self.context, ml_model),
225 ports={},
226 initial_values={}
227 )
229 def add_unit_model(self, unit_model: SimulationObject) -> None:
230 """Serialise and append a unit model using its registered adapter.
232 Args:
233 unit_model: Simulation object to convert into IDAES schema.
235 Raises:
236 Exception: If the adapter fails to serialise the unit model.
237 """
238 try:
239 adapter = unit_model.schema.idaes_adapter
240 if adapter is None: 240 ↛ 241line 240 didn't jump to line 241 because the condition on line 240 was never true
241 raise ValueError(
242 f"No IDAES adapter registered for object type {unit_model.objectType}"
243 )
244 schema = adapter.serialise(self.context, unit_model)
245 self.flowsheet.unit_models.append(schema)
246 if unit_model.objectType != "machineLearningBlock" and unit_model.MLModels.exists():
248 for ml_model in unit_model.MLModels.all():
249 self.flowsheet.unit_models.append(self.add_ml_model_properties(unit_model, ml_model))
251 except Exception as e:
252 raise Exception(
253 f"Error adding unit model {unit_model.componentName} to the flowsheet: {e}"
254 )
256 def add_expressions(self) -> None:
257 """Collect custom property expressions and expose them on the flowsheet."""
258 # expressions are stored in the property set of a group
259 # eg. the global base flowsheet object
260 simulation_object: SimulationObject
261 for simulation_object in self.context.exclude_object_type({"machineLearningBlock"}):
262 # skip machine learning blocks, their properties are handled differently. We still need to support them in future.
264 properties = simulation_object.properties
265 prop: PropertyInfo
266 for prop in properties.ContainedProperties.all():
267 if prop.key in simulation_object.schema.properties:
268 # This is a default property, we have already processed it.
269 # We only want to capture custom properties
270 continue
271 if prop.formula_incomplete:
272 # Incomplete generated formulas stay visible in the UI but
273 # cannot be represented in the IDAES expression payload.
274 continue
275 property_value = get_value_object(prop)
276 if property_value is None: 276 ↛ 277line 276 didn't jump to line 277 because the condition on line 276 was never true
277 continue
278 self.context.add_property_value_dependency(property_value)
279 if property_value.formula in (None, ""):
280 # Custom properties can also be plain user-entered values.
281 # Only formula-backed custom properties are builder expressions.
282 continue
283 self._add_expression(prop, property_value, track_dependency=False)
284 self.add_managed_expressions()
286 def add_managed_expressions(self) -> None:
287 """Serialize complete managed formulas that live outside the active group."""
289 if not self.context.has_loaded_managed_properties():
290 return
292 existing_expression_ids = {expression["id"] for expression in self.flowsheet.expressions}
293 properties = (
294 PropertyInfo.objects.filter(
295 flowsheet=self._active_flowsheet(),
296 managed=True,
297 formula_incomplete=False,
298 values__formula__isnull=False,
299 )
300 .prefetch_related("values")
301 .distinct()
302 )
303 for prop in properties:
304 property_value = get_value_object(prop)
305 if property_value is None or property_value.formula in (None, ""): 305 ↛ 306line 305 didn't jump to line 306 because the condition on line 305 was never true
306 continue
307 if property_value.id in existing_expression_ids:
308 continue
309 self._add_expression(prop, property_value)
310 existing_expression_ids.add(property_value.id)
312 def _add_expression(
313 self,
314 prop: PropertyInfo,
315 property_value: PropertyValue,
316 *,
317 track_dependency: bool = True,
318 ) -> None:
319 if track_dependency:
320 self.context.add_property_value_dependency(property_value)
321 self.flowsheet.expressions.append(
322 {
323 "id": property_value.id,
324 "name": prop.displayName,
325 "expression": convert_expression(
326 property_value.formula,
327 self.context,
328 property_value,
329 ),
330 }
331 )
333 def _active_flowsheet(self):
334 """Return the flowsheet being serialized, with or without a scenario."""
336 if self.scenario is not None:
337 return self.scenario.flowsheet
338 return self.context.group.flowsheet
340 def add_optimizations(self) -> None:
341 """Serialise scenario-level optimisation settings onto the flowsheet."""
342 # This method was originally written to return multiple optimisations.
343 # this doesn't make sense, but idaes_service hasn't been updated to only expect one.
344 # so for now, it sets optimisations to an array with one item
345 optimization = self.context.scenario
346 if optimization is None or optimization.enable_optimization is False:
347 # no optimization to add
348 return
349 sense = "minimize" if optimization.minimize else "maximize"
350 if optimization.objective is None: 350 ↛ 351line 350 didn't jump to line 351 because the condition on line 350 was never true
351 raise ValueError(
352 "Please set an objective for the optimization to minimize or maximize.")
353 objective = optimization.objective
354 validate_objective_property(objective)
355 objective_value = get_value_object(objective)
356 if objective_value is None: 356 ↛ 357line 356 didn't jump to line 357 because the condition on line 356 was never true
357 raise ValueError("The selected objective does not have a property value.")
359 degrees_of_freedom = []
360 degree_of_freedom: OptimizationDegreesOfFreedom
361 for degree_of_freedom in optimization.degreesOfFreedom.all():
362 validate_optimization_dof_property_value(degree_of_freedom.propertyValue)
363 property_value_id = degree_of_freedom.propertyValue_id
365 dof_schema = UnfixedVariableSchema(
366 id=property_value_id,
367 lower_bound=degree_of_freedom.lower_bound,
368 upper_bound=degree_of_freedom.upper_bound
369 )
370 degrees_of_freedom.append(dof_schema)
372 self.flowsheet.optimizations.append(OptimizationSchema(
373 objective=objective_value.id,
374 sense=sense,
375 unfixed_variables=degrees_of_freedom,
376 ))
378 def check_dependencies(self) -> None:
379 """Verify that all property value dependencies are serialised"""
380 if self.scenario is not None and self.scenario.enable_optimization is True:
381 # No need to check since we are serialising everything
382 return
383 serialised_property_values = self.context.serialised_property_values
384 for dependency, prop_values in self.context.property_value_dependencies.items():
385 if dependency not in serialised_property_values:
386 dependency_prop_value = PropertyValue.objects.get(id=dependency)
387 prop_info_list_str = ", ".join([f"{prop_value.get_simulation_object().componentName}/{prop_value.property.displayName}" for prop_value in prop_values])
388 raise Exception(f"Dependency property {dependency_prop_value.get_simulation_object().componentName}/{dependency_prop_value.property.displayName} is not serialised, but is required by properties {prop_info_list_str}.")
390 def create_arcs(self):
391 """Serialise stream-like objects into arc connections for the flowsheet."""
392 streams = self.context.filter_object_type(
393 {"stream", "energy_stream", "ac_stream", "humid_air_stream"})
394 for stream in streams:
395 arc_schema = arc_adapter.create_arc(self.context, stream)
397 if arc_schema is not None:
398 self.flowsheet.arcs.append(arc_schema)
401# noinspection PyUnreachableCode
402def store_properties_schema(
403 properties_schema: list[SolvedPropertyValueSchema] | None,
404 flowsheet_id: int,
405 scenario_id: int | None = None,
406 solve_index: int | None = None
407) -> None:
408 """Persist solved property values and dynamic results to the database.
410 Args:
411 properties_schema: Collection of property payloads returned by IDAES.
412 flowsheet_id: Identifier of the flowsheet whose properties were solved.
413 scenario_id: Optional scenario identifier associated with the solve.
414 solve_index: Multi-steady-state index for the stored values, if any.
415 """
416 if not properties_schema:
417 return
418 properties_schema = [
419 SolvedPropertyValueSchema(**prop) if isinstance(prop, dict) else prop
420 for prop in properties_schema
421 ]
422 # create a id->property map
423 ids = [prop.id for prop in properties_schema]
424 property_values = PropertyValue.objects.filter(id__in=ids).select_related(
425 "property"
426 )
427 prop_map = {prop.id: prop for prop in property_values}
429 property_values = []
430 property_infos = []
431 dynamic_results = []
433 for prop_schema in properties_schema:
434 prop = prop_map.get(prop_schema.id, None)
435 if prop is None: 435 ↛ 436line 435 didn't jump to line 436 because the condition on line 435 was never true
436 raise Exception(
437 f"Property {prop_schema.id} not found in the database.")
439 property_info: PropertyInfo = prop.property
440 updated_value = prop_schema.value
442 updated_value = prop_schema.value
444 from_unit = prop_schema.unit
446 is_multi_steady_state = solve_index is not None
447 is_dynamics = scenario_id != None and isinstance(
448 updated_value, list) and len(updated_value) > 1
450 # If we're doing MSS or dynamics, we need to create associated dynamic results.
451 if is_multi_steady_state or is_dynamics:
452 dynamic_result = Solution(
453 property=prop,
454 flowsheet_id=flowsheet_id,
455 solve_index=solve_index,
456 scenario_id=scenario_id
457 )
459 # If we're dealing with dynamics, the updated_value is a list of values.
460 # Otherwise, it's a single scalar value.
461 dynamic_result.values = (updated_value if isinstance(updated_value, list)
462 else [updated_value])
463 dynamic_results.append(dynamic_result)
464 continue
466 # else we can continue as normal
467 if prop_schema.unknown_units and not can_convert(
468 from_unit, property_info.unit
469 ):
470 # we don't know the category of unit_type this unit is in!
471 # default to "unknown" with a custom unit
472 property_info.unitType = "unknown"
473 property_info.unit = from_unit
474 to_unit = from_unit
475 # try to find the unit_type by looping through all
476 # the units library and checking the first unit in the unit_type
477 # to see if it can be converted
478 for unit_type in units_library.keys():
479 default_unit = units_library[unit_type][0]["value"]
480 if can_convert(from_unit, default_unit):
481 # update the unitType
482 property_info.unitType = unit_type
483 property_info.unit = default_unit
484 to_unit = default_unit
485 break
486 property_infos.append(property_info)
487 else:
488 to_unit = property_info.unit
490 # TODO: better handling of multi-dimensional indexed properties.
491 if isinstance(updated_value, list):
492 val = updated_value[0]
493 else:
494 val = updated_value
496 new_value = convert_value(val, from_unit=from_unit, to_unit=to_unit)
497 prop.value = new_value
498 property_values.append(prop)
500 with transaction.atomic():
501 # save the property values
502 PropertyValue.objects.bulk_update(property_values, ["value"])
503 Solution.objects.bulk_create(
504 dynamic_results,
505 update_conflicts=True,
506 update_fields=["values"],
507 unique_fields=["pk"],
508 )
510 # save the property infos
511 PropertyInfo.objects.bulk_update(property_infos, ["unitType", "unit"])
514def save_all_initial_values(unit_models: dict[str, Any]) -> None:
515 """Persist initial values returned from IDAES for each unit model.
517 Args:
518 unit_models: Mapping of unit model ids to serialised initial value payloads.
519 """
520 simulation_objects = {unit_op.id: unit_op for unit_op in (SimulationObject.objects
521 .filter(id__in=unit_models.keys())
522 .only("id", "initial_values")
523 )}
525 for unit_model_id, unit_model in unit_models.items():
526 simulation_object = simulation_objects.get(int(unit_model_id))
527 if simulation_object is None:
528 # Attached ML surrogate sidecars are emitted with MLModel ids, not
529 # SimulationObject ids, so they have no flowsheet object to persist
530 # initial values back onto.
531 continue
532 initial_values = unit_model
533 simulation_object.initial_values = initial_values
535 SimulationObject.objects.bulk_update(
536 simulation_objects.values(), ["initial_values"])