46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
|
# flake8: noqa
|
||
|
|
import ROOT
|
||
|
|
|
||
|
|
|
||
|
|
def preselection(
|
||
|
|
cuts: str = "",
|
||
|
|
input_file: str = None,
|
||
|
|
outfile_postfix: str = "selected",
|
||
|
|
tree_name: str = "PrDebugTrackingLosses.PrDebugTrackingTool/Tuple",
|
||
|
|
) -> str:
|
||
|
|
"""Function that apply a selection to given data.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
cuts (str, optional): String specifying the selection. Defaults to "".
|
||
|
|
input_file (str, optional): Defaults to None.
|
||
|
|
outfile_postfix (str, optional): Defaults to "selected".
|
||
|
|
tree_name (str, optional): Defaults to "PrDebugTrackingLosses.PrDebugTrackingTool/Tuple".
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
str: Path to the output file.
|
||
|
|
"""
|
||
|
|
rdf = ROOT.RDataFrame(tree_name, input_file)
|
||
|
|
rdf = rdf.Filter(cuts, "Selection")
|
||
|
|
out_file = input_file.strip(".root") + f"_{outfile_postfix}.root"
|
||
|
|
rdf.Snapshot("Selected", out_file)
|
||
|
|
return out_file
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
import argparse
|
||
|
|
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument(
|
||
|
|
"--input-file",
|
||
|
|
type=str,
|
||
|
|
help="Path to the input file",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--cuts",
|
||
|
|
type=str,
|
||
|
|
default="pt > 10 && p > 1500 && p < 100000 && isElectron == 1",
|
||
|
|
help="Cuts of the preselection",
|
||
|
|
)
|
||
|
|
args = parser.parse_args()
|
||
|
|
preselection(**vars(args))
|