Skip to content

Commit

Permalink
performance optimization for lumimask taskgraphs
Browse files Browse the repository at this point in the history
  • Loading branch information
lgray committed Mar 19, 2024
1 parent f38d7ef commit 60b0650
Showing 1 changed file with 7 additions and 4 deletions.
11 changes: 7 additions & 4 deletions src/coffea/lumi_tools/lumi_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -169,10 +169,12 @@ def __call__(self, runs, lumis):
"""

def apply(runs, lumis):
backend = awkward.backend(runs)
# fill numba typed dict
_masks = Dict.empty(key_type=types.uint32, value_type=types.uint32[:])
for k, v in self._masks.items():
_masks[k] = v
if backend != "typetracer":
for k, v in self._masks.items():
_masks[k] = v

runs_orig = runs
if isinstance(runs, awkward.highlevel.Array):
Expand All @@ -184,10 +186,11 @@ def apply(runs, lumis):
awkward.typetracer.length_zero_if_typetracer(lumis)
)
mask_out = numpy.zeros(dtype="bool", shape=runs.shape)
LumiMask._apply_run_lumi_mask_kernel(_masks, runs, lumis, mask_out)
if backend != "typetracer":
LumiMask._apply_run_lumi_mask_kernel(_masks, runs, lumis, mask_out)
if isinstance(runs_orig, awkward.Array):
mask_out = awkward.Array(mask_out)
if awkward.backend(runs_orig) == "typetracer":
if backend == "typetracer":
mask_out = awkward.Array(
mask_out.layout.to_typetracer(forget_length=True)
)
Expand Down

0 comments on commit 60b0650

Please sign in to comment.