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

1import itertools 

2from typing import TYPE_CHECKING, Literal, TypedDict 

3 

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 

15 

16 

17 

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() 

26 

27 

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 

36 

37 

38 other_port_keys = get_connected_port_keys(port.key, unitop.schema) 

39 

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)) 

47 

48 

49 propagator = CompoundPropogation(stream, unique_compounds) 

50 propagator.run_for_stream(stream) 

51 

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() 

55 

56 

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) 

68 

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) 

72 

73 

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] 

77 

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() 

87 

88 for port in decisionNode.ports.all(): 

89 if port.stream: 

90 compound_key_set.update(_get_compound_keys(port.stream)) 

91 

92 # Create propagator with all compounds 

93 compound_keys = list(compound_key_set) 

94 

95 # Update the decision node itself 

96 propagator = CompoundPropogation(decisionNode, compound_keys) 

97 propagator.run() 

98 

99 

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) 

106 

107 

108class StreamsToUpdateItem(TypedDict): 

109 add: list[str] 

110 remove: list[str] 

111 

112 

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 = [] 

126 

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) 

138 

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() 

142 

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 

147 

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 

151 

152 # Create property info 

153 propertySet = unitop.properties 

154 

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) 

163 

164 # Note that this does not set the enabled status correctly, so we need to reevaluate_properties_enabled later.. 

165 

166 

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) 

178 

179 # The propertyValues to delete should already have been deleted by now. 

180 

181 indexes = [list(other_indexes), new_indexed_items] 

182 combinations = list(itertools.product(*indexes)) 

183 property_values = [] 

184 

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 ) 

210 

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 

216 

217 self.create_property_values.extend(property_values) 

218 

219 

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 

227  

228 Args: 

229 - include_source: Whether to include the source stream in the update 

230 """ 

231 

232 unit_ops, streams = track_stream_flow(self.source_stream) 

233 

234 for unit_op in unit_ops: 

235 self._update_compounds(unit_op) 

236 

237 for current_stream in streams: 

238 self._update_compounds(current_stream) 

239 

240 # Perform database changes 

241 self._handle_database_changes() 

242 

243 objects = streams | unit_ops 

244 for simulation_object in objects: 

245 simulation_object.reevaluate_properties_enabled() 

246 

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() 

254 

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 """ 

259 

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 ) 

271