You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
44 lines
1.3 KiB
44 lines
1.3 KiB
import ROOT
|
|
|
|
|
|
def preselection(
|
|
cuts: str = "",
|
|
input_file: str = None,
|
|
outfile_postfix: str = "selected",
|
|
tree_name: str = "PrParameterisationData.PrMCDebugReconstructibleLong/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 "PrParameterisationData.PrMCDebugReconstructibleLong/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="chi2_comb < 5 && pt > 10 && p > 1500 && p < 100000 && pid != 11",
|
|
help="Cuts of the preselection",
|
|
)
|
|
args = parser.parse_args()
|
|
preselection(**vars(args))
|