jax_sn.operators.transport_operator.implementations.transport_operator_forward module

class VmapForward[source]

Bases: SingleAngleForwardStrategy

Forward via vmap over all elements.

forward(data: TransportData, x)[source]
class BatchedForward(batch_size: int = None)[source]

Bases: SingleAngleForwardStrategy

Forward via batched vmap + fori_loop. batch_size=None reproduces the non-batched version.

batch_size: int = None
forward(data: TransportData, x)[source]
forward_transport_operator_vmap(geometric_domain: GeometricDomain, fem_matrices: FEMMatrices, precomputed_geometric_data: PrecomputedGeometricData, cross_sections: CrossSectionDataSingleGroup, angle: Array, x: Array)[source]

Forward transport operator using vmap over angles.

Does not require a topological sort so angle is sufficient.

Parameters:

  • geometric_domain (GeometricDomain) – Geometric domain (containing element to element connectivity and element to face connectivity).

  • fem_matrices (FEMMatrices) – Finite element method matrices.

  • precomputed_geometric_data (PrecomputedGeometricData) – Precomputed geometric data.

  • cross_sections (CrossSectionDataSingleGroup) – Cross sections for a single energy group

  • angle (Array) – Angle direction

  • x (Array) – x vector to apply the transport operator to.

Returns:

Result of applying the forward transport operator. Shape [n_angles, n_elements, n_basis]

Return type:

Array

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

Forward transport operator using batching over elements.

The result is checkpointed in the batching. This reduces memory significantly.

Parameters:

  • geometric_domain (GeometricDomain) – Geometric domain (containing element to element connectivity and element to face connectivity).

  • fem_matrices (FEMMatrices) – Finite element method matrices.

  • precomputed_geometric_data (PrecomputedGeometricData) – Precomputed geometric data.

  • cross_sections (CrossSectionDataSingleGroup) – Cross sections for a single energy group

  • angle (Array) – Angle direction.

  • x (Array) – x vector to apply the transport operator to.

Returns:

Result of applying the forward transport operator. Shape [n_angles, n_elements, n_basis]

Return type:

Array