[REQ] Public API for Graph Coloring and Color Balancing

[shi-eric]

Summary

Expose graph coloring and color balancing functionality through a public Python API. Currently, users must access internal wp.context.runtime.core functions directly to use these features.

Motivation

Projects using Warp (including Newton) are currently forced to access internal APIs:

# Current usage pattern in Newton
num_colors = wp.context.runtime.core.wp_graph_coloring(
    graph.offsets.ptr,
    graph.indices.ptr,
    graph.offsets.shape[0] - 1,
    colors.ptr,
    algorithm  # e.g., wp.context.runtime.core.GRAPH_COLORING_MCS
)

max_min_ratio = wp.context.runtime.core.wp_balance_coloring(
    graph.offsets.ptr,
    graph.indices.ptr,
    graph.offsets.shape[0] - 1,
    colors.ptr,
    num_colors
)

This creates several issues:

1. API deprecation: With Warp v1.11's API cleanup, these internal APIs are now deprecated, but no public alternative exists.
2. Code reusability: These algorithms are implemented in Warp's native C++ code, making them difficult for downstream projects to extract or reimplement in Python without significant effort.

Proposed API

One possible approach would be to add public Python wrappers that provide a clean, Pythonic interface:

# Example of a possible public API
num_colors = wp.graph_coloring(
    offsets=graph.offsets,
    indices=graph.indices,
    colors=colors,
    algorithm="mcs"  # or "greedy"
)

max_min_ratio = wp.balance_coloring(
    offsets=graph.offsets,
    indices=graph.indices,
    colors=colors,
    num_colors=num_colors
)

Implementation Notes

• The underlying C++ implementations already exist and are exposed via ctypes
• Need to add Python wrappers that:
  • Accept Warp arrays instead of raw pointers
  • Provide string-based algorithm selection instead of requiring enum constants
  • Include proper docstrings with parameter descriptions and usage examples
  • Handle error cases gracefully

References

• Current usage in Newton: newton/_src/sim/graph_coloring.py
• Related to Warp v1.11 API cleanup: Internal APIs now emit deprecation warnings