From fd334b5ebbc9600a28c16080ea04db6def039fe7 Mon Sep 17 00:00:00 2001 From: beckermr Date: Tue, 12 May 2026 06:16:11 -0500 Subject: [PATCH 1/8] fix: round properly --- jax_galsim/core/draw.py | 16 ++++++++++++++-- jax_galsim/gsobject.py | 6 ++++++ jax_galsim/image.py | 33 ++++++++++++++++++++++++--------- tests/GalSim | 2 +- 4 files changed, 45 insertions(+), 12 deletions(-) diff --git a/jax_galsim/core/draw.py b/jax_galsim/core/draw.py index 4ca931a3..d3f456c3 100644 --- a/jax_galsim/core/draw.py +++ b/jax_galsim/core/draw.py @@ -34,7 +34,14 @@ def draw_by_xValue( ) # Apply the flux scaling - im = (im * flux_scaling).astype(image.dtype) + im *= flux_scaling + + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed + if jnp.issubdtype(im.dtype, jnp.floating) and jnp.issubdtype( + image.dtype, jnp.integer + ): + im = jnp.around(im) # Return an image return Image(array=im, bounds=image.bounds, wcs=image.wcs, _check_bounds=False) @@ -53,7 +60,12 @@ def draw_by_kValue(gsobject, image, jacobian=jnp.eye(2)): im = jax.vmap(lambda *args: gsobject._kValue(PositionD(*args)))( coords[..., 0], coords[..., 1] ) - im = (im).astype(image.dtype) + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed + if jnp.issubdtype(im.dtype, jnp.floating) and jnp.issubdtype( + image.dtype, jnp.integer + ): + im = jnp.around(im) # Return an image return Image(array=im, bounds=image.bounds, wcs=image.wcs, _check_bounds=False) diff --git a/jax_galsim/gsobject.py b/jax_galsim/gsobject.py index c14077fe..1a512926 100644 --- a/jax_galsim/gsobject.py +++ b/jax_galsim/gsobject.py @@ -808,6 +808,12 @@ def drawReal(self, image, add_to_image=False): ) im1 = self._drawReal(image) temp = im1.subImage(image.bounds) + + if jnp.issubdtype(temp.array.dtype, jnp.floating) and jnp.issubdtype( + image.array.dtype, jnp.integer + ): + temp.array = jnp.around(temp.array) + if add_to_image: image._array = image._array.at[...].add(temp._array) else: diff --git a/jax_galsim/image.py b/jax_galsim/image.py index 25fdd7ec..872e22db 100644 --- a/jax_galsim/image.py +++ b/jax_galsim/image.py @@ -150,7 +150,9 @@ def __init__(self, *args, **kwargs): dtype = array.dtype.type if dtype in self._alias_dtypes: dtype = self._alias_dtypes[dtype] - array = array.astype(dtype) + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed + array = _safe_cast(array, jnp.issubdtype(dtype, jnp.integer), dtype) elif dtype not in self._valid_dtypes: raise _galsim.GalSimValueError( "Invalid dtype of provided array.", @@ -158,7 +160,9 @@ def __init__(self, *args, **kwargs): self._valid_dtypes, ) else: - array = array.astype(dtype) + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed + array = _safe_cast(array, jnp.issubdtype(dtype, jnp.integer), dtype) # Be careful here: we have to watch out for little-endian / big-endian issues. # The path of least resistance is to check whether the array.dtype is equal to the # native one (using the dtype.isnative flag), and if not, make a new array that has a @@ -206,7 +210,7 @@ def __init__(self, *args, **kwargs): if init_value: self._array = self._array.at[...].add(init_value) elif array is not None: - self._array = array.view(dtype=self._dtype) + self._array = array.view() nrow, ncol = array.shape if not has_tracers(xmin) and not has_tracers(ymin): self._bounds = BoundsI( @@ -239,7 +243,10 @@ def __init__(self, *args, **kwargs): # e.g. im = ImageF(...) # im2 = ImageD(im) self._dtype = dtype - self._array = image.array.astype(self._dtype) + + self._array = _safe_cast( + image.array, jnp.issubdtype(self._dtype, jnp.integer), self._dtype + ) else: self._array = jnp.zeros(shape=(1, 1), dtype=self._dtype) self._bounds = BoundsI() @@ -591,7 +598,9 @@ def setSubImage(self, bounds, rhs): j1 = bounds.xmin - self.xmin j2 = bounds.xmax - self.xmin + 1 self._array = self._array.at[i1:i2, j1:j2].set( - jnp.astype(rhs.array, self.dtype) + _safe_cast( + rhs.array, jnp.issubdtype(self.dtype, jnp.integer), self.dtype + ) ) else: start_inds = ( @@ -600,7 +609,9 @@ def setSubImage(self, bounds, rhs): ) self._array = jax.lax.dynamic_update_slice( self.array, - jnp.astype(rhs.array, self.dtype), + _safe_cast( + rhs.array, jnp.issubdtype(self.dtype, jnp.integer), self.dtype + ), start_inds, ) @@ -947,7 +958,7 @@ def view( # Recast the array type if necessary if dtype != self.array.dtype: - array = self.array.astype(dtype) + array = _safe_cast(self.array, jnp.issubdtype(dtype, jnp.integer), dtype) elif contiguous: # this is a noop since all jax arrays are contiguous pass @@ -1109,7 +1120,11 @@ def _invertSelf(self): self._array, ) self._array = self._array.at[...].set( - (jnp.where(msk, 0.0, 1.0 / safe_array)).astype(self._array.dtype) + _safe_cast( + (jnp.where(msk, 0.0, 1.0 / safe_array)), + jnp.issubdtype(self._array.dtype, jnp.integer), + self._array.dtype, + ) ) @implements(_galsim.Image.replaceNegative) @@ -1289,7 +1304,7 @@ def _Image(array, bounds, wcs): ret._dtype = array.dtype.type if ret._dtype in Image._alias_dtypes: ret._dtype = Image._alias_dtypes[ret._dtype] - array = array.astype(ret._dtype) + array = _safe_cast(array, jnp.issubdtype(ret._dtype, jnp.integer), ret._dtype) ret._array = array ret._bounds = bounds return ret diff --git a/tests/GalSim b/tests/GalSim index fee09cae..5724050f 160000 --- a/tests/GalSim +++ b/tests/GalSim @@ -1 +1 @@ -Subproject commit fee09cae81d9aeaaa3f37aae9556010ec06fee62 +Subproject commit 5724050f6efe8ae766b15f5aec845a05666bbbe7 From 52922c17a20c3fa45b49f7af6f059bafc9acb6ee Mon Sep 17 00:00:00 2001 From: beckermr Date: Tue, 12 May 2026 06:55:56 -0500 Subject: [PATCH 2/8] fix: do not cast in draewing by k value --- jax_galsim/core/draw.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/jax_galsim/core/draw.py b/jax_galsim/core/draw.py index d3f456c3..181bf464 100644 --- a/jax_galsim/core/draw.py +++ b/jax_galsim/core/draw.py @@ -60,12 +60,7 @@ def draw_by_kValue(gsobject, image, jacobian=jnp.eye(2)): im = jax.vmap(lambda *args: gsobject._kValue(PositionD(*args)))( coords[..., 0], coords[..., 1] ) - # jax-galsim's rounding of float-to-int is platform dependent - # so we explicitly round to ints if needed - if jnp.issubdtype(im.dtype, jnp.floating) and jnp.issubdtype( - image.dtype, jnp.integer - ): - im = jnp.around(im) + im = im.astype(image.dtype) # Return an image return Image(array=im, bounds=image.bounds, wcs=image.wcs, _check_bounds=False) From 38eddfab11d1417700f0fb22043132cf2b6dd311 Mon Sep 17 00:00:00 2001 From: beckermr Date: Tue, 12 May 2026 07:02:33 -0500 Subject: [PATCH 3/8] fix: more rounding --- jax_galsim/gsobject.py | 10 ++++++++++ jax_galsim/image.py | 30 ++++++++++++++++++++++++++++++ jax_galsim/interpolatedimage.py | 11 +++++++++-- 3 files changed, 49 insertions(+), 2 deletions(-) diff --git a/jax_galsim/gsobject.py b/jax_galsim/gsobject.py index 1a512926..3d72dbab 100644 --- a/jax_galsim/gsobject.py +++ b/jax_galsim/gsobject.py @@ -812,6 +812,8 @@ def drawReal(self, image, add_to_image=False): if jnp.issubdtype(temp.array.dtype, jnp.floating) and jnp.issubdtype( image.array.dtype, jnp.integer ): + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed temp.array = jnp.around(temp.array) if add_to_image: @@ -932,6 +934,14 @@ def drawFFT_finish(self, image, kimage, wrap_size, add_to_image): real_image = Image( bounds=breal, array=real_image_arr, dtype=image.dtype, wcs=image.wcs ) + + if jnp.issubdtype(real_image.array.dtype, jnp.floating) and jnp.issubdtype( + image.array.dtype, jnp.integer + ): + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed + real_image.array = jnp.around(real_image.array) + # Add (a portion of) this to the original image. temp = real_image.subImage(image.bounds) if add_to_image: diff --git a/jax_galsim/image.py b/jax_galsim/image.py index 872e22db..7b2bff36 100644 --- a/jax_galsim/image.py +++ b/jax_galsim/image.py @@ -244,6 +244,8 @@ def __init__(self, *args, **kwargs): # im2 = ImageD(im) self._dtype = dtype + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed self._array = _safe_cast( image.array, jnp.issubdtype(self._dtype, jnp.integer), self._dtype ) @@ -372,6 +374,8 @@ def array(self): @array.setter def array(self, other): + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed self._array = self._array.at[...].set( _safe_cast(other, self.isinteger, self.array.dtype) ) @@ -597,6 +601,8 @@ def setSubImage(self, bounds, rhs): i2 = bounds.ymax - self.ymin + 1 j1 = bounds.xmin - self.xmin j2 = bounds.xmax - self.xmin + 1 + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed self._array = self._array.at[i1:i2, j1:j2].set( _safe_cast( rhs.array, jnp.issubdtype(self.dtype, jnp.integer), self.dtype @@ -609,6 +615,8 @@ def setSubImage(self, bounds, rhs): ) self._array = jax.lax.dynamic_update_slice( self.array, + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed _safe_cast( rhs.array, jnp.issubdtype(self.dtype, jnp.integer), self.dtype ), @@ -915,6 +923,8 @@ def copyFrom(self, rhs): def _copyFrom(self, rhs): """Same as copyFrom, but no sanity checks.""" self._array = self._array.at[...].set( + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed _safe_cast(rhs._array, self.isinteger, self.array.dtype) ) @@ -958,6 +968,8 @@ def view( # Recast the array type if necessary if dtype != self.array.dtype: + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed array = _safe_cast(self.array, jnp.issubdtype(dtype, jnp.integer), dtype) elif contiguous: # this is a noop since all jax arrays are contiguous @@ -1120,6 +1132,8 @@ def _invertSelf(self): self._array, ) self._array = self._array.at[...].set( + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed _safe_cast( (jnp.where(msk, 0.0, 1.0 / safe_array)), jnp.issubdtype(self._array.dtype, jnp.integer), @@ -1304,6 +1318,8 @@ def _Image(array, bounds, wcs): ret._dtype = array.dtype.type if ret._dtype in Image._alias_dtypes: ret._dtype = Image._alias_dtypes[ret._dtype] + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed array = _safe_cast(array, jnp.issubdtype(ret._dtype, jnp.integer), ret._dtype) ret._array = array ret._bounds = bounds @@ -1443,6 +1459,8 @@ def Image_iadd(self, other): if dt == self.array.dtype: self._array = self.array.at[...].add(a) else: + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed self._array = self.array.at[...].set( _safe_cast(self.array + a, self.isinteger, self.array.dtype) ) @@ -1473,6 +1491,8 @@ def Image_isub(self, other): if dt == self.array.dtype: self._array = self.array.at[...].subtract(a) else: + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed self._array = self.array.at[...].set( _safe_cast(self.array - a, self.isinteger, self.array.dtype) ) @@ -1499,6 +1519,8 @@ def Image_imul(self, other): if dt == self.array.dtype: self._array = self.array.at[...].multiply(a) else: + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed self._array = self.array.at[...].set( _safe_cast(self.array * a, self.isinteger, self.array.dtype) ) @@ -1531,6 +1553,8 @@ def Image_idiv(self, other): # back to an integer array. So for integers (or mixed types), don't use /=. self._array = self.array.at[...].divide(a) else: + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed self._array = self.array.at[...].set( _safe_cast(self.array / a, self.isinteger, self.array.dtype) ) @@ -1562,6 +1586,8 @@ def Image_ifloordiv(self, other): if dt == self.array.dtype: self._array = self.array.at[...].set(self.array // a) else: + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed self._array = self.array.at[...].set( _safe_cast(self.array // a, self.isinteger, self.array.dtype) ) @@ -1593,6 +1619,8 @@ def Image_imod(self, other): if dt == self.array.dtype: self._array = self.array.at[...].set(self.array % a) else: + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed self._array = self.array.at[...].set( _safe_cast(self.array % a, self.isinteger, self.array.dtype) ) @@ -1610,6 +1638,8 @@ def Image_ipow(self, other): if not self.isinteger or isinstance(other, int): self._array = self.array.at[...].power(other) else: + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed self._array = self.array.at[...].set( _safe_cast(self.array**other, self.isinteger, self.array.dtype) ) diff --git a/jax_galsim/interpolatedimage.py b/jax_galsim/interpolatedimage.py index bd4556ef..f6558c12 100644 --- a/jax_galsim/interpolatedimage.py +++ b/jax_galsim/interpolatedimage.py @@ -792,7 +792,14 @@ def _drawReal(self, image, jac=None, offset=(0.0, 0.0), flux_scaling=1.0): ) # Apply the flux scaling - im = (im * flux_scaling).astype(image.dtype) + im *= flux_scaling + + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed + if jnp.issubdtype(im.dtype, jnp.floating) and jnp.issubdtype( + image.dtype, jnp.integer + ): + im = jnp.around(im) # Return an image return Image(array=im, bounds=image.bounds, wcs=image.wcs, _check_bounds=False) @@ -817,7 +824,7 @@ def _drawKImage(self, image, jac=None): self._x_interpolant, self._k_interpolant, ) - im = (im).astype(image.dtype) + im = im.astype(image.dtype) # Return an image return Image(array=im, bounds=image.bounds, wcs=image.wcs, _check_bounds=False) From 434fe2968880fa9aa910df80345eb3946747f919 Mon Sep 17 00:00:00 2001 From: beckermr Date: Tue, 12 May 2026 07:09:10 -0500 Subject: [PATCH 4/8] fix: one more spot --- jax_galsim/wcs.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/jax_galsim/wcs.py b/jax_galsim/wcs.py index 5eeccae5..20e31eb3 100644 --- a/jax_galsim/wcs.py +++ b/jax_galsim/wcs.py @@ -516,7 +516,16 @@ def _makeSkyImage(self, image, sky_level, color): dvdy = 0.5 * (v[2:ny, 1 : nx - 1] - v[0 : ny - 2, 1 : nx - 1]) area = jnp.abs(dudx * dvdy - dvdx * dudy) - image._array = image._array.at[...].set((area * sky_level).astype(image.dtype)) + im = area * sky_level + + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed + if jnp.issubdtype(im.dtype, jnp.floating) and jnp.issubdtype( + image.dtype, jnp.integer + ): + im = jnp.around(im) + + image._array = image._array.at[...].set(im) # Each class should define the __eq__ function. Then __ne__ is obvious. def __ne__(self, other): From 3611ec0dce7659ef28d2b80ccc3071bc13a1177f Mon Sep 17 00:00:00 2001 From: beckermr Date: Tue, 12 May 2026 07:10:29 -0500 Subject: [PATCH 5/8] doc: add todo for later --- jax_galsim/photon_array.py | 1 + 1 file changed, 1 insertion(+) diff --git a/jax_galsim/photon_array.py b/jax_galsim/photon_array.py index 2070f981..58846d5a 100644 --- a/jax_galsim/photon_array.py +++ b/jax_galsim/photon_array.py @@ -1000,6 +1000,7 @@ def _add_photons_to_image(x, y, flux, xmin, ymin, arr): # dropped and negative indices wrap around good = (xinds >= 0) & (xinds < arr.shape[1]) & (yinds >= 0) & (yinds < arr.shape[0]) _flux = jnp.where(good, flux, 0.0) + # TODO: what dp we do with rounding here? _arr = arr.at[yinds, xinds].add(_flux.astype(arr.dtype)) return _arr, _flux.sum() From 50f588bf5b7fe4a5f63f0c3e8419c46c3beb98bf Mon Sep 17 00:00:00 2001 From: beckermr Date: Tue, 12 May 2026 09:58:01 -0500 Subject: [PATCH 6/8] fix: convert to double to then round --- jax_galsim/photon_array.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/jax_galsim/photon_array.py b/jax_galsim/photon_array.py index 58846d5a..336c28a7 100644 --- a/jax_galsim/photon_array.py +++ b/jax_galsim/photon_array.py @@ -1000,8 +1000,14 @@ def _add_photons_to_image(x, y, flux, xmin, ymin, arr): # dropped and negative indices wrap around good = (xinds >= 0) & (xinds < arr.shape[1]) & (yinds >= 0) & (yinds < arr.shape[0]) _flux = jnp.where(good, flux, 0.0) - # TODO: what dp we do with rounding here? - _arr = arr.at[yinds, xinds].add(_flux.astype(arr.dtype)) + + # jax-galsim's rounding of float-to-int is platform dependent + # so we explicitly round to ints if needed + if jnp.issubdtype(arr.dtype, jnp.integer): + _arr = arr.astype(float).at[yinds, xinds].add(_flux.astype(float)) + _arr = jnp.around(_arr).astype(arr.dtype) + else: + _arr = arr.at[yinds, xinds].add(_flux.astype(arr.dtype)) return _arr, _flux.sum() From f06d6c0e9c9da4ad2f088d0e41a1139b7a7d2071 Mon Sep 17 00:00:00 2001 From: beckermr Date: Tue, 12 May 2026 10:17:37 -0500 Subject: [PATCH 7/8] test: add test for inverting self and dtypes --- tests/jax/test_int_float_dtype_handling.py | 23 ++++++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 tests/jax/test_int_float_dtype_handling.py diff --git a/tests/jax/test_int_float_dtype_handling.py b/tests/jax/test_int_float_dtype_handling.py new file mode 100644 index 00000000..e6aee58b --- /dev/null +++ b/tests/jax/test_int_float_dtype_handling.py @@ -0,0 +1,23 @@ +import galsim as _galsim +import jax.numpy as jnp +import numpy as np + +import jax_galsim + + +def test_int_float_dtype_handling_invertSelf(): + gim = _galsim.Image(np.arange(20).reshape(4, 5), dtype=int) + gim.invertSelf() + + assert gim[1, 1] == 0 + assert gim[2, 1] == 1 + assert gim[4, 4] == 0 + + jgim = jax_galsim.Image(jnp.arange(20).reshape(4, 5), dtype=int) + jgim.invertSelf() + + assert jgim[1, 1] == 0 + assert jgim[2, 1] == 1 + assert jgim[4, 4] == 0 + + np.testing.assert_array_equal(gim.array, jgim.array) From 5b7836e5d7300a7835cb34b49fa3d47eb22fcb69 Mon Sep 17 00:00:00 2001 From: beckermr Date: Tue, 12 May 2026 10:19:06 -0500 Subject: [PATCH 8/8] test: update submodule to latest test --- tests/GalSim | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/GalSim b/tests/GalSim index 5724050f..2ceb789d 160000 --- a/tests/GalSim +++ b/tests/GalSim @@ -1 +1 @@ -Subproject commit 5724050f6efe8ae766b15f5aec845a05666bbbe7 +Subproject commit 2ceb789db4c04644739c58bd629ee9fd79ae1ab8