diff --git a/.all-contributorsrc b/.all-contributorsrc
index 78194e308..0a43d7903 100644
--- a/.all-contributorsrc
+++ b/.all-contributorsrc
@@ -175,6 +175,24 @@
"contributions": [
"code"
]
+ },
+ {
+ "login": "Lara813",
+ "name": "lamarkus",
+ "avatar_url": "https://avatars.githubusercontent.com/u/67378765?v=4",
+ "profile": "https://github.com/Lara813",
+ "contributions": [
+ "code"
+ ]
+ },
+ {
+ "login": "LennertGriesing",
+ "name": "Lennert Griesing",
+ "avatar_url": "https://avatars.githubusercontent.com/u/115483673?v=4",
+ "profile": "https://github.com/LennertGriesing",
+ "contributions": [
+ "code"
+ ]
}
],
"commitType": "docs"
diff --git a/README.md b/README.md
index 43abbe91b..5ed8abf74 100644
--- a/README.md
+++ b/README.md
@@ -154,6 +154,8 @@ For a better overview of the tasks that are triggered by the commands below, che
Ana Andrade
💻 👀 |
philippgadow
💻 |
Lukas Schaller
💻 |
+ 
lamarkus
💻 |
+ 
Lennert Griesing
💻 |
diff --git a/analysis_templates/cms_minimal/__cf_module_name__/columnflow_patches.py b/analysis_templates/cms_minimal/__cf_module_name__/columnflow_patches.py
index 4a0eba031..e86c55350 100644
--- a/analysis_templates/cms_minimal/__cf_module_name__/columnflow_patches.py
+++ b/analysis_templates/cms_minimal/__cf_module_name__/columnflow_patches.py
@@ -4,8 +4,6 @@
Collection of patches of underlying columnflow tasks.
"""
-import os
-
import law
from columnflow.util import memoize
@@ -17,19 +15,8 @@
def patch_bundle_repo_exclude_files():
from columnflow.tasks.framework.remote import BundleRepo
- # get the relative path to CF_BASE
- cf_rel = os.path.relpath(os.environ["CF_BASE"], os.environ["__cf_short_name_uc___BASE"])
-
- # amend exclude files to start with the relative path to CF_BASE
- exclude_files = [os.path.join(cf_rel, path) for path in BundleRepo.exclude_files]
-
- # add additional files
- exclude_files.extend([
- "docs", "tests", "data", "assets", ".law", ".setups", ".data", ".github",
- ])
-
- # overwrite them
- BundleRepo.exclude_files[:] = exclude_files
+ # add additional files to exclude
+ BundleRepo.exclude_files += ["docs"]
logger.debug("patched exclude_files of cf.BundleRepo")
diff --git a/analysis_templates/cms_minimal/__cf_module_name__/config/analysis___cf_short_name_lc__.py b/analysis_templates/cms_minimal/__cf_module_name__/config/analysis___cf_short_name_lc__.py
index 5fbc11d9b..063f8816a 100644
--- a/analysis_templates/cms_minimal/__cf_module_name__/config/analysis___cf_short_name_lc__.py
+++ b/analysis_templates/cms_minimal/__cf_module_name__/config/analysis___cf_short_name_lc__.py
@@ -308,23 +308,20 @@
cfg.add_variable(
name="event",
expression="event",
- binning=(1, 0.0, 1.0e9),
+ binning=(100, 0.0, 1.0e9),
x_title="Event number",
- discrete_x=True,
)
cfg.add_variable(
name="run",
expression="run",
- binning=(1, 100000.0, 500000.0),
+ binning=(100, 100000.0, 500000.0),
x_title="Run number",
- discrete_x=True,
)
cfg.add_variable(
name="lumi",
expression="luminosityBlock",
- binning=(1, 0.0, 5000.0),
+ binning=(100, 0.0, 5000.0),
x_title="Luminosity block",
- discrete_x=True,
)
cfg.add_variable(
name="n_jet",
diff --git a/analysis_templates/cms_minimal/law.cfg b/analysis_templates/cms_minimal/law.cfg
index 35a233f00..7496ee29f 100644
--- a/analysis_templates/cms_minimal/law.cfg
+++ b/analysis_templates/cms_minimal/law.cfg
@@ -209,6 +209,7 @@ cache_cleanup: $CF_WLCG_CACHE_CLEANUP
cache_max_size: 15GB
cache_global_lock: True
cache_mtime_patience: -1
+rucio_report_access: True
[wlcg_fs_global_redirector]
@@ -220,3 +221,4 @@ cache_cleanup: $CF_WLCG_CACHE_CLEANUP
cache_max_size: 15GB
cache_global_lock: True
cache_mtime_patience: -1
+rucio_report_access: True
diff --git a/bin/cf_inspect b/bin/cf_inspect
index 2b35c5449..fdefffe38 100755
--- a/bin/cf_inspect
+++ b/bin/cf_inspect
@@ -6,10 +6,10 @@ action () {
local this_dir="$( cd "$( dirname "${this_file}" )" && pwd )"
# check arguments
- # [ "$#" -eq 0 ] && {
- # echo "ERROR: at least one file must be provided"
- # return 1
- # }
+ [ "$#" -eq 0 ] && {
+ echo "ERROR: at least one file must be provided"
+ return 1
+ }
# determine the sandbox to use
local cf_inspect_sandbox="${CF_INSPECT_SANDBOX:-venv_columnar_dev}"
diff --git a/bin/cf_inspect.py b/bin/cf_inspect.py
index 8e5465508..598d95b31 100644
--- a/bin/cf_inspect.py
+++ b/bin/cf_inspect.py
@@ -13,8 +13,8 @@
import pickle
import awkward as ak
-import numpy as np # noqa
+from columnflow.columnar_util import update_ak_array, ChunkedIOHandler
from columnflow.util import ipython_shell
from columnflow.types import Any
@@ -35,11 +35,9 @@ def _load_parquet(fname: str, **kwargs) -> ak.Array:
def _load_nano_root(fname: str, treepath: str | None = None, **kwargs) -> ak.Array:
import uproot
- import coffea.nanoevents
-
- source = uproot.open(fname)
# get the default treepath
+ source = uproot.open(fname)
if treepath is None:
for treepath in ["events", "Events"] + list(source.keys()):
treepath = treepath.split(";", 1)[0]
@@ -48,24 +46,13 @@ def _load_nano_root(fname: str, treepath: str | None = None, **kwargs) -> ak.Arr
break
else:
raise ValueError(f"no default treepath determined in {fname}")
- try:
- return coffea.nanoevents.NanoEventsFactory.from_root(
- source,
- treepath=treepath,
- delayed=False,
- runtime_cache=None,
- persistent_cache=None,
- ).events()
- except:
- return uproot.open(fname)
-
- return coffea.nanoevents.NanoEventsFactory.from_root(
- source,
- treepath=treepath,
- mode="eager",
- runtime_cache=None,
- persistent_cache=None,
- ).events()
+
+ # load the tree
+ tree = source[treepath]
+
+ # use the io handler to load arrays and attach nano schema
+ chunk_pos = ChunkedIOHandler.create_chunk_position(tree.num_entries, tree.num_entries, 0)
+ return ChunkedIOHandler.read_coffea_root(tree, chunk_pos)
def load(fname: str, **kwargs) -> Any:
@@ -112,12 +99,47 @@ def list_content(data: Any) -> None:
)
ap.register("action", "help", argparse._HelpAction)
- ap.add_argument("files", metavar="FILE", nargs="*", help="one or more supported files")
- ap.add_argument("--events", "-e", action="store_true", help="assume files to contain event info")
- ap.add_argument("--hists", "-h", action="store_true", help="assume files to contain histograms")
- ap.add_argument("--treepath", "-t", type=str, help="name of the tree in ROOT files")
- ap.add_argument("--list", "-l", action="store_true", help="list contents of the loaded file")
- ap.add_argument("--help", action="help", help="show this help message and exit")
+ ap.add_argument(
+ "files",
+ metavar="FILE",
+ nargs="+",
+ help="one or more supported files",
+ )
+ ap.add_argument(
+ "--events",
+ "-e",
+ action="store_true",
+ help="assume files to contain event info",
+ )
+ ap.add_argument(
+ "--hists",
+ "-h",
+ action="store_true",
+ help="assume files to contain histograms",
+ )
+ ap.add_argument(
+ "--treepath",
+ "-t",
+ type=str,
+ help="name of the tree in ROOT files",
+ )
+ ap.add_argument(
+ "--list",
+ "-l",
+ action="store_true",
+ help="list contents of the loaded file",
+ )
+ ap.add_argument(
+ "--merge-columns",
+ "-m",
+ action="store_true",
+ help="merge columns of multiple files into one array",
+ )
+ ap.add_argument(
+ "--help",
+ action="help",
+ help="show this help message and exit",
+ )
args = ap.parse_args()
@@ -125,6 +147,9 @@ def list_content(data: Any) -> None:
"treepath": args.treepath,
}
objects = [load(fname, **load_kwargs) for fname in args.files]
+
+ if len(objects) > 1 and args.merge_columns:
+ objects = [update_ak_array(objects[0], *objects[1:])]
if len(objects) == 1:
objects = objects[0]
print("file content loaded into variable 'objects'")
diff --git a/columnflow/calibration/__init__.py b/columnflow/calibration/__init__.py
index f0ed046bf..3774e28ac 100644
--- a/columnflow/calibration/__init__.py
+++ b/columnflow/calibration/__init__.py
@@ -6,32 +6,46 @@
from __future__ import annotations
+import copy
import inspect
import law
-from columnflow.util import DerivableMeta
from columnflow.columnar_util import TaskArrayFunction
-from columnflow.types import Callable, Sequence, Any
+from columnflow.util import DerivableMeta, UNSET
+from columnflow.types import Callable, Sequence, Any, UNSET_TYPE
class TaskArrayFunctionWithCalibratorRequirements(TaskArrayFunction):
require_calibrators: Sequence[str] | set[str] | None = None
+ def __init__(self, *args, **kwargs):
+ if "require_calibrators" in kwargs or self.__class__.require_calibrators is None:
+ kwargs["require_calibrators"] = kwargs.get("require_calibrators") or []
+ elif isinstance(self.__class__.require_calibrators, (list, tuple)):
+ kwargs["require_calibrators"] = copy.copy(self.__class__.require_calibrators)
+
+ super().__init__(*args, **kwargs)
+
def _req_calibrator(self, task: law.Task, calibrator: str) -> Any:
# hook to customize how required calibrators are requested
from columnflow.tasks.calibration import CalibrateEvents
return CalibrateEvents.req_other_calibrator(task, calibrator=calibrator)
def requires_func(self, task: law.Task, reqs: dict, **kwargs) -> None:
+ super().requires_func(task=task, reqs=reqs, **kwargs)
+
# no requirements for workflows in pilot mode
if callable(getattr(task, "is_workflow", None)) and task.is_workflow() and getattr(task, "pilot", False):
return
# add required calibrators when set
if (calibs := self.require_calibrators):
- reqs["required_calibrators"] = {calib: self._req_calibrator(task, calib) for calib in calibs}
+ reqs["required_calibrators"] = {
+ calib: self._req_calibrator(task, calib)
+ for calib in law.util.make_unique(calibs)
+ }
def setup_func(
self,
@@ -41,6 +55,8 @@ def setup_func(
reader_targets: law.util.InsertableDict,
**kwargs,
) -> None:
+ super().setup_func(task=task, reqs=reqs, inputs=inputs, reader_targets=reader_targets, **kwargs)
+
if "required_calibrators" in inputs:
for calib, inp in inputs["required_calibrators"].items():
reader_targets[f"required_calibrator_{calib}"] = inp["columns"]
@@ -62,9 +78,9 @@ def calibrator(
cls,
func: Callable | None = None,
bases: tuple = (),
- mc_only: bool = False,
- data_only: bool = False,
- require_calibrators: Sequence[str] | set[str] | None = None,
+ mc_only: bool | UNSET_TYPE = UNSET,
+ data_only: bool | UNSET_TYPE = UNSET,
+ require_calibrators: Sequence[str] | set[str] | None | UNSET_TYPE = UNSET,
**kwargs,
) -> DerivableMeta | Callable:
"""
@@ -88,13 +104,13 @@ def calibrator(
"""
def decorator(func: Callable) -> DerivableMeta:
# create the class dict
- cls_dict = {
- **kwargs,
- "call_func": func,
- "mc_only": mc_only,
- "data_only": data_only,
- "require_calibrators": require_calibrators,
- }
+ cls_dict = {**kwargs, "call_func": func}
+ if mc_only is not UNSET:
+ cls_dict["mc_only"] = mc_only
+ if data_only is not UNSET:
+ cls_dict["data_only"] = data_only
+ if require_calibrators is not UNSET:
+ cls_dict["require_calibrators"] = require_calibrators
# get the module name
frame = inspect.stack()[1]
@@ -113,8 +129,7 @@ def update_cls_dict(cls_name, cls_dict, get_attr):
raise Exception(f"calibrator {cls_name} received both mc_only and data_only")
if (mc_only or data_only) and cls_dict.get("skip_func"):
raise Exception(
- f"calibrator {cls_name} received custom skip_func, but either mc_only or "
- "data_only are set",
+ f"calibrator {cls_name} received custom skip_func, but either mc_only or data_only are set",
)
if "skip_func" not in cls_dict:
diff --git a/columnflow/calibration/cms/egamma.py b/columnflow/calibration/cms/egamma.py
index 54993bf01..50dd0f108 100644
--- a/columnflow/calibration/cms/egamma.py
+++ b/columnflow/calibration/cms/egamma.py
@@ -97,7 +97,7 @@ def get_inputs(corrector, **additional_variables):
# get scaled energy error
smear = self.smear_syst_corrector.evaluate("smear", *get_inputs(self.smear_syst_corrector, pt=pt_scaled))
- energy_err_scaled = (((coll.energyErr)**2 + (coll.energy * smear)**2) * scale)**0.5
+ energy_err_scaled = (((coll.energyErr)**2 + (coll.energy * smear)**2))**0.5 * scale
# store columns
events = set_ak_column_f32(events, f"{self.collection_name}.pt", pt_scaled)
@@ -127,7 +127,7 @@ def apply_smearing(syst):
smear_factor = 1.0 + smear * rnd
pt_smeared = coll.pt * smear_factor
# get smeared energy error
- energy_err_smeared = (((coll.energyErr)**2 + (coll.energy * smear)**2) * smear_factor)**0.5
+ energy_err_smeared = (((coll.energyErr)**2 + (coll.energy * smear)**2))**0.5 * smear_factor
# return both
return pt_smeared, energy_err_smeared
@@ -160,6 +160,8 @@ def apply_smearing(syst):
@_egamma_scale_smear.init
def _egamma_scale_smear_init(self: Calibrator, **kwargs) -> None:
+ super(_egamma_scale_smear, self).init_func(**kwargs)
+
# store the config
self.cfg = self.get_scale_smear_config()
@@ -182,6 +184,8 @@ def _egamma_scale_smear_init(self: Calibrator, **kwargs) -> None:
@_egamma_scale_smear.requires
def _egamma_scale_smear_requires(self, task: law.Task, reqs: dict[str, DotDict[str, Any]], **kwargs) -> None:
+ super(_egamma_scale_smear, self).requires_func(task=task, reqs=reqs, **kwargs)
+
if "external_files" in reqs:
return
@@ -198,6 +202,14 @@ def _egamma_scale_smear_setup(
reader_targets: law.util.InsertableDict,
**kwargs,
) -> None:
+ super(_egamma_scale_smear, self).setup_func(
+ task=task,
+ reqs=reqs,
+ inputs=inputs,
+ reader_targets=reader_targets,
+ **kwargs,
+ )
+
# get and load the correction file
corr_file = self.get_correction_file(reqs["external_files"].files)
corr_set = load_correction_set(corr_file)
diff --git a/columnflow/calibration/cms/jets.py b/columnflow/calibration/cms/jets.py
index 09f381770..7337757dd 100644
--- a/columnflow/calibration/cms/jets.py
+++ b/columnflow/calibration/cms/jets.py
@@ -14,7 +14,7 @@
from columnflow.calibration.util import ak_random, propagate_met, sum_transverse
from columnflow.production.util import attach_coffea_behavior
from columnflow.util import UNSET, maybe_import, DotDict, load_correction_set
-from columnflow.columnar_util import set_ak_column, layout_ak_array, optional_column as optional
+from columnflow.columnar_util import set_ak_column, layout_ak_array, optional_column as optional, ak_concatenate_safe
from columnflow.types import TYPE_CHECKING, Any
np = maybe_import("numpy")
@@ -36,6 +36,7 @@
def get_evaluators(
correction_set: correctionlib.highlevel.CorrectionSet,
names: list[str],
+ attrs: list[dict[str, Any]] | None = None,
) -> list[Any]:
"""
Helper function to get a list of correction evaluators from a
@@ -45,6 +46,7 @@ def get_evaluators(
:param correction_set: evaluator provided by :external+correctionlib:doc:`index`
:param names: List of names of corrections to be applied
+ :param: attrs: List of dictionaries containing attributes to be added to each evaluator.
:raises RuntimeError: If a requested correction in *names* is not available
:return: List of compounded corrections, see
:external+correctionlib:py:class:`correctionlib.highlevel.CorrectionSet`
@@ -59,13 +61,27 @@ def get_evaluators(
f"\n - {name}" for name in sorted(available_keys)
))
+ if attrs and len(attrs) != len(names):
+ raise ValueError(
+ f"number of attribute dictionaries ({len(attrs)}) does not match number of evaluator names ({len(attrs)})",
+ )
+
# retrieve the evaluators
- return [
- correction_set.compound[name]
- if name in correction_set.compound
- else correction_set[name]
- for name in names
- ]
+ evaluators = []
+ for i, name in enumerate(names):
+ e = (
+ correction_set.compound[name]
+ if name in correction_set.compound
+ else correction_set[name]
+ )
+ # attach attributes if given
+ if attrs:
+ for attr, value in attrs[i].items():
+ setattr(e, attr, value)
+ # save
+ evaluators.append(e)
+
+ return evaluators
def ak_evaluate(evaluator: correctionlib.highlevel.Correction, *args) -> float:
@@ -251,6 +267,8 @@ def get_jec_config_default(self: Calibrator) -> DotDict:
get_jec_file=get_jerc_file_default,
# function to determine the jec configuration dict
get_jec_config=get_jec_config_default,
+ # function to update variables before jec corrector call
+ update_corrector_variables=(lambda self, corrector, variables: variables),
)
def jec(
self: Calibrator,
@@ -333,9 +351,6 @@ def jec(
events = set_ak_column_f32(events, f"{jet_name}.pt_raw", events[jet_name].pt * (1 - events[jet_name].rawFactor))
events = set_ak_column_f32(events, f"{jet_name}.mass_raw", events[jet_name].mass * (1 - events[jet_name].rawFactor))
- # run number for each jet required in 2023 L2L3Residual jec corrections for data
- jet_run = ak.broadcast_arrays(events.run, events[jet_name].pt_raw)[0]
-
def correct_jets(*, pt, eta, phi, area, rho, run, evaluator_key="jec"):
# variable naming convention
variable_map = {
@@ -350,13 +365,15 @@ def correct_jets(*, pt, eta, phi, area, rho, run, evaluator_key="jec"):
# apply all correctors sequentially, updating the pt each time
full_correction = ak.ones_like(pt, dtype=np.float32)
for corrector in self.evaluators[evaluator_key]:
+ # optionally update variables for this corrector call
+ _variable_map = variable_map
+ if callable(self.update_corrector_variables):
+ _variable_map = variable_map.copy()
+ _variable_map = self.update_corrector_variables(corrector, _variable_map)
# determine correct inputs (change depending on corrector)
- inputs = [
- variable_map[inp.name]
- for inp in corrector.inputs
- ]
+ inputs = [_variable_map[inp.name] for inp in corrector.inputs]
correction = ak_evaluate(corrector, *inputs)
- # update pt for subsequent correctors
+ # update pt in original variable map for subsequent correctors
variable_map["JetPt"] = variable_map["JetPt"] * correction
full_correction = full_correction * correction
@@ -380,7 +397,7 @@ def correct_jets(*, pt, eta, phi, area, rho, run, evaluator_key="jec"):
phi=events[jet_name].phi,
area=events[jet_name].area,
rho=rho,
- run=jet_run,
+ run=events.run,
evaluator_key="jec_subset_type1_met",
)
@@ -403,7 +420,7 @@ def correct_jets(*, pt, eta, phi, area, rho, run, evaluator_key="jec"):
phi=events[jet_name].phi,
area=events[jet_name].area,
rho=rho,
- run=jet_run,
+ run=events.run,
evaluator_key="jec",
)
@@ -489,6 +506,8 @@ def correct_jets(*, pt, eta, phi, area, rho, run, evaluator_key="jec"):
@jec.init
def jec_init(self: Calibrator, **kwargs) -> None:
+ super(jec, self).init_func(**kwargs)
+
jec_cfg = self.get_jec_config()
sources = self.uncertainty_sources
@@ -531,6 +550,8 @@ def jec_requires(
reqs: dict[str, DotDict[str, Any]],
**kwargs,
) -> None:
+ super(jec, self).requires_func(task=task, reqs=reqs, **kwargs)
+
if "external_files" in reqs:
return
@@ -599,6 +620,8 @@ def jec_setup(
:param inputs: Additional inputs, currently not used
:param reader_targets: TODO: add documentation
"""
+ super(jec, self).setup_func(task=task, reqs=reqs, inputs=inputs, reader_targets=reader_targets, **kwargs)
+
# import the correction sets from the external file
jec_file = self.get_jec_file(reqs["external_files"].files)
correction_set = load_correction_set(jec_file)
@@ -644,8 +667,16 @@ def make_jme_keys(names, jec=jec_cfg, is_data=self.dataset_inst.is_data):
# store the evaluators
self.evaluators = {
- "jec": get_evaluators(correction_set, jec_keys),
- "jec_subset_type1_met": get_evaluators(correction_set, jec_keys_subset_type1_met),
+ "jec": get_evaluators(
+ correction_set,
+ jec_keys,
+ attrs=[{"level": level} for level in jec_cfg.levels],
+ ),
+ "jec_subset_type1_met": get_evaluators(
+ correction_set,
+ jec_keys_subset_type1_met,
+ attrs=[{"level": level} for level in jec_cfg.levels_for_type1_met],
+ ),
"junc": dict(zip(self.uncertainty_sources, get_evaluators(correction_set, junc_keys))),
}
@@ -867,11 +898,11 @@ def jer(self: Calibrator, events: ak.Array, **kwargs) -> ak.Array:
# array with all JER scale factor variations as an additional axis
# (note: axis needs to be regular for broadcasting to work correctly)
- jer = ak.concatenate(
+ jer = ak_concatenate_safe(
[jer[v][..., None] for v in self.jer_variations + self.jec_variations],
axis=-1,
)
- jersf = ak.concatenate(
+ jersf = ak_concatenate_safe(
[jersf[v][..., None] for v in self.jer_variations + self.jec_variations],
axis=-1,
)
@@ -912,7 +943,7 @@ def jer(self: Calibrator, events: ak.Array, **kwargs) -> ak.Array:
else:
# concatenate varied pt values for broadcasting
pt_names = ["pt" for _ in self.jer_variations] + [f"pt_{jec_var}" for jec_var in self.jec_variations]
- match_pt = ak.concatenate([events[jet_name][pt_name][..., None] for pt_name in pt_names], axis=-1)
+ match_pt = ak_concatenate_safe([events[jet_name][pt_name][..., None] for pt_name in pt_names], axis=-1)
pt_relative_diff = 1 - matched_gen_jet.pt / match_pt
# test if matched gen jets are within 3 * resolution
@@ -927,9 +958,6 @@ def jer(self: Calibrator, events: ak.Array, **kwargs) -> ak.Array:
# otherwise take the stochastic ones
smear_factors = ak.where(is_matched_pt, smear_factors_scaling, smear_factors_stochastic)
- # ensure array with correctionlib output 'Nan' are set to 0.0 in the next line
- smear_factors = ak.nan_to_none(smear_factors)
-
# ensure array is not nullable (avoid ambiguity on Arrow/Parquet conversion)
smear_factors = ak.fill_none(smear_factors, 0.0)
@@ -1007,6 +1035,8 @@ def jer(self: Calibrator, events: ak.Array, **kwargs) -> ak.Array:
@jer.init
def jer_init(self: Calibrator, **kwargs) -> None:
+ super(jer, self).init_func(**kwargs)
+
# add jec_cfg for applying nominal smearing to jec variations
jec_cfg = self.get_jec_config()
jec_sources = self.jec_uncertainty_sources
@@ -1071,6 +1101,8 @@ def jer_requires(
reqs: dict[str, DotDict[str, Any]],
**kwargs,
) -> None:
+ super(jer, self).requires_func(task=task, reqs=reqs, **kwargs)
+
if "external_files" in reqs:
return
@@ -1119,6 +1151,8 @@ def jer_setup(
:param inputs: Additional inputs, currently not used.
:param reader_targets: TODO: add documentation.
"""
+ super(jer, self).setup_func(task=task, reqs=reqs, inputs=inputs, reader_targets=reader_targets, **kwargs)
+
# import the correction sets from the external file
jer_file = self.get_jer_file(reqs["external_files"].files)
correction_set = load_correction_set(jer_file)
@@ -1190,6 +1224,8 @@ def jets(self: Calibrator, events: ak.Array, **kwargs) -> ak.Array:
@jets.init
def jets_init(self: Calibrator, **kwargs) -> None:
+ super(jets, self).init_func(**kwargs)
+
# create custom jec and jer calibrators, using the jet name as the identifying value
def get_attrs(attrs):
cls_dict = {}
diff --git a/columnflow/calibration/cms/met.py b/columnflow/calibration/cms/met.py
index 9774c8006..4dbf7e463 100644
--- a/columnflow/calibration/cms/met.py
+++ b/columnflow/calibration/cms/met.py
@@ -51,7 +51,7 @@ def requires_func(self, task: law.Task, reqs: dict[str, DotDict[str, Any]], **kw
@dataclass
class METPhiConfigRun2:
- correction_set_template = r"{variable}_metphicorr_pfmet_{data_source}"
+ correction_set_template: str = r"{variable}_metphicorr_pfmet_{data_source}"
met_name: str = "MET"
keep_uncorrected: bool = False
@@ -127,6 +127,8 @@ def met_phi_run2(self: Calibrator, events: ak.Array, **kwargs) -> ak.Array:
@met_phi_run2.init
def met_phi_run2_init(self: Calibrator, **kwargs) -> None:
+ super(met_phi_run2, self).init_func(**kwargs)
+
self.met_config = self.get_met_config()
# set used columns
@@ -147,6 +149,8 @@ def met_phi_run2_setup(
reader_targets: law.util.InsertableDict,
**kwargs,
) -> None:
+ super(_met_phi_base, self).setup_func(task=task, reqs=reqs, inputs=inputs, reader_targets=reader_targets, **kwargs)
+
# create the pt and phi correctors
met_file = self.get_met_file(reqs["external_files"].files)
correction_set = load_correction_set(met_file)
@@ -281,6 +285,8 @@ def met_phi(self: Calibrator, events: ak.Array, **kwargs) -> ak.Array:
@met_phi.init
def met_phi_init(self: Calibrator, **kwargs) -> None:
+ super(met_phi, self).init_func(**kwargs)
+
self.met_config = self.get_met_config()
# set used columns
@@ -304,6 +310,8 @@ def met_phi_setup(
reader_targets: law.util.InsertableDict,
**kwargs,
) -> None:
+ super(met_phi, self).setup_func(task=task, reqs=reqs, inputs=inputs, reader_targets=reader_targets, **kwargs)
+
# load the corrector
met_file = self.get_met_file(reqs["external_files"].files)
correction_set = load_correction_set(met_file)
diff --git a/columnflow/calibration/cms/muon.py b/columnflow/calibration/cms/muon.py
index d096c94f1..5f76daaac 100644
--- a/columnflow/calibration/cms/muon.py
+++ b/columnflow/calibration/cms/muon.py
@@ -127,7 +127,6 @@ def muon_sr(
events.event,
events.luminosityBlock,
self.muon_correction_set,
- rnd_gen="np",
nested=True,
)
events = set_ak_column_f32(events, "Muon.pt", pt_scale_res_corr)
@@ -170,6 +169,8 @@ def muon_sr(
@muon_sr.init
def muon_sr_init(self: Calibrator, **kwargs) -> None:
+ super(muon_sr, self).init_func(**kwargs)
+
self.muon_cfg = self.get_muon_sr_config()
# add produced columns with unceratinties if requested
@@ -189,6 +190,8 @@ def muon_sr_requires(
reqs: dict[str, DotDict[str, Any]],
**kwargs,
) -> None:
+ super(muon_sr, self).requires_func(task=task, reqs=reqs, **kwargs)
+
if "external_files" in reqs:
return
@@ -205,6 +208,8 @@ def muon_sr_setup(
reader_targets: law.util.InsertableDict,
**kwargs,
) -> None:
+ super(muon_sr, self).setup_func(task=task, reqs=reqs, inputs=inputs, reader_targets=reader_targets, **kwargs)
+
# load the correction set
muon_sr_file = self.get_muon_sr_file(reqs["external_files"].files)
self.muon_correction_set = load_correction_set(muon_sr_file)
diff --git a/columnflow/calibration/cms/tau.py b/columnflow/calibration/cms/tau.py
index 897ebea4f..1b6e894ec 100644
--- a/columnflow/calibration/cms/tau.py
+++ b/columnflow/calibration/cms/tau.py
@@ -15,7 +15,7 @@
from columnflow.calibration import Calibrator, calibrator
from columnflow.calibration.util import propagate_met
from columnflow.util import maybe_import, load_correction_set, DotDict
-from columnflow.columnar_util import TAFConfig, set_ak_column, flat_np_view, ak_copy
+from columnflow.columnar_util import TAFConfig, set_ak_column, flat_np_view, layout_ak_array, full_like
from columnflow.types import Any
ak = maybe_import("awkward")
@@ -101,77 +101,77 @@ def tec(
if self.dataset_inst.is_data:
raise ValueError("attempt to apply tau energy corrections in data")
- # the correction tool only supports flat arrays, so convert inputs to flat np view first
- pt = flat_np_view(events.Tau.pt, axis=1)
- mass = flat_np_view(events.Tau.mass, axis=1)
- eta = flat_np_view(events.Tau.eta, axis=1)
- dm = flat_np_view(events.Tau.decayMode, axis=1)
- match = flat_np_view(events.Tau.genPartFlav, axis=1)
-
- # get the scale factors for the four supported decay modes
+ # create mask to select taus with supported decay modes
+ match = events.Tau.genPartFlav
+ dm = events.Tau.decayMode
dm_mask = (dm == 0) | (dm == 1) | (dm == 10) | (dm == 11)
- # prepare arguments
+ # prepare inputs
variable_map = {
- "pt": pt[dm_mask],
- "eta": eta[dm_mask],
+ "pt": events.Tau.pt[dm_mask],
+ "eta": events.Tau.eta[dm_mask],
"dm": dm[dm_mask],
"genmatch": match[dm_mask],
"id": self.tec_cfg.tagger,
**self.tec_cfg.corrector_kwargs,
}
- args = tuple(
+ inputs = [
variable_map[inp.name] for inp in self.tec_corrector.inputs
if inp.name in variable_map
- )
+ ]
+
+ # helper to get scales
+ def get_scales(syst):
+ scales = flat_np_view(full_like(dm_mask, 1.0, dtype=np.float32))
+ scales[flat_np_view(dm_mask)] = flat_np_view(self.tec_corrector(*inputs, syst))
+ return layout_ak_array(scales, events.Tau)
# nominal correction
- scales_nom = np.ones_like(dm_mask, dtype=np.float32)
- scales_nom[dm_mask] = self.tec_corrector(*args, "nom")
+ scales_nom = get_scales("nom")
# varied corrections
if self.with_uncertainties:
- scales_up = np.ones_like(dm_mask, dtype=np.float32)
- scales_up[dm_mask] = self.tec_corrector(*args, "up")
- scales_down = np.ones_like(dm_mask, dtype=np.float32)
- scales_down[dm_mask] = self.tec_corrector(*args, "down")
+ scales_up = get_scales("up")
+ scales_down = get_scales("down")
# custom adjustment: reset where the matching value is unhandled
reset_mask = (match < 1) | (match > 5)
- scales_nom[reset_mask] = 1.0
- if self.with_uncertainties:
- scales_up[reset_mask] = 1.0
- scales_down[reset_mask] = 1.0
+ if ak.any(reset_mask):
+ scales_nom = ak.where(reset_mask, 1.0, scales_nom)
+ if self.with_uncertainties:
+ scales_up = ak.where(reset_mask, 1.0, scales_up)
+ scales_down = ak.where(reset_mask, 1.0, scales_down)
- # create varied collections per decay mode
+ # create varied collections per decay mode, gen match, and direction
if self.with_uncertainties:
for (match_mask, match_name), _dm, (direction, scales) in itertools.product(
- [(match == 5, "jet"), ((match == 1) | (match == 3), "e")],
+ [(match == 5, "tau"), ((match == 1) | (match == 3), "e"), ((match == 2) | (match == 4), "mu")],
[0, 1, 10, 11],
[("up", scales_up), ("down", scales_down)],
):
# copy pt and mass
- pt_varied = ak_copy(events.Tau.pt)
- mass_varied = ak_copy(events.Tau.mass)
- pt_view = flat_np_view(pt_varied, axis=1)
- mass_view = flat_np_view(mass_varied, axis=1)
+ pt_flat = flat_np_view(events.Tau.pt, copy=True)
+ mass_flat = flat_np_view(events.Tau.mass, copy=True)
- # correct pt and mass for taus with that gen match and decay mode
+ # correct pt and mass using
+ # - varied scale values for taus with that gen match and decay mode
+ # - nominal scale values for the rest
mask = match_mask & (dm == _dm)
- pt_view[mask] *= scales[mask]
- mass_view[mask] *= scales[mask]
+ flat_scales = flat_np_view(ak.where(mask, scales, scales_nom))
+ pt_flat *= flat_scales
+ mass_flat *= flat_scales
# save columns
postfix = f"tec_{match_name}_dm{_dm}_{direction}"
- events = set_ak_column_f32(events, f"Tau.pt_{postfix}", pt_varied)
- events = set_ak_column_f32(events, f"Tau.mass_{postfix}", mass_varied)
+ events = set_ak_column_f32(events, f"Tau.pt_{postfix}", layout_ak_array(pt_flat, events.Tau))
+ events = set_ak_column_f32(events, f"Tau.mass_{postfix}", layout_ak_array(mass_flat, events.Tau))
# propagate changes to MET
if self.propagate_met:
met_pt_varied, met_phi_varied = propagate_met(
events.Tau.pt,
events.Tau.phi,
- pt_varied,
+ events.Tau[f"pt_{postfix}"],
events.Tau.phi,
events[self.met_name].pt,
events[self.met_name].phi,
@@ -180,11 +180,9 @@ def tec(
events = set_ak_column_f32(events, f"{self.met_name}.phi_{postfix}", met_phi_varied)
# apply the nominal correction
- # note: changes are applied to the views and directly propagate to the original ak arrays
- # and do not need to be inserted into the events chunk again
tau_sum_before = events.Tau.sum(axis=1)
- pt *= scales_nom
- mass *= scales_nom
+ events = set_ak_column_f32(events, "Tau.pt", events.Tau.pt * scales_nom)
+ events = set_ak_column_f32(events, "Tau.mass", events.Tau.mass * scales_nom)
# propagate changes to MET
if self.propagate_met:
@@ -204,6 +202,8 @@ def tec(
@tec.init
def tec_init(self: Calibrator, **kwargs) -> None:
+ super(tec, self).init_func(**kwargs)
+
self.tec_cfg = self.get_tec_config()
# add nominal met columns of propagating nominal tec
@@ -219,7 +219,7 @@ def tec_init(self: Calibrator, **kwargs) -> None:
src_fields += [f"{self.met_name}.{var}" for var in ["pt", "phi"]]
self.produces |= {
- f"{field}_tec_{{jet,e}}_dm{{0,1,10,11}}_{{up,down}}"
+ f"{field}_tec_{{tau,e,mu}}_dm{{0,1,10,11}}_{{up,down}}"
for field in src_fields
}
@@ -231,6 +231,8 @@ def tec_requires(
reqs: dict[str, DotDict[str, Any]],
**kwargs,
) -> None:
+ super(tec, self).requires_func(task=task, reqs=reqs, **kwargs)
+
if "external_files" in reqs:
return
@@ -247,6 +249,8 @@ def tec_setup(
reader_targets: law.util.InsertableDict,
**kwargs,
) -> None:
+ super(tec, self).setup_func(task=task, reqs=reqs, inputs=inputs, reader_targets=reader_targets, **kwargs)
+
# create the tec corrector
tau_file = self.get_tau_file(reqs["external_files"].files)
self.tec_corrector = load_correction_set(tau_file)[self.tec_cfg.correction_set]
diff --git a/columnflow/calibration/util.py b/columnflow/calibration/util.py
index b95e0b20a..f66a56e9b 100644
--- a/columnflow/calibration/util.py
+++ b/columnflow/calibration/util.py
@@ -58,7 +58,7 @@ def sum_transverse(pt: ak.Array, phi: ak.Array) -> tuple[ak.Array, ak.Array]:
def propagate_met(
- jet_pt1: (ak.Array),
+ jet_pt1: ak.Array,
jet_phi1: ak.Array,
jet_pt2: ak.Array,
jet_phi2: ak.Array,
diff --git a/columnflow/categorization/__init__.py b/columnflow/categorization/__init__.py
index cd98e7d9a..c1e182221 100644
--- a/columnflow/categorization/__init__.py
+++ b/columnflow/categorization/__init__.py
@@ -39,10 +39,7 @@ def categorizer(
"""
def decorator(func: Callable) -> DerivableMeta:
# create the class dict
- cls_dict = {
- **kwargs,
- "call_func": func,
- }
+ cls_dict = {**kwargs, "call_func": func}
# get the module name
frame = inspect.stack()[1]
diff --git a/columnflow/cms_util.py b/columnflow/cms_util.py
index 2e283009f..0fa72971a 100644
--- a/columnflow/cms_util.py
+++ b/columnflow/cms_util.py
@@ -20,6 +20,9 @@
#: Default root path to CAT metadata.
cat_metadata_root = "/cvmfs/cms-griddata.cern.ch/cat/metadata"
+#: Default URL of CAT metadata website.
+cat_metadata_url = "https://cms-analysis-corrections.docs.cern.ch"
+
@dataclasses.dataclass
class CATSnapshot:
@@ -73,6 +76,7 @@ class CATInfo:
pog_directories: dict[str, str] = dataclasses.field(default_factory=dict)
metadata_root: ClassVar[str] = cat_metadata_root
+ metadata_url: ClassVar[str] = cat_metadata_url
def get_era_directory(self, pog: str = "") -> str:
"""
@@ -90,6 +94,26 @@ def get_era_directory(self, pog: str = "") -> str:
era = self.pog_eras.get(pog.lower(), self.era) if pog else self.era
return f"Run{self.run}-{era}-NanoAODv{self.vnano}"
+ def get_era_url(self, pog: str = "", timestamp: str = "") -> str:
+ """
+ Returns the URL to the era directory on the CAT metadata website. When a *pog* or *timestamp* is given, the URL
+ will point to the specific POG era directory or timestamp subdirectory, respectively.
+
+ :param pog: Optional POG name to get the era URL for.
+ :param timestamp: Optional timestamp to get the URL for.
+ """
+ url_parts = [
+ self.metadata_url,
+ "corrections_era",
+ self.get_era_directory(pog=pog),
+ ]
+ if pog:
+ url_parts.append(pog.upper())
+ if timestamp:
+ url_parts.append(timestamp)
+
+ return "/".join(url_parts)
+
def get_file(self, pog: str, *paths: str | pathlib.Path) -> str:
"""
Returns the full path to a specific file or directory defined by *paths* in the CAT metadata structure for a
diff --git a/columnflow/columnar_util.py b/columnflow/columnar_util.py
index 00878e579..477626d32 100644
--- a/columnflow/columnar_util.py
+++ b/columnflow/columnar_util.py
@@ -14,6 +14,7 @@
import math
import time
import enum
+import uuid
import inspect
import threading
import dataclasses
@@ -25,15 +26,18 @@
import law
import order as od
-from columnflow.types import Sequence, Callable, Any, T, Generator, Hashable
from columnflow.util import (
UNSET, maybe_import, classproperty, DotDict, DerivableMeta, CachedDerivableMeta, Derivable, pattern_matcher,
get_source_code, real_path, freeze, get_docs_url,
)
+from columnflow.types import Sequence, Callable, Any, T, Generator, Hashable, TypeAlias, Union, TYPE_CHECKING
np = maybe_import("numpy")
ak = maybe_import("awkward")
uproot = maybe_import("uproot")
+if TYPE_CHECKING:
+ coffea = maybe_import("coffea")
+ maybe_import("coffea.nanoevents")
# loggers
@@ -965,7 +969,7 @@ def _create_matcher(routes):
ak_array = set_ak_column(
ak_array,
route,
- ak.concatenate((route.apply(ak_array), route.apply(other)), axis=-1),
+ ak_concatenate_safe((route.apply(ak_array), route.apply(other)), axis=-1),
)
elif do_add(route):
# add and reassign
@@ -1111,9 +1115,19 @@ def sorted_ak_to_root(
if not os.path.exists(os.path.dirname(path)):
os.makedirs(os.path.dirname(path))
+ # helper to linearize nested arrays
+ def pack(ak_array, prefix=""):
+ packed = {}
+ for field in ak_array.fields:
+ if ak_array[field].fields:
+ packed.update(pack(ak_array[field], prefix + field + "_"))
+ else:
+ packed[prefix + field] = ak_array[field]
+ return packed
+
# create the file
f = uproot.recreate(path)
- f[tree_name] = {field: ak_array[field] for field in ak_array.fields}
+ f[tree_name] = pack(ak_array)
f.close()
@@ -1172,6 +1186,57 @@ def mask_from_indices(indices: np.array | ak.Array, layout_array: ak.Array) -> a
return layout_ak_array(flat_mask, layout_array)
+def embed_with_mask(
+ ak_array: ak.Array,
+ embed_mask: ak.Array,
+ layout_array: ak.Array,
+ *,
+ value: Any = UNSET,
+ dtype: Any = None,
+ flat: bool = False,
+) -> ak.Array:
+ """
+ Embeds the values of an *ak_array* into a new array shaped like *layout_array* at locations defined by
+ *embed_mask*. The missing values are taken from *layout_array* when *value* is not set, and set to *value* with a
+ custom *dtype* otherwise.
+
+ Example:
+
+ :code-block:: python
+
+ data = ak.Array([1, 2])
+ embed_mask = [True, False, False, True]
+ layout_array = [w, x, y, z]
+ value = 99
+
+ embed_with_mask(data, embed_mask, layout_array, value=99)
+ # -> [1, 99, 99, 2]
+
+ :param ak_array: The array with data to embed into a new, larger array.
+ :param embed_mask: The boolean mask defining the locations to embed the values of *ak_array*.
+ :param layout_array: The array defining the layout of the output array.
+ :param value: The value to fill in for missing values when *value* is set.
+ :param dtype: The data type to use when *value* is set.
+ :param flat: Whether to return the output array in a flattened form.
+ :return: A new array shapes like *layout_array* with values of *ak_array* at locations defined by *embed_mask*.
+ """
+ # define the new array in a flattened form
+ flat_out_array = flat_np_view(
+ ak.copy(layout_array)
+ if value is UNSET
+ else full_like(layout_array, value, dtype=dtype),
+ )
+
+ # define the location mask based on the embed mask
+ broadcasted_embed_mask = ak.fill_none(full_like(layout_array, True, dtype=bool), True, axis=1) & embed_mask
+
+ # insert flat data array into the output array
+ flat_out_array[flat_np_view(broadcasted_embed_mask)] = flat_np_view(ak_array)
+
+ # optionally restructure the output array
+ return flat_out_array if flat else layout_ak_array(flat_out_array, layout_array)
+
+
def full_like(layout_array: ak.Array, value: Any, *, dtype: Any = None, **kwargs) -> ak.Array:
"""
Creates an awkward array with the same layout as *layout_array* and fills it with a constant
@@ -1300,6 +1365,10 @@ def attach_coffea_behavior(
continue
coll = ak_array[name]
+ # when info is a string, interpret as the collection in default_coffea_collections
+ if isinstance(info, str):
+ info = default_coffea_collections[info]
+
# when a check_attr is defined, do nothing in case it already exists
if info.get("check_attr") and getattr(coll, info["check_attr"], None) is not None:
continue
@@ -1335,12 +1404,22 @@ def layout_ak_array(data_array: np.array | ak.Array, layout_array: ak.Array) ->
return ak.unflatten(ak.flatten(data_array, axis=None), ak.num(layout_array, axis=1), axis=0)
-def flat_np_view(ak_array: ak.Array, axis: int | None = None) -> np.array:
+def flat_np_view(ak_array: ak.Array, axis: int | None = None, copy: bool = False) -> np.array:
"""
Takes an *ak_array* and returns a fully flattened numpy view. The flattening is applied along
*axis*. See *ak.flatten* for more info.
+
+ If *copy* is *True*, a full copy of the numpy array is returned instead of a view.
+
+ .. note::
+
+ Note that the returned numpy array is a view and in-place value assignments will not necessarily be propagated
+ back to the underlying awkward array. The conditions under which propagation can occur are limited: the
+ underlying array should not be the immediate results of a masked array and it should not have optional types.
+ However, this is implementation dependent and might change with future releases of awkward.
"""
- return np.asarray(ak.flatten(ak_array, axis=axis))
+ np_func = np.array if copy else np.asarray
+ return np_func(ak.flatten(ak_array, axis=axis))
def ak_copy(ak_array: ak.Array) -> ak.Array:
@@ -1352,6 +1431,123 @@ def ak_copy(ak_array: ak.Array) -> ak.Array:
return layout_ak_array(np.array(ak.flatten(ak_array)), ak_array)
+def ak_concatenate_safe(arrays: Sequence[ak.Array], *args, **kwargs) -> ak.Array:
+ """
+ Safe version of ``ak.concatenate`` that ensures that all arrays in *arrays* are materialized and that their masks
+ have been fully applied.
+
+ .. note::
+
+ This is necessary to avoid concatenation of masked arrays whose original data is still referrenced internally by
+ awkward and that would be concatenated as well, leading to needlessly high memory consumption.
+
+ :param arrays: The arrays to concatenate.
+ :param args: Additional positional arguments forwarded to ``ak.concatenate``.
+ :param kwargs: Additional keyword arguments forwarded to ``ak.concatenate``.
+ :return: The concatenated array.
+ """
+ return ak.concatenate(list(map(ak.to_packed, arrays)), *args, **kwargs)
+
+
+def slice_across_chunks(
+ slice_funcs: Sequence[Callable[[int, int], ak.Array | np.ndarray]],
+ sizes: Sequence[int],
+ entry_start: int | None = None,
+ entry_stop: int | None = None,
+) -> ak.Array | np.ndarray:
+ """
+ Loads multiple array chunks, concatenates them and returns the result. Each chunk is loaded by its own slicing
+ function in *slice_funcs* which receives the start and stop entry of the chunk to load. The total sizes of the
+ chunks are given in *sizes* and used for internal range calculation. By default, the full range of all chunks is
+ loaded, but this can be limited by setting *entry_start* and *entry_stop* to the global start and stop entry to
+ load, respectively.
+
+ :param slice_funcs: The slicing/loading function that provide chunks.
+ :param sizes: The total sizes per chunks for internal range calculation.
+ :param entry_start: The global start entry to load. Defaults to the start of the first chunk.
+ :param entry_stop: The global stop entry to load. Defaults to the end of the last chunk.
+ :return: The concatenated array of all loaded chunks.
+ """
+ if len(slice_funcs) != len(sizes):
+ raise ValueError(f"slice_funcs and sizes must have the same length, but got {len(slice_funcs)} and {len(sizes)}")
+
+ # collect arrays according to chunk ranges
+ arrays = []
+ chunk_ranges = law.util.chunk_slice_ranges(sizes, entry_start, entry_stop)
+ for r, slice_func in zip(chunk_ranges, slice_funcs):
+ if r is not None:
+ start, stop = r
+ arrays.append(slice_func(start, stop))
+
+ if not arrays:
+ raise ValueError(f"no array chunks loaded, ranges were {chunk_ranges}")
+
+ # check if and how to concatenate
+ if len(arrays) == 1:
+ array = arrays[0]
+ elif any(isinstance(arr, ak.Array) for arr in arrays):
+ # at least one array is awkward, so concatenate safely
+ array = ak_concatenate_safe(arrays, axis=0)
+ else:
+ array = np.concatenate(arrays, axis=0)
+ del arrays
+
+ return array
+
+
+class ArraySource(dict):
+ """
+ Wrapper around a mapping from column name to arrays, and implementing the minimal interface to be used as an array
+ source for :py:meth:`~coffea.nanoevents.NanoEventsFactory.from_preloaded`.
+ """
+
+ def __init__(self, ak_array: ak.Array | dict[str, ak.Array], object_path: str = "/Events;1") -> None:
+ # decompose ak array if needed
+ if not isinstance(ak_array, dict):
+ ak_array = flatten_ak_array(ak_array)
+ if not ak_array:
+ raise ValueError("cannot create ArraySource from empty input array")
+
+ # initialize the dict
+ super().__init__(ak_array)
+
+ # store metadata
+ self.metadata = {
+ "uuid": str(uuid.uuid4()),
+ "object_path": object_path,
+ "num_rows": len(self[list(self.keys())[0]]),
+ }
+
+
+def attach_nano_schema(
+ ak_array: ak.Array,
+ schema_cls: Any = None,
+) -> ak.Array:
+ """
+ Attaches coffea's nano schema to an *ak_array* and returns it. This is required to make use of the full nano
+ behavior.
+
+ :param ak_array: The input array.
+ :param schema_cls: The nano schema class to attach. Defaults to ``coffea.nanoevents.schemas.NanoAODSchema``.
+ :return: The array with the schema attached.
+ """
+ import coffea.nanoevents
+
+ # default schema_cls
+ if schema_cls is None:
+ schema_cls = coffea.nanoevents.schemas.NanoAODSchema
+
+ # wrap array into source object
+ source = ArraySource(ak_array)
+
+ # build the factory and return events
+ factory = coffea.nanoevents.NanoEventsFactory.from_preloaded(
+ source,
+ schemaclass=schema_cls,
+ )
+ return factory.events()
+
+
class RouteFilter(object):
"""
Shallow helper class that handles the filtering of routes in an awkward array that do (not) match those in
@@ -1638,7 +1834,7 @@ def my_other_func_init(self):
check_used_columns = True
check_produced_columns = True
_dependency_sets = {"uses", "produces"}
- log_runtime = law.config.get_expanded_boolean("analysis", "log_array_function_runtime", False)
+ log_runtime = law.config.get_expanded_bool("analysis", "log_array_function_runtime", False)
# flags for declaring inputs (via uses) or outputs (via produces)
class IOFlag(enum.Flag):
@@ -1946,7 +2142,7 @@ def deferred_init(self, instance_cache: dict | None = None) -> None:
self.pre_init_func()
# create dependencies once
- instance_cache = instance_cache or {}
+ instance_cache = instance_cache.copy() if instance_cache else {}
self.create_dependencies(instance_cache)
# run this instance's init function which might update dependent classes
@@ -3067,205 +3263,6 @@ def get_next(self) -> tuple | None:
return task
-class DaskArrayReader(object):
- """
- Class that wraps a dask_awkward array and handles chunked reading via splitting and merging of
- materialized partitions. To allow memory efficient caching in case of overlaps between
- partitions on disk and chunks to be read (possibly with different sizes) this process is
- implemented as a one-time-only read operation. Hence, in situations where particular chunks need
- to be read more than once, another instance of this class should be used.
- """
-
- class MaterializationStrategy(enum.Flag):
- """
- Flag to define which materialization strategy to follow.
- """
- SLICES = enum.auto()
- PARTITIONS = enum.auto()
-
- def __init__(
- self,
- path: str,
- open_options: dict | None = None,
- materialization_strategy: MaterializationStrategy = MaterializationStrategy.PARTITIONS,
- ) -> None:
- super().__init__()
-
- open_options = open_options or {}
-
- # backwards compatibility: when row group splitting is enbaled (the default), detect whether
- # the input file was written with support for that; a file does not support splitting in
- # case nested nodes separated by "*.list.element.*" (rather than "*.list.item.*") are found
- # (to be removed in the future)
- if open_options.get("split_row_groups"):
- try:
- nodes = ak.ak_from_parquet.metadata(path)[0]
- except:
- logger.error(f"unable to read {path}")
- raise
- cre = re.compile(r"^.+\.list\.element(|\..+)$")
- if any(map(cre.match, nodes)):
- logger.warning(
- f"while reading input parquet file from {path}, 'split_row_groups' is enabled "
- "but the file does not support it; it was probably created with an older "
- "version of columnflow which did not make use of this feature; row group "
- "splitting will be disabled, potentially leading to suboptimal performance",
- )
- open_options["split_row_groups"] = False
-
- # open the file
- import dask_awkward as dak
- self.dak_array = dak.from_parquet(path, **open_options)
- self.path = path
-
- # strategy dependent attributes and setup
- self.materialization_strategy = materialization_strategy
- if materialization_strategy == self.MaterializationStrategy.PARTITIONS:
- # fixed mapping of chunk to partition indices, created in _materialize_via_partitions
- self.chunk_to_partitions = {}
-
- # mapping of partition indices to cache information (chunks still to be handled and a
- # cached array) that changes during the read process in _materialize_via_partitions
- self.partition_cache = {
- p: DotDict(chunks=set(), array=None)
- for p in range(self.dak_array.npartitions)
- }
-
- # locks to protect against RCs during read operations by different threads
- self.chunk_to_partitions_lock = threading.Lock()
- self.partition_locks = {p: threading.Lock() for p in range(self.dak_array.npartitions)}
-
- def __del__(self) -> None:
- self.close()
-
- def __len__(self) -> int:
- if not self.dak_array.known_divisions:
- self.dak_array.eager_compute_divisions()
- return len(self.dak_array)
-
- @property
- def closed(self) -> bool:
- return self.dak_array is None
-
- def close(self) -> None:
- # free memory
- self.dak_array = None
- if getattr(self, "partition_cache", None):
- self.partition_cache.clear()
-
- def materialize(self, *args, **kwargs) -> ak.Array:
- """
- Materializes (reads from disk) a slice of the array using the configured
- :py:attr:`materialization_strategy`. All *args* and *kwargs* are forwarded to the internal
- implementations.
-
- :return: The materialized array.
- """
- if self.materialization_strategy == self.MaterializationStrategy.SLICES:
- return self._materialize_via_slices(*args, **kwargs)
-
- if self.materialization_strategy == self.MaterializationStrategy.PARTITIONS:
- return self._materialize_via_partitions(*args, **kwargs)
-
- raise NotImplementedError(
- f"unknown materialization strategy {self.materialization_strategy}",
- )
-
- def _materialize_via_slices(
- self,
- *,
- chunk_index: int,
- entry_start: int,
- entry_stop: int,
- max_chunk_size: int,
- ) -> ak.Array:
- # NOTE: calling `compute()` will materialize the data from the full dataset in memory
- # even if we only return a slice of the data.
- return self.dak_array[entry_start:entry_stop].compute()
-
- def _materialize_via_partitions(
- self,
- *,
- chunk_index: int,
- entry_start: int,
- entry_stop: int,
- max_chunk_size: int,
- ) -> ak.Array:
- # slicing on dak arrays was not supported for a long time, i.e. arr[start:stop].compute()
- # used to raise a DaskAwkwardNotImplemented, but it seems to be supported now, so the
- # following code might be obsolete, but it is kept for now as a fallback and as a potential
- # alternative in case the slicing implementation is not as efficient as the partition one,
- # reasons:
- # 1. in cf, there are typically no parquet files from external sources, but they are all
- # created within cf and thus, they typical partition sizes exactly match the desired
- # chunk size (as they were created in a chunked way and eventually merged), leading to
- # zero overhead and no caching / overlap issues
- # 2. it seems far more performant to read full partitions from disk rather than parts of
- # them (if possible at all) since meta data might have to be read in any case
- # strategy: read from disk with granularity given by partition divisions
- # - use chunk info to determine which partitions need to be read
- # - guard each read operation of a partition by locks
- # - add materialized partitions that might overlap with another chunk in a temporary cache
- # - remove cached partitions eagerly once it becomes clear that no chunk will need it
-
- # fill the chunk -> partitions mapping once
- with self.chunk_to_partitions_lock:
- if not self.chunk_to_partitions:
- if not self.dak_array.known_divisions:
- self.dak_array.eager_compute_divisions()
- divs = tuple(map(int, self.dak_array.divisions))
- # note: a hare-and-tortoise algorithm could be possible to get the mapping with less
- # than n^2 complexity, but for our case with ~30 chunks this should be ok (for now)
- n_chunks = int(math.ceil(len(self) / max_chunk_size))
- # in case there are no entries, ensure that at least one empty chunk is created
- for _chunk_index in range(max(n_chunks, 1)):
- _entry_start = _chunk_index * max_chunk_size
- _entry_stop = min(_entry_start + max_chunk_size, len(self))
- partitions = []
- for p, (p_start, p_stop) in enumerate(zip(divs[:-1], divs[1:])):
- # note: check strict increase of chunk size to accommodate zero-length size
- if p_stop <= _entry_start < _entry_stop:
- continue
- if p_start >= _entry_stop > _entry_start:
- break
- partitions.append(p)
- self.partition_cache[p].chunks.add(_chunk_index)
- self.chunk_to_partitions[_chunk_index] = partitions
-
- # read partitions one at a time and store parts that make up the chunk for concatenation
- parts = []
- for p in self.chunk_to_partitions[chunk_index]:
- # obtain the array
- with self.partition_locks[p]:
- # remove this chunk from the list of chunks to be handled
- self.partition_cache[p].chunks.remove(chunk_index)
-
- if self.partition_cache[p].array is None:
- arr = self.dak_array.partitions[p].compute()
- # add to cache when there is a chunk left that will need it
- if self.partition_cache[p].chunks:
- self.partition_cache[p].array = arr
- else:
- arr = self.partition_cache[p].array
- # remove from cache when there is no chunk left that would need it
- if not self.partition_cache[p].chunks:
- self.partition_cache[p].array = None
-
- # add part for concatenation using entry info
- div_start, div_stop = self.dak_array.divisions[p:p + 2]
- part_start = max(entry_start - div_start, 0)
- part_stop = min(entry_stop - div_start, div_stop - div_start)
- parts.append(arr[part_start:part_stop])
-
- # construct the full array
- arr = parts[0] if len(parts) == 1 else ak.concatenate(parts, axis=0)
-
- # cleanup
- del parts
-
- return arr
-
-
class ChunkedParquetReader(object):
"""
Class that wraps a parquet file containing an awkward array and handles chunked reading via splitting and merging of
@@ -3276,12 +3273,10 @@ class ChunkedParquetReader(object):
def __init__(self, path: str, open_options: dict | None = None) -> None:
super().__init__()
- if not open_options:
- open_options = {}
# store attributes
self.path = path
- self.open_options = open_options.copy()
+ self.open_options = open_options.copy() if open_options else {}
# open and store meta data with updated open options
# (when closing the reader, this attribute is set to None)
@@ -3292,7 +3287,7 @@ def __init__(self, path: str, open_options: dict | None = None) -> None:
try:
self.metadata = ak.metadata_from_parquet(path, **meta_options)
except:
- logger.error(f"unable to read {path}")
+ logger.error(f"unable to read meta data with options {meta_options} from {path}")
raise
# extract row group sizes for chunked reading
@@ -3380,7 +3375,7 @@ def materialize(
self.group_cache[g].chunks.remove(chunk_index)
if self.group_cache[g].array is None:
- arr = ak.from_parquet(self.path, row_groups=[g], **self.open_options)
+ arr = law.awkward.from_parquet(self.path, row_groups=[g], **self.open_options)
# add to cache when there is a chunk left that will need it
if self.group_cache[g].chunks:
self.group_cache[g].array = arr
@@ -3397,7 +3392,7 @@ def materialize(
parts.append(arr[part_start:part_stop])
# construct the full array
- arr = parts[0] if len(parts) == 1 else ak.concatenate(parts, axis=0)
+ arr = ak.to_packed(parts[0]) if len(parts) == 1 else ak_concatenate_safe(parts, axis=0)
# cleanup
del parts
@@ -3405,6 +3400,16 @@ def materialize(
return arr
+# alias to summarize possible root handler input types for read_*_root methods below
+RootHandlerInput: TypeAlias = Union[
+ str,
+ uproot.ReadOnlyDirectory,
+ uproot.TTree,
+ tuple[str, str],
+ tuple[uproot.ReadOnlyDirectory, str],
+]
+
+
class ChunkedIOHandler(object):
"""
Allows reading one or multiple files and iterating through chunks of their content with
@@ -3498,7 +3503,7 @@ def __init__(
read_options: dict | Sequence[dict] | None = None,
read_columns: set | Sequence[set] | None = None,
iter_message: str = "handling chunk {pos.index}",
- debug: bool = law.config.get_expanded_boolean("analysis", "chunked_io_debug", False),
+ debug: bool = law.config.get_expanded_bool("analysis", "chunked_io_debug", False),
):
super().__init__()
@@ -3612,20 +3617,22 @@ def get_source_handler(
- "dask_awkward_parquet"
"""
if source_type is None:
- if isinstance(source, uproot.ReadOnlyDirectory):
+ # perform checks with the first source in case a list is given
+ first_source = source[0] if isinstance(source, list) and source else source
+ if isinstance(first_source, uproot.ReadOnlyDirectory):
# uproot file
source_type = "uproot_root"
- elif isinstance(source, str):
+ elif isinstance(first_source, str):
# file path, guess based on extension
- if source.endswith(".root"):
+ if first_source.endswith(".root"):
# priotize coffea nano events
source_type = "coffea_root"
- elif source.endswith(".parquet"):
+ elif first_source.endswith(".parquet"):
# prioritize non-dask awkward reader
source_type = "awkward_parquet"
if not source_type:
- raise Exception(f"could not determine source_type from source '{source}'")
+ raise Exception(f"could not determine source_type from source '{first_source}'")
if source_type == "uproot_root":
return cls.SourceHandler(
@@ -3641,13 +3648,6 @@ def get_source_handler(
cls.close_coffea_root,
cls.read_coffea_root,
)
- if source_type == "coffea_parquet":
- return cls.SourceHandler(
- source_type,
- cls.open_coffea_parquet,
- cls.close_coffea_parquet,
- cls.read_coffea_parquet,
- )
if source_type == "awkward_parquet":
return cls.SourceHandler(
source_type,
@@ -3668,258 +3668,150 @@ def get_source_handler(
@classmethod
def open_uproot_root(
cls,
- source: (
- str |
- uproot.ReadOnlyDirectory |
- tuple[str, str] |
- tuple[uproot.ReadOnlyDirectory, str]
- ),
+ source: RootHandlerInput | list[RootHandlerInput],
+ *,
open_options: dict | None = None,
read_columns: set[str | Route] | None = None,
- ) -> tuple[uproot.TTree, int]:
- """
- Opens an uproot tree from a root file at *source* and returns a 2-tuple *(tree, entries)*.
- *source* can be the path of the file, an already opened, readable uproot file (assuming the
- tree is called "Events"), or a 2-tuple whose second item defines the name of the tree to be
- loaded. When a new file is opened, it receives *open_options*. Passing *read_columns* has no
- effect.
- """
- tree_name = "Events"
- if isinstance(source, tuple) and len(source) == 2:
- source, tree_name = source
- if isinstance(source, str):
- # default open options
- open_options = open_options or {}
- open_options["object_cache"] = None
- open_options["array_cache"] = None
- open_options.setdefault("decompression_executor", None)
- open_options.setdefault("interpretation_executor", None)
- f = uproot.open(source, **open_options)
- tree = f[tree_name]
- elif isinstance(source, uproot.ReadOnlyDirectory):
- tree = source[tree_name]
- else:
+ ) -> tuple[uproot.TTree | list[uproot.TTree], int]:
+ """
+ Opens one or multiple uproot trees from root files at *source* and returns a 2-tuple *(tree(s), entries)*.
+ *source* can be the path of the file, an already opened, readable uproot file (assuming the tree is called
+ "Events"), an uproot tree, a 2-tuple whose second item defines the name of the tree to be loaded, or a list of
+ these. When a new file is opened, it receives *open_options*. Passing *read_columns* has no effect.
+ """
+ # default open options
+ open_options = open_options.copy() if open_options else {}
+ open_options["object_cache"] = None
+ open_options["array_cache"] = None
+ open_options.setdefault("decompression_executor", None)
+ open_options.setdefault("interpretation_executor", None)
+
+ # helper to extract the correct tree from the source
+ def detect_tree(source: RootHandlerInput) -> uproot.TTree:
+ if isinstance(source, uproot.TTree):
+ return source
+ tree_name = "Events"
+ if isinstance(source, tuple) and len(source) == 2:
+ source, tree_name = source
+ if isinstance(source, str):
+ f = uproot.open(source, **open_options)
+ tree = f[tree_name]
+ elif isinstance(source, uproot.ReadOnlyDirectory):
+ tree = source[tree_name]
+ else:
+ raise Exception(f"'{source}' cannot be opened as uproot_root")
+ return tree
+
+ # convert source to list to treat cases identically
+ is_multi = isinstance(source, list)
+ if not source:
raise Exception(f"'{source}' cannot be opened as uproot_root")
+ trees = list(map(detect_tree, law.util.make_list(source)))
+ n_entries = sum(tree.num_entries for tree in trees)
- return (tree, tree.num_entries)
+ return (trees if is_multi else trees[0], n_entries)
@classmethod
def close_uproot_root(
cls,
- source_object: uproot.TTree,
+ source_object: uproot.TTree | list[uproot.TTree],
) -> None:
"""
Closes the file that contains the TTree *source_object*.
"""
- f = getattr(source_object, "file", None)
- if f is not None:
- f.close()
+ for tree in law.util.make_list(source_object):
+ f = getattr(tree, "file", None)
+ if f is not None:
+ f.close()
@classmethod
def read_uproot_root(
cls,
- source_object: uproot.TTree,
+ source_object: uproot.TTree | list[uproot.TTree],
chunk_pos: ChunkPosition,
+ *,
read_options: dict | None = None,
read_columns: set[str | Route] | None = None,
) -> ak.Array:
"""
- Given an uproot TTree *source_object*, returns an awkward array chunk referred to by
- *chunk_pos*. *read_options* are passed to ``uproot.TTree.arrays``. *read_columns* are
- converted to strings and, if not already present, added as field ``filter_name`` to
- *read_options*.
+ Given one or multiple uproot TTrees *source_object*, returns an awkward array chunk referred to by *chunk_pos*.
+ *read_options* are passed to ``uproot.TTree.arrays``. *read_columns* are converted to strings and, if not
+ already present, added as field ``filter_name`` to *read_options*.
"""
# default read options
- read_options = read_options or {}
+ read_options = read_options.copy() if read_options else {}
read_options["array_cache"] = None
+ read_options.pop("how", None)
# inject read_columns
if read_columns and "filter_name" not in read_options:
- filter_name = [Route(s).string_nano_column for s in read_columns]
- read_options["filter_name"] = filter_name
-
- chunk = source_object.arrays(
- entry_start=chunk_pos.entry_start,
- entry_stop=chunk_pos.entry_stop,
- **read_options,
- )
-
- return chunk
-
- @classmethod
- def open_coffea_root(
- cls,
- source: (
- str |
- uproot.ReadOnlyDirectory |
- tuple[str, str] |
- tuple[uproot.ReadOnlyDirectory, str]
- ),
- open_options: dict | None = None,
- read_columns: set[str | Route] | None = None,
- ) -> tuple[tuple[uproot.ReadOnlyDirectory, str], int]:
- """
- Opens an uproot file at *source* for subsequent processing with coffea and returns a 2-tuple
- *((uproot file, tree name), tree entries)*. *source* can be the path of the file, an already
- opened, readable uproot file (assuming the tree is called "Events"), or a 2-tuple whose
- second item defines the name of the tree to be loaded. *open_options* are forwarded to
- ``uproot.open`` if a new file is opened. Passing *read_columns* has no effect.
- """
- tree_name = "Events"
- if isinstance(source, tuple) and len(source) == 2:
- source, tree_name = source
- if isinstance(source, str):
- # default open options
- open_options = open_options or {}
- open_options["object_cache"] = None
- open_options["array_cache"] = None
- source = uproot.open(source, **open_options)
- tree = source[tree_name]
- elif isinstance(source, uproot.ReadOnlyDirectory):
- tree = source[tree_name]
- else:
- raise Exception(f"'{source}' cannot be opened as coffea_root")
-
- return ((source, tree_name), tree.num_entries)
-
- @classmethod
- def close_coffea_root(
- cls,
- source_object: tuple[str | uproot.ReadOnlyDirectory, str],
- ) -> None:
- """
- Closes a ROOT file referred to by the first item in a 2-tuple *source_object* if it is a
- ``ReadOnlyDirectory``. In case a string is passed, this method does nothing.
- """
- close = getattr(source_object[0], "close", None)
- if callable(close):
- close()
-
- @classmethod
- def read_coffea_root(
- cls,
- source_object: tuple[str | uproot.ReadOnlyDirectory, str],
- chunk_pos: ChunkPosition,
- read_options: dict | None = None,
- read_columns: set[str | Route] | None = None,
- ) -> ak.Array:
- """
- Given a file location or opened uproot file, and a tree name in a 2-tuple *source_object*,
- returns an awkward array chunk referred to by *chunk_pos*, assuming nanoAOD structure.
- *read_options* are passed to ``coffea.nanoevents.NanoEventsFactory.from_root``.
- *read_columns* are converted to strings and, if not already present, added as nested fields
- ``iteritems_options.filter_name`` to *read_options*.
- """
- import coffea.nanoevents
-
- # default read options
- read_options = read_options or {}
- read_options["mode"] = "eager"
- read_options["runtime_cache"] = None
- read_options["persistent_cache"] = None
-
- # inject read_columns
- if read_columns and (
- "iteritems_options" not in read_options or
- "filter_name" not in read_options["iteritems_options"]
- ):
- filter_name = [Route(s).string_nano_column for s in read_columns]
-
+ routes = list(map(Route, read_columns))
+ filter_name = [r.string_nano_column for r in routes]
# add names prefixed with an 'n' to the list of columns to read
# (needed to construct the nested list structure of jagged columns)
- maybe_jagged_fields = {Route(s)[0] for s in read_columns}
- filter_name.extend(
- f"n{field}"
- for field in maybe_jagged_fields
- )
+ maybe_jagged = {r[0] for r in routes if len(r) > 1}
+ filter_name.extend(f"n{f}" for f in maybe_jagged)
+ # add to read options
+ read_options["filter_name"] = filter_name
- # filter on these column names when reading
- read_options.setdefault("iteritems_options", {})["filter_name"] = filter_name
+ # helper to read from a single tree
+ def load(tree, entry_start, entry_stop):
+ return tree.arrays(entry_start=entry_start, entry_stop=entry_stop, **read_options)
- # read the events chunk into memory
- _source_object, tree_name = source_object
- chunk = coffea.nanoevents.NanoEventsFactory.from_root(
- _source_object,
- treepath=tree_name,
- entry_start=chunk_pos.entry_start,
- entry_stop=chunk_pos.entry_stop,
- **read_options,
- ).events()
+ # distinguish multi and single source cases
+ if isinstance(source_object, list):
+ tree_sizes = [tree.num_entries for tree in source_object]
+ slice_funcs = [partial(load, tree) for tree in source_object]
+ chunk = slice_across_chunks(slice_funcs, tree_sizes, chunk_pos.entry_start, chunk_pos.entry_stop)
+ else:
+ chunk = load(source_object, chunk_pos.entry_start, chunk_pos.entry_stop)
return chunk
@classmethod
- def open_coffea_parquet(
+ def open_coffea_root(
cls,
- source: str,
+ source: RootHandlerInput | list[RootHandlerInput],
+ *,
open_options: dict | None = None,
read_columns: set[str | Route] | None = None,
- ) -> tuple[str, int]:
+ ) -> tuple[tuple[uproot.ReadOnlyDirectory | list[uproot.ReadOnlyDirectory], str], int]:
"""
- Given a parquet file located at *source*, returns a 2-tuple *(source, entries)*. Passing
- *open_options* or *read_columns* has no effect.
+ Same as :py:meth:`open_uproot_root`.
"""
- import pyarrow.parquet as pq
-
- return (source, pq.ParquetFile(source).metadata.num_rows)
+ return cls.open_uproot_root(source, open_options=open_options, read_columns=read_columns)
@classmethod
- def close_coffea_parquet(
+ def close_coffea_root(
cls,
- source_object: str,
+ source_object: uproot.TTree | list[uproot.TTree],
) -> None:
"""
- This is a placeholder method and has no effect.
+ Same as :py:meth:`close_uproot_root`.
"""
- return
+ return cls.close_uproot_root(source_object)
@classmethod
- def read_coffea_parquet(
+ def read_coffea_root(
cls,
- source_object: str,
+ source_object: uproot.TTree | list[uproot.TTree],
chunk_pos: ChunkPosition,
+ *,
read_options: dict | None = None,
read_columns: set[str | Route] | None = None,
) -> ak.Array:
"""
- Given a the location of a parquet file *source_object*, returns an awkward array chunk
- referred to by *chunk_pos*, assuming nanoAOD structure. *read_options* are passed to
- ``coffea.nanoevents.NanoEventsFactory.from_parquet``. *read_columns* are converted to
- strings and, if not already present, added as nested field
- ``parquet_options.read_dictionary`` to *read_options*.
+ Same as :py:meth:`read_uproot_root`, but attaches coffea's nano schema to the returned array.
"""
- import coffea.nanoevents
-
- # default read options
- read_options = read_options or {}
- read_options["mode"] = "eager"
- read_options["runtime_cache"] = None
- read_options["persistent_cache"] = None
-
- # inject read_columns
- if read_columns and (
- "parquet_options" not in read_options or
- "read_dictionary" not in read_options["parquet_options"]
- ):
- read_dictionary = [Route(s).string_column for s in read_columns]
-
- # add names prefixed with an 'n' to the list of columns to read
- # (needed to construct the nested list structure of jagged columns)
- maybe_jagged_fields = {Route(s)[0] for s in read_columns}
- read_dictionary.extend(
- f"n{field}"
- for field in maybe_jagged_fields
- )
+ chunk = cls.read_uproot_root(
+ source_object=source_object,
+ chunk_pos=chunk_pos,
+ read_options=read_options,
+ read_columns=read_columns,
+ )
- read_options.setdefault("parquet_options", {})["read_dictionary"] = read_dictionary
-
- # read the events chunk into memory
- chunk = coffea.nanoevents.NanoEventsFactory.from_parquet(
- source_object,
- entry_start=chunk_pos.entry_start,
- entry_stop=chunk_pos.entry_stop,
- **read_options,
- ).events()
+ # attach nano schema
+ chunk = attach_nano_schema(chunk)
return chunk
@@ -3941,7 +3833,7 @@ def open_awkward_parquet(
raise Exception(f"'{source}' cannot be opened as awkward_parquet")
# default open options
- open_options = open_options or {}
+ open_options = open_options.copy() if open_options else {}
# inject read_columns
if read_columns and "columns" not in open_options:
@@ -3989,7 +3881,7 @@ def open_dask_awkward_parquet(
source: str,
open_options: dict | None = None,
read_columns: set[str | Route] | None = None,
- ) -> tuple[DaskArrayReader, int]:
+ ) -> tuple["DaskArrayReader", int]: # noqa: F821
"""
Opens a parquet file saved at *source*, loads the content as an dask awkward array, wrapped by a
:py:class:`DaskArrayReader`, and returns a 2-tuple *(reader, length)*.
@@ -4009,32 +3901,31 @@ def open_dask_awkward_parquet(
open_options["columns"] = filter_name
# load the array wrapper
- reader = DaskArrayReader(source, open_options)
+ reader = DaskArrayReader(source, open_options) # noqa: F821
return (reader, len(reader))
@classmethod
def close_dask_awkward_parquet(
cls,
- source_object: DaskArrayReader,
+ source_object: "DaskArrayReader", # noqa: F821
) -> None:
"""
- Closes the dask array wrapper referred to by *source_object*.
+ Closes the chunked parquet reader referred to by *source_object*.
"""
source_object.close()
@classmethod
def read_dask_awkward_parquet(
cls,
- source_object: DaskArrayReader,
+ source_object: ChunkedParquetReader,
chunk_pos: ChunkedIOHandler.ChunkPosition,
read_options: dict | None = None,
read_columns: set[str | Route] | None = None,
) -> ak.Array:
"""
- Given a :py:class:`DaskArrayReader` *source_object*, returns the chunk referred to by
- *chunk_pos* as a full copy loaded into memory. Passing neither *read_options* nor
- *read_columns* has an effect.
+ Given a :py:class:`ChunkedParquetReader` *source_object*, returns the chunk referred to by *chunk_pos* as a
+ full copy loaded into memory. Passing neither *read_options* nor *read_columns* has an effect.
"""
# get the materialized ak array for that chunk
return source_object.materialize(
diff --git a/columnflow/config_util.py b/columnflow/config_util.py
index d728f5800..7b4646756 100644
--- a/columnflow/config_util.py
+++ b/columnflow/config_util.py
@@ -9,15 +9,16 @@
__all__ = []
import re
+import pathlib
import dataclasses
import itertools
-from collections import OrderedDict, defaultdict
+import collections
import law
import order as od
-from columnflow.util import maybe_import, get_docs_url
from columnflow.columnar_util import flat_np_view, layout_ak_array
+from columnflow.util import maybe_import, get_docs_url, CacheBase, PersistentCache
from columnflow.types import Callable, Any, Sequence, Literal
ak = maybe_import("awkward")
@@ -49,7 +50,7 @@ def get_events_from_categories(
f"{get_events_from_categories.__name__} requires the 'category_ids' field to be present",
)
- categories = law.util.make_list(categories)
+ categories = law.util.make_unique(law.util.make_list(categories))
if config_inst:
# get category insts
categories = [config_inst.get_category(cat) for cat in categories]
@@ -249,7 +250,7 @@ def get_datasets_from_process(
return law.util.make_unique(dataset_insts)
# at this point, strategy is exclusive or exclusive_strict
- dataset_insts_dict: OrderedDict[str, od.Dataset] = OrderedDict()
+ dataset_insts_dict: collections.OrderedDict[str, od.Dataset] = collections.OrderedDict()
for process_inst, _, child_insts in root_inst.walk_processes(include_self=True, algo="dfs_post"):
# check if child processes have matched datasets already
if child_insts:
@@ -263,7 +264,7 @@ def get_datasets_from_process(
continue
# at this point, the process itself must be checked,
# so remove potentially found datasets of children
- dataset_insts_dict = OrderedDict({
+ dataset_insts_dict = collections.OrderedDict({
child_inst: _dataset_insts
for child_inst, _dataset_insts in dataset_insts_dict.items()
if child_inst not in child_insts
@@ -371,7 +372,7 @@ def group_shifts(
An exception is raised in case a shift source is represented only by its up or down shift.
"""
nominal = None
- grouped = defaultdict(lambda: [None, None])
+ grouped = collections.defaultdict(lambda: [None, None])
up_sources = set()
down_sources = set()
@@ -399,22 +400,43 @@ def expand_shift_sources(shifts: Sequence[str] | set[str]) -> list[str]:
.. code-block:: python
- expand_shift_sources(["jes", "jer_up"])
- # -> ["jes_up", "jes_down", "jer_up"]
+ expand_shift_sources(["jes", "jer_up", "nominal"])
+ # -> ["jes_up", "jes_down", "jer_up", "nominal"]
"""
_shifts = []
for shift in shifts:
- try:
- od.Shift.split_name(shift)
+ if shift == od.Shift.NOMINAL:
_shifts.append(shift)
- except ValueError as e:
- if not isinstance(shift, str):
- raise e
- _shifts.extend([f"{shift}_{od.Shift.UP}", f"{shift}_{od.Shift.DOWN}"])
+ else:
+ try:
+ od.Shift.split_name(shift)
+ _shifts.append(shift)
+ except ValueError as e:
+ if not isinstance(shift, str):
+ raise e
+ _shifts.extend([f"{shift}_{od.Shift.UP}", f"{shift}_{od.Shift.DOWN}"])
return law.util.make_unique(_shifts)
+class CategoryIDCache(PersistentCache):
+
+ def __init__(self, cache_path: od.Config | str | pathlib.Path | law.FileSystemFileTarget) -> None:
+ # when a config object is given, store the cache in the law home dir
+ if isinstance(cache_path, od.Config):
+ cache_path = law.config.law_home_path(f"category_ids_{cache_path.name}.json")
+
+ super().__init__(cache_path)
+
+ def open(self) -> CategoryIDCache:
+ super().open()
+
+ # synchronize the maximum category id counter in order
+ od.Category._max_id = max(self.cache.values(), default=0)
+
+ return self
+
+
def create_category_id(
config: od.Config,
category_name: str,
@@ -433,7 +455,7 @@ def create_category_id(
subsequently in an array, please be aware that values 8 or more require a ``np.int64``.
"""
# create the hash
- h = law.util.create_hash((config.name, config.id, category_name, salt), l=hash_len)
+ h = law.util.create_hash((config.name, config.id, category_name, salt), hash_len)
h = int(h, base=16)
# add an offset to ensure that are hashes are above a threshold
@@ -449,27 +471,50 @@ def add_category(
**kwargs,
) -> od.Category:
"""
- Creates a :py:class:`order.Category` instance by forwarding all *kwargs* to its constructor,
- adds it to a *parent* object. such as a :py:class:`order.Config` or an other
- :py:class:`order.Category`, and returns it. When *kwargs* do not contain a field *id*,
- :py:func:`create_category_id` is used to create one.
+ Creates a :py:class:`order.Category` instance by forwarding all *kwargs* to its constructor, adds it to a *parent*
+ object. such as a :py:class:`order.Config` or an other :py:class:`order.Category`, and returns it.
+
+ When *kwargs* do not contain a field *id*, :py:func:`create_category_id` is used to create one. Additionally, if the
+ field *id* is given as a :py:class:`CacheBase` object, it is used to cache the created id for the given category
+ name.
:param config: :py:class:`order.Config` object for which the category is created.
:param parent: Parent object to which the category is added. If *None*, *config* is used.
:param kwargs: Keyword arguments forwarded to the category constructor.
:return: The newly created category instance.
"""
+ # name must be given
if "name" not in kwargs:
fields = ",".join(map(str, kwargs))
raise ValueError(f"a field 'name' is required to create a category, got '{fields}'")
- if "id" not in kwargs:
+ # default id
+ if kwargs.get("id") is None:
kwargs["id"] = create_category_id(config, kwargs["name"])
+ # id can be a cache object
+ id_cache = None
+ if isinstance(kwargs["id"], CacheBase):
+ id_cache = kwargs["id"]
+ # use the cached id, or use "+" to auto-increment and store the id below
+ kwargs["id"] = (
+ id_cache.get(kwargs["name"])
+ if (id_cached := id_cache.has(kwargs["name"]))
+ else od.UniqueObject.AUTO_ID
+ )
+
+ # default parent
if parent is None:
parent = config
- return parent.add_category(**kwargs)
+ # create the category
+ category = parent.add_category(**kwargs)
+
+ # cache if requested
+ if id_cache and not id_cached:
+ id_cache.set(category.name, category.id)
+
+ return category
@dataclasses.dataclass
@@ -526,14 +571,7 @@ def create_category_combinations(
columnflow to determine whether the summation over specific categories is valid or may result in under- or
over-counting when combining leaf categories. These checks may be performed by other functions and tools based
on information derived from groups and stored in auxiliary fields of the newly created categories.
- Given a *config* object and sequences of *categories* in a dict, creates all combinations of possible leaf
- categories at different depths, connects them with parent - child relations (see :py:class:`order.Category`) and
- returns the number of newly created categories.
- *categories* should be a dictionary that maps string names to :py:class:`CategoryGroup` objects which are thin
- wrappers around sequences of categories (objects or names). Group names (dictionary keys) are used as keyword
- arguments in a callable *name_fn* that is supposed to return the name of newly created categories (see example
- below).
All intermediate layers of categories can be built and connected automatically to one another by parent - child
category relations. The exact behavior is controlled by *parent_mode*:
@@ -551,7 +589,8 @@ def create_category_combinations(
*kwargs_fn*. This function is called with the categories (in a dictionary, mapped to the sequence names as given in
*categories*) that contribute to the newly created category and should return a dictionary. If the fields ``"id"``
and ``"selection"`` are missing, they are filled with reasonable defaults leading to a auto-generated, deterministic
- id and a list of all parent selection statements.
+ id and a list of all parent selection statements. Similar to :py:func:`add_category`, the field ``"id"`` can be set
+ to a :py:class:`CacheBase` object to cache the created id for the given category name.
If the name of a new category is already known to *config* it is skipped unless *skip_existing* is *False*. In
addition, *skip_fn* can be a callable that receives a dictionary mapping group names to categories that represents
@@ -676,9 +715,24 @@ def kwargs_fn(categories):
if "selection" not in kwargs:
kwargs["selection"] = [c.selection for c in root_cats.values()]
+ # use cache id if given
+ id_cache = None
+ if isinstance(kwargs["id"], CacheBase):
+ id_cache = kwargs["id"]
+ # use the cached id, or use "+" to auto-increment and store the id below
+ kwargs["id"] = (
+ id_cache.get(cat_name)
+ if (id_cached := id_cache.has(cat_name))
+ else od.UniqueObject.AUTO_ID
+ )
+
# create the new category
cat = od.Category(name=cat_name, **kwargs)
- created_categories[cat_name] = cat
+ created_categories[cat.name] = cat
+
+ # cache if requested
+ if id_cache and not id_cached:
+ id_cache.set(cat.name, cat.id)
# add a tag to denote this category was auto-created
cat.add_tag("auto_created_by_combinations")
@@ -687,9 +741,9 @@ def kwargs_fn(categories):
if isinstance(kwargs["id"], int):
if kwargs["id"] in unique_ids_cache:
matching_cat = config.get_category(kwargs["id"])
- if matching_cat.name != cat_name:
+ if matching_cat.name != cat.name:
raise ValueError(
- f"non-unique category id '{kwargs['id']}' for '{cat_name}' has already been used for "
+ f"non-unique category id '{kwargs['id']}' for '{cat.name}' has already been used for "
f"category '{matching_cat.name}'",
)
unique_ids_cache.add(kwargs["id"])
diff --git a/columnflow/hist_util.py b/columnflow/hist_util.py
index 4efd3c73d..ce9335c09 100644
--- a/columnflow/hist_util.py
+++ b/columnflow/hist_util.py
@@ -9,12 +9,14 @@
__all__ = []
import functools
+import operator
+
import law
import order as od
-from columnflow.columnar_util import flat_np_view
+from columnflow.columnar_util import flat_np_view, layout_ak_array
from columnflow.util import maybe_import
-from columnflow.types import TYPE_CHECKING, Any, Sequence
+from columnflow.types import TYPE_CHECKING, Any, Sequence, Callable
np = maybe_import("numpy")
ak = maybe_import("awkward")
@@ -71,10 +73,59 @@ def allows_shift(ax) -> bool:
# correct last bin values
for ax in correct_last_bin_axes:
- right_egde_mask = ak.flatten(data[ax.name], axis=None) == ax.edges[-1]
+ flat_values = flat_np_view(data[ax.name])
+ right_egde_mask = flat_values == ax.edges[-1]
if np.any(right_egde_mask):
- data[ax.name] = ak.copy(data[ax.name])
- flat_np_view(data[ax.name])[right_egde_mask] -= ax.widths[-1] * 1e-5
+ flat_values = ak.where(right_egde_mask, flat_values - ax.widths[-1] * 1e-5, flat_values)
+ data[ax.name] = (
+ flat_values
+ if data[ax.name].ndim == 1
+ else layout_ak_array(flat_values, data[ax.name])
+ )
+
+ # check if conversion to records is needed
+ arr_types = (ak.Array, np.ndarray)
+ vals = list(data.values())
+ convert = (
+ # values is a mixture of singular and array types
+ (any(isinstance(v, arr_types) for v in vals) and not all(isinstance(v, arr_types) for v in vals)) or
+ # values contain at least one array with more than one dimension
+ any(isinstance(v, arr_types) and v.ndim != 1 for v in vals)
+ )
+
+ # actual conversion
+ if convert:
+ # in case there is more than one variable axis that will be filled with data with ndim > 1, we need to build
+ # object level combinations
+ var_axes = [
+ ax for ax in h.axes
+ if isinstance(ax, (hist.axis.Variable, hist.axis.Integer)) and ax.name in data and data[ax.name].ndim > 1
+ ]
+ if len(var_axes) > 1:
+ # build pairwise combinations of obj level variables
+ combi_key = "cf_hist_fill_obj_lvl_combinations"
+ data[combi_key] = ak.zip({ax.name: data.pop(ax.name) for ax in var_axes})
+
+ # build event level combinations of all remaining axes
+ arrays = ak.cartesian(data)
+
+ # unpack entries of the object level combinations
+ unpacked_dict = {}
+ for field in arrays.fields:
+ if field == combi_key:
+ for ax in var_axes:
+ unpacked_dict[ax.name] = arrays[field][ax.name]
+ else:
+ unpacked_dict[field] = arrays[field]
+ unpacked_arrays = ak.zip(unpacked_dict)
+ # flatten
+ data = {field: ak.flatten(unpacked_arrays[field]) for field in unpacked_arrays.fields}
+ del arrays, unpacked_arrays, unpacked_dict
+
+ else:
+ arrays = ak.flatten(ak.cartesian(data))
+ data = {field: arrays[field] for field in arrays.fields}
+ del arrays
# check if conversion to records is needed
arr_types = (ak.Array, np.ndarray)
@@ -93,15 +144,22 @@ def allows_shift(ax) -> bool:
del arrays
# fill
- h.fill(**fill_kwargs, **data)
+ try:
+ h.fill(**fill_kwargs, **data)
+ except MemoryError as e:
+ msg = f"error while filling histogram\n{h!r}\nwith data {data}\nand kwargs {fill_kwargs}"
+ raise MemoryError(msg) from e
-def add_hist_axis(histogram: hist.Hist, variable_inst: od.Variable) -> hist.Hist:
+def add_hist_axis(
+ h: hist.quick_construct.QuickConstruct,
+ variable_inst: od.Variable,
+) -> hist.quick_construct.QuickConstruct:
"""
Add an axis to a histogram based on a variable instance. The axis_type is chosen based on the variable instance's
"axis_type" auxiliary.
- :param histogram: The histogram to add the axis to.
+ :param h: The histogram to add the axis to.
:param variable_inst: The variable instance to use for the axis.
:return: The histogram with the added axis.
"""
@@ -121,17 +179,17 @@ def add_hist_axis(histogram: hist.Hist, variable_inst: od.Variable) -> hist.Hist
axis_type = variable_inst.x("axis_type", default_axis_type).lower()
if axis_type in {"variable", "var"}:
- return histogram.Var(variable_inst.bin_edges, **axis_kwargs)
+ return h.Variable(variable_inst.bin_edges, **axis_kwargs)
if axis_type in {"integer", "int"}:
- return histogram.Integer(
+ return h.Integer(
int(variable_inst.bin_edges[0]),
int(variable_inst.bin_edges[-1]),
**axis_kwargs,
)
if axis_type in {"boolean", "bool"}:
- return histogram.Boolean(**axis_kwargs)
+ return h.Boolean(**axis_kwargs)
if axis_type in {"intcategory", "intcat"}:
binning = (
@@ -140,16 +198,16 @@ def add_hist_axis(histogram: hist.Hist, variable_inst: od.Variable) -> hist.Hist
else []
)
axis_kwargs.setdefault("growth", True)
- return histogram.IntCat(binning, **axis_kwargs)
+ return h.IntCategory(binning, **axis_kwargs)
if axis_type in {"strcategory", "strcat"}:
axis_kwargs.setdefault("growth", True)
- return histogram.StrCat([], **axis_kwargs)
+ return h.StrCategory([], **axis_kwargs)
if axis_type in {"regular", "reg"}:
if not variable_inst.even_binning:
logger.warning("regular axis with uneven binning is not supported, using first and last bin edge instead")
- return histogram.Regular(
+ return h.Regular(
variable_inst.n_bins,
variable_inst.bin_edges[0],
variable_inst.bin_edges[-1],
@@ -214,7 +272,7 @@ def copy_axis(axis: hist.axis.AxesMixin, **kwargs: dict[str, Any]) -> hist.axis.
def create_hist_from_variables(
*variable_insts,
- categorical_axes: tuple[tuple[str, str]] | None = None,
+ categorical_axes: Sequence[tuple[str, str] | tuple[str, str, list[str]]] | None = None,
weight: bool = True,
storage: str | None = None,
) -> hist.Hist:
@@ -224,11 +282,12 @@ def create_hist_from_variables(
# additional category axes
if categorical_axes:
- for name, axis_type in categorical_axes:
+ for tpl in categorical_axes:
+ name, axis_type, axis_values = (tpl + ([],))[:3]
if axis_type in ("intcategory", "intcat"):
- histogram = histogram.IntCat([], name=name, growth=True)
+ histogram = histogram.IntCat(axis_values, name=name, growth=True)
elif axis_type in ("strcategory", "strcat"):
- histogram = histogram.StrCat([], name=name, growth=True)
+ histogram = histogram.StrCat(axis_values, name=name, growth=True)
else:
raise ValueError(f"unknown axis type '{axis_type}' in argument 'categorical_axes'")
@@ -242,35 +301,32 @@ def create_hist_from_variables(
storage = "weight" if weight else "double"
else:
storage = storage.lower()
+
+ # actual creation
if storage == "weight":
- histogram = histogram.Weight()
+ h = histogram.Weight()
elif storage == "double":
- histogram = histogram.Double()
+ h = histogram.Double()
else:
raise ValueError(f"unknown storage type '{storage}'")
- return histogram
-
-
-create_columnflow_hist = functools.partial(create_hist_from_variables, categorical_axes=(
- # axes that are used in columnflow tasks per default
- # (NOTE: "category" axis is filled as int, but transformed to str afterwards)
- ("category", "intcat"),
- ("process", "intcat"),
- ("shift", "strcat"),
-))
+ return h
def translate_hist_intcat_to_strcat(
h: hist.Hist,
axis_name: str,
- id_map: dict[int, str],
+ id_map: Callable[[int], str] | dict[int, str],
) -> hist.Hist:
import hist
+ # wrap id_map in a callable if it is given as a dict
+ if isinstance(id_map, dict):
+ id_map = functools.partial(operator.getitem, id_map)
+
out_axes = [
ax if ax.name != axis_name else hist.axis.StrCategory(
- [id_map[v] for v in list(ax)],
+ [id_map(v) for v in list(ax)],
name=ax.name,
label=ax.label,
growth=ax.traits.growth,
@@ -316,7 +372,7 @@ def update_ax_labels(hists: list[hist.Hist], config_inst: od.Config, variable_na
:param hists: List of histograms to update.
:param config_inst: Configuration instance containing variable definitions.
:param variable_name: Name of the variable to update labels for, formatted as a string
- with variable names separated by hyphens (e.g., "var1-var2").
+ with variable names separated by hyphens (e.g., "var1-var2").
:raises ValueError: If a variable name is not found in the histogram axes.
"""
labels = {}
@@ -366,3 +422,201 @@ def sum_hists(hists: Sequence[hist.Hist]) -> hist.Hist:
h_sum = h_sum + (h if h_aligned_labels is None else h_aligned_labels)
return h_sum
+
+
+def select_category_bins(
+ h: hist.Hist,
+ categories: od.Category | str | Sequence[od.Category | str],
+ use_leaves: bool = False,
+ prefer_parents: bool = False,
+ reduce: bool = False,
+) -> hist.Hist:
+ """
+ Given a histogram *h* that has a categorical "category" axis, this function selects bins on this axis according to
+ *categories*.
+
+ Depending on the types of the provided *categories*, the behavior can be more intricate. If *categories* are given
+ as strings, bins with matching names are selected with pattern support. If *categories* are
+ :py:class:`~order.Category` objects, the selection is similar (without pattern support). However, if *use_leaves* is
+ set to *True*, the selection is done based on the leaf categories of the provided *categories* instead. For
+ additional control in cases where the requested categories *and* their leaves might be present in this "category"
+ axis, *prefer_parents* can be altered to decide which of the two cases to prefer.
+
+ If *reduce* is set to *True*, the selected bins are summed before the resulting histogram is returned.
+
+ :param h: The histogram to select from. Must have a categorical "category" axis.
+ :param categories: The categories to select. Can be given as strings or :py:class:`~order.Category` objects.
+ :param use_leaves: Whether to select based on the leaf categories of the provided *categories* instead of the
+ provided *categories* themselves (if they are :py:class:`~order.Category` objects).
+ :param prefer_parents: Whether to prefer selection based on the provided *categories* themselves over their leaf
+ categories in cases where both are present in the "category" axis. Defaults to the value of *use_leaves*.
+ :param reduce: Whether to sum the selected bins before returning the resulting histogram.
+ :return: The resulting histogram.
+ """
+ import hist
+
+ # get the category axis
+ axis = h.axes["category"]
+
+ # obtain category names to select based on the provided types
+ categories = law.util.make_unique(law.util.make_list(categories))
+ if all(isinstance(c, str) for c in categories):
+ # interpret categories as sequence of patterns
+ selected_names = [name for name in axis if law.util.multi_match(name, categories, mode=any)]
+
+ elif all(isinstance(c, od.Category) for c in categories):
+ # categories are objects, so check leaf usage
+ if use_leaves:
+ selected_names = []
+ for c in categories:
+ if prefer_parents and c.name in axis:
+ selected_names.append(c.name)
+ else:
+ leaves = c.get_leaf_categories() or [c]
+ selected_names.extend(l.name for l in leaves if l.name in axis)
+ else:
+ selected_names = [c.name for c in categories if c.name in axis]
+
+ else:
+ raise ValueError(f"categories must be given as all strings or od.Category objects, got {categories}")
+
+ # finally select and optionally reduce
+ h = h[{"category": [hist.loc(name) for name in selected_names]}]
+ if reduce:
+ h = h[{"category": sum}]
+
+ return h
+
+
+def ensure_bin_exists(
+ h: hist.Hist,
+ axis_name: str,
+ value: Any,
+ add_growth: bool = True,
+) -> hist.Hist:
+ """
+ Takes a histogram *h* and ensures that a certain *value* is a valid bin in a categorical axis
+ named *axis_name*. Optionally, *growth* can be activated on the selected axis.
+
+ :param h: The histogram.
+ :param axis_name: The name of the categorical axis to ammend.
+ :param value: The value for which a bin should be ensured on the axis.
+ :param add_growth: If growth should be enabled on the selected axis.
+ :return: The updated histogram.
+ """
+ import hist
+
+ # work on a copy
+ h = h.copy()
+
+ # find the axis and its index
+ for axis_index, axis in enumerate(h.axes):
+ if axis.name == axis_name:
+ break
+ else:
+ raise ValueError(f"no axis named '{axis_name}' found in histogram: {h}")
+
+ # for now, only adjustments to categorical axes are allowed
+ categorical_types = (hist.axis.IntCategory, hist.axis.StrCategory)
+ if not isinstance(axis, categorical_types):
+ raise TypeError(f"axis '{axis_name}' must be categorical, but found {axis.__class__.__name__}")
+
+ # helper to create advanced slice object using axis_index
+ def create_slice(where: Any) -> tuple:
+ return tuple(
+ axis_index * [slice(None, None)] + # axes before target axis
+ [where] + # last (= newest) position on amended axis
+ (h.ndim - axis_index - 1) * [slice(None, None)], # remaining axes
+ )
+
+ # check that growth is enabled
+ if not axis.traits.growth:
+ if not add_growth:
+ raise ValueError(f"cannot extend values axis '{axis_name}' without growth enabled")
+ # enable growth
+ axis = copy_axis(axis, growth=True)
+ new_axes = [
+ (axis if ax.name == axis.name else copy_axis(ax))
+ for ax in h.axes
+ ]
+ # take the initial data but drop the overflow bin on the requested axis
+ # (which always exists on categorical axes that don't have growth enabled)
+ data = h.view(flow=True)[create_slice(slice(0, -1))]
+ h = hist.Hist(*new_axes, storage=h.storage_type(), data=data)
+
+ # nothing to do in case the value already exists
+ if value in axis:
+ return h
+
+ # strategy: fill dummy values, then zero them in view afterwards
+ # get bin centers of first bin in other axes
+ dummy_values = {
+ ax.name: (list(ax)[0] if isinstance(ax, categorical_types) else ax.centers[0])
+ for ax in h.axes if ax != axis
+ }
+ h.fill(**(dummy_values | {axis_name: value}))
+ # now zero array
+ view = h.view()
+ zero_pos = create_slice(-1)
+ if isinstance(h.storage_type(), (hist.storage.Weight, hist.storage.WeightedMean)):
+ view.value[zero_pos] = 0
+ view.variance[zero_pos] = 0
+ else:
+ view[zero_pos] = 0
+
+ return h
+
+
+def merge_axis_bins(
+ h: hist.Hist,
+ axis_name: str,
+ merged_bin_name: str,
+ bins_to_merge: Sequence[str],
+ remove_bins: bool = True,
+) -> hist.Hist:
+ """
+ This function merges categorical bins given in *bins_to_merge* onto the bin *merged_bin_name*.
+
+ - If *merged_bin_name* does not exist yet in *axis_name*, it is initialized with empty content.
+ - If *merged_bin_name* exists, it must be included in *bins_to_merge* to avoid merging into an existing bin.
+
+ :param h: The histogram.
+ :param axis_name: The name of the categorical axis to ammend.
+ :param merged_bin_name: The name of the merged bin which is optionally created.
+ :param bins_to_merge: The bins to merge and subsequently remove.
+ :param remove_bins: Optionally remove bins in *bins_to_merge* after merging.
+ :return: The updated histogram.
+ """
+ import hist
+
+ # sanity checks
+ for bin_name in bins_to_merge:
+ if bin_name not in h.axes[axis_name]:
+ raise ValueError(f"bin '{bin_name}' not found in axis '{axis_name}'")
+
+ if merged_bin_name in h.axes[axis_name] and merged_bin_name not in bins_to_merge:
+ raise ValueError(
+ f"bin '{merged_bin_name}' already exists but is not found in bins to be merged '{bins_to_merge}', "
+ "leading to unclear bin content handling",
+ )
+
+ # ensure merged_bin_name exists in axis, it may or may not be already filled
+ h = ensure_bin_exists(h, axis_name, merged_bin_name)
+
+ # add up bin contents
+ for bin_name in bins_to_merge:
+ # avoid double counting in case merged_bin_name is also in bins_to_merge
+ if bin_name != merged_bin_name:
+ h[{axis_name: merged_bin_name}] += h[{axis_name: bin_name}]
+
+ # remove unneeded bins
+ if remove_bins:
+ bins_to_remove = set(bins_to_merge) - {merged_bin_name}
+ h = h[{
+ axis_name: [
+ hist.loc(bin_name) for bin_name in h.axes[axis_name]
+ if bin_name not in bins_to_remove
+ ],
+ }]
+
+ return h
diff --git a/columnflow/histogramming/__init__.py b/columnflow/histogramming/__init__.py
index f8b76ff20..2cc1f28e6 100644
--- a/columnflow/histogramming/__init__.py
+++ b/columnflow/histogramming/__init__.py
@@ -12,8 +12,8 @@
import order as od
from columnflow.production import TaskArrayFunctionWithProducerRequirements
-from columnflow.util import DerivableMeta, maybe_import
-from columnflow.types import TYPE_CHECKING, Any, Callable, Sequence
+from columnflow.util import DerivableMeta, maybe_import, UNSET
+from columnflow.types import TYPE_CHECKING, Any, Callable, Sequence, UNSET_TYPE
if TYPE_CHECKING:
hist = maybe_import("hist")
@@ -66,9 +66,9 @@ def hist_producer(
cls,
func: Callable | None = None,
bases: tuple = (),
- mc_only: bool = False,
- data_only: bool = False,
- require_producers: Sequence[str] | set[str] | None = None,
+ mc_only: bool | UNSET_TYPE = UNSET,
+ data_only: bool | UNSET_TYPE = UNSET,
+ require_producers: Sequence[str] | set[str] | None | UNSET_TYPE = UNSET,
**kwargs,
) -> DerivableMeta | Callable:
"""
@@ -92,13 +92,13 @@ def hist_producer(
"""
def decorator(func: Callable) -> DerivableMeta:
# create the class dict
- cls_dict = {
- **kwargs,
- "call_func": func,
- "mc_only": mc_only,
- "data_only": data_only,
- "require_producers": require_producers,
- }
+ cls_dict = {**kwargs, "call_func": func}
+ if mc_only is not UNSET:
+ cls_dict["mc_only"] = mc_only
+ if data_only is not UNSET:
+ cls_dict["data_only"] = data_only
+ if require_producers is not UNSET:
+ cls_dict["require_producers"] = require_producers
# get the module name
frame = inspect.stack()[1]
@@ -119,8 +119,8 @@ def update_cls_dict(cls_name, cls_dict, get_attr):
if mc_only or data_only:
if cls_dict.get("skip_func"):
raise Exception(
- f"hist producer {cls_name} received custom skip_func, but either mc_only or data_only "
- "are set",
+ f"hist producer {cls_name} received custom skip_func, but either mc_only or data_only are "
+ "set",
)
if "skip_func" not in cls_dict:
diff --git a/columnflow/histogramming/default.py b/columnflow/histogramming/default.py
index c702b0d87..8bd51de31 100644
--- a/columnflow/histogramming/default.py
+++ b/columnflow/histogramming/default.py
@@ -6,6 +6,8 @@
from __future__ import annotations
+import functools
+
import law
import order as od
@@ -46,11 +48,11 @@ def cf_default_create_hist(
"""
return create_hist_from_variables(
*variables,
- categorical_axes=(
+ categorical_axes=[
("category", "intcat"),
("process", "intcat"),
- ("shift", "intcat"),
- ),
+ ("shift", "intcat", [0]),
+ ],
weight=True,
)
@@ -60,6 +62,28 @@ def cf_default_fill_hist(self: HistProducer, h: hist.Hist, data: dict[str, Any],
"""
Fill the histogram with the data.
"""
+ # in case multiple variable axes are given that refer to data arrays with more than one dimension (i.e. nested),
+ # check if they are broadcasting-compatible since otherwise, the full combinatorics of values would be fille which
+ # is not supported by fill_hist in its default implementation
+ import hist
+
+ var_axes = [
+ ax for ax in h.axes
+ if isinstance(ax, (hist.axis.Variable, hist.axis.Integer)) and ax.name in data and data[ax.name].ndim > 1
+ ]
+ if len(var_axes) > 1:
+ ref_counts = ak.count(data[var_axes[0].name], axis=1)
+ for ax in var_axes[1:]:
+ counts = ak.count(data[ax.name], axis=1)
+ if not ak.all(counts == ref_counts):
+ err = (
+ "detected multiple variable axes with data to be filled that is not broadcasting-compatible:\n" +
+ "\n - ".join(f"{ax.name}: {data[ax.name]}" for ax in var_axes) +
+ "please use a custom histogram producer whose fill_hist implementation supports the desired "
+ "filling logic including combinatorics or custom broadcasting"
+ )
+ raise ValueError(err)
+
fill_hist(h, data, last_edge_inclusive=task.last_edge_inclusive)
@@ -73,14 +97,20 @@ def cf_default_post_process_hist(self: HistProducer, h: hist.Hist, task: law.Tas
# translate axes
if "category" in axis_names:
- category_map = {cat.id: cat.name for cat in self.config_inst.get_leaf_categories()}
- h = translate_hist_intcat_to_strcat(h, "category", category_map)
+ @functools.cache
+ def get_category_name(cat_id: int) -> str:
+ return self.config_inst.get_category(cat_id).name
+ h = translate_hist_intcat_to_strcat(h, "category", get_category_name)
if "process" in axis_names:
- process_map = {proc_id: self.config_inst.get_process(proc_id).name for proc_id in h.axes["process"]}
- h = translate_hist_intcat_to_strcat(h, "process", process_map)
+ @functools.cache
+ def get_process_name(process_id: int) -> str:
+ return self.config_inst.get_process(process_id).name
+ h = translate_hist_intcat_to_strcat(h, "process", get_process_name)
if "shift" in axis_names:
- shift_map = {task.global_shift_inst.id: task.global_shift_inst.name}
- h = translate_hist_intcat_to_strcat(h, "shift", shift_map)
+ @functools.cache
+ def get_shift_name(shift_id: int) -> str:
+ return self.config_inst.get_shift(shift_id).name
+ h = translate_hist_intcat_to_strcat(h, "shift", get_shift_name)
return h
@@ -132,7 +162,9 @@ def all_weights(self: HistProducer, events: ak.Array, **kwargs) -> ak.Array:
@all_weights.init
-def all_weights_init(self: HistProducer) -> None:
+def all_weights_init(self: HistProducer, **kwargs) -> None:
+ super(all_weights, self).init_func(**kwargs)
+
weight_columns = set()
if self.dataset_inst.is_data:
diff --git a/columnflow/inference/__init__.py b/columnflow/inference/__init__.py
index 0023c1352..1457dcf2f 100644
--- a/columnflow/inference/__init__.py
+++ b/columnflow/inference/__init__.py
@@ -257,6 +257,7 @@ class InferenceModel(Derivable, metaclass=InferenceModelMeta):
mc_datasets: [hh_ggf]
scale: 1.0
is_dynamic: False
+ skip_if_empty: True
parameters:
- name: lumi
type: rate_gauss
@@ -286,6 +287,7 @@ class InferenceModel(Derivable, metaclass=InferenceModelMeta):
mc_datasets: [tt_sl, tt_dl, tt_fh]
scale: 1.0
is_dynamic: False
+ skip_if_empty: True
parameters:
- name: lumi
type: rate_gauss
@@ -374,10 +376,7 @@ def inference_model(
"""
def decorator(func: Callable) -> Type[T]:
# create the class dict
- cls_dict = {
- **kwargs,
- "init_func": func,
- }
+ cls_dict = {**kwargs, "init_func": func}
# create the subclass
subclass = cls.derive(func.__name__, bases=bases, cls_dict=cls_dict)
@@ -465,6 +464,7 @@ def process_spec(
config_data: dict[str, DotDict] | None = None,
scale: float | int = 1.0,
is_dynamic: bool = False,
+ skip_if_empty: bool = True,
) -> DotDict:
"""
Returns a dictionary representing a process, forwarding all arguments.
@@ -476,6 +476,7 @@ def process_spec(
:param scale: A float value to scale the process, defaulting to 1.0.
:param is_dynamic: A boolean flag deciding whether this process is dynamic, i.e., whether it is created
on-the-fly.
+ :param skip_if_empty: A boolean flag deciding whether this process should be skipped if input hists are empty.
:returns: A dictionary representing the process.
"""
return DotDict([
@@ -488,6 +489,7 @@ def process_spec(
)),
("scale", float(scale)),
("is_dynamic", bool(is_dynamic)),
+ ("skip_if_empty", bool(skip_if_empty)),
("parameters", []),
])
@@ -554,7 +556,7 @@ def category_config_spec(
data_datasets: Sequence[str] | None = None,
) -> DotDict:
"""
- Returns a dictionary representing configuration specific data, forwarding all arguments.
+ Returns a dictionary representing configuration specific data for a category, forwarding all arguments.
:param category: The name of the source category in the config to use.
:param variable: The name of the variable in the config to use.
@@ -574,7 +576,7 @@ def process_config_spec(
mc_datasets: Sequence[str] | None = None,
) -> DotDict:
"""
- Returns a dictionary representing configuration specific data, forwarding all arguments.
+ Returns a dictionary representing configuration specific data for a process, forwarding all arguments.
:param process: The name of the process in the config to use.
:param mc_datasets: List of names or patterns of datasets in the config to use for mc.
@@ -591,7 +593,7 @@ def parameter_config_spec(
shift_source: str | None = None,
) -> DotDict:
"""
- Returns a dictionary representing configuration specific data, forwarding all arguments.
+ Returns a dictionary representing configuration specific data for a parameter, forwarding all arguments.
:param shift_source: The name of a systematic shift source in the config.
:returns: A dictionary representing parameter specific settings.
@@ -703,7 +705,7 @@ def get_category(
only_name: bool = False,
match_mode: Callable = any,
silent: bool = False,
- ) -> DotDict | str:
+ ) -> DotDict | str | None:
"""
Returns a single category whose name matches *category*. *category* can be a string, a
pattern, or sequence of them. An exception is raised if no or more than one category is
@@ -874,7 +876,7 @@ def get_process(
match_mode: Callable = any,
category_match_mode: Callable = any,
silent: bool = False,
- ) -> DotDict | str:
+ ) -> DotDict | str | None:
"""
Returns a single process whose name matches *process*, and optionally, whose category's
name matches *category*. Both *process* and *category* can be a string, a pattern, or
@@ -1152,7 +1154,7 @@ def get_parameter(
process_match_mode: Callable = any,
only_name: bool = False,
silent: bool = False,
- ) -> DotDict | str:
+ ) -> DotDict | str | None:
"""
Returns a single parameter whose name matches *parameter*, and optionally, whose category's
and process' name matches *category* and *process*. All three, *parameter*, *process* and
diff --git a/columnflow/inference/cms/datacard.py b/columnflow/inference/cms/datacard.py
index 373f0875a..3e0d8fc3d 100644
--- a/columnflow/inference/cms/datacard.py
+++ b/columnflow/inference/cms/datacard.py
@@ -17,6 +17,8 @@
from columnflow.util import DotDict, maybe_import, real_path, ensure_dir, safe_div, maybe_int
from columnflow.types import TYPE_CHECKING, Sequence, Any, Union, Hashable
+np = maybe_import("numpy")
+
if TYPE_CHECKING:
hist = maybe_import("hist")
@@ -151,6 +153,53 @@ def validate_model(cls, inference_model_inst: InferenceModel, silent: bool = Fal
return True
+ @classmethod
+ def validate_histograms(cls, histograms: DatacardHists, silent: bool = False) -> bool:
+ import hist
+
+ # validate structure of histograms, shape keys and histogram types
+ errors: list[str] = []
+ for cat_name, proc_hists in histograms.items():
+ if not isinstance(cat_name, str):
+ errors.append(f"category name key '{cat_name}' is not a string")
+ for proc_name, config_hists in proc_hists.items():
+ if not isinstance(proc_name, str):
+ errors.append(f"process name '{proc_name}' in category '{cat_name}' is not a string")
+ for config_name, shift_hists in config_hists.items():
+ if not isinstance(config_name, str):
+ errors.append(
+ f"config name '{config_name}' for process '{proc_name}' in category '{cat_name}' is not a "
+ f"string",
+ )
+ for shift_key, h in shift_hists.items():
+ # shift_key must be nominal or a tuple of (param_name, "up|down")
+ if (
+ shift_key != "nominal" and
+ (
+ not isinstance(shift_key, (tuple, list)) or
+ len(shift_key) != 2 or
+ shift_key[1] not in {"up", "down"}
+ )
+ ):
+ errors.append(
+ f"invalid shift key '{shift_key}' in config '{config_name}' for process "
+ f"'{proc_name}' in category '{cat_name}'",
+ )
+ if not isinstance(h, hist.Hist):
+ errors.append(
+ f"histogram for shift '{shift_key}' in config '{config_name}' for process "
+ f"'{proc_name}' in category '{cat_name}' is not a hist.Hist instance",
+ )
+
+ # handle errors
+ if errors:
+ if silent:
+ return False
+ errors_repr = "\n - ".join(errors)
+ raise ValueError(f"datacard histograms invalid, reasons:\n - {errors_repr}")
+
+ return True
+
def __init__(
self,
inference_model_inst: InferenceModel,
@@ -168,13 +217,13 @@ def __init__(
self.histograms = histograms
self.rate_precision = rate_precision
self.effect_precision = effect_precision
- self.effect_precision = effect_precision
self.effect_from_shape_if_flat_max_outlier = effect_from_shape_if_flat_max_outlier
self.effect_from_shape_if_flat_max_deviation = effect_from_shape_if_flat_max_deviation
self.asymmetrize_if_large_threshold = asymmetrize_if_large_threshold
- # validate the inference model
+ # validate the inference model and histograms
self.validate_model(self.inference_model_inst)
+ self.validate_histograms(self.histograms)
def write(
self,
@@ -200,7 +249,7 @@ def write(
cat_objects = [self.inference_model_inst.get_category(cat_name) for cat_name in rates]
# prepare blocks and lines to write
- blocks = DotDict()
+ blocks: DotDict[str, list] = DotDict()
separators = set()
empty_lines = set()
@@ -270,8 +319,10 @@ def write(
types = set()
effects = []
for cat_name, proc_name in flat_rates:
+ cat_obj = self.inference_model_inst.get_category(category=cat_name)
+ proc_obj = self.inference_model_inst.get_process(category=cat_name, process=proc_name)
param_obj = self.inference_model_inst.get_parameter(
- param_name,
+ parameter=param_name,
category=cat_name,
process=proc_name,
silent=True,
@@ -304,7 +355,20 @@ def write(
if param_obj.effect_precision <= 0
else param_obj.effect_precision
)
- rnd = lambda f: round(f, effect_precision)
+
+ def rnd(f: float | int) -> float:
+ r = round(f, effect_precision)
+
+ # warn in case the precision is too low for the effect
+ if abs(1.0 - f) < 10**(-effect_precision):
+ logger.warning(
+ f"the effect value '{f}' is rounded to '{r}' which probably leads to loosing its intended "
+ f"impact; consider choosing an effect precision higher than the current value of "
+ f"{effect_precision} for the paremeter '{param_name}' acting on process '{proc_name}' in "
+ f"category '{cat_name}'",
+ )
+
+ return r
# update and transform effects
if param_obj.type.is_rate:
@@ -373,7 +437,7 @@ def write(
effect = 1.0
# custom hook to modify the effect
- effect = self.modify_parameter_effect(cat_name, proc_name, param_obj, effect)
+ effect = self.modify_parameter_effect(cat_obj, proc_obj, param_obj, effect)
# encode the effect
if isinstance(effect, (int, float)):
@@ -404,7 +468,7 @@ def write(
type_str = "shape"
elif types == {ParameterType.rate_gauss, ParameterType.shape}:
# when mixing lnN and shape effects, combine expects the "shape?" type and makes the actual decision
- # dependend on the presence of shape variations in the accompaying shape files, see
+ # dependent on the presence of shape variations in the accompanying shape files, see
# https://cms-analysis.github.io/HiggsAnalysis-CombinedLimit/v10.2.X/part2/settinguptheanalysis/?h=shape%3F#template-shape-uncertainties # noqa
type_str = "shape?"
if not type_str:
@@ -480,6 +544,9 @@ def write(
if blocks.mc_stats:
blocks.mc_stats = self.align_lines(list(blocks.mc_stats))
+ # allow modification before writing via hook
+ blocks, separators, empty_lines = self.modify_before_write(blocks, separators, empty_lines)
+
# write the blocks
with open(datacard_path, "w") as f:
for block_name, lines in blocks.items():
@@ -554,7 +621,7 @@ def handle_flow(cat_obj, h, name):
return
# warn in case of flow content
- if cat_obj.flow_strategy == FlowStrategy.warn:
+ if cat_obj.flow_strategy in {FlowStrategy.warn, FlowStrategy.move}:
if underflow[0]:
logger.warning_once(
f"underflow_warn_{self.inference_model_inst.cls_name}_{cat_obj.name}_{name}",
@@ -567,6 +634,9 @@ def handle_flow(cat_obj, h, name):
f"overflow content detected in category '{cat_obj.name}' for histogram "
f"'{name}' ({overflow[0] / view.value.sum() * 100:.1f}% of integral)",
)
+
+ # stop here in case of warn-only
+ if cat_obj.flow_strategy == FlowStrategy.warn:
return
# here, we can already remove overflow values
@@ -603,7 +673,8 @@ def fill_empty(cat_obj, h):
_effects = effects[cat_name] = OrderedDict()
for proc_name, config_hists in proc_hists.items():
# skip if process is not known to category
- if not self.inference_model_inst.has_process(process=proc_name, category=cat_name):
+ proc_obj = self.inference_model_inst.get_process(process=proc_name, category=cat_name, silent=True)
+ if not proc_obj:
continue
# defer the handling of data to the end
@@ -618,7 +689,7 @@ def fill_empty(cat_obj, h):
if not hists:
continue
- # helper to sum over them for a given shift key and an optional fallback
+ # helper to sum over histograms for a given shift key and an optional fallback
def get_hist_sum(key: Hashable, fallback_key: Hashable | None = None) -> hist.Hist:
def get(hd: dict[Hashable, hist.Hist]) -> hist.Hist:
if key in hd:
@@ -630,6 +701,14 @@ def get(hd: dict[Hashable, hist.Hist]) -> hist.Hist:
)
return sum_hists(map(get, hists))
+ # optionally skip the process under specific conditions
+ if (skip_reason := self.check_skip_process(cat_obj, proc_obj, get_hist_sum("nominal"))):
+ skip_msg = f"skipping process '{proc_name}' in category '{cat_name}'"
+ if not isinstance(skip_reason, bool):
+ skip_msg += f", reason: {skip_reason}"
+ skip_msg = logger.info(skip_msg)
+ continue
+
# helper to extract sum of hists, apply scale, handle flow and fill empty bins
def load(
hist_name: str,
@@ -643,7 +722,6 @@ def load(
return h
# get the process scale (usually 1)
- proc_obj = self.inference_model_inst.get_process(proc_name, category=cat_name)
scale = proc_obj.scale
# nominal shape
@@ -688,18 +766,15 @@ def load(
h_down = h_nom.copy() * f_down
h_up = h_nom.copy() * f_up
else:
- # just extract the shapes
- h_down = sum_hists((param_obj.name, "down"), "nominal") * scale
- h_up = sum_hists((param_obj.name, "up"), "nominal") * scale
# just extract the shapes from the inputs
- h_down = load(down_name, (param_obj.name, "down"), "nominal", scale=scale)
- h_up = load(up_name, (param_obj.name, "up"), "nominal", scale=scale)
+ h_down = load(down_name, (param_obj.name, "down"), fallback_key="nominal", scale=scale)
+ h_up = load(up_name, (param_obj.name, "up"), fallback_key="nominal", scale=scale)
elif param_obj.type.is_rate:
if param_obj.transformations.any_from_shape:
# just extract the shapes
- h_down = load(down_name, (param_obj.name, "down"), "nominal", scale=scale)
- h_up = load(up_name, (param_obj.name, "up"), "nominal", scale=scale)
+ h_down = load(down_name, (param_obj.name, "down"), fallback_key="nominal", scale=scale)
+ h_up = load(up_name, (param_obj.name, "up"), fallback_key="nominal", scale=scale)
# in case the transformation is effect_from_shape_if_flat, and any of the two variations
# do not qualify as "flat", convert the parameter to shape-type and drop all transformations
@@ -794,8 +869,8 @@ def load(
# custom hook to modify the shapes
h_nom, h_down, h_up = self.modify_parameter_shape(
- cat_name,
- proc_name,
+ cat_obj,
+ proc_obj,
param_obj,
h_nom,
h_down,
@@ -825,7 +900,7 @@ def load(
if proc_name in proc_hists:
h_data.extend([hd["nominal"] for hd in proc_hists[proc_name].values()])
else:
- logger.warning(f"process '{proc_name}' not found in histograms for created fake data, skipping")
+ logger.warning(f"process '{proc_name}' not found in histograms for creating fake data, skipping")
if not h_data:
proc_str = ",".join(map(str, cat_obj.data_from_processes))
raise Exception(f"none of requested processes '{proc_str}' found to create fake data")
@@ -922,10 +997,26 @@ def align_rates_and_parameters(
return lines[:n_rate_lines], lines[n_rate_lines:]
+ def modify_before_write(
+ self,
+ blocks: DotDict[str, list],
+ separators: set[str],
+ empty_lines: set[str],
+ ) -> tuple[DotDict[str, list], set[str], set[str]]:
+ """
+ Hook to modify the datacard blocks, empty lines and separators before they are written to the datacard file.
+
+ :param blocks: Datacard blocks.
+ :param separators: Set of block names after which a separator line should be inserted.
+ :param empty_lines: Set of block names after which an empty line should be inserted.
+ :returns: The modified datacard blocks, separators and empty lines.
+ """
+ return blocks, separators, empty_lines
+
def modify_parameter_effect(
self,
- category: str,
- process: str,
+ cat_obj: DotDict,
+ proc_obj: DotDict,
param_obj: DotDict,
effect: float | tuple[float, float],
) -> float | tuple[float, float]:
@@ -933,8 +1024,8 @@ def modify_parameter_effect(
Custom hook to modify the effect of a parameter on a given category and process before it is encoded into the
datacard. By default, this does nothing and simply returns the given effect.
- :param category: The category name.
- :param process: The process name.
+ :param cat_obj: The category object, following :py:meth:`columnflow.inference.InferenceModel.category_spec`.
+ :param proc_obj: The process object, following :py:meth:`columnflow.inference.InferenceModel.process_spec`.
:param param_obj: The parameter object, following :py:meth:`columnflow.inference.InferenceModel.parameter_spec`.
:param effect: The effect value(s) to be modified.
:returns: The modified effect value(s).
@@ -943,8 +1034,8 @@ def modify_parameter_effect(
def modify_parameter_shape(
self,
- category: str,
- process: str,
+ cat_obj: DotDict,
+ proc_obj: DotDict,
param_obj: DotDict,
h_nom: hist.Hist,
h_down: hist.Hist,
@@ -954,8 +1045,8 @@ def modify_parameter_shape(
Custom hook to modify the nominal and varied (down, up) shapes of a parameter on a given category and process
before they are saved to the shapes file. By default, this does nothing and simply returns the given histograms.
- :param category: The category name.
- :param process: The process name.
+ :param cat_obj: The category object, following :py:meth:`columnflow.inference.InferenceModel.category_spec`.
+ :param proc_obj: The process object, following :py:meth:`columnflow.inference.InferenceModel.process_spec`.
:param param_obj: The parameter object, following :py:meth:`columnflow.inference.InferenceModel.parameter_spec`.
:param h_nom: The nominal histogram.
:param h_down: The down-varied histogram.
@@ -963,3 +1054,23 @@ def modify_parameter_shape(
:returns: The modified nominal and varied (down, up) histograms.
"""
return h_nom, h_down, h_up
+
+ def check_skip_process(
+ self,
+ cat_obj: DotDict,
+ proc_obj: DotDict,
+ h: hist.Hist,
+ ) -> bool | str:
+ """
+ Custom hook to check if a process in a given category should be skipped entirely based on the nominal histogram
+ and the process and category objects. If a string is returned, it is added to the log message as a reason for
+ skipping the process.
+
+ :param cat_obj: The category object, following :py:meth:`columnflow.inference.InferenceModel.category_spec`.
+ :param proc_obj: The process object, following :py:meth:`columnflow.inference.InferenceModel.process_spec`.
+ :param h: The nominal histogram for the process in the category.
+ :returns: Whether to skip the process, and optionally a reason for skipping.
+ """
+ if proc_obj.skip_if_empty and np.all(h.view().value == 0):
+ return "nominal histogram is empty"
+ return False
diff --git a/columnflow/plotting/plot_all.py b/columnflow/plotting/plot_all.py
index 7f8d67ad6..4c8a83665 100644
--- a/columnflow/plotting/plot_all.py
+++ b/columnflow/plotting/plot_all.py
@@ -34,8 +34,6 @@ def draw_stat_error_bands(
norm: float | Sequence | np.ndarray = 1.0,
**kwargs,
) -> None:
- import hist
-
assert len(h.axes) == 1
# compute relative statistical errors
@@ -49,7 +47,7 @@ def draw_stat_error_bands(
baseline[np.isnan(baseline)] = 0.0
bar_kwargs = {
- "x": h.axes[0].centers - (h.axes[0].edges[0] if type(h.axes[0]) is hist.axes.Integer else 0),
+ "x": h.axes[0].centers,
"bottom": baseline * (1 - rel_stat_error),
"height": baseline * 2 * rel_stat_error,
"width": h.axes[0].edges[1:] - h.axes[0].edges[:-1],
@@ -279,6 +277,7 @@ def draw_errorbars(
h: hist.Hist,
norm: float | Sequence | np.ndarray = 1.0,
error_type: str = "poisson_unweighted",
+ density: bool = False,
**kwargs,
) -> None:
import hist
@@ -303,7 +302,7 @@ def draw_errorbars(
"Error bars calculation only implemented for histograms with storage type WeightedSum "
"either change the Histogram storage_type or set yerr manually",
)
- yerr = calculate_stat_error(h, error_type)
+ yerr = calculate_stat_error(h, error_type, density=density)
# normalize yerr to the histogram = error propagation on standard deviation
yerr = abs(yerr / norm)
# replace inf with nan for any bin where norm = 0 and calculate_stat_error returns a non zero value
@@ -392,29 +391,29 @@ def plot_all(
# check if required fields are present
if "method" not in cfg:
raise ValueError(f"no method given in plot_cfg entry {key}")
- if "hist" not in cfg:
- raise ValueError(f"no histogram(s) given in plot_cfg entry {key}")
# invoke the method
method = cfg["method"]
- h = cfg["hist"]
-
- def get_method(method):
- return method if callable(method) else plot_methods[method]
-
- get_method(method)(ax, h, **cfg.get("kwargs", {}))
+ method_func = method if callable(method) else plot_methods[method]
+ args = (ax, cfg["hist"]) if "hist" in cfg else (ax,)
+ method_func(*args, **cfg.get("kwargs", {}))
# repeat for ratio axes if configured
if not skip_ratio and "ratio_kwargs" in cfg:
# take ratio_method if the ratio plot requires a different plotting method
method = cfg.get("ratio_method", method)
- get_method(method)(rax, h, **cfg.get("ratio_kwargs", {}))
+ method_func = method if callable(method) else plot_methods[method]
+ args = (rax, cfg["hist"]) if "hist" in cfg else (rax,)
+ method_func(*args, **cfg.get("ratio_kwargs", {}))
# axis styling
ax_kwargs = {
"ylabel": "Counts",
"xlabel": "variable",
"yscale": "linear",
+ "xticklabelformat": {"style": "sci", "useMathText": True},
+ "yticklabelformat": {"style": "sci", "useMathText": True},
+ "yoffsettext": {"horizontalalignment": "right", "fontsize": 18},
}
# some default ylim settings based on yscale
@@ -439,6 +438,7 @@ def get_method(method):
"ylabel": "Ratio",
"xlabel": "Variable",
"yscale": "linear",
+ "xticklabelformat": {"style": "sci", "useMathText": True},
}
rax_kwargs.update(style_config.get("rax_cfg", {}))
diff --git a/columnflow/plotting/plot_functions_1d.py b/columnflow/plotting/plot_functions_1d.py
index 34b6d02a7..192a893f4 100644
--- a/columnflow/plotting/plot_functions_1d.py
+++ b/columnflow/plotting/plot_functions_1d.py
@@ -91,6 +91,7 @@ def plot_variable_stack(
hists,
shape_norm=shape_norm,
shift_insts=shift_insts,
+ density=density,
**kwargs,
)
@@ -326,7 +327,7 @@ def plot_shifted_variable(
if legend_title:
default_style_config["legend_cfg"]["title"] = legend_title
if shape_norm:
- style_config["ax_cfg"]["ylabel"] = "Normalized entries"
+ default_style_config["ax_cfg"]["ylabel"] = "Normalized entries"
style_config = law.util.merge_dicts(
default_style_config,
process_style_config,
diff --git a/columnflow/plotting/plot_util.py b/columnflow/plotting/plot_util.py
index c680cc46a..3f2044bdf 100644
--- a/columnflow/plotting/plot_util.py
+++ b/columnflow/plotting/plot_util.py
@@ -476,10 +476,23 @@ def prepare_style_config(
# unit format on axes (could be configurable)
unit_format = "{title} [{unit}]"
+ if density:
+ ylabel = variable_inst.get_full_y_title(
+ bin_width=False,
+ unit=variable_inst.unit or "unit",
+ unit_format="{title} / {unit}",
+ )
+ else:
+ ylabel = variable_inst.get_full_y_title(
+ bin_width=False,
+ unit=False,
+ unit_format=unit_format,
+ )
+
style_config = {
"ax_cfg": {
"xlim": xlim,
- "ylabel": variable_inst.get_full_y_title(bin_width=False, unit=False, unit_format=unit_format),
+ "ylabel": ylabel,
"xlabel": variable_inst.get_full_x_title(unit_format=unit_format),
"yscale": yscale,
"xscale": "log" if variable_inst.log_x else "linear",
@@ -530,6 +543,7 @@ def prepare_stack_plot_config(
shape_norm: bool | None = False,
hide_stat_errors: bool | None = None,
shift_insts: Sequence[od.Shift] | None = None,
+ density: bool = False,
**kwargs,
) -> OrderedDict:
"""
@@ -579,8 +593,6 @@ def prepare_stack_plot_config(
# setup plotting configs
plot_config = OrderedDict()
- # take first (non-underflow) bin
- # shape_norm_func = kwargs.get("shape_norm_func", lambda h, shape_norm: h.values()[0] if shape_norm else 1)
shape_norm_func = kwargs.get("shape_norm_func", lambda h, shape_norm: sum(h.values()) if shape_norm else 1)
# draw stack
@@ -661,6 +673,7 @@ def prepare_stack_plot_config(
"norm": data_norm,
"label": data_label or "Data",
"error_type": "poisson_unweighted",
+ "density": density,
},
}
@@ -668,6 +681,7 @@ def prepare_stack_plot_config(
plot_config["data"]["ratio_kwargs"] = {
"norm": h_mc.values() * data_norm / mc_norm,
"error_type": "poisson_unweighted",
+ "density": density,
}
# suppress error bars by overriding `yerr`
@@ -687,7 +701,8 @@ def split_ax_kwargs(kwargs: dict[str, Any]) -> tuple[dict[str, Any], dict[str, A
set_kwargs, other_kwargs = {}, {}
other_keys = {
"xmajorticks", "xminorticks", "xmajorticklabels", "xminorticklabels", "xloc", "xrotation",
- "ymajorticks", "yminorticks", "yloc", "yrotation",
+ "ymajorticks", "yminorticks", "yloc", "yrotation", "xticklabelformat", "yticklabelformat",
+ "xoffsettext", "yoffsettext",
}
for key, value in kwargs.items():
(other_kwargs if key in other_keys else set_kwargs)[key] = value
@@ -726,6 +741,24 @@ def apply_ax_kwargs(ax: plt.Axes, kwargs: dict[str, Any]) -> None:
ax.tick_params(axis="x", labelrotation=other_kwargs.get("xrotation"))
if other_kwargs.get("yrotation") is not None:
ax.tick_params(axis="y", labelrotation=other_kwargs.get("yrotation"))
+ if other_kwargs.get("xticklabelformat") is not None:
+ _kwargs = other_kwargs["xticklabelformat"].copy()
+ if not callable(getattr(ax.xaxis.major.formatter, "set_scientific", None)):
+ if _kwargs.get("style") == "sci":
+ _kwargs.pop("style")
+ _kwargs.pop("useMathText", None)
+ ax.ticklabel_format(axis="x", **_kwargs)
+ if other_kwargs.get("yticklabelformat") is not None:
+ _kwargs = other_kwargs["yticklabelformat"].copy()
+ if not callable(getattr(ax.yaxis.major.formatter, "set_scientific", None)):
+ if _kwargs.get("style") == "sci":
+ _kwargs.pop("style")
+ _kwargs.pop("useMathText", None)
+ ax.ticklabel_format(axis="y", **_kwargs)
+ if other_kwargs.get("xoffsettext") is not None:
+ ax.xaxis.get_offset_text().set(**other_kwargs["xoffsettext"])
+ if other_kwargs.get("yoffsettext") is not None:
+ ax.yaxis.get_offset_text().set(**other_kwargs["yoffsettext"])
def get_position(minimum: float, maximum: float, factor: float = 1.4, logscale: bool = False) -> float:
@@ -1257,7 +1290,7 @@ def remove_label_placeholders(
return re.sub(f"__{sel}__", "", label)
-def calculate_stat_error(h: hist.Hist, error_type: str) -> np.ndarray:
+def calculate_stat_error(h: hist.Hist, error_type: str, density: bool = True) -> np.ndarray:
"""
Calculate the error to be plotted for the given histogram *h*.
Supported error types are:
@@ -1266,6 +1299,11 @@ def calculate_stat_error(h: hist.Hist, error_type: str) -> np.ndarray:
- "poisson_unweighted": the plotted error is the poisson error for each bin
- "poisson_weighted": the plotted error is the poisson error for each bin, weighted by the variance
"""
+ # undo density if needed
+ if density:
+ area = functools.reduce(operator.mul, h.axes.widths)
+ h = h * area
+
# determine the error type
if error_type == "variance":
yerr = h.view().variance ** 0.5
@@ -1299,4 +1337,10 @@ def calculate_stat_error(h: hist.Hist, error_type: str) -> np.ndarray:
else:
raise ValueError(f"unknown error type '{error_type}'")
+ # re-apply density if needed
+ if density:
+ area = functools.reduce(operator.mul, h.axes.widths)
+ h = h / area
+ yerr = yerr / area
+
return yerr
diff --git a/columnflow/production/__init__.py b/columnflow/production/__init__.py
index 03ff6faf9..147188a5e 100644
--- a/columnflow/production/__init__.py
+++ b/columnflow/production/__init__.py
@@ -6,32 +6,46 @@
from __future__ import annotations
+import copy
import inspect
import law
-from columnflow.util import DerivableMeta
from columnflow.columnar_util import TaskArrayFunction
-from columnflow.types import Callable, Sequence, Any
+from columnflow.util import DerivableMeta, UNSET
+from columnflow.types import Callable, Sequence, Any, UNSET_TYPE
class TaskArrayFunctionWithProducerRequirements(TaskArrayFunction):
require_producers: Sequence[str] | set[str] | None = None
+ def __init__(self, *args, **kwargs):
+ if "require_producers" in kwargs or self.__class__.require_producers is None:
+ kwargs["require_producers"] = kwargs.get("require_producers") or []
+ elif isinstance(self.__class__.require_producers, (list, tuple)):
+ kwargs["require_producers"] = copy.copy(self.__class__.require_producers)
+
+ super().__init__(*args, **kwargs)
+
def _req_producer(self, task: law.Task, producer: str) -> Any:
# hook to customize how required producers are requested
from columnflow.tasks.production import ProduceColumns
return ProduceColumns.req_other_producer(task, producer=producer)
def requires_func(self, task: law.Task, reqs: dict, **kwargs) -> None:
+ super().requires_func(task=task, reqs=reqs, **kwargs)
+
# no requirements for workflows in pilot mode
if callable(getattr(task, "is_workflow", None)) and task.is_workflow() and getattr(task, "pilot", False):
return
# add required producers when set
if (prods := self.require_producers):
- reqs["required_producers"] = {prod: self._req_producer(task, prod) for prod in prods}
+ reqs["required_producers"] = {
+ prod: self._req_producer(task, prod)
+ for prod in law.util.make_unique(prods)
+ }
def setup_func(
self,
@@ -41,6 +55,8 @@ def setup_func(
reader_targets: law.util.InsertableDict,
**kwargs,
) -> None:
+ super().setup_func(task=task, reqs=reqs, inputs=inputs, reader_targets=reader_targets, **kwargs)
+
if "required_producers" in inputs:
for prod, inp in inputs["required_producers"].items():
reader_targets[f"required_producer_{prod}"] = inp["columns"]
@@ -62,9 +78,9 @@ def producer(
cls,
func: Callable | None = None,
bases: tuple = (),
- mc_only: bool = False,
- data_only: bool = False,
- require_producers: Sequence[str] | set[str] | None = None,
+ mc_only: bool | UNSET_TYPE = UNSET,
+ data_only: bool | UNSET_TYPE = UNSET,
+ require_producers: Sequence[str] | set[str] | None | UNSET_TYPE = UNSET,
**kwargs,
) -> DerivableMeta | Callable:
"""
@@ -89,13 +105,13 @@ def producer(
"""
def decorator(func: Callable) -> DerivableMeta:
# create the class dict
- cls_dict = {
- **kwargs,
- "call_func": func,
- "mc_only": mc_only,
- "data_only": data_only,
- "require_producers": require_producers,
- }
+ cls_dict = {**kwargs, "call_func": func}
+ if mc_only is not UNSET:
+ cls_dict["mc_only"] = mc_only
+ if data_only is not UNSET:
+ cls_dict["data_only"] = data_only
+ if require_producers is not UNSET:
+ cls_dict["require_producers"] = require_producers
# get the module name
frame = inspect.stack()[1]
@@ -114,8 +130,7 @@ def update_cls_dict(cls_name, cls_dict, get_attr):
raise Exception(f"producer {cls_name} received both mc_only and data_only")
if (mc_only or data_only) and cls_dict.get("skip_func"):
raise Exception(
- f"producer {cls_name} received custom skip_func, but either mc_only or "
- "data_only are set",
+ f"producer {cls_name} received custom skip_func, but either mc_only or data_only are set",
)
if "skip_func" not in cls_dict:
diff --git a/columnflow/production/categories.py b/columnflow/production/categories.py
index d825caf45..b0bf7ca5a 100644
--- a/columnflow/production/categories.py
+++ b/columnflow/production/categories.py
@@ -14,8 +14,7 @@
from columnflow.categorization import Categorizer
from columnflow.production import Producer, producer
from columnflow.util import maybe_import
-from columnflow.columnar_util import set_ak_column
-from columnflow.columnar_util_Ghent import safe_concatenate
+from columnflow.columnar_util import set_ak_column, ak_concatenate_safe
np = maybe_import("numpy")
ak = maybe_import("awkward")
@@ -58,7 +57,7 @@ def category_ids(
category_ids.append(ak.singletons(ak.nan_to_none(ids)))
# combine
- category_ids = safe_concatenate(category_ids, axis=1)
+ category_ids = ak_concatenate_safe(category_ids, axis=1)
# save, optionally on a target events array
if target_events is None:
@@ -70,6 +69,8 @@ def category_ids(
@category_ids.init
def category_ids_init(self: Producer, **kwargs) -> None:
+ super(category_ids, self).init_func(**kwargs)
+
# store a mapping from leaf category to categorizer classes for faster lookup
self.categorizer_map = {}
diff --git a/columnflow/production/cms/btag.py b/columnflow/production/cms/btag.py
index 27c6ab023..d11afd96a 100644
--- a/columnflow/production/cms/btag.py
+++ b/columnflow/production/cms/btag.py
@@ -6,14 +6,16 @@
from __future__ import annotations
-from dataclasses import dataclass, field
+import dataclasses
import law
+import order as od
from columnflow.production import Producer, producer
+from columnflow.columnar_util import set_ak_column, DotDict, TAFConfig, EMPTY_FLOAT
+from columnflow.hist_util import sum_hists, merge_axis_bins
from columnflow.util import maybe_import, load_correction_set
-from columnflow.columnar_util import set_ak_column, flat_np_view, layout_ak_array, DotDict
-from columnflow.types import Any
+from columnflow.types import Any, Callable, Sequence
np = maybe_import("numpy")
ak = maybe_import("awkward")
@@ -22,13 +24,13 @@
logger = law.logger.get_logger(__name__)
-@dataclass
-class BTagSFConfig:
+@dataclasses.dataclass
+class BTagSFConfig(TAFConfig):
correction_set: str
sources: list[str]
jec_sources: list[str]
discriminator: str = "" # when empty, set in post init based on correction set
- corrector_kwargs: dict[str, Any] = field(default_factory=dict)
+ corrector_kwargs: dict[str, Any] = dataclasses.field(default_factory=dict)
def __post_init__(self):
cs = self.correction_set.lower()
@@ -156,9 +158,6 @@ def btag_weights(
f"known values are {','.join(known_log_modes)}",
)
- # get the total number of jets in the chunk
- n_jets_all = ak.sum(ak.num(events.Jet, axis=1))
-
# check that the b-tag score is not negative for all jets considered in the SF calculation
discr = events.Jet[self.btag_config.discriminator]
jets_negative_b_score = discr[jet_mask] < 0
@@ -186,14 +185,14 @@ def btag_weights(
# set jet mask to False when b_score is negative
jet_mask = (discr >= 0) & (True if jet_mask is Ellipsis else jet_mask)
- # get flat inputs, evaluated at jet_mask
- flavor = flat_np_view(events.Jet.hadronFlavour[jet_mask], axis=1)
- abs_eta = flat_np_view(abs(events.Jet.eta[jet_mask]), axis=1)
- pt = flat_np_view(events.Jet.pt[jet_mask], axis=1)
- discr_flat = flat_np_view(discr[jet_mask], axis=1)
+ # get inputs, evaluated at jet_mask
+ flavor = events.Jet.hadronFlavour[jet_mask]
+ abs_eta = abs(events.Jet.eta[jet_mask])
+ pt = events.Jet.pt[jet_mask]
+ discr = discr[jet_mask]
# fix edge cases where the discriminator is non-finite
- discr_flat[~np.isfinite(discr_flat)] = 0
+ discr = ak.where(np.isfinite(discr), discr, 0.0)
# helper to create and store the weight
def add_weight(syst_name, syst_direction, column_name):
@@ -212,32 +211,19 @@ def add_weight(syst_name, syst_direction, column_name):
"flavor": flavor[flavor_mask],
"abseta": abs_eta[flavor_mask],
"pt": pt[flavor_mask],
- "discriminant": discr_flat[flavor_mask],
+ "discriminant": discr[flavor_mask],
**self.btag_config.corrector_kwargs,
}
- # get the flat scale factors
- sf_flat = self.btag_sf_corrector(*(
+ # get the scale factors
+ sf = self.btag_sf_corrector(*(
variable_map[inp.name]
for inp in self.btag_sf_corrector.inputs
))
- # insert them into an array of ones whose length corresponds to the total number of jets
- sf_flat_all = np.ones(n_jets_all, dtype=np.float32)
- if jet_mask is Ellipsis:
- indices = flavor_mask
- else:
- indices = flat_np_view(jet_mask)
- if flavor_mask is not Ellipsis:
- indices = np.where(indices)[0][flavor_mask]
- sf_flat_all[indices] = sf_flat
-
- # enforce the correct shape and create the product over all jets per event
- sf = layout_ak_array(sf_flat_all, events.Jet.pt)
-
if negative_b_score_action == "remove":
# set the weight to 0 for jets with negative btag score
- sf = ak.where(jets_negative_b_score, 0, sf)
+ sf = ak.where(jets_negative_b_score, 0.0, sf)
weight = ak.prod(sf, axis=1, mask_identity=False)
@@ -287,6 +273,7 @@ def btag_weights_post_init(self: Producer, task: law.Task, **kwargs) -> None:
# 2. when the nominal shift is requested, the central weight and all variations related to the
# method-intrinsic shifts are produced
# 3. when any other shift is requested, only create the central weight column
+ super(btag_weights, self).post_init_func(task=task, **kwargs)
# NOTE: we currently setup the produced columns only during the post_init. This means
# that the `produces` of this Producer will be empty during task initialization, meaning
@@ -340,6 +327,8 @@ def btag_weights_requires(
reqs: dict[str, DotDict[str, Any]],
**kwargs,
) -> None:
+ super(btag_weights, self).requires_func(task=task, reqs=reqs, **kwargs)
+
if "external_files" in reqs:
return
@@ -356,6 +345,370 @@ def btag_weights_setup(
reader_targets: law.util.InsertableDict,
**kwargs,
) -> None:
+ super(btag_weights, self).setup_func(task=task, reqs=reqs, inputs=inputs, reader_targets=reader_targets, **kwargs)
+
# load the btag sf corrector
btag_file = self.get_btag_file(reqs["external_files"].files)
self.btag_sf_corrector = load_correction_set(btag_file)[self.btag_config.correction_set]
+
+
+@dataclasses.dataclass
+class BTagWPSFConfig(TAFConfig):
+ # name of the jet collection
+ jet_name: str = "Jet"
+ # name of the b-tag score column
+ btag_column: str = "btagUParTAK4B"
+ # name of the correction set to load
+ correction_set: str = "UParTAK4_merged"
+ # mapping of working point names and thresholds
+ # ! values are merely examples and should be overwritten
+ btag_wps: dict[str, float] = dataclasses.field(default_factory=lambda: {
+ "loose": 0.0246,
+ "medium": 0.1272,
+ "tight": 0.4648,
+ "xtight": 0.6298,
+ "xxtight": 0.9739,
+ })
+ # edges for histogram re-binning when set
+ # ! note that, when given, these edges need to be a valid subset of the original bin edges of the counting hists
+ pt_edges: tuple[float, ...] | None = None
+ abs_eta_edges: tuple[float, ...] | None = None
+ # same for merging working points in a dictionary like {"xtight": ["xtight", "xxtight"]} to merge xtight and xxtight
+ # and call the resulting bin on the "wp" axis "xtight"
+ wp_merging: dict[str, Sequence[str]] = dataclasses.field(default_factory=dict)
+ # key of the tagging counts histogram to load from MergeSelectionStats output
+ hist_key: str = "btag_wp_counts"
+ # name of the weight column to produce
+ weight_name: str = "btag_weight"
+ # mapping of systematic variations to postfixes to be added to the weight name
+ systs: dict[str, str] = dataclasses.field(default_factory=dict)
+ # a function that, given a dataset_inst, returns other dataset_inst's (or their names) whose counting histograms
+ # should be summed, can also be a sequence of lists containing dataset_inst's (converted to callable in init)
+ # (see "binning recommendations" in ref 1 below)
+ dataset_groups: (
+ Callable[[od.Dataset], Sequence[od.Dataset | str] | None] |
+ Sequence[Sequence[od.Dataset | str]] |
+ None
+ ) = None
+
+ def __post_init__(self) -> None:
+ # ensure group_datasets is always a callable
+ if not self.dataset_groups:
+ self.dataset_groups = lambda dataset_inst: [dataset_inst]
+ elif not callable(self.dataset_groups):
+ # convert nested sequence to list of sets for faster lookup
+ groups = list(set(seq) for seq in self.dataset_groups)
+ def dataset_groups(dataset_inst: od.Dataset) -> list[od.Dataset] | None:
+ for group in groups:
+ if dataset_inst in group or dataset_inst.name in group:
+ return list(group)
+ return None
+ self.dataset_groups = dataset_groups
+
+
+@producer(
+ # only run on mc
+ mc_only=True,
+ # function to determine the correction file
+ get_btag_wp_file=(lambda self, external_files: external_files.btag_wp_sf_corr),
+ # function to configure how to retrieve the BTagWPSFConfig
+ get_btag_wp_sf_config=(lambda self: self.config_inst.x.btag_wp_sf_config),
+)
+def btag_wp_weights(
+ self: Producer,
+ events: ak.Array,
+ task: law.Task,
+ jet_mask: ak.Array | type(Ellipsis) = Ellipsis,
+ **kwargs,
+) -> ak.Array:
+ """
+ B-tag scale factor weight producer using the fixed working point method. Requires an external file in the config as
+ under ``btag_wp_sf_corr``:
+
+ .. code-block:: python
+
+ cfg.x.external_files = DotDict.wrap({
+ "btag_wp_sf_corr": "/afs/cern.ch/work/m/mrieger/public/hbt/external_files/custom_btv_files/btag_merged_2024.json.gz", # noqa
+ })
+
+ *get_btag_wp_file* can be adapted in a subclass in case it is stored differently in the external files.
+
+ To configure the tagger, working points, column names and other settings, a :py:class:`BTagWPSFConfig` object should
+ be registered in the config like the following.
+
+ .. code-block:: python
+
+ cfg.x.btag_wp_sf_config = BTagWPSFConfig(
+ jet_name="Jet",
+ btag_column="btagUParTAK4B",
+ correction_set="UParTAK4_merged",
+ )
+
+ *get_btag_wp_sf_config* can be adapted in a subclass in case it is stored differently in the config.
+
+ Resources:
+ 1. https://btv-wiki.docs.cern.ch/PerformanceCalibration/fixedWPSFRecommendations/#recommendations-for-fixed-working-point-sfs # noqa
+ 2. https://cms-analysis-corrections.docs.cern.ch/corrections_era/Run3-24CDEReprocessingFGHIPrompt-Summer24-NanoAODv15/BTV/2026-01-30/#btagging_preliminaryjsongz # noqa
+ 3. https://indico.cern.ch/event/1583955/contributions/6771046/attachments/3176162/5648591/BTVreportHIGPAG_18112025.pdf # noqa
+ """
+ # get inputs, evaluated at jet_mask
+ discr = events[self.cfg.jet_name][self.cfg.btag_column][jet_mask]
+ flavor = events[self.cfg.jet_name].hadronFlavour[jet_mask]
+ abs_eta = abs(events[self.cfg.jet_name].eta[jet_mask])
+ pt = events[self.cfg.jet_name].pt[jet_mask]
+
+ # helpers to get the sf and efficiencies
+ def get_sf_and_eff(
+ wp: str,
+ variable_map: dict[str, ak.Array],
+ check_zeros: bool,
+ ) -> tuple[ak.Array, ak.Array]:
+ variable_map = variable_map.copy()
+
+ # add BTV compatible working point name
+ variable_map["working_point"] = self.convert_wp_str[wp]
+
+ # get scale factor
+ sf = self.btag_wp_sf_corrector(
+ *(variable_map[inp.name] for inp in self.btag_wp_sf_corrector.inputs),
+ )
+
+ # add further variables needed by the efficiency corrector
+ variable_map["wp"] = wp
+
+ # get efficiency
+ eff = self.wp_eff_corrector(
+ *(variable_map[inp.name] for inp in self.wp_eff_corrector.inputs),
+ )
+
+ # complain in cases of empty efficiency values
+ if ak.any(empty_mask := eff == EMPTY_FLOAT):
+ vals = list(zip(
+ ak.flatten(variable_map["flavor"][empty_mask]),
+ ak.flatten(variable_map["pt"][empty_mask]),
+ ak.flatten(variable_map["abs_eta"][empty_mask]),
+ ))
+ err = (
+ f"encountered empty values for {len(vals)} jet(s) during the lookup of jet tagging efficiencies at "
+ f"working point '{wp}'; either update the binning of the efficiency map to have coarser binning, or "
+ "define dataset groups whose tagging counts should be combined and summed such that the per-bin "
+ "statistic is sufficiently large; empty values found for the following (flavor, pt, abs_eta) values: "
+ "\n" + "\n".join(f" - ({f}, {p:.4f}, {e:.4f})" for f, p, e in vals)
+ )
+ raise ValueError(err)
+
+ # complain in cases of zero efficiency values
+ if check_zeros and ak.any(zero_mask := eff == 0):
+ vals = list(zip(
+ ak.flatten(variable_map["flavor"][zero_mask]),
+ ak.flatten(variable_map["pt"][zero_mask]),
+ ak.flatten(variable_map["abs_eta"][zero_mask]),
+ ))
+ err = (
+ f"encountered zero values for {len(vals)} jet(s) during the lookup of jet tagging efficiencies at "
+ f"working point '{wp}'; either update the binning of the efficiency map to have coarser binning, or "
+ "define dataset groups whose tagging counts should be combined and summed such that the per-bin "
+ "statistic is sufficiently large; empty values found for the following (flavor, pt, abs_eta) values: "
+ "\n" + "\n".join(f" - ({f}, {p:.4f}, {e:.4f})" for f, p, e in vals)
+ )
+ raise ValueError(err)
+
+ return sf, eff
+
+ # helper to create and store the weight
+ def add_wp_weight(syst_name: str, column_name: str) -> ak.Array:
+ # initialize b-tag event weight with ones
+ btag_weight = ak.ones_like(events.event, dtype=np.float32)
+
+ # apply WP masks to each jet in the event falling between the WP thresholds
+ wp_names = list(self.sorted_wps.keys())
+ for lower_key, upper_key in zip([None] + wp_names, wp_names + [None]):
+ jet_mask = (
+ (discr >= self.sorted_wps.get(lower_key, -np.inf)) &
+ (discr < self.sorted_wps.get(upper_key, np.inf))
+ )
+
+ # prepare arguments with specific WP jet_mask applied
+ variable_map = {
+ "systematic": syst_name,
+ "flavor": flavor[jet_mask],
+ "abseta": abs_eta[jet_mask],
+ "abs_eta": abs_eta[jet_mask],
+ "pt": pt[jet_mask],
+ }
+
+ # handle three cases for sf calculation
+ if lower_key is None: # jets that fail lowest WP
+ sf_fail, wp_eff_fail = get_sf_and_eff(upper_key, variable_map, False)
+ sf_pass, wp_eff_pass = 1.0, 1.0
+ elif upper_key is None: # jets that pass highest WP
+ sf_fail, wp_eff_fail = 0.0, 0.0
+ sf_pass, wp_eff_pass = get_sf_and_eff(lower_key, variable_map, True)
+ else: # jets that pass one WP but fail the next one
+ sf_fail, wp_eff_fail = get_sf_and_eff(upper_key, variable_map, False)
+ sf_pass, wp_eff_pass = get_sf_and_eff(lower_key, variable_map, True)
+
+ # calculate scale factor term per jet, depending on WP efficiency
+ denominator = wp_eff_pass - wp_eff_fail
+ denominator = ak.where(denominator <= 0, 1.0, denominator)
+ sf_term = (sf_pass * wp_eff_pass - sf_fail * wp_eff_fail) / denominator
+
+ # handle edge cases where the SF becomes negative (unphysical) or zero (killing the event)
+ sf_term = ak.where(sf_term <= 0, 1.0, sf_term)
+
+ # calculate final b-tag event weight as a product of the individual jet scale factor terms
+ btag_weight = btag_weight * ak.prod(sf_term, axis=1, mask_identity=False)
+
+ # add the column and return
+ return set_ak_column(events, column_name, btag_weight, value_type=np.float32)
+
+ # always produce the central weight
+ events = add_wp_weight("central", self.cfg.weight_name)
+
+ # produce variations when the nominal shift was requested
+ if task.global_shift_inst.is_nominal:
+ for syst_name, col_postfix in self.cfg.systs.items():
+ events = add_wp_weight(syst_name, f"{self.cfg.weight_name}_{col_postfix}")
+
+ return events
+
+
+@btag_wp_weights.init
+def btag_wp_weights_init(self: Producer, **kwargs) -> None:
+ super(btag_wp_weights, self).init_func(**kwargs)
+
+ self.cfg = self.get_btag_wp_sf_config()
+
+ # keep a list of all dataset insts whose tagging counts should be grouped (summed)
+ group = self.cfg.dataset_groups(self.dataset_inst)
+ log_id = f"{self.cls_name}_{self.config_inst.name}_{self.dataset_inst.name}_group_lookup"
+ if group:
+ # make sure to deal with instances rather than names
+ group = [
+ (dataset if isinstance(dataset, od.Dataset) else self.config_inst.get_dataset(dataset))
+ for dataset in group
+ ]
+ logger.debug_once(
+ log_id,
+ f"{self.cls_name}: found {len(group)} dataset(s) for grouping tagging counts for requested dataset "
+ f"'{self.dataset_inst.name}': {', '.join(dataset.name for dataset in group)}",
+ )
+ else:
+ logger.warning_once(
+ log_id,
+ f"{self.cls_name}: found no dataset group for requested dataset '{self.dataset_inst.name}', please check "
+ f"the 'dataset_groups' setting in the {self.cfg.__class__.__name__} object",
+ )
+ group = [self.dataset_inst]
+ self.dataset_group = group
+
+
+@btag_wp_weights.post_init
+def btag_wp_weights_post_init(self: Producer, task: law.Task, **kwargs) -> None:
+ super(btag_wp_weights, self).post_init_func(task=task, **kwargs)
+
+ # add used columns
+ self.uses.add(f"{self.cfg.jet_name}.{{pt,eta,phi,mass,hadronFlavour,{self.cfg.btag_column}}}")
+
+ # produce the nominal weight
+ self.produces.add(self.cfg.weight_name)
+
+ # produce varied weights only when the nominal shift is requested
+ if task.global_shift_inst.is_nominal:
+ for col_postfix in self.cfg.systs.values():
+ self.produces.add(f"{self.cfg.weight_name}_{col_postfix}")
+
+
+@btag_wp_weights.requires
+def btag_wp_weights_requires(
+ self: Producer,
+ task: law.Task,
+ reqs: dict[str, DotDict[str, Any]],
+ **kwargs,
+) -> None:
+ super(btag_wp_weights, self).requires_func(task=task, reqs=reqs, **kwargs)
+
+ if "external_files" not in reqs:
+ from columnflow.tasks.external import BundleExternalFiles
+ reqs["external_files"] = BundleExternalFiles.req(task)
+
+ if "grouped_selection_stats" not in reqs:
+ from columnflow.tasks.selection import MergeSelectionStats
+ reqs["grouped_selection_stats"] = {
+ dataset_inst.name: MergeSelectionStats.req_different_branching(
+ task,
+ dataset=dataset_inst.name,
+ branch=-1 if task.is_workflow() else 0,
+ )
+ for dataset_inst in self.dataset_group
+ }
+
+
+@btag_wp_weights.setup
+def btag_wp_weights_setup(
+ self: Producer,
+ task: law.Task,
+ reqs: dict[str, DotDict[str, Any]],
+ inputs: dict[str, Any],
+ reader_targets: law.util.InsertableDict,
+ **kwargs,
+) -> None:
+ super(btag_wp_weights, self).setup_func(
+ task=task,
+ reqs=reqs,
+ inputs=inputs,
+ reader_targets=reader_targets,
+ **kwargs,
+ )
+
+ import hist
+ import correctionlib as clib
+ import correctionlib.convert
+
+ # load the btag sf corrector
+ btag_wp_file = self.get_btag_wp_file(reqs["external_files"].files)
+ self.btag_wp_sf_corrector = load_correction_set(btag_wp_file)[self.cfg.correction_set]
+
+ # load the count histograms and compute efficiencies
+ counts = sum_hists([
+ inps["hists"].load(formatter="pickle")[self.cfg.hist_key]
+ for inps in inputs["grouped_selection_stats"].values()
+ ])
+ # optionally rebin pt and abs_eta axes
+ if self.cfg.pt_edges:
+ counts = counts[{"pt": hist.rebin(edges=self.cfg.pt_edges)}]
+ if self.cfg.abs_eta_edges:
+ counts = counts[{"abs_eta": hist.rebin(edges=self.cfg.abs_eta_edges)}]
+ # merge wp bins if necessary
+ for merged_wp, wps in self.cfg.wp_merging.items():
+ counts = merge_axis_bins(counts, "wp", merged_wp, wps)
+ # get the total counts
+ counts_total = counts[{"wp": "total"}].view()
+ # compute efficiencies
+ effs = counts[{"wp": [wp for wp in counts.axes["wp"] if wp != "total"]}]
+ eff_values = effs.view()
+ eff_values[...] /= counts_total[..., None]
+ # bins with a count of zero will contain nan's, so replace them with a placeholder
+ eff_values[counts_total == 0] = EMPTY_FLOAT
+
+ # convert to clib corrector
+ effs.name = "btag_efficiencies"
+ effs.label = "eff"
+ cset = clib.convert.from_histogram(effs)
+ cset.data.flow = "clamp" # clamps the pt and abs eta axis
+ self.wp_eff_corrector = cset.to_evaluator()
+
+ # dictionary for BTV naming scheme compatibility
+ self.convert_wp_str = {
+ "loose": "L",
+ "medium": "M",
+ "tight": "T",
+ "xtight": "XT",
+ "xxtight": "XXT",
+ }
+
+ # sort the WPs therehold values and make the WP strings all lowercase
+ self.sorted_wps = dict(sorted(
+ ((k.lower(), v) for k, v in self.cfg.btag_wps.items()),
+ key=(lambda tpl: tpl[1]),
+ ))
diff --git a/columnflow/production/cms/dy.py b/columnflow/production/cms/dy.py
index 9e618c007..e389af043 100644
--- a/columnflow/production/cms/dy.py
+++ b/columnflow/production/cms/dy.py
@@ -139,7 +139,7 @@ def dy_weights(self: Producer, events: ak.Array, **kwargs) -> ak.Array:
.. code-block:: python
cfg.x.external_files = DotDict.wrap({
- "dy_weight_sf": "/afs/cern.ch/work/m/mrieger/public/mirrors/external_files/DY_pTll_weights_v2.json.gz", # noqa
+ "dy_weight_sf": "/afs/cern.ch/work/m/mrieger/public/mirrors/external_files/hbt_corrections.json.gz", # noqa
})
*get_dy_weight_file* can be adapted in a subclass in case it is stored differently in the external files.
@@ -170,19 +170,19 @@ def dy_weights(self: Producer, events: ak.Array, **kwargs) -> ak.Array:
if callable(self.dy_config.get_njets):
variable_map["njets"] = self.dy_config.get_njets(self, events)
if callable(self.dy_config.get_nbtags):
- variable_map["nbtags"] = self.dy_config.get_nbtags(self, events)
- # for compatibility
- variable_map["ntags"] = variable_map["nbtags"]
+ variable_map["ntags"] = self.dy_config.get_nbtags(self, events)
# initializing the list of weight variations (called syst in the dy files)
- systs = [("nom", "")]
+ systs = []
# add specific uncertainties or additional systs
if self.dy_config.unc_correction:
+ systs.append(("nom", ""))
for i in range(self.n_unc):
for direction in ["up", "down"]:
systs.append((f"{direction}{i + 1}", f"_{direction}{i + 1}"))
elif self.dy_config.systs:
+ systs.append(("nominal", ""))
for syst in self.dy_config.systs:
systs.append((syst, f"_{syst}"))
@@ -201,8 +201,10 @@ def dy_weights(self: Producer, events: ak.Array, **kwargs) -> ak.Array:
@dy_weights.init
-def dy_weights_init(self: Producer) -> None:
- if self.config_inst.campaign.x.year not in {2022, 2023}:
+def dy_weights_init(self: Producer, **kwargs) -> None:
+ super(dy_weights, self).init_func(**kwargs)
+
+ if self.config_inst.campaign.x.year not in {2022, 2023, 2024}:
raise NotImplementedError(
f"campaign year {self.config_inst.campaign.x.year} is not yet supported by {self.cls_name}",
)
@@ -226,10 +228,12 @@ def dy_weights_init(self: Producer) -> None:
@dy_weights.requires
-def dy_weights_requires(self: Producer, task: law.Task, reqs: dict) -> None:
+def dy_weights_requires(self: Producer, task: law.Task, reqs: dict, **kwargs) -> None:
"""
Adds the requirements needed the underlying task to derive the Drell-Yan weights into *reqs*.
"""
+ super(dy_weights, self).requires_func(task=task, reqs=reqs, **kwargs)
+
if "external_files" in reqs:
return
@@ -244,12 +248,15 @@ def dy_weights_setup(
reqs: dict,
inputs: dict,
reader_targets: law.util.InsertableDict,
+ **kwargs,
) -> None:
"""
Loads the Drell-Yan weight calculator from the external files bundle and saves them in the py:attr:`dy_corrector`
attribute for simpler access in the actual callable. The number of uncertainties is calculated, per era, by another
correcter in the external file and is saved in the py:attr:`dy_unc_corrector` attribute.
"""
+ super(dy_weights, self).setup_func(task=task, reqs=reqs, inputs=inputs, reader_targets=reader_targets, **kwargs)
+
bundle = reqs["external_files"]
# import all correctors from the external file
@@ -450,7 +457,9 @@ def recoil_corrected_met(self: Producer, events: ak.Array, **kwargs) -> ak.Array
@recoil_corrected_met.init
-def recoil_corrected_met_init(self: Producer) -> None:
+def recoil_corrected_met_init(self: Producer, **kwargs) -> None:
+ super(recoil_corrected_met, self).init_func(**kwargs)
+
if self.njet_column:
self.uses.add(f"{self.njet_column}")
else:
@@ -458,8 +467,9 @@ def recoil_corrected_met_init(self: Producer) -> None:
@recoil_corrected_met.requires
-def recoil_corrected_met_requires(self: Producer, task: law.Task, reqs: dict) -> None:
- # Ensure that external files are bundled.
+def recoil_corrected_met_requires(self: Producer, task: law.Task, reqs: dict, **kwargs) -> None:
+ super(recoil_corrected_met, self).requires_func(task=task, reqs=reqs, **kwargs)
+
if "external_files" in reqs:
return
@@ -474,7 +484,16 @@ def recoil_corrected_met_setup(
reqs: dict,
inputs: dict,
reader_targets: law.util.InsertableDict,
+ **kwargs,
) -> None:
+ super(recoil_corrected_met, self).setup_func(
+ task=task,
+ reqs=reqs,
+ inputs=inputs,
+ reader_targets=reader_targets,
+ **kwargs,
+ )
+
# load the correction set
bundle = reqs["external_files"]
correction_set = load_correction_set(self.get_dy_recoil_file(bundle.files))
diff --git a/columnflow/production/cms/electron.py b/columnflow/production/cms/electron.py
index 89b739daa..159470e7b 100644
--- a/columnflow/production/cms/electron.py
+++ b/columnflow/production/cms/electron.py
@@ -12,7 +12,7 @@
from columnflow.production import Producer, producer
from columnflow.util import maybe_import, load_correction_set, DotDict
-from columnflow.columnar_util import set_ak_column, full_like, flat_np_view
+from columnflow.columnar_util import set_ak_column, full_like, flat_np_view, layout_ak_array
from columnflow.types import Any, Callable
np = maybe_import("numpy")
@@ -61,7 +61,7 @@ def new(cls, obj: ElectronSFConfig | tuple[str, str, str]) -> ElectronSFConfig:
use_supercluster_eta=True,
# name of the saved weight column
weight_name="electron_weight",
- supported_versions={1, 2, 3},
+ supported_versions={1, 2, 3, 4, 5},
)
def electron_weights(
self: Producer,
@@ -151,8 +151,7 @@ def electron_weights(
sf = self.electron_sf_corrector.evaluate(*inputs)
elif isinstance(wp, dict):
# mapping of wps to masks, evaluate per wp and combine
- sf = full_like(eta, 1.0)
- sf_flat = flat_np_view(sf)
+ sf_flat = flat_np_view(full_like(eta, 1.0))
for _wp, mask_fn in wp.items():
mask = mask_fn(variable_map)
variable_map_syst_wp = variable_map_syst | {"WorkingPoint": _wp}
@@ -166,6 +165,7 @@ def electron_weights(
for inp in self.electron_sf_corrector.inputs
]
sf_flat[flat_np_view(mask)] = flat_np_view(self.electron_sf_corrector.evaluate(*inputs))
+ sf = layout_ak_array(sf_flat, eta)
else:
raise ValueError(f"unsupported working point type {type(variable_map['WorkingPoint'])}")
@@ -180,6 +180,8 @@ def electron_weights(
@electron_weights.init
def electron_weights_init(self: Producer, **kwargs) -> None:
+ super(electron_weights, self).init_func(**kwargs)
+
# add the product of nominal and up/down variations to produced columns
self.produces.add(f"{self.weight_name}{{,_up,_down}}")
@@ -191,6 +193,8 @@ def electron_weights_requires(
reqs: dict[str, DotDict[str, Any]],
**kwargs,
) -> None:
+ super(electron_weights, self).requires_func(task=task, reqs=reqs, **kwargs)
+
if "external_files" in reqs:
return
@@ -207,6 +211,14 @@ def electron_weights_setup(
reader_targets: law.util.InsertableDict,
**kwargs,
) -> None:
+ super(electron_weights, self).setup_func(
+ task=task,
+ reqs=reqs,
+ inputs=inputs,
+ reader_targets=reader_targets,
+ **kwargs,
+ )
+
self.electron_config = self.get_electron_config()
# load the corrector
diff --git a/columnflow/production/cms/gen_particles.py b/columnflow/production/cms/gen_particles.py
index 294af1a81..7fe3fba57 100644
--- a/columnflow/production/cms/gen_particles.py
+++ b/columnflow/production/cms/gen_particles.py
@@ -10,7 +10,7 @@
import law
from columnflow.production import Producer, producer
-from columnflow.columnar_util import set_ak_column
+from columnflow.columnar_util import set_ak_column, ak_concatenate_safe
from columnflow.util import UNSET, maybe_import
np = maybe_import("numpy")
@@ -106,7 +106,7 @@ def gen_top_lookup(self: Producer, events: ak.Array, strict: bool = True, **kwar
# sort them so that down-type quarks and charged leptons (odd pdgIds) come first, followed by up-type quarks and
# neutrinos (even pdgIds), then add back the remaining ones
w_children_hard = w_children_hard[ak.argsort(-(w_children_hard.pdgId % 2), axis=2)]
- w_children = ak.concatenate([w_children_hard, w_children_rest], axis=2)
+ w_children = ak_concatenate_safe([w_children_hard, w_children_rest], axis=2)
# further distinguish tau decays in w_children
w_tau_children = ak.drop_none(w_children[abs(w_children.pdgId) == 15].distinctChildrenDeep)
@@ -115,7 +115,7 @@ def gen_top_lookup(self: Producer, events: ak.Array, strict: bool = True, **kwar
w_tau_photon_mask = w_tau_children.pdgId == 22
w_tau_rest = w_tau_children[~(w_tau_nu_mask | w_tau_photon_mask)]
w_tau_rest = w_tau_rest[ak.argsort(abs(w_tau_rest.pdgId), axis=3, ascending=True)]
- w_tau_children = ak.concatenate(
+ w_tau_children = ak_concatenate_safe(
[w_tau_children[w_tau_nu_mask], w_tau_rest, w_tau_children[w_tau_photon_mask]],
axis=3,
)
@@ -214,9 +214,9 @@ def gen_higgs_lookup(self: Producer, events: ak.Array, strict: bool = True, **kw
tau_w = ak.where(tau_has_rest, tau_w_rest, tau_w)
tau_w = set_ak_column(tau_w, "pdgId", ak.values_astype(-24 * np.sign(tau.pdgId), np.int32))
# combine nu and w again
- tau_nuw = ak.concatenate([tau_nu[..., None], tau_w[..., None]], axis=3)
+ tau_nuw = ak_concatenate_safe([tau_nu[..., None], tau_w[..., None]], axis=3)
# define w children
- tau_w_children = ak.concatenate(
+ tau_w_children = ak_concatenate_safe(
[tau_children[tau_rest_mask], ak.drop_none(ak.firsts(tau_children[tau_w_mask], axis=3).children)],
axis=2,
)
@@ -335,9 +335,9 @@ def gen_dy_lookup(self: Producer, events: ak.Array, strict: bool = True, **kwarg
tau_w = ak.where(tau_has_rest, tau_w_rest, tau_w)
tau_w = set_ak_column(tau_w, "pdgId", ak.values_astype(-24 * np.sign(tau.pdgId), np.int32))
# combine nu and w again
- tau_nuw = ak.concatenate([tau_nu[..., None], tau_w[..., None]], axis=2)
+ tau_nuw = ak_concatenate_safe([tau_nu[..., None], tau_w[..., None]], axis=2)
# define w children
- tau_w_children = ak.concatenate(
+ tau_w_children = ak_concatenate_safe(
[tau_children[tau_rest_mask], ak.drop_none(ak.firsts(tau_children[tau_w_mask], axis=2).children)],
axis=1,
)
diff --git a/columnflow/production/cms/jet.py b/columnflow/production/cms/jet.py
index b37736df8..b062f4f6e 100644
--- a/columnflow/production/cms/jet.py
+++ b/columnflow/production/cms/jet.py
@@ -12,7 +12,7 @@
from columnflow.production import Producer, producer
from columnflow.util import maybe_import, load_correction_set
-from columnflow.columnar_util import set_ak_column, layout_ak_array, flat_np_view
+from columnflow.columnar_util import set_ak_column, layout_ak_array, flat_np_view, ak_concatenate_safe
np = maybe_import("numpy")
ak = maybe_import("awkward")
@@ -109,7 +109,7 @@ def jet_id(self: Producer, events: ak.Array, **kwargs) -> ak.Array:
jet_id_flat[valid_mask] |= id_flag << (pass_bit - 1)
# apply correct layout
- jet_id = layout_ak_array(jet_id_flat, events[self.jet_name].eta)
+ jet_id = layout_ak_array(jet_id_flat, events[self.jet_name])
# store them
events = set_ak_column(events, f"{self.jet_name}.jetId", jet_id, value_type=np.uint8)
@@ -122,6 +122,8 @@ def jet_id_init(self: Producer, **kwargs) -> None:
"""
Dynamically add the names of the used and produced columns depending on the jet name.
"""
+ super(jet_id, self).init_func(**kwargs)
+
self.jet_columns = ["eta", "chHEF", "neHEF", "chEmEF", "neEmEF", "muEF", "chMultiplicity", "neMultiplicity"]
self.uses.update(f"{self.jet_name}.{col}" for col in self.jet_columns)
self.produces.add(f"{self.jet_name}.jetId")
@@ -132,6 +134,8 @@ def jet_id_requires(self: Producer, task: law.Task, reqs: dict, **kwargs) -> Non
"""
Adds the requirements needed the underlying task to recompute the jet id into *reqs*.
"""
+ super(jet_id, self).requires_func(task=task, reqs=reqs, **kwargs)
+
if "external_files" in reqs:
return
@@ -151,6 +155,8 @@ def jet_id_setup(
"""
Sets up the correction sets needed for the jet id using the external files.
"""
+ super(jet_id, self).setup_func(task=task, reqs=reqs, inputs=inputs, reader_targets=reader_targets, **kwargs)
+
bundle = reqs["external_files"]
# get the jet id config
@@ -208,8 +214,8 @@ def msoftdrop(self: Producer, events: ak.Array, **kwargs) -> ak.Array:
valid_subjet = padded_subjet[valid_subjet_idxs]
valid_subjets.append(ak.singletons(valid_subjet, axis=-1))
- # merge lists for all sibjet index columns
- valid_subjets = ak.concatenate(valid_subjets, axis=-1)
+ # merge lists for all subjet index columns
+ valid_subjets = ak_concatenate_safe(valid_subjets, axis=-1)
# attach coffea behavior so we can do LV arithmetic
valid_subjets = ak.with_name(
@@ -243,6 +249,8 @@ def msoftdrop_init(self: Producer, **kwargs) -> None:
"""
Dynamically add `uses` and `produces`.
"""
+ super(msoftdrop, self).init_func(**kwargs)
+
# input columns
self.uses |= {
f"{collection}.{var}"
@@ -262,6 +270,7 @@ def msoftdrop_init(self: Producer, **kwargs) -> None:
@msoftdrop.setup
def msoftdrop_setup(self: Producer, task: law.Task, reqs: dict, **kwargs) -> None:
- import coffea
+ super(msoftdrop, self).setup_func(task=task, reqs=reqs, **kwargs)
+ import coffea
self.nano_behavior = coffea.nanoevents.NanoAODSchema.behavior()
diff --git a/columnflow/production/cms/muon.py b/columnflow/production/cms/muon.py
index 762e6b544..659755836 100644
--- a/columnflow/production/cms/muon.py
+++ b/columnflow/production/cms/muon.py
@@ -56,6 +56,8 @@ def __post_init__(self):
# name of the saved weight column
weight_name="muon_weight",
supported_versions={1, 2},
+ # function to update variables before corrector call
+ update_corrector_variables=(lambda self, corrector, variables: variables),
)
def muon_weights(
self: Producer,
@@ -112,12 +114,17 @@ def muon_weights(
("systdown", "_down"),
]:
# get the inputs for this type of variation
- variable_map_syst = {
+ _variable_map = {
**variable_map,
- "scale_factors": "nominal" if syst == "sf" else syst, # syst key in 2022
- "ValType": syst, # syst key in 2017
+ "scale_factors": "nominal" if syst == "sf" else syst,
+ "ValType": syst,
}
- inputs = [variable_map_syst[inp.name] for inp in self.muon_sf_corrector.inputs]
+
+ # optionally update variables for this corrector call
+ if callable(self.update_corrector_variables):
+ _variable_map = self.update_corrector_variables(self.muon_sf_corrector, _variable_map)
+
+ inputs = [_variable_map[inp.name] for inp in self.muon_sf_corrector.inputs]
sf = self.muon_sf_corrector.evaluate(*inputs)
# create the product over all muons in one event
@@ -131,6 +138,8 @@ def muon_weights(
@muon_weights.init
def muon_weights_init(self: Producer, **kwargs) -> None:
+ super(muon_weights, self).init_func(**kwargs)
+
# add the product of nominal and up/down variations to produced columns
self.produces.add(f"{self.weight_name}{{,_up,_down}}")
@@ -142,6 +151,8 @@ def muon_weights_requires(
reqs: dict[str, DotDict[str, Any]],
**kwargs,
) -> None:
+ super(muon_weights, self).requires_func(task=task, reqs=reqs, **kwargs)
+
if "external_files" in reqs:
return
@@ -158,6 +169,8 @@ def muon_weights_setup(
reader_targets: law.util.InsertableDict,
**kwargs,
) -> None:
+ super(muon_weights, self).setup_func(task=task, reqs=reqs, inputs=inputs, reader_targets=reader_targets, **kwargs)
+
bundle = reqs["external_files"]
# load the corrector
diff --git a/columnflow/production/cms/pdf.py b/columnflow/production/cms/pdf.py
index 6d70a2aa6..8d660d4db 100644
--- a/columnflow/production/cms/pdf.py
+++ b/columnflow/production/cms/pdf.py
@@ -227,6 +227,8 @@ def pdf_weights(
@pdf_weights.init
def pdf_weight_init(self: Producer, **kwargs) -> None:
+ super(pdf_weights, self).init_func(**kwargs)
+
# add produced columns: nominal+all, or nominal+up+down
self.produces.add("pdf_weight{,s}" if self.store_all_weights else "pdf_weight{,_up,_down}")
self.produces.add("alphas_weight{,_up,_down}")
diff --git a/columnflow/production/cms/pileup.py b/columnflow/production/cms/pileup.py
index 6675e8d04..f87bd819c 100644
--- a/columnflow/production/cms/pileup.py
+++ b/columnflow/production/cms/pileup.py
@@ -9,8 +9,8 @@
import law
from columnflow.production import Producer, producer
-from columnflow.util import maybe_import, DotDict
from columnflow.columnar_util import set_ak_column, fill_at
+from columnflow.util import maybe_import, DotDict, load_correction_set
from columnflow.types import Any
np = maybe_import("numpy")
@@ -113,6 +113,8 @@ def pu_weight_requires(
"""
Adds the requirements needed the underlying task to derive the pileup weights into *reqs*.
"""
+ super(pu_weight, self).requires_func(task=task, reqs=reqs, **kwargs)
+
if "external_files" in reqs:
return
@@ -133,16 +135,12 @@ def pu_weight_setup(
Loads the pileup calculator from the external files bundle and saves them in the
py:attr:`pileup_corrector` attribute for simpler access in the actual callable.
"""
+ super(pu_weight, self).setup_func(task=task, reqs=reqs, inputs=inputs, reader_targets=reader_targets, **kwargs)
+
bundle = reqs["external_files"]
# create the corrector
- import correctionlib
- correctionlib.highlevel.Correction.__call__ = correctionlib.highlevel.Correction.evaluate
- correction_set = correctionlib.CorrectionSet.from_string(
- self.get_pileup_file(bundle.files).load(formatter="gzip").decode("utf-8"),
- )
-
- # check
+ correction_set = load_correction_set(self.get_pileup_file(bundle.files))
if len(correction_set.keys()) != 1:
raise Exception("Expected exactly one type of pileup correction")
@@ -185,6 +183,8 @@ def pu_weights_from_columnflow_requires(
"""
Adds the requirements needed the underlying task to derive the pileup weights into *reqs*.
"""
+ super(pu_weights_from_columnflow, self).requires_func(task=task, reqs=reqs, **kwargs)
+
if "pu_weights" in reqs:
return
@@ -205,4 +205,12 @@ def pu_weights_from_columnflow_setup(
Loads the pileup weights added through the requirements and saves them in the
py:attr:`pu_weights` attribute for simpler access in the actual callable.
"""
+ super(pu_weights_from_columnflow, self).setup_func(
+ task=task,
+ reqs=reqs,
+ inputs=inputs,
+ reader_targets=reader_targets,
+ **kwargs,
+ )
+
self.pu_weights = ak.zip(inputs["pu_weights"].load(formatter="json"))
diff --git a/columnflow/production/cms/scale.py b/columnflow/production/cms/scale.py
index 0a44fefa2..43e5d705b 100644
--- a/columnflow/production/cms/scale.py
+++ b/columnflow/production/cms/scale.py
@@ -36,6 +36,8 @@ def setup_func(
reader_targets: law.util.InsertableDict,
**kwargs,
) -> None:
+ super().setup_func(task=task, reqs=reqs, inputs=inputs, reader_targets=reader_targets, **kwargs)
+
# named weight indices
self.indices_9 = DotDict(
mur_down_muf_down=0,
diff --git a/columnflow/production/cms/seeds.py b/columnflow/production/cms/seeds.py
index f8834b8db..5bc9d2037 100644
--- a/columnflow/production/cms/seeds.py
+++ b/columnflow/production/cms/seeds.py
@@ -139,6 +139,8 @@ def deterministic_event_seeds_init(self, **kwargs) -> None:
Producer initialization that adds columns to the set of *used* columns based on the
*event_columns*, *object_count_columns*, and *object_columns* lists.
"""
+ super(deterministic_event_seeds, self).init_func(**kwargs)
+
# add used columns
for column in self.event_columns + self.object_count_columns + self.object_columns:
self.uses.add(Route(column))
@@ -156,6 +158,14 @@ def deterministic_event_seeds_setup(
"""
Setup function that defines conventions methods needed during the producer function.
"""
+ super(deterministic_event_seeds, self).setup_func(
+ task=task,
+ reqs=reqs,
+ inputs=inputs,
+ reader_targets=reader_targets,
+ **kwargs,
+ )
+
# store primes in array
self.primes = np.array(primes, dtype=np.uint64)
@@ -230,6 +240,7 @@ def call_func(self, events: ak.Array, **kwargs) -> ak.Array:
return events
def init_func(self, **kwargs) -> None:
+ super().init_func(**kwargs)
self.uses |= {f"{self.object_field}.pt"}
self.produces |= {f"{self.object_field}.deterministic_seed"}
@@ -249,6 +260,8 @@ def setup_func(
:param inputs: Dictionary for inputs (not used).
:param reader_targets: Dictionary for additional column to retrieve (not used).
"""
+ super().setup_func(task=task, reqs=reqs, inputs=inputs, reader_targets=reader_targets, **kwargs)
+
# store primes in array
self.primes = np.array(primes, dtype=np.uint64)
diff --git a/columnflow/production/cms/top_pt_weight.py b/columnflow/production/cms/top_pt_weight.py
index 8207414d2..1c3942379 100644
--- a/columnflow/production/cms/top_pt_weight.py
+++ b/columnflow/production/cms/top_pt_weight.py
@@ -127,7 +127,9 @@ def top_pt_weight(self: Producer, events: ak.Array, **kwargs) -> ak.Array:
@top_pt_weight.init
-def top_pt_weight_init(self: Producer) -> None:
+def top_pt_weight_init(self: Producer, **kwargs) -> None:
+ super(top_pt_weight, self).init_func(**kwargs)
+
# store the top pt weight config
self.cfg = self.get_top_pt_weight_config()
if not isinstance(self.cfg, (TopPtWeightFromDataConfig, TopPtWeightFromTheoryConfig)):
diff --git a/columnflow/production/matching.py b/columnflow/production/matching.py
index fadec5161..73c50ccce 100644
--- a/columnflow/production/matching.py
+++ b/columnflow/production/matching.py
@@ -86,6 +86,8 @@ def delta_r_matcher_init(self: Producer, **kwargs) -> None:
"""
Dynamically add `uses` and `produces`
"""
+ super(delta_r_matcher, self).init_func(**kwargs)
+
# input columns
self.uses |= {
f"{collection}.{var}"
diff --git a/columnflow/production/normalization.py b/columnflow/production/normalization.py
index c118b3468..c0808b358 100644
--- a/columnflow/production/normalization.py
+++ b/columnflow/production/normalization.py
@@ -349,6 +349,8 @@ def normalization_weights_init(self: Producer, **kwargs) -> None:
"""
Initializes the normalization weights producer by setting up the normalization weight column.
"""
+ super(normalization_weights, self).init_func(**kwargs)
+
# declare the weight name to be a produced column
self.produces.add(self.weight_name)
@@ -376,6 +378,8 @@ def normalization_weights_requires(
"""
Adds the requirements needed by the underlying py:attr:`task` to access selection stats into *reqs*.
"""
+ super(normalization_weights, self).requires_func(task=task, reqs=reqs, **kwargs)
+
# check that all datasets are known
for dataset in self.required_datasets:
if not self.config_inst.has_dataset(dataset):
@@ -410,10 +414,15 @@ def normalization_weights_setup(
This weight is defined as the product of the luminosity, the cross section, divided by the sum of event
weights per process.
- py: attr: `known_process_ids`: A set of all process ids that are known by the lookup table.
- - py: attr: `process_weight_table`: A sparse array serving as a lookup table for the calculated process weights.
- This weight is defined as the product of the luminosity, the cross section, divided by the sum of event
- weights per process.
"""
+ super(normalization_weights, self).setup_func(
+ task=task,
+ reqs=reqs,
+ inputs=inputs,
+ reader_targets=reader_targets,
+ **kwargs,
+ )
+
import scipy.sparse
# load the selection stats
@@ -489,7 +498,7 @@ def extract_stats(*update_funcs):
)
# setup the event weight lookup table
- process_weight_table = scipy.sparse.lil_matrix((max(process_ids) + 1, 1), dtype=np.float32)
+ process_weight_table = scipy.sparse.dok_matrix((max(process_ids) + 1, 1), dtype=np.float32)
def fill_weight_table(process_inst: od.Process, xsec: float, sum_weights: float) -> None:
if sum_weights == 0:
@@ -507,7 +516,13 @@ def fill_weight_table(process_inst: od.Process, xsec: float, sum_weights: float)
# prepare info for the inclusive dataset
inclusive_proc = self.inclusive_dataset.processes.get_first()
- inclusive_xsec = inclusive_proc.get_xsec(self.config_inst.campaign.ecm).nominal
+ try:
+ inclusive_xsec = inclusive_proc.get_xsec(self.config_inst.campaign.ecm).nominal
+ except KeyError as e:
+ raise KeyError(
+ f"no cross section registered for inclusive process {inclusive_proc} for center-of-mass energy of "
+ f"{self.config_inst.campaign.ecm}",
+ ) from e
# compute the weight the inclusive dataset would have on its own without stitching
if self.allow_stitching and self.dataset_inst == self.inclusive_dataset:
diff --git a/columnflow/production/util.py b/columnflow/production/util.py
index 938876282..31053f3f4 100644
--- a/columnflow/production/util.py
+++ b/columnflow/production/util.py
@@ -8,10 +8,17 @@
import functools
-from columnflow.types import Iterable, Sequence, Union
+import law
+
from columnflow.production import Producer, producer
+from columnflow.columnar_util import (
+ TaskArrayFunction,
+ Route,
+ set_ak_column,
+ attach_coffea_behavior as attach_coffea_behavior_fn,
+)
from columnflow.util import maybe_import
-from columnflow.columnar_util import attach_coffea_behavior as attach_coffea_behavior_fn
+from columnflow.types import Iterable, Sequence, Union, Callable
ak = maybe_import("awkward")
@@ -202,3 +209,37 @@ def delta_r_match_multiple(
# return either index or four-vector of best match
best_match = best_match_idxs if as_index else dst_lvs[best_match_idxs]
return best_match, dst_lvs_filtered
+
+
+def transfer_produced_columns(
+ func: TaskArrayFunction,
+ src_array: ak.Array,
+ dst_array: ak.Array,
+ filter_routes: Sequence[str] | set[str] | Callable[[Route], bool] | None = None,
+) -> ak.Array:
+ """
+ Transfers all columns produced by a :py:class:`TaskArrayFunction` from a source array *src_array* to a destination
+ array *dst_array*. Optionally, only columns produced for certain routes can be transferred by specifying
+ *filter_routes*, which can be a sequence or set of route names or patterns, or a callable that receives a
+ :py:class:`Route`.
+
+ :param func: :py:class:`TaskArrayFunction` that produced columns in *src_array*.
+ :param src_array: Source array containing the produced columns.
+ :param dst_array: Destination array to which the produced columns are transferred.
+ :param filter_routes: Optional filter.
+ :return: Destination array with transferred columns.
+ """
+ # prepare filtering
+ if not filter_routes:
+ filter_routes = lambda r: True
+ elif not callable(filter_routes):
+ patterns = set(filter_routes)
+ filter_routes = lambda r: law.util.multi_match(str(r), patterns, mode=any)
+
+ # start transferring
+ for r in func.produced_columns:
+ if not filter_routes(r):
+ continue
+ dst_array = set_ak_column(dst_array, r, r.apply(src_array))
+
+ return dst_array
diff --git a/columnflow/reduction/default.py b/columnflow/reduction/default.py
index 21c9da51c..11083c9d0 100644
--- a/columnflow/reduction/default.py
+++ b/columnflow/reduction/default.py
@@ -31,6 +31,8 @@ def cf_default_keep_columns(self: Reducer, events: ak.Array, selection: ak.Array
@cf_default_keep_columns.post_init
def cf_default_keep_columns_post_init(self: Reducer, task: law.Task, **kwargs) -> None:
+ super(cf_default_keep_columns, self).post_init_func(task=task, **kwargs)
+
for c in self.config_inst.x.keep_columns.get(task.task_family, ["*"]):
self.produces.update(task._expand_keep_column(c))
@@ -40,6 +42,8 @@ def cf_default_keep_columns_post_init(self: Reducer, task: law.Task, **kwargs) -
check_used_columns=False,
# whether to add cf_default_keep_columns as a dependency to achieve backwards compatibility
add_keep_columns=True,
+ # whether to register the shifts of the upstream selector as shifts of this reducer
+ mirror_selector_shifts=True,
)
def cf_default(self: Reducer, events: ak.Array, selection: ak.Array, task: law.Task, **kwargs) -> ak.Array:
# build the event mask
@@ -57,13 +61,21 @@ def cf_default(self: Reducer, events: ak.Array, selection: ak.Array, task: law.T
@cf_default.init
def cf_default_init(self: Reducer, **kwargs) -> None:
+ super(cf_default, self).init_func(**kwargs)
+
if self.add_keep_columns:
self.uses.add(cf_default_keep_columns.PRODUCES)
self.produces.add(cf_default_keep_columns.PRODUCES)
+ # mirror selector shifts
+ if self.mirror_selector_shifts and "selector_shifts" in self.inst_dict:
+ self.shifts |= self.selector_shifts
+
@cf_default.post_init
def cf_default_post_init(self: Reducer, task: law.Task, **kwargs) -> None:
+ super(cf_default, self).post_init_func(task=task, **kwargs)
+
# the updates to used columns are only necessary if the task invokes the reducer
if not task.invokes_reducer:
return
diff --git a/columnflow/selection/__init__.py b/columnflow/selection/__init__.py
index 1f4368fc0..53fd98a07 100644
--- a/columnflow/selection/__init__.py
+++ b/columnflow/selection/__init__.py
@@ -13,8 +13,8 @@
import order as od
from columnflow.calibration import TaskArrayFunctionWithCalibratorRequirements
-from columnflow.util import maybe_import, DotDict, DerivableMeta
-from columnflow.types import Callable, T, Sequence
+from columnflow.util import maybe_import, DotDict, DerivableMeta, UNSET
+from columnflow.types import Callable, T, Sequence, UNSET_TYPE
ak = maybe_import("awkward")
np = maybe_import("numpy")
@@ -46,9 +46,9 @@ def selector(
cls,
func: Callable | None = None,
bases=(),
- mc_only: bool = False,
- data_only: bool = False,
- require_calibrators: Sequence[str] | set[str] | None = None,
+ mc_only: bool | UNSET_TYPE = UNSET,
+ data_only: bool | UNSET_TYPE = UNSET,
+ require_calibrators: Sequence[str] | set[str] | None | UNSET_TYPE = UNSET,
**kwargs,
) -> DerivableMeta | Callable:
"""
@@ -72,13 +72,13 @@ def selector(
"""
def decorator(func: Callable) -> DerivableMeta:
# create the class dict
- cls_dict = {
- **kwargs,
- "call_func": func,
- "mc_only": mc_only,
- "data_only": data_only,
- "require_calibrators": require_calibrators,
- }
+ cls_dict = {**kwargs, "call_func": func}
+ if mc_only is not UNSET:
+ cls_dict["mc_only"] = mc_only
+ if data_only is not UNSET:
+ cls_dict["data_only"] = data_only
+ if require_calibrators is not UNSET:
+ cls_dict["require_calibrators"] = require_calibrators
# get the module name
frame = inspect.stack()[1]
@@ -97,8 +97,7 @@ def update_cls_dict(cls_name, cls_dict, get_attr):
raise Exception(f"selector {cls_name} received both mc_only and data_only")
if (mc_only or data_only) and cls_dict.get("skip_func"):
raise Exception(
- f"selector {cls_name} received custom skip_func, but either mc_only or "
- "data_only are set",
+ f"selector {cls_name} received custom skip_func, but either mc_only or data_only are set",
)
if "skip_func" not in cls_dict:
diff --git a/columnflow/selection/cms/btag.py b/columnflow/selection/cms/btag.py
new file mode 100644
index 000000000..fce483c27
--- /dev/null
+++ b/columnflow/selection/cms/btag.py
@@ -0,0 +1,109 @@
+# coding: utf-8
+
+"""
+B-tag selection methods.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+
+from columnflow.selection import Selector, selector
+from columnflow.columnar_util import TAFConfig
+from columnflow.util import maybe_import, DotDict
+from columnflow.types import TYPE_CHECKING
+
+np = maybe_import("numpy")
+ak = maybe_import("awkward")
+
+if TYPE_CHECKING:
+ hist = maybe_import("hist")
+
+
+@dataclasses.dataclass
+class BTagWPCountConfig(TAFConfig):
+ # name of the jet collection
+ jet_name: str = "Jet"
+ # name of the b-tag score column
+ btag_column: str = "btagUParTAK4B"
+ # mapping of working point names and thresholds
+ # ! values are merely examples and should be overwritten
+ btag_wps: dict[str, float] = dataclasses.field(default_factory=lambda: {
+ "loose": 0.0246,
+ "medium": 0.1272,
+ "tight": 0.4648,
+ "xtight": 0.6298,
+ "xxtight": 0.9739,
+ })
+ # edges for histogram binning
+ # ! values are merely examples and should be overwritten
+ pt_edges: tuple[float, ...] = (0, 20, 30, 50, 70, 100, 140, 200, 300, 600, 10_000)
+ abs_eta_edges: tuple[float, ...] = (0.0, 1.0, 1.5, 2.0, 5.0)
+ # key of the histogram to save in selector hists
+ hist_key: str = "btag_wp_counts"
+
+
+@selector(
+ # function to configure how to retrieve the BTagWPCountConfig
+ get_btag_wp_count_config=(lambda self: self.config_inst.x.btag_wp_count_config),
+ # only run on mc
+ mc_only=True,
+ # unexpose so that it cannot be called from the command line
+ exposed=False,
+)
+def fill_btag_wp_count_hists(
+ self: Selector,
+ events: ak.Array,
+ event_mask: ak.Array,
+ jet_mask: ak.Array,
+ hists: DotDict[str, hist.Hist],
+ **kwargs,
+) -> None:
+ """
+ Selector that fills numbers of selected jets passing different b-tagging working points into histograms. There is no
+ return value as the created histogram is stored in *hists* in-place.
+
+ The counts can be used later on to compute b-tagging efficiencies, required by scale factor calculations such as
+ https://btv-wiki.docs.cern.ch/PerformanceCalibration/fixedWPSFRecommendations/#b-tagging-efficiencies-in-simulation.
+ """
+ import hist
+
+ # select jets and reduce to selected events only
+ jets = events[self.cfg.jet_name][jet_mask][event_mask]
+
+ # flatten necessary columns
+ pt = ak.flatten(jets.pt, axis=None)
+ abs_eta = abs(ak.flatten(jets.eta, axis=None))
+ flavor = ak.flatten(jets.hadronFlavour, axis=None)
+ btag_score = ak.flatten(jets[self.cfg.btag_column], axis=None)
+
+ # create empty histogram
+ h = hist.Hist(
+ hist.axis.Variable(self.cfg.pt_edges, name="pt", label="pt"),
+ hist.axis.Variable(self.cfg.abs_eta_edges, name="abs_eta", label="abs_eta"),
+ hist.axis.IntCategory(categories=[], growth=True, name="flavor", label="flavor"),
+ hist.axis.StrCategory(["total", *self.cfg.btag_wps.keys()], name="wp", label="wp"),
+ storage=hist.storage.Double(),
+ )
+
+ # fill total wp bin
+ h.fill(pt=pt, abs_eta=abs_eta, flavor=flavor, wp="total")
+
+ # fill wp bins
+ for wp, threshold in self.cfg.btag_wps.items():
+ mask = btag_score >= threshold
+ h.fill(pt=pt[mask], abs_eta=abs_eta[mask], flavor=flavor[mask], wp=wp)
+
+ # store the histogram
+ hists[self.cfg.hist_key] = h
+
+
+@fill_btag_wp_count_hists.init
+def fill_btag_wp_count_hists_init(self: Selector, **kwargs) -> None:
+ super(fill_btag_wp_count_hists, self).init_func(**kwargs)
+
+ # retrieve and store the config
+ self.cfg: BTagWPCountConfig = self.get_btag_wp_count_config()
+
+ # add used columns
+ self.uses.add(f"{self.cfg.jet_name}.{{pt,eta,phi,mass,hadronFlavour,{self.cfg.btag_column}}}")
diff --git a/columnflow/selection/cms/jets.py b/columnflow/selection/cms/jets.py
index 20c008778..8a5e693a9 100644
--- a/columnflow/selection/cms/jets.py
+++ b/columnflow/selection/cms/jets.py
@@ -11,7 +11,7 @@
from columnflow.selection import Selector, SelectionResult, selector
from columnflow.util import maybe_import, load_correction_set, DotDict
-from columnflow.columnar_util import set_ak_column, flat_np_view, optional_column as optional
+from columnflow.columnar_util import set_ak_column, flat_np_view, layout_ak_array
from columnflow.types import Any
np = maybe_import("numpy")
@@ -22,11 +22,9 @@
@selector(
- uses={
- "Jet.{pt,eta,phi,mass,jetId,chEmEF}", optional("Jet.puId"),
- "Muon.{pt,eta,phi,mass,isPFcand}",
- },
+ # all used columns are registered in init below
produces={"Jet.veto_map_mask"},
+ use_lepton_veto_id=None, # depends on the "run" info of the underlying campaign of the config
get_veto_map_file=(lambda self, external_files: external_files.jet_veto_map),
)
def jet_veto_map(
@@ -35,14 +33,16 @@ def jet_veto_map(
**kwargs,
) -> tuple[ak.Array, SelectionResult]:
"""
- Selector that applies the Jet Veto Map to the jets and stores the result as a new column ``Jet.veto_maps``.
- Additionally, the ``jet_veto_map`` step is added to the SelectionResult that masks events containing
- jets from the veto map, which is the recommended way to use the veto map.
- For users that only want to remove the jets from the veto map, the ``veto_map_jets`` object
- is added to the SelectionResult.
+ Selector that applies the jet veto map to the jets and stores the result as a new column ``Jet.veto_map_mask``.
+ Additionally, the ``jet_veto_map`` step is added to the SelectionResult that masks events containing jets from the
+ veto map, which is the recommended way to use the veto map.
- Requires an external file in the config
- under ``jet_veto_map``:
+ A minimal selection is applied to jets that are fed into the veto map check as documented in [1]. However, this
+ recommendation depends on the data taking period and uses either the proximity to muons in the event (run 2), or the
+ so-called "tightLepVeto" (run 3). When *use_lepton_veto_id* is *None*, the decision is based on the ``run``
+ auxiliary field of the campaign of the config. Otherwise, when *True*, the "tightLepVeto" is used.
+
+ Requires an external file in the config under ``jet_veto_map``:
.. code-block:: python
@@ -52,32 +52,41 @@ def jet_veto_map(
*get_veto_map_file* can be adapted in a subclass in case it is stored differently in the external files.
- documentation: https://cms-jerc.web.cern.ch/Recommendations/#jet-veto-maps
+ Resources:
+
+ 1. https://cms-jerc.web.cern.ch/Recommendations/#jet-veto-maps
+ 2. https://cms-talk.web.cern.ch/t/updated-jet-selection-criterion-for-jet-veto-map/130527
"""
+ # jet selection
jet = events.Jet
- muon = events.Muon[events.Muon.isPFcand]
-
- # loose jet selection
jet_mask = (
(jet.pt > 15) &
- (jet.jetId >= 2) & # tight id
- (jet.chEmEF < 0.9) &
- ak.all(events.Jet.metric_table(muon) >= 0.2, axis=2)
+ ((jet.chEmEF + jet.neEmEF) < 0.9)
)
+ # fold in veto id or manually filter against muons
+ if self.use_lepton_veto_id:
+ jet_mask = jet_mask & (jet.jetId & (1 << 2) != 0) # third bit is tightLepVeto
+ else:
+ muon = events.Muon[events.Muon.isPFcand]
+ jet_mask = jet_mask & (
+ (jet.jetId & (1 << 1) != 0) & # second bit is tight
+ ak.all(events.Jet.metric_table(muon) >= 0.2, axis=2)
+ )
+
# apply loose Jet puId in Run 2 to jets with pt below 50 GeV
if self.config_inst.campaign.x.run == 2:
- jet_pu_mask = (events.Jet.puId >= 4) | (events.Jet.pt >= 50)
- jet_mask = jet_mask & jet_pu_mask
-
- jet_phi = jet.phi
- jet_eta = jet.eta
+ jet_mask = jet_mask & (
+ (events.Jet.puId >= 4) |
+ (events.Jet.pt >= 50)
+ )
# for some reason, math.pi is not included in the ranges, so we need to subtract a small number
pi = math.pi - 1e-10
# values outside [-pi, pi] are not included, so we need to clip the phi values
- phi_outside_range = abs(jet.phi) > pi
+ jet_phi = jet.phi
+ phi_outside_range = abs(jet_phi) > pi
if ak.any(phi_outside_range):
# warn in severe cases
if ak.any(severe := abs(jet_phi[phi_outside_range]) >= 3.15):
@@ -86,23 +95,24 @@ def jet_veto_map(
"with phi values outside [-pi, pi] that will be clipped",
)
jet_phi = ak.where(
- np.abs(jet.phi) > pi,
- jet.phi - 2 * pi * np.sign(jet.phi),
- jet.phi,
+ np.abs(jet_phi) > pi,
+ jet_phi - 2 * pi * np.sign(jet_phi),
+ jet_phi,
)
# values outside [-5.19, 5.19] are not included, so we need to clip the eta values
- eta_outside_range = np.abs(jet.eta) > 5.19
+ jet_eta = jet.eta
+ eta_outside_range = np.abs(jet_eta) > 5.19
if ak.any(eta_outside_range):
- jet_eta = ak.where(
- np.abs(jet.eta) > 5.19,
- 5.19 * np.sign(jet.eta),
- jet.eta,
- )
logger.warning(
- f"Jet eta values {jet.eta[eta_outside_range][ak.any(eta_outside_range, axis=1)]} outside [-5.19, 5.19] "
- f"({ak.sum(eta_outside_range)} in total) "
- f"detected and set to {jet_eta[eta_outside_range][ak.any(eta_outside_range, axis=1)]}",
+ f"jet eta values {jet_eta[eta_outside_range][ak.any(eta_outside_range, axis=1)]} outside [-5.19, 5.19] "
+ f"({ak.sum(eta_outside_range)} in total) detected and set to "
+ f"{jet_eta[eta_outside_range][ak.any(eta_outside_range, axis=1)]}",
+ )
+ jet_eta = ak.where(
+ np.abs(jet_eta) > 5.19,
+ 5.19 * np.sign(jet_eta),
+ jet_eta,
)
# evalute the veto map only for selected jets
@@ -121,9 +131,10 @@ def jet_veto_map(
inputs = [variable_map[inp.name] for inp in self.veto_map.inputs]
veto_mask_sel = veto_mask_sel & ~(self.veto_map(*inputs) != 0)
- # insert back into full jet mask in-place
- flat_jet_mask = flat_np_view(jet_mask)
+ # combine again with jet mask
+ flat_jet_mask = flat_np_view(jet_mask, copy=True)
flat_jet_mask[flat_jet_mask] = ak.flatten(veto_mask_sel)
+ jet_mask = layout_ak_array(flat_jet_mask, events.Jet)
# store the per-jet veto mask for further processing
# note: to be consistent with conventions, the exported values should be True for passing jets
@@ -137,6 +148,26 @@ def jet_veto_map(
return events, results
+@jet_veto_map.init
+def get_veto_map_init(self: Selector, **kwargs) -> None:
+ super(jet_veto_map, self).init_func(**kwargs)
+
+ # get the default for use_lepton_veto_id
+ if self.use_lepton_veto_id is None:
+ self.use_lepton_veto_id = self.config_inst.campaign.x.run == 3
+
+ # always read specific jet columns
+ self.uses.add("Jet.{pt,eta,phi,mass,jetId,chEmEF,neEmEF}")
+
+ # read puId in run 2 for additional cut
+ if self.config_inst.campaign.x.run == 2:
+ self.uses.add("Jet.puId")
+
+ # read muon columns when not using the veto id
+ if not self.use_lepton_veto_id:
+ self.uses.add("Muon.{pt,eta,phi,mass,isPFcand}")
+
+
@jet_veto_map.requires
def jet_veto_map_requires(
self: Selector,
@@ -144,6 +175,8 @@ def jet_veto_map_requires(
reqs: dict[str, DotDict[str, Any]],
**kwargs,
) -> None:
+ super(jet_veto_map, self).requires_func(task=task, reqs=reqs, **kwargs)
+
if "external_files" in reqs:
return
@@ -160,6 +193,8 @@ def jet_veto_map_setup(
reader_targets: law.util.InsertableDict,
**kwargs,
) -> None:
+ super(jet_veto_map, self).setup_func(task=task, reqs=reqs, inputs=inputs, reader_targets=reader_targets, **kwargs)
+
# create the corrector
map_file = self.get_veto_map_file(reqs["external_files"].files)
correction_set = load_correction_set(map_file)
diff --git a/columnflow/selection/cms/json_filter.py b/columnflow/selection/cms/json_filter.py
index 6eddb84d1..27a851b20 100644
--- a/columnflow/selection/cms/json_filter.py
+++ b/columnflow/selection/cms/json_filter.py
@@ -98,6 +98,8 @@ def json_filter_requires(
reqs: dict[str, DotDict[str, Any]],
**kwargs,
) -> None:
+ super(json_filter, self).requires_func(task=task, reqs=reqs, **kwargs)
+
if "external_files" in reqs:
return
@@ -122,6 +124,8 @@ def json_filter_setup(
:param inputs: Additional inputs, currently not used
:param reader_targets: Additional targets, currently not used
"""
+ super(json_filter, self).setup_func(task=task, reqs=reqs, inputs=inputs, reader_targets=reader_targets, **kwargs)
+
import scipy.sparse
bundle = reqs["external_files"]
@@ -134,7 +138,7 @@ def json_filter_setup(
max_run = max(map(int, json.keys()))
# build lookup table
- self.run_ls_lookup = scipy.sparse.lil_matrix((max_run + 1, max_ls + 1), dtype=bool)
+ self.run_ls_lookup = scipy.sparse.dok_matrix((max_run + 1, max_ls + 1), dtype=bool)
for run, ls_ranges in json.items():
run = int(run)
for ls_range in ls_ranges:
diff --git a/columnflow/selection/cms/met_filters.py b/columnflow/selection/cms/met_filters.py
index 3033b8641..ded216ace 100644
--- a/columnflow/selection/cms/met_filters.py
+++ b/columnflow/selection/cms/met_filters.py
@@ -88,10 +88,12 @@ def met_filters(
@met_filters.init
def met_filters_init(self: Selector, **kwargs) -> None:
- met_filters = self.get_met_filters()
- if isinstance(met_filters, dict):
- met_filters = met_filters[self.dataset_inst.data_source]
+ super(met_filters, self).init_func(**kwargs)
+
+ filters = self.get_met_filters()
+ if isinstance(filters, dict):
+ filters = filters[self.dataset_inst.data_source]
# store filters as an attribute for faster lookup
- self.met_filters = set(met_filters)
+ self.met_filters = set(filters)
self.uses |= self.met_filters
diff --git a/columnflow/selection/empty.py b/columnflow/selection/empty.py
index 563846e47..b0b6e234c 100644
--- a/columnflow/selection/empty.py
+++ b/columnflow/selection/empty.py
@@ -98,6 +98,8 @@ def empty_init(self: Selector, **kwargs) -> None:
:raises ValueError: If the inclusive category cannot be found.
"""
+ super(empty, self).init_func(**kwargs)
+
# do nothing when category ids are set
if self.category_ids is not None:
return
diff --git a/columnflow/selection/stats.py b/columnflow/selection/stats.py
index fdd84a5e1..ba5ed06ae 100644
--- a/columnflow/selection/stats.py
+++ b/columnflow/selection/stats.py
@@ -222,6 +222,14 @@ def increment_stats_setup(
reader_targets: law.util.InsertableDict,
**kwargs,
) -> None:
+ super(increment_stats, self).setup_func(
+ task=task,
+ reqs=reqs,
+ inputs=inputs,
+ reader_targets=reader_targets,
+ **kwargs,
+ )
+
# flags to descibe "number" and "sum" fields
self.NUM, self.SUM = range(2)
diff --git a/columnflow/tasks/calibration.py b/columnflow/tasks/calibration.py
index 8b58be61f..1f39f6baa 100644
--- a/columnflow/tasks/calibration.py
+++ b/columnflow/tasks/calibration.py
@@ -129,17 +129,19 @@ def run(self):
write_columns = self.calibrator_inst.produced_columns
route_filter = RouteFilter(keep=write_columns)
- # let the lfn_task prepare the nano file (basically determine a good pfn)
- [(lfn_index, input_file)] = lfn_task.iter_nano_files(self)
+ # let the lfn_task locate and prepare the nano file(s)
+ nano_input = [nano_target for _, nano_target in lfn_task.iter_nano_files(self)]
+ if len(nano_input) == 1:
+ nano_input = nano_input[0]
# prepare inputs for localization
with law.localize_file_targets(
- [input_file, *reader_targets.values()],
+ [nano_input, *reader_targets.values()],
mode="r",
) as inps:
# iterate over chunks
for (events, *cols), pos in self.iter_chunked_io(
- [inp.abspath for inp in inps],
+ law.util.map_struct(law.target.file.get_path, inps),
source_type=["coffea_root"] + (len(inps) - 1) * [None],
read_columns=len(inps) * [read_columns],
chunk_size=self.calibrator_inst.get_min_chunk_size(),
@@ -162,8 +164,8 @@ def run(self):
self.raise_if_not_finite(events)
# save as parquet via a thread in the same pool
- chunk = tmp_dir.child(f"file_{lfn_index}_{pos.index}.parquet", type="f")
- output_chunks[(lfn_index, pos.index)] = chunk
+ chunk = tmp_dir.child(f"file_{pos.index}.parquet", type="f")
+ output_chunks[pos.index] = chunk
self.chunked_io.queue(sorted_ak_to_parquet, (events, chunk.abspath))
# teardown the calibrator
diff --git a/columnflow/tasks/cms/base.py b/columnflow/tasks/cms/base.py
index a0152b035..c24fdf81c 100644
--- a/columnflow/tasks/cms/base.py
+++ b/columnflow/tasks/cms/base.py
@@ -51,6 +51,7 @@ class CrabWorkflow(RemoteWorkflowMixin, law.cms.CrabWorkflow):
crab_forward_env_variables = {
"CF_CERN_USER": "cf_cern_user",
"CF_STORE_NAME": "cf_store_name",
+ "CF_PYTHON_VERSION": "cf_python_version",
}
# upstream requirements
diff --git a/columnflow/tasks/cms/external.py b/columnflow/tasks/cms/external.py
index 89d527244..d1f1a1df7 100644
--- a/columnflow/tasks/cms/external.py
+++ b/columnflow/tasks/cms/external.py
@@ -193,21 +193,25 @@ def run(self):
# loop through configs
for config_inst in self.config_insts:
+ if not (cat_info := config_inst.x("cat_info", None)):
+ self.logger.warning(f"no 'cat_info' entry found in config '{config_inst.name}', skipping")
+ continue
+
with self.publish_step(
f"checking CAT metadata updates for config '{law.util.colored(config_inst.name, style='bright')}' in "
- f"{config_inst.x.cat_info.metadata_root}",
+ f"{cat_info.metadata_root}",
):
newest_dates = {}
updated_any = False
- for pog, date_str in config_inst.x.cat_info.snapshot.items():
+ for pog, date_str in cat_info.snapshot.items():
if not date_str:
continue
# get all versions in the cat directory, split by date numbers
pog_era_dir = os.path.join(
- config_inst.x.cat_info.metadata_root,
+ cat_info.metadata_root,
pog.upper(),
- config_inst.x.cat_info.get_era_directory(pog),
+ cat_info.get_era_directory(pog),
)
if not os.path.isdir(pog_era_dir):
self.logger.warning(f"CAT metadata directory '{pog_era_dir}' does not exist, skipping")
@@ -226,7 +230,7 @@ def run(self):
updated_any = True
self.publish_message(
f"found newer {law.util.colored(pog.upper(), color='cyan')} snapshot: {date_str} -> "
- f"{latest_date_str} ({os.path.join(pog_era_dir, latest_date_str)})",
+ f"{latest_date_str} ({cat_info.get_era_url(pog, latest_date_str)})",
)
else:
newest_dates[pog] = date_str
diff --git a/columnflow/tasks/cms/inference.py b/columnflow/tasks/cms/inference.py
index abf8ec2ec..d3cf43690 100644
--- a/columnflow/tasks/cms/inference.py
+++ b/columnflow/tasks/cms/inference.py
@@ -15,6 +15,7 @@
from columnflow.tasks.framework.inference import SerializeInferenceModelBase
from columnflow.tasks.histograms import MergeHistograms
from columnflow.inference.cms.datacard import DatacardWriter
+from columnflow.hist_util import select_category_bins
from columnflow.types import TYPE_CHECKING
if TYPE_CHECKING:
@@ -114,9 +115,7 @@ def run(self):
for config_inst in _input_hists.keys():
config_data = cat_obj.config_data.get(config_inst.name)
- # determine leaf categories to gather
category_inst = config_inst.get_category(category)
- leaf_category_insts = category_inst.get_leaf_categories() or [category_inst]
# eagerly remove data histograms in case data is supposed to be faked from mc processes
if cat_obj.data_from_processes:
@@ -159,14 +158,13 @@ def run(self):
continue
# select relevant categories
- h_proc = h_proc[{
- "category": [
- hist.loc(c.name)
- for c in leaf_category_insts
- if c.name in h_proc.axes["category"]
- ],
- }]
- h_proc = h_proc[{"category": sum}]
+ h_proc = select_category_bins(
+ h_proc,
+ category_inst,
+ use_leaves=True,
+ prefer_parents=True,
+ reduce=True,
+ )
# create the nominal hist
datacard_hists[cat_obj.name].setdefault(proc_obj.name, {}).setdefault(config_inst.name, {})
diff --git a/columnflow/tasks/cutflow.py b/columnflow/tasks/cutflow.py
index 6a515c07e..7030bcd6d 100644
--- a/columnflow/tasks/cutflow.py
+++ b/columnflow/tasks/cutflow.py
@@ -11,41 +11,32 @@
import law
import order as od
-from columnflow.types import Any
-
from columnflow.tasks.framework.base import (
Requirements, AnalysisTask, ShiftTask, wrapper_factory, RESOLVE_DEFAULT,
)
from columnflow.tasks.framework.mixins import (
CalibratorsMixin, SelectorMixin, VariablesMixin, CategoriesMixin, ChunkedIOMixin,
- DatasetsProcessesMixin, ProducerMixin,
+ DatasetsProcessesMixin, MergeHistogramMixin,
CalibratorClassesMixin, SelectorClassMixin,
- MergeHistogramMixin,
)
from columnflow.tasks.framework.plotting import (
PlotBase, PlotBase1D, PlotBase2D, ProcessPlotSettingMixin, VariablePlotSettingMixin,
)
from columnflow.tasks.framework.remote import RemoteWorkflow
from columnflow.tasks.framework.decorators import view_output_plots
-from columnflow.tasks.framework.parameters import DerivableInstParameter, last_edge_inclusive_inst
-from columnflow.tasks.external import GetDatasetLFNs
-from columnflow.tasks.selection import SelectEvents
-from columnflow.tasks.calibration import CalibrateEvents
-from columnflow.util import DotDict, dev_sandbox, maybe_import
-from columnflow.hist_util import create_columnflow_hist, translate_hist_intcat_to_strcat
-
-
-np = maybe_import("numpy")
+from columnflow.tasks.framework.parameters import last_edge_inclusive_inst
+from columnflow.tasks.selection import MergeSelectionMasks
+from columnflow.hist_util import create_hist_from_variables, translate_hist_intcat_to_strcat, select_category_bins
+from columnflow.util import DotDict, dev_sandbox
class SelectorStepsMixin:
selector_steps_all = ("ALL",)
- # overwrite selector steps to use default resolution
selector_steps = law.CSVParameter(
default=(RESOLVE_DEFAULT,),
description="a subset of steps of the selector to apply; uses all steps when empty; "
- f"Set to {selector_steps_all[0]} to apply all."
+ f"Set to {selector_steps_all[0]} to apply all; "
"default: value of config.x.default_selector_steps",
brace_expand=True,
parse_empty=True,
@@ -70,7 +61,6 @@ class _CreateCutflowHistograms(
class CreateCutflowHistograms(_CreateCutflowHistograms):
- # strategy for handling selector steps not defined by selectors
missing_selector_step_strategy = luigi.ChoiceParameter(
significant=False,
default=law.config.get_default("analysis", "missing_selector_step_strategy", "raise"),
@@ -86,62 +76,24 @@ class CreateCutflowHistograms(_CreateCutflowHistograms):
steps_variable = od.Variable(name="step", aux={"axis_type": "strcategory"})
last_edge_inclusive = last_edge_inclusive_inst
sandbox = dev_sandbox(law.config.get("analysis", "default_columnar_sandbox"))
+ selector_steps_order_sensitive = True
initial_step = "Initial"
default_variables = ("event", "cf_*")
missing_column_alias_strategy = "original"
- norm_weights_producer = "normalization_weights"
- norm_weight_producer_inst = DerivableInstParameter(
- default=None,
- visibility=luigi.parameter.ParameterVisibility.PRIVATE,
- )
-
- exclude_params_index = {"norm_weight_producer_inst"}
- exclude_params_repr = {"norm_weight_producer_inst"}
- exclude_params_sandbox = {"norm_weight_producer_inst"}
- exclude_params_remote_workflow = {"norm_weight_producer_inst"}
-
# upstream requirements
reqs = Requirements(
RemoteWorkflow.reqs,
- GetDatasetLFNs=GetDatasetLFNs,
- CalibrateEvents=CalibrateEvents,
- SelectEvents=SelectEvents,
+ MergeSelectionMasks=MergeSelectionMasks,
)
- @classmethod
- def get_producer_dict(cls, params: dict[str, Any]) -> dict[str, Any]:
- return cls.get_array_function_dict(params)
-
- build_producer_inst = ProducerMixin.build_producer_inst
-
- @classmethod
- def resolve_instances(cls, params: dict[str, Any], shifts) -> dict[str, Any]:
- if not params.get("norm_weight_producer_inst"):
- params["norm_weight_producer_inst"] = cls.build_producer_inst(cls.norm_weights_producer, params)
-
- params = super().resolve_instances(params, shifts)
-
- return params
+ def create_branch_map(self):
+ # dummy branch map
+ return [None]
def workflow_requires(self):
reqs = super().workflow_requires()
- reqs["lfns"] = self.reqs.GetDatasetLFNs.req(self)
- if not self.pilot:
- reqs["calibrations"] = [
- self.reqs.CalibrateEvents.req(self, calibrator=calibrator_inst.cls_name)
- for calibrator_inst in self.calibrator_insts
- if calibrator_inst.produced_columns
- ]
- reqs["selection"] = self.reqs.SelectEvents.req(self)
- else:
- # pass-through pilot workflow requirements of upstream task
- t = self.reqs.SelectEvents.req(self)
- reqs = law.util.merge_dicts(reqs, t.workflow_requires(), inplace=True)
-
- if self.dataset_inst.is_mc:
- reqs["normalization"] = self.norm_weight_producer_inst.run_requires(task=self)
-
+ reqs["selection"] = self.reqs.MergeSelectionMasks.req(self, tree_index=0, _exclude={"branches"})
return reqs
def requires(self):
@@ -193,6 +145,12 @@ def run(self):
for cat in self.config_inst.get_leaf_categories()
}
+ # get IDs and names of all leaf categories
+ leaf_category_map = {
+ cat.id: cat.name
+ for cat in self.config_inst.get_leaf_categories()
+ }
+
# create a temp dir for saving intermediate files
tmp_dir = law.LocalDirectoryTarget(is_tmp=True)
tmp_dir.touch()
@@ -254,9 +212,14 @@ def prepare_hists(steps):
# create histogram of not already existing
if var_key not in histograms:
- histograms[var_key] = create_columnflow_hist(
+ histograms[var_key] = create_hist_from_variables(
self.steps_variable,
*variable_insts,
+ categorical_axes=[
+ ("category", "intcat"),
+ ("process", "intcat"),
+ ("shift", "strcat"),
+ ],
)
# let the lfn_task prepare the nano file (basically determine a good pfn)
@@ -290,7 +253,7 @@ def prepare_hists(steps):
events = self.norm_weight_producer_inst(events)
# overwrite steps if not defined yet
- if steps == self.selector_steps_all:
+ if not steps:
steps = sel.steps.fields
# prepare histograms and exprepssions once
@@ -405,7 +368,6 @@ class MergeCutflowHistograms(
):
sandbox = dev_sandbox(law.config.get("analysis", "default_columnar_sandbox"))
- # upstream requirements
reqs = Requirements(
RemoteWorkflow.reqs,
CreateHistograms=CreateCutflowHistograms,
@@ -428,14 +390,14 @@ class _PlotCutflowBase(
law.LocalWorkflow,
RemoteWorkflow,
):
- resolution_task_cls = MergeCutflowHistograms
+ resolution_task_cls = CreateCutflowHistograms
single_config = True
class PlotCutflowBase(
_PlotCutflowBase,
):
- selector_steps = MergeCutflowHistograms.selector_steps
+ selector_steps = CreateCutflowHistograms.selector_steps
sandbox = dev_sandbox(law.config.get("analysis", "default_columnar_sandbox"))
@@ -553,7 +515,6 @@ def run(self):
# prepare config objects
category_inst = self.config_inst.get_category(self.branch_data)
- leaf_category_insts = [category_inst] + (category_inst.get_leaf_categories() or [])
sub_process_insts = {
proc: [sub for sub, _, _ in proc.walk_processes(include_self=True)]
for proc in process_insts
@@ -606,16 +567,8 @@ def run(self):
_hists = OrderedDict()
for process_inst in sorted(hists, key=process_insts.index):
h = hists[process_inst]
- # selections
- h = h[{
- "category": [
- hist.loc(c.name)
- for c in leaf_category_insts
- if c.name in h.axes["category"]
- ],
- }]
- # reductions
- h = h[{"category": sum, self.variable: sum}]
+ # select and reduce categories
+ h = select_category_bins(h, category_inst, use_leaves=True, prefer_parents=True, reduce=True)
# store
_hists[process_inst] = h
hists = _hists
@@ -739,7 +692,6 @@ def run(self):
for var_name in variable_tuple
]
category_inst = self.config_inst.get_category(self.branch_data.category)
- leaf_category_insts = [category_inst] + (category_inst.get_leaf_categories() or [])
sub_process_insts = {
process_inst: [sub for sub, _, _ in process_inst.walk_processes(include_self=True)]
for process_inst in process_insts
@@ -792,16 +744,8 @@ def run(self):
_hists = OrderedDict()
for process_inst in sorted(hists, key=process_insts.index):
h = hists[process_inst]
- # selections
- h = h[{
- "category": [
- hist.loc(c.name)
- for c in leaf_category_insts
- if c.name in h.axes["category"]
- ],
- }]
- # reductions
- h = h[{"category": sum}]
+ # select and reduce categories
+ h = select_category_bins(h, category_inst, use_leaves=True, prefer_parents=True, reduce=True)
# store
_hists[process_inst] = h
hists = _hists
diff --git a/columnflow/tasks/external.py b/columnflow/tasks/external.py
index 0d901bc54..588fff67c 100644
--- a/columnflow/tasks/external.py
+++ b/columnflow/tasks/external.py
@@ -7,22 +7,23 @@
from __future__ import annotations
import os
+import glob
import time
import shutil
import subprocess
-from dataclasses import dataclass, field
-import glob
+import dataclasses
+import copy
import luigi
import law
import order as od
-from columnflow import env_is_local
+from columnflow import env_is_local, flavor as cf_flavor
from columnflow.tasks.framework.base import AnalysisTask, ConfigTask, DatasetTask, wrapper_factory
from columnflow.tasks.framework.parameters import user_parameter_inst
from columnflow.tasks.framework.decorators import only_local_env
-from columnflow.util import wget, DotDict
-from columnflow.types import Sequence
+from columnflow.util import wget, DotDict, UNSET
+from columnflow.types import Sequence, ClassVar
logger = law.logger.get_logger(__name__)
@@ -197,8 +198,8 @@ def iter_nano_files(
:param lfn_indices: List of indices of LFNs that are processed by this *task* instance, defaults to None
:param eager_lookup: Look at the next fs if stat takes too long, defaults to 1
:param skip_fallback: Skip the fallback mechanism to fetch the LFN, defaults to False
- :raises TypeError: If *task* is not of type :external+law:py:class:`~law.workflow.base.BaseWorkflow` or not
- a task analyzing a single branch in the task tree
+ :raises TypeError: If *task* is not of type :external+law:py:class:`~law.workflow.base.BaseWorkflow` or not a
+ task analyzing a single branch in the task tree
:raises Exception: If current task is not complete as indicated with ``self.complete()``
:raises ValueError: If no fs is provided at call and none can be found in either the config instance or the law
config.
@@ -233,8 +234,6 @@ def iter_nano_files(
# loop
for lfn_index in lfn_indices:
- task.publish_message(f"handling file {lfn_index}")
-
# get the lfn of the file referenced by this file index
lfn = str(lfns[lfn_index])
@@ -267,7 +266,7 @@ def iter_nano_files(
input_stat = input_file.exists(stat=True)
duration = time.perf_counter() - t1
i += 1
- logger.info(f"lfn {lfn} does{'' if input_stat else ' not'} exist at fs {selected_fs}")
+ logger.info(f"lfn {lfn} (lfn {lfn_index}) does{'' if input_stat else ' not'} exist at fs {selected_fs}")
# when the stat query took longer than some duration, eagerly try the next fs
# and check if it responds faster and if so, take it instead
@@ -293,16 +292,41 @@ def iter_nano_files(
# stop when the stat was successful at this point
if input_stat:
task.publish_message(
- f"using fs {selected_fs}, stat responded in "
- f"{law.util.human_duration(seconds=duration)}",
+ f"using fs {selected_fs}, stat responded in {law.util.human_duration(seconds=duration)}",
)
break
else:
- raise Exception(f"lfn {lfn} not found at any remote fs {fs}")
+ raise Exception(f"lfn {lfn} (lfn {lfn_index}) not found at any remote fs {fs}")
# log the file size
input_size = law.util.human_bytes(input_stat.st_size, fmt=True)
- task.publish_message(f"lfn {lfn}, size is {input_size}")
+ task.publish_message(f"lfn {lfn} (lfn {lfn_index}), size is {input_size}")
+
+ # when cf is run in cms flavor, access to central files must be reported to a database for bookkeeping
+ if cf_flavor == "cms":
+ # the selected fs must have an option "rucio_report_access" in the config, which should be either a bool
+ # or the site name (rse) to report the access for
+ report_key = "rucio_report_access"
+ report_val = law.config.get_expanded(selected_fs, report_key, default=UNSET)
+ if report_val is UNSET:
+ raise Exception(
+ f"configuration section for selected fs '{selected_fs}' is missing an entry '{report_key}' "
+ "which should be either a boolean flag or the name of a site (rse) to report the access for; "
+ "this is required for reporting access to central files which is necessary for CMS bookkeeping",
+ )
+ # check if the value is a bool, and otherwise assume a valid rse name
+ rse = None
+ try:
+ report_val = law.config.Config.instance()._convert_to_boolean(report_val)
+ except ValueError:
+ rse = report_val
+ if not law.cms.Site.validate(rse):
+ raise ValueError(
+ f"entry '{report_key}' for selected fs '{selected_fs}' does not refer to a valid site ",
+ f"name: {rse}",
+ )
+ if report_val:
+ law.cms.rucio_report_access(lfn, rse=rse)
yield (lfn_index, input_file)
@@ -389,7 +413,7 @@ def _fetch_lfn_fallback(
)
-@dataclass
+@dataclasses.dataclass
class ExternalFile:
"""
Container object to define an external file resource that is understood by (e.g.)
@@ -405,12 +429,11 @@ class ExternalFile:
"""
location: str
- subpaths: dict[str, str] = field(default_factory=dict)
+ subpaths: dict[str, str] = dataclasses.field(default_factory=dict)
version: str = "v1"
- def __str__(self) -> str:
- sub = (" / " + ",".join(f"{n}={p}" for n, p in self.subpaths.items())) if self.subpaths else ""
- return f"{self.location}{sub} ({self.version})"
+ single: bool = dataclasses.field(init=False, default=False)
+ single_key: ClassVar[str] = "_single_key"
@classmethod
def new(cls, resource: ExternalFile | str | tuple[str] | tuple[str, str]) -> ExternalFile:
@@ -429,6 +452,29 @@ def new(cls, resource: ExternalFile | str | tuple[str] | tuple[str, str]) -> Ext
return cls(location=resource[0], version=resource[1])
raise ValueError(f"invalid resource type and format: {resource}")
+ def __post_init__(self) -> None:
+ # convert different types of subpaths to dict
+ if isinstance(self.subpaths, str):
+ self.subpaths = DotDict({self.single_key: self.subpaths})
+ self.single = True
+ elif isinstance(self.subpaths, (list, tuple)):
+ self.subpaths = DotDict(zip(enumerate(self.subpaths)))
+ else:
+ self.subpaths = DotDict.wrap(copy.deepcopy(self.subpaths))
+ # remove None's
+ for key in list(self.subpaths.keys()):
+ if self.subpaths[key] is None:
+ del self.subpaths[key]
+
+ def __str__(self) -> str:
+ sub = ""
+ if self.subpaths:
+ if self.single:
+ sub = f"/{self.subpaths[self.single_key]}"
+ else:
+ sub = " / " + ",".join(f"{n}={p}" for n, p in self.subpaths.items())
+ return f"{self.location}{sub} ({self.version})"
+
def __getattr__(self, attr: str) -> str:
if attr in self.subpaths:
return self.subpaths[attr]
@@ -497,6 +543,11 @@ def create_unique_basename(cls, path: str | ExternalFile) -> str | dict[str, str
# path must be an ExternalFile
if path.subpaths:
+ # single mode
+ if path.single:
+ return cls.create_unique_basename(os.path.join(path.location, path.subpaths[path.single_key]))
+
+ # multiple subpaths
return type(path.subpaths)(
(name, cls.create_unique_basename(os.path.join(path.location, subpath)))
for name, subpath in path.subpaths.items()
@@ -567,19 +618,18 @@ def trace_transfer_output(self, output):
@law.decorator.notify
@law.decorator.log
- @law.decorator.safe_output
def run(self):
outputs = self.output()
# remove the bundle if recreating
- if outputs["bundle"].exists() and self.recreate:
+ if self.recreate and outputs["bundle"].exists():
outputs["bundle"].remove()
# bundle only if needed
if not outputs["bundle"].exists():
if not env_is_local:
raise RuntimeError(
- f"the output bundle {outputs['bundle'].basename} is missing, but cannot be created in non-local "
+ f"the output bundle {self.single_output().abspath} is missing, but cannot be created in non-local "
"environments",
)
@@ -629,9 +679,15 @@ def fetch_file(ext_file, counter=[0]):
else:
scratch_src = scratch_dst
# copy all subpaths
- basenames = self.create_unique_basename(ext_file)
- for name, subpath in ext_file.subpaths.items():
- fetch(os.path.join(scratch_src, subpath), os.path.join(tmp_dir.abspath, basenames[name]))
+ if ext_file.single:
+ fetch(
+ os.path.join(scratch_src, ext_file.subpaths[ext_file.single_key]),
+ os.path.join(tmp_dir.abspath, self.create_unique_basename(ext_file)),
+ )
+ else:
+ basenames = self.create_unique_basename(ext_file)
+ for name, subpath in ext_file.subpaths.items():
+ fetch(os.path.join(scratch_src, subpath), os.path.join(tmp_dir.abspath, basenames[name]))
else:
# copy directly to the bundle dir
src = ext_file.location
@@ -657,19 +713,36 @@ def fetch_file(ext_file, counter=[0]):
# transfer the result
self.transfer(tmp, outputs["bundle"])
- # unpack the bundle to have local files available
- with self.publish_step(f"unpacking to {outputs['local_files'].dir.abspath} ..."):
+ # remove all local files if recreating or if only existing partially to do a full refresh
+ local_files_exist = outputs["local_files"].exists()
+ if (self.recreate and local_files_exist) or not local_files_exist:
outputs["local_files"].dir.remove()
- bundle = outputs["bundle"]
- if isinstance(bundle, law.FileCollection):
- bundle = bundle.random_target()
- bundle.load(outputs["local_files"].dir, formatter="tar")
-
- # check if unpacked files/directories are described by the correct target class
- for target in outputs["local_files"]._flat_target_list:
- mismatch = (
- (isinstance(target, law.FileSystemFileTarget) and not os.path.isfile(target.abspath)) or
- (isinstance(target, law.FileSystemDirectoryTarget) and not os.path.isdir(target.abspath))
- )
- if mismatch:
- raise Exception(f"mismatching file/directory type of unpacked target {target!r}")
+ local_files_exist = False
+
+ # unpack the bundle to have local files available if needed
+ if not local_files_exist:
+ with self.publish_step(f"unpacking to {outputs['local_files'].dir.abspath} ..."):
+ bundle = outputs["bundle"]
+ if isinstance(bundle, law.FileCollection):
+ bundle = bundle.random_target()
+ bundle.load(outputs["local_files"].dir, formatter="tar")
+
+ # check if unpacked files/directories are described by the correct target class
+ for target in outputs["local_files"]._flat_target_list:
+ mismatch = (
+ (isinstance(target, law.FileSystemFileTarget) and not os.path.isfile(target.abspath)) or
+ (isinstance(target, law.FileSystemDirectoryTarget) and not os.path.isdir(target.abspath))
+ )
+ if mismatch:
+ raise Exception(f"mismatching file/directory type of unpacked target {target!r}")
+
+
+BundleExternalFilesWrapper = wrapper_factory(
+ base_cls=AnalysisTask,
+ require_cls=BundleExternalFiles,
+ enable=["configs", "skip_configs"],
+ attributes={"version": None},
+ docs="""
+Wrapper task trigger the BundleExternalFiles task for multiple configs.
+""",
+)
diff --git a/columnflow/tasks/framework/base.py b/columnflow/tasks/framework/base.py
index 4315af4e4..74e609426 100644
--- a/columnflow/tasks/framework/base.py
+++ b/columnflow/tasks/framework/base.py
@@ -8,6 +8,7 @@
import os
import abc
+import math
import enum
import importlib
import itertools
@@ -197,7 +198,10 @@ def req_params(cls, inst: AnalysisTask, **kwargs) -> dict[str, Any]:
not law.parser.global_cmdline_values().get(f"{cls.task_family}_version") and
(
cls.task_family != inst.task_family or
- freeze(cls.get_config_lookup_keys(params)) != freeze(inst.get_config_lookup_keys(inst))
+ (
+ freeze(cls.get_config_lookup_keys(params, significant=True)) !=
+ freeze(inst.get_config_lookup_keys(inst, significant=True))
+ )
)
):
default_version = cls.get_default_version(inst, params)
@@ -352,12 +356,14 @@ def _get_default_version(
def get_config_lookup_keys(
cls,
inst_or_params: AnalysisTask | dict[str, Any],
+ significant: bool = False,
) -> law.util.InsertiableDict:
"""
Returns a dictionary with keys that can be used to lookup state specific values in a config or dictionary, such
as default task versions or output locations.
:param inst_or_params: The tasks instance or its parameters.
+ :param significant: Whether only significant keys should be returned.
:return: A dictionary with keys that can be used for nested lookup.
"""
keys = law.util.InsertableDict()
@@ -556,7 +562,8 @@ def has_obj(name: str) -> bool:
object_names.append(name)
elif groups_str and name in (object_groups := container.x(groups_str, {})):
# a key in the object group dict
- lookup.extend(list(object_groups[name]))
+ for entry in list(object_groups[name]):
+ lookup.extend(law.util.brace_expand(entry))
elif accept_patterns:
# must eventually be a pattern, perform an object traversal
found = []
@@ -582,6 +589,7 @@ def resolve_config_default(
container: str | od.AuxDataMixin | Sequence[od.AuxDataMixin],
default_str: str | None = None,
multi_strategy: str = "first",
+ debug: bool = False,
) -> Any | list[Any] | dict[od.AuxDataMixin, Any]:
"""
Resolves a given parameter value *param*, checks if it should be placed with a default value when empty, and in
@@ -702,11 +710,12 @@ def default_selector(task_cls, config_inst, task_params) -> str:
return params
if multi_strategy == "first":
return params[container[0]]
- # NOTE: in there two strategies, we loose all order information
- if multi_strategy == "union":
- return list(set.union(*map(set, params.values())))
- if multi_strategy == "intersection":
- return list(set.intersection(*map(set, params.values())))
+ if multi_strategy in {"union", "intersection"}:
+ union = law.util.make_unique(sum(map(list, params.values()), []))
+ if multi_strategy == "union":
+ return union
+ # for intersection, use ordered union as index for sorting
+ return sorted(set.intersection(*map(set, params.values())), key=union.index)
# "same", so check that values are identical
first = params[container[0]]
if not all(params[c] == first for c in container[1:]):
@@ -724,7 +733,7 @@ def resolve_config_default_and_groups(
groups_str: str,
default_str: str | None = None,
multi_strategy: str = "first",
- debug=False,
+ debug: bool = False,
) -> Any | list[Any] | dict[od.AuxDataMixin, Any]:
"""
This method is similar to :py:meth:`~.resolve_config_default` in that it checks if a parameter value *param* is
@@ -819,7 +828,7 @@ def resolve_config_default_and_groups(
raise Exception(
f"definition of '{groups_str}' contains circular references involving group '{value}'",
)
- lookup.extendleft(law.util.make_list(param_groups[value]))
+ lookup.extendleft(law.util.make_list(param_groups[value])[::-1])
handled_groups.add(value)
else:
_values.append(value)
@@ -832,10 +841,12 @@ def resolve_config_default_and_groups(
return values
if multi_strategy == "first":
return values[container[0]]
- if multi_strategy == "union":
- return list(set.union(*map(set, values.values())))
- if multi_strategy == "intersection":
- return list(set.intersection(*map(set, values.values())))
+ if multi_strategy in {"union", "intersection"}:
+ union = law.util.make_unique(sum(map(list, values.values()), []))
+ if multi_strategy == "union":
+ return union
+ # for intersection, use ordered union as index for sorting
+ return sorted(set.intersection(*map(set, values.values())), key=union.index)
# "same", so check that values are identical
first = values[container[0]]
if not all(values[c] == first for c in container[1:]):
@@ -883,7 +894,7 @@ def build_repr(
not (0 < max_len < (len(r) + sum(map(len, objects[:max_count])) + len(sep) * max_count + hash_len))
):
r += sep.join(objects[:max_count])
- r += f"{sep}{law.util.create_hash(objects[max_count:], l=hash_len)}"
+ r += f"{sep}{law.util.create_hash(objects[max_count:], hash_len)}"
else:
r += sep.join(objects)
else:
@@ -891,7 +902,7 @@ def build_repr(
# handle overall truncation
if max_len > 0 and len(r) > max_len:
- r = f"{r[:max_len - hash_len - len(sep)]}{sep}{law.util.create_hash(r, l=hash_len)}"
+ r = f"{r[:max_len - hash_len - len(sep)]}{sep}{law.util.create_hash(r, hash_len)}"
return r
@@ -1132,6 +1143,20 @@ def target(self, *path, **kwargs) -> law.LocalTarget | law.wlcg.WLCGTarget | law
raise Exception(f"cannot determine output location based on '{location}'")
+ def pilot_workflow_requires(self, task: law.Task) -> Any:
+ """
+ Helper for situtations where *this* task is a workflow with ``--pilot`` activated to decide if an upstream task
+ itself should be required, or its own upstream dependencies.
+
+ :param task: The task to be required.
+ :return: Either the task itself or the result of its :py:meth:`~.workflow_requires` method.
+ """
+ return (
+ task.workflow_requires()
+ if self.pilot and isinstance(task, law.BaseWorkflow) and task.is_workflow()
+ else task
+ )
+
def get_parquet_writer_opts(self, repeating_values: bool = False) -> dict[str, Any]:
"""
Returns an option dictionary that can be passed as *writer_opts* to :py:meth:`~law.pyarrow.merge_parquet_task`,
@@ -1492,8 +1517,9 @@ def _get_default_version(
def get_config_lookup_keys(
cls,
inst_or_params: ConfigTask | dict[str, Any],
+ significant: bool = False,
) -> law.util.InsertiableDict:
- keys = super().get_config_lookup_keys(inst_or_params)
+ keys = super().get_config_lookup_keys(inst_or_params, significant=significant)
# add the config name in front of the task family
config = (
@@ -1673,8 +1699,9 @@ def resolve_shifts(cls, params: dict[str, Any]) -> dict[str, Any]:
def get_config_lookup_keys(
cls,
inst_or_params: ShiftTask | dict[str, Any],
+ significant: bool = False,
) -> law.util.InsertiableDict:
- keys = super().get_config_lookup_keys(inst_or_params)
+ keys = super().get_config_lookup_keys(inst_or_params, significant=significant)
# add the (global) shift name
shift = (
@@ -1780,8 +1807,9 @@ def get_known_shifts(
def get_config_lookup_keys(
cls,
inst_or_params: DatasetTask | dict[str, Any],
+ significant: bool = False,
) -> law.util.InsertiableDict:
- keys = super().get_config_lookup_keys(inst_or_params)
+ keys = super().get_config_lookup_keys(inst_or_params, significant=significant)
# add the dataset name before the shift name
dataset = (
@@ -1838,19 +1866,30 @@ def file_merging_factor(self) -> int:
Consecutive merging steps are not handled yet.
"""
n_files = self.dataset_info_inst.n_files
+ file_merging = self.file_merging
- if isinstance(self.file_merging, int):
+ if isinstance(file_merging, int):
# interpret the file_merging attribute as the merging factor itself
# zero means "merge all in one"
- if self.file_merging < 0:
- raise ValueError(f"invalid file_merging value {self.file_merging}")
- n_merge = n_files if self.file_merging == 0 else self.file_merging
+ if file_merging < 0:
+ raise ValueError(f"invalid file_merging value {file_merging}")
+ n_merge = n_files if file_merging == 0 else file_merging
else:
# no merging at all
n_merge = 1
return n_merge
+ @property
+ def n_merged_files(self) -> int:
+ """
+ Returns the number of files that are expected after merging, making use of :py:attr:`file_merging_factor` which
+ can depend on dynamic, dataset-dependent information.
+ """
+ n_files = self.dataset_info_inst.n_files
+ n_merge = self.file_merging_factor
+ return int(math.ceil((1.0 * n_files / n_merge)))
+
def create_branch_map(self):
"""
Define the branch map for when this task is used as a workflow. By default, use the merging
@@ -1858,8 +1897,8 @@ def create_branch_map(self):
branches to one or more input file indices. E.g. `1 -> [3, 4, 5]` would mean that branch 1
is simultaneously handling input file indices 3, 4 and 5.
"""
- n_merge = self.file_merging_factor
n_files = self.dataset_info_inst.n_files
+ n_merge = self.file_merging_factor
# use iter_chunks which splits a list of length n_files into chunks of maximum size n_merge
chunks = law.util.iter_chunks(n_files, n_merge)
@@ -1999,6 +2038,7 @@ def wrapper_factory(
cls_name: str | None = None,
attributes: dict | None = None,
docs: str | None = None,
+ port_parameters: bool | Sequence[str] = True,
) -> law.task.base.Register:
"""
Factory function creating wrapper task classes, inheriting from *base_cls* and
@@ -2050,6 +2090,8 @@ class MyTask(DatasetTask):
class
:param docs: Manually set the documentation string `__doc__` of the new :py:class:`~law.task.base.WrapperTask` class
instance
+ :param port_parameters: Whether to port the parameters of the `require_cls`. When a sequence of strings is passed,
+ only parameters with these names are ported.
:raises ValueError: If a parameter provided with `enable` is not in the list of known parameters
:raises TypeError: If any parameter in `enable` is incompatible with the :py:class:`~law.task.base.WrapperTask`
class instance or the inheritance structure of corresponding classes
@@ -2315,8 +2357,42 @@ def update_wrapper_params(self, params):
# overwrite __name__
Wrapper.__name__ = cls_name or f"{require_cls.__name__}Wrapper"
+ # use same task family
+ Wrapper.task_namespace = require_cls.task_namespace
+
# set docs
if docs:
Wrapper.__docs__ = docs
+ # port parameters from require_cls
+ if port_parameters:
+ # define which parameters to port
+ upstream_params = dict(require_cls.get_params())
+ if isinstance(port_parameters, bool):
+ port_params = (
+ # start from all non-private upstream parameters
+ set(
+ name for name, param in upstream_params.items()
+ if param.visibility != luigi.parameter.ParameterVisibility.PRIVATE
+ ) -
+ # skip existing parameters
+ set(dict(Wrapper.get_params())) -
+ # skip interactive parameters
+ set(require_cls.interactive_params) -
+ # skip with some heuristics
+ {"config", "dataset", "shift", "effective_workflow", "local_shift", "known_shifts"}
+ )
+ else:
+ # take sequence as is, but check for existence
+ port_params = law.util.make_unique(port_parameters)
+ for name in port_params:
+ if name not in upstream_params:
+ raise ValueError(
+ f"cannot port parameter '{name}' to '{Wrapper.__name__}': not existing in "
+ f"'{require_cls.__name__}'",
+ )
+ # actual porting
+ for name in port_params:
+ setattr(Wrapper, name, upstream_params[name])
+
return Wrapper
diff --git a/columnflow/tasks/framework/histograms.py b/columnflow/tasks/framework/histograms.py
index 81a8682b1..4ce9343b3 100644
--- a/columnflow/tasks/framework/histograms.py
+++ b/columnflow/tasks/framework/histograms.py
@@ -15,6 +15,7 @@
HistProducerClassMixin, VariablesMixin, DatasetsProcessesMixin, CategoriesMixin, ShiftSourcesMixin,
)
from columnflow.tasks.histograms import MergeHistograms, MergeShiftedHistograms
+from columnflow.hist_util import select_category_bins
from columnflow.util import dev_sandbox, maybe_import
from columnflow.types import TYPE_CHECKING
@@ -107,59 +108,29 @@ def slice_histogram(
def flatten_nested_list(nested_list):
return [item for sublist in nested_list for item in sublist]
- selection_dict = {}
+ # work on a copy
+ h = histogram.copy()
+ # select processes
if processes:
- # transform into lists if necessary
- processes = law.util.make_list(processes)
- # get all sub processes
-
- process_insts = list(map(config_inst.get_process, processes))
+ process_insts = list(map(config_inst.get_process, law.util.make_unique(law.util.make_list(processes))))
sub_process_insts = set(flatten_nested_list([
[sub for sub, _, _ in proc.walk_processes(include_self=True)]
for proc in process_insts
]))
- selection_dict["process"] = [
- hist.loc(p.name)
- for p in sub_process_insts
- if p.name in histogram.axes["process"]
- ]
+ h = h[{"process": [hist.loc(p.name) for p in sub_process_insts if p.name in histogram.axes["process"]]}]
+
+ # select categories
if categories:
- # transform into lists if necessary
- categories = law.util.make_list(categories)
-
- # get all leaf categories
- category_insts = list(map(config_inst.get_category, categories))
- leaf_category_insts = set(flatten_nested_list([
- category_inst.get_leaf_categories() or [category_inst]
- for category_inst in category_insts
- ]))
- selection_dict["category"] = [
- hist.loc(c.name)
- for c in leaf_category_insts
- if c.name in histogram.axes["category"]
- ]
+ h = select_category_bins(h, categories, use_leaves=True, prefer_parents=True, reduce=False)
+ # select shifts
if shifts:
- # transform into lists if necessary
- shifts = law.util.make_list(shifts)
-
- # get all shift instances
- shift_insts = [config_inst.get_shift(shift) for shift in shifts]
- selection_dict["shift"] = [
- hist.loc(s.name)
- for s in shift_insts
- if s.name in histogram.axes["shift"]
- ]
-
- # work on a copy
- h = histogram.copy()
-
- # axis selections
- h = h[selection_dict]
+ shift_insts = [config_inst.get_shift(shift) for shift in law.util.make_unique(law.util.make_list(shifts))]
+ h = h[{"shift": [hist.loc(s.name) for s in shift_insts if s.name in histogram.axes["shift"]]}]
+ # optional axis reductions
if reduce_axes:
- # axis reductions
h = h[{"process": sum, "category": sum, "shift": sum}]
return h
diff --git a/columnflow/tasks/framework/inference.py b/columnflow/tasks/framework/inference.py
index 508b67e17..346997856 100644
--- a/columnflow/tasks/framework/inference.py
+++ b/columnflow/tasks/framework/inference.py
@@ -6,8 +6,6 @@
from __future__ import annotations
-import pickle
-
import law
import order as od
@@ -116,6 +114,8 @@ def combined_config_data(self) -> dict[od.ConfigInst, dict[str, dict | set]]:
config_inst: {
# all variables used in this config in any datacard category
"variables": set(),
+ # set of all used categories
+ "categories": set(),
# plain set of names of real data datasets
"data_datasets": set(),
# per mc dataset name, the set of shift sources and the names processes to be extracted from them
@@ -126,7 +126,7 @@ def combined_config_data(self) -> dict[od.ConfigInst, dict[str, dict | set]]:
# iterate over all model categories
for cat_obj in self.inference_model_inst.categories:
- # keep track of per-category information for consistency checks
+ # keep track of per-category information across configs for consistency checks
variables = set()
categories = set()
@@ -135,8 +135,9 @@ def combined_config_data(self) -> dict[od.ConfigInst, dict[str, dict | set]]:
for config_inst in config_insts:
data = config_data[config_inst]
- # variables
+ # variables and categories
data["variables"].add(cat_obj.config_data[config_inst.name].variable)
+ data["categories"].add(cat_obj.config_data[config_inst.name].category)
# data datasets, but only if
# - data in that category is not faked from mc processes, or
@@ -193,29 +194,38 @@ def create_branch_map(self):
# dummy branch map
return {0: None}
+ def _hist_requirement(self, **kwargs):
+ return self.reqs.MergeShiftedHistograms.req_different_branching(self, **kwargs)
+
def _hist_requirements(self, **kwargs):
# gather data from inference model to define requirements in the structure
# config_name -> dataset_name -> MergeHistogramsTask
reqs = {}
for config_inst, data in self.combined_config_data.items():
reqs[config_inst.name] = {}
+ # ensure that all variables exist
+ for var_name in data["variables"]:
+ if not config_inst.has_variable(var_name):
+ raise ValueError(
+ f"config '{config_inst.name}' does not have variable '{var_name}' required by inference model "
+ f"'{self.inference_model}'",
+ )
# mc datasets
for dataset_name in sorted(data["mc_datasets"]):
- reqs[config_inst.name][dataset_name] = self.reqs.MergeShiftedHistograms.req_different_branching(
- self,
+ reqs[config_inst.name][dataset_name] = self._hist_requirement(
config=config_inst.name,
dataset=dataset_name,
- shift_sources=tuple(sorted(data["mc_datasets"][dataset_name]["shift_sources"])),
+ shift_sources=("nominal",) + tuple(sorted(data["mc_datasets"][dataset_name]["shift_sources"])),
variables=tuple(sorted(data["variables"])),
**kwargs,
)
- # data datasets
+
+ # data datasets, no shift sources so not chunked
for dataset_name in sorted(data["data_datasets"]):
- reqs[config_inst.name][dataset_name] = self.reqs.MergeShiftedHistograms.req_different_branching(
- self,
+ reqs[config_inst.name][dataset_name] = self._hist_requirement(
config=config_inst.name,
dataset=dataset_name,
- shift_sources=(),
+ shift_sources=("nominal",),
variables=tuple(sorted(data["variables"])),
**kwargs,
)
@@ -244,47 +254,52 @@ def load_process_hists(
with self.publish_step(f"extracting '{variable}' for config {config_inst.name} ..."):
for dataset_name, process_names in dataset_processes.items():
- # open the histogram and work on a copy
- inp = inputs[dataset_name]["collection"][0]["hists"][variable]
- try:
- h = inp.load(formatter="pickle").copy()
- except pickle.UnpicklingError as e:
- raise Exception(
- f"failed to load '{variable}' histogram for dataset '{dataset_name}' in config "
- f"'{config_inst.name}' from {inp.abspath}",
- ) from e
-
- # determine processes to extract
- process_insts = [config_inst.get_process(name) for name in process_names]
-
- # loop over all proceses assigned to this dataset
- for process_inst in process_insts:
- # gather all subprocesses for a full query later
- sub_process_insts = [sub for sub, _, _ in process_inst.walk_processes(include_self=True)]
-
- # there must be at least one matching sub process
- if not any(p.name in h.axes["process"] for p in sub_process_insts):
- raise Exception(f"no '{variable}' histograms found for process '{process_inst.name}'")
-
- # select and reduce over relevant processes
- h_proc = h[{
- "process": [hist.loc(p.name) for p in sub_process_insts if p.name in h.axes["process"]],
- }]
- h_proc = h_proc[{"process": sum}]
-
- # additional custom reductions
- h_proc = self.modify_process_hist(
- config_inst=config_inst,
- process_inst=process_inst,
- variable=variable,
- h=h_proc,
- )
+ # loop through inputs
+ for inp in inputs[dataset_name]["collection"].targets.values():
+ # open the histogram
+ try:
+ h = inp["hists"][variable].load(formatter="pickle")
+ except Exception as e:
+ raise Exception(
+ f"failed to load '{variable}' histogram for dataset '{dataset_name}' in config "
+ f"'{config_inst.name}' from {inp.abspath}",
+ ) from e
+
+ # determine processes to extract
+ process_insts = [config_inst.get_process(name) for name in process_names]
+
+ # loop over all proceses assigned to this dataset
+ for process_inst in process_insts:
+ # gather all subprocesses for a full query later
+ sub_process_insts = [sub for sub, _, _ in process_inst.walk_processes(include_self=True)]
+
+ # only perform subprocess selection if there are any on the process axis
+ if len(h.axes["process"]):
+ # then, there must be at least one matching sub process
+ if not any(p.name in h.axes["process"] for p in sub_process_insts):
+ raise Exception(f"no '{variable}' histograms found for process '{process_inst.name}'")
+
+ # select and reduce over relevant processes
+ h_proc = h[{
+ "process": [hist.loc(p.name) for p in sub_process_insts if p.name in h.axes["process"]],
+ }]
+ h_proc = h_proc[{"process": sum}]
+ else:
+ h_proc = h[{"process": sum}]
+
+ # additional custom reductions
+ h_proc = self.modify_process_hist(
+ config_inst=config_inst,
+ process_inst=process_inst,
+ variable=variable,
+ h=h_proc,
+ )
- # store it
- if process_inst in hists:
- hists[process_inst] += h_proc
- else:
- hists[process_inst] = h_proc
+ # store it
+ if process_inst in hists:
+ hists[process_inst] += h_proc
+ else:
+ hists[process_inst] = h_proc
return hists
@@ -300,6 +315,7 @@ def modify_process_hist(
:param config_inst: The config instance the histogram belongs to.
:param process_inst: The process instance the histogram belongs to.
+ :param variable: The variable name the histogram corresponds to.
:param h: The histogram to modify.
:return: The modified histogram.
"""
diff --git a/columnflow/tasks/framework/mixins.py b/columnflow/tasks/framework/mixins.py
index 36b99d5ad..290e7311c 100644
--- a/columnflow/tasks/framework/mixins.py
+++ b/columnflow/tasks/framework/mixins.py
@@ -25,9 +25,9 @@
from columnflow.ml import MLModel
from columnflow.inference import InferenceModel
from columnflow.columnar_util import Route, ColumnCollection, ChunkedIOHandler, TaskArrayFunction
+from columnflow.config_util import expand_shift_sources
from columnflow.util import maybe_import, DotDict, get_docs_url, get_code_url
from columnflow.types import Callable
-from columnflow.timing import Timer
np = maybe_import("numpy")
ak = maybe_import("awkward")
@@ -97,8 +97,9 @@ def req_params(cls, inst: law.Task, **kwargs) -> dict[str, Any]:
def get_config_lookup_keys(
cls,
inst_or_params: CalibratorClassMixin | dict[str, Any],
+ significant: bool = False,
) -> law.util.InsertiableDict:
- keys = super().get_config_lookup_keys(inst_or_params)
+ keys = super().get_config_lookup_keys(inst_or_params, significant=significant)
# add the calibrator name
calibrator = (
@@ -314,8 +315,9 @@ def store_parts(self) -> law.util.InsertableDict:
def get_config_lookup_keys(
cls,
inst_or_params: CalibratorClassesMixin | dict[str, Any],
+ significant: bool = False,
) -> law.util.InsertiableDict:
- keys = super().get_config_lookup_keys(inst_or_params)
+ keys = super().get_config_lookup_keys(inst_or_params, significant=significant)
# add the calibrator names
calibrators = (
@@ -500,8 +502,9 @@ def req_params(cls, inst: law.Task, **kwargs) -> dict[str, Any]:
def get_config_lookup_keys(
cls,
inst_or_params: SelectorClassMixin | dict[str, Any],
+ significant: bool = False,
) -> law.util.InsertiableDict:
- keys = super().get_config_lookup_keys(inst_or_params)
+ keys = super().get_config_lookup_keys(inst_or_params, significant=significant)
# add the selector name
selector = (
@@ -716,8 +719,9 @@ def req_params(cls, inst: law.Task, **kwargs) -> dict[str, Any]:
def get_config_lookup_keys(
cls,
inst_or_params: ReducerClassMixin | dict[str, Any],
+ significant: bool = False,
) -> law.util.InsertiableDict:
- keys = super().get_config_lookup_keys(inst_or_params)
+ keys = super().get_config_lookup_keys(inst_or_params, significant=significant)
# add the reducer name
reducer = (
@@ -767,7 +771,13 @@ class ReducerMixin(ArrayFunctionInstanceMixin, ReducerClassMixin):
@classmethod
def get_reducer_dict(cls, params: dict[str, Any]) -> dict[str, Any]:
- return cls.get_array_function_dict(params)
+ d = cls.get_array_function_dict(params)
+
+ # special case: add selector shifts
+ if (selector_inst := params.get("selector_inst")):
+ d["selector_shifts"] = selector_inst.all_shifts
+
+ return d
@classmethod
def build_reducer_inst(
@@ -909,8 +919,9 @@ def req_params(cls, inst: law.Task, **kwargs) -> dict[str, Any]:
def get_config_lookup_keys(
cls,
inst_or_params: ProducerClassMixin | dict[str, Any],
+ significant: bool = False,
) -> law.util.InsertiableDict:
- keys = super().get_config_lookup_keys(inst_or_params)
+ keys = super().get_config_lookup_keys(inst_or_params, significant=significant)
# add the producer name
producer = (
@@ -1126,8 +1137,9 @@ def store_parts(self) -> law.util.InsertableDict:
def get_config_lookup_keys(
cls,
inst_or_params: ProducerClassesMixin | dict[str, Any],
+ significant: bool = False,
) -> law.util.InsertiableDict:
- keys = super().get_config_lookup_keys(inst_or_params)
+ keys = super().get_config_lookup_keys(inst_or_params, significant=significant)
# add the producer names
producers = (
@@ -1344,8 +1356,9 @@ def events_used_in_training(
def get_config_lookup_keys(
cls,
inst_or_params: MLModelMixinBase | dict[str, Any],
+ significant: bool = False,
) -> law.util.InsertiableDict:
- keys = super().get_config_lookup_keys(inst_or_params)
+ keys = super().get_config_lookup_keys(inst_or_params, significant=significant)
# add the ml model name
ml_model = (
@@ -1689,8 +1702,9 @@ def find_keep_columns(self, collection: ColumnCollection) -> set[Route]:
def get_config_lookup_keys(
cls,
inst_or_params: MLModelsMixin | dict[str, Any],
+ significant: bool = False,
) -> law.util.InsertiableDict:
- keys = super().get_config_lookup_keys(inst_or_params)
+ keys = super().get_config_lookup_keys(inst_or_params, significant=significant)
# add the ml model names
ml_models = (
@@ -1754,8 +1768,9 @@ def req_params(cls, inst: law.Task, **kwargs) -> dict[str, Any]:
def get_config_lookup_keys(
cls,
inst_or_params: HistProducerClassMixin | dict[str, Any],
+ significant: bool = False,
) -> law.util.InsertiableDict:
- keys = super().get_config_lookup_keys(inst_or_params)
+ keys = super().get_config_lookup_keys(inst_or_params, significant=significant)
# add the hist producer name
producer = (
@@ -2036,6 +2051,7 @@ class CategoriesMixin(ConfigTask):
default_categories = None
allow_empty_categories = False
+ sort_categories = True
@classmethod
def resolve_param_values_post_init(cls, params: dict[str, Any]) -> dict[str, Any]:
@@ -2073,15 +2089,26 @@ def resolve_param_values_post_init(cls, params: dict[str, Any]) -> dict[str, Any
if not categories and not cls.allow_empty_categories:
raise ValueError(f"no categories found matching {params['categories']}")
+ # sort them
+ if cls.sort_categories:
+ categories = sorted(categories)
+
params["categories"] = tuple(categories)
return params
+ @classmethod
+ def _categories_repr(cls, categories: Sequence[str]) -> str:
+ # single category representation
+ if len(categories) == 1:
+ return cls.build_repr(categories[0])
+
+ # full representation
+ return cls.build_repr(categories, prepend_count=True)
+
@property
def categories_repr(self) -> str:
- if len(self.categories) == 1:
- return self.build_repr(self.categories[0])
- return self.build_repr(self.categories, prepend_count=True)
+ return self._categories_repr(self.categories)
class VariablesMixin(ConfigTask):
@@ -2124,7 +2151,7 @@ def resolve_param_values_post_init(cls, params: dict[str, Any]) -> dict[str, Any
multi_strategy="union",
)
# since there can be multi-dimensional variables, resolve each part separately
- resolved_variables = set()
+ resolved_variables = []
for variable in variables:
resolved_parts = [
cls.find_config_objects(
@@ -2137,8 +2164,8 @@ def resolve_param_values_post_init(cls, params: dict[str, Any]) -> dict[str, Any
for part in cls.split_multi_variable(variable)
]
# build combinatrics
- resolved_variables.update(map(cls.join_multi_variable, itertools.product(*resolved_parts)))
- variables = resolved_variables
+ resolved_variables.extend(map(cls.join_multi_variable, itertools.product(*resolved_parts)))
+ variables = law.util.make_unique(resolved_variables)
# when still empty, fallback to using all known variables
if not variables:
@@ -2168,6 +2195,15 @@ def join_multi_variable(cls, variables: Sequence[str]) -> str:
"""
return "-".join(map(str, variables))
+ @classmethod
+ def _variables_repr(cls, variables: Sequence[str]) -> str:
+ # simplified representation for single source
+ if len(variables) == 1:
+ return cls.build_repr(variables[0])
+
+ # full representation
+ return cls.build_repr(sorted(variables), prepend_count=True)
+
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
@@ -2179,9 +2215,7 @@ def __init__(self, *args, **kwargs) -> None:
@property
def variables_repr(self) -> str:
- if len(self.variables) == 1:
- return self.build_repr(self.variables[0])
- return self.build_repr(sorted(self.variables), prepend_count=True)
+ return self._variables_repr(self.variables)
class DatasetsMixin(ConfigTask):
@@ -2209,7 +2243,6 @@ class DatasetsMixin(ConfigTask):
@classmethod
def modify_task_attributes(cls) -> None:
super().modify_task_attributes()
- # single/multi config adjustments in case the switch has been specified
if isinstance(cls.single_config, bool) and getattr(cls, "datasets_multi", None) is not None:
if not cls.has_single_config():
cls.datasets = cls.datasets_multi
@@ -2218,9 +2251,6 @@ def modify_task_attributes(cls) -> None:
@classmethod
def resolve_datasets(cls, config_inst: od.Config, datasets: Any) -> tuple[list[str], list[str]]:
- """
- helper to resolve processes and datasets for one config
- """
if datasets:
datasets_orig = datasets
datasets = cls.find_config_objects(
@@ -2231,29 +2261,22 @@ def resolve_datasets(cls, config_inst: od.Config, datasets: Any) -> tuple[list[s
)
if not datasets and not cls.allow_empty_datasets:
raise ValueError(f"no datasets found matching {datasets_orig}")
-
return datasets
@classmethod
def resolve_param_values_pre_init(cls, params: dict[str, Any]) -> dict[str, Any]:
params = super().resolve_param_values_pre_init(params)
- # get processes and datasets
single_config = cls.has_single_config()
datasets = (params.get("datasets", law.no_value),) if single_config else params.get("datasets", ())
-
- # "broadcast" to match number of configs
config_insts = params.get("config_insts")
- # perform resolution per config
multi_datasets = []
for config_inst, _datasets in zip(config_insts, datasets):
_datasets = cls.resolve_datasets(config_inst, _datasets)
multi_datasets.append(tuple(_datasets) if _datasets != law.no_value else None)
params["datasets"] = multi_datasets[0] if single_config else tuple(multi_datasets)
-
- # store instances
params["dataset_insts"] = {
config_inst: datasets and list(map(config_inst.get_dataset, datasets))
for config_inst, datasets in zip(config_insts, multi_datasets)
@@ -2267,7 +2290,6 @@ def resolve_instances(cls, params: dict[str, Any], shifts: TaskShifts) -> dict[s
cls.get_known_shifts(params, shifts)
- # we loop over all configs/datasets, but return initial params
for i, config_inst in enumerate(params["config_insts"]):
if cls.has_single_config():
datasets = params["datasets"]
@@ -2275,7 +2297,6 @@ def resolve_instances(cls, params: dict[str, Any], shifts: TaskShifts) -> dict[s
datasets = params["datasets"][i]
for dataset in datasets:
- # NOTE: we need to copy here, because otherwise taf inits will only be triggered once
_params = {
**params,
"config_inst": config_inst,
@@ -2287,27 +2308,18 @@ def resolve_instances(cls, params: dict[str, Any], shifts: TaskShifts) -> dict[s
cls.resolution_task_cls.get_known_shifts(_params, shifts)
params["known_shifts"] = shifts
-
return params
@classmethod
def get_known_shifts(
- cls,
- params: dict[str, Any],
- shifts: TaskShifts,
+ cls,
+ params: dict[str, Any],
+ shifts: TaskShifts,
) -> None:
- """
- Updates the set of known *shifts* implemented by *this* and upstream tasks.
-
- :param params: Dictionary of task parameters.
- :param shifts: TaskShifts object to adjust.
- """
- # add shifts of all datasets to upstream ones
for config_inst, dataset_insts in params["dataset_insts"].items():
for dataset_inst in dataset_insts:
if dataset_inst.is_mc:
shifts.upstream |= set(dataset_inst.info.keys())
-
super().get_known_shifts(params, shifts)
@property
@@ -2372,8 +2384,8 @@ def resolve_param_values_pre_init(cls, params: dict[str, Any]) -> dict[str, Any]
# helper to resolve processes and datasets for one config
def resolve(config_inst: od.Config, processes: Any, datasets: Any) -> tuple[list[str], list[str]]:
+ processes_orig = processes
if processes != law.no_value:
- processes_orig = processes
if processes:
processes = cls.find_config_objects(
names=processes,
@@ -2412,6 +2424,21 @@ def resolve(config_inst: od.Config, processes: Any, datasets: Any) -> tuple[list
object_cls=od.Dataset,
groups_str="dataset_groups",
)
+ # reduce processes to those present in selected datasets when none were given initially
+ if datasets and processes and processes_orig in {law.no_value, ()}:
+ def use_process_for_dataset(process_inst: od.Process, dataset_inst: od.Dataset) -> bool:
+ return (
+ dataset_inst.has_process(process_inst) or
+ process_inst.has_process(dataset_inst.processes.get_first())
+ )
+ dataset_insts = list(map(config_inst.get_dataset, datasets))
+ processes = tuple(
+ process for process in processes
+ if any(
+ use_process_for_dataset(config_inst.get_process(process), dataset_inst)
+ for dataset_inst in dataset_insts
+ )
+ )
elif processes and processes != law.no_value:
# pick all datasets that contain any of the requested (sub)processes
sub_process_insts = sum((
@@ -2524,12 +2551,30 @@ class ShiftSourcesMixin(ConfigTask):
shift_sources = law.CSVParameter(
default=(),
description="comma-separated shift source names (without direction) or patterns to select; can also be the key "
- "of a mapping defined in the 'shift_group' auxiliary data of the config; default: ()",
+ "of a mapping defined in the 'shift_group' auxiliary data of the config; empty default",
brace_expand=True,
parse_empty=True,
)
allow_empty_shift_sources = False
+ sort_shift_sources = True
+ enforce_nominal_shift_source = False
+ remove_nominal_shift_source = False
+
+ @classmethod
+ def modify_param_values(cls, params: dict[str, Any]) -> dict[str, Any]:
+ params = super().modify_param_values(params)
+
+ # enforce/remove nominal shift source
+ if params.get("shift_sources") is not None:
+ if cls.enforce_nominal_shift_source and "nominal" not in params["shift_sources"]:
+ params["shift_sources"] = ("nominal",) + tuple(params["shift_sources"])
+ elif cls.remove_nominal_shift_source and "nominal" in params["shift_sources"]:
+ params["shift_sources"] = tuple(
+ source for source in params["shift_sources"] if source != "nominal"
+ )
+
+ return params
@classmethod
def resolve_param_values_post_init(cls, params: dict[str, Any]) -> dict[str, Any]:
@@ -2538,7 +2583,7 @@ def resolve_param_values_post_init(cls, params: dict[str, Any]) -> dict[str, Any
# resolve shift sources
if (container := cls._get_config_container(params)) and "shift_sources" in params:
shifts = cls.find_config_objects(
- names=cls.expand_shift_sources(params["shift_sources"]),
+ names=expand_shift_sources(params["shift_sources"]),
container=container,
object_cls=od.Shift,
groups_str="shift_groups",
@@ -2550,12 +2595,12 @@ def resolve_param_values_post_init(cls, params: dict[str, Any]) -> dict[str, Any
if shifts:
sources = cls.reduce_shifts(shifts)
- # # reduce shifts based on known shifts
+ # reduce shifts based on known shifts
if "known_shifts" not in params:
raise ValueError("known_shifts must be set before resolving shift sources")
sources = [
source for source in sources
- if (
+ if source == "nominal" or (
f"{source}_up" in params["known_shifts"].upstream and
f"{source}_down" in params["known_shifts"].upstream
)
@@ -2565,15 +2610,15 @@ def resolve_param_values_post_init(cls, params: dict[str, Any]) -> dict[str, Any
if not sources and not cls.allow_empty_shift_sources:
raise ValueError(f"no shifts found matching {params['shift_sources']}")
+ # potentially sort them
+ if cls.sort_shift_sources:
+ sources = sorted(sources)
+
# store them
params["shift_sources"] = tuple(sources)
return params
- @classmethod
- def expand_shift_sources(cls, sources: Sequence[str] | set[str]) -> list[str]:
- return sum(([f"{s}_up", f"{s}_down"] for s in sources), [])
-
@classmethod
def reduce_shifts(cls, shifts: Sequence[str] | set[str]) -> list[str]:
return list(set(od.Shift.split_name(shift)[0] for shift in shifts))
@@ -2581,15 +2626,41 @@ def reduce_shifts(cls, shifts: Sequence[str] | set[str]) -> list[str]:
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
- self.shifts = self.expand_shift_sources(self.shift_sources)
+ self.shifts = expand_shift_sources(self.shift_sources)
+
+ @classmethod
+ def _shift_sources_repr(
+ cls,
+ shift_sources: tuple[str, ...] | None,
+ enforce_nominal_shift_source: bool = False,
+ ) -> str:
+ if not shift_sources:
+ return "none"
+
+ # when nominal is the only source, but it is enforced to be present, also show "none" as in "no additional" ones
+ if enforce_nominal_shift_source and tuple(shift_sources) == ("nominal",):
+ return "none"
+
+ # sort shift sources, moving nominal to front if present, but dropping it if enforced
+ sorted_sources = sorted(shift_sources)
+ if "nominal" in sorted_sources:
+ sorted_sources.remove("nominal")
+ if not enforce_nominal_shift_source:
+ sorted_sources.insert(0, "nominal")
+
+ # simplified representation for single source
+ if len(sorted_sources) == 1:
+ return cls.build_repr(sorted_sources[0])
+
+ # full representation
+ return cls.build_repr(sorted_sources, prepend_count=True)
@property
def shift_sources_repr(self) -> str:
- if not self.shift_sources:
- return "none"
- if len(self.shift_sources) == 1:
- return self.build_repr(self.shift_sources[0])
- return self.build_repr(sorted(self.shift_sources), prepend_count=True)
+ return self._shift_sources_repr(
+ self.shift_sources,
+ enforce_nominal_shift_source=self.enforce_nominal_shift_source,
+ )
def store_parts(self) -> law.util.InsertableDict:
parts = super().store_parts()
@@ -2604,7 +2675,7 @@ class DatasetShiftSourcesMixin(ShiftSourcesMixin, DatasetTask):
effective_shift = None
allow_empty_shift = True
- # allow empty sources, i.e., using only nominal
+ # allow empty sources
allow_empty_shift_sources = True
@@ -2873,17 +2944,14 @@ def output(self):
@law.decorator.log
def run(self):
- # preare inputs and outputs
inputs = self.input()["collection"]
outputs = self.output()
- # load input histograms
hists = [
inp["hists"].load(formatter="pickle")
for inp in self.iter_progress(inputs.targets.values(), len(inputs), reach=(0, 50))
]
- # create a separate file per output variable
variable_names = list(hists[0].keys())
for variable_name in self.iter_progress(variable_names, len(variable_names), reach=(50, 100)):
self.publish_message(f"merging histograms for '{variable_name}'")
@@ -2892,89 +2960,5 @@ def run(self):
merged = sum(variable_hists[1:], variable_hists[0].copy())
outputs["hists"][variable_name].dump(merged, formatter="pickle")
- # optionally remove inputs
if self.remove_previous:
inputs.remove()
-
-
-class ParamsCacheMixin:
-
- # the get_param_values is called again for every value of this parameter (config by default included)
- cache_param_sep = ["shift", "shift_sources"]
-
- # dict to store cached params for different tasks
- cache_param_values = dict()
-
- time_get_params = luigi.BoolParameter(
- default=False,
- description="Whether to report the time taken to load get the parameters",
- )
-
- no_cached_params = luigi.BoolParameter(
- default=False,
- description="Return normal call to get_param_values and report if it differs from cached output",
- )
-
- check_cached_params = luigi.BoolParameter(
- default=False,
- description="Return normal call to get_param_values and report if it differs from cached output",
- )
-
- @classmethod
- def cache_tag(cls, kwargs):
- tag = cls.__name__
- for p in ["config"] + list(cls.cache_param_sep):
- if kwargs.get(p, None):
- tag += f"__{p}_{kwargs[p]}"
- return tag
-
- @classmethod
- def get_param_values(cls, params, args, kwargs):
-
- tag = cls.cache_tag(kwargs)
- cache = cls.cache_param_values.setdefault(tag, [])
- if cache:
- # return cached params but overwrite branch(es), dataset, and output/status options
- dct_update = dict(branch=int(kwargs.get("branch", -1)))
- if "dataset" in kwargs:
- dct_update["dataset"] = kwargs["dataset"]
- cached_out = cache[0] | dct_update | {
- k: kwargs.get(k, ()) for k in
- ["print_output", "branches", "print_status", "remove_output", "fetch_output"]
- }
-
- # call the normal get_param_values first time, or by request
- if any([
- check_cached_params := kwargs.get("check_cached_params", False),
- no_cached_params := kwargs.get("no_cached_params", False),
- first_call := not cache,
- ]):
- if time_get_params := kwargs.get("time_get_params", False):
- tmr = Timer(f"{tag} get_param_values")
-
- out = super().get_param_values(params, args, kwargs)
-
- if time_get_params:
- tmr("end of call")
-
- dct_out = dict(out)
-
- # store first call to cache
- if first_call:
- cache.append(dct_out)
- cls.cache_param_values[cls.cache_tag(dct_out)] = cache
- # compare cached output to normal call
- elif check_cached_params:
- for key, value in dct_out.items():
- cached_val = cached_out.get(key, None)
- if isinstance(value, (tuple, list)) and isinstance(cached_val, (tuple, list)):
- value = tuple(value)
- cached_val = tuple(cached_val)
- if cached_val != value:
- logger.warning(
- "normal call to get_param_values yielded output that differs from cached output\n"
- f"Normal call -> {key} [key]: {value}\n"
- f"Cached call -> {key} [key]: {cached_val}",
- )
-
- return out if no_cached_params or first_call else list(cached_out.items())
diff --git a/columnflow/tasks/framework/plotting.py b/columnflow/tasks/framework/plotting.py
index a743220e1..e0146c64f 100644
--- a/columnflow/tasks/framework/plotting.py
+++ b/columnflow/tasks/framework/plotting.py
@@ -93,21 +93,21 @@ def resolve_param_values(cls, params):
# NOTE: we currently assume that general_settings defaults and groups are the same for all
# config instances
if "general_settings" in params:
- settings = params["general_settings"]
- # when empty and default general_settings are defined, use them instead
- if not settings and config_inst.x("default_general_settings", ()):
- settings = config_inst.x("default_general_settings", ())
+ settings = params["general_settings"] or {}
+ if settings:
+ groups = config_inst.x("general_settings_groups", {})
+ settings = groups.get(list(settings.keys())[0], settings)
if isinstance(settings, tuple):
settings = cls.general_settings.parse(settings)
- # when general_settings are a key to a general_settings_groups, use them instead
- groups = config_inst.x("general_settings_groups", {})
- if settings and list(settings.keys())[0] in groups.keys():
- settings = groups[list(settings.keys())[0]]
- if isinstance(settings, tuple):
- settings = cls.general_settings.parse(settings)
+ # add defaults
+ default_settings = config_inst.x("default_general_settings", ())
+ if isinstance(default_settings, tuple):
+ default_settings = cls.general_settings.parse(default_settings)
+ if default_settings:
+ settings = law.util.merge_dicts(default_settings, settings, deep=True)
- params["general_settings"] = settings
+ params["general_settings"] = DotDict.wrap(settings)
return params
@@ -275,7 +275,7 @@ def update_plot_kwargs(self, kwargs: dict) -> dict:
# update style_config
style_config = kwargs.get("style_config", {})
if isinstance(custom_style_config, dict) and isinstance(style_config, dict):
- style_config = law.util.merge_dicts(style_config, custom_style_config)
+ style_config = law.util.merge_dicts(style_config, custom_style_config, deep=True)
kwargs["style_config"] = style_config
# update other defaults
@@ -473,7 +473,10 @@ def get_plot_parameters(self) -> DotDict:
return params
-class VariableSettingMixin():
+class VariableSettingMixin(
+ PlotBase,
+ VariablesMixin,
+):
"""
Mixin class for applying variable settings
"""
diff --git a/columnflow/tasks/framework/remote.py b/columnflow/tasks/framework/remote.py
index da2c59c2c..7c50d5746 100644
--- a/columnflow/tasks/framework/remote.py
+++ b/columnflow/tasks/framework/remote.py
@@ -22,6 +22,12 @@
from columnflow.types import Any
+# helpers to join paths relative to some base directories
+_cf_path = lambda *p: os.path.join(os.environ["CF_BASE"], *map(str, p))
+_repo_path = lambda *p: os.path.join(os.environ["CF_REPO_BASE"], *map(str, p))
+_cf_repo_path = lambda *p: [_cf_path(*p), _repo_path(*p)]
+
+
class BundleRepo(AnalysisTask, law.git.BundleGitRepository, law.tasks.TransferLocalFile):
replicas = luigi.IntParameter(
@@ -32,16 +38,19 @@ class BundleRepo(AnalysisTask, law.git.BundleGitRepository, law.tasks.TransferLo
version = None
exclude_files = [
- "docs",
- "tests",
- "data",
- "assets",
- ".law",
- ".setups",
- ".data",
- ".github",
- # also make sure that CF specific files that are not part of
- # the repository are excluded
+ # excluded from cf _and_ analysis repos
+ *_cf_repo_path("docs"),
+ *_cf_repo_path("tests"),
+ *_cf_repo_path("data"),
+ *_cf_repo_path("tmp"),
+ *_cf_repo_path(".data"),
+ *_cf_repo_path(".github"),
+ # excluded from cf repo
+ _cf_path("assets"),
+ # excluded from analysis repo
+ _repo_path(".law", "cms"),
+ _repo_path(".setups"),
+ # also make sure that CF specific files that are not part of the repository are excluded
os.environ["CF_STORE_LOCAL"],
os.environ["CF_SOFTWARE_BASE"],
os.environ["CF_VENV_BASE"],
@@ -49,7 +58,8 @@ class BundleRepo(AnalysisTask, law.git.BundleGitRepository, law.tasks.TransferLo
]
include_files = [
- "law_user.cfg",
+ _repo_path("law_user.cfg"),
+ _repo_path(".law"),
]
def get_repo_path(self):
@@ -392,17 +402,19 @@ def __init__(self, *args, **kwargs) -> None:
def get_config_lookup_keys(
cls,
inst_or_params: RemoteWorkflowMixin | dict[str, Any],
+ significant: bool = False,
) -> law.util.InsertiableDict:
- keys = super().get_config_lookup_keys(inst_or_params)
+ keys = super().get_config_lookup_keys(inst_or_params, significant=significant)
# add the pilot flag
- pilot = (
- inst_or_params.get("pilot")
- if isinstance(inst_or_params, dict)
- else getattr(inst_or_params, "pilot", None)
- )
- if pilot not in (law.NO_STR, None, ""):
- keys["pilot"] = f"pilot_{pilot}"
+ if not significant:
+ pilot = (
+ inst_or_params.get("pilot")
+ if isinstance(inst_or_params, dict)
+ else getattr(inst_or_params, "pilot", None)
+ )
+ if pilot not in (law.NO_STR, None, ""):
+ keys["pilot"] = f"pilot_{pilot}"
return keys
@@ -482,7 +494,7 @@ def add_bundle_render_variables(
:param config: The job :py:class:`law.BaseJobFileFactory.Config` whose render variables
should be set.
"""
- join_bash = lambda seq: " ".join(map('"{}"'.format, seq))
+ join_bash = lambda seq: " ".join(map(str, seq))
def get_bundle_info(task):
uris = task.output().dir.uri(base_name="filecopy", return_all=True)
@@ -560,15 +572,23 @@ def add_common_configs(
)
# forward voms proxy
- if voms and not law.config.get_expanded_boolean("analysis", "skip_ensure_proxy", False):
+ if voms:
+ # when skipping the check, still send it if the proxy exists and is valid, otherwise enforce it
+ skip_check = law.config.get_expanded_bool("analysis", "skip_ensure_proxy", False)
vomsproxy_file = law.wlcg.get_vomsproxy_file()
- if not law.wlcg.check_vomsproxy_validity(proxy_file=vomsproxy_file):
- raise Exception("voms proxy not valid, submission aborted")
- config.input_files["vomsproxy_file"] = law.JobInputFile(
- vomsproxy_file,
- share=True,
- render=False,
- )
+ vomsproxy_exists = os.path.isfile(vomsproxy_file)
+ vomsproxy_valid = vomsproxy_exists and law.wlcg.check_vomsproxy_validity(proxy_file=vomsproxy_file)
+ if not skip_check:
+ if not vomsproxy_exists:
+ raise Exception(f"voms proxy '{vomsproxy_file}' does not exist, submission aborted")
+ if not vomsproxy_valid:
+ raise Exception(f"voms proxy '{vomsproxy_file}' not valid, submission aborted")
+ if vomsproxy_valid:
+ config.input_files["vomsproxy_file"] = law.JobInputFile(
+ vomsproxy_file,
+ share=True,
+ render=False,
+ )
# forward kerberos proxy
if kerberos and "KRB5CCNAME" in os.environ:
@@ -654,7 +674,7 @@ def handle_scheduler_message(self, msg, _attr_value=None):
_default_htcondor_flavor = law.config.get_expanded("analysis", "htcondor_flavor", law.NO_STR)
-_default_htcondor_share_software = law.config.get_expanded_boolean("analysis", "htcondor_share_software", False)
+_default_htcondor_share_software = law.config.get_expanded_bool("analysis", "htcondor_share_software", False)
_default_htcondor_memory = law.util.parse_bytes(
law.config.get_expanded("analysis", "htcondor_memory", law.NO_FLOAT),
input_unit="GB",
@@ -754,8 +774,12 @@ class HTCondorWorkflow(RemoteWorkflowMixin, law.htcondor.HTCondorWorkflow):
"CF_STORE_NAME": "cf_store_name",
"CF_STORE_LOCAL": "cf_store_local",
"CF_LOCAL_SCHEDULER": "cf_local_scheduler",
+ "CF_PYTHON_VERSION": "cf_python_version",
}
+ # whether to show a memory summary histogram after workflow completion
+ show_memory_summary_hist = True
+
# upstream requirements
reqs = Requirements(
BundleRepo=BundleRepo,
@@ -869,6 +893,7 @@ def htcondor_job_config(self, config, job_num, branches):
# request memory
if self.htcondor_memory is not None and self.htcondor_memory > 0:
+ config.custom_content.append(("RequestMemory", f"{self.htcondor_memory} Gb"))
config.custom_content.append(("Request_Memory", f"{self.htcondor_memory} Gb"))
# request disk space
@@ -886,8 +911,9 @@ def htcondor_job_config(self, config, job_num, branches):
"cf_remote_lcg_setup_force",
"1" if law.config.get_expanded_bool("job", "remote_lcg_setup_force") else "",
)
- if self.htcondor_share_software:
- config.render_variables["cf_software_base"] = os.environ["CF_SOFTWARE_BASE"]
+ config.render_variables["cf_htcondor_share_software"] = str(self.htcondor_share_software).lower()
+ config.render_variables["cf_conda_base"] = os.environ["CF_CONDA_BASE"]
+ config.render_variables["cf_venv_base"] = os.environ["CF_VENV_BASE"]
# forward env variables
for ev, rv in self.htcondor_forward_env_variables.items():
@@ -904,6 +930,15 @@ def htcondor_destination_info(self, info: dict[str, str]) -> dict[str, str]:
info = self.common_destination_info(info)
return info
+ def htcondor_post_poll_callback(self, success, duration, summary_kwargs=None):
+ from law.workflow.remote import log_job_memory_summary
+
+ # prepare kwargs to forward
+ summary_kwargs = summary_kwargs or {}
+ summary_kwargs.setdefault("use_uniplot", self.show_memory_summary_hist)
+
+ log_job_memory_summary(self.workflow_proxy.job_data, log=self.logger.info, **summary_kwargs)
+
_default_slurm_flavor = law.config.get_expanded("analysis", "slurm_flavor", "maxwell")
_default_slurm_partition = law.config.get_expanded("analysis", "slurm_partition", "cms-uhh")
@@ -954,6 +989,7 @@ class SlurmWorkflow(RemoteWorkflowMixin, law.slurm.SlurmWorkflow):
"CF_STORE_NAME": "cf_store_name",
"CF_STORE_LOCAL": "cf_store_local",
"CF_LOCAL_SCHEDULER": "cf_local_scheduler",
+ "CF_PYTHON_VERSION": "cf_python_version",
}
# upstream requirements
diff --git a/columnflow/tasks/framework/remote_bootstrap.sh b/columnflow/tasks/framework/remote_bootstrap.sh
index 4d17e4c8f..2fb298c54 100755
--- a/columnflow/tasks/framework/remote_bootstrap.sh
+++ b/columnflow/tasks/framework/remote_bootstrap.sh
@@ -14,7 +14,10 @@ bootstrap_htcondor_standalone() {
export CF_CERN_USER_FIRSTCHAR="${CF_CERN_USER:0:1}"
export CF_REPO_BASE="${LAW_JOB_HOME}/repo"
export CF_DATA="${LAW_JOB_HOME}/cf_data"
- export CF_SOFTWARE_BASE="{{cf_software_base}}"
+ export CF_SOFTWARE_BASE="${CF_DATA}/software"
+ export CF_CONDA_BASE="{{cf_conda_base}}"
+ export CF_VENV_BASE="{{cf_venv_base}}"
+ export CF_PYTHON_VERSION="{{cf_python_version}}"
export CF_STORE_NAME="{{cf_store_name}}"
export CF_STORE_LOCAL="{{cf_store_local}}"
export CF_LOCAL_SCHEDULER="{{cf_local_scheduler}}"
@@ -26,9 +29,15 @@ bootstrap_htcondor_standalone() {
# also move it to the /tmp/x509up_u location as some packages expect
# the file to be at this path and do not respect the X509_USER_PROXY env var
local tmp_x509="/tmp/x509up_u$( id -u )"
- [ ! -f "${tmp_x509}" ] && cp "${X509_USER_PROXY}" "${tmp_x509}"
+ if [ ! -f "${tmp_x509}" ]; then
+ cp "${X509_USER_PROXY}" "${tmp_x509}" || {
+ >&2 echo "could not copy X509_USER_PROXY file from ${X509_USER_PROXY} to ${tmp_x509}, stopping job"
+ return "1"
+ }
+ fi
fi
- local sharing_software="$( [ -z "{{cf_software_base}}" ] && echo "false" || echo "true" )"
+ local share_software="{{cf_htcondor_share_software}}"
+ share_software="${share_software:-false}"
local lcg_setup="{{cf_remote_lcg_setup}}"
lcg_setup="${lcg_setup:-/cvmfs/grid.cern.ch/alma9-ui-test/etc/profile.d/setup-alma9-test.sh}"
local force_lcg_setup="$( [ -z "{{cf_remote_lcg_setup_force}}" ] && echo "false" || echo "true" )"
@@ -43,17 +52,42 @@ bootstrap_htcondor_standalone() {
export CAPATH="/cvmfs/grid.cern.ch/etc/grid-security/certificates"
fi
- # fallback to a default path when the externally given software base is empty or inaccessible
- local fetch_software="true"
- if [ -z "${CF_SOFTWARE_BASE}" ]; then
- export CF_SOFTWARE_BASE="${CF_DATA}/software"
- elif [ ! -d "${CF_SOFTWARE_BASE}/conda" ] || [ ! -x "${CF_SOFTWARE_BASE}/conda" ]; then
- echo "software base directory ${CF_SOFTWARE_BASE} was configured to be shared,"
- echo "but conda subdirectory is either missing or not accessible"
- export CF_SOFTWARE_BASE="${CF_DATA}/software"
+ # prepare sharing of conda and venv setup
+ local fetch_conda_bundle="true"
+ if ${share_software}; then
+ # check conda
+ if [ -z "${CF_CONDA_BASE}" ]; then
+ >&2 echo "cf_htcondor_share_software is set to true, but cf_conda_base is empty"
+ return "1"
+ fi
+ if [ ! -d "${CF_CONDA_BASE}" ] || [ ! -x "${CF_CONDA_BASE}" ]; then
+ >&2 echo "cf_htcondor_share_software is set to true, but cf_conda_base is not accessible: ${CF_CONDA_BASE}"
+ return "1"
+ fi
+ fetch_conda_bundle="false"
+ echo "detected existing conda setup at ${CF_CONDA_BASE}"
+
+ # check venvs
+ if [ -z "${CF_VENV_BASE}" ]; then
+ >&2 echo "cf_htcondor_share_software is set to true, but cf_venv_base is empty"
+ return "1"
+ fi
+ if [ ! -d "${CF_VENV_BASE}" ] || [ ! -x "${CF_VENV_BASE}" ]; then
+ >&2 echo "cf_htcondor_share_software is set to true, but cf_venv_base is not accessible: ${CF_VENV_BASE}"
+ return "1"
+ fi
+ echo "detected existing venvs at ${CF_VENV_BASE}"
else
- fetch_software="false"
- echo "found existing software at ${CF_SOFTWARE_BASE}"
+ # force local bases
+ export CF_CONDA_BASE="${CF_SOFTWARE_BASE}/conda"
+ export CF_VENV_BASE="${CF_SOFTWARE_BASE}/venvs"
+ # set variables required by _setup_venv script to load and setup venvs on-the-fly
+ export CF_JOB_BASH_SANDBOX_URIS="{{cf_bash_sandbox_uris}}"
+ export CF_JOB_BASH_SANDBOX_PATTERNS="{{cf_bash_sandbox_patterns}}"
+ export CF_JOB_BASH_SANDBOX_NAMES="{{cf_bash_sandbox_names}}"
+ export CF_JOB_CMSSW_SANDBOX_URIS="{{cf_cmssw_sandbox_uris}}"
+ export CF_JOB_CMSSW_SANDBOX_PATTERNS="{{cf_cmssw_sandbox_patterns}}"
+ export CF_JOB_CMSSW_SANDBOX_NAMES="{{cf_cmssw_sandbox_names}}"
fi
# when gfal is not available, check that the lcg_setup file exists
@@ -78,12 +112,12 @@ bootstrap_htcondor_standalone() {
echo "done sourcing wlcg tools"
# load and unpack the software bundle, then source it
- if ${fetch_software}; then
+ if ${fetch_conda_bundle}; then
(
echo -e "\nfetching software bundle ..."
{ ${skip_lcg_setup} || source "${lcg_setup}" ""; } &&
- mkdir -p "${CF_SOFTWARE_BASE}/conda" &&
- cd "${CF_SOFTWARE_BASE}/conda" &&
+ mkdir -p "${CF_CONDA_BASE}" &&
+ cd "${CF_CONDA_BASE}" &&
GFAL_PYTHONBIN="$( which python3 )" law_wlcg_get_file '{{cf_software_uris}}' '{{cf_software_pattern}}' "software.tgz" &&
tar -xzf "software.tgz" &&
rm "software.tgz" &&
@@ -103,16 +137,6 @@ bootstrap_htcondor_standalone() {
echo "done fetching repository bundle"
) || return "$?"
- # export variables used in cf setup script on-the-fly to load sandboxes
- if ! ${sharing_software}; then
- export CF_JOB_BASH_SANDBOX_URIS="{{cf_bash_sandbox_uris}}"
- export CF_JOB_BASH_SANDBOX_PATTERNS="{{cf_bash_sandbox_patterns}}"
- export CF_JOB_BASH_SANDBOX_NAMES="{{cf_bash_sandbox_names}}"
- export CF_JOB_CMSSW_SANDBOX_URIS="{{cf_cmssw_sandbox_uris}}"
- export CF_JOB_CMSSW_SANDBOX_PATTERNS="{{cf_cmssw_sandbox_patterns}}"
- export CF_JOB_CMSSW_SANDBOX_NAMES="{{cf_cmssw_sandbox_names}}"
- fi
-
# optional custom command before the setup is sourced
{{cf_pre_setup_command}}
@@ -136,6 +160,7 @@ bootstrap_slurm() {
export CF_SLURM_FLAVOR="{{cf_slurm_flavor}}"
export CF_REPO_BASE="{{cf_repo_base}}"
export CF_WLCG_CACHE_ROOT="${LAW_JOB_HOME}/cf_wlcg_cache"
+ export CF_PYTHON_VERSION="{{cf_python_version}}"
export KRB5CCNAME="FILE:{{kerberosproxy_file}}"
[ ! -z "{{vomsproxy_file}}" ] && export X509_USER_PROXY="{{vomsproxy_file}}"
@@ -162,6 +187,9 @@ bootstrap_crab() {
export CF_REPO_BASE="${LAW_JOB_HOME}/repo"
export CF_DATA="${LAW_JOB_HOME}/cf_data"
export CF_SOFTWARE_BASE="${CF_DATA}/software"
+ export CF_CONDA_BASE="${CF_SOFTWARE_BASE}/conda"
+ export CF_VENV_BASE="${CF_SOFTWARE_BASE}/venvs"
+ export CF_PYTHON_VERSION="{{cf_python_version}}"
export CF_STORE_NAME="{{cf_store_name}}"
export CF_STORE_LOCAL="${CF_DATA}/${CF_STORE_NAME}"
export CF_WLCG_CACHE_ROOT="${LAW_JOB_HOME}/cf_wlcg_cache"
@@ -191,8 +219,8 @@ bootstrap_crab() {
(
echo -e "\nfetching software bundle ..."
{ ${skip_lcg_setup} || source "${lcg_setup}" ""; } &&
- mkdir -p "${CF_SOFTWARE_BASE}/conda" &&
- cd "${CF_SOFTWARE_BASE}/conda" &&
+ mkdir -p "${CF_CONDA_BASE}" &&
+ cd "${CF_CONDA_BASE}" &&
GFAL_PYTHONBIN="$( which python3 )" law_wlcg_get_file '{{cf_software_uris}}' '{{cf_software_pattern}}' "software.tgz" &&
tar -xzf "software.tgz" &&
rm "software.tgz" &&
diff --git a/columnflow/tasks/histograms.py b/columnflow/tasks/histograms.py
index f57e18b14..f597e2964 100644
--- a/columnflow/tasks/histograms.py
+++ b/columnflow/tasks/histograms.py
@@ -22,7 +22,12 @@
from columnflow.tasks.production import ProduceColumns
from columnflow.tasks.ml import MLEvaluation
from columnflow.hist_util import update_ax_labels, sum_hists
-from columnflow.util import dev_sandbox
+from columnflow.config_util import expand_shift_sources
+from columnflow.util import maybe_import, dev_sandbox
+from columnflow.types import TYPE_CHECKING
+
+if TYPE_CHECKING:
+ hist = maybe_import("hist")
class VariablesMixinWorkflow(
@@ -89,6 +94,24 @@ def check_histogram_compatibility(cls, h) -> None:
def workflow_requires(self):
reqs = super().workflow_requires()
+ # depending on pilot flag, add upstream workflows or pass-through their own requirements only
+ reqs["producers"] = list(map(self.pilot_workflow_requires, (
+ self.reqs.ProduceColumns.req(
+ self,
+ producer=producer_inst.cls_name,
+ producer_inst=producer_inst,
+ )
+ for producer_inst in self.producer_insts
+ if producer_inst.produced_columns
+ )))
+ reqs["ml"] = list(map(self.pilot_workflow_requires, (
+ self.reqs.MLEvaluation.req(self, ml_model=ml_model_inst.cls_name)
+ for ml_model_inst in self.ml_model_insts
+ )))
+
+ # add hist producer dependent requirements
+ reqs["hist_producer"] = law.util.make_unique(law.util.flatten(self.hist_producer_inst.run_requires(task=self)))
+
# require the full merge forest
reqs["events"] = self.reqs.ProvideReducedEvents.req(self)
@@ -119,7 +142,7 @@ def workflow_requires(self):
return reqs
def requires(self):
- reqs = {"events": self.reqs.ProvideReducedEvents.req(self)}
+ reqs = {}
if self.producer_insts:
reqs["producers"] = [
@@ -142,6 +165,9 @@ def requires(self):
self.hist_producer_inst.run_requires(task=self),
))
+ # require merged events
+ reqs["events"] = self.reqs.ProvideReducedEvents.req(self)
+
return reqs
workflow_condition = ReducedEventsUser.workflow_condition.copy()
@@ -159,7 +185,7 @@ def run(self):
import numpy as np
import awkward as ak
from columnflow.columnar_util import (
- Route, update_ak_array, add_ak_aliases, has_ak_column, attach_coffea_behavior,
+ Route, update_ak_array, add_ak_aliases, has_ak_column, attach_coffea_behavior, ak_concatenate_safe,
)
from columnflow.columnar_util_Ghent import remove_obj_overlap
@@ -256,7 +282,7 @@ def run(self):
continue
# merge category ids and check that they are defined as leaf categories
- category_ids = ak.concatenate(
+ category_ids = ak_concatenate_safe(
[Route(c).apply(events) for c in self.category_id_columns],
axis=-1,
)
@@ -471,8 +497,11 @@ def run(self):
if not self.hist_producer_inst.skip_compatibility_check:
CreateHistograms.check_histogram_compatibility(merged)
+ # do not overwrite permissions when the file was already existing
+ perm = 0 if outputs["hists"][variable_name].exists() else None
+
# write the output
- outputs["hists"][variable_name].dump(merged, formatter="pickle")
+ outputs["hists"][variable_name].dump(merged, perm=perm, formatter="pickle")
# optionally remove inputs
if self.remove_previous:
@@ -503,11 +532,24 @@ class _MergeShiftedHistograms(
class MergeShiftedHistograms(_MergeShiftedHistograms):
+ shift_source_chunk_size = luigi.IntParameter(
+ default=law.NO_INT,
+ description="maximum number of shift sources to be processed by a single branch task; when > 0, sources are "
+ "split into chunks of that size and processed separately by multiple branch tasks; shifts in histograms are "
+ "consequently split across multiple output files; when empty or <= 0, no shift source splitting is applied and "
+ "there will only be a single branch; default: empty",
+ )
+
sandbox = dev_sandbox(law.config.get("analysis", "default_columnar_sandbox"))
# use the MergeHistograms task to trigger upstream TaskArrayFunction initialization
resolution_task_cls = MergeHistograms
+ output_collection_cls = law.FileCollection
+
+ # do not allow empty shift sources
+ allow_empty_shift_sources = False
+
# upstream requirements
reqs = Requirements(
RemoteWorkflow.reqs,
@@ -515,15 +557,15 @@ class MergeShiftedHistograms(_MergeShiftedHistograms):
)
def create_branch_map(self):
- # create a dummy branch map so that this task can run as a job
- return {0: None}
+ # each chunk of shift sources is assigned its own branch
+ return list(law.util.iter_chunks(self.shift_sources, self.shift_source_chunk_size))
def workflow_requires(self):
reqs = super().workflow_requires()
if not self.pilot:
# add nominal and both directions per shift source
- for shift in ["nominal"] + self.shifts:
+ for shift in expand_shift_sources(self.shift_sources):
reqs[shift] = self.reqs.MergeHistograms.req(self, shift=shift, _prefer_cli={"variables"})
return reqs
@@ -531,17 +573,32 @@ def workflow_requires(self):
def requires(self):
return {
shift: self.reqs.MergeHistograms.req(self, shift=shift, _prefer_cli={"variables"})
- for shift in ["nominal"] + self.shifts
+ for shift in expand_shift_sources(self.branch_data)
}
+ def store_parts(self) -> law.util.InsertableDict:
+ parts = super().store_parts()
+
+ if self.is_branch() and self.shift_source_chunk_size > 0:
+ shifts_repr = self._shift_sources_repr(self.branch_data, self.enforce_nominal_shift_source)
+ parts["shift_sources"] = f"shifts__{shifts_repr}"
+
+ return parts
+
def output(self):
return {
"hists": law.SiblingFileCollection({
- variable_name: self.target(f"hists__{variable_name}.pickle")
- for variable_name in self.variables
+ variable: self.target(f"hists__{variable}.pickle")
+ for variable in self.variables
}),
}
+ def control_output_postfix(self) -> str:
+ postfix = super().control_output_postfix()
+ if self.shift_source_chunk_size > 0:
+ postfix += f"__sscs{self.shift_source_chunk_size}"
+ return postfix
+
@law.decorator.notify
@law.decorator.log
def run(self):
@@ -549,21 +606,60 @@ def run(self):
inputs = self.input()
outputs = self.output()["hists"].targets
- for variable_name, outp in self.iter_progress(outputs.items(), len(outputs)):
- with self.publish_step(f"merging histograms for '{variable_name}' ..."):
- # load hists
- variable_hists = [
- coll["hists"].targets[variable_name].load(formatter="pickle")
- for coll in inputs.values()
- ]
+ shifts_msg = ""
+ if self.shift_source_chunk_size > 0:
+ shifts_msg = f" for {len(self.branch_data)} shift sources {','.join(self.branch_data)}"
- # update axis labels from variable insts for consistency
+ for variable_name, outp in self.iter_progress(outputs.items(), len(outputs)):
+ # verbosely load input histograms
+ with self.publish_step(f"loading '{variable_name}' histograms{shifts_msg} ..."):
+ variable_hists = []
+ for shift_name, inp in inputs.items():
+ try:
+ h = inp["hists"].targets[variable_name].load(formatter="pickle")
+ except:
+ self.logger.error(
+ f"cannot read file {inp['hists'].targets[variable_name].abspath} for with inputs for "
+ f"shift '{shift_name}'",
+ )
+ raise
+
+ # invoke modification hook, then save
+ h = self.modify_input_hist(shift_name, variable_name, h)
+ variable_hists.append(h)
+
+ # merge and write the output
+ with self.publish_step(f"merging '{variable_name}' histograms{shifts_msg} ..."):
update_ax_labels(variable_hists, self.config_inst, variable_name)
-
- # merge and write the output
merged = sum_hists(variable_hists)
+
+ # invoke modification hook, then save
+ merged = self.modify_merged_hist(variable_name, merged)
outp.dump(merged, formatter="pickle")
+ def modify_input_hist(self, shift: str, variable: str, h: hist.Hist) -> hist.Hist:
+ """
+ Hook to modify an input histogram for a shift after it has been loaded. This can be helpful to reduce memory
+ early on.
+
+ :param shift: The shift name the histogram corresponds to.
+ :param variable: The variable name the histogram corresponds to.
+ :param h: The histogram to modify.
+ :return: The modified histogram.
+ """
+ return h
+
+ def modify_merged_hist(self, variable: str, h: hist.Hist) -> hist.Hist:
+ """
+ Hook to modify the histogram merged for different shifts right before it is saved. This can be helpful to
+ reduce memory early on.
+
+ :param variable: The variable name the histogram corresponds to.
+ :param h: The histogram to modify.
+ :return: The modified histogram.
+ """
+ return h
+
MergeShiftedHistogramsWrapper = wrapper_factory(
base_cls=AnalysisTask,
diff --git a/columnflow/tasks/ml.py b/columnflow/tasks/ml.py
index 5ae06ec18..94a741997 100644
--- a/columnflow/tasks/ml.py
+++ b/columnflow/tasks/ml.py
@@ -33,7 +33,6 @@
from columnflow.tasks.reduction import ReducedEventsUser
from columnflow.tasks.production import ProduceColumns
from columnflow.util import dev_sandbox, safe_div, DotDict, maybe_import
-from columnflow.columnar_util import set_ak_column
ak = maybe_import("awkward")
@@ -83,29 +82,28 @@ def __init__(self, *args, **kwargs):
def workflow_requires(self):
reqs = super().workflow_requires()
- # require the full merge forest
- reqs["events"] = self.reqs.ProvideReducedEvents.req(self)
-
# add producer dependent requirements
if self.preparation_producer_inst:
reqs["preparation_producer"] = self.preparation_producer_inst.run_requires(task=self)
- # add producers to requirements
- if not self.pilot and self.producer_insts:
- reqs["producers"] = [
- self.reqs.ProduceColumns.req(
- self,
- producer=producer_inst.cls_name,
- producer_inst=producer_inst,
- )
- for producer_inst in self.producer_insts
- if producer_inst.produced_columns
- ]
+ # depending on pilot flag, add upstream workflows or pass-through their own requirements only
+ reqs["producers"] = list(map(self.pilot_workflow_requires, (
+ self.reqs.ProduceColumns.req(
+ self,
+ producer=producer_inst.cls_name,
+ producer_inst=producer_inst,
+ )
+ for producer_inst in self.producer_insts
+ if producer_inst.produced_columns
+ )))
+
+ # require the full merge forest
+ reqs["events"] = self.reqs.ProvideReducedEvents.req(self)
return reqs
def requires(self):
- reqs = {"events": self.reqs.ProvideReducedEvents.req(self)}
+ reqs = {}
if self.preparation_producer_inst:
reqs["preparation_producer"] = self.preparation_producer_inst.run_requires(task=self)
@@ -121,6 +119,9 @@ def requires(self):
if producer_inst.produced_columns
]
+ # require merged events
+ reqs["events"] = self.reqs.ProvideReducedEvents.req(self)
+
return reqs
workflow_condition = ReducedEventsUser.workflow_condition.copy()
@@ -144,7 +145,7 @@ def output(self):
@on_failure(callback=lambda task: task.teardown_preaparation_producer_inst())
def run(self):
from columnflow.columnar_util import (
- Route, RouteFilter, sorted_ak_to_parquet, update_ak_array, add_ak_aliases,
+ Route, RouteFilter, sorted_ak_to_parquet, update_ak_array, add_ak_aliases, set_ak_column,
)
from columnflow.columnar_util_Ghent import remove_obj_overlap
@@ -628,8 +629,6 @@ def workflow_requires(self):
configs=(self.config_inst.name,),
)
- reqs["events"] = self.reqs.ProvideReducedEvents.req(self)
-
# add producer dependent requirements
if self.preparation_producer_inst:
reqs["preparation_producer"] = self.preparation_producer_inst.run_requires(task=self)
@@ -645,17 +644,21 @@ def workflow_requires(self):
if producer_inst.produced_columns
]
+ # require the full merge forest
+ reqs["events"] = self.reqs.ProvideReducedEvents.req(self)
+
return reqs
def requires(self):
- reqs = {
- "models": self.reqs.MLTraining.req_different_branching(
- self,
- configs=(self.config_inst.name,),
- branch=-1,
- ),
- "events": self.reqs.ProvideReducedEvents.req(self, _exclude=self.exclude_params_branch),
- }
+ reqs = {}
+
+ # add models
+ reqs["models"] = self.reqs.MLTraining.req_different_branching(
+ self,
+ configs=(self.config_inst.name,),
+ branch=-1,
+ )
+
if self.preparation_producer_inst:
reqs["preparation_producer"] = self.preparation_producer_inst.run_requires(task=self)
@@ -670,6 +673,9 @@ def requires(self):
if producer_inst.produced_columns
]
+ # require merged events
+ reqs["events"] = self.reqs.ProvideReducedEvents.req(self, _exclude=self.exclude_params_branch)
+
return reqs
workflow_condition = ReducedEventsUser.workflow_condition.copy()
@@ -685,7 +691,7 @@ def output(self):
@on_failure(callback=lambda task: task.teardown_preparation_producer_inst())
def run(self):
from columnflow.columnar_util import (
- Route, RouteFilter, sorted_ak_to_parquet, update_ak_array, add_ak_aliases,
+ Route, RouteFilter, sorted_ak_to_parquet, update_ak_array, add_ak_aliases, set_ak_column,
)
# prepare inputs and outputs
@@ -995,6 +1001,8 @@ def prepare_inputs(self: PlotMLResultsBase) -> dict[str, ak.Array]:
:return: dict[str, ak.Array]: A dictionary with the dataset names as keys and
the corresponding predictions as values.
"""
+ from columnflow.columnar_util import ak_concatenate_safe
+
category_inst = self.config_inst.get_category(self.branch_data.category)
leaf_category_insts = category_inst.get_leaf_categories() or [category_inst]
process_insts = list(map(self.config_inst.get_process, self.processes))
@@ -1012,7 +1020,7 @@ def prepare_inputs(self: PlotMLResultsBase) -> dict[str, ak.Array]:
"which is not implemented yet.",
)
- events = ak.from_parquet(inp["mlcolumns"].abspath)
+ events = law.awkward.from_parquet(inp["mlcolumns"].abspath)
# masking with leaf categories
category_mask = False
@@ -1028,7 +1036,7 @@ def prepare_inputs(self: PlotMLResultsBase) -> dict[str, ak.Array]:
if not self.plot_sub_processes:
if process_inst.name in all_events.keys():
- all_events[process_inst.name] = ak.concatenate([
+ all_events[process_inst.name] = ak_concatenate_safe([
all_events[process_inst.name], getattr(events, self.ml_model),
])
else:
@@ -1045,7 +1053,7 @@ def prepare_inputs(self: PlotMLResultsBase) -> dict[str, ak.Array]:
process_mask = ak.where(events.process_ids == sub_process.id, True, False)
if sub_process.name in all_events.keys():
- all_events[sub_process.name] = ak.concatenate([
+ all_events[sub_process.name] = ak_concatenate_safe([
all_events[sub_process.name],
getattr(events[process_mask], self.ml_model),
])
diff --git a/columnflow/tasks/plotting.py b/columnflow/tasks/plotting.py
index ba91c0492..92c097303 100644
--- a/columnflow/tasks/plotting.py
+++ b/columnflow/tasks/plotting.py
@@ -5,6 +5,7 @@
"""
import itertools
+import functools
from collections import OrderedDict, defaultdict
from abc import abstractmethod
@@ -27,7 +28,7 @@
from columnflow.tasks.histograms import MergeHistograms, MergeShiftedHistograms
from columnflow.util import DotDict, dev_sandbox, dict_add_strict
from columnflow.hist_util import add_missing_shifts
-from columnflow.config_util import get_shift_from_configs
+from columnflow.config_util import get_shift_from_configs, expand_shift_sources
class _PlotVariablesBase(
@@ -51,6 +52,11 @@ class _PlotVariablesBase(
class PlotVariablesBase(_PlotVariablesBase):
+ multi_variable = luigi.BoolParameter(
+ default=False,
+ description="whether a single plot for all variables should be created; this requires that the used plot "
+ "function accepts a nested dictionary with all variable and process histograms as an input; default: False",
+ )
bypass_branch_requirements = luigi.BoolParameter(
default=False,
description="whether to skip branch requirements and only use that of the workflow; default: False",
@@ -72,11 +78,12 @@ def store_parts(self) -> law.util.InsertableDict:
return parts
def create_branch_map(self):
- return [
- DotDict({"category": cat_name, "variable": var_name})
- for cat_name in sorted(self.categories)
- for var_name in sorted(self.variables)
- ]
+ keys = ["category"]
+ seqs = [self.categories]
+ if not self.multi_variable:
+ keys.append("variable")
+ seqs.append(self.variables)
+ return [DotDict(zip(keys, vals)) for vals in itertools.product(*seqs)]
def workflow_requires(self):
reqs = super().workflow_requires()
@@ -118,15 +125,20 @@ def get_config_process_map(self) -> tuple[dict[od.Config, dict[od.Process, dict[
process_insts = [config_inst.get_process(p) for p in self.processes[i]]
dataset_insts = [config_inst.get_dataset(d) for d in self.datasets[i]]
- requested_shifts_per_dataset: dict[od.Dataset, list[od.Shift]] = {}
+ requested_shifts_per_dataset: dict[od.Dataset, list[str]] = {}
for dataset_inst in dataset_insts:
_req = reqs[config_inst.name][dataset_inst.name]
- if hasattr(_req, "shift") and _req.shift:
+ if isinstance(_req, ShiftTask) and _req.shift:
# when a shift is found, use it
requested_shifts = [_req.shift]
+ elif isinstance(_req, ShiftSourcesMixin):
+ # when no shift is found, check for shift sources and expand to up/down variations
+ requested_shifts = expand_shift_sources(_req.shift_sources)
else:
- # when no shift is found, check upstream requirements
- requested_shifts = [sub_req.shift for sub_req in _req.requires().values()]
+ raise Exception(
+ f"no shift or shift source found in requirements for dataset {dataset_inst.name} "
+ f"of config {config_inst.name}",
+ )
requested_shifts_per_dataset[dataset_inst] = requested_shifts
@@ -373,7 +385,10 @@ def plot_parts(self) -> law.util.InsertableDict:
parts["processes"] = f"proc_{self.processes_repr}"
parts["category"] = f"cat_{self.branch_data.category}"
- parts["variable"] = f"var_{self.branch_data.variable}"
+ if self.multi_variable:
+ parts["variables"] = f"vars_{self.variables_repr}"
+ else:
+ parts["variable"] = f"var_{self.branch_data.variable}"
hooks_repr = self.hist_hooks_repr
if hooks_repr:
@@ -481,6 +496,9 @@ class PlotVariablesBaseMultiShifts(
"the plot, the process_inst label is used; empty default",
)
+ # always ensure the nominal shift is present in shift sources
+ enforce_nominal_shift_source = True
+
# whether this task creates a single plot combining all shifts or one plot per shift
combine_shifts = True
@@ -510,19 +528,21 @@ def requires(self):
if self.is_branch() and self.bypass_branch_requirements:
return reqs
- req_cls = lambda dataset_name, config_inst: (
- self.reqs.MergeShiftedHistograms
- if config_inst.get_dataset(dataset_name).is_mc
- else self.reqs.MergeHistograms
- )
+ def hist_req(config_inst, dataset_name, **kwargs):
+ # return simple merged histograms for data
+ if config_inst.get_dataset(dataset_name).is_data:
+ return self.reqs.MergeHistograms.req(self, **kwargs)
+ # for mc, return shifted histograms
+ return self.reqs.MergeShiftedHistograms.req(self, **kwargs)
for config_inst, datasets in zip(self.config_insts, self.datasets):
reqs[config_inst.name] = {}
for d in datasets:
if d not in config_inst.datasets:
continue
- reqs[config_inst.name][d] = req_cls(d, config_inst).req(
- self,
+ reqs[config_inst.name][d] = hist_req(
+ config_inst,
+ d,
config=config_inst.name,
dataset=d,
branch=-1,
@@ -537,7 +557,10 @@ def plot_parts(self) -> law.util.InsertableDict:
parts["processes"] = f"proc_{self.processes_repr}"
parts["category"] = f"cat_{self.branch_data.category}"
- parts["variable"] = f"var_{self.branch_data.variable}"
+ if self.multi_variable:
+ parts["variables"] = f"vars_{self.variables_repr}"
+ else:
+ parts["variable"] = f"var_{self.branch_data.variable}"
# shift source or sources
parts["shift_source"] = (
@@ -557,20 +580,16 @@ def output(self):
"plots": [self.target(name) for name in self.get_plot_names("plot")],
}
- def get_plot_shifts(self):
+ def get_plot_shifts(self) -> list[od.Shift]:
# only to be called by branch tasks
if self.is_workflow():
raise Exception("calls to get_plots_shifts are forbidden for workflow tasks")
# gather sources, and expand to up/down shifts
sources = self.shift_sources if self.combine_shifts else [self.branch_data.shift_source]
- shifts = []
- for source in sources:
- shifts.append(get_shift_from_configs(self.config_insts, f"{source}_{od.Shift.UP}"))
- shifts.append(get_shift_from_configs(self.config_insts, f"{source}_{od.Shift.DOWN}"))
+ shifts = list(map(functools.partial(get_shift_from_configs, self.config_insts), expand_shift_sources(sources)))
- # add nominal
- return [self.config_inst.get_shift("nominal"), *shifts]
+ return shifts
def get_plot_parameters(self):
# convert parameters to usable values during plotting
@@ -595,6 +614,7 @@ class PlotShiftedVariablesPerShift1D(
):
# this tasks creates one plot per shift
combine_shifts = False
+
plot_function = PlotBase.plot_function.copy(
default="columnflow.plotting.plot_functions_1d.plot_shifted_variable",
add_default_to_description=True,
@@ -607,6 +627,7 @@ class PlotShiftedVariablesPerConfig1D(
):
# force this one to be a local workflow
workflow = "local"
+
output_collection_cls = law.NestedSiblingFileCollection
def requires(self):
diff --git a/columnflow/tasks/production.py b/columnflow/tasks/production.py
index bb5f8d090..1ecade920 100644
--- a/columnflow/tasks/production.py
+++ b/columnflow/tasks/production.py
@@ -9,12 +9,13 @@
import luigi
import law
-from columnflow.tasks.framework.base import Requirements, AnalysisTask, wrapper_factory
+from columnflow.tasks.framework.base import Requirements, AnalysisTask, wrapper_factory, RESOLVE_DEFAULT
from columnflow.tasks.framework.mixins import ProducerMixin, ChunkedIOMixin
from columnflow.tasks.framework.remote import RemoteWorkflow
from columnflow.tasks.framework.decorators import on_failure
from columnflow.tasks.reduction import ReducedEventsUser
from columnflow.util import dev_sandbox
+from columnflow.types import Any
class _ProduceColumns(
@@ -46,20 +47,20 @@ class ProduceColumns(_ProduceColumns):
def workflow_requires(self):
reqs = super().workflow_requires()
- # require the full merge forest
- reqs["events"] = self.reqs.ProvideReducedEvents.req(self)
-
# add producer dependent requirements
reqs["producer"] = law.util.make_unique(law.util.flatten(self.producer_inst.run_requires(task=self)))
+ # require the full merge forest
+ reqs["events"] = self.reqs.ProvideReducedEvents.req(self)
+
return reqs
def requires(self):
return {
- "events": self.reqs.ProvideReducedEvents.req(self),
"producer": law.util.make_unique(law.util.flatten(
self.producer_inst.run_requires(task=self),
)),
+ "events": self.reqs.ProvideReducedEvents.req(self),
}
workflow_condition = ReducedEventsUser.workflow_condition.copy()
@@ -195,16 +196,36 @@ def run(self):
enable=["configs", "skip_configs", "datasets", "skip_datasets", "shifts", "skip_shifts"],
)
_ProduceColumnsWrapperBase.exclude_index = True
+delattr(_ProduceColumnsWrapperBase, "producer")
class ProduceColumnsWrapper(_ProduceColumnsWrapperBase):
producers = law.CSVParameter(
- default=(),
- description="names of producers to use; if empty, the default producer is used",
+ default=(RESOLVE_DEFAULT,),
+ description="comma-separated names of producers to be applied; default: value of the 'default_producer' "
+ "analysis aux",
brace_expand=True,
+ parse_empty=True,
)
+ @classmethod
+ def modify_param_values(cls, params: dict[str, Any]) -> dict[str, Any]:
+ params = super().modify_param_values(params)
+
+ if params.get("analysis_inst"):
+ config_insts = [params["analysis_inst"].get_config(name) for name in params["configs"]]
+ params["producers"] = cls.resolve_config_default_and_groups(
+ param=params.get("producers"),
+ task_params=params,
+ container=config_insts,
+ default_str="default_producer",
+ groups_str="producer_groups",
+ multi_strategy="same",
+ )
+
+ return params
+
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
diff --git a/columnflow/tasks/reduction.py b/columnflow/tasks/reduction.py
index 1d62a5d71..4c206779c 100644
--- a/columnflow/tasks/reduction.py
+++ b/columnflow/tasks/reduction.py
@@ -18,7 +18,7 @@
from columnflow.tasks.framework.decorators import on_failure
from columnflow.tasks.external import GetDatasetLFNs
from columnflow.tasks.selection import CalibrateEvents, SelectEvents
-from columnflow.util import maybe_import, ensure_proxy, dev_sandbox, safe_div
+from columnflow.util import maybe_import, ensure_proxy, dev_sandbox, safe_div, DotDict
from columnflow.types import Any
ak = maybe_import("awkward")
@@ -63,21 +63,17 @@ def workflow_requires(self):
reqs["lfns"] = self.reqs.GetDatasetLFNs.req(self)
- if not self.pilot:
- reqs["calibrations"] = [
- self.reqs.CalibrateEvents.req(
- self,
- calibrator=calibrator_inst.cls_name,
- calibrator_inst=calibrator_inst,
- )
- for calibrator_inst in self.calibrator_insts
- if calibrator_inst.produced_columns
- ]
- reqs["selection"] = self.reqs.SelectEvents.req(self)
- else:
- # pass-through pilot workflow requirements of upstream task
- t = self.reqs.SelectEvents.req(self)
- law.util.merge_dicts(reqs, t.workflow_requires(), inplace=True)
+ # depending on pilot flag, add upstream workflows or pass-through their own requirements only
+ reqs["calibrations"] = list(map(self.pilot_workflow_requires, (
+ self.reqs.CalibrateEvents.req(
+ self,
+ calibrator=calibrator_inst.cls_name,
+ calibrator_inst=calibrator_inst,
+ )
+ for calibrator_inst in self.calibrator_insts
+ if calibrator_inst.produced_columns
+ )))
+ reqs["selection"] = self.pilot_workflow_requires(self.reqs.SelectEvents.req(self))
# add reducer dependent requirements
reqs["reducer"] = law.util.make_unique(law.util.flatten(self.reducer_inst.run_requires(task=self)))
@@ -141,6 +137,18 @@ def run(self):
inputs=luigi.task.getpaths(reducer_reqs),
)
+ # special case for reducers: issue a warning in case the upstream selector has shifts registered that are not
+ # known to the reducer, meaning that requested shifts would be known as global but not local ones, leading to
+ # the nominal behavior in the event chunk loop below, especially regarding alias handling
+ if (missing_reducer_shifts := self.selector_inst.all_shifts - self.reducer_inst.all_shifts):
+ self.logger.warning(
+ f"the upstream selector '{self.selector_inst.cls_name}' has {len(missing_reducer_shifts)} shifts "
+ f"registered that are not known to this reducer '{self.reducer_inst.cls_name}'; please check your "
+ "reducer as this is probably a misconfiguration and can lead to mismatches between event selection and "
+ "reduction, especially when shift-specific aliases are to be applied; missing shifts:\n"
+ f"{', '.join(sorted(missing_reducer_shifts))}",
+ )
+
# create a temp dir for saving intermediate files
tmp_dir = law.LocalDirectoryTarget(is_tmp=True)
tmp_dir.touch()
@@ -176,11 +184,13 @@ def run(self):
n_all = 0
n_reduced = 0
- # let the lfn_task prepare the nano file (basically determine a good pfn)
- [(lfn_index, input_file)] = lfn_task.iter_nano_files(self)
+ # let the lfn_task locate and prepare the nano file(s)
+ nano_input = [nano_target for _, nano_target in lfn_task.iter_nano_files(self)]
+ if len(nano_input) == 1:
+ nano_input = nano_input[0]
# collect input targets
- input_targets = [input_file]
+ input_targets = [nano_input]
input_targets.append(inputs["selection"]["results"])
input_targets.extend([inp["columns"] for inp in inputs["calibrations"]])
if self.selector_inst.produced_columns:
@@ -191,7 +201,7 @@ def run(self):
with law.localize_file_targets(input_targets, mode="r") as inps:
# iterate over chunks of events and diffs
for (events, sel, *diffs), pos in self.iter_chunked_io(
- [inp.abspath for inp in inps],
+ law.util.map_struct(law.target.file.get_path, inps),
source_type=["coffea_root"] + (len(inps) - 1) * ["awkward_parquet"],
read_columns=[read_columns, read_sel_columns] + (len(inps) - 2) * [read_columns],
chunk_size=self.reducer_inst.get_min_chunk_size(),
@@ -230,7 +240,7 @@ def run(self):
self.raise_if_not_finite(events)
# save as parquet via a thread in the same pool
- chunk = tmp_dir.child(f"file_{lfn_index}_{pos.index}.parquet", type="f")
+ chunk = tmp_dir.child(f"file_{pos.index}.parquet", type="f")
output_chunks[pos.index] = chunk
self.chunked_io.queue(sorted_ak_to_parquet, (ak.to_packed(events), chunk.abspath))
@@ -328,6 +338,14 @@ def resolve_param_values(cls, params: dict[str, Any]) -> dict[str, Any]:
return params
+ def __init__(self, *args, **kwargs) -> None:
+ super().__init__(*args, **kwargs)
+
+ # cap n_inputs
+ n_merged_files = self.n_merged_files
+ if self.n_inputs < 0 or self.n_inputs > n_merged_files:
+ self.n_inputs = n_merged_files
+
def create_branch_map(self):
# single branch without payload
return {0: None}
@@ -402,13 +420,17 @@ def get_avg_std(values):
stats["max_size_merged"] = self.merged_size * 1024**2 # MB to bytes
# determine the number of files after merging, allowing a possible ~15% increase per file
+ # TODO: merging: check n_files usage
n_total = self.dataset_info_inst.n_files
if n_total > 1:
- extrapolation = n_total / n
- n_merged_files = extrapolation * stats["tot_size"] / stats["max_size_merged"]
+ # get the expected number of files after merging
+ n_merged_files = n_total / n * stats["tot_size"] / stats["max_size_merged"]
+ # round using some heuristics
rnd = math.ceil if n_merged_files % 1.0 > 0.15 else math.floor
n_merged_files = max(int(rnd(n_merged_files)), 1)
+ # determine the merging factor and adjust the number of merged files accordingly for numeric edge cases
stats["merge_factor"] = max(math.ceil(n_total / n_merged_files), 1)
+ n_merged_files = max(math.ceil(n_total / stats["merge_factor"]), 1)
else:
# trivial case, no merging needed
n_merged_files = 1
@@ -420,13 +442,15 @@ def get_avg_std(values):
# print them
self.publish_message(f" stats of {n} input files ".center(40, "-"))
self.publish_message(f"average size: {law.util.human_bytes(stats['avg_size'], fmt=True)}")
- deviation = stats["std_size"] / stats["avg_size"]
- self.publish_message(f"deviation : {deviation * 100:.2f}% (std/avg)")
+ self.publish_message(f"deviation : {law.util.human_bytes(stats['std_size'], fmt=True)}")
self.publish_message(" merging info ".center(40, "-"))
self.publish_message(f"target size : {self.merged_size} MB")
self.publish_message(f"merging : {stats['merge_factor']} into 1")
self.publish_message(f"files before: {n_total}")
self.publish_message(f"files after : {n_merged_files}")
+ tot_size = stats["avg_size"] * n_total
+ rel_err = stats["std_size"] / stats["avg_size"]
+ self.publish_message(f"total size : {law.util.human_bytes(tot_size, fmt=True)} +- {rel_err * 100:.1f} %")
self.publish_message(40 * "-")
@@ -491,6 +515,7 @@ def workflow_requires(self):
reqs["stats"] = self.reqs.MergeReductionStats.req_different_branching(self)
reqs["events"] = self.reqs.ReduceEvents.req_different_branching(
self,
+ # TODO: merging: check n_files usage
branches=((0, self.dataset_info_inst.n_files),),
)
return reqs
@@ -527,6 +552,12 @@ def run(self):
target_row_group_size=self.merging_row_group_size,
)
+ # log the number of events after merging
+ n_events = output.load(formatter="parquet").metadata.num_rows
+ self.publish_message(f"merged file contains {n_events:_} events")
+ if not n_events:
+ self.logger.warning("the merged file contains no events")
+
# optionally remove initial inputs
if not self.keep_reduced_events and self.is_leaf():
with self.publish_step("removing reduced inputs ..."):
@@ -582,6 +613,7 @@ def _resolve_workflow_parameters(cls, params):
@law.workflow_property(setter=True, cache=True, empty_value=0)
def file_merging(self):
+ # TODO: merging: check n_files usage
if self.skip_merging or self.dataset_info_inst.n_files == 1:
return 1
@@ -609,6 +641,7 @@ def workflow_requires(self):
# - otherwise, always require the reduction stats as they are needed to make a decision
# - when merging is forced, require it
# - otherwise, and if the merging is already known, require either reduced or merged events
+ # TODO: merging: check n_files usage
if self.skip_merging or (not self.force_merging and self.dataset_info_inst.n_files == 1):
# reduced events are used directly without having to look into the file merging factor
if not self.pilot:
@@ -629,11 +662,16 @@ def workflow_requires(self):
elif file_merging == 1 and not self.pilot:
reqs["events"] = self._req_reduced_events()
+ # move reduction stats requirement to the end
+ if "reduction_stats" in reqs:
+ reqs.move_to_end("reduction_stats")
+
return reqs
def requires(self):
# same as for workflow requirements without optional pilot check
- reqs = {}
+ reqs = DotDict()
+ # TODO: merging: check n_files usage
if self.skip_merging or (not self.force_merging and self.dataset_info_inst.n_files == 1):
reqs["events"] = self._req_reduced_events()
else:
@@ -648,6 +686,10 @@ def requires(self):
elif file_merging == 1:
reqs["events"] = self._req_reduced_events()
+ # move reduction stats requirement to the end
+ if "reduction_stats" in reqs:
+ reqs.move_to_end("reduction_stats")
+
return reqs
@workflow_condition.output
@@ -661,6 +703,7 @@ def output(self):
def _yield_dynamic_deps(self):
# do nothing if a decision was pre-set in which case requirements were already triggered
+ # TODO: merging: check n_files usage
if self.skip_merging or (not self.force_merging and self.dataset_info_inst.n_files == 1):
return
diff --git a/columnflow/tasks/selection.py b/columnflow/tasks/selection.py
index cc2b32fae..968cc2a3f 100644
--- a/columnflow/tasks/selection.py
+++ b/columnflow/tasks/selection.py
@@ -10,16 +10,15 @@
import luigi
import law
-from columnflow.types import Any
-
from columnflow.tasks.framework.base import Requirements, AnalysisTask, wrapper_factory
from columnflow.tasks.framework.mixins import CalibratorsMixin, SelectorMixin, ChunkedIOMixin, ProducerMixin
from columnflow.tasks.framework.remote import RemoteWorkflow
from columnflow.tasks.framework.decorators import on_failure
+from columnflow.tasks.framework.parameters import DerivableInstParameter
from columnflow.tasks.external import GetDatasetLFNs
from columnflow.tasks.calibration import CalibrateEvents
from columnflow.util import maybe_import, ensure_proxy, dev_sandbox, safe_div, DotDict
-from columnflow.tasks.framework.parameters import DerivableInstParameter
+from columnflow.types import Any
from columnflow.production import Producer
np = maybe_import("numpy")
@@ -77,20 +76,16 @@ def workflow_requires(self):
reqs["lfns"] = self.reqs.GetDatasetLFNs.req(self)
- if not self.pilot:
- reqs["calibrations"] = [
- self.reqs.CalibrateEvents.req(
- self,
- calibrator=calibrator_inst.cls_name,
- calibrator_inst=calibrator_inst,
- )
- for calibrator_inst in self.calibrator_insts
- if calibrator_inst.produced_columns
- ]
- elif self.calibrator_insts:
- # pass-through pilot workflow requirements of upstream task
- t = self.reqs.CalibrateEvents.req(self)
- law.util.merge_dicts(reqs, t.workflow_requires(), inplace=True)
+ # depending on pilot flag, add upstream workflows or pass-through their own requirements only
+ reqs["calibrations"] = list(map(self.pilot_workflow_requires, (
+ self.reqs.CalibrateEvents.req(
+ self,
+ calibrator=calibrator_inst.cls_name,
+ calibrator_inst=calibrator_inst,
+ )
+ for calibrator_inst in self.calibrator_insts
+ if calibrator_inst.produced_columns
+ )))
# add selector dependent requirements
reqs["selector"] = law.util.make_unique(law.util.flatten(self.selector_inst.run_requires(task=self)))
@@ -196,13 +191,15 @@ def run(self):
write_columns |= self.selector_inst.produced_columns
route_filter = RouteFilter(keep=write_columns)
- # let the lfn_task prepare the nano file (basically determine a good pfn)
- [(lfn_index, input_file)] = lfn_task.iter_nano_files(self)
+ # let the lfn_task locate and prepare the nano file(s)
+ nano_input = [nano_target for _, nano_target in lfn_task.iter_nano_files(self)]
+ if len(nano_input) == 1:
+ nano_input = nano_input[0]
# prepare inputs for localization
with law.localize_file_targets(
[
- input_file,
+ nano_input,
*(inp["columns"] for inp in inputs["calibrations"]),
*reader_targets.values(),
],
@@ -211,7 +208,7 @@ def run(self):
# iterate over chunks of events and diffs
n_calib = len(inputs["calibrations"])
for (events, *cols), pos in self.iter_chunked_io(
- [inp.abspath for inp in inps],
+ law.util.map_struct(law.target.file.get_path, inps),
source_type=["coffea_root"] + ["awkward_parquet"] * n_calib + [None] * n_ext,
read_columns=[read_columns] * (1 + n_calib + n_ext),
chunk_size=self.selector_inst.get_min_chunk_size(),
@@ -224,7 +221,7 @@ def run(self):
events = update_ak_array(events, *cols)
# add veto
- events = self.veto_producer(events, file=input_file)
+ events = self.veto_producer(events, file=nano_input)
# add aliases
events = add_ak_aliases(
@@ -253,8 +250,8 @@ def run(self):
self.raise_if_not_finite(results_array[~events.veto])
# save results as parquet via a thread in the same pool
- chunk = tmp_dir.child(f"res_{lfn_index}_{pos.index}.parquet", type="f")
- result_chunks[(lfn_index, pos.index)] = chunk
+ chunk = tmp_dir.child(f"res_{pos.index}.parquet", type="f")
+ result_chunks[pos.index] = chunk
self.chunked_io.queue(sorted_ak_to_parquet, (results_array, chunk.abspath))
# remove columns
@@ -271,8 +268,8 @@ def run(self):
self.raise_if_not_finite(events[~veto])
# save additional columns as parquet via a thread in the same pool
- chunk = tmp_dir.child(f"cols_{lfn_index}_{pos.index}.parquet", type="f")
- column_chunks[(lfn_index, pos.index)] = chunk
+ chunk = tmp_dir.child(f"cols_{pos.index}.parquet", type="f")
+ column_chunks[pos.index] = chunk
self.chunked_io.queue(sorted_ak_to_parquet, (events, chunk.abspath))
# teardown the selector
diff --git a/columnflow/tasks/union.py b/columnflow/tasks/union.py
index e73784175..5da5afd26 100644
--- a/columnflow/tasks/union.py
+++ b/columnflow/tasks/union.py
@@ -63,30 +63,28 @@ class UniteColumns(_UniteColumns):
def workflow_requires(self):
reqs = super().workflow_requires()
+ # depending on pilot flag, add upstream workflows or pass-through their own requirements only
+ reqs["producers"] = list(map(self.pilot_workflow_requires, (
+ self.reqs.ProduceColumns.req(
+ self,
+ producer=producer_inst.cls_name,
+ producer_inst=producer_inst,
+ )
+ for producer_inst in self.producer_insts
+ if producer_inst.produced_columns
+ )))
+ reqs["ml"] = list(map(self.pilot_workflow_requires, (
+ self.reqs.MLEvaluation.req(self, ml_model=m)
+ for m in self.ml_models
+ )))
+
# require the full merge forest
reqs["events"] = self.reqs.ProvideReducedEvents.req(self)
- if not self.pilot:
- if self.producer_insts:
- reqs["producers"] = [
- self.reqs.ProduceColumns.req(
- self,
- producer=producer_inst.cls_name,
- producer_inst=producer_inst,
- )
- for producer_inst in self.producer_insts
- if producer_inst.produced_columns
- ]
- if self.ml_model_insts:
- reqs["ml"] = [
- self.reqs.MLEvaluation.req(self, ml_model=m)
- for m in self.ml_models
- ]
-
return reqs
def requires(self):
- reqs = {"events": self.reqs.ProvideReducedEvents.req(self)}
+ reqs = {}
if self.producer_insts:
reqs["producers"] = [
@@ -104,6 +102,9 @@ def requires(self):
for m in self.ml_models
]
+ # require merged events
+ reqs["events"] = self.reqs.ProvideReducedEvents.req(self)
+
return reqs
workflow_condition = ReducedEventsUser.workflow_condition.copy()
diff --git a/columnflow/tasks/yields.py b/columnflow/tasks/yields.py
index 55aec3c98..ff5dc58c1 100644
--- a/columnflow/tasks/yields.py
+++ b/columnflow/tasks/yields.py
@@ -18,6 +18,7 @@
)
from columnflow.tasks.framework.remote import RemoteWorkflow
from columnflow.tasks.histograms import MergeHistograms
+from columnflow.hist_util import select_category_bins
from columnflow.util import dev_sandbox, try_int
@@ -200,14 +201,7 @@ def run(self):
processes.append(process_inst)
for category_inst in category_insts:
- leaf_category_insts = category_inst.get_leaf_categories() or [category_inst]
-
- h_cat = h[{"category": [
- hist.loc(c.name)
- for c in leaf_category_insts
- if c.name in h.axes["category"]
- ]}]
- h_cat = h_cat[{"category": sum}]
+ h_cat = select_category_bins(h, category_inst, use_leaves=True, prefer_parents=True, reduce=True)
value = Number(h_cat.value)
if not self.skip_uncertainties:
diff --git a/columnflow/types.py b/columnflow/types.py
index 05d783daa..cb34b3f07 100644
--- a/columnflow/types.py
+++ b/columnflow/types.py
@@ -23,7 +23,7 @@
from types import ModuleType, GeneratorType, GenericAlias # noqa
from typing import ( # noqa
TYPE_CHECKING, Any, Union, TypeVar, ClassVar, Sequence, Callable, Generator, TextIO, Iterable, Hashable, Type,
- Literal,
+ Literal, Protocol, runtime_checkable,
)
from typing_extensions import Annotated, _AnnotatedAlias as AnnotatedType, TypeAlias # noqa
@@ -31,3 +31,6 @@
#: Generic type variable, more stringent than Any.
T = TypeVar("T")
+
+#: Type of the UNSET attribute in util
+UNSET_TYPE = object
diff --git a/columnflow/util.py b/columnflow/util.py
index 2db95eaf8..de58d8c03 100644
--- a/columnflow/util.py
+++ b/columnflow/util.py
@@ -17,6 +17,7 @@
import threading
import subprocess
import importlib
+import pathlib
import fnmatch
import inspect
import pprint
@@ -30,7 +31,7 @@
import luigi
from columnflow import env_is_dev, env_is_remote, docs_url, github_url
-from columnflow.types import Callable, Any, Sequence, Union, ModuleType, Type, T, Hashable
+from columnflow.types import Callable, Any, Sequence, Union, ModuleType, Type, T, Hashable, Protocol, runtime_checkable
#: Placeholder for an unset value.
@@ -362,7 +363,7 @@ def before_call():
return None
# do nothing when explicitly skipped by the law config
- if law.config.get_expanded_boolean("analysis", "skip_ensure_proxy", False):
+ if law.config.get_expanded_bool("analysis", "skip_ensure_proxy", False):
return None
# check the proxy validity
@@ -1090,7 +1091,79 @@ def __str__(self) -> str:
return str(self.value)
-def load_correction_set(target: law.FileSystemFileTarget) -> Any:
+@runtime_checkable
+class CacheBase(Protocol):
+
+ def has(self, key: Hashable) -> bool:
+ ...
+
+ def get(self, key: Hashable) -> Any:
+ ...
+
+ def set(self, key: Hashable, value: Any) -> None:
+ ...
+
+
+class PersistentCache(CacheBase):
+
+ def __init__(self, cache_path: str | pathlib.Path | law.FileSystemFileTarget) -> None:
+ super().__init__()
+
+ self.cache_target = (
+ cache_path
+ if isinstance(cache_path, law.FileSystemFileTarget)
+ else law.LocalFileTarget(cache_path)
+ )
+
+ # state
+ self.cache: dict[Hashable, Any] = {}
+ self.opened: bool = False
+ self.modified: bool = False
+
+ def __enter__(self) -> PersistentCache:
+ return self.open()
+
+ def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+ self.close()
+
+ def open(self) -> PersistentCache:
+ if self.opened:
+ raise Exception(f"{self.__class__.__name__} is already opened")
+
+ # when existing, overwrite cache in memory with that from file
+ self.cache.clear()
+ if self.cache_target.exists():
+ self.cache |= self.cache_target.load(formatter="json")
+
+ self.modified = False
+ self.opened = True
+
+ return self
+
+ def close(self) -> None:
+ if not self.opened:
+ raise Exception(f"{self.__class__.__name__} is not opened")
+
+ # write to file when modified
+ if self.modified:
+ self.cache_target.dump(self.cache, formatter="json", indent=2)
+
+ self.modified = False
+ self.opened = False
+
+ def has(self, key: Hashable) -> bool:
+ return key in self.cache
+
+ def get(self, key: Hashable) -> Any:
+ return self.cache[key]
+
+ def set(self, key: Hashable, value: Any) -> None:
+ if not self.modified and (not self.has(key) or self.get(key) != value):
+ self.modified = True
+ self.cache[key] = value
+
+
+def load_correction_set(target: law.FileSystemFileTarget | str) -> Any:
"""
Loads a correction set using the correctionlib from a file *target*.
"""
@@ -1099,6 +1172,10 @@ def load_correction_set(target: law.FileSystemFileTarget) -> Any:
# extend the Correction object
correctionlib.highlevel.Correction.__call__ = correctionlib.highlevel.Correction.evaluate
+ # convert str to target
+ if isinstance(target, str):
+ target = law.LocalFileTarget(os.path.abspath(target))
+
# use the path when the input file is a normal json
if target.ext() == "json":
return correctionlib.CorrectionSet.from_file(target.abspath)
diff --git a/create_analysis.sh b/create_analysis.sh
index df7af9e09..daa272d4a 100755
--- a/create_analysis.sh
+++ b/create_analysis.sh
@@ -199,7 +199,7 @@ create_analysis() {
return "1"
fi
- echo_color cyan "checking out analysis tempate to ${cf_analysis_base}"
+ echo_color cyan "checking out analysis template to ${cf_analysis_base}"
if ${debug}; then
cp -r "${this_dir}/analysis_templates/${cf_analysis_flavor}" "${cf_analysis_base}"
diff --git a/law.cfg b/law.cfg
index 2ba8f2ba3..a5a46ddbb 100644
--- a/law.cfg
+++ b/law.cfg
@@ -215,6 +215,7 @@ cache_cleanup: $CF_WLCG_CACHE_CLEANUP
cache_max_size: 15GB
cache_global_lock: True
cache_mtime_patience: -1
+rucio_report_access: T2_DE_DESY
[wlcg_fs_infn_redirector]
@@ -226,6 +227,7 @@ cache_cleanup: $CF_WLCG_CACHE_CLEANUP
cache_max_size: 15GB
cache_global_lock: True
cache_mtime_patience: -1
+rucio_report_access: True
[wlcg_fs_fnal_redirector]
@@ -237,6 +239,7 @@ cache_cleanup: $CF_WLCG_CACHE_CLEANUP
cache_max_size: 15GB
cache_global_lock: True
cache_mtime_patience: -1
+rucio_report_access: True
[wlcg_fs_global_redirector]
@@ -248,6 +251,7 @@ cache_cleanup: $CF_WLCG_CACHE_CLEANUP
cache_max_size: 15GB
cache_global_lock: True
cache_mtime_patience: -1
+rucio_report_access: True
[luigi_core]
diff --git a/sandboxes/_setup_venv.sh b/sandboxes/_setup_venv.sh
index fc45460e1..2e92233ae 100644
--- a/sandboxes/_setup_venv.sh
+++ b/sandboxes/_setup_venv.sh
@@ -203,7 +203,10 @@ setup_venv() {
fi
# create the pending_flag to express that the venv state might be changing
- [ ! -f "${pending_flag_file}" ] && touch "${pending_flag_file}"
+ mark_pending() {
+ [ -f "${pending_flag_file}" ] && return "0"
+ touch "${pending_flag_file}"
+ }
clear_pending() {
rm -f "${pending_flag_file}"
}
@@ -220,6 +223,7 @@ setup_venv() {
if [ "${current_version}" != "${venv_version}" ]; then
if [ "${mode}" = "update" ]; then
# remove the venv in case an update is requested
+ mark_pending || return "$?"
echo "removing current installation at ${install_path_repr} (mode '${mode}', installed version ${current_version}, requested version ${venv_version})"
rm -rf "${install_path}"
@@ -248,6 +252,9 @@ setup_venv() {
# install if not existing
if [ ! -f "${CF_SANDBOX_FLAG_FILE}" ]; then
+ mark_pending || return "$?"
+
+
echo -n "$( cf_color cyan "installing venv" )"
echo -n " $( cf_color cyan_bright "${CF_VENV_NAME}" )"
echo " $( cf_color cyan "from ${sandbox_file} at ${install_path}" )"
@@ -328,7 +335,14 @@ setup_venv() {
law_wlcg_get_file "${sandbox_uris[i]}" "${sandbox_patterns[i]}" "bundle.tgz" &&
tar -xzf "bundle.tgz"
) || return "$?"
+
found_sandbox="true"
+
+ # let the home variable in pyvenv.cfg point to the conda bin directory
+ sed -i -r \
+ "s|^(home = ).+/bin/?$|\1$CF_CONDA_BASE\/bin|" \
+ "${install_path}/pyvenv.cfg"
+
break
done
if ! ${found_sandbox}; then
@@ -337,11 +351,6 @@ setup_venv() {
fi
fi
- # let the home variable in pyvenv.cfg point to the conda bin directory
- sed -i -r \
- "s|^(home = ).+/bin/?$|\1$CF_CONDA_BASE\/bin|" \
- "${install_path}/pyvenv.cfg"
-
# activate it
source "${install_path}/bin/activate" "" || return "$?"
diff --git a/sandboxes/cf.txt b/sandboxes/cf.txt
index 7c910e0ef..aea9e830b 100644
--- a/sandboxes/cf.txt
+++ b/sandboxes/cf.txt
@@ -1,7 +1,8 @@
# version 15
luigi~=3.6.0
-scinum~=2.2.1
+scinum~=2.2.2
six~=1.17.0
pyyaml~=6.0.2
-typing_extensions~=4.12.2
+typing_extensions~=4.13.0
+tabulate~=0.9.0
diff --git a/sandboxes/cmssw_columnar.sh b/sandboxes/cmssw_columnar.sh
index cdc8f80bd..a91e02c24 100644
--- a/sandboxes/cmssw_columnar.sh
+++ b/sandboxes/cmssw_columnar.sh
@@ -10,8 +10,8 @@ action() {
# set variables and source the generic CMSSW setup
export CF_SANDBOX_FILE="${CF_SANDBOX_FILE:-${this_file}}"
- export CF_SCRAM_ARCH="el9_amd64_gcc11"
- export CF_CMSSW_VERSION="CMSSW_13_0_19"
+ export CF_SCRAM_ARCH="el9_amd64_gcc12"
+ export CF_CMSSW_VERSION="CMSSW_14_1_9"
export CF_CMSSW_ENV_NAME="$( basename "${this_file%.sh}" )"
export CF_CMSSW_FLAG="1" # increment when content changed
diff --git a/sandboxes/cmssw_default.sh b/sandboxes/cmssw_default.sh
index d2e31eb15..2b4f823a7 100644
--- a/sandboxes/cmssw_default.sh
+++ b/sandboxes/cmssw_default.sh
@@ -10,8 +10,8 @@ action() {
# set variables and source the generic CMSSW setup
export CF_SANDBOX_FILE="${CF_SANDBOX_FILE:-${this_file}}"
- export CF_SCRAM_ARCH="el9_amd64_gcc11"
- export CF_CMSSW_VERSION="CMSSW_13_0_19"
+ export CF_SCRAM_ARCH="el9_amd64_gcc12"
+ export CF_CMSSW_VERSION="CMSSW_14_1_9"
export CF_CMSSW_ENV_NAME="$( basename "${this_file%.sh}" )"
export CF_CMSSW_FLAG="1" # increment when content changed
diff --git a/sandboxes/columnar.txt b/sandboxes/columnar.txt
index 36cbda25c..270c1b874 100644
--- a/sandboxes/columnar.txt
+++ b/sandboxes/columnar.txt
@@ -1,4 +1,4 @@
-# version 18
+# version 19
# exact versions for core array packages
awkward==2.8.9
@@ -8,6 +8,7 @@ correctionlib==2.7.0
coffea==2025.9.0
# minimum versions for general packages
+rucio~=40.3.0
zstandard~=0.25.0
lz4~=4.4.4
xxhash~=3.5.0
diff --git a/sandboxes/dev.txt b/sandboxes/dev.txt
index 16fd687de..00e0455de 100644
--- a/sandboxes/dev.txt
+++ b/sandboxes/dev.txt
@@ -11,3 +11,4 @@ pymarkdownlnt~=0.9.32
uniplot~=0.21.4
pipdeptree~=2.28.0
mermaidmro~=0.2.1
+psutil~=7.1.3
diff --git a/setup.sh b/setup.sh
index bdcf9337b..31faf846a 100644
--- a/setup.sh
+++ b/setup.sh
@@ -37,6 +37,8 @@ setup_columnflow() {
# serves as a default for e.g. $CF_SOFTWARE_BASE, $CF_CMSSW_BASE, $CF_JOB_BASE and
# $CF_VENV_BASE which can, however, potentially point to a different directory. Queried
# during the interactive setup.
+ # CF_PYTHON_VERSION
+ # The python version to use for the conda and venv installations. Defaults to "3.9".
# CF_SOFTWARE_BASE
# The directory where general software is installed. Might point to $CF_DATA/software.
# Queried during the interactive setup.
@@ -275,8 +277,22 @@ cf_setup_common_variables() {
export LC_ALL="${LC_ALL:-en_US.UTF-8}"
fi
- # proxy
+ # variables for external tools and libraries
+ export PYTHONWARNINGS="${PYTHONWARNINGS:-ignore}"
+ export PYTHONNOUSERSITE="${PYTHONNOUSERSITE:-1}"
+ export GLOBUS_THREAD_MODEL="${GLOBUS_THREAD_MODEL:-none}"
+ export VIRTUAL_ENV_DISABLE_PROMPT="${VIRTUAL_ENV_DISABLE_PROMPT:-1}"
+ export MYPROXY_SERVER="${MYPROXY_SERVER:-myproxy.cern.ch}"
export X509_USER_PROXY="${X509_USER_PROXY:-/tmp/x509up_u$( id -u )}"
+ export X509_CERT_DIR="${X509_CERT_DIR:-/cvmfs/grid.cern.ch/etc/grid-security/certificates}"
+ export X509_VOMS_DIR="${X509_VOMS_DIR:-/cvmfs/grid.cern.ch/etc/grid-security/vomsdir}"
+ export X509_VOMSES="${X509_VOMSES:-/cvmfs/grid.cern.ch/etc/grid-security/vomses}"
+ export VOMS_USERCONF="${VOMS_USERCONF:-${X509_VOMSES}}"
+ ulimit -s unlimited
+ if [ "${CF_FLAVOR}" = "cms" ]; then
+ export RUCIO_ACCOUNT="${RUCIO_ACCOUNT:-${CF_CERN_USER}}"
+ export RUCIO_HOME="${RUCIO_HOME:-"/cvmfs/cms.cern.ch/rucio/x86_64/rhel9/py3/current"}"
+ fi
# overwrite some variables in remote and ci jobs
if ${CF_REMOTE_ENV}; then
@@ -310,15 +326,17 @@ cf_setup_common_variables() {
export CF_SLURM_FLAVOR="${CF_SLURM_FLAVOR:-${cf_slurm_flavor_default}}"
export CF_SLURM_PARTITION="${CF_SLURM_PARTITION:-${cf_slurm_partition_default}}"
+ # notification variables
+ export CF_MATTERMOST_HOOK_URL="${CF_MATTERMOST_HOOK_URL:-}"
+ export CF_MATTERMOST_CHANNEL="${CF_MATTERMOST_CHANNEL:-}"
+
# show a warning in case no CF_REPO_BASE_ALIAS is set
if [ -z "${CF_REPO_BASE_ALIAS}" ]; then
cf_color yellow "the variable CF_REPO_BASE_ALIAS is unset"
cf_color yellow "please consider setting it to the name of the variable that refers to your analysis base directory"
fi
- # notification variables
- export CF_MATTERMOST_HOOK_URL="${CF_MATTERMOST_HOOK_URL:-}"
- export CF_MATTERMOST_CHANNEL="${CF_MATTERMOST_CHANNEL:-}"
+ return "0"
}
cf_show_banner() {
@@ -376,7 +394,7 @@ cf_setup_interactive_common_variables() {
query CF_VENV_SETUP_MODE_UPDATE "Automatically update virtual envs if needed" "false"
[ "${CF_VENV_SETUP_MODE_UPDATE}" != "true" ] && export_and_save CF_VENV_SETUP_MODE "update"
unset CF_VENV_SETUP_MODE_UPDATE
- query CF_INTERACTIVE_VENV_FILE "Custom venv setup fill to use for interactive work instead of 'cf_dev'" "" "''"
+ query CF_INTERACTIVE_VENV_FILE "Custom venv setup file to use for interactive work instead of 'cf_dev'" "" "''"
query CF_LOCAL_SCHEDULER "Use a local scheduler for law tasks" "true"
if [ "${CF_LOCAL_SCHEDULER}" != "true" ]; then
@@ -529,6 +547,8 @@ cf_setup_software_stack() {
# The base directory were virtual envs are installed.
#
# Optional environments variables:
+ # CF_PYTHON_VERSION
+ # The python version for the conda and venv installations. Defaults to "3.9".
# CF_REMOTE_ENV
# When true-ish, the software stack is sourced but not built.
# CF_LOCAL_ENV
@@ -562,7 +582,6 @@ cf_setup_software_stack() {
local setup_name="${1}"
local setup_is_default="false"
[ "${setup_name}" = "default" ] && setup_is_default="true"
- local pyv="${CF_PYTHON_VERSION:-3.9}"
local conda_arch="${CF_CONDA_ARCH:-linux-64}"
local ret
@@ -572,13 +591,16 @@ cf_setup_software_stack() {
setopt globdots
fi
+ # default python version
+ export CF_PYTHON_VERSION="${CF_PYTHON_VERSION:-3.9}"
+
# empty the PYTHONPATH, except for specific customizable paths
export PYTHONPATH="${CF_INITIAL_PYTHONPATH:-}"
# persistent PATH and PYTHONPATH parts that should be
# priotized over any additions made in sandboxes
export CF_PERSISTENT_PATH="${CF_BASE}/bin:${CF_BASE}/modules/law/bin"
- export CF_PERSISTENT_PYTHONPATH="${CF_BASE}:${CF_BASE}/bin:${CF_BASE}/modules/law:${CF_BASE}/modules/order"
+ export CF_PERSISTENT_PYTHONPATH="${CF_BASE}:${CF_BASE}/bin:${CF_BASE}/modules/law:${CF_BASE}/modules/law/src:${CF_BASE}/modules/order"
# flavor specific paths
if [ ! -z "${CF_FLAVOR}" ]; then
@@ -590,22 +612,12 @@ cf_setup_software_stack() {
export PYTHONPATH="${CF_PERSISTENT_PYTHONPATH}:${PYTHONPATH}"
# also add the python path of the venv to be installed to propagate changes to any outer venv
- export CF_CONDA_PYTHONPATH="${CF_CONDA_BASE}/lib/python${pyv}/site-packages"
+ export CF_CONDA_PYTHONPATH="${CF_CONDA_BASE}/lib/python${CF_PYTHON_VERSION}/site-packages"
export PYTHONPATH="${PYTHONPATH}:${CF_CONDA_PYTHONPATH}"
- # update paths and flags
+ # update mamba variables
export MAMBA_ROOT_PREFIX="${CF_CONDA_BASE}"
export MAMBA_EXE="${MAMBA_ROOT_PREFIX}/bin/micromamba"
- export PYTHONWARNINGS="${PYTHONWARNINGS:-ignore}"
- export PYTHONNOUSERSITE="${PYTHONNOUSERSITE:-1}"
- export GLOBUS_THREAD_MODEL="${GLOBUS_THREAD_MODEL:-none}"
- export VIRTUAL_ENV_DISABLE_PROMPT="${VIRTUAL_ENV_DISABLE_PROMPT:-1}"
- export MYPROXY_SERVER="${MYPROXY_SERVER:-myproxy.cern.ch}"
- export X509_CERT_DIR="${X509_CERT_DIR:-/cvmfs/grid.cern.ch/etc/grid-security/certificates}"
- export X509_VOMS_DIR="${X509_VOMS_DIR:-/cvmfs/grid.cern.ch/etc/grid-security/vomsdir}"
- export X509_VOMSES="${X509_VOMSES:-/cvmfs/grid.cern.ch/etc/grid-security/vomses}"
- export VOMS_USERCONF="${VOMS_USERCONF:-${X509_VOMSES}}"
- ulimit -s unlimited
#
# setup in local envs (not remote)
@@ -656,7 +668,7 @@ EOF
# initialize micromamba
source "${CF_CONDA_BASE}/etc/profile.d/micromamba.sh" "" || return "$?"
micromamba activate || return "$?"
- echo "initialized conda with $( cf_color magenta "micromamba" ) interface and $( cf_color magenta "python ${pyv}" )"
+ echo "initialized conda with $( cf_color magenta "micromamba" ) interface and $( cf_color magenta "python ${CF_PYTHON_VERSION}" )"
# install packages
if ${conda_missing}; then
@@ -667,7 +679,7 @@ EOF
libgcc \
bash \
zsh \
- "python=${pyv}" \
+ "python=${CF_PYTHON_VERSION}" \
git \
git-lfs \
gfal2 \
@@ -771,7 +783,7 @@ EOF
# initialize conda
source "${CF_CONDA_BASE}/etc/profile.d/micromamba.sh" "" || return "$?"
micromamba activate || return "$?"
- echo "initialized conda with $( cf_color magenta "micromamba" ) interface and $( cf_color magenta "python ${pyv}" )"
+ echo "initialized conda with $( cf_color magenta "micromamba" ) interface and $( cf_color magenta "python ${CF_PYTHON_VERSION}" )"
# source the production sandbox
source "${CF_BASE}/sandboxes/cf.sh" "" "no"
diff --git a/tests/test_base_tasks.py b/tests/test_base_tasks.py
index 030bb538c..8ac54ae56 100644
--- a/tests/test_base_tasks.py
+++ b/tests/test_base_tasks.py
@@ -143,14 +143,14 @@ def test_resolve_config_default(self):
groups_str="producer_groups",
multi_strategy="first",
)
- self.assertEqual(resolved_producer_groups, ("b", "a", "B", "d", "c")) # TODO: order reversed
+ self.assertEqual(resolved_producer_groups, ("a", "b", "B", "c", "d"))
# multi config
for multi_strategy, expected_producer in (
("all", {self.config_inst1: ("A", "B", "C"), self.config_inst2: ("B", "C", "D")}),
("first", ("A", "B", "C")),
- ("union", ("A", "B", "C", "D")),
- ("intersection", ("B", "C")),
+ ("union", ["A", "B", "C", "D"]),
+ ("intersection", ["B", "C"]),
):
resolved_producer = AnalysisTask.resolve_config_default(
param=(RESOLVE_DEFAULT,),
@@ -159,8 +159,7 @@ def test_resolve_config_default(self):
default_str="default_producer",
multi_strategy=multi_strategy,
)
- # TODO: remove set() when order is fixed
- self.assertEqual(set(resolved_producer), set(expected_producer))
+ self.assertEqual(resolved_producer, expected_producer)
# "same" strategy
resolved_calibrator = AnalysisTask.resolve_config_default(
@@ -231,8 +230,7 @@ def test_resolve_categories(self):
"categories": input_categories,
}
resolved_params = CategoriesMixin.modify_param_values(params=input_params)
- # TODO: remove set() when order is fixed
- self.assertEqual(set(resolved_params["categories"]), set(expected_categories))
+ self.assertEqual(resolved_params["categories"], expected_categories)
def test_resolve_variables(self):
# testing with single config
@@ -247,8 +245,7 @@ def test_resolve_variables(self):
"variables": input_variables,
}
resolved_params = VariablesMixin.modify_param_values(params=input_params)
- # TODO: remove set() when order is fixed
- self.assertEqual(set(resolved_params["variables"]), set(expected_variables))
+ self.assertEqual(resolved_params["variables"], expected_variables)
def test_resolve_datasets_processes(self):
DatasetsProcessesMixin.single_config = False
diff --git a/tests/test_hist_util.py b/tests/test_hist_util.py
index 9049a0d03..927f0a478 100644
--- a/tests/test_hist_util.py
+++ b/tests/test_hist_util.py
@@ -6,10 +6,7 @@
import unittest
from columnflow.util import maybe_import
-from columnflow.hist_util import (
- add_hist_axis, create_hist_from_variables, create_columnflow_hist,
- translate_hist_intcat_to_strcat,
-)
+from columnflow.hist_util import add_hist_axis, create_hist_from_variables, translate_hist_intcat_to_strcat
import order as od
np = maybe_import("numpy")
@@ -96,7 +93,14 @@ def test_create_hist_from_variables(self):
self.assertEqual(histogram, histogram_manually)
# test with default categorical axes
- histogram = create_columnflow_hist(*self.variable_examples)
+ histogram = create_hist_from_variables(
+ *self.variable_examples,
+ categorical_axes=[
+ ("category", "intcat"),
+ ("process", "intcat"),
+ ("shift", "strcat"),
+ ],
+ )
expected_default_axes = {
"category": hist.axis.IntCategory,
diff --git a/tests/test_inference.py b/tests/test_inference.py
index 16e4161f9..6248fe70c 100644
--- a/tests/test_inference.py
+++ b/tests/test_inference.py
@@ -32,6 +32,7 @@ def test_process_spec(self):
},
scale=scale,
is_dynamic=is_dynamic,
+ skip_if_empty=True,
parameters=[],
)