diff --git a/jax_galsim/gsobject.py b/jax_galsim/gsobject.py index 641bbdf3..fe41760a 100644 --- a/jax_galsim/gsobject.py +++ b/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 ) diff --git a/jax_galsim/image.py b/jax_galsim/image.py index 50e1a7ed..f3ba8f0c 100644 --- a/jax_galsim/image.py +++ b/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] diff --git a/jax_galsim/interpolatedimage.py b/jax_galsim/interpolatedimage.py index a2a6bd9e..e197e7fa 100644 --- a/jax_galsim/interpolatedimage.py +++ b/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"]: diff --git a/tests/jax/test_dtype_consistency.py b/tests/jax/test_dtype_consistency.py new file mode 100644 index 00000000..8ecbb5bd --- /dev/null +++ b/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