Implementation of k-space convolution

[EiffL]

This PR adds the ability to perform Fourier-based PSF convolutions. Most of this code comes from @jecampagne.

[EiffL]

@ismael-mendoza could you check all the add/sub/mul things in image.py?

[ismael-mendoza]

Somethings that are missing so far:

• Adding the pixel convolution (see TODO)
• no wrapping

[ismael-mendoza]

just rebased with the pixel branch and uncommented the TODO, it seems like pixels convolutions are working and agree with Galsim for the example here:

# Define a galsim galaxy as the sum of two objects
obj1 = galsim.Gaussian(half_light_radius=1.)
obj2 = galsim.Exponential(half_light_radius=0.5)

# Rescale the flux of one object
obj2 = obj2.withFlux(0.4)

# Sum the two components of my galaxy
gal = obj1 + obj2

psf = galsim.Gaussian(half_light_radius=2.)

obs_gal = Convolution(gal, psf)

[EiffL]

Awesome :-D

[EiffL]

Well, so, for some reason there appears to be an offset in the fft drawing method, also I still am not sure whether the rotations work correctly, this should be tested

[EiffL]

So, I did 2 additional things:

• I changed the way that drawImage works so that (even if it's not recommended to do so) the old expected galsim behaviour of being able to update an image will work:
  

  JAX-GalSim/jax_galsim/gsobject.py

  Line 685 in 477c2ce

  image_in._array = image._array

• I updated the GalSim branch we use for testing to this branch: https://github.com/GalSim-developers/GalSim/tree/fix_tests
  where we can make edits like making the np.testing function array-friendly when desired, with high likelihood that the GalSim Devs will be ok with merging these small modifs. We should use that mechanism to fix tests so that they work for us, when these fixes are really minimal

[EiffL]

One quick note before leaving for the airport, the reason the Fourier transform is behaving in a funny way is that I found a bug yesterday here:

JAX-GalSim/jax_galsim/gsobject.py

Line 812 in 477c2ce

jnp.fft.irfft2(kimg_shift, breal.numpyShape())

[EiffL]

ok, so with the latest fixes, the drawImage (method='fft') wotks again, but I don't think it works as well as what we had the other day (which was actually buggy as it was using arrays of the wrong size, but thanks to JAX not checking for bounds, ended up working nevertheless).

We currently have only one test that fails, but it's a pretty important one which tests whether shear, rotations, and shifts work properly. For sure it should get sorted before we merge.

That's probably as far as I can go on this project for now unfortunately :-(

[ismael-mendoza]

@EiffL - thanks for all your work on this. Could you elaborate a bit on why you think the fft method is not working as well as before?

[EiffL]

If you plot residuals vs galsim there are some ugly aliasing patterns

[EiffL]

Ok :-) I think I fixed all the main problems. All the tests pass now, and I also added a few ones to do some basic checks against the reference GalSim in temporary_tests.py
It's pretty cool, you can test jax_galsim against galsim with this sort of scheme:

def func(galsim):
      gal =  galsim.Gaussian(flux=1, sigma=1).shear(e1=0.1, e2=0.)
      pixel = galsim.Pixel(scale=0.2, flux=1.0)
      conv = galsim.Convolve([gal, pixel])
      return conv.drawImage(
          nx=32, ny=32, scale=0.2, method="sb").array

np.testing.assert_array_almost_equal(
      func(ref_galsim),  func(jax_galsim),
      5,
      err_msg="Using GSObject Convolve([psf,pixel]) disagrees with expected result")

So, anyway, I think this PR is fully ready for review @ismael-mendoza

[ismael-mendoza]

Amazing, thanks @EiffL ! Will review soon, very interested in hearing about the fix at some point.

[ismael-mendoza]

It looks very good! All the tests are similarly passing for me. Although I do have some questions below but not sure if it stop merging.

@jecampagne feel free to chime in if you are able to since a lot of the code was borrowed from you.

One last outstanding thing if we should add the draw methods for the Exponential here or in a different PR (EDIT: looking at wrong file)

[ismael-mendoza]

@EiffL did you update something in Galsim?

[EiffL]

Hi @ismael-mendoza , yes I updated to a specific galsim branch where we can push updates for the tests:
https://github.com/GalSim-developers/GalSim/tree/fix_tests

[jecampagne]

Hi @EiffL @ismael-mendoza you did an amazing progress and if I've contributed to this development then I am pleased having not waste your time.

[EiffL]

Thanks @jecampagne :-) I've mostly copy/pasted your code over and did a bunch of reformatting, I credited you in the corresponding commits so that your code contributions will be tracked in the project history :-)

[ismael-mendoza]

LGTM! I think it's ready to merge :D

[jecampagne]

Looks ok me too.

[EiffL]

Wouhou! Merging :-D