Closed EiffL closed 3 years ago
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.
This is amazing!!!
For the TensorFlow implementation, maybe this will be helpful: https://github.com/tensorflow/tensorflow/blob/0b9c9497f1c7f924f3fbeffcf9426d5174a68600/tensorflow/python/eager/backprop.py#L510
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:
@tf.custom_gradient
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.
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....
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:
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