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.unitOpData 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 

11 

12def track_stream_flow( 

13 stream: "SimulationObject", 

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

15 """ 

16 Tracks the flow of a stream through the system.  

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

18 Arguments: 

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

20 Returns: 

21 A tuple containing: 

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

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

24 

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

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

27 figuring out which ports  

28 """ 

29 return _track_streams( 

30 stream, 

31 expand_ports=get_connected_ports, 

32 include_unit_op_for_port=lambda _port: True, 

33 ) 

34 

35 

36def track_downstream_stream_flow( 

37 stream: "SimulationObject", 

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

39 """ 

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

41 

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

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

44 ports in the same property-package port group. 

45 """ 

46 return _track_streams( 

47 stream, 

48 expand_ports=get_downstream_connected_ports, 

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

50 ) 

51 

52 

53def _track_streams( 

54 stream: "SimulationObject", 

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

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

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

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

59 unit_ops = set() 

60 streams = set() 

61 

62 remaining_streams = [stream] 

63 

64 

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

66 if stream.objectType == SimulationObjectClass.DecisionNode: 

67 # it's not actually a stream lol 

68 decision_node = stream 

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

70 unit_ops.add(decision_node) 

71 # add its outlets instead. 

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

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

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

75 remaining_streams.append(port.stream) 

76 

77 while remaining_streams: 

78 current_stream = remaining_streams.pop(0) 

79 streams.add(current_stream) 

80 stream_ports = current_stream.connectedPorts 

81 port: Port 

82 for port in stream_ports.all(): 

83 unit_op: "SimulationObject" = port.unitOp 

84 

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

86 if unit_op.objectType == SimulationObjectClass.Translator: 

87 continue 

88 

89 if include_unit_op_for_port(port): 

90 unit_ops.add(unit_op) 

91 

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

93 ports_to_visit = expand_ports(port) 

94 for port in ports_to_visit: 

95 stream = port.stream 

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

97 remaining_streams.append(stream) 

98 

99 return unit_ops, streams 

100 

101 

102def get_connected_ports(port: "Port"): 

103 """ 

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

105 and represent the same "flow". 

106 """ 

107 unit_op: "SimulationObject" = port.unitOp 

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

109 

110 ports_to_visit = unit_op.ports.filter( 

111 key__in=connected_port_keys 

112 ) 

113 return ports_to_visit 

114 

115 

116def get_downstream_connected_ports(port: "Port"): 

117 """ 

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

119 the same property-package flow. 

120 """ 

121 unit_op: "SimulationObject" = port.unitOp 

122 if port.direction != ConType.Inlet: 

123 return unit_op.ports.none() 

124 

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

126 ports_to_visit = unit_op.ports.filter( 

127 key__in=connected_port_keys, 

128 direction=ConType.Outlet, 

129 ) 

130 return ports_to_visit