{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Writting coffea Schema files\n", "\n", "The interpretation of the TTree data is configurable via schema objects.\n", "\n", "Schema teaches the event processor how to group variables into collections, so\n", "operations can be run over entire collection at once. And we can also define\n", "some handy\n", "[behaviors](https://awkward-array.readthedocs.io/en/latest/ak.behavior.html)\n", "for a specific collection in a schema.\n", "\n", "In this demo, we will create our own dummy schema and implement our own dummy\n", "behavior.\n", "\n", "Let's look at the content of dummy_nanoevents.root file with `BaseSchema`.\n", "We'll show how to construct a schema that can be used to interpret this root\n", "file later.\n", "\n", "The events object can be instantiated as follows:\n" ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "scrolled": true }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "['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']\n" ] } ], "source": [ "from coffea.nanoevents import NanoEventsFactory, BaseSchema, NanoAODSchema\n", "fname = \"../../samples/dummy_nanoevents.root\"\n", "events = NanoEventsFactory.from_root(\n", " fname,\n", " schemaclass=BaseSchema\n", " ).events()\n", "print(events.fields)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now we can copy the skeleton of a schema class:" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "class DummySchema(BaseSchema):\n", " \"\"\"\n", " \"\"\"\n", " def __init__(self, base_form):\n", " super().__init__(base_form)\n", " self._form[\"contents\"] = self._build_collections(self._form[\"contents\"])\n", "\n", " def _build_collections(self, branch_forms):\n", " output = {}\n", " return output\n", "\n", " @property\n", " def behavior(self):\n", " from coffea.nanoevents.methods import base, vector\n", " behavior = {}\n", " behavior.update(base.behavior)\n", " behavior.update(vector.behavior)\n", " return behavior" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "As you can see, this schema is so simple it is not useful currently. If we\n", "call the events again with our own schema, we'll find it contains nothing." ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "[]" ] }, "execution_count": 3, "metadata": {}, "output_type": "execute_result" } ], "source": [ "events = NanoEventsFactory.from_root(\n", " fname,\n", " schemaclass=DummySchema\n", " ).events()\n", "events.fields" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Create collections\n", "\n", "In schema, the `branch_forms` is a python dictionary used to define branch\n", "grouping.\n", "\n", "By default, using `BaseSchema`, it will be completely flat:\n", "\n", "```python\n", "branch_forms={\n", " \"particle_pt\":{},\n", " \"particle_eta\":{},\n", " \"particle_phi\":{},\n", " \"particle_mass\":{},\n", " ...\n", "}\n", "```\n", "\n", "Just as you have seen when we open the dummy root file with `BaseSchema`, all\n", "the branches are listed when do `print(events.fields)`. These branches are\n", "independent to each other. For example when you do a selection on\n", "`particle_pt`, the branch `particle_eta` will not be affected.\n", "\n", "We typically like to do selections collectively. If some particles don't pass\n", "the selection on `particle_pt`, their `particle_eta` should also be eliminated\n", "from the list.\n", "\n", "I.e. We would like to put some branches into the same collection, as what\n", "follows:\n", "\n", "```python\n", "new_branch_forms={\n", " \"particle\": schemas.zip_forms({\n", " \"pt\" : branch_forms[\"particle_pt\"],\n", " \"eta\" : branch_forms[\"particle_eta\"],\n", " \"phi\" : branch_forms[\"particle_phi\"],\n", " \"mass\" : branch_forms[\"particle_mass\"],\n", " })\n", "}\n", "```\n", "\n", "So when we want to call `particle_pt`, we actually do `particle.pt`. And we can\n", "apply selections on the `particle` level so `particle_pt` and `particle_eta`\n", "are selected at the same time now.\n", "\n", "All of this is to be implemented in the `Schema._build_collections` method.\n", "\n", "For example, let's add the `Electron` collection to our schema. To do this we\n", "need to import `zip_forms`. " ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [], "source": [ "from coffea.nanoevents.schemas import zip_forms\n", "class DummySchema(BaseSchema):\n", " \"\"\"\n", " \"\"\"\n", " def __init__(self, base_form):\n", " super().__init__(base_form)\n", " self._form[\"contents\"] = self._build_collections(self._form[\"contents\"])\n", "\n", " def _build_collections(self, branch_forms):\n", " output = {}\n", " output[\"Electron\"] = zip_forms(\n", " {\n", " \"pt\" : branch_forms[\"Electron_pt\"],\n", " \"eta\" : branch_forms [\"Electron_eta\"] ,\n", " \"phi\": branch_forms[\"Electron_phi\"],\n", " \"mass\": branch_forms[\"Electron_mass\"],\n", " #\"xx\": branch_forms[\"Electron_xx\"],\n", " },\n", " \"Electron\",\n", " )\n", " return output\n", "\n", " @property\n", " def behavior(self):\n", " from coffea.nanoevents.methods import base, vector\n", " behavior = {}\n", " behavior.update(base.behavior)\n", " behavior.update(vector.behavior)\n", " return behavior" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now we successfully created a schema with one collection `Electron`. It will be\n", "able to recognize branches with name `Electron_pt, Electron_eta, Electron_phi,\n", "Electron_mass`. Try to call the `events` again." ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "['Electron']\n", "['pt', 'eta', 'phi', 'mass']\n" ] } ], "source": [ "events = NanoEventsFactory.from_root(\n", " fname,\n", " schemaclass=DummySchema\n", " ).events()\n", "print(events.fields)\n", "print(events.Electron.fields)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Congratulations, we have now collected the fields related to an electron into a\n", "common collection! We can use the mask and do selection on the whole collection\n", "at once now:" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[[49.3, 38.1], [48.9, 43.9, 50.8, 43.9], ... 40.4], [44.6, 45.5, 49.2, 58.3, 51.9]]\n", "[[0.366, 0.437], [0.149, 0.236, 0.472, ... [0.442, 0.854, 0.344, 0.156, 0.564]]\n" ] } ], "source": [ "mask = (events.Electron.pt>3) & (events.Electron.pt<60)\n", "good_elec = events.Electron[mask]\n", "print(good_elec.pt)\n", "print(good_elec.eta)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "However, if you put some unknown branches to the collection, errors will be returned.\n", "For example, uncomment the following line in `DummySchema`:\n", "```python\n", "\"xx\": branch_forms[\"Electron_xx\"],\n", "```\n", "Run the above code again, you will see:\n", "```bash\n", "KeyError: 'Electron_xx'\n", "```\n", "\n", "Of course, we can make a long list and manually group all the TBranches. But\n", "use naming conventions would enable you to write the schema in a very neat way.\n", "Like what we have in the dummy root file, branches are named as\n", "`object_varible`. So we can define some collections and put the corresponding\n", "TBranches into them.\n", "\n", "4 collections are defined below in `mixins`:" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [], "source": [ "class DummySchema(BaseSchema):\n", " \"\"\"\n", " \"\"\"\n", " mixins = {\n", " 'Electron': \"PtEtaPhiMLorentzVector\",\n", " 'Muon': 'PtEtaPhiMLorentzVector',\n", " 'Jet': 'PtEtaPhiELorentzVector',\n", " 'Subjet': 'PtEtaPhiELorentzVector',\n", " }\n", " def __init__(self, base_form):\n", " super().__init__(base_form)\n", " self._form[\"contents\"] = self._build_collections(self._form[\"contents\"])\n", "\n", " def _build_collections(self, branch_forms):\n", " output = {}\n", " ## Making the basic\n", " for name in self.mixins:\n", " mixin = self.mixins.get(name, \"NanoCollection\")\n", " output[name] = zip_forms({\n", " k[len(name) + 1 :]: branch_forms[k]\n", " for k in branch_forms\n", " if k.startswith(name + \"_\")\n", " },\n", " name,\n", " record_name=mixin)\n", " return output\n", "\n", " @property\n", " def behavior(self):\n", " from coffea.nanoevents.methods import base, vector\n", " behavior = {}\n", " behavior.update(base.behavior)\n", " behavior.update(vector.behavior)\n", " return behavior" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We defined `Electron, Muon, Jet, Subjet` in the above DummySchema, and we used\n", "`for k in branch_forms` to search all the branches start with `{collection}_`\n", "and put them into the corresponding collections. Now open the dummy root file\n", "again, we can see the defined collections." ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "['Electron', 'Muon', 'Jet', 'Subjet']\n", "['pt', 'eta', 'phi', 'mass', 'charge', 'flag', 'dxy', 'dxyErr', 'dz', 'dzErr']\n", "[[48.7, 33.9], [26.6, 40.2, 29.8, 31.6], ... 47.3, 24.6], [33, 30.2, 47.8, 31.5, 45]]\n", "[[1.21, 1.15], [1.42, 1.34, 1.12, 1.47], ... 1.56], [1.14, 0.805, 1.23, 1.42, 1.03]]\n" ] } ], "source": [ "events = NanoEventsFactory.from_root(\n", " fname,\n", " schemaclass=DummySchema\n", " ).events()\n", "print(events.fields)\n", "print(events.Electron.fields)\n", "print(events.Electron.px)\n", "print(events.Electron.theta)" ] }, { "cell_type": "markdown", "id": "d4441136", "metadata": { "lines_to_next_cell": 0 }, "source": [ " Maybe you already noticed that we printed something called\n", "`events.Electron.px` and `events.Electron.theta` in the previous block, which\n", "don't exist in the dummy root file at all. How were they created? This is\n", "actually the `behavior` this collection has, which we'll talk about right now.\n", "\n", "Look at the collection list:" ] }, { "cell_type": "code", "execution_count": null, "id": "662ba9f8", "metadata": { "lines_to_next_cell": 0 }, "outputs": [], "source": [ "mixins = {\n", " 'Electron': \"PtEtaPhiMLorentzVector\",\n", " 'Muon': 'PtEtaPhiMLorentzVector',\n", " 'Jet': 'PtEtaPhiELorentzVector',\n", " 'Subjet': 'PtEtaPhiELorentzVector',\n", " }" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "`PtEtaPhiMLorentzVector` is the name of `behavior` that each collection has. It\n", "is defined\n", "[here](https://github.com/yihui-lai/coffea/blob/351cc727845ab83a8e31a193dc06e534bedb97fe/coffea/nanoevents/methods/vector.py#L497).\n", "And we imported it through `from coffea.nanoevents.methods import base, vector`" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n", "\n", "## Create behavior\n", "\n", "Aside from putting different branches into collections, we can also add\n", "`behavior` to collections. This means additional awkward arrays are generated\n", "on-the-fly via predefined algorithm. Like we can get `px, theta` previously.\n", "\n", "A bunch of other common physics behaviors are already provided in coffea, and\n", "you can find them in\n", "[methods](https://github.com/CoffeaTeam/coffea/tree/a95401cad91e88ceac47a4c693068bc4cbc7d338/coffea/nanoevents/methods).\n", "\n", "To write our own coffea behavior, we need to define the `behavior` class. In\n", "the following code, we create `DummyBehavior`, which has one function\n", "`plus1()`. It returns the `particle.mass+1` when you call `particle.plus1`.\n" ] }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [], "source": [ "import awkward1\n", "dummybehavior={}\n", "@awkward1.mixin_class(dummybehavior)\n", "class DummyBehavior:\n", " @property\n", " def plus1(self):\n", " return self.mass+1\n", "\n", "class DummySchema(BaseSchema):\n", " \"\"\"\n", " \"\"\"\n", " mixins = {\n", " 'Electron': \"DummyBehavior\",\n", " 'Muon': 'PtEtaPhiMLorentzVector',\n", " 'Jet': 'PtEtaPhiELorentzVector',\n", " 'Subjet': 'PtEtaPhiELorentzVector',\n", " }\n", " def __init__(self, base_form):\n", " super().__init__(base_form)\n", " self._form[\"contents\"] = self._build_collections(self._form[\"contents\"])\n", "\n", " def _build_collections(self, branch_forms):\n", " output = {}\n", " ## Making the basic\n", " for name in self.mixins:\n", " mixin = self.mixins.get(name, \"NanoCollection\")\n", " output[name] = zip_forms({\n", " k[len(name) + 1 :]: branch_forms[k]\n", " for k in branch_forms\n", " if k.startswith(name + \"_\")\n", " },\n", " name,\n", " record_name=mixin)\n", " return output\n", "\n", " @property\n", " def behavior(self):\n", " from coffea.nanoevents.methods import base, vector\n", " behavior = {}\n", " behavior.update(base.behavior)\n", " behavior.update(vector.behavior)\n", " behavior.update(dummybehavior)\n", " return behavior" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The `behavior` of `Electron` is changed to `DummyBehavior` in the above\n", "DummySchema. Also the `dummybehavior` is included in the `DummySchema` through\n", "` behavior.update(dummybehavior)`. Now try our customer behavior:" ] }, { "cell_type": "code", "execution_count": 13, "metadata": { "scrolled": true }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "['Electron', 'Muon', 'Jet', 'Subjet']\n", "['pt', 'eta', 'phi', 'mass', 'charge', 'flag', 'dxy', 'dxyErr', 'dz', 'dzErr']\n", "[0.00051, 0.00051]\n", "[1, 1]\n" ] } ], "source": [ "events = NanoEventsFactory.from_root(\n", " fname,\n", " schemaclass=DummySchema\n", " ).events()\n", "print(events.fields)\n", "print(events.Electron.fields)\n", "print(events.Electron.mass[0])\n", "print(events.Electron.plus1[0])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Python did the rounding so the output is not exactly 1.00051.\n", "\n", "Since we changed the `behavior`, `print(events.Electron.theta)` should not work\n", "now." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "print(events.Electron.theta)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Nested jagged array implementation\n", "\n", "A jagged array is an array of arrays. Each array is not guaranteed to be of the\n", "same size. You could have one parent array that each of its component is\n", "another array.\n", "\n", "But why do we need such kind of data structure? It will be clear after we take\n", "a look at the dummy_nanoevents.cc file. For convenience, I pasted the related\n", "codes below:\n", "\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "```C++\n", "#include \"TFile.h\"\n", "#include \"TRandom.h\"\n", "#include \"TTree.h\"\n", "\n", "void\n", "dummy_nanoevent()\n", "{\n", " UInt_t run, event, luminosityBlock;\n", "\n", " std::vector Jets_Pt;\n", " std::vector Jets_Eta;\n", " std::vector Jets_Phi;\n", " std::vector Jets_E;\n", " std::vector Jets_ID;\n", " std::vector Jets_SubjetsCounts;\n", " std::vector Jets_subjet_Pt;\n", " std::vector Jets_subjet_Eta;\n", " std::vector Jets_subjet_Phi;\n", " std::vector Jets_subjet_E;\n", "\n", " TFile* Tfile = Tfile = TFile::Open( \"dummy_nanoevents.root\", \"RECREATE\" );\n", " TTree* ttree = new TTree( \"Events\", \"\" );\n", "\n", " ttree->Branch( \"run\", &run, \"run/I\" );\n", " ttree->Branch( \"luminosityBlock\", &luminosityBlock, \"luminosityBlock/I\" );\n", " ttree->Branch( \"event\", &event, \"event/I\" );\n", "\n", " ttree->Branch( \"Jet_pt\", &Jets_Pt );\n", " ttree->Branch( \"Jet_eta\", &Jets_Eta );\n", " ttree->Branch( \"Jet_phi\", &Jets_Phi );\n", " ttree->Branch( \"Jet_energy\", &Jets_E );\n", " ttree->Branch( \"Jet_ID\", &Jets_ID );\n", " ttree->Branch( \"Jet_SubjetsCounts\", &Jets_SubjetsCounts );\n", " ttree->Branch( \"Subjet_pt\", &Jets_subjet_Pt );\n", " ttree->Branch( \"Subjet_eta\", &Jets_subjet_Eta );\n", " ttree->Branch( \"Subjet_phi\", &Jets_subjet_Phi );\n", " ttree->Branch( \"Subjet_energy\", &Jets_subjet_E );\n", "\n", " for( Int_t ev = 0; ev < 100; ev++ ){\n", " run = 1;\n", " event = ev;\n", " luminosityBlock = 1000 * ev;\n", " Int_t njet = Int_t( 3 * gRandom->Rndm() + 1 );\n", "\n", " Jets_Pt.clear();\n", " Jets_Eta.clear();\n", " Jets_Phi.clear();\n", " Jets_E.clear();\n", " Jets_ID.clear();\n", " Jets_SubjetsCounts.clear();\n", " Jets_subjet_Pt.clear();\n", " Jets_subjet_Eta.clear();\n", " Jets_subjet_Phi.clear();\n", " Jets_subjet_E.clear();\n", "\n", " for( int i = 0; i < njet; i++ ){\n", " Jets_Pt.push_back( 10 * gRandom->Rndm() );\n", " Jets_Eta.push_back( gRandom->Rndm() );\n", " Jets_Phi.push_back( gRandom->Rndm() );\n", " Jets_E.push_back( gRandom->Gaus( 50, 10 ) );\n", " Jets_ID.push_back( int( 7*gRandom->Rndm() ) );\n", " // subjets\n", " Int_t jets_sub = Int_t( 3 * gRandom->Rndm() );\n", " Jets_SubjetsCounts.push_back( jets_sub );\n", "\n", " for( int i = 0; i < jets_sub; i++ ){\n", " Jets_subjet_Pt.push_back( 10 * gRandom->Rndm() );\n", " Jets_subjet_Eta.push_back( gRandom->Rndm() );\n", " Jets_subjet_Phi.push_back( gRandom->Rndm() );\n", " Jets_subjet_E.push_back( gRandom->Gaus( 25, 10 ) );\n", " }\n", " }\n", "\n", " ttree->Fill();\n", " }\n", "\n", " ttree->Write();\n", "}\n", "```\n", "\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n", "In the file, we generated `{njet}` jets and each jet has `{jets_sub}` subjets. So the structure looks like this:\n", "```python\n", "'jet' : {\n", " 'pt' : [n],\n", " 'eta' : [n],\n", " 'phi' : [n],\n", " 'energy' : [n],\n", " 'ID' : [n],\n", " 'subjets' : [\n", " {'pt':[m0], 'eta':[m0], 'phi':[m0], 'energy':[m0]}, # 0\n", " {'pt':[m1], 'eta':[m1], 'phi':[m1], 'energy':[m1]}, # 1\n", " ...\n", " {'pt':[m], 'eta':[m], 'phi':[m], 'energy':[m]}, # n-1\n", " ]\n", "}\n", "```\n", "\n", "We have `n` jets here and each jet has `m` subjets. `m` does not need to be the\n", "same for all the jets, that's where the name `jagged` comes. `jet.pt` returns\n", "array with `n` numbers (`[0, 1, ... n-1]`), while `jet.subjets.pt` returns `n`\n", "arrays (`[[m0], [m1], ... [m]]`).\n", "\n", "Note that when we create the TBranches, we can define it to store vector of\n", "vectors `vector< vector >`, this is intuitively the same structure as\n", "we shown above. However, Coffea doesn't recognize this structure. If you have\n", "this type of branch in the root file, you will need to make modifications to\n", "the generation codes.\n", "\n", "We can store the jets as a flat vector (length `n`) and use another flat vector\n", "to store all the subjets (length `N`). Since one jet at least has one subjet,\n", "`N>=n` should be fulfilled. Additionally, we need another Branch\n", "`Jet_SubjetsCounts` (with length `n`) to tell us how many subjets that each jet\n", "contains.\n", "\n", "So if we have the following vectors:\n", "```\n", "Jets_energy = [10, 10.1, 10.2];\n", "Jet_SubjetsCounts = [1, 3, 2];\n", "Jets_subjet_energy = [10, 1.1, 5, 4, 5.1, 5.1];\n", "```\n", "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}]`\n", "\n", "So how to implement this structure in coffea? It is actually very easy because\n", "we only need to `import nest_jagged_forms` and follow the rules of this\n", "predefined function: " ] }, { "cell_type": "code", "execution_count": 14, "metadata": {}, "outputs": [], "source": [ "from coffea.nanoevents.schemas import nest_jagged_forms\n", "\n", "dummybehavior={}\n", "@awkward1.mixin_class(dummybehavior)\n", "class DummyBehavior:\n", " @property\n", " def plus1(self):\n", " return self.mass+1\n", "\n", "class DummySchema(BaseSchema):\n", " \"\"\"\n", " \"\"\"\n", " mixins = {\n", " 'Electron': \"DummyBehavior\",\n", " 'Muon': 'PtEtaPhiMLorentzVector',\n", " 'Jet': 'PtEtaPhiELorentzVector',\n", " 'Subjet': 'PtEtaPhiELorentzVector',\n", " }\n", " def __init__(self, base_form):\n", " super().__init__(base_form)\n", " self._form[\"contents\"] = self._build_collections(self._form[\"contents\"])\n", "\n", " def _build_collections(self, branch_forms):\n", " output = {}\n", " ## Making the basic\n", " for name in self.mixins:\n", " mixin = self.mixins.get(name, \"NanoCollection\")\n", " output[name] = zip_forms({\n", " k[len(name) + 1 :]: branch_forms[k]\n", " for k in branch_forms\n", " if k.startswith(name + \"_\")\n", " },\n", " name,\n", " record_name=mixin)\n", "\n", " nest_jagged_forms(output['Jet'],\n", " output.pop('Subjet'),\n", " 'SubjetsCounts',\n", " 'Subjets')\n", " return output\n", "\n", " @property\n", " def behavior(self):\n", " from coffea.nanoevents.methods import base, vector\n", " behavior = {}\n", " behavior.update(base.behavior)\n", " behavior.update(vector.behavior)\n", " behavior.update(dummybehavior)\n", " return behavior" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ " i\n", "[`nest_jagged_forms`](https://github.com/yihui-lai/coffea/blob/351cc727845ab83a8e31a193dc06e534bedb97fe/coffea/nanoevents/schemas/base.py#L62)\n", "is defined as follows. The first input parameter is the parent array, the\n", "second parameter is the child array. Then the third one is the counts array.\n", "The final one is the new name for the child array." ] }, { "cell_type": "code", "execution_count": null, "id": "b00a1e9f", "metadata": { "lines_to_next_cell": 2 }, "outputs": [], "source": [ "def nest_jagged_forms(parent, child, counts_name, name):\n", " \"\"\"Place child listarray inside parent listarray as a double-jagged array\"\"\"\n", " if not parent[\"class\"].startswith(\"ListOffsetArray\"):\n", " raise ValueError\n", " if parent[\"content\"][\"class\"] != \"RecordArray\":\n", " raise ValueError\n", " if not child[\"class\"].startswith(\"ListOffsetArray\"):\n", " raise ValueError\n", " counts = parent[\"content\"][\"contents\"][counts_name]\n", " offsets = transforms.counts2offsets_form(counts)\n", " inner = listarray_form(child[\"content\"], offsets)\n", " parent[\"content\"][\"contents\"][name] = inner" ] }, { "cell_type": "markdown", "id": "731904cd", "metadata": {}, "source": [ "We can try this DummySchema now:\n", "" ] }, { "cell_type": "code", "execution_count": 15, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "['Electron', 'Muon', 'Jet']\n", "['pt', 'eta', 'phi', 'energy', 'ID', 'SubjetsCounts', 'Subjets']\n", "[51.8, 47.9, 56]\n", "[[17.4], [22.5, 24.7], [28.3, 23]]\n" ] } ], "source": [ "events = NanoEventsFactory.from_root(\n", " fname,\n", " schemaclass=DummySchema\n", " ).events()\n", "print(events.fields)\n", "print(events.Jet.fields)\n", "print(events.Jet.energy[2])\n", "print(events.Jet.Subjets.energy[2])" ] }, { "cell_type": "markdown", "metadata": { "lines_to_next_cell": 0 }, "source": [ "## Inherit from NanoEvents\n", "\n", "So far we showed how to create customer collections and behaviors. Actually, if\n", "you don't want to write a new schema, you can name the TBranch with the same\n", "naming convention as NanoAOD and use the `NanoAODSchema`. Then it will be\n", "recognized automatically.\n", "\n", "Looking at the content of [NanoAOD\n", "file](https://cms-nanoaod-integration.web.cern.ch/integration/master-cmsswmaster/mc102X_doc.html),\n", "you will need to follow the same naming if you want to use `NanoAODSchema`.\n", "\n", "By doing so, you can also make use of the `behaviors` in `NanoAODSchema`\n", "automatically, which is very convenient. " ] }, { "cell_type": "markdown", "id": "71b8b8f2", "metadata": {}, "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.8.1" } }, "nbformat": 4, "nbformat_minor": 2 }