Accelerated or differentiable GLASS

[ntessore]

This discussion is to try and unify the questions of: What about GPU or TPU support / autodiff / jax support / torch support / tensorflow support / etc.?

Short answer

All of this should be possible, but requires an amount of work that depends on where you want to make a cut; from surface-level modifications (Gaussian fields) to deep programming work (lognormal fields) and more difficult conceptual questions (galaxy catalogues).

Longer answer and to do list

GLASS has from the beginning been designed with acceleration and differentiability in mind. As a matter of fact, what was ultimately released is actually the fourth complete rewrite of GLASS, and part of each was that at some point I knew the old code would not be flexible enough to support things like autodiff or GPU computation.

So this has been very much part of the plan. In fact, I can give a reasonable to do list of what is necessary to implement e.g. JAX support in the existing library functions.

Overall, my strategy and hope is that GLASS should eventually be written in a way that is completely agnostic about the type of array it is working on, and thus support all of numpy, jax, dask, etc. natively in one code base. This could involve e.g. use of the Array API.

Gaussian fields

The first step should definitely be making Gaussian matter fields work, because lognormal ones pose a greater difficulty (see below), and this is still a good test case.

In principle, only surface-level changes to the existing code are necessary, because everything that is necessary to sample random fields is GLASS' own code, precisely for the purpose of this discussion, with the exception of healpy's spherical harmonic transforms.

What needs to change is that the functions in glass.fields must be made array agnostic, so that they work equally on e.g. numpy arrays and JAX arrays. This should not be hard. The most tricky bit would be the iternorm generator for the iterative sampling of a multivariate normal, but I actually have half a prototype for that. Everything else appears to only require figuring out how to change from the numpy namespace to something more flexible.

For sampling Gaussian random fields, that leaves the spherical harmonic transforms, and perhaps others can comment on our options.

Lognormal fields

Lognormal fields (or rather, any transformation of a Gaussian field) are difficult. The actual transformation of the Gaussian field (for lognormal, Y = exp(X + c) - 1) is easy; as above, it only requires swapping out the numpy namespace for something generic.

But lognormal fields require finding the Gaussian angular power spectra (Gls) that result in the correct realised angular power spectra (Cls) after transformation of the field. To get the desired accuracy in the Cls, we use the nonlinear solver we present in arXiv:2302-01942. That solver has two layers: One is the solver itself, which is my gaussiancl package. That is Python code, and I think it can easily be made array agnostic, because it performs simple arithmetic. The code can be seen here: https://github.com/cltools/gaussiancl/blob/main/gaussiancl.py

The problem is the second layer, which is the conversion between angular power spectra and angular correlation functions C(θ). That is nominally done by my transformcl package, but the actual numeric heavy lifting is done by a Fast Legendre Transform implemented in the flt package. That is C code using the method we present in arXiv:2302-01942, which has two steps:

1. A discrete cosine transform to convert C(θ) to its Fourier coefficients.
2. A matrix multiplication to convert the Fourier coefficients to the Legendre coefficients.

The implementation of that can be found here: https://github.com/ntessore/flt/blob/main/flt.pyx

Making this same code array agnostic would require (1) either a generic DCT implementation or a DCT implementation for each array type, and (2) coding up the matrix multiplication in an array agnostic manner.

The latter point is not trivial; the reason that it is done in C is that the matrices are too big to construct explicitly, and we instead compute them iteratively at the same time as doing the matrix multiplication itself. This happens here: https://github.com/ntessore/flt/blob/main/dctdlt.c

Weak lensing

This should be a simple case of removing a few numpy namespace uses. The actual operation performed is a weighted sum of the matter fields, so it should always work with whatever array type is being used.

Galaxies

Instead of going into detail about the individual steps in the galaxies sector, let me ask a question: How can we make

n = rng.poisson(lamda)

differentiable?

Summary

I think we should be able to get Gaussian random fields and lensing working more or less right away.

If that looks promising, maybe we can come up with a plan for how to make the Gaussian angular power spectra work.

[jasonmcewen]

Thanks for the clear summary @ntessore.

For the spherical harmonic transforms, we've just released a new JAX code: s2fft. This is a very early release; we're still polishing the code and writing the related paper.

For the fast Legendre transforms, it seems there is already a DCT in JAX. Could the iterative matrix construction and then application be ported to Python? Otherwise, we could perhaps consider alterative approaches to the Legendre transform, perhaps leveraging components from spherical harmonic transforms algorithms.

[ntessore]

The s2fft looks great @jasonmcewen, congratulations!

With that code being available already, perhaps we can set out on the plan that @EiffL and I talked about last week: get the Gaussian random fields with lensing to work in a differentiable way.

In the core library, we could factor a generate_alms function out of generate_gaussian, and make the group of iternorm, cls2cov, multalm work with any array type.

On the JAX side, we could then create a glass.jax.generate_gaussian function that calls the s2fft transform instead of the healpy one.

[CosmoMatt]

@ntessore for this application it seems we only need the inverse spin-spherical harmonic transform? Currently, the healpix support from s2fft should be exact for inverse but we haven't yet implemented the optimisation iterations that are necessary to get good accuracy for the forward transform with healpix sampling. Of course, if you will need the forward transform too we can bump this to the top of our priorities!

[ntessore]

Hi @CosmoMatt, we use both transforms to compute shear from convergence for weak lensing, which is a convolution in real space.

[ntessore]

Let me reply to this as a new comment:

For the fast Legendre transforms, it seems there is already a DCT in JAX.

That's good, so we nominally only have to figure out a big matrix multiplication.

Could the iterative matrix construction and then application be ported to Python?

Porting it is not a problem as such, but when doing validation tests, the matrix was routinely 100,000 x 100,000 (and a lot of accuracy tests were 10,000,000 x 10,000,000).

However, maybe one can exploit the structure of the matrix better. It is a triangular matrix where every second off-diagonal is zero. Maybe one could look at efficient triangular x vector multiplications in Python, and try to make one compatible with the recurrence? The one we use goes one row down the diagonal and across all columns, but maybe one can go down across two columns and down an entire diagonal as well.

Otherwise, we could perhaps consider alterative approaches to the Legendre transform, perhaps leveraging components from spherical harmonic transforms algorithms.

I have no particular attachment to our current method, but we did come up with it because no existing alternative worked well enough. All methods that I am aware of map between $n$ values $C_l$ and $2n$ values $C(\theta)$. This is not invertible, and you can have situations where you

• compute $C(\theta)$ from $C_l$,
• apply a transformation to $C'(\theta) = f(C(\theta))$,
• compute $C_l'$ from $C'(\theta)$,

and end up with a result where mapping $C_l'$ back produces a different $C'(\theta)$ that is not the transform of $C(\theta)$. (To see that this must happen, sometimes, just look at the dimensions of the spaces: $n \mapsto 2n \mapsto n$.)

[ntessore]

It was helpfully pointed out to me by @nstarman that the array-api-compat would already allow us to add universal array support to most of the core library.

So let's try and get something done in that direction. I have opened issue glass-dev/glass#67 for more specific implementation details. Hopefully, we can get a first draft PR soon.

@EiffL, let me know if you or others need to be in the loop.

[ntessore]

I have started work on this in glass-dev/glass#68. It was surprisingly easy to get set up. With the mechanism proposed in the PR, you would only need to add a glass._array_api.jax module that imports the JAX array namespace (should be added to array-api-compat eventually).

The best thing is that we can even move the spherical harmonic transform definitions into these array namespaces, so that xp.sht will call either glass._array_api.numpy.sht (uses numpy) or glass._array_api.jax.sht (uses s2fft) depending on the array namespace being used.

[ntessore]

This idea required quite a bit more work than that, because the random number generation and HEALPix functionality are tightly coupled to the array type.

Nevertheless, I think I found a workable mechanism, using a broader compatibility layer in the form of new glass.array, glass.random, and glass.healpix modules.

I have started experimenting with this on the compat branch. There is a skeletal glass.compat.jax module (just for experimenting, that goes into its own glass-compat-jax package). You would enable JAX support by importing the module before loading anything else:

>>> import glass.compat.jax

Afterwards, the compatibility layer modules point to the JAX modules:

>>> from glass import array, random, healpix
>>> array
<module 'glass.compat.jax.array'>
>>> random
<module 'glass.compat.jax.random'>
>>> healpix
<module 'glass.compat.jax.healpix'>

Which means that the JAX support can easily provide its own implementation of e.g. healpix.map2alm.

[ntessore]

We have a prototype of a working compatibility layer! It consists of the glass.generic module in the GLASS core library, which contains functions

• array_namespace to get the array namespace for given objects,
• random_namespace to get a namespace of functions for random sampling, and
• healpix_namespace to get a namespace of HEALPix routines.

The individual array implementations are packaged into individual modules. As a test, I made:

• glass-compat-numpy for numpy support, using
  • array_api_compat for arrays,
  • numpy.random for random sampling, and
  • healpix/healpy for HEALPix.
• glass-compat-jax for JAX support, which is currently empty.

This seems to work well in practice. As a test, I have adapted the glass.galaxies module, which makes straightforward use of the generic functions, without any difficulty.

In particular @EiffL / @jasonmcewen if you fill the modules in glass.compat.jax we have a working prototype with JAX support!

[jasonmcewen]

Sounds like great progress @ntessore !

Just to let you know, we’ve added custom gradients to s2fft . We were seeing a large memory overhead for high resolutions when running backward passes with autodiff (due to the nature of reverse mode autodiff). The custom gradients have eliminated this and we can now run at high resolutions.

We're still working on some further improvements for HEALPix (interations in transforms to imporve accuracy, faster jit compile time) but we could integrate things as they stand now. Further improvements in s2fft should then be automatically picked up glass.

I'm not quite sure how things should be added in glass.compat.jax but perhaps we could have a call with @CosmoMatt to discuss further or even do some pair programming?

[ntessore]

Sounds good @jasonmcewen and @CosmoMatt!

Let's definitely set up a call offline. In the meantime, I will try and give you a rough idea of how I imagine the API might work. This is all tentative, of course.

In each array namespace, there is a healpix module with the same standardised functions. For example, in the glass.compat.numpy.healpix module, I currently import all pixel functions from the healpix package, and the map2alm and alm2map functions from healpy.

To get jax support, you have to add the same functions to the glass.compat.jax.healpix module. I imagine that your function signatures differ from the healpy ones, so you need to wrap them into the healpy form, if that is what we adopt as the healpix API. Alternatively, we might wrap both healpy and s2fft into a common API which is the superset of both interfaces.