jax_sn.domain.cross_sections module

Dataclasses for cross section data representation. It contains matrices for total and scattering cross sections, both for single and multiple energy groups.

\[\sigma_t\]

\[\sigma_{s,l}^{g' \to g}\]

class CrossSectionDataSingleGroup(total: Array, scattering: Array)

Bases: Module

Dataclass to hold cross section data for a single energy group. Batched over leading dimensions

total: Array

total cross section, batched arbitrarily over leading dimensions. Shape is […].

scattering: Array

scattering cross section moments, batched arbitrarily over leading dimensions. Shape is […, n_L]

class CrossSectionData(total: Array, scattering: Array)

Bases: Module

Dataclass to hold cross section data for multiple energy groups. Batched over leading dimensions.

total: Array

total cross section, batched arbitrarily over leading dimensions. Shape is […, n_energy_groups]

scattering: Array

scattering cross section moments, batched arbitrarily over leading dimensions. Shape is […, n_L, n_energy_groups, n_energy_groups]. Here, \(\sigma^{l}_{g' \leftarrow g}\) is stored at index […, l, g’, g] In other words, the first energy group index corresponds to the outgoing energy group, and the second to the incoming energy group.

classmethod zero(n_eg: int, n_L: int)

Generates a CrossSectionData object with zero total and scattering cross sections.

Parameters:

• n_eg (int) – Number of energy groups.

• n_L (int) – Number of scattering moments (L).

classmethod constant(total: float, n_eg: int, n_L: int)

Generates a CrossSectionData object with constant total cross section and zero scattering cross sections.

Parameters:

• total (float) – Value for the total cross section.

• n_eg (int) – Number of energy groups.

• n_L (int) – Number of scattering moments (L).

classmethod from_list(cross_section_data: List[CrossSectionData])

Constructs a batched CrossSectionData object from a list of CrossSectionData objects.

Parameters:

cross_section_data (List[CrossSectionData]) – List of CrossSectionData objects to batch together.

assert_valid()

Asserts that the cross section data is valid, checking dimensions and consistency.

assert_valid_batched()

Asserts that the batched cross section data is valid, checking dimensions and consistency of the leading batch dimensions additionally.

assert_n_energy_groups(n_energy_groups: int)

Asserts that the cross section data has at least n_energy_groups energy groups. :type n_energy_groups: int :param n_energy_groups: Number of energy groups to check against.

Return type:

bool

property downscattering_only: bool

convert_to_downscattering_only()

Converts the cross section data to a downscattering only cross section data by zeroing all upscattering cross sections.

Return type:

CrossSectionData

class CrossSectionsIndexed(indices: Array, cross_section_data: CrossSectionData)

Bases: Module

Dataclass to hold indexed cross section data for multiple energy groups. Each element has an index into the cross section data. This avoids repeating cross section data for each element.

indices: Array

Indices into the cross section data for each element. Shape is [n_elements]

cross_section_data: CrossSectionData

The cross section data for the indexed cross sections. Batched over n_cross_sections (where it should hold that n_cross_sections >= max(indices)+1])

map_to_full_array(energy_group: int)

Maps the indexed cross sections to a full array for the given energy group. :type energy_group: int :param energy_group: Energy group index to map to.

Return type:

CrossSectionDataSingleGroup

property n_energy_groups

Returns the number of energy groups in the cross section data.

property l_max

Returns the maximum scattering moment L in the cross section data.