Lifted convolutions

The generalised lifted spherical convolution (Sec. V).

For a kernel \(\Psi\) on the rotation group and spins \(s, s'\), the four operators of Definition V.1 are

\[\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 \(\star\) the group convolution in correlation form, \((\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):

\[\begin{split}(\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'} .\end{split}\]

Every bounded rotation-equivariant map of the four types is of this form, and acts through compact weights (Theorem V.5): a multiplier \(c_\ell\), a column \(u^{\ell}_n\), a row \(v^{\ell}_k\) or a block \(B^{\ell}\). This module applies the operators in both parameterisations; 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).

s2conv.convolutions.lifted.spin_multiplier(flm, c, s_in: int, s_out: int) → Array

A spin-changing multiplier, \(({}_{s'}g)_{\ell m} = c_\ell\, {}_sf_{\ell m}\) (Theorem V.5).

This is the general bounded rotation-equivariant map from spin \(s\) to spin \(s'\); degrees \(\ell < \max(|s|, |s'|)\) are absent from the target and are annihilated.

Parameters:

• flm – Spin-\(s\) harmonic coefficients, [..., L, 2L - 1].

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

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

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

Returns:

Spin-\(s'\) harmonic coefficients, [..., L, 2L - 1].

Return type:

jnp.ndarray

s2conv.convolutions.lifted.spin_to_group_weights(flm, u, s_in: int) → Array

An equivariant map from spin \(s\) to the rotation group, \(F^{\ell}_{mn} = u^{\ell}_n\, {}_sf_{\ell m}\) (Theorem V.5(c)).

Parameters:

• flm – Spin-\(s\) harmonic coefficients, [..., L, 2L - 1].

• u – Column weights [..., 2N - 1, L].

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

Returns:

Wigner coefficients, [..., 2N - 1, L, 2L - 1].

Return type:

jnp.ndarray

s2conv.convolutions.lifted.group_to_spin_weights(flmn, v, s_out: int) → Array

An equivariant map from the rotation group to spin \(s'\), \(g_{\ell m} = \sum_k v^{\ell}_k F^{\ell}_{mk}\) (Theorem V.5(c)).

Parameters:

• flmn – Wigner coefficients, [..., 2N - 1, L, 2L - 1].

• v – Row weights [..., 2N - 1, L].

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

Returns:

Spin-\(s'\) harmonic coefficients, [..., L, 2L - 1].

Return type:

jnp.ndarray

s2conv.convolutions.lifted.group_to_group_weights(flmn, B) → Array

An equivariant map on the rotation group, \(G^{\ell} = F^{\ell} B^{\ell}\) (Theorem V.5(c)).

Parameters:

• flmn – Wigner coefficients, [..., 2N_in - 1, L, 2L - 1].

• B – Block weights [..., L, 2N_in - 1, 2N_out - 1].

Returns:

Wigner coefficients, [..., 2N_out - 1, L, 2L - 1].

Return type:

jnp.ndarray

s2conv.convolutions.lifted.spin_to_spin(flm, kernel, s_in: int, s_out: int) → Array

\(\mathcal{C}_{s\to s'}[\Psi]\), which reads the single entry \(\Psi^{\ell}_{-s',-s}\) (Eq. 43).

Parameters:

• flm – Spin-\(s\) harmonic coefficients, [..., L, 2L - 1].

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

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

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

Returns:

Spin-\(s'\) harmonic coefficients, [..., L, 2L - 1].

Return type:

jnp.ndarray

s2conv.convolutions.lifted.spin_to_group(flm, kernel, s_in: int, N_out: int = None) → Array

\(\mathcal{C}_{s\to G}[\Psi]\), which reads the column \(k = -s\) (Eq. 42).

Parameters:

• flm – Spin-\(s\) harmonic coefficients, [..., L, 2L - 1].

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

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

• N_out (int, optional) – Azimuthal band-limit of the output. Defaults to \(L\).

Returns:

Wigner coefficients, [..., 2N_out - 1, L, 2L - 1].

Return type:

jnp.ndarray

s2conv.convolutions.lifted.group_to_spin(flmn, kernel, s_out: int) → Array

\(\mathcal{C}_{G\to s'}[\Psi]\), which reads the row \(n = -s'\) (Eq. 44).

Parameters:

• flmn – Wigner coefficients, [..., 2N - 1, L, 2L - 1].

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

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

Returns:

Spin-\(s'\) harmonic coefficients, [..., L, 2L - 1].

Return type:

jnp.ndarray

s2conv.convolutions.lifted.group_to_group(flmn, kernel, N_out: int = None) → Array

\(\mathcal{C}_{G\to G}[\Psi] = \Psi \star\,\cdot\,\), the group convolution (Eqs. 20, 45).

Parameters:

• flmn – Wigner coefficients, [..., 2N_in - 1, L, 2L - 1].

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

• N_out (int, optional) – Azimuthal band-limit of the output. Defaults to \(L\).

Returns:

Wigner coefficients, [..., 2N_out - 1, L, 2L - 1].

Return type:

jnp.ndarray

s2conv.convolutions.lifted.group_convolution(flmn, kernel, N_out: int = None) → Array

\(\mathcal{C}_{G\to G}[\Psi] = \Psi \star\,\cdot\,\), the group convolution (Eqs. 20, 45).

Parameters:

• flmn – Wigner coefficients, [..., 2N_in - 1, L, 2L - 1].

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

• N_out (int, optional) – Azimuthal band-limit of the output. Defaults to \(L\).

Returns:

Wigner coefficients, [..., 2N_out - 1, L, 2L - 1].

Return type:

jnp.ndarray

s2conv.convolutions.lifted.lifted_convolution(x, kernel, s_in, s_out, N_out: int = None) → Array

The generalised lifted convolution \(\mathcal{C}_{s\to s'}[\Psi]\) of Definition V.1.

Parameters:

• x – Spin-\(s\) harmonic coefficients [..., L, 2L - 1] or, if s_in == "G", Wigner coefficients [..., 2N - 1, L, 2L - 1].

• kernel – The kernel \(\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 \(L\).

Returns:

Harmonic or Wigner coefficients of the output.

Return type:

jnp.ndarray