diff --git a/jax_galsim/core/draw.py b/jax_galsim/core/draw.py index 4ca931a3..181bf464 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,7 @@ 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) + im = im.astype(image.dtype) # 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..3d72dbab 100644 --- a/jax_galsim/gsobject.py +++ b/jax_galsim/gsobject.py @@ -808,6 +808,14 @@ 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 + ): + # 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: image._array = image._array.at[...].add(temp._array) else: @@ -926,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 25fdd7ec..7b2bff36 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,12 @@ def __init__(self, *args, **kwargs): # e.g. im = ImageF(...) # im2 = ImageD(im) self._dtype = dtype - self._array = image.array.astype(self._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 + ) else: self._array = jnp.zeros(shape=(1, 1), dtype=self._dtype) self._bounds = BoundsI() @@ -365,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) ) @@ -590,8 +601,12 @@ 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( - jnp.astype(rhs.array, self.dtype) + _safe_cast( + rhs.array, jnp.issubdtype(self.dtype, jnp.integer), self.dtype + ) ) else: start_inds = ( @@ -600,7 +615,11 @@ def setSubImage(self, bounds, rhs): ) self._array = jax.lax.dynamic_update_slice( self.array, - jnp.astype(rhs.array, self.dtype), + # 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 + ), start_inds, ) @@ -904,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) ) @@ -947,7 +968,9 @@ def view( # Recast the array type if necessary if dtype != self.array.dtype: - array = self.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(self.array, jnp.issubdtype(dtype, jnp.integer), dtype) elif contiguous: # this is a noop since all jax arrays are contiguous pass @@ -1109,7 +1132,13 @@ def _invertSelf(self): self._array, ) self._array = self._array.at[...].set( - (jnp.where(msk, 0.0, 1.0 / safe_array)).astype(self._array.dtype) + # 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), + self._array.dtype, + ) ) @implements(_galsim.Image.replaceNegative) @@ -1289,7 +1318,9 @@ 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) + # 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 return ret @@ -1428,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) ) @@ -1458,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) ) @@ -1484,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) ) @@ -1516,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) ) @@ -1547,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) ) @@ -1578,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) ) @@ -1595,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) diff --git a/jax_galsim/photon_array.py b/jax_galsim/photon_array.py index 2070f981..336c28a7 100644 --- a/jax_galsim/photon_array.py +++ b/jax_galsim/photon_array.py @@ -1000,7 +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) - _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() 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): diff --git a/tests/GalSim b/tests/GalSim index fee09cae..2ceb789d 160000 --- a/tests/GalSim +++ b/tests/GalSim @@ -1 +1 @@ -Subproject commit fee09cae81d9aeaaa3f37aae9556010ec06fee62 +Subproject commit 2ceb789db4c04644739c58bd629ee9fd79ae1ab8 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)