warp.utils.graph_coloring_get_groups#

warp.utils.graph_coloring_get_groups(
node_colors,
color_count,
return_wp_array=True,
device='cpu',
)[source]#

Convert node colors into per-color groups.

This converts a node_colors array to a tuple of arrays, where each array contains the IDs of nodes with a certain color. The output can be used to process nodes in color order after graph_coloring_assign() and graph_coloring_balance().

Parameters:
  • node_colors (array) – A 1D array of shape (node_count,) containing color assignments. Must be a CPU array with int32 dtype.

  • color_count (int) – The number of colors in the coloring.

  • return_wp_array (bool) – Whether to return Warp arrays instead of numpy arrays.

  • device (Device | str | None) – Warp device for returned arrays when return_wp_array is True.

Returns:

A tuple of arrays. Each array contains the node IDs assigned to a color. Numpy arrays are returned by default, or Warp arrays when return_wp_array is True.

Return type:

tuple[ndarray, …] | tuple[array, …]