Coverage for backend/django/flowsheetInternals/unitops/models/compound_propogation.py: 96%
136 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-07-22 05:22 +0000
« prev ^ index » next coverage.py v7.10.7, created at 2026-07-22 05:22 +0000
1import itertools
2from typing import TYPE_CHECKING, Literal, TypedDict
4from core.auxiliary.models.IndexedItem import IndexedItem, IndexChoices
5from core.auxiliary.models.PropertyValue import PropertyValue, PropertyValueIntermediate
6from core.auxiliary.models.PropertyInfo import PropertyInfo
7from flowsheetInternals.unitops.models.flow_tracking import track_stream_flow
8from flowsheetInternals.unitops.config.config_methods import get_connected_port_keys
9from flowsheetInternals.unitops.services.edit_operations.recorder import (
10 tracked_bulk_create,
11)
12if TYPE_CHECKING:
13 from flowsheetInternals.unitops.models.SimulationObject import SimulationObject
14 from flowsheetInternals.unitops.models.Port import Port
18def update_compounds_on_set(stream: "SimulationObject", expected_keys: list[str]) -> None:
19 """
20 Directly sets the compounds in a stream to match the expected compounds.
21 """
22 # Preserve the first-seen order while ignoring duplicate compound keys.
23 unique_expected_keys = list(dict.fromkeys(expected_keys))
24 propagator = CompoundPropogation(stream, unique_expected_keys)
25 propagator.run()
28def update_compounds_on_add_stream(port: "Port", stream: "SimulationObject") -> None:
29 """
30 Handles compound propogation when a stream is added to a port.
31 """
32 # no mole_frac_comp -> probably a bus
33 if not stream.properties.ContainedProperties.filter(key="mole_frac_comp").exists():
34 return
35 unitop = port.unitOp
38 other_port_keys = get_connected_port_keys(port.key, unitop.schema)
40 other_ports = unitop.ports.filter(
41 key__in=other_port_keys,
42 )
43 streams = [port.stream for port in other_ports if port.stream]
44 unique_compounds = set()
45 for neighbour_stream in streams:
46 unique_compounds.update(_get_compound_keys(neighbour_stream))
49 propagator = CompoundPropogation(stream, unique_compounds)
50 propagator.run_for_stream(stream)
52 if neighbour_stream is not None: 52 ↛ exitline 52 didn't return from function 'update_compounds_on_add_stream' because the condition on line 52 was always true
53 stream.propertyPackageType = neighbour_stream.propertyPackageType
54 stream.save()
57def update_compounds_on_merge(inlet_stream: "SimulationObject", outlet_stream: "SimulationObject") -> None:
58 """
59 Handles compound propogation when two streams are merged.
60 """
61 expected_keys = _get_compound_keys(outlet_stream)
62 other_expected_keys = _get_compound_keys(inlet_stream)
63 # combine the lists, removing duplicates
64 expected_keys.extend(other_expected_keys)
65 expected_keys = list(set(expected_keys))
66 propogator = CompoundPropogation(inlet_stream, expected_keys)
67 propogator.run(include_source=False)
69 from flowsheetInternals.unitops.models.property_package_propogation import propogate_property_package
70 inlet_stream.propertyPackageType = outlet_stream.propertyPackageType
71 propogate_property_package(inlet_stream)
74def _get_compound_keys(stream: "SimulationObject") -> list[str]:
75 indexed_items = IndexedItem.objects.filter(owner=stream, type=IndexChoices.Compound)
76 return [item.key for item in indexed_items]
78# "and" should never be in a function name,
79# because one function should do one thing.
80# TODO: Split this into two seperate functions.
81def update_decision_node_and_propagate(decisionNode: "SimulationObject", updated_via_right_click: bool = False):
82 """
83 Look through all the decision node's inlet streams to find a sum of the compounds
84 Then call propagate for every outlet stream of the decision node
85 """
86 compound_key_set = set()
88 for port in decisionNode.ports.all():
89 if port.stream:
90 compound_key_set.update(_get_compound_keys(port.stream))
92 # Create propagator with all compounds
93 compound_keys = list(compound_key_set)
95 # Update the decision node itself
96 propagator = CompoundPropogation(decisionNode, compound_keys)
97 propagator.run()
100def run_for_stream(stream: "SimulationObject", expected_compounds: list[str]) -> None:
101 """
102 Runs the compound propogation for a specific stream (doesn't propogate to other streams).
103 """
104 propagator = CompoundPropogation(stream, expected_compounds)
105 propagator.run_for_stream(stream)
108class StreamsToUpdateItem(TypedDict):
109 add: list[str]
110 remove: list[str]
113class CompoundPropogation:
114 """
115 This class handles the addition and removal of chemical compounds in process streams.
116 It manages how compounds propagate through connected equipment (like pumps, mixers, etc.)
117 """
118 def __init__(self, stream: "SimulationObject", expected_compounds: list[str]) -> None:
119 self.source_stream: "SimulationObject" = stream
120 self.expected_compounds: list[str] = expected_compounds
121 self.create_property_values = []
122 self.create_indexed_items = []
123 self.create_property_value_indexed_items = []
124 self.delete_property_values = []
125 self.delete_indexed_items = []
127 def _update_compounds(self, unitop: "SimulationObject") -> None:
128 """
129 Updates the compounds in an object, deleting any that are no longer present
130 and adding any that are now present.
131 """
132 new_indexed_items = []
133 if "compound" in unitop.schema.indexSets:
134 # Remove any compounds that are not in the list of expected compounds
135 compound_indexed_items = IndexedItem.objects.filter(owner=unitop, type=IndexChoices.Compound)
136 indexed_items_to_remove = compound_indexed_items.exclude(key__in=self.expected_compounds)
137 indexed_items_to_keep = compound_indexed_items.filter(key__in=self.expected_compounds)
139 property_values_to_remove = PropertyValue.objects.filter(indexedItems__in=indexed_items_to_remove)
140 property_values_to_remove.delete()
141 indexed_items_to_remove.delete()
143 present_compounds: list[str] = indexed_items_to_keep.values_list('key', flat=True)
144 for compound in self.expected_compounds:
145 if compound not in present_compounds:
146 # If the compound is not already present, create a new indexed item for it
148 new_index = IndexedItem(owner=unitop, key=compound, displayName=compound, type=IndexChoices.Compound, flowsheet_state=unitop.flowsheet_state)
149 new_indexed_items.append(new_index) # List of new indexed items for this unit op
150 self.create_indexed_items.append(new_index) # List of all new indexed items for all unit ops
152 # Create property info
153 propertySet = unitop.properties
155 # Update compound-dependent properties ( e.g compound separator split fraction)
156 for property in propertySet.ContainedProperties.all ():
157 if property.key not in unitop.schema.properties: 157 ↛ 158line 157 didn't jump to line 158 because the condition on line 157 was never true
158 continue # This is a custom property so we don't have to worry about adding compounds to it.
159 property_schema = unitop.schema.properties[property.key]
160 index_sets = property_schema.indexSets
161 if index_sets is not None and "compound" in index_sets:
162 self._property_add_remove(unitop,property,new_indexed_items)
164 # Note that this does not set the enabled status correctly, so we need to reevaluate_properties_enabled later..
167 def _property_add_remove(self, unitop:"SimulationObject", propertyInfo: "PropertyInfo",new_indexed_items: list[IndexedItem]) -> None:
168 """
169 Adds or removes the indexes and property values from the propertyInfo
170 """
171 # get the other indexes for this property
172 schema = unitop.schema.properties[propertyInfo.key]
173 group = schema.propertySetGroup
174 enabled = propertyInfo.key in unitop.schema.propertySetGroups[group].stateVars
175 other_indexes = IndexedItem.objects.filter(
176 owner=unitop, type__in=schema.indexSets
177 ).exclude(type=IndexChoices.Compound)
179 # The propertyValues to delete should already have been deleted by now.
181 indexes = [list(other_indexes), new_indexed_items]
182 combinations = list(itertools.product(*indexes))
183 property_values = []
185 if schema.sumToOne:
186 # Leave the last combination as disabled, as they can be calculated from the other compounds.
187 cutoff = len(combinations) - 1 - len(self.create_indexed_items)
188 else:
189 # all combinations should be enabled
190 cutoff = len(combinations)
191 for index, idxes in enumerate(combinations):
192 if index <= cutoff:
193 property_value = PropertyValue(value=schema.value, property=propertyInfo, enabled=enabled, flowsheet_state=unitop.flowsheet_state)
194 else:
195 property_value = PropertyValue(value=schema.value, property=propertyInfo, enabled=False, flowsheet_state=unitop.flowsheet_state)
196 for indexed_item in idxes:
197 self.create_property_value_indexed_items.append(
198 PropertyValueIntermediate(propertyvalue=property_value, indexeditem=indexed_item)
199 )
200 property_values.append(property_value)
201 if len(combinations) == 0:
202 # There are no other indexes other than for this compound. (e.g mole_frac_comp)
203 # just create a property value for each indexed item
204 for indexed_item in new_indexed_items:
205 property_value = PropertyValue(value=schema.value, property=propertyInfo, enabled=enabled, flowsheet_state=unitop.flowsheet_state)
206 property_values.append(property_value)
207 self.create_property_value_indexed_items.append(
208 PropertyValueIntermediate(propertyvalue=property_value, indexeditem=indexed_item)
209 )
211 if propertyInfo.key == "mole_frac_comp" and len(self.expected_compounds) == 1:
212 # If this is a mole_frac_comp property and there is only one compound, set the value to 1.0
213 for property_value in property_values:
214 property_value.value = 1.0
215 property_value.displayValue = 1.0
217 self.create_property_values.extend(property_values)
220 def run(self, include_source: bool = True) -> None:
221 """
222 Main method that orchestrates the entire compound update process:
223 1. Identifies all affected streams
224 2. Prepares database changes (additions and removals)
225 3. Updates the database in bulk for better performance
226 4. Handles read/write permissions for properties
228 Args:
229 - include_source: Whether to include the source stream in the update
230 """
232 unit_ops, streams = track_stream_flow(self.source_stream)
234 for unit_op in unit_ops:
235 self._update_compounds(unit_op)
237 for current_stream in streams:
238 self._update_compounds(current_stream)
240 # Perform database changes
241 self._handle_database_changes()
243 objects = streams | unit_ops
244 for simulation_object in objects:
245 simulation_object.reevaluate_properties_enabled()
247 def run_for_stream(self, stream: "SimulationObject") -> None:
248 """
249 Runs the compound propogation only for a specific stream (doesn't propogate to other streams).
250 """
251 self._update_compounds(stream)
252 self._handle_database_changes()
253 stream.reevaluate_properties_enabled()
255 def _handle_database_changes(self) -> None: # First
256 """
257 Handles the database changes for the compound propogation based on the streams_to_update dictionary.
258 """
260 # Perform bulk operations
261 if self.delete_property_values: 261 ↛ 262line 261 didn't jump to line 262 because the condition on line 261 was never true
262 PropertyValue.objects.filter(id__in=self.delete_property_values).delete()
263 IndexedItem.objects.filter(id__in=self.delete_indexed_items).delete()
264 if self.create_property_values:
265 tracked_bulk_create(PropertyValue.objects, self.create_property_values)
266 tracked_bulk_create(IndexedItem.objects, self.create_indexed_items)
267 tracked_bulk_create(
268 PropertyValueIntermediate.objects,
269 self.create_property_value_indexed_items,
270 )