jax_sn.solution_domain

class Element[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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]

class PrecomputedGeometricData(vol_jacobians: Array, vol_jacobians_inv: Array, vol_jacobians_det: Array, face_jacobians_det: Array)[source]

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 FEMMatrices(mass_matrix: Array, integration_coeff: Array, derivative_matrix: Array, outflow_matrices: Array, inflow_matrices: Array)[source]

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)[source]
property n_basis: int
class SolutionDomain(spatial_data: SpatialData)[source]

Bases: Module

Convenience class to bundle all data required to solve the transport equation on a domain Is not fully JIT-compatible due to the Element class. Its methods are also available as standalone functions.

spatial_data: SpatialData
classmethod from_element_and_domain(element: Element, domain: Domain)[source]
Return type:

SolutionDomain

property element: Element
interpolate(points: Array, solution: Array)[source]
property domain: Domain
property precomputed_geometric_data: PrecomputedGeometricData
property fem_matrices: FEMMatrices
property n_basis: int
property n_elements: int
property dimension: int
property elements_basis_shape: ShapeDtypeStruct
map_cross_sections_to_full_array(energy_group: int)[source]
class SpatialData(domain: Domain, fem_matrices: FEMMatrices, precomputed_geometric_data: PrecomputedGeometricData, element: Element)[source]

Bases: Module

Bundles all data required to solve the transport equation on a domain in a JIT-compatible way.

domain: Domain

The geometric domain on which to solve the transport equation.

fem_matrices: FEMMatrices

The finite element matrices for the domain.

precomputed_geometric_data: PrecomputedGeometricData

The precomputed geometric data for the domain.

element: Element

The finite element used to compute fem_matrices/precomputed_geometric_data. Static: not JIT-compatible (can involve e.g. a non-jax Basix element), and never referenced by the differentiable operator machinery – kept only so a SolutionDomain can be reconstructed (e.g. for interpolation) from data that has passed through operator creation.

property n_basis: int
property n_elements: int
property elements_basis_shape: ShapeDtypeStruct

The [n_elements, n_basis] structure operators batch over (identified with \(F\)).

compute_setup_for_solution_domain(element: Element, domain: GeometricDomain)[source]

Computes the FEM matrices and precomputed geometric data for the given element and domain.

Parameters:
Return type:

Tuple[FEMMatrices, PrecomputedGeometricData]

Returns:

  • fem_matrices (FEMMatrices) – The finite element matrices for the given element.

  • precomputed_data (PrecomputedGeometricData) – The precomputed geometric data for the given element and domain.

class BasixLagrangianSimplex(degree: int, dimension: int)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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.

Submodules