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

Adding TLorentzVectorArray #95

Closed jeffkrupa closed 3 years ago

jeffkrupa commented 3 years ago

Hi, I am defining two TLorentzVectorArray and wanting to take the sum. I can add mu_p4 to itself, or met_p4 to itself, but I cannot add them to each other without it blowing up. I'm using awkward1 and uproot_methods3. What am I doing wrong?

>>> candidatemuon = ak.firsts(events.Muon)
>>> mu_p4 = TLorentzVectorArray.from_ptetaphim(ak.fill_none(candidatemuon.pt,0),ak.fill_none(candidatemuon.eta,0),ak.fill_none(candidatemuon.phi,0),ak.fill_none(candidatemuon.mass,0))
>>> met_p4 = TLorentzVectorArray.from_ptetaphim(ak.from_iter([[v] for v in events.MET.pt]),ak.from_iter( [[v] for v in np.zeros(len(events))]), ak.from_iter([[v] for v in events.MET.phi]), ak.from_iter([[v] for v in np.zeros(len(events))]) )
>>> mu_p4.shape, met_p4.shape
((100000,), (100000,))
>>> mu_p4 + met_p4
MemoryError: Unable to allocate 74.5 GiB for an array with shape (100000, 100000) and data type float64
jpivarski commented 3 years ago

You can't use Awkward 1 with uproot-methods (any version). Lorentz vectors in the Awkward 1 era will be handled by the Vector package, which is going to get more fully developed in March, or by coffea.nanoevents.methods.vector in the interim.

I'm guessing that what's happening here is that Awkward 0 doesn't recognize an Awkward 1 array and is iterating over it, or that some broadcasting rule applies and makes it want to generate a huge array (74.5 GiB).