jax_sn.domain.graph_generation module

Module for graph generation utilities in the context of geometric domains. This contains internal functions to compute inflow faces and element levels based on angles.

compute_face_inflow(geometric_domain: GeometricDomain, angle: Array)

Return type:

Array

compute_levels_domain(geometric_domain: GeometricDomain, angle: Array, n_max_iterations: int, while_loop: bool)

Computes computation levels for elements in a geometric domain given an angle. Each element’s level is determined by the maximum level of its upstream elements plus one (zero upstream elements -> 0 level).

See also compute_levels_iterative for a safer alternative that checks for cycles.

Parameters:

geometric_domainGeometricDomain

The geometric domain containing the mesh and adjacency information.

anglejnp.ndarray

A 1D array representing the angle for which to compute the levels.

n_max_iterationsint

The maximum number of iterations to perform. Default is 1500.

while_loopbool

If True, use a jax.while_loop to compute levels; otherwise, use jax.lax.scan.

Returns:

: jnp.ndarray

An array of shape (n_elements,) containing the computed levels for each element.

Return type:

Array

compute_levels_domain_safe(geometric_domain: GeometricDomain, angle: Array, n_iterations_per_loop: int)

Computes computation levels for elements in a geometric domain given an angle. 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 safe iterative approach that checks for cycles in the graph.

See also compute_levels_domain() for a jittable version (also for i.e. vmaps) that does not check for cycles or compute_levels_domain_safe_jax() for a jittable version that checks for cycles.

Parameters:

geometric_domainGeometricDomain

The geometric domain containing the mesh and adjacency information.

anglejnp.ndarray

A 1D array representing the angle for which to compute the levels.

n_max_iterationsint

The maximum number of iterations to perform.

Returns:

: jnp.ndarray

An array of shape (n_elements,) containing the computed levels for each element.

Return type:

Array

compute_levels_domain_safe_jax(geometric_domain: GeometricDomain, angle: Array, n_iterations_per_loop: int)

Computes computation levels for elements in a geometric domain given an angle. 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 safe iterative approach that checks for cycles in the graph.

Parameters:

geometric_domainGeometricDomain

The geometric domain containing the mesh and adjacency information.

anglejnp.ndarray

A 1D array representing the angle for which to compute the levels.

n_max_iterationsint

The maximum number of iterations to perform.

Returns:

: jnp.ndarray

An array of shape (n_elements,) containing the computed levels for each element.

Return type:

Array