jax_sn.domain.geometric_domain module

compute_face_normals(vertices: Array, connectivity: Array, face_template: Array)[source]

The canonical face-normal derivation: the single place that knows how face_normals follows from vertices. Traceable, so it can be re-derived under jit/grad when the vertices move.

Parameters:

  • vertices (Array)

  • connectivity (Array)

  • face_template (Array)

Returns:

face_normals

Return type:

Array

compute_topology(connectivity: Array, face_template: Array)[source]

The mesh adjacency: etoe, etof and face_perm_idx.

Depends on the connectivity alone, never on the vertices – which is what lets a deformable parametrisation move the vertices without redoing this (the expensive, non-differentiable part of the domain setup).

Parameters:

  • connectivity (Array)

  • face_template (Array)

Returns:

etoe, etof, face_perm_idx

Return type:

jnp.ndarray, each [n_elements, n_faces_per_element]

class GeometricDomain(vertices: Array, connectivity: Array, face_normals: Array, etoe: Array, etof: Array, face_perm_idx: Array, face_template: Array)[source]

Bases: Module

Module representing a geometric domain. These contain precomputed connectivity, face normals, and adjacency information.

vertices: Array

Vertices of the domain. Shape is [n_vertices, n_dim]

connectivity: Array

Connectivity of the domain. Shape is [n_elements, n_vertices_per_element]

face_normals: Array

Face normals of the domain. Shape is [n_elements, n_faces_per_element, n_dim]

etoe: Array

Element-to-element adjacency. Shape is [n_elements, n_faces_per_element]. -1 indicates a boundary face (unconnected to another element).

etof: Array

Element-to-face adjacency. Shape is [n_elements, n_faces_per_element]. Each face has a unique index in the domain.

face_perm_idx: Array

Face permutation indices. Shape is [n_elements, n_faces_per_element]. Indicates how the face vertices are permuted between adjacent elements.

face_template: Array

Face template used to define faces of elements. Shape is [n_faces_per_element, n_vertices_per_face].

classmethod from_mesh(mesh: Tuple[Array, Array], face_template: Array)[source]

Constructs a GeometricDomain from a mesh and a face template (see jax_sn.tools.mesh_tools._face_template for an example template).

Parameters:

  • mesh (Tuple[Array, Array]) – A tuple containing vertices and connectivity of the mesh. vertices is a jnp.ndarray of shape [n_vertices, n_dim] and connectivity is a jnp.ndarray of shape [n_elements, n_vertices_per_element].

  • face_template (Array) – A 2D array representing the face template.

Returns:

domain – The constructed GeometricDomain.

Return type:

GeometricDomain

classmethod from_mesh_numpy(mesh: Tuple[Array, Array], face_template: Array)[source]

Constructs a GeometricDomain from a mesh and a face template (see jax_sn.tools.mesh_tools._face_template for an example template).

Uses a numpy implementation of the domain setup. This avoids compilation overheads of especially the lexsort, which can take significant time.

Parameters:

  • mesh (Tuple[Array, Array]) – A tuple containing vertices and connectivity of the mesh. vertices is a jnp.ndarray of shape [n_vertices, n_dim] and connectivity is a jnp.ndarray of shape [n_elements, n_vertices_per_element].

  • face_template (Array) – A 2D array representing the face template.

Returns:

domain – The constructed GeometricDomain.

Return type:

GeometricDomain

property n_elements

Number of elements in the domain

property n_faces_per_element

Number of faces per element in the domain

containing_cell(points: Array)[source]

Containing cell for each point in points. Uses a cached BVH, so cannot be used in a jitted function

Parameters:

points (Array) – The points for which to find the containing cell.

Returns:

element_indices – The indices of the containing cell for each point. -1 if no containing cell is found

Return type:

jnp.ndarray [n_points]

class Domain(vertices: Array, connectivity: Array, face_normals: Array, etoe: Array, etof: Array, face_perm_idx: Array, face_template: Array, cross_sections_indexed: CrossSectionsIndexed)[source]

Bases: GeometricDomain

Geometric domain with indexed cross sections. This is seperated from GeometricDomain because many algorithms do not depend on cross sections and therefore we want to keep the cross sections distinct.

cross_sections_indexed: CrossSectionsIndexed

Indexed cross sections for the domain.

classmethod from_mesh_and_cross_sections(mesh: Tuple[Array, Array], face_template: Array, cross_sections_indexed: CrossSectionsIndexed)[source]

Constructs a Domain from a mesh, face template, and indexed cross sections. Checks compatibility between the mesh and cross sections.

Parameters:

  • mesh (Tuple[Array, Array]) – A tuple containing vertices and connectivity of the mesh. vertices is a jnp.ndarray of shape [n_vertices, n_dim] and connectivity is a jnp.ndarray of shape [n_elements, n_vertices_per_element].

  • face_template (Array) – A 2D array representing the face template.

  • cross_sections_indexed (CrossSectionsIndexed) – Indexed cross sections for the domain.

Returns:

domain – The constructed Domain.

Return type:

Domain

classmethod from_mesh_and_cross_sections_numpy(mesh: Tuple[Array, Array], face_template: Array, cross_sections_indexed: CrossSectionsIndexed)[source]

Constructs a Domain from a mesh, face template, and indexed cross sections. Checks compatibility between the mesh and cross sections. Uses a numpy implementation of the domain setup. This avoids compilation overheads of especially the lexsort, which can take significant time.

Parameters:

  • mesh (Tuple[Array, Array]) – A tuple containing vertices and connectivity of the mesh. vertices is a jnp.ndarray of shape [n_vertices, n_dim] and connectivity is a jnp.ndarray of shape [n_elements, n_vertices_per_element].

  • face_template (Array) – A 2D array representing the face template.

  • cross_sections_indexed (CrossSectionsIndexed) – Indexed cross sections for the domain.

Returns:

domain – The constructed Domain.

Return type:

Domain

classmethod from_geometric_domain_and_cross_sections(geometric_domain: GeometricDomain, cross_sections_indexed: CrossSectionsIndexed)[source]

Constructs a Domain from a GeometricDomain and indexed cross sections. Checks compatibility between the geometric domain and cross sections.

Parameters:
Returns:

domain – The constructed Domain.

Return type:

Domain