Equivariance errors#

Equivariance errors under random rotations, following the protocols of the literature.

For an operator \(T\), a signal \(f\) and a rotation \(\mathcal{R}\), the equivariance defect is \(\lVert T(\mathcal{R}f) - \mathcal{R}(Tf)\rVert\), the \(L^2\) distance between rotating the input and rotating the output. It is reported relative to either of two norms:

  • the rotated output, \(\lVert\mathcal{R}(Tf)\rVert = \lVert Tf\rVert\), which does not depend on the rotation (rotated_output, the default in this package); or

  • the output of the rotated input, \(\lVert T(\mathcal{R}f)\rVert\), as in Cobb et al. (2021) and Ocampo et al. (2023) (output_of_rotated_input).

The two agree for an equivariant operator, and differ at relative order of the error itself otherwise. Averaged over random signals and rotations, for instance 100 of each as in Cobb et al., they give the mean relative equivariance error

\[\epsilon(T) = \frac{1}{N_f N_g} \sum_{i=1}^{N_f} \sum_{j=1}^{N_g} \frac{\lVert T(\mathcal{R}_{g_j} f_i) - \mathcal{R}_{g_j}(T f_i)\rVert}{\lVert \mathcal{R}_{g_j}(T f_i)\rVert}.\]

equivariance_errors() handles operators between harmonic coefficients: spin spherical harmonic coefficients [L, 2L - 1] or Wigner coefficients [2N - 1, L, 2L - 1]. Rotations act on both through the Wigner matrices, on the first Wigner index for functions on the rotation group (left translation), and the norms follow from Parseval’s relation, exactly for band-limited outputs. equivariance_errors_on_points() handles operators whose outputs are not band-limited, such as evaluations on a section, by comparing the outputs in real space on a Gauss-Legendre grid. Both are JIT compiled and vectorised over rotations.

class s2conv.utils.equivariance.EquivarianceErrors(rotated_output: Array, output_of_rotated_input: Array)[source]#

Relative equivariance errors, [n_signals, n_rotations], in both normalisations.

rotated_output holds \(\lVert T(\mathcal{R}f) - \mathcal{R}(Tf)\rVert / \lVert\mathcal{R}(Tf)\rVert\), and output_of_rotated_input holds \(\lVert T(\mathcal{R}f) - \mathcal{R}(Tf)\rVert / \lVert T(\mathcal{R}f)\rVert\), the normalisation of Cobb et al. (2021) and Ocampo et al. (2023).

rotated_output: Array#

Alias for field number 0

output_of_rotated_input: Array#

Alias for field number 1

s2conv.utils.equivariance.random_signals(key, n: int, L: int, spin: int = 0, N: int = None) → Array[source]#

Random band-limited signals with standard normal harmonic coefficients.

The real and imaginary parts of each coefficient are independent standard normal variables; coefficients outside the band-limit, or with \(\ell < |s|\), are zero.

Parameters:

  • key – A JAX random key.

  • n (int) – Number of signals.

  • L (int) – Harmonic band-limit.

  • spin (int, optional) – Spin \(s\) of fields on the sphere. Defaults to 0.

  • N (int, optional) – If given, draw functions on the rotation group with azimuthal band-limit \(N\) instead of fields on the sphere.

Returns:

Harmonic coefficients [n, L, 2L - 1] or Wigner coefficients [n, 2N - 1, L, 2L - 1].

Return type:

jnp.ndarray

s2conv.utils.equivariance.random_rotations(key, n: int) → tuple[source]#

Euler angles of rotations drawn from the Haar measure on the rotation group.

Parameters:

  • key – A JAX random key.

  • n (int) – Number of rotations.

Returns:

Arrays \((\alpha, \beta, \gamma)\), each of shape [n].

Return type:

tuple

s2conv.utils.equivariance.quadrature(L: int) → tuple[source]#

Gauss-Legendre nodes on the sphere and their weights, flattened.

The \(L\) nodes in \(\cos\theta\) and \(2L - 1\) equispaced longitudes, the Gauss-Legendre sampling of s2fft, integrate \(|g|^2\) exactly for a spin field \(g\) of band-limit \(L\); the nodes avoid the poles.

Parameters:

L (int) – Band-limit up to which the squared modulus is integrated exactly.

Returns:

Colatitudes, longitudes and weights, each of shape [L (2L - 1)].

Return type:

tuple

s2conv.utils.equivariance.equivariance_errors(operator: Callable, signals, rotations, batch_size: int = 100) → EquivarianceErrors[source]#

Relative equivariance errors of an operator between harmonic coefficients.

For every signal \(f_i\) and rotation \(g_j\), compares \(T(\mathcal{R}_{g_j} f_i)\) with \(\mathcal{R}_{g_j}(T f_i)\), rotating harmonic coefficients by the Wigner matrices (Eq. 10), Wigner coefficients by left translation, and measuring both in \(L^2\) by Parseval’s relation, \(\sum_{\ell m} |f_{\ell m}|^2\) on the sphere and \(\sum_\ell \frac{2\ell+1}{8\pi^2} \sum_{mn} |F^{\ell}_{mn}|^2\) on the rotation group.

Parameters:

  • operator – Function mapping one input, [L, 2L - 1] or [2N - 1, L, 2L - 1], to one output of either form. It must be a JAX function of its input.

  • signals – Inputs stacked along a leading axis, as from random_signals().

  • rotations – Euler angles \((\alpha, \beta, \gamma)\), arrays of equal length, as from random_rotations().

  • batch_size (int, optional) – Number of rotations evaluated together, which bounds memory. Defaults to 100.

Returns:

The errors, [n_signals, n_rotations], in both normalisations.

Return type:

EquivarianceErrors

s2conv.utils.equivariance.equivariance_errors_on_points(evaluator: Callable, signals, rotations, spin_out: int = 0, L_quad: int = None, batch_size: int = 10) → EquivarianceErrors[source]#

Relative equivariance errors of an operator with outputs on the sphere, compared in real space.

For outputs that are not band-limited, such as evaluations on a section (Proposition VII.1), compares \(T(\mathcal{R}f)(\omega)\) with \((\mathcal{R}\,Tf)(\omega) = e^{-\mathrm{i}s'\vartheta}(Tf)(R^{-1}\omega)\) (Eq. 12), with \(\vartheta\) the third Euler angle of \(R^{-1}s(\omega)\), and integrates on a Gauss-Legendre grid (quadrature()).

Parameters:

  • evaluator – Function mapping one input [L, 2L - 1] to a function of (theta, phi) returning the output, a spin-\(s'\) field in the coordinate frame, at those points. It must be written with JAX.

  • signals – Inputs [n, L, 2L - 1].

  • rotations – Euler angles \((\alpha, \beta, \gamma)\), arrays of equal length.

  • spin_out (int, optional) – Spin \(s'\) of the output. Defaults to 0.

  • L_quad (int, optional) – Band-limit of the quadrature. Defaults to \(2L\), which oversamples band-limited outputs by two.

  • batch_size (int, optional) – Number of rotations evaluated together. Defaults to 10.

Returns:

The errors, [n_signals, n_rotations], in both normalisations.

Return type:

EquivarianceErrors

s2conv.utils.equivariance.equivariance_errors_of_section(kernel, s_in: int, N: int, signals, rotations, chi: Callable = None, L_quad: int = None, batch_size: int = 10, method: str = 'jax') → EquivarianceErrors[source]#

Relative equivariance errors of an evaluation on a section, from its spin components.

The evaluation on a section \(E_\sigma[\Psi]\,{}_sf = \sum_{|s'|<N} e^{-\mathrm{i}s'\chi_\sigma}(2\pi)^{-1/2}\mathcal{C}_{s\to s'}[\Psi]\,{}_sf\) (Proposition VII.1, Eq. 62) is a scalar field in the coordinate frame that is not band-limited. Both \(E_\sigma(\mathcal{R}f)\) and \(\mathcal{R}(E_\sigma f)\) are sums of band-limited spin components weighted by frame phases: the components of \(\mathcal{R}f\), and the rotated components of \(f\), since \(g(R^{-1}\omega) = e^{\mathrm{i}s'\vartheta}(\mathcal{R}g)(\omega)\) (Eq. 12). Both are therefore evaluated by inverse spin transforms on the Gauss-Legendre grid of band-limit L_quad, rather than point by point, and the errors follow by quadrature. The components are band-limited, so their values on the grid are exact; only the quadrature of the error, which is not band-limited, is approximate. This is much faster than equivariance_errors_on_points() at large band-limits.

Standard correlation is chi=None with the kernel \(\Lambda_0\psi\); the commutative anisotropic convolution of Sadeghi et al. is chi = lambda theta, phi: np.pi - phi with the kernel \(\Lambda_0 h^*\).

Parameters:

  • kernel – The kernel \(\Psi\) (see s2conv.kernels).

  • s_in (int) – Spin \(s\) of the input.

  • N (int) – Azimuthal band-limit: the spin components \(|s'| < N\) are kept.

  • signals – Inputs [n, L, 2L - 1].

  • rotations – Euler angles \((\alpha, \beta, \gamma)\), arrays of equal length.

  • chi (callable, optional) – The twist \(\chi_\sigma(\theta, \varphi)\) of the section, written with JAX; None for the coordinate section.

  • L_quad (int, optional) – Band-limit of the quadrature grid, at least \(L\). Defaults to \(L\); a larger value refines the quadrature of the error.

  • batch_size (int, optional) – Number of rotations evaluated together. Defaults to 10.

  • method (str, optional) – Transform backend, one of s2conv.transforms.METHODS. Defaults to "jax", accurate for \(N \le 8\).

Returns:

The errors, [n_signals, n_rotations], in both normalisations.

Return type:

EquivarianceErrors

s2conv.utils.equivariance.summarise(errors: EquivarianceErrors) → dict[source]#

Mean, standard deviation and maximum of the errors in each normalisation.

Parameters:

errors (EquivarianceErrors) – Errors from equivariance_errors() or equivariance_errors_on_points().

Returns:

{normalisation: {"mean": ..., "std": ..., "max": ...}}, with float values.

Return type:

dict