Created
May 30, 2025 18:53
-
-
Save lpsinger/f419160bc2326e185c3ae620e6895430 to your computer and use it in GitHub Desktop.
Call METIS from Python
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import subprocess | |
| from pathlib import Path | |
| import tempfile | |
| from numpy import typing as npt | |
| def gpmetis( | |
| edge_weights_matrix: npt.ArrayLike, | |
| n_partitions: int, | |
| seed: int = 1, | |
| verbose: bool = False, | |
| ): | |
| """Partition a graph into disjoint subgraphs using METIS. | |
| Parameters | |
| ---------- | |
| edge_weights_matrix | |
| A square matrix of integer edge weights. For an unweighted graph, just | |
| pass an adjacency matrix consisting of zeros and ones. | |
| n_partitions | |
| The desired number of partitions. There may be fewer partitions | |
| returned. | |
| seed | |
| Random seed for METIS. | |
| verbose | |
| Set to True to print output from METIS. | |
| Returns | |
| ------- | |
| : | |
| Array of integer indices identifying the partition for each node. | |
| """ | |
| edge_weights_matrix = np.asarray(edge_weights_matrix, int) | |
| adjacency_matrix = edge_weights_matrix.astype(np.bool) | |
| if edge_weights_matrix.ndim != 2: | |
| raise ValueError("Edge weights matrix must be 2D") | |
| if edge_weights_matrix.shape[0] != edge_weights_matrix.shape[1]: | |
| raise ValueError("Edge weights matrix must be square") | |
| n_vertices = len(adjacency_matrix) | |
| n_edges = np.count_nonzero(adjacency_matrix[np.triu_indices(n_vertices, 1)]) | |
| with tempfile.TemporaryDirectory() as tmpdir: | |
| tmpdir_path = Path(tmpdir) | |
| in_path = tmpdir_path / "graph" | |
| out_path = in_path.with_suffix(f".part.{n_partitions}") | |
| with in_path.open("w") as f: | |
| print(n_vertices, n_edges, 1, file=f) | |
| for adjacency_row, edge_row in zip(adjacency_matrix, edge_weights_matrix): | |
| i = np.flatnonzero(adjacency_row) | |
| print(*chain.from_iterable(zip(i + 1, edge_row[i])), file=f) | |
| cmdline = ["gpmetis", "-seed", str(seed), in_path, str(n_partitions)] | |
| subprocess.run(cmdline, check=True, capture_output=not verbose) | |
| return np.loadtxt(out_path, dtype=np.intp) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment