Wigner FFT + IFFT is not the identity for 'DH' sampling and low bandwidths

[PhilippMisofCH]

First of all, thanks again for this nice library. While working with it, I've noticed some unexpected results in my project, which eventually boil down to the fact that applying the inverse Wigner FFT on the Wigner FFT does not yield the original input, not even at low bandwidth (L=8) and dh sampling (This is actually one of the only cases that I've tested thoroughly). In contrast, an equivalent setup for the $S^2$ transforms seems to work as expected.

So I wondered if this is indeed a problem, or it is to be expected.

I've also attached a small example where I test this out:

L = 8
sampling = 'dh'
a = s2fft.sampling.s2_samples.phis_equiang(L, sampling=sampling)
b = s2fft.sampling.s2_samples.thetas(L, sampling=sampling)
g = s2fft.sampling.s2_samples.phis_equiang(L, sampling=sampling)
f_grid = jnp.meshgrid(g, b, a, sparse=True, indexing='ij')
f_s2 = jnp.squeeze(jnp.sin(f_grid[1]) * jnp.cos(f_grid[2]))
flm = s2fft.transforms.spherical.forward_jax(f_s2, L, sampling=sampling, reality=True)
f_s2_back = s2fft.transforms.spherical.inverse_jax(flm, L, sampling=sampling)
assert jnp.allclose(f_s2, f_s2_back, atol=1e-4), ("Spherical forward + inverse transform is"
                                                  " not the identity")
                                                                                             
f_so3 = jnp.sin(f_grid[0]) * jnp.cos(f_grid[1]) * jnp.sin(f_grid[2])
flmn = s2fft.transforms.wigner.forward_jax(f_so3, L, L, sampling=sampling, reality=True)
f_back = s2fft.transforms.wigner.inverse_jax(flmn, L, L, sampling=sampling)
assert jnp.allclose(f_so3, f_back, atol=1e-2), ("Wigner forward + inverse transform is not "
                                                "the identity")

The first assertion passes, but the second one fails.

I've also looked around at other reported issues and found e.g.

• Inconsistent results for pure JAX implementation of inverse Wigner transformation #209, but this seems to only appear for higher values of N than I'm testing here
• Tests for consistency of HEALPix FFT and IFFT implementations #170 (comment) but this is specific to the HEALPix as far as I understand.

It would be great if you could tell me if I misunderstood something or if this is indeed a bug / numerical issue. Thanks in advance!

[PhilippMisofCH]

Adding on this:
At first sight it looks like one of the notebooks is demonstrating explicitly that this works, but here, one starts in the Fourier space, applies the inverse transformation and then goes back to the Fourier space, where all values agree.

Similarly, doing the following also passes the assertion:

rng = np.random.default_rng(83459)
flmn = s2fft.utils.signal_generator.generate_flmn(rng, L, L, reality=True)
f_so3 = s2fft.transforms.wigner.inverse_jax(flmn, L, L, sampling=sampling)
flmn = s2fft.transforms.wigner.forward_jax(f_so3, L, L, sampling=sampling, reality=True)
f_so3_back = s2fft.transforms.wigner.inverse_jax(flmn, L, L, sampling=sampling)
assert jnp.allclose(f_so3, f_so3_back, atol=1e-8), ("Wigner forward + inverse transform is not "
                                                "the identity")

But for an arbitrary (smooth) signal generated on $SO(3)$ this does not work as shown before.

[CosmoMatt]

Hey @PhilippMisofCH so one of the assumptions of spherical harmonic (and by extension Wigner) transforms is that the underlying signals are band-limited (see e.g. section IV.A of this paper). For signals that are not band-limited the truncation of the summation in equation 9 doesn't hold with equality.

Generally

When you generate a random signal in pixel space it isn't necessarily band-limited, and therefore the (on its face reasonable) assumption that forward followed by inverse should return the original signal doesn't hold.

Specifically

The checkerboard example you use here is a bit of a trick problem I think.

• For the spherical harmonic transform you use $f(\theta, \phi) = \sin\theta \cos\phi$. As $\theta \rightarrow 0$ the effective band-limit (maximum frequency) in $\phi \rightarrow 0$, which you can see by considering the case where $\theta = 0$ the function must necessarily be the same $\forall \phi$ as these are all the same locations on the sphere. What is making this function appear to work in the example is that $\sin\theta \rightarrow 0$ as $\theta \rightarrow 0$ and therefore $f(\theta,\phi) \rightarrow 0$ $\forall\phi$ which more closely satisfies (but by no means exactly) the effective band-limit in $\phi$ (everything is just 0). If you instead used $\cos\theta$ in $f(\theta,\phi)$ you will see much more substantial errors at both poles due to this.

• For the Wigner transform you use $f(\alpha, \beta,\gamma) = \sin\alpha\cos\beta\sin\gamma$. As per the above argument as $\beta\rightarrow0$ the $\cos\beta$ component doesn't kill the $\gamma$ and $\alpha$ dependencies, and therefore the signal fails to satisfy the effective band-limiting. Changing $\cos\beta$ to $\sin\beta$ reduces these errors somewhat, again for the reasons outlined above, but this is really just hiding the underlying band-limiting issue. If you explicitly band limit your function by

f_so3 = jnp.sin(f_grid[0]) * jnp.cos(f_grid[1]) * jnp.sin(f_grid[2])

# Apply once to bandlimit the signal.
flmn = s2fft.transforms.wigner.forward_jax(f_so3, L, N, sampling=sampling, reality=True)
f_so3 = s2fft.transforms.wigner.inverse_jax(flmn, L, N, sampling=sampling)

# Apply to bandlimited signal to get error.
flmn = s2fft.transforms.wigner.forward_jax(f_so3, L, N, sampling=sampling, reality=True)
f_so3_back = s2fft.transforms.wigner.inverse_jax(flmn, L, N, sampling=sampling)

you will find machine precision (or at least that's what I find on my side).

Sidenote

As you mention, if you increase $N$ beyond $\sim 8$ or so you will start to see errors being introduced by the recursion. We've recently implemented more stable recursions for these cases which you can pick up easily through the precompute versions of these transforms. These should be stable $\forall N \leq L$, as always if you have any issues picking this functionality up let us know :)