Adds jittable Image

[EiffL]

This PR should be rebased on main once #15 is merged, it adds an implementation of galsim.Image.

There are some questions about this though. galsim.Image are mutable objects, which really does not play nicely with JAX, it is simply not possible to assign values to array in jax without making a copy. We can follow the galsim API, but the behavior will be different, it won't be possible to assign values to a view of a larger image.

In particular, we cannot do the following in Jax-GalSim:

gal = galsim.Gaussian(sigma=1)
im = galsim.Image(nx=128, ny=128, scale=0.1)

imview = im.view()

gal.drawImage(imview)

This would not update the original im.

This will become even sharper of a question when we go to the next PR, which implements drawing functions.

Already here, there is a question of whether we want to make all Image immutable explicitly, whether we want to allow this sort of things:

image = galsim.Image(nx=128, ny=128, scale=0.1)
image.wcs = galsim.PixelScale(0.2)
image.setCenter(0.0, 0.)

these are examples of bad OO patterns that do not play nicely at all with the JAX functional spirit, but if we remove them, we break the GalSim API....

[ismael-mendoza]

preliminary comments

[ismael-mendoza]

Seems like there are several places where jitting will break

[ismael-mendoza]

Facing this problem with the _make_empty function where I believe making arrays out of arguments is not jittable/vmappable. Is there any workaround? The alternative is to enforce the only option for creating an image is to use the array but that will really break the API.

I guess we need to discuss how we want to allow the Image object to be vectorized.

[ismael-mendoza]

ToDo:

• cleanup explicit init so it's a class method
• check tree methods are correct
• remove possibility of empty image?
• fix test_image_jax.py and check it passes
• use same dtype as array if passed in if dtype is None

Questions

• Still don't understand why empty images are bad
• What should default dtype be?
• What was the point of "enable x64" again?
• Test now uses explicit init? (just like a user would)
• What to do about error messages about 'int' / 'float' type not available? (if x64 off)
• valid dtypes? jnp.float32?

[ismael-mendoza]

Tests are passing, see if CI runs ?