fix: round properly when rendering int images

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)

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:

jax_galsim/image.py

@@ -150,15 +150,19 @@ 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.",
                         array.dtype,
                         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)
         )

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)

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()
 

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):

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)