Implement affine image transformation

[EiffL]

One of the key features of GalSim is the ability to apply affine transformations to a GSObject, this is in particular needed for adding shear to a galaxy image.

In GalSim, these things are defined here for instance: https://github.com/GalSim-developers/GalSim/blob/013292ad45b6ef715acf70976e94fe21f624ee4c/galsim/transform.py#L98

And you can see here what happens when a GSObject is sheared: https://github.com/GalSim-developers/GalSim/blob/013292ad45b6ef715acf70976e94fe21f624ee4c/galsim/gsobject.py#L863 What happens is that an Affine Transformation is defined, and applied on the image.

The goal of this issue is to discuss the implementation of similar transformations within GalFlow using TensorFlow, so that the transformation can be differentiated with respect to the input Jacobian. I've found these image tools in TensorFlow Addons: https://www.tensorflow.org/addons/api_docs/python/tfa/image/transform And I think this may be exactly what we need.

[EiffL]

I made a minor correction, and it seems to be working nicely now :-)

Up to one small detail, I'm not 100% sure how to correctly compute the shift to ensure that the galaxy remains are the center of the image when we apply a shear. I came up with a formula, but it's not 100% full proof.... Although it seems to be working quite nicely as it is :-D

There is unfortunately one tiny bit of sad news. From the documentation of transform here https://www.tensorflow.org/addons/api_docs/python/tfa/image/transform

Note that gradients are not backpropagated into transformation parameters.

So... it's possible that we can't get gradients with this function for some reason :-(

[EiffL]

yep, confirmed... no grads with respect to transformation parameters :-( I've added a demo code in the notebook.

Soooooo I guess we'll have to recode this affine transformation function ^^'

[EiffL]

one small discovery, we actually dont need TFA, the function we need is a TF function in 2.4:

https://www.tensorflow.org/api_docs/python/tf/raw_ops/ImageProjectiveTransformV3

[EiffL]

And the source for these ops is here: https://github.com/tensorflow/tensorflow/blob/dec8e0b11f4f87693b67e125e67dfbc68d26c205/tensorflow/python/ops/image_ops.py#L277 https://github.com/tensorflow/tensorflow/blob/dec8e0b11f4f87693b67e125e67dfbc68d26c205/tensorflow/core/kernels/image/image_ops.cc#L169 And indeed it doesn't implement differentiation with respect to the transform parameters.... that would have been too easy

[EiffL]

Ok, so tf addons didn't pan out, but we are saved by TensorFlow Graphics, which implements similar functions, in a different way, allowing for differentiability with respect to the transformation matrix :-D Details are here: https://www.tensorflow.org/graphics/api_docs/python/tfg/image/transformer/perspective_transform

CAREFUL: this requires tfa-nightly, tfg-nightly, and TF 2.4

But I can successfully shear a model by gradient descent to recover an input shear. Here is the unsheared model: And after a few iterations of gradient descent:

so.... great success!