Plugins for position reconstruction with conditional normalizing flow

juehang commented 2 months ago

What does the code in this PR do / what does it improve?

In this PR we introduce a set of new plugins to perform probabilistic position reconstruction using a conditional normalising flow model. This position reconstruction method provides uncertainties on the reconstructed position, and the best-fit point improves on the resolution of the MLP. Details can be found in this wiki note.

Can you briefly describe how it works?

We introduce the use of a normalizing flow model for position reconstruction. This model provides both a best-fit position and uncertainties.

Can you give a minimal working example (or illustrate with a figure)?

Please include the following if applicable:

[x] Update the docstring(s)
[x] Update the documentation
[x] Tests to check the (new) code is working as desired.
[x] Does it solve one of the open issues on github?

Notes on testing

Until the automated tests pass, please mark the PR as a draft.
On the XENONnT fork we test with database access, on private forks there is no database access for security considerations.

All italic comments can be removed from this template.

LuisSanchez25 commented 2 months ago

I wonder if it would make sense to integrate PeakPositionsFlow by inheriting from PeakPositionsBaseNT since it looks like there is a lot of overlap between what the two are computing and instead just having any additional datatypes you need get called in PeakPositionsFlow? Same with the events, in an ideal world we want to have as little duplicate code as possible but I am not 100% sure that this all can work like that.

Other than that, if the plugin is working and once it passes all the tests it will be fine by me

sebvetter commented 2 months ago

One issue with that is that PeakPositionsBaseNT has a get_tf_model function and no generic get_model function. But this could of course be changed.

juehang commented 2 months ago

I wonder if it would make sense to integrate PeakPositionsFlow by inheriting from PeakPositionsBaseNT since it looks like there is a lot of overlap between what the two are computing and instead just having any additional datatypes you need get called in PeakPositionsFlow? Same with the events, in an ideal world we want to have as little duplicate code as possible but I am not 100% sure that this all can work like that.

Other than that, if the plugin is working and once it passes all the tests it will be fine by me

I found it quite difficult to inherit from PeakPositionsBaseNT because that one cannot currently deal with the extra fields for uncertainty, and if I inherited it in a naive way I would end up overloading almost everything.

LuisSanchez25 commented 2 months ago

What do you mean by overloading everything?

juehang commented 2 months ago

I just meant re-defining the methods to accommodate the new data fields

coveralls commented 2 months ago

coverage: 90.827% (-0.02%) from 90.845% when pulling 64947fa3dba35cb0f1a3ba178947ec28413275a0 on flow-posrec into ced9f51eb26ce00eed22b633acffced0476f4659 on master.

yuema137 commented 2 months ago

Hi @juehang and @sebvetter thanks for this great work! The flow-based position reconstruction undoubtfully improves the interpretability of our ML efforts a lot! The PR is generally good and I put some minor suggestions in the review above.

One thing that I'm a little bit concerned about is the speed. The flow-based model kindly gives us a position contour, unlike other algorithms which only offer deterministic outputs. This nice feature, on the other hand, usually also lead to much more computation cost. So it will be great if you could provide some kind of estimation on the speed (like compare the speed of mlp, cnn, gcn and flow model for the same dataset) so that we have a sense for it.

sebvetter commented 2 months ago

Hi @yuema137 , In this note: https://xe1t-wiki.lngs.infn.it/doku.php?id=xenon:xenonnt:juehang:flow_posrec_proposal_sr2 you can find the attached plot cnf_vs_current_timing_per_sample

The Flow model is faster than the CNN

XENONnT / straxen