Reparametrize regression model output

[jpata]

Currently, the regression model outputs pt, E, eta, sin_phi, cos_phi.

However, pt and E are highly similar, so the model has to effectively learn almost the same quantity twice, which may hurt the parameter efficiency and loss landscape.
In addition, it means that the model may not necessarily reconstruct particles with a positive mass: m^2 = E^2 - pt^2 - pz^2.

It might be useful to reparametrize the loss and the model output in alternative ways, e.g. always assuming a fixed positive mass (e.g. charged pion mass), and then predicting only e.g. pt, inferring energy using the mass relation.