Drawing image using k-spectrum Inverse FFT

[jecampagne]

Hello,
I have cooked a JAX-Galsim drawing using FFT spectrum implementing the GalSim mechanism with a pythonic rough version of the wrapping process in this context. I have set up demo nb here
https://colab.research.google.com/drive/1Up236BZlIGBXTqJ-RUZOvklyR45ZROmp?usp=sharing

The demo use starts with a simple gal= galsim.Gaussian(sigma=1) and compare the image production using real_space and fft, then the same exo is done using GalSim code.
Then the demo continues with the addition of two profiles, and once again compare real_space and fft image production as well as the GalSim production.

Now, the code I have done is a very moderate JAX version of the C++ GalSim code, and need less to say that the wrap_hermx_cols_pair and _wrap methods in image.pyshould be reimplemented using more JAX version. But at least the algo seems to work.

Final warning: In the GalSim code of drawImage when using FFT operates a Convolution with a Pixel Shape before, and this covolution is done in real space due to the sharp edge of the Pixel Box profile. This Convolution is not yet implemented in my current version.

[jecampagne]

Hi,
I Have update the image.py (wrap_hermx_cols_pair and _wrap methods) to use jax.lax control flux functions that implicitly performs a JIT compilation. The demo notebook is using this new implementation even if the API has not been changed

https://colab.research.google.com/drive/1Up236BZlIGBXTqJ-RUZOvklyR45ZROmp?usp=sharing#scrollTo=i--xWIuFXdn0