jax_sn.solution_domain.solution_domain module

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

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]