fix: second wave of #1578 fixes — 43 items across processor, nanoevents, analysis/jetmet/lookup/lumi/ml/btag tools, CI, and dead code

.github/workflows/ci.yml

@@ -246,7 +246,7 @@ jobs:
   release:
     if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')
     runs-on: ubuntu-latest
-    needs: [test, test-no-dask, test-vine]
+    needs: [test, test-dask-client, test-no-dask, test-vine]
     permissions:
       id-token: write
       attestations: write
@@ -278,7 +278,10 @@ jobs:
         password: ${{ secrets.PYPI_TOKEN }}
 
   pass:
+    if: always()
     needs: [test, test-dask-client, test-no-dask, test-vine]
     runs-on: ubuntu-latest
     steps:
+      - if: contains(needs.*.result, 'failure') || contains(needs.*.result, 'cancelled') || contains(needs.*.result, 'skipped')
+        run: exit 1
       - run: echo "All jobs passed"

codecov.yml

@@ -1,7 +1,6 @@
 comment: off
 
 codecov:
-  token: 045255bb-e0d8-4c5d-b413-9c33128a03a4
   notify:
     require_ci_to_pass: yes
 

src/coffea/analysis_tools.py

@@ -55,28 +55,6 @@ def _get_hist_class(delayed_mode):
     return hist.Hist
 
 
-def _generate_slices(array_length, max_elements=128):
-    """Generate slices to split an array into chunks of at most `max_elements` elements
-
-    Parameters
-    ----------
-    array_length : int
-        The length of the array to split
-    max_elements : int, optional
-        The maximum number of elements in each chunk. Default is 128.
-
-    Returns
-    -------
-    slices : list of slice objects
-        A list of slice objects to iterate over and split the array into chunks with at most `max_elements` elements per slice
-    """
-    slices = []
-    for start in range(0, array_length, max_elements):
-        end = min(start + max_elements, array_length)
-        slices.append(slice(start, end))
-    return slices
-
-
 def boolean_masks_to_categorical_integers(
     masks,
     insert_unmasked_as_zeros=False,
@@ -187,10 +165,12 @@ def add(self, other):
 
     def __add__(self, other):
         temp = WeightStatistics(self.sumw, self.sumw2, self.minw, self.maxw, self.n)
-        return temp.add(other)
+        temp.add(other)
+        return temp
 
     def __iadd__(self, other):
-        return self.add(other)
+        self.add(other)
+        return self
 
 
 class Weights:
@@ -394,7 +374,7 @@ def __add_multivariation_delayed(
         """Add a new weight with multiple variations in delayed mode"""
         dask_awkward = _import_dask_awkward()
 
-        if isinstance(weight, awkward.types.OptionType):
+        if isinstance(dask_awkward.type(weight), awkward.types.OptionType):
             # TODO what to do with option-type? is it representative of unknown weight
             # and we default to one or is it an invalid weight and we should never use this
             # event in the first place (0) ?
@@ -567,8 +547,8 @@ def weight(self, modifier=None):
         """
         if modifier is None:
             return self._weight
-        elif "Down" in modifier and modifier not in self._modifiers:
-            return self._weight / self._modifiers[modifier.replace("Down", "Up")]
+        elif modifier.endswith("Down") and modifier not in self._modifiers:
+            return self._weight / self._modifiers[modifier[:-4] + "Up"]
         return self._weight * self._modifiers[modifier]
 
     def partial_weight(self, include=[], exclude=[], modifier=None):
@@ -628,12 +608,17 @@ def _partial_weight(self, include, exclude, modifier=None):
 
         if modifier is None:
             return w
-        elif modifier.replace("Down", "").replace("Up", "") not in names:
+        base = (
+            modifier[:-4]
+            if modifier.endswith("Down")
+            else modifier[:-2] if modifier.endswith("Up") else modifier
+        )
+        if not any(base == n or base.startswith(n + "_") for n in names):
             raise ValueError(
                 f"Modifier {modifier} is not in the list of included weights"
             )
-        elif "Down" in modifier and modifier not in self._modifiers:
-            return w / self._modifiers[modifier.replace("Down", "Up")]
+        if modifier.endswith("Down") and modifier not in self._modifiers:
+            return w / self._modifiers[modifier[:-4] + "Up"]
         return w * self._modifiers[modifier]
 
     @property
@@ -642,7 +627,8 @@ def variations(self):
         keys = set(self._modifiers.keys())
         # add any missing 'Down' variation
         for k in self._modifiers.keys():
-            keys.add(k.replace("Up", "Down"))
+            if k.endswith("Up"):
+                keys.add(k[:-2] + "Down")
         return keys
 
 
@@ -1188,10 +1174,6 @@ def yieldhist(self, weighted=None, scale=None, categorical=None):
         if do_weighted:
             axes.append(hist.storage.Weight())
         if not self._delayed_mode and not do_categorical:
-            if categorical is not None:
-                raise NotImplementedError(
-                    "yieldhist is not implemented for non-delayed mode (v1) with categorical"
-                )
             h = hist.Hist(*axes)
             weighttofill = self._wgtev if do_weighted else self._nev
             if do_scaled:
@@ -1200,10 +1182,6 @@ def yieldhist(self, weighted=None, scale=None, categorical=None):
         elif self._delayed_mode and not do_categorical:
             dask_awkward = _import_dask_awkward()
 
-            if categorical is not None:
-                raise NotImplementedError(
-                    "yieldhist is not implemented for non-delayed mode (v1) with categorical"
-                )
             h = Hist(*axes)
 
             for i, mask in enumerate(self._masks, 1):
@@ -1772,10 +1750,6 @@ def yieldhist(self, weighted=None, scale=None, categorical=None):
         if do_weighted:
             axes.append(hist.storage.Weight())
         if not self._delayed_mode and not do_categorical:
-            if categorical is not None:
-                raise NotImplementedError(
-                    "yieldhist is not implemented for non-delayed mode (v1) with categorical"
-                )
             honecut = hist.Hist(*axes)
             hcutflow = honecut.copy()
             hcutflow.axes.name = ("cutflow",)
@@ -1789,10 +1763,6 @@ def yieldhist(self, weighted=None, scale=None, categorical=None):
         elif self._delayed_mode and not do_categorical:
             dask_awkward = _import_dask_awkward()
 
-            if categorical is not None:
-                raise NotImplementedError(
-                    "yieldhist is not implemented for non-delayed mode (v1) with categorical"
-                )
             honecut = Hist(*axes)
             hcutflow = honecut.copy()
             hcutflow.axes.name = ("cutflow",)
@@ -2291,7 +2261,6 @@ def add_multiple(self, selections, fill_value=False):
         for name, selection in selections.items():
             self.add(name, selection, fill_value)
 
-    @lru_cache
     def require(self, **names):
         """Return a mask vector corresponding to specific requirements
 
@@ -2318,6 +2287,12 @@ def require(self, **names):
         returns a boolean array where an entry is True if the corresponding entries
         ``cut1 == True``, ``cut2 == False``, and ``cut3`` arbitrary.
         """
+        # copy so a caller mutating the returned mask cannot corrupt the shared cache
+        result = self._require(**names)
+        return result.copy() if isinstance(result, numpy.ndarray) else result
+
+    @lru_cache
+    def _require(self, **names):
         for cut, v in names.items():
             if not isinstance(cut, str) or cut not in self._names:
                 raise ValueError(

src/coffea/btag_tools/btagscalefactor.py

@@ -147,31 +147,40 @@ def __init__(self, filename, workingpoint, methods="comb,comb,incl", keep_df=Fal
             )
             mapping = numpy.full(bin_low_edges[0].shape, -1)
 
-            def findbin(flavor, eta, pt, discr):
-                btvflavor = self._flavor2btvflavor[flavor]
-                for i, (fbin, ebin, pbin, dbin) in enumerate(allbins):
-                    if (
-                        btvflavor == fbin
-                        and ebin[0] <= eta < ebin[1]
-                        and pbin[0] <= pt < pbin[1]
-                        and dbin[0] <= discr < dbin[1]
-                    ):
-                        return i
-                if eta < 0:
-                    # maybe in this region we have only abseta
-                    for i, (fbin, ebin, pbin, dbin) in enumerate(allbins):
-                        if (
-                            btvflavor == fbin
-                            and -ebin[1] <= eta < -ebin[0]
-                            and pbin[0] <= pt < pbin[1]
-                            and dbin[0] <= discr < dbin[1]
-                        ):
-                            return i
-                return -1
-
-            for idx, _ in numpy.ndenumerate(mapping):
-                flavor, eta, pt, discr = (x[idx] for x in bin_low_edges)
-                mapping[idx] = findbin(flavor, eta, pt, discr)
+            fbins = numpy.array([b[0] for b in allbins])
+            eta_lo = numpy.array([b[1][0] for b in allbins])
+            eta_hi = numpy.array([b[1][1] for b in allbins])
+            pt_lo = numpy.array([b[2][0] for b in allbins])
+            pt_hi = numpy.array([b[2][1] for b in allbins])
+            discr_lo = numpy.array([b[3][0] for b in allbins])
+            discr_hi = numpy.array([b[3][1] for b in allbins])
+
+            eta_cell = bin_low_edges[1].reshape(-1, 1)
+            pt_cell = bin_low_edges[2].reshape(-1, 1)
+            discr_cell = bin_low_edges[3].reshape(-1, 1)
+            btvflavor_cell = numpy.empty(bin_low_edges[0].size, dtype=fbins.dtype)
+            for flav, btv in self._flavor2btvflavor.items():
+                btvflavor_cell[bin_low_edges[0].reshape(-1) == flav] = btv
+            btvflavor_cell = btvflavor_cell.reshape(-1, 1)
+
+            base_match = (
+                (btvflavor_cell == fbins)
+                & (pt_lo <= pt_cell)
+                & (pt_cell < pt_hi)
+                & (discr_lo <= discr_cell)
+                & (discr_cell < discr_hi)
+            )
+            match = base_match & (eta_lo <= eta_cell) & (eta_cell < eta_hi)
+            found = match.any(axis=1)
+            result = numpy.where(found, match.argmax(axis=1), -1)
+
+            # maybe in this region we have only abseta
+            abseta_match = base_match & (-eta_hi <= eta_cell) & (eta_cell < -eta_lo)
+            fallback = (~found) & (eta_cell[:, 0] < 0)
+            fallback &= abseta_match.any(axis=1)
+            result[fallback] = abseta_match.argmax(axis=1)[fallback]
+
+            mapping = result.reshape(mapping.shape)
 
             if self.workingpoint == BTagScaleFactor.RESHAPE:
                 self._corrections[syst] = dense_mapped_lookup(

src/coffea/dataset_tools/manipulations.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import ast
 import copy
 import sys
 from collections.abc import Callable
@@ -321,7 +322,7 @@ def get_failed_steps_for_dataset(
         )
 
     for failure in failures:
-        args_as_types = tuple(eval(arg) for arg in failure.args)
+        args_as_types = tuple(ast.literal_eval(arg) for arg in failure.args)
 
         fname, object_path, start, stop, is_step = args_as_types
 

src/coffea/dataset_tools/rucio_utils.py

@@ -274,8 +274,9 @@ def get_dataset_files_replicas(
                 outfiles.append(outfile)
                 outsites.append(outsite)
             elif mode == "first":
-                outfiles.append(outfile[0])
-                outsites.append(outsite[0])
+                if outfile:
+                    outfiles.append(outfile[0])
+                    outsites.append(outsite[0])
             else:
                 raise NotImplementedError(f"Mode {mode} not yet implemented!")
 
@@ -287,7 +288,7 @@ def get_dataset_files_replicas(
                 sites_counts[site] += 1
     elif mode == "first":
         for site_by_file in outsites:
-            sites_counts[site] += 1
+            sites_counts[site_by_file] += 1
 
     return outfiles, outsites, sites_counts
 

src/coffea/jetmet_tools/CorrectedJetsFactory.py

@@ -34,8 +34,11 @@ def __call__(self, array, like_what):
 
 
 def rand_gauss(item):
+    np_item = awkward.typetracer.length_one_if_typetracer(item).to_numpy()
     seeds = (
-        awkward.typetracer.length_one_if_typetracer(item).to_numpy()[[0, -1]].view("i4")
+        np_item[[0, -1]].view("i4")
+        if len(np_item)
+        else numpy.zeros(2, dtype=numpy.int32)
     )
     randomstate = numpy.random.Generator(numpy.random.PCG64(seeds))
 
@@ -148,20 +151,20 @@ def __init__(self, name_map, jec_stack):
         # from PhysicsTools/PatUtils/interface/SmearedJetProducerT.h#L283
         self.forceStochastic = False
 
-        if "ptRaw" not in name_map or name_map["ptRaw"] is None:
+        self.treat_pt_as_raw = "ptRaw" not in name_map or name_map["ptRaw"] is None
+        if self.treat_pt_as_raw:
             warnings.warn(
                 "There is no name mapping for ptRaw,"
                 " CorrectedJets will assume that <object>.pt is raw pt!"
             )
             name_map["ptRaw"] = name_map["JetPt"] + "_raw"
-        self.treat_pt_as_raw = "ptRaw" not in name_map
 
         if "massRaw" not in name_map or name_map["massRaw"] is None:
             warnings.warn(
                 "There is no name mapping for massRaw,"
-                " CorrectedJets will assume that <object>.mass is raw pt!"
+                " CorrectedJets will assume that <object>.mass is raw mass!"
             )
-            name_map["ptRaw"] = name_map["JetMass"] + "_raw"
+            name_map["massRaw"] = name_map["JetMass"] + "_raw"
 
         total_signature = set()
         for part in _stack_parts:
@@ -226,7 +229,7 @@ def build(self, injets):
         fields = awkward.fields(jets)
         if len(fields) == 0:
             raise Exception(
-                "Empty record, please pass a jet object with at least {self.real_sig} defined!"
+                f"Empty record, please pass a jet object with at least {self.real_sig} defined!"
             )
         out = awkward.flatten(jets)
         wrap = partial(
@@ -484,7 +487,7 @@ def build_variant(unc, template, jetpt, jetpt_orig, jetmass, jetmass_orig):
                     label=f"{name}",
                 )
 
-        out_parms = out.layout.parameters
+        out_parms = dict(out.layout.parameters)
         out_parms["corrected"] = True
         out = awkward.zip(
             out_dict, depth_limit=1, parameters=out_parms, behavior=out.behavior

src/coffea/jetmet_tools/FactorizedJetCorrector.py

@@ -13,8 +13,7 @@ def _checkConsistency(against, tocheck):
     else:
         if against != tocheck:
             raise Exception(
-                "Corrector for {} is mixed"
-                "with correctors for {}!".format(tocheck, against)
+                f"Corrector for {tocheck} is mixed with correctors for {against}!"
             )
     return tocheck
 

src/coffea/lookup_tools/evaluator.py

@@ -4,6 +4,7 @@
 from coffea.lookup_tools.jec_uncertainty_lookup import jec_uncertainty_lookup
 from coffea.lookup_tools.jersf_lookup import jersf_lookup
 from coffea.lookup_tools.jme_standard_function import jme_standard_function
+from coffea.lookup_tools.json_lookup import json_lookup
 from coffea.lookup_tools.rochester_lookup import rochester_lookup
 
 lookup_types = {
@@ -14,6 +15,7 @@
     "jec_uncertainty_lookup": jec_uncertainty_lookup,
     "rochester_lookup": rochester_lookup,
     "correctionlib_wrapper": correctionlib_wrapper,
+    "json_lookup": json_lookup,
 }
 
 

src/coffea/lookup_tools/extractor.py

@@ -130,8 +130,6 @@ def add_weight_sets(self, weightsdescs):
             else:
                 weights, thetype = self.extract_from_file(thefile, name)
                 self.add_weight_set(local_name, thetype, weights)
-                if thetype == "json_lookup":
-                    self._names[local_name] = 0
 
     def import_file(self, thefile):
         """