NormalizingFlow
Created by: RasmusOrsoe
This PR adds compatibility for normalizing flows to graphnet. Specifically, this PR does the following:
- Adds
NormalizingFlow
ingraphnet.models.flows
as a new generic model class which all normalizing flows are intended to inherit from. These models are assumed to output a single tensor containing the transformed coordinates and the jacobian associated with the transformation. For book-keeping reasons, the column indices that identifies which columns are transformed coordinates and jacobians, these models have propertiescoordinate_columns
andjacobian_columns
. These properties are used to split the tensor for loss computation (see diagram). - Adds a normalizing flow model called
INGA
which is similar to what is presented in https://arxiv.org/abs/1906.04032 . It relies on splines and the spline functions are currently put ingraphnet.models.flows.spline_blocks
and can be used both inside and outside the context ofNormalizingFlows
(although likely most relevant there). - Adds an example of training
INGA
underexamples/04_training/05_train_normalizing_flow.py
. Here the syntax of training the normalizing flow is shown. The example trainsINGA
to learn a mapping from the rather complicated distribution of the pulse attributes (xyz, time, charge, etc) into a multivariate gaussian distribution. - Introduces slight changes to the inheritance structure of our loss functions: the generic base-class
LossFunction
no longer has an abstract_forward
method; this is left for implementation of further problem specific subclasses. This is introduced here because not all machine learning paradigms follow our usualprediction
,target
argument structure; which is indeed the case for normalizing flows. - Created a new
LossFunction
subclass calledStandardLossFunction
which all of our current loss functions now inherit from. These loss functions all follow theprediction
,target
argument structure. - Created a new
LossFunction
subclassFlowLossFunction
which follow theprediction
,jacobian
argument structure. - Created
MultivariateGaussianFlowLoss
, subclassed ofFlowLossFunction
, which calculates the nllh of a multivariate Gaussian. - Slight changes to
StandardModel
that ties everything together. While working on this PR i noticed a few list comprehensions were hard to read; so I took the liberty of writing them out explicitly and adding a few comments here and there.
A diagram showing the new inheritance structure for loss functions is shown here:
A diagram showing the overall data flow and usage of the column indices for normalizing flows can be seen here:
After Christmas, I'll make a new PR that adds the ability use NormalizingFlows
in a generative sense and an example of how this is done. I decided not to include it here to avoid making this PR larger than it already is.