fix: raise on image views

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

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

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)

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."

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