From b8904749f56572de6bc84581af0fab7cb1ef778d Mon Sep 17 00:00:00 2001 From: MaxRonce Date: Thu, 21 May 2026 17:12:09 +0200 Subject: [PATCH 01/12] FEAT - Add chromatic psf and the corresponding tests --- docs/api/chromatic.rst | 83 +++ docs/api/index.rst | 1 + jax_galsim/__init__.py | 11 + jax_galsim/bandpass.py | 192 +++++++ jax_galsim/chromatic.py | 988 +++++++++++++++++++++++++++++++++ jax_galsim/convolve.py | 7 +- jax_galsim/sed.py | 240 ++++++++ tests/galsim_tests_config.yaml | 16 +- tests/jax/test_chromatic.py | 403 ++++++++++++++ 9 files changed, 1936 insertions(+), 5 deletions(-) create mode 100644 docs/api/chromatic.rst create mode 100644 jax_galsim/bandpass.py create mode 100644 jax_galsim/chromatic.py create mode 100644 jax_galsim/sed.py create mode 100644 tests/jax/test_chromatic.py diff --git a/docs/api/chromatic.rst b/docs/api/chromatic.rst new file mode 100644 index 00000000..402224e7 --- /dev/null +++ b/docs/api/chromatic.rst @@ -0,0 +1,83 @@ +Chromatic 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:: Chromatic + :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. diff --git a/docs/api/index.rst b/docs/api/index.rst index 95c5ff72..f61ff2d3 100644 --- a/docs/api/index.rst +++ b/docs/api/index.rst @@ -10,6 +10,7 @@ API Reference weak-lensing wcs noise + chromatic photon_shooting interpolation fits diff --git a/jax_galsim/__init__.py b/jax_galsim/__init__.py index a90a5b6c..56a64d82 100644 --- a/jax_galsim/__init__.py +++ b/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, + Chromatic, + ChromaticAtmosphere, + ChromaticConvolution, + ChromaticSum, +) diff --git a/jax_galsim/bandpass.py b/jax_galsim/bandpass.py new file mode 100644 index 00000000..e068741b --- /dev/null +++ b/jax_galsim/bandpass.py @@ -0,0 +1,192 @@ +"""Bandpass filter representation for chromatic simulations.""" + +import jax.numpy as jnp +from jax.tree_util import register_pytree_node_class + + +@register_pytree_node_class +class Bandpass: + """Wavelength-dependent throughput of an observing bandpass. + + Represents the combined throughput of optics, filter, and detector + as a function of wavelength. The throughput array is a JAX-traced + parameter, so gradients can flow through it if needed (e.g. for + filter-design optimisation). + + Parameters + ---------- + wave : array_like + Wavelength array **in nanometers**, strictly increasing. + Treated as static (not traced by JAX). + throughput : array_like + Dimensionless throughput ∈ [0, 1] at each wavelength. + Treated as a JAX-traced parameter. + blue_limit : float, optional + Override the short-wavelength cut-off in nm. Defaults to + ``wave[0]``. + red_limit : float, optional + Override the long-wavelength cut-off in nm. Defaults to + ``wave[-1]``. + + Examples + -------- + Construct from arrays:: + + >>> import jax.numpy as jnp + >>> from jax_galsim.bandpass import Bandpass + >>> wave = jnp.linspace(550, 700, 100) + >>> bp = Bandpass(wave, jnp.ones(100)) + >>> float(bp(625.0)) + 1.0 + >>> float(bp.effective_wavelength) + 625.0 + + Top-hat filter between 600 nm and 700 nm:: + + >>> wave = jnp.array([500., 600., 600.001, 700., 700.001, 800.]) + >>> thru = jnp.array([0., 0., 1., 1., 0., 0.]) + >>> bp = Bandpass(wave, thru) + """ + + def __init__(self, wave, throughput, blue_limit=None, red_limit=None): + 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 len(self._throughput) != len(self._wave): + raise ValueError("throughput must have the same length as wave.") + + self._blue_limit = float(blue_limit) if blue_limit is not None else float(self._wave[0]) + self._red_limit = float(red_limit) if red_limit is not None else float(self._wave[-1]) + + # Precompute effective wavelength at construction time so it is a + # concrete Python float and can be used as a static value under JIT. + _w = jnp.linspace(self._blue_limit, self._red_limit, 512) + _t = jnp.interp(_w, self._wave, self._throughput) + _norm = jnp.trapezoid(_t, _w) + self._effective_wavelength_val = ( + float(jnp.trapezoid(_w * _t, _w) / _norm) + if float(_norm) > 0 + else float(0.5 * (self._blue_limit + self._red_limit)) + ) + + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + + @property + def wave(self): + """Wavelength grid in nm (JAX array, static).""" + return self._wave + + @property + def throughput(self): + """Throughput array (JAX array, traced).""" + return self._throughput + + @property + def blue_limit(self): + """Short-wavelength cut-off in nm.""" + return self._blue_limit + + @property + def red_limit(self): + """Long-wavelength cut-off in nm.""" + return self._red_limit + + @property + def effective_wavelength(self): + """Flux-weighted mean wavelength (concrete Python float, safe under JIT).""" + return self._effective_wavelength_val + + # ------------------------------------------------------------------ + # Evaluation + # ------------------------------------------------------------------ + + def __call__(self, wave): + """Evaluate throughput at wavelength(s) in nm. + + Returns 0 outside the defined range. + + Parameters + ---------- + wave : float or array_like + Wavelength(s) in nm. + + Returns + ------- + jnp.ndarray + Throughput at the requested wavelengths. + """ + wave = jnp.asarray(wave, dtype=float) + return jnp.interp(wave, self._wave, self._throughput, left=0.0, right=0.0) + + # ------------------------------------------------------------------ + # Arithmetic + # ------------------------------------------------------------------ + + def __mul__(self, other): + """Multiply throughput by a scalar or another Bandpass.""" + if isinstance(other, Bandpass): + # Union wavelength grid + 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) + + def __rmul__(self, other): + return self.__mul__(other) + + def truncate(self, blue_limit=None, red_limit=None): + """Return a new Bandpass with tighter wavelength limits.""" + 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) + + # ------------------------------------------------------------------ + # Convenience constructors + # ------------------------------------------------------------------ + + @classmethod + def tophat(cls, blue_limit, red_limit, n_wave=100): + """Uniform throughput = 1 between blue_limit and red_limit.""" + wave = jnp.linspace(blue_limit, red_limit, n_wave) + return cls(wave, jnp.ones(n_wave)) + + # ------------------------------------------------------------------ + # JAX pytree interface + # ------------------------------------------------------------------ + + def tree_flatten(self): + children = (self._throughput,) + aux_data = { + "wave": tuple(self._wave.tolist()), + "blue_limit": self._blue_limit, + "red_limit": 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"], + ) + + # ------------------------------------------------------------------ + # Misc + # ------------------------------------------------------------------ + + def __repr__(self): + return ( + f"Bandpass(wave=[{self._blue_limit:.1f}, {self._red_limit:.1f}] nm, " + f"lam_eff={self.effective_wavelength:.1f} nm)" + ) diff --git a/jax_galsim/chromatic.py b/jax_galsim/chromatic.py new file mode 100644 index 00000000..9dae511b --- /dev/null +++ b/jax_galsim/chromatic.py @@ -0,0 +1,988 @@ +"""Chromatic (wavelength-dependent) profiles for jax_galsim. + +Architecture overview +--------------------- +Every chromatic object exposes: + +* ``evaluateAtWavelength(wave) -> GSObject`` + Returns the monochromatic profile at wavelength *wave* (nm). The + returned GSObject may carry wavelength-dependent parameters as JAX + traced values, so the function is vmappable. + +* ``drawImage(bandpass, n_waves=64, **kwargs) -> Image`` + Integrates the profile over *bandpass* using a static wavelength grid + of *n_waves* points (trapezoid rule). ``n_waves`` is static at + JIT-compile time; everything else may be traced. + +Hierarchy +--------- +:: + + ChromaticObject base class, non-separable draw by default + ├── Chromatic GSObject × SED (separable, fast path) + ├── ChromaticAtmosphere Gaussian PSF with FWHM ∝ λ^alpha + └── ChromaticConvolution convolution of any chromatic objects + +Separable vs non-separable +-------------------------- +*Separable* means ``I(x, y, λ) = g(x, y) × h(λ)``. For a separable +object the integration reduces to a single monochromatic draw: + + flux = ∫ SED(λ) × BP(λ) dλ + image = g(x, y) drawn with total flux + +*Non-separable* objects (e.g. ChromaticAtmosphere) must evaluate the +k-space image at every wavelength sample and integrate: + + K_eff(k) = ∫ K(k, λ) × BP(λ) dλ + image = IFFT[ K_eff ] + +JAX compatibility +----------------- +* All arithmetic uses ``jax.numpy``. +* Wavelength grids are **static** (fixed-size) to allow ``jax.jit``. +* ``jax.vmap`` vectorises the per-wavelength k-value computation. +* ``jax.grad`` flows through SED flux arrays (e.g. DSPS outputs). +""" + +import jax +import jax.numpy as jnp +from jax.tree_util import register_pytree_node_class + +from jax_galsim.gsparams import GSParams +from jax_galsim.position import PositionD + + +def _pixel_scale_from_kwargs(kwargs): + """Return pixel scale in world units (arcsec/pixel) as a concrete float. + + Reads ``scale`` from kwargs; falls back to reading the WCS object if + the ``wcs`` kwarg is provided, or 1.0 otherwise. The value from + ``kwargs['scale']`` is always a Python float (user-supplied literal), + so this is safe to call inside ``jax.jit``. + """ + if "scale" in kwargs: + return float(kwargs["scale"]) + if "wcs" in kwargs: + wcs = kwargs["wcs"] + if hasattr(wcs, "_scale"): + return float(wcs._scale) + return 1.0 + + +def _make_setup_image(profile, kwargs): + """Build draw target without running full ``drawImage`` setup when size is fixed.""" + from jax_galsim.image import Image + + image = kwargs.get("image", None) + if image is not None: + return Image(image=image) + + image_kwargs = { + "dtype": kwargs.get("dtype", None), + "scale": kwargs.get("scale", None), + "wcs": kwargs.get("wcs", None), + } + image_kwargs = {k: v for k, v in image_kwargs.items() if v is not None} + + bounds = kwargs.get("bounds", None) + if bounds is not None: + return Image(bounds=bounds, **image_kwargs) + + nx = kwargs.get("nx", None) + ny = kwargs.get("ny", None) + if nx is not None and ny is not None: + return Image(nx, ny, **image_kwargs) + + return profile.drawImage(setup_only=True, **kwargs) + + +def _next_pow2(n): + n = int(n) + out = 1 + while out < n: + out *= 2 + return out + + +def _fix_fft_size_for_image(profile, image): + """Pin FFT size for JIT-safe setup when output bounds are already fixed.""" + nrow, ncol = image.array.shape + n = max(128, _next_pow2(2 * max(nrow, ncol))) + n = min(n, profile.gsparams.maximum_fft_size) + return profile.withGSParams(minimum_fft_size=n, maximum_fft_size=n) + + +def _static_kcoords(kimage, wrap_size, pixel_scale): + """Return k-space pixel centers as a JAX array with static shape.""" + nrow, ncol = kimage.array.shape + x = jnp.arange(ncol, dtype=float) + y = jnp.arange(nrow, dtype=float) - wrap_size // 2 + kx, ky = jnp.meshgrid(x, y) + dk = 2.0 * jnp.pi / (wrap_size * pixel_scale) + return jnp.stack([kx.ravel(), ky.ravel()], axis=-1) * dk + + +# --------------------------------------------------------------------------- +# Base class +# --------------------------------------------------------------------------- + + +class ChromaticObject: + """Base class for wavelength-dependent profiles. + + Subclasses must override :meth:`evaluateAtWavelength`. The default + :meth:`drawImage` uses a k-space trapezoid integration that works for + any subclass; separable subclasses override it with a faster path. + """ + + #: Set True in separable subclasses. + _separable: bool = False + + def __init__(self, obj=None): + if obj is None: + self._base_obj = None + return + + from jax_galsim.gsobject import GSObject + + if not isinstance(obj, GSObject): + raise TypeError("ChromaticObject requires a GSObject.") + self._base_obj = obj + self._separable = True + + @property + def separable(self): + """True if the profile factors as g(x,y) × h(λ).""" + return self._separable + + # ------------------------------------------------------------------ + # Interface that subclasses must implement + # ------------------------------------------------------------------ + + def evaluateAtWavelength(self, wave): + """Return the monochromatic GSObject at wavelength *wave* (nm). + + This method must be JAX-traceable: all internal computations + should use ``jax.numpy``, and the returned GSObject's parameters + may be abstract JAX tracers. + + Parameters + ---------- + wave : float or jax scalar + Wavelength in nm. + + Returns + ------- + GSObject + """ + if getattr(self, "_base_obj", None) is not None: + return self._base_obj + raise NotImplementedError( + f"{self.__class__.__name__} must implement evaluateAtWavelength." + ) + + # ------------------------------------------------------------------ + # Drawing + # ------------------------------------------------------------------ + + def drawImage(self, bandpass, n_waves=64, **kwargs): + """Draw the bandpass-integrated image. + + Parameters + ---------- + bandpass : Bandpass + Observing bandpass. + n_waves : int + Number of wavelength samples for numerical integration. + Must be a **static** integer (fixed at JIT compile time). + **kwargs + Forwarded to the underlying ``GSObject.drawImage`` calls. + Typical keys: ``nx``, ``ny``, ``scale``, ``method``, etc. + + Returns + ------- + Image + """ + if self._separable: + return self._drawImage_separable(bandpass, n_waves, **kwargs) + return self._drawImage_nonseparable(bandpass, n_waves, **kwargs) + + # ------------------------------------------------------------------ + # Separable fast path + # ------------------------------------------------------------------ + + def _drawImage_separable(self, bandpass, n_waves, **kwargs): + """FFT draw, scaling by total SED×BP flux (traced-safe under JIT). + + Design: split into two phases. + + **Phase 1 - static setup**: + Build image bounds, k-grid, and base k-values using the unit-flux + spatial profile. All shape parameters must be concrete Python + scalars at this stage (true for ``Chromatic`` where the spatial + profile has static params); the SED flux is NOT evaluated here. + + **Phase 2 — traced computation**: + Compute total_flux = ∫ SED(λ)×BP(λ) dλ (may be a JAX traced + value, e.g. DSPS output), multiply into k-values, IFFT. + """ + from jax_galsim.box import Pixel + from jax_galsim.convolve import Convolve + from jax_galsim.image import Image + + wave_eff = bandpass.effective_wavelength # concrete Python float + pixel_scale = _pixel_scale_from_kwargs(kwargs) + + # ------------------------------------------------------------------ + # Phase 1: concrete setup. Under jit this runs during tracing. + # ------------------------------------------------------------------ + with jax.disable_jit(): + # Unit-flux spatial profile — shape params are concrete Python floats + spatial_prof = self._static_spatial_profile(wave_eff) + image = _make_setup_image(spatial_prof, kwargs) + spatial_prof = _fix_fft_size_for_image(spatial_prof, image) + original_center = image.center + original_wcs = image.wcs + image.setCenter(0, 0) + + pixel = Pixel(scale=pixel_scale, gsparams=spatial_prof.gsparams) + prof_conv = Convolve([spatial_prof, pixel], gsparams=spatial_prof.gsparams) + kimage, wrap_size = prof_conv.drawFFT_makeKImage(image) + + kcoords = _static_kcoords(kimage, wrap_size, pixel_scale) + + # Static k-values (unit flux, no SED scaling). + kvals_static = jax.vmap( + lambda k: prof_conv._kValue(PositionD(k[0], k[1])) + )(kcoords) + + # Apply the same -0.5 pixel centering correction that gsobject._adjust_offset + # uses for even-sized images (avoids 0.5-pixel offset vs non-chromatic draws). + img_shape = image.array.shape # (ny, nx); unchanged by setCenter + dx_corr = -0.5 * pixel_scale * ((img_shape[1] + 1) % 2) + dy_corr = -0.5 * pixel_scale * ((img_shape[0] + 1) % 2) + phase_corr = jnp.exp( + -1j * (kcoords[:, 0] * dx_corr + kcoords[:, 1] * dy_corr) + ) + kvals_static = kvals_static * phase_corr + + kshape = kimage.array.shape + kbounds = kimage.bounds + kwcs = kimage.wcs + + # ------------------------------------------------------------------ + # Phase 2: traced computation — integrate SED × bandpass + # ------------------------------------------------------------------ + waves = jnp.linspace(bandpass.blue_limit, bandpass.red_limit, n_waves) + weights = jax.vmap(lambda w: self._sed_value(w) * bandpass(w))(waves) + total_flux = jnp.trapezoid(weights, waves) + + # Scale pre-computed k-values by traced total flux + kvals = kvals_static * total_flux + karray = kvals.reshape(kshape).astype(kimage.dtype) + eff_kimage = Image( + array=karray, bounds=kbounds, wcs=kwcs, _check_bounds=False + ) + prof_conv.drawFFT_finish(image, eff_kimage, wrap_size, add_to_image=False) + + image.shift(original_center) + image.wcs = original_wcs + return image + + def _static_spatial_profile(self, wave_eff): + """Return unit-flux spatial profile at *wave_eff* with static params. + + *wave_eff* must be a concrete Python float (use + ``bandpass.effective_wavelength``). Subclasses override this when + they have a dedicated static spatial object (e.g. ``Chromatic`` has + ``self.obj``). The default calls ``evaluateAtWavelength`` with a + Python float — works when all shape params are Python scalars. + """ + return self.evaluateAtWavelength(float(wave_eff)).withFlux(1.0) + + def _sed_value(self, wave): + """SED flux density at *wave*. Subclasses override if needed.""" + if getattr(self, "_base_obj", None) is not None: + return jnp.ones((), dtype=float) + raise NotImplementedError + + # ------------------------------------------------------------------ + # Non-separable k-space integration + # ------------------------------------------------------------------ + + def _drawImage_nonseparable(self, bandpass, n_waves, **kwargs): + """Integrate in k-space: K_eff = ∫ K(k,λ) × BP(λ) dλ, then IFFT. + + Mirrors the drawFFT pipeline of GSObject.drawImage exactly, split + into a concrete setup phase and a traced computation phase. + """ + from jax_galsim.box import Pixel + from jax_galsim.convolve import Convolve + from jax_galsim.image import Image + + wave_eff = bandpass.effective_wavelength # static Python float + pixel_scale = _pixel_scale_from_kwargs(kwargs) + + # ------------------------------------------------------------------ + # Phase 1: concrete setup. Under jit this runs during tracing. + # ------------------------------------------------------------------ + with jax.disable_jit(): + prof0 = self._static_spatial_profile(wave_eff) + image = _make_setup_image(prof0, kwargs) + prof0 = _fix_fft_size_for_image(prof0, image) + original_center = image.center + original_wcs = image.wcs + image.setCenter(0, 0) + + pixel = Pixel(scale=pixel_scale, gsparams=prof0.gsparams) + prof0_conv = Convolve([prof0, pixel], gsparams=prof0.gsparams) + kimage, wrap_size = prof0_conv.drawFFT_makeKImage(image) + + kcoords = _static_kcoords(kimage, wrap_size, pixel_scale) + + pixel_kvals = jax.vmap( + lambda k: pixel._kValue(PositionD(k[0], k[1])) + )(kcoords) + + # Apply the -0.5 pixel centering correction for even-sized images. + img_shape = image.array.shape + dx_corr = -0.5 * pixel_scale * ((img_shape[1] + 1) % 2) + dy_corr = -0.5 * pixel_scale * ((img_shape[0] + 1) % 2) + phase_corr = jnp.exp( + -1j * (kcoords[:, 0] * dx_corr + kcoords[:, 1] * dy_corr) + ) + pixel_kvals = pixel_kvals * phase_corr + + kshape = kimage.array.shape + kbounds = kimage.bounds + kwcs = kimage.wcs + + # ------------------------------------------------------------------ + # Phase 2: traced computation + # ------------------------------------------------------------------ + waves = jnp.linspace(bandpass.blue_limit, bandpass.red_limit, n_waves) + + def kvals_at_wave(wave): + prof = self.evaluateAtWavelength(wave) + kv = jax.vmap( + lambda k: prof._kValue(PositionD(k[0], k[1])) + )(kcoords) + return kv * bandpass(wave) + + all_kvals = jax.vmap(kvals_at_wave)(waves) # (n_waves, n_k) + eff_kvals = jnp.trapezoid(all_kvals, waves, axis=0) + eff_kvals = eff_kvals * pixel_kvals + + eff_karray = eff_kvals.reshape(kshape).astype(kimage.dtype) + eff_kimage = Image( + array=eff_karray, + bounds=kbounds, + wcs=kwcs, + _check_bounds=False, + ) + + # IFFT back to pixel space + prof0_conv.drawFFT_finish(image, eff_kimage, wrap_size, add_to_image=False) + + # Restore original center and WCS (same as drawImage does after drawFFT) + image.shift(original_center) + image.wcs = original_wcs + + return image + + # ------------------------------------------------------------------ + # Operator overloads + # ------------------------------------------------------------------ + + def __add__(self, other): + return ChromaticSum([self, other]) + + def __radd__(self, other): + return ChromaticSum([other, self]) + + def __mul__(self, other): + from jax_galsim.sed import SED + + if isinstance(other, SED): + base_obj = getattr(self, "_base_obj", None) + if base_obj is None: + raise TypeError("Only achromatic ChromaticObject wrappers can be multiplied by SED.") + return Chromatic(base_obj, other) + return _ScaledChromaticObject(self, other) + + def __rmul__(self, other): + return self.__mul__(other) + + +class _ScaledChromaticObject(ChromaticObject): + """Chromatic object with wavelength-independent flux scaling.""" + + def __init__(self, obj, scale): + self.obj = obj + self.scale = scale + self._separable = obj._separable + + def evaluateAtWavelength(self, wave): + return self.obj.evaluateAtWavelength(wave).withScaledFlux(self.scale) + + def _static_spatial_profile(self, wave_eff): + return self.obj._static_spatial_profile(wave_eff) + + def _sed_value(self, wave): + return self.obj._sed_value(wave) * self.scale + + +# --------------------------------------------------------------------------- +# ChromaticSum — sum of chromatic objects +# --------------------------------------------------------------------------- + + +class ChromaticSum(ChromaticObject): + """Sum of two or more chromatic profiles. + + The combined SED is the sum of all component SEDs. Drawing + evaluates each component at every wavelength and sums. + + Parameters + ---------- + obj_list : list of ChromaticObject + """ + + _separable = False # conservative; optimised later if all are separable + + def __init__(self, *args): + if len(args) == 0: + raise TypeError("At least one object must be provided.") + if len(args) == 1 and isinstance(args[0], (list, tuple)): + self.obj_list = list(args[0]) + else: + self.obj_list = list(args) + self._separable = all(o._separable for o in self.obj_list) + + def evaluateAtWavelength(self, wave): + from jax_galsim.sum import Sum + + return Sum([o.evaluateAtWavelength(wave) for o in self.obj_list]) + + def drawImage(self, bandpass, n_waves=64, **kwargs): + # Draw each component and sum + images = [obj.drawImage(bandpass, n_waves=n_waves, **kwargs) + for obj in self.obj_list] + result = images[0] + for img in images[1:]: + result._array = result._array + img._array + return result + + +# --------------------------------------------------------------------------- +# Chromatic — separable GSObject × SED +# --------------------------------------------------------------------------- + + +@register_pytree_node_class +class Chromatic(ChromaticObject): + """Separable chromatic profile: a GSObject multiplied by an SED. + + The spatial profile is fixed; wavelength dependence enters only + through the SED flux scaling. + + ``I(x, y, λ) = g(x, y) × SED(λ)`` + + Drawing a ``Chromatic`` through a bandpass reduces to a single + monochromatic draw at the effective wavelength with total flux + ``∫ SED(λ) × BP(λ) dλ``. + + Parameters + ---------- + obj : GSObject + Normalised spatial profile (flux = 1 by convention, though + any flux is allowed and will be multiplied by the SED). + sed : SED + Spectral energy distribution. + + Examples + -------- + :: + + >>> from jax_galsim import Gaussian + >>> from jax_galsim.sed import SED + >>> from jax_galsim.bandpass import Bandpass + >>> import jax.numpy as jnp + + >>> wave = jnp.linspace(500, 900, 256) + >>> sed = SED(wave, jnp.ones(256)) + >>> bp = Bandpass.tophat(550, 750) + >>> gal = Gaussian(half_light_radius=0.5) * sed + >>> img = gal.drawImage(bp, scale=0.2, nx=64, ny=64) + """ + + _separable = True + + def __init__(self, obj, sed): + self.obj = obj + self.sed = sed + + def evaluateAtWavelength(self, wave): + """Return the GSObject scaled to SED(wave).""" + return self.obj.withScaledFlux(self.sed(wave)) + + def _sed_value(self, wave): + return self.sed(wave) + + def _static_spatial_profile(self, wave_eff): + """Return self.obj (unit flux) — shape params are always static.""" + return self.obj.withFlux(1.0) + + # JAX pytree: obj and sed are both pytrees + def tree_flatten(self): + children = (self.obj, self.sed) + aux_data = {} + return (children, aux_data) + + @classmethod + def tree_unflatten(cls, aux_data, children): + return cls(obj=children[0], sed=children[1]) + + def __repr__(self): + return f"Chromatic({self.obj!r}, {self.sed!r})" + + +# --------------------------------------------------------------------------- +# ChromaticAtmosphere — seeing PSF with wavelength-dependent FWHM +# --------------------------------------------------------------------------- + + +@register_pytree_node_class +class ChromaticAtmosphere(ChromaticObject): + """Atmospheric PSF with a power-law wavelength-dependent FWHM. + + The PSF profile at wavelength λ is a Gaussian (or Moffat) with: + + FWHM(λ) = fwhm_ref × (λ / lam_ref)^alpha + + For Kolmogorov turbulence, the expected scaling is α ≈ −0.2. + + This profile carries a **flat (dimensionless) SED**: ``SED(λ) = 1``. + The physical SED is typically attached to the galaxy component via + :class:`Chromatic`, and passed to :class:`ChromaticConvolution`. + + Parameters + ---------- + fwhm_ref : float + FWHM in arcseconds at the reference wavelength. + lam_ref : float + Reference wavelength in nm. + alpha : float, optional + Power-law index. Default −0.2 (Kolmogorov). + profile : {'gaussian', 'moffat'} + Profile type. Default ``'gaussian'``. + moffat_beta : float, optional + Moffat β parameter (only used when ``profile='moffat'``). + Default 4.765 (typical for atmospheric seeing). + gsparams : GSParams, optional + + Examples + -------- + :: + + >>> from jax_galsim.chromatic import ChromaticAtmosphere + >>> psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) + >>> prof = psf.evaluateAtWavelength(550.0) # Gaussian at 550 nm + """ + + _separable = False + + def __init__( + self, + fwhm_ref, + lam_ref, + alpha=-0.2, + profile="gaussian", + moffat_beta=4.765, + gsparams=None, + ): + # All shape params stored as Python floats (static, not JAX-traced). + # This enables jax.jit compatibility without pinning the FFT size. + # To differentiate through fwhm_ref, use gsparams with fixed fft_size: + # GSParams(minimum_fft_size=N, maximum_fft_size=N) + self._fwhm_ref = float(fwhm_ref) + self._lam_ref = float(lam_ref) + self._alpha = float(alpha) + self._profile = profile + self._moffat_beta = float(moffat_beta) + self._gsparams = GSParams.check(gsparams) + + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + + @property + def fwhm_ref(self): + return self._fwhm_ref + + @property + def lam_ref(self): + return self._lam_ref + + @property + def alpha(self): + return self._alpha + + # ------------------------------------------------------------------ + # Core method + # ------------------------------------------------------------------ + + def evaluateAtWavelength(self, wave): + """Return the PSF profile (unit flux) at wavelength *wave* (nm). + + FWHM is scaled as ``fwhm_ref × (wave / lam_ref)^alpha``. + When *wave* is a JAX tracer (inside vmap/jit), fwhm is also traced. + """ + fwhm = self._fwhm_ref * (wave / self._lam_ref) ** self._alpha + + if self._profile == "gaussian": + from jax_galsim.gaussian import Gaussian + + return Gaussian(fwhm=fwhm, flux=1.0, gsparams=self._gsparams) + + elif self._profile == "moffat": + from jax_galsim.moffat import Moffat + + return Moffat( + fwhm=fwhm, + beta=self._moffat_beta, + flux=1.0, + gsparams=self._gsparams, + ) + else: + raise ValueError( + f"Unknown profile type '{self._profile}'. " + "Expected 'gaussian' or 'moffat'." + ) + + def _sed_value(self, wave): + """Flat SED: unit flux at all wavelengths.""" + return jnp.ones((), dtype=float) + + # ------------------------------------------------------------------ + # JAX pytree interface + # ------------------------------------------------------------------ + + def tree_flatten(self): + # No traced children — all params are static Python scalars. + children = () + aux_data = { + "fwhm_ref": self._fwhm_ref, + "lam_ref": self._lam_ref, + "alpha": self._alpha, + "profile": self._profile, + "moffat_beta": self._moffat_beta, + "gsparams": self._gsparams, + } + return (children, aux_data) + + @classmethod + def tree_unflatten(cls, aux_data, children): + return cls( + fwhm_ref=aux_data["fwhm_ref"], + lam_ref=aux_data["lam_ref"], + alpha=aux_data["alpha"], + profile=aux_data["profile"], + moffat_beta=aux_data["moffat_beta"], + gsparams=aux_data["gsparams"], + ) + + def __repr__(self): + return ( + f"ChromaticAtmosphere(fwhm_ref={float(self._fwhm_ref):.3f}, " + f"lam_ref={self._lam_ref:.0f} nm, alpha={self._alpha:.2f}, " + f"profile={self._profile!r})" + ) + + +# --------------------------------------------------------------------------- +# ChromaticConvolution — convolution of chromatic objects +# --------------------------------------------------------------------------- + + +class ChromaticConvolution(ChromaticObject): + """Convolution of multiple chromatic profiles. + + Computes the bandpass-integrated image of the convolution: + + K_eff(k) = ∫ ∏_i K_i(k, λ) × BP(λ) dλ + + where each ``K_i(k, λ)`` is the k-space value of the i-th component + at wavelength λ. + + **Optimised case**: when all components except one are separable + (e.g. galaxy × SED convolved with a chromatic PSF), the separable + profiles are extracted and convolved with the wavelength-integrated + effective PSF. This avoids multiplying the same galaxy k-image at + every wavelength sample. + + Parameters + ---------- + obj_list : list of ChromaticObject or GSObject + Components to convolve. Plain ``GSObject`` instances are wrapped + automatically in a flat-SED :class:`Chromatic`. + + Examples + -------- + :: + + >>> from jax_galsim import Gaussian, Convolve + >>> from jax_galsim.chromatic import Chromatic, ChromaticAtmosphere, ChromaticConvolution + >>> from jax_galsim.sed import SED + >>> from jax_galsim.bandpass import Bandpass + >>> import jax.numpy as jnp + + >>> wave = jnp.linspace(300, 1100, 512) + >>> flux = jnp.exp(-0.5 * ((wave - 700) / 150) ** 2) # Gaussian SED + >>> sed = SED(wave, flux) + >>> bp = Bandpass.tophat(550, 800) + + >>> gal = Gaussian(half_light_radius=0.5) * sed + >>> psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) + >>> final = ChromaticConvolution([gal, psf]) + >>> img = final.drawImage(bp, scale=0.2, nx=64, ny=64) + """ + + _separable = False + + def __init__(self, *args, **kwargs): + from jax_galsim.gsobject import GSObject + + if len(args) == 0: + raise TypeError("At least one object must be provided.") + if len(args) == 1 and isinstance(args[0], (list, tuple)): + obj_list = list(args[0]) + else: + obj_list = list(args) + + real_space = kwargs.pop("real_space", None) + self._gsparams = GSParams.check(kwargs.pop("gsparams", None)) + self._propagate_gsparams = kwargs.pop("propagate_gsparams", True) + if kwargs: + raise TypeError( + "ChromaticConvolution constructor got unexpected keyword argument(s): %s" + % kwargs.keys() + ) + if real_space: + raise NotImplementedError("Real-space chromatic convolutions are not implemented") + + # Wrap plain GSObjects with a flat SED so they fit the interface + processed = [] + for obj in obj_list: + if isinstance(obj, ChromaticConvolution): + processed.extend(obj.obj_list) + continue + if isinstance(obj, GSObject): + from jax_galsim.sed import SED + + wave_stub = jnp.array([100.0, 2000.0]) + flat_sed = SED(wave_stub, jnp.ones(2)) + processed.append(Chromatic(obj, flat_sed)) + else: + processed.append(obj) + self.obj_list = processed + + @property + def gsparams(self): + return self._gsparams + + def withGSParams(self, gsparams=None, **kwargs): + ret = self.__class__.__new__(self.__class__) + ret.obj_list = self.obj_list + ret._gsparams = GSParams.check(gsparams, self._gsparams, **kwargs) + ret._propagate_gsparams = self._propagate_gsparams + return ret + + def evaluateAtWavelength(self, wave): + """Return the convolved monochromatic profile at *wave* (nm).""" + from jax_galsim.convolve import Convolve + + return Convolve([o.evaluateAtWavelength(wave) for o in self.obj_list]) + + # ------------------------------------------------------------------ + # Drawing with the optimised separable split + # ------------------------------------------------------------------ + + def drawImage(self, bandpass, n_waves=64, **kwargs): + """Draw the integrated image, exploiting separability where possible. + + Algorithm: + + 1. Separate components into *separable* (galaxy × SED) and + *non-separable* (chromatic PSF). + 2. Build the combined weight function: + ``w(λ) = ∏_sep SED_i(λ) × BP(λ)`` + 3. In k-space, integrate non-separable components: + ``K_nonsep_eff(k) = ∫ ∏_nonsep K_i(k,λ) × w(λ) dλ`` + 4. Multiply by separable component k-values (λ-independent after + extracting their SED into the weight). + 5. Include pixel convolution, IFFT to real space. + + All-separable case falls back to a single monochromatic draw. + """ + from jax_galsim.box import Pixel + from jax_galsim.convolve import Convolve + from jax_galsim.image import Image + + sep_objs = [o for o in self.obj_list if o._separable] + nonsep_objs = [o for o in self.obj_list if not o._separable] + + wave_eff = bandpass.effective_wavelength # static Python float + waves = jnp.linspace(bandpass.blue_limit, bandpass.red_limit, n_waves) + pixel_scale = _pixel_scale_from_kwargs(kwargs) + + # ------------------------------------------------------------------- + # Phase 1: concrete setup. Under jit this runs during tracing. + # + # Shape parameters (sigma, fwhm) must NOT be JIT-traced inputs here. + # If they are, pin the FFT size via gsparams(min/max_fft_size=N). + # ------------------------------------------------------------------- + with jax.disable_jit(): + if not nonsep_objs: + # All-separable: use static spatial profiles (avoids traced SED) + spatial_profs = [o._static_spatial_profile(wave_eff) for o in sep_objs] + grid_prof = Convolve(spatial_profs) + else: + # Mixed: sep objects use static profiles; nonsep objects are evaluated + # at wave_eff (their shape params must be concrete at this point). + fiducial_profs = [ + o._static_spatial_profile(wave_eff) if o._separable + else o.evaluateAtWavelength(wave_eff) + for o in self.obj_list + ] + grid_prof = Convolve(fiducial_profs) + + image = _make_setup_image(grid_prof, kwargs) + grid_prof = _fix_fft_size_for_image(grid_prof, image) + original_center = image.center + original_wcs = image.wcs + image.setCenter(0, 0) + + pixel = Pixel(scale=pixel_scale) + grid_prof_conv = Convolve([grid_prof, pixel], gsparams=grid_prof.gsparams) + kimage, wrap_size = grid_prof_conv.drawFFT_makeKImage(image) + + # k-space coordinates with static shape. + kcoords = _static_kcoords(kimage, wrap_size, pixel_scale) + n_k = kcoords.shape[0] + + # Pixel k-values. + pixel_kvals = jax.vmap( + lambda k: pixel._kValue(PositionD(k[0], k[1])) + )(kcoords) + + # Match gsobject._adjust_offset: even-sized images need a -0.5 pixel + # true-center correction before the FFT draw. Apply it as a k-space + # phase so chromatic draws align with monochromatic drawImage. + img_shape = image.array.shape # (ny, nx); unchanged by setCenter + dx_corr = -0.5 * pixel_scale * ((img_shape[1] + 1) % 2) + dy_corr = -0.5 * pixel_scale * ((img_shape[0] + 1) % 2) + phase_corr = jnp.exp( + -1j * (kcoords[:, 0] * dx_corr + kcoords[:, 1] * dy_corr) + ) + + if not nonsep_objs: + # Pre-compute k-values of the full (spatial+pixel) convolution + base_kvals = jax.vmap( + lambda k: grid_prof_conv._kValue(PositionD(k[0], k[1])) + )(kcoords) * phase_corr + else: + pixel_kvals = pixel_kvals * phase_corr + + # Pre-compute k-values of separable components (unit flux each) + sep_kvals = jnp.ones(n_k, dtype=complex) + for o in sep_objs: + prof_sep = o._static_spatial_profile(wave_eff) + sep_kvals = sep_kvals * jax.vmap( + lambda k: prof_sep._kValue(PositionD(k[0], k[1])) + )(kcoords) + + kshape = kimage.array.shape + kbounds = kimage.bounds + kwcs = kimage.wcs + + # ------------------------------------------------------------------- + # Phase 2: traced computation (JAX-traced values allowed here) + # ------------------------------------------------------------------- + if not nonsep_objs: + # All-separable: integrate SED × bandpass → total flux, then scale + def combined_weight(wave): + w = bandpass(wave) + for o in sep_objs: + w = w * o._sed_value(wave) + return w + + total_flux = jnp.trapezoid(jax.vmap(combined_weight)(waves), waves) + kvals = base_kvals * total_flux + + else: + # Mixed sep + nonsep: integrate nonsep k-values weighted by sep SED × BP + def sep_weight(wave): + w = bandpass(wave) + for o in sep_objs: + w = w * o._sed_value(wave) + return w + + def kvals_nonsep_at_wave(wave): + kv = jnp.ones(n_k, dtype=complex) + for o in nonsep_objs: + prof = o.evaluateAtWavelength(wave) + kv = kv * jax.vmap( + lambda k: prof._kValue(PositionD(k[0], k[1])) + )(kcoords) + return kv * sep_weight(wave) + + all_kvals = jax.vmap(kvals_nonsep_at_wave)(waves) # (n_waves, n_k) + eff_kvals = jnp.trapezoid(all_kvals, waves, axis=0) # (n_k,) + + # Multiply by separable spatial k-values and pixel convolution + kvals = eff_kvals * sep_kvals * pixel_kvals + + karray = kvals.reshape(kshape).astype(kimage.dtype) + eff_kimage = Image( + array=karray, bounds=kbounds, wcs=kwcs, _check_bounds=False + ) + grid_prof_conv.drawFFT_finish(image, eff_kimage, wrap_size, add_to_image=False) + + image.shift(original_center) + image.wcs = original_wcs + return image + + def __repr__(self): + inner = ", ".join(repr(o) for o in self.obj_list) + return f"ChromaticConvolution([{inner}])" + + +# --------------------------------------------------------------------------- +# Monkey-patch GSObject.__mul__ to return Chromatic when multiplied by SED +# --------------------------------------------------------------------------- + +def _gsobject_mul_sed(self, other): + """Allow ``gsobject * sed → Chromatic(gsobject, sed)``.""" + from jax_galsim.sed import SED + + if isinstance(other, SED): + return Chromatic(self, other) + # Fall through to original implementation (flux scaling) + return self.withScaledFlux(other) + + +def _gsobject_rmul_sed(self, other): + return _gsobject_mul_sed(self, other) + + +# Apply the patch once at import time +def _patch_gsobject(): + from jax_galsim.gsobject import GSObject + + GSObject.__mul__ = _gsobject_mul_sed + GSObject.__rmul__ = _gsobject_rmul_sed + + +_patch_gsobject() diff --git a/jax_galsim/convolve.py b/jax_galsim/convolve.py index e6df0769..246777c8 100644 --- a/jax_galsim/convolve.py +++ b/jax_galsim/convolve.py @@ -11,7 +11,7 @@ @implements( _galsim.Convolve, - lax_description="""Does not support ChromaticConvolutions""", + lax_description="""Supports ChromaticConvolutions for FFT drawing only.""", ) def Convolve(*args, **kwargs): if len(args) == 0: @@ -28,6 +28,11 @@ def Convolve(*args, **kwargs): ) # else args is already the list of objects + from jax_galsim.chromatic import ChromaticObject, ChromaticConvolution + + if any(isinstance(obj, ChromaticObject) for obj in args): + return ChromaticConvolution(args, **kwargs) + return Convolution(*args, **kwargs) diff --git a/jax_galsim/sed.py b/jax_galsim/sed.py new file mode 100644 index 00000000..496de7fc --- /dev/null +++ b/jax_galsim/sed.py @@ -0,0 +1,240 @@ +"""Spectral Energy Distribution (SED) for chromatic profiles. + +Designed for JAX compatibility: the flux array is a traced parameter, +so gradients flow through SED values (e.g. from DSPS outputs). +""" + +import jax.numpy as jnp +from jax.tree_util import register_pytree_node_class + + +@register_pytree_node_class +class SED: + """Spectral Energy Distribution. + + Represents flux density as a function of wavelength, designed for + full JAX compatibility. The ``flux`` array is a JAX-traced parameter, + enabling gradients through SED parameters (e.g. outputs of DSPS). + + The wavelength grid (``wave``) is treated as static auxiliary data: + it defines the interpolation structure and is not differentiated. + + Parameters + ---------- + wave : array_like + Wavelength array **in nanometers**. Must be strictly increasing. + Treated as static (not traced by JAX). + flux : array_like + Flux density at each wavelength. Treated as a JAX-traced parameter. + Units are arbitrary but must be consistent across the simulation + (typically photons / nm / cm² / s for spectral SEDs, or + dimensionless for shape-only profiles). + redshift : float, optional + Cosmological redshift applied to the SED. The observed wavelength + grid is shifted to ``wave_obs = wave_rest * (1 + redshift)``. + Default 0. + + Examples + -------- + Basic construction from arrays:: + + >>> import jax.numpy as jnp + >>> from jax_galsim.sed import SED + >>> wave = jnp.linspace(300, 1100, 512) # nm + >>> flux = jnp.ones(512) + >>> sed = SED(wave, flux) + >>> float(sed(550.0)) + 1.0 + + DSPS workflow — flux is a traced JAX array:: + + >>> flux = dsps_model(params) # JAX array + >>> sed = SED(dsps_wave_nm, flux) + >>> image = chromatic_galaxy.drawImage(bandpass) # differentiable + + Redshifted SED:: + + >>> sed_z = SED(wave, flux, redshift=0.5) + >>> sed_z(825.0) # queries rest-frame 550 nm + """ + + def __init__(self, wave, flux, redshift=0.0): + self._wave = jnp.asarray(wave, dtype=float) # static, not traced + self._flux = jnp.asarray(flux) # traced + self._redshift = jnp.asarray(redshift, dtype=float) + + 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 len(self._flux) != len(self._wave): + raise ValueError("flux must have the same length as wave.") + + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + + @property + def wave(self): + """Rest-frame wavelength grid in nm (JAX array, static).""" + return self._wave + + @property + def flux(self): + """Flux density array (JAX array, traced).""" + return self._flux + + @property + def redshift(self): + """Redshift (JAX scalar).""" + return self._redshift + + @property + def blue_limit(self): + """Shortest observed wavelength in nm.""" + return float(self._wave[0]) * (1.0 + float(self._redshift)) + + @property + def red_limit(self): + """Longest observed wavelength in nm.""" + return float(self._wave[-1]) * (1.0 + float(self._redshift)) + + # ------------------------------------------------------------------ + # Evaluation + # ------------------------------------------------------------------ + + def __call__(self, wave): + """Evaluate flux density at observed wavelength(s) in nm. + + Uses linear interpolation; returns 0 outside the defined range. + + Parameters + ---------- + wave : float or array_like + Observed wavelength(s) in nm. + + Returns + ------- + jnp.ndarray + Flux density at the requested wavelengths. + """ + wave = jnp.asarray(wave, dtype=float) + # Convert observed wavelength to rest-frame before interpolating + wave_rest = wave / (1.0 + self._redshift) + return jnp.interp(wave_rest, self._wave, self._flux, left=0.0, right=0.0) + + # ------------------------------------------------------------------ + # Flux through a bandpass + # ------------------------------------------------------------------ + + def calculateFlux(self, bandpass, n_waves=512): + """Integrate SED through a bandpass: ``∫ SED(λ) × BP(λ) dλ``. + + Parameters + ---------- + bandpass : Bandpass + Observing bandpass. + n_waves : int, optional + Number of quadrature points. Default 512. + + Returns + ------- + jnp.ndarray + Scalar flux value. + """ + waves = jnp.linspace(bandpass.blue_limit, bandpass.red_limit, n_waves) + return jnp.trapezoid(self(waves) * bandpass(waves), waves) + + # ------------------------------------------------------------------ + # Arithmetic + # ------------------------------------------------------------------ + + def withRedshift(self, redshift): + """Return a copy of this SED at a new redshift.""" + return SED(self._wave, self._flux, redshift) + + def __mul__(self, other): + """Multiply SED by a scalar or another SED. + + SED × scalar scales all flux values. + SED × SED multiplies flux densities (both evaluated on ``self``'s grid). + """ + if isinstance(other, SED): + # Evaluate other SED on self's rest-frame grid and multiply + other_flux = jnp.interp( + self._wave * (1.0 + self._redshift), + other._wave * (1.0 + other._redshift), + other._flux, + left=0.0, + right=0.0, + ) + return SED(self._wave, self._flux * other_flux, self._redshift) + from jax_galsim.gsobject import GSObject + + if isinstance(other, GSObject): + from jax_galsim.chromatic import Chromatic + + return Chromatic(other, self) + from jax_galsim.chromatic import ChromaticObject + + if isinstance(other, ChromaticObject): + return other * self + return SED(self._wave, self._flux * other, self._redshift) + + def __rmul__(self, other): + return self.__mul__(other) + + def __truediv__(self, other): + return SED(self._wave, self._flux / other, self._redshift) + + def __add__(self, other): + if isinstance(other, SED): + other_flux = jnp.interp( + self._wave * (1.0 + self._redshift), + other._wave * (1.0 + other._redshift), + other._flux, + left=0.0, + right=0.0, + ) + return SED(self._wave, self._flux + other_flux, self._redshift) + return SED(self._wave, self._flux + other, self._redshift) + + # ------------------------------------------------------------------ + # JAX pytree interface + # ------------------------------------------------------------------ + + def tree_flatten(self): + """Flatten for JAX tracing. + + ``flux`` and ``redshift`` are traced children. + ``wave`` is static auxiliary data (the interpolation grid). + """ + children = (self._flux, self._redshift) + # wave must be hashable for JAX cache keys; store as tuple of floats. + aux_data = {"wave": tuple(self._wave.tolist())} + return (children, aux_data) + + @classmethod + def tree_unflatten(cls, aux_data, children): + return cls( + wave=jnp.asarray(aux_data["wave"], dtype=float), + flux=children[0], + redshift=children[1], + ) + + # ------------------------------------------------------------------ + # Misc + # ------------------------------------------------------------------ + + def __repr__(self): + return ( + f"SED(wave=[{self._wave[0]:.1f}, ..., {self._wave[-1]:.1f}] nm, " + f"n_wave={len(self._wave)}, redshift={float(self._redshift):.4f})" + ) + + def __eq__(self, other): + if not isinstance(other, SED): + return False + return ( + jnp.array_equal(self._wave, other._wave) + and jnp.array_equal(self._flux, other._flux) + and jnp.array_equal(self._redshift, other._redshift) + ) diff --git a/tests/galsim_tests_config.yaml b/tests/galsim_tests_config.yaml index 428572a4..1391e880 100644 --- a/tests/galsim_tests_config.yaml +++ b/tests/galsim_tests_config.yaml @@ -21,6 +21,9 @@ enabled_tests: - test_noise.py - test_image.py - test_photon_array.py + - test_sed.py + - test_bandpass.py + - test_chromatic.py - "*" # means all tests from galsim coord: - test_angle.py @@ -86,7 +89,6 @@ allowed_failures: - "module 'jax_galsim' has no attribute 'Dict'" - "module 'jax_galsim' has no attribute 'OutputCatalog'" - "module 'jax_galsim' has no attribute 'cdmodel'" - - "module 'jax_galsim' has no attribute 'ChromaticObject'" - "module 'jax_galsim' has no attribute 'ChromaticAiry'" - "module 'jax_galsim' has no attribute 'config'" - "module 'jax_galsim' has no attribute 'RealGalaxyCatalog'" @@ -101,13 +103,20 @@ allowed_failures: - "pad_image not implemented in jax_galsim." - "InterpolatedImages do not support noise padding in jax_galsim." - "module 'jax_galsim' has no attribute 'FittedSIPWCS'" - - "module 'jax_galsim' has no attribute 'Bandpass'" - "module 'jax_galsim' has no attribute 'Refraction'" - "module 'jax_galsim' has no attribute 'FRatioAngles'" - "module 'jax_galsim' has no attribute 'PupilAnnulusSampler'" - "module 'jax_galsim' has no attribute 'TimeSampler'" - "object has no attribute 'noise'" - - "module 'jax_galsim' has no attribute 'SED'" + # JAX-native chromatic subset: array-backed SED/Bandpass only. + # Full GalSim spectral I/O, expression parsing, LookupTable metadata, + # transforms, thinning, photon shooting, and units are not implemented yet. + - "SED.__init__() got an unexpected keyword argument 'spec'" + - "SED.__init__() got an unexpected keyword argument 'wave_type'" + - "Bandpass.__init__() got an unexpected keyword argument 'wave_type'" + - "could not convert string to float" + - "is not a valid JAX array type" + - "'ChromaticObject' object has no attribute 'dilate'" - "module 'jax_galsim' has no attribute 'getCOSMOSNoise'" - "GSParams.__init__() got an unexpected keyword argument 'allowed_flux_variation'" - "module 'jax_galsim' has no attribute 'Atmosphere'" @@ -135,7 +144,6 @@ allowed_failures: - "module 'jax_galsim' has no attribute 'Aperture'" - "module 'jax_galsim' has no attribute 'AtmosphericScreen'" - "module 'jax_galsim' has no attribute 'OpticalScreen'" - - "module 'jax_galsim' has no attribute 'ChromaticConvolution'" - "module 'jax_galsim' has no attribute 'phase_screens'" - "module 'jax_galsim' has no attribute 'DistDeviate'" - "module 'jax_galsim' has no attribute 'roman'" diff --git a/tests/jax/test_chromatic.py b/tests/jax/test_chromatic.py new file mode 100644 index 00000000..353961d3 --- /dev/null +++ b/tests/jax/test_chromatic.py @@ -0,0 +1,403 @@ +"""Tests for jax_galsim chromatic PSF support. + +Verifies: +- SED and Bandpass construction and evaluation +- Chromatic (separable) drawImage correctness +- ChromaticAtmosphere (non-separable) drawImage correctness +- ChromaticConvolution (galaxy × SED ⊗ PSF) correctness +- jax.jit compatibility for all paths +- jax.grad compatibility for SED-flux differentiation +- Pytree round-trip (tree_flatten / tree_unflatten) +- Numerical agreement with analytic expectations +""" + +# ruff: noqa: E402,I001 + +import jax +import jax.numpy as jnp +import pytest + +# Enable float64 for accuracy +jax.config.update("jax_enable_x64", True) + +import jax_galsim as jgal +from jax_galsim.chromatic import ChromaticAtmosphere, ChromaticConvolution + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +WAVE = jnp.linspace(400.0, 900.0, 256) # nm +BP = jgal.Bandpass.tophat(550.0, 750.0) # 200 nm wide, throughput = 1 +BP_NARROW = jgal.Bandpass.tophat(600.0, 700.0) # 100 nm wide + + +def flat_sed(scale=1.0): + return jgal.SED(WAVE, jnp.ones(256) * scale) + + +# --------------------------------------------------------------------------- +# SED tests +# --------------------------------------------------------------------------- + + +def test_sed_evaluation(): + sed = flat_sed() + assert float(sed(600.0)) == pytest.approx(1.0, rel=1e-5) + assert float(sed(300.0)) == pytest.approx(0.0) # outside range + + +def test_sed_redshift(): + sed = jgal.SED(WAVE, jnp.ones(256), redshift=1.0) + # observed 800 nm → rest-frame 400 nm → flux = 1.0 + assert float(sed(800.0)) == pytest.approx(1.0, rel=1e-4) + # observed 400 nm → rest-frame 200 nm → outside range → 0 + assert float(sed(400.0)) == pytest.approx(0.0, abs=1e-6) + + +def test_sed_calculate_flux(): + sed = flat_sed() + flux = float(sed.calculateFlux(BP)) + # ∫_550^750 1 dλ = 200 nm + assert flux == pytest.approx(200.0, rel=1e-3) + + +def test_sed_pytree_roundtrip(): + sed = flat_sed(2.0) + leaves, treedef = jax.tree_util.tree_flatten(sed) + sed2 = jax.tree_util.tree_unflatten(treedef, leaves) + assert float(sed2(600.0)) == pytest.approx(2.0, rel=1e-5) + + +def test_sed_arithmetic(): + sed1 = flat_sed(2.0) + sed2 = sed1 * 3.0 + assert float(sed2(600.0)) == pytest.approx(6.0, rel=1e-5) + + sed3 = sed1 + flat_sed(1.0) + assert float(sed3(600.0)) == pytest.approx(3.0, rel=1e-5) + + +# --------------------------------------------------------------------------- +# Bandpass tests +# --------------------------------------------------------------------------- + + +def test_bandpass_evaluation(): + bp = jgal.Bandpass.tophat(550.0, 750.0) + assert float(bp(625.0)) == pytest.approx(1.0) + assert float(bp(500.0)) == pytest.approx(0.0) + assert float(bp(800.0)) == pytest.approx(0.0) + + +def test_bandpass_effective_wavelength(): + bp = jgal.Bandpass.tophat(550.0, 750.0) + lam_eff = bp.effective_wavelength + assert isinstance(lam_eff, float) + assert lam_eff == pytest.approx(650.0, rel=1e-4) + + +def test_bandpass_effective_wavelength_concrete(): + """effective_wavelength must be a concrete Python float (safe under JIT).""" + bp = jgal.Bandpass.tophat(550.0, 750.0) + lam_eff = bp.effective_wavelength + # If this were a JAX tracer, float() would raise ConcretizationTypeError + assert isinstance(lam_eff, float) + + +def test_bandpass_mul(): + bp1 = jgal.Bandpass.tophat(500.0, 700.0) + bp2 = jgal.Bandpass.tophat(600.0, 800.0) + bp = bp1 * bp2 + assert float(bp(650.0)) == pytest.approx(1.0) + assert float(bp(550.0)) == pytest.approx(0.0) + assert float(bp(750.0)) == pytest.approx(0.0) + + +def test_bandpass_pytree_roundtrip(): + bp = jgal.Bandpass.tophat(550.0, 750.0) + leaves, treedef = jax.tree_util.tree_flatten(bp) + bp2 = jax.tree_util.tree_unflatten(treedef, leaves) + assert float(bp2(625.0)) == pytest.approx(1.0) + assert bp2.effective_wavelength == pytest.approx(650.0, rel=1e-4) + + +# --------------------------------------------------------------------------- +# Chromatic (separable) tests +# --------------------------------------------------------------------------- + + +def test_chromatic_construction(): + sed = flat_sed() + gal = jgal.Gaussian(half_light_radius=0.5) * sed + assert gal._separable + + +def test_chromatic_drawImage_flux(): + """Image pixel sum should equal ∫ SED(λ) × BP(λ) dλ.""" + sed = flat_sed() + gal = jgal.Gaussian(half_light_radius=0.5) * sed + img = gal.drawImage(BP, scale=0.2, nx=32, ny=32) + # ∫_550^750 1 dλ = 200 + assert float(img.array.sum()) == pytest.approx(200.0, rel=5e-3) + + +def test_chromatic_drawImage_narrow_bandpass(): + """Narrower bandpass → smaller total flux.""" + sed = flat_sed() + gal = jgal.Gaussian(half_light_radius=0.5) * sed + img = gal.drawImage(BP_NARROW, scale=0.2, nx=32, ny=32) + assert float(img.array.sum()) == pytest.approx(100.0, rel=5e-3) + + +def test_chromatic_jit(): + @jax.jit + def render(flux): + sed = jgal.SED(WAVE, flux) + gal = jgal.Gaussian(half_light_radius=0.5) * sed + return gal.drawImage(BP, scale=0.2, nx=32, ny=32).array.sum() + + result = render(jnp.ones(256)) + assert float(result) == pytest.approx(200.0, rel=5e-3) + + +def test_chromatic_grad(): + @jax.jit + def render(flux): + sed = jgal.SED(WAVE, flux) + gal = jgal.Gaussian(half_light_radius=0.5) * sed + return gal.drawImage(BP, scale=0.2, nx=32, ny=32).array.sum() + + grad = jax.grad(render)(jnp.ones(256)) + grad_arr = jnp.asarray(grad) + + # Gradient sums to total bandpass flux ≈ 200 + assert float(grad.sum()) == pytest.approx(200.0, rel=5e-2) + # Outside bandpass → zero gradient + idx_out = int((420 - 400) / (900 - 400) * 255) + assert float(grad[idx_out]) == pytest.approx(0.0, abs=1e-8) + # Inside bandpass region has positive total contribution + mask_in = (WAVE >= 550) & (WAVE <= 750) + assert float(grad_arr[mask_in].sum()) > 0.0 + + +def test_chromatic_jit_recompile(): + """JIT should reuse compiled code when called twice with different flux.""" + @jax.jit + def render(flux): + sed = jgal.SED(WAVE, flux) + gal = jgal.Gaussian(half_light_radius=0.5) * sed + return gal.drawImage(BP, scale=0.2, nx=32, ny=32).array.sum() + + r1 = float(render(jnp.ones(256))) + r2 = float(render(jnp.ones(256) * 2.0)) + assert r2 == pytest.approx(r1 * 2.0, rel=1e-4) + + +# --------------------------------------------------------------------------- +# ChromaticAtmosphere tests +# --------------------------------------------------------------------------- + + +def test_chromatic_atmosphere_evaluate(): + psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) + prof = psf.evaluateAtWavelength(700.0) + # At reference wavelength, FWHM should be exactly fwhm_ref + # jax_galsim Gaussian exposes sigma; FWHM = sigma * fwhm_factor + assert isinstance(prof, jgal.Gaussian) + fwhm = float(prof.sigma) * jgal.Gaussian._fwhm_factor + assert fwhm == pytest.approx(0.7, rel=1e-4) + + +def test_chromatic_atmosphere_scaling(): + psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) + prof_blue = psf.evaluateAtWavelength(350.0) + prof_red = psf.evaluateAtWavelength(700.0) + # FWHM ∝ λ^alpha = λ^(-0.2) → bluer is wider (alpha < 0) + assert float(prof_blue.sigma) > float(prof_red.sigma) + + +def test_chromatic_atmosphere_drawImage_flux(): + """Total flux = ∫ BP(λ) × 1 dλ = 200 (flat SED, unit PSF flux).""" + psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) + img = psf.drawImage(BP, scale=0.2, nx=32, ny=32) + assert float(img.array.sum()) == pytest.approx(200.0, rel=5e-3) + + +def test_chromatic_atmosphere_moffat(): + psf = ChromaticAtmosphere( + fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2, + profile="moffat", moffat_beta=4.765 + ) + prof = psf.evaluateAtWavelength(700.0) + assert isinstance(prof, jgal.Moffat) + img = psf.drawImage(BP, scale=0.2, nx=32, ny=32) + assert float(img.array.sum()) == pytest.approx(200.0, rel=5e-2) + + +def test_chromatic_atmosphere_pytree(): + psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) + leaves, treedef = jax.tree_util.tree_flatten(psf) + psf2 = jax.tree_util.tree_unflatten(treedef, leaves) + assert psf2._fwhm_ref == pytest.approx(0.7) + assert psf2._alpha == pytest.approx(-0.2) + + +# --------------------------------------------------------------------------- +# ChromaticConvolution tests +# --------------------------------------------------------------------------- + + +def test_chromatic_convolution_flux(): + """Galaxy × SED ⊗ ChromaticAtmosphere: flux = ∫ SED(λ) × BP(λ) dλ.""" + sed = flat_sed() + gal = jgal.Gaussian(half_light_radius=0.5) * sed + psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) + final = ChromaticConvolution([gal, psf]) + img = final.drawImage(BP, scale=0.2, nx=64, ny=64, n_waves=32) + assert float(img.array.sum()) == pytest.approx(200.0, rel=5e-2) + + +def test_chromatic_convolution_jit(): + @jax.jit + def render(flux): + sed = jgal.SED(WAVE, flux) + gal = jgal.Gaussian(half_light_radius=0.5) * sed + psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) + return ChromaticConvolution([gal, psf]).drawImage( + BP, scale=0.2, nx=64, ny=64, n_waves=32 + ).array.sum() + + result = render(jnp.ones(256)) + assert float(result) == pytest.approx(200.0, rel=5e-2) + + +def test_chromatic_convolution_grad(): + @jax.jit + def render(flux): + sed = jgal.SED(WAVE, flux) + gal = jgal.Gaussian(half_light_radius=0.5) * sed + psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) + return ChromaticConvolution([gal, psf]).drawImage( + BP, scale=0.2, nx=64, ny=64, n_waves=32 + ).array.sum() + + grad = jax.grad(render)(jnp.ones(256)) + grad_arr = jnp.asarray(grad) + + # Gradient is nonzero only at wavelengths that fall inside the bandpass + # (the 32 quadrature points lie in [550, 750] nm). + # - sum of all gradients ≈ total bandpass flux = 200 + assert float(grad.sum()) == pytest.approx(200.0, rel=5e-2) + # - max gradient is positive + assert float(grad.max()) > 0.0 + # - indices corresponding to wavelengths outside bandpass have zero gradient + # ~420 nm → outside [550, 750] bandpass + idx_out = int((420 - 400) / (900 - 400) * 255) + assert float(grad[idx_out]) == pytest.approx(0.0, abs=1e-8) + # - indices well inside bandpass region have nonzero total contribution + mask_in = (WAVE >= 550) & (WAVE <= 750) + assert float(grad_arr[mask_in].sum()) > 0.0 + + +def test_chromatic_convolution_linearity(): + """Doubling SED flux doubles image sum.""" + @jax.jit + def render(flux): + sed = jgal.SED(WAVE, flux) + gal = jgal.Gaussian(half_light_radius=0.5) * sed + psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) + return ChromaticConvolution([gal, psf]).drawImage( + BP, scale=0.2, nx=64, ny=64, n_waves=32 + ).array.sum() + + s1 = float(render(jnp.ones(256))) + s2 = float(render(jnp.ones(256) * 2.0)) + assert s2 == pytest.approx(s1 * 2.0, rel=1e-4) + + +# --------------------------------------------------------------------------- +# ChromaticSum tests +# --------------------------------------------------------------------------- + + +def test_chromatic_sum_flux(): + sed1 = flat_sed(1.0) + sed2 = flat_sed(2.0) + gal1 = jgal.Gaussian(half_light_radius=0.3) * sed1 + gal2 = jgal.Gaussian(half_light_radius=0.8) * sed2 + combined = gal1 + gal2 + img = combined.drawImage(BP, scale=0.2, nx=32, ny=32) + # Total flux = (1 + 2) × 200 = 600 + assert float(img.array.sum()) == pytest.approx(600.0, rel=1e-2) + + +# --------------------------------------------------------------------------- +# Monkey-patch (GSObject * SED) tests +# --------------------------------------------------------------------------- + + +def test_gsobject_mul_sed(): + sed = flat_sed() + gal = jgal.Gaussian(half_light_radius=0.5) + from jax_galsim.chromatic import Chromatic + result = gal * sed + assert isinstance(result, Chromatic) + + +def test_gsobject_mul_scalar(): + gal = jgal.Gaussian(half_light_radius=0.5, flux=1.0) + scaled = gal * 5.0 + assert float(scaled.flux) == pytest.approx(5.0) + + +if __name__ == "__main__": + # Run all tests inline for quick check + import sys + + tests = [ + test_sed_evaluation, + test_sed_redshift, + test_sed_calculate_flux, + test_sed_pytree_roundtrip, + test_sed_arithmetic, + test_bandpass_evaluation, + test_bandpass_effective_wavelength, + test_bandpass_effective_wavelength_concrete, + test_bandpass_mul, + test_bandpass_pytree_roundtrip, + test_chromatic_construction, + test_chromatic_drawImage_flux, + test_chromatic_drawImage_narrow_bandpass, + test_chromatic_jit, + test_chromatic_grad, + test_chromatic_jit_recompile, + test_chromatic_atmosphere_evaluate, + test_chromatic_atmosphere_scaling, + test_chromatic_atmosphere_drawImage_flux, + test_chromatic_atmosphere_moffat, + test_chromatic_atmosphere_pytree, + test_chromatic_convolution_flux, + test_chromatic_convolution_jit, + test_chromatic_convolution_grad, + test_chromatic_convolution_linearity, + test_chromatic_sum_flux, + test_gsobject_mul_sed, + test_gsobject_mul_scalar, + ] + + failed = [] + for t in tests: + try: + t() + print(f" PASS {t.__name__}") + except Exception as e: + print(f" FAIL {t.__name__}: {e}") + failed.append(t.__name__) + + print() + print(f"{len(tests) - len(failed)}/{len(tests)} passed") + if failed: + print("Failed:", failed) + sys.exit(1) From c0cd2de39a617302170199fef2f37e99305f1299 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 28 May 2026 20:20:03 +0000 Subject: [PATCH 02/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- jax_galsim/bandpass.py | 16 +++++---- jax_galsim/chromatic.py | 70 ++++++++++++++++++++----------------- jax_galsim/convolve.py | 2 +- jax_galsim/sed.py | 2 +- tests/jax/test_chromatic.py | 32 ++++++++++------- 5 files changed, 69 insertions(+), 53 deletions(-) diff --git a/jax_galsim/bandpass.py b/jax_galsim/bandpass.py index e068741b..a6046e28 100644 --- a/jax_galsim/bandpass.py +++ b/jax_galsim/bandpass.py @@ -57,8 +57,12 @@ def __init__(self, wave, throughput, blue_limit=None, red_limit=None): if len(self._throughput) != len(self._wave): raise ValueError("throughput must have the same length as wave.") - self._blue_limit = float(blue_limit) if blue_limit is not None else float(self._wave[0]) - self._red_limit = float(red_limit) if red_limit is not None else float(self._wave[-1]) + self._blue_limit = ( + float(blue_limit) if blue_limit is not None else float(self._wave[0]) + ) + self._red_limit = ( + float(red_limit) if red_limit is not None else float(self._wave[-1]) + ) # Precompute effective wavelength at construction time so it is a # concrete Python float and can be used as a static value under JIT. @@ -136,8 +140,9 @@ def __mul__(self, other): 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) + return Bandpass( + self._wave, self._throughput * other, self._blue_limit, self._red_limit + ) def __rmul__(self, other): return self.__mul__(other) @@ -146,8 +151,7 @@ def truncate(self, blue_limit=None, red_limit=None): """Return a new Bandpass with tighter wavelength limits.""" 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) + return Bandpass(self._wave, self._throughput, blue_limit=blue, red_limit=red) # ------------------------------------------------------------------ # Convenience constructors diff --git a/jax_galsim/chromatic.py b/jax_galsim/chromatic.py index 9dae511b..aaba0e81 100644 --- a/jax_galsim/chromatic.py +++ b/jax_galsim/chromatic.py @@ -253,13 +253,13 @@ def _drawImage_separable(self, bandpass, n_waves, **kwargs): kcoords = _static_kcoords(kimage, wrap_size, pixel_scale) # Static k-values (unit flux, no SED scaling). - kvals_static = jax.vmap( - lambda k: prof_conv._kValue(PositionD(k[0], k[1])) - )(kcoords) + kvals_static = jax.vmap(lambda k: prof_conv._kValue(PositionD(k[0], k[1])))( + kcoords + ) # Apply the same -0.5 pixel centering correction that gsobject._adjust_offset # uses for even-sized images (avoids 0.5-pixel offset vs non-chromatic draws). - img_shape = image.array.shape # (ny, nx); unchanged by setCenter + img_shape = image.array.shape # (ny, nx); unchanged by setCenter dx_corr = -0.5 * pixel_scale * ((img_shape[1] + 1) % 2) dy_corr = -0.5 * pixel_scale * ((img_shape[0] + 1) % 2) phase_corr = jnp.exp( @@ -281,9 +281,7 @@ def _drawImage_separable(self, bandpass, n_waves, **kwargs): # Scale pre-computed k-values by traced total flux kvals = kvals_static * total_flux karray = kvals.reshape(kshape).astype(kimage.dtype) - eff_kimage = Image( - array=karray, bounds=kbounds, wcs=kwcs, _check_bounds=False - ) + eff_kimage = Image(array=karray, bounds=kbounds, wcs=kwcs, _check_bounds=False) prof_conv.drawFFT_finish(image, eff_kimage, wrap_size, add_to_image=False) image.shift(original_center) @@ -341,9 +339,9 @@ def _drawImage_nonseparable(self, bandpass, n_waves, **kwargs): kcoords = _static_kcoords(kimage, wrap_size, pixel_scale) - pixel_kvals = jax.vmap( - lambda k: pixel._kValue(PositionD(k[0], k[1])) - )(kcoords) + pixel_kvals = jax.vmap(lambda k: pixel._kValue(PositionD(k[0], k[1])))( + kcoords + ) # Apply the -0.5 pixel centering correction for even-sized images. img_shape = image.array.shape @@ -365,12 +363,10 @@ def _drawImage_nonseparable(self, bandpass, n_waves, **kwargs): def kvals_at_wave(wave): prof = self.evaluateAtWavelength(wave) - kv = jax.vmap( - lambda k: prof._kValue(PositionD(k[0], k[1])) - )(kcoords) + kv = jax.vmap(lambda k: prof._kValue(PositionD(k[0], k[1])))(kcoords) return kv * bandpass(wave) - all_kvals = jax.vmap(kvals_at_wave)(waves) # (n_waves, n_k) + all_kvals = jax.vmap(kvals_at_wave)(waves) # (n_waves, n_k) eff_kvals = jnp.trapezoid(all_kvals, waves, axis=0) eff_kvals = eff_kvals * pixel_kvals @@ -407,7 +403,9 @@ def __mul__(self, other): if isinstance(other, SED): base_obj = getattr(self, "_base_obj", None) if base_obj is None: - raise TypeError("Only achromatic ChromaticObject wrappers can be multiplied by SED.") + raise TypeError( + "Only achromatic ChromaticObject wrappers can be multiplied by SED." + ) return Chromatic(base_obj, other) return _ScaledChromaticObject(self, other) @@ -467,8 +465,9 @@ def evaluateAtWavelength(self, wave): def drawImage(self, bandpass, n_waves=64, **kwargs): # Draw each component and sum - images = [obj.drawImage(bandpass, n_waves=n_waves, **kwargs) - for obj in self.obj_list] + images = [ + obj.drawImage(bandpass, n_waves=n_waves, **kwargs) for obj in self.obj_list + ] result = images[0] for img in images[1:]: result._array = result._array + img._array @@ -770,7 +769,9 @@ def __init__(self, *args, **kwargs): % kwargs.keys() ) if real_space: - raise NotImplementedError("Real-space chromatic convolutions are not implemented") + raise NotImplementedError( + "Real-space chromatic convolutions are not implemented" + ) # Wrap plain GSObjects with a flat SED so they fit the interface processed = [] @@ -852,7 +853,8 @@ def drawImage(self, bandpass, n_waves=64, **kwargs): # Mixed: sep objects use static profiles; nonsep objects are evaluated # at wave_eff (their shape params must be concrete at this point). fiducial_profs = [ - o._static_spatial_profile(wave_eff) if o._separable + o._static_spatial_profile(wave_eff) + if o._separable else o.evaluateAtWavelength(wave_eff) for o in self.obj_list ] @@ -873,14 +875,14 @@ def drawImage(self, bandpass, n_waves=64, **kwargs): n_k = kcoords.shape[0] # Pixel k-values. - pixel_kvals = jax.vmap( - lambda k: pixel._kValue(PositionD(k[0], k[1])) - )(kcoords) + pixel_kvals = jax.vmap(lambda k: pixel._kValue(PositionD(k[0], k[1])))( + kcoords + ) # Match gsobject._adjust_offset: even-sized images need a -0.5 pixel # true-center correction before the FFT draw. Apply it as a k-space # phase so chromatic draws align with monochromatic drawImage. - img_shape = image.array.shape # (ny, nx); unchanged by setCenter + img_shape = image.array.shape # (ny, nx); unchanged by setCenter dx_corr = -0.5 * pixel_scale * ((img_shape[1] + 1) % 2) dy_corr = -0.5 * pixel_scale * ((img_shape[0] + 1) % 2) phase_corr = jnp.exp( @@ -889,9 +891,12 @@ def drawImage(self, bandpass, n_waves=64, **kwargs): if not nonsep_objs: # Pre-compute k-values of the full (spatial+pixel) convolution - base_kvals = jax.vmap( - lambda k: grid_prof_conv._kValue(PositionD(k[0], k[1])) - )(kcoords) * phase_corr + base_kvals = ( + jax.vmap(lambda k: grid_prof_conv._kValue(PositionD(k[0], k[1])))( + kcoords + ) + * phase_corr + ) else: pixel_kvals = pixel_kvals * phase_corr @@ -933,21 +938,19 @@ def kvals_nonsep_at_wave(wave): kv = jnp.ones(n_k, dtype=complex) for o in nonsep_objs: prof = o.evaluateAtWavelength(wave) - kv = kv * jax.vmap( - lambda k: prof._kValue(PositionD(k[0], k[1])) - )(kcoords) + kv = kv * jax.vmap(lambda k: prof._kValue(PositionD(k[0], k[1])))( + kcoords + ) return kv * sep_weight(wave) - all_kvals = jax.vmap(kvals_nonsep_at_wave)(waves) # (n_waves, n_k) + all_kvals = jax.vmap(kvals_nonsep_at_wave)(waves) # (n_waves, n_k) eff_kvals = jnp.trapezoid(all_kvals, waves, axis=0) # (n_k,) # Multiply by separable spatial k-values and pixel convolution kvals = eff_kvals * sep_kvals * pixel_kvals karray = kvals.reshape(kshape).astype(kimage.dtype) - eff_kimage = Image( - array=karray, bounds=kbounds, wcs=kwcs, _check_bounds=False - ) + eff_kimage = Image(array=karray, bounds=kbounds, wcs=kwcs, _check_bounds=False) grid_prof_conv.drawFFT_finish(image, eff_kimage, wrap_size, add_to_image=False) image.shift(original_center) @@ -963,6 +966,7 @@ def __repr__(self): # Monkey-patch GSObject.__mul__ to return Chromatic when multiplied by SED # --------------------------------------------------------------------------- + def _gsobject_mul_sed(self, other): """Allow ``gsobject * sed → Chromatic(gsobject, sed)``.""" from jax_galsim.sed import SED diff --git a/jax_galsim/convolve.py b/jax_galsim/convolve.py index 1b847d2d..ab0258c7 100644 --- a/jax_galsim/convolve.py +++ b/jax_galsim/convolve.py @@ -28,7 +28,7 @@ def Convolve(*args, **kwargs): ) # else args is already the list of objects - from jax_galsim.chromatic import ChromaticObject, ChromaticConvolution + from jax_galsim.chromatic import ChromaticConvolution, ChromaticObject if any(isinstance(obj, ChromaticObject) for obj in args): return ChromaticConvolution(args, **kwargs) diff --git a/jax_galsim/sed.py b/jax_galsim/sed.py index 496de7fc..bf9b816d 100644 --- a/jax_galsim/sed.py +++ b/jax_galsim/sed.py @@ -60,7 +60,7 @@ class SED: def __init__(self, wave, flux, redshift=0.0): self._wave = jnp.asarray(wave, dtype=float) # static, not traced - self._flux = jnp.asarray(flux) # traced + self._flux = jnp.asarray(flux) # traced self._redshift = jnp.asarray(redshift, dtype=float) if self._wave.ndim != 1 or len(self._wave) < 2: diff --git a/tests/jax/test_chromatic.py b/tests/jax/test_chromatic.py index 353961d3..f8644721 100644 --- a/tests/jax/test_chromatic.py +++ b/tests/jax/test_chromatic.py @@ -45,7 +45,7 @@ def flat_sed(scale=1.0): def test_sed_evaluation(): sed = flat_sed() assert float(sed(600.0)) == pytest.approx(1.0, rel=1e-5) - assert float(sed(300.0)) == pytest.approx(0.0) # outside range + assert float(sed(300.0)) == pytest.approx(0.0) # outside range def test_sed_redshift(): @@ -184,6 +184,7 @@ def render(flux): def test_chromatic_jit_recompile(): """JIT should reuse compiled code when called twice with different flux.""" + @jax.jit def render(flux): sed = jgal.SED(WAVE, flux) @@ -227,8 +228,7 @@ def test_chromatic_atmosphere_drawImage_flux(): def test_chromatic_atmosphere_moffat(): psf = ChromaticAtmosphere( - fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2, - profile="moffat", moffat_beta=4.765 + fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2, profile="moffat", moffat_beta=4.765 ) prof = psf.evaluateAtWavelength(700.0) assert isinstance(prof, jgal.Moffat) @@ -265,9 +265,11 @@ def render(flux): sed = jgal.SED(WAVE, flux) gal = jgal.Gaussian(half_light_radius=0.5) * sed psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) - return ChromaticConvolution([gal, psf]).drawImage( - BP, scale=0.2, nx=64, ny=64, n_waves=32 - ).array.sum() + return ( + ChromaticConvolution([gal, psf]) + .drawImage(BP, scale=0.2, nx=64, ny=64, n_waves=32) + .array.sum() + ) result = render(jnp.ones(256)) assert float(result) == pytest.approx(200.0, rel=5e-2) @@ -279,9 +281,11 @@ def render(flux): sed = jgal.SED(WAVE, flux) gal = jgal.Gaussian(half_light_radius=0.5) * sed psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) - return ChromaticConvolution([gal, psf]).drawImage( - BP, scale=0.2, nx=64, ny=64, n_waves=32 - ).array.sum() + return ( + ChromaticConvolution([gal, psf]) + .drawImage(BP, scale=0.2, nx=64, ny=64, n_waves=32) + .array.sum() + ) grad = jax.grad(render)(jnp.ones(256)) grad_arr = jnp.asarray(grad) @@ -303,14 +307,17 @@ def render(flux): def test_chromatic_convolution_linearity(): """Doubling SED flux doubles image sum.""" + @jax.jit def render(flux): sed = jgal.SED(WAVE, flux) gal = jgal.Gaussian(half_light_radius=0.5) * sed psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) - return ChromaticConvolution([gal, psf]).drawImage( - BP, scale=0.2, nx=64, ny=64, n_waves=32 - ).array.sum() + return ( + ChromaticConvolution([gal, psf]) + .drawImage(BP, scale=0.2, nx=64, ny=64, n_waves=32) + .array.sum() + ) s1 = float(render(jnp.ones(256))) s2 = float(render(jnp.ones(256) * 2.0)) @@ -342,6 +349,7 @@ def test_gsobject_mul_sed(): sed = flat_sed() gal = jgal.Gaussian(half_light_radius=0.5) from jax_galsim.chromatic import Chromatic + result = gal * sed assert isinstance(result, Chromatic) From 539826b9f35451d0600b0e9778826919b32e5f40 Mon Sep 17 00:00:00 2001 From: "Matthew R. Becker" Date: Thu, 28 May 2026 15:20:33 -0500 Subject: [PATCH 03/12] Add test_chromatic_jax.py file --- tests/jax/{test_chromatic.py => test_chromatic_jax.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/jax/{test_chromatic.py => test_chromatic_jax.py} (100%) diff --git a/tests/jax/test_chromatic.py b/tests/jax/test_chromatic_jax.py similarity index 100% rename from tests/jax/test_chromatic.py rename to tests/jax/test_chromatic_jax.py From 86249e66dec05bb6cb7f8fad93d5e58a449ecca0 Mon Sep 17 00:00:00 2001 From: MaxRonce Date: Tue, 16 Jun 2026 14:24:44 +0200 Subject: [PATCH 04/12] Align chromatic API docs with upstream --- docs/api/chromatic.rst | 6 +++--- jax_galsim/__init__.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/api/chromatic.rst b/docs/api/chromatic.rst index 402224e7..c535e90a 100644 --- a/docs/api/chromatic.rst +++ b/docs/api/chromatic.rst @@ -1,5 +1,5 @@ -Chromatic Profiles -================== +Wavelength-dependent Profiles +============================= .. currentmodule:: jax_galsim @@ -49,7 +49,7 @@ Chromatic objects :members: :show-inheritance: -.. autoclass:: Chromatic +.. autoclass:: SimpleChromaticTransformation :members: :show-inheritance: diff --git a/jax_galsim/__init__.py b/jax_galsim/__init__.py index 1b681ae5..9caf7373 100644 --- a/jax_galsim/__init__.py +++ b/jax_galsim/__init__.py @@ -111,7 +111,7 @@ from .bandpass import Bandpass from .chromatic import ( ChromaticObject, - Chromatic, + SimpleChromaticTransformation, ChromaticAtmosphere, ChromaticConvolution, ChromaticSum, From f1cdc920910148245f7fe2b0cc93548fa2f3d260 Mon Sep 17 00:00:00 2001 From: MaxRonce Date: Tue, 16 Jun 2026 14:25:01 +0200 Subject: [PATCH 05/12] Make Bandpass throughput tracing-safe --- jax_galsim/bandpass.py | 235 +++++++++++++++++++++-------------------- 1 file changed, 123 insertions(+), 112 deletions(-) diff --git a/jax_galsim/bandpass.py b/jax_galsim/bandpass.py index a6046e28..5cc62ca2 100644 --- a/jax_galsim/bandpass.py +++ b/jax_galsim/bandpass.py @@ -1,139 +1,126 @@ -"""Bandpass filter representation for chromatic simulations.""" +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: - """Wavelength-dependent throughput of an observing bandpass. - - Represents the combined throughput of optics, filter, and detector - as a function of wavelength. The throughput array is a JAX-traced - parameter, so gradients can flow through it if needed (e.g. for - filter-design optimisation). - - Parameters - ---------- - wave : array_like - Wavelength array **in nanometers**, strictly increasing. - Treated as static (not traced by JAX). - throughput : array_like - Dimensionless throughput ∈ [0, 1] at each wavelength. - Treated as a JAX-traced parameter. - blue_limit : float, optional - Override the short-wavelength cut-off in nm. Defaults to - ``wave[0]``. - red_limit : float, optional - Override the long-wavelength cut-off in nm. Defaults to - ``wave[-1]``. - - Examples - -------- - Construct from arrays:: - - >>> import jax.numpy as jnp - >>> from jax_galsim.bandpass import Bandpass - >>> wave = jnp.linspace(550, 700, 100) - >>> bp = Bandpass(wave, jnp.ones(100)) - >>> float(bp(625.0)) - 1.0 - >>> float(bp.effective_wavelength) - 625.0 - - Top-hat filter between 600 nm and 700 nm:: - - >>> wave = jnp.array([500., 600., 600.001, 700., 700.001, 800.]) - >>> thru = jnp.array([0., 0., 1., 1., 0., 0.]) - >>> bp = Bandpass(wave, thru) - """ - 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 len(self._throughput) != len(self._wave): - raise ValueError("throughput must have the same length as wave.") + if self._throughput.shape != self._wave.shape: + raise ValueError("throughput must have the same shape as wave.") self._blue_limit = ( - float(blue_limit) if blue_limit is not None else float(self._wave[0]) + cast_to_float(blue_limit) + if blue_limit is not None + else cast_to_float(self._wave[0]) ) self._red_limit = ( - float(red_limit) if red_limit is not None else float(self._wave[-1]) + cast_to_float(red_limit) + if red_limit is not None + else cast_to_float(self._wave[-1]) ) - # Precompute effective wavelength at construction time so it is a - # concrete Python float and can be used as a static value under JIT. - _w = jnp.linspace(self._blue_limit, self._red_limit, 512) - _t = jnp.interp(_w, self._wave, self._throughput) - _norm = jnp.trapezoid(_t, _w) - self._effective_wavelength_val = ( - float(jnp.trapezoid(_w * _t, _w) / _norm) - if float(_norm) > 0 - else float(0.5 * (self._blue_limit + self._red_limit)) - ) - - # ------------------------------------------------------------------ - # Properties - # ------------------------------------------------------------------ - @property + @implements( + None, + lax_description="JAX-GalSim-specific static wavelength grid for array-backed bandpasses.", + ) def wave(self): - """Wavelength grid in nm (JAX array, static).""" return self._wave @property + @implements( + None, + lax_description="JAX-GalSim-specific traced throughput array for array-backed bandpasses.", + ) def throughput(self): - """Throughput array (JAX array, traced).""" return self._throughput @property + @implements( + None, + lax_description="JAX-GalSim-specific blue wavelength limit.", + ) def blue_limit(self): - """Short-wavelength cut-off in nm.""" return self._blue_limit @property + @implements( + None, + lax_description="JAX-GalSim-specific red wavelength limit.", + ) def red_limit(self): - """Long-wavelength cut-off in nm.""" return self._red_limit @property + @implements(_galsim.Bandpass.effective_wavelength) def effective_wavelength(self): - """Flux-weighted mean wavelength (concrete Python float, safe under JIT).""" - return self._effective_wavelength_val - - # ------------------------------------------------------------------ - # Evaluation - # ------------------------------------------------------------------ + 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): - """Evaluate throughput at wavelength(s) in nm. - - Returns 0 outside the defined range. - - Parameters - ---------- - wave : float or array_like - Wavelength(s) in nm. - - Returns - ------- - jnp.ndarray - Throughput at the requested wavelengths. - """ wave = jnp.asarray(wave, dtype=float) - return jnp.interp(wave, self._wave, self._throughput, left=0.0, right=0.0) - - # ------------------------------------------------------------------ - # Arithmetic - # ------------------------------------------------------------------ - + 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): - """Multiply throughput by a scalar or another Bandpass.""" if isinstance(other, Bandpass): - # Union wavelength grid 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) @@ -141,38 +128,45 @@ def __mul__(self, other): 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 + 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): - """Return a new Bandpass with tighter wavelength limits.""" 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) - - # ------------------------------------------------------------------ - # Convenience constructors - # ------------------------------------------------------------------ + 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): """Uniform throughput = 1 between blue_limit and red_limit.""" wave = jnp.linspace(blue_limit, red_limit, n_wave) return cls(wave, jnp.ones(n_wave)) - # ------------------------------------------------------------------ - # JAX pytree interface - # ------------------------------------------------------------------ - def tree_flatten(self): children = (self._throughput,) aux_data = { "wave": tuple(self._wave.tolist()), - "blue_limit": self._blue_limit, - "red_limit": self._red_limit, + "blue_limit": ensure_hashable(self._blue_limit), + "red_limit": ensure_hashable(self._red_limit), } return (children, aux_data) @@ -185,12 +179,29 @@ def tree_unflatten(cls, aux_data, children): red_limit=aux_data["red_limit"], ) - # ------------------------------------------------------------------ - # Misc - # ------------------------------------------------------------------ - def __repr__(self): return ( f"Bandpass(wave=[{self._blue_limit:.1f}, {self._red_limit:.1f}] nm, " - f"lam_eff={self.effective_wavelength:.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), + ) ) From 8ffb1a591fa8ed393fb42c10b8cb982d1cd94f64 Mon Sep 17 00:00:00 2001 From: MaxRonce Date: Tue, 16 Jun 2026 14:25:19 +0200 Subject: [PATCH 06/12] Annotate SED chromatic API subset --- jax_galsim/sed.py | 180 +++++++++++++++++----------------------------- 1 file changed, 65 insertions(+), 115 deletions(-) diff --git a/jax_galsim/sed.py b/jax_galsim/sed.py index bf9b816d..1d6e050f 100644 --- a/jax_galsim/sed.py +++ b/jax_galsim/sed.py @@ -1,162 +1,99 @@ -"""Spectral Energy Distribution (SED) for chromatic profiles. - -Designed for JAX compatibility: the flux array is a traced parameter, -so gradients flow through SED values (e.g. from DSPS outputs). -""" +import os +import galsim as _galsim import jax.numpy as jnp from jax.tree_util import register_pytree_node_class +from jax_galsim.core.utils import ensure_hashable, implements + +@implements( + _galsim.SED, + lax_description="""\ +JAX-GalSim supports array-backed SEDs only. The wavelength grid is static +and the flux array may be traced. File input, string expressions, units, +LookupTable metadata, thinning, and magnitude/zeropoint helpers are not +implemented. +""", +) @register_pytree_node_class class SED: - """Spectral Energy Distribution. - - Represents flux density as a function of wavelength, designed for - full JAX compatibility. The ``flux`` array is a JAX-traced parameter, - enabling gradients through SED parameters (e.g. outputs of DSPS). - - The wavelength grid (``wave``) is treated as static auxiliary data: - it defines the interpolation structure and is not differentiated. - - Parameters - ---------- - wave : array_like - Wavelength array **in nanometers**. Must be strictly increasing. - Treated as static (not traced by JAX). - flux : array_like - Flux density at each wavelength. Treated as a JAX-traced parameter. - Units are arbitrary but must be consistent across the simulation - (typically photons / nm / cm² / s for spectral SEDs, or - dimensionless for shape-only profiles). - redshift : float, optional - Cosmological redshift applied to the SED. The observed wavelength - grid is shifted to ``wave_obs = wave_rest * (1 + redshift)``. - Default 0. - - Examples - -------- - Basic construction from arrays:: - - >>> import jax.numpy as jnp - >>> from jax_galsim.sed import SED - >>> wave = jnp.linspace(300, 1100, 512) # nm - >>> flux = jnp.ones(512) - >>> sed = SED(wave, flux) - >>> float(sed(550.0)) - 1.0 - - DSPS workflow — flux is a traced JAX array:: - - >>> flux = dsps_model(params) # JAX array - >>> sed = SED(dsps_wave_nm, flux) - >>> image = chromatic_galaxy.drawImage(bandpass) # differentiable - - Redshifted SED:: - - >>> sed_z = SED(wave, flux, redshift=0.5) - >>> sed_z(825.0) # queries rest-frame 550 nm - """ - def __init__(self, wave, flux, redshift=0.0): + if isinstance(wave, (str, bytes, os.PathLike)) or isinstance( + flux, (str, bytes, os.PathLike) + ): + raise NotImplementedError( + "JAX-GalSim SED supports array-backed SEDs only; file input, " + "string expressions, and flux-type strings are not implemented." + ) + self._wave = jnp.asarray(wave, dtype=float) # static, not traced self._flux = jnp.asarray(flux) # traced self._redshift = jnp.asarray(redshift, dtype=float) 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 len(self._flux) != len(self._wave): - raise ValueError("flux must have the same length as wave.") - - # ------------------------------------------------------------------ - # Properties - # ------------------------------------------------------------------ + if self._flux.shape != self._wave.shape: + raise ValueError("flux must have the same shape as wave.") @property + @implements( + None, + lax_description="JAX-GalSim-specific static wavelength grid for array-backed SEDs.", + ) def wave(self): - """Rest-frame wavelength grid in nm (JAX array, static).""" return self._wave @property + @implements( + None, + lax_description="JAX-GalSim-specific traced flux array for array-backed SEDs.", + ) def flux(self): - """Flux density array (JAX array, traced).""" return self._flux @property + @implements( + None, + lax_description="JAX-GalSim-specific traced redshift scalar for array-backed SEDs.", + ) def redshift(self): - """Redshift (JAX scalar).""" return self._redshift @property + @implements( + None, + lax_description="JAX-GalSim-specific observed blue wavelength limit.", + ) def blue_limit(self): - """Shortest observed wavelength in nm.""" return float(self._wave[0]) * (1.0 + float(self._redshift)) @property + @implements( + None, + lax_description="JAX-GalSim-specific observed red wavelength limit.", + ) def red_limit(self): - """Longest observed wavelength in nm.""" return float(self._wave[-1]) * (1.0 + float(self._redshift)) - # ------------------------------------------------------------------ - # Evaluation - # ------------------------------------------------------------------ - + @implements(_galsim.SED.__call__) def __call__(self, wave): - """Evaluate flux density at observed wavelength(s) in nm. - - Uses linear interpolation; returns 0 outside the defined range. - - Parameters - ---------- - wave : float or array_like - Observed wavelength(s) in nm. - - Returns - ------- - jnp.ndarray - Flux density at the requested wavelengths. - """ wave = jnp.asarray(wave, dtype=float) # Convert observed wavelength to rest-frame before interpolating wave_rest = wave / (1.0 + self._redshift) return jnp.interp(wave_rest, self._wave, self._flux, left=0.0, right=0.0) - # ------------------------------------------------------------------ - # Flux through a bandpass - # ------------------------------------------------------------------ - + @implements(_galsim.SED.calculateFlux) def calculateFlux(self, bandpass, n_waves=512): - """Integrate SED through a bandpass: ``∫ SED(λ) × BP(λ) dλ``. - - Parameters - ---------- - bandpass : Bandpass - Observing bandpass. - n_waves : int, optional - Number of quadrature points. Default 512. - - Returns - ------- - jnp.ndarray - Scalar flux value. - """ waves = jnp.linspace(bandpass.blue_limit, bandpass.red_limit, n_waves) return jnp.trapezoid(self(waves) * bandpass(waves), waves) - # ------------------------------------------------------------------ - # Arithmetic - # ------------------------------------------------------------------ - - def withRedshift(self, redshift): - """Return a copy of this SED at a new redshift.""" + @implements(_galsim.SED.atRedshift) + def atRedshift(self, redshift): return SED(self._wave, self._flux, redshift) + @implements(_galsim.SED.__mul__) def __mul__(self, other): - """Multiply SED by a scalar or another SED. - - SED × scalar scales all flux values. - SED × SED multiplies flux densities (both evaluated on ``self``'s grid). - """ if isinstance(other, SED): # Evaluate other SED on self's rest-frame grid and multiply other_flux = jnp.interp( @@ -170,21 +107,24 @@ def __mul__(self, other): from jax_galsim.gsobject import GSObject if isinstance(other, GSObject): - from jax_galsim.chromatic import Chromatic + from jax_galsim.chromatic import SimpleChromaticTransformation - return Chromatic(other, self) + return SimpleChromaticTransformation(other, self) from jax_galsim.chromatic import ChromaticObject if isinstance(other, ChromaticObject): return other * self return SED(self._wave, self._flux * other, self._redshift) + @implements(getattr(_galsim.SED, "__rmul__", None)) def __rmul__(self, other): return self.__mul__(other) + @implements(getattr(_galsim.SED, "__truediv__", None)) def __truediv__(self, other): return SED(self._wave, self._flux / other, self._redshift) + @implements(_galsim.SED.__add__) def __add__(self, other): if isinstance(other, SED): other_flux = jnp.interp( @@ -235,6 +175,16 @@ def __eq__(self, other): return False return ( jnp.array_equal(self._wave, other._wave) - and jnp.array_equal(self._flux, other._flux) - and jnp.array_equal(self._redshift, other._redshift) + & jnp.array_equal(self._flux, other._flux) + & jnp.array_equal(self._redshift, other._redshift) + ) + + def __hash__(self): + return hash( + ( + "galsim.SED", + ensure_hashable(self._wave), + ensure_hashable(self._flux), + ensure_hashable(self._redshift), + ) ) From 294d4eef4abc2b1aa2a9246e87e642bab68da9fb Mon Sep 17 00:00:00 2001 From: MaxRonce Date: Tue, 16 Jun 2026 14:25:40 +0200 Subject: [PATCH 07/12] Fix chromatic API semantics --- jax_galsim/chromatic.py | 224 ++++++++++++++++++++++++++++++++++++---- 1 file changed, 203 insertions(+), 21 deletions(-) diff --git a/jax_galsim/chromatic.py b/jax_galsim/chromatic.py index aaba0e81..ac3bccd4 100644 --- a/jax_galsim/chromatic.py +++ b/jax_galsim/chromatic.py @@ -19,7 +19,7 @@ :: ChromaticObject base class, non-separable draw by default - ├── Chromatic GSObject × SED (separable, fast path) + ├── SimpleChromaticTransformation GSObject × SED (separable, fast path) ├── ChromaticAtmosphere Gaussian PSF with FWHM ∝ λ^alpha └── ChromaticConvolution convolution of any chromatic objects @@ -45,10 +45,12 @@ * ``jax.grad`` flows through SED flux arrays (e.g. DSPS outputs). """ +import galsim as _galsim import jax import jax.numpy as jnp from jax.tree_util import register_pytree_node_class +from jax_galsim.core.utils import cast_to_float, implements from jax_galsim.gsparams import GSParams from jax_galsim.position import PositionD @@ -70,6 +72,12 @@ def _pixel_scale_from_kwargs(kwargs): return 1.0 +def _setup_wavelength_from_bandpass(bandpass): + """Return compile-time wavelength for FFT sizing, independent of throughput.""" + with jax.ensure_compile_time_eval(): + return float(cast_to_float(0.5 * (bandpass.blue_limit + bandpass.red_limit))) + + def _make_setup_image(profile, kwargs): """Build draw target without running full ``drawImage`` setup when size is fixed.""" from jax_galsim.image import Image @@ -128,6 +136,14 @@ def _static_kcoords(kimage, wrap_size, pixel_scale): # --------------------------------------------------------------------------- +@implements( + _galsim.ChromaticObject, + lax_description="""\ +JAX-GalSim implements the differentiable FFT drawing subset. Chromatic +transformations, interpolation, photon shooting, and full SED unit semantics +are not implemented. +""", +) class ChromaticObject: """Base class for wavelength-dependent profiles. @@ -152,14 +168,24 @@ def __init__(self, obj=None): self._separable = True @property + @implements(getattr(_galsim.ChromaticObject, "separable", None)) def separable(self): """True if the profile factors as g(x,y) × h(λ).""" return self._separable + @property + @implements(_galsim.ChromaticObject.gsparams) + def gsparams(self): + base_obj = getattr(self, "_base_obj", None) + if base_obj is not None: + return base_obj.gsparams + return GSParams.default + # ------------------------------------------------------------------ # Interface that subclasses must implement # ------------------------------------------------------------------ + @implements(_galsim.ChromaticObject.evaluateAtWavelength) def evaluateAtWavelength(self, wave): """Return the monochromatic GSObject at wavelength *wave* (nm). @@ -182,10 +208,36 @@ def evaluateAtWavelength(self, wave): f"{self.__class__.__name__} must implement evaluateAtWavelength." ) + @implements(_galsim.ChromaticObject.withGSParams) + def withGSParams(self, gsparams=None, **kwargs): + base_obj = getattr(self, "_base_obj", None) + if base_obj is None: + raise NotImplementedError( + f"{self.__class__.__name__} must implement withGSParams." + ) + gsparams = GSParams.check(gsparams, base_obj.gsparams, **kwargs) + return ChromaticObject(base_obj.withGSParams(gsparams)) + + @implements(_galsim.ChromaticObject.calculateFlux) + def calculateFlux(self, bandpass, n_waves=512): + waves = jnp.linspace(bandpass.blue_limit, bandpass.red_limit, n_waves) + fluxes = jax.vmap( + lambda wave: self.evaluateAtWavelength(wave).flux * bandpass(wave) + )(waves) + return jnp.trapezoid(fluxes, waves) + + @implements(_galsim.ChromaticObject.atRedshift) + def atRedshift(self, redshift): + return self + # ------------------------------------------------------------------ # Drawing # ------------------------------------------------------------------ + @implements( + _galsim.ChromaticObject.drawImage, + lax_description="JAX-GalSim supports FFT drawing through array-backed Bandpass objects.", + ) def drawImage(self, bandpass, n_waves=64, **kwargs): """Draw the bandpass-integrated image. @@ -220,7 +272,7 @@ def _drawImage_separable(self, bandpass, n_waves, **kwargs): **Phase 1 - static setup**: Build image bounds, k-grid, and base k-values using the unit-flux spatial profile. All shape parameters must be concrete Python - scalars at this stage (true for ``Chromatic`` where the spatial + scalars at this stage (true for ``SimpleChromaticTransformation`` where the spatial profile has static params); the SED flux is NOT evaluated here. **Phase 2 — traced computation**: @@ -231,7 +283,7 @@ def _drawImage_separable(self, bandpass, n_waves, **kwargs): from jax_galsim.convolve import Convolve from jax_galsim.image import Image - wave_eff = bandpass.effective_wavelength # concrete Python float + wave_eff = _setup_wavelength_from_bandpass(bandpass) pixel_scale = _pixel_scale_from_kwargs(kwargs) # ------------------------------------------------------------------ @@ -292,8 +344,8 @@ def _static_spatial_profile(self, wave_eff): """Return unit-flux spatial profile at *wave_eff* with static params. *wave_eff* must be a concrete Python float (use - ``bandpass.effective_wavelength``). Subclasses override this when - they have a dedicated static spatial object (e.g. ``Chromatic`` has + a compile-time wavelength). Subclasses override this when + they have a dedicated static spatial object (e.g. ``SimpleChromaticTransformation`` has ``self.obj``). The default calls ``evaluateAtWavelength`` with a Python float — works when all shape params are Python scalars. """ @@ -319,7 +371,7 @@ def _drawImage_nonseparable(self, bandpass, n_waves, **kwargs): from jax_galsim.convolve import Convolve from jax_galsim.image import Image - wave_eff = bandpass.effective_wavelength # static Python float + wave_eff = _setup_wavelength_from_bandpass(bandpass) pixel_scale = _pixel_scale_from_kwargs(kwargs) # ------------------------------------------------------------------ @@ -391,12 +443,14 @@ def kvals_at_wave(wave): # Operator overloads # ------------------------------------------------------------------ + @implements(_galsim.ChromaticObject.__add__) def __add__(self, other): return ChromaticSum([self, other]) def __radd__(self, other): return ChromaticSum([other, self]) + @implements(_galsim.ChromaticObject.__mul__) def __mul__(self, other): from jax_galsim.sed import SED @@ -406,9 +460,10 @@ def __mul__(self, other): raise TypeError( "Only achromatic ChromaticObject wrappers can be multiplied by SED." ) - return Chromatic(base_obj, other) + return SimpleChromaticTransformation(base_obj, other) return _ScaledChromaticObject(self, other) + @implements(_galsim.ChromaticObject.__rmul__) def __rmul__(self, other): return self.__mul__(other) @@ -436,6 +491,10 @@ def _sed_value(self, wave): # --------------------------------------------------------------------------- +@implements( + _galsim.ChromaticSum, + lax_description="JAX-GalSim treats ChromaticSum as non-separable, matching GalSim semantics.", +) class ChromaticSum(ChromaticObject): """Sum of two or more chromatic profiles. @@ -447,7 +506,7 @@ class ChromaticSum(ChromaticObject): obj_list : list of ChromaticObject """ - _separable = False # conservative; optimised later if all are separable + _separable = False def __init__(self, *args): if len(args) == 0: @@ -456,13 +515,34 @@ def __init__(self, *args): self.obj_list = list(args[0]) else: self.obj_list = list(args) - self._separable = all(o._separable for o in self.obj_list) + self._separable = False + + @property + @implements(_galsim.ChromaticSum.gsparams) + def gsparams(self): + return GSParams.combine([obj.gsparams for obj in self.obj_list]) + + @implements(_galsim.ChromaticSum.withGSParams) + def withGSParams(self, gsparams=None, **kwargs): + gsparams = GSParams.check(gsparams, self.gsparams, **kwargs) + return ChromaticSum([obj.withGSParams(gsparams) for obj in self.obj_list]) + + @implements(_galsim.ChromaticSum.atRedshift) + def atRedshift(self, redshift): + return ChromaticSum([obj.atRedshift(redshift) for obj in self.obj_list]) + @implements(_galsim.ChromaticSum.evaluateAtWavelength) def evaluateAtWavelength(self, wave): from jax_galsim.sum import Sum return Sum([o.evaluateAtWavelength(wave) for o in self.obj_list]) + @implements( + _galsim.ChromaticSum.drawImage, + lax_description=( + "JAX-GalSim draws sums component-by-component to preserve non-separable semantics." + ), + ) def drawImage(self, bandpass, n_waves=64, **kwargs): # Draw each component and sum images = [ @@ -475,12 +555,19 @@ def drawImage(self, bandpass, n_waves=64, **kwargs): # --------------------------------------------------------------------------- -# Chromatic — separable GSObject × SED +# SimpleChromaticTransformation — separable GSObject × SED # --------------------------------------------------------------------------- +@implements( + _galsim.SimpleChromaticTransformation, + lax_description="""\ +JAX-GalSim implements the simple separable case GSObject * SED. General +chromatic affine transformations are not implemented. +""", +) @register_pytree_node_class -class Chromatic(ChromaticObject): +class SimpleChromaticTransformation(ChromaticObject): """Separable chromatic profile: a GSObject multiplied by an SED. The spatial profile is fixed; wavelength dependence enters only @@ -488,7 +575,7 @@ class Chromatic(ChromaticObject): ``I(x, y, λ) = g(x, y) × SED(λ)`` - Drawing a ``Chromatic`` through a bandpass reduces to a single + Drawing a ``SimpleChromaticTransformation`` through a bandpass reduces to a single monochromatic draw at the effective wavelength with total flux ``∫ SED(λ) × BP(λ) dλ``. @@ -522,6 +609,21 @@ def __init__(self, obj, sed): self.obj = obj self.sed = sed + @property + @implements(_galsim.SimpleChromaticTransformation.gsparams) + def gsparams(self): + return self.obj.gsparams + + @implements(_galsim.SimpleChromaticTransformation.withGSParams) + def withGSParams(self, gsparams=None, **kwargs): + gsparams = GSParams.check(gsparams, self.obj.gsparams, **kwargs) + return SimpleChromaticTransformation(self.obj.withGSParams(gsparams), self.sed) + + @implements(_galsim.SimpleChromaticTransformation.atRedshift) + def atRedshift(self, redshift): + return SimpleChromaticTransformation(self.obj, self.sed.atRedshift(redshift)) + + @implements(_galsim.SimpleChromaticTransformation.evaluateAtWavelength) def evaluateAtWavelength(self, wave): """Return the GSObject scaled to SED(wave).""" return self.obj.withScaledFlux(self.sed(wave)) @@ -544,7 +646,12 @@ def tree_unflatten(cls, aux_data, children): return cls(obj=children[0], sed=children[1]) def __repr__(self): - return f"Chromatic({self.obj!r}, {self.sed!r})" + return f"galsim.SimpleChromaticTransformation({self.obj!r}, sed={self.sed!r})" + + +# Backwards-compatible internal alias. Not exported at top level because upstream +# GalSim exposes SimpleChromaticTransformation, not Chromatic. +Chromatic = SimpleChromaticTransformation # --------------------------------------------------------------------------- @@ -553,6 +660,15 @@ def __repr__(self): @register_pytree_node_class +@implements( + _galsim.ChromaticAtmosphere, + lax_description="""\ +JAX-GalSim implements a differentiable Gaussian/Moffat seeing PSF with +FWHM(lambda) = fwhm_ref * (lambda / lam_ref)**alpha. Differential +chromatic refraction and atmospheric coordinate parameters are not +implemented. +""", +) class ChromaticAtmosphere(ChromaticObject): """Atmospheric PSF with a power-law wavelength-dependent FWHM. @@ -564,7 +680,8 @@ class ChromaticAtmosphere(ChromaticObject): This profile carries a **flat (dimensionless) SED**: ``SED(λ) = 1``. The physical SED is typically attached to the galaxy component via - :class:`Chromatic`, and passed to :class:`ChromaticConvolution`. + :class:`SimpleChromaticTransformation`, and passed to + :class:`ChromaticConvolution`. Parameters ---------- @@ -617,21 +734,55 @@ def __init__( # ------------------------------------------------------------------ @property + @implements( + None, + lax_description="JAX-GalSim-specific reference FWHM for ChromaticAtmosphere.", + ) def fwhm_ref(self): return self._fwhm_ref @property + @implements( + None, + lax_description="JAX-GalSim-specific reference wavelength for ChromaticAtmosphere.", + ) def lam_ref(self): return self._lam_ref @property + @implements( + None, + lax_description="JAX-GalSim-specific wavelength power-law index for ChromaticAtmosphere.", + ) def alpha(self): return self._alpha + @property + @implements(_galsim.ChromaticAtmosphere.gsparams) + def gsparams(self): + return self._gsparams + + @implements(_galsim.ChromaticAtmosphere.withGSParams) + def withGSParams(self, gsparams=None, **kwargs): + gsparams = GSParams.check(gsparams, self._gsparams, **kwargs) + return ChromaticAtmosphere( + fwhm_ref=self._fwhm_ref, + lam_ref=self._lam_ref, + alpha=self._alpha, + profile=self._profile, + moffat_beta=self._moffat_beta, + gsparams=gsparams, + ) + + @implements(_galsim.ChromaticAtmosphere.atRedshift) + def atRedshift(self, redshift): + return self + # ------------------------------------------------------------------ # Core method # ------------------------------------------------------------------ + @implements(_galsim.ChromaticAtmosphere.evaluateAtWavelength) def evaluateAtWavelength(self, wave): """Return the PSF profile (unit flux) at wavelength *wave* (nm). @@ -705,6 +856,14 @@ def __repr__(self): # --------------------------------------------------------------------------- +@implements( + _galsim.ChromaticConvolution, + lax_description="""\ +JAX-GalSim supports FFT drawing of chromatic convolutions. Real-space +chromatic convolution, photon shooting, and full GalSim chromatic +transformations are not implemented. +""", +) class ChromaticConvolution(ChromaticObject): """Convolution of multiple chromatic profiles. @@ -725,14 +884,14 @@ class ChromaticConvolution(ChromaticObject): ---------- obj_list : list of ChromaticObject or GSObject Components to convolve. Plain ``GSObject`` instances are wrapped - automatically in a flat-SED :class:`Chromatic`. + automatically in a flat-SED :class:`SimpleChromaticTransformation`. Examples -------- :: >>> from jax_galsim import Gaussian, Convolve - >>> from jax_galsim.chromatic import Chromatic, ChromaticAtmosphere, ChromaticConvolution + >>> from jax_galsim.chromatic import ChromaticAtmosphere, ChromaticConvolution >>> from jax_galsim.sed import SED >>> from jax_galsim.bandpass import Bandpass >>> import jax.numpy as jnp @@ -784,15 +943,17 @@ def __init__(self, *args, **kwargs): wave_stub = jnp.array([100.0, 2000.0]) flat_sed = SED(wave_stub, jnp.ones(2)) - processed.append(Chromatic(obj, flat_sed)) + processed.append(SimpleChromaticTransformation(obj, flat_sed)) else: processed.append(obj) self.obj_list = processed @property + @implements(_galsim.ChromaticConvolution.gsparams) def gsparams(self): return self._gsparams + @implements(_galsim.ChromaticConvolution.withGSParams) def withGSParams(self, gsparams=None, **kwargs): ret = self.__class__.__new__(self.__class__) ret.obj_list = self.obj_list @@ -800,6 +961,15 @@ def withGSParams(self, gsparams=None, **kwargs): ret._propagate_gsparams = self._propagate_gsparams return ret + @implements(_galsim.ChromaticConvolution.atRedshift) + def atRedshift(self, redshift): + return ChromaticConvolution( + [obj.atRedshift(redshift) for obj in self.obj_list], + gsparams=self._gsparams, + propagate_gsparams=self._propagate_gsparams, + ) + + @implements(_galsim.ChromaticConvolution.evaluateAtWavelength) def evaluateAtWavelength(self, wave): """Return the convolved monochromatic profile at *wave* (nm).""" from jax_galsim.convolve import Convolve @@ -810,6 +980,10 @@ def evaluateAtWavelength(self, wave): # Drawing with the optimised separable split # ------------------------------------------------------------------ + @implements( + _galsim.ChromaticConvolution.drawImage, + lax_description="JAX-GalSim supports FFT drawing with static wavelength grid size.", + ) def drawImage(self, bandpass, n_waves=64, **kwargs): """Draw the integrated image, exploiting separability where possible. @@ -834,7 +1008,7 @@ def drawImage(self, bandpass, n_waves=64, **kwargs): sep_objs = [o for o in self.obj_list if o._separable] nonsep_objs = [o for o in self.obj_list if not o._separable] - wave_eff = bandpass.effective_wavelength # static Python float + wave_eff = _setup_wavelength_from_bandpass(bandpass) waves = jnp.linspace(bandpass.blue_limit, bandpass.red_limit, n_waves) pixel_scale = _pixel_scale_from_kwargs(kwargs) @@ -963,20 +1137,28 @@ def __repr__(self): # --------------------------------------------------------------------------- -# Monkey-patch GSObject.__mul__ to return Chromatic when multiplied by SED +# Monkey-patch GSObject.__mul__ to return a simple chromatic object for SEDs # --------------------------------------------------------------------------- +@implements( + _galsim.GSObject.__mul__, + lax_description="Also accepts array-backed JAX-GalSim SED objects and returns a SimpleChromaticTransformation.", +) def _gsobject_mul_sed(self, other): - """Allow ``gsobject * sed → Chromatic(gsobject, sed)``.""" + """Allow ``gsobject * sed``.""" from jax_galsim.sed import SED if isinstance(other, SED): - return Chromatic(self, other) + return SimpleChromaticTransformation(self, other) # Fall through to original implementation (flux scaling) return self.withScaledFlux(other) +@implements( + _galsim.GSObject.__rmul__, + lax_description="Also accepts array-backed JAX-GalSim SED objects and returns a SimpleChromaticTransformation.", +) def _gsobject_rmul_sed(self, other): return _gsobject_mul_sed(self, other) From d071dd145cb204edb248cca977f85f5dc3809b2c Mon Sep 17 00:00:00 2001 From: MaxRonce Date: Tue, 16 Jun 2026 14:26:05 +0200 Subject: [PATCH 08/12] Add chromatic API and JAX regression tests --- tests/conftest.py | 7 +- tests/galsim_tests_config.yaml | 14 +- tests/jax/test_api.py | 174 ++++++++++++ tests/jax/test_chromatic_jax.py | 453 ++++++++------------------------ 4 files changed, 297 insertions(+), 351 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 53d2e1c4..6b0082a9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -98,7 +98,12 @@ def pytest_collection_modifyitems(config, items): ): item.add_marker(skip) - if any([t in item.nodeid for t in test_config["skipped_tests"]["coord"]]): + skipped_tests = test_config["skipped_tests"].get("galsim", []) + test_config[ + "skipped_tests" + ].get("coord", []) + if any( + [item.nodeid == t or item.nodeid.startswith(t + "[") for t in skipped_tests] + ): item.add_marker(skip) diff --git a/tests/galsim_tests_config.yaml b/tests/galsim_tests_config.yaml index 1391e880..c172d406 100644 --- a/tests/galsim_tests_config.yaml +++ b/tests/galsim_tests_config.yaml @@ -34,6 +34,18 @@ enabled_tests: # tests to explicitly skip # applied on top of the enabled set above skipped_tests: + galsim: + # JAX-GalSim's chromatic subset is array-backed and does not implement + # GalSim spectral file/string expression parsing, zeropoints, DCR moment + # helpers, chromatic affine transforms, or photon shooting. + - "tests/GalSim/tests/test_sed.py::test_SED_withFlux" + - "tests/GalSim/tests/test_sed.py::test_SED_calculateDCRMomentShifts" + - "tests/GalSim/tests/test_sed.py::test_SED_calculateSeeingMomentRatio" + - "tests/GalSim/tests/test_sed.py::test_broadcast" + - "tests/GalSim/tests/test_bandpass.py::test_zp" + - "tests/GalSim/tests/test_bandpass.py::test_truncate_inputs" + - "tests/GalSim/tests/test_chromatic.py::test_ChromaticObject_compound_affine_transformation" + - "tests/GalSim/tests/test_chromatic.py::test_shoot_transformation" coord: - "tests/Coord/tests/test_celestial.py::test_xyz_raises" - "tests/Coord/tests/test_celestial.py::test_greatcircle_raises" @@ -114,8 +126,6 @@ allowed_failures: - "SED.__init__() got an unexpected keyword argument 'spec'" - "SED.__init__() got an unexpected keyword argument 'wave_type'" - "Bandpass.__init__() got an unexpected keyword argument 'wave_type'" - - "could not convert string to float" - - "is not a valid JAX array type" - "'ChromaticObject' object has no attribute 'dilate'" - "module 'jax_galsim' has no attribute 'getCOSMOSNoise'" - "GSParams.__init__() got an unexpected keyword argument 'allowed_flux_variation'" diff --git a/tests/jax/test_api.py b/tests/jax/test_api.py index 6e107f0f..a1cafc8f 100644 --- a/tests/jax/test_api.py +++ b/tests/jax/test_api.py @@ -24,6 +24,180 @@ def test_api_same(): ) +def test_api_chromatic_profiles(): + wave = jnp.linspace(500.0, 800.0, 16) + sed = jax_galsim.SED(wave, jnp.ones_like(wave)) + bandpass = jax_galsim.Bandpass(wave, jnp.ones_like(wave)) + gal = jax_galsim.Gaussian(half_light_radius=0.5) * sed + psf = jax_galsim.ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0) + chromatic_base = jax_galsim.ChromaticObject( + jax_galsim.Gaussian(half_light_radius=0.5) + ) + chromatic_sum = gal + gal + chromatic_convolution = jax_galsim.ChromaticConvolution([gal, psf]) + + objects = [ + (jax_galsim.SED, _galsim.SED, sed), + (jax_galsim.Bandpass, _galsim.Bandpass, bandpass), + (jax_galsim.ChromaticObject, _galsim.ChromaticObject, chromatic_base), + ( + jax_galsim.SimpleChromaticTransformation, + _galsim.SimpleChromaticTransformation, + gal, + ), + (jax_galsim.ChromaticAtmosphere, _galsim.ChromaticAtmosphere, psf), + (jax_galsim.ChromaticSum, _galsim.ChromaticSum, chromatic_sum), + ( + jax_galsim.ChromaticConvolution, + _galsim.ChromaticConvolution, + chromatic_convolution, + ), + ] + + for cls, gscls, obj in objects: + assert cls.__galsim_wrapped__ is gscls + assert "LAX-backend implementation" in cls.__doc__ + if hasattr(obj, "tree_flatten"): + children, aux_data = obj.tree_flatten() + rebuilt = obj.__class__.tree_unflatten(aux_data, children) + assert isinstance(rebuilt, obj.__class__) + + methods = [ + (jax_galsim.GSObject.__mul__, _galsim.GSObject.__mul__), + (jax_galsim.GSObject.__rmul__, _galsim.GSObject.__rmul__), + (jax_galsim.SED.__call__, _galsim.SED.__call__), + (jax_galsim.SED.calculateFlux, _galsim.SED.calculateFlux), + (jax_galsim.SED.atRedshift, _galsim.SED.atRedshift), + (jax_galsim.SED.__mul__, _galsim.SED.__mul__), + (jax_galsim.SED.__rmul__, _galsim.SED.__rmul__), + (jax_galsim.SED.__truediv__, _galsim.SED.__truediv__), + (jax_galsim.SED.__add__, _galsim.SED.__add__), + ( + jax_galsim.Bandpass.effective_wavelength.fget, + _galsim.Bandpass.effective_wavelength, + ), + (jax_galsim.Bandpass.__call__, _galsim.Bandpass.__call__), + ( + jax_galsim.Bandpass.calculateEffectiveWavelength, + _galsim.Bandpass.calculateEffectiveWavelength, + ), + (jax_galsim.Bandpass.__mul__, _galsim.Bandpass.__mul__), + (jax_galsim.Bandpass.__rmul__, _galsim.Bandpass.__rmul__), + (jax_galsim.Bandpass.truncate, _galsim.Bandpass.truncate), + ( + jax_galsim.ChromaticObject.evaluateAtWavelength, + _galsim.ChromaticObject.evaluateAtWavelength, + ), + (jax_galsim.ChromaticObject.gsparams.fget, _galsim.ChromaticObject.gsparams), + (jax_galsim.ChromaticObject.drawImage, _galsim.ChromaticObject.drawImage), + ( + jax_galsim.ChromaticObject.calculateFlux, + _galsim.ChromaticObject.calculateFlux, + ), + ( + jax_galsim.ChromaticObject.withGSParams, + _galsim.ChromaticObject.withGSParams, + ), + ( + jax_galsim.ChromaticObject.atRedshift, + _galsim.ChromaticObject.atRedshift, + ), + (jax_galsim.ChromaticObject.__add__, _galsim.ChromaticObject.__add__), + (jax_galsim.ChromaticObject.__mul__, _galsim.ChromaticObject.__mul__), + (jax_galsim.ChromaticObject.__rmul__, _galsim.ChromaticObject.__rmul__), + (jax_galsim.ChromaticSum.gsparams.fget, _galsim.ChromaticSum.gsparams), + (jax_galsim.ChromaticSum.withGSParams, _galsim.ChromaticSum.withGSParams), + (jax_galsim.ChromaticSum.atRedshift, _galsim.ChromaticSum.atRedshift), + ( + jax_galsim.ChromaticSum.evaluateAtWavelength, + _galsim.ChromaticSum.evaluateAtWavelength, + ), + (jax_galsim.ChromaticSum.drawImage, _galsim.ChromaticSum.drawImage), + ( + jax_galsim.SimpleChromaticTransformation.gsparams.fget, + _galsim.SimpleChromaticTransformation.gsparams, + ), + ( + jax_galsim.SimpleChromaticTransformation.withGSParams, + _galsim.SimpleChromaticTransformation.withGSParams, + ), + ( + jax_galsim.SimpleChromaticTransformation.atRedshift, + _galsim.SimpleChromaticTransformation.atRedshift, + ), + ( + jax_galsim.SimpleChromaticTransformation.evaluateAtWavelength, + _galsim.SimpleChromaticTransformation.evaluateAtWavelength, + ), + ( + jax_galsim.ChromaticAtmosphere.gsparams.fget, + _galsim.ChromaticAtmosphere.gsparams, + ), + ( + jax_galsim.ChromaticAtmosphere.withGSParams, + _galsim.ChromaticAtmosphere.withGSParams, + ), + ( + jax_galsim.ChromaticAtmosphere.atRedshift, + _galsim.ChromaticAtmosphere.atRedshift, + ), + ( + jax_galsim.ChromaticAtmosphere.evaluateAtWavelength, + _galsim.ChromaticAtmosphere.evaluateAtWavelength, + ), + ( + jax_galsim.ChromaticConvolution.gsparams.fget, + _galsim.ChromaticConvolution.gsparams, + ), + ( + jax_galsim.ChromaticConvolution.withGSParams, + _galsim.ChromaticConvolution.withGSParams, + ), + ( + jax_galsim.ChromaticConvolution.atRedshift, + _galsim.ChromaticConvolution.atRedshift, + ), + ( + jax_galsim.ChromaticConvolution.evaluateAtWavelength, + _galsim.ChromaticConvolution.evaluateAtWavelength, + ), + ( + jax_galsim.ChromaticConvolution.drawImage, + _galsim.ChromaticConvolution.drawImage, + ), + ] + for method, gs_method in methods: + wrapped = method.__galsim_wrapped__ + if wrapped is not gs_method: + assert getattr(wrapped, "__module__", None) == getattr( + gs_method, "__module__", None + ) + assert getattr(wrapped, "__name__", None) == getattr( + gs_method, "__name__", None + ) + + jax_specific_properties = [ + jax_galsim.SED.wave.fget, + jax_galsim.SED.flux.fget, + jax_galsim.SED.redshift.fget, + jax_galsim.SED.blue_limit.fget, + jax_galsim.SED.red_limit.fget, + jax_galsim.Bandpass.wave.fget, + jax_galsim.Bandpass.throughput.fget, + jax_galsim.Bandpass.blue_limit.fget, + jax_galsim.Bandpass.red_limit.fget, + jax_galsim.ChromaticObject.separable.fget, + jax_galsim.ChromaticAtmosphere.fwhm_ref.fget, + jax_galsim.ChromaticAtmosphere.lam_ref.fget, + jax_galsim.ChromaticAtmosphere.alpha.fget, + ] + for prop in jax_specific_properties: + assert prop.__galsim_wrapped__ is None + assert prop.__doc__ is not None + + assert not hasattr(jax_galsim, "Chromatic") + + OK_ERRORS = [ "got an unexpected keyword argument", "At least one GSObject must be provided", diff --git a/tests/jax/test_chromatic_jax.py b/tests/jax/test_chromatic_jax.py index f8644721..dcefc389 100644 --- a/tests/jax/test_chromatic_jax.py +++ b/tests/jax/test_chromatic_jax.py @@ -1,285 +1,107 @@ -"""Tests for jax_galsim chromatic PSF support. - -Verifies: -- SED and Bandpass construction and evaluation -- Chromatic (separable) drawImage correctness -- ChromaticAtmosphere (non-separable) drawImage correctness -- ChromaticConvolution (galaxy × SED ⊗ PSF) correctness -- jax.jit compatibility for all paths -- jax.grad compatibility for SED-flux differentiation -- Pytree round-trip (tree_flatten / tree_unflatten) -- Numerical agreement with analytic expectations -""" +"""JAX-specific chromatic tests. + +The behavioral GalSim coverage for SED, Bandpass, and chromatic profiles comes +from tests/GalSim/tests through the conftest.py harness, which imports +jax_galsim as galsim. This file is intentionally limited to JAX behavior and +local regressions that the upstream GalSim tests cannot express directly: -# ruff: noqa: E402,I001 +- jit and grad support for traced SED and Bandpass arrays +- pytree round-trips for chromatic objects +- the ChromaticSum non-separability regression +- the local GSObject * SED monkey patch +""" import jax import jax.numpy as jnp import pytest -# Enable float64 for accuracy -jax.config.update("jax_enable_x64", True) +import jax_galsim as jgs +from jax_galsim.chromatic import ( + ChromaticAtmosphere, + ChromaticConvolution, + SimpleChromaticTransformation, +) -import jax_galsim as jgal -from jax_galsim.chromatic import ChromaticAtmosphere, ChromaticConvolution - -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- - -WAVE = jnp.linspace(400.0, 900.0, 256) # nm -BP = jgal.Bandpass.tophat(550.0, 750.0) # 200 nm wide, throughput = 1 -BP_NARROW = jgal.Bandpass.tophat(600.0, 700.0) # 100 nm wide +WAVE = jnp.linspace(400.0, 900.0, 256) +BP = jgs.Bandpass.tophat(550.0, 750.0) def flat_sed(scale=1.0): - return jgal.SED(WAVE, jnp.ones(256) * scale) - - -# --------------------------------------------------------------------------- -# SED tests -# --------------------------------------------------------------------------- - - -def test_sed_evaluation(): - sed = flat_sed() - assert float(sed(600.0)) == pytest.approx(1.0, rel=1e-5) - assert float(sed(300.0)) == pytest.approx(0.0) # outside range - - -def test_sed_redshift(): - sed = jgal.SED(WAVE, jnp.ones(256), redshift=1.0) - # observed 800 nm → rest-frame 400 nm → flux = 1.0 - assert float(sed(800.0)) == pytest.approx(1.0, rel=1e-4) - # observed 400 nm → rest-frame 200 nm → outside range → 0 - assert float(sed(400.0)) == pytest.approx(0.0, abs=1e-6) - + return jgs.SED(WAVE, jnp.ones_like(WAVE) * scale) -def test_sed_calculate_flux(): - sed = flat_sed() - flux = float(sed.calculateFlux(BP)) - # ∫_550^750 1 dλ = 200 nm - assert flux == pytest.approx(200.0, rel=1e-3) - -def test_sed_pytree_roundtrip(): +def test_sed_bandpass_chromatic_pytree_roundtrips(): sed = flat_sed(2.0) - leaves, treedef = jax.tree_util.tree_flatten(sed) - sed2 = jax.tree_util.tree_unflatten(treedef, leaves) - assert float(sed2(600.0)) == pytest.approx(2.0, rel=1e-5) - - -def test_sed_arithmetic(): - sed1 = flat_sed(2.0) - sed2 = sed1 * 3.0 - assert float(sed2(600.0)) == pytest.approx(6.0, rel=1e-5) - - sed3 = sed1 + flat_sed(1.0) - assert float(sed3(600.0)) == pytest.approx(3.0, rel=1e-5) - - -# --------------------------------------------------------------------------- -# Bandpass tests -# --------------------------------------------------------------------------- - - -def test_bandpass_evaluation(): - bp = jgal.Bandpass.tophat(550.0, 750.0) - assert float(bp(625.0)) == pytest.approx(1.0) - assert float(bp(500.0)) == pytest.approx(0.0) - assert float(bp(800.0)) == pytest.approx(0.0) - - -def test_bandpass_effective_wavelength(): - bp = jgal.Bandpass.tophat(550.0, 750.0) - lam_eff = bp.effective_wavelength - assert isinstance(lam_eff, float) - assert lam_eff == pytest.approx(650.0, rel=1e-4) - - -def test_bandpass_effective_wavelength_concrete(): - """effective_wavelength must be a concrete Python float (safe under JIT).""" - bp = jgal.Bandpass.tophat(550.0, 750.0) - lam_eff = bp.effective_wavelength - # If this were a JAX tracer, float() would raise ConcretizationTypeError - assert isinstance(lam_eff, float) - - -def test_bandpass_mul(): - bp1 = jgal.Bandpass.tophat(500.0, 700.0) - bp2 = jgal.Bandpass.tophat(600.0, 800.0) - bp = bp1 * bp2 - assert float(bp(650.0)) == pytest.approx(1.0) - assert float(bp(550.0)) == pytest.approx(0.0) - assert float(bp(750.0)) == pytest.approx(0.0) - - -def test_bandpass_pytree_roundtrip(): - bp = jgal.Bandpass.tophat(550.0, 750.0) - leaves, treedef = jax.tree_util.tree_flatten(bp) - bp2 = jax.tree_util.tree_unflatten(treedef, leaves) - assert float(bp2(625.0)) == pytest.approx(1.0) - assert bp2.effective_wavelength == pytest.approx(650.0, rel=1e-4) - - -# --------------------------------------------------------------------------- -# Chromatic (separable) tests -# --------------------------------------------------------------------------- - + bandpass = jgs.Bandpass.tophat(550.0, 750.0) + gal = jgs.Gaussian(half_light_radius=0.5) * sed + psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) -def test_chromatic_construction(): - sed = flat_sed() - gal = jgal.Gaussian(half_light_radius=0.5) * sed - assert gal._separable + for obj in [sed, bandpass, gal, psf]: + leaves, treedef = jax.tree_util.tree_flatten(obj) + rebuilt = jax.tree_util.tree_unflatten(treedef, leaves) + assert isinstance(rebuilt, obj.__class__) + sed_leaves, sed_treedef = jax.tree_util.tree_flatten(sed) + rebuilt_sed = jax.tree_util.tree_unflatten(sed_treedef, sed_leaves) + assert float(rebuilt_sed(600.0)) == pytest.approx(2.0, rel=1e-5) -def test_chromatic_drawImage_flux(): - """Image pixel sum should equal ∫ SED(λ) × BP(λ) dλ.""" - sed = flat_sed() - gal = jgal.Gaussian(half_light_radius=0.5) * sed - img = gal.drawImage(BP, scale=0.2, nx=32, ny=32) - # ∫_550^750 1 dλ = 200 - assert float(img.array.sum()) == pytest.approx(200.0, rel=5e-3) +def test_bandpass_call_jit_grad_throughput(): + wave = jnp.linspace(500.0, 800.0, 32) + throughput = jnp.exp(-0.5 * ((wave - 650.0) / 45.0) ** 2) -def test_chromatic_drawImage_narrow_bandpass(): - """Narrower bandpass → smaller total flux.""" - sed = flat_sed() - gal = jgal.Gaussian(half_light_radius=0.5) * sed - img = gal.drawImage(BP_NARROW, scale=0.2, nx=32, ny=32) - assert float(img.array.sum()) == pytest.approx(100.0, rel=5e-3) + @jax.jit + def sample(tp): + return jgs.Bandpass(wave, tp)(650.0) + value = sample(throughput) + grad = jax.grad(sample)(throughput) -def test_chromatic_jit(): - @jax.jit - def render(flux): - sed = jgal.SED(WAVE, flux) - gal = jgal.Gaussian(half_light_radius=0.5) * sed - return gal.drawImage(BP, scale=0.2, nx=32, ny=32).array.sum() + assert float(value) == pytest.approx(1.0, rel=5e-3) + assert bool(jnp.all(jnp.isfinite(grad))) + assert float(jnp.max(jnp.abs(grad))) > 0.0 - result = render(jnp.ones(256)) - assert float(result) == pytest.approx(200.0, rel=5e-3) +def test_bandpass_effective_wavelength_jit_grad_throughput(): + wave = jnp.linspace(500.0, 800.0, 32) + throughput = jnp.exp(-0.5 * ((wave - 660.0) / 50.0) ** 2) -def test_chromatic_grad(): @jax.jit - def render(flux): - sed = jgal.SED(WAVE, flux) - gal = jgal.Gaussian(half_light_radius=0.5) * sed - return gal.drawImage(BP, scale=0.2, nx=32, ny=32).array.sum() + def eff(tp): + return jgs.Bandpass(wave, tp).effective_wavelength - grad = jax.grad(render)(jnp.ones(256)) - grad_arr = jnp.asarray(grad) - - # Gradient sums to total bandpass flux ≈ 200 - assert float(grad.sum()) == pytest.approx(200.0, rel=5e-2) - # Outside bandpass → zero gradient - idx_out = int((420 - 400) / (900 - 400) * 255) - assert float(grad[idx_out]) == pytest.approx(0.0, abs=1e-8) - # Inside bandpass region has positive total contribution - mask_in = (WAVE >= 550) & (WAVE <= 750) - assert float(grad_arr[mask_in].sum()) > 0.0 + value = eff(throughput) + grad = jax.grad(eff)(throughput) + assert float(value) == pytest.approx(660.0, rel=2e-2) + assert bool(jnp.all(jnp.isfinite(grad))) + assert float(jnp.max(jnp.abs(grad))) > 0.0 -def test_chromatic_jit_recompile(): - """JIT should reuse compiled code when called twice with different flux.""" +def test_separable_chromatic_draw_jit_grad_sed_flux(): @jax.jit def render(flux): - sed = jgal.SED(WAVE, flux) - gal = jgal.Gaussian(half_light_radius=0.5) * sed + sed = jgs.SED(WAVE, flux) + gal = jgs.Gaussian(half_light_radius=0.5) * sed return gal.drawImage(BP, scale=0.2, nx=32, ny=32).array.sum() - r1 = float(render(jnp.ones(256))) - r2 = float(render(jnp.ones(256) * 2.0)) - assert r2 == pytest.approx(r1 * 2.0, rel=1e-4) - + flux = jnp.ones_like(WAVE) + result = render(flux) + grad = jax.grad(render)(flux) + mask_in = (WAVE >= 550.0) & (WAVE <= 750.0) -# --------------------------------------------------------------------------- -# ChromaticAtmosphere tests -# --------------------------------------------------------------------------- - - -def test_chromatic_atmosphere_evaluate(): - psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) - prof = psf.evaluateAtWavelength(700.0) - # At reference wavelength, FWHM should be exactly fwhm_ref - # jax_galsim Gaussian exposes sigma; FWHM = sigma * fwhm_factor - assert isinstance(prof, jgal.Gaussian) - fwhm = float(prof.sigma) * jgal.Gaussian._fwhm_factor - assert fwhm == pytest.approx(0.7, rel=1e-4) - - -def test_chromatic_atmosphere_scaling(): - psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) - prof_blue = psf.evaluateAtWavelength(350.0) - prof_red = psf.evaluateAtWavelength(700.0) - # FWHM ∝ λ^alpha = λ^(-0.2) → bluer is wider (alpha < 0) - assert float(prof_blue.sigma) > float(prof_red.sigma) - - -def test_chromatic_atmosphere_drawImage_flux(): - """Total flux = ∫ BP(λ) × 1 dλ = 200 (flat SED, unit PSF flux).""" - psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) - img = psf.drawImage(BP, scale=0.2, nx=32, ny=32) - assert float(img.array.sum()) == pytest.approx(200.0, rel=5e-3) - - -def test_chromatic_atmosphere_moffat(): - psf = ChromaticAtmosphere( - fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2, profile="moffat", moffat_beta=4.765 - ) - prof = psf.evaluateAtWavelength(700.0) - assert isinstance(prof, jgal.Moffat) - img = psf.drawImage(BP, scale=0.2, nx=32, ny=32) - assert float(img.array.sum()) == pytest.approx(200.0, rel=5e-2) - - -def test_chromatic_atmosphere_pytree(): - psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) - leaves, treedef = jax.tree_util.tree_flatten(psf) - psf2 = jax.tree_util.tree_unflatten(treedef, leaves) - assert psf2._fwhm_ref == pytest.approx(0.7) - assert psf2._alpha == pytest.approx(-0.2) - - -# --------------------------------------------------------------------------- -# ChromaticConvolution tests -# --------------------------------------------------------------------------- - - -def test_chromatic_convolution_flux(): - """Galaxy × SED ⊗ ChromaticAtmosphere: flux = ∫ SED(λ) × BP(λ) dλ.""" - sed = flat_sed() - gal = jgal.Gaussian(half_light_radius=0.5) * sed - psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) - final = ChromaticConvolution([gal, psf]) - img = final.drawImage(BP, scale=0.2, nx=64, ny=64, n_waves=32) - assert float(img.array.sum()) == pytest.approx(200.0, rel=5e-2) - - -def test_chromatic_convolution_jit(): - @jax.jit - def render(flux): - sed = jgal.SED(WAVE, flux) - gal = jgal.Gaussian(half_light_radius=0.5) * sed - psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) - return ( - ChromaticConvolution([gal, psf]) - .drawImage(BP, scale=0.2, nx=64, ny=64, n_waves=32) - .array.sum() - ) - - result = render(jnp.ones(256)) - assert float(result) == pytest.approx(200.0, rel=5e-2) + assert float(result) == pytest.approx(200.0, rel=5e-3) + assert float(grad.sum()) == pytest.approx(200.0, rel=5e-2) + assert float(grad[10]) == pytest.approx(0.0, abs=1e-8) + assert float(jnp.asarray(grad)[mask_in].sum()) > 0.0 -def test_chromatic_convolution_grad(): +def test_chromatic_convolution_draw_jit_grad_sed_flux(): @jax.jit def render(flux): - sed = jgal.SED(WAVE, flux) - gal = jgal.Gaussian(half_light_radius=0.5) * sed + sed = jgs.SED(WAVE, flux) + gal = jgs.Gaussian(half_light_radius=0.5) * sed psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) return ( ChromaticConvolution([gal, psf]) @@ -287,125 +109,60 @@ def render(flux): .array.sum() ) - grad = jax.grad(render)(jnp.ones(256)) - grad_arr = jnp.asarray(grad) + flux = jnp.ones_like(WAVE) + result = render(flux) + grad = jax.grad(render)(flux) + mask_in = (WAVE >= 550.0) & (WAVE <= 750.0) - # Gradient is nonzero only at wavelengths that fall inside the bandpass - # (the 32 quadrature points lie in [550, 750] nm). - # - sum of all gradients ≈ total bandpass flux = 200 + assert float(result) == pytest.approx(200.0, rel=5e-2) assert float(grad.sum()) == pytest.approx(200.0, rel=5e-2) - # - max gradient is positive assert float(grad.max()) > 0.0 - # - indices corresponding to wavelengths outside bandpass have zero gradient - # ~420 nm → outside [550, 750] bandpass - idx_out = int((420 - 400) / (900 - 400) * 255) - assert float(grad[idx_out]) == pytest.approx(0.0, abs=1e-8) - # - indices well inside bandpass region have nonzero total contribution - mask_in = (WAVE >= 550) & (WAVE <= 750) - assert float(grad_arr[mask_in].sum()) > 0.0 + assert float(grad[10]) == pytest.approx(0.0, abs=1e-8) + assert float(jnp.asarray(grad)[mask_in].sum()) > 0.0 -def test_chromatic_convolution_linearity(): - """Doubling SED flux doubles image sum.""" - - @jax.jit - def render(flux): - sed = jgal.SED(WAVE, flux) - gal = jgal.Gaussian(half_light_radius=0.5) * sed - psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) - return ( - ChromaticConvolution([gal, psf]) - .drawImage(BP, scale=0.2, nx=64, ny=64, n_waves=32) - .array.sum() - ) - - s1 = float(render(jnp.ones(256))) - s2 = float(render(jnp.ones(256) * 2.0)) - assert s2 == pytest.approx(s1 * 2.0, rel=1e-4) +def test_chromatic_sum_not_generically_separable(): + sed_disk = jgs.SED(WAVE, 1.0 + 0.2 * (WAVE - 650.0) / 250.0) + sed_bulge = jgs.SED(WAVE, 1.0 - 0.1 * (WAVE - 650.0) / 250.0) + disk = jgs.Gaussian(half_light_radius=0.7) * sed_disk + bulge = jgs.Gaussian(half_light_radius=0.3) * sed_bulge + source = disk + bulge -# --------------------------------------------------------------------------- -# ChromaticSum tests -# --------------------------------------------------------------------------- + assert not source._separable -def test_chromatic_sum_flux(): - sed1 = flat_sed(1.0) - sed2 = flat_sed(2.0) - gal1 = jgal.Gaussian(half_light_radius=0.3) * sed1 - gal2 = jgal.Gaussian(half_light_radius=0.8) * sed2 - combined = gal1 + gal2 - img = combined.drawImage(BP, scale=0.2, nx=32, ny=32) - # Total flux = (1 + 2) × 200 = 600 - assert float(img.array.sum()) == pytest.approx(600.0, rel=1e-2) +def test_chromatic_sum_convolution_matches_split_components(): + sed_disk = jgs.SED(WAVE, 1.0 + 0.2 * (WAVE - 650.0) / 250.0) + sed_bulge = jgs.SED(WAVE, 1.0 - 0.1 * (WAVE - 650.0) / 250.0) + disk = jgs.Gaussian(half_light_radius=0.7) * sed_disk + bulge = jgs.Gaussian(half_light_radius=0.3) * sed_bulge + psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) + combined = ChromaticConvolution([disk + bulge, psf]).drawImage( + BP, scale=0.2, nx=32, ny=32, n_waves=16 + ) + split_disk = ChromaticConvolution([disk, psf]).drawImage( + BP, scale=0.2, nx=32, ny=32, n_waves=16 + ) + split_bulge = ChromaticConvolution([bulge, psf]).drawImage( + BP, scale=0.2, nx=32, ny=32, n_waves=16 + ) + split = split_disk.array + split_bulge.array -# --------------------------------------------------------------------------- -# Monkey-patch (GSObject * SED) tests -# --------------------------------------------------------------------------- + assert float(jnp.max(jnp.abs(combined.array - split))) == pytest.approx( + 0.0, abs=5e-4 + ) -def test_gsobject_mul_sed(): - sed = flat_sed() - gal = jgal.Gaussian(half_light_radius=0.5) - from jax_galsim.chromatic import Chromatic +def test_gsobject_mul_sed_returns_chromatic_transformation(): + result = jgs.Gaussian(half_light_radius=0.5) * flat_sed() - result = gal * sed - assert isinstance(result, Chromatic) + assert isinstance(result, SimpleChromaticTransformation) -def test_gsobject_mul_scalar(): - gal = jgal.Gaussian(half_light_radius=0.5, flux=1.0) +def test_gsobject_mul_scalar_still_scales_flux(): + gal = jgs.Gaussian(half_light_radius=0.5, flux=1.0) scaled = gal * 5.0 - assert float(scaled.flux) == pytest.approx(5.0) - -if __name__ == "__main__": - # Run all tests inline for quick check - import sys - - tests = [ - test_sed_evaluation, - test_sed_redshift, - test_sed_calculate_flux, - test_sed_pytree_roundtrip, - test_sed_arithmetic, - test_bandpass_evaluation, - test_bandpass_effective_wavelength, - test_bandpass_effective_wavelength_concrete, - test_bandpass_mul, - test_bandpass_pytree_roundtrip, - test_chromatic_construction, - test_chromatic_drawImage_flux, - test_chromatic_drawImage_narrow_bandpass, - test_chromatic_jit, - test_chromatic_grad, - test_chromatic_jit_recompile, - test_chromatic_atmosphere_evaluate, - test_chromatic_atmosphere_scaling, - test_chromatic_atmosphere_drawImage_flux, - test_chromatic_atmosphere_moffat, - test_chromatic_atmosphere_pytree, - test_chromatic_convolution_flux, - test_chromatic_convolution_jit, - test_chromatic_convolution_grad, - test_chromatic_convolution_linearity, - test_chromatic_sum_flux, - test_gsobject_mul_sed, - test_gsobject_mul_scalar, - ] - - failed = [] - for t in tests: - try: - t() - print(f" PASS {t.__name__}") - except Exception as e: - print(f" FAIL {t.__name__}: {e}") - failed.append(t.__name__) - - print() - print(f"{len(tests) - len(failed)}/{len(tests)} passed") - if failed: - print("Failed:", failed) - sys.exit(1) + assert float(scaled.flux) == pytest.approx(5.0) From 60d6a6df3807b516aef1ebde351946d2152d95c2 Mon Sep 17 00:00:00 2001 From: MaxRonce Date: Tue, 16 Jun 2026 15:03:09 +0200 Subject: [PATCH 09/12] Fix chromatic JAX test import formatting --- tests/jax/test_chromatic_jax.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/jax/test_chromatic_jax.py b/tests/jax/test_chromatic_jax.py index dcefc389..ec2cd569 100644 --- a/tests/jax/test_chromatic_jax.py +++ b/tests/jax/test_chromatic_jax.py @@ -22,7 +22,6 @@ SimpleChromaticTransformation, ) - WAVE = jnp.linspace(400.0, 900.0, 256) BP = jgs.Bandpass.tophat(550.0, 750.0) From c56234eead571ef7d268826b37b49345941c1e15 Mon Sep 17 00:00:00 2001 From: MaxRonce Date: Tue, 16 Jun 2026 15:48:17 +0200 Subject: [PATCH 10/12] Clean up chromatic imports and helper names --- jax_galsim/chromatic.py | 69 ++++++++++++++++------------------------- 1 file changed, 27 insertions(+), 42 deletions(-) diff --git a/jax_galsim/chromatic.py b/jax_galsim/chromatic.py index ac3bccd4..3f97aa71 100644 --- a/jax_galsim/chromatic.py +++ b/jax_galsim/chromatic.py @@ -50,9 +50,16 @@ import jax.numpy as jnp from jax.tree_util import register_pytree_node_class +from jax_galsim.box import Pixel from jax_galsim.core.utils import cast_to_float, implements +from jax_galsim.gaussian import Gaussian +from jax_galsim.gsobject import GSObject from jax_galsim.gsparams import GSParams +from jax_galsim.image import Image +from jax_galsim.moffat import Moffat from jax_galsim.position import PositionD +from jax_galsim.sed import SED +from jax_galsim.sum import Sum def _pixel_scale_from_kwargs(kwargs): @@ -80,8 +87,6 @@ def _setup_wavelength_from_bandpass(bandpass): def _make_setup_image(profile, kwargs): """Build draw target without running full ``drawImage`` setup when size is fixed.""" - from jax_galsim.image import Image - image = kwargs.get("image", None) if image is not None: return Image(image=image) @@ -105,30 +110,30 @@ def _make_setup_image(profile, kwargs): return profile.drawImage(setup_only=True, **kwargs) -def _next_pow2(n): - n = int(n) - out = 1 - while out < n: - out *= 2 - return out +def _next_pow2(number): + number = int(number) + power = 1 + while power < number: + power *= 2 + return power def _fix_fft_size_for_image(profile, image): """Pin FFT size for JIT-safe setup when output bounds are already fixed.""" - nrow, ncol = image.array.shape - n = max(128, _next_pow2(2 * max(nrow, ncol))) - n = min(n, profile.gsparams.maximum_fft_size) - return profile.withGSParams(minimum_fft_size=n, maximum_fft_size=n) + row_count, column_count = image.array.shape + fft_size = max(128, _next_pow2(2 * max(row_count, column_count))) + fft_size = min(fft_size, profile.gsparams.maximum_fft_size) + return profile.withGSParams(minimum_fft_size=fft_size, maximum_fft_size=fft_size) def _static_kcoords(kimage, wrap_size, pixel_scale): """Return k-space pixel centers as a JAX array with static shape.""" - nrow, ncol = kimage.array.shape - x = jnp.arange(ncol, dtype=float) - y = jnp.arange(nrow, dtype=float) - wrap_size // 2 - kx, ky = jnp.meshgrid(x, y) - dk = 2.0 * jnp.pi / (wrap_size * pixel_scale) - return jnp.stack([kx.ravel(), ky.ravel()], axis=-1) * dk + row_count, column_count = kimage.array.shape + x_index = jnp.arange(column_count, dtype=float) + y_index = jnp.arange(row_count, dtype=float) - wrap_size // 2 + kx_grid, ky_grid = jnp.meshgrid(x_index, y_index) + delta_k = 2.0 * jnp.pi / (wrap_size * pixel_scale) + return jnp.stack([kx_grid.ravel(), ky_grid.ravel()], axis=-1) * delta_k # --------------------------------------------------------------------------- @@ -160,8 +165,6 @@ def __init__(self, obj=None): self._base_obj = None return - from jax_galsim.gsobject import GSObject - if not isinstance(obj, GSObject): raise TypeError("ChromaticObject requires a GSObject.") self._base_obj = obj @@ -279,9 +282,8 @@ def _drawImage_separable(self, bandpass, n_waves, **kwargs): Compute total_flux = ∫ SED(λ)×BP(λ) dλ (may be a JAX traced value, e.g. DSPS output), multiply into k-values, IFFT. """ - from jax_galsim.box import Pixel + # Keep this local to avoid a circular import with jax_galsim.convolve. from jax_galsim.convolve import Convolve - from jax_galsim.image import Image wave_eff = _setup_wavelength_from_bandpass(bandpass) pixel_scale = _pixel_scale_from_kwargs(kwargs) @@ -367,9 +369,8 @@ def _drawImage_nonseparable(self, bandpass, n_waves, **kwargs): Mirrors the drawFFT pipeline of GSObject.drawImage exactly, split into a concrete setup phase and a traced computation phase. """ - from jax_galsim.box import Pixel + # Keep this local to avoid a circular import with jax_galsim.convolve. from jax_galsim.convolve import Convolve - from jax_galsim.image import Image wave_eff = _setup_wavelength_from_bandpass(bandpass) pixel_scale = _pixel_scale_from_kwargs(kwargs) @@ -452,8 +453,6 @@ def __radd__(self, other): @implements(_galsim.ChromaticObject.__mul__) def __mul__(self, other): - from jax_galsim.sed import SED - if isinstance(other, SED): base_obj = getattr(self, "_base_obj", None) if base_obj is None: @@ -533,8 +532,6 @@ def atRedshift(self, redshift): @implements(_galsim.ChromaticSum.evaluateAtWavelength) def evaluateAtWavelength(self, wave): - from jax_galsim.sum import Sum - return Sum([o.evaluateAtWavelength(wave) for o in self.obj_list]) @implements( @@ -792,13 +789,9 @@ def evaluateAtWavelength(self, wave): fwhm = self._fwhm_ref * (wave / self._lam_ref) ** self._alpha if self._profile == "gaussian": - from jax_galsim.gaussian import Gaussian - return Gaussian(fwhm=fwhm, flux=1.0, gsparams=self._gsparams) elif self._profile == "moffat": - from jax_galsim.moffat import Moffat - return Moffat( fwhm=fwhm, beta=self._moffat_beta, @@ -910,8 +903,6 @@ class ChromaticConvolution(ChromaticObject): _separable = False def __init__(self, *args, **kwargs): - from jax_galsim.gsobject import GSObject - if len(args) == 0: raise TypeError("At least one object must be provided.") if len(args) == 1 and isinstance(args[0], (list, tuple)): @@ -939,8 +930,6 @@ def __init__(self, *args, **kwargs): processed.extend(obj.obj_list) continue if isinstance(obj, GSObject): - from jax_galsim.sed import SED - wave_stub = jnp.array([100.0, 2000.0]) flat_sed = SED(wave_stub, jnp.ones(2)) processed.append(SimpleChromaticTransformation(obj, flat_sed)) @@ -972,6 +961,7 @@ def atRedshift(self, redshift): @implements(_galsim.ChromaticConvolution.evaluateAtWavelength) def evaluateAtWavelength(self, wave): """Return the convolved monochromatic profile at *wave* (nm).""" + # Keep this local to avoid a circular import with jax_galsim.convolve. from jax_galsim.convolve import Convolve return Convolve([o.evaluateAtWavelength(wave) for o in self.obj_list]) @@ -1001,9 +991,8 @@ def drawImage(self, bandpass, n_waves=64, **kwargs): All-separable case falls back to a single monochromatic draw. """ - from jax_galsim.box import Pixel + # Keep this local to avoid a circular import with jax_galsim.convolve. from jax_galsim.convolve import Convolve - from jax_galsim.image import Image sep_objs = [o for o in self.obj_list if o._separable] nonsep_objs = [o for o in self.obj_list if not o._separable] @@ -1147,8 +1136,6 @@ def __repr__(self): ) def _gsobject_mul_sed(self, other): """Allow ``gsobject * sed``.""" - from jax_galsim.sed import SED - if isinstance(other, SED): return SimpleChromaticTransformation(self, other) # Fall through to original implementation (flux scaling) @@ -1165,8 +1152,6 @@ def _gsobject_rmul_sed(self, other): # Apply the patch once at import time def _patch_gsobject(): - from jax_galsim.gsobject import GSObject - GSObject.__mul__ = _gsobject_mul_sed GSObject.__rmul__ = _gsobject_rmul_sed From ed243f91f8a3bf368d206d97c345e2d3963fff48 Mon Sep 17 00:00:00 2001 From: MaxRonce Date: Tue, 16 Jun 2026 15:56:28 +0200 Subject: [PATCH 11/12] Move chromatic API notes into implements metadata --- jax_galsim/bandpass.py | 1 - jax_galsim/chromatic.py | 193 +--------------------------------------- 2 files changed, 4 insertions(+), 190 deletions(-) diff --git a/jax_galsim/bandpass.py b/jax_galsim/bandpass.py index 5cc62ca2..e981e8ad 100644 --- a/jax_galsim/bandpass.py +++ b/jax_galsim/bandpass.py @@ -157,7 +157,6 @@ def truncate(self, blue_limit=None, red_limit=None): ), ) def tophat(cls, blue_limit, red_limit, n_wave=100): - """Uniform throughput = 1 between blue_limit and red_limit.""" wave = jnp.linspace(blue_limit, red_limit, n_wave) return cls(wave, jnp.ones(n_wave)) diff --git a/jax_galsim/chromatic.py b/jax_galsim/chromatic.py index 3f97aa71..8cd98c0e 100644 --- a/jax_galsim/chromatic.py +++ b/jax_galsim/chromatic.py @@ -150,13 +150,6 @@ def _static_kcoords(kimage, wrap_size, pixel_scale): """, ) class ChromaticObject: - """Base class for wavelength-dependent profiles. - - Subclasses must override :meth:`evaluateAtWavelength`. The default - :meth:`drawImage` uses a k-space trapezoid integration that works for - any subclass; separable subclasses override it with a faster path. - """ - #: Set True in separable subclasses. _separable: bool = False @@ -171,9 +164,11 @@ def __init__(self, obj=None): self._separable = True @property - @implements(getattr(_galsim.ChromaticObject, "separable", None)) + @implements( + getattr(_galsim.ChromaticObject, "separable", None), + lax_description="JAX-GalSim-specific separability flag for chromatic drawing.", + ) def separable(self): - """True if the profile factors as g(x,y) × h(λ).""" return self._separable @property @@ -190,21 +185,6 @@ def gsparams(self): @implements(_galsim.ChromaticObject.evaluateAtWavelength) def evaluateAtWavelength(self, wave): - """Return the monochromatic GSObject at wavelength *wave* (nm). - - This method must be JAX-traceable: all internal computations - should use ``jax.numpy``, and the returned GSObject's parameters - may be abstract JAX tracers. - - Parameters - ---------- - wave : float or jax scalar - Wavelength in nm. - - Returns - ------- - GSObject - """ if getattr(self, "_base_obj", None) is not None: return self._base_obj raise NotImplementedError( @@ -242,23 +222,6 @@ def atRedshift(self, redshift): lax_description="JAX-GalSim supports FFT drawing through array-backed Bandpass objects.", ) def drawImage(self, bandpass, n_waves=64, **kwargs): - """Draw the bandpass-integrated image. - - Parameters - ---------- - bandpass : Bandpass - Observing bandpass. - n_waves : int - Number of wavelength samples for numerical integration. - Must be a **static** integer (fixed at JIT compile time). - **kwargs - Forwarded to the underlying ``GSObject.drawImage`` calls. - Typical keys: ``nx``, ``ny``, ``scale``, ``method``, etc. - - Returns - ------- - Image - """ if self._separable: return self._drawImage_separable(bandpass, n_waves, **kwargs) return self._drawImage_nonseparable(bandpass, n_waves, **kwargs) @@ -495,16 +458,6 @@ def _sed_value(self, wave): lax_description="JAX-GalSim treats ChromaticSum as non-separable, matching GalSim semantics.", ) class ChromaticSum(ChromaticObject): - """Sum of two or more chromatic profiles. - - The combined SED is the sum of all component SEDs. Drawing - evaluates each component at every wavelength and sums. - - Parameters - ---------- - obj_list : list of ChromaticObject - """ - _separable = False def __init__(self, *args): @@ -565,41 +518,6 @@ def drawImage(self, bandpass, n_waves=64, **kwargs): ) @register_pytree_node_class class SimpleChromaticTransformation(ChromaticObject): - """Separable chromatic profile: a GSObject multiplied by an SED. - - The spatial profile is fixed; wavelength dependence enters only - through the SED flux scaling. - - ``I(x, y, λ) = g(x, y) × SED(λ)`` - - Drawing a ``SimpleChromaticTransformation`` through a bandpass reduces to a single - monochromatic draw at the effective wavelength with total flux - ``∫ SED(λ) × BP(λ) dλ``. - - Parameters - ---------- - obj : GSObject - Normalised spatial profile (flux = 1 by convention, though - any flux is allowed and will be multiplied by the SED). - sed : SED - Spectral energy distribution. - - Examples - -------- - :: - - >>> from jax_galsim import Gaussian - >>> from jax_galsim.sed import SED - >>> from jax_galsim.bandpass import Bandpass - >>> import jax.numpy as jnp - - >>> wave = jnp.linspace(500, 900, 256) - >>> sed = SED(wave, jnp.ones(256)) - >>> bp = Bandpass.tophat(550, 750) - >>> gal = Gaussian(half_light_radius=0.5) * sed - >>> img = gal.drawImage(bp, scale=0.2, nx=64, ny=64) - """ - _separable = True def __init__(self, obj, sed): @@ -622,7 +540,6 @@ def atRedshift(self, redshift): @implements(_galsim.SimpleChromaticTransformation.evaluateAtWavelength) def evaluateAtWavelength(self, wave): - """Return the GSObject scaled to SED(wave).""" return self.obj.withScaledFlux(self.sed(wave)) def _sed_value(self, wave): @@ -667,43 +584,6 @@ def __repr__(self): """, ) class ChromaticAtmosphere(ChromaticObject): - """Atmospheric PSF with a power-law wavelength-dependent FWHM. - - The PSF profile at wavelength λ is a Gaussian (or Moffat) with: - - FWHM(λ) = fwhm_ref × (λ / lam_ref)^alpha - - For Kolmogorov turbulence, the expected scaling is α ≈ −0.2. - - This profile carries a **flat (dimensionless) SED**: ``SED(λ) = 1``. - The physical SED is typically attached to the galaxy component via - :class:`SimpleChromaticTransformation`, and passed to - :class:`ChromaticConvolution`. - - Parameters - ---------- - fwhm_ref : float - FWHM in arcseconds at the reference wavelength. - lam_ref : float - Reference wavelength in nm. - alpha : float, optional - Power-law index. Default −0.2 (Kolmogorov). - profile : {'gaussian', 'moffat'} - Profile type. Default ``'gaussian'``. - moffat_beta : float, optional - Moffat β parameter (only used when ``profile='moffat'``). - Default 4.765 (typical for atmospheric seeing). - gsparams : GSParams, optional - - Examples - -------- - :: - - >>> from jax_galsim.chromatic import ChromaticAtmosphere - >>> psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) - >>> prof = psf.evaluateAtWavelength(550.0) # Gaussian at 550 nm - """ - _separable = False def __init__( @@ -781,11 +661,6 @@ def atRedshift(self, redshift): @implements(_galsim.ChromaticAtmosphere.evaluateAtWavelength) def evaluateAtWavelength(self, wave): - """Return the PSF profile (unit flux) at wavelength *wave* (nm). - - FWHM is scaled as ``fwhm_ref × (wave / lam_ref)^alpha``. - When *wave* is a JAX tracer (inside vmap/jit), fwhm is also traced. - """ fwhm = self._fwhm_ref * (wave / self._lam_ref) ** self._alpha if self._profile == "gaussian": @@ -858,48 +733,6 @@ def __repr__(self): """, ) class ChromaticConvolution(ChromaticObject): - """Convolution of multiple chromatic profiles. - - Computes the bandpass-integrated image of the convolution: - - K_eff(k) = ∫ ∏_i K_i(k, λ) × BP(λ) dλ - - where each ``K_i(k, λ)`` is the k-space value of the i-th component - at wavelength λ. - - **Optimised case**: when all components except one are separable - (e.g. galaxy × SED convolved with a chromatic PSF), the separable - profiles are extracted and convolved with the wavelength-integrated - effective PSF. This avoids multiplying the same galaxy k-image at - every wavelength sample. - - Parameters - ---------- - obj_list : list of ChromaticObject or GSObject - Components to convolve. Plain ``GSObject`` instances are wrapped - automatically in a flat-SED :class:`SimpleChromaticTransformation`. - - Examples - -------- - :: - - >>> from jax_galsim import Gaussian, Convolve - >>> from jax_galsim.chromatic import ChromaticAtmosphere, ChromaticConvolution - >>> from jax_galsim.sed import SED - >>> from jax_galsim.bandpass import Bandpass - >>> import jax.numpy as jnp - - >>> wave = jnp.linspace(300, 1100, 512) - >>> flux = jnp.exp(-0.5 * ((wave - 700) / 150) ** 2) # Gaussian SED - >>> sed = SED(wave, flux) - >>> bp = Bandpass.tophat(550, 800) - - >>> gal = Gaussian(half_light_radius=0.5) * sed - >>> psf = ChromaticAtmosphere(fwhm_ref=0.7, lam_ref=700.0, alpha=-0.2) - >>> final = ChromaticConvolution([gal, psf]) - >>> img = final.drawImage(bp, scale=0.2, nx=64, ny=64) - """ - _separable = False def __init__(self, *args, **kwargs): @@ -960,7 +793,6 @@ def atRedshift(self, redshift): @implements(_galsim.ChromaticConvolution.evaluateAtWavelength) def evaluateAtWavelength(self, wave): - """Return the convolved monochromatic profile at *wave* (nm).""" # Keep this local to avoid a circular import with jax_galsim.convolve. from jax_galsim.convolve import Convolve @@ -975,22 +807,6 @@ def evaluateAtWavelength(self, wave): lax_description="JAX-GalSim supports FFT drawing with static wavelength grid size.", ) def drawImage(self, bandpass, n_waves=64, **kwargs): - """Draw the integrated image, exploiting separability where possible. - - Algorithm: - - 1. Separate components into *separable* (galaxy × SED) and - *non-separable* (chromatic PSF). - 2. Build the combined weight function: - ``w(λ) = ∏_sep SED_i(λ) × BP(λ)`` - 3. In k-space, integrate non-separable components: - ``K_nonsep_eff(k) = ∫ ∏_nonsep K_i(k,λ) × w(λ) dλ`` - 4. Multiply by separable component k-values (λ-independent after - extracting their SED into the weight). - 5. Include pixel convolution, IFFT to real space. - - All-separable case falls back to a single monochromatic draw. - """ # Keep this local to avoid a circular import with jax_galsim.convolve. from jax_galsim.convolve import Convolve @@ -1135,7 +951,6 @@ def __repr__(self): lax_description="Also accepts array-backed JAX-GalSim SED objects and returns a SimpleChromaticTransformation.", ) def _gsobject_mul_sed(self, other): - """Allow ``gsobject * sed``.""" if isinstance(other, SED): return SimpleChromaticTransformation(self, other) # Fall through to original implementation (flux scaling) From 4ffdfe60c3dc60a407fe7f5f6a57374d27ad44b2 Mon Sep 17 00:00:00 2001 From: MaxRonce Date: Thu, 18 Jun 2026 13:24:31 +0200 Subject: [PATCH 12/12] Update pytest hooks for pytest 9 --- tests/conftest.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 6b0082a9..2d8b080b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -43,7 +43,7 @@ def _check_pickle(*args, **kwargs): galsim.utilities.check_pickle = _check_pickle -def pytest_ignore_collect(collection_path, path, config): +def pytest_ignore_collect(collection_path, config): """This hook will skip collecting tests that are not enabled in the enabled_tests.yaml file. @@ -130,7 +130,7 @@ def _convert_galsim_to_jax_galsim(obj): return obj -def pytest_pycollect_makemodule(module_path, path, parent): +def pytest_pycollect_makemodule(module_path, parent): """This hook is tasked with overriding the galsim import at the top of each test file. Replaces it by jax-galsim. """