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

[beckermr]

Let's not repeat doc strings from upstream and instead use the implements decorator, with any caveats put into the lax_description keyword.

[beckermr]

Here self._throughput could be a traced array in which case this precompute step won't work.

[beckermr]

This interpolation is linear. There is a slightly better akima interpolant that is fast to compute on-the-fly. See jax_galsim.core.interpolate.

[beckermr]

I took a first pass at review. There are a few big structural things that need addressing before we can review it again.

1. All classes, methods, and functions need to be wrapped with the @implements decorator. Any changes specific to JAX need to go in the lax_description keyword.
2. Please add the new objects to the API tests in test_api.py. From my read of the code, many of those API tests will fail.
3. The Bandpass class cannot cache the effective wavelength on init as a python float since the throughput could be traced.
4. The JAX tests for Bandpass don't appear the test any gradients or JIT with respect to throughput and instead always use a tophat. I think this is possibly why the init stuff from 3 works in the tests, but won't work in general.
5. If an array/computation is meant to only be used at compile time, then we should use numpy and wrap it in an ensure_compile_time_eval block.

[beckermr]

pre-commit.ci autofix

[codspeed]

Merging this PR will improve performance by 34.71%

⚠️ Unknown Walltime execution environment detected

Using the Walltime instrument on standard Hosted Runners will lead to inconsistent data.

For the most accurate results, we recommend using CodSpeed Macro Runners: bare-metal machines fine-tuned for performance measurement consistency.

⚠️ Different runtime environments detected

Some benchmarks with significant performance changes were compared across different runtime environments,
which may affect the accuracy of the results.

Open the report in CodSpeed to investigate

⚡ 1 improved benchmark
✅ 35 untouched benchmarks

Performance Changes

	Mode	Benchmark	BASE	HEAD	Efficiency
⚡	Simulation	test_benchmark_invert_ab_noraise[run]	1,084.2 µs	804.9 µs	+34.71%

Tip

Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing MaxRonce:chromatic_psf (539826b) with main (4c1bf06)}

[MaxRonce]

I think if made the change you requsted.
We also did a first pass code quality review with @sfarrens