Class GraphSampler

taskbench/graph_sampler.py:8–203

random.seed(0)
class GraphSampler:
    def __init__(self, graph: nx.Graph = None, file_name = None):
        if file_name:
            with open(file_name, "r") as f:
                data = json.load(f)

            # Represent your graph in NetworkX
            graph = nx.DiGraph()

            # Add nodes to the graph
            if "input-type" in data["nodes"][0]:
                for node in data["nodes"]:
                    graph.add_node(node["id"], desc=node["desc"], input_type=node["input-type"], output_type=node["output-type"])
            else:
                for node in data["nodes"]:
                    graph.add_node(node["id"], desc=node["desc"], parameters=node["parameters"])

            # Add edges to the graph
            for link in data["links"]:
                graph.add_edge(link["source"], link["target"], type=link["type"])

        self.graph = graph

    def sample_subgraph_by_weight(self, number_weights, method_weights):
        method = random.choices(list(method_weights.keys()), weights=list(method_weights.values()))[0]
        if method == "single":
            tool_number = 1
        else:
            tool_number = random.choices(list(number_weights.keys()), weights=list(number_weights.values()))[0]
        return self.sample_subgraph(tool_number, sample_method=method)

    def sample_subgraph(self, num_nodes=3, sample_method="chain"):
        seed_node = random.choice(list(self.graph.nodes))
        if sample_method == "single":
            sub_G = nx.DiGraph()
            sub_G.add_node(seed_node)
            return sub_G
        elif sample_method == "chain":
            return self.sample_subgraph_chain(seed_node, num_nodes)
        elif sample_method == "dag":
            return self.sample_subgraph_dag(seed_node, num_nodes)
        else:
            raise ValueError("Invalid sample method")

    def sample_subgraph_chain(self, seed_node, num_nodes):
        # Create a list to store the sub-graph nodes
        sub_graph_nodes = [seed_node]
        head_node = seed_node
        tail_node = seed_node
        edges = []

        # Keep adding nodes until we reach the desired number
        while len(sub_graph_nodes) < num_nodes:
            # Get the neighbors of the last node in the sub-graph
            head_node_neighbors = list(self.graph.predecessors(head_node))
            tail_node_neighbors = list(self.graph.successors(tail_node))
            neighbors = head_node_neighbors + tail_node_neighbors