search

scikit-hep / uproot3-methods

Pythonic behaviors for non-I/O related ROOT classes.
BSD 3-Clause "New" or "Revised" License
21 stars 28 forks source link

boosting 4-vectors does not work on TLorentzVectorArray #62

Open lukasheinrich opened 5 years ago

lukasheinrich commented 5 years ago

Hi,

I am stumbling over a problem w/ 4-vectors in uproot_methods:

Consider the problem of having events with lorentz vectors and wanting to boost all of them into their rest frame. The operation works per-event, but not o the full full (batched) set of events

I'm looking into what the issue could be but perhaps someone is faster

x = [
    [
        [0,0,12,13],
        [0,0,4,5]
    ],
    [
        [0,0,15,17],
        [0,0,24,25]
    ],
]

mom = awkward.fromiter(x)
v = uproot_methods.TLorentzVectorArray.from_cartesian(mom[:,:,0],mom[:,:,1],mom[:,:,2],mom[:,:,3])

boost_back = (v.p3/v.p3.mag)*v.beta
print(
    [vv.boost(bb) for vv,bb in zip(v,-boost_back)]
)
print(
    v.boost(-boost_back)
)

the first print works as expected, the seconds does not and fails with 

in boost(self, p3)
    227         gamma2 = self.awkward.numpy.zeros(b2.shape, dtype=self.awkward.numpy.float64)
    228         mask = (b2 != 0)
--> 229         gamma2[mask] = (gamma[mask] - 1) / b2[mask]
    230         del mask
    231 

IndexError: too many indices for array
jpivarski commented 5 years ago

I see what's going on here: it's a casualty of scope creep. Originally, these functions were written for Numpy arrays—the idea being that you could have flat arrays of Lorentz vectors—but in most cases, they "just worked" for awkward arrays (because they share Numpy idioms) and this fact is clearly useful. However, awkward arrays don't support masked-assignment (I'm hoping to make this possible in awkward 1.0). In most of these functions, I removed masked-assignment and in-place assignment, but I must have missed this one.

The masked-assignment is intended to protect against division by zero—specifically, zero over zero (corresponding to boosting by exactly zero). To do this truly vectorially, we should let the division by zero happen and clean up afterward. Since these are zero over zero cases, the result we want to clean up is NaN. If the array is Numpy, we can mask-assign these cases; otherwise, we can use awkward's fillna method (name borrowed from Pandas).

I believe that the correct limit for this case is gamma2 → 0.5. Since you're closer to the physics of this problem, hopefully you can verify this. In fact, I don't know if we've had a customer for boosted 4-vectors yet, so checking the math would be prudent for your goals and very much appreciated for mine.

The fix is in PR #63.

jpivarski commented 4 years ago

I just noticed that I have an open PR (#63) that should fix this, and it's pretty stale. Did this fix your problem/is it no longer relevant/etc.?

jrueb commented 4 years ago

For me this problem is still relevant. I get a TypeError: (arrays of) TVector3 can only be added to/subtracted from other (arrays of) TVector3 after applying the patch. It seems that this is unrelated as it arises from a few lines further down.