use absolute value of pz to protect against minus infinity for rapidity

[nikoladze]

Thanks for providing this code! I started playing a bit with the ParticleTransformer and ran into an issue of getting negative infinite values when pz is negative and very close in magnitude to the energy. The current implementation for calculating the rapidity protects against positive pz being close to energy, but not negative ones. This PR suggests using the absolute value and multiplying by the sign afterwards.

def rapidity_old(energy, pz):
    return 0.5 * torch.log(1 + (2 * pz) / (energy - pz).clamp(min=1e-20))
    
def rapidity_new(energy, pz):
    return 0.5 * pz.sign() * torch.log(1 + (2 * pz.abs()) / (energy - pz.abs()).clamp(min=1e-20))

>>> import torch
>>> rapidity_old(torch.tensor(10.), torch.tensor(9.))
tensor(1.4722)
>>> rapidity_old(torch.tensor(10.), torch.tensor(10.))
tensor(24.5237)
>>> rapidity_old(torch.tensor(10.), torch.tensor(-9.))
tensor(-1.4722)
>>> rapidity_old(torch.tensor(10.), torch.tensor(-9.9999999))
tensor(-inf)
>>> rapidity_old(torch.tensor(10.), torch.tensor(-10.))
tensor(-inf)
>>>
>>> rapidity_new(torch.tensor(10.), torch.tensor(9.))
tensor(1.4722)
>>> rapidity_new(torch.tensor(10.), torch.tensor(10.))
tensor(24.5237)
>>> rapidity_new(torch.tensor(10.), torch.tensor(-9.))
tensor(-1.4722)
>>> rapidity_new(torch.tensor(10.), torch.tensor(-9.9999999))
tensor(-24.5237)
>>> rapidity_new(torch.tensor(10.), torch.tensor(-10.))
tensor(-24.5237)