Autostronomy / AstroPhot

A fast, flexible, automated, and differentiable astronomical image 2D forward modelling tool for precise parallel multi-wavelength photometry
https://astrophot.readthedocs.io
GNU General Public License v3.0
78 stars 9 forks source link

Poisson likelihood #168

Open ConnorStoneAstro opened 4 months ago

ConnorStoneAstro commented 4 months ago

Is your feature request related to a problem? Please describe. For data which is in the very low SNR domain, where individual photon/electron counts are discernable, the proper likelihood model is a Poisson distribution.

Describe the solution you'd like Currently, it is possible to evaluate a gaussian likelihood, it should also be possible to evaluate a Poisson likelihood. This will not be able to use the Levenberg-Marquardt optimization, however posterior samplers like NUTS will be able to function reliably.

jjsutil commented 1 month ago

Hi @ConnorStoneAstro I just saw a post, actually on Facebook, linking to /Autoastronomy/AstroPhot. Happens that I'd really like to contribute, if possible, having already done the BSc. in Astronomy (UC-Chile). This issue /feat request seems very feasible. Could you please help me finding where the gaussian likelihood evaluation (and ~ how) is done, so I can try implementing a Poisson likelihood and (wishfull thinking) pull request later on? Best,

ConnorStoneAstro commented 1 month ago

Hi @jjsutil That would be amazing if you could implement this! I agree, it's very feasible as a first contribution. Here is the line for the gaussian likelihood already implemented: https://github.com/Autostronomy/AstroPhot/blob/3d0a5f39ba5eedbdac841422788d875b80a1a455/astrophot/models/core_model.py#L205

So you would need to make a poisson version of this function that also goes in the core model (probably also rename the current one from negative log likelihood, to something like negative gaussian log likelihood). Unfortunately the Poisson likelihood wont work for the LM optimizer (as far as I know), but for the gradient optimizer, MCMC, NUTS, HMC they would all need a new parameter which lets the user select which likelihood to use. Then you would probably want to make a tutorial on how to use it, you could maybe just add this in to the fitting methods tutorial with an example of very low SNR where one needs to use a poisson likelihood.

I can help anywhere you get stuck, but it would be so cool to have this feature as there are a lot of people who do low SNR work (UV, X-Ray, Gamma ray, high redshift) that is best treated with Poisson likelihood :)

jjsutil commented 1 month ago

Thanks for the response! With a couple of days I'll be ready to go. By the way, are you doing any Unit Testing along with new features? Can you give me a hint (if applies) on how to stick to the guidelines of the projects in this matter?

ConnorStoneAstro commented 1 month ago

@jjsutil Ah good question, yes AstroPhot has tons of unit tests here: https://github.com/Autostronomy/AstroPhot/tree/main/tests

That's a good point and adding in a unit test would be really helpful! I imagine for this case you could just create a small mock image (say 50x50 pixels) then run the gradient descent on the poisson likelihood where you restrict the max iterations to something small like 10 steps. Then the test would be to check that the negative poisson likelihood got smaller.

jjsutil commented 1 month ago

Hey! @ConnorStoneAstro, hope you're doing well. Also, sorry for the delay. I've been going through a lot. Had to recall a few things before.

If I'm understanding, the idea is to get the resulting image from a stack of (high SNR and maybe not many) frames. If we assume it distributes following Poisson's, maximizing $\mathrm{ln}$ $(\mathcal{L})$ gives $\lambda$ equals the mean of the samples for each pixel. I'm assuming here each frame is a grid with photon counts or something proportional (haven't onboarded/setted astrophot yet)... so the task summarizes esentially to find the mean instead of the chi2, right? where mean $= \bar{x} =\frac{1}{n} \sum x_j$

In a nutshell, this method should get the already SExtracted source images (>1 grid of counts), cropped, and passes the "right (or chosed) model" from all the available? mask and weights comes from fitting the model? and then it should return the variance image considering all the frames and taking the mean?

Could you please explain a little on how the image, model, mask and weights work in astrophot? Is this coinciding with the paragraphs in section 2.2 $\chi^2$ Optimization here, right? I apologize before hand if it's a naive question. I want to make sure I understand the output: it should? be a grid of values with the variance of the sample, which is the mean of the counts for each pixel in this case, instead of a grid of the chi-squared distribution considering: a model?

Clarifying this, it seems simple to implement. Does it makes any sense to consider a Poisson-Gaussian distribution or "mixed" to model the noise? e.g., I was looking this optics letter.

Also, do you have any recommendation on which image(s) to use for testing?

ConnorStoneAstro commented 1 month ago

Right, so section 2.2 in the paper goes over how a likelihood is computed under the assumption that the pixels have Gaussian noise, the goal here is to redo this but with poisson noise. I think the best way to get started would be to assume the image is in units of electron counts and then for the negative_log_likelihood function in the AstroPhot_Model class, you would make a copy of the function and just change this line: https://github.com/Autostronomy/AstroPhot/blob/3d0a5f39ba5eedbdac841422788d875b80a1a455/astrophot/models/core_model.py#L239

That line is the simplest case where there is no mask and the target is a single image. If you can get that one sorted, then we can talk through the rest of the cases. So the model variable will be an image generated by astrophot (the lambda for a poisson), the data will be the actual measurement (the counts from a poisson), the weight might have a different meaning for poisson, maybe closer to gain? I'm not so sure, so that's something to figure out.

I don't think we need to consider a mixed distribution. Really all images are poisson distributed, but for high SNR we can approximate the poisson as a gaussian and it makes the math a bit easier. If someone is going through the trouble of using poisson, it's probably because their data requires the extra accuracy and they won't really get a benefit from mixing. Maybe if someone was doing multi-band imaging then some bands could be poisson (low SNR) and othe rbands could be gaussian (high SNR). But that can be a future problem.

For testing I think you should just make your own image. Just make a sersic or a gaussian model, use that to sample a poisson to get some data, then try and fit it.

I hope that helps!