warp.utils.graph_coloring_balance#

warp.utils.graph_coloring_balance(
edges,
node_colors,
color_count,
target_max_min_ratio,
*,
max_iterations=None,
)[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, when no further improvements are possible, or when max_iterations is reached.

  • max_iterations (int | None) – The non-negative maximum number of node recolorings to perform. If None, the number of graph nodes is used. Values must fit in a signed 32-bit integer. A warning is emitted if balancing stops because this limit is reached.

Returns:

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

Return type:

float

Example

import warp as wp

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

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