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