jax_sn.operators.transport_operator.implementations.transport_operator_pallas module¶

transport_inverse_sweep_batched_pallas(geometric_domain: GeometricDomain, fem_matrices: FEMMatrices, precomputed_geometric_data: PrecomputedGeometricData, cross_sections: CrossSectionDataSingleGroup, angle: Array, sweep_data: BasicSweepData, x: Array, batch_size: int)[source]¶

Perform a complete transport inverse sweep over all elements in the geometric domain using batched vmap and fori_loop in Pallas.

See transport_inverse_sweep_batched for more details. This is essentially the same functions but with padding to power-of-two sizes and using a Pallas kernel for the sweep update.

Parameters:

geometric_domain (GeometricDomain) – The geometric domain containing the mesh and element connectivity.
fem_matrices (FEMMatrices) – The finite element method matrices for the element type used in the domain.
precomputed_geometric_data (PrecomputedGeometricData) – Precomputed geometric data for the domain and element type.
cross_sections (CrossSectionDataSingleGroup) – The cross section data for the domain.
sweep_data (BasicSweepData) – The topologically sorted angle for the sweep.
x (Array) – The array to apply the sweep to.
batch_size (int) – The batch size for the batched sweep.

Returns:

The updated sweep result after performing the transport inverse sweep. Shape (n_elements, n_basis).

Return type:

jnp.ndarray

transport_inverse_sweep_complete_vmap_pallas(geometric_domain: GeometricDomain, fem_matrices: FEMMatrices, precomputed_geometric_data: PrecomputedGeometricData, cross_sections: CrossSectionDataSingleGroup, angle: Array, sweep_data: BasicSweepData, x: Array)[source]¶

Perform a complete transport inverse sweep over all elements in the geometric domain using batched vmap and fori_loop in Pallas.

See transport_inverse_sweep_complete_vmap for more details. This is essentially the same functions but with padding to power-of-two sizes and using a Pallas kernel for the sweep update.

Parameters:

geometric_domain (GeometricDomain) – The geometric domain containing the mesh and element connectivity.
fem_matrices (FEMMatrices) – The finite element method matrices for the element type used in the domain.
precomputed_geometric_data (PrecomputedGeometricData) – Precomputed geometric data for the domain and element type.
cross_sections (CrossSectionDataSingleGroup) – The cross section data for the domain.
sweep_data (BasicSweepData) – The topologically sorted angle for the sweep.
x (Array) – The array to apply the sweep to.
batch_size (int) – The batch size for the batched sweep.

Returns:

The updated sweep result after performing the transport inverse sweep. Shape (n_elements, n_basis).

Return type:

jnp.ndarray

sweep_pl_update(elements_toposort, gij_inv_all_ref, qm_i_all_ref, dij_f_p_fm_all_ref, etoe_ref, abs_omega_dot_n_Je_ref, sweep_result)[source]¶

Wrapper function to call the Pallas kernel for the sweep update.

Parameters:

elements_toposort (jnp.ndarray) – The indices of the elements in topological order for the sweep.
gij_inv_all_ref (jnp.ndarray) – The padded inverse of the gij matrices for all elements in the batch.
qm_i_all_ref (jnp.ndarray) – The padded source terms for all elements in the batch.
dij_f_p_fm_all_ref (jnp.ndarray) – The padded dij_f_p_fm matrices for all elements in the batch.
etoe_ref (jnp.ndarray) – The padded element-to-element connectivity array.
abs_omega_dot_n_Je_ref (jnp.ndarray) – The padded absolute values of the dot products of omega and the element face normals for all elements in the batch.
sweep_result (jnp.ndarray) – The buffer array to store the sweep results. Updated in place.

Returns:

The updated sweep result after performing the transport inverse sweep over the batch of elements.

Return type:

jnp.ndarray

transport_inverse_sweep_multiple_angles_batched_pallas(geometric_domain: GeometricDomain, fem_matrices: FEMMatrices, precomputed_geometric_data: PrecomputedGeometricData, cross_sections: CrossSectionDataSingleGroup, angle: Array, sweep_data: BasicSweepData, x: Array, batch_size: int)[source]¶

Perform a complete transport inverse sweep over all elements in the geometric domain for multiple angles using batched vmap and fori_loop.

vmap is used to compute all Gij_inv, QM_i and Dij_f_p_fm matrices for all elements in a vectorized manner for a fixed batch size and multiple angles. Then, a fori_loop is used to iterate over the elements in topological order to compute the sweep result in batches. The sweep kernel uses Pallas for improved performance.

Parameters:

geometric_domain (GeometricDomain) – The geometric domain data.
fem_matrices (FEMMatrices) – The finite element method matrices.
precomputed_geometric_data (PrecomputedGeometricData) – The precomputed geometric data (Jacobian determinants and inverses).
cross_sections (CrossSectionDataSingleGroup) – The cross section data.
sweep_data (BasicSweepData) – The topologically sorted angle data for multiple angles. (batch dimension is the first dimension)
x (Array) – The source data. Should also be batched over angles (batch dimension is the first dimension).
batch_size (int) – The batch size for processing elements.

Returns: