jittor_geometric.partition

Graph partitioning and distributed computing utilities.

Description: Author: lusz Date: 2024-11-15 14:13:33

class jittor_geometric.partition.ChunkManager(output_dir=None, graph_data=None)[source]

Bases: object

Parameters:

output_dir (Optional[str])

__init__(output_dir=None, graph_data=None)[source]

Initialize the ChunkManager. :type output_dir: Optional[str] :param output_dir: Path for saving files. :type graph_data: :param graph_data: Original graph data, optional (only needed for partitioning).

metis_partition(edge_index, num_nodes, num_parts)[source]

Perform graph partitioning using Metis and save partition files. :type edge_index: :param edge_index: Edge indices of the graph (Var or np.ndarray). :type num_nodes: :param num_nodes: Number of nodes in the graph. :type num_parts: :param num_parts: Number of subgraphs to partition into. :return: Partition information.

partition_to_chunk(partition_file, edge_index, edge_weight, num_nodes, num_parts)[source]

Partition the graph into multiple chunks and create a GraphChunk object for each chunk.

Parameters:
  • partition_file – File path containing partition information

  • edge_index – Edge index (two arrays representing the start and end nodes of edges)

  • edge_weight – Weights of the edges

  • num_nodes – Total number of nodes in the graph

  • num_parts – Number of chunks to partition the graph into

Returns:

List of GraphChunk objects