Writting coffea Schema files

The interpretation of the TTree data is configurable via schema objects.

Schema teaches the event processor how to group variables into collections, so operations can be run over entire collection at once. And we can also define some handy behaviors for a specific collection in a schema.

In this demo, we will create our own dummy schema and implement our own dummy behavior.

Let’s look at the content of dummy_nanoevents.root file with BaseSchema. We’ll show how to construct a schema that can be used to interpret this root file later.

The events object can be instantiated as follows:

from coffea.nanoevents import NanoEventsFactory, BaseSchema, NanoAODSchema
fname = "../../samples/dummy_nanoevents.root"
events = NanoEventsFactory.from_root(
           fname,
           schemaclass=BaseSchema
         ).events()
print(events.fields)

['run', 'luminosityBlock', 'event', 'nMuon', 'Muon_pt', 'Muon_eta', 'Muon_phi', 'Muon_mass', 'Muon_charge', 'Muon_flag', 'Muon_dxy', 'Muon_dxyErr', 'Muon_dz', 'Muon_dzErr', 'Electron_pt', 'Electron_eta', 'Electron_phi', 'Electron_mass', 'Electron_charge', 'Electron_flag', 'Electron_dxy', 'Electron_dxyErr', 'Electron_dz', 'Electron_dzErr', 'Jet_pt', 'Jet_eta', 'Jet_phi', 'Jet_energy', 'Jet_ID', 'Jet_SubjetsCounts', 'Subjet_pt', 'Subjet_eta', 'Subjet_phi', 'Subjet_energy']

Now we can copy the skeleton of a schema class:

class DummySchema(BaseSchema):
    """
    """
    def __init__(self, base_form):
        super().__init__(base_form)
        self._form["contents"] = self._build_collections(self._form["contents"])

    def _build_collections(self, branch_forms):
        output = {}
        return output

    @property
    def behavior(self):
        from coffea.nanoevents.methods import base, vector
        behavior = {}
        behavior.update(base.behavior)
        behavior.update(vector.behavior)
        return behavior

As you can see, this schema is so simple it is not useful currently. If we call the events again with our own schema, we’ll find it contains nothing.

events = NanoEventsFactory.from_root(
           fname,
           schemaclass=DummySchema
         ).events()
events.fields

[]

Create collections

In schema, the branch_forms is a python dictionary used to define branch grouping.

By default, using BaseSchema, it will be completely flat:

branch_forms={
  "particle_pt":{},
  "particle_eta":{},
  "particle_phi":{},
  "particle_mass":{},
  ...
}

Just as you have seen when we open the dummy root file with BaseSchema, all the branches are listed when do print(events.fields). These branches are independent to each other. For example when you do a selection on particle_pt, the branch particle_eta will not be affected.

We typically like to do selections collectively. If some particles don’t pass the selection on particle_pt, their particle_eta should also be eliminated from the list.

I.e. We would like to put some branches into the same collection, as what follows:

new_branch_forms={
  "particle": schemas.zip_forms({
      "pt" : branch_forms["particle_pt"],
      "eta" : branch_forms["particle_eta"],
      "phi" : branch_forms["particle_phi"],
      "mass" : branch_forms["particle_mass"],
  })
}

So when we want to call particle_pt, we actually do particle.pt. And we can apply selections on the particle level so particle_pt and particle_eta are selected at the same time now.

All of this is to be implemented in the Schema._build_collections method.

For example, let’s add the Electron collection to our schema. To do this we need to import zip_forms.

from coffea.nanoevents.schemas import zip_forms
class DummySchema(BaseSchema):
    """
    """
    def __init__(self, base_form):
        super().__init__(base_form)
        self._form["contents"] = self._build_collections(self._form["contents"])

    def _build_collections(self, branch_forms):
        output = {}
        output["Electron"] = zip_forms(
            {
                "pt" : branch_forms["Electron_pt"],
                "eta" : branch_forms ["Electron_eta"] ,
                "phi": branch_forms["Electron_phi"],
                "mass": branch_forms["Electron_mass"],
                #"xx": branch_forms["Electron_xx"],
            },
            "Electron",
        )
        return output

    @property
    def behavior(self):
        from coffea.nanoevents.methods import base, vector
        behavior = {}
        behavior.update(base.behavior)
        behavior.update(vector.behavior)
        return behavior

Now we successfully created a schema with one collection Electron. It will be able to recognize branches with name Electron_pt, Electron_eta, Electron_phi, Electron_mass. Try to call the events again.

events = NanoEventsFactory.from_root(
           fname,
           schemaclass=DummySchema
         ).events()
print(events.fields)
print(events.Electron.fields)

['Electron']
['pt', 'eta', 'phi', 'mass']

Congratulations, we have now collected the fields related to an electron into a common collection! We can use the mask and do selection on the whole collection at once now:

mask = (events.Electron.pt>3) & (events.Electron.pt<60)
good_elec = events.Electron[mask]
print(good_elec.pt)
print(good_elec.eta)

[[49.3, 38.1], [48.9, 43.9, 50.8, 43.9], ... 40.4], [44.6, 45.5, 49.2, 58.3, 51.9]]
[[0.366, 0.437], [0.149, 0.236, 0.472, ... [0.442, 0.854, 0.344, 0.156, 0.564]]

However, if you put some unknown branches to the collection, errors will be returned. For example, uncomment the following line in DummySchema:

"xx": branch_forms["Electron_xx"],

Run the above code again, you will see:

KeyError: 'Electron_xx'

Of course, we can make a long list and manually group all the TBranches. But use naming conventions would enable you to write the schema in a very neat way. Like what we have in the dummy root file, branches are named as object_varible. So we can define some collections and put the corresponding TBranches into them.

4 collections are defined below in mixins:

class DummySchema(BaseSchema):
    """
    """
    mixins = {
        'Electron': "PtEtaPhiMLorentzVector",
        'Muon': 'PtEtaPhiMLorentzVector',
        'Jet': 'PtEtaPhiELorentzVector',
        'Subjet': 'PtEtaPhiELorentzVector',
    }
    def __init__(self, base_form):
        super().__init__(base_form)
        self._form["contents"] = self._build_collections(self._form["contents"])

    def _build_collections(self, branch_forms):
        output = {}
        ## Making the basic
        for name in self.mixins:
            mixin = self.mixins.get(name, "NanoCollection")
            output[name] = zip_forms({
                k[len(name) + 1 :]: branch_forms[k]
                for k in branch_forms
                if k.startswith(name + "_")
            },
            name,
            record_name=mixin)
        return output

    @property
    def behavior(self):
        from coffea.nanoevents.methods import base, vector
        behavior = {}
        behavior.update(base.behavior)
        behavior.update(vector.behavior)
        return behavior

We defined Electron, Muon, Jet, Subjet in the above DummySchema, and we used for k in branch_forms to search all the branches start with {collection}_ and put them into the corresponding collections. Now open the dummy root file again, we can see the defined collections.

events = NanoEventsFactory.from_root(
           fname,
           schemaclass=DummySchema
         ).events()
print(events.fields)
print(events.Electron.fields)
print(events.Electron.px)
print(events.Electron.theta)

['Electron', 'Muon', 'Jet', 'Subjet']
['pt', 'eta', 'phi', 'mass', 'charge', 'flag', 'dxy', 'dxyErr', 'dz', 'dzErr']
[[48.7, 33.9], [26.6, 40.2, 29.8, 31.6], ... 47.3, 24.6], [33, 30.2, 47.8, 31.5, 45]]
[[1.21, 1.15], [1.42, 1.34, 1.12, 1.47], ... 1.56], [1.14, 0.805, 1.23, 1.42, 1.03]]

Maybe you already noticed that we printed something called events.Electron.px and events.Electron.theta in the previous block, which don’t exist in the dummy root file at all. How were they created? This is actually the behavior this collection has, which we’ll talk about right now.

Look at the collection list:

mixins = {
        'Electron': "PtEtaPhiMLorentzVector",
        'Muon': 'PtEtaPhiMLorentzVector',
        'Jet': 'PtEtaPhiELorentzVector',
        'Subjet': 'PtEtaPhiELorentzVector',
    }

PtEtaPhiMLorentzVector is the name of behavior that each collection has. It is defined here. And we imported it through from coffea.nanoevents.methods import base, vector

Create behavior

Aside from putting different branches into collections, we can also add behavior to collections. This means additional awkward arrays are generated on-the-fly via predefined algorithm. Like we can get px, theta previously.

A bunch of other common physics behaviors are already provided in coffea, and you can find them in methods.

To write our own coffea behavior, we need to define the behavior class. In the following code, we create DummyBehavior, which has one function plus1(). It returns the particle.mass+1 when you call particle.plus1.

import awkward1
dummybehavior={}
@awkward1.mixin_class(dummybehavior)
class DummyBehavior:
    @property
    def plus1(self):
        return self.mass+1

class DummySchema(BaseSchema):
    """
    """
    mixins = {
        'Electron': "DummyBehavior",
        'Muon': 'PtEtaPhiMLorentzVector',
        'Jet': 'PtEtaPhiELorentzVector',
        'Subjet': 'PtEtaPhiELorentzVector',
    }
    def __init__(self, base_form):
        super().__init__(base_form)
        self._form["contents"] = self._build_collections(self._form["contents"])

    def _build_collections(self, branch_forms):
        output = {}
        ## Making the basic
        for name in self.mixins:
            mixin = self.mixins.get(name, "NanoCollection")
            output[name] = zip_forms({
                k[len(name) + 1 :]: branch_forms[k]
                for k in branch_forms
                if k.startswith(name + "_")
            },
            name,
            record_name=mixin)
        return output

    @property
    def behavior(self):
        from coffea.nanoevents.methods import base, vector
        behavior = {}
        behavior.update(base.behavior)
        behavior.update(vector.behavior)
        behavior.update(dummybehavior)
        return behavior

The behavior of Electron is changed to DummyBehavior in the above DummySchema. Also the dummybehavior is included in the DummySchema through behavior.update(dummybehavior). Now try our customer behavior:

events = NanoEventsFactory.from_root(
           fname,
           schemaclass=DummySchema
         ).events()
print(events.fields)
print(events.Electron.fields)
print(events.Electron.mass[0])
print(events.Electron.plus1[0])

['Electron', 'Muon', 'Jet', 'Subjet']
['pt', 'eta', 'phi', 'mass', 'charge', 'flag', 'dxy', 'dxyErr', 'dz', 'dzErr']
[0.00051, 0.00051]
[1, 1]

Python did the rounding so the output is not exactly 1.00051.

Since we changed the behavior, print(events.Electron.theta) should not work now.

print(events.Electron.theta)

Nested jagged array implementation

A jagged array is an array of arrays. Each array is not guaranteed to be of the same size. You could have one parent array that each of its component is another array.

But why do we need such kind of data structure? It will be clear after we take a look at the dummy_nanoevents.cc file. For convenience, I pasted the related codes below:

#include "TFile.h"
#include "TRandom.h"
#include "TTree.h"

void
dummy_nanoevent()
{
  UInt_t run, event, luminosityBlock;

  std::vector<double> Jets_Pt;
  std::vector<double> Jets_Eta;
  std::vector<double> Jets_Phi;
  std::vector<double> Jets_E;
  std::vector<double> Jets_ID;
  std::vector<int> Jets_SubjetsCounts;
  std::vector<double> Jets_subjet_Pt;
  std::vector<double> Jets_subjet_Eta;
  std::vector<double> Jets_subjet_Phi;
  std::vector<double> Jets_subjet_E;

  TFile* Tfile = Tfile = TFile::Open( "dummy_nanoevents.root", "RECREATE" );
  TTree* ttree = new TTree( "Events", "" );

  ttree->Branch( "run",                &run,             "run/I" );
  ttree->Branch( "luminosityBlock",    &luminosityBlock, "luminosityBlock/I" );
  ttree->Branch( "event",              &event,           "event/I" );

  ttree->Branch( "Jet_pt",            &Jets_Pt );
  ttree->Branch( "Jet_eta",           &Jets_Eta );
  ttree->Branch( "Jet_phi",           &Jets_Phi );
  ttree->Branch( "Jet_energy",             &Jets_E );
  ttree->Branch( "Jet_ID",            &Jets_ID );
  ttree->Branch( "Jet_SubjetsCounts", &Jets_SubjetsCounts );
  ttree->Branch( "Subjet_pt",     &Jets_subjet_Pt );
  ttree->Branch( "Subjet_eta",    &Jets_subjet_Eta );
  ttree->Branch( "Subjet_phi",    &Jets_subjet_Phi );
  ttree->Branch( "Subjet_energy",      &Jets_subjet_E );

  for( Int_t ev = 0; ev < 100; ev++ ){
    run             = 1;
    event           = ev;
    luminosityBlock = 1000 * ev;
    Int_t njet = Int_t( 3 * gRandom->Rndm() + 1 );

    Jets_Pt.clear();
    Jets_Eta.clear();
    Jets_Phi.clear();
    Jets_E.clear();
    Jets_ID.clear();
    Jets_SubjetsCounts.clear();
    Jets_subjet_Pt.clear();
    Jets_subjet_Eta.clear();
    Jets_subjet_Phi.clear();
    Jets_subjet_E.clear();

    for( int i = 0; i < njet; i++ ){
      Jets_Pt.push_back( 10 * gRandom->Rndm() );
      Jets_Eta.push_back( gRandom->Rndm() );
      Jets_Phi.push_back( gRandom->Rndm() );
      Jets_E.push_back( gRandom->Gaus( 50, 10 ) );
      Jets_ID.push_back( int( 7*gRandom->Rndm() ) );
      // subjets
      Int_t jets_sub = Int_t( 3 * gRandom->Rndm() );
      Jets_SubjetsCounts.push_back( jets_sub );

      for( int i = 0; i < jets_sub; i++ ){
        Jets_subjet_Pt.push_back( 10 * gRandom->Rndm() );
        Jets_subjet_Eta.push_back( gRandom->Rndm() );
        Jets_subjet_Phi.push_back( gRandom->Rndm() );
        Jets_subjet_E.push_back( gRandom->Gaus( 25, 10 ) );
      }
    }

    ttree->Fill();
  }

  ttree->Write();
}

In the file, we generated {njet} jets and each jet has {jets_sub} subjets. So the structure looks like this:

'jet' : {
          'pt' : [n],
          'eta' : [n],
          'phi' : [n],
          'energy' : [n],
          'ID' : [n],
          'subjets' : [
              {'pt':[m0], 'eta':[m0], 'phi':[m0], 'energy':[m0]},   # 0
              {'pt':[m1], 'eta':[m1], 'phi':[m1], 'energy':[m1]},   # 1
              ...
              {'pt':[m], 'eta':[m], 'phi':[m], 'energy':[m]},   # n-1
          ]
}

We have n jets here and each jet has m subjets. m does not need to be the same for all the jets, that’s where the name jagged comes. jet.pt returns array with n numbers ([0, 1, ... n-1]), while jet.subjets.pt returns n arrays ([[m0], [m1], ... [m]]).

Note that when we create the TBranches, we can define it to store vector of vectors vector< vector<double> >, this is intuitively the same structure as we shown above. However, Coffea doesn’t recognize this structure. If you have this type of branch in the root file, you will need to make modifications to the generation codes.

We can store the jets as a flat vector (length n) and use another flat vector to store all the subjets (length N). Since one jet at least has one subjet, N>=n should be fulfilled. Additionally, we need another Branch Jet_SubjetsCounts (with length n) to tell us how many subjets that each jet contains.

So if we have the following vectors:

Jets_energy = [10, 10.1, 10.2];
Jet_SubjetsCounts = [1, 3, 2];
Jets_subjet_energy = [10, 1.1, 5, 4, 5.1, 5.1];

It means we have 3 jets, the mapping is: Jet: [10, 10.1, 10.2]  -> Subjets: [ {10}, {1.1, 5, 4}, {5.1, 5.1}]

So how to implement this structure in coffea? It is actually very easy because we only need to import nest_jagged_forms and follow the rules of this predefined function:

from coffea.nanoevents.schemas import nest_jagged_forms

dummybehavior={}
@awkward1.mixin_class(dummybehavior)
class DummyBehavior:
    @property
    def plus1(self):
        return self.mass+1

class DummySchema(BaseSchema):
    """
    """
    mixins = {
        'Electron': "DummyBehavior",
        'Muon': 'PtEtaPhiMLorentzVector',
        'Jet': 'PtEtaPhiELorentzVector',
        'Subjet': 'PtEtaPhiELorentzVector',
    }
    def __init__(self, base_form):
        super().__init__(base_form)
        self._form["contents"] = self._build_collections(self._form["contents"])

    def _build_collections(self, branch_forms):
        output = {}
        ## Making the basic
        for name in self.mixins:
            mixin = self.mixins.get(name, "NanoCollection")
            output[name] = zip_forms({
                k[len(name) + 1 :]: branch_forms[k]
                for k in branch_forms
                if k.startswith(name + "_")
            },
            name,
            record_name=mixin)

        nest_jagged_forms(output['Jet'],
                          output.pop('Subjet'),
                          'SubjetsCounts',
                          'Subjets')
        return output

    @property
    def behavior(self):
        from coffea.nanoevents.methods import base, vector
        behavior = {}
        behavior.update(base.behavior)
        behavior.update(vector.behavior)
        behavior.update(dummybehavior)
        return behavior

i `nest_jagged_forms <https://github.com/yihui-lai/coffea/blob/351cc727845ab83a8e31a193dc06e534bedb97fe/coffea/nanoevents/schemas/base.py#L62>`__ is defined as follows. The first input parameter is the parent array, the second parameter is the child array. Then the third one is the counts array. The final one is the new name for the child array.

def nest_jagged_forms(parent, child, counts_name, name):
    """Place child listarray inside parent listarray as a double-jagged array"""
    if not parent["class"].startswith("ListOffsetArray"):
        raise ValueError
    if parent["content"]["class"] != "RecordArray":
        raise ValueError
    if not child["class"].startswith("ListOffsetArray"):
        raise ValueError
    counts = parent["content"]["contents"][counts_name]
    offsets = transforms.counts2offsets_form(counts)
    inner = listarray_form(child["content"], offsets)
    parent["content"]["contents"][name] = inner

We can try this DummySchema now:

events = NanoEventsFactory.from_root(
           fname,
           schemaclass=DummySchema
         ).events()
print(events.fields)
print(events.Jet.fields)
print(events.Jet.energy[2])
print(events.Jet.Subjets.energy[2])

['Electron', 'Muon', 'Jet']
['pt', 'eta', 'phi', 'energy', 'ID', 'SubjetsCounts', 'Subjets']
[51.8, 47.9, 56]
[[17.4], [22.5, 24.7], [28.3, 23]]

Inherit from NanoEvents

So far we showed how to create customer collections and behaviors. Actually, if you don’t want to write a new schema, you can name the TBranch with the same naming convention as NanoAOD and use the NanoAODSchema. Then it will be recognized automatically.

Looking at the content of NanoAOD file, you will need to follow the same naming if you want to use NanoAODSchema.

By doing so, you can also make use of the behaviors in NanoAODSchema automatically, which is very convenient.