Implement a model fitting shape measurement method

[EiffL]

One of the first things we want to demonstrate is that we can reproduce the equivalent of ngmix in this framework, i.e. using a model fitting method.

This will allow us to make 2 points:

• Classical techniques like model fitting can be made autodiffable
• This metacal framework can work without any ML in it

And bonus points for using a cool implicit differentiation trick through the optimization algorithm ;-)

We have an example of that here: https://colab.research.google.com/drive/1yF1lMP7tDdyNd0HKdNEbWozkjtoiB1I-?usp=sharing Now, we can do the same thing in TensorFlow (I think) using custom gradients.

So, the goal of this issue is to implement an equivalent of the ngmix fitter algorithm. Axel pointed to this ngmix issue which might be a good place to have a reference example: https://github.com/esheldon/ngmix/issues/72

[b-remy]

I made a first experiment where I fit a Gaussian to a galaxy image and compute the ellipticity from its covariance matrix: https://github.com/CosmoStat/autometacal/blob/model_fitting/notebooks/EllipticityModelFitting.ipynb

Everything is differentiable, and the fit is done within an implicit layer as @EiffL did in his example above.

However the code is still in JAX because I couldn't find any vector-Jacobian product implementation in TensorFlow (only Jacobean-vector product here...), that is necessary for implicit layers.

[EiffL]

This is amazing!!!

[b-remy]

For the TensorFlow implementation, maybe this will be helpful: https://github.com/tensorflow/tensorflow/blob/0b9c9497f1c7f924f3fbeffcf9426d5174a68600/tensorflow/python/eager/backprop.py#L510

[b-remy]

Small increment here :-) I managed to reproduce @EiffL's JAX notebook that fits a gaussian to a gaussian, but in TensorFlow! Using custom gradient as you mentioned: https://colab.research.google.com/drive/1p8iZ2PsjqauIuaNR2o-uwcOB4yf8Cpl0?usp=sharing

I have an ongoing notebook where I use this method to estimate the ellipticity of a galaxy, but still working on:

• [x] implement gradient computation with @tf.custom_gradient
• [x] using XLA to speed the fixed point computation (necessary for gradient descent on my laptop)
• [x] write the gaussian -> ellipticity conversion method in TensorFlow.
• [ ] optional: implement an better optimization step (BFGS/LBFGS) to be able to fit a more complex model

[EiffL]

ahah! AWESOME! sorry it took me 14 hours to check it out! This is super cool :-) So we may be able to use good old TF after all :-D

So I see that you explicitly compute the jacobian, if fstar remains low dimensional, it should work ok :-) But reading the code it may look like you don't actually need the full jacobian, and instead just need to evaluate vector jacobian products? And vjps is just normal reverse mode autodiff, so I think_ you mght be able to define a nested function with custom gradients, inside the current function? But I couldn't tell how to do it exactly off the top of my head ^^'. And maybe I'm wrong.

[EiffL]

hummmm actually I see the tf.linalg.inv(tf.eye(3) - g1).... maybe I'm wrong and there is no better way to do it....