jax_sn.solution_domain.basix_fem module

create_all_combined_face_quadratures(element: FiniteElement)

Create all combined face quadratures for an element, i.e., for each face, for each permutation of the face vertices, and for each connected face. Returns a list of tuples (points_own, points_conn, weights) where points_own are the quadrature points on the own face, points_conn are the quadrature points on the connected face (after applying the permutation), and weights are the quadrature weights.

Parameters:

element (FiniteElement) – The finite element for which to create the combined face quadratures.

Returns:

A list of tuples (points_own, points_conn, weights) for each combination of face, permutation, and connected face. Ordering is

Return type:

list of tuples

class BasixLagrangianSimplex(degree: int, dimension: int)

Bases: Element

Finite element based on Basix Lagrangian elements on simplices (interval, triangle, tetrahedron).

This has a linear jacobian mapping between reference and real coordinates. Basix is used for the basis functions and quadrature generation. These do not change during the simulation, so non-jittable Basix functions can be used during setup and precomputation.

property volume_quadrature

Returns the quadrature points and weights for the element that are sufficiently accurate for the volume integrals (in other words, 2 times the maximum polynomial degree of the basis functions)

property has_linear_jacobian: bool

Returns whether the element has a linear Jacobian

property face_quadratures

Returns the quadrature points and weights for the faces of the element.

property all_connection_face_quadratures: List[Tuple[ndarray, ndarray, ndarray]]

Returns the quadrature points, and weights for all possible face connections of each face.

The order needs to be as follows (row major):

most significant index:

face number of this element

middle index:

permutation number of the face (as defined by its Lehmer code like in _get_all_permutations in src/jax_sn/tools/permutation_tools.py)

least significant index:

face number of the connected element

Purposefully not a stacked array since in principle the number of permutations can vary per face (e.g. for non-simplicial elements).

property n_faces: int

Returns the number of faces of the element.

property n_basis: int

Returns the number of basis functions of the element.

property face_template: List[Tuple]

Returns the face template of the element.

The face template is a list of tuples, where each tuple contains the indices of the vertices that make up a face.

For example, for a triangle element, the face template would be [(0, 1), (1, 2), (2, 0)].

n_vertices_per_face(face_idx)

Returns the number of vertices per face of the element.

Parameters:

face_idx (int) – The index of the face.

Return type:

int

tabulate_basis_functions_first_derivative(points)

Returns the basis functions and their first derivatives at the given points.

Parameters:

points (jnp.ndarray [n_points, n_dim]) – The points at which to evaluate the basis functions and their derivatives.

Returns:

The basis functions and their first derivatives at the given points. The n_values depends on the number of dimensions of the element. For example, for a triangle in 2D, n_derivatives = 3 (1 for the function value and 2 for the derivatives in x and y).

Return type:

jnp.ndarray [n_values, n_points, n_basis]

compatible(geometric_domain: GeometricDomain)

Checks whether the element is compatible with the given geometric domain.

This can include checks such as whether the number of faces of the element matches the number of faces in the geometric domain, whether the number of vertices per face matches, etc.

Parameters:

geometric_domain (GeometricDomain) – The geometric domain to check compatibility with.

Return type:

bool

compute_geometric_data(geometric_domain)

Computes the Jacobians for the element given a geometric domain. Must be jittable/traceable: unlike precompute_geometric_data(), this must NOT call assert_compatible() (or otherwise touch geometric_domain in a non-traceable way), since it is also used to recompute geometric data under jit/grad when the domain vertices are perturbed (e.g. by a parametrisation reconstructing the operator after a geometry deformation).

This is required to be dependent on the element: a different coordinate interpretation (i.e. a different standard element) in the type of finite element leads to a different reference-element-to-element mapping, and thus a different Jacobian.

Parameters:

geometric_domain (GeometricDomain) – The geometric domain containing the mesh and adjacency information. Here, we only use geometric_domain.vertices and geometric_domain.connectivity.

Returns:

• PrecomputedGeometricData – The precomputed geometric data for all elements [does not depend on direction or energy group so is the same for all angles and energy groups]

local_coordinates(real_coordinates: Array, element_indices: Array, geometric_domain: GeometricDomain, precomputed_geometric_data: PrecomputedGeometricData)

Computes the local coordinates of the given points in the given elements.

Parameters:

• points (jnp.ndarray [n_points, n_dim]) – The global coordinates of the points to compute local coordinates for.

• element_indices (Array) – The indices of the elements containing the points.

• geometric_domain (GeometricDomain) – The geometric domain containing the mesh and adjacency information.

Returns:

The local coordinates of the points in the given elements.

Return type:

jnp.ndarray [n_points, n_dim]

real_coordinates(local_coordinates: Array, element_indices: Array, geometric_domain: GeometricDomain, precomputed_geometric_data: PrecomputedGeometricData)

Computes the real coordinates of the given local coordinates in the given elements.

Parameters:

• local_coordinates (Array) – The local coordinates of the points to compute real coordinates for.

• element_indices (Array) – The indices of the elements containing the points.

• geometric_domain (GeometricDomain) – The geometric domain containing the mesh and adjacency information.

Returns:

The real coordinates of the points in the given elements.

Return type:

jnp.ndarray [n_points, n_dim]

volume_quadrature_order(n_deg: int)

Returns the quadrature points and weights for the element that are sufficiently accurate for the volume integrals of polynomials up to degree n_deg.

Parameters:

n_deg (int) – The maximum polynomial degree to integrate exactly.

Return type:

Tuple[Array, Array]

Returns:

• points (jnp.ndarray [n_quadrature_points, n_dim]) – The quadrature points.

• weights (jnp.ndarray [n_quadrature_points]) – The quadrature weights.