Can I recalculate the energy of an awkward array of Vectors by declaring a new mass value?

My question is about the Vector module in scikit-hep.

https://vector.readthedocs.io/en/latest/index.html

I have an awkward array of vectors and I'd like to set the mass of all of them to be a common value. For example, I can do this with a single vector object.

x = vector.obj(pt=2, eta=1.5, phi=1, energy=10)
y = x.from_rhophietatau(rho=x.rho, eta=x.eta, phi=x.phi, tau=20)

print(f"{x.mass:6.3f}  {x.pt}  {x.eta}  {x.phi}  {x.energy:6.2f}")
print(f"{y.mass:6.3f}  {y.pt}  {y.eta}  {y.phi}  {y.energy:6.2f}")

Output

 8.824  2  1.5  1   10.00
20.000  2  1.5  1   20.55

But suppose I want to do this with an awkward array of vectors?

Let me start with some starter code from this previous question:

Using awkward-array with zip/unzip with two different physics objects

First, I'll get an input file

curl http://opendata.cern.ch/record/12361/files/SMHiggsToZZTo4L.root --output SMHiggsToZZTo4L.root

Then I'll make use of the code from the answer to that question:

import numpy as np
import matplotlib.pylab as plt

import uproot
import awkward as ak

import vector
vector.register_awkward()


infile = uproot.open("/tmp/SMHiggsToZZTo4L.root")

muon_branch_arrays = infile["Events"].arrays(filter_name="Muon_*")
electron_branch_arrays = infile["Events"].arrays(filter_name="Electron_*")

muons = ak.zip({
    "pt": muon_branch_arrays["Muon_pt"],
    "phi": muon_branch_arrays["Muon_phi"],
    "eta": muon_branch_arrays["Muon_eta"],
    "mass": muon_branch_arrays["Muon_mass"],
    "charge": muon_branch_arrays["Muon_charge"],
}, with_name="Momentum4D")

quads = ak.combinations(muons, 4)
mu1, mu2, mu3, mu4 = ak.unzip(quads)

p4 = mu1 + mu2 + mu3 + mu4

The type of p4 is <class 'vector._backends.awkward_.MomentumArray4D'>. Is there a way to set all the masses of the p4 objects to be, for example, 125? While this is not exactly my analysis, I need to do something similar where I will then use p4 to boost the muX objects to the CM frame of p4 and look at some relative angles. But I need to set the mass of p4 to be a constant value.

Is this possible? Thanks!

Matt

Solution

This is a well written question, thank you for the effort!

The answer here is yes, you can set a new value for the mass! One would do this updating the mass field using ak.with_field, or using the subscript __setitem__ operator, e.g.

p4['mass'] = 125.0

This will internally call ak.with_field, which you could also use e.g.

p4 = ak.with_field(p4, 125.0, "mass")

and broadcasts the 125.0 value against the rest of the array.

It is sometimes more convenient to use the vector Awkward constructors, as it is slightly less typing:

muons = vector.zip({
    'pt': muon_branch_arrays['Muon_pt'],
    'phi': muon_branch_arrays['Muon_phi'],
    'eta': muon_branch_arrays['Muon_eta'],
    'charge': muon_branch_arrays['Muon_charge'],
    'mass': muon_branch_arrays['Muon_mass'],
})

vector determines what kind of array you are building from the field names. This provides a good opportunity to highlight something important: vector supports aliases for fields, e.g. tau ↔ mass. If you compare the fields of the muons array above with the array you built with ak.zip, you'll notice that my muons array has fields ['rho', 'phi', 'eta', 'tau', 'charge'] whilst your muon array has fields ['pt', 'phi', 'eta', 'mass', 'charge']. What's happening here is that vector is canonicalising the field names. This means that, were you to build an array in this manner, you'd want to use p4['tau'] = 125.0 instead of p4['mass'] = 125.0.

This would also be apparent if you transformed your muons array in any way, e.g. double_muons = muons + muons. You'd find that the result loses the charge field, and has tau instead of mass and rho instead of pt. So, something to be mindful of if you need to set a field.

The reason that p4['mass'] = 125 works, but pt['mass'][:] = 125 does not is because of how Awkward Array is designed. Whilst Awkward Arrays are immutable, this is only half of the story. You can already see that there is some kind of mutability - we can modify a field in-place. This works because, whilst the underlying "layouts" from which Arrays are built do not allow users to modify their values, the high level ak.Array can be given a new layout. This is what __setitem__ does under the hood, i.e. https://github.com/scikit-hep/awkward/blob/72c9edd55b9c4611ffc46952cda4cf9920a91315/src/awkward/highlevel.py#L1062-L1063