From c459c74dc39ca26b3d539ec11a97870c760a640d Mon Sep 17 00:00:00 2001 From: Biswajit Biswas Date: Tue, 22 Sep 2026 12:02:09 -0500 Subject: [PATCH] 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)