From c459c74dc39ca26b3d539ec11a97870c760a640d Mon Sep 17 00:00:00 2001 From: Biswajit Biswas Date: Tue, 22 Sep 2026 12:02:09 -0500 Subject: [PATCH 01/11] keeping float32 --- jax_galsim/image.py | 8 +++++++- jax_galsim/interpolatedimage.py | 3 ++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/jax_galsim/image.py b/jax_galsim/image.py index 50e1a7ed..913612e1 100644 --- a/jax_galsim/image.py +++ b/jax_galsim/image.py @@ -877,9 +877,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 diff --git a/jax_galsim/interpolatedimage.py b/jax_galsim/interpolatedimage.py index 589c27a9..65913c3d 100644 --- a/jax_galsim/interpolatedimage.py +++ b/jax_galsim/interpolatedimage.py @@ -634,7 +634,8 @@ def _image(self): "InterpolatedImages do not support 'depixelize' in jax_galsim." ) else: - image = self._jax_children[0].copy(dtype=float) + # Preserve the input dtype + image = self._jax_children[0].copy() if self._jax_aux_data["_recenter_image"]: image.setCenter(0, 0) From 5223137d082d935944c276361813098ddcd78f41 Mon Sep 17 00:00:00 2001 From: Biswajit Biswas Date: Thu, 29 Oct 2026 14:13:37 -0700 Subject: [PATCH 02/11] keeping dtype consistent --- jax_galsim/image.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/jax_galsim/image.py b/jax_galsim/image.py index 913612e1..96a0bca1 100644 --- a/jax_galsim/image.py +++ b/jax_galsim/image.py @@ -959,9 +959,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( From 30a2f4fd7cb6d815a8e4e519d9544c006286aaa9 Mon Sep 17 00:00:00 2001 From: Biswajit Biswas Date: Thu, 29 Oct 2026 14:14:00 -0700 Subject: [PATCH 03/11] add option to set dtype for drawKImage --- jax_galsim/gsobject.py | 31 +++++++++++++++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) diff --git a/jax_galsim/gsobject.py b/jax_galsim/gsobject.py index 641bbdf3..8158cd55 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 ) From 7161edc89ce49168a7397f904d3f4173a3529cb7 Mon Sep 17 00:00:00 2001 From: Biswajit Biswas Date: Thu, 29 Oct 2026 16:05:09 -0700 Subject: [PATCH 04/11] dtype tests --- tests/jax/test_dtype_consistency.py | 79 +++++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) create mode 100644 tests/jax/test_dtype_consistency.py diff --git a/tests/jax/test_dtype_consistency.py b/tests/jax/test_dtype_consistency.py new file mode 100644 index 00000000..a1a2a68b --- /dev/null +++ b/tests/jax/test_dtype_consistency.py @@ -0,0 +1,79 @@ +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 From 94978ea2531aeb039abe8f4d28f625edcfbf260e Mon Sep 17 00:00:00 2001 From: Biswajit Biswas Date: Thu, 29 Oct 2026 16:10:40 -0700 Subject: [PATCH 05/11] minor update --- tests/jax/test_dtype_consistency.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/jax/test_dtype_consistency.py b/tests/jax/test_dtype_consistency.py index a1a2a68b..b368e813 100644 --- a/tests/jax/test_dtype_consistency.py +++ b/tests/jax/test_dtype_consistency.py @@ -1,4 +1,3 @@ -import jax import jax.numpy as jnp import numpy as np import pytest From f15efe1e6a0ce996b4783ede69febedb5b57fe18 Mon Sep 17 00:00:00 2001 From: Biswajit Biswas Date: Tue, 22 Sep 2026 12:02:09 -0500 Subject: [PATCH 06/11] keeping float32 --- jax_galsim/image.py | 8 +++++++- jax_galsim/interpolatedimage.py | 3 ++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/jax_galsim/image.py b/jax_galsim/image.py index 50e1a7ed..913612e1 100644 --- a/jax_galsim/image.py +++ b/jax_galsim/image.py @@ -877,9 +877,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 diff --git a/jax_galsim/interpolatedimage.py b/jax_galsim/interpolatedimage.py index a2a6bd9e..9fba8631 100644 --- a/jax_galsim/interpolatedimage.py +++ b/jax_galsim/interpolatedimage.py @@ -634,7 +634,8 @@ def _image(self): "InterpolatedImages do not support 'depixelize' in jax_galsim." ) else: - image = self._jax_children[0].copy(dtype=float) + # Preserve the input dtype + image = self._jax_children[0].copy() if self._jax_aux_data["_recenter_image"]: image.setCenter(0, 0) From 104d3e29fe5c35007bd8c1069fa97f7211c0899e Mon Sep 17 00:00:00 2001 From: Ismael Mendoza <11745764+ismael-mendoza@users.noreply.github.com> Date: Mon, 26 Oct 2026 10:18:08 -0700 Subject: [PATCH 07/11] ismael : some investigation on dtypes --- jax_galsim/image.py | 4 ++- scripts/check_dtype_consistency.py | 55 ++++++++++++++++++++++++++++++ tests/jax/test_image_jax.py | 18 ++++++++++ 3 files changed, 76 insertions(+), 1 deletion(-) create mode 100644 scripts/check_dtype_consistency.py diff --git a/jax_galsim/image.py b/jax_galsim/image.py index 913612e1..c07b6b96 100644 --- a/jax_galsim/image.py +++ b/jax_galsim/image.py @@ -269,6 +269,8 @@ def __init__(self, *args, **kwargs): raise TypeError("wcs parameters must be a galsim.BaseWCS instance") self.wcs = wcs + 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""" @@ -1291,7 +1293,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/scripts/check_dtype_consistency.py b/scripts/check_dtype_consistency.py new file mode 100644 index 00000000..746045f6 --- /dev/null +++ b/scripts/check_dtype_consistency.py @@ -0,0 +1,55 @@ +import jax +import jax.numpy as jnp + +from jax_galsim.image import Image +from jax_galsim.interpolatedimage import InterpolatedImage + + +def describe_image(label, image): + print(f"{label}: Image.dtype={image.dtype}, Image.array.dtype={image.array.dtype}") + + +def check_mode(enable_x64): + jax.config.update("jax_enable_x64", enable_x64) + print(f"\njax_enable_x64={enable_x64}") + + image = Image(2, 2, dtype=jnp.float64) + describe_image("requested float64", image) + + leaves, treedef = jax.tree_util.tree_flatten(image) + rebuilt = jax.tree_util.tree_unflatten(treedef, leaves) + describe_image("after PyTree roundtrip", rebuilt) + + copied = image.copy() + describe_image("after copy", copied) + print() + + float32_image = Image(jnp.ones((4, 4), dtype=jnp.float32), scale=1.0) + fourier_image = float32_image.calculate_fft() + inverse_image = fourier_image.calculate_inverse_fft() + describe_image("float32 FFT", fourier_image) + describe_image("inverse FFT", inverse_image) + print() + + profile = InterpolatedImage(float32_image, scale=1.0) + implementation = profile._original + print( + "InterpolatedImage: source array.dtype=" + f"{float32_image.array.dtype}, profile image array.dtype=" + f"{implementation._image.array.dtype}, image flux dtype=" + f"{implementation._image_flux.dtype}, flux ratio dtype=" + f"{jnp.asarray(implementation._flux_ratio).dtype}" + ) + + +def main(): + original_x64 = jax.config.read("jax_enable_x64") + try: + check_mode(False) + check_mode(True) + finally: + jax.config.update("jax_enable_x64", original_x64) + + +if __name__ == "__main__": + main() diff --git a/tests/jax/test_image_jax.py b/tests/jax/test_image_jax.py index db14bb8c..a88e6eb4 100644 --- a/tests/jax/test_image_jax.py +++ b/tests/jax/test_image_jax.py @@ -1,3 +1,5 @@ +import jax +import jax.numpy as jnp import numpy as np import pytest @@ -8,3 +10,19 @@ def test_image_jax_view_raises(): im = jax_galsim.ImageD(np.arange(20).reshape(4, 5)) with pytest.raises(NotImplementedError): im.view() + + +def test_image_dtype_tracks_storage_with_x64_disabled(): + original_x64 = jax.config.read("jax_enable_x64") + try: + jax.config.update("jax_enable_x64", False) + with pytest.warns(UserWarning, match="truncated to dtype float32"): + 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 + finally: + jax.config.update("jax_enable_x64", original_x64) From edfa4fdacda5b18a874fbbdcdb21f942a3637fab Mon Sep 17 00:00:00 2001 From: Biswajit Biswas Date: Fri, 30 Oct 2026 13:12:10 -0700 Subject: [PATCH 08/11] pre-commit fix [skip-ci] --- jax_galsim/gsobject.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jax_galsim/gsobject.py b/jax_galsim/gsobject.py index 8158cd55..fe41760a 100644 --- a/jax_galsim/gsobject.py +++ b/jax_galsim/gsobject.py @@ -989,7 +989,7 @@ def drawFFT(self, image, add_to_image=False): @implements( _galsim.GSObject.drawKImage, lax_description="""\ -The JAX-GalSim version of ``drawKImage`` accepts an extra ``dtype`` to set +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. From 085193fc263cfbec95c85fd48f9349fc74ac0a3d Mon Sep 17 00:00:00 2001 From: Ismael Mendoza <11745764+ismael-mendoza@users.noreply.github.com> Date: Fri, 30 Oct 2026 13:33:03 -0700 Subject: [PATCH 09/11] make sure dtype array and image are consistnet even with float64 turned off --- jax_galsim/image.py | 1 + scripts/check_dtype_consistency.py | 55 ----------------------------- tests/jax/test_dtype_consistency.py | 11 ++++++ tests/jax/test_image_jax.py | 16 --------- 4 files changed, 12 insertions(+), 71 deletions(-) delete mode 100644 scripts/check_dtype_consistency.py diff --git a/jax_galsim/image.py b/jax_galsim/image.py index 6d6c3983..f3ba8f0c 100644 --- a/jax_galsim/image.py +++ b/jax_galsim/image.py @@ -269,6 +269,7 @@ 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 diff --git a/scripts/check_dtype_consistency.py b/scripts/check_dtype_consistency.py deleted file mode 100644 index 746045f6..00000000 --- a/scripts/check_dtype_consistency.py +++ /dev/null @@ -1,55 +0,0 @@ -import jax -import jax.numpy as jnp - -from jax_galsim.image import Image -from jax_galsim.interpolatedimage import InterpolatedImage - - -def describe_image(label, image): - print(f"{label}: Image.dtype={image.dtype}, Image.array.dtype={image.array.dtype}") - - -def check_mode(enable_x64): - jax.config.update("jax_enable_x64", enable_x64) - print(f"\njax_enable_x64={enable_x64}") - - image = Image(2, 2, dtype=jnp.float64) - describe_image("requested float64", image) - - leaves, treedef = jax.tree_util.tree_flatten(image) - rebuilt = jax.tree_util.tree_unflatten(treedef, leaves) - describe_image("after PyTree roundtrip", rebuilt) - - copied = image.copy() - describe_image("after copy", copied) - print() - - float32_image = Image(jnp.ones((4, 4), dtype=jnp.float32), scale=1.0) - fourier_image = float32_image.calculate_fft() - inverse_image = fourier_image.calculate_inverse_fft() - describe_image("float32 FFT", fourier_image) - describe_image("inverse FFT", inverse_image) - print() - - profile = InterpolatedImage(float32_image, scale=1.0) - implementation = profile._original - print( - "InterpolatedImage: source array.dtype=" - f"{float32_image.array.dtype}, profile image array.dtype=" - f"{implementation._image.array.dtype}, image flux dtype=" - f"{implementation._image_flux.dtype}, flux ratio dtype=" - f"{jnp.asarray(implementation._flux_ratio).dtype}" - ) - - -def main(): - original_x64 = jax.config.read("jax_enable_x64") - try: - check_mode(False) - check_mode(True) - finally: - jax.config.update("jax_enable_x64", original_x64) - - -if __name__ == "__main__": - main() diff --git a/tests/jax/test_dtype_consistency.py b/tests/jax/test_dtype_consistency.py index b368e813..8ecbb5bd 100644 --- a/tests/jax/test_dtype_consistency.py +++ b/tests/jax/test_dtype_consistency.py @@ -1,3 +1,4 @@ +import jax import jax.numpy as jnp import numpy as np import pytest @@ -76,3 +77,13 @@ def test_drawkimage_image_kwarg_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 diff --git a/tests/jax/test_image_jax.py b/tests/jax/test_image_jax.py index a88e6eb4..50ee8e6a 100644 --- a/tests/jax/test_image_jax.py +++ b/tests/jax/test_image_jax.py @@ -10,19 +10,3 @@ def test_image_jax_view_raises(): im = jax_galsim.ImageD(np.arange(20).reshape(4, 5)) with pytest.raises(NotImplementedError): im.view() - - -def test_image_dtype_tracks_storage_with_x64_disabled(): - original_x64 = jax.config.read("jax_enable_x64") - try: - jax.config.update("jax_enable_x64", False) - with pytest.warns(UserWarning, match="truncated to dtype float32"): - 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 - finally: - jax.config.update("jax_enable_x64", original_x64) From 70133519bd869fe690d032466886ee1c4167e988 Mon Sep 17 00:00:00 2001 From: Biswajit Biswas Date: Fri, 30 Oct 2026 13:38:14 -0700 Subject: [PATCH 10/11] going back to GalSim default --- jax_galsim/interpolatedimage.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/jax_galsim/interpolatedimage.py b/jax_galsim/interpolatedimage.py index 9fba8631..e197e7fa 100644 --- a/jax_galsim/interpolatedimage.py +++ b/jax_galsim/interpolatedimage.py @@ -634,8 +634,8 @@ def _image(self): "InterpolatedImages do not support 'depixelize' in jax_galsim." ) else: - # Preserve the input dtype - image = self._jax_children[0].copy() + # 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"]: image.setCenter(0, 0) From 08123368bc86b1f32013e4a1f143a95c7ba6cb44 Mon Sep 17 00:00:00 2001 From: Ismael Mendoza <11745764+ismael-mendoza@users.noreply.github.com> Date: Fri, 30 Oct 2026 13:39:29 -0700 Subject: [PATCH 11/11] unncessary imports --- tests/jax/test_image_jax.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/jax/test_image_jax.py b/tests/jax/test_image_jax.py index 50ee8e6a..db14bb8c 100644 --- a/tests/jax/test_image_jax.py +++ b/tests/jax/test_image_jax.py @@ -1,5 +1,3 @@ -import jax -import jax.numpy as jnp import numpy as np import pytest