Skip to content

Commit

Permalink
Merge branch 'master' into master
Browse files Browse the repository at this point in the history
  • Loading branch information
@lgray
lgray authored Aug 21, 2024
2 parents 2f5ca58 + 246f134 commit ea7c869
Show file tree
Hide file tree
Showing 8 changed files with 110 additions and 170 deletions.
4 changes: 2 additions & 2 deletions binder/applying_corrections.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -439,11 +439,11 @@
"jets['pt_raw'] = (1 - jets['rawFactor']) * jets['pt']\n",
"jets['mass_raw'] = (1 - jets['rawFactor']) * jets['mass']\n",
"jets['pt_gen'] = ak.values_astype(ak.fill_none(jets.matched_gen.pt, 0), np.float32)\n",
"jets['rho'] = ak.broadcast_arrays(events.fixedGridRhoFastjetAll, jets.pt)[0]\n",
"jets['PU_rho'] = ak.broadcast_arrays(events.fixedGridRhoFastjetAll, jets.pt)[0]\n",
"name_map['ptGenJet'] = 'pt_gen'\n",
"name_map['ptRaw'] = 'pt_raw'\n",
"name_map['massRaw'] = 'mass_raw'\n",
"name_map['Rho'] = 'rho'\n",
"name_map['Rho'] = 'PU_rho'\n",
" \n",
"corrector = FactorizedJetCorrector(\n",
" Fall17_17Nov2017_V32_MC_L2Relative_AK4PFPuppi=evaluator['Fall17_17Nov2017_V32_MC_L2Relative_AK4PFPuppi'],\n",
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ classifiers = [
]
dependencies = [
"awkward>=2.6.7",
"uproot>=5.3.0,!=5.3.3,!=5.3.4,!=5.3.5",
"uproot>=5.3.10",
"dask[array]>=2024.3.0;python_version>'3.8'",
"dask[array]>=2023.4.0;python_version<'3.9'",
"dask-awkward>=2024.3.0",
Expand Down
14 changes: 10 additions & 4 deletions src/coffea/nanoevents/methods/physlite.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,14 @@ def namefcn(self):


def _element_link(target_collection, eventindex, index, key):
# check to make sure both the target_collection and the eventindex are both dask-awkward arrays or both awkward arrays
if isinstance(target_collection, dask_awkward.Array) != isinstance(
eventindex, dask_awkward.Array
):
raise ValueError(
"element linking must be done on two dask_awkward arrays or two akward arrays not a mix of the two"
)

global_index = _get_global_index(target_collection, eventindex, index)
global_index = awkward.where(key != 0, global_index, -1)
return target_collection._apply_global_index(global_index)
Expand Down Expand Up @@ -277,14 +285,12 @@ def trackParticle(self, dask_array):

@dask_property
def caloClusters(self):
return _element_link_method(
self, "caloClusterLinks", "CaloCalTopoClusters", None
)
return _element_link_method(self, "caloClusterLinks", "egammaClusters", None)

@caloClusters.dask
def caloClusters(self, dask_array):
return _element_link_method(
self, "caloClusterLinks", "CaloCalTopoClusters", dask_array
self, "caloClusterLinks", "egammaClusters", dask_array
)


Expand Down
160 changes: 0 additions & 160 deletions src/coffea/nanoevents/methods/vector.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -513,90 +513,6 @@ class PtEtaPhiMLorentzVector(LorentzVector):
Some additional properties are overridden for performance
"""

@property
def E(self):
"""Alias for `t`"""
return self.t

@property
def pt(self):
"""Alias for `r`"""
return self["pt"]

@property
def eta(self):
r"""Pseudorapidity
:math:`-\ln\tan(\theta/2) = \text{arcsinh}(z/r)`
"""
return self["eta"]

@property
def phi(self):
r"""Azimuthal angle relative to X axis in XY plane
:math:`\text{arctan2}(y, x)`
"""
return self["phi"]

@property
def mass(self):
r"""Invariant mass (+, -, -, -)
:math:`\sqrt{t^2-x^2-y^2-z^2}`
"""
return self["mass"]

@property
def rho(self):
r"""Distance from origin in 3D
:math:`\sqrt{x^2+y^2+z^2} = \sqrt{r^2+z^2}`
"""
return self.pt * numpy.cosh(self.eta)

@property
def theta(self):
r"""Inclination angle from XY plane
:math:`\text{arctan2}(r, z) = 2\text{arctan}(e^{-\eta})`
"""
return 2 * numpy.arctan(numpy.exp(-self.eta))

@property
def r(self):
r"""Distance from origin in XY plane
:math:`\sqrt{x^2+y^2} = \rho \sin(\theta)`
"""
return self.pt

@property
def z(self):
r"""Cartesian z value
:math:`\rho \cos(\theta) = r \sinh(\eta)`
"""
return self.pt * numpy.sinh(self.eta)

@property
def t(self):
r"""Cartesian time component
:math:`\sqrt{\rho^2+m^2}`
"""
return numpy.hypot(self.rho, self.mass)

@property
def rho2(self):
"""Squared `rho`"""
return self.rho * self.rho

@property
def mass2(self):
"""Squared `mass`"""
return self.mass * self.mass

@awkward.mixin_class_method(numpy.multiply, {numbers.Number})
def multiply(self, other):
"""Multiply this vector by a scalar elementwise using `x`, `y`, `z`, and `t` components
Expand Down Expand Up @@ -644,82 +560,6 @@ class PtEtaPhiELorentzVector(LorentzVector):
Some additional properties are overridden for performance
"""

@property
def E(self):
"""Alias for `t`"""
return self.t

@property
def pt(self):
"""Alias for `r`"""
return self["pt"]

@property
def eta(self):
r"""Pseudorapidity
:math:`-\ln\tan(\theta/2) = \text{arcsinh}(z/r)`
"""
return self["eta"]

@property
def phi(self):
r"""Azimuthal angle relative to X axis in XY plane
:math:`\text{arctan2}(y, x)`
"""
return self["phi"]

@property
def energy(self):
"""Alias for `t`"""
return self["energy"]

@property
def t(self):
r"""Cartesian time component
:math:`\sqrt{\rho^2+m^2}`
"""
return self["energy"]

@property
def rho(self):
r"""Distance from origin in 3D
:math:`\sqrt{x^2+y^2+z^2} = \sqrt{r^2+z^2}`
"""
return self.pt * numpy.cosh(self.eta)

@property
def theta(self):
r"""Inclination angle from XY plane
:math:`\text{arctan2}(r, z) = 2\text{arctan}(e^{-\eta})`
"""
return 2 * numpy.arctan(numpy.exp(-self.eta))

@property
def r(self):
r"""Distance from origin in XY plane
:math:`\sqrt{x^2+y^2} = \rho \sin(\theta)`
"""
return self.pt

@property
def z(self):
r"""Cartesian z value
:math:`r \sinh(\eta)`
"""
return self.pt * numpy.sinh(self.eta)

@property
def rho2(self):
"""Squared `rho`"""
return self.rho * self.rho

@awkward.mixin_class_method(numpy.multiply, {numbers.Number})
def multiply(self, other):
"""Multiply this vector by a scalar elementwise using `x`, `y`, `z`, and `t` components
Expand Down
2 changes: 1 addition & 1 deletion src/coffea/nanoevents/schemas/physlite.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ class PHYSLITESchema(BaseSchema):
"GSFTrackParticles": "TrackParticle",
"InDetTrackParticles": "TrackParticle",
"MuonSpectrometerTrackParticles": "TrackParticle",
"CaloCalTopoClusters": "NanoCollection",
"egammaClusters": "NanoCollection",
}
"""Default configuration for mixin types, based on the collection name.
Expand Down
Binary file removed tests/samples/DAOD_PHYSLITE_21.2.108.0.art.pool.root
View file
Binary file not shown.
Binary file added tests/samples/PHYSLITE_example.root
View file
Binary file not shown.
98 changes: 96 additions & 2 deletions tests/test_nanoevents_physlite.py
View file
Original file line number Diff line number Diff line change
@@ -1,17 +1,20 @@
import json
import os

import awkward as ak
import dask
import pytest

from coffea.nanoevents import NanoEventsFactory, PHYSLITESchema


def _events():
path = os.path.abspath("tests/samples/DAOD_PHYSLITE_21.2.108.0.art.pool.root")
def _events(filter=None):
path = os.path.abspath("tests/samples/PHYSLITE_example.root")
factory = NanoEventsFactory.from_root(
{path: "CollectionTree"},
schemaclass=PHYSLITESchema,
delayed=True,
uproot_options=dict(filter_name=filter),
)
return factory.events()

Expand All @@ -26,6 +29,9 @@ def test_load_single_field_of_linked(events):
events.Electrons.caloClusters.calE.compute()


@pytest.mark.skip(
reason="temporarily disabled because of uproot issue #1267 https://github.com/scikit-hep/uproot5/issues/1267"
)
@pytest.mark.parametrize("do_slice", [False, True])
def test_electron_track_links(events, do_slice):
if do_slice:
Expand All @@ -39,3 +45,91 @@ def test_electron_track_links(events, do_slice):
event.GSFTrackParticles[track_index].z0
== trackParticles[i][j][link_index].z0
)


def mock_empty(form, behavior={}):
return ak.Array(
form.length_zero_array(),
behavior=behavior,
)


def test_electron_forms():
def filter_name(name):
return name in [
"AnalysisElectronsAuxDyn.pt",
"AnalysisElectronsAuxDyn.eta",
"AnalysisElectronsAuxDyn.phi",
"AnalysisElectronsAuxDyn.m",
]

events = _events(filter_name)

mocked, _, _ = ak.to_buffers(mock_empty(events.form))

expected_json = {
"class": "RecordArray",
"fields": ["Electrons"],
"contents": [
{
"class": "ListOffsetArray",
"offsets": "i64",
"content": {
"class": "RecordArray",
"fields": ["pt", "_eventindex", "eta", "phi", "m"],
"contents": [
{
"class": "NumpyArray",
"primitive": "float32",
"inner_shape": [],
"parameters": {"__doc__": "AnalysisElectronsAuxDyn.pt"},
"form_key": "node3",
},
{
"class": "NumpyArray",
"primitive": "int64",
"inner_shape": [],
"parameters": {},
"form_key": "node4",
},
{
"class": "NumpyArray",
"primitive": "float32",
"inner_shape": [],
"parameters": {"__doc__": "AnalysisElectronsAuxDyn.eta"},
"form_key": "node5",
},
{
"class": "NumpyArray",
"primitive": "float32",
"inner_shape": [],
"parameters": {"__doc__": "AnalysisElectronsAuxDyn.phi"},
"form_key": "node6",
},
{
"class": "NumpyArray",
"primitive": "float32",
"inner_shape": [],
"parameters": {"__doc__": "AnalysisElectronsAuxDyn.m"},
"form_key": "node7",
},
],
"parameters": {
"__record__": "Electron",
"collection_name": "Electrons",
},
"form_key": "node2",
},
"parameters": {},
"form_key": "node1",
}
],
"parameters": {
"__doc__": "CollectionTree",
"__record__": "NanoEvents",
"metadata": {},
},
"form_key": "node0",
}

assert json.dumps(expected_json) == mocked.to_json()

0 comments on commit ea7c869

Please sign in to comment.