jax_sn.tools.graph_tools module¶
- compute_levels_iterative(dependency_matrix: Array, n_iterations_per_loop: int)[source]¶
Compute levels of elements in a graph defined by dependency_matrix using a safe while loop. Each element’s level is determined by the maximum level of its upstream elements plus one (zero upstream elements -> 0 level).
This function uses a Python loop to iteratively compute levels until all levels are computed or no new levels can be computed. This uses static number of iterations for each JAX-compiled step to avoid recompilation and to be able to use jax.lax.scan. This checks for cycles in the graph and raises an error if a cycle is detected.
For large enough graphs and n_iterations_per_loop, the Python overhead is insignificant compared to the JAX computation time. Since this function is safer than the other two level computation functions, it is recommended to use this function when possible.
- Parameters:
dependency_matrix (
Array) – A (n_elements, n_faces_per_element) array where each row contains the indices of upstream elements for each element. Use n_elements as a dummy index for faces that do not connect to any upstream element.n_iterations_per_loop (
int) – The number of iterations to perform in each JAX-compiled step.
- Return type:
Array- Returns:
An array of shape (n_elements,) containing the computed levels for each element.
- compute_levels_iterative_jax(dependency_matrix: Array, n_iterations_per_loop: int)[source]¶
Compute levels of elements in a graph defined by dependency_matrix using a safe while loop. Each element’s level is determined by the maximum level of its upstream elements plus one (zero upstream elements -> 0 level).
Uses a while loop and equinox runtime errors to check for cycles in the graph.
- Parameters:
dependency_matrix (
Array) – A (n_elements, n_faces_per_element) array where each row contains the indices of upstream elements for each element. Use n_elements as a dummy index for faces that do not connect to any upstream element.n_iterations_per_loop (
int) – The number of iterations to perform in each JAX-compiled step.
- Return type:
Array- Returns:
An array of shape (n_elements,) containing the computed levels for each element.