warp.utils.graph_coloring_balance#

warp.utils.graph_coloring_balance( edges, node_colors, color_count, target_max_min_ratio, )[source]#

Balance the sizes of color groups in a graph coloring.

This function adjusts the color assignments to make the color groups more balanced in size, which improves load balancing when processing nodes in parallel by color. It attempts to move nodes between color groups while maintaining the validity of the coloring (no adjacent nodes share a color).

Parameters:

edges (array) – A 2D array of shape (edge_count, 2) containing edge pairs, where each row [i, j] represents an undirected edge between nodes i and j. Must be a CPU array with int32 dtype.
node_colors (array) – A 1D array of shape (node_count,) containing the current color assignments. Will be modified in-place with the balanced coloring. Must be a CPU array with int32 dtype.
color_count (int) – The number of colors in the current coloring (as returned by graph_coloring_assign()).
target_max_min_ratio (float) – The target ratio between the largest and smallest color group sizes. The algorithm will stop when this ratio is achieved or when no further improvements are possible.

Returns:

The actual max/min ratio achieved after balancing. This may be higher than the target if the graph structure prevents further balancing.

Return type:

float

Example

import warp as wp

edges = wp.array([[0, 1], [1, 2], [2, 3]], dtype=wp.int32, device="cpu")
colors = wp.empty(4, dtype=wp.int32, device="cpu")

color_count = wp.utils.graph_coloring_assign(edges, colors, wp.utils.GraphColoringAlgorithm.MCS)
ratio = wp.utils.graph_coloring_balance(edges, colors, color_count, 1.1)