Serval

Main Library

Core.py

TripleSHIntegrator

A class for performing integrals of the products of three functions on the sphere expressed as their spherical harmonic transformations.

This class provides methods for this integration performed with a variety of techniques using Gaunt coefficients and grid-based approaches. It supports caching for efficient repeated computations when changing only one of the functions. Additionally it allows for computing the Fourier transform of these integrals in an azimuthal rotation angle of one of the functions (always the 1st) around the others, ie. the \(m\)-modes.

That is,

\[ \begin{align} \mathcal{V} &= \int\diff\Omega~f_1(\dirvec)f_2(\dirvec)f_3(\dirvec) \\ &= \sum_{\substack{\ell_1 \ell_2 \ell_3 \\ m_1 m_2 m_3}} a_{\ell_1 m_1}a_{\ell_2 m_2}a_{\ell_3 m_3} \int\diff\Omega~Y_{\ell_1 m_1}(\dirvec) Y_{\ell_2 m_2}(\dirvec) Y_{\ell_3 m_3}(\dirvec), \end{align} \]

where

\[ f_1(\dirvec) = \sum_{\ell_1=0}^{l_{1,\mathrm{max}}} \sum_{m_1=-\ell_1}^{\ell_1} a_{\ell_1 m_1} Y_{\ell_1 m_1}(\dirvec) \]

etc.

For the \(m\)-modes:

\[ \begin{align} \mathcal{V}_{m_1} &= \sum_{\substack{\ell_1 \ell_2 \ell_3 \\ m_2 m_3}} a_{\ell_1 m_1}a_{\ell_2 m_2}a_{\ell_3 m_3} \int\diff\Omega~Y_{\ell_1 m_1}(\dirvec) Y_{\ell_2 m_2}(\dirvec) Y_{\ell_3 m_3}(\dirvec), \\ &= \int \frac{\diff\phi}{2\pi} \exp[-im_1\phi] \int\diff\Omega~f_1( \matr{\mathcal{R}}_Z(-\phi)\dirvec)f_2(\dirvec)f_3(\dirvec) \end{align} \]

where \(\matr{\mathcal{R}}_Z(-\phi)\) is a passive basis rotation about the polar axis of the sphere.

Notes

The index of the functions only matters in differentiating from the others the first function/set of coefficients as the integration can be split up in \(m_1\) values. Additionally integrals can be performed rotating this function azimuthally around the others. That is, for the purposes of this code, this represents the sky.

Parameters:

Name	Type	Description	Default
l1max	int	Maximum spherical harmonic degree for the first set of coefficients.	required
l2max	int	Maximum spherical harmonic degree for the second set of coefficients.	required
l3max	int	Maximum spherical harmonic degree for the third set of coefficients.	required
absm1_limits	tuple[int, int | None]	Inclusive lower and exclusive upper limits for the absolute value of the m1 index. Defaults to (0, None), ie, all m1 values.	required
generate_cache_on_init	bool	Whether to generate the Gaunt coefficient cache during initialization, by default False.	required

Attributes:

Name	Type	Description
l1max	int	Maximum spherical harmonic degree for the first set of coefficients.
l2max	int	Maximum spherical harmonic degree for the second set of coefficients.
l3max	int	Maximum spherical harmonic degree for the third set of coefficients.
absm1_limits	tuple[int, int | None]	Inclusive lower and exclusive upper limits for the absolute value of the m1 index.
generate_cache_on_init	bool	Whether to generate the Gaunt coefficient cache during initialization.
m1_values	list[int]	List of m1 values that integrations are performed for.
m1_index_map	dict[int, int]	Mapping of m1 values to their indices in a completed set of m1 values.

Methods:

Name	Description
generate_gaunt_cache	Generates the Gaunt coefficient cache.
clear_gaunt_cache	Clears the Gaunt coefficient cache.
sparse_gaunt_einsum	Performs sparse tensor contraction using computed Gaunt coefficients.
grid_integrate	Performs integration using a grid-based approach.
gaunt_integrate	Performs integration using a direct sum over precomputed Gaunt coefficients.
linear_gaunt_integrator_12	Computes a linear integrator for fixed first and second sets of coefficients.
generate_integrator_cache_12	batch_parallel_mode="channel") Generates and caches the linear integrator for fixed first and second sets of coefficients.
batch_gaunt_integrate_cached_12	Uses a integrator cache for fixed first and second sets of coefficients to compute integrals over a batch of the third set of coefficients.
linear_gaunt_integrator_23	Computes a linear integrator for fixed second and third sets of coefficients.
generate_integrator_cache_23	Generates and caches the linear integrator for fixed second and third sets of coefficients.
batch_gaunt_integrate_cached_12	Uses a integrator cache for fixed second and third sets of coefficients to compute integrals over a bacth of the first set of coefficients.

Source code in src/serval/core.py

@define(eq=False)
class TripleSHIntegrator:
    """
    A class for performing integrals of the products of three functions on the
    sphere expressed as their spherical harmonic transformations.

    This class provides methods for this integration performed with a variety
    of techniques using Gaunt coefficients and grid-based approaches. It
    supports caching for efficient repeated computations when changing only
    one of the functions. Additionally it allows for computing the Fourier
    transform of these integrals in an azimuthal rotation angle of one of
    the functions (always the 1st) around the others, ie. the $m$-modes.

    That is,

    $$
    \\begin{align}
    \mathcal{V} &= \int\diff\Omega~f_1(\dirvec)f_2(\dirvec)f_3(\dirvec) \\\\
                &= \sum_{\substack{\ell_1 \ell_2 \ell_3 \\\\ m_1 m_2 m_3}}
                  a_{\ell_1 m_1}a_{\ell_2 m_2}a_{\ell_3 m_3}
                  \int\diff\Omega~Y_{\ell_1 m_1}(\dirvec)
                  Y_{\ell_2 m_2}(\dirvec) Y_{\ell_3 m_3}(\dirvec),
    \end{align}
    $$

    where

    $$
    f_1(\dirvec) = \sum_{\ell_1=0}^{l_{1,\mathrm{max}}}
                   \sum_{m_1=-\ell_1}^{\ell_1} a_{\ell_1 m_1} Y_{\ell_1 m_1}(\dirvec)
    $$

    etc.

    For the $m$-modes:

    $$
    \\begin{align}
    \mathcal{V}_{m_1} &= \sum_{\substack{\ell_1 \ell_2 \ell_3 \\\\ m_2 m_3}}
                        a_{\ell_1 m_1}a_{\ell_2 m_2}a_{\ell_3 m_3}
                        \int\diff\Omega~Y_{\ell_1 m_1}(\dirvec)
                        Y_{\ell_2 m_2}(\dirvec) Y_{\ell_3 m_3}(\dirvec), \\\\
                      &=
                        \int \\frac{\diff\phi}{2\pi} \exp[-im_1\phi] \int\diff\Omega~f_1(
                        \matr{\mathcal{R}}_Z(-\phi)\dirvec)f_2(\dirvec)f_3(\dirvec)
    \end{align}
    $$

    where $\matr{\mathcal{R}}_Z(-\phi)$ is a passive basis rotation about the polar axis
    of the sphere.

    Notes
    -----
    The index of the functions only matters in differentiating from the others the first
    function/set of coefficients as the integration can be split up in $m_1$ values.
    Additionally integrals can be performed rotating this function azimuthally around
    the others. That is, for the purposes of this code, this represents the sky.


    Parameters
    ----------
    l1max : int
        Maximum spherical harmonic degree for the first set of coefficients.
    l2max : int
        Maximum spherical harmonic degree for the second set of coefficients.
    l3max : int
        Maximum spherical harmonic degree for the third set of coefficients.
    absm1_limits : tuple[int, int | None], optional
        Inclusive lower and exclusive upper limits for the absolute value of
        the m1 index. Defaults to (0, None), ie, all m1 values.
    generate_cache_on_init : bool, optional
        Whether to generate the Gaunt coefficient cache during initialization,
        by default False.

    Attributes
    ----------
    l1max : int
        Maximum spherical harmonic degree for the first set of coefficients.
    l2max : int
        Maximum spherical harmonic degree for the second set of coefficients.
    l3max : int
        Maximum spherical harmonic degree for the third set of coefficients.
    absm1_limits : tuple[int, int | None]
        Inclusive lower and exclusive upper limits for the absolute value of
        the m1 index.
    generate_cache_on_init : bool
        Whether to generate the Gaunt coefficient cache during initialization.
    m1_values : list[int]
        List of m1 values that integrations are performed for.
    m1_index_map : dict[int, int]
        Mapping of m1 values to their indices in a completed set of m1 values.

    Methods
    -------
    generate_gaunt_cache()
        Generates the Gaunt coefficient cache.
    clear_gaunt_cache()
        Clears the Gaunt coefficient cache.
    sparse_gaunt_einsum(sparse_alm1, sparse_alm2, sparse_alm3, sum_m1=False, gaunts=None)
        Performs sparse tensor contraction using computed Gaunt coefficients.
    grid_integrate(alm1, alm2, alm3, sum_m1=False)
        Performs integration using a grid-based approach.
    gaunt_integrate(alm1, alm2, alm3, sum_m1=False)
        Performs integration using a direct sum over precomputed Gaunt coefficients.
    linear_gaunt_integrator_12(alm1, alm2, contract3=None, sum_m1=False)
        Computes a linear integrator for fixed first and second sets of coefficients.
    generate_integrator_cache_12(alm1, alm2, contract3=None, release_gaunt_cache=True,
                                 batch_parallel_mode="channel")
        Generates and caches the linear integrator for fixed first and second sets of
        coefficients.
    batch_gaunt_integrate_cached_12(alm3, integrator_cache=None, sum_m1=False)
        Uses a integrator cache for fixed first and second sets of coefficients to
        compute integrals over a batch of the third set of coefficients.
    linear_gaunt_integrator_23(alm2, alm3, sum_m1=False)
        Computes a linear integrator for fixed second and third sets of coefficients.
    generate_integrator_cache_23(alm2, alm3)
        Generates and caches the linear integrator for fixed second and third sets of
        coefficients.
    batch_gaunt_integrate_cached_12(alm1, integrator_cache=None, sum_m1=False)
        Uses a integrator cache for fixed second and third sets of coefficients to
        compute integrals over a bacth of the first set of coefficients.
    """

    l1max: int
    l2max: int
    l3max: int

    absm1_limits: tuple[int, int | None] = (0, None)
    generate_cache_on_init: bool = False

    linear_integrator_cache_12: npt.NDArray[np.complex128] | None = field(
        init=False, default=None
    )
    linear_integrator_cache_23: npt.NDArray[np.complex128] | None = field(
        init=False, default=None
    )
    gaunt_cache: CachingGauntContractor | None = field(init=False, default=None)
    m1_values: list[int] = field(init=False)
    m1_index_map: dict[int, int] = field(init=False)
    m1_global_index: list[int] = field(init=False)

    def __attrs_post_init__(self):
        if self.absm1_limits[1] is None:
            self.absm1_limits = (self.absm1_limits[0], self.l1max + 1)
        if self.generate_cache_on_init:
            self.generate_gaunt_cache()
        self.m1_values, self.m1_index_map = m1_indexing(
            self.l1max, self.absm1_limits[0], self.absm1_limits[1]
        )
        self.m1_global_index = list(self.m1_index_map.keys())

    def generate_gaunt_cache(self):
        """Generates the Gaunt coefficient cache."""
        self.gaunt_cache = CachingGauntContractor(
            self.l1max,
            self.l2max,
            self.l3max,
            absm1_lower=self.absm1_limits[0],
            absm1_upper=self.absm1_limits[1],
        )

    def clear_gaunt_cache(self):
        """Clears the Gaunt coefficient cache."""
        self.gaunt_cache = None

    def sparse_gaunt_einsum(
        self,
        sparse_alm1: sparse.SparseArray,
        sparse_alm2: sparse.SparseArray,
        sparse_alm3: sparse.SparseArray,
        sum_m1: bool = False,
        gaunts: sparse.COO | None = None,
    ) -> sparse.COO:
        """
        Perform the triple integral using a sparse tensor contraction.

        Compute the triple integral by generating or using a sparse
        representation of the Gaunt coefficients and performing the sum
        with a sparse einsum. This is mostly a cross-checking function
        and is not performant.

        Parameters
        ----------
        sparse_alm1 : sparse.SparseArray
            Sparse represenation of the first set of coefficients with shape (l1max+1, Nm1).
        sparse_alm2 : sparse.SparseArray
            Sparse represenation of the second set of coefficients with shape (l2max+1, 2*l2max+1).
        sparse_alm3 : sparse.SparseArray
            Sparse represenation of the third set of coefficients (l3max+1, 2*l3max+1).
        sum_m1 : bool, optional
            Whether to sum over the m1's before output, by default False.
        gaunts : sparse.COO | None, optional
            Sparse representation of the Gaunt coefficients if already.
            computed, by default None.

        Returns
        -------
        npt.NDArray[np.complex128] | np.complex128
            If sum_m1 = False, Nm1 array of the m-modes, otherwise, scalar
            sum over all m-modes.
        """
        if gaunts is None:
            gaunts = gaunts_coo(self.l1max, self.l2max, self.l3max)
        if sum_m1:
            return sparse.einsum(
                "ijklmn, ij, kl, mn -> ", gaunts, sparse_alm1, sparse_alm2, sparse_alm3
            )
        else:
            return sparse.to_numpy(
                sparse.einsum(
                    "ijklmn, ij, kl, mn -> j",
                    gaunts,
                    sparse_alm1,
                    sparse_alm2,
                    sparse_alm3,
                )
            )

    def grid_integrate(
        self,
        alm1: npt.NDArray[np.complex128],
        alm2: npt.NDArray[np.complex128],
        alm3: npt.NDArray[np.complex128],
        sum_m1: bool = False,
    ) -> npt.NDArray[np.complex128]:
        """Perform the triple integral using a grid-based approach.

        Here the product of the second and third fields are evaluated on a
        consistent resolution grid, decomposed to conjugate spherical
        harmonics and multipled by the first set of coefficients. This more
        standard m-mode approach is performant for single integrals
        but, as implemented here, can't cache intermediate results and
        therefore not suitable for batching.


        Parameters
        ---------
        alm1 : npt.NDArray[np.complex128]
            The first set of spherical harmonic coefficients with shape (l1max+1, Nm1).
        alm2 : npt.NDArray[np.complex128]
            The second set of spherical harmonic coefficients with shape (l2max+1, 2*l2max+1).
        alm3 : npt.NDArray[np.complex128]
            The third set of spherical harmonic coefficients with shape (l3max+1, 2*l3max+1).
        sum_m1 : bool, optional
            Whether to sum over the m1's before output, by default False.

        Returns
        -------
        npt.NDArray[np.complex128] | np.complex128
            If sum_m1 = False, Nm1 array of the m-modes, otherwise, scalar
            sum over all m-modes.
        """
        grid_lmax = int(np.ceil(self.l1max + self.l2max + self.l3max))
        alm1, alm2, alm3 = broadcast_bandlimits(alm1, alm2, alm3, lmax=grid_lmax)
        static_grid = grid_template(grid_lmax)
        static_grid.data[:] = 1.0
        static_grid *= array_synthesis(alm2) * array_synthesis(alm3)
        static_alm = analysis(static_grid, conjugate=True, return_sparse=False)
        out_mmodes = compute_mmodes(static_alm, alm1)
        return out_mmodes[np.asarray(self.m1_values) + grid_lmax]

    def gaunt_integrate(
        self,
        alm1: npt.NDArray[np.complex128],
        alm2: npt.NDArray[np.complex128],
        alm3: npt.NDArray[np.complex128],
        sum_m1: bool = False,
    ) -> npt.NDArray[np.complex128] | float:
        """Perform the triple integral by doing and inplace
        sum-product over gaunt coefficients.

        $$
        \mathcal{V} = \sum_{\substack{\ell_1 \ell_2 \ell_3 \\\\ m_1 m_2 m_3}}
                    \mathcal{G}^{~~\ell_1~~\ell_2~~\ell_3}_{m_1 m_2 m_3}~
                    a_{\ell_1 m_1}a_{\ell_2 m_2}a_{\ell_3 m_3}.
        $$

        Here the sum is performed inplace with no caching.

        Parameters
        ----------
        alm1 : npt.NDArray[np.complex128]
            The first set of spherical harmonic coefficients with shape (l1max+1, Nm1).
        alm2 : npt.NDArray[np.complex128]
            The second set of spherical harmonic coefficients with shape (l2max+1, 2*l2max+1).
        alm3 : npt.NDArray[np.complex128]
            The third set of spherical harmonic coefficients with shape (l3max+1, 2*l3max+1).
        sum_m1 : bool, optional
            Whether to sum over the m1's before output, by default False.

        Returns
        -------
        npt.NDArray[np.complex128] | np.complex128
            If sum_m1 = False, Nm1 array of the m-modes, otherwise, scalar
            sum over all m-modes.
        """
        return gaunt_dot123(
            alm1=alm1,
            alm2=alm2,
            alm3=alm3,
            sum_m1=sum_m1,
            absm1_lower=self.absm1_limits[0],
            absm1_upper=self.absm1_limits[1],
        )

    def linear_gaunt_integrator_12(
        self,
        alm1: npt.NDArray[np.complex128],
        alm2: npt.NDArray[np.complex128],
        contract3: npt.NDArray[np.complex128] | None = None,
        sum_m1: bool = False,
    ) -> npt.NDArray[np.complex128]:
        """Generate a linear operator that, with fixed first and second set of coefficients
        performs the integral when sum-producted with the third set of coefficients.

        $$
        \matr{L}_{m_1}^{\ell_3 m_3} =
                    \sum_{\substack{\ell_1 \ell_2 \\\\ m_1 m_2}}
                    \mathcal{G}^{~~\ell_1~~\ell_2~~\ell_3}_{m_1 m_2 m_3}~
                    a_{\ell_1 m_1}a_{\ell_2 m_2}.
        $$

        If requested, additionally contract with a linear operator along the
        m3 axis (e.g. a rotation with a Wigner-d matrix),

        $$
        \matr{L}_{m_1}^{\ell_3 m^\prime_3} = \sum_{m_3} \matr{W}^{\ell_3 m^\prime_3}_{m_3}
                                             \matr{L}_{m_1}^{\ell_3 m_3}
        $$

        Parameters
        ----------
        alm1 : npt.NDArray[np.complex128]
            The first set of spherical harmonic coefficients with shape (l1max+1, Nm1).
        alm2 : npt.NDArray[np.complex128]
            The second set of spherical harmonic coefficients with shape (l2max+1, 2*l2max+1).
        contract3 : npt.NDArray[np.complex128] | None, optional
            An optional linear operator along the m3 axis, shape
            (l3max+1, 2*l3max+1, Nm3prime<=2*l3max+1), by default None
        sum_m1 : bool, optional
            Whether to sum over the m1's before output, by default False.

        Returns
        -------
        npt.NDArray[np.complex128]
            By default (Nm1, Nl3, nm3) linear operator. If sum_m1=True, an (Nl3, Nm3) array
            where the m1 axis has been summed over. If contract3 is present, this will instead
            be an (Nm1, Nl3, Nm3prime) or (Nl3, Nm3prime) array resulting from the subsequent
            contraction on the m3 axis.
        """
        tensor = gaunt_dot12(
            alm1=alm1,
            alm2=alm2,
            l3max=self.l3max,
            sum_m1=sum_m1,
            absm1_lower=self.absm1_limits[0],
            absm1_upper=self.absm1_limits[1],
        )
        if contract3 is not None:
            return integrator12_contract3(tensor, contract3)
        else:
            return tensor

    def generate_integrator_cache_12(
        self,
        alm1: npt.NDArray[np.complex128],
        alm2: npt.NDArray[np.complex128],
        contract3: npt.NDArray[np.complex128] | None = None,
        release_gaunt_cache: bool = True,
        batch_parallel_mode: Literal["channel", "gaunt"] = "channel",
    ):
        """_summary_

        Parameters
        ----------
        alm1 : npt.NDArray[np.complex128]
            _description_
        alm2 : npt.NDArray[np.complex128]
            _description_
        contract3 : npt.NDArray[np.complex128] | None, optional
            _description_, by default None
        release_gaunt_cache : bool, optional
            _description_, by default True
        batch_parallel_mode : Literal["channel", "gaunt"], optional
            _description_, by default "channel"
        """
        if self.gaunt_cache is None:
            self.generate_gaunt_cache()
        assert type(self.gaunt_cache) is CachingGauntContractor
        if alm1.ndim < 3:
            alm1 = alm1[None, ...]
        if alm2.ndim < 3:
            alm2 = alm2[None, ...]
        result = np.zeros(
            (alm1.shape[0], len(self.m1_values), self.l3max + 1, 2 * self.l3max + 1),
            dtype=np.complex128,
        )
        if batch_parallel_mode == "channel":
            self.gaunt_cache.batch_parallel_batch_dot12(alm1, alm2, result)
        else:
            self.gaunt_cache.gaunt_parallel_batch_dot12(alm1, alm2, result)
        if release_gaunt_cache:
            self.clear_gaunt_cache()
        if contract3 is not None:
            result = integrator12_contract3(result, contract3)
        if result.shape[0] == 1:
            self.linear_integrator_cache_12 = result[0, ...]
        else:
            self.linear_integrator_cache_12 = result

    def batch_gaunt_integrate_cached_12(
        self,
        alm3: npt.NDArray[np.complex128],
        integrator_cache: npt.NDArray[np.complex128] | None = None,
        sum_m1: bool = False,
    ):
        if integrator_cache is None:
            if self.linear_integrator_cache_12 is None:
                raise ValueError("Integrator cache must be provided or pre-computed")
            else:
                tensor = self.linear_integrator_cache_12
        else:
            tensor = integrator_cache
        result = integrator12_dot3(tensor, alm3)
        if sum_m1:
            return np.sum(result, axis=0)
        else:
            return result

    def linear_gaunt_integrator_23(
        self,
        alm2: npt.NDArray[np.complex128],
        alm3: npt.NDArray[np.complex128],
        sum_m1: bool = False,
    ) -> npt.NDArray[np.complex128]:
        return gaunt_dot23(
            alm2=alm2,
            alm3=alm3,
            l1max=self.l1max,
            sum_m1=sum_m1,
            absm1_lower=self.absm1_limits[0],
            absm1_upper=self.absm1_limits[1],
        )

    def generate_integrator_cache_23(
        self,
        alm2: npt.NDArray[np.complex128],
        alm3: npt.NDArray[np.complex128],
        release_gaunt_cache: bool = True,
        batch_parallel_mode: Literal["channel", "gaunt"] = "channel",
    ):
        if self.gaunt_cache is None:
            self.generate_gaunt_cache()
        assert type(self.gaunt_cache) is CachingGauntContractor
        if alm2.ndim < 3:
            alm2 = alm2[None, ...]
        if alm3.ndim < 3:
            alm3 = alm3[None, ...]
        result = np.zeros(
            (alm2.shape[0], len(self.m1_values), self.l1max + 1), dtype=np.complex128
        )
        if batch_parallel_mode == "channel":
            self.gaunt_cache.batch_parallel_batch_dot23(alm2, alm3, result)
        else:
            self.gaunt_cache.gaunt_parallel_batch_dot23(alm2, alm3, result)
        if release_gaunt_cache:
            self.clear_gaunt_cache()
        if result.shape[0] == 1:
            self.linear_integrator_cache_23 = result[0, ...]
        else:
            self.linear_integrator_cache_23 = result

    def batch_gaunt_integrate_cached_23(
        self,
        alm1: npt.NDArray[np.complex128],
        integrator_cache: npt.NDArray[np.complex128] | None = None,
        sum_m1: bool = False,
    ):
        if integrator_cache is None:
            if self.linear_integrator_cache_23 is None:
                raise ValueError("Integrator cache must be provided or pre-computed")
            else:
                tensor = self.linear_integrator_cache_23
        else:
            tensor = integrator_cache
        result = integrator23_dot1(tensor, alm1)
        if sum_m1:
            return np.sum(result, axis=0)
        else:
            return result

clear_gaunt_cache()

Clears the Gaunt coefficient cache.

Source code in src/serval/core.py

def clear_gaunt_cache(self):
    """Clears the Gaunt coefficient cache."""
    self.gaunt_cache = None

gaunt_integrate(alm1, alm2, alm3, sum_m1=False)

Perform the triple integral by doing and inplace sum-product over gaunt coefficients.

\[ \mathcal{V} = \sum_{\substack{\ell_1 \ell_2 \ell_3 \\ m_1 m_2 m_3}} \mathcal{G}^{~~\ell_1~~\ell_2~~\ell_3}_{m_1 m_2 m_3}~ a_{\ell_1 m_1}a_{\ell_2 m_2}a_{\ell_3 m_3}. \]

Here the sum is performed inplace with no caching.

Parameters:

Name	Type	Description	Default
alm1	NDArray[complex128]	The first set of spherical harmonic coefficients with shape (l1max+1, Nm1).	required
alm2	NDArray[complex128]	The second set of spherical harmonic coefficients with shape (l2max+1, 2*l2max+1).	required
alm3	NDArray[complex128]	The third set of spherical harmonic coefficients with shape (l3max+1, 2*l3max+1).	required
sum_m1	bool	Whether to sum over the m1's before output, by default False.	False

Returns:

Type	Description
NDArray[complex128] | complex128	If sum_m1 = False, Nm1 array of the m-modes, otherwise, scalar sum over all m-modes.

Source code in src/serval/core.py

def gaunt_integrate(
    self,
    alm1: npt.NDArray[np.complex128],
    alm2: npt.NDArray[np.complex128],
    alm3: npt.NDArray[np.complex128],
    sum_m1: bool = False,
) -> npt.NDArray[np.complex128] | float:
    """Perform the triple integral by doing and inplace
    sum-product over gaunt coefficients.

    $$
    \mathcal{V} = \sum_{\substack{\ell_1 \ell_2 \ell_3 \\\\ m_1 m_2 m_3}}
                \mathcal{G}^{~~\ell_1~~\ell_2~~\ell_3}_{m_1 m_2 m_3}~
                a_{\ell_1 m_1}a_{\ell_2 m_2}a_{\ell_3 m_3}.
    $$

    Here the sum is performed inplace with no caching.

    Parameters
    ----------
    alm1 : npt.NDArray[np.complex128]
        The first set of spherical harmonic coefficients with shape (l1max+1, Nm1).
    alm2 : npt.NDArray[np.complex128]
        The second set of spherical harmonic coefficients with shape (l2max+1, 2*l2max+1).
    alm3 : npt.NDArray[np.complex128]
        The third set of spherical harmonic coefficients with shape (l3max+1, 2*l3max+1).
    sum_m1 : bool, optional
        Whether to sum over the m1's before output, by default False.

    Returns
    -------
    npt.NDArray[np.complex128] | np.complex128
        If sum_m1 = False, Nm1 array of the m-modes, otherwise, scalar
        sum over all m-modes.
    """
    return gaunt_dot123(
        alm1=alm1,
        alm2=alm2,
        alm3=alm3,
        sum_m1=sum_m1,
        absm1_lower=self.absm1_limits[0],
        absm1_upper=self.absm1_limits[1],
    )

generate_gaunt_cache()

Generates the Gaunt coefficient cache.

Source code in src/serval/core.py

def generate_gaunt_cache(self):
    """Generates the Gaunt coefficient cache."""
    self.gaunt_cache = CachingGauntContractor(
        self.l1max,
        self.l2max,
        self.l3max,
        absm1_lower=self.absm1_limits[0],
        absm1_upper=self.absm1_limits[1],
    )

generate_integrator_cache_12(alm1, alm2, contract3=None, release_gaunt_cache=True, batch_parallel_mode='channel')

summary

Parameters:

Name	Type	Description	Default
alm1	NDArray[complex128]	description	required
alm2	NDArray[complex128]	description	required
contract3	NDArray[complex128] | None	description, by default None	None
release_gaunt_cache	bool	description, by default True	True
batch_parallel_mode	Literal['channel', 'gaunt']	description, by default "channel"	'channel'

Source code in src/serval/core.py

def generate_integrator_cache_12(
    self,
    alm1: npt.NDArray[np.complex128],
    alm2: npt.NDArray[np.complex128],
    contract3: npt.NDArray[np.complex128] | None = None,
    release_gaunt_cache: bool = True,
    batch_parallel_mode: Literal["channel", "gaunt"] = "channel",
):
    """_summary_

    Parameters
    ----------
    alm1 : npt.NDArray[np.complex128]
        _description_
    alm2 : npt.NDArray[np.complex128]
        _description_
    contract3 : npt.NDArray[np.complex128] | None, optional
        _description_, by default None
    release_gaunt_cache : bool, optional
        _description_, by default True
    batch_parallel_mode : Literal["channel", "gaunt"], optional
        _description_, by default "channel"
    """
    if self.gaunt_cache is None:
        self.generate_gaunt_cache()
    assert type(self.gaunt_cache) is CachingGauntContractor
    if alm1.ndim < 3:
        alm1 = alm1[None, ...]
    if alm2.ndim < 3:
        alm2 = alm2[None, ...]
    result = np.zeros(
        (alm1.shape[0], len(self.m1_values), self.l3max + 1, 2 * self.l3max + 1),
        dtype=np.complex128,
    )
    if batch_parallel_mode == "channel":
        self.gaunt_cache.batch_parallel_batch_dot12(alm1, alm2, result)
    else:
        self.gaunt_cache.gaunt_parallel_batch_dot12(alm1, alm2, result)
    if release_gaunt_cache:
        self.clear_gaunt_cache()
    if contract3 is not None:
        result = integrator12_contract3(result, contract3)
    if result.shape[0] == 1:
        self.linear_integrator_cache_12 = result[0, ...]
    else:
        self.linear_integrator_cache_12 = result

grid_integrate(alm1, alm2, alm3, sum_m1=False)

Perform the triple integral using a grid-based approach.

Here the product of the second and third fields are evaluated on a consistent resolution grid, decomposed to conjugate spherical harmonics and multipled by the first set of coefficients. This more standard m-mode approach is performant for single integrals but, as implemented here, can't cache intermediate results and therefore not suitable for batching.

Parameters:

Name	Type	Description	Default
alm1	NDArray[complex128]	The first set of spherical harmonic coefficients with shape (l1max+1, Nm1).	required
alm2	NDArray[complex128]	The second set of spherical harmonic coefficients with shape (l2max+1, 2*l2max+1).	required
alm3	NDArray[complex128]	The third set of spherical harmonic coefficients with shape (l3max+1, 2*l3max+1).	required
sum_m1	bool	Whether to sum over the m1's before output, by default False.	False

Returns:

Type	Description
NDArray[complex128] | complex128	If sum_m1 = False, Nm1 array of the m-modes, otherwise, scalar sum over all m-modes.

Source code in src/serval/core.py

def grid_integrate(
    self,
    alm1: npt.NDArray[np.complex128],
    alm2: npt.NDArray[np.complex128],
    alm3: npt.NDArray[np.complex128],
    sum_m1: bool = False,
) -> npt.NDArray[np.complex128]:
    """Perform the triple integral using a grid-based approach.

    Here the product of the second and third fields are evaluated on a
    consistent resolution grid, decomposed to conjugate spherical
    harmonics and multipled by the first set of coefficients. This more
    standard m-mode approach is performant for single integrals
    but, as implemented here, can't cache intermediate results and
    therefore not suitable for batching.


    Parameters
    ---------
    alm1 : npt.NDArray[np.complex128]
        The first set of spherical harmonic coefficients with shape (l1max+1, Nm1).
    alm2 : npt.NDArray[np.complex128]
        The second set of spherical harmonic coefficients with shape (l2max+1, 2*l2max+1).
    alm3 : npt.NDArray[np.complex128]
        The third set of spherical harmonic coefficients with shape (l3max+1, 2*l3max+1).
    sum_m1 : bool, optional
        Whether to sum over the m1's before output, by default False.

    Returns
    -------
    npt.NDArray[np.complex128] | np.complex128
        If sum_m1 = False, Nm1 array of the m-modes, otherwise, scalar
        sum over all m-modes.
    """
    grid_lmax = int(np.ceil(self.l1max + self.l2max + self.l3max))
    alm1, alm2, alm3 = broadcast_bandlimits(alm1, alm2, alm3, lmax=grid_lmax)
    static_grid = grid_template(grid_lmax)
    static_grid.data[:] = 1.0
    static_grid *= array_synthesis(alm2) * array_synthesis(alm3)
    static_alm = analysis(static_grid, conjugate=True, return_sparse=False)
    out_mmodes = compute_mmodes(static_alm, alm1)
    return out_mmodes[np.asarray(self.m1_values) + grid_lmax]

linear_gaunt_integrator_12(alm1, alm2, contract3=None, sum_m1=False)

Generate a linear operator that, with fixed first and second set of coefficients performs the integral when sum-producted with the third set of coefficients.

\[ \matr{L}_{m_1}^{\ell_3 m_3} = \sum_{\substack{\ell_1 \ell_2 \\ m_1 m_2}} \mathcal{G}^{~~\ell_1~~\ell_2~~\ell_3}_{m_1 m_2 m_3}~ a_{\ell_1 m_1}a_{\ell_2 m_2}. \]

If requested, additionally contract with a linear operator along the m3 axis (e.g. a rotation with a Wigner-d matrix),

\[ \matr{L}_{m_1}^{\ell_3 m^\prime_3} = \sum_{m_3} \matr{W}^{\ell_3 m^\prime_3}_{m_3} \matr{L}_{m_1}^{\ell_3 m_3} \]

Parameters:

Name	Type	Description	Default
alm1	NDArray[complex128]	The first set of spherical harmonic coefficients with shape (l1max+1, Nm1).	required
alm2	NDArray[complex128]	The second set of spherical harmonic coefficients with shape (l2max+1, 2*l2max+1).	required
contract3	NDArray[complex128] | None	An optional linear operator along the m3 axis, shape (l3max+1, 2l3max+1, Nm3prime<=2l3max+1), by default None	None
sum_m1	bool	Whether to sum over the m1's before output, by default False.	False

Returns:

Type	Description
NDArray[complex128]	By default (Nm1, Nl3, nm3) linear operator. If sum_m1=True, an (Nl3, Nm3) array where the m1 axis has been summed over. If contract3 is present, this will instead be an (Nm1, Nl3, Nm3prime) or (Nl3, Nm3prime) array resulting from the subsequent contraction on the m3 axis.

Source code in src/serval/core.py

def linear_gaunt_integrator_12(
    self,
    alm1: npt.NDArray[np.complex128],
    alm2: npt.NDArray[np.complex128],
    contract3: npt.NDArray[np.complex128] | None = None,
    sum_m1: bool = False,
) -> npt.NDArray[np.complex128]:
    """Generate a linear operator that, with fixed first and second set of coefficients
    performs the integral when sum-producted with the third set of coefficients.

    $$
    \matr{L}_{m_1}^{\ell_3 m_3} =
                \sum_{\substack{\ell_1 \ell_2 \\\\ m_1 m_2}}
                \mathcal{G}^{~~\ell_1~~\ell_2~~\ell_3}_{m_1 m_2 m_3}~
                a_{\ell_1 m_1}a_{\ell_2 m_2}.
    $$

    If requested, additionally contract with a linear operator along the
    m3 axis (e.g. a rotation with a Wigner-d matrix),

    $$
    \matr{L}_{m_1}^{\ell_3 m^\prime_3} = \sum_{m_3} \matr{W}^{\ell_3 m^\prime_3}_{m_3}
                                         \matr{L}_{m_1}^{\ell_3 m_3}
    $$

    Parameters
    ----------
    alm1 : npt.NDArray[np.complex128]
        The first set of spherical harmonic coefficients with shape (l1max+1, Nm1).
    alm2 : npt.NDArray[np.complex128]
        The second set of spherical harmonic coefficients with shape (l2max+1, 2*l2max+1).
    contract3 : npt.NDArray[np.complex128] | None, optional
        An optional linear operator along the m3 axis, shape
        (l3max+1, 2*l3max+1, Nm3prime<=2*l3max+1), by default None
    sum_m1 : bool, optional
        Whether to sum over the m1's before output, by default False.

    Returns
    -------
    npt.NDArray[np.complex128]
        By default (Nm1, Nl3, nm3) linear operator. If sum_m1=True, an (Nl3, Nm3) array
        where the m1 axis has been summed over. If contract3 is present, this will instead
        be an (Nm1, Nl3, Nm3prime) or (Nl3, Nm3prime) array resulting from the subsequent
        contraction on the m3 axis.
    """
    tensor = gaunt_dot12(
        alm1=alm1,
        alm2=alm2,
        l3max=self.l3max,
        sum_m1=sum_m1,
        absm1_lower=self.absm1_limits[0],
        absm1_upper=self.absm1_limits[1],
    )
    if contract3 is not None:
        return integrator12_contract3(tensor, contract3)
    else:
        return tensor

sparse_gaunt_einsum(sparse_alm1, sparse_alm2, sparse_alm3, sum_m1=False, gaunts=None)

Perform the triple integral using a sparse tensor contraction.

Compute the triple integral by generating or using a sparse representation of the Gaunt coefficients and performing the sum with a sparse einsum. This is mostly a cross-checking function and is not performant.

Parameters:

Name	Type	Description	Default
sparse_alm1	SparseArray	Sparse represenation of the first set of coefficients with shape (l1max+1, Nm1).	required
sparse_alm2	SparseArray	Sparse represenation of the second set of coefficients with shape (l2max+1, 2*l2max+1).	required
sparse_alm3	SparseArray	Sparse represenation of the third set of coefficients (l3max+1, 2*l3max+1).	required
sum_m1	bool	Whether to sum over the m1's before output, by default False.	False
gaunts	COO | None	Sparse representation of the Gaunt coefficients if already. computed, by default None.	None

Returns:

Type	Description
NDArray[complex128] | complex128	If sum_m1 = False, Nm1 array of the m-modes, otherwise, scalar sum over all m-modes.

Source code in src/serval/core.py

def sparse_gaunt_einsum(
    self,
    sparse_alm1: sparse.SparseArray,
    sparse_alm2: sparse.SparseArray,
    sparse_alm3: sparse.SparseArray,
    sum_m1: bool = False,
    gaunts: sparse.COO | None = None,
) -> sparse.COO:
    """
    Perform the triple integral using a sparse tensor contraction.

    Compute the triple integral by generating or using a sparse
    representation of the Gaunt coefficients and performing the sum
    with a sparse einsum. This is mostly a cross-checking function
    and is not performant.

    Parameters
    ----------
    sparse_alm1 : sparse.SparseArray
        Sparse represenation of the first set of coefficients with shape (l1max+1, Nm1).
    sparse_alm2 : sparse.SparseArray
        Sparse represenation of the second set of coefficients with shape (l2max+1, 2*l2max+1).
    sparse_alm3 : sparse.SparseArray
        Sparse represenation of the third set of coefficients (l3max+1, 2*l3max+1).
    sum_m1 : bool, optional
        Whether to sum over the m1's before output, by default False.
    gaunts : sparse.COO | None, optional
        Sparse representation of the Gaunt coefficients if already.
        computed, by default None.

    Returns
    -------
    npt.NDArray[np.complex128] | np.complex128
        If sum_m1 = False, Nm1 array of the m-modes, otherwise, scalar
        sum over all m-modes.
    """
    if gaunts is None:
        gaunts = gaunts_coo(self.l1max, self.l2max, self.l3max)
    if sum_m1:
        return sparse.einsum(
            "ijklmn, ij, kl, mn -> ", gaunts, sparse_alm1, sparse_alm2, sparse_alm3
        )
    else:
        return sparse.to_numpy(
            sparse.einsum(
                "ijklmn, ij, kl, mn -> j",
                gaunts,
                sparse_alm1,
                sparse_alm2,
                sparse_alm3,
            )
        )

Rotate.py

Sht.py

Uitls.py

Gaunt Library

Core.py

gaunt_dot12(alm1, alm2, l3max, sum_m1=False, absm1_lower=None, absm1_upper=None)

Compute a projector for alm3 for the integral of the triple product of spherical harmonics by computing the Gaunt co-efficients in-place and sum-producting over l1, l2, m2 and m1 (if requested).

TODO add formula.

Parameters:

Name	Type	Description	Default
alm1	NDArray[complex128]	Spherical harmonic co-efficients in l1, m1.	required
alm2	NDArray[complex128]	Spherical harmonic co-efficients in l2, m2.	required
l3max	int	Maximum l3 to compute co-efficients up to.	required
sum_m1	bool	If True, also sum over m1, otherwise return m1-modes for the m1 range specified, by default False	False
absm1_lower	int | None	Lower limit in |m1| to use, by default None, ie. |m1| >= 0.	None
absm1_upper	int | None	Upper limit in |m1| to use, by default None, ie. |m1| <= m1max = l1max.	None

Returns:

Type	Description
NDArray[complex128]	m1-mode alm3 projector as a numpy array of shape (Nm1, l3max+1, 2l3max+1) if sum_m1 is False or (l3max+1, 2l3max+1) is sum_m1 is True.

Source code in src/serval/gaunt/core.py

def gaunt_dot12(
    alm1: npt.NDArray[np.complex128],
    alm2: npt.NDArray[np.complex128],
    l3max: int,
    sum_m1: bool = False,
    absm1_lower: int | None = None,
    absm1_upper: int | None = None,
) -> npt.NDArray[np.complex128]:
    """Compute a projector for alm3 for the integral of the triple product
    of spherical harmonics by computing the Gaunt co-efficients in-place
    and sum-producting over l1, l2, m2 and m1 (if requested).

    TODO add formula.

    Parameters
    ----------
    alm1 : npt.NDArray[np.complex128]
        Spherical harmonic co-efficients in l1, m1.
    alm2 : npt.NDArray[np.complex128]
        Spherical harmonic co-efficients in l2, m2.
    l3max : int
        Maximum l3 to compute co-efficients up to.
    sum_m1 : bool, optional
        If True, also sum over m1, otherwise return m1-modes for the m1
        range specified, by default False
    absm1_lower : int | None, optional
        Lower limit in |m1| to use, by default None, ie. |m1| >= 0.
    absm1_upper : int | None, optional
        Upper limit in |m1| to use, by default None, ie. |m1| <= m1max = l1max.

    Returns
    -------
    npt.NDArray[np.complex128]
        m1-mode alm3 projector as a numpy array of shape (Nm1, l3max+1, 2*l3max+1)
        if sum_m1 is False or (l3max+1, 2*l3max+1) is sum_m1 is True.
    """
    if absm1_lower is None:
        absm1_lower = 0
    if absm1_upper is None:
        absm1_upper = alm1.shape[0]
    nm1 = len(_gaunt.m1_indexing(alm1.shape[0] - 1, absm1_lower, absm1_upper)[0])
    result = np.zeros((nm1, l3max + 1, 2 * l3max + 1), dtype=np.complex128)
    _gaunt.inplace_dot12(alm1, alm2, result, l3max, absm1_lower, absm1_upper)
    if sum_m1:
        return np.sum(result, axis=0)
    else:
        return result

gaunt_dot123(alm1, alm2, alm3, sum_m1=False, absm1_lower=None, absm1_upper=None)

Compute the integral of the triple product of spherical harmonics by computing the Gaunt co-efficients in-place and sum-producting over all harmonic degrees and orders except m1, unless requested.

TODO add formula.

Parameters:

Name	Type	Description	Default
alm1	NDArray[complex128]	Spherical harmonic co-efficients in l1, m1.	required
alm2	NDArray[complex128]	Spherical harmonic co-efficients in l2, m2.	required
alm3	NDArray[complex128]	Spherical harmonic co-efficients in l3, m3.	required
sum_m1	bool	If True, also sum over m1, otherwise return m1-modes for the m1 range specified, by default False	False
absm1_lower	int | None	Lower limit in |m1| to use, by default None, ie. |m1| >= 0.	None
absm1_upper	int | None	Upper limit in |m1| to use, by default None, ie. |m1| <= m1max = l1max.	None

Returns:

Type	Description
NDArray[complex128] | float	Numpy array, shape (Nm1), of m1-modes or their sum if sum_m1 is True.

Source code in src/serval/gaunt/core.py

def gaunt_dot123(
    alm1: npt.NDArray[np.complex128],
    alm2: npt.NDArray[np.complex128],
    alm3: npt.NDArray[np.complex128],
    sum_m1: bool = False,
    absm1_lower: int | None = None,
    absm1_upper: int | None = None,
) -> npt.NDArray[np.complex128] | float:
    """Compute the integral of the triple product of spherical harmonics
    by computing the Gaunt co-efficients in-place and sum-producting over all
    harmonic degrees and orders except m1, unless requested.

    TODO add formula.

    Parameters
    ----------
    alm1 : npt.NDArray[np.complex128]
        Spherical harmonic co-efficients in l1, m1.
    alm2 : npt.NDArray[np.complex128]
        Spherical harmonic co-efficients in l2, m2.
    alm3 : npt.NDArray[np.complex128]
        Spherical harmonic co-efficients in l3, m3.
    sum_m1 : bool, optional
        If True, also sum over m1, otherwise return m1-modes for the m1
        range specified, by default False
    absm1_lower : int | None, optional
        Lower limit in |m1| to use, by default None, ie. |m1| >= 0.
    absm1_upper : int | None, optional
        Upper limit in |m1| to use, by default None, ie. |m1| <= m1max = l1max.

    Returns
    -------
    npt.NDArray[np.complex128] | float
        Numpy array, shape (Nm1), of m1-modes or their sum if sum_m1 is True.
    """
    if absm1_lower is None:
        absm1_lower = 0
    if absm1_upper is None:
        absm1_upper = alm1.shape[0]
    nm1 = len(_gaunt.m1_indexing(alm1.shape[0] - 1, absm1_lower, absm1_upper)[0])
    result = np.zeros(nm1, dtype=np.complex128)
    _gaunt.inplace_dot123(alm1, alm2, alm3, result, absm1_lower, absm1_upper)
    if sum_m1:
        return result.sum()
    else:
        return result

gaunt_dot23(alm2, alm3, l1max, sum_m1=False, absm1_lower=None, absm1_upper=None)

Compute a projector for alm1 for the integral of the triple product of spherical harmonics by computing the Gaunt co-efficients in-place and sum-producting over l2, m2, l3, m3 and m1 (if requested).

TODO add formula.

Parameters:

Name	Type	Description	Default
alm2	NDArray[complex128]	Spherical harmonic co-efficients in l2, m2.	required
alm3	NDArray[complex128]	Spherical harmonic co-efficients in l3, m3.	required
l1max	int	Maximum l1 to compute co-efficients up to.	required
sum_m1	bool	If True, also sum over m1, otherwise return m1-modes for the m1 range specified, by default False	False
absm1_lower	int | None	Lower limit in |m1| to use, by default None, ie. |m1| >= 0.	None
absm1_upper	int | None	Upper limit in |m1| to use, by default None, ie. |m1| <= m1max = l1max.	None

Returns:

Type	Description
NDArray[complex128]	m1-mode alm1 projector as a numpy array of shape (Nm1, l1max+1) if sum_m1 is False or (l1max+1,) is sum_m1 is True.

Source code in src/serval/gaunt/core.py

def gaunt_dot23(
    alm2: npt.NDArray[np.complex128],
    alm3: npt.NDArray[np.complex128],
    l1max: int,
    sum_m1: bool = False,
    absm1_lower: int | None = None,
    absm1_upper: int | None = None,
) -> npt.NDArray[np.complex128]:
    """Compute a projector for alm1 for the integral of the triple product
    of spherical harmonics by computing the Gaunt co-efficients in-place
    and sum-producting over l2, m2, l3, m3 and m1 (if requested).

    TODO add formula.

    Parameters
    ----------
    alm2 : npt.NDArray[np.complex128]
        Spherical harmonic co-efficients in l2, m2.
    alm3 : npt.NDArray[np.complex128]
        Spherical harmonic co-efficients in l3, m3.
    l1max : int
        Maximum l1 to compute co-efficients up to.
    sum_m1 : bool, optional
        If True, also sum over m1, otherwise return m1-modes for the m1
        range specified, by default False
    absm1_lower : int | None, optional
        Lower limit in |m1| to use, by default None, ie. |m1| >= 0.
    absm1_upper : int | None, optional
        Upper limit in |m1| to use, by default None, ie. |m1| <= m1max = l1max.

    Returns
    -------
    npt.NDArray[np.complex128]
        m1-mode alm1 projector as a numpy array of shape (Nm1, l1max+1)
        if sum_m1 is False or (l1max+1,) is sum_m1 is True.
    """
    if absm1_lower is None:
        absm1_lower = 0
    if absm1_upper is None:
        absm1_upper = l1max + 1
    nm1 = len(_gaunt.m1_indexing(l1max, absm1_lower, absm1_upper)[0])
    result = np.zeros((nm1, l1max + 1), dtype=np.complex128)
    _gaunt.inplace_dot23(alm2, alm3, result, l1max, absm1_lower, absm1_upper)
    if sum_m1:
        return np.sum(result, axis=0)
    else:
        return result

gaunts_coo(l1max, l2max, l3max)

Compute a sparse representation of all Gaunt co-efficients up to maximum harmonic degrees l1max, l2max and l3max.

Parameters:

Name	Type	Description	Default
l1max	int	Maximum harmonic degree l1.	required
l2max	int	Maximum harmonic degree l2.	required
l3max	int	Maximum harmonic degree l3.	required

Returns:

Type	Description
COO	COO formatted sparse array of the Gaunt co-efficients.

Source code in src/serval/gaunt/core.py

def gaunts_coo(l1max: int, l2max: int, l3max: int) -> sparse.COO:
    """Compute a sparse representation of all Gaunt co-efficients up
    to maximum harmonic degrees l1max, l2max and l3max.

    Parameters
    ----------
    l1max : int
        Maximum harmonic degree l1.
    l2max : int
        Maximum harmonic degree l2.
    l3max : int
        Maximum harmonic degree l3.

    Returns
    -------
    sparse.COO
        COO formatted sparse array of the Gaunt co-efficients.
    """
    coords, data, shape = _gaunt.gaunts_coo(l1max, l2max, l3max)
    coords = np.array(coords)
    coords = coords.view(dtype=coords.dtype[0]).reshape(-1, 6).T
    data = np.array(data)
    return sparse.COO(
        coords=coords,
        data=data,
        shape=shape,
        has_duplicates=False,
        sorted=True,
        fill_value=0.0,
    )

gaunts_l1m1_coo(l1, absm1, l1max, l2max, l3max)

Compute a sparse representation of all Gaunt co-efficients with given harmonic degree l1 and harmonic order |m1|, up to maximum harmonic degrees l1max, l2max and l3max. The shape of the coordinates is the same as for the full gaunt coeffcients (gaunts_coo). The only difference is that the first two indices are l1 and m1, set to the input values.

Parameters:

Name	Type	Description	Default
l1	int	Harmonic degree l1.	required
absm1	int	Harmonic order |m1|.	required
l1max	int	Maximum harmonic degree l1.	required
l2max	int	Maximum harmonic degree l2.	required
l3max	int	Maximum harmonic degree l3.	required

Returns:

Type	Description
COO	COO formatted sparse array of the Gaunt co-efficients.

Source code in src/serval/gaunt/core.py

def gaunts_l1m1_coo(
    l1: int, absm1: int, l1max: int, l2max: int, l3max: int
) -> sparse.COO:
    """Compute a sparse representation of all Gaunt co-efficients with given harmonic degree
    l1 and harmonic order |m1|, up to maximum harmonic degrees l1max, l2max and l3max. The shape
    of the coordinates is the same as for the full gaunt coeffcients (gaunts_coo).
    The only difference is that the first two indices are l1 and m1, set to the input values.

    Parameters
    ----------
    l1 : int
        Harmonic degree l1.
    absm1 : int
        Harmonic order |m1|.
    l1max : int
        Maximum harmonic degree l1.
    l2max : int
        Maximum harmonic degree l2.
    l3max : int
        Maximum harmonic degree l3.

    Returns
    -------
    sparse.COO
        COO formatted sparse array of the Gaunt co-efficients.
    """
    # Why is l1max needed...?
    coords, data, shape = _gaunt.gaunts_l1m1_coo(l1, absm1, l1max, l2max, l3max)
    coords = np.array(coords)
    coords = coords.view(dtype=coords.dtype[0]).reshape(-1, 6).T
    data = np.array(data)
    return sparse.COO(
        coords=coords,
        data=data,
        shape=shape,
        has_duplicates=False,
        sorted=True,
        fill_value=0.0,
    )

integrator12_contract3(int12, contract3)

(m1 l3 m3) (l3 m3' m3) -> (m1 l3 m3')

Source code in src/serval/gaunt/core.py

def integrator12_contract3(
    int12: npt.NDArray[np.complex128], contract3: npt.NDArray[np.complex128]
) -> npt.NDArray[np.complex128]:
    """(m1 l3 m3) (l3 m3' m3) -> (m1 l3 m3')"""
    # Check m3' axis smaller than m3 axis
    if int12.ndim == 3:  # Has no frequency axis
        result = np.zeros(
            (int12.shape[0], contract3.shape[0], contract3.shape[1]),
            dtype=np.complex128,
        )
        _contract.integrator12_contract3(int12, contract3, result)
        return result
    elif int12.ndim == 4:  # Has frequency axis
        result = np.zeros(
            (int12.shape[0] * int12.shape[1], contract3.shape[0], contract3.shape[1]),
            dtype=np.complex128,
        )
        _contract.integrator12_contract3(
            int12.reshape((-1,) + int12.shape[2:]), contract3, result
        )
        return result.reshape(
            (int12.shape[0], int12.shape[1], contract3.shape[0], contract3.shape[1])
        )
    else:
        raise ValueError

single_gaunt(l1, l2, l3, m1, m2)

Computes a single Gaunt co-efficient for given harmonic degrees and orders. This uses Wigner-3j family algorithms so it is not efficient for computing many Gaunts co-efficients. Primarily for testing purposes m3 is determined by m1 + m2 + m3 = 0.

Parameters:

Name	Type	Description	Default
l1	int	Harmonic degree l1.	required
l2	int	Harmonic degree l2.	required
l3	int	Harmonic degree l3.	required
m1	int	Harmonic order m1.	required
m2	int	Harmonic order m2.	required

Returns:

Type	Description
float	Computed Gaunt co-efficient.

Source code in src/serval/gaunt/core.py

def single_gaunt(l1: int, l2: int, l3: int, m1: int, m2: int) -> float:
    """Computes a single Gaunt co-efficient for given harmonic degrees
    and orders. This uses Wigner-3j family algorithms so it is not
    efficient for computing many Gaunts co-efficients. Primarily for
    testing purposes m3 is determined by m1 + m2 + m3 = 0.

    Parameters
    ----------
    l1 : int
        Harmonic degree l1.
    l2 : int
        Harmonic degree l2.
    l3 : int
        Harmonic degree l3.
    m1 : int
        Harmonic order m1.
    m2 : int
        Harmonic order m2.

    Returns
    -------
    float
        Computed Gaunt co-efficient.
    """
    triangle = abs(l1 - l3) <= l2 <= abs(l1 + l3)
    parity = (l1 + l2 + l3) % 2 != 0
    if (
        not triangle
        or l1 < 0
        or l2 < 0
        or l3 < 0
        or abs(m1) > l1
        or abs(m2) > l2
        or abs(m1 + m2) > l3
        or parity      # parity condition: the sum of the harmonic degrees must be even
    ):
        return 0.0
    return _gaunt.single_gaunt(l1, l2, l3, m1, m2)

wigner_3jj(l2, l3, m2, m3)

Compute the family of non-zero Wigner-3j terms for harmonic degrees l2 and l3 for harmonic degrees m1, m2, m3 = -m1 -m2.

Parameters:

Name	Type	Description	Default
l2	int	Harmonic degree l2.	required
l3	int	Harmonic degree l3.	required
m2	int	Harmonic order m2.	required
m3	int	Harmonic order m3.	required

Returns:

Type	Description
tuple[int, NDArray[float64]]	Tuple of first non-zero harmonic degree l1min and array of Wigner-3j values of length l1max - l1min. If there are no non-zero elements of the Wigner-3j family, returns (-1, np.array([])).

Source code in src/serval/gaunt/core.py

def wigner_3jj(
    l2: int, l3: int, m2: int, m3: int
) -> tuple[int, npt.NDArray[np.float64]]:
    """Compute the family of non-zero Wigner-3j terms for harmonic degrees l2 and l3
    for harmonic degrees m1, m2, m3 = -m1 -m2.

    Parameters
    ----------
    l2 : int
        Harmonic degree l2.
    l3 : int
        Harmonic degree l3.
    m2 : int
        Harmonic order m2.
    m3 : int
        Harmonic order m3.

    Returns
    -------
    tuple[int, npt.NDArray[np.float64]]
        Tuple of first non-zero harmonic degree l1min and array of Wigner-3j values of
        length l1max - l1min. If there are no non-zero elements of the Wigner-3j family,
        returns (-1, np.array([])).
    """
    # In case there is an error in the result, pyshtools Wigner3j can also be used.
    # However, it carries extra zeros that need to be removed:
    # Wigner_3j = pysh.utils.Wigner3j(l2, l3, -m2 - m3, m2, m3)[0]
    # lmin = max(abs(l2 - l3), abs(-m2 -m3))
    # lmax = l2 + l3
    # lnum = lmax - lmin + 1
    # return lmin, Wigner_3j[: lnum + 1 :]
    if m2 == 0 and m3 == 0:
        return wigner_3jj_000(l2, l3)
    if l2 < 0 or l3 < 0 or abs(m2) > l2 or abs(m3) > l3:
        return (-1, np.array([]))
    result = _wigner.wigner_3jj_nochecks(l2, l3, m2, m3)
    return result[0], np.array(result[1])

wigner_3jj_000(l2, l3)

Compute the family of non-zero Wigner-3j terms for harmonic degrees l2 and l3 for harmonic degrees m1 = m2 = m3 = 0.

Parameters:

Name	Type	Description	Default
l2	int	Harmonic degree l2.	required
l3	int	Harmonic degree l3.	required

Returns:

Type	Description
tuple[int, NDArray[float64]]	Tuple of first non-zero harmonic order l1min and array of Wigner-3j values of length l1max - l1min. If there are no non-zero elements of the Wigner-3j family, returns (-1, np.array([])).

Source code in src/serval/gaunt/core.py

def wigner_3jj_000(l2: int, l3: int) -> tuple[int, npt.NDArray[np.float64]]:
    """Compute the family of non-zero Wigner-3j terms for harmonic degrees l2 and l3
    for harmonic degrees m1 = m2 = m3 = 0.

    Parameters
    ----------
    l2 : int
        Harmonic degree l2.
    l3 : int
        Harmonic degree l3.

    Returns
    -------
    tuple[int, npt.NDArray[np.float64]]
        Tuple of first non-zero harmonic order l1min and array of Wigner-3j values of
        length l1max - l1min. If there are no non-zero elements of the Wigner-3j family,
        returns (-1, np.array([])).
    """
    if l2 < 0 or l3 < 0:
        return (-1, np.array([]))
    result = _wigner.wigner_3jj_000_nochecks(l2, l3)
    return result[0], np.array(result[1])

wigner_3jm(l1, l2, l3, m1)

Compute the family of Wigner-3j terms for harmonic degrees l1, l2 and l3 all non-zero terms with harmonic order m1.

Parameters:

Name	Type	Description	Default
l1	int	Harmonic degree l1.	required
l2	int	Harmonic degree l2.	required
l3	int	Harmonic degree l3.	required
m1	int	Harmonic order m1.	required

Returns:

Type	Description
tuple[int, NDArray[float64]]	Tuple of first non-zero harmonic order m2min and array of Wigner-3j values of length l2max - l2min. If there are no non-zero elements of the Wigner-3j family, returns (-1, np.array([])).

Source code in src/serval/gaunt/core.py

def wigner_3jm(
    l1: int, l2: int, l3: int, m1: int
) -> tuple[int, npt.NDArray[np.float64]]:
    """Compute the family of Wigner-3j terms for harmonic degrees l1, l2 and l3
    all non-zero terms with harmonic order m1.

    Parameters
    ----------
    l1 : int
        Harmonic degree l1.
    l2 : int
        Harmonic degree l2.
    l3 : int
        Harmonic degree l3.
    m1 : int
        Harmonic order m1.

    Returns
    -------
    tuple[int, npt.NDArray[np.float64]]
        Tuple of first non-zero harmonic order m2min and array of Wigner-3j values of
        length l2max - l2min. If there are no non-zero elements of the Wigner-3j family,
        returns (-1, np.array([])).
    """
    triangle = abs(l1 - l3) <= l2 <= abs(l1 + l3)
    if (
        not triangle
        or l1 < 0
        or l2 < 0
        or l3 < 0
        or abs(m1) > l1
        or not (abs(l1 - l3) < l2 < l1 + l3)
    ):
        return (-1, np.array([]))
    result = _wigner.wigner_3jm_nochecks(l1, l2, l3, m1)
    return result[0], np.array(result[1])

Gaunt.cpp

These are the c++ functions:

Wigner.cpp

These are the c++ functions:

\[ E = mc^2 \]