float32 consistency

jax_galsim/gsobject.py

@@ -986,7 +986,15 @@ def drawFFT(self, image, add_to_image=False):
         kimage = self._drawKImage(kimage)
         return self.drawFFT_finish(image, kimage, wrap_size, add_to_image)
 
-    @implements(_galsim.GSObject.drawKImage)
+    @implements(
+        _galsim.GSObject.drawKImage,
+        lax_description="""\
+The JAX-GalSim version of ``drawKImage`` accepts an extra ``dtype`` to set
+the complex dtype of k-space-image when ``image`` is not provided
+(e.g. ``numpy.complex64``). If ``dtype`` is None, it defaults to ``complex128``
+when ``jax_enable_x64`` is enabled and ``complex64`` otherwise.
+""",
+    )
     def drawKImage(
         self,
         image=None,
@@ -998,10 +1006,18 @@ def drawKImage(
         recenter=True,
         bandpass=None,
         setup_only=False,
+        dtype=None,
     ):
         from jax_galsim.image import Image
         from jax_galsim.wcs import PixelScale
 
+        if dtype is not None:
+            dtype = Image._alias_dtypes.get(dtype, dtype)
+            if not jnp.issubdtype(dtype, jnp.complexfloating):
+                raise _galsim.GalSimValueError(
+                    "dtype must be complex for drawKImage", dtype
+                )
+
         # Make sure provided image is complex
         if image is not None:
             if not isinstance(image, Image):
@@ -1010,6 +1026,13 @@ def drawKImage(
             if not image.iscomplex:
                 raise _galsim.GalSimValueError("Provided image must be complex", image)
 
+            if dtype is not None and image.dtype != dtype:
+                raise _galsim.GalSimIncompatibleValuesError(
+                    "Cannot specify dtype != image.dtype if image is provided",
+                    dtype=dtype,
+                    image=image,
+                )
+
         # Possibly get the scale from image.
         if image is not None and scale is None:
             # Grab the scale to use from the image.
@@ -1036,7 +1059,11 @@ def drawKImage(
         # So make that profile.
         if image is None or not image.bounds.isDefined():
             real_prof = PixelScale(dx).profileToImage(self)
-            dtype = np.complex128 if image is None else image.dtype
+            if image is not None:
+                dtype = image.dtype
+            elif dtype is None:
+                # complex128 is not available when x64 is disabled
+                dtype = np.complex128 if jax.config.jax_enable_x64 else np.complex64
             image = real_prof._setup_image(
                 image, nx, ny, bounds, add_to_image, dtype, center=None, odd=True
             )

jax_galsim/image.py

@@ -269,6 +269,9 @@ def __init__(self, *args, **kwargs):
                 raise TypeError("wcs parameters must be a galsim.BaseWCS instance")
             self.wcs = wcs
 
+        # make sure they are always consistent even when user specifies unavailable dtype
+        self._dtype = self._array.dtype.type
+
     @staticmethod
     def _get_xmin_ymin(array, kwargs, check_bounds=True):
         """A helper function for parsing xmin, ymin, bounds options with a given array"""
@@ -877,9 +880,15 @@ def calculate_fft(self):
         # dk = 2pi / (N dk)
         dk = jnp.pi / (No2 * dx)
 
+        # Use complex64 for float32 inputs, complex128 for float64
+        if self.dtype == jnp.float32:
+            fft_dtype = np.complex64
+        else:
+            fft_dtype = np.complex128
+
         out = Image(
             BoundsI(xmin=0, deltax=No2 + 1, ymin=-No2, deltay=2 * No2),
-            dtype=np.complex128,
+            dtype=fft_dtype,
             scale=dk,
         )
         # we shift the image before and after the FFT to match the layout of the modes
@@ -953,9 +962,15 @@ def calculate_inverse_fft(self):
         # inplace FFT. Thus we do not use the
         # padding.
 
+        # Use float32 for complex64 input, float64 for complex128 input
+        if self.dtype == jnp.complex64:
+            ifft_dtype = jnp.float32
+        else:
+            ifft_dtype = jnp.float64
+
         out = Image(
             bounds=BoundsI(xmin=-No2, deltax=2 * No2, ymin=-No2, deltay=2 * No2),
-            dtype=float,
+            dtype=ifft_dtype,
             scale=dx,
             # we shift the image before and after the FFT to match the layout used by galsim
             array=jnp.fft.fftshift(
@@ -1285,7 +1300,7 @@ def tree_unflatten(cls, aux_data, children):
         obj._array = children[0]
         obj.wcs = children[1]
         obj._bounds = children[2]
-        obj._dtype = aux_data["dtype"]
+        obj._dtype = obj._array.dtype.type
         obj._is_const = aux_data["isconst"]
         if len(children) > 3:
             obj.added_flux = children[3]

jax_galsim/interpolatedimage.py

@@ -634,6 +634,7 @@ def _image(self):
                 "InterpolatedImages do not support 'depixelize' in jax_galsim."
             )
         else:
+            # This does not preserve the input dtype to maintain consistency with Galsim
             image = self._jax_children[0].copy(dtype=float)
 
         if self._jax_aux_data["_recenter_image"]:

tests/jax/test_dtype_consistency.py

@@ -0,0 +1,89 @@
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+import jax_galsim
+
+
+def _gaussian_image(dtype, n=32, scale=0.2):
+    return jax_galsim.Gaussian(sigma=1.0).drawImage(
+        nx=n, ny=n, scale=scale, dtype=dtype, method="no_pixel"
+    )
+
+
+@pytest.mark.parametrize(
+    "dtype,fft_dtype,ifft_dtype",
+    [
+        (jnp.float32, jnp.complex64, jnp.float32),
+        (jnp.int32, jnp.complex128, jnp.float64),
+    ],
+)
+def test_calculate_fft_dtype(dtype, fft_dtype, ifft_dtype):
+    im = jax_galsim.Image(
+        np.arange(64).reshape(8, 8).astype(dtype), scale=0.2, dtype=dtype
+    )
+    kim = im.calculate_fft()
+    assert kim.dtype == fft_dtype
+    assert kim.array.dtype == fft_dtype
+
+    inv_fft_im = kim.calculate_inverse_fft()
+    assert inv_fft_im.dtype == ifft_dtype
+    assert inv_fft_im.array.dtype == ifft_dtype
+
+
+@pytest.mark.parametrize(
+    "dtype,kdtype",
+    [
+        (jnp.float32, jnp.complex64),
+    ],
+)
+def test_interpolated_image_preserves_dtype(dtype, kdtype):
+    im = _gaussian_image(dtype)
+    ii = jax_galsim.InterpolatedImage(im, x_interpolant="lanczos15")
+    assert ii._original._image.dtype == dtype
+    assert ii._original._kim.dtype == kdtype
+
+
+@pytest.mark.parametrize(
+    "dtype,expected",
+    [
+        (np.complex64, jnp.complex64),
+        (jnp.complex64, jnp.complex64),
+        (np.complex128, jnp.complex128),
+        (jnp.complex128, jnp.complex128),
+        (complex, jnp.complex128),
+    ],
+)
+def test_drawkimage_dtype_kwarg(dtype, expected):
+    obj = jax_galsim.Gaussian(sigma=1.0)
+    kim = obj.drawKImage(nx=32, ny=32, scale=0.1, dtype=dtype)
+    assert kim.dtype == expected
+    assert kim.array.dtype == expected
+
+
+@pytest.mark.parametrize("dtype", [jnp.complex64, jnp.complex128])
+def test_drawkimage_image_kwarg_dtype(dtype):
+    obj = jax_galsim.Gaussian(sigma=1.0)
+    image_in = jax_galsim.Image(32, 32, scale=0.1, dtype=dtype)
+
+    # below dtype argument shout be ignored
+    kim = obj.drawKImage(image=image_in)
+    for im in (image_in, kim):
+        assert im.dtype == dtype
+        assert im.array.dtype == dtype
+
+    kim = obj.drawKImage(image=image_in, dtype=dtype)
+    for im in (image_in, kim):
+        assert im.dtype == dtype
+        assert im.array.dtype == dtype
+
+
+@pytest.mark.test_in_float32
+def test_image_and_array_dtype_are_always_consistent():
+    image = jax_galsim.Image(2, 2, dtype=jnp.float64)
+    assert image.dtype == image.array.dtype.type
+
+    leaves, treedef = jax.tree_util.tree_flatten(image)
+    rebuilt = jax.tree_util.tree_unflatten(treedef, leaves)
+    assert rebuilt.dtype == rebuilt.array.dtype.type