| 50 | raise ValueError("Invalid sample method") |
| 51 | |
| 52 | def sample_subgraph_chain(self, seed_node, num_nodes): |
| 53 | # Create a list to store the sub-graph nodes |
| 54 | sub_graph_nodes = [seed_node] |
| 55 | head_node = seed_node |
| 56 | tail_node = seed_node |
| 57 | edges = [] |
| 58 | |
| 59 | # Keep adding nodes until we reach the desired number |
| 60 | while len(sub_graph_nodes) < num_nodes: |
| 61 | # Get the neighbors of the last node in the sub-graph |
| 62 | head_node_neighbors = list(self.graph.predecessors(head_node)) |
| 63 | tail_node_neighbors = list(self.graph.successors(tail_node)) |
| 64 | neighbors = head_node_neighbors + tail_node_neighbors |
| 65 | |
| 66 | # If the node has neighbors, randomly select one and add it to the sub-graph |
| 67 | if len(neighbors) > 0: |
| 68 | neighbor = random.choice(neighbors) |
| 69 | if neighbor not in sub_graph_nodes: |
| 70 | if neighbor in head_node_neighbors: |
| 71 | sub_graph_nodes.insert(0, neighbor) |
| 72 | edges.insert(0, (neighbor, head_node)) |
| 73 | head_node = neighbor |
| 74 | else: |
| 75 | sub_graph_nodes.append(neighbor) |
| 76 | edges.append((tail_node, neighbor)) |
| 77 | tail_node = neighbor |
| 78 | else: |
| 79 | break |
| 80 | |
| 81 | # Create the sub-graph |
| 82 | sub_G = nx.DiGraph() |
| 83 | sub_G.add_nodes_from(sub_graph_nodes) |
| 84 | sub_G.add_edges_from(edges) |
| 85 | |
| 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 |