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

[MaxRonce]

Should go with the issue #251

python -m pytest -c /home/maxime/src/galsim/JAX-GalSim_chromaticPSF/pyproject.toml tests/GalSim/tests/test_sed.py tests/GalSim/tests/test_bandpass.py tests/GalSim/tests/test_chromatic.py -q

Tests seams to pass and results are similar to original galsim

Element by element testing and benchmark was done using this script :

bench_dsps_chromatic.py

The 10e-5 error on the bottom left plot seams to be only numerical bias due to the NWAVES for the fft and not an actuel implementation bias

Some class diagram before and after :

Regular Galsim

Jax Galsim

[MaxRonce]

Know issue

ChromaticSum incorrectly treated as separable when summing separable components with different spatial profiles

ChromaticSum currently marks itself as separable when all child components are separable:

self._separable = all(o._separable for o in self.obj_list)

This is not generally correct. A sum of separable chromatic objects is not necessarily separable.

Ex : physical bulge + disk model:

disk(x, y) * SED_disk(lambda) + bulge(x, y) * SED_bulge(lambda)

is not separable in general, even though each component is individually separable. It is separable only in special cases, e.g. if the spatial profiles are identical or if the SEDs are identical up to a scalar factor.

Observed behavior

When building a model like:

disk = disk_profile * sed_disk
bulge = bulge_profile * sed_bulge
source = disk + bulge
model = ChromaticConvolution([source, psf])
image = model.drawImage(bandpass, ...)

source becomes a ChromaticSum.

Because both disk and bulge are separable, ChromaticSum sets:

self._separable = True

Then ChromaticConvolution.drawImage() takes the optimized separable path and calls:

o._sed_value(wave)

on the ChromaticSum.

But ChromaticSum does not implement _sed_value(), so the base implementation raises:

NotImplementedError

Expected behavior

A sum of separable chromatic components should only be marked separable if the sum itself can be written as:

spatial_profile(x, y) * spectral_weight(lambda)

A generic bulge + disk model cannot be represented this way.

Fixes ??

Always treat ChromaticSum as non-separable:

class ChromaticSum(ChromaticObject):
    _separable = False

or inside __init__:

self._separable = False

This is simple and robust, but may lose optimization opportunities for rare truly separable sums.

More precise fix

Only mark ChromaticSum as separable when the sum can actually be factored into:

g(x, y) * h(lambda)

That would require implementing compatible _sed_value() and _static_spatial_profile() behavior only for cases where the children share an equivalent spatial profile or equivalent SED structure.

However, for a generic bulge + disk model with distinct spatial profiles and distinct SEDs, ChromaticSum should remain non-separable.

Current fix

For testing purpose I used this solution :

(disk + bulge) (*) psf = disk (*) psf + bulge (*) psf

In code:

disk_image = ChromaticConvolution([disk, psf]).drawImage(bandpass, ...)
bulge_image = ChromaticConvolution([bulge, psf]).drawImage(bandpass, ...)
image = disk_image + bulge_image

This is mathematically correct, but it would be preferable for ChromaticSum / ChromaticConvolution to handle the generic case properly