flat calling structure for trackParticle(s) behavior methods

scikit-hep · Aug 22, 2023 · bcc7dea · bcc7dea
1 parent afbc8c5
commit bcc7dea
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/src/coffea/nanoevents/methods/physlite.py b/src/coffea/nanoevents/methods/physlite.py
@@ -204,12 +204,12 @@ def trackParticles(self, _dask_array_=None):
 
     @property
     def trackParticle(self, _dask_array_=None):
-        if _dask_array_ is not None:
-            self = _dask_array_  # TODO: is this what i should be doing?
-        trackParticles = self.trackParticles
-        return self.trackParticles[
-            tuple([slice(None) for i in range(trackParticles.ndim - 1)] + [0])
-        ]
+        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):