diff --git a/jax_galsim/__init__.py b/jax_galsim/__init__.py index 80cbe04b..d548c48f 100644 --- a/jax_galsim/__init__.py +++ b/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 diff --git a/jax_galsim/des/__init__.py b/jax_galsim/des/__init__.py new file mode 100644 index 00000000..4fc9e9b5 --- /dev/null +++ b/jax_galsim/des/__init__.py @@ -0,0 +1 @@ +from .des_psfex import DES_PSFEx diff --git a/jax_galsim/des/des_psfex.py b/jax_galsim/des/des_psfex.py new file mode 100644 index 00000000..35474008 --- /dev/null +++ b/jax_galsim/des/des_psfex.py @@ -0,0 +1,239 @@ +# 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 +from jax.tree_util import register_pytree_node_class + +from jax_galsim._pyfits import pyfits +from jax_galsim.core.utils import ( + cast_numpy_array_to_native_byte_order, + cast_to_float, + 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, the data read from the PSFEx file (the PCA basis and the +polynomial zero points and scales) together with the ``wcs`` are traced +children, so a batch of PSFEx models can be stacked and evaluated in a single +``jit``/``vmap`` call. Only ``fit_order`` and ``fit_size`` are static +auxiliary data, since they set the number of polynomial terms and hence the +shapes of the traced arrays; batching therefore applies to models sharing a +polynomial degree, which is the usual case within one instrument and reduction. + +``file_name`` is deliberately *not* part of the PyTree, so that models read +from different files share a tree structure and can be batched together. An +instance rebuilt by ``tree_unflatten`` has ``file_name`` set to ``None``. +""" + + +@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)) + + # The basis and the polynomial zero points/scales are traced children, + # so that models from different files can be batched together. PSFEx + # stores the cube big-endian, which JAX will not accept, so convert to + # the native byte order first. + self.basis = jnp.asarray(cast_numpy_array_to_native_byte_order(basis)) + # fit_order/fit_size stay static: they set the number of polynomial + # terms and hence the shapes of the traced arrays. + self.fit_order = int(pol_deg) + self.fit_size = int(psf_axis3) + self.x_zero = cast_to_float(pol_zero1) + self.y_zero = cast_to_float(pol_zero2) + self.x_scale = cast_to_float(pol_scal1) + self.y_scale = cast_to_float(pol_scal2) + self.sample_scale = cast_to_float(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) + ] + ) + return jnp.tensordot(P, 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). The + # leading 1 takes x's dtype so the concatenation cannot silently promote + # the result to a wider type than the input position. + x = jnp.asarray(x) + return jnp.concatenate( + [ + jnp.ones((1,), dtype=x.dtype), + jnp.cumprod(jnp.full((self.fit_order,), x)), + ] + ) + + def tree_flatten(self): + """Flatten into traced children and static auxiliary data. + + The PSFEx data (basis, polynomial zero points and scales) and the + ``wcs`` are traced children, so that models read from different files + can be stacked and evaluated in one ``jit``/``vmap`` call. Only + ``fit_order``/``fit_size`` are auxiliary, since they fix the number of + polynomial terms and hence the shapes of the traced arrays. + ``file_name`` is not part of the tree, so that models from different + files share a tree structure. + """ + children = ( + self.wcs, + self.basis, + self.x_zero, + self.y_zero, + self.x_scale, + self.y_scale, + self.sample_scale, + ) + aux_data = {"fit_order": self.fit_order, "fit_size": self.fit_size} + 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. ``file_name`` is not + carried through the tree and is set to ``None``. + """ + obj = object.__new__(cls) + ( + obj.wcs, + obj.basis, + obj.x_zero, + obj.y_zero, + obj.x_scale, + obj.y_scale, + obj.sample_scale, + ) = children + obj.fit_order = aux_data["fit_order"] + obj.fit_size = aux_data["fit_size"] + obj.file_name = None + return obj diff --git a/tests/conftest.py b/tests/conftest.py index 53d2e1c4..7560990b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,6 +8,7 @@ import inspect # noqa: E402 import os # noqa: E402 +import types # noqa: E402 from functools import lru_cache, partial # noqa: E402 from unittest.mock import patch # noqa: E402 @@ -156,6 +157,24 @@ def pytest_pycollect_makemodule(module_path, path, parent): ) and hasattr(module.obj, "setup"): module.obj.setup() + if str(module_path).endswith("tests/GalSim/tests/test_des.py"): + # test_psf reads an optional example catalog inside a + # ``try: ... except OSError:`` block, falling back to hard-coded + # reference values when that (not required) example data is absent. + # jax_galsim does not implement ``Catalog``, so the lookup raises + # AttributeError instead of OSError and aborts the test before it + # reaches the DES_PSFEx checks. Give this module its own namespace in + # which ``Catalog`` triggers the upstream fallback, so the PSFEx model + # is actually exercised. The real jax_galsim module is left untouched. + _test_des_galsim = types.ModuleType("jax_galsim_for_test_des") + _test_des_galsim.__dict__.update(__import__("jax_galsim").__dict__) + + def _catalog_not_implemented(*args, **kwargs): + raise OSError("jax_galsim does not implement galsim.Catalog") + + _test_des_galsim.Catalog = _catalog_not_implemented + module.obj.galsim = _test_des_galsim + # Overwrites galsim in the galsim_test_helpers module for k, v in module.obj.__dict__.items(): if ( diff --git a/tests/galsim_tests_config.yaml b/tests/galsim_tests_config.yaml index 428572a4..32333c27 100644 --- a/tests/galsim_tests_config.yaml +++ b/tests/galsim_tests_config.yaml @@ -114,6 +114,11 @@ allowed_failures: - "module 'jax_galsim' has no attribute 'RandomWalk'" - "module 'jax_galsim' has no attribute 'hsm'" - "module 'jax_galsim' has no attribute 'des'" + # jax_galsim.des implements DES_PSFEx only; the MEDS and shapelet parts of + # the GalSim des module are not ported. + - "module 'jax_galsim.des' has no attribute 'MultiExposureObject'" + - "module 'jax_galsim.des' has no attribute 'WriteMEDS'" + - "module 'jax_galsim.des' has no attribute 'DES_Shapelet'" - "'Image' object has no attribute 'applyNonlinearity'" - "'Image' object has no attribute 'addReciprocityFailure'" - "'Image' object has no attribute 'quantize'" diff --git a/tests/jax/test_des_psfex_jax.py b/tests/jax/test_des_psfex_jax.py new file mode 100644 index 00000000..5f3b6360 --- /dev/null +++ b/tests/jax/test_des_psfex_jax.py @@ -0,0 +1,186 @@ +import os +import shutil +import tempfile + +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, + ) + + # file_name is not carried through the tree, so a rebuilt instance loses it + assert rebuilt.file_name is None + + # Pass the object itself as an argument to a jitted function. Two separate + # instances share a tree structure, so the second call reuses the trace. + 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)) + + +@requires_des_data +@timer +def test_des_psfex_vmap_over_multiple_psf_models(): + """A batch of *different* PSFEx models can be evaluated in one jit/vmap. + + The PSFEx data is traced, so several models -- including ones read from + different files -- stack into a single PyTree and are evaluated by one + compiled kernel. This is the case where a simulation draws galaxies whose + PSFs come from different exposures. Note that ``fit_order``/``fit_size`` + are static, so a batch must share a polynomial degree. + """ + with tempfile.TemporaryDirectory() as tmpdir: + # A second file on disk, so the two models genuinely differ by file + # name (which must not be part of the tree for this to work). + other_name = "other_psfcat.psf" + shutil.copyfile( + os.path.join(DES_DATA_DIR, PSFEX_FILE), os.path.join(tmpdir, other_name) + ) + + psf_a = galsim.des.DES_PSFEx(PSFEX_FILE, dir=DES_DATA_DIR) + psf_b = galsim.des.DES_PSFEx(other_name, dir=tmpdir) + # Only one PSFEx file ships with the test data, so perturb the second + # model's basis to stand in for a different exposure's solution. + psf_b.basis = psf_b.basis * 1.05 + + # Stacking would raise if any per-model data were static auxiliary data. + batch = jax.tree_util.tree_map(lambda *xs: jnp.stack(xs), psf_a, psf_b) + + xs = jnp.array([456.0, 1024.0]) + ys = jnp.array([789.0, 2048.0]) + out = jax.jit( + jax.vmap(lambda model, x, y: model.getPSFArray(galsim.PositionD(x, y))) + )(batch, xs, ys) + + assert out.shape[0] == 2 + + # Each batch element matches that model evaluated on its own. + for i, (model, x, y) in enumerate( + ((psf_a, 456.0, 789.0), (psf_b, 1024.0, 2048.0)) + ): + np.testing.assert_allclose( + np.asarray(out[i]), + np.asarray(model.getPSFArray(galsim.PositionD(x, y))), + rtol=0, + atol=1e-6, + ) + + # The two models must give different answers, i.e. the basis really is + # batched over rather than baked in as a constant. + assert not np.allclose(np.asarray(out[0]), np.asarray(out[1])) + + +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() + test_des_psfex_vmap_over_multiple_psf_models() + print("all DES_PSFEx tests passed")