gh-487: reorder alms in _generate_grf to glass ordering

[Saransh-cpp]

Description

The PR reorder alms in _generate_grf to glass ordering before operating on them and then reorders it back to healpix before passing into the yield function.

The code works but it is slowing down examples instead of making them faster. I think I am doing something wrong here. @ntessore could you look at this and point me in the right direction?

Fixes: #487

Changelog entry

Checks

• Is your code passing linting?
• Is your code passing tests?
• Have you added additional tests (if required)?
• Have you modified/extended the documentation (if required)?
• Have you added a one-liner changelog entry above (if required)?

[Saransh-cpp]

See #530 for the mypy failure.

[Saransh-cpp]

Should the functions return a numpy array instead?

[ntessore]

Thanks for taking this on @Saransh-cpp!

I think we crossed wires with two similar but separate things going on at the same time:

• The ordering of the spectra. Here we are using GLASS ordering and we want to continue GLASS ordering. The two new functions are only useful for legacy mode interactions between GLASS and other HEALPix-based codes.
• The ordering of the alms. Here we are using HEALPix ordering at the moment, and the idea is to switch this around.

This switch is fully internal and there are no observable changes to anything from the outside.

The way our alms = $a_{lm}$ are sorted now is keeping $m$ fixes while $l$ increases, so first $m = 0$, $l = 0, 1, 2, \ldots$, then $m = 1$, $l = 1, 2, \ldots$, then $m = 2$, $l = 2, \ldots$, etc. This is encoded in how multalm() works.

After the change, our internal generation would keep $l$ fixed while $m$ increases, so first $l = 0$, $m = 0$, then $l = 1$, $m = 0, 1$, then $l = 2$, $m = 0, 1, 2$, etc. Note how this corresponds to switching from iterating over an upper triangular matrix to a lower triangular matrix.

The goal of this internal reordering is that we will be able to split everything up in blocks of $l$ values with different correlation lengths.

After this internal reordering, we still need to send our alms to healpy for the spherical harmonic transform. So there needs to be an internal reordering.

[paddyroddy]

Can you ask me again for a review once the science is worked out? I'm not fully sure what is going on here

[ntessore]

Oh, I just realised this is marked as a fix for #487. That issue is about an unrelated change. The reordering of the alms in the GRF sampling that we discussed is a purely internal change for the time being. Maybe #487 will never be done.

[Saransh-cpp]

@ntessore could you please take a look at this again and let me know if this is the right direction?

multalm is definitely slower now because of the dynamic indexing. It would really benefit from using ragged arrays for accessing alms in a NumPy-like way, but we can discuss that in the JAX project. Just using plain NumPy arrays would eat up extra memory because of the matrix structure.

[Saransh-cpp]

The function now assumes that the incoming alms are in the order -

[00, 10, 11, 20, 21, 22, 30, 31, 32, 33, ...]

Therefore, it now multiplies alm with bl in the following way -

bl := [1, 10, 100, 1000]

[0, 10, 20, 30] * bl → [0, 100, 2000, 30000]
[11, 21, 31] * bl[:3] → [11, 210, 3100]
[22, 32] * bl[:2] → [22, 320]
[33] * bl[:1] → [33]

[ntessore]

I don't think that's right: bl[0] is to be multiplied with alm where l=0, so just alm[0:1]. bl[1] is to be multiplied with alm where l=1, so alm[1:3], and so on.

Long version of the loop:

stop = 0
for ell in range(n):
    start, stop = stop, stop + ell + 1
    alm[start : stop] *= bl[ell]

Short version of the loop:

for ell in range(n):
    alm[ell * (ell + 1) // 2 : (ell + 1) * (ell + 2) // 2] *= bl[ell]

[Saransh-cpp]

Ah, I got the multiplication wrong! That's why I was confused about why we are switching to this order. Thanks for the correction!

[Saransh-cpp]

IIRC, the alms should be converted back to healpix ordering before being passed into hp functions. I have used the glass.glass_to_healpix_spectra because it does the ordering in the same way.

[ntessore]

If it does the reordering in the same way, that's "by accident", and we should not bind ourselves to an unrelated function that happens to do the same thing.

For reordering the alm:

Have: $[a_{00}, a_{10}, a_{11}, a_{20}, a_{21}, a_{22}, \ldots]$
Want: $[a_{00}, a_{10}, a_{20}, \ldots, a_{11}, a_{21}, \ldots, a_{22}, \ldots]$

In the first array, the element with indices $(\ell, m)$ appears in position $\ell (\ell + 1) / 2 + m$. This is what you have in the function above.

def glass_to_healpix_alm(alm):
    n = inv_triangle_number(alm.size)  # we need that function back :)
    ell = np.arange(n)
    out = []
    for m in ell:
        out.append(alm[ell[m:] * (ell[m:] + 1) // 2 + m])
    return np.concatenate(out)

[Saransh-cpp]

Thanks! Should I just revert #527 so that we don't use functions named for spectra inside functions meant for alm?

[ntessore]

Maybe we should add the private function back in, then make two helper functions with clear purpose (and error msg) instead of using it directly?

[Saransh-cpp]

Sounds good! I will wrap that up in another PR, which can be merged before this one to make everything cleaner.

[Saransh-cpp]

Okay, so the calculations are not correct, but I am not sure where it is going wrong. I will debug tomorrow -

cosmic_shear.ipynb

lensing.ipynb

[ntessore]

We should really set a timeout on the examples test 😬

[ntessore]

Let's make this more obvious

Suggested change

	inp = np.array([11, 22, 21, 33, 32, 31, 44, 43, 42, 41])
	out = glass.fields._glass_to_healpix_alm(inp)
	np.testing.assert_array_equal(
	out, np.array([11, 22, 33, 44, 21, 32, 43, 31, 42, 41])
	)
	inp = np.array([00, 10, 11, 20, 21, 22, 30, 31, 32, 33])
	out = glass.fields._glass_to_healpix_alm(inp)
	np.testing.assert_array_equal(
	out, np.array([00, 10, 20, 30, 11, 21, 31, 22, 32, 33])
	)