s2conv 0.1.0__py3-none-any.whl

s2conv/__init__.py

@@ -0,0 +1,35 @@
+r"""
+s2conv: differentiable and accelerated spherical convolutions with JAX.
+
+s2conv implements the generalised lifted spherical convolution, in which a spin-:math:`s`
+field is lifted to the rotation group, convolved there with a kernel :math:`\Psi` and
+projected back to the sphere through the adjoint of the lift for a possibly different
+spin :math:`s'`, and the spherical convolutions of the literature that it unifies.  The
+harmonic transforms are those of s2fft, which warns on import if JAX is not using 64-bit
+precision.
+"""
+
+from . import convolutions, kernels, lifting, transforms, utils
+from .convolutions import (
+    classical,
+    diagonal,
+    group_convolution,
+    group_to_group,
+    group_to_spin,
+    lifted,
+    lifted_convolution,
+    physics,
+    sampled,
+    section,
+    spin_multiplier,
+    spin_to_group,
+    spin_to_spin,
+)
+from .convolutions.sampled import convolve, convolve_multiplier
+from .lifting import lift, project
+from .transforms import forward, inverse, wigner_forward, wigner_inverse
+
+try:
+    from ._version import __version__
+except ImportError:  # pragma: no cover - only without an installed version file
+    __version__ = "unknown"

s2conv/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
+
+__commit_id__ = commit_id = 'g03ddef413'

s2conv/convolutions/__init__.py

@@ -0,0 +1,40 @@
+"""The generalised lifted convolution and the spherical convolutions it unifies."""
+
+from . import classical, diagonal, lifted, physics, sampled, section
+from .classical import (
+    axisymmetric_convolution,
+    axisymmetric_synthesis,
+    directional_convolution,
+    directional_synthesis,
+    left_convolution,
+    spin_changing_convolution,
+)
+from .diagonal import harmonic_multiplication, sifting_convolution
+from .lifted import (
+    group_convolution,
+    group_to_group,
+    group_to_group_weights,
+    group_to_spin,
+    group_to_spin_weights,
+    lifted_convolution,
+    spin_multiplier,
+    spin_to_group,
+    spin_to_group_weights,
+    spin_to_spin,
+)
+from .physics import (
+    eb_decomposition,
+    eb_to_polarisation,
+    eth,
+    ethbar,
+    kaiser_squires,
+    kaiser_squires_inverse,
+)
+from .sampled import convolve, convolve_multiplier
+from .section import (
+    commutative_anisotropic_convolution,
+    section_components,
+    section_evaluation,
+    standard_correlation,
+    standard_correlation_fourier,
+)

s2conv/convolutions/classical.py

@@ -0,0 +1,175 @@
+r"""
+The rotation-equivariant spherical convolutions of the literature (Secs. III and VI).
+
+Each is a generalised lifted convolution with a particular kernel (Table II), and is
+written here in its closed harmonic form, with the constants of the article:
+
+.. list-table::
+   :header-rows: 1
+
+   * - Convolution
+     - Lifted form
+     - Harmonic form
+   * - axisymmetric (Defs III.1, III.6)
+     - :math:`(2\pi)^{-1/2}\mathcal{C}_{s\to0}[\Lambda_s\psi]`
+     - :math:`\sqrt{4\pi/(2\ell+1)}\, {}_sf_{\ell m}\, {}_s\psi^*_{\ell0}`
+   * - left (Def. III.2)
+     - :math:`\mathcal{C}_{0\to0}[\Psi_h]`
+     - :math:`2\pi\sqrt{4\pi/(2\ell+1)}\, f_{\ell m} h_{\ell0}`
+   * - directional (Defs III.3, III.5)
+     - :math:`\mathcal{C}_{s\to G}[\Lambda_s\psi]`
+     - :math:`\frac{8\pi^2}{2\ell+1}\, {}_sf_{\ell m}\, {}_s\psi^*_{\ell n}`
+   * - synthesis (Def. III.10)
+     - :math:`\mathcal{C}_{G\to s}[\widetilde{\Lambda_s\psi}]`
+     - :math:`\sum_n W^{\ell}_{mn}\, {}_s\psi_{\ell n}`
+   * - spin-changing (Prop. VI.5)
+     - :math:`\mathcal{C}_{s\to s'}[\Psi_T]`
+     - :math:`c_\ell\, {}_sf_{\ell m}`
+
+The group convolution (Definition III.4) is :func:`s2conv.convolutions.lifted.group_convolution`.
+All functions act on harmonic coefficients in the s2fft layout, broadcast over leading
+axes, and are differentiable and JIT compilable.
+"""
+
+import jax.numpy as jnp
+import numpy as np
+
+from s2conv.convolutions.lifted import (
+    group_to_spin_weights,
+    spin_multiplier,
+    spin_to_group_weights,
+)
+from s2conv.utils.indexing import azimuthal_bandlimit, lm_mask
+
+
+def _axisymmetric_weight(L: int) -> np.ndarray:
+    return np.sqrt(4 * np.pi / (2 * np.arange(L) + 1))
+
+
+def axisymmetric_convolution(flm, psi_lm, spin: int = 0) -> jnp.ndarray:
+    r"""
+    Axisymmetric convolution :math:`(f \odot \psi)(\omega) = \langle f, \mathcal{R}_\omega \psi \rangle`.
+
+    For a scalar field (Definition III.1) or a spin-:math:`s` field (Definition III.6) with an
+    axisymmetric kernel, the output is a scalar field with
+    :math:`(f \odot \psi)_{\ell m} = \sqrt{4\pi/(2\ell+1)}\, {}_sf_{\ell m}\, {}_s\psi^*_{\ell 0}`
+    (Eqs. 14, 25).  It is :math:`(2\pi)^{-1/2}\mathcal{C}_{s\to0}[\Lambda_s\psi]`, the factor being the
+    difference between evaluation on the section and the adjoint projection
+    (Proposition VI.3).  Only the axisymmetric part :math:`\psi_{\ell0}` of the kernel enters.
+
+    Args:
+        flm: Spin-:math:`s` harmonic coefficients, ``[..., L, 2L - 1]``.
+        psi_lm: Spin-:math:`s` harmonic coefficients of the kernel, ``[L, 2L - 1]``.
+        spin (int, optional): Spin :math:`s`. Defaults to 0.
+
+    Returns:
+        jnp.ndarray: Scalar harmonic coefficients, ``[..., L, 2L - 1]``.
+
+    """
+    L = flm.shape[-2]
+    c = _axisymmetric_weight(L) * jnp.conj(psi_lm[..., :, L - 1])
+    return spin_multiplier(flm, c, spin, 0)
+
+
+def left_convolution(flm, h_lm) -> jnp.ndarray:
+    r"""
+    Left convolution of Driscoll & Healy, :math:`(f \star_{\mathrm{L}} h)(\omega) = \int f(\rho\eta)\, h(R_\rho^{-1}\omega)\, \mathrm{d}\mu(\rho)`.
+
+    Its harmonic form is :math:`2\pi\sqrt{4\pi/(2\ell+1)}\, f_{\ell m} h_{\ell 0}` (Definition III.2,
+    Eq. 16): only the zonal part of the kernel enters, and the :math:`2\pi` is the volume of
+    the fibre.  It is :math:`\mathcal{C}_{0\to0}[\Psi_h]` with :math:`\Psi_h(g) = h^*(g^{-1}\eta)`,
+    exactly (Proposition VI.4).
+
+    Args:
+        flm: Scalar harmonic coefficients, ``[..., L, 2L - 1]``.
+        h_lm: Scalar harmonic coefficients of the kernel, ``[L, 2L - 1]``.
+
+    Returns:
+        jnp.ndarray: Scalar harmonic coefficients, ``[..., L, 2L - 1]``.
+
+    """
+    L = flm.shape[-2]
+    c = 2 * np.pi * _axisymmetric_weight(L) * h_lm[..., :, L - 1]
+    return spin_multiplier(flm, c, 0, 0)
+
+
+def directional_convolution(flm, psi_lm, spin: int = 0, N: int = None) -> jnp.ndarray:
+    r"""
+    Directional convolution :math:`(f \circledast \psi)(\rho) = \langle f, \mathcal{R}_\rho\psi\rangle`, a function on the rotation group.
+
+    For scalar (Definition III.3) or spin-:math:`s` fields (Definition III.5) its Wigner
+    coefficients are :math:`\frac{8\pi^2}{2\ell+1}\, {}_sf_{\ell m}\, {}_s\psi^*_{\ell n}`
+    (Eqs. 18, 23).  It is :math:`\mathcal{C}_{s\to G}[\Lambda_s\psi]`, exactly, with no residual
+    constant (Proposition VI.1).  For a steerable kernel, azimuthally band-limited at
+    :math:`N`, the output is the tuple of :math:`2N - 1` spin-changing convolutions of
+    Proposition VI.6 (see :func:`s2conv.lifting.spin_components`).
+
+    Args:
+        flm: Spin-:math:`s` harmonic coefficients, ``[..., L, 2L - 1]``.
+        psi_lm: Spin-:math:`s` harmonic coefficients of the kernel, ``[L, 2L - 1]``.
+        spin (int, optional): Spin :math:`s`. Defaults to 0.
+        N (int, optional): Azimuthal band-limit of the output, retaining the kernel's
+            orders :math:`|n| < N`.  Defaults to :math:`L`, which is exact for any kernel.
+
+    Returns:
+        jnp.ndarray: Wigner coefficients, ``[..., 2N - 1, L, 2L - 1]``.
+
+    """
+    L = flm.shape[-2]
+    N = L if N is None else N
+    weight = 8 * np.pi**2 / (2 * np.arange(L) + 1)
+    u = (
+        jnp.conj(
+            psi_lm[..., :, L - N : L + N - 1] * lm_mask(L, spin)[:, L - N : L + N - 1]
+        )
+        * weight[:, None]
+    )
+    return spin_to_group_weights(flm, jnp.swapaxes(u, -1, -2), spin)
+
+
+def directional_synthesis(flmn, psi_lm, spin: int = 0) -> jnp.ndarray:
+    r"""
+    Haar synthesis :math:`(W \circledast^\dagger \psi)(\omega) = \int W(\rho)\, (\mathcal{R}_\rho\psi)(\omega)\, \mathrm{d}\mu(\rho)`.
+
+    Its harmonic form is :math:`\sum_n W^{\ell}_{mn}\, {}_s\psi_{\ell n}` (Definition III.10,
+    Eq. 31).  It is :math:`\mathcal{C}_{G\to s}[\widetilde{\Lambda_s\psi}]`, the adjoint of the
+    directional convolution with the same kernel (Proposition VI.2).
+
+    Args:
+        flmn: Wigner coefficients, ``[..., 2N - 1, L, 2L - 1]``.
+        psi_lm: Spin-:math:`s` harmonic coefficients of the kernel, ``[L, 2L - 1]``.
+        spin (int, optional): Spin :math:`s` of the output. Defaults to 0.
+
+    Returns:
+        jnp.ndarray: Spin-:math:`s` harmonic coefficients, ``[..., L, 2L - 1]``.
+
+    """
+    N, L = azimuthal_bandlimit(flmn), flmn.shape[-2]
+    v = (psi_lm * lm_mask(L, spin))[..., :, L - N : L + N - 1]
+    return group_to_spin_weights(flmn, jnp.swapaxes(v, -1, -2), spin)
+
+
+def axisymmetric_synthesis(flm, psi_lm, spin: int = 0) -> jnp.ndarray:
+    r"""
+    Axisymmetric synthesis :math:`\int W(\omega')\, (\mathcal{R}_{\omega'}\psi)(\omega)\, \mathrm{d}\mu(\omega')`.
+
+    The adjoint of :func:`axisymmetric_convolution`, mapping a scalar field to a spin-:math:`s`
+    field, :math:`\sqrt{4\pi/(2\ell+1)}\, W_{\ell m}\, {}_s\psi_{\ell 0}`: the first term of the
+    spin wavelet synthesis (Eq. 26).
+
+    Args:
+        flm: Scalar harmonic coefficients, ``[..., L, 2L - 1]``.
+        psi_lm: Spin-:math:`s` harmonic coefficients of the kernel, ``[L, 2L - 1]``.
+        spin (int, optional): Spin :math:`s` of the output. Defaults to 0.
+
+    Returns:
+        jnp.ndarray: Spin-:math:`s` harmonic coefficients, ``[..., L, 2L - 1]``.
+
+    """
+    L = flm.shape[-2]
+    c = _axisymmetric_weight(L) * psi_lm[..., :, L - 1]
+    return spin_multiplier(flm, c, 0, spin)
+
+
+spin_changing_convolution = spin_multiplier
+r"""Spin-changing convolution with multiplier :math:`c_\ell` (Proposition VI.5); see :func:`s2conv.convolutions.lifted.spin_multiplier`."""

s2conv/convolutions/diagonal.py

@@ -0,0 +1,41 @@
+r"""
+Operators diagonal in the spherical harmonic basis (Sec. VII B).
+
+Harmonic multiplication of Kennedy et al. and the sifting convolution of Roddy & McEwen
+multiply each coefficient by :math:`g_{\ell m}` or :math:`g^*_{\ell m}` (Definition III.9).
+They commute with every polar rotation, and are equivariant under all rotations exactly
+when :math:`g_{\ell m}` does not depend on :math:`m`, in which case they are the lifted
+convolution :math:`\mathcal{C}_{0\to0}[\Psi_T]` with :math:`c_\ell = g_{\ell0}` (Proposition VII.2).
+"""
+
+import jax.numpy as jnp
+
+
+def harmonic_multiplication(flm, glm) -> jnp.ndarray:
+    r"""
+    Harmonic multiplication :math:`(g \odot_{\mathrm{K}} f)_{\ell m} = g_{\ell m} f_{\ell m}` (Definition III.9).
+
+    Args:
+        flm: Scalar harmonic coefficients, ``[..., L, 2L - 1]``.
+        glm: Scalar harmonic coefficients of the filter, ``[..., L, 2L - 1]``.
+
+    Returns:
+        jnp.ndarray: Scalar harmonic coefficients, ``[..., L, 2L - 1]``.
+
+    """
+    return glm * flm
+
+
+def sifting_convolution(flm, glm) -> jnp.ndarray:
+    r"""
+    Sifting convolution :math:`(f \circledcirc g)_{\ell m} = f_{\ell m} g^*_{\ell m}` (Definition III.9).
+
+    Args:
+        flm: Scalar harmonic coefficients, ``[..., L, 2L - 1]``.
+        glm: Scalar harmonic coefficients of the filter, ``[..., L, 2L - 1]``.
+
+    Returns:
+        jnp.ndarray: Scalar harmonic coefficients, ``[..., L, 2L - 1]``.
+
+    """
+    return flm * jnp.conj(glm)

s2conv/convolutions/lifted.py

@@ -0,0 +1,221 @@
+r"""
+The generalised lifted spherical convolution (Sec. V).
+
+For a kernel :math:`\Psi` on the rotation group and spins :math:`s, s'`, the four operators
+of Definition V.1 are
+
+.. math::
+
+    \mathcal{C}_{s\to s'}[\Psi] = \Lambda_{s'}^\dagger (\Psi \star \Lambda_s\,\cdot\,), \quad
+    \mathcal{C}_{s\to G}[\Psi] = \Psi \star \Lambda_s\,\cdot\,, \quad
+    \mathcal{C}_{G\to s'}[\Psi] = \Lambda_{s'}^\dagger (\Psi \star\,\cdot\,), \quad
+    \mathcal{C}_{G\to G}[\Psi] = \Psi \star\,\cdot\,,
+
+with :math:`\star` the group convolution in correlation form,
+:math:`(\Psi \star F)(\rho) = \int \Psi^*(\rho^{-1}\rho') F(\rho')\, \mathrm{d}\mu(\rho')`
+(Definition III.4).  In harmonic space they read a single entry, column, row or the
+whole block of the kernel (Theorem V.2, Eqs. 42–45):
+
+.. math::
+
+    (\mathcal{C}_{s\to G}[\Psi]\, {}_sf)^{\ell}_{mn} &= (-1)^s \sqrt{\tfrac{8\pi^2}{2\ell+1}}\, {}_sf_{\ell m}\, \Psi^{\ell *}_{n,-s}, \\
+    (\mathcal{C}_{s\to s'}[\Psi]\, {}_sf)_{\ell m} &= (-1)^{s+s'}\, {}_sf_{\ell m}\, \Psi^{\ell *}_{-s',-s}, \\
+    (\mathcal{C}_{G\to s'}[\Psi]\, F)_{\ell m} &= (-1)^{s'} \sqrt{\tfrac{2\ell+1}{8\pi^2}} \sum_{m'} F^{\ell}_{mm'}\, \Psi^{\ell *}_{-s',m'}, \\
+    (\mathcal{C}_{G\to G}[\Psi]\, F)^{\ell}_{mn} &= \sum_{m'} F^{\ell}_{mm'}\, \Psi^{\ell *}_{nm'} .
+
+Every bounded rotation-equivariant map of the four types is of this form, and acts
+through compact weights (Theorem V.5): a multiplier :math:`c_\ell`, a column
+:math:`u^{\ell}_n`, a row :math:`v^{\ell}_k` or a block :math:`B^{\ell}`.  This module applies
+the operators in both parameterisations; :mod:`s2conv.kernels` converts between them.
+All functions act on harmonic coefficients in the s2fft layout, broadcast over leading
+axes, and are differentiable and JIT compilable (with spins and band-limits static).
+"""
+
+import jax.numpy as jnp
+
+from s2conv import kernels
+from s2conv.utils.indexing import azimuthal_bandlimit, lm_mask, wigner_mask
+
+# ------------------------------------------------------------------------------ weight form
+
+
+def spin_multiplier(flm, c, s_in: int, s_out: int) -> jnp.ndarray:
+    r"""
+    A spin-changing multiplier, :math:`({}_{s'}g)_{\ell m} = c_\ell\, {}_sf_{\ell m}` (Theorem V.5).
+
+    This is the general bounded rotation-equivariant map from spin :math:`s` to spin
+    :math:`s'`; degrees :math:`\ell < \max(|s|, |s'|)` are absent from the target and are
+    annihilated.
+
+    Args:
+        flm: Spin-:math:`s` harmonic coefficients, ``[..., L, 2L - 1]``.
+        c: Multiplier :math:`c_\ell`, ``[..., L]``.
+        s_in (int): Input spin :math:`s`.
+        s_out (int): Output spin :math:`s'`.
+
+    Returns:
+        jnp.ndarray: Spin-:math:`s'` harmonic coefficients, ``[..., L, 2L - 1]``.
+
+    """
+    L = flm.shape[-2]
+    return jnp.asarray(c)[..., :, None] * flm * lm_mask(L, max(abs(s_in), abs(s_out)))
+
+
+def spin_to_group_weights(flm, u, s_in: int) -> jnp.ndarray:
+    r"""
+    An equivariant map from spin :math:`s` to the rotation group, :math:`F^{\ell}_{mn} = u^{\ell}_n\, {}_sf_{\ell m}` (Theorem V.5(c)).
+
+    Args:
+        flm: Spin-:math:`s` harmonic coefficients, ``[..., L, 2L - 1]``.
+        u: Column weights ``[..., 2N - 1, L]``.
+        s_in (int): Input spin :math:`s`.
+
+    Returns:
+        jnp.ndarray: Wigner coefficients, ``[..., 2N - 1, L, 2L - 1]``.
+
+    """
+    L, N = flm.shape[-2], (u.shape[-2] + 1) // 2
+    return (
+        u[..., :, :, None]
+        * (flm * lm_mask(L, s_in))[..., None, :, :]
+        * wigner_mask(L, N)
+    )
+
+
+def group_to_spin_weights(flmn, v, s_out: int) -> jnp.ndarray:
+    r"""
+    An equivariant map from the rotation group to spin :math:`s'`, :math:`g_{\ell m} = \sum_k v^{\ell}_k F^{\ell}_{mk}` (Theorem V.5(c)).
+
+    Args:
+        flmn: Wigner coefficients, ``[..., 2N - 1, L, 2L - 1]``.
+        v: Row weights ``[..., 2N - 1, L]``.
+        s_out (int): Output spin :math:`s'`.
+
+    Returns:
+        jnp.ndarray: Spin-:math:`s'` harmonic coefficients, ``[..., L, 2L - 1]``.
+
+    """
+    L = flmn.shape[-2]
+    return jnp.einsum("...kl,...klm->...lm", v, flmn) * lm_mask(L, s_out)
+
+
+def group_to_group_weights(flmn, B) -> jnp.ndarray:
+    r"""
+    An equivariant map on the rotation group, :math:`G^{\ell} = F^{\ell} B^{\ell}` (Theorem V.5(c)).
+
+    Args:
+        flmn: Wigner coefficients, ``[..., 2N_in - 1, L, 2L - 1]``.
+        B: Block weights ``[..., L, 2N_in - 1, 2N_out - 1]``.
+
+    Returns:
+        jnp.ndarray: Wigner coefficients, ``[..., 2N_out - 1, L, 2L - 1]``.
+
+    """
+    L, N_out = flmn.shape[-2], (B.shape[-1] + 1) // 2
+    return jnp.einsum("...klm,...lkn->...nlm", flmn, B) * wigner_mask(L, N_out)
+
+
+# ------------------------------------------------------------------------------ kernel form
+
+
+def spin_to_spin(flm, kernel, s_in: int, s_out: int) -> jnp.ndarray:
+    r"""
+    :math:`\mathcal{C}_{s\to s'}[\Psi]`, which reads the single entry :math:`\Psi^{\ell}_{-s',-s}` (Eq. 43).
+
+    Args:
+        flm: Spin-:math:`s` harmonic coefficients, ``[..., L, 2L - 1]``.
+        kernel: The kernel :math:`\Psi`, ``[2N - 1, L, 2L - 1]``.
+        s_in (int): Input spin :math:`s`.
+        s_out (int): Output spin :math:`s'`.
+
+    Returns:
+        jnp.ndarray: Spin-:math:`s'` harmonic coefficients, ``[..., L, 2L - 1]``.
+
+    """
+    return spin_multiplier(flm, kernels.to_multiplier(kernel, s_in, s_out), s_in, s_out)
+
+
+def spin_to_group(flm, kernel, s_in: int, N_out: int = None) -> jnp.ndarray:
+    r"""
+    :math:`\mathcal{C}_{s\to G}[\Psi]`, which reads the column :math:`k = -s` (Eq. 42).
+
+    Args:
+        flm: Spin-:math:`s` harmonic coefficients, ``[..., L, 2L - 1]``.
+        kernel: The kernel :math:`\Psi`, ``[2N - 1, L, 2L - 1]``.
+        s_in (int): Input spin :math:`s`.
+        N_out (int, optional): Azimuthal band-limit of the output.  Defaults to :math:`L`.
+
+    Returns:
+        jnp.ndarray: Wigner coefficients, ``[..., 2N_out - 1, L, 2L - 1]``.
+
+    """
+    N_out = flm.shape[-2] if N_out is None else N_out
+    return spin_to_group_weights(flm, kernels.to_column(kernel, s_in, N_out), s_in)
+
+
+def group_to_spin(flmn, kernel, s_out: int) -> jnp.ndarray:
+    r"""
+    :math:`\mathcal{C}_{G\to s'}[\Psi]`, which reads the row :math:`n = -s'` (Eq. 44).
+
+    Args:
+        flmn: Wigner coefficients, ``[..., 2N - 1, L, 2L - 1]``.
+        kernel: The kernel :math:`\Psi`, ``[2N_k - 1, L, 2L - 1]``.
+        s_out (int): Output spin :math:`s'`.
+
+    Returns:
+        jnp.ndarray: Spin-:math:`s'` harmonic coefficients, ``[..., L, 2L - 1]``.
+
+    """
+    v = kernels.to_row(kernel, s_out, azimuthal_bandlimit(flmn))
+    return group_to_spin_weights(flmn, v, s_out)
+
+
+def group_to_group(flmn, kernel, N_out: int = None) -> jnp.ndarray:
+    r"""
+    :math:`\mathcal{C}_{G\to G}[\Psi] = \Psi \star\,\cdot\,`, the group convolution (Eqs. 20, 45).
+
+    Args:
+        flmn: Wigner coefficients, ``[..., 2N_in - 1, L, 2L - 1]``.
+        kernel: The kernel :math:`\Psi`, ``[2N_k - 1, L, 2L - 1]``.
+        N_out (int, optional): Azimuthal band-limit of the output.  Defaults to :math:`L`.
+
+    Returns:
+        jnp.ndarray: Wigner coefficients, ``[..., 2N_out - 1, L, 2L - 1]``.
+
+    """
+    N_out = flmn.shape[-2] if N_out is None else N_out
+    B = kernels.to_block(kernel, azimuthal_bandlimit(flmn), N_out)
+    return group_to_group_weights(flmn, B)
+
+
+group_convolution = group_to_group
+
+
+def lifted_convolution(x, kernel, s_in, s_out, N_out: int = None) -> jnp.ndarray:
+    r"""
+    The generalised lifted convolution :math:`\mathcal{C}_{s\to s'}[\Psi]` of Definition V.1.
+
+    Args:
+        x: Spin-:math:`s` harmonic coefficients ``[..., L, 2L - 1]`` or, if
+            ``s_in == "G"``, Wigner coefficients ``[..., 2N - 1, L, 2L - 1]``.
+        kernel: The kernel :math:`\Psi`, ``[2N_k - 1, L, 2L - 1]``.
+        s_in (int or "G"): Input spin, or ``"G"`` for an input on the rotation group.
+        s_out (int or "G"): Output spin, or ``"G"`` for an output on the rotation group.
+        N_out (int, optional): Azimuthal band-limit of an output on the rotation group.
+            Defaults to :math:`L`.
+
+    Returns:
+        jnp.ndarray: Harmonic or Wigner coefficients of the output.
+
+    """
+    if s_in == "G":
+        return (
+            group_to_group(x, kernel, N_out)
+            if s_out == "G"
+            else group_to_spin(x, kernel, s_out)
+        )
+    return (
+        spin_to_group(x, kernel, s_in, N_out)
+        if s_out == "G"
+        else spin_to_spin(x, kernel, s_in, s_out)
+    )