How to copy form one ROOT file to another (assuming NanoAOD-like structure)

[jpivarski]

Edit: See the next comment below. Now there's a package that provides a convenience function for this!

This is such a common request that there ought to be a simple function that just copies TBranches. (Does anyone have time to write such a function?) Meanwhile, this is a rough guide.

If you just try to read from Uproot into an Awkward Array and write that same Awkward Array back, it doesn't work:

>>> import skhep_testdata
>>> import uproot
>>> tree = uproot.open(skhep_testdata("uproot-HZZ.root").data_path())["events"]
>>> arrays = tree.arrays(filter_name=["Jet_*", "Muon_*"])
>>> arrays
<Array [{Jet_Px: [], Jet_Py: [], ...}, ...] type='2421 * {Jet_Px: var * flo...'>
>>> arrays.type.show()
2421 * {
    Jet_Px: var * float32,
    Jet_Py: var * float32,
    Jet_Pz: var * float32,
    Jet_E: var * float32,
    Jet_btag: var * float32,
    Jet_ID: var * bool,
    Muon_Px: var * float32,
    Muon_Py: var * float32,
    Muon_Pz: var * float32,
    Muon_E: var * float32,
    Muon_Charge: var * int32,
    Muon_Iso: var * float32
}

and now writing:

>>> outfile = uproot.recreate("/tmp/out.root")
>>> outfile["tree"] = arrays
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/jpivarski/irishep/uproot5/src/uproot/writing/writable.py", line 973, in __setitem__
    self.update({where: what})
  File "/home/jpivarski/irishep/uproot5/src/uproot/writing/writable.py", line 1546, in update
    uproot.writing.identify.add_to_directory(v, name, directory, streamers)
  File "/home/jpivarski/irishep/uproot5/src/uproot/writing/identify.py", line 151, in add_to_directory
    tree = directory.mktree(name, metadata)
  File "/home/jpivarski/irishep/uproot5/src/uproot/writing/writable.py", line 1297, in mktree
    directory._cascading.add_tree(
  File "/home/jpivarski/irishep/uproot5/src/uproot/writing/_cascade.py", line 1788, in add_tree
    tree = uproot.writing._cascadetree.Tree(
  File "/home/jpivarski/irishep/uproot5/src/uproot/writing/_cascadetree.py", line 284, in __init__
    raise TypeError(
TypeError: fields of a record must be NumPy types, though the record itself may be in a jagged array

    field 'Jet_Px' has type var * float32

It was a giant usability oversight for me to not realize that someone might want to write the same array they just read, without modification.

The error here is because the output-writer (assigning to outfile["tree"]) is expecting a dict, and arrays is an ak.Array.

Okay, we can make a dict. We can get all of the record field names (e.g. "Jet_Px") with ak.fields and all of the field values (e.g. the array corresponding to jet $p_T$) with ak.unzip.

This is able to write...

>>> outfile["tree"] = dict(zip(ak.fields(arrays), ak.unzip(arrays)))

...but it's probably not what you want.

>>> outfile["tree"].show()
name                 | typename                 | interpretation                
---------------------+--------------------------+-------------------------------
nJet_Px              | int32_t                  | AsDtype('>i4')
Jet_Px               | float[]                  | AsJagged(AsDtype('>f4'))
nJet_Py              | int32_t                  | AsDtype('>i4')
Jet_Py               | float[]                  | AsJagged(AsDtype('>f4'))
nJet_Pz              | int32_t                  | AsDtype('>i4')
Jet_Pz               | float[]                  | AsJagged(AsDtype('>f4'))
nJet_E               | int32_t                  | AsDtype('>i4')
Jet_E                | float[]                  | AsJagged(AsDtype('>f4'))
nJet_btag            | int32_t                  | AsDtype('>i4')
Jet_btag             | float[]                  | AsJagged(AsDtype('>f4'))
nJet_ID              | int32_t                  | AsDtype('>i4')
Jet_ID               | bool[]                   | AsJagged(AsDtype('bool'))
nMuon_Px             | int32_t                  | AsDtype('>i4')
Muon_Px              | float[]                  | AsJagged(AsDtype('>f4'))
nMuon_Py             | int32_t                  | AsDtype('>i4')
Muon_Py              | float[]                  | AsJagged(AsDtype('>f4'))
nMuon_Pz             | int32_t                  | AsDtype('>i4')
Muon_Pz              | float[]                  | AsJagged(AsDtype('>f4'))
nMuon_E              | int32_t                  | AsDtype('>i4')
Muon_E               | float[]                  | AsJagged(AsDtype('>f4'))
nMuon_Charge         | int32_t                  | AsDtype('>i4')
Muon_Charge          | int32_t[]                | AsJagged(AsDtype('>i4'))
nMuon_Iso            | int32_t                  | AsDtype('>i4')
Muon_Iso             | float[]                  | AsJagged(AsDtype('>f4'))

It's because the Awkward Array doesn't know that all jet attributes have the same number of values per event and all muon attributes have the same number of values per event, and so Uproot has to make a separate counter for each.

You can inform Uproot that these arrays should share counters by putting them into the same records:

>>> jets = ak.zip({name[4:]: array for name, array in zip(ak.fields(arrays), ak.unzip(arrays)) if name.startswith("Jet_")})
>>> jets.type.show()
2421 * var * {
    Px: float32,
    Py: float32,
    Pz: float32,
    E: float32,
    btag: float32,
    ID: bool
}
>>> muons = ak.zip({name[5:]: array for name, array in zip(ak.fields(arrays), ak.unzip(arrays)) if name.startswith("Muon_")})
>>> muons.type.show()
2421 * var * {
    Px: float32,
    Py: float32,
    Pz: float32,
    E: float32,
    Charge: int32,
    Iso: float32
}

(If they were not compatible, the ak.zip would fail.) The important thing now is that the var is outside the record in the type descriptor: all of these attributes have the same raggedness (and different between jets and muons).

Now Uproot has enough information to make only one counter for each particle type.

>>> outfile["tree"] = {"Jet": jets, "Muon": muons}
>>> outfile["tree"].show()
name                 | typename                 | interpretation                
---------------------+--------------------------+-------------------------------
nJet                 | int32_t                  | AsDtype('>i4')
Jet_Px               | float[]                  | AsJagged(AsDtype('>f4'))
Jet_Py               | float[]                  | AsJagged(AsDtype('>f4'))
Jet_Pz               | float[]                  | AsJagged(AsDtype('>f4'))
Jet_E                | float[]                  | AsJagged(AsDtype('>f4'))
Jet_btag             | float[]                  | AsJagged(AsDtype('>f4'))
Jet_ID               | bool[]                   | AsJagged(AsDtype('bool'))
nMuon                | int32_t                  | AsDtype('>i4')
Muon_Px              | float[]                  | AsJagged(AsDtype('>f4'))
Muon_Py              | float[]                  | AsJagged(AsDtype('>f4'))
Muon_Pz              | float[]                  | AsJagged(AsDtype('>f4'))
Muon_E               | float[]                  | AsJagged(AsDtype('>f4'))
Muon_Charge          | int32_t[]                | AsJagged(AsDtype('>i4'))
Muon_Iso             | float[]                  | AsJagged(AsDtype('>f4'))

and that's probably what you want.

[jpivarski]

The hepconvert package contains high-level functions to do this sort of thing, so that you don't have to write scripts like the above anymore.