Add DES_PSFEx (PSFEx model reader) to jax_galsim.des

jax_galsim/__init__.py

@@ -102,6 +102,7 @@
 from . import bessel
 from . import fits
 from . import integ
+from . import des
 
 # this one is specific to jax_galsim
 from . import core

jax_galsim/des/__init__.py

@@ -0,0 +1 @@
+from .des_psfex import DES_PSFEx

jax_galsim/des/des_psfex.py

@@ -0,0 +1,225 @@
+# This is a JAX port of galsim.des.des_psfex (galsim/des/des_psfex.py).
+# The reading of the PSFEx file is unchanged host-side I/O; the per-position
+# PSF evaluation (getPSFArray) is reimplemented in JAX so it can be jitted,
+# vmapped, and differentiated with respect to the image position.
+import os
+
+import galsim as _galsim
+import galsim.des  # noqa: F401  (populates _galsim.des for @implements below)
+import jax.numpy as jnp
+import numpy as np
+from jax.tree_util import register_pytree_node_class
+
+from jax_galsim._pyfits import pyfits
+from jax_galsim.core.utils import ensure_hashable, implements
+from jax_galsim.errors import GalSimIncompatibleValuesError
+from jax_galsim.fits import FitsHeader
+from jax_galsim.image import Image
+from jax_galsim.interpolant import Lanczos
+from jax_galsim.interpolatedimage import InterpolatedImage
+from jax_galsim.wcs import readFromFitsHeader
+
+LAX_DES_PSFEX = """\
+The JAX-GalSim version of ``DES_PSFEx`` does not register itself with the
+GalSim config framework (the ``des_psfex`` input type and ``DES_PSFEx`` object
+type are not available), since JAX-GalSim does not implement config processing.
+
+As a PyTree, all of the data read from the PSFEx file (the PCA basis and the
+polynomial fit parameters) is static auxiliary data stored as NumPy arrays,
+and only the ``wcs`` is traced. Autodiff is therefore supported with respect
+to the ``image_pos`` argument (via ``getPSFArray``) and the ``wcs``, but not
+with respect to the fixed PSFEx calibration data.
+"""
+
+
+@implements(_galsim.des.DES_PSFEx, lax_description=LAX_DES_PSFEX, module="galsim.des")
+@register_pytree_node_class
+class DES_PSFEx:
+    _req_params = {"file_name": str}
+    _opt_params = {"dir": str, "image_file_name": str}
+    _single_params = []
+    _takes_rng = False
+
+    def __init__(self, file_name, image_file_name=None, wcs=None, dir=None):
+        if dir:
+            if not isinstance(file_name, str):
+                raise TypeError("file_name must be a string")
+            file_name = os.path.join(dir, file_name)
+            if image_file_name is not None:
+                image_file_name = os.path.join(dir, image_file_name)
+        self.file_name = file_name
+        if image_file_name:
+            if wcs is not None:
+                raise GalSimIncompatibleValuesError(
+                    "Cannot provide both image_file_name and wcs",
+                    image_file_name=image_file_name,
+                    wcs=wcs,
+                )
+            header = FitsHeader(file_name=image_file_name)
+            wcs, origin = readFromFitsHeader(header)
+            self.wcs = wcs
+        elif wcs:
+            self.wcs = wcs
+        else:
+            self.wcs = None
+        self.read()
+
+    def read(self):
+        if isinstance(self.file_name, str):
+            hdu_list = pyfits.open(self.file_name)
+            hdu = hdu_list[1]
+        else:
+            hdu = self.file_name
+            hdu_list = None
+        pol_naxis = hdu.header["POLNAXIS"]
+
+        pol_name1 = hdu.header["POLNAME1"]
+        pol_name2 = hdu.header["POLNAME2"]
+
+        pol_zero1 = hdu.header["POLZERO1"]
+        pol_zero2 = hdu.header["POLZERO2"]
+        pol_scal1 = hdu.header["POLSCAL1"]
+        pol_scal2 = hdu.header["POLSCAL2"]
+
+        pol_ngrp = hdu.header["POLNGRP"]
+        pol_group1 = hdu.header["POLGRP1"]
+        pol_group2 = hdu.header["POLGRP2"]
+        pol_deg = hdu.header["POLDEG1"]
+
+        psf_naxis = hdu.header["PSFNAXIS"]
+        psf_axis1 = hdu.header["PSFAXIS1"]
+        psf_axis2 = hdu.header["PSFAXIS2"]
+        psf_axis3 = hdu.header["PSFAXIS3"]
+        psf_samp = hdu.header["PSF_SAMP"]
+
+        basis = hdu.data.field("PSF_MASK")[0]
+
+        if hdu_list:
+            hdu_list.close()
+
+        try:
+            assert pol_naxis == 2
+            assert pol_name1.startswith("X") and pol_name1.endswith("IMAGE")
+            assert pol_name2.startswith("Y") and pol_name2.endswith("IMAGE")
+            assert pol_ngrp == 1
+            assert pol_group1 == 1
+            assert pol_group2 == 1
+            assert psf_naxis == 3
+            assert psf_axis3 == ((pol_deg + 1) * (pol_deg + 2)) // 2
+            assert basis.shape[0] == psf_axis3
+            assert basis.shape[1] == psf_axis2
+            assert basis.shape[2] == psf_axis1
+        except AssertionError as e:
+            raise OSError("PSFEx file %s is not as expected.\n%r" % (self.file_name, e))
+
+        # All data read from the PSFEx file is fixed calibration and is kept as
+        # static NumPy (it becomes auxiliary PyTree data, not traced leaves).
+        # PSFEx stores the cube as big-endian float32; cast to native float32
+        # (galsim casts the combined array to float32 anyway).
+        self.basis = np.ascontiguousarray(basis, dtype=np.float32)
+        self.fit_order = int(pol_deg)
+        self.fit_size = int(psf_axis3)
+        self.x_zero = pol_zero1
+        self.y_zero = pol_zero2
+        self.x_scale = pol_scal1
+        self.y_scale = pol_scal2
+        self.sample_scale = psf_samp
+
+    @implements(_galsim.des.DES_PSFEx.getSampleScale)
+    def getSampleScale(self):
+        return self.sample_scale
+
+    @implements(_galsim.des.DES_PSFEx.getLocalWCS)
+    def getLocalWCS(self, image_pos):
+        if self.wcs:
+            return self.wcs.local(image_pos)
+        else:
+            return None
+
+    @implements(_galsim.des.DES_PSFEx.getPSF)
+    def getPSF(self, image_pos, gsparams=None):
+        im = Image(self.getPSFArray(image_pos))
+        psf = InterpolatedImage(
+            im,
+            scale=self.sample_scale,
+            flux=1,
+            x_interpolant=Lanczos(3),
+            gsparams=gsparams,
+        )
+        if self.wcs:
+            psf = self.wcs.toWorld(psf, image_pos=image_pos)
+        return psf
+
+    @implements(_galsim.des.DES_PSFEx.getPSFArray)
+    def getPSFArray(self, image_pos):
+        xto = self._powers((image_pos.x - self.x_zero) / self.x_scale)
+        yto = self._powers((image_pos.y - self.y_zero) / self.y_scale)
+        order = self.fit_order
+        # order is a static Python int, so this comprehension is unrolled at
+        # trace time; it mirrors galsim's ordering of the polynomial terms.
+        P = jnp.stack(
+            [
+                xto[nx] * yto[ny]
+                for ny in range(order + 1)
+                for nx in range(order + 1 - ny)
+            ]
+        )
+        # basis is static NumPy; jnp folds it in as a compile-time constant.
+        return jnp.tensordot(P, jnp.asarray(self.basis), (0, 0)).astype(jnp.float32)
+
+    def _powers(self, x):
+        # JAX-safe replacement for galsim's ``np.empty`` + in-place loop: build
+        # [1, x, x**2, ..., x**order] via a cumulative product (same recurrence
+        # as galsim, but without an in-place update, which JAX forbids).
+        return jnp.concatenate(
+            [
+                jnp.ones((1,), dtype=jnp.result_type(float)),
+                jnp.cumprod(jnp.full((self.fit_order,), x)),
+            ]
+        )
+
+    def tree_flatten(self):
+        """Flatten into traced children and static auxiliary data.
+
+        Only ``wcs`` is traced. All of the data read from the PSFEx file is
+        auxiliary; the basis array is passed through ``ensure_hashable`` so the
+        PyTree metadata stays hashable when instances are used as arguments to
+        transformed functions.
+        """
+        children = (self.wcs,)
+        aux_data = {
+            "file_name": self.file_name,
+            "basis": ensure_hashable(jnp.asarray(self.basis)),
+            "basis_shape": self.basis.shape,
+            "fit_order": self.fit_order,
+            "fit_size": self.fit_size,
+            "x_zero": ensure_hashable(self.x_zero),
+            "y_zero": ensure_hashable(self.y_zero),
+            "x_scale": ensure_hashable(self.x_scale),
+            "y_scale": ensure_hashable(self.y_scale),
+            "sample_scale": ensure_hashable(self.sample_scale),
+        }
+        return children, aux_data
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        """Rebuild an instance without re-reading the file.
+
+        ``__init__`` opens the PSFEx file, so (following ``CelestialCoord`` /
+        ``Image``) we construct via ``object.__new__`` and restore attributes
+        directly from the flattened representation.
+        """
+        obj = object.__new__(cls)
+        (obj.wcs,) = children
+        obj.file_name = aux_data["file_name"]
+        obj.basis = np.asarray(aux_data["basis"], dtype=np.float32).reshape(
+            aux_data["basis_shape"]
+        )
+        obj.fit_order = aux_data["fit_order"]
+        obj.fit_size = aux_data["fit_size"]
+        obj.x_zero = aux_data["x_zero"]
+        obj.y_zero = aux_data["y_zero"]
+        obj.x_scale = aux_data["x_scale"]
+        obj.y_scale = aux_data["y_scale"]
+        obj.sample_scale = aux_data["sample_scale"]
+        return obj

tests/jax/test_des_psfex_jax.py

@@ -0,0 +1,128 @@
+import os
+
+import galsim as _galsim
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+from galsim.utilities import timer
+
+import jax_galsim as galsim
+
+DES_DATA_DIR = os.path.join(
+    os.path.dirname(__file__), "..", "GalSim", "tests", "des_data"
+)
+PSFEX_FILE = "DECam_00154912_12_psfcat.psf"
+
+# A few positions spread across the DECam chip.
+POSITIONS = [(100.0, 100.0), (456.0, 789.0), (1024.0, 2048.0), (1700.0, 3500.0)]
+
+
+def _have_des_data():
+    return os.path.isfile(os.path.join(DES_DATA_DIR, PSFEX_FILE))
+
+
+requires_des_data = pytest.mark.skipif(
+    not _have_des_data(),
+    reason="DES test data (tests/GalSim submodule) not available",
+)
+
+
+@requires_des_data
+@timer
+def test_des_psfex_getPSFArray_vs_galsim():
+    """The interpolated PSF array should match reference GalSim."""
+    ref = _galsim.des.DES_PSFEx(PSFEX_FILE, dir=DES_DATA_DIR)
+    jgs = galsim.des.DES_PSFEx(PSFEX_FILE, dir=DES_DATA_DIR)
+
+    assert jgs.fit_order == ref.fit_order
+    assert jgs.fit_size == ref.fit_size
+    np.testing.assert_allclose(jgs.sample_scale, ref.sample_scale)
+
+    for x, y in POSITIONS:
+        a = np.asarray(jgs.getPSFArray(galsim.PositionD(x, y)))
+        b = ref.getPSFArray(_galsim.PositionD(x, y))
+        # float32 interpolation, so compare at ~single precision.
+        np.testing.assert_allclose(a, b, rtol=0, atol=1e-6)
+
+
+@requires_des_data
+@timer
+def test_des_psfex_getPSF_drawn_image_vs_galsim():
+    """The effective-PSF image (drawn with no_pixel) should match GalSim."""
+    ref = _galsim.des.DES_PSFEx(PSFEX_FILE, dir=DES_DATA_DIR)
+    jgs = galsim.des.DES_PSFEx(PSFEX_FILE, dir=DES_DATA_DIR)
+
+    for x, y in POSITIONS:
+        # PSFEx PSFs already include the pixel, so draw with method='no_pixel'.
+        gimg = ref.getPSF(_galsim.PositionD(x, y)).drawImage(
+            nx=25, ny=25, scale=0.2, method="no_pixel"
+        )
+        jimg = jgs.getPSF(galsim.PositionD(x, y)).drawImage(
+            nx=25, ny=25, scale=0.2, method="no_pixel"
+        )
+        np.testing.assert_allclose(
+            np.asarray(jimg.array), gimg.array, rtol=0, atol=1e-6
+        )
+
+
+@requires_des_data
+@timer
+def test_des_psfex_pytree_roundtrip_and_traced_arg():
+    """DES_PSFEx is a registered PyTree: it round-trips through flatten/
+    unflatten (without re-reading the file) and can be passed as an argument to
+    a transformed function, including two distinct-but-equal instances."""
+    jgs = galsim.des.DES_PSFEx(PSFEX_FILE, dir=DES_DATA_DIR)
+
+    leaves, treedef = jax.tree_util.tree_flatten(jgs)
+    rebuilt = jax.tree_util.tree_unflatten(treedef, leaves)
+    np.testing.assert_array_equal(np.asarray(rebuilt.basis), np.asarray(jgs.basis))
+    for x, y in POSITIONS:
+        np.testing.assert_allclose(
+            np.asarray(rebuilt.getPSFArray(galsim.PositionD(x, y))),
+            np.asarray(jgs.getPSFArray(galsim.PositionD(x, y))),
+            rtol=0,
+            atol=1e-6,
+        )
+
+    # Pass the object itself as a jitted argument. Using two distinct instances
+    # exercises the hashability of the (auxiliary) PSFEx data in the treedef.
+    f = jax.jit(lambda obj, x, y: obj.getPSFArray(galsim.PositionD(x, y)))
+    jgs2 = galsim.des.DES_PSFEx(PSFEX_FILE, dir=DES_DATA_DIR)
+    a1 = f(jgs, 456.0, 789.0)
+    a2 = f(jgs2, 456.0, 789.0)
+    np.testing.assert_allclose(np.asarray(a1), np.asarray(a2), rtol=0, atol=1e-6)
+
+
+@requires_des_data
+@timer
+def test_des_psfex_is_jittable_vmappable_differentiable():
+    """getPSFArray should support jit, vmap, and grad over the image position."""
+    jgs = galsim.des.DES_PSFEx(PSFEX_FILE, dir=DES_DATA_DIR)
+
+    def psf_sum(x, y):
+        return jnp.sum(jgs.getPSFArray(galsim.PositionD(x, y)))
+
+    # jit
+    jitted = jax.jit(lambda x, y: jgs.getPSFArray(galsim.PositionD(x, y)))
+    arr = jitted(456.0, 789.0)
+    ref = jgs.getPSFArray(galsim.PositionD(456.0, 789.0))
+    np.testing.assert_allclose(np.asarray(arr), np.asarray(ref), rtol=0, atol=1e-6)
+
+    # vmap over a batch of positions
+    xs = jnp.array([p[0] for p in POSITIONS])
+    ys = jnp.array([p[1] for p in POSITIONS])
+    batched = jax.vmap(lambda x, y: jgs.getPSFArray(galsim.PositionD(x, y)))(xs, ys)
+    assert batched.shape[0] == len(POSITIONS)
+
+    # grad w.r.t. position must be finite (the PSF is differentiable in position)
+    gx = jax.grad(psf_sum, argnums=0)(456.0, 789.0)
+    assert np.isfinite(float(gx))
+
+
+if __name__ == "__main__":
+    test_des_psfex_getPSFArray_vs_galsim()
+    test_des_psfex_getPSF_drawn_image_vs_galsim()
+    test_des_psfex_pytree_roundtrip_and_traced_arg()
+    test_des_psfex_is_jittable_vmappable_differentiable()
+    print("all DES_PSFEx tests passed")