FEAT - Add chromatic psf and the corresponding tests (Issue #251)

docs/api/chromatic.rst

@@ -0,0 +1,83 @@
+Wavelength-dependent Profiles
+=============================
+
+.. currentmodule:: jax_galsim
+
+JAX-GalSim supports a JAX-native subset of GalSim chromatic rendering.  The
+core use case is a wavelength-dependent PSF convolved with a source whose SED
+is a traced JAX array:
+
+.. code-block:: python
+
+   import jax
+   import jax.numpy as jnp
+   import jax_galsim as galsim
+
+   wave = jnp.linspace(400.0, 900.0, 256)
+   sed = galsim.SED(wave, jnp.ones_like(wave))
+   bandpass = galsim.Bandpass.tophat(550.0, 750.0)
+
+   gal = galsim.Gaussian(half_light_radius=0.5) * sed
+   psf = galsim.ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2)
+   final = galsim.ChromaticConvolution([gal, psf])
+   image = final.drawImage(bandpass, nx=64, ny=64, scale=0.2, n_waves=32)
+
+Separable chromatic objects, such as ``GSObject * SED``, are rendered by
+integrating only the scalar spectral weight, then drawing one unit-flux spatial
+profile.  Non-separable chromatic objects, such as a chromatic PSF whose size
+changes with wavelength, are rendered by integrating their Fourier-space values
+over a fixed wavelength grid.
+
+The wavelength grid size is static, while SED flux values are traced.  This
+keeps the rendering path compatible with ``jax.jit`` and ``jax.grad``.
+
+Spectral objects
+----------------
+
+.. autoclass:: SED
+   :members:
+   :show-inheritance:
+
+.. autoclass:: Bandpass
+   :members:
+   :show-inheritance:
+
+Chromatic objects
+-----------------
+
+.. autoclass:: ChromaticObject
+   :members:
+   :show-inheritance:
+
+.. autoclass:: SimpleChromaticTransformation
+   :members:
+   :show-inheritance:
+
+.. autoclass:: ChromaticAtmosphere
+   :members:
+   :show-inheritance:
+
+.. autoclass:: ChromaticConvolution
+   :members:
+   :show-inheritance:
+
+.. autoclass:: ChromaticSum
+   :members:
+   :show-inheritance:
+
+Compatibility notes
+-------------------
+
+``galsim.Convolve`` dispatches to ``ChromaticConvolution`` when any input is
+chromatic.  Multiplication follows GalSim's common pattern:
+
+.. code-block:: python
+
+   chromatic_gal = galsim.Gaussian(half_light_radius=0.5) * sed
+   same_object = sed * galsim.Gaussian(half_light_radius=0.5)
+
+The current implementation focuses on differentiable array-backed SEDs,
+array-backed bandpasses, separable chromatic sources, and non-separable
+Gaussian/Moffat atmospheric PSFs.  Full GalSim spectral I/O, lookup-table
+metadata, Airy/OpticalPSF chromatic optics, and photon shooting are still
+outside this JAX-native subset.

docs/api/index.rst

@@ -10,6 +10,7 @@ API Reference
    weak-lensing
    wcs
    noise
+   chromatic
    photon_shooting
    interpolation
    fits

jax_galsim/__init__.py

@@ -105,3 +105,14 @@
 
 # this one is specific to jax_galsim
 from . import core
+
+# Chromatic profiles
+from .sed import SED
+from .bandpass import Bandpass
+from .chromatic import (
+    ChromaticObject,
+    SimpleChromaticTransformation,
+    ChromaticAtmosphere,
+    ChromaticConvolution,
+    ChromaticSum,
+)

jax_galsim/bandpass.py

@@ -0,0 +1,206 @@
+import os
+
+import galsim as _galsim
+import jax.numpy as jnp
+from jax.tree_util import register_pytree_node_class
+
+from jax_galsim.core.interpolate import akima_interp, akima_interp_coeffs
+from jax_galsim.core.utils import cast_to_float, ensure_hashable, implements
+
+
+@implements(
+    _galsim.Bandpass,
+    lax_description="""\
+JAX-GalSim supports array-backed bandpasses only.  The wavelength grid is
+static and the throughput array may be traced.  File input, string
+expressions, zeropoints, units, and LookupTable metadata are not implemented.
+""",
+)
+@register_pytree_node_class
+class Bandpass:
+    def __init__(self, wave, throughput, blue_limit=None, red_limit=None):
+        if isinstance(wave, (str, bytes, os.PathLike)) or isinstance(
+            throughput, (str, bytes, os.PathLike)
+        ):
+            raise NotImplementedError(
+                "JAX-GalSim Bandpass supports array-backed bandpasses only; "
+                "file input, string expressions, and unit strings are not implemented."
+            )
+
+        self._wave = jnp.asarray(wave, dtype=float)
+        self._throughput = jnp.asarray(throughput)
+
+        if self._wave.ndim != 1 or len(self._wave) < 2:
+            raise ValueError("wave must be a 1-D array with at least 2 elements.")
+        if self._throughput.shape != self._wave.shape:
+            raise ValueError("throughput must have the same shape as wave.")
+
+        self._blue_limit = (
+            cast_to_float(blue_limit)
+            if blue_limit is not None
+            else cast_to_float(self._wave[0])
+        )
+        self._red_limit = (
+            cast_to_float(red_limit)
+            if red_limit is not None
+            else cast_to_float(self._wave[-1])
+        )
+
+    @property
+    @implements(
+        None,
+        lax_description="JAX-GalSim-specific static wavelength grid for array-backed bandpasses.",
+    )
+    def wave(self):
+        return self._wave
+
+    @property
+    @implements(
+        None,
+        lax_description="JAX-GalSim-specific traced throughput array for array-backed bandpasses.",
+    )
+    def throughput(self):
+        return self._throughput
+
+    @property
+    @implements(
+        None,
+        lax_description="JAX-GalSim-specific blue wavelength limit.",
+    )
+    def blue_limit(self):
+        return self._blue_limit
+
+    @property
+    @implements(
+        None,
+        lax_description="JAX-GalSim-specific red wavelength limit.",
+    )
+    def red_limit(self):
+        return self._red_limit
+
+    @property
+    @implements(_galsim.Bandpass.effective_wavelength)
+    def effective_wavelength(self):
+        return self.calculateEffectiveWavelength()
+
+    @implements(
+        _galsim.Bandpass.calculateEffectiveWavelength,
+        lax_description=(
+            "The JAX implementation returns a traced scalar when throughput is traced."
+        ),
+    )
+    def calculateEffectiveWavelength(self, precise=False):
+        if precise:
+            raise NotImplementedError(
+                "JAX-GalSim Bandpass does not support precise=True."
+            )
+        waves = jnp.linspace(self._blue_limit, self._red_limit, 512)
+        throughput = self(waves)
+        norm = jnp.trapezoid(throughput, waves)
+        midpoint = 0.5 * (self._blue_limit + self._red_limit)
+        safe_norm = jnp.where(norm > 0, norm, 1.0)
+        return jnp.where(
+            norm > 0,
+            jnp.trapezoid(waves * throughput, waves) / safe_norm,
+            midpoint,
+        )
+
+    @implements(_galsim.Bandpass.__call__)
+    def __call__(self, wave):
+        wave = jnp.asarray(wave, dtype=float)
+        if len(self._wave) >= 5:
+            coeffs = akima_interp_coeffs(self._wave, self._throughput, use_jax=True)
+            throughput = akima_interp(wave, self._wave, self._throughput, coeffs)
+        else:
+            throughput = jnp.interp(
+                wave, self._wave, self._throughput, left=0.0, right=0.0
+            )
+        in_band = (wave >= self._blue_limit) & (wave <= self._red_limit)
+        return jnp.where(in_band, throughput, 0.0)
+
+    @implements(_galsim.Bandpass.__mul__)
+    def __mul__(self, other):
+        if isinstance(other, Bandpass):
+            wave = jnp.unique(jnp.concatenate([self._wave, other._wave]))
+            t = jnp.interp(wave, self._wave, self._throughput, left=0.0, right=0.0)
+            t2 = jnp.interp(wave, other._wave, other._throughput, left=0.0, right=0.0)
+            blue = max(self._blue_limit, other._blue_limit)
+            red = min(self._red_limit, other._red_limit)
+            return Bandpass(wave, t * t2, blue_limit=blue, red_limit=red)
+        return Bandpass(
+            self._wave,
+            self._throughput * other,
+            self._blue_limit,
+            self._red_limit,
+        )
+
+    @implements(getattr(_galsim.Bandpass, "__rmul__", None))
+    def __rmul__(self, other):
+        return self.__mul__(other)
+
+    @implements(_galsim.Bandpass.truncate)
+    def truncate(self, blue_limit=None, red_limit=None):
+        blue = blue_limit if blue_limit is not None else self._blue_limit
+        red = red_limit if red_limit is not None else self._red_limit
+        return Bandpass(
+            self._wave,
+            self._throughput,
+            blue_limit=blue,
+            red_limit=red,
+        )
+
+    @classmethod
+    @implements(
+        None,
+        lax_description=(
+            "Construct a unit-throughput tophat bandpass on a static JAX wavelength grid."
+        ),
+    )
+    def tophat(cls, blue_limit, red_limit, n_wave=100):
+        wave = jnp.linspace(blue_limit, red_limit, n_wave)
+        return cls(wave, jnp.ones(n_wave))
+
+    def tree_flatten(self):
+        children = (self._throughput,)
+        aux_data = {
+            "wave": tuple(self._wave.tolist()),
+            "blue_limit": ensure_hashable(self._blue_limit),
+            "red_limit": ensure_hashable(self._red_limit),
+        }
+        return (children, aux_data)
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        return cls(
+            wave=jnp.asarray(aux_data["wave"], dtype=float),
+            throughput=children[0],
+            blue_limit=aux_data["blue_limit"],
+            red_limit=aux_data["red_limit"],
+        )
+
+    def __repr__(self):
+        return (
+            f"Bandpass(wave=[{self._blue_limit:.1f}, {self._red_limit:.1f}] nm, "
+            f"lam_eff={float(self.effective_wavelength):.1f} nm)"
+        )
+
+    def __eq__(self, other):
+        if not isinstance(other, Bandpass):
+            return False
+        return (
+            jnp.array_equal(self._wave, other._wave)
+            & jnp.array_equal(self._throughput, other._throughput)
+            & jnp.array_equal(self._blue_limit, other._blue_limit)
+            & jnp.array_equal(self._red_limit, other._red_limit)
+        )
+
+    def __hash__(self):
+        return hash(
+            (
+                "galsim.Bandpass",
+                ensure_hashable(self._wave),
+                ensure_hashable(self._throughput),
+                ensure_hashable(self._blue_limit),
+                ensure_hashable(self._red_limit),
+            )
+        )