Coverage for backend/django/idaes_factory/idaes_factory_context.py: 94%
122 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-07-22 05:22 +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._loaded_simulation_object_ids: set[int] | None = None
84 self._solve_index_data_cells_by_column: dict[int, float] | None = None
85 self._dynamic_data_cells_by_column: dict[int, list[float]] | None = None
86 self.load()
89 def load(self):
90 """Prefetch all simulation objects and related data for the flowsheet."""
91 # load all associated unit models, streams, property sets,
92 # properties, ports, etc. into the context
94 # db calls: simulation_objects, 1 for each prefetch_related\
95 sim_objs = self.group.get_recursive_simulation_objects()
96 property_values = PropertyValue.objects.select_related(
97 "controlManipulated",
98 "controlSetPoint",
99 "controlSetPoint__manipulated__property",
100 "controlManipulated__setPoint",
101 ).prefetch_related(
102 "indexedItems",
103 "controlSetPoint__manipulated__property__values"
104 )
105 if self.scenario is not None:
106 # Keep scenario data-column lookup in the existing property-value
107 # prefetch query. Querying DataColumn from each property adapter
108 # creates an N+1 during IDAES request serialisation.
109 property_values = property_values.annotate(
110 scenario_data_column_id=Subquery(
111 DataColumn.objects.filter(
112 scenario_id=self.scenario.id,
113 property_value_id=OuterRef("pk"),
114 )
115 .order_by("pk")
116 .values("pk")[:1]
117 )
118 )
120 self.simulation_objects = sim_objs.filter(
121 is_deleted=False
122 ).select_related(
123 "properties",
124 "recycleConnection",
125 ).prefetch_related(
126 Prefetch(
127 "properties__ContainedProperties",
128 queryset=PropertyInfo.objects
129 .select_related("recycleConnection")
130 .prefetch_related(
131 Prefetch("values", queryset=property_values)
132 )
133 ),
134 "ports",
135 "connectedPorts",
136 )
137 self._loaded_simulation_object_ids = None
139 def has_loaded_simulation_object(self, object_id: int | None) -> bool:
140 """Return whether a simulation object is loaded for the current solve group."""
141 if object_id is None: 141 ↛ 142line 141 didn't jump to line 142 because the condition on line 141 was never true
142 return False
143 if self._loaded_simulation_object_ids is None:
144 self._loaded_simulation_object_ids = {
145 simulation_object.id
146 for simulation_object in self.simulation_objects
147 }
148 return object_id in self._loaded_simulation_object_ids
150 def iter_loaded_property_infos(self):
151 """Yield property infos from the prefetched simulation object graph."""
152 for simulation_object in self.simulation_objects:
153 property_set = getattr(simulation_object, "properties", None)
154 if property_set is None: 154 ↛ 155line 154 didn't jump to line 155 because the condition on line 154 was never true
155 continue
156 yield from property_set.ContainedProperties.all()
158 def has_loaded_managed_properties(self) -> bool:
159 """Return whether the loaded context contains managed properties."""
160 return any(prop.managed for prop in self.iter_loaded_property_infos())
162 def get_data_column_id(self, property_value: PropertyValue) -> int | None:
163 """Return the scenario data column bound to a property value, if any."""
164 if hasattr(property_value, "scenario_data_column_id"):
165 return property_value.scenario_data_column_id
166 if self.scenario is None:
167 return None
169 # Fallback for property values that were not loaded by the annotated
170 # context prefetch. The normal adapter path stays query-free.
171 return (
172 DataColumn.objects.filter(
173 scenario_id=self.scenario.id,
174 property_value_id=property_value.id,
175 )
176 .values_list("id", flat=True)
177 .first()
178 )
180 def get_solve_index_data_cell_value(self, data_column_id: int) -> float:
181 """Return the cached data-cell value for the current MSS solve row."""
182 if self.solve_index is None: 182 ↛ 183line 182 didn't jump to line 183 because the condition on line 182 was never true
183 raise ValueError("A solve index is required to read MSS data cells.")
184 if self._solve_index_data_cells_by_column is None:
185 self._solve_index_data_cells_by_column = {
186 cell.data_column_id: cell.value
187 for cell in DataCell.objects.filter(data_row=self.solve_index)
188 }
189 return self._solve_index_data_cells_by_column[data_column_id]
191 def get_dynamic_data_cell_values(self, data_column_id: int) -> list[float]:
192 """Return cached dynamic input values grouped by data column."""
193 if self.scenario is None: 193 ↛ 194line 193 didn't jump to line 194 because the condition on line 193 was never true
194 return []
195 if self._dynamic_data_cells_by_column is None:
196 data_cells_by_column: dict[int, list[float]] = {}
197 for cell in DataCell.objects.filter(
198 data_column__scenario=self.scenario
199 ).order_by("data_row__index", "created_at"):
200 data_cells_by_column.setdefault(cell.data_column_id, []).append(
201 cell.value
202 )
203 self._dynamic_data_cells_by_column = data_cells_by_column
204 return self._dynamic_data_cells_by_column.get(data_column_id, [])
206 def track_property_value(self, property_value_id: int):
207 self.serialised_property_values[property_value_id] = True
209 def _add_dependency(self, depends_on_id: int, property_value: PropertyValue):
210 if depends_on_id not in self.property_value_dependencies:
211 self.property_value_dependencies[depends_on_id] = set()
212 self.property_value_dependencies[depends_on_id].add(property_value)
214 def add_property_value_dependency(self, property_value: PropertyValue):
215 """Register a dependency between property values for tracking purposes."""
216 self.track_property_value(property_value.id)
217 # Manipulated by
218 manipulated_by = getattr(property_value, "controlManipulated", None)
219 if manipulated_by is not None:
220 depends_on_id = manipulated_by.setPoint.id
221 self._add_dependency(depends_on_id, property_value)
223 # Formula dependencies
224 self.add_expression_dependency(property_value)
227 def add_expression_dependency(self, property_value: PropertyValue):
228 expression = property_value.formula
229 if not expression:
230 return
231 converted_expression = get_expression_dependencies(expression, self, property_value)
232 for dependency in converted_expression:
233 self._add_dependency(dependency, property_value) # add self-dependency to ensure it's tracked
235 def is_dynamic(self) -> bool:
236 """Return whether the bound scenario has dynamics enabled."""
237 if self.scenario is not None:
238 return self.scenario.enable_dynamics
239 else:
240 return False
242 # Updates the solve index so that the alread-loaded context can be used for multiple solves
243 # (Useful for multi-steady state)
244 def update_solve_index(self, index: int):
245 """Rebind the context to a different solve index for multi-solves.
247 Args:
248 index: Solve index associated with the current flowsheet.
249 """
250 row_filters = {
251 "index": index,
252 "flowsheet_state_id": self.group.flowsheet_state_id,
253 }
254 if self.scenario is not None: 254 ↛ 256line 254 didn't jump to line 256 because the condition on line 254 was always true
255 row_filters["scenario_id"] = self.scenario.id
256 self.solve_index = DataRow.objects.get(**row_filters)
257 self._solve_index_data_cells_by_column = None
259 def get_simulation_object(self, obj_id: int) -> SimulationObject:
260 """Return a simulation object by id using the prefetched queryset.
262 Args:
263 obj_id: Primary key of the simulation object to retrieve.
265 Returns:
266 The matching `SimulationObject` instance from the context cache.
267 """
268 return queryset_lookup.get_simulation_object(self.simulation_objects, obj_id)
271 def filter_object_type(self, include: set[str] = set()) -> list[SimulationObject]:
272 """Filter cached simulation objects to the provided set of type keys.
274 Args:
275 include: Unit operation type identifiers that should be returned.
277 Returns:
278 List of `SimulationObject` instances matching the requested types.
279 """
280 return queryset_lookup.filter_simulation_objects(self.simulation_objects, include)
283 def exclude_object_type(self, exclude: set[str]) -> list[SimulationObject]:
284 """Filter cached simulation objects by excluding specific type keys.
286 Args:
287 exclude: Unit operation type identifiers that should be omitted.
289 Returns:
290 List of `SimulationObject` instances not belonging to the excluded types.
291 """
292 return queryset_lookup.exclude_simulation_objects(self.simulation_objects, exclude)
294 def get_property(self, property_set: PropertySet, key: str) -> PropertyInfo:
295 """Retrieve a property from a property set by its key.
297 Args:
298 property_set: Property set that contains the requested property.
299 key: Identifier for the property within the set.
301 Returns:
302 The `PropertyInfo` instance found in the given property set.
303 """
304 return queryset_lookup.get_property(property_set, key)
306 def get_property_value(self, property: PropertyInfo, indexes: list[Any] | None = None) -> PropertyValue:
307 """Retrieve a property value, optionally constrained by index selections.
309 Args:
310 property: Property whose value object should be fetched.
311 indexes: Optional list of indices to select a specific value entry.
313 Returns:
314 The `PropertyValue` matching the property and index configuration.
315 """
316 return queryset_lookup.get_value_object(property, indexes)