Add DES_PSFEx (PSFEx model reader) to jax_galsim.des

[AdamField118]

Add DES_PSFEx to jax_galsim.des

Hi! Nice to meet you, this is my first PR ever, happy to adjust anything to match your conventions.

This ports galsim.des.DES_PSFEx so a PSFEx model (*_psfcat.psf) can be read and evaluated entirely within JAX-GalSim. Right now, if you want an empirical PSF you have to drop back to reference GalSim just to build it; this closes that gap. I hit it while benchmarking a weak lensing pipeline's dataset generation before and after porting to jax-galsim. DES_PSFEx was the last thing keeping regular GalSim as a dependency.

What it does

• read() is host-side FITS I/O, unchanged from GalSim (parses the PSFEx header + PCA basis cube).
• getPSFArray() is the per-position evaluation of the interpolated PSF. It's reimplemented in JAX so it's jit-able, vmap-able over image positions, and differentiable w.r.t. the image position. The one non-obvious change: GalSim builds the [1, x, x^2, ...] powers with an in-place np.empty loop, which tracing forbids, so I use a cumulative product with the same recurrence.
• getPSF() wraps the array in an InterpolatedImage exactly as GalSim does (Lanczos(3), scale = PSF_SAMP) and applies the WCS if one was provided.

Scope (intentionally narrow)

Just DES_PSFEx, not DES_Shapelet, and not the config-framework registration (des_psfex input type / DES_PSFEx object type), since JAX-GalSim doesn't have a config system. That's noted in the lax_description. Happy to follow up with more in a separate PR if you'd like it.

Testing

tests/jax/test_des_psfex_jax.py compares against reference GalSim using the existing DECam_00154912_12_psfcat.psf in the tests/GalSim submodule:

• getPSFArray and the drawn effective-PSF image match GalSim to ~single precision (the PSFEx bases are float32). Drawn with method='no_pixel', since PSFEx PSFs already include the pixel.
• getPSFArray is verified jit-able, vmap-able across a batch of positions, and differentiable in position (finite gradient).

ruff check and ruff format are clean.

Open questions

• Happy to move the module, rename things, or adjust the test tolerances to match your house style.
• Let me know if you'd prefer the config registration included after all, or DES_Shapelet alongside it.

[beckermr]

This class needs to be registered as a pytree and we need think about what kinds of autodiff we'd allow.

[AdamField118]

I'll register it as a pytree. My idea of the split: basis, the polynomial zero/scale constants, and wcs as traced children (keeping it consistent with how other jax_galsim objects treat numeric params), with fit_order/fit_size as static aux_data since they set the loop unrolling and array shapes. tree_unflatten will rebuild via object.__new__ to skip the file read, matching what CelestialCoord/Image do.

On autodiff scope: I think the physically meaningful gradient is w.r.t. image_pos (where on the focal plane the PSF is evaluated), which already works and comes through the method argument rather than the pytree. Differentiating w.r.t. basis/wcs/calibration constants falls out structurally once they're leaves, but I don't see a concrete use case for those. I'd be happy to lock any of them into aux_data if you'd rather keep the differentiable surface narrow. What's your preference?