diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 628df33dd..fe8453b8d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -119,7 +119,7 @@ jobs: touch build/html/.nojekyll - name: Deploy documentation if: github.event_name == 'push' && matrix.os == 'ubuntu-latest' && matrix.python-version == 3.11 - uses: crazy-max/ghaction-github-pages@v3 + uses: crazy-max/ghaction-github-pages@v4 with: target_branch: gh-pages build_dir: docs/build/html diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a4d511b07..d535cd79e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -12,7 +12,7 @@ ci: repos: - repo: https://github.com/psf/black - rev: 23.7.0 + rev: 23.9.1 hooks: - id: black @@ -37,7 +37,7 @@ repos: - id: trailing-whitespace - repo: https://github.com/asottile/pyupgrade - rev: v3.9.0 + rev: v3.11.0 hooks: - id: pyupgrade args: ["--py38-plus"] @@ -48,7 +48,7 @@ repos: - id: setup-cfg-fmt - repo: https://github.com/pycqa/flake8 - rev: 6.0.0 + rev: 6.1.0 hooks: - id: flake8 exclude: coffea/processor/templates diff --git a/src/coffea/analysis_tools.py b/src/coffea/analysis_tools.py index 66b92fe2b..facf14e97 100644 --- a/src/coffea/analysis_tools.py +++ b/src/coffea/analysis_tools.py @@ -418,7 +418,7 @@ def variations(self): class NminusOneToNpz: - """Object to be returned by NmiusOne.to_npz()""" + """Object to be returned by NminusOne.to_npz()""" def __init__(self, file, labels, nev, masks, saver): self._file = file @@ -494,11 +494,17 @@ def maskscutflow(self): return self._maskscutflow def compute(self): - self._nevonecut = list(dask.compute(*self._nevonecut)) - self._nevcutflow = list(dask.compute(*self._nevcutflow)) - self._masksonecut = list(dask.compute(*self._masksonecut)) - self._maskscutflow = list(dask.compute(*self._maskscutflow)) - numpy.savez( + self._nevonecut, self._nevcutflow = dask.compute( + self._nevonecut, self._nevcutflow + ) + self._masksonecut, self._maskscutflow = dask.compute( + self._masksonecut, self._maskscutflow + ) + self._nevonecut = list(self._nevonecut) + self._nevcutflow = list(self._nevcutflow) + self._masksonecut = list(self._masksonecut) + self._maskscutflow = list(self._maskscutflow) + self._saver( self._file, labels=self._labels, nevonecut=self._nevonecut, @@ -538,7 +544,7 @@ def result(self): labels = ["initial"] + [f"N - {i}" for i in self._names] + ["N"] return NminusOneResult(labels, self._nev, self._masks) - def to_npz(self, file, compressed=False, compute=True): + def to_npz(self, file, compressed=False, compute=False): """Saves the results of the N-1 selection to a .npz file Parameters @@ -554,7 +560,7 @@ def to_npz(self, file, compressed=False, compute=True): compute : bool, optional Whether to immediately start writing or to return an object that the user can choose when to start writing by calling compute(). - Default is True. + Default is False. Returns ------- @@ -580,22 +586,29 @@ def print(self): """Prints the statistics of the N-1 selection""" if self._delayed_mode: + warnings.warn( + "Printing the N-1 selection statistics is going to compute dask_awkward objects." + ) self._nev = list(dask.compute(*self._nev)) + nev = self._nev print("N-1 selection stats:") for i, name in enumerate(self._names): - print( - f"Ignoring {name:<20}: pass = {nev[i+1]:<20}\ - all = {nev[0]:<20}\ - -- eff = {nev[i+1]*100/nev[0]:.1f} %" + stats = ( + f"Ignoring {name:<20}" + f"pass = {nev[i+1]:<20}" + f"all = {nev[0]:<20}" + f"-- eff = {nev[i+1]*100/nev[0]:.1f} %" ) + print(stats) - if True: - print( - f"All cuts {'':<20}: pass = {nev[-1]:<20}\ - all = {nev[0]:<20}\ - -- eff = {nev[-1]*100/nev[0]:.1f} %" - ) + stats_all = ( + f"All cuts {'':<20}" + f"pass = {nev[-1]:<20}" + f"all = {nev[0]:<20}" + f"-- eff = {nev[-1]*100/nev[0]:.1f} %" + ) + print(stats_all) def yieldhist(self): """Returns the N-1 selection yields as a ``hist.Hist`` object @@ -610,13 +623,13 @@ def yieldhist(self): labels = ["initial"] + [f"N - {i}" for i in self._names] + ["N"] if not self._delayed_mode: h = hist.Hist(hist.axis.Integer(0, len(labels), name="N-1")) - h.fill(numpy.arange(len(labels)), weight=self._nev) + h.fill(numpy.arange(len(labels), dtype=int), weight=self._nev) else: h = hist.dask.Hist(hist.axis.Integer(0, len(labels), name="N-1")) for i, weight in enumerate(self._masks, 1): h.fill(dask_awkward.full_like(weight, i, dtype=int), weight=weight) - h.fill(dask_awkward.zeros_like(weight)) + h.fill(dask_awkward.zeros_like(weight, dtype=int)) return h, labels @@ -712,7 +725,7 @@ def plot_vars( hist.axis.Integer(0, len(labels), name="N-1"), ) arr = awkward.flatten(var) - h.fill(arr, awkward.zeros_like(arr)) + h.fill(arr, awkward.zeros_like(arr, dtype=int)) for i, mask in enumerate(self.result().masks, 1): arr = awkward.flatten(var[mask]) h.fill(arr, awkward.full_like(arr, i, dtype=int)) @@ -725,7 +738,7 @@ def plot_vars( hist.axis.Integer(0, len(labels), name="N-1"), ) arr = dask_awkward.flatten(var) - h.fill(arr, dask_awkward.zeros_like(arr)) + h.fill(arr, dask_awkward.zeros_like(arr, dtype=int)) for i, mask in enumerate(self.result().masks, 1): arr = dask_awkward.flatten(var[mask]) h.fill(arr, dask_awkward.full_like(arr, i, dtype=int)) @@ -780,7 +793,7 @@ def result(self): self._maskscutflow, ) - def to_npz(self, file, compressed=False, compute=True): + def to_npz(self, file, compressed=False, compute=False): """Saves the results of the cutflow to a .npz file Parameters @@ -796,7 +809,7 @@ def to_npz(self, file, compressed=False, compute=True): compute : bool, optional Whether to immediately start writing or to return an object that the user can choose when to start writing by calling compute(). - Default is True. + Default is False. Returns ------- @@ -824,19 +837,27 @@ def print(self): """Prints the statistics of the Cutflow""" if self._delayed_mode: - self._nevonecut = list(dask.compute(*self._nevonecut)) - self._nevcutflow = list(dask.compute(*self._nevcutflow)) + warnings.warn( + "Printing the cutflow statistics is going to compute dask_awkward objects." + ) + self._nevonecut, self._nevcutflow = dask.compute( + self._nevonecut, self._nevcutflow + ) + nevonecut = self._nevonecut nevcutflow = self._nevcutflow + print("Cutflow stats:") for i, name in enumerate(self._names): - print( - f"Cut {name:<20}: pass = {nevonecut[i+1]:<20}\ - cumulative pass = {nevcutflow[i+1]:<20}\ - all = {nevonecut[0]:<20}\ - -- eff = {nevonecut[i+1]*100/nevonecut[0]:.1f} %\ - -- cumulative eff = {nevcutflow[i+1]*100/nevcutflow[0]:.1f} %" + stats = ( + f"Cut {name:<20}:" + f"pass = {nevonecut[i+1]:<20}" + f"cumulative pass = {nevcutflow[i+1]:<20}" + f"all = {nevonecut[0]:<20}" + f"-- eff = {nevonecut[i+1]*100/nevonecut[0]:.1f} %{'':<20}" + f"-- cumulative eff = {nevcutflow[i+1]*100/nevcutflow[0]:.1f} %" ) + print(stats) def yieldhist(self): """Returns the cutflow yields as ``hist.Hist`` objects @@ -856,8 +877,8 @@ def yieldhist(self): honecut = hist.Hist(hist.axis.Integer(0, len(labels), name="onecut")) hcutflow = honecut.copy() hcutflow.axes.name = ("cutflow",) - honecut.fill(numpy.arange(len(labels)), weight=self._nevonecut) - hcutflow.fill(numpy.arange(len(labels)), weight=self._nevcutflow) + honecut.fill(numpy.arange(len(labels), dtype=int), weight=self._nevonecut) + hcutflow.fill(numpy.arange(len(labels), dtype=int), weight=self._nevcutflow) else: honecut = hist.dask.Hist(hist.axis.Integer(0, len(labels), name="onecut")) @@ -868,12 +889,12 @@ def yieldhist(self): honecut.fill( dask_awkward.full_like(weight, i, dtype=int), weight=weight ) - honecut.fill(dask_awkward.zeros_like(weight)) + honecut.fill(dask_awkward.zeros_like(weight, dtype=int)) for i, weight in enumerate(self._maskscutflow, 1): hcutflow.fill( dask_awkward.full_like(weight, i, dtype=int), weight=weight ) - hcutflow.fill(dask_awkward.zeros_like(weight)) + hcutflow.fill(dask_awkward.zeros_like(weight, dtype=int)) return honecut, hcutflow, labels @@ -975,8 +996,8 @@ def plot_vars( hcutflow.axes.name = name, "cutflow" arr = awkward.flatten(var) - honecut.fill(arr, awkward.zeros_like(arr)) - hcutflow.fill(arr, awkward.zeros_like(arr)) + honecut.fill(arr, awkward.zeros_like(arr, dtype=int)) + hcutflow.fill(arr, awkward.zeros_like(arr, dtype=int)) for i, mask in enumerate(self.result().masksonecut, 1): arr = awkward.flatten(var[mask]) @@ -998,8 +1019,8 @@ def plot_vars( hcutflow.axes.name = name, "cutflow" arr = dask_awkward.flatten(var) - honecut.fill(arr, dask_awkward.zeros_like(arr)) - hcutflow.fill(arr, dask_awkward.zeros_like(arr)) + honecut.fill(arr, dask_awkward.zeros_like(arr, dtype=int)) + hcutflow.fill(arr, dask_awkward.zeros_like(arr, dtype=int)) for i, mask in enumerate(self.result().masksonecut, 1): arr = dask_awkward.flatten(var[mask]) diff --git a/src/coffea/nanoevents/methods/physlite.py b/src/coffea/nanoevents/methods/physlite.py index c0efcdc39..751c5d03f 100644 --- a/src/coffea/nanoevents/methods/physlite.py +++ b/src/coffea/nanoevents/methods/physlite.py @@ -2,6 +2,7 @@ from numbers import Number import awkward +import dask_awkward import numpy from coffea.nanoevents.methods import base, vector @@ -38,7 +39,30 @@ def _element_link(target_collection, eventindex, index, key): return target_collection._apply_global_index(global_index) +def _element_link_method(self, link_name, target_name, _dask_array_): + if _dask_array_ is not None: + target = _dask_array_.behavior["__original_array__"]()[target_name] + links = _dask_array_[link_name] + return _element_link( + target, + _dask_array_._eventindex, + links.m_persIndex, + links.m_persKey, + ) + links = self[link_name] + return _element_link( + self._events()[target_name], + self._eventindex, + links.m_persIndex, + links.m_persKey, + ) + + def _element_link_multiple(events, obj, link_field, with_name=None): + # currently not working in dask because: + # - we don't know the resulting type beforehand + # - also not the targets, so no way to find out which columns to load? + # - could consider to treat the case of truth collections by just loading all truth columns link = obj[link_field] key = link.m_persKey index = link.m_persIndex @@ -64,22 +88,46 @@ def where(unique_keys): return out -def _get_target_offsets(offsets, event_index): +def _get_target_offsets(load_column, event_index): + if isinstance(load_column, dask_awkward.Array) and isinstance( + event_index, dask_awkward.Array + ): + # wrap in map_partitions if dask arrays + return dask_awkward.map_partitions( + _get_target_offsets, load_column, event_index + ) + + offsets = load_column.layout.offsets.data + if isinstance(event_index, Number): return offsets[event_index] + # let the necessary column optimization know that we need to load this + # column to get the offsets + if awkward.backend(load_column) == "typetracer": + awkward.typetracer.touch_data(load_column) + + # necessary to stick it into the `NumpyArray` constructor + # if typetracer is passed through + offsets = awkward.typetracer.length_zero_if_typetracer( + load_column.layout.offsets.data + ) + def descend(layout, depth, **kwargs): if layout.purelist_depth == 1: return awkward.contents.NumpyArray(offsets)[layout] - return awkward.transform(descend, event_index) + return awkward.transform(descend, event_index.layout) def _get_global_index(target, eventindex, index): - load_column = target[ - target.fields[0] - ] # awkward is eager-mode now (will need to dask this) - target_offsets = _get_target_offsets(load_column.layout.offsets, eventindex) + for field in target.fields: + # fetch first column to get offsets from + # (but try to avoid the double-jagged ones if possible) + load_column = target[field] + if load_column.ndim < 3: + break + target_offsets = _get_target_offsets(load_column, eventindex) return target_offsets + index @@ -140,12 +188,12 @@ class Muon(Particle): """ @property - def trackParticle(self): - return _element_link( - self._events().CombinedMuonTrackParticles, - self._eventindex, - self["combinedTrackParticleLink.m_persIndex"], - self["combinedTrackParticleLink.m_persKey"], + def trackParticle(self, _dask_array_=None): + return _element_link_method( + self, + "combinedTrackParticleLink", + "CombinedMuonTrackParticles", + _dask_array_, ) @@ -159,21 +207,25 @@ class Electron(Particle): """ @property - def trackParticles(self): - links = self.trackParticleLinks - return _element_link( - self._events().GSFTrackParticles, - self._eventindex, - links.m_persIndex, - links.m_persKey, + def trackParticles(self, _dask_array_=None): + return _element_link_method( + self, "trackParticleLinks", "GSFTrackParticles", _dask_array_ ) @property - def trackParticle(self): - trackParticles = self.trackParticles - return self.trackParticles[ - tuple([slice(None) for i in range(trackParticles.ndim - 1)] + [0]) - ] + def trackParticle(self, _dask_array_=None): + trackParticles = _element_link_method( + self, "trackParticleLinks", "GSFTrackParticles", _dask_array_ + ) + # Ellipsis (..., 0) slicing not supported yet by dask_awkward + slicer = tuple([slice(None) for i in range(trackParticles.ndim - 1)] + [0]) + return trackParticles[slicer] + + @property + def caloClusters(self, _dask_array_=None): + return _element_link_method( + self, "caloClusterLinks", "CaloCalTopoClusters", _dask_array_ + ) _set_repr_name("Electron") diff --git a/src/coffea/nanoevents/schemas/base.py b/src/coffea/nanoevents/schemas/base.py index 09812eee0..8a1f2251e 100644 --- a/src/coffea/nanoevents/schemas/base.py +++ b/src/coffea/nanoevents/schemas/base.py @@ -105,7 +105,6 @@ class BaseSchema: """ __dask_capable__ = True - behavior = {} def __init__(self, base_form, *args, **kwargs): params = dict(base_form.get("parameters", {})) diff --git a/src/coffea/nanoevents/schemas/physlite.py b/src/coffea/nanoevents/schemas/physlite.py index 1b9b89205..368dc7d8a 100644 --- a/src/coffea/nanoevents/schemas/physlite.py +++ b/src/coffea/nanoevents/schemas/physlite.py @@ -53,6 +53,7 @@ class PHYSLITESchema(BaseSchema): "GSFTrackParticles": "TrackParticle", "InDetTrackParticles": "TrackParticle", "MuonSpectrometerTrackParticles": "TrackParticle", + "CaloCalTopoClusters": "NanoCollection", } """Default configuration for mixin types, based on the collection name. @@ -79,7 +80,12 @@ def _build_collections(self, branch_forms): key_fields = key.split("/")[-1].split(".") top_key = key_fields[0] sub_key = ".".join(key_fields[1:]) - objname = top_key.replace("Analysis", "").replace("AuxDyn", "") + if ak_form["class"] == "RecordArray" and not ak_form["fields"]: + # skip empty records (e.g. the branches ending in "." only containing the base class) + continue + objname = ( + top_key.replace("Analysis", "").replace("AuxDyn", "").replace("Aux", "") + ) zip_groups[objname].append(((key, sub_key), ak_form)) @@ -97,6 +103,10 @@ def _build_collections(self, branch_forms): # zip the forms contents = {} for objname, keys_and_form in zip_groups.items(): + if len(keys_and_form) == 1: + # don't zip if there is only one item + contents[objname] = keys_and_form[0][1] + continue to_zip = {} for (key, sub_key), form in keys_and_form: if "." in sub_key: @@ -118,14 +128,21 @@ def _build_collections(self, branch_forms): to_zip, objname, self.mixins.get(objname, None), - bypass=True, - ) - content = contents[objname]["content"] - content["parameters"] = dict( - content.get("parameters", {}), collection_name=objname + bypass=False, ) except NotImplementedError: warnings.warn(f"Can't zip collection {objname}") + if "content" in contents[objname]: + # in this case we were able to zip everything together to a ListOffsetArray(RecordArray) + assert "List" in contents[objname]["class"] + content = contents[objname]["content"] + else: + # in this case this was not possible (e.g. because we also had non-list fields) + assert contents[objname]["class"] == "RecordArray" + content = contents[objname] + content["parameters"] = dict( + content.get("parameters", {}), collection_name=objname + ) return contents @staticmethod diff --git a/tests/test_analysis_tools.py b/tests/test_analysis_tools.py index 1e8c46ec1..bb3221432 100644 --- a/tests/test_analysis_tools.py +++ b/tests/test_analysis_tools.py @@ -513,14 +513,14 @@ def test_packed_selection_nminusone(): ): assert np.all(mask == truth) - nminusone.to_npz("nminusone.npz", compressed=False) + nminusone.to_npz("nminusone.npz", compressed=False).compute() with np.load("nminusone.npz") as file: assert np.all(file["labels"] == labels) assert np.all(file["nev"] == nev) assert np.all(file["masks"] == masks) os.remove("nminusone.npz") - nminusone.to_npz("nminusone.npz", compressed=True) + nminusone.to_npz("nminusone.npz", compressed=True).compute() with np.load("nminusone.npz") as file: assert np.all(file["labels"] == labels) assert np.all(file["nev"] == nev) @@ -619,7 +619,7 @@ def test_packed_selection_cutflow(): ): assert np.all(mask == truth) - cutflow.to_npz("cutflow.npz", compressed=False) + cutflow.to_npz("cutflow.npz", compressed=False).compute() with np.load("cutflow.npz") as file: assert np.all(file["labels"] == labels) assert np.all(file["nevonecut"] == nevonecut) @@ -628,7 +628,7 @@ def test_packed_selection_cutflow(): assert np.all(file["maskscutflow"] == maskscutflow) os.remove("cutflow.npz") - cutflow.to_npz("cutflow.npz", compressed=True) + cutflow.to_npz("cutflow.npz", compressed=True).compute() with np.load("cutflow.npz") as file: assert np.all(file["labels"] == labels) assert np.all(file["nevonecut"] == nevonecut) @@ -854,14 +854,14 @@ def test_packed_selection_nminusone_dak(optimization_enabled): ): assert np.all(mask.compute() == truth.compute()) - nminusone.to_npz("nminusone.npz", compressed=False) + nminusone.to_npz("nminusone.npz", compressed=False).compute() with np.load("nminusone.npz") as file: assert np.all(file["labels"] == labels) assert np.all(file["nev"] == list(dask.compute(*nev))) assert np.all(file["masks"] == list(dask.compute(*masks))) os.remove("nminusone.npz") - nminusone.to_npz("nminusone.npz", compressed=True) + nminusone.to_npz("nminusone.npz", compressed=True).compute() with np.load("nminusone.npz") as file: assert np.all(file["labels"] == labels) assert np.all(file["nev"] == list(dask.compute(*nev))) @@ -978,7 +978,7 @@ def test_packed_selection_cutflow_dak(optimization_enabled): ): assert np.all(mask.compute() == truth.compute()) - cutflow.to_npz("cutflow.npz", compressed=False) + cutflow.to_npz("cutflow.npz", compressed=False).compute() with np.load("cutflow.npz") as file: assert np.all(file["labels"] == labels) assert np.all(file["nevonecut"] == list(dask.compute(*nevonecut))) @@ -987,7 +987,7 @@ def test_packed_selection_cutflow_dak(optimization_enabled): assert np.all(file["maskscutflow"] == list(dask.compute(*maskscutflow))) os.remove("cutflow.npz") - cutflow.to_npz("cutflow.npz", compressed=True) + cutflow.to_npz("cutflow.npz", compressed=True).compute() with np.load("cutflow.npz") as file: assert np.all(file["labels"] == labels) assert np.all(file["nevonecut"] == list(dask.compute(*nevonecut))) @@ -1109,14 +1109,14 @@ def test_packed_selection_nminusone_dak_uproot_only(optimization_enabled): ): assert np.all(mask.compute() == truth.compute()) - nminusone.to_npz("nminusone.npz", compressed=False) + nminusone.to_npz("nminusone.npz", compressed=False).compute() with np.load("nminusone.npz") as file: assert np.all(file["labels"] == labels) assert np.all(file["nev"] == list(dask.compute(*nev))) assert np.all(file["masks"] == list(dask.compute(*masks))) os.remove("nminusone.npz") - nminusone.to_npz("nminusone.npz", compressed=True) + nminusone.to_npz("nminusone.npz", compressed=True).compute() with np.load("nminusone.npz") as file: assert np.all(file["labels"] == labels) assert np.all(file["nev"] == list(dask.compute(*nev))) @@ -1233,7 +1233,7 @@ def test_packed_selection_cutflow_dak_uproot_only(optimization_enabled): ): assert np.all(mask.compute() == truth.compute()) - cutflow.to_npz("cutflow.npz", compressed=False) + cutflow.to_npz("cutflow.npz", compressed=False).compute() with np.load("cutflow.npz") as file: assert np.all(file["labels"] == labels) assert np.all(file["nevonecut"] == list(dask.compute(*nevonecut))) @@ -1242,7 +1242,7 @@ def test_packed_selection_cutflow_dak_uproot_only(optimization_enabled): assert np.all(file["maskscutflow"] == list(dask.compute(*maskscutflow))) os.remove("cutflow.npz") - cutflow.to_npz("cutflow.npz", compressed=True) + cutflow.to_npz("cutflow.npz", compressed=True).compute() with np.load("cutflow.npz") as file: assert np.all(file["labels"] == labels) assert np.all(file["nevonecut"] == list(dask.compute(*nevonecut))) diff --git a/tests/test_nanoevents_physlite.py b/tests/test_nanoevents_physlite.py index f82471198..95f58491d 100644 --- a/tests/test_nanoevents_physlite.py +++ b/tests/test_nanoevents_physlite.py @@ -1,19 +1,17 @@ import os -import numpy as np +import dask import pytest from coffea.nanoevents import NanoEventsFactory, PHYSLITESchema -pytestmark = pytest.mark.skip(reason="uproot is upset with this file...") - def _events(): path = os.path.abspath("tests/samples/DAOD_PHYSLITE_21.2.108.0.art.pool.root") factory = NanoEventsFactory.from_root( {path: "CollectionTree"}, schemaclass=PHYSLITESchema, - permit_dask=False, + permit_dask=True, ) return factory.events() @@ -23,54 +21,21 @@ def events(): return _events() +def test_load_single_field_of_linked(events): + with dask.config.set({"awkward.raise-failed-meta": True}): + events.Electrons.caloClusters.calE.compute() + + @pytest.mark.parametrize("do_slice", [False, True]) def test_electron_track_links(events, do_slice): if do_slice: - events = events[np.random.randint(2, size=len(events)).astype(bool)] - for event in events: - for electron in event.Electrons: + events = events[::2] + trackParticles = events.Electrons.trackParticles.compute() + for i, event in enumerate(events[["Electrons", "GSFTrackParticles"]].compute()): + for j, electron in enumerate(event.Electrons): for link_index, link in enumerate(electron.trackParticleLinks): track_index = link.m_persIndex - print(track_index) - print(event.GSFTrackParticles) - print(electron.trackParticleLinks) - print(electron.trackParticles) - assert ( event.GSFTrackParticles[track_index].z0 - == electron.trackParticles[link_index].z0 - ) - - -# from MetaData/EventFormat -_hash_to_target_name = { - 13267281: "TruthPhotons", - 342174277: "TruthMuons", - 368360608: "TruthNeutrinos", - 375408000: "TruthTaus", - 394100163: "TruthElectrons", - 614719239: "TruthBoson", - 660928181: "TruthTop", - 779635413: "TruthBottom", -} - - -def test_truth_links_toplevel(events): - children_px = events.TruthBoson.children.px - for i_event, event in enumerate(events): - for i_particle, particle in enumerate(event.TruthBoson): - for i_link, link in enumerate(particle.childLinks): - assert ( - event[_hash_to_target_name[link.m_persKey]][link.m_persIndex].px - == children_px[i_event][i_particle][i_link] - ) - - -def test_truth_links(events): - for i_event, event in enumerate(events): - for i_particle, particle in enumerate(event.TruthBoson): - for i_link, link in enumerate(particle.childLinks): - assert ( - event[_hash_to_target_name[link.m_persKey]][link.m_persIndex].px - == particle.children[i_link].px + == trackParticles[i][j][link_index].z0 )