jax_sn.solution_domain.base_fem module

class FEMMatrices(mass_matrix: Array, integration_coeff: Array, derivative_matrix: Array, outflow_matrices: Array, inflow_matrices: Array)

Bases: Module

Convenience class to bundle all FEM matrices required for the solution of the transport equation on a domain This is static wrt the domain geometry since it is only a function of the element type and basis functions.s

mass_matrix: Array

Mass matrix of the reference element:

\[H_{ij}= \int_R H_i H_j dV\]

The mass matrix of a non-reference element with linear Jacobian can be obtained by scaling this matrix with the determinant of the Jacobian:

\[H_{e,ij} = |J_e| H_{ij}\]

integration_coeff: Array

Integration coefficient vector of the reference element:

\[I_{i}= \int_R H_i dV\]

The integration coefficient vector of a non-reference element with linear Jacobian can be obtained by scaling this vector with the determinant of the Jacobian:

\[I_{e,i} = |J_e| I_{i}\]

derivative_matrix: Array

Derivative matrix of the reference element:

\[B_{\alpha, ij}= \int_R \frac{\partial H_i}{\partial r_\alpha} H_j dV\]

where \(r_\alpha\) is the reference coordinate in dimension \(\alpha\).

The derivative matrix of a non-reference element with linear Jacobian can be obtained by transforming this matrix with the inverse of the Jacobian:

\[B_{e,\beta, ij} = \sum_{\alpha} |J_e| (\mathbf{J}_e^{-1})_{\alpha, \beta} B_{\alpha, ij}\]

such that the required finite element integral in real coordinates can be computed as:

\[\int_e H_j \Omega \cdot \nabla H_i = \sum_{\beta} \Omega_{\beta} \sum_{\alpha} |J_e| (\mathbf{J}_e^{-1})_{\alpha, \beta} B_{\alpha, ij}\]

outflow_matrices: Array

Outflow face matrices of the reference element:

\[C_{\gamma ij}= \int_{\gamma} H_i H_j dS\]

The outflow face matrices of a non-reference element with linear Jacobian can be obtained by scaling this matrix with the determinant of the face Jacobian:

\[\Omega \cdot n_{\gamma} \int_\gamma H_i H_j dA = |J_{e,\gamma}| C_{\gamma ij} \Omega \cdot n_{\gamma}\]

inflow_matrices: Array

Inflow face connection matrices of the reference element:

\[D_{\gamma_1, p, \gamma_2, ij}= \int_{\gamma_1} H_i \overline{H}_j dS\]

where :p: indicates the permutation of the face \(\gamma_2\) to match face \(\gamma_1\).

The inflow face connection matrices of a non-reference element with linear Jacobian can be obtained by scaling this matrix with the determinant of the face Jacobian:

\[\Omega \cdot n_{\gamma_1} \int_{\gamma_1} H_i \overline{H}_j dA = |J_{e,\gamma_1}| D_{\gamma_1, p, \gamma_2, ij} \Omega \cdot n_{\gamma_1}\]

classmethod from_element(element: Element)

property n_basis: int

class PrecomputedGeometricData(vol_jacobians: Array, vol_jacobians_inv: Array, vol_jacobians_det: Array, face_jacobians_det: Array)

Bases: Module

Convenience class to bundle all precomputed geometric data required for the solution of the transport equation on a domain This is specific to a given domain geometry. Therefore, it is separate from FEMMatrices.

vol_jacobians: Array

Volume Jacobians for all elements. Computed as:

\[\mathbf{J}_{e,ij} = \frac{\partial x_i}{\partial r_j}\]

where \(\mathbf{x}\) are the real coordinates in the element \(e\), and \(\mathbf{r}\) are the reference element coordinates.

Shape [n_elements, n_dim, n_dim].

vol_jacobians_inv: Array

Inverse of the volume Jacobians for all elements. Shape [n_elements, n_dim, n_dim].

vol_jacobians_det: Array

Determinant of the volume Jacobians for all elements. Shape [n_elements].

face_jacobians_det: Array

Determinant of the face Jacobians for all elements and faces. Shape [n_elements, n_faces_per_element].

property n_elements: int

class Element

Bases: ABC

Base class for finite elements in jax-sn. This provides functions for computing the finite element matrices, as well as for mapping between local and real coordinates.

Methods to tabulate basis functions, get quadrature points and weights, and other element-specific operations are defined here as abstract methods to be implemented by concrete element classes. These do not necessarily need to be jittable, since they are typically only called during setup and precomputation. They can thus use non-jittable libraries such as basix. The only exception is the precompute_geometric_data method; this needs to be jittable and traceable to allow for derivative computations with respect to the geometry.

abstract property volume_quadrature: Tuple[Array, Array]

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)

abstract property has_linear_jacobian: bool

Returns whether the element has a linear Jacobian

abstract property face_quadratures: Tuple[Array, Array]

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

abstract property all_connection_face_quadratures: List[Tuple[Array, Array, Array]]

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).

abstract property n_faces: int

Returns the number of faces of the element.

abstract property n_basis: int

Returns the number of basis functions of the element.

abstract 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)].

abstractmethod n_vertices_per_face(face_idx: int)

Returns the number of vertices per face of the element.

Parameters:

face_idx (int) – The index of the face.

Return type:

int

abstractmethod tabulate_basis_functions_first_derivative(points: Array)

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

Parameters:

points (Array) – 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:

Array

abstractmethod 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

abstractmethod compute_geometric_data(geometric_domain: GeometricDomain)

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.

Return type:

PrecomputedGeometricData

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]

abstractmethod 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:

Array

abstractmethod 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]

abstractmethod 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.

assert_compatible(geometric_domain: GeometricDomain)

Asserts that the element is compatible with the given geometric domain.

Parameters:

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

precompute_geometric_data(geometric_domain: GeometricDomain)

Asserts that geometric_domain is compatible with this element, then computes its geometric data. NOT jittable (the compatibility check is not traceable) – this is meant for the initial, untraced setup of a SolutionDomain. If you need to recompute geometric data for an already-validated, possibly-traced domain (e.g. after perturbing vertices under jit/grad), call compute_geometric_data() directly instead.

Parameters:

geometric_domain (GeometricDomain) – The geometric domain to precompute geometric data for.

Returns:

The precomputed geometric data for all elements.

Return type:

PrecomputedGeometricData

interpolate_from_local_coordinate(local_coordinates: Array, element_indices: Array, values_array: Array)

Interpolates the given values at the given local coordinates in the given elements.

Parameters:

• local_coordinates (Array) – The local coordinates of the points to interpolate.

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

• values_array (Array) – The values of each basis function for each element.

Returns:

The interpolated values at the given local coordinates in the given elements.

Return type:

jnp.ndarray [n_points]

interpolate_from_real_coordinate(real_coordinates: Array, element_indices: Array, geometric_domain: GeometricDomain, precomputed_geometric_data: PrecomputedGeometricData, values_array: Array)

Interpolates the given values at the given real coordinates in the given elements.

Parameters:

• real_coordinates (Array) – The real coordinates of the points to interpolate.

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

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

• precomputed_geometric_data (PrecomputedGeometricData) – The precomputed geometric data for the elements.

• values_array (Array) – The values of each basis functions for each element.

Returns:

The interpolated values at the given real coordinates in the given elements.

Return type:

jnp.ndarray [n_points]

compute_volume_matrices_linear_jacobian(element: Element)

Computes the mass matrix, integration coefficient vector, and derivative matrix for an element with a linear Jacobian.

Parameters:

element (Element) – The finite element for which to compute the matrices.

Returns:

• mass_matrix (jnp.ndarray) – The mass matrix of shape [n_basis, n_basis].

• integration_coeff (jnp.ndarray) – The integration coefficient vector of shape [n_basis].

• derivative_matrix (jnp.ndarray) – The derivative matrix of shape [n_basis, n_basis].

compute_face_matrices_linear_jacobian(element: Element)

Computes the face inflow and outflow matrices for an element with a linear Jacobian.

Parameters:

element (Element) – The finite element for which to compute the face matrices.

Returns:

• face_outflow_matrices (jnp.ndarray) – The face outflow matrices of shape [n_faces, n_basis, n_basis].

• face_inflow_matrices (jnp.ndarray) – The face inflow matrices of shape [n_faces, n_basis, n_basis].