Reimplement Gauss Moments from ngmix with tensorflow

[andrevitorelli]

We want this example: https://github.com/esheldon/ngmix/blob/master/examples/metacal/metacal.py

working with functions that have the exact same code from ngmix, but where operations are substituted by tensorflow operations so that they can be automatically differentiated.

[EiffL]

Oh wow, @andrevitorelli I just poked my head in your restructure branch, and I'm finding a copy/paste of a huge chunk of ngmix. Are you sure this is what we want?

We only talked about reimplementing the moments functions.

I'm not saying necessarily that's bad, vbut that's a huge change of scope, what's the reasoning behind it?

[andrevitorelli]

Yes, @EiffL , so... I don't intend to have this large scope. I was testing things by altering the original code as less as possible (to be sure I didn't introduce any mistakes that would invalidate the check), but the dependency tree took me out of scope. I'm already trimming down. Thanks!

[EiffL]

right right, but we only want to be able to measure moments and get the same result as ngmix, we don't want any of the rest of the ngmix logic, which is not going to work well with TF

[EiffL]

so for instance, I may be wrong, but I don't think we want any of the gmix stuf

[EiffL]

So for instance if I introspect the code a little bit it looks like this is one of the main files:

https://github.com/esheldon/ngmix/blob/master/ngmix/gmix/gmix_nb.py

And the get_weighted_sums is one of the functions we want to reproduce.

[andrevitorelli]

It is still very obscure to me how it works. get_weighted_sums works on a pixel structure that contains wcs transformations which we don't need for what we're developing at the moment... ie. this numba function depends on the structure of the inputs. When tracing it's use, you arrive at the observation.pixels, and then pixels module, which adds a lot in complexity to what we do. I'll try using them without going through observations to see if I get it right, but just because other than that, our development is stalled waiting tfa interpolation solutions.

[esheldon]

The pixel structure holds a bunch of information about each pixel, for example its location in the u, v plane, the image value, uncertainty, etc.

If you have constant noise and identity transformation between pixels and u, v coords you can simplify this function a lot

[EiffL]

Yep, so maybe let's take a step back @andrevitorelli

What do we want? We want to have a TensorFlow function that implements gaussian moments, and that, in the limit of our current asumptions (no wcs, constant noise, centerd objects, etc.) will return the same output (e1,e2,T, and maybe s2n) as ngmix.GaussMom.go() when applied on a given postage observation.

How would this look like? This will not be a line-by-line translation of ngmix, we can look at the ngmix code to see what operations are needed, but we need to reimplement it from the ground up in the "TensorFlow way". So, for instance, we don't want to translate get_weighted_sums but implement a function that does nothing more than computing an equivalent Gaussian weigthed sum of the pixels of the stamp.

And at the end, the test we will want to do will be something very similar our test_get_metcal, so something like:

# ngmix part
obs = make_data(rng=rng, noise=args['noise'], shear=shear_true)
fitter = ngmix.gaussmom.GaussMom(fwhm=weight_fwhm)
ngmix_results = fitter.go(obs) # this will give us the data we want to compare to

# TensorFlow part
im = obs.image.reshape(1,45,45).astype('float32')
tf_results = autometacal.gaussian_moments(im, fwhm=weight_fwhm, ...) # maybe some extra args, like noise or 

# Testing part
assert .... 

How to get there? We want to compute Gaussian weighted moments, so this is not going to be very far from our existing toy function get_ellipticity for instance in this notebook.

So, you would just need to follow the chain of steps in ngmix.GaussMom.go and check we are doing the same thing in the limit of simple noise, no wcs, etc.

Let me know if that makes sense :-) and don't hesitate if you have questions! Happy to look at prototypes along the way

[EiffL]

closed by #32