jax_sn.domain

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

class GeometricDomain(vertices: Array, connectivity: Array, face_normals: Array, etoe: Array, etof: Array, face_perm_idx: Array, face_template: Array)

Bases: Module

Module representing a geometric domain. These contain precomputed connectivity, face normals, and adjacency information.

vertices: Array

Vertices of the domain. Shape is [n_vertices, n_dim]

connectivity: Array

Connectivity of the domain. Shape is [n_elements, n_vertices_per_element]

face_normals: Array

Face normals of the domain. Shape is [n_elements, n_faces_per_element, n_dim]

etoe: Array

Element-to-element adjacency. Shape is [n_elements, n_faces_per_element]. -1 indicates a boundary face (unconnected to another element).

etof: Array

Element-to-face adjacency. Shape is [n_elements, n_faces_per_element]. Each face has a unique index in the domain.

face_perm_idx: Array

Face permutation indices. Shape is [n_elements, n_faces_per_element]. Indicates how the face vertices are permuted between adjacent elements.

face_template: Array

Face template used to define faces of elements. Shape is [n_faces_per_element, n_vertices_per_face].

classmethod from_mesh(mesh: Tuple[Array, Array], face_template: Array)

Constructs a GeometricDomain from a mesh and a face template (see jax_sn.tools.mesh_tools._face_template for an example template).

Parameters:

• mesh (Tuple[Array, Array]) – A tuple containing vertices and connectivity of the mesh. vertices is a jnp.ndarray of shape [n_vertices, n_dim] and connectivity is a jnp.ndarray of shape [n_elements, n_vertices_per_element].

• face_template (Array) – A 2D array representing the face template.

Returns:

domain – The constructed GeometricDomain.

Return type:

GeometricDomain

classmethod from_mesh_numpy(mesh: Tuple[Array, Array], face_template: Array)

Constructs a GeometricDomain from a mesh and a face template (see jax_sn.tools.mesh_tools._face_template for an example template).

Uses a numpy implementation of the domain setup. This avoids compilation overheads of especially the lexsort, which can take significant time.

Parameters:

• mesh (Tuple[Array, Array]) – A tuple containing vertices and connectivity of the mesh. vertices is a jnp.ndarray of shape [n_vertices, n_dim] and connectivity is a jnp.ndarray of shape [n_elements, n_vertices_per_element].

• face_template (Array) – A 2D array representing the face template.

Returns:

domain – The constructed GeometricDomain.

Return type:

GeometricDomain

property n_elements

Number of elements in the domain

property n_faces_per_element

Number of faces per element in the domain

containing_cell(points: Array)

Containing cell for each point in points. Uses a cached BVH, so cannot be used in a jitted function

Parameters:

points (Array) – The points for which to find the containing cell.

Returns:

element_indices – The indices of the containing cell for each point. -1 if no containing cell is found

Return type:

jnp.ndarray [n_points]

class Domain(vertices: Array, connectivity: Array, face_normals: Array, etoe: Array, etof: Array, face_perm_idx: Array, face_template: Array, cross_sections_indexed: CrossSectionsIndexed)

Bases: GeometricDomain

Geometric domain with indexed cross sections. This is seperated from GeometricDomain because many algorithms do not depend on cross sections and therefore we want to keep the cross sections distinct.

cross_sections_indexed: CrossSectionsIndexed

Indexed cross sections for the domain.

classmethod from_mesh_and_cross_sections(mesh: Tuple[Array, Array], face_template: Array, cross_sections_indexed: CrossSectionsIndexed)

Constructs a Domain from a mesh, face template, and indexed cross sections. Checks compatibility between the mesh and cross sections.

Parameters:

• mesh (Tuple[Array, Array]) – A tuple containing vertices and connectivity of the mesh. vertices is a jnp.ndarray of shape [n_vertices, n_dim] and connectivity is a jnp.ndarray of shape [n_elements, n_vertices_per_element].

• face_template (Array) – A 2D array representing the face template.

• cross_sections_indexed (CrossSectionsIndexed) – Indexed cross sections for the domain.

Returns:

domain – The constructed Domain.

Return type:

Domain

classmethod from_mesh_and_cross_sections_numpy(mesh: Tuple[Array, Array], face_template: Array, cross_sections_indexed: CrossSectionsIndexed)

Constructs a Domain from a mesh, face template, and indexed cross sections. Checks compatibility between the mesh and cross sections. Uses a numpy implementation of the domain setup. This avoids compilation overheads of especially the lexsort, which can take significant time.

Parameters:

• mesh (Tuple[Array, Array]) – A tuple containing vertices and connectivity of the mesh. vertices is a jnp.ndarray of shape [n_vertices, n_dim] and connectivity is a jnp.ndarray of shape [n_elements, n_vertices_per_element].

• face_template (Array) – A 2D array representing the face template.

• cross_sections_indexed (CrossSectionsIndexed) – Indexed cross sections for the domain.

Returns:

domain – The constructed Domain.

Return type:

Domain

classmethod from_geometric_domain_and_cross_sections(geometric_domain: GeometricDomain, cross_sections_indexed: CrossSectionsIndexed)

Constructs a Domain from a GeometricDomain and indexed cross sections. Checks compatibility between the geometric domain and cross sections.

Parameters:

• geometric_domain (GeometricDomain) – The geometric domain.

• cross_sections_indexed (CrossSectionsIndexed) – Indexed cross sections for the domain.

Returns:

domain – The constructed Domain.

Return type:

Domain

Submodules

• jax_sn.domain.cross_sections module
  • CrossSectionDataSingleGroup
  • CrossSectionData
  • CrossSectionsIndexed
• jax_sn.domain.domain_generation module
  • create_cartesian_setup()
• jax_sn.domain.geometric_domain module
  • compute_face_normals()
  • compute_topology()
  • GeometricDomain
  • Domain
• jax_sn.domain.graph_generation module
  • compute_face_inflow()
  • compute_levels_domain()
  • compute_levels_domain_safe()
  • compute_levels_domain_safe_jax()