jax_sn.tools.mesh_tools module

create_face_centroids_from_simplices(simplices: Array, face_template: Array)

Create centroids of each face in the simplex given by the face_template.

Fully vectorized.

Parameters:

• simplices (Array) – Array representing the vertices of the simplices.

• face_template (Array) – Array defining which nodes belong to each face of the simplex.

Returns:

Array of centroids for each face of the simplex.

Return type:

Array

compute_face_normal_vectors_from_3d_simplices(simplices: Array, face_template: Array)

Create outward pointing normal vectors for each face of a 3D simplex.

Parameters:

simplicesjnp.ndarray […, n_nodes_per_element, n_dim]

Array representing the vertices of the simplices. Is zero-padded to 3D if n_dim < 3.

face_templatejnp.ndarray [n_faces_per_element, n_vertices_per_face]

Array defining which nodes belong to each face of the simplex.

Returns:

: jnp.ndarray […, n_faces_per_element, 3]

Array of outward pointing normal vectors for each face of the simplex.

Return type:

Array

extract_elements_from_mesh(mesh, element_indices)

Extract a subset of elements from a mesh, remapping vertex indices.

Parameters:

• mesh (Tuple[jnp.ndarray, jnp.ndarray]) – The input mesh as (vertices, connectivity)

• element_indices (jnp.ndarray [n_selected_elements]) – The indices of the elements to extract

extract_boundaries_from_mesh(mesh, boundary_index, face_template)

Extract boundary faces from a mesh.

Parameters:

• mesh (Tuple[jnp.ndarray, jnp.ndarray]) – The input mesh as (vertices, connectivity)

• boundary_index (int or jnp.ndarray [n_elements]) – If int, the index of the boundary face to extract. If array, an array with one entry per element indicating which boundary face to extract for each element.

• face_template (jnp.ndarray [n_faces_per_element, n_vertices_per_face]) – The face template for the element type

create_bvh(vertices: Array, connectivity: Array)

Create a linear bounding volume hierarchy (LBVH) for the given mesh.

Parameters:

• vertices (Array) – The vertices of the mesh

• connectivity (Array) – The connectivity of the mesh

Returns:

bvh – The constructed LBVH for the mesh

Return type:

jax_sbgeom.jax_utils.raytracing.BVH

create_bvh_cached(vertices: Array, connectivity: Array)

Cached version of create_bvh, keyed on the identity of the arrays rather than on the identity of the containing domain. A domain is frequently rebuilt around the same mesh arrays (eqx.tree_at on the cross sections, Domain.from_geometric_domain_and_cross_sections, a pytree flatten/unflatten round trip), and all of those still hit this cache.

Entries are evicted once either array is garbage collected, so a cached BVH never keeps its mesh alive, and an id cannot be recycled while its entry is still present.

Building a BVH is a host-side operation on concrete arrays, so this cannot be used inside jitted functions.

Parameters:

• vertices (Array) – The vertices of the mesh

• connectivity (Array) – The connectivity of the mesh

Returns:

bvh – The constructed LBVH for the mesh, built on the first call for these arrays and reused afterwards.

Return type:

jax_sbgeom.jax_utils.raytracing.BVH

probe_bvh(bvh, points: Array)

Probe a BVH with given points.

Parameters:

• bvh (jax_sbgeom.jax_utils.raytracing.BVH) – The BVH to probe

• points (Array) – The points to probe the BVH with

Returns:

possible_hitsjnp.ndarray[n_points, n_max_hits]

The possible hits for each point. Each entry contains the indices of the elements that may contain the point. -1 if no hit.