Kernels

Kernels of the generalised lifted convolution and their compact weights.

A kernel \(\Psi\) is a function on the rotation group, stored by its Wigner coefficients in the s2fft layout, kernel[..., N - 1 + k, el, L - 1 + n] holding \(\Psi^{\ell}_{nk}\), with shape [..., 2N - 1, L, 2L - 1]. Its azimuthal band-limit \(N\) bounds the second index \(k\), the one along which a lifted spin-\(s\) input is supported (\(k = -s\)).

Each operator reads only part of the kernel block (Theorem V.2, Fig. 5): the entry \((-s', -s)\) for \(\mathcal{C}_{s\to s'}\), the column \(k = -s\) for \(\mathcal{C}_{s\to G}\), the row \(n = -s'\) for \(\mathcal{C}_{G\to s'}\) and the whole block for \(\mathcal{C}_{G\to G}\). The corresponding compact weights of Theorem V.5 are, with \(a_\ell = (2\ell+1)/(8\pi^2)\),

\[c_\ell = (-1)^{s+s'} \Psi^{\ell *}_{-s',-s}, \quad u^{\ell}_n = (-1)^s a_\ell^{-1/2} \Psi^{\ell *}_{n,-s}, \quad v^{\ell}_k = (-1)^{s'} a_\ell^{1/2} \Psi^{\ell *}_{-s',k}, \quad B^{\ell} = \Psi^{\ell\dagger},\]

stored as c[..., el], u[..., N - 1 + n, el], v[..., N - 1 + k, el] and B[..., el, N_in - 1 + k, N_out - 1 + n].

s2conv.kernels.lifted(psi_lm, spin: int, N: int = None) → Array

The lifted spherical kernel \(\Lambda_s \psi\) (Lemma IV.7(c)).

A spin-\(s\) kernel on the sphere lifted like a signal occupies the column \(k = -s\), \((\Lambda_s\psi)^{\ell}_{n,-s} = (-1)^s \sqrt{8\pi^2/(2\ell+1)}\, {}_s\psi_{\ell n}\). With \(\mathcal{C}_{s\to G}\) it gives the directional convolutions (Proposition VI.1) and, for axisymmetric \(\psi\), the axisymmetric convolutions (Proposition VI.3).

Parameters:

• psi_lm – Spin-\(s\) harmonic coefficients of the kernel, [..., L, 2L - 1].

• spin (int) – Spin \(s\).

• N (int, optional) – Azimuthal band-limit of the kernel. Defaults to \(|s| + 1\).

Returns:

The kernel, [..., 2N - 1, L, 2L - 1].

Return type:

jnp.ndarray

s2conv.kernels.flipped(kernel, N: int = None) → Array

The flipped kernel \(\widetilde\Psi(g) = \Psi^*(g^{-1})\), with blocks \(\Psi^{\ell\dagger}\).

It appears in the adjoints of Theorem V.2; for instance Haar synthesis is \(\mathcal{C}_{G\to s}[\widetilde{\Lambda_s\psi}]\) (Proposition VI.2).

Parameters:

• kernel – The kernel \(\Psi\), [..., 2N_k - 1, L, 2L - 1].

• N (int, optional) – Azimuthal band-limit of the flipped kernel, which holds the first index of \(\Psi\). Defaults to \(L\), which is exact.

Returns:

The kernel \(\widetilde\Psi\), [..., 2N - 1, L, 2L - 1].

Return type:

jnp.ndarray

s2conv.kernels.left(h_lm, N: int = None) → Array

The kernel \(\Psi_h(g) = h^*(g^{-1}\eta)\) of the left convolution (Proposition VI.4).

It is invariant under left multiplication by the stabiliser, so its block has the single row \(n = 0\), \((\Psi_h)^{\ell}_{0k} = 2\pi \sqrt{4\pi/(2\ell+1)}\, h^*_{\ell k}\).

Parameters:

• h_lm – Harmonic coefficients of the scalar kernel \(h\), [..., L, 2L - 1].

• N (int, optional) – Azimuthal band-limit of the kernel. Defaults to \(L\), which is exact; \(\mathcal{C}_{0\to0}\) reads only \(k = 0\).

Returns:

The kernel, [..., 2N - 1, L, 2L - 1].

Return type:

jnp.ndarray

s2conv.kernels.multiplier(c, s_in: int, s_out: int, N: int = None) → Array

The canonical kernel of a spin-changing multiplier (Theorem V.5(a)).

Every bounded rotation-equivariant map from spin \(s\) to spin \(s'\) multiplies by \(c_\ell\) and is \(\mathcal{C}_{s\to s'}[\Psi_T]\) with the single entry \((\Psi_T)^{\ell}_{-s',-s} = (-1)^{s+s'} c^*_\ell\) (Proposition VI.5).

Parameters:

• c – Multiplier \(c_\ell\), [..., L].

• s_in (int) – Input spin \(s\).

• s_out (int) – Output spin \(s'\).

• N (int, optional) – Azimuthal band-limit of the kernel. Defaults to \(|s| + 1\).

Returns:

The kernel, [..., 2N - 1, L, 2L - 1].

Return type:

jnp.ndarray

s2conv.kernels.identity(L: int, spin: int, N: int = None) → Array

The single entry \(\Psi^{\ell}_{-s,-s} = 1\), the kernel of the lift and projection.

It gives \(\mathcal{C}_{s\to G} = \Lambda_s\), \(\mathcal{C}_{G\to s} = \Lambda_s^\dagger\) and \(\mathcal{C}_{s\to s} = \mathrm{Id}\) (Corollary V.6). It is a distribution, not a square-integrable kernel (Theorem V.5(b)).

Parameters:

• L (int) – Harmonic band-limit.

• spin (int) – Spin \(s\).

• N (int, optional) – Azimuthal band-limit of the kernel. Defaults to \(|s| + 1\).

Returns:

The kernel, [2N - 1, L, 2L - 1].

Return type:

jnp.ndarray

s2conv.kernels.steerable(multipliers: Mapping, s_in: int, N: int = None) → Array

A kernel for a tuple of spin-changing convolutions (Proposition VI.6).

With \(\mathcal{C}_{s\to G}\) the column \(k = -s\) supported on \(|n| < N\) produces the \(2N - 1\) spin components \(\mathcal{C}_{s\to s'}\), \(|s'| < N\); this constructs that column from their multipliers.

Parameters:

• multipliers (Mapping) – Map from output spin \(s'\) to multiplier \(c_\ell\).

• s_in (int) – Input spin \(s\).

• N (int, optional) – Azimuthal band-limit of the kernel. Defaults to \(|s| + 1\).

Returns:

The kernel, [2N - 1, L, 2L - 1].

Return type:

jnp.ndarray

s2conv.kernels.to_multiplier(kernel, s_in: int, s_out: int) → Array

The multiplier \(c_\ell = (-1)^{s+s'} \Psi^{\ell *}_{-s',-s}\) of \(\mathcal{C}_{s\to s'}[\Psi]\) (Eq. 43).

Parameters:

• kernel – The kernel, [..., 2N - 1, L, 2L - 1].

• s_in (int) – Input spin \(s\).

• s_out (int) – Output spin \(s'\).

Returns:

\(c_\ell\), [..., L], zero for \(\ell < \max(|s|, |s'|)\).

Return type:

jnp.ndarray

s2conv.kernels.to_column(kernel, s_in: int, N_out: int) → Array

The weights \(u^{\ell}_n = (-1)^s a_\ell^{-1/2} \Psi^{\ell *}_{n,-s}\) of \(\mathcal{C}_{s\to G}[\Psi]\) (Eq. 53).

Parameters:

• kernel – The kernel, [..., 2N - 1, L, 2L - 1].

• s_in (int) – Input spin \(s\).

• N_out (int) – Azimuthal band-limit of the output.

Returns:

\(u\), [..., 2N_out - 1, L].

Return type:

jnp.ndarray

s2conv.kernels.to_row(kernel, s_out: int, N_in: int) → Array

The weights \(v^{\ell}_k = (-1)^{s'} a_\ell^{1/2} \Psi^{\ell *}_{-s',k}\) of \(\mathcal{C}_{G\to s'}[\Psi]\) (Eq. 53).

Parameters:

• kernel – The kernel, [..., 2N - 1, L, 2L - 1].

• s_out (int) – Output spin \(s'\).

• N_in (int) – Azimuthal band-limit of the input.

Returns:

\(v\), [..., 2N_in - 1, L].

Return type:

jnp.ndarray

s2conv.kernels.to_block(kernel, N_in: int, N_out: int) → Array

The weights \(B^{\ell} = \Psi^{\ell\dagger}\) of \(\mathcal{C}_{G\to G}[\Psi]\) (Eq. 53).

Parameters:

• kernel – The kernel, [..., 2N - 1, L, 2L - 1].

• N_in (int) – Azimuthal band-limit of the input.

• N_out (int) – Azimuthal band-limit of the output.

Returns:

\(B\), [..., L, 2N_in - 1, 2N_out - 1], with B[..., el, N_in - 1 + k, N_out - 1 + n] \(= \Psi^{\ell *}_{nk}\).

Return type:

jnp.ndarray

s2conv.kernels.from_multiplier(c, s_in: int, s_out: int, N: int = None) → Array

Alias of multiplier(): the canonical kernel of a multiplier.

s2conv.kernels.from_column(u, s_in: int, N: int = None) → Array

The canonical kernel realising column weights \(u\) (Theorem V.5(c)).

Parameters:

• u – Weights [..., 2N_out - 1, L].

• s_in (int) – Input spin \(s\).

• N (int, optional) – Azimuthal band-limit of the kernel. Defaults to \(|s| + 1\).

Returns:

A kernel whose only nonzero column is \(k = -s\).

Return type:

jnp.ndarray

s2conv.kernels.from_row(v, s_out: int) → Array

The canonical kernel realising row weights \(v\) (Theorem V.5(c)).

Parameters:

• v – Weights [..., 2N_in - 1, L].

• s_out (int) – Output spin \(s'\).

Returns:

A kernel [..., 2N_in - 1, L, 2L - 1] whose only nonzero row is \(n = -s'\).

Return type:

jnp.ndarray

s2conv.kernels.from_block(B) → Array

The kernel \(\Psi^{\ell} = B^{\ell\dagger}\) realising block weights (Theorem V.5(c)).

Parameters:

B – Weights [..., L, 2N_in - 1, 2N_out - 1].

Returns:

The kernel [..., 2N_in - 1, L, 2L - 1].

Return type:

jnp.ndarray

s2conv.kernels.compose(kernel2, kernel1) → Array

The kernel of a composition, \(\Psi_3^{\ell} = \Psi_2^{\ell} \Psi_1^{\ell}\) (Proposition V.3).

\(\mathcal{C}_{G\to G}[\Psi_2]\, \mathcal{C}_{G\to G}[\Psi_1] = \mathcal{C}_{G\to G}[\Psi_3]\), and likewise with a lift or projection at either end. The product runs over the intermediate index, restricted to the azimuthal band-limit of kernel2.

Parameters:

• kernel2 – The kernel applied second, [..., 2N_2 - 1, L, 2L - 1].

• kernel1 – The kernel applied first, [..., 2N_1 - 1, L, 2L - 1].

Returns:

\(\Psi_3\), [..., 2N_1 - 1, L, 2L - 1].

Return type:

jnp.ndarray