feat: rotation matrices using s2fft

[lgrcia]

Using the s2fft Python package to compute the Wigner D-matrices used to rotate the spherical harmonics. See also #140

[dfm]

The issue here must be related to s2fft! @lgrcia, can you try to work out which inputs we're passing in the test that hangs to isolate exactly what s2fft call we're executing. It would be interesting to know if we can get the same hang using just s2fft and no jaxoplanet. In that case we can report upstream and see what they say.

[lgrcia]

I think I identified where the problem is from.

Here is a way to reproduce it on macos

import numpy as np
from jaxoplanet.experimental.starry.rotation import dot_rotation_matrix

l_max = 5
theta = 0
ident = np.eye(l_max**2 + 2 * l_max + 1)
expected = dot_rotation_matrix(l_max, 0.0, 0.0, 1.0, theta)(ident)
calc = dot_rotation_matrix(l_max, None, None, 1.0, theta)(ident)

This runs ok. But then, when l_max is changed in the same runtime:

l_max = 6
ident = np.eye(l_max**2 + 2 * l_max + 1)
expected = dot_rotation_matrix(l_max, 0.0, 0.0, 1.0, theta)(ident)

it freezes.

I think it has to do in how I (or s2fft) combine the static arguments in the jitted functions from s2fft.utils.rotation and s2fft.utils (see also the jitted signature from jaxoplanet.experimental.starry.s2fft_rotation.py:compute_rotation_matrices). But I honestly don't totally get why it is only a problem on macos.

I don't really understand why it behaves like this but a workaround for me is to avoid decorating the s2fft rotation functions with jit. So I copied all required functions (we only need 100 lines of python from s2fft) and removed the s2fft dependency, for now.

I'm down to understand the problem better before reintroducing s2fft as a dependency.

[lgrcia]

The issue here must be related to s2fft! @lgrcia, can you try to work out which inputs we're passing in the test that hangs to isolate exactly what s2fft call we're executing. It would be interesting to know if we can get the same hang using just s2fft and no jaxoplanet. In that case we can report upstream and see what they say.

@dfm, here is a way to reproduce only with s2fft:

import jax
from functools import partial
from s2fft.utils.rotation import generate_rotate_dls


@partial(jax.jit, static_argnums=(0,))
def f(deg, alpha):
    return generate_rotate_dls(deg, alpha)


_ = f(5, 0.0)  # this executes fine
_ = f(10, 0.0)  # this freezes

I might open an issue but I think this is not a proper use of this function given the static args (see https://github.com/astro-informatics/s2fft/blob/main/s2fft/utils/rotation.py#L75). Any idea why this would happen?

I understand that each test run in separate python instances should pass. So could the issue be due to how pytest runners are dispatched on macOS?

[dfm]

It's fascinating to me that that happens and that your solution works! I don't see any reason why the beta parameter should be labelled as static, but it also seems like it should crash to nest these static_argnums incompatibly like this...

Regardless: I think this is a good "fix"!

[lgrcia]

Thanks! All done