Use ragged arrays for alms

[ntessore]

Add your issue here

Inspired by s2fft, I think it would be great to use ragged alm arrays. The first axis would be $\ell = 0, \ldots, \ell_{\max}$, the second axis would be ragged with $m = 0, \ldots, \ell$.

[ntessore]

See also #382.

[Saransh-cpp]

Is s2fft already using ragged arrays from some package right now? I could not find anything relevant. Their paper says -

From a software engineering perspective, this raises some further complications. Ideally, the FFTs of Eq. (31) and Eq. (32) would be implemented as a single function call to jax.numpy.fft across a regular array. However, jax.numpy.fft does not support ragged arrays and therefore each ring within the polar region must computed by individual function calls within a scheduled loop.

but I am not sure where the ragged array (that they want to feed in jax.numpy.fft) is coming from.

[ntessore]

You can see here the format s2fft uses to store their alms:

$$\mathtt{alm} = \begin{bmatrix} 0 & 0 & flm_{(0,0)} & 0 & 0 \\\ 0 & flm_{(1,-1)} & flm_{(1,0)} & flm_{(1,1)} & 0 \\\ flm_{(2,-2)} & flm_{(2,-1)} & flm_{(2,0)} & flm_{(2,1)} & flm_{(2,2)} \end{bmatrix}$$

This is always for complex fields, so there are negative $m$ values present.

We would continue to use pairs of $E$ and $B$ mode alms instead, so no negative $m$ values. Our arrays could therefore look like this:

$$\mathtt{alm} = \begin{bmatrix} flm_{(0,0)} & 0 & 0 \\\ flm_{(1,0)} & flm_{(1,1)} & 0 \\\ flm_{(2,0)} & flm_{(2,1)} & flm_{(2,2)} \end{bmatrix}$$

Representing this as a ragged array makes it very easy to apply functions of $\ell$, e.g., when random sampling: just multiply over the leading axis. Selecting individual $m$ values would also very naturally do the right thing, e.g., alm[:, 2:] would have two empty rows, since there are no $m < 2$ values at $\ell = 0, 1$.

The only major issue here is that ragged doesn't support JAX 😞

[Saransh-cpp]

I see, thanks! I did notice this format, but the indexing looked very different from the one in GLASS, so I was skeptical.

ragged supports GPU and JAX through Awkward Array. GPU support is extensive (including Numba + CUDA), but JAX support is limited (mostly due to restrictions from JAX). There are ongoing talks of writing a custom autodiff backend for Awkward and they might do it this summer.

I am not very sure how this is to be implemented but I am expecting that the alms should start looking like gls in the function arguments -

gls = [gl_00,
       gl_11, gl_10,
       gl_22, gl_21, gl_20,
       ...]

but yes, it will become more clear once I have a quick meet with you after FOSDEM.