fix: use adaptive precision for moffat computation (#181)

* fix: use adaptive precision for moffat computation

* feat: add pytest mark and option to run specific float32 tests

* restart tests

Author: beckermr

.github/workflows/python_package.yaml

@@ -74,6 +74,13 @@ jobs:
             --clean-durations \
             --retries 1
 
+      - name: Test with pytest in float32
+        if: ${{ matrix.group == '1' }}
+        run: |
+          pytest \
+            -vv \
+            --test-in-float32
+
       - name: Upload test durations
         uses: actions/upload-artifact@v6
         with:

jax_galsim/bessel.py

@@ -174,7 +174,7 @@ def si(x):
 
 def _sqrt1px2(x):
     """Numerically stable computation of sqrt(1 + x^2)."""
-    eps = jnp.finfo(jnp.float64).eps
+    eps = jnp.finfo(x.dtype).eps
     return jnp.where(
         jnp.abs(x) * jnp.sqrt(eps) <= 1.0,
         jnp.exp(0.5 * jnp.log1p(x * x)),
@@ -244,7 +244,7 @@ def _temme_series_kve(v, z):
     Assumes |v| < 0.5 and |z| <= 2 for fast convergence.
     Returns exponentially scaled values: Kv(v,z)*exp(z).
     """
-    tol = jnp.finfo(jnp.float64).eps
+    tol = jnp.finfo(z.dtype).eps
 
     coeff1, coeff2, gamma1pv_inv, gamma1mv_inv = _evaluate_temme_coeffs(v)
 
@@ -308,7 +308,7 @@ def _continued_fraction_kve(v, z):
     Assumes |v| < 0.5 and |z| > 2.
     Returns exponentially scaled values: Kv(v,z)*exp(z).
     """
-    tol = jnp.finfo(jnp.float64).eps
+    tol = jnp.finfo(z.dtype).eps
     max_iterations = 1000
 
     initial_numerator = v * v - 0.25

tests/conftest.py

@@ -1,11 +1,13 @@
 # Define the accuracy for running the tests
+import sys
+
 import jax
 
-jax.config.update("jax_enable_x64", True)
+if "--test-in-float32" not in sys.argv:
+    jax.config.update("jax_enable_x64", True)
 
 import inspect  # noqa: E402
 import os  # noqa: E402
-import sys  # noqa: E402
 from functools import lru_cache, partial  # noqa: E402
 from unittest.mock import patch  # noqa: E402
 
@@ -216,3 +218,29 @@ def pytest_runtest_logreport(report):
             report.outcome = "allowed failure"
 
     yield report
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--test-in-float32",
+        action="store_true",
+        default=False,
+        help="Run tests in float32 instead of float64.",
+    )
+
+
+def pytest_configure(config):
+    config.addinivalue_line(
+        "markers", "test_in_float32: mark test to run in float32 tests."
+    )
+
+
+def pytest_report_header(config):
+    return [f"JAX float64 enabled: {jax.config.read('jax_enable_x64')}"]
+
+
+def pytest_runtest_setup(item):
+    if "test_in_float32" not in item.keywords and item.config.getoption(
+        "--test-in-float32"
+    ):
+        pytest.skip("Skipped test not marked to be run in float32.")

tests/jax/test_moffat_comp_galsim.py

@@ -1,5 +1,7 @@
 import galsim as _galsim
+import jax.numpy as jnp
 import numpy as np
+import pytest
 
 import jax_galsim as galsim
 
@@ -54,3 +56,31 @@ def test_moffat_comp_galsim_maxk():
                 psf.kValue(1.0, 0.0), gpsf.kValue(1.0, 0.0), rtol=1e-5
             )
             np.testing.assert_allclose(gpsf.maxk, psf.maxk, rtol=0.25, atol=0)
+
+
+@pytest.mark.test_in_float32
+def test_moffat_conv_nan_float32():
+    # test case from https://github.com/GalSim-developers/JAX-GalSim/issues/179
+    gal_flux = 1.0e5  # 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 = 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 = galsim.Moffat(beta=psf_beta, flux=1.0, half_light_radius=psf_re)
+
+    # Final profile is the convolution of these.
+    final = galsim.Convolve([gal, psf])
+
+    img_arr = final.drawImage(scale=pixel_scale, dtype=jnp.float32).array
+
+    assert jnp.all(jnp.isfinite(img_arr))