jax_sn.quadrature_set

class QuadratureSetBase(angles: Array, weights: Array)

Bases: Module

Abstract base for quadrature sets. Holds the shared state – the angles and weights – and defines the spherical-harmonics interface that every quadrature set must provide: the unweighted and the weighted real spherical harmonics evaluated at the set’s angles.

The two harmonic methods are left abstract: the meaning of the angles/weights is shared, but how the (weighted) harmonics are computed from them is up to each concrete subclass (a reduced set, for instance, accumulates the weights over collapsed angles).

angles: Array

The angles of the quadrature set in Cartesian coordinates. Shape [n_angles, dimension].

weights: Array

The weights associated with each angle. Shape [n_angles]. Should sum to 4pi (such that \(\sum_{i} w_i 1 = \int d\Omega = 4 \pi\)).

property n_angles: int

abstractmethod wk_ylm_real(l_max: int)

Compute the weighted real spherical harmonic expansion coefficients \(\sqrt{w_k} Y_{lm}(\Omega_k)\) for this quadrature set up to degree l_max. See compute_3d_wk_yklm_real_from_quadrature_set() for details.

Return type:

Array

abstractmethod ylm_real(l_max: int)

Compute the real spherical harmonic expansion coefficients \(Y_{lm}(\Omega_k)\) for this quadrature set up to degree l_max. See compute_3d_yklm_real_from_quadrature_set() for details.

Return type:

Array

class QuadratureSet(angles: Array, weights: Array)

Bases: QuadratureSetBase

The standard concrete quadrature set: the stored weights are exactly the integration weights for the stored angles, so the (weighted) harmonics are evaluated directly on them. Additionally supports single-angle indexing via __getitem__.

wk_ylm_real(l_max: int)

Compute the weighted real spherical harmonic expansion coefficients \(\sqrt{w_k} Y_{lm}(\Omega_k)\) for this quadrature set up to degree l_max. See compute_3d_wk_yklm_real_from_quadrature_set() for details.

Return type:

Array

ylm_real(l_max: int)

Compute the real spherical harmonic expansion coefficients \(Y_{lm}(\Omega_k)\) for this quadrature set up to degree l_max. See compute_3d_yklm_real_from_quadrature_set() for details.

Return type:

Array

class QuadratureSetReduced(angles: Array, weights: Array, angles_original: Array, weights_original: Array, original_angles_indices: Array)

Bases: QuadratureSetBase

Represents a reduced quadrature set, where only unique angles are stored. The original quadrature set can be reconstructed from the unique angles and the indices to reconstruct the original set from the unique set.

Note that this does not subclass QuadratureSet: a reduced set cannot return a single angle/weight pair (the weights are accumulated over collapsed angles), so it does not support __getitem__. It only implements the spherical-harmonics interface of QuadratureSetBase.

angles_original: Array

weights_original: Array

original_angles_indices: Array

classmethod from_quadrature_set(quadrature_set: QuadratureSetBase, n_dim: Literal[1, 2, 3], rounding_tolerance: int = 10)

ylm_real(l_max: int)

Compute the real spherical harmonics \(Y_{lm}(\Omega_k)\) at the reduced angles. The harmonics are unweighted, so they are evaluated directly on the reduced angles.

Return type:

Array

wk_ylm_real(l_max: int)

Compute the weighted real spherical harmonic expansion coefficients \(\sqrt{w_k} Y_{lm}(\Omega_k)\) for this quadrature set up to degree l_max. See compute_3d_wk_yklm_real_from_quadrature_set() for details.

Return type:

Array

create_tn_quadrature_set(N: int)

Creates a TN quadrature set with N levels per octant. See [1].

Number of angles is:

\[8 N^2.\]

[1] Thurgood, C. P., Pollard, A., & Becker, H. A. (1995). The TN quadrature set for the discrete ordinates method.

Parameters:

N (int) – Number of levels per octant.

Return type:

QuadratureSet

Returns:

• total_angles (jnp.ndarray) – Array of shape [8 * N**2, 3] containing the discrete ordinate angles.

• total_weights (jnp.ndarray) – Array of shape [8 * N**2] containing the weights.

compute_3d_yklm_real_from_quadrature_set(quadrature_set: QuadratureSetBase, l_max: int)

Compute the real spherical harmonic expansion coefficients \(Y_{lm}\) for a given quadrature set up to degree l_max.

See _ylm_real for the definition of the real spherical harmonics used. Computes

\[Y_{lm}(\Omega_k)\]

for all angles \(\Omega_k\) in the quadrature set.

This uses a three-dimensional quadrature set, i.e., the angles are given as (x,y,z) coordinates on the unit sphere. For lower-dimensional quadrature sets, the symmetry on the unit sphere could be taken into account to reduce the number of angles.

Parameters:

• quadrature_set (QuadratureSetBase) – The quadrature set containing angles and weights.

• l_max (int) – Maximum degree of the spherical harmonics.

Return type:

Array

compute_3d_wk_yklm_real_from_quadrature_set(quadrature_set: QuadratureSetBase, l_max: int)

Compute the weighted real spherical harmonic expansion coefficients \(w_k Y_{lm}(\Omega_k)\) for a given quadrature set up to degree l_max.

See _ylm_real for the definition of the real spherical harmonics used.

Computes

\[w_k Y_{lm}(\Omega_k)\]

for all angles \(\Omega_k\) in the quadrature set with weights \(w_k\).

This uses a three-dimensional quadrature set, i.e., the angles are given as (x,y,z) coordinates on the unit sphere. For lower-dimensional quadrature sets, the symmetry on the unit sphere can be taken into account to reduce the number of angles.

Parameters:

• quadrature_set (QuadratureSetBase) – The quadrature set containing angles and weights.

• l_max (int) – Maximum degree of the spherical harmonics.

Return type:

Array

Submodules

• jax_sn.quadrature_set.quadrature_set_base module
  • QuadratureSetBase
  • QuadratureSet
  • QuadratureSetReduced
  • compute_3d_yklm_real_from_quadrature_set()
  • compute_3d_wk_yklm_real_from_quadrature_set()
• jax_sn.quadrature_set.tn module
  • create_tn_quadrature_set()