[paddyroddy]

Description

This allows users to use glass.harmonics.transform and glass.harmonics.inverse_transform with either healpy or s2fft inputs. S2FFT depends on JAX necessary so I have created an optional dependency group jax that includes both. I've added some tests to check the values are consistent between healpy and s2fft for NumPy inputs. I've also added a test to check what happens if S2FFT isn't installed.

Currently blocked because S2FFT isn't working on 3.13/3.14, see astro-informatics/s2fft#413.

Closes: #1077

Changelog entry

Added: JAX-native scalar and spin spherical harmonic transforms
Added: A glass[jax] install extra
Changed: nside is now mandatory in glass.harmonics.transform

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)?

[Copilot]

Copilot review overview

🟡 Changes recommended

The coefficient-order contract, CPU weighting behavior, GPU residency, and benchmark dependencies have unresolved issues.

Review effort: Balanced
Findings: 1 · 3

Open (4)

• Wrappers use healpy m-major order instead of GLASS ell-major order · New
• Preserve pixel weighting in NumPy convergence transforms · New
• Use backend-native harmonic multiplication for JAX GPU tracing · New
• Include S2FFT in benchmark and ROCm dependencies · New

What changed in this PR

Adds JAX-backed spherical harmonic transforms via S2FFT and migrates relevant field-generation and lensing paths.

Changes:

• Adds NumPy/JAX transform wrappers and S2FFT dependencies.
• Migrates scalar harmonic transforms in fields and lensing.
• Adds cross-backend transform tests.

File	Description
glass/​harmonics.py	Adds forward and inverse SHT wrappers.
glass/​lensing.py	Uses the new wrappers for convergence transforms.
glass/​fields.py	Uses the inverse wrapper for field generation.
glass/​healpix.py	Corrects map2alm return typing.
noxfile.py	Installs S2FFT for JAX test sessions.
pyproject.toml	Adds S2FFT to development groups.
tests/​core/​test_harmonics.py	Tests JAX transforms against NumPy.
tests/​core/​test_healpix.py	Clarifies optional-backend detection.
tests/​fixtures/​helper_classes.py	Generates valid real-map harmonic coefficients.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[paddyroddy]

3 is the default for methods healpix and healpix_jax but not for jax for some reason https://astro-informatics.github.io/s2fft/api/transforms/spin_spherical_transform.html#s2fft.transforms.spherical.forward. Required for good accuracy.

[matt-graham]

The rationale for this is method="healpix_jax" with sampling="healpy" wraps healpy and so matches the default used there while the support for iter in the numpy and jax methods was only added latterly and so the original default behaviour of using no iterative refinement was retained when iter support was added for these methods. It could make sense to make a breaking change and default to iter = 3 for sampling = "healpix" with any method though.

[paddyroddy]

I do think that might make sense. People using S2FFT for the JAX benefits coming from a HEALPix configuration might get thrown.

[Copilot]

Copilot review overview

🟡 Changes recommended

JAX transforms still have unsupported bandlimits and CPU fallbacks, while runtime dependency and test expectations are incomplete.

Review effort: Balanced
Findings: 3

Open (3)

• Declare S2FFT as a supported runtime dependency or extra · New
• Enable pixel weights in the reference call · New
• Preserve pixel weighting in NumPy convergence transforms

Resolved since last review (3)

• Wrappers use healpy m-major order instead of GLASS ell-major order
• Include S2FFT in benchmark and ROCm dependencies
• Use backend-native harmonic multiplication for JAX GPU tracing

Previously missed (4)

In code that hasn't changed since last review

JAX forward transform fails when lmax is below 2*nside

glass/​harmonics.py:146

S2FFT also enforces L >= 2 * nside for the forward HEALPix transform, but this path passes lmax + 1 unchanged. Therefore valid calls such as from_convergence(..., lmax=8) at nside=32 work with NumPy but fail with an assertion for JAX. Mirror the inverse path by transforming at an effective bandlimit of at least 2 * nside, then crop the 2-D coefficients back to the requested bandlimit before packing them.

Derive NSIDE from shape instead of converting JAX arrays to NumPy

glass/​lensing.py:297

Before reaching this JAX transform, hp.get_nside(kappa) uses numpy_fallback, which calls np.asarray(kappa). That copies the entire device map to the host in eager execution and raises a tracer-conversion error under jax.jit, so even the potential-only path is not fully GPU-resident or traceable. NSIDE can be derived from the static map shape without converting the array.

Use a backend-native spin transform for JAX deflection and shear

glass/​lensing.py:330

The deflection path still reaches hp.alm2map_spin at line 334 (and the shear path at line 357); its numpy_fallback converts these JAX coefficients to NumPy and performs the transform on the CPU. Thus the full lensing workflow requested by #1077 remains GPU-incompatible whenever deflection or shear is requested. A backend-native spin inverse transform is needed before this issue can be considered closed.

Document transform and inverse_transform in the harmonics API

glass/​harmonics.py:56

This new public wrapper (and transform) is absent from docs/modules/harmonics.rst, whose API reference still lists only multalm. Add both transform functions there so the feature introduced by this PR appears in the generated user documentation.

[connoraird]

Looks good. Just one thought

[connoraird]

I'm not sure what I personally prefer but should we be defining np as a fixture since that's what we're doing for jnp?

[paddyroddy]

I had this wonder too. The question is whether there is much point or benefit? For array_api_strict and jax we're actually doing something. We wouldn't need to configure NumPy at all. Like even pytest.importorskip is unnecessary at the moment as NumPy is always required.

[connoraird]

I guess the only benefit is np won't be available so we'd know if we accidentally call it an somehow don't notice it in the code.

[paddyroddy]

There would still be the risk that we have the fixture and forget we're using it. I'm happy to do it though.

[connoraird]

True, I don't know either way so I'd say leave as is and merge

[paddyroddy]

Thoughts on f63a8d2?

[connoraird]

I think I do prefer it with the np fixture