jax_sn.tools.mesh_tools_numpy module

NumPy implementations of mesh setup functions from mesh_tools.py.

Avoids JAX compilation overhead during mesh initialisation. The main bottleneck in the JAX version is the argsort / lexsort inside _compute_adjacency and the subsequent _topologicalsorted_angle scan, which each require XLA compilation the first time they are called.

The functions here are drop-in replacements that return plain NumPy arrays; callers can wrap the results in jnp.array(…) as needed.

compute_topology_numpy(connectivity: ndarray, face_template: ndarray)[source]

NumPy drop-in for jax_sn.domain.geometric_domain.compute_topology.

Depends on the connectivity alone, never on the vertices.

Parameters:

  • connectivity (ndarray)

  • face_template (ndarray)

Returns:

etoe, etof, face_perm_idx

Return type:

Tuple[ndarray, ndarray, ndarray]