Coverage for backend/django/idaes_factory/idaes_factory_context.py: 95%
114 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-06-23 21:51 +0000
1from typing import Any
3from django.db import transaction
4from django.db.models import QuerySet, Prefetch, OuterRef, Subquery
6from ahuora_builder_types.flowsheet_schema import PropertyPackageType
7from idaes_factory.adapters.convert_expression import get_expression_dependencies
8from core.auxiliary.models import PropertyInfo, PropertySet
9from core.auxiliary.models.PropertyValue import PropertyValue
10from core.auxiliary.models.DataCell import DataCell
11from core.auxiliary.models.DataColumn import DataColumn
12from core.auxiliary.models.DataRow import DataRow
13from flowsheetInternals.graphicData.models.groupingModel import Grouping
14from flowsheetInternals.unitops.models.SimulationObject import SimulationObject
15from . import queryset_lookup
16from core.auxiliary.models.Scenario import Scenario
19ParameterName = str
20ParameterValue = float
21LiveSolveParams = dict[ParameterName,ParameterValue]
22DependencyId = int
23PropertyValueDependencies = dict[DependencyId, set[PropertyValue]]
24"""
25Parameters for solving. Used when solving live, to specify the values that each of the input columns in the multi-steady state simulation should be set to.
26Key is the key of the column, should be one of the input columns used.
27Value is a float value to set it to.
28"""
32class IdaesFactoryContext:
33 """
34 The IdaesFactoryContext class provides:
36 - Efficient retrieval of related objects from the
37 database, for both serialisation and reloading
38 of the flowsheet
39 - A container for context variables such as time
40 steps, property packages, etc.
42 It is used within idaes_factory to provide context
43 through the various adapter classes.
44 """
47 def __init__(
48 self,
49 group_id: int,
50 time_steps: list[int] = [0],
51 time_step_size: float = 1.0,
52 require_variables_fixed: bool = False,
53 solve_index: int | None = None, # If it's from multi-steady state, index to store values at
54 scenario : Scenario = None,
55 ) -> None:
56 """Initialise the context with eager-loaded flowsheet data.
58 Args:
59 group_id: Identifier of the flowsheet being prepared.
60 time_steps: Time indices that adapters should serialise.
61 require_variables_fixed: Whether adapters must enforce fixed variables.
62 solve_index: Optional multi-steady-state solve index to bind.
63 scenario: Scenario providing state and dynamics configuration.
64 """
65 self.group_id = group_id
66 self.group = Grouping.objects.get(id=group_id)
67 self.simulation_objects: QuerySet[SimulationObject] = None
69 if solve_index is not None:
70 self.solve_index = DataRow.objects.get(index=solve_index, scenario_id=scenario.id)
71 else:
72 self.solve_index = None
74 # context vars
75 self.time_steps = time_steps
76 self.time_step_size = time_step_size
77 self.require_variables_fixed = require_variables_fixed
78 self.property_packages: list[PropertyPackageType] = []
79 self.expression_values = {}
80 self.scenario = scenario
81 self.property_value_dependencies: PropertyValueDependencies = {} # dict mapping property value id to the set of property value ids that depend on it (i.e. controlManipulated or formula dependencies)
82 self.serialised_property_values = {} # set of property value ids that have been serialised
83 self._solve_index_data_cells_by_column: dict[int, float] | None = None
84 self._dynamic_data_cells_by_column: dict[int, list[float]] | None = None
85 self.load()
88 def load(self):
89 """Prefetch all simulation objects and related data for the flowsheet."""
90 # load all associated unit models, streams, property sets,
91 # properties, ports, etc. into the context
93 # db calls: simulation_objects, 1 for each prefetch_related\
94 sim_objs = self.group.get_recursive_simulation_objects()
95 property_values = PropertyValue.objects.select_related(
96 "controlManipulated",
97 "controlSetPoint",
98 "controlSetPoint__manipulated__property",
99 "controlManipulated__setPoint",
100 ).prefetch_related(
101 "indexedItems",
102 "controlSetPoint__manipulated__property__values"
103 )
104 if self.scenario is not None:
105 # Keep scenario data-column lookup in the existing property-value
106 # prefetch query. Querying DataColumn from each property adapter
107 # creates an N+1 during IDAES request serialisation.
108 property_values = property_values.annotate(
109 scenario_data_column_id=Subquery(
110 DataColumn.objects.filter(
111 scenario_id=self.scenario.id,
112 property_value_id=OuterRef("pk"),
113 )
114 .order_by("pk")
115 .values("pk")[:1]
116 )
117 )
119 self.simulation_objects = sim_objs.filter(
120 is_deleted=False
121 ).select_related(
122 "properties",
123 "recycleConnection",
124 ).prefetch_related(
125 Prefetch(
126 "properties__ContainedProperties",
127 queryset=PropertyInfo.objects
128 .select_related("recycleConnection")
129 .prefetch_related(
130 Prefetch("values", queryset=property_values)
131 )
132 ),
133 "ports",
134 "connectedPorts",
135 )
137 def iter_loaded_property_infos(self):
138 """Yield property infos from the prefetched simulation object graph."""
139 for simulation_object in self.simulation_objects:
140 property_set = getattr(simulation_object, "properties", None)
141 if property_set is None: 141 ↛ 142line 141 didn't jump to line 142 because the condition on line 141 was never true
142 continue
143 yield from property_set.ContainedProperties.all()
145 def has_loaded_managed_properties(self) -> bool:
146 """Return whether the loaded context contains managed properties."""
147 return any(prop.managed for prop in self.iter_loaded_property_infos())
149 def get_data_column_id(self, property_value: PropertyValue) -> int | None:
150 """Return the scenario data column bound to a property value, if any."""
151 if hasattr(property_value, "scenario_data_column_id"):
152 return property_value.scenario_data_column_id
153 if self.scenario is None:
154 return None
156 # Fallback for property values that were not loaded by the annotated
157 # context prefetch. The normal adapter path stays query-free.
158 return (
159 DataColumn.objects.filter(
160 scenario_id=self.scenario.id,
161 property_value_id=property_value.id,
162 )
163 .values_list("id", flat=True)
164 .first()
165 )
167 def get_solve_index_data_cell_value(self, data_column_id: int) -> float:
168 """Return the cached data-cell value for the current MSS solve row."""
169 if self.solve_index is None: 169 ↛ 170line 169 didn't jump to line 170 because the condition on line 169 was never true
170 raise ValueError("A solve index is required to read MSS data cells.")
171 if self._solve_index_data_cells_by_column is None:
172 self._solve_index_data_cells_by_column = {
173 cell.data_column_id: cell.value
174 for cell in DataCell.objects.filter(data_row=self.solve_index)
175 }
176 return self._solve_index_data_cells_by_column[data_column_id]
178 def get_dynamic_data_cell_values(self, data_column_id: int) -> list[float]:
179 """Return cached dynamic input values grouped by data column."""
180 if self.scenario is None: 180 ↛ 181line 180 didn't jump to line 181 because the condition on line 180 was never true
181 return []
182 if self._dynamic_data_cells_by_column is None:
183 data_cells_by_column: dict[int, list[float]] = {}
184 for cell in DataCell.objects.filter(
185 data_column__scenario=self.scenario
186 ).order_by("data_row__index", "created_at"):
187 data_cells_by_column.setdefault(cell.data_column_id, []).append(
188 cell.value
189 )
190 self._dynamic_data_cells_by_column = data_cells_by_column
191 return self._dynamic_data_cells_by_column.get(data_column_id, [])
193 def track_property_value(self, property_value_id: int):
194 self.serialised_property_values[property_value_id] = True
196 def _add_dependency(self, depends_on_id: int, property_value: PropertyValue):
197 if depends_on_id not in self.property_value_dependencies:
198 self.property_value_dependencies[depends_on_id] = set()
199 self.property_value_dependencies[depends_on_id].add(property_value)
201 def add_property_value_dependency(self, property_value: PropertyValue):
202 """Register a dependency between property values for tracking purposes."""
203 self.track_property_value(property_value.id)
204 # Manipulated by
205 manipulated_by = getattr(property_value, "controlManipulated", None)
206 if manipulated_by is not None:
207 depends_on_id = manipulated_by.setPoint.id
208 self._add_dependency(depends_on_id, property_value)
210 # Formula dependencies
211 self.add_expression_dependency(property_value)
214 def add_expression_dependency(self, property_value: PropertyValue):
215 expression = property_value.formula
216 if not expression:
217 return
218 converted_expression = get_expression_dependencies(expression, self, property_value)
219 for dependency in converted_expression:
220 self._add_dependency(dependency, property_value) # add self-dependency to ensure it's tracked
222 def is_dynamic(self) -> bool:
223 """Return whether the bound scenario has dynamics enabled."""
224 if self.scenario is not None:
225 return self.scenario.enable_dynamics
226 else:
227 return False
229 # Updates the solve index so that the alread-loaded context can be used for multiple solves
230 # (Useful for multi-steady state)
231 def update_solve_index(self, index: int):
232 """Rebind the context to a different solve index for multi-solves.
234 Args:
235 index: Solve index associated with the current flowsheet.
236 """
237 row_filters = {"index": index, "flowsheet_id": self.group.flowsheet.id}
238 if self.scenario is not None: 238 ↛ 240line 238 didn't jump to line 240 because the condition on line 238 was always true
239 row_filters["scenario_id"] = self.scenario.id
240 self.solve_index = DataRow.objects.get(**row_filters)
241 self._solve_index_data_cells_by_column = None
243 def get_simulation_object(self, obj_id: int) -> SimulationObject:
244 """Return a simulation object by id using the prefetched queryset.
246 Args:
247 obj_id: Primary key of the simulation object to retrieve.
249 Returns:
250 The matching `SimulationObject` instance from the context cache.
251 """
252 return queryset_lookup.get_simulation_object(self.simulation_objects, obj_id)
255 def filter_object_type(self, include: set[str] = set()) -> list[SimulationObject]:
256 """Filter cached simulation objects to the provided set of type keys.
258 Args:
259 include: Unit operation type identifiers that should be returned.
261 Returns:
262 List of `SimulationObject` instances matching the requested types.
263 """
264 return queryset_lookup.filter_simulation_objects(self.simulation_objects, include)
267 def exclude_object_type(self, exclude: set[str]) -> list[SimulationObject]:
268 """Filter cached simulation objects by excluding specific type keys.
270 Args:
271 exclude: Unit operation type identifiers that should be omitted.
273 Returns:
274 List of `SimulationObject` instances not belonging to the excluded types.
275 """
276 return queryset_lookup.exclude_simulation_objects(self.simulation_objects, exclude)
278 def get_property(self, property_set: PropertySet, key: str) -> PropertyInfo:
279 """Retrieve a property from a property set by its key.
281 Args:
282 property_set: Property set that contains the requested property.
283 key: Identifier for the property within the set.
285 Returns:
286 The `PropertyInfo` instance found in the given property set.
287 """
288 return queryset_lookup.get_property(property_set, key)
290 def get_property_value(self, property: PropertyInfo, indexes: list[Any] | None = None) -> PropertyValue:
291 """Retrieve a property value, optionally constrained by index selections.
293 Args:
294 property: Property whose value object should be fetched.
295 indexes: Optional list of indices to select a specific value entry.
297 Returns:
298 The `PropertyValue` matching the property and index configuration.
299 """
300 return queryset_lookup.get_value_object(property, indexes)