Coverage for backend/django/flowsheetInternals/unitops/models/flow_tracking.py: 100%

1from flowsheetInternals.unitops.config.config_methods import get_connected_port_keys 

2from core.auxiliary.enums.unitOpGraphics import ConType 

3from core.auxiliary.enums import SimulationObjectClass 

4from typing import Callable, Iterable, TYPE_CHECKING 

5 

6if TYPE_CHECKING: 

7 from flowsheetInternals.unitops.models.SimulationObject import SimulationObject 

8 from flowsheetInternals.unitops.models.Port import Port 

9 

10 

11DOWNSTREAM_TRACKING_TERMINAL_UNIT_OP_TYPES = { 

12 SimulationObjectClass.Mixer, 

13 SimulationObjectClass.Splitter, 

14 SimulationObjectClass.Header, 

15 SimulationObjectClass.SimpleHeader, 

16} 

17 

18 

19def track_stream_flow( 

20 stream: "SimulationObject", 

21) -> tuple[set["SimulationObject"], set["SimulationObject"]]: 

22 """ 

23 Tracks the flow of a stream through the system.  

24 Stops at translator blocks, or recycles, where the property package has to be re-defined. 

25 Arguments: 

26 Stream: The stream simulation object to start from. Can also be a decision node. 

27 Returns: 

28 A tuple containing: 

29 - A list of all unit operations that the fluid goes through 

30 - A list of all streams that have properties of the fluid, 

31 

32 This is done by iterating down through the outlets of the stream and the unit ops of the stream. 

33 The Unit Operations in their config file have a propertyPackagePorts parameter that helps with  

34 figuring out which ports  

35 """ 

36 return _track_streams( 

37 stream, 

38 expand_ports=get_connected_ports, 

39 include_unit_op_for_port=lambda _port: True, 

40 stop_at_port=lambda _port: False, 

41 ) 

42 

43 

44def track_downstream_stream_flow( 

45 stream: "SimulationObject", 

46) -> tuple[set["SimulationObject"], set["SimulationObject"]]: 

47 """ 

48 Track only streams downstream of the source stream for visual highlighting. 

49 

50 This reuses the same traversal machinery as compound propagation, but only 

51 follows a stream into unit operation inlet ports and then out through outlet 

52 ports in the same property-package port group. 

53 """ 

54 return _track_streams( 

55 stream, 

56 expand_ports=get_downstream_connected_ports, 

57 include_unit_op_for_port=lambda port: port.direction == ConType.Inlet, 

58 stop_at_port=should_stop_downstream_tracking_at_port, 

59 ) 

60 

61 

62def _track_streams( 

63 stream: "SimulationObject", 

64 expand_ports: Callable[["Port"], Iterable["Port"]], 

65 include_unit_op_for_port: Callable[["Port"], bool], 

66 stop_at_port: Callable[["Port"], bool], 

67) -> tuple[set["SimulationObject"], set["SimulationObject"]]: 

68 """Traverse streams using a caller-provided port expansion rule.""" 

69 unit_ops = set() 

70 streams = set() 

71 

72 remaining_streams = [stream] 

73 

74 

75 #If we're starting with a decision node, get all the streams in its outlet. 

76 if stream.objectType == SimulationObjectClass.DecisionNode: 

77 # it's not actually a stream lol 

78 decision_node = stream 

79 remaining_streams = [] # remove it from the list of streams 

80 unit_ops.add(decision_node) 

81 # add its outlets instead. 

82 for port in decision_node.ports.filter(direction=ConType.Outlet): 

83 if port.stream is not None and port.stream not in streams: 

84 # If the port has a stream, add it to the remaining streams. 

85 remaining_streams.append(port.stream) 

86 

87 while remaining_streams: 

88 current_stream = remaining_streams.pop(0) 

89 streams.add(current_stream) 

90 stream_ports = current_stream.connectedPorts 

91 port: Port 

92 for port in stream_ports.all(): 

93 unit_op: "SimulationObject" = port.unitOp 

94 

95 # If it's a translator, we stop tracking flow (separate property package and compounds on other side.) 

96 if unit_op.objectType == SimulationObjectClass.Translator: 

97 continue 

98 

99 if include_unit_op_for_port(port): 

100 unit_ops.add(unit_op) 

101 

102 if stop_at_port(port): 

103 continue 

104 

105 # Recursively visit any other streams that we haven't already visited. 

106 ports_to_visit = expand_ports(port) 

107 for port in ports_to_visit: 

108 stream = port.stream 

109 if stream is not None and stream not in streams and stream not in remaining_streams: 

110 remaining_streams.append(stream) 

111 

112 return unit_ops, streams 

113 

114 

115def get_connected_ports(port: "Port"): 

116 """ 

117 From this port, find any other ports that are also connected to this unit operation, 

118 and represent the same "flow". 

119 """ 

120 unit_op: "SimulationObject" = port.unitOp 

121 connected_port_keys = get_connected_port_keys(port.key, unit_op.schema) 

122 

123 ports_to_visit = unit_op.ports.filter( 

124 key__in=connected_port_keys 

125 ) 

126 return ports_to_visit 

127 

128 

129def should_stop_downstream_tracking_at_port(port: "Port") -> bool: 

130 """ 

131 Stop visual downstream tracking at flow combining/splitting unit ops. 

132 

133 Compound/property propagation still uses ``track_stream_flow`` and can cross 

134 these unit operations. The canvas tracking path stops here so users can 

135 explicitly start a new tracked flow from whichever outlet they care about. 

136 """ 

137 return ( 

138 port.direction == ConType.Inlet 

139 and port.unitOp.objectType in DOWNSTREAM_TRACKING_TERMINAL_UNIT_OP_TYPES 

140 ) 

141 

142 

143def get_downstream_connected_ports(port: "Port"): 

144 """ 

145 From an inlet port, find outlet ports on the same unit operation that carry 

146 the same property-package flow. 

147 """ 

148 unit_op: "SimulationObject" = port.unitOp 

149 if port.direction != ConType.Inlet: 

150 return unit_op.ports.none() 

151 

152 connected_port_keys = get_connected_port_keys(port.key, unit_op.schema) 

153 ports_to_visit = unit_op.ports.filter( 

154 key__in=connected_port_keys, 

155 direction=ConType.Outlet, 

156 ) 

157 return ports_to_visit