jax_sn.operators.transport_operator.implementations.transport_operator_jax module

Contains pure JAX implementations of the single-angle transport operator sweep functions. Both pure VMAP (vmap single-element inversion) and batched (batched single-element inversion).

class VmapSweep[source]

Bases: SingleAngleInversionStrategy

Inverse sweep with all per-element matrices computed via vmap (high memory).

invert(data: TransportData, b)[source]
classmethod vmappable()[source]

Whether an outer vmap over this inversion vectorises.

class BatchedSweep(batch_size: int = None)[source]

Bases: SingleAngleInversionStrategy

Inverse sweep over batches of elements. batch_size=None reproduces the non-batched version.

batch_size: int = None
invert(data: TransportData, b)[source]
classmethod vmappable()[source]

Whether an outer vmap over this inversion vectorises.

transport_inverse_sweep_batched(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.

vmap is used to compute all Gij_inv, and Dij_f_p_fm matrices for all elements in a vectorized manner for a fixed batch size. Then, a fori_loop is used to iterate over the elements in topological order to compute the sweep result in batches. This takes significantly less memory than using the complete vmap version.

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.

  • x (Array) – The x vector to apply the transport inverse sweep to.

  • batch_size (int) – The batch size to use for the batched vmap computation.

Returns:

sweep_result – The computed sweep result array after performing the transport sweep.

Return type:

jnp.ndarray [n_elements, n_basis]

transport_inverse_sweep_complete_vmap(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 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. Then, a fori_loop is used to iterate over the elements in topological order to compute the sweep result. This takes significant memory due to the large intermediate matrices.

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.

  • x (Array) – The x vector to apply the transport inverse sweep to.

Returns:

sweep_result – The computed sweep result array after performing the transport sweep.

Return type:

Array

gij_kernel_idx_in_precomp(elem_idx: int, fem_matrices: FEMMatrices, precomputed_geometric_data: PrecomputedGeometricData, cross_sections: CrossSectionDataSingleGroup, angle: Array, inflow: Array, abs_omega_dot_n_Je: Array)[source]
Internal function to compute the Gij matrix for a single element index:

Gij = σ_t Jv ∫H_i H_j dV + Jv ∫H_i (Ω·∇H_j) dV + ∑_{f} |Ω·n_f| Je_f ∫H_i H_j dA

Parameters:

  • elem_idx (int) – The element index to compute the Gij matrix for.

  • fem_matrices (FEMMatrices) – The finite element method matrices.

  • precomputed_geometric_data (PrecomputedGeometricData) – The precomputed geometric data (Jacobian determinants and inverses). Since these are already precomputed, we do not pass a matrix that has ordered data for this specific direction. This would be significantly more memory intensive.

  • angle (Array) – The direction vector.

  • inflow (Array) – Boolean array indicating inflow faces for the specified element. This has the advantage of allowing us to batch it by providing it for each element in the batch only (instead of all elements). Since we have to compute it separately for each direction anyway, this does not involve additional copying (as it would be for the precomputed jacobian data).

  • abs_omega_dot_n_Je (Array) – The absolute value of the dot product of omega and face normals multiplied by face Jacobian determinants for the specified element. Similarly to inflow, we pass only for this specific element here.

  • cross_sections (CrossSectionDataSingleGroup) – The cross section data. total attribute must have [n_elements] shape (indexed in). This is used to get σ_t for the specific element. This is not precomputed per direction to save memory.

  • Returns

gij_kernel_idx_in_precomp_vmap(*args, **kwargs)
gij_inv_kernel_idx_in_precomp(elem_idx: int, fem_matrices: FEMMatrices, precomputed_geometric_data: PrecomputedGeometricData, cross_sections: CrossSectionDataSingleGroup, angle: Array, inflow: Array, abs_omega_dot_n_Je: Array)[source]
Compute the inverse of the Gij matrix for a single element index.

Gij = σ_t Jv ∫H_i H_j dV + Jv ∫H_i (Ω·∇H_j) dV + ∑_{f} |Ω·n_f| Je_f ∫H_i H_j dA

See also gij_kernel_idx_in_precomp for explanation of the parameters. :returns: gij_inv – The computed inverse Gij matrix for the specified element. :rtype: jnp.ndarray [n_basis, n_basis]

gij_inv_kernel_idx_in_precomp_vmap(*args, **kwargs)
dij_f_p_fm_kernel_idx_in_precomp(elem_idx: int, fem_matrices: FEMMatrices, inflow: Array, etof: Array, face_perm_idx: Array)[source]
Compute the Dij_f_p_fm matrices for a single element index:

Dij_f_p_fm = |Ω·n_f| J_f ∫ H_i H_j dA

See gij_kernel_idx_in_precomp for explanation on why we pass inflow as a single array instead of an indexed array.

Parameters:

  • elem_idx (int) – The element index to compute the Dij_f_p_fm matrices for.

  • fem_matrices (FEMMatrices) – The finite element method matrices.

  • inflow (Array) – Boolean array indicating inflow faces for the specified element.

  • etof (Array) – The element-to-face connectivity array.

  • face_perm_idx (Array) – The face permutation index array.

Returns:

dij_f_p_fm – The computed Dij_f_p_fm matrices for the specified element.

Return type:

jnp.ndarray

dij_f_p_fm_kernel_idx_in_precomp_vmap(elem_idx: int, fem_matrices: FEMMatrices, inflow: Array, etof: Array, face_perm_idx: Array)

Vectorized version of dij_f_p_fm_kernel_idx_in_precomp. Takes similar arguments as dij_f_p_fm_kernel_idx_in_precomp but with additional array axes over which dij_f_p_fm_kernel_idx_in_precomp is mapped.

Original documentation:

Compute the Dij_f_p_fm matrices for a single element index:

Dij_f_p_fm = |Ω·n_f| J_f ∫ H_i H_j dA

See gij_kernel_idx_in_precomp for explanation on why we pass inflow as a single array instead of an indexed array.

elem_idxint

The element index to compute the Dij_f_p_fm matrices for.

fem_matricesFEMMatrices

The finite element method matrices.

inflowjnp.ndarray [n_faces_per_element]

Boolean array indicating inflow faces for the specified element.

etofjnp.ndarray [n_elements, n_faces_per_element]

The element-to-face connectivity array.

face_perm_idxjnp.ndarray [n_elements, n_faces_per_element]

The face permutation index array.

dij_f_p_fmjnp.ndarray

The computed Dij_f_p_fm matrices for the specified element.