| 86 | return sub_G |
| 87 | |
| 88 | def sample_subgraph_dag(self, seed_node, num_nodes): |
| 89 | # Create a list to store the sub-graph nodes |
| 90 | sub_graph_nodes = [seed_node] |
| 91 | edges = [] |
| 92 | |
| 93 | # Keep adding nodes until we reach the desired number |
| 94 | while len(sub_graph_nodes) < num_nodes: |
| 95 | # Randomly select a node from the current sub-graph |
| 96 | node = random.choice(sub_graph_nodes) |
| 97 | # prec_neighbors = list(self.graph.predecessors(node)) |
| 98 | succ_neighbors = list(self.graph.successors(node)) |
| 99 | |
| 100 | if "input_type" in self.graph.nodes[node]: |
| 101 | # filter exisiting income edge type |
| 102 | prec_neighbors = [] |
| 103 | input_type = list(self.graph.nodes[node]["input_type"]) |
| 104 | all_in_edges = list(self.graph.in_edges(node, data=True)) |
| 105 | for edge in edges: |
| 106 | for ref_edge in all_in_edges: |
| 107 | if edge[0] == ref_edge[0] and edge[1] == ref_edge[1]: |
| 108 | input_type.remove(ref_edge[2]["type"]) |
| 109 | for edge in all_in_edges: |
| 110 | if edge[2]["type"] in input_type: |
| 111 | prec_neighbors.append(edge[0]) |
| 112 | else: |
| 113 | prec_neighbors = list(self.graph.predecessors(node)) |
| 114 | |
| 115 | neighbors = prec_neighbors + succ_neighbors |
| 116 | |
| 117 | # If the node has neighbors, randomly select one and add it to the sub-graph |
| 118 | if neighbors: |
| 119 | neighbor = random.choice(neighbors) |
| 120 | if neighbor not in sub_graph_nodes: |
| 121 | if neighbor in prec_neighbors: |
| 122 | edges.append((neighbor, node)) |
| 123 | else: |
| 124 | edges.append((node, neighbor)) |
| 125 | sub_graph_nodes.append(neighbor) |
| 126 | # If the node has no neighbors, select a new node from the original graph |
| 127 | else: |
| 128 | node = random.choice(list(self.graph.nodes)) |
| 129 | if node not in sub_graph_nodes: |
| 130 | sub_graph_nodes.append(node) |
| 131 | |
| 132 | # Create the sub-graph |
| 133 | sub_G = nx.DiGraph() |
| 134 | sub_G.add_nodes_from(sub_graph_nodes) |
| 135 | sub_G.add_edges_from(edges) |
| 136 | |
| 137 | return sub_G |
| 138 | |
| 139 | def sample_subgraph_random_walk(self, seed_node, num_nodes): |
| 140 | # Create a list to store the sub-graph nodes |