Moffat profile drawing returns NaN (`demo2.py`)

[b-remy]

Running GalSim's demo2.py with JAX-GalSim fails and it seems to come from the Moffat profile. Here's a minimal script showing the error

import jax_galsim

gal_flux = 1.e5    # counts
gal_r0 = 2.7       # arcsec
g1 = 0.1           #
g2 = 0.2           #
psf_beta = 5       #
psf_re = 1.0       # arcsec
pixel_scale = 0.2  # arcsec / pixel

# Define the galaxy profile.
gal = jax_galsim.Exponential(flux=gal_flux, scale_radius=gal_r0)

# Shear the galaxy by some value.
gal = gal.shear(g1=g1, g2=g2)

# Define the PSF profile.
psf = jax_galsim.Moffat(beta=psf_beta, flux=1., half_light_radius=psf_re)

# Final profile is the convolution of these.
final = jax_galsim.Convolve([gal, psf])

print(gal.drawImage(scale=pixel_scale).array)
"""
[[3.8022862e-04 3.8958775e-04 4.0335540e-04 ... 6.2590830e-06
  7.4280219e-06 6.6112520e-06]
 [4.0020858e-04 4.1460019e-04 4.2499771e-04 ... 7.5581743e-06
  6.7426722e-06 7.9196307e-06]
 [4.2534867e-04 4.3639768e-04 4.5194468e-04 ... 6.8268273e-06
  7.9949132e-06 7.2239072e-06]
 ...
 [7.2855555e-06 8.0750069e-06 6.9304629e-06 ... 4.5205365e-04
  4.3652899e-04 4.2551477e-04]
 [7.2642902e-06 6.0641905e-06 6.8564354e-06 ... 4.2461159e-04
  4.1419003e-04 3.9978177e-04]
 [6.7501096e-06 7.5687421e-06 6.3996608e-06 ... 4.0379266e-04
  3.9000134e-04 3.8062737e-04]]
"""

print(psf.drawImage(scale=pixel_scale).array)
"""
[[nan nan nan ... nan nan nan]
 [nan nan nan ... nan nan nan]
 [nan nan nan ... nan nan nan]
 ...
 [nan nan nan ... nan nan nan]
 [nan nan nan ... nan nan nan]
 [nan nan nan ... nan nan nan]]
"""

print(final.drawImage(scale=pixel_scale).array)
"""
[[nan nan nan ... nan nan nan]
 [nan nan nan ... nan nan nan]
 [nan nan nan ... nan nan nan]
 ...
 [nan nan nan ... nan nan nan]
 [nan nan nan ... nan nan nan]
 [nan nan nan ... nan nan nan]]
"""

[EiffL]

sighhhh.... bessel functions

[EiffL]

@b-remy could you try with this branch? #180

[EiffL]

This is fixed in #180.

Root cause: The convergence tolerances in the Temme series and Steed's continued fraction algorithms were hardcoded to jnp.finfo(jnp.float64).eps (~1e-16). When running in float32 (the default without JAX_ENABLE_X64=1), the iterative algorithms could never converge to that precision within the fixed iteration count, causing accumulated numerical garbage to produce NaN.

Fix: Changed the three tolerance/epsilon values to use jnp.finfo(z.dtype).eps, so they adapt to the actual computation dtype (~1e-7 for float32, ~1e-16 for float64).

Verified that the exact reproduction script from the issue now works correctly in float32:

• psf.drawImage(scale=0.2) — no NaN, center value matches float64 result
• final.drawImage(scale=0.2) (Exponential + Moffat convolution) — no NaN

[b-remy]

Yes, this is solved with #180 !