diff --git a/jax_galsim/fits.py b/jax_galsim/fits.py index 1b67f03a..8c7837f9 100644 --- a/jax_galsim/fits.py +++ b/jax_galsim/fits.py @@ -19,7 +19,7 @@ def _maybe_convert_and_warn(image): "The dtype of the input image is not supported by jax_galsim. " "Converting to float64." ) - _image = image.view(dtype=jnp.float64) + _image = image.copy(dtype=jnp.float64) if hasattr(image, "header"): _image.header = image.header return _image diff --git a/jax_galsim/image.py b/jax_galsim/image.py index 7b2bff36..3e8c0e69 100644 --- a/jax_galsim/image.py +++ b/jax_galsim/image.py @@ -18,14 +18,11 @@ IMAGE_LAX_DOCS = """\ Contrary to GalSim native Image, this implementation does not support -sharing of the underlying numpy array between different Images or Views. +sharing of the underlying array between different Images or views. This is due to the fact that in JAX numpy arrays are immutable, so any -operation applied to this Image will create a new ``jnp.ndarray``. - -In particular the following methods will create a copy of the Image: - -- ``Image.view()`` -- ``Image.subImage()`` +operation applied to this Image will create a new ``jnp.ndarray``. Making +a view via ``.view()`` will raise an error. Instead, use the ``.copy()`` +method. The ``Image.subImage()`` method will return a copy. """ @@ -487,9 +484,72 @@ def imag(self): def conjugate(self): return self.__class__(self.array.conjugate(), bounds=self.bounds, wcs=self.wcs) - @implements(_galsim.Image.copy) - def copy(self): - return self.__class__(self.array.copy(), bounds=self.bounds, wcs=self.wcs) + @implements( + _galsim.Image.copy, + lax_description=( + "JAX-GalSim supports extra keyword arguments to ``.copy`` so " + "that users can make copies of images while also changing the image " + "properties (e.g., the wcs). The extra keywords behave exactly like " + "those of ``Image.view``." + ), + ) + def copy( + self, + scale=None, + wcs=None, + origin=None, + center=None, + dtype=None, + make_const=False, + contiguous=False, + ): + if origin is not None and center is not None: + raise _galsim.GalSimIncompatibleValuesError( + "Cannot provide both center and origin", center=center, origin=origin + ) + + if scale is not None: + if wcs is not None: + raise _galsim.GalSimIncompatibleValuesError( + "Cannot provide both scale and wcs", scale=scale, wcs=wcs + ) + wcs = PixelScale(scale) + elif wcs is not None: + if not isinstance(wcs, BaseWCS): + raise TypeError("wcs parameters must be a galsim.BaseWCS instance") + else: + wcs = self.wcs + + # Figure out the dtype for the return Image + dtype = dtype if dtype else self.dtype + + # If currently empty, just return a new empty image. + if not self.bounds.isDefined(): + return Image(wcs=wcs, dtype=dtype, make_const=make_const) + + # Recast the array type if necessary + array = self.array.copy() + if dtype != array.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 contiguous: + # this is a noop since all jax arrays are contiguous + pass + else: + # do nothing here since we made copy above + pass + + # Make the return Image - already made copy above + ret = self.__class__(array, bounds=self.bounds, wcs=wcs, make_const=make_const) + + # Update the origin if requested + if origin is not None: + ret.setOrigin(origin) + elif center is not None: + ret.setCenter(center) + + return ret @implements(_galsim.Image.get_pixel_centers) def get_pixel_centers(self): @@ -930,7 +990,10 @@ def _copyFrom(self, rhs): @implements( _galsim.Image.view, - lax_description="Contrary to GalSim, this will create a copy of the orginal image.", + lax_description=( + "JAX-GalSim does not support image views. This " + "method will raise an error if called." + ), ) def view( self, @@ -942,51 +1005,9 @@ def view( make_const=False, contiguous=False, ): - if origin is not None and center is not None: - raise _galsim.GalSimIncompatibleValuesError( - "Cannot provide both center and origin", center=center, origin=origin - ) - - if scale is not None: - if wcs is not None: - raise _galsim.GalSimIncompatibleValuesError( - "Cannot provide both scale and wcs", scale=scale, wcs=wcs - ) - wcs = PixelScale(scale) - elif wcs is not None: - if not isinstance(wcs, BaseWCS): - raise TypeError("wcs parameters must be a galsim.BaseWCS instance") - else: - wcs = self.wcs - - # Figure out the dtype for the return Image - dtype = dtype if dtype else self.dtype - - # If currently empty, just return a new empty image. - if not self.bounds.isDefined(): - return Image(wcs=wcs, dtype=dtype, make_const=make_const) - - # 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 - pass - else: - array = self.array - - # Make the return Image - ret = self.__class__(array, bounds=self.bounds, wcs=wcs, make_const=make_const) - - # Update the origin if requested - if origin is not None: - ret.setOrigin(origin) - elif center is not None: - ret.setCenter(center) - - return ret + raise NotImplementedError( + "JAX-GalSim does not support views of images! Use ``.copy`` instead." + ) @implements(_galsim.Image.shift) def shift(self, *args, **kwargs): diff --git a/jax_galsim/interpolatedimage.py b/jax_galsim/interpolatedimage.py index f6558c12..0424e9d8 100644 --- a/jax_galsim/interpolatedimage.py +++ b/jax_galsim/interpolatedimage.py @@ -610,7 +610,7 @@ def _image(self): "InterpolatedImages do not support 'depixelize' in jax_galsim." ) else: - image = self._jax_children[0].view(dtype=float) + image = self._jax_children[0].copy(dtype=float) if self._jax_aux_data["_recenter_image"]: image.setCenter(0, 0) diff --git a/tests/GalSim b/tests/GalSim index 2ceb789d..fadb1494 160000 --- a/tests/GalSim +++ b/tests/GalSim @@ -1 +1 @@ -Subproject commit 2ceb789db4c04644739c58bd629ee9fd79ae1ab8 +Subproject commit fadb1494a8599af328922719b62b421602d6f8e7 diff --git a/tests/galsim_tests_config.yaml b/tests/galsim_tests_config.yaml index 2f96d201..46a64c18 100644 --- a/tests/galsim_tests_config.yaml +++ b/tests/galsim_tests_config.yaml @@ -146,3 +146,4 @@ allowed_failures: - "module 'jax_galsim' has no attribute 'zernike'" - "Invalid TFORM4: 1PE(7)" # see https://github.com/astropy/astropy/issues/15477 - "JAX-GalSim does not support truncated Moffat" + - "JAX-GalSim does not support views of images! Use ``.copy`` instead." diff --git a/tests/jax/test_image_jax.py b/tests/jax/test_image_jax.py new file mode 100644 index 00000000..db14bb8c --- /dev/null +++ b/tests/jax/test_image_jax.py @@ -0,0 +1,10 @@ +import numpy as np +import pytest + +import jax_galsim + + +def test_image_jax_view_raises(): + im = jax_galsim.ImageD(np.arange(20).reshape(4, 5)) + with pytest.raises(NotImplementedError): + im.view()